Merge branch 'main' of https://github.com/NVIDIA-NeMo/Megatron-Bridge into rghadia/nt3-ultra-gtp-cg

Signed-off-by: rghadia <rghadia@nvidia.com>

Author: ghadiaravi13

.main.commit

@@ -1 +1 @@
-e0c5e18673bf4d4fd89cf87c4f1afbe521fe297f
+e30d863a6d82a3c45f2d71355dc73fbf5610967b

3rdparty/Megatron-LM

@@ -5393,21 +5393,21 @@
           "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": {
             "moe_dispatcher": "hybridep",
             "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": {

docs/_static/model-verification/catalog-v1.json

@@ -115,6 +115,12 @@ family's Megatron Core `ModelBuilder`, such as `GPTModelBuilder` or
 while other model families migrate. Calling the legacy provider API for a
 builder-backed family emits a deprecation warning.
 
+The builder path is opt-in per family: a bridge opts in by setting
+`MODEL_CONFIG_CLASS`, which defaults to `None`. For a family that has not
+migrated, `get_model_config()` raises `ModelConfigNotSupportedError` and
+`to_megatron_provider()` remains the supported path, without a deprecation
+warning.
+
 ```python
 from megatron.bridge import AutoBridge
 

docs/bridge-guide.md

@@ -408,12 +408,12 @@ Choose a workflow, precision, and exact recorded combination. The command and ex
             <span>Command</span>
             <button type="button" class="verification-copy-command">Copy</button>
           </div>
-          <pre><code class="language-bash">./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 &#x27;dataset.hf_dataset.split=&quot;train_1M[:10%]&quot;&#x27; &#x27;dataset.hf_dataset.load_kwargs=&#123;revision:&quot;469216e3f46f4dacf476b382e192485ea51a143e&quot;&#125;&#x27; 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</code></pre>
+          <pre><code class="language-bash">./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 &#x27;dataset.hf_dataset.split=&quot;train_1M[:10%]&quot;&#x27; &#x27;dataset.hf_dataset.load_kwargs=&#123;revision:&quot;469216e3f46f4dacf476b382e192485ea51a143e&quot;&#125;&#x27; 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</code></pre>
         </div>
       </section>
       <section class="verification-expected-result">
         <h5>Expected result</h5>
-        <p>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.</p>
+        <p>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.</p>
       </section>
     </article>
     <article id="nemotron-3-ultra-550b-a55b-sft-gb200" class="verification-model-detail" data-entry-detail="nemotron-3-ultra-550b-a55b-sft-gb200" tabindex="-1">

docs/fern/versions/nightly/pages/models/nemotron/nemotron3-ultra.mdx

@@ -414,12 +414,12 @@ Choose a workflow, precision, and exact recorded combination. The command and ex
             <span>Command</span>
             <button type="button" class="verification-copy-command">Copy</button>
           </div>
-          <pre><code class="language-bash">./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 &#x27;dataset.hf_dataset.split=&quot;train_1M[:10%]&quot;&#x27; &#x27;dataset.hf_dataset.load_kwargs={revision:&quot;469216e3f46f4dacf476b382e192485ea51a143e&quot;}&#x27; 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</code></pre>
+          <pre><code class="language-bash">./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 &#x27;dataset.hf_dataset.split=&quot;train_1M[:10%]&quot;&#x27; &#x27;dataset.hf_dataset.load_kwargs={revision:&quot;469216e3f46f4dacf476b382e192485ea51a143e&quot;}&#x27; 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</code></pre>
         </div>
       </section>
       <section class="verification-expected-result">
         <h5>Expected result</h5>
-        <p>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.
+        <p>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.
 </p>
       </section>
     </article>

docs/models/nemotron/nemotron3-ultra.md

@@ -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

examples/model_verification_cards/nemotron-3-ultra-550b-a55b/card.yaml

@@ -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

examples/models/nemotron/nemotron_3_ultra/README.md

@@ -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 "======================================"

examples/models/nemotron/nemotron_3_ultra/slurm_sft.sh

@@ -171,7 +171,7 @@ requires-dist = []
 # MCore pin's metadata so a lock-only operation never invokes the build backend.
 [[tool.uv.dependency-metadata]]
 name = "transformer-engine"
-version = "2.20.2+6ea2a74a"
+version = "2.20.2+44f31bbe"
 requires-dist = [
     "pydantic<2.14",
     "importlib-metadata>=1.0",