DeepSpeed 优化器问题

Question

我正在尝试使用 deepspeed 优化器来让我能够将更多内容放入内存并加快我的训练速度。我一直在优化器方面遇到这个问题，但我不确定是什么原因造成的。错误如下：

Traceback (most recent call last):
  File "02232022_kat_repr_train.py", line 89, in <module>
    first_exp.run_experiment()
  File "experiments.py", line 28, in run_experiment
    self.trainer.train(self.hyperparameters["epochs"], self.hyperparameters["eval_period"])
  File "trainer.py", line 90, in train
    self.optimizer.step()
  File "/anaconda3/envs/mmlm/lib/python3.8/site-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 1633, in step
    self.check_overflow()
  File "/anaconda3/envs/mmlm/lib/python3.8/site-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 1913, in check_overflow
    self._check_overflow(partition_gradients)
  File "/anaconda3/envs/mmlm/lib/python3.8/site-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 1818, in _check_overflow
    self.overflow = self.has_overflow(partition_gradients)
  File "/anaconda3/envs/mmlm/lib/python3.8/site-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 1837, in has_overflow
    overflow = self.local_overflow if self.cpu_offload else self.has_overflow_partitioned_grads_serial(
  File "/anaconda3/envs/mmlm/lib/python3.8/site-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 1830, in has_overflow_partitioned_grads_serial
    for j, grad in enumerate(self.averaged_gradients[i]):
KeyError: 0

我的训练代码非常标准：

outputs = self.model(**batch)
loss = self.loss_function(outputs)
loss = self.loss_aggregation(loss)
epoch_loss += loss
loss.backward()
self.optimizer.backward(loss)
self.model.zero_grad()

我像这样初始化deepspeed： self.model, self.optimizer, _, self.lr_scheduler = ds.initialize(model=self.model, config_params=self.deepspeed_config , optimizationr=self.optimizer, lr_scheduler=self.lr_scheduler)

我正在使用此处的示例配置：https://github.com/microsoft/DeepSpeedExamples/blob/ master/Megatron-LM/scripts/ds_zero2_config.json

有谁知道为什么我会遇到这个问题？

谢谢你！