From 5bb259b37d62b7ce0f963175030d69b1dc343acc Mon Sep 17 00:00:00 2001 From: Cory Ye Date: Fri, 21 Aug 2026 20:17:22 -0700 Subject: [PATCH 01/19] Pass pre-expanded media tokens to MLLM. Signed-off-by: Cory Ye --- nemo_rl/models/generation/megatron/config.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/nemo_rl/models/generation/megatron/config.py b/nemo_rl/models/generation/megatron/config.py index 98d0733a80e..f07ca54bc25 100644 --- a/nemo_rl/models/generation/megatron/config.py +++ b/nemo_rl/models/generation/megatron/config.py @@ -44,6 +44,8 @@ class MCoreGenerationSpecificArgs(TypedDict): # - 'block': graphs are owned at the enclosing block (TransformerBlock / HybridBlock). # Only meaningful when cuda_graph_impl='local'. inference_cuda_graph_scope: NotRequired[str] + # Required for EP>1 + local CUDA graphs. + moe_pad_experts_for_cuda_graph_inference: NotRequired[bool] materialize_only_last_token_logits: bool enable_chunked_prefill: bool From 62da428740872c11bfb337ea627127a5ee9f8a02 Mon Sep 17 00:00:00 2001 From: Cory Ye Date: Wed, 26 Aug 2026 14:03:20 -0700 Subject: [PATCH 02/19] Add tests and lint. Signed-off-by: Cory Ye --- .../L1_Functional_Tests_Megatron_Omni.sh | 38 +++++ .../nemotron_omni_clevr_megatron_1n4g.sh | 145 +++++++++++++++++ .../nemotron_omni_gym_video_megatron_1n4g.sh | 154 ++++++++++++++++++ 3 files changed, 337 insertions(+) create mode 100755 tests/functional/L1_Functional_Tests_Megatron_Omni.sh create mode 100755 tests/functional/nemotron_omni_clevr_megatron_1n4g.sh create mode 100755 tests/functional/nemotron_omni_gym_video_megatron_1n4g.sh diff --git a/tests/functional/L1_Functional_Tests_Megatron_Omni.sh b/tests/functional/L1_Functional_Tests_Megatron_Omni.sh new file mode 100755 index 00000000000..1d95ae917ed --- /dev/null +++ b/tests/functional/L1_Functional_Tests_Megatron_Omni.sh @@ -0,0 +1,38 @@ +#!/bin/bash +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +set -xeuo pipefail + +SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &> /dev/null && pwd) +PROJECT_ROOT=$(realpath "${SCRIPT_DIR}/../..") + +cd "${PROJECT_ROOT}" + +GPU_COUNT=$(nvidia-smi --query-gpu=index --format=csv,noheader | wc -l) +if (( GPU_COUNT < 4 )); then + echo "SKIP: Nemotron Omni functional tests require at least four GPUs" + exit 0 +fi + +# The recipes are intentionally 1n4g even when the CI runner exposes eight GPUs. +export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3}" + +time uv run --no-sync bash ./tests/functional/nemotron_omni_clevr_megatron_1n4g.sh +time uv run --no-sync bash ./tests/functional/nemotron_omni_gym_video_megatron_1n4g.sh + +cd "${PROJECT_ROOT}/tests" +if compgen -G ".coverage*" > /dev/null; then + coverage combine .coverage* +fi diff --git a/tests/functional/nemotron_omni_clevr_megatron_1n4g.sh b/tests/functional/nemotron_omni_clevr_megatron_1n4g.sh new file mode 100755 index 00000000000..49a02dc97e1 --- /dev/null +++ b/tests/functional/nemotron_omni_clevr_megatron_1n4g.sh @@ -0,0 +1,145 @@ +#!/usr/bin/env bash +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +set -euo pipefail + +SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd) +PROJECT_ROOT=$(realpath "${SCRIPT_DIR}/../..") + +if [[ -z "${HF_TOKEN:-}" ]]; then + echo "SKIP: HF_TOKEN is required for the Omni checkpoint" + exit 0 +fi + +GPU_COUNT=$(nvidia-smi --query-gpu=index --format=csv,noheader | wc -l) +if (( GPU_COUNT < 4 )); then + echo "SKIP: Omni CLEVR Megatron smoke requires at least four visible GPUs" + exit 0 +fi +DETECTED_CUDA_ARCH=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader -i 0) +export TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST:-${DETECTED_CUDA_ARCH}}" +MEGATRON_TRANSFORMER_IMPL="${MEGATRON_TRANSFORMER_IMPL:-inference_optimized}" +MEGATRON_CUDA_GRAPH_IMPL="${MEGATRON_CUDA_GRAPH_IMPL:-local}" +if [[ "${MEGATRON_TRANSFORMER_IMPL}" != "inference_optimized" && + "${MEGATRON_CUDA_GRAPH_IMPL}" == "local" ]]; then + MOE_PAD_EXPERTS_FOR_CG=true +else + MOE_PAD_EXPERTS_FOR_CG=false +fi + +EXP_NAME=$(basename "$0" .sh) +EXP_DIR="${SCRIPT_DIR}/${EXP_NAME}" +LOG_DIR="${EXP_DIR}/logs" +DATA_ROOT="${EXP_DIR}/data" +TRAIN_PATH="${DATA_ROOT}/train.jsonl" +VAL_PATH="${DATA_ROOT}/val.jsonl" +JSON_METRICS="${EXP_DIR}/metrics.json" +RUN_LOG="${EXP_DIR}/run.log" +rm -rf "${EXP_DIR}" +mkdir -p "${LOG_DIR}" "${DATA_ROOT}" + +cd "${PROJECT_ROOT}" +export PYTHONPATH="${PROJECT_ROOT}:${PYTHONPATH:-}" + +# Use a tiny local image dataset. Downloading the full 70K CLEVR training split +# adds several minutes to a one-step smoke and does not improve E2E coverage. +TRAIN_PATH="${TRAIN_PATH}" VAL_PATH="${VAL_PATH}" uv run --no-sync python - <<'PY' +import base64 +import io +import json +import os + +from PIL import Image + +buffer = io.BytesIO() +Image.new("RGB", (224, 224), color="red").save(buffer, format="PNG") +image_url = "data:image/png;base64," + base64.b64encode(buffer.getvalue()).decode() + +def sample(index: int) -> dict: + return { + "messages": [ + { + "role": "user", + "content": [ + {"type": "image", "image": image_url}, + { + "type": "text", + "text": f"Sample {index}: What color is the image?", + }, + ], + }, + {"role": "assistant", "content": "red"}, + ] + } + +for path, count in ((os.environ["TRAIN_PATH"], 64), (os.environ["VAL_PATH"], 2)): + with open(path, "w") as output: + for index in range(count): + output.write(json.dumps(sample(index)) + "\n") +PY + +uv run --no-sync python examples/run_vlm_grpo.py \ + --config examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.yaml \ + cluster.num_nodes=1 \ + cluster.gpus_per_node=4 \ + policy.megatron_cfg.env_vars.TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST}" \ + policy.megatron_cfg.tensor_model_parallel_size=4 \ + policy.megatron_cfg.expert_model_parallel_size=4 \ + policy.megatron_cfg.expert_tensor_parallel_size=1 \ + policy.megatron_cfg.context_parallel_size=1 \ + policy.megatron_cfg.sequence_parallel=true \ + policy.megatron_cfg.activation_checkpointing=true \ + policy.megatron_cfg.optimizer.optimizer_cpu_offload=false \ + policy.megatron_cfg.optimizer.optimizer_offload_fraction=0.0 \ + ++policy.megatron_cfg.optimizer.exp_avg_dtype=bfloat16 \ + ++policy.megatron_cfg.optimizer.exp_avg_sq_dtype=bfloat16 \ + ++policy.megatron_cfg.optimizer.store_param_remainders=true \ + policy.generation.backend=megatron \ + policy.generation.colocated.enabled=true \ + policy.generation.colocated.resources.num_nodes=1 \ + policy.generation.colocated.resources.gpus_per_node=4 \ + policy.generation.max_new_tokens=128 \ + policy.generation.mcore_generation_config.tensor_model_parallel_size=4 \ + policy.generation.mcore_generation_config.expert_model_parallel_size=4 \ + policy.generation.mcore_generation_config.expert_tensor_parallel_size=1 \ + ++policy.generation.mcore_generation_config.context_parallel_size=1 \ + ++policy.generation.mcore_generation_config.moe_router_dtype=fp32 \ + policy.generation.mcore_generation_config.transformer_impl="${MEGATRON_TRANSFORMER_IMPL}" \ + policy.generation.mcore_generation_config.sequence_parallel=true \ + policy.generation.mcore_generation_config.refit_backend=nccl \ + policy.generation.mcore_generation_config.cuda_graph_impl="${MEGATRON_CUDA_GRAPH_IMPL}" \ + policy.generation.mcore_generation_config.inference_cuda_graph_scope=block \ + policy.generation.mcore_generation_config.num_cuda_graphs=-1 \ + policy.generation.mcore_generation_config.use_cuda_graphs_for_non_decode_steps=false \ + policy.generation.mcore_generation_config.moe_pad_experts_for_cuda_graph_inference="${MOE_PAD_EXPERTS_FOR_CG}" \ + policy.generation.mcore_generation_config.enable_chunked_prefill=true \ + ++policy.generation.mcore_generation_config.async_sched_mode=async \ + data.train.dataset_name=ResponseDataset \ + ++data.train.data_path="${TRAIN_PATH}" \ + data.train.split=train \ + data.validation.dataset_name=ResponseDataset \ + ++data.validation.data_path="${VAL_PATH}" \ + data.validation.split=train \ + data.num_workers=0 \ + grpo.async_grpo.enabled=true \ + grpo.async_grpo.max_trajectory_age_steps=2 \ + grpo.async_grpo.in_flight_weight_updates=true \ + grpo.num_prompts_per_step=1 \ + grpo.num_generations_per_prompt=2 \ + grpo.max_num_steps=1 \ + grpo.val_period=0 \ + grpo.val_at_start=false \ + grpo.val_at_end=false \ + policy.train_global_batch_size=2 \ + policy.train_micro_batch_size=1 \ + logger.tensorboard_enabled=true \ + logger.log_dir="${LOG_DIR}" \ + logger.wandb_enabled=false \ + logger.monitor_gpus=false \ + checkpointing.enabled=false \ + "$@" 2>&1 | tee "${RUN_LOG}" + +uv run --no-sync tests/json_dump_tb_logs.py "${LOG_DIR}" --output_path "${JSON_METRICS}" +uv run --no-sync tests/check_metrics.py "${JSON_METRICS}" \ + 'max(data["train/loss"]) < 1e6' \ + 'min(data["train/loss"]) > -1e6' diff --git a/tests/functional/nemotron_omni_gym_video_megatron_1n4g.sh b/tests/functional/nemotron_omni_gym_video_megatron_1n4g.sh new file mode 100755 index 00000000000..798d3b735fa --- /dev/null +++ b/tests/functional/nemotron_omni_gym_video_megatron_1n4g.sh @@ -0,0 +1,154 @@ +#!/usr/bin/env bash +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +set -euo pipefail + +SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd) +PROJECT_ROOT=$(realpath "${SCRIPT_DIR}/../..") + +if [[ -z "${HF_TOKEN:-}" ]]; then + echo "SKIP: HF_TOKEN is required for the Omni checkpoint" + exit 0 +fi + +GPU_COUNT=$(nvidia-smi --query-gpu=index --format=csv,noheader | wc -l) +if (( GPU_COUNT < 4 )); then + echo "SKIP: Omni Gym-video Megatron smoke requires at least four GPUs" + exit 0 +fi +DETECTED_CUDA_ARCH=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader -i 0) +export TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST:-${DETECTED_CUDA_ARCH}}" +MEGATRON_TRANSFORMER_IMPL="${MEGATRON_TRANSFORMER_IMPL:-inference_optimized}" +MOE_PAD_EXPERTS_FOR_CG=false + +EXP_NAME=$(basename "$0" .sh) +EXP_DIR="${SCRIPT_DIR}/${EXP_NAME}" +LOG_DIR="${EXP_DIR}/logs" +DATA_ROOT="${EXP_DIR}/data" +VIDEO_PATH="${DATA_ROOT}/red.mp4" +RAW_TRAIN_PATH="${DATA_ROOT}/train-raw.jsonl" +RAW_VAL_PATH="${DATA_ROOT}/val-raw.jsonl" +TRAIN_PATH="${DATA_ROOT}/train-gym.jsonl" +VAL_PATH="${DATA_ROOT}/val-gym.jsonl" +JSON_METRICS="${EXP_DIR}/metrics.json" +RUN_LOG="${EXP_DIR}/run.log" +rm -rf "${EXP_DIR}" +mkdir -p "${LOG_DIR}" "${DATA_ROOT}" + +cd "${PROJECT_ROOT}" +export PYTHONPATH="${PROJECT_ROOT}:${PYTHONPATH:-}" +export NRL_VIDEO_BACKEND=torchcodec +export NRL_VIDEO_SAMPLING_STYLE=nemotron_vl +export NRL_VIDEO_TEMPORAL_PATCH_SIZE=2 + +bash tools/install_audio_deps.sh +ffmpeg -hide_banner -loglevel error -y \ + -f lavfi -i color=c=red:s=224x224:r=8:d=2 \ + -c:v libx264 -pix_fmt yuv420p "${VIDEO_PATH}" + +for sample_id in $(seq 1 64); do + jq -nc \ + --arg prompt "Sample ${sample_id}: What color fills the video? A. Red B. Blue" \ + --arg video "${VIDEO_PATH}" \ + '{prompt: $prompt, video: $video, answer: "A", verifier: "mcqa"}' +done > "${RAW_TRAIN_PATH}" +for sample_id in $(seq 1 2); do + jq -nc \ + --arg prompt "Validation ${sample_id}: What color fills the video? A. Red B. Blue" \ + --arg video "${VIDEO_PATH}" \ + '{prompt: $prompt, video: $video, answer: "A", verifier: "mcqa"}' +done > "${RAW_VAL_PATH}" + +uv run --no-sync examples/nemo_gym/prepare_video_dataset.py convert \ + --input "${RAW_TRAIN_PATH}" \ + --output "${TRAIN_PATH}" +uv run --no-sync examples/nemo_gym/prepare_video_dataset.py convert \ + --input "${RAW_VAL_PATH}" \ + --output "${VAL_PATH}" + +uv run --no-sync python examples/nemo_gym/run_grpo_nemo_gym.py \ + --config examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-16n8g-megatron-tp4ep4-async-gym-video.v1.yaml \ + cluster.num_nodes=1 \ + cluster.gpus_per_node=4 \ + policy.megatron_cfg.env_vars.TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST}" \ + policy.megatron_cfg.tensor_model_parallel_size=2 \ + policy.megatron_cfg.pipeline_model_parallel_size=1 \ + policy.megatron_cfg.expert_model_parallel_size=2 \ + policy.megatron_cfg.expert_tensor_parallel_size=1 \ + policy.megatron_cfg.context_parallel_size=1 \ + policy.megatron_cfg.sequence_parallel=true \ + policy.megatron_cfg.activation_checkpointing=true \ + policy.megatron_cfg.optimizer.optimizer_cpu_offload=false \ + policy.megatron_cfg.optimizer.optimizer_offload_fraction=0.0 \ + ++policy.megatron_cfg.optimizer.exp_avg_dtype=bfloat16 \ + ++policy.megatron_cfg.optimizer.exp_avg_sq_dtype=bfloat16 \ + ++policy.megatron_cfg.optimizer.store_param_remainders=true \ + policy.generation.backend=megatron \ + ++policy.generation.bad_words=null \ + policy.generation.colocated.enabled=false \ + policy.generation.colocated.resources.num_nodes=1 \ + policy.generation.colocated.resources.gpus_per_node=2 \ + policy.generation.max_new_tokens=128 \ + policy.generation.mcore_generation_config.expose_http_server=true \ + policy.generation.mcore_generation_config.tensor_model_parallel_size=2 \ + policy.generation.mcore_generation_config.expert_model_parallel_size=2 \ + policy.generation.mcore_generation_config.expert_tensor_parallel_size=1 \ + ++policy.generation.mcore_generation_config.context_parallel_size=1 \ + ++policy.generation.mcore_generation_config.moe_router_dtype=fp32 \ + policy.generation.mcore_generation_config.transformer_impl="${MEGATRON_TRANSFORMER_IMPL}" \ + policy.generation.mcore_generation_config.sequence_parallel=true \ + policy.generation.mcore_generation_config.refit_backend=nccl \ + policy.generation.mcore_generation_config.buffer_size_gb=8 \ + policy.generation.mcore_generation_config.cuda_graph_impl=none \ + policy.generation.mcore_generation_config.inference_cuda_graph_scope=none \ + policy.generation.mcore_generation_config.num_cuda_graphs=0 \ + policy.generation.mcore_generation_config.use_cuda_graphs_for_non_decode_steps=false \ + ++policy.generation.mcore_generation_config.moe_pad_experts_for_cuda_graph_inference="${MOE_PAD_EXPERTS_FOR_CG}" \ + policy.generation.mcore_generation_config.enable_chunked_prefill=true \ + ++policy.generation.mcore_generation_config.async_sched_mode=async \ + policy.generation.mcore_generation_config.enable_prefix_caching=true \ + policy.generation.mcore_generation_config.max_model_len=4096 \ + policy.generation.mcore_generation_config.max_tokens=4096 \ + ++policy.generation.mcore_generation_config.video_num_frames=8 \ + ++policy.generation.mcore_generation_config.video_temporal_patch_size=2 \ + ++policy.generation.mcore_generation_config.video_target_num_patches=256 \ + policy.max_total_sequence_length=4096 \ + +data.default.num_frames=8 \ + +data.default.video_sampling_style=nemotron_vl \ + +data.default.video_temporal_patch_size=2 \ + +data.default.min_generation_tokens=128 \ + data.default.video_target_num_patches=256 \ + data.train.data_path="${TRAIN_PATH}" \ + data.validation.data_path="${VAL_PATH}" \ + grpo.deduplicate_multimodal_data=false \ + grpo.async_grpo.enabled=true \ + grpo.async_grpo.max_trajectory_age_steps=2 \ + grpo.async_grpo.in_flight_weight_updates=true \ + grpo.num_prompts_per_step=1 \ + grpo.num_generations_per_prompt=2 \ + grpo.max_num_steps=1 \ + grpo.val_period=0 \ + grpo.val_at_start=false \ + grpo.val_at_end=false \ + policy.train_global_batch_size=2 \ + policy.train_micro_batch_size=1 \ + logger.tensorboard_enabled=true \ + logger.log_dir="${LOG_DIR}" \ + logger.wandb_enabled=false \ + logger.monitor_gpus=false \ + checkpointing.enabled=false \ + "$@" 2>&1 | tee "${RUN_LOG}" + +uv run --no-sync tests/json_dump_tb_logs.py "${LOG_DIR}" --output_path "${JSON_METRICS}" + +RECORDED_STEP=$(jq -r \ + 'if has("train/loss") then (."train/loss" | keys | map(tonumber) | max // 0) else 0 end' \ + "${JSON_METRICS}") +if (( RECORDED_STEP < 1 )); then + echo "[ERROR] Expected at least one completed Gym-video training step" + exit 1 +fi + +uv run --no-sync tests/check_metrics.py "${JSON_METRICS}" \ + 'max(data["train/loss"]) < 1e6' \ + 'min(data["train/loss"]) > -1e6' From 2152f02b706c5178c1ef18375cf4d1bfcdcb93ca Mon Sep 17 00:00:00 2001 From: Cory Ye Date: Sun, 30 Aug 2026 16:10:19 -0700 Subject: [PATCH 03/19] Address more review feedback. Signed-off-by: Cory Ye --- .../L1_Functional_Tests_Megatron_Omni.sh | 38 ----- .../nemotron_omni_clevr_megatron_1n4g.sh | 145 ----------------- .../nemotron_omni_gym_video_megatron_1n4g.sh | 154 ------------------ ...0ba3b-clevr-1n4g-megatron_generation.v1.sh | 34 ++++ 4 files changed, 34 insertions(+), 337 deletions(-) delete mode 100755 tests/functional/L1_Functional_Tests_Megatron_Omni.sh delete mode 100755 tests/functional/nemotron_omni_clevr_megatron_1n4g.sh delete mode 100755 tests/functional/nemotron_omni_gym_video_megatron_1n4g.sh create mode 100755 tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.sh diff --git a/tests/functional/L1_Functional_Tests_Megatron_Omni.sh b/tests/functional/L1_Functional_Tests_Megatron_Omni.sh deleted file mode 100755 index 1d95ae917ed..00000000000 --- a/tests/functional/L1_Functional_Tests_Megatron_Omni.sh +++ /dev/null @@ -1,38 +0,0 @@ -#!/bin/bash -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -set -xeuo pipefail - -SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &> /dev/null && pwd) -PROJECT_ROOT=$(realpath "${SCRIPT_DIR}/../..") - -cd "${PROJECT_ROOT}" - -GPU_COUNT=$(nvidia-smi --query-gpu=index --format=csv,noheader | wc -l) -if (( GPU_COUNT < 4 )); then - echo "SKIP: Nemotron Omni functional tests require at least four GPUs" - exit 0 -fi - -# The recipes are intentionally 1n4g even when the CI runner exposes eight GPUs. -export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3}" - -time uv run --no-sync bash ./tests/functional/nemotron_omni_clevr_megatron_1n4g.sh -time uv run --no-sync bash ./tests/functional/nemotron_omni_gym_video_megatron_1n4g.sh - -cd "${PROJECT_ROOT}/tests" -if compgen -G ".coverage*" > /dev/null; then - coverage combine .coverage* -fi diff --git a/tests/functional/nemotron_omni_clevr_megatron_1n4g.sh b/tests/functional/nemotron_omni_clevr_megatron_1n4g.sh deleted file mode 100755 index 49a02dc97e1..00000000000 --- a/tests/functional/nemotron_omni_clevr_megatron_1n4g.sh +++ /dev/null @@ -1,145 +0,0 @@ -#!/usr/bin/env bash -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. - -set -euo pipefail - -SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd) -PROJECT_ROOT=$(realpath "${SCRIPT_DIR}/../..") - -if [[ -z "${HF_TOKEN:-}" ]]; then - echo "SKIP: HF_TOKEN is required for the Omni checkpoint" - exit 0 -fi - -GPU_COUNT=$(nvidia-smi --query-gpu=index --format=csv,noheader | wc -l) -if (( GPU_COUNT < 4 )); then - echo "SKIP: Omni CLEVR Megatron smoke requires at least four visible GPUs" - exit 0 -fi -DETECTED_CUDA_ARCH=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader -i 0) -export TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST:-${DETECTED_CUDA_ARCH}}" -MEGATRON_TRANSFORMER_IMPL="${MEGATRON_TRANSFORMER_IMPL:-inference_optimized}" -MEGATRON_CUDA_GRAPH_IMPL="${MEGATRON_CUDA_GRAPH_IMPL:-local}" -if [[ "${MEGATRON_TRANSFORMER_IMPL}" != "inference_optimized" && - "${MEGATRON_CUDA_GRAPH_IMPL}" == "local" ]]; then - MOE_PAD_EXPERTS_FOR_CG=true -else - MOE_PAD_EXPERTS_FOR_CG=false -fi - -EXP_NAME=$(basename "$0" .sh) -EXP_DIR="${SCRIPT_DIR}/${EXP_NAME}" -LOG_DIR="${EXP_DIR}/logs" -DATA_ROOT="${EXP_DIR}/data" -TRAIN_PATH="${DATA_ROOT}/train.jsonl" -VAL_PATH="${DATA_ROOT}/val.jsonl" -JSON_METRICS="${EXP_DIR}/metrics.json" -RUN_LOG="${EXP_DIR}/run.log" -rm -rf "${EXP_DIR}" -mkdir -p "${LOG_DIR}" "${DATA_ROOT}" - -cd "${PROJECT_ROOT}" -export PYTHONPATH="${PROJECT_ROOT}:${PYTHONPATH:-}" - -# Use a tiny local image dataset. Downloading the full 70K CLEVR training split -# adds several minutes to a one-step smoke and does not improve E2E coverage. -TRAIN_PATH="${TRAIN_PATH}" VAL_PATH="${VAL_PATH}" uv run --no-sync python - <<'PY' -import base64 -import io -import json -import os - -from PIL import Image - -buffer = io.BytesIO() -Image.new("RGB", (224, 224), color="red").save(buffer, format="PNG") -image_url = "data:image/png;base64," + base64.b64encode(buffer.getvalue()).decode() - -def sample(index: int) -> dict: - return { - "messages": [ - { - "role": "user", - "content": [ - {"type": "image", "image": image_url}, - { - "type": "text", - "text": f"Sample {index}: What color is the image?", - }, - ], - }, - {"role": "assistant", "content": "red"}, - ] - } - -for path, count in ((os.environ["TRAIN_PATH"], 64), (os.environ["VAL_PATH"], 2)): - with open(path, "w") as output: - for index in range(count): - output.write(json.dumps(sample(index)) + "\n") -PY - -uv run --no-sync python examples/run_vlm_grpo.py \ - --config examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.yaml \ - cluster.num_nodes=1 \ - cluster.gpus_per_node=4 \ - policy.megatron_cfg.env_vars.TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST}" \ - policy.megatron_cfg.tensor_model_parallel_size=4 \ - policy.megatron_cfg.expert_model_parallel_size=4 \ - policy.megatron_cfg.expert_tensor_parallel_size=1 \ - policy.megatron_cfg.context_parallel_size=1 \ - policy.megatron_cfg.sequence_parallel=true \ - policy.megatron_cfg.activation_checkpointing=true \ - policy.megatron_cfg.optimizer.optimizer_cpu_offload=false \ - policy.megatron_cfg.optimizer.optimizer_offload_fraction=0.0 \ - ++policy.megatron_cfg.optimizer.exp_avg_dtype=bfloat16 \ - ++policy.megatron_cfg.optimizer.exp_avg_sq_dtype=bfloat16 \ - ++policy.megatron_cfg.optimizer.store_param_remainders=true \ - policy.generation.backend=megatron \ - policy.generation.colocated.enabled=true \ - policy.generation.colocated.resources.num_nodes=1 \ - policy.generation.colocated.resources.gpus_per_node=4 \ - policy.generation.max_new_tokens=128 \ - policy.generation.mcore_generation_config.tensor_model_parallel_size=4 \ - policy.generation.mcore_generation_config.expert_model_parallel_size=4 \ - policy.generation.mcore_generation_config.expert_tensor_parallel_size=1 \ - ++policy.generation.mcore_generation_config.context_parallel_size=1 \ - ++policy.generation.mcore_generation_config.moe_router_dtype=fp32 \ - policy.generation.mcore_generation_config.transformer_impl="${MEGATRON_TRANSFORMER_IMPL}" \ - policy.generation.mcore_generation_config.sequence_parallel=true \ - policy.generation.mcore_generation_config.refit_backend=nccl \ - policy.generation.mcore_generation_config.cuda_graph_impl="${MEGATRON_CUDA_GRAPH_IMPL}" \ - policy.generation.mcore_generation_config.inference_cuda_graph_scope=block \ - policy.generation.mcore_generation_config.num_cuda_graphs=-1 \ - policy.generation.mcore_generation_config.use_cuda_graphs_for_non_decode_steps=false \ - policy.generation.mcore_generation_config.moe_pad_experts_for_cuda_graph_inference="${MOE_PAD_EXPERTS_FOR_CG}" \ - policy.generation.mcore_generation_config.enable_chunked_prefill=true \ - ++policy.generation.mcore_generation_config.async_sched_mode=async \ - data.train.dataset_name=ResponseDataset \ - ++data.train.data_path="${TRAIN_PATH}" \ - data.train.split=train \ - data.validation.dataset_name=ResponseDataset \ - ++data.validation.data_path="${VAL_PATH}" \ - data.validation.split=train \ - data.num_workers=0 \ - grpo.async_grpo.enabled=true \ - grpo.async_grpo.max_trajectory_age_steps=2 \ - grpo.async_grpo.in_flight_weight_updates=true \ - grpo.num_prompts_per_step=1 \ - grpo.num_generations_per_prompt=2 \ - grpo.max_num_steps=1 \ - grpo.val_period=0 \ - grpo.val_at_start=false \ - grpo.val_at_end=false \ - policy.train_global_batch_size=2 \ - policy.train_micro_batch_size=1 \ - logger.tensorboard_enabled=true \ - logger.log_dir="${LOG_DIR}" \ - logger.wandb_enabled=false \ - logger.monitor_gpus=false \ - checkpointing.enabled=false \ - "$@" 2>&1 | tee "${RUN_LOG}" - -uv run --no-sync tests/json_dump_tb_logs.py "${LOG_DIR}" --output_path "${JSON_METRICS}" -uv run --no-sync tests/check_metrics.py "${JSON_METRICS}" \ - 'max(data["train/loss"]) < 1e6' \ - 'min(data["train/loss"]) > -1e6' diff --git a/tests/functional/nemotron_omni_gym_video_megatron_1n4g.sh b/tests/functional/nemotron_omni_gym_video_megatron_1n4g.sh deleted file mode 100755 index 798d3b735fa..00000000000 --- a/tests/functional/nemotron_omni_gym_video_megatron_1n4g.sh +++ /dev/null @@ -1,154 +0,0 @@ -#!/usr/bin/env bash -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. - -set -euo pipefail - -SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd) -PROJECT_ROOT=$(realpath "${SCRIPT_DIR}/../..") - -if [[ -z "${HF_TOKEN:-}" ]]; then - echo "SKIP: HF_TOKEN is required for the Omni checkpoint" - exit 0 -fi - -GPU_COUNT=$(nvidia-smi --query-gpu=index --format=csv,noheader | wc -l) -if (( GPU_COUNT < 4 )); then - echo "SKIP: Omni Gym-video Megatron smoke requires at least four GPUs" - exit 0 -fi -DETECTED_CUDA_ARCH=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader -i 0) -export TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST:-${DETECTED_CUDA_ARCH}}" -MEGATRON_TRANSFORMER_IMPL="${MEGATRON_TRANSFORMER_IMPL:-inference_optimized}" -MOE_PAD_EXPERTS_FOR_CG=false - -EXP_NAME=$(basename "$0" .sh) -EXP_DIR="${SCRIPT_DIR}/${EXP_NAME}" -LOG_DIR="${EXP_DIR}/logs" -DATA_ROOT="${EXP_DIR}/data" -VIDEO_PATH="${DATA_ROOT}/red.mp4" -RAW_TRAIN_PATH="${DATA_ROOT}/train-raw.jsonl" -RAW_VAL_PATH="${DATA_ROOT}/val-raw.jsonl" -TRAIN_PATH="${DATA_ROOT}/train-gym.jsonl" -VAL_PATH="${DATA_ROOT}/val-gym.jsonl" -JSON_METRICS="${EXP_DIR}/metrics.json" -RUN_LOG="${EXP_DIR}/run.log" -rm -rf "${EXP_DIR}" -mkdir -p "${LOG_DIR}" "${DATA_ROOT}" - -cd "${PROJECT_ROOT}" -export PYTHONPATH="${PROJECT_ROOT}:${PYTHONPATH:-}" -export NRL_VIDEO_BACKEND=torchcodec -export NRL_VIDEO_SAMPLING_STYLE=nemotron_vl -export NRL_VIDEO_TEMPORAL_PATCH_SIZE=2 - -bash tools/install_audio_deps.sh -ffmpeg -hide_banner -loglevel error -y \ - -f lavfi -i color=c=red:s=224x224:r=8:d=2 \ - -c:v libx264 -pix_fmt yuv420p "${VIDEO_PATH}" - -for sample_id in $(seq 1 64); do - jq -nc \ - --arg prompt "Sample ${sample_id}: What color fills the video? A. Red B. Blue" \ - --arg video "${VIDEO_PATH}" \ - '{prompt: $prompt, video: $video, answer: "A", verifier: "mcqa"}' -done > "${RAW_TRAIN_PATH}" -for sample_id in $(seq 1 2); do - jq -nc \ - --arg prompt "Validation ${sample_id}: What color fills the video? A. Red B. Blue" \ - --arg video "${VIDEO_PATH}" \ - '{prompt: $prompt, video: $video, answer: "A", verifier: "mcqa"}' -done > "${RAW_VAL_PATH}" - -uv run --no-sync examples/nemo_gym/prepare_video_dataset.py convert \ - --input "${RAW_TRAIN_PATH}" \ - --output "${TRAIN_PATH}" -uv run --no-sync examples/nemo_gym/prepare_video_dataset.py convert \ - --input "${RAW_VAL_PATH}" \ - --output "${VAL_PATH}" - -uv run --no-sync python examples/nemo_gym/run_grpo_nemo_gym.py \ - --config examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-16n8g-megatron-tp4ep4-async-gym-video.v1.yaml \ - cluster.num_nodes=1 \ - cluster.gpus_per_node=4 \ - policy.megatron_cfg.env_vars.TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST}" \ - policy.megatron_cfg.tensor_model_parallel_size=2 \ - policy.megatron_cfg.pipeline_model_parallel_size=1 \ - policy.megatron_cfg.expert_model_parallel_size=2 \ - policy.megatron_cfg.expert_tensor_parallel_size=1 \ - policy.megatron_cfg.context_parallel_size=1 \ - policy.megatron_cfg.sequence_parallel=true \ - policy.megatron_cfg.activation_checkpointing=true \ - policy.megatron_cfg.optimizer.optimizer_cpu_offload=false \ - policy.megatron_cfg.optimizer.optimizer_offload_fraction=0.0 \ - ++policy.megatron_cfg.optimizer.exp_avg_dtype=bfloat16 \ - ++policy.megatron_cfg.optimizer.exp_avg_sq_dtype=bfloat16 \ - ++policy.megatron_cfg.optimizer.store_param_remainders=true \ - policy.generation.backend=megatron \ - ++policy.generation.bad_words=null \ - policy.generation.colocated.enabled=false \ - policy.generation.colocated.resources.num_nodes=1 \ - policy.generation.colocated.resources.gpus_per_node=2 \ - policy.generation.max_new_tokens=128 \ - policy.generation.mcore_generation_config.expose_http_server=true \ - policy.generation.mcore_generation_config.tensor_model_parallel_size=2 \ - policy.generation.mcore_generation_config.expert_model_parallel_size=2 \ - policy.generation.mcore_generation_config.expert_tensor_parallel_size=1 \ - ++policy.generation.mcore_generation_config.context_parallel_size=1 \ - ++policy.generation.mcore_generation_config.moe_router_dtype=fp32 \ - policy.generation.mcore_generation_config.transformer_impl="${MEGATRON_TRANSFORMER_IMPL}" \ - policy.generation.mcore_generation_config.sequence_parallel=true \ - policy.generation.mcore_generation_config.refit_backend=nccl \ - policy.generation.mcore_generation_config.buffer_size_gb=8 \ - policy.generation.mcore_generation_config.cuda_graph_impl=none \ - policy.generation.mcore_generation_config.inference_cuda_graph_scope=none \ - policy.generation.mcore_generation_config.num_cuda_graphs=0 \ - policy.generation.mcore_generation_config.use_cuda_graphs_for_non_decode_steps=false \ - ++policy.generation.mcore_generation_config.moe_pad_experts_for_cuda_graph_inference="${MOE_PAD_EXPERTS_FOR_CG}" \ - policy.generation.mcore_generation_config.enable_chunked_prefill=true \ - ++policy.generation.mcore_generation_config.async_sched_mode=async \ - policy.generation.mcore_generation_config.enable_prefix_caching=true \ - policy.generation.mcore_generation_config.max_model_len=4096 \ - policy.generation.mcore_generation_config.max_tokens=4096 \ - ++policy.generation.mcore_generation_config.video_num_frames=8 \ - ++policy.generation.mcore_generation_config.video_temporal_patch_size=2 \ - ++policy.generation.mcore_generation_config.video_target_num_patches=256 \ - policy.max_total_sequence_length=4096 \ - +data.default.num_frames=8 \ - +data.default.video_sampling_style=nemotron_vl \ - +data.default.video_temporal_patch_size=2 \ - +data.default.min_generation_tokens=128 \ - data.default.video_target_num_patches=256 \ - data.train.data_path="${TRAIN_PATH}" \ - data.validation.data_path="${VAL_PATH}" \ - grpo.deduplicate_multimodal_data=false \ - grpo.async_grpo.enabled=true \ - grpo.async_grpo.max_trajectory_age_steps=2 \ - grpo.async_grpo.in_flight_weight_updates=true \ - grpo.num_prompts_per_step=1 \ - grpo.num_generations_per_prompt=2 \ - grpo.max_num_steps=1 \ - grpo.val_period=0 \ - grpo.val_at_start=false \ - grpo.val_at_end=false \ - policy.train_global_batch_size=2 \ - policy.train_micro_batch_size=1 \ - logger.tensorboard_enabled=true \ - logger.log_dir="${LOG_DIR}" \ - logger.wandb_enabled=false \ - logger.monitor_gpus=false \ - checkpointing.enabled=false \ - "$@" 2>&1 | tee "${RUN_LOG}" - -uv run --no-sync tests/json_dump_tb_logs.py "${LOG_DIR}" --output_path "${JSON_METRICS}" - -RECORDED_STEP=$(jq -r \ - 'if has("train/loss") then (."train/loss" | keys | map(tonumber) | max // 0) else 0 end' \ - "${JSON_METRICS}") -if (( RECORDED_STEP < 1 )); then - echo "[ERROR] Expected at least one completed Gym-video training step" - exit 1 -fi - -uv run --no-sync tests/check_metrics.py "${JSON_METRICS}" \ - 'max(data["train/loss"]) < 1e6' \ - 'min(data["train/loss"]) > -1e6' diff --git a/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.sh b/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.sh new file mode 100755 index 00000000000..f8d2e75cdf8 --- /dev/null +++ b/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.sh @@ -0,0 +1,34 @@ +#!/bin/bash +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +source $SCRIPT_DIR/common.env + +# TODO(@cspades): Run and validate this functional test, then add golden +# convergence metrics before enabling it in a recurring suite. + +# ===== BEGIN CONFIG ===== +NUM_NODES=1 +GPUS_PER_NODE=4 +STEPS_PER_RUN=4 +MAX_STEPS=4 +NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) +NUM_MINUTES=120 +# ===== END CONFIG ===== + +exit_if_max_steps_reached + +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 + +uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS From 6cf8a441d74cc1e80bd6ede3a0a21f5f29a0de59 Mon Sep 17 00:00:00 2001 From: Cory Ye Date: Thu, 27 Aug 2026 21:10:36 -0700 Subject: [PATCH 04/19] Support multimodal Megatron-Inference with NeMo-RL TransferQueue (V2). Signed-off-by: Cory Ye --- ...g-megatron-single-controller-async.v1.yaml | 51 ++++++++++++++++++ ...g-megatron-single-controller-async.v1.yaml | 46 ++++++++++++++++ ...g-megatron-single-controller-async.v1.yaml | 53 +++++++++++++++++++ examples/run_grpo_single_controller.py | 11 +++- .../single_controller_utils/setup.py | 36 ++++++++++--- nemo_rl/data_plane/interfaces.py | 9 +++- nemo_rl/experience/payload.py | 1 + nemo_rl/experience/rollout_manager.py | 15 ++++-- tests/unit/data_plane/test_kvbatchmeta.py | 28 ++++++++++ tests/unit/experience/test_rollout_manager.py | 51 ++++++++++++++++++ .../unit/single_controller/test_entrypoint.py | 29 +++++++++- tests/unit/single_controller/test_setup.py | 24 ++++++++- 12 files changed, 338 insertions(+), 16 deletions(-) create mode 100644 examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron-single-controller-async.v1.yaml create mode 100644 examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-single-controller-async.v1.yaml create mode 100644 examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-1n4g-megatron-single-controller-async.v1.yaml diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron-single-controller-async.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron-single-controller-async.v1.yaml new file mode 100644 index 00000000000..2705ab46b5b --- /dev/null +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron-single-controller-async.v1.yaml @@ -0,0 +1,51 @@ +# NeMo-RL v2 SingleController overlay for one-node, non-colocated Omni CLEVR. +defaults: ./vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.yaml + +grpo: + async_grpo: null + val_period: 0 + val_at_start: false + val_at_end: false + +data_plane: + enabled: true + impl: transfer_queue + backend: simple + claim_meta_poll_interval_s: 0.5 + simple: + num_storage_units: 2 + +async_rl: + sampler: + name: windowed + max_staleness_versions: 2 + recompute_kv_cache_after_weight_updates: false + min_groups_for_streaming_train: ${grpo.num_prompts_per_step} + max_inflight_prompts: ${grpo.num_prompts_per_step} + max_buffered_rollouts: 32 + diagnostics: true + +policy: + generation: + backend: megatron + colocated: + enabled: false + resources: + num_nodes: 1 + gpus_per_node: 2 + mcore_generation_config: + transformer_impl: inference_optimized + moe_router_dtype: fp32 + moe_pad_experts_for_cuda_graph_inference: false + cuda_graph_impl: local + inference_cuda_graph_scope: block + num_cuda_graphs: -1 + use_cuda_graphs_for_non_decode_steps: false + enable_chunked_prefill: true + async_sched_mode: async + kv_cache_management_mode: persist + refit_backend: nccl + +cluster: + num_nodes: 1 + gpus_per_node: 4 diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-single-controller-async.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-single-controller-async.v1.yaml new file mode 100644 index 00000000000..22bc06daf50 --- /dev/null +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-single-controller-async.v1.yaml @@ -0,0 +1,46 @@ +# NeMo-RL v2 SingleController overlay for eight-node, non-colocated Omni CLEVR. +defaults: ./vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron_generation.v1.yaml + +grpo: + async_grpo: null + val_period: 0 + val_at_start: false + val_at_end: false + +data_plane: + enabled: true + +async_rl: + sampler: + name: windowed + max_staleness_versions: 2 + recompute_kv_cache_after_weight_updates: false + min_groups_for_streaming_train: ${grpo.num_prompts_per_step} + max_inflight_prompts: ${grpo.num_prompts_per_step} + max_buffered_rollouts: 32 + diagnostics: true + +policy: + generation: + backend: megatron + colocated: + enabled: false + resources: + num_nodes: 6 + gpus_per_node: 4 + mcore_generation_config: + transformer_impl: inference_optimized + moe_router_dtype: fp32 + moe_pad_experts_for_cuda_graph_inference: false + cuda_graph_impl: local + inference_cuda_graph_scope: block + num_cuda_graphs: -1 + use_cuda_graphs_for_non_decode_steps: false + enable_chunked_prefill: true + async_sched_mode: async + kv_cache_management_mode: persist + refit_backend: nccl + +cluster: + num_nodes: 8 + gpus_per_node: 4 diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-1n4g-megatron-single-controller-async.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-1n4g-megatron-single-controller-async.v1.yaml new file mode 100644 index 00000000000..f3b8e8aee68 --- /dev/null +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-1n4g-megatron-single-controller-async.v1.yaml @@ -0,0 +1,53 @@ +# NeMo-RL v2 SingleController overlay for one-node, non-colocated Omni VSTAT. +defaults: ./vlm_grpo-nemotron-omni-30ba3b-16n8g-megatron-tp4ep4-async-gym-video.v1.yaml + +grpo: + async_grpo: null + val_period: 0 + val_at_start: false + val_at_end: false + deduplicate_multimodal_data: false + overlong_filtering: false + +data_plane: + enabled: true + impl: transfer_queue + backend: simple + claim_meta_poll_interval_s: 0.5 + simple: + num_storage_units: 2 + +async_rl: + sampler: + name: windowed + max_staleness_versions: 2 + recompute_kv_cache_after_weight_updates: false + min_groups_for_streaming_train: ${grpo.num_prompts_per_step} + max_inflight_prompts: ${grpo.num_prompts_per_step} + max_buffered_rollouts: 32 + diagnostics: true + +policy: + generation: + backend: megatron + colocated: + enabled: false + resources: + num_nodes: 1 + gpus_per_node: 2 + mcore_generation_config: + transformer_impl: inference_optimized + moe_router_dtype: fp32 + moe_pad_experts_for_cuda_graph_inference: false + cuda_graph_impl: local + inference_cuda_graph_scope: block + num_cuda_graphs: -1 + use_cuda_graphs_for_non_decode_steps: false + enable_chunked_prefill: true + async_sched_mode: async + kv_cache_management_mode: persist + refit_backend: nccl + +cluster: + num_nodes: 1 + gpus_per_node: 4 diff --git a/examples/run_grpo_single_controller.py b/examples/run_grpo_single_controller.py index 3a151016a2a..8456f51e7d1 100644 --- a/examples/run_grpo_single_controller.py +++ b/examples/run_grpo_single_controller.py @@ -126,7 +126,12 @@ def main() -> None: maybe_configure_data_plane_env(config.data_plane) init_ray() - tokenizer = get_tokenizer(config.policy["tokenizer"]) + processor = None + if config.policy.get("is_vlm", False): + processor = get_tokenizer(config.policy["tokenizer"], get_processor=True) + tokenizer = processor.tokenizer + else: + tokenizer = get_tokenizer(config.policy["tokenizer"]) assert config.policy["generation"] is not None, ( "A generation config is required for SC-driven async GRPO" ) @@ -144,7 +149,9 @@ def main() -> None: if bool(config.env.get("should_use_nemo_gym")): setup_nemo_gym_config(config, tokenizer) - actor_args, setup_timing_metrics = setup_single_controller(config, tokenizer) + actor_args, setup_timing_metrics = setup_single_controller( + config, tokenizer, processor=processor + ) print("🚀 Launching SingleControllerActor") sc = SingleControllerActor.remote( diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py index d753bc602c9..a5e56d2f79e 100644 --- a/nemo_rl/algorithms/single_controller_utils/setup.py +++ b/nemo_rl/algorithms/single_controller_utils/setup.py @@ -70,6 +70,10 @@ ) from nemo_rl.algorithms.utils import set_seed from nemo_rl.data.collate_fn import rl_collate_fn +from nemo_rl.data.multimodal_utils import ( + PACKED_MULTIMODAL_FIELDS, + PER_TOKEN_MULTIMODAL_FIELDS, +) from nemo_rl.data.utils import load_dataloader_state, setup_response_data from nemo_rl.data_plane import ( DATA_PLANE_CHECKPOINT_SCHEMA_VERSION, @@ -1074,6 +1078,8 @@ def setup_single_controller( # ========================== # TODO: add validate dataset wiring. use_nemo_gym = should_use_nemo_gym(master_config) + data_tokenizer = processor if processor is not None else tokenizer + is_vlm = processor is not None if use_nemo_gym and generation_config["backend"] not in ("vllm", "megatron"): raise NotImplementedError( "SC NeMo-Gym integration currently supports the vllm and megatron backends only; got " @@ -1084,13 +1090,18 @@ def setup_single_controller( if use_nemo_gym: # NeMo-Gym creates the env actor outside setup_response_data; we wire # it in after generation is up (it needs the OpenAI server URLs). - response_data = setup_response_data(tokenizer, data_config, env_configs=None) + response_data = setup_response_data( + data_tokenizer, data_config, env_configs=None, is_vlm=is_vlm + ) assert len(response_data) == 2 dataset, _val_dataset = response_data env_handles: dict[str, EnvironmentInterface] = {} else: response_data = setup_response_data( - tokenizer, data_config, env_configs=master_config.env + data_tokenizer, + data_config, + env_configs=master_config.env, + is_vlm=is_vlm, ) assert len(response_data) == 4 dataset, _val_dataset, env_handles, _val_env_handles = response_data @@ -1154,6 +1165,7 @@ def setup_single_controller( megatron_reserved_url = None megatron_port_holder = None reserved_http_server_port = None + weight_synchronizer: Optional[WeightSynchronizer] = None if megatron_backend: generation_config["model_name"] = master_config.policy["model_name"] @@ -1313,7 +1325,6 @@ def _build_generation_then_trainer( build_tasks["trainer"] = _build_trainer_and_value # Submit build tasks and get results - weight_synchronizer: Optional[WeightSynchronizer] = None try: with ThreadPoolExecutor(max_workers=len(build_tasks)) as executor: submitted = {k: executor.submit(fn) for k, fn in build_tasks.items()} @@ -1341,6 +1352,7 @@ def _build_generation_then_trainer( inference_cluster=inference_cluster, refit_buffer_size_gb=policy_config.get("refit_buffer_size_gb"), ) + generation.weight_synchronizer = weight_synchronizer weight_synchronizer.init_communicator() setup_timing_metrics.collective_init_time_s = time.perf_counter() - t0 t0 = time.perf_counter() @@ -1427,12 +1439,21 @@ def _build_generation_then_trainer( # SingleController reuses one partition for the run. Warm every known # tensor field before rollout, policy, and teacher writers become # concurrent; TransferQueue otherwise registers field names lazily. + partition_fields = fields_with_optional_routed_experts( + SC_ROLLOUT_SCHEMA_FIELDS, + enabled=router_replay_enabled(policy_config), + ) + if processor is not None: + partition_fields.extend( + field + for field in sorted( + PACKED_MULTIMODAL_FIELDS | PER_TOKEN_MULTIMODAL_FIELDS + ) + if field not in partition_fields + ) dp_client.register_partition( partition_id=partition_id, - fields=fields_with_optional_routed_experts( - SC_ROLLOUT_SCHEMA_FIELDS, - enabled=router_replay_enabled(policy_config), - ), + fields=partition_fields, num_samples=( master_config.async_rl.max_buffered_rollouts * algo_cfg.num_generations_per_prompt @@ -1509,6 +1530,7 @@ def _build_generation_then_trainer( refit_buffer_size_gb=policy_config.get("refit_buffer_size_gb"), refit_timeout_s=master_config.async_rl.generation_fleet_health.refit_timeout_s, ) + generation.weight_synchronizer = weight_synchronizer weight_synchronizer.init_communicator() setup_timing_metrics.collective_init_time_s = time.perf_counter() - t0 diff --git a/nemo_rl/data_plane/interfaces.py b/nemo_rl/data_plane/interfaces.py index ce8084c4d3e..4c33a6d0856 100644 --- a/nemo_rl/data_plane/interfaces.py +++ b/nemo_rl/data_plane/interfaces.py @@ -312,7 +312,7 @@ def slice(self, start: int, stop: int) -> "KVBatchMeta": ) def concat(self, *others: "KVBatchMeta") -> "KVBatchMeta": - """Append ``others`` to ``self``. All metas must share ``partition_id``.""" + """Append ``others`` and union their fields in first-seen order.""" if any(o.partition_id != self.partition_id for o in others): raise ValueError("KVBatchMeta.concat: partition_ids must match") all_m = (self, *others) @@ -325,9 +325,14 @@ def concat(self, *others: "KVBatchMeta") -> "KVBatchMeta": ) all_have_tags = all(m.tags is not None for m in all_m) tags = [t for m in all_m for t in (m.tags or [])] if all_have_tags else None - return self._replace( + merged_fields = list( + dict.fromkeys(field for meta in all_m for field in (meta.fields or [])) + ) + result = self._replace( sample_ids=sample_ids, sequence_lengths=seq_lens, tags=tags ) + result.fields = merged_fields or None + return result def drop(self, indices: "Sequence[int]") -> "KVBatchMeta | None": """Complement of :meth:`subset`. Returns ``None`` when all rows are dropped.""" diff --git a/nemo_rl/experience/payload.py b/nemo_rl/experience/payload.py index c981bc08c55..6ed9b408ee9 100644 --- a/nemo_rl/experience/payload.py +++ b/nemo_rl/experience/payload.py @@ -176,6 +176,7 @@ def record_to_train_batch( if include_message_violation_fields: train_data[INVALID_TOOL_CALL_MASK] = flat[INVALID_TOOL_CALL_MASK] train_data[MALFORMED_THINKING_MASK] = flat[MALFORMED_THINKING_MASK] + train_data.update(flat.get_multimodal_dict(as_tensors=False)) return BatchedDataDict[Any](train_data) diff --git a/nemo_rl/experience/rollout_manager.py b/nemo_rl/experience/rollout_manager.py index b2e16b7cd06..501c85f3a32 100644 --- a/nemo_rl/experience/rollout_manager.py +++ b/nemo_rl/experience/rollout_manager.py @@ -31,6 +31,7 @@ TQReplayBuffer, ) from nemo_rl.data.interfaces import DatumSpec, LLMMessageLogType +from nemo_rl.data.llm_message_utils import batched_message_log_to_flat_message from nemo_rl.data_plane.schema import MASK_SAMPLE from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.environments.interfaces import EnvironmentInterface @@ -633,9 +634,14 @@ async def _generate_response( Returns: Tuple of (assistant_message, input_lengths, gen_metrics) """ - # Prepare generation input - input_ids = torch.cat([m["token_ids"] for m in message_log]).unsqueeze(0) - input_lengths = torch.tensor([input_ids.shape[1]], dtype=torch.int32) + # Flatten both tokens and model-ready multimodal inputs. Building this + # from token_ids alone leaves expanded media placeholders in the prompt + # without the pixel tensors Megatron needs to project. + flat_messages, input_lengths = batched_message_log_to_flat_message( + [message_log], + pad_value_dict={"token_ids": self._tokenizer.pad_token_id}, + ) + input_ids = flat_messages["token_ids"] generation_input_data = BatchedDataDict[GenerationDatumSpec]( { "input_ids": input_ids, @@ -643,6 +649,9 @@ async def _generate_response( "stop_strings": [stop_strings], } ) + generation_input_data.update( + flat_messages.get_multimodal_dict(as_tensors=False) + ) # Generate response # TODO: update generate_async to return a single item directly diff --git a/tests/unit/data_plane/test_kvbatchmeta.py b/tests/unit/data_plane/test_kvbatchmeta.py index a8dc3bc822d..1fbabd6faa9 100644 --- a/tests/unit/data_plane/test_kvbatchmeta.py +++ b/tests/unit/data_plane/test_kvbatchmeta.py @@ -247,6 +247,34 @@ def test_tags_none_when_either_side_missing_in_concat(): assert with_tags.concat(without).tags is None +def test_concat_unions_payload_fields_in_first_seen_order(): + text = KVBatchMeta( + partition_id="p", + task_name="train", + sample_ids=["a"], + fields=["input_ids", "input_lengths"], + ) + multimodal = KVBatchMeta( + partition_id="p", + task_name="train", + sample_ids=["b"], + fields=[ + "input_ids", + "pixel_values", + "__nrl_packed_tensor_meta__pixel_values", + ], + ) + + joined = text.concat(multimodal) + + assert joined.fields == [ + "input_ids", + "input_lengths", + "pixel_values", + "__nrl_packed_tensor_meta__pixel_values", + ] + + # ── Realistic tags from the rollout-shapes helper ── diff --git a/tests/unit/experience/test_rollout_manager.py b/tests/unit/experience/test_rollout_manager.py index fe8726250e4..b1cc003d0c9 100644 --- a/tests/unit/experience/test_rollout_manager.py +++ b/tests/unit/experience/test_rollout_manager.py @@ -29,6 +29,7 @@ import tempfile import uuid from copy import deepcopy +from types import SimpleNamespace import pytest import torch @@ -40,6 +41,7 @@ from nemo_rl.data.collate_fn import rl_collate_fn from nemo_rl.data.datasets.response_datasets import NemoGymDataset from nemo_rl.data.interfaces import DatumSpec +from nemo_rl.data.multimodal_utils import PackedTensor from nemo_rl.data.processors import nemo_gym_data_processor from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.experience.interfaces import ( @@ -92,6 +94,55 @@ async def apply(): return _run(apply()) +def test_generate_response_forwards_message_log_media_to_generation() -> None: + captured: dict[str, BatchedDataDict] = {} + + class _Generation: + async def generate_async(self, data): + captured["data"] = data + input_len = int(data["input_lengths"][0]) + yield 0, BatchedDataDict( + { + "output_ids": torch.cat( + (data["input_ids"], torch.tensor([[42]])), dim=1 + ), + "unpadded_sequence_lengths": torch.tensor([input_len + 1]), + "logprobs": torch.zeros(1, input_len + 1), + } + ) + + manager = object.__new__(RolloutManager) + manager._policy_generation = _Generation() + manager._tokenizer = SimpleNamespace( + pad_token_id=0, + decode=lambda *_args, **_kwargs: "answer", + ) + manager._timeouts = SimpleNamespace(generation_s=10.0) + pixel_values = PackedTensor(torch.ones(2, 3, 4, 4), dim_to_pack=0) + imgs_sizes = PackedTensor(torch.tensor([[4, 4], [4, 4]]), dim_to_pack=0) + message_log = [ + { + "role": "user", + "content": "image", + "token_ids": torch.tensor([1, 2, 3]), + "pixel_values": pixel_values, + "imgs_sizes": imgs_sizes, + } + ] + + _run(manager._generate_response(message_log, None)) + + generation_data = captured["data"] + assert isinstance(generation_data["pixel_values"], PackedTensor) + assert isinstance(generation_data["imgs_sizes"], PackedTensor) + assert torch.equal( + generation_data["pixel_values"].as_tensor(), pixel_values.as_tensor() + ) + assert torch.equal( + generation_data["imgs_sizes"].as_tensor(), imgs_sizes.as_tensor() + ) + + class _FakeBuffer: """Minimal TQReplayBuffer stand-in that records reserve/commit calls.""" diff --git a/tests/unit/single_controller/test_entrypoint.py b/tests/unit/single_controller/test_entrypoint.py index c4436936db8..c3b931e037f 100644 --- a/tests/unit/single_controller/test_entrypoint.py +++ b/tests/unit/single_controller/test_entrypoint.py @@ -89,7 +89,7 @@ def main_context(monkeypatch: pytest.MonkeyPatch) -> SimpleNamespace: monkeypatch.setattr( run_grpo_single_controller, "setup_single_controller", - lambda *_args: (actor_args, SetupTimingMetrics()), + lambda *_args, **_kwargs: (actor_args, SetupTimingMetrics()), ) monkeypatch.setattr( run_grpo_single_controller.SingleControllerActor, @@ -170,3 +170,30 @@ def test_main_configures_generation_for_trained_mtp( assert ( main_context.config.policy["generation"] is main_context.configured_generation ) + + +def test_main_passes_processor_for_vlm( + main_context: SimpleNamespace, + monkeypatch: pytest.MonkeyPatch, +) -> None: + processor = SimpleNamespace(tokenizer="vlm-tokenizer") + get_tokenizer = MagicMock(return_value=processor) + setup_single_controller = MagicMock( + return_value=(main_context.actor_args, SetupTimingMetrics()) + ) + main_context.config.policy["is_vlm"] = True + monkeypatch.setattr(run_grpo_single_controller, "get_tokenizer", get_tokenizer) + monkeypatch.setattr( + run_grpo_single_controller, + "setup_single_controller", + setup_single_controller, + ) + + run_grpo_single_controller.main() + + get_tokenizer.assert_called_once_with( + main_context.config.policy["tokenizer"], get_processor=True + ) + setup_single_controller.assert_called_once_with( + main_context.config, "vlm-tokenizer", processor=processor + ) diff --git a/tests/unit/single_controller/test_setup.py b/tests/unit/single_controller/test_setup.py index 9b1ad8e5cc8..ba0b075eb8e 100644 --- a/tests/unit/single_controller/test_setup.py +++ b/tests/unit/single_controller/test_setup.py @@ -1052,6 +1052,26 @@ def test_env_handles_sourced_from_setup_response_data(self, patched_factories): assert call_kwargs["env_configs"] == {"math": math_env_cfg} assert actor_args.env_handles is patched_factories["env_handles"] + def test_vlm_processor_used_for_data_and_environment_setup( + self, patched_factories + ): + mc = _make_master_config(env={"clevr-cogent": {"some": "value"}}) + tokenizer = MagicMock(pad_token_id=0) + processor = MagicMock(tokenizer=tokenizer) + processor.model_input_names = ["input_ids", "pixel_values", "image_grid_thw"] + + actor_args, _ = setup_single_controller(mc, tokenizer, processor=processor) + + call_args, call_kwargs = patched_factories["setup_response_data"].call_args + assert call_args[0] is processor + assert call_kwargs["env_configs"] == { + "clevr-cogent": {"some": "value"} + } + assert call_kwargs["is_vlm"] is True + warmup_fields = actor_args.dp_client.register_partition.call_args.kwargs["fields"] + assert "pixel_values" in warmup_fields + assert "__nrl_packed_tensor_meta__pixel_values" in warmup_fields + def test_weight_sync_factory_args(self, patched_factories): """create_weight_synchronizer receives policy / generation / topology.""" mc = _make_master_config(colocated=False, backend="vllm") @@ -1533,13 +1553,15 @@ def _spinup_gym(**_): mock_megatron.return_value.finish_generation.assert_called_once_with() if gym: # Gym spins up on the reserved URL, before the served-address - # cross-check — so the mismatch leg sees it too. + # cross-check — so the mismatch leg sees it too. The initial refit + # must happen during that wait because it starts Megatron's server. _, spinup_kwargs = mock_spinup.call_args assert spinup_kwargs["base_urls"] == [reserved_url] # The initial refit ran in setup, against the collective brought up # there; the served-address check reads the URLs it populated. weight_sync.init_communicator.assert_called_once_with() weight_sync.sync_weights.assert_called_once_with() + assert mock_megatron.return_value.weight_synchronizer is weight_sync else: mock_spinup.assert_not_called() # Native: the actor's startup sync performs the initial refit. From 7bce7df44c36427d709cf14d5f3e099106f78aa7 Mon Sep 17 00:00:00 2001 From: Cory Ye Date: Sat, 29 Aug 2026 13:23:29 -0700 Subject: [PATCH 05/19] WAR fix for multimodal request checkpoint embedding recompute. Signed-off-by: Cory Ye --- ...g-megatron-single-controller-async.v1.yaml | 8 +-- ...g-megatron-single-controller-async.v1.yaml | 15 +++-- ...g-megatron-single-controller-async.v1.yaml | 8 +-- ...g-megatron-single-controller-async.v1.yaml | 54 ++++++++++++++++++ .../single_controller_utils/setup.py | 1 + .../generation/megatron/megatron_worker.py | 57 +++++++++++++++++-- .../models/policy/test_megatron_worker.py | 25 ++++++++ 7 files changed, 150 insertions(+), 18 deletions(-) create mode 100644 examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-8n4g-megatron-single-controller-async.v1.yaml diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron-single-controller-async.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron-single-controller-async.v1.yaml index 2705ab46b5b..ae9150e691d 100644 --- a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron-single-controller-async.v1.yaml +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron-single-controller-async.v1.yaml @@ -17,12 +17,12 @@ data_plane: async_rl: sampler: - name: windowed - max_staleness_versions: 2 + name: in_order + max_lookahead_versions: 1 recompute_kv_cache_after_weight_updates: false min_groups_for_streaming_train: ${grpo.num_prompts_per_step} - max_inflight_prompts: ${grpo.num_prompts_per_step} - max_buffered_rollouts: 32 + max_inflight_prompts: ${mul:${grpo.num_prompts_per_step}, 2} + max_buffered_rollouts: ${mul:${grpo.num_prompts_per_step}, 2} diagnostics: true policy: diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-single-controller-async.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-single-controller-async.v1.yaml index 22bc06daf50..be31246f1b4 100644 --- a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-single-controller-async.v1.yaml +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-single-controller-async.v1.yaml @@ -9,16 +9,21 @@ grpo: data_plane: enabled: true + impl: transfer_queue + backend: simple + claim_meta_poll_interval_s: 0.5 + simple: + num_storage_units: 16 async_rl: sampler: - name: windowed - max_staleness_versions: 2 + name: in_order + max_lookahead_versions: 1 recompute_kv_cache_after_weight_updates: false min_groups_for_streaming_train: ${grpo.num_prompts_per_step} - max_inflight_prompts: ${grpo.num_prompts_per_step} - max_buffered_rollouts: 32 - diagnostics: true + max_inflight_prompts: ${mul:${grpo.num_prompts_per_step}, 2} + max_buffered_rollouts: ${mul:${grpo.num_prompts_per_step}, 2} + diagnostics: false policy: generation: diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-1n4g-megatron-single-controller-async.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-1n4g-megatron-single-controller-async.v1.yaml index f3b8e8aee68..27b83e431d9 100644 --- a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-1n4g-megatron-single-controller-async.v1.yaml +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-1n4g-megatron-single-controller-async.v1.yaml @@ -19,12 +19,12 @@ data_plane: async_rl: sampler: - name: windowed - max_staleness_versions: 2 + name: in_order + max_lookahead_versions: 1 recompute_kv_cache_after_weight_updates: false min_groups_for_streaming_train: ${grpo.num_prompts_per_step} - max_inflight_prompts: ${grpo.num_prompts_per_step} - max_buffered_rollouts: 32 + max_inflight_prompts: ${mul:${grpo.num_prompts_per_step}, 2} + max_buffered_rollouts: ${mul:${grpo.num_prompts_per_step}, 2} diagnostics: true policy: diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-8n4g-megatron-single-controller-async.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-8n4g-megatron-single-controller-async.v1.yaml new file mode 100644 index 00000000000..6214e10aa2f --- /dev/null +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-8n4g-megatron-single-controller-async.v1.yaml @@ -0,0 +1,54 @@ +# NeMo-RL v2 SingleController overlay for eight-node, non-colocated Omni VSTAT. +defaults: ./vlm_grpo-nemotron-omni-30ba3b-16n8g-megatron-tp4ep4-async-gym-video.v1.yaml + +grpo: + async_grpo: null + val_period: 0 + val_at_start: false + val_at_end: false + deduplicate_multimodal_data: false + overlong_filtering: false + +data_plane: + enabled: true + impl: transfer_queue + backend: simple + claim_meta_poll_interval_s: 0.5 + simple: + num_storage_units: 16 + +async_rl: + sampler: + name: in_order + max_lookahead_versions: 1 + recompute_kv_cache_after_weight_updates: false + min_groups_for_streaming_train: ${grpo.num_prompts_per_step} + max_inflight_prompts: ${mul:${grpo.num_prompts_per_step}, 2} + max_buffered_rollouts: ${mul:${grpo.num_prompts_per_step}, 2} + diagnostics: false + +policy: + generation: + backend: megatron + colocated: + enabled: false + resources: + num_nodes: 6 + gpus_per_node: 4 + mcore_generation_config: + transformer_impl: inference_optimized + moe_router_dtype: fp32 + moe_pad_experts_for_cuda_graph_inference: false + cuda_graph_impl: local + inference_cuda_graph_scope: block + num_cuda_graphs: -1 + use_cuda_graphs_for_non_decode_steps: false + enable_chunked_prefill: true + enable_prefix_caching: false + async_sched_mode: async + kv_cache_management_mode: persist + refit_backend: nccl + +cluster: + num_nodes: 8 + gpus_per_node: 4 diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py index a5e56d2f79e..1f5c9c3ae83 100644 --- a/nemo_rl/algorithms/single_controller_utils/setup.py +++ b/nemo_rl/algorithms/single_controller_utils/setup.py @@ -1351,6 +1351,7 @@ def _build_generation_then_trainer( train_cluster=train_cluster, inference_cluster=inference_cluster, refit_buffer_size_gb=policy_config.get("refit_buffer_size_gb"), + refit_timeout_s=master_config.async_rl.generation_fleet_health.refit_timeout_s, ) generation.weight_synchronizer = weight_synchronizer weight_synchronizer.init_communicator() diff --git a/nemo_rl/models/generation/megatron/megatron_worker.py b/nemo_rl/models/generation/megatron/megatron_worker.py index 882544e4347..aae8221b36c 100644 --- a/nemo_rl/models/generation/megatron/megatron_worker.py +++ b/nemo_rl/models/generation/megatron/megatron_worker.py @@ -504,15 +504,55 @@ async def _sleep_engine(self): self.inference_client.suspend_engines() await self.dynamic_inference_engine.wait_until(EngineState.SUSPENDED) + def _move_retained_vlm_media(self, device: torch.device) -> None: + """Move raw media retained by active Megatron requests to ``device``.""" + engine = self.dynamic_inference_engine + if engine is None: + return + + moved_tensors: dict[int, torch.Tensor] = {} + for entry in getattr(engine, "requests", {}).values(): + record = getattr(entry, "record", None) + if not record: + continue + request = record[-1] + imgs = getattr(request, "imgs", None) + if not torch.is_tensor(imgs) or imgs.device == device: + continue + + # Multiple sampled completions can retain the same source image. + # Preserve that sharing rather than allocating one copy per request. + moved = moved_tensors.get(id(imgs)) + if moved is None: + moved = imgs.to(device=device) + moved_tensors[id(imgs)] = moved + request.imgs = moved + def _wake(self) -> None: """Resume + unpause the engine. No-op if already awake.""" if not self._inference_engine_asleep: return - future = asyncio.run_coroutine_threadsafe( - self._wake_engine(), self._inference_loop - ) - future.result() - torch.distributed.barrier() + + # Megatron invalidates trainable vision projections during refit and + # recomputes them while resuming. Raw media retained by an active request + # may live on CPU, so make it device-local for that refresh. Move it back + # after the engine reaches RUNNING to avoid pinning image batches in HBM. + # HACK(@cspades): Revert this H2D-D2H code after this is merged: + # https://github.com/NVIDIA/Megatron-LM/pull/6976 + cuda_device = torch.device("cuda", torch.cuda.current_device()) + self._move_retained_vlm_media(cuda_device) + try: + future = asyncio.run_coroutine_threadsafe( + self._wake_engine(), self._inference_loop + ) + future.result() + torch.distributed.barrier() + except BaseException: + # Best-effort cleanup if resume fails before _wake_engine reaches + # its normal post-refresh offload point. + self._move_retained_vlm_media(torch.device("cpu")) + raise + self._inference_engine_asleep = False print(f"[Rank {self.rank}] resumed inference engine") @@ -521,6 +561,13 @@ async def _wake_engine(self): self.inference_client.resume_engines() await self.dynamic_inference_engine.wait_until(EngineState.RESUMED) + # DynamicInferenceEngine.resume() refreshes stale VLM embeddings before + # publishing RESUMED. The engine is still paused here, so raw media can + # be returned to CPU before any request is allowed to execute. + # HACK(@cspades): Revert this H2D-D2H code after this is merged: + # https://github.com/NVIDIA/Megatron-LM/pull/6976 + self._move_retained_vlm_media(torch.device("cpu")) + if torch.distributed.get_rank() == 0: self.inference_client.unpause_engines() await self.dynamic_inference_engine.wait_until(EngineState.RUNNING) diff --git a/tests/unit/models/policy/test_megatron_worker.py b/tests/unit/models/policy/test_megatron_worker.py index 78db03db202..29eaa03c9a4 100644 --- a/tests/unit/models/policy/test_megatron_worker.py +++ b/tests/unit/models/policy/test_megatron_worker.py @@ -822,6 +822,31 @@ def test_prepare_for_generation_disables_param_gather_hook_before_wake( assert model.config.flash_decode is False +def test_move_retained_vlm_media_preserves_shared_tensors() -> None: + from nemo_rl.models.generation.megatron.megatron_worker import ( + MegatronGenerationMixin, + ) + + shared_imgs = torch.ones(2) + first_request = SimpleNamespace(imgs=shared_imgs) + second_request = SimpleNamespace(imgs=shared_imgs) + text_request = SimpleNamespace() + worker = MegatronGenerationMixin() + worker.dynamic_inference_engine = SimpleNamespace( + requests={ + 1: SimpleNamespace(record=[first_request]), + 2: SimpleNamespace(record=[second_request]), + 3: SimpleNamespace(record=[text_request]), + } + ) + + worker._move_retained_vlm_media(torch.device("meta")) + + assert first_request.imgs.device.type == "meta" + assert second_request.imgs is first_request.imgs + assert not hasattr(text_request, "imgs") + + def create_megatron_test_config( model_name: str, tp: int = 1, From 2ee0d377fd6ebdd0c17e6db55aba870f8d6636a1 Mon Sep 17 00:00:00 2001 From: Cory Ye Date: Mon, 31 Aug 2026 20:11:42 -0700 Subject: [PATCH 06/19] Address review feedback and incorporate code from #2957. Signed-off-by: Cory Ye --- ...g-megatron-single-controller-async.v1.yaml | 14 +++- ...g-megatron-single-controller-async.v1.yaml | 6 ++ nemo_rl/data/multimodal_utils.py | 17 +++-- nemo_rl/experience/payload.py | 13 ++++ .../generation/megatron/megatron_worker.py | 57 ++--------------- nemo_rl/models/value/tq_value.py | 4 ++ tests/test_suites/disabled.txt | 6 ++ tests/test_suites/nightly.txt | 3 + ...n4g-megatron-single-controller-async.v1.sh | 64 +++++++++++++++++++ ...n4g-megatron-single-controller-async.v1.sh | 46 +++++++++++++ ...n4g-megatron-single-controller-async.v1.sh | 48 ++++++++++++++ ...n4g-megatron-single-controller-async.v1.sh | 48 ++++++++++++++ tests/unit/data_plane/test_kvbatchmeta.py | 4 +- tests/unit/experience/test_payload.py | 35 ++++++++++ tests/unit/experience/test_rollout_manager.py | 3 +- tests/unit/single_controller/test_setup.py | 10 ++- 16 files changed, 314 insertions(+), 64 deletions(-) create mode 100755 tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron-single-controller-async.v1.sh create mode 100755 tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-single-controller-async.v1.sh create mode 100755 tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-1n4g-megatron-single-controller-async.v1.sh create mode 100755 tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-8n4g-megatron-single-controller-async.v1.sh diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron-single-controller-async.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron-single-controller-async.v1.yaml index ae9150e691d..0840a470d9b 100644 --- a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron-single-controller-async.v1.yaml +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron-single-controller-async.v1.yaml @@ -1,7 +1,9 @@ # NeMo-RL v2 SingleController overlay for one-node, non-colocated Omni CLEVR. -defaults: ./vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.yaml +defaults: ./vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron_generation.v1.yaml grpo: + num_prompts_per_step: 2 + num_generations_per_prompt: 8 async_grpo: null val_period: 0 val_at_start: false @@ -26,8 +28,13 @@ async_rl: diagnostics: true policy: + train_global_batch_size: ${mul:${grpo.num_prompts_per_step}, ${grpo.num_generations_per_prompt}} + megatron_cfg: + tensor_model_parallel_size: 2 + expert_model_parallel_size: 2 generation: backend: megatron + max_new_tokens: 512 colocated: enabled: false resources: @@ -35,6 +42,8 @@ policy: gpus_per_node: 2 mcore_generation_config: transformer_impl: inference_optimized + tensor_model_parallel_size: 2 + expert_model_parallel_size: 2 moe_router_dtype: fp32 moe_pad_experts_for_cuda_graph_inference: false cuda_graph_impl: local @@ -46,6 +55,9 @@ policy: kv_cache_management_mode: persist refit_backend: nccl +data: + num_workers: 0 + cluster: num_nodes: 1 gpus_per_node: 4 diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-1n4g-megatron-single-controller-async.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-1n4g-megatron-single-controller-async.v1.yaml index 27b83e431d9..4eca9a9e08f 100644 --- a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-1n4g-megatron-single-controller-async.v1.yaml +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-1n4g-megatron-single-controller-async.v1.yaml @@ -28,6 +28,9 @@ async_rl: diagnostics: true policy: + megatron_cfg: + tensor_model_parallel_size: 2 + expert_model_parallel_size: 2 generation: backend: megatron colocated: @@ -37,6 +40,8 @@ policy: gpus_per_node: 2 mcore_generation_config: transformer_impl: inference_optimized + tensor_model_parallel_size: 2 + expert_model_parallel_size: 2 moe_router_dtype: fp32 moe_pad_experts_for_cuda_graph_inference: false cuda_graph_impl: local @@ -44,6 +49,7 @@ policy: num_cuda_graphs: -1 use_cuda_graphs_for_non_decode_steps: false enable_chunked_prefill: true + enable_prefix_caching: false async_sched_mode: async kv_cache_management_mode: persist refit_backend: nccl diff --git a/nemo_rl/data/multimodal_utils.py b/nemo_rl/data/multimodal_utils.py index e9bebd4e447..c2b79d73c63 100644 --- a/nemo_rl/data/multimodal_utils.py +++ b/nemo_rl/data/multimodal_utils.py @@ -1093,6 +1093,16 @@ def encode_multimodal_for_wire( ) +# Model inputs some remote-code processors omit from ``model_input_names`` even +# though their forward requires them. Keep extraction and TQ schema warmup aligned. +UNDECLARED_MULTIMODAL_MODEL_INPUTS = ( + "imgs_sizes", + "num_frames", + "pixel_values_flat", + "image_num_patches", +) + + def get_multimodal_keys_from_processor(processor) -> list[str]: """Get keys of the multimodal data that can be used as model inputs. @@ -1216,12 +1226,7 @@ def extract_multimodal_model_inputs( # TODO(rohitrango): Let ProcessorInterface declare model-specific media inputs. # Some remote-code processors omit these inputs from model_input_names even # though their model forward requires them. - for key in ( - "imgs_sizes", - "num_frames", - "pixel_values_flat", - "image_num_patches", - ): + for key in UNDECLARED_MULTIMODAL_MODEL_INPUTS: if key in processed and key not in multimodal_keys: multimodal_keys.append(key) for key in multimodal_keys: diff --git a/nemo_rl/experience/payload.py b/nemo_rl/experience/payload.py index 6ed9b408ee9..e1f4de8431b 100644 --- a/nemo_rl/experience/payload.py +++ b/nemo_rl/experience/payload.py @@ -22,6 +22,10 @@ from tensordict import TensorDict from nemo_rl.data.interfaces import LLMMessageLogType, VLMMessageLogType +from nemo_rl.data.multimodal_utils import ( + encode_multimodal_for_wire, + multimodal_row_tags, +) from nemo_rl.data_plane.codec import pack_jagged_fields from nemo_rl.data_plane.column_io import TOKEN_ALIGNED_FIELDS from nemo_rl.data_plane.schema import ( @@ -207,16 +211,25 @@ def pack_payload( if isinstance(v, torch.Tensor) or (isinstance(v, np.ndarray) and v.dtype == object) } + multimodal = BatchedDataDict[Any](train_batch).get_multimodal_dict( + as_tensors=False + ) + for key, value in multimodal.items(): + wire_value = encode_multimodal_for_wire(key, value) + if wire_value is not None: + tensor_fields[key] = wire_value fields_td = pack_jagged_fields( tensor_fields, lengths=lengths, token_aligned_fields=TOKEN_ALIGNED_FIELDS ) sample_ids = [f"{group_id}_g{i}" for i in range(n)] violations = train_batch.get(_VIOLATION_COUNTS_KEY, [{}] * n) + multimodal_tags = multimodal_row_tags(multimodal, n) or [{} for _ in range(n)] tags = [ { "weight_version": weight_version, "prompt_idx": prompt_idx, **violations[i], + **multimodal_tags[i], } for i in range(n) ] diff --git a/nemo_rl/models/generation/megatron/megatron_worker.py b/nemo_rl/models/generation/megatron/megatron_worker.py index aae8221b36c..882544e4347 100644 --- a/nemo_rl/models/generation/megatron/megatron_worker.py +++ b/nemo_rl/models/generation/megatron/megatron_worker.py @@ -504,55 +504,15 @@ async def _sleep_engine(self): self.inference_client.suspend_engines() await self.dynamic_inference_engine.wait_until(EngineState.SUSPENDED) - def _move_retained_vlm_media(self, device: torch.device) -> None: - """Move raw media retained by active Megatron requests to ``device``.""" - engine = self.dynamic_inference_engine - if engine is None: - return - - moved_tensors: dict[int, torch.Tensor] = {} - for entry in getattr(engine, "requests", {}).values(): - record = getattr(entry, "record", None) - if not record: - continue - request = record[-1] - imgs = getattr(request, "imgs", None) - if not torch.is_tensor(imgs) or imgs.device == device: - continue - - # Multiple sampled completions can retain the same source image. - # Preserve that sharing rather than allocating one copy per request. - moved = moved_tensors.get(id(imgs)) - if moved is None: - moved = imgs.to(device=device) - moved_tensors[id(imgs)] = moved - request.imgs = moved - def _wake(self) -> None: """Resume + unpause the engine. No-op if already awake.""" if not self._inference_engine_asleep: return - - # Megatron invalidates trainable vision projections during refit and - # recomputes them while resuming. Raw media retained by an active request - # may live on CPU, so make it device-local for that refresh. Move it back - # after the engine reaches RUNNING to avoid pinning image batches in HBM. - # HACK(@cspades): Revert this H2D-D2H code after this is merged: - # https://github.com/NVIDIA/Megatron-LM/pull/6976 - cuda_device = torch.device("cuda", torch.cuda.current_device()) - self._move_retained_vlm_media(cuda_device) - try: - future = asyncio.run_coroutine_threadsafe( - self._wake_engine(), self._inference_loop - ) - future.result() - torch.distributed.barrier() - except BaseException: - # Best-effort cleanup if resume fails before _wake_engine reaches - # its normal post-refresh offload point. - self._move_retained_vlm_media(torch.device("cpu")) - raise - + future = asyncio.run_coroutine_threadsafe( + self._wake_engine(), self._inference_loop + ) + future.result() + torch.distributed.barrier() self._inference_engine_asleep = False print(f"[Rank {self.rank}] resumed inference engine") @@ -561,13 +521,6 @@ async def _wake_engine(self): self.inference_client.resume_engines() await self.dynamic_inference_engine.wait_until(EngineState.RESUMED) - # DynamicInferenceEngine.resume() refreshes stale VLM embeddings before - # publishing RESUMED. The engine is still paused here, so raw media can - # be returned to CPU before any request is allowed to execute. - # HACK(@cspades): Revert this H2D-D2H code after this is merged: - # https://github.com/NVIDIA/Megatron-LM/pull/6976 - self._move_retained_vlm_media(torch.device("cpu")) - if torch.distributed.get_rank() == 0: self.inference_client.unpause_engines() await self.dynamic_inference_engine.wait_until(EngineState.RUNNING) diff --git a/nemo_rl/models/value/tq_value.py b/nemo_rl/models/value/tq_value.py index 9760248c822..6582960ec70 100644 --- a/nemo_rl/models/value/tq_value.py +++ b/nemo_rl/models/value/tq_value.py @@ -110,6 +110,8 @@ def get_values_from_meta( timer: Optional timer for nested get_values measurements. """ spa, dba = self._packing_args("logprob_mb_tokens") + # The critic is text-only (built with is_vlm=False), so media columns + # are deliberately not fetched here. value_meta = self._isolated_meta( meta, fields=list(VALUE_SEED_FIELDS), @@ -165,6 +167,8 @@ def train_from_meta( micro_batch_size = mbs or self.cfg["train_micro_batch_size"] spa, dba = self._packing_args("train_mb_tokens") + # The critic is text-only (built with is_vlm=False), so media columns + # are deliberately not fetched here. train_meta = self._isolated_meta( meta, fields=list(DP_VALUE_TRAIN_FIELDS), diff --git a/tests/test_suites/disabled.txt b/tests/test_suites/disabled.txt index fcbfd7e18e1..592348c897c 100644 --- a/tests/test_suites/disabled.txt +++ b/tests/test_suites/disabled.txt @@ -31,6 +31,12 @@ tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-click-2n8g-megatron-t # functional test before moving it to a recurring suite. tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-count-1n4g-megatron_generation.v1.sh +# TODO(@cspades): Multimodal SingleController/TransferQueue recipes awaiting prepared +# fixtures and validation. The 1n4g CLEVR sibling remains enabled as an L1 smoke test. +tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-single-controller-async.v1.sh +tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-1n4g-megatron-single-controller-async.v1.sh +tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-8n4g-megatron-single-controller-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/nightly.txt b/tests/test_suites/nightly.txt index 93171964fad..8910ac1c4f6 100644 --- a/tests/test_suites/nightly.txt +++ b/tests/test_suites/nightly.txt @@ -58,6 +58,9 @@ tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8.v tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-2n8g-megatron-tp8ep8.v1.sh tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-megatron-tp8ep16.v1.sh +# L1 multimodal SingleController/TransferQueue smoke coverage +tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron-single-controller-async.v1.sh + # Functional Qwen3.5-35B VLM GRPO run # The AutoModel variant is re-enabled with the vLLM 0.25.1 bump (no longer hits # https://github.com/vllm-project/vllm/issues/36237). The Megatron variant still diff --git a/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron-single-controller-async.v1.sh b/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron-single-controller-async.v1.sh new file mode 100755 index 00000000000..52e9c9ab9cc --- /dev/null +++ b/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron-single-controller-async.v1.sh @@ -0,0 +1,64 @@ +#!/bin/bash +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +source "$SCRIPT_DIR/common.env" + +# Two steps exercise CP=1 multimodal generation, TQ policy training, weight +# refit, and a post-refit rollout without treating this smoke test as a +# convergence run. CP>1 + multimodal + TQ remains untested. +# ===== BEGIN CONFIG ===== +NUM_NODES=1 +GPUS_PER_NODE=4 +STEPS_PER_RUN=2 +MAX_STEPS=2 +NUM_RUNS=1 +NUM_MINUTES=120 +# ===== END CONFIG ===== + +exit_if_max_steps_reached + +cd "$PROJECT_ROOT" + +uv run examples/run_grpo_single_controller.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" + +uv run tests/json_dump_tb_logs.py "$LOG_DIR" --output_path "$JSON_METRICS" + +RECORDED_STEP=$(jq -r \ + 'if has("train/loss") then (."train/loss" | keys | map(tonumber) | max // 0) else 0 end' \ + "$JSON_METRICS") +if [[ "$RECORDED_STEP" -lt 1 ]]; then + echo "[ERROR] Expected at least one completed training step" + exit 1 +fi + +uv run tests/check_metrics.py "$JSON_METRICS" \ + 'all_finite(data["train/token_mult_prob_error"])' \ + 'max(data["train/loss"]) < 1000000.0' \ + 'min(data["train/loss"]) > -1000000.0' + +rm -rf "$CKPT_DIR" diff --git a/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-single-controller-async.v1.sh b/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-single-controller-async.v1.sh new file mode 100755 index 00000000000..48ba6451ce1 --- /dev/null +++ b/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-single-controller-async.v1.sh @@ -0,0 +1,46 @@ +#!/bin/bash +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +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=2 +MAX_STEPS=2 +NUM_RUNS=1 +NUM_MINUTES=120 +# ===== END CONFIG ===== + +exit_if_max_steps_reached + +cd "$PROJECT_ROOT" + +uv run examples/run_grpo_single_controller.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" + +uv run tests/json_dump_tb_logs.py "$LOG_DIR" --output_path "$JSON_METRICS" diff --git a/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-1n4g-megatron-single-controller-async.v1.sh b/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-1n4g-megatron-single-controller-async.v1.sh new file mode 100755 index 00000000000..65703c42e74 --- /dev/null +++ b/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-1n4g-megatron-single-controller-async.v1.sh @@ -0,0 +1,48 @@ +#!/bin/bash +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +source "$SCRIPT_DIR/common.env" + +# TODO(@cspades): Prepare a self-contained VSTAT fixture and validate this +# driver before moving it from disabled.txt into a recurring suite. +# ===== BEGIN CONFIG ===== +NUM_NODES=1 +GPUS_PER_NODE=4 +STEPS_PER_RUN=2 +MAX_STEPS=2 +NUM_RUNS=1 +NUM_MINUTES=120 +# ===== END CONFIG ===== + +exit_if_max_steps_reached + +cd "$PROJECT_ROOT" + +uv run examples/run_grpo_single_controller.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" + +uv run tests/json_dump_tb_logs.py "$LOG_DIR" --output_path "$JSON_METRICS" diff --git a/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-8n4g-megatron-single-controller-async.v1.sh b/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-8n4g-megatron-single-controller-async.v1.sh new file mode 100755 index 00000000000..3f65e783dbc --- /dev/null +++ b/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-8n4g-megatron-single-controller-async.v1.sh @@ -0,0 +1,48 @@ +#!/bin/bash +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +source "$SCRIPT_DIR/common.env" + +# TODO(@cspades): Prepare a self-contained VSTAT fixture and validate this +# driver before moving it from disabled.txt into a recurring suite. +# ===== BEGIN CONFIG ===== +NUM_NODES=8 +GPUS_PER_NODE=4 +STEPS_PER_RUN=2 +MAX_STEPS=2 +NUM_RUNS=1 +NUM_MINUTES=120 +# ===== END CONFIG ===== + +exit_if_max_steps_reached + +cd "$PROJECT_ROOT" + +uv run examples/run_grpo_single_controller.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" + +uv run tests/json_dump_tb_logs.py "$LOG_DIR" --output_path "$JSON_METRICS" diff --git a/tests/unit/data_plane/test_kvbatchmeta.py b/tests/unit/data_plane/test_kvbatchmeta.py index 1fbabd6faa9..49272eaea53 100644 --- a/tests/unit/data_plane/test_kvbatchmeta.py +++ b/tests/unit/data_plane/test_kvbatchmeta.py @@ -261,7 +261,7 @@ def test_concat_unions_payload_fields_in_first_seen_order(): fields=[ "input_ids", "pixel_values", - "__nrl_packed_tensor_meta__pixel_values", + "image_grid_thw", ], ) @@ -271,7 +271,7 @@ def test_concat_unions_payload_fields_in_first_seen_order(): "input_ids", "input_lengths", "pixel_values", - "__nrl_packed_tensor_meta__pixel_values", + "image_grid_thw", ] diff --git a/tests/unit/experience/test_payload.py b/tests/unit/experience/test_payload.py index 49de9e044c9..04755999404 100644 --- a/tests/unit/experience/test_payload.py +++ b/tests/unit/experience/test_payload.py @@ -16,6 +16,8 @@ import torch +from nemo_rl.data.multimodal_utils import PackedTensor +from nemo_rl.data_plane.codec import materialize from nemo_rl.data_plane.schema import ( INVALID_TOOL_CALL_MASK, MALFORMED_THINKING_MASK, @@ -243,6 +245,39 @@ def test_record_to_train_batch_omits_routed_experts_when_absent() -> None: assert "routed_experts" not in fields +def test_multimodal_packed_tensor_round_trips_through_tq_payload() -> None: + completions = [ + _completion(route_start=10, reward=1.0, with_routes=False), + _completion(route_start=30, reward=2.0, with_routes=False), + ] + media = torch.arange(8, dtype=torch.float32).reshape(2, 4) + completions[0].message_log[0]["pixel_values"] = PackedTensor(media, dim_to_pack=0) + + train_batch = record_to_train_batch( + _record(completions), + pad_value_dict={"token_ids": 0, "input_ids": 0}, + include_message_violation_fields=False, + ) + assert isinstance(train_batch["pixel_values"], PackedTensor) + + _, fields, tags = pack_payload( + train_batch, + weight_version=3, + group_id="group", + prompt_idx=17, + ) + assert "pixel_values" in fields + assert tags[0]["pixel_values__row_shapes"]["shapes"] == [[2, 4]] + assert "pixel_values__row_shapes" not in tags[1] + + restored = materialize(fields, tags=tags) + restored_media = restored["pixel_values"] + assert isinstance(restored_media, PackedTensor) + assert len(restored_media) == 2 + assert restored_media.logical_segment_counts_by_row() == [1, 0] + assert torch.equal(restored_media.as_tensor(), media) + + def test_record_to_train_batch_carries_raw_masks_without_applying_them() -> None: record = _record( [ diff --git a/tests/unit/experience/test_rollout_manager.py b/tests/unit/experience/test_rollout_manager.py index b1cc003d0c9..75ee79132a0 100644 --- a/tests/unit/experience/test_rollout_manager.py +++ b/tests/unit/experience/test_rollout_manager.py @@ -53,6 +53,7 @@ ) from nemo_rl.experience.rollout_manager import ( AsyncNemoGymRolloutImpl, + AsyncRolloutImpl, RolloutManager, RolloutOutcome, RolloutRetryPolicy, @@ -111,7 +112,7 @@ async def generate_async(self, data): } ) - manager = object.__new__(RolloutManager) + manager = object.__new__(AsyncRolloutImpl) manager._policy_generation = _Generation() manager._tokenizer = SimpleNamespace( pad_token_id=0, diff --git a/tests/unit/single_controller/test_setup.py b/tests/unit/single_controller/test_setup.py index ba0b075eb8e..23401f3139c 100644 --- a/tests/unit/single_controller/test_setup.py +++ b/tests/unit/single_controller/test_setup.py @@ -1069,8 +1069,14 @@ def test_vlm_processor_used_for_data_and_environment_setup( } assert call_kwargs["is_vlm"] is True warmup_fields = actor_args.dp_client.register_partition.call_args.kwargs["fields"] - assert "pixel_values" in warmup_fields - assert "__nrl_packed_tensor_meta__pixel_values" in warmup_fields + for field in ( + "pixel_values", + "image_grid_thw", + "imgs_sizes", + "num_frames", + "mm_token_type_ids", + ): + assert field in warmup_fields def test_weight_sync_factory_args(self, patched_factories): """create_weight_synchronizer receives policy / generation / topology.""" From f7e892e272635d49f2d40bda9ea168150c71c592 Mon Sep 17 00:00:00 2001 From: Cory Ye Date: Thu, 3 Sep 2026 22:15:22 -0700 Subject: [PATCH 07/19] Add SC L1 and nightly tests. Signed-off-by: Cory Ye --- ...8g-megatron-tp4ep4-async-gym-video.v1.yaml | 16 ++ ...g-megatron-single-controller-async.v1.yaml | 63 ------- ...g-megatron-single-controller-async.v1.yaml | 1 - ...g-megatron-single-controller-async.v1.yaml | 59 ------ ...g-megatron-single-controller-async.v1.yaml | 54 ------ nemo_rl/distributed/virtual_cluster.py | 8 +- ...s_GB200_Megatron_Omni_Single_Controller.sh | 39 ++++ ...i_clevr_megatron_single_controller_1n2g.sh | 174 ++++++++++++++++++ .../nemotron_omni_gym_video_megatron_1n2g.sh | 6 +- ...m_video_megatron_single_controller_1n2g.sh | 174 ++++++++++++++++++ tests/test_suites/disabled.txt | 6 - tests/test_suites/nightly.txt | 3 - tests/test_suites/nightly_gb200.txt | 1 + ...n4g-megatron-single-controller-async.v1.sh | 64 ------- ...0ba3b-clevr-1n4g-megatron_generation.v1.sh | 34 ---- ...n4g-megatron-single-controller-async.v1.sh | 20 +- ...n4g-megatron-single-controller-async.v1.sh | 48 ----- ...n4g-megatron-single-controller-async.v1.sh | 48 ----- .../unit/distributed/test_virtual_cluster.py | 25 +++ 19 files changed, 455 insertions(+), 388 deletions(-) delete mode 100644 examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron-single-controller-async.v1.yaml delete mode 100644 examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-1n4g-megatron-single-controller-async.v1.yaml delete mode 100644 examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-8n4g-megatron-single-controller-async.v1.yaml create mode 100755 tests/functional/L1_Functional_Tests_GB200_Megatron_Omni_Single_Controller.sh create mode 100755 tests/functional/nemotron_omni_clevr_megatron_single_controller_1n2g.sh create mode 100755 tests/functional/nemotron_omni_gym_video_megatron_single_controller_1n2g.sh delete mode 100755 tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron-single-controller-async.v1.sh delete mode 100755 tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.sh delete mode 100755 tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-1n4g-megatron-single-controller-async.v1.sh delete mode 100755 tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-8n4g-megatron-single-controller-async.v1.sh diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-16n8g-megatron-tp4ep4-async-gym-video.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-16n8g-megatron-tp4ep4-async-gym-video.v1.yaml index 20799328e5d..e5896928d53 100644 --- a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-16n8g-megatron-tp4ep4-async-gym-video.v1.yaml +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-16n8g-megatron-tp4ep4-async-gym-video.v1.yaml @@ -49,6 +49,19 @@ policy: generation: bad_words: [] mcore_generation_config: + buffer_size_gb: 8 + transformer_impl: inference_optimized + cuda_graph_impl: local + inference_cuda_graph_scope: block + num_cuda_graphs: -1 + use_cuda_graphs_for_non_decode_steps: false + moe_pad_experts_for_cuda_graph_inference: false + moe_router_dtype: fp32 + expose_http_server: true + vision_embedding_cache_max_bytes: 536870912 + video_num_frames: ${data.default.num_frames} + video_temporal_patch_size: ${data.default.video_temporal_patch_size} + video_target_num_patches: ${data.default.video_target_num_patches} image_dynamic_resolution: true logprobs_mode: raw_logprobs megatron_inference_wrapper: megatron.core.inference.model_inference_wrappers.multimodal.nemotron_omni_inference_wrapper.NemotronOmniInferenceWrapper @@ -91,6 +104,9 @@ data: max_input_seq_length: ${policy.max_total_sequence_length} num_workers: 0 default: + num_frames: 32 + video_sampling_style: nemotron_vl + video_temporal_patch_size: 2 video_target_num_patches: 1024 video_maintain_aspect_ratio: true env: diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron-single-controller-async.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron-single-controller-async.v1.yaml deleted file mode 100644 index 0840a470d9b..00000000000 --- a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron-single-controller-async.v1.yaml +++ /dev/null @@ -1,63 +0,0 @@ -# NeMo-RL v2 SingleController overlay for one-node, non-colocated Omni CLEVR. -defaults: ./vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron_generation.v1.yaml - -grpo: - num_prompts_per_step: 2 - num_generations_per_prompt: 8 - async_grpo: null - val_period: 0 - val_at_start: false - val_at_end: false - -data_plane: - enabled: true - impl: transfer_queue - backend: simple - claim_meta_poll_interval_s: 0.5 - simple: - num_storage_units: 2 - -async_rl: - sampler: - name: in_order - max_lookahead_versions: 1 - recompute_kv_cache_after_weight_updates: false - min_groups_for_streaming_train: ${grpo.num_prompts_per_step} - max_inflight_prompts: ${mul:${grpo.num_prompts_per_step}, 2} - max_buffered_rollouts: ${mul:${grpo.num_prompts_per_step}, 2} - diagnostics: true - -policy: - train_global_batch_size: ${mul:${grpo.num_prompts_per_step}, ${grpo.num_generations_per_prompt}} - megatron_cfg: - tensor_model_parallel_size: 2 - expert_model_parallel_size: 2 - generation: - backend: megatron - max_new_tokens: 512 - colocated: - enabled: false - resources: - num_nodes: 1 - gpus_per_node: 2 - mcore_generation_config: - transformer_impl: inference_optimized - tensor_model_parallel_size: 2 - expert_model_parallel_size: 2 - moe_router_dtype: fp32 - moe_pad_experts_for_cuda_graph_inference: false - cuda_graph_impl: local - inference_cuda_graph_scope: block - num_cuda_graphs: -1 - use_cuda_graphs_for_non_decode_steps: false - enable_chunked_prefill: true - async_sched_mode: async - kv_cache_management_mode: persist - refit_backend: nccl - -data: - num_workers: 0 - -cluster: - num_nodes: 1 - gpus_per_node: 4 diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-single-controller-async.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-single-controller-async.v1.yaml index be31246f1b4..fef0f0c8053 100644 --- a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-single-controller-async.v1.yaml +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-single-controller-async.v1.yaml @@ -23,7 +23,6 @@ async_rl: min_groups_for_streaming_train: ${grpo.num_prompts_per_step} max_inflight_prompts: ${mul:${grpo.num_prompts_per_step}, 2} max_buffered_rollouts: ${mul:${grpo.num_prompts_per_step}, 2} - diagnostics: false policy: generation: diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-1n4g-megatron-single-controller-async.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-1n4g-megatron-single-controller-async.v1.yaml deleted file mode 100644 index 4eca9a9e08f..00000000000 --- a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-1n4g-megatron-single-controller-async.v1.yaml +++ /dev/null @@ -1,59 +0,0 @@ -# NeMo-RL v2 SingleController overlay for one-node, non-colocated Omni VSTAT. -defaults: ./vlm_grpo-nemotron-omni-30ba3b-16n8g-megatron-tp4ep4-async-gym-video.v1.yaml - -grpo: - async_grpo: null - val_period: 0 - val_at_start: false - val_at_end: false - deduplicate_multimodal_data: false - overlong_filtering: false - -data_plane: - enabled: true - impl: transfer_queue - backend: simple - claim_meta_poll_interval_s: 0.5 - simple: - num_storage_units: 2 - -async_rl: - sampler: - name: in_order - max_lookahead_versions: 1 - recompute_kv_cache_after_weight_updates: false - min_groups_for_streaming_train: ${grpo.num_prompts_per_step} - max_inflight_prompts: ${mul:${grpo.num_prompts_per_step}, 2} - max_buffered_rollouts: ${mul:${grpo.num_prompts_per_step}, 2} - diagnostics: true - -policy: - megatron_cfg: - tensor_model_parallel_size: 2 - expert_model_parallel_size: 2 - generation: - backend: megatron - colocated: - enabled: false - resources: - num_nodes: 1 - gpus_per_node: 2 - mcore_generation_config: - transformer_impl: inference_optimized - tensor_model_parallel_size: 2 - expert_model_parallel_size: 2 - moe_router_dtype: fp32 - moe_pad_experts_for_cuda_graph_inference: false - cuda_graph_impl: local - inference_cuda_graph_scope: block - num_cuda_graphs: -1 - use_cuda_graphs_for_non_decode_steps: false - enable_chunked_prefill: true - enable_prefix_caching: false - async_sched_mode: async - kv_cache_management_mode: persist - refit_backend: nccl - -cluster: - num_nodes: 1 - gpus_per_node: 4 diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-8n4g-megatron-single-controller-async.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-8n4g-megatron-single-controller-async.v1.yaml deleted file mode 100644 index 6214e10aa2f..00000000000 --- a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-8n4g-megatron-single-controller-async.v1.yaml +++ /dev/null @@ -1,54 +0,0 @@ -# NeMo-RL v2 SingleController overlay for eight-node, non-colocated Omni VSTAT. -defaults: ./vlm_grpo-nemotron-omni-30ba3b-16n8g-megatron-tp4ep4-async-gym-video.v1.yaml - -grpo: - async_grpo: null - val_period: 0 - val_at_start: false - val_at_end: false - deduplicate_multimodal_data: false - overlong_filtering: false - -data_plane: - enabled: true - impl: transfer_queue - backend: simple - claim_meta_poll_interval_s: 0.5 - simple: - num_storage_units: 16 - -async_rl: - sampler: - name: in_order - max_lookahead_versions: 1 - recompute_kv_cache_after_weight_updates: false - min_groups_for_streaming_train: ${grpo.num_prompts_per_step} - max_inflight_prompts: ${mul:${grpo.num_prompts_per_step}, 2} - max_buffered_rollouts: ${mul:${grpo.num_prompts_per_step}, 2} - diagnostics: false - -policy: - generation: - backend: megatron - colocated: - enabled: false - resources: - num_nodes: 6 - gpus_per_node: 4 - mcore_generation_config: - transformer_impl: inference_optimized - moe_router_dtype: fp32 - moe_pad_experts_for_cuda_graph_inference: false - cuda_graph_impl: local - inference_cuda_graph_scope: block - num_cuda_graphs: -1 - use_cuda_graphs_for_non_decode_steps: false - enable_chunked_prefill: true - enable_prefix_caching: false - async_sched_mode: async - kv_cache_management_mode: persist - refit_backend: nccl - -cluster: - num_nodes: 8 - gpus_per_node: 4 diff --git a/nemo_rl/distributed/virtual_cluster.py b/nemo_rl/distributed/virtual_cluster.py index e65ef299dc3..d7366e76122 100644 --- a/nemo_rl/distributed/virtual_cluster.py +++ b/nemo_rl/distributed/virtual_cluster.py @@ -318,7 +318,13 @@ def init_ray(log_dir: Optional[str] = None) -> None: if _k.startswith(("PMIX_", "PMI_", "MPI_", "OMPI_", "SLURM_")): os.environ.pop(_k, None) - env_vars = dict(os.environ) + # Ray actors deserialize constructor arguments before importing NeMo-RL. + # Put Hugging Face's generated ``transformers_modules`` package on the + # cluster-wide PYTHONPATH so trust_remote_code objects can be unpickled at + # that boundary. This covers both V1 worker groups and direct V2/SC actors. + from nemo_rl.utils.venvs import add_hf_modules_cache_to_pythonpath + + env_vars = add_hf_modules_cache_to_pythonpath(dict(os.environ)) env_vars.pop("RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES", None) runtime_env = { diff --git a/tests/functional/L1_Functional_Tests_GB200_Megatron_Omni_Single_Controller.sh b/tests/functional/L1_Functional_Tests_GB200_Megatron_Omni_Single_Controller.sh new file mode 100755 index 00000000000..ca6ec20bf28 --- /dev/null +++ b/tests/functional/L1_Functional_Tests_GB200_Megatron_Omni_Single_Controller.sh @@ -0,0 +1,39 @@ +#!/bin/bash +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +set -xeuo pipefail + +SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &> /dev/null && pwd) +PROJECT_ROOT=$(realpath "${SCRIPT_DIR}/../..") + +cd "${PROJECT_ROOT}" + +GPU_COUNT=$(nvidia-smi --query-gpu=index --format=csv,noheader | wc -l) +if (( GPU_COUNT < 2 )); then + echo "SKIP: Nemotron Omni SingleController functional tests require at least two GB200 GPUs" + exit 0 +fi + +# SingleController is non-colocated: one GPU trains the frozen-decoder policy +# and one GPU hosts Megatron generation. +export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1}" + +time uv run --no-sync bash ./tests/functional/nemotron_omni_clevr_megatron_single_controller_1n2g.sh +time uv run --no-sync bash ./tests/functional/nemotron_omni_gym_video_megatron_single_controller_1n2g.sh + +cd "${PROJECT_ROOT}/tests" +if compgen -G ".coverage*" > /dev/null; then + coverage combine .coverage* +fi diff --git a/tests/functional/nemotron_omni_clevr_megatron_single_controller_1n2g.sh b/tests/functional/nemotron_omni_clevr_megatron_single_controller_1n2g.sh new file mode 100755 index 00000000000..9de2f6ba706 --- /dev/null +++ b/tests/functional/nemotron_omni_clevr_megatron_single_controller_1n2g.sh @@ -0,0 +1,174 @@ +#!/usr/bin/env bash +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +set -euo pipefail + +SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd) +PROJECT_ROOT=$(realpath "${SCRIPT_DIR}/../..") + +if [[ -z "${HF_TOKEN:-}" ]]; then + echo "SKIP: HF_TOKEN is required for the Omni checkpoint" + exit 0 +fi + +GPU_COUNT=$(nvidia-smi --query-gpu=index --format=csv,noheader | wc -l) +if (( GPU_COUNT < 2 )); then + echo "SKIP: Omni CLEVR SingleController smoke requires at least two visible GPUs" + exit 0 +fi +DETECTED_CUDA_ARCH=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader -i 0) +export TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST:-${DETECTED_CUDA_ARCH}}" +MEGATRON_TRANSFORMER_IMPL="${MEGATRON_TRANSFORMER_IMPL:-inference_optimized}" +MEGATRON_CUDA_GRAPH_IMPL="${MEGATRON_CUDA_GRAPH_IMPL:-local}" +if [[ "${MEGATRON_CUDA_GRAPH_IMPL}" == "local" ]]; then + INFERENCE_CUDA_GRAPH_SCOPE=block + NUM_CUDA_GRAPHS=-1 +else + INFERENCE_CUDA_GRAPH_SCOPE=none + NUM_CUDA_GRAPHS=0 +fi +if [[ "${MEGATRON_TRANSFORMER_IMPL}" != "inference_optimized" && + "${MEGATRON_CUDA_GRAPH_IMPL}" == "local" ]]; then + MOE_PAD_EXPERTS_FOR_CG=true +else + MOE_PAD_EXPERTS_FOR_CG=false +fi + +EXP_NAME=$(basename "$0" .sh) +EXP_DIR="${SCRIPT_DIR}/${EXP_NAME}" +LOG_DIR="${EXP_DIR}/logs" +DATA_ROOT="${EXP_DIR}/data" +TRAIN_PATH="${DATA_ROOT}/train.jsonl" +VAL_PATH="${DATA_ROOT}/val.jsonl" +JSON_METRICS="${EXP_DIR}/metrics.json" +RUN_LOG="${EXP_DIR}/run.log" +rm -rf "${EXP_DIR}" +mkdir -p "${LOG_DIR}" "${DATA_ROOT}" + +cd "${PROJECT_ROOT}" +export PYTHONPATH="${PROJECT_ROOT}:${PYTHONPATH:-}" + +# Match the non-SingleController L1 fixture. +TRAIN_PATH="${TRAIN_PATH}" VAL_PATH="${VAL_PATH}" uv run --no-sync python - <<'PY' +import base64 +import io +import json +import os + +from PIL import Image + +buffer = io.BytesIO() +Image.new("RGB", (224, 224), color="red").save(buffer, format="PNG") +image_url = "data:image/png;base64," + base64.b64encode(buffer.getvalue()).decode() + + +def sample(index: int) -> dict: + return { + "messages": [ + { + "role": "user", + "content": [ + {"type": "image", "image": image_url}, + { + "type": "text", + "text": f"Sample {index}: What color is the image?", + }, + ], + }, + {"role": "assistant", "content": "red"}, + ] + } + + +for path, count in ((os.environ["TRAIN_PATH"], 64), (os.environ["VAL_PATH"], 2)): + with open(path, "w") as output: + for index in range(count): + output.write(json.dumps(sample(index)) + "\n") +PY + +# SingleController requires disaggregated generation. One frozen-decoder model +# fits on each GB200, so split the two visible GPUs 1 trainer + 1 generator. +uv run --no-sync python examples/run_grpo_single_controller.py \ + --config examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron_generation.v1.yaml \ + cluster.num_nodes=1 \ + cluster.gpus_per_node=2 \ + ++cluster.segment_size=1 \ + policy.megatron_cfg.env_vars.TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST}" \ + policy.megatron_cfg.tensor_model_parallel_size=1 \ + policy.megatron_cfg.expert_model_parallel_size=1 \ + policy.megatron_cfg.expert_tensor_parallel_size=1 \ + policy.megatron_cfg.context_parallel_size=1 \ + policy.megatron_cfg.sequence_parallel=true \ + policy.megatron_cfg.activation_checkpointing=true \ + ++policy.megatron_cfg.freeze_config.freeze_language_model=true \ + +policy.megatron_cfg.bias_dropout_fusion=false \ + policy.megatron_cfg.optimizer.optimizer_cpu_offload=false \ + policy.megatron_cfg.optimizer.optimizer_offload_fraction=0.0 \ + ++policy.megatron_cfg.optimizer.params_dtype=bfloat16 \ + ++policy.megatron_cfg.optimizer.main_grads_dtype=bfloat16 \ + ++policy.megatron_cfg.optimizer.main_params_dtype=float16 \ + ++policy.megatron_cfg.optimizer.exp_avg_dtype=bfloat16 \ + ++policy.megatron_cfg.optimizer.exp_avg_sq_dtype=bfloat16 \ + ++policy.megatron_cfg.optimizer.store_param_remainders=false \ + policy.generation.backend=megatron \ + policy.generation.colocated.enabled=false \ + policy.generation.colocated.resources.num_nodes=1 \ + policy.generation.colocated.resources.gpus_per_node=1 \ + policy.generation.max_new_tokens=128 \ + policy.generation.mcore_generation_config.tensor_model_parallel_size=1 \ + policy.generation.mcore_generation_config.expert_model_parallel_size=1 \ + policy.generation.mcore_generation_config.expert_tensor_parallel_size=1 \ + ++policy.generation.mcore_generation_config.context_parallel_size=1 \ + ++policy.generation.mcore_generation_config.moe_router_dtype=fp32 \ + policy.generation.mcore_generation_config.transformer_impl="${MEGATRON_TRANSFORMER_IMPL}" \ + policy.generation.mcore_generation_config.sequence_parallel=true \ + policy.generation.mcore_generation_config.refit_backend=nccl \ + policy.generation.mcore_generation_config.buffer_size_gb=2 \ + policy.generation.mcore_generation_config.cuda_graph_impl="${MEGATRON_CUDA_GRAPH_IMPL}" \ + policy.generation.mcore_generation_config.inference_cuda_graph_scope="${INFERENCE_CUDA_GRAPH_SCOPE}" \ + policy.generation.mcore_generation_config.num_cuda_graphs="${NUM_CUDA_GRAPHS}" \ + policy.generation.mcore_generation_config.use_cuda_graphs_for_non_decode_steps=false \ + policy.generation.mcore_generation_config.moe_pad_experts_for_cuda_graph_inference="${MOE_PAD_EXPERTS_FOR_CG}" \ + policy.generation.mcore_generation_config.enable_chunked_prefill=true \ + ++policy.generation.mcore_generation_config.async_sched_mode=async \ + policy.generation.mcore_generation_config.max_model_len=1024 \ + policy.generation.mcore_generation_config.max_tokens=1024 \ + policy.max_total_sequence_length=1024 \ + data.train.dataset_name=ResponseDataset \ + ++data.train.data_path="${TRAIN_PATH}" \ + data.train.split=train \ + data.validation.dataset_name=ResponseDataset \ + ++data.validation.data_path="${VAL_PATH}" \ + data.validation.split=train \ + data.num_workers=0 \ + grpo.async_grpo=null \ + grpo.num_prompts_per_step=1 \ + grpo.num_generations_per_prompt=2 \ + grpo.max_num_steps=1 \ + grpo.val_period=0 \ + grpo.val_at_start=false \ + grpo.val_at_end=false \ + policy.train_global_batch_size=2 \ + policy.train_micro_batch_size=1 \ + ++data_plane.enabled=true \ + ++data_plane.impl=transfer_queue \ + ++data_plane.backend=simple \ + ++data_plane.claim_meta_poll_interval_s=0.5 \ + ++data_plane.simple.num_storage_units=2 \ + ++async_rl.sampler.name=in_order \ + ++async_rl.sampler.max_lookahead_versions=1 \ + ++async_rl.recompute_kv_cache_after_weight_updates=false \ + ++async_rl.min_groups_for_streaming_train=1 \ + ++async_rl.max_inflight_prompts=2 \ + ++async_rl.max_buffered_rollouts=2 \ + logger.tensorboard_enabled=true \ + logger.log_dir="${LOG_DIR}" \ + logger.wandb_enabled=false \ + logger.monitor_gpus=false \ + checkpointing.enabled=false \ + "$@" 2>&1 | tee "${RUN_LOG}" + +uv run --no-sync tests/json_dump_tb_logs.py "${LOG_DIR}" --output_path "${JSON_METRICS}" +uv run --no-sync tests/check_metrics.py "${JSON_METRICS}" \ + 'max(data["train/gen_kl_error"]) < 0.05' \ + 'all_finite(data["train/reward"])' diff --git a/tests/functional/nemotron_omni_gym_video_megatron_1n2g.sh b/tests/functional/nemotron_omni_gym_video_megatron_1n2g.sh index 3fd6aad2675..54d84dd1b0d 100755 --- a/tests/functional/nemotron_omni_gym_video_megatron_1n2g.sh +++ b/tests/functional/nemotron_omni_gym_video_megatron_1n2g.sh @@ -133,9 +133,9 @@ uv run --no-sync python examples/nemo_gym/run_grpo_nemo_gym.py \ ++policy.generation.mcore_generation_config.video_temporal_patch_size=2 \ ++policy.generation.mcore_generation_config.video_target_num_patches=256 \ policy.max_total_sequence_length=1024 \ - +data.default.num_frames=8 \ - +data.default.video_sampling_style=nemotron_vl \ - +data.default.video_temporal_patch_size=2 \ + data.default.num_frames=8 \ + data.default.video_sampling_style=nemotron_vl \ + data.default.video_temporal_patch_size=2 \ +data.default.min_generation_tokens=128 \ data.default.video_target_num_patches=256 \ data.train.data_path="${TRAIN_PATH}" \ diff --git a/tests/functional/nemotron_omni_gym_video_megatron_single_controller_1n2g.sh b/tests/functional/nemotron_omni_gym_video_megatron_single_controller_1n2g.sh new file mode 100755 index 00000000000..86d24273d8f --- /dev/null +++ b/tests/functional/nemotron_omni_gym_video_megatron_single_controller_1n2g.sh @@ -0,0 +1,174 @@ +#!/usr/bin/env bash +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +set -euo pipefail + +SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd) +PROJECT_ROOT=$(realpath "${SCRIPT_DIR}/../..") + +if [[ -z "${HF_TOKEN:-}" ]]; then + echo "SKIP: HF_TOKEN is required for the Omni checkpoint" + exit 0 +fi + +GPU_COUNT=$(nvidia-smi --query-gpu=index --format=csv,noheader | wc -l) +if (( GPU_COUNT < 2 )); then + echo "SKIP: Omni Gym-video SingleController smoke requires at least two GPUs" + exit 0 +fi +DETECTED_CUDA_ARCH=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader -i 0) +export TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST:-${DETECTED_CUDA_ARCH}}" +MEGATRON_TRANSFORMER_IMPL="${MEGATRON_TRANSFORMER_IMPL:-inference_optimized}" +MEGATRON_CUDA_GRAPH_IMPL="${MEGATRON_CUDA_GRAPH_IMPL:-local}" +if [[ "${MEGATRON_CUDA_GRAPH_IMPL}" == "local" ]]; then + INFERENCE_CUDA_GRAPH_SCOPE=block + NUM_CUDA_GRAPHS=-1 +else + INFERENCE_CUDA_GRAPH_SCOPE=none + NUM_CUDA_GRAPHS=0 +fi +if [[ "${MEGATRON_TRANSFORMER_IMPL}" != "inference_optimized" && + "${MEGATRON_CUDA_GRAPH_IMPL}" == "local" ]]; then + MOE_PAD_EXPERTS_FOR_CG=true +else + MOE_PAD_EXPERTS_FOR_CG=false +fi + +EXP_NAME=$(basename "$0" .sh) +EXP_DIR="${SCRIPT_DIR}/${EXP_NAME}" +LOG_DIR="${EXP_DIR}/logs" +DATA_ROOT="${EXP_DIR}/data" +VIDEO_PATH="${DATA_ROOT}/red.mp4" +RAW_TRAIN_PATH="${DATA_ROOT}/train-raw.jsonl" +RAW_VAL_PATH="${DATA_ROOT}/val-raw.jsonl" +TRAIN_PATH="${DATA_ROOT}/train-gym.jsonl" +VAL_PATH="${DATA_ROOT}/val-gym.jsonl" +JSON_METRICS="${EXP_DIR}/metrics.json" +RUN_LOG="${EXP_DIR}/run.log" +rm -rf "${EXP_DIR}" +mkdir -p "${LOG_DIR}" "${DATA_ROOT}" + +cd "${PROJECT_ROOT}" +export PYTHONPATH="${PROJECT_ROOT}:${PYTHONPATH:-}" +export NRL_VIDEO_BACKEND=torchcodec +export NRL_VIDEO_SAMPLING_STYLE=nemotron_vl +export NRL_VIDEO_TEMPORAL_PATCH_SIZE=2 + +bash tools/install_audio_deps.sh +ffmpeg -hide_banner -loglevel error -y \ + -f lavfi -i color=c=red:s=224x224:r=8:d=2 \ + -c:v libx264 -pix_fmt yuv420p "${VIDEO_PATH}" + +for sample_id in $(seq 1 64); do + jq -nc \ + --arg prompt "Sample ${sample_id}: What color fills the video? A. Red B. Blue" \ + --arg video "${VIDEO_PATH}" \ + '{prompt: $prompt, video: $video, answer: "A", verifier: "mcqa"}' +done > "${RAW_TRAIN_PATH}" +for sample_id in $(seq 1 2); do + jq -nc \ + --arg prompt "Validation ${sample_id}: What color fills the video? A. Red B. Blue" \ + --arg video "${VIDEO_PATH}" \ + '{prompt: $prompt, video: $video, answer: "A", verifier: "mcqa"}' +done > "${RAW_VAL_PATH}" + +uv run --no-sync examples/nemo_gym/prepare_video_dataset.py convert \ + --input "${RAW_TRAIN_PATH}" \ + --output "${TRAIN_PATH}" +uv run --no-sync examples/nemo_gym/prepare_video_dataset.py convert \ + --input "${RAW_VAL_PATH}" \ + --output "${VAL_PATH}" + +# SingleController requires disaggregated generation. One frozen-decoder model +# fits on each GB200, so split the two visible GPUs 1 trainer + 1 generator. +uv run --no-sync python examples/run_grpo_single_controller.py \ + --config examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-16n8g-megatron-tp4ep4-async-gym-video.v1.yaml \ + cluster.num_nodes=1 \ + cluster.gpus_per_node=2 \ + ++cluster.segment_size=1 \ + policy.megatron_cfg.env_vars.TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST}" \ + policy.megatron_cfg.tensor_model_parallel_size=1 \ + policy.megatron_cfg.pipeline_model_parallel_size=1 \ + policy.megatron_cfg.expert_model_parallel_size=1 \ + policy.megatron_cfg.expert_tensor_parallel_size=1 \ + policy.megatron_cfg.context_parallel_size=1 \ + policy.megatron_cfg.sequence_parallel=true \ + policy.megatron_cfg.activation_checkpointing=true \ + ++policy.megatron_cfg.freeze_config.freeze_language_model=true \ + policy.megatron_cfg.optimizer.optimizer_cpu_offload=false \ + policy.megatron_cfg.optimizer.optimizer_offload_fraction=0.0 \ + ++policy.megatron_cfg.optimizer.params_dtype=bfloat16 \ + ++policy.megatron_cfg.optimizer.main_grads_dtype=bfloat16 \ + ++policy.megatron_cfg.optimizer.main_params_dtype=float16 \ + ++policy.megatron_cfg.optimizer.exp_avg_dtype=bfloat16 \ + ++policy.megatron_cfg.optimizer.exp_avg_sq_dtype=bfloat16 \ + ++policy.megatron_cfg.optimizer.store_param_remainders=false \ + policy.generation.backend=megatron \ + ++policy.generation.bad_words=null \ + policy.generation.colocated.enabled=false \ + policy.generation.colocated.resources.num_nodes=1 \ + policy.generation.colocated.resources.gpus_per_node=1 \ + policy.generation.max_new_tokens=128 \ + policy.generation.mcore_generation_config.expose_http_server=true \ + policy.generation.mcore_generation_config.tensor_model_parallel_size=1 \ + policy.generation.mcore_generation_config.expert_model_parallel_size=1 \ + policy.generation.mcore_generation_config.expert_tensor_parallel_size=1 \ + ++policy.generation.mcore_generation_config.context_parallel_size=1 \ + ++policy.generation.mcore_generation_config.moe_router_dtype=fp32 \ + policy.generation.mcore_generation_config.transformer_impl="${MEGATRON_TRANSFORMER_IMPL}" \ + policy.generation.mcore_generation_config.sequence_parallel=true \ + policy.generation.mcore_generation_config.refit_backend=nccl \ + policy.generation.mcore_generation_config.buffer_size_gb=2 \ + policy.generation.mcore_generation_config.cuda_graph_impl="${MEGATRON_CUDA_GRAPH_IMPL}" \ + policy.generation.mcore_generation_config.inference_cuda_graph_scope="${INFERENCE_CUDA_GRAPH_SCOPE}" \ + policy.generation.mcore_generation_config.num_cuda_graphs="${NUM_CUDA_GRAPHS}" \ + policy.generation.mcore_generation_config.use_cuda_graphs_for_non_decode_steps=false \ + ++policy.generation.mcore_generation_config.moe_pad_experts_for_cuda_graph_inference="${MOE_PAD_EXPERTS_FOR_CG}" \ + policy.generation.mcore_generation_config.enable_chunked_prefill=true \ + ++policy.generation.mcore_generation_config.async_sched_mode=async \ + policy.generation.mcore_generation_config.enable_prefix_caching=true \ + policy.generation.mcore_generation_config.max_model_len=1024 \ + policy.generation.mcore_generation_config.max_tokens=1024 \ + ++policy.generation.mcore_generation_config.video_num_frames=8 \ + ++policy.generation.mcore_generation_config.video_temporal_patch_size=2 \ + ++policy.generation.mcore_generation_config.video_target_num_patches=256 \ + policy.max_total_sequence_length=1024 \ + data.default.num_frames=8 \ + data.default.video_sampling_style=nemotron_vl \ + data.default.video_temporal_patch_size=2 \ + +data.default.min_generation_tokens=128 \ + data.default.video_target_num_patches=256 \ + data.train.data_path="${TRAIN_PATH}" \ + data.validation.data_path="${VAL_PATH}" \ + grpo.deduplicate_multimodal_data=false \ + grpo.async_grpo=null \ + grpo.num_prompts_per_step=1 \ + grpo.num_generations_per_prompt=2 \ + grpo.max_num_steps=1 \ + grpo.val_period=0 \ + grpo.val_at_start=false \ + grpo.val_at_end=false \ + policy.train_global_batch_size=2 \ + policy.train_micro_batch_size=1 \ + ++data_plane.enabled=true \ + ++data_plane.impl=transfer_queue \ + ++data_plane.backend=simple \ + ++data_plane.claim_meta_poll_interval_s=0.5 \ + ++data_plane.simple.num_storage_units=2 \ + ++async_rl.sampler.name=in_order \ + ++async_rl.sampler.max_lookahead_versions=1 \ + ++async_rl.recompute_kv_cache_after_weight_updates=false \ + ++async_rl.min_groups_for_streaming_train=1 \ + ++async_rl.max_inflight_prompts=2 \ + ++async_rl.max_buffered_rollouts=2 \ + logger.tensorboard_enabled=true \ + logger.log_dir="${LOG_DIR}" \ + logger.wandb_enabled=false \ + logger.monitor_gpus=false \ + checkpointing.enabled=false \ + "$@" 2>&1 | tee "${RUN_LOG}" + +uv run --no-sync tests/json_dump_tb_logs.py "${LOG_DIR}" --output_path "${JSON_METRICS}" +uv run --no-sync tests/check_metrics.py "${JSON_METRICS}" \ + 'max(data["train/gen_kl_error"]) < 0.05' \ + 'all_finite(data["train/reward"])' diff --git a/tests/test_suites/disabled.txt b/tests/test_suites/disabled.txt index 592348c897c..fcbfd7e18e1 100644 --- a/tests/test_suites/disabled.txt +++ b/tests/test_suites/disabled.txt @@ -31,12 +31,6 @@ tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-click-2n8g-megatron-t # functional test before moving it to a recurring suite. tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-count-1n4g-megatron_generation.v1.sh -# TODO(@cspades): Multimodal SingleController/TransferQueue recipes awaiting prepared -# fixtures and validation. The 1n4g CLEVR sibling remains enabled as an L1 smoke test. -tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-single-controller-async.v1.sh -tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-1n4g-megatron-single-controller-async.v1.sh -tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-8n4g-megatron-single-controller-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/nightly.txt b/tests/test_suites/nightly.txt index 8910ac1c4f6..93171964fad 100644 --- a/tests/test_suites/nightly.txt +++ b/tests/test_suites/nightly.txt @@ -58,9 +58,6 @@ tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8.v tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-2n8g-megatron-tp8ep8.v1.sh tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-megatron-tp8ep16.v1.sh -# L1 multimodal SingleController/TransferQueue smoke coverage -tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron-single-controller-async.v1.sh - # Functional Qwen3.5-35B VLM GRPO run # The AutoModel variant is re-enabled with the vLLM 0.25.1 bump (no longer hits # https://github.com/vllm-project/vllm/issues/36237). The Megatron variant still diff --git a/tests/test_suites/nightly_gb200.txt b/tests/test_suites/nightly_gb200.txt index 9ab02ba24ce..d2c8c9acf63 100644 --- a/tests/test_suites/nightly_gb200.txt +++ b/tests/test_suites/nightly_gb200.txt @@ -32,6 +32,7 @@ tests/test_suites/llm/grpo-qwen3-30ba3b-4n4g-megatron-qa-nvfp4-w4a4-real.sh tests/test_suites/vlm/vlm_grpo-qwen2.5-vl-3b-instruct-clevr-1n4g-dtensor2tp1.v1.sh tests/test_suites/vlm/vlm_grpo-qwen2.5-vl-3b-instruct-clevr-1n4g-megatrontp1.v1.sh tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron_generation.v1.sh +tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-single-controller-async.v1.sh # Deepscaler (short tests) tests/test_suites/llm/grpo-deepscaler-1.5b-1n4g-8K.sh diff --git a/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron-single-controller-async.v1.sh b/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron-single-controller-async.v1.sh deleted file mode 100755 index 52e9c9ab9cc..00000000000 --- a/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron-single-controller-async.v1.sh +++ /dev/null @@ -1,64 +0,0 @@ -#!/bin/bash -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) -source "$SCRIPT_DIR/common.env" - -# Two steps exercise CP=1 multimodal generation, TQ policy training, weight -# refit, and a post-refit rollout without treating this smoke test as a -# convergence run. CP>1 + multimodal + TQ remains untested. -# ===== BEGIN CONFIG ===== -NUM_NODES=1 -GPUS_PER_NODE=4 -STEPS_PER_RUN=2 -MAX_STEPS=2 -NUM_RUNS=1 -NUM_MINUTES=120 -# ===== END CONFIG ===== - -exit_if_max_steps_reached - -cd "$PROJECT_ROOT" - -uv run examples/run_grpo_single_controller.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" - -uv run tests/json_dump_tb_logs.py "$LOG_DIR" --output_path "$JSON_METRICS" - -RECORDED_STEP=$(jq -r \ - 'if has("train/loss") then (."train/loss" | keys | map(tonumber) | max // 0) else 0 end' \ - "$JSON_METRICS") -if [[ "$RECORDED_STEP" -lt 1 ]]; then - echo "[ERROR] Expected at least one completed training step" - exit 1 -fi - -uv run tests/check_metrics.py "$JSON_METRICS" \ - 'all_finite(data["train/token_mult_prob_error"])' \ - 'max(data["train/loss"]) < 1000000.0' \ - 'min(data["train/loss"]) > -1000000.0' - -rm -rf "$CKPT_DIR" diff --git a/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.sh b/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.sh deleted file mode 100755 index f8d2e75cdf8..00000000000 --- a/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.sh +++ /dev/null @@ -1,34 +0,0 @@ -#!/bin/bash -SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) -source $SCRIPT_DIR/common.env - -# TODO(@cspades): Run and validate this functional test, then add golden -# convergence metrics before enabling it in a recurring suite. - -# ===== BEGIN CONFIG ===== -NUM_NODES=1 -GPUS_PER_NODE=4 -STEPS_PER_RUN=4 -MAX_STEPS=4 -NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) -NUM_MINUTES=120 -# ===== END CONFIG ===== - -exit_if_max_steps_reached - -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 - -uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS diff --git a/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-single-controller-async.v1.sh b/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-single-controller-async.v1.sh index 48ba6451ce1..d0b8676b1a7 100755 --- a/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-single-controller-async.v1.sh +++ b/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-single-controller-async.v1.sh @@ -16,12 +16,16 @@ SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) source "$SCRIPT_DIR/common.env" +# Match the non-SingleController convergence run while exercising the +# SingleController/TransferQueue orchestration path. + # ===== BEGIN CONFIG ===== NUM_NODES=8 GPUS_PER_NODE=4 -STEPS_PER_RUN=2 -MAX_STEPS=2 -NUM_RUNS=1 +SEGMENT_SIZE=2 +STEPS_PER_RUN=50 +MAX_STEPS=50 +NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) NUM_MINUTES=120 # ===== END CONFIG ===== @@ -32,15 +36,23 @@ cd "$PROJECT_ROOT" uv run examples/run_grpo_single_controller.py \ --config "$CONFIG_PATH" \ grpo.max_num_steps="$MAX_STEPS" \ + policy.megatron_cfg.scheduler.lr_warmup_iters=10 \ 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.enabled=False \ checkpointing.checkpoint_dir="$CKPT_DIR" \ "$@" \ 2>&1 | tee "$RUN_LOG" uv run tests/json_dump_tb_logs.py "$LOG_DIR" --output_path "$JSON_METRICS" + +uv run tests/check_metrics.py "$JSON_METRICS" \ + 'all_finite(data["train/loss"])' \ + 'all_finite(data["train/grad_norm"])' \ + 'min(data["train/grad_norm"]) > 0' \ + 'all_finite(data["train/token_mult_prob_error"])' \ + 'mean(data["train/reward"], range_start=-10) > 0.6' diff --git a/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-1n4g-megatron-single-controller-async.v1.sh b/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-1n4g-megatron-single-controller-async.v1.sh deleted file mode 100755 index 65703c42e74..00000000000 --- a/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-1n4g-megatron-single-controller-async.v1.sh +++ /dev/null @@ -1,48 +0,0 @@ -#!/bin/bash -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) -source "$SCRIPT_DIR/common.env" - -# TODO(@cspades): Prepare a self-contained VSTAT fixture and validate this -# driver before moving it from disabled.txt into a recurring suite. -# ===== BEGIN CONFIG ===== -NUM_NODES=1 -GPUS_PER_NODE=4 -STEPS_PER_RUN=2 -MAX_STEPS=2 -NUM_RUNS=1 -NUM_MINUTES=120 -# ===== END CONFIG ===== - -exit_if_max_steps_reached - -cd "$PROJECT_ROOT" - -uv run examples/run_grpo_single_controller.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" - -uv run tests/json_dump_tb_logs.py "$LOG_DIR" --output_path "$JSON_METRICS" diff --git a/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-8n4g-megatron-single-controller-async.v1.sh b/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-8n4g-megatron-single-controller-async.v1.sh deleted file mode 100755 index 3f65e783dbc..00000000000 --- a/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-8n4g-megatron-single-controller-async.v1.sh +++ /dev/null @@ -1,48 +0,0 @@ -#!/bin/bash -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) -source "$SCRIPT_DIR/common.env" - -# TODO(@cspades): Prepare a self-contained VSTAT fixture and validate this -# driver before moving it from disabled.txt into a recurring suite. -# ===== BEGIN CONFIG ===== -NUM_NODES=8 -GPUS_PER_NODE=4 -STEPS_PER_RUN=2 -MAX_STEPS=2 -NUM_RUNS=1 -NUM_MINUTES=120 -# ===== END CONFIG ===== - -exit_if_max_steps_reached - -cd "$PROJECT_ROOT" - -uv run examples/run_grpo_single_controller.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" - -uv run tests/json_dump_tb_logs.py "$LOG_DIR" --output_path "$JSON_METRICS" diff --git a/tests/unit/distributed/test_virtual_cluster.py b/tests/unit/distributed/test_virtual_cluster.py index 4a17a69fcff..a832b69e5ac 100644 --- a/tests/unit/distributed/test_virtual_cluster.py +++ b/tests/unit/distributed/test_virtual_cluster.py @@ -283,6 +283,31 @@ def test_maybe_configure_data_plane_env_then_init_ray_threads_env_vars(): assert env_vars["MC_ENABLE_DEST_DEVICE_AFFINITY"] == "1" +def test_init_ray_adds_hf_modules_cache_to_cluster_pythonpath(): + """Direct actors must import trust_remote_code classes while unpickling.""" + from nemo_rl.distributed.virtual_cluster import init_ray + + with ( + patch("ray.init") as mock_ray_init, + patch("ray.cluster_resources") as mock_cluster_resources, + ): + mock_cluster_resources.return_value = {"nrl_tag_0": 1} + env = { + "CUDA_VISIBLE_DEVICES": "0", + "HF_MODULES_CACHE": "/hf/modules", + "PYTHONPATH": "/project", + } + with patch.dict(os.environ, env, clear=True): + init_ray() + + env_vars = mock_ray_init.call_args_list[0][1]["runtime_env"]["env_vars"] + assert env_vars["HF_MODULES_CACHE"] == "/hf/modules" + assert env_vars["PYTHONPATH"].split(os.pathsep) == [ + "/hf/modules", + "/project", + ] + + def test_init_ray_alone_has_no_data_plane_awareness(): """Every non-data-plane launcher's call (bare init_ray(), no preceding maybe_configure_data_plane_env) must not touch mooncake env vars -- From 556807963e1426fb4e2480f36819325e7a418472 Mon Sep 17 00:00:00 2001 From: Cory Ye Date: Fri, 4 Sep 2026 11:02:32 -0700 Subject: [PATCH 08/19] Fix rebase errors to reenable V2 multimodal with MCore. Signed-off-by: Cory Ye --- nemo_rl/data_plane/worker_mixin.py | 8 +++---- .../policy/workers/megatron_policy_worker.py | 7 +++++++ .../policy/test_megatron_split_state.py | 21 +++++++++++++++++++ 3 files changed, 31 insertions(+), 5 deletions(-) diff --git a/nemo_rl/data_plane/worker_mixin.py b/nemo_rl/data_plane/worker_mixin.py index 92a01628ecd..a5309d453f8 100644 --- a/nemo_rl/data_plane/worker_mixin.py +++ b/nemo_rl/data_plane/worker_mixin.py @@ -319,11 +319,9 @@ def setup_data_plane(self, cfg: DataPlaneRuntimeConfig) -> None: # entrypoints then delegate to the same ``train`` / ``get_logprobs`` # that carry the flag into ``models/megatron/data.py``. # - # ``train_microbatch_presharded`` is the exception: it lands in - # ``_train_microbatch_body``, which passes none of the capability flags - # and never attaches the media-token validity mask. That path is - # SingleController-only, so ``train_microbatch`` raises for a - # multimodal model rather than training on rows it mis-describes. + # ``train_microbatch_presharded`` lands in ``_train_microbatch_body``, + # which applies the same media-token validity mask and model packing/CP + # capability flags as the regular training path. if self._dp_client is not None: return self._route_fallback_counts = Counter() diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index 3b1dae699fe..c350c3b2c10 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -1455,6 +1455,9 @@ def train_microbatch( explicitly in ``finish_train_step``. Returns nothing: gradients land in ``param.main_grad`` and per-microbatch metrics accumulate in the open-step state until ``finish_train_step`` surfaces them. + + Multimodal validity-mask and model-owned packing/CP behavior match the + regular ``train`` path. """ state = self._assert_step_open() try: @@ -1541,6 +1544,10 @@ def _train_microbatch_body( # call carries one DP slice; the iterator subdivides into pipeline # microbatches. attach_media_token_validity_mask(data, self.media_placeholder_token_id) +<<<<<<< HEAD +======= + +>>>>>>> f1369705e (Fix rebase errors to reenable V2 multimodal with MCore.) ( data_iterator, num_microbatches, diff --git a/tests/unit/models/policy/test_megatron_split_state.py b/tests/unit/models/policy/test_megatron_split_state.py index c3e023bc4e1..990397b4bc4 100644 --- a/tests/unit/models/policy/test_megatron_split_state.py +++ b/tests/unit/models/policy/test_megatron_split_state.py @@ -410,6 +410,27 @@ def test_finish_without_begin_raises(self, mock_module_symbols): class TestTrainMicrobatch: + def test_forwards_multimodal_iterator_capabilities(self, mock_module_symbols): + from nemo_rl.algorithms.loss.interfaces import LossType + + w = _make_worker(LossType.TOKEN_LEVEL) + w.media_placeholder_token_id = 42 + w.delegate_pack_to_model = True + w.delegate_mtp_loss_mask_to_model = True + batch = _fake_batch() + + with patch( + f"{WORKER_MOD}.attach_media_token_validity_mask" + ) as attach_validity_mask: + w.begin_train_step(loss_fn=w._test_loss_fn) + w.train_microbatch(batch) + + attach_validity_mask.assert_called_once_with(batch, 42) + iterator_kwargs = mock_module_symbols["gmi"].call_args.kwargs + assert iterator_kwargs["delegate_pack_to_model"] is True + assert iterator_kwargs["delegate_mtp_loss_mask_to_model"] is True + assert iterator_kwargs["model_slices_context_parallel_inputs"] is False + def test_wraps_forward_backward_in_no_sync(self, mock_module_symbols): """The single most important assertion in this file. Without the no_sync wrap, mcore DDP dispatches a per-call cross-DP reduce on From e4ee67a2ed29e0f251b10511829b9bd45c7da943 Mon Sep 17 00:00:00 2001 From: Cory Ye Date: Fri, 4 Sep 2026 18:11:07 -0700 Subject: [PATCH 09/19] Get rid of rebase artifacts. Signed-off-by: Cory Ye --- ...8g-megatron-tp4ep4-async-gym-video.v1.yaml | 5 -- ...g-megatron-single-controller-async.v1.yaml | 32 ----------- examples/run_grpo_single_controller.py | 2 +- nemo_rl/data/multimodal_utils.py | 4 +- nemo_rl/data/processors.py | 54 +++---------------- nemo_rl/distributed/virtual_cluster.py | 4 +- nemo_rl/experience/payload.py | 4 +- nemo_rl/models/generation/megatron/config.py | 5 +- nemo_rl/models/value/tq_value.py | 4 -- tests/unit/experience/test_payload.py | 2 +- .../test_rollout_generation_failures.py | 2 + tests/unit/experience/test_rollout_manager.py | 19 ++++--- .../test_rollout_manager_router_replay.py | 2 + .../models/policy/test_megatron_worker.py | 25 --------- tests/unit/single_controller/test_setup.py | 12 ++--- 15 files changed, 38 insertions(+), 138 deletions(-) diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-16n8g-megatron-tp4ep4-async-gym-video.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-16n8g-megatron-tp4ep4-async-gym-video.v1.yaml index e5896928d53..4dafdca0185 100644 --- a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-16n8g-megatron-tp4ep4-async-gym-video.v1.yaml +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-16n8g-megatron-tp4ep4-async-gym-video.v1.yaml @@ -50,14 +50,9 @@ policy: bad_words: [] mcore_generation_config: buffer_size_gb: 8 - transformer_impl: inference_optimized - cuda_graph_impl: local - inference_cuda_graph_scope: block - num_cuda_graphs: -1 use_cuda_graphs_for_non_decode_steps: false moe_pad_experts_for_cuda_graph_inference: false moe_router_dtype: fp32 - expose_http_server: true vision_embedding_cache_max_bytes: 536870912 video_num_frames: ${data.default.num_frames} video_temporal_patch_size: ${data.default.video_temporal_patch_size} diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-single-controller-async.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-single-controller-async.v1.yaml index fef0f0c8053..02392c7f296 100644 --- a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-single-controller-async.v1.yaml +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-single-controller-async.v1.yaml @@ -1,20 +1,13 @@ -# NeMo-RL v2 SingleController overlay for eight-node, non-colocated Omni CLEVR. defaults: ./vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron_generation.v1.yaml - grpo: async_grpo: null val_period: 0 val_at_start: false val_at_end: false - data_plane: enabled: true - impl: transfer_queue - backend: simple - claim_meta_poll_interval_s: 0.5 simple: num_storage_units: 16 - async_rl: sampler: name: in_order @@ -23,28 +16,3 @@ async_rl: min_groups_for_streaming_train: ${grpo.num_prompts_per_step} max_inflight_prompts: ${mul:${grpo.num_prompts_per_step}, 2} max_buffered_rollouts: ${mul:${grpo.num_prompts_per_step}, 2} - -policy: - generation: - backend: megatron - colocated: - enabled: false - resources: - num_nodes: 6 - gpus_per_node: 4 - mcore_generation_config: - transformer_impl: inference_optimized - moe_router_dtype: fp32 - moe_pad_experts_for_cuda_graph_inference: false - cuda_graph_impl: local - inference_cuda_graph_scope: block - num_cuda_graphs: -1 - use_cuda_graphs_for_non_decode_steps: false - enable_chunked_prefill: true - async_sched_mode: async - kv_cache_management_mode: persist - refit_backend: nccl - -cluster: - num_nodes: 8 - gpus_per_node: 4 diff --git a/examples/run_grpo_single_controller.py b/examples/run_grpo_single_controller.py index 8456f51e7d1..b06b7ccf9ce 100644 --- a/examples/run_grpo_single_controller.py +++ b/examples/run_grpo_single_controller.py @@ -127,7 +127,7 @@ def main() -> None: init_ray() processor = None - if config.policy.get("is_vlm", False): + if config.policy.get("is_vlm"): processor = get_tokenizer(config.policy["tokenizer"], get_processor=True) tokenizer = processor.tokenizer else: diff --git a/nemo_rl/data/multimodal_utils.py b/nemo_rl/data/multimodal_utils.py index c2b79d73c63..c9d3d6b1a1b 100644 --- a/nemo_rl/data/multimodal_utils.py +++ b/nemo_rl/data/multimodal_utils.py @@ -1094,7 +1094,9 @@ def encode_multimodal_for_wire( # Model inputs some remote-code processors omit from ``model_input_names`` even -# though their forward requires them. Keep extraction and TQ schema warmup aligned. +# though their forward requires them. Consumed by +# ``extract_multimodal_model_inputs``; membership here does NOT imply the field +# is wire-registered (see ``PACKED_/PER_TOKEN_MULTIMODAL_FIELDS``). UNDECLARED_MULTIMODAL_MODEL_INPUTS = ( "imgs_sizes", "num_frames", diff --git a/nemo_rl/data/processors.py b/nemo_rl/data/processors.py index ef8744afb7b..77cb14dae2f 100644 --- a/nemo_rl/data/processors.py +++ b/nemo_rl/data/processors.py @@ -460,9 +460,8 @@ def vlm_hf_data_processor( from nemo_rl.data.datasets.response_datasets.refcoco import format_refcoco_dataset from nemo_rl.data.multimodal_utils import ( PackedTensor, - get_dim_to_pack_along, + extract_multimodal_model_inputs, get_multimodal_default_settings_from_processor, - get_multimodal_keys_from_processor, resolve_to_image, uses_image_placeholder, ) @@ -608,50 +607,13 @@ def vlm_hf_data_processor( # add this for backward compatibility user_message["token_ids"] = message["input_ids"][0] - # add all keys and values to the user message, and the list of keys - multimodal_keys = list(get_multimodal_keys_from_processor(processor)) - # Current Nemotron Omni processors emit imgs_sizes. Historical MMPR - # checkpoints instead emit a batch of fixed-size image tiles and only - # declare pixel_values. Treat each tile as one dynamic-resolution image so - # the Nemotron Omni path can patchify it and preserve the processor's exact - # placeholder count. - if ( - uses_placeholder - and "pixel_values" in message - and "imgs_sizes" not in message - and message["pixel_values"].ndim == 4 - ): - pixel_values = message["pixel_values"] - num_tiles, _, height, width = pixel_values.shape - message["imgs_sizes"] = torch.tensor( - [[height, width]] * num_tiles, dtype=torch.long - ) - - # imgs_sizes is not always declared in model_input_names by bundled image - # processors, so append it explicitly when present. RADIO uses temporal - # patching even for still images and requires one num_frames=1 entry per - # image/tile. - if "imgs_sizes" in message and "imgs_sizes" not in multimodal_keys: - multimodal_keys.append("imgs_sizes") - if "imgs_sizes" in message and "num_frames" not in message: - message["num_frames"] = torch.ones(len(message["imgs_sizes"]), dtype=torch.long) - if "num_frames" in message and "num_frames" not in multimodal_keys: - multimodal_keys.append("num_frames") - for key in multimodal_keys: - if key in message: - user_message[key] = PackedTensor( - message[key], - dim_to_pack=get_dim_to_pack_along(processor, key), - pad_to_max_shape=uses_placeholder and key == "pixel_values", - ) - - # specifically for gemma, we need to add token_type_ids to the user message as a sequence-type value - if "token_type_ids" in message: - user_message["token_type_ids"] = message["token_type_ids"][0] - - # for qwen2.5-vl (transformers>=5.3), mm_token_type_ids tells the model which tokens are text/image/video for 3D RoPE - if "mm_token_type_ids" in message: - user_message["mm_token_type_ids"] = message["mm_token_type_ids"][0] + # Single source of truth for media extraction: the MMPR imgs_sizes + # fallback, RADIO num_frames synthesis, PackedTensor wrapping (incl. the + # imgs_sizes int32 cast) and the gemma / qwen2.5-vl sequence-type maps all + # live in ``extract_multimodal_model_inputs``, which the NeMo-Gym path also + # uses. One implementation is what stops the two paths from handing the + # same model differently-typed inputs. + user_message.update(extract_multimodal_model_inputs(processor, message)) ### append to user message message_log.append(user_message) diff --git a/nemo_rl/distributed/virtual_cluster.py b/nemo_rl/distributed/virtual_cluster.py index d7366e76122..0d2d7455b5c 100644 --- a/nemo_rl/distributed/virtual_cluster.py +++ b/nemo_rl/distributed/virtual_cluster.py @@ -27,6 +27,8 @@ ) from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy +from nemo_rl.utils.venvs import add_hf_modules_cache_to_pythonpath + logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) @@ -322,8 +324,6 @@ def init_ray(log_dir: Optional[str] = None) -> None: # Put Hugging Face's generated ``transformers_modules`` package on the # cluster-wide PYTHONPATH so trust_remote_code objects can be unpickled at # that boundary. This covers both V1 worker groups and direct V2/SC actors. - from nemo_rl.utils.venvs import add_hf_modules_cache_to_pythonpath - env_vars = add_hf_modules_cache_to_pythonpath(dict(os.environ)) env_vars.pop("RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES", None) diff --git a/nemo_rl/experience/payload.py b/nemo_rl/experience/payload.py index e1f4de8431b..1a66cdac0d7 100644 --- a/nemo_rl/experience/payload.py +++ b/nemo_rl/experience/payload.py @@ -211,9 +211,7 @@ def pack_payload( if isinstance(v, torch.Tensor) or (isinstance(v, np.ndarray) and v.dtype == object) } - multimodal = BatchedDataDict[Any](train_batch).get_multimodal_dict( - as_tensors=False - ) + multimodal = BatchedDataDict[Any](train_batch).get_multimodal_dict(as_tensors=False) for key, value in multimodal.items(): wire_value = encode_multimodal_for_wire(key, value) if wire_value is not None: diff --git a/nemo_rl/models/generation/megatron/config.py b/nemo_rl/models/generation/megatron/config.py index f07ca54bc25..eb3a41f0258 100644 --- a/nemo_rl/models/generation/megatron/config.py +++ b/nemo_rl/models/generation/megatron/config.py @@ -44,7 +44,8 @@ class MCoreGenerationSpecificArgs(TypedDict): # - 'block': graphs are owned at the enclosing block (TransformerBlock / HybridBlock). # Only meaningful when cuda_graph_impl='local'. inference_cuda_graph_scope: NotRequired[str] - # Required for EP>1 + local CUDA graphs. + # Required for EP>1 + inference CUDA graphs, except when using the + # `inference_optimized` transformer implementation. moe_pad_experts_for_cuda_graph_inference: NotRequired[bool] materialize_only_last_token_logits: bool @@ -93,8 +94,6 @@ class MCoreGenerationSpecificArgs(TypedDict): # FP8/MXFP8 for the dedicated (non-colocated) inference model; # merged into its `megatron_cfg` by `merged_inference_megatron_cfg`. fp8_cfg: NotRequired[Fp8Config] - # Merged into megatron_cfg for gen workers; required for EP>1 + local CUDA graphs. - moe_pad_experts_for_cuda_graph_inference: NotRequired[bool] class MCoreGenerationConfig(GenerationConfig): diff --git a/nemo_rl/models/value/tq_value.py b/nemo_rl/models/value/tq_value.py index 6582960ec70..9760248c822 100644 --- a/nemo_rl/models/value/tq_value.py +++ b/nemo_rl/models/value/tq_value.py @@ -110,8 +110,6 @@ def get_values_from_meta( timer: Optional timer for nested get_values measurements. """ spa, dba = self._packing_args("logprob_mb_tokens") - # The critic is text-only (built with is_vlm=False), so media columns - # are deliberately not fetched here. value_meta = self._isolated_meta( meta, fields=list(VALUE_SEED_FIELDS), @@ -167,8 +165,6 @@ def train_from_meta( micro_batch_size = mbs or self.cfg["train_micro_batch_size"] spa, dba = self._packing_args("train_mb_tokens") - # The critic is text-only (built with is_vlm=False), so media columns - # are deliberately not fetched here. train_meta = self._isolated_meta( meta, fields=list(DP_VALUE_TRAIN_FIELDS), diff --git a/tests/unit/experience/test_payload.py b/tests/unit/experience/test_payload.py index 04755999404..62a85a992f7 100644 --- a/tests/unit/experience/test_payload.py +++ b/tests/unit/experience/test_payload.py @@ -268,7 +268,7 @@ def test_multimodal_packed_tensor_round_trips_through_tq_payload() -> None: ) assert "pixel_values" in fields assert tags[0]["pixel_values__row_shapes"]["shapes"] == [[2, 4]] - assert "pixel_values__row_shapes" not in tags[1] + assert tags[1]["pixel_values__row_shapes"]["shapes"] == [] restored = materialize(fields, tags=tags) restored_media = restored["pixel_values"] diff --git a/tests/unit/experience/test_rollout_generation_failures.py b/tests/unit/experience/test_rollout_generation_failures.py index f6317d3a8ec..3b191494959 100644 --- a/tests/unit/experience/test_rollout_generation_failures.py +++ b/tests/unit/experience/test_rollout_generation_failures.py @@ -88,6 +88,8 @@ def __init__(self, input_ids: torch.Tensor) -> None: class _FakeTokenizer: + pad_token_id = 0 + def decode(self, ids, skip_special_tokens=True): del skip_special_tokens return f"<{len(ids)} tokens>" diff --git a/tests/unit/experience/test_rollout_manager.py b/tests/unit/experience/test_rollout_manager.py index 75ee79132a0..e344e76192c 100644 --- a/tests/unit/experience/test_rollout_manager.py +++ b/tests/unit/experience/test_rollout_manager.py @@ -102,14 +102,17 @@ class _Generation: async def generate_async(self, data): captured["data"] = data input_len = int(data["input_lengths"][0]) - yield 0, BatchedDataDict( - { - "output_ids": torch.cat( - (data["input_ids"], torch.tensor([[42]])), dim=1 - ), - "unpadded_sequence_lengths": torch.tensor([input_len + 1]), - "logprobs": torch.zeros(1, input_len + 1), - } + yield ( + 0, + BatchedDataDict( + { + "output_ids": torch.cat( + (data["input_ids"], torch.tensor([[42]])), dim=1 + ), + "unpadded_sequence_lengths": torch.tensor([input_len + 1]), + "logprobs": torch.zeros(1, input_len + 1), + } + ), ) manager = object.__new__(AsyncRolloutImpl) diff --git a/tests/unit/experience/test_rollout_manager_router_replay.py b/tests/unit/experience/test_rollout_manager_router_replay.py index 6f5390364bb..ec41573cc37 100644 --- a/tests/unit/experience/test_rollout_manager_router_replay.py +++ b/tests/unit/experience/test_rollout_manager_router_replay.py @@ -38,6 +38,8 @@ def _fallback_routes(count: int) -> torch.Tensor: class _FakeTokenizer: + pad_token_id = 0 + def decode(self, token_ids: torch.Tensor, skip_special_tokens: bool) -> str: del token_ids, skip_special_tokens return "generated" diff --git a/tests/unit/models/policy/test_megatron_worker.py b/tests/unit/models/policy/test_megatron_worker.py index 29eaa03c9a4..78db03db202 100644 --- a/tests/unit/models/policy/test_megatron_worker.py +++ b/tests/unit/models/policy/test_megatron_worker.py @@ -822,31 +822,6 @@ def test_prepare_for_generation_disables_param_gather_hook_before_wake( assert model.config.flash_decode is False -def test_move_retained_vlm_media_preserves_shared_tensors() -> None: - from nemo_rl.models.generation.megatron.megatron_worker import ( - MegatronGenerationMixin, - ) - - shared_imgs = torch.ones(2) - first_request = SimpleNamespace(imgs=shared_imgs) - second_request = SimpleNamespace(imgs=shared_imgs) - text_request = SimpleNamespace() - worker = MegatronGenerationMixin() - worker.dynamic_inference_engine = SimpleNamespace( - requests={ - 1: SimpleNamespace(record=[first_request]), - 2: SimpleNamespace(record=[second_request]), - 3: SimpleNamespace(record=[text_request]), - } - ) - - worker._move_retained_vlm_media(torch.device("meta")) - - assert first_request.imgs.device.type == "meta" - assert second_request.imgs is first_request.imgs - assert not hasattr(text_request, "imgs") - - def create_megatron_test_config( model_name: str, tp: int = 1, diff --git a/tests/unit/single_controller/test_setup.py b/tests/unit/single_controller/test_setup.py index 23401f3139c..e098658f6f3 100644 --- a/tests/unit/single_controller/test_setup.py +++ b/tests/unit/single_controller/test_setup.py @@ -1052,9 +1052,7 @@ def test_env_handles_sourced_from_setup_response_data(self, patched_factories): assert call_kwargs["env_configs"] == {"math": math_env_cfg} assert actor_args.env_handles is patched_factories["env_handles"] - def test_vlm_processor_used_for_data_and_environment_setup( - self, patched_factories - ): + def test_vlm_processor_used_for_data_and_environment_setup(self, patched_factories): mc = _make_master_config(env={"clevr-cogent": {"some": "value"}}) tokenizer = MagicMock(pad_token_id=0) processor = MagicMock(tokenizer=tokenizer) @@ -1064,11 +1062,11 @@ def test_vlm_processor_used_for_data_and_environment_setup( call_args, call_kwargs = patched_factories["setup_response_data"].call_args assert call_args[0] is processor - assert call_kwargs["env_configs"] == { - "clevr-cogent": {"some": "value"} - } + assert call_kwargs["env_configs"] == {"clevr-cogent": {"some": "value"}} assert call_kwargs["is_vlm"] is True - warmup_fields = actor_args.dp_client.register_partition.call_args.kwargs["fields"] + warmup_fields = actor_args.dp_client.register_partition.call_args.kwargs[ + "fields" + ] for field in ( "pixel_values", "image_grid_thw", From 11f2eb749288f6df25a765d556cc9f2cd383c420 Mon Sep 17 00:00:00 2001 From: Cory Ye Date: Sat, 5 Sep 2026 23:02:21 -0700 Subject: [PATCH 10/19] Restore Gym revision required by multimodal branch Signed-off-by: Cory Ye --- 3rdparty/Gym-workspace/Gym | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3rdparty/Gym-workspace/Gym b/3rdparty/Gym-workspace/Gym index fd5e84d6b1c..c3bac96314a 160000 --- a/3rdparty/Gym-workspace/Gym +++ b/3rdparty/Gym-workspace/Gym @@ -1 +1 @@ -Subproject commit fd5e84d6b1c485c80e7ae61553bbd485611c03b4 +Subproject commit c3bac96314a59f28b896f597eb9845d175bb0252 From 91286e14c5d908d00d7bff8d40b0b425309a7c52 Mon Sep 17 00:00:00 2001 From: Cory Ye Date: Sat, 5 Sep 2026 23:08:50 -0700 Subject: [PATCH 11/19] Revert "Restore Gym revision required by multimodal branch" This reverts commit ff519e1646864eba304e28f5109cf9c774b8c2f0. Signed-off-by: Cory Ye --- 3rdparty/Gym-workspace/Gym | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3rdparty/Gym-workspace/Gym b/3rdparty/Gym-workspace/Gym index c3bac96314a..fd5e84d6b1c 160000 --- a/3rdparty/Gym-workspace/Gym +++ b/3rdparty/Gym-workspace/Gym @@ -1 +1 @@ -Subproject commit c3bac96314a59f28b896f597eb9845d175bb0252 +Subproject commit fd5e84d6b1c485c80e7ae61553bbd485611c03b4 From 319aaa1e201831d3a450d89366f50cda13806507 Mon Sep 17 00:00:00 2001 From: Cory Ye Date: Sun, 6 Sep 2026 10:56:53 -0700 Subject: [PATCH 12/19] Address more review comments. Signed-off-by: Cory Ye --- docs/guides/single-controller.md | 3 +++ .../single_controller_utils/setup.py | 26 ++++++++++--------- nemo_rl/data/multimodal_utils.py | 2 +- nemo_rl/data_plane/interfaces.py | 17 +++++++++++- nemo_rl/environments/nemo_gym_multimodal.py | 6 ----- nemo_rl/experience/interfaces.py | 1 + nemo_rl/experience/payload.py | 21 ++++++++++----- nemo_rl/experience/rollout_manager.py | 3 +++ nemo_rl/experience/rollout_reassembler.py | 7 +++-- .../experience/rollout_reassembler_actor.py | 3 +++ ...L1_Functional_Tests_GB200_Megatron_Omni.sh | 18 +++++++++++-- ...s_GB200_Megatron_Omni_Single_Controller.sh | 18 +++++++++++-- tests/unit/data/datasets/test_mmpr_tiny.py | 2 +- .../data_plane/test_rollout_reassembler.py | 3 ++- tests/unit/environments/test_nemo_gym.py | 1 + tests/unit/experience/test_payload.py | 23 +++++++++++++++- tests/unit/experience/test_rollout_manager.py | 20 +++++++++++--- .../test_rollout_reassembler_actor.py | 1 + tests/unit/single_controller/test_setup.py | 21 ++++++++------- 19 files changed, 148 insertions(+), 48 deletions(-) diff --git a/docs/guides/single-controller.md b/docs/guides/single-controller.md index d50bac5d280..391ff704df8 100644 --- a/docs/guides/single-controller.md +++ b/docs/guides/single-controller.md @@ -234,6 +234,9 @@ Do not carry `max_num_epochs: -1` across either. [ppo.md](./ppo.md#asynchronous- The SC path is still under active development. Feature gaps are tracked in [issue #2625](https://github.com/NVIDIA-NeMo/RL/issues/2625). Notable items: +- Multimodal/VLM GRPO is supported with Megatron generation. Set + `policy.is_vlm: true`; see the + [CLEVR Single-Controller recipe](../../examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-single-controller-async.v1.yaml). - Multi-Teacher On-Policy Distillation (MOPD) is supported for text-only NeMo Gym rollouts; multimodal/VLM MOPD is not yet supported. See [Multi-Teacher On-Policy Distillation](../about/algorithms/mopd.md#running-mopd). diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py index 1f5c9c3ae83..6043d07fbe6 100644 --- a/nemo_rl/algorithms/single_controller_utils/setup.py +++ b/nemo_rl/algorithms/single_controller_utils/setup.py @@ -70,10 +70,7 @@ ) from nemo_rl.algorithms.utils import set_seed from nemo_rl.data.collate_fn import rl_collate_fn -from nemo_rl.data.multimodal_utils import ( - PACKED_MULTIMODAL_FIELDS, - PER_TOKEN_MULTIMODAL_FIELDS, -) +from nemo_rl.data.multimodal_utils import WIRE_MULTIMODAL_FIELDS from nemo_rl.data.utils import load_dataloader_state, setup_response_data from nemo_rl.data_plane import ( DATA_PLANE_CHECKPOINT_SCHEMA_VERSION, @@ -1447,9 +1444,7 @@ def _build_generation_then_trainer( if processor is not None: partition_fields.extend( field - for field in sorted( - PACKED_MULTIMODAL_FIELDS | PER_TOKEN_MULTIMODAL_FIELDS - ) + for field in sorted(WIRE_MULTIMODAL_FIELDS) if field not in partition_fields ) dp_client.register_partition( @@ -1479,13 +1474,20 @@ def _build_generation_then_trainer( ) group_size = algo_cfg.num_generations_per_prompt num_rollout_samples = master_config.async_rl.max_buffered_rollouts * group_size + partition_fields = fields_with_optional_routed_experts( + DP_TRAIN_FIELDS, + enabled=r3_enabled + and not token_capture_cfg.defer_routed_experts_to_policy, + ) + if processor is not None: + partition_fields.extend( + field + for field in sorted(WIRE_MULTIMODAL_FIELDS) + if field not in partition_fields + ) dp_client.register_partition( partition_id=partition_id, - fields=fields_with_optional_routed_experts( - DP_TRAIN_FIELDS, - enabled=r3_enabled - and not token_capture_cfg.defer_routed_experts_to_policy, - ), + fields=partition_fields, num_samples=num_rollout_samples, consumer_tasks=["prev_lp", "ref_lp", "train"], grpo_group_size=group_size, diff --git a/nemo_rl/data/multimodal_utils.py b/nemo_rl/data/multimodal_utils.py index c9d3d6b1a1b..4f047bb21e9 100644 --- a/nemo_rl/data/multimodal_utils.py +++ b/nemo_rl/data/multimodal_utils.py @@ -1240,7 +1240,7 @@ def extract_multimodal_model_inputs( f"Processor model input {key!r} must be a torch.Tensor, got " f"{type(value).__name__}." ) - if key == "imgs_sizes": + if key in ("imgs_sizes", "num_frames"): value = value.to(dtype=torch.int32) extracted[key] = PackedTensor( value, diff --git a/nemo_rl/data_plane/interfaces.py b/nemo_rl/data_plane/interfaces.py index 4c33a6d0856..9f527c4b2c8 100644 --- a/nemo_rl/data_plane/interfaces.py +++ b/nemo_rl/data_plane/interfaces.py @@ -312,7 +312,22 @@ def slice(self, start: int, stop: int) -> "KVBatchMeta": ) def concat(self, *others: "KVBatchMeta") -> "KVBatchMeta": - """Append ``others`` and union their fields in first-seen order.""" + """Append metadata from the same partition. + + Sample IDs are concatenated in argument order, while fields are + unioned in first-seen order. Sequence lengths and tags are retained + only when every input provides them. + + Args: + *others: Metadata batches whose ``partition_id`` matches this + batch. + + Returns: + A new metadata batch containing all input rows. + + Raises: + ValueError: If any input has a different ``partition_id``. + """ if any(o.partition_id != self.partition_id for o in others): raise ValueError("KVBatchMeta.concat: partition_ids must match") all_m = (self, *others) diff --git a/nemo_rl/environments/nemo_gym_multimodal.py b/nemo_rl/environments/nemo_gym_multimodal.py index 3e65b88ba69..4987288f576 100644 --- a/nemo_rl/environments/nemo_gym_multimodal.py +++ b/nemo_rl/environments/nemo_gym_multimodal.py @@ -33,7 +33,6 @@ PackedTensor, extract_input_media_sources_from_responses_messages, extract_multimodal_model_inputs, - get_dim_to_pack_along, get_responses_content_part_url, image_to_data_url, media_sources_equal, @@ -824,11 +823,6 @@ def nemo_gym_example_to_video_datum_spec( if "imgs_sizes" in processed and "num_frames" not in processed: processed["num_frames"] = torch.tensor([len(frame_items)], dtype=torch.int32) user_message.update(extract_multimodal_model_inputs(processor, processed)) - if "num_frames" in processed: - user_message["num_frames"] = PackedTensor( - processed["num_frames"].to(dtype=torch.int32), - dim_to_pack=get_dim_to_pack_along(processor, "num_frames"), - ) length = len(user_message["token_ids"]) loss_multiplier = 1.0 diff --git a/nemo_rl/experience/interfaces.py b/nemo_rl/experience/interfaces.py index 9457d98658a..29320e93f2b 100644 --- a/nemo_rl/experience/interfaces.py +++ b/nemo_rl/experience/interfaces.py @@ -61,3 +61,4 @@ class PromptGroupRecord: metadata: dict[str, Any] completions: list["Completion"] rollout_metrics: dict[str, Any] + loss_multiplier: float = 1.0 diff --git a/nemo_rl/experience/payload.py b/nemo_rl/experience/payload.py index 1a66cdac0d7..aa1c17dfe24 100644 --- a/nemo_rl/experience/payload.py +++ b/nemo_rl/experience/payload.py @@ -109,10 +109,12 @@ def record_to_train_batch( flags for configured advantage penalties. Returns: - BatchedDataDict with input_ids, input_lengths, generation_logprobs, - token_mask, an all-ones sample_mask, the raw mask_sample and truncated - flags, prompt_ids_for_adv, total_reward, violation counts, and optional - routed experts and message-violation masks. + BatchedDataDict with input IDs and lengths, generation log probabilities, + token and prompt-level sample masks, raw ``mask_sample`` and ``truncated`` + flags, prompt IDs for advantage computation, rewards, and violation + counts. Optional fields include routed experts, message-violation masks, + and any packed or per-token multimodal model inputs carried by the + completions. """ # Lazy imports: grpo and llm_message_utils transitively pull # experience.rollouts, so importing at module top risks a cycle. @@ -161,7 +163,9 @@ def record_to_train_batch( ) mask_sample = _mask_sample_flags(c.env_extras for c in completions) truncated = torch.tensor([c.truncated for c in completions], dtype=torch.bool) - sample_mask = torch.ones(n, dtype=torch.float32) + sample_mask = torch.full( + (n,), float(record.loss_multiplier), dtype=torch.float32 + ) train_data: dict[str, Any] = { "input_ids": flat["token_ids"], @@ -200,8 +204,11 @@ def pack_payload( prompt_idx: Stable dataset prompt index stamped on every row's tag. Returns: - sample_ids of the form {group_id}_g{i}, a jagged-packed TensorDict, and per-row - tags carrying weight_version plus any per-row violation counts. + Sample IDs of the form ``{group_id}_g{i}``, a jagged-packed TensorDict + containing tensor fields and encoded multimodal wire fields, and + per-row tags. Tags carry the weight version, prompt index, violation + counts, and ``__row_shapes`` metadata required to reconstruct + packed multimodal rows. """ lengths = train_batch["input_lengths"] n = int(lengths.shape[0]) diff --git a/nemo_rl/experience/rollout_manager.py b/nemo_rl/experience/rollout_manager.py index 501c85f3a32..e5b5286fbab 100644 --- a/nemo_rl/experience/rollout_manager.py +++ b/nemo_rl/experience/rollout_manager.py @@ -455,6 +455,7 @@ async def run_rollout( metadata={"task_name": input_sample["task_name"]}, completions=completions, rollout_metrics=rollout_metrics, + loss_multiplier=float(input_sample.get("loss_multiplier", 1.0)), ) async def _run_single_rollout( @@ -886,6 +887,7 @@ async def run_rollout( metadata={"task_name": "nemo_gym"}, completions=completions, rollout_metrics=rollout_metrics, + loss_multiplier=float(input_sample.get("loss_multiplier", 1.0)), ) def _validate_init_params(self) -> None: @@ -1935,6 +1937,7 @@ async def _generate_for_finalization_attempt( fallback_weight_version=start_version, prompt_idx=record.prompt_idx, mask_sample=mask_sample, + loss_multiplier=record.loss_multiplier, ) from nemo_rl.experience.rollout_reassembler_actor import ( assert_metadata_only, diff --git a/nemo_rl/experience/rollout_reassembler.py b/nemo_rl/experience/rollout_reassembler.py index 551586d8c70..2b5818c8a2d 100644 --- a/nemo_rl/experience/rollout_reassembler.py +++ b/nemo_rl/experience/rollout_reassembler.py @@ -364,6 +364,7 @@ def finalize_group( mask_sample: list[bool], fallback_weight_version: int, prompt_idx: int, + loss_multiplier: float = 1.0, ) -> FinalizedGroup: """Publish exactly N canonical rows for one prompt group. @@ -374,7 +375,9 @@ def finalize_group( advantage-stage flag the native ``pack_payload`` path emits from each ``Completion``; it rides along unchanged so the train pump's environment masking reads the same field on both paths (placeholder - rows already train nothing through ``sample_mask`` 0). ``truncated`` + rows already train nothing through ``sample_mask`` 0). + ``loss_multiplier`` supplies the dataset-level weight for every valid + row, matching the ordinary ``record_to_train_batch`` path. ``truncated`` is not carried from the dispatcher -- the receipt path has no real tokens to measure it from at dispatch time -- so it is computed here instead, from each row's rebuilt length against ``max_seq_len``. @@ -511,7 +514,7 @@ def finalize_group( input_ids[i, :length] = torch.tensor(row.token_ids, dtype=torch.int64) token_mask[i, :length] = torch.tensor(row.token_mask, dtype=torch.float32) logprobs[i, :length] = torch.tensor(row.logprobs, dtype=torch.float32) - sample_mask[i] = 1.0 + sample_mask[i] = float(loss_multiplier) train_batch = { "input_ids": input_ids, diff --git a/nemo_rl/experience/rollout_reassembler_actor.py b/nemo_rl/experience/rollout_reassembler_actor.py index 61505f49d08..57d91953d29 100644 --- a/nemo_rl/experience/rollout_reassembler_actor.py +++ b/nemo_rl/experience/rollout_reassembler_actor.py @@ -63,6 +63,8 @@ class ReassemblyRequest: # train pump reads the same ``mask_sample`` field as the native path # (SingleController reads it unconditionally). mask_sample: tuple[bool, ...] + # Dataset-level loss weight shared by every completion in this prompt group. + loss_multiplier: float = 1.0 @dataclass(frozen=True) @@ -152,6 +154,7 @@ def finalize(self, request: ReassemblyRequest) -> FinalizedGroup: mask_sample=list(request.mask_sample), fallback_weight_version=request.fallback_weight_version, prompt_idx=request.prompt_idx, + loss_multiplier=request.loss_multiplier, ) assert_metadata_only(result) return result diff --git a/tests/functional/L1_Functional_Tests_GB200_Megatron_Omni.sh b/tests/functional/L1_Functional_Tests_GB200_Megatron_Omni.sh index c8f916ef053..023f0257752 100755 --- a/tests/functional/L1_Functional_Tests_GB200_Megatron_Omni.sh +++ b/tests/functional/L1_Functional_Tests_GB200_Megatron_Omni.sh @@ -20,6 +20,20 @@ PROJECT_ROOT=$(realpath "${SCRIPT_DIR}/../..") cd "${PROJECT_ROOT}" +# run_test [fast] +# - "run_test fast " = always runs (both fast and full modes) +# - "run_test " = only runs in full mode; skipped when FAST=1 +run_test() { + if [[ "$1" == "fast" ]]; then + shift + time "$@" + elif [[ "${FAST:-0}" == "1" ]]; then + echo "FAST: Skipping: $*" + else + time "$@" + fi +} + GPU_COUNT=$(nvidia-smi --query-gpu=index --format=csv,noheader | wc -l) if (( GPU_COUNT < 2 )); then echo "SKIP: Nemotron Omni functional tests require at least two GB200 GPUs" @@ -29,8 +43,8 @@ fi # Both tests colocate TP2/EP2 training and generation on two GB200 GPUs. export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1}" -time uv run --no-sync bash ./tests/functional/nemotron_omni_clevr_megatron_1n2g.sh -time uv run --no-sync bash ./tests/functional/nemotron_omni_gym_video_megatron_1n2g.sh +run_test fast uv run --no-sync bash ./tests/functional/nemotron_omni_clevr_megatron_1n2g.sh +run_test fast uv run --no-sync bash ./tests/functional/nemotron_omni_gym_video_megatron_1n2g.sh cd "${PROJECT_ROOT}/tests" if compgen -G ".coverage*" > /dev/null; then diff --git a/tests/functional/L1_Functional_Tests_GB200_Megatron_Omni_Single_Controller.sh b/tests/functional/L1_Functional_Tests_GB200_Megatron_Omni_Single_Controller.sh index ca6ec20bf28..36f15552ddd 100755 --- a/tests/functional/L1_Functional_Tests_GB200_Megatron_Omni_Single_Controller.sh +++ b/tests/functional/L1_Functional_Tests_GB200_Megatron_Omni_Single_Controller.sh @@ -20,6 +20,20 @@ PROJECT_ROOT=$(realpath "${SCRIPT_DIR}/../..") cd "${PROJECT_ROOT}" +# run_test [fast] +# - "run_test fast " = always runs (both fast and full modes) +# - "run_test " = only runs in full mode; skipped when FAST=1 +run_test() { + if [[ "$1" == "fast" ]]; then + shift + time "$@" + elif [[ "${FAST:-0}" == "1" ]]; then + echo "FAST: Skipping: $*" + else + time "$@" + fi +} + GPU_COUNT=$(nvidia-smi --query-gpu=index --format=csv,noheader | wc -l) if (( GPU_COUNT < 2 )); then echo "SKIP: Nemotron Omni SingleController functional tests require at least two GB200 GPUs" @@ -30,8 +44,8 @@ fi # and one GPU hosts Megatron generation. export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1}" -time uv run --no-sync bash ./tests/functional/nemotron_omni_clevr_megatron_single_controller_1n2g.sh -time uv run --no-sync bash ./tests/functional/nemotron_omni_gym_video_megatron_single_controller_1n2g.sh +run_test fast uv run --no-sync bash ./tests/functional/nemotron_omni_clevr_megatron_single_controller_1n2g.sh +run_test fast uv run --no-sync bash ./tests/functional/nemotron_omni_gym_video_megatron_single_controller_1n2g.sh cd "${PROJECT_ROOT}/tests" if compgen -G ".coverage*" > /dev/null; then diff --git a/tests/unit/data/datasets/test_mmpr_tiny.py b/tests/unit/data/datasets/test_mmpr_tiny.py index 7ec4c34a538..d4e0da4b969 100644 --- a/tests/unit/data/datasets/test_mmpr_tiny.py +++ b/tests/unit/data/datasets/test_mmpr_tiny.py @@ -307,7 +307,7 @@ def test_historical_tiled_processor_gets_media_metadata(self, tiny_image_path): torch.tensor([[224, 224], [224, 224], [224, 224]]), ) assert torch.equal( - user_message["num_frames"].as_tensor(), torch.ones(3, dtype=torch.long) + user_message["num_frames"].as_tensor(), torch.ones(3, dtype=torch.int32) ) def test_prompted_text_contains_boxed_literal_and_no_raw_dataset_string( diff --git a/tests/unit/data_plane/test_rollout_reassembler.py b/tests/unit/data_plane/test_rollout_reassembler.py index bd6b965c2f9..27816aad807 100644 --- a/tests/unit/data_plane/test_rollout_reassembler.py +++ b/tests/unit/data_plane/test_rollout_reassembler.py @@ -215,6 +215,7 @@ def test_finalize_group_publishes_n_rows_with_placeholder(tq_client, partitions) mask_sample=[True, False], fallback_weight_version=9, prompt_idx=0, + loss_multiplier=0.25, ) assert not finalized.dropped assert finalized.meta is not None @@ -229,7 +230,7 @@ def test_finalize_group_publishes_n_rows_with_placeholder(tq_client, partitions) rows = _fetch_rows(tq_client, rollout_ids) sample_mask = torch.as_tensor(rows["sample_mask"]).flatten() - assert sample_mask.tolist() == [1.0, 0.0] + assert sample_mask.tolist() == [0.25, 0.0] input_ids = torch.as_tensor(rows["input_ids"][0]).flatten() assert input_ids[:valid_len].tolist() == expected.token_ids # Placeholder borrows the valid sibling's prompt for baseline grouping. diff --git a/tests/unit/environments/test_nemo_gym.py b/tests/unit/environments/test_nemo_gym.py index 51cb57c6f0f..0c57afd293f 100644 --- a/tests/unit/environments/test_nemo_gym.py +++ b/tests/unit/environments/test_nemo_gym.py @@ -507,6 +507,7 @@ def apply_chat_template(self, messages, *, tokenize, **kwargs): assert datum is not None user_message = datum["message_log"][0] assert user_message["num_frames"].as_tensor().tolist() == [4] + assert user_message["num_frames"].as_tensor().dtype == torch.int32 assert user_message["imgs_sizes"].as_tensor().dtype == torch.int32 extra_env_info = datum["extra_env_info"] outbound_content = extra_env_info["responses_create_params"]["input"][0]["content"] diff --git a/tests/unit/experience/test_payload.py b/tests/unit/experience/test_payload.py index 62a85a992f7..df879700e8b 100644 --- a/tests/unit/experience/test_payload.py +++ b/tests/unit/experience/test_payload.py @@ -84,7 +84,9 @@ def _completion( ) -def _record(completions: list[Completion]) -> PromptGroupRecord: +def _record( + completions: list[Completion], *, loss_multiplier: float = 1.0 +) -> PromptGroupRecord: return PromptGroupRecord( prompt_idx=0, prompt=[ @@ -98,6 +100,7 @@ def _record(completions: list[Completion]) -> PromptGroupRecord: metadata={"task_name": "test"}, completions=completions, rollout_metrics={}, + loss_multiplier=loss_multiplier, ) @@ -322,6 +325,24 @@ def test_record_to_train_batch_carries_raw_masks_without_applying_them() -> None assert torch.equal(fields["truncated"], train_batch["truncated"]) +def test_record_to_train_batch_broadcasts_prompt_loss_multiplier() -> None: + record = _record( + [ + _completion(route_start=10, reward=1.0), + _completion(route_start=30, reward=2.0), + ], + loss_multiplier=0.0, + ) + + train_batch = record_to_train_batch( + record, + pad_value_dict={"token_ids": 0, "input_ids": 0}, + include_message_violation_fields=False, + ) + + assert torch.equal(train_batch["sample_mask"], torch.zeros(2)) + + def _failed_completion() -> Completion: """A trajectory whose first generation raised: prompt only, no routes.""" return Completion( diff --git a/tests/unit/experience/test_rollout_manager.py b/tests/unit/experience/test_rollout_manager.py index e344e76192c..ae43d0c973a 100644 --- a/tests/unit/experience/test_rollout_manager.py +++ b/tests/unit/experience/test_rollout_manager.py @@ -1020,7 +1020,10 @@ def test_async_rollout_manager( - completions hold independent (not aliased) message_log objects """ vllm_generation, tokenizer, task_to_env, _, _ = multi_step_setup_vllm_async - input_sample = single_multi_step_calculator_input_sample + input_sample = { + **single_multi_step_calculator_input_sample, + "loss_multiplier": 0.25, + } num_generations = 2 max_seq_len = 1024 max_rollout_turns = input_sample["extra_env_info"]["max_steps"] + 1 @@ -1044,6 +1047,7 @@ def test_async_rollout_manager( f"Expected {num_generations} completions, got {len(record.completions)}" ) assert record.prompt_idx == input_sample["idx"] + assert record.loss_multiplier == input_sample["loss_multiplier"] for i, completion in enumerate(record.completions): assert isinstance(completion, Completion) @@ -1298,6 +1302,7 @@ def test_async_nemo_gym_rollout_manager( f"Expected {num_generations} completions, got {len(record.completions)}" ) assert record.prompt_idx == 0 + assert record.loss_multiplier == single_prompt["loss_multiplier"] for i, completion in enumerate(record.completions): assert isinstance(completion, Completion) @@ -1499,7 +1504,9 @@ def reserve( ) -def _receipt_record(rollout_ids, receipts, instance_configs=None): +def _receipt_record( + rollout_ids, receipts, instance_configs=None, *, loss_multiplier=1.0 +): instance_configs = instance_configs or [None] * len(rollout_ids) completions = [ Completion( @@ -1522,6 +1529,7 @@ def _receipt_record(rollout_ids, receipts, instance_configs=None): metadata={"task_name": "nemo_gym"}, completions=completions, rollout_metrics={}, + loss_multiplier=loss_multiplier, ) @@ -1560,6 +1568,7 @@ async def run_rollout(self, _sample, *, rollout_ids=None): rollout_ids, [{"rollout_id": rid} for rid in rollout_ids], instance_configs=instance_configs, + loss_multiplier=float(_sample.get("loss_multiplier", 1.0)), ) mgr._impl = _CaptureImpl() @@ -1585,7 +1594,11 @@ def test_mints_ids_and_returns_metadata_request(self): buf = _FakeCaptureBuffer() mgr = _make_capture_manager(buf) - request = _run(mgr.generate_for_finalization({"prompt": "p"}, target_step=5)) + request = _run( + mgr.generate_for_finalization( + {"prompt": "p", "loss_multiplier": 0.25}, target_step=5 + ) + ) # Rollout ids were minted from the reserved group id and threaded # end to end: reserve -> impl -> metadata-only actor request. @@ -1598,6 +1611,7 @@ def test_mints_ids_and_returns_metadata_request(self): assert [r["rollout_id"] for r in request.receipts] == expected_ids assert request.rewards == (0.5, 0.5) assert request.mask_sample == (False, False) + assert request.loss_multiplier == 0.25 assert request.fallback_weight_version == 7 # Finalization and commit are exclusively owned by the controller's # actor-pool path; the manager leaves the reservation unready. diff --git a/tests/unit/experience/test_rollout_reassembler_actor.py b/tests/unit/experience/test_rollout_reassembler_actor.py index e73e6d14eb9..f84da80170b 100644 --- a/tests/unit/experience/test_rollout_reassembler_actor.py +++ b/tests/unit/experience/test_rollout_reassembler_actor.py @@ -100,6 +100,7 @@ def test_rpc_dataclass_fields_are_classified() -> None: "fallback_weight_version", "prompt_idx", "mask_sample", + "loss_multiplier", } assert {f.name for f in fields(FinalizedGroup)} == { "meta", diff --git a/tests/unit/single_controller/test_setup.py b/tests/unit/single_controller/test_setup.py index e098658f6f3..eb93a0086cf 100644 --- a/tests/unit/single_controller/test_setup.py +++ b/tests/unit/single_controller/test_setup.py @@ -59,6 +59,7 @@ from nemo_rl.algorithms.single_controller_utils.config import ( validate_single_controller_config, ) +from nemo_rl.data.multimodal_utils import WIRE_MULTIMODAL_FIELDS from nemo_rl.data_plane import DATA_PLANE_CHECKPOINT_SCHEMA_VERSION from nemo_rl.data_plane.schema import SC_ROLLOUT_SCHEMA_FIELDS from nemo_rl.experience.rollouts import EffortLevelsConfig @@ -1067,14 +1068,7 @@ def test_vlm_processor_used_for_data_and_environment_setup(self, patched_factori warmup_fields = actor_args.dp_client.register_partition.call_args.kwargs[ "fields" ] - for field in ( - "pixel_values", - "image_grid_thw", - "imgs_sizes", - "num_frames", - "mm_token_type_ids", - ): - assert field in warmup_fields + assert WIRE_MULTIMODAL_FIELDS <= set(warmup_fields) def test_weight_sync_factory_args(self, patched_factories): """create_weight_synchronizer receives policy / generation / topology.""" @@ -1227,6 +1221,8 @@ def test_token_capture_always_creates_finalizer_actor_pool(self, patched_factori None, ) fake_actors = [MagicMock(name=f"finalizer_{index}") for index in range(3)] + tokenizer = MagicMock(pad_token_id=9) + processor = MagicMock(tokenizer=tokenizer) with ( patch.object(sc_setup_mod, "should_use_nemo_gym", return_value=True), @@ -1239,7 +1235,9 @@ def test_token_capture_always_creates_finalizer_actor_pool(self, patched_factori return_value=fake_actors, ) as mock_create_finalizer_actors, ): - actor_args, _ = setup_single_controller(mc, MagicMock(pad_token_id=9)) + actor_args, _ = setup_single_controller( + mc, tokenizer, processor=processor + ) (actor_dp_config, actor_config), actor_kwargs = ( mock_create_finalizer_actors.call_args @@ -1251,6 +1249,11 @@ def test_token_capture_always_creates_finalizer_actor_pool(self, patched_factori assert actor_kwargs == {"num_workers": 3} assert actor_args.finalizer_actors == fake_actors assert not hasattr(actor_args.rollout_manager, "_finalizer") + partition_calls = actor_args.dp_client.register_partition.call_args_list + assert WIRE_MULTIMODAL_FIELDS <= set(partition_calls[0].kwargs["fields"]) + assert WIRE_MULTIMODAL_FIELDS.isdisjoint( + partition_calls[1].kwargs["fields"] + ) def test_setup_timing_populated_for_noncolocated_vllm(self, patched_factories): """Non-colocated vLLM records every per-phase field.""" From f87affb49c57a3fceaf7ab87e7306893036c313d Mon Sep 17 00:00:00 2001 From: Cory Ye Date: Sun, 6 Sep 2026 18:33:33 -0700 Subject: [PATCH 13/19] Strengthen reserved port handoff test Signed-off-by: Cory Ye --- .../unit/models/generation/test_megatron_generation_parse.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/unit/models/generation/test_megatron_generation_parse.py b/tests/unit/models/generation/test_megatron_generation_parse.py index e177cf74914..40114f75095 100644 --- a/tests/unit/models/generation/test_megatron_generation_parse.py +++ b/tests/unit/models/generation/test_megatron_generation_parse.py @@ -236,6 +236,11 @@ def test_http_server_port_reservation(monkeypatch): assert holder._sock.fileno() == -1 assert reserved_socket.getsockname()[1] == port + # Still accepting after the holder closed its copy: the port was + # never released across the handoff. + with socket.create_connection(("127.0.0.1", port), timeout=5): + pass + # MCore closes the handed-off fd and gives every frontend replica # its own SO_REUSEPORT listener. Such a listener can join the reuse # group while this test stub still holds the adopted duplicate. From 56acacf0fbb669c685edca84c75a7156e3c0f9c2 Mon Sep 17 00:00:00 2001 From: Cory Ye Date: Sun, 6 Sep 2026 19:09:10 -0700 Subject: [PATCH 14/19] Lint. Signed-off-by: Cory Ye --- nemo_rl/algorithms/single_controller_utils/setup.py | 3 +-- nemo_rl/experience/payload.py | 4 +--- tests/unit/single_controller/test_setup.py | 8 ++------ 3 files changed, 4 insertions(+), 11 deletions(-) diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py index 6043d07fbe6..1893a115216 100644 --- a/nemo_rl/algorithms/single_controller_utils/setup.py +++ b/nemo_rl/algorithms/single_controller_utils/setup.py @@ -1476,8 +1476,7 @@ def _build_generation_then_trainer( num_rollout_samples = master_config.async_rl.max_buffered_rollouts * group_size partition_fields = fields_with_optional_routed_experts( DP_TRAIN_FIELDS, - enabled=r3_enabled - and not token_capture_cfg.defer_routed_experts_to_policy, + enabled=r3_enabled and not token_capture_cfg.defer_routed_experts_to_policy, ) if processor is not None: partition_fields.extend( diff --git a/nemo_rl/experience/payload.py b/nemo_rl/experience/payload.py index aa1c17dfe24..c6de8933b86 100644 --- a/nemo_rl/experience/payload.py +++ b/nemo_rl/experience/payload.py @@ -163,9 +163,7 @@ def record_to_train_batch( ) mask_sample = _mask_sample_flags(c.env_extras for c in completions) truncated = torch.tensor([c.truncated for c in completions], dtype=torch.bool) - sample_mask = torch.full( - (n,), float(record.loss_multiplier), dtype=torch.float32 - ) + sample_mask = torch.full((n,), float(record.loss_multiplier), dtype=torch.float32) train_data: dict[str, Any] = { "input_ids": flat["token_ids"], diff --git a/tests/unit/single_controller/test_setup.py b/tests/unit/single_controller/test_setup.py index eb93a0086cf..0aef3a117bc 100644 --- a/tests/unit/single_controller/test_setup.py +++ b/tests/unit/single_controller/test_setup.py @@ -1235,9 +1235,7 @@ def test_token_capture_always_creates_finalizer_actor_pool(self, patched_factori return_value=fake_actors, ) as mock_create_finalizer_actors, ): - actor_args, _ = setup_single_controller( - mc, tokenizer, processor=processor - ) + actor_args, _ = setup_single_controller(mc, tokenizer, processor=processor) (actor_dp_config, actor_config), actor_kwargs = ( mock_create_finalizer_actors.call_args @@ -1251,9 +1249,7 @@ def test_token_capture_always_creates_finalizer_actor_pool(self, patched_factori assert not hasattr(actor_args.rollout_manager, "_finalizer") partition_calls = actor_args.dp_client.register_partition.call_args_list assert WIRE_MULTIMODAL_FIELDS <= set(partition_calls[0].kwargs["fields"]) - assert WIRE_MULTIMODAL_FIELDS.isdisjoint( - partition_calls[1].kwargs["fields"] - ) + assert WIRE_MULTIMODAL_FIELDS.isdisjoint(partition_calls[1].kwargs["fields"]) def test_setup_timing_populated_for_noncolocated_vllm(self, patched_factories): """Non-colocated vLLM records every per-phase field.""" From 81c99e80c09f0e40b18ff7b006a37faa6cbd8486 Mon Sep 17 00:00:00 2001 From: Cory Ye Date: Sun, 6 Sep 2026 19:58:46 -0700 Subject: [PATCH 15/19] Fix rebase error. Signed-off-by: Cory Ye --- nemo_rl/models/policy/workers/megatron_policy_worker.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index c350c3b2c10..6a83b4036dc 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -1544,10 +1544,6 @@ def _train_microbatch_body( # call carries one DP slice; the iterator subdivides into pipeline # microbatches. attach_media_token_validity_mask(data, self.media_placeholder_token_id) -<<<<<<< HEAD -======= - ->>>>>>> f1369705e (Fix rebase errors to reenable V2 multimodal with MCore.) ( data_iterator, num_microbatches, From 2b064d20c15fcf57c03d1d0fba6670f79277d1bb Mon Sep 17 00:00:00 2001 From: Cory Ye Date: Sun, 6 Sep 2026 23:53:40 -0700 Subject: [PATCH 16/19] Fix flaky unit tests and test coverage. Signed-off-by: Cory Ye --- tests/unit/experience/test_payload.py | 37 ++++++++++++++++++- tests/unit/experience/test_rollout_manager.py | 17 ++++++++- .../test_rollout_reassembler_actor.py | 32 +++++++++++++++- .../models/value/test_dtensor_value_worker.py | 4 +- tests/unit/single_controller/test_setup.py | 29 +++++++++++++-- .../single_controller/test_watchdog_pump.py | 2 +- 6 files changed, 110 insertions(+), 11 deletions(-) diff --git a/tests/unit/experience/test_payload.py b/tests/unit/experience/test_payload.py index df879700e8b..0ba33dd4325 100644 --- a/tests/unit/experience/test_payload.py +++ b/tests/unit/experience/test_payload.py @@ -281,6 +281,30 @@ def test_multimodal_packed_tensor_round_trips_through_tq_payload() -> None: assert torch.equal(restored_media.as_tensor(), media) +def test_per_token_multimodal_field_is_packed_with_sequence_lengths() -> None: + train_batch = { + "input_lengths": torch.tensor([3, 2], dtype=torch.int32), + "input_ids": torch.tensor([[10, 11, 12], [20, 21, 0]]), + "token_type_ids": torch.tensor([[0, 1, 1], [0, 1, 0]]), + } + + _, fields, tags = pack_payload( + train_batch, + weight_version=3, + group_id="group", + prompt_idx=17, + ) + + assert [row.tolist() for row in fields["token_type_ids"].unbind()] == [ + [0, 1, 1], + [0, 1], + ] + assert tags == [ + {"weight_version": 3, "prompt_idx": 17}, + {"weight_version": 3, "prompt_idx": 17}, + ] + + def test_record_to_train_batch_carries_raw_masks_without_applying_them() -> None: record = _record( [ @@ -331,7 +355,7 @@ def test_record_to_train_batch_broadcasts_prompt_loss_multiplier() -> None: _completion(route_start=10, reward=1.0), _completion(route_start=30, reward=2.0), ], - loss_multiplier=0.0, + loss_multiplier=0.25, ) train_batch = record_to_train_batch( @@ -340,7 +364,16 @@ def test_record_to_train_batch_broadcasts_prompt_loss_multiplier() -> None: include_message_violation_fields=False, ) - assert torch.equal(train_batch["sample_mask"], torch.zeros(2)) + expected = torch.full((2,), 0.25) + assert torch.equal(train_batch["sample_mask"], expected) + + _, fields, _ = pack_payload( + train_batch, + weight_version=3, + group_id="group", + prompt_idx=17, + ) + assert torch.equal(fields["sample_mask"], expected) def _failed_completion() -> Completion: diff --git a/tests/unit/experience/test_rollout_manager.py b/tests/unit/experience/test_rollout_manager.py index ae43d0c973a..ea7f78c9e29 100644 --- a/tests/unit/experience/test_rollout_manager.py +++ b/tests/unit/experience/test_rollout_manager.py @@ -131,12 +131,22 @@ async def generate_async(self, data): "token_ids": torch.tensor([1, 2, 3]), "pixel_values": pixel_values, "imgs_sizes": imgs_sizes, - } + }, + { + "role": "assistant", + "content": "follow-up", + "token_ids": torch.tensor([4, 5]), + }, ] - _run(manager._generate_response(message_log, None)) + assistant_message, input_lengths, _ = _run( + manager._generate_response(message_log, [""]) + ) generation_data = captured["data"] + assert generation_data["input_ids"].tolist() == [[1, 2, 3, 4, 5]] + assert generation_data["input_lengths"].tolist() == [5] + assert generation_data["stop_strings"] == [[""]] assert isinstance(generation_data["pixel_values"], PackedTensor) assert isinstance(generation_data["imgs_sizes"], PackedTensor) assert torch.equal( @@ -145,6 +155,9 @@ async def generate_async(self, data): assert torch.equal( generation_data["imgs_sizes"].as_tensor(), imgs_sizes.as_tensor() ) + assert input_lengths.tolist() == [5] + assert assistant_message["content"] == "answer" + assert assistant_message["token_ids"].tolist() == [42] class _FakeBuffer: diff --git a/tests/unit/experience/test_rollout_reassembler_actor.py b/tests/unit/experience/test_rollout_reassembler_actor.py index f84da80170b..4a21747e228 100644 --- a/tests/unit/experience/test_rollout_reassembler_actor.py +++ b/tests/unit/experience/test_rollout_reassembler_actor.py @@ -15,7 +15,8 @@ from __future__ import annotations -from dataclasses import fields +from dataclasses import fields, replace +from unittest.mock import MagicMock import pytest import torch @@ -25,6 +26,7 @@ from nemo_rl.experience.rollout_reassembler_actor import ( _FORBIDDEN_RPC_KEYS, ReassemblyRequest, + RolloutReassemblerActor, assert_metadata_only, ) @@ -71,6 +73,34 @@ def test_finalizer_request_and_result_are_metadata_only() -> None: assert_metadata_only(result) +def test_finalize_forwards_loss_multiplier_to_reassembler() -> None: + actor_cls = RolloutReassemblerActor.__ray_metadata__.modified_class + actor = object.__new__(actor_cls) + actor._finalizer = MagicMock() + result = FinalizedGroup( + meta=None, + group_min_wv=4, + group_max_wv=4, + staging_keys=[], + dropped=True, + drop_reason="test", + ) + actor._finalizer.finalize_group.return_value = result + request = replace(_request(), loss_multiplier=0.25) + + assert actor.finalize(request) is result + actor._finalizer.finalize_group.assert_called_once_with( + "group", + ["group_g0"], + [request.receipts[0]], + [1.0], + mask_sample=[False], + fallback_weight_version=4, + prompt_idx=0, + loss_multiplier=0.25, + ) + + @pytest.mark.parametrize( "payload", [ diff --git a/tests/unit/models/value/test_dtensor_value_worker.py b/tests/unit/models/value/test_dtensor_value_worker.py index 5ca4129652f..5ab1c429d57 100644 --- a/tests/unit/models/value/test_dtensor_value_worker.py +++ b/tests/unit/models/value/test_dtensor_value_worker.py @@ -546,7 +546,9 @@ def test_value_worker_train_decreases_loss(value_setup): losses.append(float(loss_tensor.mean().item())) value.finish_training() - assert losses[-1] <= losses[0] + 1e-3, ( + # This is a small fixed-batch smoke, not a convergence test. Allow minor + # optimizer/model-version jitter while still catching a meaningful loss jump. + assert losses[-1] <= losses[0] + 2e-3, ( f"Value loss should not increase after 3 steps; got {losses}" ) diff --git a/tests/unit/single_controller/test_setup.py b/tests/unit/single_controller/test_setup.py index 0aef3a117bc..25d695f5ad0 100644 --- a/tests/unit/single_controller/test_setup.py +++ b/tests/unit/single_controller/test_setup.py @@ -401,8 +401,11 @@ def __init__(self, **kwargs): assert teacher_topology is None -def test_single_controller_mopd_recipe_resolves_to_runtime_contract(): +def test_single_controller_mopd_recipe_resolves_to_runtime_contract(monkeypatch): """The inherited recipe resolves exactly as the SC entrypoint consumes it.""" + # The parent recipe locates its fixture data below HF_HOME. This test only + # validates config resolution, so it needs a stable path, not real data. + monkeypatch.setenv("HF_HOME", "/tmp/nemo-rl-test-hf") register_omegaconf_resolvers() repo_root = Path(__file__).resolve().parents[3] recipe = repo_root / ( @@ -1046,11 +1049,14 @@ def test_env_handles_sourced_from_setup_response_data(self, patched_factories): """setup_response_data receives master_config.env and supplies env handles.""" math_env_cfg = {"some": "value"} mc = _make_master_config(env={"math": math_env_cfg}) + tokenizer = MagicMock(pad_token_id=0) - actor_args, _ = setup_single_controller(mc, MagicMock(pad_token_id=0)) + actor_args, _ = setup_single_controller(mc, tokenizer) - _, call_kwargs = patched_factories["setup_response_data"].call_args + call_args, call_kwargs = patched_factories["setup_response_data"].call_args + assert call_args[0] is tokenizer assert call_kwargs["env_configs"] == {"math": math_env_cfg} + assert call_kwargs["is_vlm"] is False assert actor_args.env_handles is patched_factories["env_handles"] def test_vlm_processor_used_for_data_and_environment_setup(self, patched_factories): @@ -1073,6 +1079,7 @@ def test_vlm_processor_used_for_data_and_environment_setup(self, patched_factori def test_weight_sync_factory_args(self, patched_factories): """create_weight_synchronizer receives policy / generation / topology.""" mc = _make_master_config(colocated=False, backend="vllm") + mc.async_rl.generation_fleet_health.refit_timeout_s = 42.0 tokenizer = MagicMock(pad_token_id=0) setup_single_controller(mc, tokenizer) @@ -1082,6 +1089,11 @@ def test_weight_sync_factory_args(self, patched_factories): assert factory_kwargs["generation"] is patched_factories["fake_gen"] assert factory_kwargs["generation_backend"] == "vllm" assert factory_kwargs["colocated"] is False + assert factory_kwargs["refit_timeout_s"] == 42.0 + assert ( + patched_factories["fake_gen"].weight_synchronizer + is patched_factories["create_weight_synchronizer"].return_value + ) def test_custom_partition_id(self, patched_factories): mc = _make_master_config() @@ -1185,8 +1197,13 @@ def test_nemo_gym_wires_env_handle(self, patched_factories): patch.object(sc_setup_mod, "router_replay_enabled", return_value=False), ): tokenizer = MagicMock(pad_token_id=0) - actor_args, _ = setup_single_controller(mc, tokenizer) + processor = MagicMock(tokenizer=tokenizer) + actor_args, _ = setup_single_controller(mc, tokenizer, processor=processor) + data_args, data_kwargs = patched_factories["setup_response_data"].call_args + assert data_args[0] is processor + assert data_kwargs["env_configs"] is None + assert data_kwargs["is_vlm"] is True mock_spinup.assert_called_once_with( env_configs=mc.env, base_urls=patched_factories["fake_gen"].dp_openai_server_base_urls, @@ -1199,6 +1216,10 @@ def test_nemo_gym_wires_env_handle(self, patched_factories): token_capture=None, ) assert actor_args.env_handles["nemo_gym"] is fake_gym_actor + warmup_fields = actor_args.dp_client.register_partition.call_args.kwargs[ + "fields" + ] + assert WIRE_MULTIMODAL_FIELDS <= set(warmup_fields) def test_token_capture_always_creates_finalizer_actor_pool(self, patched_factories): mc = _make_master_config(backend="vllm") diff --git a/tests/unit/single_controller/test_watchdog_pump.py b/tests/unit/single_controller/test_watchdog_pump.py index 0ae1f0ce77d..b15659508de 100644 --- a/tests/unit/single_controller/test_watchdog_pump.py +++ b/tests/unit/single_controller/test_watchdog_pump.py @@ -339,7 +339,7 @@ def test_a_confirmed_death_stands_the_trainers_deadline_down(self): shard_count=2, policy=FleetHealthPolicy(unhealthy_threshold=99) ) ctrl = self._with_fleet(monitor, worker_alive=[True, False]) - asyncio.run(_run_probe_ticks(ctrl, 1)) + asyncio.run(ctrl._probe_generation_fleet()) assert ctrl._stood_down == [0, 1], ( "every policy worker must be told to stand its refit deadline down once a " f"generation shard is confirmed gone; saw {ctrl._stood_down}" From 32fbaefbd4d0192617b7f92e9dc208199259ec30 Mon Sep 17 00:00:00 2001 From: Cory Ye Date: Mon, 7 Sep 2026 07:59:57 -0700 Subject: [PATCH 17/19] Add a digest to the RL venv cache to automatically refresh stale venv's. Signed-off-by: Cory Ye --- nemo_rl/utils/venvs.py | 116 ++++++++++++++++++++------- tests/unit/utils/test_venvs.py | 138 +++++++++++++++++++++++++++++++++ 2 files changed, 226 insertions(+), 28 deletions(-) diff --git a/nemo_rl/utils/venvs.py b/nemo_rl/utils/venvs.py index 65f6d8b0930..b71d12cdfd4 100644 --- a/nemo_rl/utils/venvs.py +++ b/nemo_rl/utils/venvs.py @@ -17,7 +17,7 @@ import shutil import subprocess import time -from functools import lru_cache +from hashlib import sha256 from pathlib import Path import ray @@ -26,6 +26,8 @@ dir_path = os.path.dirname(os.path.abspath(__file__)) git_root = os.path.abspath(os.path.join(dir_path, "../..")) DEFAULT_VENV_DIR = os.path.join(git_root, "venvs") +_VENV_SPEC_FILE = ".nemo_rl_venv_spec" +_VENV_BUILD_LOCK_SUFFIX = ".STARTED_ENV_BUILDER" logger = logging.getLogger(__name__) @@ -49,7 +51,40 @@ def add_hf_modules_cache_to_pythonpath(env_vars: dict[str, str]) -> dict[str, st return result -@lru_cache(maxsize=None) +def _venv_fingerprint(py_executable: str) -> str: + """Fingerprint the requested uv tier and its resolved project dependencies.""" + digest = sha256() + digest.update(py_executable.encode("utf-8")) + for filename in ("pyproject.toml", "uv.lock"): + path = Path(git_root) / filename + digest.update(filename.encode("utf-8")) + try: + digest.update(path.read_bytes()) + except FileNotFoundError: + digest.update(b"") + return digest.hexdigest() + + +def _venv_matches_spec(venv_path: Path, py_executable: str) -> bool: + """Return whether a cached venv matches the uv tier and dependency lock.""" + python_path = venv_path / "bin" / "python" + spec_path = venv_path / _VENV_SPEC_FILE + if not python_path.exists() or not spec_path.is_file(): + return False + try: + return spec_path.read_text(encoding="utf-8") == _venv_fingerprint(py_executable) + except OSError: + return False + + +def _write_venv_spec(venv_path: Path, py_executable: str) -> None: + """Atomically mark a venv valid after all of its sync commands succeed.""" + spec_path = venv_path / _VENV_SPEC_FILE + temporary_path = venv_path / f"{_VENV_SPEC_FILE}.{os.getpid()}.tmp" + temporary_path.write_text(_venv_fingerprint(py_executable), encoding="utf-8") + os.replace(temporary_path, spec_path) + + def create_local_venv( py_executable: str, venv_name: str, force_rebuild: bool = False ) -> str: @@ -58,8 +93,10 @@ def create_local_venv( The output can be used as a py_executable for a Ray worker assuming the worker nodes also have access to the same file system as the head node. - This function is cached to avoid multiple calls to uv to create the same venv, - which avoids duplicate logging. + A fingerprint of the requested uv invocation, ``pyproject.toml``, and + ``uv.lock`` is persisted inside the venv. Reusing an actor name after its + extras or resolved dependencies change automatically resyncs the environment + instead of trusting an incompatible cached interpreter. Args: py_executable (str): Command to run with the virtual environment (e.g., "uv.sh run --locked") @@ -77,26 +114,40 @@ def create_local_venv( # # You can override this location by setting the NEMO_RL_VENV_DIR environment variable - NEMO_RL_VENV_DIR = os.path.normpath( + nemo_rl_venv_dir = os.path.normpath( os.environ.get("NEMO_RL_VENV_DIR", DEFAULT_VENV_DIR) ) - logger.info(f"NEMO_RL_VENV_DIR is set to {NEMO_RL_VENV_DIR}.") + logger.info(f"NEMO_RL_VENV_DIR is set to {nemo_rl_venv_dir}.") # Create the venv directory if it doesn't exist - os.makedirs(NEMO_RL_VENV_DIR, exist_ok=True) + os.makedirs(nemo_rl_venv_dir, exist_ok=True) # Full path to the virtual environment - venv_path = os.path.join(NEMO_RL_VENV_DIR, venv_name) + venv_path = Path(nemo_rl_venv_dir) / venv_name + python_path = venv_path / "bin" / "python" - # Force rebuild if requested - if force_rebuild and os.path.exists(venv_path): + # Build or retrieve the venv. + if force_rebuild and venv_path.exists(): + # Force rebuild if requested. logger.info(f"Force rebuilding venv at {venv_path}") shutil.rmtree(venv_path) + elif _venv_matches_spec(venv_path, py_executable): + # If the cached venv spec matches the current requirements + # computed from pyproject.toml, uv.lock, etc. then use the + # cached venv. + logger.info(f"Using compatible cached venv at {venv_path}") + return str(python_path) + elif venv_path.exists(): + # Rebuild / update the venv. + logger.info( + "Refreshing venv at %s because its cached uv specification changed", + venv_path, + ) logger.info(f"Creating new venv at {venv_path}") # Create the virtual environment - uv_venv_cmd = ["uv", "venv", "--allow-existing", venv_path] + uv_venv_cmd = ["uv", "venv", "--allow-existing", str(venv_path)] subprocess.run(uv_venv_cmd, check=True) # Execute the command with the virtual environment @@ -105,7 +156,7 @@ def create_local_venv( # one call to this in the driver. It is not safe to use this in a multi-process # context. # https://docs.astral.sh/uv/concepts/projects/config/#project-environment-path - env["UV_PROJECT_ENVIRONMENT"] = venv_path + env["UV_PROJECT_ENVIRONMENT"] = str(venv_path) # Split the py_executable into command and arguments exec_cmd = shlex.split(py_executable) @@ -116,9 +167,12 @@ def create_local_venv( subprocess.run(["uv", "sync", "--directory", git_root], env=env, check=True) subprocess.run(exec_cmd, env=env, check=True) + if not python_path.exists(): + raise RuntimeError(f"uv completed without creating {python_path}") + _write_venv_spec(venv_path, py_executable) + # Return the path to the python executable in the virtual environment - python_path = os.path.join(venv_path, "bin", "python") - return python_path + return str(python_path) # Ray-based helper to create a virtual environment on each Ray node @@ -127,37 +181,43 @@ def _env_builder( py_executable: str, venv_name: str, node_idx: int, force_rebuild: bool = False ): # Check if another node is already building - NEMO_RL_VENV_DIR = os.path.normpath( + nemo_rl_venv_dir = os.path.normpath( os.environ.get("NEMO_RL_VENV_DIR", DEFAULT_VENV_DIR) ) - venv_path = Path(NEMO_RL_VENV_DIR) / venv_name + venv_path = Path(nemo_rl_venv_dir) / venv_name python_path = venv_path / "bin" / "python" - started_file = venv_path / "STARTED_ENV_BUILDER" + started_file = Path(f"{venv_path}{_VENV_BUILD_LOCK_SUFFIX}") # Skip early return if force_rebuild is True - if not force_rebuild and python_path.exists(): + if not force_rebuild and _venv_matches_spec(venv_path, py_executable): logger.info(f"Using existing venv at {venv_path}") return str(python_path) # Sleep to stagger node startup time.sleep(1 * node_idx) - if started_file.exists(): + try: + started_file.parent.mkdir(parents=True, exist_ok=True) + lock_fd = os.open(started_file, os.O_CREAT | os.O_EXCL | os.O_WRONLY) + os.close(lock_fd) + owns_build_lock = True + except FileExistsError: + owns_build_lock = False + + if not owns_build_lock: # Another node is already building, wait for completion logger.info( f"Node {node_idx}: Another node is building {venv_name}, skipping..." ) - # Wait for the venv to be ready (check for python executable) - python_path = venv_path / "bin" / "python" - while not python_path.exists(): + while started_file.exists(): time.sleep(1) - return str(python_path) - - # Create the venv directory if needed - venv_path.mkdir(parents=True, exist_ok=True) + if _venv_matches_spec(venv_path, py_executable): + return str(python_path) + raise RuntimeError( + f"Venv builder for {venv_name} finished without producing a " + f"compatible environment at {venv_path}" + ) - # Touch the started file to signal we're building - started_file.touch() try: # Create the virtual environment on this node return create_local_venv(py_executable, venv_name, force_rebuild=force_rebuild) diff --git a/tests/unit/utils/test_venvs.py b/tests/unit/utils/test_venvs.py index 2635689385b..639048e806e 100644 --- a/tests/unit/utils/test_venvs.py +++ b/tests/unit/utils/test_venvs.py @@ -13,9 +13,12 @@ # limitations under the License. import os import subprocess +from pathlib import Path from tempfile import TemporaryDirectory from unittest.mock import patch +import pytest + from nemo_rl.utils.venvs import ( add_hf_modules_cache_to_pythonpath, create_local_venv, @@ -24,6 +27,13 @@ from tests.unit.conftest import TEST_ASSETS_DIR +def _mock_uv_venv(command, **_kwargs): + if command[:2] == ["uv", "venv"]: + python_path = Path(command[-1]) / "bin" / "python" + python_path.parent.mkdir(parents=True, exist_ok=True) + python_path.touch() + + def test_create_local_venv(): # The temporary directory is created within the project. # For some reason, creating a virtual environment outside of the project @@ -54,6 +64,134 @@ def test_create_local_venv(): assert "Sphinx package is installed" in result.stdout +def test_create_local_venv_refreshes_when_uv_spec_changes(tmp_path): + vllm = "uv run --locked --extra vllm" + vllm_gym = "uv run --locked --extra vllm --extra nemo_gym" + venv_path = tmp_path / "worker" + + with ( + patch.dict(os.environ, {"NEMO_RL_VENV_DIR": str(tmp_path)}), + patch( + "nemo_rl.utils.venvs.subprocess.run", side_effect=_mock_uv_venv + ) as mock_run, + ): + create_local_venv(vllm, "worker") + first_fingerprint = (venv_path / ".nemo_rl_venv_spec").read_text() + + mock_run.reset_mock() + create_local_venv(vllm, "worker") + mock_run.assert_not_called() + + create_local_venv(vllm_gym, "worker") + assert mock_run.call_count == 3 + assert (venv_path / ".nemo_rl_venv_spec").read_text() != first_fingerprint + + +def test_create_local_venv_refreshes_legacy_cache_without_spec(tmp_path): + venv_path = tmp_path / "worker" + python_path = venv_path / "bin" / "python" + python_path.parent.mkdir(parents=True) + python_path.touch() + + with ( + patch.dict(os.environ, {"NEMO_RL_VENV_DIR": str(tmp_path)}), + patch( + "nemo_rl.utils.venvs.subprocess.run", side_effect=_mock_uv_venv + ) as mock_run, + ): + create_local_venv("uv run --locked --extra vllm", "worker") + + assert mock_run.call_count == 3 + assert (venv_path / ".nemo_rl_venv_spec").is_file() + + +def test_create_local_venv_refreshes_when_lockfile_changes(tmp_path): + project_path = tmp_path / "project" + project_path.mkdir() + (project_path / "pyproject.toml").write_text("[project]\nname = 'test'\n") + lock_path = project_path / "uv.lock" + lock_path.write_text("version = 1\n") + venv_root = tmp_path / "venvs" + venv_path = venv_root / "worker" + py_executable = "uv run --locked --extra vllm" + + with ( + patch.dict(os.environ, {"NEMO_RL_VENV_DIR": str(venv_root)}), + patch("nemo_rl.utils.venvs.git_root", str(project_path)), + patch( + "nemo_rl.utils.venvs.subprocess.run", side_effect=_mock_uv_venv + ) as mock_run, + ): + create_local_venv(py_executable, "worker") + first_fingerprint = (venv_path / ".nemo_rl_venv_spec").read_text() + + mock_run.reset_mock() + lock_path.write_text("version = 2\n") + create_local_venv(py_executable, "worker") + + assert mock_run.call_count == 3 + assert (venv_path / ".nemo_rl_venv_spec").read_text() != first_fingerprint + + +def test_create_local_venv_refreshes_when_pyproject_changes(tmp_path): + project_path = tmp_path / "project" + project_path.mkdir() + pyproject_path = project_path / "pyproject.toml" + pyproject_path.write_text("[project]\nname = 'test'\n") + (project_path / "uv.lock").write_text("version = 1\n") + venv_root = tmp_path / "venvs" + venv_path = venv_root / "worker" + py_executable = "uv run --locked --extra vllm" + + with ( + patch.dict(os.environ, {"NEMO_RL_VENV_DIR": str(venv_root)}), + patch("nemo_rl.utils.venvs.git_root", str(project_path)), + patch( + "nemo_rl.utils.venvs.subprocess.run", side_effect=_mock_uv_venv + ) as mock_run, + ): + create_local_venv(py_executable, "worker") + first_fingerprint = (venv_path / ".nemo_rl_venv_spec").read_text() + + mock_run.reset_mock() + pyproject_path.write_text( + "[project]\nname = 'test'\ndependencies = ['new-dependency']\n" + ) + create_local_venv(py_executable, "worker") + + assert mock_run.call_count == 3 + assert (venv_path / ".nemo_rl_venv_spec").read_text() != first_fingerprint + + +def test_create_local_venv_does_not_cache_failed_sync(tmp_path): + venv_path = tmp_path / "worker" + + def fail_sync(command, **kwargs): + _mock_uv_venv(command, **kwargs) + if command[:2] == ["uv", "sync"]: + raise subprocess.CalledProcessError(2, command) + + with ( + patch.dict(os.environ, {"NEMO_RL_VENV_DIR": str(tmp_path)}), + patch("nemo_rl.utils.venvs.subprocess.run", side_effect=fail_sync), + pytest.raises(subprocess.CalledProcessError), + ): + create_local_venv("uv run --locked --extra vllm", "worker") + + assert not (venv_path / ".nemo_rl_venv_spec").exists() + + with ( + patch.dict(os.environ, {"NEMO_RL_VENV_DIR": str(tmp_path)}), + patch( + "nemo_rl.utils.venvs.subprocess.run", side_effect=_mock_uv_venv + ) as mock_run, + ): + create_local_venv("uv run --locked --extra vllm", "worker") + + assert mock_run.call_count == 3 + assert (venv_path / ".nemo_rl_venv_spec").is_file() + + def test_add_hf_modules_cache_to_pythonpath(): result = add_hf_modules_cache_to_pythonpath( { From 3e547fa9166dbf9311651882e7d191837ffbbb66 Mon Sep 17 00:00:00 2001 From: Terry Kong Date: Mon, 7 Sep 2026 10:09:31 -0700 Subject: [PATCH 18/19] fix(sc): give vLLM workers the nemo_gym extra by default instead of swapping the env at runtime Token capture swapped the VllmAsyncGenerationWorker registry entry to VLLM_GYM at setup time. Worker venvs are cached by actor class name, so a venv prebuilt with plain --extra vllm (which is what the Dockerfile bakes) was reused as-is and the nemo_gym import failed. This is why the SingleController L1 job fails even on a freshly built container. Make VLLM_EXECUTABLE use VLLM_GYM so the baked venv already has nemo_gym, and drop the runtime registry override and the error wrapper that described the old behavior. Signed-off-by: Terry Kong --- .../single_controller_utils/setup.py | 33 +++---------------- .../ray_actor_environment_registry.py | 7 +++- nemo_rl/distributed/virtual_cluster.py | 8 +++-- 3 files changed, 17 insertions(+), 31 deletions(-) diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py index 1893a115216..a3411f11bf0 100644 --- a/nemo_rl/algorithms/single_controller_utils/setup.py +++ b/nemo_rl/algorithms/single_controller_utils/setup.py @@ -998,9 +998,10 @@ def setup_single_controller( policy_config["pretrained_checkpoint"] = checkpointing_pretrained # Token capture: validate the supported combination loudly at setup - # (NeMo-Gym rollout path, vLLM backend, async_engine=true) and give - # capture-enabled vLLM workers a venv that carries nemo_gym (the - # worker hosts Gym's capture core + adapter in-process). + # (NeMo-Gym rollout path, vLLM backend, async_engine=true). The vLLM + # worker venv always carries nemo_gym (see VLLM_EXECUTABLE in + # ray_actor_environment_registry.py), so nothing here needs to change the + # worker's environment. token_capture_cfg = master_config.token_capture if token_capture_cfg.enabled: if not should_use_nemo_gym(master_config): @@ -1021,14 +1022,6 @@ def setup_single_controller( "policy.generation.vllm_cfg.async_engine=true (the capture " "host is the worker's in-process HTTP server)" ) - from nemo_rl.distributed.ray_actor_environment_registry import ( - ACTOR_ENVIRONMENT_REGISTRY, - ) - from nemo_rl.distributed.virtual_cluster import PY_EXECUTABLES - - ACTOR_ENVIRONMENT_REGISTRY[ - "nemo_rl.models.generation.vllm.vllm_worker_async.VllmAsyncGenerationWorker" - ] = PY_EXECUTABLES.VLLM_GYM # Fill the derived ledger-hosting fields (see TokenCaptureConfig): a # per-run control-plane bearer token and the process-shared capture @@ -1501,23 +1494,7 @@ def _build_generation_then_trainer( # Host Gym's capture core in every vLLM DP leader (in-worker DP # client + TQTokenSink + the single install_capture call), and give # workers the initial weight version to stamp on captured calls. - try: - generation.setup_token_capture( - dp_config, token_capture_cfg.staging_partition - ) - except Exception as error: - if "No module named 'nemo_gym'" in str(error): - # Worker venvs are cached by actor class name - # (nemo_rl/utils/venvs.py), so a venv prebuilt before token - # capture predates the nemo_gym extra and is reused as-is. - raise RuntimeError( - "token_capture.enabled requires nemo_gym inside the vLLM " - "worker venv, but the cached worker venv predates it. " - "Rebuild worker venvs (NRL_FORCE_REBUILD_VENVS=true) or " - "delete $NEMO_RL_VENV_DIR/nemo_rl.models.generation.vllm." - "vllm_worker_async.VllmAsyncGenerationWorker and rerun." - ) from error - raise + generation.setup_token_capture(dp_config, token_capture_cfg.staging_partition) generation.set_rollout_weight_version(0) if weight_synchronizer is None: diff --git a/nemo_rl/distributed/ray_actor_environment_registry.py b/nemo_rl/distributed/ray_actor_environment_registry.py index 1d27c8ff413..a2788019ab4 100644 --- a/nemo_rl/distributed/ray_actor_environment_registry.py +++ b/nemo_rl/distributed/ray_actor_environment_registry.py @@ -17,8 +17,13 @@ from nemo_rl.distributed.virtual_cluster import PY_EXECUTABLES USE_SYSTEM_EXECUTABLE = os.environ.get("NEMO_RL_PY_EXECUTABLES_SYSTEM", "0") == "1" +# vLLM workers always get the vllm + nemo_gym extras. Token capture +# (token_capture.enabled) needs nemo_gym inside the worker, and worker venvs +# are cached by actor class name, so the extras must be fixed here rather than +# swapped in at runtime (a venv prebuilt with plain `--extra vllm` would be +# reused as-is and the nemo_gym import would fail). VLLM_EXECUTABLE = ( - PY_EXECUTABLES.SYSTEM if USE_SYSTEM_EXECUTABLE else PY_EXECUTABLES.VLLM + PY_EXECUTABLES.SYSTEM if USE_SYSTEM_EXECUTABLE else PY_EXECUTABLES.VLLM_GYM ) SGLANG_EXECUTABLE = ( PY_EXECUTABLES.SYSTEM if USE_SYSTEM_EXECUTABLE else PY_EXECUTABLES.SGLANG diff --git a/nemo_rl/distributed/virtual_cluster.py b/nemo_rl/distributed/virtual_cluster.py index 0d2d7455b5c..3f8ae284a62 100644 --- a/nemo_rl/distributed/virtual_cluster.py +++ b/nemo_rl/distributed/virtual_cluster.py @@ -75,8 +75,12 @@ class PY_EXECUTABLES: # Use NeMo-Gym dependencies NEMO_GYM = f"uv run --locked --extra nemo_gym --directory {git_root}" - # vLLM worker hosting Gym's token capture (token_capture.enabled): the - # worker imports nemo_gym's dependency-free capture core + vLLM adapter. + # Default env for the vLLM generation workers (see + # ray_actor_environment_registry.py). It carries nemo_gym so the worker can + # host Gym's token capture (token_capture.enabled) without swapping the + # worker's env at runtime: worker venvs are cached by actor class name, so + # a venv prebuilt with plain `--extra vllm` would be reused as-is and the + # nemo_gym import would fail. VLLM_GYM = f"uv run --locked --extra vllm --extra nemo_gym --directory {git_root}" # Use NeMo-RL direct dependencies and SGLang. From 82a0dc56284a916635275673586221cc6c7de237 Mon Sep 17 00:00:00 2001 From: Cory Ye Date: Mon, 7 Sep 2026 10:31:59 -0700 Subject: [PATCH 19/19] Revert the cache digest and refresh. Signed-off-by: Cory Ye --- nemo_rl/utils/venvs.py | 116 +++++++-------------------- tests/unit/utils/test_venvs.py | 138 --------------------------------- 2 files changed, 28 insertions(+), 226 deletions(-) diff --git a/nemo_rl/utils/venvs.py b/nemo_rl/utils/venvs.py index b71d12cdfd4..65f6d8b0930 100644 --- a/nemo_rl/utils/venvs.py +++ b/nemo_rl/utils/venvs.py @@ -17,7 +17,7 @@ import shutil import subprocess import time -from hashlib import sha256 +from functools import lru_cache from pathlib import Path import ray @@ -26,8 +26,6 @@ dir_path = os.path.dirname(os.path.abspath(__file__)) git_root = os.path.abspath(os.path.join(dir_path, "../..")) DEFAULT_VENV_DIR = os.path.join(git_root, "venvs") -_VENV_SPEC_FILE = ".nemo_rl_venv_spec" -_VENV_BUILD_LOCK_SUFFIX = ".STARTED_ENV_BUILDER" logger = logging.getLogger(__name__) @@ -51,40 +49,7 @@ def add_hf_modules_cache_to_pythonpath(env_vars: dict[str, str]) -> dict[str, st return result -def _venv_fingerprint(py_executable: str) -> str: - """Fingerprint the requested uv tier and its resolved project dependencies.""" - digest = sha256() - digest.update(py_executable.encode("utf-8")) - for filename in ("pyproject.toml", "uv.lock"): - path = Path(git_root) / filename - digest.update(filename.encode("utf-8")) - try: - digest.update(path.read_bytes()) - except FileNotFoundError: - digest.update(b"") - return digest.hexdigest() - - -def _venv_matches_spec(venv_path: Path, py_executable: str) -> bool: - """Return whether a cached venv matches the uv tier and dependency lock.""" - python_path = venv_path / "bin" / "python" - spec_path = venv_path / _VENV_SPEC_FILE - if not python_path.exists() or not spec_path.is_file(): - return False - try: - return spec_path.read_text(encoding="utf-8") == _venv_fingerprint(py_executable) - except OSError: - return False - - -def _write_venv_spec(venv_path: Path, py_executable: str) -> None: - """Atomically mark a venv valid after all of its sync commands succeed.""" - spec_path = venv_path / _VENV_SPEC_FILE - temporary_path = venv_path / f"{_VENV_SPEC_FILE}.{os.getpid()}.tmp" - temporary_path.write_text(_venv_fingerprint(py_executable), encoding="utf-8") - os.replace(temporary_path, spec_path) - - +@lru_cache(maxsize=None) def create_local_venv( py_executable: str, venv_name: str, force_rebuild: bool = False ) -> str: @@ -93,10 +58,8 @@ def create_local_venv( The output can be used as a py_executable for a Ray worker assuming the worker nodes also have access to the same file system as the head node. - A fingerprint of the requested uv invocation, ``pyproject.toml``, and - ``uv.lock`` is persisted inside the venv. Reusing an actor name after its - extras or resolved dependencies change automatically resyncs the environment - instead of trusting an incompatible cached interpreter. + This function is cached to avoid multiple calls to uv to create the same venv, + which avoids duplicate logging. Args: py_executable (str): Command to run with the virtual environment (e.g., "uv.sh run --locked") @@ -114,40 +77,26 @@ def create_local_venv( # # You can override this location by setting the NEMO_RL_VENV_DIR environment variable - nemo_rl_venv_dir = os.path.normpath( + NEMO_RL_VENV_DIR = os.path.normpath( os.environ.get("NEMO_RL_VENV_DIR", DEFAULT_VENV_DIR) ) - logger.info(f"NEMO_RL_VENV_DIR is set to {nemo_rl_venv_dir}.") + logger.info(f"NEMO_RL_VENV_DIR is set to {NEMO_RL_VENV_DIR}.") # Create the venv directory if it doesn't exist - os.makedirs(nemo_rl_venv_dir, exist_ok=True) + os.makedirs(NEMO_RL_VENV_DIR, exist_ok=True) # Full path to the virtual environment - venv_path = Path(nemo_rl_venv_dir) / venv_name - python_path = venv_path / "bin" / "python" + venv_path = os.path.join(NEMO_RL_VENV_DIR, venv_name) - # Build or retrieve the venv. - if force_rebuild and venv_path.exists(): - # Force rebuild if requested. + # Force rebuild if requested + if force_rebuild and os.path.exists(venv_path): logger.info(f"Force rebuilding venv at {venv_path}") shutil.rmtree(venv_path) - elif _venv_matches_spec(venv_path, py_executable): - # If the cached venv spec matches the current requirements - # computed from pyproject.toml, uv.lock, etc. then use the - # cached venv. - logger.info(f"Using compatible cached venv at {venv_path}") - return str(python_path) - elif venv_path.exists(): - # Rebuild / update the venv. - logger.info( - "Refreshing venv at %s because its cached uv specification changed", - venv_path, - ) logger.info(f"Creating new venv at {venv_path}") # Create the virtual environment - uv_venv_cmd = ["uv", "venv", "--allow-existing", str(venv_path)] + uv_venv_cmd = ["uv", "venv", "--allow-existing", venv_path] subprocess.run(uv_venv_cmd, check=True) # Execute the command with the virtual environment @@ -156,7 +105,7 @@ def create_local_venv( # one call to this in the driver. It is not safe to use this in a multi-process # context. # https://docs.astral.sh/uv/concepts/projects/config/#project-environment-path - env["UV_PROJECT_ENVIRONMENT"] = str(venv_path) + env["UV_PROJECT_ENVIRONMENT"] = venv_path # Split the py_executable into command and arguments exec_cmd = shlex.split(py_executable) @@ -167,12 +116,9 @@ def create_local_venv( subprocess.run(["uv", "sync", "--directory", git_root], env=env, check=True) subprocess.run(exec_cmd, env=env, check=True) - if not python_path.exists(): - raise RuntimeError(f"uv completed without creating {python_path}") - _write_venv_spec(venv_path, py_executable) - # Return the path to the python executable in the virtual environment - return str(python_path) + python_path = os.path.join(venv_path, "bin", "python") + return python_path # Ray-based helper to create a virtual environment on each Ray node @@ -181,43 +127,37 @@ def _env_builder( py_executable: str, venv_name: str, node_idx: int, force_rebuild: bool = False ): # Check if another node is already building - nemo_rl_venv_dir = os.path.normpath( + NEMO_RL_VENV_DIR = os.path.normpath( os.environ.get("NEMO_RL_VENV_DIR", DEFAULT_VENV_DIR) ) - venv_path = Path(nemo_rl_venv_dir) / venv_name + venv_path = Path(NEMO_RL_VENV_DIR) / venv_name python_path = venv_path / "bin" / "python" - started_file = Path(f"{venv_path}{_VENV_BUILD_LOCK_SUFFIX}") + started_file = venv_path / "STARTED_ENV_BUILDER" # Skip early return if force_rebuild is True - if not force_rebuild and _venv_matches_spec(venv_path, py_executable): + if not force_rebuild and python_path.exists(): logger.info(f"Using existing venv at {venv_path}") return str(python_path) # Sleep to stagger node startup time.sleep(1 * node_idx) - try: - started_file.parent.mkdir(parents=True, exist_ok=True) - lock_fd = os.open(started_file, os.O_CREAT | os.O_EXCL | os.O_WRONLY) - os.close(lock_fd) - owns_build_lock = True - except FileExistsError: - owns_build_lock = False - - if not owns_build_lock: + if started_file.exists(): # Another node is already building, wait for completion logger.info( f"Node {node_idx}: Another node is building {venv_name}, skipping..." ) - while started_file.exists(): + # Wait for the venv to be ready (check for python executable) + python_path = venv_path / "bin" / "python" + while not python_path.exists(): time.sleep(1) - if _venv_matches_spec(venv_path, py_executable): - return str(python_path) - raise RuntimeError( - f"Venv builder for {venv_name} finished without producing a " - f"compatible environment at {venv_path}" - ) + return str(python_path) + + # Create the venv directory if needed + venv_path.mkdir(parents=True, exist_ok=True) + # Touch the started file to signal we're building + started_file.touch() try: # Create the virtual environment on this node return create_local_venv(py_executable, venv_name, force_rebuild=force_rebuild) diff --git a/tests/unit/utils/test_venvs.py b/tests/unit/utils/test_venvs.py index 639048e806e..2635689385b 100644 --- a/tests/unit/utils/test_venvs.py +++ b/tests/unit/utils/test_venvs.py @@ -13,12 +13,9 @@ # limitations under the License. import os import subprocess -from pathlib import Path from tempfile import TemporaryDirectory from unittest.mock import patch -import pytest - from nemo_rl.utils.venvs import ( add_hf_modules_cache_to_pythonpath, create_local_venv, @@ -27,13 +24,6 @@ from tests.unit.conftest import TEST_ASSETS_DIR -def _mock_uv_venv(command, **_kwargs): - if command[:2] == ["uv", "venv"]: - python_path = Path(command[-1]) / "bin" / "python" - python_path.parent.mkdir(parents=True, exist_ok=True) - python_path.touch() - - def test_create_local_venv(): # The temporary directory is created within the project. # For some reason, creating a virtual environment outside of the project @@ -64,134 +54,6 @@ def test_create_local_venv(): assert "Sphinx package is installed" in result.stdout -def test_create_local_venv_refreshes_when_uv_spec_changes(tmp_path): - vllm = "uv run --locked --extra vllm" - vllm_gym = "uv run --locked --extra vllm --extra nemo_gym" - venv_path = tmp_path / "worker" - - with ( - patch.dict(os.environ, {"NEMO_RL_VENV_DIR": str(tmp_path)}), - patch( - "nemo_rl.utils.venvs.subprocess.run", side_effect=_mock_uv_venv - ) as mock_run, - ): - create_local_venv(vllm, "worker") - first_fingerprint = (venv_path / ".nemo_rl_venv_spec").read_text() - - mock_run.reset_mock() - create_local_venv(vllm, "worker") - mock_run.assert_not_called() - - create_local_venv(vllm_gym, "worker") - assert mock_run.call_count == 3 - assert (venv_path / ".nemo_rl_venv_spec").read_text() != first_fingerprint - - -def test_create_local_venv_refreshes_legacy_cache_without_spec(tmp_path): - venv_path = tmp_path / "worker" - python_path = venv_path / "bin" / "python" - python_path.parent.mkdir(parents=True) - python_path.touch() - - with ( - patch.dict(os.environ, {"NEMO_RL_VENV_DIR": str(tmp_path)}), - patch( - "nemo_rl.utils.venvs.subprocess.run", side_effect=_mock_uv_venv - ) as mock_run, - ): - create_local_venv("uv run --locked --extra vllm", "worker") - - assert mock_run.call_count == 3 - assert (venv_path / ".nemo_rl_venv_spec").is_file() - - -def test_create_local_venv_refreshes_when_lockfile_changes(tmp_path): - project_path = tmp_path / "project" - project_path.mkdir() - (project_path / "pyproject.toml").write_text("[project]\nname = 'test'\n") - lock_path = project_path / "uv.lock" - lock_path.write_text("version = 1\n") - venv_root = tmp_path / "venvs" - venv_path = venv_root / "worker" - py_executable = "uv run --locked --extra vllm" - - with ( - patch.dict(os.environ, {"NEMO_RL_VENV_DIR": str(venv_root)}), - patch("nemo_rl.utils.venvs.git_root", str(project_path)), - patch( - "nemo_rl.utils.venvs.subprocess.run", side_effect=_mock_uv_venv - ) as mock_run, - ): - create_local_venv(py_executable, "worker") - first_fingerprint = (venv_path / ".nemo_rl_venv_spec").read_text() - - mock_run.reset_mock() - lock_path.write_text("version = 2\n") - create_local_venv(py_executable, "worker") - - assert mock_run.call_count == 3 - assert (venv_path / ".nemo_rl_venv_spec").read_text() != first_fingerprint - - -def test_create_local_venv_refreshes_when_pyproject_changes(tmp_path): - project_path = tmp_path / "project" - project_path.mkdir() - pyproject_path = project_path / "pyproject.toml" - pyproject_path.write_text("[project]\nname = 'test'\n") - (project_path / "uv.lock").write_text("version = 1\n") - venv_root = tmp_path / "venvs" - venv_path = venv_root / "worker" - py_executable = "uv run --locked --extra vllm" - - with ( - patch.dict(os.environ, {"NEMO_RL_VENV_DIR": str(venv_root)}), - patch("nemo_rl.utils.venvs.git_root", str(project_path)), - patch( - "nemo_rl.utils.venvs.subprocess.run", side_effect=_mock_uv_venv - ) as mock_run, - ): - create_local_venv(py_executable, "worker") - first_fingerprint = (venv_path / ".nemo_rl_venv_spec").read_text() - - mock_run.reset_mock() - pyproject_path.write_text( - "[project]\nname = 'test'\ndependencies = ['new-dependency']\n" - ) - create_local_venv(py_executable, "worker") - - assert mock_run.call_count == 3 - assert (venv_path / ".nemo_rl_venv_spec").read_text() != first_fingerprint - - -def test_create_local_venv_does_not_cache_failed_sync(tmp_path): - venv_path = tmp_path / "worker" - - def fail_sync(command, **kwargs): - _mock_uv_venv(command, **kwargs) - if command[:2] == ["uv", "sync"]: - raise subprocess.CalledProcessError(2, command) - - with ( - patch.dict(os.environ, {"NEMO_RL_VENV_DIR": str(tmp_path)}), - patch("nemo_rl.utils.venvs.subprocess.run", side_effect=fail_sync), - pytest.raises(subprocess.CalledProcessError), - ): - create_local_venv("uv run --locked --extra vllm", "worker") - - assert not (venv_path / ".nemo_rl_venv_spec").exists() - - with ( - patch.dict(os.environ, {"NEMO_RL_VENV_DIR": str(tmp_path)}), - patch( - "nemo_rl.utils.venvs.subprocess.run", side_effect=_mock_uv_venv - ) as mock_run, - ): - create_local_venv("uv run --locked --extra vllm", "worker") - - assert mock_run.call_count == 3 - assert (venv_path / ".nemo_rl_venv_spec").is_file() - - def test_add_hf_modules_cache_to_pythonpath(): result = add_hf_modules_cache_to_pythonpath( {