From af7f475b6d983a8498b76aa1788348ceb25859aa Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Fri, 7 Aug 2026 04:16:24 +0000 Subject: [PATCH 1/2] fix(mfsdp): refresh V2 compute weights after optimizer step Signed-off-by: Jingyue Wu --- megatron/core/optimizer/fully_sharded_optimizer.py | 14 +++++++++++++- .../distributed/mfsdp_v2/test_mcore_adapter.py | 4 ++-- 2 files changed, 15 insertions(+), 3 deletions(-) diff --git a/megatron/core/optimizer/fully_sharded_optimizer.py b/megatron/core/optimizer/fully_sharded_optimizer.py index 18c2354dcb2..b15deafa961 100644 --- a/megatron/core/optimizer/fully_sharded_optimizer.py +++ b/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 ( + get_containing_parameter_group, +) from ..transformer.module import MegatronModule from .grad_scaler import MegatronGradScaler from .optimizer import MixedPrecisionOptimizer @@ -120,7 +123,16 @@ 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.""" + # TODO(wujingyue): Reuse experimental.checkpoint.resync_compute_weights after #6024 merges. + seen_parameter_groups = set() + for parameter in self.get_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() def _copy_model_params_to_main_params(self, state_dict=None) -> None: """No-op: model loads already write into MFSDP v2's main weights.""" diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_mcore_adapter.py b/tests/unit_tests/distributed/mfsdp_v2/test_mcore_adapter.py index a132ff88139..18f236d6bb6 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_mcore_adapter.py +++ b/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) From 296a795e6940a95a91886166aba9d8e3493f879f Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Mon, 10 Aug 2026 15:28:39 +0000 Subject: [PATCH 2/2] refactor(mfsdp): share compute-weight synchronization Extract sync_model_weights_from_main_weights() so the checkpoint load path and the fully_shard_optimizer() step post-hook refresh compute weights through one implementation instead of two. The post-hook previously collected parameter groups into a set and iterated it. FsdpParameterGroup is @dataclass(frozen=True, eq=False), so it hashes by id() and set iteration order varies per process. sync_model_weight_from_main_weight() issues a collective whenever main_weight and model_weight are distinct buffers, so ranks could issue the per-group collectives in different orders and deadlock. Reusing the shared helper fixes that: it iterates the parameters, whose order is identical on every rank, and uses a set only to skip groups already synced. Reproduced with an HFSDP configuration (optimizer [Flat, Flat], parameter [Replicate, Flat]) over three parameter groups on 4 GPUs, which hung until this change; it now completes in ~9s. Signed-off-by: Jingyue Wu Co-Authored-By: Claude Opus 5 (1M context) --- .../megatron_fsdp/experimental/checkpoint.py | 21 ++----------------- .../megatron_fsdp/experimental/optimizer.py | 16 +++++--------- .../experimental/parameter_group.py | 17 +++++++++++++++ .../core/optimizer/fully_sharded_optimizer.py | 12 ++--------- 4 files changed, 26 insertions(+), 40 deletions(-) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/checkpoint.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/checkpoint.py index c9dc44b04db..7502cedb93a 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/checkpoint.py +++ b/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()) 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 f1617141569..c5e4b28f211 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/optimizer.py +++ b/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) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py index 6728af5691e..251af7509df 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/parameter_group.py +++ b/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.""" diff --git a/megatron/core/optimizer/fully_sharded_optimizer.py b/megatron/core/optimizer/fully_sharded_optimizer.py index b15deafa961..9c70c5b6f86 100644 --- a/megatron/core/optimizer/fully_sharded_optimizer.py +++ b/megatron/core/optimizer/fully_sharded_optimizer.py @@ -9,7 +9,7 @@ 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 ( - get_containing_parameter_group, + sync_model_weights_from_main_weights, ) from ..transformer.module import MegatronModule from .grad_scaler import MegatronGradScaler @@ -124,15 +124,7 @@ def _copy_model_grads_to_main_grads(self) -> None: def _copy_main_params_to_model_params(self) -> None: """Refresh MFSDP V2 compute weights after updating optimizer weights.""" - # TODO(wujingyue): Reuse experimental.checkpoint.resync_compute_weights after #6024 merges. - seen_parameter_groups = set() - for parameter in self.get_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() + 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."""