Support HSDP deferred DP-outer gradient reduction

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

@@ -449,7 +449,10 @@ 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.
+                reduce_op = "avg" if placement.reduce_op == dist.ReduceOp.AVG else "sum"
+                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

@@ -247,7 +247,7 @@ 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.reduce_gradients(self.context.is_last_microbatch)
         self._reshard_parameter_groups()
         self.context.enqueue_release(self)
         if self.is_root():

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

@@ -25,7 +25,7 @@
 
 from ..mixed_precision import MixedPrecisionPolicy
 from .dbuffer import DBuffer
-from .placement import Partial, Placements, Replicate, changed_mesh_axis
+from .placement import Partial, Placement, Placements, Replicate, changed_mesh_axis
 
 _CONTAINING_PARAMETER_GROUP_ATTR = "_mfsdp_parameter_group"
 
@@ -151,13 +151,6 @@ 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}."
-                )
         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
@@ -245,8 +238,41 @@ 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 _reduce_grad_axis(
+        self, partial_grad: DBuffer, target_placements: tuple[Placement, ...]
+    ) -> DBuffer:
+        """Reduce ``partial_grad`` to ``target_placements`` over its one changed axis.
+
+        HSDP and HFSDP change a single mesh axis between the all-Partial gradient
+        and its sharded resting placement, so ``changed_mesh_axis`` enforces the
+        single-axis contract.
+        """
+        axis = changed_mesh_axis(partial_grad.placements, target_placements)
+        if axis is None:
+            raise RuntimeError("FSDP gradient reduction requires a changed placement axis.")
+        if self._symm_mem_pool is not None:
+            partial_grad.rendezvous(axis)
+        reduced = partial_grad.redistribute(target_placements)
+        # Symmetric-memory reduce-scatter runs as SUM; rescale to recover the mean.
+        if self._symm_mem_pool is not None:
+            reduced.local_buffer.div_(self.mesh.size(axis))
+        return reduced
+
+    def reduce_gradients(self, is_last_microbatch: bool = True) -> None:
+        """Reduce full local gradients into sharded parameter gradients.
+
+        main_grad rests DP-outer-Partial (Partial where main_weight is Replicate)
+        between microbatches: each backward reduce-scatters DP-inner into it and
+        accumulates through the standard zero_grad contract (``set_to_none=True``
+        clears the sharded grads so main_grad is overwritten; ``set_to_none=False``
+        keeps them as views into a zeroed main_grad so they accumulate). The last
+        microbatch all-reduces the DP-outer axes in place, finalizing main_grad to
+        main_weight's placements so ``.grad`` -- a view into main_grad -- is the
+        fully reduced gradient before ``optimizer.step()``. A finalized main_grad
+        marks a completed step, so the next reduction resets it back to the
+        accumulation placement. With every axis Flat (plain DP) main_grad is
+        already finalized and this is the previous every-backward behavior.
+        """
         assert self.main_grad is not None
 
         def has_grad(parameters: Iterable[nn.Parameter]) -> bool:
@@ -268,42 +294,54 @@ def has_grad(parameters: Iterable[nn.Parameter]) -> bool:
             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.
+        # Preserve AVG semantics by reducing SUM and scaling the output in _reduce_grad_axis.
         partial_op = dist.ReduceOp.AVG if self._symm_mem_pool is None else dist.ReduceOp.SUM
         with self._symmetric_memory_context():
             partial_grad = DBuffer.distribute_tensors(
                 grads, mesh=self.mesh, placements=[Partial(partial_op)] * self.mesh.ndim
             )
 
-        # 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)
-        can_reduce_into_main_grad = (
-            not has_sharded_grads and partial_grad.dtype == self.main_grad.dtype
+        # DP-outer axes (Replicate in main_weight) accumulate as Partial and are
+        # all-reduced only on the last microbatch.
+        deferred_axes = [
+            axis
+            for axis, placement in enumerate(self.main_weight.placements)
+            if isinstance(placement, Replicate)
+        ]
+        if deferred_axes and self._symm_mem_pool is not None:
+            raise NotImplementedError(
+                "Symmetric-memory gradient reduction does not support deferred DP-outer axes."
+            )
+        accumulation_placements = tuple(
+            Partial(partial_op) if isinstance(placement, Replicate) else placement
+            for placement in self.main_weight.placements
         )
-        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
-        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)
-            if grad_divisor != 1:
-                self.main_grad.local_buffer.div_(grad_divisor)
+
+        # A finalized main_grad (placements == main_weight) means the previous step
+        # is done; reset it to the DP-outer-Partial accumulation placement.
+        if self.main_grad.placements == self.main_weight.placements:
+            self.main_grad.placements = accumulation_placements
+        reduced = self._reduce_grad_axis(partial_grad, accumulation_placements)
+
+        # set_to_none=True cleared the sharded grads -> overwrite main_grad and
+        # install .grad; else they are views into a zeroed main_grad -> accumulate.
+        if has_grad(self.sharded_parameters):
+            self.main_grad.local_buffer.add_(reduced.local_buffer)
         else:
-            reduced_grad = partial_grad.redistribute(self.main_grad.placements)
-            if grad_divisor != 1:
-                reduced_grad.local_buffer.div_(grad_divisor)
-            if has_sharded_grads:
-                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.main_grad.local_buffer.copy_(reduced.local_buffer)
+            for index, sharded_parameter in enumerate(self.sharded_parameters):
+                sharded_parameter.grad = self.main_grad.get_dtensor(index)
+
+        if is_last_microbatch:
+            for axis in deferred_axes:
+                dist.all_reduce(
+                    self.main_grad.local_buffer,
+                    op=dist.ReduceOp.AVG,
+                    group=self.mesh.get_group(axis),
+                )
+            self.main_grad.placements = self.main_weight.placements
+            for index, sharded_parameter in enumerate(self.sharded_parameters):
+                sharded_parameter.grad = self.main_grad.get_dtensor(index)
 
         for parameter in self.unsharded_parameters:
             parameter.grad = None

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"
 
@@ -168,6 +183,77 @@ 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)
+    baseline = TinyModel().to(device)
+    model = TinyModel().to(device)
+    model.load_state_dict(baseline.state_dict())
+
+    fully_shard(model.fc1, mesh=mesh, placements=_hsdp_placements())
+    fully_shard(model.fc2, 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, 8, device=device)
+    target = torch.randn(num_microbatches, micro_batch_size, 4, 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_nested_fully_shard_excludes_child_owned_parameters(distributed_setup):
     """An outer FSDP unit owns direct parameters but not nested child-unit parameters."""
     world_size = distributed_setup.world_size