Add SchedulePolicy.unshard_all_planes_in_backward for MXFP8 recomputation in MFSDP v2

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

@@ -32,7 +32,7 @@
 from .context import FsdpContext
 from .countdown import Countdown
 from .execution_order import ExecutionOrder
-from .module_utils import get_parameter_owner
+from .module_utils import get_parameter_owner, is_in_backward
 from .parameter_group import (
     FsdpParameter,
     FsdpParameterGroup,
@@ -44,11 +44,6 @@
 from .schedule_policy import SchedulePolicy
 
 
-def _is_in_backward() -> bool:
-    """Return whether the current thread is executing an autograd GraphTask."""
-    return torch._C._current_graph_task_id() != -1
-
-
 class FsdpModule:
     """Mixin attached to modules managed by the minimal FSDP path."""
 
@@ -273,7 +268,7 @@ def pre_forward(self) -> None:
         context.ensure_finalized()
         # A reentrant checkpoint recomputes before the child module's backward-pre
         # hook runs. The active autograd GraphTask identifies that recomputation.
-        is_recomputing = self.phase is Phase.BACKWARD or _is_in_backward()
+        is_recomputing = self.phase is Phase.BACKWARD or is_in_backward()
         if self.phase is not Phase.BACKWARD:
             self.phase = Phase.FORWARD
         # forward/backward each span multiple lifecycle methods (pre_forward ->
@@ -360,13 +355,18 @@ def _unshard_parameter_groups(self, phase: Phase | None) -> None:
             raise ValueError(
                 f"Unsharding requires a FORWARD, BACKWARD, or None phase, got {phase!r}."
             )
+        # Recomputation that does not re-enter this module's pre_forward reads MXFP8 rowwise
+        # planes in backward, so gather them together with columnwise ones. Other dtype
+        # weights serve both phases already.
+        if phase is Phase.BACKWARD and self._schedule_policy.unshard_all_planes_in_backward:
+            phase = None
         with torch.cuda.stream(self.context.allgather_stream):
             for group in self._parameter_groups:
                 group.unshard_parameters(phase)
 
     def post_forward(self) -> None:
         """Return parameters to their sharded resting state after forward compute."""
-        is_recomputing = self.phase is Phase.BACKWARD or _is_in_backward()
+        is_recomputing = self.phase is Phase.BACKWARD or is_in_backward()
         if is_recomputing:
             self._release_unsharded_storage(rowwise_only=True)
         else:

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

@@ -12,11 +12,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Utilities for parameter ownership and metadata in FSDP modules."""
+"""Utilities for parameter ownership, metadata, and autograd state in FSDP modules."""
 
+import torch
 from torch import nn
 
 
+def is_in_backward() -> bool:
+    """Return whether the current thread is executing an autograd GraphTask."""
+    return torch._C._current_graph_task_id() != -1
+
+
 def get_parameter_owner(root_module: nn.Module, parameter_fqn: str) -> tuple[nn.Module, str]:
     """Resolve a root-module-relative parameter FQN to its direct owner."""
     module_name, separator, parameter_name = parameter_fqn.rpartition(".")

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

@@ -33,12 +33,19 @@ class SchedulePolicy:
             window, including separate GraphTasks and delayed weight gradients.
             None defaults to one callback per owned trainable parameter. The scheduler
             must declare the complete count before sharding the unit.
+        unshard_all_planes_in_backward: Whether backward unshards and prefetches of this
+            module gather MXFP8 rowwise planes along with columnwise ones. Enable it when
+            activations recomputed in backward use this module's MXFP8 weights without
+            re-entering this module's forward hooks, e.g. a checkpointed submodule.
+            Otherwise, backward unshards gather only columnwise planes, and the
+            recomputation fails in Transformer Engine.
     """
 
     forward_prefetch_size: int | None = None
     backward_prefetch_size: int | None = None
     register_hooks: bool = True
     grad_accumulation_count: int | None = None
+    unshard_all_planes_in_backward: bool = False
 
     def __post_init__(self) -> None:
         """Validate non-negative prefetch budgets."""

tests/unit_tests/distributed/mfsdp_v2/test_quantization.py

@@ -18,12 +18,16 @@
 from megatron.core.distributed.fsdp.src.megatron_fsdp.experimental import (
     FsdpContext,
     Placements,
+    SchedulePolicy,
     fully_shard,
     fully_shard_context,
     fully_shard_optimizer,
     microbatch,
 )
 from megatron.core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer import DBuffer
+from megatron.core.distributed.fsdp.src.megatron_fsdp.experimental.module_utils import (
+    is_in_backward,
+)
 from megatron.core.distributed.fsdp.src.megatron_fsdp.experimental.placement import (
     BlockAtomic,
     ShardLayout,
@@ -57,6 +61,27 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return te_checkpoint(self.fc2, F.gelu(x), use_reentrant=self.use_reentrant)
 
 
+class _RecomputedBlockWithOwnedProj(nn.Module):
+    """Block that checkpoints only its ``proj`` submodule.
+
+    When the block is an FsdpModule, the checkpoint recomputes ``proj`` without re-entering
+    the block's forward hooks, like MultiTokenPredictionLayer.eh_proj under the root.
+    """
+
+    def __init__(self, device, use_reentrant: bool | None) -> None:
+        super().__init__()
+        self.proj = te.Linear(64, 64, params_dtype=torch.bfloat16, device=device)
+        self.mlp = _make_mlp(device)
+        self.use_reentrant = use_reentrant
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Run the block, checkpointing ``proj`` unless use_reentrant is None."""
+        if self.use_reentrant is None:
+            return self.mlp(self.proj(x))
+        # TE's checkpoint restores the MXFP8 autocast state during recomputation.
+        return self.mlp(te_checkpoint(self.proj, x, use_reentrant=self.use_reentrant))
+
+
 _ZERO3_PLACEMENTS = Placements(
     dp_axes=[0], parameter=[Shard(0)], gradient=[Shard(0)], optimizer=[Shard(0)]
 )
@@ -105,6 +130,55 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 )
 
 
+def _unsharded_mxfp8_weight(module):
+    """Return the unsharded buffer of the module's single MXFP8 parameter group."""
+    groups = [g for g in module.parameter_groups if isinstance(g.model_weight, QuantizedDBuffer)]
+    assert len(groups) == 1, f"Expected one MXFP8 parameter group, got {len(groups)}."
+    return groups[0]._unsharded_model_weight
+
+
+def _is_rowwise_allocated(module):
+    """Return whether the module's unsharded MXFP8 rowwise planes are allocated."""
+    return _unsharded_mxfp8_weight(module).is_rowwise_allocated
+
+
+def _is_columnwise_allocated(module):
+    """Return whether the module's unsharded MXFP8 columnwise planes are allocated."""
+    return _unsharded_mxfp8_weight(module).is_columnwise_allocated
+
+
+def _check_recompute_matches_no_recompute(distributed_setup, recipe, build, use_reentrant):
+    """Train builds without and with recomputation and check they match bit-exactly.
+
+    ``build(use_reentrant)`` returns a fully sharded model and its optimizer.
+    """
+    torch.manual_seed(1234 + distributed_setup.rank)
+    inputs = torch.randn(3, 32, 64, dtype=torch.bfloat16, device=distributed_setup.device)
+
+    def train(model, optimizer):
+        """Return per-step losses and sharded gradients, and the final sharded weights."""
+        losses, grads = [], []
+        for x in inputs:
+            optimizer.zero_grad(set_to_none=True)
+            # Reentrant checkpointing produces gradients only for inputs requiring grad.
+            x = x.clone().requires_grad_()
+            with te.autocast(recipe=recipe):
+                output = model(x)
+            loss = output.float().square().mean()
+            loss.backward()
+            losses.append(loss.detach())
+            grads.append([parameter.grad.to_local().clone() for parameter in model.parameters()])
+            optimizer.step()
+        weights = [parameter.to_local().clone() for parameter in model.parameters()]
+        return losses, grads, weights
+
+    expected_losses, expected_grads, expected_weights = train(*build(None))
+    losses, grads, weights = train(*build(use_reentrant))
+    torch.testing.assert_close(losses, expected_losses, rtol=0, atol=0)
+    torch.testing.assert_close(grads, expected_grads, rtol=0, atol=0)
+    torch.testing.assert_close(weights, expected_weights, rtol=0, atol=0)
+
+
 def _check_mxfp8_mlp_matches_reference(
     distributed_setup,
     mesh,
@@ -245,14 +319,6 @@ def test_mxfp8_unshard_gathers_rowwise_for_forward_and_columnwise_for_backward(d
         fully_shard(model.fc2, mesh=mesh, placements=_ZERO3_PLACEMENTS)
         fully_shard(model, mesh=mesh, placements=_ZERO3_PLACEMENTS)
 
-    def allocated_planes(module):
-        """Return whether the module's unsharded (rowwise, columnwise) planes are allocated."""
-        [group] = [
-            g for g in module.parameter_groups if isinstance(g.model_weight, QuantizedDBuffer)
-        ]
-        unsharded = group._unsharded_model_weight
-        return unsharded.is_rowwise_allocated, unsharded.is_columnwise_allocated
-
     def installed_planes(module):
         """Return whether the module's weight carries (rowwise, columnwise) data."""
         weight = module.weight
@@ -265,14 +331,16 @@ def installed_planes(module):
     def check_forward(_module, _args):
         # fc1 runs forward GEMMs and prefetches fc2 for its forward GEMMs.
         for module in (model.fc1, model.fc2):
-            assert allocated_planes(module) == (True, False)
+            assert _is_rowwise_allocated(module)
+            assert not _is_columnwise_allocated(module)
             assert installed_planes(module) == (True, False)
         checked_phases.append("forward")
 
     def check_backward(_module, _grad_output):
         # fc2 runs data-gradient GEMMs and prefetches fc1 for its data-gradient GEMMs.
         for module in (model.fc2, model.fc1):
-            assert allocated_planes(module) == (False, True)
+            assert not _is_rowwise_allocated(module)
+            assert _is_columnwise_allocated(module)
             assert installed_planes(module) == (False, True)
         checked_phases.append("backward")
 
@@ -286,7 +354,8 @@ def check_backward(_module, _grad_output):
 
     assert checked_phases == ["forward", "backward"]
     for module in (model.fc1, model.fc2):
-        assert allocated_planes(module) == (False, False)
+        assert not _is_rowwise_allocated(module)
+        assert not _is_columnwise_allocated(module)
 
 
 @pytest.mark.launch_on_gb200
@@ -296,14 +365,7 @@ def check_backward(_module, _grad_output):
 def test_mxfp8_activation_recompute_matches_no_recompute(
     distributed_setup, use_reentrant, placements
 ):
-    """Recomputation regathers rowwise planes in backward without changing numerics.
-
-    With sharded compute weights, backward prefetch gathers only columnwise planes,
-    so recomputed forwards must gather rowwise planes on demand. Replicated compute
-    weights need no gather and keep both planes installed. Shapes need padded
-    scales, which covers reinstalling padded scale copies on the Parameter TE saved
-    for backward.
-    """
+    """Recomputation regathers rowwise planes in backward without changing numerics."""
     device = distributed_setup.device
     if distributed_setup.world_size < 2:
         pytest.skip("Distributed MXFP8 coverage requires at least two ranks.")
@@ -323,28 +385,81 @@ def build(use_reentrant):
         fully_shard_optimizer(optimizer)
         return model, optimizer
 
-    torch.manual_seed(1234 + distributed_setup.rank)
-    inputs = torch.randn(3, 32, 64, dtype=torch.bfloat16, device=device)
+    _check_recompute_matches_no_recompute(distributed_setup, recipe, build, use_reentrant)
 
-    def train(model, optimizer):
-        """Return per-step losses and sharded gradients, and the final sharded weights."""
-        losses, grads = [], []
-        for x in inputs:
-            optimizer.zero_grad(set_to_none=True)
-            # Reentrant checkpointing produces gradients only for inputs requiring grad.
-            x = x.clone().requires_grad_()
-            with te.autocast(recipe=recipe):
-                output = model(x)
-            loss = output.float().square().mean()
+
+@pytest.mark.launch_on_gb200
+@requires_mxfp8
+@pytest.mark.parametrize("use_reentrant", [False, True], ids=["non_reentrant", "reentrant"])
+def test_mxfp8_recompute_inside_fsdp_module_matches_no_recompute(distributed_setup, use_reentrant):
+    """Checks recomputing ``proj`` inside the opted-in root matches no recompute bit-exactly."""
+    device = distributed_setup.device
+    if distributed_setup.world_size < 2:
+        pytest.skip("Distributed MXFP8 coverage requires at least two ranks.")
+
+    recipe = MXFP8BlockScaling()
+    mesh = init_device_mesh(device.type, (distributed_setup.world_size,))
+    policy = SchedulePolicy(unshard_all_planes_in_backward=True)
+
+    def build(use_reentrant):
+        with te.quantized_model_init(recipe=recipe, preserve_high_precision_init_val=True):
+            torch.manual_seed(2026)
+            model = _RecomputedBlockWithOwnedProj(device, use_reentrant)
+        with fully_shard_context(FsdpContext(device=device)):
+            fully_shard(model, mesh=mesh, placements=_ZERO3_PLACEMENTS, schedule_policy=policy)
+        optimizer = FusedAdam(model.parameters(), lr=0.01)
+        fully_shard_optimizer(optimizer)
+        return model, optimizer
+
+    _check_recompute_matches_no_recompute(distributed_setup, recipe, build, use_reentrant)
+
+
+@pytest.mark.launch_on_gb200
+@requires_mxfp8
+@pytest.mark.parametrize("use_reentrant", [False, True], ids=["non_reentrant", "reentrant"])
+@pytest.mark.parametrize(
+    "unshard_all_planes_in_backward", [False, True], ids=["on_demand", "all_planes_in_backward"]
+)
+def test_mxfp8_backward_unshard_gathers_rowwise_for_recompute_inside_fsdp_module(
+    distributed_setup, use_reentrant, unshard_all_planes_in_backward
+):
+    """Checks recomputing ``proj`` inside the root fails without the opt-in.
+
+    With the opt-in, checks both planes are ready during recomputation and released after
+    backward.
+    """
+    device = distributed_setup.device
+    if distributed_setup.world_size < 2:
+        pytest.skip("Distributed MXFP8 coverage requires at least two ranks.")
+
+    recipe = MXFP8BlockScaling()
+    with te.quantized_model_init(recipe=recipe, preserve_high_precision_init_val=True):
+        model = _RecomputedBlockWithOwnedProj(device, use_reentrant)
+    mesh = init_device_mesh(device.type, (distributed_setup.world_size,))
+    policy = SchedulePolicy(unshard_all_planes_in_backward=unshard_all_planes_in_backward)
+    with fully_shard_context(FsdpContext(device=device)):
+        fully_shard(model, mesh=mesh, placements=_ZERO3_PLACEMENTS, schedule_policy=policy)
+
+    # Reentrant checkpointing produces gradients only for inputs requiring grad.
+    x = torch.randn(32, 64, dtype=torch.bfloat16, device=device, requires_grad=True)
+    with te.autocast(recipe=recipe):
+        output = model(x)
+    loss = output.float().square().mean()
+    if not unshard_all_planes_in_backward:
+        with pytest.raises(Exception, match="row-scaled"):
             loss.backward()
-            losses.append(loss.detach())
-            grads.append([parameter.grad.to_local().clone() for parameter in model.parameters()])
-            optimizer.step()
-        weights = [parameter.to_local().clone() for parameter in model.parameters()]
-        return losses, grads, weights
+        return
 
-    expected_losses, expected_grads, expected_weights = train(*build(None))
-    losses, grads, weights = train(*build(use_reentrant))
-    torch.testing.assert_close(losses, expected_losses, rtol=0, atol=0)
-    torch.testing.assert_close(grads, expected_grads, rtol=0, atol=0)
-    torch.testing.assert_close(weights, expected_weights, rtol=0, atol=0)
+    recomputed = []
+
+    def check_recompute(_module, _args):
+        if is_in_backward():
+            assert _is_rowwise_allocated(model)
+            assert _is_columnwise_allocated(model)
+            recomputed.append(True)
+
+    model.proj.register_forward_pre_hook(check_recompute)
+    loss.backward()
+    assert recomputed == [True]
+    assert not _is_rowwise_allocated(model)
+    assert not _is_columnwise_allocated(model)