From b415a91bc8464abc0cd141be62c56b085382d9f5 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Mon, 8 Jun 2026 19:07:21 -0700 Subject: [PATCH 01/10] feat(nemotronh): add Nemotron 3 Ultra examples Signed-off-by: Chen Cui --- docs/models/nemotron/index.md | 2 + docs/models/nemotron/nemotron3-ultra.md | 11 + examples/models/nemotron/nemotron_3/README.md | 15 +- .../nemotron/nemotron_3/ultra/README.md | 243 ++++++++++++++ .../nemotron/nemotron_3/ultra/conversion.sh | 41 +++ .../nemotron_3/ultra/pack_data_job.sh | 62 ++++ .../nemotron_3/ultra/slurm_conversion.sh | 120 +++++++ .../nemotron_3/ultra/slurm_inference.sh | 138 ++++++++ .../nemotron/nemotron_3/ultra/slurm_peft.sh | 179 ++++++++++ .../nemotron_3/ultra/slurm_pretrain.sh | 180 ++++++++++ .../nemotron/nemotron_3/ultra/slurm_sft.sh | 186 +++++++++++ .../bridge/recipes/nemotronh/__init__.py | 9 + .../recipes/nemotronh/nemotron_3_ultra.py | 314 ++++++++++++++++++ .../nemotronh/test_nemotron_3_ultra.py | 116 +++++++ 14 files changed, 1615 insertions(+), 1 deletion(-) create mode 100644 docs/models/nemotron/nemotron3-ultra.md create mode 100644 examples/models/nemotron/nemotron_3/ultra/README.md create mode 100755 examples/models/nemotron/nemotron_3/ultra/conversion.sh create mode 100755 examples/models/nemotron/nemotron_3/ultra/pack_data_job.sh create mode 100755 examples/models/nemotron/nemotron_3/ultra/slurm_conversion.sh create mode 100755 examples/models/nemotron/nemotron_3/ultra/slurm_inference.sh create mode 100755 examples/models/nemotron/nemotron_3/ultra/slurm_peft.sh create mode 100755 examples/models/nemotron/nemotron_3/ultra/slurm_pretrain.sh create mode 100755 examples/models/nemotron/nemotron_3/ultra/slurm_sft.sh create mode 100644 src/megatron/bridge/recipes/nemotronh/nemotron_3_ultra.py create mode 100644 tests/unit_tests/recipes/nemotronh/test_nemotron_3_ultra.py diff --git a/docs/models/nemotron/index.md b/docs/models/nemotron/index.md index 66089cf1a2..28d978789e 100644 --- a/docs/models/nemotron/index.md +++ b/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) | diff --git a/docs/models/nemotron/nemotron3-ultra.md b/docs/models/nemotron/nemotron3-ultra.md new file mode 100644 index 0000000000..44df39ec4f --- /dev/null +++ b/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). diff --git a/examples/models/nemotron/nemotron_3/README.md b/examples/models/nemotron/nemotron_3/README.md index c2446a0369..cfefc5aba9 100644 --- a/examples/models/nemotron/nemotron_3/README.md +++ b/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. diff --git a/examples/models/nemotron/nemotron_3/ultra/README.md b/examples/models/nemotron/nemotron_3/ultra/README.md new file mode 100644 index 0000000000..54c5ec9788 --- /dev/null +++ b/examples/models/nemotron/nemotron_3/ultra/README.md @@ -0,0 +1,243 @@ +# 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. + +## 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=2 PP=12 EP=32` script default; validate on your DCLM mount before long runs | 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=32`, selective recompute on `moe+layernorm+core_attn+moe_act` | 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. + +## Day-0 Code + +Start from the NeMo 26.04.01 container: + +```bash +docker run --rm -it --gpus all --shm-size=24g \ + -w /workspace \ + nvcr.io/nvidia/nemo:26.04.01 \ + bash +``` + +Inside the container, clone the Day-0 Megatron Bridge branch and initialize the +bundled Megatron-Core submodule: + +```bash +export WORKSPACE=/workspace +export MEGATRON_BRIDGE_REPO=https://github.com/NVIDIA-NeMo/Megatron-Bridge.git +export MEGATRON_BRIDGE_BRANCH=nemotron_3_ultra + +cd "$WORKSPACE" +git clone --branch "$MEGATRON_BRIDGE_BRANCH" "$MEGATRON_BRIDGE_REPO" +cd Megatron-Bridge + +git submodule update --init --recursive --depth 1 3rdparty/Megatron-LM + +uv lock +uv sync --extra ssm +``` + +Run dependency sync during setup. The Slurm job scripts use `uv run --no-sync` +so they do not download dependencies or mutate the Python environment during +allocated jobs. + +Verify that `megatron.core` and `megatron.bridge` resolve to this checkout: + +```bash +uv run python - <<'PY' +import megatron.bridge +import megatron.core +from megatron.bridge.recipes.nemotronh import nemotron_3_ultra_pretrain_config + +print("core:", list(megatron.core.__path__)) +print("bridge:", list(megatron.bridge.__path__)) +print("recipe:", nemotron_3_ultra_pretrain_config.__name__) +PY +``` + +Expected paths should include: + +```text +/workspace/Megatron-Bridge/3rdparty/Megatron-LM/megatron/core +/workspace/Megatron-Bridge/src/megatron/bridge +``` + +## 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=2 PP=12 EP=32` 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`. + +## 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=32`, selective recompute on + `moe+layernorm+core_attn+moe_act`. 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. diff --git a/examples/models/nemotron/nemotron_3/ultra/conversion.sh b/examples/models/nemotron/nemotron_3/ultra/conversion.sh new file mode 100755 index 0000000000..be701694cb --- /dev/null +++ b/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 diff --git a/examples/models/nemotron/nemotron_3/ultra/pack_data_job.sh b/examples/models/nemotron/nemotron_3/ultra/pack_data_job.sh new file mode 100755 index 0000000000..2243ca7ac7 --- /dev/null +++ b/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= +#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 diff --git a/examples/models/nemotron/nemotron_3/ultra/slurm_conversion.sh b/examples/models/nemotron/nemotron_3/ultra/slurm_conversion.sh new file mode 100755 index 0000000000..83f7d4f2eb --- /dev/null +++ b/examples/models/nemotron/nemotron_3/ultra/slurm_conversion.sh @@ -0,0 +1,120 @@ +#!/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. + +# ============================================================================== +# Nemotron 3 Ultra Distributed GPU Checkpoint Import +# +# Use this path when host RAM is not large enough for CPU checkpoint import +# and the checkpoint must be materialized across GPUs. +# +# Usage: +# 1. Modify the #SBATCH directives for your cluster. +# 2. Set CONTAINER_IMAGE and optional CONTAINER_MOUNTS. +# 3. Submit: sbatch slurm_conversion.sh +# ============================================================================== + +#SBATCH --job-name=nemotron-ultra-gpu-import +#SBATCH --nodes=6 +#SBATCH --ntasks-per-node=8 +#SBATCH --gpus-per-node=8 +#SBATCH --time=04:00:00 +#SBATCH --account= +#SBATCH --partition=batch +#SBATCH --output=logs/nemotron_ultra_import_%j.log +#SBATCH --exclusive + +set -euo pipefail + +CONTAINER_IMAGE=${CONTAINER_IMAGE:-} +CONTAINER_MOUNTS=${CONTAINER_MOUNTS:-} +WORKDIR=${WORKDIR:-/opt/Megatron-Bridge} +MODEL_HOME=${MODEL_HOME:-${WORKSPACE:-/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} + +TP=${TP:-1} +PP=${PP:-6} +EP=${EP:-8} +ETP=${ETP:-1} +GPUS_PER_NODE=${GPUS_PER_NODE:-8} + +[ -n "${HF_HOME:-}" ] && export HF_HOME +[ -n "${UV_CACHE_DIR:-}" ] && export UV_CACHE_DIR +export TORCH_NCCL_AVOID_RECORD_STREAMS=1 +export NCCL_NVLS_ENABLE=0 +export NCCL_DEBUG=${NCCL_DEBUG:-WARN} +export NCCL_TIMEOUT=${NCCL_TIMEOUT:-1800000} +export PYTORCH_CUDA_ALLOC_CONF=${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True} +export CUDA_DEVICE_MAX_CONNECTIONS=${CUDA_DEVICE_MAX_CONNECTIONS:-1} + +if [ -z "$CONTAINER_IMAGE" ]; then + echo "ERROR: CONTAINER_IMAGE must be set." + exit 1 +fi + +if [ "$((TP * PP * EP))" -ne "$((SLURM_JOB_NUM_NODES * GPUS_PER_NODE))" ]; then + echo "ERROR: TP*PP*EP must equal nodes*GPUS_PER_NODE for this script." + echo "TP=$TP PP=$PP EP=$EP nodes=$SLURM_JOB_NUM_NODES GPUS_PER_NODE=$GPUS_PER_NODE" + exit 2 +fi + +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 3 +fi + +mkdir -p logs "$(dirname "$MEGATRON_MODEL_PATH")" + +MASTER_ADDR=$(python3 - <<'PY' +import os +import re + +nodelist = os.environ.get("SLURM_NODELIST", "") +match = re.match(r"([A-Za-z0-9_-]+)\[(\d+)", nodelist) +print(match.group(1) + match.group(2) if match else nodelist.split(",")[0]) +PY +) +MASTER_PORT=$((18000 + SLURM_JOB_ID % 40000)) +export MASTER_ADDR MASTER_PORT HF_MODEL_PATH MEGATRON_MODEL_PATH TP PP EP ETP GPUS_PER_NODE WORKDIR + +echo "Nemotron 3 Ultra distributed GPU import" +echo "Job ${SLURM_JOB_ID} nodes=${SLURM_JOB_NUM_NODES} GPUs/node=${GPUS_PER_NODE} TP=${TP} PP=${PP} EP=${EP} ETP=${ETP}" +echo "HF_MODEL_PATH=${HF_MODEL_PATH}" +echo "MEGATRON_MODEL_PATH=${MEGATRON_MODEL_PATH}" + +SRUN_CMD=(srun --mpi=pmix --no-kill --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:-}" +export RANK="${SLURM_PROCID}" +export WORLD_SIZE="${SLURM_NTASKS}" +export LOCAL_RANK="${SLURM_LOCALID}" +export LOCAL_WORLD_SIZE="$GPUS_PER_NODE" + +uv run --no-sync python examples/conversion/convert_checkpoints_multi_gpu.py \ + import \ + --hf-model "$HF_MODEL_PATH" \ + --megatron-path "$MEGATRON_MODEL_PATH" \ + --tp "$TP" --pp "$PP" --ep "$EP" --etp "$ETP" \ + --torch-dtype bfloat16 +' + +echo IMPORT_DONE diff --git a/examples/models/nemotron/nemotron_3/ultra/slurm_inference.sh b/examples/models/nemotron/nemotron_3/ultra/slurm_inference.sh new file mode 100755 index 0000000000..c9e820757f --- /dev/null +++ b/examples/models/nemotron/nemotron_3/ultra/slurm_inference.sh @@ -0,0 +1,138 @@ +#!/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. + +# ============================================================================== +# Nemotron 3 Ultra Inference (4 nodes / 32 GPUs via Slurm) +# +# Usage: +# 1. Set CONTAINER_IMAGE, CONTAINER_MOUNTS, and cache/token environment variables. +# 2. Optionally set HF_MODEL_PATH to a local Hugging Face snapshot. +# 3. Submit with: sbatch examples/models/nemotron/nemotron_3/ultra/slurm_inference.sh +# ============================================================================== + +#SBATCH --job-name=nemotron-ultra-inference +#SBATCH --nodes=4 +#SBATCH --ntasks-per-node=8 +#SBATCH --gpus-per-node=8 +#SBATCH --time=04:00:00 +#SBATCH --account= +#SBATCH --partition=batch +#SBATCH --output=logs/nemotron_ultra_inference_%j.log +#SBATCH --exclusive + +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} +MEGATRON_MODEL_PATH=${MEGATRON_MODEL_PATH:-} +PROMPT=${PROMPT:-"Solve 2x + 3 = 11. Show the reasoning briefly."} +MAX_NEW_TOKENS=${MAX_NEW_TOKENS:-64} +KV_CACHE_BUFFER_SIZE_GB=${KV_CACHE_BUFFER_SIZE_GB:-4} +INFERENCE_MOE_TOKEN_DISPATCHER_TYPE=${INFERENCE_MOE_TOKEN_DISPATCHER_TYPE:-nccl} + +TP=${TP:-1} +PP=${PP:-4} +EP=${EP:-8} +ETP=${ETP:-1} +GPUS_PER_NODE=${GPUS_PER_NODE:-8} + +[ -n "${HF_HOME:-}" ] && export HF_HOME +[ -n "${UV_CACHE_DIR:-}" ] && export UV_CACHE_DIR +[ -n "${NEMO_HOME:-}" ] && export NEMO_HOME +export TORCH_NCCL_AVOID_RECORD_STREAMS=1 +export NCCL_NVLS_ENABLE=0 +export NCCL_DEBUG=${NCCL_DEBUG:-WARN} +export NCCL_TIMEOUT=${NCCL_TIMEOUT:-1800000} +export PYTORCH_CUDA_ALLOC_CONF=${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True} +export CUDA_DEVICE_MAX_CONNECTIONS=${CUDA_DEVICE_MAX_CONNECTIONS:-1} + +if [ -z "$CONTAINER_IMAGE" ]; then + echo "ERROR: CONTAINER_IMAGE must be set." + exit 1 +fi + +TOTAL_GPUS=$((SLURM_JOB_NUM_NODES * GPUS_PER_NODE)) +EXPERT_MESH=$((EP * ETP)) +if [ "$EXPERT_MESH" -gt "$TP" ]; then + STAGE_GPUS=$EXPERT_MESH +else + STAGE_GPUS=$TP +fi +MIN_GPUS=$((PP * STAGE_GPUS)) +if [ "$((TOTAL_GPUS % MIN_GPUS))" -ne 0 ]; then + echo "ERROR: nodes*GPUS_PER_NODE must be a multiple of PP*max(TP,EP*ETP) for MoE inference." + echo "TP=$TP PP=$PP EP=$EP ETP=$ETP nodes=$SLURM_JOB_NUM_NODES GPUS_PER_NODE=$GPUS_PER_NODE min_gpus=$MIN_GPUS total_gpus=$TOTAL_GPUS" + exit 2 +fi + +mkdir -p logs + +if [ -z "${COORDINATOR_HOST:-}" ]; then + if command -v scontrol >/dev/null 2>&1 && [ -n "${SLURM_NODELIST:-}" ]; then + COORDINATOR_HOST=$(scontrol show hostnames "$SLURM_NODELIST" | head -n 1) + else + COORDINATOR_HOST=$(hostname) + fi +fi +if [ -z "$COORDINATOR_HOST" ]; then + echo "ERROR: unable to resolve COORDINATOR_HOST from SLURM_NODELIST=${SLURM_NODELIST:-}." + exit 3 +fi + +export HF_MODEL_PATH MEGATRON_MODEL_PATH PROMPT MAX_NEW_TOKENS KV_CACHE_BUFFER_SIZE_GB INFERENCE_MOE_TOKEN_DISPATCHER_TYPE +export TP PP EP ETP GPUS_PER_NODE WORKDIR +export COORDINATOR_HOST + +echo "Nemotron 3 Ultra inference" +echo "Job ${SLURM_JOB_ID} nodes=${SLURM_JOB_NUM_NODES} GPUs/node=${GPUS_PER_NODE} TP=${TP} PP=${PP} EP=${EP} ETP=${ETP}" +echo "HF_MODEL_PATH=${HF_MODEL_PATH}" +echo "MEGATRON_MODEL_PATH=${MEGATRON_MODEL_PATH:-}" +echo "KV_CACHE_BUFFER_SIZE_GB=${KV_CACHE_BUFFER_SIZE_GB}" +echo "COORDINATOR_HOST=${COORDINATOR_HOST}" + +SRUN_CMD=(srun --mpi=pmix --no-kill --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" + +rm -f /opt/venv/lib/python3.12/site-packages/__editable__*megatron*.pth \ + /opt/venv/lib/python3.12/site-packages/__editable__*megatron*.py 2>/dev/null || true + +export PYTHONPATH="$WORKDIR/src:$WORKDIR/3rdparty/Megatron-LM:${PYTHONPATH:-}" + +MEGATRON_MODEL_ARGS=() +if [ -n "${MEGATRON_MODEL_PATH:-}" ]; then + MEGATRON_MODEL_ARGS=(--megatron_model_path "$MEGATRON_MODEL_PATH") +fi + +uv run --no-sync python scripts/inference/text_generation.py \ + --hf_model_path "$HF_MODEL_PATH" \ + "${MEGATRON_MODEL_ARGS[@]}" \ + --prompt "$PROMPT" \ + --max_new_tokens "$MAX_NEW_TOKENS" \ + --kv_cache_buffer_size_gb "$KV_CACHE_BUFFER_SIZE_GB" \ + --tp "$TP" --pp "$PP" --ep "$EP" --etp "$ETP" \ + --use-coordinator \ + --coordinator-host "$COORDINATOR_HOST" \ + --inference-moe-token-dispatcher-type "$INFERENCE_MOE_TOKEN_DISPATCHER_TYPE" \ + --distributed-timeout-minutes 90 +' diff --git a/examples/models/nemotron/nemotron_3/ultra/slurm_peft.sh b/examples/models/nemotron/nemotron_3/ultra/slurm_peft.sh new file mode 100755 index 0000000000..f84a8640d6 --- /dev/null +++ b/examples/models/nemotron/nemotron_3/ultra/slurm_peft.sh @@ -0,0 +1,179 @@ +#!/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. + +# ============================================================================== +# Nemotron 3 Ultra OpenMath LoRA PEFT +# +# Usage: +# 1. Modify the #SBATCH directives for your cluster. +# 2. Set CONTAINER_IMAGE and optional CONTAINER_MOUNTS. +# 3. Submit: sbatch slurm_peft.sh +# ============================================================================== + +#SBATCH --job-name=nemotron-ultra-openmath-peft +#SBATCH --nodes=4 +#SBATCH --ntasks-per-node=8 +#SBATCH --gpus-per-node=8 +#SBATCH --time=24:00:00 +#SBATCH --account= +#SBATCH --partition=batch +#SBATCH --output=logs/nemotron_ultra_openmath_peft_%j.log +#SBATCH --exclusive + +set -euo pipefail + +# ============================================================================== +# CONFIGURATION +# ============================================================================== + +WORKSPACE=${WORKSPACE:-/workspace} +WORKDIR=${WORKDIR:-/opt/Megatron-Bridge} +MODEL_HOME=${MODEL_HOME:-${WORKSPACE}/models/nvidia} + +HF_MODEL_PATH=${HF_MODEL_PATH:-nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16} +PRETRAINED_CHECKPOINT=${PRETRAINED_CHECKPOINT:-${MODEL_HOME}/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16-megatron} +MODEL_NAME=nemotron_3_ultra +DATASET_NAME=openmathinstruct2 +RECIPE_NAME=nemotron_3_ultra_peft_openmathinstruct2_packed_config +PEFT_SCHEME=lora + +SEQ_LENGTH=4096 +TRAIN_ITERS=1000 +GLOBAL_BATCH_SIZE=128 +MICRO_BATCH_SIZE=1 +EVAL_ITERS=32 +EVAL_INTERVAL=50 +LR_WARMUP_ITERS=250 +LR_DECAY_ITERS=$TRAIN_ITERS +LR=1e-4 +MIN_LR=1e-5 +SAVE_INTERVAL=${SAVE_INTERVAL:-$TRAIN_ITERS} +ASYNC_SAVE=${ASYNC_SAVE:-True} +ASYNC_STRATEGY=${ASYNC_STRATEGY:-nvrx} +LOG_INTERVAL=1 + +TP=${TP:-2} +PP=${PP:-4} +EP=${EP:-8} +ETP=${ETP:-1} +CP=${CP:-1} +SP=${SP:-True} +GPUS_PER_NODE=${GPUS_PER_NODE:-8} + +RECOMPUTE_GRANULARITY=${RECOMPUTE_GRANULARITY:-selective} +RECOMPUTE_MODULES=${RECOMPUTE_MODULES:-"[moe,layernorm,core_attn,moe_act,mlp,shared_experts]"} +RECOMPUTE_TAG=${RECOMPUTE_TAG:-recompute_selective_moe_layernorm_core_attn_moe_act_mlp_shared_experts} + +WANDB_ENTITY=${WANDB_ENTITY:-nvidia-nemo-fw-public} +WANDB_PROJECT=${WANDB_PROJECT:-megatron-bridge-nemotron-ultra} +WANDB_MODE=${WANDB_MODE:-disabled} + +CONTAINER_IMAGE=${CONTAINER_IMAGE:-} +CONTAINER_MOUNTS=${CONTAINER_MOUNTS:-} + +# ============================================================================== +# Environment Setup +# ============================================================================== + +[ -n "${HF_HOME:-}" ] && export HF_HOME +[ -n "${NEMO_HOME:-}" ] && export NEMO_HOME +[ -n "${UV_CACHE_DIR:-}" ] && export UV_CACHE_DIR +export WANDB_MODE +export TORCH_NCCL_AVOID_RECORD_STREAMS=1 +export NCCL_NVLS_ENABLE=0 +export NCCL_DEBUG=${NCCL_DEBUG:-WARN} +export NCCL_TIMEOUT=${NCCL_TIMEOUT:-1800000} +export PYTORCH_CUDA_ALLOC_CONF=${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True} +export CUDA_DEVICE_MAX_CONNECTIONS=${CUDA_DEVICE_MAX_CONNECTIONS:-1} + +# ============================================================================== +# Job Execution +# ============================================================================== + +mkdir -p logs + +if [ -z "$CONTAINER_IMAGE" ]; then + echo "ERROR: CONTAINER_IMAGE must be set." + exit 1 +fi + +if [ "$WANDB_MODE" = "online" ] && [ -z "${WANDB_API_KEY:-}" ]; then + echo "ERROR: WANDB_API_KEY must be visible in the submit environment for online W&B logging." + exit 2 +fi + +SAVE_DIR="${WORKSPACE}/results/${MODEL_NAME}_${DATASET_NAME}_peft_tp${TP}_pp${PP}_ep${EP}_sp${SP}_cp${CP}_${RECOMPUTE_TAG}_${SLURM_JOB_ID}" +WANDB_EXP_NAME="${MODEL_NAME}_${DATASET_NAME}_peft_tp${TP}_pp${PP}_ep${EP}_${RECOMPUTE_TAG}_${SLURM_JOB_ID}" + +CLI_OVERRIDES="\ + checkpoint.pretrained_checkpoint=${PRETRAINED_CHECKPOINT} \ + checkpoint.save=${SAVE_DIR} \ + checkpoint.save_interval=${SAVE_INTERVAL} \ + checkpoint.async_save=${ASYNC_SAVE} \ + checkpoint.async_strategy=${ASYNC_STRATEGY} \ + train.train_iters=${TRAIN_ITERS} \ + train.global_batch_size=${GLOBAL_BATCH_SIZE} \ + train.micro_batch_size=${MICRO_BATCH_SIZE} \ + validation.eval_interval=${EVAL_INTERVAL} \ + validation.eval_iters=${EVAL_ITERS} \ + scheduler.lr_warmup_iters=${LR_WARMUP_ITERS} \ + scheduler.lr_decay_iters=${LR_DECAY_ITERS} \ + optimizer.lr=${LR} \ + optimizer.min_lr=${MIN_LR} \ + logger.log_interval=${LOG_INTERVAL} \ + logger.tensorboard_dir=${SAVE_DIR}/tb_logs \ + logger.wandb_entity=${WANDB_ENTITY} \ + logger.wandb_project=${WANDB_PROJECT} \ + logger.wandb_exp_name=${WANDB_EXP_NAME} \ + logger.wandb_save_dir=${SAVE_DIR}/wandb \ + model.tensor_model_parallel_size=${TP} \ + model.pipeline_model_parallel_size=${PP} \ + model.expert_model_parallel_size=${EP} \ + model.expert_tensor_parallel_size=${ETP} \ + model.sequence_parallel=${SP} \ + model.context_parallel_size=${CP} \ + model.seq_length=${SEQ_LENGTH} \ + model.recompute_granularity=${RECOMPUTE_GRANULARITY} \ + model.recompute_modules=${RECOMPUTE_MODULES} \ + dist.distributed_timeout_minutes=90" + +CMD="cd ${WORKDIR} && mkdir -p ${WORKSPACE}/results ${SAVE_DIR}/wandb ${SAVE_DIR}/tb_logs && \ +export PYTHONPATH=${WORKDIR}/src:${WORKDIR}/3rdparty/Megatron-LM:\${PYTHONPATH:-} && \ +uv run --no-sync python scripts/training/run_recipe.py \ +--recipe ${RECIPE_NAME} --peft_scheme ${PEFT_SCHEME} --seq_length ${SEQ_LENGTH} --hf_path ${HF_MODEL_PATH} \ +${CLI_OVERRIDES}" + +SRUN_CMD="srun --mpi=pmix --no-kill --container-image=${CONTAINER_IMAGE} --no-container-mount-home" +if [ -n "$CONTAINER_MOUNTS" ]; then + SRUN_CMD="${SRUN_CMD} --container-mounts=${CONTAINER_MOUNTS}" +fi + +echo "======================================" +echo "Nemotron 3 Ultra OpenMath LoRA PEFT" +echo "======================================" +echo "Job ID: ${SLURM_JOB_ID}" +echo "Nodes: ${SLURM_JOB_NUM_NODES}" +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 "Recompute: ${RECOMPUTE_GRANULARITY} ${RECOMPUTE_MODULES}" +echo "Async save: ${ASYNC_SAVE} (${ASYNC_STRATEGY})" +echo "Save dir: ${SAVE_DIR}" +echo "W&B: ${WANDB_ENTITY}/${WANDB_PROJECT} (${WANDB_MODE})" +echo "======================================" + +$SRUN_CMD bash -c "$CMD" + +echo OPENMATH_PEFT_DONE diff --git a/examples/models/nemotron/nemotron_3/ultra/slurm_pretrain.sh b/examples/models/nemotron/nemotron_3/ultra/slurm_pretrain.sh new file mode 100755 index 0000000000..ac01321daa --- /dev/null +++ b/examples/models/nemotron/nemotron_3/ultra/slurm_pretrain.sh @@ -0,0 +1,180 @@ +#!/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. + +# ============================================================================== +# Nemotron 3 Ultra DCLM Pretraining +# +# Usage: +# 1. Modify the #SBATCH directives for your cluster. +# 2. Set CONTAINER_IMAGE and DCLM_DATA_DIR. +# 3. Optionally set CONTAINER_MOUNTS, WORKSPACE, WORKDIR, or HF_MODEL_PATH. +# 4. Submit: sbatch slurm_pretrain.sh +# ============================================================================== + +#SBATCH --job-name=nemotron-ultra-dclm-pretrain +#SBATCH --nodes=48 +#SBATCH --ntasks-per-node=8 +#SBATCH --gpus-per-node=8 +#SBATCH --time=24:00:00 +#SBATCH --account= +#SBATCH --partition=batch +#SBATCH --output=logs/nemotron_ultra_dclm_pretrain_%j.log +#SBATCH --exclusive + +set -euo pipefail + +# Required for most clusters: +CONTAINER_IMAGE=${CONTAINER_IMAGE:-} +DCLM_DATA_DIR=${DCLM_DATA_DIR:-} + +# Optional environment-specific paths: +WORKSPACE=${WORKSPACE:-/workspace} +WORKDIR=${WORKDIR:-/opt/Megatron-Bridge} +CONTAINER_MOUNTS=${CONTAINER_MOUNTS:-} +HF_MODEL_PATH=${HF_MODEL_PATH:-nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16} +EXTRA_OVERRIDES=${EXTRA_OVERRIDES:-} + +# Starter profile: +MODEL_NAME=nemotron_3_ultra +RECIPE_NAME=nemotron_3_ultra_pretrain_config +SEQ_LENGTH=4096 +DCLM_PATTERN=${DCLM_PATTERN:-dclm_01_*_text_document.bin} +DCLM_CACHE="${WORKSPACE}/data_cache/dclm" +TRAIN_ITERS=${TRAIN_ITERS:-1000} +SAVE_INTERVAL=${SAVE_INTERVAL:-250} +TP=${TP:-2} +PP=${PP:-12} +EP=${EP:-32} +ETP=${ETP:-1} +CP=${CP:-1} +SP=${SP:-True} +RECOMPUTE_MODULES=${RECOMPUTE_MODULES:-"[moe,layernorm,core_attn,moe_act,mlp,shared_experts]"} +GPUS_PER_NODE=${GPUS_PER_NODE:-8} +SAVE_DIR="${WORKSPACE}/results/${MODEL_NAME}_dclm_pretrain_tp${TP}_pp${PP}_ep${EP}_selective_${SLURM_JOB_ID}" +WANDB_ENTITY=${WANDB_ENTITY:-nvidia-nemo-fw-public} +WANDB_PROJECT=${WANDB_PROJECT:-megatron-bridge-nemotron-ultra} +WANDB_EXP_NAME="${MODEL_NAME}_dclm_pretrain_tp${TP}_pp${PP}_ep${EP}_selective_${SLURM_JOB_ID}" +WANDB_MODE=${WANDB_MODE:-disabled} + +[ -n "${HF_HOME:-}" ] && export HF_HOME +[ -n "${NEMO_HOME:-}" ] && export NEMO_HOME +[ -n "${UV_CACHE_DIR:-}" ] && export UV_CACHE_DIR +export WANDB_MODE +export TORCH_NCCL_AVOID_RECORD_STREAMS=1 +export NCCL_NVLS_ENABLE=0 +export NCCL_DEBUG=${NCCL_DEBUG:-WARN} +export NCCL_TIMEOUT=${NCCL_TIMEOUT:-1800000} +export PYTORCH_CUDA_ALLOC_CONF=${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True} +export CUDA_DEVICE_MAX_CONNECTIONS=${CUDA_DEVICE_MAX_CONNECTIONS:-1} + +mkdir -p logs + +if [ -z "$CONTAINER_IMAGE" ]; then + echo "ERROR: set CONTAINER_IMAGE to a valid sqsh image." + exit 1 +fi + +if [ -z "$DCLM_DATA_DIR" ]; then + echo "ERROR: set DCLM_DATA_DIR to a preprocessed DCLM directory." + exit 2 +fi + +if [ "$WANDB_MODE" = "online" ] && [ -z "${WANDB_API_KEY:-}" ]; then + echo "ERROR: WANDB_API_KEY must be visible in the submit environment for online W&B logging." + exit 3 +fi + +export DCLM_DATA_DIR DCLM_PATTERN DCLM_CACHE EXTRA_OVERRIDES HF_MODEL_PATH RECIPE_NAME SAVE_DIR SEQ_LENGTH +export TRAIN_ITERS SAVE_INTERVAL TP PP EP ETP CP SP RECOMPUTE_MODULES GPUS_PER_NODE +export WANDB_ENTITY WANDB_PROJECT WANDB_EXP_NAME WORKDIR WORKSPACE + +CMD=' +set -euo pipefail +cd "$WORKDIR" +mkdir -p "$WORKSPACE/results" "$SAVE_DIR/wandb" "$SAVE_DIR/tb_logs" "$DCLM_CACHE" +export PYTHONPATH="$WORKDIR/src:$WORKDIR/3rdparty/Megatron-LM:${PYTHONPATH:-}" + +BLEND_PATHS="" +shopt -s nullglob +for BIN_PATH in "$DCLM_DATA_DIR"/$DCLM_PATTERN; do + PREFIX=${BIN_PATH%.bin} + BLEND_PATHS="${BLEND_PATHS}\"${PREFIX}\"," +done +shopt -u nullglob +BLEND_PATHS="${BLEND_PATHS%,}" + +if [ -z "$BLEND_PATHS" ]; then + echo "ERROR: no DCLM shards matching ${DCLM_DATA_DIR}/${DCLM_PATTERN}" + exit 4 +fi + +uv run --no-sync python scripts/training/run_recipe.py \ + --recipe "$RECIPE_NAME" \ + --dataset llm-pretrain \ + --seq_length "$SEQ_LENGTH" \ + --hf_path "$HF_MODEL_PATH" \ + checkpoint.save="$SAVE_DIR" \ + checkpoint.save_interval="$SAVE_INTERVAL" \ + train.train_iters="$TRAIN_ITERS" \ + train.global_batch_size=128 \ + train.micro_batch_size=1 \ + validation.eval_interval=100 \ + validation.eval_iters=10 \ + logger.log_interval=1 \ + logger.tensorboard_dir="$SAVE_DIR/tb_logs" \ + logger.wandb_entity="$WANDB_ENTITY" \ + logger.wandb_project="$WANDB_PROJECT" \ + logger.wandb_exp_name="$WANDB_EXP_NAME" \ + logger.wandb_save_dir="$SAVE_DIR/wandb" \ + model.tensor_model_parallel_size="$TP" \ + model.pipeline_model_parallel_size="$PP" \ + model.expert_model_parallel_size="$EP" \ + model.expert_tensor_parallel_size="$ETP" \ + model.sequence_parallel="$SP" \ + model.context_parallel_size="$CP" \ + model.seq_length="$SEQ_LENGTH" \ + dataset.sequence_length="$SEQ_LENGTH" \ + model.recompute_granularity=selective \ + model.recompute_modules=${RECOMPUTE_MODULES} \ + dist.distributed_timeout_minutes=90 \ + "dataset.blend=[[${BLEND_PATHS}],null]" \ + dataset.split=\"9999,8,2\" \ + dataset.path_to_cache="$DCLM_CACHE" \ + ${EXTRA_OVERRIDES} +' + +SRUN_CMD="srun --mpi=pmix --no-kill --container-image=${CONTAINER_IMAGE} --no-container-mount-home" +if [ -n "$CONTAINER_MOUNTS" ]; then + SRUN_CMD="${SRUN_CMD} --container-mounts=${CONTAINER_MOUNTS}" +fi + +echo "======================================" +echo "Nemotron 3 Ultra DCLM Pretraining" +echo "======================================" +echo "Job ID: ${SLURM_JOB_ID}" +echo "Nodes: ${SLURM_JOB_NUM_NODES}" +echo "GPUs/node: ${GPUS_PER_NODE}" +echo "Recipe: ${RECIPE_NAME}" +echo "HF model: ${HF_MODEL_PATH}" +echo "DCLM data: ${DCLM_DATA_DIR}" +echo "Parallelism: TP=${TP} PP=${PP} EP=${EP} ETP=${ETP} CP=${CP} SP=${SP}" +echo "Recompute: selective ${RECOMPUTE_MODULES}" +echo "Save dir: ${SAVE_DIR}" +echo "W&B: ${WANDB_ENTITY}/${WANDB_PROJECT} (${WANDB_MODE})" +echo "======================================" + +$SRUN_CMD bash -c "$CMD" + +echo DCLM_PRETRAIN_DONE diff --git a/examples/models/nemotron/nemotron_3/ultra/slurm_sft.sh b/examples/models/nemotron/nemotron_3/ultra/slurm_sft.sh new file mode 100755 index 0000000000..76450b17f4 --- /dev/null +++ b/examples/models/nemotron/nemotron_3/ultra/slurm_sft.sh @@ -0,0 +1,186 @@ +#!/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. + +# ============================================================================== +# Nemotron 3 Ultra OpenMath Full SFT +# +# Usage: +# 1. Modify the #SBATCH directives for your cluster. +# 2. Set CONTAINER_IMAGE and optional CONTAINER_MOUNTS. +# 3. Submit: sbatch slurm_sft.sh +# ============================================================================== + +#SBATCH --job-name=nemotron-ultra-openmath-sft +#SBATCH --nodes=48 +#SBATCH --ntasks-per-node=8 +#SBATCH --gpus-per-node=8 +#SBATCH --time=24:00:00 +#SBATCH --account= +#SBATCH --partition=batch +#SBATCH --output=logs/nemotron_ultra_openmath_sft_%j.log +#SBATCH --exclusive + +set -euo pipefail + +# ============================================================================== +# CONFIGURATION +# ============================================================================== + +WORKSPACE=${WORKSPACE:-/workspace} +WORKDIR=${WORKDIR:-/opt/Megatron-Bridge} +MODEL_HOME=${MODEL_HOME:-${WORKSPACE}/models/nvidia} + +HF_MODEL_PATH=${HF_MODEL_PATH:-nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16} +PRETRAINED_CHECKPOINT=${PRETRAINED_CHECKPOINT:-${MODEL_HOME}/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16-megatron} +MODEL_NAME=nemotron_3_ultra +DATASET_NAME=openmathinstruct2 +RECIPE_NAME=nemotron_3_ultra_sft_openmathinstruct2_packed_config + +SEQ_LENGTH=4096 +TRAIN_ITERS=1000 +GLOBAL_BATCH_SIZE=128 +MICRO_BATCH_SIZE=1 +EVAL_ITERS=32 +EVAL_INTERVAL=50 +LR_WARMUP_ITERS=250 +LR_DECAY_ITERS=$TRAIN_ITERS +LR=5e-6 +MIN_LR=5e-7 +SAVE_INTERVAL=$TRAIN_ITERS +LOG_INTERVAL=1 + +TP=${TP:-2} +PP=${PP:-12} +EP=${EP:-32} +ETP=${ETP:-1} +CP=${CP:-1} +SP=${SP:-True} +GPUS_PER_NODE=${GPUS_PER_NODE:-8} + +RECOMPUTE_GRANULARITY=selective +RECOMPUTE_MODULES="[moe,layernorm,core_attn,moe_act]" +RECOMPUTE_TAG=recompute_selective + +WANDB_ENTITY=${WANDB_ENTITY:-nvidia-nemo-fw-public} +WANDB_PROJECT=${WANDB_PROJECT:-megatron-bridge-nemotron-ultra} +WANDB_MODE=${WANDB_MODE:-disabled} + +CONTAINER_IMAGE=${CONTAINER_IMAGE:-} +CONTAINER_MOUNTS=${CONTAINER_MOUNTS:-} + +# ============================================================================== +# Environment Setup +# ============================================================================== + +[ -n "${HF_HOME:-}" ] && export HF_HOME +[ -n "${NEMO_HOME:-}" ] && export NEMO_HOME +[ -n "${UV_CACHE_DIR:-}" ] && export UV_CACHE_DIR +export WANDB_MODE +export TORCH_NCCL_AVOID_RECORD_STREAMS=1 +export NCCL_NVLS_ENABLE=0 +export NCCL_DEBUG=${NCCL_DEBUG:-WARN} +export NCCL_TIMEOUT=${NCCL_TIMEOUT:-1800000} +export PYTORCH_CUDA_ALLOC_CONF=${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True} +export CUDA_DEVICE_MAX_CONNECTIONS=${CUDA_DEVICE_MAX_CONNECTIONS:-1} + +# ============================================================================== +# Job Execution +# ============================================================================== + +mkdir -p logs + +if [ -z "$CONTAINER_IMAGE" ]; then + echo "ERROR: CONTAINER_IMAGE must be set." + exit 1 +fi + +if [ "$WANDB_MODE" = "online" ] && [ -z "${WANDB_API_KEY:-}" ]; then + echo "ERROR: WANDB_API_KEY must be visible in the submit environment for online W&B logging." + exit 2 +fi + +SAVE_DIR="${WORKSPACE}/results/${MODEL_NAME}_${DATASET_NAME}_sft_tp${TP}_pp${PP}_ep${EP}_sp${SP}_cp${CP}_${RECOMPUTE_TAG}_${SLURM_JOB_ID}" +WANDB_EXP_NAME="${MODEL_NAME}_${DATASET_NAME}_sft_tp${TP}_pp${PP}_ep${EP}_${RECOMPUTE_TAG}_${SLURM_JOB_ID}" + +CLI_OVERRIDES="\ + checkpoint.pretrained_checkpoint=${PRETRAINED_CHECKPOINT} \ + checkpoint.save=${SAVE_DIR} \ + checkpoint.save_interval=${SAVE_INTERVAL} \ + train.train_iters=${TRAIN_ITERS} \ + train.global_batch_size=${GLOBAL_BATCH_SIZE} \ + train.micro_batch_size=${MICRO_BATCH_SIZE} \ + validation.eval_interval=${EVAL_INTERVAL} \ + validation.eval_iters=${EVAL_ITERS} \ + scheduler.lr_warmup_iters=${LR_WARMUP_ITERS} \ + scheduler.lr_decay_iters=${LR_DECAY_ITERS} \ + optimizer.lr=${LR} \ + optimizer.min_lr=${MIN_LR} \ + logger.log_interval=${LOG_INTERVAL} \ + logger.tensorboard_dir=${SAVE_DIR}/tb_logs \ + logger.wandb_entity=${WANDB_ENTITY} \ + logger.wandb_project=${WANDB_PROJECT} \ + logger.wandb_exp_name=${WANDB_EXP_NAME} \ + logger.wandb_save_dir=${SAVE_DIR}/wandb \ + model.tensor_model_parallel_size=${TP} \ + model.pipeline_model_parallel_size=${PP} \ + model.expert_model_parallel_size=${EP} \ + model.expert_tensor_parallel_size=${ETP} \ + model.sequence_parallel=${SP} \ + model.context_parallel_size=${CP} \ + model.seq_length=${SEQ_LENGTH} \ + model.recompute_granularity=${RECOMPUTE_GRANULARITY} \ + model.recompute_modules=${RECOMPUTE_MODULES} \ + dist.distributed_timeout_minutes=90" + +CMD="cd ${WORKDIR} && mkdir -p ${WORKSPACE}/results ${SAVE_DIR}/wandb ${SAVE_DIR}/tb_logs && \ +export PYTHONPATH=${WORKDIR}/src:${WORKDIR}/3rdparty/Megatron-LM:\${PYTHONPATH:-} && \ +uv run --no-sync python scripts/training/run_recipe.py \ +--recipe ${RECIPE_NAME} --seq_length ${SEQ_LENGTH} --hf_path ${HF_MODEL_PATH} \ +${CLI_OVERRIDES}" + +SRUN_CMD="srun --mpi=pmix --no-kill --container-image=${CONTAINER_IMAGE} --no-container-mount-home" +if [ -n "$CONTAINER_MOUNTS" ]; then + SRUN_CMD="${SRUN_CMD} --container-mounts=${CONTAINER_MOUNTS}" +fi + +echo "======================================" +echo "Nemotron 3 Ultra OpenMath Full SFT" +echo "======================================" +echo "Job ID: ${SLURM_JOB_ID}" +echo "Nodes: ${SLURM_JOB_NUM_NODES}" +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 "Recompute: ${RECOMPUTE_GRANULARITY} ${RECOMPUTE_MODULES}" +echo "Save dir: ${SAVE_DIR}" +echo "W&B: ${WANDB_ENTITY}/${WANDB_PROJECT} (${WANDB_MODE})" +echo "======================================" + +$SRUN_CMD bash -c "$CMD" + +LATEST_ITER_FILE="${SAVE_DIR}/latest_checkpointed_iteration.txt" +if [ -f "$LATEST_ITER_FILE" ]; then + LATEST_ITER=$(tr -d '[:space:]' < "$LATEST_ITER_FILE") + if [[ "$LATEST_ITER" =~ ^[0-9]+$ ]]; then + LATEST_DIR=$(printf "iter_%07d" "$LATEST_ITER") + find "$SAVE_DIR" -mindepth 1 -maxdepth 1 -type d -name "iter_*" ! -name "$LATEST_DIR" -exec rm -rf {} + + else + echo "Skipping SFT intermediate checkpoint cleanup: latest checkpoint marker is not numeric: ${LATEST_ITER}" + fi +else + echo "Skipping SFT intermediate checkpoint cleanup: missing ${LATEST_ITER_FILE}" +fi + +echo OPENMATH_SFT_DONE diff --git a/src/megatron/bridge/recipes/nemotronh/__init__.py b/src/megatron/bridge/recipes/nemotronh/__init__.py index 46ae25a152..91bfe4feeb 100644 --- a/src/megatron/bridge/recipes/nemotronh/__init__.py +++ b/src/megatron/bridge/recipes/nemotronh/__init__.py @@ -26,6 +26,11 @@ nemotron_3_super_pretrain_config, nemotron_3_super_sft_config, ) +from megatron.bridge.recipes.nemotronh.nemotron_3_ultra import ( + nemotron_3_ultra_peft_openmathinstruct2_packed_config, + nemotron_3_ultra_pretrain_config, + nemotron_3_ultra_sft_openmathinstruct2_packed_config, +) from megatron.bridge.recipes.nemotronh.nemotron_nano_v2 import ( nemotron_nano_9b_v2_peft_config, nemotron_nano_9b_v2_pretrain_config, @@ -81,4 +86,8 @@ "nemotron_3_super_pretrain_config", "nemotron_3_super_sft_config", "nemotron_3_super_peft_config", + # Nemotron 3 Ultra models + "nemotron_3_ultra_pretrain_config", + "nemotron_3_ultra_sft_openmathinstruct2_packed_config", + "nemotron_3_ultra_peft_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 new file mode 100644 index 0000000000..00c5d31c00 --- /dev/null +++ b/src/megatron/bridge/recipes/nemotronh/nemotron_3_ultra.py @@ -0,0 +1,314 @@ +# 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. + +import torch + +from megatron.bridge import AutoBridge +from megatron.bridge.peft.base import PEFT +from megatron.bridge.recipes.common import _peft_common, _pretrain_common, _sft_common +from megatron.bridge.recipes.utils.finetune_utils import default_openmathinstruct2_config, default_peft_config +from megatron.bridge.training.config import ConfigContainer + + +NEMOTRON_3_ULTRA_HF_MODEL_ID = "nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16" +NEMOTRON_3_ULTRA_TOKENIZER_NAME = "nvidia--NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16" + + +def nemotron_3_ultra_pretrain_config( + *, + hf_path: str | None = None, + seq_length: int = 8192, +) -> ConfigContainer: + """Return a pre-training config for Nemotron 3 Ultra. + + Args: + hf_path: Optional Hugging Face model ID or local snapshot path. + seq_length: Sequence length for model and dataset settings. + + Returns: + Pre-training configuration for Nemotron 3 Ultra. + """ + cfg = _pretrain_common() + model_source = hf_path or NEMOTRON_3_ULTRA_HF_MODEL_ID + + cfg.model = AutoBridge.from_hf_pretrained(model_source).to_megatron_provider(load_weights=False) + cfg.model.tensor_model_parallel_size = 1 + cfg.model.pipeline_model_parallel_size = 3 + 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 = 8 + cfg.model.pipeline_model_parallel_layout = None + cfg.model.seq_length = seq_length + cfg.model.apply_rope_fusion = False + cfg.model.attention_backend = "fused" + cfg.model.gradient_accumulation_fusion = True + cfg.model.init_method_std = 0.014 + cfg.model.use_fused_weighted_squared_relu = True + cfg.model.calculate_per_token_loss = True + cfg.model.moe_token_dispatcher_type = "alltoall" + cfg.model.moe_shared_expert_overlap = False + cfg.model.moe_flex_dispatcher_backend = "hybridep" + cfg.model.cuda_graph_impl = "none" + cfg.model.cuda_graph_scope = [] + cfg.model.mtp_num_layers = 2 + cfg.model.keep_mtp_spec_in_bf16 = True + cfg.model.mtp_loss_scaling_factor = 0.3 + cfg.model.mtp_use_repeated_layer = True + cfg.model.use_te_rng_tracker = True + + cfg.tokenizer.tokenizer_model = model_source + cfg.dataset.seq_length = seq_length + cfg.dataset.blend = None + cfg.dataset.num_workers = 1 + cfg.dataset.mmap_bin_files = False + + cfg.train.train_iters = 39735 + cfg.train.global_batch_size = 3072 + cfg.train.micro_batch_size = 1 + cfg.train.manual_gc = False + cfg.train.manual_gc_interval = 0 + cfg.validation.eval_interval = 1000 + + cfg.model.transformer_impl = "transformer_engine" + cfg.model.cross_entropy_fusion_impl = "te" + cfg.mixed_precision = "bf16_mixed" + + cfg.optimizer.lr = 2.5e-4 + cfg.optimizer.min_lr = 2.5e-4 + cfg.optimizer.weight_decay = 0.1 + cfg.optimizer.adam_beta1 = 0.9 + cfg.optimizer.adam_beta2 = 0.95 + cfg.optimizer.adam_eps = 1e-8 + cfg.scheduler.lr_warmup_iters = 0 + cfg.scheduler.start_weight_decay = 0.1 + cfg.scheduler.end_weight_decay = 0.1 + cfg.scheduler.lr_decay_style = "constant" + + cfg.checkpoint.save_interval = 200 + cfg.checkpoint.ckpt_assume_constant_structure = True + cfg.checkpoint.dist_ckpt_strictness = "log_all" + cfg.checkpoint.async_save = True + + cfg.ddp.overlap_grad_reduce = True + cfg.ddp.overlap_param_gather = True + cfg.ddp.check_for_nan_in_grad = True + cfg.ddp.use_distributed_optimizer = True + cfg.ddp.average_in_collective = False + + return cfg + + +def nemotron_3_ultra_sft_openmathinstruct2_packed_config( + *, + hf_path: str | None = None, + seq_length: int = 4096, +) -> ConfigContainer: + """Return a packed OpenMathInstruct-2 full SFT config for Nemotron 3 Ultra. + + Args: + hf_path: Optional Hugging Face model ID or local snapshot path. + seq_length: Packed sequence length. + + Returns: + Full-parameter SFT configuration for OpenMathInstruct-2. + """ + cfg = _sft_common() + model_source = hf_path or NEMOTRON_3_ULTRA_HF_MODEL_ID + + cfg.model = AutoBridge.from_hf_pretrained(model_source).to_megatron_provider(load_weights=False) + cfg.model.tensor_model_parallel_size = 2 + cfg.model.pipeline_model_parallel_size = 6 + 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.pipeline_model_parallel_layout = None + cfg.model.seq_length = seq_length + cfg.model.apply_rope_fusion = False + cfg.model.attention_backend = "fused" + cfg.model.gradient_accumulation_fusion = True + cfg.model.init_method_std = 0.014 + cfg.model.use_fused_weighted_squared_relu = True + cfg.model.calculate_per_token_loss = True + cfg.model.moe_token_dispatcher_type = "alltoall" + cfg.model.moe_shared_expert_overlap = False + cfg.model.moe_flex_dispatcher_backend = "hybridep" + cfg.model.cuda_graph_impl = "none" + cfg.model.cuda_graph_scope = [] + cfg.model.mtp_num_layers = 2 + cfg.model.keep_mtp_spec_in_bf16 = True + cfg.model.mtp_loss_scaling_factor = 0.3 + cfg.model.mtp_use_repeated_layer = True + cfg.model.use_te_rng_tracker = True + 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.tokenizer.tokenizer_model = model_source + cfg.dataset = default_openmathinstruct2_config(seq_length=seq_length, packed_sequence=True) + if cfg.dataset.packed_sequence_specs is not None: + cfg.dataset.packed_sequence_specs.packed_sequence_size = seq_length + cfg.dataset.packed_sequence_specs.tokenizer_model_name = NEMOTRON_3_ULTRA_TOKENIZER_NAME + + cfg.train.train_iters = 1000 + cfg.train.global_batch_size = 128 + cfg.train.micro_batch_size = 1 + cfg.validation.eval_interval = 50 + cfg.validation.eval_iters = 32 + + cfg.optimizer.lr = 5e-6 + cfg.optimizer.min_lr = 5e-7 + cfg.optimizer.adam_beta1 = 0.9 + cfg.optimizer.adam_beta2 = 0.98 + cfg.optimizer.adam_eps = 1e-8 + cfg.optimizer.weight_decay = 0.1 + cfg.scheduler.start_weight_decay = 0.1 + cfg.scheduler.end_weight_decay = 0.1 + cfg.scheduler.lr_decay_style = "cosine" + cfg.scheduler.lr_warmup_iters = 250 + cfg.scheduler.lr_decay_iters = 1000 + + cfg.checkpoint.save_interval = 250 + cfg.checkpoint.ckpt_format = "torch_dist" + cfg.checkpoint.dist_ckpt_strictness = "log_all" + cfg.checkpoint.ckpt_assume_constant_structure = True + cfg.checkpoint.async_save = True + + cfg.logger.log_interval = 1 + cfg.rng.seed = 5678 + + cfg.ddp.check_for_nan_in_grad = True + cfg.ddp.grad_reduce_in_fp32 = True + cfg.ddp.overlap_grad_reduce = True + cfg.ddp.overlap_param_gather = True + cfg.ddp.use_distributed_optimizer = True + + return cfg + + +def nemotron_3_ultra_peft_openmathinstruct2_packed_config( + *, + peft: str | PEFT | None = "lora", + hf_path: str | None = None, + seq_length: int = 4096, +) -> ConfigContainer: + """Return a packed OpenMathInstruct-2 PEFT config for Nemotron 3 Ultra. + + Args: + peft: PEFT scheme, PEFT instance, or "none". + hf_path: Optional Hugging Face model ID or local snapshot path. + seq_length: Packed sequence length. + + Returns: + PEFT configuration for OpenMathInstruct-2. + """ + cfg = _peft_common() + model_source = hf_path or NEMOTRON_3_ULTRA_HF_MODEL_ID + + cfg.model = AutoBridge.from_hf_pretrained(model_source).to_megatron_provider(load_weights=False) + cfg.model.tensor_model_parallel_size = 2 + cfg.model.pipeline_model_parallel_size = 4 + 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 = 8 + cfg.model.pipeline_model_parallel_layout = None + cfg.model.seq_length = seq_length + cfg.model.apply_rope_fusion = False + cfg.model.attention_backend = "fused" + cfg.model.gradient_accumulation_fusion = True + cfg.model.init_method_std = 0.014 + cfg.model.use_fused_weighted_squared_relu = True + cfg.model.calculate_per_token_loss = True + cfg.model.moe_token_dispatcher_type = "alltoall" + cfg.model.moe_shared_expert_overlap = False + cfg.model.moe_flex_dispatcher_backend = "hybridep" + cfg.model.cuda_graph_impl = "none" + cfg.model.cuda_graph_scope = [] + cfg.model.mtp_num_layers = 2 + cfg.model.keep_mtp_spec_in_bf16 = True + cfg.model.mtp_loss_scaling_factor = 0.3 + cfg.model.mtp_use_repeated_layer = True + cfg.model.use_te_rng_tracker = True + 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", "mlp", "shared_experts"] + + target_modules = [ + "linear_qkv", + "linear_proj", + "linear_fc1", + "linear_fc2", + "in_proj", + "out_proj", + ] + cfg.peft = default_peft_config(peft, target_modules=target_modules) + + cfg.tokenizer.tokenizer_model = model_source + cfg.dataset = default_openmathinstruct2_config(seq_length=seq_length, packed_sequence=True) + if cfg.dataset.packed_sequence_specs is not None: + cfg.dataset.packed_sequence_specs.packed_sequence_size = seq_length + cfg.dataset.packed_sequence_specs.tokenizer_model_name = NEMOTRON_3_ULTRA_TOKENIZER_NAME + + cfg.train.train_iters = 1000 + cfg.train.global_batch_size = 128 + cfg.train.micro_batch_size = 1 + cfg.validation.eval_interval = 50 + cfg.validation.eval_iters = 32 + + cfg.optimizer.lr = 1e-4 + cfg.optimizer.min_lr = 1e-5 + cfg.optimizer.adam_beta1 = 0.9 + cfg.optimizer.adam_beta2 = 0.98 + cfg.optimizer.adam_eps = 1e-8 + cfg.optimizer.weight_decay = 0.1 + cfg.scheduler.start_weight_decay = 0.1 + cfg.scheduler.end_weight_decay = 0.1 + cfg.scheduler.lr_decay_style = "cosine" + cfg.scheduler.lr_warmup_iters = 250 + cfg.scheduler.lr_decay_iters = 1000 + + cfg.checkpoint.save_interval = 250 + cfg.checkpoint.ckpt_format = "torch_dist" + cfg.checkpoint.dist_ckpt_strictness = "log_all" + cfg.checkpoint.ckpt_assume_constant_structure = True + cfg.checkpoint.async_save = True + cfg.checkpoint.async_strategy = "nvrx" + + cfg.logger.log_interval = 1 + cfg.rng.seed = 5678 + + cfg.ddp.check_for_nan_in_grad = True + cfg.ddp.grad_reduce_in_fp32 = True + cfg.ddp.overlap_grad_reduce = True + cfg.ddp.overlap_param_gather = True + cfg.ddp.use_distributed_optimizer = True + + return cfg + + +__all__ = [ + "nemotron_3_ultra_pretrain_config", + "nemotron_3_ultra_sft_openmathinstruct2_packed_config", + "nemotron_3_ultra_peft_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 new file mode 100644 index 0000000000..99ab262e6d --- /dev/null +++ b/tests/unit_tests/recipes/nemotronh/test_nemotron_3_ultra.py @@ -0,0 +1,116 @@ +# 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. + +import importlib + +import pytest + +from megatron.bridge.recipes.nemotronh.nemotron_3_ultra import ( + NEMOTRON_3_ULTRA_TOKENIZER_NAME, + nemotron_3_ultra_peft_openmathinstruct2_packed_config, + nemotron_3_ultra_sft_openmathinstruct2_packed_config, +) + + +class _FakeUltraProvider: + """Fake model provider for testing recipe field overrides without HF Hub I/O.""" + + def __init__(self) -> None: + self.vocab_size = 256 + + def finalize(self) -> None: + return None + + +class _FakeAutoBridge: + """Fake AutoBridge that returns an Ultra provider without loading a model.""" + + @classmethod + def from_hf_pretrained(cls, *args, **kwargs): + return cls() + + def to_megatron_provider(self, *args, **kwargs): + return _FakeUltraProvider() + + +@pytest.fixture(autouse=True) +def _patch_autobridge(monkeypatch: pytest.MonkeyPatch) -> None: + """Patch AutoBridge in the recipe module to avoid Hugging Face access.""" + mod = importlib.import_module("megatron.bridge.recipes.nemotronh.nemotron_3_ultra") + monkeypatch.setattr(mod, "AutoBridge", _FakeAutoBridge) + + +@pytest.mark.unit +def test_openmath_sft_uses_initial_parallelism_values() -> 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.sequence_parallel is True + assert cfg.model.virtual_pipeline_model_parallel_size is 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.train.train_iters == 1000 + assert cfg.train.global_batch_size == 128 + assert cfg.dataset.dataset_name == "nvidia/OpenMathInstruct-2" + assert cfg.dataset.packed_sequence_specs.packed_sequence_size == 4096 + assert cfg.dataset.packed_sequence_specs.tokenizer_model_name == NEMOTRON_3_ULTRA_TOKENIZER_NAME + + +@pytest.mark.unit +def test_openmath_peft_uses_validated_parallelism_values() -> None: + cfg = nemotron_3_ultra_peft_openmathinstruct2_packed_config() + + assert cfg.model.tensor_model_parallel_size == 2 + assert cfg.model.pipeline_model_parallel_size == 4 + assert cfg.model.expert_model_parallel_size == 8 + assert cfg.model.sequence_parallel is True + assert cfg.model.virtual_pipeline_model_parallel_size is 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", "mlp", "shared_experts"] + + assert cfg.optimizer.lr == 1e-4 + assert cfg.optimizer.min_lr == 1e-5 + assert cfg.train.train_iters == 1000 + assert cfg.train.global_batch_size == 128 + + +@pytest.mark.unit +def test_openmath_peft_none_disables_adapter() -> None: + cfg = nemotron_3_ultra_peft_openmathinstruct2_packed_config(peft="none") + assert cfg.peft is None + + +@pytest.mark.unit +def test_openmath_peft_recompute_modules_are_not_shared() -> None: + cfg = nemotron_3_ultra_peft_openmathinstruct2_packed_config() + cfg.model.recompute_modules.append("sentinel") + + fresh_cfg = nemotron_3_ultra_peft_openmathinstruct2_packed_config() + assert fresh_cfg.model.recompute_modules == ["moe", "layernorm", "core_attn", "moe_act", "mlp", "shared_experts"] + + +@pytest.mark.unit +def test_openmath_sft_recompute_modules_are_not_shared() -> None: + cfg = nemotron_3_ultra_sft_openmathinstruct2_packed_config() + 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"] From 2f7475008c108eebc4e82ff3f09cc2a5017c750b Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Mon, 8 Jun 2026 19:27:08 -0700 Subject: [PATCH 02/10] docs(nemotronh): remove Ultra day-0 setup Signed-off-by: Chen Cui --- .../nemotron/nemotron_3/ultra/README.md | 54 ------------------- 1 file changed, 54 deletions(-) diff --git a/examples/models/nemotron/nemotron_3/ultra/README.md b/examples/models/nemotron/nemotron_3/ultra/README.md index 54c5ec9788..ff95279dcd 100644 --- a/examples/models/nemotron/nemotron_3/ultra/README.md +++ b/examples/models/nemotron/nemotron_3/ultra/README.md @@ -55,60 +55,6 @@ 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. -## Day-0 Code - -Start from the NeMo 26.04.01 container: - -```bash -docker run --rm -it --gpus all --shm-size=24g \ - -w /workspace \ - nvcr.io/nvidia/nemo:26.04.01 \ - bash -``` - -Inside the container, clone the Day-0 Megatron Bridge branch and initialize the -bundled Megatron-Core submodule: - -```bash -export WORKSPACE=/workspace -export MEGATRON_BRIDGE_REPO=https://github.com/NVIDIA-NeMo/Megatron-Bridge.git -export MEGATRON_BRIDGE_BRANCH=nemotron_3_ultra - -cd "$WORKSPACE" -git clone --branch "$MEGATRON_BRIDGE_BRANCH" "$MEGATRON_BRIDGE_REPO" -cd Megatron-Bridge - -git submodule update --init --recursive --depth 1 3rdparty/Megatron-LM - -uv lock -uv sync --extra ssm -``` - -Run dependency sync during setup. The Slurm job scripts use `uv run --no-sync` -so they do not download dependencies or mutate the Python environment during -allocated jobs. - -Verify that `megatron.core` and `megatron.bridge` resolve to this checkout: - -```bash -uv run python - <<'PY' -import megatron.bridge -import megatron.core -from megatron.bridge.recipes.nemotronh import nemotron_3_ultra_pretrain_config - -print("core:", list(megatron.core.__path__)) -print("bridge:", list(megatron.bridge.__path__)) -print("recipe:", nemotron_3_ultra_pretrain_config.__name__) -PY -``` - -Expected paths should include: - -```text -/workspace/Megatron-Bridge/3rdparty/Megatron-LM/megatron/core -/workspace/Megatron-Bridge/src/megatron/bridge -``` - ## Checkpoint Conversion Use [conversion.sh](conversion.sh) for CPU checkpoint import when the node has From 0cdae372c66bdbf9ca186f100935dc0ef3ce4d0b Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Tue, 9 Jun 2026 15:36:20 -0700 Subject: [PATCH 03/10] test(nemotronh): cover Ultra pretrain recipe Signed-off-by: Chen Cui --- .../nemotronh/test_nemotron_3_ultra.py | 26 +++++++++++++++++++ 1 file changed, 26 insertions(+) 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 99ab262e6d..4af09ac867 100644 --- a/tests/unit_tests/recipes/nemotronh/test_nemotron_3_ultra.py +++ b/tests/unit_tests/recipes/nemotronh/test_nemotron_3_ultra.py @@ -19,6 +19,7 @@ from megatron.bridge.recipes.nemotronh.nemotron_3_ultra import ( NEMOTRON_3_ULTRA_TOKENIZER_NAME, nemotron_3_ultra_peft_openmathinstruct2_packed_config, + nemotron_3_ultra_pretrain_config, nemotron_3_ultra_sft_openmathinstruct2_packed_config, ) @@ -51,6 +52,31 @@ def _patch_autobridge(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(mod, "AutoBridge", _FakeAutoBridge) +@pytest.mark.unit +def test_pretrain_uses_initial_parallelism_values() -> None: + cfg = nemotron_3_ultra_pretrain_config() + + assert cfg.model.tensor_model_parallel_size == 1 + assert cfg.model.pipeline_model_parallel_size == 3 + assert cfg.model.expert_model_parallel_size == 8 + assert cfg.model.sequence_parallel is True + assert cfg.model.virtual_pipeline_model_parallel_size is None + assert cfg.model.mtp_num_layers == 2 + assert cfg.model.mtp_loss_scaling_factor == 0.3 + assert cfg.model.mtp_use_repeated_layer is True + + assert cfg.optimizer.lr == 2.5e-4 + assert cfg.optimizer.min_lr == 2.5e-4 + assert cfg.optimizer.weight_decay == 0.1 + assert cfg.scheduler.lr_decay_style == "constant" + assert cfg.scheduler.lr_warmup_iters == 0 + + assert cfg.train.global_batch_size == 3072 + assert cfg.train.micro_batch_size == 1 + assert cfg.dataset.seq_length == 8192 + assert cfg.dataset.blend is None + + @pytest.mark.unit def test_openmath_sft_uses_initial_parallelism_values() -> None: cfg = nemotron_3_ultra_sft_openmathinstruct2_packed_config() From 79094ec516fc8acae8200230d11de99c57cdc9be Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Mon, 15 Jun 2026 13:55:43 -0700 Subject: [PATCH 04/10] fix(nemotronh): clean up Ultra mainline examples Signed-off-by: Chen Cui --- README.md | 2 +- examples/models/nemotron/nemotron_3/ultra/slurm_peft.sh | 4 ++-- examples/models/nemotron/nemotron_3/ultra/slurm_sft.sh | 6 +++--- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index 6b7a09d716..77a4936e8e 100644 --- a/README.md +++ b/README.md @@ -11,7 +11,7 @@ ## 📣 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. diff --git a/examples/models/nemotron/nemotron_3/ultra/slurm_peft.sh b/examples/models/nemotron/nemotron_3/ultra/slurm_peft.sh index f84a8640d6..f8fa1dd389 100755 --- a/examples/models/nemotron/nemotron_3/ultra/slurm_peft.sh +++ b/examples/models/nemotron/nemotron_3/ultra/slurm_peft.sh @@ -50,13 +50,13 @@ RECIPE_NAME=nemotron_3_ultra_peft_openmathinstruct2_packed_config PEFT_SCHEME=lora SEQ_LENGTH=4096 -TRAIN_ITERS=1000 +TRAIN_ITERS=${TRAIN_ITERS:-1000} GLOBAL_BATCH_SIZE=128 MICRO_BATCH_SIZE=1 EVAL_ITERS=32 EVAL_INTERVAL=50 LR_WARMUP_ITERS=250 -LR_DECAY_ITERS=$TRAIN_ITERS +LR_DECAY_ITERS=${LR_DECAY_ITERS:-$TRAIN_ITERS} LR=1e-4 MIN_LR=1e-5 SAVE_INTERVAL=${SAVE_INTERVAL:-$TRAIN_ITERS} diff --git a/examples/models/nemotron/nemotron_3/ultra/slurm_sft.sh b/examples/models/nemotron/nemotron_3/ultra/slurm_sft.sh index 76450b17f4..459b014a3c 100755 --- a/examples/models/nemotron/nemotron_3/ultra/slurm_sft.sh +++ b/examples/models/nemotron/nemotron_3/ultra/slurm_sft.sh @@ -49,16 +49,16 @@ DATASET_NAME=openmathinstruct2 RECIPE_NAME=nemotron_3_ultra_sft_openmathinstruct2_packed_config SEQ_LENGTH=4096 -TRAIN_ITERS=1000 +TRAIN_ITERS=${TRAIN_ITERS:-1000} GLOBAL_BATCH_SIZE=128 MICRO_BATCH_SIZE=1 EVAL_ITERS=32 EVAL_INTERVAL=50 LR_WARMUP_ITERS=250 -LR_DECAY_ITERS=$TRAIN_ITERS +LR_DECAY_ITERS=${LR_DECAY_ITERS:-$TRAIN_ITERS} LR=5e-6 MIN_LR=5e-7 -SAVE_INTERVAL=$TRAIN_ITERS +SAVE_INTERVAL=${SAVE_INTERVAL:-$TRAIN_ITERS} LOG_INTERVAL=1 TP=${TP:-2} From acd0e75bc06f4ff36ac764900322085a2a7c32c3 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Mon, 15 Jun 2026 17:02:46 -0700 Subject: [PATCH 05/10] fix(model): update Nemotron Ultra example defaults Signed-off-by: Chen Cui --- .../nemotron/nemotron_3/ultra/README.md | 15 ++++---- .../nemotron_3/ultra/slurm_inference.sh | 29 +++++++-------- .../nemotron_3/ultra/slurm_pretrain.sh | 37 ++++++++++++++----- .../nemotron/nemotron_3/ultra/slurm_sft.sh | 28 +++++++++++--- 4 files changed, 72 insertions(+), 37 deletions(-) diff --git a/examples/models/nemotron/nemotron_3/ultra/README.md b/examples/models/nemotron/nemotron_3/ultra/README.md index ff95279dcd..fc65a102e2 100644 --- a/examples/models/nemotron/nemotron_3/ultra/README.md +++ b/examples/models/nemotron/nemotron_3/ultra/README.md @@ -45,8 +45,8 @@ environment variables to the GB200 values in this table. | --- | --- | --- | | 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=2 PP=12 EP=32` script default; validate on your DCLM mount before long runs | 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=32`, selective recompute on `moe+layernorm+core_attn+moe_act` | 48 nodes, `TP=2 PP=3 EP=32`, selective recompute on `moe+layernorm+core_attn+moe_act` | +| 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. @@ -111,8 +111,8 @@ 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=2 PP=12 EP=32` on 8xH100 nodes. On 4xGB200 nodes, use 24 nodes with -`TP=2 PP=3 EP=32`. +`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 @@ -158,9 +158,10 @@ 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=32`, selective recompute on - `moe+layernorm+core_attn+moe_act`. This is the current H100 starting point - for 4096-token packed OpenMath SFT. +- 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: diff --git a/examples/models/nemotron/nemotron_3/ultra/slurm_inference.sh b/examples/models/nemotron/nemotron_3/ultra/slurm_inference.sh index c9e820757f..be5ff72756 100755 --- a/examples/models/nemotron/nemotron_3/ultra/slurm_inference.sh +++ b/examples/models/nemotron/nemotron_3/ultra/slurm_inference.sh @@ -82,28 +82,16 @@ fi mkdir -p logs -if [ -z "${COORDINATOR_HOST:-}" ]; then - if command -v scontrol >/dev/null 2>&1 && [ -n "${SLURM_NODELIST:-}" ]; then - COORDINATOR_HOST=$(scontrol show hostnames "$SLURM_NODELIST" | head -n 1) - else - COORDINATOR_HOST=$(hostname) - fi -fi -if [ -z "$COORDINATOR_HOST" ]; then - echo "ERROR: unable to resolve COORDINATOR_HOST from SLURM_NODELIST=${SLURM_NODELIST:-}." - exit 3 -fi - export HF_MODEL_PATH MEGATRON_MODEL_PATH PROMPT MAX_NEW_TOKENS KV_CACHE_BUFFER_SIZE_GB INFERENCE_MOE_TOKEN_DISPATCHER_TYPE export TP PP EP ETP GPUS_PER_NODE WORKDIR -export COORDINATOR_HOST +[ -n "${COORDINATOR_HOST:-}" ] && export COORDINATOR_HOST echo "Nemotron 3 Ultra inference" echo "Job ${SLURM_JOB_ID} nodes=${SLURM_JOB_NUM_NODES} GPUs/node=${GPUS_PER_NODE} TP=${TP} PP=${PP} EP=${EP} ETP=${ETP}" echo "HF_MODEL_PATH=${HF_MODEL_PATH}" echo "MEGATRON_MODEL_PATH=${MEGATRON_MODEL_PATH:-}" echo "KV_CACHE_BUFFER_SIZE_GB=${KV_CACHE_BUFFER_SIZE_GB}" -echo "COORDINATOR_HOST=${COORDINATOR_HOST}" +echo "COORDINATOR_HOST=${COORDINATOR_HOST:-}" SRUN_CMD=(srun --mpi=pmix --no-kill --container-image="${CONTAINER_IMAGE}" --no-container-mount-home) if [ -n "$CONTAINER_MOUNTS" ]; then @@ -124,6 +112,17 @@ if [ -n "${MEGATRON_MODEL_PATH:-}" ]; then MEGATRON_MODEL_ARGS=(--megatron_model_path "$MEGATRON_MODEL_PATH") fi +COORDINATOR_ARGS=() +if [ -z "${COORDINATOR_HOST:-}" ]; then + COORDINATOR_HOST=$(python3 - <<'"'"'PY'"'"' +import socket + +print(socket.gethostbyname(socket.gethostname())) +PY +) +fi +COORDINATOR_ARGS=(--coordinator-host "$COORDINATOR_HOST") + uv run --no-sync python scripts/inference/text_generation.py \ --hf_model_path "$HF_MODEL_PATH" \ "${MEGATRON_MODEL_ARGS[@]}" \ @@ -132,7 +131,7 @@ uv run --no-sync python scripts/inference/text_generation.py \ --kv_cache_buffer_size_gb "$KV_CACHE_BUFFER_SIZE_GB" \ --tp "$TP" --pp "$PP" --ep "$EP" --etp "$ETP" \ --use-coordinator \ - --coordinator-host "$COORDINATOR_HOST" \ + "${COORDINATOR_ARGS[@]}" \ --inference-moe-token-dispatcher-type "$INFERENCE_MOE_TOKEN_DISPATCHER_TYPE" \ --distributed-timeout-minutes 90 ' diff --git a/examples/models/nemotron/nemotron_3/ultra/slurm_pretrain.sh b/examples/models/nemotron/nemotron_3/ultra/slurm_pretrain.sh index ac01321daa..d06a3f5f62 100755 --- a/examples/models/nemotron/nemotron_3/ultra/slurm_pretrain.sh +++ b/examples/models/nemotron/nemotron_3/ultra/slurm_pretrain.sh @@ -54,18 +54,24 @@ DCLM_PATTERN=${DCLM_PATTERN:-dclm_01_*_text_document.bin} DCLM_CACHE="${WORKSPACE}/data_cache/dclm" TRAIN_ITERS=${TRAIN_ITERS:-1000} SAVE_INTERVAL=${SAVE_INTERVAL:-250} -TP=${TP:-2} +TP=${TP:-4} PP=${PP:-12} -EP=${EP:-32} +EP=${EP:-16} ETP=${ETP:-1} CP=${CP:-1} SP=${SP:-True} -RECOMPUTE_MODULES=${RECOMPUTE_MODULES:-"[moe,layernorm,core_attn,moe_act,mlp,shared_experts]"} +RECOMPUTE_GRANULARITY=${RECOMPUTE_GRANULARITY:-full} +RECOMPUTE_METHOD=${RECOMPUTE_METHOD:-uniform} +if [ -z "${RECOMPUTE_MODULES+x}" ]; then + RECOMPUTE_MODULES="" +fi +RECOMPUTE_NUM_LAYERS=${RECOMPUTE_NUM_LAYERS:-1} +RECOMPUTE_TAG=${RECOMPUTE_TAG:-recompute_full_uniform1} GPUS_PER_NODE=${GPUS_PER_NODE:-8} -SAVE_DIR="${WORKSPACE}/results/${MODEL_NAME}_dclm_pretrain_tp${TP}_pp${PP}_ep${EP}_selective_${SLURM_JOB_ID}" +SAVE_DIR="${WORKSPACE}/results/${MODEL_NAME}_dclm_pretrain_tp${TP}_pp${PP}_ep${EP}_${RECOMPUTE_TAG}_${SLURM_JOB_ID}" WANDB_ENTITY=${WANDB_ENTITY:-nvidia-nemo-fw-public} WANDB_PROJECT=${WANDB_PROJECT:-megatron-bridge-nemotron-ultra} -WANDB_EXP_NAME="${MODEL_NAME}_dclm_pretrain_tp${TP}_pp${PP}_ep${EP}_selective_${SLURM_JOB_ID}" +WANDB_EXP_NAME="${MODEL_NAME}_dclm_pretrain_tp${TP}_pp${PP}_ep${EP}_${RECOMPUTE_TAG}_${SLURM_JOB_ID}" WANDB_MODE=${WANDB_MODE:-disabled} [ -n "${HF_HOME:-}" ] && export HF_HOME @@ -78,6 +84,7 @@ export NCCL_DEBUG=${NCCL_DEBUG:-WARN} export NCCL_TIMEOUT=${NCCL_TIMEOUT:-1800000} export PYTORCH_CUDA_ALLOC_CONF=${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True} export CUDA_DEVICE_MAX_CONNECTIONS=${CUDA_DEVICE_MAX_CONNECTIONS:-1} +export PYTHONWARNINGS="${PYTHONWARNINGS:+${PYTHONWARNINGS},}ignore:The AccumulateGrad node:UserWarning,ignore:The pad token id in the tokenizer collides:UserWarning" mkdir -p logs @@ -97,7 +104,8 @@ if [ "$WANDB_MODE" = "online" ] && [ -z "${WANDB_API_KEY:-}" ]; then fi export DCLM_DATA_DIR DCLM_PATTERN DCLM_CACHE EXTRA_OVERRIDES HF_MODEL_PATH RECIPE_NAME SAVE_DIR SEQ_LENGTH -export TRAIN_ITERS SAVE_INTERVAL TP PP EP ETP CP SP RECOMPUTE_MODULES GPUS_PER_NODE +export TRAIN_ITERS SAVE_INTERVAL TP PP EP ETP CP SP RECOMPUTE_GRANULARITY RECOMPUTE_METHOD RECOMPUTE_MODULES +export RECOMPUTE_NUM_LAYERS GPUS_PER_NODE export WANDB_ENTITY WANDB_PROJECT WANDB_EXP_NAME WORKDIR WORKSPACE CMD=' @@ -120,6 +128,17 @@ if [ -z "$BLEND_PATHS" ]; then exit 4 fi +OPTIONAL_OVERRIDES=() +if [ -n "$RECOMPUTE_METHOD" ]; then + OPTIONAL_OVERRIDES+=(model.recompute_method="$RECOMPUTE_METHOD") +fi +if [ -n "$RECOMPUTE_MODULES" ]; then + OPTIONAL_OVERRIDES+=(model.recompute_modules="$RECOMPUTE_MODULES") +fi +if [ -n "$RECOMPUTE_NUM_LAYERS" ]; then + OPTIONAL_OVERRIDES+=(model.recompute_num_layers="$RECOMPUTE_NUM_LAYERS") +fi + uv run --no-sync python scripts/training/run_recipe.py \ --recipe "$RECIPE_NAME" \ --dataset llm-pretrain \ @@ -146,8 +165,8 @@ uv run --no-sync python scripts/training/run_recipe.py \ model.context_parallel_size="$CP" \ model.seq_length="$SEQ_LENGTH" \ dataset.sequence_length="$SEQ_LENGTH" \ - model.recompute_granularity=selective \ - model.recompute_modules=${RECOMPUTE_MODULES} \ + model.recompute_granularity="$RECOMPUTE_GRANULARITY" \ + "${OPTIONAL_OVERRIDES[@]}" \ dist.distributed_timeout_minutes=90 \ "dataset.blend=[[${BLEND_PATHS}],null]" \ dataset.split=\"9999,8,2\" \ @@ -170,7 +189,7 @@ echo "Recipe: ${RECIPE_NAME}" echo "HF model: ${HF_MODEL_PATH}" echo "DCLM data: ${DCLM_DATA_DIR}" echo "Parallelism: TP=${TP} PP=${PP} EP=${EP} ETP=${ETP} CP=${CP} SP=${SP}" -echo "Recompute: selective ${RECOMPUTE_MODULES}" +echo "Recompute: ${RECOMPUTE_GRANULARITY} ${RECOMPUTE_METHOD:-} ${RECOMPUTE_MODULES}" echo "Save dir: ${SAVE_DIR}" echo "W&B: ${WANDB_ENTITY}/${WANDB_PROJECT} (${WANDB_MODE})" echo "======================================" diff --git a/examples/models/nemotron/nemotron_3/ultra/slurm_sft.sh b/examples/models/nemotron/nemotron_3/ultra/slurm_sft.sh index 459b014a3c..02ca86e9f9 100755 --- a/examples/models/nemotron/nemotron_3/ultra/slurm_sft.sh +++ b/examples/models/nemotron/nemotron_3/ultra/slurm_sft.sh @@ -63,15 +63,19 @@ LOG_INTERVAL=1 TP=${TP:-2} PP=${PP:-12} -EP=${EP:-32} +EP=${EP:-16} ETP=${ETP:-1} CP=${CP:-1} SP=${SP:-True} GPUS_PER_NODE=${GPUS_PER_NODE:-8} -RECOMPUTE_GRANULARITY=selective -RECOMPUTE_MODULES="[moe,layernorm,core_attn,moe_act]" -RECOMPUTE_TAG=recompute_selective +RECOMPUTE_GRANULARITY=${RECOMPUTE_GRANULARITY:-full} +RECOMPUTE_METHOD=${RECOMPUTE_METHOD:-uniform} +if [ -z "${RECOMPUTE_MODULES+x}" ]; then + RECOMPUTE_MODULES="" +fi +RECOMPUTE_NUM_LAYERS=${RECOMPUTE_NUM_LAYERS:-1} +RECOMPUTE_TAG=${RECOMPUTE_TAG:-recompute_full_uniform1} WANDB_ENTITY=${WANDB_ENTITY:-nvidia-nemo-fw-public} WANDB_PROJECT=${WANDB_PROJECT:-megatron-bridge-nemotron-ultra} @@ -79,6 +83,7 @@ WANDB_MODE=${WANDB_MODE:-disabled} CONTAINER_IMAGE=${CONTAINER_IMAGE:-} CONTAINER_MOUNTS=${CONTAINER_MOUNTS:-} +EXTRA_OVERRIDES=${EXTRA_OVERRIDES:-} # ============================================================================== # Environment Setup @@ -94,6 +99,7 @@ export NCCL_DEBUG=${NCCL_DEBUG:-WARN} export NCCL_TIMEOUT=${NCCL_TIMEOUT:-1800000} export PYTORCH_CUDA_ALLOC_CONF=${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True} export CUDA_DEVICE_MAX_CONNECTIONS=${CUDA_DEVICE_MAX_CONNECTIONS:-1} +export PYTHONWARNINGS="${PYTHONWARNINGS:+${PYTHONWARNINGS},}ignore:The AccumulateGrad node:UserWarning" # ============================================================================== # Job Execution @@ -141,9 +147,19 @@ CLI_OVERRIDES="\ model.context_parallel_size=${CP} \ model.seq_length=${SEQ_LENGTH} \ model.recompute_granularity=${RECOMPUTE_GRANULARITY} \ - model.recompute_modules=${RECOMPUTE_MODULES} \ dist.distributed_timeout_minutes=90" +if [ -n "$RECOMPUTE_METHOD" ]; then + CLI_OVERRIDES="${CLI_OVERRIDES} model.recompute_method=${RECOMPUTE_METHOD}" +fi +if [ -n "$RECOMPUTE_MODULES" ]; then + CLI_OVERRIDES="${CLI_OVERRIDES} model.recompute_modules=${RECOMPUTE_MODULES}" +fi +if [ -n "$RECOMPUTE_NUM_LAYERS" ]; then + CLI_OVERRIDES="${CLI_OVERRIDES} model.recompute_num_layers=${RECOMPUTE_NUM_LAYERS}" +fi +CLI_OVERRIDES="${CLI_OVERRIDES} ${EXTRA_OVERRIDES}" + CMD="cd ${WORKDIR} && mkdir -p ${WORKSPACE}/results ${SAVE_DIR}/wandb ${SAVE_DIR}/tb_logs && \ export PYTHONPATH=${WORKDIR}/src:${WORKDIR}/3rdparty/Megatron-LM:\${PYTHONPATH:-} && \ uv run --no-sync python scripts/training/run_recipe.py \ @@ -163,7 +179,7 @@ echo "Nodes: ${SLURM_JOB_NUM_NODES}" 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 "Recompute: ${RECOMPUTE_GRANULARITY} ${RECOMPUTE_MODULES}" +echo "Recompute: ${RECOMPUTE_GRANULARITY} ${RECOMPUTE_METHOD:-} ${RECOMPUTE_MODULES}" echo "Save dir: ${SAVE_DIR}" echo "W&B: ${WANDB_ENTITY}/${WANDB_PROJECT} (${WANDB_MODE})" echo "======================================" From dcb42c40211976ada5e1faad40e7dc2b31ec6080 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Mon, 15 Jun 2026 18:13:31 -0700 Subject: [PATCH 06/10] docs(model): clarify Ultra checkpoint storage defaults Signed-off-by: Chen Cui --- examples/models/nemotron/nemotron_3/ultra/README.md | 8 ++++++-- .../models/nemotron/nemotron_3/ultra/slurm_pretrain.sh | 2 +- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/examples/models/nemotron/nemotron_3/ultra/README.md b/examples/models/nemotron/nemotron_3/ultra/README.md index fc65a102e2..6d11e9262f 100644 --- a/examples/models/nemotron/nemotron_3/ultra/README.md +++ b/examples/models/nemotron/nemotron_3/ultra/README.md @@ -32,7 +32,9 @@ 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. +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 @@ -120,7 +122,9 @@ 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`. +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 diff --git a/examples/models/nemotron/nemotron_3/ultra/slurm_pretrain.sh b/examples/models/nemotron/nemotron_3/ultra/slurm_pretrain.sh index d06a3f5f62..d58f27daa6 100755 --- a/examples/models/nemotron/nemotron_3/ultra/slurm_pretrain.sh +++ b/examples/models/nemotron/nemotron_3/ultra/slurm_pretrain.sh @@ -53,7 +53,7 @@ SEQ_LENGTH=4096 DCLM_PATTERN=${DCLM_PATTERN:-dclm_01_*_text_document.bin} DCLM_CACHE="${WORKSPACE}/data_cache/dclm" TRAIN_ITERS=${TRAIN_ITERS:-1000} -SAVE_INTERVAL=${SAVE_INTERVAL:-250} +SAVE_INTERVAL=${SAVE_INTERVAL:-1000} TP=${TP:-4} PP=${PP:-12} EP=${EP:-16} From 48cbe9132c956365b085f85fc096ed2a16095fb8 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Tue, 16 Jun 2026 10:20:40 -0700 Subject: [PATCH 07/10] fix(model): address Nemotron Ultra review feedback Signed-off-by: Chen Cui --- docs/models/nemotron/nemotron3-ultra.md | 4 ++-- src/megatron/bridge/recipes/nemotronh/nemotron_3_ultra.py | 2 ++ .../unit_tests/recipes/nemotronh/test_nemotron_3_ultra.py | 7 +++++++ 3 files changed, 11 insertions(+), 2 deletions(-) diff --git a/docs/models/nemotron/nemotron3-ultra.md b/docs/models/nemotron/nemotron3-ultra.md index 44df39ec4f..7545f97d03 100644 --- a/docs/models/nemotron/nemotron3-ultra.md +++ b/docs/models/nemotron/nemotron3-ultra.md @@ -4,8 +4,8 @@ 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. +Hugging Face to Megatron conversion, inference, DCLM pretraining, +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). diff --git a/src/megatron/bridge/recipes/nemotronh/nemotron_3_ultra.py b/src/megatron/bridge/recipes/nemotronh/nemotron_3_ultra.py index 00c5d31c00..9ae09aa717 100644 --- a/src/megatron/bridge/recipes/nemotronh/nemotron_3_ultra.py +++ b/src/megatron/bridge/recipes/nemotronh/nemotron_3_ultra.py @@ -102,6 +102,7 @@ def nemotron_3_ultra_pretrain_config( cfg.checkpoint.ckpt_assume_constant_structure = True cfg.checkpoint.dist_ckpt_strictness = "log_all" cfg.checkpoint.async_save = True + cfg.checkpoint.async_strategy = "nvrx" cfg.ddp.overlap_grad_reduce = True cfg.ddp.overlap_param_gather = True @@ -190,6 +191,7 @@ def nemotron_3_ultra_sft_openmathinstruct2_packed_config( cfg.checkpoint.dist_ckpt_strictness = "log_all" cfg.checkpoint.ckpt_assume_constant_structure = True cfg.checkpoint.async_save = True + cfg.checkpoint.async_strategy = "nvrx" cfg.logger.log_interval = 1 cfg.rng.seed = 5678 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 4af09ac867..8802278ae7 100644 --- a/tests/unit_tests/recipes/nemotronh/test_nemotron_3_ultra.py +++ b/tests/unit_tests/recipes/nemotronh/test_nemotron_3_ultra.py @@ -70,6 +70,9 @@ def test_pretrain_uses_initial_parallelism_values() -> None: assert cfg.optimizer.weight_decay == 0.1 assert cfg.scheduler.lr_decay_style == "constant" assert cfg.scheduler.lr_warmup_iters == 0 + assert cfg.checkpoint.async_save is True + assert cfg.checkpoint.async_strategy == "nvrx" + assert cfg.checkpoint.save_interval == 200 assert cfg.train.global_batch_size == 3072 assert cfg.train.micro_batch_size == 1 @@ -93,6 +96,8 @@ def test_openmath_sft_uses_initial_parallelism_values() -> None: assert cfg.train.train_iters == 1000 assert cfg.train.global_batch_size == 128 + assert cfg.checkpoint.async_save is True + assert cfg.checkpoint.async_strategy == "nvrx" assert cfg.dataset.dataset_name == "nvidia/OpenMathInstruct-2" assert cfg.dataset.packed_sequence_specs.packed_sequence_size == 4096 assert cfg.dataset.packed_sequence_specs.tokenizer_model_name == NEMOTRON_3_ULTRA_TOKENIZER_NAME @@ -116,6 +121,8 @@ def test_openmath_peft_uses_validated_parallelism_values() -> None: assert cfg.optimizer.min_lr == 1e-5 assert cfg.train.train_iters == 1000 assert cfg.train.global_batch_size == 128 + assert cfg.checkpoint.async_save is True + assert cfg.checkpoint.async_strategy == "nvrx" @pytest.mark.unit From 8179720f81b54307ac6d025859fdb65fe08b4e93 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Tue, 16 Jun 2026 10:32:00 -0700 Subject: [PATCH 08/10] fix(model): use validated Ultra async checkpoint strategies Signed-off-by: Chen Cui --- src/megatron/bridge/recipes/nemotronh/nemotron_3_ultra.py | 4 ++-- tests/unit_tests/recipes/nemotronh/test_nemotron_3_ultra.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/megatron/bridge/recipes/nemotronh/nemotron_3_ultra.py b/src/megatron/bridge/recipes/nemotronh/nemotron_3_ultra.py index 9ae09aa717..476c1cda8a 100644 --- a/src/megatron/bridge/recipes/nemotronh/nemotron_3_ultra.py +++ b/src/megatron/bridge/recipes/nemotronh/nemotron_3_ultra.py @@ -102,7 +102,7 @@ def nemotron_3_ultra_pretrain_config( cfg.checkpoint.ckpt_assume_constant_structure = True cfg.checkpoint.dist_ckpt_strictness = "log_all" cfg.checkpoint.async_save = True - cfg.checkpoint.async_strategy = "nvrx" + cfg.checkpoint.async_strategy = "mcore" cfg.ddp.overlap_grad_reduce = True cfg.ddp.overlap_param_gather = True @@ -191,7 +191,7 @@ def nemotron_3_ultra_sft_openmathinstruct2_packed_config( cfg.checkpoint.dist_ckpt_strictness = "log_all" cfg.checkpoint.ckpt_assume_constant_structure = True cfg.checkpoint.async_save = True - cfg.checkpoint.async_strategy = "nvrx" + cfg.checkpoint.async_strategy = "mcore" cfg.logger.log_interval = 1 cfg.rng.seed = 5678 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 8802278ae7..24db06ef70 100644 --- a/tests/unit_tests/recipes/nemotronh/test_nemotron_3_ultra.py +++ b/tests/unit_tests/recipes/nemotronh/test_nemotron_3_ultra.py @@ -71,7 +71,7 @@ def test_pretrain_uses_initial_parallelism_values() -> None: assert cfg.scheduler.lr_decay_style == "constant" assert cfg.scheduler.lr_warmup_iters == 0 assert cfg.checkpoint.async_save is True - assert cfg.checkpoint.async_strategy == "nvrx" + assert cfg.checkpoint.async_strategy == "mcore" assert cfg.checkpoint.save_interval == 200 assert cfg.train.global_batch_size == 3072 @@ -97,7 +97,7 @@ def test_openmath_sft_uses_initial_parallelism_values() -> None: assert cfg.train.train_iters == 1000 assert cfg.train.global_batch_size == 128 assert cfg.checkpoint.async_save is True - assert cfg.checkpoint.async_strategy == "nvrx" + assert cfg.checkpoint.async_strategy == "mcore" assert cfg.dataset.dataset_name == "nvidia/OpenMathInstruct-2" assert cfg.dataset.packed_sequence_specs.packed_sequence_size == 4096 assert cfg.dataset.packed_sequence_specs.tokenizer_model_name == NEMOTRON_3_ULTRA_TOKENIZER_NAME From 088b9ba9117252c5a9b6a87042c75afb986bc7ec Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Tue, 16 Jun 2026 11:35:33 -0700 Subject: [PATCH 09/10] docs(model): drop top-level Ultra README update Signed-off-by: Chen Cui --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 77a4936e8e..6b7a09d716 100644 --- a/README.md +++ b/README.md @@ -11,7 +11,7 @@ ## 📣 News -- [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. +- [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. - [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. From 821deb2cda1207f3b647fe28e5c435ca7a673ba3 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Tue, 16 Jun 2026 13:38:01 -0700 Subject: [PATCH 10/10] fix(model): enable HybridEP dispatcher for Ultra recipes Signed-off-by: Chen Cui --- src/megatron/bridge/recipes/nemotronh/nemotron_3_ultra.py | 6 +++--- tests/unit_tests/recipes/nemotronh/test_nemotron_3_ultra.py | 6 ++++++ 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/src/megatron/bridge/recipes/nemotronh/nemotron_3_ultra.py b/src/megatron/bridge/recipes/nemotronh/nemotron_3_ultra.py index 476c1cda8a..f99c8fa30f 100644 --- a/src/megatron/bridge/recipes/nemotronh/nemotron_3_ultra.py +++ b/src/megatron/bridge/recipes/nemotronh/nemotron_3_ultra.py @@ -59,7 +59,7 @@ def nemotron_3_ultra_pretrain_config( cfg.model.init_method_std = 0.014 cfg.model.use_fused_weighted_squared_relu = True cfg.model.calculate_per_token_loss = True - cfg.model.moe_token_dispatcher_type = "alltoall" + cfg.model.moe_token_dispatcher_type = "flex" cfg.model.moe_shared_expert_overlap = False cfg.model.moe_flex_dispatcher_backend = "hybridep" cfg.model.cuda_graph_impl = "none" @@ -147,7 +147,7 @@ def nemotron_3_ultra_sft_openmathinstruct2_packed_config( cfg.model.init_method_std = 0.014 cfg.model.use_fused_weighted_squared_relu = True cfg.model.calculate_per_token_loss = True - cfg.model.moe_token_dispatcher_type = "alltoall" + cfg.model.moe_token_dispatcher_type = "flex" cfg.model.moe_shared_expert_overlap = False cfg.model.moe_flex_dispatcher_backend = "hybridep" cfg.model.cuda_graph_impl = "none" @@ -241,7 +241,7 @@ def nemotron_3_ultra_peft_openmathinstruct2_packed_config( cfg.model.init_method_std = 0.014 cfg.model.use_fused_weighted_squared_relu = True cfg.model.calculate_per_token_loss = True - cfg.model.moe_token_dispatcher_type = "alltoall" + cfg.model.moe_token_dispatcher_type = "flex" cfg.model.moe_shared_expert_overlap = False cfg.model.moe_flex_dispatcher_backend = "hybridep" cfg.model.cuda_graph_impl = "none" 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 24db06ef70..3546790849 100644 --- a/tests/unit_tests/recipes/nemotronh/test_nemotron_3_ultra.py +++ b/tests/unit_tests/recipes/nemotronh/test_nemotron_3_ultra.py @@ -59,6 +59,8 @@ def test_pretrain_uses_initial_parallelism_values() -> None: assert cfg.model.tensor_model_parallel_size == 1 assert cfg.model.pipeline_model_parallel_size == 3 assert cfg.model.expert_model_parallel_size == 8 + assert cfg.model.moe_token_dispatcher_type == "flex" + assert cfg.model.moe_flex_dispatcher_backend == "hybridep" assert cfg.model.sequence_parallel is True assert cfg.model.virtual_pipeline_model_parallel_size is None assert cfg.model.mtp_num_layers == 2 @@ -87,6 +89,8 @@ def test_openmath_sft_uses_initial_parallelism_values() -> None: 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.moe_token_dispatcher_type == "flex" + assert cfg.model.moe_flex_dispatcher_backend == "hybridep" assert cfg.model.sequence_parallel is True assert cfg.model.virtual_pipeline_model_parallel_size is None assert cfg.model.recompute_granularity == "selective" @@ -110,6 +114,8 @@ def test_openmath_peft_uses_validated_parallelism_values() -> None: assert cfg.model.tensor_model_parallel_size == 2 assert cfg.model.pipeline_model_parallel_size == 4 assert cfg.model.expert_model_parallel_size == 8 + assert cfg.model.moe_token_dispatcher_type == "flex" + assert cfg.model.moe_flex_dispatcher_backend == "hybridep" assert cfg.model.sequence_parallel is True assert cfg.model.virtual_pipeline_model_parallel_size is None assert cfg.model.recompute_granularity == "selective"