From 823b90ac5d7b0c2f3517fd9a2bfc29e114222e5c Mon Sep 17 00:00:00 2001 From: Malay Nagda Date: Fri, 9 Oct 2026 16:15:54 +0530 Subject: [PATCH 1/9] perf(recipe): default GB200 MoE perf recipes to NCCL EP and grouped tensors Switch the GB200 MoE perf recipes that nemo-ci benchmarks on HybridEP to the NCCL EP flex dispatcher backend with moe_use_grouped_tensor=True, matching the GB300 recipes moved in #6171 and GPT-OSS 120B GB200 in #6261. The NCCL EP environment sets NCCL_EP_HT_EM_PULL_PUSH=1 and drops the HybridEP topology names. Switched: DeepSeek V3 256-GPU BF16/MXFP8/NVFP4/MXFP8 large-scale, Kimi K2 MXFP8, Nemotron 3 Super BF16/MXFP8/NVFP4, Nemotron 3 Ultra MXFP8, Nemotron 3 Nano BF16/MXFP8, Nemotron 3.5 Lightning BF16/MXFP8, Qwen3 30B-A3B 8-GPU BF16/MXFP8 and 32-GPU BF16, Qwen3 235B-A22B 256-GPU BF16/MXFP8/NVFP4/MXFP8 large-scale. Recipes whose GB300 counterpart stays on HybridEP (DeepSeek V4 Flash, Nemotron 3 Ultra NVFP4) and GB200 recipes without a nemo-ci benchmark keep HybridEP. DeepSeek V3 NVFP4 now builds on a shared HybridEP BF16 base so the NCCL EP overlay is applied explicitly per recipe. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Malay Nagda --- src/megatron/bridge/perf_recipes/_common.py | 2 +- .../deepseek/gb200/deepseek_v3.py | 59 +++++----- .../bridge/perf_recipes/kimi/gb200/kimi_k2.py | 14 ++- .../perf_recipes/nemotronh/gb200/nemotronh.py | 100 ++++++++-------- .../perf_recipes/qwen/gb200/qwen3_moe.py | 84 ++++++------- .../recipes/test_perf_config_integration.py | 7 +- ...st_bseries_nvfp4_full_iteration_recipes.py | 8 +- .../recipes/test_deepseek_recipes.py | 8 +- .../recipes/test_gb200_ncclep_perf_recipes.py | 110 ++++++++++++++++++ .../test_moe_full_iteration_recipes.py | 2 +- ...est_nemotron_3_5_lightning_perf_recipes.py | 19 ++- .../recipes/test_nemotronh_recipes.py | 15 ++- .../recipes/test_perf_recipe_environment.py | 4 +- tests/unit_tests/recipes/test_qwen_recipes.py | 14 ++- .../recipes/test_r050_perf_recipe_ports.py | 14 ++- 15 files changed, 314 insertions(+), 146 deletions(-) create mode 100644 tests/unit_tests/recipes/test_gb200_ncclep_perf_recipes.py diff --git a/src/megatron/bridge/perf_recipes/_common.py b/src/megatron/bridge/perf_recipes/_common.py index 97b4b71548..31db3bde78 100644 --- a/src/megatron/bridge/perf_recipes/_common.py +++ b/src/megatron/bridge/perf_recipes/_common.py @@ -96,7 +96,7 @@ def _benchmark_common( def _enable_ncclep(cfg: ConfigContainer) -> None: - """Switch a flex-dispatcher recipe to the NCCL EP backend (GB300 MoE default). + """Switch a flex-dispatcher recipe to the NCCL EP backend (GB200 and GB300 MoE default). Call after ``_benchmark_common``. Only the dispatch stack changes; parallelism, precision, recompute and the CUDA graph mode stay as the recipe set them. The calling recipe still declares diff --git a/src/megatron/bridge/perf_recipes/deepseek/gb200/deepseek_v3.py b/src/megatron/bridge/perf_recipes/deepseek/gb200/deepseek_v3.py index 9ecdc99160..cc17b9bbd5 100644 --- a/src/megatron/bridge/perf_recipes/deepseek/gb200/deepseek_v3.py +++ b/src/megatron/bridge/perf_recipes/deepseek/gb200/deepseek_v3.py @@ -13,6 +13,7 @@ # limitations under the License. """GB200 performance recipes for DeepSeek V3.""" +from megatron.bridge.perf_recipes._common import _enable_ncclep from megatron.bridge.perf_recipes.deepseek.common import ( ConfigContainer, _benchmark_common, @@ -26,8 +27,8 @@ from megatron.bridge.perf_recipes.environment import COMMON_PERF_ENV_VARS -def deepseek_v3_pretrain_256gpu_gb200_bf16_config() -> ConfigContainer: - """DeepSeek V3 pretrain: 256× GB200, BF16.""" +def _build_deepseek_v3_gb200_bf16() -> ConfigContainer: + """Shared HybridEP BF16 base for the DeepSeek V3 256-GPU GB200 BF16 and NVFP4 recipes.""" cfg = deepseek_v3_pretrain_config() cfg.mixed_precision = _perf_precision("bf16") _deepseek_v3_common(cfg) @@ -52,6 +53,16 @@ def deepseek_v3_pretrain_256gpu_gb200_bf16_config() -> ConfigContainer: set_deepseek_v3_pipeline_model_parallel_layout(cfg.model) _benchmark_common(cfg) + return cfg + + +def deepseek_v3_pretrain_256gpu_gb200_bf16_config() -> ConfigContainer: + """DeepSeek V3 pretrain: 256× GB200, BF16, NCCL EP.""" + cfg = _build_deepseek_v3_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. + 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, @@ -67,11 +78,8 @@ def deepseek_v3_pretrain_256gpu_gb200_bf16_config() -> ConfigContainer: "TORCH_CPP_LOG_LEVEL": "ERROR", # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # 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, + # NCCL EP dispatcher mode. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, @@ -133,7 +141,7 @@ def deepseek_v3_pretrain_256gpu_gb200_fp8cs_config() -> ConfigContainer: def deepseek_v3_pretrain_256gpu_gb200_fp8mx_config() -> ConfigContainer: - """DeepSeek V3 pretrain: 256× GB200, MXFP8.""" + """DeepSeek V3 pretrain: 256× GB200, MXFP8, NCCL EP.""" cfg = deepseek_v3_pretrain_config() cfg.mixed_precision = _perf_precision("fp8_mx") _deepseek_v3_common(cfg) @@ -161,6 +169,10 @@ def deepseek_v3_pretrain_256gpu_gb200_fp8mx_config() -> ConfigContainer: _enable_deepseek_full_iteration(cfg) cfg.model.fp8_output_proj = True cfg.mixed_precision.fp8_dot_product_attention = False + _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, @@ -172,11 +184,8 @@ def deepseek_v3_pretrain_256gpu_gb200_fp8mx_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 0, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # 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, + # NCCL EP dispatcher mode. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "CUDNNFE_CLUSTER_OVERLAP_MARGIN": 8, "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, @@ -192,8 +201,8 @@ def deepseek_v3_pretrain_256gpu_gb200_fp8mx_config() -> ConfigContainer: def deepseek_v3_pretrain_256gpu_gb200_nvfp4_config() -> ConfigContainer: - """DeepSeek V3 pretrain: 256× GB200, NVFP4 with optimizer CUDA graph.""" - cfg = deepseek_v3_pretrain_256gpu_gb200_bf16_config() + """DeepSeek V3 pretrain: 256× GB200, NVFP4 with optimizer CUDA graph, NCCL EP.""" + cfg = _build_deepseek_v3_gb200_bf16() cfg.mixed_precision = _perf_precision("nvfp4") _enable_deepseek_full_iteration(cfg) cfg.model.fp8_output_proj = False @@ -204,6 +213,10 @@ def deepseek_v3_pretrain_256gpu_gb200_nvfp4_config() -> ConfigContainer: cfg.model.mla_down_proj_fusion = True cfg.model.moe_paged_stash_buffer_size_factor_cuda = 1.1 cfg.model.recompute_modules = ["mla_up_proj", "mlp"] + _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, @@ -215,11 +228,8 @@ def deepseek_v3_pretrain_256gpu_gb200_nvfp4_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # 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, + # NCCL EP dispatcher mode. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, @@ -235,7 +245,7 @@ def deepseek_v3_pretrain_256gpu_gb200_nvfp4_config() -> ConfigContainer: def deepseek_v3_pretrain_256gpu_gb200_fp8mx_large_scale_config() -> ConfigContainer: - """DeepSeek V3 pretrain: 256× GB200, MXFP8, large-scale proxy (GBS=256).""" + """DeepSeek V3 pretrain: 256× GB200, MXFP8, large-scale proxy (GBS=256), NCCL EP.""" cfg = deepseek_v3_pretrain_256gpu_gb200_fp8mx_config() cfg.train.global_batch_size = 256 # Keep process settings next to the recipe so users can see the exact benchmark environment. @@ -249,11 +259,8 @@ def deepseek_v3_pretrain_256gpu_gb200_fp8mx_large_scale_config() -> ConfigContai "TORCH_NCCL_AVOID_RECORD_STREAMS": 0, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # 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, + # NCCL EP dispatcher mode (inherited from the 256-GPU MXFP8 recipe). + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "CUDNNFE_CLUSTER_OVERLAP_MARGIN": 8, "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, diff --git a/src/megatron/bridge/perf_recipes/kimi/gb200/kimi_k2.py b/src/megatron/bridge/perf_recipes/kimi/gb200/kimi_k2.py index a064edbaad..0c5dadd9f3 100644 --- a/src/megatron/bridge/perf_recipes/kimi/gb200/kimi_k2.py +++ b/src/megatron/bridge/perf_recipes/kimi/gb200/kimi_k2.py @@ -13,6 +13,7 @@ # 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, @@ -138,7 +139,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.""" + """Kimi K2 pretrain: 256× GB200, MXFP8, NCCL EP.""" cfg = kimi_k2_pretrain_config() cfg.mixed_precision = _perf_precision("fp8_mx") cfg.mixed_precision.reuse_grad_buf_for_mxfp8_param_ag = False @@ -173,6 +174,10 @@ 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, @@ -184,11 +189,8 @@ 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, - # 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, + # NCCL EP dispatcher mode. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, diff --git a/src/megatron/bridge/perf_recipes/nemotronh/gb200/nemotronh.py b/src/megatron/bridge/perf_recipes/nemotronh/gb200/nemotronh.py index 4372401e6c..005a64a974 100644 --- a/src/megatron/bridge/perf_recipes/nemotronh/gb200/nemotronh.py +++ b/src/megatron/bridge/perf_recipes/nemotronh/gb200/nemotronh.py @@ -13,6 +13,7 @@ # limitations under the License. """GB200 performance recipes for NemotronH and Nemotron 3.""" +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.nemotronh.common import ( _TE_QUANT_CFG_PATH, @@ -76,7 +77,7 @@ def nemotronh_56b_pretrain_64gpu_gb200_fp8cs_config() -> ConfigContainer: def nemotron_3_super_pretrain_64gpu_gb200_bf16_config() -> ConfigContainer: - """Nemotron 3 Super pretrain: 64× GB200, BF16.""" + """Nemotron 3 Super pretrain: 64× GB200, BF16, NCCL EP.""" cfg = nemotron_3_super_pretrain_config() cfg.mixed_precision = _perf_precision("bf16") @@ -98,6 +99,10 @@ def nemotron_3_super_pretrain_64gpu_gb200_bf16_config() -> ConfigContainer: cfg.model.cuda_graph_scope = ["attn", "mamba", "moe_router", "moe_preprocess"] _apply_nemotron_3_super_perf_defaults(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, @@ -109,11 +114,8 @@ def nemotron_3_super_pretrain_64gpu_gb200_bf16_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # 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, + # NCCL EP dispatcher mode. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, @@ -122,7 +124,7 @@ def nemotron_3_super_pretrain_64gpu_gb200_bf16_config() -> ConfigContainer: def nemotron_3_super_pretrain_64gpu_gb200_fp8mx_config() -> ConfigContainer: - """Nemotron 3 Super pretrain: 64× GB200, MXFP8 with full-iteration CUDA graph.""" + """Nemotron 3 Super pretrain: 64× GB200, MXFP8 with full-iteration CUDA graph, NCCL EP.""" cfg = nemotron_3_super_pretrain_config() cfg.mixed_precision = _perf_precision("fp8_mx") @@ -144,6 +146,10 @@ def nemotron_3_super_pretrain_64gpu_gb200_fp8mx_config() -> ConfigContainer: _apply_nemotron_3_super_perf_defaults(cfg) _enable_nemotron_3_super_full_iteration(cfg) 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, @@ -155,11 +161,8 @@ def nemotron_3_super_pretrain_64gpu_gb200_fp8mx_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 0, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # HybridEP topology for the target system. - "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 16, - "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, - "NVLINK_DOMAIN_SIZE": 72, - "USE_MNNVL": 1, + # NCCL EP dispatcher mode. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "CUDNNFE_CLUSTER_OVERLAP_MARGIN": 8, "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, @@ -170,7 +173,7 @@ def nemotron_3_super_pretrain_64gpu_gb200_fp8mx_config() -> ConfigContainer: def nemotron_3_super_pretrain_64gpu_gb200_nvfp4_config() -> ConfigContainer: - """Nemotron 3 Super pretrain: 64× GB200, NVFP4 with full-iteration CUDA graph.""" + """Nemotron 3 Super pretrain: 64× GB200, NVFP4 with full-iteration CUDA graph, NCCL EP.""" cfg = nemotron_3_super_pretrain_config() cfg.mixed_precision = _nemotron_3_super_nvfp4_precision() @@ -193,6 +196,10 @@ def nemotron_3_super_pretrain_64gpu_gb200_nvfp4_config() -> ConfigContainer: _apply_nemotron_3_super_perf_defaults(cfg) _enable_nemotron_3_super_full_iteration(cfg) cfg.mixed_precision.fp8_dot_product_attention = False + _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, @@ -204,11 +211,8 @@ def nemotron_3_super_pretrain_64gpu_gb200_nvfp4_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 0, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # HybridEP topology for the target system. - "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 16, - "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, - "NVLINK_DOMAIN_SIZE": 72, - "USE_MNNVL": 1, + # NCCL EP dispatcher mode. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "CUDNNFE_CLUSTER_OVERLAP_MARGIN": 8, "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, @@ -224,14 +228,13 @@ def nemotron_3_ultra_pretrain_256gpu_gb200_fp8mx_config() -> ConfigContainer: """Nemotron 3 Ultra (550B-A55B LatentMoE) pretrain: 256× GB200, MXFP8, Megatron-FSDP (HSDP). TP2 + SP (due to smaller GB200 HBM) / PP1 / CP1 / EP64 / ETP1, GBS 256 / MBS 1, seq 8192, BF16 + MXFP8 mixed - precision, HybridEP flex dispatcher, CuteDSL fused grouped MLP, selective recompute + + precision, NCCL EP flex dispatcher, CuteDSL fused grouped MLP, selective recompute + fine-grained activation offload of the expert MLP, MTP=2. """ num_gpus = 256 expert_model_parallel_size = 64 global_batch_size = 256 - hybrid_ep_ranks_per_nvlink_domain = 64 cfg = nemotron_3_ultra_pretrain_config() cfg.mixed_precision = _perf_precision("fp8_mx") @@ -274,6 +277,10 @@ def nemotron_3_ultra_pretrain_256gpu_gb200_fp8mx_config() -> ConfigContainer: cfg.model.recompute_granularity = "selective" cfg.model.recompute_modules = ["moe_act"] + _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, @@ -288,11 +295,8 @@ def nemotron_3_ultra_pretrain_256gpu_gb200_fp8mx_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # HybridEP topology for the target system. - "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": hybrid_ep_ranks_per_nvlink_domain, - "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, - "NVLINK_DOMAIN_SIZE": 72, - "USE_MNNVL": 1, + # NCCL EP dispatcher mode. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, @@ -397,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.""" + """Nemotron 3 Nano pretrain: 8× GB200, BF16, NCCL EP.""" cfg = nemotron_3_nano_pretrain_config() _apply_nemotron_3_nano_perf_defaults(cfg) cfg.mixed_precision = _perf_precision("bf16") @@ -421,6 +425,10 @@ 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, @@ -432,11 +440,8 @@ 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, - # 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, + # NCCL EP dispatcher mode. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, @@ -475,11 +480,15 @@ 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.""" + """Nemotron 3 Nano pretrain: 8× GB200, MXFP8, NCCL EP.""" 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, @@ -491,11 +500,8 @@ 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, - # 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, + # NCCL EP dispatcher mode. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "CUDNNFE_CLUSTER_OVERLAP_MARGIN": 8, "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, @@ -562,7 +568,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.""" + """Nemotron 3.5 Lightning pretrain: 8× GB200, BF16, NCCL EP (inherited from Nemotron 3 Nano).""" cfg = nemotron_3_nano_pretrain_8gpu_gb200_bf16_config() cfg.model.recompute_modules = [] cfg.model.mtp_num_layers = 2 @@ -581,10 +587,8 @@ def nemotron_3_5_lightning_pretrain_8gpu_gb200_bf16_config() -> ConfigContainer: "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, "NCCL_NVLS_ENABLE": 0, - "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 8, - "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, - "NVLINK_DOMAIN_SIZE": 72, - "USE_MNNVL": 1, + # NCCL EP dispatcher mode (inherited from the Nemotron 3 Nano BF16 recipe). + "NCCL_EP_HT_EM_PULL_PUSH": 1, "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, "NVTE_NORM_BWD_USE_CUDNN": 1, @@ -609,11 +613,15 @@ def _build_nemotron_3_5_lightning_gb200_mxfp8() -> ConfigContainer: def nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_config() -> ConfigContainer: - """Nemotron 3.5 Lightning pretrain: 8× GB200, MXFP8.""" + """Nemotron 3.5 Lightning pretrain: 8× GB200, MXFP8, NCCL EP.""" cfg = _build_nemotron_3_5_lightning_gb200_mxfp8() cfg.model.recompute_modules = [] cfg.model.use_transformer_engine_op_fuser = True 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 cfg.env_vars = { **COMMON_PERF_ENV_VARS, "CUDA_DEVICE_MAX_CONNECTIONS": 32, @@ -621,10 +629,8 @@ def nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_config() -> ConfigContainer "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, "NCCL_NVLS_ENABLE": 0, - "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 8, - "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, - "NVLINK_DOMAIN_SIZE": 72, - "USE_MNNVL": 1, + # NCCL EP dispatcher mode. + "NCCL_EP_HT_EM_PULL_PUSH": 1, "CUDNNFE_CLUSTER_OVERLAP_MARGIN": 8, "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_CUTEDSL_FUSED_GROUPED_MLP": 1, diff --git a/src/megatron/bridge/perf_recipes/qwen/gb200/qwen3_moe.py b/src/megatron/bridge/perf_recipes/qwen/gb200/qwen3_moe.py index 8f5bdc3f76..e101e54caf 100644 --- a/src/megatron/bridge/perf_recipes/qwen/gb200/qwen3_moe.py +++ b/src/megatron/bridge/perf_recipes/qwen/gb200/qwen3_moe.py @@ -15,6 +15,7 @@ import torch +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.qwen.common import ( CommOverlapConfig, @@ -196,7 +197,7 @@ def qwen3_235b_a22b_pretrain_64gpu_gb200_fp8mx_config() -> ConfigContainer: def qwen3_30b_a3b_pretrain_8gpu_gb200_bf16_config() -> ConfigContainer: - """Qwen3 30B-A3B pretrain: 8× GB200, BF16, EP=8.""" + """Qwen3 30B-A3B pretrain: 8× GB200, BF16, EP=8, NCCL EP.""" cfg = qwen3_30b_a3b_pretrain_config() cfg.mixed_precision = _perf_precision("bf16") cfg.model.bias_activation_fusion = True @@ -227,6 +228,10 @@ def qwen3_30b_a3b_pretrain_8gpu_gb200_bf16_config() -> ConfigContainer: cfg.comm_overlap = CommOverlapConfig(tp_comm_overlap=False) _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, @@ -238,11 +243,8 @@ def qwen3_30b_a3b_pretrain_8gpu_gb200_bf16_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # 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, + # NCCL EP dispatcher mode. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, @@ -305,7 +307,7 @@ def qwen3_30b_a3b_pretrain_8gpu_gb200_fp8cs_config() -> ConfigContainer: def qwen3_30b_a3b_pretrain_8gpu_gb200_fp8mx_config() -> ConfigContainer: - """Qwen3 30B-A3B pretrain benchmark: GB200 MXFP8 recipe plus benchmark overrides.""" + """Qwen3 30B-A3B pretrain benchmark: GB200 MXFP8 recipe plus benchmark overrides, NCCL EP.""" cfg = _qwen3_30b_a3b_pretrain_8gpu_gb200_fp8mx_config() _benchmark_common(cfg) @@ -313,6 +315,10 @@ def qwen3_30b_a3b_pretrain_8gpu_gb200_fp8mx_config() -> ConfigContainer: # Full-iteration graphs remain benchmark-only because they alter natural-routing auxiliary loss. _enable_hybridep_full_iteration_mxfp8(cfg) cfg.model.recompute_modules = [] + _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, @@ -324,11 +330,8 @@ def qwen3_30b_a3b_pretrain_8gpu_gb200_fp8mx_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 0, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # 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, + # NCCL EP dispatcher mode. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "CUDNNFE_CLUSTER_OVERLAP_MARGIN": 8, "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, @@ -339,7 +342,7 @@ def qwen3_30b_a3b_pretrain_8gpu_gb200_fp8mx_config() -> ConfigContainer: def qwen3_235b_a22b_pretrain_256gpu_gb200_bf16_config() -> ConfigContainer: - """Qwen3 235B-A22B pretrain: 256× GB200, BF16, PP=8 EP=8.""" + """Qwen3 235B-A22B pretrain: 256× GB200, BF16, PP=8 EP=8, NCCL EP.""" cfg = qwen3_235b_a22b_pretrain_config() cfg.mixed_precision = _perf_precision("bf16") cfg.model.bias_activation_fusion = True @@ -370,6 +373,10 @@ def qwen3_235b_a22b_pretrain_256gpu_gb200_bf16_config() -> ConfigContainer: cfg.comm_overlap = CommOverlapConfig(tp_comm_overlap=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, @@ -385,11 +392,8 @@ def qwen3_235b_a22b_pretrain_256gpu_gb200_bf16_config() -> ConfigContainer: "TORCH_CPP_LOG_LEVEL": "ERROR", # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # 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, + # NCCL EP dispatcher mode. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, @@ -452,7 +456,7 @@ def qwen3_235b_a22b_pretrain_256gpu_gb200_fp8cs_config() -> ConfigContainer: def qwen3_235b_a22b_pretrain_256gpu_gb200_fp8mx_config() -> ConfigContainer: - """Qwen3 235B-A22B pretrain: 256× GB200, MXFP8, PP=8 EP=32.""" + """Qwen3 235B-A22B pretrain: 256× GB200, MXFP8, PP=8 EP=32, NCCL EP.""" cfg = qwen3_235b_a22b_pretrain_config() cfg.mixed_precision = _perf_precision("fp8_mx") cfg.model.bias_activation_fusion = True @@ -479,6 +483,10 @@ def qwen3_235b_a22b_pretrain_256gpu_gb200_fp8mx_config() -> ConfigContainer: _benchmark_common(cfg) _enable_hybridep_full_iteration_mxfp8(cfg) cfg.model.recompute_modules = [] + _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, @@ -490,11 +498,8 @@ def qwen3_235b_a22b_pretrain_256gpu_gb200_fp8mx_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 0, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # HybridEP topology for the target system. - "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 32, - "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, - "NVLINK_DOMAIN_SIZE": 72, - "USE_MNNVL": 1, + # NCCL EP dispatcher mode. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "CUDNNFE_CLUSTER_OVERLAP_MARGIN": 8, "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, @@ -505,7 +510,7 @@ def qwen3_235b_a22b_pretrain_256gpu_gb200_fp8mx_config() -> ConfigContainer: def qwen3_235b_a22b_pretrain_256gpu_gb200_fp8mx_large_scale_config() -> ConfigContainer: - """Qwen3 235B A22B pretrain: 256× GB200, FP8-MX, large-scale proxy (GBS=512).""" + """Qwen3 235B A22B pretrain: 256× GB200, FP8-MX, large-scale proxy (GBS=512), NCCL EP.""" cfg = qwen3_235b_a22b_pretrain_256gpu_gb200_fp8mx_config() cfg.train.global_batch_size = 512 # Keep process settings next to the recipe so users can see the exact benchmark environment. @@ -519,11 +524,8 @@ def qwen3_235b_a22b_pretrain_256gpu_gb200_fp8mx_large_scale_config() -> ConfigCo "TORCH_NCCL_AVOID_RECORD_STREAMS": 0, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # HybridEP topology for the target system. - "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 32, - "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, - "NVLINK_DOMAIN_SIZE": 72, - "USE_MNNVL": 1, + # NCCL EP dispatcher mode (inherited from the 256-GPU MXFP8 recipe). + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "CUDNNFE_CLUSTER_OVERLAP_MARGIN": 8, "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, @@ -589,13 +591,17 @@ def qwen3_235b_a22b_pretrain_64gpu_gb200_nvfp4_full_iteration_config() -> Config def qwen3_235b_a22b_pretrain_256gpu_gb200_nvfp4_config() -> ConfigContainer: - """Qwen3 235B-A22B pretrain: 256× GB200, NVFP4, full-iteration CG and A2A overlap.""" + """Qwen3 235B-A22B pretrain: 256× GB200, NVFP4, full-iteration CG and A2A overlap, NCCL EP.""" cfg = qwen3_235b_a22b_pretrain_256gpu_gb200_fp8cs_config() cfg.mixed_precision = _perf_precision("nvfp4") cfg.comm_overlap.tp_comm_overlap = False cfg.model.virtual_pipeline_model_parallel_size = 3 _enable_hybridep_full_iteration_nvfp4(cfg) cfg.model.recompute_modules = [] + _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, @@ -607,11 +613,8 @@ def qwen3_235b_a22b_pretrain_256gpu_gb200_nvfp4_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 0, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # HybridEP topology for the target system. - "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 32, - "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, - "NVLINK_DOMAIN_SIZE": 72, - "USE_MNNVL": 1, + # NCCL EP dispatcher mode. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "CUDNNFE_CLUSTER_OVERLAP_MARGIN": 8, "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, @@ -652,7 +655,7 @@ 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.""" + """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) # Keep process settings next to the recipe so users can see the exact benchmark environment. cfg.env_vars = { @@ -665,11 +668,8 @@ 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, - # 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, + # NCCL EP dispatcher mode (inherited from the 8-GPU BF16 recipe). + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, diff --git a/tests/functional_tests/test_groups/recipes/test_perf_config_integration.py b/tests/functional_tests/test_groups/recipes/test_perf_config_integration.py index 7e136ba5f6..a65955fa8e 100644 --- a/tests/functional_tests/test_groups/recipes/test_perf_config_integration.py +++ b/tests/functional_tests/test_groups/recipes/test_perf_config_integration.py @@ -191,8 +191,7 @@ def test_workload_base_config_derives_from_flat_recipe(self): ), { "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, - "NVLINK_DOMAIN_SIZE": 72, - "USE_MNNVL": 1, + "NCCL_EP_HT_EM_PULL_PUSH": 1, "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, }, ), @@ -221,14 +220,14 @@ def test_workload_base_config_derives_from_flat_recipe(self): "megatron.bridge.perf_recipes.nemotronh.gb200.nemotronh:" "nemotron_3_super_pretrain_64gpu_gb200_bf16_config" ), - {"NVTE_FWD_LAYERNORM_SM_MARGIN": 20, "NVLINK_DOMAIN_SIZE": 72, "USE_MNNVL": 1}, + {"NVTE_FWD_LAYERNORM_SM_MARGIN": 20, "NCCL_EP_HT_EM_PULL_PUSH": 1}, ), ( "qwen", ("megatron.bridge.perf_recipes.qwen.gb200.qwen3_moe:qwen3_30b_a3b_pretrain_8gpu_gb200_fp8mx_config"), { "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, - "NVLINK_DOMAIN_SIZE": 72, + "NCCL_EP_HT_EM_PULL_PUSH": 1, "TORCH_NCCL_AVOID_RECORD_STREAMS": 0, }, ), diff --git a/tests/unit_tests/recipes/test_bseries_nvfp4_full_iteration_recipes.py b/tests/unit_tests/recipes/test_bseries_nvfp4_full_iteration_recipes.py index cfbb44e41f..aa115da2a2 100644 --- a/tests/unit_tests/recipes/test_bseries_nvfp4_full_iteration_recipes.py +++ b/tests/unit_tests/recipes/test_bseries_nvfp4_full_iteration_recipes.py @@ -30,8 +30,12 @@ pytestmark = pytest.mark.unit -# GB300 perf recipes dispatch through NCCL EP (no HybridEP topology in their environment). -_NCCLEP_RECIPES = {qwen3_235b_a22b_pretrain_256gpu_gb300_nvfp4_config} +# The 256-GPU GB200 and GB300 perf recipes dispatch through NCCL EP (no HybridEP topology in their +# environment); the 64-GPU full-iteration variants keep HybridEP. +_NCCLEP_RECIPES = { + qwen3_235b_a22b_pretrain_256gpu_gb200_nvfp4_config, + qwen3_235b_a22b_pretrain_256gpu_gb300_nvfp4_config, +} @pytest.fixture(autouse=True) diff --git a/tests/unit_tests/recipes/test_deepseek_recipes.py b/tests/unit_tests/recipes/test_deepseek_recipes.py index bbf03377e2..f3f885fa61 100644 --- a/tests/unit_tests/recipes/test_deepseek_recipes.py +++ b/tests/unit_tests/recipes/test_deepseek_recipes.py @@ -29,6 +29,7 @@ from megatron.core.transformer.pipeline_parallel_layer_layout import PipelineParallelLayerLayout from megatron.bridge.models.model_provider import ModelProviderMixin +from megatron.bridge.perf_recipes.environment import HYBRID_EP_ENV_NAMES from megatron.bridge.recipes.deepseek import ( set_deepseek_v3_pipeline_model_parallel_layout, set_deepseek_v4_pipeline_model_parallel_layout, @@ -187,9 +188,10 @@ def test_model_and_perf_deepseek_recipes_bake_environment_defaults(monkeypatch: assert "QUANTIZATION_TYPE_DEBUG" not in recipe_config.env_vars assert "TORCHINDUCTOR_WORKER_START" not in perf_config.env_vars assert "QUANTIZATION_TYPE_DEBUG" not in perf_config.env_vars - assert perf_config.env_vars["NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN"] == 64 - assert perf_config.env_vars["NVLINK_DOMAIN_SIZE"] == 72 - assert perf_config.env_vars["USE_MNNVL"] == 1 + # The GB200 perf recipe dispatches through NCCL EP, so it carries no HybridEP topology. + assert perf_config.model.moe_flex_dispatcher_backend == "ncclep" + assert perf_config.env_vars["NCCL_EP_HT_EM_PULL_PUSH"] == 1 + assert perf_config.env_vars.keys().isdisjoint(HYBRID_EP_ENV_NAMES) def test_deepseek_v3_pipeline_layout_can_place_mtp_in_standalone_stage(): diff --git a/tests/unit_tests/recipes/test_gb200_ncclep_perf_recipes.py b/tests/unit_tests/recipes/test_gb200_ncclep_perf_recipes.py new file mode 100644 index 0000000000..8811e34bd6 --- /dev/null +++ b/tests/unit_tests/recipes/test_gb200_ncclep_perf_recipes.py @@ -0,0 +1,110 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Every GB200 MoE perf recipe switched to NCCL EP carries the full NCCL EP stack, and the GB200 recipes +left on HybridEP, several of which share a base with a switched recipe, do not pick it up.""" + +import importlib + +import pytest + +from megatron.bridge.perf_recipes.environment import HYBRID_EP_ENV_NAMES +from tests.unit_tests.recipes.recipe_test_utils import patch_recipe_construction_dependencies + + +pytestmark = pytest.mark.unit + +_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. +_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. +_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"), + ("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_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"), +) + + +# Import the recipe modules at collection time: the offline-construction fixture patches ``AutoBridge`` in +# every recipe module that is already loaded, so a lazy import inside a test would escape the patch. +_RECIPE_MODULES = { + module_name: importlib.import_module(f"megatron.bridge.perf_recipes.{module_name}") + for module_name in sorted({name for name, _ in (*_GB200_NCCLEP_RECIPES, *_GB200_HYBRIDEP_RECIPES)}) +} + + +def _build(module_name: str, factory_name: str): + return getattr(_RECIPE_MODULES[module_name], factory_name)() + + +@pytest.fixture(autouse=True) +def _keep_recipe_construction_offline(monkeypatch: pytest.MonkeyPatch) -> None: + patch_recipe_construction_dependencies(monkeypatch) + + +@pytest.mark.parametrize(("module_name", "factory_name"), _GB200_NCCLEP_RECIPES, ids=lambda value: value) +def test_gb200_moe_perf_recipes_default_to_the_nccl_ep_stack(module_name: str, factory_name: str) -> None: + cfg = _build(module_name, factory_name) + + assert cfg.model.moe_token_dispatcher_type == "flex" + assert cfg.model.moe_flex_dispatcher_backend == "ncclep" + # Device-side expert token counts; the legacy grouped-MLP path host-syncs tokens_per_expert per layer. + assert cfg.model.moe_use_grouped_tensor is True + # The plugin runs in hierarchical-topology pull mode, and nothing reads the HybridEP tuning anymore. + assert {name: cfg.env_vars[name] for name in _NCCL_EP_ENV} == _NCCL_EP_ENV + assert cfg.env_vars.keys().isdisjoint(HYBRID_EP_ENV_NAMES) + + +@pytest.mark.parametrize(("module_name", "factory_name"), _GB200_HYBRIDEP_RECIPES, ids=lambda value: value) +def test_gb200_hybridep_recipes_do_not_inherit_nccl_ep(module_name: str, factory_name: str) -> None: + cfg = _build(module_name, factory_name) + + assert cfg.model.moe_token_dispatcher_type == "flex" + assert cfg.model.moe_flex_dispatcher_backend == "hybridep" + assert cfg.env_vars.keys().isdisjoint(_NCCL_EP_ENV) + assert cfg.env_vars.keys() >= HYBRID_EP_ENV_NAMES diff --git a/tests/unit_tests/recipes/test_moe_full_iteration_recipes.py b/tests/unit_tests/recipes/test_moe_full_iteration_recipes.py index 23904cf812..cbc800f4ac 100644 --- a/tests/unit_tests/recipes/test_moe_full_iteration_recipes.py +++ b/tests/unit_tests/recipes/test_moe_full_iteration_recipes.py @@ -48,7 +48,7 @@ def _assert_full_iteration_hybridep(cfg: ConfigContainer) -> None: assert cfg.ddp.check_for_nan_in_grad is False assert cfg.rerun_state_machine.check_for_nan_in_loss is False - # GB300 recipes default to the NCCL EP dispatch stack; every other GPU keeps HybridEP. + # GB200 and GB300 recipes default to the NCCL EP dispatch stack; every other GPU keeps HybridEP. if cfg.model.moe_flex_dispatcher_backend == "ncclep": assert cfg.model.moe_use_grouped_tensor is True assert cfg.env_vars["NCCL_EP_HT_EM_PULL_PUSH"] == 1 diff --git a/tests/unit_tests/recipes/test_nemotron_3_5_lightning_perf_recipes.py b/tests/unit_tests/recipes/test_nemotron_3_5_lightning_perf_recipes.py index fd84cbf23a..dacfb00881 100644 --- a/tests/unit_tests/recipes/test_nemotron_3_5_lightning_perf_recipes.py +++ b/tests/unit_tests/recipes/test_nemotron_3_5_lightning_perf_recipes.py @@ -99,8 +99,11 @@ nemotron_3_5_lightning_pretrain_8gpu_gb300_fp8mx_config, nemotron_3_5_lightning_pretrain_8gpu_gb300_nvfp4_config, ) -# GB300 perf recipes covered by nemo-ci dispatch through NCCL EP (no HybridEP topology in their environment). +# GB200 and GB300 perf recipes covered by nemo-ci dispatch through NCCL EP (no HybridEP topology in their +# environment). _NCCLEP_RECIPES = ( + nemotron_3_5_lightning_pretrain_8gpu_gb200_bf16_config, + nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_config, nemotron_3_5_lightning_pretrain_8gpu_gb300_bf16_config, nemotron_3_5_lightning_pretrain_8gpu_gb300_fp8mx_config, ) @@ -531,6 +534,9 @@ def test_bf16_perf_recipes_share_training_workload() -> None: gb200_cfg.model.moe_permute_fusion_into_hybridep = h100_cfg.model.moe_permute_fusion_into_hybridep gb200_cfg.model.moe_hybridep_num_sms = h100_cfg.model.moe_hybridep_num_sms gb200_cfg.model.moe_flex_dispatcher_num_sms = h100_cfg.model.moe_flex_dispatcher_num_sms + # The GB200 benchmark dispatches through NCCL EP with device-side expert counts. + gb200_cfg.model.moe_flex_dispatcher_backend = h100_cfg.model.moe_flex_dispatcher_backend + gb200_cfg.model.moe_use_grouped_tensor = h100_cfg.model.moe_use_grouped_tensor gb200_cfg.optimizer.use_precision_aware_optimizer = h100_cfg.optimizer.use_precision_aware_optimizer gb200_cfg.optimizer.exp_avg_dtype = h100_cfg.optimizer.exp_avg_dtype gb200_cfg.optimizer.exp_avg_sq_dtype = h100_cfg.optimizer.exp_avg_sq_dtype @@ -585,8 +591,15 @@ def test_gb200_perf_recipe_topology(recipe_factory: Callable[[], ConfigContainer assert cfg.model.seq_length == 8192 assert cfg.dataset.seq_length == 8192 assert cfg.model.moe_hybridep_num_sms == 16 - assert cfg.env_vars["NVLINK_DOMAIN_SIZE"] == 72 - assert cfg.env_vars["USE_MNNVL"] == 1 + if recipe_factory in _NCCLEP_RECIPES: + assert cfg.model.moe_flex_dispatcher_backend == "ncclep" + assert cfg.model.moe_use_grouped_tensor is True + assert cfg.env_vars["NCCL_EP_HT_EM_PULL_PUSH"] == 1 + assert "NVLINK_DOMAIN_SIZE" not in cfg.env_vars + else: + assert cfg.model.moe_flex_dispatcher_backend == "hybridep" + assert cfg.env_vars["NVLINK_DOMAIN_SIZE"] == 72 + assert cfg.env_vars["USE_MNNVL"] == 1 if recipe_factory is nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_config: assert cfg.model.cuda_graph_impl == "transformer_engine" diff --git a/tests/unit_tests/recipes/test_nemotronh_recipes.py b/tests/unit_tests/recipes/test_nemotronh_recipes.py index 8fea02dda5..13d70f2e3c 100644 --- a/tests/unit_tests/recipes/test_nemotronh_recipes.py +++ b/tests/unit_tests/recipes/test_nemotronh_recipes.py @@ -26,6 +26,7 @@ import pytest import torch +from megatron.bridge.perf_recipes.environment import HYBRID_EP_ENV_NAMES from megatron.bridge.training.utils.omegaconf_utils import OverridesError, process_config_with_overrides from tests.unit_tests.recipes.recipe_test_utils import patch_recipe_module_global @@ -251,7 +252,6 @@ def test_nemotron_3_super_64gpu_gb200_matches_benchmark_hardware_configuration() "context_parallel_size", "expert_tensor_parallel_size", "expert_model_parallel_size", - "moe_flex_dispatcher_backend", "moe_hybridep_num_sms", "moe_token_dispatcher_type", "moe_shared_expert_overlap", @@ -270,7 +270,18 @@ def test_nemotron_3_super_64gpu_gb200_matches_benchmark_hardware_configuration() assert training_cfg.train.global_batch_size == benchmark_cfg.train.global_batch_size == 512 assert training_cfg.train.micro_batch_size == benchmark_cfg.train.micro_batch_size == 1 - assert training_cfg.env_vars == benchmark_cfg.env_vars + + # The GB200 benchmark dispatches through NCCL EP; the training recipe keeps HybridEP. Apart from the + # dispatcher topology settings, both declare the same process environment. + assert training_cfg.model.moe_flex_dispatcher_backend == "hybridep" + assert benchmark_cfg.model.moe_flex_dispatcher_backend == "ncclep" + assert benchmark_cfg.model.moe_use_grouped_tensor is True + assert training_cfg.env_vars.keys() >= HYBRID_EP_ENV_NAMES + assert benchmark_cfg.env_vars["NCCL_EP_HT_EM_PULL_PUSH"] == 1 + dispatcher_env_names = HYBRID_EP_ENV_NAMES | {"NCCL_EP_HT_EM_PULL_PUSH"} + assert {name: value for name, value in training_cfg.env_vars.items() if name not in dispatcher_env_names} == { + name: value for name, value in benchmark_cfg.env_vars.items() if name not in dispatcher_env_names + } assert training_cfg.model.moe_router_force_load_balancing is False assert benchmark_cfg.model.moe_router_force_load_balancing is True diff --git a/tests/unit_tests/recipes/test_perf_recipe_environment.py b/tests/unit_tests/recipes/test_perf_recipe_environment.py index 91a2832067..9606fc71e2 100644 --- a/tests/unit_tests/recipes/test_perf_recipe_environment.py +++ b/tests/unit_tests/recipes/test_perf_recipe_environment.py @@ -248,10 +248,8 @@ def test_explicit_environment_invariants_across_all_flat_recipes(): "deepseek/gb200/deepseek_v3.py", "deepseek_v3_pretrain_256gpu_gb200_bf16_config", { - "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 64, - "NVLINK_DOMAIN_SIZE": 72, + "NCCL_EP_HT_EM_PULL_PUSH": 1, "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, - "USE_MNNVL": 1, }, ), ( diff --git a/tests/unit_tests/recipes/test_qwen_recipes.py b/tests/unit_tests/recipes/test_qwen_recipes.py index 631b8d701a..b8bce64916 100644 --- a/tests/unit_tests/recipes/test_qwen_recipes.py +++ b/tests/unit_tests/recipes/test_qwen_recipes.py @@ -27,6 +27,7 @@ import pytest import torch +from megatron.bridge.perf_recipes.environment import HYBRID_EP_ENV_NAMES from tests.unit_tests.recipes.recipe_test_utils import patch_recipe_module_global @@ -795,16 +796,25 @@ def test_qwen3_30b_a3b_gb200_fp8mx_perf_recipe_uses_main_recipe( assert perf_cfg.model.sequence_parallel == main_cfg.model.sequence_parallel assert perf_cfg.train.global_batch_size == main_cfg.train.global_batch_size assert perf_cfg.train.micro_batch_size == main_cfg.train.micro_batch_size - assert perf_cfg.model.moe_flex_dispatcher_backend == main_cfg.model.moe_flex_dispatcher_backend assert perf_cfg.model.moe_token_dispatcher_type == main_cfg.model.moe_token_dispatcher_type assert perf_cfg.model.moe_a2a_overlap == main_cfg.model.moe_a2a_overlap assert perf_cfg.comm_overlap == main_cfg.comm_overlap - assert perf_cfg.env_vars == main_cfg.env_vars assert perf_cfg.model.offload_modules == main_cfg.model.offload_modules == [] # Benchmark-only policy remains outside the main recipe. assert main_cfg.model.moe_router_force_load_balancing is False assert perf_cfg.model.moe_router_force_load_balancing is True + # The benchmark dispatches through NCCL EP; the main recipe keeps HybridEP. Apart from the + # dispatcher topology settings, both declare the same process environment. + assert main_cfg.model.moe_flex_dispatcher_backend == "hybridep" + assert perf_cfg.model.moe_flex_dispatcher_backend == "ncclep" + assert perf_cfg.model.moe_use_grouped_tensor is True + assert main_cfg.env_vars.keys() >= HYBRID_EP_ENV_NAMES + assert perf_cfg.env_vars["NCCL_EP_HT_EM_PULL_PUSH"] == 1 + dispatcher_env_names = HYBRID_EP_ENV_NAMES | {"NCCL_EP_HT_EM_PULL_PUSH"} + assert {name: value for name, value in perf_cfg.env_vars.items() if name not in dispatcher_env_names} == { + name: value for name, value in main_cfg.env_vars.items() if name not in dispatcher_env_names + } assert main_cfg.train.train_iters == 100 assert perf_cfg.train.train_iters == 50 assert main_cfg.train.global_batch_size == 512 diff --git a/tests/unit_tests/recipes/test_r050_perf_recipe_ports.py b/tests/unit_tests/recipes/test_r050_perf_recipe_ports.py index b6f6519725..3ba80fa709 100644 --- a/tests/unit_tests/recipes/test_r050_perf_recipe_ports.py +++ b/tests/unit_tests/recipes/test_r050_perf_recipe_ports.py @@ -43,8 +43,10 @@ def _keep_recipe_construction_offline(monkeypatch: pytest.MonkeyPatch) -> None: patch_recipe_construction_dependencies(monkeypatch) -def _assert_full_iteration_hybridep_mxfp8(cfg, *, fp8_dot_product_attention: bool = True) -> None: - """Check the normalized r0.5.0 full-iteration HybridEP settings.""" +def _assert_full_iteration_hybridep_mxfp8( + cfg, *, fp8_dot_product_attention: bool = True, dispatcher_backend: str = "hybridep" +) -> None: + """Check the normalized r0.5.0 full-iteration settings; NCCL EP recipes keep the rest of the stack.""" assert cfg.model.cuda_graph_impl == "full_iteration" assert cfg.model.cuda_graph_scope == [] assert cfg.model.use_te_rng_tracker is True @@ -56,7 +58,7 @@ def _assert_full_iteration_hybridep_mxfp8(cfg, *, fp8_dot_product_attention: boo assert cfg.model.moe_paged_stash_buffer_size_factor_cuda == 1.2 assert cfg.model.moe_paged_stash_buffer_size_factor_cpu == 1.0 assert cfg.model.moe_shared_expert_overlap is False - assert cfg.model.moe_flex_dispatcher_backend == "hybridep" + assert cfg.model.moe_flex_dispatcher_backend == dispatcher_backend assert cfg.model.moe_token_dispatcher_type == "flex" assert cfg.model.moe_hybridep_num_sms == 32 assert cfg.model.high_priority_a2a_comm_stream is True @@ -131,7 +133,11 @@ def test_qwen_large_scale_recipes_only_override_global_batch_size(recipe) -> Non def test_deepseek_v3_gb200_large_scale_matches_r050_fp8mx_base() -> None: cfg = deepseek_v3_pretrain_256gpu_gb200_fp8mx_large_scale_config() - _assert_full_iteration_hybridep_mxfp8(cfg, fp8_dot_product_attention=False) + # GB200 keeps the r0.5.0 full-iteration MXFP8 stack but dispatches through NCCL EP. + _assert_full_iteration_hybridep_mxfp8(cfg, fp8_dot_product_attention=False, dispatcher_backend="ncclep") + assert cfg.model.moe_use_grouped_tensor is True + assert cfg.env_vars["NCCL_EP_HT_EM_PULL_PUSH"] == 1 + assert "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN" not in cfg.env_vars assert cfg.train.global_batch_size == 256 assert cfg.model.pipeline_model_parallel_size == 4 assert cfg.model.virtual_pipeline_model_parallel_size == 4 From 268f9127b48c3a882acb8f28215ca91da690d677 Mon Sep 17 00:00:00 2001 From: Malay Nagda Date: Fri, 9 Oct 2026 16:31:38 +0530 Subject: [PATCH 2/9] 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 Signed-off-by: Malay Nagda --- .../bridge/perf_recipes/kimi/gb200/kimi_k2.py | 14 ++++---- .../perf_recipes/nemotronh/gb200/nemotronh.py | 34 ++++++++++--------- .../perf_recipes/qwen/gb200/qwen3_moe.py | 21 ++++++++---- .../recipes/test_gb200_ncclep_perf_recipes.py | 15 ++++---- ...est_nemotron_3_5_lightning_perf_recipes.py | 13 +++++-- 5 files changed, 56 insertions(+), 41 deletions(-) diff --git a/src/megatron/bridge/perf_recipes/kimi/gb200/kimi_k2.py b/src/megatron/bridge/perf_recipes/kimi/gb200/kimi_k2.py index 0c5dadd9f3..a064edbaad 100644 --- a/src/megatron/bridge/perf_recipes/kimi/gb200/kimi_k2.py +++ b/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, diff --git a/src/megatron/bridge/perf_recipes/nemotronh/gb200/nemotronh.py b/src/megatron/bridge/perf_recipes/nemotronh/gb200/nemotronh.py index 005a64a974..73ecddd473 100644 --- a/src/megatron/bridge/perf_recipes/nemotronh/gb200/nemotronh.py +++ b/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,6 +578,10 @@ 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, @@ -587,7 +589,7 @@ def nemotron_3_5_lightning_pretrain_8gpu_gb200_bf16_config() -> ConfigContainer: "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, diff --git a/src/megatron/bridge/perf_recipes/qwen/gb200/qwen3_moe.py b/src/megatron/bridge/perf_recipes/qwen/gb200/qwen3_moe.py index e101e54caf..6de91218f1 100644 --- a/src/megatron/bridge/perf_recipes/qwen/gb200/qwen3_moe.py +++ b/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, diff --git a/tests/unit_tests/recipes/test_gb200_ncclep_perf_recipes.py b/tests/unit_tests/recipes/test_gb200_ncclep_perf_recipes.py index 8811e34bd6..5fc4c27951 100644 --- a/tests/unit_tests/recipes/test_gb200_ncclep_perf_recipes.py +++ b/tests/unit_tests/recipes/test_gb200_ncclep_perf_recipes.py @@ -26,26 +26,21 @@ _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"), @@ -53,17 +48,21 @@ ) # 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"), diff --git a/tests/unit_tests/recipes/test_nemotron_3_5_lightning_perf_recipes.py b/tests/unit_tests/recipes/test_nemotron_3_5_lightning_perf_recipes.py index dacfb00881..ec1f4fbd9a 100644 --- a/tests/unit_tests/recipes/test_nemotron_3_5_lightning_perf_recipes.py +++ b/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 From 27c0f0fb31127db7969207e5b106f6ea52dcbe2a Mon Sep 17 00:00:00 2001 From: Malay Nagda Date: Fri, 9 Oct 2026 17:19:16 +0530 Subject: [PATCH 3/9] perf(recipe): use NCCL EP for GPT-OSS 120B BF16 and DeepSeek V4 Flash on GB200 Switch two more GB200 MoE perf recipes to the NCCL EP flex dispatcher backend with moe_use_grouped_tensor=True and NCCL_EP_HT_EM_PULL_PUSH=1: - GPT-OSS 120B 64-GPU BF16, which previously ran the alltoall dispatcher inherited from its base recipe. EP A2A overlap and delayed wgrad are kept. - DeepSeek V4 Flash 128-GPU MXFP8, which previously ran HybridEP. The DeepSeek V4 Flash GB300 recipe derives from the GB200 one, and the VR200 and B300 recipes derive from GB300, so the GB200 settings move into a shared HybridEP builder that GB300 now uses; those recipes keep HybridEP. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Malay Nagda --- .../deepseek/gb200/deepseek_v4.py | 21 ++++++++++++------- .../deepseek/gb300/deepseek_v4.py | 4 ++-- .../perf_recipes/gpt_oss/gb200/gpt_oss.py | 8 ++++++- .../test_deepseek_v4_flash_perf_recipes.py | 11 ++++++---- .../recipes/test_gb200_ncclep_perf_recipes.py | 2 ++ 5 files changed, 32 insertions(+), 14 deletions(-) diff --git a/src/megatron/bridge/perf_recipes/deepseek/gb200/deepseek_v4.py b/src/megatron/bridge/perf_recipes/deepseek/gb200/deepseek_v4.py index 126ca64ad2..2854b7e61d 100644 --- a/src/megatron/bridge/perf_recipes/deepseek/gb200/deepseek_v4.py +++ b/src/megatron/bridge/perf_recipes/deepseek/gb200/deepseek_v4.py @@ -18,7 +18,7 @@ from megatron.bridge.models.deepseek.deepseek_v4_bridge import ( set_deepseek_v4_pipeline_model_parallel_layout, ) -from megatron.bridge.perf_recipes._common import _benchmark_common +from megatron.bridge.perf_recipes._common import _benchmark_common, _enable_ncclep from megatron.bridge.perf_recipes.environment import COMMON_PERF_ENV_VARS from megatron.bridge.recipes.deepseek.gb200.deepseek_v4 import ( deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config, @@ -27,8 +27,8 @@ from megatron.bridge.utils.cuda_graph import set_full_iteration_cuda_graph -def deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_config() -> ConfigContainer: - """DeepSeek V4 Flash pretrain: 128× GB200, MXFP8, full-iteration CUDA graph.""" +def _build_deepseek_v4_flash_gb200_fp8mx() -> ConfigContainer: + """Shared HybridEP MXFP8 base for the DeepSeek V4 Flash GB200 recipe and its GB300 derivative.""" cfg = deepseek_v4_flash_pretrain_64gpu_gb200_fp8mx_config() cfg.model.tensor_model_parallel_size = 1 @@ -111,6 +111,16 @@ def deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_config() -> ConfigContainer: cfg.checkpoint.load_optim = False cfg.checkpoint.load_rng = False cfg.checkpoint.save_optim = False + return cfg + + +def deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_config() -> ConfigContainer: + """DeepSeek V4 Flash pretrain: 128× GB200, MXFP8, full-iteration CUDA graph, NCCL EP.""" + cfg = _build_deepseek_v4_flash_gb200_fp8mx() + _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, @@ -119,10 +129,7 @@ def deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_config() -> ConfigContainer: "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True,graph_capture_record_stream_reuse:True", "TORCH_NCCL_AVOID_RECORD_STREAMS": 0, "NCCL_NVLS_ENABLE": 0, - "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 32, - "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, - "NVLINK_DOMAIN_SIZE": 72, - "USE_MNNVL": 1, + "NCCL_EP_HT_EM_PULL_PUSH": 1, "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_CUTEDSL_FUSED_GROUPED_MLP": 1, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, diff --git a/src/megatron/bridge/perf_recipes/deepseek/gb300/deepseek_v4.py b/src/megatron/bridge/perf_recipes/deepseek/gb300/deepseek_v4.py index 0cba00e8cb..82514d2685 100644 --- a/src/megatron/bridge/perf_recipes/deepseek/gb300/deepseek_v4.py +++ b/src/megatron/bridge/perf_recipes/deepseek/gb300/deepseek_v4.py @@ -20,7 +20,7 @@ ) from megatron.bridge.perf_recipes._common import _benchmark_common from megatron.bridge.perf_recipes.deepseek.gb200.deepseek_v4 import ( - deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_config, + _build_deepseek_v4_flash_gb200_fp8mx, ) from megatron.bridge.perf_recipes.environment import COMMON_PERF_ENV_VARS from megatron.bridge.recipes.deepseek.gb300.deepseek_v4 import ( @@ -31,7 +31,7 @@ def deepseek_v4_flash_pretrain_128gpu_gb300_fp8mx_config() -> ConfigContainer: """DeepSeek V4 Flash pretrain: 128× GB300, MXFP8.""" - cfg = deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_config() + cfg = _build_deepseek_v4_flash_gb200_fp8mx() cfg.model.expert_model_parallel_size = 32 cfg.train.micro_batch_size = 2 diff --git a/src/megatron/bridge/perf_recipes/gpt_oss/gb200/gpt_oss.py b/src/megatron/bridge/perf_recipes/gpt_oss/gb200/gpt_oss.py index 79e9c411f4..48fe1c37bd 100644 --- a/src/megatron/bridge/perf_recipes/gpt_oss/gb200/gpt_oss.py +++ b/src/megatron/bridge/perf_recipes/gpt_oss/gb200/gpt_oss.py @@ -180,7 +180,7 @@ def gpt_oss_20b_pretrain_512gpu_gb200_fp8mx_config() -> ConfigContainer: def gpt_oss_120b_pretrain_64gpu_gb200_bf16_config() -> ConfigContainer: - """GPT-OSS 120B pretrain: 64× GB200, BF16, GBS=1280.""" + """GPT-OSS 120B pretrain: 64× GB200, BF16, GBS=1280, NCCL EP.""" cfg = gpt_oss_120b_pretrain_config() cfg.mixed_precision = _perf_precision("bf16") cfg.model.moe_router_fusion = True @@ -201,6 +201,10 @@ def gpt_oss_120b_pretrain_64gpu_gb200_bf16_config() -> ConfigContainer: cfg.comm_overlap.overlap_moe_expert_parallel_comm = 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, @@ -212,6 +216,8 @@ def gpt_oss_120b_pretrain_64gpu_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, # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, diff --git a/tests/unit_tests/recipes/test_deepseek_v4_flash_perf_recipes.py b/tests/unit_tests/recipes/test_deepseek_v4_flash_perf_recipes.py index 7f8d5ae988..03b90faea7 100644 --- a/tests/unit_tests/recipes/test_deepseek_v4_flash_perf_recipes.py +++ b/tests/unit_tests/recipes/test_deepseek_v4_flash_perf_recipes.py @@ -26,6 +26,7 @@ deepseek_v4_flash_pretrain_128gpu_gb300_fp8mx_config, deepseek_v4_flash_pretrain_128gpu_vr200_fp8mx_config, ) +from megatron.bridge.perf_recipes.environment import HYBRID_EP_ENV_NAMES from megatron.bridge.utils.cuda_graph import cuda_graph_module_names, is_full_iteration_cuda_graph from tests.unit_tests.recipes.recipe_test_utils import patch_recipe_construction_dependencies @@ -55,7 +56,8 @@ def test_deepseek_v4_flash_128gpu_gb200_fp8mx_config() -> None: assert cfg.train.micro_batch_size == 1 assert cfg.model.moe_token_dispatcher_type == "flex" - assert cfg.model.moe_flex_dispatcher_backend == "hybridep" + assert cfg.model.moe_flex_dispatcher_backend == "ncclep" + assert cfg.model.moe_use_grouped_tensor is True assert cfg.model.moe_router_force_load_balancing is True assert cfg.model.moe_flex_dispatcher_num_sms == 32 assert cfg.model.moe_hybridep_num_sms is None @@ -115,9 +117,8 @@ def test_deepseek_v4_flash_128gpu_gb200_fp8mx_config() -> None: assert cfg.env_vars["NVTE_FWD_LAYERNORM_SM_MARGIN"] == 20 assert cfg.env_vars["NVTE_BWD_LAYERNORM_SM_MARGIN"] == 20 assert cfg.env_vars["NVTE_CUTEDSL_FUSED_GROUPED_MLP"] == 1 - assert cfg.env_vars["NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN"] == 32 - assert cfg.env_vars["NVLINK_DOMAIN_SIZE"] == 72 - assert cfg.env_vars["USE_MNNVL"] == 1 + assert cfg.env_vars["NCCL_EP_HT_EM_PULL_PUSH"] == 1 + assert cfg.env_vars.keys().isdisjoint(HYBRID_EP_ENV_NAMES) def test_deepseek_v4_flash_128gpu_gb300_fp8mx_config() -> None: @@ -130,6 +131,8 @@ def test_deepseek_v4_flash_128gpu_gb300_fp8mx_config() -> None: assert cfg.train.global_batch_size == 2048 assert cfg.train.micro_batch_size == 2 assert cfg.model.recompute_modules == ["mla_up_proj"] + # GB300 builds on the shared HybridEP base, not on the NCCL EP GB200 recipe. + assert cfg.model.moe_flex_dispatcher_backend == "hybridep" assert cfg.model.moe_flex_dispatcher_num_sms == 32 assert cfg.model.moe_hybridep_num_sms is None assert is_full_iteration_cuda_graph(cfg.model) diff --git a/tests/unit_tests/recipes/test_gb200_ncclep_perf_recipes.py b/tests/unit_tests/recipes/test_gb200_ncclep_perf_recipes.py index 5fc4c27951..144f38a150 100644 --- a/tests/unit_tests/recipes/test_gb200_ncclep_perf_recipes.py +++ b/tests/unit_tests/recipes/test_gb200_ncclep_perf_recipes.py @@ -32,6 +32,8 @@ ("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"), + ("deepseek.gb200.deepseek_v4", "deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_config"), + ("gpt_oss.gb200.gpt_oss", "gpt_oss_120b_pretrain_64gpu_gb200_bf16_config"), ("gpt_oss.gb200.gpt_oss", "gpt_oss_120b_pretrain_64gpu_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"), From 6a880fbab8e386817f4eed175cdfc3d0ad32ad77 Mon Sep 17 00:00:00 2001 From: Malay Nagda Date: Wed, 14 Oct 2026 16:08:09 +0530 Subject: [PATCH 4/9] perf(recipe): keep GPT-OSS 120B and Qwen3 30B-A3B GB200 perf recipes unchanged Drop the GPT-OSS 120B and Qwen3 30B-A3B parts of the GB200 NCCL EP switch. These recipes, their environments and their tests match main again: - GPT-OSS 120B 64-GPU BF16 runs the alltoall dispatcher again. - Qwen3 30B-A3B 8-GPU BF16 and MXFP8 run HybridEP again. The 32-GPU BF16 recipe derives from the 8-GPU recipe as on main, and the shared BF16 builder is removed. The GB200 NCCL EP recipe test no longer lists these recipes. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Malay Nagda --- .../perf_recipes/gpt_oss/gb200/gpt_oss.py | 8 +---- .../perf_recipes/qwen/gb200/qwen3_moe.py | 36 ++++++++----------- .../recipes/test_perf_config_integration.py | 2 +- .../recipes/test_gb200_ncclep_perf_recipes.py | 19 ++++------ tests/unit_tests/recipes/test_qwen_recipes.py | 14 ++------ 5 files changed, 24 insertions(+), 55 deletions(-) diff --git a/src/megatron/bridge/perf_recipes/gpt_oss/gb200/gpt_oss.py b/src/megatron/bridge/perf_recipes/gpt_oss/gb200/gpt_oss.py index 48fe1c37bd..79e9c411f4 100644 --- a/src/megatron/bridge/perf_recipes/gpt_oss/gb200/gpt_oss.py +++ b/src/megatron/bridge/perf_recipes/gpt_oss/gb200/gpt_oss.py @@ -180,7 +180,7 @@ def gpt_oss_20b_pretrain_512gpu_gb200_fp8mx_config() -> ConfigContainer: def gpt_oss_120b_pretrain_64gpu_gb200_bf16_config() -> ConfigContainer: - """GPT-OSS 120B pretrain: 64× GB200, BF16, GBS=1280, NCCL EP.""" + """GPT-OSS 120B pretrain: 64× GB200, BF16, GBS=1280.""" cfg = gpt_oss_120b_pretrain_config() cfg.mixed_precision = _perf_precision("bf16") cfg.model.moe_router_fusion = True @@ -201,10 +201,6 @@ def gpt_oss_120b_pretrain_64gpu_gb200_bf16_config() -> ConfigContainer: cfg.comm_overlap.overlap_moe_expert_parallel_comm = 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, @@ -216,8 +212,6 @@ def gpt_oss_120b_pretrain_64gpu_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, # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, diff --git a/src/megatron/bridge/perf_recipes/qwen/gb200/qwen3_moe.py b/src/megatron/bridge/perf_recipes/qwen/gb200/qwen3_moe.py index 6de91218f1..4bfd389098 100644 --- a/src/megatron/bridge/perf_recipes/qwen/gb200/qwen3_moe.py +++ b/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 _build_qwen3_30b_a3b_gb200_bf16() -> ConfigContainer: - """Shared HybridEP BF16 base for the Qwen3 30B-A3B 8-GPU and 32-GPU GB200 recipes.""" +def qwen3_30b_a3b_pretrain_8gpu_gb200_bf16_config() -> ConfigContainer: + """Qwen3 30B-A3B pretrain: 8× GB200, BF16, EP=8.""" cfg = qwen3_30b_a3b_pretrain_config() cfg.mixed_precision = _perf_precision("bf16") cfg.model.bias_activation_fusion = True @@ -228,16 +228,6 @@ def _build_qwen3_30b_a3b_gb200_bf16() -> 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. - 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, @@ -249,8 +239,11 @@ def qwen3_30b_a3b_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, @@ -313,7 +306,7 @@ def qwen3_30b_a3b_pretrain_8gpu_gb200_fp8cs_config() -> ConfigContainer: def qwen3_30b_a3b_pretrain_8gpu_gb200_fp8mx_config() -> ConfigContainer: - """Qwen3 30B-A3B pretrain benchmark: GB200 MXFP8 recipe plus benchmark overrides, NCCL EP.""" + """Qwen3 30B-A3B pretrain benchmark: GB200 MXFP8 recipe plus benchmark overrides.""" cfg = _qwen3_30b_a3b_pretrain_8gpu_gb200_fp8mx_config() _benchmark_common(cfg) @@ -321,10 +314,6 @@ def qwen3_30b_a3b_pretrain_8gpu_gb200_fp8mx_config() -> ConfigContainer: # Full-iteration graphs remain benchmark-only because they alter natural-routing auxiliary loss. _enable_hybridep_full_iteration_mxfp8(cfg) cfg.model.recompute_modules = [] - _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, @@ -336,8 +325,11 @@ def qwen3_30b_a3b_pretrain_8gpu_gb200_fp8mx_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 0, # 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, @@ -662,7 +654,7 @@ 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.""" - cfg = _with_global_batch_size(_build_qwen3_30b_a3b_gb200_bf16(), 2048) + cfg = _with_global_batch_size(qwen3_30b_a3b_pretrain_8gpu_gb200_bf16_config(), 2048) # Keep process settings next to the recipe so users can see the exact benchmark environment. cfg.env_vars = { **COMMON_PERF_ENV_VARS, diff --git a/tests/functional_tests/test_groups/recipes/test_perf_config_integration.py b/tests/functional_tests/test_groups/recipes/test_perf_config_integration.py index a65955fa8e..d0d87dafa6 100644 --- a/tests/functional_tests/test_groups/recipes/test_perf_config_integration.py +++ b/tests/functional_tests/test_groups/recipes/test_perf_config_integration.py @@ -227,7 +227,7 @@ def test_workload_base_config_derives_from_flat_recipe(self): ("megatron.bridge.perf_recipes.qwen.gb200.qwen3_moe:qwen3_30b_a3b_pretrain_8gpu_gb200_fp8mx_config"), { "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, - "NCCL_EP_HT_EM_PULL_PUSH": 1, + "NVLINK_DOMAIN_SIZE": 72, "TORCH_NCCL_AVOID_RECORD_STREAMS": 0, }, ), diff --git a/tests/unit_tests/recipes/test_gb200_ncclep_perf_recipes.py b/tests/unit_tests/recipes/test_gb200_ncclep_perf_recipes.py index 144f38a150..74c2468705 100644 --- a/tests/unit_tests/recipes/test_gb200_ncclep_perf_recipes.py +++ b/tests/unit_tests/recipes/test_gb200_ncclep_perf_recipes.py @@ -11,8 +11,9 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Every GB200 MoE perf recipe switched to NCCL EP carries the full NCCL EP stack, and the GB200 recipes -left on HybridEP, several of which share a base with a switched recipe, do not pick it up.""" +"""The GB200 DeepSeek, Nemotron 3 and Qwen3 235B perf recipes switched to NCCL EP carry the full NCCL EP +stack, and the GB200 recipes left on HybridEP, several of which share a base with a switched recipe, do not +pick it up.""" import importlib @@ -26,23 +27,19 @@ _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. +# (module under megatron.bridge.perf_recipes, recipe factory) for the GB200 recipes switched 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"), ("deepseek.gb200.deepseek_v4", "deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_config"), - ("gpt_oss.gb200.gpt_oss", "gpt_oss_120b_pretrain_64gpu_gb200_bf16_config"), - ("gpt_oss.gb200.gpt_oss", "gpt_oss_120b_pretrain_64gpu_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_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_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"), @@ -50,8 +47,8 @@ ) # GB200 recipes that stay on HybridEP. Most build on the same base or parent recipe as a recipe above (the -# 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. +# shared DeepSeek V3 BF16 builder, the Nemotron 3 Nano parents of the Lightning recipes, the Qwen3 235B FP8-CS +# parent of its NVFP4 recipe), 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"), @@ -62,10 +59,6 @@ ("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"), ) diff --git a/tests/unit_tests/recipes/test_qwen_recipes.py b/tests/unit_tests/recipes/test_qwen_recipes.py index de038fa3b8..a692f5c830 100644 --- a/tests/unit_tests/recipes/test_qwen_recipes.py +++ b/tests/unit_tests/recipes/test_qwen_recipes.py @@ -27,7 +27,6 @@ import pytest import torch -from megatron.bridge.perf_recipes.environment import HYBRID_EP_ENV_NAMES from tests.unit_tests.recipes.recipe_test_utils import patch_recipe_module_global @@ -800,25 +799,16 @@ def test_qwen3_30b_a3b_gb200_fp8mx_perf_recipe_uses_main_recipe( assert perf_cfg.model.sequence_parallel == main_cfg.model.sequence_parallel assert perf_cfg.train.global_batch_size == main_cfg.train.global_batch_size assert perf_cfg.train.micro_batch_size == main_cfg.train.micro_batch_size + assert perf_cfg.model.moe_flex_dispatcher_backend == main_cfg.model.moe_flex_dispatcher_backend assert perf_cfg.model.moe_token_dispatcher_type == main_cfg.model.moe_token_dispatcher_type assert perf_cfg.model.moe_a2a_overlap == main_cfg.model.moe_a2a_overlap assert perf_cfg.comm_overlap == main_cfg.comm_overlap + assert perf_cfg.env_vars == main_cfg.env_vars assert perf_cfg.model.offload_modules == main_cfg.model.offload_modules == [] # Benchmark-only policy remains outside the main recipe. assert main_cfg.model.moe_router_force_load_balancing is False assert perf_cfg.model.moe_router_force_load_balancing is True - # The benchmark dispatches through NCCL EP; the main recipe keeps HybridEP. Apart from the - # dispatcher topology settings, both declare the same process environment. - assert main_cfg.model.moe_flex_dispatcher_backend == "hybridep" - assert perf_cfg.model.moe_flex_dispatcher_backend == "ncclep" - assert perf_cfg.model.moe_use_grouped_tensor is True - assert main_cfg.env_vars.keys() >= HYBRID_EP_ENV_NAMES - assert perf_cfg.env_vars["NCCL_EP_HT_EM_PULL_PUSH"] == 1 - dispatcher_env_names = HYBRID_EP_ENV_NAMES | {"NCCL_EP_HT_EM_PULL_PUSH"} - assert {name: value for name, value in perf_cfg.env_vars.items() if name not in dispatcher_env_names} == { - name: value for name, value in main_cfg.env_vars.items() if name not in dispatcher_env_names - } assert main_cfg.train.train_iters == 100 assert perf_cfg.train.train_iters == 50 assert main_cfg.train.global_batch_size == 512 From 270ed1078440e325ca340a6f06d802b29176c24f Mon Sep 17 00:00:00 2001 From: Malay Nagda Date: Fri, 16 Oct 2026 13:03:18 +0530 Subject: [PATCH 5/9] perf(recipe): switch VRNVL72 MoE perf recipes from HybridEP to NCCL EP Default 41 VRNVL72 MoE perf recipes to the NCCL EP dispatch stack used on GB200 and GB300: the ncclep flex backend, device-side expert token counts (moe_use_grouped_tensor) and NCCL_EP_HT_EM_PULL_PUSH=1 in place of the HybridEP topology environment. The switch covers DeepSeek V3 (including the 64- and 4-GPU proxies) and V4 Flash, GPT-OSS 20B NVFP4 and 120B MXFP8, Kimi K2, Nemotron 3 Nano and Super, Nemotron 3.5 Lightning, Qwen3 235B and 30B-A3B, and Qwen3.5-VL 35B, 122B and 397B. The Nemotron 3 Ultra recipes were already on NCCL EP. Only the dispatch stack changes: parallelism, precision, recompute, CUDA graphs and NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 on the DeepSeek recipes stay as they were. The Qwen3.5-VL MXFP8 recipes drop their HybridEP-only permute-fusion and preprocessing-SM settings, which NCCL EP does not read. GPT-OSS 20B 8-GPU MXFP8 runs EP=1 and stays on HybridEP; GPT-OSS 120B BF16 stays on alltoall. The VRNVL72 recipes join the GB300 NCCL EP test, and the GB200 NCCL EP test now expects Nemotron 3 Ultra NVFP4 on NCCL EP, which main switched. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Malay Nagda --- .../deepseek/vrnvl72/deepseek_v3.py | 135 +++++++++-------- .../deepseek/vrnvl72/deepseek_v4.py | 10 +- .../perf_recipes/gpt_oss/vrnvl72/gpt_oss.py | 40 ++--- .../perf_recipes/kimi/vrnvl72/kimi_k2.py | 27 ++-- .../nemotronh/vrnvl72/nemotronh.py | 100 +++++++------ .../perf_recipes/qwen/vrnvl72/qwen3_moe.py | 78 +++++----- .../perf_recipes/qwen_vl/vrnvl72/qwen35_vl.py | 138 +++++++++--------- .../recipes/test_deepseek_v3_proxy_recipes.py | 3 +- .../test_deepseek_v4_flash_perf_recipes.py | 11 +- .../recipes/test_gb200_ncclep_perf_recipes.py | 2 +- .../recipes/test_gb300_ncclep_perf_recipes.py | 71 ++++++--- .../test_moe_full_iteration_recipes.py | 5 +- ...est_nemotron_3_5_lightning_perf_recipes.py | 20 +-- .../test_qwen35_vl_workload_base_configs.py | 22 ++- 14 files changed, 360 insertions(+), 302 deletions(-) diff --git a/src/megatron/bridge/perf_recipes/deepseek/vrnvl72/deepseek_v3.py b/src/megatron/bridge/perf_recipes/deepseek/vrnvl72/deepseek_v3.py index cc16ca9564..bd166d96d1 100644 --- a/src/megatron/bridge/perf_recipes/deepseek/vrnvl72/deepseek_v3.py +++ b/src/megatron/bridge/perf_recipes/deepseek/vrnvl72/deepseek_v3.py @@ -13,6 +13,7 @@ # limitations under the License. """VRNVL72 performance recipes for DeepSeek V3.""" +from megatron.bridge.perf_recipes._common import _enable_ncclep from megatron.bridge.perf_recipes.deepseek.common import ( ConfigContainer, _benchmark_common, @@ -32,7 +33,7 @@ def deepseek_v3_pretrain_128gpu_vrnvl72_bf16_config() -> ConfigContainer: - """DeepSeek V3 pretrain: 128× VRNVL72, BF16.""" + """DeepSeek V3 pretrain: 128× VRNVL72, BF16, NCCL EP.""" cfg = deepseek_v3_pretrain_config() cfg.mixed_precision = _perf_precision("bf16") _deepseek_v3_common(cfg) @@ -57,6 +58,10 @@ def deepseek_v3_pretrain_128gpu_vrnvl72_bf16_config() -> ConfigContainer: set_deepseek_v3_pipeline_model_parallel_layout(cfg.model) _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, @@ -68,11 +73,8 @@ def deepseek_v3_pretrain_128gpu_vrnvl72_bf16_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # 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, + # NCCL EP dispatcher mode and one GPU per rank. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, @@ -83,7 +85,7 @@ def deepseek_v3_pretrain_128gpu_vrnvl72_bf16_config() -> ConfigContainer: def deepseek_v3_pretrain_128gpu_vrnvl72_fp8cs_config() -> ConfigContainer: - """DeepSeek V3 pretrain: 128× VRNVL72, FP8 current-scaling.""" + """DeepSeek V3 pretrain: 128× VRNVL72, FP8 current-scaling, NCCL EP.""" cfg = deepseek_v3_pretrain_config() cfg.mixed_precision = _perf_precision("fp8_cs") _deepseek_v3_common(cfg) @@ -108,6 +110,10 @@ def deepseek_v3_pretrain_128gpu_vrnvl72_fp8cs_config() -> ConfigContainer: set_deepseek_v3_pipeline_model_parallel_layout(cfg.model) _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, @@ -119,11 +125,8 @@ def deepseek_v3_pretrain_128gpu_vrnvl72_fp8cs_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # 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, + # NCCL EP dispatcher mode and one GPU per rank. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, @@ -134,7 +137,7 @@ def deepseek_v3_pretrain_128gpu_vrnvl72_fp8cs_config() -> ConfigContainer: def deepseek_v3_pretrain_128gpu_vrnvl72_fp8mx_config() -> ConfigContainer: - """DeepSeek V3 pretrain: 128× VRNVL72, MXFP8.""" + """DeepSeek V3 pretrain: 128× VRNVL72, MXFP8, NCCL EP.""" cfg = deepseek_v3_pretrain_config() cfg.mixed_precision = _perf_precision("fp8_mx") _deepseek_v3_common(cfg) @@ -157,6 +160,10 @@ def deepseek_v3_pretrain_128gpu_vrnvl72_fp8mx_config() -> ConfigContainer: _benchmark_common(cfg) _enable_deepseek_full_iteration(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 cfg.model.fp8_output_proj = False cfg.mixed_precision.fp8_dot_product_attention = False # Keep process settings next to the recipe so users can see the exact benchmark environment. @@ -170,11 +177,8 @@ def deepseek_v3_pretrain_128gpu_vrnvl72_fp8mx_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 0, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # 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, + # NCCL EP dispatcher mode and one GPU per rank. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "CUDNNFE_CLUSTER_OVERLAP_MARGIN": 8, "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, @@ -190,7 +194,7 @@ def deepseek_v3_pretrain_128gpu_vrnvl72_fp8mx_config() -> ConfigContainer: def deepseek_v3_pretrain_128gpu_vrnvl72_nvfp4_config() -> ConfigContainer: - """DeepSeek V3 pretrain: 128× VRNVL72, NVFP4.""" + """DeepSeek V3 pretrain: 128× VRNVL72, NVFP4, NCCL EP.""" cfg = deepseek_v3_pretrain_config() cfg.mixed_precision = _perf_precision("nvfp4") _deepseek_v3_common(cfg) @@ -215,6 +219,10 @@ def deepseek_v3_pretrain_128gpu_vrnvl72_nvfp4_config() -> ConfigContainer: set_deepseek_v3_pipeline_model_parallel_layout(cfg.model) _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, @@ -226,11 +234,8 @@ def deepseek_v3_pretrain_128gpu_vrnvl72_nvfp4_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # 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, + # NCCL EP dispatcher mode and one GPU per rank. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, @@ -243,8 +248,12 @@ def deepseek_v3_pretrain_128gpu_vrnvl72_nvfp4_config() -> ConfigContainer: def deepseek_v3_pretrain_256gpu_vrnvl72_bf16_config() -> ConfigContainer: - """DeepSeek V3 pretrain: 256× VRNVL72, BF16 (alias of GB300).""" + """DeepSeek V3 pretrain: 256× VRNVL72, BF16, NCCL EP (alias of GB300).""" cfg = _build_deepseek_v3_gb300_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. + cfg.model.moe_use_grouped_tensor = True cfg.model.cuda_graph_scope = ["attn"] # Keep process settings next to the recipe so users can see the exact benchmark environment. @@ -258,11 +267,8 @@ def deepseek_v3_pretrain_256gpu_vrnvl72_bf16_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # 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, + # NCCL EP dispatcher mode and one GPU per rank. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, @@ -273,8 +279,12 @@ def deepseek_v3_pretrain_256gpu_vrnvl72_bf16_config() -> ConfigContainer: def deepseek_v3_pretrain_256gpu_vrnvl72_fp8cs_config() -> ConfigContainer: - """DeepSeek V3 pretrain: 256× VRNVL72, FP8-CS (alias of GB300).""" + """DeepSeek V3 pretrain: 256× VRNVL72, FP8-CS, NCCL EP (alias of GB300).""" cfg = deepseek_v3_pretrain_256gpu_gb300_fp8cs_config() + _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, @@ -286,11 +296,8 @@ def deepseek_v3_pretrain_256gpu_vrnvl72_fp8cs_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # HybridEP topology for the target system. - "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 32, - "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, - "NVLINK_DOMAIN_SIZE": 72, - "USE_MNNVL": 1, + # NCCL EP dispatcher mode and one GPU per rank. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, @@ -301,8 +308,12 @@ def deepseek_v3_pretrain_256gpu_vrnvl72_fp8cs_config() -> ConfigContainer: def deepseek_v3_pretrain_256gpu_vrnvl72_fp8mx_config() -> ConfigContainer: - """DeepSeek V3 pretrain: 256× VRNVL72, FP8-MX (alias of GB300).""" + """DeepSeek V3 pretrain: 256× VRNVL72, FP8-MX, NCCL EP (alias of GB300).""" cfg = _build_deepseek_v3_gb300_fp8mx() + _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, @@ -314,11 +325,8 @@ def deepseek_v3_pretrain_256gpu_vrnvl72_fp8mx_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 0, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # HybridEP topology for the target system. - "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 32, - "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, - "NVLINK_DOMAIN_SIZE": 72, - "USE_MNNVL": 1, + # NCCL EP dispatcher mode and one GPU per rank. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "CUDNNFE_CLUSTER_OVERLAP_MARGIN": 8, "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, @@ -359,10 +367,8 @@ def deepseek_v3_pretrain_64gpu_vrnvl72_fp8mx_proxy_config() -> ConfigContainer: "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True,graph_capture_record_stream_reuse:True", "TORCH_NCCL_AVOID_RECORD_STREAMS": 0, "NCCL_NVLS_ENABLE": 0, - "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 32, - "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, - "NVLINK_DOMAIN_SIZE": 72, - "USE_MNNVL": 1, + # NCCL EP dispatcher mode and one GPU per rank. + "NCCL_EP_HT_EM_PULL_PUSH": 1, "CUDNNFE_CLUSTER_OVERLAP_MARGIN": 8, "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_CUTEDSL_FUSED_GROUPED_MLP": 1, @@ -375,8 +381,12 @@ def deepseek_v3_pretrain_64gpu_vrnvl72_fp8mx_proxy_config() -> ConfigContainer: def deepseek_v3_pretrain_256gpu_vrnvl72_nvfp4_config() -> ConfigContainer: - """DeepSeek V3 pretrain: 256× VRNVL72, NVFP4 with full-iteration CUDA graph.""" + """DeepSeek V3 pretrain: 256× VRNVL72, NVFP4 with full-iteration CUDA graph, NCCL EP.""" cfg = _build_deepseek_v3_gb300_nvfp4() + _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 = { @@ -389,11 +399,8 @@ def deepseek_v3_pretrain_256gpu_vrnvl72_nvfp4_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # HybridEP topology for the target system. - "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 32, - "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, - "NVLINK_DOMAIN_SIZE": 72, - "USE_MNNVL": 1, + # NCCL EP dispatcher mode and one GPU per rank. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, @@ -433,10 +440,8 @@ def deepseek_v3_pretrain_64gpu_vrnvl72_nvfp4_proxy_config() -> ConfigContainer: "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True,graph_capture_record_stream_reuse:True", "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, "NCCL_NVLS_ENABLE": 0, - "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 32, - "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, - "NVLINK_DOMAIN_SIZE": 72, - "USE_MNNVL": 1, + # NCCL EP dispatcher mode and one GPU per rank. + "NCCL_EP_HT_EM_PULL_PUSH": 1, "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, @@ -481,7 +486,7 @@ def deepseek_v3_pretrain_4gpu_vrnvl72_fp8mx_proxy_config() -> ConfigContainer: Matches ``deepseek_v3_pretrain_4gpu_vrnvl72_nvfp4_proxy_config`` in everything except the quantization recipe, so the two profiles can be compared directly. Per-layer shapes and per-rank expert work follow the 256-GPU parent: TP=1, PP=1, EP=4, MBS=1, GBS=128, one MTP - layer, full-iteration CUDA graphs, HybridEP and combined-1F1B MoE overlap. Attention and all + layer, full-iteration CUDA graphs, NCCL EP and combined-1F1B MoE overlap. Attention and all GEMMs run in MXFP8; the LM head stays in BF16. This proxy does not reproduce at-scale communication, convergence or memory use. """ @@ -500,11 +505,8 @@ def deepseek_v3_pretrain_4gpu_vrnvl72_fp8mx_proxy_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 0, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # HybridEP topology: the whole EP=4 group sits in one node. - "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 4, - "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, - "NVLINK_DOMAIN_SIZE": 72, - "USE_MNNVL": 1, + # NCCL EP dispatcher mode and one GPU per rank. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "CUDNNFE_CLUSTER_OVERLAP_MARGIN": 8, "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, @@ -525,7 +527,7 @@ def deepseek_v3_pretrain_4gpu_vrnvl72_nvfp4_proxy_config() -> ConfigContainer: Matches ``deepseek_v3_pretrain_4gpu_vrnvl72_fp8mx_proxy_config`` in everything except the quantization recipe, so the two profiles can be compared directly. Per-layer shapes and per-rank expert work follow the 256-GPU parent: TP=1, PP=1, EP=4, MBS=1, GBS=128, one MTP - layer, full-iteration CUDA graphs, HybridEP and combined-1F1B MoE overlap. GEMMs run in + layer, full-iteration CUDA graphs, NCCL EP and combined-1F1B MoE overlap. GEMMs run in NVFP4, attention in MXFP8 E4M3 as in the MXFP8 proxy, and the LM head in BF16. This proxy does not reproduce at-scale communication, convergence or memory use. """ @@ -544,11 +546,8 @@ def deepseek_v3_pretrain_4gpu_vrnvl72_nvfp4_proxy_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 0, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # HybridEP topology: the whole EP=4 group sits in one node. - "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 4, - "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, - "NVLINK_DOMAIN_SIZE": 72, - "USE_MNNVL": 1, + # NCCL EP dispatcher mode and one GPU per rank. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "CUDNNFE_CLUSTER_OVERLAP_MARGIN": 8, "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, diff --git a/src/megatron/bridge/perf_recipes/deepseek/vrnvl72/deepseek_v4.py b/src/megatron/bridge/perf_recipes/deepseek/vrnvl72/deepseek_v4.py index 1ed0d3bf3a..25a3ab0dbe 100644 --- a/src/megatron/bridge/perf_recipes/deepseek/vrnvl72/deepseek_v4.py +++ b/src/megatron/bridge/perf_recipes/deepseek/vrnvl72/deepseek_v4.py @@ -13,6 +13,7 @@ # limitations under the License. """VRNVL72 performance recipes for DeepSeek V4.""" +from megatron.bridge.perf_recipes._common import _enable_ncclep from megatron.bridge.perf_recipes.deepseek.gb300.deepseek_v4 import ( deepseek_v4_flash_pretrain_128gpu_gb300_fp8mx_config, ) @@ -21,8 +22,9 @@ def deepseek_v4_flash_pretrain_128gpu_vrnvl72_fp8mx_config() -> ConfigContainer: - """DeepSeek V4 Flash pretrain: 128× VRNVL72, MXFP8 (alias of GB300).""" + """DeepSeek V4 Flash pretrain: 128× VRNVL72, MXFP8, NCCL EP (alias of GB300).""" cfg = deepseek_v4_flash_pretrain_128gpu_gb300_fp8mx_config() + _enable_ncclep(cfg) # Rubin's grouped GLU kernel does not support DeepSeek V4's SwiGLU clamp parameters. This is a # benchmark-only override that preserves the lower-memory CutDSL fusion and full CUDA graph. @@ -38,10 +40,8 @@ def deepseek_v4_flash_pretrain_128gpu_vrnvl72_fp8mx_config() -> ConfigContainer: "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True,graph_capture_record_stream_reuse:True", "TORCH_NCCL_AVOID_RECORD_STREAMS": 0, "NCCL_NVLS_ENABLE": 0, - "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 32, - "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, - "NVLINK_DOMAIN_SIZE": 72, - "USE_MNNVL": 1, + # NCCL EP dispatcher mode and one GPU per rank. + "NCCL_EP_HT_EM_PULL_PUSH": 1, "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_CUTEDSL_FUSED_GROUPED_MLP": 1, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, diff --git a/src/megatron/bridge/perf_recipes/gpt_oss/vrnvl72/gpt_oss.py b/src/megatron/bridge/perf_recipes/gpt_oss/vrnvl72/gpt_oss.py index dbc5450e9c..40543f2353 100644 --- a/src/megatron/bridge/perf_recipes/gpt_oss/vrnvl72/gpt_oss.py +++ b/src/megatron/bridge/perf_recipes/gpt_oss/vrnvl72/gpt_oss.py @@ -13,6 +13,7 @@ # limitations under the License. """VRNVL72 performance recipes for GPT-OSS.""" +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.gpt_oss.common import ( ConfigContainer, @@ -31,8 +32,12 @@ def gpt_oss_20b_pretrain_8gpu_vrnvl72_nvfp4_config() -> ConfigContainer: - """GPT-OSS 20B pretrain: 8× VRNVL72, NVFP4 (alias of GB300).""" + """GPT-OSS 20B pretrain: 8× VRNVL72, NVFP4, NCCL EP (alias of GB300).""" cfg = gpt_oss_20b_pretrain_8gpu_gb300_nvfp4_config() + _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, @@ -43,11 +48,8 @@ def gpt_oss_20b_pretrain_8gpu_vrnvl72_nvfp4_config() -> ConfigContainer: # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 1, "NCCL_CTA_POLICY": 1, - # HybridEP topology for the target system. - "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 4, - "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, - "NVLINK_DOMAIN_SIZE": 72, - "USE_MNNVL": 1, + # NCCL EP dispatcher mode and one GPU per rank. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, @@ -109,7 +111,7 @@ def gpt_oss_20b_pretrain_8gpu_vrnvl72_fp8mx_config() -> ConfigContainer: def gpt_oss_20b_pretrain_64gpu_vrnvl72_nvfp4_config() -> ConfigContainer: - """GPT-OSS 20B pretrain: 64× VRNVL72, NVFP4.""" + """GPT-OSS 20B pretrain: 64× VRNVL72, NVFP4, NCCL EP.""" cfg = gpt_oss_20b_pretrain_config() cfg.mixed_precision = _gpt_oss_20b_nvfp4_precision() @@ -124,6 +126,10 @@ def gpt_oss_20b_pretrain_64gpu_vrnvl72_nvfp4_config() -> ConfigContainer: _benchmark_common(cfg) _apply_gpt_oss_20b_common_configs(cfg) _apply_gpt_oss_20b_transformer_engine_graph_configs(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 cfg.model.cuda_graph_warmup_steps = 1 cfg.optimizer.lr = 0.0006 @@ -141,11 +147,8 @@ def gpt_oss_20b_pretrain_64gpu_vrnvl72_nvfp4_config() -> ConfigContainer: # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 1, "NCCL_CTA_POLICY": 1, - # HybridEP topology for the target system. - "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 4, - "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, - "NVLINK_DOMAIN_SIZE": 72, - "USE_MNNVL": 1, + # NCCL EP dispatcher mode and one GPU per rank. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, @@ -180,8 +183,12 @@ def gpt_oss_120b_pretrain_64gpu_vrnvl72_bf16_config() -> ConfigContainer: def gpt_oss_120b_pretrain_64gpu_vrnvl72_fp8mx_config() -> ConfigContainer: - """GPT-OSS 120B pretrain: 64× VRNVL72, FP8-MX (alias of GB300).""" + """GPT-OSS 120B pretrain: 64× VRNVL72, FP8-MX, NCCL EP (alias of GB300).""" cfg = _build_gpt_oss_120b_gb300_fp8mx() + _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, @@ -193,11 +200,8 @@ def gpt_oss_120b_pretrain_64gpu_vrnvl72_fp8mx_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 0, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # HybridEP topology for the target system. - "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 16, - "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, - "NVLINK_DOMAIN_SIZE": 72, - "USE_MNNVL": 1, + # NCCL EP dispatcher mode and one GPU per rank. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "CUDNNFE_CLUSTER_OVERLAP_MARGIN": 8, "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, diff --git a/src/megatron/bridge/perf_recipes/kimi/vrnvl72/kimi_k2.py b/src/megatron/bridge/perf_recipes/kimi/vrnvl72/kimi_k2.py index 208168dda5..abf7aaa48f 100644 --- a/src/megatron/bridge/perf_recipes/kimi/vrnvl72/kimi_k2.py +++ b/src/megatron/bridge/perf_recipes/kimi/vrnvl72/kimi_k2.py @@ -13,6 +13,7 @@ # limitations under the License. """VRNVL72 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, @@ -24,8 +25,12 @@ def kimi_k2_pretrain_256gpu_vrnvl72_bf16_config() -> ConfigContainer: - """Kimi K2 pretrain: 256× VRNVL72, BF16 (alias of GB300).""" + """Kimi K2 pretrain: 256× VRNVL72, BF16, NCCL EP (alias of GB300).""" cfg = kimi_k2_pretrain_256gpu_gb300_bf16_config() + _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, @@ -37,11 +42,8 @@ def kimi_k2_pretrain_256gpu_vrnvl72_bf16_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # 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, + # NCCL EP dispatcher mode and one GPU per rank. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, @@ -50,8 +52,12 @@ def kimi_k2_pretrain_256gpu_vrnvl72_bf16_config() -> ConfigContainer: def kimi_k2_pretrain_256gpu_vrnvl72_fp8mx_config() -> ConfigContainer: - """Kimi K2 pretrain: 256× VRNVL72, MXFP8 (alias of GB300).""" + """Kimi K2 pretrain: 256× VRNVL72, MXFP8, NCCL EP (alias of GB300).""" cfg = _build_kimi_k2_gb300_fp8mx() + _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, @@ -63,11 +69,8 @@ def kimi_k2_pretrain_256gpu_vrnvl72_fp8mx_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # 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, + # NCCL EP dispatcher mode and one GPU per rank. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, diff --git a/src/megatron/bridge/perf_recipes/nemotronh/vrnvl72/nemotronh.py b/src/megatron/bridge/perf_recipes/nemotronh/vrnvl72/nemotronh.py index 473373a02c..293f6158d2 100644 --- a/src/megatron/bridge/perf_recipes/nemotronh/vrnvl72/nemotronh.py +++ b/src/megatron/bridge/perf_recipes/nemotronh/vrnvl72/nemotronh.py @@ -13,6 +13,7 @@ # limitations under the License. """VRNVL72 performance recipes for Nemotron 3 and 3.5.""" +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.nemotronh.common import ( ConfigContainer, @@ -33,8 +34,12 @@ def nemotron_3_nano_pretrain_8gpu_vrnvl72_bf16_config() -> ConfigContainer: - """Nemotron 3 Nano pretrain: 8× VRNVL72, BF16 (alias of GB300).""" + """Nemotron 3 Nano pretrain: 8× VRNVL72, BF16, NCCL EP (alias of GB300).""" cfg = _build_nemotron_3_nano_gb300_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. + 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, @@ -46,11 +51,8 @@ def nemotron_3_nano_pretrain_8gpu_vrnvl72_bf16_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # 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, + # NCCL EP dispatcher mode and one GPU per rank. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, @@ -62,8 +64,12 @@ def nemotron_3_nano_pretrain_8gpu_vrnvl72_bf16_config() -> ConfigContainer: def nemotron_3_nano_pretrain_8gpu_vrnvl72_fp8mx_config() -> ConfigContainer: - """Nemotron 3 Nano pretrain: 8× VRNVL72, FP8-MX (alias of GB300).""" + """Nemotron 3 Nano pretrain: 8× VRNVL72, FP8-MX, NCCL EP (alias of GB300).""" cfg = _build_nemotron_3_nano_gb300_mxfp8() + _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.model.use_transformer_engine_op_fuser = True cfg.model.moe_mlp_glu_interleave_size = 32 cfg.mixed_precision.fp8_dot_product_attention = True @@ -78,11 +84,8 @@ def nemotron_3_nano_pretrain_8gpu_vrnvl72_fp8mx_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # 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, + # NCCL EP dispatcher mode and one GPU per rank. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "CUDNNFE_CLUSTER_OVERLAP_MARGIN": 8, "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, @@ -96,8 +99,12 @@ def nemotron_3_nano_pretrain_8gpu_vrnvl72_fp8mx_config() -> ConfigContainer: def nemotron_3_nano_pretrain_8gpu_vrnvl72_nvfp4_config() -> ConfigContainer: - """Nemotron 3 Nano pretrain: 8× VRNVL72, NVFP4 (alias of GB300).""" + """Nemotron 3 Nano pretrain: 8× VRNVL72, NVFP4, NCCL EP (alias of GB300).""" cfg = nemotron_3_nano_pretrain_8gpu_gb300_nvfp4_config() + _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, @@ -109,11 +116,8 @@ def nemotron_3_nano_pretrain_8gpu_vrnvl72_nvfp4_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # 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, + # NCCL EP dispatcher mode and one GPU per rank. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, @@ -127,8 +131,12 @@ def nemotron_3_nano_pretrain_8gpu_vrnvl72_nvfp4_config() -> ConfigContainer: def nemotron_3_super_pretrain_64gpu_vrnvl72_bf16_config() -> ConfigContainer: - """Nemotron 3 Super pretrain: 64× VRNVL72, BF16 (alias of GB300).""" + """Nemotron 3 Super pretrain: 64× VRNVL72, BF16, NCCL EP (alias of GB300).""" cfg = _build_nemotron_3_super_gb300_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. + 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, @@ -140,11 +148,8 @@ def nemotron_3_super_pretrain_64gpu_vrnvl72_bf16_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # 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, + # NCCL EP dispatcher mode and one GPU per rank. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, @@ -153,8 +158,12 @@ def nemotron_3_super_pretrain_64gpu_vrnvl72_bf16_config() -> ConfigContainer: def nemotron_3_super_pretrain_64gpu_vrnvl72_fp8mx_config() -> ConfigContainer: - """Nemotron 3 Super pretrain: 64× VRNVL72, MXFP8 with full-iteration CUDA graph.""" + """Nemotron 3 Super pretrain: 64× VRNVL72, MXFP8 with full-iteration CUDA graph, NCCL EP.""" cfg = _build_nemotron_3_super_gb300_mxfp8() + _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.mixed_precision.fp8_dot_product_attention = True cfg.model.expert_model_parallel_size = 8 @@ -169,11 +178,8 @@ def nemotron_3_super_pretrain_64gpu_vrnvl72_fp8mx_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 0, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # 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, + # NCCL EP dispatcher mode and one GPU per rank. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "CUDNNFE_CLUSTER_OVERLAP_MARGIN": 8, "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, @@ -184,8 +190,12 @@ def nemotron_3_super_pretrain_64gpu_vrnvl72_fp8mx_config() -> ConfigContainer: def nemotron_3_super_pretrain_64gpu_vrnvl72_nvfp4_config() -> ConfigContainer: - """Nemotron 3 Super pretrain: 64× VRNVL72, NVFP4 with full-iteration CUDA graph.""" + """Nemotron 3 Super pretrain: 64× VRNVL72, NVFP4 with full-iteration CUDA graph, NCCL EP.""" cfg = _build_nemotron_3_super_gb300_nvfp4() + _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.model.expert_model_parallel_size = 16 # Keep process settings next to the recipe so users can see the exact benchmark environment. @@ -199,11 +209,8 @@ def nemotron_3_super_pretrain_64gpu_vrnvl72_nvfp4_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 0, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # HybridEP topology for the target system. - "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 16, - "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, - "NVLINK_DOMAIN_SIZE": 72, - "USE_MNNVL": 1, + # NCCL EP dispatcher mode and one GPU per rank. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "CUDNNFE_CLUSTER_OVERLAP_MARGIN": 8, "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, @@ -294,8 +301,12 @@ def nemotron_3_ultra_pretrain_256gpu_vrnvl72_nvfp4_config(*, num_gpus: int = 256 def nemotron_3_5_lightning_pretrain_8gpu_vrnvl72_bf16_config() -> ConfigContainer: - """Nemotron 3.5 Lightning pretrain: 8× VRNVL72, BF16 (alias of GB300).""" + """Nemotron 3.5 Lightning pretrain: 8× VRNVL72, BF16, NCCL EP (alias of GB300).""" cfg = _build_nemotron_3_5_lightning_gb300_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. + cfg.model.moe_use_grouped_tensor = True # Offload MoE expert activations to host memory: the current Rubin base image ships a Transformer # Engine without TE PR #3273, whose fused cross entropy holds an extra fp32 gradient per loss head. # Revert this offload once the Rubin image picks up TE PR #3273: @@ -311,10 +322,8 @@ def nemotron_3_5_lightning_pretrain_8gpu_vrnvl72_bf16_config() -> ConfigContaine "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True,graph_capture_record_stream_reuse:True", "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, "NCCL_NVLS_ENABLE": 0, - "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 8, - "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, - "NVLINK_DOMAIN_SIZE": 72, - "USE_MNNVL": 1, + # NCCL EP dispatcher mode and one GPU per rank. + "NCCL_EP_HT_EM_PULL_PUSH": 1, "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_CPU_OFFLOAD_V1": 1, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, @@ -325,11 +334,12 @@ def nemotron_3_5_lightning_pretrain_8gpu_vrnvl72_bf16_config() -> ConfigContaine def nemotron_3_5_lightning_pretrain_8gpu_vrnvl72_fp8mx_config() -> ConfigContainer: - """Nemotron 3.5 Lightning pretrain: 8× VRNVL72, MXFP8 with full-iteration CUDA graph.""" + """Nemotron 3.5 Lightning pretrain: 8× VRNVL72, MXFP8 with full-iteration CUDA graph, NCCL EP.""" cfg = _build_nemotron_3_5_lightning_gb300_mxfp8() cfg.model.recompute_modules = [] cfg.model.use_transformer_engine_op_fuser = True cfg.mixed_precision.fp8_dot_product_attention = True + _enable_ncclep(cfg) _enable_nemotron_3_5_lightning_full_iteration(cfg) # Keep the VRNVL72 launch environment explicit instead of inheriting it from GB300. @@ -340,10 +350,8 @@ def nemotron_3_5_lightning_pretrain_8gpu_vrnvl72_fp8mx_config() -> ConfigContain "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True,graph_capture_record_stream_reuse:True", "TORCH_NCCL_AVOID_RECORD_STREAMS": 0, "NCCL_NVLS_ENABLE": 0, - "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 8, - "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, - "NVLINK_DOMAIN_SIZE": 72, - "USE_MNNVL": 1, + # NCCL EP dispatcher mode and one GPU per rank. + "NCCL_EP_HT_EM_PULL_PUSH": 1, "CUDNNFE_CLUSTER_OVERLAP_MARGIN": 8, "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_CUTEDSL_FUSED_GROUPED_MLP": 1, diff --git a/src/megatron/bridge/perf_recipes/qwen/vrnvl72/qwen3_moe.py b/src/megatron/bridge/perf_recipes/qwen/vrnvl72/qwen3_moe.py index ac548cb2fe..8aae311f52 100644 --- a/src/megatron/bridge/perf_recipes/qwen/vrnvl72/qwen3_moe.py +++ b/src/megatron/bridge/perf_recipes/qwen/vrnvl72/qwen3_moe.py @@ -13,6 +13,7 @@ # limitations under the License. """VRNVL72 performance recipes for Qwen3 MoE.""" +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.qwen.common import ConfigContainer from megatron.bridge.perf_recipes.qwen.gb300.qwen3_moe import ( @@ -26,8 +27,12 @@ def qwen3_235b_a22b_pretrain_256gpu_vrnvl72_bf16_config() -> ConfigContainer: - """Qwen3 235B A22B pretrain: 256× VRNVL72, BF16 (alias of GB300).""" + """Qwen3 235B A22B pretrain: 256× VRNVL72, BF16, NCCL EP (alias of GB300).""" cfg = _build_qwen3_235b_a22b_gb300_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. + cfg.model.moe_use_grouped_tensor = True cfg.model.cuda_graph_scope = ["moe_router", "moe_preprocess"] # Keep process settings next to the recipe so users can see the exact benchmark environment. @@ -41,11 +46,8 @@ def qwen3_235b_a22b_pretrain_256gpu_vrnvl72_bf16_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # HybridEP topology for the target system. - "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 32, - "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, - "NVLINK_DOMAIN_SIZE": 72, - "USE_MNNVL": 1, + # NCCL EP dispatcher mode and one GPU per rank. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, @@ -54,8 +56,12 @@ def qwen3_235b_a22b_pretrain_256gpu_vrnvl72_bf16_config() -> ConfigContainer: def qwen3_235b_a22b_pretrain_256gpu_vrnvl72_fp8mx_config() -> ConfigContainer: - """Qwen3 235B A22B pretrain: 256× VRNVL72, FP8-MX (alias of GB300).""" + """Qwen3 235B A22B pretrain: 256× VRNVL72, FP8-MX, NCCL EP (alias of GB300).""" cfg = _build_qwen3_235b_a22b_gb300_fp8mx() + _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, @@ -67,11 +73,8 @@ def qwen3_235b_a22b_pretrain_256gpu_vrnvl72_fp8mx_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 0, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # HybridEP topology for the target system. - "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 32, - "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, - "NVLINK_DOMAIN_SIZE": 72, - "USE_MNNVL": 1, + # NCCL EP dispatcher mode and one GPU per rank. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "CUDNNFE_CLUSTER_OVERLAP_MARGIN": 8, "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, @@ -82,8 +85,12 @@ def qwen3_235b_a22b_pretrain_256gpu_vrnvl72_fp8mx_config() -> ConfigContainer: def qwen3_235b_a22b_pretrain_256gpu_vrnvl72_nvfp4_config() -> ConfigContainer: - """Qwen3 235B A22B pretrain: 256× VRNVL72, NVFP4, full-iteration CG and A2A overlap.""" + """Qwen3 235B A22B pretrain: 256× VRNVL72, NVFP4, NCCL EP, full-iteration CG and A2A overlap.""" cfg = _build_qwen3_235b_a22b_gb300_nvfp4() + _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 = { @@ -96,11 +103,8 @@ def qwen3_235b_a22b_pretrain_256gpu_vrnvl72_nvfp4_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 0, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # HybridEP topology for the target system. - "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 32, - "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, - "NVLINK_DOMAIN_SIZE": 72, - "USE_MNNVL": 1, + # NCCL EP dispatcher mode and one GPU per rank. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "CUDNNFE_CLUSTER_OVERLAP_MARGIN": 8, "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, @@ -113,8 +117,12 @@ def qwen3_235b_a22b_pretrain_256gpu_vrnvl72_nvfp4_config() -> ConfigContainer: def qwen3_30b_a3b_pretrain_8gpu_vrnvl72_bf16_config() -> ConfigContainer: - """Qwen3 30B-A3B pretrain: 8× VRNVL72, BF16 (alias of GB300).""" + """Qwen3 30B-A3B pretrain: 8× VRNVL72, BF16, NCCL EP (alias of GB300).""" cfg = _build_qwen3_30b_a3b_gb300_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. + 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, @@ -126,11 +134,8 @@ def qwen3_30b_a3b_pretrain_8gpu_vrnvl72_bf16_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # 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, + # NCCL EP dispatcher mode and one GPU per rank. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, @@ -139,8 +144,12 @@ def qwen3_30b_a3b_pretrain_8gpu_vrnvl72_bf16_config() -> ConfigContainer: def qwen3_30b_a3b_pretrain_8gpu_vrnvl72_fp8mx_config() -> ConfigContainer: - """Qwen3 30B-A3B pretrain: 8× VRNVL72, FP8-MX (alias of GB300).""" + """Qwen3 30B-A3B pretrain: 8× VRNVL72, FP8-MX, NCCL EP (alias of GB300).""" cfg = _build_qwen3_30b_a3b_gb300_fp8mx() + _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.train.micro_batch_size = 4 @@ -159,11 +168,8 @@ def qwen3_30b_a3b_pretrain_8gpu_vrnvl72_fp8mx_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 0, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # 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, + # NCCL EP dispatcher mode and one GPU per rank. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "CUDNNFE_CLUSTER_OVERLAP_MARGIN": 8, "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, @@ -174,12 +180,16 @@ def qwen3_30b_a3b_pretrain_8gpu_vrnvl72_fp8mx_config() -> ConfigContainer: def qwen3_30b_a3b_pretrain_8gpu_vrnvl72_nvfp4_config() -> ConfigContainer: - """Qwen3 30B-A3B pretrain: 8× VRNVL72, NVFP4 (same layout as GB300). + """Qwen3 30B-A3B pretrain: 8× VRNVL72, NVFP4, NCCL EP (same layout as GB300). NVFP4's fp4_param_gather path is incompatible with TP comm overlap, so it is disabled here. """ cfg = qwen3_30b_a3b_pretrain_8gpu_gb300_nvfp4_config() + _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, @@ -187,10 +197,8 @@ def qwen3_30b_a3b_pretrain_8gpu_vrnvl72_nvfp4_config() -> ConfigContainer: "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, "NCCL_NVLS_ENABLE": 0, - "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 8, - "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, - "NVLINK_DOMAIN_SIZE": 72, - "USE_MNNVL": 1, + # NCCL EP dispatcher mode and one GPU per rank. + "NCCL_EP_HT_EM_PULL_PUSH": 1, "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, "NVTE_USE_FAST_MATH": 1, diff --git a/src/megatron/bridge/perf_recipes/qwen_vl/vrnvl72/qwen35_vl.py b/src/megatron/bridge/perf_recipes/qwen_vl/vrnvl72/qwen35_vl.py index c796054d28..d7dbe76284 100644 --- a/src/megatron/bridge/perf_recipes/qwen_vl/vrnvl72/qwen35_vl.py +++ b/src/megatron/bridge/perf_recipes/qwen_vl/vrnvl72/qwen35_vl.py @@ -13,6 +13,7 @@ # limitations under the License. """VRNVL72 performance recipes for Qwen3.5-VL.""" +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.qwen_vl.common import ConfigContainer, _select_gdn_kernel_backend from megatron.bridge.perf_recipes.qwen_vl.gb300.qwen35_vl import ( @@ -29,8 +30,12 @@ def qwen35_vl_35b_a3b_pretrain_8gpu_vrnvl72_bf16_config() -> ConfigContainer: - """Qwen3.5-VL 35B-A3B pretrain: 8× VRNVL72, BF16 (alias of GB300).""" + """Qwen3.5-VL 35B-A3B pretrain: 8× VRNVL72, BF16, NCCL EP (alias of GB300).""" cfg = _build_qwen35_vl_35b_a3b_gb300_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. + cfg.model.moe_use_grouped_tensor = True # The GB300 base selects the cuDNN GDN kernel. On VRNVL72 it was measured only with MXFP8 (against FLA # at 35B), so this precision keeps FLA until it is measured too. _select_gdn_kernel_backend(cfg, "fla") @@ -45,11 +50,8 @@ def qwen35_vl_35b_a3b_pretrain_8gpu_vrnvl72_bf16_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # 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, + # NCCL EP dispatcher mode and one GPU per rank. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, @@ -58,8 +60,12 @@ def qwen35_vl_35b_a3b_pretrain_8gpu_vrnvl72_bf16_config() -> ConfigContainer: def qwen35_vl_35b_a3b_pretrain_8gpu_vrnvl72_fp8cs_config() -> ConfigContainer: - """Qwen3.5-VL 35B-A3B pretrain: 8× VRNVL72, FP8-CS (alias of GB300).""" + """Qwen3.5-VL 35B-A3B pretrain: 8× VRNVL72, FP8-CS, NCCL EP (alias of GB300).""" cfg = qwen35_vl_35b_a3b_pretrain_8gpu_gb300_fp8cs_config() + _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 # The GB300 base selects the cuDNN GDN kernel. On VRNVL72 it was measured only with MXFP8 (against FLA # at 35B), so this precision keeps FLA until it is measured too. _select_gdn_kernel_backend(cfg, "fla") @@ -74,11 +80,8 @@ def qwen35_vl_35b_a3b_pretrain_8gpu_vrnvl72_fp8cs_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # 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, + # NCCL EP dispatcher mode and one GPU per rank. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, @@ -87,16 +90,19 @@ def qwen35_vl_35b_a3b_pretrain_8gpu_vrnvl72_fp8cs_config() -> ConfigContainer: def qwen35_vl_35b_a3b_pretrain_8gpu_vrnvl72_fp8mx_config() -> ConfigContainer: - """Qwen3.5-VL 35B-A3B pretrain: 8× VRNVL72, MXFP8, EP=4, HybridEP. + """Qwen3.5-VL 35B-A3B pretrain: 8× VRNVL72, MXFP8, EP=4, NCCL EP. Starts from the GB300 MXFP8 recipe (CuTe DSL fused grouped MLP, cuDNN GDN, per-layer CUDA graphs, forced load balancing) and applies EP 4, MBS 4 and the kernel settings of the measured VRNVL72 runs, at GBS 1024. """ cfg = qwen35_vl_35b_a3b_pretrain_8gpu_gb300_fp8mx_config() + _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 - # EP 4 and MBS 4 as in the measured VRNVL72 runs: each expert-parallel group is one 4-GPU node - # (NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN below must equal EP). GBS 1024 gives 32 microbatches - # per step on 8 GPUs. + # EP 4 and MBS 4 as in the measured VRNVL72 runs: each expert-parallel group is one 4-GPU node. + # GBS 1024 gives 32 microbatches per step on 8 GPUs. cfg.model.expert_model_parallel_size = 4 cfg.train.micro_batch_size = 4 cfg.train.global_batch_size = 1024 @@ -113,12 +119,8 @@ def qwen35_vl_35b_a3b_pretrain_8gpu_vrnvl72_fp8mx_config() -> ConfigContainer: cfg.model.use_grouped_gemm_for_shared_expert = True cfg.model.moe_shared_expert_glu_interleave_size = 32 - # HybridEP: permute outside the dispatcher (already the default; stated explicitly because fusing - # the permute into HybridEP measured slower on VRNVL72), 16 SMs for preprocessing as in the measured - # runs, and overlap the dispatch backward with the expert weight-gradient GEMMs. Megatron-Core does - # not allow this overlap together with EP A2A overlap or delay_wgrad_compute. - cfg.model.moe_permute_fusion_into_hybridep = False - cfg.model.moe_hybridep_num_sms_preprocessing = 16 + # Overlap the dispatch backward with the expert weight-gradient GEMMs, as in the measured runs. + # Megatron-Core does not allow this overlap together with EP A2A overlap or delay_wgrad_compute. cfg.model.overlap_dispatch_backward_with_experts_wgrad = True # Vision-encoder CUDA graph. The cap is the number of vision patches per microbatch: a longer @@ -146,11 +148,8 @@ def qwen35_vl_35b_a3b_pretrain_8gpu_vrnvl72_fp8mx_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 0, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # HybridEP topology for the target system: one rank per GPU of the EP 4 group. - "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 4, - "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, - "NVLINK_DOMAIN_SIZE": 72, - "USE_MNNVL": 1, + # NCCL EP dispatcher mode and one GPU per rank. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 0, "NVTE_FWD_LAYERNORM_SM_MARGIN": 0, @@ -165,8 +164,12 @@ def qwen35_vl_35b_a3b_pretrain_8gpu_vrnvl72_fp8mx_config() -> ConfigContainer: def qwen35_vl_122b_a10b_pretrain_32gpu_vrnvl72_bf16_config() -> ConfigContainer: - """Qwen3.5-VL 122B-A10B pretrain: 32× VRNVL72, BF16 (alias of GB300).""" + """Qwen3.5-VL 122B-A10B pretrain: 32× VRNVL72, BF16, NCCL EP (alias of GB300).""" cfg = qwen35_vl_122b_a10b_pretrain_32gpu_gb300_bf16_config() + _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, @@ -178,11 +181,8 @@ def qwen35_vl_122b_a10b_pretrain_32gpu_vrnvl72_bf16_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # HybridEP topology for the target system. - "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 32, - "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, - "NVLINK_DOMAIN_SIZE": 72, - "USE_MNNVL": 1, + # NCCL EP dispatcher mode and one GPU per rank. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, @@ -191,8 +191,12 @@ def qwen35_vl_122b_a10b_pretrain_32gpu_vrnvl72_bf16_config() -> ConfigContainer: def qwen35_vl_122b_a10b_pretrain_32gpu_vrnvl72_fp8cs_config() -> ConfigContainer: - """Qwen3.5-VL 122B-A10B pretrain: 32× VRNVL72, FP8-CS (alias of GB300).""" + """Qwen3.5-VL 122B-A10B pretrain: 32× VRNVL72, FP8-CS, NCCL EP (alias of GB300).""" cfg = qwen35_vl_122b_a10b_pretrain_32gpu_gb300_fp8cs_config() + _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, @@ -204,11 +208,8 @@ def qwen35_vl_122b_a10b_pretrain_32gpu_vrnvl72_fp8cs_config() -> ConfigContainer "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # HybridEP topology for the target system. - "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 32, - "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, - "NVLINK_DOMAIN_SIZE": 72, - "USE_MNNVL": 1, + # NCCL EP dispatcher mode and one GPU per rank. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, @@ -217,8 +218,12 @@ def qwen35_vl_122b_a10b_pretrain_32gpu_vrnvl72_fp8cs_config() -> ConfigContainer def qwen35_vl_122b_a10b_pretrain_32gpu_vrnvl72_fp8mx_config() -> ConfigContainer: - """Qwen3.5-VL 122B-A10B pretrain: 32× VRNVL72, FP8-MX (alias of GB300).""" + """Qwen3.5-VL 122B-A10B pretrain: 32× VRNVL72, FP8-MX, NCCL EP (alias of GB300).""" cfg = qwen35_vl_122b_a10b_pretrain_32gpu_gb300_fp8mx_config() + _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, @@ -230,11 +235,8 @@ def qwen35_vl_122b_a10b_pretrain_32gpu_vrnvl72_fp8mx_config() -> ConfigContainer "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # HybridEP topology for the target system. - "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 32, - "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, - "NVLINK_DOMAIN_SIZE": 72, - "USE_MNNVL": 1, + # NCCL EP dispatcher mode and one GPU per rank. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, @@ -243,8 +245,12 @@ def qwen35_vl_122b_a10b_pretrain_32gpu_vrnvl72_fp8mx_config() -> ConfigContainer def qwen35_vl_397b_a17b_pretrain_64gpu_vrnvl72_bf16_config() -> ConfigContainer: - """Qwen3.5-VL 397B-A17B pretrain: 64× VRNVL72, BF16 (alias of GB300).""" + """Qwen3.5-VL 397B-A17B pretrain: 64× VRNVL72, BF16, NCCL EP (alias of GB300).""" cfg = qwen35_vl_397b_a17b_pretrain_64gpu_gb300_bf16_config() + _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 # The GB300 base selects the cuDNN GDN kernel. On VRNVL72 it was measured only with MXFP8 (against FLA # at 35B), so this precision keeps FLA until it is measured too. _select_gdn_kernel_backend(cfg, "fla") @@ -259,11 +265,8 @@ def qwen35_vl_397b_a17b_pretrain_64gpu_vrnvl72_bf16_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # HybridEP topology for the target system. - "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 32, - "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, - "NVLINK_DOMAIN_SIZE": 72, - "USE_MNNVL": 1, + # NCCL EP dispatcher mode and one GPU per rank. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, @@ -272,8 +275,12 @@ def qwen35_vl_397b_a17b_pretrain_64gpu_vrnvl72_bf16_config() -> ConfigContainer: def qwen35_vl_397b_a17b_pretrain_64gpu_vrnvl72_fp8cs_config() -> ConfigContainer: - """Qwen3.5-VL 397B-A17B pretrain: 64× VRNVL72, FP8-CS (alias of GB300).""" + """Qwen3.5-VL 397B-A17B pretrain: 64× VRNVL72, FP8-CS, NCCL EP (alias of GB300).""" cfg = qwen35_vl_397b_a17b_pretrain_64gpu_gb300_fp8cs_config() + _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 # The GB300 base selects the cuDNN GDN kernel. On VRNVL72 it was measured only with MXFP8 (against FLA # at 35B), so this precision keeps FLA until it is measured too. _select_gdn_kernel_backend(cfg, "fla") @@ -288,11 +295,8 @@ def qwen35_vl_397b_a17b_pretrain_64gpu_vrnvl72_fp8cs_config() -> ConfigContainer "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # HybridEP topology for the target system. - "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 32, - "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, - "NVLINK_DOMAIN_SIZE": 72, - "USE_MNNVL": 1, + # NCCL EP dispatcher mode and one GPU per rank. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, @@ -301,12 +305,16 @@ def qwen35_vl_397b_a17b_pretrain_64gpu_vrnvl72_fp8cs_config() -> ConfigContainer def qwen35_vl_397b_a17b_pretrain_64gpu_vrnvl72_fp8mx_config() -> ConfigContainer: - """Qwen3.5-VL 397B-A17B pretrain: 64× VRNVL72, MXFP8, EP=32, HybridEP. + """Qwen3.5-VL 397B-A17B pretrain: 64× VRNVL72, MXFP8, EP=32, NCCL EP. Starts from the GB300 MXFP8 recipe (cuDNN GDN, per-layer CUDA graphs, forced load balancing) and adds the CuTe DSL fused grouped MLP and the settings of the measured VRNVL72 run. """ cfg = qwen35_vl_397b_a17b_pretrain_64gpu_gb300_fp8mx_config() + _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 # MBS 1, as in the GB300 base; the vision cap below is set for it. cfg.train.micro_batch_size = 1 @@ -329,10 +337,6 @@ def qwen35_vl_397b_a17b_pretrain_64gpu_vrnvl72_fp8mx_config() -> ConfigContainer cfg.model.recompute_granularity = "selective" cfg.model.recompute_modules = ["gdn_norm_out"] - # HybridEP: permute outside the dispatcher. This is already the default; it is stated explicitly - # because fusing the permute into HybridEP measured slower on VRNVL72. - cfg.model.moe_permute_fusion_into_hybridep = False - # Vision-encoder CUDA graph with the cap for MBS 1: each 256x256 mock image reaches the encoder as # 784 patches (see the 35B recipe above). Set the cap to MBS x images per sample x patches per image # whenever one of them changes. As in the 35B recipe, the graphed encoder runs without the 2D rotary @@ -352,14 +356,8 @@ def qwen35_vl_397b_a17b_pretrain_64gpu_vrnvl72_fp8mx_config() -> ConfigContainer "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # HybridEP topology for the target system. - "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 32, - "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, - "NVLINK_DOMAIN_SIZE": 72, - "USE_MNNVL": 1, - # HybridEP dispatch and preprocessing chunks of the measured run (HybridEP's default is 64). - "NUM_OF_TOKENS_PER_CHUNK_DISPATCH_API": 128, - "NUM_OF_TOKENS_PER_CHUNK_PREPROCESSING_API": 128, + # NCCL EP dispatcher mode and one GPU per rank. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 16, "NVTE_FWD_LAYERNORM_SM_MARGIN": 16, diff --git a/tests/unit_tests/recipes/test_deepseek_v3_proxy_recipes.py b/tests/unit_tests/recipes/test_deepseek_v3_proxy_recipes.py index 4193f41dff..fbbfcff081 100644 --- a/tests/unit_tests/recipes/test_deepseek_v3_proxy_recipes.py +++ b/tests/unit_tests/recipes/test_deepseek_v3_proxy_recipes.py @@ -121,7 +121,8 @@ def test_proxy_preserves_parent_except_depth_and_layout(precision: str) -> None: assert proxy.comm_overlap.overlap_moe_expert_parallel_comm is True assert proxy.comm_overlap.delay_wgrad_compute is True assert proxy.env_vars["NVTE_CUTEDSL_FUSED_GROUPED_MLP"] == 1 - assert model.moe_flex_dispatcher_backend == "hybridep" + assert model.moe_flex_dispatcher_backend == "ncclep" + assert proxy.env_vars["NCCL_EP_HT_EM_PULL_PUSH"] == 1 assert model.mtp_num_layers == 1 assert model.recompute_modules == [] diff --git a/tests/unit_tests/recipes/test_deepseek_v4_flash_perf_recipes.py b/tests/unit_tests/recipes/test_deepseek_v4_flash_perf_recipes.py index cd8373ef4b..2353950a51 100644 --- a/tests/unit_tests/recipes/test_deepseek_v4_flash_perf_recipes.py +++ b/tests/unit_tests/recipes/test_deepseek_v4_flash_perf_recipes.py @@ -149,13 +149,22 @@ def test_deepseek_v4_flash_128gpu_vrnvl72_fp8mx_disables_clamp_for_cutedsl() -> assert vrnvl72_cfg.env_vars["NVTE_CUTEDSL_FUSED_GROUPED_MLP"] == 1 assert vrnvl72_cfg.env_vars["NVTE_ALLOW_NONDETERMINISTIC_ALGO"] == 0 assert is_full_iteration_cuda_graph(vrnvl72_cfg.model) + # VRNVL72 dispatches through NCCL EP; the GB300 recipe stays on the shared HybridEP base. + assert vrnvl72_cfg.model.moe_flex_dispatcher_backend == "ncclep" + assert vrnvl72_cfg.env_vars["NCCL_EP_HT_EM_PULL_PUSH"] == 1 + assert vrnvl72_cfg.env_vars.keys().isdisjoint(HYBRID_EP_ENV_NAMES) + assert gb300_cfg.model.moe_flex_dispatcher_backend == "hybridep" gb300_cfg.model.activation_func_clamp_value = None gb300_cfg.model.use_transformer_engine_op_fuser = True gb300_cfg.model.moe_use_grouped_tensor = True + gb300_cfg.model.moe_flex_dispatcher_backend = "ncclep" assert vars(vrnvl72_cfg.model) == vars(gb300_cfg.model) + # The environments differ only in the dispatcher settings. + expected_env = {name: value for name, value in gb300_cfg.env_vars.items() if name not in HYBRID_EP_ENV_NAMES} + assert vrnvl72_cfg.env_vars == {**expected_env, "NCCL_EP_HT_EM_PULL_PUSH": 1} for field in fields(vrnvl72_cfg): - if field.name != "model": + if field.name not in ("model", "env_vars"): assert getattr(vrnvl72_cfg, field.name) == getattr(gb300_cfg, field.name), field.name diff --git a/tests/unit_tests/recipes/test_gb200_ncclep_perf_recipes.py b/tests/unit_tests/recipes/test_gb200_ncclep_perf_recipes.py index 74c2468705..1b0dd3806d 100644 --- a/tests/unit_tests/recipes/test_gb200_ncclep_perf_recipes.py +++ b/tests/unit_tests/recipes/test_gb200_ncclep_perf_recipes.py @@ -38,6 +38,7 @@ ("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_ultra_pretrain_256gpu_gb200_nvfp4_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_235b_a22b_pretrain_256gpu_gb200_bf16_config"), @@ -55,7 +56,6 @@ ("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"), diff --git a/tests/unit_tests/recipes/test_gb300_ncclep_perf_recipes.py b/tests/unit_tests/recipes/test_gb300_ncclep_perf_recipes.py index 1e0c3a69ac..0827a72363 100644 --- a/tests/unit_tests/recipes/test_gb300_ncclep_perf_recipes.py +++ b/tests/unit_tests/recipes/test_gb300_ncclep_perf_recipes.py @@ -11,9 +11,8 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Every GB300 MoE perf recipe switched to NCCL EP carries the full NCCL EP stack, and the VRNVL72 -aliases built on the shared HybridEP bases stay on HybridEP. VRNVL72 aliases built on an NCCL EP base -carry the same NCCL EP stack.""" +"""Every GB300 and VRNVL72 MoE perf recipe switched to NCCL EP carries the full NCCL EP stack, and the +VRNVL72 MoE recipes left off NCCL EP do not pick it up.""" import importlib @@ -57,34 +56,59 @@ ("qwen_vl.gb300.qwen35_vl", "qwen35_vl_35b_a3b_pretrain_8gpu_gb300_bf16_config"), ) -# VRNVL72 aliases that share the GB300 NCCL EP builder instead of a HybridEP base. +# VRNVL72 MoE recipes on NCCL EP. The Nemotron 3 Ultra aliases share the GB300 NCCL EP builders; the rest +# switch their GB300 HybridEP base, a GB300 HybridEP recipe, or their own HybridEP setup to NCCL EP. _VRNVL72_NCCLEP_RECIPES = ( - ("nemotronh.vrnvl72.nemotronh", "nemotron_3_ultra_pretrain_256gpu_vrnvl72_fp8mx_config"), - ("nemotronh.vrnvl72.nemotronh", "nemotron_3_ultra_pretrain_256gpu_vrnvl72_nvfp4_config"), -) - -# VRNVL72 aliases of the same recipes: they reuse the shared HybridEP builders and must not pick up NCCL EP. -_VRNVL72_ALIASES = ( + ("deepseek.vrnvl72.deepseek_v3", "deepseek_v3_pretrain_128gpu_vrnvl72_bf16_config"), + ("deepseek.vrnvl72.deepseek_v3", "deepseek_v3_pretrain_128gpu_vrnvl72_fp8cs_config"), + ("deepseek.vrnvl72.deepseek_v3", "deepseek_v3_pretrain_128gpu_vrnvl72_fp8mx_config"), + ("deepseek.vrnvl72.deepseek_v3", "deepseek_v3_pretrain_128gpu_vrnvl72_nvfp4_config"), ("deepseek.vrnvl72.deepseek_v3", "deepseek_v3_pretrain_256gpu_vrnvl72_bf16_config"), + ("deepseek.vrnvl72.deepseek_v3", "deepseek_v3_pretrain_256gpu_vrnvl72_fp8cs_config"), ("deepseek.vrnvl72.deepseek_v3", "deepseek_v3_pretrain_256gpu_vrnvl72_fp8mx_config"), ("deepseek.vrnvl72.deepseek_v3", "deepseek_v3_pretrain_256gpu_vrnvl72_nvfp4_config"), + ("deepseek.vrnvl72.deepseek_v3", "deepseek_v3_pretrain_64gpu_vrnvl72_fp8mx_proxy_config"), + ("deepseek.vrnvl72.deepseek_v3", "deepseek_v3_pretrain_64gpu_vrnvl72_nvfp4_proxy_config"), + ("deepseek.vrnvl72.deepseek_v3", "deepseek_v3_pretrain_4gpu_vrnvl72_fp8mx_proxy_config"), + ("deepseek.vrnvl72.deepseek_v3", "deepseek_v3_pretrain_4gpu_vrnvl72_nvfp4_proxy_config"), + ("deepseek.vrnvl72.deepseek_v4", "deepseek_v4_flash_pretrain_128gpu_vrnvl72_fp8mx_config"), + ("gpt_oss.vrnvl72.gpt_oss", "gpt_oss_20b_pretrain_8gpu_vrnvl72_nvfp4_config"), + ("gpt_oss.vrnvl72.gpt_oss", "gpt_oss_20b_pretrain_64gpu_vrnvl72_nvfp4_config"), ("gpt_oss.vrnvl72.gpt_oss", "gpt_oss_120b_pretrain_64gpu_vrnvl72_fp8mx_config"), + ("kimi.vrnvl72.kimi_k2", "kimi_k2_pretrain_256gpu_vrnvl72_bf16_config"), ("kimi.vrnvl72.kimi_k2", "kimi_k2_pretrain_256gpu_vrnvl72_fp8mx_config"), ("nemotronh.vrnvl72.nemotronh", "nemotron_3_nano_pretrain_8gpu_vrnvl72_bf16_config"), + ("nemotronh.vrnvl72.nemotronh", "nemotron_3_nano_pretrain_8gpu_vrnvl72_fp8mx_config"), + ("nemotronh.vrnvl72.nemotronh", "nemotron_3_nano_pretrain_8gpu_vrnvl72_nvfp4_config"), ("nemotronh.vrnvl72.nemotronh", "nemotron_3_super_pretrain_64gpu_vrnvl72_bf16_config"), + ("nemotronh.vrnvl72.nemotronh", "nemotron_3_super_pretrain_64gpu_vrnvl72_fp8mx_config"), ("nemotronh.vrnvl72.nemotronh", "nemotron_3_super_pretrain_64gpu_vrnvl72_nvfp4_config"), + ("nemotronh.vrnvl72.nemotronh", "nemotron_3_ultra_pretrain_256gpu_vrnvl72_fp8mx_config"), + ("nemotronh.vrnvl72.nemotronh", "nemotron_3_ultra_pretrain_256gpu_vrnvl72_nvfp4_config"), ("nemotronh.vrnvl72.nemotronh", "nemotron_3_5_lightning_pretrain_8gpu_vrnvl72_bf16_config"), ("nemotronh.vrnvl72.nemotronh", "nemotron_3_5_lightning_pretrain_8gpu_vrnvl72_fp8mx_config"), - ("deepseek.vrnvl72.deepseek_v3", "deepseek_v3_pretrain_64gpu_vrnvl72_fp8mx_proxy_config"), - ("deepseek.vrnvl72.deepseek_v3", "deepseek_v3_pretrain_64gpu_vrnvl72_nvfp4_proxy_config"), - ("deepseek.vrnvl72.deepseek_v3", "deepseek_v3_pretrain_4gpu_vrnvl72_fp8mx_proxy_config"), - ("deepseek.vrnvl72.deepseek_v3", "deepseek_v3_pretrain_4gpu_vrnvl72_nvfp4_proxy_config"), ("qwen.vrnvl72.qwen3_moe", "qwen3_235b_a22b_pretrain_256gpu_vrnvl72_bf16_config"), ("qwen.vrnvl72.qwen3_moe", "qwen3_235b_a22b_pretrain_256gpu_vrnvl72_fp8mx_config"), ("qwen.vrnvl72.qwen3_moe", "qwen3_235b_a22b_pretrain_256gpu_vrnvl72_nvfp4_config"), ("qwen.vrnvl72.qwen3_moe", "qwen3_30b_a3b_pretrain_8gpu_vrnvl72_bf16_config"), ("qwen.vrnvl72.qwen3_moe", "qwen3_30b_a3b_pretrain_8gpu_vrnvl72_fp8mx_config"), + ("qwen.vrnvl72.qwen3_moe", "qwen3_30b_a3b_pretrain_8gpu_vrnvl72_nvfp4_config"), ("qwen_vl.vrnvl72.qwen35_vl", "qwen35_vl_35b_a3b_pretrain_8gpu_vrnvl72_bf16_config"), + ("qwen_vl.vrnvl72.qwen35_vl", "qwen35_vl_35b_a3b_pretrain_8gpu_vrnvl72_fp8cs_config"), + ("qwen_vl.vrnvl72.qwen35_vl", "qwen35_vl_35b_a3b_pretrain_8gpu_vrnvl72_fp8mx_config"), + ("qwen_vl.vrnvl72.qwen35_vl", "qwen35_vl_122b_a10b_pretrain_32gpu_vrnvl72_bf16_config"), + ("qwen_vl.vrnvl72.qwen35_vl", "qwen35_vl_122b_a10b_pretrain_32gpu_vrnvl72_fp8cs_config"), + ("qwen_vl.vrnvl72.qwen35_vl", "qwen35_vl_122b_a10b_pretrain_32gpu_vrnvl72_fp8mx_config"), + ("qwen_vl.vrnvl72.qwen35_vl", "qwen35_vl_397b_a17b_pretrain_64gpu_vrnvl72_bf16_config"), + ("qwen_vl.vrnvl72.qwen35_vl", "qwen35_vl_397b_a17b_pretrain_64gpu_vrnvl72_fp8cs_config"), + ("qwen_vl.vrnvl72.qwen35_vl", "qwen35_vl_397b_a17b_pretrain_64gpu_vrnvl72_fp8mx_config"), +) + +# VRNVL72 MoE recipes left off NCCL EP: GPT-OSS 20B MXFP8 runs EP=1, so it has no expert-parallel traffic, +# and GPT-OSS 120B BF16 uses alltoall, like the GB300 recipe it aliases. +_VRNVL72_NON_NCCLEP_RECIPES = ( + ("gpt_oss.vrnvl72.gpt_oss", "gpt_oss_20b_pretrain_8gpu_vrnvl72_fp8mx_config"), + ("gpt_oss.vrnvl72.gpt_oss", "gpt_oss_120b_pretrain_64gpu_vrnvl72_bf16_config"), ) @@ -92,7 +116,16 @@ # every recipe module that is already loaded, so a lazy import inside a test would escape the patch. _RECIPE_MODULES = { module_name: importlib.import_module(f"megatron.bridge.perf_recipes.{module_name}") - for module_name in sorted({name for name, _ in (*_GB300_NCCLEP_RECIPES, *_VRNVL72_ALIASES)}) + for module_name in sorted( + { + name + for name, _ in ( + *_GB300_NCCLEP_RECIPES, + *_VRNVL72_NCCLEP_RECIPES, + *_VRNVL72_NON_NCCLEP_RECIPES, + ) + } + ) } @@ -131,12 +164,12 @@ def test_gb300_deepseek_hsdp_uses_hybridep() -> None: assert cfg.env_vars.keys().isdisjoint(_NCCL_EP_ENV) -@pytest.mark.parametrize(("module_name", "factory_name"), _VRNVL72_ALIASES, ids=lambda value: value) -def test_vrnvl72_aliases_do_not_inherit_nccl_ep(module_name: str, factory_name: str) -> None: +@pytest.mark.parametrize(("module_name", "factory_name"), _VRNVL72_NON_NCCLEP_RECIPES, ids=lambda value: value) +def test_vrnvl72_recipes_left_off_nccl_ep_do_not_inherit_it(module_name: str, factory_name: str) -> None: cfg = _build(module_name, factory_name) - # The aliases keep the dispatcher their HybridEP base had upstream (flex/hybridep, or alltoall for the - # DeepSeek BF16 and NVFP4 recipes); the GB300 NCCL EP overlay must never leak into them. + # These recipes keep their own dispatcher (flex/hybridep or alltoall); the NCCL EP overlay must never + # leak into them. assert cfg.model.moe_flex_dispatcher_backend != "ncclep" assert cfg.env_vars.keys().isdisjoint(_NCCL_EP_ENV) if cfg.model.moe_token_dispatcher_type == "flex": diff --git a/tests/unit_tests/recipes/test_moe_full_iteration_recipes.py b/tests/unit_tests/recipes/test_moe_full_iteration_recipes.py index ecc26fb43d..6b8a281893 100644 --- a/tests/unit_tests/recipes/test_moe_full_iteration_recipes.py +++ b/tests/unit_tests/recipes/test_moe_full_iteration_recipes.py @@ -52,7 +52,7 @@ def _assert_full_iteration_hybridep(cfg: ConfigContainer) -> None: assert cfg.ddp.check_for_nan_in_grad is False assert cfg.rerun_state_machine.check_for_nan_in_loss is False - # GB200 and GB300 recipes default to the NCCL EP dispatch stack; every other GPU keeps HybridEP. + # GB200, GB300 and VRNVL72 recipes default to the NCCL EP dispatch stack; every other GPU keeps HybridEP. if cfg.model.moe_flex_dispatcher_backend == "ncclep": assert cfg.model.moe_use_grouped_tensor is True assert cfg.env_vars["NCCL_EP_HT_EM_PULL_PUSH"] == 1 @@ -89,7 +89,6 @@ def test_qwen3_235b_vrnvl72_nvfp4_uses_full_iteration_stack() -> None: assert cfg.comm_overlap.overlap_moe_expert_parallel_comm is True assert cfg.comm_overlap.delay_wgrad_compute is True assert cfg.env_vars["TORCH_NCCL_AVOID_RECORD_STREAMS"] == 0 - assert cfg.env_vars["NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN"] == 32 assert ( cfg.model.tensor_model_parallel_size, cfg.model.pipeline_model_parallel_size, @@ -112,7 +111,6 @@ def test_deepseek_v3_256gpu_vrnvl72_nvfp4_uses_full_iteration_stack() -> None: assert cfg.comm_overlap.delay_wgrad_compute is True assert cfg.model.recompute_modules == [] assert cfg.env_vars["TORCH_NCCL_AVOID_RECORD_STREAMS"] == 1 - assert cfg.env_vars["NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN"] == 32 assert cfg.env_vars["NVTE_DPA_FP8_RECIPE"] == "MXFP8BlockScaling" assert cfg.env_vars["NVTE_DPA_FP8_FORMAT"] == "E4M3" assert ( @@ -136,7 +134,6 @@ def test_deepseek_v3_4gpu_vrnvl72_proxies_use_full_iteration_stack(recipe: Calla assert cfg.comm_overlap.overlap_moe_expert_parallel_comm is True assert cfg.comm_overlap.delay_wgrad_compute is True assert cfg.env_vars["TORCH_NCCL_AVOID_RECORD_STREAMS"] == 0 - assert cfg.env_vars["NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN"] == cfg.model.expert_model_parallel_size == 4 assert ( cfg.model.tensor_model_parallel_size, cfg.model.pipeline_model_parallel_size, diff --git a/tests/unit_tests/recipes/test_nemotron_3_5_lightning_perf_recipes.py b/tests/unit_tests/recipes/test_nemotron_3_5_lightning_perf_recipes.py index 5baaa46e08..8d3947fadd 100644 --- a/tests/unit_tests/recipes/test_nemotron_3_5_lightning_perf_recipes.py +++ b/tests/unit_tests/recipes/test_nemotron_3_5_lightning_perf_recipes.py @@ -21,7 +21,6 @@ import pytest 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, @@ -100,13 +99,15 @@ nemotron_3_5_lightning_pretrain_8gpu_gb300_fp8mx_config, nemotron_3_5_lightning_pretrain_8gpu_gb300_nvfp4_config, ) -# GB200 and GB300 perf recipes covered by nemo-ci dispatch through NCCL EP (no HybridEP topology in their -# environment). +# The GB200, GB300 and VRNVL72 BF16 and MXFP8 perf recipes dispatch through NCCL EP (no HybridEP topology in +# their environment). _NCCLEP_RECIPES = ( nemotron_3_5_lightning_pretrain_8gpu_gb200_bf16_config, nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_config, nemotron_3_5_lightning_pretrain_8gpu_gb300_bf16_config, nemotron_3_5_lightning_pretrain_8gpu_gb300_fp8mx_config, + nemotron_3_5_lightning_pretrain_8gpu_vrnvl72_bf16_config, + nemotron_3_5_lightning_pretrain_8gpu_vrnvl72_fp8mx_config, ) _GB_FSDP_RECIPES = ( nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_fsdp_config, @@ -658,16 +659,17 @@ def test_vrnvl72_perf_recipes_match_gb300_configs( vrnvl72_factory: Callable[[], ConfigContainer], gb300_factory: Callable[[], ConfigContainer], ) -> None: - """VRNVL72 Nemotron 3.5 recipes match their GB300 baselines up to the GB300 NCCL EP overlay. + """VRNVL72 Nemotron 3.5 recipes match their GB300 baselines apart from the launch environment. - The GB300 recipes default to NCCL EP (dispatcher backend, device-side expert counts and the - NCCL EP process setting); the VRNVL72 aliases keep the shared HybridEP base. The VRNVL72 BF16 - recipe also offloads expert activations. Everything else must be identical. + Both platforms dispatch through NCCL EP with device-side expert counts. The VRNVL72 BF16 recipe + also offloads expert activations. Everything else must be identical. """ vrnvl72_cfg = vrnvl72_factory() gb300_cfg = gb300_factory() - assert vrnvl72_cfg.model.moe_flex_dispatcher_backend == "hybridep" + assert vrnvl72_cfg.model.moe_flex_dispatcher_backend == "ncclep" assert gb300_cfg.model.moe_flex_dispatcher_backend == "ncclep" + assert vrnvl72_cfg.env_vars["NCCL_EP_HT_EM_PULL_PUSH"] == 1 + assert vrnvl72_cfg.env_vars.keys().isdisjoint(HYBRID_EP_ENV_NAMES) if vrnvl72_factory is nemotron_3_5_lightning_pretrain_8gpu_vrnvl72_bf16_config: assert vrnvl72_cfg.model.fine_grained_activation_offloading is True @@ -677,8 +679,6 @@ def test_vrnvl72_perf_recipes_match_gb300_configs( vrnvl72_cfg.model.fine_grained_activation_offloading = gb300_cfg.model.fine_grained_activation_offloading vrnvl72_cfg.model.offload_modules = gb300_cfg.model.offload_modules - _enable_ncclep(vrnvl72_cfg) - vrnvl72_cfg.model.moe_use_grouped_tensor = True vrnvl72_cfg.env_vars = gb300_cfg.env_vars assert vrnvl72_cfg == gb300_cfg diff --git a/tests/unit_tests/scripts/performance/test_qwen35_vl_workload_base_configs.py b/tests/unit_tests/scripts/performance/test_qwen35_vl_workload_base_configs.py index f6c272fb3d..4dbc4edc7a 100644 --- a/tests/unit_tests/scripts/performance/test_qwen35_vl_workload_base_configs.py +++ b/tests/unit_tests/scripts/performance/test_qwen35_vl_workload_base_configs.py @@ -367,19 +367,17 @@ def test_qwen35_vl_122b_gb300_keeps_cuda_graphs_disabled(monkeypatch: pytest.Mon qwen35_vl_397b_a17b_pretrain_64gpu_gb300_bf16_config, qwen35_vl_397b_a17b_pretrain_64gpu_gb300_fp8cs_config, qwen35_vl_397b_a17b_pretrain_64gpu_gb300_fp8mx_config, - qwen35_vl_397b_a17b_pretrain_64gpu_vrnvl72_bf16_config, - qwen35_vl_397b_a17b_pretrain_64gpu_vrnvl72_fp8cs_config, - qwen35_vl_397b_a17b_pretrain_64gpu_vrnvl72_fp8mx_config, ], ) def test_qwen35_vl_397b_hybridep_domain_matches_expert_parallel_size( recipe_fn: Callable, monkeypatch: pytest.MonkeyPatch ) -> None: - """The VRNVL72 397B recipes inherit EP from GB300 but restate the HybridEP env inline. + """The GB300 397B recipes state the HybridEP env inline, next to their EP size. HybridEP asserts that the EP group size is divisible by ``NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN``, and the launcher only re-derives that env when a CLI override changes EP, so the two must stay in sync here. + The VRNVL72 397B recipes run NCCL EP and carry no HybridEP env. """ patch_recipe_construction_dependencies(monkeypatch) @@ -465,11 +463,11 @@ def test_qwen35_vl_35b_vrnvl72_fp8mx_uses_the_vrnvl72_settings(monkeypatch: pyte assert config.model.use_grouped_gemm_for_shared_expert is True assert config.model.moe_shared_expert_glu_interleave_size == 32 - assert config.model.moe_flex_dispatcher_backend == "hybridep" - assert config.model.moe_permute_fusion_into_hybridep is False - assert config.model.moe_hybridep_num_sms_preprocessing == 16 + assert config.model.moe_flex_dispatcher_backend == "ncclep" + assert config.model.moe_use_grouped_tensor is True assert config.model.overlap_dispatch_backward_with_experts_wgrad is True - assert config.env_vars["NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN"] == config.model.expert_model_parallel_size + assert config.env_vars["NCCL_EP_HT_EM_PULL_PUSH"] == 1 + assert not any(name.startswith(("NUM_OF_HYBRID_EP_", "NUM_OF_TOKENS_PER_CHUNK_")) for name in config.env_vars) assert config.model.cuda_graph_impl == "transformer_engine" assert set(cuda_graph_module_names(config.model)) == {"attn", "moe_router", "moe_preprocess"} @@ -499,7 +497,8 @@ def test_qwen35_vl_397b_vrnvl72_fp8mx_uses_the_vrnvl72_settings(monkeypatch: pyt _assert_cutedsl_grouped_mlp(config) assert config.model.recompute_granularity == "selective" assert config.model.recompute_modules == ["gdn_norm_out"] - assert config.model.moe_permute_fusion_into_hybridep is False + assert config.model.moe_flex_dispatcher_backend == "ncclep" + assert config.model.moe_use_grouped_tensor is True assert getattr(config.model, "use_grouped_gemm_for_shared_expert", False) is False assert getattr(config.model, "overlap_dispatch_backward_with_experts_wgrad", False) is False @@ -507,9 +506,8 @@ def test_qwen35_vl_397b_vrnvl72_fp8mx_uses_the_vrnvl72_settings(monkeypatch: pyt assert set(cuda_graph_module_names(config.model)) == {"attn", "moe_router", "moe_preprocess"} _assert_vision_graph_covers_one_microbatch(config) - assert config.env_vars["NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN"] == 32 - assert config.env_vars["NUM_OF_TOKENS_PER_CHUNK_DISPATCH_API"] == 128 - assert config.env_vars["NUM_OF_TOKENS_PER_CHUNK_PREPROCESSING_API"] == 128 + assert config.env_vars["NCCL_EP_HT_EM_PULL_PUSH"] == 1 + assert not any(name.startswith(("NUM_OF_HYBRID_EP_", "NUM_OF_TOKENS_PER_CHUNK_")) for name in config.env_vars) assert config.env_vars["NVTE_FWD_LAYERNORM_SM_MARGIN"] == 16 assert config.env_vars["NVTE_BWD_LAYERNORM_SM_MARGIN"] == 16 assert "NVTE_NORM_FWD_USE_CUDNN" not in config.env_vars From c25b2da0bc2fc16997c67874360ad96a1ad93ffb Mon Sep 17 00:00:00 2001 From: Malay Nagda Date: Fri, 16 Oct 2026 13:06:38 +0530 Subject: [PATCH 6/9] docs(recipe): list VRNVL72 in the NCCL EP helper docstring Co-Authored-By: Claude Opus 5.5 Signed-off-by: Malay Nagda --- src/megatron/bridge/perf_recipes/_common.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/megatron/bridge/perf_recipes/_common.py b/src/megatron/bridge/perf_recipes/_common.py index 31db3bde78..32f48f8aeb 100644 --- a/src/megatron/bridge/perf_recipes/_common.py +++ b/src/megatron/bridge/perf_recipes/_common.py @@ -96,7 +96,7 @@ def _benchmark_common( def _enable_ncclep(cfg: ConfigContainer) -> None: - """Switch a flex-dispatcher recipe to the NCCL EP backend (GB200 and GB300 MoE default). + """Switch a flex-dispatcher recipe to the NCCL EP backend (GB200, GB300 and VRNVL72 MoE default). Call after ``_benchmark_common``. Only the dispatch stack changes; parallelism, precision, recompute and the CUDA graph mode stay as the recipe set them. The calling recipe still declares From 7ed58d47448eaee42431c47a552469de0de27fc7 Mon Sep 17 00:00:00 2001 From: Malay Nagda Date: Fri, 16 Oct 2026 13:20:10 +0530 Subject: [PATCH 7/9] perf(recipe): keep untracked VRNVL72 MoE perf recipes on HybridEP Restore the VRNVL72 GPT-OSS 20B NVFP4, Kimi K2, Nemotron 3 Nano, FP8-CS (DeepSeek V3 and Qwen3.5-VL) and Qwen3.5-VL 122B perf recipes to main, so only the tracked VRNVL72 benchmarks switch to NCCL EP. Their code is now identical to main. The GB300 NCCL EP test lists these recipes with the VRNVL72 recipes that stay off NCCL EP, and the Qwen3.5-VL 397B HybridEP domain test covers the VRNVL72 FP8-CS recipe again. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Malay Nagda --- .../deepseek/vrnvl72/deepseek_v3.py | 26 ++++---- .../perf_recipes/gpt_oss/vrnvl72/gpt_oss.py | 26 ++++---- .../perf_recipes/kimi/vrnvl72/kimi_k2.py | 27 ++++---- .../nemotronh/vrnvl72/nemotronh.py | 39 +++++------ .../perf_recipes/qwen_vl/vrnvl72/qwen35_vl.py | 65 +++++++++---------- .../recipes/test_gb300_ncclep_perf_recipes.py | 34 +++++----- .../test_qwen35_vl_workload_base_configs.py | 5 +- 7 files changed, 104 insertions(+), 118 deletions(-) diff --git a/src/megatron/bridge/perf_recipes/deepseek/vrnvl72/deepseek_v3.py b/src/megatron/bridge/perf_recipes/deepseek/vrnvl72/deepseek_v3.py index bd166d96d1..cea8d2d4f4 100644 --- a/src/megatron/bridge/perf_recipes/deepseek/vrnvl72/deepseek_v3.py +++ b/src/megatron/bridge/perf_recipes/deepseek/vrnvl72/deepseek_v3.py @@ -85,7 +85,7 @@ def deepseek_v3_pretrain_128gpu_vrnvl72_bf16_config() -> ConfigContainer: def deepseek_v3_pretrain_128gpu_vrnvl72_fp8cs_config() -> ConfigContainer: - """DeepSeek V3 pretrain: 128× VRNVL72, FP8 current-scaling, NCCL EP.""" + """DeepSeek V3 pretrain: 128× VRNVL72, FP8 current-scaling.""" cfg = deepseek_v3_pretrain_config() cfg.mixed_precision = _perf_precision("fp8_cs") _deepseek_v3_common(cfg) @@ -110,10 +110,6 @@ def deepseek_v3_pretrain_128gpu_vrnvl72_fp8cs_config() -> ConfigContainer: set_deepseek_v3_pipeline_model_parallel_layout(cfg.model) _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, @@ -125,8 +121,11 @@ def deepseek_v3_pretrain_128gpu_vrnvl72_fp8cs_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # NCCL EP dispatcher mode and one GPU per rank. - "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, @@ -279,12 +278,8 @@ def deepseek_v3_pretrain_256gpu_vrnvl72_bf16_config() -> ConfigContainer: def deepseek_v3_pretrain_256gpu_vrnvl72_fp8cs_config() -> ConfigContainer: - """DeepSeek V3 pretrain: 256× VRNVL72, FP8-CS, NCCL EP (alias of GB300).""" + """DeepSeek V3 pretrain: 256× VRNVL72, FP8-CS (alias of GB300).""" cfg = deepseek_v3_pretrain_256gpu_gb300_fp8cs_config() - _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, @@ -296,8 +291,11 @@ def deepseek_v3_pretrain_256gpu_vrnvl72_fp8cs_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # NCCL EP dispatcher mode and one GPU per rank. - "NCCL_EP_HT_EM_PULL_PUSH": 1, + # HybridEP topology for the target system. + "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 32, + "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, diff --git a/src/megatron/bridge/perf_recipes/gpt_oss/vrnvl72/gpt_oss.py b/src/megatron/bridge/perf_recipes/gpt_oss/vrnvl72/gpt_oss.py index 40543f2353..3a2a21f7c0 100644 --- a/src/megatron/bridge/perf_recipes/gpt_oss/vrnvl72/gpt_oss.py +++ b/src/megatron/bridge/perf_recipes/gpt_oss/vrnvl72/gpt_oss.py @@ -32,12 +32,8 @@ def gpt_oss_20b_pretrain_8gpu_vrnvl72_nvfp4_config() -> ConfigContainer: - """GPT-OSS 20B pretrain: 8× VRNVL72, NVFP4, NCCL EP (alias of GB300).""" + """GPT-OSS 20B pretrain: 8× VRNVL72, NVFP4 (alias of GB300).""" cfg = gpt_oss_20b_pretrain_8gpu_gb300_nvfp4_config() - _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, @@ -48,8 +44,11 @@ def gpt_oss_20b_pretrain_8gpu_vrnvl72_nvfp4_config() -> ConfigContainer: # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 1, "NCCL_CTA_POLICY": 1, - # NCCL EP dispatcher mode and one GPU per rank. - "NCCL_EP_HT_EM_PULL_PUSH": 1, + # HybridEP topology for the target system. + "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 4, + "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, @@ -111,7 +110,7 @@ def gpt_oss_20b_pretrain_8gpu_vrnvl72_fp8mx_config() -> ConfigContainer: def gpt_oss_20b_pretrain_64gpu_vrnvl72_nvfp4_config() -> ConfigContainer: - """GPT-OSS 20B pretrain: 64× VRNVL72, NVFP4, NCCL EP.""" + """GPT-OSS 20B pretrain: 64× VRNVL72, NVFP4.""" cfg = gpt_oss_20b_pretrain_config() cfg.mixed_precision = _gpt_oss_20b_nvfp4_precision() @@ -126,10 +125,6 @@ def gpt_oss_20b_pretrain_64gpu_vrnvl72_nvfp4_config() -> ConfigContainer: _benchmark_common(cfg) _apply_gpt_oss_20b_common_configs(cfg) _apply_gpt_oss_20b_transformer_engine_graph_configs(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 cfg.model.cuda_graph_warmup_steps = 1 cfg.optimizer.lr = 0.0006 @@ -147,8 +142,11 @@ def gpt_oss_20b_pretrain_64gpu_vrnvl72_nvfp4_config() -> ConfigContainer: # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 1, "NCCL_CTA_POLICY": 1, - # NCCL EP dispatcher mode and one GPU per rank. - "NCCL_EP_HT_EM_PULL_PUSH": 1, + # HybridEP topology for the target system. + "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 4, + "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, diff --git a/src/megatron/bridge/perf_recipes/kimi/vrnvl72/kimi_k2.py b/src/megatron/bridge/perf_recipes/kimi/vrnvl72/kimi_k2.py index abf7aaa48f..208168dda5 100644 --- a/src/megatron/bridge/perf_recipes/kimi/vrnvl72/kimi_k2.py +++ b/src/megatron/bridge/perf_recipes/kimi/vrnvl72/kimi_k2.py @@ -13,7 +13,6 @@ # limitations under the License. """VRNVL72 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, @@ -25,12 +24,8 @@ def kimi_k2_pretrain_256gpu_vrnvl72_bf16_config() -> ConfigContainer: - """Kimi K2 pretrain: 256× VRNVL72, BF16, NCCL EP (alias of GB300).""" + """Kimi K2 pretrain: 256× VRNVL72, BF16 (alias of GB300).""" cfg = kimi_k2_pretrain_256gpu_gb300_bf16_config() - _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, @@ -42,8 +37,11 @@ def kimi_k2_pretrain_256gpu_vrnvl72_bf16_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # NCCL EP dispatcher mode and one GPU per rank. - "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, @@ -52,12 +50,8 @@ def kimi_k2_pretrain_256gpu_vrnvl72_bf16_config() -> ConfigContainer: def kimi_k2_pretrain_256gpu_vrnvl72_fp8mx_config() -> ConfigContainer: - """Kimi K2 pretrain: 256× VRNVL72, MXFP8, NCCL EP (alias of GB300).""" + """Kimi K2 pretrain: 256× VRNVL72, MXFP8 (alias of GB300).""" cfg = _build_kimi_k2_gb300_fp8mx() - _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, @@ -69,8 +63,11 @@ def kimi_k2_pretrain_256gpu_vrnvl72_fp8mx_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # NCCL EP dispatcher mode and one GPU per rank. - "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, diff --git a/src/megatron/bridge/perf_recipes/nemotronh/vrnvl72/nemotronh.py b/src/megatron/bridge/perf_recipes/nemotronh/vrnvl72/nemotronh.py index 293f6158d2..e8f42a3aea 100644 --- a/src/megatron/bridge/perf_recipes/nemotronh/vrnvl72/nemotronh.py +++ b/src/megatron/bridge/perf_recipes/nemotronh/vrnvl72/nemotronh.py @@ -34,12 +34,8 @@ def nemotron_3_nano_pretrain_8gpu_vrnvl72_bf16_config() -> ConfigContainer: - """Nemotron 3 Nano pretrain: 8× VRNVL72, BF16, NCCL EP (alias of GB300).""" + """Nemotron 3 Nano pretrain: 8× VRNVL72, BF16 (alias of GB300).""" cfg = _build_nemotron_3_nano_gb300_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. - 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, @@ -51,8 +47,11 @@ def nemotron_3_nano_pretrain_8gpu_vrnvl72_bf16_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # NCCL EP dispatcher mode and one GPU per rank. - "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, @@ -64,12 +63,8 @@ def nemotron_3_nano_pretrain_8gpu_vrnvl72_bf16_config() -> ConfigContainer: def nemotron_3_nano_pretrain_8gpu_vrnvl72_fp8mx_config() -> ConfigContainer: - """Nemotron 3 Nano pretrain: 8× VRNVL72, FP8-MX, NCCL EP (alias of GB300).""" + """Nemotron 3 Nano pretrain: 8× VRNVL72, FP8-MX (alias of GB300).""" cfg = _build_nemotron_3_nano_gb300_mxfp8() - _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.model.use_transformer_engine_op_fuser = True cfg.model.moe_mlp_glu_interleave_size = 32 cfg.mixed_precision.fp8_dot_product_attention = True @@ -84,8 +79,11 @@ def nemotron_3_nano_pretrain_8gpu_vrnvl72_fp8mx_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # NCCL EP dispatcher mode and one GPU per rank. - "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, @@ -99,12 +97,8 @@ def nemotron_3_nano_pretrain_8gpu_vrnvl72_fp8mx_config() -> ConfigContainer: def nemotron_3_nano_pretrain_8gpu_vrnvl72_nvfp4_config() -> ConfigContainer: - """Nemotron 3 Nano pretrain: 8× VRNVL72, NVFP4, NCCL EP (alias of GB300).""" + """Nemotron 3 Nano pretrain: 8× VRNVL72, NVFP4 (alias of GB300).""" cfg = nemotron_3_nano_pretrain_8gpu_gb300_nvfp4_config() - _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, @@ -116,8 +110,11 @@ def nemotron_3_nano_pretrain_8gpu_vrnvl72_nvfp4_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # NCCL EP dispatcher mode and one GPU per rank. - "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, diff --git a/src/megatron/bridge/perf_recipes/qwen_vl/vrnvl72/qwen35_vl.py b/src/megatron/bridge/perf_recipes/qwen_vl/vrnvl72/qwen35_vl.py index d7dbe76284..4867e6c85a 100644 --- a/src/megatron/bridge/perf_recipes/qwen_vl/vrnvl72/qwen35_vl.py +++ b/src/megatron/bridge/perf_recipes/qwen_vl/vrnvl72/qwen35_vl.py @@ -60,12 +60,8 @@ def qwen35_vl_35b_a3b_pretrain_8gpu_vrnvl72_bf16_config() -> ConfigContainer: def qwen35_vl_35b_a3b_pretrain_8gpu_vrnvl72_fp8cs_config() -> ConfigContainer: - """Qwen3.5-VL 35B-A3B pretrain: 8× VRNVL72, FP8-CS, NCCL EP (alias of GB300).""" + """Qwen3.5-VL 35B-A3B pretrain: 8× VRNVL72, FP8-CS (alias of GB300).""" cfg = qwen35_vl_35b_a3b_pretrain_8gpu_gb300_fp8cs_config() - _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 # The GB300 base selects the cuDNN GDN kernel. On VRNVL72 it was measured only with MXFP8 (against FLA # at 35B), so this precision keeps FLA until it is measured too. _select_gdn_kernel_backend(cfg, "fla") @@ -80,8 +76,11 @@ def qwen35_vl_35b_a3b_pretrain_8gpu_vrnvl72_fp8cs_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # NCCL EP dispatcher mode and one GPU per rank. - "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, @@ -164,12 +163,8 @@ def qwen35_vl_35b_a3b_pretrain_8gpu_vrnvl72_fp8mx_config() -> ConfigContainer: def qwen35_vl_122b_a10b_pretrain_32gpu_vrnvl72_bf16_config() -> ConfigContainer: - """Qwen3.5-VL 122B-A10B pretrain: 32× VRNVL72, BF16, NCCL EP (alias of GB300).""" + """Qwen3.5-VL 122B-A10B pretrain: 32× VRNVL72, BF16 (alias of GB300).""" cfg = qwen35_vl_122b_a10b_pretrain_32gpu_gb300_bf16_config() - _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, @@ -181,8 +176,11 @@ def qwen35_vl_122b_a10b_pretrain_32gpu_vrnvl72_bf16_config() -> ConfigContainer: "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # NCCL EP dispatcher mode and one GPU per rank. - "NCCL_EP_HT_EM_PULL_PUSH": 1, + # HybridEP topology for the target system. + "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 32, + "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, @@ -191,12 +189,8 @@ def qwen35_vl_122b_a10b_pretrain_32gpu_vrnvl72_bf16_config() -> ConfigContainer: def qwen35_vl_122b_a10b_pretrain_32gpu_vrnvl72_fp8cs_config() -> ConfigContainer: - """Qwen3.5-VL 122B-A10B pretrain: 32× VRNVL72, FP8-CS, NCCL EP (alias of GB300).""" + """Qwen3.5-VL 122B-A10B pretrain: 32× VRNVL72, FP8-CS (alias of GB300).""" cfg = qwen35_vl_122b_a10b_pretrain_32gpu_gb300_fp8cs_config() - _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, @@ -208,8 +202,11 @@ def qwen35_vl_122b_a10b_pretrain_32gpu_vrnvl72_fp8cs_config() -> ConfigContainer "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # NCCL EP dispatcher mode and one GPU per rank. - "NCCL_EP_HT_EM_PULL_PUSH": 1, + # HybridEP topology for the target system. + "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 32, + "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, @@ -218,12 +215,8 @@ def qwen35_vl_122b_a10b_pretrain_32gpu_vrnvl72_fp8cs_config() -> ConfigContainer def qwen35_vl_122b_a10b_pretrain_32gpu_vrnvl72_fp8mx_config() -> ConfigContainer: - """Qwen3.5-VL 122B-A10B pretrain: 32× VRNVL72, FP8-MX, NCCL EP (alias of GB300).""" + """Qwen3.5-VL 122B-A10B pretrain: 32× VRNVL72, FP8-MX (alias of GB300).""" cfg = qwen35_vl_122b_a10b_pretrain_32gpu_gb300_fp8mx_config() - _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, @@ -235,8 +228,11 @@ def qwen35_vl_122b_a10b_pretrain_32gpu_vrnvl72_fp8mx_config() -> ConfigContainer "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # NCCL EP dispatcher mode and one GPU per rank. - "NCCL_EP_HT_EM_PULL_PUSH": 1, + # HybridEP topology for the target system. + "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 32, + "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, @@ -275,12 +271,8 @@ def qwen35_vl_397b_a17b_pretrain_64gpu_vrnvl72_bf16_config() -> ConfigContainer: def qwen35_vl_397b_a17b_pretrain_64gpu_vrnvl72_fp8cs_config() -> ConfigContainer: - """Qwen3.5-VL 397B-A17B pretrain: 64× VRNVL72, FP8-CS, NCCL EP (alias of GB300).""" + """Qwen3.5-VL 397B-A17B pretrain: 64× VRNVL72, FP8-CS (alias of GB300).""" cfg = qwen35_vl_397b_a17b_pretrain_64gpu_gb300_fp8cs_config() - _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 # The GB300 base selects the cuDNN GDN kernel. On VRNVL72 it was measured only with MXFP8 (against FLA # at 35B), so this precision keeps FLA until it is measured too. _select_gdn_kernel_backend(cfg, "fla") @@ -295,8 +287,11 @@ def qwen35_vl_397b_a17b_pretrain_64gpu_vrnvl72_fp8cs_config() -> ConfigContainer "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # NCCL EP dispatcher mode and one GPU per rank. - "NCCL_EP_HT_EM_PULL_PUSH": 1, + # HybridEP topology for the target system. + "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 32, + "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, diff --git a/tests/unit_tests/recipes/test_gb300_ncclep_perf_recipes.py b/tests/unit_tests/recipes/test_gb300_ncclep_perf_recipes.py index 0827a72363..38dc3caa6f 100644 --- a/tests/unit_tests/recipes/test_gb300_ncclep_perf_recipes.py +++ b/tests/unit_tests/recipes/test_gb300_ncclep_perf_recipes.py @@ -57,14 +57,12 @@ ) # VRNVL72 MoE recipes on NCCL EP. The Nemotron 3 Ultra aliases share the GB300 NCCL EP builders; the rest -# switch their GB300 HybridEP base, a GB300 HybridEP recipe, or their own HybridEP setup to NCCL EP. +# switch their GB300 HybridEP base or recipe, or their own HybridEP setup, to NCCL EP. _VRNVL72_NCCLEP_RECIPES = ( ("deepseek.vrnvl72.deepseek_v3", "deepseek_v3_pretrain_128gpu_vrnvl72_bf16_config"), - ("deepseek.vrnvl72.deepseek_v3", "deepseek_v3_pretrain_128gpu_vrnvl72_fp8cs_config"), ("deepseek.vrnvl72.deepseek_v3", "deepseek_v3_pretrain_128gpu_vrnvl72_fp8mx_config"), ("deepseek.vrnvl72.deepseek_v3", "deepseek_v3_pretrain_128gpu_vrnvl72_nvfp4_config"), ("deepseek.vrnvl72.deepseek_v3", "deepseek_v3_pretrain_256gpu_vrnvl72_bf16_config"), - ("deepseek.vrnvl72.deepseek_v3", "deepseek_v3_pretrain_256gpu_vrnvl72_fp8cs_config"), ("deepseek.vrnvl72.deepseek_v3", "deepseek_v3_pretrain_256gpu_vrnvl72_fp8mx_config"), ("deepseek.vrnvl72.deepseek_v3", "deepseek_v3_pretrain_256gpu_vrnvl72_nvfp4_config"), ("deepseek.vrnvl72.deepseek_v3", "deepseek_v3_pretrain_64gpu_vrnvl72_fp8mx_proxy_config"), @@ -72,14 +70,7 @@ ("deepseek.vrnvl72.deepseek_v3", "deepseek_v3_pretrain_4gpu_vrnvl72_fp8mx_proxy_config"), ("deepseek.vrnvl72.deepseek_v3", "deepseek_v3_pretrain_4gpu_vrnvl72_nvfp4_proxy_config"), ("deepseek.vrnvl72.deepseek_v4", "deepseek_v4_flash_pretrain_128gpu_vrnvl72_fp8mx_config"), - ("gpt_oss.vrnvl72.gpt_oss", "gpt_oss_20b_pretrain_8gpu_vrnvl72_nvfp4_config"), - ("gpt_oss.vrnvl72.gpt_oss", "gpt_oss_20b_pretrain_64gpu_vrnvl72_nvfp4_config"), ("gpt_oss.vrnvl72.gpt_oss", "gpt_oss_120b_pretrain_64gpu_vrnvl72_fp8mx_config"), - ("kimi.vrnvl72.kimi_k2", "kimi_k2_pretrain_256gpu_vrnvl72_bf16_config"), - ("kimi.vrnvl72.kimi_k2", "kimi_k2_pretrain_256gpu_vrnvl72_fp8mx_config"), - ("nemotronh.vrnvl72.nemotronh", "nemotron_3_nano_pretrain_8gpu_vrnvl72_bf16_config"), - ("nemotronh.vrnvl72.nemotronh", "nemotron_3_nano_pretrain_8gpu_vrnvl72_fp8mx_config"), - ("nemotronh.vrnvl72.nemotronh", "nemotron_3_nano_pretrain_8gpu_vrnvl72_nvfp4_config"), ("nemotronh.vrnvl72.nemotronh", "nemotron_3_super_pretrain_64gpu_vrnvl72_bf16_config"), ("nemotronh.vrnvl72.nemotronh", "nemotron_3_super_pretrain_64gpu_vrnvl72_fp8mx_config"), ("nemotronh.vrnvl72.nemotronh", "nemotron_3_super_pretrain_64gpu_vrnvl72_nvfp4_config"), @@ -94,21 +85,30 @@ ("qwen.vrnvl72.qwen3_moe", "qwen3_30b_a3b_pretrain_8gpu_vrnvl72_fp8mx_config"), ("qwen.vrnvl72.qwen3_moe", "qwen3_30b_a3b_pretrain_8gpu_vrnvl72_nvfp4_config"), ("qwen_vl.vrnvl72.qwen35_vl", "qwen35_vl_35b_a3b_pretrain_8gpu_vrnvl72_bf16_config"), - ("qwen_vl.vrnvl72.qwen35_vl", "qwen35_vl_35b_a3b_pretrain_8gpu_vrnvl72_fp8cs_config"), ("qwen_vl.vrnvl72.qwen35_vl", "qwen35_vl_35b_a3b_pretrain_8gpu_vrnvl72_fp8mx_config"), - ("qwen_vl.vrnvl72.qwen35_vl", "qwen35_vl_122b_a10b_pretrain_32gpu_vrnvl72_bf16_config"), - ("qwen_vl.vrnvl72.qwen35_vl", "qwen35_vl_122b_a10b_pretrain_32gpu_vrnvl72_fp8cs_config"), - ("qwen_vl.vrnvl72.qwen35_vl", "qwen35_vl_122b_a10b_pretrain_32gpu_vrnvl72_fp8mx_config"), ("qwen_vl.vrnvl72.qwen35_vl", "qwen35_vl_397b_a17b_pretrain_64gpu_vrnvl72_bf16_config"), - ("qwen_vl.vrnvl72.qwen35_vl", "qwen35_vl_397b_a17b_pretrain_64gpu_vrnvl72_fp8cs_config"), ("qwen_vl.vrnvl72.qwen35_vl", "qwen35_vl_397b_a17b_pretrain_64gpu_vrnvl72_fp8mx_config"), ) -# VRNVL72 MoE recipes left off NCCL EP: GPT-OSS 20B MXFP8 runs EP=1, so it has no expert-parallel traffic, -# and GPT-OSS 120B BF16 uses alltoall, like the GB300 recipe it aliases. +# VRNVL72 MoE recipes that stay off NCCL EP: GPT-OSS 20B, Kimi K2, Nemotron 3 Nano, the FP8-CS recipes and +# Qwen3.5-VL 122B stay on HybridEP, and GPT-OSS 120B BF16 uses alltoall, like the GB300 recipe it aliases. _VRNVL72_NON_NCCLEP_RECIPES = ( + ("deepseek.vrnvl72.deepseek_v3", "deepseek_v3_pretrain_128gpu_vrnvl72_fp8cs_config"), + ("deepseek.vrnvl72.deepseek_v3", "deepseek_v3_pretrain_256gpu_vrnvl72_fp8cs_config"), + ("gpt_oss.vrnvl72.gpt_oss", "gpt_oss_20b_pretrain_8gpu_vrnvl72_nvfp4_config"), ("gpt_oss.vrnvl72.gpt_oss", "gpt_oss_20b_pretrain_8gpu_vrnvl72_fp8mx_config"), + ("gpt_oss.vrnvl72.gpt_oss", "gpt_oss_20b_pretrain_64gpu_vrnvl72_nvfp4_config"), ("gpt_oss.vrnvl72.gpt_oss", "gpt_oss_120b_pretrain_64gpu_vrnvl72_bf16_config"), + ("kimi.vrnvl72.kimi_k2", "kimi_k2_pretrain_256gpu_vrnvl72_bf16_config"), + ("kimi.vrnvl72.kimi_k2", "kimi_k2_pretrain_256gpu_vrnvl72_fp8mx_config"), + ("nemotronh.vrnvl72.nemotronh", "nemotron_3_nano_pretrain_8gpu_vrnvl72_bf16_config"), + ("nemotronh.vrnvl72.nemotronh", "nemotron_3_nano_pretrain_8gpu_vrnvl72_fp8mx_config"), + ("nemotronh.vrnvl72.nemotronh", "nemotron_3_nano_pretrain_8gpu_vrnvl72_nvfp4_config"), + ("qwen_vl.vrnvl72.qwen35_vl", "qwen35_vl_35b_a3b_pretrain_8gpu_vrnvl72_fp8cs_config"), + ("qwen_vl.vrnvl72.qwen35_vl", "qwen35_vl_122b_a10b_pretrain_32gpu_vrnvl72_bf16_config"), + ("qwen_vl.vrnvl72.qwen35_vl", "qwen35_vl_122b_a10b_pretrain_32gpu_vrnvl72_fp8cs_config"), + ("qwen_vl.vrnvl72.qwen35_vl", "qwen35_vl_122b_a10b_pretrain_32gpu_vrnvl72_fp8mx_config"), + ("qwen_vl.vrnvl72.qwen35_vl", "qwen35_vl_397b_a17b_pretrain_64gpu_vrnvl72_fp8cs_config"), ) diff --git a/tests/unit_tests/scripts/performance/test_qwen35_vl_workload_base_configs.py b/tests/unit_tests/scripts/performance/test_qwen35_vl_workload_base_configs.py index 4dbc4edc7a..5a7567b299 100644 --- a/tests/unit_tests/scripts/performance/test_qwen35_vl_workload_base_configs.py +++ b/tests/unit_tests/scripts/performance/test_qwen35_vl_workload_base_configs.py @@ -367,17 +367,18 @@ def test_qwen35_vl_122b_gb300_keeps_cuda_graphs_disabled(monkeypatch: pytest.Mon qwen35_vl_397b_a17b_pretrain_64gpu_gb300_bf16_config, qwen35_vl_397b_a17b_pretrain_64gpu_gb300_fp8cs_config, qwen35_vl_397b_a17b_pretrain_64gpu_gb300_fp8mx_config, + qwen35_vl_397b_a17b_pretrain_64gpu_vrnvl72_fp8cs_config, ], ) def test_qwen35_vl_397b_hybridep_domain_matches_expert_parallel_size( recipe_fn: Callable, monkeypatch: pytest.MonkeyPatch ) -> None: - """The GB300 397B recipes state the HybridEP env inline, next to their EP size. + """The GB300 397B recipes and the VRNVL72 FP8-CS alias state the HybridEP env inline. HybridEP asserts that the EP group size is divisible by ``NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN``, and the launcher only re-derives that env when a CLI override changes EP, so the two must stay in sync here. - The VRNVL72 397B recipes run NCCL EP and carry no HybridEP env. + The VRNVL72 397B BF16 and MXFP8 recipes run NCCL EP and carry no HybridEP env. """ patch_recipe_construction_dependencies(monkeypatch) From 5f8300eaafc380c508b5a96668188b972ffb4231 Mon Sep 17 00:00:00 2001 From: Malay Nagda Date: Fri, 16 Oct 2026 13:29:12 +0530 Subject: [PATCH 8/9] perf(recipe): drop the deterministic TE pin from DeepSeek NCCL EP perf recipes Remove NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 from the 16 DeepSeek perf recipes this PR runs on NCCL EP: GB200 DeepSeek V3 BF16, MXFP8, NVFP4 and MXFP8 large-scale and DeepSeek V4 Flash, and the VRNVL72 DeepSeek V3 BF16, MXFP8 and NVFP4 recipes (128 and 256 GPUs), their 64- and 4-GPU proxies, and DeepSeek V4 Flash. This matches the GB300 DeepSeek V3 recipes. Recent Transformer Engine accumulates the routing-probability gradient of its op-fuser grouped MLP with atomics and raises when deterministic kernels are requested. DeepSeek V4 Flash on GB200 failed in its first backward pass this way in nemo-ci. The DeepSeek recipes that stay on HybridEP keep the pin. The environment invariant test now requires DeepSeek NCCL EP recipes to leave the setting unpinned, and the GB200 and VRNVL72 DeepSeek V4 Flash tests check it. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Malay Nagda --- .../perf_recipes/deepseek/gb200/deepseek_v3.py | 8 -------- .../perf_recipes/deepseek/gb200/deepseek_v4.py | 1 - .../deepseek/vrnvl72/deepseek_v3.py | 18 ------------------ .../deepseek/vrnvl72/deepseek_v4.py | 1 - .../recipes/test_perf_config_integration.py | 6 +----- .../test_deepseek_v4_flash_perf_recipes.py | 9 ++++++--- .../recipes/test_perf_recipe_environment.py | 16 ++++++++-------- 7 files changed, 15 insertions(+), 44 deletions(-) diff --git a/src/megatron/bridge/perf_recipes/deepseek/gb200/deepseek_v3.py b/src/megatron/bridge/perf_recipes/deepseek/gb200/deepseek_v3.py index cc17b9bbd5..2710624f51 100644 --- a/src/megatron/bridge/perf_recipes/deepseek/gb200/deepseek_v3.py +++ b/src/megatron/bridge/perf_recipes/deepseek/gb200/deepseek_v3.py @@ -83,8 +83,6 @@ def deepseek_v3_pretrain_256gpu_gb200_bf16_config() -> ConfigContainer: # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, - # Keep DeepSeek kernel selection aligned with the measured baseline. - "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, } return cfg @@ -194,8 +192,6 @@ def deepseek_v3_pretrain_256gpu_gb200_fp8mx_config() -> ConfigContainer: # Use cuDNN LayerNorm for this measured baseline. "NVTE_NORM_BWD_USE_CUDNN": 1, "NVTE_NORM_FWD_USE_CUDNN": 1, - # Keep DeepSeek kernel selection aligned with the measured baseline. - "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, } return cfg @@ -233,8 +229,6 @@ def deepseek_v3_pretrain_256gpu_gb200_nvfp4_config() -> ConfigContainer: # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, - # Keep DeepSeek kernel selection aligned with the measured baseline. - "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, # NVFP4 fast-math path. "NVTE_USE_FAST_MATH": 1, "NVTE_CUTEDSL_FUSED_GROUPED_MLP": 1, @@ -269,8 +263,6 @@ def deepseek_v3_pretrain_256gpu_gb200_fp8mx_large_scale_config() -> ConfigContai # Use cuDNN LayerNorm for this measured baseline. "NVTE_NORM_BWD_USE_CUDNN": 1, "NVTE_NORM_FWD_USE_CUDNN": 1, - # Keep DeepSeek kernel selection aligned with the measured baseline. - "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, } return cfg diff --git a/src/megatron/bridge/perf_recipes/deepseek/gb200/deepseek_v4.py b/src/megatron/bridge/perf_recipes/deepseek/gb200/deepseek_v4.py index 2854b7e61d..9a6eb10dca 100644 --- a/src/megatron/bridge/perf_recipes/deepseek/gb200/deepseek_v4.py +++ b/src/megatron/bridge/perf_recipes/deepseek/gb200/deepseek_v4.py @@ -135,6 +135,5 @@ def deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_config() -> ConfigContainer: "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, "NVTE_NORM_BWD_USE_CUDNN": 1, "NVTE_NORM_FWD_USE_CUDNN": 1, - "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, } return cfg diff --git a/src/megatron/bridge/perf_recipes/deepseek/vrnvl72/deepseek_v3.py b/src/megatron/bridge/perf_recipes/deepseek/vrnvl72/deepseek_v3.py index cea8d2d4f4..0c782dd803 100644 --- a/src/megatron/bridge/perf_recipes/deepseek/vrnvl72/deepseek_v3.py +++ b/src/megatron/bridge/perf_recipes/deepseek/vrnvl72/deepseek_v3.py @@ -78,8 +78,6 @@ def deepseek_v3_pretrain_128gpu_vrnvl72_bf16_config() -> ConfigContainer: # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, - # Keep DeepSeek kernel selection aligned with the measured baseline. - "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, } return cfg @@ -186,8 +184,6 @@ def deepseek_v3_pretrain_128gpu_vrnvl72_fp8mx_config() -> ConfigContainer: # Use cuDNN LayerNorm for this measured baseline. "NVTE_NORM_BWD_USE_CUDNN": 1, "NVTE_NORM_FWD_USE_CUDNN": 1, - # Keep DeepSeek kernel selection aligned with the measured baseline. - "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, } return cfg @@ -238,8 +234,6 @@ def deepseek_v3_pretrain_128gpu_vrnvl72_nvfp4_config() -> ConfigContainer: # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, - # Keep DeepSeek kernel selection aligned with the measured baseline. - "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, # NVFP4 fast-math path. "NVTE_USE_FAST_MATH": 1, } @@ -271,8 +265,6 @@ def deepseek_v3_pretrain_256gpu_vrnvl72_bf16_config() -> ConfigContainer: # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, - # Keep DeepSeek kernel selection aligned with the measured baseline. - "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, } return cfg @@ -333,8 +325,6 @@ def deepseek_v3_pretrain_256gpu_vrnvl72_fp8mx_config() -> ConfigContainer: # Use cuDNN LayerNorm for this measured baseline. "NVTE_NORM_BWD_USE_CUDNN": 1, "NVTE_NORM_FWD_USE_CUDNN": 1, - # Keep DeepSeek kernel selection aligned with the measured baseline. - "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, } return cfg @@ -373,7 +363,6 @@ def deepseek_v3_pretrain_64gpu_vrnvl72_fp8mx_proxy_config() -> ConfigContainer: "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, "NVTE_NORM_BWD_USE_CUDNN": 1, "NVTE_NORM_FWD_USE_CUDNN": 1, - "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, } return cfg @@ -402,8 +391,6 @@ def deepseek_v3_pretrain_256gpu_vrnvl72_nvfp4_config() -> ConfigContainer: # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, - # Keep DeepSeek kernel selection aligned with the measured baseline. - "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, # NVFP4 fast-math path. "NVTE_USE_FAST_MATH": 1, "NVTE_CUTEDSL_FUSED_GROUPED_MLP": 1, @@ -442,7 +429,6 @@ def deepseek_v3_pretrain_64gpu_vrnvl72_nvfp4_proxy_config() -> ConfigContainer: "NCCL_EP_HT_EM_PULL_PUSH": 1, "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, - "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, "NVTE_USE_FAST_MATH": 1, "NVTE_CUTEDSL_FUSED_GROUPED_MLP": 1, "NVTE_DPA_FP8_RECIPE": "MXFP8BlockScaling", @@ -513,8 +499,6 @@ def deepseek_v3_pretrain_4gpu_vrnvl72_fp8mx_proxy_config() -> ConfigContainer: # Use cuDNN LayerNorm in both 4-GPU proxies. "NVTE_NORM_BWD_USE_CUDNN": 1, "NVTE_NORM_FWD_USE_CUDNN": 1, - # Keep DeepSeek kernel selection aligned with the 256-GPU parent recipes. - "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, } return cfg @@ -554,8 +538,6 @@ def deepseek_v3_pretrain_4gpu_vrnvl72_nvfp4_proxy_config() -> ConfigContainer: # Use cuDNN LayerNorm in both 4-GPU proxies. "NVTE_NORM_BWD_USE_CUDNN": 1, "NVTE_NORM_FWD_USE_CUDNN": 1, - # Keep DeepSeek kernel selection aligned with the 256-GPU parent recipes. - "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, # NVFP4 fast-math path. "NVTE_USE_FAST_MATH": 1, # Run attention in MXFP8 E4M3, as the MXFP8 proxy does natively. diff --git a/src/megatron/bridge/perf_recipes/deepseek/vrnvl72/deepseek_v4.py b/src/megatron/bridge/perf_recipes/deepseek/vrnvl72/deepseek_v4.py index 25a3ab0dbe..254e0923cd 100644 --- a/src/megatron/bridge/perf_recipes/deepseek/vrnvl72/deepseek_v4.py +++ b/src/megatron/bridge/perf_recipes/deepseek/vrnvl72/deepseek_v4.py @@ -47,6 +47,5 @@ def deepseek_v4_flash_pretrain_128gpu_vrnvl72_fp8mx_config() -> ConfigContainer: "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, "NVTE_NORM_BWD_USE_CUDNN": 1, "NVTE_NORM_FWD_USE_CUDNN": 1, - "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, } return cfg diff --git a/tests/functional_tests/test_groups/recipes/test_perf_config_integration.py b/tests/functional_tests/test_groups/recipes/test_perf_config_integration.py index d0d87dafa6..6001e8b1da 100644 --- a/tests/functional_tests/test_groups/recipes/test_perf_config_integration.py +++ b/tests/functional_tests/test_groups/recipes/test_perf_config_integration.py @@ -189,11 +189,7 @@ def test_workload_base_config_derives_from_flat_recipe(self): "megatron.bridge.perf_recipes.deepseek.gb200.deepseek_v3:" "deepseek_v3_pretrain_256gpu_gb200_bf16_config" ), - { - "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, - "NCCL_EP_HT_EM_PULL_PUSH": 1, - "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, - }, + {"NVTE_FWD_LAYERNORM_SM_MARGIN": 20, "NCCL_EP_HT_EM_PULL_PUSH": 1}, ), ( "gpt_oss", diff --git a/tests/unit_tests/recipes/test_deepseek_v4_flash_perf_recipes.py b/tests/unit_tests/recipes/test_deepseek_v4_flash_perf_recipes.py index 2353950a51..94cfea3c21 100644 --- a/tests/unit_tests/recipes/test_deepseek_v4_flash_perf_recipes.py +++ b/tests/unit_tests/recipes/test_deepseek_v4_flash_perf_recipes.py @@ -119,6 +119,8 @@ def test_deepseek_v4_flash_128gpu_gb200_fp8mx_config() -> None: assert cfg.env_vars["NVTE_CUTEDSL_FUSED_GROUPED_MLP"] == 1 assert cfg.env_vars["NCCL_EP_HT_EM_PULL_PUSH"] == 1 assert cfg.env_vars.keys().isdisjoint(HYBRID_EP_ENV_NAMES) + # TE's op-fuser grouped MLP raises when deterministic kernels are requested. + assert "NVTE_ALLOW_NONDETERMINISTIC_ALGO" not in cfg.env_vars def test_deepseek_v4_flash_128gpu_gb300_fp8mx_config() -> None: @@ -147,7 +149,7 @@ def test_deepseek_v4_flash_128gpu_vrnvl72_fp8mx_disables_clamp_for_cutedsl() -> assert vrnvl72_cfg.model.use_transformer_engine_op_fuser is True assert vrnvl72_cfg.model.moe_use_grouped_tensor is True assert vrnvl72_cfg.env_vars["NVTE_CUTEDSL_FUSED_GROUPED_MLP"] == 1 - assert vrnvl72_cfg.env_vars["NVTE_ALLOW_NONDETERMINISTIC_ALGO"] == 0 + assert "NVTE_ALLOW_NONDETERMINISTIC_ALGO" not in vrnvl72_cfg.env_vars assert is_full_iteration_cuda_graph(vrnvl72_cfg.model) # VRNVL72 dispatches through NCCL EP; the GB300 recipe stays on the shared HybridEP base. assert vrnvl72_cfg.model.moe_flex_dispatcher_backend == "ncclep" @@ -160,8 +162,9 @@ def test_deepseek_v4_flash_128gpu_vrnvl72_fp8mx_disables_clamp_for_cutedsl() -> gb300_cfg.model.moe_use_grouped_tensor = True gb300_cfg.model.moe_flex_dispatcher_backend = "ncclep" assert vars(vrnvl72_cfg.model) == vars(gb300_cfg.model) - # The environments differ only in the dispatcher settings. - expected_env = {name: value for name, value in gb300_cfg.env_vars.items() if name not in HYBRID_EP_ENV_NAMES} + # The environments differ only in the dispatcher settings and the GB300 determinism pin. + gb300_only_env_names = HYBRID_EP_ENV_NAMES | {"NVTE_ALLOW_NONDETERMINISTIC_ALGO"} + expected_env = {name: value for name, value in gb300_cfg.env_vars.items() if name not in gb300_only_env_names} assert vrnvl72_cfg.env_vars == {**expected_env, "NCCL_EP_HT_EM_PULL_PUSH": 1} for field in fields(vrnvl72_cfg): if field.name not in ("model", "env_vars"): diff --git a/tests/unit_tests/recipes/test_perf_recipe_environment.py b/tests/unit_tests/recipes/test_perf_recipe_environment.py index 80e88feefa..35c2899c83 100644 --- a/tests/unit_tests/recipes/test_perf_recipe_environment.py +++ b/tests/unit_tests/recipes/test_perf_recipe_environment.py @@ -227,10 +227,13 @@ def test_explicit_environment_invariants_across_all_flat_recipes(): if function_name.endswith("_deterministic_config"): # Bit-exact recipes pin deterministic TE kernels. assert environment["NVTE_ALLOW_NONDETERMINISTIC_ALGO"] == "0" - elif path.parent.name == "gb300" and path.stem == "deepseek_v3": - # The DeepSeek V3 GB300 recipes run the TE op-fuser grouped MLP, whose quantized scale - # gradient is accumulated with atomics; TE raises when bit-exact execution is requested, - # so these recipes must not pin NVTE_ALLOW_NONDETERMINISTIC_ALGO=0. + elif ( + path.parent.name == "gb300" and path.stem == "deepseek_v3" + ) or "NCCL_EP_HT_EM_PULL_PUSH" in environment: + # The DeepSeek V3 GB300 recipes and the DeepSeek NCCL EP recipes on other GPUs do not pin + # NVTE_ALLOW_NONDETERMINISTIC_ALGO=0: the TE op-fuser grouped MLP that many of them run + # accumulates its quantized scale gradient with atomics, and TE raises when bit-exact + # execution is requested. assert "NVTE_ALLOW_NONDETERMINISTIC_ALGO" not in environment else: assert environment["NVTE_ALLOW_NONDETERMINISTIC_ALGO"] == 0 @@ -257,10 +260,7 @@ def test_explicit_environment_invariants_across_all_flat_recipes(): ( "deepseek/gb200/deepseek_v3.py", "deepseek_v3_pretrain_256gpu_gb200_bf16_config", - { - "NCCL_EP_HT_EM_PULL_PUSH": 1, - "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, - }, + {"NCCL_EP_HT_EM_PULL_PUSH": 1}, ), ( "llama/h100/llama3.py", From 43801ef4251e73dad4476d9402d1db026c9aa6cc Mon Sep 17 00:00:00 2001 From: Malay Nagda Date: Fri, 16 Oct 2026 13:36:18 +0530 Subject: [PATCH 9/9] Revert "perf(recipe): drop the deterministic TE pin from DeepSeek NCCL EP perf recipes" This reverts commit 5f8300eaafc380c508b5a96668188b972ffb4231. The DeepSeek NCCL EP perf recipes keep NVTE_ALLOW_NONDETERMINISTIC_ALGO=0, as on main. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Malay Nagda --- .../perf_recipes/deepseek/gb200/deepseek_v3.py | 8 ++++++++ .../perf_recipes/deepseek/gb200/deepseek_v4.py | 1 + .../deepseek/vrnvl72/deepseek_v3.py | 18 ++++++++++++++++++ .../deepseek/vrnvl72/deepseek_v4.py | 1 + .../recipes/test_perf_config_integration.py | 6 +++++- .../test_deepseek_v4_flash_perf_recipes.py | 9 +++------ .../recipes/test_perf_recipe_environment.py | 16 ++++++++-------- 7 files changed, 44 insertions(+), 15 deletions(-) diff --git a/src/megatron/bridge/perf_recipes/deepseek/gb200/deepseek_v3.py b/src/megatron/bridge/perf_recipes/deepseek/gb200/deepseek_v3.py index 2710624f51..cc17b9bbd5 100644 --- a/src/megatron/bridge/perf_recipes/deepseek/gb200/deepseek_v3.py +++ b/src/megatron/bridge/perf_recipes/deepseek/gb200/deepseek_v3.py @@ -83,6 +83,8 @@ def deepseek_v3_pretrain_256gpu_gb200_bf16_config() -> ConfigContainer: # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, + # Keep DeepSeek kernel selection aligned with the measured baseline. + "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, } return cfg @@ -192,6 +194,8 @@ def deepseek_v3_pretrain_256gpu_gb200_fp8mx_config() -> ConfigContainer: # Use cuDNN LayerNorm for this measured baseline. "NVTE_NORM_BWD_USE_CUDNN": 1, "NVTE_NORM_FWD_USE_CUDNN": 1, + # Keep DeepSeek kernel selection aligned with the measured baseline. + "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, } return cfg @@ -229,6 +233,8 @@ def deepseek_v3_pretrain_256gpu_gb200_nvfp4_config() -> ConfigContainer: # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, + # Keep DeepSeek kernel selection aligned with the measured baseline. + "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, # NVFP4 fast-math path. "NVTE_USE_FAST_MATH": 1, "NVTE_CUTEDSL_FUSED_GROUPED_MLP": 1, @@ -263,6 +269,8 @@ def deepseek_v3_pretrain_256gpu_gb200_fp8mx_large_scale_config() -> ConfigContai # Use cuDNN LayerNorm for this measured baseline. "NVTE_NORM_BWD_USE_CUDNN": 1, "NVTE_NORM_FWD_USE_CUDNN": 1, + # Keep DeepSeek kernel selection aligned with the measured baseline. + "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, } return cfg diff --git a/src/megatron/bridge/perf_recipes/deepseek/gb200/deepseek_v4.py b/src/megatron/bridge/perf_recipes/deepseek/gb200/deepseek_v4.py index 9a6eb10dca..2854b7e61d 100644 --- a/src/megatron/bridge/perf_recipes/deepseek/gb200/deepseek_v4.py +++ b/src/megatron/bridge/perf_recipes/deepseek/gb200/deepseek_v4.py @@ -135,5 +135,6 @@ def deepseek_v4_flash_pretrain_128gpu_gb200_fp8mx_config() -> ConfigContainer: "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, "NVTE_NORM_BWD_USE_CUDNN": 1, "NVTE_NORM_FWD_USE_CUDNN": 1, + "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, } return cfg diff --git a/src/megatron/bridge/perf_recipes/deepseek/vrnvl72/deepseek_v3.py b/src/megatron/bridge/perf_recipes/deepseek/vrnvl72/deepseek_v3.py index 0c782dd803..cea8d2d4f4 100644 --- a/src/megatron/bridge/perf_recipes/deepseek/vrnvl72/deepseek_v3.py +++ b/src/megatron/bridge/perf_recipes/deepseek/vrnvl72/deepseek_v3.py @@ -78,6 +78,8 @@ def deepseek_v3_pretrain_128gpu_vrnvl72_bf16_config() -> ConfigContainer: # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, + # Keep DeepSeek kernel selection aligned with the measured baseline. + "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, } return cfg @@ -184,6 +186,8 @@ def deepseek_v3_pretrain_128gpu_vrnvl72_fp8mx_config() -> ConfigContainer: # Use cuDNN LayerNorm for this measured baseline. "NVTE_NORM_BWD_USE_CUDNN": 1, "NVTE_NORM_FWD_USE_CUDNN": 1, + # Keep DeepSeek kernel selection aligned with the measured baseline. + "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, } return cfg @@ -234,6 +238,8 @@ def deepseek_v3_pretrain_128gpu_vrnvl72_nvfp4_config() -> ConfigContainer: # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, + # Keep DeepSeek kernel selection aligned with the measured baseline. + "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, # NVFP4 fast-math path. "NVTE_USE_FAST_MATH": 1, } @@ -265,6 +271,8 @@ def deepseek_v3_pretrain_256gpu_vrnvl72_bf16_config() -> ConfigContainer: # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, + # Keep DeepSeek kernel selection aligned with the measured baseline. + "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, } return cfg @@ -325,6 +333,8 @@ def deepseek_v3_pretrain_256gpu_vrnvl72_fp8mx_config() -> ConfigContainer: # Use cuDNN LayerNorm for this measured baseline. "NVTE_NORM_BWD_USE_CUDNN": 1, "NVTE_NORM_FWD_USE_CUDNN": 1, + # Keep DeepSeek kernel selection aligned with the measured baseline. + "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, } return cfg @@ -363,6 +373,7 @@ def deepseek_v3_pretrain_64gpu_vrnvl72_fp8mx_proxy_config() -> ConfigContainer: "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, "NVTE_NORM_BWD_USE_CUDNN": 1, "NVTE_NORM_FWD_USE_CUDNN": 1, + "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, } return cfg @@ -391,6 +402,8 @@ def deepseek_v3_pretrain_256gpu_vrnvl72_nvfp4_config() -> ConfigContainer: # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, + # Keep DeepSeek kernel selection aligned with the measured baseline. + "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, # NVFP4 fast-math path. "NVTE_USE_FAST_MATH": 1, "NVTE_CUTEDSL_FUSED_GROUPED_MLP": 1, @@ -429,6 +442,7 @@ def deepseek_v3_pretrain_64gpu_vrnvl72_nvfp4_proxy_config() -> ConfigContainer: "NCCL_EP_HT_EM_PULL_PUSH": 1, "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, + "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, "NVTE_USE_FAST_MATH": 1, "NVTE_CUTEDSL_FUSED_GROUPED_MLP": 1, "NVTE_DPA_FP8_RECIPE": "MXFP8BlockScaling", @@ -499,6 +513,8 @@ def deepseek_v3_pretrain_4gpu_vrnvl72_fp8mx_proxy_config() -> ConfigContainer: # Use cuDNN LayerNorm in both 4-GPU proxies. "NVTE_NORM_BWD_USE_CUDNN": 1, "NVTE_NORM_FWD_USE_CUDNN": 1, + # Keep DeepSeek kernel selection aligned with the 256-GPU parent recipes. + "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, } return cfg @@ -538,6 +554,8 @@ def deepseek_v3_pretrain_4gpu_vrnvl72_nvfp4_proxy_config() -> ConfigContainer: # Use cuDNN LayerNorm in both 4-GPU proxies. "NVTE_NORM_BWD_USE_CUDNN": 1, "NVTE_NORM_FWD_USE_CUDNN": 1, + # Keep DeepSeek kernel selection aligned with the 256-GPU parent recipes. + "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, # NVFP4 fast-math path. "NVTE_USE_FAST_MATH": 1, # Run attention in MXFP8 E4M3, as the MXFP8 proxy does natively. diff --git a/src/megatron/bridge/perf_recipes/deepseek/vrnvl72/deepseek_v4.py b/src/megatron/bridge/perf_recipes/deepseek/vrnvl72/deepseek_v4.py index 254e0923cd..25a3ab0dbe 100644 --- a/src/megatron/bridge/perf_recipes/deepseek/vrnvl72/deepseek_v4.py +++ b/src/megatron/bridge/perf_recipes/deepseek/vrnvl72/deepseek_v4.py @@ -47,5 +47,6 @@ def deepseek_v4_flash_pretrain_128gpu_vrnvl72_fp8mx_config() -> ConfigContainer: "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, "NVTE_NORM_BWD_USE_CUDNN": 1, "NVTE_NORM_FWD_USE_CUDNN": 1, + "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, } return cfg diff --git a/tests/functional_tests/test_groups/recipes/test_perf_config_integration.py b/tests/functional_tests/test_groups/recipes/test_perf_config_integration.py index 6001e8b1da..d0d87dafa6 100644 --- a/tests/functional_tests/test_groups/recipes/test_perf_config_integration.py +++ b/tests/functional_tests/test_groups/recipes/test_perf_config_integration.py @@ -189,7 +189,11 @@ def test_workload_base_config_derives_from_flat_recipe(self): "megatron.bridge.perf_recipes.deepseek.gb200.deepseek_v3:" "deepseek_v3_pretrain_256gpu_gb200_bf16_config" ), - {"NVTE_FWD_LAYERNORM_SM_MARGIN": 20, "NCCL_EP_HT_EM_PULL_PUSH": 1}, + { + "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, + "NCCL_EP_HT_EM_PULL_PUSH": 1, + "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, + }, ), ( "gpt_oss", diff --git a/tests/unit_tests/recipes/test_deepseek_v4_flash_perf_recipes.py b/tests/unit_tests/recipes/test_deepseek_v4_flash_perf_recipes.py index 94cfea3c21..2353950a51 100644 --- a/tests/unit_tests/recipes/test_deepseek_v4_flash_perf_recipes.py +++ b/tests/unit_tests/recipes/test_deepseek_v4_flash_perf_recipes.py @@ -119,8 +119,6 @@ def test_deepseek_v4_flash_128gpu_gb200_fp8mx_config() -> None: assert cfg.env_vars["NVTE_CUTEDSL_FUSED_GROUPED_MLP"] == 1 assert cfg.env_vars["NCCL_EP_HT_EM_PULL_PUSH"] == 1 assert cfg.env_vars.keys().isdisjoint(HYBRID_EP_ENV_NAMES) - # TE's op-fuser grouped MLP raises when deterministic kernels are requested. - assert "NVTE_ALLOW_NONDETERMINISTIC_ALGO" not in cfg.env_vars def test_deepseek_v4_flash_128gpu_gb300_fp8mx_config() -> None: @@ -149,7 +147,7 @@ def test_deepseek_v4_flash_128gpu_vrnvl72_fp8mx_disables_clamp_for_cutedsl() -> assert vrnvl72_cfg.model.use_transformer_engine_op_fuser is True assert vrnvl72_cfg.model.moe_use_grouped_tensor is True assert vrnvl72_cfg.env_vars["NVTE_CUTEDSL_FUSED_GROUPED_MLP"] == 1 - assert "NVTE_ALLOW_NONDETERMINISTIC_ALGO" not in vrnvl72_cfg.env_vars + assert vrnvl72_cfg.env_vars["NVTE_ALLOW_NONDETERMINISTIC_ALGO"] == 0 assert is_full_iteration_cuda_graph(vrnvl72_cfg.model) # VRNVL72 dispatches through NCCL EP; the GB300 recipe stays on the shared HybridEP base. assert vrnvl72_cfg.model.moe_flex_dispatcher_backend == "ncclep" @@ -162,9 +160,8 @@ def test_deepseek_v4_flash_128gpu_vrnvl72_fp8mx_disables_clamp_for_cutedsl() -> gb300_cfg.model.moe_use_grouped_tensor = True gb300_cfg.model.moe_flex_dispatcher_backend = "ncclep" assert vars(vrnvl72_cfg.model) == vars(gb300_cfg.model) - # The environments differ only in the dispatcher settings and the GB300 determinism pin. - gb300_only_env_names = HYBRID_EP_ENV_NAMES | {"NVTE_ALLOW_NONDETERMINISTIC_ALGO"} - expected_env = {name: value for name, value in gb300_cfg.env_vars.items() if name not in gb300_only_env_names} + # The environments differ only in the dispatcher settings. + expected_env = {name: value for name, value in gb300_cfg.env_vars.items() if name not in HYBRID_EP_ENV_NAMES} assert vrnvl72_cfg.env_vars == {**expected_env, "NCCL_EP_HT_EM_PULL_PUSH": 1} for field in fields(vrnvl72_cfg): if field.name not in ("model", "env_vars"): diff --git a/tests/unit_tests/recipes/test_perf_recipe_environment.py b/tests/unit_tests/recipes/test_perf_recipe_environment.py index 35c2899c83..80e88feefa 100644 --- a/tests/unit_tests/recipes/test_perf_recipe_environment.py +++ b/tests/unit_tests/recipes/test_perf_recipe_environment.py @@ -227,13 +227,10 @@ def test_explicit_environment_invariants_across_all_flat_recipes(): if function_name.endswith("_deterministic_config"): # Bit-exact recipes pin deterministic TE kernels. assert environment["NVTE_ALLOW_NONDETERMINISTIC_ALGO"] == "0" - elif ( - path.parent.name == "gb300" and path.stem == "deepseek_v3" - ) or "NCCL_EP_HT_EM_PULL_PUSH" in environment: - # The DeepSeek V3 GB300 recipes and the DeepSeek NCCL EP recipes on other GPUs do not pin - # NVTE_ALLOW_NONDETERMINISTIC_ALGO=0: the TE op-fuser grouped MLP that many of them run - # accumulates its quantized scale gradient with atomics, and TE raises when bit-exact - # execution is requested. + elif path.parent.name == "gb300" and path.stem == "deepseek_v3": + # The DeepSeek V3 GB300 recipes run the TE op-fuser grouped MLP, whose quantized scale + # gradient is accumulated with atomics; TE raises when bit-exact execution is requested, + # so these recipes must not pin NVTE_ALLOW_NONDETERMINISTIC_ALGO=0. assert "NVTE_ALLOW_NONDETERMINISTIC_ALGO" not in environment else: assert environment["NVTE_ALLOW_NONDETERMINISTIC_ALGO"] == 0 @@ -260,7 +257,10 @@ def test_explicit_environment_invariants_across_all_flat_recipes(): ( "deepseek/gb200/deepseek_v3.py", "deepseek_v3_pretrain_256gpu_gb200_bf16_config", - {"NCCL_EP_HT_EM_PULL_PUSH": 1}, + { + "NCCL_EP_HT_EM_PULL_PUSH": 1, + "NVTE_ALLOW_NONDETERMINISTIC_ALGO": 0, + }, ), ( "llama/h100/llama3.py",