Detect and fix most _init_weights() issues - make it work for composite models

[Cyrilvallez]

What does this PR do?

This is a follow-up of #36963.

This PR makes _init_weights work seamlessly with composite models. Until this point, composite models would only use the _init_weights of the outer-most PreTrainedModel wrapper, leading to errors or skipped modules. Now, sub-models are correctly initialized according to their own _init_weights, without any overhead. This is increasingly important as most recent models are now multimodal.
Without this change, every composite model would have to recurse a second time on all sub-models explicitly in the outer-most _init_weights, which is extremely error prone and inefficient. E.g., we would need to do one or the other of the following in the outer-most _init_weights:

# FIRST BAD OPTION

def _init_weights(self, module):
    std = self.config.initializer_range
       
    # for each module in the model, check the whole module list of the submodel (very inefficient)
    if module in self.vision_tower.modules():
        self.vision_tower._init_weights(module)

    # similar for the other sub-model
    elif module in self.language_model.modules():
        self.language_model._init_weights(module)

    # usual init block for only the modules external to the sub-models
    elif isinstance(module, nn.Linear):
        ...

# OR EQUALLY INEFFICIENT

def _init_weights(self, module):
    std = self.config.initializer_range
       
    # Here, as `apply` is depth-first graph traversal, every module will be initialized a first time, then re-initialized
    # a second time (extremely inefficient as well)
    if module is self.vision_tower:
        self.vision_tower.apply(self.vision_tower._init_weights)

    # similar for the other sub-model
    elif module is self.language_model:
        self.language_model.apply(self.language_model._init_weights)

    # usual init block for only the modules external to the sub-models
    elif isinstance(module, nn.Linear):
        ...

This PR allows to simply do

def _init_weights(self, module):
    std = self.config.initializer_range

    # usual init block for only the modules external to the sub-models
    if isinstance(module, nn.Linear):
        ...

and have all submodels correctly initialized automatically.

Also, enforce torch.no_grad() for initialization, which was not the case before and would slow down the process.

Finally, fix the _init_weights of a LOT of models, the most important ones (the most recent ones, and the ones with the flag _supports_cache_class=True) for now.
The reason not to do them all is simply that there are too much to fix. Almost all models in the library have broken _ init_weights 🙃
We'll patch incrementally. In the meantime, the added test will enforce that new models are correct.

[github-actions]

Hi 👋, thank you for opening this pull request! The pull request is converted to draft by default. The CI will be paused while the PR is in draft mode. When it is ready for review, please click the Ready for review button (at the bottom of the PR page). This will assign reviewers and trigger CI.

[ydshieh]

Before a huge refactorization, could you review this one and hopefully we can merge it ?🙏

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[Cyrilvallez]

cc @ArthurZucker, I highlighted most important changes. Other changes are just the _init_weights fixing on all the models

[Cyrilvallez]

This is the most important change to review @ArthurZucker. It's the most efficient and elegant way to handle it, as we only need to traverse modules once. However, it requires to hot-patch torch.nn.Module, which is a bummer but fine IMO.
Other options to avoid doing so all require to traverse the modules several times (at least 2 times) which is less efficient.

[Cyrilvallez]

Because there are too many models to patch otherwise, I'm just enforcing it for most recent/most important ones for now. Having the test like this will enforce that new models have correct init schemes, while we patch older models. Just using copyright date as a proxy of added time as it's the easiest way that came to mind

[Cyrilvallez]

run-slow: llama, mistral, mistral3, qwen2_5_vl

[github-actions]

This comment contains run-slow, running the specified jobs: This comment contains run-slow, running the specified jobs:

models: ['models/llama', 'models/mistral', 'models/mistral3', 'models/qwen2_5_vl']
quantizations: [] ...

[Cyrilvallez]

Slow tests are similar as main, other tests are hub timeouts and flaky test for Qwen2.5 Omni
Merging

[ArthurZucker]

LGTM, not sure how much more efficient it is, but it should be a lot! 🤗