diff --git a/scripts/performance/dump_perf_configs.py b/scripts/performance/dump_perf_configs.py index ae8ee0768f..a69a6217da 100644 --- a/scripts/performance/dump_perf_configs.py +++ b/scripts/performance/dump_perf_configs.py @@ -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"), diff --git a/src/megatron/bridge/perf_recipes/nemotronh/__init__.py b/src/megatron/bridge/perf_recipes/nemotronh/__init__.py index 2606829e67..5aeb9fa435 100644 --- a/src/megatron/bridge/perf_recipes/nemotronh/__init__.py +++ b/src/megatron/bridge/perf_recipes/nemotronh/__init__.py @@ -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, diff --git a/src/megatron/bridge/perf_recipes/nemotronh/gb200/nemotronh.py b/src/megatron/bridge/perf_recipes/nemotronh/gb200/nemotronh.py index 1b0c8ca4dc..f4dbf11710 100644 --- a/src/megatron/bridge/perf_recipes/nemotronh/gb200/nemotronh.py +++ b/src/megatron/bridge/perf_recipes/nemotronh/gb200/nemotronh.py @@ -523,7 +523,6 @@ def nemotron_3_5_lightning_pretrain_8gpu_gb200_fp8mx_config() -> ConfigContainer """Nemotron 3.5 Lightning pretrain: 8× GB200, MXFP8.""" cfg = _build_nemotron_3_5_lightning_gb200_mxfp8() cfg.model.use_transformer_engine_op_fuser = True - cfg.model.moe_mlp_glu_interleave_size = 32 cfg.mixed_precision.fp8_dot_product_attention = True cfg.env_vars = { **COMMON_PERF_ENV_VARS, diff --git a/src/megatron/bridge/perf_recipes/nemotronh/gb300/nemotronh.py b/src/megatron/bridge/perf_recipes/nemotronh/gb300/nemotronh.py index 5ce3921891..97c0b27c0a 100644 --- a/src/megatron/bridge/perf_recipes/nemotronh/gb300/nemotronh.py +++ b/src/megatron/bridge/perf_recipes/nemotronh/gb300/nemotronh.py @@ -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() @@ -524,3 +529,151 @@ 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.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 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 71e12b7ead..394dfe8d4e 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 @@ -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, ) @@ -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 = ( ( @@ -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"), @@ -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", + ), ) @@ -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: @@ -188,14 +234,22 @@ 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 assert cfg.model.use_transformer_engine_op_fuser is True - assert cfg.model.moe_mlp_glu_interleave_size == 32 + assert cfg.model.moe_mlp_glu_interleave_size is None assert cfg.model.high_priority_a2a_comm_stream is False assert cfg.model.moe_hybridep_num_sms_preprocessing == 108 assert cfg.mixed_precision.fp8_dot_product_attention is True @@ -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 @@ -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