Support transformers 5: call post_init() and recompute non-persistent buffers

#84

With transformers 5.x, AutoModelForCausalLM.from_pretrained("vikhyatk/moondream2", trust_remote_code=True) either fails to load or loads a model that generates garbage. This PR adds the two hooks transformers 5 expects from model classes to HfMoondream.

1. Loading on CUDA/MPS fails

AttributeError: 'HfMoondream' object has no attribute 'all_tied_weights_keys'. Did you mean: '_tied_weights_keys'?

Transformers 5 sets all_tied_weights_keys and other loading metadata in PreTrainedModel.post_init(), which every model's __init__ is expected to call. HfMoondream.__init__ now calls it.

2. The model loads but generates garbage

Even when loading succeeds (e.g. on CPU), query() returns noise such as "].append].append…". Transformers 5 builds the model on the meta device, and non-persistent buffers aren't in the checkpoint, so they're materialized uninitialized and must be recomputed in _init_weights(). Moondream has two, MoondreamModel.attn_mask and text.freqs_cis. HfMoondream._init_weights() now recomputes both, using the same logic as their constructors.

Transformers maintainers confirmed both requirements for remote code in huggingface/transformers#43883 (post_init) and huggingface/transformers#43644 (non-persistent buffers).

Testing

Loaded with AutoModelForCausalLM.from_pretrained(..., trust_remote_code=True, device_map={"": "mps"}, dtype=torch.float16) on transformers 5.17.0 and 4.57.6. On both, the loaded buffers are identical to an eager (non-meta) build, and query() / caption() return correct answers. Behavior on 4.57.6 is unchanged from the current main.

Ready to merge
This branch is ready to get merged automatically.

Sign up or log in to comment