diff --git a/3rdparty/Gym-workspace/Gym b/3rdparty/Gym-workspace/Gym index 473f446f71e..e9fd77cc31c 160000 --- a/3rdparty/Gym-workspace/Gym +++ b/3rdparty/Gym-workspace/Gym @@ -1 +1 @@ -Subproject commit 473f446f71ec7c1243eb1517fe2440a2b37fe68b +Subproject commit e9fd77cc31cdb511c8cdd02a7d19577d19c08fc9 diff --git a/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge b/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge index 573e088c9c6..bb0d604d2ab 160000 --- a/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge +++ b/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge @@ -1 +1 @@ -Subproject commit 573e088c9c6740082c39744e03dc5b009e730ed4 +Subproject commit bb0d604d2ab255eac29d18358a95c3661989465a diff --git a/ehsan_scripts/README.md b/ehsan_scripts/README.md new file mode 100644 index 00000000000..d0d3281de88 --- /dev/null +++ b/ehsan_scripts/README.md @@ -0,0 +1,118 @@ +# Nemotron Omni video GRPO launchers + +This directory contains the reproducible Slurm launchers used to validate +video Gym training: + +- `run_2n_sync.sh`: 2-node synchronous GRPO on the `interactive` partition + and `nemotron_edge_omni` account. +- `run_16n_async.sh`: 16-node asynchronous GRPO on either the + `nemotron_edge_omni` or `nemotron_omni_vision` account. + +The canonical training and data configuration is kept with the NeMo RL +recipes: + +- `examples/nemo_gym/grpo_nemotron_omni_30ba3b_video_sync.yaml` +- `examples/nemo_gym/grpo_nemotron_omni_30ba3b_video_async.yaml` +- `examples/nemo_gym/prepare_video_dataset.py` + +Both launchers use cached-video JSONL input. They intentionally keep +`grpo.max_num_steps=-1` and `grpo.seq_logprob_error_threshold=null`; the +four-hour Slurm allocation, rather than a GRPO step cap, ends a validation +run. They also enable W&B and checkpointing. No credential, user-specific +path, temporary source overlay, or version-mismatch bypass is embedded in +the scripts. + +The 16-node asynchronous recipe uses four prompts with sixteen generations +per prompt (global batch 64). Its policy uses the validated TP4/EP4 two-node +topology and the other fourteen nodes run generation. Thus only 16/128 GPUs +are idle during the initial rollout, below the batch scheduler's 25-percent +idle reaper threshold. This does not cap generations or GRPO training steps. + +The Megatron policy uses MBridge's canonical `NemotronOmniModel` expanded- +sequence contract. V2-labeled Nano Omni MoE checkpoints are routed through the +canonical bridge while dense V2 checkpoints retain their legacy behavior. The +launchers key the converted Megatron checkpoint cache by the exact MBridge +commit so a checkpoint produced by the retired LLaVA path cannot be reused. + +## Prepare the cached-video dataset + +The recipe expects training and validation JSONL files plus a media root. +The JSONL may be generated from a supported source dataset with: + +```bash +uv run examples/nemo_gym/prepare_video_dataset.py --help +``` + +Set `NEMO_RL_VIDEO_MEDIA_ROOT` to the filesystem prefix from which the video +paths in the JSONL can be resolved. No dataset content is committed to this +repository. + +## Required environment + +Set these variables before launching either job: + +```bash +export CONTAINER=/path/to/nemo-rl-vllm-0.25.1.sqsh +export MOUNTS=/lustre:/lustre +export NEMO_RL_MODEL=/path/to/nemotron-omni-checkpoint +export NEMO_RL_CHAT_TEMPLATE="${NEMO_RL_MODEL}/chat_template.jinja" +export NEMO_RL_VIDEO_TRAIN_JSONL=/path/to/cached_video_train.jsonl +export NEMO_RL_VIDEO_VAL_JSONL=/path/to/cached_video_validation.jsonl +export NEMO_RL_VIDEO_MEDIA_ROOT=/lustre +export NEMO_RL_RUN_ROOT=/path/to/training-output +export WANDB_API_KEY=... +export WANDB_ENTITY=... +export WANDB_PROJECT=... +``` + +The launchers install the stock vLLM `0.25.1` wheel declared by the project +into a node-local overlay and verify its version before training. The base +container must provide `/opt/rl_main_vg_runtime.env`, the NeMo RL +environment at `/opt/nemo_rl_venv`, the project metadata at `/opt/nemo-rl`, +and the video decoding dependencies. They bootstrap the exact `uv` release +declared by `docker/Dockerfile`, then install the interpreter declared by +`.python-version` into shared, versioned directories before Ray starts. This +keeps actor environments aligned with the checked-out branch when a compatible +base container has older `uv` or Python patch releases. A branch-locked main +environment is also materialized node-locally so the Ray cluster, driver, and +actor environments all use that same Python release. A runtime Ray dispatcher +in the output directory ensures the launcher's initial cleanup can use the base +image while cluster startup uses the branch-locked environment; it is not added +to the repository. After the initial cleanup, node setup installs that dispatcher +at the container-local Ray CLI path and preserves the image CLI as its fallback. +This ensures `ray.sub` starts the cluster with the same interpreter as the driver +even when the base image has an older Python patch release. Runtime dependency +synchronization uses the committed lock file in locked mode, so setup fails on +dependency drift instead of rewriting `uv.lock` in the source tree. Gym child +services use a run-scoped node-local venv root keyed by the commit and Python +release, preventing the recipe's venv reuse from selecting image-baked services +created with a different Python patch release. +The large uv package cache is also commit-keyed and node-local, avoiding shared +Lustre quota exhaustion while CUDA and PyTorch wheels are extracted on each node. +Credentials must come from the caller's environment or an approved secret +mechanism; do not add them to these scripts. +The launchers create a UTC timestamp run ID by default. Set +`NEMO_RL_RUN_ID` explicitly only when a stable W&B/checkpoint run name is +required, such as for an intentional resume. + +## Launch 2-node synchronous validation + +```bash +bash ehsan_scripts/run_2n_sync.sh +``` + +The launcher submits a four-hour job. To change only the Slurm duration, +set `SBATCH_TIME` before invoking it. Do not use the duration to introduce a +GRPO step limit. + +## Launch 16-node asynchronous validation + +Select one approved account and launch: + +```bash +export SBATCH_ACCOUNT=nemotron_edge_omni +bash ehsan_scripts/run_16n_async.sh +``` + +`SBATCH_ACCOUNT=nemotron_omni_vision` is also supported. If jobs are raced +between accounts, cancel the duplicate immediately after one starts. diff --git a/ehsan_scripts/run_16n_async.sh b/ehsan_scripts/run_16n_async.sh new file mode 100755 index 00000000000..02663d9b0cc --- /dev/null +++ b/ehsan_scripts/run_16n_async.sh @@ -0,0 +1,249 @@ +#!/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 -euo pipefail + +readonly SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +readonly REPO_ROOT="$(git -C "${SCRIPT_DIR}" rev-parse --show-toplevel)" +readonly SBATCH_PARTITION="batch_block1,backfill" + +require_env() { + local name="$1" + if [[ -z "${!name:-}" ]]; then + echo "Required environment variable is unset: ${name}" >&2 + exit 2 + fi +} + +for name in \ + CONTAINER \ + MOUNTS \ + NEMO_RL_CHAT_TEMPLATE \ + NEMO_RL_MODEL \ + NEMO_RL_VIDEO_TRAIN_JSONL \ + NEMO_RL_VIDEO_VAL_JSONL \ + NEMO_RL_VIDEO_MEDIA_ROOT \ + NEMO_RL_RUN_ROOT \ + SBATCH_ACCOUNT \ + WANDB_API_KEY \ + WANDB_ENTITY \ + WANDB_PROJECT; do + require_env "${name}" +done + +case "${SBATCH_ACCOUNT}" in + nemotron_edge_omni | nemotron_omni_vision) ;; + *) + echo "SBATCH_ACCOUNT must be nemotron_edge_omni or nemotron_omni_vision" >&2 + exit 2 + ;; +esac + +mkdir -p "${NEMO_RL_RUN_ROOT}" + +NEMO_RL_RUN_ID="${NEMO_RL_RUN_ID:-$(date -u +%Y%m%dT%H%M%SZ)}" +NEMO_RL_RUN_DATE="${NEMO_RL_RUN_DATE:-$(date -u +%m%d)}" + +export BASE_LOG_DIR="${NEMO_RL_RUN_ROOT}" +export CONTAINER +export GPUS_PER_NODE=8 +export MOUNTS +export NEMO_RL_CHAT_TEMPLATE +export NEMO_RL_EXPECTED_COMMIT="$(git -C "${REPO_ROOT}" rev-parse HEAD)" +export NEMO_RL_EXPECTED_GYM_COMMIT="$( + git -C "${REPO_ROOT}" ls-tree HEAD 3rdparty/Gym-workspace/Gym | awk '{print $3}' +)" +export NEMO_RL_EXPECTED_MBRIDGE_COMMIT="$( + git -C "${REPO_ROOT}" ls-tree HEAD 3rdparty/Megatron-Bridge-workspace/Megatron-Bridge | awk '{print $3}' +)" +export NEMO_RL_MODEL +export NEMO_RL_PYTHON_VERSION="$(tr -d '[:space:]' < "${REPO_ROOT}/.python-version")" +export NEMO_RL_REPO="${REPO_ROOT}" +export NEMO_RL_RUN_ID +export NEMO_RL_RUN_DATE +export NEMO_RL_RUN_ROOT +export NEMO_RL_VIDEO_MEDIA_ROOT +export NEMO_RL_VIDEO_TRAIN_JSONL +export NEMO_RL_VIDEO_VAL_JSONL +export NRL_RUN_PREFIX="rl_main_vg16_async_video_nemotron_omni_model_tp4_unlimited" +export RAY_TMPDIR=/tmp/ray +export NEMO_RL_UV_VERSION="$(awk -F= '/^ARG UV_VERSION=/{print $2; exit}' "${REPO_ROOT}/docker/Dockerfile")" +export NEMO_RL_UV_DIR="${NEMO_RL_RUN_ROOT}/uv/${NEMO_RL_UV_VERSION}" +export NEMO_RL_MAIN_VENV="/tmp/nemorl-main-${NEMO_RL_EXPECTED_COMMIT:0:12}" +export NEMO_RL_UV_CACHE_DIR="/tmp/nemorl-uv-cache-${NEMO_RL_EXPECTED_COMMIT:0:12}" +export NEMO_RL_RAY_WRAPPER_DIR="${NEMO_RL_RUN_ROOT}/runtime/${NEMO_RL_EXPECTED_COMMIT:0:12}/bin" +export UV_PYTHON_INSTALL_DIR="${NEMO_RL_RUN_ROOT}/uv_python/${NEMO_RL_PYTHON_VERSION}" +export WANDB_API_KEY +export WANDB_ENTITY +export WANDB_PROJECT +mkdir -p "${NEMO_RL_RAY_WRAPPER_DIR}" +readonly RAY_WRAPPER="${NEMO_RL_RAY_WRAPPER_DIR}/ray" +readonly RAY_WRAPPER_TMP="${RAY_WRAPPER}.tmp.$$" +printf '%s\n' \ + '#!/bin/bash' \ + 'set -euo pipefail' \ + 'if [[ -x "${NEMO_RL_MAIN_VENV}/bin/ray" ]]; then' \ + ' exec "${NEMO_RL_MAIN_VENV}/bin/ray" "$@"' \ + 'fi' \ + 'exec /opt/nemo_rl_venv/bin/ray.nemorl-base "$@"' \ + > "${RAY_WRAPPER_TMP}" +chmod 755 "${RAY_WRAPPER_TMP}" +mv -f "${RAY_WRAPPER_TMP}" "${RAY_WRAPPER}" +export PATH="${NEMO_RL_RAY_WRAPPER_DIR}:${NEMO_RL_MAIN_VENV}/bin:${NEMO_RL_UV_DIR}:${PATH}" +export VIRTUAL_ENV="${NEMO_RL_MAIN_VENV}" +export UV_PROJECT_ENVIRONMENT="${NEMO_RL_MAIN_VENV}" + +read -r -d '' SETUP_COMMAND <<'SETUP_EOF' || true +set -euo pipefail + +source /opt/rl_main_vg_runtime.env +mkdir -p "${NEMO_RL_UV_DIR}" +flock "${NEMO_RL_UV_DIR}.lock" bash -c ' + set -euo pipefail + if [[ ! -x "${NEMO_RL_UV_DIR}/uv" ]]; then + curl --retry 3 --retry-delay 2 -LsSf \ + "https://astral.sh/uv/${NEMO_RL_UV_VERSION}/install.sh" | \ + env UV_INSTALL_DIR="${NEMO_RL_UV_DIR}" UV_NO_MODIFY_PATH=1 sh + fi +' +export PATH="${NEMO_RL_UV_DIR}:/root/.local/bin:/opt/nemo_rl_venv/bin:${PATH}" +uv --version +mkdir -p "${UV_PYTHON_INSTALL_DIR}" "${NEMO_RL_UV_CACHE_DIR}" +flock "${UV_PYTHON_INSTALL_DIR}.lock" \ + uv python install "${NEMO_RL_PYTHON_VERSION}" +uv python find "${NEMO_RL_PYTHON_VERSION}" +export VIRTUAL_ENV="${NEMO_RL_MAIN_VENV}" +export UV_PROJECT_ENVIRONMENT="${NEMO_RL_MAIN_VENV}" +UV_CACHE_DIR="${NEMO_RL_UV_CACHE_DIR}" uv sync \ + --directory "${NEMO_RL_REPO}" \ + --locked \ + --no-install-project +readonly CONTAINER_RAY=/opt/nemo_rl_venv/bin/ray +readonly CONTAINER_RAY_BASE=/opt/nemo_rl_venv/bin/ray.nemorl-base +if [[ ! -x "${CONTAINER_RAY_BASE}" ]]; then + cp -p "${CONTAINER_RAY}" "${CONTAINER_RAY_BASE}" +fi +install -m 755 \ + "${NEMO_RL_RAY_WRAPPER_DIR}/ray" \ + "${CONTAINER_RAY}.tmp.$$" +mv -f "${CONTAINER_RAY}.tmp.$$" "${CONTAINER_RAY}" +export PATH="${NEMO_RL_MAIN_VENV}/bin:${NEMO_RL_UV_DIR}:/root/.local/bin:${PATH}" +python -c \ + 'import sys; assert sys.version_info[:3] == tuple(map(int, sys.argv[1].split("."))), sys.version' \ + "${NEMO_RL_PYTHON_VERSION}" +ray --version +readonly VLLM_OVERLAY="/tmp/nemorl-stock-vllm-0.25.1-${NEMO_RL_EXPECTED_COMMIT:0:12}" +readonly VLLM_WHEEL="https://github.com/vllm-project/vllm/releases/download/v0.25.1/vllm-0.25.1-cp38-abi3-manylinux_2_28_x86_64.whl" +mkdir -p "${VLLM_OVERLAY}" +UV_CACHE_DIR="${NEMO_RL_UV_CACHE_DIR}" uv pip install \ + --target "${VLLM_OVERLAY}" \ + --reinstall \ + --no-deps \ + "vllm @ ${VLLM_WHEEL}" +SETUP_EOF +export SETUP_COMMAND + +read -r -d '' COMMAND <<'COMMAND_EOF' || true +set -euo pipefail + +source /opt/rl_main_vg_runtime.env +export NEMO_GYM_VENV_DIR="/tmp/nemorl-gym-${NEMO_RL_EXPECTED_COMMIT:0:12}-${NEMO_RL_PYTHON_VERSION}-${NEMO_RL_RUN_ID}" +export PATH="${NEMO_RL_MAIN_VENV}/bin:${NEMO_RL_UV_DIR}:/root/.local/bin:${PATH}" +export VIRTUAL_ENV="${NEMO_RL_MAIN_VENV}" +export UV_PROJECT_ENVIRONMENT="${NEMO_RL_MAIN_VENV}" +export HF_HOME="${NEMO_RL_RUN_ROOT}/hf_home" +export HF_MODULES_CACHE="${HF_HOME}/modules" +export NRL_MEGATRON_CHECKPOINT_DIR="${HF_HOME}/nemo_rl-${NEMO_RL_EXPECTED_MBRIDGE_COMMIT:0:12}" +export NRL_FORCE_REBUILD_VENVS=true +export NRL_VIDEO_BACKEND=torchcodec +export NRL_VIDEO_SAMPLING_STYLE=nemotron_vl +export NRL_VIDEO_SFT_MAX_FRAMES=32 +export NRL_VIDEO_SFT_MIN_FRAMES=32 +export NRL_VIDEO_TEMPORAL_PATCH_SIZE=2 +export TORCH_CUDA_ARCH_LIST=9.0 +export VLLM_MAMBA_BACKEND=flashinfer +export VLLM_RAY_EXTRA_ENV_VARS_TO_COPY="PYTHONPATH,NEMO_RL_VIDEO_MEDIA_ROOT,NRL_VIDEO_BACKEND,NRL_VIDEO_SAMPLING_STYLE,NRL_VIDEO_TEMPORAL_PATCH_SIZE" +export VLLM_VIDEO_LOADER_BACKEND=nemotron_vl +unset NRL_IGNORE_VERSION_MISMATCH + +readonly VLLM_OVERLAY="/tmp/nemorl-stock-vllm-0.25.1-${NEMO_RL_EXPECTED_COMMIT:0:12}" +export NEMO_GYM_EXTRA_ROOTS="${NEMO_RL_REPO}/3rdparty/Gym-workspace/Gym" +export PYTHONPATH="${VLLM_OVERLAY}:${NEMO_RL_REPO}:${NEMO_GYM_EXTRA_ROOTS}:${NEMO_RL_REPO}/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge/src:${NEMO_RL_REPO}/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge/3rdparty/Megatron-LM:${HF_MODULES_CACHE}${PYTHONPATH:+:${PYTHONPATH}}" + +readonly RUN_NAME="${NRL_RUN_PREFIX}_${NEMO_RL_RUN_DATE}_${NEMO_RL_RUN_ID}" +readonly CHECKPOINT_DIR="${NEMO_RL_RUN_ROOT}/checkpoints/${RUN_NAME}" +mkdir -p "${CHECKPOINT_DIR}" "${HF_MODULES_CACHE}" "${NRL_MEGATRON_CHECKPOINT_DIR}" +cd "${NEMO_RL_REPO}" +test "$(git rev-parse HEAD)" = "${NEMO_RL_EXPECTED_COMMIT}" +test "$(git -C 3rdparty/Gym-workspace/Gym rev-parse HEAD)" = "${NEMO_RL_EXPECTED_GYM_COMMIT}" +test "$(git -C 3rdparty/Megatron-Bridge-workspace/Megatron-Bridge rev-parse HEAD)" = "${NEMO_RL_EXPECTED_MBRIDGE_COMMIT}" +python -c \ + 'import sys; assert sys.version_info[:3] == tuple(map(int, sys.argv[1].split("."))), sys.version' \ + "${NEMO_RL_PYTHON_VERSION}" + +python \ + examples/nemo_gym/run_grpo_nemo_gym.py \ + --config examples/nemo_gym/grpo_nemotron_omni_30ba3b_video_async.yaml \ + policy.model_name="${NEMO_RL_MODEL}" \ + policy.tokenizer.name="${NEMO_RL_MODEL}" \ + policy.tokenizer.chat_template="${NEMO_RL_CHAT_TEMPLATE}" \ + policy.generation.vllm_cfg.http_server_serving_chat_kwargs.chat_template="${NEMO_RL_CHAT_TEMPLATE}" \ + grpo.max_num_steps=-1 \ + grpo.max_num_epochs=1000000 \ + grpo.seq_logprob_error_threshold=null \ + grpo.num_prompts_per_step=4 \ + grpo.num_generations_per_prompt=16 \ + grpo.async_grpo.enabled=true \ + grpo.async_grpo.max_trajectory_age_steps=1 \ + grpo.async_grpo.in_flight_weight_updates=true \ + policy.train_global_batch_size=64 \ + policy.max_total_sequence_length=32768 \ + policy.megatron_cfg.tensor_model_parallel_size=4 \ + policy.megatron_cfg.expert_model_parallel_size=4 \ + policy.megatron_cfg.moe_shared_expert_overlap=false \ + policy.megatron_cfg.optimizer.optimizer_cpu_offload=true \ + policy.megatron_cfg.optimizer.optimizer_offload_fraction=1.0 \ + policy.generation.max_new_tokens=16000 \ + policy.generation.vllm_cfg.tensor_parallel_size=4 \ + policy.generation.vllm_cfg.max_model_len=32768 \ + policy.generation.vllm_kwargs.max_num_batched_tokens=32768 \ + policy.generation.vllm_kwargs.max_num_seqs=1 \ + policy.generation.colocated.enabled=false \ + policy.generation.colocated.resources.num_nodes=14 \ + cluster.num_nodes=16 \ + cluster.gpus_per_node=8 \ + checkpointing.enabled=true \ + checkpointing.checkpoint_dir="${CHECKPOINT_DIR}" \ + logger.log_dir="${NEMO_RL_RUN_ROOT}/${RUN_NAME}/training" \ + logger.wandb_enabled=true \ + logger.wandb.project="${WANDB_PROJECT}" \ + +logger.wandb.entity="${WANDB_ENTITY}" \ + logger.wandb.name="${RUN_NAME}" +COMMAND_EOF +export COMMAND + +cd "${REPO_ROOT}" +sbatch \ + --account="${SBATCH_ACCOUNT}" \ + --partition="${SBATCH_PARTITION}" \ + --nodes=16 \ + --ntasks=16 \ + --ntasks-per-node=1 \ + --gpus-per-node=8 \ + --time="${SBATCH_TIME:-04:00:00}" \ + --job-name="rl_main_vg16_async_omni_video_${SBATCH_ACCOUNT}" \ + --output="${NEMO_RL_RUN_ROOT}/slurm-16n-async-%j.out" \ + ray.sub diff --git a/ehsan_scripts/run_2n_sync.sh b/ehsan_scripts/run_2n_sync.sh new file mode 100755 index 00000000000..c16a02ec4c0 --- /dev/null +++ b/ehsan_scripts/run_2n_sync.sh @@ -0,0 +1,235 @@ +#!/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 -euo pipefail + +readonly SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +readonly REPO_ROOT="$(git -C "${SCRIPT_DIR}" rev-parse --show-toplevel)" +readonly SBATCH_ACCOUNT="nemotron_edge_omni" +readonly SBATCH_PARTITION="interactive" + +require_env() { + local name="$1" + if [[ -z "${!name:-}" ]]; then + echo "Required environment variable is unset: ${name}" >&2 + exit 2 + fi +} + +for name in \ + CONTAINER \ + MOUNTS \ + NEMO_RL_CHAT_TEMPLATE \ + NEMO_RL_MODEL \ + NEMO_RL_VIDEO_TRAIN_JSONL \ + NEMO_RL_VIDEO_VAL_JSONL \ + NEMO_RL_VIDEO_MEDIA_ROOT \ + NEMO_RL_RUN_ROOT \ + WANDB_API_KEY \ + WANDB_ENTITY \ + WANDB_PROJECT; do + require_env "${name}" +done + +mkdir -p "${NEMO_RL_RUN_ROOT}" + +NEMO_RL_RUN_ID="${NEMO_RL_RUN_ID:-$(date -u +%Y%m%dT%H%M%SZ)}" +NEMO_RL_RUN_DATE="${NEMO_RL_RUN_DATE:-$(date -u +%m%d)}" + +export BASE_LOG_DIR="${NEMO_RL_RUN_ROOT}" +export CONTAINER +export GPUS_PER_NODE=8 +export MOUNTS +export NEMO_RL_CHAT_TEMPLATE +export NEMO_RL_EXPECTED_COMMIT="$(git -C "${REPO_ROOT}" rev-parse HEAD)" +export NEMO_RL_EXPECTED_GYM_COMMIT="$( + git -C "${REPO_ROOT}" ls-tree HEAD 3rdparty/Gym-workspace/Gym | awk '{print $3}' +)" +export NEMO_RL_EXPECTED_MBRIDGE_COMMIT="$( + git -C "${REPO_ROOT}" ls-tree HEAD 3rdparty/Megatron-Bridge-workspace/Megatron-Bridge | awk '{print $3}' +)" +export NEMO_RL_MODEL +export NEMO_RL_PYTHON_VERSION="$(tr -d '[:space:]' < "${REPO_ROOT}/.python-version")" +export NEMO_RL_REPO="${REPO_ROOT}" +export NEMO_RL_RUN_ID +export NEMO_RL_RUN_DATE +export NEMO_RL_RUN_ROOT +export NEMO_RL_VIDEO_MEDIA_ROOT +export NEMO_RL_VIDEO_TRAIN_JSONL +export NEMO_RL_VIDEO_VAL_JSONL +export NRL_RUN_PREFIX="rl_main_vg2_sync_video_nemotron_omni_model_unlimited" +export RAY_TMPDIR=/tmp/ray +export NEMO_RL_UV_VERSION="$(awk -F= '/^ARG UV_VERSION=/{print $2; exit}' "${REPO_ROOT}/docker/Dockerfile")" +export NEMO_RL_UV_DIR="${NEMO_RL_RUN_ROOT}/uv/${NEMO_RL_UV_VERSION}" +export NEMO_RL_MAIN_VENV="/tmp/nemorl-main-${NEMO_RL_EXPECTED_COMMIT:0:12}" +export NEMO_RL_UV_CACHE_DIR="/tmp/nemorl-uv-cache-${NEMO_RL_EXPECTED_COMMIT:0:12}" +export NEMO_RL_RAY_WRAPPER_DIR="${NEMO_RL_RUN_ROOT}/runtime/${NEMO_RL_EXPECTED_COMMIT:0:12}/bin" +export UV_PYTHON_INSTALL_DIR="${NEMO_RL_RUN_ROOT}/uv_python/${NEMO_RL_PYTHON_VERSION}" +export WANDB_API_KEY +export WANDB_ENTITY +export WANDB_PROJECT +mkdir -p "${NEMO_RL_RAY_WRAPPER_DIR}" +readonly RAY_WRAPPER="${NEMO_RL_RAY_WRAPPER_DIR}/ray" +readonly RAY_WRAPPER_TMP="${RAY_WRAPPER}.tmp.$$" +printf '%s\n' \ + '#!/bin/bash' \ + 'set -euo pipefail' \ + 'if [[ -x "${NEMO_RL_MAIN_VENV}/bin/ray" ]]; then' \ + ' exec "${NEMO_RL_MAIN_VENV}/bin/ray" "$@"' \ + 'fi' \ + 'exec /opt/nemo_rl_venv/bin/ray.nemorl-base "$@"' \ + > "${RAY_WRAPPER_TMP}" +chmod 755 "${RAY_WRAPPER_TMP}" +mv -f "${RAY_WRAPPER_TMP}" "${RAY_WRAPPER}" +export PATH="${NEMO_RL_RAY_WRAPPER_DIR}:${NEMO_RL_MAIN_VENV}/bin:${NEMO_RL_UV_DIR}:${PATH}" +export VIRTUAL_ENV="${NEMO_RL_MAIN_VENV}" +export UV_PROJECT_ENVIRONMENT="${NEMO_RL_MAIN_VENV}" + +read -r -d '' SETUP_COMMAND <<'SETUP_EOF' || true +set -euo pipefail + +source /opt/rl_main_vg_runtime.env +mkdir -p "${NEMO_RL_UV_DIR}" +flock "${NEMO_RL_UV_DIR}.lock" bash -c ' + set -euo pipefail + if [[ ! -x "${NEMO_RL_UV_DIR}/uv" ]]; then + curl --retry 3 --retry-delay 2 -LsSf \ + "https://astral.sh/uv/${NEMO_RL_UV_VERSION}/install.sh" | \ + env UV_INSTALL_DIR="${NEMO_RL_UV_DIR}" UV_NO_MODIFY_PATH=1 sh + fi +' +export PATH="${NEMO_RL_UV_DIR}:/root/.local/bin:/opt/nemo_rl_venv/bin:${PATH}" +uv --version +mkdir -p "${UV_PYTHON_INSTALL_DIR}" "${NEMO_RL_UV_CACHE_DIR}" +flock "${UV_PYTHON_INSTALL_DIR}.lock" \ + uv python install "${NEMO_RL_PYTHON_VERSION}" +uv python find "${NEMO_RL_PYTHON_VERSION}" +export VIRTUAL_ENV="${NEMO_RL_MAIN_VENV}" +export UV_PROJECT_ENVIRONMENT="${NEMO_RL_MAIN_VENV}" +UV_CACHE_DIR="${NEMO_RL_UV_CACHE_DIR}" uv sync \ + --directory "${NEMO_RL_REPO}" \ + --locked \ + --no-install-project +readonly CONTAINER_RAY=/opt/nemo_rl_venv/bin/ray +readonly CONTAINER_RAY_BASE=/opt/nemo_rl_venv/bin/ray.nemorl-base +if [[ ! -x "${CONTAINER_RAY_BASE}" ]]; then + cp -p "${CONTAINER_RAY}" "${CONTAINER_RAY_BASE}" +fi +install -m 755 \ + "${NEMO_RL_RAY_WRAPPER_DIR}/ray" \ + "${CONTAINER_RAY}.tmp.$$" +mv -f "${CONTAINER_RAY}.tmp.$$" "${CONTAINER_RAY}" +export PATH="${NEMO_RL_MAIN_VENV}/bin:${NEMO_RL_UV_DIR}:/root/.local/bin:${PATH}" +python -c \ + 'import sys; assert sys.version_info[:3] == tuple(map(int, sys.argv[1].split("."))), sys.version' \ + "${NEMO_RL_PYTHON_VERSION}" +ray --version +readonly VLLM_OVERLAY="/tmp/nemorl-stock-vllm-0.25.1-${NEMO_RL_EXPECTED_COMMIT:0:12}" +readonly VLLM_WHEEL="https://github.com/vllm-project/vllm/releases/download/v0.25.1/vllm-0.25.1-cp38-abi3-manylinux_2_28_x86_64.whl" +mkdir -p "${VLLM_OVERLAY}" +UV_CACHE_DIR="${NEMO_RL_UV_CACHE_DIR}" uv pip install \ + --target "${VLLM_OVERLAY}" \ + --reinstall \ + --no-deps \ + "vllm @ ${VLLM_WHEEL}" +SETUP_EOF +export SETUP_COMMAND + +read -r -d '' COMMAND <<'COMMAND_EOF' || true +set -euo pipefail + +source /opt/rl_main_vg_runtime.env +export NEMO_GYM_VENV_DIR="/tmp/nemorl-gym-${NEMO_RL_EXPECTED_COMMIT:0:12}-${NEMO_RL_PYTHON_VERSION}-${NEMO_RL_RUN_ID}" +export PATH="${NEMO_RL_MAIN_VENV}/bin:${NEMO_RL_UV_DIR}:/root/.local/bin:${PATH}" +export VIRTUAL_ENV="${NEMO_RL_MAIN_VENV}" +export UV_PROJECT_ENVIRONMENT="${NEMO_RL_MAIN_VENV}" +export HF_HOME="${NEMO_RL_RUN_ROOT}/hf_home" +export HF_MODULES_CACHE="${HF_HOME}/modules" +export NRL_MEGATRON_CHECKPOINT_DIR="${HF_HOME}/nemo_rl-${NEMO_RL_EXPECTED_MBRIDGE_COMMIT:0:12}" +export NRL_FORCE_REBUILD_VENVS=true +export NRL_VIDEO_BACKEND=torchcodec +export NRL_VIDEO_SAMPLING_STYLE=nemotron_vl +export NRL_VIDEO_SFT_MAX_FRAMES=32 +export NRL_VIDEO_SFT_MIN_FRAMES=32 +export NRL_VIDEO_TEMPORAL_PATCH_SIZE=2 +export TORCH_CUDA_ARCH_LIST=9.0 +export VLLM_MAMBA_BACKEND=flashinfer +export VLLM_RAY_EXTRA_ENV_VARS_TO_COPY="PYTHONPATH,NEMO_RL_VIDEO_MEDIA_ROOT,NRL_VIDEO_BACKEND,NRL_VIDEO_SAMPLING_STYLE,NRL_VIDEO_TEMPORAL_PATCH_SIZE" +export VLLM_VIDEO_LOADER_BACKEND=nemotron_vl +unset NRL_IGNORE_VERSION_MISMATCH + +readonly VLLM_OVERLAY="/tmp/nemorl-stock-vllm-0.25.1-${NEMO_RL_EXPECTED_COMMIT:0:12}" +export NEMO_GYM_EXTRA_ROOTS="${NEMO_RL_REPO}/3rdparty/Gym-workspace/Gym" +export PYTHONPATH="${VLLM_OVERLAY}:${NEMO_RL_REPO}:${NEMO_GYM_EXTRA_ROOTS}:${NEMO_RL_REPO}/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge/src:${NEMO_RL_REPO}/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge/3rdparty/Megatron-LM:${HF_MODULES_CACHE}${PYTHONPATH:+:${PYTHONPATH}}" + +readonly RUN_NAME="${NRL_RUN_PREFIX}_${NEMO_RL_RUN_DATE}_${NEMO_RL_RUN_ID}" +readonly CHECKPOINT_DIR="${NEMO_RL_RUN_ROOT}/checkpoints/${RUN_NAME}" +mkdir -p "${CHECKPOINT_DIR}" "${HF_MODULES_CACHE}" "${NRL_MEGATRON_CHECKPOINT_DIR}" +cd "${NEMO_RL_REPO}" +test "$(git rev-parse HEAD)" = "${NEMO_RL_EXPECTED_COMMIT}" +test "$(git -C 3rdparty/Gym-workspace/Gym rev-parse HEAD)" = "${NEMO_RL_EXPECTED_GYM_COMMIT}" +test "$(git -C 3rdparty/Megatron-Bridge-workspace/Megatron-Bridge rev-parse HEAD)" = "${NEMO_RL_EXPECTED_MBRIDGE_COMMIT}" +python -c \ + 'import sys; assert sys.version_info[:3] == tuple(map(int, sys.argv[1].split("."))), sys.version' \ + "${NEMO_RL_PYTHON_VERSION}" + +python \ + examples/nemo_gym/run_grpo_nemo_gym.py \ + --config examples/nemo_gym/grpo_nemotron_omni_30ba3b_video_sync.yaml \ + policy.model_name="${NEMO_RL_MODEL}" \ + policy.tokenizer.name="${NEMO_RL_MODEL}" \ + policy.tokenizer.chat_template="${NEMO_RL_CHAT_TEMPLATE}" \ + policy.generation.vllm_cfg.http_server_serving_chat_kwargs.chat_template="${NEMO_RL_CHAT_TEMPLATE}" \ + grpo.max_num_steps=-1 \ + grpo.max_num_epochs=1000000 \ + grpo.seq_logprob_error_threshold=null \ + grpo.num_prompts_per_step=2 \ + grpo.num_generations_per_prompt=2 \ + policy.train_global_batch_size=4 \ + policy.max_total_sequence_length=32768 \ + policy.megatron_cfg.tensor_model_parallel_size=4 \ + policy.megatron_cfg.expert_model_parallel_size=4 \ + policy.generation.max_new_tokens=16000 \ + policy.generation.vllm_cfg.tensor_parallel_size=4 \ + policy.generation.vllm_cfg.max_model_len=32768 \ + policy.generation.vllm_kwargs.max_num_batched_tokens=32768 \ + policy.generation.vllm_kwargs.max_num_seqs=1 \ + policy.generation.colocated.enabled=true \ + policy.generation.colocated.resources.num_nodes=2 \ + cluster.num_nodes=2 \ + cluster.gpus_per_node=8 \ + checkpointing.enabled=true \ + checkpointing.checkpoint_dir="${CHECKPOINT_DIR}" \ + logger.log_dir="${NEMO_RL_RUN_ROOT}/${RUN_NAME}/training" \ + logger.wandb_enabled=true \ + logger.wandb.project="${WANDB_PROJECT}" \ + +logger.wandb.entity="${WANDB_ENTITY}" \ + logger.wandb.name="${RUN_NAME}" +COMMAND_EOF +export COMMAND + +cd "${REPO_ROOT}" +sbatch \ + --account="${SBATCH_ACCOUNT}" \ + --partition="${SBATCH_PARTITION}" \ + --nodes=2 \ + --ntasks=2 \ + --ntasks-per-node=1 \ + --gpus-per-node=8 \ + --time="${SBATCH_TIME:-04:00:00}" \ + --job-name=rl_main_vg2_sync_omni_video \ + --output="${NEMO_RL_RUN_ROOT}/slurm-2n-sync-%j.out" \ + ray.sub diff --git a/examples/configs/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml index 062abddaa4c..5c2cc719435 100644 --- a/examples/configs/grpo_math_1B.yaml +++ b/examples/configs/grpo_math_1B.yaml @@ -27,6 +27,7 @@ grpo: use_dynamic_sampling: false dynamic_sampling_max_gen_batches: 10 batch_multiplier: 1 + deduplicate_multimodal_data: false reward_shaping: enabled: false overlong_buffer_length: 128 diff --git a/examples/nemo_gym/grpo_nemotron_omni_30ba3b_video_async.yaml b/examples/nemo_gym/grpo_nemotron_omni_30ba3b_video_async.yaml new file mode 100644 index 00000000000..cb8c7ffaf2c --- /dev/null +++ b/examples/nemo_gym/grpo_nemotron_omni_30ba3b_video_async.yaml @@ -0,0 +1,177 @@ +# Static-video, one-turn async GRPO for Nemotron Omni on 16 nodes. +# +# Prepare JSONL rows with: +# python examples/nemo_gym/prepare_video_dataset.py convert \ +# --input source.jsonl --output video-gym.jsonl --require-video \ +# --skip-missing-local-videos +# Export NEMO_RL_VIDEO_TRAIN_JSONL, NEMO_RL_VIDEO_VAL_JSONL, and the common +# NEMO_RL_VIDEO_MEDIA_ROOT allowed for vLLM file-URL access before launch. +# Policy-side preprocessing also requires NRL_VIDEO_BACKEND=torchcodec and +# NRL_VIDEO_SAMPLING_STYLE=nemotron_vl to match the validated rollout contract. + +# Keep the video recipe on the maintained NeMo-Gym base. The former +# Nemotron-Omni Gym-V smoke overlay was scratch configuration and was removed +# from the integration branch. +defaults: grpo_nanov3.yaml + +grpo: + # Keep the initial rollout below the batch scheduler's idle-policy reaper + # window while retaining 16 generations for each GRPO prompt. + num_prompts_per_step: 4 + num_generations_per_prompt: 16 + num_val_generations_per_prompt: 1 + # Run until the configured dataset epochs are exhausted. + max_num_steps: -1 + max_rollout_turns: 1 + val_period: 500 + val_batch_size: null + # Report raw logprob error without masking or dropping high-error samples. + seq_logprob_error_threshold: null + deduplicate_multimodal_data: true + async_grpo: + enabled: true + max_trajectory_age_steps: 1 + in_flight_weight_updates: true + +policy: + model_name: "nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16" + is_vlm: true + train_global_batch_size: 64 + max_total_sequence_length: 24576 + logprob_chunk_size: 2048 + refit_buffer_size_gb: 1 + # Multimodal payloads are already deduplicated per prompt. Keep the source + # video recipe's conservative unpacked path until packed video training is + # validated independently. + sequence_packing: + enabled: false + tokenizer: + name: ${policy.model_name} + # Declare the key so launchers can override it with a local .jinja file. + # "default" preserves the model-bundled template when no override is used. + chat_template: default + chat_template_kwargs: + enable_thinking: true + truncate_history_thinking: false + video: + num_frames: 32 + megatron_cfg: + # H100 nodes use CUDA compute capability 9.0. Megatron requires this + # explicitly when the runtime container does not provide it. + env_vars: + TORCH_CUDA_ARCH_LIST: "9.0" + # Frozen RADIO features keep vLLM's multimodal cache valid across refits. + freeze_vision_model: true + freeze_vision_projection: true + # Match the validated source run: the checkpoint can contain MTP weights, + # but GRPO trains only the next-token objective here. + mtp_num_layers: 0 + mtp_use_repeated_layer: true + mtp_detach_heads: true + mtp_loss_scaling_factor: 0.0 + tensor_model_parallel_size: 4 + pipeline_model_parallel_size: 1 + # Reuse the two-node policy topology verified by the synchronous control. + # Fourteen generation nodes leave only 12.5% of the allocation idle while + # an uncapped long rollout is still in flight. + expert_model_parallel_size: 4 + context_parallel_size: 1 + moe_shared_expert_overlap: false + optimizer: + optimizer_cpu_offload: true + optimizer_offload_fraction: 1.0 + radio_force_cpe_eval_mode: true + clear_memory_caches_before_refit: true + distributed_data_parallel_config: + overlap_grad_reduce: false + overlap_param_gather: false + scheduler: + lr_warmup_iters: 0 + # Keep sequence padding consistent with the validated Omni configuration. + make_sequence_length_divisible_by: 32 + generation: + max_new_tokens: ${policy.max_total_sequence_length} + bad_words: [] + vllm_cfg: + # Both vision modules are frozen, so refits do not invalidate encoder + # outputs and in-flight async requests can safely remain active. + reset_mm_cache_after_refit: false + logprobs_mode: raw_logprobs + # Match the TP4 policy topology. TP8 vLLM inference can diverge from the + # TP4 Megatron logprobs on rare tokens, while TP4 preserves raw-logprob + # parity without masking or dropping samples. + tensor_parallel_size: 4 + gpu_memory_utilization: 0.6 + max_model_len: ${policy.max_total_sequence_length} + enforce_eager: true + enable_prefix_caching: false + skip_tokenizer_init: false + env_vars: + # Decode raw video bytes with the same TorchCodec/Nemotron sampling + # contract used by policy-side preprocessing. + NRL_VIDEO_BACKEND: "torchcodec" + NRL_VIDEO_SAMPLING_STYLE: "nemotron_vl" + # Keep recovered short-video frame counts divisible by RADIO's temporal + # tubelet size in both rollout and policy preprocessing. + NRL_VIDEO_TEMPORAL_PATCH_SIZE: "2" + VLLM_VIDEO_LOADER_BACKEND: "nemotron_vl" + http_server_serving_chat_kwargs: + # None delegates to the initialized tokenizer's template. Declaring the + # key also permits a normal Hydra override for local model checkpoints. + chat_template: null + chat_template_content_format: string + enable_auto_tools: true + tool_parser: qwen3_coder + reasoning_parser: nemotron_v3 + vllm_kwargs: + allowed_local_media_path: "${oc.env:NEMO_RL_VIDEO_MEDIA_ROOT}" + mm_processor_cache_gb: 0 + max_num_seqs: 1 + limit_mm_per_prompt: + # Cached frames are represented as one native video for vLLM while + # policy preprocessing consumes the same lossless PNG sequence. + image: 32 + video: + count: 1 + num_frames: 32 + max_num_batched_tokens: ${policy.max_total_sequence_length} + enable_chunked_prefill: false + disable_custom_all_reduce: true + attention_backend: FLASH_ATTN + attention_config: + use_trtllm_attention: false + colocated: + enabled: false + resources: + gpus_per_node: 8 + num_nodes: 14 + +data: + max_input_seq_length: ${policy.max_total_sequence_length} + # Video decoding imports the pinned vLLM media loader; keep it in the driver + # instead of importing vLLM for the first time after a DataLoader fork. + num_workers: 0 + default: + num_frames: 32 + video_temporal_patch_size: 2 + video_target_num_patches: 1024 + video_maintain_aspect_ratio: true + min_generation_tokens: 2000 + train: + data_path: "${oc.env:NEMO_RL_VIDEO_TRAIN_JSONL}" + validation: + data_path: "${oc.env:NEMO_RL_VIDEO_VAL_JSONL}" + +env: + should_log_nemo_gym_responses: false + nemo_gym: + skip_venv_if_present: true + # The source blend's video-only subset is multiple-choice, so only the + # policy model and MCQA grader are required. + config_paths: + - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml + - resources_servers/mcqa/configs/mcqa.yaml + +cluster: + gpus_per_node: 8 + num_nodes: 16 diff --git a/examples/nemo_gym/grpo_nemotron_omni_30ba3b_video_sync.yaml b/examples/nemo_gym/grpo_nemotron_omni_30ba3b_video_sync.yaml new file mode 100644 index 00000000000..0ecdcea2534 --- /dev/null +++ b/examples/nemo_gym/grpo_nemotron_omni_30ba3b_video_sync.yaml @@ -0,0 +1,62 @@ +# 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. + +# Two-node synchronous Gym video GRPO validation overlay. +defaults: grpo_nemotron_omni_30ba3b_video_async.yaml + +grpo: + num_prompts_per_step: 2 + num_generations_per_prompt: 2 + # Run until the configured dataset epochs are exhausted. + max_num_steps: -1 + async_grpo: + enabled: false + in_flight_weight_updates: false + +loss_fn: + use_importance_sampling_correction: false + +policy: + train_global_batch_size: 4 + max_total_sequence_length: 16384 + refit_buffer_size_gb: 5 + sequence_packing: + enabled: false + megatron_cfg: + # H100 nodes use CUDA compute capability 9.0. Megatron requires this + # explicitly when the runtime container does not provide it. + env_vars: + TORCH_CUDA_ARCH_LIST: "9.0" + tensor_model_parallel_size: 4 + expert_model_parallel_size: 4 + context_parallel_size: 1 + moe_shared_expert_overlap: false + optimizer: + optimizer_cpu_offload: true + optimizer_offload_fraction: 1.0 + generation: + max_new_tokens: ${policy.max_total_sequence_length} + vllm_cfg: + # Match the TP4 policy topology so raw generation and training logprobs + # use the validated numerical decomposition. + tensor_parallel_size: 4 + max_model_len: ${policy.max_total_sequence_length} + vllm_kwargs: + max_num_batched_tokens: ${policy.max_total_sequence_length} + colocated: + enabled: true + +cluster: + gpus_per_node: 8 + num_nodes: 2 diff --git a/examples/nemo_gym/prepare_video_dataset.py b/examples/nemo_gym/prepare_video_dataset.py new file mode 100644 index 00000000000..36301d6f49a --- /dev/null +++ b/examples/nemo_gym/prepare_video_dataset.py @@ -0,0 +1,295 @@ +# 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. +"""Convert and validate one-video-per-row NeMo-Gym JSONL datasets.""" + +import argparse +import json +import re +from collections import Counter +from pathlib import Path +from typing import Any +from urllib.parse import unquote, urlparse + + +VERIFIER_TO_AGENT = { + "mcqa": "mcqa_simple_agent", + "multiple-choice": "mcqa_simple_agent", + "multiple_choice": "mcqa_simple_agent", + "mathruler": "math_with_judge_simple_agent", + "math-with-judge": "math_with_judge_simple_agent", + "math_with_judge": "math_with_judge_simple_agent", +} +OPTION_RE = re.compile( + r"(?:^|\n|\s)(?:\(([A-Ja-j])\)|([A-Ja-j])[.)::])\s+(.+?)" + r"(?=(?:\s|\n)(?:\([A-Ja-j]\)|[A-Ja-j][.)::])\s+|\Z)", + re.S, +) + + +def _read_jsonl(path: Path): + with path.open(encoding="utf-8") as input_file: + for line_number, line in enumerate(input_file, start=1): + if line.strip(): + yield line_number, json.loads(line) + + +def _video_parts(row: dict[str, Any]) -> list[dict[str, Any]]: + parts = [] + for message in row.get("responses_create_params", {}).get("input", []): + content = message.get("content", []) if isinstance(message, dict) else [] + if not isinstance(content, list): + continue + parts.extend( + part + for part in content + if isinstance(part, dict) + and part.get("type") in ("input_video", "video", "video_url") + ) + return parts + + +def _part_source(part: dict[str, Any]) -> str: + value = part.get("video_url") or part.get("video") or part.get("url") + if isinstance(value, dict): + value = value.get("url") or value.get("path") + return value if isinstance(value, str) else "" + + +def _as_list(value: Any) -> list[Any]: + if value is None: + return [] + if isinstance(value, list): + return [item for item in value if item] + return [value] if value else [] + + +def _clean_question(question: str) -> str: + question = question.strip() + for token in ("", "