HelloDeepSpeed not reproducible

Hi DeepSpeed community,

I was trying to run the HelloDeepSpeed example with a AWS p3.16x instance (8 v100 gpus). However, I was hitting this issue:

deepspeed train_bert_ds.py --checkpoint_dir .

  File "train_bert_ds.py", line 809, in <module>
    fire.Fire(train)
  File "/home/ubuntu/anaconda3/envs/pytorch_latest_p37/lib/python3.7/site-packages/fire/core.py", line 141, in Fire
    component_trace = _Fire(component, args, parsed_flag_args, context, name)
  File "/home/ubuntu/anaconda3/envs/pytorch_latest_p37/lib/python3.7/site-packages/fire/core.py", line 471, in _Fire
    target=component.__name__)
  File "/home/ubuntu/anaconda3/envs/pytorch_latest_p37/lib/python3.7/site-packages/fire/core.py", line 681, in _CallAndUpdateTrace
    component = fn(*varargs, **kwargs)
  File "train_bert_ds.py", line 783, in train
    model.backward(loss)
  File "/home/ubuntu/anaconda3/envs/pytorch_latest_p37/lib/python3.7/site-packages/deepspeed/runtime/engine.py", line 1722, in backward
    self.allreduce_gradients()
  File "/home/ubuntu/anaconda3/envs/pytorch_latest_p37/lib/python3.7/site-packages/deepspeed/runtime/engine.py", line 1651, in allreduce_gradients
    pipeline_parallel=self.pipeline_parallelism)
  File "/home/ubuntu/anaconda3/envs/pytorch_latest_p37/lib/python3.7/site-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 601, in reduce_gradients
    self.overlapping_partition_gradients_reduce_epilogue()
  File "/home/ubuntu/anaconda3/envs/pytorch_latest_p37/lib/python3.7/site-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 756, in overlapping_partition_gradients_reduce_epilogue
    self.independent_gradient_partition_epilogue()
  File "/home/ubuntu/anaconda3/envs/pytorch_latest_p37/lib/python3.7/site-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 643, in independent_gradient_partition_epilogue
    self.reduce_ipg_grads()
  File "/home/ubuntu/anaconda3/envs/pytorch_latest_p37/lib/python3.7/site-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 1240, in reduce_ipg_grads
    self.copy_grads_in_partition(param)
  File "/home/ubuntu/anaconda3/envs/pytorch_latest_p37/lib/python3.7/site-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 1158, in copy_grads_in_partition
    self.set_norm_for_param_grad_in_gpu(param)
  File "/home/ubuntu/anaconda3/envs/pytorch_latest_p37/lib/python3.7/site-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 1080, in set_norm_for_param_grad_in_gpu
    [i, source_offset, dest_offset, num_elements] = self.grad_position[param_id]
KeyError: 44

DeepSpeed was pip installed and the torch version does not seem to matter
Would appreciate any insights!

Reproduce deepspeed train_bert_ds.py --checkpoint_dir . on the reported AWS p3.16x setup, then inspect train_bert_ds.py around lines 783 and 809 alongside the DeepSpeed Zero stage files named in the traceback. Done means the HelloDeepSpeed example completes without the KeyError: 44, or the issue is narrowed to a documented version or environment constraint.