Improve FSDP wrapper implementation

Author: cbalioglu

src/fairseq2/models/_handler.py

@@ -30,7 +30,8 @@
     UnknownModelArchitectureError,
     model_asset_card_error,
 )
-from fairseq2.nn.data_parallel import load_with_sdp_gang
+from fairseq2.models.fsdp import apply_fsdp_to_transformer
+from fairseq2.nn.data_parallel import FsdpGranularity, FsdpWrapper, load_with_sdp_gang
 from fairseq2.nn.utils.module import (
     load_state_dict,
     reset_non_persistent_buffers,
@@ -79,6 +80,11 @@ def load_from_path(
     @abstractmethod
     def compile(self, model: Module, config: object) -> Module: ...
 
+    @abstractmethod
+    def apply_fsdp(
+        self, model: Module, granularity: FsdpGranularity, wrapper: FsdpWrapper
+    ) -> Module: ...
+
     @property
     @abstractmethod
     def family(self) -> str: ...
@@ -461,6 +467,7 @@ def _do_create(
 
         return model
 
+    @override
     def compile(self, model: Module, config: object) -> Module:
         if self._torch_compiler is None:
             raise NotSupportedError(
@@ -479,6 +486,12 @@ def compile(self, model: Module, config: object) -> Module:
 
         return self._torch_compiler(model, config)
 
+    @override
+    def apply_fsdp(
+        self, model: Module, granularity: FsdpGranularity, wrapper: FsdpWrapper
+    ) -> Module:
+        return apply_fsdp_to_transformer(model, granularity, wrapper)
+
     @property
     @override
     def family(self) -> str:

src/fairseq2/models/fsdp.py

@@ -6,57 +6,37 @@
 
 from __future__ import annotations
 
-from functools import partial
-from typing import Literal
-
-from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
-    CheckpointWrapper,
-)
-from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
 from torch.nn import Module
 
-from fairseq2.nn.data_parallel import FSDPWrapPolicy
-from fairseq2.nn.transformer import (
-    TransformerDecoder,
-    TransformerDecoderLayer,
-    TransformerEncoder,
-    TransformerEncoderLayer,
-)
-
-
-def get_fsdp_wrap_policy(
-    model: Module, wrap_granularity: Literal["layer", "stack", "model"] = "layer"
-) -> tuple[FSDPWrapPolicy | None, list[Module] | None]:
-    """Return the FSDP wrap policy for ``model`` along with ignored modules.
-
-    :param model: The model to be wrapped.
-    :param wrap_granularity: The granularity at which to wrap modules of ``model``.
-          - 'layer': Wraps individual layers (e.g. :class:`TransformerDecoderLayer`).
-          - 'stack': Wraps layer stacks (e.g. :class:`TransformerDecoder`).
-          - 'model': Wraps ``model``.
-    """
-    if wrap_granularity == "model":
-        return None, None
-
-    kls: set[type[Module]]
-
-    if wrap_granularity == "stack":
-        kls = {TransformerEncoder, TransformerDecoder}
-    elif wrap_granularity == "layer":
-        kls = {TransformerEncoderLayer, TransformerDecoderLayer}
-
-        # We make the assumption that if the model uses activation checkpointing,
-        # it is at the layer granularity.
-        for m in model.modules():
-            if isinstance(m, CheckpointWrapper):
-                kls = {CheckpointWrapper}
-
-                break
-    else:
-        raise ValueError(
-            f"`wrap_granularity` must be 'layer', 'stack', or 'model', but is '{wrap_granularity}' instead."
-        )
-
-    wrap_policy = partial(transformer_auto_wrap_policy, transformer_layer_cls=kls)
-
-    return wrap_policy, None
+from fairseq2.nn.data_parallel import FsdpGranularity, FsdpWrapper
+from fairseq2.nn.transformer import TransformerDecoder, TransformerEncoder
+
+
+def apply_fsdp_to_transformer(
+    model: Module, granularity: FsdpGranularity, wrapper: FsdpWrapper
+) -> Module:
+    if granularity == "model":
+        return wrapper(model)
+
+    children = list(model.named_children())
+
+    for name, child in children:
+        if isinstance(child, (TransformerEncoder, TransformerDecoder)):
+            if granularity == "stack":
+                model.register_module(name, wrapper(child))
+            else:
+                layers = list(child.layers.named_children())
+
+                for idx, (layer_name, layer) in enumerate(layers):
+                    # We don't need to reshard the last layer since we will
+                    # immediately gather it for the backward pass.
+                    if idx < len(layers) - 1:
+                        reshard_after_forward = None
+                    else:
+                        reshard_after_forward = False
+
+                    child.layers.register_module(
+                        layer_name, wrapper(layer, reshard_after_forward)
+                    )
+
+    return model

src/fairseq2/nn/data_parallel/__init__.py

@@ -10,20 +10,12 @@
 from fairseq2.nn.data_parallel._error import (
     DistributedSetupError as DistributedSetupError,
 )
+from fairseq2.nn.data_parallel._fsdp import FsdpApplier as FsdpApplier
+from fairseq2.nn.data_parallel._fsdp import FsdpGranularity as FsdpGranularity
 from fairseq2.nn.data_parallel._fsdp import (
-    FSDP_LOW_MEMORY_POLICY as FSDP_LOW_MEMORY_POLICY,
+    FsdpParameterInitializer as FsdpParameterInitializer,
 )
-from fairseq2.nn.data_parallel._fsdp import (
-    FSDP_STANDARD_MEMORY_POLICY as FSDP_STANDARD_MEMORY_POLICY,
-)
-from fairseq2.nn.data_parallel._fsdp import (
-    FSDP_VERY_LOW_MEMORY_POLICY as FSDP_VERY_LOW_MEMORY_POLICY,
-)
-from fairseq2.nn.data_parallel._fsdp import FSDPMemoryPolicy as FSDPMemoryPolicy
-from fairseq2.nn.data_parallel._fsdp import (
-    FSDPParameterInitializer as FSDPParameterInitializer,
-)
-from fairseq2.nn.data_parallel._fsdp import FSDPWrapPolicy as FSDPWrapPolicy
+from fairseq2.nn.data_parallel._fsdp import FsdpWrapper as FsdpWrapper
 from fairseq2.nn.data_parallel._fsdp import (
     fsdp_local_state_dict as fsdp_local_state_dict,
 )