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 6f0e6c5170f..4433e093d84 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py @@ -14,6 +14,7 @@ """Module mixin for the minimal Megatron-FSDP path.""" +import enum from collections.abc import Callable from typing import Literal, cast from weakref import ref @@ -28,6 +29,11 @@ from .placement import Placements +def _is_in_backward() -> bool: + """Return whether the current thread is executing an autograd GraphTask.""" + return torch._C._current_graph_task_id() != -1 + + class FsdpContext: """Runtime stream and prefetch state shared by FSDP roots constructed together.""" @@ -119,6 +125,13 @@ def post_backward_final_callback() -> None: class FsdpModule: """Mixin attached to modules managed by the minimal FSDP path.""" + class Phase(enum.Enum): + """Lifecycle phase of this FsdpModule.""" + + RESTING = enum.auto() + FORWARD = enum.auto() + BACKWARD = enum.auto() + # Name relative to the root FSDP module from named_modules(). # Root uses "" and None means uninitialized. _name: str | None @@ -131,6 +144,9 @@ class FsdpModule: # ``None`` lets pre_forward enqueue an all-gather unless an earlier FsdpModule # already prefetched this module. _unshard_event: torch.cuda.Event | None + # Backward-pre hook sets this to BACKWARD before activation recomputation + # can run. Forward and backward hooks own all other transitions. + _phase: Phase def __init__( self, @@ -145,6 +161,7 @@ def __init__( self._is_root = False self._name = None self._unshard_event = None + self._phase = FsdpModule.Phase.RESTING owned_parameters = _collect_owned_parameters(self) assert tuple(placements.dp_axes) == tuple( range(mesh.ndim) @@ -237,6 +254,16 @@ def pre_forward(self) -> None: """ context = self.context context.ensure_finalized() + # post_forward() resets the phase after a non-recomputed forward, so a + # FORWARD phase here means this forward-pre hook ran while the previous + # forward was still in progress. + assert self._phase is not FsdpModule.Phase.FORWARD + # A reentrant checkpoint recomputes before the child module's backward-pre + # hook can set its phase. Its forward still runs inside the active autograd + # GraphTask, which is the signal PyTorch FSDP2 uses as well. + is_recomputing = self._phase is FsdpModule.Phase.BACKWARD or _is_in_backward() + if not is_recomputing: + self._phase = FsdpModule.Phase.FORWARD torch.cuda.nvtx.range_push(self._nvtx_label("forward")) self._num_ready_grad_parameters = 0 allgather_stream = context.allgather_stream @@ -251,9 +278,13 @@ def pre_forward(self) -> None: # issued afterwards, so it is free to run concurrently with this FsdpModule). current_stream.wait_event(self._unshard_event) - next_module = context.forward_order.next_item(self) - if next_module is not None: - next_module._unshard_parameter_groups() + # Activation recomputation runs forward hooks inside backward. Do not + # prefetch the next module in forward order: its backward may already + # be complete, so no later backward hook would reshard it. + if not is_recomputing: + next_module = context.forward_order.next_item(self) + if next_module is not None: + next_module._unshard_parameter_groups() def _unshard_parameter_groups(self) -> None: """Unshard this FsdpModule's parameter groups on the all-gather stream. @@ -274,7 +305,13 @@ def _unshard_parameter_groups(self) -> None: def post_forward(self) -> None: """Return parameters to their sharded resting state after forward compute.""" - self._reshard_parameter_groups() + # Recomputed parameters are consumed immediately by this module's + # backward. Keep them materialized to avoid an unnecessary all-gather; + # post_backward() will reshard them after gradient reduction. + is_recomputing = self._phase is FsdpModule.Phase.BACKWARD or _is_in_backward() + if not is_recomputing: + self._reshard_parameter_groups() + self._phase = FsdpModule.Phase.RESTING torch.cuda.nvtx.range_pop() def _reshard_parameter_groups(self) -> None: @@ -297,6 +334,7 @@ def _reshard_parameter_groups(self) -> None: def pre_backward(self) -> None: """Prepare full parameters and prefetch the next FsdpModule in backward order.""" + self._phase = FsdpModule.Phase.BACKWARD torch.cuda.nvtx.range_push(self._nvtx_label("backward")) context = self.context current_stream = context.current_stream() @@ -323,6 +361,7 @@ def post_backward(self) -> None: """Reduce gradients and return parameters to their sharded resting state.""" self._reduce_gradient_groups() self._reshard_parameter_groups() + self._phase = FsdpModule.Phase.RESTING torch.cuda.nvtx.range_pop() def _reduce_gradient_groups(self) -> 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 357d765f63e..863becf8b77 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py @@ -11,6 +11,7 @@ from torch.distributed.device_mesh import DeviceMesh, init_device_mesh from torch.distributed.tensor import DTensor from torch.profiler import ProfilerActivity, profile +from torch.utils.checkpoint import checkpoint from megatron.core.distributed.fsdp.src.megatron_fsdp.experimental import ( Flat, @@ -22,6 +23,7 @@ fully_shard_optimizer, microbatch, ) +from megatron.core.distributed.fsdp.src.megatron_fsdp.experimental.module import FsdpModule from megatron.core.distributed.fsdp.src.megatron_fsdp.mixed_precision import MixedPrecisionPolicy from tests.unit_tests.distributed.mfsdp_v2.profiler_utils import ( collect_linked_kernels, @@ -45,6 +47,19 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return self.fc2(self.relu(self.fc1(x))) +class CheckpointedTinyModel(TinyModel): + """Tiny model that activation-checkpoints each shardable module.""" + + def __init__(self, use_reentrant: bool) -> None: + super().__init__() + self.use_reentrant = use_reentrant + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Run each linear layer through activation checkpointing.""" + x = checkpoint(self.fc1, x, use_reentrant=self.use_reentrant) + return checkpoint(self.fc2, self.relu(x), use_reentrant=self.use_reentrant) + + class NestedModel(nn.Module): """Model with direct and child-owned parameters.""" @@ -217,6 +232,48 @@ def train(model, optimizer, log_prefix) -> list[torch.Tensor]: ) +@pytest.mark.parametrize("use_reentrant", [False, True], ids=["non_reentrant", "reentrant"]) +def test_fully_shard_activation_recompute_reshards_parameters(distributed_setup, use_reentrant): + """Activation recomputation should leave every FSDP module resharded. + + Backward completes ``fc2`` before recomputing ``fc1``. Without suppressing + forward prefetch during recomputation, ``fc1`` unshards ``fc2`` again after + its backward hook has run, leaving ``fc2.weight`` as an unsharded Parameter + instead of a sharded DTensor at the end of backward. + """ + world_size = distributed_setup.world_size + device = distributed_setup.device + + mesh = init_device_mesh(device.type, (world_size,)) + model = CheckpointedTinyModel(use_reentrant=use_reentrant).to(device) + with fully_shard_context(device=device): + fully_shard(model.fc1, mesh=mesh, placements=_flat_placements()) + fully_shard(model.fc2, mesh=mesh, placements=_flat_placements()) + fully_shard(model, mesh=mesh, placements=_flat_placements()) + + x = torch.randn(2, 8, device=device, requires_grad=True) + model(x).sum().backward() + + # Without the forward-prefetch suppression, ``fc1``'s recomputed forward + # would unshard ``fc2`` after ``fc2``'s backward already resharded it, + # leaving an unsharded Parameter here. + assert isinstance(model.fc1.weight, DTensor) + assert isinstance(model.fc2.weight, DTensor) + + # Backward completes each module before recomputing the previous one, so + # every module-local phase must be cleared after its matching backward. + assert model._phase is FsdpModule.Phase.RESTING + assert model.fc1._phase is FsdpModule.Phase.RESTING + assert model.fc2._phase is FsdpModule.Phase.RESTING + + # A second forward after backward runs in the forward phase again, so + # forward-order prefetch resumes and the module phases return to resting. + model(x).sum().backward() + assert model._phase is FsdpModule.Phase.RESTING + assert model.fc1._phase is FsdpModule.Phase.RESTING + assert model.fc2._phase is FsdpModule.Phase.RESTING + + @pytest.mark.parametrize("set_to_none", [True, False]) @pytest.mark.parametrize("num_microbatches", [1, 3]) def test_hsdp_losses_match_baseline(distributed_setup, num_microbatches, set_to_none):