Repository: microsoft/DeepSpeed

Issue:
[BUG] Error on loading saved optimizer after training (zero-3)
**Describe the bug**
I am training a large model via huggingface trainer and deepspeed. I am using ZeRO-3 (config attached below). The training works fine but once the code tries to reload the saved optimizer for evaluation, I get the following error (below is the traceback).

The train and evaluation code works fine if I train with `python train.py`. OPTIMIZER USED: adafactor

My env:
```
python3.8
cuda-11.3
torch==1.11.0+cu113
transformers==4.18.0
deepspeed==0.6.3
```

ZeRO-3 config
```
{
    "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",
    "steps_per_print": 2000,
    "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
```

**To Reproduce**
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.

**Expected behavior**
The model should be able to reload the optimizer w/o errors.

**ds_report output**
```
--------------------------------------------------
DeepSpeed C++/CUDA extension op report
--------------------------------------------------
NOTE: Ops not installed will be just-in-time (JIT) compiled at
      runtime if needed. Op compatibility means that your system
      meet the required dependencies to JIT install the op.
--------------------------------------------------
JIT compiled ops requires ninja
ninja .................. [OKAY]
--------------------------------------------------
op name ................ installed .. compatible
--------------------------------------------------
cpu_adam ............... [NO] ....... [OKAY]
cpu_adagrad ............ [NO] ....... [OKAY]
fused_adam ............. [NO] ....... [OKAY]
fused_lamb ............. [NO] ....... [OKAY]
 [WARNING]  please install triton==1.0.0 if you want to use sparse attention
sparse_attn ............ [NO] ....... [NO]
transformer ............ [NO] ....... [OKAY]
stochastic_transformer . [NO] ....... [OKAY]
async_io ............... [NO] ....... [OKAY]
utils .................. [NO] ....... [OKAY]
quantizer .............. [NO] ....... [OKAY]
transformer_inference .. [NO] ....... [OKAY]
--------------------------------------------------
DeepSpeed general environment info:
torch install path ............... ['/home/user/venv/lib/python3.8/site-packages/torch']
torch version .................... 1.11.0+cu113
torch cuda version ............... 11.3
torch hip version ................ None
nvcc version ..................... 11.3
deepspeed install path ........... ['/home/user/venv/lib/python3.8/site-packages/deepspeed']
deepspeed info ................... 0.6.3, unknown, unknown
deepspeed wheel compiled w. ...... torch 1.11, cuda 11.3
```

**System info (please complete the following information):**
 - OS: Ubuntu 18.04.1
 - GPU count and types [e.g. 8 machines with NVIDIA v100 each]
 - Python version 3.8

**Launcher context**
Are you launching your experiment with the `deepspeed` launcher, MPI, or something else? `deepspeed train.py config.json`

**Docker context**
Are you using a specific docker image that you can share? No



Discussion:
Also, wondering why does the deepspeed reload for evaluation after training?

```
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/ubuntu/.cache/torch_extensions/py38_cu113 as PyTorch extensions root...
Using /home/ubuntu/.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/ubuntu/.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/ubuntu/.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/ubuntu/.cache/torch_extensions/py38_cu113 as PyTorch extensions root...No modifications detected for re-loaded extension module utils, skipping build step...

Using /home/ubuntu/.cache/torch_extensions/py38_cu113 as PyTorch extensions root...Using /home/ubuntu/.cache/torch_extensions/py38_cu113 as PyTorch extensions root...Loading extension module utils...
```

The above code looks like deepspeed is reloading again. Is this expected and if so, can someone shed some light on whats happening?
@tjruwase I have been scratching my head around this for the last few days. Can you help in this regard? I also tried other optimizers (`adam_hf`, `adam_torch`, etc) but still the same error.
## Script to reproduce the error

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"
}
```
@base-y, sorry for the late response. However, I see that you have raised this issue on the HF side. I am working with them on it. 

Resolving patch:
diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py
--- a/deepspeed/runtime/engine.py
+++ b/deepspeed/runtime/engine.py
@@ -343,6 +343,10 @@ def __init__(
         self.flatten = util_ops.flatten
         self.unflatten = util_ops.unflatten
 
+    def destroy(self):
+        if self.optimizer is not None and hasattr(self.optimizer, 'destroy'):
+            self.optimizer.destroy()
+
     def _get_model_parameters(self):
         if self.autotuning_profile_model_info():
             self.autotuning_model_info = {}
diff --git a/deepspeed/runtime/zero/stage3.py b/deepspeed/runtime/zero/stage3.py
--- a/deepspeed/runtime/zero/stage3.py
+++ b/deepspeed/runtime/zero/stage3.py
@@ -359,7 +359,12 @@ def __init__(self,
 
         self.persistent_parameters = self.persistent_parameters()
 
+        self.forward_hooks = []
+        self.backward_hooks = []
         self.setup_zero_stage3_hooks()
+        print_rank_0(
+            f'Created module hooks: forward = {len(self.forward_hooks)}, backward = {len(self.backward_hooks)}',
+            force=False)
 
         #resetting ds_tensor just in case parameters have been changed after initialization
         #example .half() or .to()
@@ -526,6 +531,23 @@ def __init__(self,
         if dist.get_rank(group=self.dp_process_group) == 0:
             see_memory_usage(f"After initializing ZeRO optimizer", force=True)
 
+    def destroy(self):
+        self._remove_module_hooks()
+
+    def _remove_module_hooks(self):
+        num_forward_hooks = len(self.forward_hooks)
+        num_backward_hooks = len(self.backward_hooks)
+
+        for hook in self.forward_hooks:
+            hook.remove()
+
+        for hook in self.backward_hooks:
+            hook.remove()
+
+        print_rank_0(
+            f'Deleted module hooks: forward = {num_forward_hooks}, backward = {num_backward_hooks}',
+            force=False)
+
     def _setup_for_real_optimizer(self):
         see_memory_usage("Before creating fp32 partitions", force=False)
         self._create_fp32_partitions()
@@ -1201,15 +1223,20 @@ def _run_after_backward_function(sub_module):
                                           inputs)
 
         # Pre forward hook
-        module.register_forward_pre_hook(_pre_forward_module_hook)
+        self.forward_hooks.append(
+            module.register_forward_pre_hook(_pre_forward_module_hook))
+
         # Post forward hook
-        module.register_forward_hook(_post_forward_module_hook)
+        self.forward_hooks.append(
+            module.register_forward_hook(_post_forward_module_hook))
 
         # Pre backward hook
-        module.register_forward_hook(_pre_backward_module_hook)
+        self.backward_hooks.append(
+            module.register_forward_hook(_pre_backward_module_hook))
 
         # post backward hook
-        module.register_forward_pre_hook(_post_backward_module_hook)
+        self.backward_hooks.append(
+            module.register_forward_pre_hook(_post_backward_module_hook))
 
     @torch.no_grad()
     def pre_sub_module_forward_function(self, sub_module):


Test patch: