Weights not tied when loading `from_pretrained` with a wrapped model

[bryant1410]

System Info

• transformers version: 4.53.3
• Platform: Linux-5.15.134***-x86_64-with-glibc2.35
• Python version: 3.10.18
• Huggingface_hub version: 0.33.5
• Safetensors version: 0.5.3
• Accelerate version: 1.9.0
• Accelerate config: not found
• DeepSpeed version: not installed
• PyTorch version (accelerator?): 2.7.1+cu126 (CUDA)
• Tensorflow version (GPU?): not installed (NA)
• Flax version (CPU?/GPU?/TPU?): not installed (NA)
• Jax version: not installed
• JaxLib version: not installed
• Using distributed or parallel set-up in script?: no
• Using GPU in script?: yes
• GPU type: NVIDIA A100-SXM4-40GB

Who can help?

@ArthurZucker @Cyrilvallez

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

Run:

import transformers

class Config(transformers.PretrainedConfig):
    pass

class Model(transformers.PreTrainedModel):
    config_class = Config

    def __init__(self, config):
        super().__init__(config)
        self.another_model = transformers.AutoModelForSeq2SeqLM.from_pretrained("t5-base")


a = Model(Config())
print(a.another_model.lm_head.weight)
a.save_pretrained("/tmp/abc")

b = Model.from_pretrained("/tmp/abc")
print(b.another_model.lm_head.weight)

and you'll get the output:

Parameter containing:
tensor([[ -0.7539,   0.5977,  -2.4375,  ...,   1.2500,  -0.7891,   3.5156],
        [ 11.3750,  -4.8750,   9.0625,  ...,   4.8438,  14.3750,  -5.7812],
        [-16.6250,  11.0625, -20.8750,  ...,  10.6875,  22.2500,  25.0000],
        ...,
        [  2.2344,   6.7500, -11.0625,  ..., -11.3125,  13.5625,  16.6250],
        [  4.2500,   5.1250, -12.2500,  ..., -11.9375,  13.5000,  17.0000],
        [  4.0625,   6.9688, -12.2500,  ..., -11.3750,  11.9375,  16.6250]],
       requires_grad=True)
Some weights of T5ForConditionalGeneration were not initialized from the model checkpoint at t5-base and are newly initialized: ['lm_head.weight']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
Some weights of Model were not initialized from the model checkpoint at /tmp/abc and are newly initialized: ['another_model.lm_head.weight']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
Parameter containing:
tensor([[0., 0., 0.,  ..., 0., 0., 0.],
        [0., 0., 0.,  ..., 0., 0., 0.],
        [0., 0., 0.,  ..., 0., 0., 0.],
        ...,
        [0., 0., 0.,  ..., 0., 0., 0.],
        [0., 0., 0.,  ..., 0., 0., 0.],
        [0., 0., 0.,  ..., 0., 0., 0.]], requires_grad=True)

Expected behavior

That the last print statement outputs the same as the first one, instead of zeros:

tensor([[ -0.7539,   0.5977,  -2.4375,  ...,   1.2500,  -0.7891,   3.5156],
        [ 11.3750,  -4.8750,   9.0625,  ...,   4.8438,  14.3750,  -5.7812],
        [-16.6250,  11.0625, -20.8750,  ...,  10.6875,  22.2500,  25.0000],
        ...,
        [  2.2344,   6.7500, -11.0625,  ..., -11.3125,  13.5625,  16.6250],
        [  4.2500,   5.1250, -12.2500,  ..., -11.9375,  13.5000,  17.0000],
        [  4.0625,   6.9688, -12.2500,  ..., -11.3750,  11.9375,  16.6250]],
       requires_grad=True)

[bryant1410]

I have investigated why this happens and why it didn't happen before.

The library has been saving in safetensors by default for a couple of years. This version doesn't save tied weights.. Then, the tied weights aren't loaded, but after being tied they get the correct values at least (but still not tied, which is wrong).

low_cpu_mem_usage=True became the default a few months ago. Then, wrapped models (such as Model.another_model in the example) don't load the actual weights but only place the weights in the meta device, and then at the end they are actually loaded and the weights tied. The problem is that tie_weights() is called for the wrapper model, but this one doesn't transitively call tie_weights() for its submodules (but it does call _tie_weights(). The issue is that the following code block before this transitive call won't be called transitively:

transformers/src/transformers/modeling_utils.py

Lines 2950 to 2961 in 3fd456b

	self = getattr(self, self.base_model_prefix)
	tied_weights = self._tie_encoder_decoder_weights(
	self.encoder, self.decoder, self.base_model_prefix, "encoder"
	)
	# Setting a dynamic variable instead of `_tied_weights_keys` because it's a class
	# attributed not an instance member, therefore modifying it will modify the entire class
	# Leading to issues on subsequent calls by different tests or subsequent calls.
	self._dynamic_tied_weights_keys = tied_weights
	for module in self.modules():
	if hasattr(module, "_tie_weights"):
	module._tie_weights()

IMHO the solution would be to call tie_weights transitively instead of _tie_weights.

[Rocketknight1]

Issue seems real and important, but everyone's a little toasted after the GPT-OSS launch! cc @ArthurZucker @Cyrilvallez as core maintainers whenever you're recovered 😅

[Cyrilvallez]

Will have a look probably tomorrow!

[Cyrilvallez]

Hey! I just opened a PR here #39996 to fix it.
However, note that your way of doing is wrong, we never initialize submodels with from_pretrained for a variety of reasons, only with from_config. So you should mostly update your code to be

import transformers
from transformers import AutoConfig

class Config(transformers.PretrainedConfig):
    pass

class Model(transformers.PreTrainedModel):
    config_class = Config

    def __init__(self, config):
        super().__init__(config)
        self.another_model = transformers.AutoModelForSeq2SeqLM.from_config(AutoConfig.from_pretrained("t5-base"))


a = Model(Config())
a.another_model = transformers.AutoModelForSeq2SeqLM.from_pretrained("t5-base")
print(a.another_model.lm_head.weight)
a.save_pretrained("/tmp/abc")

b = Model.from_pretrained("/tmp/abc")
print(b.another_model.lm_head.weight)

where you set the submodel afterwards, before saving it (then later reloading will be much easier)

[bryant1410]

Cool to hear!

we never initialize submodels with from_pretrained for a variety of reasons,

Thanks for the feedback. Can you elaborate on why?

[Cyrilvallez]

Oups sorry for the delay. Well, it's highly inefficient as you will call from_pretrained twice without any reason, which is a costly operation. I.e., you will first run all the loading logic to in the end keep the weights on meta device. This could lead to very surprising behavior somewhere, although I think with the work we put in from_pretrained it seem to now bo robust enough so that it still works (i.e. at the end of the internal from_pretrained, all weights are still on meta, which is required afterwards) 🤗

And if you pass any arguments to the inner from_pretrained, then behavior is completely undefined (i.e. if you try setting a device_map in the inner call or something)