GPTNeoForCausalLM: resuming Trainer from checkpoint causes Missing key(s) in state_dict: "lm_head.weight"
Author: xusky69
Pull request: no

## Environment info

- `transformers` version: 4.6.0.dev0 (also happens with pip 4.5.1)
- Platform: Linux-4.19.112+-x86_64-with-Ubuntu-18.04-bionic (Google Colab)
- Python version: 3.7.10
- PyTorch version (GPU?): 1.8.1+cu101 (True)
- Tensorflow version (GPU?): Not installed
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: No

### Who can help

- gpt2: @patrickvonplaten, @LysandreJik
- trainer: @sgugger

## Information

Resuming training from a `Trainer` checkpoint for `GPTNeoForCausalLM` causes the following runtime error:
```
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-14-3b03205cdcc2> in <module>()
      2 ### %%%%%%%%%%%%%%%%%%%%%%%% TRAINING %%%%%%%%%%%%%%%%%%%%%%%%% ###
      3 ### %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% ###
----> 4 trainer.train(checkpoint)

1 frames
/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py in load_state_dict(self, state_dict, strict)
   1222         if len(error_msgs) > 0:
   1223             raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
-> 1224                                self.__class__.__name__, "\n\t".join(error_msgs)))
   1225         return _IncompatibleKeys(missing_keys, unexpected_keys)
   1226 

RuntimeError: Error(s) in loading state_dict for GPTNeoForCausalLM:
	Missing key(s) in state_dict: "lm_head.weight". 
```

This happens with the 125M model, havent tested with 1.3b an 2.7b. Loadding the model manually using `.from_pretrained()` and commenting the following lines in `/transformers/trainer.py`

```
             else:
                # We load the model state dict on the CPU to avoid an OOM error.
                state_dict = torch.load(os.path.join(resume_from_checkpoint, WEIGHTS_NAME), map_location="cpu")
                # If the model is on the GPU, it still works!
                self.model.load_state_dict(state_dict)
```

Allows me to resume training.

## To reproduce

Steps to reproduce the behavior:

1. Initialize training via `Trainer` for `GPTNeoForCausalLM` and save a checkpoint
2. Reset env and try to resume training from such checkpoint

<!-- If you have code snippets, error messages, stack traces please provide them here as well.
     Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
     Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.-->

## Expected behavior

For the training to resume correctly

Comments:
@xusky69 , does the fix work for you? I'm still getting an error upon Gpt-neo training:
```bash
...
  File "huggingface/transformers_local/src/transformers/trainer.py", line 1366, in train
    self.model.load_state_dict(state_dict)
  File "huggingface-SJGCx2Wk/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1224, in load_state_dict
    self.__class__.__name__, "\n\t".join(error_msgs)))
RuntimeError: Error(s) in loading state_dict for GPTNeoForCausalLM:
        Missing key(s) in state_dict: "lm_head.weight".
```




The traceback shows you are not using a version of the library that has the fix (the line `self.model.load_state_dict(state_dict)` has been changed in the PR mentioned). Make sure to use a source install or upgrade to the latest release (4.6.0).

@sgugger , thanks for the response! I use the latest version, therefore the script fails in a different place (line 1365 – when the best model is loaded, but the PR fixes initial loading from the checkpoint). I've created a [new one](https://github.com/huggingface/transformers/pull/11718) – could you please take a look?