Support HSDP deferred DP-outer gradient reduction

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

@@ -292,8 +292,9 @@ def redistribute(
         """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.
+        Flat -> Replicate, Partial -> Replicate, Partial -> Flat,
+        Replicate -> Flat, and Replicate -> Partial. Other placement changes are
+        intentionally unsupported.
         """
         new_placements = tuple(new_placements)
         if len(new_placements) != self.mesh.ndim:
@@ -322,6 +323,23 @@ def redistribute(
             return self.reduce_scatter(axis, new_placement, out=out)
         if isinstance(old_placement, Replicate) and isinstance(new_placement, Flat):
             return self.scatter(axis, new_placement, out=out)
+        if isinstance(old_placement, Replicate) and isinstance(new_placement, Partial):
+            # Replicate and Partial share the same local layout, so relabel the
+            # buffer without communication. Value-preserving for AVG only -- the
+            # mean of identical per-rank locals is that value; SUM would need a
+            # 1/axis_size rescale, which no caller needs.
+            if new_placement.reduce_op != dist.ReduceOp.AVG:
+                raise NotImplementedError(
+                    "Replicate -> Partial redistribute supports AVG only, got "
+                    f"{new_placement.reduce_op!r}."
+                )
+            if out is not None:
+                raise NotImplementedError(
+                    "Replicate -> Partial redistribute does not support an out buffer."
+                )
+            return DBuffer.from_local(
+                self.local_buffer, self.mesh, new_placements, self.layout.tensor_shapes
+            )
         raise NotImplementedError(
             "Unsupported DBuffer placement transition on axis "
             f"{axis}: {old_placement!r} -> {new_placement!r}."
@@ -452,7 +470,17 @@ def get_dtensor(self, index: int) -> DTensor:
             elif isinstance(placement, Flat):
                 torch_placements.append(dist_tensor.Shard(0))
             elif isinstance(placement, Partial):
-                raise ValueError("Partial DBuffer placements cannot be represented as DTensor.")
+                # main_grad backs .grad while it rests DP-outer-Partial between
+                # microbatches, so a Partial placement must round-trip to a DTensor.
+                if placement.reduce_op == dist.ReduceOp.AVG:
+                    reduce_op = "avg"
+                elif placement.reduce_op == dist.ReduceOp.SUM:
+                    reduce_op = "sum"
+                else:
+                    raise ValueError(
+                        f"Unsupported Partial reduce op for DTensor: {placement.reduce_op!r}."
+                    )
+                torch_placements.append(dist_tensor.Partial(reduce_op))
             else:
                 raise TypeError(f"Unsupported placement for DTensor conversion: {placement!r}.")
 

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

@@ -341,7 +341,7 @@ def _reduce_gradient_groups(self) -> None:
 
             reduce_scatter_stream.wait_stream(current_stream)
             with torch.cuda.stream(reduce_scatter_stream):
-                group.reduce_partial_gradients(partial_grad)
+                group.reduce_partial_gradients(partial_grad, self.context.is_last_microbatch)
 
     @property
     def parameter_groups(self) -> tuple[FsdpParameterGroup, ...]:

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

@@ -150,13 +150,9 @@ def __init__(
                 "main_grad is built from main_weight tensor shapes on the same mesh, "
                 "and DBuffer layouts are deterministic from those shapes and mesh size."
             )
-            if self.main_grad.placements != self.main_weight.placements:
-                raise ValueError(
-                    "FSDP temporarily requires main_grad and main_weight to have the same "
-                    "placements until HSDP/HFSDP support is implemented. "
-                    f"Got main_grad placements {self.main_grad.placements} and "
-                    f"main_weight placements {self.main_weight.placements}."
-                )
+            # main_grad rests here (DP-outer-Partial for HSDP) between microbatches and
+            # is finalized to main_weight's placements after the last microbatch.
+            self._accumulation_placements = main_grad_placements
         sharded_parameters: list[nn.Parameter] = []
         unsharded_parameters: list[nn.Parameter] = []
         main_grad_dtype = self.main_grad.dtype if self.main_grad is not None else None
@@ -251,6 +247,12 @@ def release_unsharded_storage(self) -> None:
         # so keep the shared storage-release path.
         self._unsharded_model_weight.release_storage()
 
+    def _install_sharded_grads(self) -> None:
+        """Point each sharded parameter's grad at main_grad's current DTensor view."""
+        assert self.main_grad is not None
+        for index, sharded_parameter in enumerate(self.sharded_parameters):
+            sharded_parameter.grad = self.main_grad.get_dtensor(index)
+
     def allocate_partial_grad_buffer(self) -> DBuffer:
         """Allocate the unreduced reduce-scatter input buffer."""
         assert self.main_grad is not None
@@ -279,8 +281,18 @@ def copy_gradients_to_partial_buffer(self, partial_grad: DBuffer) -> None:
             partial_grad.get_local_tensor(index).copy_(parameter.grad)
             parameter.grad = None
 
-    def reduce_partial_gradients(self, partial_grad: DBuffer) -> None:
-        """Reduce a packed partial gradient buffer into sharded parameter gradients."""
+    def reduce_partial_gradients(
+        self, partial_grad: DBuffer, is_last_microbatch: bool = True
+    ) -> None:
+        """Reduce a packed partial gradient buffer into sharded parameter gradients.
+
+        For HSDP main_grad rests DP-outer-Partial (Partial where main_weight is
+        Replicate) between microbatches, accumulating each backward through the
+        standard zero_grad contract; the last microbatch reduces the DP-outer axes,
+        finalizing main_grad to main_weight's placements so ``.grad`` is the fully
+        reduced gradient before ``optimizer.step()``. With every axis Flat (plain
+        DP) main_grad already rests finalized.
+        """
         assert self.main_grad is not None
 
         def has_grad(parameters: tuple[nn.Parameter, ...]) -> bool:
@@ -295,10 +307,20 @@ def has_grad(parameters: tuple[nn.Parameter, ...]) -> bool:
                 raise RuntimeError("FSDP sharded gradients must be either all set or all None.")
             return has_any_grad
 
-        # zero_grad(set_to_none=True) clears sharded parameter grads, so the next
+        # zero_grad(set_to_none=True) clears sharded parameter grads, so this
         # 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)
+
+        # A non-accumulation main_grad means the previous step finalized it; this
+        # only happens on the first microbatch. Redistribute it back to the
+        # DP-outer-Partial accumulation placement -- a metadata relabel for HSDP,
+        # and a fresh reduce-scattered buffer for HFSDP in the future.
+        if self.main_grad.placements != self._accumulation_placements:
+            self.main_grad = self.main_grad.redistribute(self._accumulation_placements)
+            if has_sharded_grads:
+                self._install_sharded_grads()
+
         can_reduce_into_main_grad = (
             not has_sharded_grads and partial_grad.dtype == self.main_grad.dtype
         )
@@ -321,10 +343,14 @@ def has_grad(parameters: tuple[nn.Parameter, ...]) -> bool:
                 self.main_grad.local_buffer.add_(reduced_grad.local_buffer)
             else:
                 self.main_grad.local_buffer.copy_(reduced_grad.local_buffer)
-
         if not has_sharded_grads:
-            for index, parameter in enumerate(self.sharded_parameters):
-                parameter.grad = self.main_grad.get_dtensor(index)
+            self._install_sharded_grads()
+
+        if is_last_microbatch:
+            # Finalize the deferred DP-outer reduction (all-reduce for HSDP,
+            # reduce-scatter for HFSDP) and install the sharded parameter gradients.
+            self.main_grad = self.main_grad.redistribute(self.main_weight.placements)
+            self._install_sharded_grads()
 
 
 def _get_parameter_owner(module: nn.Module, name: str) -> tuple[nn.Module, str]:

tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py

@@ -6,14 +6,17 @@
 
 import pytest
 import torch
+import torch.distributed as dist
 from torch import nn
 from torch.distributed.device_mesh import init_device_mesh
 from torch.distributed.tensor import DTensor
 from torch.profiler import ProfilerActivity, profile
 
 from megatron.core.distributed.fsdp.src.megatron_fsdp.experimental import (
     Flat,
+    Partial,
     Placements,
+    Replicate,
     fully_shard,
     microbatch,
 )
@@ -100,6 +103,18 @@ def _flat_placements() -> Placements:
     return Placements(dp_axes=[0], parameter=[Flat()], gradient=[Flat()], optimizer=[Flat()])
 
 
+def _hsdp_placements() -> Placements:
+    """HSDP: params/optimizer replicated across DP-outer (axis 0), sharded within
+    DP-inner (axis 1). main_grad rests [Partial, Flat] between microbatches and is
+    all-reduced to [Replicate, Flat] on the last microbatch."""
+    return Placements(
+        dp_axes=[0, 1],
+        parameter=[Replicate(), Flat()],
+        gradient=[Partial(dist.ReduceOp.AVG), Flat()],
+        optimizer=[Replicate(), Flat()],
+    )
+
+
 def _mb(num_bytes: int) -> str:
     return f"{num_bytes / 1024**2:.2f} MB"
 
@@ -111,6 +126,17 @@ def _events_overlap(first, second) -> bool:
     )
 
 
+def _nccl_events(cuda_events, *name_fragments):
+    """CUDA NCCL events whose name contains any of ``name_fragments`` (case-insensitive)."""
+    return [
+        event
+        for event in cuda_events
+        if "nccl" in event.name.lower()
+        and event.activity_type == "kernel"
+        and any(fragment in event.name.lower() for fragment in name_fragments)
+    ]
+
+
 @pytest.mark.parametrize("num_microbatches", [1, 3])
 def test_fully_shard_losses_match_baseline(distributed_setup, num_microbatches):
     """Minimal per-module FSDP training should match single-rank SGD."""
@@ -168,6 +194,147 @@ def train(model, optimizer, log_prefix) -> list[torch.Tensor]:
     )
 
 
+@pytest.mark.parametrize("set_to_none", [True, False])
+@pytest.mark.parametrize("num_microbatches", [1, 3])
+def test_hsdp_losses_match_baseline(distributed_setup, num_microbatches, set_to_none):
+    """HSDP (DP-outer replicated, DP-inner sharded) training should match single-rank SGD.
+
+    Gradients reduce-scatter within DP-inner every backward and accumulate into
+    main_grad; the DP-outer all-reduce runs only on the last microbatch, scoped
+    via ``microbatch(...)``. Every rank sees identical data, so the averaged
+    gradient equals the single-rank gradient and losses must match. Both
+    ``zero_grad`` modes are covered: ``set_to_none=True`` overwrites main_grad,
+    ``set_to_none=False`` accumulates into a zeroed main_grad.
+    """
+    rank = distributed_setup.rank
+    world_size = distributed_setup.world_size
+    device = distributed_setup.device
+    if world_size < 4 or world_size % 2 != 0:
+        pytest.skip("This test requires an even number of at least 4 ranks for a 2-D DP mesh.")
+
+    outer_size = 2
+    inner_size = world_size // outer_size
+    mesh = init_device_mesh(
+        device.type, (outer_size, inner_size), mesh_dim_names=("dp_outer", "dp_inner")
+    )
+    torch.manual_seed(1234)
+    dim = 8
+    baseline = MultiChildModel(dim=dim, num_children=2).to(device)
+    model = MultiChildModel(dim=dim, num_children=2).to(device)
+    model.load_state_dict(baseline.state_dict())
+
+    # Shard the child layers, then the model, so the children share a root context
+    # and reduce through the overlap path instead of as independent roots.
+    for layer in model.layers:
+        fully_shard(layer, mesh=mesh, placements=_hsdp_placements())
+    fully_shard(model, mesh=mesh, placements=_hsdp_placements())
+    baseline_optimizer = torch.optim.SGD(baseline.parameters(), lr=0.05)
+    optimizer = torch.optim.SGD(model.parameters(), lr=0.05)
+
+    micro_batch_size = 2
+    x = torch.randn(num_microbatches, micro_batch_size, dim, device=device)
+    target = torch.randn(num_microbatches, micro_batch_size, dim, device=device)
+    microbatches = tuple(zip(x.unbind(), target.unbind()))
+
+    def train(model, optimizer, log_prefix) -> list[torch.Tensor]:
+        losses = []
+        for step in range(5):
+            optimizer.zero_grad(set_to_none=set_to_none)
+
+            for microbatch_index, (microbatch_x, microbatch_target) in enumerate(microbatches):
+                is_last = microbatch_index == num_microbatches - 1
+                with microbatch(model, is_last=is_last):
+                    loss = torch.nn.functional.mse_loss(model(microbatch_x), microbatch_target)
+                    (loss / num_microbatches).backward()
+                losses.append(loss.detach())
+                logger.debug(
+                    "%s train parity: rank=%s, step=%s, microbatch=%s, loss=%s",
+                    log_prefix,
+                    rank,
+                    step,
+                    microbatch_index,
+                    loss,
+                )
+
+            optimizer.step()
+        return losses
+
+    baseline_losses = train(baseline, baseline_optimizer, "Baseline")
+    sharded_losses = train(model, optimizer, "HSDP")
+
+    torch.testing.assert_close(
+        torch.stack(sharded_losses),
+        torch.stack(baseline_losses),
+        msg="HSDP losses did not match baseline losses.",
+    )
+
+
+def test_hsdp_defers_dp_outer_allreduce_to_last_microbatch(distributed_setup):
+    """HSDP reduce-scatters DP-inner every microbatch but all-reduces DP-outer once.
+
+    ``fully_shard(model)`` makes the child units share a root context so their
+    reductions run through the overlap path rather than as independent roots.
+    Counting NCCL events over a multi-microbatch step, the DP-inner reduce-scatter
+    fires once per microbatch per group while the DP-outer all-reduce that
+    finalizes main_grad fires only on the last microbatch, so the reduce-scatter
+    count is exactly ``num_microbatches`` times the all-reduce count. This asserts
+    on event counts only, not numerics.
+    """
+    world_size = distributed_setup.world_size
+    device = distributed_setup.device
+    if world_size < 4 or world_size % 2 != 0:
+        pytest.skip("This test requires an even number of at least 4 ranks for a 2-D DP mesh.")
+
+    outer_size = 2
+    inner_size = world_size // outer_size
+    mesh = init_device_mesh(
+        device.type, (outer_size, inner_size), mesh_dim_names=("dp_outer", "dp_inner")
+    )
+    torch.manual_seed(1234)
+    dim = 8
+    num_children = 2
+    model = MultiChildModel(dim=dim, num_children=num_children).to(device)
+    for layer in model.layers:
+        fully_shard(layer, mesh=mesh, placements=_hsdp_placements())
+    fully_shard(model, mesh=mesh, placements=_hsdp_placements())
+    optimizer = torch.optim.SGD(model.parameters(), lr=0.05)
+
+    num_microbatches = 3
+    micro_batch_size = 2
+    x = torch.randn(num_microbatches, micro_batch_size, dim, device=device)
+    target = torch.randn(num_microbatches, micro_batch_size, dim, device=device)
+    microbatches = tuple(zip(x.unbind(), target.unbind()))
+
+    def train_one_step() -> None:
+        optimizer.zero_grad(set_to_none=True)
+        for microbatch_index, (microbatch_x, microbatch_target) in enumerate(microbatches):
+            is_last = microbatch_index == num_microbatches - 1
+            with microbatch(model, is_last=is_last):
+                loss = torch.nn.functional.mse_loss(model(microbatch_x), microbatch_target)
+                (loss / num_microbatches).backward()
+        optimizer.step()
+
+    train_one_step()
+    torch.cuda.synchronize(device)
+
+    with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
+        train_one_step()
+        torch.cuda.synchronize(device)
+
+    cuda_events = [event for event in prof.events() if event.device_type.name == "CUDA"]
+    reduce_scatter_events = _nccl_events(cuda_events, "reducescatter", "reduce_scatter")
+    all_reduce_events = _nccl_events(cuda_events, "allreduce")
+    # One DP-outer all-reduce per parameter group -- each child layer plus the
+    # root unit's bias -- fired only on the last microbatch. Plain DP fires none.
+    assert len(all_reduce_events) == num_children + 1, [event.name for event in cuda_events]
+    # DP-inner reduce-scatter runs every microbatch; the DP-outer all-reduce runs
+    # only on the last, so the counts differ by exactly the microbatch factor.
+    assert len(reduce_scatter_events) == len(all_reduce_events) * num_microbatches, (
+        f"Expected reduce-scatter ({len(reduce_scatter_events)}) to be {num_microbatches}x "
+        f"the DP-outer all-reduce count ({len(all_reduce_events)})."
+    )
+
+
 def test_nested_fully_shard_excludes_child_owned_parameters(distributed_setup):
     """An outer FsdpModule owns direct parameters but not nested child FsdpModule parameters."""
     world_size = distributed_setup.world_size
@@ -353,26 +520,9 @@ def train_one_iteration() -> None:
         # drop the CUDA events.
         torch.cuda.synchronize(device)
 
-    # Keep only real device kernels. NCCL also emits per-collective GPU user
-    # annotations (e.g. "nccl:_reduce_scatter_base") that the profiler reports as
-    # CUDA events. Filtering by activity type avoids counting those annotations as
-    # kernels without relying on their current naming convention.
-    cuda_events = [
-        event
-        for event in prof.events()
-        if event.device_type.name == "CUDA" and event.activity_type == "kernel"
-    ]
-    all_gather_events = [
-        event
-        for event in cuda_events
-        if "nccl" in event.name.lower() and "allgather" in event.name.lower()
-    ]
-    reduce_scatter_events = [
-        event
-        for event in cuda_events
-        if "nccl" in event.name.lower()
-        and ("reducescatter" in event.name.lower() or "reduce_scatter" in event.name.lower())
-    ]
+    cuda_events = [event for event in prof.events() if event.device_type.name == "CUDA"]
+    all_gather_events = _nccl_events(cuda_events, "allgather")
+    reduce_scatter_events = _nccl_events(cuda_events, "reducescatter", "reduce_scatter")
     # GEMM device-kernel names vary across CUDA/cuBLAS versions and GPU archs
     # (e.g. "*gemm*", "cutlass*", "cublas*", and cuBLASLt's Hopper "nvjet_sm90_*").
     gemm_events = [