From 2812aabfcbf2b319d6622441eb7fd1a42c9c64db Mon Sep 17 00:00:00 2001 From: tongliu Date: Mon, 11 May 2026 01:21:08 -0700 Subject: [PATCH] refactor(mfsdp): simplify mixed precision policy defaults --- .../distributed/fsdp/mcore_fsdp_adapter.py | 16 +++++++++++-- .../fully_shard_rewrite/dp_buffer.py | 8 +++---- .../fully_shard_rewrite/fully_shard.py | 6 +---- .../fully_shard_rewrite/mixed_precision.py | 24 +++---------------- 4 files changed, 21 insertions(+), 33 deletions(-) diff --git a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py b/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py index 977f0419ba4..80289f0bde3 100644 --- a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py +++ b/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py @@ -53,7 +53,7 @@ ) from megatron.core.distributed.fsdp.src.megatron_fsdp.fully_shard_rewrite import FSDPModule from megatron.core.distributed.fsdp.src.megatron_fsdp.fully_shard_rewrite.mixed_precision import ( - build_fully_shard_mixed_precision_policy, + FullyShardMixedPrecisionPolicy, ) HAVE_MEGATRON_FSDP = True @@ -241,7 +241,19 @@ def _init_with_fully_shard( edp_mesh = _init_dp_mesh(pg_collection, edp=True) dp_mesh = _init_dp_mesh(pg_collection, edp=False) - fully_shard_mp_policy = build_fully_shard_mixed_precision_policy(ddp_config) + fully_shard_mp_policy = FullyShardMixedPrecisionPolicy( + main_params_dtype=ddp_config.megatron_fsdp_main_params_dtype, + main_grads_dtype=( + torch.float32 + if ddp_config.grad_reduce_in_fp32 + else ddp_config.megatron_fsdp_main_grads_dtype + ), + grad_comm_dtype=( + torch.float32 + if ddp_config.grad_reduce_in_fp32 + else ddp_config.megatron_fsdp_grad_comm_dtype + ), + ) kwargs = { "mp_policy": fully_shard_mp_policy, "enable_unshard_prefetch": ddp_config.overlap_param_gather, diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/fully_shard_rewrite/dp_buffer.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/fully_shard_rewrite/dp_buffer.py index 9bf7c8a7e15..5c2f6cb7a8c 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/fully_shard_rewrite/dp_buffer.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/fully_shard_rewrite/dp_buffer.py @@ -549,11 +549,9 @@ def reduce_grad(self, grad_comm_dtype: Optional[torch.dtype] = None): output=reduced_grad_shard, input=comm_input, group=self.dp_group, op=op ) - # Accumulate the reduced shard into the persistent local shard buffer. The - # reduce-scatter output must not alias the full input buffer, otherwise the - # collective can clobber its own input and silently corrupt gradients. - if reduced_grad_shard.dtype != self.dtype: - reduced_grad_shard = reduced_grad_shard.to(self.dtype) + # Accumulate into the persistent local shard buffer. The reduce-scatter output + # must not alias the full input buffer, otherwise the collective can clobber its + # own input and silently corrupt gradients. local_grad_shard += reduced_grad_shard diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/fully_shard_rewrite/fully_shard.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/fully_shard_rewrite/fully_shard.py index c2ff6ce5207..4452d1b321a 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/fully_shard_rewrite/fully_shard.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/fully_shard_rewrite/fully_shard.py @@ -56,11 +56,7 @@ def fully_shard( "Please do not call fully_shard on the same module more than once." ) if mp_policy is None: - mp_policy = FullyShardMixedPrecisionPolicy( - main_params_dtype=None, - main_grads_dtype=None, - grad_comm_dtype=None, - ) + mp_policy = FullyShardMixedPrecisionPolicy() cls = module.__class__ new_cls = type(f"FSDP{cls.__name__}", (FSDPModule, cls), {}) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/fully_shard_rewrite/mixed_precision.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/fully_shard_rewrite/mixed_precision.py index 71db4aaabab..c9e7c43c3e6 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/fully_shard_rewrite/mixed_precision.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/fully_shard_rewrite/mixed_precision.py @@ -2,9 +2,8 @@ """Mixed precision policy helpers for the Megatron-FSDP fully_shard rewrite path. -The adapter builds the v2 policy here and passes it to ``fully_shard``. Keep -FP8-specific policy decisions in this module instead of spreading them through -the adapter or ``ParameterGroup``. +This module owns the v2 policy data model. Translation from Megatron/MCore +config objects belongs in the adapter layer. """ from dataclasses import dataclass @@ -17,23 +16,6 @@ class FullyShardMixedPrecisionPolicy: """Mixed precision dtype policy owned by the v2 ``fully_shard`` path.""" - main_params_dtype: Optional[torch.dtype] = torch.float32 + main_params_dtype: Optional[torch.dtype] = None main_grads_dtype: Optional[torch.dtype] = None grad_comm_dtype: Optional[torch.dtype] = None - - -def build_fully_shard_mixed_precision_policy(ddp_config) -> FullyShardMixedPrecisionPolicy: - """Build the v2 mixed precision policy from Megatron's DDP/FSDP config.""" - - if ddp_config.grad_reduce_in_fp32: - main_grads_dtype = torch.float32 - grad_comm_dtype = torch.float32 - else: - main_grads_dtype = ddp_config.megatron_fsdp_main_grads_dtype - grad_comm_dtype = ddp_config.megatron_fsdp_grad_comm_dtype - - return FullyShardMixedPrecisionPolicy( - main_params_dtype=ddp_config.megatron_fsdp_main_params_dtype, - main_grads_dtype=main_grads_dtype, - grad_comm_dtype=grad_comm_dtype, - )