Support configurable parameter, gradient, and optimizer sharding in experimental FSDP v2

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

@@ -109,13 +109,13 @@ def step_post_hook(
             set_grad(parameter, original_grad)
         casted_grads.clear()
 
-        fsdp_parameter_groups: set[FsdpParameterGroup] = set()
+        fsdp_parameter_groups: dict[FsdpParameterGroup, None] = {}
         for optimizer_group in hooked_optimizer.param_groups:
             for parameter in optimizer_group["params"]:
                 parameter_group = get_containing_parameter_group(parameter)
                 if parameter_group is None:
                     continue
-                fsdp_parameter_groups.add(parameter_group)
+                fsdp_parameter_groups[parameter_group] = None
 
         for parameter_group in fsdp_parameter_groups:
             parameter_group.sync_model_weight_from_main_weight()

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

@@ -24,7 +24,7 @@
 
 from ..mixed_precision import MixedPrecisionPolicy
 from .dbuffer import DBuffer
-from .placement import Partial, Placements, Replicate, changed_mesh_axis
+from .placement import Flat, Partial, Placements, Replicate, changed_mesh_axis
 
 _CONTAINING_PARAMETER_GROUP_ATTR = "_mfsdp_parameter_group"
 
@@ -44,6 +44,7 @@ class FsdpParameterGroup:
     mesh: DeviceMesh
     dtype: torch.dtype
     requires_grad: bool
+    is_first_microbatch: bool
     main_weight: DBuffer
     model_weight: DBuffer
     main_grad: DBuffer | None
@@ -85,6 +86,7 @@ def __init__(
         first_parameter = next(iter(parameters.values()))
         self.dtype = first_parameter.dtype
         self.requires_grad = first_parameter.requires_grad
+        self.is_first_microbatch = True
         for name, parameter in parameters.items():
             if parameter.dtype != self.dtype:
                 raise ValueError(
@@ -146,13 +148,11 @@ def __init__(
                 dtype=grad_dtype,
                 device=self.main_weight.device,
             )
+            self.main_grad.local_buffer.zero_()
             assert self.main_grad.layout == self.main_weight.layout, (
                 "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
@@ -222,12 +222,10 @@ def unshard_parameters(self) -> None:
         gather_axis = changed_mesh_axis(
             self.model_weight.placements, self._unsharded_model_weight.placements
         )
-        if gather_axis is None:
-            raise RuntimeError("FSDP parameter unshard requires a changed placement axis.")
         with torch.autograd._unsafe_preserve_version_counter(
             self._unsharded_model_weight.local_buffer
         ):
-            if self._symm_mem_pool is not None:
+            if self._symm_mem_pool is not None and gather_axis is not None:
                 self._unsharded_model_weight.rendezvous(gather_axis)
             self.model_weight.redistribute(
                 self._unsharded_model_weight.placements, out=self._unsharded_model_weight
@@ -283,71 +281,58 @@ def reduce_partial_gradients(
     ) -> None:
         """Reduce a packed partial gradient buffer into sharded parameter gradients.
 
-        For HSDP main_grad rests DP-outer-Partial (Partial where main_weight is
-        Replicate) between microbatches, accumulating each backward through the
-        standard zero_grad contract; the last microbatch reduces the DP-outer axes,
-        finalizing main_grad to main_weight's placements so ``.grad`` is the fully
-        reduced gradient before ``optimizer.step()``. With every axis Flat (plain
-        DP) main_grad already rests finalized.
+        main_grad remains in the gradient placements across optimizer steps. The
+        last microbatch produces a separate optimizer gradient when its placements
+        differ from main_grad.
         """
         assert self.main_grad is not None
 
-        def has_grad(parameters: tuple[nn.Parameter, ...]) -> bool:
-            has_any_grad = False
-            has_any_missing_grad = False
-            for parameter in parameters:
-                if parameter.grad is None:
-                    has_any_missing_grad = True
-                else:
-                    has_any_grad = True
-            if has_any_grad and has_any_missing_grad:
-                raise RuntimeError("FSDP sharded gradients must be either all set or all None.")
-            return has_any_grad
-
-        # zero_grad(set_to_none=True) clears sharded parameter grads, so this
-        # backward can reduce directly into main_grad. zero_grad(set_to_none=False)
-        # leaves sharded grads installed, so this backward accumulates into main_grad.
-        has_sharded_grads = has_grad(self.sharded_parameters)
-
-        # A non-accumulation main_grad means the previous step finalized it; this
-        # only happens on the first microbatch. Redistribute it back to the
-        # DP-outer-Partial accumulation placement -- a metadata relabel for HSDP,
-        # and a fresh reduce-scattered buffer for HFSDP in the future.
-        if self.main_grad.placements != self._accumulation_placements:
-            self.main_grad = self.main_grad.redistribute(self._accumulation_placements)
-
-        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.")
-        partial_reduce_op = partial_grad.placements[reduce_axis].reduce_op
-        grad_divisor = self.mesh.size(reduce_axis) if partial_reduce_op == dist.ReduceOp.SUM else 1
-        if self._symm_mem_pool is not None:
-            partial_grad.rendezvous(reduce_axis)
-        if can_reduce_into_main_grad:
-            partial_grad.redistribute(self.main_grad.placements, out=self.main_grad)
-            if grad_divisor != 1:
-                self.main_grad.local_buffer.div_(grad_divisor)
+            assert isinstance(self.main_grad.placements[-1], Partial), (
+                "The last placement must be Partial"
+            )
+            if self.is_first_microbatch:
+                self.main_grad.local_buffer.copy_(partial_grad.local_buffer)
+            else:
+                self.main_grad.local_buffer.add_(partial_grad.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)
+            can_reduce_into_main_grad = (
+                self.is_first_microbatch and partial_grad.dtype == self.main_grad.dtype
+            )
+            partial_reduce_op = partial_grad.placements[reduce_axis].reduce_op
+            grad_divisor = self.mesh.size(reduce_axis) if partial_reduce_op == dist.ReduceOp.SUM else 1
+            if self._symm_mem_pool is not None:
+                partial_grad.rendezvous(reduce_axis)
+            if can_reduce_into_main_grad:
+                partial_grad.redistribute(self.main_grad.placements, out=self.main_grad)
+                if grad_divisor != 1:
+                    self.main_grad.local_buffer.div_(grad_divisor)
             else:
-                self.main_grad.local_buffer.copy_(reduced_grad.local_buffer)
+                reduced_grad = partial_grad.redistribute(self.main_grad.placements)
+                if grad_divisor != 1:
+                    reduced_grad.local_buffer.div_(grad_divisor)
+                if self.is_first_microbatch:
+                    self.main_grad.local_buffer.copy_(reduced_grad.local_buffer)
+                else:
+                    self.main_grad.local_buffer.add_(reduced_grad.local_buffer)
 
+        optimizer_grad = self.main_grad
         if is_last_microbatch:
-            # Finalize the deferred DP-outer reduction (all-reduce for HSDP,
-            # reduce-scatter for HFSDP) before binding the sharded parameter grads.
-            self.main_grad = self.main_grad.redistribute(self.main_weight.placements)
+            optimizer_grad = self.main_grad.redistribute(self.main_weight.placements)
+            if optimizer_grad is not self.main_grad:
+                self.main_grad.local_buffer.zero_()
 
-        # Make each sharded parameter's .grad consistent with the final main_grad.
+        # Install the accumulation grad between microbatches and the finalized grad
+        # for optimizer.step() after the last microbatch.
         for index, sharded_parameter in enumerate(self.sharded_parameters):
-            sharded_parameter.grad = self.main_grad.get_dtensor(index)
+            sharded_parameter.grad = optimizer_grad.get_dtensor(index)
 
+        if is_last_microbatch:
+            # Reset for the first microbatch of the next accumulation cycle.
+            self.is_first_microbatch = True
+        else:
+            self.is_first_microbatch = False
 
 def _get_parameter_owner(module: nn.Module, name: str) -> tuple[nn.Module, str]:
     """Resolve a root-module-relative parameter FQN to its direct owner."""

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

@@ -67,7 +67,9 @@ def changed_mesh_axis(
     for axis, (old_placement, new_placement) in enumerate(
         zip(old_placements, new_placements, strict=True)
     ):
-        if old_placement == new_placement:
+        if old_placement == new_placement or (
+            isinstance(old_placement, Partial) and isinstance(new_placement, Partial)
+        ):
             continue
         if changed_axis is not None:
             raise NotImplementedError(

tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py

@@ -108,6 +108,19 @@ def _flat_placements() -> Placements:
     return Placements(dp_axes=[0], parameter=[Flat()], gradient=[Flat()], optimizer=[Flat()])
 
 
+def _strategy_placements(sharding_strategy: str) -> Placements:
+    placements = {
+        "no_shard": (Replicate(), Partial(dist.ReduceOp.AVG), Replicate()),
+        "optim": (Replicate(), Partial(dist.ReduceOp.AVG), Flat()),
+        "optim_grads": (Replicate(), Flat(), Flat()),
+        "optim_grads_params": (Flat(), Flat(), Flat()),
+    }
+    parameter, gradient, optimizer = placements[sharding_strategy]
+    return Placements(
+        dp_axes=[0], parameter=[parameter], gradient=[gradient], optimizer=[optimizer]
+    )
+
+
 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
@@ -132,9 +145,14 @@ def _mb(num_bytes: int) -> str:
 _GEMM_OP_NAME_SUBSTRING = "aten::mm"
 
 
+@pytest.mark.parametrize(
+    "sharding_strategy", ["no_shard", "optim", "optim_grads", "optim_grads_params"]
+)
 @pytest.mark.parametrize("num_microbatches", [1, 3])
-def test_fully_shard_sgd_losses_match_baseline(distributed_setup, num_microbatches):
-    """Minimal per-module FSDP training should match single-rank SGD."""
+def test_fully_shard_sgd_losses_match_baseline(
+    distributed_setup, num_microbatches, sharding_strategy
+):
+    """Every supported sharding strategy should match single-rank SGD."""
     rank = distributed_setup.rank
     world_size = distributed_setup.world_size
     device = distributed_setup.device
@@ -147,10 +165,12 @@ def test_fully_shard_sgd_losses_match_baseline(distributed_setup, num_microbatch
     model = TinyModel().to(device)
     model.load_state_dict(baseline.state_dict())
 
-    fully_shard(model.fc1, mesh=mesh, placements=_flat_placements())
-    fully_shard(model.fc2, mesh=mesh, placements=_flat_placements())
+    placements = _strategy_placements(sharding_strategy)
+    fully_shard(model.fc1, mesh=mesh, placements=placements)
+    fully_shard(model.fc2, mesh=mesh, placements=placements)
     baseline_optimizer = torch.optim.SGD(baseline.parameters(), lr=0.05)
     optimizer = torch.optim.SGD(model.parameters(), lr=0.05)
+    fully_shard_optimizer(optimizer)
 
     micro_batch_size = 2
     x = torch.randn(num_microbatches, micro_batch_size, 8, device=device)
@@ -162,19 +182,19 @@ def train(model, optimizer, log_prefix) -> list[torch.Tensor]:
         for step in range(5):
             optimizer.zero_grad()
 
-            for microbatch, (microbatch_x, microbatch_target) in enumerate(microbatches):
-                loss = torch.nn.functional.mse_loss(model(microbatch_x), microbatch_target)
-                losses.append(loss.detach())
-                logger.debug(
-                    "%s train parity: rank=%s, step=%s, microbatch=%s, loss=%s",
-                    log_prefix,
-                    rank,
-                    step,
-                    microbatch,
-                    loss,
-                )
-
-                (loss / num_microbatches).backward()
+            for microbatch_index, (microbatch_x, microbatch_target) in enumerate(microbatches):
+                with microbatch(model, is_last=microbatch_index == num_microbatches - 1):
+                    loss = torch.nn.functional.mse_loss(model(microbatch_x), microbatch_target)
+                    losses.append(loss.detach())
+                    logger.debug(
+                        "%s train parity: rank=%s, step=%s, microbatch=%s, loss=%s",
+                        log_prefix,
+                        rank,
+                        step,
+                        microbatch_index,
+                        loss,
+                    )
+                    (loss / num_microbatches).backward()
 
             optimizer.step()
         return losses
@@ -669,8 +689,9 @@ def test_backward_averages_across_dp_and_accumulates_across_calls(distributed_se
     fully_shard(model, mesh=mesh, placements=_flat_placements())
 
     x = torch.full((1, 1), float(rank + 1), device=device)
-    model(x).sum().backward()
-    model(x).sum().backward()
+    with microbatch(model, is_last=False):
+        model(x).sum().backward()
+        model(x).sum().backward()
 
     assert isinstance(model.weight.grad, DTensor)
     local_grad = model.weight.grad.to_local()