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
-
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
-
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):