XLA ddp trainer

Author: justusschock

src/lightning/pytorch/strategies/xla.py

@@ -11,8 +11,6 @@
 # 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.
-import io
-import os
 from typing import TYPE_CHECKING, Any, Optional, Union
 
 import torch
@@ -21,20 +19,16 @@
 from typing_extensions import override
 
 import lightning.pytorch as pl
-from lightning.fabric.accelerators.xla import _XLA_AVAILABLE, _XLA_GREATER_EQUAL_2_1
 from lightning.fabric.plugins import CheckpointIO, Precision, XLACheckpointIO
 from lightning.fabric.plugins.environments import XLAEnvironment
 from lightning.fabric.strategies import _StrategyRegistry
-from lightning.fabric.utilities.optimizer import _optimizers_to_device
+from lightning.fabric.utilities.imports import _raise_enterprise_not_available
 from lightning.fabric.utilities.types import _PATH, ReduceOp
 from lightning.pytorch.plugins import XLAPrecision
 from lightning.pytorch.plugins.io.wrapper import _WrappingCheckpointIO
 from lightning.pytorch.strategies.ddp import DDPStrategy
 from lightning.pytorch.strategies.launchers.xla import _XLALauncher
 from lightning.pytorch.strategies.strategy import TBroadcast
-from lightning.pytorch.trainer.states import TrainerFn
-from lightning.pytorch.utilities import find_shared_parameters, set_shared_parameters
-from lightning.pytorch.utilities.rank_zero import rank_zero_only
 
 if TYPE_CHECKING:
     from torch_xla.distributed.parallel_loader import MpDeviceLoader
@@ -56,8 +50,6 @@ def __init__(
         sync_module_states: bool = True,
         **_: Any,
     ) -> None:
-        if not _XLA_AVAILABLE:
-            raise ModuleNotFoundError(str(_XLA_AVAILABLE))
         super().__init__(
             accelerator=accelerator,
             parallel_devices=parallel_devices,
@@ -66,9 +58,12 @@ def __init__(
             precision_plugin=precision_plugin,
             start_method="fork",
         )
-        self.debug = debug
-        self._launched = False
-        self._sync_module_states = sync_module_states
+        _raise_enterprise_not_available()
+        from pytorch_lightning_enterprise.strategies.xla.ddp import XLAStrategyTrainer as EnterpriseXLAStrategy
+
+        self.xla_strategy_impl = EnterpriseXLAStrategy(
+            outer_object=self, debug=debug, sync_module_states=sync_module_states
+        )
 
     @property
     @override
@@ -105,145 +100,64 @@ def precision_plugin(self, precision_plugin: Optional[Precision]) -> None:
     @property
     @override
     def root_device(self) -> torch.device:
-        if not self._launched:
-            raise RuntimeError("Accessing the XLA device before processes have spawned is not allowed.")
-        import torch_xla.core.xla_model as xm
-
-        return xm.xla_device()
+        return self.xla_strategy_impl.root_device
 
     @property
     @override
     def global_rank(self) -> int:
-        return super().global_rank if self._launched else 0
+        return self.xla_strategy_impl.global_rank
 
     @property
     @override
     def local_rank(self) -> int:
-        return super().local_rank if self._launched else 0
+        return self.xla_strategy_impl.local_rank
 
     @property
     @override
     def node_rank(self) -> int:
-        return super().node_rank if self._launched else 0
+        return self.xla_strategy_impl.node_rank
 
     @property
     @override
     def world_size(self) -> int:
-        return super().world_size if self._launched else 1
+        return self.xla_strategy_impl.world_size
 
     @override
     def _configure_launcher(self) -> None:
         self._launcher = _XLALauncher(self)
 
     @override
     def setup(self, trainer: "pl.Trainer") -> None:
-        assert self.accelerator is not None
-        self.accelerator.setup(trainer)
-
-        if self.debug:
-            os.environ["PT_XLA_DEBUG"] = "1"
-
-        assert self.model is not None
-        self.precision_plugin.convert_module(self.model)
-
-        shared_params = find_shared_parameters(self.model)
-        self.model_to_device()
-        set_shared_parameters(self.model, shared_params)
-
-        self.model = self._setup_model(self.model)
-
-        if self._sync_module_states:
-            if _XLA_GREATER_EQUAL_2_1:
-                from torch_xla.core.xla_model import broadcast_master_param
-            else:
-                from torch_xla.experimental.pjrt import broadcast_master_param
-
-            broadcast_master_param(self.model)
-
-        if trainer.state.fn == TrainerFn.FITTING:
-            self.setup_optimizers(trainer)
-        self.setup_precision_plugin()
-        if trainer.state.fn == TrainerFn.FITTING:
-            _optimizers_to_device(self.optimizers, self.root_device)
+        return self.xla_strategy_impl.setup(trainer=trainer)
 
     @override
     def _setup_model(self, model: Module) -> Module:  # type: ignore
-        return model
+        return self.xla_strategy_impl._setup_model(model=model)
 
     @property
     @override
     def distributed_sampler_kwargs(self) -> dict[str, int]:
-        return {"num_replicas": self.world_size, "rank": self.global_rank}
+        return self.xla_strategy_impl.distributed_sampler_kwargs
 
     @override
     def process_dataloader(self, dataloader: object) -> "MpDeviceLoader":
-        from torch_xla.distributed.parallel_loader import MpDeviceLoader
-
-        if isinstance(dataloader, MpDeviceLoader):
-            # dataloader is already wrapped by MpDeviceLoader
-            return dataloader
-
-        dataloader = MpDeviceLoader(dataloader, self.root_device)
-        # Mimic interface to torch.utils.data.DataLoader
-        dataloader.dataset = dataloader._loader.dataset
-        dataloader.batch_sampler = getattr(dataloader._loader, "batch_sampler", None)
-        return dataloader
+        return self.xla_strategy_impl.process_dataloader(dataloader=dataloader)
 
     @override
     def configure_ddp(self) -> None:
-        pass
+        return self.xla_strategy_impl.configure_ddp()
 
     @override
     def model_to_device(self) -> None:
-        assert self.model is not None
-        self.model = self.model.to(self.root_device)
+        return self.xla_strategy_impl.model_to_device()
 
     @override
     def barrier(self, name: Optional[str] = None, *args: Any, **kwargs: Any) -> None:
-        if not self._launched:
-            return
-
-        import torch_xla.core.xla_model as xm
-
-        if name is None:
-            # `None` is not supported: "TypeError: _xla_rendezvous(): incompatible function arguments"
-            name = ""
-        xm.rendezvous(name)
+        return self.xla_strategy_impl.barrier(name=name, *args, **kwargs)
 
     @override
     def broadcast(self, obj: TBroadcast, src: int = 0) -> TBroadcast:
-        if not self._launched:
-            return obj
-
-        import torch_xla.core.xla_model as xm
-
-        is_tensor = isinstance(obj, Tensor)
-        if is_tensor:
-            if obj.dim() == 0:
-                obj = obj.unsqueeze(0)
-            original_device = obj.device
-            # XLA distributed requires that the data is on the XLA device
-            obj = obj.to(self.root_device)
-        else:
-            # support for arbitrary pickle-ables
-            buffer = io.BytesIO()
-            torch.save(obj, buffer)
-            obj = torch.tensor(  # type: ignore[assignment]
-                bytearray(buffer.getbuffer()), device=self.root_device, dtype=torch.float
-            )
-
-        obj = [obj]
-        xm.collective_broadcast(obj, root_ordinal=src)
-        obj = obj[0]
-
-        if not is_tensor:
-            # this will preserve the dtype and device of any tensors
-            buffer = io.BytesIO(obj.cpu().byte().numpy())
-            obj = torch.load(buffer)
-        else:
-            obj = obj.to(original_device)
-
-        return obj
+        return self.xla_strategy_impl.broadcast(obj=obj, src=src)
 
     @override
     def reduce(
@@ -252,60 +166,27 @@ def reduce(
         group: Optional[Any] = None,
         reduce_op: Optional[Union[ReduceOp, str]] = "mean",
     ) -> Tensor:
-        if not isinstance(output, Tensor):
-            output = torch.tensor(output, device=self.root_device)
-
-        invalid_reduce_op = isinstance(reduce_op, ReduceOp) and reduce_op != ReduceOp.SUM
-        invalid_reduce_op_str = isinstance(reduce_op, str) and reduce_op.lower() not in ("sum", "mean", "avg")
-        if invalid_reduce_op or invalid_reduce_op_str:
-            raise ValueError(
-                "Currently, the XLAStrategy only supports `sum`, `mean`, `avg` for the reduce operation, got:"
-                f" {reduce_op}"
-            )
-
-        import torch_xla.core.xla_model as xm
-
-        output = xm.mesh_reduce("reduce", output, sum)
-
-        if isinstance(reduce_op, str) and reduce_op.lower() in ("avg", "mean"):
-            output = output / self.world_size
-
-        return output
+        return self.xla_strategy_impl.reduce(output=output, group=group, reduce_op=reduce_op)
 
     @override
     def setup_environment(self) -> None:
-        self._launched = True
-        super().setup_environment()
+        return self.xla_strategy_impl.setup_environment()
 
     @override
     def setup_distributed(self) -> None:
-        assert self.parallel_devices is not None
-        if len(self.parallel_devices) == 1:
-            # spawning only 1 device with PjRT is not supported:
-            # https://github.com/Lightning-AI/pytorch-lightning/pull/17408#discussion_r1170671732
-            raise NotImplementedError(
-                "The `XLAStrategy` does not support running on a single device with the PjRT runtime."
-                " Try using all devices or the `SingleDeviceXLAStrategy` strategy"
-            )
-        rank_zero_only.rank = self.global_rank
+        return self.xla_strategy_impl.setup_distributed()
 
     @override
     def set_world_ranks(self) -> None:
-        # accessing global_rank will initialize the XLA computation client. since this is called outside of the spawned
-        # processes (by the accelerator connector), we cannot run the code that would normally be here.
-        # instead it's done in `setup_distributed`
-        pass
+        return self.xla_strategy_impl.set_world_ranks()
 
     @override
     def save_checkpoint(
         self, checkpoint: dict[str, Any], filepath: _PATH, storage_options: Optional[Any] = None
     ) -> None:
-        import torch_xla.core.xla_model as xm
-
-        # sync any pending lazy tensors on all ranks before saving to prevent potential collective hangs
-        xm.mark_step()
-        # save on global rank zero only
-        super().save_checkpoint(checkpoint, filepath, storage_options=storage_options)
+        return self.xla_strategy_impl.save_checkpoint(
+            checkpoint=checkpoint, filepath=filepath, storage_options=storage_options
+        )
 
     @override
     def remove_checkpoint(self, filepath: _PATH) -> None:
@@ -315,8 +196,7 @@ def remove_checkpoint(self, filepath: _PATH) -> None:
             filepath: Path to checkpoint
 
         """
-        if self.local_rank == 0:
-            self.checkpoint_io.remove_checkpoint(filepath)
+        return self.xla_strategy_impl.remove_checkpoint(filepath=filepath)
 
     @override
     def all_gather(self, tensor: Tensor, group: Optional[Any] = None, sync_grads: bool = False) -> Tensor:
@@ -330,29 +210,11 @@ def all_gather(self, tensor: Tensor, group: Optional[Any] = None, sync_grads: bo
             A tensor of shape (world_size, ...)
 
         """
-        if not self._launched:
-            return tensor
-        if not isinstance(tensor, Tensor):
-            raise NotImplementedError(
-                f"`{type(self).__name__}.all_gather` is only implemented for tensors. Given {tensor}"
-            )
-        if tensor.dim() == 0:
-            tensor = tensor.unsqueeze(0)
-        original_device = tensor.device
-        tensor = tensor.to(self.root_device)
-
-        import torch_xla.core.functions as xf
-        import torch_xla.core.xla_model as xm
-
-        tensor = xf.all_gather(tensor) if sync_grads else xm.all_gather(tensor)
-        tensor = tensor.to(original_device)
-        return tensor
+        return self.xla_strategy_impl.all_gather(tensor=tensor, group=group, sync_grads=sync_grads)
 
     @override
     def teardown(self) -> None:
-        super().teardown()
-        self._launched = False  # after the Trainer finishes, we aren't inside the spawned region
-        os.environ.pop("PT_XLA_DEBUG", None)
+        return self.xla_strategy_impl.teardown()
 
     @classmethod
     @override