Error on loading saved optimizer after training (zero-3)
### System Info

```shell
Platform: Ubuntu 18.04.1
python3.8.0
cuda-11.3
torch==1.11.0+cu113 (GPU)
transformers==4.18.0
deepspeed==0.6.3
huggingface_hub version: 0.5.1
```


### Who can help?

@sgugger, @patrickvonplaten 

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

Use any code which trains and evaluates a model with huggingface trainer (e.g https://github.com/ElementAI/picard/blob/main/seq2seq/run_seq2seq.py#L216-L267). Use `save_steps=1` in config. Train for few epochs and evaluate. An error is thrown when the model is trying to load the optimizer after training.

### OPTIMIZER USED: adafactor (issue also occurs with `adaw_hf`, `adamw_torch`)

ZeRO-3 config (used the same from hf page)
```
{
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        },
        "offload_param": {
            "device": "cpu",
            "pin_memory": true
        },
        "overlap_comm": true,
        "contiguous_gradients": true,
        "sub_group_size": 1e9,
        "reduce_bucket_size": "auto",
        "stage3_prefetch_bucket_size": "auto",
        "stage3_param_persistence_threshold": "auto",
        "stage3_max_live_parameters": 1e9,
        "stage3_max_reuse_distance": 1e9,
        "stage3_gather_16bit_weights_on_model_save": true
    },
    "gradient_accumulation_steps": "auto",
    "gradient_clipping": "auto",
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
}
```


Traceback:

```
Training completed. Do not forget to share your model on huggingface.co/models =)


Loading best model from ./output/checkpoint-12 (score: 70.1).
[2022-05-06 14:15:47,319] [INFO] [logging.py:69:log_dist] [Rank 0] DeepSpeed info: version=0.6.3, git-hash=unknown, git-branch=unknown
[2022-05-06 14:15:47,323] [INFO] [engine.py:278:__init__] DeepSpeed Flops Profiler Enabled: False
[2022-05-06 14:15:47,323] [INFO] [engine.py:1042:_configure_optimizer] Removing param_group that has no 'params' in the client Optimizer
[2022-05-06 14:15:47,323] [INFO] [engine.py:1048:_configure_optimizer] Using client Optimizer as basic optimizer
[2022-05-06 14:15:47,324] [INFO] [engine.py:1064:_configure_optimizer] DeepSpeed Basic Optimizer = Adafactor
[2022-05-06 14:15:47,324] [INFO] [utils.py:52:is_zero_supported_optimizer] Checking ZeRO support for optimizer=Adafactor type=<class 'transformers.optimization.Adafactor'>
[2022-05-06 14:15:47,324] [WARNING] [engine.py:1077:_configure_optimizer] **** You are using ZeRO with an untested optimizer, proceed with caution *****
[2022-05-06 14:15:47,324] [INFO] [logging.py:69:log_dist] [Rank 0] Creating fp16 ZeRO stage 3 optimizer
[2022-05-06 14:15:47,324] [INFO] [engine.py:1362:_configure_zero_optimizer] Initializing ZeRO Stage 3
Using /home/user/.cache/torch_extensions/py38_cu113 as PyTorch extensions root...
Using /home/user/.cache/torch_extensions/py38_cu113 as PyTorch extensions root...
[2022-05-06 14:15:47,325] [INFO] [stage3.py:273:__init__] Reduce bucket size 262144
Using /home/user/.cache/torch_extensions/py38_cu113 as PyTorch extensions root...
No modifications detected for re-loaded extension module utils, skipping build step...
[2022-05-06 14:15:47,325] [INFO] [stage3.py:274:__init__] Allgather bucket size 235929.6
Loading extension module utils...
Using /home/user/.cache/torch_extensions/py38_cu113 as PyTorch extensions root...
No modifications detected for re-loaded extension module utils, skipping build step...
Loading extension module utils...
Using /home/user/.cache/torch_extensions/py38_cu113 as PyTorch extensions root...No modifications detected for re-loaded extension module utils, skipping build step...

Using /home/user/.cache/torch_extensions/py38_cu113 as PyTorch extensions root...Using /home/user/.cache/torch_extensions/py38_cu113 as PyTorch extensions root...Loading extension module utils...


Time to load utils op: 0.0006072521209716797 seconds
No modifications detected for re-loaded extension module utils, skipping build step...
Loading extension module utils...
Time to load utils op: 0.0006115436553955078 seconds
Traceback (most recent call last):
Time to load utils op: 0.0006079673767089844 secondsNo modifications detected for re-loaded extension module utils, skipping build step...
  File "train.py", line 227, in <module>

No modifications detected for re-loaded extension module utils, skipping build step...Loading extension module utils...
No modifications detected for re-loaded extension module utils, skipping build step...
Loading extension module utils...

Time to load utils op: 0.0005943775177001953 secondsLoading extension module utils...

Using /home/user/.cache/torch_extensions/py38_cu113 as PyTorch extensions root...
Traceback (most recent call last):
  File "train.py", line 227, in <module>
Traceback (most recent call last):
  File "train.py", line 227, in <module>
Traceback (most recent call last):
Time to load utils op: 0.0006155967712402344 seconds    
  File "train.py", line 227, in <module>
main()Time to load utils op: 0.0006160736083984375 seconds

Time to load utils op: 0.0006232261657714844 seconds  File "train.py", line 176, in main

No modifications detected for re-loaded extension module utils, skipping build step...
Loading extension module utils...
    main()
        main()train_result = trainer.train(resume_from_checkpoint=checkpoint)
  File "train.py", line 176, in main

Traceback (most recent call last):
  File "train.py", line 176, in main
  File "/home/user/venv/lib/python3.8/site-packages/transformers/trainer.py", line 1546, in train
  File "train.py", line 227, in <module>
    Traceback (most recent call last):
main()Traceback (most recent call last):
  File "train.py", line 227, in <module>

  File "train.py", line 227, in <module>
Time to load utils op: 0.0004525184631347656 seconds  File "train.py", line 176, in main

        train_result = trainer.train(resume_from_checkpoint=checkpoint)train_result = trainer.train(resume_from_checkpoint=checkpoint)

  File "/home/user/venv/lib/python3.8/site-packages/transformers/trainer.py", line 1546, in train
  File "/home/user/venv/lib/python3.8/site-packages/transformers/trainer.py", line 1546, in train
        train_result = trainer.train(resume_from_checkpoint=checkpoint)main()

  File "/home/user/venv/lib/python3.8/site-packages/transformers/trainer.py", line 1546, in train
  File "train.py", line 176, in main
Traceback (most recent call last):
  File "train.py", line 227, in <module>
        main()main()

  File "train.py", line 176, in main
  File "train.py", line 176, in main
    train_result = trainer.train(resume_from_checkpoint=checkpoint)
  File "/home/user/venv/lib/python3.8/site-packages/transformers/trainer.py", line 1546, in train
    train_result = trainer.train(resume_from_checkpoint=checkpoint)    
train_result = trainer.train(resume_from_checkpoint=checkpoint)  File "/home/user/venv/lib/python3.8/site-packages/transformers/trainer.py", line 1546, in train
    
main()
  File "/home/user/venv/lib/python3.8/site-packages/transformers/trainer.py", line 1546, in train
  File "train.py", line 176, in main
    train_result = trainer.train(resume_from_checkpoint=checkpoint)
  File "/home/user/venv/lib/python3.8/site-packages/transformers/trainer.py", line 1546, in train
                deepspeed_engine, optimizer, lr_scheduler = deepspeed_reinit(self)    deepspeed_engine, optimizer, lr_scheduler = deepspeed_reinit(self)    deepspeed_engine, optimizer, lr_scheduler = deepspeed_reinit(self)deepspeed_engine, optimizer, lr_scheduler = deepspeed_reinit(self)
deepspeed_engine, optimizer, lr_scheduler = deepspeed_reinit(self)    
deepspeed_engine, optimizer, lr_scheduler = deepspeed_reinit(self)


  File "/home/user/venv/lib/python3.8/site-packages/transformers/deepspeed.py", line 374, in deepspeed_reinit
  File "/home/user/venv/lib/python3.8/site-packages/transformers/deepspeed.py", line 374, in deepspeed_reinit
deepspeed_engine, optimizer, lr_scheduler = deepspeed_reinit(self)    
  File "/home/user/venv/lib/python3.8/site-packages/transformers/deepspeed.py", line 374, in deepspeed_reinit
  File "/home/user/venv/lib/python3.8/site-packages/transformers/deepspeed.py", line 374, in deepspeed_reinit
  File "/home/user/venv/lib/python3.8/site-packages/transformers/deepspeed.py", line 374, in deepspeed_reinit

deepspeed_engine, optimizer, lr_scheduler = deepspeed_reinit(self)  File "/home/user/venv/lib/python3.8/site-packages/transformers/deepspeed.py", line 374, in deepspeed_reinit

  File "/home/user/venv/lib/python3.8/site-packages/transformers/deepspeed.py", line 374, in deepspeed_reinit
  File "/home/user/venv/lib/python3.8/site-packages/transformers/deepspeed.py", line 374, in deepspeed_reinit
        deepspeed_engine, optimizer, _, lr_scheduler = deepspeed.initialize(**trainer.deepspeed_initialize_kwargs)deepspeed_engine, optimizer, _, lr_scheduler = deepspeed.initialize(**trainer.deepspeed_initialize_kwargs)    

deepspeed_engine, optimizer, _, lr_scheduler = deepspeed.initialize(**trainer.deepspeed_initialize_kwargs)  File "/home/user/venv/lib/python3.8/site-packages/deepspeed/__init__.py", line 119, in initialize
  File "/home/user/venv/lib/python3.8/site-packages/deepspeed/__init__.py", line 119, in initialize

    deepspeed_engine, optimizer, _, lr_scheduler = deepspeed.initialize(**trainer.deepspeed_initialize_kwargs)          File "/home/user/venv/lib/python3.8/site-packages/deepspeed/__init__.py", line 119, in initialize

    deepspeed_engine, optimizer, _, lr_scheduler = deepspeed.initialize(**trainer.deepspeed_initialize_kwargs)deepspeed_engine, optimizer, _, lr_scheduler = deepspeed.initialize(**trainer.deepspeed_initialize_kwargs)    deepspeed_engine, optimizer, _, lr_scheduler = deepspeed.initialize(**trainer.deepspeed_initialize_kwargs)
  File "/home/user/venv/lib/python3.8/site-packages/deepspeed/__init__.py", line 119, in initialize

  File "/home/user/venv/lib/python3.8/site-packages/deepspeed/__init__.py", line 119, in initialize
deepspeed_engine, optimizer, _, lr_scheduler = deepspeed.initialize(**trainer.deepspeed_initialize_kwargs)
  File "/home/user/venv/lib/python3.8/site-packages/deepspeed/__init__.py", line 119, in initialize

          File "/home/user/venv/lib/python3.8/site-packages/deepspeed/__init__.py", line 119, in initialize
engine = DeepSpeedEngine(args=args,  File "/home/user/venv/lib/python3.8/site-packages/deepspeed/__init__.py", line 119, in initialize
engine = DeepSpeedEngine(args=args,    

engine = DeepSpeedEngine(args=args,
  File "/home/user/venv/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 294, in __init__
  File "/home/user/venv/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 294, in __init__
  File "/home/user/venv/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 294, in __init__
    engine = DeepSpeedEngine(args=args,
      File "/home/user/venv/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 294, in __init__
    engine = DeepSpeedEngine(args=args,engine = DeepSpeedEngine(args=args,
    
  File "/home/user/venv/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 294, in __init__
engine = DeepSpeedEngine(args=args,      File "/home/user/venv/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 294, in __init__

engine = DeepSpeedEngine(args=args,  File "/home/user/venv/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 294, in __init__

  File "/home/user/venv/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 294, in __init__
        self._configure_optimizer(optimizer, model_parameters)self._configure_optimizer(optimizer, model_parameters)
        
          File "/home/user/venv/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 1080, in _configure_optimizer
    self._configure_optimizer(optimizer, model_parameters)self._configure_optimizer(optimizer, model_parameters)self._configure_optimizer(optimizer, model_parameters)  File "/home/user/venv/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 1080, in _configure_optimizer
    self._configure_optimizer(optimizer, model_parameters)self._configure_optimizer(optimizer, model_parameters)


self._configure_optimizer(optimizer, model_parameters)

  File "/home/user/venv/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 1080, in _configure_optimizer
  File "/home/user/venv/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 1080, in _configure_optimizer

  File "/home/user/venv/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 1080, in _configure_optimizer
  File "/home/user/venv/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 1080, in _configure_optimizer
  File "/home/user/venv/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 1080, in _configure_optimizer
  File "/home/user/venv/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 1080, in _configure_optimizer
                self.optimizer = self._configure_zero_optimizer(basic_optimizer)self.optimizer = self._configure_zero_optimizer(basic_optimizer)    self.optimizer = self._configure_zero_optimizer(basic_optimizer)self.optimizer = self._configure_zero_optimizer(basic_optimizer)

    self.optimizer = self._configure_zero_optimizer(basic_optimizer)

      File "/home/user/venv/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 1365, in _configure_zero_optimizer
  File "/home/user/venv/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 1365, in _configure_zero_optimizer
self.optimizer = self._configure_zero_optimizer(basic_optimizer)
      File "/home/user/venv/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 1365, in _configure_zero_optimizer
  File "/home/user/venv/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 1365, in _configure_zero_optimizer
self.optimizer = self._configure_zero_optimizer(basic_optimizer)
self.optimizer = self._configure_zero_optimizer(basic_optimizer)  File "/home/user/venv/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 1365, in _configure_zero_optimizer

  File "/home/user/venv/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 1365, in _configure_zero_optimizer

  File "/home/user/venv/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 1365, in _configure_zero_optimizer
  File "/home/user/venv/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 1365, in _configure_zero_optimizer
            optimizer = DeepSpeedZeroOptimizer_Stage3(        optimizer = DeepSpeedZeroOptimizer_Stage3(    optimizer = DeepSpeedZeroOptimizer_Stage3(        
optimizer = DeepSpeedZeroOptimizer_Stage3(optimizer = DeepSpeedZeroOptimizer_Stage3(
optimizer = DeepSpeedZeroOptimizer_Stage3(
optimizer = DeepSpeedZeroOptimizer_Stage3(optimizer = DeepSpeedZeroOptimizer_Stage3(  File "/home/user/venv/lib/python3.8/site-packages/deepspeed/runtime/zero/stage3.py", line 293, in __init__


  File "/home/user/venv/lib/python3.8/site-packages/deepspeed/runtime/zero/stage3.py", line 293, in __init__

  File "/home/user/venv/lib/python3.8/site-packages/deepspeed/runtime/zero/stage3.py", line 293, in __init__


  File "/home/user/venv/lib/python3.8/site-packages/deepspeed/runtime/zero/stage3.py", line 293, in __init__
  File "/home/user/venv/lib/python3.8/site-packages/deepspeed/runtime/zero/stage3.py", line 293, in __init__
  File "/home/user/venv/lib/python3.8/site-packages/deepspeed/runtime/zero/stage3.py", line 293, in __init__
  File "/home/user/venv/lib/python3.8/site-packages/deepspeed/runtime/zero/stage3.py", line 293, in __init__
  File "/home/user/venv/lib/python3.8/site-packages/deepspeed/runtime/zero/stage3.py", line 293, in __init__
            self.dtype = self.optimizer.param_groups[0]['params'][0].dtypeself.dtype = self.optimizer.param_groups[0]['params'][0].dtypeself.dtype = self.optimizer.param_groups[0]['params'][0].dtype
    
    
self.dtype = self.optimizer.param_groups[0]['params'][0].dtype            self.dtype = self.optimizer.param_groups[0]['params'][0].dtypeIndexError
IndexErrorself.dtype = self.optimizer.param_groups[0]['params'][0].dtypeself.dtype = self.optimizer.param_groups[0]['params'][0].dtypeIndexErrorself.dtype = self.optimizer.param_groups[0]['params'][0].dtype
: : 

: IndexError
list index out of rangelist index out of rangeIndexErrorlist index out of range: IndexErrorIndexError

: IndexError
list index out of range: : list index out of range: 
list index out of rangelist index out of range
list index out of range
```

### Expected behavior

```shell
The code should be able to reload the optimizer w/o errors.
```

I don't think DeepSpeed is compatible with the use of Adafactor, cc @stas00 <|||||>@sgugger oh, but I tried using other optimizers too (like `adamw_hf`, `adam_torch`, etc) which also resulted in the same error. <|||||>I can handle this ticket, @sgugger 

@base-y, thank you for the report - how can I reproduce this situation on a small example with public data/model?
<|||||>Hi @stas00 , I havent tried to reproduce the error with public datasets. But I think to reproduce the error, you could take any seq2seq model (like T5-small) and train it on any seq2seq dataset and make sure to have atleast one checkpoint saved before the training ends. I used deepspeed config mentioned above. Here was my thought on how to reproduce the error (mentioned above): 

## Reproduction

Use any code which trains and evaluates a model with huggingface trainer (e.g https://github.com/ElementAI/picard/blob/main/seq2seq/run_seq2seq.py#L216-L267). Use `save_steps=1` in config (any `save_steps` will work as long there was atleast one saved ckpt before the training ends). Train for few epochs and evaluate. An error is thrown when the model is trying to load the optimizer after training.

> Make sure to set the optimizer to ADAFACTOR.

Please let me know if any other information is required.<|||||>This is not how: please provide a way for us to reproduce the problem works.

We unfortunately don't have time to go and figure out things that aren't part of `transformers` directly - so telling us to go and train and write a script, etc. we won't do.

What you need to do is to write a small reproducible example that we can copy and paste to reproduce the problem with.

Once the problem is reproduced then we can start working on fixing it or delegating it to Deepspeed as the problem could be coming from there. But let's understand it here first.

Thanks.<|||||>One other thing I could recommend to try is not to use the  load best model feature but to try to just train for a few steps, and then exit and normally resume from the saved deepspeed checkpoint and see if it works. If it fails then chances are very high you need to report it to the Deepspeed repo instead.
<|||||>I wrote a script which takes in SQUAD dataset and trains for 2 steps (32 batch size). Though the entire dataset is downloaded, I am only loading 10 samples for train and validation set to get the error asap.

The code takes ~1min to run and throw the error.

```
from dataclasses import field, dataclass
from typing import Optional

import datasets
from transformers import T5Tokenizer, T5ForConditionalGeneration, Trainer, HfArgumentParser, TrainingArguments

args_dict = {
  "model_name_or_path": 't5-small',
  "max_source_length": 512,
  "max_target_length": 16,
  "per_gpu_train_batch_size": 8,
  "per_gpu_eval_batch_size": 8,
  "gradient_accumulation_steps": 4,
  "learning_rate": 1e-4,
  "num_train_epochs": 2,
  "do_train": True,
  "do_eval": True,
  "optim": "adafactor",
  "evaluation_strategy": "steps",
  "eval_steps": 1,
  "save_strategy": "steps",
  "save_steps": 1,
  "load_best_model_at_end": True,
  "output_dir": './output',
  "cache_dir": './cache',
  "max_steps": 2,
  "deepspeed": "ds_config.json"
}


@dataclass
class ModelArguments:
    """
    Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.
    """

    model_name_or_path: str = field(
        metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}
    )
    tokenizer_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
    )
    cache_dir: Optional[str] = field(
        default=None, metadata={"help": "Where do you want to store the pretrained models downloaded from s3"}
    )


@dataclass
class DataTrainingArguments:
    """
    Arguments pertaining to what data we are going to input our model for training and eval.
    """
    train_file_path: Optional[str] = field(
        default='train_data.pt',
        metadata={"help": "Path for cached train dataset"},
    )
    valid_file_path: Optional[str] = field(
        default='valid_data.pt',
        metadata={"help": "Path for cached valid dataset"},
    )
    max_len: Optional[int] = field(
        default=512,
        metadata={"help": "Max input length for the source text"},
    )
    target_max_len: Optional[int] = field(
        default=32,
        metadata={"help": "Max input length for the target text"},
    )


def _add_eos_to_examples(example):
    example['input_text'] = 'question: %s  context: %s' % (example['question'], example['context'])
    example['target_text'] = '%s' % example['answers']['text'][0]
    return example


def _convert_to_features(example_batch):
    input_encodings = tokenizer.batch_encode_plus(example_batch['input_text'], pad_to_max_length=True, max_length=512, truncation=True)
    target_encodings = tokenizer.batch_encode_plus(example_batch['target_text'], pad_to_max_length=True, max_length=16, truncation=True)

    encodings = {
        'input_ids': input_encodings['input_ids'],
        'attention_mask': input_encodings['attention_mask'],
        'labels':  target_encodings['input_ids'],
    }

    return encodings


def get_dataset():
    train_dataset = datasets.load_dataset('squad', split='train[10:20]')
    valid_dataset = datasets.load_dataset('squad', split='validation[:10]')

    train_dataset = train_dataset.map(_add_eos_to_examples)
    # map convert_to_features batch wise
    train_dataset = train_dataset.map(_convert_to_features, batched=True)

    valid_dataset = valid_dataset.map(_add_eos_to_examples, load_from_cache_file=False)
    valid_dataset = valid_dataset.map(_convert_to_features, batched=True, load_from_cache_file=False)

    # set the tensor type and the columns which the dataset should return
    columns = ['input_ids', 'attention_mask', 'labels']
    train_dataset.set_format(type='torch', columns=columns)
    valid_dataset.set_format(type='torch', columns=columns)
    return train_dataset, valid_dataset


if __name__ == "__main__":
    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))
    model_args, data_args, training_args = parser.parse_dict(args_dict)

    tokenizer = T5Tokenizer.from_pretrained('t5-small', cache_dir=model_args.cache_dir)
    model = T5ForConditionalGeneration.from_pretrained('t5-small', cache_dir='./cache')

    train_dataset, valid_dataset = get_dataset()
    print('loading done')

    # Initialize our Trainer
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=valid_dataset,
    )

    # Training
    if training_args.do_train:
        print('Training....')
        trainer.train()
        trainer.save_model()

    # Evaluation
    if training_args.do_eval:
        results = {}
        print("*** Evaluate ***")
        eval_output = trainer.evaluate()

        output_eval_file = "./output/eval_results.txt"
        with open(output_eval_file, "w") as writer:
            print("***** Eval results *****")
            for key in sorted(eval_output.keys()):
                print("  %s = %s", key, str(eval_output[key]))
                writer.write("%s = %s\n" % (key, str(eval_output[key])))

        results.update(eval_output)
```

And here is the `ds_config.json` file:
```
{
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        },
        "offload_param": {
            "device": "cpu",
            "pin_memory": true
        },
        "overlap_comm": true,
        "contiguous_gradients": true,
        "sub_group_size": 1e9,
        "reduce_bucket_size": "auto",
        "stage3_prefetch_bucket_size": "auto",
        "stage3_param_persistence_threshold": "auto",
        "stage3_max_live_parameters": 1e9,
        "stage3_max_reuse_distance": 1e9,
        "stage3_gather_16bit_weights_on_model_save": true
    },
    "gradient_accumulation_steps": "auto",
    "gradient_clipping": "auto",
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto"
}
```

Hope you find this useful.<|||||>super! that's very helpful, @base-y 

I'm able to reproduce the failure:

```
$ CUDA_VISIBLE_DEVICES=0 python -m torch.distributed.run --nproc_per_node=1  --master_addr='127.0.0.1' --master_port=9901 test.py
[...]
Traceback (most recent call last):
  File "test.py", line 129, in <module>
    trainer.train()
  File "/mnt/nvme0/code/huggingface/transformers-master/src/transformers/trainer.py", line 1545, in train
    self._load_best_model()
  File "/mnt/nvme0/code/huggingface/transformers-master/src/transformers/trainer.py", line 1608, in _load_best_model
    deepspeed_engine, optimizer, lr_scheduler = deepspeed_reinit(self)
  File "/mnt/nvme0/code/huggingface/transformers-master/src/transformers/deepspeed.py", line 374, in deepspeed_reinit
    deepspeed_engine, optimizer, _, lr_scheduler = deepspeed.initialize(**trainer.deepspeed_initialize_kwargs)
  File "/mnt/nvme0/code/github/00optimize/deepspeed/deepspeed/__init__.py", line 119, in initialize
    engine = DeepSpeedEngine(args=args,
  File "/mnt/nvme0/code/github/00optimize/deepspeed/deepspeed/runtime/engine.py", line 295, in __init__
    self._configure_optimizer(optimizer, model_parameters)
  File "/mnt/nvme0/code/github/00optimize/deepspeed/deepspeed/runtime/engine.py", line 1081, in _configure_optimizer
    self.optimizer = self._configure_zero_optimizer(basic_optimizer)
  File "/mnt/nvme0/code/github/00optimize/deepspeed/deepspeed/runtime/engine.py", line 1366, in _configure_zero_optimizer
    optimizer = DeepSpeedZeroOptimizer_Stage3(
  File "/mnt/nvme0/code/github/00optimize/deepspeed/deepspeed/runtime/zero/stage3.py", line 610, in __init__
    self.dtype = self.optimizer.param_groups[0]['params'][0].dtype
IndexError: list index out of range
```

I will try to analyze this later today or tomorrow.
<|||||>I solved the training part, but run into a new problem in eval.

```
Traceback (most recent call last):
  File "test.py", line 136, in <module>
    eval_output = trainer.evaluate()
  File "/mnt/nvme0/code/huggingface/transformers-master/src/transformers/trainer.py", line 2420, in evaluate
    output = eval_loop(
  File "/mnt/nvme0/code/huggingface/transformers-master/src/transformers/trainer.py", line 2597, in evaluation_loop
    loss, logits, labels = self.prediction_step(model, inputs, prediction_loss_only, ignore_keys=ignore_keys)
  File "/mnt/nvme0/code/huggingface/transformers-master/src/transformers/trainer.py", line 2838, in prediction_step
    loss, outputs = self.compute_loss(model, inputs, return_outputs=True)
  File "/mnt/nvme0/code/huggingface/transformers-master/src/transformers/trainer.py", line 2177, in compute_loss
    outputs = model(**inputs)
  File "/home/stas/anaconda3/envs/py38-pt111/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/mnt/nvme0/code/github/00optimize/deepspeed/deepspeed/utils/nvtx.py", line 11, in wrapped_fn
    return func(*args, **kwargs)
  File "/mnt/nvme0/code/github/00optimize/deepspeed/deepspeed/runtime/engine.py", line 1573, in forward
    loss = self.module(*inputs, **kwargs)
  File "/home/stas/anaconda3/envs/py38-pt111/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1128, in _call_impl
    result = forward_call(*input, **kwargs)
  File "/mnt/nvme0/code/huggingface/transformers-master/src/transformers/models/t5/modeling_t5.py", line 1603, in forward
    encoder_outputs = self.encoder(
  File "/home/stas/anaconda3/envs/py38-pt111/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1128, in _call_impl
    result = forward_call(*input, **kwargs)
  File "/mnt/nvme0/code/huggingface/transformers-master/src/transformers/models/t5/modeling_t5.py", line 1035, in forward
    layer_outputs = layer_module(
  File "/home/stas/anaconda3/envs/py38-pt111/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1128, in _call_impl
    result = forward_call(*input, **kwargs)
  File "/mnt/nvme0/code/huggingface/transformers-master/src/transformers/models/t5/modeling_t5.py", line 670, in forward
    self_attention_outputs = self.layer[0](
  File "/home/stas/anaconda3/envs/py38-pt111/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1128, in _call_impl
    result = forward_call(*input, **kwargs)
  File "/mnt/nvme0/code/huggingface/transformers-master/src/transformers/models/t5/modeling_t5.py", line 575, in forward
    normed_hidden_states = self.layer_norm(hidden_states)
  File "/home/stas/anaconda3/envs/py38-pt111/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1117, in _call_impl
    result = hook(self, input)
  File "/mnt/nvme0/code/github/00optimize/deepspeed/deepspeed/utils/nvtx.py", line 11, in wrapped_fn
    return func(*args, **kwargs)
  File "/mnt/nvme0/code/github/00optimize/deepspeed/deepspeed/runtime/zero/stage3.py", line 1096, in _pre_forward_module_hook
    self.pre_sub_module_forward_function(module)
  File "/home/stas/anaconda3/envs/py38-pt111/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/mnt/nvme0/code/github/00optimize/deepspeed/deepspeed/runtime/zero/stage3.py", line 1214, in pre_sub_module_forward_function
    param_coordinator.fetch_sub_module(sub_module)
  File "/mnt/nvme0/code/github/00optimize/deepspeed/deepspeed/utils/nvtx.py", line 11, in wrapped_fn
    return func(*args, **kwargs)
  File "/home/stas/anaconda3/envs/py38-pt111/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/mnt/nvme0/code/github/00optimize/deepspeed/deepspeed/runtime/zero/partitioned_param_coordinator.py", line 279, in fetch_sub_module
    assert param.ds_status == ZeroParamStatus.AVAILABLE, param.ds_summary()
AssertionError: {'id': 6, 'status': 'INFLIGHT', 'numel': 0, 'ds_numel': 512, 'shape': (0,), 'ds_shape': (512,), 'requires_grad': True, 'grad_shape': None, 'persist': True, 'active_sub_modules': {292}}
```

the training issue has to do with the hack that requires us to re-init the deepspeed engine every time the model is reloaded. So I'm trying a different approach.

You can try this fix meanwhile for the training assert you reported in this issue: https://github.com/huggingface/transformers/pull/17151

Please let me know if that fixes the problem you reported originally. Thanks.

I will keep you posted when I get the whole thing working.<|||||>The 2nd problem is a bug introduced by https://github.com/microsoft/DeepSpeed/pull/1916

So to unblock you use my PR and roll back deepspeed to `deepspeed==0.6.4`

I will report this other issue to Deepspeed.<|||||>So to update, the 2nd problem happens not because of the PR I quoted above but because a new deepspeed engine init, that we are forced to do to overcome this issue: https://github.com/microsoft/DeepSpeed/issues/1612, creates a new set of forward hooks, but it doesn't delete the old ones - and since the model persists hence the breakage. 

So Tunji is going to create a PR that gives us a method to remove the old hooks.

So I will shortly have 2 PRs for you to test with which should solve this problem altogether.<|||||>@base-y, ok - Tunji and I have a working solution - needed work in both repos.

if it's not merged yet, please use these 2 branches:

- https://github.com/microsoft/DeepSpeed/pull/1947
- https://github.com/huggingface/transformers/pull/17151

and it should all work.

Please let me know if you have tried it and it works for you.

If I don't hear back from you I will merge the `transformers` fix once deepspeed merges their side and make a new release.
<|||||>Hey, sorry for the delayed response. Sure, I will install huggingface and deepspeed locally from the PR branches and check if it works asap.<|||||>Hi @stas00 , I reran the training code and there are no errors this time around. Thank you!<|||||>Perfect. Thank you for validating, @base-y 

I will merge the HF PR once the Deepspeed merges their side and makes a new release.

cc: @tjruwase <|||||>FYI, the deepspeed side has been merged, so just waiting for their new release and will then merge the HF PR. <|||||>Hi, I am using Huggingface version = 4.25.1 and Deepspeed = 0.8.0 and still facing this issue. Is there something I am missing?<|||||>Hi, smitanannaware

as a long time has passed since this Issue, please kindly file a new Issue and provide full information about your specific case including most importantly the full traceback and if possible an easy way for us to reproduce it - and please tag me on that issue. Thank you.