forward trainer xla single device

Author: justusschock

src/lightning/pytorch/strategies/single_xla.py

@@ -11,23 +11,18 @@
 # 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 os
 from typing import Optional, Union
 
-import torch
 from typing_extensions import override
 
 import lightning.pytorch as pl
-from lightning.fabric.accelerators.xla import _XLA_AVAILABLE
 from lightning.fabric.plugins import CheckpointIO, Precision, XLACheckpointIO
 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 _DEVICE
 from lightning.pytorch.plugins.io.wrapper import _WrappingCheckpointIO
 from lightning.pytorch.plugins.precision.xla import XLAPrecision
 from lightning.pytorch.strategies.single_device import SingleDeviceStrategy
-from lightning.pytorch.trainer.states import TrainerFn
-from lightning.pytorch.utilities import find_shared_parameters, set_shared_parameters
 
 
 class SingleDeviceXLAStrategy(SingleDeviceStrategy):
@@ -41,20 +36,18 @@ def __init__(
         precision_plugin: Optional[XLAPrecision] = None,
         debug: bool = False,
     ):
-        if not _XLA_AVAILABLE:
-            raise ModuleNotFoundError(str(_XLA_AVAILABLE))
-        if isinstance(device, torch.device):
-            # unwrap the `torch.device` in favor of `xla_device`
-            device = device.index
-        import torch_xla.core.xla_model as xm
-
         super().__init__(
             accelerator=accelerator,
-            device=xm.xla_device(device),
+            device=device,
             checkpoint_io=checkpoint_io,
             precision_plugin=precision_plugin,
         )
-        self.debug = debug
+        _raise_enterprise_not_available()
+        from pytorch_lightning_enterprise.strategies.xla.single import (
+            SingleDeviceXLAStrategyTrainer as EnterpriseSingleDeviceXLAStrategy,
+        )
+
+        self.single_xla_strategy_impl = EnterpriseSingleDeviceXLAStrategy(outer_object=self, device=device, debug=debug)
 
     @property
     @override
@@ -90,26 +83,7 @@ def precision_plugin(self, precision_plugin: Optional[Precision]) -> None:
 
     @override
     def setup(self, trainer: "pl.Trainer") -> None:
-        if self.debug:
-            os.environ["PT_XLA_DEBUG"] = str(1)
-
-        assert self.accelerator is not None
-        self.accelerator.setup(trainer)
-
-        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 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.single_xla_strategy_impl.setup(trainer=trainer)
 
     @classmethod
     @override
@@ -118,5 +92,4 @@ def register_strategies(cls, strategy_registry: _StrategyRegistry) -> None:
 
     @override
     def teardown(self) -> None:
-        super().teardown()
-        os.environ.pop("PT_XLA_DEBUG", None)
+        return self.single_xla_strategy_impl.teardown()