From ab5f6aa6a7f63e5f2a4377b7ad4d8f223aaaf21c Mon Sep 17 00:00:00 2001 From: Kamran Jafari Date: Mon, 12 Oct 2026 11:44:59 -0700 Subject: [PATCH] fix(recipe)!: run Nemotron 3 Ultra H100 SFT on the documented 48-node layout (#6376) Signed-off-by: Kamran Jafari Signed-off-by: Chen Cui Co-authored-by: Claude Opus 5.5 Co-authored-by: Chen Cui Signed-off-by: NeMo Bot --- .../model-verification/catalog-v1.json | 10 +-- .../pages/models/nemotron/nemotron3-ultra.mdx | 4 +- docs/models/nemotron/nemotron3-ultra.md | 4 +- .../nemotron-3-ultra-550b-a55b/card.yaml | 6 +- .../nemotron/nemotron_3_ultra/README.md | 19 ++-- .../nemotron/nemotron_3_ultra/slurm_sft.sh | 15 ++-- .../bridge/recipes/nemotronh/h100/__init__.py | 2 +- .../nemotronh/h100/nemotron_3_ultra.py | 61 +++++++++++-- .../recipes/nemotronh/nemotron_3_ultra.py | 2 +- .../nemotronh/test_nemotron_3_ultra.py | 88 +++++++++++++++++-- tests/unit_tests/recipes/recipe_test_utils.py | 9 ++ .../recipes/test_nemotronh_recipes.py | 23 ++++- 12 files changed, 202 insertions(+), 41 deletions(-) diff --git a/docs/_static/model-verification/catalog-v1.json b/docs/_static/model-verification/catalog-v1.json index 23cbe9b650..6d7be3736a 100644 --- a/docs/_static/model-verification/catalog-v1.json +++ b/docs/_static/model-verification/catalog-v1.json @@ -5393,13 +5393,13 @@ "command_topologies": [ { "gpus_per_node": "8", - "nodes": "24", - "recipe": "nemotron_3_ultra_sft_192gpu_h100_bf16_openmathinstruct2_packed_config", - "total_gpus": 192 + "nodes": "48", + "recipe": "nemotron_3_ultra_sft_384gpu_h100_bf16_openmathinstruct2_packed_config", + "total_gpus": 384 } ], "commands": [ - "./scripts/training/train.sh --nodes 24 --gpus-per-node 8 --recipe nemotron_3_ultra_sft_192gpu_h100_bf16_openmathinstruct2_packed_config --mode sft --wait --pretrained_checkpoint work/model-verification/nemotron-3-ultra-550b-a55b/imported-megatron/iter_0000000 --max_steps 100 --warmup_iters 10 scheduler.lr_decay_iters=100 'dataset.hf_dataset.split=\"train_1M[:10%]\"' 'dataset.hf_dataset.load_kwargs={revision:\"469216e3f46f4dacf476b382e192485ea51a143e\"}' dataset.hf_output_root=work/data/openmathinstruct2-ultra-4096-v1 dataset.seed=1234 dataset.offline_packing_specs.pad_seq_to_mult=2 validation.eval_iters=0 validation.eval_interval=0 --save_dir work/model-verification/nemotron-3-ultra-550b-a55b/sft-checkpoints --save_interval 100 ddp.check_for_nan_in_grad=true ddp.check_for_large_grads=true rerun_state_machine.check_for_nan_in_loss=true logger.log_interval=1 logger.log_throughput=true logger.save_config_filepath=work/model-verification/nemotron-3-ultra-550b-a55b/sft-resolved.yaml logger.tensorboard_dir=null" + "./scripts/training/train.sh --nodes 48 --gpus-per-node 8 --recipe nemotron_3_ultra_sft_384gpu_h100_bf16_openmathinstruct2_packed_config --mode sft --wait --pretrained_checkpoint work/model-verification/nemotron-3-ultra-550b-a55b/imported-megatron/iter_0000000 --max_steps 100 --warmup_iters 10 scheduler.lr_decay_iters=100 'dataset.hf_dataset.split=\"train_1M[:10%]\"' 'dataset.hf_dataset.load_kwargs={revision:\"469216e3f46f4dacf476b382e192485ea51a143e\"}' dataset.hf_output_root=work/data/openmathinstruct2-ultra-4096-v1 dataset.seed=1234 dataset.offline_packing_specs.pad_seq_to_mult=2 validation.eval_iters=0 validation.eval_interval=0 --save_dir work/model-verification/nemotron-3-ultra-550b-a55b/sft-checkpoints --save_interval 100 ddp.check_for_nan_in_grad=true ddp.check_for_large_grads=true rerun_state_machine.check_for_nan_in_loss=true logger.log_interval=1 logger.log_throughput=true logger.save_config_filepath=work/model-verification/nemotron-3-ultra-550b-a55b/sft-resolved.yaml logger.tensorboard_dir=null" ], "dimensions": {}, "enabled_features": { @@ -5407,7 +5407,7 @@ "sequence_packing": "offline" }, "entry_id": "nemotron-3-ultra-550b-a55b-sft-h100", - "expected_result": "The public 192-H100 packed OpenMathInstruct-2 full-SFT recipe must use the pinned dataset revision and bounded train_1M first-ten-percent selection, complete 100 optimizer steps at sequence length 4096 and GBS/MBS 128/1, and preserve pad_seq_to_mult 2. This published Ultra contract is support verification rather than the 8K/GBS8 qwen3_30b_a3b_convergence_v2 cohort. Loss must remain finite with no skipped or NaN iterations, all five metrics, a persisted post-setup configuration, audited supervised-token counts, and a complete reloadable final checkpoint.\n", + "expected_result": "The public 384-H100 packed OpenMathInstruct-2 full-SFT recipe must use the pinned dataset revision and bounded train_1M first-ten-percent selection, complete 100 optimizer steps at sequence length 4096 and GBS/MBS 128/1, and preserve pad_seq_to_mult 2. This published Ultra contract is support verification rather than the 8K/GBS8 qwen3_30b_a3b_convergence_v2 cohort. Loss must remain finite with no skipped or NaN iterations, all five metrics, a persisted post-setup configuration, audited supervised-token counts, and a complete reloadable final checkpoint.\n", "hardware": "H100", "last_verified": null, "metrics": { diff --git a/docs/fern/versions/nightly/pages/models/nemotron/nemotron3-ultra.mdx b/docs/fern/versions/nightly/pages/models/nemotron/nemotron3-ultra.mdx index ca05305440..6c48db4319 100644 --- a/docs/fern/versions/nightly/pages/models/nemotron/nemotron3-ultra.mdx +++ b/docs/fern/versions/nightly/pages/models/nemotron/nemotron3-ultra.mdx @@ -408,12 +408,12 @@ Choose a workflow, precision, and exact recorded combination. The command and ex Command - 
./scripts/training/train.sh --nodes 24 --gpus-per-node 8 --recipe nemotron_3_ultra_sft_192gpu_h100_bf16_openmathinstruct2_packed_config --mode sft --wait --pretrained_checkpoint work/model-verification/nemotron-3-ultra-550b-a55b/imported-megatron/iter_0000000 --max_steps 100 --warmup_iters 10 scheduler.lr_decay_iters=100 'dataset.hf_dataset.split="train_1M[:10%]"' 'dataset.hf_dataset.load_kwargs={revision:"469216e3f46f4dacf476b382e192485ea51a143e"}' dataset.hf_output_root=work/data/openmathinstruct2-ultra-4096-v1 dataset.seed=1234 dataset.offline_packing_specs.pad_seq_to_mult=2 validation.eval_iters=0 validation.eval_interval=0 --save_dir work/model-verification/nemotron-3-ultra-550b-a55b/sft-checkpoints --save_interval 100 ddp.check_for_nan_in_grad=true ddp.check_for_large_grads=true rerun_state_machine.check_for_nan_in_loss=true logger.log_interval=1 logger.log_throughput=true logger.save_config_filepath=work/model-verification/nemotron-3-ultra-550b-a55b/sft-resolved.yaml logger.tensorboard_dir=null
+ 
./scripts/training/train.sh --nodes 48 --gpus-per-node 8 --recipe nemotron_3_ultra_sft_384gpu_h100_bf16_openmathinstruct2_packed_config --mode sft --wait --pretrained_checkpoint work/model-verification/nemotron-3-ultra-550b-a55b/imported-megatron/iter_0000000 --max_steps 100 --warmup_iters 10 scheduler.lr_decay_iters=100 'dataset.hf_dataset.split="train_1M[:10%]"' 'dataset.hf_dataset.load_kwargs={revision:"469216e3f46f4dacf476b382e192485ea51a143e"}' dataset.hf_output_root=work/data/openmathinstruct2-ultra-4096-v1 dataset.seed=1234 dataset.offline_packing_specs.pad_seq_to_mult=2 validation.eval_iters=0 validation.eval_interval=0 --save_dir work/model-verification/nemotron-3-ultra-550b-a55b/sft-checkpoints --save_interval 100 ddp.check_for_nan_in_grad=true ddp.check_for_large_grads=true rerun_state_machine.check_for_nan_in_loss=true logger.log_interval=1 logger.log_throughput=true logger.save_config_filepath=work/model-verification/nemotron-3-ultra-550b-a55b/sft-resolved.yaml logger.tensorboard_dir=null
Expected result
-

The public 192-H100 packed OpenMathInstruct-2 full-SFT recipe must use the pinned dataset revision and bounded train_1M first-ten-percent selection, complete 100 optimizer steps at sequence length 4096 and GBS/MBS 128/1, and preserve pad_seq_to_mult 2. This published Ultra contract is support verification rather than the 8K/GBS8 qwen3_30b_a3b_convergence_v2 cohort. Loss must remain finite with no skipped or NaN iterations, all five metrics, a persisted post-setup configuration, audited supervised-token counts, and a complete reloadable final checkpoint.

+

The public 384-H100 packed OpenMathInstruct-2 full-SFT recipe must use the pinned dataset revision and bounded train_1M first-ten-percent selection, complete 100 optimizer steps at sequence length 4096 and GBS/MBS 128/1, and preserve pad_seq_to_mult 2. This published Ultra contract is support verification rather than the 8K/GBS8 qwen3_30b_a3b_convergence_v2 cohort. Loss must remain finite with no skipped or NaN iterations, all five metrics, a persisted post-setup configuration, audited supervised-token counts, and a complete reloadable final checkpoint.

diff --git a/docs/models/nemotron/nemotron3-ultra.md b/docs/models/nemotron/nemotron3-ultra.md index e04d66cdf6..7205a7cfd4 100644 --- a/docs/models/nemotron/nemotron3-ultra.md +++ b/docs/models/nemotron/nemotron3-ultra.md @@ -414,12 +414,12 @@ Choose a workflow, precision, and exact recorded combination. The command and ex Command - 
./scripts/training/train.sh --nodes 24 --gpus-per-node 8 --recipe nemotron_3_ultra_sft_192gpu_h100_bf16_openmathinstruct2_packed_config --mode sft --wait --pretrained_checkpoint work/model-verification/nemotron-3-ultra-550b-a55b/imported-megatron/iter_0000000 --max_steps 100 --warmup_iters 10 scheduler.lr_decay_iters=100 'dataset.hf_dataset.split="train_1M[:10%]"' 'dataset.hf_dataset.load_kwargs={revision:"469216e3f46f4dacf476b382e192485ea51a143e"}' dataset.hf_output_root=work/data/openmathinstruct2-ultra-4096-v1 dataset.seed=1234 dataset.offline_packing_specs.pad_seq_to_mult=2 validation.eval_iters=0 validation.eval_interval=0 --save_dir work/model-verification/nemotron-3-ultra-550b-a55b/sft-checkpoints --save_interval 100 ddp.check_for_nan_in_grad=true ddp.check_for_large_grads=true rerun_state_machine.check_for_nan_in_loss=true logger.log_interval=1 logger.log_throughput=true logger.save_config_filepath=work/model-verification/nemotron-3-ultra-550b-a55b/sft-resolved.yaml logger.tensorboard_dir=null
+ 
./scripts/training/train.sh --nodes 48 --gpus-per-node 8 --recipe nemotron_3_ultra_sft_384gpu_h100_bf16_openmathinstruct2_packed_config --mode sft --wait --pretrained_checkpoint work/model-verification/nemotron-3-ultra-550b-a55b/imported-megatron/iter_0000000 --max_steps 100 --warmup_iters 10 scheduler.lr_decay_iters=100 'dataset.hf_dataset.split="train_1M[:10%]"' 'dataset.hf_dataset.load_kwargs={revision:"469216e3f46f4dacf476b382e192485ea51a143e"}' dataset.hf_output_root=work/data/openmathinstruct2-ultra-4096-v1 dataset.seed=1234 dataset.offline_packing_specs.pad_seq_to_mult=2 validation.eval_iters=0 validation.eval_interval=0 --save_dir work/model-verification/nemotron-3-ultra-550b-a55b/sft-checkpoints --save_interval 100 ddp.check_for_nan_in_grad=true ddp.check_for_large_grads=true rerun_state_machine.check_for_nan_in_loss=true logger.log_interval=1 logger.log_throughput=true logger.save_config_filepath=work/model-verification/nemotron-3-ultra-550b-a55b/sft-resolved.yaml logger.tensorboard_dir=null
Expected result
-

The public 192-H100 packed OpenMathInstruct-2 full-SFT recipe must use the pinned dataset revision and bounded train_1M first-ten-percent selection, complete 100 optimizer steps at sequence length 4096 and GBS/MBS 128/1, and preserve pad_seq_to_mult 2. This published Ultra contract is support verification rather than the 8K/GBS8 qwen3_30b_a3b_convergence_v2 cohort. Loss must remain finite with no skipped or NaN iterations, all five metrics, a persisted post-setup configuration, audited supervised-token counts, and a complete reloadable final checkpoint. +

The public 384-H100 packed OpenMathInstruct-2 full-SFT recipe must use the pinned dataset revision and bounded train_1M first-ten-percent selection, complete 100 optimizer steps at sequence length 4096 and GBS/MBS 128/1, and preserve pad_seq_to_mult 2. This published Ultra contract is support verification rather than the 8K/GBS8 qwen3_30b_a3b_convergence_v2 cohort. Loss must remain finite with no skipped or NaN iterations, all five metrics, a persisted post-setup configuration, audited supervised-token counts, and a complete reloadable final checkpoint.

diff --git a/examples/model_verification_cards/nemotron-3-ultra-550b-a55b/card.yaml b/examples/model_verification_cards/nemotron-3-ultra-550b-a55b/card.yaml index ad11a17604..7e920f2848 100644 --- a/examples/model_verification_cards/nemotron-3-ultra-550b-a55b/card.yaml +++ b/examples/model_verification_cards/nemotron-3-ultra-550b-a55b/card.yaml @@ -269,8 +269,8 @@ items: sequence_packing: offline moe_dispatcher: hybridep command: > - ./scripts/training/train.sh --nodes 24 --gpus-per-node 8 - --recipe nemotron_3_ultra_sft_192gpu_h100_bf16_openmathinstruct2_packed_config + ./scripts/training/train.sh --nodes 48 --gpus-per-node 8 + --recipe nemotron_3_ultra_sft_384gpu_h100_bf16_openmathinstruct2_packed_config --mode sft --wait --pretrained_checkpoint work/model-verification/nemotron-3-ultra-550b-a55b/imported-megatron/iter_0000000 --max_steps 100 @@ -296,7 +296,7 @@ items: last_10_steps_model_tflops_per_gpu_avg: null last_10_steps_tokens_per_second_per_gpu_avg: null expected_result: > - The public 192-H100 packed OpenMathInstruct-2 full-SFT recipe must use + The public 384-H100 packed OpenMathInstruct-2 full-SFT recipe must use the pinned dataset revision and bounded train_1M first-ten-percent selection, complete 100 optimizer steps at sequence length 4096 and GBS/MBS 128/1, and preserve pad_seq_to_mult 2. This published Ultra diff --git a/examples/models/nemotron/nemotron_3_ultra/README.md b/examples/models/nemotron/nemotron_3_ultra/README.md index 03c4d00873..c8beacc1b3 100644 --- a/examples/models/nemotron/nemotron_3_ultra/README.md +++ b/examples/models/nemotron/nemotron_3_ultra/README.md @@ -48,7 +48,7 @@ environment variables to the GB200 values in this table. | Checkpoint import | 1 node with `convert.sh --device cpu` | 6 nodes with `convert.sh --device gpu`, `TP=1 PP=6 EP=4` | | Base inference | 4 nodes, `TP=1 PP=4 EP=8`, `KV_CACHE_BUFFER_SIZE_GB=4` | 3 nodes, `TP=1 PP=3 EP=4` | | DCLM pretraining | 48 nodes, `TP=4 PP=12 EP=16`, full uniform recompute with `RECOMPUTE_GRANULARITY=full RECOMPUTE_METHOD=uniform RECOMPUTE_NUM_LAYERS=1 RECOMPUTE_MODULES=""` | 24 nodes, `TP=2 PP=3 EP=32`, selective recompute on `moe+layernorm+core_attn+moe_act+mlp+shared_experts` | -| OpenMath SFT | 48 nodes, `TP=2 PP=12 EP=16`, full uniform recompute with `RECOMPUTE_GRANULARITY=full RECOMPUTE_METHOD=uniform RECOMPUTE_NUM_LAYERS=1 RECOMPUTE_MODULES=""` | 48 nodes, `TP=2 PP=3 EP=32`, selective recompute on `moe+layernorm+core_attn+moe_act` | +| OpenMath SFT | 48 nodes, `TP=2 PP=12 EP=16`, the recipe's selective recompute on `moe+core_attn+moe_act` | 48 nodes, `TP=2 PP=3 EP=32`, the recipe's selective recompute on `moe+core_attn+moe_act` | | OpenMath PEFT | 4 nodes, `TP=2 PP=4 EP=8`, selective recompute on `moe+layernorm+core_attn+moe_act+mlp+shared_experts` | 4 nodes, `TP=2 PP=1 EP=16`, selective recompute on `moe+layernorm+core_attn+moe_act` | These are bring-up and convergence starting points, not universal optima. @@ -173,17 +173,22 @@ Current OpenMath starting points are: - PEFT: 4 nodes, `TP=2 PP=4 EP=8`, selective recompute on `moe+layernorm+core_attn+moe_act+mlp+shared_experts`. -- Full SFT: 48 nodes, `TP=2 PP=12 EP=16`, full uniform recompute with - `RECOMPUTE_GRANULARITY=full RECOMPUTE_METHOD=uniform - RECOMPUTE_NUM_LAYERS=1 RECOMPUTE_MODULES=""`. This is the current H100 - starting point for 4096-token packed OpenMath SFT. +- Full SFT: 48 nodes, `TP=2 PP=12 EP=16`, the recipe's selective recompute on + `moe+core_attn+moe_act`. These are the recipe defaults and the current H100 + starting point for 4096-token packed OpenMath SFT. On Transformer Engine 2.18 + and later, a TE norm that is a pipeline stage's first backward operation fails + with `invalid device context`, so the recipe does not recompute the pre-MLP + norms and sets `NVTE_NORM_FWD_USE_CUDNN=1` and `NVTE_NORM_BWD_USE_CUDNN=1` to + use cuDNN normalization kernels. The recipe places 5 decoder layers on the + last pipeline stage, which also holds the output layer and MTP, and 9 to 11 on + the others; when changing `PP`, also override `model.hybrid_layer_pattern`. For 4xGB200 nodes: - PEFT: 4 nodes, `TP=2 PP=1 EP=16`, selective recompute on `moe+layernorm+core_attn+moe_act`. -- Full SFT: 48 nodes, `TP=2 PP=3 EP=32`, selective recompute on - `moe+layernorm+core_attn+moe_act`. +- Full SFT: 48 nodes, `TP=2 PP=3 EP=32`, the recipe's selective recompute on + `moe+core_attn+moe_act`. Advanced VPP, pipeline-layout, and recompute sweeps are intentionally left out of these starter scripts; add those overrides only for targeted performance diff --git a/examples/models/nemotron/nemotron_3_ultra/slurm_sft.sh b/examples/models/nemotron/nemotron_3_ultra/slurm_sft.sh index 51f6251867..6b92dd5928 100755 --- a/examples/models/nemotron/nemotron_3_ultra/slurm_sft.sh +++ b/examples/models/nemotron/nemotron_3_ultra/slurm_sft.sh @@ -91,13 +91,14 @@ if [ -z "$PAD_SEQ_TO_MULT" ]; then fi fi -RECOMPUTE_GRANULARITY=${RECOMPUTE_GRANULARITY:-full} -RECOMPUTE_METHOD=${RECOMPUTE_METHOD:-uniform} +# Leave these empty to use the recipe's selective recompute; set them to override it. +RECOMPUTE_GRANULARITY=${RECOMPUTE_GRANULARITY:-} +RECOMPUTE_METHOD=${RECOMPUTE_METHOD:-} if [ -z "${RECOMPUTE_MODULES+x}" ]; then RECOMPUTE_MODULES="" fi -RECOMPUTE_NUM_LAYERS=${RECOMPUTE_NUM_LAYERS:-1} -RECOMPUTE_TAG=${RECOMPUTE_TAG:-recompute_full_uniform1} +RECOMPUTE_NUM_LAYERS=${RECOMPUTE_NUM_LAYERS:-} +RECOMPUTE_TAG=${RECOMPUTE_TAG:-recompute_${RECOMPUTE_GRANULARITY:-recipe}} WANDB_ENTITY=${WANDB_ENTITY:-nvidia-nemo-fw-public} WANDB_PROJECT=${WANDB_PROJECT:-megatron-bridge-nemotron-ultra} @@ -171,9 +172,11 @@ CLI_OVERRIDES="\ dataset.seq_length=${SEQ_LENGTH} \ dataset.offline_packing_specs.packed_sequence_size=${SEQ_LENGTH} \ dataset.offline_packing_specs.pad_seq_to_mult=${PAD_SEQ_TO_MULT} \ - model.recompute_granularity=${RECOMPUTE_GRANULARITY} \ dist.distributed_timeout_minutes=90" +if [ -n "$RECOMPUTE_GRANULARITY" ]; then + CLI_OVERRIDES="${CLI_OVERRIDES} model.recompute_granularity=${RECOMPUTE_GRANULARITY}" +fi if [ -n "$RECOMPUTE_METHOD" ]; then CLI_OVERRIDES="${CLI_OVERRIDES} model.recompute_method=${RECOMPUTE_METHOD}" fi @@ -205,7 +208,7 @@ echo "GPUs/node: ${GPUS_PER_NODE}" echo "Recipe: ${RECIPE_NAME}" echo "Parallelism: TP=${TP} PP=${PP} EP=${EP} ETP=${ETP} CP=${CP} SP=${SP}" echo "Packed pad_seq_to_mult: ${PAD_SEQ_TO_MULT}" -echo "Recompute: ${RECOMPUTE_GRANULARITY} ${RECOMPUTE_METHOD:-} ${RECOMPUTE_MODULES}" +echo "Recompute: ${RECOMPUTE_GRANULARITY:-recipe default} ${RECOMPUTE_METHOD:-} ${RECOMPUTE_MODULES}" echo "Save dir: ${SAVE_DIR}" echo "W&B: ${WANDB_ENTITY}/${WANDB_PROJECT} (${WANDB_MODE})" echo "======================================" diff --git a/src/megatron/bridge/recipes/nemotronh/h100/__init__.py b/src/megatron/bridge/recipes/nemotronh/h100/__init__.py index 260b649a75..c0512ca195 100644 --- a/src/megatron/bridge/recipes/nemotronh/h100/__init__.py +++ b/src/megatron/bridge/recipes/nemotronh/h100/__init__.py @@ -38,5 +38,5 @@ "nemotron_3_ultra_peft_32gpu_h100_bf16_openmathinstruct2_packed_config", "nemotron_3_ultra_pretrain_24gpu_h100_bf16_config", "nemotron_3_ultra_pretrain_256gpu_h100_bf16_fsdp_config", - "nemotron_3_ultra_sft_192gpu_h100_bf16_openmathinstruct2_packed_config", + "nemotron_3_ultra_sft_384gpu_h100_bf16_openmathinstruct2_packed_config", ] diff --git a/src/megatron/bridge/recipes/nemotronh/h100/nemotron_3_ultra.py b/src/megatron/bridge/recipes/nemotronh/h100/nemotron_3_ultra.py index 02d33e7712..cb660f5494 100644 --- a/src/megatron/bridge/recipes/nemotronh/h100/nemotron_3_ultra.py +++ b/src/megatron/bridge/recipes/nemotronh/h100/nemotron_3_ultra.py @@ -27,6 +27,34 @@ NEMOTRON_3_ULTRA_TOKENIZER_NAME = "nvidia--NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16" NEMOTRON_3_ULTRA_PRETRAIN_SEQ_LENGTH = 8192 NEMOTRON_3_ULTRA_OPENMATHINSTRUCT2_SEQ_LENGTH = 4096 +# Decoder layers on each of the 12 pipeline stages of the 384-GPU SFT recipe. The last stage also +# holds the output layer and the two-depth MTP block, so it takes 5 layers and stages 8-10 take +# the extra ones. +NEMOTRON_3_ULTRA_SFT_LAYERS_PER_STAGE = (9, 9, 9, 9, 9, 9, 9, 9, 10, 10, 11, 5) + + +def _split_hybrid_layer_pattern(pattern: str, layers_per_stage: tuple[int, ...]) -> str: + """Insert pipeline stage separators into a hybrid layer pattern. + + Args: + pattern: Hybrid layer pattern without stage separators, optionally followed by an MTP suffix after "/". + layers_per_stage: Number of decoder layers on each pipeline stage, in order. + + Returns: + The decoder pattern with "|" between stages, followed by the unchanged MTP suffix. + + Raises: + ValueError: If the pattern already has stage separators or the stage sizes do not cover it exactly. + """ + main, separator, mtp = pattern.partition("/") + if "|" in main or any(count < 1 for count in layers_per_stage) or sum(layers_per_stage) != len(main): + raise ValueError(f"Cannot split a {len(main)}-layer hybrid pattern into stages of {layers_per_stage}.") + stages = [] + start = 0 + for count in layers_per_stage: + stages.append(main[start : start + count]) + start += count + return "|".join(stages) + separator + mtp def nemotron_3_ultra_pretrain_24gpu_h100_bf16_config() -> ConfigContainer: @@ -234,8 +262,11 @@ def nemotron_3_ultra_pretrain_256gpu_h100_bf16_fsdp_config() -> ConfigContainer: return cfg -def nemotron_3_ultra_sft_192gpu_h100_bf16_openmathinstruct2_packed_config() -> ConfigContainer: - """Return a packed OpenMathInstruct-2 full SFT config for Nemotron 3 Ultra. +def nemotron_3_ultra_sft_384gpu_h100_bf16_openmathinstruct2_packed_config() -> ConfigContainer: + """Return a packed OpenMathInstruct-2 full SFT config for Nemotron 3 Ultra on 384 H100 GPUs. + + The layout matches the H100 full-SFT starting point documented in the Ultra + examples: 48 nodes at TP2/PP12/EP16 with selective activation recompute. Returns: Full-parameter SFT configuration for OpenMathInstruct-2. @@ -244,14 +275,17 @@ def nemotron_3_ultra_sft_192gpu_h100_bf16_openmathinstruct2_packed_config() -> C cfg.model = AutoBridge.from_hf_pretrained(NEMOTRON_3_ULTRA_HF_MODEL_ID).to_megatron_provider(load_weights=False) cfg.model.tensor_model_parallel_size = 2 - cfg.model.pipeline_model_parallel_size = 6 + cfg.model.pipeline_model_parallel_size = 12 cfg.model.pipeline_dtype = torch.bfloat16 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.expert_model_parallel_size = 32 + cfg.model.expert_model_parallel_size = 16 cfg.model.pipeline_model_parallel_layout = None + cfg.model.hybrid_layer_pattern = _split_hybrid_layer_pattern( + cfg.model.hybrid_layer_pattern, NEMOTRON_3_ULTRA_SFT_LAYERS_PER_STAGE + ) cfg.model.seq_length = NEMOTRON_3_ULTRA_OPENMATHINSTRUCT2_SEQ_LENGTH cfg.model.apply_rope_fusion = False cfg.model.attention_backend = "fused" @@ -269,10 +303,14 @@ def nemotron_3_ultra_sft_192gpu_h100_bf16_openmathinstruct2_packed_config() -> C cfg.model.mtp_loss_scaling_factor = 0.3 cfg.model.mtp_use_repeated_layer = True cfg.model.use_te_rng_tracker = True + # Selective recompute fits the 48-node layout. The standalone pre-MLP norms are not + # recomputed: on Transformer Engine 2.18 and later, a recomputed TE norm that is a pipeline + # stage's first backward operation fails with "invalid device context" on GPU 0, and most + # stages of this layout end with an MoE layer. cfg.model.recompute_granularity = "selective" cfg.model.recompute_method = None cfg.model.recompute_num_layers = None - cfg.model.recompute_modules = ["moe", "layernorm", "core_attn", "moe_act"] + cfg.model.recompute_modules = ["moe", "core_attn", "moe_act"] cfg.tokenizer.tokenizer_model = NEMOTRON_3_ULTRA_HF_MODEL_ID cfg.dataset = default_openmathinstruct2_config( @@ -317,9 +355,20 @@ def nemotron_3_ultra_sft_192gpu_h100_bf16_openmathinstruct2_packed_config() -> C cfg.ddp.overlap_param_gather = True cfg.ddp.use_distributed_optimizer = True + # The first step pays one-time startup costs (HybridEP setup, kernel compilation) + # stage by stage through the twelve-stage pipeline, so stage 0 can wait longer than + # the 10-minute default timeout for its first backward. + cfg.dist.distributed_timeout_minutes = 90 + # Keep the complete process environment visible on the recipe. cfg.env_vars = { **COMMON_RECIPE_ENV_VARS, + # One hardware queue, as Megatron-Core expects for TP>1 with sequence parallelism on Hopper. + "CUDA_DEVICE_MAX_CONNECTIONS": 1, + # cuDNN normalization kernels bypass TE's NVRTC norm path, which fails with "invalid device + # context" when a norm is the first backward operation on GPU 0 (TE 2.18 and later). + "NVTE_NORM_BWD_USE_CUDNN": 1, + "NVTE_NORM_FWD_USE_CUDNN": 1, # HybridEP topology for this recipe. "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 8, "NVLINK_DOMAIN_SIZE": 8, @@ -440,5 +489,5 @@ def nemotron_3_ultra_peft_32gpu_h100_bf16_openmathinstruct2_packed_config( "nemotron_3_ultra_peft_32gpu_h100_bf16_openmathinstruct2_packed_config", "nemotron_3_ultra_pretrain_24gpu_h100_bf16_config", "nemotron_3_ultra_pretrain_256gpu_h100_bf16_fsdp_config", - "nemotron_3_ultra_sft_192gpu_h100_bf16_openmathinstruct2_packed_config", + "nemotron_3_ultra_sft_384gpu_h100_bf16_openmathinstruct2_packed_config", ] diff --git a/src/megatron/bridge/recipes/nemotronh/nemotron_3_ultra.py b/src/megatron/bridge/recipes/nemotronh/nemotron_3_ultra.py index 4a116b9874..ff71c52c18 100644 --- a/src/megatron/bridge/recipes/nemotronh/nemotron_3_ultra.py +++ b/src/megatron/bridge/recipes/nemotronh/nemotron_3_ultra.py @@ -27,7 +27,7 @@ nemotron_3_ultra_pretrain_24gpu_h100_bf16_config as nemotron_3_ultra_pretrain_config, ) from megatron.bridge.recipes.nemotronh.h100.nemotron_3_ultra import ( - nemotron_3_ultra_sft_192gpu_h100_bf16_openmathinstruct2_packed_config as nemotron_3_ultra_sft_openmathinstruct2_packed_config, + nemotron_3_ultra_sft_384gpu_h100_bf16_openmathinstruct2_packed_config as nemotron_3_ultra_sft_openmathinstruct2_packed_config, ) diff --git a/tests/unit_tests/recipes/nemotronh/test_nemotron_3_ultra.py b/tests/unit_tests/recipes/nemotronh/test_nemotron_3_ultra.py index 4e2e7201d1..28d250f86b 100644 --- a/tests/unit_tests/recipes/nemotronh/test_nemotron_3_ultra.py +++ b/tests/unit_tests/recipes/nemotronh/test_nemotron_3_ultra.py @@ -26,7 +26,9 @@ nemotron_3_ultra_pretrain_256gpu_gb200_bf16_ep16_config, ) from megatron.bridge.recipes.nemotronh.h100.nemotron_3_ultra import ( + _split_hybrid_layer_pattern, nemotron_3_ultra_pretrain_256gpu_h100_bf16_fsdp_config, + nemotron_3_ultra_sft_384gpu_h100_bf16_openmathinstruct2_packed_config, ) from megatron.bridge.recipes.nemotronh.nemotron_3_ultra import ( NEMOTRON_3_ULTRA_TOKENIZER_NAME, @@ -34,7 +36,10 @@ nemotron_3_ultra_pretrain_config, nemotron_3_ultra_sft_openmathinstruct2_packed_config, ) -from tests.unit_tests.recipes.recipe_test_utils import patch_recipe_module_global +from tests.unit_tests.recipes.recipe_test_utils import ( + NEMOTRON_3_ULTRA_HYBRID_LAYER_PATTERN, + patch_recipe_module_global, +) class _FakeUltraProvider: @@ -42,6 +47,7 @@ class _FakeUltraProvider: def __init__(self) -> None: self.vocab_size = 256 + self.hybrid_layer_pattern = NEMOTRON_3_ULTRA_HYBRID_LAYER_PATTERN def finalize(self) -> None: return None @@ -216,12 +222,25 @@ def test_vr200_nvfp4_perf_recipe_uses_nvl72_ultra_topology() -> None: @pytest.mark.unit -def test_openmath_sft_uses_initial_parallelism_values() -> None: +def test_openmath_sft_recipe_name_matches_its_384gpu_layout() -> None: + assert recipes.nemotron_3_ultra_sft_384gpu_h100_bf16_openmathinstruct2_packed_config is ( + nemotron_3_ultra_sft_384gpu_h100_bf16_openmathinstruct2_packed_config + ) + assert nemotron_3_ultra_sft_openmathinstruct2_packed_config is ( + nemotron_3_ultra_sft_384gpu_h100_bf16_openmathinstruct2_packed_config + ) + assert not hasattr(recipes, "nemotron_3_ultra_sft_192gpu_h100_bf16_openmathinstruct2_packed_config") + + +@pytest.mark.unit +def test_openmath_sft_uses_48_node_layout() -> None: cfg = nemotron_3_ultra_sft_openmathinstruct2_packed_config() assert cfg.model.tensor_model_parallel_size == 2 - assert cfg.model.pipeline_model_parallel_size == 6 - assert cfg.model.expert_model_parallel_size == 32 + assert cfg.model.pipeline_model_parallel_size == 12 + assert cfg.model.expert_model_parallel_size == 16 + assert cfg.model.expert_tensor_parallel_size == 1 + assert cfg.model.context_parallel_size == 1 assert cfg.model.moe_token_dispatcher_type == "flex" assert cfg.model.moe_flex_dispatcher_backend == "hybridep" assert cfg.model.sequence_parallel is True @@ -229,7 +248,52 @@ def test_openmath_sft_uses_initial_parallelism_values() -> None: assert cfg.model.recompute_granularity == "selective" assert cfg.model.recompute_method is None assert cfg.model.recompute_num_layers is None - assert cfg.model.recompute_modules == ["moe", "layernorm", "core_attn", "moe_act"] + assert cfg.model.recompute_modules == ["moe", "core_attn", "moe_act"] + stages = cfg.model.hybrid_layer_pattern.split("|") + assert [len(stage) for stage in stages] == [9, 9, 9, 9, 9, 9, 9, 9, 10, 10, 11, 5] + assert "".join(stages) == NEMOTRON_3_ULTRA_HYBRID_LAYER_PATTERN + # A recomputed TE norm as a stage's first backward op fails on TE >= 2.18 (invalid device context). + assert "layernorm" not in cfg.model.recompute_modules + assert cfg.dist.distributed_timeout_minutes == 90 + assert cfg.env_vars["CUDA_DEVICE_MAX_CONNECTIONS"] == 1 + # cuDNN norms avoid TE's NVRTC norm path and its first-backward "invalid device context" failure. + assert cfg.env_vars["NVTE_NORM_FWD_USE_CUDNN"] == 1 + assert cfg.env_vars["NVTE_NORM_BWD_USE_CUDNN"] == 1 + + # 48 nodes x 8 GPUs: dense DP 16 keeps 8 microbatches per step, and expert DP 2 shards expert optimizer state. + world_size = 384 + model = cfg.model + data_parallel_size = world_size // ( + model.tensor_model_parallel_size * model.pipeline_model_parallel_size * model.context_parallel_size + ) + expert_data_parallel_size = world_size // ( + model.pipeline_model_parallel_size * model.expert_model_parallel_size * model.expert_tensor_parallel_size + ) + assert data_parallel_size == 16 + assert expert_data_parallel_size == 2 + assert cfg.train.global_batch_size % (cfg.train.micro_batch_size * data_parallel_size) == 0 + + +@pytest.mark.unit +def test_openmath_sft_layout_change_keeps_training_contract() -> None: + cfg = nemotron_3_ultra_sft_openmathinstruct2_packed_config() + + assert cfg.train.global_batch_size == 128 + assert cfg.train.micro_batch_size == 1 + assert cfg.model.seq_length == 4096 + assert cfg.model.pipeline_dtype == torch.bfloat16 + assert cfg.optimizer.lr == 5e-6 + assert cfg.optimizer.min_lr == 5e-7 + assert cfg.optimizer.adam_beta1 == 0.9 + assert cfg.optimizer.adam_beta2 == 0.98 + assert cfg.optimizer.adam_eps == 1e-8 + assert cfg.optimizer.weight_decay == 0.1 + assert cfg.scheduler.lr_decay_style == "cosine" + assert cfg.scheduler.lr_warmup_iters == 250 + assert cfg.scheduler.lr_decay_iters == 1000 + assert cfg.rng.seed == 5678 + assert cfg.ddp.grad_reduce_in_fp32 is True + assert cfg.ddp.use_distributed_optimizer is True assert cfg.train.train_iters == 1000 assert cfg.train.global_batch_size == 128 @@ -284,4 +348,16 @@ def test_openmath_sft_recompute_modules_are_not_shared() -> None: cfg.model.recompute_modules.append("sentinel") fresh_cfg = nemotron_3_ultra_sft_openmathinstruct2_packed_config() - assert fresh_cfg.model.recompute_modules == ["moe", "layernorm", "core_attn", "moe_act"] + assert fresh_cfg.model.recompute_modules == ["moe", "core_attn", "moe_act"] + + +@pytest.mark.unit +def test_split_hybrid_layer_pattern_keeps_mtp_suffix() -> None: + assert _split_hybrid_layer_pattern("MEM*E/*E", (2, 3)) == "ME|M*E/*E" + + +@pytest.mark.unit +@pytest.mark.parametrize(("pattern", "layers_per_stage"), [("MEM*E", (2, 2)), ("ME|M*E", (2, 3)), ("MEM*E", (5, 0))]) +def test_split_hybrid_layer_pattern_rejects_mismatched_stages(pattern: str, layers_per_stage: tuple[int, ...]) -> None: + with pytest.raises(ValueError, match="Cannot split"): + _split_hybrid_layer_pattern(pattern, layers_per_stage) diff --git a/tests/unit_tests/recipes/recipe_test_utils.py b/tests/unit_tests/recipes/recipe_test_utils.py index d9254c3d38..4bc3df3fce 100644 --- a/tests/unit_tests/recipes/recipe_test_utils.py +++ b/tests/unit_tests/recipes/recipe_test_utils.py @@ -117,6 +117,12 @@ def exported_recipe_factory_keys(module: ModuleType) -> set[tuple[str, str]]: } +# Nemotron 3 Ultra's 108-layer decoder pattern from the Hugging Face config. +NEMOTRON_3_ULTRA_HYBRID_LAYER_PATTERN = ( + "MEMEMEM*EMEMEM*EMEMEMEM*EMEMEMEM*EMEMEM*EMEMEMEM*EMEMEMEM*EMEMEM*EMEMEMEM*EMEMEMEM*EMEMEM*EMEMEMEM*EMEMEMEME" +) + + class _OfflineModelProvider: """Mutable provider with the fields HF-backed recipes read before writing. @@ -195,6 +201,9 @@ def to_megatron_provider(self, *args: object, **kwargs: object) -> _OfflineModel if model_id in {"zai-org/glm-5.2", "zai-org/glm-5.3"}: provider.dsa_indexer_topk_freq = 4 provider.dsa_indexer_skip_topk_offset = 3 + if model_id == "nvidia/nvidia-nemotron-3-ultra-550b-a55b-bf16": + provider.hybrid_layer_pattern = NEMOTRON_3_ULTRA_HYBRID_LAYER_PATTERN + provider.num_layers = len(provider.hybrid_layer_pattern) if "deepseek-v4-" in model_id: logical_layers = 61 if "deepseek-v4-pro" in model_id else 43 provider.hybrid_layer_pattern = "WEWE" + "".join( diff --git a/tests/unit_tests/recipes/test_nemotronh_recipes.py b/tests/unit_tests/recipes/test_nemotronh_recipes.py index cc8901dd4d..f1c62da2e3 100644 --- a/tests/unit_tests/recipes/test_nemotronh_recipes.py +++ b/tests/unit_tests/recipes/test_nemotronh_recipes.py @@ -27,7 +27,10 @@ import torch from megatron.bridge.training.utils.omegaconf_utils import OverridesError, process_config_with_overrides -from tests.unit_tests.recipes.recipe_test_utils import patch_recipe_module_global +from tests.unit_tests.recipes.recipe_test_utils import ( + NEMOTRON_3_ULTRA_HYBRID_LAYER_PATTERN, + patch_recipe_module_global, +) _nemotronh_module = importlib.import_module("megatron.bridge.recipes.nemotronh") @@ -72,6 +75,21 @@ def to_megatron_provider(self, *args, **kwargs): return _FakeModelProvider() +class _FakeUltraModelProvider(_FakeModelProvider): + """Fake provider with Nemotron 3 Ultra's decoder pattern, which the SFT recipe splits into stages.""" + + def __init__(self) -> None: + super().__init__() + self.hybrid_layer_pattern = NEMOTRON_3_ULTRA_HYBRID_LAYER_PATTERN + + +class _FakeUltraAutoBridge(_FakeAutoBridge): + """Return the Nemotron 3 Ultra fake provider.""" + + def to_megatron_provider(self, *args, **kwargs): + return _FakeUltraModelProvider() + + @pytest.fixture(autouse=True) def _patch_hf_backed_recipe_providers(monkeypatch: pytest.MonkeyPatch) -> None: """Keep AutoBridge-backed recipe construction deterministic and offline.""" @@ -81,10 +99,11 @@ def _patch_hf_backed_recipe_providers(monkeypatch: pytest.MonkeyPatch) -> None: "megatron.bridge.recipes.nemotronh.h100.nemotron_3_nano", "megatron.bridge.recipes.nemotronh.h100.nemotron_3_super", "megatron.bridge.recipes.nemotronh.nemotron_3_super", - "megatron.bridge.recipes.nemotronh.nemotron_3_ultra", ): module = importlib.import_module(module_name) patch_recipe_module_global(monkeypatch, module, "AutoBridge", _FakeAutoBridge) + ultra_module = importlib.import_module("megatron.bridge.recipes.nemotronh.nemotron_3_ultra") + patch_recipe_module_global(monkeypatch, ultra_module, "AutoBridge", _FakeUltraAutoBridge) def _assert_basic_config(cfg):