Fix: Don't tie weights when checkpoint has different values
Author: diodiogod
Pull request: yes

## 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

Comments:
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!

@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.

@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.

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

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

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! 🙏