works. i'm still worthy

Author: deependujha

src/lightning/fabric/strategies/fsdp.py

@@ -31,6 +31,7 @@
 from lightning_utilities.core.imports import RequirementCache
 from lightning_utilities.core.rank_zero import rank_zero_only as utils_rank_zero_only
 from torch import Tensor
+from torch.distributed.tensor import DTensor
 from torch.nn import Module
 from torch.optim import Optimizer
 from typing_extensions import TypeGuard, override
@@ -795,6 +796,10 @@ def _optimizer_has_flat_params(optimizer: Optimizer) -> bool:
     )
 
 
+def _optimizer_has_dtensor_params(optimizer: Optimizer) -> bool:
+    return any(isinstance(param, DTensor) for group in optimizer.param_groups for param in group["params"])
+
+
 def _get_sharded_state_dict_context(module: Module) -> Generator[None, None, None]:
     from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
     from torch.distributed.fsdp.api import ShardedOptimStateDictConfig, ShardedStateDictConfig, StateDictType

src/lightning/pytorch/plugins/__init__.py

@@ -8,6 +8,7 @@
 from lightning.pytorch.plugins.precision.deepspeed import DeepSpeedPrecision
 from lightning.pytorch.plugins.precision.double import DoublePrecision
 from lightning.pytorch.plugins.precision.fsdp import FSDPPrecision
+from lightning.pytorch.plugins.precision.fsdp2 import FSDP2Precision
 from lightning.pytorch.plugins.precision.half import HalfPrecision
 from lightning.pytorch.plugins.precision.precision import Precision
 from lightning.pytorch.plugins.precision.transformer_engine import TransformerEnginePrecision
@@ -28,6 +29,7 @@
     "Precision",
     "TransformerEnginePrecision",
     "FSDPPrecision",
+    "FSDP2Precision",
     "XLAPrecision",
     "LayerSync",
     "TorchSyncBatchNorm",

src/lightning/pytorch/strategies/__init__.py

@@ -18,6 +18,7 @@
 from lightning.pytorch.strategies.ddp import DDPStrategy
 from lightning.pytorch.strategies.deepspeed import DeepSpeedStrategy
 from lightning.pytorch.strategies.fsdp import FSDPStrategy
+from lightning.pytorch.strategies.fsdp2 import FSDP2Strategy
 from lightning.pytorch.strategies.model_parallel import ModelParallelStrategy
 from lightning.pytorch.strategies.parallel import ParallelStrategy
 from lightning.pytorch.strategies.single_device import SingleDeviceStrategy
@@ -32,6 +33,7 @@
     "DDPStrategy",
     "DeepSpeedStrategy",
     "FSDPStrategy",
+    "FSDP2Strategy",
     "ModelParallelStrategy",
     "ParallelStrategy",
     "SingleDeviceStrategy",

src/lightning/pytorch/strategies/fsdp2.py

@@ -41,7 +41,7 @@
     _distributed_checkpoint_load,
     _distributed_checkpoint_save,
     _move_torchmetrics_to_device,
-    _optimizer_has_flat_params,
+    _optimizer_has_dtensor_params,
 )
 from lightning.fabric.utilities.distributed import (
     _distributed_is_initialized,
@@ -139,6 +139,7 @@ def __init__(
         self.mp_policy = _init_fsdp2_mp_policy(mp_policy)
 
         self.device_mesh = device_mesh
+        self.kwargs = kwargs
 
     @property
     @override
@@ -249,12 +250,19 @@ def _setup_model(self, model: Module) -> Module:
                 )
 
             log.debug(f"setting up FSDP model with device id: {self.root_device.index}, kwargs: {self.kwargs}")
+            if isinstance(self.device_mesh, tuple):
+                from torch.distributed.device_mesh import DeviceMesh
+
+                self.device_mesh = DeviceMesh("cuda", self.device_mesh)
+
+            if self.mp_policy is None:
+                raise ValueError("`mp_policy` cannot be None when calling `fully_shard`.")
+
             fully_shard(
                 module=model,
                 mesh=self.device_mesh,
                 mp_policy=self.mp_policy,
                 offload_policy=self.cpu_offload,
-                cpu_offload=self.cpu_offload,
             )
 
             if is_on_meta_device:
@@ -321,7 +329,7 @@ def setup_optimizers(self, trainer: "pl.Trainer") -> None:
                 raise
             invalid_params_error = True
 
-        if invalid_params_error or any(not _optimizer_has_flat_params(optimizer) for optimizer in self.optimizers):
+        if invalid_params_error or any(not _optimizer_has_dtensor_params(optimizer) for optimizer in self.optimizers):
             # We avoid this limitation by setting `use_orig_params=True`
             raise ValueError(
                 "The optimizer does not seem to reference any FSDP parameters. HINT: Make sure to create the"
@@ -428,7 +436,7 @@ def register_strategies(cls, strategy_registry: _StrategyRegistry) -> None:
         cls._registered_strategies.append("fsdp2")
 
         strategy_registry.register(
-            "fsdp_cpu_offload",
+            "fsdp2_cpu_offload",
             cls,
             description="FSDP2 training with Full Sharding and CPU Offloading",
             cpu_offload=True,

src/lightning/pytorch/trainer/connectors/accelerator_connector.py

@@ -42,6 +42,7 @@
     CheckpointIO,
     DeepSpeedPrecision,
     DoublePrecision,
+    FSDP2Precision,
     FSDPPrecision,
     HalfPrecision,
     MixedPrecision,
@@ -53,6 +54,7 @@
 from lightning.pytorch.strategies import (
     DDPStrategy,
     DeepSpeedStrategy,
+    FSDP2Strategy,
     FSDPStrategy,
     ModelParallelStrategy,
     ParallelStrategy,
@@ -493,6 +495,8 @@ def _check_and_init_precision(self) -> Precision:
             return DeepSpeedPrecision(self._precision_flag)  # type: ignore[arg-type]
         if isinstance(self.strategy, FSDPStrategy):
             return FSDPPrecision(self._precision_flag)  # type: ignore[arg-type]
+        if isinstance(self.strategy, FSDP2Strategy):
+            return FSDP2Precision(self._precision_flag)  # type: ignore[arg-type]
         if self._precision_flag in ("16-true", "bf16-true"):
             return HalfPrecision(self._precision_flag)  # type: ignore
         if self._precision_flag == "32-true":