perf(recipe): keep Kimi K2, Nemotron 3 Nano and Qwen3 30B 32-GPU GB200 on HybridEP

Revert the NCCL EP switch for the GB200 Kimi K2 MXFP8, Nemotron 3 Nano
BF16/MXFP8 and Qwen3 30B-A3B 32-GPU BF16 perf recipes; they keep their
HybridEP configuration and environment from main.

Nemotron 3.5 Lightning BF16 still derives from the Nano BF16 recipe, so it now
applies the NCCL EP overlay itself. The Qwen3 30B-A3B 8-GPU and 32-GPU BF16
recipes share a new HybridEP builder so only the 8-GPU recipe switches.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Malay Nagda <malayn@nvidia.com>

Author: malay-nagda

src/megatron/bridge/perf_recipes/kimi/gb200/kimi_k2.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 """GB200 performance recipes for Kimi K2."""
 
-from megatron.bridge.perf_recipes._common import _enable_ncclep
 from megatron.bridge.perf_recipes.environment import COMMON_PERF_ENV_VARS
 from megatron.bridge.perf_recipes.kimi.common import (
     ConfigContainer,
@@ -139,7 +138,7 @@ def kimi_k2_pretrain_256gpu_gb200_fp8cs_config() -> ConfigContainer:
 
 
 def kimi_k2_pretrain_256gpu_gb200_fp8mx_config() -> ConfigContainer:
-    """Kimi K2 pretrain: 256× GB200, MXFP8, NCCL EP."""
+    """Kimi K2 pretrain: 256× GB200, MXFP8."""
     cfg = kimi_k2_pretrain_config()
     cfg.mixed_precision = _perf_precision("fp8_mx")
     cfg.mixed_precision.reuse_grad_buf_for_mxfp8_param_ag = False
@@ -174,10 +173,6 @@ def kimi_k2_pretrain_256gpu_gb200_fp8mx_config() -> ConfigContainer:
     cfg.comm_overlap.overlap_grad_reduce = True
 
     _benchmark_common(cfg)
-    _enable_ncclep(cfg)
-    # Device-side expert token counts: the legacy grouped MLP path syncs tokens_per_expert to the
-    # host every layer, which serializes the CPU behind the GPU when dispatch is fast.
-    cfg.model.moe_use_grouped_tensor = True
     # Keep process settings next to the recipe so users can see the exact benchmark environment.
     cfg.env_vars = {
         **COMMON_PERF_ENV_VARS,
@@ -189,8 +184,11 @@ def kimi_k2_pretrain_256gpu_gb200_fp8mx_config() -> ConfigContainer:
         "TORCH_NCCL_AVOID_RECORD_STREAMS": 1,
         # NCCL user-buffer and launch settings.
         "NCCL_NVLS_ENABLE": 0,
-        # NCCL EP dispatcher mode.
-        "NCCL_EP_HT_EM_PULL_PUSH": 1,
+        # HybridEP topology for the target system.
+        "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 64,
+        "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128,
+        "NVLINK_DOMAIN_SIZE": 72,
+        "USE_MNNVL": 1,
         # Transformer Engine overlap settings for this model.
         "NVTE_BWD_LAYERNORM_SM_MARGIN": 20,
         "NVTE_FWD_LAYERNORM_SM_MARGIN": 20,

src/megatron/bridge/perf_recipes/nemotronh/gb200/nemotronh.py

@@ -401,7 +401,7 @@ def nemotron_3_ultra_pretrain_256gpu_gb200_nvfp4_config() -> ConfigContainer:
 
 
 def nemotron_3_nano_pretrain_8gpu_gb200_bf16_config() -> ConfigContainer:
-    """Nemotron 3 Nano pretrain: 8× GB200, BF16, NCCL EP."""
+    """Nemotron 3 Nano pretrain: 8× GB200, BF16."""
     cfg = nemotron_3_nano_pretrain_config()
     _apply_nemotron_3_nano_perf_defaults(cfg)
     cfg.mixed_precision = _perf_precision("bf16")
@@ -425,10 +425,6 @@ def nemotron_3_nano_pretrain_8gpu_gb200_bf16_config() -> ConfigContainer:
 
     _benchmark_common(cfg)
     cfg.model.moe_hybridep_num_sms = 16
-    _enable_ncclep(cfg)
-    # Device-side expert token counts: the legacy grouped MLP path syncs tokens_per_expert to the
-    # host every layer, which serializes the CPU behind the GPU when dispatch is fast.
-    cfg.model.moe_use_grouped_tensor = True
     # Keep process settings next to the recipe so users can see the exact benchmark environment.
     cfg.env_vars = {
         **COMMON_PERF_ENV_VARS,
@@ -440,8 +436,11 @@ def nemotron_3_nano_pretrain_8gpu_gb200_bf16_config() -> ConfigContainer:
         "TORCH_NCCL_AVOID_RECORD_STREAMS": 1,
         # NCCL user-buffer and launch settings.
         "NCCL_NVLS_ENABLE": 0,
-        # NCCL EP dispatcher mode.
-        "NCCL_EP_HT_EM_PULL_PUSH": 1,
+        # HybridEP topology for the target system.
+        "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 8,
+        "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128,
+        "NVLINK_DOMAIN_SIZE": 72,
+        "USE_MNNVL": 1,
         # Transformer Engine overlap settings for this model.
         "NVTE_BWD_LAYERNORM_SM_MARGIN": 20,
         "NVTE_FWD_LAYERNORM_SM_MARGIN": 20,
@@ -480,15 +479,11 @@ def _build_nemotron_3_nano_gb200_mxfp8() -> ConfigContainer:
 
 
 def nemotron_3_nano_pretrain_8gpu_gb200_fp8mx_config() -> ConfigContainer:
-    """Nemotron 3 Nano pretrain: 8× GB200, MXFP8, NCCL EP."""
+    """Nemotron 3 Nano pretrain: 8× GB200, MXFP8."""
     cfg = _build_nemotron_3_nano_gb200_mxfp8()
     cfg.model.use_transformer_engine_op_fuser = True
     cfg.model.moe_mlp_glu_interleave_size = 32
     cfg.mixed_precision.fp8_dot_product_attention = True
-    _enable_ncclep(cfg)
-    # Device-side expert token counts: the legacy grouped MLP path syncs tokens_per_expert to the
-    # host every layer, which serializes the CPU behind the GPU when dispatch is fast.
-    cfg.model.moe_use_grouped_tensor = True
     # Keep process settings next to the recipe so users can see the exact benchmark environment.
     cfg.env_vars = {
         **COMMON_PERF_ENV_VARS,
@@ -500,8 +495,11 @@ def nemotron_3_nano_pretrain_8gpu_gb200_fp8mx_config() -> ConfigContainer:
         "TORCH_NCCL_AVOID_RECORD_STREAMS": 1,
         # NCCL user-buffer and launch settings.
         "NCCL_NVLS_ENABLE": 0,
-        # NCCL EP dispatcher mode.
-        "NCCL_EP_HT_EM_PULL_PUSH": 1,
+        # HybridEP topology for the target system.
+        "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 8,
+        "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128,
+        "NVLINK_DOMAIN_SIZE": 72,
+        "USE_MNNVL": 1,
         # Transformer Engine overlap settings for this model.
         "CUDNNFE_CLUSTER_OVERLAP_MARGIN": 8,
         "NVTE_BWD_LAYERNORM_SM_MARGIN": 20,
@@ -568,7 +566,7 @@ def nemotron_3_nano_pretrain_8gpu_gb200_nvfp4_config() -> ConfigContainer:
 
 
 def nemotron_3_5_lightning_pretrain_8gpu_gb200_bf16_config() -> ConfigContainer:
-    """Nemotron 3.5 Lightning pretrain: 8× GB200, BF16, NCCL EP (inherited from Nemotron 3 Nano)."""
+    """Nemotron 3.5 Lightning pretrain: 8× GB200, BF16, NCCL EP."""
     cfg = nemotron_3_nano_pretrain_8gpu_gb200_bf16_config()
     cfg.model.recompute_modules = []
     cfg.model.mtp_num_layers = 2
@@ -580,14 +578,18 @@ def nemotron_3_5_lightning_pretrain_8gpu_gb200_bf16_config() -> ConfigContainer:
     cfg.model.hf_model_revision = _NEMOTRON_3_5_LIGHTNING_MODEL_REVISION
     cfg.tokenizer.tokenizer_model = _NEMOTRON_3_5_LIGHTNING_MODEL_ID
     cfg.tokenizer.hf_tokenizer_kwargs = {"revision": _NEMOTRON_3_5_LIGHTNING_MODEL_REVISION}
+    _enable_ncclep(cfg)
+    # Device-side expert token counts: the legacy grouped MLP path syncs tokens_per_expert to the
+    # host every layer, which serializes the CPU behind the GPU when dispatch is fast.
+    cfg.model.moe_use_grouped_tensor = True
     cfg.env_vars = {
         **COMMON_PERF_ENV_VARS,
         "CUDA_DEVICE_MAX_CONNECTIONS": 32,
         "NCCL_GRAPH_REGISTER": 0,
         "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True",
         "TORCH_NCCL_AVOID_RECORD_STREAMS": 1,
         "NCCL_NVLS_ENABLE": 0,
-        # NCCL EP dispatcher mode (inherited from the Nemotron 3 Nano BF16 recipe).
+        # NCCL EP dispatcher mode.
         "NCCL_EP_HT_EM_PULL_PUSH": 1,
         "NVTE_BWD_LAYERNORM_SM_MARGIN": 20,
         "NVTE_FWD_LAYERNORM_SM_MARGIN": 20,

src/megatron/bridge/perf_recipes/qwen/gb200/qwen3_moe.py

@@ -196,8 +196,8 @@ def qwen3_235b_a22b_pretrain_64gpu_gb200_fp8mx_config() -> ConfigContainer:
     return cfg
 
 
-def qwen3_30b_a3b_pretrain_8gpu_gb200_bf16_config() -> ConfigContainer:
-    """Qwen3 30B-A3B pretrain: 8× GB200, BF16, EP=8, NCCL EP."""
+def _build_qwen3_30b_a3b_gb200_bf16() -> ConfigContainer:
+    """Shared HybridEP BF16 base for the Qwen3 30B-A3B 8-GPU and 32-GPU GB200 recipes."""
     cfg = qwen3_30b_a3b_pretrain_config()
     cfg.mixed_precision = _perf_precision("bf16")
     cfg.model.bias_activation_fusion = True
@@ -228,6 +228,12 @@ def qwen3_30b_a3b_pretrain_8gpu_gb200_bf16_config() -> ConfigContainer:
     cfg.comm_overlap = CommOverlapConfig(tp_comm_overlap=False)
 
     _benchmark_common(cfg)
+    return cfg
+
+
+def qwen3_30b_a3b_pretrain_8gpu_gb200_bf16_config() -> ConfigContainer:
+    """Qwen3 30B-A3B pretrain: 8× GB200, BF16, EP=8, NCCL EP."""
+    cfg = _build_qwen3_30b_a3b_gb200_bf16()
     _enable_ncclep(cfg)
     # Device-side expert token counts: the legacy grouped MLP path syncs tokens_per_expert to the
     # host every layer, which serializes the CPU behind the GPU when dispatch is fast.
@@ -655,8 +661,8 @@ def qwen3_30b_a3b_pretrain_8gpu_gb200_nvfp4_config() -> ConfigContainer:
 
 
 def qwen3_30b_a3b_pretrain_32gpu_gb200_bf16_config() -> ConfigContainer:
-    """Qwen3 30B-A3B pretrain: 32× GB200, BF16, legacy-scaled GBS, NCCL EP."""
-    cfg = _with_global_batch_size(qwen3_30b_a3b_pretrain_8gpu_gb200_bf16_config(), 2048)
+    """Qwen3 30B-A3B pretrain: 32× GB200, BF16, legacy-scaled GBS."""
+    cfg = _with_global_batch_size(_build_qwen3_30b_a3b_gb200_bf16(), 2048)
     # Keep process settings next to the recipe so users can see the exact benchmark environment.
     cfg.env_vars = {
         **COMMON_PERF_ENV_VARS,
@@ -668,8 +674,11 @@ def qwen3_30b_a3b_pretrain_32gpu_gb200_bf16_config() -> ConfigContainer:
         "TORCH_NCCL_AVOID_RECORD_STREAMS": 1,
         # NCCL user-buffer and launch settings.
         "NCCL_NVLS_ENABLE": 0,
-        # NCCL EP dispatcher mode (inherited from the 8-GPU BF16 recipe).
-        "NCCL_EP_HT_EM_PULL_PUSH": 1,
+        # HybridEP topology for the target system.
+        "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 8,
+        "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128,
+        "NVLINK_DOMAIN_SIZE": 72,
+        "USE_MNNVL": 1,
         # Transformer Engine overlap settings for this model.
         "NVTE_BWD_LAYERNORM_SM_MARGIN": 20,
         "NVTE_FWD_LAYERNORM_SM_MARGIN": 20,

tests/unit_tests/recipes/test_gb200_ncclep_perf_recipes.py

@@ -26,44 +26,43 @@
 
 _NCCL_EP_ENV = {"NCCL_EP_HT_EM_PULL_PUSH": 1}
 
-# (module under megatron.bridge.perf_recipes, recipe factory) for the GB200 recipes defaulted to NCCL EP: the
-# nemo-ci GB200 benchmarks whose GB300 counterparts already run NCCL EP.
+# (module under megatron.bridge.perf_recipes, recipe factory) for the GB200 recipes defaulted to NCCL EP.
 _GB200_NCCLEP_RECIPES = (
     ("deepseek.gb200.deepseek_v3", "deepseek_v3_pretrain_256gpu_gb200_bf16_config"),
     ("deepseek.gb200.deepseek_v3", "deepseek_v3_pretrain_256gpu_gb200_fp8mx_config"),
     ("deepseek.gb200.deepseek_v3", "deepseek_v3_pretrain_256gpu_gb200_nvfp4_config"),
     ("deepseek.gb200.deepseek_v3", "deepseek_v3_pretrain_256gpu_gb200_fp8mx_large_scale_config"),
     ("gpt_oss.gb200.gpt_oss", "gpt_oss_120b_pretrain_64gpu_gb200_fp8mx_config"),
-    ("kimi.gb200.kimi_k2", "kimi_k2_pretrain_256gpu_gb200_fp8mx_config"),
     ("nemotronh.gb200.nemotronh", "nemotron_3_super_pretrain_64gpu_gb200_bf16_config"),
     ("nemotronh.gb200.nemotronh", "nemotron_3_super_pretrain_64gpu_gb200_fp8mx_config"),
     ("nemotronh.gb200.nemotronh", "nemotron_3_super_pretrain_64gpu_gb200_nvfp4_config"),
     ("nemotronh.gb200.nemotronh", "nemotron_3_ultra_pretrain_256gpu_gb200_fp8mx_config"),
-    ("nemotronh.gb200.nemotronh", "nemotron_3_nano_pretrain_8gpu_gb200_bf16_config"),
-    ("nemotronh.gb200.nemotronh", "nemotron_3_nano_pretrain_8gpu_gb200_fp8mx_config"),
     ("nemotronh.gb200.nemotronh", "nemotron_3_5_lightning_pretrain_8gpu_gb200_bf16_config"),
     ("nemotronh.gb200.nemotronh", "nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_config"),
     ("qwen.gb200.qwen3_moe", "qwen3_30b_a3b_pretrain_8gpu_gb200_bf16_config"),
     ("qwen.gb200.qwen3_moe", "qwen3_30b_a3b_pretrain_8gpu_gb200_fp8mx_config"),
-    ("qwen.gb200.qwen3_moe", "qwen3_30b_a3b_pretrain_32gpu_gb200_bf16_config"),
     ("qwen.gb200.qwen3_moe", "qwen3_235b_a22b_pretrain_256gpu_gb200_bf16_config"),
     ("qwen.gb200.qwen3_moe", "qwen3_235b_a22b_pretrain_256gpu_gb200_fp8mx_config"),
     ("qwen.gb200.qwen3_moe", "qwen3_235b_a22b_pretrain_256gpu_gb200_fp8mx_large_scale_config"),
     ("qwen.gb200.qwen3_moe", "qwen3_235b_a22b_pretrain_256gpu_gb200_nvfp4_config"),
 )
 
 # GB200 recipes that stay on HybridEP. Most build on the same base or parent recipe as a recipe above (the
-# shared DeepSeek BF16 builder, the Nano/Lightning MXFP8 base, the FP8-CS parents of the NVFP4 recipes), so
-# the NCCL EP overlay must not leak into them.
+# shared DeepSeek and Qwen3 30B BF16 builders, the Nemotron 3 Nano parents of the Lightning recipes, the FP8-CS
+# parents of the NVFP4 recipes), so the NCCL EP overlay must not leak into them.
 _GB200_HYBRIDEP_RECIPES = (
     ("deepseek.gb200.deepseek_v3", "deepseek_v3_pretrain_256gpu_gb200_fp8cs_config"),
     ("deepseek.gb200.deepseek_v3", "deepseek_v3_pretrain_256gpu_gb200_fp8mx_partial_cg_dev_config"),
+    ("kimi.gb200.kimi_k2", "kimi_k2_pretrain_256gpu_gb200_fp8mx_config"),
+    ("nemotronh.gb200.nemotronh", "nemotron_3_nano_pretrain_8gpu_gb200_bf16_config"),
+    ("nemotronh.gb200.nemotronh", "nemotron_3_nano_pretrain_8gpu_gb200_fp8mx_config"),
     ("nemotronh.gb200.nemotronh", "nemotron_3_ultra_pretrain_256gpu_gb200_nvfp4_config"),
     ("nemotronh.gb200.nemotronh", "nemotron_3_nano_pretrain_8gpu_gb200_nvfp4_config"),
     ("nemotronh.gb200.nemotronh", "nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_fsdp_config"),
     ("nemotronh.gb200.nemotronh", "nemotron_3_5_lightning_pretrain_8gpu_gb200_nvfp4_config"),
     ("qwen.gb200.qwen3_moe", "qwen3_30b_a3b_pretrain_8gpu_gb200_fp8cs_config"),
     ("qwen.gb200.qwen3_moe", "qwen3_30b_a3b_pretrain_8gpu_gb200_nvfp4_config"),
+    ("qwen.gb200.qwen3_moe", "qwen3_30b_a3b_pretrain_32gpu_gb200_bf16_config"),
     ("qwen.gb200.qwen3_moe", "qwen3_30b_a3b_pretrain_32gpu_gb200_fp8cs_config"),
     ("qwen.gb200.qwen3_moe", "qwen3_235b_a22b_pretrain_256gpu_gb200_fp8cs_config"),
     ("qwen.gb200.qwen3_moe", "qwen3_235b_a22b_pretrain_64gpu_gb200_nvfp4_full_iteration_config"),

tests/unit_tests/recipes/test_nemotron_3_5_lightning_perf_recipes.py

@@ -22,6 +22,7 @@
 import torch
 
 from megatron.bridge.perf_recipes._common import _enable_ncclep
+from megatron.bridge.perf_recipes.environment import HYBRID_EP_ENV_NAMES
 from megatron.bridge.perf_recipes.nemotronh import (
     nemotron_3_5_lightning_pretrain_8gpu_b200_bf16_config,
     nemotron_3_5_lightning_pretrain_8gpu_b200_fp8mx_config,
@@ -350,16 +351,22 @@ def test_nemotron_3_5_perf_recipes_inherit_nemotron_3_policy(
     cfg = recipe_factory()
     base_cfg = base_recipe_factory()
 
+    base_env = base_cfg.env_vars
+    if recipe_factory in _NCCLEP_RECIPES and base_cfg.model.moe_flex_dispatcher_backend == "hybridep":
+        # The GB200 Lightning benchmarks dispatch through NCCL EP while their Nemotron 3 Nano bases keep HybridEP.
+        base_env = {name: value for name, value in base_env.items() if name not in HYBRID_EP_ENV_NAMES}
+        base_env["NCCL_EP_HT_EM_PULL_PUSH"] = 1
+
     if recipe_factory in _B_MXFP8_RECIPES:
-        assert cfg.env_vars == {**base_cfg.env_vars, "NVTE_CUTEDSL_FUSED_GROUPED_MLP": 1}
+        assert cfg.env_vars == {**base_env, "NVTE_CUTEDSL_FUSED_GROUPED_MLP": 1}
     elif recipe_factory in _FULL_ITERATION_MXFP8_RECIPES:
         assert cfg.env_vars == {
-            **base_cfg.env_vars,
+            **base_env,
             "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True,graph_capture_record_stream_reuse:True",
             "TORCH_NCCL_AVOID_RECORD_STREAMS": 0,
         }
     elif recipe_factory is not nemotron_3_5_lightning_pretrain_16gpu_h100_bf16_config:
-        assert cfg.env_vars == base_cfg.env_vars
+        assert cfg.env_vars == base_env
     assert cfg.model.calculate_per_token_loss == base_cfg.model.calculate_per_token_loss
     assert cfg.model.use_te_rng_tracker == base_cfg.model.use_te_rng_tracker
     assert cfg.tokenizer.tokenizer_model != base_cfg.tokenizer.tokenizer_model