Add fully_shard_optimizer for mixed-precision FSDP

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",
 ]

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)

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

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="same device and the same 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())
+    )