[BUG] Problem in training step with ZERO optimization. in has_overflow_partitioned_grads_serial for j, grad in enumerate(self.averaged_gradients[i]): KeyError: 0

Below is the Traceback (most recent call last):
File "parallel_gpt.py", line 81, in
time_.append(train_gpt2_with_deepspeed(batch_len))
File "parallel_gpt.py", line 72, in train_gpt2_with_deepspeed
model_engine.step()
File "/usr/local/lib/python3.8/dist-packages/deepspeed/runtime/engine.py", line 2116, in step
self._take_model_step(lr_kwargs)
File "/usr/local/lib/python3.8/dist-packages/deepspeed/runtime/engine.py", line 2022, in _take_model_step
self.optimizer.step()
File "/usr/local/lib/python3.8/dist-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 1714, in step
self.check_overflow()
File "/usr/local/lib/python3.8/dist-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 1965, in check_overflow
self._check_overflow(partition_gradients)
File "/usr/local/lib/python3.8/dist-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 1866, in _check_overflow
self.overflow = self.has_overflow(partition_gradients)
File "/usr/local/lib/python3.8/dist-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 1885, in has_overflow
overflow = self.local_overflow if self.cpu_offload else self.has_overflow_partitioned_grads_serial()
File "/usr/local/lib/python3.8/dist-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 1878, in has_overflow_partitioned_grads_serial
for j, grad in enumerate(self.averaged_gradients[i]):
KeyError: 0

Installation step
DS_BUILD_CPU_ADAM=1 DS_BUILD_CPU_ADAGRAD=1 DS_BUILD_FUSEDLAMB=1 DS_BUILD_TRANSFORMER=1 DS_BUILD_ASYNC_IO=1 BUILD_UTILS=1 pip install --no-cache-dir deepspeed -U

ds_report output:
[2023-09-25 14:56:22,579] [INFO] [real_accelerator.py:158: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

async_io ............... [NO] ....... [OKAY]
fused_adam ............. [NO] ....... [OKAY]
cpu_adam ............... [YES] ...... [OKAY]
cpu_adagrad ............ [YES] ...... [OKAY]
fused_lamb ............. [NO] ....... [OKAY]
quantizer .............. [NO] ....... [OKAY]
random_ltd ............. [NO] ....... [OKAY]
[WARNING] sparse_attn requires a torch version >= 1.5 and < 2.0 but detected 2.0
[WARNING] using untested triton version (2.0.0), only 1.0.0 is known to be compatible
sparse_attn ............ [NO] ....... [NO]
spatial_inference ...... [NO] ....... [OKAY]
transformer ............ [YES] ...... [OKAY]
stochastic_transformer . [NO] ....... [OKAY]
transformer_inference .. [NO] ....... [OKAY]

DeepSpeed general environment info:
torch install path ............... ['/usr/local/lib/python3.8/dist-packages/torch']
torch version .................... 2.0.1+cu117
deepspeed install path ........... ['/usr/local/lib/python3.8/dist-packages/deepspeed']
deepspeed info ................... 0.10.3, unknown, unknown
torch cuda version ............... 11.7
torch hip version ................ None
nvcc version ..................... 11.1
deepspeed wheel compiled w. ...... torch 2.0, cuda 11.7
shared memory (/dev/shm) size .... 503.84 GB

config file:
{
"train_batch_size": 4,
// "train_micro_batch_size_per_gpu": 64s,
"steps_per_print": 100,
/"gradient_accumulation_steps": 4,
"optimizer": {
"type": "Adam",
"params": {
"lr": 3e-4,
"betas": [0.8, 0.999],
"eps": 1e-8,
"weight_decay": 3e-7
}
},
"fp16": {
"enabled": true,
"loss_scale": 0,
"loss_scale_window": 1000,
"hysteresis": 2,
"min_loss_scale": 1
},

"zero_optimization": {
"stage": 1,
"allgather_partitions": true,
"allgather_bucket_size": 5e8,
"overlap_comm": true,
"reduce_scatter": true,
"reduce_bucket_size": 5e8,
"contiguous_gradients" : false,
"offload_param": {
"device": "cpu",
},

"stage3_max_live_parameters" : 1e9,
"stage3_max_reuse_distance" : 1e9,
"stage3_prefetch_bucket_size" : 5e8,
"stage3_param_persistence_threshold" : 1e6,
"sub_group_size" : 1e12,
"elastic_checkpoint" : false,
"stage3_gather_16bit_weights_on_model_save": true,
"ignore_unused_parameters": false,
"round_robin_gradients": true,
"zero_hpz_partition_size": 1,
"zero_quantized_weights": false,
"zero_quantized_gradients": false
}

Launcher context
deepspeed --include localhost:4,5 parallel_gpt.py

Start with the traceback entry point in deepspeed/runtime/zero/stage_1_and_2.py, especially has_overflow_partitioned_grads_serial and the surrounding overflow checks. Reproduce with parallel_gpt.py and the supplied ZeRO configuration, then compare the reported averaged_gradients access with the failing setup. Done means the training step no longer raises KeyError: 0 under this configuration.