checkpoint breaks with deepspeed 
## 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.3.3
- Platform: linux
- Python version: 3.7
- 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: @LysandreJik

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:

- nlp datasets: [different repo](https://github.com/huggingface/nlp)
- 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

 -->
deepspeed: @stas00

## Information
Dear @stas00 
Having your permission, I opened up this bug, you are really my only hope with this issue and I truly appreciate your help. Thank you very much.

I am using mt5 model, I modified it with adding adapters layers. The problem arises when:
* loading checkpoints from the model trained with deepspeed 

The tasks I am working on is:
* paraphrase detection using paws-x dataset on mt5 model 
 
## To reproduce

Steps to reproduce the behavior:

```
git clone git@github.com:dorost1234/codes.git
conda create --name deepspeed python=3.7
conda activate deepspeed
conda install pytorch torchvision torchaudio cudatoolkit=11.1 -c pytorch -c conda-forge
python setup.py develop
pip install deepspeed 
```


running the codes:
```
deepspeed run_seq2seq.py  configs/test.json
```

I save a checkpoint every 10 steps, the output would look like the below:

```
After first checkpoint I kill the codes, here is the output:
onfiguration saved in outputs/checkpoint-10/config.json                                                                                                                                       
Model weights saved in outputs/checkpoint-10/pytorch_model.bin
[2021-03-20 15:18:45,897] [INFO] [logging.py:60:log_dist] [Rank 0] Saving model checkpoint: outputs/checkpoint-10/global_step10/mp_rank_00_model_states.pt
[2021-03-20 15:18:51,783] [INFO] [engine.py:1680:_save_zero_checkpoint] zero checkpoint saved outputs/checkpoint-10/global_step10/zero_pp_rank_0_mp_rank_00optim_states.pt
Configuration saved in outputs/config.json
Model weights saved in outputs/pytorch_model.bin
```

Then, I contunue training with running the command again:
```
deepspeed run_seq2seq.py  configs/test.json
```

once loading the checkpoint, it cannot load it with deepspeed:
```
successfully loaded 1 ZeRO state_dicts for rank 0
Traceback (most recent call last):
  File "run_seq2seq.py", line 512, in <module>
    main()
  File "run_seq2seq.py", line 476, in main
    train_result = trainer.train(resume_from_checkpoint=checkpoint)
  File "/users/dara/dev/debug_codes/seq2seq/third_party/trainers/trainer.py", line 780, in train
    self._load_optimizer_and_scheduler(resume_from_checkpoint)
  File "/users/dara/dev/debug_codes/seq2seq/third_party/trainers/trainer.py", line 1169, in _load_optimizer_and_scheduler
    self.deepspeed.load_checkpoint(checkpoint, load_optimizer_states=True, load_lr_scheduler_states=True)
  File "/users/dara/libs/anaconda3/envs/deepspeed/lib/python3.7/site-packages/deepspeed/runtime/engine.py", line 1416, in load_checkpoint
    load_optimizer_states=load_optimizer_states)
  File "/users/dara/libs/anaconda3/envs/deepspeed/lib/python3.7/site-packages/deepspeed/runtime/engine.py", line 1488, in _load_zero_checkpoint
    load_from_fp32_weights=self.zero_load_from_fp32_weights())
  File "/users/dara/libs/anaconda3/envs/deepspeed/lib/python3.7/site-packages/deepspeed/runtime/zero/stage2.py", line 1844, in load_state_dict
    self._restore_base_optimizer_state(state_dict_list)
  File "/users/dara/libs/anaconda3/envs/deepspeed/lib/python3.7/site-packages/deepspeed/runtime/zero/stage2.py", line 1805, in _restore_base_optimizer_state
    self.optimizer.state[p][key].data.copy_(saved.data)
RuntimeError: The size of tensor a (302612288) must match the size of tensor b (129296512) at non-singleton dimension 0
Killing subprocess 23829
Traceback (most recent call last):
  File "/users/dara/libs/anaconda3/envs/deepspeed/lib/python3.7/runpy.py", line 193, in _run_module_as_main
    "__main__", mod_spec)
  File "/users/dara/libs/anaconda3/envs/deepspeed/lib/python3.7/runpy.py", line 85, in _run_code
    exec(code, run_globals)
  File "/users/dara/libs/anaconda3/envs/deepspeed/lib/python3.7/site-packages/deepspeed/launcher/launch.py", line 171, in <module>
    main()
  File "/users/dara/libs/anaconda3/envs/deepspeed/lib/python3.7/site-packages/deepspeed/launcher/launch.py", line 161, in main
    sigkill_handler(signal.SIGTERM, None)  # not coming back
  File "/users/dara/libs/anaconda3/envs/deepspeed/lib/python3.7/site-packages/deepspeed/launcher/launch.py", line 139, in sigkill_handler
    raise subprocess.CalledProcessError(returncode=last_return_code, cmd=cmd)
subprocess.CalledProcessError: Command '['/users/dara/anaconda3/envs/deepspeed/bin/python', '-u', 'run_seq2seq.py', '--local_rank=0', 'configs/test.json']' returned non-zero exit status 1.
```



## Expected behavior

being able to continue training from the saved checkpoints 
I'm able to reproduce your issue, @dorost1234 

So this particular error happens because you change something in the model structure after the checkpoint is saved and before it's resumed. For example at the resume point your model is different than what it was just before you saved the checkpoint.

The reason you encountered this problem with deepspeed and did not without it, is because deepspeed by default saves the optimizer state and resumes from it. So here is a quick hack to overcome this in short-term:

```
--- a/seq2seq/third_party/trainers/trainer.py
+++ b/seq2seq/third_party/trainers/trainer.py
@@ -1166,7 +1166,7 @@ class Trainer:

         if self.deepspeed:
             # Not sure how to check if there is a saved deepspeed checkpoint, but since it just return None if it fails to find a deepspeed checkpoint this is sort of a check-n-load function
-            self.deepspeed.load_checkpoint(checkpoint, load_optimizer_states=True, load_lr_scheduler_states=True)
+            self.deepspeed.load_checkpoint(checkpoint, load_optimizer_states=False, load_lr_scheduler_states=False)

     def hyperparameter_search(
         self,

diff --git a/seq2seq/ds_config.json b/seq2seq/ds_config.json
index 18ce5a3..44b5cc0 100644
--- a/seq2seq/ds_config.json
+++ b/seq2seq/ds_config.json
@@ -15,7 +15,8 @@
        "reduce_scatter": true,
        "reduce_bucket_size": 2e8,
        "contiguous_gradients": true,
-       "cpu_offload": false
+       "cpu_offload": false,
+       "load_from_fp32_weights": false
    },

    "zero_allow_untested_optimizer": true,

```

Basically, we are telling deepspeed not to resume the optimizer/scheduler states and ignore the fp32 weights as well. The last one is not great and may not do what you want, as it's likely to impact the precision.

I added some debug code to deepspeed and the key it fails on in your traceback is "exp_avg" . If it helps here is what I did:
```
diff --git a/deepspeed/runtime/zero/stage2.py b/deepspeed/runtime/zero/stage2.py
index e0ca4f0..02d904b 100755
--- a/deepspeed/runtime/zero/stage2.py
+++ b/deepspeed/runtime/zero/stage2.py
@@ -1802,7 +1802,11 @@ class FP16_DeepSpeedZeroOptimizer(object):
             p = group['params'][0]
             for key, saved in base_optimizer_group_states[i].items():
                 if torch.is_tensor(self.optimizer.state[p][key]):
-                    self.optimizer.state[p][key].data.copy_(saved.data)
+                    try:
+                        self.optimizer.state[p][key].data.copy_(saved.data)
+                    except:
+                        print(f"failed with key={key}")
+                        raise
                 else:
                     self.optimizer.state[p][key] = saved
```

------------

Also: unrelated to this particular issue, the code base your forked from is quite outdated and many deepspeed-related bug fixes were applied since then, so I highly recommend that you sync your fork with the current trainer. If possible try to subclass the Trainer rather than hacking it directly, so you always get the most up-to-date code base you inherit from. Deepspeed is very fresh and expect many more changes happening in the next few weeks/months.

-----------

Now to actually solving the problem.

Study your code and see where your model's structure gets modified between its creation and the checkpoint saving. Trainer tries to resume from the checkpoint as soon as `train()` starts and it appears that at that point the model is different structurally (dimensions are different) then it is later when it gets saved. So you need your model to be in an identical shape during resume and saving points.

Please let me know if this helps.

I'd attack it as simple as dumping the model's param dimensions just before the checkpoint is loaded and just before it's saved, comparing the two - most likely finding the mismatch and then going forward from loading or backward from saving in the code and finding a spot where you modify the model. Then move the modification to before you resume from the checkpoint. i.e. before you call `train()`.

----------

a minor comment - your repro instructions are 95% working, I had to do a few fixes to make it work (e.g. your `setup.py` is slightly broken), so it's always a good idea to re-validate that it's actually reproducible ;) But it's all good, I figured it out. It was very helpful to have what you provided to reproduce the problem.<|||||>Dear @stas00 
Thank you very much for taking your precious time to look into this issue to assist me, I am indebted to you, and for all the incredible job you do. About reproducibility, I very honestly created a new condo environment and tested it before sending it out, and was working on my side, please accept my sincere apologies for any shortcomings and if I missed something without realizing it. 

I will investigate the issue with the great pointer you shared with me and I will keep this issue updated.

Thank you very much again for the great help. <|||||>Dear @stas00 
I will be indebted to also ask also about the nans I get with deepspeed just with the same code you run the loss is nan. I very much appreciate any suggestion you might have for me to try to resolve the issue I face with deepspeed, I am using mt5-small. It would be a great help to me if I could use the great work you have done in deepspeed in huggingface repo and overcome the nan issue. Thank you.<|||||>can we close this one now? we are dealing with nans at https://github.com/huggingface/transformers/pull/10956/files<|||||>Dear Stas
I unfortunately could not still figure this out. I am a bit confused by
where fp16 casting is applied in MT5, specially with the new PR, where this
is disabled in forward path. To me, in the middle there is some fp16
casting, which causes this, but I could not figure this out so far, I was
wondering if you could give me more time. I appreciate any hints. thanks

On Tue, Mar 30, 2021 at 12:52 AM Stas Bekman ***@***.***>
wrote:

> can we close this one now? we are dealing with nans at
> https://github.com/huggingface/transformers/pull/10956/files
>
> —
> You are receiving this because you were mentioned.
> Reply to this email directly, view it on GitHub
> <https://github.com/huggingface/transformers/issues/10821#issuecomment-809771691>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AS37NMTDDH2FJGNOVF6ELVDTGEAELANCNFSM4ZQQFMIA>
> .
>
<|||||>I had a closer look and  - you have a very dated trainer code with multiple bugs that have been fixed since then - you will need to update your code base to `transformers` master. I'm pretty sure that once you had it synced this problem will no longer be there.

Please do update me if this is still not the case and then we will fix it then. Thank you!

<|||||>Dear @stas00 
Thank you very much for the help. Much appreciated.
I upgraded the codes to the last version of the codes in huggingface repository and I am still having the same issue. 
I will make an updated the repository asap and keep you updated on this.
Thank you very much.<|||||>Yes, let me know when you have a repo I can reproduce the issue with. Thank you.<|||||>Hi @stas00 
I finally found this bug, this is the issue reported also here https://github.com/huggingface/transformers/issues/11294 
I was freezing some parameters, and during checkpoiting, since huggingface codes does not handle the freezed parameters properly and has a bug currently regarding this, those parameters were not freezed when its loads from the checkpoints, and caused the difference in the number of parameters for the deepspeed.
thanks a lot for all the hints and helps on this.