Update tests in strategies directory in preparation for #11040 (#12467)

Author: daniellepintz

tests/strategies/test_bagua_strategy.py

@@ -118,4 +118,4 @@ def test_bagua_not_available(monkeypatch):
     monkeypatch.setattr(imports, "_BAGUA_AVAILABLE", False)
     with mock.patch("torch.cuda.device_count", return_value=1):
         with pytest.raises(MisconfigurationException, match="you must have `Bagua` installed"):
-            Trainer(strategy="bagua", gpus=1)
+            Trainer(strategy="bagua", accelerator="gpu", devices=1)

tests/strategies/test_common.py

@@ -26,9 +26,9 @@
 @pytest.mark.parametrize(
     "trainer_kwargs",
     (
-        pytest.param(dict(gpus=1), marks=RunIf(min_gpus=1)),
-        pytest.param(dict(strategy="dp", gpus=2), marks=RunIf(min_gpus=2)),
-        pytest.param(dict(strategy="ddp_spawn", gpus=2), marks=RunIf(min_gpus=2)),
+        pytest.param(dict(accelerator="gpu", devices=1), marks=RunIf(min_gpus=1)),
+        pytest.param(dict(strategy="dp", accelerator="gpu", devices=2), marks=RunIf(min_gpus=2)),
+        pytest.param(dict(strategy="ddp_spawn", accelerator="gpu", devices=2), marks=RunIf(min_gpus=2)),
     ),
 )
 def test_evaluate(tmpdir, trainer_kwargs):

tests/strategies/test_ddp.py

@@ -109,7 +109,7 @@ def setup(self, stage: Optional[str] = None) -> None:
             raise SystemExit()
 
     model = TestModel()
-    trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=True, strategy="ddp", gpus=1)
+    trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=True, strategy="ddp", accelerator="gpu", devices=1)
     with pytest.deprecated_call(match="Environment variable `PL_TORCH_DISTRIBUTED_BACKEND` was deprecated in v1.6"):
         with pytest.raises(SystemExit):
             trainer.fit(model)

tests/strategies/test_ddp_fully_sharded_with_full_state_dict.py

@@ -34,7 +34,9 @@ def test_invalid_on_cpu(tmpdir):
 @RunIf(fairscale_fully_sharded=True)
 def test_fsdp_with_sharded_amp(device_count_mock, mock_cuda_available, tmpdir):
     """Test to ensure that plugin native amp plugin is correctly chosen when using sharded."""
-    trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=True, strategy="fsdp", gpus=1, precision=16)
+    trainer = Trainer(
+        default_root_dir=tmpdir, fast_dev_run=True, strategy="fsdp", accelerator="gpu", devices=1, precision=16
+    )
     assert isinstance(trainer.strategy, DDPFullyShardedStrategy)
     assert isinstance(trainer.strategy.precision_plugin, FullyShardedNativeMixedPrecisionPlugin)
 

tests/strategies/test_ddp_spawn_strategy.py

@@ -69,7 +69,7 @@ def test_ddp_cpu():
 def test_ddp_spawn_extra_parameters(tmpdir):
     """Tests if device is set correctly when training for DDPSpawnStrategy and tests add_to_queue/get_from_queue
     with Lightning Module (deprecated way)."""
-    trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=True, gpus=2, strategy="ddp_spawn")
+    trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=True, accelerator="gpu", devices=2, strategy="ddp_spawn")
 
     assert isinstance(trainer.strategy, DDPSpawnStrategy)
     assert trainer.strategy.root_device == torch.device("cuda:0")
@@ -103,7 +103,9 @@ def test_ddp_spawn_add_get_queue(tmpdir):
     """Tests add_to_queue/get_from_queue with DDPSpawnStrategy."""
 
     ddp_spawn_strategy = TestDDPSpawnStrategy()
-    trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=True, num_processes=2, strategy=ddp_spawn_strategy)
+    trainer = Trainer(
+        default_root_dir=tmpdir, fast_dev_run=True, accelerator="cpu", devices=2, strategy=ddp_spawn_strategy
+    )
 
     val: float = 1.0
     val_name: str = "val_acc"
@@ -138,7 +140,7 @@ def on_predict_start(self) -> None:
 @RunIf(skip_windows=True, skip_49370=True, skip_hanging_spawn=True)
 def test_ddp_spawn_configure_ddp(tmpdir):
     """Tests with ddp spawn strategy."""
-    trainer = Trainer(default_root_dir=tmpdir, num_processes=2, strategy="ddp_spawn", fast_dev_run=True)
+    trainer = Trainer(default_root_dir=tmpdir, accelerator="cpu", devices=2, strategy="ddp_spawn", fast_dev_run=True)
 
     model = BoringModelDDP()
 

tests/strategies/test_ddp_strategy.py

@@ -87,7 +87,8 @@ def creates_processes_externally(self):
     trainer = Trainer(
         default_root_dir=tmpdir,
         strategy="ddp",
-        num_processes=2,
+        accelerator="cpu",
+        devices=2,
         plugins=[WronglyImplementedEnvironment()],
     )
     with pytest.raises(
@@ -138,7 +139,7 @@ def test_ddp_dont_configure_sync_batchnorm(trainer_fn):
     model = BoringModelGPU()
     model.layer = torch.nn.BatchNorm1d(10)
     ddp_strategy = DDPStrategy()
-    trainer = Trainer(gpus=1, strategy=ddp_strategy, sync_batchnorm=True)
+    trainer = Trainer(accelerator="gpu", devices=1, strategy=ddp_strategy, sync_batchnorm=True)
     trainer.state.fn = trainer_fn
     trainer.strategy.connect(model)
     trainer.lightning_module.trainer = trainer

tests/strategies/test_dp.py

@@ -72,7 +72,8 @@ def test_multi_gpu_early_stop_dp(tmpdir):
         max_epochs=50,
         limit_train_batches=10,
         limit_val_batches=10,
-        gpus=[0, 1],
+        accelerator="gpu",
+        devices=[0, 1],
         strategy="dp",
     )
 
@@ -88,7 +89,8 @@ def test_multi_gpu_model_dp(tmpdir):
         max_epochs=1,
         limit_train_batches=10,
         limit_val_batches=10,
-        gpus=[0, 1],
+        accelerator="gpu",
+        devices=[0, 1],
         strategy="dp",
         enable_progress_bar=False,
     )
@@ -166,7 +168,7 @@ def transfer_batch_to_device(self, batch, device, dataloader_idx):
             batch = batch.to(device)
             return batch
 
-    trainer_options = dict(default_root_dir=tmpdir, max_steps=7, gpus=[0, 1], strategy="dp")
+    trainer_options = dict(default_root_dir=tmpdir, max_steps=7, accelerator="gpu", devices=[0, 1], strategy="dp")
 
     trainer = Trainer(**trainer_options)
     model = CustomModel()
@@ -208,7 +210,8 @@ def test_dp_training_step_dict(tmpdir):
     trainer = pl.Trainer(
         default_root_dir=tmpdir,
         fast_dev_run=True,
-        gpus=2,
+        accelerator="gpu",
+        devices=2,
         strategy="dp",
     )
     trainer.fit(model)

tests/strategies/test_sharded_strategy.py

@@ -22,7 +22,14 @@
 def test_ddp_sharded_precision_16_clip_gradients(mock_oss_clip_grad_norm, clip_val, tmpdir):
     """Ensure that clip gradients is only called if the value is greater than 0."""
     model = BoringModel()
-    trainer = Trainer(strategy="ddp_sharded", gpus=1, precision=16, fast_dev_run=True, gradient_clip_val=clip_val)
+    trainer = Trainer(
+        strategy="ddp_sharded",
+        accelerator="gpu",
+        devices=1,
+        precision=16,
+        fast_dev_run=True,
+        gradient_clip_val=clip_val,
+    )
     trainer.fit(model)
     if clip_val > 0:
         mock_oss_clip_grad_norm.assert_called()
@@ -46,15 +53,15 @@ def test_sharded_ddp_choice(tmpdir, strategy, expected):
 )
 def test_ddp_choice_sharded_amp(tmpdir, strategy, expected):
     """Test to ensure that plugin native amp plugin is correctly chosen when using sharded."""
-    trainer = Trainer(fast_dev_run=True, gpus=1, precision=16, strategy=strategy)
+    trainer = Trainer(fast_dev_run=True, accelerator="gpu", devices=1, precision=16, strategy=strategy)
     assert isinstance(trainer.strategy, expected)
 
 
 @RunIf(skip_windows=True, fairscale=True)
 def test_ddp_sharded_strategy_checkpoint_cpu(tmpdir):
     """Test to ensure that checkpoint is saved correctly."""
     model = BoringModel()
-    trainer = Trainer(strategy="ddp_sharded_spawn", num_processes=2, fast_dev_run=True)
+    trainer = Trainer(strategy="ddp_sharded_spawn", accelerator="cpu", devices=2, fast_dev_run=True)
 
     trainer.fit(model)
 
@@ -103,7 +110,7 @@ def test_ddp_sharded_strategy_finetune(tmpdir):
 def test_ddp_sharded_strategy_fit_ckpt_path(tmpdir):
     """Test to ensure that resuming from checkpoint works."""
     model = BoringModel()
-    trainer = Trainer(strategy="ddp_sharded_spawn", num_processes=2, fast_dev_run=True)
+    trainer = Trainer(strategy="ddp_sharded_spawn", accelerator="cpu", devices=2, fast_dev_run=True)
 
     trainer.fit(model)
 
@@ -112,7 +119,7 @@ def test_ddp_sharded_strategy_fit_ckpt_path(tmpdir):
 
     model = BoringModel()
 
-    trainer = Trainer(strategy="ddp_sharded_spawn", num_processes=2, fast_dev_run=True)
+    trainer = Trainer(strategy="ddp_sharded_spawn", accelerator="cpu", devices=2, fast_dev_run=True)
 
     trainer.fit(model, ckpt_path=checkpoint_path)
 
@@ -141,7 +148,7 @@ def test_ddp_sharded_strategy_fit_ckpt_path_downsize_gpus(tmpdir):
 def test_ddp_sharded_strategy_fit_ckpt_path_gpu_to_cpu(tmpdir):
     """Test to ensure that resuming from checkpoint works when going from GPUs- > CPU."""
     model = BoringModel()
-    trainer = Trainer(strategy="ddp_sharded_spawn", gpus=1, fast_dev_run=True)
+    trainer = Trainer(strategy="ddp_sharded_spawn", accelerator="gpu", devices=1, fast_dev_run=True)
 
     trainer.fit(model)
 
@@ -150,13 +157,16 @@ def test_ddp_sharded_strategy_fit_ckpt_path_gpu_to_cpu(tmpdir):
 
     model = BoringModel()
 
-    trainer = Trainer(strategy="ddp_sharded_spawn", num_processes=2, fast_dev_run=True)
+    trainer = Trainer(strategy="ddp_sharded_spawn", accelerator="cpu", devices=2, fast_dev_run=True)
 
     trainer.fit(model, ckpt_path=checkpoint_path)
 
 
 @RunIf(skip_windows=True, standalone=True, fairscale=True)
-@pytest.mark.parametrize("trainer_kwargs", (dict(num_processes=2), pytest.param(dict(gpus=2), marks=RunIf(min_gpus=2))))
+@pytest.mark.parametrize(
+    "trainer_kwargs",
+    (dict(accelerator="cpu", devices=2), pytest.param(dict(accelerator="gpu", devices=2), marks=RunIf(min_gpus=2))),
+)
 def test_ddp_sharded_strategy_test_multigpu(tmpdir, trainer_kwargs):
     """Test to ensure we can use validate and test without fit."""
     model = BoringModel()