From de7c519c9009b65288f6d515e2dd7f01f82c6866 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Wed, 24 Jun 2026 18:32:55 +0000 Subject: [PATCH 01/15] docs: add Nemotron 3 Ultra launch guide Signed-off-by: Anish Mahishi --- docs/guides/nemotron-3-ultra.md | 638 ++++++++++++++++++++++++++ docs/index.md | 8 + ultra_launch.sh | 778 ++++++++++++++++++++++++++++++++ 3 files changed, 1424 insertions(+) create mode 100644 docs/guides/nemotron-3-ultra.md create mode 100755 ultra_launch.sh diff --git a/docs/guides/nemotron-3-ultra.md b/docs/guides/nemotron-3-ultra.md new file mode 100644 index 0000000000..9fc49907d6 --- /dev/null +++ b/docs/guides/nemotron-3-ultra.md @@ -0,0 +1,638 @@ +# Nemotron 3 Ultra + +**Technical Report:** [NVIDIA Nemotron 3 Ultra Technical Report](https://research.nvidia.com/labs/nemotron/files/NVIDIA-Nemotron-3-Ultra-Technical-Report.pdf) + +This guide explains how to post-train the Nemotron 3 Ultra model with NeMo RL on +**GB200 NVL72** (ARM64 / aarch64) hardware. + +## Overview + +Nemotron 3 Ultra is post-trained with a multi-stage pipeline that mixes +Reinforcement Learning with Verifiable Rewards (RLVR), RLHF, and Multi-Teacher +On-Policy Distillation (MOPD) from a panel of specialised teacher models. The +main stages are: + +1. **Student RLVR** — produces the student policy from the supervised + fine-tuning (SFT) checkpoint using GRPO with verifiable rewards. +2. **Teacher training** — trains a panel of teachers (the Student RLVR policy + itself serves as the general teacher, alongside specialised teachers such as + reasoning, instruction-following/abstention, RLHF chat, and SWE). +3. **MOPD** — on-policy distillation from the student into the teacher panel. + +Every stage shares the same launcher (`ultra_launch.sh`) and a per-stage YAML +config under `examples/configs/ultra/`. + +### Checkpoint flow + +The pipeline starts from an SFT checkpoint and produces a panel of teacher +checkpoints that, together with the Student RLVR output, feed MOPD. The Student +RLVR policy itself serves as the general teacher: + +``` + ┌─────┐ + │ SFT │ + └──┬──┘ + v + ┌──────────────┐ + │ Student RLVR │ + └──┬───────────┘ + │ + │ ┌────────────────────┐ + ├──>│ General Teacher │──┐ + │ └────────────────────┘ │ + │ ┌────────────────────┐ │ + ├──>│ Reasoning Teacher │──┤ + │ └────────────────────┘ │ + │ ┌────────────────────┐ │ + ├──>│ RLHF Teacher │──┤ + │ └────────────────────┘ │ + │ ┌────────────────────┐ │ + ├──>│ IFBench Teacher │──┤ + │ └────────────────────┘ │ + │ ┌────────────────────┐ │ + ├──>│ SWE Teacher │──┤ + │ └────────────────────┘ │ + │ │ + v v + ┌─────────┐ ┌──────────┐ + │ Student │──────────────>│ MOPD │ + └─────────┘ └──────────┘ +``` + +## Container + +These are the instructions to build the training container. The image targets +**aarch64 (arm64)** for GB200 NVL72 nodes and bundles a custom vLLM build +required by Ultra. + +From the root of the repo: + +```bash +docker buildx build \ + --progress=plain \ + -f docker/Dockerfile \ + --target release \ + -t nemo-rl-ultra:arm64 \ + --build-context nemo-rl=. \ + --build-arg MAX_JOBS=8 \ + --build-arg SKIP_SGLANG_BUILD=1 \ + --build-arg BUILD_CUSTOM_VLLM=1 \ + --build-arg BUILD_CUSTOM_VLLM_URL=https://github.com/TomerBN-Nvidia/vllm.git \ + --build-arg BUILD_CUSTOM_VLLM_REF=ultra-rl-v0.17 \ + --build-arg BUILD_CUSTOM_VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.17.0/vllm-0.17.0-cp38-abi3-manylinux_2_31_aarch64.whl \ + . +``` + +Build args: +- `BUILD_CUSTOM_VLLM=1` with `BUILD_CUSTOM_VLLM_URL` / `BUILD_CUSTOM_VLLM_REF` — + builds the Ultra vLLM fork; `BUILD_CUSTOM_VLLM_PRECOMPILED_WHEEL_LOCATION` + points at the matching upstream aarch64 wheel so the build reuses precompiled + kernels instead of compiling from source. +- `SKIP_SGLANG_BUILD=1` — Ultra runs on vLLM; skip the SGLang build. +- `MAX_JOBS` — parallel build jobs; tune to your machine. +- `--build-context nemo-rl=.` — build from your local checkout (otherwise the + Dockerfile pulls `NVIDIA-NeMo/RL.git#main`). + +To run on the cluster with Slurm, convert the image to a squashfs (`.sqsh`) +with [enroot](https://github.com/NVIDIA/enroot): + +```bash +enroot import -o nemo-rl-container.sqsh dockerd://nemo-rl-ultra:arm64 +``` + +Pass the resulting image as `CONTAINER` in every launch command below (shown as +`CONTAINER=/path/to/nemo-rl-container` — a `.sqsh` path, or a registry image URI +if you're not using enroot). All Ultra stages run from this single image. + +## Download and prepare the data + +The training blends are published as +[`nvidia/Nemotron-RL-Ultra-Training-Blends`](https://huggingface.co/datasets/nvidia/Nemotron-RL-Ultra-Training-Blends), +one JSONL per stage: `rlvr1`, `rlvr2`, `ifbench`, `rlhf`, `reasoning`, `swe`, and +`mopd`. Math rows that originate from `BytedTsinghua-SIA/DAPO-Math-17k` and +`Skywork/Skywork-OR1-RL-Data` ship as placeholders; the bundled +`fill_placeholders.py` restores them from the original datasets on Hugging Face. + +```bash +export DATA_DIR=/path/to/ultra/data + +# 1. Download the blends + fill_placeholders.py +huggingface-cli download nvidia/Nemotron-RL-Ultra-Training-Blends \ + --repo-type dataset --local-dir ultra-blends + +# 2. Restore the DAPO / Skywork placeholders into $DATA_DIR +cd ultra-blends +./fill_placeholders.py --input-dir . --output-dir "$DATA_DIR" # requires uv +``` + +This produces `$DATA_DIR/{rlvr1,rlvr2,ifbench,rlhf,reasoning,swe,mopd}.jsonl`. Hold +out the last 100 rows of each blend as a validation split — the launch commands +below consume `.train.jsonl` as `TRAIN_PATH` and `.val.jsonl` as +`VAL_PATH`: + +```bash +cd "$DATA_DIR" +for name in rlvr1 rlvr2 ifbench rlhf reasoning swe mopd; do + head -n -100 "$name.jsonl" > "$name.train.jsonl" + tail -n 100 "$name.jsonl" > "$name.val.jsonl" +done +``` + +By electing to use the external datasets you are responsible for confirming their +licenses are fit for your intended use. + +The SWE stage additionally requires per-instance `.sif` container images built +from SWE-Gym and SWE-rebench-V2 — see [SWE Teacher](#swe-teacher) for the build +steps. + +For now, each stage takes a JSONL training file and a JSONL validation file +(see [Launch script](#launch-script) below). + +## Prepare the code + +```bash +git clone --recursive -b ultra-v3 https://github.com/NVIDIA-NeMo/RL.git +cd RL +``` + +## Prepare the starting checkpoint + +If your starting checkpoint comes from Hugging Face or is the result of SFT with +Megatron-Bridge, it is a **transformers v5** checkpoint. NeMo RL runs an older +**transformers v4**, so convert the checkpoint to a v4-compatible version before +using it as `MODEL_PATH`. + +The converter rewrites `config.json`, adds the v4 modeling files, and symlinks +the weight shards back to the source checkpoint. + +```bash +python examples/converters/ultra/convert_ultra_ckpt_t5_to_t4.py \ + --source /path/to/ultra_sft_checkpoint_v5 \ + --output /path/to/ultra_sft_checkpoint_v4 \ + --force +``` + +Use the converted directory (`/path/to/ultra_sft_checkpoint_v4`) as `MODEL_PATH` +in the launch commands below. The converter copies the bundled v4 NemotronH +modeling files (`configuration_nemotron_h.py` / `modeling_nemotron_h.py` in `examples/converters/ultra/`) into the output +automatically (the default `--runtime-source`); pass `--runtime-source ` to +override. + +## Build the sandbox container + +Several [Gym](https://github.com/NVIDIA-NeMo/Gym) environments used during +training (notably `ns_tools` for stateful Python execution with math +verification, and `math_formal_lean` for Lean4 proof verification) rely on a +sandbox container. Build it from the +[NeMo-Skills Dockerfile](https://github.com/NVIDIA-NeMo/Skills/blob/main/dockerfiles/Dockerfile.sandbox): + +```bash +git clone https://github.com/NVIDIA-NeMo/Skills.git +cd Skills +git checkout b620e79 # Skills commit pinned for the Ultra release +docker build -t nemo-skills-sandbox:latest -f dockerfiles/Dockerfile.sandbox . +``` + +For SLURM clusters using [enroot](https://github.com/NVIDIA/enroot), convert +to a `.sqsh`: + +```bash +enroot import -o nemo-skills-sandbox.sqsh dockerd://nemo-skills-sandbox:latest +``` + +## Launch script + +Every stage is submitted with `ultra_launch.sh` at the repo root. The launcher +handles SLURM submission, code snapshotting, persistent cache management, and +container mounts — stage-specific hyperparameters (batch size, advantage clip, +MoE parallelism, learning rate) live in the per-stage YAML. + +Set the following before each `bash ultra_launch.sh` invocation: + +| Variable | Purpose | +|---|---| +| `EXP_NAME` | Job name, W&B run name, and the suffix for output directories. Must be unique per run; same name across resubmissions resumes from the latest checkpoint. | +| `CONFIG_PATH` | Path to the per-stage YAML config (e.g. `examples/configs/ultra/student_rlvr.yaml`). | +| `MODEL_PATH` | Initial policy checkpoint (HuggingFace repo id or local path). Student RLVR starts from the Ultra SFT checkpoint; the teacher stages start from the Student RLVR checkpoint; MOPD starts from the student (the Student RLVR checkpoint). | +| `TRAIN_PATH` | Training data JSONL. | +| `VAL_PATH` | Validation data JSONL. | +| `CONTAINER` | NeMo RL container image (`.sqsh` path or registry image URI). | +| `SANDBOX_CONTAINER` | Sandbox image from [Build the sandbox container](#build-the-sandbox-container). | +| `PERSISTENT_CACHE` | Directory on a shared filesystem (e.g. Lustre) where vLLM/Triton/Inductor compile caches are persisted across runs. | +| `EXTRA_MOUNTS` | Comma-separated `host:container` mount pairs for any shared filesystems holding your data, model checkpoints, and `PERSISTENT_CACHE` (e.g. `EXTRA_MOUNTS=/lustre:/lustre,/scratch:/scratch`). | +| `SLURM_PARTITION`, `SLURM_ACCOUNT` | Your SLURM cluster credentials. | +| `GENRM_MODEL` | GenRM judge: HF repo id or local path. Default is [`nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-GenRM`](https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-GenRM). Served in-cluster from the gym pool at TP=4 DP=4 (16 GPUs). Required unless `GENRM_BASE_URL` is set. | +| `GENRM_BASE_URL` | _Optional._ URL of a separately-deployed GenRM service (e.g. `http://genrm-host:9213/v1`). If set, routes judging to that endpoint and ignores `GENRM_MODEL`. Useful when sharing a single GenRM deployment across many training jobs. | +| `NL2BASH_JUDGE_MODEL` | NL2Bash / general-purpose judge: HF repo id or local path. Default judge is `Qwen/Qwen3-235B-A22B-Instruct-2507-FP8`. | +| `SAFETY_JUDGE_MODEL` | Content-safety judge: HF repo id or local path. Default is [`nvidia/Nemotron-Content-Safety-Reasoning-4B`](https://huggingface.co/nvidia/Nemotron-Content-Safety-Reasoning-4B). | + +Optional knobs: + +| Variable | Default | Purpose | +|---|---|---| +| `WALLTIME` | `4:00:00` | SLURM `--time` | +| `SLURM_QOS`, `SLURM_RESERVATION`, `EXCLUDE_NODES` | _empty_ | Optional SLURM flags | +| `NUM_TRAIN_NODES`, `NUM_GEN_NODES`, `NUM_GYM_NODES` | `64`, `172`, `20` | GB200 4-GPU node split. Total must be a multiple of 16. | +| `ENABLE_MTP_INFERENCE` | `0` | Set to `1` to enable MTP speculative decoding for vLLM | +| `NRL_MAX_STEPS` | _from YAML_ | Override `grpo.max_num_steps` | +| `WANDB_API_KEY`, `WANDB_PROJ`, `WANDB_ENTITY` | _unset_ / `nemotron-3-ultra` / _unset_ | W&B logging is disabled if `WANDB_API_KEY` is unset | +| `HF_HOME`, `HF_TOKEN` | _unset_ | Shared HuggingFace cache and gated-model token | +| `USE_SNAPSHOT` | `1` | Snapshot the source tree at submission time | +| `DRY_RUN` | `0` | Set to `1` to print the resolved `TRAIN_CMD` without submitting | + +## Stage 1 — Student RLVR + +GRPO with verifiable rewards on the Ultra SFT checkpoint. + +Student RLVR is split into two phases that share the same `EXP_NAME` so the +second phase resumes from the first phase's checkpoint: + +| | Phase 1 | Phase 2 | +|---|---|---| +| Config | `examples/configs/ultra/student_rlvr1.yaml` | `examples/configs/ultra/student_rlvr2.yaml` | +| `max_total_sequence_length` | 49,152 | 65,536 | +| Steps in this phase | ~128 | ~50 | +| `NRL_MAX_STEPS` to set | `128` | `178` (= 128 + 50) | +| `group_answer_length_penalty_coeff` | 0.15 | 0.08 | + +Both phases share TP=8, EP=64, CP=8, PP=1, GBS=8192 (512 prompts × 16 +generations), advantage clip ±20, and the 256-node cluster shape. + +### Phase 1 — 49k context, 128 steps + +```bash +EXP_NAME=ultra-student-rlvr \ +CONFIG_PATH=examples/configs/ultra/student_rlvr1.yaml \ +ENABLE_MTP_INFERENCE=1 \ +NRL_MAX_STEPS=128 \ +MODEL_PATH=/path/to/ultra_sft_checkpoint \ +TRAIN_PATH=$DATA_DIR/rlvr1.train.jsonl \ +VAL_PATH=$DATA_DIR/rlvr1.val.jsonl \ +CONTAINER=/path/to/nemo-rl-container \ +SANDBOX_CONTAINER=/path/to/nemo-skills-sandbox.sqsh \ +PERSISTENT_CACHE=/path/to/persistent/cache \ +EXTRA_MOUNTS=/lustre:/lustre \ +SLURM_PARTITION=$SLURM_PARTITION \ +SLURM_ACCOUNT=$SLURM_ACCOUNT \ +GENRM_MODEL=nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-GenRM \ +NL2BASH_JUDGE_MODEL=Qwen/Qwen3-235B-A22B-Instruct-2507-FP8 \ +SAFETY_JUDGE_MODEL=nvidia/Nemotron-Content-Safety-Reasoning-4B \ +WANDB_API_KEY=$WANDB_API_KEY \ +HF_HOME=/path/to/hf_cache \ +HF_TOKEN=$HF_TOKEN \ +bash ultra_launch.sh +``` + +### Phase 2 — 65k context, ~50 more steps + +Same env vars as Phase 1 but swap the config, raise `NRL_MAX_STEPS`, and point +at the Phase 2 training data file. Keep `EXP_NAME` and `RESULTS_DIR` identical +to Phase 1 so `CheckpointManager` auto-resumes from the latest Phase 1 +checkpoint. + +```bash +EXP_NAME=ultra-student-rlvr \ +CONFIG_PATH=examples/configs/ultra/student_rlvr2.yaml \ +ENABLE_MTP_INFERENCE=1 \ +NRL_MAX_STEPS=178 \ +MODEL_PATH=/path/to/ultra_sft_checkpoint \ +TRAIN_PATH=$DATA_DIR/rlvr2.train.jsonl \ +VAL_PATH=$DATA_DIR/rlvr2.val.jsonl \ +CONTAINER=/path/to/nemo-rl-container \ +SANDBOX_CONTAINER=/path/to/nemo-skills-sandbox.sqsh \ +PERSISTENT_CACHE=/path/to/persistent/cache \ +EXTRA_MOUNTS=/lustre:/lustre \ +SLURM_PARTITION=$SLURM_PARTITION \ +SLURM_ACCOUNT=$SLURM_ACCOUNT \ +GENRM_MODEL=nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-GenRM \ +NL2BASH_JUDGE_MODEL=Qwen/Qwen3-235B-A22B-Instruct-2507-FP8 \ +SAFETY_JUDGE_MODEL=nvidia/Nemotron-Content-Safety-Reasoning-4B \ +WANDB_API_KEY=$WANDB_API_KEY \ +HF_HOME=/path/to/hf_cache \ +HF_TOKEN=$HF_TOKEN \ +bash ultra_launch.sh +``` + +Note: Phase 1 and Phase 2 use different training blends (`rlvr1.train.jsonl` and +`rlvr2.train.jsonl`); see the [Download and prepare the data](#download-and-prepare-the-data) +section above. + +The launcher reports the experiment directory layout, sample monitoring +commands, and (on submission) the SLURM job id. + +## Stage 2 — Teacher training + +The teacher panel is a set of specialised RL runs that each start from the +Student RLVR output (Stage 1) (see [Checkpoint flow](#checkpoint-flow) above). Each teacher +runs independently with its own YAML config under `examples/configs/ultra/`. + +The teachers don't depend on each other and can run in parallel. + +### IFBench Teacher + +RLHF teacher specializing in instruction following, abstention, and refusal +behavior. Trained at 49k context with a smaller batch (`GBS=2048`) and lower +learning rate (`lr=2.5e-6`) than the student RLVR stage. + +**Config:** `examples/configs/ultra/ifbench_teacher.yaml` +- TP=8, EP=64, CP=8, PP=1 +- `max_total_sequence_length=49152` +- `train_global_batch_size=2048`, `num_prompts_per_step=128`, `num_generations_per_prompt=16` +- Learning rate `2.5e-6` constant +- Default cluster shape: 80 nodes (32 training + 28 vLLM + 20 Gym) + +```bash +EXP_NAME=ultra-ifbench-teacher \ +CONFIG_PATH=examples/configs/ultra/ifbench_teacher.yaml \ +MODEL_PATH=/path/to/student_rlvr_output \ +TRAIN_PATH=$DATA_DIR/ifbench.train.jsonl \ +VAL_PATH=$DATA_DIR/ifbench.val.jsonl \ +NUM_TRAIN_NODES=32 \ +NUM_GEN_NODES=28 \ +NUM_GYM_NODES=20 \ +ENABLE_MTP_INFERENCE=1 \ +CONTAINER=/path/to/nemo-rl-container \ +SANDBOX_CONTAINER=/path/to/nemo-skills-sandbox.sqsh \ +PERSISTENT_CACHE=/path/to/persistent/cache \ +EXTRA_MOUNTS=/lustre:/lustre \ +SLURM_PARTITION=$SLURM_PARTITION \ +SLURM_ACCOUNT=$SLURM_ACCOUNT \ +GENRM_MODEL=nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-GenRM \ +NL2BASH_JUDGE_MODEL=Qwen/Qwen3-235B-A22B-Instruct-2507-FP8 \ +SAFETY_JUDGE_MODEL=nvidia/Nemotron-Content-Safety-Reasoning-4B \ +WANDB_API_KEY=$WANDB_API_KEY \ +HF_HOME=/path/to/hf_cache \ +HF_TOKEN=$HF_TOKEN \ +bash ultra_launch.sh +``` + +### RLHF Teacher + +General-purpose RLHF teacher trained against the pairwise GenRM comparison +signal alone. Same training shape as the IFBench teacher (cluster, batch, +learning rate, context). + +**Config:** `examples/configs/ultra/rlhf_teacher.yaml` +- TP=8, EP=64, CP=8, PP=1 +- `max_total_sequence_length=49152` +- `train_global_batch_size=2048`, `num_prompts_per_step=128`, `num_generations_per_prompt=16` +- Learning rate `2.5e-6` constant +- Default cluster shape: 64 nodes (32 training + 28 vLLM + 4 Gym) + +```bash +EXP_NAME=ultra-rlhf-teacher \ +CONFIG_PATH=examples/configs/ultra/rlhf_teacher.yaml \ +MODEL_PATH=/path/to/student_rlvr_output \ +TRAIN_PATH=$DATA_DIR/rlhf.train.jsonl \ +VAL_PATH=$DATA_DIR/rlhf.val.jsonl \ +NUM_TRAIN_NODES=32 \ +NUM_GEN_NODES=28 \ +NUM_GYM_NODES=4 \ +ENABLE_MTP_INFERENCE=1 \ +CONTAINER=/path/to/nemo-rl-container \ +SANDBOX_CONTAINER=/path/to/nemo-skills-sandbox.sqsh \ +PERSISTENT_CACHE=/path/to/persistent/cache \ +EXTRA_MOUNTS=/lustre:/lustre \ +SLURM_PARTITION=$SLURM_PARTITION \ +SLURM_ACCOUNT=$SLURM_ACCOUNT \ +GENRM_MODEL=nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-GenRM \ +WANDB_API_KEY=$WANDB_API_KEY \ +HF_HOME=/path/to/hf_cache \ +HF_TOKEN=$HF_TOKEN \ +bash ultra_launch.sh +``` + +This is the teacher referred to as `NRL_CHAT_TEACHER1` / `NRL_RLHF_TEACHER` in +the MOPD config — it provides the `genrm_simple_agent` and +`genrm_simple_agent_reasoning_off` teacher signals. + +### Reasoning Teacher + +General reasoning teacher. The training data subsamples an RLVR +blend, where every prompt is graded by the `equivalence_llm_judge` agent +(LLM-judge equivalence over freeform short answers). The output checkpoint +serves the `code_gen`, `ns_tools`, `math_with_judge`, +`equivalence_llm_judge`, and `mcqa` agent slots in MOPD — one checkpoint, +many roles. + +**Config:** `examples/configs/ultra/reasoning_teacher.yaml` +- TP=8, EP=32, CP=8, PP=1 (half the expert parallelism of Student RLVR) +- `max_total_sequence_length=65536` +- `train_global_batch_size=2048`, `num_prompts_per_step=128`, `num_generations_per_prompt=16` +- Learning rate `3.0e-6` constant +- `max_num_epochs=10` — small sub-sampled dataset, multiple passes expected +- Default cluster shape: 128 nodes (64 training + 60 vLLM + 4 Gym) + +```bash +EXP_NAME=ultra-reasoning-teacher \ +CONFIG_PATH=examples/configs/ultra/reasoning_teacher.yaml \ +ENABLE_MTP_INFERENCE=1 \ +MODEL_PATH=/path/to/student_rlvr_output \ +TRAIN_PATH=$DATA_DIR/reasoning.train.jsonl \ +VAL_PATH=$DATA_DIR/reasoning.val.jsonl \ +NUM_TRAIN_NODES=64 \ +NUM_GEN_NODES=60 \ +NUM_GYM_NODES=4 \ +CONTAINER=/path/to/nemo-rl-container \ +SANDBOX_CONTAINER=/path/to/nemo-skills-sandbox.sqsh \ +PERSISTENT_CACHE=/path/to/persistent/cache \ +EXTRA_MOUNTS=/lustre:/lustre \ +SLURM_PARTITION=$SLURM_PARTITION \ +SLURM_ACCOUNT=$SLURM_ACCOUNT \ +NL2BASH_JUDGE_MODEL=Qwen/Qwen3-235B-A22B-Instruct-2507-FP8 \ +WANDB_API_KEY=$WANDB_API_KEY \ +HF_HOME=/path/to/hf_cache \ +HF_TOKEN=$HF_TOKEN \ +bash ultra_launch.sh +``` + +### SWE Teacher + +Software-engineering RLVR teacher trained against code execution. The +`swe_agents` agent runs the policy's candidate fixes inside apptainer (`.sif`) +container images for each SWE-Gym / SWE-rebench-V2 instance and rewards +the rollout based on test pass/fail. + +**Config:** `examples/configs/ultra/swe_teacher.yaml` +- TP=8, EP=32, CP=32, PP=1 (large CP for long context) +- `max_total_sequence_length=196608` (192k context) +- `train_global_batch_size=512`, `num_prompts_per_step=32`, `num_generations_per_prompt=16` +- Learning rate `3.0e-6` constant, advantage clip ±100, `max_num_epochs=4` +- Default cluster shape: 128 nodes (64 training + 64 vLLM + 0 Gym) + +### Building the SIF images + +Each rollout runs inside a per-instance apptainer (`.sif`) image resolved from +`${SIF_DIR}` via the agent's `container_formatter`. Since this recipe targets +GB200 (aarch64) and the upstream SWE images ship for x86 only, the images must +be rebuilt for ARM — one per instance — then converted to `.sif`. The resulting +directory is reused across runs. + +The released SWE blend draws from two benchmarks: + +| Benchmark | HF dataset | `.sif` path under `${SIF_DIR}` | Instances | +|---|---|---|---| +| SWE-Gym | `SWE-Gym/SWE-Gym` | `swegym/sweb.eval.arm64.{instance_id}.sif` | 206 | +| SWE-rebench-V2 | `nebius/SWE-rebench-V2` | `swerebench/{instance_id}.sif` | 7,610 | + +**1. Prerequisites.** Docker, Apptainer, and `uv` on the build host. Build on an +ARM64 (GB200) node so images are natively aarch64 with no emulation. You also +need a container **registry** to publish to: set `REGISTRY` to its endpoint and +`docker login` to it first. The build scripts push every per-instance image +there, and the conversion step (3) pulls them back to produce the `.sif` files — +so the registry must be reachable from both the build and the convert hosts. + +Install Apptainer on the build host via the official PPA, pinned to the version +the training container ships (so the `.sif` format matches — see +`docker/install_apptainer.sh`): + +```bash +sudo add-apt-repository -y ppa:apptainer/ppa +sudo apt-get update +sudo apt-get install -y "apptainer=1.5.0-2-1~$(. /etc/os-release && echo "$VERSION_CODENAME")" +``` + +> **Storage.** The registry accumulates ~7,800 images, plus the converted `.sif` +> set on the build host. The pushed images can be deleted once every `.sif` has +> been built. + +**2. Build the per-instance images.** +```bash +git clone -b ultra-v3 https://github.com/nujoug/swe-gym-arm-build +git clone -b ultra-v3 https://github.com/nujoug/swe-rebench-v2-arm-build +export REGISTRY=registry.example.com/ultra-swe # your registry endpoint (docker login first) + +# --- SWE-Gym (206 images) --- +# No verify gate: this wrapper builds + pushes only. After it finishes, drop any +# instance listed under handoff/failed_instances/ before using the images. +cd swe-gym-arm-build +uv venv && source .venv/bin/activate && uv pip install -e . +python scripts/batch_build_push.py \ + --dataset SWE-Gym/SWE-Gym --split train \ + --instance_ids_file swe_gym_instance_ids.txt \ + --registry "${REGISTRY}/swe-gym" --push_env_images \ + --max_workers 8 --state_file build_push_state.json +deactivate; cd .. + +# --- SWE-rebench-V2 (7,610 images) --- +# Omitting --skip-eval enables the verify gate: each image is built, the gold +# patch is applied, the tests are run, and only images whose FAIL_TO_PASS / +# PASS_TO_PASS transition matches are published (others land in +# handoff/failed_instances/). +cd swe-rebench-v2-arm-build +uv venv && source .venv/bin/activate && uv pip install -r requirements.txt +# Build the per-language base images first (Go/Java/Rust/Python/… environments +# that the instance images layer on top of). +# Note: expect one known failure when building scala_base (its Dockerfile runs `foundryup`, which 403s on the GitHub API). +python3 scripts/build_all_arm_bases.py --platform linux/arm64 --keep-going --max-workers 4 --skip-existing +python3 scripts/prepare_ready_tasks.py --hf-dataset nebius/SWE-rebench-V2 --output ready_tasks.json +python3 scripts/build_eval_cleanup.py \ + --json ready_tasks.json --platform linux/arm64 --max-workers 8 \ + --report-json eval_report.json --skip-done \ + --gitlab-registry "${REGISTRY}/swerebenchv2" +deactivate; cd .. +``` + +**3. Convert to `.sif` and lay out `${SIF_DIR}`.** `build_swe_sif_images.py` pulls +each published image and runs `apptainer build` under the exact filename the +recipe expects — SWE-Gym from the `swe-gym:sweb.eval.arm64.` tags, and +SWE-rebench-V2 from the verified (`passed_match`) instances in `eval_report.json`. +It skips images already converted and continues past any that are missing (e.g. +instances that failed to build), recording them in `${SIF_DIR}/missing_instances.txt`. +Run it from the NeMo RL repo root: +```bash +export SIF_DIR=/path/to/sif/images +python examples/nemo_gym/build_swe_sif_images.py \ + --registry "${REGISTRY}" --sif-dir "${SIF_DIR}" \ + --swe-gym-ids /path/to/swe-gym-arm-build/swe_gym_instance_ids.txt \ + --rebench-report /path/to/swe-rebench-v2-arm-build/eval_report.json +``` + +With `${SIF_DIR}` populated, launch the SWE teacher: + +```bash +EXP_NAME=ultra-swe-teacher \ +CONFIG_PATH=examples/configs/ultra/swe_teacher.yaml \ +ENABLE_MTP_INFERENCE=1 \ +MODEL_PATH=/path/to/student_rlvr_output \ +TRAIN_PATH=$DATA_DIR/swe.train.jsonl \ +VAL_PATH=$DATA_DIR/swe.val.jsonl \ +NUM_TRAIN_NODES=64 \ +NUM_GEN_NODES=64 \ +NUM_GYM_NODES=0 \ +SIF_DIR=/path/to/sif/images \ +CONTAINER=/path/to/nemo-rl-container \ +SANDBOX_CONTAINER=/path/to/nemo-skills-sandbox.sqsh \ +PERSISTENT_CACHE=/path/to/persistent/cache \ +EXTRA_MOUNTS=/lustre:/lustre \ +SLURM_PARTITION=$SLURM_PARTITION \ +SLURM_ACCOUNT=$SLURM_ACCOUNT \ +WANDB_API_KEY=$WANDB_API_KEY \ +HF_HOME=/path/to/hf_cache \ +HF_TOKEN=$HF_TOKEN \ +bash ultra_launch.sh +``` + +## Stage 3 — MOPD + +Multi-Teacher On-Policy Distillation. The Student RLVR output is the student; each +Gym agent is routed to one of the Stage 2 teacher checkpoints. Trains the +student to match per-agent teacher distributions. + +**Config:** `examples/configs/ultra/mopd.yaml` +- TP=8, EP=64, CP=8, PP=1, max context 192k +- Teacher parallelism: TP=8, CP=2, EP=16, 4 nodes per teacher +- Routing: agent → teacher checkpoint baked into the YAML via + `${_teachers.}` references; only `_teachers.general` is required and + every other slot falls back to it. +- Default cluster shape: 224 nodes (64 training + 128 vLLM + 12 Gym + 20 teachers). + +### Teacher mapping + +| Logical slot | Path source | MOPD agents it serves | +|---|---|---| +| `general` | Student RLVR output | `lc_judge_simple_agent`, fallback for unset slots below | +| `rlhf` | RLHF Teacher | `genrm_simple_agent`, `genrm_simple_agent_reasoning_off` | +| `ifbench` | IFBench Teacher | `instruction_following_simple_agent`, `abstention_simple_agent`, `multichallenge_simple_agent` | +| `reasoning` | Reasoning Teacher | `math_with_judge_simple_agent`, `equivalence_llm_judge_simple_agent`, `mcqa_simple_agent`, `ns_tools_simple_agent`, `code_gen_simple_agent` | +| `swe` | SWE Teacher | all `swe_pivot_*`, `terminal_multi_harness_{opencode,agent006,codex}`, `droid_pivot_*`, `structured_outputs_v3_simple_agent`, `freeform_formatting_simple_agent`, `citation_format_simple_agent` | + +### Launch + +Pass `STAGE_TYPE=mopd` to the launcher to enable the teacher-pool node math +and the `_teachers.X` Hydra overrides. `NRL_GENERAL_TEACHER_PATH` is required; +the other four teacher paths are optional and fall back to general when +unset. + +```bash +STAGE_TYPE=mopd \ +EXP_NAME=ultra-mopd-stage1 \ +CONFIG_PATH=examples/configs/ultra/mopd.yaml \ +ENABLE_MTP_INFERENCE=1 \ +MODEL_PATH=/path/to/student_rlvr_output \ +TRAIN_PATH=$DATA_DIR/mopd.train.jsonl \ +VAL_PATH=$DATA_DIR/mopd.val.jsonl \ +NUM_TRAIN_NODES=64 \ +NUM_GEN_NODES=128 \ +NUM_GYM_NODES=12 \ +NUM_UNIQUE_TEACHERS=5 \ +NUM_NODES_PER_TEACHER=4 \ +NRL_GENERAL_TEACHER_PATH=/path/to/student_rlvr_output \ +NRL_RLHF_TEACHER_PATH=/path/to/rlhf_teacher_output \ +NRL_IFBENCH_TEACHER_PATH=/path/to/ifbench_teacher_output \ +NRL_REASONING_TEACHER_PATH=/path/to/reasoning_teacher_output \ +NRL_SWE_TEACHER_PATH=/path/to/swe_teacher_output \ +CONTAINER=/path/to/nemo-rl-container \ +SANDBOX_CONTAINER=/path/to/nemo-skills-sandbox.sqsh \ +PERSISTENT_CACHE=/path/to/persistent/cache \ +EXTRA_MOUNTS=/lustre:/lustre \ +SLURM_PARTITION=$SLURM_PARTITION \ +SLURM_ACCOUNT=$SLURM_ACCOUNT \ +GENRM_MODEL=nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-GenRM \ +NL2BASH_JUDGE_MODEL=Qwen/Qwen3-235B-A22B-Instruct-2507-FP8 \ +SAFETY_JUDGE_MODEL=nvidia/Nemotron-Content-Safety-Reasoning-4B \ +WANDB_API_KEY=$WANDB_API_KEY \ +HF_HOME=/path/to/hf_cache \ +HF_TOKEN=$HF_TOKEN \ +bash ultra_launch.sh +``` diff --git a/docs/index.md b/docs/index.md index 58c2f13fb7..951f663171 100644 --- a/docs/index.md +++ b/docs/index.md @@ -100,6 +100,13 @@ Reproduce DeepscaleR results with NeMo RL using GRPO on mathematical reasoning t Step-by-step guide for supervised fine-tuning on the OpenMathInstruct2 dataset. ::: +:::{grid-item-card} {octicon}`rocket` Nemotron 3 Ultra +:link: guides/nemotron-3-ultra +:link-type: doc + +Post-train Nemotron 3 Ultra with RLVR, teacher training, and MOPD stages on GB200 NVL72 hardware. +::: + :::{grid-item-card} {octicon}`stack` Environments :link: guides/environments :link-type: doc @@ -278,6 +285,7 @@ guides/sft-openmathinstruct2.md guides/nemotron-3-nano.md guides/nemotron-3-nano-omni.md guides/nemotron-3-super.md +guides/nemotron-3-ultra.md adding-new-models.md guides/sft.md guides/dpo.md diff --git a/ultra_launch.sh b/ultra_launch.sh new file mode 100755 index 0000000000..c51458c2dc --- /dev/null +++ b/ultra_launch.sh @@ -0,0 +1,778 @@ +#!/bin/bash +set -euo pipefail + +# ============================================================================= +# ultra_launch.sh +# +# Public launcher for Nemotron 3 Ultra post-training stages on a SLURM cluster. +# +# Each training stage (Student RLVR, teacher RLVR/RLHF stages, MOPD) has a +# matching YAML config under examples/configs/ultra/. The stage-specific +# hyperparameters (batch size, advantage clip, MoE parallelism, etc.) live +# in the YAML; this launcher only handles orchestration: SLURM submission, +# code snapshotting, persistent cache management, container mounts, and the +# Hydra overrides that vary per run (data paths, model checkpoint, judge +# endpoints, log directories). +# +# Usage: +# +# EXP_NAME=ultra-student-rlvr-001 \ +# CONFIG_PATH=examples/configs/ultra/student_rlvr.yaml \ +# MODEL_PATH=/path/to/sft_checkpoint \ +# TRAIN_PATH=/path/to/train.jsonl \ +# VAL_PATH=/path/to/val.jsonl \ +# CONTAINER=nvcr.io/nvidia/nemo-rl: \ +# SANDBOX_CONTAINER=/path/to/nemo-skills-sandbox.sqsh \ +# PERSISTENT_CACHE=/path/to/persistent/cache \ +# SLURM_PARTITION=batch \ +# SLURM_ACCOUNT=your_account \ +# GENRM_MODEL= # Or set GENRM_BASE_URL to use a remote service +# NL2BASH_JUDGE_MODEL=Qwen/Qwen3-235B-A22B-Instruct-2507-FP8 \ +# SAFETY_JUDGE_MODEL=/path/to/safety_checkpoint \ +# bash ultra_launch.sh +# +# Optional knobs: +# WALLTIME=4:00:00 Slurm --time +# SLURM_QOS=normal Slurm --qos +# SLURM_RESERVATION= Slurm --reservation +# SLURM_DEPENDENCY= Extra Slurm dependency, merged with +# singleton (e.g. afterany:) +# EXCLUDE_NODES= Slurm --exclude +# NUM_TRAIN_NODES=64 Training (Megatron) nodes +# NUM_GEN_NODES=172 vLLM generation nodes +# NUM_GYM_NODES=20 NeMo Gym (judge) nodes +# ENABLE_MTP_INFERENCE=0 1 to enable MTP speculative decoding +# NUM_SPECULATIVE_TOKENS=5 MTP speculative tokens +# MAX_NUM_BATCHED_TOKENS=8480 vLLM max batched tokens (MTP) +# NRL_MAX_STEPS= Override grpo.max_num_steps +# EXTRA_MOUNTS= Comma-separated host:container pairs +# USE_SNAPSHOT=1 Snapshot source tree at submission +# DRY_RUN=0 1 to print TRAIN_CMD and exit +# INTERACTIVE=0 1 to bring up Ray and idle for attach +# (no training driver) for debugging +# INTERACTIVE_WAIT=1 0 to submit and return immediately +# INTERACTIVE_WALLTIME= override WALLTIME for the interactive alloc +# HF_HOME= HuggingFace cache root (recommended) +# HF_TOKEN= HuggingFace API token +# WANDB_API_KEY= Weights & Biases API key +# WANDB_PROJ=nemotron-3-ultra W&B project +# WANDB_ENTITY= W&B entity +# +# Hydra overrides are forwarded verbatim as positional arguments: +# bash ultra_launch.sh policy.megatron_cfg.optimizer.lr=1e-6 grpo.val_period=50 +# +# GB200 NVL72 nodes have 4 GPUs each. SLURM total = NUM_TRAIN + NUM_GEN + NUM_GYM +# and must be a multiple of SEGMENT_SIZE (default 16, one NVLink domain group). +# ============================================================================= + +# ============================================================================= +# Required environment +# ============================================================================= +: "${EXP_NAME:?EXP_NAME is required (used for job name, W&B run, checkpoint/log dirs)}" +: "${CONFIG_PATH:?CONFIG_PATH is required (e.g. examples/configs/ultra/student_rlvr.yaml)}" +: "${MODEL_PATH:?MODEL_PATH is required (initial policy checkpoint, HF repo id or local path)}" +: "${TRAIN_PATH:?TRAIN_PATH is required (training data jsonl path)}" +: "${VAL_PATH:?VAL_PATH is required (validation data jsonl path)}" +: "${CONTAINER:?CONTAINER is required (NGC image URI or .sqsh path)}" +: "${SANDBOX_CONTAINER:?SANDBOX_CONTAINER is required (nemo-skills sandbox image)}" +: "${PERSISTENT_CACHE:?PERSISTENT_CACHE is required (Lustre dir for vLLM/Triton/Inductor caches)}" +: "${SLURM_PARTITION:?SLURM_PARTITION is required}" +: "${SLURM_ACCOUNT:?SLURM_ACCOUNT is required}" +# Judge models are recipe-specific. Most teachers (student RLVR, IFBench, RLHF, +# Reasoning) need all three (GenRM, NL2Bash, Safety). The SWE teacher uses +# code-execution rewards and needs none of them. Set per recipe; unset vars +# skip the corresponding override. +NL2BASH_JUDGE_MODEL="${NL2BASH_JUDGE_MODEL:-}" +SAFETY_JUDGE_MODEL="${SAFETY_JUDGE_MODEL:-}" +GENRM_BASE_URL="${GENRM_BASE_URL:-}" +GENRM_MODEL="${GENRM_MODEL:-}" +GENRM_OVERRIDE="" +if [[ -n "${GENRM_BASE_URL}" ]]; then + GENRM_OVERRIDE="env.nemo_gym.genrm_model.responses_api_models.genrm_model.base_url=${GENRM_BASE_URL}" +elif [[ -n "${GENRM_MODEL}" ]]; then + GENRM_OVERRIDE="env.nemo_gym.genrm_model.responses_api_models.genrm_model.model=${GENRM_MODEL}" +fi + +# SIF_DIR: for the SWE teacher recipe — directory containing apptainer .sif +# images for SWE-Bench / SWE-Gym / R2E-Gym instances. The yaml's +# container_formatter uses `${sif_dir}/...` paths. Unset for non-SWE recipes. +SIF_DIR="${SIF_DIR:-}" + +if [[ ! -f "${CONFIG_PATH}" ]]; then + echo "ERROR: CONFIG_PATH does not exist: ${CONFIG_PATH}" >&2 + exit 1 +fi + +# ============================================================================= +# Project root and code root +# ============================================================================= +PROJECT_ROOT=$(realpath "$PWD") +cd "${PROJECT_ROOT}" + +# ============================================================================= +# Job identity — fixed name for singleton. +# Slurm --dependency=singleton serialises queued submissions with the same name +# so a resubmission after preemption resumes from the latest checkpoint instead +# of running in parallel. +# ============================================================================= +JOB_NAME="${EXP_NAME}" + +# ============================================================================= +# Output directories +# ============================================================================= +RESULTS_DIR="${RESULTS_DIR:-results/${EXP_NAME}}" +CHECKPOINT_DIR="${CHECKPOINT_DIR:-${RESULTS_DIR}/checkpoints}" + +# Per-submission dirs for logs and Slurm output (timestamped for history). +RUN_DIR="${RESULTS_DIR}/runs/$(date +%Y%m%d-%H%M)" +LOG_DIR="${RUN_DIR}/logs" +SLURM_LOG_DIR="${RUN_DIR}/slurm" +mkdir -p "${CHECKPOINT_DIR}" "${LOG_DIR}" "${SLURM_LOG_DIR}" +ln -sfn "$(realpath "${RUN_DIR}")" "${RESULTS_DIR}/runs/latest" + +# ray.sub reads BASE_LOG_DIR and creates $BASE_LOG_DIR/$SLURM_JOB_ID-logs/ for +# ray infrastructure logs (ray-head.log, ray-driver.log, ray-worker-*.log, +# topology probes, attach scripts, etc.). +export BASE_LOG_DIR="${BASE_LOG_DIR:-${RESULTS_DIR}/ray_logs}" + +# ============================================================================= +# SLURM configuration +# ============================================================================= +WALLTIME="${WALLTIME:-4:00:00}" +SLURM_QOS="${SLURM_QOS:-}" +SLURM_RESERVATION="${SLURM_RESERVATION:-}" +EXCLUDE_NODES="${EXCLUDE_NODES:-}" +# INTERACTIVE=1 brings up the Ray cluster and idles for attachment (no training +# driver), so you can run/debug the recipe by hand. INTERACTIVE_WAIT=1 (default) +# blocks until Ray is ready; INTERACTIVE_WALLTIME overrides WALLTIME for the alloc. +INTERACTIVE="${INTERACTIVE:-0}" +INTERACTIVE_WAIT="${INTERACTIVE_WAIT:-1}" +# If set (format DD:HH:MM:SS), training stops early to reserve time for a final +# checkpoint save before walltime. Unset to use the YAML's default and let +# slurm walltime end the job naturally — fine when each step checkpoints. +CHECKPOINTING_SAVE_BY="${CHECKPOINTING_SAVE_BY:-}" + +# ============================================================================= +# Container & mounts +# ============================================================================= +export CONTAINER +MOUNTS="${MOUNTS:-}" + +# GB200 NVL72: fixed at 4 GPUs/node. +export GPUS_PER_NODE=4 +export CPUS_PER_WORKER="${CPUS_PER_WORKER:-144}" + +# ============================================================================= +# HuggingFace configuration +# ============================================================================= +if [[ -n "${HF_HOME:-}" ]]; then + export HF_HOME + export HF_HUB_CACHE="${HF_HUB_CACHE:-${HF_HOME}/hub}" + export HF_DATASETS_CACHE="${HF_DATASETS_CACHE:-${HF_HOME}/hub}" +else + echo "[WARN] HF_HOME is not set — HuggingFace will use the default cache (~/.cache/huggingface) per-node." >&2 +fi + +# ============================================================================= +# W&B configuration +# ============================================================================= +WANDB_PROJ="${WANDB_PROJ:-nemotron-3-ultra}" +WANDB_NAME="${EXP_NAME}" +WANDB_ENABLED=False +if [[ -n "${WANDB_API_KEY:-}" ]]; then + export WANDB_API_KEY + WANDB_ENABLED=True + if [[ -n "${WANDB_ENTITY:-}" ]]; then + export WANDB_ENTITY + fi +else + echo "[WARN] WANDB_API_KEY is not set — W&B logging will be disabled." >&2 +fi + +# ============================================================================= +# Training overrides +# ============================================================================= +NRL_MAX_STEPS="${NRL_MAX_STEPS:-}" + +# ============================================================================= +# MTP speculative decoding (optional) +# ============================================================================= +ENABLE_MTP_INFERENCE="${ENABLE_MTP_INFERENCE:-0}" +NUM_SPECULATIVE_TOKENS="${NUM_SPECULATIVE_TOKENS:-5}" +MAX_NUM_BATCHED_TOKENS="${MAX_NUM_BATCHED_TOKENS:-8480}" +MTP_EXTRA_ARGS="" +if [[ "${ENABLE_MTP_INFERENCE}" == "1" ]]; then + MTP_EXTRA_ARGS="\ +++policy.generation.vllm_cfg.enable_prefix_caching=true \ +++policy.generation.vllm_kwargs.enable_chunked_prefill=true \ +++policy.generation.vllm_kwargs.max_num_batched_tokens=${MAX_NUM_BATCHED_TOKENS} \ +++policy.generation.vllm_kwargs.mamba_cache_mode=align \ +~policy.generation.vllm_kwargs.compilation_config.cudagraph_capture_sizes \ +++policy.generation.vllm_kwargs.speculative_config.num_speculative_tokens=${NUM_SPECULATIVE_TOKENS} \ +++policy.generation.vllm_kwargs.speculative_config.method=mtp" + echo "MTP speculative decoding ENABLED (num_speculative_tokens=${NUM_SPECULATIVE_TOKENS})" +fi + +# ============================================================================= +# Job shape — defaults match the 256-node student_rlvr.yaml +# +# Training: 64 nodes ( 256 GPUs) — Megatron training backend +# vLLM: 172 nodes ( 688 GPUs) — async generation, EP=8 instances at TP=8 +# Gym: 20 nodes ( 80 GPUs) — judges (GenRM, NL2Bash, Safety) +# +# Override via NUM_TRAIN_NODES / NUM_GEN_NODES / NUM_GYM_NODES. +# +# For STAGE_TYPE=mopd, additional teacher nodes are allocated for the +# non-colocated teacher panel: NUM_UNIQUE_TEACHERS × NUM_NODES_PER_TEACHER. +# ============================================================================= +NUM_TRAIN_NODES="${NUM_TRAIN_NODES:-64}" +NUM_GEN_NODES="${NUM_GEN_NODES:-172}" +NUM_GYM_NODES="${NUM_GYM_NODES:-20}" + +STAGE_TYPE="${STAGE_TYPE:-grpo}" +NUM_TEACHER_NODES=0 +MOPD_OVERRIDES="" +if [[ "${STAGE_TYPE}" == "mopd" ]]; then + : "${NRL_GENERAL_TEACHER_PATH:?NRL_GENERAL_TEACHER_PATH is required for STAGE_TYPE=mopd (path to the Student RLVR output checkpoint)}" + NUM_UNIQUE_TEACHERS="${NUM_UNIQUE_TEACHERS:-5}" + NUM_NODES_PER_TEACHER="${NUM_NODES_PER_TEACHER:-4}" + NUM_TEACHER_NODES=$((NUM_UNIQUE_TEACHERS * NUM_NODES_PER_TEACHER)) + TEACHER_TP="${TEACHER_TP:-8}" + TEACHER_CP="${TEACHER_CP:-2}" + TEACHER_PP="${TEACHER_PP:-1}" + TEACHER_EP="${TEACHER_EP:-16}" + + # _teachers.general is required; other slots fall back via the YAML's + # interpolation. Pass only the slots the user explicitly set. + MOPD_OVERRIDES="_teachers.general=${NRL_GENERAL_TEACHER_PATH}" + for _slot in RLHF IFBENCH REASONING SWE; do + _var="NRL_${_slot}_TEACHER_PATH" + _val="${!_var:-}" + if [[ -n "${_val}" ]]; then + MOPD_OVERRIDES="${MOPD_OVERRIDES} _teachers.$(echo ${_slot} | tr A-Z a-z)=${_val}" + fi + done + + # Teacher parallelism + per-teacher node count + MOPD_OVERRIDES="${MOPD_OVERRIDES} \ +on_policy_distillation.non_colocated_teachers.default_teacher_cfg.tensor_model_parallel_size=${TEACHER_TP} \ +on_policy_distillation.non_colocated_teachers.default_teacher_cfg.context_parallel_size=${TEACHER_CP} \ +on_policy_distillation.non_colocated_teachers.default_teacher_cfg.pipeline_model_parallel_size=${TEACHER_PP} \ +on_policy_distillation.non_colocated_teachers.default_teacher_cfg.expert_model_parallel_size=${TEACHER_EP} \ +on_policy_distillation.non_colocated_teachers.default_teacher_cfg.num_nodes=${NUM_NODES_PER_TEACHER}" + + echo "MOPD: ${NUM_UNIQUE_TEACHERS} teacher pools × ${NUM_NODES_PER_TEACHER} nodes = ${NUM_TEACHER_NODES} teacher nodes" +fi + +NUM_ACTOR_NODES=$((NUM_TRAIN_NODES + NUM_GEN_NODES + NUM_TEACHER_NODES)) +NUM_TOTAL_NODES=$((NUM_ACTOR_NODES + NUM_GYM_NODES)) + +if (( NUM_TRAIN_NODES <= 0 )); then + echo "ERROR: NUM_TRAIN_NODES must be > 0 (got ${NUM_TRAIN_NODES})" >&2; exit 1 +fi +if (( NUM_GEN_NODES <= 0 )); then + echo "ERROR: NUM_GEN_NODES must be > 0 (got ${NUM_GEN_NODES})" >&2; exit 1 +fi +if (( NUM_GYM_NODES < 0 )); then + echo "ERROR: NUM_GYM_NODES must be >= 0 (got ${NUM_GYM_NODES})" >&2; exit 1 +fi + +# GB200 NVL72 topology: 18 nodes per NVLink domain, allocate in groups of 16. +SEGMENT_SIZE="${SEGMENT_SIZE:-16}" +if (( NUM_TOTAL_NODES < SEGMENT_SIZE )); then + echo "ERROR: NUM_TOTAL_NODES=${NUM_TOTAL_NODES} < SEGMENT_SIZE=${SEGMENT_SIZE}" >&2 + exit 1 +fi +if (( NUM_TOTAL_NODES % SEGMENT_SIZE != 0 )); then + echo "ERROR: NUM_TOTAL_NODES=${NUM_TOTAL_NODES} is not divisible by SEGMENT_SIZE=${SEGMENT_SIZE}." >&2 + echo " Training=${NUM_TRAIN_NODES} + Generation=${NUM_GEN_NODES} + Gym=${NUM_GYM_NODES} + Teachers=${NUM_TEACHER_NODES} = ${NUM_TOTAL_NODES}" >&2 + echo " Adjust node counts so the total is a multiple of ${SEGMENT_SIZE}." >&2 + exit 1 +fi + +# ============================================================================= +# NeMo Skills sandbox (for math_formal_lean, ns_tools, etc.) +# ============================================================================= +export SANDBOX_CONTAINER +export SANDBOX_COMMAND="${SANDBOX_COMMAND:-/start-with-nginx.sh}" +export NEMO_SKILLS_SANDBOX_PORT="${NEMO_SKILLS_SANDBOX_PORT:-6000}" + +# ============================================================================= +# Ray log sync +# ============================================================================= +export RAY_LOG_SYNC_FREQUENCY="${RAY_LOG_SYNC_FREQUENCY:-60}" + +CODE_ROOT="/opt/nemo-rl" +VLLM_ENV_SOURCE="source /opt/nemo-rl/3rdparty/vllm/nemo-rl.env && " + +# ============================================================================= +# Persistent cache directories +# ============================================================================= +# Lustre holds the warm persistent cache. At job start, SETUP_COMMAND clears +# stale /tmp caches then seeds node-local /tmp from Lustre. JIT writes go to +# /tmp to avoid Lustre metadata contention from parallel compilation. +_vllm_cache_precision="bf16" +CACHE_READ_DIR="${PERSISTENT_CACHE}/cache_read" +CACHE_WRITE_DIR="${PERSISTENT_CACHE}/cache_write" +LUSTRE_VLLM_CACHE="${CACHE_WRITE_DIR}/vllm_compile_cache_${_vllm_cache_precision}" +LUSTRE_FLASHINFER_CUBIN_CACHE="${PERSISTENT_CACHE}/flashinfer_cubins" +FLASHINFER_CUBIN_CACHE="/tmp/nemo_rl_flashinfer_cubins" +FLASHINFER_WS_BASE="${PERSISTENT_CACHE}/flashinfer_workspace" +LUSTRE_INDUCTOR_CACHE="${PERSISTENT_CACHE}/inductor_cache" +LUSTRE_TRITON_CACHE="${PERSISTENT_CACHE}/triton_cache" +NRL_VLLM_LOCAL_CACHE_DIR="/tmp/nemo_rl_vllm_cache" +NRL_VLLM_CACHE_SEED_DIR="/tmp/nemo_rl_vllm_cache_warm" +INDUCTOR_CACHE_DIR="/tmp/nemo_rl_inductor_cache" +TRITON_CACHE_DIR="/tmp/nemo_rl_triton_cache" +CACHE_SYNC_FREQUENCY="${CACHE_SYNC_FREQUENCY:-0}" + +export LUSTRE_VLLM_CACHE +export LUSTRE_INDUCTOR_CACHE +export LUSTRE_TRITON_CACHE +export CACHE_READ_DIR +export CACHE_WRITE_DIR +export NRL_VLLM_LOCAL_CACHE_DIR +export INDUCTOR_CACHE_DIR +export TRITON_CACHE_DIR +export CACHE_SYNC_FREQUENCY + +mkdir -p "${LUSTRE_FLASHINFER_CUBIN_CACHE}" "${FLASHINFER_WS_BASE}" \ + "${LUSTRE_INDUCTOR_CACHE}" "${LUSTRE_TRITON_CACHE}" \ + "${CACHE_READ_DIR}" "${CACHE_WRITE_DIR}" + +# Read path : cache_read/*.tar.zst — compute nodes extract tarballs (hundreds of concurrent reads) +# Write path : cache_write/*/ — sidecar rsyncs individual files (one sequential writer) +# Splitting reads (tarball) from writes (directory) avoids Lustre MDT invalidation storms +# and lets rsync accumulate the union of all roles' kernels across jobs. +for _name in inductor_cache triton_cache; do + _write_dir="${CACHE_WRITE_DIR}/${_name}" + _old_dir="${PERSISTENT_CACHE}/${_name}" + + # One-time migration: move legacy dir → cache_write/ (instant rename, same FS) + if ([ ! -d "$_write_dir" ] || [ -z "$(ls -A "$_write_dir" 2>/dev/null)" ]) \ + && [ -d "$_old_dir" ] && [ -n "$(ls -A "$_old_dir" 2>/dev/null)" ]; then + [ -d "$_write_dir" ] && rmdir "$_write_dir" 2>/dev/null + mv "$_old_dir" "$_write_dir" 2>/dev/null \ + && echo "[CACHE] Moved legacy ${_name}/ → cache_write/${_name}/" \ + || echo "[CACHE] Failed to move legacy ${_name}/" + fi +done + +# vLLM: migrate the most recent legacy seed dir → cache_write/ (one-time, instant rename) +_vllm_write="${CACHE_WRITE_DIR}/vllm_compile_cache_${_vllm_cache_precision}" +_vllm_read_tar="${CACHE_READ_DIR}/vllm_compile_cache_${_vllm_cache_precision}.tar.zst" + +if [ ! -d "$_vllm_write" ] || [ -z "$(ls -A "$_vllm_write" 2>/dev/null)" ]; then + _best="$(ls -1dt \ + "${PERSISTENT_CACHE}/vllm_compile_cache_${_vllm_cache_precision}" \ + "${PERSISTENT_CACHE}/vllm_compile_cache_${_vllm_cache_precision}_"* \ + 2>/dev/null \ + | while IFS= read -r d; do + [ -d "$d" ] && [ -n "$(ls -A "$d" 2>/dev/null)" ] && echo "$d" && break + done + )" || true + if [ -n "$_best" ]; then + [ -d "$_vllm_write" ] && rmdir "$_vllm_write" 2>/dev/null || true + mv "$_best" "$_vllm_write" 2>/dev/null \ + && echo "[CACHE] Moved $(basename "$_best") → cache_write/vllm_compile_cache_${_vllm_cache_precision}/" \ + || echo "[CACHE] Failed to move vLLM cache" + fi +fi + +# Purge redundant legacy vLLM cache directories. +# The old sidecar wrote every vLLM seed as a separate directory on Lustre +# (e.g. vllm_compile_cache_bf16_2058, _3072, ...). With cache_write/ + tarball, +# only cache_write/vllm_compile_cache_{precision}/ matters. All seed copies are +# content-addressed duplicates — safe to remove after migration. +_purge_count=0 +for _d in "${PERSISTENT_CACHE}/vllm_compile_cache_${_vllm_cache_precision}" \ + "${PERSISTENT_CACHE}/vllm_compile_cache_${_vllm_cache_precision}_"*; do + [ -d "$_d" ] || continue + rm -rf "$_d" 2>/dev/null && (( _purge_count++ )) || true +done +for _d in "${PERSISTENT_CACHE}"/vllm_compile_cache_[0-9]*/; do + [ -d "$_d" ] || continue + rm -rf "$_d" 2>/dev/null && (( _purge_count++ )) || true +done +for _d in "${PERSISTENT_CACHE}/vllm_compile_cache" \ + "${PERSISTENT_CACHE}/vllm_compile_cache_warm"; do + [ -d "$_d" ] || continue + rm -rf "$_d" 2>/dev/null && (( _purge_count++ )) || true +done +if (( _purge_count > 0 )); then + echo "[CACHE] Purged ${_purge_count} redundant legacy vLLM cache directories from ${PERSISTENT_CACHE}/" +fi + +# ============================================================================= +# Code snapshot +# ============================================================================= +# Snapshot the git-tracked source tree so the code is frozen at submission time. +# This guarantees we know exactly which code was used for a given experiment. +# Set USE_SNAPSHOT=0 to skip (runs from container built-in or live checkout). +# Interactive mode defaults to the live checkout for fast iteration; batch snapshots. +if [[ "${INTERACTIVE}" == "1" ]]; then + USE_SNAPSHOT="${USE_SNAPSHOT:-0}" +else + USE_SNAPSHOT="${USE_SNAPSHOT:-1}" +fi + +if [[ "${USE_SNAPSHOT}" == "1" ]]; then + if [[ ! -f "${PROJECT_ROOT}/tools/code_snapshot.sh" ]]; then + echo "ERROR: tools/code_snapshot.sh not found at ${PROJECT_ROOT}/tools/code_snapshot.sh" >&2 + echo " Set USE_SNAPSHOT=0 to run from the live checkout instead." >&2 + exit 1 + fi + SNAPSHOT_DIR=$(bash "${PROJECT_ROOT}/tools/code_snapshot.sh" "${JOB_NAME}") + + if [[ -d "${PROJECT_ROOT}/3rdparty/vllm" ]] && [[ ! -e "${SNAPSHOT_DIR}/3rdparty/vllm" ]]; then + mkdir -p "${SNAPSHOT_DIR}/3rdparty" + ln -s "${PROJECT_ROOT}/3rdparty/vllm" "${SNAPSHOT_DIR}/3rdparty/vllm" + fi + + echo "Code snapshot: ${SNAPSHOT_DIR}" + OVERLAY_SOURCE="${SNAPSHOT_DIR}" +else + OVERLAY_SOURCE="${PROJECT_ROOT}" +fi + +# ============================================================================= +# Container mounts +# ============================================================================= +# By default, nemo_rl (Python package) and examples/configs (YAML configs) from +# the code snapshot are overlaid into the container. Everything else uses the +# container's built-in code at /opt/nemo-rl. +# +# To overlay additional components (e.g. a local Megatron-LM checkout), pass +# EXTRA_MOUNTS as a comma-separated list of host:container pairs: +# +# EXTRA_MOUNTS="/path/to/Megatron-LM:/opt/nemo-rl/3rdparty/Megatron-LM-workspace/Megatron-LM" bash ultra_launch.sh +# +# Container paths for reference: +# /opt/nemo-rl/nemo_rl — Python package +# /opt/nemo-rl/examples/configs — YAML configs +# /opt/nemo-rl/3rdparty/Megatron-LM-workspace/Megatron-LM — Megatron-LM +# /opt/nemo-rl/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge — Megatron-Bridge +# /opt/nemo-rl/3rdparty/Gym-workspace/Gym — NeMo-Gym +# /opt/nemo-rl/3rdparty/vllm — vLLM +# ============================================================================= +_append_mount() { + if [[ -z "${MOUNTS}" ]]; then + MOUNTS="$1" + else + MOUNTS="${MOUNTS},$1" + fi +} + +if [[ -d "${OVERLAY_SOURCE}/nemo_rl" ]]; then + _append_mount "${OVERLAY_SOURCE}/nemo_rl:/opt/nemo-rl/nemo_rl" + echo " Mount: nemo_rl → /opt/nemo-rl/nemo_rl" +fi +if [[ -d "${OVERLAY_SOURCE}/examples/configs" ]]; then + _append_mount "${OVERLAY_SOURCE}/examples/configs:/opt/nemo-rl/examples/configs" + echo " Mount: configs → /opt/nemo-rl/examples/configs" +fi +if [[ -d "${OVERLAY_SOURCE}/3rdparty/Gym-workspace/Gym" ]]; then + _append_mount "${OVERLAY_SOURCE}/3rdparty/Gym-workspace/Gym:/opt/nemo-rl/3rdparty/Gym-workspace/Gym" + echo " Mount: Gym → /opt/nemo-rl/3rdparty/Gym-workspace/Gym" +fi + +if [[ "${USE_SNAPSHOT}" == "1" ]]; then + _append_mount "${SNAPSHOT_DIR}:${SNAPSHOT_DIR}" +fi + +if [[ -n "${EXTRA_MOUNTS:-}" ]]; then + _append_mount "${EXTRA_MOUNTS}" + echo " Extra mounts: ${EXTRA_MOUNTS}" +fi + +export MOUNTS + +# ============================================================================= +# Resolve ray.sub +# ============================================================================= +RAY_SUB="${RAY_SUB:-${PROJECT_ROOT}/ray.sub}" +if [[ ! -f "${RAY_SUB}" ]]; then + echo "ERROR: ray.sub not found at ${RAY_SUB}" >&2 + exit 1 +fi + +# ============================================================================= +# Per-node cache seeding (SETUP_COMMAND) +# ============================================================================= +# Triton, Inductor, and FlashInfer cubins compile/download to node-local /tmp to +# avoid Lustre race conditions and file lock contention during concurrent JIT +# compilation. To avoid cold-start penalties, we seed /tmp from a warm Lustre +# cache before Ray starts. +# +# IMPORTANT: Stale /tmp caches from previous jobs can cause hangs (e.g. the +# Triton bundler skipping non-empty temp dirs). We rm -rf /tmp caches first, +# then seed fresh from Lustre. +# ============================================================================= +read -r -d '' SETUP_COMMAND </dev/null 2>&1 || { apt-get update -qq && apt-get install -y -qq zstd; } 2>/dev/null || true +echo "[CACHE SEED] Clearing stale /tmp caches and seeding from Lustre..." +WARM_SEED="${NRL_VLLM_CACHE_SEED_DIR}" +LOCAL_IND="${INDUCTOR_CACHE_DIR}" +LOCAL_TRI="${TRITON_CACHE_DIR}" +CACHE_READ="${CACHE_READ_DIR}" + +# vLLM caches are per-instance (VLLM_CACHE_ROOT_{seed}). Clear ALL from prior jobs. +rm -rf /tmp/nemo_rl_vllm_cache /tmp/nemo_rl_vllm_cache_* +rm -rf "\$LOCAL_IND" "\$LOCAL_TRI" +mkdir -p "\$LOCAL_IND" "\$LOCAL_TRI" + +_seed_cache() { + local tarball="\$1" local_dir="\$2" name="\$3" + if [ -f "\$tarball" ]; then + tar --zstd -xf "\$tarball" -C "\$local_dir" \ + && echo "[CACHE SEED] \$name: seeded from tarball (\$(du -sh "\$local_dir" 2>/dev/null | cut -f1))" \ + || echo "[CACHE SEED] \$name: tarball extract failed (non-fatal)" + else + echo "[CACHE SEED] \$name: no warm cache on Lustre yet" + fi +} + +# Seed vLLM compile cache from cache_read/ tarball (one per precision). +rm -rf "\$WARM_SEED" +_vllm_tar="\$CACHE_READ/vllm_compile_cache_${_vllm_cache_precision}.tar.zst" +if [ -f "\$_vllm_tar" ]; then + mkdir -p "\$WARM_SEED" + tar --zstd -xf "\$_vllm_tar" -C "\$WARM_SEED" \ + && echo "[CACHE SEED] vLLM (${_vllm_cache_precision}): seeded from tarball (\$(du -sh "\$WARM_SEED" 2>/dev/null | cut -f1))" \ + || echo "[CACHE SEED] vLLM: tarball extract failed (non-fatal)" +else + echo "[CACHE SEED] vLLM: no warm cache on Lustre yet" +fi + +_seed_cache "\$CACHE_READ/inductor_cache.tar.zst" "\$LOCAL_IND" "Inductor" +_seed_cache "\$CACHE_READ/triton_cache.tar.zst" "\$LOCAL_TRI" "Triton" + +echo "[CACHE SEED] Done." +SETUPEOF +export SETUP_COMMAND + +# ============================================================================= +# Build the training command +# ============================================================================= +# Stage-specific hyperparameters (batch sizes, advantage clip, MoE parallelism, +# learning rate, etc.) live in CONFIG_PATH. The launcher only passes the +# per-run overrides: cluster shape, paths, judge endpoints, logging. +# ============================================================================= +TRAIN_CMD="cd ${CODE_ROOT} && date ; \ +${VLLM_ENV_SOURCE}\ +OMP_NUM_THREADS=16 \ +RAY_DEDUP_LOGS=1 \ +WANDB_INIT_TIMEOUT=300 \ +VLLM_CACHE_ROOT=${NRL_VLLM_LOCAL_CACHE_DIR} \ +NRL_VLLM_CACHE_SEED_DIR=${NRL_VLLM_CACHE_SEED_DIR} \ +DG_JIT_CACHE_DIR=${NRL_VLLM_LOCAL_CACHE_DIR}/deep_gemm \ +TORCHINDUCTOR_CACHE_DIR=${INDUCTOR_CACHE_DIR} \ +TRITON_CACHE_DIR=${TRITON_CACHE_DIR} \ +UV_CACHE_DIR=${PERSISTENT_CACHE}/uv \ +RAY_ENABLE_UV_RUN_RUNTIME_ENV=0 \ +UV_HTTP_TIMEOUT=10 \ +VLLM_USE_FLASHINFER_MOE_FP8=1 \ +VLLM_FLASHINFER_MOE_BACKEND=latency \ +NRL_VLLM_ASYNC_TIMEOUT_SECONDS=1800 \ +NRL_WG_USE_RAY_REF=1 \ +HF_HOME=${HF_HOME:-} \ +HF_TOKEN=${HF_TOKEN:-} \ +NRL_USE_FASTOKENS=${NRL_USE_FASTOKENS:-1} \ +uv run ./examples/nemo_gym/run_grpo_nemo_gym.py \ +--config ${CONFIG_PATH} \ +policy.model_name=${MODEL_PATH} \ +cluster.num_nodes=${NUM_ACTOR_NODES} \ +policy.generation.colocated.resources.num_nodes=${NUM_GEN_NODES} \ +checkpointing.checkpoint_dir=${CHECKPOINT_DIR} \ +${CHECKPOINTING_SAVE_BY:+checkpointing.checkpoint_must_save_by=${CHECKPOINTING_SAVE_BY}} \ +data.train.data_path=${TRAIN_PATH} \ +data.validation.data_path=${VAL_PATH} \ +${GENRM_OVERRIDE:+${GENRM_OVERRIDE}} \ +${NL2BASH_JUDGE_MODEL:+env.nemo_gym.nl2bash_judge_model.responses_api_models.local_vllm_model.model=${NL2BASH_JUDGE_MODEL}} \ +${SAFETY_JUDGE_MODEL:+env.nemo_gym.safety_judge_model.responses_api_models.local_vllm_model.model=${SAFETY_JUDGE_MODEL}} \ +${SIF_DIR:+sif_dir=${SIF_DIR}} \ +env.nemo_gym.nemo_gym_log_dir=${LOG_DIR}/nemo_gym \ +logger.log_dir=${LOG_DIR} \ +logger.wandb_enabled=${WANDB_ENABLED} \ +logger.wandb.name=${WANDB_NAME} \ +logger.wandb.project=${WANDB_PROJ} \ +${NRL_MAX_STEPS:+grpo.max_num_steps=${NRL_MAX_STEPS}} \ +${MTP_EXTRA_ARGS} \ +${MOPD_OVERRIDES} \ +${*}" + +export COMMAND="${TRAIN_CMD}" + +# ============================================================================= +# Summary +# ============================================================================= +echo "" +echo "================================================================" +echo " Nemotron 3 Ultra — ${EXP_NAME} (${NUM_TOTAL_NODES}-node)" +echo "================================================================" +echo " Job name: ${JOB_NAME} (singleton — only one runs at a time)" +echo " Config: ${CONFIG_PATH}" +echo " Nodes: ${NUM_TOTAL_NODES} total (segment=${SEGMENT_SIZE})" +echo " Training: ${NUM_TRAIN_NODES} ($((NUM_TRAIN_NODES * GPUS_PER_NODE)) GPUs)" +echo " vLLM gen: ${NUM_GEN_NODES} ($((NUM_GEN_NODES * GPUS_PER_NODE)) GPUs)" +echo " Gym: ${NUM_GYM_NODES} ($((NUM_GYM_NODES * GPUS_PER_NODE)) GPUs)" +if (( NUM_TEACHER_NODES > 0 )); then +echo " Teachers: ${NUM_TEACHER_NODES} ($((NUM_TEACHER_NODES * GPUS_PER_NODE)) GPUs)" +fi +echo " Walltime: ${WALLTIME}" +echo "" +echo " Checkpoints: ${CHECKPOINT_DIR} (stable — auto-resumes across jobs)" +echo " Run dir: ${RUN_DIR}" +echo " Logs: ${LOG_DIR}" +echo " Slurm logs: ${SLURM_LOG_DIR}" +echo " W&B: ${WANDB_PROJ} / ${WANDB_NAME} (enabled=${WANDB_ENABLED})" +echo "" +echo " Model: ${MODEL_PATH}" +echo " Train data: ${TRAIN_PATH}" +echo " Val data: ${VAL_PATH}" +echo " Container: ${CONTAINER}" +echo " Sandbox: ${SANDBOX_CONTAINER}" +if [[ "${USE_SNAPSHOT}" == "1" ]]; then +echo " Snapshot: ${SNAPSHOT_DIR}" +fi +echo "" +echo " Monitor: squeue -u \$USER -n ${JOB_NAME}" +echo " Logs: tail -f ${SLURM_LOG_DIR}/*.out" +echo " Latest: ls -la ${RESULTS_DIR}/runs/latest" +echo "" +echo "================================================================" +echo "" + +# ============================================================================= +# Record code provenance in the run directory +# ============================================================================= +{ + echo "timestamp: $(date -Iseconds)" + echo "branch: $(git -C "${PROJECT_ROOT}" rev-parse --abbrev-ref HEAD 2>/dev/null || echo unknown)" + echo "commit: $(git -C "${PROJECT_ROOT}" rev-parse HEAD 2>/dev/null || echo unknown)" + echo "dirty: $(git -C "${PROJECT_ROOT}" status --porcelain 2>/dev/null | head -20)" + echo "snapshot: ${USE_SNAPSHOT}" + if [[ "${USE_SNAPSHOT}" == "1" ]]; then + echo "snapshot_dir: ${SNAPSHOT_DIR}" + fi + echo "container: ${CONTAINER}" + echo "config: ${CONFIG_PATH}" + echo "command: ${TRAIN_CMD}" +} > "${RUN_DIR}/provenance.txt" + +# ============================================================================= +# Dry-run mode: print everything, don't submit +# ============================================================================= +DRY_RUN="${DRY_RUN:-0}" +if [[ "${DRY_RUN}" == "1" ]]; then + echo "DRY_RUN=1 — printing TRAIN_CMD and exiting without submission." + echo "" + echo "--- TRAIN_CMD ---" + echo "${TRAIN_CMD}" + echo "--- end ---" + exit 0 +fi + +# ============================================================================= +# Interactive mode: bring up Ray and idle for attachment (no training driver) +# ============================================================================= +# With COMMAND empty, ray.sub starts the Ray cluster, writes -attach.sh, +# then idles. We save the driver command to -run-cmd.sh so you can attach +# and run it by hand, edit it, and re-run without requeueing. +if [[ "${INTERACTIVE}" == "1" ]]; then + unset COMMAND 2>/dev/null || true # empty COMMAND -> ray.sub idle/interactive mode + WALLTIME="${INTERACTIVE_WALLTIME:-${WALLTIME}}" + + echo "" + echo "================================================================" + echo " INTERACTIVE MODE — ${NUM_TOTAL_NODES}-node allocation (walltime ${WALLTIME})" + echo " Ray will start and idle until you attach." + echo "================================================================" + + SBATCH_OUTPUT=$(sbatch \ + --nodes="${NUM_TOTAL_NODES}" \ + --account="${SLURM_ACCOUNT}" \ + --job-name="interactive-${JOB_NAME}" \ + --partition="${SLURM_PARTITION}" \ + --time="${WALLTIME}" \ + --gres=gpu:${GPUS_PER_NODE} \ + --exclusive \ + --mem=0 \ + --segment="${SEGMENT_SIZE}" \ + --output="${SLURM_LOG_DIR}/%j.out" \ + --error="${SLURM_LOG_DIR}/%j.err" \ + ${SLURM_QOS:+--qos="${SLURM_QOS}"} \ + ${EXCLUDE_NODES:+--exclude="${EXCLUDE_NODES}"} \ + ${SLURM_RESERVATION:+--reservation="${SLURM_RESERVATION}"} \ + "${RAY_SUB}") + echo "${SBATCH_OUTPUT}" + JOB_ID=$(echo "${SBATCH_OUTPUT}" | grep -oP '\d+$') + [[ -z "${JOB_ID}" ]] && { echo "ERROR: could not parse job ID from sbatch output." >&2; exit 1; } + + LAUNCH_DIR="$(pwd)" + ATTACH_SCRIPT="${LAUNCH_DIR}/${JOB_ID}-attach.sh" + CMD_FILE="${LAUNCH_DIR}/${JOB_ID}-run-cmd.sh" + cat > "${CMD_FILE}" </dev/null || true) + [[ -z "${state}" ]] && { echo " Job ${JOB_ID} left the queue. Check: sacct -j ${JOB_ID}"; exit 1; } + [[ "${state}" != "${prev_state}" ]] && { echo " [$(date +%H:%M:%S)] state: ${state}"; prev_state="${state}"; } + sleep 15 + done + echo "" + echo " Ray is ready — attach: bash ${ATTACH_SCRIPT}" + fi + exit 0 +fi + +# ============================================================================= +# Submit +# ============================================================================= +# Always serialise same-name submissions via singleton; optionally chain after +# another job with SLURM_DEPENDENCY (e.g. "afterany:3044848" or "afterok:JOBID"). +SLURM_DEPENDENCY="${SLURM_DEPENDENCY:-}" +DEPENDENCY="singleton" +[[ -n "${SLURM_DEPENDENCY}" ]] && DEPENDENCY="singleton,${SLURM_DEPENDENCY}" + +SBATCH_OUTPUT=$(sbatch \ + --nodes="${NUM_TOTAL_NODES}" \ + --account="${SLURM_ACCOUNT}" \ + --job-name="${JOB_NAME}" \ + --partition="${SLURM_PARTITION}" \ + --time="${WALLTIME}" \ + --gres=gpu:${GPUS_PER_NODE} \ + --exclusive \ + --mem=0 \ + --dependency="${DEPENDENCY}" \ + --segment="${SEGMENT_SIZE}" \ + --output="${SLURM_LOG_DIR}/%j.out" \ + --error="${SLURM_LOG_DIR}/%j.err" \ + ${SLURM_QOS:+--qos="${SLURM_QOS}"} \ + ${EXCLUDE_NODES:+--exclude="${EXCLUDE_NODES}"} \ + ${SLURM_RESERVATION:+--reservation="${SLURM_RESERVATION}"} \ + "${RAY_SUB}") + +echo "${SBATCH_OUTPUT}" +JOB_ID=$(echo "${SBATCH_OUTPUT}" | grep -oP '\d+$') + +if [[ -n "${JOB_ID}" ]]; then + echo "" + echo " Ray logs: ${BASE_LOG_DIR}/${JOB_ID}-logs/" + echo "" +fi From 0720271b344f5b04ea4aff573b5083a236a05e56 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Wed, 24 Jun 2026 18:16:34 +0000 Subject: [PATCH 02/15] add original ultra-v3 recipes Signed-off-by: Anish Mahishi --- examples/configs/ultra/ifbench_teacher.yaml | 654 ++++++++++++++ examples/configs/ultra/mopd.yaml | 810 ++++++++++++++++++ examples/configs/ultra/reasoning_teacher.yaml | 500 +++++++++++ examples/configs/ultra/rlhf_teacher.yaml | 520 +++++++++++ examples/configs/ultra/student_rlvr1.yaml | 673 +++++++++++++++ examples/configs/ultra/student_rlvr2.yaml | 674 +++++++++++++++ examples/configs/ultra/swe_teacher.yaml | 476 ++++++++++ 7 files changed, 4307 insertions(+) create mode 100644 examples/configs/ultra/ifbench_teacher.yaml create mode 100644 examples/configs/ultra/mopd.yaml create mode 100644 examples/configs/ultra/reasoning_teacher.yaml create mode 100644 examples/configs/ultra/rlhf_teacher.yaml create mode 100644 examples/configs/ultra/student_rlvr1.yaml create mode 100644 examples/configs/ultra/student_rlvr2.yaml create mode 100644 examples/configs/ultra/swe_teacher.yaml diff --git a/examples/configs/ultra/ifbench_teacher.yaml b/examples/configs/ultra/ifbench_teacher.yaml new file mode 100644 index 0000000000..badeb263b3 --- /dev/null +++ b/examples/configs/ultra/ifbench_teacher.yaml @@ -0,0 +1,654 @@ +# ============================================================================= +# Nemotron 3 Ultra — IFBench Teacher (instruction-following / abstention RLHF) +# ============================================================================= +# RLHF teacher specializing in instruction following, abstention, and +# omniscience-style refusal. Starts from the Student RLVR output and trains +# at 49k context with a smaller batch (GBS=2048) and lower LR (2.5e-6) than +# the student RLVR recipe. +# +# Cluster shape (80 nodes × 4 GPUs on GB200 NVL72 = 320 GPUs): +# Training 32 / vLLM 28 / Gym 20. Override via NUM_TRAIN_NODES, etc. +# Gym is sized at 20 nodes (= 80 GPUs) so GenRM can serve locally +# (TP=4 DP=4 = 16 GPUs) alongside the NL2Bash and Safety judges. +# +# Training parallelism: TP=8, EP=64, CP=8, PP=1, SP=true. +# vLLM parallelism: TP=8, EP=8 (EP=TP keeps vllm_dp_size=1; EP>TP is blocked +# by https://github.com/NVIDIA-NeMo/RL/issues/1101). +# ============================================================================= + +# ============================================================================= +# Cluster — read by ultra_launch.sh to size the SLURM allocation +# ============================================================================= +cluster: + gpus_per_node: 4 + num_nodes: 60 + segment_size: 16 + +# ============================================================================= +# Checkpointing — checkpoint_dir is overridden by the launcher +# ============================================================================= +checkpointing: + enabled: true + checkpoint_dir: "results/ultra_student_rlvr" + metric_name: "val:total_reward/mean" + higher_is_better: true + keep_top_k: 1000000 + save_period: 10 + ft_keep_latest_k: 1 + ft_save_period: 1 + checkpoint_must_save_by: "00:47:35:00" + model_save_format: "safetensors" + save_consolidated: false + +# ============================================================================= +# GRPO Algorithm +# ============================================================================= +grpo: + num_prompts_per_step: 128 + num_generations_per_prompt: 16 + num_val_generations_per_prompt: 2 + max_rollout_turns: 1 + max_num_epochs: 1 + max_num_steps: 1000000 + normalize_rewards: true + use_leave_one_out_baseline: true + advantage_clip_low: -20 + advantage_clip_high: 20 + val_period: 10000 + val_at_start: false + val_at_end: false + overlong_filtering: false + max_val_samples: null + val_batch_size: 256 + seed: 42 + + use_dynamic_sampling: false + dynamic_sampling_max_gen_batches: 10 + batch_multiplier: 1 + + penalize_invalid_tool_call: true + invalid_tool_call_advantage: -5.0 + + reward_shaping: + enabled: false + overlong_buffer_length: 128 + overlong_buffer_penalty: 1 + max_response_length: ${policy.max_total_sequence_length} + stop_properly_penalty_coef: null + reward_scaling: + enabled: false + source_min: 0.0 + source_max: 1.0 + target_min: 0.0 + target_max: 1.0 + + async_grpo: + enabled: true + max_trajectory_age_steps: 1 + in_flight_weight_updates: true + recompute_kv_cache_after_weight_updates: false + + use_best_at_k: false + best_at_k_k: 8 + best_at_k_m: 1000 + + use_combined_training: false + combined_training_weight_mode: "auto" + combined_training_best_at_k_weight: 0.2 + combined_training_pass_at_1_weight: 1.0 + + dynamic_sampling_oversample_ratio: 1.0 + seq_logprob_error_threshold: 2 + +# ============================================================================= +# Loss Function +# ============================================================================= +loss_fn: + reference_policy_kl_penalty: 0.0 + reference_policy_kl_type: "k3" + kl_input_clamp_value: null + kl_output_clamp_value: null + + ratio_clip_min: 0.2 + ratio_clip_max: 0.28 + ratio_clip_c: null + use_on_policy_kl_approximation: true + use_importance_sampling_correction: true + truncated_importance_sampling_ratio: 5 + truncated_importance_sampling_ratio_min: 0.2 + truncated_importance_sampling_type: tis + sequence_level_importance_ratios: false + token_level_loss: true + force_on_policy_ratio: true + use_kl_in_reward: false + +# ============================================================================= +# Policy — model_name is overridden by the launcher (MODEL_PATH env var) +# ============================================================================= +policy: + model_name: null + tokenizer: + name: ${policy.model_name} + chat_template_kwargs: null + hf_config_overrides: {} + + train_global_batch_size: 2048 + train_micro_batch_size: 1 + generation_batch_size: 64 + logprob_batch_size: 1 + max_total_sequence_length: 49152 + precision: "bfloat16" + logprob_chunk_size: 2048 + offload_optimizer_for_logprob: false + + dtensor_cfg: + _v2: true + enabled: false + cpu_offload: false + sequence_parallel: false + activation_checkpointing: false + tensor_parallel_size: 1 + context_parallel_size: 1 + custom_parallel_plan: null + + megatron_cfg: + enabled: true + empty_unused_memory_level: 2 + activation_checkpointing: true + + # TP=8 spans 2 GB200 nodes (4 GPUs each) via intra-rack NVLink. + tensor_model_parallel_size: 8 + expert_tensor_parallel_size: 1 + # EP=64: 512 experts / 64 = 8 experts per EP rank + # All-to-all spans 64 GPUs (16 nodes), fits within one NVLink domain + expert_model_parallel_size: 64 + pipeline_model_parallel_size: 1 + num_layers_in_first_pipeline_stage: null + num_layers_in_last_pipeline_stage: null + context_parallel_size: 8 + pipeline_dtype: ${policy.precision} + sequence_parallel: true + + # MoE + freeze_moe_router: true + moe_router_dtype: "fp32" + moe_router_load_balancing_type: "none" + moe_router_bias_update_rate: 1.0e-3 + moe_router_enable_expert_bias: true + moe_permute_fusion: true + moe_enable_deepep: false + moe_token_dispatcher_type: "flex" + moe_flex_dispatcher_backend: "hybridep" + moe_hybridep_num_sms: 32 + moe_aux_loss_coeff: 0.0 + moe_shared_expert_overlap: false + + # Compute + apply_rope_fusion: true + use_fused_weighted_squared_relu: true + bias_activation_fusion: false + defer_fp32_logits: true + + # Logging + track_moe_metrics: true + moe_per_layer_logging: true + do_not_average_loss: true + cp_normalize: true + calculate_per_token_loss: true + scale_loss_by_dp_cp_size: false + + # MTP — disabled by default; enable via ENABLE_MTP_INFERENCE=1 in the launcher + mtp_loss_scaling_factor: 0.3 + mtp_use_repeated_layer: true + mtp_num_layers: 5 + mtp_detach_heads: true + + optimizer: + optimizer: "adam" + lr: 2.5e-6 + min_lr: 2.5e-6 + weight_decay: 0.0 + bf16: true + fp16: false + params_dtype: "float32" + + adam_beta1: 0.9 + adam_beta2: 0.999 + adam_eps: 1e-8 + + sgd_momentum: 0.9 + + use_distributed_optimizer: true + use_precision_aware_optimizer: true + + clip_grad: ${policy.max_grad_norm} + + optimizer_cpu_offload: false + optimizer_offload_fraction: 0.0 + + scheduler: + start_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} + end_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} + weight_decay_incr_style: "constant" + lr_decay_style: "constant" + lr_decay_iters: null + lr_warmup_iters: 10 + lr_warmup_init: 4e-7 + override_opt_param_scheduler: true + + distributed_data_parallel_config: + grad_reduce_in_fp32: false + overlap_grad_reduce: false + overlap_param_gather: true + average_in_collective: false + use_custom_fsdp: false + data_parallel_sharding_strategy: "optim_grads_params" + + fp8_cfg: + enabled: false + fp8: "e4m3" + fp8_recipe: "mxfp8" + fp8_param: false + + first_last_layers_bf16: true + num_layers_at_start_in_bf16: 1 + num_layers_at_end_in_bf16: 1 + + use_gloo_process_groups: false + + checkpoint: + async_save: true + ckpt_assume_constant_structure: true + fully_parallel_save_process_group: "ep_dp" + fully_parallel_load_process_group: "ep_dp" + fully_parallel_load_exchange_algo: "broadcast" + + env_vars: null + + # --------------------------------------------------------------------------- + # Sequence Packing + # --------------------------------------------------------------------------- + dynamic_batching: + enabled: false + train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}} + sequence_length_round: 64 + + sequence_packing: + enabled: true + train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}} + algorithm: "modified_first_fit_decreasing" + sequence_length_round: 64 + fuse_loss: true + + make_sequence_length_divisible_by: ${policy.megatron_cfg.tensor_model_parallel_size} + max_grad_norm: 1.0 + optimizer: null + scheduler: null + + # --------------------------------------------------------------------------- + # Generation (vLLM) — Non-colocated, async + # --------------------------------------------------------------------------- + generation: + port_range_low: 3000 + port_range_high: 4999 + backend: "vllm" + max_new_tokens: ${policy.max_total_sequence_length} + temperature: 1.0 + top_p: 1.0 + top_k: null + stop_token_ids: null + stop_strings: null + # TP=8 EP=8: EP=TP so vllm_dp_size=1, async_engine=true works with NeMo Gym. + # Same memory as EP=1 but uses all-to-all for expert routing. + # EP > TP blocked by https://github.com/NVIDIA-NeMo/RL/issues/1101. + vllm_cfg: + async_engine: true + precision: ${policy.precision} + kv_cache_dtype: "auto" + tensor_parallel_size: 8 + pipeline_parallel_size: 1 + expert_parallel_size: 8 + gpu_memory_utilization: 0.85 + max_model_len: ${policy.max_total_sequence_length} + enforce_eager: false + use_deep_gemm: false + num_last_layers_in_bf16: 0 + num_first_layers_in_bf16: 0 + enable_vllm_metrics_logger: true + vllm_metrics_logger_interval: 0.5 + expose_http_server: true + skip_tokenizer_init: false + http_server_serving_chat_kwargs: + enable_auto_tools: true + tool_parser: qwen3_coder + reasoning_parser: nano_v3 + reasoning_parser_plugin: nemo_rl/utils/nano_v3_reasoning_parser.py + + vllm_kwargs: + attention_backend: FLASH_ATTN + mamba_ssm_cache_dtype: "float32" + compilation_config: + cudagraph_capture_sizes: [1,2,4,8,16,32,64] + pass_config: + fuse_allreduce_rms: false + + colocated: + enabled: false + resources: + gpus_per_node: 4 + num_nodes: 28 # Overridden by launch script via NUM_GEN_NODES + +# ============================================================================= +# Data — data_path is overridden by the launcher (TRAIN_PATH / VAL_PATH) +# ============================================================================= +data: + max_input_seq_length: null + shuffle: false + num_workers: 1 + train: + data_path: null + validation: + data_path: null + default: + dataset_name: NemoGymDataset + env_name: "nemo_gym" + prompt_file: null + system_prompt_file: null + processor: "nemo_gym_data_processor" + +# ============================================================================= +# Environment — NeMo Gym + Judge Models +# +# Judge model paths and the GenRM service URL are set by the launcher: +# - genrm_model.base_url ← GENRM_BASE_URL +# - nl2bash_judge_model.model ← NL2BASH_JUDGE_MODEL +# - safety_judge_model.model ← SAFETY_JUDGE_MODEL +# ============================================================================= +env: + should_use_nemo_gym: true + should_log_nemo_gym_responses: true + nemo_gym: + nemo_gym_log_dir: "logs/nemo_gym" + skip_venv_if_present: true + port_range_low: 5000 + port_range_high: 5999 + invalid_tool_call_patterns: + - "" + - "" + - "" + - "" + thinking_tags: + - "" + - "" + # Only the resources servers exercised by the IFBench training data are + # loaded — abstention, GenRM, multichallenge, instruction_following, + # inverse_if. Production loads the full ~35-entry list for all teacher + # stages; we trim to keep gym startup focused. + config_paths: + - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml + - resources_servers/abstention/configs/abstention.yaml + - resources_servers/genrm_compare/configs/genrm_compare.yaml + - resources_servers/instruction_following/configs/instruction_following.yaml + - resources_servers/inverse_if/configs/inverse_if.yaml + - resources_servers/multichallenge/configs/multichallenge.yaml + + policy_model: + responses_api_models: + vllm_model: + num_workers: 16 + num_groups_nemo_rl: ${add:${grpo.async_grpo.max_trajectory_age_steps}, 1} + # Model-specific (Nemotron chat template): keep reasoning in multi-turn + # history so the trajectory grows monotonically for token-in-token-out RL. + chat_template_kwargs: + enable_thinking: true + truncate_history_thinking: false + sequential_reasoning_allowed: false + policy_model_reasoning_off: + _copy: policy_model + responses_api_models: + vllm_model: + num_workers: 4 + num_groups_nemo_rl: ${add:${grpo.async_grpo.max_trajectory_age_steps}, 1} + chat_template_kwargs: + enable_thinking: false + + abstention: + resources_servers: + abstention: + judge_model_server: + type: responses_api_models + name: nl2bash_judge_model + judge_responses_create_params: + max_output_tokens: 8192 + abstention_reward: 0.9 + + # Safety judge: TP=4 ensures each placement group claims a full GB200 node, + # avoiding GPU fragmentation that can block larger-TP models. + jailbreak_detection: + resources_servers: + jailbreak_detection: + judge_model_server: + type: responses_api_models + name: safety_judge_model + + safety_judge_model: + responses_api_models: + local_vllm_model: + entrypoint: app.py + model: null # Set by launcher (SAFETY_JUDGE_MODEL) + return_token_id_information: false + uses_reasoning_parser: false + debug: true + vllm_serve_env_vars: + VLLM_RAY_DP_PACK_STRATEGY: strict + + vllm_serve_kwargs: + attention_backend: TRITON_ATTN + tensor_parallel_size: 4 + data_parallel_size: 2 + data_parallel_size_local: 1 + pipeline_parallel_size: 1 + gpu_memory_utilization: 0.85 + max_model_len: 96000 + model_loader_extra_config: + enable_multithread_load: true + num_threads: 112 + compilation_config: + cudagraph_capture_sizes: [1,2,4,8,16] + + # NL2Bash / general judge: TP=4 on GB200 192GB + nl2bash_judge_model: + responses_api_models: + local_vllm_model: + entrypoint: app.py + model: null # Set by launcher (NL2BASH_JUDGE_MODEL) + return_token_id_information: false + uses_reasoning_parser: false + debug: true + vllm_serve_env_vars: + VLLM_RAY_DP_PACK_STRATEGY: strict + + vllm_serve_kwargs: + attention_backend: FLASH_ATTN + tensor_parallel_size: 4 + data_parallel_size: 4 + data_parallel_size_local: 1 + pipeline_parallel_size: 1 + enable_expert_parallel: true + enable_auto_tool_choice: true + tool_call_parser: hermes + gpu_memory_utilization: 0.85 + max_model_len: 131072 + enable_prefix_caching: true + enable_chunked_prefill: true + model_loader_extra_config: + enable_multithread_load: true + num_threads: 112 + compilation_config: + cudagraph_capture_sizes: [1,2,4,8,16,32] + server_env: + NCCL_MNNVL_ENABLE: "0" + + inverse_if: + resources_servers: + inverse_if: + judge_model_server: + type: responses_api_models + name: nl2bash_judge_model + + multichallenge: + resources_servers: + multichallenge: + judge_model_server: + type: responses_api_models + name: nl2bash_judge_model + judge_responses_create_params: + max_output_tokens: 8192 + + equivalence_llm_judge: + resources_servers: + equivalence_llm_judge: + judge_model_server: + name: nl2bash_judge_model + judge_responses_create_params: + max_output_tokens: 8192 + + # GenRM compare: deployed as a separate service; URL set by launcher + genrm_compare_resources_server: + resources_servers: + genrm_compare: + num_rollouts_per_prompt: ${grpo.num_generations_per_prompt} + genrm_model_server: + type: responses_api_models + name: genrm_model + genrm_responses_create_params: + max_output_tokens: 24576 + temperature: 1.0 + top_p: 0.95 + comparison_strategy: "circular" + num_judges_per_comparison: 1 + use_principle: true + default_principle: "You will be given one or more evaluation criteria (rubrics).\nEvaluate both responses on EACH criterion individually first, then synthesize an overall judgment.\nCriteria:\n\n1. Please act as an impartial judge and evaluate the quality of the responses provided by two AI assistants to the user prompt. Begin your evaluation by generating your own answer to the prompt. You must provide your answer before judging any answers. When evaluating the assistants' answers, compare both assistants' answers with your answer. You must identify and correct any mistakes or inaccurate information. Then consider if the assistant's answers are helpful, relevant, and concise. Helpful means the answer correctly responds to the prompt or follows the instructions. Note when user prompt has any ambiguity or more than one interpretation, it is more helpful and appropriate to ask for clarifications or more information from the user than providing an answer based on assumptions. Relevant means all parts of the response closely connect or are appropriate to what is being asked. Concise means the response is clear and not verbose or excessive. Then consider the creativity and novelty of the assistant's answers when needed. Finally, identify any missing important information in the assistants' answers that would be beneficial to include when responding to the user prompt." + aggregator_method: "simple_tiebreaker" + reasoning_bonus: 0.5 + answer_bonus: 0.5 + top_percentile: 0.2 + genrm_parse_retries: 1 + group_reasoning_length_penalty_coeff: 0.1 + use_golden_anchor: false + group_answer_length_penalty_coeff: 0.1 + group_style_penalty_coeff: 0.0 + default_score: 3.0 + default_ranking: 3.5 + + # Launcher sets either `base_url` (route to remote GenRM service) or + # `model` (serve locally using vllm_serve_kwargs below). See the guide. + genrm_model: + responses_api_models: + genrm_model: + entrypoint: app.py + api_key: dummy_key + model: null + uses_reasoning_parser: true + return_token_id_information: false + debug: true + ray_worker_py_executable: /opt/ray_venvs/nemo_rl.models.generation.vllm.vllm_worker_async.VllmAsyncGenerationWorker/bin/python + vllm_serve_env_vars: + VLLM_RAY_DP_PACK_STRATEGY: strict + + vllm_serve_kwargs: + tensor_parallel_size: 4 + data_parallel_size: 4 + data_parallel_size_local: 1 + pipeline_parallel_size: 1 + reasoning_parser: deepseek_r1 + gpu_memory_utilization: 0.85 + max_model_len: 131072 + enable_prefix_caching: true + enable_chunked_prefill: true + model_loader_extra_config: + enable_multithread_load: true + num_threads: 112 + load_format: instanttensor + compilation_config: + cudagraph_capture_sizes: [1,2,4,8,16,32] + server_env: + NCCL_MNNVL_ENABLE: "0" + + lc_judge: + resources_servers: + equivalence_llm_judge: + judge_model_server: + name: nl2bash_judge_model + judge_responses_create_params: + max_output_tokens: 8192 + + math_with_judge: + resources_servers: + math_with_judge: + judge_model_server: + name: nl2bash_judge_model + judge_responses_create_params: + max_output_tokens: 8192 + should_use_judge: true + + code_gen: + resources_servers: + code_gen: + num_processes: 2048 + unit_test_timeout_secs: 10 + debug: false + + math_formal_lean_refinement_agent: + responses_api_agents: + proof_refinement_agent: + max_correction_turns: 2 + +# ============================================================================= +# Logger — log_dir, wandb.name, wandb.project are overridden by the launcher +# ============================================================================= +logger: + log_dir: "logs/ultra_student_rlvr" + num_val_samples_to_print: 0 + wandb_enabled: false + tensorboard_enabled: false + mlflow_enabled: false + monitor_gpus: true + swanlab_enabled: false + wandb: + project: "nemotron-3-ultra" + name: "ultra-student-rlvr" + tensorboard: {} + mlflow: + experiment_name: "nemotron-3-ultra" + run_name: "ultra-student-rlvr" + gpu_monitoring: + collection_interval: 10 + flush_interval: 10 + +# ============================================================================= +# Effort Levels +# ============================================================================= +effort_levels: + low_string: "{reasoning effort: efficient}" + low_weight: 0.1 + low_penalty: 1 + low_ub: 15000 + +# ============================================================================= +# Token IDs (model-specific) +# ============================================================================= +token_ids: + eos: 2 + think_open: 12 + think_close: 13 + +# ============================================================================= +# Reward Penalties (set reward to 0 when triggered) +# ============================================================================= +penalize_duplicated_reasoning: true +penalize_empty_final_answer: true +penalize_eos_token: true +penalize_malformed_think_tag: true diff --git a/examples/configs/ultra/mopd.yaml b/examples/configs/ultra/mopd.yaml new file mode 100644 index 0000000000..33ec40d077 --- /dev/null +++ b/examples/configs/ultra/mopd.yaml @@ -0,0 +1,810 @@ +# ============================================================================= +# Nemotron 3 Ultra — MOPD Stage 1 (on-policy distillation from teacher panel) +# ============================================================================= +# Multi-On-Policy Distillation. The Student RLVR output is the student; each +# Gym agent is routed to one of the teacher checkpoints below via +# `_teachers.`. Missing teachers fall back to `_teachers.general` +# (Student RLVR output) so the recipe runs even without the full panel. +# +# Required teacher (set via launcher → `_teachers.general=`): +# - general (= Student RLVR output; also serves as the MOPD student) +# +# Optional teachers (each defaults to `${_teachers.general}` when unset): +# - rlhf (RLHF Teacher output) +# - ifbench (IFBench Teacher output; also serves omniscience/multichallenge) +# - reasoning (Reasoning Teacher output; also serves ns_tools/code/math/mcqa) +# - swe (SWE Teacher output; also serves agentic and pivot/tool-use roles) +# +# Production MOPD also distinguishes agentic-safety / gdpval / search / +# terminalbench / tau teachers. We fall those back to `general` until their +# recipes are ready. +# +# Cluster shape (256 nodes × 4 GPUs on GB200 NVL72 = 1024 GPUs): +# Training 64 / vLLM 142 / Gym 10 / Teachers 5×4=20 (with deduplication 4 +# actual replicas when only `general` is set; up to 5 replicas when all +# teacher overrides are passed). Plus a 20-node buffer for one extra unique +# teacher in stage-2 / future use → 256 total. +# +# Training parallelism: TP=8, EP=64, CP=8, PP=1, SP=true. +# Teacher parallelism (`non_colocated_teachers.default_teacher_cfg`): TP=8, +# CP=2, EP=16, 4 nodes per teacher. +# vLLM parallelism: TP=8, EP=8. +# ============================================================================= + +# ============================================================================= +# Cluster — overridden by launch script +# ============================================================================= +cluster: + gpus_per_node: 4 + num_nodes: 256 + segment_size: 16 + +# ============================================================================= +# Checkpointing +# ============================================================================= +checkpointing: + enabled: true + checkpoint_dir: "results/grpo_ultra_v3" + metric_name: "val:total_reward/mean" + higher_is_better: true + keep_top_k: 1000000 + save_period: 6 + ft_keep_latest_k: 1 + ft_save_period: 1 + checkpoint_must_save_by: "00:23:30:00" + model_save_format: "safetensors" + save_consolidated: false + +# ============================================================================= +# GRPO Algorithm +# ============================================================================= +grpo: + num_prompts_per_step: 1024 + num_generations_per_prompt: 1 + num_val_generations_per_prompt: 2 + max_rollout_turns: 1 + max_num_epochs: 1 + max_num_steps: 1000000 + normalize_rewards: true + use_leave_one_out_baseline: true + advantage_clip_low: -20 + advantage_clip_high: 20 + val_period: -1 + val_at_start: false + val_at_end: false + overlong_filtering: false + max_val_samples: null + val_batch_size: 256 + seed: 42 + + use_dynamic_sampling: false + dynamic_sampling_max_gen_batches: 10 + batch_multiplier: 1 + + penalize_invalid_tool_call: false + invalid_tool_call_advantage: -5.0 + + reward_shaping: + enabled: false + overlong_buffer_length: 128 + overlong_buffer_penalty: 1 + max_response_length: ${policy.max_total_sequence_length} + stop_properly_penalty_coef: null + reward_scaling: + enabled: false + source_min: 0.0 + source_max: 1.0 + target_min: 0.0 + target_max: 1.0 + + async_grpo: + enabled: true + max_trajectory_age_steps: 1 + in_flight_weight_updates: true + recompute_kv_cache_after_weight_updates: false + + # OPD: use OPD advantage estimator + adv_estimator: + name: opd + + + use_best_at_k: false + best_at_k_k: 8 + best_at_k_m: 1000 + + use_combined_training: false + combined_training_weight_mode: "auto" + combined_training_best_at_k_weight: 0.2 + combined_training_pass_at_1_weight: 1.0 + + dynamic_sampling_oversample_ratio: 1.0 + seq_logprob_error_threshold: 2 + +# ============================================================================= +# Loss Function +# ============================================================================= +loss_fn: + reference_policy_kl_penalty: 0.0 + reference_policy_kl_type: "k3" + kl_input_clamp_value: null + kl_output_clamp_value: null + + ratio_clip_min: 0.2 + ratio_clip_max: 0.28 + ratio_clip_c: null + use_on_policy_kl_approximation: true + use_importance_sampling_correction: true + truncated_importance_sampling_ratio: 5 + truncated_importance_sampling_ratio_min: 0.2 + truncated_importance_sampling_type: tis + sequence_level_importance_ratios: false + token_level_loss: true + force_on_policy_ratio: true + use_kl_in_reward: false + +# ============================================================================= +# Policy +# ============================================================================= +policy: + model_name: null # set by the launcher via MODEL_PATH + tokenizer: + name: ${policy.model_name} + chat_template_kwargs: null + hf_config_overrides: {} + + train_global_batch_size: 1024 + train_micro_batch_size: 1 + generation_batch_size: 64 + logprob_batch_size: 1 + max_total_sequence_length: 196608 + precision: "bfloat16" + logprob_chunk_size: 2048 + offload_optimizer_for_logprob: false + + dtensor_cfg: + _v2: true + enabled: false + cpu_offload: false + sequence_parallel: false + activation_checkpointing: false + tensor_parallel_size: 1 + context_parallel_size: 1 + custom_parallel_plan: null + + megatron_cfg: + enabled: true + empty_unused_memory_level: 2 + activation_checkpointing: true + + # TP=8 spans 2 GB200 nodes (4 GPUs each) via intra-rack NVLink. + tensor_model_parallel_size: 8 + expert_tensor_parallel_size: 1 + # EP=64: 512 experts / 64 = 8 experts per EP rank + # All-to-all spans 64 GPUs (16 nodes), fits within one NVLink domain + expert_model_parallel_size: 64 + pipeline_model_parallel_size: 1 + num_layers_in_first_pipeline_stage: null + num_layers_in_last_pipeline_stage: null + context_parallel_size: 32 + pipeline_dtype: ${policy.precision} + sequence_parallel: true + + # MoE + freeze_moe_router: true + moe_router_dtype: "fp32" + moe_router_load_balancing_type: "none" + moe_router_bias_update_rate: 1.0e-3 + moe_router_enable_expert_bias: true + moe_permute_fusion: true + moe_enable_deepep: false + moe_token_dispatcher_type: "flex" + moe_flex_dispatcher_backend: "hybridep" + moe_hybridep_num_sms: 32 + moe_aux_loss_coeff: 0.0 + moe_shared_expert_overlap: false + + # Compute + apply_rope_fusion: true + use_fused_weighted_squared_relu: true + bias_activation_fusion: false + defer_fp32_logits: true + + # Logging + track_moe_metrics: true + moe_per_layer_logging: true + do_not_average_loss: true + cp_normalize: true + calculate_per_token_loss: true + scale_loss_by_dp_cp_size: false + + # MTP — disabled + mtp_loss_scaling_factor: 0.3 + mtp_use_repeated_layer: true + mtp_num_layers: 5 + mtp_detach_heads: true + + optimizer: + optimizer: "adam" + lr: 2.0e-6 + min_lr: 2.0e-6 + weight_decay: 0.0 + bf16: true + fp16: false + params_dtype: "float32" + + adam_beta1: 0.9 + adam_beta2: 0.999 + adam_eps: 1e-8 + + sgd_momentum: 0.9 + + use_distributed_optimizer: true + use_precision_aware_optimizer: true + + clip_grad: ${policy.max_grad_norm} + + optimizer_cpu_offload: false + optimizer_offload_fraction: 0.0 + + scheduler: + start_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} + end_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} + weight_decay_incr_style: "constant" + lr_decay_style: "constant" + lr_decay_iters: null + lr_warmup_iters: 10 + lr_warmup_init: 3e-7 + override_opt_param_scheduler: true + + distributed_data_parallel_config: + grad_reduce_in_fp32: false + overlap_grad_reduce: false + overlap_param_gather: true + average_in_collective: false + use_custom_fsdp: false + data_parallel_sharding_strategy: "optim_grads_params" + + # FP8 — disabled for bf16 runs. Enable for mxfp8 validation. + fp8_cfg: + enabled: false + fp8: "e4m3" + fp8_recipe: "mxfp8" + fp8_param: false + + first_last_layers_bf16: true + num_layers_at_start_in_bf16: 1 + num_layers_at_end_in_bf16: 1 + + use_gloo_process_groups: false + + checkpoint: + async_save: true + ckpt_assume_constant_structure: true + fully_parallel_save_process_group: "ep_dp" + fully_parallel_load_process_group: "ep_dp" + fully_parallel_load_exchange_algo: "broadcast" + + env_vars: null + + # --------------------------------------------------------------------------- + # Sequence Packing + # --------------------------------------------------------------------------- + dynamic_batching: + enabled: false + train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}} + sequence_length_round: 64 + + sequence_packing: + enabled: true + train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}} + algorithm: "modified_first_fit_decreasing" + sequence_length_round: 64 + fuse_loss: true + + make_sequence_length_divisible_by: ${policy.megatron_cfg.tensor_model_parallel_size} + max_grad_norm: 1.0 + optimizer: null + scheduler: null + + # --------------------------------------------------------------------------- + # Generation (vLLM) — Non-colocated, async + # --------------------------------------------------------------------------- + generation: + port_range_low: 3000 + port_range_high: 4999 + backend: "vllm" + max_new_tokens: ${policy.max_total_sequence_length} + temperature: 1.0 + top_p: 1.0 + top_k: null + stop_token_ids: null + stop_strings: null + # TP=8 EP=8: EP=TP so vllm_dp_size=1, async_engine=true works with NeMo Gym. + # Same memory as EP=1 but uses all-to-all for expert routing. + # EP > TP blocked by https://github.com/NVIDIA-NeMo/RL/issues/1101. + vllm_cfg: + async_engine: true + precision: ${policy.precision} + kv_cache_dtype: "auto" + tensor_parallel_size: 8 + pipeline_parallel_size: 1 + expert_parallel_size: 8 + gpu_memory_utilization: 0.85 + max_model_len: ${policy.max_total_sequence_length} + enforce_eager: false + use_deep_gemm: false + num_last_layers_in_bf16: 0 + num_first_layers_in_bf16: 0 + enable_vllm_metrics_logger: true + vllm_metrics_logger_interval: 0.5 + expose_http_server: true + skip_tokenizer_init: false + http_server_serving_chat_kwargs: + enable_auto_tools: true + tool_parser: qwen3_coder + reasoning_parser: nano_v3 + reasoning_parser_plugin: nemo_rl/utils/nano_v3_reasoning_parser.py + + vllm_kwargs: + attention_backend: FLASH_ATTN + mamba_ssm_cache_dtype: "float32" + compilation_config: + cudagraph_capture_sizes: [1,2,4,8,16,32,64] + pass_config: + fuse_allreduce_rms: false + + colocated: + enabled: false + resources: + gpus_per_node: 4 + num_nodes: 182 # Overridden by launch script + +# ============================================================================= +# Data +# ============================================================================= +data: + max_input_seq_length: null + shuffle: false + num_workers: 1 + train: + data_path: null # Set by launch script + validation: + data_path: null # Set by launch script + default: + dataset_name: NemoGymDataset + env_name: "nemo_gym" + prompt_file: null + system_prompt_file: null + processor: "nemo_gym_data_processor" + +# ============================================================================= +# Environment — NeMo Gym + Judge Models +# ============================================================================= +env: + should_use_nemo_gym: true + # true: skip expensive train_data_step*.jsonl (recommended for large Gym runs); false: write full jsonl. + should_log_nemo_gym_responses: true + nemo_gym: + nemo_gym_log_dir: "logs/nemo_gym" + skip_venv_if_present: true + port_range_low: 5000 + port_range_high: 5999 + # MOPD trains from teacher logprobs, so skip env verifiers/unit tests during rollout. + skip_verification: false + skip_verification_reward: 0.0 + invalid_tool_call_patterns: + - "" + - "" + - "" + - "" + thinking_tags: + - "" + - "" + config_paths: + - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml + # - responses_api_agents/swe_agents/configs/swebench_openhands_training.yaml + - resources_servers/math_with_judge/configs/math_with_judge.yaml + - resources_servers/code_gen/configs/code_gen.yaml + - resources_servers/workplace_assistant/configs/workplace_assistant.yaml + - resources_servers/mcqa/configs/mcqa.yaml + - resources_servers/instruction_following/configs/instruction_following.yaml + - resources_servers/equivalence_llm_judge/configs/lc_judge.yaml + - resources_servers/calendar/configs/calendar.yaml + - resources_servers/genrm_compare/configs/genrm_compare.yaml + - resources_servers/equivalence_llm_judge/configs/nl2bash-equivalency.yaml + - resources_servers/equivalence_llm_judge/configs/equivalence_llm_judge.yaml + - resources_servers/single_step_tool_use_with_argument_comparison/configs/single_step_tool_use_with_argument_comparison.yaml + - resources_servers/reasoning_gym/configs/reasoning_gym.yaml + - resources_servers/terminus_judge/configs/terminus_judge_string_only.yaml + - resources_servers/ns_tools/configs/ns_tools.yaml + - resources_servers/math_formal_lean/configs/math_formal_lean_multi_turn.yaml + # swerl_gen disabled: requires Apptainer/Singularity (not available on aarch64) + # - resources_servers/swerl_gen/configs/swerl_gen.yaml + - resources_servers/multichallenge/configs/multichallenge.yaml + - resources_servers/inverse_if/configs/inverse_if.yaml + - resources_servers/single_step_tool_use_with_argument_comparison/configs/search_pivot_single_step_tool_use_with_argument_comparison.yaml + - resources_servers/single_step_tool_use_with_argument_comparison/configs/toolcall_schema_single_step_tool_use_with_argument_comparison.yaml + - resources_servers/single_step_tool_use_with_argument_comparison/configs/swe_pivot_single_step_tool_use_with_argument_comparison.yaml + - resources_servers/abstention/configs/abstention.yaml + - resources_servers/nvarc/configs/inductive.yaml + - resources_servers/nvarc/configs/transductive.yaml + - resources_servers/single_step_tool_use_with_argument_comparison/configs/droid_pivot_single_step_tool_use_with_argument_comparison.yaml + - resources_servers/equivalence_rule/configs/lc.yaml + - resources_servers/ether0/configs/ether0.yaml + - resources_servers/structured_outputs/configs/structured_outputs_json_yaml_xml_v1.yaml + - resources_servers/structured_outputs/configs/structured_outputs_v3.yaml + - resources_servers/format_verification/configs/freeform_formatting.yaml + - resources_servers/format_verification/configs/citation_format.yaml + - resources_servers/rdkit_chemistry/configs/rdkit_chemistry.yaml + - resources_servers/jailbreak_detection/configs/jailbreak_detection_nemotron_combined_reward_tp8.yaml + - resources_servers/indirect_prompt_injection/configs/indirect_prompt_injection.yaml + - resources_servers/swe_pivot/configs/swe_pivot.yaml + - resources_servers/terminal_multi_harness/configs/terminal_multi_harness_stirrup.yaml + - resources_servers/terminal_multi_harness/configs/terminal_multi_harness_agent006.yaml + - resources_servers/terminal_multi_harness/configs/terminal_multi_harness_opencode.yaml + - resources_servers/terminal_multi_harness/configs/terminal_multi_harness_codex.yaml + - resources_servers/format_verification/configs/citation_format.yaml + - resources_servers/format_verification/configs/freeform_formatting.yaml + + # Increase the num workers to help with the serialization/deserialization overhead of token IDs + policy_model: + responses_api_models: + vllm_model: + num_workers: 16 + num_groups_nemo_rl: ${add:${grpo.async_grpo.max_trajectory_age_steps}, 1} + # Model-specific (Nemotron chat template): keep reasoning in multi-turn + # history so the trajectory grows monotonically for token-in-token-out RL. + chat_template_kwargs: + enable_thinking: true + truncate_history_thinking: false + sequential_reasoning_allowed: false + policy_model_reasoning_off: + _copy: policy_model + responses_api_models: + vllm_model: + num_workers: 4 + num_groups_nemo_rl: ${add:${grpo.async_grpo.max_trajectory_age_steps}, 1} + chat_template_kwargs: + enable_thinking: false + + abstention: + resources_servers: + abstention: + judge_model_server: + type: responses_api_models + name: nl2bash_judge_model + judge_responses_create_params: + max_output_tokens: 8192 + + # Safety Model: 4B — TP=4 ensures each PG claims a full node, + # avoiding GPU fragmentation that can block larger-TP models. + jailbreak_detection: + resources_servers: + jailbreak_detection: + judge_model_server: + type: responses_api_models + name: safety_judge_model + + safety_judge_model: + responses_api_models: + local_vllm_model: + entrypoint: app.py + model: null # Set by launch script + return_token_id_information: false + uses_reasoning_parser: false + debug: true + vllm_serve_env_vars: + VLLM_RAY_DP_PACK_STRATEGY: strict + + vllm_serve_kwargs: + attention_backend: TRITON_ATTN + tensor_parallel_size: 4 + data_parallel_size: 2 + data_parallel_size_local: 1 + pipeline_parallel_size: 1 + gpu_memory_utilization: 0.85 + max_model_len: 96000 + model_loader_extra_config: + enable_multithread_load: true + num_threads: 112 + compilation_config: + cudagraph_capture_sizes: [1,2,4,8,16] + + + # nl2bash / General Judge: TP=4 on GB200 192GB + nl2bash_judge_model: + responses_api_models: + local_vllm_model: + entrypoint: app.py + model: null # Set by launch script + return_token_id_information: false + uses_reasoning_parser: false + debug: true + vllm_serve_env_vars: + VLLM_RAY_DP_PACK_STRATEGY: strict + + vllm_serve_kwargs: + attention_backend: FLASH_ATTN + tensor_parallel_size: 4 + data_parallel_size: 4 + data_parallel_size_local: 1 + pipeline_parallel_size: 1 + enable_expert_parallel: true + enable_auto_tool_choice: true + tool_call_parser: hermes + gpu_memory_utilization: 0.85 + max_model_len: 131072 + model_loader_extra_config: + enable_multithread_load: true + num_threads: 112 + compilation_config: + cudagraph_capture_sizes: [1,2,4,8,16,32] + server_env: + NCCL_MNNVL_ENABLE: "0" + + inverse_if: + resources_servers: + inverse_if: + judge_model_server: + type: responses_api_models + name: nl2bash_judge_model + + multichallenge: + resources_servers: + multichallenge: + judge_model_server: + type: responses_api_models + name: nl2bash_judge_model + judge_responses_create_params: + max_output_tokens: 8192 + + equivalence_llm_judge: + resources_servers: + equivalence_llm_judge: + judge_model_server: + name: nl2bash_judge_model + judge_responses_create_params: + max_output_tokens: 8192 + + # GenRM: TP=4 on GB200 192GB + genrm_compare_resources_server: + resources_servers: + genrm_compare: + num_rollouts_per_prompt: ${grpo.num_generations_per_prompt} + genrm_model_server: + type: responses_api_models + name: genrm_model + genrm_responses_create_params: + max_output_tokens: 16384 + temperature: 0.6 + top_p: 0.95 + comparison_strategy: "circular" + num_judges_per_comparison: 1 + use_principle: true + default_principle: "You will be given one or more evaluation criteria (rubrics).\nEvaluate both responses on EACH criterion individually first, then synthesize an overall judgment.\nCriteria:\n\n1. Please act as an impartial judge and evaluate the quality of the responses provided by two AI assistants to the user prompt. Begin your evaluation by generating your own answer to the prompt. You must provide your answer before judging any answers. When evaluating the assistants' answers, compare both assistants' answers with your answer. You must identify and correct any mistakes or inaccurate information. Then consider if the assistant's answers are helpful, relevant, and concise. Helpful means the answer correctly responds to the prompt or follows the instructions. Note when user prompt has any ambiguity or more than one interpretation, it is more helpful and appropriate to ask for clarifications or more information from the user than providing an answer based on assumptions. Relevant means all parts of the response closely connect or are appropriate to what is being asked. Concise means the response is clear and not verbose or excessive. Then consider the creativity and novelty of the assistant's answers when needed. Finally, identify any missing important information in the assistants' answers that would be beneficial to include when responding to the user prompt." + aggregator_method: "simple_tiebreaker" + reasoning_bonus: 0.5 + answer_bonus: 0.5 + top_percentile: 0.2 + genrm_parse_retries: 1 + group_reasoning_length_penalty_coeff: 0.1 + group_answer_length_penalty_coeff: 0.2 + group_style_penalty_coeff: 0.0 + default_score: 3.0 + default_ranking: 3.5 + + genrm_model: + responses_api_models: + genrm_model: + entrypoint: app.py + api_key: dummy_key + model: null # Set by launch script + uses_reasoning_parser: true + return_token_id_information: false + debug: true + ray_worker_py_executable: /opt/ray_venvs/nemo_rl.models.generation.vllm.vllm_worker_async.VllmAsyncGenerationWorker/bin/python + vllm_serve_env_vars: + VLLM_RAY_DP_PACK_STRATEGY: strict + + vllm_serve_kwargs: + tensor_parallel_size: 4 + data_parallel_size: 4 + data_parallel_size_local: 1 + pipeline_parallel_size: 1 + reasoning_parser: deepseek_r1 + gpu_memory_utilization: 0.85 + max_model_len: 60000 + enable_prefix_caching: true + enable_chunked_prefill: true + model_loader_extra_config: + enable_multithread_load: true + num_threads: 112 + load_format: instanttensor + compilation_config: + cudagraph_capture_sizes: [1,2,4,8,16,32] + server_env: + NCCL_MNNVL_ENABLE: "0" + + lc_judge: + resources_servers: + equivalence_llm_judge: + judge_model_server: + name: nl2bash_judge_model + judge_responses_create_params: + max_output_tokens: 8192 + + math_with_judge: + resources_servers: + math_with_judge: + judge_model_server: + name: nl2bash_judge_model + judge_responses_create_params: + max_output_tokens: 8192 + should_use_judge: true + + code_gen: + resources_servers: + code_gen: + num_processes: 2048 + unit_test_timeout_secs: 10 + debug: false + + math_formal_lean_refinement_agent: + responses_api_agents: + proof_refinement_agent: + max_correction_turns: 2 + + # swe_agents_train: + # responses_api_agents: + # swe_agents: + # agent_max_turns: 200 + # concurrency: ${mul:2, ${mul:${grpo.num_prompts_per_step}, ${grpo.num_generations_per_prompt}}} + # swebench_agent_timeout: 3600 + # run_with_mixed_prompts: true + # dataset_path: ${data.train_jsonl_fpath} + # apptainer_memory_limit_mb: 65536 + # container_formatter: + # - "${SIF_DIR}/swegym/sweb.eval.arm64.{instance_id}.sif" + # - "${SIF_DIR}/swerebench/{instance_id}.sif" + + # swe_agents_val: + # responses_api_agents: + # swe_agents: + # agent_max_turns: 200 + # concurrency: ${mul:2, ${mul:${grpo.num_prompts_per_step}, ${grpo.num_generations_per_prompt}}} + # swebench_agent_timeout: 3600 + # dataset_path: ${data.validation_jsonl_fpath} + # apptainer_memory_limit_mb: 65536 + # container_formatter: ${env.nemo_gym.swe_agents_train.responses_api_agents.swe_agents.container_formatter} + +# ============================================================================= +# Logger +# ============================================================================= +logger: + log_dir: "logs" + num_val_samples_to_print: 0 + wandb_enabled: false + tensorboard_enabled: false + mlflow_enabled: false + monitor_gpus: true + swanlab_enabled: false + wandb: + project: "grpo-ultra-v3" + name: "grpo-ultra-v3-256n" + tensorboard: {} + mlflow: + experiment_name: "grpo-ultra-v3" + run_name: "grpo-ultra-v3-256n" + gpu_monitoring: + collection_interval: 10 + flush_interval: 10 + +# ============================================================================= +# Effort Levels +# ============================================================================= +effort_levels: + low_string: "{reasoning effort: efficient}" + low_weight: 0.1 + low_penalty: 1 + low_ub: 15000 + +# ============================================================================= +# Token IDs (model-specific, used by token-based penalties) +# ============================================================================= +token_ids: + eos: 2 # + think_open: 12 # + think_close: 13 # + +# ============================================================================= +# Reward Penalties (set reward to 0 when triggered) +# ============================================================================= +penalize_duplicated_reasoning: true # reasoning content == final answer +penalize_empty_final_answer: true # last message output has empty content +penalize_eos_token: true # eos token appears in generation +penalize_malformed_think_tag: true # /<\/think> count != 1 per turn + +# ============================================================================= +# Teacher panel — paths set by the launcher. +# +# `_teachers.general` is required; the other slots fall back to it when unset. +# In the simplified public flow: +# - general = Student RLVR output (also the MOPD student) +# - rlhf = RLHF Teacher output (genrm_* agents) +# - ifbench = IFBench Teacher output (instruction_following, abstention, multichallenge) +# - reasoning = Reasoning Teacher output (math, mcqa, equivalence, ns_tools, code, lc_judge) +# - swe = SWE Teacher output (swe_*, terminal_multi_harness_*, droid_*, structured_outputs_v3, freeform/citation) +# +# Agents below map to one of these five teacher slots. Production also has +# distinct agentic_safety / gdpval / search / terminalbench / tau teachers; +# until those recipes are ready they fall back to `general` here. +# ============================================================================= +_teachers: + general: ??? # Required — set by launcher via NRL_GENERAL_TEACHER_PATH + rlhf: ${_teachers.general} + ifbench: ${_teachers.general} + reasoning: ${_teachers.general} + swe: ${_teachers.general} + +# ============================================================================= +# OPD: on-policy distillation with non-colocated teachers +# ============================================================================= +on_policy_distillation: + enabled: true + use_orm_advantage: false + orm_advantage_weight: 0.0 + teacher_model_by_agent_name: + # Reasoning / math / mcqa / equivalence / lc-judge / ns-tools / code + math_with_judge_simple_agent: ${_teachers.reasoning} + equivalence_llm_judge_simple_agent: ${_teachers.reasoning} + mcqa_simple_agent: ${_teachers.reasoning} + lc_judge_simple_agent: ${_teachers.general} + ns_tools_simple_agent: ${_teachers.reasoning} + code_gen_simple_agent: ${_teachers.reasoning} + + # IFBench / abstention / multichallenge (all share IFBench teacher) + instruction_following_simple_agent: ${_teachers.ifbench} + abstention_simple_agent: ${_teachers.ifbench} + multichallenge_simple_agent: ${_teachers.ifbench} + + # RLHF (GenRM-judged agents) + genrm_simple_agent: ${_teachers.rlhf} + genrm_simple_agent_reasoning_off: ${_teachers.rlhf} + + # SWE / Agentic — production aliases AGENTIC_TEACHER and SWE_TEACHER to + # the same checkpoint; we follow that convention here. + swe_pivot_single_step_tool_use_with_argument_comparison_agent: ${_teachers.swe} + swe_pivot_tool_simulation_agent: ${_teachers.swe} + structured_outputs_v3_simple_agent: ${_teachers.swe} + terminal_multi_harness_opencode_agent: ${_teachers.swe} + terminal_multi_harness_agent006_agent: ${_teachers.swe} + terminal_multi_harness_codex_agent: ${_teachers.swe} + droid_pivot_single_step_tool_use_with_argument_comparison_agent: ${_teachers.swe} + freeform_formatting_simple_agent: ${_teachers.swe} + citation_format_simple_agent: ${_teachers.swe} + + # Falling back to general until dedicated recipes exist + terminal_multi_harness_stirrup_agent: ${_teachers.general} # (prod: gdpval teacher) + terminus_judge_string_only_simple_agent: ${_teachers.general} # (prod: terminalbench teacher) + search_pivot_single_step_tool_use_with_argument_comparison_agent: ${_teachers.general} # (prod: search teacher) + single_step_tool_use_with_argument_comparison_agent: ${_teachers.general} # (prod: tau teacher) + indirect_prompt_injection_simple_agent: ${_teachers.general} # (prod: agentic_safety teacher) + + default_teacher_alias: lc_judge_simple_agent + deduplicate_shared_teacher_checkpoints: true + non_colocated_teachers: + enabled: true + default_teacher_cfg: + tensor_model_parallel_size: 8 + pipeline_model_parallel_size: 1 + expert_model_parallel_size: 16 + context_parallel_size: 2 + num_nodes: 4 + gpus_per_node: 4 + precision: bf16 + micro_batch_size: 1 + megatron_cfg_overrides: + moe_token_dispatcher_type: alltoall + moe_flex_dispatcher_backend: deepep diff --git a/examples/configs/ultra/reasoning_teacher.yaml b/examples/configs/ultra/reasoning_teacher.yaml new file mode 100644 index 0000000000..3b43e0f8c8 --- /dev/null +++ b/examples/configs/ultra/reasoning_teacher.yaml @@ -0,0 +1,500 @@ +# ============================================================================= +# Nemotron 3 Ultra — Reasoning Teacher (non-STEM RLVR, GBS=2048, 65k context) +# ============================================================================= +# General reasoning teacher trained against verifiable rewards on a +# non-STEM RLVR data blend (math, MCQA, equivalence-judging, code, NS-tools). +# The output checkpoint serves the `code_gen`, `ns_tools`, `math_with_judge`, +# `equivalence_llm_judge`, and `mcqa` agent slots in MOPD. +# +# Starts from the Student RLVR output. Smaller batch (GBS=2048) and smaller +# expert-model parallelism (EP=32) than the Student RLVR recipe. +# +# Cluster shape (128 nodes × 4 GPUs on GB200 NVL72 = 512 GPUs): +# Training 64 / vLLM 54 / Gym 10. Override via NUM_TRAIN_NODES, etc. +# Gym is sized to fit local GenRM (TP=4 DP=4 = 16 GPUs), NL2Bash judge +# (TP=4 DP=4 = 16 GPUs), and Safety judge (TP=4 DP=2 = 8 GPUs); 40 GPUs +# total = exactly 10 nodes, no headroom. +# +# Training parallelism: TP=8, EP=32, CP=8, PP=1, SP=true. +# vLLM parallelism: TP=8, EP=8 (EP=TP keeps vllm_dp_size=1; EP>TP is blocked +# by https://github.com/NVIDIA-NeMo/RL/issues/1101). +# ============================================================================= + +# ============================================================================= +# Cluster — read by ultra_launch.sh to size the SLURM allocation +# ============================================================================= +cluster: + gpus_per_node: 4 + num_nodes: 124 + segment_size: 16 + +# ============================================================================= +# Checkpointing — checkpoint_dir is overridden by the launcher +# ============================================================================= +checkpointing: + enabled: true + checkpoint_dir: "results/ultra_student_rlvr" + metric_name: "val:total_reward/mean" + higher_is_better: true + keep_top_k: 1000000 + save_period: 5 + ft_keep_latest_k: 1 + ft_save_period: 1 + checkpoint_must_save_by: "00:47:35:00" + model_save_format: "safetensors" + save_consolidated: false + +# ============================================================================= +# GRPO Algorithm +# ============================================================================= +grpo: + num_prompts_per_step: 128 + num_generations_per_prompt: 16 + num_val_generations_per_prompt: 2 + max_rollout_turns: 1 + max_num_epochs: 10 + max_num_steps: 1000000 + normalize_rewards: true + use_leave_one_out_baseline: true + advantage_clip_low: -20 + advantage_clip_high: 20 + val_period: 10000 + val_at_start: false + val_at_end: false + overlong_filtering: false + max_val_samples: null + val_batch_size: 256 + seed: 42 + + use_dynamic_sampling: false + dynamic_sampling_max_gen_batches: 10 + batch_multiplier: 1 + + penalize_invalid_tool_call: true + invalid_tool_call_advantage: -5.0 + + reward_shaping: + enabled: false + overlong_buffer_length: 128 + overlong_buffer_penalty: 1 + max_response_length: ${policy.max_total_sequence_length} + stop_properly_penalty_coef: null + reward_scaling: + enabled: false + source_min: 0.0 + source_max: 1.0 + target_min: 0.0 + target_max: 1.0 + + async_grpo: + enabled: true + max_trajectory_age_steps: 1 + in_flight_weight_updates: true + recompute_kv_cache_after_weight_updates: false + + use_best_at_k: false + best_at_k_k: 8 + best_at_k_m: 1000 + + use_combined_training: false + combined_training_weight_mode: "auto" + combined_training_best_at_k_weight: 0.2 + combined_training_pass_at_1_weight: 1.0 + + dynamic_sampling_oversample_ratio: 1.0 + seq_logprob_error_threshold: 2 + +# ============================================================================= +# Loss Function +# ============================================================================= +loss_fn: + reference_policy_kl_penalty: 0.0 + reference_policy_kl_type: "k3" + kl_input_clamp_value: null + kl_output_clamp_value: null + + ratio_clip_min: 0.2 + ratio_clip_max: 0.28 + ratio_clip_c: null + use_on_policy_kl_approximation: true + use_importance_sampling_correction: true + truncated_importance_sampling_ratio: 5 + truncated_importance_sampling_ratio_min: 0.2 + truncated_importance_sampling_type: tis + sequence_level_importance_ratios: false + token_level_loss: true + force_on_policy_ratio: true + use_kl_in_reward: false + +# ============================================================================= +# Policy — model_name is overridden by the launcher (MODEL_PATH env var) +# ============================================================================= +policy: + model_name: null + tokenizer: + name: ${policy.model_name} + chat_template_kwargs: null + hf_config_overrides: {} + + train_global_batch_size: 2048 + train_micro_batch_size: 1 + generation_batch_size: 64 + logprob_batch_size: 1 + max_total_sequence_length: 65536 + precision: "bfloat16" + logprob_chunk_size: 2048 + offload_optimizer_for_logprob: false + + dtensor_cfg: + _v2: true + enabled: false + cpu_offload: false + sequence_parallel: false + activation_checkpointing: false + tensor_parallel_size: 1 + context_parallel_size: 1 + custom_parallel_plan: null + + megatron_cfg: + enabled: true + empty_unused_memory_level: 2 + activation_checkpointing: true + + # TP=8 spans 2 GB200 nodes (4 GPUs each) via intra-rack NVLink. + tensor_model_parallel_size: 8 + expert_tensor_parallel_size: 1 + # EP=32: 512 experts / 32 = 16 experts per EP rank + expert_model_parallel_size: 32 + pipeline_model_parallel_size: 1 + num_layers_in_first_pipeline_stage: null + num_layers_in_last_pipeline_stage: null + context_parallel_size: 8 + pipeline_dtype: ${policy.precision} + sequence_parallel: true + + # MoE + freeze_moe_router: true + moe_router_dtype: "fp32" + moe_router_load_balancing_type: "none" + moe_router_bias_update_rate: 1.0e-3 + moe_router_enable_expert_bias: true + moe_permute_fusion: true + moe_enable_deepep: false + moe_token_dispatcher_type: "flex" + moe_flex_dispatcher_backend: "hybridep" + moe_hybridep_num_sms: 32 + moe_aux_loss_coeff: 0.0 + moe_shared_expert_overlap: false + + # Compute + apply_rope_fusion: true + use_fused_weighted_squared_relu: true + bias_activation_fusion: false + defer_fp32_logits: true + + # Logging + track_moe_metrics: true + moe_per_layer_logging: true + do_not_average_loss: true + cp_normalize: true + calculate_per_token_loss: true + scale_loss_by_dp_cp_size: false + + # MTP — disabled by default; enable via ENABLE_MTP_INFERENCE=1 in the launcher + mtp_loss_scaling_factor: 0.3 + mtp_use_repeated_layer: true + mtp_num_layers: 5 + mtp_detach_heads: true + + optimizer: + optimizer: "adam" + lr: 3.0e-6 + min_lr: 3.0e-6 + weight_decay: 0.0 + bf16: true + fp16: false + params_dtype: "float32" + + adam_beta1: 0.9 + adam_beta2: 0.999 + adam_eps: 1e-8 + + sgd_momentum: 0.9 + + use_distributed_optimizer: true + use_precision_aware_optimizer: true + + clip_grad: ${policy.max_grad_norm} + + optimizer_cpu_offload: false + optimizer_offload_fraction: 0.0 + + scheduler: + start_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} + end_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} + weight_decay_incr_style: "constant" + lr_decay_style: "constant" + lr_decay_iters: null + lr_warmup_iters: 10 + lr_warmup_init: 3e-7 + override_opt_param_scheduler: true + + distributed_data_parallel_config: + grad_reduce_in_fp32: false + overlap_grad_reduce: false + overlap_param_gather: true + average_in_collective: false + use_custom_fsdp: false + data_parallel_sharding_strategy: "optim_grads_params" + + fp8_cfg: + enabled: false + fp8: "e4m3" + fp8_recipe: "mxfp8" + fp8_param: false + + first_last_layers_bf16: true + num_layers_at_start_in_bf16: 1 + num_layers_at_end_in_bf16: 1 + + use_gloo_process_groups: false + + checkpoint: + async_save: true + ckpt_assume_constant_structure: true + fully_parallel_save_process_group: "ep_dp" + fully_parallel_load_process_group: "ep_dp" + fully_parallel_load_exchange_algo: "broadcast" + + env_vars: null + + # --------------------------------------------------------------------------- + # Sequence Packing + # --------------------------------------------------------------------------- + dynamic_batching: + enabled: false + train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}} + sequence_length_round: 64 + + sequence_packing: + enabled: true + train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}} + algorithm: "modified_first_fit_decreasing" + sequence_length_round: 64 + fuse_loss: true + + make_sequence_length_divisible_by: ${policy.megatron_cfg.tensor_model_parallel_size} + max_grad_norm: 1.0 + optimizer: null + scheduler: null + + # --------------------------------------------------------------------------- + # Generation (vLLM) — Non-colocated, async + # --------------------------------------------------------------------------- + generation: + port_range_low: 3000 + port_range_high: 4999 + backend: "vllm" + max_new_tokens: ${policy.max_total_sequence_length} + temperature: 1.0 + top_p: 1.0 + top_k: null + stop_token_ids: null + stop_strings: null + # TP=8 EP=8: EP=TP so vllm_dp_size=1, async_engine=true works with NeMo Gym. + # Same memory as EP=1 but uses all-to-all for expert routing. + # EP > TP blocked by https://github.com/NVIDIA-NeMo/RL/issues/1101. + vllm_cfg: + async_engine: true + precision: ${policy.precision} + kv_cache_dtype: "auto" + tensor_parallel_size: 8 + pipeline_parallel_size: 1 + expert_parallel_size: 8 + gpu_memory_utilization: 0.85 + max_model_len: ${policy.max_total_sequence_length} + enforce_eager: false + use_deep_gemm: false + num_last_layers_in_bf16: 0 + num_first_layers_in_bf16: 0 + enable_vllm_metrics_logger: true + vllm_metrics_logger_interval: 0.5 + expose_http_server: true + skip_tokenizer_init: false + http_server_serving_chat_kwargs: + enable_auto_tools: true + tool_parser: qwen3_coder + reasoning_parser: nano_v3 + reasoning_parser_plugin: nemo_rl/utils/nano_v3_reasoning_parser.py + + vllm_kwargs: + attention_backend: FLASH_ATTN + mamba_ssm_cache_dtype: "float32" + compilation_config: + cudagraph_capture_sizes: [1,2,4,8,16,32,64] + pass_config: + fuse_allreduce_rms: false + + colocated: + enabled: false + resources: + gpus_per_node: 4 + num_nodes: 54 # Overridden by launch script via NUM_GEN_NODES + +# ============================================================================= +# Data — data_path is overridden by the launcher (TRAIN_PATH / VAL_PATH) +# ============================================================================= +data: + max_input_seq_length: null + shuffle: false + num_workers: 1 + train: + data_path: null + validation: + data_path: null + default: + dataset_name: NemoGymDataset + env_name: "nemo_gym" + prompt_file: null + system_prompt_file: null + processor: "nemo_gym_data_processor" + +# ============================================================================= +# Environment — NeMo Gym + NL2Bash judge (only judge needed for Reasoning data) +# +# Judge path is set by the launcher: +# - nl2bash_judge_model.model ← NL2BASH_JUDGE_MODEL +# ============================================================================= +env: + should_use_nemo_gym: true + should_log_nemo_gym_responses: true + nemo_gym: + nemo_gym_log_dir: "logs/nemo_gym" + skip_venv_if_present: true + port_range_low: 5000 + port_range_high: 5999 + invalid_tool_call_patterns: + - "" + - "" + - "" + - "" + thinking_tags: + - "" + - "" + # Reasoning data only exercises equivalence_llm_judge_simple_agent. + # Production loads the full ~35-entry list; we trim to keep gym startup + # focused. + config_paths: + - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml + - resources_servers/equivalence_llm_judge/configs/equivalence_llm_judge.yaml + + policy_model: + responses_api_models: + vllm_model: + num_workers: 16 + num_groups_nemo_rl: ${add:${grpo.async_grpo.max_trajectory_age_steps}, 1} + # Model-specific (Nemotron chat template): keep reasoning in multi-turn + # history so the trajectory grows monotonically for token-in-token-out RL. + chat_template_kwargs: + enable_thinking: true + truncate_history_thinking: false + sequential_reasoning_allowed: false + policy_model_reasoning_off: + _copy: policy_model + responses_api_models: + vllm_model: + num_workers: 4 + num_groups_nemo_rl: ${add:${grpo.async_grpo.max_trajectory_age_steps}, 1} + chat_template_kwargs: + enable_thinking: false + + # NL2Bash judge: TP=4 on GB200 192GB + nl2bash_judge_model: + responses_api_models: + local_vllm_model: + entrypoint: app.py + model: null # Set by launcher (NL2BASH_JUDGE_MODEL) + return_token_id_information: false + uses_reasoning_parser: false + debug: true + vllm_serve_env_vars: + VLLM_RAY_DP_PACK_STRATEGY: strict + + vllm_serve_kwargs: + attention_backend: FLASH_ATTN + tensor_parallel_size: 4 + data_parallel_size: 4 + data_parallel_size_local: 1 + pipeline_parallel_size: 1 + enable_expert_parallel: true + enable_auto_tool_choice: true + tool_call_parser: hermes + gpu_memory_utilization: 0.85 + max_model_len: 131072 + enable_prefix_caching: true + enable_chunked_prefill: true + model_loader_extra_config: + enable_multithread_load: true + num_threads: 112 + compilation_config: + cudagraph_capture_sizes: [1,2,4,8,16,32] + server_env: + NCCL_MNNVL_ENABLE: "0" + + equivalence_llm_judge: + resources_servers: + equivalence_llm_judge: + judge_model_server: + name: nl2bash_judge_model + judge_responses_create_params: + max_output_tokens: 8192 + + +# ============================================================================= +# Logger — log_dir, wandb.name, wandb.project are overridden by the launcher +# ============================================================================= +logger: + log_dir: "logs/ultra_student_rlvr" + num_val_samples_to_print: 0 + wandb_enabled: false + tensorboard_enabled: false + mlflow_enabled: false + monitor_gpus: true + swanlab_enabled: false + wandb: + project: "nemotron-3-ultra" + name: "ultra-student-rlvr" + tensorboard: {} + mlflow: + experiment_name: "nemotron-3-ultra" + run_name: "ultra-student-rlvr" + gpu_monitoring: + collection_interval: 10 + flush_interval: 10 + +# ============================================================================= +# Effort Levels +# ============================================================================= +effort_levels: + low_string: "{reasoning effort: efficient}" + low_weight: 0.1 + low_penalty: 1 + low_ub: 15000 + +# ============================================================================= +# Token IDs (model-specific) +# ============================================================================= +token_ids: + eos: 2 + think_open: 12 + think_close: 13 + +# ============================================================================= +# Reward Penalties (set reward to 0 when triggered) +# ============================================================================= +penalize_duplicated_reasoning: true +penalize_empty_final_answer: true +penalize_eos_token: true +penalize_malformed_think_tag: true diff --git a/examples/configs/ultra/rlhf_teacher.yaml b/examples/configs/ultra/rlhf_teacher.yaml new file mode 100644 index 0000000000..97ffb0e521 --- /dev/null +++ b/examples/configs/ultra/rlhf_teacher.yaml @@ -0,0 +1,520 @@ +# ============================================================================= +# Nemotron 3 Ultra — RLHF Teacher +# ============================================================================= +# General-purpose RLHF teacher trained against GenRM comparison signal with +# all length-penalty and reasoning/answer-bonus terms zeroed out — the +# training signal is the pure pairwise GenRM judgment. Starts from the +# Student RLVR output. Same training shape as the IFBench teacher (GBS=2048, +# 49k context, lr=2.5e-6, 32+28+20 cluster). +# +# Cluster shape (80 nodes × 4 GPUs on GB200 NVL72 = 320 GPUs): +# Training 32 / vLLM 28 / Gym 20. Override via NUM_TRAIN_NODES, etc. +# Gym is sized at 20 nodes (= 80 GPUs) so GenRM can serve locally +# (TP=4 DP=4 = 16 GPUs) alongside the NL2Bash and Safety judges. +# +# Training parallelism: TP=8, EP=64, CP=8, PP=1, SP=true. +# vLLM parallelism: TP=8, EP=8 (EP=TP keeps vllm_dp_size=1; EP>TP is blocked +# by https://github.com/NVIDIA-NeMo/RL/issues/1101). +# ============================================================================= + +# ============================================================================= +# Cluster — read by ultra_launch.sh to size the SLURM allocation +# ============================================================================= +cluster: + gpus_per_node: 4 + num_nodes: 60 + segment_size: 16 + +# ============================================================================= +# Checkpointing — checkpoint_dir is overridden by the launcher +# ============================================================================= +checkpointing: + enabled: true + checkpoint_dir: "results/ultra_student_rlvr" + metric_name: "val:total_reward/mean" + higher_is_better: true + keep_top_k: 1000000 + save_period: 10 + ft_keep_latest_k: 1 + ft_save_period: 1 + checkpoint_must_save_by: "00:47:35:00" + model_save_format: "safetensors" + save_consolidated: false + +# ============================================================================= +# GRPO Algorithm +# ============================================================================= +grpo: + num_prompts_per_step: 128 + num_generations_per_prompt: 16 + num_val_generations_per_prompt: 2 + max_rollout_turns: 1 + max_num_epochs: 1 + max_num_steps: 1000000 + normalize_rewards: true + use_leave_one_out_baseline: true + advantage_clip_low: -20 + advantage_clip_high: 20 + val_period: 10000 + val_at_start: false + val_at_end: false + overlong_filtering: false + max_val_samples: null + val_batch_size: 256 + seed: 42 + + use_dynamic_sampling: false + dynamic_sampling_max_gen_batches: 10 + batch_multiplier: 1 + + penalize_invalid_tool_call: true + invalid_tool_call_advantage: -5.0 + + reward_shaping: + enabled: false + overlong_buffer_length: 128 + overlong_buffer_penalty: 1 + max_response_length: ${policy.max_total_sequence_length} + stop_properly_penalty_coef: null + reward_scaling: + enabled: false + source_min: 0.0 + source_max: 1.0 + target_min: 0.0 + target_max: 1.0 + + async_grpo: + enabled: true + max_trajectory_age_steps: 1 + in_flight_weight_updates: true + recompute_kv_cache_after_weight_updates: false + + use_best_at_k: false + best_at_k_k: 8 + best_at_k_m: 1000 + + use_combined_training: false + combined_training_weight_mode: "auto" + combined_training_best_at_k_weight: 0.2 + combined_training_pass_at_1_weight: 1.0 + + dynamic_sampling_oversample_ratio: 1.0 + seq_logprob_error_threshold: 2 + +# ============================================================================= +# Loss Function +# ============================================================================= +loss_fn: + reference_policy_kl_penalty: 0.0 + reference_policy_kl_type: "k3" + kl_input_clamp_value: null + kl_output_clamp_value: null + + ratio_clip_min: 0.2 + ratio_clip_max: 0.28 + ratio_clip_c: null + use_on_policy_kl_approximation: true + use_importance_sampling_correction: true + truncated_importance_sampling_ratio: 5 + truncated_importance_sampling_ratio_min: 0.2 + truncated_importance_sampling_type: tis + sequence_level_importance_ratios: false + token_level_loss: true + force_on_policy_ratio: true + use_kl_in_reward: false + +# ============================================================================= +# Policy — model_name is overridden by the launcher (MODEL_PATH env var) +# ============================================================================= +policy: + model_name: null + tokenizer: + name: ${policy.model_name} + chat_template_kwargs: null + hf_config_overrides: {} + + train_global_batch_size: 2048 + train_micro_batch_size: 1 + generation_batch_size: 64 + logprob_batch_size: 1 + max_total_sequence_length: 49152 + precision: "bfloat16" + logprob_chunk_size: 2048 + offload_optimizer_for_logprob: false + + dtensor_cfg: + _v2: true + enabled: false + cpu_offload: false + sequence_parallel: false + activation_checkpointing: false + tensor_parallel_size: 1 + context_parallel_size: 1 + custom_parallel_plan: null + + megatron_cfg: + enabled: true + empty_unused_memory_level: 2 + activation_checkpointing: true + + # TP=8 spans 2 GB200 nodes (4 GPUs each) via intra-rack NVLink. + tensor_model_parallel_size: 8 + expert_tensor_parallel_size: 1 + # EP=64: 512 experts / 64 = 8 experts per EP rank + # All-to-all spans 64 GPUs (16 nodes), fits within one NVLink domain + expert_model_parallel_size: 64 + pipeline_model_parallel_size: 1 + num_layers_in_first_pipeline_stage: null + num_layers_in_last_pipeline_stage: null + context_parallel_size: 8 + pipeline_dtype: ${policy.precision} + sequence_parallel: true + + # MoE + freeze_moe_router: true + moe_router_dtype: "fp32" + moe_router_load_balancing_type: "none" + moe_router_bias_update_rate: 1.0e-3 + moe_router_enable_expert_bias: true + moe_permute_fusion: true + moe_enable_deepep: false + moe_token_dispatcher_type: "flex" + moe_flex_dispatcher_backend: "hybridep" + moe_hybridep_num_sms: 32 + moe_aux_loss_coeff: 0.0 + moe_shared_expert_overlap: false + + # Compute + apply_rope_fusion: true + use_fused_weighted_squared_relu: true + bias_activation_fusion: false + defer_fp32_logits: true + + # Logging + track_moe_metrics: true + moe_per_layer_logging: true + do_not_average_loss: true + cp_normalize: true + calculate_per_token_loss: true + scale_loss_by_dp_cp_size: false + + # MTP — disabled by default; enable via ENABLE_MTP_INFERENCE=1 in the launcher + mtp_loss_scaling_factor: 0.3 + mtp_use_repeated_layer: true + mtp_num_layers: 5 + mtp_detach_heads: true + + optimizer: + optimizer: "adam" + lr: 2.5e-6 + min_lr: 2.5e-6 + weight_decay: 0.0 + bf16: true + fp16: false + params_dtype: "float32" + + adam_beta1: 0.9 + adam_beta2: 0.999 + adam_eps: 1e-8 + + sgd_momentum: 0.9 + + use_distributed_optimizer: true + use_precision_aware_optimizer: true + + clip_grad: ${policy.max_grad_norm} + + optimizer_cpu_offload: false + optimizer_offload_fraction: 0.0 + + scheduler: + start_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} + end_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} + weight_decay_incr_style: "constant" + lr_decay_style: "constant" + lr_decay_iters: null + lr_warmup_iters: 10 + lr_warmup_init: 4e-7 + override_opt_param_scheduler: true + + distributed_data_parallel_config: + grad_reduce_in_fp32: false + overlap_grad_reduce: false + overlap_param_gather: true + average_in_collective: false + use_custom_fsdp: false + data_parallel_sharding_strategy: "optim_grads_params" + + fp8_cfg: + enabled: false + fp8: "e4m3" + fp8_recipe: "mxfp8" + fp8_param: false + + first_last_layers_bf16: true + num_layers_at_start_in_bf16: 1 + num_layers_at_end_in_bf16: 1 + + use_gloo_process_groups: false + + checkpoint: + async_save: true + ckpt_assume_constant_structure: true + fully_parallel_save_process_group: "ep_dp" + fully_parallel_load_process_group: "ep_dp" + fully_parallel_load_exchange_algo: "broadcast" + + env_vars: null + + # --------------------------------------------------------------------------- + # Sequence Packing + # --------------------------------------------------------------------------- + dynamic_batching: + enabled: false + train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}} + sequence_length_round: 64 + + sequence_packing: + enabled: true + train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}} + algorithm: "modified_first_fit_decreasing" + sequence_length_round: 64 + fuse_loss: true + + make_sequence_length_divisible_by: ${policy.megatron_cfg.tensor_model_parallel_size} + max_grad_norm: 1.0 + optimizer: null + scheduler: null + + # --------------------------------------------------------------------------- + # Generation (vLLM) — Non-colocated, async + # --------------------------------------------------------------------------- + generation: + port_range_low: 3000 + port_range_high: 4999 + backend: "vllm" + max_new_tokens: ${policy.max_total_sequence_length} + temperature: 1.0 + top_p: 1.0 + top_k: null + stop_token_ids: null + stop_strings: null + # TP=8 EP=8: EP=TP so vllm_dp_size=1, async_engine=true works with NeMo Gym. + # Same memory as EP=1 but uses all-to-all for expert routing. + # EP > TP blocked by https://github.com/NVIDIA-NeMo/RL/issues/1101. + vllm_cfg: + async_engine: true + precision: ${policy.precision} + kv_cache_dtype: "auto" + tensor_parallel_size: 8 + pipeline_parallel_size: 1 + expert_parallel_size: 8 + gpu_memory_utilization: 0.85 + max_model_len: ${policy.max_total_sequence_length} + enforce_eager: false + use_deep_gemm: false + num_last_layers_in_bf16: 0 + num_first_layers_in_bf16: 0 + enable_vllm_metrics_logger: true + vllm_metrics_logger_interval: 0.5 + expose_http_server: true + skip_tokenizer_init: false + http_server_serving_chat_kwargs: + enable_auto_tools: true + tool_parser: qwen3_coder + reasoning_parser: nano_v3 + reasoning_parser_plugin: nemo_rl/utils/nano_v3_reasoning_parser.py + + vllm_kwargs: + attention_backend: FLASH_ATTN + mamba_ssm_cache_dtype: "float32" + compilation_config: + cudagraph_capture_sizes: [1,2,4,8,16,32,64] + pass_config: + fuse_allreduce_rms: false + + colocated: + enabled: false + resources: + gpus_per_node: 4 + num_nodes: 28 # Overridden by launch script via NUM_GEN_NODES + +# ============================================================================= +# Data — data_path is overridden by the launcher (TRAIN_PATH / VAL_PATH) +# ============================================================================= +data: + max_input_seq_length: null + shuffle: false + num_workers: 1 + train: + data_path: null + validation: + data_path: null + default: + dataset_name: NemoGymDataset + env_name: "nemo_gym" + prompt_file: null + system_prompt_file: null + processor: "nemo_gym_data_processor" + +# ============================================================================= +# Environment — NeMo Gym + GenRM (only judge needed for RLHF data) +# +# GenRM path/URL is set by the launcher: +# - genrm_model.base_url ← GENRM_BASE_URL (remote service) +# - genrm_model.model ← GENRM_MODEL (local serve) +# ============================================================================= +env: + should_use_nemo_gym: true + should_log_nemo_gym_responses: true + nemo_gym: + nemo_gym_log_dir: "logs/nemo_gym" + skip_venv_if_present: true + port_range_low: 5000 + port_range_high: 5999 + invalid_tool_call_patterns: + - "" + - "" + - "" + - "" + thinking_tags: + - "" + - "" + # RLHF data only uses genrm_simple_agent and genrm_simple_agent_reasoning_off + # → only GenRM compare is needed. Production loads the full ~35-entry + # list across teacher stages; we trim to keep gym startup focused. + config_paths: + - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml + - resources_servers/genrm_compare/configs/genrm_compare.yaml + + policy_model: + responses_api_models: + vllm_model: + num_workers: 16 + num_groups_nemo_rl: ${add:${grpo.async_grpo.max_trajectory_age_steps}, 1} + # Model-specific (Nemotron chat template): keep reasoning in multi-turn + # history so the trajectory grows monotonically for token-in-token-out RL. + chat_template_kwargs: + enable_thinking: true + truncate_history_thinking: false + sequential_reasoning_allowed: false + policy_model_reasoning_off: + _copy: policy_model + responses_api_models: + vllm_model: + num_workers: 4 + num_groups_nemo_rl: ${add:${grpo.async_grpo.max_trajectory_age_steps}, 1} + chat_template_kwargs: + enable_thinking: false + + # GenRM compare: deployed as a separate service; URL set by launcher + genrm_compare_resources_server: + resources_servers: + genrm_compare: + num_rollouts_per_prompt: ${grpo.num_generations_per_prompt} + genrm_model_server: + type: responses_api_models + name: genrm_model + genrm_responses_create_params: + max_output_tokens: 24576 + temperature: 1.0 + top_p: 0.95 + comparison_strategy: "circular" + num_judges_per_comparison: 1 + use_principle: true + default_principle: "You will be given one or more evaluation criteria (rubrics).\nEvaluate both responses on EACH criterion individually first, then synthesize an overall judgment.\nCriteria:\n\n1. Please act as an impartial judge and evaluate the quality of the responses provided by two AI assistants to the user prompt. Begin your evaluation by generating your own answer to the prompt. You must provide your answer before judging any answers. When evaluating the assistants' answers, compare both assistants' answers with your answer. You must identify and correct any mistakes or inaccurate information. Then consider if the assistant's answers are helpful, relevant, and concise. Helpful means the answer correctly responds to the prompt or follows the instructions. Note when user prompt has any ambiguity or more than one interpretation, it is more helpful and appropriate to ask for clarifications or more information from the user than providing an answer based on assumptions. Relevant means all parts of the response closely connect or are appropriate to what is being asked. Concise means the response is clear and not verbose or excessive. Then consider the creativity and novelty of the assistant's answers when needed. Finally, identify any missing important information in the assistants' answers that would be beneficial to include when responding to the user prompt." + aggregator_method: "simple_tiebreaker" + reasoning_bonus: 0.0 + answer_bonus: 0.0 + top_percentile: 0.2 + genrm_parse_retries: 1 + group_reasoning_length_penalty_coeff: 0.0 + use_golden_anchor: false + group_answer_length_penalty_coeff: 0.0 + group_style_penalty_coeff: 0.0 + default_score: 3.0 + default_ranking: 3.5 + + # Launcher sets either `base_url` (route to remote GenRM service) or + # `model` (serve locally using vllm_serve_kwargs below). See the guide. + genrm_model: + responses_api_models: + genrm_model: + entrypoint: app.py + api_key: dummy_key + model: null + uses_reasoning_parser: true + return_token_id_information: false + debug: true + ray_worker_py_executable: /opt/ray_venvs/nemo_rl.models.generation.vllm.vllm_worker_async.VllmAsyncGenerationWorker/bin/python + vllm_serve_env_vars: + VLLM_RAY_DP_PACK_STRATEGY: strict + + vllm_serve_kwargs: + tensor_parallel_size: 4 + data_parallel_size: 4 + data_parallel_size_local: 1 + pipeline_parallel_size: 1 + reasoning_parser: deepseek_r1 + gpu_memory_utilization: 0.85 + max_model_len: 131072 + enable_prefix_caching: true + enable_chunked_prefill: true + model_loader_extra_config: + enable_multithread_load: true + num_threads: 112 + load_format: instanttensor + compilation_config: + cudagraph_capture_sizes: [1,2,4,8,16,32] + server_env: + NCCL_MNNVL_ENABLE: "0" + + +# ============================================================================= +# Logger — log_dir, wandb.name, wandb.project are overridden by the launcher +# ============================================================================= +logger: + log_dir: "logs/ultra_student_rlvr" + num_val_samples_to_print: 0 + wandb_enabled: false + tensorboard_enabled: false + mlflow_enabled: false + monitor_gpus: true + swanlab_enabled: false + wandb: + project: "nemotron-3-ultra" + name: "ultra-student-rlvr" + tensorboard: {} + mlflow: + experiment_name: "nemotron-3-ultra" + run_name: "ultra-student-rlvr" + gpu_monitoring: + collection_interval: 10 + flush_interval: 10 + +# ============================================================================= +# Effort Levels +# ============================================================================= +effort_levels: + low_string: "{reasoning effort: efficient}" + low_weight: 0.1 + low_penalty: 1 + low_ub: 15000 + +# ============================================================================= +# Token IDs (model-specific) +# ============================================================================= +token_ids: + eos: 2 + think_open: 12 + think_close: 13 + +# ============================================================================= +# Reward Penalties (set reward to 0 when triggered) +# ============================================================================= +penalize_duplicated_reasoning: true +penalize_empty_final_answer: true +penalize_eos_token: true +penalize_malformed_think_tag: true diff --git a/examples/configs/ultra/student_rlvr1.yaml b/examples/configs/ultra/student_rlvr1.yaml new file mode 100644 index 0000000000..692bc6066b --- /dev/null +++ b/examples/configs/ultra/student_rlvr1.yaml @@ -0,0 +1,673 @@ +# ============================================================================= +# Nemotron 3 Ultra — Student RLVR (Phase 1: 49k context, ~128 steps) +# ============================================================================= +# GRPO with verifiable rewards for the Ultra student model. Phase 1 runs at +# 49k max context to warm up training; Phase 2 (`student_rlvr2.yaml`) resumes +# the same checkpoint at 65k context for an additional ~50 steps. +# +# Cluster shape (256 nodes × 4 GPUs on GB200 NVL72 = 1024 GPUs): +# Training 64 / vLLM 172 / Gym 20. Override via NUM_TRAIN_NODES, etc. +# +# Training parallelism: TP=8, EP=64, CP=8, PP=1, SP=true. +# vLLM parallelism: TP=8, EP=8 (EP=TP keeps vllm_dp_size=1; EP>TP is blocked +# by https://github.com/NVIDIA-NeMo/RL/issues/1101). +# ============================================================================= + +# ============================================================================= +# Cluster — read by ultra_launch.sh to size the SLURM allocation +# ============================================================================= +cluster: + gpus_per_node: 4 + num_nodes: 256 + segment_size: 16 + +# ============================================================================= +# Checkpointing — checkpoint_dir is overridden by the launcher +# ============================================================================= +checkpointing: + enabled: true + checkpoint_dir: "results/ultra_student_rlvr" + metric_name: "val:total_reward/mean" + higher_is_better: true + keep_top_k: 1000000 + save_period: 8 + ft_keep_latest_k: 1 + ft_save_period: 1 + checkpoint_must_save_by: "00:47:35:00" + model_save_format: "safetensors" + save_consolidated: false + +# ============================================================================= +# GRPO Algorithm +# ============================================================================= +grpo: + num_prompts_per_step: 512 + num_generations_per_prompt: 16 + num_val_generations_per_prompt: 2 + max_rollout_turns: 1 + max_num_epochs: 1 + max_num_steps: 1000000 + normalize_rewards: true + use_leave_one_out_baseline: true + advantage_clip_low: -20 + advantage_clip_high: 20 + val_period: 10000 + val_at_start: false + val_at_end: false + overlong_filtering: false + max_val_samples: null + val_batch_size: 256 + seed: 42 + + use_dynamic_sampling: false + dynamic_sampling_max_gen_batches: 10 + batch_multiplier: 1 + + penalize_invalid_tool_call: true + invalid_tool_call_advantage: -5.0 + + reward_shaping: + enabled: false + overlong_buffer_length: 128 + overlong_buffer_penalty: 1 + max_response_length: ${policy.max_total_sequence_length} + stop_properly_penalty_coef: null + reward_scaling: + enabled: false + source_min: 0.0 + source_max: 1.0 + target_min: 0.0 + target_max: 1.0 + + async_grpo: + enabled: true + max_trajectory_age_steps: 1 + in_flight_weight_updates: true + recompute_kv_cache_after_weight_updates: false + + use_best_at_k: false + best_at_k_k: 8 + best_at_k_m: 1000 + + use_combined_training: false + combined_training_weight_mode: "auto" + combined_training_best_at_k_weight: 0.2 + combined_training_pass_at_1_weight: 1.0 + + dynamic_sampling_oversample_ratio: 1.0 + seq_logprob_error_threshold: 2 + +# ============================================================================= +# Loss Function +# ============================================================================= +loss_fn: + reference_policy_kl_penalty: 0.0 + reference_policy_kl_type: "k3" + kl_input_clamp_value: null + kl_output_clamp_value: null + + ratio_clip_min: 0.2 + ratio_clip_max: 0.28 + ratio_clip_c: null + use_on_policy_kl_approximation: true + use_importance_sampling_correction: true + truncated_importance_sampling_ratio: 5 + truncated_importance_sampling_ratio_min: 0.2 + truncated_importance_sampling_type: tis + sequence_level_importance_ratios: false + token_level_loss: true + force_on_policy_ratio: true + use_kl_in_reward: false + +# ============================================================================= +# Policy — model_name is overridden by the launcher (MODEL_PATH env var) +# ============================================================================= +policy: + model_name: null + tokenizer: + name: ${policy.model_name} + chat_template_kwargs: null + hf_config_overrides: {} + + train_global_batch_size: 8192 + train_micro_batch_size: 1 + generation_batch_size: 64 + logprob_batch_size: 1 + max_total_sequence_length: 49152 + precision: "bfloat16" + logprob_chunk_size: 2048 + offload_optimizer_for_logprob: false + + dtensor_cfg: + _v2: true + enabled: false + cpu_offload: false + sequence_parallel: false + activation_checkpointing: false + tensor_parallel_size: 1 + context_parallel_size: 1 + custom_parallel_plan: null + + megatron_cfg: + enabled: true + empty_unused_memory_level: 2 + activation_checkpointing: true + + # TP=8 spans 2 GB200 nodes (4 GPUs each) via intra-rack NVLink. + tensor_model_parallel_size: 8 + expert_tensor_parallel_size: 1 + # EP=64: 512 experts / 64 = 8 experts per EP rank + # All-to-all spans 64 GPUs (16 nodes), fits within one NVLink domain + expert_model_parallel_size: 64 + pipeline_model_parallel_size: 1 + num_layers_in_first_pipeline_stage: null + num_layers_in_last_pipeline_stage: null + context_parallel_size: 8 + pipeline_dtype: ${policy.precision} + sequence_parallel: true + + # MoE + freeze_moe_router: true + moe_router_dtype: "fp32" + moe_router_load_balancing_type: "none" + moe_router_bias_update_rate: 1.0e-3 + moe_router_enable_expert_bias: true + moe_permute_fusion: true + moe_enable_deepep: false + moe_token_dispatcher_type: "flex" + moe_flex_dispatcher_backend: "hybridep" + moe_hybridep_num_sms: 32 + moe_aux_loss_coeff: 0.0 + moe_shared_expert_overlap: false + + # Compute + apply_rope_fusion: true + use_fused_weighted_squared_relu: true + bias_activation_fusion: false + defer_fp32_logits: true + + # Logging + track_moe_metrics: true + moe_per_layer_logging: true + do_not_average_loss: true + cp_normalize: true + calculate_per_token_loss: true + scale_loss_by_dp_cp_size: false + + # MTP — disabled by default; enable via ENABLE_MTP_INFERENCE=1 in the launcher + mtp_loss_scaling_factor: 0.3 + mtp_use_repeated_layer: true + mtp_num_layers: 5 + mtp_detach_heads: true + + optimizer: + optimizer: "adam" + lr: 4.0e-6 + min_lr: 4.0e-6 + weight_decay: 0.0 + bf16: true + fp16: false + params_dtype: "float32" + + adam_beta1: 0.9 + adam_beta2: 0.999 + adam_eps: 1e-8 + + sgd_momentum: 0.9 + + use_distributed_optimizer: true + use_precision_aware_optimizer: true + + clip_grad: ${policy.max_grad_norm} + + optimizer_cpu_offload: false + optimizer_offload_fraction: 0.0 + + scheduler: + start_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} + end_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} + weight_decay_incr_style: "constant" + lr_decay_style: "constant" + lr_decay_iters: null + lr_warmup_iters: 10 + lr_warmup_init: 4e-7 + override_opt_param_scheduler: true + + distributed_data_parallel_config: + grad_reduce_in_fp32: false + overlap_grad_reduce: false + overlap_param_gather: true + average_in_collective: false + use_custom_fsdp: false + data_parallel_sharding_strategy: "optim_grads_params" + + fp8_cfg: + enabled: false + fp8: "e4m3" + fp8_recipe: "mxfp8" + fp8_param: false + + first_last_layers_bf16: true + num_layers_at_start_in_bf16: 1 + num_layers_at_end_in_bf16: 1 + + use_gloo_process_groups: false + + checkpoint: + async_save: true + ckpt_assume_constant_structure: true + fully_parallel_save_process_group: "ep_dp" + fully_parallel_load_process_group: "ep_dp" + fully_parallel_load_exchange_algo: "broadcast" + + env_vars: null + + # --------------------------------------------------------------------------- + # Sequence Packing + # --------------------------------------------------------------------------- + dynamic_batching: + enabled: false + train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}} + sequence_length_round: 64 + + sequence_packing: + enabled: true + train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}} + algorithm: "modified_first_fit_decreasing" + sequence_length_round: 64 + fuse_loss: true + + make_sequence_length_divisible_by: ${policy.megatron_cfg.tensor_model_parallel_size} + max_grad_norm: 1.0 + optimizer: null + scheduler: null + + # --------------------------------------------------------------------------- + # Generation (vLLM) — Non-colocated, async + # --------------------------------------------------------------------------- + generation: + port_range_low: 3000 + port_range_high: 4999 + backend: "vllm" + max_new_tokens: 49152 + temperature: 1.0 + top_p: 1.0 + top_k: null + stop_token_ids: null + stop_strings: null + # TP=8 EP=8: EP=TP so vllm_dp_size=1, async_engine=true works with NeMo Gym. + # Same memory as EP=1 but uses all-to-all for expert routing. + # EP > TP blocked by https://github.com/NVIDIA-NeMo/RL/issues/1101. + vllm_cfg: + async_engine: true + precision: ${policy.precision} + kv_cache_dtype: "auto" + tensor_parallel_size: 8 + pipeline_parallel_size: 1 + expert_parallel_size: 8 + gpu_memory_utilization: 0.85 + max_model_len: 49152 + enforce_eager: false + use_deep_gemm: false + num_last_layers_in_bf16: 0 + num_first_layers_in_bf16: 0 + enable_vllm_metrics_logger: true + vllm_metrics_logger_interval: 0.5 + expose_http_server: true + skip_tokenizer_init: false + http_server_serving_chat_kwargs: + enable_auto_tools: true + tool_parser: qwen3_coder + reasoning_parser: nano_v3 + reasoning_parser_plugin: nemo_rl/utils/nano_v3_reasoning_parser.py + + vllm_kwargs: + attention_backend: FLASH_ATTN + mamba_ssm_cache_dtype: "float32" + compilation_config: + cudagraph_capture_sizes: [1,2,4,8,16,32,64] + pass_config: + fuse_allreduce_rms: false + + colocated: + enabled: false + resources: + gpus_per_node: 4 + num_nodes: 172 # Overridden by launch script via NUM_GEN_NODES + +# ============================================================================= +# Data — data_path is overridden by the launcher (TRAIN_PATH / VAL_PATH) +# ============================================================================= +data: + max_input_seq_length: null + shuffle: false + num_workers: 1 + train: + data_path: null + validation: + data_path: null + default: + dataset_name: NemoGymDataset + env_name: "nemo_gym" + prompt_file: null + system_prompt_file: null + processor: "nemo_gym_data_processor" + +# ============================================================================= +# Environment — NeMo Gym + Judge Models +# +# Judge model paths and the GenRM service URL are set by the launcher: +# - genrm_model.base_url ← GENRM_BASE_URL +# - nl2bash_judge_model.model ← NL2BASH_JUDGE_MODEL +# - safety_judge_model.model ← SAFETY_JUDGE_MODEL +# ============================================================================= +env: + should_use_nemo_gym: true + should_log_nemo_gym_responses: true + nemo_gym: + nemo_gym_log_dir: "logs/nemo_gym" + skip_venv_if_present: true + port_range_low: 5000 + port_range_high: 5999 + invalid_tool_call_patterns: + - "" + - "" + - "" + - "" + thinking_tags: + - "" + - "" + config_paths: + - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml + - resources_servers/math_with_judge/configs/math_with_judge.yaml + - resources_servers/code_gen/configs/code_gen.yaml + - resources_servers/workplace_assistant/configs/workplace_assistant.yaml + - resources_servers/mcqa/configs/mcqa.yaml + - resources_servers/instruction_following/configs/instruction_following.yaml + - resources_servers/equivalence_llm_judge/configs/lc_judge.yaml + - resources_servers/calendar/configs/calendar.yaml + - resources_servers/genrm_compare/configs/genrm_compare.yaml + - resources_servers/equivalence_llm_judge/configs/nl2bash-equivalency.yaml + - resources_servers/equivalence_llm_judge/configs/equivalence_llm_judge.yaml + - resources_servers/single_step_tool_use_with_argument_comparison/configs/single_step_tool_use_with_argument_comparison.yaml + - resources_servers/reasoning_gym/configs/reasoning_gym.yaml + - resources_servers/terminus_judge/configs/terminus_judge_string_only.yaml + - resources_servers/ns_tools/configs/ns_tools.yaml + - resources_servers/math_formal_lean/configs/math_formal_lean_multi_turn.yaml + - resources_servers/multichallenge/configs/multichallenge.yaml + - resources_servers/inverse_if/configs/inverse_if.yaml + - resources_servers/single_step_tool_use_with_argument_comparison/configs/search_pivot_single_step_tool_use_with_argument_comparison.yaml + - resources_servers/single_step_tool_use_with_argument_comparison/configs/toolcall_schema_single_step_tool_use_with_argument_comparison.yaml + - resources_servers/single_step_tool_use_with_argument_comparison/configs/swe_pivot_single_step_tool_use_with_argument_comparison.yaml + - resources_servers/abstention/configs/abstention.yaml + - resources_servers/nvarc/configs/inductive.yaml + - resources_servers/nvarc/configs/transductive.yaml + - resources_servers/single_step_tool_use_with_argument_comparison/configs/droid_pivot_single_step_tool_use_with_argument_comparison.yaml + - resources_servers/equivalence_rule/configs/lc.yaml + - resources_servers/ether0/configs/ether0.yaml + - resources_servers/structured_outputs/configs/structured_outputs_json_yaml_xml_v1.yaml + - resources_servers/structured_outputs/configs/structured_outputs_v3.yaml + - resources_servers/format_verification/configs/freeform_formatting.yaml + - resources_servers/format_verification/configs/citation_format.yaml + - resources_servers/rdkit_chemistry/configs/rdkit_chemistry.yaml + - resources_servers/jailbreak_detection/configs/jailbreak_detection_nemotron_combined_reward_tp8.yaml + - resources_servers/indirect_prompt_injection/configs/indirect_prompt_injection.yaml + + policy_model: + responses_api_models: + vllm_model: + num_workers: 16 + num_groups_nemo_rl: ${add:${grpo.async_grpo.max_trajectory_age_steps}, 1} + # Model-specific (Nemotron chat template): keep reasoning in multi-turn + # history so the trajectory grows monotonically for token-in-token-out RL. + chat_template_kwargs: + enable_thinking: true + truncate_history_thinking: false + sequential_reasoning_allowed: false + policy_model_reasoning_off: + _copy: policy_model + responses_api_models: + vllm_model: + num_workers: 4 + num_groups_nemo_rl: ${add:${grpo.async_grpo.max_trajectory_age_steps}, 1} + chat_template_kwargs: + enable_thinking: false + + abstention: + resources_servers: + abstention: + judge_model_server: + type: responses_api_models + name: nl2bash_judge_model + judge_responses_create_params: + max_output_tokens: 8192 + + # Safety judge: TP=4 ensures each placement group claims a full GB200 node, + # avoiding GPU fragmentation that can block larger-TP models. + jailbreak_detection: + resources_servers: + jailbreak_detection: + judge_model_server: + type: responses_api_models + name: safety_judge_model + + safety_judge_model: + responses_api_models: + local_vllm_model: + entrypoint: app.py + model: null # Set by launcher (SAFETY_JUDGE_MODEL) + return_token_id_information: false + uses_reasoning_parser: false + debug: true + vllm_serve_env_vars: + VLLM_RAY_DP_PACK_STRATEGY: strict + + vllm_serve_kwargs: + attention_backend: TRITON_ATTN + tensor_parallel_size: 4 + data_parallel_size: 2 + data_parallel_size_local: 1 + pipeline_parallel_size: 1 + gpu_memory_utilization: 0.85 + max_model_len: 96000 + model_loader_extra_config: + enable_multithread_load: true + num_threads: 112 + compilation_config: + cudagraph_capture_sizes: [1,2,4,8,16] + + # NL2Bash / general judge: TP=4 on GB200 192GB + nl2bash_judge_model: + responses_api_models: + local_vllm_model: + entrypoint: app.py + model: null # Set by launcher (NL2BASH_JUDGE_MODEL) + return_token_id_information: false + uses_reasoning_parser: false + debug: true + vllm_serve_env_vars: + VLLM_RAY_DP_PACK_STRATEGY: strict + + vllm_serve_kwargs: + attention_backend: FLASH_ATTN + tensor_parallel_size: 4 + data_parallel_size: 4 + data_parallel_size_local: 1 + pipeline_parallel_size: 1 + enable_expert_parallel: true + enable_auto_tool_choice: true + tool_call_parser: hermes + gpu_memory_utilization: 0.85 + max_model_len: 131072 + enable_prefix_caching: true + enable_chunked_prefill: true + model_loader_extra_config: + enable_multithread_load: true + num_threads: 112 + compilation_config: + cudagraph_capture_sizes: [1,2,4,8,16,32] + server_env: + NCCL_MNNVL_ENABLE: "0" + + inverse_if: + resources_servers: + inverse_if: + judge_model_server: + type: responses_api_models + name: nl2bash_judge_model + + multichallenge: + resources_servers: + multichallenge: + judge_model_server: + type: responses_api_models + name: nl2bash_judge_model + judge_responses_create_params: + max_output_tokens: 8192 + + equivalence_llm_judge: + resources_servers: + equivalence_llm_judge: + judge_model_server: + name: nl2bash_judge_model + judge_responses_create_params: + max_output_tokens: 8192 + + # GenRM compare: deployed as a separate service; URL set by launcher + genrm_compare_resources_server: + resources_servers: + genrm_compare: + num_rollouts_per_prompt: ${grpo.num_generations_per_prompt} + genrm_model_server: + type: responses_api_models + name: genrm_model + genrm_responses_create_params: + max_output_tokens: 24576 + temperature: 1.0 + top_p: 0.95 + comparison_strategy: "circular" + num_judges_per_comparison: 1 + use_principle: true + default_principle: "You will be given one or more evaluation criteria (rubrics).\nEvaluate both responses on EACH criterion individually first, then synthesize an overall judgment.\nCriteria:\n\n1. Please act as an impartial judge and evaluate the quality of the responses provided by two AI assistants to the user prompt. Begin your evaluation by generating your own answer to the prompt. You must provide your answer before judging any answers. When evaluating the assistants' answers, compare both assistants' answers with your answer. You must identify and correct any mistakes or inaccurate information. Then consider if the assistant's answers are helpful, relevant, and concise. Helpful means the answer correctly responds to the prompt or follows the instructions. Note when user prompt has any ambiguity or more than one interpretation, it is more helpful and appropriate to ask for clarifications or more information from the user than providing an answer based on assumptions. Relevant means all parts of the response closely connect or are appropriate to what is being asked. Concise means the response is clear and not verbose or excessive. Then consider the creativity and novelty of the assistant's answers when needed. Finally, identify any missing important information in the assistants' answers that would be beneficial to include when responding to the user prompt." + aggregator_method: "simple_tiebreaker" + reasoning_bonus: 0.5 + answer_bonus: 0.5 + top_percentile: 0.2 + genrm_parse_retries: 1 + group_reasoning_length_penalty_coeff: 0.1 + group_answer_length_penalty_coeff: 0.25 + group_style_penalty_coeff: 0.0 + default_score: 3.0 + default_ranking: 3.5 + + # Launcher sets either `base_url` (route to remote GenRM service) or + # `model` (serve locally using vllm_serve_kwargs below). See the guide. + genrm_model: + responses_api_models: + genrm_model: + entrypoint: app.py + api_key: dummy_key + model: null + uses_reasoning_parser: true + return_token_id_information: false + debug: true + ray_worker_py_executable: /opt/ray_venvs/nemo_rl.models.generation.vllm.vllm_worker_async.VllmAsyncGenerationWorker/bin/python + vllm_serve_env_vars: + VLLM_RAY_DP_PACK_STRATEGY: strict + + vllm_serve_kwargs: + tensor_parallel_size: 4 + data_parallel_size: 4 + data_parallel_size_local: 1 + pipeline_parallel_size: 1 + reasoning_parser: deepseek_r1 + gpu_memory_utilization: 0.85 + max_model_len: 131072 + enable_prefix_caching: true + enable_chunked_prefill: true + model_loader_extra_config: + enable_multithread_load: true + num_threads: 112 + load_format: instanttensor + compilation_config: + cudagraph_capture_sizes: [1,2,4,8,16,32] + server_env: + NCCL_MNNVL_ENABLE: "0" + + lc_judge: + resources_servers: + equivalence_llm_judge: + judge_model_server: + name: nl2bash_judge_model + judge_responses_create_params: + max_output_tokens: 8192 + + math_with_judge: + resources_servers: + math_with_judge: + judge_model_server: + name: nl2bash_judge_model + judge_responses_create_params: + max_output_tokens: 8192 + should_use_judge: true + + code_gen: + resources_servers: + code_gen: + num_processes: 2048 + unit_test_timeout_secs: 10 + debug: false + + math_formal_lean_refinement_agent: + responses_api_agents: + proof_refinement_agent: + max_correction_turns: 2 + +# ============================================================================= +# Logger — log_dir, wandb.name, wandb.project are overridden by the launcher +# ============================================================================= +logger: + log_dir: "logs/ultra_student_rlvr" + num_val_samples_to_print: 0 + wandb_enabled: false + tensorboard_enabled: false + mlflow_enabled: false + monitor_gpus: true + swanlab_enabled: false + wandb: + project: "nemotron-3-ultra" + name: "ultra-student-rlvr" + tensorboard: {} + mlflow: + experiment_name: "nemotron-3-ultra" + run_name: "ultra-student-rlvr" + gpu_monitoring: + collection_interval: 10 + flush_interval: 10 + +# ============================================================================= +# Effort Levels +# ============================================================================= +effort_levels: + low_string: "{reasoning effort: efficient}" + low_weight: 0.1 + low_penalty: 1 + low_ub: 15000 + +# ============================================================================= +# Token IDs (model-specific) +# ============================================================================= +token_ids: + eos: 2 + think_open: 12 + think_close: 13 + +# ============================================================================= +# Reward Penalties (set reward to 0 when triggered) +# ============================================================================= +penalize_duplicated_reasoning: true +penalize_empty_final_answer: true +penalize_eos_token: true +penalize_malformed_think_tag: true diff --git a/examples/configs/ultra/student_rlvr2.yaml b/examples/configs/ultra/student_rlvr2.yaml new file mode 100644 index 0000000000..2f46aa402c --- /dev/null +++ b/examples/configs/ultra/student_rlvr2.yaml @@ -0,0 +1,674 @@ +# ============================================================================= +# Nemotron 3 Ultra — Student RLVR (Phase 2: 65k context, ~50 steps) +# ============================================================================= +# Phase 2 of Student RLVR. Resumes from Phase 1's checkpoint +# (`student_rlvr1.yaml`, ~128 steps at 49k context) and continues for ~50 +# more steps at 65k context. Run with the same EXP_NAME / RESULTS_DIR as +# Phase 1 so `CheckpointManager` auto-resumes. +# +# Cluster shape (256 nodes × 4 GPUs on GB200 NVL72 = 1024 GPUs): +# Training 64 / vLLM 172 / Gym 20. Override via NUM_TRAIN_NODES, etc. +# +# Training parallelism: TP=8, EP=64, CP=8, PP=1, SP=true. +# vLLM parallelism: TP=8, EP=8 (EP=TP keeps vllm_dp_size=1; EP>TP is blocked +# by https://github.com/NVIDIA-NeMo/RL/issues/1101). +# ============================================================================= + +# ============================================================================= +# Cluster — read by ultra_launch.sh to size the SLURM allocation +# ============================================================================= +cluster: + gpus_per_node: 4 + num_nodes: 256 + segment_size: 16 + +# ============================================================================= +# Checkpointing — checkpoint_dir is overridden by the launcher +# ============================================================================= +checkpointing: + enabled: true + checkpoint_dir: "results/ultra_student_rlvr" + metric_name: "val:total_reward/mean" + higher_is_better: true + keep_top_k: 1000000 + save_period: 8 + ft_keep_latest_k: 1 + ft_save_period: 1 + checkpoint_must_save_by: "00:47:35:00" + model_save_format: "safetensors" + save_consolidated: false + +# ============================================================================= +# GRPO Algorithm +# ============================================================================= +grpo: + num_prompts_per_step: 512 + num_generations_per_prompt: 16 + num_val_generations_per_prompt: 2 + max_rollout_turns: 1 + max_num_epochs: 1 + max_num_steps: 1000000 + normalize_rewards: true + use_leave_one_out_baseline: true + advantage_clip_low: -20 + advantage_clip_high: 20 + val_period: 10000 + val_at_start: false + val_at_end: false + overlong_filtering: false + max_val_samples: null + val_batch_size: 256 + seed: 42 + + use_dynamic_sampling: false + dynamic_sampling_max_gen_batches: 10 + batch_multiplier: 1 + + penalize_invalid_tool_call: true + invalid_tool_call_advantage: -5.0 + + reward_shaping: + enabled: false + overlong_buffer_length: 128 + overlong_buffer_penalty: 1 + max_response_length: ${policy.max_total_sequence_length} + stop_properly_penalty_coef: null + reward_scaling: + enabled: false + source_min: 0.0 + source_max: 1.0 + target_min: 0.0 + target_max: 1.0 + + async_grpo: + enabled: true + max_trajectory_age_steps: 1 + in_flight_weight_updates: true + recompute_kv_cache_after_weight_updates: false + + use_best_at_k: false + best_at_k_k: 8 + best_at_k_m: 1000 + + use_combined_training: false + combined_training_weight_mode: "auto" + combined_training_best_at_k_weight: 0.2 + combined_training_pass_at_1_weight: 1.0 + + dynamic_sampling_oversample_ratio: 1.0 + seq_logprob_error_threshold: 2 + +# ============================================================================= +# Loss Function +# ============================================================================= +loss_fn: + reference_policy_kl_penalty: 0.0 + reference_policy_kl_type: "k3" + kl_input_clamp_value: null + kl_output_clamp_value: null + + ratio_clip_min: 0.2 + ratio_clip_max: 0.28 + ratio_clip_c: null + use_on_policy_kl_approximation: true + use_importance_sampling_correction: true + truncated_importance_sampling_ratio: 5 + truncated_importance_sampling_ratio_min: 0.2 + truncated_importance_sampling_type: tis + sequence_level_importance_ratios: false + token_level_loss: true + force_on_policy_ratio: true + use_kl_in_reward: false + +# ============================================================================= +# Policy — model_name is overridden by the launcher (MODEL_PATH env var) +# ============================================================================= +policy: + model_name: null + tokenizer: + name: ${policy.model_name} + chat_template_kwargs: null + hf_config_overrides: {} + + train_global_batch_size: 8192 + train_micro_batch_size: 1 + generation_batch_size: 64 + logprob_batch_size: 1 + max_total_sequence_length: 65536 + precision: "bfloat16" + logprob_chunk_size: 2048 + offload_optimizer_for_logprob: false + + dtensor_cfg: + _v2: true + enabled: false + cpu_offload: false + sequence_parallel: false + activation_checkpointing: false + tensor_parallel_size: 1 + context_parallel_size: 1 + custom_parallel_plan: null + + megatron_cfg: + enabled: true + empty_unused_memory_level: 2 + activation_checkpointing: true + + # TP=8 spans 2 GB200 nodes (4 GPUs each) via intra-rack NVLink. + tensor_model_parallel_size: 8 + expert_tensor_parallel_size: 1 + # EP=64: 512 experts / 64 = 8 experts per EP rank + # All-to-all spans 64 GPUs (16 nodes), fits within one NVLink domain + expert_model_parallel_size: 64 + pipeline_model_parallel_size: 1 + num_layers_in_first_pipeline_stage: null + num_layers_in_last_pipeline_stage: null + context_parallel_size: 8 + pipeline_dtype: ${policy.precision} + sequence_parallel: true + + # MoE + freeze_moe_router: true + moe_router_dtype: "fp32" + moe_router_load_balancing_type: "none" + moe_router_bias_update_rate: 1.0e-3 + moe_router_enable_expert_bias: true + moe_permute_fusion: true + moe_enable_deepep: false + moe_token_dispatcher_type: "flex" + moe_flex_dispatcher_backend: "hybridep" + moe_hybridep_num_sms: 32 + moe_aux_loss_coeff: 0.0 + moe_shared_expert_overlap: false + + # Compute + apply_rope_fusion: true + use_fused_weighted_squared_relu: true + bias_activation_fusion: false + defer_fp32_logits: true + + # Logging + track_moe_metrics: true + moe_per_layer_logging: true + do_not_average_loss: true + cp_normalize: true + calculate_per_token_loss: true + scale_loss_by_dp_cp_size: false + + # MTP — disabled by default; enable via ENABLE_MTP_INFERENCE=1 in the launcher + mtp_loss_scaling_factor: 0.3 + mtp_use_repeated_layer: true + mtp_num_layers: 5 + mtp_detach_heads: true + + optimizer: + optimizer: "adam" + lr: 4.0e-6 + min_lr: 4.0e-6 + weight_decay: 0.0 + bf16: true + fp16: false + params_dtype: "float32" + + adam_beta1: 0.9 + adam_beta2: 0.999 + adam_eps: 1e-8 + + sgd_momentum: 0.9 + + use_distributed_optimizer: true + use_precision_aware_optimizer: true + + clip_grad: ${policy.max_grad_norm} + + optimizer_cpu_offload: false + optimizer_offload_fraction: 0.0 + + scheduler: + start_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} + end_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} + weight_decay_incr_style: "constant" + lr_decay_style: "constant" + lr_decay_iters: null + lr_warmup_iters: 10 + lr_warmup_init: 4e-7 + override_opt_param_scheduler: true + + distributed_data_parallel_config: + grad_reduce_in_fp32: false + overlap_grad_reduce: false + overlap_param_gather: true + average_in_collective: false + use_custom_fsdp: false + data_parallel_sharding_strategy: "optim_grads_params" + + fp8_cfg: + enabled: false + fp8: "e4m3" + fp8_recipe: "mxfp8" + fp8_param: false + + first_last_layers_bf16: true + num_layers_at_start_in_bf16: 1 + num_layers_at_end_in_bf16: 1 + + use_gloo_process_groups: false + + checkpoint: + async_save: true + ckpt_assume_constant_structure: true + fully_parallel_save_process_group: "ep_dp" + fully_parallel_load_process_group: "ep_dp" + fully_parallel_load_exchange_algo: "broadcast" + + env_vars: null + + # --------------------------------------------------------------------------- + # Sequence Packing + # --------------------------------------------------------------------------- + dynamic_batching: + enabled: false + train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}} + sequence_length_round: 64 + + sequence_packing: + enabled: true + train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}} + algorithm: "modified_first_fit_decreasing" + sequence_length_round: 64 + fuse_loss: true + + make_sequence_length_divisible_by: ${policy.megatron_cfg.tensor_model_parallel_size} + max_grad_norm: 1.0 + optimizer: null + scheduler: null + + # --------------------------------------------------------------------------- + # Generation (vLLM) — Non-colocated, async + # --------------------------------------------------------------------------- + generation: + port_range_low: 3000 + port_range_high: 4999 + backend: "vllm" + max_new_tokens: 65536 + temperature: 1.0 + top_p: 1.0 + top_k: null + stop_token_ids: null + stop_strings: null + # TP=8 EP=8: EP=TP so vllm_dp_size=1, async_engine=true works with NeMo Gym. + # Same memory as EP=1 but uses all-to-all for expert routing. + # EP > TP blocked by https://github.com/NVIDIA-NeMo/RL/issues/1101. + vllm_cfg: + async_engine: true + precision: ${policy.precision} + kv_cache_dtype: "auto" + tensor_parallel_size: 8 + pipeline_parallel_size: 1 + expert_parallel_size: 8 + gpu_memory_utilization: 0.85 + max_model_len: 65536 + enforce_eager: false + use_deep_gemm: false + num_last_layers_in_bf16: 0 + num_first_layers_in_bf16: 0 + enable_vllm_metrics_logger: true + vllm_metrics_logger_interval: 0.5 + expose_http_server: true + skip_tokenizer_init: false + http_server_serving_chat_kwargs: + enable_auto_tools: true + tool_parser: qwen3_coder + reasoning_parser: nano_v3 + reasoning_parser_plugin: nemo_rl/utils/nano_v3_reasoning_parser.py + + vllm_kwargs: + attention_backend: FLASH_ATTN + mamba_ssm_cache_dtype: "float32" + compilation_config: + cudagraph_capture_sizes: [1,2,4,8,16,32,64] + pass_config: + fuse_allreduce_rms: false + + colocated: + enabled: false + resources: + gpus_per_node: 4 + num_nodes: 172 # Overridden by launch script via NUM_GEN_NODES + +# ============================================================================= +# Data — data_path is overridden by the launcher (TRAIN_PATH / VAL_PATH) +# ============================================================================= +data: + max_input_seq_length: null + shuffle: false + num_workers: 1 + train: + data_path: null + validation: + data_path: null + default: + dataset_name: NemoGymDataset + env_name: "nemo_gym" + prompt_file: null + system_prompt_file: null + processor: "nemo_gym_data_processor" + +# ============================================================================= +# Environment — NeMo Gym + Judge Models +# +# Judge model paths and the GenRM service URL are set by the launcher: +# - genrm_model.base_url ← GENRM_BASE_URL +# - nl2bash_judge_model.model ← NL2BASH_JUDGE_MODEL +# - safety_judge_model.model ← SAFETY_JUDGE_MODEL +# ============================================================================= +env: + should_use_nemo_gym: true + should_log_nemo_gym_responses: true + nemo_gym: + nemo_gym_log_dir: "logs/nemo_gym" + skip_venv_if_present: true + port_range_low: 5000 + port_range_high: 5999 + invalid_tool_call_patterns: + - "" + - "" + - "" + - "" + thinking_tags: + - "" + - "" + config_paths: + - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml + - resources_servers/math_with_judge/configs/math_with_judge.yaml + - resources_servers/code_gen/configs/code_gen.yaml + - resources_servers/workplace_assistant/configs/workplace_assistant.yaml + - resources_servers/mcqa/configs/mcqa.yaml + - resources_servers/instruction_following/configs/instruction_following.yaml + - resources_servers/equivalence_llm_judge/configs/lc_judge.yaml + - resources_servers/calendar/configs/calendar.yaml + - resources_servers/genrm_compare/configs/genrm_compare.yaml + - resources_servers/equivalence_llm_judge/configs/nl2bash-equivalency.yaml + - resources_servers/equivalence_llm_judge/configs/equivalence_llm_judge.yaml + - resources_servers/single_step_tool_use_with_argument_comparison/configs/single_step_tool_use_with_argument_comparison.yaml + - resources_servers/reasoning_gym/configs/reasoning_gym.yaml + - resources_servers/terminus_judge/configs/terminus_judge_string_only.yaml + - resources_servers/ns_tools/configs/ns_tools.yaml + - resources_servers/math_formal_lean/configs/math_formal_lean_multi_turn.yaml + - resources_servers/multichallenge/configs/multichallenge.yaml + - resources_servers/inverse_if/configs/inverse_if.yaml + - resources_servers/single_step_tool_use_with_argument_comparison/configs/search_pivot_single_step_tool_use_with_argument_comparison.yaml + - resources_servers/single_step_tool_use_with_argument_comparison/configs/toolcall_schema_single_step_tool_use_with_argument_comparison.yaml + - resources_servers/single_step_tool_use_with_argument_comparison/configs/swe_pivot_single_step_tool_use_with_argument_comparison.yaml + - resources_servers/abstention/configs/abstention.yaml + - resources_servers/nvarc/configs/inductive.yaml + - resources_servers/nvarc/configs/transductive.yaml + - resources_servers/single_step_tool_use_with_argument_comparison/configs/droid_pivot_single_step_tool_use_with_argument_comparison.yaml + - resources_servers/equivalence_rule/configs/lc.yaml + - resources_servers/ether0/configs/ether0.yaml + - resources_servers/structured_outputs/configs/structured_outputs_json_yaml_xml_v1.yaml + - resources_servers/structured_outputs/configs/structured_outputs_v3.yaml + - resources_servers/format_verification/configs/freeform_formatting.yaml + - resources_servers/format_verification/configs/citation_format.yaml + - resources_servers/rdkit_chemistry/configs/rdkit_chemistry.yaml + - resources_servers/jailbreak_detection/configs/jailbreak_detection_nemotron_combined_reward_tp8.yaml + - resources_servers/indirect_prompt_injection/configs/indirect_prompt_injection.yaml + + policy_model: + responses_api_models: + vllm_model: + num_workers: 16 + num_groups_nemo_rl: ${add:${grpo.async_grpo.max_trajectory_age_steps}, 1} + # Model-specific (Nemotron chat template): keep reasoning in multi-turn + # history so the trajectory grows monotonically for token-in-token-out RL. + chat_template_kwargs: + enable_thinking: true + truncate_history_thinking: false + sequential_reasoning_allowed: false + policy_model_reasoning_off: + _copy: policy_model + responses_api_models: + vllm_model: + num_workers: 4 + num_groups_nemo_rl: ${add:${grpo.async_grpo.max_trajectory_age_steps}, 1} + chat_template_kwargs: + enable_thinking: false + + abstention: + resources_servers: + abstention: + judge_model_server: + type: responses_api_models + name: nl2bash_judge_model + judge_responses_create_params: + max_output_tokens: 8192 + + # Safety judge: TP=4 ensures each placement group claims a full GB200 node, + # avoiding GPU fragmentation that can block larger-TP models. + jailbreak_detection: + resources_servers: + jailbreak_detection: + judge_model_server: + type: responses_api_models + name: safety_judge_model + + safety_judge_model: + responses_api_models: + local_vllm_model: + entrypoint: app.py + model: null # Set by launcher (SAFETY_JUDGE_MODEL) + return_token_id_information: false + uses_reasoning_parser: false + debug: true + vllm_serve_env_vars: + VLLM_RAY_DP_PACK_STRATEGY: strict + + vllm_serve_kwargs: + attention_backend: TRITON_ATTN + tensor_parallel_size: 4 + data_parallel_size: 2 + data_parallel_size_local: 1 + pipeline_parallel_size: 1 + gpu_memory_utilization: 0.85 + max_model_len: 96000 + model_loader_extra_config: + enable_multithread_load: true + num_threads: 112 + compilation_config: + cudagraph_capture_sizes: [1,2,4,8,16] + + # NL2Bash / general judge: TP=4 on GB200 192GB + nl2bash_judge_model: + responses_api_models: + local_vllm_model: + entrypoint: app.py + model: null # Set by launcher (NL2BASH_JUDGE_MODEL) + return_token_id_information: false + uses_reasoning_parser: false + debug: true + vllm_serve_env_vars: + VLLM_RAY_DP_PACK_STRATEGY: strict + + vllm_serve_kwargs: + attention_backend: FLASH_ATTN + tensor_parallel_size: 4 + data_parallel_size: 4 + data_parallel_size_local: 1 + pipeline_parallel_size: 1 + enable_expert_parallel: true + enable_auto_tool_choice: true + tool_call_parser: hermes + gpu_memory_utilization: 0.85 + max_model_len: 131072 + enable_prefix_caching: true + enable_chunked_prefill: true + model_loader_extra_config: + enable_multithread_load: true + num_threads: 112 + compilation_config: + cudagraph_capture_sizes: [1,2,4,8,16,32] + server_env: + NCCL_MNNVL_ENABLE: "0" + + inverse_if: + resources_servers: + inverse_if: + judge_model_server: + type: responses_api_models + name: nl2bash_judge_model + + multichallenge: + resources_servers: + multichallenge: + judge_model_server: + type: responses_api_models + name: nl2bash_judge_model + judge_responses_create_params: + max_output_tokens: 8192 + + equivalence_llm_judge: + resources_servers: + equivalence_llm_judge: + judge_model_server: + name: nl2bash_judge_model + judge_responses_create_params: + max_output_tokens: 8192 + + # GenRM compare: deployed as a separate service; URL set by launcher + genrm_compare_resources_server: + resources_servers: + genrm_compare: + num_rollouts_per_prompt: ${grpo.num_generations_per_prompt} + genrm_model_server: + type: responses_api_models + name: genrm_model + genrm_responses_create_params: + max_output_tokens: 24576 + temperature: 1.0 + top_p: 0.95 + comparison_strategy: "circular" + num_judges_per_comparison: 1 + use_principle: true + default_principle: "You will be given one or more evaluation criteria (rubrics).\nEvaluate both responses on EACH criterion individually first, then synthesize an overall judgment.\nCriteria:\n\n1. Please act as an impartial judge and evaluate the quality of the responses provided by two AI assistants to the user prompt. Begin your evaluation by generating your own answer to the prompt. You must provide your answer before judging any answers. When evaluating the assistants' answers, compare both assistants' answers with your answer. You must identify and correct any mistakes or inaccurate information. Then consider if the assistant's answers are helpful, relevant, and concise. Helpful means the answer correctly responds to the prompt or follows the instructions. Note when user prompt has any ambiguity or more than one interpretation, it is more helpful and appropriate to ask for clarifications or more information from the user than providing an answer based on assumptions. Relevant means all parts of the response closely connect or are appropriate to what is being asked. Concise means the response is clear and not verbose or excessive. Then consider the creativity and novelty of the assistant's answers when needed. Finally, identify any missing important information in the assistants' answers that would be beneficial to include when responding to the user prompt." + aggregator_method: "simple_tiebreaker" + reasoning_bonus: 0.5 + answer_bonus: 0.5 + top_percentile: 0.2 + genrm_parse_retries: 1 + group_reasoning_length_penalty_coeff: 0.1 + group_answer_length_penalty_coeff: 0.08 + group_style_penalty_coeff: 0.0 + default_score: 3.0 + default_ranking: 3.5 + + # Launcher sets either `base_url` (route to remote GenRM service) or + # `model` (serve locally using vllm_serve_kwargs below). See the guide. + genrm_model: + responses_api_models: + genrm_model: + entrypoint: app.py + api_key: dummy_key + model: null + uses_reasoning_parser: true + return_token_id_information: false + debug: true + ray_worker_py_executable: /opt/ray_venvs/nemo_rl.models.generation.vllm.vllm_worker_async.VllmAsyncGenerationWorker/bin/python + vllm_serve_env_vars: + VLLM_RAY_DP_PACK_STRATEGY: strict + + vllm_serve_kwargs: + tensor_parallel_size: 4 + data_parallel_size: 4 + data_parallel_size_local: 1 + pipeline_parallel_size: 1 + reasoning_parser: deepseek_r1 + gpu_memory_utilization: 0.85 + max_model_len: 131072 + enable_prefix_caching: true + enable_chunked_prefill: true + model_loader_extra_config: + enable_multithread_load: true + num_threads: 112 + load_format: instanttensor + compilation_config: + cudagraph_capture_sizes: [1,2,4,8,16,32] + server_env: + NCCL_MNNVL_ENABLE: "0" + + lc_judge: + resources_servers: + equivalence_llm_judge: + judge_model_server: + name: nl2bash_judge_model + judge_responses_create_params: + max_output_tokens: 8192 + + math_with_judge: + resources_servers: + math_with_judge: + judge_model_server: + name: nl2bash_judge_model + judge_responses_create_params: + max_output_tokens: 8192 + should_use_judge: true + + code_gen: + resources_servers: + code_gen: + num_processes: 2048 + unit_test_timeout_secs: 10 + debug: false + + math_formal_lean_refinement_agent: + responses_api_agents: + proof_refinement_agent: + max_correction_turns: 2 + +# ============================================================================= +# Logger — log_dir, wandb.name, wandb.project are overridden by the launcher +# ============================================================================= +logger: + log_dir: "logs/ultra_student_rlvr" + num_val_samples_to_print: 0 + wandb_enabled: false + tensorboard_enabled: false + mlflow_enabled: false + monitor_gpus: true + swanlab_enabled: false + wandb: + project: "nemotron-3-ultra" + name: "ultra-student-rlvr" + tensorboard: {} + mlflow: + experiment_name: "nemotron-3-ultra" + run_name: "ultra-student-rlvr" + gpu_monitoring: + collection_interval: 10 + flush_interval: 10 + +# ============================================================================= +# Effort Levels +# ============================================================================= +effort_levels: + low_string: "{reasoning effort: efficient}" + low_weight: 0.1 + low_penalty: 1 + low_ub: 15000 + +# ============================================================================= +# Token IDs (model-specific) +# ============================================================================= +token_ids: + eos: 2 + think_open: 12 + think_close: 13 + +# ============================================================================= +# Reward Penalties (set reward to 0 when triggered) +# ============================================================================= +penalize_duplicated_reasoning: true +penalize_empty_final_answer: true +penalize_eos_token: true +penalize_malformed_think_tag: true diff --git a/examples/configs/ultra/swe_teacher.yaml b/examples/configs/ultra/swe_teacher.yaml new file mode 100644 index 0000000000..69b8620796 --- /dev/null +++ b/examples/configs/ultra/swe_teacher.yaml @@ -0,0 +1,476 @@ +# ============================================================================= +# Nemotron 3 Ultra — SWE Teacher (code-execution RLVR, GBS=512, 192k context) +# ============================================================================= +# RLVR teacher for software-engineering benchmarks. Trains the policy against +# multi-turn coding agents (`swe_agents`) that execute candidate fixes inside +# apptainer (.sif) container images for SWE-Bench, SWE-Gym, etc., +# and reward based on test pass/fail — no LLM judges in this stage. +# +# Trained at 192k context with a +# small batch (GBS=512, PPS=32, GPP=16), large CP=32, and EP=32. +# +# Cluster shape (128 nodes × 4 GPUs on GB200 NVL72 = 512 GPUs): +# Training 64 / vLLM 64 / Gym 0 (no LLM judges). Override via +# NUM_TRAIN_NODES, NUM_GEN_NODES, NUM_GYM_NODES. +# +# SIF images: the SWE agent's `container_formatter` paths use the top-level +# `sif_dir` Hydra key — set via `SIF_DIR=...` in the launcher. The released +# blend uses two benchmarks, so the directory needs `swerebench/` and `swegym/` +# subdirectories with per-instance `{instance_id}.sif` files. See the public +# guide for how to build these for ARM64 (GB200). +# ============================================================================= +# +# - gpus_per_node: 4 +# - TP: 8, CP: 8, EP: 32, PP: 1 +# - vLLM TP: 8 +# - Non-colocated async inference with 64 generation nodes +# - Max sequence length: 131072 +# ============================================================================= + +# ============================================================================= +# Cluster — overridden by launch script for parallelism changes +# ============================================================================= +cluster: + gpus_per_node: 4 + num_nodes: 128 + segment_size: 16 + +# ============================================================================= +# SIF image root — required, set by the launcher via SIF_DIR. The SWE agent's +# container_formatter (in env.nemo_gym.swe_agents_train) resolves +# `${sif_dir}//{instance_id}.sif` per SWE benchmark family. +# ============================================================================= +sif_dir: ??? + +# ============================================================================= +# Checkpointing +# ============================================================================= +checkpointing: + enabled: true + checkpoint_dir: "results/grpo_ultra_e2e" + metric_name: "val:total_reward/mean" + higher_is_better: true + keep_top_k: 1000000 + save_period: 5 + ft_keep_latest_k: 1 + ft_save_period: 1 + checkpoint_must_save_by: "00:03:30:00" + model_save_format: "safetensors" + save_consolidated: false + +# ============================================================================= +# GRPO Algorithm +# ============================================================================= +grpo: + num_prompts_per_step: 32 + num_generations_per_prompt: 16 + num_val_generations_per_prompt: 2 + max_rollout_turns: 1 + max_num_epochs: 4 + max_num_steps: 1000000 + normalize_rewards: true + use_leave_one_out_baseline: true + advantage_clip_low: -100 + advantage_clip_high: 100 + val_period: 10000 + val_at_start: false + val_at_end: false + overlong_filtering: false + max_val_samples: null + val_batch_size: 256 + seed: 42 + + use_dynamic_sampling: false + dynamic_sampling_max_gen_batches: 10 + batch_multiplier: 1 + + penalize_invalid_tool_call: true + invalid_tool_call_advantage: -5.0 + penalize_malformed_thinking: true + malformed_thinking_advantage: -5.0 + + reward_shaping: + enabled: false + overlong_buffer_length: 128 + overlong_buffer_penalty: 1 + max_response_length: ${policy.max_total_sequence_length} + stop_properly_penalty_coef: null + reward_scaling: + enabled: false + source_min: 0.0 + source_max: 1.0 + target_min: 0.0 + target_max: 1.0 + + async_grpo: + enabled: true + max_trajectory_age_steps: 1 + in_flight_weight_updates: true + recompute_kv_cache_after_weight_updates: false + + use_best_at_k: false + best_at_k_k: 8 + best_at_k_m: 1000 + + use_combined_training: false + combined_training_weight_mode: "auto" + combined_training_best_at_k_weight: 0.2 + combined_training_pass_at_1_weight: 1.0 + + dynamic_sampling_oversample_ratio: 1.0 + seq_logprob_error_threshold: 2 + +# ============================================================================= +# Loss Function +# ============================================================================= +loss_fn: + reference_policy_kl_penalty: 0.0 + reference_policy_kl_type: "k3" + kl_input_clamp_value: null + kl_output_clamp_value: null + + ratio_clip_min: 0.2 + ratio_clip_max: 0.28 + ratio_clip_c: null + use_on_policy_kl_approximation: true + use_importance_sampling_correction: true + truncated_importance_sampling_ratio: 5 + truncated_importance_sampling_ratio_min: 0.2 + truncated_importance_sampling_type: tis + sequence_level_importance_ratios: false + token_level_loss: true + force_on_policy_ratio: true + use_kl_in_reward: false + +# ============================================================================= +# Policy +# ============================================================================= +policy: + model_name: null + tokenizer: + name: ${policy.model_name} + chat_template_kwargs: null + hf_config_overrides: {} + + train_global_batch_size: 512 + train_micro_batch_size: 1 + generation_batch_size: 64 + logprob_batch_size: 1 + max_total_sequence_length: 196608 + precision: "bfloat16" + logprob_chunk_size: 1024 + offload_optimizer_for_logprob: false + + dtensor_cfg: + _v2: true + enabled: false + cpu_offload: false + sequence_parallel: false + activation_checkpointing: false + tensor_parallel_size: 1 + context_parallel_size: 1 + custom_parallel_plan: null + + megatron_cfg: + enabled: true + empty_unused_memory_level: 2 + activation_checkpointing: true + + # TP=8 spans 2 GB200 nodes (4 GPUs each) via intra-rack NVLink. + tensor_model_parallel_size: 8 + expert_tensor_parallel_size: 1 + expert_model_parallel_size: 32 + pipeline_model_parallel_size: 1 + num_layers_in_first_pipeline_stage: null + num_layers_in_last_pipeline_stage: null + context_parallel_size: 32 + pipeline_dtype: ${policy.precision} + sequence_parallel: true + + # MoE + freeze_moe_router: true + moe_router_dtype: "fp32" + moe_router_load_balancing_type: "none" + moe_router_bias_update_rate: 1.0e-3 + moe_router_enable_expert_bias: true + moe_permute_fusion: true + moe_enable_deepep: false + moe_token_dispatcher_type: "flex" + moe_flex_dispatcher_backend: "hybridep" + moe_hybridep_num_sms: 32 + moe_aux_loss_coeff: 0.0 + moe_shared_expert_overlap: false + use_gloo_process_groups: false + + # Compute + apply_rope_fusion: true + use_fused_weighted_squared_relu: true + bias_activation_fusion: false + defer_fp32_logits: true + + # Logging + track_moe_metrics: true + moe_per_layer_logging: true + do_not_average_loss: true + cp_normalize: true + calculate_per_token_loss: true + scale_loss_by_dp_cp_size: false + + # MTP — disabled + mtp_loss_scaling_factor: 0.0 + mtp_use_repeated_layer: true + mtp_num_layers: 0 + mtp_detach_heads: true + + optimizer: + optimizer: "adam" + lr: 3.0e-6 + min_lr: 3.0e-6 + weight_decay: 0.0 + bf16: true + fp16: false + params_dtype: "float32" + + adam_beta1: 0.9 + adam_beta2: 0.999 + adam_eps: 1e-8 + + sgd_momentum: 0.9 + + use_distributed_optimizer: true + use_precision_aware_optimizer: true + + clip_grad: ${policy.max_grad_norm} + + optimizer_cpu_offload: false + optimizer_offload_fraction: 0.0 + + scheduler: + start_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} + end_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} + weight_decay_incr_style: "constant" + lr_decay_style: "constant" + lr_decay_iters: null + lr_warmup_iters: 10 + lr_warmup_init: 3e-7 + override_opt_param_scheduler: true + + distributed_data_parallel_config: + grad_reduce_in_fp32: false + overlap_grad_reduce: false + overlap_param_gather: true + average_in_collective: false + use_custom_fsdp: false + data_parallel_sharding_strategy: "optim_grads_params" + + # FP8 — overridden by precision recipe in launch script + fp8_cfg: + enabled: false + fp8: "e4m3" + fp8_recipe: "mxfp8" + fp8_param: false + + first_last_layers_bf16: true + num_layers_at_start_in_bf16: 1 + num_layers_at_end_in_bf16: 1 + + checkpoint: + async_save: true + ckpt_assume_constant_structure: true + fully_parallel_save_process_group: "ep_dp" + fully_parallel_load_process_group: "ep_dp" + fully_parallel_load_exchange_algo: "broadcast" + + env_vars: null + + # --------------------------------------------------------------------------- + # Sequence Packing + # --------------------------------------------------------------------------- + dynamic_batching: + enabled: false + train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}} + sequence_length_round: 64 + + sequence_packing: + enabled: true + train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}} + algorithm: "modified_first_fit_decreasing" + sequence_length_round: 64 + fuse_loss: true + + make_sequence_length_divisible_by: ${policy.megatron_cfg.tensor_model_parallel_size} + max_grad_norm: 1.0 + optimizer: null + scheduler: null + + # --------------------------------------------------------------------------- + # Generation (vLLM) — Non-colocated, async + # --------------------------------------------------------------------------- + generation: + port_range_low: 11001 + port_range_high: 15000 + backend: "vllm" + max_new_tokens: ${policy.max_total_sequence_length} + temperature: 1.0 + top_p: 1.0 + top_k: null + stop_token_ids: null + stop_strings: null + vllm_cfg: + async_engine: true + precision: ${policy.precision} + kv_cache_dtype: "auto" + tensor_parallel_size: 8 + pipeline_parallel_size: 1 + expert_parallel_size: 1 + gpu_memory_utilization: 0.8 + max_model_len: ${policy.max_total_sequence_length} + enforce_eager: false + use_deep_gemm: false + num_last_layers_in_bf16: 0 + num_first_layers_in_bf16: 0 + enable_vllm_metrics_logger: true + vllm_metrics_logger_interval: 0.5 + expose_http_server: true + skip_tokenizer_init: false + http_server_serving_chat_kwargs: + enable_auto_tools: true + tool_parser: qwen3_coder + reasoning_parser: nano_v3 + reasoning_parser_plugin: nemo_rl/utils/nano_v3_reasoning_parser.py + + vllm_kwargs: + attention_backend: FLASH_ATTN + mamba_ssm_cache_dtype: "float32" + compilation_config: + cudagraph_capture_sizes: [1,2,4,8,16,32,64] + pass_config: + fuse_allreduce_rms: false + + colocated: + enabled: false + resources: + gpus_per_node: 4 + num_nodes: 64 + +# ============================================================================= +# Data +# ============================================================================= +data: + max_input_seq_length: null + shuffle: false + num_workers: 1 + train: + data_path: null # Set by launch script + validation: + data_path: null # Set by launch script + default: + dataset_name: NemoGymDataset + env_name: "nemo_gym" + prompt_file: null + system_prompt_file: null + processor: "nemo_gym_data_processor" + +# ============================================================================= +# Environment — NeMo Gym + SWE Agents +# ============================================================================= +env: + should_use_nemo_gym: true + should_log_nemo_gym_responses: true + nemo_gym: + nemo_gym_log_dir: "logs/nemo_gym" + skip_venv_if_present: false + num_gpu_nodes: 0 + port_range_low: 15001 + port_range_high: 20000 + invalid_tool_call_patterns: + - "" + - "" + - "" + - "" + thinking_tags: + - "" + - "" + config_paths: + - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml + - responses_api_agents/swe_agents/configs/swebench_openhands_training.yaml + # Model-specific (Nemotron chat template): keep reasoning in the multi-turn SWE + # agent history so the trajectory grows monotonically for token-in-token-out RL. + policy_model: + responses_api_models: + vllm_model: + chat_template_kwargs: + enable_thinking: true + truncate_history_thinking: false + sequential_reasoning_allowed: false + swe_agents_train: + responses_api_agents: + swe_agents: + agent_max_turns: 200 + concurrency: ${mul:2, ${mul:${grpo.num_prompts_per_step}, ${grpo.num_generations_per_prompt}}} + swebench_agent_timeout: 3600 + run_with_mixed_prompts: true + dataset_path: ${data.train.data_path} + apptainer_memory_limit_mb: 65536 + container_formatter: + - "${sif_dir}/swerebench/{instance_id}.sif" + - "${sif_dir}/swegym/sweb.eval.arm64.{instance_id}.sif" + swe_agents_val: + responses_api_agents: + swe_agents: + agent_max_turns: 200 + concurrency: ${mul:2, ${mul:${grpo.num_prompts_per_step}, ${grpo.num_generations_per_prompt}}} + swebench_agent_timeout: 3600 + dataset_path: ${data.validation.data_path} + apptainer_memory_limit_mb: 65536 + container_formatter: ${env.nemo_gym.swe_agents_train.responses_api_agents.swe_agents.container_formatter} + use_absolute_ip: true + +# ============================================================================= +# Logger +# ============================================================================= +logger: + log_dir: "logs" + num_val_samples_to_print: 0 + wandb_enabled: false + tensorboard_enabled: false + mlflow_enabled: false + monitor_gpus: true + swanlab_enabled: false + wandb: + project: "ultra-v3-swe-e2e" + name: "grpo-ultra-e2e" + tensorboard: {} + mlflow: + experiment_name: "grpo-ultra-e2e" + run_name: "grpo-ultra-e2e" + gpu_monitoring: + collection_interval: 10 + flush_interval: 10 + +# ============================================================================= +# Effort Levels +# ============================================================================= +effort_levels: + low_string: "{reasoning effort: efficient}" + low_weight: 0.1 + low_penalty: 1 + low_ub: 15000 + +# ============================================================================= +# Token IDs (model-specific, used by token-based penalties) +# ============================================================================= +token_ids: + eos: 2 # + think_open: 12 # + think_close: 13 # + +# ============================================================================= +# Reward Penalties (set reward to 0 when triggered) +# ============================================================================= +penalize_duplicated_reasoning: true # reasoning content == final answer +penalize_empty_final_answer: true # last message output has empty content +penalize_eos_token: true # eos token appears in generation +penalize_malformed_think_tag: true # /<\/think> count != 1 per turn From 242e79583e2a15c90365939195de3dc199a84290 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Wed, 24 Jun 2026 18:18:05 +0000 Subject: [PATCH 03/15] WIP: port ultra configs to main Signed-off-by: Anish Mahishi --- examples/configs/ultra/ifbench_teacher.yaml | 33 ++++++++-------- .../ultra/{mopd.yaml => mopd.yaml.disabled} | 38 ++++++++++--------- examples/configs/ultra/reasoning_teacher.yaml | 26 ++++++------- examples/configs/ultra/rlhf_teacher.yaml | 26 ++++++------- examples/configs/ultra/student_rlvr1.yaml | 31 +++++++-------- examples/configs/ultra/student_rlvr2.yaml | 31 +++++++-------- examples/configs/ultra/swe_teacher.yaml | 21 +++++----- 7 files changed, 97 insertions(+), 109 deletions(-) rename examples/configs/ultra/{mopd.yaml => mopd.yaml.disabled} (97%) diff --git a/examples/configs/ultra/ifbench_teacher.yaml b/examples/configs/ultra/ifbench_teacher.yaml index badeb263b3..42124e417b 100644 --- a/examples/configs/ultra/ifbench_teacher.yaml +++ b/examples/configs/ultra/ifbench_teacher.yaml @@ -187,6 +187,7 @@ policy: apply_rope_fusion: true use_fused_weighted_squared_relu: true bias_activation_fusion: false + gradient_accumulation_fusion: false defer_fp32_logits: true # Logging @@ -282,7 +283,7 @@ policy: sequence_length_round: 64 fuse_loss: true - make_sequence_length_divisible_by: ${policy.megatron_cfg.tensor_model_parallel_size} + make_sequence_length_divisible_by: ${mul:${policy.megatron_cfg.context_parallel_size}, ${mul:2, ${policy.megatron_cfg.tensor_model_parallel_size}}} max_grad_norm: 1.0 optimizer: null scheduler: null @@ -320,11 +321,11 @@ policy: vllm_metrics_logger_interval: 0.5 expose_http_server: true skip_tokenizer_init: false + reasoning_parser_plugin: nemo_rl/models/generation/vllm/reasoning_parsers/nano_v3_reasoning_parser.py http_server_serving_chat_kwargs: enable_auto_tools: true tool_parser: qwen3_coder reasoning_parser: nano_v3 - reasoning_parser_plugin: nemo_rl/utils/nano_v3_reasoning_parser.py vllm_kwargs: attention_backend: FLASH_ATTN @@ -347,6 +348,7 @@ data: max_input_seq_length: null shuffle: false num_workers: 1 + use_multiple_dataloader: false train: data_path: null validation: @@ -372,6 +374,12 @@ env: nemo_gym: nemo_gym_log_dir: "logs/nemo_gym" skip_venv_if_present: true + # NeMo-RL-side effort shaping; stripped before forwarding to Gym. + effort_levels: + low_string: "{reasoning effort: efficient}" + low_weight: 0.1 + low_penalty: 1 + low_ub: 15000 port_range_low: 5000 port_range_high: 5999 invalid_tool_call_patterns: @@ -443,9 +451,9 @@ env: debug: true vllm_serve_env_vars: VLLM_RAY_DP_PACK_STRATEGY: strict + VLLM_ATTENTION_BACKEND: TRITON_ATTN vllm_serve_kwargs: - attention_backend: TRITON_ATTN tensor_parallel_size: 4 data_parallel_size: 2 data_parallel_size_local: 1 @@ -469,9 +477,10 @@ env: debug: true vllm_serve_env_vars: VLLM_RAY_DP_PACK_STRATEGY: strict + VLLM_ATTENTION_BACKEND: FLASH_ATTN + NCCL_MNNVL_ENABLE: "0" vllm_serve_kwargs: - attention_backend: FLASH_ATTN tensor_parallel_size: 4 data_parallel_size: 4 data_parallel_size_local: 1 @@ -488,8 +497,6 @@ env: num_threads: 112 compilation_config: cudagraph_capture_sizes: [1,2,4,8,16,32] - server_env: - NCCL_MNNVL_ENABLE: "0" inverse_if: resources_servers: @@ -532,12 +539,12 @@ env: use_principle: true default_principle: "You will be given one or more evaluation criteria (rubrics).\nEvaluate both responses on EACH criterion individually first, then synthesize an overall judgment.\nCriteria:\n\n1. Please act as an impartial judge and evaluate the quality of the responses provided by two AI assistants to the user prompt. Begin your evaluation by generating your own answer to the prompt. You must provide your answer before judging any answers. When evaluating the assistants' answers, compare both assistants' answers with your answer. You must identify and correct any mistakes or inaccurate information. Then consider if the assistant's answers are helpful, relevant, and concise. Helpful means the answer correctly responds to the prompt or follows the instructions. Note when user prompt has any ambiguity or more than one interpretation, it is more helpful and appropriate to ask for clarifications or more information from the user than providing an answer based on assumptions. Relevant means all parts of the response closely connect or are appropriate to what is being asked. Concise means the response is clear and not verbose or excessive. Then consider the creativity and novelty of the assistant's answers when needed. Finally, identify any missing important information in the assistants' answers that would be beneficial to include when responding to the user prompt." aggregator_method: "simple_tiebreaker" + use_golden_anchor: false reasoning_bonus: 0.5 answer_bonus: 0.5 top_percentile: 0.2 genrm_parse_retries: 1 group_reasoning_length_penalty_coeff: 0.1 - use_golden_anchor: false group_answer_length_penalty_coeff: 0.1 group_style_penalty_coeff: 0.0 default_score: 3.0 @@ -557,6 +564,7 @@ env: ray_worker_py_executable: /opt/ray_venvs/nemo_rl.models.generation.vllm.vllm_worker_async.VllmAsyncGenerationWorker/bin/python vllm_serve_env_vars: VLLM_RAY_DP_PACK_STRATEGY: strict + NCCL_MNNVL_ENABLE: "0" vllm_serve_kwargs: tensor_parallel_size: 4 @@ -574,8 +582,6 @@ env: load_format: instanttensor compilation_config: cudagraph_capture_sizes: [1,2,4,8,16,32] - server_env: - NCCL_MNNVL_ENABLE: "0" lc_judge: resources_servers: @@ -628,15 +634,6 @@ logger: collection_interval: 10 flush_interval: 10 -# ============================================================================= -# Effort Levels -# ============================================================================= -effort_levels: - low_string: "{reasoning effort: efficient}" - low_weight: 0.1 - low_penalty: 1 - low_ub: 15000 - # ============================================================================= # Token IDs (model-specific) # ============================================================================= diff --git a/examples/configs/ultra/mopd.yaml b/examples/configs/ultra/mopd.yaml.disabled similarity index 97% rename from examples/configs/ultra/mopd.yaml rename to examples/configs/ultra/mopd.yaml.disabled index 33ec40d077..a6b43520c7 100644 --- a/examples/configs/ultra/mopd.yaml +++ b/examples/configs/ultra/mopd.yaml.disabled @@ -1,3 +1,8 @@ +# NOTE: Disabled on main because this config requires OPD/MOPD runtime support +# from the ultra-v3 branch (OPDAdvantageEstimator, on_policy_distillation, +# and non-colocated teacher worker groups). Rename to .yaml only after that +# support lands on main. + # ============================================================================= # Nemotron 3 Ultra — MOPD Stage 1 (on-policy distillation from teacher panel) # ============================================================================= @@ -207,6 +212,7 @@ policy: apply_rope_fusion: true use_fused_weighted_squared_relu: true bias_activation_fusion: false + gradient_accumulation_fusion: false defer_fp32_logits: true # Logging @@ -303,7 +309,7 @@ policy: sequence_length_round: 64 fuse_loss: true - make_sequence_length_divisible_by: ${policy.megatron_cfg.tensor_model_parallel_size} + make_sequence_length_divisible_by: ${mul:${policy.megatron_cfg.context_parallel_size}, ${mul:2, ${policy.megatron_cfg.tensor_model_parallel_size}}} max_grad_norm: 1.0 optimizer: null scheduler: null @@ -341,11 +347,11 @@ policy: vllm_metrics_logger_interval: 0.5 expose_http_server: true skip_tokenizer_init: false + reasoning_parser_plugin: nemo_rl/models/generation/vllm/reasoning_parsers/nano_v3_reasoning_parser.py http_server_serving_chat_kwargs: enable_auto_tools: true tool_parser: qwen3_coder reasoning_parser: nano_v3 - reasoning_parser_plugin: nemo_rl/utils/nano_v3_reasoning_parser.py vllm_kwargs: attention_backend: FLASH_ATTN @@ -368,6 +374,7 @@ data: max_input_seq_length: null shuffle: false num_workers: 1 + use_multiple_dataloader: false train: data_path: null # Set by launch script validation: @@ -389,6 +396,12 @@ env: nemo_gym: nemo_gym_log_dir: "logs/nemo_gym" skip_venv_if_present: true + # NeMo-RL-side effort shaping; stripped before forwarding to Gym. + effort_levels: + low_string: "{reasoning effort: efficient}" + low_weight: 0.1 + low_penalty: 1 + low_ub: 15000 port_range_low: 5000 port_range_high: 5999 # MOPD trains from teacher logprobs, so skip env verifiers/unit tests during rollout. @@ -497,9 +510,9 @@ env: debug: true vllm_serve_env_vars: VLLM_RAY_DP_PACK_STRATEGY: strict + VLLM_ATTENTION_BACKEND: TRITON_ATTN vllm_serve_kwargs: - attention_backend: TRITON_ATTN tensor_parallel_size: 4 data_parallel_size: 2 data_parallel_size_local: 1 @@ -511,7 +524,7 @@ env: num_threads: 112 compilation_config: cudagraph_capture_sizes: [1,2,4,8,16] - + # nl2bash / General Judge: TP=4 on GB200 192GB nl2bash_judge_model: @@ -524,9 +537,10 @@ env: debug: true vllm_serve_env_vars: VLLM_RAY_DP_PACK_STRATEGY: strict + VLLM_ATTENTION_BACKEND: FLASH_ATTN + NCCL_MNNVL_ENABLE: "0" vllm_serve_kwargs: - attention_backend: FLASH_ATTN tensor_parallel_size: 4 data_parallel_size: 4 data_parallel_size_local: 1 @@ -541,8 +555,6 @@ env: num_threads: 112 compilation_config: cudagraph_capture_sizes: [1,2,4,8,16,32] - server_env: - NCCL_MNNVL_ENABLE: "0" inverse_if: resources_servers: @@ -607,6 +619,7 @@ env: ray_worker_py_executable: /opt/ray_venvs/nemo_rl.models.generation.vllm.vllm_worker_async.VllmAsyncGenerationWorker/bin/python vllm_serve_env_vars: VLLM_RAY_DP_PACK_STRATEGY: strict + NCCL_MNNVL_ENABLE: "0" vllm_serve_kwargs: tensor_parallel_size: 4 @@ -624,8 +637,6 @@ env: load_format: instanttensor compilation_config: cudagraph_capture_sizes: [1,2,4,8,16,32] - server_env: - NCCL_MNNVL_ENABLE: "0" lc_judge: resources_servers: @@ -701,15 +712,6 @@ logger: collection_interval: 10 flush_interval: 10 -# ============================================================================= -# Effort Levels -# ============================================================================= -effort_levels: - low_string: "{reasoning effort: efficient}" - low_weight: 0.1 - low_penalty: 1 - low_ub: 15000 - # ============================================================================= # Token IDs (model-specific, used by token-based penalties) # ============================================================================= diff --git a/examples/configs/ultra/reasoning_teacher.yaml b/examples/configs/ultra/reasoning_teacher.yaml index 3b43e0f8c8..a7261cb582 100644 --- a/examples/configs/ultra/reasoning_teacher.yaml +++ b/examples/configs/ultra/reasoning_teacher.yaml @@ -190,6 +190,7 @@ policy: apply_rope_fusion: true use_fused_weighted_squared_relu: true bias_activation_fusion: false + gradient_accumulation_fusion: false defer_fp32_logits: true # Logging @@ -285,7 +286,7 @@ policy: sequence_length_round: 64 fuse_loss: true - make_sequence_length_divisible_by: ${policy.megatron_cfg.tensor_model_parallel_size} + make_sequence_length_divisible_by: ${mul:${policy.megatron_cfg.context_parallel_size}, ${mul:2, ${policy.megatron_cfg.tensor_model_parallel_size}}} max_grad_norm: 1.0 optimizer: null scheduler: null @@ -323,11 +324,11 @@ policy: vllm_metrics_logger_interval: 0.5 expose_http_server: true skip_tokenizer_init: false + reasoning_parser_plugin: nemo_rl/models/generation/vllm/reasoning_parsers/nano_v3_reasoning_parser.py http_server_serving_chat_kwargs: enable_auto_tools: true tool_parser: qwen3_coder reasoning_parser: nano_v3 - reasoning_parser_plugin: nemo_rl/utils/nano_v3_reasoning_parser.py vllm_kwargs: attention_backend: FLASH_ATTN @@ -350,6 +351,7 @@ data: max_input_seq_length: null shuffle: false num_workers: 1 + use_multiple_dataloader: false train: data_path: null validation: @@ -373,6 +375,12 @@ env: nemo_gym: nemo_gym_log_dir: "logs/nemo_gym" skip_venv_if_present: true + # NeMo-RL-side effort shaping; stripped before forwarding to Gym. + effort_levels: + low_string: "{reasoning effort: efficient}" + low_weight: 0.1 + low_penalty: 1 + low_ub: 15000 port_range_low: 5000 port_range_high: 5999 invalid_tool_call_patterns: @@ -421,9 +429,10 @@ env: debug: true vllm_serve_env_vars: VLLM_RAY_DP_PACK_STRATEGY: strict + VLLM_ATTENTION_BACKEND: FLASH_ATTN + NCCL_MNNVL_ENABLE: "0" vllm_serve_kwargs: - attention_backend: FLASH_ATTN tensor_parallel_size: 4 data_parallel_size: 4 data_parallel_size_local: 1 @@ -440,8 +449,6 @@ env: num_threads: 112 compilation_config: cudagraph_capture_sizes: [1,2,4,8,16,32] - server_env: - NCCL_MNNVL_ENABLE: "0" equivalence_llm_judge: resources_servers: @@ -474,15 +481,6 @@ logger: collection_interval: 10 flush_interval: 10 -# ============================================================================= -# Effort Levels -# ============================================================================= -effort_levels: - low_string: "{reasoning effort: efficient}" - low_weight: 0.1 - low_penalty: 1 - low_ub: 15000 - # ============================================================================= # Token IDs (model-specific) # ============================================================================= diff --git a/examples/configs/ultra/rlhf_teacher.yaml b/examples/configs/ultra/rlhf_teacher.yaml index 97ffb0e521..ad9b690eec 100644 --- a/examples/configs/ultra/rlhf_teacher.yaml +++ b/examples/configs/ultra/rlhf_teacher.yaml @@ -188,6 +188,7 @@ policy: apply_rope_fusion: true use_fused_weighted_squared_relu: true bias_activation_fusion: false + gradient_accumulation_fusion: false defer_fp32_logits: true # Logging @@ -283,7 +284,7 @@ policy: sequence_length_round: 64 fuse_loss: true - make_sequence_length_divisible_by: ${policy.megatron_cfg.tensor_model_parallel_size} + make_sequence_length_divisible_by: ${mul:${policy.megatron_cfg.context_parallel_size}, ${mul:2, ${policy.megatron_cfg.tensor_model_parallel_size}}} max_grad_norm: 1.0 optimizer: null scheduler: null @@ -321,11 +322,11 @@ policy: vllm_metrics_logger_interval: 0.5 expose_http_server: true skip_tokenizer_init: false + reasoning_parser_plugin: nemo_rl/models/generation/vllm/reasoning_parsers/nano_v3_reasoning_parser.py http_server_serving_chat_kwargs: enable_auto_tools: true tool_parser: qwen3_coder reasoning_parser: nano_v3 - reasoning_parser_plugin: nemo_rl/utils/nano_v3_reasoning_parser.py vllm_kwargs: attention_backend: FLASH_ATTN @@ -348,6 +349,7 @@ data: max_input_seq_length: null shuffle: false num_workers: 1 + use_multiple_dataloader: false train: data_path: null validation: @@ -372,6 +374,12 @@ env: nemo_gym: nemo_gym_log_dir: "logs/nemo_gym" skip_venv_if_present: true + # NeMo-RL-side effort shaping; stripped before forwarding to Gym. + effort_levels: + low_string: "{reasoning effort: efficient}" + low_weight: 0.1 + low_penalty: 1 + low_ub: 15000 port_range_low: 5000 port_range_high: 5999 invalid_tool_call_patterns: @@ -426,12 +434,12 @@ env: use_principle: true default_principle: "You will be given one or more evaluation criteria (rubrics).\nEvaluate both responses on EACH criterion individually first, then synthesize an overall judgment.\nCriteria:\n\n1. Please act as an impartial judge and evaluate the quality of the responses provided by two AI assistants to the user prompt. Begin your evaluation by generating your own answer to the prompt. You must provide your answer before judging any answers. When evaluating the assistants' answers, compare both assistants' answers with your answer. You must identify and correct any mistakes or inaccurate information. Then consider if the assistant's answers are helpful, relevant, and concise. Helpful means the answer correctly responds to the prompt or follows the instructions. Note when user prompt has any ambiguity or more than one interpretation, it is more helpful and appropriate to ask for clarifications or more information from the user than providing an answer based on assumptions. Relevant means all parts of the response closely connect or are appropriate to what is being asked. Concise means the response is clear and not verbose or excessive. Then consider the creativity and novelty of the assistant's answers when needed. Finally, identify any missing important information in the assistants' answers that would be beneficial to include when responding to the user prompt." aggregator_method: "simple_tiebreaker" + use_golden_anchor: false reasoning_bonus: 0.0 answer_bonus: 0.0 top_percentile: 0.2 genrm_parse_retries: 1 group_reasoning_length_penalty_coeff: 0.0 - use_golden_anchor: false group_answer_length_penalty_coeff: 0.0 group_style_penalty_coeff: 0.0 default_score: 3.0 @@ -451,6 +459,7 @@ env: ray_worker_py_executable: /opt/ray_venvs/nemo_rl.models.generation.vllm.vllm_worker_async.VllmAsyncGenerationWorker/bin/python vllm_serve_env_vars: VLLM_RAY_DP_PACK_STRATEGY: strict + NCCL_MNNVL_ENABLE: "0" vllm_serve_kwargs: tensor_parallel_size: 4 @@ -468,8 +477,6 @@ env: load_format: instanttensor compilation_config: cudagraph_capture_sizes: [1,2,4,8,16,32] - server_env: - NCCL_MNNVL_ENABLE: "0" # ============================================================================= @@ -494,15 +501,6 @@ logger: collection_interval: 10 flush_interval: 10 -# ============================================================================= -# Effort Levels -# ============================================================================= -effort_levels: - low_string: "{reasoning effort: efficient}" - low_weight: 0.1 - low_penalty: 1 - low_ub: 15000 - # ============================================================================= # Token IDs (model-specific) # ============================================================================= diff --git a/examples/configs/ultra/student_rlvr1.yaml b/examples/configs/ultra/student_rlvr1.yaml index 692bc6066b..aae77f431f 100644 --- a/examples/configs/ultra/student_rlvr1.yaml +++ b/examples/configs/ultra/student_rlvr1.yaml @@ -184,6 +184,7 @@ policy: apply_rope_fusion: true use_fused_weighted_squared_relu: true bias_activation_fusion: false + gradient_accumulation_fusion: false defer_fp32_logits: true # Logging @@ -279,7 +280,7 @@ policy: sequence_length_round: 64 fuse_loss: true - make_sequence_length_divisible_by: ${policy.megatron_cfg.tensor_model_parallel_size} + make_sequence_length_divisible_by: ${mul:${policy.megatron_cfg.context_parallel_size}, ${mul:2, ${policy.megatron_cfg.tensor_model_parallel_size}}} max_grad_norm: 1.0 optimizer: null scheduler: null @@ -317,11 +318,11 @@ policy: vllm_metrics_logger_interval: 0.5 expose_http_server: true skip_tokenizer_init: false + reasoning_parser_plugin: nemo_rl/models/generation/vllm/reasoning_parsers/nano_v3_reasoning_parser.py http_server_serving_chat_kwargs: enable_auto_tools: true tool_parser: qwen3_coder reasoning_parser: nano_v3 - reasoning_parser_plugin: nemo_rl/utils/nano_v3_reasoning_parser.py vllm_kwargs: attention_backend: FLASH_ATTN @@ -344,6 +345,7 @@ data: max_input_seq_length: null shuffle: false num_workers: 1 + use_multiple_dataloader: false train: data_path: null validation: @@ -369,6 +371,12 @@ env: nemo_gym: nemo_gym_log_dir: "logs/nemo_gym" skip_venv_if_present: true + # NeMo-RL-side effort shaping; stripped before forwarding to Gym. + effort_levels: + low_string: "{reasoning effort: efficient}" + low_weight: 0.1 + low_penalty: 1 + low_ub: 15000 port_range_low: 5000 port_range_high: 5999 invalid_tool_call_patterns: @@ -463,9 +471,9 @@ env: debug: true vllm_serve_env_vars: VLLM_RAY_DP_PACK_STRATEGY: strict + VLLM_ATTENTION_BACKEND: TRITON_ATTN vllm_serve_kwargs: - attention_backend: TRITON_ATTN tensor_parallel_size: 4 data_parallel_size: 2 data_parallel_size_local: 1 @@ -489,9 +497,10 @@ env: debug: true vllm_serve_env_vars: VLLM_RAY_DP_PACK_STRATEGY: strict + VLLM_ATTENTION_BACKEND: FLASH_ATTN + NCCL_MNNVL_ENABLE: "0" vllm_serve_kwargs: - attention_backend: FLASH_ATTN tensor_parallel_size: 4 data_parallel_size: 4 data_parallel_size_local: 1 @@ -508,8 +517,6 @@ env: num_threads: 112 compilation_config: cudagraph_capture_sizes: [1,2,4,8,16,32] - server_env: - NCCL_MNNVL_ENABLE: "0" inverse_if: resources_servers: @@ -576,6 +583,7 @@ env: ray_worker_py_executable: /opt/ray_venvs/nemo_rl.models.generation.vllm.vllm_worker_async.VllmAsyncGenerationWorker/bin/python vllm_serve_env_vars: VLLM_RAY_DP_PACK_STRATEGY: strict + NCCL_MNNVL_ENABLE: "0" vllm_serve_kwargs: tensor_parallel_size: 4 @@ -593,8 +601,6 @@ env: load_format: instanttensor compilation_config: cudagraph_capture_sizes: [1,2,4,8,16,32] - server_env: - NCCL_MNNVL_ENABLE: "0" lc_judge: resources_servers: @@ -647,15 +653,6 @@ logger: collection_interval: 10 flush_interval: 10 -# ============================================================================= -# Effort Levels -# ============================================================================= -effort_levels: - low_string: "{reasoning effort: efficient}" - low_weight: 0.1 - low_penalty: 1 - low_ub: 15000 - # ============================================================================= # Token IDs (model-specific) # ============================================================================= diff --git a/examples/configs/ultra/student_rlvr2.yaml b/examples/configs/ultra/student_rlvr2.yaml index 2f46aa402c..c66492c31a 100644 --- a/examples/configs/ultra/student_rlvr2.yaml +++ b/examples/configs/ultra/student_rlvr2.yaml @@ -185,6 +185,7 @@ policy: apply_rope_fusion: true use_fused_weighted_squared_relu: true bias_activation_fusion: false + gradient_accumulation_fusion: false defer_fp32_logits: true # Logging @@ -280,7 +281,7 @@ policy: sequence_length_round: 64 fuse_loss: true - make_sequence_length_divisible_by: ${policy.megatron_cfg.tensor_model_parallel_size} + make_sequence_length_divisible_by: ${mul:${policy.megatron_cfg.context_parallel_size}, ${mul:2, ${policy.megatron_cfg.tensor_model_parallel_size}}} max_grad_norm: 1.0 optimizer: null scheduler: null @@ -318,11 +319,11 @@ policy: vllm_metrics_logger_interval: 0.5 expose_http_server: true skip_tokenizer_init: false + reasoning_parser_plugin: nemo_rl/models/generation/vllm/reasoning_parsers/nano_v3_reasoning_parser.py http_server_serving_chat_kwargs: enable_auto_tools: true tool_parser: qwen3_coder reasoning_parser: nano_v3 - reasoning_parser_plugin: nemo_rl/utils/nano_v3_reasoning_parser.py vllm_kwargs: attention_backend: FLASH_ATTN @@ -345,6 +346,7 @@ data: max_input_seq_length: null shuffle: false num_workers: 1 + use_multiple_dataloader: false train: data_path: null validation: @@ -370,6 +372,12 @@ env: nemo_gym: nemo_gym_log_dir: "logs/nemo_gym" skip_venv_if_present: true + # NeMo-RL-side effort shaping; stripped before forwarding to Gym. + effort_levels: + low_string: "{reasoning effort: efficient}" + low_weight: 0.1 + low_penalty: 1 + low_ub: 15000 port_range_low: 5000 port_range_high: 5999 invalid_tool_call_patterns: @@ -464,9 +472,9 @@ env: debug: true vllm_serve_env_vars: VLLM_RAY_DP_PACK_STRATEGY: strict + VLLM_ATTENTION_BACKEND: TRITON_ATTN vllm_serve_kwargs: - attention_backend: TRITON_ATTN tensor_parallel_size: 4 data_parallel_size: 2 data_parallel_size_local: 1 @@ -490,9 +498,10 @@ env: debug: true vllm_serve_env_vars: VLLM_RAY_DP_PACK_STRATEGY: strict + VLLM_ATTENTION_BACKEND: FLASH_ATTN + NCCL_MNNVL_ENABLE: "0" vllm_serve_kwargs: - attention_backend: FLASH_ATTN tensor_parallel_size: 4 data_parallel_size: 4 data_parallel_size_local: 1 @@ -509,8 +518,6 @@ env: num_threads: 112 compilation_config: cudagraph_capture_sizes: [1,2,4,8,16,32] - server_env: - NCCL_MNNVL_ENABLE: "0" inverse_if: resources_servers: @@ -577,6 +584,7 @@ env: ray_worker_py_executable: /opt/ray_venvs/nemo_rl.models.generation.vllm.vllm_worker_async.VllmAsyncGenerationWorker/bin/python vllm_serve_env_vars: VLLM_RAY_DP_PACK_STRATEGY: strict + NCCL_MNNVL_ENABLE: "0" vllm_serve_kwargs: tensor_parallel_size: 4 @@ -594,8 +602,6 @@ env: load_format: instanttensor compilation_config: cudagraph_capture_sizes: [1,2,4,8,16,32] - server_env: - NCCL_MNNVL_ENABLE: "0" lc_judge: resources_servers: @@ -648,15 +654,6 @@ logger: collection_interval: 10 flush_interval: 10 -# ============================================================================= -# Effort Levels -# ============================================================================= -effort_levels: - low_string: "{reasoning effort: efficient}" - low_weight: 0.1 - low_penalty: 1 - low_ub: 15000 - # ============================================================================= # Token IDs (model-specific) # ============================================================================= diff --git a/examples/configs/ultra/swe_teacher.yaml b/examples/configs/ultra/swe_teacher.yaml index 69b8620796..43c38b23c2 100644 --- a/examples/configs/ultra/swe_teacher.yaml +++ b/examples/configs/ultra/swe_teacher.yaml @@ -206,6 +206,7 @@ policy: apply_rope_fusion: true use_fused_weighted_squared_relu: true bias_activation_fusion: false + gradient_accumulation_fusion: false defer_fp32_logits: true # Logging @@ -300,7 +301,7 @@ policy: sequence_length_round: 64 fuse_loss: true - make_sequence_length_divisible_by: ${policy.megatron_cfg.tensor_model_parallel_size} + make_sequence_length_divisible_by: ${mul:${policy.megatron_cfg.context_parallel_size}, ${mul:2, ${policy.megatron_cfg.tensor_model_parallel_size}}} max_grad_norm: 1.0 optimizer: null scheduler: null @@ -335,11 +336,11 @@ policy: vllm_metrics_logger_interval: 0.5 expose_http_server: true skip_tokenizer_init: false + reasoning_parser_plugin: nemo_rl/models/generation/vllm/reasoning_parsers/nano_v3_reasoning_parser.py http_server_serving_chat_kwargs: enable_auto_tools: true tool_parser: qwen3_coder reasoning_parser: nano_v3 - reasoning_parser_plugin: nemo_rl/utils/nano_v3_reasoning_parser.py vllm_kwargs: attention_backend: FLASH_ATTN @@ -362,6 +363,7 @@ data: max_input_seq_length: null shuffle: false num_workers: 1 + use_multiple_dataloader: false train: data_path: null # Set by launch script validation: @@ -382,6 +384,12 @@ env: nemo_gym: nemo_gym_log_dir: "logs/nemo_gym" skip_venv_if_present: false + # NeMo-RL-side effort shaping; stripped before forwarding to Gym. + effort_levels: + low_string: "{reasoning effort: efficient}" + low_weight: 0.1 + low_penalty: 1 + low_ub: 15000 num_gpu_nodes: 0 port_range_low: 15001 port_range_high: 20000 @@ -450,15 +458,6 @@ logger: collection_interval: 10 flush_interval: 10 -# ============================================================================= -# Effort Levels -# ============================================================================= -effort_levels: - low_string: "{reasoning effort: efficient}" - low_weight: 0.1 - low_penalty: 1 - low_ub: 15000 - # ============================================================================= # Token IDs (model-specific, used by token-based penalties) # ============================================================================= From 0afe4afb473ea66723188f7ab5a72934f207f9f9 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Tue, 30 Jun 2026 11:28:31 -0700 Subject: [PATCH 04/15] fix(ultra): update runtime config compatibility Signed-off-by: Anish Mahishi --- examples/configs/ultra/ifbench_teacher.yaml | 1 + examples/configs/ultra/mopd.yaml.disabled | 1 + examples/configs/ultra/reasoning_teacher.yaml | 1 + examples/configs/ultra/rlhf_teacher.yaml | 1 + examples/configs/ultra/student_rlvr1.yaml | 1 + examples/configs/ultra/student_rlvr2.yaml | 1 + examples/configs/ultra/swe_teacher.yaml | 1 + ultra_launch.sh | 17 ++++++++++++++++- 8 files changed, 23 insertions(+), 1 deletion(-) diff --git a/examples/configs/ultra/ifbench_teacher.yaml b/examples/configs/ultra/ifbench_teacher.yaml index 42124e417b..78415dde8a 100644 --- a/examples/configs/ultra/ifbench_teacher.yaml +++ b/examples/configs/ultra/ifbench_teacher.yaml @@ -39,6 +39,7 @@ checkpointing: checkpoint_must_save_by: "00:47:35:00" model_save_format: "safetensors" save_consolidated: false + save_optimizer: true # ============================================================================= # GRPO Algorithm diff --git a/examples/configs/ultra/mopd.yaml.disabled b/examples/configs/ultra/mopd.yaml.disabled index a6b43520c7..3d827adb99 100644 --- a/examples/configs/ultra/mopd.yaml.disabled +++ b/examples/configs/ultra/mopd.yaml.disabled @@ -59,6 +59,7 @@ checkpointing: checkpoint_must_save_by: "00:23:30:00" model_save_format: "safetensors" save_consolidated: false + save_optimizer: true # ============================================================================= # GRPO Algorithm diff --git a/examples/configs/ultra/reasoning_teacher.yaml b/examples/configs/ultra/reasoning_teacher.yaml index a7261cb582..74fd6077d6 100644 --- a/examples/configs/ultra/reasoning_teacher.yaml +++ b/examples/configs/ultra/reasoning_teacher.yaml @@ -43,6 +43,7 @@ checkpointing: checkpoint_must_save_by: "00:47:35:00" model_save_format: "safetensors" save_consolidated: false + save_optimizer: true # ============================================================================= # GRPO Algorithm diff --git a/examples/configs/ultra/rlhf_teacher.yaml b/examples/configs/ultra/rlhf_teacher.yaml index ad9b690eec..d06c5659ec 100644 --- a/examples/configs/ultra/rlhf_teacher.yaml +++ b/examples/configs/ultra/rlhf_teacher.yaml @@ -40,6 +40,7 @@ checkpointing: checkpoint_must_save_by: "00:47:35:00" model_save_format: "safetensors" save_consolidated: false + save_optimizer: true # ============================================================================= # GRPO Algorithm diff --git a/examples/configs/ultra/student_rlvr1.yaml b/examples/configs/ultra/student_rlvr1.yaml index aae77f431f..b732668896 100644 --- a/examples/configs/ultra/student_rlvr1.yaml +++ b/examples/configs/ultra/student_rlvr1.yaml @@ -36,6 +36,7 @@ checkpointing: checkpoint_must_save_by: "00:47:35:00" model_save_format: "safetensors" save_consolidated: false + save_optimizer: true # ============================================================================= # GRPO Algorithm diff --git a/examples/configs/ultra/student_rlvr2.yaml b/examples/configs/ultra/student_rlvr2.yaml index c66492c31a..b4a21203e3 100644 --- a/examples/configs/ultra/student_rlvr2.yaml +++ b/examples/configs/ultra/student_rlvr2.yaml @@ -37,6 +37,7 @@ checkpointing: checkpoint_must_save_by: "00:47:35:00" model_save_format: "safetensors" save_consolidated: false + save_optimizer: true # ============================================================================= # GRPO Algorithm diff --git a/examples/configs/ultra/swe_teacher.yaml b/examples/configs/ultra/swe_teacher.yaml index 43c38b23c2..3756753a6a 100644 --- a/examples/configs/ultra/swe_teacher.yaml +++ b/examples/configs/ultra/swe_teacher.yaml @@ -57,6 +57,7 @@ checkpointing: checkpoint_must_save_by: "00:03:30:00" model_save_format: "safetensors" save_consolidated: false + save_optimizer: true # ============================================================================= # GRPO Algorithm diff --git a/ultra_launch.sh b/ultra_launch.sh index c51458c2dc..ef595b02bc 100755 --- a/ultra_launch.sh +++ b/ultra_launch.sh @@ -47,6 +47,8 @@ set -euo pipefail # NRL_MAX_STEPS= Override grpo.max_num_steps # EXTRA_MOUNTS= Comma-separated host:container pairs # USE_SNAPSHOT=1 Snapshot source tree at submission +# USE_CUSTOM_VLLM=1 1 to require the Ultra vLLM fork; +# 0 to use the container's regular vLLM # DRY_RUN=0 1 to print TRAIN_CMD and exit # INTERACTIVE=0 1 to bring up Ray and idle for attach # (no training driver) for debugging @@ -303,7 +305,19 @@ export NEMO_SKILLS_SANDBOX_PORT="${NEMO_SKILLS_SANDBOX_PORT:-6000}" export RAY_LOG_SYNC_FREQUENCY="${RAY_LOG_SYNC_FREQUENCY:-60}" CODE_ROOT="/opt/nemo-rl" -VLLM_ENV_SOURCE="source /opt/nemo-rl/3rdparty/vllm/nemo-rl.env && " +USE_CUSTOM_VLLM="${USE_CUSTOM_VLLM:-1}" +case "${USE_CUSTOM_VLLM}" in + 1) + VLLM_ENV_SOURCE="source /opt/nemo-rl/3rdparty/vllm/nemo-rl.env && " + ;; + 0) + VLLM_ENV_SOURCE="" + ;; + *) + echo "ERROR: USE_CUSTOM_VLLM must be 0 or 1, got: ${USE_CUSTOM_VLLM}" >&2 + exit 1 + ;; +esac # ============================================================================= # Persistent cache directories @@ -631,6 +645,7 @@ echo " Model: ${MODEL_PATH}" echo " Train data: ${TRAIN_PATH}" echo " Val data: ${VAL_PATH}" echo " Container: ${CONTAINER}" +echo " Custom vLLM: ${USE_CUSTOM_VLLM}" echo " Sandbox: ${SANDBOX_CONTAINER}" if [[ "${USE_SNAPSHOT}" == "1" ]]; then echo " Snapshot: ${SNAPSHOT_DIR}" From 4b3540194b630a1177881b0a12bf9252cd06a815 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Sun, 5 Jul 2026 10:00:51 -0700 Subject: [PATCH 05/15] fix(ultra): update Ultra recipe config compatibility Signed-off-by: Anish Mahishi --- examples/configs/ultra/ifbench_teacher.yaml | 22 +++++++++---------- .../ultra/{mopd.yaml.disabled => mopd.yaml} | 22 +++++++++---------- examples/configs/ultra/reasoning_teacher.yaml | 22 +++++++++---------- examples/configs/ultra/rlhf_teacher.yaml | 22 +++++++++---------- examples/configs/ultra/student_rlvr1.yaml | 22 +++++++++---------- examples/configs/ultra/student_rlvr2.yaml | 22 +++++++++---------- examples/configs/ultra/swe_teacher.yaml | 21 ++++++++---------- 7 files changed, 69 insertions(+), 84 deletions(-) rename examples/configs/ultra/{mopd.yaml.disabled => mopd.yaml} (98%) diff --git a/examples/configs/ultra/ifbench_teacher.yaml b/examples/configs/ultra/ifbench_teacher.yaml index 78415dde8a..25b3a6a0b0 100644 --- a/examples/configs/ultra/ifbench_teacher.yaml +++ b/examples/configs/ultra/ifbench_teacher.yaml @@ -375,6 +375,7 @@ env: nemo_gym: nemo_gym_log_dir: "logs/nemo_gym" skip_venv_if_present: true + num_gpu_nodes: 20 # NeMo-RL-side effort shaping; stripped before forwarding to Gym. effort_levels: low_string: "{reasoning effort: efficient}" @@ -635,18 +636,15 @@ logger: collection_interval: 10 flush_interval: 10 -# ============================================================================= -# Token IDs (model-specific) -# ============================================================================= -token_ids: - eos: 2 - think_open: 12 - think_close: 13 - # ============================================================================= # Reward Penalties (set reward to 0 when triggered) # ============================================================================= -penalize_duplicated_reasoning: true -penalize_empty_final_answer: true -penalize_eos_token: true -penalize_malformed_think_tag: true +reward_penalties: + penalize_duplicated_reasoning: true + penalize_empty_final_answer: true + penalize_unwanted_tokens: true + penalize_malformed_think_tag: true + token_ids: + unwanted: [2] + think_open: 12 + think_close: 13 diff --git a/examples/configs/ultra/mopd.yaml.disabled b/examples/configs/ultra/mopd.yaml similarity index 98% rename from examples/configs/ultra/mopd.yaml.disabled rename to examples/configs/ultra/mopd.yaml index 3d827adb99..081ec8219b 100644 --- a/examples/configs/ultra/mopd.yaml.disabled +++ b/examples/configs/ultra/mopd.yaml @@ -397,6 +397,7 @@ env: nemo_gym: nemo_gym_log_dir: "logs/nemo_gym" skip_venv_if_present: true + num_gpu_nodes: 12 # NeMo-RL-side effort shaping; stripped before forwarding to Gym. effort_levels: low_string: "{reasoning effort: efficient}" @@ -713,21 +714,18 @@ logger: collection_interval: 10 flush_interval: 10 -# ============================================================================= -# Token IDs (model-specific, used by token-based penalties) -# ============================================================================= -token_ids: - eos: 2 # - think_open: 12 # - think_close: 13 # - # ============================================================================= # Reward Penalties (set reward to 0 when triggered) # ============================================================================= -penalize_duplicated_reasoning: true # reasoning content == final answer -penalize_empty_final_answer: true # last message output has empty content -penalize_eos_token: true # eos token appears in generation -penalize_malformed_think_tag: true # /<\/think> count != 1 per turn +reward_penalties: + penalize_duplicated_reasoning: true # reasoning content == final answer + penalize_empty_final_answer: true # last message output has empty content + penalize_unwanted_tokens: true # unwanted token appears in generation + penalize_malformed_think_tag: true # /<\/think> count != 1 per turn + token_ids: + unwanted: [2] # + think_open: 12 # + think_close: 13 # # ============================================================================= # Teacher panel — paths set by the launcher. diff --git a/examples/configs/ultra/reasoning_teacher.yaml b/examples/configs/ultra/reasoning_teacher.yaml index 74fd6077d6..26e708ebc2 100644 --- a/examples/configs/ultra/reasoning_teacher.yaml +++ b/examples/configs/ultra/reasoning_teacher.yaml @@ -376,6 +376,7 @@ env: nemo_gym: nemo_gym_log_dir: "logs/nemo_gym" skip_venv_if_present: true + num_gpu_nodes: 10 # NeMo-RL-side effort shaping; stripped before forwarding to Gym. effort_levels: low_string: "{reasoning effort: efficient}" @@ -482,18 +483,15 @@ logger: collection_interval: 10 flush_interval: 10 -# ============================================================================= -# Token IDs (model-specific) -# ============================================================================= -token_ids: - eos: 2 - think_open: 12 - think_close: 13 - # ============================================================================= # Reward Penalties (set reward to 0 when triggered) # ============================================================================= -penalize_duplicated_reasoning: true -penalize_empty_final_answer: true -penalize_eos_token: true -penalize_malformed_think_tag: true +reward_penalties: + penalize_duplicated_reasoning: true + penalize_empty_final_answer: true + penalize_unwanted_tokens: true + penalize_malformed_think_tag: true + token_ids: + unwanted: [2] + think_open: 12 + think_close: 13 diff --git a/examples/configs/ultra/rlhf_teacher.yaml b/examples/configs/ultra/rlhf_teacher.yaml index d06c5659ec..2b0e0c42b8 100644 --- a/examples/configs/ultra/rlhf_teacher.yaml +++ b/examples/configs/ultra/rlhf_teacher.yaml @@ -375,6 +375,7 @@ env: nemo_gym: nemo_gym_log_dir: "logs/nemo_gym" skip_venv_if_present: true + num_gpu_nodes: 20 # NeMo-RL-side effort shaping; stripped before forwarding to Gym. effort_levels: low_string: "{reasoning effort: efficient}" @@ -502,18 +503,15 @@ logger: collection_interval: 10 flush_interval: 10 -# ============================================================================= -# Token IDs (model-specific) -# ============================================================================= -token_ids: - eos: 2 - think_open: 12 - think_close: 13 - # ============================================================================= # Reward Penalties (set reward to 0 when triggered) # ============================================================================= -penalize_duplicated_reasoning: true -penalize_empty_final_answer: true -penalize_eos_token: true -penalize_malformed_think_tag: true +reward_penalties: + penalize_duplicated_reasoning: true + penalize_empty_final_answer: true + penalize_unwanted_tokens: true + penalize_malformed_think_tag: true + token_ids: + unwanted: [2] + think_open: 12 + think_close: 13 diff --git a/examples/configs/ultra/student_rlvr1.yaml b/examples/configs/ultra/student_rlvr1.yaml index b732668896..41a207c4b8 100644 --- a/examples/configs/ultra/student_rlvr1.yaml +++ b/examples/configs/ultra/student_rlvr1.yaml @@ -372,6 +372,7 @@ env: nemo_gym: nemo_gym_log_dir: "logs/nemo_gym" skip_venv_if_present: true + num_gpu_nodes: 20 # NeMo-RL-side effort shaping; stripped before forwarding to Gym. effort_levels: low_string: "{reasoning effort: efficient}" @@ -654,18 +655,15 @@ logger: collection_interval: 10 flush_interval: 10 -# ============================================================================= -# Token IDs (model-specific) -# ============================================================================= -token_ids: - eos: 2 - think_open: 12 - think_close: 13 - # ============================================================================= # Reward Penalties (set reward to 0 when triggered) # ============================================================================= -penalize_duplicated_reasoning: true -penalize_empty_final_answer: true -penalize_eos_token: true -penalize_malformed_think_tag: true +reward_penalties: + penalize_duplicated_reasoning: true + penalize_empty_final_answer: true + penalize_unwanted_tokens: true + penalize_malformed_think_tag: true + token_ids: + unwanted: [2] + think_open: 12 + think_close: 13 diff --git a/examples/configs/ultra/student_rlvr2.yaml b/examples/configs/ultra/student_rlvr2.yaml index b4a21203e3..a791cb305e 100644 --- a/examples/configs/ultra/student_rlvr2.yaml +++ b/examples/configs/ultra/student_rlvr2.yaml @@ -373,6 +373,7 @@ env: nemo_gym: nemo_gym_log_dir: "logs/nemo_gym" skip_venv_if_present: true + num_gpu_nodes: 20 # NeMo-RL-side effort shaping; stripped before forwarding to Gym. effort_levels: low_string: "{reasoning effort: efficient}" @@ -655,18 +656,15 @@ logger: collection_interval: 10 flush_interval: 10 -# ============================================================================= -# Token IDs (model-specific) -# ============================================================================= -token_ids: - eos: 2 - think_open: 12 - think_close: 13 - # ============================================================================= # Reward Penalties (set reward to 0 when triggered) # ============================================================================= -penalize_duplicated_reasoning: true -penalize_empty_final_answer: true -penalize_eos_token: true -penalize_malformed_think_tag: true +reward_penalties: + penalize_duplicated_reasoning: true + penalize_empty_final_answer: true + penalize_unwanted_tokens: true + penalize_malformed_think_tag: true + token_ids: + unwanted: [2] + think_open: 12 + think_close: 13 diff --git a/examples/configs/ultra/swe_teacher.yaml b/examples/configs/ultra/swe_teacher.yaml index 3756753a6a..caa4fa7209 100644 --- a/examples/configs/ultra/swe_teacher.yaml +++ b/examples/configs/ultra/swe_teacher.yaml @@ -459,18 +459,15 @@ logger: collection_interval: 10 flush_interval: 10 -# ============================================================================= -# Token IDs (model-specific, used by token-based penalties) -# ============================================================================= -token_ids: - eos: 2 # - think_open: 12 # - think_close: 13 # - # ============================================================================= # Reward Penalties (set reward to 0 when triggered) # ============================================================================= -penalize_duplicated_reasoning: true # reasoning content == final answer -penalize_empty_final_answer: true # last message output has empty content -penalize_eos_token: true # eos token appears in generation -penalize_malformed_think_tag: true # /<\/think> count != 1 per turn +reward_penalties: + penalize_duplicated_reasoning: true # reasoning content == final answer + penalize_empty_final_answer: true # last message output has empty content + penalize_unwanted_tokens: true # unwanted token appears in generation + penalize_malformed_think_tag: true # /<\/think> count != 1 per turn + token_ids: + unwanted: [2] # + think_open: 12 # + think_close: 13 # From fdcc7176c8799b8c72d6ac3ea1617e884a63f108 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Mon, 6 Jul 2026 10:09:18 -0700 Subject: [PATCH 06/15] test(ultra): cover Ultra config composition Signed-off-by: Anish Mahishi --- nemo_rl/utils/config.py | 6 +++ tests/unit/utils/test_config.py | 71 +++++++++++++++++++++++++++++++-- 2 files changed, 74 insertions(+), 3 deletions(-) diff --git a/nemo_rl/utils/config.py b/nemo_rl/utils/config.py index a911e9289a..8573246e67 100644 --- a/nemo_rl/utils/config.py +++ b/nemo_rl/utils/config.py @@ -32,6 +32,10 @@ def merge_with_override( ) -> DictConfig: """Merge configs with support for _override_ marker to completely override sections.""" for key in list(override_config.keys()): + # Keep mandatory values (``???``) unresolved while composing configs. + # A child config or a CLI override may provide them after inheritance. + if OmegaConf.is_missing(override_config, key): + continue if isinstance(override_config[key], DictConfig): if override_config[key].get("_override_", False): # remove the _override_ marker @@ -190,6 +194,8 @@ def parse_hydra_overrides(cfg: DictConfig, overrides: list[str]) -> DictConfig: def register_omegaconf_resolvers() -> None: """Register shared OmegaConf resolvers used in configs.""" + if not OmegaConf.has_resolver("add"): + OmegaConf.register_new_resolver("add", lambda a, b: a + b) if not OmegaConf.has_resolver("mul"): OmegaConf.register_new_resolver("mul", lambda a, b: a * b) if not OmegaConf.has_resolver("div"): diff --git a/tests/unit/utils/test_config.py b/tests/unit/utils/test_config.py index 7c893befa3..399ee8ea6c 100644 --- a/tests/unit/utils/test_config.py +++ b/tests/unit/utils/test_config.py @@ -15,8 +15,20 @@ from pathlib import Path import pytest +from omegaconf import OmegaConf -from nemo_rl.utils.config import load_config +from nemo_rl.utils.config import load_config, register_omegaconf_resolvers + +REPO_ROOT = Path(__file__).resolve().parents[3] +ULTRA_CONFIG_PATHS = [ + "examples/configs/ultra/student_rlvr1.yaml", + "examples/configs/ultra/student_rlvr2.yaml", + "examples/configs/ultra/ifbench_teacher.yaml", + "examples/configs/ultra/reasoning_teacher.yaml", + "examples/configs/ultra/rlhf_teacher.yaml", + "examples/configs/ultra/swe_teacher.yaml", + "examples/configs/ultra/mopd.yaml", +] @pytest.fixture @@ -175,6 +187,20 @@ def test_nested_inheritance(temp_config_dir): assert config.child_only.value == 300 # Child-only value exists +def test_inheritance_preserves_missing_mandatory_value(temp_config_dir): + """Test that a mandatory parent value can be supplied after inheritance.""" + create_test_config(temp_config_dir, "parent.yaml", "required: ???") + child_path = create_test_config( + temp_config_dir, + "child.yaml", + "defaults: parent.yaml", + ) + + config = load_config(child_path) + + assert OmegaConf.is_missing(config, "required") + + def test_interpolation(temp_config_dir): """Test that interpolation works with inherited configs.""" # Create parent config @@ -198,10 +224,49 @@ def test_interpolation(temp_config_dir): assert config.derived.value == 43 # Interpolation uses child's base_value +def test_add_resolver(): + """Test the arithmetic resolver used by Ultra configs.""" + register_omegaconf_resolvers() + config = OmegaConf.create({"value": "${add:2,3}"}) + + assert config.value == 5 + + +@pytest.mark.parametrize("config_path", ULTRA_CONFIG_PATHS) +def test_ultra_configs_satisfy_current_grpo_contract(config_path): + """Ensure Ultra configs compose with all fields required by current GRPO.""" + from nemo_rl.algorithms.grpo import MasterConfig + from nemo_rl.utils.checkpoint import CheckpointManager + + register_omegaconf_resolvers() + config = load_config(REPO_ROOT / config_path) + + # These values are intentionally supplied by ultra_launch.sh at runtime. + config.policy.model_name = "test-model" + for split in ("train", "validation"): + datasets = config.data.get(split) + if datasets is None: + continue + if not OmegaConf.is_list(datasets): + datasets = [datasets] + for dataset in datasets: + if "data_path" in dataset: + dataset.data_path = "/tmp/test-data.jsonl" + + if OmegaConf.is_missing(config, "sif_dir"): + config["sif_dir"] = "/tmp/test-sifs" + if "_teachers" in config and OmegaConf.is_missing(config["_teachers"], "general"): + config["_teachers"]["general"] = "/tmp/test-teacher" + + resolved = OmegaConf.to_container(config, resolve=True) + master_config = MasterConfig.model_validate(resolved) + + assert master_config.checkpointing["save_optimizer"] is True + CheckpointManager(master_config.checkpointing) + + def test_parse_hydra_overrides(): """Test parsing and applying Hydra overrides.""" - from omegaconf import OmegaConf - from nemo_rl.utils.config import OverridesError, parse_hydra_overrides # Create initial config From 222406c6f18b557b6eb4c37ed034c0755ceaa785 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Wed, 15 Jul 2026 16:52:37 -0700 Subject: [PATCH 07/15] fix: ultra recipe fixes for main Signed-off-by: Anish Mahishi --- examples/configs/ultra/ifbench_teacher.yaml | 18 +++++- examples/configs/ultra/mopd.yaml | 29 ++++++--- examples/configs/ultra/reasoning_teacher.yaml | 11 +++- examples/configs/ultra/rlhf_teacher.yaml | 12 +++- examples/configs/ultra/student_rlvr1.yaml | 24 +++++-- examples/configs/ultra/student_rlvr2.yaml | 25 ++++++-- examples/configs/ultra/swe_teacher.yaml | 7 ++- ultra_launch.sh | 62 +++++++++++++++++-- 8 files changed, 160 insertions(+), 28 deletions(-) diff --git a/examples/configs/ultra/ifbench_teacher.yaml b/examples/configs/ultra/ifbench_teacher.yaml index 25b3a6a0b0..d1f7c0f2b8 100644 --- a/examples/configs/ultra/ifbench_teacher.yaml +++ b/examples/configs/ultra/ifbench_teacher.yaml @@ -331,8 +331,13 @@ policy: vllm_kwargs: attention_backend: FLASH_ATTN mamba_ssm_cache_dtype: "float32" + # Hybrid Mamba: one state slot per running seq; the vLLM default (1024) can + # exceed available slots. 256 matches the Ultra GB200 serving reference + # (Gym nemotron_3_ultra_dev_nemorl_gb200.yaml). + max_num_seqs: 256 compilation_config: cudagraph_capture_sizes: [1,2,4,8,16,32,64] + cudagraph_mode: PIECEWISE pass_config: fuse_allreduce_rms: false @@ -453,9 +458,12 @@ env: debug: true vllm_serve_env_vars: VLLM_RAY_DP_PACK_STRATEGY: strict - VLLM_ATTENTION_BACKEND: TRITON_ATTN + VLLM_USE_DEEP_GEMM: "0" + VLLM_MOE_USE_DEEP_GEMM: "0" vllm_serve_kwargs: + attention_backend: TRITON_ATTN + language_model_only: true tensor_parallel_size: 4 data_parallel_size: 2 data_parallel_size_local: 1 @@ -479,10 +487,14 @@ env: debug: true vllm_serve_env_vars: VLLM_RAY_DP_PACK_STRATEGY: strict - VLLM_ATTENTION_BACKEND: FLASH_ATTN - NCCL_MNNVL_ENABLE: "0" + VLLM_USE_FLASHINFER_MOE_FP16: "0" + VLLM_USE_FLASHINFER_MOE_FP8: "0" + VLLM_USE_DEEP_GEMM: "0" + VLLM_MOE_USE_DEEP_GEMM: "0" + NCCL_MNNVL_ENABLE: "1" vllm_serve_kwargs: + attention_backend: TRITON_ATTN tensor_parallel_size: 4 data_parallel_size: 4 data_parallel_size_local: 1 diff --git a/examples/configs/ultra/mopd.yaml b/examples/configs/ultra/mopd.yaml index 081ec8219b..59c8277959 100644 --- a/examples/configs/ultra/mopd.yaml +++ b/examples/configs/ultra/mopd.yaml @@ -357,8 +357,13 @@ policy: vllm_kwargs: attention_backend: FLASH_ATTN mamba_ssm_cache_dtype: "float32" + # Hybrid Mamba: one state slot per running seq; the vLLM default (1024) can + # exceed available slots. 256 matches the Ultra GB200 serving reference + # (Gym nemotron_3_ultra_dev_nemorl_gb200.yaml). + max_num_seqs: 256 compilation_config: cudagraph_capture_sizes: [1,2,4,8,16,32,64] + cudagraph_mode: PIECEWISE pass_config: fuse_allreduce_rms: false @@ -512,9 +517,12 @@ env: debug: true vllm_serve_env_vars: VLLM_RAY_DP_PACK_STRATEGY: strict - VLLM_ATTENTION_BACKEND: TRITON_ATTN + VLLM_USE_DEEP_GEMM: "0" + VLLM_MOE_USE_DEEP_GEMM: "0" vllm_serve_kwargs: + attention_backend: TRITON_ATTN + language_model_only: true tensor_parallel_size: 4 data_parallel_size: 2 data_parallel_size_local: 1 @@ -539,10 +547,14 @@ env: debug: true vllm_serve_env_vars: VLLM_RAY_DP_PACK_STRATEGY: strict - VLLM_ATTENTION_BACKEND: FLASH_ATTN - NCCL_MNNVL_ENABLE: "0" + VLLM_USE_FLASHINFER_MOE_FP16: "0" + VLLM_USE_FLASHINFER_MOE_FP8: "0" + VLLM_USE_DEEP_GEMM: "0" + VLLM_MOE_USE_DEEP_GEMM: "0" + NCCL_MNNVL_ENABLE: "1" vllm_serve_kwargs: + attention_backend: TRITON_ATTN tensor_parallel_size: 4 data_parallel_size: 4 data_parallel_size_local: 1 @@ -582,7 +594,8 @@ env: judge_responses_create_params: max_output_tokens: 8192 - # GenRM: TP=4 on GB200 192GB + # GenRM: TP=4 shards attention within each GB200 node; DP=4 extends the + # MoE expert group across all 16 GPUs while `strict` keeps each DP rank local. genrm_compare_resources_server: resources_servers: genrm_compare: @@ -621,7 +634,9 @@ env: ray_worker_py_executable: /opt/ray_venvs/nemo_rl.models.generation.vllm.vllm_worker_async.VllmAsyncGenerationWorker/bin/python vllm_serve_env_vars: VLLM_RAY_DP_PACK_STRATEGY: strict - NCCL_MNNVL_ENABLE: "0" + NCCL_MNNVL_ENABLE: "1" + VLLM_USE_DEEP_GEMM: "0" + VLLM_MOE_USE_DEEP_GEMM: "0" vllm_serve_kwargs: tensor_parallel_size: 4 @@ -629,14 +644,14 @@ env: data_parallel_size_local: 1 pipeline_parallel_size: 1 reasoning_parser: deepseek_r1 - gpu_memory_utilization: 0.85 + gpu_memory_utilization: 0.95 max_model_len: 60000 enable_prefix_caching: true enable_chunked_prefill: true model_loader_extra_config: enable_multithread_load: true num_threads: 112 - load_format: instanttensor + load_format: auto compilation_config: cudagraph_capture_sizes: [1,2,4,8,16,32] diff --git a/examples/configs/ultra/reasoning_teacher.yaml b/examples/configs/ultra/reasoning_teacher.yaml index 26e708ebc2..c406529f6e 100644 --- a/examples/configs/ultra/reasoning_teacher.yaml +++ b/examples/configs/ultra/reasoning_teacher.yaml @@ -334,8 +334,13 @@ policy: vllm_kwargs: attention_backend: FLASH_ATTN mamba_ssm_cache_dtype: "float32" + # Hybrid Mamba: one state slot per running seq; the vLLM default (1024) can + # exceed available slots. 256 matches the Ultra GB200 serving reference + # (Gym nemotron_3_ultra_dev_nemorl_gb200.yaml). + max_num_seqs: 256 compilation_config: cudagraph_capture_sizes: [1,2,4,8,16,32,64] + cudagraph_mode: PIECEWISE pass_config: fuse_allreduce_rms: false @@ -431,10 +436,14 @@ env: debug: true vllm_serve_env_vars: VLLM_RAY_DP_PACK_STRATEGY: strict - VLLM_ATTENTION_BACKEND: FLASH_ATTN NCCL_MNNVL_ENABLE: "0" + VLLM_USE_FLASHINFER_MOE_FP16: "0" + VLLM_USE_FLASHINFER_MOE_FP8: "0" + VLLM_USE_DEEP_GEMM: "0" + VLLM_MOE_USE_DEEP_GEMM: "0" vllm_serve_kwargs: + attention_backend: TRITON_ATTN tensor_parallel_size: 4 data_parallel_size: 4 data_parallel_size_local: 1 diff --git a/examples/configs/ultra/rlhf_teacher.yaml b/examples/configs/ultra/rlhf_teacher.yaml index 2b0e0c42b8..0d75545a1d 100644 --- a/examples/configs/ultra/rlhf_teacher.yaml +++ b/examples/configs/ultra/rlhf_teacher.yaml @@ -332,8 +332,13 @@ policy: vllm_kwargs: attention_backend: FLASH_ATTN mamba_ssm_cache_dtype: "float32" + # Hybrid Mamba: one state slot per running seq; the vLLM default (1024) can + # exceed available slots. 256 matches the Ultra GB200 serving reference + # (Gym nemotron_3_ultra_dev_nemorl_gb200.yaml). + max_num_seqs: 256 compilation_config: cudagraph_capture_sizes: [1,2,4,8,16,32,64] + cudagraph_mode: PIECEWISE pass_config: fuse_allreduce_rms: false @@ -461,9 +466,12 @@ env: ray_worker_py_executable: /opt/ray_venvs/nemo_rl.models.generation.vllm.vllm_worker_async.VllmAsyncGenerationWorker/bin/python vllm_serve_env_vars: VLLM_RAY_DP_PACK_STRATEGY: strict - NCCL_MNNVL_ENABLE: "0" + NCCL_MNNVL_ENABLE: "1" + VLLM_USE_DEEP_GEMM: "0" + VLLM_MOE_USE_DEEP_GEMM: "0" vllm_serve_kwargs: + attention_backend: TRITON_ATTN tensor_parallel_size: 4 data_parallel_size: 4 data_parallel_size_local: 1 @@ -476,7 +484,7 @@ env: model_loader_extra_config: enable_multithread_load: true num_threads: 112 - load_format: instanttensor + load_format: auto compilation_config: cudagraph_capture_sizes: [1,2,4,8,16,32] diff --git a/examples/configs/ultra/student_rlvr1.yaml b/examples/configs/ultra/student_rlvr1.yaml index 41a207c4b8..94c1deb4fd 100644 --- a/examples/configs/ultra/student_rlvr1.yaml +++ b/examples/configs/ultra/student_rlvr1.yaml @@ -328,8 +328,13 @@ policy: vllm_kwargs: attention_backend: FLASH_ATTN mamba_ssm_cache_dtype: "float32" + # Hybrid Mamba: one state slot per running seq; the vLLM default (1024) can + # exceed available slots. 256 matches the Ultra GB200 serving reference + # (Gym nemotron_3_ultra_dev_nemorl_gb200.yaml). + max_num_seqs: 256 compilation_config: cudagraph_capture_sizes: [1,2,4,8,16,32,64] + cudagraph_mode: PIECEWISE pass_config: fuse_allreduce_rms: false @@ -473,9 +478,11 @@ env: debug: true vllm_serve_env_vars: VLLM_RAY_DP_PACK_STRATEGY: strict - VLLM_ATTENTION_BACKEND: TRITON_ATTN + VLLM_USE_DEEP_GEMM: "0" + VLLM_MOE_USE_DEEP_GEMM: "0" vllm_serve_kwargs: + attention_backend: TRITON_ATTN tensor_parallel_size: 4 data_parallel_size: 2 data_parallel_size_local: 1 @@ -499,10 +506,14 @@ env: debug: true vllm_serve_env_vars: VLLM_RAY_DP_PACK_STRATEGY: strict - VLLM_ATTENTION_BACKEND: FLASH_ATTN - NCCL_MNNVL_ENABLE: "0" + VLLM_USE_FLASHINFER_MOE_FP16: "0" + VLLM_USE_FLASHINFER_MOE_FP8: "0" + VLLM_USE_DEEP_GEMM: "0" + VLLM_MOE_USE_DEEP_GEMM: "0" + NCCL_MNNVL_ENABLE: "1" vllm_serve_kwargs: + attention_backend: FLASH_ATTN tensor_parallel_size: 4 data_parallel_size: 4 data_parallel_size_local: 1 @@ -585,9 +596,12 @@ env: ray_worker_py_executable: /opt/ray_venvs/nemo_rl.models.generation.vllm.vllm_worker_async.VllmAsyncGenerationWorker/bin/python vllm_serve_env_vars: VLLM_RAY_DP_PACK_STRATEGY: strict - NCCL_MNNVL_ENABLE: "0" + NCCL_MNNVL_ENABLE: "1" + VLLM_USE_DEEP_GEMM: "0" + VLLM_MOE_USE_DEEP_GEMM: "0" vllm_serve_kwargs: + attention_backend: TRITON_ATTN tensor_parallel_size: 4 data_parallel_size: 4 data_parallel_size_local: 1 @@ -600,7 +614,7 @@ env: model_loader_extra_config: enable_multithread_load: true num_threads: 112 - load_format: instanttensor + load_format: auto compilation_config: cudagraph_capture_sizes: [1,2,4,8,16,32] diff --git a/examples/configs/ultra/student_rlvr2.yaml b/examples/configs/ultra/student_rlvr2.yaml index a791cb305e..89ce63b0d1 100644 --- a/examples/configs/ultra/student_rlvr2.yaml +++ b/examples/configs/ultra/student_rlvr2.yaml @@ -329,8 +329,13 @@ policy: vllm_kwargs: attention_backend: FLASH_ATTN mamba_ssm_cache_dtype: "float32" + # Hybrid Mamba: one state slot per running seq; the vLLM default (1024) can + # exceed available slots. 256 matches the Ultra GB200 serving reference + # (Gym nemotron_3_ultra_dev_nemorl_gb200.yaml). + max_num_seqs: 256 compilation_config: cudagraph_capture_sizes: [1,2,4,8,16,32,64] + cudagraph_mode: PIECEWISE pass_config: fuse_allreduce_rms: false @@ -474,9 +479,12 @@ env: debug: true vllm_serve_env_vars: VLLM_RAY_DP_PACK_STRATEGY: strict - VLLM_ATTENTION_BACKEND: TRITON_ATTN + VLLM_USE_DEEP_GEMM: "0" + VLLM_MOE_USE_DEEP_GEMM: "0" vllm_serve_kwargs: + attention_backend: TRITON_ATTN + language_model_only: true tensor_parallel_size: 4 data_parallel_size: 2 data_parallel_size_local: 1 @@ -500,10 +508,14 @@ env: debug: true vllm_serve_env_vars: VLLM_RAY_DP_PACK_STRATEGY: strict - VLLM_ATTENTION_BACKEND: FLASH_ATTN - NCCL_MNNVL_ENABLE: "0" + VLLM_USE_FLASHINFER_MOE_FP16: "0" + VLLM_USE_FLASHINFER_MOE_FP8: "0" + VLLM_USE_DEEP_GEMM: "0" + VLLM_MOE_USE_DEEP_GEMM: "0" + NCCL_MNNVL_ENABLE: "1" vllm_serve_kwargs: + attention_backend: TRITON_ATTN tensor_parallel_size: 4 data_parallel_size: 4 data_parallel_size_local: 1 @@ -586,9 +598,12 @@ env: ray_worker_py_executable: /opt/ray_venvs/nemo_rl.models.generation.vllm.vllm_worker_async.VllmAsyncGenerationWorker/bin/python vllm_serve_env_vars: VLLM_RAY_DP_PACK_STRATEGY: strict - NCCL_MNNVL_ENABLE: "0" + NCCL_MNNVL_ENABLE: "1" + VLLM_USE_DEEP_GEMM: "0" + VLLM_MOE_USE_DEEP_GEMM: "0" vllm_serve_kwargs: + attention_backend: TRITON_ATTN tensor_parallel_size: 4 data_parallel_size: 4 data_parallel_size_local: 1 @@ -601,7 +616,7 @@ env: model_loader_extra_config: enable_multithread_load: true num_threads: 112 - load_format: instanttensor + load_format: auto compilation_config: cudagraph_capture_sizes: [1,2,4,8,16,32] diff --git a/examples/configs/ultra/swe_teacher.yaml b/examples/configs/ultra/swe_teacher.yaml index caa4fa7209..cc129c192a 100644 --- a/examples/configs/ultra/swe_teacher.yaml +++ b/examples/configs/ultra/swe_teacher.yaml @@ -221,7 +221,7 @@ policy: # MTP — disabled mtp_loss_scaling_factor: 0.0 mtp_use_repeated_layer: true - mtp_num_layers: 0 + mtp_num_layers: 5 mtp_detach_heads: true optimizer: @@ -346,8 +346,13 @@ policy: vllm_kwargs: attention_backend: FLASH_ATTN mamba_ssm_cache_dtype: "float32" + # Hybrid Mamba: one state slot per running seq; the vLLM default (1024) can + # exceed available slots. 256 matches the Ultra GB200 serving reference + # (Gym nemotron_3_ultra_dev_nemorl_gb200.yaml). + max_num_seqs: 256 compilation_config: cudagraph_capture_sizes: [1,2,4,8,16,32,64] + cudagraph_mode: PIECEWISE pass_config: fuse_allreduce_rms: false diff --git a/ultra_launch.sh b/ultra_launch.sh index ef595b02bc..a77f23d782 100755 --- a/ultra_launch.sh +++ b/ultra_launch.sh @@ -33,7 +33,8 @@ set -euo pipefail # # Optional knobs: # WALLTIME=4:00:00 Slurm --time -# SLURM_QOS=normal Slurm --qos +# SLURM_QOS= Slurm --qos; defaults to short when +# WALLTIME is under two hours # SLURM_RESERVATION= Slurm --reservation # SLURM_DEPENDENCY= Extra Slurm dependency, merged with # singleton (e.g. afterany:) @@ -144,6 +145,54 @@ WALLTIME="${WALLTIME:-4:00:00}" SLURM_QOS="${SLURM_QOS:-}" SLURM_RESERVATION="${SLURM_RESERVATION:-}" EXCLUDE_NODES="${EXCLUDE_NODES:-}" + +slurm_walltime_seconds() { + local value="$1" + local days=0 + local -a fields + + if [[ "${value}" == *-* ]]; then + days="${value%%-*}" + value="${value#*-}" + fi + [[ "${days}" =~ ^[0-9]+$ ]] || return 1 + + IFS=: read -r -a fields <<< "${value}" + for field in "${fields[@]}"; do + [[ "${field}" =~ ^[0-9]+$ ]] || return 1 + done + + case "${#fields[@]}" in + 1) + if (( days > 0 )); then + echo $((10#${days} * 86400 + 10#${fields[0]} * 3600)) + else + echo $((10#${fields[0]} * 60)) + fi + ;; + 2) + if (( days > 0 )); then + echo $((10#${days} * 86400 + 10#${fields[0]} * 3600 + 10#${fields[1]} * 60)) + else + echo $((10#${fields[0]} * 60 + 10#${fields[1]})) + fi + ;; + 3) + echo $((10#${days} * 86400 + 10#${fields[0]} * 3600 + 10#${fields[1]} * 60 + 10#${fields[2]})) + ;; + *) return 1 ;; + esac +} + +if [[ -z "${SLURM_QOS}" ]]; then + if WALLTIME_SECONDS="$(slurm_walltime_seconds "${WALLTIME}")"; then + if (( WALLTIME_SECONDS < 2 * 60 * 60 )); then + SLURM_QOS=short + fi + else + echo "[WARN] Could not parse WALLTIME=${WALLTIME}; leaving SLURM_QOS unset." >&2 + fi +fi # INTERACTIVE=1 brings up the Ray cluster and idles for attachment (no training # driver), so you can run/debug the recipe by hand. INTERACTIVE_WAIT=1 (default) # blocks until Ray is ready; INTERACTIVE_WALLTIME overrides WALLTIME for the alloc. @@ -160,8 +209,9 @@ CHECKPOINTING_SAVE_BY="${CHECKPOINTING_SAVE_BY:-}" export CONTAINER MOUNTS="${MOUNTS:-}" -# GB200 NVL72: fixed at 4 GPUs/node. -export GPUS_PER_NODE=4 +# GB200 NVL72 defaults to 4 GPUs/node. Allow H100 smoke configs to request +# their native 8-GPU node shape through the launch environment. +export GPUS_PER_NODE="${GPUS_PER_NODE:-4}" export CPUS_PER_WORKER="${CPUS_PER_WORKER:-144}" # ============================================================================= @@ -582,7 +632,8 @@ NRL_VLLM_CACHE_SEED_DIR=${NRL_VLLM_CACHE_SEED_DIR} \ DG_JIT_CACHE_DIR=${NRL_VLLM_LOCAL_CACHE_DIR}/deep_gemm \ TORCHINDUCTOR_CACHE_DIR=${INDUCTOR_CACHE_DIR} \ TRITON_CACHE_DIR=${TRITON_CACHE_DIR} \ -UV_CACHE_DIR=${PERSISTENT_CACHE}/uv \ +UV_CACHE_DIR=/tmp/nemo-gym-uv-cache-\${SLURM_JOB_ID:-default} \ +UV_LOCK_TIMEOUT=1800 \ RAY_ENABLE_UV_RUN_RUNTIME_ENV=0 \ UV_HTTP_TIMEOUT=10 \ VLLM_USE_FLASHINFER_MOE_FP8=1 \ @@ -597,6 +648,7 @@ uv run ./examples/nemo_gym/run_grpo_nemo_gym.py \ policy.model_name=${MODEL_PATH} \ cluster.num_nodes=${NUM_ACTOR_NODES} \ policy.generation.colocated.resources.num_nodes=${NUM_GEN_NODES} \ +env.nemo_gym.num_gpu_nodes=${NUM_GYM_NODES} \ checkpointing.checkpoint_dir=${CHECKPOINT_DIR} \ ${CHECKPOINTING_SAVE_BY:+checkpointing.checkpoint_must_save_by=${CHECKPOINTING_SAVE_BY}} \ data.train.data_path=${TRAIN_PATH} \ @@ -719,6 +771,7 @@ if [[ "${INTERACTIVE}" == "1" ]]; then ${SLURM_QOS:+--qos="${SLURM_QOS}"} \ ${EXCLUDE_NODES:+--exclude="${EXCLUDE_NODES}"} \ ${SLURM_RESERVATION:+--reservation="${SLURM_RESERVATION}"} \ + --comment='{"OccupiedIdleGPUsJobReaper":{"exemptIdleTimeMins":"60","reason":"interactive","description":"interactive debugging"}}' \ "${RAY_SUB}") echo "${SBATCH_OUTPUT}" JOB_ID=$(echo "${SBATCH_OUTPUT}" | grep -oP '\d+$') @@ -781,6 +834,7 @@ SBATCH_OUTPUT=$(sbatch \ ${SLURM_QOS:+--qos="${SLURM_QOS}"} \ ${EXCLUDE_NODES:+--exclude="${EXCLUDE_NODES}"} \ ${SLURM_RESERVATION:+--reservation="${SLURM_RESERVATION}"} \ + --comment='{"OccupiedIdleGPUsJobReaper":{"exemptIdleTimeMins":"60","reason":"interactive","description":"interactive debugging"}}' \ "${RAY_SUB}") echo "${SBATCH_OUTPUT}" From 51493e24c92f920751097b09f05575e2d75ddbd7 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Wed, 22 Jul 2026 23:50:58 -0400 Subject: [PATCH 08/15] fix(ultra): use resolvable config placeholders and correct recipe references Address review feedback on the Ultra recipes PR: - Replace `sif_dir: ???` / `_teachers.general: ???` with concrete placeholders so example configs resolve under `OmegaConf.to_container(resolve=True)`, fixing the blocking CI regression in test_config_validation.py. Fail-fast is preserved at the launch boundary: ultra_launch.sh now hard-requires SIF_DIR whenever the selected config interpolates `${sif_dir}` (mopd already guards NRL_GENERAL_TEACHER_PATH). - Correct nonexistent `student_rlvr.yaml` references to `student_rlvr1.yaml` in the launch guide and ultra_launch.sh. - Drop the tautological `save_optimizer is True` assert in the Ultra config contract test, keeping the meaningful schema + CheckpointManager checks. - Use a non-"interactive" reaper reason for the batch training submit. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Anish Mahishi --- docs/guides/nemotron-3-ultra.md | 2 +- examples/configs/ultra/mopd.yaml | 2 +- examples/configs/ultra/swe_teacher.yaml | 10 ++++++---- tests/unit/utils/test_config.py | 5 +++-- ultra_launch.sh | 15 +++++++++++---- 5 files changed, 22 insertions(+), 12 deletions(-) diff --git a/docs/guides/nemotron-3-ultra.md b/docs/guides/nemotron-3-ultra.md index 9fc49907d6..3f873a1c25 100644 --- a/docs/guides/nemotron-3-ultra.md +++ b/docs/guides/nemotron-3-ultra.md @@ -212,7 +212,7 @@ Set the following before each `bash ultra_launch.sh` invocation: | Variable | Purpose | |---|---| | `EXP_NAME` | Job name, W&B run name, and the suffix for output directories. Must be unique per run; same name across resubmissions resumes from the latest checkpoint. | -| `CONFIG_PATH` | Path to the per-stage YAML config (e.g. `examples/configs/ultra/student_rlvr.yaml`). | +| `CONFIG_PATH` | Path to the per-stage YAML config (e.g. `examples/configs/ultra/student_rlvr1.yaml`). | | `MODEL_PATH` | Initial policy checkpoint (HuggingFace repo id or local path). Student RLVR starts from the Ultra SFT checkpoint; the teacher stages start from the Student RLVR checkpoint; MOPD starts from the student (the Student RLVR checkpoint). | | `TRAIN_PATH` | Training data JSONL. | | `VAL_PATH` | Validation data JSONL. | diff --git a/examples/configs/ultra/mopd.yaml b/examples/configs/ultra/mopd.yaml index 59c8277959..553d404af9 100644 --- a/examples/configs/ultra/mopd.yaml +++ b/examples/configs/ultra/mopd.yaml @@ -758,7 +758,7 @@ reward_penalties: # until those recipes are ready they fall back to `general` here. # ============================================================================= _teachers: - general: ??? # Required — set by launcher via NRL_GENERAL_TEACHER_PATH + general: /path/to/general_teacher # Placeholder so the config resolves; required at runtime — set by launcher via NRL_GENERAL_TEACHER_PATH rlhf: ${_teachers.general} ifbench: ${_teachers.general} reasoning: ${_teachers.general} diff --git a/examples/configs/ultra/swe_teacher.yaml b/examples/configs/ultra/swe_teacher.yaml index cc129c192a..994973a7cb 100644 --- a/examples/configs/ultra/swe_teacher.yaml +++ b/examples/configs/ultra/swe_teacher.yaml @@ -36,11 +36,13 @@ cluster: segment_size: 16 # ============================================================================= -# SIF image root — required, set by the launcher via SIF_DIR. The SWE agent's -# container_formatter (in env.nemo_gym.swe_agents_train) resolves -# `${sif_dir}//{instance_id}.sif` per SWE benchmark family. +# SIF image root — required at runtime, set by the launcher via SIF_DIR. The SWE +# agent's container_formatter (in env.nemo_gym.swe_agents_train) resolves +# `${sif_dir}//{instance_id}.sif` per SWE benchmark family. The value +# below is a placeholder so the config resolves at load/validation time; +# ultra_launch.sh hard-requires SIF_DIR and overrides it for real runs. # ============================================================================= -sif_dir: ??? +sif_dir: /path/to/sif_images # ============================================================================= # Checkpointing diff --git a/tests/unit/utils/test_config.py b/tests/unit/utils/test_config.py index 399ee8ea6c..913a56a790 100644 --- a/tests/unit/utils/test_config.py +++ b/tests/unit/utils/test_config.py @@ -259,9 +259,10 @@ def test_ultra_configs_satisfy_current_grpo_contract(config_path): config["_teachers"]["general"] = "/tmp/test-teacher" resolved = OmegaConf.to_container(config, resolve=True) - master_config = MasterConfig.model_validate(resolved) - assert master_config.checkpointing["save_optimizer"] is True + # The real contract checks: the config validates against GRPO's MasterConfig + # schema and the checkpointing block is accepted by CheckpointManager. + master_config = MasterConfig.model_validate(resolved) CheckpointManager(master_config.checkpointing) diff --git a/ultra_launch.sh b/ultra_launch.sh index a77f23d782..2bfd4c9595 100755 --- a/ultra_launch.sh +++ b/ultra_launch.sh @@ -17,7 +17,7 @@ set -euo pipefail # Usage: # # EXP_NAME=ultra-student-rlvr-001 \ -# CONFIG_PATH=examples/configs/ultra/student_rlvr.yaml \ +# CONFIG_PATH=examples/configs/ultra/student_rlvr1.yaml \ # MODEL_PATH=/path/to/sft_checkpoint \ # TRAIN_PATH=/path/to/train.jsonl \ # VAL_PATH=/path/to/val.jsonl \ @@ -72,7 +72,7 @@ set -euo pipefail # Required environment # ============================================================================= : "${EXP_NAME:?EXP_NAME is required (used for job name, W&B run, checkpoint/log dirs)}" -: "${CONFIG_PATH:?CONFIG_PATH is required (e.g. examples/configs/ultra/student_rlvr.yaml)}" +: "${CONFIG_PATH:?CONFIG_PATH is required (e.g. examples/configs/ultra/student_rlvr1.yaml)}" : "${MODEL_PATH:?MODEL_PATH is required (initial policy checkpoint, HF repo id or local path)}" : "${TRAIN_PATH:?TRAIN_PATH is required (training data jsonl path)}" : "${VAL_PATH:?VAL_PATH is required (validation data jsonl path)}" @@ -106,6 +106,13 @@ if [[ ! -f "${CONFIG_PATH}" ]]; then exit 1 fi +# The SWE teacher recipe interpolates `${sif_dir}/...` paths at runtime. The +# exemplar config carries only a placeholder, so hard-require SIF_DIR whenever +# the selected config actually uses it (mirrors the mopd teacher-path guard). +if grep -q '${sif_dir}' "${CONFIG_PATH}"; then + : "${SIF_DIR:?SIF_DIR is required for the SWE recipe (directory of apptainer .sif images)}" +fi + # ============================================================================= # Project root and code root # ============================================================================= @@ -266,7 +273,7 @@ if [[ "${ENABLE_MTP_INFERENCE}" == "1" ]]; then fi # ============================================================================= -# Job shape — defaults match the 256-node student_rlvr.yaml +# Job shape — defaults match the 256-node student_rlvr1.yaml # # Training: 64 nodes ( 256 GPUs) — Megatron training backend # vLLM: 172 nodes ( 688 GPUs) — async generation, EP=8 instances at TP=8 @@ -834,7 +841,7 @@ SBATCH_OUTPUT=$(sbatch \ ${SLURM_QOS:+--qos="${SLURM_QOS}"} \ ${EXCLUDE_NODES:+--exclude="${EXCLUDE_NODES}"} \ ${SLURM_RESERVATION:+--reservation="${SLURM_RESERVATION}"} \ - --comment='{"OccupiedIdleGPUsJobReaper":{"exemptIdleTimeMins":"60","reason":"interactive","description":"interactive debugging"}}' \ + --comment='{"OccupiedIdleGPUsJobReaper":{"exemptIdleTimeMins":"60","reason":"batch","description":"batch training run"}}' \ "${RAY_SUB}") echo "${SBATCH_OUTPUT}" From d337d77b6bc2839fe7096609563ba171c4b21ac7 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Fri, 24 Jul 2026 21:02:31 -0700 Subject: [PATCH 09/15] fix: fix genrm related configs Signed-off-by: Anish Mahishi --- examples/configs/ultra/mopd.yaml | 5 ++++ examples/configs/ultra/rlhf_teacher.yaml | 34 +++++++++++++++-------- examples/configs/ultra/student_rlvr2.yaml | 31 ++++++++++++++------- 3 files changed, 49 insertions(+), 21 deletions(-) diff --git a/examples/configs/ultra/mopd.yaml b/examples/configs/ultra/mopd.yaml index 553d404af9..d500beff38 100644 --- a/examples/configs/ultra/mopd.yaml +++ b/examples/configs/ultra/mopd.yaml @@ -637,6 +637,11 @@ env: NCCL_MNNVL_ENABLE: "1" VLLM_USE_DEEP_GEMM: "0" VLLM_MOE_USE_DEEP_GEMM: "0" + # # Qwen-235B expert dims aren't FlashInfer-TRTLLM-MoE-compatible at TP=4 + # # (K % blockK != 0). Disable the FlashInfer MoE path (inherited on from + # # the launcher's VLLM_USE_FLASHINFER_MOE_FP8=1), matching nl2bash_judge_model. + VLLM_USE_FLASHINFER_MOE_FP8: "0" + VLLM_USE_FLASHINFER_MOE_FP16: "0" vllm_serve_kwargs: tensor_parallel_size: 4 diff --git a/examples/configs/ultra/rlhf_teacher.yaml b/examples/configs/ultra/rlhf_teacher.yaml index 0d75545a1d..8d507a8238 100644 --- a/examples/configs/ultra/rlhf_teacher.yaml +++ b/examples/configs/ultra/rlhf_teacher.yaml @@ -452,8 +452,9 @@ env: default_score: 3.0 default_ranking: 3.5 - # Launcher sets either `base_url` (route to remote GenRM service) or - # `model` (serve locally using vllm_serve_kwargs below). See the guide. + # Gym-managed native-DP GenRM. The companion launcher sets `model` to the + # same checkpoint used by the external-service control. TP=4 x DP=16 uses + # sixteen four-GPU nodes, with each replica contained within one node. genrm_model: responses_api_models: genrm_model: @@ -465,28 +466,39 @@ env: debug: true ray_worker_py_executable: /opt/ray_venvs/nemo_rl.models.generation.vllm.vllm_worker_async.VllmAsyncGenerationWorker/bin/python vllm_serve_env_vars: + # Keep each TP=4 replica on one four-GPU GB200 node. VLLM_RAY_DP_PACK_STRATEGY: strict NCCL_MNNVL_ENABLE: "1" - VLLM_USE_DEEP_GEMM: "0" - VLLM_MOE_USE_DEEP_GEMM: "0" + VLLM_FLASHINFER_ALLREDUCE_BACKEND: trtllm + VLLM_ALLREDUCE_USE_SYMM_MEM: "0" vllm_serve_kwargs: - attention_backend: TRITON_ATTN tensor_parallel_size: 4 - data_parallel_size: 4 + data_parallel_size: 16 data_parallel_size_local: 1 pipeline_parallel_size: 1 - reasoning_parser: deepseek_r1 - gpu_memory_utilization: 0.85 - max_model_len: 131072 + distributed_executor_backend: ray + data_parallel_backend: ray + api_server_count: 1 + enable_expert_parallel: false + dtype: bfloat16 + kv_cache_dtype: fp8 + gpu_memory_utilization: 0.95 + max_model_len: 262144 + max_num_seqs: 256 + max_num_batched_tokens: 8192 enable_prefix_caching: true enable_chunked_prefill: true + trust_remote_code: true + enable_auto_tool_choice: true + tool_call_parser: qwen3_coder model_loader_extra_config: enable_multithread_load: true - num_threads: 112 + num_threads: 96 load_format: auto compilation_config: - cudagraph_capture_sizes: [1,2,4,8,16,32] + pass_config: + fuse_allreduce_rms: false # ============================================================================= diff --git a/examples/configs/ultra/student_rlvr2.yaml b/examples/configs/ultra/student_rlvr2.yaml index 89ce63b0d1..48fba94c72 100644 --- a/examples/configs/ultra/student_rlvr2.yaml +++ b/examples/configs/ultra/student_rlvr2.yaml @@ -517,7 +517,7 @@ env: vllm_serve_kwargs: attention_backend: TRITON_ATTN tensor_parallel_size: 4 - data_parallel_size: 4 + data_parallel_size: 2 data_parallel_size_local: 1 pipeline_parallel_size: 1 enable_expert_parallel: true @@ -597,28 +597,39 @@ env: debug: true ray_worker_py_executable: /opt/ray_venvs/nemo_rl.models.generation.vllm.vllm_worker_async.VllmAsyncGenerationWorker/bin/python vllm_serve_env_vars: + # Keep each TP=4 replica on one four-GPU GB200 node. VLLM_RAY_DP_PACK_STRATEGY: strict NCCL_MNNVL_ENABLE: "1" - VLLM_USE_DEEP_GEMM: "0" - VLLM_MOE_USE_DEEP_GEMM: "0" + VLLM_FLASHINFER_ALLREDUCE_BACKEND: trtllm + VLLM_ALLREDUCE_USE_SYMM_MEM: "0" vllm_serve_kwargs: - attention_backend: TRITON_ATTN tensor_parallel_size: 4 - data_parallel_size: 4 + data_parallel_size: 16 data_parallel_size_local: 1 pipeline_parallel_size: 1 - reasoning_parser: deepseek_r1 - gpu_memory_utilization: 0.85 - max_model_len: 131072 + distributed_executor_backend: ray + data_parallel_backend: ray + api_server_count: 1 + enable_expert_parallel: false + dtype: bfloat16 + kv_cache_dtype: fp8 + gpu_memory_utilization: 0.95 + max_model_len: 262144 + max_num_seqs: 256 + max_num_batched_tokens: 8192 enable_prefix_caching: true enable_chunked_prefill: true + trust_remote_code: true + enable_auto_tool_choice: true + tool_call_parser: qwen3_coder model_loader_extra_config: enable_multithread_load: true - num_threads: 112 + num_threads: 96 load_format: auto compilation_config: - cudagraph_capture_sizes: [1,2,4,8,16,32] + pass_config: + fuse_allreduce_rms: false lc_judge: resources_servers: From 46263ecf642551d2095cba2c38602e8e49df23bf Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Tue, 28 Jul 2026 18:26:41 -0400 Subject: [PATCH 10/15] refactor(ultra): relocate Ultra recipes under examples/nemo_gym/nemotron-3-ultra Move the 7 Ultra recipe YAMLs and ultra_launch.sh from examples/configs/ultra/ and the repo root into examples/nemo_gym/nemotron-3-ultra/, mirroring the existing nemotron-3-super layout (recipe configs + launcher co-located in one nemo_gym recipe dir). Update all path references so everything resolves: - tests/unit/utils/test_config.py: repoint ULTRA_CONFIG_PATHS to the new dir (otherwise test_ultra_configs_satisfy_current_grpo_contract errors on missing files). - ultra_launch.sh: update usage comments, the CONFIG_PATH example, and the invocation hints to the new script path (run from the repo root, matching super_launch.sh). - docs/guides/nemotron-3-ultra.md: update config paths and the launch command. The configs now sit outside examples/configs/**, so test_config_validation.py no longer globs them (same as the super recipe); the Ultra contract test still covers them via ULTRA_CONFIG_PATHS. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Anish Mahishi --- docs/guides/nemotron-3-ultra.md | 51 ++++++++++--------- .../nemotron-3-ultra}/ifbench_teacher.yaml | 0 .../nemotron-3-ultra}/mopd.yaml | 0 .../nemotron-3-ultra}/reasoning_teacher.yaml | 0 .../nemotron-3-ultra}/rlhf_teacher.yaml | 0 .../nemotron-3-ultra}/student_rlvr1.yaml | 0 .../nemotron-3-ultra}/student_rlvr2.yaml | 0 .../nemotron-3-ultra}/swe_teacher.yaml | 0 .../nemo_gym/nemotron-3-ultra/ultra_launch.sh | 12 ++--- tests/unit/utils/test_config.py | 14 ++--- 10 files changed, 39 insertions(+), 38 deletions(-) rename examples/{configs/ultra => nemo_gym/nemotron-3-ultra}/ifbench_teacher.yaml (100%) rename examples/{configs/ultra => nemo_gym/nemotron-3-ultra}/mopd.yaml (100%) rename examples/{configs/ultra => nemo_gym/nemotron-3-ultra}/reasoning_teacher.yaml (100%) rename examples/{configs/ultra => nemo_gym/nemotron-3-ultra}/rlhf_teacher.yaml (100%) rename examples/{configs/ultra => nemo_gym/nemotron-3-ultra}/student_rlvr1.yaml (100%) rename examples/{configs/ultra => nemo_gym/nemotron-3-ultra}/student_rlvr2.yaml (100%) rename examples/{configs/ultra => nemo_gym/nemotron-3-ultra}/swe_teacher.yaml (100%) rename ultra_launch.sh => examples/nemo_gym/nemotron-3-ultra/ultra_launch.sh (98%) diff --git a/docs/guides/nemotron-3-ultra.md b/docs/guides/nemotron-3-ultra.md index 3f873a1c25..59c240fb1e 100644 --- a/docs/guides/nemotron-3-ultra.md +++ b/docs/guides/nemotron-3-ultra.md @@ -20,7 +20,7 @@ main stages are: 3. **MOPD** — on-policy distillation from the student into the teacher panel. Every stage shares the same launcher (`ultra_launch.sh`) and a per-stage YAML -config under `examples/configs/ultra/`. +config under `examples/nemo_gym/nemotron-3-ultra/`. ### Checkpoint flow @@ -202,17 +202,18 @@ enroot import -o nemo-skills-sandbox.sqsh dockerd://nemo-skills-sandbox:latest ## Launch script -Every stage is submitted with `ultra_launch.sh` at the repo root. The launcher +Every stage is submitted with `examples/nemo_gym/nemotron-3-ultra/ultra_launch.sh`, +run from the repo root. The launcher handles SLURM submission, code snapshotting, persistent cache management, and container mounts — stage-specific hyperparameters (batch size, advantage clip, MoE parallelism, learning rate) live in the per-stage YAML. -Set the following before each `bash ultra_launch.sh` invocation: +Set the following before each `bash examples/nemo_gym/nemotron-3-ultra/ultra_launch.sh` invocation: | Variable | Purpose | |---|---| | `EXP_NAME` | Job name, W&B run name, and the suffix for output directories. Must be unique per run; same name across resubmissions resumes from the latest checkpoint. | -| `CONFIG_PATH` | Path to the per-stage YAML config (e.g. `examples/configs/ultra/student_rlvr1.yaml`). | +| `CONFIG_PATH` | Path to the per-stage YAML config (e.g. `examples/nemo_gym/nemotron-3-ultra/student_rlvr1.yaml`). | | `MODEL_PATH` | Initial policy checkpoint (HuggingFace repo id or local path). Student RLVR starts from the Ultra SFT checkpoint; the teacher stages start from the Student RLVR checkpoint; MOPD starts from the student (the Student RLVR checkpoint). | | `TRAIN_PATH` | Training data JSONL. | | `VAL_PATH` | Validation data JSONL. | @@ -249,7 +250,7 @@ second phase resumes from the first phase's checkpoint: | | Phase 1 | Phase 2 | |---|---|---| -| Config | `examples/configs/ultra/student_rlvr1.yaml` | `examples/configs/ultra/student_rlvr2.yaml` | +| Config | `examples/nemo_gym/nemotron-3-ultra/student_rlvr1.yaml` | `examples/nemo_gym/nemotron-3-ultra/student_rlvr2.yaml` | | `max_total_sequence_length` | 49,152 | 65,536 | | Steps in this phase | ~128 | ~50 | | `NRL_MAX_STEPS` to set | `128` | `178` (= 128 + 50) | @@ -262,7 +263,7 @@ generations), advantage clip ±20, and the 256-node cluster shape. ```bash EXP_NAME=ultra-student-rlvr \ -CONFIG_PATH=examples/configs/ultra/student_rlvr1.yaml \ +CONFIG_PATH=examples/nemo_gym/nemotron-3-ultra/student_rlvr1.yaml \ ENABLE_MTP_INFERENCE=1 \ NRL_MAX_STEPS=128 \ MODEL_PATH=/path/to/ultra_sft_checkpoint \ @@ -280,7 +281,7 @@ SAFETY_JUDGE_MODEL=nvidia/Nemotron-Content-Safety-Reasoning-4B \ WANDB_API_KEY=$WANDB_API_KEY \ HF_HOME=/path/to/hf_cache \ HF_TOKEN=$HF_TOKEN \ -bash ultra_launch.sh +bash examples/nemo_gym/nemotron-3-ultra/ultra_launch.sh ``` ### Phase 2 — 65k context, ~50 more steps @@ -292,7 +293,7 @@ checkpoint. ```bash EXP_NAME=ultra-student-rlvr \ -CONFIG_PATH=examples/configs/ultra/student_rlvr2.yaml \ +CONFIG_PATH=examples/nemo_gym/nemotron-3-ultra/student_rlvr2.yaml \ ENABLE_MTP_INFERENCE=1 \ NRL_MAX_STEPS=178 \ MODEL_PATH=/path/to/ultra_sft_checkpoint \ @@ -310,7 +311,7 @@ SAFETY_JUDGE_MODEL=nvidia/Nemotron-Content-Safety-Reasoning-4B \ WANDB_API_KEY=$WANDB_API_KEY \ HF_HOME=/path/to/hf_cache \ HF_TOKEN=$HF_TOKEN \ -bash ultra_launch.sh +bash examples/nemo_gym/nemotron-3-ultra/ultra_launch.sh ``` Note: Phase 1 and Phase 2 use different training blends (`rlvr1.train.jsonl` and @@ -324,7 +325,7 @@ commands, and (on submission) the SLURM job id. The teacher panel is a set of specialised RL runs that each start from the Student RLVR output (Stage 1) (see [Checkpoint flow](#checkpoint-flow) above). Each teacher -runs independently with its own YAML config under `examples/configs/ultra/`. +runs independently with its own YAML config under `examples/nemo_gym/nemotron-3-ultra/`. The teachers don't depend on each other and can run in parallel. @@ -334,7 +335,7 @@ RLHF teacher specializing in instruction following, abstention, and refusal behavior. Trained at 49k context with a smaller batch (`GBS=2048`) and lower learning rate (`lr=2.5e-6`) than the student RLVR stage. -**Config:** `examples/configs/ultra/ifbench_teacher.yaml` +**Config:** `examples/nemo_gym/nemotron-3-ultra/ifbench_teacher.yaml` - TP=8, EP=64, CP=8, PP=1 - `max_total_sequence_length=49152` - `train_global_batch_size=2048`, `num_prompts_per_step=128`, `num_generations_per_prompt=16` @@ -343,7 +344,7 @@ learning rate (`lr=2.5e-6`) than the student RLVR stage. ```bash EXP_NAME=ultra-ifbench-teacher \ -CONFIG_PATH=examples/configs/ultra/ifbench_teacher.yaml \ +CONFIG_PATH=examples/nemo_gym/nemotron-3-ultra/ifbench_teacher.yaml \ MODEL_PATH=/path/to/student_rlvr_output \ TRAIN_PATH=$DATA_DIR/ifbench.train.jsonl \ VAL_PATH=$DATA_DIR/ifbench.val.jsonl \ @@ -363,7 +364,7 @@ SAFETY_JUDGE_MODEL=nvidia/Nemotron-Content-Safety-Reasoning-4B \ WANDB_API_KEY=$WANDB_API_KEY \ HF_HOME=/path/to/hf_cache \ HF_TOKEN=$HF_TOKEN \ -bash ultra_launch.sh +bash examples/nemo_gym/nemotron-3-ultra/ultra_launch.sh ``` ### RLHF Teacher @@ -372,7 +373,7 @@ General-purpose RLHF teacher trained against the pairwise GenRM comparison signal alone. Same training shape as the IFBench teacher (cluster, batch, learning rate, context). -**Config:** `examples/configs/ultra/rlhf_teacher.yaml` +**Config:** `examples/nemo_gym/nemotron-3-ultra/rlhf_teacher.yaml` - TP=8, EP=64, CP=8, PP=1 - `max_total_sequence_length=49152` - `train_global_batch_size=2048`, `num_prompts_per_step=128`, `num_generations_per_prompt=16` @@ -381,7 +382,7 @@ learning rate, context). ```bash EXP_NAME=ultra-rlhf-teacher \ -CONFIG_PATH=examples/configs/ultra/rlhf_teacher.yaml \ +CONFIG_PATH=examples/nemo_gym/nemotron-3-ultra/rlhf_teacher.yaml \ MODEL_PATH=/path/to/student_rlvr_output \ TRAIN_PATH=$DATA_DIR/rlhf.train.jsonl \ VAL_PATH=$DATA_DIR/rlhf.val.jsonl \ @@ -399,7 +400,7 @@ GENRM_MODEL=nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-GenRM \ WANDB_API_KEY=$WANDB_API_KEY \ HF_HOME=/path/to/hf_cache \ HF_TOKEN=$HF_TOKEN \ -bash ultra_launch.sh +bash examples/nemo_gym/nemotron-3-ultra/ultra_launch.sh ``` This is the teacher referred to as `NRL_CHAT_TEACHER1` / `NRL_RLHF_TEACHER` in @@ -415,7 +416,7 @@ serves the `code_gen`, `ns_tools`, `math_with_judge`, `equivalence_llm_judge`, and `mcqa` agent slots in MOPD — one checkpoint, many roles. -**Config:** `examples/configs/ultra/reasoning_teacher.yaml` +**Config:** `examples/nemo_gym/nemotron-3-ultra/reasoning_teacher.yaml` - TP=8, EP=32, CP=8, PP=1 (half the expert parallelism of Student RLVR) - `max_total_sequence_length=65536` - `train_global_batch_size=2048`, `num_prompts_per_step=128`, `num_generations_per_prompt=16` @@ -425,7 +426,7 @@ many roles. ```bash EXP_NAME=ultra-reasoning-teacher \ -CONFIG_PATH=examples/configs/ultra/reasoning_teacher.yaml \ +CONFIG_PATH=examples/nemo_gym/nemotron-3-ultra/reasoning_teacher.yaml \ ENABLE_MTP_INFERENCE=1 \ MODEL_PATH=/path/to/student_rlvr_output \ TRAIN_PATH=$DATA_DIR/reasoning.train.jsonl \ @@ -443,7 +444,7 @@ NL2BASH_JUDGE_MODEL=Qwen/Qwen3-235B-A22B-Instruct-2507-FP8 \ WANDB_API_KEY=$WANDB_API_KEY \ HF_HOME=/path/to/hf_cache \ HF_TOKEN=$HF_TOKEN \ -bash ultra_launch.sh +bash examples/nemo_gym/nemotron-3-ultra/ultra_launch.sh ``` ### SWE Teacher @@ -453,7 +454,7 @@ Software-engineering RLVR teacher trained against code execution. The container images for each SWE-Gym / SWE-rebench-V2 instance and rewards the rollout based on test pass/fail. -**Config:** `examples/configs/ultra/swe_teacher.yaml` +**Config:** `examples/nemo_gym/nemotron-3-ultra/swe_teacher.yaml` - TP=8, EP=32, CP=32, PP=1 (large CP for long context) - `max_total_sequence_length=196608` (192k context) - `train_global_batch_size=512`, `num_prompts_per_step=32`, `num_generations_per_prompt=16` @@ -552,7 +553,7 @@ With `${SIF_DIR}` populated, launch the SWE teacher: ```bash EXP_NAME=ultra-swe-teacher \ -CONFIG_PATH=examples/configs/ultra/swe_teacher.yaml \ +CONFIG_PATH=examples/nemo_gym/nemotron-3-ultra/swe_teacher.yaml \ ENABLE_MTP_INFERENCE=1 \ MODEL_PATH=/path/to/student_rlvr_output \ TRAIN_PATH=$DATA_DIR/swe.train.jsonl \ @@ -570,7 +571,7 @@ SLURM_ACCOUNT=$SLURM_ACCOUNT \ WANDB_API_KEY=$WANDB_API_KEY \ HF_HOME=/path/to/hf_cache \ HF_TOKEN=$HF_TOKEN \ -bash ultra_launch.sh +bash examples/nemo_gym/nemotron-3-ultra/ultra_launch.sh ``` ## Stage 3 — MOPD @@ -579,7 +580,7 @@ Multi-Teacher On-Policy Distillation. The Student RLVR output is the student; ea Gym agent is routed to one of the Stage 2 teacher checkpoints. Trains the student to match per-agent teacher distributions. -**Config:** `examples/configs/ultra/mopd.yaml` +**Config:** `examples/nemo_gym/nemotron-3-ultra/mopd.yaml` - TP=8, EP=64, CP=8, PP=1, max context 192k - Teacher parallelism: TP=8, CP=2, EP=16, 4 nodes per teacher - Routing: agent → teacher checkpoint baked into the YAML via @@ -607,7 +608,7 @@ unset. ```bash STAGE_TYPE=mopd \ EXP_NAME=ultra-mopd-stage1 \ -CONFIG_PATH=examples/configs/ultra/mopd.yaml \ +CONFIG_PATH=examples/nemo_gym/nemotron-3-ultra/mopd.yaml \ ENABLE_MTP_INFERENCE=1 \ MODEL_PATH=/path/to/student_rlvr_output \ TRAIN_PATH=$DATA_DIR/mopd.train.jsonl \ @@ -634,5 +635,5 @@ SAFETY_JUDGE_MODEL=nvidia/Nemotron-Content-Safety-Reasoning-4B \ WANDB_API_KEY=$WANDB_API_KEY \ HF_HOME=/path/to/hf_cache \ HF_TOKEN=$HF_TOKEN \ -bash ultra_launch.sh +bash examples/nemo_gym/nemotron-3-ultra/ultra_launch.sh ``` diff --git a/examples/configs/ultra/ifbench_teacher.yaml b/examples/nemo_gym/nemotron-3-ultra/ifbench_teacher.yaml similarity index 100% rename from examples/configs/ultra/ifbench_teacher.yaml rename to examples/nemo_gym/nemotron-3-ultra/ifbench_teacher.yaml diff --git a/examples/configs/ultra/mopd.yaml b/examples/nemo_gym/nemotron-3-ultra/mopd.yaml similarity index 100% rename from examples/configs/ultra/mopd.yaml rename to examples/nemo_gym/nemotron-3-ultra/mopd.yaml diff --git a/examples/configs/ultra/reasoning_teacher.yaml b/examples/nemo_gym/nemotron-3-ultra/reasoning_teacher.yaml similarity index 100% rename from examples/configs/ultra/reasoning_teacher.yaml rename to examples/nemo_gym/nemotron-3-ultra/reasoning_teacher.yaml diff --git a/examples/configs/ultra/rlhf_teacher.yaml b/examples/nemo_gym/nemotron-3-ultra/rlhf_teacher.yaml similarity index 100% rename from examples/configs/ultra/rlhf_teacher.yaml rename to examples/nemo_gym/nemotron-3-ultra/rlhf_teacher.yaml diff --git a/examples/configs/ultra/student_rlvr1.yaml b/examples/nemo_gym/nemotron-3-ultra/student_rlvr1.yaml similarity index 100% rename from examples/configs/ultra/student_rlvr1.yaml rename to examples/nemo_gym/nemotron-3-ultra/student_rlvr1.yaml diff --git a/examples/configs/ultra/student_rlvr2.yaml b/examples/nemo_gym/nemotron-3-ultra/student_rlvr2.yaml similarity index 100% rename from examples/configs/ultra/student_rlvr2.yaml rename to examples/nemo_gym/nemotron-3-ultra/student_rlvr2.yaml diff --git a/examples/configs/ultra/swe_teacher.yaml b/examples/nemo_gym/nemotron-3-ultra/swe_teacher.yaml similarity index 100% rename from examples/configs/ultra/swe_teacher.yaml rename to examples/nemo_gym/nemotron-3-ultra/swe_teacher.yaml diff --git a/ultra_launch.sh b/examples/nemo_gym/nemotron-3-ultra/ultra_launch.sh similarity index 98% rename from ultra_launch.sh rename to examples/nemo_gym/nemotron-3-ultra/ultra_launch.sh index 2bfd4c9595..3d4b9db6a0 100755 --- a/ultra_launch.sh +++ b/examples/nemo_gym/nemotron-3-ultra/ultra_launch.sh @@ -7,7 +7,7 @@ set -euo pipefail # Public launcher for Nemotron 3 Ultra post-training stages on a SLURM cluster. # # Each training stage (Student RLVR, teacher RLVR/RLHF stages, MOPD) has a -# matching YAML config under examples/configs/ultra/. The stage-specific +# matching YAML config under examples/nemo_gym/nemotron-3-ultra/. The stage-specific # hyperparameters (batch size, advantage clip, MoE parallelism, etc.) live # in the YAML; this launcher only handles orchestration: SLURM submission, # code snapshotting, persistent cache management, container mounts, and the @@ -17,7 +17,7 @@ set -euo pipefail # Usage: # # EXP_NAME=ultra-student-rlvr-001 \ -# CONFIG_PATH=examples/configs/ultra/student_rlvr1.yaml \ +# CONFIG_PATH=examples/nemo_gym/nemotron-3-ultra/student_rlvr1.yaml \ # MODEL_PATH=/path/to/sft_checkpoint \ # TRAIN_PATH=/path/to/train.jsonl \ # VAL_PATH=/path/to/val.jsonl \ @@ -29,7 +29,7 @@ set -euo pipefail # GENRM_MODEL= # Or set GENRM_BASE_URL to use a remote service # NL2BASH_JUDGE_MODEL=Qwen/Qwen3-235B-A22B-Instruct-2507-FP8 \ # SAFETY_JUDGE_MODEL=/path/to/safety_checkpoint \ -# bash ultra_launch.sh +# bash examples/nemo_gym/nemotron-3-ultra/ultra_launch.sh # # Optional knobs: # WALLTIME=4:00:00 Slurm --time @@ -62,7 +62,7 @@ set -euo pipefail # WANDB_ENTITY= W&B entity # # Hydra overrides are forwarded verbatim as positional arguments: -# bash ultra_launch.sh policy.megatron_cfg.optimizer.lr=1e-6 grpo.val_period=50 +# bash examples/nemo_gym/nemotron-3-ultra/ultra_launch.sh policy.megatron_cfg.optimizer.lr=1e-6 grpo.val_period=50 # # GB200 NVL72 nodes have 4 GPUs each. SLURM total = NUM_TRAIN + NUM_GEN + NUM_GYM # and must be a multiple of SEGMENT_SIZE (default 16, one NVLink domain group). @@ -72,7 +72,7 @@ set -euo pipefail # Required environment # ============================================================================= : "${EXP_NAME:?EXP_NAME is required (used for job name, W&B run, checkpoint/log dirs)}" -: "${CONFIG_PATH:?CONFIG_PATH is required (e.g. examples/configs/ultra/student_rlvr1.yaml)}" +: "${CONFIG_PATH:?CONFIG_PATH is required (e.g. examples/nemo_gym/nemotron-3-ultra/student_rlvr1.yaml)}" : "${MODEL_PATH:?MODEL_PATH is required (initial policy checkpoint, HF repo id or local path)}" : "${TRAIN_PATH:?TRAIN_PATH is required (training data jsonl path)}" : "${VAL_PATH:?VAL_PATH is required (validation data jsonl path)}" @@ -516,7 +516,7 @@ fi # To overlay additional components (e.g. a local Megatron-LM checkout), pass # EXTRA_MOUNTS as a comma-separated list of host:container pairs: # -# EXTRA_MOUNTS="/path/to/Megatron-LM:/opt/nemo-rl/3rdparty/Megatron-LM-workspace/Megatron-LM" bash ultra_launch.sh +# EXTRA_MOUNTS="/path/to/Megatron-LM:/opt/nemo-rl/3rdparty/Megatron-LM-workspace/Megatron-LM" bash examples/nemo_gym/nemotron-3-ultra/ultra_launch.sh # # Container paths for reference: # /opt/nemo-rl/nemo_rl — Python package diff --git a/tests/unit/utils/test_config.py b/tests/unit/utils/test_config.py index 913a56a790..9d82e89b30 100644 --- a/tests/unit/utils/test_config.py +++ b/tests/unit/utils/test_config.py @@ -21,13 +21,13 @@ REPO_ROOT = Path(__file__).resolve().parents[3] ULTRA_CONFIG_PATHS = [ - "examples/configs/ultra/student_rlvr1.yaml", - "examples/configs/ultra/student_rlvr2.yaml", - "examples/configs/ultra/ifbench_teacher.yaml", - "examples/configs/ultra/reasoning_teacher.yaml", - "examples/configs/ultra/rlhf_teacher.yaml", - "examples/configs/ultra/swe_teacher.yaml", - "examples/configs/ultra/mopd.yaml", + "examples/nemo_gym/nemotron-3-ultra/student_rlvr1.yaml", + "examples/nemo_gym/nemotron-3-ultra/student_rlvr2.yaml", + "examples/nemo_gym/nemotron-3-ultra/ifbench_teacher.yaml", + "examples/nemo_gym/nemotron-3-ultra/reasoning_teacher.yaml", + "examples/nemo_gym/nemotron-3-ultra/rlhf_teacher.yaml", + "examples/nemo_gym/nemotron-3-ultra/swe_teacher.yaml", + "examples/nemo_gym/nemotron-3-ultra/mopd.yaml", ] From 21c6972232d2ecf45945729404e8895d16b58179 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Wed, 29 Jul 2026 15:41:46 -0400 Subject: [PATCH 11/15] chore(ultra): tune vLLM generation/judge settings and reaper reason Runtime tuning for the GB200 Ultra recipes: - Point judge models (genrm/nl2bash) at the async vLLM worker venv via ray_worker_py_executable, matching the generation path. - Add max_num_seqs=256 and max_num_batched_tokens=8192 to vLLM serve kwargs. - Enable NCCL_MNNVL_ENABLE and the trtllm FlashInfer allreduce backend (VLLM_FLASHINFER_ALLREDUCE_BACKEND=trtllm, VLLM_ALLREDUCE_USE_SYMM_MEM=0); route DP through Ray (distributed_executor_backend/data_parallel_backend=ray). - Switch load_format to auto and replace cudagraph_capture_sizes with pass_config.fuse_allreduce_rms=false in compilation_config. - student_rlvr2: lower gpu_memory_utilization 0.85 -> 0.80 to leave headroom for the packed buffer during Megatron->vLLM refits, and document the NL2BASH_BASE_URL / GENRM_MODEL launcher overrides. - ultra_launch.sh: set the batch reaper reason to "other" (was "batch"). Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Anish Mahishi --- .../nemotron-3-ultra/ifbench_teacher.yaml | 17 +++++++++++++---- examples/nemo_gym/nemotron-3-ultra/mopd.yaml | 6 ++++++ .../nemotron-3-ultra/reasoning_teacher.yaml | 4 +++- .../nemotron-3-ultra/student_rlvr1.yaml | 13 ++++++++++--- .../nemotron-3-ultra/student_rlvr2.yaml | 18 ++++++++++++------ .../nemo_gym/nemotron-3-ultra/ultra_launch.sh | 2 +- 6 files changed, 45 insertions(+), 15 deletions(-) diff --git a/examples/nemo_gym/nemotron-3-ultra/ifbench_teacher.yaml b/examples/nemo_gym/nemotron-3-ultra/ifbench_teacher.yaml index d1f7c0f2b8..1340e506e3 100644 --- a/examples/nemo_gym/nemotron-3-ultra/ifbench_teacher.yaml +++ b/examples/nemo_gym/nemotron-3-ultra/ifbench_teacher.yaml @@ -485,6 +485,7 @@ env: return_token_id_information: false uses_reasoning_parser: false debug: true + ray_worker_py_executable: /opt/ray_venvs/nemo_rl.models.generation.vllm.vllm_worker_async.VllmAsyncGenerationWorker/bin/python vllm_serve_env_vars: VLLM_RAY_DP_PACK_STRATEGY: strict VLLM_USE_FLASHINFER_MOE_FP16: "0" @@ -504,6 +505,7 @@ env: tool_call_parser: hermes gpu_memory_utilization: 0.85 max_model_len: 131072 + max_num_seqs: 256 enable_prefix_caching: true enable_chunked_prefill: true model_loader_extra_config: @@ -578,24 +580,31 @@ env: ray_worker_py_executable: /opt/ray_venvs/nemo_rl.models.generation.vllm.vllm_worker_async.VllmAsyncGenerationWorker/bin/python vllm_serve_env_vars: VLLM_RAY_DP_PACK_STRATEGY: strict - NCCL_MNNVL_ENABLE: "0" + NCCL_MNNVL_ENABLE: "1" + VLLM_FLASHINFER_ALLREDUCE_BACKEND: trtllm + VLLM_ALLREDUCE_USE_SYMM_MEM: "0" vllm_serve_kwargs: tensor_parallel_size: 4 data_parallel_size: 4 data_parallel_size_local: 1 pipeline_parallel_size: 1 + distributed_executor_backend: ray + data_parallel_backend: ray reasoning_parser: deepseek_r1 gpu_memory_utilization: 0.85 max_model_len: 131072 enable_prefix_caching: true enable_chunked_prefill: true + max_num_seqs: 256 + max_num_batched_tokens: 8192 model_loader_extra_config: enable_multithread_load: true - num_threads: 112 - load_format: instanttensor + num_threads: 96 + load_format: auto compilation_config: - cudagraph_capture_sizes: [1,2,4,8,16,32] + pass_config: + fuse_allreduce_rms: false lc_judge: resources_servers: diff --git a/examples/nemo_gym/nemotron-3-ultra/mopd.yaml b/examples/nemo_gym/nemotron-3-ultra/mopd.yaml index d500beff38..959dd088b6 100644 --- a/examples/nemo_gym/nemotron-3-ultra/mopd.yaml +++ b/examples/nemo_gym/nemotron-3-ultra/mopd.yaml @@ -545,6 +545,7 @@ env: return_token_id_information: false uses_reasoning_parser: false debug: true + ray_worker_py_executable: /opt/ray_venvs/nemo_rl.models.generation.vllm.vllm_worker_async.VllmAsyncGenerationWorker/bin/python vllm_serve_env_vars: VLLM_RAY_DP_PACK_STRATEGY: strict VLLM_USE_FLASHINFER_MOE_FP16: "0" @@ -564,6 +565,9 @@ env: tool_call_parser: hermes gpu_memory_utilization: 0.85 max_model_len: 131072 + max_num_seqs: 256 + enable_prefix_caching: true + enable_chunked_prefill: true model_loader_extra_config: enable_multithread_load: true num_threads: 112 @@ -651,6 +655,8 @@ env: reasoning_parser: deepseek_r1 gpu_memory_utilization: 0.95 max_model_len: 60000 + max_num_seqs: 256 + max_num_batched_tokens: 8192 enable_prefix_caching: true enable_chunked_prefill: true model_loader_extra_config: diff --git a/examples/nemo_gym/nemotron-3-ultra/reasoning_teacher.yaml b/examples/nemo_gym/nemotron-3-ultra/reasoning_teacher.yaml index c406529f6e..f46e1693d6 100644 --- a/examples/nemo_gym/nemotron-3-ultra/reasoning_teacher.yaml +++ b/examples/nemo_gym/nemotron-3-ultra/reasoning_teacher.yaml @@ -434,9 +434,10 @@ env: return_token_id_information: false uses_reasoning_parser: false debug: true + ray_worker_py_executable: /opt/ray_venvs/nemo_rl.models.generation.vllm.vllm_worker_async.VllmAsyncGenerationWorker/bin/python vllm_serve_env_vars: VLLM_RAY_DP_PACK_STRATEGY: strict - NCCL_MNNVL_ENABLE: "0" + NCCL_MNNVL_ENABLE: "1" VLLM_USE_FLASHINFER_MOE_FP16: "0" VLLM_USE_FLASHINFER_MOE_FP8: "0" VLLM_USE_DEEP_GEMM: "0" @@ -453,6 +454,7 @@ env: tool_call_parser: hermes gpu_memory_utilization: 0.85 max_model_len: 131072 + max_num_seqs: 256 enable_prefix_caching: true enable_chunked_prefill: true model_loader_extra_config: diff --git a/examples/nemo_gym/nemotron-3-ultra/student_rlvr1.yaml b/examples/nemo_gym/nemotron-3-ultra/student_rlvr1.yaml index 94c1deb4fd..39ce9dbf0d 100644 --- a/examples/nemo_gym/nemotron-3-ultra/student_rlvr1.yaml +++ b/examples/nemo_gym/nemotron-3-ultra/student_rlvr1.yaml @@ -504,6 +504,7 @@ env: return_token_id_information: false uses_reasoning_parser: false debug: true + ray_worker_py_executable: /opt/ray_venvs/nemo_rl.models.generation.vllm.vllm_worker_async.VllmAsyncGenerationWorker/bin/python vllm_serve_env_vars: VLLM_RAY_DP_PACK_STRATEGY: strict VLLM_USE_FLASHINFER_MOE_FP16: "0" @@ -523,6 +524,7 @@ env: tool_call_parser: hermes gpu_memory_utilization: 0.85 max_model_len: 131072 + max_num_seqs: 256 enable_prefix_caching: true enable_chunked_prefill: true model_loader_extra_config: @@ -597,8 +599,8 @@ env: vllm_serve_env_vars: VLLM_RAY_DP_PACK_STRATEGY: strict NCCL_MNNVL_ENABLE: "1" - VLLM_USE_DEEP_GEMM: "0" - VLLM_MOE_USE_DEEP_GEMM: "0" + VLLM_FLASHINFER_ALLREDUCE_BACKEND: trtllm + VLLM_ALLREDUCE_USE_SYMM_MEM: "0" vllm_serve_kwargs: attention_backend: TRITON_ATTN @@ -606,9 +608,13 @@ env: data_parallel_size: 4 data_parallel_size_local: 1 pipeline_parallel_size: 1 + distributed_executor_backend: ray + data_parallel_backend: ray reasoning_parser: deepseek_r1 gpu_memory_utilization: 0.85 max_model_len: 131072 + max_num_seqs: 256 + max_num_batched_tokens: 8192 enable_prefix_caching: true enable_chunked_prefill: true model_loader_extra_config: @@ -616,7 +622,8 @@ env: num_threads: 112 load_format: auto compilation_config: - cudagraph_capture_sizes: [1,2,4,8,16,32] + pass_config: + fuse_allreduce_rms: false lc_judge: resources_servers: diff --git a/examples/nemo_gym/nemotron-3-ultra/student_rlvr2.yaml b/examples/nemo_gym/nemotron-3-ultra/student_rlvr2.yaml index 48fba94c72..a72b77485f 100644 --- a/examples/nemo_gym/nemotron-3-ultra/student_rlvr2.yaml +++ b/examples/nemo_gym/nemotron-3-ultra/student_rlvr2.yaml @@ -310,7 +310,9 @@ policy: tensor_parallel_size: 8 pipeline_parallel_size: 1 expert_parallel_size: 8 - gpu_memory_utilization: 0.85 + # Leave enough headroom for the temporary packed buffer used during + # Megatron-to-vLLM weight refits. + gpu_memory_utilization: 0.80 max_model_len: 65536 enforce_eager: false use_deep_gemm: false @@ -367,10 +369,11 @@ data: # ============================================================================= # Environment — NeMo Gym + Judge Models # -# Judge model paths and the GenRM service URL are set by the launcher: -# - genrm_model.base_url ← GENRM_BASE_URL -# - nl2bash_judge_model.model ← NL2BASH_JUDGE_MODEL -# - safety_judge_model.model ← SAFETY_JUDGE_MODEL +# Judge model paths and external service URLs are set by the launcher: +# - genrm_model.model/base_url ← GENRM_MODEL / GENRM_BASE_URL +# - nl2bash_judge_model.model ← NL2BASH_JUDGE_MODEL +# - nl2bash_judge_model.base_url ← NL2BASH_BASE_URL +# - safety_judge_model.model ← SAFETY_JUDGE_MODEL # ============================================================================= env: should_use_nemo_gym: true @@ -497,7 +500,8 @@ env: compilation_config: cudagraph_capture_sizes: [1,2,4,8,16] - # NL2Bash / general judge: TP=4 on GB200 192GB + # NL2Bash / general judge: use NL2BASH_BASE_URL for an external service, + # otherwise launch TP=4 locally on GB200 192GB. nl2bash_judge_model: responses_api_models: local_vllm_model: @@ -506,6 +510,7 @@ env: return_token_id_information: false uses_reasoning_parser: false debug: true + ray_worker_py_executable: /opt/ray_venvs/nemo_rl.models.generation.vllm.vllm_worker_async.VllmAsyncGenerationWorker/bin/python vllm_serve_env_vars: VLLM_RAY_DP_PACK_STRATEGY: strict VLLM_USE_FLASHINFER_MOE_FP16: "0" @@ -525,6 +530,7 @@ env: tool_call_parser: hermes gpu_memory_utilization: 0.85 max_model_len: 131072 + max_num_seqs: 256 enable_prefix_caching: true enable_chunked_prefill: true model_loader_extra_config: diff --git a/examples/nemo_gym/nemotron-3-ultra/ultra_launch.sh b/examples/nemo_gym/nemotron-3-ultra/ultra_launch.sh index 3d4b9db6a0..a956127206 100755 --- a/examples/nemo_gym/nemotron-3-ultra/ultra_launch.sh +++ b/examples/nemo_gym/nemotron-3-ultra/ultra_launch.sh @@ -841,7 +841,7 @@ SBATCH_OUTPUT=$(sbatch \ ${SLURM_QOS:+--qos="${SLURM_QOS}"} \ ${EXCLUDE_NODES:+--exclude="${EXCLUDE_NODES}"} \ ${SLURM_RESERVATION:+--reservation="${SLURM_RESERVATION}"} \ - --comment='{"OccupiedIdleGPUsJobReaper":{"exemptIdleTimeMins":"60","reason":"batch","description":"batch training run"}}' \ + --comment='{"OccupiedIdleGPUsJobReaper":{"exemptIdleTimeMins":"60","reason":"other","description":"batch training run"}}' \ "${RAY_SUB}") echo "${SBATCH_OUTPUT}" From 917fb6e77f9cecbaa42b4ddc28edd01ea9c4e8e1 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Wed, 29 Jul 2026 17:30:20 -0400 Subject: [PATCH 12/15] refactor(ultra): drop custom vLLM fork path from launcher MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Ultra runs on the container's stock vLLM, so remove the USE_CUSTOM_VLLM switch and everything that only served the fork: the case block that sourced 3rdparty/vllm/nemo-rl.env, the ${VLLM_ENV_SOURCE} prefix in TRAIN_CMD, the 3rdparty/vllm code-snapshot symlink, the status echo, and the related header/mount comments. The vLLM compile-cache and runtime env (VLLM_CACHE_ROOT, FlashInfer MoE, etc.) are untouched — they serve the stock vLLM too. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Anish Mahishi --- .../nemo_gym/nemotron-3-ultra/ultra_launch.sh | 23 ------------------- 1 file changed, 23 deletions(-) diff --git a/examples/nemo_gym/nemotron-3-ultra/ultra_launch.sh b/examples/nemo_gym/nemotron-3-ultra/ultra_launch.sh index a956127206..f06e952b21 100755 --- a/examples/nemo_gym/nemotron-3-ultra/ultra_launch.sh +++ b/examples/nemo_gym/nemotron-3-ultra/ultra_launch.sh @@ -48,8 +48,6 @@ set -euo pipefail # NRL_MAX_STEPS= Override grpo.max_num_steps # EXTRA_MOUNTS= Comma-separated host:container pairs # USE_SNAPSHOT=1 Snapshot source tree at submission -# USE_CUSTOM_VLLM=1 1 to require the Ultra vLLM fork; -# 0 to use the container's regular vLLM # DRY_RUN=0 1 to print TRAIN_CMD and exit # INTERACTIVE=0 1 to bring up Ray and idle for attach # (no training driver) for debugging @@ -362,19 +360,6 @@ export NEMO_SKILLS_SANDBOX_PORT="${NEMO_SKILLS_SANDBOX_PORT:-6000}" export RAY_LOG_SYNC_FREQUENCY="${RAY_LOG_SYNC_FREQUENCY:-60}" CODE_ROOT="/opt/nemo-rl" -USE_CUSTOM_VLLM="${USE_CUSTOM_VLLM:-1}" -case "${USE_CUSTOM_VLLM}" in - 1) - VLLM_ENV_SOURCE="source /opt/nemo-rl/3rdparty/vllm/nemo-rl.env && " - ;; - 0) - VLLM_ENV_SOURCE="" - ;; - *) - echo "ERROR: USE_CUSTOM_VLLM must be 0 or 1, got: ${USE_CUSTOM_VLLM}" >&2 - exit 1 - ;; -esac # ============================================================================= # Persistent cache directories @@ -495,11 +480,6 @@ if [[ "${USE_SNAPSHOT}" == "1" ]]; then fi SNAPSHOT_DIR=$(bash "${PROJECT_ROOT}/tools/code_snapshot.sh" "${JOB_NAME}") - if [[ -d "${PROJECT_ROOT}/3rdparty/vllm" ]] && [[ ! -e "${SNAPSHOT_DIR}/3rdparty/vllm" ]]; then - mkdir -p "${SNAPSHOT_DIR}/3rdparty" - ln -s "${PROJECT_ROOT}/3rdparty/vllm" "${SNAPSHOT_DIR}/3rdparty/vllm" - fi - echo "Code snapshot: ${SNAPSHOT_DIR}" OVERLAY_SOURCE="${SNAPSHOT_DIR}" else @@ -524,7 +504,6 @@ fi # /opt/nemo-rl/3rdparty/Megatron-LM-workspace/Megatron-LM — Megatron-LM # /opt/nemo-rl/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge — Megatron-Bridge # /opt/nemo-rl/3rdparty/Gym-workspace/Gym — NeMo-Gym -# /opt/nemo-rl/3rdparty/vllm — vLLM # ============================================================================= _append_mount() { if [[ -z "${MOUNTS}" ]]; then @@ -630,7 +609,6 @@ export SETUP_COMMAND # per-run overrides: cluster shape, paths, judge endpoints, logging. # ============================================================================= TRAIN_CMD="cd ${CODE_ROOT} && date ; \ -${VLLM_ENV_SOURCE}\ OMP_NUM_THREADS=16 \ RAY_DEDUP_LOGS=1 \ WANDB_INIT_TIMEOUT=300 \ @@ -704,7 +682,6 @@ echo " Model: ${MODEL_PATH}" echo " Train data: ${TRAIN_PATH}" echo " Val data: ${VAL_PATH}" echo " Container: ${CONTAINER}" -echo " Custom vLLM: ${USE_CUSTOM_VLLM}" echo " Sandbox: ${SANDBOX_CONTAINER}" if [[ "${USE_SNAPSHOT}" == "1" ]]; then echo " Snapshot: ${SNAPSHOT_DIR}" From a99f247da0a0961618b226c7cbd7dcdf10f007da Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Wed, 29 Jul 2026 17:30:32 -0400 Subject: [PATCH 13/15] docs(ultra): use stock/nightly container, target main, sync recipe params - Container section: Ultra runs on the stock NeMo RL container (upstream aarch64 vLLM wheel); document pulling a prebuilt nightly from NGC and drop the BUILD_CUSTOM_VLLM build args and the USE_CUSTOM_VLLM note. - Prepare the code: clone -b main instead of the ultra-v3 branch. - Sync recipe parameters to the configs: student RLVR phase-1 group_answer_length_penalty_coeff 0.15 -> 0.25, MOPD CP 8 -> 32, and the RLHF/Reasoning teacher default cluster shapes (Gym node counts). Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Anish Mahishi --- docs/guides/nemotron-3-ultra.md | 41 +++++++++++++++++---------------- 1 file changed, 21 insertions(+), 20 deletions(-) diff --git a/docs/guides/nemotron-3-ultra.md b/docs/guides/nemotron-3-ultra.md index 59c240fb1e..a71d3bafaf 100644 --- a/docs/guides/nemotron-3-ultra.md +++ b/docs/guides/nemotron-3-ultra.md @@ -61,11 +61,19 @@ RLVR policy itself serves as the general teacher: ## Container -These are the instructions to build the training container. The image targets -**aarch64 (arm64)** for GB200 NVL72 nodes and bundles a custom vLLM build -required by Ultra. +Ultra runs on the stock NeMo RL container. The vLLM version is whatever `pyproject.toml` pins (currently the upstream +aarch64 wheel), so the only requirement is an **aarch64 (arm64)** image for +GB200 NVL72 nodes. -From the root of the repo: +The quickest option is to pull a prebuilt nightly image from +[NGC](https://registry.ngc.nvidia.com/orgs/nvidia/containers/nemo-rl/tags) and +skip the build entirely: + +```bash +docker pull nvcr.io/nvidia/nemo-rl: +``` + +To build it yourself instead, from the root of the repo: ```bash docker buildx build \ @@ -76,23 +84,16 @@ docker buildx build \ --build-context nemo-rl=. \ --build-arg MAX_JOBS=8 \ --build-arg SKIP_SGLANG_BUILD=1 \ - --build-arg BUILD_CUSTOM_VLLM=1 \ - --build-arg BUILD_CUSTOM_VLLM_URL=https://github.com/TomerBN-Nvidia/vllm.git \ - --build-arg BUILD_CUSTOM_VLLM_REF=ultra-rl-v0.17 \ - --build-arg BUILD_CUSTOM_VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.17.0/vllm-0.17.0-cp38-abi3-manylinux_2_31_aarch64.whl \ . ``` Build args: -- `BUILD_CUSTOM_VLLM=1` with `BUILD_CUSTOM_VLLM_URL` / `BUILD_CUSTOM_VLLM_REF` — - builds the Ultra vLLM fork; `BUILD_CUSTOM_VLLM_PRECOMPILED_WHEEL_LOCATION` - points at the matching upstream aarch64 wheel so the build reuses precompiled - kernels instead of compiling from source. - `SKIP_SGLANG_BUILD=1` — Ultra runs on vLLM; skip the SGLang build. - `MAX_JOBS` — parallel build jobs; tune to your machine. - `--build-context nemo-rl=.` — build from your local checkout (otherwise the Dockerfile pulls `NVIDIA-NeMo/RL.git#main`). + To run on the cluster with Slurm, convert the image to a squashfs (`.sqsh`) with [enroot](https://github.com/NVIDIA/enroot): @@ -151,7 +152,7 @@ For now, each stage takes a JSONL training file and a JSONL validation file ## Prepare the code ```bash -git clone --recursive -b ultra-v3 https://github.com/NVIDIA-NeMo/RL.git +git clone --recursive -b main https://github.com/NVIDIA-NeMo/RL.git cd RL ``` @@ -254,7 +255,7 @@ second phase resumes from the first phase's checkpoint: | `max_total_sequence_length` | 49,152 | 65,536 | | Steps in this phase | ~128 | ~50 | | `NRL_MAX_STEPS` to set | `128` | `178` (= 128 + 50) | -| `group_answer_length_penalty_coeff` | 0.15 | 0.08 | +| `group_answer_length_penalty_coeff` | 0.25 | 0.08 | Both phases share TP=8, EP=64, CP=8, PP=1, GBS=8192 (512 prompts × 16 generations), advantage clip ±20, and the 256-node cluster shape. @@ -378,7 +379,7 @@ learning rate, context). - `max_total_sequence_length=49152` - `train_global_batch_size=2048`, `num_prompts_per_step=128`, `num_generations_per_prompt=16` - Learning rate `2.5e-6` constant -- Default cluster shape: 64 nodes (32 training + 28 vLLM + 4 Gym) +- Default cluster shape: 80 nodes (32 training + 28 vLLM + 20 Gym) ```bash EXP_NAME=ultra-rlhf-teacher \ @@ -388,7 +389,7 @@ TRAIN_PATH=$DATA_DIR/rlhf.train.jsonl \ VAL_PATH=$DATA_DIR/rlhf.val.jsonl \ NUM_TRAIN_NODES=32 \ NUM_GEN_NODES=28 \ -NUM_GYM_NODES=4 \ +NUM_GYM_NODES=20 \ ENABLE_MTP_INFERENCE=1 \ CONTAINER=/path/to/nemo-rl-container \ SANDBOX_CONTAINER=/path/to/nemo-skills-sandbox.sqsh \ @@ -422,7 +423,7 @@ many roles. - `train_global_batch_size=2048`, `num_prompts_per_step=128`, `num_generations_per_prompt=16` - Learning rate `3.0e-6` constant - `max_num_epochs=10` — small sub-sampled dataset, multiple passes expected -- Default cluster shape: 128 nodes (64 training + 60 vLLM + 4 Gym) +- Default cluster shape: 128 nodes (64 training + 54 vLLM + 10 Gym) ```bash EXP_NAME=ultra-reasoning-teacher \ @@ -432,8 +433,8 @@ MODEL_PATH=/path/to/student_rlvr_output \ TRAIN_PATH=$DATA_DIR/reasoning.train.jsonl \ VAL_PATH=$DATA_DIR/reasoning.val.jsonl \ NUM_TRAIN_NODES=64 \ -NUM_GEN_NODES=60 \ -NUM_GYM_NODES=4 \ +NUM_GEN_NODES=54 \ +NUM_GYM_NODES=10 \ CONTAINER=/path/to/nemo-rl-container \ SANDBOX_CONTAINER=/path/to/nemo-skills-sandbox.sqsh \ PERSISTENT_CACHE=/path/to/persistent/cache \ @@ -581,7 +582,7 @@ Gym agent is routed to one of the Stage 2 teacher checkpoints. Trains the student to match per-agent teacher distributions. **Config:** `examples/nemo_gym/nemotron-3-ultra/mopd.yaml` -- TP=8, EP=64, CP=8, PP=1, max context 192k +- TP=8, EP=64, CP=32, PP=1, max context 192k - Teacher parallelism: TP=8, CP=2, EP=16, 4 nodes per teacher - Routing: agent → teacher checkpoint baked into the YAML via `${_teachers.}` references; only `_teachers.general` is required and From b65d3fdf85584289fe89d7c57a192148c8a3df8d Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Wed, 29 Jul 2026 17:43:26 -0400 Subject: [PATCH 14/15] docs: note about external genrm service Signed-off-by: Anish Mahishi --- docs/guides/nemotron-3-ultra.md | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/docs/guides/nemotron-3-ultra.md b/docs/guides/nemotron-3-ultra.md index a71d3bafaf..2b4481a8a2 100644 --- a/docs/guides/nemotron-3-ultra.md +++ b/docs/guides/nemotron-3-ultra.md @@ -228,6 +228,14 @@ Set the following before each `bash examples/nemo_gym/nemotron-3-ultra/ultra_lau | `NL2BASH_JUDGE_MODEL` | NL2Bash / general-purpose judge: HF repo id or local path. Default judge is `Qwen/Qwen3-235B-A22B-Instruct-2507-FP8`. | | `SAFETY_JUDGE_MODEL` | Content-safety judge: HF repo id or local path. Default is [`nvidia/Nemotron-Content-Safety-Reasoning-4B`](https://huggingface.co/nvidia/Nemotron-Content-Safety-Reasoning-4B). | +> **Serving GenRM as a standalone service.** The GenRM judge does not have to run +> inside the training job. You can bring it up separately — any OpenAI-compatible +> vLLM endpoint, or the external-GenRM service launcher under `tools/external_genrm/` +> (runs the judge fleet in a second Slurm hetgroup behind a load balancer) — and +> point the run at it with `GENRM_BASE_URL=http://:/v1`. Judging is then +> routed to that endpoint instead of being served from the gym pool, which frees +> those GPUs and lets one GenRM deployment back many training runs. + Optional knobs: | Variable | Default | Purpose | From 31c3bf7f07647b97de0c1334aabfbb46feeb74bc Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Wed, 29 Jul 2026 20:02:46 -0400 Subject: [PATCH 15/15] =?UTF-8?q?docs(ultra):=20address=20review=20?= =?UTF-8?q?=E2=80=94=20drop=20stale=20MOPD=20header,=20fix=20GenRM=20DP=20?= =?UTF-8?q?claim?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - mopd.yaml: remove the "Disabled on main" header note. OPD/MOPD runtime support (OPDAdvantageEstimator, non-colocated teachers) is already on main and the recipe ships as an active .yaml, so the note told readers not to use a shipped recipe. - nemotron-3-ultra.md: GENRM_MODEL is served at TP=4 with DP varying by stage (DP=16 in the RLHF and student RLVR2 stages), not a fixed DP=4 / 16 GPUs. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Anish Mahishi --- docs/guides/nemotron-3-ultra.md | 2 +- examples/nemo_gym/nemotron-3-ultra/mopd.yaml | 5 ----- 2 files changed, 1 insertion(+), 6 deletions(-) diff --git a/docs/guides/nemotron-3-ultra.md b/docs/guides/nemotron-3-ultra.md index 2b4481a8a2..3eed7fd3ea 100644 --- a/docs/guides/nemotron-3-ultra.md +++ b/docs/guides/nemotron-3-ultra.md @@ -223,7 +223,7 @@ Set the following before each `bash examples/nemo_gym/nemotron-3-ultra/ultra_lau | `PERSISTENT_CACHE` | Directory on a shared filesystem (e.g. Lustre) where vLLM/Triton/Inductor compile caches are persisted across runs. | | `EXTRA_MOUNTS` | Comma-separated `host:container` mount pairs for any shared filesystems holding your data, model checkpoints, and `PERSISTENT_CACHE` (e.g. `EXTRA_MOUNTS=/lustre:/lustre,/scratch:/scratch`). | | `SLURM_PARTITION`, `SLURM_ACCOUNT` | Your SLURM cluster credentials. | -| `GENRM_MODEL` | GenRM judge: HF repo id or local path. Default is [`nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-GenRM`](https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-GenRM). Served in-cluster from the gym pool at TP=4 DP=4 (16 GPUs). Required unless `GENRM_BASE_URL` is set. | +| `GENRM_MODEL` | GenRM judge: HF repo id or local path. Default is [`nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-GenRM`](https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-GenRM). Served in-cluster from the gym pool at TP=4 (DP varies by stage). Required unless `GENRM_BASE_URL` is set. | | `GENRM_BASE_URL` | _Optional._ URL of a separately-deployed GenRM service (e.g. `http://genrm-host:9213/v1`). If set, routes judging to that endpoint and ignores `GENRM_MODEL`. Useful when sharing a single GenRM deployment across many training jobs. | | `NL2BASH_JUDGE_MODEL` | NL2Bash / general-purpose judge: HF repo id or local path. Default judge is `Qwen/Qwen3-235B-A22B-Instruct-2507-FP8`. | | `SAFETY_JUDGE_MODEL` | Content-safety judge: HF repo id or local path. Default is [`nvidia/Nemotron-Content-Safety-Reasoning-4B`](https://huggingface.co/nvidia/Nemotron-Content-Safety-Reasoning-4B). | diff --git a/examples/nemo_gym/nemotron-3-ultra/mopd.yaml b/examples/nemo_gym/nemotron-3-ultra/mopd.yaml index 959dd088b6..b7190133c7 100644 --- a/examples/nemo_gym/nemotron-3-ultra/mopd.yaml +++ b/examples/nemo_gym/nemotron-3-ultra/mopd.yaml @@ -1,8 +1,3 @@ -# NOTE: Disabled on main because this config requires OPD/MOPD runtime support -# from the ultra-v3 branch (OPDAdvantageEstimator, on_policy_distillation, -# and non-colocated teacher worker groups). Rename to .yaml only after that -# support lands on main. - # ============================================================================= # Nemotron 3 Ultra — MOPD Stage 1 (on-policy distillation from teacher panel) # =============================================================================