`push_to_hub()` for Llama 3.1 8B doesn't save `lm_head.weight` tensor
Author: wizeng23
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.49.0
- Platform: Linux-6.8.0-1015-gcp-x86_64-with-glibc2.35
- Python version: 3.10.13
- Huggingface_hub version: 0.30.1
- Safetensors version: 0.5.3
- Accelerate version: 1.2.1
- Accelerate config: 	not found
- DeepSpeed version: not installed
- PyTorch version (GPU?): 2.5.1+cu124 (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 A100-SXM4-40GB

### Who can help?

@ArthurZucker 

### Information

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

### Tasks

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

### Reproduction

```
import torch
import transformers
model = transformers.AutoModel.from_pretrained("meta-llama/Llama-3.1-8B", torch_dtype=torch.bfloat16)
model.push_to_hub('wizeng23/Llama-test')
tokenizer = transformers.AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B")
tokenizer.push_to_hub('wizeng23/Llama-test')
```

### Expected behavior

I'd expect the model weights to be completely unchanged when saving the model. However, it seems the `lm_head.weight` is not saved at all. `model-00004-of-00004.safetensors` for Llama 3.1 8B 1.17GB, while in the saved model, it's 117MB: https://huggingface.co/wizeng23/Llama-test/tree/main. I checked the save tensor file, and the only difference is the missing lm head tensor (shape [128256, 4096]); this is 500M params, which seems to fully account for the missing size.

Comments:
If the weights are tied, we don't save lm head. When loading the model, an `embed_tokens` weight can be used for both: embedding and the head. Is loading back raising warnings like "Some weights are not initialized from ckpt"?

Hi! @wizeng23 , just as @zucchini-nlp mentioned, sometimes tied weights is the reason that lm_head cannot be found in `model.state_dict().keys()`. You may check [this comment](https://github.com/huggingface/transformers/issues/36598#issuecomment-2728111315).

However, for your issue, it seems tied weights is not the cause. You are loading llama 3.1 8B model with `AutoModel.from_pretrained`. As a result, the type of your model is `<class 'transformers.models.llama.modeling_llama.LlamaModel'>`, which does not have `lm_head`. You may want to verify with the following code:
```python
import os
import torch
import transformers

model_path = "meta-llama/Llama-3.1-8B"

llama_model = transformers.AutoModel.from_pretrained(model_path, torch_dtype=torch.bfloat16)

# check if lm_head is in the named_parameters of LlamaForCausalLM
print(any("lm_head" in name for name, _ in llama_model.named_parameters()))
print(type(llama_model))

# save LlamaModel and check the size of the safetensors
llama_model.save_pretrained('./ckpt/llama_model')
print('model-00004-of-00004.safetensors size: {:.2e} bytes'.format(
    os.path.getsize('./ckpt/llama_model/model-00004-of-00004.safetensors')
))

llama_for_causal_lm = transformers.AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.bfloat16)

# check if lm_head is in the named_parameters of LlamaForCausalLM
print(any("lm_head" in name for name, _ in llama_for_causal_lm.named_parameters()))
print([name for name, _ in llama_for_causal_lm.named_parameters() if "lm_head" in name])

# check if lm_head is in the state_dict of LlamaForCausalLM
print(any("lm_head" in key for key in llama_for_causal_lm.state_dict().keys()))
print([key for key in llama_for_causal_lm.state_dict().keys() if "lm_head" in key])

# save LlamaForCausalLM and check the size of the safetensors
llama_for_causal_lm.save_pretrained('./ckpt/llama_for_causal_lm')
print('model-00004-of-00004.safetensors size: {:.2e} bytes'.format(
    os.path.getsize('./ckpt/llama_for_causal_lm/model-00004-of-00004.safetensors')
))

# load LlamaModel and LlamaForCausalLM from the save path
llama_model = transformers.AutoModel.from_pretrained('./ckpt/llama_model', torch_dtype=torch.bfloat16)
llama_for_causal_lm = transformers.AutoModelForCausalLM.from_pretrained('./ckpt/llama_model', torch_dtype=torch.bfloat16)
```
### Expected Results
LlamaModel
- size of `model-00004-of-00004.safetensors` is 117MB
- does not raise a warning when loading from './ckpt/llama_model'

LlamaForCausalLM
- size of `model-00004-of-00004.safetensors` is 1.17GB
- raise a warning when loading from './ckpt/llama_model': `Some weights of LlamaForCausalLM were not initialized from the model checkpoint at ./ckpt/llama_model and are newly initialized: ['lm_head.weight']`




Thanks for your answer @Zephyr271828 ! To summarize, this is user error on my part, and I should be using `AutoModelForCausalLM` instead of `AutoModel`.

> Thanks for your answer [@Zephyr271828](https://github.com/Zephyr271828) ! To summarize, this is user error on my part, and I should be using `AutoModelForCausalLM` instead of `AutoModel`.

Glad it helps!