Overlap experimental Megatron-FSDP communication with two process groups

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

@@ -284,13 +284,21 @@ def cast(self, dtype: torch.dtype) -> "DBuffer":
         return destination
 
     def redistribute(
-        self, new_placements: Iterable[Placement], *, out: "DBuffer | None" = None
+        self,
+        new_placements: Iterable[Placement],
+        *,
+        out: "DBuffer | None" = None,
+        reduce_group: "dist.ProcessGroup | None" = None,
     ) -> "DBuffer":
         """Redistribute this buffer to ``new_placements``.
 
         This dispatcher supports the one-axis transitions:
         Flat -> Replicate, Partial -> Replicate, Partial -> Flat, and
         Replicate -> Flat. Other placement changes are intentionally unsupported.
+
+        ``reduce_group`` optionally overrides the process group for the
+        Partial -> Flat reduce-scatter, letting the caller run reduce-scatter on a
+        communicator separate from the mesh's default (all-gather) group.
         """
         new_placements = tuple(new_placements)
         if len(new_placements) != self.mesh.ndim:
@@ -316,7 +324,7 @@ def redistribute(
         if isinstance(old_placement, Partial) and isinstance(new_placement, Replicate):
             return self.allreduce(axis, out=out)
         if isinstance(old_placement, Partial) and isinstance(new_placement, Flat):
-            return self.reduce_scatter(axis, new_placement, out=out)
+            return self.reduce_scatter(axis, new_placement, out=out, group=reduce_group)
         if isinstance(old_placement, Replicate) and isinstance(new_placement, Flat):
             return self.scatter(axis, new_placement, out=out)
         raise NotImplementedError(
@@ -359,9 +367,18 @@ def allreduce(self, mesh_axis: int, *, out: "DBuffer | None" = None) -> "DBuffer
         return out
 
     def reduce_scatter(
-        self, mesh_axis: int, new_placement: Placement, *, out: "DBuffer | None" = None
+        self,
+        mesh_axis: int,
+        new_placement: Placement,
+        *,
+        out: "DBuffer | None" = None,
+        group: "dist.ProcessGroup | None" = None,
     ) -> "DBuffer":
-        """Reduce-scatter a Partial axis into ``new_placement``."""
+        """Reduce-scatter a Partial axis into ``new_placement``.
+
+        ``group`` defaults to the mesh's axis group; pass a dedicated communicator
+        to run reduce-scatter separately from all-gather.
+        """
         axis = mesh_axis
         if not isinstance(new_placement, Flat):
             raise NotImplementedError("DBuffer currently supports reduce_scatter() to Flat only.")
@@ -377,7 +394,7 @@ def reduce_scatter(
             output=out.local_buffer,
             input=self.local_buffer,
             op=partial_placement.reduce_op,
-            group=self.mesh.get_group(axis),
+            group=group if group is not None else self.mesh.get_group(axis),
         )
         return out
 

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

@@ -20,6 +20,7 @@
 from typing import Literal, cast
 
 import torch
+import torch.distributed as dist
 from torch import nn
 from torch.distributed import DeviceMesh
 
@@ -37,9 +38,11 @@ class DelayedRelease:
 
 
 class FsdpContext:
-    """Runtime state, stream, and release scheduler shared by one FSDP subtree."""
+    """Runtime state, streams, and release scheduler shared by one FSDP subtree."""
 
     allgather_stream: torch.cuda.Stream
+    reduce_scatter_stream: torch.cuda.Stream
+    reduce_scatter_group: dist.ProcessGroup
     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
@@ -57,8 +60,27 @@ 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._post_backward_callback_queued = False
         with torch.cuda.device(device):
             self.allgather_stream = torch.cuda.Stream()
+            self.reduce_scatter_stream = torch.cuda.Stream()
+        # Dedicated NCCL communicator for reduce-scatter, distinct from the mesh's
+        # all-gather group. A separate communicator lets an eager reduce-scatter run
+        # concurrently with the next unit's all-gather instead of contending on one
+        # comm's ordered stream -- which is what makes it safe to drop the delayed
+        # reduction and launch reduce-scatter immediately in post_backward. new_group
+        # is collective, so every rank reaches this on the first forward through the
+        # root. Prototype: 1D DP mesh only (the general case needs one dedicated group
+        # per mesh-axis subgroup, mirroring DeviceMesh construction).
+        mesh = root_module._parameter_groups[0].mesh
+        if mesh.ndim != 1:
+            raise NotImplementedError(
+                "Prototype dedicated reduce-scatter communicator supports a 1D DP mesh "
+                f"only; got a {mesh.ndim}D mesh."
+            )
+        self.reduce_scatter_group = dist.new_group(
+            ranks=dist.get_process_group_ranks(mesh.get_group(0))
+        )
 
     def enqueue_release(self, module: "FsdpModule") -> None:
         """Queue a module's unsharded storage for delayed release."""
@@ -77,6 +99,36 @@ def drain_delayed_releases(self, target_length: int) -> None:
                     self.allgather_stream.wait_event(delayed_release.consumer_event)
                 delayed_release.module.release_unsharded_storage()
 
+    def queue_post_backward_callback(self) -> None:
+        """Queue an end-of-backward callback so the optimizer waits for reduce-scatter."""
+        if self._post_backward_callback_queued:
+            return
+
+        self._post_backward_callback_queued = True
+        try:
+            torch.autograd.Variable._execution_engine.queue_callback(self.finalize_reductions)
+        except RuntimeError as error:
+            self._post_backward_callback_queued = False
+            if str(error) != "Final callbacks can only be installed during backward pass.":
+                raise
+            self.finalize_reductions()
+
+    def finalize_reductions(self) -> None:
+        """Order the default stream (the optimizer's) after all eager reduce-scatters.
+
+        Reduce-scatters are launched eagerly in ``post_backward`` on
+        ``reduce_scatter_stream``; their input and output buffers are allocated on and
+        consumed by that stream, so the CUDA caching allocator keeps the storage alive
+        until each collective completes -- no explicit pending-buffer retention needed.
+        This single stream barrier makes the default stream, where the optimizer reads
+        the reduced gradients, wait for those reductions to finish.
+        """
+        try:
+            default_stream = torch.cuda.current_stream(self.reduce_scatter_stream.device)
+            default_stream.wait_stream(self.reduce_scatter_stream)
+        finally:
+            self._post_backward_callback_queued = False
+
 
 class FsdpModule:
     """Mixin attached to modules managed by the minimal FSDP path."""
@@ -245,16 +297,49 @@ def pre_backward(self) -> None:
 
     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 _reduce_gradient_groups(self) -> None:
+        # Eagerly launch each group's reduce-scatter on the dedicated reduce-scatter
+        # communicator/stream as soon as its gradients are packed. Because that comm is
+        # separate from the all-gather group, an eager reduce-scatter does not serialize
+        # against the next unit's all-gather, so no deferral to the next pre_backward and
+        # no prepared/pending buffer bookkeeping are needed -- only the end-of-backward
+        # barrier in finalize_reductions. Delayed *releases* are unchanged.
+        context = self.context
+        default_stream = torch.cuda.current_stream(context.reduce_scatter_stream.device)
+        scheduled_reduction = False
+        for group in self._parameter_groups:
+            if not group.requires_grad:
+                continue
+            with torch.cuda.stream(context.reduce_scatter_stream):
+                partial_grad = group.allocate_partial_grad_buffer()
+
+            # Pack on the default stream (where grads are produced), so reads of
+            # parameter.grad and the subsequent `.grad = None` stay same-stream.
+            default_stream.wait_stream(context.reduce_scatter_stream)
+            group.copy_gradients_to_partial_buffer(partial_grad)
+
+            # Reduce-scatter waits for the pack, then runs eagerly on the reduce-scatter
+            # stream. partial_grad was allocated on that stream, so dropping the reference
+            # here is safe: the caching allocator retains the storage until the collective
+            # completes.
+            context.reduce_scatter_stream.wait_stream(default_stream)
+            with torch.cuda.stream(context.reduce_scatter_stream):
+                group.reduce_partial_gradients(
+                    partial_grad, reduce_group=context.reduce_scatter_group
+                )
+            scheduled_reduction = True
+
+        if scheduled_reduction:
+            context.queue_post_backward_callback()
+
     def release_unsharded_storage(self) -> None:
         """Release unsharded storage owned by this FSDP unit."""
         for group in self._parameter_groups:

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

@@ -14,7 +14,6 @@
 
 """Parameter-group runtime state for the minimal Megatron-FSDP path."""
 
-from collections.abc import Iterable
 from contextlib import nullcontext
 
 import torch
@@ -87,6 +86,11 @@ def __init__(
         self.dtype = first_parameter.dtype
         self.requires_grad = first_parameter.requires_grad
         for name, parameter in parameters.items():
+            if parameter.is_meta:
+                raise ValueError(
+                    f"Expected parameter {name!r} to be materialized before "
+                    "FsdpParameterGroup construction."
+                )
             if parameter.dtype != self.dtype:
                 raise ValueError(
                     f"Expected parameter {name!r} to have dtype {self.dtype}, "
@@ -245,55 +249,76 @@ def release_unsharded_storage(self) -> None:
         # so keep the shared storage-release path.
         self._unsharded_model_weight.release_storage()
 
-    def reduce_gradients(self) -> None:
-        """Reduce full local gradients into sharded parameter gradients."""
+    def allocate_partial_grad_buffer(self) -> DBuffer:
+        """Allocate the unreduced reduce-scatter input buffer."""
         assert self.main_grad is not None
 
-        def has_grad(parameters: Iterable[nn.Parameter]) -> bool:
-            has_any_grad = False
-            has_any_missing_grad = False
-            for parameter in parameters:
-                if parameter.grad is None:
-                    has_any_missing_grad = True
-                else:
-                    has_any_grad = True
-            if has_any_grad and has_any_missing_grad:
-                raise RuntimeError("FSDP sharded gradients must be either all set or all None.")
-            return has_any_grad
-
-        grads: list[torch.Tensor] = []
-        for name, parameter in zip(self.parameter_names, self.unsharded_parameters, strict=True):
-            if parameter.grad is None:
-                raise RuntimeError(f"Missing gradient for FSDP parameter {name!r}.")
-            grads.append(parameter.grad)
-
         # NCCL symmetric-memory reduce-scatter only selects the symmetric kernel for SUM today.
         # Preserve AVG semantics by reducing SUM and scaling the output below.
         partial_op = dist.ReduceOp.AVG if self._symm_mem_pool is None else dist.ReduceOp.SUM
+        grads = self._require_unsharded_grads()
         with self._symmetric_memory_context():
-            partial_grad = DBuffer.distribute_tensors(
-                grads, mesh=self.mesh, placements=[Partial(partial_op)] * self.mesh.ndim
+            return DBuffer(
+                mesh=self.mesh,
+                placements=[Partial(partial_op)] * self.mesh.ndim,
+                tensor_shapes=tuple(grad.shape for grad in grads),
+                dtype=grads[0].dtype,
+                device=grads[0].device,
             )
 
+    def copy_gradients_to_partial_buffer(self, partial_grad: DBuffer) -> None:
+        """Pack full local gradients into an existing reduce-scatter input buffer."""
+        grads = self._require_unsharded_grads()
+        # This packs per-parameter grads into the reduce-scatter input buffer. A future
+        # fused-wgrad path can avoid this copy by writing directly into those buffer views.
+        for index, grad in enumerate(grads):
+            partial_grad.get_local_tensor(index).copy_(grad)
+        for parameter in self.unsharded_parameters:
+            parameter.grad = None
+
+    def reduce_partial_gradients(
+        self, partial_grad: DBuffer, *, reduce_group: "dist.ProcessGroup | None" = None
+    ) -> None:
+        """Reduce a packed partial gradient buffer into sharded parameter gradients.
+
+        ``reduce_group`` optionally routes the reduce-scatter through a dedicated
+        communicator so it can run concurrently with all-gather. All temporaries
+        (``partial_grad`` and any intermediate ``reduced_grad``) are allocated on and
+        consumed by the reduce-scatter stream, so the caller can drop its reference on
+        return: the caching allocator keeps the storage alive until the same-stream
+        reduce-scatter completes -- no explicit pending-buffer bookkeeping needed.
+        """
+        assert self.main_grad is not None
+
         # zero_grad(set_to_none=True) clears sharded parameter grads, so the next
         # backward can reduce directly into main_grad. zero_grad(set_to_none=False)
         # leaves sharded grads installed, so this backward accumulates into main_grad.
-        has_sharded_grads = has_grad(self.sharded_parameters)
+        has_sharded_grads = self._sharded_parameters_have_grad()
         can_reduce_into_main_grad = (
             not has_sharded_grads and partial_grad.dtype == self.main_grad.dtype
         )
         reduce_axis = changed_mesh_axis(partial_grad.placements, self.main_grad.placements)
         if reduce_axis is None:
             raise RuntimeError("FSDP gradient reduction requires a changed placement axis.")
-        grad_divisor = self.mesh.size(reduce_axis) if partial_op == dist.ReduceOp.SUM else 1
+        partial_placement = partial_grad.placements[reduce_axis]
+        assert isinstance(partial_placement, Partial)
+        grad_divisor = (
+            self.mesh.size(reduce_axis)
+            if partial_placement.reduce_op == dist.ReduceOp.SUM
+            else 1
+        )
         if self._symm_mem_pool is not None:
             partial_grad.rendezvous(reduce_axis)
         if can_reduce_into_main_grad:
-            partial_grad.redistribute(self.main_grad.placements, out=self.main_grad)
+            partial_grad.redistribute(
+                self.main_grad.placements, out=self.main_grad, reduce_group=reduce_group
+            )
             if grad_divisor != 1:
                 self.main_grad.local_buffer.div_(grad_divisor)
         else:
-            reduced_grad = partial_grad.redistribute(self.main_grad.placements)
+            reduced_grad = partial_grad.redistribute(
+                self.main_grad.placements, reduce_group=reduce_group
+            )
             if grad_divisor != 1:
                 reduced_grad.local_buffer.div_(grad_divisor)
             if has_sharded_grads:
@@ -305,8 +330,25 @@ def has_grad(parameters: Iterable[nn.Parameter]) -> bool:
             for index, parameter in enumerate(self.sharded_parameters):
                 parameter.grad = self.main_grad.get_dtensor(index)
 
-        for parameter in self.unsharded_parameters:
-            parameter.grad = None
+    def _require_unsharded_grads(self) -> tuple[torch.Tensor, ...]:
+        grads: list[torch.Tensor] = []
+        for name, parameter in zip(self.parameter_names, self.unsharded_parameters, strict=True):
+            if parameter.grad is None:
+                raise RuntimeError(f"Missing gradient for FSDP parameter {name!r}.")
+            grads.append(parameter.grad)
+        return tuple(grads)
+
+    def _sharded_parameters_have_grad(self) -> bool:
+        has_any_grad = False
+        has_any_missing_grad = False
+        for parameter in self.sharded_parameters:
+            if parameter.grad is None:
+                has_any_missing_grad = True
+            else:
+                has_any_grad = True
+        if has_any_grad and has_any_missing_grad:
+            raise RuntimeError("FSDP sharded gradients must be either all set or all None.")
+        return has_any_grad
 
 
 def _get_parameter_owner(module: nn.Module, name: str) -> tuple[nn.Module, str]: