text: 
  
  TITLE
  checkpoint breaks with deepspeed 
  
  COMMENTS
  9
  
  REACTIONS
  +1: 0
  -1: 0
  laugh: 0
  hooray: 0
  heart: 0
  rocket: 0
  eyes: 0
  
  BODY
  ## 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 