fix(fsdp): drop stale MFSDP v2 guard rails for clip_grad, precision-aware optimizer, HSDP, and optimizer CUDA graphs (#5933)

Signed-off-by: Jingyue Wu <wujingyue@gmail.com>
Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

Author: wujingyue

src/megatron/bridge/models/conversion/utils.py

@@ -46,15 +46,18 @@ def unwrap_model(model, module_instances=None):
     if module_instances is None:
         from megatron.core.distributed import DistributedDataParallel as DDP
         from megatron.core.distributed import TorchFullyShardedDataParallel as torch_FSDP
+        from megatron.core.distributed.fsdp.mcore_fsdp_adapter import (
+            FullyShardedDataParallelV1,
+            FullyShardedDataParallelV2,
+        )
         from megatron.core.distributed.fsdp.src.megatron_fsdp.megatron_fsdp import MegatronFSDP
         from megatron.core.transformer.module import Float16Module
 
-        from megatron.bridge.training.fsdp_compat import MEGATRON_FSDP_TYPES
-
         module_instances = (
             DDP,
             torch_FSDP,
-            *MEGATRON_FSDP_TYPES,
+            FullyShardedDataParallelV1,
+            FullyShardedDataParallelV2,
             Float16Module,
             MegatronFSDP,
         )

src/megatron/bridge/training/config.py

@@ -71,7 +71,6 @@
 from megatron.bridge.peft.base import PEFT
 from megatron.bridge.training.comm_overlap import CommOverlapConfig
 from megatron.bridge.training.flex_dispatcher_backend import validate_flex_dispatcher_backend
-from megatron.bridge.training.fsdp_compat import MCORE_HAS_MEGATRON_FSDP_V2
 from megatron.bridge.training.mixed_precision import MixedPrecisionConfig, get_mixed_precision_config
 from megatron.bridge.training.tokenizers.config import TokenizerConfig
 from megatron.bridge.training.utils.config_utils import _ConfigContainerBase as Container
@@ -1082,8 +1081,8 @@ def _validate_and_apply_megatron_fsdp_configs(self) -> None:
         self.dist.use_megatron_fsdp = True
         self.ddp.use_megatron_fsdp = True
 
-        megatron_fsdp_version = getattr(self.ddp, "megatron_fsdp_version", 1)
-        if not MCORE_HAS_MEGATRON_FSDP_V2 or megatron_fsdp_version == 1:
+        megatron_fsdp_version = self.ddp.megatron_fsdp_version
+        if megatron_fsdp_version == 1:
             self._validate_and_apply_megatron_fsdp_v1_configs()
         elif megatron_fsdp_version == 2:
             self._validate_and_apply_megatron_fsdp_v2_configs()
@@ -1173,32 +1172,20 @@ def _validate_and_apply_megatron_fsdp_v2_configs(self) -> None:
             raise ValueError("MFSDP V2 does not support use_tp_pp_dp_mapping.")
         if self.rng.data_parallel_random_init:
             raise ValueError("MFSDP V2 does not support data_parallel_random_init.")
-        if self.ddp.num_distributed_optimizer_instances != 1:
-            raise ValueError("MFSDP V2 does not currently support HSDP.")
-        if self.ddp.outer_dp_sharding_strategy != "no_shard":
-            raise ValueError("MFSDP V2 does not currently support outer DP sharding.")
         if self.checkpoint.save is not None or self.checkpoint.load is not None:
             raise ValueError("MFSDP V2 checkpoint save and load are not yet supported.")
         if self.checkpoint.pretrained_checkpoint is not None:
             raise ValueError("MFSDP V2 checkpoint loading is not yet supported.")
         if self.optimizer.loss_scale is not None:
             raise ValueError("MFSDP V2 does not support loss scaling.")
-        if self.optimizer.clip_grad > 0.0:
-            raise ValueError("MFSDP V2 does not currently support gradient clipping.")
-        if self.optimizer.use_precision_aware_optimizer:
-            raise ValueError("MFSDP V2 does not support precision-aware optimizer.")
         if self.optimizer.optimizer_cpu_offload:
             raise ValueError("MFSDP V2 does not support optimizer CPU offload.")
         if self.optimizer.use_layer_wise_distributed_optimizer:
             raise ValueError("MFSDP V2 does not support layer-wise distributed optimizer.")
-        if self.optimizer.optimizer_cuda_graph:
-            raise ValueError("MFSDP V2 does not support optimizer CUDA graphs.")
         if self.model.calculate_per_token_loss:
             raise ValueError("MFSDP V2 does not support per-token loss normalization.")
         if self.model.fp8 or self.model.fp4 or self.ddp.fp8_param_gather or self.ddp.fp4_param_gather:
             raise ValueError("MFSDP V2 does not support FP8 or FP4.")
-        if self.model.cuda_graph_impl != "none" or self.ddp.megatron_fsdp_cuda_graph_mode:
-            raise ValueError("MFSDP V2 does not support CUDA graphs.")
 
         self.ddp.data_parallel_sharding_strategy = "optim_grads_params"
         self.ddp.use_distributed_optimizer = False

src/megatron/bridge/training/fsdp_compat.py

@@ -27,6 +27,10 @@
 from megatron.core import tensor_parallel
 from megatron.core.config import set_experimental_flag
 from megatron.core.distributed import DistributedDataParallel, DistributedDataParallelConfig, finalize_model_grads
+from megatron.core.distributed.fsdp.mcore_fsdp_adapter import (
+    FullyShardedDataParallelV1,
+    FullyShardedDataParallelV2,
+)
 from megatron.core.jit import disable_jit_fuser
 from megatron.core.optimizer import MegatronOptimizer
 from megatron.core.optimizer_param_scheduler import OptimizerParamScheduler
@@ -58,7 +62,6 @@
     maybe_load_dataloader_state,
 )
 from megatron.bridge.training.config import ConfigContainer
-from megatron.bridge.training.fsdp_compat import MEGATRON_FSDP_TYPES
 from megatron.bridge.training.gtp import (
     classify_gtp_remat_chains,
     configure_gtp_remat,
@@ -632,10 +635,19 @@ def _update_model_config_funcs(
     peft_enabled: bool = False,
 ) -> None:
     """Update model config sync funcs based on initialized model."""
-    if isinstance(model[0], (DistributedDataParallel, *MEGATRON_FSDP_TYPES)) and ddp_config.overlap_grad_reduce:
+    # DDP and MFSDP v1 only reduce during backward when overlap_grad_reduce is on, so
+    # without it there is nothing to suppress. MFSDP v2 always reduces in backward, so
+    # no_sync is how the schedule marks a non-final microbatch -- a correctness
+    # requirement, not an overlap optimization (mirrors Megatron-LM #7186). Without it,
+    # every backward finalizes the DP-outer axis and only the last microbatch's gradient
+    # reaches the optimizer.
+    if isinstance(model[0], FullyShardedDataParallelV2) or (
+        isinstance(model[0], (DistributedDataParallel, FullyShardedDataParallelV1))
+        and ddp_config.overlap_grad_reduce
+    ):
         assert model_config.no_sync_func is None, (
-            "When overlap_grad_reduce is True, config.no_sync_func must be None; "
-            "a custom no_sync_func is not supported when overlapping grad-reduce"
+            "config.no_sync_func must be None when the wrapper supplies its own no_sync "
+            "(overlap_grad_reduce, or Megatron-FSDP v2); a custom no_sync_func is not supported"
         )
         model_config.no_sync_func = [model_chunk.no_sync for model_chunk in model]
         if len(model) == 1:

src/megatron/bridge/training/setup.py

@@ -24,6 +24,10 @@
 import torch
 import torch.profiler
 from megatron.core.distributed import DistributedDataParallel as DDP
+from megatron.core.distributed.fsdp.mcore_fsdp_adapter import (
+    FullyShardedDataParallelV1,
+    FullyShardedDataParallelV2,
+)
 from megatron.core.full_cuda_graph import FullCudaGraphWrapper
 from megatron.core.num_microbatches_calculator import (
     get_current_global_batch_size,
@@ -73,7 +77,6 @@
 from megatron.bridge.training.config import ConfigContainer
 from megatron.bridge.training.eval import evaluate_and_print_results
 from megatron.bridge.training.forward_step_func_types import ForwardStepCallable
-from megatron.bridge.training.fsdp_compat import MEGATRON_FSDP_TYPES
 from megatron.bridge.training.gtp import get_data_distribution_group
 from megatron.bridge.training.initialize import destroy_global_state
 from megatron.bridge.training.nvrx_straggler import (
@@ -1741,7 +1744,7 @@ def _maybe_register_fsdp_buffers(
     ):
         print_rank_0("[Megatron-FSDP] Registering FSDP communication buffers manually")
         for model_chunk in model:
-            if isinstance(model_chunk, MEGATRON_FSDP_TYPES) and getattr(
+            if isinstance(model_chunk, (FullyShardedDataParallelV1, FullyShardedDataParallelV2)) and getattr(
                 model_chunk.ddp_config, "fsdp_manual_registration", False
             ):
                 fsdp_param_and_grad_buffer = getattr(model_chunk, "param_and_grad_buffer", None)

src/megatron/bridge/training/train.py

@@ -19,6 +19,7 @@
 import pytest
 import torch
 import torch.nn.functional as F
+from megatron.core.tensor_parallel.random import initialize_rng_tracker
 from megatron.core.transformer.enums import AttnBackend
 
 from megatron.bridge.models.gpt_provider import GPTModelProvider
@@ -38,7 +39,6 @@
     ValidationConfig,
     runtime_config_update,
 )
-from megatron.bridge.training.fsdp_compat import MCORE_HAS_MEGATRON_FSDP_V2
 from megatron.bridge.training.gpt_step import forward_step
 from megatron.bridge.training.pretrain import pretrain
 from megatron.bridge.training.state import GlobalState
@@ -422,19 +422,47 @@ def test_fsdp_pretrain_basic(self):
         torch.distributed.barrier()
 
     @pytest.mark.run_only_on("GPU")
-    @pytest.mark.skipif(not MCORE_HAS_MEGATRON_FSDP_V2, reason="MFSDP v2 is unavailable in this MCore revision")
     def test_fsdp_v2_dense_hybrid_pretrain_smoke(self):
-        """Train a dense two-layer HybridModel with the experimental MFSDP V2 path."""
+        """Train a dense two-layer HybridModel with MFSDP V2 in eager mode."""
         initialize_distributed()
         torch.distributed.barrier()
 
         cfg = create_fsdp_config_container(
             seq_length=128,
             train_iters=10,
-            optimizer={"clip_grad": 0.0},
+            optimizer={"use_precision_aware_optimizer": True},
+        )
+        cfg.model = create_dense_hybrid_smoke_model_config()
+        cfg.ddp.megatron_fsdp_version = 2
+
+        pretrain(cfg, forward_step)
+
+        torch.distributed.barrier()
+
+    @pytest.mark.run_only_on("GPU")
+    def test_fsdp_v2_dense_hybrid_cuda_graph_pretrain_smoke(self):
+        """Train a dense HybridModel with MFSDP V2 full-iteration and optimizer graphs."""
+        initialize_distributed()
+        torch.distributed.barrier()
+
+        # Replace the tracker left by eager tests with MCore's graph-safe RNG tracker.
+        initialize_rng_tracker(use_cudagraphable_rng=True, force_reset=True)
+
+        cfg = create_fsdp_config_container(
+            seq_length=128,
+            train_iters=10,
+            optimizer={"use_precision_aware_optimizer": True, "optimizer_cuda_graph": True},
         )
         cfg.model = create_dense_hybrid_smoke_model_config()
         cfg.ddp.megatron_fsdp_version = 2
+        # Full-iteration graphs over a dense model: Megatron-LM #7075 made these work
+        # under MFSDP v2. MoE stays out, its dispatch shapes are not capturable.
+        cfg.model.cuda_graph_impl = "full_iteration"
+        # The loss NaN check reads a GPU value on the host and cannot be captured.
+        cfg.rerun_state_machine.check_for_nan_in_loss = False
+        # Capturable TE FusedAdam requires the main-gradient and main-weight dtypes
+        # to match (https://github.com/NVIDIA/TransformerEngine/issues/3358).
+        cfg.ddp.megatron_fsdp_main_grads_dtype = torch.float32
 
         pretrain(cfg, forward_step)
 
@@ -455,6 +483,32 @@ def test_fsdp_v2_moe_ep2_pretrain_smoke(self):
         cfg.ddp.megatron_fsdp_version = 2
 
         pretrain(cfg, forward_step)
+
+        torch.distributed.barrier()
+
+    @pytest.mark.run_only_on("GPU")
+    @pytest.mark.parametrize("outer_dp_sharding_strategy", ["no_shard", "optim"], ids=["hsdp", "hfsdp"])
+    def test_fsdp_v2_hybrid_dp_pretrain_smoke(self, outer_dp_sharding_strategy):
+        """Train the dense HybridModel over an outer data-parallel axis.
+
+        no_shard outer over the optim_grads_params inner axis is HSDP, optim outer is HFSDP.
+        """
+        initialize_distributed()
+        torch.distributed.barrier()
+
+        cfg = create_fsdp_config_container(
+            seq_length=128,
+            train_iters=10,
+            ddp={
+                "num_distributed_optimizer_instances": 2,
+                "outer_dp_sharding_strategy": outer_dp_sharding_strategy,
+            },
+        )
+        cfg.model = create_dense_hybrid_smoke_model_config()
+        cfg.ddp.megatron_fsdp_version = 2
+
+        pretrain(cfg, forward_step)
+
         torch.distributed.barrier()
 
     @pytest.mark.run_only_on("GPU")

tests/functional_tests/test_groups/megatron_fsdp/test_megatron_fsdp.py

@@ -27,8 +27,12 @@ class FSDPV2(FSDPV1):
         pass
 
     monkeypatch.setattr(
-        "megatron.bridge.training.fsdp_compat.MEGATRON_FSDP_TYPES",
-        (FSDPV1, FSDPV2),
+        "megatron.core.distributed.fsdp.mcore_fsdp_adapter.FullyShardedDataParallelV1",
+        FSDPV1,
+    )
+    monkeypatch.setattr(
+        "megatron.core.distributed.fsdp.mcore_fsdp_adapter.FullyShardedDataParallelV2",
+        FSDPV2,
     )
 
     module = torch.nn.Linear(1, 1)

tests/unit_tests/models/test_conversion_unwrap_utils.py

@@ -20,9 +20,9 @@
 from unittest.mock import Mock, patch
 
 import pytest
+from megatron.core.distributed.fsdp.mcore_fsdp_adapter import FullyShardedDataParallelV1
 from megatron.core.optimizer.distrib_optimizer import DistributedOptimizer
 
-from megatron.bridge.training.fsdp_compat import MEGATRON_FSDP_TYPES
 from megatron.bridge.training.state import GlobalState
 from megatron.bridge.training.train import (
     _delete_cuda_graphs,
@@ -312,7 +312,7 @@ def test_maybe_register_fsdp_buffers_execution(self):
         config.ddp.fsdp_manual_registration = True
 
         # Mock model chunk
-        model_chunk = Mock(spec=MEGATRON_FSDP_TYPES[0])
+        model_chunk = Mock(spec=FullyShardedDataParallelV1)
         # Mock ddp_config on the chunk
         model_chunk.ddp_config = Mock()
         model_chunk.ddp_config.fsdp_manual_registration = True