Bug: lm_head weight not tied in Mistral3ForConditionalGeneration (affects AutoModelForImageTextToText)
Author: aswin00000
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 5.0.0.dev0
- Platform: Linux-5.15.133+-x86_64-with-glibc2.35
- Python version: 3.12.0
- PyTorch version: 2.9.0+cu126
- CUDA/cuDNN version: 12.6
- GPU: Tesla T4 (compute capability 7.5)

### Who can help?

@zucchini-nlp 
@ArthurZucker 
@amyeroberts

### Information

- [x] The official example scripts
- [ ] My own modified scripts

### Tasks

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

### Reproduction

### Issue Description

When loading `Mistral3ForConditionalGeneration` via `AutoModelForImageTextToText` for model `mistralai/Ministral-3-3B-Instruct-2512`, the `lm_head.weight` parameter is not properly tied to input embeddings, causing random initialization and gibberish outputs.

### Model Loading Warning
```
Mistral3ForConditionalGeneration LOAD REPORT from: mistralai/Ministral-3-3B-Instruct-2512
Key            | Status  | 
---------------+---------+-
lm_head.weight | MISSING | 

Notes:
- MISSING: those params were newly initialized because missing from the checkpoint.
```

### Code to Reproduce
```python
from transformers import AutoModelForImageTextToText, AutoProcessor
import torch

# Load model
model = AutoModelForImageTextToText.from_pretrained(
    "mistralai/Ministral-3-3B-Instruct-2512",
    device_map="auto",
    dtype=torch.bfloat16,
    trust_remote_code=True
)

processor = AutoProcessor.from_pretrained(
    "mistralai/Ministral-3-3B-Instruct-2512",
    trust_remote_code=True
)

# Simple test
messages = [{"role": "user", "content": "What is 2+2?"}]
text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = processor(text=text, return_tensors="pt").to(model.device)

outputs = model.generate(**inputs, max_new_tokens=50, do_sample=False)
response = processor.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)

print(f"Response: {response}")
```

### Actual Output (Gibberish)
```
Response: klasse код 분위 francês Хри maréchal Río ormai**』（وامل ingenierobuבו Miguelbu complejobu complejo حب Descriptionbu complejo...
```

### Root Cause

The checkpoint file (`model.safetensors`) contains only `language_model.embed_tokens.weight` and not `lm_head.weight`, as these weights are meant to be tied (shared). The automatic weight tying mechanism in `Mistral3ForConditionalGeneration` is not working.

**Verification:**
```python
from safetensors import safe_open

with safe_open("model.safetensors", framework="pt") as f:
    keys = list(f.keys())
    lm_head_keys = [k for k in keys if 'lm_head' in k]
    print(f"LM head keys: {lm_head_keys}")  # Output: []
```

### Workaround

Manually tying the weights fixes the issue:
```python
# After loading the model:
embed_tokens = model.model.language_model.embed_tokens.weight
model.lm_head.weight = torch.nn.Parameter(embed_tokens)

# Now generation works correctly
# Response: "2 + 2 equals **4**."  ✓
```

### Proposed Fix

Add proper weight tying to `transformers/models/mistral3/modeling_mistral3.py`:
```python
def tie_weights(self):
    """Tie the weights between input embeddings and output embeddings."""
    super().tie_weights()
    if hasattr(self.model, 'language_model') and hasattr(self.model.language_model, 'embed_tokens'):
        self._tie_or_clone_weights(self.lm_head, self.model.language_model.embed_tokens)
```

### Additional Context

- Affects all Ministral-3 models (3B, 8B, 14B variants)
- Breaks fine-tuning workflows as model starts with corrupted weights
- Model structure shows:
  - `model.language_model.embed_tokens.weight`: Shape `[131072, 3072]` ✓ Properly loaded
  - `lm_head.weight`: Shape `[131072, 3072]` ❌ Randomly initialized
  - These should be the **same tensor** (tied weights)

### Expected behavior

The model should:

1. **Automatically tie weights** during loading: `lm_head.weight` should reference `language_model.embed_tokens.weight`
2. **Generate coherent text** for the test prompt "What is 2+2?"
   - Expected output: `"2 + 2 equals 4"` or similar coherent response
   - Current output: Gibberish (random tokens from corrupted lm_head)

The weight tying should happen automatically in the model's `tie_weights()` method, similar to how other transformer models (BERT, GPT-2, LLaMA) handle this.

Comments:
Nice find...

Good catch, the root reason is the absence of `tie_word_embeddings` with default value in model's config (in https://github.com/huggingface/transformers/blob/main/src/transformers/models/mistral3/configuration_mistral3.py). I am not 100% sure if all mistral VLMs tie weights

The solution is to check official ckpt of mistral VLM on the hub and if all of them tie weights, we need to add `tie_word_embeddings` attr in config with default value of `True`. Would you like to submit a PR?

Thanks for the quick response and identifying the root cause! 
Give me a bit of time to:
1. Check all the official Ministral model checkpoints (3B, 8B, 14B)
2. Verify they all follow the same pattern
3. Prepare and submit the PR with the fix

I'll update here once I have the results!

I've completed a thorough investigation of all Ministral-3 models with verification.

### **Summary Table:**

| Model | Config | Checkpoint | Checkpoint File | Status |
|-------|--------|------------|-----------------|--------|
| Ministral-3-3B-Instruct-2512 | `text_config.tie_word_embeddings = true` | Missing `lm_head` | `model.safetensors` | ❌ Bug |
| Ministral-3-8B-Instruct-2512 | `tie_word_embeddings = false` | Has both weights | `consolidated.safetensors` | ✅ OK |
| Ministral-3-14B-Instruct-2512 | `tie_word_embeddings = false` | Has both weights | `consolidated.safetensors` | ✅ OK |

### **Detailed Findings:**

#### **3B Model (HuggingFace format):**
- ✅ Config: `text_config.tie_word_embeddings = true` (nested in text_config)
- ✅ Checkpoint: Only has `language_model.model.embed_tokens.weight`
- ❌ Checkpoint: Missing `lm_head.weight` (should be tied to embeddings)
- ❌ **Bug:** Model doesn't read `config.text_config.tie_word_embeddings` → doesn't tie weights → gibberish output

#### **8B/14B Models (Mistral native format):**
- ✅ Config: `tie_word_embeddings = false` (top-level)
- ✅ Checkpoint: Has BOTH `tok_embeddings.weight` AND `output.weight`
- ✅ **Working correctly:** No tying needed, both weights present

### **Root Cause:**

The `Mistral3ForConditionalGeneration` class doesn't check `config.text_config.tie_word_embeddings`. It only looks for `config.tie_word_embeddings` (top-level), which doesn't exist in the 3B model config.


PR with this fix. Should I proceed? @zucchini-nlp


Resolved!