Repository: huggingface/transformers

Issue:
GPTNeoForCausalLM: resuming Trainer from checkpoint causes Missing key(s) in state_dict: "lm_head.weight"
## 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


Discussion:
@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).

Resolving patch:
diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py
--- a/src/transformers/trainer.py
+++ b/src/transformers/trainer.py
@@ -1059,18 +1059,7 @@ def train(
                 # 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!
-                load_result = self.model.load_state_dict(state_dict, strict=False)
-                if len(load_result.missing_keys) != 0:
-                    if load_result.missing_keys == self.model._keys_to_ignore_on_save:
-                        self.model.tie_weights()
-                    else:
-                        logger.warn(
-                            f"There were missing keys in the checkpoint model loaded: {load_result.missing_keys}."
-                        )
-                if len(load_result.unexpected_keys) != 0:
-                    logger.warn(
-                        f"There were unexpected keys in the checkpoint model loaded: {load_result.unexpected_keys}."
-                    )
+                self._load_state_dict_in_model(state_dict)
 
         # If model was re-initialized, put it on the right device and update self.model_wrapped
         if model_reloaded:
@@ -1363,7 +1352,7 @@ def train(
             # We load the model state dict on the CPU to avoid an OOM error.
             state_dict = torch.load(os.path.join(self.state.best_model_checkpoint, WEIGHTS_NAME), map_location="cpu")
             # If the model is on the GPU, it still works!
-            self.model.load_state_dict(state_dict)
+            self._load_state_dict_in_model(state_dict)
 
             if self.deepspeed:
                 self.deepspeed.load_checkpoint(
@@ -1385,6 +1374,17 @@ def train(
 
         return TrainOutput(self.state.global_step, self._total_loss_scalar / self.state.global_step, metrics)
 
+    def _load_state_dict_in_model(self, state_dict):
+        load_result = self.model.load_state_dict(state_dict, strict=False)
+
+        if len(load_result.missing_keys) != 0:
+            if set(load_result.missing_keys) == set(self.model._keys_to_ignore_on_save):
+                self.model.tie_weights()
+            else:
+                logger.warn(f"There were missing keys in the checkpoint model loaded: {load_result.missing_keys}.")
+        if len(load_result.unexpected_keys) != 0:
+            logger.warn(f"There were unexpected keys in the checkpoint model loaded: {load_result.unexpected_keys}.")
+
     def _maybe_log_save_evaluate(self, tr_loss, model, trial, epoch):
         if self.control.should_log:
             logs: Dict[str, float] = {}


Test patch:
diff --git a/tests/test_modeling_common.py b/tests/test_modeling_common.py
--- a/tests/test_modeling_common.py
+++ b/tests/test_modeling_common.py
@@ -180,7 +180,8 @@ def test_save_load__keys_to_ignore_on_save(self):
                 # Test we can load the state dict in the model, necessary for the checkpointing API in Trainer.
                 load_result = model.load_state_dict(state_dict_saved, strict=False)
                 self.assertTrue(
-                    len(load_result.missing_keys) == 0 or load_result.missing_keys == model._keys_to_ignore_on_save
+                    len(load_result.missing_keys) == 0
+                    or set(load_result.missing_keys) == set(model._keys_to_ignore_on_save)
                 )
                 self.assertTrue(len(load_result.unexpected_keys) == 0)