Replace lazy MFSDP context initialization with an explicit scope

megatron/core/distributed/fsdp/mcore_fsdp_adapter.py

@@ -57,6 +57,7 @@
         Flat,
         Placements,
         fully_shard,
+        fully_shard_context,
     )
 
     HAVE_MEGATRON_FSDP = True
@@ -563,19 +564,22 @@ def __init__(
         placements = Placements(
             dp_axes=[0], parameter=[Flat()], gradient=[Flat()], optimizer=[Flat()]
         )
-        for submodule in reversed(list(module.modules())):
-            if submodule is module:
-                # The root is always sharded after selected child units so it is not
-                # wrapped twice when its type also appears in fsdp_unit_modules.
-                continue
-            if any(isinstance(submodule, module_type) for module_type in fsdp_unit_modules):
-                fully_shard(
-                    submodule,
-                    mesh=mesh,
-                    placements=placements,
-                    mixed_precision_policy=self.mp_policy,
-                )
-        fully_shard(module, mesh=mesh, placements=placements, mixed_precision_policy=self.mp_policy)
+        with fully_shard_context(device=device):
+            for submodule in reversed(list(module.modules())):
+                if submodule is module:
+                    # The root is always sharded after selected child units so it is not
+                    # wrapped twice when its type also appears in fsdp_unit_modules.
+                    continue
+                if any(isinstance(submodule, module_type) for module_type in fsdp_unit_modules):
+                    fully_shard(
+                        submodule,
+                        mesh=mesh,
+                        placements=placements,
+                        mixed_precision_policy=self.mp_policy,
+                    )
+            fully_shard(
+                module, mesh=mesh, placements=placements, mixed_precision_policy=self.mp_policy
+            )
         super().__init__(config=config, module=module)
 
     @staticmethod

megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/__init__.py

@@ -15,7 +15,7 @@
 """Experimental Megatron-FSDP implementation."""
 
 from .dbuffer import DBuffer
-from .fully_shard import fully_shard, microbatch
+from .fully_shard import fully_shard, fully_shard_context, microbatch
 from .optimizer import fully_shard_optimizer
 from .placement import Flat, Partial, Placement, Placements, Replicate
 
@@ -27,6 +27,7 @@
     "Placements",
     "Replicate",
     "fully_shard",
+    "fully_shard_context",
     "fully_shard_optimizer",
     "microbatch",
 ]

megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/fully_shard.py

@@ -17,14 +17,48 @@
 import dataclasses
 from collections.abc import Iterator
 from contextlib import contextmanager
+from contextvars import ContextVar
 
+import torch
 from torch import nn
 from torch.distributed import DeviceMesh
 
 from ..mixed_precision import MixedPrecisionPolicy
 from .module import FsdpContext, FsdpModule
 from .placement import MeshAxis, Placements
 
+_FSDP_CONTEXT = ContextVar[FsdpContext | None]("megatron_fsdp_context", default=None)
+
+
+@contextmanager
+def fully_shard_context(device: torch.device | None = None) -> Iterator[FsdpContext]:
+    """Construct FSDP modules that share runtime streams and prefetch orders.
+
+    Independent roots are ordered by their root-level ``fully_shard`` calls.
+    Construction must finish before any of the registered modules run forward.
+
+    Args:
+        device: CUDA device on which to create communication streams. Defaults to
+            the current CUDA device.
+    """
+    if _FSDP_CONTEXT.get() is not None:
+        raise RuntimeError("fully_shard_context does not support nesting.")
+
+    device = device or torch.device("cuda", torch.cuda.current_device())
+    if device.type != "cuda":
+        raise ValueError(f"fully_shard_context requires a CUDA device, got {device}.")
+
+    context = FsdpContext(device=device)
+    token = _FSDP_CONTEXT.set(context)
+    try:
+        yield context
+    except Exception:
+        raise
+    else:
+        context.finalize()
+    finally:
+        _FSDP_CONTEXT.reset(token)
+
 
 def fully_shard(
     module: nn.Module,
@@ -49,6 +83,9 @@ def fully_shard(
     """
     if isinstance(module, FsdpModule):
         raise ValueError("This module is already managed by FSDP.")
+    context = _FSDP_CONTEXT.get()
+    if context is None:
+        raise RuntimeError("fully_shard must run inside fully_shard_context.")
 
     placements = _normalize_placements(mesh, placements)
     mixed_precision_policy = mixed_precision_policy or MixedPrecisionPolicy()
@@ -58,6 +95,7 @@ def fully_shard(
         assert isinstance(module, FsdpModule)
         FsdpModule.__init__(
             module,
+            context=context,
             mesh=mesh,
             placements=placements,
             mixed_precision_policy=mixed_precision_policy,
@@ -90,28 +128,25 @@ def _axis_index(mesh: DeviceMesh, axis: MeshAxis) -> int:
 
 
 @contextmanager
-def microbatch(module: nn.Module, is_last: bool) -> Iterator[None]:
+def microbatch(context: FsdpContext, is_last: bool) -> Iterator[None]:
     """Mark an FSDP microbatch as the last accumulation microbatch.
 
     At present, this is only needed for HSDP/HFSDP gradient accumulation, so
     FSDP finalizes gradients only on the last backward. Plain all-Flat data
     parallelism finalizes gradients on every backward and does not need it.
 
     Args:
-        module: Module tree whose FSDP roots should use this microbatch state.
+        context: FSDP context whose roots should use this microbatch state.
         is_last: Whether forwards in this scope are for the last microbatch.
     """
-    contexts: list[FsdpContext] = []
-    _collect_fsdp_contexts(module, contexts)
-    previous_states = [(context, context.is_last_microbatch) for context in contexts]
-    for context in contexts:
-        context.is_last_microbatch = is_last
+    context.ensure_finalized()
+    previous_state = context.is_last_microbatch
+    context.is_last_microbatch = is_last
 
     try:
         yield
     finally:
-        for context, is_last_microbatch in previous_states:
-            context.is_last_microbatch = is_last_microbatch
+        context.is_last_microbatch = previous_state
 
 
 def _attach_mixin(module: nn.Module) -> None:
@@ -120,13 +155,3 @@ def _attach_mixin(module: nn.Module) -> None:
     module_cls = module.__class__
     fsdp_cls = type(f"ExperimentalFsdp{module_cls.__name__}", (FsdpModule, module_cls), {})
     module.__class__ = fsdp_cls
-
-
-def _collect_fsdp_contexts(module: nn.Module, contexts: list[FsdpContext]) -> None:
-    if isinstance(module, FsdpModule):
-        module._lazy_init_context()
-        contexts.append(module.context)
-        return
-
-    for child in module.children():
-        _collect_fsdp_contexts(child, contexts)

megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py

@@ -28,36 +28,76 @@
 
 
 class FsdpContext:
-    """Runtime stream and prefetch state shared by one FSDP subtree."""
+    """Runtime stream and prefetch state shared by FSDP roots constructed together."""
 
     allgather_stream: torch.cuda.Stream
     reduce_scatter_stream: torch.cuda.Stream
     # HFSDP/HSDP need explicit last-microbatch state. First-microbatch state is
     # unnecessary because it can be detected when ``model_weight``, after syncing
     # from ``main_weight``, has placements different from ``Placements.optimizer``.
     is_last_microbatch: bool
-    root_module: "FsdpModule"
     # Static orders used to drive all-gather prefetch. We may want to switch to
     # capturing runtime order if static module order proves too fragile. Each
     # FsdpModule tracks its own materialized state via ``FsdpModule._unshard_event``.
     forward_order: IndexedOrder["FsdpModule"]
     backward_order: IndexedOrder["FsdpModule"]
 
-    def __init__(self, device: torch.device, root_module: "FsdpModule") -> None:
-        """Create rank-local runtime state for a root FSDP subtree.
+    def __init__(self, device: torch.device) -> None:
+        """Create rank-local runtime state for FSDP modules on ``device``.
 
         Args:
             device: Device on which this context schedules communication.
-            root_module: Outermost module that owns this context.
         """
-        self.root_module = root_module
         self.is_last_microbatch = True
         self.forward_order = IndexedOrder()
         self.backward_order = IndexedOrder()
+        self._registered_modules: list[FsdpModule] = []
+        self._is_finalized = False
         with torch.cuda.device(device):
             self.allgather_stream = torch.cuda.Stream()
             self.reduce_scatter_stream = torch.cuda.Stream()
 
+    def register_module(self, module: "FsdpModule") -> None:
+        """Register a module constructed in this context."""
+        if self._is_finalized:
+            raise RuntimeError("Cannot register an FSDP module after its context is finalized.")
+        self._registered_modules.append(module)
+
+    def finalize(self) -> None:
+        """Finalize roots, names, and cross-root prefetch orders."""
+        if self._is_finalized:
+            raise RuntimeError("FSDP context is already finalized.")
+
+        visited: set[FsdpModule] = set()
+        root_modules: list[FsdpModule] = []
+        for module in reversed(self._registered_modules):
+            if module in visited:
+                continue
+            root_modules.append(module)
+            _collect_fsdp_modules_under(cast(nn.Module, module), visited)
+        root_modules.reverse()
+        for root_module in root_modules:
+            root_module._is_root = True
+
+        for root_module in root_modules:
+            for name, module in cast(nn.Module, root_module).named_modules():
+                if not isinstance(module, FsdpModule):
+                    continue
+                module._name = name
+                self.forward_order.append(module)
+
+        for root_module in reversed(root_modules):
+            _collect_backward_order(cast(nn.Module, root_module), self.backward_order)
+
+        self._is_finalized = True
+
+    def ensure_finalized(self) -> None:
+        """Raise if construction has not completed for this context."""
+        if not self._is_finalized:
+            raise RuntimeError(
+                "FSDP context is not finalized. Exit fully_shard_context before running forward."
+            )
+
     def current_stream(self) -> torch.cuda.Stream:
         """Current stream on this context's device."""
         return torch.cuda.current_stream(self.allgather_stream.device)
@@ -84,7 +124,8 @@ class FsdpModule:
     # Root uses "" and None means uninitialized.
     _name: str | None
     _parameter_groups: tuple[FsdpParameterGroup, ...]
-    _context: FsdpContext | None
+    _context: FsdpContext
+    _is_root: bool
     _ready_grad_parameters: set[nn.Parameter]
     _num_trainable_parameters: int
     # Event recorded after this FsdpModule's full parameters are materialized.
@@ -94,13 +135,15 @@ class FsdpModule:
 
     def __init__(
         self,
+        context: FsdpContext,
         mesh: DeviceMesh,
         placements: Placements,
         mixed_precision_policy: MixedPrecisionPolicy,
         use_symm_mem: bool = False,
     ) -> None:
         """Initialize FSDP runtime state on an already-constructed module."""
-        self._context = None
+        self._context = context
+        self._is_root = False
         self._name = None
         self._unshard_event = None
         owned_parameters = _collect_owned_parameters(self)
@@ -124,60 +167,11 @@ def __init__(
             len(group.sharded_parameters) for group in self._parameter_groups if group.requires_grad
         )
         self._register_hooks()
-
-    def _lazy_init_context(self) -> None:
-        """Initialize one shared runtime context for this FSDP root subtree.
-
-        MFSDP v2 requires users to apply ``fully_shard`` bottom-up, so child FSDP
-        modules are constructed before their eventual root module is constructed.
-        This method resolves the root lazily on the first forward through the
-        outermost FSDP module and shares that one context with every FSDP
-        descendant.
-
-        Alternatives considered:
-        - Eagerly initialize contexts during ``fully_shard``. When a parent is
-          sharded, we could create a new root context and reassign it to all
-          descendant FSDP modules. This creates transient child contexts that are
-          never used if the parent is later sharded, and each parent shard must
-          walk its descendants again, making nested sharding quadratic.
-        - Store an ``is_root`` field on each FSDP module. ``fully_shard`` could
-          mark newly sharded modules as roots and clear that flag on descendant
-          FSDP modules when a parent is sharded. This avoids creating unused
-          contexts but moves root tracking onto every FSDP module, adding
-          per-module state that must stay consistent with the final sharded
-          module hierarchy.
-        """
-        if self._context is not None:
-            return
-
-        root_module = cast(nn.Module, self)
-        first_parameter = next(root_module.parameters(), None)
-        if first_parameter is None:
-            raise RuntimeError("FSDP root module requires at least one parameter in its subtree.")
-
-        context = FsdpContext(device=first_parameter.device, root_module=self)
-        # named_modules() yields FsdpModules in registration order, which is the static
-        # forward execution order used to prefetch the next FsdpModule's all-gather.
-        for submodule_name, submodule in root_module.named_modules():
-            if not isinstance(submodule, FsdpModule):
-                continue
-            if submodule._context is not None:
-                raise RuntimeError(
-                    "FSDP context is already initialized for a descendant module. "
-                    "Run forward through the root FSDP module first."
-                )
-            submodule._context = context
-            submodule._name = submodule_name
-            context.forward_order.append(submodule)
-
-        # Backward starts from the root pre-backward hook before visiting child
-        # subtrees in reverse module order.
-        _collect_backward_order(root_module, context.backward_order)
+        context.register_module(self)
 
     @property
     def context(self) -> FsdpContext:
-        """Return the initialized runtime context."""
-        assert self._context is not None
+        """Return the FSDP context."""
         return self._context
 
     @property
@@ -189,8 +183,8 @@ def name(self) -> str:
         return name
 
     def is_root(self) -> bool:
-        """Return whether this module is the outermost FsdpModule in its context."""
-        return self.context.root_module is self
+        """Return whether this module is an outermost FsdpModule in its context."""
+        return self._is_root
 
     def _register_hooks(self) -> None:
         module = cast(nn.Module, self)
@@ -227,10 +221,10 @@ def pre_forward(self) -> None:
         While this FsdpModule computes, we issue the next FsdpModule's all-gather
         on the comm stream, so ``AG_{i+1}`` is launched before ``F_i`` finishes.
         """
-        self._lazy_init_context()
+        context = self.context
+        context.ensure_finalized()
         torch.cuda.nvtx.range_push(self._nvtx_label("forward"))
         self._ready_grad_parameters.clear()
-        context = self.context
         allgather_stream = context.allgather_stream
         current_stream = context.current_stream()
 
@@ -348,8 +342,19 @@ def _nvtx_label(self, phase: Literal["forward", "backward"]) -> str:
         return f"MFSDP {name} {phase}"
 
 
+def _collect_fsdp_modules_under(module: nn.Module, modules: set["FsdpModule"]) -> None:
+    """Collect FSDP modules reachable from ``module``."""
+    if isinstance(module, FsdpModule):
+        if module in modules:
+            return
+        modules.add(module)
+
+    for child in module.children():
+        _collect_fsdp_modules_under(child, modules)
+
+
 def _collect_backward_order(module: nn.Module, order: IndexedOrder["FsdpModule"]) -> None:
-    """Collect FsdpModules in static backward prefetch order."""
+    """Collect one root's static backward prefetch order."""
     if isinstance(module, FsdpModule):
         order.append(module)