Unable to resume training from checkpoint on TPU v3-8
Author: finiteautomata
Pull request: no

I'm facing a similar issue as #11326. When trying to resume training from checkpoint on TPUs, it crashes with error message `ValueError: loaded state dict contains a parameter group that doesn't match the size of optimizer's group` 

## Environment info

- `transformers` version: 4.10.0.dev0
- Platform: Linux-4.19.0-17-cloud-amd64-x86_64-with-debian-10.10
- Python version: 3.7.10
- PyTorch version (GPU?): 1.9.0+cu102 (False)
- 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?: No, TPU v3 on GCP
- Using distributed or parallel set-up in script?: yes, `xla_spawn`

### Who can help

@sgugger 

## Information

Model I am using: RoBERTa model 

The problem arises when using:

[* ] my own modified scripts: 

A barely modified version of `run_mlm.py` (check [here](https://gist.github.com/finiteautomata/bef480d508d12e2028fdeae19a92b350))

## To reproduce

Steps to reproduce the behavior:

1. Run `python xla_spawn.py run_mlm.py config.json`
2. Save checkpoint
3. Run again `python xla_spawn.py run_mlm.py config.json` (with `resume_from_checkpoint` set to `true`)

```config.json
{
    "train_dir": "data/tweets/train",
    "eval_dir": "data/tweets/test",

    "pad_to_max_length": true,
    "max_seq_length": 128,
    "tokenize_on_the_fly": true,

    "do_train": true,
    "do_eval": true,
    "seed": 123456,
    "max_steps": 225000,
    "eval_steps": 6000,
    "save_steps": 1500,

    "max_eval_samples": 150000,
    "logging_steps": 200,
    "logging_strategy": "steps",
    "logging_dir": "./logs/",
    "evaluation_strategy": "steps",


    "config_name": "models/twerto-base-uncased",
    "tokenizer_name": "models/twerto-base-uncased",
    "output_dir": "models/twerto-base-uncased-trained",

    "tokenization_batch_size": 81920,

    "weight_decay": 0.01,
    "adam_beta1": 0.9,
    "adam_beta2": 0.98,
    "adam_epsilon": 1e-6,
    "learning_rate": 6e-4,
    "max_grad_norm": 0,

    "warmup_ratio": 0.06,

    "resume_from_checkpoint": true,
    "ignore_data_skip": true,

    "per_device_train_batch_size": 128,
    "per_device_eval_batch_size": 128,
    "gradient_accumulation_steps": 4
}
```

### Error trace

```python
[INFO|trainer.py:1053] 2021-08-05 14:43:49,852 >> Loading model from models/twerto-base-uncased-trained/checkpoint-21000).
Traceback (most recent call last):
  File "bin/run_mlm.py", line 507, in <module>
    main(time.time())
  File "bin/run_mlm.py", line 449, in main
    train_result = trainer.train(resume_from_checkpoint=checkpoint)
  File "/home/jmperez/.cache/pypoetry/virtualenvs/finetune-vs-scratch-gHiQbun3-py3.7/lib/python3.7/site-packages/transformers/trainer.py", line 1153, in train
    self._load_optimizer_and_scheduler(resume_from_checkpoint)
  File "/home/jmperez/.cache/pypoetry/virtualenvs/finetune-vs-scratch-gHiQbun3-py3.7/lib/python3.7/site-packages/transformers/trainer.py", line 1612, in _load_optimizer_and_scheduler
    self.optimizer.load_state_dict(optimizer_state)
  File "/home/jmperez/.cache/pypoetry/virtualenvs/finetune-vs-scratch-gHiQbun3-py3.7/lib/python3.7/site-packages/torch/optim/optimizer.py", line 145, in load_state_dict
    raise ValueError("loaded state dict contains a parameter group "
ValueError: loaded state dict contains a parameter group that doesn't match the size of optimizer's group
```

## Expected behavior

Training resumes from checkpoint


Comments:
It looks like you are not using the `run_mlm` script but a modified version of it, as there are parameters you are passing that are not in this script. Could you share your modified version?

Sure. This is the modified version => https://gist.github.com/finiteautomata/cb1fba94202c1535d2a516eef2215baf

Main changes are that an extra seed and using a custom `IterableDataset`. Running without `xla_spawn.py` seems to be yielding the same error

This is the model configuration (`models/twerto-base-uncased/config.json`)

```json
{
  "attention_probs_dropout_prob": 0.1,
  "bos_token_id": 0,
  "eos_token_id": 2,
  "gradient_checkpointing": false,
  "hidden_act": "gelu",
  "hidden_dropout_prob": 0.1,
  "hidden_size": 768,
  "initializer_range": 0.02,
  "intermediate_size": 3072,
  "layer_norm_eps": 1e-12,
  "max_position_embeddings": 130,
  "model_type": "roberta",
  "num_attention_heads": 12,
  "num_hidden_layers": 12,
  "pad_token_id": 1,
  "position_embedding_type": "absolute",
  "transformers_version": "4.9.1",
  "type_vocab_size": 1,
  "use_cache": true,
  "vocab_size": 30000
}
```

Update: you can repeat this file a couple of times and change `train_dir`/`test_dir` to point to the directory containing it

https://gist.github.com/finiteautomata/38bf8893ad0035e7001653a91a5f7ec3

After some digging, I gathered some extra information of the crash. It seems that the first saved params (`saved_groups[0]["params"]`) contains 77 elements, while the new optimizer has 76. This raises the exception

```python
   1149         if delay_optimizer_creation:
   1150             self.create_optimizer_and_scheduler(num_training_steps=max_steps)
   1151 
   1152         # Check if saved optimizer or scheduler states exist
-> 1153         self._load_optimizer_and_scheduler(resume_from_checkpoint)
   1154 
   1155         # important: at this point:
   1156         # self.model         is the Transformers Model
   1157         # self.model_wrapped is DDP(Transformers Model), Deepspeed(Transformers Model), etc.
   1158 
   1159         # Train!
   1160         num_examples = (
   1161             self.num_examples(train_dataloader) if train_dataset_is_sized else total_train_batch_size * args.max_steps
   1162         )

  /home/jmperez/.cache/pypoetry/virtualenvs/finetune-vs-scratch-gHiQbun3-py3.7/lib/python3.7/site-packages/transformers/trainer.py(1612)_load_optimizer_and_scheduler()
   1602             if is_torch_tpu_available():
   1603                 # On TPU we have to take some extra precautions to properly load the states on the right device.
   1604                 optimizer_state = torch.load(os.path.join(checkpoint, "optimizer.pt"), map_location="cpu")
   1605                 with warnings.catch_warnings(record=True) as caught_warnings:
   1606                     lr_scheduler_state = torch.load(os.path.join(checkpoint, "scheduler.pt"), map_location="cpu")
   1607                 reissue_pt_warnings(caught_warnings)
   1608 
   1609                 xm.send_cpu_data_to_device(optimizer_state, self.args.device)
   1610                 xm.send_cpu_data_to_device(lr_scheduler_state, self.args.device)
   1611                 import ipdb; ipdb.set_trace()
-> 1612                 self.optimizer.load_state_dict(optimizer_state)
   1613                 self.lr_scheduler.load_state_dict(lr_scheduler_state)
   1614             else:
   1615                 map_location = "cpu" if is_sagemaker_mp_enabled() else self.args.device
   1616                 self.optimizer.load_state_dict(
   1617                     torch.load(os.path.join(checkpoint, "optimizer.pt"), map_location=map_location)
   1618                 )
   1619                 with warnings.catch_warnings(record=True) as caught_warnings:
   1620                     self.lr_scheduler.load_state_dict(torch.load(os.path.join(checkpoint, "scheduler.pt")))
   1621                 reissue_pt_warnings(caught_warnings)

> /home/jmperez/.cache/pypoetry/virtualenvs/finetune-vs-scratch-gHiQbun3-py3.7/lib/python3.7/site-packages/torch/optim/optimizer.py(144)load_state_dict()
    134         state_dict = deepcopy(state_dict)
    135         # Validate the state_dict
    136         groups = self.param_groups
    137         saved_groups = state_dict['param_groups']
    138 
    139         if len(groups) != len(saved_groups):
    140             raise ValueError("loaded state dict has a different number of "
    141                              "parameter groups")
    142         param_lens = (len(g['params']) for g in groups)
    143         saved_lens = (len(g['params']) for g in saved_groups)
--> 144         if any(p_len != s_len for p_len, s_len in zip(param_lens, saved_lens)):
    145             raise ValueError("loaded state dict contains a parameter group "
    146                              "that doesn't match the size of optimizer's group")
    147 
    148         # Update the state
    149         id_map = {old_id: p for old_id, p in
    150                   zip(chain.from_iterable((g['params'] for g in saved_groups)),
    151                       chain.from_iterable((g['params'] for g in groups)))}
    152 
    153         def cast(param, value):

ipdb> len(saved_groups[1]["params"])
127
ipdb> len(groups[1]["params"])
127
ipdb> len(saved_groups[0]["params"])
77
ipdb> len(groups[0]["params"])
76
ipdb> saved_groups[0]["params"]
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76]
ipdb> type(groups[0]["params"])
<class 'list'>
ipdb> groups[0]["params"][0]
Parameter containing:
tensor([[ 0.0166,  0.0198,  0.0155,  ...,  0.0200, -0.0159,  0.0022],
        [ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],
        [ 0.0469,  0.0027, -0.0038,  ...,  0.0280,  0.0718,  0.0199],
        ...,
        [ 0.0048,  0.0189, -0.0068,  ..., -0.0642, -0.0060,  0.0320],
        [-0.0138, -0.0080,  0.0119,  ...,  0.0585, -0.0214, -0.0042],
        [ 0.0244,  0.0121, -0.0498,  ..., -0.0162, -0.0110, -0.0159]],
       device='xla:1', requires_grad=True)

```

What's the architecture used? It could be a model that adds some parameters during training for some reason (`twerto-base-uncased-trained` does not help me ;-) )

It is a `RobertaForMaskedLM`

I changed and re-ran everything changing the `transformers_version` (I noticed there was a mismatch between environment version and the one in the config file) and adding the architecture with no success.

```json
{
    "architectures": [
      "RobertaForMaskedLM"
    ],
    "attention_probs_dropout_prob": 0.1,
    "bos_token_id": 0,
    "classifier_dropout": null,
    "eos_token_id": 2,
    "gradient_checkpointing": false,
    "hidden_act": "gelu",
    "hidden_dropout_prob": 0.1,
    "hidden_size": 768,
    "initializer_range": 0.02,
    "intermediate_size": 3072,
    "layer_norm_eps": 1e-12,
    "max_position_embeddings": 130,
    "model_type": "roberta",
    "num_attention_heads": 12,
    "num_hidden_layers": 12,
    "pad_token_id": 1,
    "position_embedding_type": "absolute",
    "torch_dtype": "float32",
    "transformers_version": "4.10.0.dev0",
    "type_vocab_size": 1,
    "use_cache": true,
    "vocab_size": 30000
  }
```

Ok, it seems that `RoBERTa` is the problem here. Changing

```json
{
    "architectures": [
      "BertForMaskedLM"
    ],
   # ...
    "model_type": "bert",
  }
```

enables checkpoint recovery. 

Is this problem in my config or is this a bug?

Good to know it's specific to Roberta! I think it may be due to some parameter being created dynamically during training. Will investigate more tomorrow.

More digging: when trying to reload the checkpoint, it seems that the missing parameter name (that is, the one that the new optimizer is not willing to load) is `['lm_head.decoder.weight']`. 

A Colab notebook reproducing the error on TPU, without any custom script

https://colab.research.google.com/drive/1GvOktm36m3Q43KWLv681QU8VydubOTAQ?usp=sharing

This notebook is barely the same but using GPUs, and it works

https://colab.research.google.com/drive/1GMUgpSNIAdGTk9mOk6ua5pACj-Qlgs9v?usp=sharing

So the problem is `RoBERTa`+TPUs

I have taken a deep dive into this issue, and it made me discover that all the weight tying in Transformers was thrown into the bins the moment the model is placed on an XLA device, which is why your state dict comports more tensors than your model expects.

#13030 should fix the issue.

Great work @sgugger. I can confirm that the notebook in the last comment now reloads the checkpoint. 

Just one extra question: should I use a saved checkpoint with the previous code or is it now useless? I'm not sure if there was a problem during training too

You should definitely start from scratch (sorry) as your previous trainings don't have the proper weights for the decoder (they are not saved since they are supposed to be the same as the embeddings so you can't even retrieve them)

Great! Thanks again