how to load the weight of decoder.embed_tokens.weight seperately from the shared weight?
Author: SoSongzhi
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.46.3
- Platform: Linux-6.8.0-49-generic-x86_64-with-glibc2.17
- Python version: 3.8.20
- Huggingface_hub version: 0.26.2
- Safetensors version: 0.4.5
- Accelerate version: 1.0.1
- Accelerate config: 	not found
- PyTorch version (GPU?): 2.4.1+cu121 (True)
- 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?: <fill in>
- Using GPU in script?: <fill in>
- GPU type: NVIDIA RTX A4000


### Who can help?

@ArthurZucker @muellerzr  @SunMarc


### Information

- [ ] The official example scripts
- [X] 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

when i use t5 1.1 on seq2seq task, which has 59744 source vocab size and only 32 target vocab size. And To correctly use softmax to calculate each token's probality and score on 32 candidates so I  set model.lm_head as below: 
```python
torch.nn.Linear(config.d_model,target_vocab_size=32,bias=False). 

And everything looks good when the model is training. But after training, I load the safetensor as below:
checkpoint_path = "./resultstest/checkpoint-100"
config = T5Config.from_pretrained("./onlychangelmhead/checkpoint-100/config.json")
model = T5ForConditionalGeneration(config)
model.lm_head = torch.nn.Linear(config.d_model,target_vocab_size,bias=False)
state_dict = load_file(f"{checkpoint_path}/model.safetensors")
model.load_state_dict(state_dict, strict=True) 
```

And the issue comes as:
```
Traceback (most recent call last):
  File "bs_based_on_massdic_failed.py", line 110, in <module>
    model.load_state_dict(state_dict, strict=True)
  File "/home/zhi/anaconda3/envs/peptide_completion/lib/python3.8/site-packages/torch/nn/modules/module.py", line 2215, in load_state_dict
    raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
RuntimeError: Error(s) in loading state_dict for T5ForConditionalGeneration:
	Missing key(s) in state_dict: "encoder.embed_tokens.weight". 
	size mismatch for decoder.embed_tokens.weight: copying a param with shape torch.Size([32, 768]) from checkpoint, the shape in current model is torch.Size([59744, 768]).
```

when I try to print the safetensors' shape it shows that the `lm_head. weight` looks fine as size of `[32, 768]`, but with no `decoder.embeded_tokens` or the way I load the safetensor can not load the embeded_tokens's weight from shared weight properly(I guess). So how can I fix that problem to correctly feat this model on my exact target vocab size as 32 but not same as the source vocab's size. It would be very appreciate if you can reply. Best.

    

### Expected behavior

Use t5 1.1 to feat on 32 target vocab size task. And load the safetensor properly.

Comments:
I would advise you to first load correctly the model then modify the nn.Linear if needed. Also why are you using T5ForConditionalGeneration with a 32 target vocab size ? The generation will be poor. 





> I would advise you to first load correctly the model then modify the nn.Linear if needed. Also why are you using T5ForConditionalGeneration with a 32 target vocab size ? The generation will be poor.

Hi, First thanks for your reply. Because I change the lmhead first and when I entirly load the checkpoint as an instance. It has a bug tell me the mismatch of lmhead, because that lmhead is not included in config. So how can i first load the model then modify the nn.Linear? Does it mean that I don't change the lmhead when training?
 @SunMarc

To modify the linear you should be using `_get_resized_lm_head` ! 
https://github.com/huggingface/transformers/blob/ab0cad6d43df04dd9b7ab876fc41a18fd90bcef0/src/transformers/modeling_utils.py#L2267-L2300

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.