fix(mfsdp): refresh V2 compute weights after optimizer step

megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/checkpoint.py

@@ -41,28 +41,11 @@
 )
 
 from ..uneven_dtensor import preprocess_state_dict_for_uneven_dtensor
-from .module import FsdpModule
+from .parameter_group import sync_model_weights_from_main_weights
 
 __all__ = ["save_checkpoint", "load_checkpoint"]
 
 
-def _sync_model_weight_from_main_weight(model: torch.nn.Module) -> None:
-    """Refresh every FSDP group's compute weights from its (loaded) main weights.
-
-    A load writes into the ``main_weight``-backed sharded DTensors. When mixed precision keeps a
-    separate lower-precision compute buffer, that buffer is stale until the next forward pre-hook
-    would resync it; doing it here makes the post-load state deterministic. It is a no-op when the
-    compute buffer aliases the main buffer.
-
-    Args:
-        model: Root module (or any module tree) containing ``FsdpModule`` instances.
-    """
-    for module in model.modules():
-        if isinstance(module, FsdpModule):
-            for parameter_group in module.parameter_groups:
-                parameter_group.sync_model_weight_from_main_weight()
-
-
 def _init_optimizer_state(optimizer: torch.optim.Optimizer) -> None:
     """Allocate optimizer state so a DCP load has DTensors to fill.
 
@@ -139,4 +122,4 @@ def load_checkpoint(
     set_model_state_dict(model, model_state_dict)
     set_optimizer_state_dict(model, optimizer, optimizer_state_dict)
     if sync_model_weights:
-        _sync_model_weight_from_main_weight(model)
+        sync_model_weights_from_main_weights(model.parameters())

megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/optimizer.py

@@ -14,12 +14,13 @@
 
 """Optimizer adapter for the minimal Megatron-FSDP path."""
 
+from itertools import chain
 from typing import Any, NamedTuple
 
 import torch
 from torch import nn
 
-from .parameter_group import FsdpParameterGroup, get_containing_parameter_group
+from .parameter_group import get_containing_parameter_group, sync_model_weights_from_main_weights
 
 
 def fully_shard_optimizer(
@@ -109,16 +110,9 @@ def step_post_hook(
             set_grad(parameter, original_grad)
         casted_grads.clear()
 
-        fsdp_parameter_groups: set[FsdpParameterGroup] = set()
-        for optimizer_group in hooked_optimizer.param_groups:
-            for parameter in optimizer_group["params"]:
-                parameter_group = get_containing_parameter_group(parameter)
-                if parameter_group is None:
-                    continue
-                fsdp_parameter_groups.add(parameter_group)
-
-        for parameter_group in fsdp_parameter_groups:
-            parameter_group.sync_model_weight_from_main_weight()
+        sync_model_weights_from_main_weights(
+            chain.from_iterable(group["params"] for group in hooked_optimizer.param_groups)
+        )
 
     optimizer.register_step_pre_hook(step_pre_hook)
     optimizer.register_step_post_hook(step_post_hook)

megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py

@@ -14,6 +14,7 @@
 
 """Parameter-group runtime state for the minimal Megatron-FSDP path."""
 
+from collections.abc import Iterable
 from contextlib import nullcontext
 from dataclasses import dataclass
 from weakref import ReferenceType, ref
@@ -41,6 +42,22 @@ def get_containing_parameter_group(parameter: nn.Parameter) -> "FsdpParameterGro
     return parameter_group_ref()
 
 
+def sync_model_weights_from_main_weights(parameters: Iterable[nn.Parameter]) -> None:
+    """Refresh MFSDP compute weights for parameter groups represented by ``parameters``.
+
+    Parameters outside the experimental MFSDP path are ignored. A parameter group
+    may own multiple parameters, but its compute-weight buffer is refreshed once.
+    """
+    seen_parameter_groups = set()
+    for parameter in parameters:
+        if (parameter_group := get_containing_parameter_group(parameter)) is None:
+            continue
+        if parameter_group in seen_parameter_groups:
+            continue
+        seen_parameter_groups.add(parameter_group)
+        parameter_group.sync_model_weight_from_main_weight()
+
+
 @dataclass(frozen=True, eq=False)
 class FsdpParameter:
     """One physical parameter and its FSDP runtime representations."""

megatron/core/optimizer/fully_sharded_optimizer.py

@@ -8,6 +8,9 @@
 
 from ..config_logger import has_config_logger_enabled, log_config_to_disk
 from ..dist_checkpointing.mapping import ShardedStateDict
+from ..distributed.fsdp.src.megatron_fsdp.experimental.parameter_group import (
+    sync_model_weights_from_main_weights,
+)
 from ..transformer.module import MegatronModule
 from .grad_scaler import MegatronGradScaler
 from .optimizer import MixedPrecisionOptimizer
@@ -120,7 +123,8 @@ def _copy_model_grads_to_main_grads(self) -> None:
         """No-op: MFSDP v2 reduces directly into optimizer-visible sharded grads."""
 
     def _copy_main_params_to_model_params(self) -> None:
-        """No-op: MFSDP v2 currently syncs compute weights in its forward pre-hook."""
+        """Refresh MFSDP V2 compute weights after updating optimizer weights."""
+        sync_model_weights_from_main_weights(self.get_parameters())
 
     def _copy_model_params_to_main_params(self, state_dict=None) -> None:
         """No-op: model loads already write into MFSDP v2's main weights."""

tests/unit_tests/distributed/mfsdp_v2/test_mcore_adapter.py

@@ -164,7 +164,7 @@ def test_build_train_and_step(self):
                 torch.randn(8, 2, config.hidden_size, device="cuda", dtype=torch.bfloat16)
                 for _ in range(2)
             ]
-            for _ in range(3)
+            for _ in range(10)
         ]
 
         reference_losses = []
@@ -198,4 +198,4 @@ def test_build_train_and_step(self):
         reference_losses = torch.stack(reference_losses)
         assert torch.isfinite(losses).all()
         assert torch.isfinite(reference_losses).all()
-        torch.testing.assert_close(losses, reference_losses, rtol=1e-2, atol=0)
+        torch.testing.assert_close(losses, reference_losses, rtol=1e-3, atol=0)