Overlap FSDP communication with compute

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

@@ -0,0 +1,53 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+"""Ordered sequence with indexed item lookup."""
+
+from collections.abc import Iterator
+from typing import Generic, TypeVar
+
+T = TypeVar("T")
+
+
+class IndexedOrder(Generic[T]):
+    """Insertion order with constant-time successor lookup by item."""
+
+    def __init__(self) -> None:
+        """Create an empty indexed order."""
+        self._items: list[T] = []
+        self._index_by_item: dict[T, int] = {}
+
+    def append(self, item: T) -> None:
+        """Append ``item`` to the order.
+
+        Args:
+            item: Item to append.
+
+        Raises:
+            ValueError: If ``item`` is already present in the order.
+        """
+        if item in self._index_by_item:
+            raise ValueError("IndexedOrder does not support duplicate items.")
+        self._index_by_item[item] = len(self._items)
+        self._items.append(item)
+
+    def __iter__(self) -> Iterator[T]:
+        """Iterate over items in order."""
+        return iter(self._items)
+
+    def next_item(self, item: T) -> T | None:
+        """Return the item that follows ``item``, if any."""
+        index = self._index_by_item[item]
+        next_index = index + 1
+        return self._items[next_index] if next_index < len(self._items) else None

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

@@ -14,8 +14,6 @@
 
 """Module mixin for the minimal Megatron-FSDP path."""
 
-import dataclasses
-from collections import deque
 from collections.abc import Callable
 from typing import Literal, cast
 
@@ -24,28 +22,26 @@
 from torch.distributed import DeviceMesh
 
 from ..mixed_precision import MixedPrecisionPolicy
+from .indexed_order import IndexedOrder
 from .parameter_group import FsdpParameterGroup, contained_in_parameter_group
 from .placement import MeshAxis, Placements
 
 
-@dataclasses.dataclass(frozen=True)
-class DelayedRelease:
-    """A module whose unsharded storage can be released after its consumer event."""
-
-    consumer_event: torch.cuda.Event | None
-    module: "FsdpModule"
-
-
 class FsdpContext:
-    """Runtime state, stream, and release scheduler shared by one FSDP subtree."""
+    """Runtime stream and prefetch state shared by one FSDP subtree."""
 
     allgather_stream: torch.cuda.Stream
-    delayed_releases: deque[DelayedRelease]
+    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.
@@ -56,26 +52,29 @@ def __init__(self, device: torch.device, root_module: "FsdpModule") -> None:
         """
         self.root_module = root_module
         self.is_last_microbatch = True
-        self.delayed_releases = deque()
+        self.forward_order = IndexedOrder()
+        self.backward_order = IndexedOrder()
         with torch.cuda.device(device):
             self.allgather_stream = torch.cuda.Stream()
+            self.reduce_scatter_stream = torch.cuda.Stream()
+
+    def current_stream(self) -> torch.cuda.Stream:
+        """Current stream on this context's device."""
+        return torch.cuda.current_stream(self.allgather_stream.device)
+
+    def register_post_backward_final_callback(self) -> None:
+        """Register this root context's final callback for the current backward.
 
-    def enqueue_release(self, module: "FsdpModule") -> None:
-        """Queue a module's unsharded storage for delayed release."""
-        consumer_event = torch.cuda.current_stream(self.allgather_stream.device).record_event()
-        self.delayed_releases.append(DelayedRelease(consumer_event=consumer_event, module=module))
+        Root ``post_backward()`` means only that root-owned parameters have
+        accumulated gradients; it may run before descendant reductions, or not
+        run at all when the root owns no trainable parameters. Waiting at
+        autograd completion orders consumers after every descendant reduction.
+        """
 
-    def drain_delayed_releases(self, target_length: int) -> None:
-        """Release queued module storages FIFO until the queue reaches ``target_length``."""
-        if target_length < 0:
-            raise ValueError(f"target_length must be non-negative, got {target_length}.")
+        def post_backward_final_callback() -> None:
+            self.current_stream().wait_stream(self.reduce_scatter_stream)
 
-        while len(self.delayed_releases) > target_length:
-            delayed_release = self.delayed_releases.popleft()
-            with torch.cuda.stream(self.allgather_stream):
-                if delayed_release.consumer_event is not None:
-                    self.allgather_stream.wait_event(delayed_release.consumer_event)
-                delayed_release.module.release_unsharded_storage()
+        torch.autograd.Variable._execution_engine.queue_callback(post_backward_final_callback)
 
 
 class FsdpModule:
@@ -88,6 +87,10 @@ class FsdpModule:
     _context: FsdpContext | None
     _ready_grad_parameters: set[nn.Parameter]
     _num_trainable_parameters: int
+    # Event recorded after this FsdpModule's full parameters are materialized.
+    # ``None`` lets pre_forward enqueue an all-gather unless an earlier FsdpModule
+    # already prefetched this module.
+    _unshard_event: torch.cuda.Event | None
 
     def __init__(
         self,
@@ -99,6 +102,7 @@ def __init__(
         """Initialize FSDP runtime state on an already-constructed module."""
         self._context = None
         self._name = None
+        self._unshard_event = None
         owned_parameters = _collect_owned_parameters(self)
         axis_indices = tuple(_axis_index(mesh, axis) for axis in placements.dp_axes)
         assert axis_indices == tuple(
@@ -147,8 +151,11 @@ def _lazy_init_context(self) -> None:
         if self._context is not None:
             return
 
+        root_module = cast(nn.Module, self)
         context = FsdpContext(device=self._parameter_groups[0].main_weight.device, root_module=self)
-        for submodule_name, submodule in cast(nn.Module, self).named_modules():
+        # 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:
@@ -158,6 +165,11 @@ def _lazy_init_context(self) -> None:
                 )
             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)
 
     @property
     def context(self) -> FsdpContext:
@@ -207,64 +219,125 @@ def grad_hook(_parameter: nn.Parameter) -> None:
         return grad_hook
 
     def pre_forward(self) -> None:
-        """Prepare full parameters for forward compute."""
+        """Prepare full parameters for forward compute and prefetch the next FsdpModule.
+
+        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()
         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()
+
         if self.is_root():
-            allgather_stream = self.context.allgather_stream
-            allgather_stream.wait_stream(torch.cuda.current_stream(allgather_stream.device))
+            allgather_stream.wait_stream(current_stream)
+
         self._unshard_parameter_groups(sync_model_weight=True)
+        assert self._unshard_event is not None
+        # Compute waits only for this FsdpModule's all-gather (the prefetch below is
+        # issued afterwards, so it is free to run concurrently with this FsdpModule).
+        current_stream.wait_event(self._unshard_event)
+
+        next_module = context.forward_order.next_item(self)
+        if next_module is not None:
+            next_module._unshard_parameter_groups(sync_model_weight=True)
 
     def _unshard_parameter_groups(self, *, sync_model_weight: bool) -> None:
-        """Materialize full parameters for this FsdpModule."""
-        self.context.drain_delayed_releases(target_length=1)
+        """Unshard this FsdpModule's parameter groups on the all-gather stream.
 
-        allgather_stream = self.context.allgather_stream
-        current_stream = torch.cuda.current_stream(allgather_stream.device)
+        If ``_unshard_event`` is already set, this FsdpModule was already
+        unsharded or prefetched and this method is a no-op. Otherwise, this
+        method records ``_unshard_event`` after materialization so compute
+        can wait without depending on later release work.
+        """
+        if self._unshard_event is not None:
+            return
 
+        allgather_stream = self.context.allgather_stream
         with torch.cuda.stream(allgather_stream):
             for group in self._parameter_groups:
                 if sync_model_weight:
                     # TODO: After NVIDIA/Megatron-LM#5411 lands, move this sync to the
                     # optimizer post-step hook instead of running it every microbatch.
                     group.sync_model_weight_from_main_weight()
                 group.unshard_parameters()
-        current_stream.wait_stream(allgather_stream)
+            self._unshard_event = allgather_stream.record_event()
 
     def post_forward(self) -> None:
         """Return parameters to their sharded resting state after forward compute."""
         self._reshard_parameter_groups()
-        self.context.enqueue_release(self)
-        if self.is_root():
-            self.context.drain_delayed_releases(target_length=0)
         torch.cuda.nvtx.range_pop()
 
     def _reshard_parameter_groups(self) -> None:
+        """Reshard parameter groups and release unsharded storage after compute.
+
+        This method clears ``_unshard_event`` after queuing the release, so
+        future users enqueue a fresh all-gather.
+        """
         for group in self._parameter_groups:
             group.reshard_parameters()
 
+        allgather_stream = self.context.allgather_stream
+        allgather_stream.wait_stream(self.context.current_stream())
+        # Release on the all-gather stream where unsharded storage was allocated,
+        # so no record_stream() call is required for the storage.
+        with torch.cuda.stream(allgather_stream):
+            for group in self._parameter_groups:
+                group.release_unsharded_storage()
+            self._unshard_event = None
+
     def pre_backward(self) -> None:
-        """Prepare full parameters for backward compute."""
+        """Prepare full parameters and prefetch the next FsdpModule in backward order."""
         torch.cuda.nvtx.range_push(self._nvtx_label("backward"))
+        context = self.context
+        current_stream = context.current_stream()
+        if self.is_root():
+            context.register_post_backward_final_callback()
+            # Fork the reduce-scatter stream from the current stream once, at the
+            # start of backward, so every module's post-backward reduce-scatter is
+            # part of any active CUDA-graph capture. A stream only joins the
+            # capture via this wait_stream edge; without it the first allocation on
+            # the reduce-scatter stream falls back to a raw cudaMalloc, which is
+            # illegal during capture. Later modules are covered by the post-copy
+            # fork each preceding module issues before its collective.
+            context.reduce_scatter_stream.wait_stream(current_stream)
+
         self._unshard_parameter_groups(sync_model_weight=False)
+        assert self._unshard_event is not None
+        current_stream.wait_event(self._unshard_event)
+
+        next_module = context.backward_order.next_item(self)
+        if next_module is not None:
+            next_module._unshard_parameter_groups(sync_model_weight=False)
 
     def post_backward(self) -> None:
         """Reduce gradients and return parameters to their sharded resting state."""
-        for group in self._parameter_groups:
-            if group.requires_grad:
-                group.reduce_gradients()
+        self._reduce_gradient_groups()
         self._reshard_parameter_groups()
-        self.context.enqueue_release(self)
-        if self.is_root():
-            self.context.drain_delayed_releases(target_length=0)
         self._ready_grad_parameters.clear()
         torch.cuda.nvtx.range_pop()
 
-    def release_unsharded_storage(self) -> None:
-        """Release unsharded storage owned by this FsdpModule."""
+    def _reduce_gradient_groups(self) -> None:
+        """Pack gradients and immediately launch their reduce-scatters."""
+        context = self.context
+        reduce_scatter_stream = context.reduce_scatter_stream
+        current_stream = context.current_stream()
+
         for group in self._parameter_groups:
-            group.release_unsharded_storage()
+            if not group.requires_grad:
+                continue
+
+            with torch.cuda.stream(reduce_scatter_stream):
+                partial_grad = group.allocate_partial_grad_buffer()
+
+            current_stream.wait_stream(reduce_scatter_stream)
+            group.copy_gradients_to_partial_buffer(partial_grad)
+
+            reduce_scatter_stream.wait_stream(current_stream)
+            with torch.cuda.stream(reduce_scatter_stream):
+                group.reduce_partial_gradients(partial_grad)
 
     @property
     def parameter_groups(self) -> tuple[FsdpParameterGroup, ...]:
@@ -276,6 +349,15 @@ def _nvtx_label(self, phase: Literal["forward", "backward"]) -> str:
         return f"MFSDP {name} {phase}"
 
 
+def _collect_backward_order(module: nn.Module, order: IndexedOrder["FsdpModule"]) -> None:
+    """Collect FsdpModules in static backward prefetch order."""
+    if isinstance(module, FsdpModule):
+        order.append(module)
+
+    for child in reversed(list(module.children())):
+        _collect_backward_order(child, order)
+
+
 def _axis_index(mesh: DeviceMesh, axis: MeshAxis) -> int:
     if isinstance(axis, int):
         axis_index = axis