[BUG] gan stage2 training TypeError: 'NoneType' object is not subscriptable

Describe the bug
I tested the training code training/gan on the DeepSpeedExamples repository and it ran normally with the default configuration. However, when I modified the config JSON to stage2 for testing, an error occurred. I don't know why

To Reproduce
Steps to reproduce the behavior:

1. clone and goto https://github.com/microsoft/DeepSpeedExamples/tree/master/training/gan
2. i change gan_deepspeed_config.json
   to stage2

{
   "train_batch_size": 32,
   "train_micro_batch_size_per_gpu": 16,
   "gradient_accumulation_steps": 2,
   "zero_optimization": {
       "stage": 2,
       "offload_optimizer": {
           "device": "cpu"
       },
       "offload_param": {
           "device": "cpu",
           "pin_memory": true
       },
       "overlap_comm": true,
       "contiguous_gradients": true 
   },
   "steps_per_print": 1,
   "optimizer": {
       "type": "Adam",
       "params": {
           "lr": 0.001
       }
   }
}

4. change adam to DeepSpeedCPUAdam, line 106

# setup optimizer
# optimizerD = torch.optim.Adam(netD.parameters(), lr=args.lr, betas=(args.beta1, 0.999))
# optimizerG = torch.optim.Adam(netG.parameters(), lr=args.lr, betas=(args.beta1, 0.999))
from deepspeed.ops.adam import DeepSpeedCPUAdam
optimizerD = DeepSpeedCPUAdam(netD.parameters(), lr=args.lr, betas=(args.beta1, 0.999))
optimizerG = DeepSpeedCPUAdam(netG.parameters(), lr=args.lr, betas=(args.beta1, 0.999))
 

6. See error
   [2024-01-09 16:49:10,135] [INFO] [config.py:988:print] zero_force_ds_cpu_optimizer .. True
   [2024-01-09 16:49:10,135] [INFO] [config.py:988:print] zero_optimization_stage ...... 2
   [2024-01-09 16:49:10,135] [INFO] [config.py:974:print_user_config] json = {
   "train_batch_size": 32,
   "train_micro_batch_size_per_gpu": 16,
   "gradient_accumulation_steps": 2,
   "zero_optimization": {
   "stage": 2,
   "offload_optimizer": {
   "device": "cpu"
   },
   "offload_param": {
   "device": "cpu",
   "pin_memory": true
   },
   "overlap_comm": true,
   "contiguous_gradients": true
   },
   "steps_per_print": 1,
   "optimizer": {
   "type": "Adam",
   "params": {
   "lr": 0.001
   }
   }
   }
   Traceback (most recent call last):
   File "/root/DeepSpeedExamples/training/gan/gan_deepspeed_train.py", line 185, in
   main()
   File "/root/DeepSpeedExamples/training/gan/gan_deepspeed_train.py", line 182, in main
   train(args)
   File "/root/DeepSpeedExamples/training/gan/gan_deepspeed_train.py", line 152, in train
   model_engineG.backward(errG)
   File "/root/miniconda3/lib/python3.10/site-packages/deepspeed/utils/nvtx.py", line 15, in wrapped_fn
   ret_val = func(*args, **kwargs)
   File "/root/miniconda3/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 1955, in backward
   self.optimizer.backward(loss, retain_graph=retain_graph)
   File "/root/miniconda3/lib/python3.10/site-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 2019, in backward
   self.loss_scaler.backward(loss.float(), retain_graph=retain_graph)
   File "/root/miniconda3/lib/python3.10/site-packages/deepspeed/runtime/fp16/loss_scaler.py", line 63, in backward
   scaled_loss.backward(retain_graph=retain_graph)
   File "/root/miniconda3/lib/python3.10/site-packages/torch/_tensor.py", line 492, in backward
   torch.autograd.backward(
   File "/root/miniconda3/lib/python3.10/site-packages/torch/autograd/init.py", line 251, in backward
   Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass
   File "/root/miniconda3/lib/python3.10/site-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 865, in reduce_partition_and_remove_grads
   self.reduce_ready_partitions_and_remove_grads(param, i)
   File "/root/miniconda3/lib/python3.10/site-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 1377, in reduce_ready_partitions_and_remove_grads
   self.reduce_independent_p_g_buckets_and_remove_grads(param, i)
   File "/root/miniconda3/lib/python3.10/site-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 910, in reduce_independent_p_g_buckets_and_remove_grads
   new_grad_tensor = self.ipg_buffer[self.ipg_index].narrow(0, self.elements_in_ipg_bucket, param.numel())
   TypeError: 'NoneType' object is not subscriptable

Expected behavior
It should be running normally

ds_report output

[2024-01-09 17:32:21,320] [INFO] [real_accelerator.py:161:get_accelerator] Setting ds_accelerator to cuda (auto detect)
--------------------------------------------------
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
--------------------------------------------------
 [WARNING]  async_io requires the dev libaio .so object and headers but these were not found.
 [WARNING]  async_io: please install the libaio-dev package with apt
 [WARNING]  If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found.
async_io ............... [NO] ....... [NO]
fused_adam ............. [NO] ....... [OKAY]
cpu_adam ............... [NO] ....... [OKAY]
cpu_adagrad ............ [NO] ....... [OKAY]
cpu_lion ............... [NO] ....... [OKAY]
 [WARNING]  Please specify the CUTLASS repo directory as environment variable $CUTLASS_PATH
evoformer_attn ......... [NO] ....... [NO]
fused_lamb ............. [NO] ....... [OKAY]
fused_lion ............. [NO] ....... [OKAY]
inference_core_ops ..... [NO] ....... [OKAY]
cutlass_ops ............ [NO] ....... [OKAY]
quantizer .............. [NO] ....... [OKAY]
ragged_device_ops ...... [NO] ....... [OKAY]
ragged_ops ............. [NO] ....... [OKAY]
random_ltd ............. [NO] ....... [OKAY]
 [WARNING]  sparse_attn requires a torch version >= 1.5 and < 2.0 but detected 2.1
 [WARNING]  using untested triton version (2.1.0), only 1.0.0 is known to be compatible
sparse_attn ............ [NO] ....... [NO]
spatial_inference ...... [NO] ....... [OKAY]
transformer ............ [NO] ....... [OKAY]
stochastic_transformer . [NO] ....... [OKAY]
transformer_inference .. [NO] ....... [OKAY]
--------------------------------------------------
DeepSpeed general environment info:
torch install path ............... ['/root/miniconda3/lib/python3.10/site-packages/torch']
torch version .................... 2.1.2+cu121
deepspeed install path ........... ['/root/miniconda3/lib/python3.10/site-packages/deepspeed']
deepspeed info ................... 0.12.6, unknown, unknown
torch cuda version ............... 12.1
torch hip version ................ None
nvcc version ..................... 12.1
deepspeed wheel compiled w. ...... torch 2.1, cuda 12.1
shared memory (/dev/shm) size .... 60.00 GB

Screenshots

System info (please complete the following information):

• OS: Ubuntu 22.04

• GPU: 1x 4090

• Python version

• Python 3.10.8

Launcher context
deepspeed gan_deepspeed_train.py --dataset cifar10 --cuda --deepspeed_config gan_deepspeed_config.json --tensorboard_path './runs/deepspeed'

Docker context

Additional context

Reproduce the failure from training/gan using gan_deepspeed_train.py with the shown stage-2 gan_deepspeed_config.json, then inspect the backward call at line 152 and the reported DeepSpeed ZeRO stage 1/2 stack. Compare the working default configuration with stage 2 and verify that training completes normally without the NoneType error.