diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-tp1ep8-async.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-tp1ep8-async.v1.yaml new file mode 100644 index 00000000000..0d5682bc295 --- /dev/null +++ b/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 diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-tp1ep8.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-tp1ep8.v1.yaml new file mode 100644 index 00000000000..ce1818158ce --- /dev/null +++ b/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 diff --git a/tests/test_suites/disabled.txt b/tests/test_suites/disabled.txt index 2ad14580107..9c5dd5a9b14 100644 --- a/tests/test_suites/disabled.txt +++ b/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 diff --git a/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-tp1ep8-async.v1.sh b/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-tp1ep8-async.v1.sh new file mode 100755 index 00000000000..df84c24bc5a --- /dev/null +++ b/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 diff --git a/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-tp1ep8.v1.sh b/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-tp1ep8.v1.sh new file mode 100755 index 00000000000..df84c24bc5a --- /dev/null +++ b/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