Skip to content

perf(recipe): default GB200 MoE perf recipes to NCCL EP and grouped tensors #6373

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
malay-nagda wants to merge 5 commits into
base: main
Choose a base branch
from
+285 −116
Draft

perf(recipe): default GB200 MoE perf recipes to NCCL EP and grouped tensors #6373

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
823b90a
perf(recipe): default GB200 MoE perf recipes to NCCL EP and grouped t…
malay-nagda Oct 9, 2026
268f912
perf(recipe): keep Kimi K2, Nemotron 3 Nano and Qwen3 30B 32-GPU GB20…
malay-nagda Oct 9, 2026
27c0f0f
perf(recipe): use NCCL EP for GPT-OSS 120B BF16 and DeepSeek V4 Flash…
malay-nagda Oct 9, 2026
8d53019
Merge branch 'main' into malay/gb200-ncclep
malay-nagda Oct 14, 2026
6a880fb
perf(recipe): keep GPT-OSS 120B and Qwen3 30B-A3B GB200 perf recipes …
malay-nagda Oct 14, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion src/megatron/bridge/perf_recipes/_common.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -96,7 +96,7 @@ def _benchmark_common(


def _enable_ncclep(cfg: ConfigContainer) -> None:
"""Switch a flex-dispatcher recipe to the NCCL EP backend (GB300 MoE default).
"""Switch a flex-dispatcher recipe to the NCCL EP backend (GB200 and GB300 MoE default).

Call after ``_benchmark_common``. Only the dispatch stack changes; parallelism, precision,
recompute and the CUDA graph mode stay as the recipe set them. The calling recipe still declares
Expand Down
59 changes: 33 additions & 26 deletions src/megatron/bridge/perf_recipes/deepseek/gb200/deepseek_v3.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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)
Expand All @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand All @@ -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
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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.
Expand All @@ -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,
Expand Down
21 changes: 14 additions & 7 deletions src/megatron/bridge/perf_recipes/deepseek/gb200/deepseek_v4.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand Down
4 changes: 2 additions & 2 deletions src/megatron/bridge/perf_recipes/deepseek/gb300/deepseek_v4.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand All @@ -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
Expand Down
63 changes: 35 additions & 28 deletions src/megatron/bridge/perf_recipes/nemotronh/gb200/nemotronh.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,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")

Expand All @@ -64,6 +64,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,
Expand All @@ -75,11 +79,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,
Expand All @@ -88,7 +89,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")

Expand All @@ -110,6 +111,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,
Expand All @@ -121,11 +126,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,
Expand All @@ -136,7 +138,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()

Expand All @@ -159,6 +161,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,
Expand All @@ -170,11 +176,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,
Expand Down Expand Up @@ -550,7 +553,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
Expand All @@ -562,17 +565,19 @@ def nemotron_3_5_lightning_pretrain_8gpu_gb200_bf16_config() -> ConfigContainer:
cfg.model.hf_model_revision = _NEMOTRON_3_5_LIGHTNING_MODEL_REVISION
cfg.tokenizer.tokenizer_model = _NEMOTRON_3_5_LIGHTNING_MODEL_ID
cfg.tokenizer.hf_tokenizer_kwargs = {"revision": _NEMOTRON_3_5_LIGHTNING_MODEL_REVISION}
_enable_ncclep(cfg)
# Device-side expert token counts: the legacy grouped MLP path syncs tokens_per_expert to the
# host every layer, which serializes the CPU behind the GPU when dispatch is fast.
cfg.model.moe_use_grouped_tensor = True
cfg.env_vars = {
**COMMON_PERF_ENV_VARS,
"CUDA_DEVICE_MAX_CONNECTIONS": 32,
"NCCL_GRAPH_REGISTER": 0,
"PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True",
"TORCH_NCCL_AVOID_RECORD_STREAMS": 1,
"NCCL_NVLS_ENABLE": 0,
"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,
Expand All @@ -597,22 +602,24 @@ 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,
"NCCL_GRAPH_REGISTER": 0,
"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,
Expand Down
Loading
Loading