diff --git a/docs/guides/nemotron-3-ultra.md b/docs/guides/nemotron-3-ultra.md new file mode 100644 index 0000000000..3eed7fd3ea --- /dev/null +++ b/docs/guides/nemotron-3-ultra.md @@ -0,0 +1,648 @@ +# 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/nemo_gym/nemotron-3-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 + +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. + +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 \ + --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 args: +- `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 main 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 `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 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/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. | +| `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 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). | + +> **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 | +|---|---|---| +| `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/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) | +| `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. + +### Phase 1 — 49k context, 128 steps + +```bash +EXP_NAME=ultra-student-rlvr \ +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 \ +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 examples/nemo_gym/nemotron-3-ultra/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/nemo_gym/nemotron-3-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 examples/nemo_gym/nemotron-3-ultra/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/nemo_gym/nemotron-3-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/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` +- 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/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 \ +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 examples/nemo_gym/nemotron-3-ultra/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/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` +- Learning rate `2.5e-6` constant +- Default cluster shape: 80 nodes (32 training + 28 vLLM + 20 Gym) + +```bash +EXP_NAME=ultra-rlhf-teacher \ +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 \ +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 \ +WANDB_API_KEY=$WANDB_API_KEY \ +HF_HOME=/path/to/hf_cache \ +HF_TOKEN=$HF_TOKEN \ +bash examples/nemo_gym/nemotron-3-ultra/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/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` +- Learning rate `3.0e-6` constant +- `max_num_epochs=10` — small sub-sampled dataset, multiple passes expected +- Default cluster shape: 128 nodes (64 training + 54 vLLM + 10 Gym) + +```bash +EXP_NAME=ultra-reasoning-teacher \ +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 \ +VAL_PATH=$DATA_DIR/reasoning.val.jsonl \ +NUM_TRAIN_NODES=64 \ +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 \ +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 examples/nemo_gym/nemotron-3-ultra/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/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` +- 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/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 \ +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 examples/nemo_gym/nemotron-3-ultra/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/nemo_gym/nemotron-3-ultra/mopd.yaml` +- 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 + 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/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 \ +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 examples/nemo_gym/nemotron-3-ultra/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/examples/nemo_gym/nemotron-3-ultra/ifbench_teacher.yaml b/examples/nemo_gym/nemotron-3-ultra/ifbench_teacher.yaml new file mode 100644 index 0000000000..1340e506e3 --- /dev/null +++ b/examples/nemo_gym/nemotron-3-ultra/ifbench_teacher.yaml @@ -0,0 +1,671 @@ +# ============================================================================= +# 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 + save_optimizer: true + +# ============================================================================= +# 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 + gradient_accumulation_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: ${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 + + # --------------------------------------------------------------------------- + # 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 + 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 + + 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 + + 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 + use_multiple_dataloader: false + 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 + num_gpu_nodes: 20 + # 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: + - "" + - "" + - "" + - "" + 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_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 + 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 + 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" + 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 + 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 + max_num_seqs: 256 + 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] + + 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" + 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 + 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 + 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: 96 + load_format: auto + compilation_config: + pass_config: + fuse_allreduce_rms: false + + 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 + +# ============================================================================= +# Reward Penalties (set reward to 0 when triggered) +# ============================================================================= +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/nemo_gym/nemotron-3-ultra/mopd.yaml b/examples/nemo_gym/nemotron-3-ultra/mopd.yaml new file mode 100644 index 0000000000..b7190133c7 --- /dev/null +++ b/examples/nemo_gym/nemotron-3-ultra/mopd.yaml @@ -0,0 +1,832 @@ +# ============================================================================= +# 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 + save_optimizer: true + +# ============================================================================= +# 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 + gradient_accumulation_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: ${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 + + # --------------------------------------------------------------------------- + # 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 + 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 + + 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 + + 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 + use_multiple_dataloader: false + 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 + num_gpu_nodes: 12 + # 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. + 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_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 + 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 + 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" + 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 + 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 + max_num_seqs: 256 + 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] + + 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 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: + 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 + 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 + data_parallel_size: 4 + data_parallel_size_local: 1 + pipeline_parallel_size: 1 + 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: + enable_multithread_load: true + num_threads: 112 + load_format: auto + compilation_config: + cudagraph_capture_sizes: [1,2,4,8,16,32] + + 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 + +# ============================================================================= +# Reward Penalties (set reward to 0 when triggered) +# ============================================================================= +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. +# +# `_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: /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} + 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/nemo_gym/nemotron-3-ultra/reasoning_teacher.yaml b/examples/nemo_gym/nemotron-3-ultra/reasoning_teacher.yaml new file mode 100644 index 0000000000..f46e1693d6 --- /dev/null +++ b/examples/nemo_gym/nemotron-3-ultra/reasoning_teacher.yaml @@ -0,0 +1,508 @@ +# ============================================================================= +# 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 + save_optimizer: true + +# ============================================================================= +# 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 + gradient_accumulation_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: ${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 + + # --------------------------------------------------------------------------- + # 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 + 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 + + 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 + + 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 + use_multiple_dataloader: false + 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 + num_gpu_nodes: 10 + # 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: + - "" + - "" + - "" + - "" + 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 + 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: "1" + 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 + 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 + max_num_seqs: 256 + 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] + + 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 + +# ============================================================================= +# Reward Penalties (set reward to 0 when triggered) +# ============================================================================= +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/nemo_gym/nemotron-3-ultra/rlhf_teacher.yaml b/examples/nemo_gym/nemotron-3-ultra/rlhf_teacher.yaml new file mode 100644 index 0000000000..8d507a8238 --- /dev/null +++ b/examples/nemo_gym/nemotron-3-ultra/rlhf_teacher.yaml @@ -0,0 +1,537 @@ +# ============================================================================= +# 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 + save_optimizer: true + +# ============================================================================= +# 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 + gradient_accumulation_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: ${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 + + # --------------------------------------------------------------------------- + # 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 + 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 + + 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 + + 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 + use_multiple_dataloader: false + 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 + num_gpu_nodes: 20 + # 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: + - "" + - "" + - "" + - "" + 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" + 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 + group_answer_length_penalty_coeff: 0.0 + group_style_penalty_coeff: 0.0 + default_score: 3.0 + default_ranking: 3.5 + + # 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: + 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: + # Keep each TP=4 replica on one four-GPU GB200 node. + VLLM_RAY_DP_PACK_STRATEGY: strict + 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: 16 + data_parallel_size_local: 1 + pipeline_parallel_size: 1 + 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: 96 + load_format: auto + compilation_config: + pass_config: + fuse_allreduce_rms: false + + +# ============================================================================= +# 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 + +# ============================================================================= +# Reward Penalties (set reward to 0 when triggered) +# ============================================================================= +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/nemo_gym/nemotron-3-ultra/student_rlvr1.yaml b/examples/nemo_gym/nemotron-3-ultra/student_rlvr1.yaml new file mode 100644 index 0000000000..39ce9dbf0d --- /dev/null +++ b/examples/nemo_gym/nemotron-3-ultra/student_rlvr1.yaml @@ -0,0 +1,690 @@ +# ============================================================================= +# 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 + save_optimizer: true + +# ============================================================================= +# 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 + gradient_accumulation_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: ${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 + + # --------------------------------------------------------------------------- + # 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 + 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 + + 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 + + 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 + use_multiple_dataloader: false + 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 + num_gpu_nodes: 20 + # 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: + - "" + - "" + - "" + - "" + 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_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 + 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 + 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" + 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 + 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 + max_num_seqs: 256 + 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] + + 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 + NCCL_MNNVL_ENABLE: "1" + 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_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: + enable_multithread_load: true + num_threads: 112 + load_format: auto + compilation_config: + pass_config: + fuse_allreduce_rms: false + + 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 + +# ============================================================================= +# Reward Penalties (set reward to 0 when triggered) +# ============================================================================= +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/nemo_gym/nemotron-3-ultra/student_rlvr2.yaml b/examples/nemo_gym/nemotron-3-ultra/student_rlvr2.yaml new file mode 100644 index 0000000000..a72b77485f --- /dev/null +++ b/examples/nemo_gym/nemotron-3-ultra/student_rlvr2.yaml @@ -0,0 +1,702 @@ +# ============================================================================= +# 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 + save_optimizer: true + +# ============================================================================= +# 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 + gradient_accumulation_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: ${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 + + # --------------------------------------------------------------------------- + # 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 + # 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 + 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 + 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 + + 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 + + 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 + use_multiple_dataloader: false + 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 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 + should_log_nemo_gym_responses: true + 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}" + low_weight: 0.1 + low_penalty: 1 + low_ub: 15000 + 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_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 + 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: 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: + entrypoint: app.py + model: null # Set by launcher (NL2BASH_JUDGE_MODEL) + 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" + 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: 2 + 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 + max_num_seqs: 256 + 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] + + 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: + # Keep each TP=4 replica on one four-GPU GB200 node. + VLLM_RAY_DP_PACK_STRATEGY: strict + 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: 16 + data_parallel_size_local: 1 + pipeline_parallel_size: 1 + 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: 96 + load_format: auto + compilation_config: + pass_config: + fuse_allreduce_rms: false + + 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 + +# ============================================================================= +# Reward Penalties (set reward to 0 when triggered) +# ============================================================================= +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/nemo_gym/nemotron-3-ultra/swe_teacher.yaml b/examples/nemo_gym/nemotron-3-ultra/swe_teacher.yaml new file mode 100644 index 0000000000..994973a7cb --- /dev/null +++ b/examples/nemo_gym/nemotron-3-ultra/swe_teacher.yaml @@ -0,0 +1,480 @@ +# ============================================================================= +# 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 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: /path/to/sif_images + +# ============================================================================= +# 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 + save_optimizer: true + +# ============================================================================= +# 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 + gradient_accumulation_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: 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 — 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: ${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 + + # --------------------------------------------------------------------------- + # 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 + 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 + + 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 + + colocated: + enabled: false + resources: + gpus_per_node: 4 + num_nodes: 64 + +# ============================================================================= +# Data +# ============================================================================= +data: + max_input_seq_length: null + shuffle: false + num_workers: 1 + use_multiple_dataloader: false + 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 + # 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 + 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 + +# ============================================================================= +# Reward Penalties (set reward to 0 when triggered) +# ============================================================================= +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 # diff --git a/examples/nemo_gym/nemotron-3-ultra/ultra_launch.sh b/examples/nemo_gym/nemotron-3-ultra/ultra_launch.sh new file mode 100755 index 0000000000..f06e952b21 --- /dev/null +++ b/examples/nemo_gym/nemotron-3-ultra/ultra_launch.sh @@ -0,0 +1,831 @@ +#!/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/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 +# Hydra overrides that vary per run (data paths, model checkpoint, judge +# endpoints, log directories). +# +# Usage: +# +# EXP_NAME=ultra-student-rlvr-001 \ +# 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 \ +# 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 examples/nemo_gym/nemotron-3-ultra/ultra_launch.sh +# +# Optional knobs: +# WALLTIME=4:00:00 Slurm --time +# 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:) +# 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 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). +# ============================================================================= + +# ============================================================================= +# 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/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)}" +: "${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 + +# 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 +# ============================================================================= +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:-}" + +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. +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 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}" + +# ============================================================================= +# 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_rlvr1.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" + +# ============================================================================= +# 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}") + + 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 examples/nemo_gym/nemotron-3-ultra/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 +# ============================================================================= +_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 ; \ +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=/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 \ +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} \ +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} \ +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}"} \ + --comment='{"OccupiedIdleGPUsJobReaper":{"exemptIdleTimeMins":"60","reason":"interactive","description":"interactive debugging"}}' \ + "${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}"} \ + --comment='{"OccupiedIdleGPUsJobReaper":{"exemptIdleTimeMins":"60","reason":"other","description":"batch training run"}}' \ + "${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 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..9d82e89b30 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/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", +] @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,50 @@ 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) + + # 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) + + 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