feat(nemotronh): add Nemotron 3 Ultra recipes and examples

README.md

@@ -11,7 +11,7 @@
 </div>
 
 ## 📣 News
-- [06/04/2026] [**NVIDIA Nemotron 3 Ultra**](https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16) is now public! Day-0 support for the 550B-A55B hybrid Mamba-Transformer MoE model is available on the [`nemotron_3_ultra`](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/nemotron_3_ultra) branch, including checkpoint conversion, inference, SFT, PEFT (LoRA), and pretraining examples. Read the [NVIDIA Technical Blog](https://developer.nvidia.com/blog/nvidia-nemotron-3-ultra-powers-faster-more-efficient-reasoning-for-long-running-agents/) and see the [examples README](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/nemotron_3_ultra/examples/models/nemotron/nemotron_3/ultra/README.md) for the full walkthrough.
+- [06/04/2026] [**NVIDIA Nemotron 3 Ultra**](https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16) is now public! Megatron Bridge supports the 550B-A55B hybrid Mamba-Transformer MoE model on **main**, including checkpoint conversion, inference, SFT, PEFT (LoRA), and pretraining examples. Read the [NVIDIA Technical Blog](https://developer.nvidia.com/blog/nvidia-nemotron-3-ultra-powers-faster-more-efficient-reasoning-for-long-running-agents/) and see the [examples README](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/models/nemotron/nemotron_3/ultra/README.md) for the full walkthrough.
 
 - [05/28/2026] [**Step-3.7-Flash**](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/examples/models/stepfun/step37) is now merged on **main**! See the [examples README](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/examples/models/stepfun/step37/README.md) for sft training details.
 

docs/models/nemotron/index.md

@@ -9,6 +9,7 @@ llama-nemotron.md
 nemotronh.md
 nemotron3-nano.md
 nemotron3-super.md
+nemotron3-ultra.md
 nemotron-nano-v2-vl.md
 nemotron-3-omni.md
 ```
@@ -19,6 +20,7 @@ nemotron-3-omni.md
 | Nemotron H and Nemotron Nano v2 | [nemotronh.md](nemotronh.md) |
 | Nemotron-3 Nano | [nemotron3-nano.md](nemotron3-nano.md) |
 | Nemotron-3 Super | [nemotron3-super.md](nemotron3-super.md) |
+| Nemotron-3 Ultra | [nemotron3-ultra.md](nemotron3-ultra.md) |
 | Nemotron Nano V2 VL | [nemotron-nano-v2-vl.md](nemotron-nano-v2-vl.md) |
 | Nemotron-3 Nano Omni | [nemotron-3-omni.md](nemotron-3-omni.md) |
 

docs/models/nemotron/nemotron3-ultra.md

@@ -0,0 +1,11 @@
+# Nemotron 3 Ultra
+
+[Nemotron 3 Ultra](https://docs.nvidia.com/nemotron/nightly/usage-cookbook/Nemotron-3-Ultra-Base/README.html)
+is a 550B total / A55B active hybrid Mamba-Transformer MoE model.
+
+Megatron Bridge provides Nemotron 3 Ultra recipes and examples for
+Hugging Face to Megatron conversion, inference, packed OpenMathInstruct-2
+full SFT, and packed OpenMathInstruct-2 LoRA PEFT.
+
+Use the main example README for setup and scripts:
+[`examples/models/nemotron/nemotron_3/ultra/README.md`](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/models/nemotron/nemotron_3/ultra/README.md).

examples/models/nemotron/nemotron_3/README.md

@@ -6,6 +6,7 @@ This directory contains example scripts for Nemotron 3 language models:
 |-------|-----------|-------------------|--------------|
 | Nemotron 3 Nano | 30B | A3B | [nano/](nano/) |
 | Nemotron 3 Super | 120B | A12B | [super/](super/) |
+| Nemotron 3 Ultra | 550B | A55B | [ultra/](ultra/) |
 
 ## Workspace Configuration
 
@@ -21,7 +22,7 @@ Directory structure:
 
 ## Checkpoint Conversion
 
-Each model has its own conversion script: [nano/conversion.sh](nano/conversion.sh), [super/conversion.sh](super/conversion.sh).
+Nano and Super have conversion scripts: [nano/conversion.sh](nano/conversion.sh), [super/conversion.sh](super/conversion.sh). Ultra has Slurm examples for multi-node conversion, inference, and OpenMath training; see [ultra/](ultra/) and [Ultra documentation](../../../../docs/models/nemotron/nemotron3-ultra.md).
 
 ## Training Recipes
 
@@ -37,6 +38,11 @@ Available recipes:
 - `nemotron_3_super_sft_config`: Supervised fine-tuning
 - `nemotron_3_super_peft_config`: PEFT with LoRA support
 
+**Ultra** ([source](../../../../src/megatron/bridge/recipes/nemotronh/nemotron_3_ultra.py)):
+- `nemotron_3_ultra_pretrain_config`: Pretraining
+- `nemotron_3_ultra_sft_openmathinstruct2_packed_config`: Packed OpenMathInstruct-2 SFT
+- `nemotron_3_ultra_peft_openmathinstruct2_packed_config`: Packed OpenMathInstruct-2 PEFT
+
 Before training, ensure the following are configured:
 1. **Container Image**: Set `CONTAINER_IMAGE` in the SLURM scripts to your container path
 2. **Container Mounts**: (optional) Set `CONTAINER_MOUNTS` for data and workspace directories
@@ -55,6 +61,13 @@ See the SLURM scripts in [nano/](nano/): [slurm_pretrain.sh](nano/slurm_pretrain
 
 See the SLURM scripts in [super/](super/): [slurm_pretrain.sh](super/slurm_pretrain.sh), [slurm_sft.sh](super/slurm_sft.sh), [slurm_peft.sh](super/slurm_peft.sh).
 
+### Ultra
+
+See [ultra/slurm_inference.sh](ultra/slurm_inference.sh) for the 4-node inference pattern.
+For OpenMath training, use [ultra/slurm_sft.sh](ultra/slurm_sft.sh) and
+[ultra/slurm_peft.sh](ultra/slurm_peft.sh), which default to the current
+OpenMath tuning starting points.
+
 ## Evaluation
 
 Coming soon.

examples/models/nemotron/nemotron_3/ultra/README.md

@@ -0,0 +1,194 @@
+# Nemotron 3 Ultra Examples
+
+This directory contains examples for Nemotron 3 Ultra conversion, inference,
+DCLM pretraining, packed OpenMathInstruct-2 full SFT, and packed
+OpenMathInstruct-2 LoRA PEFT.
+
+Nemotron 3 Ultra is a 550B total / A55B active hybrid Mamba-Transformer MoE
+model. See the
+[Nemotron 3 Ultra Base model guide](https://docs.nvidia.com/nemotron/nightly/usage-cookbook/Nemotron-3-Ultra-Base/README.html)
+for model details.
+
+## Workspace Configuration
+
+The scripts use `WORKSPACE` as the base directory for checkpoints, packed data,
+and results. Defaults:
+
+```bash
+export WORKSPACE=/workspace
+export MODEL_HOME=${WORKSPACE}/models/nvidia
+export HF_MODEL_PATH=nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16
+export MEGATRON_MODEL_PATH=${MODEL_HOME}/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16-megatron
+export PRETRAINED_CHECKPOINT=${MEGATRON_MODEL_PATH}
+```
+
+Use shared filesystems for multi-node jobs:
+
+```bash
+export HF_HOME=${WORKSPACE}/cache/hf
+export NEMO_HOME=${WORKSPACE}/cache/nemo
+export UV_CACHE_DIR=${WORKSPACE}/cache/uv
+```
+
+The BF16 Hugging Face cache and imported Megatron checkpoint are each about
+1.1 TB. Reserve at least 2.5 TB for model storage before starting checkpoint
+conversion, plus additional space for training outputs and logs. Full-model
+training checkpoints can each require several TB, so set `WORKSPACE` to a
+filesystem with enough quota before running SFT or pretraining.
+
+## Hardware Starting Points
+
+The checked-in Slurm scripts default to 8xH100 nodes unless noted below.
+When running on 4xGB200 nodes, update the `#SBATCH --nodes`,
+`#SBATCH --ntasks-per-node`, `#SBATCH --gpus-per-node`, and parallelism
+environment variables to the GB200 values in this table.
+
+| Workflow | 8xH100 nodes | 4xGB200 nodes |
+| --- | --- | --- |
+| Checkpoint import | 1 node with [conversion.sh](conversion.sh), CPU import path | 6 nodes with [slurm_conversion.sh](slurm_conversion.sh), `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 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.
+Keep `TP` within a node-local NVLink domain and scale with `PP`, `EP`, and
+data parallelism when moving between hardware. For MoE sizing, the minimum GPU
+count is `PP * max(TP * CP, EP * ETP)`, then additional GPUs increase dense
+DP and expert DP.
+
+## Checkpoint Conversion
+
+Use [conversion.sh](conversion.sh) for CPU checkpoint import when the node has
+enough host RAM to materialize Nemotron 3 Ultra, for example an 8xH100 node.
+This is the preferred path when available because it avoids distributed GPU
+memory pressure during import.
+
+```bash
+bash conversion.sh
+```
+
+Set these variables for your environment:
+
+- `WORKSPACE`
+- `HF_HOME`
+- `UV_CACHE_DIR`
+- `HF_MODEL_PATH`
+- `MEGATRON_MODEL_PATH`
+
+Use [slurm_conversion.sh](slurm_conversion.sh) for distributed GPU checkpoint
+import when host RAM is not large enough, for example a 4xGB200 setup with less
+than 1 TB of host RAM. The checked-in distributed example defaults to 6
+8-GPU nodes with `TP=1 PP=6 EP=8`; for 4xGB200, use 6 nodes and set
+`#SBATCH --ntasks-per-node=4`, `#SBATCH --gpus-per-node=4`, and
+`TP=1 PP=6 EP=4`.
+
+```bash
+sbatch slurm_conversion.sh
+```
+
+Set these variables for your environment:
+
+- `CONTAINER_IMAGE`
+- `CONTAINER_MOUNTS`
+- `WORKDIR`
+- `WORKSPACE`
+- `HF_MODEL_PATH`
+- `MEGATRON_MODEL_PATH`
+
+## Inference
+
+Use [slurm_inference.sh](slurm_inference.sh) for 4-node text generation with
+`TP=1 PP=4 EP=8`. The script defaults `KV_CACHE_BUFFER_SIZE_GB=4` to keep
+the inference KV/context buffer within H100 memory for the default prompt
+lengths. On 4xGB200 nodes, use 3 nodes with `TP=1 PP=3 EP=4`.
+
+```bash
+sbatch slurm_inference.sh
+```
+
+Set `MEGATRON_MODEL_PATH` to generate from an imported Megatron checkpoint.
+Leave it unset to load from the Hugging Face checkpoint path.
+
+## DCLM Pretraining
+
+Use [slurm_pretrain.sh](slurm_pretrain.sh) for DCLM pretraining with
+`TP=4 PP=12 EP=16` and full uniform recompute on 8xH100 nodes. On
+4xGB200 nodes, use 24 nodes with `TP=2 PP=3 EP=32`.
+
+```bash
+sbatch slurm_pretrain.sh
+```
+
+Set `DCLM_DATA_DIR` to a preprocessed DCLM directory containing
+`*_text_document.bin` / `*_text_document.idx` files. The script defaults to
+matching `dclm_01_*_text_document.bin`. Async checkpoint saving is enabled by
+the recipe; the script defaults `SAVE_INTERVAL=1000` to save one checkpoint for
+the default 1000-iteration starter run.
+
+## OpenMath Packed Data
+
+Pre-pack OpenMath data before training:
+
+```bash
+sbatch pack_data_job.sh
+```
+
+Use the same `SEQ_LENGTH`, `HF_MODEL_PATH`, and `NEMO_HOME` for packing and
+training. `NEMO_HOME` must point at a shared filesystem visible on all nodes.
+
+## Training
+
+PEFT:
+
+```bash
+sbatch slurm_peft.sh
+```
+
+Full SFT:
+
+```bash
+sbatch slurm_sft.sh
+```
+
+The scripts default to OpenMath convergence settings: `TRAIN_ITERS=1000`,
+`GLOBAL_BATCH_SIZE=128`, `SEQ_LENGTH=4096`, and `LR_WARMUP_ITERS=250`. W&B
+logging is disabled by default. SFT and PEFT save at the final training
+iteration by default; the SFT script removes older intermediate `iter_*`
+checkpoints after a successful run to avoid retaining multiple full-model
+checkpoints.
+
+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.
+
+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`.
+
+Advanced VPP, pipeline-layout, and recompute sweeps are intentionally left out
+of these starter scripts; add those overrides only for targeted performance
+experiments.
+
+## W&B
+
+W&B logging is disabled by default:
+
+```bash
+WANDB_ENTITY=nvidia-nemo-fw-public
+WANDB_PROJECT=megatron-bridge-nemotron-ultra
+WANDB_MODE=disabled
+```
+
+To enable online W&B logging, set `WANDB_MODE=online` and make `WANDB_API_KEY`
+visible in the submit environment.
+
+Run names include model, OpenMath, mode, TP/PP/EP, recompute, and Slurm job ID.

examples/models/nemotron/nemotron_3/ultra/conversion.sh

@@ -0,0 +1,41 @@
+#!/usr/bin/env bash
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+set -euo pipefail
+
+WORKSPACE=${WORKSPACE:-/workspace}
+MODEL_HOME=${MODEL_HOME:-${WORKSPACE}/models/nvidia}
+HF_MODEL_PATH=${HF_MODEL_PATH:-nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16}
+MEGATRON_MODEL_PATH=${MEGATRON_MODEL_PATH:-${MODEL_HOME}/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16-megatron}
+
+[ -n "${HF_HOME:-}" ] && export HF_HOME
+[ -n "${UV_CACHE_DIR:-}" ] && export UV_CACHE_DIR
+
+mkdir -p "$(dirname "$MEGATRON_MODEL_PATH")"
+
+if [ -e "${MEGATRON_MODEL_PATH}/latest_checkpointed_iteration.txt" ] || [ -e "${MEGATRON_MODEL_PATH}/latest_train_state.pt" ]; then
+    echo "ERROR: target already contains a Megatron checkpoint: ${MEGATRON_MODEL_PATH}"
+    exit 1
+fi
+
+echo "Nemotron 3 Ultra CPU import"
+echo "HF_MODEL_PATH=${HF_MODEL_PATH}"
+echo "MEGATRON_MODEL_PATH=${MEGATRON_MODEL_PATH}"
+
+uv run --no-sync python examples/conversion/convert_checkpoints.py import \
+    --hf-model "$HF_MODEL_PATH" \
+    --megatron-path "$MEGATRON_MODEL_PATH" \
+    --torch-dtype bfloat16 \
+    --device-map cpu

examples/models/nemotron/nemotron_3/ultra/pack_data_job.sh

@@ -0,0 +1,62 @@
+#!/bin/bash
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+#SBATCH --job-name=nemotron-ultra-pack
+#SBATCH --nodes=1
+#SBATCH --ntasks-per-node=1
+#SBATCH --cpus-per-task=16
+#SBATCH --time=04:00:00
+#SBATCH --account=<your-account>
+#SBATCH --partition=cpu
+#SBATCH --output=logs/nemotron_ultra_pack_%j.log
+
+set -euo pipefail
+
+CONTAINER_IMAGE=${CONTAINER_IMAGE:-}
+CONTAINER_MOUNTS=${CONTAINER_MOUNTS:-}
+WORKDIR=${WORKDIR:-/opt/Megatron-Bridge}
+HF_MODEL_PATH=${HF_MODEL_PATH:-nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16}
+RECIPE_NAME=${RECIPE_NAME:-nemotron_3_ultra_sft_openmathinstruct2_packed_config}
+SEQ_LENGTH=${SEQ_LENGTH:-4096}
+
+[ -n "${HF_HOME:-}" ] && export HF_HOME
+[ -n "${NEMO_HOME:-}" ] && export NEMO_HOME
+[ -n "${UV_CACHE_DIR:-}" ] && export UV_CACHE_DIR
+export WORKDIR HF_MODEL_PATH RECIPE_NAME SEQ_LENGTH
+
+if [ -z "$CONTAINER_IMAGE" ]; then
+    echo "ERROR: CONTAINER_IMAGE must be set."
+    exit 1
+fi
+
+mkdir -p logs
+
+SRUN_CMD=(srun --mpi=pmix --container-image="${CONTAINER_IMAGE}" --no-container-mount-home)
+if [ -n "$CONTAINER_MOUNTS" ]; then
+    SRUN_CMD+=(--container-mounts="${CONTAINER_MOUNTS}")
+fi
+
+"${SRUN_CMD[@]}" bash -c '
+set -euo pipefail
+cd "$WORKDIR"
+export PYTHONPATH="$WORKDIR/src:$WORKDIR/3rdparty/Megatron-LM:${PYTHONPATH:-}"
+
+uv run --no-sync python scripts/training/pack_sft_data.py \
+    --recipe "$RECIPE_NAME" \
+    --seq-length "$SEQ_LENGTH" \
+    --hf-path "$HF_MODEL_PATH"
+'
+
+echo PACK_DATA_DONE