Add forward all-gather overlap

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

@@ -14,6 +14,8 @@
 
 """Module mixin for the minimal Megatron-FSDP path."""
 
+import dataclasses
+from collections import deque
 from collections.abc import Callable
 from typing import Literal, cast
 
@@ -26,23 +28,54 @@
 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 shared by one experimental FSDP subtree."""
+    """Runtime state, stream, and release scheduler shared by one FSDP subtree."""
 
+    allgather_stream: torch.cuda.Stream
+    delayed_releases: deque[DelayedRelease]
     # 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"
 
-    def __init__(self, root_module: "FsdpModule") -> None:
+    def __init__(self, device: torch.device, root_module: "FsdpModule") -> None:
         """Create rank-local runtime state for a root FSDP subtree.
 
         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.delayed_releases = deque()
+        with torch.cuda.device(device):
+            self.allgather_stream = torch.cuda.Stream()
+
+    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))
+
+    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}.")
+
+        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()
 
 
 class FsdpModule:
@@ -114,7 +147,7 @@ def _lazy_init_context(self) -> None:
         if self._context is not None:
             return
 
-        context = FsdpContext(root_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():
             if not isinstance(submodule, FsdpModule):
                 continue
@@ -172,33 +205,64 @@ def pre_forward(self) -> None:
         self._lazy_init_context()
         torch.cuda.nvtx.range_push(self._nvtx_label("forward"))
         self._ready_grad_parameters.clear()
-        for group in self._parameter_groups:
-            group.sync_model_weight_from_main_weight()
-            group.unshard_parameters()
+        if self.is_root():
+            allgather_stream = self.context.allgather_stream
+            allgather_stream.wait_stream(torch.cuda.current_stream(allgather_stream.device))
+        self._unshard_parameter_groups(sync_model_weight=True)
+
+    def _unshard_parameter_groups(self, *, sync_model_weight: bool) -> None:
+        """Materialize full parameters for this FSDP unit."""
+        self.context.drain_delayed_releases(target_length=1)
+
+        allgather_stream = self.context.allgather_stream
+        current_stream = torch.cuda.current_stream(allgather_stream.device)
+
+        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)
 
     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:
         for group in self._parameter_groups:
             group.reshard_parameters()
-        torch.cuda.nvtx.range_pop()
 
     def pre_backward(self) -> None:
         """Prepare full parameters for backward compute."""
         torch.cuda.nvtx.range_push(self._nvtx_label("backward"))
-        for group in self._parameter_groups:
-            group.unshard_parameters()
+        self._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()
-            group.reshard_parameters()
+        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 FSDP unit."""
+        for group in self._parameter_groups:
+            group.release_unsharded_storage()
+
+    @property
     def parameter_groups(self) -> tuple[FsdpParameterGroup, ...]:
-        """Return parameter groups owned by this FSDP unit."""
+        """Parameter groups owned by this FSDP unit."""
         return self._parameter_groups
 
     def _nvtx_label(self, phase: Literal["forward", "backward"]) -> str:

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

@@ -232,12 +232,17 @@ def unshard_parameters(self) -> None:
     def reshard_parameters(self) -> None:
         """Install sharded DTensor parameters on the owning modules."""
         self._switch_to_sharded_parameters()
-        # At post-backward time, replacing unsharded parameter .data with size-0
-        # empty tensors would also be safe: autograd has consumed the saved
-        # forward views. That alternative is not much cleaner than releasing
-        # this storage, and splitting post-forward and post-backward reshard
-        # behavior would make the caller code less clean, so keep the shared
-        # storage-release path.
+
+    def release_unsharded_storage(self) -> None:
+        """Release this group's full-parameter storage."""
+        # This method is shared by the post-forward and post-backward release
+        # paths. Post-forward must release storage because autograd may have
+        # saved forward views into the unsharded parameters. Post-backward could
+        # replace unsharded parameter .data with size-0 empty tensors, instead
+        # of releasing storage, because autograd has consumed those saved views.
+        # That alternative is not much cleaner, and splitting post-forward and
+        # post-backward reshard behavior would make the caller code less clean,
+        # so keep the shared storage-release path.
         self._unsharded_model_weight.release_storage()
 
     def reduce_gradients(self) -> None:

tests/unit_tests/distributed/mfsdp_v2/test_context.py

@@ -83,7 +83,7 @@ def test_two_child_subtrees_then_parent_collapse_to_one_context(distributed_setu
 
 
 def test_sibling_roots_without_parent_keep_separate_contexts(distributed_setup):
-    """Independent FSDP roots should not share runtime state."""
+    """Independent FSDP roots should not share runtime scheduling state."""
     device = distributed_setup.device
 
     mesh = init_device_mesh(device.type, (distributed_setup.world_size,))