deepspeed + grad acumm (#9622)

examples/seq2seq/test_finetune_trainer.py

@@ -112,6 +112,11 @@ class TestFinetuneTrainer(TestCasePlus):
     def test_finetune_trainer_deepspeed(self):
         self.finetune_trainer_quick(deepspeed=True)
 
+    @require_torch_multi_gpu
+    @require_deepspeed
+    def test_finetune_trainer_deepspeed_grad_acum(self):
+        self.finetune_trainer_quick(deepspeed=True, extra_args_str="--gradient_accumulation_steps 2")
+
     @slow
     def test_finetune_trainer_slow(self):
         # There is a missing call to __init__process_group somewhere

src/transformers/trainer.py

@@ -931,7 +931,9 @@ class Trainer:
                             )
 
                     # Optimizer step
-                    if is_torch_tpu_available():
+                    if self.deepspeed:
+                        self.deepspeed.step()
+                    elif is_torch_tpu_available():
                         xm.optimizer_step(self.optimizer)
                     elif self.use_amp:
                         self.scaler.step(self.optimizer)