From 5b8177fffae4012545b8495daae8568ab575b4a4 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Sat, 27 Jun 2026 06:18:47 +0000 Subject: [PATCH 01/12] Implement experimental fully_shard_optimizer hook adapter Signed-off-by: Jingyue Wu --- .../megatron_fsdp/experimental/__init__.py | 2 + .../megatron_fsdp/experimental/optimizer.py | 107 ++++++++++++++++ .../distributed/mfsdp_v2/test_fully_shard.py | 114 ++++++++++++++++++ 3 files changed, 223 insertions(+) create mode 100644 megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/optimizer.py 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..409dfbb9d8f --- /dev/null +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/optimizer.py @@ -0,0 +1,107 @@ +# 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 dataclasses import dataclass +from typing import Any, TypeVar + +import torch +from torch import nn + +from .parameter_group import contained_in_parameter_group + +_OptimizerT = TypeVar("_OptimizerT", bound=torch.optim.Optimizer) + + +def fully_shard_optimizer(optimizer: _OptimizerT) -> 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. + + Args: + optimizer: Optimizer instance to adapt in place. + + """ + if not isinstance(optimizer, torch.optim.Optimizer): + raise TypeError( + "fully_shard_optimizer expected a torch.optim.Optimizer instance, " + f"got {optimizer!r}." + ) + + @dataclass + class CastedGrad: + parameter: nn.Parameter + original_grad: torch.Tensor + + 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 + + original_grad = parameter.grad + casted_grads.append(CastedGrad(parameter, original_grad)) + + # 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 = parameter.dtype + parameter.grad = original_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 casted_grad in casted_grads: + parameter = casted_grad.parameter + original_grad = casted_grad.original_grad + parameter.grad = None + parameter.grad_dtype = original_grad.dtype + parameter.grad = 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..ffe95c24c5e 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 @@ -465,6 +466,119 @@ def train_iteration() -> torch.Tensor: with pytest.raises(AssertionError): torch.testing.assert_close(second_loss, first_loss) +def test_fully_shard_optimizer_keeps_optimizer_instance_and_sharded_params(distributed_setup): + """The adapter should reuse the optimizer instance and keep sharded parameters.""" + 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,)) + model = TinyModel().to(device) + fully_shard(model.fc1, mesh=mesh, placements=_flat_placements()) + fully_shard(model.fc2, mesh=mesh, placements=_flat_placements()) + fsdp_parameters = tuple( + parameter + for parameter_group in (*model.fc1.parameter_groups(), *model.fc2.parameter_groups()) + for parameter in parameter_group.sharded_parameters + ) + + adam = torch.optim.Adam(model.parameters(), lr=0.01) + result = fully_shard_optimizer(adam) + + assert result is None + assert isinstance(adam, torch.optim.Adam) + + optimizer_parameters = tuple(adam.param_groups[0]["params"]) + assert len(optimizer_parameters) == len(fsdp_parameters) + for optimizer_parameter, fsdp_parameter in zip( + optimizer_parameters, fsdp_parameters, strict=True + ): + assert optimizer_parameter is fsdp_parameter + + +@pytest.mark.parametrize("optimizer_cls", [torch.optim.Adam, torch.optim.AdamW]) +def test_fully_shard_optimizer_adam_casts_mixed_precision_grads( + distributed_setup, optimizer_cls +): + """Adam optimizers should step on default mixed-precision FSDP parameters.""" + world_size = distributed_setup.world_size + device = distributed_setup.device + if world_size < 2: + pytest.skip("This test requires at least 2 ranks.") + if device.type != "cuda": + pytest.skip("Adam mixed-precision optimizer coverage requires CUDA.") + + 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()) + + fsdp_parameter_groups = (*model.fc1.parameter_groups(), *model.fc2.parameter_groups()) + fsdp_parameters = tuple( + parameter + for parameter_group in fsdp_parameter_groups + for parameter in parameter_group.sharded_parameters + ) + for parameter_group in fsdp_parameter_groups: + assert parameter_group.main_grad is not None + assert parameter_group.main_grad.dtype == torch.bfloat16 + for parameter in fsdp_parameters: + assert parameter.dtype == torch.float32 + + extra_parameter = nn.Parameter(torch.ones((), device=device)) + optimizer = optimizer_cls([{"params": model.parameters()}, {"params": [extra_parameter]}], lr=0.01) + fully_shard_optimizer(optimizer) + + optimizer_fsdp_parameters = tuple(optimizer.param_groups[0]["params"]) + assert len(optimizer_fsdp_parameters) == len(fsdp_parameters) + for optimizer_parameter, fsdp_parameter in zip( + optimizer_fsdp_parameters, fsdp_parameters, strict=True + ): + assert optimizer_parameter is fsdp_parameter + + x = torch.randn(6, 8, device=device, dtype=torch.bfloat16) + target = torch.randn(6, 4, device=device, dtype=torch.bfloat16) + losses = [] + optimizer.zero_grad(set_to_none=True) + + for _ in range(3): + loss = torch.nn.functional.mse_loss(model(x).float(), target.float()) + losses.append(loss.detach()) + loss.backward() + + for parameter in fsdp_parameters: + assert parameter.grad is not None + assert parameter.grad.dtype == torch.bfloat16 + assert parameter.grad.dtype != parameter.dtype + assert parameter.grad_dtype == torch.bfloat16 + + optimizer.step() + + for parameter in fsdp_parameters: + assert parameter.grad is not None + assert parameter.grad.dtype == torch.bfloat16 + assert parameter.grad_dtype == torch.bfloat16 + for parameter_group in fsdp_parameter_groups: + assert parameter_group.main_grad is not None + assert parameter_group.main_grad.dtype == torch.bfloat16 + + extra_parameter.grad = torch.ones_like(extra_parameter) + optimizer.zero_grad(set_to_none=False) + + for parameter in fsdp_parameters: + assert parameter.grad is not None + torch.testing.assert_close( + parameter.grad.to_local(), torch.zeros_like(parameter.grad.to_local()) + ) + assert parameter.grad_dtype == torch.bfloat16 + assert extra_parameter.grad is not None + torch.testing.assert_close(extra_parameter.grad, torch.zeros_like(extra_parameter)) + + with pytest.raises(AssertionError): + torch.testing.assert_close(losses[-1], losses[0]) + def test_microbatch_scopes_child_contexts(distributed_setup): """microbatch() should scope FSDP child contexts under an unwrapped parent.""" From a74ec3f918f45093323e6bf3ea48a18b2b18aef3 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Sat, 27 Jun 2026 06:18:47 +0000 Subject: [PATCH 02/12] Refine experimental fully_shard_optimizer adapter Signed-off-by: Jingyue Wu --- .../megatron_fsdp/experimental/optimizer.py | 44 +++++++------------ 1 file changed, 16 insertions(+), 28 deletions(-) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/optimizer.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/optimizer.py index 409dfbb9d8f..d289485f44d 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/optimizer.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/optimizer.py @@ -14,18 +14,15 @@ """Optimizer adapter for the minimal Megatron-FSDP path.""" -from dataclasses import dataclass -from typing import Any, TypeVar +from typing import Any, NamedTuple import torch from torch import nn from .parameter_group import contained_in_parameter_group -_OptimizerT = TypeVar("_OptimizerT", bound=torch.optim.Optimizer) - -def fully_shard_optimizer(optimizer: _OptimizerT) -> None: +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 @@ -36,17 +33,19 @@ def fully_shard_optimizer(optimizer: _OptimizerT) -> None: optimizer: Optimizer instance to adapt in place. """ - if not isinstance(optimizer, torch.optim.Optimizer): - raise TypeError( - "fully_shard_optimizer expected a torch.optim.Optimizer instance, " - f"got {optimizer!r}." - ) - - @dataclass - class CastedGrad: + class CastedGrad(NamedTuple): 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( @@ -79,15 +78,8 @@ def step_pre_hook( if parameter.grad.dtype == parameter.dtype: continue - original_grad = parameter.grad - casted_grads.append(CastedGrad(parameter, original_grad)) - - # 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 = parameter.dtype - parameter.grad = original_grad.to(dtype=parameter.dtype) + 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, @@ -95,12 +87,8 @@ def step_post_hook( kwargs: dict[str, Any], ) -> None: del hooked_optimizer, args, kwargs - for casted_grad in casted_grads: - parameter = casted_grad.parameter - original_grad = casted_grad.original_grad - parameter.grad = None - parameter.grad_dtype = original_grad.dtype - parameter.grad = original_grad + for parameter, original_grad in casted_grads: + set_grad(parameter, original_grad) casted_grads.clear() optimizer.register_step_pre_hook(step_pre_hook) From f1a5934703315d99e1411a83d32897b4e294283a Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Sat, 27 Jun 2026 06:18:47 +0000 Subject: [PATCH 03/12] Refine experimental FSDP optimizer behavior tests Signed-off-by: Jingyue Wu --- .../distributed/mfsdp_v2/test_fully_shard.py | 120 +++++------------- 1 file changed, 33 insertions(+), 87 deletions(-) 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 ffe95c24c5e..b43df60aa8c 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py @@ -113,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,120 +466,66 @@ def train_iteration() -> torch.Tensor: with pytest.raises(AssertionError): torch.testing.assert_close(second_loss, first_loss) -def test_fully_shard_optimizer_keeps_optimizer_instance_and_sharded_params(distributed_setup): - """The adapter should reuse the optimizer instance and keep sharded parameters.""" +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,)) - model = TinyModel().to(device) + 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()) - fsdp_parameters = tuple( - parameter - for parameter_group in (*model.fc1.parameter_groups(), *model.fc2.parameter_groups()) - for parameter in parameter_group.sharded_parameters - ) - adam = torch.optim.Adam(model.parameters(), lr=0.01) - result = fully_shard_optimizer(adam) + baseline_optimizer = torch.optim.Adam(baseline.parameters(), lr=0.01) + optimizer = torch.optim.Adam(model.parameters(), lr=0.01) + fully_shard_optimizer(optimizer) - assert result is None - assert isinstance(adam, torch.optim.Adam) + x = torch.randn(3, 8, device=device, dtype=torch.bfloat16) + target = torch.randn(3, 4, device=device, dtype=torch.bfloat16) - optimizer_parameters = tuple(adam.param_groups[0]["params"]) - assert len(optimizer_parameters) == len(fsdp_parameters) - for optimizer_parameter, fsdp_parameter in zip( - optimizer_parameters, fsdp_parameters, strict=True - ): - assert optimizer_parameter is fsdp_parameter + 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() -@pytest.mark.parametrize("optimizer_cls", [torch.optim.Adam, torch.optim.AdamW]) -def test_fully_shard_optimizer_adam_casts_mixed_precision_grads( - distributed_setup, optimizer_cls -): - """Adam optimizers should step on default mixed-precision FSDP parameters.""" +def test_fully_shard_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 if world_size < 2: pytest.skip("This test requires at least 2 ranks.") - if device.type != "cuda": - pytest.skip("Adam mixed-precision optimizer coverage requires CUDA.") - 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()) - - fsdp_parameter_groups = (*model.fc1.parameter_groups(), *model.fc2.parameter_groups()) - fsdp_parameters = tuple( - parameter - for parameter_group in fsdp_parameter_groups - for parameter in parameter_group.sharded_parameters - ) - for parameter_group in fsdp_parameter_groups: - assert parameter_group.main_grad is not None - assert parameter_group.main_grad.dtype == torch.bfloat16 - for parameter in fsdp_parameters: - assert parameter.dtype == torch.float32 - - extra_parameter = nn.Parameter(torch.ones((), device=device)) - optimizer = optimizer_cls([{"params": model.parameters()}, {"params": [extra_parameter]}], lr=0.01) - fully_shard_optimizer(optimizer) - - optimizer_fsdp_parameters = tuple(optimizer.param_groups[0]["params"]) - assert len(optimizer_fsdp_parameters) == len(fsdp_parameters) - for optimizer_parameter, fsdp_parameter in zip( - optimizer_fsdp_parameters, fsdp_parameters, strict=True - ): - assert optimizer_parameter is fsdp_parameter + optimizer = torch.optim.Adam(model.parameters(), lr=0.01) x = torch.randn(6, 8, device=device, dtype=torch.bfloat16) target = torch.randn(6, 4, device=device, dtype=torch.bfloat16) - losses = [] optimizer.zero_grad(set_to_none=True) + loss = torch.nn.functional.mse_loss(model(x).float(), target.float()) + loss.backward() - for _ in range(3): - loss = torch.nn.functional.mse_loss(model(x).float(), target.float()) - losses.append(loss.detach()) - loss.backward() - - for parameter in fsdp_parameters: - assert parameter.grad is not None - assert parameter.grad.dtype == torch.bfloat16 - assert parameter.grad.dtype != parameter.dtype - assert parameter.grad_dtype == torch.bfloat16 - + with pytest.raises( + RuntimeError, + match="same device and the same dtype", + ): optimizer.step() - for parameter in fsdp_parameters: - assert parameter.grad is not None - assert parameter.grad.dtype == torch.bfloat16 - assert parameter.grad_dtype == torch.bfloat16 - for parameter_group in fsdp_parameter_groups: - assert parameter_group.main_grad is not None - assert parameter_group.main_grad.dtype == torch.bfloat16 - - extra_parameter.grad = torch.ones_like(extra_parameter) - optimizer.zero_grad(set_to_none=False) - - for parameter in fsdp_parameters: - assert parameter.grad is not None - torch.testing.assert_close( - parameter.grad.to_local(), torch.zeros_like(parameter.grad.to_local()) - ) - assert parameter.grad_dtype == torch.bfloat16 - assert extra_parameter.grad is not None - torch.testing.assert_close(extra_parameter.grad, torch.zeros_like(extra_parameter)) - - with pytest.raises(AssertionError): - torch.testing.assert_close(losses[-1], losses[0]) - - def test_microbatch_scopes_child_contexts(distributed_setup): """microbatch() should scope FSDP child contexts under an unwrapped parent.""" world_size = distributed_setup.world_size From 3dce5f4e318b2e2ea9aa1a98ff7612856df79549 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Sat, 27 Jun 2026 06:18:47 +0000 Subject: [PATCH 04/12] Document fully_shard_optimizer adapter alternatives Signed-off-by: Jingyue Wu --- .../fsdp/src/megatron_fsdp/experimental/optimizer.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/optimizer.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/optimizer.py index d289485f44d..b3aee69abc0 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/optimizer.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/optimizer.py @@ -29,9 +29,15 @@ def fully_shard_optimizer(optimizer: torch.optim.Optimizer) -> None: 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. + Args: optimizer: Optimizer instance to adapt in place. - """ class CastedGrad(NamedTuple): parameter: nn.Parameter From a211993c28593a8283b0a7cb37518dedd1e09b40 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Sat, 27 Jun 2026 06:18:47 +0000 Subject: [PATCH 05/12] Fix fully_shard_optimizer lint formatting Signed-off-by: Jingyue Wu --- .../fsdp/src/megatron_fsdp/experimental/optimizer.py | 11 +++++------ .../distributed/mfsdp_v2/test_fully_shard.py | 7 +++---- 2 files changed, 8 insertions(+), 10 deletions(-) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/optimizer.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/optimizer.py index b3aee69abc0..5434d6359f5 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/optimizer.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/optimizer.py @@ -39,7 +39,10 @@ def fully_shard_optimizer(optimizer: torch.optim.Optimizer) -> None: 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 @@ -55,9 +58,7 @@ def set_grad(parameter: nn.Parameter, grad: torch.Tensor) -> None: casted_grads: list[CastedGrad] = [] def step_pre_hook( - hooked_optimizer: torch.optim.Optimizer, - args: tuple[Any, ...], - kwargs: dict[str, Any], + 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: @@ -88,9 +89,7 @@ def step_pre_hook( 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], + 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: 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 b43df60aa8c..45733e88f60 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py @@ -466,6 +466,7 @@ def train_iteration() -> torch.Tensor: with pytest.raises(AssertionError): 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 @@ -520,12 +521,10 @@ def test_fully_shard_adam_without_adapter_raises_precision_error(distributed_set loss = torch.nn.functional.mse_loss(model(x).float(), target.float()) loss.backward() - with pytest.raises( - RuntimeError, - match="same device and the same dtype", - ): + with pytest.raises(RuntimeError, match="same device and the same dtype"): 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 From 98615f21ddab3d3b5a32db6d7dc152ff7d5e82bf Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Sat, 27 Jun 2026 06:18:47 +0000 Subject: [PATCH 06/12] Document optimizer precision alternatives Signed-off-by: Jingyue Wu --- .../fsdp/src/megatron_fsdp/experimental/optimizer.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/optimizer.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/optimizer.py index 5434d6359f5..3f7745da9ff 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/optimizer.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/optimizer.py @@ -35,6 +35,14 @@ def fully_shard_optimizer(optimizer: torch.optim.Optimizer) -> None: - 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. From e283ececcf8cee989403b57a27cc3ad1150cb508 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Sat, 27 Jun 2026 06:22:02 +0000 Subject: [PATCH 07/12] Move optimizer failure test Signed-off-by: Jingyue Wu --- .../distributed/mfsdp_v2/test_fully_shard.py | 23 -------- .../distributed/mfsdp_v2/test_optimizer.py | 55 +++++++++++++++++++ 2 files changed, 55 insertions(+), 23 deletions(-) create mode 100644 tests/unit_tests/distributed/mfsdp_v2/test_optimizer.py 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 45733e88f60..896f834a7f7 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py @@ -502,29 +502,6 @@ def test_fully_shard_adam_mixed_precision_losses_match_baseline(distributed_setu optimizer.step() -def test_fully_shard_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 - 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) - 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) - target = torch.randn(6, 4, device=device, dtype=torch.bfloat16) - optimizer.zero_grad(set_to_none=True) - loss = torch.nn.functional.mse_loss(model(x).float(), target.float()) - loss.backward() - - with pytest.raises(RuntimeError, match="same device and the same dtype"): - 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..6facf8d7239 --- /dev/null +++ b/tests/unit_tests/distributed/mfsdp_v2/test_optimizer.py @@ -0,0 +1,55 @@ +# 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 megatron.core.distributed.fsdp.src.megatron_fsdp.experimental import ( + Flat, + Placements, + fully_shard, +) + + +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_fully_shard_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 + 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) + 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) + target = torch.randn(6, 4, device=device, dtype=torch.bfloat16) + optimizer.zero_grad(set_to_none=True) + loss = torch.nn.functional.mse_loss(model(x).float(), target.float()) + loss.backward() + + with pytest.raises(RuntimeError, match="same device and the same dtype"): + optimizer.step() From 8fc3947a256c1400b0e22efc198ffc948638a277 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Sat, 27 Jun 2026 22:59:12 +0000 Subject: [PATCH 08/12] Rename Adam optimizer failure test Signed-off-by: Jingyue Wu --- tests/unit_tests/distributed/mfsdp_v2/test_optimizer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_optimizer.py b/tests/unit_tests/distributed/mfsdp_v2/test_optimizer.py index 6facf8d7239..d3f8de5a291 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_optimizer.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_optimizer.py @@ -32,7 +32,7 @@ def _flat_placements() -> Placements: return Placements(dp_axes=[0], parameter=[Flat()], gradient=[Flat()], optimizer=[Flat()]) -def test_fully_shard_adam_without_adapter_raises_precision_error(distributed_setup): +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 From 239fddef029828b7f149d9f17b31cff1ea8970a6 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Tue, 7 Jul 2026 21:54:51 +0000 Subject: [PATCH 09/12] Add FusedAdam MFSDP optimizer test Signed-off-by: Jingyue Wu --- .../distributed/mfsdp_v2/test_optimizer.py | 54 +++++++++++++++++-- 1 file changed, 50 insertions(+), 4 deletions(-) diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_optimizer.py b/tests/unit_tests/distributed/mfsdp_v2/test_optimizer.py index d3f8de5a291..24f5cf94416 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_optimizer.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_optimizer.py @@ -6,12 +6,14 @@ 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): @@ -36,8 +38,6 @@ 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 - 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) model = TinyModel().to(device=device, dtype=torch.bfloat16) @@ -46,10 +46,56 @@ def test_adam_without_adapter_raises_precision_error(distributed_setup): optimizer = torch.optim.Adam(model.parameters(), lr=0.01) x = torch.randn(6, 8, device=device, dtype=torch.bfloat16) - target = torch.randn(6, 4, device=device, dtype=torch.bfloat16) optimizer.zero_grad(set_to_none=True) - loss = torch.nn.functional.mse_loss(model(x).float(), target.float()) + 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_param_grad_dtype_mismatch(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 == torch.float32 + assert parameter.grad.dtype == torch.bfloat16 + + 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()) + ) From 71a086679682bd7615ef550e9777f51c8cbd5c3d Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Tue, 7 Jul 2026 23:14:01 +0000 Subject: [PATCH 10/12] Test MFSDP optimizer grad dtype mismatch Signed-off-by: Jingyue Wu --- .../distributed/mfsdp_v2/test_optimizer.py | 123 +++++++++++++----- 1 file changed, 94 insertions(+), 29 deletions(-) diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_optimizer.py b/tests/unit_tests/distributed/mfsdp_v2/test_optimizer.py index 24f5cf94416..0c5af8ede1d 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_optimizer.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_optimizer.py @@ -5,7 +5,7 @@ import pytest import torch from torch import nn -from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.device_mesh import DeviceMesh, init_device_mesh from transformer_engine.pytorch.optimizers import FusedAdam from megatron.core.distributed.fsdp.src.megatron_fsdp.experimental import ( @@ -14,6 +14,8 @@ fully_shard, ) from megatron.core.distributed.fsdp.src.megatron_fsdp.mixed_precision import MixedPrecisionPolicy +from megatron.core.optimizer.optimizer import MixedPrecisionOptimizer +from megatron.core.optimizer.optimizer_config import OptimizerConfig class TinyModel(nn.Module): @@ -30,43 +32,55 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return self.fc2(self.relu(self.fc1(x))) -def _flat_placements() -> Placements: - return Placements(dp_axes=[0], parameter=[Flat()], gradient=[Flat()], optimizer=[Flat()]) +class ParamGradCastingMixedPrecisionOptimizer(MixedPrecisionOptimizer): + """Test optimizer that makes param.grad dtype-compatible with param.""" + def __init__(self, optimizer: torch.optim.Optimizer, config: OptimizerConfig) -> None: + super().__init__(optimizer, config, grad_scaler=None, init_state_fn=lambda *_args: None) + self.is_stub_optimizer = False -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) + def zero_grad(self, set_to_none: bool = True) -> None: + self.optimizer.zero_grad(set_to_none=set_to_none) - x = torch.randn(6, 8, device=device, dtype=torch.bfloat16) - optimizer.zero_grad(set_to_none=True) - loss = model(x).sum() - loss.backward() + def prepare_grads(self) -> bool: + for parameter in self.get_parameters(): + grad = parameter.grad.to(dtype=parameter.dtype) + parameter.grad = None + parameter.grad_dtype = grad.dtype + parameter.grad = grad + # Return False because this test optimizer does not check for inf/nan. + return False - with pytest.raises(RuntimeError, match="same device and the same dtype"): - optimizer.step() + def step_with_ready_grads(self) -> bool: + self.optimizer.step() + return True + def reload_model_params(self, state_dict=None) -> None: + pass -def test_fused_adam_without_adapter_accepts_param_grad_dtype_mismatch(distributed_setup): - """TE FusedAdam should handle mixed-precision FSDP grads without the adapter.""" - world_size = distributed_setup.world_size - device = distributed_setup.device + def state_dict(self): + raise NotImplementedError - mesh = init_device_mesh(device.type, (world_size,)) + def load_state_dict(self, state_dict) -> None: + raise NotImplementedError + + def sharded_state_dict(self, *args, **kwargs): + raise NotImplementedError + + +def _flat_placements() -> Placements: + return Placements(dp_axes=[0], parameter=[Flat()], gradient=[Flat()], optimizer=[Flat()]) + + +def _build_model_with_param_grad_dtype_mismatch( + mesh: DeviceMesh, device: torch.device +) -> TinyModel: 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, + main_params_dtype=torch.float32, main_grads_dtype=torch.bfloat16 ) fully_shard( model.fc1, @@ -80,17 +94,68 @@ def test_fused_adam_without_adapter_accepts_param_grad_dtype_mismatch(distribute placements=_flat_placements(), mixed_precision_policy=mixed_precision_policy, ) - optimizer = FusedAdam(model.parameters(), lr=0.01) + return model + + +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,)) + model = _build_model_with_param_grad_dtype_mismatch(mesh, device) + optimizer = torch.optim.Adam(model.parameters(), lr=0.01) + optimizer.zero_grad(set_to_none=True) x = torch.randn(6, 8, device=device, dtype=torch.bfloat16) + 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,)) + model = _build_model_with_param_grad_dtype_mismatch(mesh, device) + optimizer = FusedAdam(model.parameters(), lr=0.01) + optimizer.zero_grad(set_to_none=True) + x = torch.randn(6, 8, device=device, dtype=torch.bfloat16) 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()) + ) + + +def test_mixed_precision_optimizer_with_adam_casts_mismatched_grads(distributed_setup): + """A MixedPrecisionOptimizer subclass can make vanilla Adam dtype-compatible.""" + world_size = distributed_setup.world_size + device = distributed_setup.device + mesh = init_device_mesh(device.type, (world_size,)) + model = _build_model_with_param_grad_dtype_mismatch(mesh, device) + base_optimizer = torch.optim.Adam(model.parameters(), lr=0.01) + optimizer_config = OptimizerConfig(optimizer="adam", lr=0.01, bf16=True, clip_grad=0.0) + optimizer = ParamGradCastingMixedPrecisionOptimizer(base_optimizer, optimizer_config) + + optimizer.zero_grad(set_to_none=True) + x = torch.randn(6, 8, device=device, dtype=torch.bfloat16) + loss = model(x).sum() + loss.backward() for parameter in model.parameters(): assert parameter.grad is not None - assert parameter.dtype == torch.float32 - assert parameter.grad.dtype == torch.bfloat16 + assert parameter.dtype != parameter.grad.dtype params_before_step = [parameter.detach().clone() for parameter in model.parameters()] optimizer.step() From dfbce460a537444defcaaa8f2e47ebf82a2a2a23 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Wed, 8 Jul 2026 00:22:05 +0000 Subject: [PATCH 11/12] Remove custom MFSDP optimizer test Signed-off-by: Jingyue Wu --- .../distributed/mfsdp_v2/test_optimizer.py | 117 ++++-------------- 1 file changed, 25 insertions(+), 92 deletions(-) diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_optimizer.py b/tests/unit_tests/distributed/mfsdp_v2/test_optimizer.py index 0c5af8ede1d..b2d01442233 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_optimizer.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_optimizer.py @@ -5,7 +5,7 @@ import pytest import torch from torch import nn -from torch.distributed.device_mesh import DeviceMesh, init_device_mesh +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 ( @@ -14,8 +14,6 @@ fully_shard, ) from megatron.core.distributed.fsdp.src.megatron_fsdp.mixed_precision import MixedPrecisionPolicy -from megatron.core.optimizer.optimizer import MixedPrecisionOptimizer -from megatron.core.optimizer.optimizer_config import OptimizerConfig class TinyModel(nn.Module): @@ -32,49 +30,36 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return self.fc2(self.relu(self.fc1(x))) -class ParamGradCastingMixedPrecisionOptimizer(MixedPrecisionOptimizer): - """Test optimizer that makes param.grad dtype-compatible with param.""" - - def __init__(self, optimizer: torch.optim.Optimizer, config: OptimizerConfig) -> None: - super().__init__(optimizer, config, grad_scaler=None, init_state_fn=lambda *_args: None) - self.is_stub_optimizer = False - - def zero_grad(self, set_to_none: bool = True) -> None: - self.optimizer.zero_grad(set_to_none=set_to_none) - - def prepare_grads(self) -> bool: - for parameter in self.get_parameters(): - grad = parameter.grad.to(dtype=parameter.dtype) - parameter.grad = None - parameter.grad_dtype = grad.dtype - parameter.grad = grad - # Return False because this test optimizer does not check for inf/nan. - return False - - def step_with_ready_grads(self) -> bool: - self.optimizer.step() - return True - - def reload_model_params(self, state_dict=None) -> None: - pass +def _flat_placements() -> Placements: + return Placements(dp_axes=[0], parameter=[Flat()], gradient=[Flat()], optimizer=[Flat()]) - def state_dict(self): - raise NotImplementedError - def load_state_dict(self, state_dict) -> None: - raise NotImplementedError +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) - def sharded_state_dict(self, *args, **kwargs): - raise NotImplementedError + 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 _flat_placements() -> Placements: - return Placements(dp_axes=[0], parameter=[Flat()], gradient=[Flat()], optimizer=[Flat()]) +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 -def _build_model_with_param_grad_dtype_mismatch( - mesh: DeviceMesh, device: torch.device -) -> TinyModel: + 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 @@ -94,65 +79,13 @@ def _build_model_with_param_grad_dtype_mismatch( placements=_flat_placements(), mixed_precision_policy=mixed_precision_policy, ) - return model - - -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,)) - model = _build_model_with_param_grad_dtype_mismatch(mesh, device) - optimizer = torch.optim.Adam(model.parameters(), lr=0.01) - - optimizer.zero_grad(set_to_none=True) - x = torch.randn(6, 8, device=device, dtype=torch.bfloat16) - 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,)) - model = _build_model_with_param_grad_dtype_mismatch(mesh, device) optimizer = FusedAdam(model.parameters(), lr=0.01) - optimizer.zero_grad(set_to_none=True) x = torch.randn(6, 8, device=device, dtype=torch.bfloat16) - 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()) - ) - - -def test_mixed_precision_optimizer_with_adam_casts_mismatched_grads(distributed_setup): - """A MixedPrecisionOptimizer subclass can make vanilla Adam dtype-compatible.""" - world_size = distributed_setup.world_size - device = distributed_setup.device - mesh = init_device_mesh(device.type, (world_size,)) - model = _build_model_with_param_grad_dtype_mismatch(mesh, device) - base_optimizer = torch.optim.Adam(model.parameters(), lr=0.01) - optimizer_config = OptimizerConfig(optimizer="adam", lr=0.01, bf16=True, clip_grad=0.0) - optimizer = ParamGradCastingMixedPrecisionOptimizer(base_optimizer, optimizer_config) - optimizer.zero_grad(set_to_none=True) - x = torch.randn(6, 8, device=device, dtype=torch.bfloat16) loss = model(x).sum() loss.backward() + for parameter in model.parameters(): assert parameter.grad is not None assert parameter.dtype != parameter.grad.dtype From 542c84ae3992a30f29f5ce77e6867cd44958194f Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Mon, 20 Jul 2026 21:44:41 +0000 Subject: [PATCH 12/12] Relax MFSDP optimizer dtype assertion Signed-off-by: Jingyue Wu --- tests/unit_tests/distributed/mfsdp_v2/test_optimizer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_optimizer.py b/tests/unit_tests/distributed/mfsdp_v2/test_optimizer.py index b2d01442233..3f3b636dd19 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_optimizer.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_optimizer.py @@ -50,7 +50,7 @@ def test_adam_without_adapter_raises_precision_error(distributed_setup): loss = model(x).sum() loss.backward() - with pytest.raises(RuntimeError, match="same device and the same dtype"): + with pytest.raises(RuntimeError, match="dtype"): optimizer.step()