From 60956871ecabf695043316cdcd89d27b6f4e8f7b Mon Sep 17 00:00:00 2001
From: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
Date: Fri, 10 Jul 2026 19:02:45 +0400
Subject: [PATCH 01/13] Support HSDP grad reduce in experimental mFSDP
 fully_shard

Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
---
 .../src/megatron_fsdp/experimental/module.py  |   2 +-
 .../experimental/parameter_group.py           | 136 +++++++++++++++---
 .../distributed/mfsdp_v2/test_fully_shard.py  |  83 +++++++++++
 3 files changed, 202 insertions(+), 19 deletions(-)

diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py
index 3c97fda3242..90b8afb0830 100644
--- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py
+++ b/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():
diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py
index eeec848416b..115e4263997 100644
--- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py
+++ b/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,22 @@ 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:
+            if not _grad_placements_reduce_to_weight(
+                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. "
+                    "FSDP requires main_grad placements to equal main_weight placements, or "
+                    "to defer a reduction by being Partial on an axis where main_weight is "
+                    "Replicate (HSDP DP-outer accumulation). "
                     f"Got main_grad placements {self.main_grad.placements} and "
                     f"main_weight placements {self.main_weight.placements}."
                 )
+            # When main_grad placements differ from main_weight, main_grad rests at a
+            # DP-outer Partial accumulation state: each backward reduce-scatters DP-inner
+            # into it, and the deferred DP-outer reduction runs on the last microbatch.
+            self._defers_grad_reduction = self.main_grad.placements != self.main_weight.placements
+            self._grad_accumulating = False
+            self._reduced_grad: DBuffer | None = None
         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 +254,82 @@ 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_partial_grad(
+        self, partial_grad: DBuffer, partial_op: dist.ReduceOp.RedOpType, *, out: "DBuffer | None"
+    ) -> DBuffer:
+        """Reduce an all-Partial gradient buffer to ``main_grad``'s placements.
+
+        ``DBuffer.redistribute`` changes one mesh axis per call, so a 2-D DP mesh
+        (HSDP/HFSDP) composes the reduction axis by axis. Innermost axes are
+        reduced first: each reduce-scatter shrinks the buffer before the next
+        collective, and every intermediate placement keeps ``Flat`` a suffix
+        (``[Partial, Flat]`` is valid, ``[Flat, Partial]`` is not). The outermost
+        changed axis is reduced last so it can write directly into ``out``.
+        """
+        target = self.main_grad.placements
+        changed_axes = [
+            axis
+            for axis in range(self.mesh.ndim)
+            if partial_grad.placements[axis] != target[axis]
+        ]
+        if not changed_axes:
+            raise RuntimeError("FSDP gradient reduction requires a changed placement axis.")
+        if self._symm_mem_pool is not None and len(changed_axes) > 1:
+            raise NotImplementedError(
+                "Symmetric-memory gradient reduction supports a single mesh axis; "
+                f"got changed axes {changed_axes}."
+            )
+
+        current = partial_grad
+        for axis in reversed(changed_axes):
+            placements = list(current.placements)
+            placements[axis] = target[axis]
+            if self._symm_mem_pool is not None:
+                current.rendezvous(axis)
+            is_outermost_change = axis == changed_axes[0]
+            current = current.redistribute(placements, out=out if is_outermost_change else None)
+            if partial_op == dist.ReduceOp.SUM:
+                current.local_buffer.div_(self.mesh.size(axis))
+        return current
+
+    def _accumulate_deferred_grad(
+        self, partial_grad: DBuffer, partial_op: dist.ReduceOp.RedOpType, is_last_microbatch: bool
+    ) -> None:
+        """Accumulate DP-inner-reduced grads, deferring the DP-outer reduction.
+
+        main_grad rests at its DP-outer-Partial placement. Each backward
+        reduce-scatters DP-inner into main_grad (resetting it on the first
+        microbatch, accumulating afterwards). Only the last microbatch reduces
+        DP-outer into ``main_weight``'s placement and installs the resulting
+        sharded parameter gradients.
+        """
+        assert self.main_grad is not None
+        inner_reduced = self._reduce_partial_grad(partial_grad, partial_op, out=None)
+        if self._grad_accumulating:
+            self.main_grad.local_buffer.add_(inner_reduced.local_buffer)
+        else:
+            self.main_grad.local_buffer.copy_(inner_reduced.local_buffer)
+            self._grad_accumulating = True
+
+        if not is_last_microbatch:
+            return
+
+        # Reduce DP-outer into the optimizer's placement; keep the buffer alive for
+        # optimizer.step() through both this reference and the sharded grad DTensors.
+        self._reduced_grad = self.main_grad.redistribute(self.main_weight.placements)
+        for index, sharded_parameter in enumerate(self.sharded_parameters):
+            sharded_parameter.grad = self._reduced_grad.get_dtensor(index)
+        self._grad_accumulating = False
+
+    def reduce_gradients(self, is_last_microbatch: bool = True) -> None:
+        """Reduce full local gradients into sharded parameter gradients.
+
+        For a plain DP mesh (or matching gradient/optimizer placements) the
+        reduction runs fully every backward. When main_grad defers a DP-outer
+        reduction (HSDP: main_grad Partial where main_weight is Replicate), each
+        backward reduce-scatters DP-inner and accumulates into main_grad, and only
+        the last microbatch reduces DP-outer and installs the sharded gradients.
+        """
         assert self.main_grad is not None
 
         def has_grad(parameters: Iterable[nn.Parameter]) -> bool:
@@ -275,6 +358,12 @@ def has_grad(parameters: Iterable[nn.Parameter]) -> bool:
                 grads, mesh=self.mesh, placements=[Partial(partial_op)] * self.mesh.ndim
             )
 
+        if self._defers_grad_reduction:
+            self._accumulate_deferred_grad(partial_grad, partial_op, is_last_microbatch)
+            for parameter in self.unsharded_parameters:
+                parameter.grad = None
+            return
+
         # 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.
@@ -282,20 +371,10 @@ def has_grad(parameters: Iterable[nn.Parameter]) -> bool:
         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
-        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)
+            self._reduce_partial_grad(partial_grad, partial_op, out=self.main_grad)
         else:
-            reduced_grad = partial_grad.redistribute(self.main_grad.placements)
-            if grad_divisor != 1:
-                reduced_grad.local_buffer.div_(grad_divisor)
+            reduced_grad = self._reduce_partial_grad(partial_grad, partial_op, out=None)
             if has_sharded_grads:
                 self.main_grad.local_buffer.add_(reduced_grad.local_buffer)
             else:
@@ -314,3 +393,24 @@ def _get_parameter_owner(module: nn.Module, name: str) -> tuple[nn.Module, str]:
     module_name, separator, parameter_name = name.rpartition(".")
     owner = module.get_submodule(module_name) if separator else module
     return owner, parameter_name
+
+
+def _grad_placements_reduce_to_weight(
+    grad_placements: tuple[Placement, ...], weight_placements: tuple[Placement, ...]
+) -> bool:
+    """Return whether main_grad placements reduce to main_weight placements.
+
+    They may be equal, or main_grad may be ``Partial`` on an axis where
+    main_weight is ``Replicate`` (an HSDP DP-outer reduction deferred to the last
+    microbatch). All other axes must match so the buffers share a layout and
+    local size.
+    """
+    if len(grad_placements) != len(weight_placements):
+        return False
+    for grad_placement, weight_placement in zip(grad_placements, weight_placements):
+        if grad_placement == weight_placement:
+            continue
+        if isinstance(grad_placement, Partial) and isinstance(weight_placement, Replicate):
+            continue
+        return False
+    return True
diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py
index 229cd0bff4b..db76204760e 100644
--- a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py
+++ b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py
@@ -6,6 +6,7 @@
 
 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
@@ -13,7 +14,9 @@
 
 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); gradients rest DP-outer-Partial and reduce 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,74 @@ def train(model, optimizer, log_prefix) -> list[torch.Tensor]:
     )
 
 
+@pytest.mark.parametrize("num_microbatches", [1, 3])
+def test_hsdp_losses_match_baseline(distributed_setup, num_microbatches):
+    """HSDP (DP-outer replicated, DP-inner sharded) training should match single-rank SGD.
+
+    Gradients reduce-scatter within DP-inner every backward and accumulate at a
+    DP-outer-Partial resting state; the DP-outer reduction 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.
+    """
+    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()
+
+            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

From 724868598161718b5ddff579bfa58a7a73fda91f Mon Sep 17 00:00:00 2001
From: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
Date: Fri, 10 Jul 2026 20:11:30 +0400
Subject: [PATCH 02/13] Restrict deferred-grad guard to one axis and release
 stale reduced-grad buffer

Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
---
 .../megatron_fsdp/experimental/parameter_group.py  | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py
index 115e4263997..2ef8697083c 100644
--- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py
+++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py
@@ -308,6 +308,10 @@ def _accumulate_deferred_grad(
         if self._grad_accumulating:
             self.main_grad.local_buffer.add_(inner_reduced.local_buffer)
         else:
+            # First microbatch of a step: reset the accumulator and release the
+            # previous step's reduced-grad buffer, whose sharded grads have been
+            # cleared by optimizer.zero_grad().
+            self._reduced_grad = None
             self.main_grad.local_buffer.copy_(inner_reduced.local_buffer)
             self._grad_accumulating = True
 
@@ -400,17 +404,21 @@ def _grad_placements_reduce_to_weight(
 ) -> bool:
     """Return whether main_grad placements reduce to main_weight placements.
 
-    They may be equal, or main_grad may be ``Partial`` on an axis where
+    They may be equal, or main_grad may be ``Partial`` on a single axis where
     main_weight is ``Replicate`` (an HSDP DP-outer reduction deferred to the last
     microbatch). All other axes must match so the buffers share a layout and
-    local size.
+    local size. At most one deferred axis is allowed because the deferred
+    reduction finalizes with a single-axis ``DBuffer.redistribute``; multi-axis
+    deferral is not yet supported.
     """
     if len(grad_placements) != len(weight_placements):
         return False
+    deferred_axes = 0
     for grad_placement, weight_placement in zip(grad_placements, weight_placements):
         if grad_placement == weight_placement:
             continue
         if isinstance(grad_placement, Partial) and isinstance(weight_placement, Replicate):
+            deferred_axes += 1
             continue
         return False
-    return True
+    return deferred_axes <= 1

From 3d271f3f852e33ec682ea46d4f17b3204e802627 Mon Sep 17 00:00:00 2001
From: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
Date: Sat, 11 Jul 2026 10:33:54 +0400
Subject: [PATCH 03/13] Remove over-defensive main_grad/main_weight placement
 guard

Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
---
 .../experimental/parameter_group.py           | 38 +------------------
 1 file changed, 2 insertions(+), 36 deletions(-)

diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py
index 2ef8697083c..117d446c464 100644
--- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py
+++ b/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, Placement, Placements, Replicate, changed_mesh_axis
+from .placement import Partial, Placements, Replicate, changed_mesh_axis
 
 _CONTAINING_PARAMETER_GROUP_ATTR = "_mfsdp_parameter_group"
 
@@ -151,19 +151,10 @@ 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 not _grad_placements_reduce_to_weight(
-                self.main_grad.placements, self.main_weight.placements
-            ):
-                raise ValueError(
-                    "FSDP requires main_grad placements to equal main_weight placements, or "
-                    "to defer a reduction by being Partial on an axis where main_weight is "
-                    "Replicate (HSDP DP-outer accumulation). "
-                    f"Got main_grad placements {self.main_grad.placements} and "
-                    f"main_weight placements {self.main_weight.placements}."
-                )
             # When main_grad placements differ from main_weight, main_grad rests at a
             # DP-outer Partial accumulation state: each backward reduce-scatters DP-inner
             # into it, and the deferred DP-outer reduction runs on the last microbatch.
+            # Unsupported placement combinations surface from DBuffer.redistribute.
             self._defers_grad_reduction = self.main_grad.placements != self.main_weight.placements
             self._grad_accumulating = False
             self._reduced_grad: DBuffer | None = None
@@ -397,28 +388,3 @@ def _get_parameter_owner(module: nn.Module, name: str) -> tuple[nn.Module, str]:
     module_name, separator, parameter_name = name.rpartition(".")
     owner = module.get_submodule(module_name) if separator else module
     return owner, parameter_name
-
-
-def _grad_placements_reduce_to_weight(
-    grad_placements: tuple[Placement, ...], weight_placements: tuple[Placement, ...]
-) -> bool:
-    """Return whether main_grad placements reduce to main_weight placements.
-
-    They may be equal, or main_grad may be ``Partial`` on a single axis where
-    main_weight is ``Replicate`` (an HSDP DP-outer reduction deferred to the last
-    microbatch). All other axes must match so the buffers share a layout and
-    local size. At most one deferred axis is allowed because the deferred
-    reduction finalizes with a single-axis ``DBuffer.redistribute``; multi-axis
-    deferral is not yet supported.
-    """
-    if len(grad_placements) != len(weight_placements):
-        return False
-    deferred_axes = 0
-    for grad_placement, weight_placement in zip(grad_placements, weight_placements):
-        if grad_placement == weight_placement:
-            continue
-        if isinstance(grad_placement, Partial) and isinstance(weight_placement, Replicate):
-            deferred_axes += 1
-            continue
-        return False
-    return deferred_axes <= 1

From 4ed16aff00157c2ea61b6d7d5ad4ae7cc4577cac Mon Sep 17 00:00:00 2001
From: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
Date: Mon, 13 Jul 2026 16:00:47 +0400
Subject: [PATCH 04/13] Make deferred DP-outer gradient reduction stateless

Rest main_grad at main_weight placements so it backs .grad and reuses the zero_grad/set_to_none reset; reduce-scatter the sharded axes every backward and all-reduce the replicated DP-outer axis in place only on the last microbatch. Drop the per-group state and multi-axis reduction helper; cover both set_to_none modes in the HSDP parity test.

Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
---
 .../experimental/parameter_group.py           | 161 +++++++-----------
 .../distributed/mfsdp_v2/test_fully_shard.py  |  25 +--
 2 files changed, 70 insertions(+), 116 deletions(-)

diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py
index 117d446c464..c55fcb140f6 100644
--- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py
+++ b/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."
             )
-            # When main_grad placements differ from main_weight, main_grad rests at a
-            # DP-outer Partial accumulation state: each backward reduce-scatters DP-inner
-            # into it, and the deferred DP-outer reduction runs on the last microbatch.
-            # Unsupported placement combinations surface from DBuffer.redistribute.
-            self._defers_grad_reduction = self.main_grad.placements != self.main_weight.placements
-            self._grad_accumulating = False
-            self._reduced_grad: DBuffer | None = None
         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,85 +238,37 @@ def release_unsharded_storage(self) -> None:
         # so keep the shared storage-release path.
         self._unsharded_model_weight.release_storage()
 
-    def _reduce_partial_grad(
-        self, partial_grad: DBuffer, partial_op: dist.ReduceOp.RedOpType, *, out: "DBuffer | None"
+    def _reduce_grad_axis(
+        self, partial_grad: DBuffer, target_placements: tuple[Placement, ...]
     ) -> DBuffer:
-        """Reduce an all-Partial gradient buffer to ``main_grad``'s placements.
-
-        ``DBuffer.redistribute`` changes one mesh axis per call, so a 2-D DP mesh
-        (HSDP/HFSDP) composes the reduction axis by axis. Innermost axes are
-        reduced first: each reduce-scatter shrinks the buffer before the next
-        collective, and every intermediate placement keeps ``Flat`` a suffix
-        (``[Partial, Flat]`` is valid, ``[Flat, Partial]`` is not). The outermost
-        changed axis is reduced last so it can write directly into ``out``.
-        """
-        target = self.main_grad.placements
-        changed_axes = [
-            axis
-            for axis in range(self.mesh.ndim)
-            if partial_grad.placements[axis] != target[axis]
-        ]
-        if not changed_axes:
-            raise RuntimeError("FSDP gradient reduction requires a changed placement axis.")
-        if self._symm_mem_pool is not None and len(changed_axes) > 1:
-            raise NotImplementedError(
-                "Symmetric-memory gradient reduction supports a single mesh axis; "
-                f"got changed axes {changed_axes}."
-            )
+        """Reduce ``partial_grad`` to ``target_placements`` over its one changed axis.
 
-        current = partial_grad
-        for axis in reversed(changed_axes):
-            placements = list(current.placements)
-            placements[axis] = target[axis]
-            if self._symm_mem_pool is not None:
-                current.rendezvous(axis)
-            is_outermost_change = axis == changed_axes[0]
-            current = current.redistribute(placements, out=out if is_outermost_change else None)
-            if partial_op == dist.ReduceOp.SUM:
-                current.local_buffer.div_(self.mesh.size(axis))
-        return current
-
-    def _accumulate_deferred_grad(
-        self, partial_grad: DBuffer, partial_op: dist.ReduceOp.RedOpType, is_last_microbatch: bool
-    ) -> None:
-        """Accumulate DP-inner-reduced grads, deferring the DP-outer reduction.
-
-        main_grad rests at its DP-outer-Partial placement. Each backward
-        reduce-scatters DP-inner into main_grad (resetting it on the first
-        microbatch, accumulating afterwards). Only the last microbatch reduces
-        DP-outer into ``main_weight``'s placement and installs the resulting
-        sharded parameter gradients.
+        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.
         """
-        assert self.main_grad is not None
-        inner_reduced = self._reduce_partial_grad(partial_grad, partial_op, out=None)
-        if self._grad_accumulating:
-            self.main_grad.local_buffer.add_(inner_reduced.local_buffer)
-        else:
-            # First microbatch of a step: reset the accumulator and release the
-            # previous step's reduced-grad buffer, whose sharded grads have been
-            # cleared by optimizer.zero_grad().
-            self._reduced_grad = None
-            self.main_grad.local_buffer.copy_(inner_reduced.local_buffer)
-            self._grad_accumulating = True
-
-        if not is_last_microbatch:
-            return
-
-        # Reduce DP-outer into the optimizer's placement; keep the buffer alive for
-        # optimizer.step() through both this reference and the sharded grad DTensors.
-        self._reduced_grad = self.main_grad.redistribute(self.main_weight.placements)
-        for index, sharded_parameter in enumerate(self.sharded_parameters):
-            sharded_parameter.grad = self._reduced_grad.get_dtensor(index)
-        self._grad_accumulating = False
+        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.
 
-        For a plain DP mesh (or matching gradient/optimizer placements) the
-        reduction runs fully every backward. When main_grad defers a DP-outer
-        reduction (HSDP: main_grad Partial where main_weight is Replicate), each
-        backward reduce-scatters DP-inner and accumulates into main_grad, and only
-        the last microbatch reduces DP-outer and installs the sharded gradients.
+        Sharded (Flat) mesh axes are reduce-scattered into main_grad every
+        backward and accumulated through the standard zero_grad contract:
+        ``set_to_none=True`` clears the sharded grads so main_grad is overwritten,
+        while ``set_to_none=False`` keeps them as views into a zeroed main_grad so
+        they accumulate. Replicated (DP-outer) mesh axes are all-reduced only on
+        the last microbatch, in place, so ``.grad`` -- a view into main_grad --
+        becomes the fully reduced gradient before ``optimizer.step()``. With every
+        axis Flat (plain DP) this is the previous every-backward behavior.
         """
         assert self.main_grad is not None
 
@@ -346,38 +291,46 @@ 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
             )
 
-        if self._defers_grad_reduction:
-            self._accumulate_deferred_grad(partial_grad, partial_op, is_last_microbatch)
-            for parameter in self.unsharded_parameters:
-                parameter.grad = None
-            return
-
-        # 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
+        # Reduce-scatter the sharded (Flat) axes now; keep replicated (DP-outer)
+        # axes Partial so their all-reduce can be deferred to the last microbatch.
+        deferred_axes = [
+            axis
+            for axis, placement in enumerate(self.main_grad.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_grad.placements
         )
-        if can_reduce_into_main_grad:
-            self._reduce_partial_grad(partial_grad, partial_op, out=self.main_grad)
+        reduced = self._reduce_grad_axis(partial_grad, accumulation_placements)
+
+        # set_to_none=True cleared the sharded grads -> overwrite main_grad; else
+        # they are views into a zeroed main_grad -> accumulate into it.
+        if has_grad(self.sharded_parameters):
+            self.main_grad.local_buffer.add_(reduced.local_buffer)
         else:
-            reduced_grad = self._reduce_partial_grad(partial_grad, partial_op, out=None)
-            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),
+                )
 
         for parameter in self.unsharded_parameters:
             parameter.grad = None
diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py
index db76204760e..587cbe2f611 100644
--- a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py
+++ b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py
@@ -6,7 +6,6 @@
 
 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
@@ -14,7 +13,6 @@
 
 from megatron.core.distributed.fsdp.src.megatron_fsdp.experimental import (
     Flat,
-    Partial,
     Placements,
     Replicate,
     fully_shard,
@@ -104,13 +102,13 @@ def _flat_placements() -> Placements:
 
 
 def _hsdp_placements() -> Placements:
-    """HSDP: params/optimizer replicated across DP-outer (axis 0), sharded within
-    DP-inner (axis 1); gradients rest DP-outer-Partial and reduce on the last
-    microbatch."""
+    """HSDP: params/gradients/optimizer replicated across DP-outer (axis 0),
+    sharded within DP-inner (axis 1). main_grad rests DP-outer-replicated and
+    backs .grad; the DP-outer all-reduce is deferred to the last microbatch."""
     return Placements(
         dp_axes=[0, 1],
         parameter=[Replicate(), Flat()],
-        gradient=[Partial(dist.ReduceOp.AVG), Flat()],
+        gradient=[Replicate(), Flat()],
         optimizer=[Replicate(), Flat()],
     )
 
@@ -183,14 +181,17 @@ 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):
+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 at a
-    DP-outer-Partial resting state; the DP-outer reduction 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.
+    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
@@ -221,7 +222,7 @@ def test_hsdp_losses_match_baseline(distributed_setup, num_microbatches):
     def train(model, optimizer, log_prefix) -> list[torch.Tensor]:
         losses = []
         for step in range(5):
-            optimizer.zero_grad()
+            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

From dac14f2283dfa6cf3e54ea1889dc221b3925982f Mon Sep 17 00:00:00 2001
From: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
Date: Tue, 14 Jul 2026 14:27:39 +0400
Subject: [PATCH 05/13] Rest main_grad DP-outer-Partial between microbatches
 for HSDP

main_grad rests at its DP-outer-Partial gradient placement between microbatches and is all-reduced to main_weight's placements on the last microbatch, per HSDP semantics. Extend DBuffer.get_dtensor to represent Partial so main_grad backs .grad throughout, reusing the existing zero_grad/set_to_none accumulation. Cover both set_to_none modes in the HSDP parity test.

Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
---
 .../src/megatron_fsdp/experimental/dbuffer.py |  5 ++-
 .../experimental/parameter_group.py           | 39 ++++++++++++-------
 .../distributed/mfsdp_v2/test_fully_shard.py  | 10 +++--
 3 files changed, 35 insertions(+), 19 deletions(-)

diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py
index 9b6e6dc44c3..9ac00cbe5a7 100644
--- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py
+++ b/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}.")
 
diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py
index c55fcb140f6..19bd05ccc0e 100644
--- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py
+++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py
@@ -261,14 +261,17 @@ def _reduce_grad_axis(
     def reduce_gradients(self, is_last_microbatch: bool = True) -> None:
         """Reduce full local gradients into sharded parameter gradients.
 
-        Sharded (Flat) mesh axes are reduce-scattered into main_grad every
-        backward and accumulated through the standard zero_grad contract:
-        ``set_to_none=True`` clears the sharded grads so main_grad is overwritten,
-        while ``set_to_none=False`` keeps them as views into a zeroed main_grad so
-        they accumulate. Replicated (DP-outer) mesh axes are all-reduced only on
-        the last microbatch, in place, so ``.grad`` -- a view into main_grad --
-        becomes the fully reduced gradient before ``optimizer.step()``. With every
-        axis Flat (plain DP) this is the previous every-backward behavior.
+        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
 
@@ -298,11 +301,11 @@ def has_grad(parameters: Iterable[nn.Parameter]) -> bool:
                 grads, mesh=self.mesh, placements=[Partial(partial_op)] * self.mesh.ndim
             )
 
-        # Reduce-scatter the sharded (Flat) axes now; keep replicated (DP-outer)
-        # axes Partial so their all-reduce can be deferred to the last microbatch.
+        # 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_grad.placements)
+            for axis, placement in enumerate(self.main_weight.placements)
             if isinstance(placement, Replicate)
         ]
         if deferred_axes and self._symm_mem_pool is not None:
@@ -311,12 +314,17 @@ def has_grad(parameters: Iterable[nn.Parameter]) -> bool:
             )
         accumulation_placements = tuple(
             Partial(partial_op) if isinstance(placement, Replicate) else placement
-            for placement in self.main_grad.placements
+            for placement in self.main_weight.placements
         )
+
+        # 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; else
-        # they are views into a zeroed main_grad -> accumulate into it.
+        # 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:
@@ -331,6 +339,9 @@ def has_grad(parameters: Iterable[nn.Parameter]) -> bool:
                     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
diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py
index 587cbe2f611..bc2ba4faf64 100644
--- a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py
+++ b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py
@@ -6,6 +6,7 @@
 
 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
@@ -13,6 +14,7 @@
 
 from megatron.core.distributed.fsdp.src.megatron_fsdp.experimental import (
     Flat,
+    Partial,
     Placements,
     Replicate,
     fully_shard,
@@ -102,13 +104,13 @@ def _flat_placements() -> Placements:
 
 
 def _hsdp_placements() -> Placements:
-    """HSDP: params/gradients/optimizer replicated across DP-outer (axis 0),
-    sharded within DP-inner (axis 1). main_grad rests DP-outer-replicated and
-    backs .grad; the DP-outer all-reduce is deferred to the last microbatch."""
+    """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=[Replicate(), Flat()],
+        gradient=[Partial(dist.ReduceOp.AVG), Flat()],
         optimizer=[Replicate(), Flat()],
     )
 

From 6570526748474d69557498dea2c797c198e14e99 Mon Sep 17 00:00:00 2001
From: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
Date: Wed, 15 Jul 2026 19:42:35 +0400
Subject: [PATCH 06/13] Address review: store accumulation placement,
 redistribute finalize, harden get_dtensor reduce op

Store the initial gradient placement instead of recomputing it each backward; finalize main_grad via DBuffer.redistribute (covers the HFSDP reduce-scatter case and the buffer-size change, and uses the Partial's own reduce op); convert Partial reduce ops explicitly in get_dtensor, raising on unsupported ops instead of silently defaulting to sum.

Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
---
 .../src/megatron_fsdp/experimental/dbuffer.py |  9 ++++-
 .../experimental/parameter_group.py           | 35 ++++++++-----------
 2 files changed, 22 insertions(+), 22 deletions(-)

diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py
index 9ac00cbe5a7..66ab751c87b 100644
--- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py
+++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py
@@ -451,7 +451,14 @@ def get_dtensor(self, index: int) -> DTensor:
             elif isinstance(placement, Partial):
                 # 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"
+                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}.")
diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py
index 19bd05ccc0e..70ee35197ad 100644
--- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py
+++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py
@@ -151,6 +151,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."
             )
+            # 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
@@ -301,27 +304,21 @@ def has_grad(parameters: Iterable[nn.Parameter]) -> bool:
                 grads, mesh=self.mesh, placements=[Partial(partial_op)] * self.mesh.ndim
             )
 
-        # 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:
+        # main_grad accumulates DP-outer-Partial (Partial where main_weight is
+        # Replicate); symmetric memory does not support that deferral yet.
+        if (
+            self._accumulation_placements != self.main_weight.placements
+            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
-        )
 
         # 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)
+            self.main_grad.placements = self._accumulation_placements
+        reduced = self._reduce_grad_axis(partial_grad, self._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.
@@ -333,13 +330,9 @@ def has_grad(parameters: Iterable[nn.Parameter]) -> bool:
                 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
+            # 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)
             for index, sharded_parameter in enumerate(self.sharded_parameters):
                 sharded_parameter.grad = self.main_grad.get_dtensor(index)
 

From ba659e8631c570e2419ec195a5e3dc1e50fdc6f5 Mon Sep 17 00:00:00 2001
From: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
Date: Thu, 16 Jul 2026 08:38:19 +0400
Subject: [PATCH 07/13] Address review: drop symm-mem guard, DRY grad install,
 tidy reset

Remove the symmetric-memory deferred-axis guard (the DP-outer finalize is a plain redistribute); extract _install_sharded_grads; and reset main_grad to the accumulation placement only when it differs (the first microbatch), refreshing the sharded grads if they are set.

Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
---
 .../experimental/parameter_group.py           | 33 +++++++++----------
 1 file changed, 16 insertions(+), 17 deletions(-)

diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py
index 70ee35197ad..7ddc6cf025e 100644
--- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py
+++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py
@@ -261,6 +261,12 @@ def _reduce_grad_axis(
             reduced.local_buffer.div_(self.mesh.size(axis))
         return reduced
 
+    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 reduce_gradients(self, is_last_microbatch: bool = True) -> None:
         """Reduce full local gradients into sharded parameter gradients.
 
@@ -304,20 +310,15 @@ def has_grad(parameters: Iterable[nn.Parameter]) -> bool:
                 grads, mesh=self.mesh, placements=[Partial(partial_op)] * self.mesh.ndim
             )
 
-        # main_grad accumulates DP-outer-Partial (Partial where main_weight is
-        # Replicate); symmetric memory does not support that deferral yet.
-        if (
-            self._accumulation_placements != self.main_weight.placements
-            and self._symm_mem_pool is not None
-        ):
-            raise NotImplementedError(
-                "Symmetric-memory gradient reduction does not support deferred DP-outer axes."
-            )
-
-        # 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:
+        # A non-accumulation main_grad means the previous step finalized it; this
+        # only happens on the first microbatch. Reset it to the DP-outer-Partial
+        # accumulation placement -- a metadata relabel for HSDP since the buffer is
+        # overwritten below. (HFSDP will instead need a fresh buffer here, since its
+        # finalized buffer was reduce-scattered and is smaller.)
+        if self.main_grad.placements != self._accumulation_placements:
             self.main_grad.placements = self._accumulation_placements
+            if has_grad(self.sharded_parameters):
+                self._install_sharded_grads()
         reduced = self._reduce_grad_axis(partial_grad, self._accumulation_placements)
 
         # set_to_none=True cleared the sharded grads -> overwrite main_grad and
@@ -326,15 +327,13 @@ def has_grad(parameters: Iterable[nn.Parameter]) -> bool:
             self.main_grad.local_buffer.add_(reduced.local_buffer)
         else:
             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)
+            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)
-            for index, sharded_parameter in enumerate(self.sharded_parameters):
-                sharded_parameter.grad = self.main_grad.get_dtensor(index)
+            self._install_sharded_grads()
 
         for parameter in self.unsharded_parameters:
             parameter.grad = None

From bd293b3464a0fb12522d85b60db7561d274a560d Mon Sep 17 00:00:00 2001
From: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
Date: Thu, 16 Jul 2026 13:27:19 +0400
Subject: [PATCH 08/13] Address review: Replicate->Partial redistribute,
 restore reduce-into-main_grad

Add Replicate->Partial as the fifth DBuffer.redistribute transition and use it to reset main_grad instead of assigning to .placements; reduce straight into main_grad on the overwrite path via out=; rename _reduce_grad_axis to _reduce_partial_gradients.

Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
---
 .../src/megatron_fsdp/experimental/dbuffer.py | 37 +++++++++++++++-
 .../experimental/parameter_group.py           | 43 +++++++++++--------
 2 files changed, 60 insertions(+), 20 deletions(-)

diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py
index 66ab751c87b..df8cd870c4a 100644
--- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py
+++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py
@@ -289,8 +289,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:
@@ -319,6 +320,8 @@ 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):
+            return self.replicate_to_partial(axis, new_placement, out=out)
         raise NotImplementedError(
             "Unsupported DBuffer placement transition on axis "
             f"{axis}: {old_placement!r} -> {new_placement!r}."
@@ -417,6 +420,36 @@ def scatter(
         out.local_buffer.copy_(local_slice)
         return out
 
+    def replicate_to_partial(
+        self, mesh_axis: int, new_placement: Placement, *, out: "DBuffer | None" = None
+    ) -> "DBuffer":
+        """Relabel a Replicate axis as Partial (local rescale, no communication).
+
+        A Replicate value equals the AVG reduction of identical per-rank locals,
+        so Replicate -> Partial(AVG) is a pure relabel. SUM reduces those locals
+        to ``value * axis_size``, so scale by ``1 / axis_size`` to preserve it.
+        """
+        axis = mesh_axis
+        if not isinstance(new_placement, Partial):
+            raise ValueError(f"replicate_to_partial() requires a Partial target on axis {axis!r}.")
+        if not isinstance(self.placements[axis], Replicate):
+            raise ValueError(
+                f"replicate_to_partial() requires Replicate placement on axis {axis!r}."
+            )
+
+        placements = list(self.placements)
+        placements[axis] = new_placement
+        _validate_placements(placements)
+        out = self._create_or_validate_out(placements, out)
+        out.local_buffer.copy_(self.local_buffer)
+        if new_placement.reduce_op == dist.ReduceOp.SUM:
+            out.local_buffer.div_(self.mesh.size(axis))
+        elif new_placement.reduce_op != dist.ReduceOp.AVG:
+            raise NotImplementedError(
+                f"replicate_to_partial() supports SUM and AVG, got {new_placement.reduce_op!r}."
+            )
+        return out
+
     def get_local_tensor(self, index: int) -> torch.Tensor:
         """Return this rank's local view for logical tensor ``index``.
 
diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py
index 7ddc6cf025e..8b198226282 100644
--- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py
+++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py
@@ -241,21 +241,24 @@ def release_unsharded_storage(self) -> None:
         # so keep the shared storage-release path.
         self._unsharded_model_weight.release_storage()
 
-    def _reduce_grad_axis(
-        self, partial_grad: DBuffer, target_placements: tuple[Placement, ...]
+    def _reduce_partial_gradients(
+        self,
+        partial_grad: DBuffer,
+        target_placements: tuple[Placement, ...],
+        *,
+        out: "DBuffer | None" = None,
     ) -> DBuffer:
-        """Reduce ``partial_grad`` to ``target_placements`` over its one changed axis.
+        """Reduce ``partial_grad`` to ``target_placements``.
 
         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.
+        and its sharded resting placement.
         """
         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)
+        reduced = partial_grad.redistribute(target_placements, out=out)
         # 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))
@@ -275,7 +278,7 @@ def reduce_gradients(self, is_last_microbatch: bool = True) -> None:
         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
+        microbatch reduces the DP-outer axes, 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
@@ -303,7 +306,7 @@ 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 in _reduce_grad_axis.
+        # Preserve AVG semantics by reducing SUM and scaling the output in _reduce_partial_gradients.
         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(
@@ -311,22 +314,26 @@ def has_grad(parameters: Iterable[nn.Parameter]) -> bool:
             )
 
         # A non-accumulation main_grad means the previous step finalized it; this
-        # only happens on the first microbatch. Reset it to the DP-outer-Partial
-        # accumulation placement -- a metadata relabel for HSDP since the buffer is
-        # overwritten below. (HFSDP will instead need a fresh buffer here, since its
-        # finalized buffer was reduce-scattered and is smaller.)
+        # only happens on the first microbatch. Redistribute it back to the
+        # DP-outer-Partial accumulation placement -- a local 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.placements = self._accumulation_placements
+            self.main_grad = self.main_grad.redistribute(self._accumulation_placements)
             if has_grad(self.sharded_parameters):
                 self._install_sharded_grads()
-        reduced = self._reduce_grad_axis(partial_grad, self._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.
+        # set_to_none=True cleared the sharded grads -> reduce straight into
+        # main_grad and install .grad; else they are views into main_grad, so
+        # reduce into a temporary and accumulate.
         if has_grad(self.sharded_parameters):
-            self.main_grad.local_buffer.add_(reduced.local_buffer)
+            reduced_grad = self._reduce_partial_gradients(
+                partial_grad, self._accumulation_placements
+            )
+            self.main_grad.local_buffer.add_(reduced_grad.local_buffer)
         else:
-            self.main_grad.local_buffer.copy_(reduced.local_buffer)
+            self._reduce_partial_gradients(
+                partial_grad, self._accumulation_placements, out=self.main_grad
+            )
             self._install_sharded_grads()
 
         if is_last_microbatch:

From 325f338f3c3f944acaadc665422470059bdf6c4b Mon Sep 17 00:00:00 2001
From: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
Date: Thu, 16 Jul 2026 23:15:32 +0400
Subject: [PATCH 09/13] Address review: zero-copy replicate_to_partial, dedup,
 DP-outer event test

Make replicate_to_partial a from_local relabel instead of an allocate-and-copy; hoist has_sharded_grads above the reset; add test_hsdp_defers_dp_outer_allreduce_to_last_microbatch counting NCCL events to verify the DP-outer all-reduce only fires on the last microbatch.

Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
---
 .../src/megatron_fsdp/experimental/dbuffer.py | 29 ++++----
 .../experimental/parameter_group.py           | 15 ++--
 .../distributed/mfsdp_v2/test_fully_shard.py  | 69 +++++++++++++++++++
 3 files changed, 92 insertions(+), 21 deletions(-)

diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py
index df8cd870c4a..9a9e27a4fab 100644
--- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py
+++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py
@@ -423,11 +423,13 @@ def scatter(
     def replicate_to_partial(
         self, mesh_axis: int, new_placement: Placement, *, out: "DBuffer | None" = None
     ) -> "DBuffer":
-        """Relabel a Replicate axis as Partial (local rescale, no communication).
+        """Reinterpret a Replicate axis as Partial without moving data.
 
-        A Replicate value equals the AVG reduction of identical per-rank locals,
-        so Replicate -> Partial(AVG) is a pure relabel. SUM reduces those locals
-        to ``value * axis_size``, so scale by ``1 / axis_size`` to preserve it.
+        Replicate and Partial share the same local layout, so this relabels the
+        existing local buffer rather than copying it, and issues no communication.
+        It is value-preserving for AVG -- the mean of identical per-rank locals is
+        that value. SUM would need a ``1 / axis_size`` rescale, so it is left
+        unsupported until a caller needs it.
         """
         axis = mesh_axis
         if not isinstance(new_placement, Partial):
@@ -436,19 +438,18 @@ def replicate_to_partial(
             raise ValueError(
                 f"replicate_to_partial() requires Replicate placement on axis {axis!r}."
             )
+        if new_placement.reduce_op != dist.ReduceOp.AVG:
+            raise NotImplementedError(
+                f"replicate_to_partial() supports AVG only, got {new_placement.reduce_op!r}."
+            )
+        if out is not None:
+            raise NotImplementedError("replicate_to_partial() does not support an out buffer.")
 
         placements = list(self.placements)
         placements[axis] = new_placement
-        _validate_placements(placements)
-        out = self._create_or_validate_out(placements, out)
-        out.local_buffer.copy_(self.local_buffer)
-        if new_placement.reduce_op == dist.ReduceOp.SUM:
-            out.local_buffer.div_(self.mesh.size(axis))
-        elif new_placement.reduce_op != dist.ReduceOp.AVG:
-            raise NotImplementedError(
-                f"replicate_to_partial() supports SUM and AVG, got {new_placement.reduce_op!r}."
-            )
-        return out
+        return DBuffer.from_local(
+            self.local_buffer, self.mesh, placements, self.layout.tensor_shapes
+        )
 
     def get_local_tensor(self, index: int) -> torch.Tensor:
         """Return this rank's local view for logical tensor ``index``.
diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py
index d3a8c568e31..5f2430d1a7f 100644
--- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py
+++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py
@@ -300,19 +300,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 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 local relabel for HSDP, and
-        # a fresh reduce-scattered buffer for HFSDP in the future.
+        # 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_grad(self.sharded_parameters):
+            if has_sharded_grads:
                 self._install_sharded_grads()
 
-        # 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)
         can_reduce_into_main_grad = (
             not has_sharded_grads and partial_grad.dtype == self.main_grad.dtype
         )
diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py
index 920f191ccfb..d4c91e275c0 100644
--- a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py
+++ b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py
@@ -254,6 +254,75 @@ def train(model, optimizer, log_prefix) -> list[torch.Tensor]:
     )
 
 
+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.
+
+    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)
+    model = TinyModel().to(device)
+    fully_shard(model.fc1, mesh=mesh, placements=_hsdp_placements())
+    fully_shard(model.fc2, 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, 8, device=device)
+    target = torch.randn(num_microbatches, micro_batch_size, 4, 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 = [
+        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())
+    ]
+    all_reduce_events = [
+        event
+        for event in cuda_events
+        if "nccl" in event.name.lower() and "allreduce" in event.name.lower()
+    ]
+    # HSDP must trigger the DP-outer all-reduce that plain DP never issues.
+    assert all_reduce_events, [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

From 8bb6e1850b18e1b48b715cc29579256d8b148115 Mon Sep 17 00:00:00 2001
From: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
Date: Thu, 16 Jul 2026 23:21:52 +0400
Subject: [PATCH 10/13] Exercise root-unit overlap in the HSDP event-count test

Use MultiChildModel with fully_shard(model) so the child units share a root context and reduce through the overlap path, addressing the earlier request to fully_shard(model) as well.

Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
---
 .../distributed/mfsdp_v2/test_fully_shard.py       | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py
index d4c91e275c0..ddc3945e5dc 100644
--- a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py
+++ b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py
@@ -257,6 +257,8 @@ def train(model, optimizer, log_prefix) -> list[torch.Tensor]:
 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
@@ -274,15 +276,17 @@ def test_hsdp_defers_dp_outer_allreduce_to_last_microbatch(distributed_setup):
         device.type, (outer_size, inner_size), mesh_dim_names=("dp_outer", "dp_inner")
     )
     torch.manual_seed(1234)
-    model = TinyModel().to(device)
-    fully_shard(model.fc1, mesh=mesh, placements=_hsdp_placements())
-    fully_shard(model.fc2, mesh=mesh, placements=_hsdp_placements())
+    dim = 8
+    model = MultiChildModel(dim=dim, num_children=2).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, 8, device=device)
-    target = torch.randn(num_microbatches, micro_batch_size, 4, device=device)
+    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:

From d0f8185e73bc2f0f8fe3f1a73d7cece0b158c4b9 Mon Sep 17 00:00:00 2001
From: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
Date: Thu, 16 Jul 2026 23:30:21 +0400
Subject: [PATCH 11/13] Exercise root-unit overlap in the HSDP parity test too

Switch test_hsdp_losses_match_baseline to MultiChildModel with fully_shard(model) so its child layers reduce under a shared root context, matching the fully_shard(model) request anchored on the parity test.

Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
---
 .../distributed/mfsdp_v2/test_fully_shard.py     | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py
index ddc3945e5dc..04e01accbfe 100644
--- a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py
+++ b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py
@@ -207,18 +207,22 @@ def test_hsdp_losses_match_baseline(distributed_setup, num_microbatches, set_to_
         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)
+    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())
 
-    fully_shard(model.fc1, mesh=mesh, placements=_hsdp_placements())
-    fully_shard(model.fc2, mesh=mesh, placements=_hsdp_placements())
+    # 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, 8, device=device)
-    target = torch.randn(num_microbatches, micro_batch_size, 4, device=device)
+    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]:

From 8b18bde0f2faa8d294df8afd08ace51ee92eaa2f Mon Sep 17 00:00:00 2001
From: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
Date: Fri, 17 Jul 2026 08:47:35 +0400
Subject: [PATCH 12/13] Address review: inline Replicate->Partial relabel,
 share event helper, stricter assert

Inline the Replicate->Partial relabel into DBuffer.redistribute and remove replicate_to_partial; extract _nccl_events and use it from both the event-count and overlap tests; assert the exact DP-outer all-reduce count (num_children + 1).

Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
---
 .../src/megatron_fsdp/experimental/dbuffer.py | 48 +++++++------------
 .../distributed/mfsdp_v2/test_fully_shard.py  | 44 ++++++++---------
 2 files changed, 35 insertions(+), 57 deletions(-)

diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py
index 9a9e27a4fab..8021d5619bb 100644
--- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py
+++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/dbuffer.py
@@ -321,7 +321,22 @@ def redistribute(
         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):
-            return self.replicate_to_partial(axis, new_placement, out=out)
+            # 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}."
@@ -420,37 +435,6 @@ def scatter(
         out.local_buffer.copy_(local_slice)
         return out
 
-    def replicate_to_partial(
-        self, mesh_axis: int, new_placement: Placement, *, out: "DBuffer | None" = None
-    ) -> "DBuffer":
-        """Reinterpret a Replicate axis as Partial without moving data.
-
-        Replicate and Partial share the same local layout, so this relabels the
-        existing local buffer rather than copying it, and issues no communication.
-        It is value-preserving for AVG -- the mean of identical per-rank locals is
-        that value. SUM would need a ``1 / axis_size`` rescale, so it is left
-        unsupported until a caller needs it.
-        """
-        axis = mesh_axis
-        if not isinstance(new_placement, Partial):
-            raise ValueError(f"replicate_to_partial() requires a Partial target on axis {axis!r}.")
-        if not isinstance(self.placements[axis], Replicate):
-            raise ValueError(
-                f"replicate_to_partial() requires Replicate placement on axis {axis!r}."
-            )
-        if new_placement.reduce_op != dist.ReduceOp.AVG:
-            raise NotImplementedError(
-                f"replicate_to_partial() supports AVG only, got {new_placement.reduce_op!r}."
-            )
-        if out is not None:
-            raise NotImplementedError("replicate_to_partial() does not support an out buffer.")
-
-        placements = list(self.placements)
-        placements[axis] = new_placement
-        return DBuffer.from_local(
-            self.local_buffer, self.mesh, placements, self.layout.tensor_shapes
-        )
-
     def get_local_tensor(self, index: int) -> torch.Tensor:
         """Return this rank's local view for logical tensor ``index``.
 
diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py
index 04e01accbfe..d978b04e945 100644
--- a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py
+++ b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py
@@ -126,6 +126,16 @@ 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 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."""
@@ -281,7 +291,8 @@ def test_hsdp_defers_dp_outer_allreduce_to_last_microbatch(distributed_setup):
     )
     torch.manual_seed(1234)
     dim = 8
-    model = MultiChildModel(dim=dim, num_children=2).to(device)
+    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())
@@ -310,19 +321,11 @@ def train_one_step() -> None:
         torch.cuda.synchronize(device)
 
     cuda_events = [event for event in prof.events() if event.device_type.name == "CUDA"]
-    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())
-    ]
-    all_reduce_events = [
-        event
-        for event in cuda_events
-        if "nccl" in event.name.lower() and "allreduce" in event.name.lower()
-    ]
-    # HSDP must trigger the DP-outer all-reduce that plain DP never issues.
-    assert all_reduce_events, [event.name for event in cuda_events]
+    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, (
@@ -510,17 +513,8 @@ def train_one_iteration() -> None:
         torch.cuda.synchronize(device)
 
     cuda_events = [event for event in prof.events() if event.device_type.name == "CUDA"]
-    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())
-    ]
+    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 = [

From 96764cddaaf9cd60ace7834c81bf246ff9fd4424 Mon Sep 17 00:00:00 2001
From: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
Date: Sat, 18 Jul 2026 01:14:36 +0400
Subject: [PATCH 13/13] Count only NCCL kernel events in _nccl_events

Filter on activity_type == 'kernel' so the profiler's gpu_user_annotation duplicates aren't double-counted, and format the boolean (per @wujingyue's suggestion).

Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
---
 tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py
index d978b04e945..4e4818beb2a 100644
--- a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py
+++ b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py
@@ -132,6 +132,7 @@ def _nccl_events(cuda_events, *name_fragments):
         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)
     ]