diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/__init__.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/__init__.py index bc9118598d1..bae27be831c 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/__init__.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/__init__.py @@ -16,6 +16,7 @@ from .dbuffer import DBuffer from .fully_shard import fully_shard, microbatch +from .optimizer import fully_shard_optimizer from .placement import Flat, Partial, Placement, Placements, Replicate __all__ = [ @@ -26,5 +27,6 @@ "Placements", "Replicate", "fully_shard", + "fully_shard_optimizer", "microbatch", ] diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/optimizer.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/optimizer.py new file mode 100644 index 00000000000..3f7745da9ff --- /dev/null +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/optimizer.py @@ -0,0 +1,108 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Optimizer adapter for the minimal Megatron-FSDP path.""" + +from typing import Any, NamedTuple + +import torch +from torch import nn + +from .parameter_group import contained_in_parameter_group + + +def fully_shard_optimizer(optimizer: torch.optim.Optimizer) -> None: + """Attach FSDP-aware step hooks to an optimizer instance. + + The adapted optimizer preserves its existing parameter groups and only adds + temporary gradient casting around optimizer steps for FSDP sharded + parameters whose data dtype differs from their grad dtype. + + Alternatives considered: + - Monkey-patching optimizer methods directly on the instance. This is + more invasive and harder to compose than hooks. + - Generating an FSDP-specific subclass per ``torch.optim.Optimizer``. + This adds extra class-generation machinery, but would let us + instrument ``zero_grad`` and ``__init__`` as well as ``step`` if needed. + - Casting from ``main_grad.dtype`` to ``main_weight.dtype`` after the + last microbatch and casting back before the first microbatch. This + should be done from a root post-backward callback if needed later, so + users do not need to call ``fully_shard_optimizer`` on an existing + ``torch.optim.Optimizer``. + - Letting the user set ``main_weight`` and ``main_grad`` to the same + dtype. This is enough for an FSDP2 drop-in replacement path and lets + optimizers stay unaware of FSDP precision handling. + + Args: + optimizer: Optimizer instance to adapt in place. + """ + + class CastedGrad(NamedTuple): + """Original grad tensor temporarily replaced during an optimizer step.""" + + parameter: nn.Parameter + original_grad: torch.Tensor + + def set_grad(parameter: nn.Parameter, grad: torch.Tensor) -> None: + """Install a grad with matching grad_dtype on a sharded parameter.""" + # Clear the existing grad before switching grad_dtype; the sharded + # parameter cannot advertise a new grad dtype while the old grad + # object with the previous dtype is still attached. + parameter.grad = None + parameter.grad_dtype = grad.dtype + parameter.grad = grad + + casted_grads: list[CastedGrad] = [] + + def step_pre_hook( + hooked_optimizer: torch.optim.Optimizer, args: tuple[Any, ...], kwargs: dict[str, Any] + ) -> None: + closure = kwargs.get("closure") + if closure is None and len(args) > 1: + closure = args[1] + if closure is not None: + # Step hooks run outside the base optimizer step, but closures run inside it. + # We need to cast grads after the closure materializes them and before the + # optimizer consumes them, which this hook-only adapter cannot intercept. + raise NotImplementedError( + "fully_shard_optimizer does not support optimizer.step closures." + ) + assert not casted_grads + for group in hooked_optimizer.param_groups: + for parameter in group["params"]: + if not isinstance(parameter, nn.Parameter): + raise TypeError( + "fully_shard_optimizer expected optimizer param groups to contain " + f"nn.Parameter values, got {type(parameter)!r}." + ) + if not contained_in_parameter_group(parameter): + continue + if parameter.grad is None: + continue + if parameter.grad.dtype == parameter.dtype: + continue + + casted_grads.append(CastedGrad(parameter, parameter.grad)) + set_grad(parameter, parameter.grad.to(dtype=parameter.dtype)) + + def step_post_hook( + hooked_optimizer: torch.optim.Optimizer, args: tuple[Any, ...], kwargs: dict[str, Any] + ) -> None: + del hooked_optimizer, args, kwargs + for parameter, original_grad in casted_grads: + set_grad(parameter, original_grad) + casted_grads.clear() + + optimizer.register_step_pre_hook(step_pre_hook) + optimizer.register_step_post_hook(step_post_hook) 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..896f834a7f7 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py @@ -15,6 +15,7 @@ Flat, Placements, fully_shard, + fully_shard_optimizer, microbatch, ) from megatron.core.distributed.fsdp.src.megatron_fsdp.mixed_precision import MixedPrecisionPolicy @@ -112,7 +113,7 @@ def _events_overlap(first, second) -> bool: @pytest.mark.parametrize("num_microbatches", [1, 3]) -def test_fully_shard_losses_match_baseline(distributed_setup, num_microbatches): +def test_fully_shard_sgd_losses_match_baseline(distributed_setup, num_microbatches): """Minimal per-module FSDP training should match single-rank SGD.""" rank = distributed_setup.rank world_size = distributed_setup.world_size @@ -466,6 +467,41 @@ def train_iteration() -> torch.Tensor: torch.testing.assert_close(second_loss, first_loss) +def test_fully_shard_adam_mixed_precision_losses_match_baseline(distributed_setup): + """Mixed-precision FSDP Adam should track an unsharded Adam baseline.""" + world_size = distributed_setup.world_size + device = distributed_setup.device + if world_size < 2: + pytest.skip("This test requires at least 2 ranks.") + mesh = init_device_mesh(device.type, (world_size,)) + torch.manual_seed(2026) + baseline = TinyModel().to(device=device, dtype=torch.bfloat16) + model = TinyModel().to(device=device, dtype=torch.bfloat16) + 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()) + + baseline_optimizer = torch.optim.Adam(baseline.parameters(), lr=0.01) + optimizer = torch.optim.Adam(model.parameters(), lr=0.01) + fully_shard_optimizer(optimizer) + + x = torch.randn(3, 8, device=device, dtype=torch.bfloat16) + target = torch.randn(3, 4, device=device, dtype=torch.bfloat16) + + for _ in range(3): + baseline_optimizer.zero_grad() + optimizer.zero_grad() + + baseline_loss = torch.nn.functional.mse_loss(baseline(x).float(), target.float()) + loss = torch.nn.functional.mse_loss(model(x).float(), target.float()) + torch.testing.assert_close(loss, baseline_loss, rtol=0, atol=3e-3) + + baseline_loss.backward() + loss.backward() + baseline_optimizer.step() + optimizer.step() + + def test_microbatch_scopes_child_contexts(distributed_setup): """microbatch() should scope FSDP child contexts under an unwrapped parent.""" world_size = distributed_setup.world_size diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_optimizer.py b/tests/unit_tests/distributed/mfsdp_v2/test_optimizer.py new file mode 100644 index 00000000000..3f3b636dd19 --- /dev/null +++ b/tests/unit_tests/distributed/mfsdp_v2/test_optimizer.py @@ -0,0 +1,99 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Unit tests for Megatron-FSDP optimizer behavior.""" + +import pytest +import torch +from torch import nn +from torch.distributed.device_mesh import init_device_mesh +from transformer_engine.pytorch.optimizers import FusedAdam + +from megatron.core.distributed.fsdp.src.megatron_fsdp.experimental import ( + Flat, + Placements, + fully_shard, +) +from megatron.core.distributed.fsdp.src.megatron_fsdp.mixed_precision import MixedPrecisionPolicy + + +class TinyModel(nn.Module): + """Small model with two separately shardable units.""" + + def __init__(self) -> None: + super().__init__() + self.fc1 = nn.Linear(8, 16) + self.relu = nn.ReLU() + self.fc2 = nn.Linear(16, 4) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Run the tiny model.""" + return self.fc2(self.relu(self.fc1(x))) + + +def _flat_placements() -> Placements: + return Placements(dp_axes=[0], parameter=[Flat()], gradient=[Flat()], optimizer=[Flat()]) + + +def test_adam_without_adapter_raises_precision_error(distributed_setup): + """Raw Adam should fail on mixed-precision FSDP parameters without the adapter.""" + world_size = distributed_setup.world_size + device = distributed_setup.device + mesh = init_device_mesh(device.type, (world_size,)) + torch.manual_seed(2026) + model = TinyModel().to(device=device, dtype=torch.bfloat16) + fully_shard(model.fc1, mesh=mesh, placements=_flat_placements()) + fully_shard(model.fc2, mesh=mesh, placements=_flat_placements()) + optimizer = torch.optim.Adam(model.parameters(), lr=0.01) + + x = torch.randn(6, 8, device=device, dtype=torch.bfloat16) + optimizer.zero_grad(set_to_none=True) + loss = model(x).sum() + loss.backward() + + with pytest.raises(RuntimeError, match="dtype"): + optimizer.step() + + +def test_fused_adam_without_adapter_accepts_mismatched_grads(distributed_setup): + """TE FusedAdam should handle mixed-precision FSDP grads without the adapter.""" + world_size = distributed_setup.world_size + device = distributed_setup.device + + mesh = init_device_mesh(device.type, (world_size,)) + torch.manual_seed(2026) + model = TinyModel().to(device=device, dtype=torch.bfloat16) + # These are the defaults, but spell them out so the test clearly exercises + # mismatched parameter and gradient precision. + mixed_precision_policy = MixedPrecisionPolicy( + main_params_dtype=torch.float32, main_grads_dtype=torch.bfloat16 + ) + fully_shard( + model.fc1, + mesh=mesh, + placements=_flat_placements(), + mixed_precision_policy=mixed_precision_policy, + ) + fully_shard( + model.fc2, + mesh=mesh, + placements=_flat_placements(), + mixed_precision_policy=mixed_precision_policy, + ) + optimizer = FusedAdam(model.parameters(), lr=0.01) + + x = torch.randn(6, 8, device=device, dtype=torch.bfloat16) + optimizer.zero_grad(set_to_none=True) + loss = model(x).sum() + loss.backward() + + for parameter in model.parameters(): + assert parameter.grad is not None + assert parameter.dtype != parameter.grad.dtype + + params_before_step = [parameter.detach().clone() for parameter in model.parameters()] + optimizer.step() + + assert any( + not torch.equal(parameter_before, parameter.detach()) + for parameter_before, parameter in zip(params_before_step, model.parameters()) + )