Update DeepSpeed version, fix failing tests (#9898)

Author: Sean Naren

.azure-pipelines/gpu-tests.yml

@@ -51,7 +51,7 @@ jobs:
     - bash: |
         python -c "fname = 'requirements/extra.txt' ; lines = [line for line in open(fname).readlines() if 'horovod' not in line] ; open(fname, 'w').writelines(lines)"
         pip install fairscale>=0.3.4
-        pip install "deepspeed==0.4.3" # FIXME: bug with >= 0.4.4
+        pip install deepspeed==0.5.4
         pip install . --requirement requirements/devel.txt
         pip list
       displayName: 'Install dependencies'

tests/models/test_hooks.py

@@ -416,18 +416,35 @@ def _predict_batch(trainer, model, batches):
         return out
 
 
+@RunIf(deepspeed=True, min_gpus=1, special=True)
+def test_trainer_model_hook_system_fit_deepspeed_automatic_optimization(tmpdir):
+    _run_trainer_model_hook_system_fit(
+        dict(gpus=1, precision=16, plugins="deepspeed"), tmpdir, automatic_optimization=True
+    )
+
+
+@RunIf(deepspeed=True, min_gpus=1, special=True)
+def test_trainer_model_hook_system_fit_deepspeed_manual_optimization(tmpdir):
+    _run_trainer_model_hook_system_fit(
+        dict(gpus=1, precision=16, plugins="deepspeed"), tmpdir, automatic_optimization=False
+    )
+
+
 @pytest.mark.parametrize(
     "kwargs",
     [
         {},
         # these precision plugins modify the optimization flow, so testing them explicitly
-        pytest.param(dict(gpus=1, precision=16, plugins="deepspeed"), marks=RunIf(deepspeed=True, min_gpus=1)),
         pytest.param(dict(gpus=1, precision=16, amp_backend="native"), marks=RunIf(min_gpus=1)),
         pytest.param(dict(gpus=1, precision=16, amp_backend="apex"), marks=RunIf(amp_apex=True, min_gpus=1)),
     ],
 )
 @pytest.mark.parametrize("automatic_optimization", (True, False))
 def test_trainer_model_hook_system_fit(tmpdir, kwargs, automatic_optimization):
+    _run_trainer_model_hook_system_fit(kwargs, tmpdir, automatic_optimization)
+
+
+def _run_trainer_model_hook_system_fit(kwargs, tmpdir, automatic_optimization):
     called = []
 
     class TestModel(HookedModel):
@@ -459,14 +476,11 @@ def training_step(self, batch, batch_idx):
         callbacks=[callback],
         **kwargs,
     )
-
     assert called == [
         dict(name="Callback.on_init_start", args=(trainer,)),
         dict(name="Callback.on_init_end", args=(trainer,)),
     ]
-
     trainer.fit(model)
-
     saved_ckpt = {
         "callbacks": ANY,
         "epoch": 1,
@@ -481,7 +495,6 @@ def training_step(self, batch, batch_idx):
     elif kwargs.get("amp_backend") == "apex":
         saved_ckpt["amp_scaling_state"] = ANY
     device = torch.device("cuda:0" if "gpus" in kwargs else "cpu")
-
     expected = [
         dict(name="Callback.on_init_start", args=(trainer,)),
         dict(name="Callback.on_init_end", args=(trainer,)),

tests/plugins/test_deepspeed_plugin.py

@@ -213,20 +213,21 @@ def backward(self, loss: Tensor, optimizer: Optimizer, optimizer_idx: int, *args
         trainer.fit(model)
 
 
-@RunIf(min_gpus=1, deepspeed=True, special=True)
+@RunIf(min_gpus=1, deepspeed=True)
 @pytest.mark.parametrize(
     ["dataset_cls", "value"],
     [(RandomDataset, "auto"), (RandomDataset, 10), (RandomIterableDataset, "auto"), (RandomIterableDataset, 10)],
 )
-def test_deepspeed_auto_batch_size_config_select(tmpdir, dataset_cls, value):
+@mock.patch("deepspeed.init_distributed", autospec=True)
+def test_deepspeed_auto_batch_size_config_select(mock_deepspeed_distributed, tmpdir, dataset_cls, value):
     """Test to ensure that the batch size is correctly set as expected for deepspeed logging purposes."""
 
     class TestModel(BoringModel):
         def train_dataloader(self):
             return DataLoader(dataset_cls(32, 64))
 
     class AssertCallback(Callback):
-        def on_train_start(self, trainer, pl_module) -> None:
+        def setup(self, trainer, pl_module, stage: Optional[str] = None) -> None:
             assert isinstance(trainer.accelerator.training_type_plugin, DeepSpeedPlugin)
             config = trainer.accelerator.training_type_plugin.config
 
@@ -855,7 +856,7 @@ def test_deepspeed_multigpu_no_schedulers(tmpdir):
     _assert_save_model_is_equal(model, tmpdir, trainer)
 
 
-@RunIf(min_gpus=1, deepspeed=True)
+@RunIf(min_gpus=1, deepspeed=True, special=True)
 def test_deepspeed_skip_backward_raises(tmpdir):
     class TestModel(BoringModel):
         def training_step(self, batch, batch_idx):