feat(training): support MFSDP V2 expert parallelism

Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>

Author: wujingyue

src/megatron/bridge/training/config.py

@@ -1140,7 +1140,6 @@ def _validate_and_apply_megatron_fsdp_v2_configs(self) -> None:
             "tensor_model_parallel_size",
             "pipeline_model_parallel_size",
             "context_parallel_size",
-            "expert_model_parallel_size",
         )
         configured_parallelisms = [
             f"{name}={getattr(self.model, name)}"
@@ -1149,11 +1148,11 @@ def _validate_and_apply_megatron_fsdp_v2_configs(self) -> None:
         ]
         if configured_parallelisms:
             raise ValueError(
-                "MFSDP V2 currently supports DP-only training; unsupported settings: "
-                + ", ".join(configured_parallelisms)
+                "MFSDP V2 requires TP=PP=CP=1; unsupported settings: " + ", ".join(configured_parallelisms)
             )
-        if self.model.num_moe_experts is not None:
-            raise ValueError("MFSDP V2 does not currently support MoE models.")
+        if self.model.expert_model_parallel_size > 1:
+            if self.model.num_moe_experts is None:
+                raise ValueError("MFSDP V2 expert parallelism requires an MoE model.")
         if self.model.virtual_pipeline_model_parallel_size is not None:
             raise ValueError("MFSDP V2 does not currently support multiple model chunks.")
         if self.dist.use_tp_pp_dp_mapping:

tests/functional_tests/test_groups/megatron_fsdp/test_megatron_fsdp.py

@@ -19,6 +19,7 @@
 import pytest
 import torch
 import torch.nn.functional as F
+from megatron.core.transformer.enums import AttnBackend
 
 from megatron.bridge.models.gpt_provider import GPTModelProvider
 from megatron.bridge.models.hybrid.hybrid_provider import HybridModelProvider
@@ -111,6 +112,19 @@ class DenseHybridSmokeModelProvider(HybridModelProvider):
     gradient_accumulation_fusion: bool = False
 
 
+@dataclass
+class MLAMoEHybridSmokeModelProvider(HybridModelProvider):
+    """Small MLA/MoE HybridModel configuration for the MFSDP V2 EP smoke test."""
+
+    attention_backend: AttnBackend = AttnBackend.auto
+    seq_length: int = 128
+    hidden_size: int = 128
+    multi_latent_attention: bool = True
+    hybrid_layer_pattern: str = "+E"
+    num_moe_experts: int = 4
+    expert_model_parallel_size: int = 2
+
+
 def create_fsdp_model_config(seq_length: int, bf16: bool = True, **kwargs) -> Llama3FSDPTestModelProvider:
     """Create a standardized FSDP model configuration."""
     base_config = {
@@ -427,6 +441,23 @@ def test_fsdp_v2_dense_hybrid_pretrain_smoke(self):
 
         torch.distributed.barrier()
 
+    @pytest.mark.run_only_on("GPU")
+    def test_fsdp_v2_mla_moe_ep2_pretrain_smoke(self):
+        """Train a small MLA/MoE HybridModel with MFSDP V2 and EP=2."""
+        initialize_distributed()
+        torch.distributed.barrier()
+
+        cfg = create_fsdp_config_container(
+            seq_length=128,
+            train_iters=10,
+            optimizer={"clip_grad": 0.0},
+        )
+        cfg.model = MLAMoEHybridSmokeModelProvider()
+        cfg.ddp.megatron_fsdp_version = 2
+
+        pretrain(cfg, forward_step)
+        torch.distributed.barrier()
+
     @pytest.mark.run_only_on("GPU")
     def test_fsdp_pretrain_save_resume(self, tmp_path):
         """