Serious issue with `device_map='balanced'` on GPT-2
Author: eric-mitchell
Pull request: no

### System Info

- `transformers` version: 4.29.1
- Platform: Linux-5.4.0-153-generic-x86_64-with-glibc2.29
- Python version: 3.8.10
- Huggingface_hub version: 0.16.4
- Safetensors version: 0.3.1
- PyTorch version (GPU?): 2.0.1+cu117 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: no

### Who can help?

@ArthurZucker @younesbelkada 

### 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

Loading GPT-2 with `device_map='balanced'` silently fails to load the pre-trained parameters for the LM head. On transformers `4.29.1`, there is no warning; if I upgrade to `4.31.0`, there is a warning that the LM head is not using pre-trained weights:
```
Some weights of GPT2LMHeadModel were not initialized from the model checkpoint at gpt2-xl 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.
```

The problem seems to be hardware-specific (a dual-A100 machine does **not** show this problem, but a dual-A6000 machine does).

Repro:

```
import transformers

m = transformers.GPT2LMHeadModel.from_pretrained("gpt2-xl", device_map='balanced', cache_dir='/scr/em7')
tokenizer = transformers.AutoTokenizer.from_pretrained("gpt2-xl", cache_dir='/scr/em7')

inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
inputs['labels'] = inputs.input_ids.clone()
outputs = m(**inputs)
loss = outputs.loss
print(loss)
```

Running this twice on a dual-A6000 machine with transformers 4.29.1, I get
```
tensor(10.8342, grad_fn=<ToCopyBackward0>)
```
both times. On a dual-A100 machine, I get the expected value of
```
tensor(4.0143, grad_fn=<ToCopyBackward0>)
```

Simply removing the `device_map='balanced'` argument produces the correct output in all cases. If I had to guess, I'd say there is a bug in `PretrainedModel.from_pretrained()` and/or `PretrainedModel._load_pretrained_model`.

### Expected behavior

The loss should be the same (`tensor(4.0143, grad_fn=<ToCopyBackward0>)`) on all machines and model configurations.

Comments:
cc @sgugger (especially the user says `Simply removing the device_map='balanced' argument produces the correct output in all cases`)

cc @SunMarc who was investigating the same thing on another (or maybe the same?) model.

Hi @eric-mitchell , thanks for reporting. Concerning the warning in `4.31.0`, this is solved in the latest version `4.32.0` with this [PR](https://github.com/huggingface/transformers/pull/25101).
For the hardware issue, can you try with the latest version `4.32.0` ? I don't have a A6000 right now but I will try to reproduce this issue asap. To summarize, on a A6000 + transformers `4.29.0` , you get `tensor(10.8342, grad_fn=<ToCopyBackward0>) `with `device_map="auto"` and `tensor(4.0143, grad_fn=<ToCopyBackward0>) `without `device_map`, is that right ? 

@SunMarc That's right, removing `device_map='balanced'` fixes the problem.

I re-ran my repro snippet with `4.32.0` on a dual-A6000 machine, and the problem seems resolved (I get the expected output, same as A100).

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.