diff --git a/examples/configs/recipes/llm/performance/dapo-deepseek-v3-64n8g.yaml b/examples/configs/recipes/llm/performance/dapo-deepseek-v3-64n8g.yaml index 056c15294a0..5ef1591ecde 100644 --- a/examples/configs/recipes/llm/performance/dapo-deepseek-v3-64n8g.yaml +++ b/examples/configs/recipes/llm/performance/dapo-deepseek-v3-64n8g.yaml @@ -44,6 +44,7 @@ policy: empty_unused_memory_level: 2 enabled: true activation_checkpointing: true + moe_grouped_gemm: true tensor_model_parallel_size: 8 expert_model_parallel_size: 32 pipeline_model_parallel_size: 8 diff --git a/examples/configs/recipes/llm/performance/grpo-deepseek-v3-32n8g.yaml b/examples/configs/recipes/llm/performance/grpo-deepseek-v3-32n8g.yaml index 35903949ad2..a3911db41a9 100644 --- a/examples/configs/recipes/llm/performance/grpo-deepseek-v3-32n8g.yaml +++ b/examples/configs/recipes/llm/performance/grpo-deepseek-v3-32n8g.yaml @@ -26,6 +26,7 @@ policy: pipeline_model_parallel_size: 16 expert_model_parallel_size: 16 activation_checkpointing: true + moe_grouped_gemm: true num_layers_in_first_pipeline_stage: 3 num_layers_in_last_pipeline_stage: 2 apply_rope_fusion: false diff --git a/examples/configs/recipes/llm/performance/grpo-qwen3-235b-16n8g.yaml b/examples/configs/recipes/llm/performance/grpo-qwen3-235b-16n8g.yaml index aecdabba73b..1228db60bb2 100644 --- a/examples/configs/recipes/llm/performance/grpo-qwen3-235b-16n8g.yaml +++ b/examples/configs/recipes/llm/performance/grpo-qwen3-235b-16n8g.yaml @@ -29,6 +29,7 @@ policy: context_parallel_size: 2 expert_model_parallel_size: 16 activation_checkpointing: true + moe_grouped_gemm: true num_layers_in_first_pipeline_stage: 11 num_layers_in_last_pipeline_stage: 11 defer_fp32_logits: true diff --git a/examples/configs/recipes/llm/performance/grpo-qwen3-30ba3b-4n4g.yaml b/examples/configs/recipes/llm/performance/grpo-qwen3-30ba3b-4n4g.yaml index 9292a72439c..947b2d1b1c9 100644 --- a/examples/configs/recipes/llm/performance/grpo-qwen3-30ba3b-4n4g.yaml +++ b/examples/configs/recipes/llm/performance/grpo-qwen3-30ba3b-4n4g.yaml @@ -22,6 +22,7 @@ policy: pipeline_model_parallel_size: 1 expert_model_parallel_size: 16 sequence_parallel: false + moe_grouped_gemm: true optimizer: lr: 3.0e-07 min_lr: 3.0e-08 diff --git a/examples/configs/recipes/llm/performance/grpo-qwen3-30ba3b-4n8g-40K.yaml b/examples/configs/recipes/llm/performance/grpo-qwen3-30ba3b-4n8g-40K.yaml index 3b4f22ffbd4..a8e130f8530 100644 --- a/examples/configs/recipes/llm/performance/grpo-qwen3-30ba3b-4n8g-40K.yaml +++ b/examples/configs/recipes/llm/performance/grpo-qwen3-30ba3b-4n8g-40K.yaml @@ -24,6 +24,7 @@ policy: expert_model_parallel_size: 8 sequence_parallel: true context_parallel_size: 8 + moe_grouped_gemm: true optimizer: lr: 3.0e-07 min_lr: 3.0e-08 diff --git a/examples/configs/recipes/llm/performance/grpo-qwen3-30ba3b-4n8g.yaml b/examples/configs/recipes/llm/performance/grpo-qwen3-30ba3b-4n8g.yaml index 21ddcc6bd36..6eda477ba1b 100644 --- a/examples/configs/recipes/llm/performance/grpo-qwen3-30ba3b-4n8g.yaml +++ b/examples/configs/recipes/llm/performance/grpo-qwen3-30ba3b-4n8g.yaml @@ -22,6 +22,7 @@ policy: pipeline_model_parallel_size: 1 expert_model_parallel_size: 8 sequence_parallel: false + moe_grouped_gemm: true optimizer: lr: 3.0e-07 min_lr: 3.0e-08 diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index 1651dde9883..2c946dfd8a0 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -630,6 +630,9 @@ def _apply_moe_config(model_cfg: Any, config: PolicyConfig) -> None: model_cfg.moe_permute_fusion = config["megatron_cfg"]["moe_permute_fusion"] + if "moe_grouped_gemm" in config["megatron_cfg"]: + model_cfg.moe_grouped_gemm = config["megatron_cfg"]["moe_grouped_gemm"] + def _apply_mtp_config(model_cfg: Any, config: PolicyConfig) -> None: if "mtp_num_layers" in config["megatron_cfg"]: diff --git a/nemo_rl/models/policy/__init__.py b/nemo_rl/models/policy/__init__.py index 2119ddb9722..9c6303519d9 100644 --- a/nemo_rl/models/policy/__init__.py +++ b/nemo_rl/models/policy/__init__.py @@ -237,6 +237,10 @@ class MegatronConfig(TypedDict): moe_token_dispatcher_type: str # Can be used only with 'alltoall' token dispatcher moe_shared_expert_overlap: bool + # Enable grouped GEMM for MoE experts via CUTLASS. Significant throughput + # gain when multiple experts are assigned per rank (num_local_experts > 1). + # Requires TE >= 1.11.0 for FP8 and Ampere (sm_80) or newer. + moe_grouped_gemm: NotRequired[bool] # HybridEP settings for MoE expert parallelism (requires moe_token_dispatcher_type='flex') # See: https://github.com/deepseek-ai/DeepEP/tree/hybrid-ep moe_flex_dispatcher_backend: NotRequired[str] diff --git a/tests/unit/models/megatron/test_megatron_setup.py b/tests/unit/models/megatron/test_megatron_setup.py index e73003e91bc..543348681f7 100644 --- a/tests/unit/models/megatron/test_megatron_setup.py +++ b/tests/unit/models/megatron/test_megatron_setup.py @@ -620,6 +620,20 @@ def test_moe_configuration(self): assert model_cfg.moe_token_dispatcher_type == "alltoall" assert model_cfg.moe_shared_expert_overlap is True + @staticmethod + def _base_moe_megatron_cfg() -> dict: + return { + "expert_tensor_parallel_size": 2, + "expert_model_parallel_size": 4, + "moe_router_dtype": "float32", + "moe_router_load_balancing_type": "none", + "moe_router_bias_update_rate": 0.0, + "moe_permute_fusion": True, + "moe_enable_deepep": False, + "moe_token_dispatcher_type": "alltoall", + "moe_shared_expert_overlap": True, + } + @staticmethod def _base_moe_cfg(**overrides): cfg = { @@ -636,6 +650,45 @@ def _base_moe_cfg(**overrides): cfg.update(overrides) return {"megatron_cfg": cfg} + @pytest.mark.parametrize("moe_grouped_gemm", [True, False]) + def test_moe_grouped_gemm_explicit(self, moe_grouped_gemm): + """moe_grouped_gemm is applied when present in config.""" + from nemo_rl.models.megatron.setup import _apply_moe_config + + model_cfg = MagicMock() + megatron_cfg = self._base_moe_megatron_cfg() + megatron_cfg["moe_grouped_gemm"] = moe_grouped_gemm + config = {"megatron_cfg": megatron_cfg} + + _apply_moe_config(model_cfg, config) + + assert model_cfg.moe_grouped_gemm is moe_grouped_gemm + + def test_moe_grouped_gemm_absent_keeps_default(self): + """Absent key leaves the attr unset on the model cfg.""" + from nemo_rl.models.megatron.setup import _apply_moe_config + + # spec lists everything _apply_moe_config writes so we can detect + # whether the moe_grouped_gemm branch fires. + model_cfg = MagicMock( + spec=[ + "expert_tensor_parallel_size", + "expert_model_parallel_size", + "moe_router_dtype", + "moe_router_load_balancing_type", + "moe_router_bias_update_rate", + "moe_permute_fusion", + "moe_enable_deepep", + "moe_token_dispatcher_type", + "moe_shared_expert_overlap", + ] + ) + config = {"megatron_cfg": self._base_moe_megatron_cfg()} + + _apply_moe_config(model_cfg, config) + + assert not hasattr(model_cfg, "moe_grouped_gemm") + def test_hybridep_env_vars_auto_set_with_warning(self, monkeypatch): """HybridEP backend with no env config: auto-set env vars and emit warnings.""" from nemo_rl.models.megatron.setup import _apply_moe_config