Support mixed precision with CPU driver

ray_lightning/ray_ddp.py

@@ -9,7 +9,6 @@
 import torch
 
 import pytorch_lightning as pl
-from pytorch_lightning.accelerators import CPUAccelerator
 from pytorch_lightning.plugins import DDPSpawnPlugin
 from pytorch_lightning import _logger as log, LightningModule
 from pytorch_lightning.trainer.states import TrainerFn
@@ -23,7 +22,7 @@
 
 from ray_lightning.session import init_session
 from ray_lightning.util import process_results, to_state_stream, \
-    load_state_stream
+    load_state_stream, swap_accelerator
 from ray_lightning.tune import TUNE_INSTALLED, is_session_enabled
 
 
@@ -173,22 +172,7 @@ def setup(self):
             ray.get([w.execute.remote(self.init_hook) for w in self.workers])
 
     def setup_environment(self) -> None:
-        # Swap out the accelerator if necessary.
-        # This is needed to support CPU head with GPU workers or Ray Client.
-        current_accelerator = self.lightning_module.trainer.accelerator
-        if self.use_gpu and isinstance(current_accelerator, CPUAccelerator):
-            from weakref import proxy
-            from ray_lightning.util import DelayedGPUAccelerator
-            precision_plugin = current_accelerator.precision_plugin
-            new_accelerator = DelayedGPUAccelerator(
-                precision_plugin=precision_plugin, training_type_plugin=self)
-            self.lightning_module.trainer._accelerator_connector \
-                ._training_type_plugin = \
-                proxy(new_accelerator.training_type_plugin)
-            self.lightning_module.trainer._accelerator_connector \
-                ._precision_plugin = proxy(new_accelerator.precision_plugin)
-            self.lightning_module.trainer._accelerator_connector.accelerator \
-                = new_accelerator
+        swap_accelerator(self)
 
     def _setup_env_vars(self):
         # Get rank 0 worker address and port for DDP connection.

ray_lightning/ray_horovod.py

@@ -1,6 +1,5 @@
 import torch
 import pytorch_lightning as pl
-from pytorch_lightning.accelerators import CPUAccelerator
 from pytorch_lightning.plugins import HorovodPlugin
 from pytorch_lightning.utilities import rank_zero_only
 
@@ -11,7 +10,7 @@
 
 from ray_lightning.session import init_session
 from ray_lightning.util import process_results, Unavailable, to_state_stream, \
-    load_state_stream
+    load_state_stream, swap_accelerator
 from ray_lightning.tune import TUNE_INSTALLED, is_session_enabled
 
 try:
@@ -126,22 +125,7 @@ def setup(self):
         self.executor.start(executable_cls=get_executable_cls())
 
     def setup_environment(self) -> None:
-        # Swap out the accelerator if necessary.
-        # This is needed to support CPU head with GPU workers or Ray Client.
-        current_accelerator = self.lightning_module.trainer.accelerator
-        if self.use_gpu and isinstance(current_accelerator, CPUAccelerator):
-            from weakref import proxy
-            from ray_lightning.util import DelayedGPUAccelerator
-            precision_plugin = current_accelerator.precision_plugin
-            new_accelerator = DelayedGPUAccelerator(
-                precision_plugin=precision_plugin, training_type_plugin=self)
-            self.lightning_module.trainer._accelerator_connector \
-                ._training_type_plugin = \
-                proxy(new_accelerator.training_type_plugin)
-            self.lightning_module.trainer._accelerator_connector \
-                ._precision_plugin = proxy(new_accelerator.precision_plugin)
-            self.lightning_module.trainer._accelerator_connector.accelerator \
-                = new_accelerator
+        swap_accelerator(self)
 
     def pre_dispatch(self):
         """All pre-dispatch logic should be done in train_remote instead."""

ray_lightning/tests/test_ddp.py

@@ -42,6 +42,7 @@ def start_ray_client_server_2_cpus():
     ray.init(num_cpus=2)
     with ray_start_client_server() as client:
         yield client
+    ray.shutdown()
 
 
 @pytest.fixture

ray_lightning/tests/test_ddp_gpu.py

@@ -43,6 +43,18 @@ def test_train(tmpdir, ray_start_2_gpus, num_workers):
     train_test(trainer, model)
 
 
+@pytest.mark.skipif(
+    torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
+def test_train_mixed_precision(tmpdir, ray_start_2_gpus):
+    """Tests if training works with mixed precision."""
+    model = BoringModel()
+    plugin = RayPlugin(num_workers=2, use_gpu=True)
+    trainer = get_trainer(tmpdir, plugins=[plugin], gpus=1, precision=16)
+    # Make sure PTL doesn't automatically replace with bf16.
+    assert trainer.precision == 16
+    train_test(trainer, model)
+
+
 @pytest.mark.skipif(
     torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
 @pytest.mark.parametrize("num_workers", [1, 2])

ray_lightning/tests/test_tune.py

@@ -25,17 +25,23 @@ def ray_start_4_cpus_4_gpus():
     ray.shutdown()
 
 
-def train_func(dir, plugin, callbacks=None):
+def train_func(dir, plugin, callbacks=None, amp=False):
     def _inner_train(config):
         model = BoringModel()
         trainer = get_trainer(
             dir,
             callbacks=callbacks,
             plugins=[plugin],
             checkpoint_callback=False,
+            gpus=1 if amp else 0,
+            precision=16 if amp else 32,
             **config)
         trainer.fit(model)
 
+        if amp:
+            # Make sure PTL doesn't automatically replace with bf16
+            assert trainer.precision == 16
+
     return _inner_train
 
 
@@ -104,3 +110,19 @@ def test_checkpoint_horovod_gpu(tmpdir, ray_start_4_cpus_4_gpus):
     """Tests if Tune checkpointing works with HorovodRayAccelerator."""
     plugin = HorovodRayPlugin(num_workers=2, use_gpu=True)
     checkpoint_test(tmpdir, plugin)
+
+
+def tune_test_mixed_precision(dir, plugin):
+    tune.run(
+        train_func(dir, plugin),
+        resources_per_trial=get_tune_resources(
+            num_workers=plugin.num_workers, use_gpu=plugin.use_gpu),
+        num_samples=2)
+
+
+@pytest.mark.skipif(
+    torch.cuda.device_count() < 4, reason="test requires multi-GPU machine")
+def test_tune_mixed_precision_ddp_gpu(tmpdir, ray_start_4_cpus_4_gpus):
+    """Tests if Tune works with mixed precision."""
+    plugin = RayPlugin(num_workers=2, use_gpu=True)
+    tune_test_mixed_precision(tmpdir, plugin)

ray_lightning/util.py

@@ -6,6 +6,7 @@
 from pytorch_lightning import Trainer
 
 import ray
+from pytorch_lightning.plugins import TrainingTypePlugin
 
 
 class DelayedGPUAccelerator(GPUAccelerator):
@@ -37,6 +38,27 @@ def on_train_start(self) -> None:
         super(DelayedGPUAccelerator, self).on_train_start()
 
 
+def swap_accelerator(plugin: TrainingTypePlugin):
+    # Swap out the accelerator if necessary.
+    # This is needed to support CPU head with GPU workers or Ray Client.
+    # This is also needed to support GPU-only optimizations like mixed
+    # precision when using CPU head with GPU workers or Ray Client.
+    current_accelerator = plugin.lightning_module.trainer.accelerator
+
+    if plugin.use_gpu:
+        from weakref import proxy
+        precision_plugin = current_accelerator.precision_plugin
+        new_accelerator = DelayedGPUAccelerator(
+            precision_plugin=precision_plugin, training_type_plugin=plugin)
+        plugin.lightning_module.trainer._accelerator_connector \
+            ._training_type_plugin = \
+            proxy(new_accelerator.training_type_plugin)
+        plugin.lightning_module.trainer._accelerator_connector \
+            ._precision_plugin = proxy(new_accelerator.precision_plugin)
+        plugin.lightning_module.trainer._accelerator_connector.accelerator \
+            = new_accelerator
+
+
 class Unavailable:
     """No object should be instance of this class"""