From 2d66113e0f2be452edbab074dc345d19bb52f924 Mon Sep 17 00:00:00 2001 From: Ravi Ghadia Date: Wed, 29 Jul 2026 13:29:15 -0700 Subject: [PATCH 1/3] feat(perf): add ultra pretrain configurations for Nemotron 3 on GB200 Signed-off-by: Ravi Ghadia --- .../bridge/perf_recipes/nemotronh/__init__.py | 2 + .../perf_recipes/nemotronh/gb200/nemotronh.py | 155 ++++++++++++++++++ 2 files changed, 157 insertions(+) diff --git a/src/megatron/bridge/perf_recipes/nemotronh/__init__.py b/src/megatron/bridge/perf_recipes/nemotronh/__init__.py index 9e2bb0ad99..10cd014e8a 100644 --- a/src/megatron/bridge/perf_recipes/nemotronh/__init__.py +++ b/src/megatron/bridge/perf_recipes/nemotronh/__init__.py @@ -25,6 +25,8 @@ nemotron_3_super_pretrain_64gpu_gb200_bf16_config, nemotron_3_super_pretrain_64gpu_gb200_fp8mx_config, nemotron_3_super_pretrain_64gpu_gb200_nvfp4_config, + nemotron_3_ultra_pretrain_8gpu_gb200_fp8mx_config, + nemotron_3_ultra_pretrain_256gpu_gb200_fp8mx_config, nemotronh_56b_pretrain_64gpu_gb200_fp8cs_config, ) from megatron.bridge.perf_recipes.nemotronh.gb300.nemotronh import ( diff --git a/src/megatron/bridge/perf_recipes/nemotronh/gb200/nemotronh.py b/src/megatron/bridge/perf_recipes/nemotronh/gb200/nemotronh.py index a9f9a73b7e..65c3a7f1f1 100644 --- a/src/megatron/bridge/perf_recipes/nemotronh/gb200/nemotronh.py +++ b/src/megatron/bridge/perf_recipes/nemotronh/gb200/nemotronh.py @@ -18,12 +18,15 @@ _TE_QUANT_CFG_PATH, ConfigContainer, _apply_nemotron_3_super_perf_defaults, + _apply_nemotron_3_ultra_fsdp_hsdp, + _apply_nemotron_3_ultra_perf_defaults, _benchmark_common, _nemotron_3_super_nvfp4_precision, _perf_precision, load_quantization_recipe, nemotron_3_nano_pretrain_config, nemotron_3_super_pretrain_config, + nemotron_3_ultra_pretrain_config, nemotronh_56b_pretrain_config, ) @@ -206,6 +209,158 @@ def nemotron_3_super_pretrain_64gpu_gb200_nvfp4_config() -> ConfigContainer: return cfg +def _nemotron_3_ultra_gb200_fp8mx_config( + *, + num_gpus: int, + expert_model_parallel_size: int, + global_batch_size: int, + hybrid_ep_ranks_per_nvlink_domain: int, +) -> ConfigContainer: + """Shared builder for Nemotron 3 Ultra GB200 MXFP8 Megatron-FSDP perf recipes. + + Same Megatron-FSDP (HSDP) layout as the GB300 recipe, but with TP2 + sequence + parallelism instead of TP1. GB200 carries less HBM per GPU, and TP alone would not + help much: it shards the internal GEMM dimensions but leaves the residual stream and + the norm inputs/outputs replicated. Sequence parallelism is what shards those, so the + two are enabled together. + """ + cfg = nemotron_3_ultra_pretrain_config() + cfg.mixed_precision = _perf_precision("fp8_mx") + + _apply_nemotron_3_ultra_perf_defaults(cfg) + + # Apply HSDP / FSDP dtype overrides last so they win over the generic defaults. + _apply_nemotron_3_ultra_fsdp_hsdp(cfg, num_gpus=num_gpus) + + # Parallelism + cfg.model.tensor_model_parallel_size = 2 + cfg.model.pipeline_model_parallel_size = 1 + cfg.model.virtual_pipeline_model_parallel_size = None + cfg.model.context_parallel_size = 1 + cfg.model.sequence_parallel = True + cfg.model.expert_tensor_parallel_size = 1 + cfg.model.pipeline_model_parallel_layout = None + cfg.model.seq_length = 8192 + + # Only tensors larger than 500MB are offloaded, which + # approximates offloading the moe_act input for seq 8192 / MBS 1. + cfg.model.min_offloaded_tensor_size = 500_000_000 + + # MXFP8 requires router padding for quantization. + cfg.model.moe_router_padding_for_quantization = True + + # GPU-count specific overrides of the canonical (256-GPU / EP64) defaults. + cfg.model.expert_model_parallel_size = expert_model_parallel_size + cfg.train.global_batch_size = global_batch_size + + # Fine-grained activation offloading. Requires NVTE_CPU_OFFLOAD_V1=1 in the + # launch environment (set in this recipe's env_vars). + # NOTE: also requires setting the min_offloaded_tensor_size to avoid CPU OOM issues + cfg.model.fine_grained_activation_offloading = True + cfg.model.offload_modules = ["fused_group_mlp"] + cfg.model.fine_grained_offloading_max_inflight_offloads = 1 + + # Selective recompute of the MoE activation + # recomputes the activation output of the MoE expert MLP, while FC1 output (activation input) is saved and offloaded to cpu + cfg.model.recompute_granularity = "selective" + cfg.model.recompute_modules = ["moe_act"] + + # full iteration cuda graph + cfg.model.cuda_graph_impl = "full_iteration" # "transformer_engine" + cfg.model.cuda_graph_scope = [] # ["attn", "mamba", "moe_router", "moe_preprocess"] + cfg.ddp.megatron_fsdp_cuda_graph_mode = True + cfg.ddp.fsdp_all_gather_in_start_param_sync = False + cfg.model.moe_expert_rank_capacity_factor = 1.2 + # cfg.model.moe_paged_stash = True + # cfg.model.moe_paged_stash_buffer_size_factor_cpu = 1.0 + # cfg.model.moe_paged_stash_buffer_size_factor_cuda = 1.2 + + # Keep process settings next to the recipe so users can see the exact benchmark environment. + cfg.env_vars = { + **COMMON_PERF_ENV_VARS, + # CUDA stream scheduling for this model and parallel layout. + "CUDA_DEVICE_MAX_CONNECTIONS": 32, + # CUDA graph and allocator behavior for this recipe. + "NCCL_GRAPH_REGISTER": 0, + # graph_capture_record_stream_reuse is required here, not an + # optimization: fine_grained_activation_offloading calls record_stream() + # on every D2H staging tensor, and cudaEventQuery is illegal during graph + # capture, so without this the allocator cannot recycle any of them until + # capture ends. Offloading then frees nothing on the capture step and + # reserved memory overshoots the eager peak (OOM at 256 GPUs / EP64). The + # option uses the captured DAG topology instead of events, which makes + # the one-time capture slower. TORCH_NCCL_AVOID_RECORD_STREAMS below + # covers only the NCCL buffers, not the offload path. + "PYTORCH_CUDA_ALLOC_CONF": ("expandable_segments:True,graph_capture_record_stream_reuse:True"), + "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, + # NCCL user-buffer and launch settings. + "NCCL_NVLS_ENABLE": 0, + # HybridEP topology for the target system. The per-domain rank count has to track + # EP: HybridEP splits the all-to-all into an intra-domain and an inter-domain leg + # using this value, so leaving it at the 256-GPU 64 on a smaller job would + # describe a domain larger than the job itself. + "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": hybrid_ep_ranks_per_nvlink_domain, + "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, + "NVLINK_DOMAIN_SIZE": 72, + "USE_MNNVL": 1, + # Transformer Engine overlap settings for this model. + "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, + "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, + # Required by fine_grained_activation_offloading (TE >= 2.10.0) to avoid + # offloading weights; + "NVTE_CPU_OFFLOAD_V1": 1, + # Enable TE's CuteDSL fused grouped MLP kernel (sm100+). Required by the + # op fuser + fused weighted squared-ReLU with moe_act activation recompute + # (ScaledSReLU(activation_recompute_in_mlp=True) only runs on this path). + "NVTE_CUTEDSL_FUSED_GROUPED_MLP": 1, + } + return cfg + + +def nemotron_3_ultra_pretrain_256gpu_gb200_fp8mx_config() -> ConfigContainer: + """Nemotron 3 Ultra (550B-A55B LatentMoE) pretrain: 256× GB200, MXFP8, Megatron-FSDP (HSDP). + + TP2 + SP / PP1 / CP1 / EP64 / ETP1, GBS 256 / MBS 1, seq 8192, BF16 + MXFP8 mixed + precision, HybridEP flex dispatcher, CuteDSL fused grouped MLP, selective recompute + + fine-grained activation offload of the expert MLP, MTP=2. + + This is the GB300 recipe's layout with TP raised from 1 to 2 and sequence parallelism + turned on, to fit GB200's smaller HBM. Megatron-FSDP shards params, grads and + optimizer state within each 64-GPU NVLink domain and replicates optimizer-sharded + across the four domains. + """ + return _nemotron_3_ultra_gb200_fp8mx_config( + num_gpus=256, + expert_model_parallel_size=64, + global_batch_size=256, + hybrid_ep_ranks_per_nvlink_domain=64, + ) + + +def nemotron_3_ultra_pretrain_8gpu_gb200_fp8mx_config() -> ConfigContainer: + """Nemotron 3 Ultra pipeclean: 8× GB200, MXFP8, Megatron-FSDP. + + Shrunk-grid counterpart of the 256-GPU recipe, for validating that the full feature + stack launches and trains before committing a full-scale job. EP drops 64 -> 4 so the + expert grid (``ETP1 x EP4`` = 4) divides an 8-GPU world, and GBS drops 256 -> 8. With + TP2 the batch splits over DP = 8 / 2 = 4, so GBS 8 / MBS 1 runs 2 microbatches. + + Below one NVLink domain ``_apply_nemotron_3_ultra_fsdp_hsdp`` resolves + ``num_distributed_optimizer_instances`` to 1, which turns HSDP off and correspondingly + sets ``outer_dp_sharding_strategy`` to ``no_shard``. + + This recipe only rescales parallelism: the model is still full-size Nemotron 3 Ultra, + which does not fit on 8 GPUs. Pair it with model-shrinking overrides (see + ``perf_bash_scripts/nt3_ultra_gb200/toy_run_perf_test_nemotron_3_ultra_gb200_fp8mx.sh``). + """ + return _nemotron_3_ultra_gb200_fp8mx_config( + num_gpus=8, + expert_model_parallel_size=4, + global_batch_size=8, + hybrid_ep_ranks_per_nvlink_domain=8, + ) + + def nemotron_3_nano_pretrain_8gpu_gb200_bf16_config() -> ConfigContainer: """Nemotron 3 Nano pretrain: 8× GB200, BF16.""" cfg = nemotron_3_nano_pretrain_config() From 580d640d4009775089d342a1aa827112de4baf43 Mon Sep 17 00:00:00 2001 From: Ravi Ghadia Date: Mon, 3 Aug 2026 19:51:57 -0700 Subject: [PATCH 2/3] refactor(perf): remove unused ultra pretrain config and update related documentation - Removed the `nemotron_3_ultra_pretrain_8gpu_gb200_fp8mx_config` function as it was only used for toy runs. - Updated the `nemotron_3_ultra_pretrain_256gpu_gb200_fp8mx_config` function to clarify configuration details and adjusted the minimum offloaded tensor size to match the config for NT3 Ultra MXFP8 on 256xGB200. - Cleaned up comments and improved documentation for clarity. Signed-off-by: Ravi Ghadia --- .../bridge/perf_recipes/nemotronh/__init__.py | 1 - .../perf_recipes/nemotronh/gb200/nemotronh.py | 105 ++++-------------- 2 files changed, 20 insertions(+), 86 deletions(-) diff --git a/src/megatron/bridge/perf_recipes/nemotronh/__init__.py b/src/megatron/bridge/perf_recipes/nemotronh/__init__.py index ff1661d64b..785860600b 100644 --- a/src/megatron/bridge/perf_recipes/nemotronh/__init__.py +++ b/src/megatron/bridge/perf_recipes/nemotronh/__init__.py @@ -29,7 +29,6 @@ nemotron_3_super_pretrain_64gpu_gb200_bf16_config, nemotron_3_super_pretrain_64gpu_gb200_fp8mx_config, nemotron_3_super_pretrain_64gpu_gb200_nvfp4_config, - nemotron_3_ultra_pretrain_8gpu_gb200_fp8mx_config, nemotron_3_ultra_pretrain_256gpu_gb200_fp8mx_config, nemotronh_56b_pretrain_64gpu_gb200_fp8cs_config, ) diff --git a/src/megatron/bridge/perf_recipes/nemotronh/gb200/nemotronh.py b/src/megatron/bridge/perf_recipes/nemotronh/gb200/nemotronh.py index 3bc921785d..c9452d4683 100644 --- a/src/megatron/bridge/perf_recipes/nemotronh/gb200/nemotronh.py +++ b/src/megatron/bridge/perf_recipes/nemotronh/gb200/nemotronh.py @@ -215,21 +215,19 @@ def nemotron_3_super_pretrain_64gpu_gb200_nvfp4_config() -> ConfigContainer: return cfg -def _nemotron_3_ultra_gb200_fp8mx_config( - *, - num_gpus: int, - expert_model_parallel_size: int, - global_batch_size: int, - hybrid_ep_ranks_per_nvlink_domain: int, -) -> ConfigContainer: - """Shared builder for Nemotron 3 Ultra GB200 MXFP8 Megatron-FSDP perf recipes. - - Same Megatron-FSDP (HSDP) layout as the GB300 recipe, but with TP2 + sequence - parallelism instead of TP1. GB200 carries less HBM per GPU, and TP alone would not - help much: it shards the internal GEMM dimensions but leaves the residual stream and - the norm inputs/outputs replicated. Sequence parallelism is what shards those, so the - two are enabled together. +def nemotron_3_ultra_pretrain_256gpu_gb200_fp8mx_config() -> ConfigContainer: + """Nemotron 3 Ultra (550B-A55B LatentMoE) pretrain: 256× GB200, MXFP8, Megatron-FSDP (HSDP). + + TP2 + SP (due to smaller GB200 HBM) / PP1 / CP1 / EP64 / ETP1, GBS 256 / MBS 1, seq 8192, BF16 + MXFP8 mixed + precision, HybridEP flex dispatcher, CuteDSL fused grouped MLP, selective recompute + + fine-grained activation offload of the expert MLP, MTP=2. """ + + num_gpus = 256 + expert_model_parallel_size = 64 + global_batch_size = 256 + hybrid_ep_ranks_per_nvlink_domain = 64 + cfg = nemotron_3_ultra_pretrain_config() cfg.mixed_precision = _perf_precision("fp8_mx") @@ -248,9 +246,9 @@ def _nemotron_3_ultra_gb200_fp8mx_config( cfg.model.pipeline_model_parallel_layout = None cfg.model.seq_length = 8192 - # Only tensors larger than 500MB are offloaded, which - # approximates offloading the moe_act input for seq 8192 / MBS 1. - cfg.model.min_offloaded_tensor_size = 500_000_000 + # Only tensors larger than 350M elements are offloaded, which + # approximates offloading the moe_act (pre-activation input) for seq 8192/2 (due to SP) / MBS 1. + cfg.model.min_offloaded_tensor_size = 350_000_000 # MXFP8 requires router padding for quantization. cfg.model.moe_router_padding_for_quantization = True @@ -261,7 +259,7 @@ def _nemotron_3_ultra_gb200_fp8mx_config( # Fine-grained activation offloading. Requires NVTE_CPU_OFFLOAD_V1=1 in the # launch environment (set in this recipe's env_vars). - # NOTE: also requires setting the min_offloaded_tensor_size to avoid CPU OOM issues + # NOTE: also requires setting the min_offloaded_tensor_size to selectively offload moe_act of the fused_group_mlp, to avoid CPU OOM issues cfg.model.fine_grained_activation_offloading = True cfg.model.offload_modules = ["fused_group_mlp"] cfg.model.fine_grained_offloading_max_inflight_offloads = 1 @@ -271,16 +269,6 @@ def _nemotron_3_ultra_gb200_fp8mx_config( cfg.model.recompute_granularity = "selective" cfg.model.recompute_modules = ["moe_act"] - # full iteration cuda graph - cfg.model.cuda_graph_impl = "full_iteration" # "transformer_engine" - cfg.model.cuda_graph_scope = [] # ["attn", "mamba", "moe_router", "moe_preprocess"] - cfg.ddp.megatron_fsdp_cuda_graph_mode = True - cfg.ddp.fsdp_all_gather_in_start_param_sync = False - cfg.model.moe_expert_rank_capacity_factor = 1.2 - # cfg.model.moe_paged_stash = True - # cfg.model.moe_paged_stash_buffer_size_factor_cpu = 1.0 - # cfg.model.moe_paged_stash_buffer_size_factor_cuda = 1.2 - # Keep process settings next to the recipe so users can see the exact benchmark environment. cfg.env_vars = { **COMMON_PERF_ENV_VARS, @@ -288,23 +276,14 @@ def _nemotron_3_ultra_gb200_fp8mx_config( "CUDA_DEVICE_MAX_CONNECTIONS": 32, # CUDA graph and allocator behavior for this recipe. "NCCL_GRAPH_REGISTER": 0, - # graph_capture_record_stream_reuse is required here, not an - # optimization: fine_grained_activation_offloading calls record_stream() - # on every D2H staging tensor, and cudaEventQuery is illegal during graph - # capture, so without this the allocator cannot recycle any of them until - # capture ends. Offloading then frees nothing on the capture step and - # reserved memory overshoots the eager peak (OOM at 256 GPUs / EP64). The - # option uses the captured DAG topology instead of events, which makes - # the one-time capture slower. TORCH_NCCL_AVOID_RECORD_STREAMS below - # covers only the NCCL buffers, not the offload path. + # TODO: graph_capture_record_stream_reuse might be potentially useful + # when enabling CG, because it allows the freed-up memory buffers of the offloaded tensors + # to be reused during the CG capture, thus keeping the peak memory usage lower. "PYTORCH_CUDA_ALLOC_CONF": ("expandable_segments:True,graph_capture_record_stream_reuse:True"), "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # HybridEP topology for the target system. The per-domain rank count has to track - # EP: HybridEP splits the all-to-all into an intra-domain and an inter-domain leg - # using this value, so leaving it at the 256-GPU 64 on a smaller job would - # describe a domain larger than the job itself. + # HybridEP topology for the target system. "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": hybrid_ep_ranks_per_nvlink_domain, "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, "NVLINK_DOMAIN_SIZE": 72, @@ -323,50 +302,6 @@ def _nemotron_3_ultra_gb200_fp8mx_config( return cfg -def nemotron_3_ultra_pretrain_256gpu_gb200_fp8mx_config() -> ConfigContainer: - """Nemotron 3 Ultra (550B-A55B LatentMoE) pretrain: 256× GB200, MXFP8, Megatron-FSDP (HSDP). - - TP2 + SP / PP1 / CP1 / EP64 / ETP1, GBS 256 / MBS 1, seq 8192, BF16 + MXFP8 mixed - precision, HybridEP flex dispatcher, CuteDSL fused grouped MLP, selective recompute + - fine-grained activation offload of the expert MLP, MTP=2. - - This is the GB300 recipe's layout with TP raised from 1 to 2 and sequence parallelism - turned on, to fit GB200's smaller HBM. Megatron-FSDP shards params, grads and - optimizer state within each 64-GPU NVLink domain and replicates optimizer-sharded - across the four domains. - """ - return _nemotron_3_ultra_gb200_fp8mx_config( - num_gpus=256, - expert_model_parallel_size=64, - global_batch_size=256, - hybrid_ep_ranks_per_nvlink_domain=64, - ) - - -def nemotron_3_ultra_pretrain_8gpu_gb200_fp8mx_config() -> ConfigContainer: - """Nemotron 3 Ultra pipeclean: 8× GB200, MXFP8, Megatron-FSDP. - - Shrunk-grid counterpart of the 256-GPU recipe, for validating that the full feature - stack launches and trains before committing a full-scale job. EP drops 64 -> 4 so the - expert grid (``ETP1 x EP4`` = 4) divides an 8-GPU world, and GBS drops 256 -> 8. With - TP2 the batch splits over DP = 8 / 2 = 4, so GBS 8 / MBS 1 runs 2 microbatches. - - Below one NVLink domain ``_apply_nemotron_3_ultra_fsdp_hsdp`` resolves - ``num_distributed_optimizer_instances`` to 1, which turns HSDP off and correspondingly - sets ``outer_dp_sharding_strategy`` to ``no_shard``. - - This recipe only rescales parallelism: the model is still full-size Nemotron 3 Ultra, - which does not fit on 8 GPUs. Pair it with model-shrinking overrides (see - ``perf_bash_scripts/nt3_ultra_gb200/toy_run_perf_test_nemotron_3_ultra_gb200_fp8mx.sh``). - """ - return _nemotron_3_ultra_gb200_fp8mx_config( - num_gpus=8, - expert_model_parallel_size=4, - global_batch_size=8, - hybrid_ep_ranks_per_nvlink_domain=8, - ) - - def nemotron_3_nano_pretrain_8gpu_gb200_bf16_config() -> ConfigContainer: """Nemotron 3 Nano pretrain: 8× GB200, BF16.""" cfg = nemotron_3_nano_pretrain_config() From ab1ae70726973d31f67e71e9fab590d23369572b Mon Sep 17 00:00:00 2001 From: Malay Nagda Date: Tue, 4 Aug 2026 15:00:45 +0530 Subject: [PATCH 3/3] test(recipe): allow local constants in environment checks Signed-off-by: Malay Nagda --- .../recipes/test_perf_recipe_environment.py | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/tests/unit_tests/recipes/test_perf_recipe_environment.py b/tests/unit_tests/recipes/test_perf_recipe_environment.py index 1a6cf1c6a9..201184a38a 100644 --- a/tests/unit_tests/recipes/test_perf_recipe_environment.py +++ b/tests/unit_tests/recipes/test_perf_recipe_environment.py @@ -72,6 +72,15 @@ def _function(path: Path, function_name: str) -> ast.FunctionDef: def _explicit_environment(path: Path, function_name: str) -> dict[str, str | int | float | bool]: """Read the literal env mapping written in a flat recipe builder.""" function = _function(path, function_name) + local_constants = { + node.targets[0].id: node.value.value + for node in function.body + if isinstance(node, ast.Assign) + and len(node.targets) == 1 + and isinstance(node.targets[0], ast.Name) + and isinstance(node.value, ast.Constant) + and isinstance(node.value.value, (str, int, float, bool)) + } assignments = [ node for node in function.body @@ -93,7 +102,12 @@ def _explicit_environment(path: Path, function_name: str) -> dict[str, str | int assert isinstance(value, ast.Name) and value.id == "COMMON_PERF_ENV_VARS" common_expansions += 1 continue - result[ast.literal_eval(key)] = ast.literal_eval(value) + if isinstance(value, ast.Name): + assert value.id in local_constants + env_value = local_constants[value.id] + else: + env_value = ast.literal_eval(value) + result[ast.literal_eval(key)] = env_value assert common_expansions == 1 return result