refactor(mfsdp): simplify mixed precision policy defaults

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,

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
 
 

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), {})

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,
-    )