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 1cb9dc2ddaf..20322e21b68 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module.py +++ b/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: diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module_utils.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module_utils.py index 06fddd7f348..e47c2af3b7e 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/module_utils.py +++ b/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(".") diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/schedule_policy.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/schedule_policy.py index 27511a1f71f..fcdf1d623de 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/schedule_policy.py +++ b/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.""" diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_quantization.py b/tests/unit_tests/distributed/mfsdp_v2/test_quantization.py index 7e1f1b9db98..7a6393e147e 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_quantization.py +++ b/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)