Repository: huggingface/transformers

Issue:
serious bug with trainer.py when restarting the training from a checkpoint
## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->

- `transformers` version: 4.5.1
- Platform: Linux
- Python version: 3.8
- PyTorch version (GPU?): 1.8
- Tensorflow version (GPU?): - 
- Using GPU in script?: - 
- Using distributed or parallel set-up in script?: - 

### Who can help
<!-- Your issue will be replied to more quickly if you can figure out the right person to tag with @
 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- albert, bert, xlm: @LysandreJik
- blenderbot, bart, marian, pegasus, encoderdecoder,  t5: @patrickvonplaten, @patil-suraj
- longformer, reformer, transfoxl, xlnet: @patrickvonplaten
- fsmt: @stas00
- funnel: @sgugger
- gpt2: @patrickvonplaten, @LysandreJik
- rag: @patrickvonplaten, @lhoestq
- tensorflow: @Rocketknight1

Library:

- benchmarks: @patrickvonplaten
- deepspeed: @stas00
- ray/raytune: @richardliaw, @amogkam
- text generation: @patrickvonplaten
- tokenizers: @LysandreJik
- trainer: @sgugger
- pipelines: @LysandreJik

Documentation: @sgugger

Model hub:

- for issues with a model report at https://discuss.huggingface.co/ and tag the model's creator.

HF projects:

- datasets: [different repo](https://github.com/huggingface/datasets)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Examples:

- maintained examples (not research project or legacy): @sgugger, @patil-suraj
- research_projects/bert-loses-patience: @JetRunner
- research_projects/distillation: @VictorSanh

 -->

trainer: @sgugger, @patil-suraj

## Information

Hi, I see this serious issue with trainer.py class, let please consider run_translation.py script [1] after you define the model, let freeze the encoder, or wrap the model in a class. So one can modify the model after this line https://github.com/huggingface/transformers/blob/d9c62047a8d75e18d2849d345ab3394875a712ef/examples/seq2seq/run_translation.py#L331 

Then, during the training, one can stop the training, and now would like to continue the training from the place it is stopped, if you print the number of parameters inside trainer.py, right before this line:

https://github.com/huggingface/transformers/blob/d9c62047a8d75e18d2849d345ab3394875a712ef/src/transformers/trainer.py#L1062

like this 
```
for n,p in model.named_parameters():
   if p.requires_grad:
       print(n)
```

what would we see? We see all parameters are there, even  the ones we made frozen, this is a serious bug that if the user modify the model after creation, those modifications are not considered when restarting the training, could you kindly have a look?
thanks 

[1] https://github.com/huggingface/transformers/blob/master/examples/seq2seq/run_translation.py 


## Expected behavior

The user should be able to continue training the modified model as they are modified. 


Discussion:
You can also consider the modification in finetune_trainer.py

https://github.com/huggingface/transformers/blob/1c06240e1b3477728129bb58e7b6c7734bb5074e/examples/legacy/seq2seq/finetune_trainer.py#L233

If you freeze some parameters as done in the line above, those would not be there when you load the model and restarting the training, this is really a serious issue, thanks for the help 

Here is the minimal code to generate this bug, we make a model, we freeze, then we save it (as done in trainer checkpoint), then we load it (as done in train() in trainer.py) and we see if the feezed parameters are there or not

```
from transformers import T5ForConditionalGeneration
from typing import Optional 
import torch
import os

# This is copied from trainer.py
def _save(model, output_dir: Optional[str] = None):
        os.makedirs(output_dir, exist_ok=True)
        print(f"Saving model checkpoint to {output_dir}")
        # Save a trained model and configuration using `save_pretrained()`.
        # They can then be reloaded using `from_pretrained()`
        state_dict = model.state_dict()
        model.save_pretrained(output_dir, state_dict=state_dict)

def print_num_parameters(model):
   for n,p in model.named_parameters():
      if (p.requires_grad):
        print("n ", n) 

def freeze_params(model):
   for n,p in model.named_parameters():
      p.requires_grad = False


model = T5ForConditionalGeneration.from_pretrained("t5-base")
freeze_params(model)
print("#### parameters before saving ####")
print_num_parameters(model)
_save(model, "temp_model")

# Now lets load the model as done in trainer from the checkpoint.
model = model.from_pretrained("temp_model")

# Now lets print the number of parameters
print("#### parameters after saving ####")
print_num_parameters(model)
```

surprisingly, no, the freezed parameters are not freezed anymore after loading the checkpoint:
```
#### parameters before saving ####
Saving model checkpoint to temp_model
#### parameters after saving ####
n  shared.weight
n  encoder.block.0.layer.0.SelfAttention.q.weight
n  encoder.block.0.layer.0.SelfAttention.k.weight
n  encoder.block.0.layer.0.SelfAttention.v.weight
n  encoder.block.0.layer.0.SelfAttention.o.weight
n  encoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight
n  encoder.block.0.layer.0.layer_norm.weight
n  encoder.block.0.layer.1.DenseReluDense.wi.weight
n  encoder.block.0.layer.1.DenseReluDense.wo.weight
n  encoder.block.0.layer.1.layer_norm.weight
n  encoder.block.1.layer.0.SelfAttention.q.weight
n  encoder.block.1.layer.0.SelfAttention.k.weight
n  encoder.block.1.layer.0.SelfAttention.v.weight
n  encoder.block.1.layer.0.SelfAttention.o.weight
n  encoder.block.1.layer.0.layer_norm.weight
n  encoder.block.1.layer.1.DenseReluDense.wi.weight
n  encoder.block.1.layer.1.DenseReluDense.wo.weight
n  encoder.block.1.layer.1.layer_norm.weight
n  encoder.block.2.layer.0.SelfAttention.q.weight
n  encoder.block.2.layer.0.SelfAttention.k.weight
n  encoder.block.2.layer.0.SelfAttention.v.weight
n  encoder.block.2.layer.0.SelfAttention.o.weight
n  encoder.block.2.layer.0.layer_norm.weight
n  encoder.block.2.layer.1.DenseReluDense.wi.weight
n  encoder.block.2.layer.1.DenseReluDense.wo.weight
n  encoder.block.2.layer.1.layer_norm.weight
n  encoder.block.3.layer.0.SelfAttention.q.weight
n  encoder.block.3.layer.0.SelfAttention.k.weight
n  encoder.block.3.layer.0.SelfAttention.v.weight
n  encoder.block.3.layer.0.SelfAttention.o.weight
n  encoder.block.3.layer.0.layer_norm.weight
n  encoder.block.3.layer.1.DenseReluDense.wi.weight
n  encoder.block.3.layer.1.DenseReluDense.wo.weight
n  encoder.block.3.layer.1.layer_norm.weight
n  encoder.block.4.layer.0.SelfAttention.q.weight
n  encoder.block.4.layer.0.SelfAttention.k.weight
n  encoder.block.4.layer.0.SelfAttention.v.weight
n  encoder.block.4.layer.0.SelfAttention.o.weight
n  encoder.block.4.layer.0.layer_norm.weight
n  encoder.block.4.layer.1.DenseReluDense.wi.weight
n  encoder.block.4.layer.1.DenseReluDense.wo.weight
n  encoder.block.4.layer.1.layer_norm.weight
n  encoder.block.5.layer.0.SelfAttention.q.weight
n  encoder.block.5.layer.0.SelfAttention.k.weight
n  encoder.block.5.layer.0.SelfAttention.v.weight
n  encoder.block.5.layer.0.SelfAttention.o.weight
n  encoder.block.5.layer.0.layer_norm.weight
n  encoder.block.5.layer.1.DenseReluDense.wi.weight
n  encoder.block.5.layer.1.DenseReluDense.wo.weight
n  encoder.block.5.layer.1.layer_norm.weight
n  encoder.block.6.layer.0.SelfAttention.q.weight
n  encoder.block.6.layer.0.SelfAttention.k.weight
n  encoder.block.6.layer.0.SelfAttention.v.weight
n  encoder.block.6.layer.0.SelfAttention.o.weight
n  encoder.block.6.layer.0.layer_norm.weight
n  encoder.block.6.layer.1.DenseReluDense.wi.weight
n  encoder.block.6.layer.1.DenseReluDense.wo.weight
n  encoder.block.6.layer.1.layer_norm.weight
n  encoder.block.7.layer.0.SelfAttention.q.weight
n  encoder.block.7.layer.0.SelfAttention.k.weight
n  encoder.block.7.layer.0.SelfAttention.v.weight
n  encoder.block.7.layer.0.SelfAttention.o.weight
n  encoder.block.7.layer.0.layer_norm.weight
n  encoder.block.7.layer.1.DenseReluDense.wi.weight
n  encoder.block.7.layer.1.DenseReluDense.wo.weight
n  encoder.block.7.layer.1.layer_norm.weight
n  encoder.block.8.layer.0.SelfAttention.q.weight
n  encoder.block.8.layer.0.SelfAttention.k.weight
n  encoder.block.8.layer.0.SelfAttention.v.weight
n  encoder.block.8.layer.0.SelfAttention.o.weight
n  encoder.block.8.layer.0.layer_norm.weight
n  encoder.block.8.layer.1.DenseReluDense.wi.weight
n  encoder.block.8.layer.1.DenseReluDense.wo.weight
n  encoder.block.8.layer.1.layer_norm.weight
n  encoder.block.9.layer.0.SelfAttention.q.weight
n  encoder.block.9.layer.0.SelfAttention.k.weight
n  encoder.block.9.layer.0.SelfAttention.v.weight
n  encoder.block.9.layer.0.SelfAttention.o.weight
n  encoder.block.9.layer.0.layer_norm.weight
n  encoder.block.9.layer.1.DenseReluDense.wi.weight
n  encoder.block.9.layer.1.DenseReluDense.wo.weight
n  encoder.block.9.layer.1.layer_norm.weight
n  encoder.block.10.layer.0.SelfAttention.q.weight
n  encoder.block.10.layer.0.SelfAttention.k.weight
n  encoder.block.10.layer.0.SelfAttention.v.weight
n  encoder.block.10.layer.0.SelfAttention.o.weight
n  encoder.block.10.layer.0.layer_norm.weight
n  encoder.block.10.layer.1.DenseReluDense.wi.weight
n  encoder.block.10.layer.1.DenseReluDense.wo.weight
n  encoder.block.10.layer.1.layer_norm.weight
n  encoder.block.11.layer.0.SelfAttention.q.weight
n  encoder.block.11.layer.0.SelfAttention.k.weight
n  encoder.block.11.layer.0.SelfAttention.v.weight
n  encoder.block.11.layer.0.SelfAttention.o.weight
n  encoder.block.11.layer.0.layer_norm.weight
n  encoder.block.11.layer.1.DenseReluDense.wi.weight
n  encoder.block.11.layer.1.DenseReluDense.wo.weight
n  encoder.block.11.layer.1.layer_norm.weight
n  encoder.final_layer_norm.weight
n  decoder.block.0.layer.0.SelfAttention.q.weight
n  decoder.block.0.layer.0.SelfAttention.k.weight
n  decoder.block.0.layer.0.SelfAttention.v.weight
n  decoder.block.0.layer.0.SelfAttention.o.weight
n  decoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight
n  decoder.block.0.layer.0.layer_norm.weight
n  decoder.block.0.layer.1.EncDecAttention.q.weight
n  decoder.block.0.layer.1.EncDecAttention.k.weight
n  decoder.block.0.layer.1.EncDecAttention.v.weight
n  decoder.block.0.layer.1.EncDecAttention.o.weight
n  decoder.block.0.layer.1.layer_norm.weight
n  decoder.block.0.layer.2.DenseReluDense.wi.weight
n  decoder.block.0.layer.2.DenseReluDense.wo.weight
n  decoder.block.0.layer.2.layer_norm.weight
n  decoder.block.1.layer.0.SelfAttention.q.weight
n  decoder.block.1.layer.0.SelfAttention.k.weight
n  decoder.block.1.layer.0.SelfAttention.v.weight
n  decoder.block.1.layer.0.SelfAttention.o.weight
n  decoder.block.1.layer.0.layer_norm.weight
n  decoder.block.1.layer.1.EncDecAttention.q.weight
n  decoder.block.1.layer.1.EncDecAttention.k.weight
n  decoder.block.1.layer.1.EncDecAttention.v.weight
n  decoder.block.1.layer.1.EncDecAttention.o.weight
n  decoder.block.1.layer.1.layer_norm.weight
n  decoder.block.1.layer.2.DenseReluDense.wi.weight
n  decoder.block.1.layer.2.DenseReluDense.wo.weight
n  decoder.block.1.layer.2.layer_norm.weight
n  decoder.block.2.layer.0.SelfAttention.q.weight
n  decoder.block.2.layer.0.SelfAttention.k.weight
n  decoder.block.2.layer.0.SelfAttention.v.weight
n  decoder.block.2.layer.0.SelfAttention.o.weight
n  decoder.block.2.layer.0.layer_norm.weight
n  decoder.block.2.layer.1.EncDecAttention.q.weight
n  decoder.block.2.layer.1.EncDecAttention.k.weight
n  decoder.block.2.layer.1.EncDecAttention.v.weight
n  decoder.block.2.layer.1.EncDecAttention.o.weight
n  decoder.block.2.layer.1.layer_norm.weight
n  decoder.block.2.layer.2.DenseReluDense.wi.weight
n  decoder.block.2.layer.2.DenseReluDense.wo.weight
n  decoder.block.2.layer.2.layer_norm.weight
n  decoder.block.3.layer.0.SelfAttention.q.weight
n  decoder.block.3.layer.0.SelfAttention.k.weight
n  decoder.block.3.layer.0.SelfAttention.v.weight
n  decoder.block.3.layer.0.SelfAttention.o.weight
n  decoder.block.3.layer.0.layer_norm.weight
n  decoder.block.3.layer.1.EncDecAttention.q.weight
n  decoder.block.3.layer.1.EncDecAttention.k.weight
n  decoder.block.3.layer.1.EncDecAttention.v.weight
n  decoder.block.3.layer.1.EncDecAttention.o.weight
n  decoder.block.3.layer.1.layer_norm.weight
n  decoder.block.3.layer.2.DenseReluDense.wi.weight
n  decoder.block.3.layer.2.DenseReluDense.wo.weight
n  decoder.block.3.layer.2.layer_norm.weight
n  decoder.block.4.layer.0.SelfAttention.q.weight
n  decoder.block.4.layer.0.SelfAttention.k.weight
n  decoder.block.4.layer.0.SelfAttention.v.weight
n  decoder.block.4.layer.0.SelfAttention.o.weight
n  decoder.block.4.layer.0.layer_norm.weight
n  decoder.block.4.layer.1.EncDecAttention.q.weight
n  decoder.block.4.layer.1.EncDecAttention.k.weight
n  decoder.block.4.layer.1.EncDecAttention.v.weight
n  decoder.block.4.layer.1.EncDecAttention.o.weight
n  decoder.block.4.layer.1.layer_norm.weight
n  decoder.block.4.layer.2.DenseReluDense.wi.weight
n  decoder.block.4.layer.2.DenseReluDense.wo.weight
n  decoder.block.4.layer.2.layer_norm.weight
n  decoder.block.5.layer.0.SelfAttention.q.weight
n  decoder.block.5.layer.0.SelfAttention.k.weight
n  decoder.block.5.layer.0.SelfAttention.v.weight
n  decoder.block.5.layer.0.SelfAttention.o.weight
n  decoder.block.5.layer.0.layer_norm.weight
n  decoder.block.5.layer.1.EncDecAttention.q.weight
n  decoder.block.5.layer.1.EncDecAttention.k.weight
n  decoder.block.5.layer.1.EncDecAttention.v.weight
n  decoder.block.5.layer.1.EncDecAttention.o.weight
n  decoder.block.5.layer.1.layer_norm.weight
n  decoder.block.5.layer.2.DenseReluDense.wi.weight
n  decoder.block.5.layer.2.DenseReluDense.wo.weight
n  decoder.block.5.layer.2.layer_norm.weight
n  decoder.block.6.layer.0.SelfAttention.q.weight
n  decoder.block.6.layer.0.SelfAttention.k.weight
n  decoder.block.6.layer.0.SelfAttention.v.weight
n  decoder.block.6.layer.0.SelfAttention.o.weight
n  decoder.block.6.layer.0.layer_norm.weight
n  decoder.block.6.layer.1.EncDecAttention.q.weight
n  decoder.block.6.layer.1.EncDecAttention.k.weight
n  decoder.block.6.layer.1.EncDecAttention.v.weight
n  decoder.block.6.layer.1.EncDecAttention.o.weight
n  decoder.block.6.layer.1.layer_norm.weight
n  decoder.block.6.layer.2.DenseReluDense.wi.weight
n  decoder.block.6.layer.2.DenseReluDense.wo.weight
n  decoder.block.6.layer.2.layer_norm.weight
n  decoder.block.7.layer.0.SelfAttention.q.weight
n  decoder.block.7.layer.0.SelfAttention.k.weight
n  decoder.block.7.layer.0.SelfAttention.v.weight
n  decoder.block.7.layer.0.SelfAttention.o.weight
n  decoder.block.7.layer.0.layer_norm.weight
n  decoder.block.7.layer.1.EncDecAttention.q.weight
n  decoder.block.7.layer.1.EncDecAttention.k.weight
n  decoder.block.7.layer.1.EncDecAttention.v.weight
n  decoder.block.7.layer.1.EncDecAttention.o.weight
n  decoder.block.7.layer.1.layer_norm.weight
n  decoder.block.7.layer.2.DenseReluDense.wi.weight
n  decoder.block.7.layer.2.DenseReluDense.wo.weight
n  decoder.block.7.layer.2.layer_norm.weight
n  decoder.block.8.layer.0.SelfAttention.q.weight
n  decoder.block.8.layer.0.SelfAttention.k.weight
n  decoder.block.8.layer.0.SelfAttention.v.weight
n  decoder.block.8.layer.0.SelfAttention.o.weight
n  decoder.block.8.layer.0.layer_norm.weight
n  decoder.block.8.layer.1.EncDecAttention.q.weight
n  decoder.block.8.layer.1.EncDecAttention.k.weight
n  decoder.block.8.layer.1.EncDecAttention.v.weight
n  decoder.block.8.layer.1.EncDecAttention.o.weight
n  decoder.block.8.layer.1.layer_norm.weight
n  decoder.block.8.layer.2.DenseReluDense.wi.weight
n  decoder.block.8.layer.2.DenseReluDense.wo.weight
n  decoder.block.8.layer.2.layer_norm.weight
n  decoder.block.9.layer.0.SelfAttention.q.weight
n  decoder.block.9.layer.0.SelfAttention.k.weight
n  decoder.block.9.layer.0.SelfAttention.v.weight
n  decoder.block.9.layer.0.SelfAttention.o.weight
n  decoder.block.9.layer.0.layer_norm.weight
n  decoder.block.9.layer.1.EncDecAttention.q.weight
n  decoder.block.9.layer.1.EncDecAttention.k.weight
n  decoder.block.9.layer.1.EncDecAttention.v.weight
n  decoder.block.9.layer.1.EncDecAttention.o.weight
n  decoder.block.9.layer.1.layer_norm.weight
n  decoder.block.9.layer.2.DenseReluDense.wi.weight
n  decoder.block.9.layer.2.DenseReluDense.wo.weight
n  decoder.block.9.layer.2.layer_norm.weight
n  decoder.block.10.layer.0.SelfAttention.q.weight
n  decoder.block.10.layer.0.SelfAttention.k.weight
n  decoder.block.10.layer.0.SelfAttention.v.weight
n  decoder.block.10.layer.0.SelfAttention.o.weight
n  decoder.block.10.layer.0.layer_norm.weight
n  decoder.block.10.layer.1.EncDecAttention.q.weight
n  decoder.block.10.layer.1.EncDecAttention.k.weight
n  decoder.block.10.layer.1.EncDecAttention.v.weight
n  decoder.block.10.layer.1.EncDecAttention.o.weight
n  decoder.block.10.layer.1.layer_norm.weight
n  decoder.block.10.layer.2.DenseReluDense.wi.weight
n  decoder.block.10.layer.2.DenseReluDense.wo.weight
n  decoder.block.10.layer.2.layer_norm.weight
n  decoder.block.11.layer.0.SelfAttention.q.weight
n  decoder.block.11.layer.0.SelfAttention.k.weight
n  decoder.block.11.layer.0.SelfAttention.v.weight
n  decoder.block.11.layer.0.SelfAttention.o.weight
n  decoder.block.11.layer.0.layer_norm.weight
n  decoder.block.11.layer.1.EncDecAttention.q.weight
n  decoder.block.11.layer.1.EncDecAttention.k.weight
n  decoder.block.11.layer.1.EncDecAttention.v.weight
n  decoder.block.11.layer.1.EncDecAttention.o.weight
n  decoder.block.11.layer.1.layer_norm.weight
n  decoder.block.11.layer.2.DenseReluDense.wi.weight
n  decoder.block.11.layer.2.DenseReluDense.wo.weight
n  decoder.block.11.layer.2.layer_norm.weight
n  decoder.final_layer_norm.weight
```

Your end example is not surprising: you are re-loading a new model so of course modifications done on the original model are erased.
We'll look at how we can fix this inside the trainer, to load the weights differently instead of using `from_pretrained`.
Hi @sgugger  thank you for the response, my intention from the example was showing that the procedure happening inside the trainer if the user resume training from a modified model. thank you. 
@sgugger  I also see some more issues with the trainer.py when I load the model from the checkpoint, the model which previously was training fine on the GPU gets out of memory issue, there must be a leakge of memory during the loading from a checkpoint, shall I make a separate ticket for this issue? thanks 
If it's because you add frozen parameters previously, that would explain the OOM error.

Resolving patch:
diff --git a/src/transformers/configuration_utils.py b/src/transformers/configuration_utils.py
--- a/src/transformers/configuration_utils.py
+++ b/src/transformers/configuration_utils.py
@@ -271,7 +271,7 @@ def __init__(self, **kwargs):
         self._name_or_path = str(kwargs.pop("name_or_path", ""))
 
         # Drop the transformers version info
-        kwargs.pop("transformers_version", None)
+        self.transformers_version = kwargs.pop("transformers_version", None)
 
         # Additional attributes without default values
         for key, value in kwargs.items():
diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py
--- a/src/transformers/trainer.py
+++ b/src/transformers/trainer.py
@@ -53,9 +53,12 @@
 from torch.utils.data.distributed import DistributedSampler
 from torch.utils.data.sampler import RandomSampler, SequentialSampler
 
+from . import __version__
+from .configuration_utils import PretrainedConfig
 from .data.data_collator import DataCollator, DataCollatorWithPadding, default_data_collator
 from .dependency_versions_check import dep_version_check
 from .file_utils import (
+    CONFIG_NAME,
     WEIGHTS_NAME,
     is_apex_available,
     is_datasets_available,
@@ -997,14 +1000,23 @@ def train(
 
             logger.info(f"Loading model from {resume_from_checkpoint}).")
 
+            if os.path.isfile(os.path.join(resume_from_checkpoint, CONFIG_NAME)):
+                config = PretrainedConfig.from_json_file(os.path.join(resume_from_checkpoint, CONFIG_NAME))
+                checkpoint_version = config.transformers_version
+                if checkpoint_version is not None and checkpoint_version != __version__:
+                    logger.warn(
+                        f"You are resuming training from a checkpoint trained with {checkpoint_version} of "
+                        f"Transformers but your current version is {__version__}. This is not recommended and could "
+                        "yield to errors or unwanted behaviors."
+                    )
+
             if self.deepspeed:
                 # will be resumed in deepspeed_init
                 pass
-            elif isinstance(self.model, PreTrainedModel):
-                self.model = self.model.from_pretrained(resume_from_checkpoint)
-                model_reloaded = True
             else:
-                state_dict = torch.load(os.path.join(resume_from_checkpoint, WEIGHTS_NAME))
+                # 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)
 
         # If model was re-initialized, put it on the right device and update self.model_wrapped
@@ -1281,13 +1293,10 @@ def train(
             logger.info(
                 f"Loading best model from {self.state.best_model_checkpoint} (score: {self.state.best_metric})."
             )
-            if isinstance(self.model, PreTrainedModel):
-                self.model = self.model.from_pretrained(self.state.best_model_checkpoint)
-                if self.place_model_on_device:
-                    self.model = self.model.to(args.device)
-            else:
-                state_dict = torch.load(os.path.join(self.state.best_model_checkpoint, WEIGHTS_NAME))
-                self.model.load_state_dict(state_dict)
+            # 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)
 
             if self.deepspeed:
                 self.deepspeed.load_checkpoint(


Test patch:
diff --git a/tests/test_trainer.py b/tests/test_trainer.py
--- a/tests/test_trainer.py
+++ b/tests/test_trainer.py
@@ -725,6 +725,46 @@ def test_resume_training_with_gradient_accumulation(self):
             self.assertEqual(b, b1)
             self.check_trainer_state_are_the_same(state, state1)
 
+    def test_resume_training_with_frozen_params(self):
+        if torch.cuda.device_count() > 2:
+            # This test will fail for more than 2 GPUs since the batch size will get bigger and with the number of
+            # save_steps, the checkpoint will resume training at epoch 2 or more (so the data seen by the model
+            # won't be the same since the training dataloader is shuffled).
+            return
+        with tempfile.TemporaryDirectory() as tmpdir:
+            trainer = get_regression_trainer(
+                output_dir=tmpdir,
+                train_len=128,
+                per_device_train_batch_size=4,
+                save_steps=5,
+                learning_rate=0.1,
+            )
+            trainer.model.a.requires_grad_(False)
+            trainer.train()
+            (a, b) = trainer.model.a.item(), trainer.model.b.item()
+            state = dataclasses.asdict(trainer.state)
+
+            checkpoint = os.path.join(tmpdir, "checkpoint-5")
+
+            # Reinitialize trainer
+            trainer = get_regression_trainer(
+                output_dir=tmpdir,
+                train_len=128,
+                per_device_train_batch_size=4,
+                save_steps=5,
+                learning_rate=0.1,
+            )
+            trainer.model.a.requires_grad_(False)
+
+            trainer.train(resume_from_checkpoint=checkpoint)
+
+            self.assertFalse(trainer.model.a.requires_grad)
+            (a1, b1) = trainer.model.a.item(), trainer.model.b.item()
+            state1 = dataclasses.asdict(trainer.state)
+            self.assertEqual(a, a1)
+            self.assertEqual(b, b1)
+            self.check_trainer_state_are_the_same(state, state1)
+
     def test_load_best_model_at_end(self):
         total = int(self.n_epochs * 64 / self.batch_size)
         with tempfile.TemporaryDirectory() as tmpdir: