Skip to content

perf(nemotron): add Nemotron 3.5 Lightning GB300 recipes #5693

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

Merged
malay-nagda merged 4 commits into from
Aug 21, 2026
Merged

perf(nemotron): add Nemotron 3.5 Lightning GB300 recipes #5693

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
1d0195d
nemotron 3.5 gb300
malay-nagda Aug 19, 2026
317b3ea
test(nemotron): validate GB300 Lightning recipes
malay-nagda Aug 19, 2026
8bbae38
no moe_mlp_glu_interleave_size
malay-nagda Aug 21, 2026
1d0cc2c
test(nemotron): update CuTeDSL interleave expectation
malay-nagda Aug 21, 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
4 changes: 4 additions & 0 deletions scripts/performance/dump_perf_configs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -277,6 +277,10 @@
("nemotronh", "nemotron_3_nano", "pretrain", 16, "h100", "bf16"),
("nemotronh", "nemotron_3_nano", "pretrain", 16, "h100", "fp8_cs"),
# Nemotron 3.5 Lightning
("nemotronh", "nemotron_3_5_lightning", "pretrain", 8, "gb300", "bf16"),
("nemotronh", "nemotron_3_5_lightning", "pretrain", 8, "gb300", "fp8_mx"),
("nemotronh", "nemotron_3_5_lightning", "pretrain", 8, "gb300", "fp8_mx", "fsdp"),
("nemotronh", "nemotron_3_5_lightning", "pretrain", 8, "gb300", "nvfp4"),
("nemotronh", "nemotron_3_5_lightning", "pretrain", 8, "gb200", "bf16"),
("nemotronh", "nemotron_3_5_lightning", "pretrain", 8, "gb200", "fp8_mx"),
("nemotronh", "nemotron_3_5_lightning", "pretrain", 8, "gb200", "nvfp4"),
Expand Down
4 changes: 4 additions & 0 deletions src/megatron/bridge/perf_recipes/nemotronh/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,10 @@
nemotronh_56b_pretrain_64gpu_gb200_fp8cs_config,
)
from megatron.bridge.perf_recipes.nemotronh.gb300.nemotronh import (
nemotron_3_5_lightning_pretrain_8gpu_gb300_bf16_config,
nemotron_3_5_lightning_pretrain_8gpu_gb300_fp8mx_config,
nemotron_3_5_lightning_pretrain_8gpu_gb300_fp8mx_fsdp_config,
nemotron_3_5_lightning_pretrain_8gpu_gb300_nvfp4_config,
nemotron_3_nano_pretrain_8gpu_gb300_bf16_config,
nemotron_3_nano_pretrain_8gpu_gb300_fp8mx_config,
nemotron_3_nano_pretrain_8gpu_gb300_nvfp4_config,
Expand Down
154 changes: 154 additions & 0 deletions src/megatron/bridge/perf_recipes/nemotronh/gb300/nemotronh.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,11 @@
from megatron.bridge.utils.cuda_graph import set_cuda_graph_modules


# Public Nemotron 3.5 Lightning checkpoint used by the Lightning recipe API.
_NEMOTRON_3_5_LIGHTNING_MODEL_ID = "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16"
_NEMOTRON_3_5_LIGHTNING_MODEL_REVISION = "b3caaabed0263651a17dc1f2d4ce97e794f76c44" # pragma: allowlist secret


def nemotronh_56b_pretrain_64gpu_gb300_fp8cs_config() -> ConfigContainer:
"""NemotronH 56B pretrain: 64× GB300, FP8 current-scaling."""
cfg = nemotronh_56b_pretrain_config()
Expand Down Expand Up @@ -524,3 +529,152 @@ def nemotronh_56b_pretrain_256gpu_gb300_fp8cs_config() -> ConfigContainer:
"NVTE_FWD_LAYERNORM_SM_MARGIN": 20,
}
return cfg


def nemotron_3_5_lightning_pretrain_8gpu_gb300_bf16_config() -> ConfigContainer:
"""Nemotron 3.5 Lightning pretrain: 8× GB300, BF16."""
cfg = nemotron_3_nano_pretrain_8gpu_gb300_bf16_config()
cfg.model.mtp_num_layers = 2
cfg.model.mtp_hybrid_override_pattern = "*E"
cfg.model.mtp_use_repeated_layer = True
cfg.model.keep_mtp_spec_in_bf16 = True
cfg.model.mtp_loss_scaling_factor = 0.3
cfg.model.hf_model_id = _NEMOTRON_3_5_LIGHTNING_MODEL_ID
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}
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,
"NVTE_BWD_LAYERNORM_SM_MARGIN": 20,
"NVTE_FWD_LAYERNORM_SM_MARGIN": 20,
"NVTE_NORM_BWD_USE_CUDNN": 1,
"NVTE_NORM_FWD_USE_CUDNN": 1,
}
return cfg


def _build_nemotron_3_5_lightning_gb300_mxfp8() -> ConfigContainer:
cfg = _build_nemotron_3_nano_gb300_mxfp8()
cfg.model.moe_hybridep_num_sms = 16
cfg.model.mtp_num_layers = 2
cfg.model.mtp_hybrid_override_pattern = "*E"
cfg.model.mtp_use_repeated_layer = True
cfg.model.keep_mtp_spec_in_bf16 = True
cfg.model.mtp_loss_scaling_factor = 0.3
cfg.model.hf_model_id = _NEMOTRON_3_5_LIGHTNING_MODEL_ID
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}
return cfg


def nemotron_3_5_lightning_pretrain_8gpu_gb300_fp8mx_config() -> ConfigContainer:
"""Nemotron 3.5 Lightning pretrain: 8× GB300, MXFP8."""
cfg = _build_nemotron_3_5_lightning_gb300_mxfp8()
cfg.model.use_transformer_engine_op_fuser = True
cfg.model.moe_mlp_glu_interleave_size = 32
Comment thread
malay-nagda marked this conversation as resolved.
Outdated
cfg.mixed_precision.fp8_dot_product_attention = 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,
"CUDNNFE_CLUSTER_OVERLAP_MARGIN": 8,
"NVTE_BWD_LAYERNORM_SM_MARGIN": 20,
"NVTE_CUTEDSL_FUSED_GROUPED_MLP": 1,
"NVTE_FWD_LAYERNORM_SM_MARGIN": 20,
"NVTE_NORM_BWD_USE_CUDNN": 1,
"NVTE_NORM_FWD_USE_CUDNN": 1,
}
return cfg


def nemotron_3_5_lightning_pretrain_8gpu_gb300_fp8mx_fsdp_config() -> ConfigContainer:
"""Nemotron 3.5 Lightning pretrain: 8× GB300, MXFP8, Megatron FSDP."""
cfg = _build_nemotron_3_5_lightning_gb300_mxfp8()

# Megatron FSDP registers module hooks that Transformer Engine CUDA graph
# capture rejects.
cfg.train.global_batch_size = 384
cfg.train.micro_batch_size = 3
cfg.model.cuda_graph_impl = "none"
set_cuda_graph_modules(cfg.model, [])

cfg.model.init_model_with_meta_device = True
cfg.mixed_precision.reuse_grad_buf_for_mxfp8_param_ag = False
cfg.dist.use_megatron_fsdp = True
cfg.ddp.use_megatron_fsdp = True
cfg.ddp.num_distributed_optimizer_instances = 1
cfg.ddp.data_parallel_sharding_strategy = "optim_grads_params"
cfg.ddp.outer_dp_sharding_strategy = "no_shard"
cfg.ddp.average_in_collective = False
cfg.ddp.keep_fp8_transpose_cache = False
cfg.ddp.reuse_grad_buf_for_mxfp8_param_ag = False
cfg.optimizer.reuse_grad_buf_for_mxfp8_param_ag = False

cfg.checkpoint.load = None
cfg.checkpoint.ckpt_format = "fsdp_dtensor"
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,
"NVTE_BWD_LAYERNORM_SM_MARGIN": 20,
"NVTE_FWD_LAYERNORM_SM_MARGIN": 20,
"NVTE_NORM_BWD_USE_CUDNN": 1,
"NVTE_NORM_FWD_USE_CUDNN": 1,
}
return cfg


def nemotron_3_5_lightning_pretrain_8gpu_gb300_nvfp4_config() -> ConfigContainer:
"""Nemotron 3.5 Lightning pretrain: 8× GB300, NVFP4."""
cfg = nemotron_3_nano_pretrain_8gpu_gb300_nvfp4_config()
cfg.model.mtp_num_layers = 2
cfg.model.mtp_hybrid_override_pattern = "*E"
cfg.model.mtp_use_repeated_layer = True
cfg.model.keep_mtp_spec_in_bf16 = True
cfg.model.mtp_loss_scaling_factor = 0.3
cfg.model.hf_model_id = _NEMOTRON_3_5_LIGHTNING_MODEL_ID
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}
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,
"NVTE_BWD_LAYERNORM_SM_MARGIN": 20,
"NVTE_FWD_LAYERNORM_SM_MARGIN": 20,
"NVTE_NORM_BWD_USE_CUDNN": 1,
"NVTE_NORM_FWD_USE_CUDNN": 1,
"NVTE_USE_FAST_MATH": 1,
}
return cfg
94 changes: 83 additions & 11 deletions tests/unit_tests/recipes/test_nemotron_3_5_lightning_perf_recipes.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,11 +26,18 @@
nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_config,
nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_fsdp_config,
nemotron_3_5_lightning_pretrain_8gpu_gb200_nvfp4_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_gb300_fp8mx_fsdp_config,
nemotron_3_5_lightning_pretrain_8gpu_gb300_nvfp4_config,
nemotron_3_5_lightning_pretrain_16gpu_h100_bf16_config,
nemotron_3_5_lightning_pretrain_16gpu_h100_fp8cs_config,
nemotron_3_nano_pretrain_8gpu_gb200_bf16_config,
nemotron_3_nano_pretrain_8gpu_gb200_fp8mx_config,
nemotron_3_nano_pretrain_8gpu_gb200_nvfp4_config,
nemotron_3_nano_pretrain_8gpu_gb300_bf16_config,
nemotron_3_nano_pretrain_8gpu_gb300_fp8mx_config,
nemotron_3_nano_pretrain_8gpu_gb300_nvfp4_config,
nemotron_3_nano_pretrain_16gpu_h100_bf16_config,
nemotron_3_nano_pretrain_16gpu_h100_fp8cs_config,
)
Expand All @@ -51,13 +58,24 @@
nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_config,
nemotron_3_5_lightning_pretrain_8gpu_gb200_nvfp4_config,
)
_GB200_FSDP_RECIPES = (nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_fsdp_config,)
_GB300_RECIPES = (
nemotron_3_5_lightning_pretrain_8gpu_gb300_bf16_config,
nemotron_3_5_lightning_pretrain_8gpu_gb300_fp8mx_config,
nemotron_3_5_lightning_pretrain_8gpu_gb300_nvfp4_config,
)
_GB_FSDP_RECIPES = (
nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_fsdp_config,
nemotron_3_5_lightning_pretrain_8gpu_gb300_fp8mx_fsdp_config,
)
_NEMOTRON_3_RECIPES = (
nemotron_3_nano_pretrain_16gpu_h100_bf16_config,
nemotron_3_nano_pretrain_16gpu_h100_fp8cs_config,
nemotron_3_nano_pretrain_8gpu_gb200_bf16_config,
nemotron_3_nano_pretrain_8gpu_gb200_fp8mx_config,
nemotron_3_nano_pretrain_8gpu_gb200_nvfp4_config,
nemotron_3_nano_pretrain_8gpu_gb300_bf16_config,
nemotron_3_nano_pretrain_8gpu_gb300_fp8mx_config,
nemotron_3_nano_pretrain_8gpu_gb300_nvfp4_config,
)
_NEMOTRON_3_5_BASE_RECIPE_PAIRS = (
(
Expand All @@ -80,6 +98,18 @@
nemotron_3_5_lightning_pretrain_8gpu_gb200_nvfp4_config,
nemotron_3_nano_pretrain_8gpu_gb200_nvfp4_config,
),
(
nemotron_3_5_lightning_pretrain_8gpu_gb300_bf16_config,
nemotron_3_nano_pretrain_8gpu_gb300_bf16_config,
),
(
nemotron_3_5_lightning_pretrain_8gpu_gb300_fp8mx_config,
nemotron_3_nano_pretrain_8gpu_gb300_fp8mx_config,
),
(
nemotron_3_5_lightning_pretrain_8gpu_gb300_nvfp4_config,
nemotron_3_nano_pretrain_8gpu_gb300_nvfp4_config,
),
)
_NEMOTRON_NANO_PERF_FACTORIES = (
("megatron.bridge.perf_recipes.nemotronh.h100.nemotronh", "nemotron_3_nano_pretrain_16gpu_h100_bf16_config"),
Expand Down Expand Up @@ -123,6 +153,22 @@
"megatron.bridge.perf_recipes.nemotronh.gb200.nemotronh",
"nemotron_3_5_lightning_pretrain_8gpu_gb200_nvfp4_config",
),
(
"megatron.bridge.perf_recipes.nemotronh.gb300.nemotronh",
"nemotron_3_5_lightning_pretrain_8gpu_gb300_bf16_config",
),
(
"megatron.bridge.perf_recipes.nemotronh.gb300.nemotronh",
"nemotron_3_5_lightning_pretrain_8gpu_gb300_fp8mx_config",
),
(
"megatron.bridge.perf_recipes.nemotronh.gb300.nemotronh",
"nemotron_3_5_lightning_pretrain_8gpu_gb300_fp8mx_fsdp_config",
),
(
"megatron.bridge.perf_recipes.nemotronh.gb300.nemotronh",
"nemotron_3_5_lightning_pretrain_8gpu_gb300_nvfp4_config",
),
)


Expand All @@ -148,7 +194,7 @@ def test_standard_perf_recipes_do_not_expose_mtp_flag(recipe_factory: Callable[[

@pytest.mark.parametrize(
"recipe_factory",
(*_H100_RECIPES, *_GB200_RECIPES, *_GB200_FSDP_RECIPES),
(*_H100_RECIPES, *_GB200_RECIPES, *_GB300_RECIPES, *_GB_FSDP_RECIPES),
ids=lambda recipe: recipe.__name__,
)
def test_perf_recipes_enable_mtp(recipe_factory: Callable[[], ConfigContainer]) -> None:
Expand Down Expand Up @@ -188,9 +234,17 @@ def test_nemotron_3_5_perf_recipes_inherit_nemotron_3_policy(
assert cfg.tokenizer.tokenizer_model != base_cfg.tokenizer.tokenizer_model


def test_gb200_mxfp8_enables_cutedsl_fusion() -> None:
"""The non-FSDP Lightning recipe enables CutDSL without MoE A2A overlap."""
cfg = nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_config()
@pytest.mark.parametrize(
"recipe_factory",
(
nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_config,
nemotron_3_5_lightning_pretrain_8gpu_gb300_fp8mx_config,
),
ids=lambda recipe: recipe.__name__,
)
def test_gb_mxfp8_enables_cutedsl_fusion(recipe_factory: Callable[[], ConfigContainer]) -> None:
"""The non-FSDP Lightning GB recipes enable CutDSL without MoE A2A overlap."""
cfg = recipe_factory()

assert cfg.env_vars["NVTE_CUTEDSL_FUSED_GROUPED_MLP"] == 1
assert cfg.env_vars["CUDNNFE_CLUSTER_OVERLAP_MARGIN"] == 8
Expand All @@ -203,9 +257,10 @@ def test_gb200_mxfp8_enables_cutedsl_fusion() -> None:
assert cfg.comm_overlap.delay_wgrad_compute is False


def test_gb200_mxfp8_fsdp_skips_cutedsl_fusion() -> None:
"""The Lightning MXFP8 FSDP variant remains outside the CutDSL tuning scope."""
cfg = nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_fsdp_config()
@pytest.mark.parametrize("recipe_factory", _GB_FSDP_RECIPES, ids=lambda recipe: recipe.__name__)
def test_gb_mxfp8_fsdp_skips_cutedsl_fusion(recipe_factory: Callable[[], ConfigContainer]) -> None:
"""The Lightning GB MXFP8 FSDP variants remain outside the CutDSL tuning scope."""
cfg = recipe_factory()

assert "NVTE_CUTEDSL_FUSED_GROUPED_MLP" not in cfg.env_vars
assert "CUDNNFE_CLUSTER_OVERLAP_MARGIN" not in cfg.env_vars
Expand Down Expand Up @@ -317,9 +372,26 @@ def test_gb200_perf_recipe_topology(recipe_factory: Callable[[], ConfigContainer
assert cfg.env_vars["USE_MNNVL"] == 1


def test_gb200_fsdp_perf_recipe_defaults() -> None:
"""The GB200 FSDP variant retains its measured 8-GPU performance settings."""
cfg = nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_fsdp_config()
@pytest.mark.parametrize("recipe_factory", _GB300_RECIPES, ids=lambda recipe: recipe.__name__)
def test_gb300_perf_recipe_topology(recipe_factory: Callable[[], ConfigContainer]) -> None:
"""GB300 Nemotron 3.5 Lightning variants retain the established performance topology."""
cfg = recipe_factory()

assert cfg.model.expert_model_parallel_size == 8
assert cfg.train.global_batch_size == 512
assert cfg.train.micro_batch_size == 4
assert cfg.model.recompute_granularity is None
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


@pytest.mark.parametrize("recipe_factory", _GB_FSDP_RECIPES, ids=lambda recipe: recipe.__name__)
def test_gb_fsdp_perf_recipe_defaults(recipe_factory: Callable[[], ConfigContainer]) -> None:
"""The GB FSDP variants retain their measured 8-GPU performance settings."""
cfg = recipe_factory()

assert cfg.train.global_batch_size == 384
assert cfg.train.micro_batch_size == 3
Expand Down
Loading