diff --git a/src/megatron/bridge/perf_recipes/_common.py b/src/megatron/bridge/perf_recipes/_common.py index 97b4b715482..32f48f8aeb2 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, 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 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 9ecdc991608..cc17b9bbd5e 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/deepseek/gb200/deepseek_v4.py b/src/megatron/bridge/perf_recipes/deepseek/gb200/deepseek_v4.py index 126ca64ad2e..2854b7e61da 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 0cba00e8cbf..82514d26859 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/deepseek/vrnvl72/deepseek_v3.py b/src/megatron/bridge/perf_recipes/deepseek/vrnvl72/deepseek_v3.py index 78421c80307..48c3b6156a1 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, @@ -134,7 +136,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 +159,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 +176,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 +193,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 +218,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 +233,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 +247,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 +266,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, @@ -301,8 +306,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 +323,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, @@ -361,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, @@ -378,8 +382,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 = { @@ -392,11 +400,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, @@ -436,10 +441,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, @@ -484,7 +487,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. """ @@ -503,11 +506,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, @@ -528,7 +528,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. """ @@ -547,11 +547,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 1ed0d3bf3ab..25a3ab0dbe8 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 dbc5450e9cf..26ee5ce191c 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, @@ -159,8 +160,12 @@ def gpt_oss_20b_pretrain_64gpu_vrnvl72_nvfp4_config() -> ConfigContainer: def gpt_oss_120b_pretrain_64gpu_vrnvl72_bf16_config() -> ConfigContainer: - """GPT-OSS 120B pretrain: 64× VRNVL72, BF16 (alias of GB300).""" + """GPT-OSS 120B pretrain: 64× VRNVL72, BF16, NCCL EP (alias of GB300).""" cfg = gpt_oss_120b_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 # Keep process settings next to the recipe so users can see the exact benchmark environment. cfg.env_vars = { **COMMON_PERF_ENV_VARS, @@ -172,6 +177,8 @@ def gpt_oss_120b_pretrain_64gpu_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, # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, @@ -180,8 +187,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 +204,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/nemotronh/gb200/nemotronh.py b/src/megatron/bridge/perf_recipes/nemotronh/gb200/nemotronh.py index 9b381aceb2c..9fdeacb1858 100644 --- a/src/megatron/bridge/perf_recipes/nemotronh/gb200/nemotronh.py +++ b/src/megatron/bridge/perf_recipes/nemotronh/gb200/nemotronh.py @@ -43,7 +43,7 @@ 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") @@ -65,6 +65,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, @@ -76,11 +80,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, @@ -89,7 +90,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") @@ -111,6 +112,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, @@ -122,11 +127,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, @@ -137,7 +139,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() @@ -160,6 +162,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, @@ -171,11 +177,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, @@ -603,7 +606,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.""" cfg = nemotron_3_nano_pretrain_8gpu_gb200_bf16_config() cfg.model.recompute_modules = [] cfg.model.mtp_num_layers = 2 @@ -616,6 +619,10 @@ def nemotron_3_5_lightning_pretrain_8gpu_gb200_bf16_config() -> ConfigContainer: ) 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, @@ -623,10 +630,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. + "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, @@ -652,11 +657,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, @@ -664,10 +673,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/nemotronh/vrnvl72/nemotronh.py b/src/megatron/bridge/perf_recipes/nemotronh/vrnvl72/nemotronh.py index 473373a02c1..e8f42a3aea5 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, @@ -127,8 +128,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 +145,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 +155,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 +175,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 +187,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 +206,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 +298,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 +319,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 +331,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 +347,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/gb200/qwen3_moe.py b/src/megatron/bridge/perf_recipes/qwen/gb200/qwen3_moe.py index 8f5bdc3f76c..4bfd3890983 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, @@ -339,7 +340,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 +371,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 +390,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 +454,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 +481,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 +496,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 +508,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 +522,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 +589,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 +611,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, 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 ac548cb2fe3..8aae311f52c 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 6afe98a3013..29d442867fa 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, @@ -87,16 +89,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 +118,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 +147,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, @@ -243,8 +241,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 +261,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, @@ -301,12 +300,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 base (cuDNN GDN, per-layer CUDA graphs, forced load balancing), not the tuned GB300 recipe, and adds the CuTe DSL fused grouped MLP and the settings of the measured VRNVL72 run. """ cfg = _build_qwen35_vl_397b_a17b_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 # MBS 1, as in the GB300 base; the vision cap below is set for it. cfg.train.micro_batch_size = 1 @@ -329,10 +332,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 +351,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/functional_tests/test_groups/recipes/test_perf_config_integration.py b/tests/functional_tests/test_groups/recipes/test_perf_config_integration.py index e84bdbd1f90..070b608ba38 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,7 +220,7 @@ 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", 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 cfbb44e41f5..aa115da2a2a 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 bbf03377e2b..f3f885fa61e 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_deepseek_v3_proxy_recipes.py b/tests/unit_tests/recipes/test_deepseek_v3_proxy_recipes.py index 2814f8dedb7..b42854dfc3b 100644 --- a/tests/unit_tests/recipes/test_deepseek_v3_proxy_recipes.py +++ b/tests/unit_tests/recipes/test_deepseek_v3_proxy_recipes.py @@ -126,7 +126,8 @@ def test_proxy_preserves_parent_except_depth_layout_and_kernel_policy(precision: 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 3d14c893025..2353950a517 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_vrnvl72_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) @@ -146,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 new file mode 100644 index 00000000000..1b0dd3806df --- /dev/null +++ b/tests/unit_tests/recipes/test_gb200_ncclep_perf_recipes.py @@ -0,0 +1,104 @@ +# 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. +"""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 + +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 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"), + ("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_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"), + ("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 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"), + ("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_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_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_gb300_ncclep_perf_recipes.py b/tests/unit_tests/recipes/test_gb300_ncclep_perf_recipes.py index 08ec0965f75..9ce1c7bd0b1 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,60 @@ ("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 base or recipe (HybridEP, or alltoall for GPT-OSS 120B BF16), 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_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_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_120b_pretrain_64gpu_vrnvl72_bf16_config"), ("gpt_oss.vrnvl72.gpt_oss", "gpt_oss_120b_pretrain_64gpu_vrnvl72_fp8mx_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_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_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_fp8mx_config"), +) + +# 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. +_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"), + ("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"), ) @@ -92,7 +117,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 +165,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 0063fc6f9cd..6b8a281893f 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 - # 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 2324742bf59..01b16ac9a3c 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,7 @@ 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, nemotron_3_5_lightning_pretrain_8gpu_b200_fp8mx_config, @@ -99,10 +99,15 @@ 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). +# 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, @@ -348,16 +353,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 @@ -532,6 +543,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 @@ -586,8 +600,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" @@ -639,16 +660,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 @@ -658,8 +680,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/recipes/test_nemotronh_recipes.py b/tests/unit_tests/recipes/test_nemotronh_recipes.py index 4ee96697e53..4b8399dc058 100644 --- a/tests/unit_tests/recipes/test_nemotronh_recipes.py +++ b/tests/unit_tests/recipes/test_nemotronh_recipes.py @@ -32,6 +32,7 @@ from megatron.bridge.models.conversion.auto_bridge import AutoBridge from megatron.bridge.models.hybrid import BridgeHybridModelConfig +from megatron.bridge.perf_recipes.environment import HYBRID_EP_ENV_NAMES from megatron.bridge.training import checkpointing from megatron.bridge.training.utils.omegaconf_utils import ( OverridesError, @@ -361,7 +362,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", @@ -374,7 +374,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 efbae4040d8..ba1699b6055 100644 --- a/tests/unit_tests/recipes/test_perf_recipe_environment.py +++ b/tests/unit_tests/recipes/test_perf_recipe_environment.py @@ -264,10 +264,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_r050_perf_recipe_ports.py b/tests/unit_tests/recipes/test_r050_perf_recipe_ports.py index b6f6519725a..3ba80fa7092 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 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 8b85465a86c..1d8da8e71e1 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 @@ -368,19 +368,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_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 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 BF16 and MXFP8 recipes run NCCL EP and carry no HybridEP env. """ patch_recipe_construction_dependencies(monkeypatch) @@ -466,11 +465,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"} @@ -500,6 +499,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_flex_dispatcher_backend == "ncclep" + assert config.model.moe_use_grouped_tensor is True # Built on the GB300 MXFP8 base, not the tuned GB300 recipe: the Adam moments stay in FP32. assert config.optimizer.use_precision_aware_optimizer is False assert config.model.moe_permute_fusion_into_hybridep is False @@ -510,9 +511,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