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

examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-tp1ep8-async.v1.yaml

@@ -0,0 +1,28 @@
+defaults: ./vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-tp1ep8.v1.yaml
+
+grpo:
+  async_grpo:
+    enabled: true
+    in_flight_weight_updates: true
+
+checkpointing:
+  checkpoint_dir: results/vlm_grpo_nemotron_omni_8n4g_megatron_async
+
+policy:
+  generation:
+    colocated:
+      enabled: false
+      resources:
+        gpus_per_node: 4
+        num_nodes: 4
+    vllm_cfg:
+      gpu_memory_utilization: 0.8
+    vllm_kwargs:
+      max_num_seqs: 128
+
+logger:
+  wandb:
+    name: nemotron-omni-8n4g-async-1-off
+
+cluster:
+  segment_size: 4

examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-tp1ep8.v1.yaml

@@ -0,0 +1,54 @@
+defaults: ./vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8.v1.yaml
+
+grpo:
+  num_prompts_per_step: 64
+
+loss_fn:
+  use_importance_sampling_correction: true
+  force_on_policy_ratio: true
+
+checkpointing:
+  checkpoint_dir: results/vlm_grpo_nemotron_omni_8n4g_megatron
+
+policy:
+  is_vlm: true
+  train_global_batch_size: 1024
+  generation_batch_size: 1024
+  logprob_chunk_size: 1024
+  make_sequence_length_divisible_by: 256
+  tokenizer:
+    chat_template_kwargs:
+      enable_thinking: true
+      truncate_history_thinking: false
+  sequence_packing:
+    train_mb_tokens: 10240
+  megatron_cfg:
+    empty_unused_memory_level: 2
+    tensor_model_parallel_size: 1
+    activation_checkpointing: false
+    defer_fp32_logits: true
+    use_fused_weighted_squared_relu: true
+  generation:
+    vllm_cfg:
+      async_engine: true
+      tensor_parallel_size: 2
+      expert_parallel_size: 2
+      enforce_eager: false
+    vllm_kwargs:
+      max_num_batched_tokens: 32768
+      max_num_seqs: 64
+      all2all_backend: flashinfer_nvlink_one_sided
+      kernel_config:
+        moe_backend: flashinfer_cutlass
+        enable_flashinfer_autotune: true
+
+logger:
+  wandb_enabled: true
+  tensorboard_enabled: false
+  wandb:
+    name: nemotron-omni-8n4g-sync
+
+cluster:
+  num_nodes: 8
+  gpus_per_node: 4
+  segment_size: 8

tests/test_suites/disabled.txt

@@ -14,6 +14,12 @@ tests/test_suites/vlm/vlm_grpo-qwen3.5-35ba3b-geo3k-2n8g-megatron-ep16.sh
 # budget has room.
 tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-click-2n8g-megatron-tp2ep8.v1.sh
 
+# Nano-Omni 32-GPU performance-reference recipes. Both topologies were
+# validated manually; keep them out of nightly CI because each reserves eight
+# GB200 nodes.
+tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-tp1ep8.v1.sh
+tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-tp1ep8-async.v1.sh
+
 # Nemotron Super Omni: 16-node topology, and the checkpoint and multimodal Gym
 # blend are too large to ship with the repo, so these are invoked manually via
 # examples/nemo_gym/nemotron-3-super-omni/super_omni_launch.sh rather than run

tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-tp1ep8-async.v1.sh

@@ -0,0 +1,42 @@
+#!/bin/bash
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd)
+source $SCRIPT_DIR/common.env
+
+# ===== BEGIN CONFIG =====
+NUM_NODES=8
+GPUS_PER_NODE=4
+STEPS_PER_RUN=10
+MAX_STEPS=10
+NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN ))  # Round up
+NUM_MINUTES=120
+# ===== END CONFIG =====
+
+exit_if_max_steps_reached
+
+# Run the experiment
+cd $PROJECT_ROOT
+uv run examples/run_vlm_grpo.py \
+    --config $CONFIG_PATH \
+    grpo.max_num_steps=$MAX_STEPS \
+    logger.log_dir=$LOG_DIR \
+    logger.wandb_enabled=True \
+    logger.wandb.project=nemo-rl \
+    logger.wandb.name=$EXP_NAME \
+    logger.monitor_gpus=True \
+    logger.tensorboard_enabled=True \
+    checkpointing.enabled=True \
+    checkpointing.checkpoint_dir=$CKPT_DIR \
+    $@ \
+    2>&1 | tee $RUN_LOG
+
+# Convert tensorboard logs to json
+uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS
+
+# Only run metrics if the target step is reached
+if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | map(tonumber) | max' $JSON_METRICS) -ge $MAX_STEPS ]]; then
+    uv run tests/check_metrics.py $JSON_METRICS \
+        'max(data["train/reward"]) > 0.5'
+
+    # Clean up checkpoint directory after successful run to save space.
+    rm -rf "$CKPT_DIR"
+fi

tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-tp1ep8.v1.sh

@@ -0,0 +1,42 @@
+#!/bin/bash
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd)
+source $SCRIPT_DIR/common.env
+
+# ===== BEGIN CONFIG =====
+NUM_NODES=8
+GPUS_PER_NODE=4
+STEPS_PER_RUN=10
+MAX_STEPS=10
+NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN ))  # Round up
+NUM_MINUTES=120
+# ===== END CONFIG =====
+
+exit_if_max_steps_reached
+
+# Run the experiment
+cd $PROJECT_ROOT
+uv run examples/run_vlm_grpo.py \
+    --config $CONFIG_PATH \
+    grpo.max_num_steps=$MAX_STEPS \
+    logger.log_dir=$LOG_DIR \
+    logger.wandb_enabled=True \
+    logger.wandb.project=nemo-rl \
+    logger.wandb.name=$EXP_NAME \
+    logger.monitor_gpus=True \
+    logger.tensorboard_enabled=True \
+    checkpointing.enabled=True \
+    checkpointing.checkpoint_dir=$CKPT_DIR \
+    $@ \
+    2>&1 | tee $RUN_LOG
+
+# Convert tensorboard logs to json
+uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS
+
+# Only run metrics if the target step is reached
+if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | map(tonumber) | max' $JSON_METRICS) -ge $MAX_STEPS ]]; then
+    uv run tests/check_metrics.py $JSON_METRICS \
+        'max(data["train/reward"]) > 0.5'
+
+    # Clean up checkpoint directory after successful run to save space.
+    rm -rf "$CKPT_DIR"
+fi