Fix: Don't tie weights when checkpoint has different values

[diodiogod]

What does this PR do?

Fixes a bug introduced in transformers 4.54.0 where tie_weights() incorrectly ties lm_head to embed_tokens even when the checkpoint contains both weights with different values.

Root Cause

The bug occurs because:

1. Model checkpoint has both lm_head.weight and model.embed_tokens.weight as separate tensors with different values
2. Model config has tie_word_embeddings=True and model code has _tied_weights_keys = ["lm_head.weight"]
3. During from_pretrained(), both weights load correctly from checkpoint
4. Then tie_weights() is called with missing_keys=None (line 4974 in modeling_utils.py)
5. The existing safety check at line 2342-2349 only runs when missing_keys is not None
6. Since missing_keys=None, the check is skipped and weights are tied unconditionally
7. This overwrites the correctly loaded lm_head.weight with embed_tokens.weight

The Fix

Adds an else branch to tie_weights() that checks if both parameters exist with different values when missing_keys=None. If they do, it skips tying and logs a warning.

Impact

This fixes any model where:

• The checkpoint has both lm_head.weight and embedding weights as separate tensors with different values
• The model config has tie_word_embeddings=True
• The model uses _tied_weights_keys

Example affected model: Wan2000/Step-Audio-EditX - this model generates wrong output (text tokens instead of audio tokens) in transformers 4.54+ due to this bug.

Testing

Tested with Step-Audio-EditX model:

Before fix (transformers 4.54-4.57.3):

• ❌ lm_head norm: 255.00 (wrong - should be 227.00)
• ❌ Weights are tied (same memory address)
• ❌ Generates text tokens instead of audio tokens

After fix:

• ✅ lm_head norm: 227.00 (correct)
• ✅ Weights are NOT tied (different memory addresses)
• ✅ Generates audio tokens correctly

Backward Compatibility

This fix is fully backward compatible:

• Models that should have tied weights continue to work (when both weights are identical or one is missing)
• Only affects the edge case where checkpoint explicitly has different values
• Adds a clear warning message to help users update their config

Fixes the regression introduced in #39339

cc @ArthurZucker @gante

[vasqu]

I'm confused why we would need this, doesn't

transformers/src/transformers/modeling_utils.py

Lines 2337 to 2349 in f6331bb

	remove_from_missing = True
	source_is_there = source_param_name not in missing_keys
	target_is_there = target_param_name not in missing_keys
	# Both are already present -> it means the config is wrong and do not reflect the actual
	# checkpoint -> let's raise a warning and do nothing
	if source_is_there and target_is_there:
	logger.warning(
	f"The tied weights mapping and config for this model specifies to tie {source_param_name} to "
	f"{target_param_name}, but both are present in the checkpoints, so we will NOT tie them. "
	"You should update the config with `tie_word_embeddings=False` to silence this warning"
	)
	# Skip to next iteration
	continue

already handle this case where both parameters are present, i.e. we skip and not tie if we find them.

[ArthurZucker]

hey! we are moving to v5, and now we only tie if the users sets tie_word_embeddings = True. so this should not be needed!

[diodiogod]

@vasqu Good question! You're right that lines 2337-2349 should handle this case, but only when missing_keys is not None.

The issue is that tie_weights() is called multiple times:

1. During from_pretrained() with missing_keys (line 4190) - this works correctly
2. Later during initialization with missing_keys=None (line 2928) - this skips the check!

The check at lines 2337-2349 is inside the if missing_keys is not None: block, so when tie_weights() is called with missing_keys=None, the check is skipped entirely.

Our fix adds the same check for the else case (when missing_keys is None). This ensures consistency: "Don't tie weights if both exist with different values" regardless of when tie_weights() is called.

Without this fix, models with wrong tie_word_embeddings=True config but different checkpoint weights get silently broken on the second tie_weights() call.

[diodiogod]

@ArthurZucker Thanks! Good to know v5 will be more explicit about tying.

However, even in v5, if a model config has tie_word_embeddings=True (incorrectly), and transformers follows that config, it would still break models like Step-Audio-EditX where the checkpoint has different weights.

The issue is:

1. Model config says tie_word_embeddings=True (wrong)
2. Checkpoint has lm_head.weight (norm=227) and embed_tokens.weight (norm=255) as separate tensors
3. Following the config breaks the model

This fix adds a safety check: "If both weights exist with different values, don't tie them" regardless of config. This prevents silent corruption when model configs are wrong.

Even in v5, wouldn't it be better to warn/skip rather than silently break a working model because of a wrong config?

The fix is minimal (18 lines) and only affects the edge case where checkpoint explicitly has different values for supposedly-tied weights.

[vasqu]

Thanks for the explanation, makes sense to me now. I overlooked the case with init_weights cc @Cyrilvallez because you initially introduced this

Would be cool if we could add a test here with any "corrupted/wrong" model.

[vasqu]

Would we receive 2 warnings in this case then? Once from the first call with missing keys and once from init weights iiuc - would be nice if we could avoid that

[vasqu]

Hmm, looking at the CI, it seems this PR breaks quite a few models. Can you check?

CI is broken so changes seem incompatible atm

[Cyrilvallez]

Sorry but I don't understand. If we are not using from_pretrained, this is a no issue. We initialize a model from scratch and we SHOULD respect the config.
If we load existing weights using from_pretrained, then we check if we find both weights and infer based on that

Any other later manual (from user) call to tie_weights SHOULD respect the config as well

[diodiogod]

After reflecting on the feedback from @ArthurZucker, @Cyrilvallez, and @vasqu, I've come to realize we were fighting the wrong battle here.

The real issue: The Step-Audio-EditX model has a bug in its config - it sets tie_word_embeddings=True but ships a checkpoint with different weights for lm_head and embed_tokens. Old transformers versions masked this bug by not strictly enforcing tying; newer versions correctly follow the config, which exposes the model's misconfiguration.

The right fix: The model should update its config to tie_word_embeddings=False. It's not transformers' responsibility to detect and work around every model's wrong config.

Our solution: We've implemented a workaround in our project (TTS Audio Suite) that detects this specific case and restores the correct weights after loading. This is the appropriate place for such defensive code.

Thank you all for your time and thoughtful feedback. I'll close this PR and instead report the bug to the Step-Audio-EditX model maintainers.

Apologies for the noise! 🙏