AttributeError: 'MolmoForCausalLM' object has no attribute 'all_tied_weights_keys'

[Kaihui-intel]

System Info

transformers==5.1.0

Who can help?

No response

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

code

from transformers import AutoModelForCausalLM
model_name = "allenai/Molmo-7B-D-0924"
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype="auto", trust_remote_code=True)

error

../../.local/lib/python3.11/site-packages/transformers/models/auto/auto_factory.py:365: in from_pretrained
    return model_class.from_pretrained(
../../.local/lib/python3.11/site-packages/transformers/modeling_utils.py:4110: in from_pretrained
    load_info = cls._finalize_load_state_dict(model, load_config, load_info)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
../../.local/lib/python3.11/site-packages/transformers/modeling_utils.py:4264: in _finalize_load_state_dict
    model.mark_tied_weights_as_initialized()
../../.local/lib/python3.11/site-packages/transformers/modeling_utils.py:4587: in mark_tied_weights_as_initialized
    for tied_param in self.all_tied_weights_keys.keys():
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = MolmoForCausalLM(
  (model): Molmo(
    (transformer): ModuleDict(
      (wte): Embedding()
      (emb_drop): Dropout(...ropout): Dropout(p=0.0, inplace=False)
      )
      (image_feature_dropout): Dropout(p=0.0, inplace=False)
    )
  )
)
name = 'all_tied_weights_keys'

    def __getattr__(self, name: str) -> Union[Tensor, "Module"]:
        if "_parameters" in self.__dict__:
            _parameters = self.__dict__["_parameters"]
            if name in _parameters:
                return _parameters[name]
        if "_buffers" in self.__dict__:
            _buffers = self.__dict__["_buffers"]
            if name in _buffers:
                return _buffers[name]
        if "_modules" in self.__dict__:
            modules = self.__dict__["_modules"]
            if name in modules:
                return modules[name]
>       raise AttributeError(
            f"'{type(self).__name__}' object has no attribute '{name}'"
        )
E       AttributeError: 'MolmoForCausalLM' object has no attribute 'all_tied_weights_keys'

../../miniforge3/envs/py311/lib/python3.11/site-packages/torch/nn/modules/module.py:1964: AttributeError

Expected behavior

The code should be work.

[Rocketknight1]

cc @Cyrilvallez remote models affected by weight tying PRs maybe?

[lordaarush]

Hey @Rocketknight1, I think i found the problem.
After #42270, all_tied_weights_keys is initialized in post_init(), but remote models (trust_remote_code=True) don't call this properly. I'm thinking of initializing it on-the-fly in _adjust_tied_keys_with_tied_pointers() by detecting tied weights via data pointers, with a safety check in mark_tied_weights_as_initialized(). Can I work on this and submit a PR?

[Rocketknight1]

I'd prefer confirmation from a core maintainer who's worked with this first! cc @Cyrilvallez @ArthurZucker

[Cyrilvallez]

Hey @lordaarush! The issue here is that tied_weights won't be handled properly if post_init is not called, and thus all_tied_weights_keys does not exist. For this reason I fully decided on the time to NOT check for the existence, so that it would crash for remote code that does not call post_init.
This is in my opinion much safer for several reasons, as post_init is very important in general. That way, while inconvenient on the moment, people can very quickly update their remote code to instead call post_init.

[lordaarush]

Thanks for explaining @Cyrilvallez! I understand now that it's by design to enforce proper post_init() usage.

I was just wondering about the user experience - since only model authors can update the remote code on the Hub, for example users are blocked from using Molmo until AllenAI updates it. Is there a way to notify model authors proactively about these issues, or do they typically find out through user reports?

Just trying to understand the workflow for these situations. Thanks!

[Cyrilvallez]

The best would be to directly open a PR for their code on the hub repo! They should in general be happy to merge!

[Cyrilvallez]

(we cannot force merge ourselves of course, even if it's really to "fix" something...)