[TPU] Add Trainer support for PyTorch XLA FSDP

src/lightning/fabric/connector.py

@@ -454,10 +454,9 @@ def _init_strategy(self) -> None:
             self.strategy = self._strategy_flag
 
     def _check_and_init_precision(self) -> Precision:
-        self._validate_precision_choice()
         if isinstance(self._precision_instance, Precision):
             return self._precision_instance
-        if isinstance(self.accelerator, XLAAccelerator):
+        if isinstance(self.strategy, (SingleDeviceXLAStrategy, XLAStrategy, XLAFSDPStrategy)):
             return XLAPrecision(self._precision_input)  # type: ignore
         if isinstance(self.strategy, DeepSpeedStrategy):
             return DeepSpeedPrecision(self._precision_input)  # type: ignore
@@ -492,18 +491,6 @@ def _check_and_init_precision(self) -> Precision:
 
         raise RuntimeError("No precision set")
 
-    def _validate_precision_choice(self) -> None:
-        """Validate the combination of choices for precision, and accelerator."""
-        if (
-            isinstance(self.accelerator, XLAAccelerator)
-            and self._precision_instance
-            and not isinstance(self._precision_instance, XLAPrecision)
-        ):
-            raise ValueError(
-                f"The `XLAAccelerator` can only be used with a `XLAPrecision` plugin,"
-                f" found: {self._precision_instance}."
-            )
-
     def _lazy_init_strategy(self) -> None:
         """Lazily set missing attributes on the previously instantiated strategy."""
         self.strategy.accelerator = self.accelerator

src/lightning/fabric/strategies/fsdp.py

@@ -305,25 +305,15 @@ def setup_optimizer(self, optimizer: Optimizer) -> Optimizer:
         flattened parameters.
 
         """
-        if _TORCH_GREATER_EQUAL_2_0:
-            return optimizer
-
-        from torch.distributed.fsdp import FlatParameter
-
-        num_groups = len(optimizer.param_groups)
-        if num_groups > 1:
+        if self._fsdp_kwargs.get("use_orig_params"):
+            return super().setup_optimizer(optimizer)
+        if not _optimizer_has_flat_params(optimizer):
+            # We avoid this limitation in PyTorch >= 2.0 by setting `use_orig_params=True`
             raise ValueError(
-                "An optimizer used with an FSDP model does not support multiple param groups."
-                f" Found {num_groups} parameter groups."
+                "The optimizer does not seem to reference any FSDP parameters. HINT: Make sure to create the optimizer"
+                " after setting up the model."
             )
-
-        if any(isinstance(param, FlatParameter) for param in optimizer.param_groups[0]["params"]):
-            return optimizer
-
-        raise ValueError(
-            "The optimizer does not seem to reference any FSDP parameters. HINT: Make sure to create the optimizer"
-            " after setting up the model."
-        )
+        return optimizer
 
     def module_to_device(self, module: Module) -> None:
         pass

src/lightning/fabric/strategies/single_xla.py

@@ -17,9 +17,8 @@
 
 from lightning.fabric.accelerators import Accelerator
 from lightning.fabric.accelerators.xla import _XLA_AVAILABLE
-from lightning.fabric.plugins.io.checkpoint_io import CheckpointIO
+from lightning.fabric.plugins import XLAPrecision
 from lightning.fabric.plugins.io.xla import XLACheckpointIO
-from lightning.fabric.plugins.precision import Precision
 from lightning.fabric.strategies import _StrategyRegistry
 from lightning.fabric.strategies.single_device import SingleDeviceStrategy
 from lightning.fabric.utilities.types import _DEVICE
@@ -32,8 +31,8 @@ def __init__(
         self,
         device: _DEVICE,
         accelerator: Optional[Accelerator] = None,
-        checkpoint_io: Optional[CheckpointIO] = None,
-        precision: Optional[Precision] = None,
+        checkpoint_io: Optional[XLACheckpointIO] = None,
+        precision: Optional[XLAPrecision] = None,
     ):
         if not _XLA_AVAILABLE:
             raise ModuleNotFoundError(str(_XLA_AVAILABLE))
@@ -50,16 +49,34 @@ def __init__(
             precision=precision,
         )
 
-    @property
-    def checkpoint_io(self) -> CheckpointIO:
-        if self._checkpoint_io is None:
-            self._checkpoint_io = XLACheckpointIO()
-        return self._checkpoint_io
+    @property  # type: ignore[override]
+    def checkpoint_io(self) -> XLACheckpointIO:
+        plugin = self._checkpoint_io
+        if plugin is not None:
+            assert isinstance(plugin, XLACheckpointIO)
+            return plugin
+        return XLACheckpointIO()
 
     @checkpoint_io.setter
-    def checkpoint_io(self, io: CheckpointIO) -> None:
+    def checkpoint_io(self, io: Optional[XLACheckpointIO]) -> None:
+        if io is not None and not isinstance(io, XLACheckpointIO):
+            raise TypeError(f"The XLA strategy can only work with the `XLACheckpointIO` plugin, found {io}")
         self._checkpoint_io = io
 
+    @property  # type: ignore[override]
+    def precision(self) -> XLAPrecision:
+        plugin = self._precision
+        if plugin is not None:
+            assert isinstance(plugin, XLAPrecision)
+            return plugin
+        return XLAPrecision("32-true")
+
+    @precision.setter
+    def precision(self, precision: Optional[XLAPrecision]) -> None:
+        if precision is not None and not isinstance(precision, XLAPrecision):
+            raise TypeError(f"The XLA strategy can only work with the `XLAPrecision` plugin, found {precision}")
+        self._precision = precision
+
     @classmethod
     def register_strategies(cls, strategy_registry: _StrategyRegistry) -> None:
         strategy_registry.register("single_xla", cls, description=cls.__name__)

src/lightning/fabric/strategies/xla.py

@@ -22,10 +22,9 @@
 
 from lightning.fabric.accelerators import Accelerator
 from lightning.fabric.accelerators.xla import _XLA_GREATER_EQUAL_2_1, _using_pjrt
+from lightning.fabric.plugins import XLAPrecision
 from lightning.fabric.plugins.environments import XLAEnvironment
-from lightning.fabric.plugins.io.checkpoint_io import CheckpointIO
 from lightning.fabric.plugins.io.xla import XLACheckpointIO
-from lightning.fabric.plugins.precision import Precision
 from lightning.fabric.strategies import ParallelStrategy, _StrategyRegistry
 from lightning.fabric.strategies.launchers.xla import _XLALauncher
 from lightning.fabric.strategies.strategy import TBroadcast
@@ -44,8 +43,8 @@ def __init__(
         self,
         accelerator: Optional[Accelerator] = None,
         parallel_devices: Optional[List[torch.device]] = None,
-        checkpoint_io: Optional[CheckpointIO] = None,
-        precision: Optional[Precision] = None,
+        checkpoint_io: Optional[XLACheckpointIO] = None,
+        precision: Optional[XLAPrecision] = None,
         sync_module_states: bool = True,
     ) -> None:
         super().__init__(
@@ -55,7 +54,6 @@ def __init__(
             checkpoint_io=checkpoint_io,
             precision=precision,
         )
-        self._checkpoint_io: Optional[CheckpointIO]
         self._backward_sync_control = None  # XLA synchronizes gradients in the optimizer.step() call
         self._launched = False
         self._sync_module_states = sync_module_states
@@ -72,16 +70,34 @@ def root_device(self) -> torch.device:
     def num_processes(self) -> int:
         return len(self.parallel_devices) if self.parallel_devices is not None else 0
 
-    @property
-    def checkpoint_io(self) -> CheckpointIO:
-        if self._checkpoint_io is None:
-            self._checkpoint_io = XLACheckpointIO()
-        return self._checkpoint_io
+    @property  # type: ignore[override]
+    def checkpoint_io(self) -> XLACheckpointIO:
+        plugin = self._checkpoint_io
+        if plugin is not None:
+            assert isinstance(plugin, XLACheckpointIO)
+            return plugin
+        return XLACheckpointIO()
 
     @checkpoint_io.setter
-    def checkpoint_io(self, io: CheckpointIO) -> None:
+    def checkpoint_io(self, io: Optional[XLACheckpointIO]) -> None:
+        if io is not None and not isinstance(io, XLACheckpointIO):
+            raise TypeError(f"The XLA strategy can only work with the `XLACheckpointIO` plugin, found {io}")
         self._checkpoint_io = io
 
+    @property  # type: ignore[override]
+    def precision(self) -> XLAPrecision:
+        plugin = self._precision
+        if plugin is not None:
+            assert isinstance(plugin, XLAPrecision)
+            return plugin
+        return XLAPrecision("32-true")
+
+    @precision.setter
+    def precision(self, precision: Optional[XLAPrecision]) -> None:
+        if precision is not None and not isinstance(precision, XLAPrecision):
+            raise TypeError(f"The XLA strategy can only work with the `XLAPrecision` plugin, found {precision}")
+        self._precision = precision
+
     @property
     def global_rank(self) -> int:
         return super().global_rank if self._launched else 0

src/lightning/fabric/strategies/xla_fsdp.py

@@ -24,10 +24,9 @@
 from torch.utils.data import DataLoader
 
 from lightning.fabric.accelerators import Accelerator
-from lightning.fabric.accelerators.xla import _using_pjrt
+from lightning.fabric.accelerators.xla import _XLA_AVAILABLE, _using_pjrt
 from lightning.fabric.plugins import XLAPrecision
 from lightning.fabric.plugins.environments import XLAEnvironment
-from lightning.fabric.plugins.io.checkpoint_io import CheckpointIO
 from lightning.fabric.plugins.io.xla import XLACheckpointIO
 from lightning.fabric.strategies import ParallelStrategy, _StrategyRegistry
 from lightning.fabric.strategies.fsdp import _apply_filter
@@ -85,22 +84,23 @@ def __init__(
         self,
         accelerator: Optional[Accelerator] = None,
         parallel_devices: Optional[List[torch.device]] = None,
-        checkpoint_io: Optional[CheckpointIO] = None,
+        checkpoint_io: Optional[XLACheckpointIO] = None,
         precision: Optional[XLAPrecision] = None,
         auto_wrap_policy: Optional[_POLICY] = None,
         activation_checkpointing_policy: Optional[_POLICY_SET] = None,
         state_dict_type: Literal["full", "sharded"] = "sharded",
         sequential_save: bool = False,
         **kwargs: Any,
     ) -> None:
+        if not _XLA_AVAILABLE:
+            raise ModuleNotFoundError(str(_XLA_AVAILABLE))
         super().__init__(
             accelerator=accelerator,
             parallel_devices=parallel_devices,
             cluster_environment=XLAEnvironment(),
             checkpoint_io=checkpoint_io,
             precision=precision,
         )
-        self._checkpoint_io: Optional[CheckpointIO]
         self._backward_sync_control = _XLAFSDPBackwardSyncControl()
 
         self._auto_wrap_policy = auto_wrap_policy
@@ -122,16 +122,34 @@ def root_device(self) -> torch.device:
     def num_processes(self) -> int:
         return len(self.parallel_devices) if self.parallel_devices is not None else 0
 
-    @property
-    def checkpoint_io(self) -> CheckpointIO:
-        if self._checkpoint_io is None:
-            self._checkpoint_io = XLACheckpointIO()
-        return self._checkpoint_io
+    @property  # type: ignore[override]
+    def checkpoint_io(self) -> XLACheckpointIO:
+        plugin = self._checkpoint_io
+        if plugin is not None:
+            assert isinstance(plugin, XLACheckpointIO)
+            return plugin
+        return XLACheckpointIO()
 
     @checkpoint_io.setter
-    def checkpoint_io(self, io: Optional[CheckpointIO]) -> None:
+    def checkpoint_io(self, io: Optional[XLACheckpointIO]) -> None:
+        if io is not None and not isinstance(io, XLACheckpointIO):
+            raise TypeError(f"The XLA strategy can only work with the `XLACheckpointIO` plugin, found {io}")
         self._checkpoint_io = io
 
+    @property  # type: ignore[override]
+    def precision(self) -> XLAPrecision:
+        plugin = self._precision
+        if plugin is not None:
+            assert isinstance(plugin, XLAPrecision)
+            return plugin
+        return XLAPrecision("32-true")
+
+    @precision.setter
+    def precision(self, precision: Optional[XLAPrecision]) -> None:
+        if precision is not None and not isinstance(precision, XLAPrecision):
+            raise TypeError(f"The XLA FSDP strategy can only work with the `XLAPrecision` plugin, found {precision}")
+        self._precision = precision
+
     @property
     def global_rank(self) -> int:
         return super().global_rank if self._launched else 0
@@ -227,21 +245,8 @@ def setup_optimizer(self, optimizer: Optimizer) -> Optimizer:
         flattened parameters.
 
         """
-        if _TORCH_GREATER_EQUAL_2_0:
-            return optimizer
-
-        from torch_xla.distributed.fsdp.xla_flatten_params_wrapper import FlatParameter
-
-        num_groups = len(optimizer.param_groups)
-        if num_groups > 1:
-            raise ValueError(
-                "An optimizer used with an XLAFSDP model does not support multiple param groups."
-                f" Found {num_groups} parameter groups."
-            )
-
-        if any(isinstance(param, FlatParameter) for param in optimizer.param_groups[0]["params"]):
+        if any(getattr(p, "_is_sharded", False) for group in optimizer.param_groups for p in group["params"]):
             return optimizer
-
         raise ValueError(
             "The optimizer does not seem to reference any XLAFSDP parameters. HINT: Make sure to create the optimizer"
             " after setting up the model."
@@ -470,7 +475,8 @@ def _save_checkpoint_shard(
             # convert the state
             if isinstance(obj, Module) and isinstance(obj, XLAFSDP):
                 converted = obj.state_dict()
-                # add shard_metadata to state
+                # add shard_metadata to state. this format is defined by
+                # https://github.com/pytorch/xla/blob/v2.1.0/torch_xla/distributed/fsdp/state_dict_utils.py#L122-L125
                 converted_state["shard_metadata"] = obj.get_shard_metadata()
             elif isinstance(obj, Optimizer):
                 converted = obj.state_dict()
@@ -561,11 +567,7 @@ def load_checkpoint(
             if len(loaded_metadata_keys):
                 for key in loaded_metadata_keys:
                     metadata[key] = sharded_ckpt[key]
-
-            # remove "shard_metadata" that is loaded in
-            if "shard_metadata" in metadata:
-                metadata.pop("shard_metadata")
-
+            metadata.pop("shard_metadata", None)
             return metadata
 
         if self._state_dict_type == "full":

src/lightning/pytorch/CHANGELOG.md

@@ -15,6 +15,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added support for the `max_size_cycle|max_size|min_size` iteration modes during evaluation ([#17163](https://github.com/Lightning-AI/lightning/pull/17163))
 - Added support for the TPU-v4 architecture ([#17227](https://github.com/Lightning-AI/lightning/pull/17227))
 - Added support for XLA's new PJRT runtime ([#17352](https://github.com/Lightning-AI/lightning/pull/17352))
+- Added support for Fully Sharded Data Parallel (FSDP) training with XLA ([#18746](https://github.com/Lightning-AI/lightning/pull/18746))
 - Check for invalid TPU device inputs ([#17227](https://github.com/Lightning-AI/lightning/pull/17227))
 - Added `XLAStrategy(sync_module_states=bool)` to control whether to broadcast the parameters to all devices ([#17522](https://github.com/Lightning-AI/lightning/pull/17522))
 - Added support for multiple optimizer parameter groups when using the FSDP strategy ([#17309](https://github.com/Lightning-AI/lightning/pull/17309))

src/lightning/pytorch/plugins/precision/xla.py

@@ -13,9 +13,10 @@
 # limitations under the License.
 import os
 from functools import partial
-from typing import Any, Callable
+from typing import Any, Callable, Union
 
 import torch
+from torch.optim import Optimizer
 from typing_extensions import get_args
 
 import lightning.pytorch as pl
@@ -59,6 +60,9 @@ def __init__(self, precision: _PRECISION_INPUT = "32-true") -> None:
         else:
             self._desired_dtype = torch.float32
 
+        # boolean flag for simplicity over an entirely new class
+        self._using_fsdp = False
+
     def optimizer_step(  # type: ignore[override]
         self,
         optimizer: Optimizable,
@@ -68,7 +72,8 @@ def optimizer_step(  # type: ignore[override]
     ) -> Any:
         import torch_xla.core.xla_model as xm
 
-        closure = partial(self._xla_wrap_closure, optimizer, closure)
+        if not self._using_fsdp:
+            closure = partial(self._reduce_gradients, optimizer, closure)
         closure = partial(self._wrap_closure, model, optimizer, closure)
         closure_result = optimizer.step(closure=closure, **kwargs)
         xm.mark_step()
@@ -87,9 +92,22 @@ def teardown(self) -> None:
         os.environ.pop("XLA_USE_BF16", None)
         os.environ.pop("XLA_USE_F16", None)
 
-    def _xla_wrap_closure(self, optimizer: Optimizable, closure: Callable[[], Any]) -> Any:
+    def _reduce_gradients(self, optimizer: Optimizable, closure: Callable[[], Any]) -> Any:
         import torch_xla.core.xla_model as xm
 
         closure_result = closure()
         xm.reduce_gradients(optimizer)
         return closure_result
+
+    def clip_grad_by_norm(self, optimizer: Optimizer, clip_val: Union[int, float]) -> None:
+        if self._using_fsdp:
+            # Not supported by us because we need a module reference, this would need to go through the Strategy
+            # as in Fabric
+            raise NotImplementedError("XLA's FSDP strategy does not support to clip gradients by norm.")
+        return super().clip_grad_by_value(optimizer, clip_val)
+
+    def clip_grad_by_value(self, optimizer: Optimizer, clip_val: Union[int, float]) -> None:
+        if self._using_fsdp:
+            # Not supported by XLA
+            raise NotImplementedError("XLA's FSDP strategy does not support to clip gradients by value.")
+        return super().clip_grad_by_value(optimizer, clip_val)

src/lightning/pytorch/strategies/__init__.py

@@ -23,6 +23,7 @@
 from lightning.pytorch.strategies.single_xla import SingleDeviceXLAStrategy  # noqa: F401
 from lightning.pytorch.strategies.strategy import Strategy
 from lightning.pytorch.strategies.xla import XLAStrategy  # noqa: F401
+from lightning.pytorch.strategies.xla_fsdp import XLAFSDPStrategy  # noqa: F401
 
 StrategyRegistry = _StrategyRegistry()
 _register_classes(StrategyRegistry, "register_strategies", sys.modules[__name__], Strategy)