diff --git a/3rdparty/Gym-workspace/Gym b/3rdparty/Gym-workspace/Gym index bb9b26de3da..c3bac96314a 160000 --- a/3rdparty/Gym-workspace/Gym +++ b/3rdparty/Gym-workspace/Gym @@ -1 +1 @@ -Subproject commit bb9b26de3dac1294faccccb5ed70f3386c16de10 +Subproject commit c3bac96314a59f28b896f597eb9845d175bb0252 diff --git a/docs/guides/models/nemotron/nemotron-3-nano-omni.md b/docs/guides/models/nemotron/nemotron-3-nano-omni.md index 1d8e1528cc5..d14a958b948 100644 --- a/docs/guides/models/nemotron/nemotron-3-nano-omni.md +++ b/docs/guides/models/nemotron/nemotron-3-nano-omni.md @@ -132,7 +132,9 @@ The Megatron backend uses a dedicated `NemotronOmniModel` supplied by Megatron B This is the same model-owned packing boundary used by maintained Megatron VLM integrations. It differs from the historical Nemotron Omni `LLaVAModel` path, which collapsed the expanded media-token sequence before packing and expanded it again inside the model. The dedicated model removes that extra representation change and allows the integration to use Megatron Bridge and Megatron-LM from their maintained main branches. -The current Megatron recipes cover Nano image-and-text GRPO. Super, video, and audio training are follow-up work and are not enabled by these recipes. +The maintained Megatron VLM recipes cover Nano image-and-text GRPO. The NeMo +Gym integration also supports the static, one-video-per-row workflow described +below. Audio and mixed audio-video rows are not supported by that workflow. ### Checkpoint compatibility @@ -159,3 +161,74 @@ CONFIG_PATH=examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g ``` The recipes keep sequence packing enabled because the model owns the packing step after multimodal embedding insertion. They also request raw generation log probabilities so that vLLM and the Megatron policy compare the same pre-processor probability values when generation constraints such as `bad_words` are active. The generation context cap prevents the processor-expanded image prompt plus generated response from exceeding the configured 8192-token context length. + +## NeMo Gym video GRPO + +Nemotron 3 Nano Omni video GRPO accepts one local video in the initial static +NeMo Gym prompt. The user message uses an `input_video` part followed by +`input_text`: + +```json +{ + "responses_create_params": { + "input": [{ + "role": "user", + "content": [ + {"type": "input_video", "video_url": "/absolute/path/clip.mp4"}, + {"type": "input_text", "text": "Which option describes the clip?"} + ] + }] + } +} +``` + +Only local paths under `policy.generation.vllm_kwargs.allowed_local_media_path` +are accepted. Environment-produced video and video introduced in a later +trajectory turn are not supported. Each row must contain exactly one video and no audio. MCQA rows +must provide real top-level `options` and one uppercase `expected_answer` +present in those options. Use +[`prepare_video_dataset.py`](../../../../examples/nemo_gym/prepare_video_dataset.py) +to convert and validate source JSONL. Raw videos are decoded with TorchCodec; +the runtime can also read externally prepared lossless frame manifests, but the +converter currently emits raw-video rows only. + +The recipes contain explicit placeholder paths. Override them when launching: + +```bash +uv run examples/nemo_gym/run_grpo_nemo_gym.py \ + --config examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-2n8g-megatron-tp4ep4-gym-video.v1.yaml \ + policy.generation.vllm_kwargs.allowed_local_media_path=/path/to/video-data \ + data.train.data_path=/path/to/train.jsonl \ + data.validation.data_path=/path/to/validation.jsonl +``` + +Policy and rollout preprocessing must use the same sampling contract. The +provided recipes set one `policy.generation.vllm_cfg.video` block for +TorchCodec-backed Nemotron sampling, 32 frames, and a temporal patch size of 2. +NeMo RL materializes those values for both policy and rollout preprocessing. +The runtime therefore needs TorchCodec and its FFmpeg dependencies; no parallel +sampling environment variables are required. + +The remaining policy-side video settings live under `data.default`: +`video_target_num_patches`, `video_maintain_aspect_ratio`, and, for generic +processors only, `min_generation_tokens`. Set +`policy.generation.vllm_cfg.reset_encoder_cache_after_weight_update: true` only +when the multimodal encoder trains; the provided frozen-vision recipes leave it +disabled. + +The synchronous and asynchronous overlays are: + +- [2-node synchronous recipe](../../../../examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-2n8g-megatron-tp4ep4-gym-video.v1.yaml) +- [16-node asynchronous recipe](../../../../examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-16n8g-megatron-tp4ep4-async-gym-video.v1.yaml) + +They require the corresponding Nemotron Omni support in Megatron Bridge and +video request/token propagation in NeMo Gym. Policy preprocessing numerically +matches unmodified stock vLLM 0.25.1; a custom vLLM fork is not required. + +Both overlays use a large positive `grpo.max_num_steps` value so it does not +bind normal training. Training still follows the existing GRPO step-limit +semantics. The recipes also leave +`grpo.seq_logprob_error_threshold: null`: logged token multiplicative +probability error (TMPE) is raw and no high-error sequence is masked. Interpret +TMPE together with reward, loss, sequence length, and refit metrics; isolated +maxima are less informative than a sustained shift in the distribution. 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 new file mode 100644 index 00000000000..2647cfc4d8a --- /dev/null +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-16n8g-megatron-tp4ep4-async-gym-video.v1.yaml @@ -0,0 +1,100 @@ +defaults: ../../../nemo_gym/grpo_nanov3.yaml +grpo: + num_prompts_per_step: 4 + val_num_generations_per_prompt: 1 + val_period: 500 + val_batch_size: null + seq_logprob_error_threshold: null + deduplicate_multimodal_data: true + async_grpo: + enabled: true + 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 + refit_buffer_size_gb: 1 + sequence_packing: + enabled: false + tokenizer: + chat_template: default + chat_template_kwargs: + enable_thinking: true + truncate_history_thinking: false + megatron_cfg: + env_vars: + TORCH_CUDA_ARCH_LIST: '9.0' + freeze_vision_model: true + freeze_vision_projection: true + 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 + expert_model_parallel_size: 4 + context_parallel_size: 1 + 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 + make_sequence_length_divisible_by: 32 + generation: + bad_words: [] + vllm_cfg: + reset_encoder_cache_after_weight_update: false + video: + sampling_style: nemotron_vl + num_frames: 32 + temporal_patch_size: 2 + logprobs_mode: raw_logprobs + gpu_memory_utilization: 0.6 + enforce_eager: true + enable_prefix_caching: false + skip_tokenizer_init: false + http_server_serving_chat_kwargs: + chat_template: null + chat_template_content_format: string + reasoning_parser: nemotron_v3 + vllm_kwargs: + allowed_local_media_path: /path/to/video-data + mm_processor_cache_gb: 0 + max_num_seqs: 1 + limit_mm_per_prompt: + 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} + num_workers: 0 + default: + video_target_num_patches: 1024 + video_maintain_aspect_ratio: true +env: + should_log_nemo_gym_responses: false + nemo_gym: + skip_venv_if_present: true + config_paths: + - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml + - resources_servers/mcqa/configs/mcqa.yaml +cluster: + num_nodes: 16 diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-2n8g-megatron-tp4ep4-gym-video.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-2n8g-megatron-tp4ep4-gym-video.v1.yaml new file mode 100644 index 00000000000..3d3d5307f5c --- /dev/null +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-2n8g-megatron-tp4ep4-gym-video.v1.yaml @@ -0,0 +1,18 @@ +defaults: vlm_grpo-nemotron-omni-30ba3b-16n8g-megatron-tp4ep4-async-gym-video.v1.yaml +grpo: + num_prompts_per_step: 2 + num_generations_per_prompt: 2 + 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 + generation: + colocated: + enabled: true +cluster: + 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..7d18a3e0685 --- /dev/null +++ b/examples/nemo_gym/prepare_video_dataset.py @@ -0,0 +1,350 @@ +# 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 ("", "