Add more trainer config and ttp register tests

Author: four4fish

tests/plugins/test_plugins_registry.py

@@ -16,6 +16,7 @@
 from pytorch_lightning import Trainer
 from pytorch_lightning.plugins import (
     CheckpointIO,
+    DDPFullyShardedPlugin,
     DDPPlugin,
     DDPShardedPlugin,
     DDPSpawnPlugin,
@@ -93,6 +94,18 @@ def test_tpu_spawn_debug_plugins_registry(tmpdir):
     assert isinstance(trainer.training_type_plugin, TPUSpawnPlugin)
 
 
+def test_fsdp_plugins_registry(tmpdir):
+
+    plugin = "fsdp"
+
+    assert plugin in TrainingTypePluginsRegistry
+    assert TrainingTypePluginsRegistry[plugin]["plugin"] == DDPFullyShardedPlugin
+
+    trainer = Trainer(strategy=plugin)
+
+    assert isinstance(trainer.training_type_plugin, DDPFullyShardedPlugin)
+
+
 @pytest.mark.parametrize(
     "plugin_name, plugin",
     [

tests/trainer/test_trainer.py

@@ -2162,13 +2162,14 @@ def training_step(self, batch, batch_idx):
             dict(strategy="ddp_spawn", num_processes=1, gpus=None),
             dict(_distrib_type=None, _device_type=DeviceType.CPU, num_gpus=0, num_processes=1),
         ),
-        (
-            dict(strategy="fsdp", gpus=2),
-            dict(_distrib_type=DistributedType.DDP_FULLY_SHARDED, _device_type=DeviceType.GPU, num_gpus=2, num_processes=1),
-        ),
         (
             dict(strategy="ddp_fully_sharded", gpus=1),
-            dict(_distrib_type=DistributedType.DDP_FULLY_SHARDED, _device_type=DeviceType.GPU, num_gpus=1, num_processes=1),
+            dict(
+                _distrib_type=DistributedType.DDP_FULLY_SHARDED,
+                _device_type=DeviceType.GPU,
+                num_gpus=1,
+                num_processes=1,
+            ),
         ),
         (
             dict(strategy=DDPSpawnPlugin(), num_processes=2, gpus=None),