`Some weights of GPT2LMHeadModel were not initialized` when specifying `device_map=torch.device("cpu")`
Author: samuela
Pull request: no

### System Info

Running
```py
AutoModelForCausalLM.from_pretrained("gpt2")
```
works for me without issue.

Running
```py
AutoModelForCausalLM.from_pretrained("gpt2", device_map=torch.device("cpu"))
```
which to should presumably do the exact same thing, gives me a warning:
```
Some weights of GPT2LMHeadModel were not initialized from the model checkpoint at gpt2 and are newly initialized: ['lm_head.weight']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
```

Why does specifying `device_map=` result in this warning/error? I'm seeing the same thing with `device_map=torch.device("cuda")` as well.

I'm using torch 2.0.1 and transformers 4.31.0.

### Who can help?

@ArthurZucker and @younesbelkada

### Information

- [X] The official example scripts
- [X] 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

Create a shell.nix file with the following contents:
```nix
let
  # Last updated: 2023-09-14. Check for new commits at https://status.nixos.org.
  pkgs = import (fetchTarball "https://github.com/NixOS/nixpkgs/archive/1697b7d480449b01111e352021f46e5879e47643.tar.gz") {
    config.allowUnfree = true;
    config.cudaSupport = true;
  };
in
pkgs.mkShell {
  buildInputs = with pkgs; [
    python3
    python3Packages.accelerate
    python3Packages.torch-bin
    python3Packages.transformers
  ];
}
```
Then run `nix-shell` followed by `python3 -c "import torch; from transformers import AutoModelForCausalLM; AutoModelForCausalLM.from_pretrained('gpt2', device_map=torch.device('cpu'))"`.

### Expected behavior

`device_map` to not result in missing weights warnings. `device_map=torch.device("cpu")` should behave the same as the default behavior.

Comments:
When we look at how GPT-2 is pretrained, we find it consists of a stack of decoders and an LM head (a linear layer and softmax) on top of this stack. When GPT-2 is saved as a pretrained model, typically only the weights for the decoders are saved.

So, when we initialize a model like `GPT2LMHeadModel`, which includes both the decoders and the LM head, the decoders get initialized with their pretrained weights, but the LM head is initialized with brand new, randomly initialized weights. This is why you might see a warning message at times during this initialization process.

Hope it helps. 

Hi @shahad-mahmud, thanks for your response. Perhaps I'm revealing my ignorance of gpt2's internals here, but isn't the "LM head" -- with linear layer and softmax -- necessary for using gpt2 as a causal language model, and in particular doing text generation?

Yes you indeed need the weights for the LM head for text generation and for this you have to fine-tune the `GPT2LMHeadModel` at first. Then using this fine-tuned model you can accomplish the text generation task.

Let's investigate this a bit more.

## The raw weights

At first we will download the GPT-2 pretrained weights and investigate which parameters it contain utilizing the following code:

1. Download the weights.

```bash
wget https://huggingface.co/gpt2/resolve/main/pytorch_model.bin
```

2. Load the weights and look at the keys.

```python
import torch

gpt2_weights = torch.load('pytorch_model.bin')
print(gpt2_weights.keys())
```

It will print the following keys:

```bash
odict_keys(['wte.weight', 'wpe.weight', 'h.0.ln_1.weight', 'h.0.ln_1.bias', 'h.0.attn.bias', 'h.0.attn.c_attn.weight', 'h.0.attn.c_attn.bias', 'h.0.attn.c_proj.weight', 'h.0.attn.c_proj.bias', 'h.0.ln_2.weight', 'h.0.ln_2.bias', 'h.0.mlp.c_fc.weight', 'h.0.mlp.c_fc.bias', 'h.0.mlp.c_proj.weight', 'h.0.mlp.c_proj.bias', 'h.1.ln_1.weight', 'h.1.ln_1.bias', 'h.1.attn.bias', 'h.1.attn.c_attn.weight', 'h.1.attn.c_attn.bias', 'h.1.attn.c_proj.weight', 'h.1.attn.c_proj.bias', 'h.1.ln_2.weight', 'h.1.ln_2.bias', 'h.1.mlp.c_fc.weight', 'h.1.mlp.c_fc.bias', 'h.1.mlp.c_proj.weight', 'h.1.mlp.c_proj.bias', 'h.2.ln_1.weight', 'h.2.ln_1.bias', 'h.2.attn.bias', 'h.2.attn.c_attn.weight', 'h.2.attn.c_attn.bias', 'h.2.attn.c_proj.weight', 'h.2.attn.c_proj.bias', 'h.2.ln_2.weight', 'h.2.ln_2.bias', 'h.2.mlp.c_fc.weight', 'h.2.mlp.c_fc.bias', 'h.2.mlp.c_proj.weight', 'h.2.mlp.c_proj.bias', 'h.3.ln_1.weight', 'h.3.ln_1.bias', 'h.3.attn.bias', 'h.3.attn.c_attn.weight', 'h.3.attn.c_attn.bias', 'h.3.attn.c_proj.weight', 'h.3.attn.c_proj.bias', 'h.3.ln_2.weight', 'h.3.ln_2.bias', 'h.3.mlp.c_fc.weight', 'h.3.mlp.c_fc.bias', 'h.3.mlp.c_proj.weight', 'h.3.mlp.c_proj.bias', 'h.4.ln_1.weight', 'h.4.ln_1.bias', 'h.4.attn.bias', 'h.4.attn.c_attn.weight', 'h.4.attn.c_attn.bias', 'h.4.attn.c_proj.weight', 'h.4.attn.c_proj.bias', 'h.4.ln_2.weight', 'h.4.ln_2.bias', 'h.4.mlp.c_fc.weight', 'h.4.mlp.c_fc.bias', 'h.4.mlp.c_proj.weight', 'h.4.mlp.c_proj.bias', 'h.5.ln_1.weight', 'h.5.ln_1.bias', 'h.5.attn.bias', 'h.5.attn.c_attn.weight', 'h.5.attn.c_attn.bias', 'h.5.attn.c_proj.weight', 'h.5.attn.c_proj.bias', 'h.5.ln_2.weight', 'h.5.ln_2.bias', 'h.5.mlp.c_fc.weight', 'h.5.mlp.c_fc.bias', 'h.5.mlp.c_proj.weight', 'h.5.mlp.c_proj.bias', 'h.6.ln_1.weight', 'h.6.ln_1.bias', 'h.6.attn.bias', 'h.6.attn.c_attn.weight', 'h.6.attn.c_attn.bias', 'h.6.attn.c_proj.weight', 'h.6.attn.c_proj.bias', 'h.6.ln_2.weight', 'h.6.ln_2.bias', 'h.6.mlp.c_fc.weight', 'h.6.mlp.c_fc.bias', 'h.6.mlp.c_proj.weight', 'h.6.mlp.c_proj.bias', 'h.7.ln_1.weight', 'h.7.ln_1.bias', 'h.7.attn.bias', 'h.7.attn.c_attn.weight', 'h.7.attn.c_attn.bias', 'h.7.attn.c_proj.weight', 'h.7.attn.c_proj.bias', 'h.7.ln_2.weight', 'h.7.ln_2.bias', 'h.7.mlp.c_fc.weight', 'h.7.mlp.c_fc.bias', 'h.7.mlp.c_proj.weight', 'h.7.mlp.c_proj.bias', 'h.8.ln_1.weight', 'h.8.ln_1.bias', 'h.8.attn.bias', 'h.8.attn.c_attn.weight', 'h.8.attn.c_attn.bias', 'h.8.attn.c_proj.weight', 'h.8.attn.c_proj.bias', 'h.8.ln_2.weight', 'h.8.ln_2.bias', 'h.8.mlp.c_fc.weight', 'h.8.mlp.c_fc.bias', 'h.8.mlp.c_proj.weight', 'h.8.mlp.c_proj.bias', 'h.9.ln_1.weight', 'h.9.ln_1.bias', 'h.9.attn.bias', 'h.9.attn.c_attn.weight', 'h.9.attn.c_attn.bias', 'h.9.attn.c_proj.weight', 'h.9.attn.c_proj.bias', 'h.9.ln_2.weight', 'h.9.ln_2.bias', 'h.9.mlp.c_fc.weight', 'h.9.mlp.c_fc.bias', 'h.9.mlp.c_proj.weight', 'h.9.mlp.c_proj.bias', 'h.10.ln_1.weight', 'h.10.ln_1.bias', 'h.10.attn.bias', 'h.10.attn.c_attn.weight', 'h.10.attn.c_attn.bias', 'h.10.attn.c_proj.weight', 'h.10.attn.c_proj.bias', 'h.10.ln_2.weight', 'h.10.ln_2.bias', 'h.10.mlp.c_fc.weight', 'h.10.mlp.c_fc.bias', 'h.10.mlp.c_proj.weight', 'h.10.mlp.c_proj.bias', 'h.11.ln_1.weight', 'h.11.ln_1.bias', 'h.11.attn.bias', 'h.11.attn.c_attn.weight', 'h.11.attn.c_attn.bias', 'h.11.attn.c_proj.weight', 'h.11.attn.c_proj.bias', 'h.11.ln_2.weight', 'h.11.ln_2.bias', 'h.11.mlp.c_fc.weight', 'h.11.mlp.c_fc.bias', 'h.11.mlp.c_proj.weight', 'h.11.mlp.c_proj.bias', 'ln_f.weight', 'ln_f.bias'])
```

Here `wte` and `wpe` are weights for word embeddings and positional embeddings respectively. All weights starting with `h` are the decoder blocks. There are 12 of them (0 to 11). Finally `ln_f` is the weights for the linear layer after the decoders.

## Model Modules

Now let's look which modules the `GPT2Model` and `GPT2LMHeadModel` contains.

### `GPT2Model` modules

```python
import transformers

gpt2_configs = transformers.AutoConfig.from_pretrained('gpt2')
gpt2 = transformers.GPT2Model(gpt2_configs)

print(vars(gpt2)['_modules'])
```

Executing it, you will get the following modules.

```bash
OrderedDict([('wte', Embedding(50257, 768)), ('wpe', Embedding(1024, 768)), ('drop', Dropout(p=0.1, inplace=False)), ('h', ModuleList(
  (0-11): 12 x GPT2Block(
    (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
    (attn): GPT2Attention(
      (c_attn): Conv1D()
      (c_proj): Conv1D()
      (attn_dropout): Dropout(p=0.1, inplace=False)
      (resid_dropout): Dropout(p=0.1, inplace=False)
    )
    (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
    (mlp): GPT2MLP(
      (c_fc): Conv1D()
      (c_proj): Conv1D()
      (act): NewGELUActivation()
      (dropout): Dropout(p=0.1, inplace=False)
    )
  )
)), ('ln_f', LayerNorm((768,), eps=1e-05, elementwise_affine=True))])
```

These modules are exactly same as the saved weights in the `pytorch_model.bin` file.

### `GPT2LMHeadModel` modules

```python
import transformers

gpt2_configs = transformers.AutoConfig.from_pretrained('gpt2')
gpt2_lm = transformers.GPT2LMHeadModel(gpt2_configs)

print(vars(gpt2_lm)['_modules'])
```

```bash
OrderedDict([('transformer', GPT2Model(
  (wte): Embedding(50257, 768)
  (wpe): Embedding(1024, 768)
  (drop): Dropout(p=0.1, inplace=False)
  (h): ModuleList(
    (0-11): 12 x GPT2Block(
      (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
      (attn): GPT2Attention(
        (c_attn): Conv1D()
        (c_proj): Conv1D()
        (attn_dropout): Dropout(p=0.1, inplace=False)
        (resid_dropout): Dropout(p=0.1, inplace=False)
      )
      (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
      (mlp): GPT2MLP(
        (c_fc): Conv1D()
        (c_proj): Conv1D()
        (act): NewGELUActivation()
        (dropout): Dropout(p=0.1, inplace=False)
      )
    )
  )
  (ln_f): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
)), ('lm_head', Linear(in_features=768, out_features=50257, bias=False))])
```

So, it has the same modules like `GPT2Model` plus the modules for `lm_head`. But there is no saved weights for this `lm_head` in the `pytorch_model.bin` file. So, you have to fine-tune your `GPT2LMHeadModel` for the text generation. Here is a tutorial from Hugging Face on text generation fine-tuning using GPT-2:  https://huggingface.co/learn/nlp-course/chapter7/6

Hope it helps. Thank you @samuela. 


Hey both! Thanks for opening the issue. It seems to me that you are just not using the latest release as I can' t reproduce this. 
The model on the hub includes no `lm_head` weights as they are most probably tied ( meaning that the lm_head weights are the transposed version of the embedding tokens

Roger that, I'll try updating to latest and see if I can repro.

Just tried on 4.33.2 and the issue is gone, thanks @ArthurZucker !