Skip to content

Commit 6124fbe

Browse files
svcnvidia-nemo-cimalay-nagda
svcnvidia-nemo-ci
and
malay-nagda
authored
cp: perf(nemotron): add Nemotron 3.5 Lightning GB300 recipes (5693) into r0.6.0 (#5734)
Signed-off-by: Malay Nagda <malayn@nvidia.com> Signed-off-by: NeMo Bot <nemo-bot@nvidia.com> Co-authored-by: malay-nagda <malayn@nvidia.com>
1 parent 1507559 commit 6124fbe

5 files changed

Lines changed: 245 additions & 13 deletions

File tree

‎scripts/performance/dump_perf_configs.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -277,6 +277,10 @@
277277
("nemotronh", "nemotron_3_nano", "pretrain", 16, "h100", "bf16"),
278278
("nemotronh", "nemotron_3_nano", "pretrain", 16, "h100", "fp8_cs"),
279279
# Nemotron 3.5 Lightning
280+
("nemotronh", "nemotron_3_5_lightning", "pretrain", 8, "gb300", "bf16"),
281+
("nemotronh", "nemotron_3_5_lightning", "pretrain", 8, "gb300", "fp8_mx"),
282+
("nemotronh", "nemotron_3_5_lightning", "pretrain", 8, "gb300", "fp8_mx", "fsdp"),
283+
("nemotronh", "nemotron_3_5_lightning", "pretrain", 8, "gb300", "nvfp4"),
280284
("nemotronh", "nemotron_3_5_lightning", "pretrain", 8, "gb200", "bf16"),
281285
("nemotronh", "nemotron_3_5_lightning", "pretrain", 8, "gb200", "fp8_mx"),
282286
("nemotronh", "nemotron_3_5_lightning", "pretrain", 8, "gb200", "nvfp4"),

‎src/megatron/bridge/perf_recipes/nemotronh/__init__.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,10 @@
3333
nemotronh_56b_pretrain_64gpu_gb200_fp8cs_config,
3434
)
3535
from megatron.bridge.perf_recipes.nemotronh.gb300.nemotronh import (
36+
nemotron_3_5_lightning_pretrain_8gpu_gb300_bf16_config,
37+
nemotron_3_5_lightning_pretrain_8gpu_gb300_fp8mx_config,
38+
nemotron_3_5_lightning_pretrain_8gpu_gb300_fp8mx_fsdp_config,
39+
nemotron_3_5_lightning_pretrain_8gpu_gb300_nvfp4_config,
3640
nemotron_3_nano_pretrain_8gpu_gb300_bf16_config,
3741
nemotron_3_nano_pretrain_8gpu_gb300_fp8mx_config,
3842
nemotron_3_nano_pretrain_8gpu_gb300_nvfp4_config,

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

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -523,7 +523,6 @@ def nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_config() -> ConfigContainer
523523
"""Nemotron 3.5 Lightning pretrain: 8× GB200, MXFP8."""
524524
cfg = _build_nemotron_3_5_lightning_gb200_mxfp8()
525525
cfg.model.use_transformer_engine_op_fuser = True
526-
cfg.model.moe_mlp_glu_interleave_size = 32
527526
cfg.mixed_precision.fp8_dot_product_attention = True
528527
cfg.env_vars = {
529528
**COMMON_PERF_ENV_VARS,

‎src/megatron/bridge/perf_recipes/nemotronh/gb300/nemotronh.py‎

Lines changed: 153 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,11 @@
3434
from megatron.bridge.utils.cuda_graph import set_cuda_graph_modules
3535

3636

37+
# Public Nemotron 3.5 Lightning checkpoint used by the Lightning recipe API.
38+
_NEMOTRON_3_5_LIGHTNING_MODEL_ID = "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16"
39+
_NEMOTRON_3_5_LIGHTNING_MODEL_REVISION = "b3caaabed0263651a17dc1f2d4ce97e794f76c44" # pragma: allowlist secret
40+
41+
3742
def nemotronh_56b_pretrain_64gpu_gb300_fp8cs_config() -> ConfigContainer:
3843
"""NemotronH 56B pretrain: 64× GB300, FP8 current-scaling."""
3944
cfg = nemotronh_56b_pretrain_config()
@@ -524,3 +529,151 @@ def nemotronh_56b_pretrain_256gpu_gb300_fp8cs_config() -> ConfigContainer:
524529
"NVTE_FWD_LAYERNORM_SM_MARGIN": 20,
525530
}
526531
return cfg
532+
533+
534+
def nemotron_3_5_lightning_pretrain_8gpu_gb300_bf16_config() -> ConfigContainer:
535+
"""Nemotron 3.5 Lightning pretrain: 8× GB300, BF16."""
536+
cfg = nemotron_3_nano_pretrain_8gpu_gb300_bf16_config()
537+
cfg.model.mtp_num_layers = 2
538+
cfg.model.mtp_hybrid_override_pattern = "*E"
539+
cfg.model.mtp_use_repeated_layer = True
540+
cfg.model.keep_mtp_spec_in_bf16 = True
541+
cfg.model.mtp_loss_scaling_factor = 0.3
542+
cfg.model.hf_model_id = _NEMOTRON_3_5_LIGHTNING_MODEL_ID
543+
cfg.model.hf_model_revision = _NEMOTRON_3_5_LIGHTNING_MODEL_REVISION
544+
cfg.tokenizer.tokenizer_model = _NEMOTRON_3_5_LIGHTNING_MODEL_ID
545+
cfg.tokenizer.hf_tokenizer_kwargs = {"revision": _NEMOTRON_3_5_LIGHTNING_MODEL_REVISION}
546+
cfg.env_vars = {
547+
**COMMON_PERF_ENV_VARS,
548+
"CUDA_DEVICE_MAX_CONNECTIONS": 32,
549+
"NCCL_GRAPH_REGISTER": 0,
550+
"PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True",
551+
"TORCH_NCCL_AVOID_RECORD_STREAMS": 1,
552+
"NCCL_NVLS_ENABLE": 0,
553+
"NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 8,
554+
"NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128,
555+
"NVLINK_DOMAIN_SIZE": 72,
556+
"USE_MNNVL": 1,
557+
"NVTE_BWD_LAYERNORM_SM_MARGIN": 20,
558+
"NVTE_FWD_LAYERNORM_SM_MARGIN": 20,
559+
"NVTE_NORM_BWD_USE_CUDNN": 1,
560+
"NVTE_NORM_FWD_USE_CUDNN": 1,
561+
}
562+
return cfg
563+
564+
565+
def _build_nemotron_3_5_lightning_gb300_mxfp8() -> ConfigContainer:
566+
cfg = _build_nemotron_3_nano_gb300_mxfp8()
567+
cfg.model.moe_hybridep_num_sms = 16
568+
cfg.model.mtp_num_layers = 2
569+
cfg.model.mtp_hybrid_override_pattern = "*E"
570+
cfg.model.mtp_use_repeated_layer = True
571+
cfg.model.keep_mtp_spec_in_bf16 = True
572+
cfg.model.mtp_loss_scaling_factor = 0.3
573+
cfg.model.hf_model_id = _NEMOTRON_3_5_LIGHTNING_MODEL_ID
574+
cfg.model.hf_model_revision = _NEMOTRON_3_5_LIGHTNING_MODEL_REVISION
575+
cfg.tokenizer.tokenizer_model = _NEMOTRON_3_5_LIGHTNING_MODEL_ID
576+
cfg.tokenizer.hf_tokenizer_kwargs = {"revision": _NEMOTRON_3_5_LIGHTNING_MODEL_REVISION}
577+
return cfg
578+
579+
580+
def nemotron_3_5_lightning_pretrain_8gpu_gb300_fp8mx_config() -> ConfigContainer:
581+
"""Nemotron 3.5 Lightning pretrain: 8× GB300, MXFP8."""
582+
cfg = _build_nemotron_3_5_lightning_gb300_mxfp8()
583+
cfg.model.use_transformer_engine_op_fuser = True
584+
cfg.mixed_precision.fp8_dot_product_attention = True
585+
cfg.env_vars = {
586+
**COMMON_PERF_ENV_VARS,
587+
"CUDA_DEVICE_MAX_CONNECTIONS": 32,
588+
"NCCL_GRAPH_REGISTER": 0,
589+
"PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True",
590+
"TORCH_NCCL_AVOID_RECORD_STREAMS": 1,
591+
"NCCL_NVLS_ENABLE": 0,
592+
"NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 8,
593+
"NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128,
594+
"NVLINK_DOMAIN_SIZE": 72,
595+
"USE_MNNVL": 1,
596+
"CUDNNFE_CLUSTER_OVERLAP_MARGIN": 8,
597+
"NVTE_BWD_LAYERNORM_SM_MARGIN": 20,
598+
"NVTE_CUTEDSL_FUSED_GROUPED_MLP": 1,
599+
"NVTE_FWD_LAYERNORM_SM_MARGIN": 20,
600+
"NVTE_NORM_BWD_USE_CUDNN": 1,
601+
"NVTE_NORM_FWD_USE_CUDNN": 1,
602+
}
603+
return cfg
604+
605+
606+
def nemotron_3_5_lightning_pretrain_8gpu_gb300_fp8mx_fsdp_config() -> ConfigContainer:
607+
"""Nemotron 3.5 Lightning pretrain: 8× GB300, MXFP8, Megatron FSDP."""
608+
cfg = _build_nemotron_3_5_lightning_gb300_mxfp8()
609+
610+
# Megatron FSDP registers module hooks that Transformer Engine CUDA graph
611+
# capture rejects.
612+
cfg.train.global_batch_size = 384
613+
cfg.train.micro_batch_size = 3
614+
cfg.model.cuda_graph_impl = "none"
615+
set_cuda_graph_modules(cfg.model, [])
616+
617+
cfg.model.init_model_with_meta_device = True
618+
cfg.mixed_precision.reuse_grad_buf_for_mxfp8_param_ag = False
619+
cfg.dist.use_megatron_fsdp = True
620+
cfg.ddp.use_megatron_fsdp = True
621+
cfg.ddp.num_distributed_optimizer_instances = 1
622+
cfg.ddp.data_parallel_sharding_strategy = "optim_grads_params"
623+
cfg.ddp.outer_dp_sharding_strategy = "no_shard"
624+
cfg.ddp.average_in_collective = False
625+
cfg.ddp.keep_fp8_transpose_cache = False
626+
cfg.ddp.reuse_grad_buf_for_mxfp8_param_ag = False
627+
cfg.optimizer.reuse_grad_buf_for_mxfp8_param_ag = False
628+
629+
cfg.checkpoint.load = None
630+
cfg.checkpoint.ckpt_format = "fsdp_dtensor"
631+
cfg.env_vars = {
632+
**COMMON_PERF_ENV_VARS,
633+
"CUDA_DEVICE_MAX_CONNECTIONS": 32,
634+
"NCCL_GRAPH_REGISTER": 0,
635+
"PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True",
636+
"TORCH_NCCL_AVOID_RECORD_STREAMS": 1,
637+
"NCCL_NVLS_ENABLE": 0,
638+
"NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 8,
639+
"NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128,
640+
"NVLINK_DOMAIN_SIZE": 72,
641+
"USE_MNNVL": 1,
642+
"NVTE_BWD_LAYERNORM_SM_MARGIN": 20,
643+
"NVTE_FWD_LAYERNORM_SM_MARGIN": 20,
644+
"NVTE_NORM_BWD_USE_CUDNN": 1,
645+
"NVTE_NORM_FWD_USE_CUDNN": 1,
646+
}
647+
return cfg
648+
649+
650+
def nemotron_3_5_lightning_pretrain_8gpu_gb300_nvfp4_config() -> ConfigContainer:
651+
"""Nemotron 3.5 Lightning pretrain: 8× GB300, NVFP4."""
652+
cfg = nemotron_3_nano_pretrain_8gpu_gb300_nvfp4_config()
653+
cfg.model.mtp_num_layers = 2
654+
cfg.model.mtp_hybrid_override_pattern = "*E"
655+
cfg.model.mtp_use_repeated_layer = True
656+
cfg.model.keep_mtp_spec_in_bf16 = True
657+
cfg.model.mtp_loss_scaling_factor = 0.3
658+
cfg.model.hf_model_id = _NEMOTRON_3_5_LIGHTNING_MODEL_ID
659+
cfg.model.hf_model_revision = _NEMOTRON_3_5_LIGHTNING_MODEL_REVISION
660+
cfg.tokenizer.tokenizer_model = _NEMOTRON_3_5_LIGHTNING_MODEL_ID
661+
cfg.tokenizer.hf_tokenizer_kwargs = {"revision": _NEMOTRON_3_5_LIGHTNING_MODEL_REVISION}
662+
cfg.env_vars = {
663+
**COMMON_PERF_ENV_VARS,
664+
"CUDA_DEVICE_MAX_CONNECTIONS": 32,
665+
"NCCL_GRAPH_REGISTER": 0,
666+
"PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True",
667+
"TORCH_NCCL_AVOID_RECORD_STREAMS": 1,
668+
"NCCL_NVLS_ENABLE": 0,
669+
"NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 8,
670+
"NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128,
671+
"NVLINK_DOMAIN_SIZE": 72,
672+
"USE_MNNVL": 1,
673+
"NVTE_BWD_LAYERNORM_SM_MARGIN": 20,
674+
"NVTE_FWD_LAYERNORM_SM_MARGIN": 20,
675+
"NVTE_NORM_BWD_USE_CUDNN": 1,
676+
"NVTE_NORM_FWD_USE_CUDNN": 1,
677+
"NVTE_USE_FAST_MATH": 1,
678+
}
679+
return cfg

‎tests/unit_tests/recipes/test_nemotron_3_5_lightning_perf_recipes.py‎

Lines changed: 84 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -26,11 +26,18 @@
2626
nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_config,
2727
nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_fsdp_config,
2828
nemotron_3_5_lightning_pretrain_8gpu_gb200_nvfp4_config,
29+
nemotron_3_5_lightning_pretrain_8gpu_gb300_bf16_config,
30+
nemotron_3_5_lightning_pretrain_8gpu_gb300_fp8mx_config,
31+
nemotron_3_5_lightning_pretrain_8gpu_gb300_fp8mx_fsdp_config,
32+
nemotron_3_5_lightning_pretrain_8gpu_gb300_nvfp4_config,
2933
nemotron_3_5_lightning_pretrain_16gpu_h100_bf16_config,
3034
nemotron_3_5_lightning_pretrain_16gpu_h100_fp8cs_config,
3135
nemotron_3_nano_pretrain_8gpu_gb200_bf16_config,
3236
nemotron_3_nano_pretrain_8gpu_gb200_fp8mx_config,
3337
nemotron_3_nano_pretrain_8gpu_gb200_nvfp4_config,
38+
nemotron_3_nano_pretrain_8gpu_gb300_bf16_config,
39+
nemotron_3_nano_pretrain_8gpu_gb300_fp8mx_config,
40+
nemotron_3_nano_pretrain_8gpu_gb300_nvfp4_config,
3441
nemotron_3_nano_pretrain_16gpu_h100_bf16_config,
3542
nemotron_3_nano_pretrain_16gpu_h100_fp8cs_config,
3643
)
@@ -51,13 +58,24 @@
5158
nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_config,
5259
nemotron_3_5_lightning_pretrain_8gpu_gb200_nvfp4_config,
5360
)
54-
_GB200_FSDP_RECIPES = (nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_fsdp_config,)
61+
_GB300_RECIPES = (
62+
nemotron_3_5_lightning_pretrain_8gpu_gb300_bf16_config,
63+
nemotron_3_5_lightning_pretrain_8gpu_gb300_fp8mx_config,
64+
nemotron_3_5_lightning_pretrain_8gpu_gb300_nvfp4_config,
65+
)
66+
_GB_FSDP_RECIPES = (
67+
nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_fsdp_config,
68+
nemotron_3_5_lightning_pretrain_8gpu_gb300_fp8mx_fsdp_config,
69+
)
5570
_NEMOTRON_3_RECIPES = (
5671
nemotron_3_nano_pretrain_16gpu_h100_bf16_config,
5772
nemotron_3_nano_pretrain_16gpu_h100_fp8cs_config,
5873
nemotron_3_nano_pretrain_8gpu_gb200_bf16_config,
5974
nemotron_3_nano_pretrain_8gpu_gb200_fp8mx_config,
6075
nemotron_3_nano_pretrain_8gpu_gb200_nvfp4_config,
76+
nemotron_3_nano_pretrain_8gpu_gb300_bf16_config,
77+
nemotron_3_nano_pretrain_8gpu_gb300_fp8mx_config,
78+
nemotron_3_nano_pretrain_8gpu_gb300_nvfp4_config,
6179
)
6280
_NEMOTRON_3_5_BASE_RECIPE_PAIRS = (
6381
(
@@ -80,6 +98,18 @@
8098
nemotron_3_5_lightning_pretrain_8gpu_gb200_nvfp4_config,
8199
nemotron_3_nano_pretrain_8gpu_gb200_nvfp4_config,
82100
),
101+
(
102+
nemotron_3_5_lightning_pretrain_8gpu_gb300_bf16_config,
103+
nemotron_3_nano_pretrain_8gpu_gb300_bf16_config,
104+
),
105+
(
106+
nemotron_3_5_lightning_pretrain_8gpu_gb300_fp8mx_config,
107+
nemotron_3_nano_pretrain_8gpu_gb300_fp8mx_config,
108+
),
109+
(
110+
nemotron_3_5_lightning_pretrain_8gpu_gb300_nvfp4_config,
111+
nemotron_3_nano_pretrain_8gpu_gb300_nvfp4_config,
112+
),
83113
)
84114
_NEMOTRON_NANO_PERF_FACTORIES = (
85115
("megatron.bridge.perf_recipes.nemotronh.h100.nemotronh", "nemotron_3_nano_pretrain_16gpu_h100_bf16_config"),
@@ -123,6 +153,22 @@
123153
"megatron.bridge.perf_recipes.nemotronh.gb200.nemotronh",
124154
"nemotron_3_5_lightning_pretrain_8gpu_gb200_nvfp4_config",
125155
),
156+
(
157+
"megatron.bridge.perf_recipes.nemotronh.gb300.nemotronh",
158+
"nemotron_3_5_lightning_pretrain_8gpu_gb300_bf16_config",
159+
),
160+
(
161+
"megatron.bridge.perf_recipes.nemotronh.gb300.nemotronh",
162+
"nemotron_3_5_lightning_pretrain_8gpu_gb300_fp8mx_config",
163+
),
164+
(
165+
"megatron.bridge.perf_recipes.nemotronh.gb300.nemotronh",
166+
"nemotron_3_5_lightning_pretrain_8gpu_gb300_fp8mx_fsdp_config",
167+
),
168+
(
169+
"megatron.bridge.perf_recipes.nemotronh.gb300.nemotronh",
170+
"nemotron_3_5_lightning_pretrain_8gpu_gb300_nvfp4_config",
171+
),
126172
)
127173

128174

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

149195
@pytest.mark.parametrize(
150196
"recipe_factory",
151-
(*_H100_RECIPES, *_GB200_RECIPES, *_GB200_FSDP_RECIPES),
197+
(*_H100_RECIPES, *_GB200_RECIPES, *_GB300_RECIPES, *_GB_FSDP_RECIPES),
152198
ids=lambda recipe: recipe.__name__,
153199
)
154200
def test_perf_recipes_enable_mtp(recipe_factory: Callable[[], ConfigContainer]) -> None:
@@ -188,24 +234,33 @@ def test_nemotron_3_5_perf_recipes_inherit_nemotron_3_policy(
188234
assert cfg.tokenizer.tokenizer_model != base_cfg.tokenizer.tokenizer_model
189235

190236

191-
def test_gb200_mxfp8_enables_cutedsl_fusion() -> None:
192-
"""The non-FSDP Lightning recipe enables CutDSL without MoE A2A overlap."""
193-
cfg = nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_config()
237+
@pytest.mark.parametrize(
238+
"recipe_factory",
239+
(
240+
nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_config,
241+
nemotron_3_5_lightning_pretrain_8gpu_gb300_fp8mx_config,
242+
),
243+
ids=lambda recipe: recipe.__name__,
244+
)
245+
def test_gb_mxfp8_enables_cutedsl_fusion(recipe_factory: Callable[[], ConfigContainer]) -> None:
246+
"""The non-FSDP Lightning GB recipes enable CutDSL without MoE A2A overlap."""
247+
cfg = recipe_factory()
194248

195249
assert cfg.env_vars["NVTE_CUTEDSL_FUSED_GROUPED_MLP"] == 1
196250
assert cfg.env_vars["CUDNNFE_CLUSTER_OVERLAP_MARGIN"] == 8
197251
assert cfg.model.use_transformer_engine_op_fuser is True
198-
assert cfg.model.moe_mlp_glu_interleave_size == 32
252+
assert cfg.model.moe_mlp_glu_interleave_size is None
199253
assert cfg.model.high_priority_a2a_comm_stream is False
200254
assert cfg.model.moe_hybridep_num_sms_preprocessing == 108
201255
assert cfg.mixed_precision.fp8_dot_product_attention is True
202256
assert cfg.comm_overlap.overlap_moe_expert_parallel_comm is False
203257
assert cfg.comm_overlap.delay_wgrad_compute is False
204258

205259

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

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

319374

320-
def test_gb200_fsdp_perf_recipe_defaults() -> None:
321-
"""The GB200 FSDP variant retains its measured 8-GPU performance settings."""
322-
cfg = nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_fsdp_config()
375+
@pytest.mark.parametrize("recipe_factory", _GB300_RECIPES, ids=lambda recipe: recipe.__name__)
376+
def test_gb300_perf_recipe_topology(recipe_factory: Callable[[], ConfigContainer]) -> None:
377+
"""GB300 Nemotron 3.5 Lightning variants retain the established performance topology."""
378+
cfg = recipe_factory()
379+
380+
assert cfg.model.expert_model_parallel_size == 8
381+
assert cfg.train.global_batch_size == 512
382+
assert cfg.train.micro_batch_size == 4
383+
assert cfg.model.recompute_granularity is None
384+
assert cfg.model.seq_length == 8192
385+
assert cfg.dataset.seq_length == 8192
386+
assert cfg.model.moe_hybridep_num_sms == 16
387+
assert cfg.env_vars["NVLINK_DOMAIN_SIZE"] == 72
388+
assert cfg.env_vars["USE_MNNVL"] == 1
389+
390+
391+
@pytest.mark.parametrize("recipe_factory", _GB_FSDP_RECIPES, ids=lambda recipe: recipe.__name__)
392+
def test_gb_fsdp_perf_recipe_defaults(recipe_factory: Callable[[], ConfigContainer]) -> None:
393+
"""The GB FSDP variants retain their measured 8-GPU performance settings."""
394+
cfg = recipe_factory()
323395

324396
assert cfg.train.global_batch_size == 384
325397
assert cfg.train.micro_batch_size == 3

0 commit comments

Comments
 (0)