Repository: ray-project/ray_lightning (language: python)

Issue:
ray ddp fails with 2 gpu workers
```python
  File "/home/ubuntu/anaconda3/envs/automm-dev-pl-latest/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1797, in __setup_profiler
    self.profiler.setup(stage=self.state.fn._setup_fn, local_rank=local_rank, log_dir=self.log_dir)
  File "/home/ubuntu/anaconda3/envs/automm-dev-pl-latest/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 2249, in log_dir
    dirpath = self.strategy.broadcast(dirpath)
  File "/home/ubuntu/anaconda3/envs/automm-dev-pl-latest/lib/python3.8/site-packages/pytorch_lightning/strategies/ddp_spawn.py", line 215, in broadcast
    torch.distributed.broadcast_object_list(obj, src, group=_group.WORLD)
  File "/home/ubuntu/anaconda3/envs/automm-dev-pl-latest/lib/python3.8/site-packages/torch/distributed/distributed_c10d.py", line 1817, in broadcast_object_list
    broadcast(object_sizes_tensor, src=src, group=group)
  File "/home/ubuntu/anaconda3/envs/automm-dev-pl-latest/lib/python3.8/site-packages/torch/distributed/distributed_c10d.py", line 1159, in broadcast
    work = default_pg.broadcast([tensor], opts)
RuntimeError: NCCL error in: ../torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp:891, internal error, NCCL version 21.0.3
ncclInternalError: Internal check failed. This is either a bug in NCCL or due to memory corruptio
```

> use this branch: https://github.com/sxjscience/autogluon/tree/kaggle_california_house and install autogluon via
> bash full_install.sh
> Afterwards, try this script: https://gist.github.com/sxjscience/53bc799e37cc0680ca9e53c2fea75cd7
> Internally, the ray strategy are constructed here: https://github.com/sxjscience/autogluon/blob/59f01b95381fba5651db17fd98fa84164ad168c2/multimodal/src/autogluon/multimodal/predictor.py#L1036-L1052 . 








Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index 7230e715..032d7ce0 100644
--- a/README.md
+++ b/README.md
@@ -1,11 +1,11 @@
 <!--$UNCOMMENT(ray-lightning)=-->
 
 # Distributed PyTorch Lightning Training on Ray
-This library adds new PyTorch Lightning plugins for distributed training using the Ray distributed computing framework.
+This library adds new PyTorch Lightning strategies for distributed training using the Ray distributed computing framework.
 
-These PyTorch Lightning Plugins on Ray enable quick and easy parallel training while still leveraging all the benefits of PyTorch Lightning and using your desired training protocol, either [PyTorch Distributed Data Parallel](https://pytorch.org/tutorials/intermediate/ddp_tutorial.html) or [Horovod](https://github.com/horovod/horovod). 
+These PyTorch Lightning strategies on Ray enable quick and easy parallel training while still leveraging all the benefits of PyTorch Lightning and using your desired training protocol, either [PyTorch Distributed Data Parallel](https://pytorch.org/tutorials/intermediate/ddp_tutorial.html) or [Horovod](https://github.com/horovod/horovod). 
 
-Once you add your plugin to the PyTorch Lightning Trainer, you can parallelize training to all the cores in your laptop, or across a massive multi-node, multi-GPU cluster with no additional code changes.
+Once you add your strategy to the PyTorch Lightning Trainer, you can parallelize training to all the cores in your laptop, or across a massive multi-node, multi-GPU cluster with no additional code changes.
 
 This library also comes with an integration with <!--$UNCOMMENT{ref}`Ray Tune <tune-main>`--><!--$REMOVE-->[Ray Tune](https://tune.io)<!--$END_REMOVE--> for distributed hyperparameter tuning experiments.
 
@@ -39,29 +39,30 @@ Here are the supported PyTorch Lightning versions:
 |---|---|
 | 0.1 | 1.4 |
 | 0.2 | 1.5 |
-| master | 1.5 |
+| 0.3 | 1.6 |
+| master | 1.6 |
 
 
-## PyTorch Distributed Data Parallel Plugin on Ray
-The `RayPlugin` provides Distributed Data Parallel training on a Ray cluster. PyTorch DDP is used as the distributed training protocol, and Ray is used to launch and manage the training worker processes.
+## PyTorch Distributed Data Parallel Strategy on Ray
+The `RayStrategy` provides Distributed Data Parallel training on a Ray cluster. PyTorch DDP is used as the distributed training protocol, and Ray is used to launch and manage the training worker processes.
 
 Here is a simplified example:
 
 ```python
 import pytorch_lightning as pl
-from ray_lightning import RayPlugin
+from ray_lightning import RayStrategy
 
 # Create your PyTorch Lightning model here.
 ptl_model = MNISTClassifier(...)
-plugin = RayPlugin(num_workers=4, num_cpus_per_worker=1, use_gpu=True)
+strategy = RayStrategy(num_workers=4, num_cpus_per_worker=1, use_gpu=True)
 
 # Don't set ``gpus`` in the ``Trainer``.
 # The actual number of GPUs is determined by ``num_workers``.
-trainer = pl.Trainer(..., plugins=[plugin])
+trainer = pl.Trainer(..., strategy=strategy)
 trainer.fit(ptl_model)
 ```
 
-Because Ray is used to launch processes, instead of the same script being called multiple times, you CAN use this plugin even in cases when you cannot use the standard `DDPPlugin` such as 
+Because Ray is used to launch processes, instead of the same script being called multiple times, you CAN use this strategy even in cases when you cannot use the standard `DDPStrategy` such as 
 - Jupyter Notebooks, Google Colab, Kaggle
 - Calling `fit` or `test` multiple times in the same script
 
@@ -94,40 +95,40 @@ Now you can run your training script on the laptop, but have it execute as if yo
 
 **Note:** When using with Ray Client, you must disable checkpointing and logging for your Trainer by setting `checkpoint_callback` and `logger` to `False`.
 
-## Horovod Plugin on Ray
-Or if you prefer to use Horovod as the distributed training protocol, use the `HorovodRayPlugin` instead.
+## Horovod Strategy on Ray
+Or if you prefer to use Horovod as the distributed training protocol, use the `HorovodRayStrategy` instead.
 
 ```python
 import pytorch_lightning as pl
-from ray_lightning import HorovodRayPlugin
+from ray_lightning import HorovodRayStrategy
 
 # Create your PyTorch Lightning model here.
 ptl_model = MNISTClassifier(...)
 
 # 2 workers, 1 CPU and 1 GPU each.
-plugin = HorovodRayPlugin(num_workers=2, use_gpu=True)
+strategy = HorovodRayStrategy(num_workers=2, use_gpu=True)
 
 # Don't set ``gpus`` in the ``Trainer``.
 # The actual number of GPUs is determined by ``num_workers``.
-trainer = pl.Trainer(..., plugins=[plugin])
+trainer = pl.Trainer(..., strategy=strategy)
 trainer.fit(ptl_model)
 ```
 
 ## Model Parallel Sharded Training on Ray
-The `RayShardedPlugin` integrates with [FairScale](https://github.com/facebookresearch/fairscale) to provide sharded DDP training on a Ray cluster.
+The `RayShardedStrategy` integrates with [FairScale](https://github.com/facebookresearch/fairscale) to provide sharded DDP training on a Ray cluster.
 With sharded training, leverage the scalability of data parallel training while drastically reducing memory usage when training large models.
 
 ```python
 import pytorch_lightning as pl
-from ray_lightning import RayShardedPlugin
+from ray_lightning import RayShardedStrategy
 
 # Create your PyTorch Lightning model here.
 ptl_model = MNISTClassifier(...)
-plugin = RayShardedPlugin(num_workers=4, num_cpus_per_worker=1, use_gpu=True)
+strategy = RayShardedStrategy(num_workers=4, num_cpus_per_worker=1, use_gpu=True)
 
 # Don't set ``gpus`` in the ``Trainer``.
 # The actual number of GPUs is determined by ``num_workers``.
-trainer = pl.Trainer(..., plugins=[plugin])
+trainer = pl.Trainer(..., strategy=strategy)
 trainer.fit(ptl_model)
 ```
 See the [Pytorch Lightning docs](https://pytorch-lightning.readthedocs.io/en/stable/advanced/model_parallel.html#sharded-training) for more information on sharded training.
@@ -140,7 +141,7 @@ Example using `ray_lightning` with Tune:
 ```python
 from ray import tune
 
-from ray_lightning import RayPlugin
+from ray_lightning import RayStrategy
 from ray_lightning.examples.ray_ddp_example import MNISTClassifier
 from ray_lightning.tune import TuneReportCallback, get_tune_resources
 
@@ -159,7 +160,7 @@ def train_mnist(config):
     trainer = pl.Trainer(
         max_epochs=4,
         callbacks=callbacks,
-        plugins=[RayPlugin(num_workers=4, use_gpu=False)])
+        strategy=[RayStrategy(num_workers=4, use_gpu=False)])
     trainer.fit(model)
     
 config = {
@@ -184,26 +185,29 @@ print("Best hyperparameters found were: ", analysis.best_config)
 **Note:** Ray Tune requires 1 additional CPU per trial to use for the Trainable driver. So the actual number of resources each trial requires is `num_workers * num_cpus_per_worker + 1`.
 
 ## FAQ
-> I see that `RayPlugin` is based off of Pytorch Lightning's `DDPSpawnPlugin`. However, doesn't the PTL team discourage the use of spawn?
+> I see that `RayStrategy` is based off of Pytorch Lightning's `DDPSpawnStrategy`. However, doesn't the PTL team discourage the use of spawn?
 
 As discussed [here](https://github.com/pytorch/pytorch/issues/51688#issuecomment-773539003), using a spawn approach instead of launch is not all that detrimental. The original factors for discouraging spawn were:
 1. not being able to use 'spawn' in a Jupyter or Colab notebook, and 
 2. not being able to use multiple workers for data loading. 
 
-Neither of these should be an issue with the `RayPlugin` due to Ray's serialization mechanisms. The only thing to keep in mind is that when using this plugin, your model does have to be serializable/pickleable.
+Neither of these should be an issue with the `RayStrategy` due to Ray's serialization mechanisms. The only thing to keep in mind is that when using this strategy, your model does have to be serializable/pickleable.
+
+> Horovod installation issue
+please see [details](./docs/horovod_faq.md)
 
 <!--$UNCOMMENT## API Reference
 
 ```{eval-rst}
-.. autoclass:: ray_lightning.RayPlugin
+.. autoclass:: ray_lightning.RayStrategy
 ```
 
 ```{eval-rst}
-.. autoclass:: ray_lightning.HorovodRayPlugin
+.. autoclass:: ray_lightning.HorovodRayStrategy
 ```
 
 ```{eval-rst}
-.. autoclass:: ray_lightning.RayShardedPlugin
+.. autoclass:: ray_lightning.RayShardedStrategy
 ```
 
 
diff --git a/docs/horovod_faq.md b/docs/horovod_faq.md
new file mode 100644
index 00000000..6644e8b8
--- /dev/null
+++ b/docs/horovod_faq.md
@@ -0,0 +1,49 @@
+# Horovod installation issue
+
+> ```
+> Extension horovod.torch has not been built: /home/ubuntu/anaconda3/envs/tensorflow2_p38/lib/python3.8/site-packages/horovod/torch/mpi_lib/_mpi_lib.cpython-38-x86_64-linux-gnu.so not found
+> If this is not expected, reinstall Horovod with HOROVOD_WITH_PYTORCH=1 to debug the build error.
+>Warning! MPI libs are missing, but python applications are still avaiable.
+> ```
+
+One might fix this issue by 
+```python
+$ pip uninstall -y horovod
+$ conda install gcc_linux-64 gxx_linux-64
+$ [flags] pip install --no-cache-dir horovod
+```
+
+from [here](https://github.com/horovod/horovod/issues/656), [here](https://github.com/tlkh/ai-lab/issues/27) and [here](https://horovod.readthedocs.io/en/stable/install_include.html)
+
+- install horovod from scratch with torch 
+
+```python
+conda create -n hd python=3.8 scipy numpy pandas -y
+conda activate hd
+conda install pytorch=1.11 torchvision torchaudio cudatoolkit=11.3 -c pytorch -y
+sudo rm -rf /usr/local/cuda
+sudo ln -s /usr/local/cuda-11.3 /usr/local/cuda
+conda install gxx_linux-64 -y
+conda install cxx-compiler=1.0 -y
+export TORCH_CUDA_ARCH_LIST="3.7;5.0;6.0;7.0;7.5;8.0"
+echo $TORCH_CUDA_ARCH_LIST
+sudo apt-get purge -y cmake
+wget -q https://github.com/Kitware/CMake/releases/download/v3.20.2/cmake-3.20.2.tar.gz
+tar -zxvf cmake-3.20.2.tar.gz
+cd cmake-3.20.2
+./bootstrap -- -DCMAKE_USE_OPENSSL=OFF
+make -j10
+sudo make install
+cmake --version
+export CUDA_TOOLKIT_ROOT_DIR=/usr/local/cuda
+export HOROVOD_NCCL_HOME=/usr/local/cuda/
+export HOROVOD_NCCL_INCLUDE=/usr/local/cuda/include
+export TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST//";8.0"/}
+export HOROVOD_BUILD_CUDA_CC_LIST=${TORCH_CUDA_ARCH_LIST//";"/","}
+export HOROVOD_BUILD_CUDA_CC_LIST=${HOROVOD_BUILD_CUDA_CC_LIST//"."/""}
+export PATH=/usr/local/cuda/bin/:$PATH
+export HOROVOD_NCCL_LIB=/usr/local/cuda/lib/
+HOROVOD_NCCL_HOME=/usr/local/cuda HOROVOD_GPU_OPERATIONS=NCCL HOROVOD_WITH_PYTORCH=1 HOROVOD_WITHOUT_TENSORFLOW=1  HOROVOD_WITHOUT_MXNET=1 HOROVOD_WITHOUT_GLOO=1 pip install --no-cache-dir horovod
+```
+
+[reference 1](https://stackoverflow.com/questions/54948216/usr-lib-x86-64-linux-gnu-libstdc-so-6-version-glibcxx-3-4-21-not-found-req) and [reference 2](https://github.com/horovod/horovod/issues/401) and [reference 3](https://github.com/Lightning-AI/lightning/issues/4472) and [reference 4](https://github.com/horovod/horovod/issues/2276) and [reference 5](https://github.com/Lightning-AI/lightning/blob/master/dockers/base-cuda/Dockerfile#L105-L121) and [reference 6](https://horovod.readthedocs.io/en/stable/gpus_include.html) and [reference 7](https://horovod.readthedocs.io/en/stable/conda_include.html) and [reference 8](https://github.com/horovod/horovod/issues/3545) and [reference 9](https://github.com/KAUST-CTL/horovod-gpu-data-science-project) and [reference 10](https://kose-y.github.io/blog/2017/12/installing-cuda-aware-mpi/)
\ No newline at end of file
diff --git a/ray_lightning/__init__.py b/ray_lightning/__init__.py
index 82401e52..5939ac89 100644
--- a/ray_lightning/__init__.py
+++ b/ray_lightning/__init__.py
@@ -1,5 +1,5 @@
-from ray_lightning.ray_ddp import RayPlugin
-from ray_lightning.ray_horovod import HorovodRayPlugin
-from ray_lightning.ray_ddp_sharded import RayShardedPlugin
+from ray_lightning.ray_ddp import RayStrategy
+from ray_lightning.ray_horovod import HorovodRayStrategy
+from ray_lightning.ray_ddp_sharded import RayShardedStrategy
 
-__all__ = ["RayPlugin", "HorovodRayPlugin", "RayShardedPlugin"]
+__all__ = ["RayStrategy", "HorovodRayStrategy", "RayShardedStrategy"]
diff --git a/ray_lightning/accelerators/__init__.py b/ray_lightning/accelerators/__init__.py
new file mode 100644
index 00000000..72b274c7
--- /dev/null
+++ b/ray_lightning/accelerators/__init__.py
@@ -0,0 +1,21 @@
+# Copyright The PyTorch Lightning team.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from pytorch_lightning.accelerators.registry import \
+    call_register_accelerators  # noqa: F401
+from ray_lightning.accelerators.delayed_gpu_accelerator import _GPUAccelerator
+
+#  these lines are to register the delayed gpu accelerator as `_gpu`
+ACCELERATORS_BASE_MODULE = "ray_lightning.accelerators"
+call_register_accelerators(ACCELERATORS_BASE_MODULE)
+
+__all__ = ["_GPUAccelerator"]
diff --git a/ray_lightning/accelerators/delayed_gpu_accelerator.py b/ray_lightning/accelerators/delayed_gpu_accelerator.py
new file mode 100644
index 00000000..2bcc3638
--- /dev/null
+++ b/ray_lightning/accelerators/delayed_gpu_accelerator.py
@@ -0,0 +1,60 @@
+# Copyright The PyTorch Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import Dict, List
+
+import torch
+
+from pytorch_lightning.accelerators import Accelerator,\
+    GPUAccelerator
+
+
+class _GPUAccelerator(GPUAccelerator):
+    """Accelerator for GPU devices.
+
+    adapted from:
+    https://github.com/Lightning-AI/lightning/blob/master/src/pytorch_lightning/accelerators/gpu.py#L43
+    but remove `torch.cuda.set_device(root_device)` in `setup_environment`
+    """
+
+    def setup_environment(self, root_device: torch.device) -> None:
+        """
+        modified: remove `torch.cuda.set_device(root_device)`
+        and call `torch.cuda.set_device(self.device)` at the later time
+        inside the `ray_launcher` or `horovod_launcher`
+        """
+        Accelerator.setup_environment(self, root_device)
+
+    @staticmethod
+    def get_parallel_devices(devices: List[int]) -> List[torch.device]:
+        """Gets parallel devices for the Accelerator."""
+        # modified: return None when no devices are available
+        if devices:
+            return [torch.device("cuda", i) for i in devices]
+        else:
+            return None
+
+    @staticmethod
+    def is_available() -> bool:
+        # modified to always return True
+        return True
+
+    @classmethod
+    def register_accelerators(cls, accelerator_registry: Dict) -> None:
+        # the delayed gpu accelerator is registered as `_gpu`
+        # in the accelerator registry
+        accelerator_registry.register(
+            "_gpu",
+            cls,
+            description=f"{cls.__class__.__name__}",
+        )
diff --git a/ray_lightning/examples/ray_ddp_example.py b/ray_lightning/examples/ray_ddp_example.py
index ce87e10b..d8683fc1 100644
--- a/ray_lightning/examples/ray_ddp_example.py
+++ b/ray_lightning/examples/ray_ddp_example.py
@@ -11,7 +11,7 @@
 import ray
 from ray import tune
 from ray_lightning.tune import TuneReportCallback, get_tune_resources
-from ray_lightning import RayPlugin
+from ray_lightning import RayStrategy
 from ray_lightning.tests.utils import LightningMNISTClassifier
 
 
@@ -73,7 +73,7 @@ def train_mnist(config,
     trainer = pl.Trainer(
         max_epochs=num_epochs,
         callbacks=callbacks,
-        plugins=[RayPlugin(num_workers=num_workers, use_gpu=use_gpu)],
+        strategy=RayStrategy(num_workers=num_workers, use_gpu=use_gpu),
         **trainer_kwargs)
     trainer.fit(model)
 
diff --git a/ray_lightning/examples/ray_ddp_sharded_example.py b/ray_lightning/examples/ray_ddp_sharded_example.py
index 72711f18..e67150c6 100644
--- a/ray_lightning/examples/ray_ddp_sharded_example.py
+++ b/ray_lightning/examples/ray_ddp_sharded_example.py
@@ -10,7 +10,7 @@
 import pytorch_lightning as pl
 from pytorch_lightning import Callback
 
-from ray_lightning import RayShardedPlugin
+from ray_lightning import RayShardedStrategy
 
 
 class CUDACallback(Callback):
@@ -53,7 +53,7 @@ def download_data():
         with FileLock(os.path.join(data_dir, ".lock")):
             MNISTDataModule(data_dir=data_dir).prepare_data()
 
-    plugin = RayShardedPlugin(
+    strategy = RayShardedStrategy(
         num_workers=num_workers, use_gpu=use_gpu, init_hook=download_data)
 
     dm = MNISTDataModule(data_dir, batch_size=batch_size)
@@ -65,7 +65,7 @@ def download_data():
         max_epochs=max_epochs,
         precision=16 if use_gpu else 32,
         callbacks=[CUDACallback()] if use_gpu else [],
-        plugins=plugin,
+        strategy=strategy,
         max_steps=max_steps)
 
     trainer.fit(model, dm)
diff --git a/ray_lightning/examples/ray_ddp_tune.py b/ray_lightning/examples/ray_ddp_tune.py
index 62bdb4b0..0388cecc 100644
--- a/ray_lightning/examples/ray_ddp_tune.py
+++ b/ray_lightning/examples/ray_ddp_tune.py
@@ -8,7 +8,7 @@
 import ray
 from ray import tune
 from ray_lightning.tune import TuneReportCallback, get_tune_resources
-from ray_lightning import RayPlugin
+from ray_lightning import RayStrategy
 from ray_lightning.tests.utils import LightningMNISTClassifier
 
 
@@ -32,12 +32,8 @@ def download_data():
         max_epochs=num_epochs,
         callbacks=callbacks,
         progress_bar_refresh_rate=0,
-        plugins=[
-            RayPlugin(
-                num_workers=num_workers,
-                use_gpu=use_gpu,
-                init_hook=download_data)
-        ])
+        strategy=RayStrategy(
+            num_workers=num_workers, use_gpu=use_gpu, init_hook=download_data))
     dm = MNISTDataModule(
         data_dir=data_dir, num_workers=1, batch_size=config["batch_size"])
     trainer.fit(model, dm)
diff --git a/ray_lightning/examples/ray_horovod_example.py b/ray_lightning/examples/ray_horovod_example.py
index 2f50912f..d48f7494 100644
--- a/ray_lightning/examples/ray_horovod_example.py
+++ b/ray_lightning/examples/ray_horovod_example.py
@@ -11,7 +11,7 @@
 import ray
 from ray import tune
 from ray_lightning.tune import TuneReportCallback, get_tune_resources
-from ray_lightning import HorovodRayPlugin
+from ray_lightning import HorovodRayStrategy
 from ray_lightning.tests.utils import LightningMNISTClassifier
 
 
@@ -73,7 +73,7 @@ def train_mnist(config,
     trainer = pl.Trainer(
         max_epochs=num_epochs,
         callbacks=callbacks,
-        plugins=[HorovodRayPlugin(num_workers=num_workers, use_gpu=use_gpu)])
+        strategy=HorovodRayStrategy(num_workers=num_workers, use_gpu=use_gpu))
     trainer.fit(model)
 
 
diff --git a/ray_lightning/launchers/__init__.py b/ray_lightning/launchers/__init__.py
new file mode 100644
index 00000000..1a427676
--- /dev/null
+++ b/ray_lightning/launchers/__init__.py
@@ -0,0 +1,4 @@
+from ray_lightning.launchers.ray_launcher import RayLauncher
+from ray_lightning.launchers.ray_horovod_launcher import RayHorovodLauncher
+
+__all__ = ["RayLauncher", "RayHorovodLauncher"]
diff --git a/ray_lightning/launchers/ray_horovod_launcher.py b/ray_lightning/launchers/ray_horovod_launcher.py
new file mode 100644
index 00000000..941b8fca
--- /dev/null
+++ b/ray_lightning/launchers/ray_horovod_launcher.py
@@ -0,0 +1,251 @@
+from typing import Callable, Any, Optional
+
+import pytorch_lightning as pl
+from pytorch_lightning.strategies.launchers import _Launcher
+from pytorch_lightning.utilities.apply_func import apply_to_collection, \
+    move_data_to_device
+import numpy as np
+import torch
+
+import ray
+from pytorch_lightning.utilities.rank_zero import rank_zero_debug
+from pytorch_lightning.strategies import Strategy
+from ray.util.queue import Queue
+
+from ray_lightning.session import init_session
+from ray_lightning.util import process_results, Unavailable, to_state_stream, \
+    load_state_stream, set_cuda_device_if_used
+
+from ray_lightning.tune import TUNE_INSTALLED, is_session_enabled
+
+try:
+    import horovod.torch as hvd
+    from horovod.ray import RayExecutor
+except (ModuleNotFoundError, ImportError):
+    HOROVOD_AVAILABLE = False
+    RayExecutor = Unavailable
+    hvd = Unavailable
+else:
+    HOROVOD_AVAILABLE = True
+
+from pytorch_lightning.utilities import rank_zero_only
+from ray_lightning.accelerators import \
+    _GPUAccelerator  # noqa: F401
+from ray_lightning.launchers.utils import _RayOutput, get_executable_cls
+
+
+class RayHorovodLauncher(_Launcher):
+    def __init__(self, strategy: "Strategy") -> None:
+        """Initialize the Ray horovod launcher."""
+        self._strategy = strategy
+        self._executor = strategy.executor
+
+        if not ray.is_initialized():
+            ray.init()
+
+        self.tune_queue = None
+
+    @property
+    def global_rank(self) -> int:
+        """Return the global rank of the current process.
+            run on the worker node.
+        """
+        if not hvd.is_initialized():
+            return 0
+        return hvd.rank()
+
+    @property
+    def local_rank(self) -> int:
+        """Return the local rank of the current process.
+            run on the worker node.
+        """
+        if not hvd.is_initialized():
+            return 0
+        return hvd.local_rank()
+
+    @property
+    def world_size(self) -> int:
+        """Return the world size of the current process.
+            run on the worker node.
+        """
+        if not hvd.is_initialized():
+            return self.num_workers
+        return hvd.size()
+
+    def is_interactive_compatible(self) -> bool:
+        """Return whether the launcher is interactive compatible."""
+        return True
+
+    def launch(self,
+               function: Callable,
+               *args: Any,
+               trainer: Optional["pl.Trainer"] = None,
+               **kwargs: Any) -> Any:
+        """Launch the function on the workers and collect the results."""
+        ray_output = self.run_function_on_workers(
+            function, *args, trainer=trainer, **kwargs)
+
+        if trainer is None:
+            raise NotImplementedError(
+                "Ray launcher does not support trainer is None!")
+        self._recover_results_in_main_process(ray_output, trainer)
+        return_value = ray_output.trainer_results
+
+        return return_value
+
+    def run_function_on_workers(self,
+                                function: Callable,
+                                *args: Any,
+                                trainer: Optional["pl.Trainer"] = None,
+                                **kwargs: Any):
+        """Run the function on the workers and collect the results.
+           `executor.run_remote` is used to launch multiple ray remote tasks
+            to distributed training the model using the horovod backend.
+        """
+
+        # put the model as the ray object
+        # this reduce the memory comsumption
+        # and remove the model temporarily from the args
+        model = trainer.model
+        model_ref = ray.put(model)
+        trainer.model = None
+        new_args = tuple([None] + list(args[1:]))
+
+        # remove the executor temporarily from the args
+        # in order to avoid the ray.get() call in the function
+        # because executor is not pickleable
+        executor = self._executor
+        self._executor = None
+        self._strategy.executor = None
+
+        executor.start(executable_cls=get_executable_cls())
+
+        if TUNE_INSTALLED and is_session_enabled():
+            # Create communication queue and send to all the workers.
+            self.tune_queue = Queue(actor_options={"num_cpus": 0})
+
+        self._futures = executor.run_remote(lambda: self._wrapping_function(
+            function, model_ref, new_args, kwargs, self.tune_queue))
+
+        # put back the executor and model
+        self._executor = executor
+        self._strategy.executor = executor
+        trainer.model = model
+
+        results = process_results(self._futures, self.tune_queue)
+        executor.shutdown()
+        self._strategy.teardown()
+
+        return results[0]
+
+    def _wrapping_function(
+            self,
+            function: Callable,
+            model_ref: Any,
+            args: Any,
+            kwargs: Any,
+            tune_queue: Queue,
+    ) -> Any:
+        """Wrapping function to run the function on the workers.
+            `_wrapping_function` is run on each remote worker.
+            `function(*args, **kwargs)` is where the actual training happens.
+        """
+
+        self._strategy.set_remote(True)
+
+        # `function` is a trainer's class method
+        # in the ray remote tasks, its object `trainer` will also
+        # be copied when the function is remoted.
+        # ALERT: passing the trainer as an argument of `_wrapping_function`
+        # does not fillfullied our purpose. Ray remote tasks will
+        # create another copy of trainer so that
+        # `function.__self__ != trainer`, in which the side effect only
+        # happens to `function.__self__` when running
+        # `function(*args, **kwargs)`
+        trainer = function.__self__
+        model = ray.get(model_ref)
+        trainer.model = model
+        args = tuple([model] + list(args[1:]))
+
+        trainer._data_connector.prepare_data()
+
+        hvd.init()
+        rank_zero_only.rank = self.global_rank
+        set_cuda_device_if_used(trainer.strategy)
+
+        # Move the model to the appropriate device.
+        trainer.strategy.model_to_device()
+
+        if tune_queue is not None:
+            # Initialize session.
+            init_session(rank=self.global_rank, queue=tune_queue)
+
+        results = function(*args, **kwargs)
+
+        if trainer is not None:
+            results = self._collect_rank_zero_results(function.__self__,
+                                                      results)
+
+        if self.local_rank == 0:
+            return move_data_to_device(results, "cpu")
+
+        return None
+
+    def _collect_rank_zero_results(self, trainer: "pl.Trainer",
+                                   results: Any) -> Optional["_RayOutput"]:
+        """Collect the results from the rank zero process."""
+        rank_zero_debug("Finalizing the ray horovod launcher environment.")
+        checkpoint_callback = trainer.checkpoint_callback
+        best_model_path = checkpoint_callback.best_model_path \
+            if checkpoint_callback else None
+
+        state_dict = trainer.lightning_module.state_dict()
+
+        if self._strategy.global_rank != 0:
+            return None
+
+        # PyTorch Lightning saves the model weights in a temp file and
+        # loads it back on the driver.
+        # This won't work in a multi-node setup though, so we return the
+        # model state stream directly.
+        model_state_stream = to_state_stream(state_dict)
+
+        # adds the `callback_metrics`
+        callback_metrics: dict = apply_to_collection(
+            trainer.callback_metrics, torch.Tensor, lambda x: x.cpu().numpy(
+            ))  # send as numpy to avoid issues with memory sharing
+
+        # Same for logged_metrics
+        logged_metrics: dict = apply_to_collection(
+            trainer.logged_metrics, torch.Tensor, lambda x: x.cpu().numpy(
+            ))  # send as numpy to avoid issues with memory sharing
+
+        return _RayOutput(best_model_path, model_state_stream, trainer.state,
+                          results, callback_metrics, logged_metrics)
+
+    def _recover_results_in_main_process(self, ray_output: "_RayOutput",
+                                         trainer: "pl.Trainer") -> None:
+        """Recover the results in the main process."""
+        # transfer back the best path to the trainer
+        if trainer.checkpoint_callback:
+            trainer.checkpoint_callback.best_model_path = str(
+                ray_output.best_model_path)
+
+        if ray_output.weights_path is not None:
+            state_stream = ray_output.weights_path
+            # DDPSpawnPlugin.__recover_child_process_weights begin
+            # Difference here is that instead of writing the model weights to a
+            # file and loading it, we use the state dict of the model directly.
+            state_dict = load_state_stream(
+                state_stream, to_gpu=self._strategy.use_gpu)
+            # Set the state for PTL using the output from remote training.
+            trainer.lightning_module.load_state_dict(state_dict)
+
+        trainer.state = ray_output.trainer_state
+
+        trainer.callback_metrics.update(
+            apply_to_collection(ray_output.callback_metrics,
+                                np.ndarray, lambda x: torch.tensor(x)))
+        trainer.logged_metrics.update(
+            apply_to_collection(ray_output.logged_metrics,
+                                np.ndarray, lambda x: torch.tensor(x)))
diff --git a/ray_lightning/launchers/ray_launcher.py b/ray_lightning/launchers/ray_launcher.py
new file mode 100644
index 00000000..b7a267b5
--- /dev/null
+++ b/ray_lightning/launchers/ray_launcher.py
@@ -0,0 +1,342 @@
+from typing import Callable, List, Any, Tuple, Optional
+
+from collections import defaultdict
+import os
+
+import pytorch_lightning as pl
+from pytorch_lightning.strategies.launchers import _Launcher
+from pytorch_lightning.utilities.apply_func import apply_to_collection,\
+    move_data_to_device
+import numpy as np
+import torch
+
+import ray
+from ray import ObjectRef
+from pytorch_lightning.utilities.rank_zero import rank_zero_debug
+from ray.util.queue import Queue
+
+from ray_lightning.session import init_session
+from ray_lightning.util import process_results, to_state_stream, \
+    load_state_stream, set_cuda_device_if_used
+from ray_lightning.tune import TUNE_INSTALLED, is_session_enabled
+from pytorch_lightning.strategies import Strategy
+from ray_lightning.launchers.utils import _RayOutput, find_free_port,\
+    RayExecutor
+
+
+class RayLauncher(_Launcher):
+    def __init__(self, strategy: "Strategy") -> None:
+        """Initializes RayLauncher."""
+        self._strategy = strategy
+        self._start_method = "ray"
+        self._workers = []
+        self._futures = []
+        self._master_addr = None
+        self._master_port = None
+
+        self._global_to_local = None
+
+        self.tune_queue = None
+
+        if not ray.is_initialized():
+            ray.init()
+
+    def is_interactive_compatible(self) -> bool:
+        """Returns True if the launcher is interactive compatible."""
+        return True
+
+    def launch(self,
+               function: Callable,
+               *args: Any,
+               trainer: Optional["pl.Trainer"] = None,
+               **kwargs: Any) -> Any:
+        """Launches the function on the workers from the driver node."""
+        self.setup_workers()
+        ray_output = self.run_function_on_workers(
+            function, *args, trainer=trainer, **kwargs)
+
+        if trainer is None:
+            raise NotImplementedError(
+                "Ray launcher does not support trainer is None!")
+        self._recover_results_in_main_process(ray_output, trainer)
+        return_value = ray_output.trainer_results
+
+        self.teardown_workers()
+        self._strategy.teardown()
+        return return_value
+
+    def setup_workers(self, tune_enabled: bool = True) -> None:
+        """Creates the Ray actors and sets up PTL Trainer environment
+            on the worker nodes.
+        """
+        self._workers = [
+            self._create_worker() for _ in range(self._strategy.num_workers)
+        ]
+        if self._strategy.init_hook:
+            ray.get([
+                w.execute.remote(self._strategy.init_hook)
+                for w in self._workers
+            ])
+
+        self._master_addr = ray.get(self._workers[0].get_node_ip.remote())
+        self._master_port = str(
+            ray.get(self._workers[0].execute.remote(find_free_port)))
+
+        # Sets environment variables for all workers.
+        # This will set the MASTER_ADDR and MASTER_PORT on each Ray actor.
+        self._setup_env_vars()
+
+        if self._strategy.use_gpu:
+            # Set the CUDA_VISIBLE_DEVICES for all workers.
+            self._share_cuda_visible_devices()
+
+        # Get the mapping from global ranks to the respective local ranks.
+        self._global_to_local = self.get_local_ranks()
+        # Todo: put model into object store?
+
+        if tune_enabled and TUNE_INSTALLED and is_session_enabled():
+            # Create communication queue and send to all the workers.
+            self.tune_queue = Queue(actor_options={"num_cpus": 0})
+
+    def _create_worker(self) -> ray.actor.ActorHandle:
+        """Creates Ray actor workers."""
+        worker = RayExecutor.options(
+            num_cpus=self._strategy.num_cpus_per_worker,
+            num_gpus=self._strategy.num_gpus_per_worker,
+            resources=self._strategy.additional_resources_per_worker).remote()
+        return worker
+
+    def teardown_workers(self):
+        """Tears down the Ray actors and PTL Trainer environment"""
+        if self.tune_queue:
+            # Shutdown the queue.
+            self.tune_queue.shutdown()
+
+        for w in self._workers:
+            ray.kill(w, no_restart=True)
+            del w
+        self._workers = []
+
+    def get_local_ranks(self) -> List[Optional[Tuple[int, int]]]:
+        """Creates a mapping of global ranks to local ranks/node ranks.
+            this method is to run on the worker nodes.
+        """
+        # Get the local ranks for all the workers and store as a list.
+        # First get the IP address of each remote worker.
+        node_ips = ray.get([w.get_node_ip.remote() for w in self._workers])
+
+        node_rank_map = {}
+        counter = 0
+        for ip in node_ips:
+            # If this is a new IP address, then increment counter.
+            if ip not in node_rank_map:
+                node_rank_map[ip] = counter
+                counter += 1
+
+        rank_counter_dict = defaultdict(int)
+        global_to_local = [None] * self._strategy.num_workers
+
+        for global_rank in range(self._strategy.num_workers):
+            ip = node_ips[global_rank]
+            global_to_local[global_rank] = (
+                rank_counter_dict[ip],  # local rank
+                node_rank_map[ip])  # node rank
+            rank_counter_dict[ip] += 1
+
+        return global_to_local
+
+    def _setup_env_vars(self):
+        """Sets environment variables for all workers."""
+        # Get rank 0 worker address and port for DDP connection.
+        os.environ["MASTER_ADDR"] = self._master_addr
+        os.environ["MASTER_PORT"] = self._master_port
+
+        # Set environment variables for remote workers.
+        keys = [
+            "PL_GLOBAL_SEED", "PL_TORCH_DISTRIBUTED_BACKEND", "MASTER_ADDR",
+            "MASTER_PORT"
+        ]
+        values = [os.getenv(k) for k in keys]
+
+        ray.get([w.set_env_vars.remote(keys, values) for w in self._workers])
+
+    def _share_cuda_visible_devices(self):
+        """Sets CUDA_VISIBLE_DEVICES on all workers.
+        For each worker, CUDA_VISIBLE_DEVICES will be set to the GPU IDs
+        visible to all workers on that worker's node.
+        This allows GPU workers on the same node to communicate with one
+        another.
+        Example:
+            Setup:
+            - Node1:
+                - Worker1: {0, 1}
+                - Worker2: {2, 3}
+            - Node2:
+                - Worker3: {0, 1}
+            CUDA_VISIBLE_DEVICES:
+            - Worker1: "0,1,2,3"
+            - Worker2: "0,1,2,3"
+            - Worker2: "0,1"
+        """
+        node_ids_and_gpu_ids = ray.get(
+            [w.get_node_and_gpu_ids.remote() for w in self._workers])
+
+        node_id_to_worker_id = defaultdict(set)
+        node_id_to_gpu_ids = defaultdict(set)
+
+        for worker_id, (node_id, gpu_ids) in enumerate(node_ids_and_gpu_ids):
+            node_id_to_worker_id[node_id].add(worker_id)
+            node_id_to_gpu_ids[node_id].update(gpu_ids)
+
+        futures = []
+        for node_id, gpu_ids in node_id_to_gpu_ids.items():
+            all_gpu_ids = ",".join([str(gpu_id) for gpu_id in gpu_ids])
+
+            def set_gpu_ids():
+                os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+                os.environ["CUDA_VISIBLE_DEVICES"] = all_gpu_ids
+
+            for worker_id in node_id_to_worker_id[node_id]:
+                futures.append(
+                    self._workers[worker_id].execute.remote(set_gpu_ids))
+        ray.get(futures)
+
+    def run_function_on_workers(self,
+                                function: Callable,
+                                *args: Any,
+                                trainer: Optional["pl.Trainer"] = None,
+                                **kwargs: Any):
+        """launch a function on all workers.
+            The actual training parts are run inside `_wrapping_function`
+        """
+        # put the model as the ray object
+        # and remove the model temporarily from the args
+        model = trainer.model
+        model_ref = ray.put(model)
+        trainer.model = None
+        new_args = tuple([None] + list(args[1:]))
+
+        # train the model and get the result to rank 0 node
+        self._futures = [
+            w.execute.remote(self._wrapping_function, i, self._global_to_local,
+                             function, model_ref, new_args, kwargs,
+                             self.tune_queue)
+            for i, w in enumerate(self._workers)
+        ]
+
+        trainer.model = model
+
+        results = process_results(self._futures, self.tune_queue)
+        return results[0]
+
+    def _wrapping_function(
+            self,
+            global_rank: int,
+            global_to_local: List[Optional[Tuple[int, int]]],
+            function: Callable,
+            model_ref: ObjectRef,
+            args: Any,
+            kwargs: Any,
+            tune_queue: Queue,
+    ) -> Any:
+        """Wraps the function to run on the workers.
+            `results = function(*args, **kwargs)` is where the
+            actual training parts are run.
+        """
+        self._strategy.set_remote(True)
+        self._strategy.set_global_to_local(global_to_local)
+
+        # `function` is a trainer's class method
+        # in the ray remote tasks, its object `trainer` will also
+        # be copied when the function is remoted.
+        # ALERT: passing the trainer as an argument of `_wrapping_function`
+        # does not fillfullied our purpose. Ray remote tasks will
+        # create another copy of trainer so that
+        # `function.__self__ != trainer`, in which the side effect only
+        # happens to `function.__self__` when running
+        # `function(*args, **kwargs)`
+        trainer = function.__self__
+        trainer.model = model_ref
+        args = tuple([model_ref] + list(args[1:]))
+
+        trainer._data_connector.prepare_data()
+        if tune_queue is not None:
+            # Initialize session.
+            init_session(rank=global_rank, queue=tune_queue)
+
+        self._strategy._worker_setup(process_idx=global_rank)
+        trainer.strategy.root_device = self._strategy.root_device
+        trainer.strategy.global_rank = self._strategy.global_rank
+        trainer.strategy.local_rank = self._strategy.local_rank
+        set_cuda_device_if_used(trainer.strategy)
+
+        results = function(*args, **kwargs)
+
+        if trainer is not None:
+            results = self._collect_rank_zero_results(trainer, results)
+
+        if trainer.strategy.local_rank == 0:
+            return move_data_to_device(results, "cpu")
+
+        trainer._teardown()
+        trainer._call_teardown_hook()
+        return None
+
+    def _collect_rank_zero_results(self, trainer: "pl.Trainer",
+                                   results: Any) -> Optional["_RayOutput"]:
+        """Collects the results from the worker node 0."""
+        rank_zero_debug("Finalizing the Ray launcher environment.")
+        checkpoint_callback = trainer.checkpoint_callback
+        best_model_path = checkpoint_callback.best_model_path \
+            if checkpoint_callback else None
+
+        state_dict = trainer.lightning_module.state_dict()
+
+        if self._strategy.global_rank != 0:
+            return None
+
+        # PyTorch Lightning saves the model weights in a temp file and
+        # loads it back on the driver.
+        # This won't work in a multi-node setup though, so we return the
+        # model state stream directly.
+        model_state_stream = to_state_stream(state_dict)
+
+        # adds the `callback_metrics`
+        callback_metrics: dict = apply_to_collection(
+            trainer.callback_metrics, torch.Tensor, lambda x: x.cpu().numpy(
+            ))  # send as numpy to avoid issues with memory sharing
+
+        # Same for logged_metrics
+        logged_metrics: dict = apply_to_collection(
+            trainer.logged_metrics, torch.Tensor, lambda x: x.cpu().numpy(
+            ))  # send as numpy to avoid issues with memory sharing
+
+        return _RayOutput(best_model_path, model_state_stream, trainer.state,
+                          results, callback_metrics, logged_metrics)
+
+    def _recover_results_in_main_process(self, ray_output: "_RayOutput",
+                                         trainer: "pl.Trainer") -> None:
+        """Recovers the results in the main process."""
+        # transfer back the best path to the trainer
+        if trainer.checkpoint_callback:
+            trainer.checkpoint_callback.best_model_path = str(
+                ray_output.best_model_path)
+
+        if ray_output.weights_path is not None:
+            state_stream = ray_output.weights_path
+            # DDPSpawnPlugin.__recover_child_process_weights begin
+            # Difference here is that instead of writing the model weights to a
+            # file and loading it, we use the state dict of the model directly.
+            state_dict = load_state_stream(
+                state_stream, to_gpu=self._strategy.use_gpu)
+            # Set the state for PTL using the output from remote training.
+            trainer.lightning_module.load_state_dict(state_dict)
+
+        trainer.state = ray_output.trainer_state
+
+        trainer.callback_metrics.update(
+            apply_to_collection(ray_output.callback_metrics,
+                                np.ndarray, lambda x: torch.tensor(x)))
+        trainer.logged_metrics.update(
+            apply_to_collection(ray_output.logged_metrics,
+                                np.ndarray, lambda x: torch.tensor(x)))
diff --git a/ray_lightning/launchers/utils.py b/ray_lightning/launchers/utils.py
new file mode 100644
index 00000000..8ecf2e29
--- /dev/null
+++ b/ray_lightning/launchers/utils.py
@@ -0,0 +1,69 @@
+from typing import Any, Optional, NamedTuple, Dict, List, Callable
+from pytorch_lightning.utilities.types import _PATH
+from pytorch_lightning.trainer.states import TrainerState
+
+from contextlib import closing
+import socket
+
+import ray
+import os
+
+
+def find_free_port():
+    """ Find a free port on the machines. """
+    with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
+        s.bind(("", 0))
+        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        return s.getsockname()[1]
+
+
+def get_executable_cls():
+    # Only used for testing purposes, currently.
+    # Only used in `ray_horovod_launcher.py`
+    # We need to override this in tests to ensure test path is set correctly.
+    return None
+
+
+@ray.remote
+class RayExecutor:
+    """A class to execute any arbitrary function remotely."""
+
+    def set_env_var(self, key: str, value: str):
+        """Set an environment variable with the provided values."""
+        if value is not None:
+            value = str(value)
+            os.environ[key] = value
+
+    def set_env_vars(self, keys: List[str], values: List[str]):
+        """Sets multiple env vars with the provided values"""
+        assert len(keys) == len(values)
+        for key, value in zip(keys, values):
+            self.set_env_var(key, value)
+
+    def get_node_ip(self):
+        """Returns the IP address of the node that this Ray actor is on."""
+        return ray.util.get_node_ip_address()
+
+    def get_node_and_gpu_ids(self):
+        return ray.get_runtime_context().node_id.hex(), ray.get_gpu_ids()
+
+    def execute(self, fn: Callable, *args, **kwargs):
+        """Execute the provided function and return the result."""
+        return fn(*args, **kwargs)
+
+
+class _RayOutput(NamedTuple):
+    """Ray output tuple with the following fields:
+       - `best_model_path`: path to the best model
+       - `weights_path`: path to the weights
+       - `trainer_state`: trainer state
+       - `trainer_results`: trainer result
+       - `callback_results`: callback result
+       - `logged_metrics`: logged metrics
+    """
+    best_model_path: Optional[_PATH]
+    weights_path: Optional[_PATH]
+    trainer_state: TrainerState
+    trainer_results: Any
+    callback_metrics: Dict[str, Any]
+    logged_metrics: Dict[str, Any]
diff --git a/ray_lightning/ray_ddp.py b/ray_lightning/ray_ddp.py
index 38e88adb..a90f4db2 100644
--- a/ray_lightning/ray_ddp.py
+++ b/ray_lightning/ray_ddp.py
@@ -1,87 +1,40 @@
 from typing import Callable, Dict, List, Union, Any, Tuple, Optional
 
-from collections import defaultdict
-from contextlib import closing
-import os
-import socket
 import warnings
 
-import numpy as np
 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
-from pytorch_lightning.utilities import rank_zero_only, rank_zero_info
-from pytorch_lightning.utilities.apply_func import apply_to_collection
-from pytorch_lightning.utilities.seed import reset_seed
+from pytorch_lightning.strategies import DDPSpawnStrategy
+from pytorch_lightning.utilities.rank_zero import rank_zero_only
 
 import ray
+from pytorch_lightning.utilities.rank_zero import rank_zero_info
+from pytorch_lightning.utilities.seed import reset_seed, log
 from ray.util import PublicAPI
-from ray.util.queue import Queue
-
-from ray_lightning.session import init_session
-from ray_lightning.util import process_results, to_state_stream, \
-    load_state_stream
-from ray_lightning.tune import TUNE_INSTALLED, is_session_enabled
-
-
-def find_free_port():
-    with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
-        s.bind(("", 0))
-        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
-        return s.getsockname()[1]
-
 
-@ray.remote
-class RayExecutor:
-    """A class to execute any arbitrary function remotely."""
+from ray_lightning.launchers import RayLauncher
+from ray_lightning.accelerators import \
+    _GPUAccelerator  # noqa: F401
 
-    def set_env_var(self, key: str, value: str):
-        """Set an environment variable with the provided values."""
-        if value is not None:
-            value = str(value)
-            os.environ[key] = value
-
-    def set_env_vars(self, keys: List[str], values: List[str]):
-        """Sets multiple env vars with the provided values"""
-        assert len(keys) == len(values)
-        for key, value in zip(keys, values):
-            self.set_env_var(key, value)
-
-    def get_node_ip(self):
-        """Returns the IP address of the node that this Ray actor is on."""
-        return ray.util.get_node_ip_address()
-
-    def get_node_and_gpu_ids(self):
-        return ray.get_runtime_context().node_id.hex(), ray.get_gpu_ids()
-
-    def execute(self, fn: Callable, *args, **kwargs):
-        """Execute the provided function and return the result."""
-        return fn(*args, **kwargs)
+import os
 
 
 @PublicAPI(stability="beta")
-class RayPlugin(DDPSpawnPlugin):
-    """Pytorch Lightning plugin for DDP training on a Ray cluster.
-
-    This plugin is used to manage distributed training using DDP and
+class RayStrategy(DDPSpawnStrategy):
+    """Pytorch Lightning strategy for DDP training on a Ray cluster.
+    This strategy is used to manage distributed training using DDP and
     Ray for process launching. Internally, the specified number of
     Ray actors are launched in the cluster and are registered as part of a
     Pytorch DDP process group. The Pytorch Lightning trainer is instantiated
     on the driver and sent to each of these training workers where training is
     executed. The distributed training protocol is handled by Pytorch DDP.
-
     Each training worker is configured to reserve ``num_cpus_per_worker``
     CPUS and 1 GPU if ``use_gpu`` is set to ``True``.
-
-    If using this plugin, you should run your code like a normal Python
+    If using this strategy, you should run your code like a normal Python
     script: ``python train.py``, and only on the head node if running in a
     distributed Ray cluster. There is no need to run this script on every
     single node.
-
     Args:
         num_workers (int): Number of training workers to use.
         num_cpus_per_worker (int): Number of CPUs per worker.
@@ -96,25 +49,21 @@ class RayPlugin(DDPSpawnPlugin):
             override the number of CPU/GPUs used by each worker.
         **ddp_kwargs: Additional arguments to pass into
             ``DistributedDataParallel`` initialization
-
     Example:
-
         .. code-block:: python
-
             import pytorch_lightning as ptl
             from ray_lightning import RayAccelerator
-
             ptl_model = MNISTClassifier(...)
-            plugin = RayPlugin(num_workers=4, cpus_per_worker=1,
+            strategy = RayStrategy(num_workers=4, cpus_per_worker=1,
                 use_gpu=True)
-
             # Don't set ``gpus`` in ``Trainer``.
             # The actual number of GPUs is determined by ``num_workers``.
-            trainer = pl.Trainer(..., plugins=[plugin])
+            trainer = pl.Trainer(..., strategy=strategy)
             trainer.fit(ptl_model)
-
     """
 
+    strategy_name = "ddp_ray"
+
     def __init__(self,
                  num_workers: int = 1,
                  num_cpus_per_worker: int = 1,
@@ -122,13 +71,11 @@ def __init__(self,
                  init_hook: Optional[Callable] = None,
                  resources_per_worker: Optional[Dict] = None,
                  **ddp_kwargs: Union[Any, Dict[str, Any]]):
-        if not ray.is_initialized():
-            ray.init()
-        self._is_remote = False
+        """Initialize the Ray strategy."""
         resources_per_worker = resources_per_worker if resources_per_worker \
             else {}
         self.nickname = "ddp_ray"
-        self.num_workers = num_workers
+        self.num_workers = int(num_workers)
         self.num_cpus_per_worker = resources_per_worker.pop(
             "CPU", num_cpus_per_worker)
 
@@ -151,256 +98,39 @@ def __init__(self,
                           "environment variable.")
 
         self.additional_resources_per_worker = resources_per_worker
-        self.workers = []
         self.init_hook = init_hook
 
         self._local_rank = 0
         self._global_rank = 0
         self._node_rank = 0
 
-        # Pop these arguments so LightningCLI doesn't complain about them being
-        # passed in twice.
-        # https://github.com/ray-project/ray_lightning/issues/151
-        # TODO(amogkam): Remove after
-        #  https://github.com/PyTorchLightning/pytorch-lightning/discussions/8561
-        #  is fixed.
-        ddp_kwargs.pop("parallel_devices", None)
-        ddp_kwargs.pop("cluster_environment", None)
+        self._is_remote = False
+        self._device = None
 
         super().__init__(
-            parallel_devices=[], cluster_environment=None, **ddp_kwargs)
-
-    def __getstate__(self):
-        d = self.__dict__.copy()
-        # Don't serialize the workers.
-        del d["workers"]
-        return d
-
-    def __setstate__(self, d):
-        d["workers"] = []
-        self.__dict__.update(d)
-
-    def _create_worker(self):
-        """Creates Ray actor."""
-        worker = RayExecutor.options(
-            num_cpus=self.num_cpus_per_worker,
-            num_gpus=self.num_gpus_per_worker,
-            resources=self.additional_resources_per_worker).remote()
-        return worker
-
-    def setup(self):
-        """Sets up PTL Trainer and creates the Ray actors."""
-        self.workers = [self._create_worker() for _ in range(self.num_workers)]
-        if self.init_hook:
-            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
-
-    def _setup_env_vars(self):
-        # Get rank 0 worker address and port for DDP connection.
-        os.environ["MASTER_ADDR"] = ray.get(
-            self.workers[0].get_node_ip.remote())
-        os.environ["MASTER_PORT"] = str(
-            ray.get(self.workers[0].execute.remote(find_free_port)))
-
-        # Set environment variables for remote workers.
-        keys = [
-            "PL_GLOBAL_SEED", "PL_TORCH_DISTRIBUTED_BACKEND", "MASTER_ADDR",
-            "MASTER_PORT"
-        ]
-        values = [os.getenv(k) for k in keys]
-        ray.get([w.set_env_vars.remote(keys, values) for w in self.workers])
-
-    def _share_cuda_visible_devices(self):
-        """Sets CUDA_VISIBLE_DEVICES on all workers.
-
-        For each worker, CUDA_VISIBLE_DEVICES will be set to the GPU IDs
-        visible to all workers on that worker's node.
-
-        This allows GPU workers on the same node to communicate with one
-        another.
-
-        Example:
-
-            Setup:
-            - Node1:
-                - Worker1: {0, 1}
-                - Worker2: {2, 3}
-            - Node2:
-                - Worker3: {0, 1}
-
-            CUDA_VISIBLE_DEVICES:
-            - Worker1: "0,1,2,3"
-            - Worker2: "0,1,2,3"
-            - Worker2: "0,1"
-
-        """
-        node_ids_and_gpu_ids = ray.get(
-            [w.get_node_and_gpu_ids.remote() for w in self.workers])
-
-        node_id_to_worker_id = defaultdict(set)
-        node_id_to_gpu_ids = defaultdict(set)
-
-        for worker_id, (node_id, gpu_ids) in enumerate(node_ids_and_gpu_ids):
-            node_id_to_worker_id[node_id].add(worker_id)
-            node_id_to_gpu_ids[node_id].update(gpu_ids)
-
-        futures = []
-        for node_id, gpu_ids in node_id_to_gpu_ids.items():
-            all_gpu_ids = ",".join([str(gpu_id) for gpu_id in gpu_ids])
-
-            def set_gpu_ids():
-                os.environ["CUDA_VISIBLE_DEVICES"] = all_gpu_ids
-
-            for worker_id in node_id_to_worker_id[node_id]:
-                futures.append(
-                    self.workers[worker_id].execute.remote(set_gpu_ids))
-        ray.get(futures)
-
-    def start_training(self, trainer):
-        results = self.execution_loop(tune_enabled=True)
-        # reset optimizers, since main process is never used for training and
-        # thus does not have a valid optim state.
-        trainer.optimizers = []
-        return results
-
-    def start_evaluating(self, trainer):
-        results = self.execution_loop(tune_enabled=False)
-        return results
-
-    def start_predicting(self, trainer):
-        results = self.execution_loop(tune_enabled=False)
-        return results
-
-    def get_local_ranks(self) -> Dict[int, Tuple[int, int]]:
-        """Creates a mapping of global ranks to local ranks/node ranks."""
-        # Get the local ranks for all the workers and store as a dict.
-        # First get the IP address of each remote worker.
-        node_ips = ray.get([w.get_node_ip.remote() for w in self.workers])
-
-        node_rank_map = {}
-        counter = 0
-        for ip in node_ips:
-            # If this is a new IP address, then increment counter.
-            if ip not in node_rank_map:
-                node_rank_map[ip] = counter
-                counter += 1
-
-        rank_counter_dict = defaultdict(int)
-        global_to_local = [None] * self.num_workers
-
-        for global_rank in range(self.num_workers):
-            ip = node_ips[global_rank]
-            global_to_local[global_rank] = (
-                rank_counter_dict[ip],  # local rank
-                node_rank_map[ip])  # node rank
-            rank_counter_dict[ip] += 1
-
-        return global_to_local
-
-    def execution_loop(self, tune_enabled: bool = True):
-        """Main execution loop for training, testing, & prediction.
-
-        Sets up the torch.distributed process group for each
-        worker. Then trigger remote training/testing/eval via
-        ``train_remote`` on each worker. If using with Ray Tune, create a
-        communication queue to retrieve intermediate results, and process
-        those results. Finally retrieve the training results from the rank 0
-        worker and return.
+            accelerator="_gpu" if use_gpu else "cpu",
+            parallel_devices=[],
+            cluster_environment=None,
+            **ddp_kwargs)
+
+    def _configure_launcher(self):
+        """Configure the Ray launcher.
+            the distributed training logic is handled by the launcher.
         """
+        self._launcher = RayLauncher(self)
 
-        # Sets environment variables for all workers.
-        # This will set the MASTER_ADDR and MASTER_PORT on each Ray actor.
-        self._setup_env_vars()
-
-        if self.use_gpu:
-            # Set the CUDA_VISIBLE_DEVICES for all workers.
-            self._share_cuda_visible_devices()
-
-        # Get the mapping from global ranks to the respective local ranks.
-        self.global_to_local = self.get_local_ranks()
-
-        model = self._model
-        model_ref = ray.put(model)
-        # Don't pickle the model when training remotely.
-        self._model = None
-
-        queue = None
-        if tune_enabled and TUNE_INSTALLED and is_session_enabled():
-            # Create communication queue and send to all the workers.
-            queue = Queue(actor_options={"num_cpus": 0})
-
-        self._futures = [
-            self.workers[i].execute.remote(self.execute_remote, model_ref, i,
-                                           queue)
-            for i in range(self.num_workers)
-        ]
-
-        process_results(self._futures, queue)
-
-        self._model = model
-        if queue:
-            # Shutdown the queue.
-            queue.shutdown()
-
-    def post_dispatch(self, trainer: "pl.Trainer"):
-        """Finalized fitting.
-
-        1. Load model weights on driver.
-        2. Shutdown DDP process group and all Ray actors.
-
-        Shutdown the DDP process group and all the Ray
-        actors."""
-
-        results = ray.get(self._futures)
-        # Get the results, checkpoint path, and model weights from worker 0.
-        results, best_path, state_stream, callback_metrics, logged_metrics \
-            = results[0]
-        self._results = results
-
-        # From DDPSpawn.get_queue
-        self.lightning_module.trainer.callback_metrics.update(
-            apply_to_collection(callback_metrics,
-                                np.ndarray, lambda x: torch.tensor(x)))
-
-        # Same for logged_metrics
-        self.lightning_module.trainer.logged_metrics.update(
-            apply_to_collection(logged_metrics,
-                                np.ndarray, lambda x: torch.tensor(x)))
-
-        # DDPSpawnPlugin.__recover_child_process_weights begin
-        # Difference here is that instead of writing the model weights to a
-        # file and loading it, we use the state dict of the model directly.
-        state_dict = load_state_stream(state_stream, to_gpu=self.use_gpu)
-        # Set the state for PTL using the output from remote training.
-        self.lightning_module.load_state_dict(state_dict)
-        if self.lightning_module.trainer.checkpoint_callback:
-            self.lightning_module.trainer.checkpoint_callback \
-                .best_model_path = best_path
-        # DDPSpawnPlugin.__recover_child_process_weights_end
-
-        for w in self.workers:
-            ray.kill(w, no_restart=True)
-            del w
-        self.workers = []
-
-    # All methods below are only executed in remote Ray workers.
+    def setup(self, trainer: "pl.Trainer") -> None:
+        """Setup the strategy."""
+        super().setup(trainer)
+
+    def set_remote(self, remote: bool):
+        """Set the remote flag. (this is useful for the remote workers)"""
+        self._is_remote = remote
+
+    def set_global_to_local(self,
+                            global_to_local: List[Optional[Tuple[int, int]]]):
+        """Set the global to local rank mapping."""
+        self.global_to_local = global_to_local
 
     def set_world_ranks(self, process_idx: int = 0):
         """Set the appropriate rank attributes for the trainer."""
@@ -415,137 +145,118 @@ def set_world_ranks(self, process_idx: int = 0):
                 self.global_rank]
 
     def _worker_setup(self, process_idx: int):
+        """Setup the workers and pytorch DDP connections."""
         reset_seed()
         self.set_world_ranks(process_idx)
         rank_zero_only.rank = self.global_rank
+        self._process_group_backend = self._get_process_group_backend()
+
+        # Copied from
+        # pytorch_lightning.utilities.distributed.init_dist_connection
+        if not torch.distributed.is_available():
+            raise RuntimeError("torch.distributed is not available. "
+                               "Cannot initialize distributed process group")
+
+        if torch.distributed.is_initialized():
+            log.debug(
+                "torch.distributed is already initialized. Exiting early")
+            return
+
+        global_rank = self.global_rank
+        world_size = self.world_size
+        torch_distributed_backend = self.torch_distributed_backend
 
         # Taken from pytorch_lightning.utilities.distributed
-        # .init_dist_connection
-        # Modified to not use cluster environment.
         if torch.distributed.is_available(
         ) and not torch.distributed.is_initialized():
-            log.info(
-                f"initializing distributed: GLOBAL_RANK: {self.global_rank}, "
-                f"MEMBER: {self.global_rank + 1}/{self.world_size}")
+            log.info(f"Initializing distributed: GLOBAL_RANK: {global_rank}, "
+                     f"MEMBER: {global_rank + 1}/{world_size}")
             torch.distributed.init_process_group(
-                self.torch_distributed_backend,
-                rank=self.global_rank,
-                world_size=self.world_size)
-
-            # on rank=0 let everyone know training is starting
-            rank_zero_info(
-                f"{'-' * 100}\n"
-                f"distributed_backend={self.torch_distributed_backend}\n"
-                f"All distributed processes registered. Starting with "
-                f"{self.world_size} processes\n"
-                f"{'-' * 100}\n")
-
-    def execute_remote(self,
-                       model: LightningModule,
-                       global_rank: int,
-                       queue: Queue = None):
-        """Train/test/eval function to be executed on each remote worker.
-
-        Modified from DDPSpawn._wrapped_function and DDPSpawn.new_process
-
-        """
-        assert isinstance(self, RayPlugin)
-        # This method should be executed remotely in each worker.
-        self._model = model
-        self.lightning_module.trainer._accelerator_connector \
-            ._training_type_plugin = self
-        self.lightning_module.trainer._accelerator_connector.accelerator \
-            .training_type_plugin = self
-
-        # TODO: See if this is necessary.
-        self.lightning_module.trainer._data_connector.prepare_data()
-
-        # Set _is_remote to True so that self.set_world_ranks will
-        # properly set the ranks.
-        self._is_remote = True
-
-        if queue is not None:
-            # Initialize session.
-            init_session(rank=global_rank, queue=queue)
-
-        self._worker_setup(process_idx=global_rank)
-
-        # Below is modified from DDPSpawn.new_process
-
-        # Move the model to the correct device.
-        self.model_to_device()
-
-        # TODO: Support syncbatchnorm.
-        # skip wrapping the model if we are not fitting as no gradients
-        # need to be exchanged.
-        trainer_fn = self.lightning_module.trainer.state.fn
-        if trainer_fn == TrainerFn.FITTING:
-            self.configure_ddp()
-
-        self.barrier()
-
-        results = self.lightning_module.trainer.run_stage()
-
-        # __transfer_distrib_spawn_state_on_fit_end start
-        if self.global_rank == 0:
-            checkpoint_callback = \
-                self.lightning_module.trainer.checkpoint_callback
-            best_model_path = checkpoint_callback.best_model_path if \
-                checkpoint_callback else None
-
-            # PyTorch Lightning saves the model weights in a temp file and
-            # loads it back on the driver.
-            # This won't work in a multi-node setup though, so we return the
-            # model state stream directly.
-            model_state_stream = to_state_stream(
-                self.lightning_module.state_dict())
-
-            # From DDPSpawn.add_to_queue
-            callback_metrics: dict = apply_to_collection(
-                self.lightning_module.trainer.callback_metrics,
-                torch.Tensor, lambda x: x.cpu().numpy(
-                ))  # send as numpy to avoid issues with memory sharing
-
-            # Same for logged_metrics
-            logged_metrics: dict = apply_to_collection(
-                self.lightning_module.trainer.logged_metrics,
-                torch.Tensor, lambda x: x.cpu().numpy(
-                ))  # send as numpy to avoid issues with memory sharing
-
-            return_val = results, best_model_path, model_state_stream, \
-                callback_metrics, logged_metrics
-        else:
-            return_val = None
-        # __transfer_distrib_spawn_state_on_fit_end end
-
-        self.lightning_module.trainer._call_teardown_hook()
-
-        return return_val
+                torch_distributed_backend,
+                rank=global_rank,
+                world_size=world_size,
+                init_method="env://")
+
+        # on rank=0 let everyone know training is starting
+        rank_zero_info(f"{'-' * 100}\n"
+                       f"distributed_backend={torch_distributed_backend}\n"
+                       f"All distributed processes registered. "
+                       f"Starting with {world_size} processes\n"
+                       f"{'-' * 100}\n")
 
     @property
     def world_size(self) -> int:
+        """Return the world size."""
         return self.num_workers
 
     @property
     def local_rank(self) -> int:
+        """Return the local rank."""
         return self._local_rank
 
+    @local_rank.setter
+    def local_rank(self, value: int):
+        """Set the local rank."""
+        self._local_rank = value
+
     @property
     def global_rank(self) -> int:
+        """Return the global rank."""
         return self._global_rank
 
+    @global_rank.setter
+    def global_rank(self, value: int):
+        """Set the global rank."""
+        self._global_rank = value
+
     @property
     def node_rank(self) -> int:
+        """Return the node rank."""
         return self._node_rank
 
     @property
     def root_device(self):
+        """Return the root device."""
+        # get the root device
+        # if the root device not set, figure it out
+        # thru `get_gpu_ids` if `use_gpu` is True
+        if self._device:
+            return self._device
         if self.use_gpu and torch.cuda.is_available():
             if self._is_remote:
+                # Gets the correct torch device to use for training.
+                # Assumes that `CUDA_VISIBLE_DEVICES` is set and is a
+                # superset of the `ray.get_gpu_ids()`.
+                # Example:
+                #     >>> # os.environ["CUDA_VISIBLE_DEVICES"] = "3,4"
+                #     >>> # ray.get_gpu_ids() == [3]
+                #     >>> # torch.cuda.is_available() == True
+                #     >>> # get_device() == torch.device("cuda:0")
+                #     >>> # os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3,4"
+                #     >>> # ray.get_gpu_ids() == [4]
+                #     >>> # torch.cuda.is_available() == True
+                #     >>> # get_device() == torch.device("cuda:4")
+                #     >>> # os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3,4,5"
+                #     >>> # ray.get_gpu_ids() == [4,5]
+                #     >>> # torch.cuda.is_available() == True
+                #     >>> # get_device() == torch.device("cuda:4")
+
                 # Adjust to support multiple GPUs per worker or fractional
                 # GPUs per worker.
-                device_id = ray.get_gpu_ids()[0]
-                return torch.device("cuda", device_id)
+                gpu_id = ray.get_gpu_ids()[0]
+                cuda_visible_str = os.environ.get("CUDA_VISIBLE_DEVICES", "")
+                if cuda_visible_str and cuda_visible_str != "NoDevFiles":
+                    cuda_visible_list = [
+                        int(dev) for dev in cuda_visible_str.split(",")
+                    ]
+                    device_id = cuda_visible_list.index(gpu_id)
+                    return torch.device("cuda", device_id)
+                else:
+                    raise RuntimeError(
+                        "CUDA_VISIBLE_DEVICES set incorrectly. "
+                        f"Got {cuda_visible_str}, expected to include "
+                        f"{gpu_id}. Did you override `CUDA_VISIBLE_DEVICES`? "
+                        "If not, please help file an issue on Github.")
             else:
                 # If the root device is requested on the driver, just return
                 # the 0th device.
@@ -553,6 +264,11 @@ def root_device(self):
         else:
             return torch.device("cpu")
 
+    @root_device.setter
+    def root_device(self, device):
+        """Set the root device."""
+        self._device = device
+
     @property
     def distributed_sampler_kwargs(self):
         """Returns the args to use for torch.data.DistributedSampler."""
@@ -562,4 +278,10 @@ def distributed_sampler_kwargs(self):
 
     @property
     def _is_single_process_single_device(self):
+        """Return True if the process is single process and single device."""
         return True
+
+    def teardown(self) -> None:
+        """Teardown the workers and pytorch DDP connections."""
+        self.accelerator = None
+        super().teardown()
diff --git a/ray_lightning/ray_ddp_sharded.py b/ray_lightning/ray_ddp_sharded.py
index b63ab1a1..b2ab41a4 100644
--- a/ray_lightning/ray_ddp_sharded.py
+++ b/ray_lightning/ray_ddp_sharded.py
@@ -1,34 +1,13 @@
-from pytorch_lightning.plugins import DDPSpawnShardedPlugin
-from pytorch_lightning.plugins.precision.sharded_native_amp import \
-    ShardedNativeMixedPrecisionPlugin
-from pytorch_lightning.utilities import _FAIRSCALE_AVAILABLE
-
-if _FAIRSCALE_AVAILABLE:
-    from fairscale.optim.grad_scaler import ShardedGradScaler
+from pytorch_lightning.strategies import DDPSpawnShardedStrategy
 
 from ray.util import PublicAPI
 
-from ray_lightning import RayPlugin
+from ray_lightning import RayStrategy
 
 
 # C3 linearization of parent classes will do breadth first since both
-# RayPlugin and DDPSpawnShardedPlugin share a common parent of DDPSpawnPlugin
+# RayStrategy and DDPSpawnShardedStrategy share
+# a common parent of DDPSpawnStrategy
 @PublicAPI(stability="beta")
-class RayShardedPlugin(RayPlugin, DDPSpawnShardedPlugin):
-    def execute_remote(self, model, global_rank, queue):
-        # Need to set self._model here otherwise self.lightning_module will
-        # return None.
-        self._model = model
-
-        # This is copied from `DDPSpawnShardedPlugin.new_process`.
-        # As of PTL 1.5, this is the only difference between
-        # `DDPSpawnShardedPlugin` and `DDPSpawnPlugin`.
-        precision_plugin = self.lightning_module.trainer.accelerator\
-            .precision_plugin
-        if isinstance(precision_plugin, ShardedNativeMixedPrecisionPlugin):
-            precision_plugin.scaler = ShardedGradScaler()
-
-        # After setting the grad scaler, we can now call the default
-        # `RayPlugin.execute_remote`.
-        return super().execute_remote(
-            model=self._model, global_rank=global_rank, queue=queue)
+class RayShardedStrategy(RayStrategy, DDPSpawnShardedStrategy):
+    strategy_name = "ddp_sharded_ray"
diff --git a/ray_lightning/ray_horovod.py b/ray_lightning/ray_horovod.py
index dc715608..96be5cde 100644
--- a/ray_lightning/ray_horovod.py
+++ b/ray_lightning/ray_horovod.py
@@ -1,18 +1,11 @@
 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
 
-import ray
-from ray import ObjectRef
 from ray.util import PublicAPI
-from ray.util.queue import Queue
 
-from ray_lightning.session import init_session
-from ray_lightning.util import process_results, Unavailable, to_state_stream, \
-    load_state_stream
-from ray_lightning.tune import TUNE_INSTALLED, is_session_enabled
+from pytorch_lightning.strategies import HorovodStrategy, ParallelStrategy
+import ray
+
+from ray_lightning.util import Unavailable
 
 try:
     import horovod.torch as hvd
@@ -24,6 +17,10 @@
 else:
     HOROVOD_AVAILABLE = True
 
+from ray_lightning.launchers import RayHorovodLauncher
+from ray_lightning.accelerators import \
+    _GPUAccelerator  # noqa: F401
+
 
 def get_executable_cls():
     # Only used for testing purposes, currently.
@@ -32,10 +29,10 @@ def get_executable_cls():
 
 
 @PublicAPI(stability="beta")
-class HorovodRayPlugin(HorovodPlugin):
-    """Pytorch Lightning Plugin for Horovod training on a Ray cluster.
+class HorovodRayStrategy(HorovodStrategy):
+    """Pytorch Lightning Strategy for Horovod training on a Ray cluster.
 
-    This plugin is used to manage distributed training on a Ray cluster
+    This strategy is used to manage distributed training on a Ray cluster
     via the Horovod training framework. Internally, the specified number of
     Ray actors are launched in the cluster and are configured as part of the
     Horovod ring. The Pytorch Lightning trainer is instantiated on the
@@ -45,7 +42,7 @@ class HorovodRayPlugin(HorovodPlugin):
     Each training worker is configured to reserve 1 CPU and if 1 GPU if
     ``use_gpu`` is set to ``True``.
 
-    If using this plugin, you should run your code like a normal Python
+    If using this strategy, you should run your code like a normal Python
     script: ``python train.py``, and not with ``horovodrun``.
 
     Args:
@@ -63,173 +60,88 @@ class HorovodRayPlugin(HorovodPlugin):
             from ray_lightning import HorovodRayPlugin
 
             ptl_model = MNISTClassifier(...)
-            plugin = HorovodRayPlugin(num_workers=2, use_gpu=True)
+            strategy = HorovodRayPlugin(num_workers=2, use_gpu=True)
 
             # Don't set ``gpus`` in ``Trainer``.
             # The actual number of GPUs is determined by ``num_workers``.
-            trainer = pl.Trainer(..., plugins=[plugin])
+            trainer = pl.Trainer(..., strategy=strategy)
             trainer.fit(ptl_model)
 
     """
+    strategy_name = "horovod_ray"
 
     def __init__(self,
                  num_workers: int,
                  num_cpus_per_worker: int = 1,
                  use_gpu: bool = False):
-
+        """Initialize HorovodRayStrategy."""
         if not HOROVOD_AVAILABLE:
-            raise RuntimeError("Please intall Horovod to use this plugin.")
+            raise RuntimeError("Please intall Horovod to use this strategy.")
         if not ray.is_initialized():
             ray.init()
-        super().__init__()
-        self.nickname = "horovod_ray"
+        ParallelStrategy.__init__(
+            self, accelerator="_gpu" if use_gpu else "cpu")
         self.num_workers = num_workers
         self.cpus_per_worker = num_cpus_per_worker
         self.use_gpu = use_gpu
         self.executor = None
+        self._exit_stack = None
+        self._local_rank = 0
 
-    def __getstate__(self):
-        d = self.__dict__.copy()
-        del d["executor"]
-        return d
+        self._is_remote = False
 
-    def __setstate__(self, d):
-        d["executor"] = None
-        self.__dict__.update(d)
+    def _configure_launcher(self):
+        """Configure the Ray launcher.
+            The horovod launcher is used to launch the Ray actors.
+        """
+        settings = RayExecutor.create_settings(timeout_s=30)
+        self.executor = RayExecutor(
+            settings,
+            num_workers=self.num_workers,
+            cpus_per_worker=self.cpus_per_worker,
+            use_gpu=self.use_gpu)
+
+        self._launcher = RayHorovodLauncher(self)
 
     @property
     def global_rank(self) -> int:
+        """Return the global rank of the current process."""
         if not hvd.is_initialized():
             return 0
         return hvd.rank()
 
     @property
     def local_rank(self) -> int:
+        """Return the local rank of the current process."""
         if not hvd.is_initialized():
             return 0
         return hvd.local_rank()
 
     @property
     def world_size(self) -> int:
+        """Return the world size of the current process."""
         if not hvd.is_initialized():
             return self.num_workers
         return hvd.size()
 
-    def setup(self):
-        """Creates the RayExecutor object."""
-        settings = RayExecutor.create_settings(timeout_s=30)
-        self.executor = RayExecutor(
-            settings,
-            num_workers=self.num_workers,
-            cpus_per_worker=self.cpus_per_worker,
-            use_gpu=self.use_gpu)
-        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
-
-    def pre_dispatch(self):
-        """All pre-dispatch logic should be done in train_remote instead."""
-        pass
-
-    def start_training(self, trainer):
-        """Main training loop.
-
-        Trigger remote training via ``train_remote`` on each
-        worker. If using with Ray Tune, create a communication queue to
-        retrieve intermediate results, and process those results. Finally
-        retrieve the training results from the rank 0 worker and return."""
-        model = self._model
-        model_ref = ray.put(model)
-        # Don't pickle the model when training remotely.
-        self._model = None
-
-        queue = None
-        if TUNE_INSTALLED and is_session_enabled():
-            # Create communication queue and send to all the workers.
-            queue = Queue(actor_options={"num_cpus": 0})
-
-        result_futures = self.executor.run_remote(
-            self.train_remote, args=[model_ref, queue])
-
-        results = process_results(result_futures, queue)
-
-        results, state_stream, best_path = results[0]
-        state_dict = load_state_stream(state_stream, to_gpu=self.use_gpu)
-        self._results = results
-        self._model = model
-        self._model.load_state_dict(state_dict)
-        self._model.trainer.accelerator.training_type_plugin = self
-        if self.lightning_module.trainer.checkpoint_callback:
-            self.lightning_module.trainer.checkpoint_callback \
-                .best_model_path = best_path
-
-        if queue:
-            # Shutdown the queue.
-            queue.shutdown()
-
-        return results
-
-    def train_remote(self, model: ObjectRef, queue: Queue = None, **kwargs):
-        """Training function to be executed on each remote worker."""
-        self._model = ray.get(model)
-        self.lightning_module.trainer._accelerator_connector\
-            ._training_type_plugin = self
-        self.lightning_module.trainer._accelerator_connector.accelerator\
-            .training_type_plugin = self
-
-        hvd.init()
-        rank_zero_only.rank = self.global_rank
-
-        if queue is not None:
-            # Initialize session.
-            init_session(rank=self.global_rank, queue=queue)
-
-        # Move the model to the appropriate device.
-        super(HorovodRayPlugin, self).model_to_device()
-
-        # TODO: Make changes in PTL to clean this up.
-        super(HorovodRayPlugin, self).pre_dispatch()
-        results = super(HorovodRayPlugin,
-                        self).start_training(self.lightning_module.trainer)
-        if self.global_rank != 0:
-            # Only want results from the first worker.
-            return None
-
-        best_model_path = None
-        if self.lightning_module.trainer.checkpoint_callback is not None:
-            best_model_path = \
-                self.lightning_module.trainer.checkpoint_callback\
-                    .best_model_path
-
-        return results, to_state_stream(self.lightning_module.state_dict()), \
-            best_model_path
-
-    def post_dispatch(self, trainer: "pl.Trainer"):
-        """Shuts down the RayExecutor."""
-        self.executor.shutdown()
+    def teardown(self) -> None:
+        """Teardown the strategy."""
+        self.join()
+        self.accelerator = None
+        super().teardown()
 
     @property
     def is_distributed(self):
+        """Return whether the strategy is distributed."""
         return True
 
+    def set_remote(self, remote: bool):
+        """Set the remote flag."""
+        self._is_remote = remote
+
     @property
     def root_device(self):
+        """Return the root device."""
         if self.use_gpu and torch.cuda.is_available():
             if hvd.is_initialized():
                 return torch.device("cuda", hvd.local_rank())
diff --git a/ray_lightning/tune.py b/ray_lightning/tune.py
index 8fc7642b..77021f9d 100644
--- a/ray_lightning/tune.py
+++ b/ray_lightning/tune.py
@@ -79,19 +79,19 @@ class TuneReportCallback(TuneCallback):
             .. code-block:: python
 
                 import pytorch_lightning as pl
-                from ray_lightning import RayPlugin
+                from ray_lightning import RayStrategy
                 from ray_lightning.tune import TuneReportCallback
 
-                # Create plugin.
-                ray_plugin = RayPlugin(num_workers=4, use_gpu=True)
+                # Create strategy.
+                ray_plugin = RayStrategy(num_workers=4, use_gpu=True)
 
                 # Report loss and accuracy to Tune after each validation epoch:
-                trainer = pl.Trainer(plugins=[ray_plugin], callbacks=[
+                trainer = pl.Trainer(strategy=[ray_plugin], callbacks=[
                     TuneReportCallback(["val_loss", "val_acc"],
                         on="validation_end")])
 
                 # Same as above, but report as `loss` and `mean_accuracy`:
-                trainer = pl.Trainer(plugins=[ray_plugin], callbacks=[
+                trainer = pl.Trainer(strategy=[ray_plugin], callbacks=[
                     TuneReportCallback(
                         {"loss": "val_loss", "mean_accuracy": "val_acc"},
                         on="validation_end")])
@@ -137,7 +137,7 @@ class _TuneCheckpointCallback(TuneCallback):
         """Distributed PyTorch Lightning to Ray Tune checkpoint callback
 
             Saves checkpoints after each validation step. To be used
-            specifically with the plugins in this library.
+            specifically with the strategies in this library.
 
             Checkpoint are currently not registered if no ``tune.report()``
             call is made afterwards. Consider using
@@ -169,7 +169,7 @@ def _create_checkpoint(checkpoint_stream, global_step: int,
         def _handle(self, trainer: Trainer, pl_module: LightningModule):
             if trainer.sanity_checking:
                 return
-            checkpoint_dict = trainer.checkpoint_connector.dump_checkpoint()
+            checkpoint_dict = trainer._checkpoint_connector.dump_checkpoint()
             # Convert to a state stream first.
             checkpoint_stream = to_state_stream(checkpoint_dict)
             global_step = trainer.global_step
@@ -183,7 +183,7 @@ class TuneReportCheckpointCallback(TuneCallback):
 
             Saves checkpoints after each validation step. Also reports metrics
             to Tune, which is needed for checkpoint registration. To be used
-            specifically with the plugins in this library.
+            specifically with the strategies in this library.
 
             Args:
                 metrics (str|list|dict): Metrics to report to Tune.
@@ -205,15 +205,15 @@ class TuneReportCheckpointCallback(TuneCallback):
             .. code-block:: python
 
                 import pytorch_lightning as pl
-                from ray_lightning import RayPlugin
+                from ray_lightning import RayStrategy
                 from ray_lightning.tune import TuneReportCheckpointCallback.
 
-                # Create the Ray plugin.
-                ray_plugin = RayPlugin()
+                # Create the Ray strategy.
+                ray_plugin = RayStrategy()
 
                 # Save checkpoint after each training batch and after each
                 # validation epoch.
-                trainer = pl.Trainer(plugins=[ray_plugin], callbacks=[
+                trainer = pl.Trainer(strategy=[ray_plugin], callbacks=[
                     TuneReportCheckpointCallback(
                         metrics={"loss": "val_loss",
                                 "mean_accuracy": "val_acc"},
diff --git a/ray_lightning/util.py b/ray_lightning/util.py
index 52718b2b..1515d3b1 100644
--- a/ray_lightning/util.py
+++ b/ray_lightning/util.py
@@ -4,6 +4,8 @@
 import torch
 from pytorch_lightning.accelerators import GPUAccelerator
 from pytorch_lightning import Trainer
+from pytorch_lightning.strategies import Strategy
+from pytorch_lightning.utilities.rank_zero import rank_zero_info
 
 import ray
 
@@ -88,3 +90,13 @@ def load_state_stream(state_stream, to_gpu):
         map_location=("cpu"
                       if not to_gpu else lambda storage, loc: storage.cuda()))
     return state_dict
+
+
+def set_cuda_device_if_used(strategy: "Strategy") -> None:
+    """Set the CUDA device to use for the root node."""
+    if strategy.use_gpu:
+        # overwrite the logger
+        rank_zero_info("GPU available: True (cuda), used: True "
+                       "(Please ignore the previous info [GPU used: False]).")
+
+        torch.cuda.set_device(strategy.root_device)
diff --git a/setup.py b/setup.py
index e795d7f7..433ab509 100644
--- a/setup.py
+++ b/setup.py
@@ -5,8 +5,8 @@
     packages=find_packages(where=".", include="ray_lightning*"),
     version="0.3.0",
     author="Ray Team",
-    description="Ray distributed plugins for Pytorch Lightning.",
-    long_description="Custom Pytorch Lightning distributed plugins "
+    description="Ray distributed strategies for Pytorch Lightning.",
+    long_description="Custom Pytorch Lightning distributed strategies "
     "built on top of distributed computing framework Ray.",
     url="https://github.com/ray-project/ray_lightning_accelerators",
-    install_requires=["pytorch-lightning==1.5.*", "ray"])
+    install_requires=["pytorch-lightning==1.6.*", "ray"])


Test patch:
diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml
index 43e2629f..547647f1 100644
--- a/.github/workflows/test.yaml
+++ b/.github/workflows/test.yaml
@@ -35,7 +35,7 @@ jobs:
           python -m pip install --upgrade pip
           python -m pip install --upgrade setuptools
           python -m pip install codecov
-          python -m pip install -U https://s3-us-west-2.amazonaws.com/ray-wheels/latest/ray-2.0.0.dev0-cp37-cp37m-manylinux2014_x86_64.whl
+          python -m pip install -U https://s3-us-west-2.amazonaws.com/ray-wheels/latest/ray-3.0.0.dev0-cp37-cp37m-manylinux2014_x86_64.whl
           if [ -f requirements-test.txt ]; then python -m pip install -r requirements-test.txt; fi
       - name: Install package
         run: |
@@ -60,7 +60,7 @@ jobs:
           python -m pip install --upgrade pip
           python -m pip install --upgrade setuptools
           python -m pip install codecov
-          python -m pip install -U https://s3-us-west-2.amazonaws.com/ray-wheels/latest/ray-2.0.0.dev0-cp37-cp37m-manylinux2014_x86_64.whl
+          python -m pip install -U https://s3-us-west-2.amazonaws.com/ray-wheels/latest/ray-3.0.0.dev0-cp37-cp37m-manylinux2014_x86_64.whl
           if [ -f requirements-test.txt ]; then python -m pip install -r requirements-test.txt; fi
           HOROVOD_WITH_GLOO=1 HOROVOD_WITHOUT_MPI=1 HOROVOD_WITHOUT_MXNET=1 pip install git+https://github.com/horovod/horovod.git
       - name: Install package
@@ -86,7 +86,7 @@ jobs:
           python -m pip install --upgrade pip
           python -m pip install --upgrade setuptools
           python -m pip install codecov
-          python -m pip install -U https://s3-us-west-2.amazonaws.com/ray-wheels/latest/ray-2.0.0.dev0-cp37-cp37m-manylinux2014_x86_64.whl
+          python -m pip install -U https://s3-us-west-2.amazonaws.com/ray-wheels/latest/ray-3.0.0.dev0-cp37-cp37m-manylinux2014_x86_64.whl
           if [ -f requirements-test.txt ]; then python -m pip install -r requirements-test.txt; fi
           HOROVOD_WITH_GLOO=1 HOROVOD_WITHOUT_MPI=1 HOROVOD_WITHOUT_MXNET=1 pip install git+https://github.com/horovod/horovod.git
       - name: Install package
diff --git a/ray_lightning/tests/test_ddp.py b/ray_lightning/tests/test_ddp.py
index 00ff732b..9c2b4c02 100644
--- a/ray_lightning/tests/test_ddp.py
+++ b/ray_lightning/tests/test_ddp.py
@@ -11,7 +11,7 @@
 import ray
 from ray.cluster_utils import Cluster
 
-from ray_lightning import RayPlugin
+from ray_lightning import RayStrategy
 from ray_lightning.tests.utils import get_trainer, train_test, \
     load_test, predict_test, BoringModel, LightningMNISTClassifier, \
     XORModel, XORDataModule
@@ -72,8 +72,8 @@ def check_num_actor():
 
     model.on_epoch_end = check_num_actor
 
-    plugin = RayPlugin(num_workers=num_workers)
-    trainer = get_trainer(tmpdir, plugins=[plugin])
+    strategy = RayStrategy(num_workers=num_workers)
+    trainer = get_trainer(tmpdir, strategy=strategy)
     trainer.fit(model)
 
 
@@ -90,18 +90,18 @@ class Node2Actor:
         def get_node_ip(self):
             return "2"
 
-    plugin = RayPlugin(num_workers=4, use_gpu=False)
+    strategy = RayStrategy(num_workers=4, use_gpu=False)
+    strategy._configure_launcher()
 
     # 2 workers on "Node 1", 2 workers on "Node 2"
-    plugin.workers = [
+    strategy._launcher._workers = [
         Node1Actor.remote(),
         Node1Actor.remote(),
         Node2Actor.remote(),
         Node2Actor.remote()
     ]
 
-    global_to_local = plugin.get_local_ranks()
-
+    global_to_local = strategy._launcher.get_local_ranks()
     assert len(global_to_local) == 4
     local_ranks = {ranks[0] for ranks in global_to_local}
     node_ranks = {ranks[1] for ranks in global_to_local}
@@ -122,7 +122,7 @@ def test_actor_creation_resources(tmpdir, ray_start_4_cpus_4_extra,
                                   num_cpus_per_worker):
     """Tests if training actors are created with custom resources."""
     model = BoringModel()
-    plugin = RayPlugin(
+    strategy = RayStrategy(
         num_workers=num_workers,
         num_cpus_per_worker=num_cpus_per_worker,
         resources_per_worker={"extra": 1})
@@ -131,49 +131,49 @@ def check_num_actor():
         assert len(ray.state.actors()) == num_workers
 
     model.on_epoch_end = check_num_actor
-    trainer = get_trainer(tmpdir, plugins=[plugin])
+    trainer = get_trainer(tmpdir, strategy=strategy)
     trainer.fit(model)
 
 
 def test_resource_override(ray_start_2_cpus):
     """Tests if CPU and GPU resources are overridden if manually passed in."""
 
-    plugin = RayPlugin(num_workers=1, num_cpus_per_worker=2, use_gpu=True)
-    assert plugin.num_cpus_per_worker == 2
-    assert plugin.use_gpu
+    strategy = RayStrategy(num_workers=1, num_cpus_per_worker=2, use_gpu=True)
+    assert strategy.num_cpus_per_worker == 2
+    assert strategy.use_gpu
 
-    plugin = RayPlugin(
+    strategy = RayStrategy(
         num_workers=1,
         num_cpus_per_worker=2,
         use_gpu=True,
         resources_per_worker={"CPU": 3})
-    assert plugin.num_cpus_per_worker == 3
-    assert plugin.use_gpu
+    assert strategy.num_cpus_per_worker == 3
+    assert strategy.use_gpu
 
-    plugin = RayPlugin(
+    strategy = RayStrategy(
         num_workers=1,
         num_cpus_per_worker=2,
         use_gpu=True,
         resources_per_worker={"GPU": 0})
-    assert plugin.num_cpus_per_worker == 2
-    assert not plugin.use_gpu
+    assert strategy.num_cpus_per_worker == 2
+    assert not strategy.use_gpu
 
-    plugin = RayPlugin(
+    strategy = RayStrategy(
         num_workers=1,
         num_cpus_per_worker=2,
         use_gpu=False,
         resources_per_worker={"GPU": 1})
-    assert plugin.num_cpus_per_worker == 2
-    assert plugin.use_gpu
+    assert strategy.num_cpus_per_worker == 2
+    assert strategy.use_gpu
 
-    plugin = RayPlugin(
+    strategy = RayStrategy(
         num_workers=1,
         num_cpus_per_worker=2,
         use_gpu=False,
         resources_per_worker={"GPU": 2})
-    assert plugin.num_cpus_per_worker == 2
-    assert plugin.num_gpus_per_worker == 2
-    assert plugin.use_gpu
+    assert strategy.num_cpus_per_worker == 2
+    assert strategy.num_gpus_per_worker == 2
+    assert strategy.use_gpu
 
 
 def test_distributed_sampler(tmpdir, ray_start_2_cpus):
@@ -205,9 +205,9 @@ def on_test_start(self, trainer, pl_module):
             assert train_sampler.num_replicas == 2
             assert train_sampler.rank == trainer.global_rank
 
-    plugin = RayPlugin(num_workers=2)
+    strategy = RayStrategy(num_workers=2)
     trainer = get_trainer(
-        tmpdir, plugins=[plugin], callbacks=[DistributedSamplerCallback()])
+        tmpdir, strategy=strategy, callbacks=[DistributedSamplerCallback()])
     trainer.fit(model)
 
 
@@ -215,8 +215,8 @@ def on_test_start(self, trainer, pl_module):
 def test_train(tmpdir, ray_start_2_cpus, num_workers):
     """Tests if training modifies model weights."""
     model = BoringModel()
-    plugin = RayPlugin(num_workers=num_workers)
-    trainer = get_trainer(tmpdir, plugins=[plugin])
+    strategy = RayStrategy(num_workers=num_workers)
+    trainer = get_trainer(tmpdir, strategy=strategy)
     train_test(trainer, model)
 
 
@@ -224,17 +224,17 @@ def test_train(tmpdir, ray_start_2_cpus, num_workers):
 def test_train_client(tmpdir, start_ray_client_server_2_cpus, num_workers):
     assert ray.util.client.ray.is_connected()
     model = BoringModel()
-    plugin = RayPlugin(num_workers=num_workers)
-    trainer = get_trainer(tmpdir, plugins=[plugin])
+    strategy = RayStrategy(num_workers=num_workers)
+    trainer = get_trainer(tmpdir, strategy=strategy)
     train_test(trainer, model)
 
 
 def test_test_with_dataloader_workers(tmpdir, ray_start_2_cpus, seed):
     """Tests trainer.test with >0 workers for data loading."""
     model = BoringModel()
-    plugin = RayPlugin(num_workers=1, use_gpu=False)
+    strategy = RayStrategy(num_workers=1, use_gpu=False)
     trainer = get_trainer(
-        tmpdir, limit_train_batches=20, max_epochs=1, plugins=[plugin])
+        tmpdir, limit_train_batches=20, max_epochs=1, strategy=strategy)
     trainer.test(model)
 
 
@@ -242,8 +242,8 @@ def test_test_with_dataloader_workers(tmpdir, ray_start_2_cpus, seed):
 def test_load(tmpdir, ray_start_2_cpus, num_workers):
     """Tests if model checkpoint can be loaded."""
     model = BoringModel()
-    plugin = RayPlugin(num_workers=num_workers, use_gpu=False)
-    trainer = get_trainer(tmpdir, plugins=[plugin])
+    strategy = RayStrategy(num_workers=num_workers, use_gpu=False)
+    trainer = get_trainer(tmpdir, strategy=strategy)
     load_test(trainer, model)
 
 
@@ -260,9 +260,9 @@ def test_predict(tmpdir, ray_start_2_cpus, seed, num_workers):
     model = LightningMNISTClassifier(config, tmpdir)
     dm = MNISTDataModule(
         data_dir=tmpdir, num_workers=1, batch_size=config["batch_size"])
-    plugin = RayPlugin(num_workers=num_workers, use_gpu=False)
+    strategy = RayStrategy(num_workers=num_workers, use_gpu=False)
     trainer = get_trainer(
-        tmpdir, limit_train_batches=20, max_epochs=1, plugins=[plugin])
+        tmpdir, limit_train_batches=20, max_epochs=1, strategy=strategy)
     predict_test(trainer, model, dm)
 
 
@@ -280,23 +280,23 @@ def test_predict_client(tmpdir, start_ray_client_server_2_cpus, seed,
     model = LightningMNISTClassifier(config, tmpdir)
     dm = MNISTDataModule(
         data_dir=tmpdir, num_workers=1, batch_size=config["batch_size"])
-    plugin = RayPlugin(num_workers=num_workers, use_gpu=False)
+    strategy = RayStrategy(num_workers=num_workers, use_gpu=False)
     trainer = get_trainer(
-        tmpdir, limit_train_batches=20, max_epochs=1, plugins=[plugin])
+        tmpdir, limit_train_batches=20, max_epochs=1, strategy=strategy)
     predict_test(trainer, model, dm)
 
 
 def test_early_stop(tmpdir, ray_start_2_cpus):
     """Tests if early stopping callback works correctly."""
     model = BoringModel()
-    plugin = RayPlugin(num_workers=1, use_gpu=False)
+    strategy = RayStrategy(num_workers=1, use_gpu=False)
     patience = 2
     early_stop = EarlyStopping(
         monitor="val_loss", patience=patience, verbose=True)
     trainer = get_trainer(
         tmpdir,
         max_epochs=500,
-        plugins=[plugin],
+        strategy=strategy,
         callbacks=[early_stop],
         num_sanity_val_steps=0,
         limit_train_batches=1.0,
@@ -311,7 +311,7 @@ def test_early_stop(tmpdir, ray_start_2_cpus):
 def test_unused_parameters(tmpdir, ray_start_2_cpus):
     """Tests if find_unused_parameters is properly passed to model."""
     model = BoringModel()
-    plugin = RayPlugin(
+    strategy = RayStrategy(
         num_workers=2, use_gpu=False, find_unused_parameters=False)
 
     class UnusedParameterCallback(Callback):
@@ -319,19 +319,21 @@ def on_train_start(self, trainer, pl_module):
             assert trainer.model.find_unused_parameters is False
 
     trainer = get_trainer(
-        tmpdir, plugins=[plugin], callbacks=[UnusedParameterCallback()])
+        tmpdir, strategy=strategy, callbacks=[UnusedParameterCallback()])
     trainer.fit(model)
 
 
 def test_metrics(tmpdir, ray_start_2_cpus):
     """Tests if metrics are returned correctly"""
     model = XORModel()
-    plugin = RayPlugin(num_workers=2, find_unused_parameters=False)
+    strategy = RayStrategy(num_workers=2, find_unused_parameters=False)
     trainer = get_trainer(
         tmpdir,
-        plugins=[plugin],
+        strategy=strategy,
         max_epochs=1,
         num_sanity_val_steps=0,
+        limit_train_batches=2,
+        limit_val_batches=2,
         reload_dataloaders_every_n_epochs=1)
     dataset = XORDataModule()
     trainer.fit(model, dataset)
diff --git a/ray_lightning/tests/test_ddp_gpu.py b/ray_lightning/tests/test_ddp_gpu.py
index 4591822f..56b42b20 100644
--- a/ray_lightning/tests/test_ddp_gpu.py
+++ b/ray_lightning/tests/test_ddp_gpu.py
@@ -8,7 +8,7 @@
 
 import ray
 
-from ray_lightning import RayPlugin
+from ray_lightning import RayStrategy
 from ray_lightning.tests.utils import get_trainer, train_test, BoringModel, \
     predict_test, LightningMNISTClassifier
 
@@ -38,8 +38,8 @@ def seed():
 def test_train(tmpdir, ray_start_2_gpus, num_workers):
     """Tests if training modifies model weights."""
     model = BoringModel()
-    plugin = RayPlugin(num_workers=num_workers, use_gpu=True)
-    trainer = get_trainer(tmpdir, plugins=[plugin])
+    strategy = RayStrategy(num_workers=num_workers, use_gpu=True)
+    trainer = get_trainer(tmpdir, strategy=strategy)
     train_test(trainer, model)
 
 
@@ -57,9 +57,9 @@ def test_predict(tmpdir, ray_start_2_gpus, seed, num_workers):
     model = LightningMNISTClassifier(config, tmpdir)
     dm = MNISTDataModule(
         data_dir=tmpdir, num_workers=1, batch_size=config["batch_size"])
-    plugin = RayPlugin(num_workers=num_workers, use_gpu=True)
+    strategy = RayStrategy(num_workers=num_workers, use_gpu=True)
     trainer = get_trainer(
-        tmpdir, limit_train_batches=20, max_epochs=1, plugins=[plugin])
+        tmpdir, limit_train_batches=20, max_epochs=1, strategy=strategy)
     predict_test(trainer, model, dm)
 
 
@@ -73,9 +73,9 @@ class CheckGPUCallback(Callback):
         def on_epoch_end(self, trainer, pl_module):
             assert next(pl_module.parameters()).is_cuda
 
-    plugin = RayPlugin(num_workers=2, use_gpu=True)
+    strategy = RayStrategy(num_workers=2, use_gpu=True)
     trainer = get_trainer(
-        tmpdir, plugins=[plugin], callbacks=[CheckGPUCallback()])
+        tmpdir, strategy=strategy, callbacks=[CheckGPUCallback()])
     trainer.fit(model)
 
 
@@ -108,17 +108,18 @@ def get_gpu_placement(current_worker_index, num_gpus_per_worker):
 
     class CheckDevicesCallback(Callback):
         def on_epoch_end(self, trainer, pl_module):
-            assert trainer.root_gpu == get_gpu_placement(
+            assert trainer.strategy.root_device.index == get_gpu_placement(
                 trainer.local_rank, num_gpus_per_worker)
-            assert trainer.root_gpu == pl_module.device.index
-            assert torch.cuda.current_device() == trainer.root_gpu
+            assert trainer.strategy.root_device.index == pl_module.device.index
+            assert torch.cuda.current_device(
+            ) == trainer.strategy.root_device.index
 
-    plugin = RayPlugin(
+    strategy = RayStrategy(
         num_workers=2,
         use_gpu=True,
         resources_per_worker={"GPU": num_gpus_per_worker})
     trainer = get_trainer(
-        tmpdir, plugins=[plugin], callbacks=[CheckDevicesCallback()])
+        tmpdir, strategy=strategy, callbacks=[CheckDevicesCallback()])
     trainer.fit(model)
 
 
@@ -131,6 +132,6 @@ def test_multi_node(tmpdir):
     ray.init("auto")
     num_gpus = ray.available_resources()["GPU"]
     model = BoringModel()
-    plugin = RayPlugin(num_workers=num_gpus, use_gpu=True)
-    trainer = get_trainer(tmpdir, plugins=[plugin])
+    strategy = RayStrategy(num_workers=num_gpus, use_gpu=True)
+    trainer = get_trainer(tmpdir, strategy=strategy)
     train_test(trainer, model)
diff --git a/ray_lightning/tests/test_ddp_sharded.py b/ray_lightning/tests/test_ddp_sharded.py
index 2b5a2e9f..9ed27c91 100644
--- a/ray_lightning/tests/test_ddp_sharded.py
+++ b/ray_lightning/tests/test_ddp_sharded.py
@@ -8,7 +8,7 @@
 
 import ray
 
-from ray_lightning import RayShardedPlugin
+from ray_lightning import RayShardedStrategy
 from ray_lightning.tests.utils import BoringModel
 
 
@@ -25,22 +25,21 @@ def seed():
 
 
 def test_ddp_choice_sharded(tmpdir, ray_start_2_cpus, seed):
-    """Tests if sharded plugin is properly recognized."""
+    """Tests if sharded strategy is properly recognized."""
 
     class CB(Callback):
         def on_fit_start(self, trainer, pl_module):
-            assert isinstance(trainer.accelerator.training_type_plugin,
-                              RayShardedPlugin)
-            raise SystemExit()
+            assert isinstance(trainer.strategy, RayShardedStrategy)
+            raise ValueError()
 
     model = BoringModel()
     trainer = Trainer(
         fast_dev_run=True,
-        plugins=[RayShardedPlugin(num_workers=2)],
+        strategy=RayShardedStrategy(num_workers=2),
         callbacks=[CB()],
     )
 
-    with pytest.raises(SystemExit):
+    with pytest.raises(ValueError):
         trainer.fit(model)
 
 
@@ -48,7 +47,7 @@ def test_ddp_sharded_plugin_checkpoint(tmpdir, ray_start_2_cpus, seed):
     """Tests if checkpoint is saved correctly."""
     model = BoringModel()
     trainer = Trainer(
-        plugins=[RayShardedPlugin(num_workers=2)],
+        strategy=RayShardedStrategy(num_workers=2),
         fast_dev_run=True,
     )
 
@@ -68,7 +67,7 @@ def test_ddp_sharded_plugin_finetune(tmpdir, ray_start_2_cpus, seed):
     """Tests if we can save and restart training."""
     model = BoringModel()
     trainer = Trainer(
-        plugins=[RayShardedPlugin(num_workers=2)],
+        strategy=RayShardedStrategy(num_workers=2),
         fast_dev_run=True,
     )
     trainer.fit(model)
@@ -86,7 +85,7 @@ def test_ddp_sharded_plugin_resume_from_checkpoint(tmpdir, ray_start_2_cpus,
     """Tests if resuming from checkpoint works."""
     model = BoringModel()
     trainer = Trainer(
-        plugins=[RayShardedPlugin(num_workers=2)],
+        strategy=RayShardedStrategy(num_workers=2),
         fast_dev_run=True,
     )
 
@@ -98,7 +97,7 @@ def test_ddp_sharded_plugin_resume_from_checkpoint(tmpdir, ray_start_2_cpus,
     model = BoringModel()
 
     trainer = Trainer(
-        plugins=[RayShardedPlugin(num_workers=2)],
+        strategy=RayShardedStrategy(num_workers=2),
         fast_dev_run=True,
         resume_from_checkpoint=checkpoint_path)
 
@@ -109,7 +108,7 @@ def test_ddp_sharded_plugin_test(tmpdir, ray_start_2_cpus, seed):
     """Tests if test works without fit."""
     model = BoringModel()
     trainer = Trainer(
-        plugins=[RayShardedPlugin(num_workers=2)],
+        strategy=RayShardedStrategy(num_workers=2),
         fast_dev_run=True,
     )
 
@@ -121,7 +120,7 @@ def test_ddp_sharded_plugin_resume_from_checkpoint_downsize(
     """Tests if we can save and resume training with less workers."""
     model = BoringModel()
     trainer = Trainer(
-        plugins=[RayShardedPlugin(num_workers=2)], fast_dev_run=True)
+        strategy=RayShardedStrategy(num_workers=2), fast_dev_run=True)
 
     trainer.fit(model)
 
@@ -131,7 +130,7 @@ def test_ddp_sharded_plugin_resume_from_checkpoint_downsize(
     model = BoringModel()
 
     trainer = Trainer(
-        plugins=[RayShardedPlugin(num_workers=1)],
+        strategy=RayShardedStrategy(num_workers=1),
         fast_dev_run=True,
         resume_from_checkpoint=checkpoint_path)
 
diff --git a/ray_lightning/tests/test_horovod.py b/ray_lightning/tests/test_horovod.py
index 775ae61a..1c8e1a13 100644
--- a/ray_lightning/tests/test_horovod.py
+++ b/ray_lightning/tests/test_horovod.py
@@ -14,7 +14,7 @@
 
 import ray
 
-from ray_lightning import HorovodRayPlugin
+from ray_lightning import HorovodRayStrategy
 from ray_lightning.tests.utils import get_trainer, BoringModel, \
     train_test, load_test, predict_test, LightningMNISTClassifier
 
@@ -49,8 +49,8 @@ def seed():
 def test_train(tmpdir, ray_start_2_cpus, seed, num_workers):
     """Tests if training modifies model weights."""
     model = BoringModel()
-    plugin = HorovodRayPlugin(num_workers=num_workers, use_gpu=False)
-    trainer = get_trainer(tmpdir, plugins=[plugin])
+    strategy = HorovodRayStrategy(num_workers=num_workers, use_gpu=False)
+    trainer = get_trainer(tmpdir, strategy=strategy)
     train_test(trainer, model)
 
 
@@ -60,8 +60,8 @@ def test_train_client(tmpdir, start_ray_client_server_2_cpus, seed,
     """Tests if training modifies model weights."""
     assert ray.util.client.ray.is_connected()
     model = BoringModel()
-    plugin = HorovodRayPlugin(num_workers=num_workers, use_gpu=False)
-    trainer = get_trainer(tmpdir, plugins=[plugin])
+    strategy = HorovodRayStrategy(num_workers=num_workers, use_gpu=False)
+    trainer = get_trainer(tmpdir, strategy=strategy)
     train_test(trainer, model)
 
 
@@ -69,8 +69,8 @@ def test_train_client(tmpdir, start_ray_client_server_2_cpus, seed,
 def test_load(tmpdir, ray_start_2_cpus, seed, num_workers):
     """Tests if model checkpoint can be loaded."""
     model = BoringModel()
-    plugin = HorovodRayPlugin(num_workers=num_workers, use_gpu=False)
-    trainer = get_trainer(tmpdir, plugins=[plugin])
+    strategy = HorovodRayStrategy(num_workers=num_workers, use_gpu=False)
+    trainer = get_trainer(tmpdir, strategy=strategy)
     load_test(trainer, model)
 
 
@@ -86,9 +86,9 @@ def test_predict(tmpdir, ray_start_2_cpus, seed, num_workers):
     model = LightningMNISTClassifier(config, tmpdir)
     dm = MNISTDataModule(
         data_dir=tmpdir, num_workers=1, batch_size=config["batch_size"])
-    plugin = HorovodRayPlugin(num_workers=num_workers, use_gpu=False)
+    strategy = HorovodRayStrategy(num_workers=num_workers, use_gpu=False)
     trainer = get_trainer(
-        tmpdir, limit_train_batches=20, max_epochs=1, plugins=[plugin])
+        tmpdir, limit_train_batches=20, max_epochs=1, strategy=strategy)
     predict_test(trainer, model, dm)
 
 
@@ -105,9 +105,9 @@ def test_predict_client(tmpdir, start_ray_client_server_2_cpus, seed,
     model = LightningMNISTClassifier(config, tmpdir)
     dm = MNISTDataModule(
         data_dir=tmpdir, num_workers=1, batch_size=config["batch_size"])
-    plugin = HorovodRayPlugin(num_workers=num_workers, use_gpu=False)
+    strategy = HorovodRayStrategy(num_workers=num_workers, use_gpu=False)
     trainer = get_trainer(
-        tmpdir, limit_train_batches=20, max_epochs=1, plugins=[plugin])
+        tmpdir, limit_train_batches=20, max_epochs=1, strategy=strategy)
     predict_test(trainer, model, dm)
 
 
@@ -117,8 +117,8 @@ def test_predict_client(tmpdir, start_ray_client_server_2_cpus, seed,
 def test_train_gpu(tmpdir, ray_start_2_gpus, seed, num_workers):
     """Tests if training modifies model weights."""
     model = BoringModel()
-    plugin = HorovodRayPlugin(num_workers=num_workers, use_gpu=True)
-    trainer = get_trainer(tmpdir, plugins=[plugin])
+    strategy = HorovodRayStrategy(num_workers=num_workers, use_gpu=True)
+    trainer = get_trainer(tmpdir, strategy=strategy)
     train_test(trainer, model)
 
 
@@ -128,8 +128,8 @@ def test_train_gpu(tmpdir, ray_start_2_gpus, seed, num_workers):
 def test_load_gpu(tmpdir, ray_start_2_gpus, seed, num_workers):
     """Tests if model checkpoint can be loaded."""
     model = BoringModel()
-    plugin = HorovodRayPlugin(num_workers=num_workers, use_gpu=True)
-    trainer = get_trainer(tmpdir, plugins=[plugin])
+    strategy = HorovodRayStrategy(num_workers=num_workers, use_gpu=True)
+    trainer = get_trainer(tmpdir, strategy=strategy)
     load_test(trainer, model)
 
 
@@ -147,7 +147,7 @@ def test_predict_gpu(tmpdir, ray_start_2_gpus, seed, num_workers):
     model = LightningMNISTClassifier(config, tmpdir)
     dm = MNISTDataModule(
         data_dir=tmpdir, num_workers=1, batch_size=config["batch_size"])
-    plugin = HorovodRayPlugin(num_workers=num_workers, use_gpu=True)
+    strategy = HorovodRayStrategy(num_workers=num_workers, use_gpu=True)
     trainer = get_trainer(
-        tmpdir, limit_train_batches=20, max_epochs=1, plugins=[plugin])
+        tmpdir, limit_train_batches=20, max_epochs=1, strategy=strategy)
     predict_test(trainer, model, dm)
diff --git a/ray_lightning/tests/test_tune.py b/ray_lightning/tests/test_tune.py
index 937ad51b..209d8d35 100644
--- a/ray_lightning/tests/test_tune.py
+++ b/ray_lightning/tests/test_tune.py
@@ -1,11 +1,10 @@
-import os
 import pytest
 
 import ray
 import torch
 from ray import tune
 
-from ray_lightning import RayPlugin, HorovodRayPlugin
+from ray_lightning import RayStrategy, HorovodRayStrategy
 from ray_lightning.tests.utils import BoringModel, get_trainer
 from ray_lightning.tune import TuneReportCallback, \
     TuneReportCheckpointCallback, get_tune_resources
@@ -25,13 +24,13 @@ def ray_start_4_cpus_4_gpus():
     ray.shutdown()
 
 
-def train_func(dir, plugin, callbacks=None):
+def train_func(dir, strategy, callbacks=None):
     def _inner_train(config):
         model = BoringModel()
         trainer = get_trainer(
             dir,
             callbacks=callbacks,
-            plugins=[plugin],
+            strategy=strategy,
             checkpoint_callback=False,
             **config)
         trainer.fit(model)
@@ -39,68 +38,79 @@ def _inner_train(config):
     return _inner_train
 
 
-def tune_test(dir, plugin):
+def tune_test(dir, strategy):
     callbacks = [TuneReportCallback(on="validation_end")]
     analysis = tune.run(
-        train_func(dir, plugin, callbacks=callbacks),
+        train_func(dir, strategy, callbacks=callbacks),
         config={"max_epochs": tune.choice([1, 2, 3])},
         resources_per_trial=get_tune_resources(
-            num_workers=plugin.num_workers, use_gpu=plugin.use_gpu),
+            num_workers=strategy.num_workers, use_gpu=strategy.use_gpu),
         num_samples=2)
-    assert all(analysis.results_df["training_iteration"] ==
-               analysis.results_df["config.max_epochs"])
+    # fix TUNE_RESULT_DELIM issue
+    config_max_epochs = analysis.results_df.get("config.max_epochs", False)
+    if config_max_epochs is False:
+        config_max_epochs = analysis.results_df.get("config/max_epochs", False)
+    assert all(analysis.results_df["training_iteration"] == config_max_epochs)
 
 
 def test_tune_iteration_ddp(tmpdir, ray_start_4_cpus):
-    """Tests if each RayPlugin runs the correct number of iterations."""
-    plugin = RayPlugin(num_workers=2, use_gpu=False)
-    tune_test(tmpdir, plugin)
+    """Tests if each RayStrategy runs the correct number of iterations."""
+    strategy = RayStrategy(num_workers=2, use_gpu=False)
+    tune_test(tmpdir, strategy)
 
 
 def test_tune_iteration_horovod(tmpdir, ray_start_4_cpus):
     """Tests if each HorovodRay trial runs the correct number of iterations."""
-    plugin = HorovodRayPlugin(num_workers=2, use_gpu=False)
-    tune_test(tmpdir, plugin)
+    strategy = HorovodRayStrategy(num_workers=2, use_gpu=False)
+    tune_test(tmpdir, strategy)
 
 
-def checkpoint_test(dir, plugin):
+def checkpoint_test(dir, strategy):
     callbacks = [TuneReportCheckpointCallback(on="validation_end")]
     analysis = tune.run(
-        train_func(dir, plugin, callbacks=callbacks),
+        train_func(dir, strategy, callbacks=callbacks),
         config={"max_epochs": 2},
         resources_per_trial=get_tune_resources(
-            num_workers=plugin.num_workers, use_gpu=plugin.use_gpu),
+            num_workers=strategy.num_workers, use_gpu=strategy.use_gpu),
         num_samples=1,
         local_dir=dir,
         log_to_file=True,
         metric="val_loss",
         mode="min")
-    assert os.path.exists(analysis.best_checkpoint)
+    assert analysis.best_checkpoint is not None
 
 
 def test_checkpoint_ddp(tmpdir, ray_start_4_cpus):
     """Tests if Tune checkpointing works with RayAccelerator."""
-    plugin = RayPlugin(num_workers=2, use_gpu=False)
-    checkpoint_test(tmpdir, plugin)
+    strategy = RayStrategy(num_workers=2, use_gpu=False)
+    checkpoint_test(tmpdir, strategy)
 
 
 def test_checkpoint_horovod(tmpdir, ray_start_4_cpus):
     """Tests if Tune checkpointing works with HorovodRayAccelerator."""
-    plugin = HorovodRayPlugin(num_workers=2, use_gpu=False)
-    checkpoint_test(tmpdir, plugin)
+    strategy = HorovodRayStrategy(num_workers=2, use_gpu=False)
+    checkpoint_test(tmpdir, strategy)
 
 
 @pytest.mark.skipif(
     torch.cuda.device_count() < 4, reason="test requires multi-GPU machine")
 def test_checkpoint_ddp_gpu(tmpdir, ray_start_4_cpus_4_gpus):
     """Tests if Tune checkpointing works with RayAccelerator."""
-    plugin = RayPlugin(num_workers=2, use_gpu=True)
-    checkpoint_test(tmpdir, plugin)
+    strategy = RayStrategy(num_workers=2, use_gpu=True)
+    checkpoint_test(tmpdir, strategy)
 
 
 @pytest.mark.skipif(
     torch.cuda.device_count() < 4, reason="test requires multi-GPU machine")
 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)
+    strategy = HorovodRayStrategy(num_workers=2, use_gpu=True)
+    checkpoint_test(tmpdir, strategy)
+
+
+@pytest.mark.skipif(
+    torch.cuda.device_count() < 4, reason="test requires multi-GPU machine")
+def test_tune_iteration_ddp_gpu(tmpdir, ray_start_4_cpus_4_gpus):
+    """Tests if each RayStrategy runs the correct number of iterations."""
+    strategy = RayStrategy(num_workers=2, use_gpu=True)
+    tune_test(tmpdir, strategy)
diff --git a/ray_lightning/tests/utils.py b/ray_lightning/tests/utils.py
index 23bdb8b0..2a5a64b3 100644
--- a/ray_lightning/tests/utils.py
+++ b/ray_lightning/tests/utils.py
@@ -6,7 +6,7 @@
 from torch.utils.data import Dataset
 
 import pytorch_lightning as pl
-from pytorch_lightning.plugins import PLUGIN_INPUT
+from pytorch_lightning.strategies import Strategy
 from pytorch_lightning import LightningModule, Callback, Trainer, \
     LightningDataModule
 
@@ -119,7 +119,7 @@ def forward(self, x):
         x = self.layer_2(x)
         x = torch.relu(x)
         x = self.layer_3(x)
-        x = F.softmax(x, dim=1)
+        x = F.log_softmax(x, dim=1)
         return x
 
     def configure_optimizers(self):
@@ -128,7 +128,7 @@ def configure_optimizers(self):
     def training_step(self, train_batch, batch_idx):
         x, y = train_batch
         logits = self.forward(x)
-        loss = F.nll_loss(logits, y)
+        loss = F.nll_loss(logits, y.long())
         acc = self.accuracy(logits, y)
         self.log("ptl/train_loss", loss)
         self.log("ptl/train_accuracy", acc)
@@ -137,7 +137,7 @@ def training_step(self, train_batch, batch_idx):
     def validation_step(self, val_batch, batch_idx):
         x, y = val_batch
         logits = self.forward(x)
-        loss = F.nll_loss(logits, y)
+        loss = F.nll_loss(logits, y.long())
         acc = self.accuracy(logits, y)
         return {"val_loss": loss, "val_accuracy": acc}
 
@@ -211,7 +211,7 @@ def val_dataloader(self):
 
 
 def get_trainer(dir,
-                plugins: List[PLUGIN_INPUT],
+                strategy: Strategy,
                 max_epochs: int = 1,
                 limit_train_batches: int = 10,
                 limit_val_batches: int = 10,
@@ -223,7 +223,7 @@ def get_trainer(dir,
     trainer = pl.Trainer(
         default_root_dir=dir,
         callbacks=callbacks,
-        plugins=plugins,
+        strategy=strategy,
         max_epochs=max_epochs,
         limit_train_batches=limit_train_batches,
         limit_val_batches=limit_val_batches,
diff --git a/requirements-test.txt b/requirements-test.txt
index d7b8f185..a26315f1 100644
--- a/requirements-test.txt
+++ b/requirements-test.txt
@@ -4,10 +4,10 @@ flake8-comprehensions
 flake8-quotes
 yapf==0.23.0
 pytest
-pytorch-lightning==1.5.9
+pytorch-lightning==1.6.4
 lightning-bolts==0.3.3
 ray[tune]
-torch==1.8.1
+torch==1.12.0
 torchmetrics
 torchvision
 protobuf<=3.20.1