feat(recipes): add optimized 32-GPU Nano-Omni GRPO recipes

[youngeunkwon0405]

What does this PR do?

Adds two production-shaped 32-GPU Nemotron-3-Nano-Omni GRPO recipes:

• synchronous colocated training on 8 nodes x 4 GPUs; and
• asynchronous 1-off training on a 16-GPU learner plus a 16-GPU rollout pool.

Both recipes inherit from the native 1-node/8-GPU Nano-Omni recipe and preserve its 8K total sequence length, 4K generation cap, and group size 16. The 32-GPU recipes use 64 prompts per step for a global rollout batch of 1024.

The shared optimized stack uses learner TP1/EP8, no activation checkpointing, a 10K packed-training token budget, fused weighted squared ReLU, and vLLM TP2/EP2 with CUDA graphs, FlashInfer CUTLASS MoE, and one-sided NVLink all-to-all. The async recipe additionally enables lag-1 GRPO with in-flight weight updates, uses a 16+16 GPU split, and raises max_num_seqs to 128.

This PR contains only the two recipe YAMLs and their required test-suite registrations. Matching suite drivers are registered as disabled because NeMo-RL requires every recipe YAML to have a same-named suite driver, while each of these recipes reserves eight GB200 nodes. The NeMo-Gym dynamic-image-padding propagation used by the private benchmark environment remains independent in #3738. The parallel image-encoding implementation remains independent in #3721.

The async recipe was validated through examples/nemo_gym/run_grpo_nemo_gym.py, which already dispatches to the async trainer. Async dispatch for the general examples/run_vlm_grpo.py entrypoint is proposed separately in #3740 and is intentionally outside this recipe PR.

Known limitation

examples/run_vlm_grpo.py does not currently dispatch to async_grpo_train when grpo.async_grpo.enabled=true. That general entrypoint fix is tracked by #3740 and intentionally remains outside this recipe-only PR. The async recipe was validated with the existing async-capable NeMo-Gym entrypoint.

Usage

# Sync colocated, 8 nodes x 4 GPUs
uv run examples/run_vlm_grpo.py \
  --config examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-tp1ep8.v1.yaml

The async 1-off recipe was launched through examples/nemo_gym/run_grpo_nemo_gym.py with the private NeMo-Gym dataset and environment supplied as launch-time overrides.

Performance

All timings below are means over warm Steps 2-10 with the same 32-GPU scale, batch 1024, 8K total sequence length, 4K generation cap, and data.

The measured optimization journey is:

• Sync: 221.98s naive -> 117.17s optimized recipe (47.2% lower) -> 81.28s with parallel image encoding (30.6% additional, 63.4% end-to-end).
• Async 1-off: 124.98s initial -> 85.12s optimized recipe (31.9% lower) -> 43.24s with parallel image encoding (49.2% additional, 65.4% end-to-end).

Configuration	Total step	Policy training	Generation exposed	Other	Step-time reduction
Naive sync eager=false (W&B)	221.98s	99.73s	108.80s	13.45s	baseline
Optimized sync recipe (W&B)	117.17s	18.31s	90.00s	8.87s	47.2% vs naive
Optimized sync recipe + parallel image encoding from #3721 (W&B)	81.28s	34.00s	39.24s	8.03s	63.4% vs naive; 30.6% incremental
Initial async 1-off (W&B)	124.98s	106.02s	17.04s	1.93s	baseline
Optimized async recipe (W&B)	85.12s	27.68s	55.50s	1.94s	31.9% vs initial async
Optimized async recipe + parallel image encoding from #3721 (W&B)	43.24s	28.09s	12.84s	2.31s	65.4% vs initial async; 49.2% incremental

The sync generation exposed column is the synchronous generation phase. For async, it is only the generation latency exposed on the learner critical path; rollout collection overlaps policy work.

The private performance launch selected NeMo Gym and supplied its dataset/environment options at launch time. It also supplied benchmark-control choices such as disabled validation, deterministic data order, and rollout filtering. Those workload-specific controls are intentionally not baked into these portable recipes. The performance-critical scale, parallelism, packing, learner, generation, and async settings in the YAMLs match the validated runs. The runs also included the separately scoped dynamic-image-padding propagation because the private examples contain heterogeneous image shapes.

The final-stack runs additionally contained the changes from #3721. The incremental values above are observed between separate benchmark runs, not a claim that image encoding alone explains every timing difference. A controlled preparation-path A/B reduced image-encoding wall time from 4.83s to 1.22s with 32 workers; #3721 uses a conservative cap of 8 workers. This recipe PR does not include either the dynamic-padding fix or the image-encoding implementation.

The integrated sync policy mean includes three intermittent slow learner steps (62.27s, 105.63s, and 59.69s); its warm policy median was 13.54s. The end-to-end mean retains those outliers.

Functional validation

The performance-critical settings represented by both YAMLs completed 10 steps in the private NeMo-Gym benchmark environment with the separate padding fix. They were then run again from the isolated final-stack branch containing #3721. Private dataset and benchmark-control overrides were supplied by the launcher rather than embedded in the portable recipe.

Recipe	Steps	Mean total reward	Mean generation KL	Max generation KL	Mean policy KL	Max policy KL	Result
Sync final stack	10	0.500946	0.001330	0.001440	0.001508	0.001680	COMPLETED 0:0
Async 1-off final stack	10	0.501827	0.001327	0.001424	0.001926	0.005674	COMPLETED 0:0

For sync, warm Steps 2-10 also produced finite mean loss 0.079477, mean gradient norm 0.016560, and mean probability ratio 1.000000001. For async, warm Steps 2-10 produced finite mean loss 0.101380, mean gradient norm 0.032467, mean probability ratio 0.999999993, and mean trajectory age 1.0. Both jobs completed all ten timing blocks without OOM or weight-update failure.

Before your PR is "Ready for review"

Pre checks:

• Read and followed the contributor guidelines
• Added minimized recipe YAMLs that pass tools/config_cli.py minimize-check
• Added matching recipe-suite drivers required by recipe accounting and registered them in disabled.txt
• Ran both recipe settings for 10 steps at their declared 8-node/4-GPU scale
• Documented the async entrypoint limitation and the validation launcher in this PR

Additional Information

The private NeMo Gym dataset/environment used for performance validation was supplied as launch-time overrides and is not embedded in these portable public recipes.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[youngeunkwon0405]

Team review of the two new 32-GPU Nano-Omni GRPO recipes (5 agents: RL expert, bug finder, test agent, devil's advocate; all findings adversarially verified against pinned submodule sources).

Great perf writeup — the W&B-linked optimization journey and 10-step functional validation made this easy to review. Both YAMLs pass tools/config_cli.py minimize-check, all overridden keys were verified to exist in the schema (including the vLLM kwargs against pinned vLLM 0.25.1 source), and the batch/parallelism/resource math all checks out.

2 blocking items (async entrypoint dispatch; recipe-accounting unit test), 1 real footgun (shared checkpoint dir), 3 low-severity consistency notes, 1 informational — all as inline comments.

Generated by Claude Code

[youngeunkwon0405]

/ok to test b00fb48

[youngeunkwon0405]

/okay to test f6091f6

[youngeunkwon0405]

@DanialTaheri, can I ask for your review of this nano omni recipe too?