diff --git a/3rdparty/Megatron-LM b/3rdparty/Megatron-LM index 24bad8e677..2857bd0a59 160000 --- a/3rdparty/Megatron-LM +++ b/3rdparty/Megatron-LM @@ -1 +1 @@ -Subproject commit 24bad8e677d22625d86ef2a54c9506b6e4992c93 +Subproject commit 2857bd0a599a45fbab5bfda3ebf9890125228a80 diff --git a/src/megatron/bridge/training/config.py b/src/megatron/bridge/training/config.py index 7e9e14a1a5..43af9d36d9 100644 --- a/src/megatron/bridge/training/config.py +++ b/src/megatron/bridge/training/config.py @@ -1138,7 +1138,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)}" @@ -1147,11 +1146,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: @@ -1168,8 +1167,6 @@ def _validate_and_apply_megatron_fsdp_v2_configs(self) -> 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: diff --git a/tests/functional_tests/test_groups/megatron_fsdp/test_megatron_fsdp.py b/tests/functional_tests/test_groups/megatron_fsdp/test_megatron_fsdp.py index e7e45e298b..bf1c3eab49 100644 --- a/tests/functional_tests/test_groups/megatron_fsdp/test_megatron_fsdp.py +++ b/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 @@ -110,6 +111,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 = { @@ -416,7 +430,6 @@ def test_fsdp_v2_dense_hybrid_pretrain_smoke(self): cfg = create_fsdp_config_container( seq_length=128, train_iters=10, - optimizer={"clip_grad": 0.0}, ) cfg.model = create_dense_hybrid_smoke_model_config() cfg.ddp.megatron_fsdp_version = 2 @@ -425,6 +438,22 @@ 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, + ) + 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): """