From b7fa6e330f8aabcdb140f2dbdcdc7f92609ddffa Mon Sep 17 00:00:00 2001 From: rohitrango Date: Mon, 3 Aug 2026 14:07:24 -0700 Subject: [PATCH 01/27] feat(nemo-gym): support multimodal rollouts with tokenizer_config plumbing and chat_template parity Folds three WIP commits (b16fca25 + 5486b57e + cd349d05): - Add tokenizer_config field to NemoGymConfig and processor reconstruction inside the actor for multi-turn multimodal postprocessing. - Add multimodal utilities (encode_images_in_examples, extract_multimodal_model_inputs, process_multimodal_chat, resolve_to_image) and their consumers in NemoGym. - Add sync single-step polish across nemo_gym.py / multimodal_utils.py / processors.py. - Maintain chat_template kwargs parity between the async vLLM worker and HF. Signed-off-by: rohitrango (cherry picked from commit 3215893f10aaf5b466435ef22e30d96800a11db6) Signed-off-by: rohitrango --- examples/nemo_gym/grpo_nanov3.yaml | 2 +- .../grpo_nemotron_omni_30ba3b_gymv_smoke.yaml | 187 ++++++++++ ...ron_omni_30ba3b_gymv_smoke_async_1off.yaml | 22 ++ .../nemo_gym/grpo_qwen25vl_gymv_smoke.yaml | 309 ++++++++++++++++ examples/nemo_gym/run_gymv_smoke.sh | 96 +++++ .../nemo_gym/run_multimodal_grpo_nemo_gym.py | 331 ++++++++++++++++++ nemo_rl/data/multimodal_utils.py | 192 +++++++++- nemo_rl/data/processors.py | 74 +--- nemo_rl/environments/nemo_gym.py | 162 ++++++++- .../generation/vllm/vllm_worker_async.py | 19 + nemo_rl/models/megatron/setup.py | 9 + nemo_rl/models/policy/__init__.py | 2 + nemo_rl/utils/packed_tensor.py | 24 +- pyproject.toml | 22 +- .../data/test_multimodal_processor_adapter.py | 177 ++++++++++ 15 files changed, 1546 insertions(+), 82 deletions(-) create mode 100644 examples/nemo_gym/grpo_nemotron_omni_30ba3b_gymv_smoke.yaml create mode 100644 examples/nemo_gym/grpo_nemotron_omni_30ba3b_gymv_smoke_async_1off.yaml create mode 100644 examples/nemo_gym/grpo_qwen25vl_gymv_smoke.yaml create mode 100755 examples/nemo_gym/run_gymv_smoke.sh create mode 100644 examples/nemo_gym/run_multimodal_grpo_nemo_gym.py create mode 100644 tests/unit/data/test_multimodal_processor_adapter.py diff --git a/examples/nemo_gym/grpo_nanov3.yaml b/examples/nemo_gym/grpo_nanov3.yaml index 8fa3175d9e4..b2073c528e3 100644 --- a/examples/nemo_gym/grpo_nanov3.yaml +++ b/examples/nemo_gym/grpo_nanov3.yaml @@ -237,7 +237,7 @@ policy: block_size_tokens: 256 # Size of each KV cache block in tokens (affects memory granularity) use_cuda_graphs_for_non_decode_steps: true # Enable CUDA graphs for prefill/context processing enable_chunked_prefill: true - max_tokens: 16384 # Maximum number of tokens to use in a single step. Analogous to vllm's max_num_batched_tokens + max_tokens: ${policy.max_total_sequence_length} # Maximum number of tokens to use in a single step. Analogous to vllm's max_num_batched_tokens kv_cache_management_mode: "persist" # KV cache lifecycle across suspend/resume. Options: "persist", "offload". To select "recompute", set grpo.async_grpo.recompute_kv_cache_after_weight_updates=true. materialize_only_last_token_logits: true num_speculative_tokens: 0 diff --git a/examples/nemo_gym/grpo_nemotron_omni_30ba3b_gymv_smoke.yaml b/examples/nemo_gym/grpo_nemotron_omni_30ba3b_gymv_smoke.yaml new file mode 100644 index 00000000000..8ca6aa78fe3 --- /dev/null +++ b/examples/nemo_gym/grpo_nemotron_omni_30ba3b_gymv_smoke.yaml @@ -0,0 +1,187 @@ +# Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16 smoke test on gym-v. +# +# Sibling of grpo_qwen25vl_gymv_smoke.yaml. Nemotron-Omni is the known-good +# multimodal + NeMo-Gym baseline in-house (see the games-bandit recipe chain); +# this file adapts it to the current mm-integration branch so the Qwen path +# and the Nemotron-Omni path can be exercised through the same code path. +# +# Both the model AND the env are multimodal here — expect real visual RL +# signal (unlike the earlier text-only Nemotron placeholder). +# +# Shape (2 nodes × 8 GPUs, non-colocated): +# - Node 1: vLLM generation, TP=8 +# - Node 2: Megatron policy training, TP=8, EP=8, PP=1, CP=1 +# +# Data: size=4 FrozenLake ablation manifest (multi-turn train, single-turn +# eval), sourced from nemo-rl-games-bandit's frozenlake_size4.yaml chain. +# +# Run with: +# uv run --locked --extra mcore --extra vllm \ +# examples/nemo_gym/run_multimodal_grpo_nemo_gym.py \ +# --config examples/nemo_gym/grpo_nemotron_omni_30ba3b_gymv_smoke.yaml + +defaults: grpo_nanov3.yaml + +grpo: + num_prompts_per_step: 1 # smoke: minimal loop + num_generations_per_prompt: 4 # gbs = 1 * 4 = 4 + num_val_generations_per_prompt: 1 + max_rollout_turns: 4 + max_num_epochs: 1 + max_num_steps: 20 # + val_period: 500 + val_at_start: false # exercise eval before training + val_at_end: false + max_val_samples: null + val_batch_size: null + async_grpo: + enabled: false # sync for the smoke run + +checkpointing: + enabled: false + checkpoint_dir: "results/grpo-nemotron-omni-30ba3b-gymv-smoke" + +policy: + model_name: "nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16" + is_vlm: true # required by run_multimodal_grpo_nemo_gym.py + tokenizer: + name: ${policy.model_name} + # Nano-Omni-Reasoning uses a native block; keep it open and + # preserve prior reasoning across turns (append-only trajectory rep). + # Mirrors the games-bandit grpo_game_rlvr_trajectory_collection.yaml recipe. + chat_template_kwargs: + enable_thinking: true + truncate_history_thinking: false + train_global_batch_size: ${mul:${grpo.num_prompts_per_step}, ${grpo.num_generations_per_prompt}} + train_micro_batch_size: 1 + logprob_batch_size: 1 + max_total_sequence_length: 4096 + + megatron_cfg: + # Full 8-GPU policy node: TP=8, EP=8, PP=1, CP=1. Overrides + # grpo_nanov3.yaml's TP=2/PP=2/CP=4 layout (tuned for a 32-node run). + tensor_model_parallel_size: 2 + expert_tensor_parallel_size: 1 + expert_model_parallel_size: 8 + pipeline_model_parallel_size: 1 + context_parallel_size: 2 + sequence_parallel: true # sp only pays off with tp + cp + activation_checkpointing: true + # RADIO CPE eval mode: keeps vision-tower positional embeddings in eval + # mode during rollout/train (required for the frozen vision path to + # produce stable features). + radio_force_cpe_eval_mode: true + # Empty the CUDA cache before the vLLM refit broadcast so the packed + # staging tensor (see nemo_rl/utils/packed_tensor.py) doesn't OOM against + # Adam m/v that materialize at the end of iter 1. + clear_memory_caches_before_refit: true + # Nemotron-Omni carries a sound_encoder / sound_projection tower that + # never sees an audio input in gym-v. Those params live in the DDP bucket + # with requires_grad=True but never receive a backward hook, so MCore's + # async grad-reduce path trips the golden-count assertion at + # param_and_grad_buffer.py:272 on iter 2's zero_grad_buffer(). Disable + # the overlap paths until the sound tower is actually frozen via a + # nemo-rl pre_wrap_hook (mirror of freeze_moe_router). This is a real + # perf cost — remove once the freeze path is wired. + distributed_data_parallel_config: + overlap_grad_reduce: false + overlap_param_gather: false + + # Flat LR throughout: skip warmup (base recipe warmed 3e-7 → 3e-6 over 10 + # iters). Combined with lr_decay_style="constant" and min_lr == lr from + # grpo_nanov3.yaml, this pins the LR at 3e-6 from step 0 onward. + scheduler: + lr_warmup_iters: 0 + + # base recipe ties this to dtensor_cfg.tensor_parallel_size; we're on + # megatron, so pin to the megatron TP directly. + make_sequence_length_divisible_by: 32 + + generation: + max_new_tokens: 4096 # per-turn cap; 4 turns fit in 8k context + # Nano-Omni-Reasoning bad_words: only the vision-side text tokens (safe, + # in-vocab). The audio-delimiter tokens (, , + # ) sit past the LM's text-logit width and cause vLLM + # out-of-bounds writes in bad_words masking — see the games-bandit + # vlm_grpo_games.yaml history comment for the full incident. + bad_words: ["", "", ""] + vllm_cfg: + # Full 8-GPU vLLM node. + tensor_parallel_size: 8 + pipeline_parallel_size: 1 + expert_parallel_size: 1 + gpu_memory_utilization: 0.8 + max_model_len: ${policy.max_total_sequence_length} + enforce_eager: true # skip CUDA-graph capture for the smoke run + # Nano-Omni + prefix caching crashes vLLM's mm cache (mm_hash miss). + # Belt-and-suspenders: also disabled via mm_processor_cache_gb=0 below. + enable_prefix_caching: false + # VLMs need the tokenizer initialized (run_multimodal_grpo_nemo_gym.py asserts). + skip_tokenizer_init: false + # Nano-Omni chat template expects string content, not the OpenAI + # {type: text, text: ...} list. Keeps the NeMo-Gym prompt-token prefix + # invariant across multi-turn rollouts. + http_server_serving_chat_kwargs: + enable_auto_tools: true + tool_parser: qwen3_coder + reasoning_parser: nemotron_v3 + # chat_template_content_format: string + vllm_kwargs: + # Bounds per-prompt image count (multi-turn appends one board / turn). + limit_mm_per_prompt: {"image": 8} + # Disable vLLM's mm processor cache — the mm-cache desync crash on + # Nano-Omni is the load-bearing fix. Must sit in vllm_kwargs, not + # vllm_cfg (see vlm_grpo_3B.yaml note). + mm_processor_cache_gb: 0 + max_num_batched_tokens: 16384 + # Nano-Omni's mamba backbone needs SSM cache in fp32 (accuracy). + mamba_ssm_cache_dtype: "float32" + colocated: + enabled: false # non-colocated: dedicated vLLM node + resources: + gpus_per_node: 8 + num_nodes: 2 + +data: + train: + # Sourced from nemo-rl-games-bandit's frozenlake_size4.yaml — the + # size=4 (num_holes=3) FrozenLake ablation train manifest. + data_path: /lustre/fsw/portfolios/coreai/users/rohitkumarj/game-rlvr-vlm-data/frozenlake_ablations/size_4/train_manifest.jsonl + validation: + # Paired eval manifest for the size=4 FrozenLake ablation. Note: rows are + # `Games/FrozenLake-singleturn-v0` with horizon_cap=1 (single-turn eval), + # while train rows are the multi-turn `Games/FrozenLake-v0`. + data_path: /lustre/fsw/portfolios/coreai/users/rohitkumarj/game-rlvr-vlm-data/frozenlake_ablations/size_4/eval_manifest.jsonl + +env: + should_use_nemo_gym: true + nemo_gym: + is_trajectory_collection: false + # Replace the base recipe's math/code/etc bundle with gym-v only. Any + # env-block overrides inherited from grpo_nanov3.yaml (math_with_judge, + # code_gen, workplace_assistant, ...) are ignored by Gym when their + # config_paths aren't loaded. + config_paths: + - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml + - environments/gym_v/config.yaml + # Agent-side horizon; keep in sync with grpo.max_rollout_turns. + # Outer group name matches the manifest's agent_ref.name = "gym_v_agent" + # (see 3rdparty/Gym-workspace/Gym/environments/gym_v/config.yaml). + gym_v_agent: + responses_api_agents: + gymv_agent: + max_steps: 1 + done_if_no_boxed_answer: true + +logger: + log_dir: "logs/grpo-nemotron-omni-30ba3b-gymv-smoke" + wandb_enabled: false + tensorboard_enabled: true + monitor_gpus: false + wandb: + project: "grpo-nemotron-omni-gymv" + name: "grpo-nemotron-omni-30ba3b-gymv-smoke" + +cluster: + gpus_per_node: 8 + num_nodes: 4 diff --git a/examples/nemo_gym/grpo_nemotron_omni_30ba3b_gymv_smoke_async_1off.yaml b/examples/nemo_gym/grpo_nemotron_omni_30ba3b_gymv_smoke_async_1off.yaml new file mode 100644 index 00000000000..db7f6edfc8f --- /dev/null +++ b/examples/nemo_gym/grpo_nemotron_omni_30ba3b_gymv_smoke_async_1off.yaml @@ -0,0 +1,22 @@ +# Async variant of grpo_nemotron_omni_30ba3b_gymv_smoke.yaml. +# +# Extends the sync smoke recipe and flips only the async_grpo block on: +# 1-step trajectory age (mirrors the LLM `-async-1off` recipes under +# examples/configs/recipes/llm/performance/). +# +# Run with: +# uv run --locked --extra mcore --extra vllm \ +# examples/nemo_gym/run_multimodal_grpo_nemo_gym.py \ +# --config examples/nemo_gym/grpo_nemotron_omni_30ba3b_gymv_smoke_async_1off.yaml + +defaults: grpo_nemotron_omni_30ba3b_gymv_smoke.yaml + +grpo: + async_grpo: + enabled: true + max_trajectory_age_steps: 1 + +logger: + log_dir: "logs/grpo-nemotron-omni-30ba3b-gymv-smoke-async" + wandb: + name: "grpo-nemotron-omni-30ba3b-gymv-smoke-async" \ No newline at end of file diff --git a/examples/nemo_gym/grpo_qwen25vl_gymv_smoke.yaml b/examples/nemo_gym/grpo_qwen25vl_gymv_smoke.yaml new file mode 100644 index 00000000000..47d5ad8c1bf --- /dev/null +++ b/examples/nemo_gym/grpo_qwen25vl_gymv_smoke.yaml @@ -0,0 +1,309 @@ +# Qwen2.5-VL-3B smoke test for the multimodal + multi-turn NeMo-Gym integration. +# +# Modeled after nemo-rl-games-bandit's frozenlake_size4.yaml chain (see +# examples/configs/frozenlake_difficulty_ablations/frozenlake_size4.yaml) but +# retargeted to this branch's live gym-v drop under +# `3rdparty/Gym-workspace/Gym/environments/gym_v/` — which bundles both the +# gym_v resources server and the gymv_agent responses-API agent. +# +# Shape (single node, 8 GPUs, non-colocated): +# - vLLM generation: 4 GPUs (TP=4) +# - Megatron policy: 4 GPUs (TP=4) +# +# Data: the bundled `environments/gym_v/data/example.jsonl` (8 rows across +# FrozenLake, GameOfLife, and a few other env families). Used for both train +# and validation so the smoke run exercises the multi-turn rollout + full +# validation path without needing an external manifest. +# +# Run with: +# uv run --locked --extra mcore --extra vllm \ +# examples/nemo_gym/run_multimodal_grpo_nemo_gym.py \ +# --config examples/nemo_gym/grpo_qwen25vl_gymv_smoke.yaml + +grpo: + num_prompts_per_step: 1 # smallest: exercises the loop, minimal cost + num_generations_per_prompt: 4 # gbs = 1 * 4 = 4 rollouts/step + num_val_generations_per_prompt: 1 + # Multi-turn: cap at 4 turns so per-rollout images stay < limit_mm_per_prompt + # and the smoke run finishes quickly even on a slow VLM. + max_rollout_turns: 4 + max_num_epochs: 1 + max_num_steps: 5 # smoke: 5 optimizer steps + normalize_rewards: true + use_leave_one_out_baseline: true + val_period: 5 + val_at_start: true # validate the eval path before training + val_at_end: false + overlong_filtering: false + advantage_clip_low: null + advantage_clip_high: null + max_val_samples: null # inferred from val dataset length + val_batch_size: null # inferred from val dataset length + seed: 42 + use_dynamic_sampling: false + batch_multiplier: 1 + reward_shaping: + enabled: false + reward_scaling: + enabled: false + seq_logprob_error_threshold: null + invalid_tool_call_advantage: null + malformed_thinking_advantage: null + async_grpo: + enabled: false # sync GRPO for the smoke run + max_trajectory_age_steps: 1 + +loss_fn: + reference_policy_kl_penalty: 0 + reference_policy_kl_type: "k3" + kl_input_clamp_value: 20.0 + kl_output_clamp_value: 10.0 + ratio_clip_min: 0.2 + ratio_clip_max: 0.2 + ratio_clip_c: null + use_on_policy_kl_approximation: false + truncated_importance_sampling_ratio: null + use_importance_sampling_correction: false + token_level_loss: true + +checkpointing: + enabled: false # smoke test — no checkpoints + checkpoint_dir: "results/grpo-qwen25vl-gymv-smoke" + metric_name: "val:total_reward/mean" + higher_is_better: true + keep_top_k: 1 + save_period: 5 + checkpoint_must_save_by: null + save_optimizer: false + +policy: + model_name: "Qwen/Qwen2.5-VL-3B-Instruct" + is_vlm: true # required by run_multimodal_grpo_nemo_gym.py + tokenizer: + name: ${policy.model_name} + chat_template_kwargs: null + hf_config_overrides: {} + # gbs must equal num_prompts_per_step * num_generations_per_prompt. + train_global_batch_size: ${mul:${grpo.num_prompts_per_step}, ${grpo.num_generations_per_prompt}} + train_micro_batch_size: 1 + logprob_batch_size: 1 + generation_batch_size: 8 + # Room for a 4-turn multi-image rollout with Qwen2.5-VL's per-image token cost. + max_total_sequence_length: 8192 + precision: "bfloat16" + logprob_chunk_size: null + offload_optimizer_for_logprob: false + + dtensor_cfg: + enabled: false # using megatron backend for the policy + + megatron_cfg: + enabled: true + checkpoint: + async_save: true + empty_unused_memory_level: 1 + activation_checkpointing: true + # 4-GPU policy world → TP=4, DP=1, PP=1. + tensor_model_parallel_size: 4 + expert_tensor_parallel_size: 1 + expert_model_parallel_size: 1 + pipeline_model_parallel_size: 1 + num_layers_in_first_pipeline_stage: null + num_layers_in_last_pipeline_stage: null + context_parallel_size: 1 + pipeline_dtype: ${policy.precision} + sequence_parallel: false # off for a dense small VLM smoke + # Qwen2.5-VL-3B is dense, but freeze the MoE knobs anyway so nothing + # unexpected fires. Router/expert paths are unused for this model. + freeze_moe_router: true + moe_router_dtype: "fp32" + moe_router_load_balancing_type: "none" + moe_router_bias_update_rate: 0.0 + apply_rope_fusion: true + # VLM: keep bias_activation_fusion off (see nano-v3 recipe rationale; + # avoids fused-kernel edge cases with the vision-tower path). + bias_activation_fusion: false + defer_fp32_logits: false + moe_permute_fusion: true + moe_enable_deepep: false + moe_token_dispatcher_type: "alltoall" + moe_shared_expert_overlap: false + gradient_accumulation_fusion: false + use_fused_weighted_squared_relu: false + + optimizer: + optimizer: "adam" + lr: 5.0e-7 + min_lr: 5.0e-7 + weight_decay: 0.0 + bf16: true + fp16: false + params_dtype: "float32" + adam_beta1: 0.9 + adam_beta2: 0.999 + adam_eps: 1e-8 + sgd_momentum: 0.9 + use_distributed_optimizer: true + use_precision_aware_optimizer: true + optimizer_cpu_offload: false + optimizer_offload_fraction: 0.0 + clip_grad: ${policy.max_grad_norm} + + scheduler: + start_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} + end_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} + weight_decay_incr_style: "constant" + lr_decay_style: "constant" + lr_decay_iters: null + # Smoke test only ever runs 5 steps; keep warmup at 0 so the + # OptimizerParamScheduler invariant `lr_warmup_iters < lr_decay_iters` + # never trips (mirrors the trajectory-collection recipe). + lr_warmup_iters: 0 + lr_warmup_init: 5.0e-8 + + distributed_data_parallel_config: + grad_reduce_in_fp32: false + overlap_grad_reduce: false + overlap_param_gather: false + use_custom_fsdp: false + data_parallel_sharding_strategy: "optim_grads_params" + + env_vars: null + + # See docs/design-docs/sequence-packing-and-dynamic-batching.md + dynamic_batching: + enabled: false + train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}} + sequence_length_round: 64 + + sequence_packing: + # Off for the smoke run — the multimodal packing path has extra invariants + # (imgs_sizes / mm_token_type_ids alignment) that we want to keep out of + # the first-cut integration signal. Turn on once the base loop is green. + enabled: false + train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}} + algorithm: "modified_first_fit_decreasing" + sequence_length_round: 64 + + make_sequence_length_divisible_by: ${policy.megatron_cfg.tensor_model_parallel_size} + max_grad_norm: 1.0 + + optimizer: null + scheduler: null + + generation: + backend: "vllm" + max_new_tokens: 1024 # per-turn cap; total ≤ max_total_sequence_length + temperature: 1.0 + top_p: 1.0 + top_k: null + stop_token_ids: null + stop_strings: null + vllm_cfg: + async_engine: true # required by NeMo-Gym setup asserts + precision: ${policy.precision} + tensor_parallel_size: 4 # 4 vLLM GPUs on the single node + pipeline_parallel_size: 1 + expert_parallel_size: 1 + gpu_memory_utilization: 0.55 # leave headroom for the mcore side of the split + max_model_len: ${policy.max_total_sequence_length} + # enforce_eager off keeps CUDA-graph compile-time impact bounded for the + # smoke run; flip to true if the run OOMs during graph capture. + enforce_eager: true + use_deep_gemm: false + num_last_layers_in_bf16: 0 + num_first_layers_in_bf16: 0 + kv_cache_dtype: "auto" + expose_http_server: true # required by NeMo-Gym setup asserts + # VLMs need the tokenizer initialized before generation + # (run_multimodal_grpo_nemo_gym.py asserts this). + skip_tokenizer_init: false + vllm_kwargs: + # Cap images per prompt at max_rollout_turns + 1 (initial obs) — each + # multi-turn step appends one board image. + limit_mm_per_prompt: {"image": 8} + # Disable vLLM's multimodal processor cache; observed hangs / mm-cache + # desyncs with it enabled on VLM RL runs. Must sit in vllm_kwargs, not + # vllm_cfg (see vlm_grpo_3B.yaml). + mm_processor_cache_gb: 0 + colocated: + enabled: false # non-colocated split + resources: + gpus_per_node: 4 # 4 GPUs dedicated to vLLM + num_nodes: 1 + +data: + # NeMo-Gym builds the real per-turn prompt server-side, so + # max_input_seq_length is not consumed by the NemoGymDataset processor. + max_input_seq_length: null + shuffle: true + num_workers: 0 + use_multiple_dataloader: false + + train: + # Sourced from nemo-rl-games-bandit's frozenlake_size4.yaml — the + # size=4 (num_holes=3) FrozenLake ablation train manifest. + data_path: /lustre/fsw/portfolios/coreai/users/rohitkumarj/game-rlvr-vlm-data/frozenlake_ablations/size_4/train_manifest.jsonl + validation: + # Paired eval manifest for the size=4 FrozenLake ablation. Note: rows are + # `Games/FrozenLake-singleturn-v0` with horizon_cap=1 (single-turn eval), + # while train rows are the multi-turn `Games/FrozenLake-v0`. + data_path: /lustre/fsw/portfolios/coreai/users/rohitkumarj/game-rlvr-vlm-data/frozenlake_ablations/size_4/eval_manifest.jsonl + default: + dataset_name: NemoGymDataset + env_name: "nemo_gym" + prompt_file: null + system_prompt_file: null + processor: "nemo_gym_data_processor" + +env: + should_use_nemo_gym: true + should_log_nemo_gym_responses: true + nemo_gym: # forwarded to NeMo-Gym as initial_global_config_dict + # Match ray.sub's port layout: Gym below the 9000 ephemeral floor, + # non-overlapping with NeMo-RL (3000-4999) or vLLM (7000-8999). + port_range_low: 5000 + port_range_high: 5999 + rollout_max_attempts_to_avoid_lp_nan: 1 + is_trajectory_collection: false # smoke test does real training + config_paths: + # for_training variant — required by NeMo-Gym + - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml + # gym_v drop: this single file wires both the gym_v resources server + # and the gymv_agent responses-API agent (see + # 3rdparty/Gym-workspace/Gym/environments/gym_v/config.yaml). + - environments/gym_v/config.yaml + # Agent-side smoke knobs. Cap agent turns at the same value as + # grpo.max_rollout_turns so both sides agree on the horizon. + # Outer group name matches the manifest's agent_ref.name = "gym_v_agent" + # (see 3rdparty/Gym-workspace/Gym/environments/gym_v/config.yaml). + gym_v_agent: + responses_api_agents: + gymv_agent: + max_steps: 4 + done_if_no_boxed_answer: true # short-circuit on malformed output for smoke signal + +logger: + log_dir: "logs/grpo-qwen25vl-gymv-smoke" + num_val_samples_to_print: 0 + wandb_enabled: false # smoke test — no external logging + tensorboard_enabled: true + mlflow_enabled: false + swanlab_enabled: false + monitor_gpus: false + wandb: + project: "grpo-qwen25vl-gymv" + name: "grpo-qwen25vl-gymv-smoke" + tensorboard: {} + mlflow: + experiment_name: "grpo-qwen25vl-gymv" + run_name: "grpo-qwen25vl-gymv-smoke" + gpu_monitoring: + collection_interval: 30 + flush_interval: 30 + +cluster: + gpus_per_node: 8 + num_nodes: 1 diff --git a/examples/nemo_gym/run_gymv_smoke.sh b/examples/nemo_gym/run_gymv_smoke.sh new file mode 100755 index 00000000000..740accddf8c --- /dev/null +++ b/examples/nemo_gym/run_gymv_smoke.sh @@ -0,0 +1,96 @@ +#!/usr/bin/env bash +# Smoke-test launcher for the gym-v multimodal recipes on a single interactive +# allocation (2 nodes × 8 GPUs). Run this from INSIDE the NeMo-RL container, +# with the repo checked out at the working directory (typically /opt/nemo-rl). +# +# Usage: +# examples/nemo_gym/run_gymv_smoke.sh qwen # Qwen2.5-VL-3B (1n, 4vllm+4mcore) +# examples/nemo_gym/run_gymv_smoke.sh nemotron # Nemotron-Omni-30B sync (2n, 1vllm+1mcore) +# examples/nemo_gym/run_gymv_smoke.sh nemotron async # Nemotron-Omni-30B async, max_trajectory_age_steps=1 +# examples/nemo_gym/run_gymv_smoke.sh nemotron async wandb # ... also sets logger.wandb_enabled=true +# examples/nemo_gym/run_gymv_smoke.sh # any recipe under examples/nemo_gym/ +# +# Extra CLI args after the recipe/mode/wandb tokens are forwarded to the +# training script, so you can layer Hydra overrides on top, e.g.: +# examples/nemo_gym/run_gymv_smoke.sh qwen grpo.max_num_steps=1 + +set -euo pipefail + +RECIPE_KEY="${1:-qwen}"; shift || true + +# Optional mode token ("sync" or "async"). Only consumed if it matches; anything +# else stays in $@ so Hydra overrides after the recipe key still work. +MODE="" +if [[ "${1:-}" == "async" || "${1:-}" == "sync" ]]; then + MODE="${1}"; shift +fi + +# Optional "wandb" token — same consume-if-matches pattern. Appends the +# Hydra override that flips wandb logging on for this run. +EXTRA_HYDRA_ARGS=() +if [[ "${1:-}" == "wandb" ]]; then + EXTRA_HYDRA_ARGS+=("logger.wandb_enabled=true") + shift +fi + +case "${RECIPE_KEY}" in + qwen) + RECIPE="examples/nemo_gym/grpo_qwen25vl_gymv_smoke.yaml" + ;; + nemotron|omni) + if [[ "${MODE}" == "async" ]]; then + RECIPE="examples/nemo_gym/grpo_nemotron_omni_30ba3b_gymv_smoke_async_1off.yaml" + else + RECIPE="examples/nemo_gym/grpo_nemotron_omni_30ba3b_gymv_smoke.yaml" + fi + ;; + *) + RECIPE="${RECIPE_KEY}" + ;; +esac + +if [[ ! -f "${RECIPE}" ]]; then + echo "error: recipe not found at ${RECIPE}" >&2 + exit 1 +fi + +# HF token — needed to download Qwen2.5-VL-3B / Nemotron-Omni checkpoints. +# Sourced from the environment; if you keep it in a dotenv, `source` it first. +if [[ -z "${HF_TOKEN:-}" ]]; then + echo "warn: HF_TOKEN unset — HF hub downloads may 401" >&2 +fi + +# Shared caches keep model weights on the mounted Lustre workspace rather than +# in the container's ephemeral rootfs. Override via env if you already have +# these pointed elsewhere. +export HF_HOME="${HF_HOME:-${PWD}/.cache/huggingface}" +export TRANSFORMERS_CACHE="${TRANSFORMERS_CACHE:-${HF_HOME}/hub}" + +# Ray's AF_UNIX socket path is capped at 107 bytes on Linux. On Lustre-rooted +# working directories (long paths) the default under $PWD/tmp overruns it and +# Ray fails to spin up. Force it under /tmp. +export RAY_TMPDIR="${RAY_TMPDIR:-/tmp/ray}" +mkdir -p "${RAY_TMPDIR}" + +# Multimodal recipes use the multimodal entry point; the text-only path uses +# run_grpo_nemo_gym.py (kept here as a fallback branch even though both +# checked-in recipes currently need the multimodal script). +if grep -q '^\s*is_vlm:\s*true' "${RECIPE}"; then + ENTRY="examples/nemo_gym/run_multimodal_grpo_nemo_gym.py" +else + ENTRY="examples/nemo_gym/run_grpo_nemo_gym.py" +fi + +echo "==> recipe: ${RECIPE}" +echo "==> entry: ${ENTRY}" +echo "==> HF_HOME=${HF_HOME}" +echo "==> RAY_TMPDIR=${RAY_TMPDIR}" +if [[ "${#EXTRA_HYDRA_ARGS[@]}" -gt 0 ]]; then + echo "==> extra: ${EXTRA_HYDRA_ARGS[*]}" +fi + +exec uv run \ + "${ENTRY}" \ + --config "${RECIPE}" \ + ${EXTRA_HYDRA_ARGS[@]+"${EXTRA_HYDRA_ARGS[@]}"} \ + "$@" diff --git a/examples/nemo_gym/run_multimodal_grpo_nemo_gym.py b/examples/nemo_gym/run_multimodal_grpo_nemo_gym.py new file mode 100644 index 00000000000..15c25423c4c --- /dev/null +++ b/examples/nemo_gym/run_multimodal_grpo_nemo_gym.py @@ -0,0 +1,331 @@ +# Copyright (c) 2025, 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. + +import argparse +import os +import pprint +import time + +# Increase the W&B single object size warning threshold. Initially 100_000 (100 KB) -> 10_000_000 (10 MB) +import wandb.util + +wandb.util.VALUE_BYTES_LIMIT = 10_000_000 + +from omegaconf import OmegaConf +from wandb import Table + +from nemo_rl.algorithms.grpo import ( + ColocatablePolicyInterface, + EnvironmentInterface, + GenerationInterface, + Logger, + MasterConfig, + StatefulDataLoader, + TokenizerType, + _should_use_nemo_gym, + grpo_train, + refit_policy_generation, + setup, +) +from nemo_rl.algorithms.utils import get_tokenizer +from nemo_rl.data.utils import setup_response_data +from nemo_rl.distributed.virtual_cluster import init_ray +from nemo_rl.environments.nemo_gym import ( + setup_nemo_gym_config, +) +from nemo_rl.experience.rollouts import run_async_nemo_gym_rollout +from nemo_rl.models.generation import configure_generation_config +from nemo_rl.utils.config import ( + load_config, + parse_hydra_overrides, + register_omegaconf_resolvers, +) +from nemo_rl.utils.logger import get_next_experiment_dir, log_container_init_timing +from nemo_rl.utils.timer import Timer + + +def parse_args() -> tuple[argparse.Namespace, list[str]]: + """Parse command line arguments.""" + parser = argparse.ArgumentParser(description="Run GRPO training with configuration") + parser.add_argument( + "--config", type=str, default=None, help="Path to YAML config file" + ) + + # Parse known args for the script + args, overrides = parser.parse_known_args() + + return args, overrides + + +# These types are directly imported from grpo_train since if something about the architecture changes we want to immediately fail. +def collect_trajectories( + policy: ColocatablePolicyInterface, + policy_generation: GenerationInterface, + val_dataloader: StatefulDataLoader, + tokenizer: TokenizerType, + val_task_to_env: dict[str, EnvironmentInterface], + logger: Logger, + master_config: MasterConfig, +) -> None: + """Run trajectory collection.""" + # common config/state items + colocated_inference = master_config.policy["generation"]["colocated"]["enabled"] + refit_policy_generation(policy, policy_generation, colocated_inference) + + log_filename = "trajectory_collection.jsonl" + + print("\n🔍 Running trajectory collection...", flush=True) + generation_config = master_config.policy["generation"] + for val_batch in val_dataloader: + nemo_gym_rollout_result = run_async_nemo_gym_rollout( + policy_generation=policy_generation, + input_batch=val_batch, + tokenizer=tokenizer, + task_to_env=val_task_to_env, + max_seq_len=master_config.policy["max_total_sequence_length"], + generation_config=generation_config, + max_rollout_turns=None, + greedy=False, + ) + + rows_to_log: list[str] = [] + for key, value in nemo_gym_rollout_result.rollout_metrics.items(): + if "full_result" not in key: + continue + + value: Table + data: list[list[str]] = value.data # (n, 1) + rows_to_log.extend(v[0] for v in data) + + logger.log_string_list_as_jsonl(rows_to_log, log_filename) + + # TODO: eventually as trajectory collection use cases exceed 4 hours, we can leverage the dataloader save functionality to resume + # And also leverage the TimeoutChecker functionality as well + + policy_generation.finish_generation() + + +def main() -> None: + """Main entry point.""" + main_start = time.perf_counter() + log_container_init_timing() + rl_init_timer = Timer(context={"worker": "driver"}) + + register_omegaconf_resolvers() + args, overrides = parse_args() + + if not args.config: + args.config = os.path.join( + os.path.dirname(__file__), + "grpo_workplace_assistant_nemotron_nano_v2_9b.yaml", + ) + + with rl_init_timer.time("config"): + config = load_config(args.config) + print(f"Loaded configuration from: {args.config}") + + if overrides: + print(f"Overrides: {overrides}") + config = parse_hydra_overrides(config, overrides) + + config = OmegaConf.to_container(config, resolve=True) + config = MasterConfig(**config) + print("Applied CLI overrides") + + # Get the next experiment directory with incremented ID + config.logger["log_dir"] = get_next_experiment_dir(config.logger["log_dir"]) + print(f"📊 Using log directory: {config.logger['log_dir']}") + if config.checkpointing["enabled"]: + print( + f"📊 Using checkpoint directory: {config.checkpointing['checkpoint_dir']}" + ) + + with rl_init_timer.time("tokenizer"): + assert config.policy.get("is_vlm", False), ( + "run_multimodal_grpo_nemo_gym.py requires `policy.is_vlm=true` in the config." + ) + processor = get_tokenizer(config.policy["tokenizer"], get_processor=True) + tokenizer = processor.tokenizer + assert config.policy["generation"] is not None, ( + "A generation config is required for GRPO" + ) + config.policy["generation"] = configure_generation_config( + config.policy["generation"], tokenizer + ) + if "vllm_cfg" in config.policy["generation"]: + assert not config.policy["generation"]["vllm_cfg"]["skip_tokenizer_init"], ( + "VLMs require tokenizer to be initialized before generation, so skip_tokenizer_init must be set to False." + ) + + # NeMo-Gym specific config setup. + setup_nemo_gym_config(config, tokenizer) + + # We assert here since this is right after the final config has been materialized. + assert _should_use_nemo_gym(config) + + # NeMo-Gym environment needs to get dp_openai_server_base_urls from policy_generation, so we don't setup env here. + with rl_init_timer.time("data"): + print("\n▶ Setting up data...") + train_dataset, val_dataset = setup_response_data( + processor, config.data, env_configs=None + ) + + # Validation dataset config setup. + if config.grpo["max_val_samples"] is not None: + raise ValueError( + """A non-null `grpo.max_val_samples` parameter is not supported. + +Gym principle is that there is no hidden data pre or post processing from you. What you see is what you get. + +The validation set you pass in will directly be used for validation with no additional preprocessing. If you want to have some number of repetitions, please include that in your dataset, via ``num_repeats``, in your dataset config and `ng_prepare_data` will prepare it accordingly.""" + ) + + if val_dataset is not None: + print( + f"Setting `grpo.max_val_samples` and `grpo.val_batch_size` to the length of the validation dataset, which is {len(val_dataset)}" + ) + config.grpo["max_val_samples"] = len(val_dataset) + config.grpo["val_batch_size"] = config.grpo["max_val_samples"] + + # Print config + print("Final config:") + pprint.pprint(config) + + with rl_init_timer.time("ray_connect"): + init_ray() + + # `is_trajectory_collection` is a NeMo-RL-side control-flow knob; pop it + # before setup() so it is not forwarded into NeMo-Gym's global config (the + # gym actor is now created inside setup()). + is_trajectory_collection = ( + config.env["nemo_gym"].pop("is_trajectory_collection", False) or False + ) + + with rl_init_timer.time("setup"): + ( + policy, + policy_generation, + nemo_gym, + cluster, + dataloader, + val_dataloader, + loss_fn, + logger, + checkpointer, + grpo_state, + master_config, + teacher_worker_groups, + alias_to_group_alias, + ) = setup(config, tokenizer, train_dataset, val_dataset, processor=processor) + + rl_init_timer.record("total", time.perf_counter() - main_start) + rl_init_metrics = rl_init_timer.get_timing_metrics(reduction_op="sum") + print("\n" + "=" * 60) + print(" " * 14 + "RL INIT TIMING BREAKDOWN") + for label, value in sorted(rl_init_metrics.items()): + if isinstance(value, (int, float)): + print(f" {label}: {value:.1f}s") + print("=" * 60 + "\n", flush=True) + + # NeMo-Gym is spun up inside setup() (overlapped with vLLM model load). + # Bind task_to_env and val_task_to_env for the nemo_gym env. + # Hardcode here to match `run_async_nemo_gym_rollout`. + task_to_env = {"nemo_gym": nemo_gym} + val_task_to_env = task_to_env + + if is_trajectory_collection: + collect_trajectories( + policy=policy, + policy_generation=policy_generation, + val_dataloader=val_dataloader, + tokenizer=tokenizer, + val_task_to_env=val_task_to_env, + logger=logger, + master_config=master_config, + ) + # Check if async mode is enabled + elif "async_grpo" in config.grpo and config.grpo["async_grpo"]["enabled"]: + # Async GRPO does not support dynamic sampling, reward scaling, or reward shaping (DAPO features) + unsupported_features = [ + "use_dynamic_sampling", + "reward_scaling", + "reward_shaping", + ] + + for feature in unsupported_features: + if feature not in config.grpo: + continue + + if feature == "use_dynamic_sampling": + if config.grpo[feature]: + raise NotImplementedError( + f"{feature} is not supported with async GRPO" + ) + else: + if config.grpo[feature]["enabled"]: + raise NotImplementedError( + f"{feature} is not supported with async GRPO" + ) + + # Async GRPO does not support multiple dataloaders + if config.data["use_multiple_dataloader"]: + raise NotImplementedError( + "use_multiple_dataloader is not supported with async GRPO" + ) + + from nemo_rl.algorithms.grpo import async_grpo_train + + print("🚀 Running async GRPO training") + + async_config = config.grpo["async_grpo"] + # Run async GRPO training + async_grpo_train( + policy=policy, + policy_generation=policy_generation, + dataloader=dataloader, + val_dataloader=val_dataloader, + tokenizer=tokenizer, + loss_fn=loss_fn, + task_to_env=task_to_env, + val_task_to_env=val_task_to_env, + logger=logger, + checkpointer=checkpointer, + grpo_save_state=grpo_state, + master_config=master_config, + max_trajectory_age_steps=async_config["max_trajectory_age_steps"], + teacher_worker_groups=teacher_worker_groups, + alias_to_group_alias=alias_to_group_alias, + ) + else: + print("🚀 Running synchronous GRPO training") + + # Run standard GRPO training + grpo_train( + policy, + policy_generation, + dataloader, + val_dataloader, + tokenizer, + loss_fn, + task_to_env, + val_task_to_env, + logger, + checkpointer, + grpo_state, + master_config, + ) + + +if __name__ == "__main__": + main() diff --git a/nemo_rl/data/multimodal_utils.py b/nemo_rl/data/multimodal_utils.py index 470618d1e74..9f10539f2c4 100644 --- a/nemo_rl/data/multimodal_utils.py +++ b/nemo_rl/data/multimodal_utils.py @@ -18,7 +18,7 @@ import re from collections import defaultdict from io import BytesIO -from typing import Any, Optional, Union +from typing import Any, Optional, Protocol, Union import requests import torch @@ -66,6 +66,86 @@ logger = logging.getLogger(__name__) +def _images_from_messages(messages: list[dict[str, Any]]) -> list[Image.Image]: + images = [] + for message in messages: + content = message.get("content") + if not isinstance(content, list): + continue + for part in content: + if isinstance(part, dict) and part.get("type") == "image": + images.append(resolve_to_image(part["image"])) + return images + + +class _HuggingFaceMultimodalProcessorAdapter: + """Adapter for processors supporting multimodal ``apply_chat_template``.""" + + def process( + self, + processor: Any, + messages: list[dict[str, Any]], + *, + add_generation_prompt: bool, + ) -> tuple[str, dict[str, Any]]: + formatted_text = processor.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=add_generation_prompt, + ) + processed = processor.apply_chat_template( + messages, + tokenize=True, + add_generation_prompt=add_generation_prompt, + return_tensors="pt", + return_dict=True, + ) + return formatted_text, dict(processed) + + +def process_multimodal_chat( + processor: Any, + messages: list[dict[str, Any]], + *, + add_generation_prompt: bool, +) -> tuple[str, dict[str, Any]]: + """Render and process multimodal chat through a registered or HF adapter. + + Processors with a nonstandard multimodal calling convention must register an + adapter. All other processors are expected to support Hugging Face's + multimodal ``apply_chat_template`` interface. + """ + adapter = _HuggingFaceMultimodalProcessorAdapter() + formatted_text, processed = adapter.process( + processor, + messages, + add_generation_prompt=add_generation_prompt, + ) + if "input_ids" not in processed: + raise ValueError( + f"{type(processor).__name__} did not return required input_ids." + ) + + images = _images_from_messages(messages) + model_inputs = extract_multimodal_model_inputs(processor, processed) + visual_keys = set( + getattr(getattr(processor, "image_processor", None), "model_input_names", []) + ) + visual_keys.update( + key + for key in get_multimodal_keys_from_processor(processor) + if any(marker in key for marker in ("image", "img", "pixel", "aspect_ratio")) + ) + visual_keys.add("imgs_sizes") + if images and not any(key in model_inputs for key in visual_keys): + raise ValueError( + f"{type(processor).__name__} processed {len(images)} image(s) but " + "returned no visual model inputs. Register a custom multimodal " + "processor adapter if this processor does not support the standard " + "Hugging Face multimodal chat-template interface." + ) + return formatted_text, processed + class PackedTensor: """Wrapper around a list of torch tensors and a dimension along which to pack the tensors. @@ -354,6 +434,75 @@ def get_dim_to_pack_along(processor, key: str) -> int: return 0 +def extract_multimodal_model_inputs( + processor: Any, processed: dict[str, Any] +) -> dict[str, PackedTensor | torch.Tensor]: + """Extract packed visual inputs and sequence-aligned auxiliary tensors. + + Multimodal inputs declared by the processor are wrapped in ``PackedTensor``. + Token-type fields remain ordinary tensors because they align with the full + language-model token sequence. + """ + input_ids = processed.get("input_ids") + if input_ids is None: + raise ValueError("Processor output is missing input_ids.") + if not isinstance(input_ids, torch.Tensor) or input_ids.ndim not in (1, 2): + raise ValueError( + "Processor input_ids must be a one- or two-dimensional torch.Tensor." + ) + if input_ids.ndim == 2 and input_ids.shape[0] != 1: + raise ValueError( + "Multimodal chat processing expects a single conversation, got " + f"input_ids shape {tuple(input_ids.shape)}." + ) + sequence_length = input_ids.shape[-1] + + extracted: dict[str, PackedTensor | torch.Tensor] = {} + multimodal_keys = list(get_multimodal_keys_from_processor(processor)) + # Some remote-code processors omit this per-image input from their declared + # model_input_names even though their model forward requires it. + if "imgs_sizes" in processed and "imgs_sizes" not in multimodal_keys: + multimodal_keys.append("imgs_sizes") + for key in multimodal_keys: + if key not in processed: + continue + value = processed[key] + if not isinstance(value, torch.Tensor): + raise ValueError( + f"Processor model input {key!r} must be a torch.Tensor, got " + f"{type(value).__name__}." + ) + if key == "imgs_sizes": + value = value.to(dtype=torch.int32) + extracted[key] = PackedTensor( + value, dim_to_pack=get_dim_to_pack_along(processor, key) + ) + + for key in ("token_type_ids", "mm_token_type_ids"): + if key not in processed: + continue + value = processed[key] + if not isinstance(value, torch.Tensor) or value.ndim not in (1, 2): + raise ValueError( + f"Processor sequence input {key!r} must be a one- or " + "two-dimensional torch.Tensor." + ) + if value.ndim == 2: + if value.shape[0] != 1: + raise ValueError( + f"Processor sequence input {key!r} must contain one " + f"conversation, got shape {tuple(value.shape)}." + ) + value = value[0] + if len(value) != sequence_length: + raise ValueError( + f"Processor sequence input {key!r} has length {len(value)}, " + f"but input_ids has length {sequence_length}." + ) + extracted[key] = value + return extracted + + def resolve_to_image(image_path_or_image: str | Image.Image) -> Image.Image: """Resolve the image path to a PIL.Image object. @@ -376,11 +525,52 @@ def resolve_to_image(image_path_or_image: str | Image.Image) -> Image.Image: header, encoded = image_path_or_image.split(",", 1) image_data = base64.b64decode(encoded) return Image.open(BytesIO(image_data)).convert("RGB") + elif image_path_or_image.startswith("file://"): + return Image.open(image_path_or_image.removeprefix("file://")).convert("RGB") else: # Handle local file path return Image.open(image_path_or_image).convert("RGB") +def image_to_data_url(image: Image.Image, fmt: str = "PNG") -> str: + """Encode a PIL Image as a base64 data URL.""" + buf = BytesIO() + image.save(buf, format=fmt) + encoded = base64.b64encode(buf.getvalue()).decode("utf-8") + return f"data:image/{fmt.lower()};base64,{encoded}" + + +def encode_images_in_examples(nemo_gym_examples: list[dict]) -> list[dict]: + """Walk examples and replace local image paths with base64 data URLs. + + Operates in-place on each example's responses_create_params.input[].content[] + items of type 'input_image'. HTTP(S) and data URLs are preserved; local + paths, including file:// URLs, are encoded as data URLs. + """ + for example in nemo_gym_examples: + input_items = example.get("responses_create_params", {}).get("input", []) + if not isinstance(input_items, list): + continue + for item in input_items: + if not isinstance(item, dict): + continue + content = item.get("content", []) + if not isinstance(content, list): + continue + for part in content: + if not isinstance(part, dict) or part.get("type") != "input_image": + continue + url = part.get("image_url", "") + if isinstance(url, dict): + url = url.get("url", "") + if not isinstance(url, str) or not url: + continue + if url.startswith(("http://", "https://", "data:")): + continue + part["image_url"] = image_to_data_url(resolve_to_image(url)) + return nemo_gym_examples + + def get_media_from_message(message: dict[str, Any]) -> dict[str, list[Any]]: """Get all media from a message log item.""" # Handle None or missing content (e.g., assistant messages with only tool_calls) diff --git a/nemo_rl/data/processors.py b/nemo_rl/data/processors.py index a6a072b0357..6608c257c64 100644 --- a/nemo_rl/data/processors.py +++ b/nemo_rl/data/processors.py @@ -30,6 +30,7 @@ VLMMessageLogType, ) from nemo_rl.data.llm_message_utils import get_formatted_message_log +from nemo_rl.data.multimodal_utils import get_multimodal_keys_from_processor TokenizerType = PreTrainedTokenizerBase @@ -460,9 +461,9 @@ 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, + process_multimodal_chat, resolve_to_image, ) @@ -551,65 +552,15 @@ def vlm_hf_data_processor( images = [resolve_to_image(image) for image in images] - # Detect processors that use placeholder style (e.g., NemotronOmni/InternVL) - # vs OpenAI content list style (e.g., Qwen-VL, Gemma). - # These processors expand tokens in __call__ but NOT in apply_chat_template, - # so we must use processor(text=..., images=...) directly. - _PLACEHOLDER_STYLE_PROCESSORS = ( - "NemotronNanoVLV2Processor", - "NemotronH_Nano_Omni_Reasoning_V3Processor", - ) - _uses_image_placeholder = type(processor).__name__ in _PLACEHOLDER_STYLE_PROCESSORS - - if _uses_image_placeholder and images: - # Convert content list to placeholder text format - image_token = getattr(processor, "image_token", "") - text_parts = [] - for content in user_message["content"]: - if content["type"] == "image": - text_parts.append(image_token) - elif content["type"] == "text": - text_parts.append(content["text"]) - user_message_for_tokenize = {"role": "user", "content": "\n".join(text_parts)} - else: - user_message_for_tokenize = user_message - - # get formatted user message - if hasattr(processor, "conversation_preprocessor"): - user_message_for_chat_template = processor.conversation_preprocessor( - user_message - ) - else: - user_message_for_chat_template = user_message_for_tokenize - - # this is the string-tokenized conversation template for the generation policy (for vllm) - string_formatted_dialog = processor.apply_chat_template( - [user_message_for_chat_template], - tokenize=False, + # Render once for vLLM and process the identical conversation for MCore. + # Registered adapters cover processors with nonstandard image placeholder + # expansion; standard Hugging Face processors use multimodal chat templates. + string_formatted_dialog, message = process_multimodal_chat( + processor, + [user_message], add_generation_prompt=True, ) - # this is the id-tokenized and image processed conversation template for the policy - if _uses_image_placeholder and images: - # Dynamic-resolution path: keep pixel_values in float32 to match vLLM's - # DynamicResolutionImageTiler bit-for-bit. vLLM stores/normalizes in - # float32 and only casts at the vision_model boundary; matching that - # rounding order tightens rollout/train logprob agreement. The model - # forward dispatches on imgs_sizes and handles the bf16 cast. - message: dict = processor( - text=string_formatted_dialog, - images=images, - return_tensors="pt", - ) - else: - message: dict = processor.apply_chat_template( - [user_message_for_tokenize], - tokenize=True, - add_generation_prompt=True, - return_tensors="pt", - return_dict=True, - ) - # 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 @@ -789,7 +740,12 @@ def nemo_gym_data_processor( max_seq_length: int | None, idx: int, ) -> DatumSpec: - """Process a datum dictionary (directly loaded from dataset) into a DatumSpec for Nemo Gym.""" + """Process a datum dictionary (directly loaded from dataset) into a DatumSpec for Nemo Gym. + + NeMo-Gym builds the real cumulative prompt server-side. Both LLM and VLM + rows therefore use a placeholder here; VLM inputs are processed once after + the complete rollout has been collected. + """ output: DatumSpec = { # load to dict format here since `Dataset` cannot handle nested structure well in `NemoGymDataset` "extra_env_info": json.loads(datum_dict["extra_env_info"]), diff --git a/nemo_rl/environments/nemo_gym.py b/nemo_rl/environments/nemo_gym.py index 30fb04954b2..12b22c9fd48 100644 --- a/nemo_rl/environments/nemo_gym.py +++ b/nemo_rl/environments/nemo_gym.py @@ -26,6 +26,17 @@ from transformers import PreTrainedTokenizerBase from nemo_rl.distributed.ray_actor_environment_registry import get_actor_python_env +from PIL import Image +from transformers import PreTrainedTokenizerBase + +from nemo_rl.data.multimodal_utils import ( + PackedTensor, + encode_images_in_examples, + get_dim_to_pack_along, + get_multimodal_keys_from_processor, + resolve_to_image, +) + from nemo_rl.distributed.virtual_cluster import ( DEFAULT_GYM_PORT_RANGE_HIGH, DEFAULT_GYM_PORT_RANGE_LOW, @@ -112,6 +123,10 @@ class NemoGymConfig(TypedDict): # Forwarded from policy.tokenizer.use_fastokens so rollout actors patch their # tokenizer consistently with the driver. Defaults to off when absent. use_fastokens: NotRequired[bool] + # Multimodal fields (populated by `setup_nemo_gym_config` when VLM is enabled). + tokenizer_config: NotRequired[ + Optional[Dict[str, Any]] + ] # For processor reconstruction inside the actor def _detect_invalid_tool_call_and_malformed_thinking( @@ -174,12 +189,119 @@ def _detect_invalid_tool_call_and_malformed_thinking( return is_invalid_tool_call, has_malformed_thinking +######################################## +# Multimodal helpers +######################################## + + +def _extract_input_images_from_message(item: dict) -> list[Image.Image]: + """Pull PIL images out of a Responses-API user-role item's content list.""" + images: list[Image.Image] = [] + content = item.get("content") or [] + if not isinstance(content, list): + return images + for part in content: + if not isinstance(part, dict): + continue + if part.get("type") not in ("input_image", "image", "image_url"): + continue + src = part.get("image") or part.get("image_url") or part.get("url") + if src is None: + continue + if isinstance(src, dict): + src = src.get("url") + if src is None: + continue + images.append(resolve_to_image(src)) + return images + + +def _index_per_turn_images( + seed_obs: list[dict], output: list[dict] +) -> list[list[Image.Image]]: + """Bin server-returned user images by the assistant turn that saw them. + + Walks the Responses-API items in order, accumulating images from user-role + items into a pending list, and flushing them into a per-turn bucket each + time a trainable assistant item (one carrying ``generation_token_ids``) is + reached. The returned list has one entry per trainable assistant turn, + aligned with the postprocess loop's ``turn_idx``. + """ + per_turn: list[list[Image.Image]] = [] + pending: list[Image.Image] = [] + for item in [*(seed_obs or []), *output]: + if item.get("role") == "user": + pending.extend(_extract_input_images_from_message(item)) + elif "generation_token_ids" in item: + per_turn.append(pending) + pending = [] + return per_turn + + +def _attach_multimodal_data_to_user_message( + user_message: dict, + *, + images: list[Image.Image], + processor: Any, +) -> None: + """Attach per-turn multimodal tensors to ``user_message``. + + The processor is only invoked to extract multimodal tensors (pixel_values, + imgs_sizes, num_patches, etc.); its text output is discarded — vLLM's + tokens remain the trajectory. We therefore feed it the minimal placeholder + text it needs to count image regions: one ``processor.image_token`` per + image. Passing the vLLM-decoded text does not work because that text + already contains expanded ``...*N...`` regions, and the + processor would try to re-expand every embedded ````. + """ + if not images or processor is None: + return + image_token = getattr(processor, "image_token", "") + processed = processor( + text=image_token * len(images), + images=images, + return_tensors="pt", + ) + multimodal_keys = list(get_multimodal_keys_from_processor(processor)) + # imgs_sizes / num_frames are not always declared in model_input_names by + # bundled image processors, so append them explicitly when present. RADIO + # (Nemotron-Omni vision encoder) uses temporal patching even for still + # images and requires one num_frames=1 entry per image/tile — mirror the + # SFT-path fix at nemo_rl/data/processors.py:589-594. + if "imgs_sizes" in processed and "imgs_sizes" not in multimodal_keys: + multimodal_keys.append("imgs_sizes") + if "imgs_sizes" in processed and "num_frames" not in processed: + processed["num_frames"] = torch.ones( + len(processed["imgs_sizes"]), dtype=torch.long + ) + if "num_frames" in processed and "num_frames" not in multimodal_keys: + multimodal_keys.append("num_frames") + for key in multimodal_keys: + if key not in processed: + continue + value = processed[key] + if key == "imgs_sizes": + value = value.to(dtype=torch.int32) + user_message[key] = PackedTensor( + value, + dim_to_pack=get_dim_to_pack_along(processor, key), + ) + + @ray.remote(max_restarts=-1, max_task_retries=-1) # pragma: no cover class NemoGym(EnvironmentInterface): """This environment class isn't really used for training. It's really meant as an integration wrapper around NeMo-Gym that hooks into the existing NeMo RL resource management via ray. So there is still one source of truth for resource management in NeMo RL.""" def __init__(self, cfg: NemoGymConfig): self.cfg = cfg + # Reconstruct the processor inside the actor (rather than serializing it + # per rollout call) for full-trajectory multimodal postprocessing. + self._processor: Optional[Any] = None + tokenizer_config = cfg.get("tokenizer_config") + if tokenizer_config: + from nemo_rl.algorithms.utils import get_tokenizer + + self._processor = get_tokenizer(tokenizer_config, get_processor=True) def _spinup(self) -> None: """Start the NeMo-Gym head server and rollout collection helper. @@ -293,6 +415,11 @@ async def run_rollouts( timer = Timer() counts_left = Counter(row["agent_ref"]["name"] for row in nemo_gym_examples) + # For multimodal runs, replace local filesystem image paths in the + # examples with base64 data URLs before shipping to vLLM. No-op when + # examples carry no `input_image` items (text-only case). + encode_images_in_examples(nemo_gym_examples) + timer.start("_run_rollouts_total") nemo_gym_result_iterator = self.rch.run_examples( examples=nemo_gym_examples, head_server_config=self.head_server_config @@ -350,12 +477,22 @@ async def run_rollouts( yield nemo_gym_row["_rowidx"], nemo_rl_result, timing_metrics def _postprocess_nemo_gym_to_nemo_rl_result( - self, nemo_gym_result: dict, tokenizer: PreTrainedTokenizerBase + self, + nemo_gym_row: dict, + nemo_gym_result: dict, + tokenizer: PreTrainedTokenizerBase, ) -> dict: assert isinstance(nemo_gym_result, dict), ( f"Hit a non-successful response when querying NeMo Gym for rollouts: {nemo_gym_result}" ) + processor = getattr(self, "_processor", None) + per_turn_images = _index_per_turn_images( + nemo_gym_result["response"].get("seed_obs") or [], + nemo_gym_result["response"]["output"], + ) + turn_idx = 0 + nemo_rl_message_log = [] seen_token_ids: List[int] = [] batch_decode_items = [] @@ -378,6 +515,7 @@ def _postprocess_nemo_gym_to_nemo_rl_result( ), f"""Non-contiguous messages found! This may be a tokenization issue where certain tokens are combined when messages are concatenated, or it may be due to part of the chat history being truncated (like if super long history is truncated or if reasoning is stripped out). Seen token IDs: {seen_token_ids} Output prompt token IDs: {output_item_dict["prompt_token_ids"]} +output prompt token ids till seen: {output_item_dict["prompt_token_ids"][: len(seen_token_ids)]} """ prompt_token_ids = output_item_dict.pop("prompt_token_ids") @@ -430,6 +568,18 @@ def _postprocess_nemo_gym_to_nemo_rl_result( if routed_experts is not None: user_message["routed_experts"] = routed_experts[prompt_start:prompt_end] nemo_rl_message_log.append(user_message) + + if processor is not None: + images_this_turn = ( + per_turn_images[turn_idx] + if turn_idx < len(per_turn_images) + else [] + ) + _attach_multimodal_data_to_user_message( + user_message, + images=images_this_turn, + processor=processor, + ) # Valid tool calls go through the structured API (tool_calls field) and get # executed by NeMo-Gym. If tool call patterns appear in the text content instead, # the call was invalid and never executed — flag it so training can penalize it. @@ -464,6 +614,7 @@ def _postprocess_nemo_gym_to_nemo_rl_result( batch_decode_items.append( (output_item_dict, prompt_token_ids, generation_token_ids) ) + turn_idx += 1 if batch_decode_items: prompt_strs = tokenizer.batch_decode( @@ -609,6 +760,13 @@ def setup_nemo_gym_config(config, tokenizer) -> None: generation_config["stop_strings"] = None generation_config["stop_token_ids"] = None + # For VLM runs, plumb the tokenizer config into the gym env config so the + # NemoGym actor can reconstruct the processor inside itself (needed for + # multi-turn multimodal postprocessing). + if config.policy.get("is_vlm", False): + env_cfg = config.env.setdefault("nemo_gym", {}) + env_cfg.setdefault("tokenizer_config", dict(config.policy["tokenizer"])) + def spinup_nemo_gym_actor( env_configs: dict[str, Any], @@ -646,6 +804,7 @@ def spinup_nemo_gym_actor( # (where the detector reads them), not part of Gym's global config. invalid_tool_call_patterns = nemo_gym_dict.pop("invalid_tool_call_patterns", None) thinking_tags = nemo_gym_dict.pop("thinking_tags", None) + tokenizer_config = nemo_gym_dict.pop("tokenizer_config", None) # Pass prebuilt cache + venv dirs through the global config so the gym reuses # image-baked venvs instead of rebuilding them. @@ -661,6 +820,7 @@ def spinup_nemo_gym_actor( base_urls=base_urls, invalid_tool_call_patterns=invalid_tool_call_patterns, thinking_tags=thinking_tags, + tokenizer_config=tokenizer_config, require_routed_experts=enable_router_replay, routed_experts_dtype=routed_experts_dtype, use_fastokens=use_fastokens, diff --git a/nemo_rl/models/generation/vllm/vllm_worker_async.py b/nemo_rl/models/generation/vllm/vllm_worker_async.py index 201d0790044..fb993f5de8a 100644 --- a/nemo_rl/models/generation/vllm/vllm_worker_async.py +++ b/nemo_rl/models/generation/vllm/vllm_worker_async.py @@ -682,6 +682,14 @@ async def create_chat_completion( assert request.temperature == generation_config["temperature"] assert request.top_p == generation_config["top_p"] + # Merge recipe-level chat_template_kwargs into the request. Client- + # provided keys win so a caller can still override per request. + if default_chat_template_kwargs: + request.chat_template_kwargs = { + **default_chat_template_kwargs, + **(request.chat_template_kwargs or {}), + } + try: generator = await openai_serving_chat.create_chat_completion( request, raw_request @@ -748,6 +756,17 @@ class NeMoRLServingTokenization(ServingTokenization): @app.post("/tokenize") async def tokenize(request: NeMoRLTokenizeRequest, raw_request: Request): + # Chat-mode tokenize also renders the chat template — inject the + # same default kwargs so /tokenize and /v1/chat/completions produce + # identical prompt tokens under multi-turn. + if default_chat_template_kwargs and hasattr( + request, "chat_template_kwargs" + ): + request.chat_template_kwargs = { + **default_chat_template_kwargs, + **(request.chat_template_kwargs or {}), + } + generator = await openai_serving_tokenization.create_tokenize( request, raw_request ) diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index 4b888cd9ea5..85df8fae8ad 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -1406,6 +1406,15 @@ def freeze_moe_router(megatron_model): # Handle VLM models if hasattr(model_module, "thinker"): model_module = model_module.thinker + # NemotronVLModel / NemotronOmniModel wrap the GPT under + # `.llava_model.language_model`; unwrap that layer first so the + # generic `.language_model.decoder.layers` walk below finds the + # MoE router. + if ( + getattr(model_module, "llava_model", None) is not None + and hasattr(model_module.llava_model, "language_model") + ): + model_module = model_module.llava_model if hasattr(model_module, "language_model"): model_module = model_module.language_model for layer in model_module.decoder.layers: diff --git a/nemo_rl/models/policy/__init__.py b/nemo_rl/models/policy/__init__.py index 6248f4b2c71..7994f37d7b6 100644 --- a/nemo_rl/models/policy/__init__.py +++ b/nemo_rl/models/policy/__init__.py @@ -562,3 +562,5 @@ class PolicyConfig(TypedDict): # If true, use standard Megatron layer specs while keeping ModelOpt # quantization enabled. Useful for faster QARL runs and logged in configs. disable_modelopt_layer_spec: NotRequired[bool] + + is_vlm: NotRequired[bool] diff --git a/nemo_rl/utils/packed_tensor.py b/nemo_rl/utils/packed_tensor.py index 01f58c55a32..c0dd856e9d6 100644 --- a/nemo_rl/utils/packed_tensor.py +++ b/nemo_rl/utils/packed_tensor.py @@ -77,14 +77,16 @@ def packed_broadcast_producer(iterator, group, src, post_iter_func): # Apply backend specific post processing and then convert to linearized uint8 tensor. # contiguous() is required because the upstream iterator may # yield non-contiguous tensors that view(...) cannot handle. - tensor = ( - post_iter_func(next(iterator)) - .contiguous() - .view(torch.uint8) - .view(-1) - ) + # 0-D tensors (e.g. BN `num_batches_tracked` counters on + # Nemotron-Omni's sound encoder) must be reshape(1)-ed + # before `.view(torch.uint8)` — Long→Byte view is illegal + # on scalars. + tensor = post_iter_func(next(iterator)).contiguous() + if tensor.dim() == 0: + tensor = tensor.reshape(1) + tensor = tensor.view(torch.uint8).view(-1) packing_tensor_list[buffer_idx].append(tensor) - packing_tensor_sizes[buffer_idx] += tensor.view(torch.uint8).numel() + packing_tensor_sizes[buffer_idx] += tensor.numel() if packing_tensor_sizes[buffer_idx] > target_packed_tensor_size: break # Pack the tensors and call broadcast collective @@ -140,11 +142,15 @@ def unpack_tensor( packed_tensor_sizes = list(map(lambda x: x[4], meta_data_list)) unpacked_tensor = packed_tensor.split_with_sizes(packed_tensor_sizes) - # unpacked_list = List[(name, torch.Tensor.view(dtype).view(*shape))] + # unpacked_list = List[(name, torch.Tensor.view(dtype).reshape(shape))] + # reshape(tuple) accepts an empty tuple for 0-D targets, whereas + # view(*shape) would call view() with no args and raise. Producer + # side reshapes 0-D tensors to (1,) before packing, and this consumer + # must reshape back to the original 0-D shape stored in meta_data. unpacked_list = [ ( meta_data_list[i][0], - tensor.view(meta_data_list[i][2]).view(*meta_data_list[i][1]), + tensor.view(meta_data_list[i][2]).reshape(tuple(meta_data_list[i][1])), ) for i, tensor in enumerate(unpacked_tensor) ] diff --git a/pyproject.toml b/pyproject.toml index 9958048b578..33e071b2d85 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -171,17 +171,17 @@ mcore = [ # sudo dpkg -i cuda-keyring_1.1-1_all.deb # sudo apt-get update # sudo apt-get install cudnn-cuda-13 - - # megatron-bridge is installed straight from the submodule's own pyproject.toml, so its - # dependencies (transformers, transformer-engine, megatron-core[dev,mlm] -> nvidia-modelopt, - # onnxscript, flash-linear-attention, ...) flow transitively and pick up upstream changes on - # every submodule bump + `uv lock`. Do not mirror them here. transformer-engine's exact - # version is still pinned globally via [tool.uv] override-dependencies. - # [te]/[ssm] activate megatron-bridge's transformer-engine and mamba-ssm/causal-conv1d extras. - # megatron-core is intentionally NOT listed: it flows from megatron-bridge's own - # dependency megatron-core[dev,mlm] and its [tool.uv.sources] path mapping, so Megatron-LM - # pyproject changes propagate on bump without any mirror here. - "megatron-bridge[te,ssm]", + # This dependency also needs to be compatible with the spec in Megatron-Bridge/pyproject.toml. + # It is specified here since we don't directly use Megatron-Bridge/pyproject.toml, but a proxy setup.py+pyproject.toml combo + # outside to allow "optionally" installing the megatron path. It's simpler to deal with transformer-engine here in the NeMo RL pyproject.toml + "transformer-engine[pytorch,core_cu13] @ git+https://github.com/NVIDIA/TransformerEngine.git@release_v2.15", + "megatron-core", + "megatron-bridge", + "megatron-energon[av-decode]~=7.0", + # Must match Megatron-Bridge main's transformers requirement (still includes GLM 5.1 support). + "transformers>=5.8.1,<5.9.0", + "nvidia-modelopt[torch]; sys_platform != 'darwin'", + "onnxscript", # Flash-attn version should be selected to satisfy both TE + vLLM requirements (xformers in particular) # https://github.com/NVIDIA/TransformerEngine/blob/v2.3/transformer_engine/pytorch/attention/dot_product_attention/utils.py#L108 # https://github.com/facebookresearch/xformers/blob/8354497deb2c04c67fbb2e2ad911e86530da0e90/xformers/ops/fmha/flash.py#L76 diff --git a/tests/unit/data/test_multimodal_processor_adapter.py b/tests/unit/data/test_multimodal_processor_adapter.py new file mode 100644 index 00000000000..5e9ccbd1c96 --- /dev/null +++ b/tests/unit/data/test_multimodal_processor_adapter.py @@ -0,0 +1,177 @@ +# 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. + +import pytest +import torch +from PIL import Image + +from nemo_rl.data.multimodal_utils import ( + PackedTensor, + extract_multimodal_model_inputs, + process_multimodal_chat, + register_multimodal_processor_adapter, +) + + +def _messages() -> list[dict]: + return [ + { + "role": "user", + "content": [ + {"type": "image", "image": Image.new("RGB", (2, 2))}, + {"type": "text", "text": "describe"}, + ], + } + ] + + +class _ImageProcessor: + model_input_names = ["pixel_values", "image_grid_thw"] + + +class _Tokenizer: + model_input_names = ["input_ids", "attention_mask"] + + +class StandardProcessor: + image_processor = _ImageProcessor() + tokenizer = _Tokenizer() + model_input_names = [ + "input_ids", + "attention_mask", + "pixel_values", + "image_grid_thw", + ] + + def __init__(self): + self.tokenized_messages = None + + def apply_chat_template( + self, messages, *, tokenize, add_generation_prompt, **kwargs + ): + assert add_generation_prompt + if not tokenize: + return "rendered" + self.tokenized_messages = messages + return { + "input_ids": torch.tensor([[1, 2, 3]]), + "pixel_values": torch.ones(2, 3, 2, 2), + "image_grid_thw": torch.tensor([[1, 2, 2]]), + "token_type_ids": torch.tensor([[0, 1, 1]]), + "mm_token_type_ids": torch.tensor([[0, 1, 1]]), + } + + +def test_standard_hf_adapter_and_model_input_extraction(): + processor = StandardProcessor() + formatted, processed = process_multimodal_chat( + processor, _messages(), add_generation_prompt=True + ) + model_inputs = extract_multimodal_model_inputs(processor, processed) + + assert formatted == "rendered" + assert processor.tokenized_messages[0]["content"][0]["type"] == "image" + assert isinstance( + processor.tokenized_messages[0]["content"][0]["image"], Image.Image + ) + assert isinstance(model_inputs["pixel_values"], PackedTensor) + assert isinstance(model_inputs["image_grid_thw"], PackedTensor) + assert model_inputs["token_type_ids"].tolist() == [0, 1, 1] + assert model_inputs["mm_token_type_ids"].tolist() == [0, 1, 1] + + +def test_smolvlm_inputs_pack_along_dimension_one(): + class SmolVLMProcessor(StandardProcessor): + pass + + processor = SmolVLMProcessor() + _, processed = process_multimodal_chat( + processor, _messages(), add_generation_prompt=True + ) + model_inputs = extract_multimodal_model_inputs(processor, processed) + assert model_inputs["pixel_values"].dim_to_pack == 1 + assert model_inputs["image_grid_thw"].dim_to_pack == 1 + + +def test_registered_nemotron_placeholder_adapter(): + class NemotronNanoVLV2Processor(StandardProcessor): + image_token = "" + + def __init__(self): + self.call = None + + def apply_chat_template(self, messages, **kwargs): + assert messages[0]["content"] == "\ndescribe" + return messages[0]["content"] + + def __call__(self, *, text, images, return_tensors): + self.call = (text, images, return_tensors) + return { + "input_ids": torch.tensor([[1, 2, 3]]), + "pixel_values": torch.ones(1, 3, 2, 2), + } + + processor = NemotronNanoVLV2Processor() + formatted, _ = process_multimodal_chat( + processor, _messages(), add_generation_prompt=True + ) + assert formatted == "\ndescribe" + assert processor.call[0] == formatted + assert len(processor.call[1]) == 1 + + +def test_custom_processor_adapter_registration(): + class CustomProcessor(StandardProcessor): + pass + + class CustomAdapter: + def process(self, processor, messages, *, add_generation_prompt): + assert add_generation_prompt + return "custom", { + "input_ids": torch.tensor([[4, 5]]), + "pixel_values": torch.ones(1, 3, 2, 2), + } + + register_multimodal_processor_adapter("CustomProcessor", CustomAdapter()) + formatted, processed = process_multimodal_chat( + CustomProcessor(), _messages(), add_generation_prompt=True + ) + assert formatted == "custom" + assert processed["input_ids"].tolist() == [[4, 5]] + + +def test_images_without_visual_model_inputs_fail_loudly(): + class MissingVisualProcessor(StandardProcessor): + def apply_chat_template( + self, messages, *, tokenize, add_generation_prompt, **kwargs + ): + if not tokenize: + return "rendered" + return {"input_ids": torch.tensor([[1, 2, 3]])} + + with pytest.raises(ValueError, match="returned no visual model inputs"): + process_multimodal_chat( + MissingVisualProcessor(), _messages(), add_generation_prompt=True + ) + + +@pytest.mark.parametrize("key", ["token_type_ids", "mm_token_type_ids"]) +def test_malformed_sequence_auxiliary_length_fails_loudly(key): + processor = StandardProcessor() + processed = { + "input_ids": torch.tensor([[1, 2, 3]]), + key: torch.tensor([[0, 1]]), + } + with pytest.raises(ValueError, match=f"{key!r} has length 2"): + extract_multimodal_model_inputs(processor, processed) From edb54a61cdfa69fabb05f1edfacf1e6727d97906 Mon Sep 17 00:00:00 2001 From: rohitrango Date: Tue, 21 Jul 2026 17:32:33 -0700 Subject: [PATCH 02/27] feat(nemo-gym): plumb multimodal rollouts through the async single-controller pathway Reword of aa04447c (async plumbing). Adds the wiring in grpo.py and run_multimodal_grpo_nemo_gym.py so that multimodal NeMo-Gym rollouts flow through the async single-controller path introduced in the multimodal foundation commit. Signed-off-by: rohitrango (cherry picked from commit c755a2511f1cfa837b1898b70fd1e96c5e476ca1) Signed-off-by: rohitrango --- .../nemo_gym/run_multimodal_grpo_nemo_gym.py | 1 + nemo_rl/algorithms/grpo.py | 27 +++++++++++++++++++ 2 files changed, 28 insertions(+) diff --git a/examples/nemo_gym/run_multimodal_grpo_nemo_gym.py b/examples/nemo_gym/run_multimodal_grpo_nemo_gym.py index 15c25423c4c..edee72b90bb 100644 --- a/examples/nemo_gym/run_multimodal_grpo_nemo_gym.py +++ b/examples/nemo_gym/run_multimodal_grpo_nemo_gym.py @@ -306,6 +306,7 @@ def main() -> None: max_trajectory_age_steps=async_config["max_trajectory_age_steps"], teacher_worker_groups=teacher_worker_groups, alias_to_group_alias=alias_to_group_alias, + processor=processor, ) else: print("🚀 Running synchronous GRPO training") diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index 4159e5d5b60..aeac5058ea2 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -1983,11 +1983,21 @@ def _preserve_router_replay_routed_experts( target["routed_experts"] = flat_messages["routed_experts"] +def _is_vlm_async_run( + master_config: MasterConfig, + processor: Optional[AutoProcessor], +) -> bool: + """Whether this async run carries multimodal tensors that must reach training.""" + return bool(master_config.policy.get("is_vlm") or processor is not None) + + def _build_async_grpo_train_data( flat_messages: BatchedDataDict, input_lengths: torch.Tensor, repeated_batch: BatchedDataDict, policy_config: PolicyConfig, + master_config: MasterConfig, + processor: Optional[AutoProcessor] = None, ) -> BatchedDataDict[ClippedPGLossDataDict]: """Build the async no-TQ policy train batch from flattened rollout messages.""" train_data = BatchedDataDict[ClippedPGLossDataDict]( @@ -2000,6 +2010,16 @@ def _build_async_grpo_train_data( } ) _preserve_router_replay_routed_experts(train_data, flat_messages, policy_config) + + extra_multimodal_data = flat_messages.get_multimodal_dict(as_tensors=False) + if _is_vlm_async_run(master_config, processor) and not extra_multimodal_data: + raise RuntimeError( + "Async GRPO: is_vlm=True (or a processor was provided) but " + "flat_messages.get_multimodal_dict() returned empty for this replay " + "batch. Check that rollout samples carry multimodal tensors and that " + "the collector is not stripping them." + ) + train_data.update(extra_multimodal_data) return train_data @@ -3879,6 +3899,7 @@ def async_grpo_train( max_trajectory_age_steps: int = 1, teacher_worker_groups: Optional[dict[str, Any]] = None, alias_to_group_alias: Optional[dict[str, str]] = None, + processor: Optional[AutoProcessor] = None, ) -> None: """Run asynchronous GRPO training with replay buffer. @@ -3896,6 +3917,9 @@ def async_grpo_train( grpo_save_state: Training state master_config: Master configuration max_trajectory_age_steps: Maximum age (in training steps) for trajectories to be used in training + processor: Optional HF processor. Required-in-effect for VLM async runs + so per-batch multimodal tensors get forwarded to policy.get_logprobs + / policy.train (see _build_async_grpo_train_data). """ # Ensure we are running with a compatible async generation backend. # Async GRPO (with in-flight weight updates) supports vLLM, Megatron, and TRT-LLM; @@ -3913,6 +3937,7 @@ def async_grpo_train( assert master_config.loss_fn.use_importance_sampling_correction, ( "Importance sampling correction must be enabled for async GRPO for good convergence due to off-policy samples!" ) + if router_replay_enabled(master_config.policy) and ( master_config.data_plane or {} ).get("enabled", False): @@ -4507,6 +4532,8 @@ def async_grpo_train( input_lengths, repeated_batch, master_config.policy, + master_config, + processor=processor, ) train_data.to("cpu") From 4cb8023330e7a96dd523fb2ae2e06d91d73644e7 Mon Sep 17 00:00:00 2001 From: rohitrango Date: Wed, 29 Jul 2026 10:35:05 -0700 Subject: [PATCH 03/27] chore(examples): add Nemotron-Omni gym-v smoke configs (tangram + polygon-naming + multi-turn tool calling) Folds four WIP commits (28ba38e9 + 81f885f6 + ecf72e4a + 71717873): - Bump seq len and generation count on the doorkey smoke config. - Add tangram env smoke config and its launcher. - Add multi-turn multimodal tool calling / polygon-naming smoke config and launcher. Signed-off-by: rohitrango (cherry picked from commit b54e6a8a13c5234dbd9467a8021faffe9d394474) Signed-off-by: rohitrango --- .../grpo_nemotron_omni_30ba3b_gymv_smoke.yaml | 6 +- ...motron_omni_30ba3b_gymv_tangram_smoke.yaml | 33 ++++ ...tron_omni_30ba3b_polygon_naming_smoke.yaml | 177 ++++++++++++++++++ examples/nemo_gym/run_gymv_smoke.sh | 6 + examples/nemo_gym/run_polygon_naming_smoke.sh | 71 +++++++ 5 files changed, 290 insertions(+), 3 deletions(-) create mode 100644 examples/nemo_gym/grpo_nemotron_omni_30ba3b_gymv_tangram_smoke.yaml create mode 100644 examples/nemo_gym/grpo_nemotron_omni_30ba3b_polygon_naming_smoke.yaml create mode 100755 examples/nemo_gym/run_polygon_naming_smoke.sh diff --git a/examples/nemo_gym/grpo_nemotron_omni_30ba3b_gymv_smoke.yaml b/examples/nemo_gym/grpo_nemotron_omni_30ba3b_gymv_smoke.yaml index 8ca6aa78fe3..f9f8edb26bc 100644 --- a/examples/nemo_gym/grpo_nemotron_omni_30ba3b_gymv_smoke.yaml +++ b/examples/nemo_gym/grpo_nemotron_omni_30ba3b_gymv_smoke.yaml @@ -24,7 +24,7 @@ defaults: grpo_nanov3.yaml grpo: num_prompts_per_step: 1 # smoke: minimal loop - num_generations_per_prompt: 4 # gbs = 1 * 4 = 4 + num_generations_per_prompt: 16 # gbs = 1 * 16 = 16 num_val_generations_per_prompt: 1 max_rollout_turns: 4 max_num_epochs: 1 @@ -55,7 +55,7 @@ policy: train_global_batch_size: ${mul:${grpo.num_prompts_per_step}, ${grpo.num_generations_per_prompt}} train_micro_batch_size: 1 logprob_batch_size: 1 - max_total_sequence_length: 4096 + max_total_sequence_length: 8192 megatron_cfg: # Full 8-GPU policy node: TP=8, EP=8, PP=1, CP=1. Overrides @@ -98,7 +98,7 @@ policy: make_sequence_length_divisible_by: 32 generation: - max_new_tokens: 4096 # per-turn cap; 4 turns fit in 8k context + max_new_tokens: ${policy.max_total_sequence_length} # per-turn cap; 4 turns fit in 8k context # Nano-Omni-Reasoning bad_words: only the vision-side text tokens (safe, # in-vocab). The audio-delimiter tokens (, , # ) sit past the LM's text-logit width and cause vLLM diff --git a/examples/nemo_gym/grpo_nemotron_omni_30ba3b_gymv_tangram_smoke.yaml b/examples/nemo_gym/grpo_nemotron_omni_30ba3b_gymv_tangram_smoke.yaml new file mode 100644 index 00000000000..5532d3b9eed --- /dev/null +++ b/examples/nemo_gym/grpo_nemotron_omni_30ba3b_gymv_tangram_smoke.yaml @@ -0,0 +1,33 @@ +# Tangram-QA variant of grpo_nemotron_omni_30ba3b_gymv_smoke.yaml. +# +# Extends the FrozenLake smoke recipe and overrides only the data paths (and +# log labels) — model, parallelism, DDP/scheduler workarounds, and the gym-v +# agent config are all inherited via `defaults:`. +# +# Tangram-QA is single-turn (env `Geometry/Tangram-QA-v0`, horizon_cap=1), +# same agent (`gym_v_agent`, `max_steps: 1`) and same \boxed{...} answer +# grammar as the FrozenLake singleturn eval, so no agent-side overrides are +# needed. +# +# Manifest source (copied + max_output_tokens stripped so the model uses the +# max_model_len - prompt_len fallback): +# manifests/game_rlvr_tangram_phase2_{train,eval}/manifest.jsonl +# → tangram_phase2/{train,eval}_manifest.jsonl +# +# Run with: +# uv run --locked --extra mcore --extra vllm \ +# examples/nemo_gym/run_multimodal_grpo_nemo_gym.py \ +# --config examples/nemo_gym/grpo_nemotron_omni_30ba3b_gymv_tangram_smoke.yaml + +defaults: grpo_nemotron_omni_30ba3b_gymv_smoke.yaml + +data: + train: + data_path: /lustre/fsw/portfolios/coreai/users/rohitkumarj/game-rlvr-vlm-data/tangram_phase2/train_manifest.jsonl + validation: + data_path: /lustre/fsw/portfolios/coreai/users/rohitkumarj/game-rlvr-vlm-data/tangram_phase2/eval_manifest.jsonl + +logger: + log_dir: "logs/grpo-nemotron-omni-30ba3b-gymv-tangram-smoke" + wandb: + name: "grpo-nemotron-omni-30ba3b-gymv-tangram-smoke" diff --git a/examples/nemo_gym/grpo_nemotron_omni_30ba3b_polygon_naming_smoke.yaml b/examples/nemo_gym/grpo_nemotron_omni_30ba3b_polygon_naming_smoke.yaml new file mode 100644 index 00000000000..ac8987e173f --- /dev/null +++ b/examples/nemo_gym/grpo_nemotron_omni_30ba3b_polygon_naming_smoke.yaml @@ -0,0 +1,177 @@ +# Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16 smoke test on polygon_naming. +# +# Sibling of grpo_nemotron_omni_30ba3b_gymv_smoke.yaml. Same model + shape; +# swaps the environment from gym-v to the multi-turn multimodal +# polygon_naming benchmark (this repo's 3rdparty/Gym-workspace/Gym/ +# resources_servers/polygon_naming), driven by multimodal_simple_agent. +# +# Rollout on each task: +# 1. /seed_session injects a user turn with two 128×128 polygon canvases. +# 2. Model calls submit_turn(answers=[[sides, colour], [sides, colour]]). +# 3. /submit_turn (turn 1) acks + injects a user turn with image 3. +# 4. Model calls submit_turn(answers=[[sides, colour]]). +# 5. /submit_turn (turn 2) returns plain text; model emits a final +# assistant message; agent loop terminates. +# 6. /verify multiset-compares against ground truth → reward 0.0 / 1.0. +# +# Shape (2 nodes × 8 GPUs, non-colocated) — identical to gym-v smoke: +# - Node 1: vLLM generation, TP=8 +# - Node 2: Megatron policy training, TP=2, EP=8, CP=2, PP=1 +# +# Data: Gym JSONLs at 3rdparty/Gym-workspace/Gym/resources_servers/ +# polygon_naming/data/{train,validation}.jsonl. Generate them first with: +# examples/nemo_gym/generate_polygon_naming_data.sh +# +# Run with: +# uv run --locked --extra mcore --extra vllm \ +# examples/nemo_gym/run_multimodal_grpo_nemo_gym.py \ +# --config examples/nemo_gym/grpo_nemotron_omni_30ba3b_polygon_naming_smoke.yaml + +defaults: grpo_nanov3.yaml + +grpo: + num_prompts_per_step: 1 + num_generations_per_prompt: 16 + num_val_generations_per_prompt: 1 + # polygon_naming completes in 3 rollout turns (2 tool calls + terminal + # assistant msg). Give one turn of headroom. + max_rollout_turns: 4 + max_num_epochs: 1 + max_num_steps: 100 + val_period: 500 + val_at_start: false + val_at_end: false + max_val_samples: null + val_batch_size: null + async_grpo: + enabled: false + +checkpointing: + enabled: false + checkpoint_dir: "results/grpo-nemotron-omni-30ba3b-polygon-naming-smoke" + +policy: + # model_name: "nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16" + model_name: "/data/models/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16" + is_vlm: true + tokenizer: + name: ${policy.model_name} + chat_template_kwargs: + enable_thinking: true + truncate_history_thinking: false + train_global_batch_size: ${mul:${grpo.num_prompts_per_step}, ${grpo.num_generations_per_prompt}} + train_micro_batch_size: 1 + logprob_batch_size: 1 + max_total_sequence_length: 8192 + + megatron_cfg: + tensor_model_parallel_size: 2 + expert_tensor_parallel_size: 1 + expert_model_parallel_size: 8 + pipeline_model_parallel_size: 1 + context_parallel_size: 2 + sequence_parallel: true + activation_checkpointing: true + radio_force_cpe_eval_mode: true + clear_memory_caches_before_refit: true + # See gym-v smoke recipe: sound tower needs frozen forward path before + # DDP grad-reduce overlap can be re-enabled. + distributed_data_parallel_config: + overlap_grad_reduce: false + overlap_param_gather: false + + scheduler: + lr_warmup_iters: 0 + + make_sequence_length_divisible_by: 32 + + generation: + max_new_tokens: ${policy.max_total_sequence_length} + bad_words: ["", "", ""] + vllm_cfg: + tensor_parallel_size: 8 + pipeline_parallel_size: 1 + expert_parallel_size: 1 + gpu_memory_utilization: 0.8 + max_model_len: ${policy.max_total_sequence_length} + enforce_eager: true + enable_prefix_caching: false + skip_tokenizer_init: false + http_server_serving_chat_kwargs: + enable_auto_tools: true + tool_parser: qwen3_coder + reasoning_parser: nano_v3 + # Required for the OpenAI-style content-parts list to be flattened + # into a string before Nemotron-Omni's Jinja template touches it. + # Without this, the assistant branch stringifies the list as a Python + # repr, breaking the turn-2 token-prefix contiguity check at + # nemo_gym.py:474 for any multi-turn rollout. + chat_template_content_format: string + # Threaded into every incoming ChatCompletionRequest by + # vllm_worker_async.py so the Jinja template sees these at render time. + # vLLM 0.20's OpenAIServingChat.__init__ has no chat_template_kwargs + # arg; NeMo-RL's wire lives in the request handler. + chat_template_kwargs: + enable_thinking: true + truncate_history_thinking: false + vllm_kwargs: + # 3 images per rollout (2 injected on turn 1, 1 injected on turn 2). + # Keep a small buffer for retries. + limit_mm_per_prompt: {"image": 6} + mm_processor_cache_gb: 0 + max_num_batched_tokens: 16384 + mamba_ssm_cache_dtype: "float32" + colocated: + enabled: false + resources: + gpus_per_node: 8 + num_nodes: 2 + +data: + train: + # Generated by examples/nemo_gym/generate_polygon_naming_data.sh. + # Path is relative to the container's WD (/opt/nemo-rl); override via + # `data.train.data_path=...` if your checkout lives elsewhere. + data_path: /opt/nemo-rl/3rdparty/Gym-workspace/Gym/resources_servers/polygon_naming/data/train.jsonl + validation: + data_path: /opt/nemo-rl/3rdparty/Gym-workspace/Gym/resources_servers/polygon_naming/data/validation.jsonl + +env: + should_use_nemo_gym: true + nemo_gym: + is_trajectory_collection: false + # Quiet-unblock: current NeMo-Gym flags polygon_naming's shipped server + # blocks (polygon_naming_resources_server / polygon_naming_multimodal_simple_agent) + # as "almost-servers" with schema-validation errors and aborts spinup by + # default. DoorKey/gym_v pass because their server blocks lack the fields + # the validator now demands. Bypass here so the smoke run reaches the + # multi-turn multimodal path; the real fix is either updating + # 3rdparty/Gym-workspace/Gym/resources_servers/polygon_naming/configs/polygon_naming.yaml + # to add the missing fields (grep the log for + # "Configuration Warnings: Almost-Servers Detected" to see which) or + # bumping the pinned NeMo-Gym to a version where they're optional. + error_on_almost_servers: false + config_paths: + - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml + - resources_servers/polygon_naming/configs/polygon_naming.yaml + # Outer group name matches the agent instance key in polygon_naming.yaml. + # multimodal_simple_agent's tool loop needs a small handful of steps: + # 1 for turn 1's submit_turn, 1 for turn 2's submit_turn, 1 for the + # final assistant message. `max_steps=4` leaves one slot of slack. + polygon_naming_multimodal_simple_agent: + responses_api_agents: + multimodal_simple_agent: + max_steps: 4 + +logger: + log_dir: "logs/grpo-nemotron-omni-30ba3b-polygon-naming-smoke" + wandb_enabled: false + tensorboard_enabled: true + monitor_gpus: false + wandb: + project: "grpo-nemotron-omni-polygon-naming" + name: "grpo-nemotron-omni-30ba3b-polygon-naming-smoke" + +cluster: + gpus_per_node: 8 + num_nodes: 4 diff --git a/examples/nemo_gym/run_gymv_smoke.sh b/examples/nemo_gym/run_gymv_smoke.sh index 740accddf8c..31ac54f3848 100755 --- a/examples/nemo_gym/run_gymv_smoke.sh +++ b/examples/nemo_gym/run_gymv_smoke.sh @@ -8,6 +8,7 @@ # examples/nemo_gym/run_gymv_smoke.sh nemotron # Nemotron-Omni-30B sync (2n, 1vllm+1mcore) # examples/nemo_gym/run_gymv_smoke.sh nemotron async # Nemotron-Omni-30B async, max_trajectory_age_steps=1 # examples/nemo_gym/run_gymv_smoke.sh nemotron async wandb # ... also sets logger.wandb_enabled=true +# examples/nemo_gym/run_gymv_smoke.sh tangram # Nemotron-Omni-30B on Tangram-QA (single-turn) # examples/nemo_gym/run_gymv_smoke.sh # any recipe under examples/nemo_gym/ # # Extra CLI args after the recipe/mode/wandb tokens are forwarded to the @@ -44,6 +45,11 @@ case "${RECIPE_KEY}" in RECIPE="examples/nemo_gym/grpo_nemotron_omni_30ba3b_gymv_smoke.yaml" fi ;; + tangram) + # Nemotron-Omni-30B on Tangram-QA. No async sibling recipe exists yet; + # MODE is silently ignored until one is added. + RECIPE="examples/nemo_gym/grpo_nemotron_omni_30ba3b_gymv_tangram_smoke.yaml" + ;; *) RECIPE="${RECIPE_KEY}" ;; diff --git a/examples/nemo_gym/run_polygon_naming_smoke.sh b/examples/nemo_gym/run_polygon_naming_smoke.sh new file mode 100755 index 00000000000..35b7a2b89da --- /dev/null +++ b/examples/nemo_gym/run_polygon_naming_smoke.sh @@ -0,0 +1,71 @@ +#!/usr/bin/env bash +# Smoke-test launcher for the polygon_naming multi-turn multimodal recipe +# on a single interactive allocation (2 nodes × 8 GPUs). Mirrors +# run_gymv_smoke.sh but points at polygon_naming + multimodal_simple_agent. +# Run this from INSIDE the NeMo-RL container, WD = /opt/nemo-rl. +# +# Usage: +# examples/nemo_gym/run_polygon_naming_smoke.sh # sync GRPO +# examples/nemo_gym/run_polygon_naming_smoke.sh wandb # + wandb logging +# examples/nemo_gym/run_polygon_naming_smoke.sh grpo.max_num_steps=1 # Hydra overrides +# +# The dataset JSONLs must exist under +# 3rdparty/Gym-workspace/Gym/resources_servers/polygon_naming/data/{train,validation}.jsonl +# Auto-generated by this script if missing (see generate_polygon_naming_data.sh +# for a standalone regenerator). + +set -euo pipefail + +# -- Optional "wandb" token -------------------------------------------------- +EXTRA_HYDRA_ARGS=() +if [[ "${1:-}" == "wandb" ]]; then + EXTRA_HYDRA_ARGS+=("logger.wandb_enabled=true") + shift +fi + +RECIPE="examples/nemo_gym/grpo_nemotron_omni_30ba3b_polygon_naming_smoke.yaml" +ENTRY="examples/nemo_gym/run_multimodal_grpo_nemo_gym.py" + +if [[ ! -f "${RECIPE}" ]]; then + echo "error: recipe not found at ${RECIPE}" >&2 + exit 1 +fi + +# HF token — needed for gated Nemotron-Omni downloads. +if [[ -z "${HF_TOKEN:-}" ]]; then + echo "warn: HF_TOKEN unset — HF hub downloads may 401" >&2 +fi + +# Shared caches on Lustre workspace, not container rootfs. +export HF_HOME="${HF_HOME:-${PWD}/.cache/huggingface}" +export TRANSFORMERS_CACHE="${TRANSFORMERS_CACHE:-${HF_HOME}/hub}" + +# Ray AF_UNIX socket path cap on Linux (107 bytes). +export RAY_TMPDIR="${RAY_TMPDIR:-/tmp/ray}" +mkdir -p "${RAY_TMPDIR}" + +# -- Auto-generate train/validation JSONLs if missing --------------------- +GYM_DATA_DIR="3rdparty/Gym-workspace/Gym/resources_servers/polygon_naming/data" +TRAIN_JSONL="${GYM_DATA_DIR}/train.jsonl" +VALIDATION_JSONL="${GYM_DATA_DIR}/validation.jsonl" + +if [[ ! -f "${TRAIN_JSONL}" || ! -f "${VALIDATION_JSONL}" ]]; then + echo "==> polygon_naming JSONLs missing — auto-generating" + examples/nemo_gym/generate_polygon_naming_data.sh +fi + +echo "==> recipe: ${RECIPE}" +echo "==> entry: ${ENTRY}" +echo "==> train: ${TRAIN_JSONL} ($(wc -l < "${TRAIN_JSONL}") rows)" +echo "==> val: ${VALIDATION_JSONL} ($(wc -l < "${VALIDATION_JSONL}") rows)" +echo "==> HF_HOME=${HF_HOME}" +echo "==> RAY_TMPDIR=${RAY_TMPDIR}" +if [[ "${#EXTRA_HYDRA_ARGS[@]}" -gt 0 ]]; then + echo "==> extra: ${EXTRA_HYDRA_ARGS[*]}" +fi + +exec uv run --locked --no-sync \ + "${ENTRY}" \ + --config "${RECIPE}" \ + ${EXTRA_HYDRA_ARGS[@]+"${EXTRA_HYDRA_ARGS[@]}"} \ + "$@" From 37098a98dad2d44e68ad8095827634ddb388a719 Mon Sep 17 00:00:00 2001 From: rohitrango Date: Thu, 30 Jul 2026 13:31:49 -0700 Subject: [PATCH 04/27] fix(nemo-gym): drop unused nemo_gym_row arg from _postprocess_nemo_gym_to_nemo_rl_result MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Commit ec8333db added `nemo_gym_row: dict` as a new positional parameter to `NemoGym._postprocess_nemo_gym_to_nemo_rl_result` but never referenced it in the function body, and neither the sole in-tree caller (`_run_rollouts_iterator` in the same file) nor the six unit-test call sites (`tests/unit/environments/test_nemo_gym{,_router_replay}.py`) were updated to pass it. As a result every sync/async multimodal smoke recipe under `examples/nemo_gym/run_gymv_smoke.sh` crashed with TypeError: NemoGym._postprocess_nemo_gym_to_nemo_rl_result() missing 1 required positional argument: 'tokenizer' as soon as the first rollout came back from the Gym HTTP server — the `nemo_gym_result` positional was being consumed by the phantom `nemo_gym_row` slot, so `tokenizer` looked missing. Since the parameter is unused, the minimal fix is to remove it and restore the original 2-arg `(nemo_gym_result, tokenizer)` signature. This lines back up with all six test call sites (which were already passing 2 args) and the production caller at nemo_gym.py:442 (which was already passing 2 args), so no other files need touching. Signed-off-by: rohitrango (cherry picked from commit a2f8fd829458e4343af8afa6e493134435646318) Signed-off-by: rohitrango --- nemo_rl/environments/nemo_gym.py | 1 - 1 file changed, 1 deletion(-) diff --git a/nemo_rl/environments/nemo_gym.py b/nemo_rl/environments/nemo_gym.py index 12b22c9fd48..a3c777867fb 100644 --- a/nemo_rl/environments/nemo_gym.py +++ b/nemo_rl/environments/nemo_gym.py @@ -478,7 +478,6 @@ async def run_rollouts( def _postprocess_nemo_gym_to_nemo_rl_result( self, - nemo_gym_row: dict, nemo_gym_result: dict, tokenizer: PreTrainedTokenizerBase, ) -> dict: From 5b43d9855712758b5737a824629970bbb4d2a549 Mon Sep 17 00:00:00 2001 From: rohitrango Date: Tue, 4 Aug 2026 08:25:23 -0700 Subject: [PATCH 05/27] delete scratchspace configs Signed-off-by: rohitrango (cherry picked from commit 7ddf223046ea6daf6ab1cbe98c02ca4ee9add978) Signed-off-by: rohitrango --- .../nemo_gym/generate_polygon_naming_data.sh | 56 ++++ .../grpo_nemotron_omni_30ba3b_gymv_smoke.yaml | 187 ----------- ...ron_omni_30ba3b_gymv_smoke_async_1off.yaml | 22 -- ...motron_omni_30ba3b_gymv_tangram_smoke.yaml | 33 -- ...tron_omni_30ba3b_polygon_naming_smoke.yaml | 177 ---------- .../nemo_gym/grpo_qwen25vl_gymv_smoke.yaml | 309 ------------------ examples/nemo_gym/run_gymv_smoke.sh | 102 ------ examples/nemo_gym/run_polygon_naming_smoke.sh | 71 ---- nemo_rl/environments/nemo_gym.py | 5 +- 9 files changed, 58 insertions(+), 904 deletions(-) create mode 100755 examples/nemo_gym/generate_polygon_naming_data.sh delete mode 100644 examples/nemo_gym/grpo_nemotron_omni_30ba3b_gymv_smoke.yaml delete mode 100644 examples/nemo_gym/grpo_nemotron_omni_30ba3b_gymv_smoke_async_1off.yaml delete mode 100644 examples/nemo_gym/grpo_nemotron_omni_30ba3b_gymv_tangram_smoke.yaml delete mode 100644 examples/nemo_gym/grpo_nemotron_omni_30ba3b_polygon_naming_smoke.yaml delete mode 100644 examples/nemo_gym/grpo_qwen25vl_gymv_smoke.yaml delete mode 100755 examples/nemo_gym/run_gymv_smoke.sh delete mode 100755 examples/nemo_gym/run_polygon_naming_smoke.sh diff --git a/examples/nemo_gym/generate_polygon_naming_data.sh b/examples/nemo_gym/generate_polygon_naming_data.sh new file mode 100755 index 00000000000..aeba8871862 --- /dev/null +++ b/examples/nemo_gym/generate_polygon_naming_data.sh @@ -0,0 +1,56 @@ +#!/usr/bin/env bash +# Generate polygon_naming train + validation JSONLs from within the +# NeMo-RL container. Wraps 3rdparty/Gym-workspace/Gym/resources_servers/ +# polygon_naming/data/generate_data.py, run through the Gym workspace's +# uv-managed venv (Pillow is not part of the NeMo-RL base env). +# +# Run from WD = /opt/nemo-rl inside the container. +# +# Usage: +# examples/nemo_gym/generate_polygon_naming_data.sh # 512 train / 64 val +# examples/nemo_gym/generate_polygon_naming_data.sh --train 128 --val 16 # custom sizes +# examples/nemo_gym/generate_polygon_naming_data.sh --seed 42 # deterministic + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +NEMO_RL_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)" +GYM_ROOT="${NEMO_RL_ROOT}/3rdparty/Gym-workspace/Gym" + +if [[ ! -f "${GYM_ROOT}/pyproject.toml" ]]; then + echo "error: Gym workspace not found at ${GYM_ROOT}" >&2 + exit 1 +fi + +TRAIN_ROWS=512 +VAL_ROWS=64 +SEED=0 + +while [[ $# -gt 0 ]]; do + case "$1" in + --train) TRAIN_ROWS="$2"; shift 2 ;; + --val) VAL_ROWS="$2"; shift 2 ;; + --seed) SEED="$2"; shift 2 ;; + *) echo "unknown arg: $1" >&2; exit 1 ;; + esac +done + +DATA_DIR="${GYM_ROOT}/resources_servers/polygon_naming/data" +GEN="${DATA_DIR}/generate_data.py" + +echo "==> gym root: ${GYM_ROOT}" +echo "==> train: ${DATA_DIR}/train.jsonl (${TRAIN_ROWS} rows, seed ${SEED})" +echo "==> val: ${DATA_DIR}/validation.jsonl (${VAL_ROWS} rows, seed $((SEED + 1)))" + +cd "${GYM_ROOT}" + +# Train and validation are drawn from different seeds so no overlap by +# construction (seeds are independent RNG streams; different num_rows +# further reduces the chance of shared rows). +uv run python "${GEN}" --num-rows "${TRAIN_ROWS}" --seed "${SEED}" \ + --output "${DATA_DIR}/train.jsonl" + +uv run python "${GEN}" --num-rows "${VAL_ROWS}" --seed "$((SEED + 1))" \ + --output "${DATA_DIR}/validation.jsonl" + +echo "==> done" diff --git a/examples/nemo_gym/grpo_nemotron_omni_30ba3b_gymv_smoke.yaml b/examples/nemo_gym/grpo_nemotron_omni_30ba3b_gymv_smoke.yaml deleted file mode 100644 index f9f8edb26bc..00000000000 --- a/examples/nemo_gym/grpo_nemotron_omni_30ba3b_gymv_smoke.yaml +++ /dev/null @@ -1,187 +0,0 @@ -# Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16 smoke test on gym-v. -# -# Sibling of grpo_qwen25vl_gymv_smoke.yaml. Nemotron-Omni is the known-good -# multimodal + NeMo-Gym baseline in-house (see the games-bandit recipe chain); -# this file adapts it to the current mm-integration branch so the Qwen path -# and the Nemotron-Omni path can be exercised through the same code path. -# -# Both the model AND the env are multimodal here — expect real visual RL -# signal (unlike the earlier text-only Nemotron placeholder). -# -# Shape (2 nodes × 8 GPUs, non-colocated): -# - Node 1: vLLM generation, TP=8 -# - Node 2: Megatron policy training, TP=8, EP=8, PP=1, CP=1 -# -# Data: size=4 FrozenLake ablation manifest (multi-turn train, single-turn -# eval), sourced from nemo-rl-games-bandit's frozenlake_size4.yaml chain. -# -# Run with: -# uv run --locked --extra mcore --extra vllm \ -# examples/nemo_gym/run_multimodal_grpo_nemo_gym.py \ -# --config examples/nemo_gym/grpo_nemotron_omni_30ba3b_gymv_smoke.yaml - -defaults: grpo_nanov3.yaml - -grpo: - num_prompts_per_step: 1 # smoke: minimal loop - num_generations_per_prompt: 16 # gbs = 1 * 16 = 16 - num_val_generations_per_prompt: 1 - max_rollout_turns: 4 - max_num_epochs: 1 - max_num_steps: 20 # - val_period: 500 - val_at_start: false # exercise eval before training - val_at_end: false - max_val_samples: null - val_batch_size: null - async_grpo: - enabled: false # sync for the smoke run - -checkpointing: - enabled: false - checkpoint_dir: "results/grpo-nemotron-omni-30ba3b-gymv-smoke" - -policy: - model_name: "nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16" - is_vlm: true # required by run_multimodal_grpo_nemo_gym.py - tokenizer: - name: ${policy.model_name} - # Nano-Omni-Reasoning uses a native block; keep it open and - # preserve prior reasoning across turns (append-only trajectory rep). - # Mirrors the games-bandit grpo_game_rlvr_trajectory_collection.yaml recipe. - chat_template_kwargs: - enable_thinking: true - truncate_history_thinking: false - train_global_batch_size: ${mul:${grpo.num_prompts_per_step}, ${grpo.num_generations_per_prompt}} - train_micro_batch_size: 1 - logprob_batch_size: 1 - max_total_sequence_length: 8192 - - megatron_cfg: - # Full 8-GPU policy node: TP=8, EP=8, PP=1, CP=1. Overrides - # grpo_nanov3.yaml's TP=2/PP=2/CP=4 layout (tuned for a 32-node run). - tensor_model_parallel_size: 2 - expert_tensor_parallel_size: 1 - expert_model_parallel_size: 8 - pipeline_model_parallel_size: 1 - context_parallel_size: 2 - sequence_parallel: true # sp only pays off with tp + cp - activation_checkpointing: true - # RADIO CPE eval mode: keeps vision-tower positional embeddings in eval - # mode during rollout/train (required for the frozen vision path to - # produce stable features). - radio_force_cpe_eval_mode: true - # Empty the CUDA cache before the vLLM refit broadcast so the packed - # staging tensor (see nemo_rl/utils/packed_tensor.py) doesn't OOM against - # Adam m/v that materialize at the end of iter 1. - clear_memory_caches_before_refit: true - # Nemotron-Omni carries a sound_encoder / sound_projection tower that - # never sees an audio input in gym-v. Those params live in the DDP bucket - # with requires_grad=True but never receive a backward hook, so MCore's - # async grad-reduce path trips the golden-count assertion at - # param_and_grad_buffer.py:272 on iter 2's zero_grad_buffer(). Disable - # the overlap paths until the sound tower is actually frozen via a - # nemo-rl pre_wrap_hook (mirror of freeze_moe_router). This is a real - # perf cost — remove once the freeze path is wired. - distributed_data_parallel_config: - overlap_grad_reduce: false - overlap_param_gather: false - - # Flat LR throughout: skip warmup (base recipe warmed 3e-7 → 3e-6 over 10 - # iters). Combined with lr_decay_style="constant" and min_lr == lr from - # grpo_nanov3.yaml, this pins the LR at 3e-6 from step 0 onward. - scheduler: - lr_warmup_iters: 0 - - # base recipe ties this to dtensor_cfg.tensor_parallel_size; we're on - # megatron, so pin to the megatron TP directly. - make_sequence_length_divisible_by: 32 - - generation: - max_new_tokens: ${policy.max_total_sequence_length} # per-turn cap; 4 turns fit in 8k context - # Nano-Omni-Reasoning bad_words: only the vision-side text tokens (safe, - # in-vocab). The audio-delimiter tokens (, , - # ) sit past the LM's text-logit width and cause vLLM - # out-of-bounds writes in bad_words masking — see the games-bandit - # vlm_grpo_games.yaml history comment for the full incident. - bad_words: ["", "", ""] - vllm_cfg: - # Full 8-GPU vLLM node. - tensor_parallel_size: 8 - pipeline_parallel_size: 1 - expert_parallel_size: 1 - gpu_memory_utilization: 0.8 - max_model_len: ${policy.max_total_sequence_length} - enforce_eager: true # skip CUDA-graph capture for the smoke run - # Nano-Omni + prefix caching crashes vLLM's mm cache (mm_hash miss). - # Belt-and-suspenders: also disabled via mm_processor_cache_gb=0 below. - enable_prefix_caching: false - # VLMs need the tokenizer initialized (run_multimodal_grpo_nemo_gym.py asserts). - skip_tokenizer_init: false - # Nano-Omni chat template expects string content, not the OpenAI - # {type: text, text: ...} list. Keeps the NeMo-Gym prompt-token prefix - # invariant across multi-turn rollouts. - http_server_serving_chat_kwargs: - enable_auto_tools: true - tool_parser: qwen3_coder - reasoning_parser: nemotron_v3 - # chat_template_content_format: string - vllm_kwargs: - # Bounds per-prompt image count (multi-turn appends one board / turn). - limit_mm_per_prompt: {"image": 8} - # Disable vLLM's mm processor cache — the mm-cache desync crash on - # Nano-Omni is the load-bearing fix. Must sit in vllm_kwargs, not - # vllm_cfg (see vlm_grpo_3B.yaml note). - mm_processor_cache_gb: 0 - max_num_batched_tokens: 16384 - # Nano-Omni's mamba backbone needs SSM cache in fp32 (accuracy). - mamba_ssm_cache_dtype: "float32" - colocated: - enabled: false # non-colocated: dedicated vLLM node - resources: - gpus_per_node: 8 - num_nodes: 2 - -data: - train: - # Sourced from nemo-rl-games-bandit's frozenlake_size4.yaml — the - # size=4 (num_holes=3) FrozenLake ablation train manifest. - data_path: /lustre/fsw/portfolios/coreai/users/rohitkumarj/game-rlvr-vlm-data/frozenlake_ablations/size_4/train_manifest.jsonl - validation: - # Paired eval manifest for the size=4 FrozenLake ablation. Note: rows are - # `Games/FrozenLake-singleturn-v0` with horizon_cap=1 (single-turn eval), - # while train rows are the multi-turn `Games/FrozenLake-v0`. - data_path: /lustre/fsw/portfolios/coreai/users/rohitkumarj/game-rlvr-vlm-data/frozenlake_ablations/size_4/eval_manifest.jsonl - -env: - should_use_nemo_gym: true - nemo_gym: - is_trajectory_collection: false - # Replace the base recipe's math/code/etc bundle with gym-v only. Any - # env-block overrides inherited from grpo_nanov3.yaml (math_with_judge, - # code_gen, workplace_assistant, ...) are ignored by Gym when their - # config_paths aren't loaded. - config_paths: - - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml - - environments/gym_v/config.yaml - # Agent-side horizon; keep in sync with grpo.max_rollout_turns. - # Outer group name matches the manifest's agent_ref.name = "gym_v_agent" - # (see 3rdparty/Gym-workspace/Gym/environments/gym_v/config.yaml). - gym_v_agent: - responses_api_agents: - gymv_agent: - max_steps: 1 - done_if_no_boxed_answer: true - -logger: - log_dir: "logs/grpo-nemotron-omni-30ba3b-gymv-smoke" - wandb_enabled: false - tensorboard_enabled: true - monitor_gpus: false - wandb: - project: "grpo-nemotron-omni-gymv" - name: "grpo-nemotron-omni-30ba3b-gymv-smoke" - -cluster: - gpus_per_node: 8 - num_nodes: 4 diff --git a/examples/nemo_gym/grpo_nemotron_omni_30ba3b_gymv_smoke_async_1off.yaml b/examples/nemo_gym/grpo_nemotron_omni_30ba3b_gymv_smoke_async_1off.yaml deleted file mode 100644 index db7f6edfc8f..00000000000 --- a/examples/nemo_gym/grpo_nemotron_omni_30ba3b_gymv_smoke_async_1off.yaml +++ /dev/null @@ -1,22 +0,0 @@ -# Async variant of grpo_nemotron_omni_30ba3b_gymv_smoke.yaml. -# -# Extends the sync smoke recipe and flips only the async_grpo block on: -# 1-step trajectory age (mirrors the LLM `-async-1off` recipes under -# examples/configs/recipes/llm/performance/). -# -# Run with: -# uv run --locked --extra mcore --extra vllm \ -# examples/nemo_gym/run_multimodal_grpo_nemo_gym.py \ -# --config examples/nemo_gym/grpo_nemotron_omni_30ba3b_gymv_smoke_async_1off.yaml - -defaults: grpo_nemotron_omni_30ba3b_gymv_smoke.yaml - -grpo: - async_grpo: - enabled: true - max_trajectory_age_steps: 1 - -logger: - log_dir: "logs/grpo-nemotron-omni-30ba3b-gymv-smoke-async" - wandb: - name: "grpo-nemotron-omni-30ba3b-gymv-smoke-async" \ No newline at end of file diff --git a/examples/nemo_gym/grpo_nemotron_omni_30ba3b_gymv_tangram_smoke.yaml b/examples/nemo_gym/grpo_nemotron_omni_30ba3b_gymv_tangram_smoke.yaml deleted file mode 100644 index 5532d3b9eed..00000000000 --- a/examples/nemo_gym/grpo_nemotron_omni_30ba3b_gymv_tangram_smoke.yaml +++ /dev/null @@ -1,33 +0,0 @@ -# Tangram-QA variant of grpo_nemotron_omni_30ba3b_gymv_smoke.yaml. -# -# Extends the FrozenLake smoke recipe and overrides only the data paths (and -# log labels) — model, parallelism, DDP/scheduler workarounds, and the gym-v -# agent config are all inherited via `defaults:`. -# -# Tangram-QA is single-turn (env `Geometry/Tangram-QA-v0`, horizon_cap=1), -# same agent (`gym_v_agent`, `max_steps: 1`) and same \boxed{...} answer -# grammar as the FrozenLake singleturn eval, so no agent-side overrides are -# needed. -# -# Manifest source (copied + max_output_tokens stripped so the model uses the -# max_model_len - prompt_len fallback): -# manifests/game_rlvr_tangram_phase2_{train,eval}/manifest.jsonl -# → tangram_phase2/{train,eval}_manifest.jsonl -# -# Run with: -# uv run --locked --extra mcore --extra vllm \ -# examples/nemo_gym/run_multimodal_grpo_nemo_gym.py \ -# --config examples/nemo_gym/grpo_nemotron_omni_30ba3b_gymv_tangram_smoke.yaml - -defaults: grpo_nemotron_omni_30ba3b_gymv_smoke.yaml - -data: - train: - data_path: /lustre/fsw/portfolios/coreai/users/rohitkumarj/game-rlvr-vlm-data/tangram_phase2/train_manifest.jsonl - validation: - data_path: /lustre/fsw/portfolios/coreai/users/rohitkumarj/game-rlvr-vlm-data/tangram_phase2/eval_manifest.jsonl - -logger: - log_dir: "logs/grpo-nemotron-omni-30ba3b-gymv-tangram-smoke" - wandb: - name: "grpo-nemotron-omni-30ba3b-gymv-tangram-smoke" diff --git a/examples/nemo_gym/grpo_nemotron_omni_30ba3b_polygon_naming_smoke.yaml b/examples/nemo_gym/grpo_nemotron_omni_30ba3b_polygon_naming_smoke.yaml deleted file mode 100644 index ac8987e173f..00000000000 --- a/examples/nemo_gym/grpo_nemotron_omni_30ba3b_polygon_naming_smoke.yaml +++ /dev/null @@ -1,177 +0,0 @@ -# Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16 smoke test on polygon_naming. -# -# Sibling of grpo_nemotron_omni_30ba3b_gymv_smoke.yaml. Same model + shape; -# swaps the environment from gym-v to the multi-turn multimodal -# polygon_naming benchmark (this repo's 3rdparty/Gym-workspace/Gym/ -# resources_servers/polygon_naming), driven by multimodal_simple_agent. -# -# Rollout on each task: -# 1. /seed_session injects a user turn with two 128×128 polygon canvases. -# 2. Model calls submit_turn(answers=[[sides, colour], [sides, colour]]). -# 3. /submit_turn (turn 1) acks + injects a user turn with image 3. -# 4. Model calls submit_turn(answers=[[sides, colour]]). -# 5. /submit_turn (turn 2) returns plain text; model emits a final -# assistant message; agent loop terminates. -# 6. /verify multiset-compares against ground truth → reward 0.0 / 1.0. -# -# Shape (2 nodes × 8 GPUs, non-colocated) — identical to gym-v smoke: -# - Node 1: vLLM generation, TP=8 -# - Node 2: Megatron policy training, TP=2, EP=8, CP=2, PP=1 -# -# Data: Gym JSONLs at 3rdparty/Gym-workspace/Gym/resources_servers/ -# polygon_naming/data/{train,validation}.jsonl. Generate them first with: -# examples/nemo_gym/generate_polygon_naming_data.sh -# -# Run with: -# uv run --locked --extra mcore --extra vllm \ -# examples/nemo_gym/run_multimodal_grpo_nemo_gym.py \ -# --config examples/nemo_gym/grpo_nemotron_omni_30ba3b_polygon_naming_smoke.yaml - -defaults: grpo_nanov3.yaml - -grpo: - num_prompts_per_step: 1 - num_generations_per_prompt: 16 - num_val_generations_per_prompt: 1 - # polygon_naming completes in 3 rollout turns (2 tool calls + terminal - # assistant msg). Give one turn of headroom. - max_rollout_turns: 4 - max_num_epochs: 1 - max_num_steps: 100 - val_period: 500 - val_at_start: false - val_at_end: false - max_val_samples: null - val_batch_size: null - async_grpo: - enabled: false - -checkpointing: - enabled: false - checkpoint_dir: "results/grpo-nemotron-omni-30ba3b-polygon-naming-smoke" - -policy: - # model_name: "nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16" - model_name: "/data/models/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16" - is_vlm: true - tokenizer: - name: ${policy.model_name} - chat_template_kwargs: - enable_thinking: true - truncate_history_thinking: false - train_global_batch_size: ${mul:${grpo.num_prompts_per_step}, ${grpo.num_generations_per_prompt}} - train_micro_batch_size: 1 - logprob_batch_size: 1 - max_total_sequence_length: 8192 - - megatron_cfg: - tensor_model_parallel_size: 2 - expert_tensor_parallel_size: 1 - expert_model_parallel_size: 8 - pipeline_model_parallel_size: 1 - context_parallel_size: 2 - sequence_parallel: true - activation_checkpointing: true - radio_force_cpe_eval_mode: true - clear_memory_caches_before_refit: true - # See gym-v smoke recipe: sound tower needs frozen forward path before - # DDP grad-reduce overlap can be re-enabled. - distributed_data_parallel_config: - overlap_grad_reduce: false - overlap_param_gather: false - - scheduler: - lr_warmup_iters: 0 - - make_sequence_length_divisible_by: 32 - - generation: - max_new_tokens: ${policy.max_total_sequence_length} - bad_words: ["", "", ""] - vllm_cfg: - tensor_parallel_size: 8 - pipeline_parallel_size: 1 - expert_parallel_size: 1 - gpu_memory_utilization: 0.8 - max_model_len: ${policy.max_total_sequence_length} - enforce_eager: true - enable_prefix_caching: false - skip_tokenizer_init: false - http_server_serving_chat_kwargs: - enable_auto_tools: true - tool_parser: qwen3_coder - reasoning_parser: nano_v3 - # Required for the OpenAI-style content-parts list to be flattened - # into a string before Nemotron-Omni's Jinja template touches it. - # Without this, the assistant branch stringifies the list as a Python - # repr, breaking the turn-2 token-prefix contiguity check at - # nemo_gym.py:474 for any multi-turn rollout. - chat_template_content_format: string - # Threaded into every incoming ChatCompletionRequest by - # vllm_worker_async.py so the Jinja template sees these at render time. - # vLLM 0.20's OpenAIServingChat.__init__ has no chat_template_kwargs - # arg; NeMo-RL's wire lives in the request handler. - chat_template_kwargs: - enable_thinking: true - truncate_history_thinking: false - vllm_kwargs: - # 3 images per rollout (2 injected on turn 1, 1 injected on turn 2). - # Keep a small buffer for retries. - limit_mm_per_prompt: {"image": 6} - mm_processor_cache_gb: 0 - max_num_batched_tokens: 16384 - mamba_ssm_cache_dtype: "float32" - colocated: - enabled: false - resources: - gpus_per_node: 8 - num_nodes: 2 - -data: - train: - # Generated by examples/nemo_gym/generate_polygon_naming_data.sh. - # Path is relative to the container's WD (/opt/nemo-rl); override via - # `data.train.data_path=...` if your checkout lives elsewhere. - data_path: /opt/nemo-rl/3rdparty/Gym-workspace/Gym/resources_servers/polygon_naming/data/train.jsonl - validation: - data_path: /opt/nemo-rl/3rdparty/Gym-workspace/Gym/resources_servers/polygon_naming/data/validation.jsonl - -env: - should_use_nemo_gym: true - nemo_gym: - is_trajectory_collection: false - # Quiet-unblock: current NeMo-Gym flags polygon_naming's shipped server - # blocks (polygon_naming_resources_server / polygon_naming_multimodal_simple_agent) - # as "almost-servers" with schema-validation errors and aborts spinup by - # default. DoorKey/gym_v pass because their server blocks lack the fields - # the validator now demands. Bypass here so the smoke run reaches the - # multi-turn multimodal path; the real fix is either updating - # 3rdparty/Gym-workspace/Gym/resources_servers/polygon_naming/configs/polygon_naming.yaml - # to add the missing fields (grep the log for - # "Configuration Warnings: Almost-Servers Detected" to see which) or - # bumping the pinned NeMo-Gym to a version where they're optional. - error_on_almost_servers: false - config_paths: - - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml - - resources_servers/polygon_naming/configs/polygon_naming.yaml - # Outer group name matches the agent instance key in polygon_naming.yaml. - # multimodal_simple_agent's tool loop needs a small handful of steps: - # 1 for turn 1's submit_turn, 1 for turn 2's submit_turn, 1 for the - # final assistant message. `max_steps=4` leaves one slot of slack. - polygon_naming_multimodal_simple_agent: - responses_api_agents: - multimodal_simple_agent: - max_steps: 4 - -logger: - log_dir: "logs/grpo-nemotron-omni-30ba3b-polygon-naming-smoke" - wandb_enabled: false - tensorboard_enabled: true - monitor_gpus: false - wandb: - project: "grpo-nemotron-omni-polygon-naming" - name: "grpo-nemotron-omni-30ba3b-polygon-naming-smoke" - -cluster: - gpus_per_node: 8 - num_nodes: 4 diff --git a/examples/nemo_gym/grpo_qwen25vl_gymv_smoke.yaml b/examples/nemo_gym/grpo_qwen25vl_gymv_smoke.yaml deleted file mode 100644 index 47d5ad8c1bf..00000000000 --- a/examples/nemo_gym/grpo_qwen25vl_gymv_smoke.yaml +++ /dev/null @@ -1,309 +0,0 @@ -# Qwen2.5-VL-3B smoke test for the multimodal + multi-turn NeMo-Gym integration. -# -# Modeled after nemo-rl-games-bandit's frozenlake_size4.yaml chain (see -# examples/configs/frozenlake_difficulty_ablations/frozenlake_size4.yaml) but -# retargeted to this branch's live gym-v drop under -# `3rdparty/Gym-workspace/Gym/environments/gym_v/` — which bundles both the -# gym_v resources server and the gymv_agent responses-API agent. -# -# Shape (single node, 8 GPUs, non-colocated): -# - vLLM generation: 4 GPUs (TP=4) -# - Megatron policy: 4 GPUs (TP=4) -# -# Data: the bundled `environments/gym_v/data/example.jsonl` (8 rows across -# FrozenLake, GameOfLife, and a few other env families). Used for both train -# and validation so the smoke run exercises the multi-turn rollout + full -# validation path without needing an external manifest. -# -# Run with: -# uv run --locked --extra mcore --extra vllm \ -# examples/nemo_gym/run_multimodal_grpo_nemo_gym.py \ -# --config examples/nemo_gym/grpo_qwen25vl_gymv_smoke.yaml - -grpo: - num_prompts_per_step: 1 # smallest: exercises the loop, minimal cost - num_generations_per_prompt: 4 # gbs = 1 * 4 = 4 rollouts/step - num_val_generations_per_prompt: 1 - # Multi-turn: cap at 4 turns so per-rollout images stay < limit_mm_per_prompt - # and the smoke run finishes quickly even on a slow VLM. - max_rollout_turns: 4 - max_num_epochs: 1 - max_num_steps: 5 # smoke: 5 optimizer steps - normalize_rewards: true - use_leave_one_out_baseline: true - val_period: 5 - val_at_start: true # validate the eval path before training - val_at_end: false - overlong_filtering: false - advantage_clip_low: null - advantage_clip_high: null - max_val_samples: null # inferred from val dataset length - val_batch_size: null # inferred from val dataset length - seed: 42 - use_dynamic_sampling: false - batch_multiplier: 1 - reward_shaping: - enabled: false - reward_scaling: - enabled: false - seq_logprob_error_threshold: null - invalid_tool_call_advantage: null - malformed_thinking_advantage: null - async_grpo: - enabled: false # sync GRPO for the smoke run - max_trajectory_age_steps: 1 - -loss_fn: - reference_policy_kl_penalty: 0 - reference_policy_kl_type: "k3" - kl_input_clamp_value: 20.0 - kl_output_clamp_value: 10.0 - ratio_clip_min: 0.2 - ratio_clip_max: 0.2 - ratio_clip_c: null - use_on_policy_kl_approximation: false - truncated_importance_sampling_ratio: null - use_importance_sampling_correction: false - token_level_loss: true - -checkpointing: - enabled: false # smoke test — no checkpoints - checkpoint_dir: "results/grpo-qwen25vl-gymv-smoke" - metric_name: "val:total_reward/mean" - higher_is_better: true - keep_top_k: 1 - save_period: 5 - checkpoint_must_save_by: null - save_optimizer: false - -policy: - model_name: "Qwen/Qwen2.5-VL-3B-Instruct" - is_vlm: true # required by run_multimodal_grpo_nemo_gym.py - tokenizer: - name: ${policy.model_name} - chat_template_kwargs: null - hf_config_overrides: {} - # gbs must equal num_prompts_per_step * num_generations_per_prompt. - train_global_batch_size: ${mul:${grpo.num_prompts_per_step}, ${grpo.num_generations_per_prompt}} - train_micro_batch_size: 1 - logprob_batch_size: 1 - generation_batch_size: 8 - # Room for a 4-turn multi-image rollout with Qwen2.5-VL's per-image token cost. - max_total_sequence_length: 8192 - precision: "bfloat16" - logprob_chunk_size: null - offload_optimizer_for_logprob: false - - dtensor_cfg: - enabled: false # using megatron backend for the policy - - megatron_cfg: - enabled: true - checkpoint: - async_save: true - empty_unused_memory_level: 1 - activation_checkpointing: true - # 4-GPU policy world → TP=4, DP=1, PP=1. - tensor_model_parallel_size: 4 - expert_tensor_parallel_size: 1 - expert_model_parallel_size: 1 - pipeline_model_parallel_size: 1 - num_layers_in_first_pipeline_stage: null - num_layers_in_last_pipeline_stage: null - context_parallel_size: 1 - pipeline_dtype: ${policy.precision} - sequence_parallel: false # off for a dense small VLM smoke - # Qwen2.5-VL-3B is dense, but freeze the MoE knobs anyway so nothing - # unexpected fires. Router/expert paths are unused for this model. - freeze_moe_router: true - moe_router_dtype: "fp32" - moe_router_load_balancing_type: "none" - moe_router_bias_update_rate: 0.0 - apply_rope_fusion: true - # VLM: keep bias_activation_fusion off (see nano-v3 recipe rationale; - # avoids fused-kernel edge cases with the vision-tower path). - bias_activation_fusion: false - defer_fp32_logits: false - moe_permute_fusion: true - moe_enable_deepep: false - moe_token_dispatcher_type: "alltoall" - moe_shared_expert_overlap: false - gradient_accumulation_fusion: false - use_fused_weighted_squared_relu: false - - optimizer: - optimizer: "adam" - lr: 5.0e-7 - min_lr: 5.0e-7 - weight_decay: 0.0 - bf16: true - fp16: false - params_dtype: "float32" - adam_beta1: 0.9 - adam_beta2: 0.999 - adam_eps: 1e-8 - sgd_momentum: 0.9 - use_distributed_optimizer: true - use_precision_aware_optimizer: true - optimizer_cpu_offload: false - optimizer_offload_fraction: 0.0 - clip_grad: ${policy.max_grad_norm} - - scheduler: - start_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} - end_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} - weight_decay_incr_style: "constant" - lr_decay_style: "constant" - lr_decay_iters: null - # Smoke test only ever runs 5 steps; keep warmup at 0 so the - # OptimizerParamScheduler invariant `lr_warmup_iters < lr_decay_iters` - # never trips (mirrors the trajectory-collection recipe). - lr_warmup_iters: 0 - lr_warmup_init: 5.0e-8 - - distributed_data_parallel_config: - grad_reduce_in_fp32: false - overlap_grad_reduce: false - overlap_param_gather: false - use_custom_fsdp: false - data_parallel_sharding_strategy: "optim_grads_params" - - env_vars: null - - # See docs/design-docs/sequence-packing-and-dynamic-batching.md - dynamic_batching: - enabled: false - train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}} - logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}} - sequence_length_round: 64 - - sequence_packing: - # Off for the smoke run — the multimodal packing path has extra invariants - # (imgs_sizes / mm_token_type_ids alignment) that we want to keep out of - # the first-cut integration signal. Turn on once the base loop is green. - enabled: false - train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}} - logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}} - algorithm: "modified_first_fit_decreasing" - sequence_length_round: 64 - - make_sequence_length_divisible_by: ${policy.megatron_cfg.tensor_model_parallel_size} - max_grad_norm: 1.0 - - optimizer: null - scheduler: null - - generation: - backend: "vllm" - max_new_tokens: 1024 # per-turn cap; total ≤ max_total_sequence_length - temperature: 1.0 - top_p: 1.0 - top_k: null - stop_token_ids: null - stop_strings: null - vllm_cfg: - async_engine: true # required by NeMo-Gym setup asserts - precision: ${policy.precision} - tensor_parallel_size: 4 # 4 vLLM GPUs on the single node - pipeline_parallel_size: 1 - expert_parallel_size: 1 - gpu_memory_utilization: 0.55 # leave headroom for the mcore side of the split - max_model_len: ${policy.max_total_sequence_length} - # enforce_eager off keeps CUDA-graph compile-time impact bounded for the - # smoke run; flip to true if the run OOMs during graph capture. - enforce_eager: true - use_deep_gemm: false - num_last_layers_in_bf16: 0 - num_first_layers_in_bf16: 0 - kv_cache_dtype: "auto" - expose_http_server: true # required by NeMo-Gym setup asserts - # VLMs need the tokenizer initialized before generation - # (run_multimodal_grpo_nemo_gym.py asserts this). - skip_tokenizer_init: false - vllm_kwargs: - # Cap images per prompt at max_rollout_turns + 1 (initial obs) — each - # multi-turn step appends one board image. - limit_mm_per_prompt: {"image": 8} - # Disable vLLM's multimodal processor cache; observed hangs / mm-cache - # desyncs with it enabled on VLM RL runs. Must sit in vllm_kwargs, not - # vllm_cfg (see vlm_grpo_3B.yaml). - mm_processor_cache_gb: 0 - colocated: - enabled: false # non-colocated split - resources: - gpus_per_node: 4 # 4 GPUs dedicated to vLLM - num_nodes: 1 - -data: - # NeMo-Gym builds the real per-turn prompt server-side, so - # max_input_seq_length is not consumed by the NemoGymDataset processor. - max_input_seq_length: null - shuffle: true - num_workers: 0 - use_multiple_dataloader: false - - train: - # Sourced from nemo-rl-games-bandit's frozenlake_size4.yaml — the - # size=4 (num_holes=3) FrozenLake ablation train manifest. - data_path: /lustre/fsw/portfolios/coreai/users/rohitkumarj/game-rlvr-vlm-data/frozenlake_ablations/size_4/train_manifest.jsonl - validation: - # Paired eval manifest for the size=4 FrozenLake ablation. Note: rows are - # `Games/FrozenLake-singleturn-v0` with horizon_cap=1 (single-turn eval), - # while train rows are the multi-turn `Games/FrozenLake-v0`. - data_path: /lustre/fsw/portfolios/coreai/users/rohitkumarj/game-rlvr-vlm-data/frozenlake_ablations/size_4/eval_manifest.jsonl - default: - dataset_name: NemoGymDataset - env_name: "nemo_gym" - prompt_file: null - system_prompt_file: null - processor: "nemo_gym_data_processor" - -env: - should_use_nemo_gym: true - should_log_nemo_gym_responses: true - nemo_gym: # forwarded to NeMo-Gym as initial_global_config_dict - # Match ray.sub's port layout: Gym below the 9000 ephemeral floor, - # non-overlapping with NeMo-RL (3000-4999) or vLLM (7000-8999). - port_range_low: 5000 - port_range_high: 5999 - rollout_max_attempts_to_avoid_lp_nan: 1 - is_trajectory_collection: false # smoke test does real training - config_paths: - # for_training variant — required by NeMo-Gym - - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml - # gym_v drop: this single file wires both the gym_v resources server - # and the gymv_agent responses-API agent (see - # 3rdparty/Gym-workspace/Gym/environments/gym_v/config.yaml). - - environments/gym_v/config.yaml - # Agent-side smoke knobs. Cap agent turns at the same value as - # grpo.max_rollout_turns so both sides agree on the horizon. - # Outer group name matches the manifest's agent_ref.name = "gym_v_agent" - # (see 3rdparty/Gym-workspace/Gym/environments/gym_v/config.yaml). - gym_v_agent: - responses_api_agents: - gymv_agent: - max_steps: 4 - done_if_no_boxed_answer: true # short-circuit on malformed output for smoke signal - -logger: - log_dir: "logs/grpo-qwen25vl-gymv-smoke" - num_val_samples_to_print: 0 - wandb_enabled: false # smoke test — no external logging - tensorboard_enabled: true - mlflow_enabled: false - swanlab_enabled: false - monitor_gpus: false - wandb: - project: "grpo-qwen25vl-gymv" - name: "grpo-qwen25vl-gymv-smoke" - tensorboard: {} - mlflow: - experiment_name: "grpo-qwen25vl-gymv" - run_name: "grpo-qwen25vl-gymv-smoke" - gpu_monitoring: - collection_interval: 30 - flush_interval: 30 - -cluster: - gpus_per_node: 8 - num_nodes: 1 diff --git a/examples/nemo_gym/run_gymv_smoke.sh b/examples/nemo_gym/run_gymv_smoke.sh deleted file mode 100755 index 31ac54f3848..00000000000 --- a/examples/nemo_gym/run_gymv_smoke.sh +++ /dev/null @@ -1,102 +0,0 @@ -#!/usr/bin/env bash -# Smoke-test launcher for the gym-v multimodal recipes on a single interactive -# allocation (2 nodes × 8 GPUs). Run this from INSIDE the NeMo-RL container, -# with the repo checked out at the working directory (typically /opt/nemo-rl). -# -# Usage: -# examples/nemo_gym/run_gymv_smoke.sh qwen # Qwen2.5-VL-3B (1n, 4vllm+4mcore) -# examples/nemo_gym/run_gymv_smoke.sh nemotron # Nemotron-Omni-30B sync (2n, 1vllm+1mcore) -# examples/nemo_gym/run_gymv_smoke.sh nemotron async # Nemotron-Omni-30B async, max_trajectory_age_steps=1 -# examples/nemo_gym/run_gymv_smoke.sh nemotron async wandb # ... also sets logger.wandb_enabled=true -# examples/nemo_gym/run_gymv_smoke.sh tangram # Nemotron-Omni-30B on Tangram-QA (single-turn) -# examples/nemo_gym/run_gymv_smoke.sh # any recipe under examples/nemo_gym/ -# -# Extra CLI args after the recipe/mode/wandb tokens are forwarded to the -# training script, so you can layer Hydra overrides on top, e.g.: -# examples/nemo_gym/run_gymv_smoke.sh qwen grpo.max_num_steps=1 - -set -euo pipefail - -RECIPE_KEY="${1:-qwen}"; shift || true - -# Optional mode token ("sync" or "async"). Only consumed if it matches; anything -# else stays in $@ so Hydra overrides after the recipe key still work. -MODE="" -if [[ "${1:-}" == "async" || "${1:-}" == "sync" ]]; then - MODE="${1}"; shift -fi - -# Optional "wandb" token — same consume-if-matches pattern. Appends the -# Hydra override that flips wandb logging on for this run. -EXTRA_HYDRA_ARGS=() -if [[ "${1:-}" == "wandb" ]]; then - EXTRA_HYDRA_ARGS+=("logger.wandb_enabled=true") - shift -fi - -case "${RECIPE_KEY}" in - qwen) - RECIPE="examples/nemo_gym/grpo_qwen25vl_gymv_smoke.yaml" - ;; - nemotron|omni) - if [[ "${MODE}" == "async" ]]; then - RECIPE="examples/nemo_gym/grpo_nemotron_omni_30ba3b_gymv_smoke_async_1off.yaml" - else - RECIPE="examples/nemo_gym/grpo_nemotron_omni_30ba3b_gymv_smoke.yaml" - fi - ;; - tangram) - # Nemotron-Omni-30B on Tangram-QA. No async sibling recipe exists yet; - # MODE is silently ignored until one is added. - RECIPE="examples/nemo_gym/grpo_nemotron_omni_30ba3b_gymv_tangram_smoke.yaml" - ;; - *) - RECIPE="${RECIPE_KEY}" - ;; -esac - -if [[ ! -f "${RECIPE}" ]]; then - echo "error: recipe not found at ${RECIPE}" >&2 - exit 1 -fi - -# HF token — needed to download Qwen2.5-VL-3B / Nemotron-Omni checkpoints. -# Sourced from the environment; if you keep it in a dotenv, `source` it first. -if [[ -z "${HF_TOKEN:-}" ]]; then - echo "warn: HF_TOKEN unset — HF hub downloads may 401" >&2 -fi - -# Shared caches keep model weights on the mounted Lustre workspace rather than -# in the container's ephemeral rootfs. Override via env if you already have -# these pointed elsewhere. -export HF_HOME="${HF_HOME:-${PWD}/.cache/huggingface}" -export TRANSFORMERS_CACHE="${TRANSFORMERS_CACHE:-${HF_HOME}/hub}" - -# Ray's AF_UNIX socket path is capped at 107 bytes on Linux. On Lustre-rooted -# working directories (long paths) the default under $PWD/tmp overruns it and -# Ray fails to spin up. Force it under /tmp. -export RAY_TMPDIR="${RAY_TMPDIR:-/tmp/ray}" -mkdir -p "${RAY_TMPDIR}" - -# Multimodal recipes use the multimodal entry point; the text-only path uses -# run_grpo_nemo_gym.py (kept here as a fallback branch even though both -# checked-in recipes currently need the multimodal script). -if grep -q '^\s*is_vlm:\s*true' "${RECIPE}"; then - ENTRY="examples/nemo_gym/run_multimodal_grpo_nemo_gym.py" -else - ENTRY="examples/nemo_gym/run_grpo_nemo_gym.py" -fi - -echo "==> recipe: ${RECIPE}" -echo "==> entry: ${ENTRY}" -echo "==> HF_HOME=${HF_HOME}" -echo "==> RAY_TMPDIR=${RAY_TMPDIR}" -if [[ "${#EXTRA_HYDRA_ARGS[@]}" -gt 0 ]]; then - echo "==> extra: ${EXTRA_HYDRA_ARGS[*]}" -fi - -exec uv run \ - "${ENTRY}" \ - --config "${RECIPE}" \ - ${EXTRA_HYDRA_ARGS[@]+"${EXTRA_HYDRA_ARGS[@]}"} \ - "$@" diff --git a/examples/nemo_gym/run_polygon_naming_smoke.sh b/examples/nemo_gym/run_polygon_naming_smoke.sh deleted file mode 100755 index 35b7a2b89da..00000000000 --- a/examples/nemo_gym/run_polygon_naming_smoke.sh +++ /dev/null @@ -1,71 +0,0 @@ -#!/usr/bin/env bash -# Smoke-test launcher for the polygon_naming multi-turn multimodal recipe -# on a single interactive allocation (2 nodes × 8 GPUs). Mirrors -# run_gymv_smoke.sh but points at polygon_naming + multimodal_simple_agent. -# Run this from INSIDE the NeMo-RL container, WD = /opt/nemo-rl. -# -# Usage: -# examples/nemo_gym/run_polygon_naming_smoke.sh # sync GRPO -# examples/nemo_gym/run_polygon_naming_smoke.sh wandb # + wandb logging -# examples/nemo_gym/run_polygon_naming_smoke.sh grpo.max_num_steps=1 # Hydra overrides -# -# The dataset JSONLs must exist under -# 3rdparty/Gym-workspace/Gym/resources_servers/polygon_naming/data/{train,validation}.jsonl -# Auto-generated by this script if missing (see generate_polygon_naming_data.sh -# for a standalone regenerator). - -set -euo pipefail - -# -- Optional "wandb" token -------------------------------------------------- -EXTRA_HYDRA_ARGS=() -if [[ "${1:-}" == "wandb" ]]; then - EXTRA_HYDRA_ARGS+=("logger.wandb_enabled=true") - shift -fi - -RECIPE="examples/nemo_gym/grpo_nemotron_omni_30ba3b_polygon_naming_smoke.yaml" -ENTRY="examples/nemo_gym/run_multimodal_grpo_nemo_gym.py" - -if [[ ! -f "${RECIPE}" ]]; then - echo "error: recipe not found at ${RECIPE}" >&2 - exit 1 -fi - -# HF token — needed for gated Nemotron-Omni downloads. -if [[ -z "${HF_TOKEN:-}" ]]; then - echo "warn: HF_TOKEN unset — HF hub downloads may 401" >&2 -fi - -# Shared caches on Lustre workspace, not container rootfs. -export HF_HOME="${HF_HOME:-${PWD}/.cache/huggingface}" -export TRANSFORMERS_CACHE="${TRANSFORMERS_CACHE:-${HF_HOME}/hub}" - -# Ray AF_UNIX socket path cap on Linux (107 bytes). -export RAY_TMPDIR="${RAY_TMPDIR:-/tmp/ray}" -mkdir -p "${RAY_TMPDIR}" - -# -- Auto-generate train/validation JSONLs if missing --------------------- -GYM_DATA_DIR="3rdparty/Gym-workspace/Gym/resources_servers/polygon_naming/data" -TRAIN_JSONL="${GYM_DATA_DIR}/train.jsonl" -VALIDATION_JSONL="${GYM_DATA_DIR}/validation.jsonl" - -if [[ ! -f "${TRAIN_JSONL}" || ! -f "${VALIDATION_JSONL}" ]]; then - echo "==> polygon_naming JSONLs missing — auto-generating" - examples/nemo_gym/generate_polygon_naming_data.sh -fi - -echo "==> recipe: ${RECIPE}" -echo "==> entry: ${ENTRY}" -echo "==> train: ${TRAIN_JSONL} ($(wc -l < "${TRAIN_JSONL}") rows)" -echo "==> val: ${VALIDATION_JSONL} ($(wc -l < "${VALIDATION_JSONL}") rows)" -echo "==> HF_HOME=${HF_HOME}" -echo "==> RAY_TMPDIR=${RAY_TMPDIR}" -if [[ "${#EXTRA_HYDRA_ARGS[@]}" -gt 0 ]]; then - echo "==> extra: ${EXTRA_HYDRA_ARGS[*]}" -fi - -exec uv run --locked --no-sync \ - "${ENTRY}" \ - --config "${RECIPE}" \ - ${EXTRA_HYDRA_ARGS[@]+"${EXTRA_HYDRA_ARGS[@]}"} \ - "$@" diff --git a/nemo_rl/environments/nemo_gym.py b/nemo_rl/environments/nemo_gym.py index a3c777867fb..2283fd00c99 100644 --- a/nemo_rl/environments/nemo_gym.py +++ b/nemo_rl/environments/nemo_gym.py @@ -217,7 +217,7 @@ def _extract_input_images_from_message(item: dict) -> list[Image.Image]: def _index_per_turn_images( - seed_obs: list[dict], output: list[dict] + output: list[dict] ) -> list[list[Image.Image]]: """Bin server-returned user images by the assistant turn that saw them. @@ -229,7 +229,7 @@ def _index_per_turn_images( """ per_turn: list[list[Image.Image]] = [] pending: list[Image.Image] = [] - for item in [*(seed_obs or []), *output]: + for item in output: if item.get("role") == "user": pending.extend(_extract_input_images_from_message(item)) elif "generation_token_ids" in item: @@ -487,7 +487,6 @@ def _postprocess_nemo_gym_to_nemo_rl_result( processor = getattr(self, "_processor", None) per_turn_images = _index_per_turn_images( - nemo_gym_result["response"].get("seed_obs") or [], nemo_gym_result["response"]["output"], ) turn_idx = 0 From faa185a36571b49d8befdbcedaa82c18c4653d1d Mon Sep 17 00:00:00 2001 From: rohitrango Date: Tue, 4 Aug 2026 09:37:52 -0700 Subject: [PATCH 06/27] consolidate into one entrypoint Signed-off-by: rohitrango (cherry picked from commit 7bf55afb43ec2f1107733c0bdcaff08ad7435e97) Signed-off-by: rohitrango --- .../nemo_gym/generate_polygon_naming_data.sh | 56 --- examples/nemo_gym/run_grpo_nemo_gym.py | 27 +- .../nemo_gym/run_multimodal_grpo_nemo_gym.py | 332 ------------------ 3 files changed, 23 insertions(+), 392 deletions(-) delete mode 100755 examples/nemo_gym/generate_polygon_naming_data.sh delete mode 100644 examples/nemo_gym/run_multimodal_grpo_nemo_gym.py diff --git a/examples/nemo_gym/generate_polygon_naming_data.sh b/examples/nemo_gym/generate_polygon_naming_data.sh deleted file mode 100755 index aeba8871862..00000000000 --- a/examples/nemo_gym/generate_polygon_naming_data.sh +++ /dev/null @@ -1,56 +0,0 @@ -#!/usr/bin/env bash -# Generate polygon_naming train + validation JSONLs from within the -# NeMo-RL container. Wraps 3rdparty/Gym-workspace/Gym/resources_servers/ -# polygon_naming/data/generate_data.py, run through the Gym workspace's -# uv-managed venv (Pillow is not part of the NeMo-RL base env). -# -# Run from WD = /opt/nemo-rl inside the container. -# -# Usage: -# examples/nemo_gym/generate_polygon_naming_data.sh # 512 train / 64 val -# examples/nemo_gym/generate_polygon_naming_data.sh --train 128 --val 16 # custom sizes -# examples/nemo_gym/generate_polygon_naming_data.sh --seed 42 # deterministic - -set -euo pipefail - -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -NEMO_RL_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)" -GYM_ROOT="${NEMO_RL_ROOT}/3rdparty/Gym-workspace/Gym" - -if [[ ! -f "${GYM_ROOT}/pyproject.toml" ]]; then - echo "error: Gym workspace not found at ${GYM_ROOT}" >&2 - exit 1 -fi - -TRAIN_ROWS=512 -VAL_ROWS=64 -SEED=0 - -while [[ $# -gt 0 ]]; do - case "$1" in - --train) TRAIN_ROWS="$2"; shift 2 ;; - --val) VAL_ROWS="$2"; shift 2 ;; - --seed) SEED="$2"; shift 2 ;; - *) echo "unknown arg: $1" >&2; exit 1 ;; - esac -done - -DATA_DIR="${GYM_ROOT}/resources_servers/polygon_naming/data" -GEN="${DATA_DIR}/generate_data.py" - -echo "==> gym root: ${GYM_ROOT}" -echo "==> train: ${DATA_DIR}/train.jsonl (${TRAIN_ROWS} rows, seed ${SEED})" -echo "==> val: ${DATA_DIR}/validation.jsonl (${VAL_ROWS} rows, seed $((SEED + 1)))" - -cd "${GYM_ROOT}" - -# Train and validation are drawn from different seeds so no overlap by -# construction (seeds are independent RNG streams; different num_rows -# further reduces the chance of shared rows). -uv run python "${GEN}" --num-rows "${TRAIN_ROWS}" --seed "${SEED}" \ - --output "${DATA_DIR}/train.jsonl" - -uv run python "${GEN}" --num-rows "${VAL_ROWS}" --seed "$((SEED + 1))" \ - --output "${DATA_DIR}/validation.jsonl" - -echo "==> done" diff --git a/examples/nemo_gym/run_grpo_nemo_gym.py b/examples/nemo_gym/run_grpo_nemo_gym.py index 89081a6f948..3798cc44f1d 100644 --- a/examples/nemo_gym/run_grpo_nemo_gym.py +++ b/examples/nemo_gym/run_grpo_nemo_gym.py @@ -152,8 +152,14 @@ def main() -> None: ) with rl_init_timer.time("tokenizer"): - # setup tokenizer - tokenizer = get_tokenizer(config.policy["tokenizer"]) + is_vlm = bool(config.policy.get("is_vlm")) + if is_vlm: + processor = get_tokenizer(config.policy["tokenizer"], get_processor=True) + tokenizer = processor.tokenizer + else: + processor = None + tokenizer = get_tokenizer(config.policy["tokenizer"]) + assert config.policy["generation"] is not None, ( "A generation config is required for GRPO" ) @@ -171,6 +177,11 @@ def main() -> None: has_refit_draft_weights=has_refit_draft_weights, trains_mtp=trains_mtp, ) + if is_vlm and "vllm_cfg" in config.policy["generation"]: + assert not config.policy["generation"]["vllm_cfg"]["skip_tokenizer_init"], ( + "VLMs require tokenizer to be initialized before generation, " + "so skip_tokenizer_init must be set to False." + ) # NeMo-Gym specific config setup. setup_nemo_gym_config(config, tokenizer) @@ -181,8 +192,9 @@ def main() -> None: # NeMo-Gym environment needs to get dp_openai_server_base_urls from policy_generation, so we don't setup env here. with rl_init_timer.time("data"): print("\n▶ Setting up data...") + data_tokenizer = processor if processor is not None else tokenizer train_dataset, val_dataset = setup_response_data( - tokenizer, config.data, env_configs=None + data_tokenizer, config.data, env_configs=None ) # Validation dataset config setup. @@ -231,7 +243,13 @@ def main() -> None: master_config, teacher_worker_groups, alias_to_group_alias, - ) = setup(config, tokenizer, train_dataset, val_dataset) + ) = setup( + config, + tokenizer, + train_dataset, + val_dataset, + processor=processor, + ) rl_init_timer.record("total", time.perf_counter() - main_start) rl_init_metrics = rl_init_timer.get_timing_metrics(reduction_op="sum") @@ -297,6 +315,7 @@ def main() -> None: max_trajectory_age_steps=config.grpo.async_grpo.max_trajectory_age_steps, teacher_worker_groups=teacher_worker_groups, alias_to_group_alias=alias_to_group_alias, + processor=processor, ) else: print("🚀 Running synchronous GRPO training") diff --git a/examples/nemo_gym/run_multimodal_grpo_nemo_gym.py b/examples/nemo_gym/run_multimodal_grpo_nemo_gym.py deleted file mode 100644 index edee72b90bb..00000000000 --- a/examples/nemo_gym/run_multimodal_grpo_nemo_gym.py +++ /dev/null @@ -1,332 +0,0 @@ -# Copyright (c) 2025, 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. - -import argparse -import os -import pprint -import time - -# Increase the W&B single object size warning threshold. Initially 100_000 (100 KB) -> 10_000_000 (10 MB) -import wandb.util - -wandb.util.VALUE_BYTES_LIMIT = 10_000_000 - -from omegaconf import OmegaConf -from wandb import Table - -from nemo_rl.algorithms.grpo import ( - ColocatablePolicyInterface, - EnvironmentInterface, - GenerationInterface, - Logger, - MasterConfig, - StatefulDataLoader, - TokenizerType, - _should_use_nemo_gym, - grpo_train, - refit_policy_generation, - setup, -) -from nemo_rl.algorithms.utils import get_tokenizer -from nemo_rl.data.utils import setup_response_data -from nemo_rl.distributed.virtual_cluster import init_ray -from nemo_rl.environments.nemo_gym import ( - setup_nemo_gym_config, -) -from nemo_rl.experience.rollouts import run_async_nemo_gym_rollout -from nemo_rl.models.generation import configure_generation_config -from nemo_rl.utils.config import ( - load_config, - parse_hydra_overrides, - register_omegaconf_resolvers, -) -from nemo_rl.utils.logger import get_next_experiment_dir, log_container_init_timing -from nemo_rl.utils.timer import Timer - - -def parse_args() -> tuple[argparse.Namespace, list[str]]: - """Parse command line arguments.""" - parser = argparse.ArgumentParser(description="Run GRPO training with configuration") - parser.add_argument( - "--config", type=str, default=None, help="Path to YAML config file" - ) - - # Parse known args for the script - args, overrides = parser.parse_known_args() - - return args, overrides - - -# These types are directly imported from grpo_train since if something about the architecture changes we want to immediately fail. -def collect_trajectories( - policy: ColocatablePolicyInterface, - policy_generation: GenerationInterface, - val_dataloader: StatefulDataLoader, - tokenizer: TokenizerType, - val_task_to_env: dict[str, EnvironmentInterface], - logger: Logger, - master_config: MasterConfig, -) -> None: - """Run trajectory collection.""" - # common config/state items - colocated_inference = master_config.policy["generation"]["colocated"]["enabled"] - refit_policy_generation(policy, policy_generation, colocated_inference) - - log_filename = "trajectory_collection.jsonl" - - print("\n🔍 Running trajectory collection...", flush=True) - generation_config = master_config.policy["generation"] - for val_batch in val_dataloader: - nemo_gym_rollout_result = run_async_nemo_gym_rollout( - policy_generation=policy_generation, - input_batch=val_batch, - tokenizer=tokenizer, - task_to_env=val_task_to_env, - max_seq_len=master_config.policy["max_total_sequence_length"], - generation_config=generation_config, - max_rollout_turns=None, - greedy=False, - ) - - rows_to_log: list[str] = [] - for key, value in nemo_gym_rollout_result.rollout_metrics.items(): - if "full_result" not in key: - continue - - value: Table - data: list[list[str]] = value.data # (n, 1) - rows_to_log.extend(v[0] for v in data) - - logger.log_string_list_as_jsonl(rows_to_log, log_filename) - - # TODO: eventually as trajectory collection use cases exceed 4 hours, we can leverage the dataloader save functionality to resume - # And also leverage the TimeoutChecker functionality as well - - policy_generation.finish_generation() - - -def main() -> None: - """Main entry point.""" - main_start = time.perf_counter() - log_container_init_timing() - rl_init_timer = Timer(context={"worker": "driver"}) - - register_omegaconf_resolvers() - args, overrides = parse_args() - - if not args.config: - args.config = os.path.join( - os.path.dirname(__file__), - "grpo_workplace_assistant_nemotron_nano_v2_9b.yaml", - ) - - with rl_init_timer.time("config"): - config = load_config(args.config) - print(f"Loaded configuration from: {args.config}") - - if overrides: - print(f"Overrides: {overrides}") - config = parse_hydra_overrides(config, overrides) - - config = OmegaConf.to_container(config, resolve=True) - config = MasterConfig(**config) - print("Applied CLI overrides") - - # Get the next experiment directory with incremented ID - config.logger["log_dir"] = get_next_experiment_dir(config.logger["log_dir"]) - print(f"📊 Using log directory: {config.logger['log_dir']}") - if config.checkpointing["enabled"]: - print( - f"📊 Using checkpoint directory: {config.checkpointing['checkpoint_dir']}" - ) - - with rl_init_timer.time("tokenizer"): - assert config.policy.get("is_vlm", False), ( - "run_multimodal_grpo_nemo_gym.py requires `policy.is_vlm=true` in the config." - ) - processor = get_tokenizer(config.policy["tokenizer"], get_processor=True) - tokenizer = processor.tokenizer - assert config.policy["generation"] is not None, ( - "A generation config is required for GRPO" - ) - config.policy["generation"] = configure_generation_config( - config.policy["generation"], tokenizer - ) - if "vllm_cfg" in config.policy["generation"]: - assert not config.policy["generation"]["vllm_cfg"]["skip_tokenizer_init"], ( - "VLMs require tokenizer to be initialized before generation, so skip_tokenizer_init must be set to False." - ) - - # NeMo-Gym specific config setup. - setup_nemo_gym_config(config, tokenizer) - - # We assert here since this is right after the final config has been materialized. - assert _should_use_nemo_gym(config) - - # NeMo-Gym environment needs to get dp_openai_server_base_urls from policy_generation, so we don't setup env here. - with rl_init_timer.time("data"): - print("\n▶ Setting up data...") - train_dataset, val_dataset = setup_response_data( - processor, config.data, env_configs=None - ) - - # Validation dataset config setup. - if config.grpo["max_val_samples"] is not None: - raise ValueError( - """A non-null `grpo.max_val_samples` parameter is not supported. - -Gym principle is that there is no hidden data pre or post processing from you. What you see is what you get. - -The validation set you pass in will directly be used for validation with no additional preprocessing. If you want to have some number of repetitions, please include that in your dataset, via ``num_repeats``, in your dataset config and `ng_prepare_data` will prepare it accordingly.""" - ) - - if val_dataset is not None: - print( - f"Setting `grpo.max_val_samples` and `grpo.val_batch_size` to the length of the validation dataset, which is {len(val_dataset)}" - ) - config.grpo["max_val_samples"] = len(val_dataset) - config.grpo["val_batch_size"] = config.grpo["max_val_samples"] - - # Print config - print("Final config:") - pprint.pprint(config) - - with rl_init_timer.time("ray_connect"): - init_ray() - - # `is_trajectory_collection` is a NeMo-RL-side control-flow knob; pop it - # before setup() so it is not forwarded into NeMo-Gym's global config (the - # gym actor is now created inside setup()). - is_trajectory_collection = ( - config.env["nemo_gym"].pop("is_trajectory_collection", False) or False - ) - - with rl_init_timer.time("setup"): - ( - policy, - policy_generation, - nemo_gym, - cluster, - dataloader, - val_dataloader, - loss_fn, - logger, - checkpointer, - grpo_state, - master_config, - teacher_worker_groups, - alias_to_group_alias, - ) = setup(config, tokenizer, train_dataset, val_dataset, processor=processor) - - rl_init_timer.record("total", time.perf_counter() - main_start) - rl_init_metrics = rl_init_timer.get_timing_metrics(reduction_op="sum") - print("\n" + "=" * 60) - print(" " * 14 + "RL INIT TIMING BREAKDOWN") - for label, value in sorted(rl_init_metrics.items()): - if isinstance(value, (int, float)): - print(f" {label}: {value:.1f}s") - print("=" * 60 + "\n", flush=True) - - # NeMo-Gym is spun up inside setup() (overlapped with vLLM model load). - # Bind task_to_env and val_task_to_env for the nemo_gym env. - # Hardcode here to match `run_async_nemo_gym_rollout`. - task_to_env = {"nemo_gym": nemo_gym} - val_task_to_env = task_to_env - - if is_trajectory_collection: - collect_trajectories( - policy=policy, - policy_generation=policy_generation, - val_dataloader=val_dataloader, - tokenizer=tokenizer, - val_task_to_env=val_task_to_env, - logger=logger, - master_config=master_config, - ) - # Check if async mode is enabled - elif "async_grpo" in config.grpo and config.grpo["async_grpo"]["enabled"]: - # Async GRPO does not support dynamic sampling, reward scaling, or reward shaping (DAPO features) - unsupported_features = [ - "use_dynamic_sampling", - "reward_scaling", - "reward_shaping", - ] - - for feature in unsupported_features: - if feature not in config.grpo: - continue - - if feature == "use_dynamic_sampling": - if config.grpo[feature]: - raise NotImplementedError( - f"{feature} is not supported with async GRPO" - ) - else: - if config.grpo[feature]["enabled"]: - raise NotImplementedError( - f"{feature} is not supported with async GRPO" - ) - - # Async GRPO does not support multiple dataloaders - if config.data["use_multiple_dataloader"]: - raise NotImplementedError( - "use_multiple_dataloader is not supported with async GRPO" - ) - - from nemo_rl.algorithms.grpo import async_grpo_train - - print("🚀 Running async GRPO training") - - async_config = config.grpo["async_grpo"] - # Run async GRPO training - async_grpo_train( - policy=policy, - policy_generation=policy_generation, - dataloader=dataloader, - val_dataloader=val_dataloader, - tokenizer=tokenizer, - loss_fn=loss_fn, - task_to_env=task_to_env, - val_task_to_env=val_task_to_env, - logger=logger, - checkpointer=checkpointer, - grpo_save_state=grpo_state, - master_config=master_config, - max_trajectory_age_steps=async_config["max_trajectory_age_steps"], - teacher_worker_groups=teacher_worker_groups, - alias_to_group_alias=alias_to_group_alias, - processor=processor, - ) - else: - print("🚀 Running synchronous GRPO training") - - # Run standard GRPO training - grpo_train( - policy, - policy_generation, - dataloader, - val_dataloader, - tokenizer, - loss_fn, - task_to_env, - val_task_to_env, - logger, - checkpointer, - grpo_state, - master_config, - ) - - -if __name__ == "__main__": - main() From 98de6f86714b6efa1bcd170cce95b83a82fa79ee Mon Sep 17 00:00:00 2001 From: rohitrango Date: Tue, 4 Aug 2026 12:46:13 -0700 Subject: [PATCH 07/27] reverted processor design to be consistent with vlm_grpo Signed-off-by: rohitrango (cherry picked from commit b7fc746db7aa03b7d6095052c3cd31fabfefc21a) Signed-off-by: rohitrango --- nemo_rl/algorithms/grpo.py | 25 +-- nemo_rl/data/multimodal_utils.py | 161 ++-------------- nemo_rl/data/processors.py | 69 +++++-- nemo_rl/environments/nemo_gym.py | 42 +++-- tests/unit/data/datasets/test_mmpr_tiny.py | 18 +- .../data/test_multimodal_image_encoding.py | 100 ++++++++++ .../data/test_multimodal_processor_adapter.py | 177 ------------------ .../environments/test_nemo_gym_mm_utils.py | 99 ++++++++++ 8 files changed, 317 insertions(+), 374 deletions(-) create mode 100644 tests/unit/data/test_multimodal_image_encoding.py delete mode 100644 tests/unit/data/test_multimodal_processor_adapter.py create mode 100644 tests/unit/environments/test_nemo_gym_mm_utils.py diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index aeac5058ea2..2f2ba76e9fe 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -1982,22 +1982,11 @@ def _preserve_router_replay_routed_experts( if router_replay_enabled(policy_config) and "routed_experts" in flat_messages: target["routed_experts"] = flat_messages["routed_experts"] - -def _is_vlm_async_run( - master_config: MasterConfig, - processor: Optional[AutoProcessor], -) -> bool: - """Whether this async run carries multimodal tensors that must reach training.""" - return bool(master_config.policy.get("is_vlm") or processor is not None) - - def _build_async_grpo_train_data( flat_messages: BatchedDataDict, input_lengths: torch.Tensor, repeated_batch: BatchedDataDict, policy_config: PolicyConfig, - master_config: MasterConfig, - processor: Optional[AutoProcessor] = None, ) -> BatchedDataDict[ClippedPGLossDataDict]: """Build the async no-TQ policy train batch from flattened rollout messages.""" train_data = BatchedDataDict[ClippedPGLossDataDict]( @@ -2010,15 +1999,8 @@ def _build_async_grpo_train_data( } ) _preserve_router_replay_routed_experts(train_data, flat_messages, policy_config) - + # update multimodal data unconditionally extra_multimodal_data = flat_messages.get_multimodal_dict(as_tensors=False) - if _is_vlm_async_run(master_config, processor) and not extra_multimodal_data: - raise RuntimeError( - "Async GRPO: is_vlm=True (or a processor was provided) but " - "flat_messages.get_multimodal_dict() returned empty for this replay " - "batch. Check that rollout samples carry multimodal tensors and that " - "the collector is not stripping them." - ) train_data.update(extra_multimodal_data) return train_data @@ -3899,7 +3881,6 @@ def async_grpo_train( max_trajectory_age_steps: int = 1, teacher_worker_groups: Optional[dict[str, Any]] = None, alias_to_group_alias: Optional[dict[str, str]] = None, - processor: Optional[AutoProcessor] = None, ) -> None: """Run asynchronous GRPO training with replay buffer. @@ -3917,9 +3898,6 @@ def async_grpo_train( grpo_save_state: Training state master_config: Master configuration max_trajectory_age_steps: Maximum age (in training steps) for trajectories to be used in training - processor: Optional HF processor. Required-in-effect for VLM async runs - so per-batch multimodal tensors get forwarded to policy.get_logprobs - / policy.train (see _build_async_grpo_train_data). """ # Ensure we are running with a compatible async generation backend. # Async GRPO (with in-flight weight updates) supports vLLM, Megatron, and TRT-LLM; @@ -4533,7 +4511,6 @@ def async_grpo_train( repeated_batch, master_config.policy, master_config, - processor=processor, ) train_data.to("cpu") diff --git a/nemo_rl/data/multimodal_utils.py b/nemo_rl/data/multimodal_utils.py index 9f10539f2c4..4cdf7a5fbca 100644 --- a/nemo_rl/data/multimodal_utils.py +++ b/nemo_rl/data/multimodal_utils.py @@ -18,7 +18,7 @@ import re from collections import defaultdict from io import BytesIO -from typing import Any, Optional, Protocol, Union +from typing import Any, Optional, Union import requests import torch @@ -45,6 +45,13 @@ "audio": ["wav", "flac", "mp3"], } +_PLACEHOLDER_STYLE_PROCESSOR_NAMES = frozenset( + { + "NemotronNanoVLV2Processor", + "NemotronH_Nano_Omni_Reasoning_V3Processor", + } +) + # different media namings maybe used in the raw dataset, # in which case, they need to be mapped to the allowed ones @@ -66,85 +73,18 @@ logger = logging.getLogger(__name__) -def _images_from_messages(messages: list[dict[str, Any]]) -> list[Image.Image]: - images = [] - for message in messages: - content = message.get("content") - if not isinstance(content, list): - continue - for part in content: - if isinstance(part, dict) and part.get("type") == "image": - images.append(resolve_to_image(part["image"])) - return images - - -class _HuggingFaceMultimodalProcessorAdapter: - """Adapter for processors supporting multimodal ``apply_chat_template``.""" - - def process( - self, - processor: Any, - messages: list[dict[str, Any]], - *, - add_generation_prompt: bool, - ) -> tuple[str, dict[str, Any]]: - formatted_text = processor.apply_chat_template( - messages, - tokenize=False, - add_generation_prompt=add_generation_prompt, - ) - processed = processor.apply_chat_template( - messages, - tokenize=True, - add_generation_prompt=add_generation_prompt, - return_tensors="pt", - return_dict=True, - ) - return formatted_text, dict(processed) +def uses_image_placeholder(processor: Any) -> bool: + """Return whether a processor requires explicit image placeholders. -def process_multimodal_chat( - processor: Any, - messages: list[dict[str, Any]], - *, - add_generation_prompt: bool, -) -> tuple[str, dict[str, Any]]: - """Render and process multimodal chat through a registered or HF adapter. + Args: + processor: Multimodal processor to classify. - Processors with a nonstandard multimodal calling convention must register an - adapter. All other processors are expected to support Hugging Face's - multimodal ``apply_chat_template`` interface. + Returns: + Whether the processor expands image placeholders through ``__call__`` + rather than tokenized ``apply_chat_template``. """ - adapter = _HuggingFaceMultimodalProcessorAdapter() - formatted_text, processed = adapter.process( - processor, - messages, - add_generation_prompt=add_generation_prompt, - ) - if "input_ids" not in processed: - raise ValueError( - f"{type(processor).__name__} did not return required input_ids." - ) - - images = _images_from_messages(messages) - model_inputs = extract_multimodal_model_inputs(processor, processed) - visual_keys = set( - getattr(getattr(processor, "image_processor", None), "model_input_names", []) - ) - visual_keys.update( - key - for key in get_multimodal_keys_from_processor(processor) - if any(marker in key for marker in ("image", "img", "pixel", "aspect_ratio")) - ) - visual_keys.add("imgs_sizes") - if images and not any(key in model_inputs for key in visual_keys): - raise ValueError( - f"{type(processor).__name__} processed {len(images)} image(s) but " - "returned no visual model inputs. Register a custom multimodal " - "processor adapter if this processor does not support the standard " - "Hugging Face multimodal chat-template interface." - ) - return formatted_text, processed + return type(processor).__name__ in _PLACEHOLDER_STYLE_PROCESSOR_NAMES class PackedTensor: @@ -434,75 +374,6 @@ def get_dim_to_pack_along(processor, key: str) -> int: return 0 -def extract_multimodal_model_inputs( - processor: Any, processed: dict[str, Any] -) -> dict[str, PackedTensor | torch.Tensor]: - """Extract packed visual inputs and sequence-aligned auxiliary tensors. - - Multimodal inputs declared by the processor are wrapped in ``PackedTensor``. - Token-type fields remain ordinary tensors because they align with the full - language-model token sequence. - """ - input_ids = processed.get("input_ids") - if input_ids is None: - raise ValueError("Processor output is missing input_ids.") - if not isinstance(input_ids, torch.Tensor) or input_ids.ndim not in (1, 2): - raise ValueError( - "Processor input_ids must be a one- or two-dimensional torch.Tensor." - ) - if input_ids.ndim == 2 and input_ids.shape[0] != 1: - raise ValueError( - "Multimodal chat processing expects a single conversation, got " - f"input_ids shape {tuple(input_ids.shape)}." - ) - sequence_length = input_ids.shape[-1] - - extracted: dict[str, PackedTensor | torch.Tensor] = {} - multimodal_keys = list(get_multimodal_keys_from_processor(processor)) - # Some remote-code processors omit this per-image input from their declared - # model_input_names even though their model forward requires it. - if "imgs_sizes" in processed and "imgs_sizes" not in multimodal_keys: - multimodal_keys.append("imgs_sizes") - for key in multimodal_keys: - if key not in processed: - continue - value = processed[key] - if not isinstance(value, torch.Tensor): - raise ValueError( - f"Processor model input {key!r} must be a torch.Tensor, got " - f"{type(value).__name__}." - ) - if key == "imgs_sizes": - value = value.to(dtype=torch.int32) - extracted[key] = PackedTensor( - value, dim_to_pack=get_dim_to_pack_along(processor, key) - ) - - for key in ("token_type_ids", "mm_token_type_ids"): - if key not in processed: - continue - value = processed[key] - if not isinstance(value, torch.Tensor) or value.ndim not in (1, 2): - raise ValueError( - f"Processor sequence input {key!r} must be a one- or " - "two-dimensional torch.Tensor." - ) - if value.ndim == 2: - if value.shape[0] != 1: - raise ValueError( - f"Processor sequence input {key!r} must contain one " - f"conversation, got shape {tuple(value.shape)}." - ) - value = value[0] - if len(value) != sequence_length: - raise ValueError( - f"Processor sequence input {key!r} has length {len(value)}, " - f"but input_ids has length {sequence_length}." - ) - extracted[key] = value - return extracted - - def resolve_to_image(image_path_or_image: str | Image.Image) -> Image.Image: """Resolve the image path to a PIL.Image object. diff --git a/nemo_rl/data/processors.py b/nemo_rl/data/processors.py index 6608c257c64..453f2f9c55f 100644 --- a/nemo_rl/data/processors.py +++ b/nemo_rl/data/processors.py @@ -30,7 +30,11 @@ VLMMessageLogType, ) from nemo_rl.data.llm_message_utils import get_formatted_message_log -from nemo_rl.data.multimodal_utils import get_multimodal_keys_from_processor +from nemo_rl.data.multimodal_utils import ( + get_dim_to_pack_along, + get_multimodal_keys_from_processor, + uses_image_placeholder, +) TokenizerType = PreTrainedTokenizerBase @@ -461,9 +465,7 @@ 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, - extract_multimodal_model_inputs, get_multimodal_default_settings_from_processor, - process_multimodal_chat, resolve_to_image, ) @@ -552,15 +554,60 @@ def vlm_hf_data_processor( images = [resolve_to_image(image) for image in images] - # Render once for vLLM and process the identical conversation for MCore. - # Registered adapters cover processors with nonstandard image placeholder - # expansion; standard Hugging Face processors use multimodal chat templates. - string_formatted_dialog, message = process_multimodal_chat( - processor, - [user_message], + # Detect processors that use placeholder style (e.g., NemotronOmni/InternVL) + # vs OpenAI content list style (e.g., Qwen-VL, Gemma). + # These processors expand tokens in __call__ but NOT in apply_chat_template, + # so we must use processor(text=..., images=...) directly. + uses_placeholder = uses_image_placeholder(processor) + + message: dict + if uses_placeholder and images: + # Convert content list to placeholder text format + image_token = getattr(processor, "image_token", "") + text_parts = [] + for content in user_message["content"]: + if content["type"] == "image": + text_parts.append(image_token) + elif content["type"] == "text": + text_parts.append(content["text"]) + user_message_for_tokenize = {"role": "user", "content": "\n".join(text_parts)} + else: + user_message_for_tokenize = user_message + + # get formatted user message + if hasattr(processor, "conversation_preprocessor"): + user_message_for_chat_template = processor.conversation_preprocessor( + user_message + ) + else: + user_message_for_chat_template = user_message_for_tokenize + + string_formatted_dialog = processor.apply_chat_template( + [user_message_for_chat_template], + tokenize=False, add_generation_prompt=True, ) + if uses_placeholder and images: + # Dynamic-resolution path: keep pixel_values in float32 to match vLLM's + # DynamicResolutionImageTiler bit-for-bit. vLLM stores/normalizes in + # float32 and only casts at the vision_model boundary; matching that + # rounding order tightens rollout/train logprob agreement. The model + # forward dispatches on imgs_sizes and handles the bf16 cast. + message = processor( + text=string_formatted_dialog, + images=images, + return_tensors="pt", + ) + else: + message = processor.apply_chat_template( + [user_message_for_tokenize], + tokenize=True, + add_generation_prompt=True, + return_tensors="pt", + return_dict=True, + ) + # 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 @@ -571,7 +618,7 @@ def vlm_hf_data_processor( # the Nemotron Omni path can patchify it and preserve the processor's exact # placeholder count. if ( - _uses_image_placeholder + uses_placeholder and "pixel_values" in message and "imgs_sizes" not in message and message["pixel_values"].ndim == 4 @@ -597,7 +644,7 @@ def vlm_hf_data_processor( user_message[key] = PackedTensor( message[key], dim_to_pack=get_dim_to_pack_along(processor, key), - pad_to_max_shape=_uses_image_placeholder and key == "pixel_values", + 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 diff --git a/nemo_rl/environments/nemo_gym.py b/nemo_rl/environments/nemo_gym.py index 2283fd00c99..6b35359467f 100644 --- a/nemo_rl/environments/nemo_gym.py +++ b/nemo_rl/environments/nemo_gym.py @@ -22,19 +22,18 @@ import ray import torch +from PIL import Image from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy from transformers import PreTrainedTokenizerBase from nemo_rl.distributed.ray_actor_environment_registry import get_actor_python_env -from PIL import Image -from transformers import PreTrainedTokenizerBase - from nemo_rl.data.multimodal_utils import ( PackedTensor, encode_images_in_examples, get_dim_to_pack_along, get_multimodal_keys_from_processor, resolve_to_image, + uses_image_placeholder, ) from nemo_rl.distributed.virtual_cluster import ( @@ -44,6 +43,7 @@ _get_node_ip_local, ) from nemo_rl.environments.interfaces import EnvironmentInterface +from nemo_rl.models.policy import TokenizerConfig from nemo_rl.utils.routed_experts_codec import decode_routed_experts from nemo_rl.utils.timer import Timer from nemo_rl.utils.venvs import create_local_venv_on_each_node @@ -125,7 +125,7 @@ class NemoGymConfig(TypedDict): use_fastokens: NotRequired[bool] # Multimodal fields (populated by `setup_nemo_gym_config` when VLM is enabled). tokenizer_config: NotRequired[ - Optional[Dict[str, Any]] + Optional[TokenizerConfig] ] # For processor reconstruction inside the actor @@ -216,9 +216,7 @@ def _extract_input_images_from_message(item: dict) -> list[Image.Image]: return images -def _index_per_turn_images( - output: list[dict] -) -> list[list[Image.Image]]: +def _index_per_turn_images(output: list[dict]) -> list[list[Image.Image]]: """Bin server-returned user images by the assistant turn that saw them. Walks the Responses-API items in order, accumulating images from user-role @@ -232,7 +230,9 @@ def _index_per_turn_images( for item in output: if item.get("role") == "user": pending.extend(_extract_input_images_from_message(item)) - elif "generation_token_ids" in item: + elif item.get( + "generation_token_ids" + ): # if the generation token ids are empty, skip appending to bucket (`verifiers_agent` and `hermes_agent` can return empty generation token ids list) per_turn.append(pending) pending = [] return per_turn @@ -262,12 +262,25 @@ def _attach_multimodal_data_to_user_message( images=images, return_tensors="pt", ) + uses_placeholder = uses_image_placeholder(processor) multimodal_keys = list(get_multimodal_keys_from_processor(processor)) + # Historical checkpoints may emit dynamic image tiles without imgs_sizes. + # Mirror the media-metadata handling in vlm_hf_data_processor. + if ( + uses_placeholder + and "pixel_values" in processed + and "imgs_sizes" not in processed + and processed["pixel_values"].ndim == 4 + ): + pixel_values = processed["pixel_values"] + num_tiles, _, height, width = pixel_values.shape + processed["imgs_sizes"] = torch.tensor( + [[height, width]] * num_tiles, dtype=torch.long + ) + # imgs_sizes / num_frames are not always declared in model_input_names by - # bundled image processors, so append them explicitly when present. RADIO - # (Nemotron-Omni vision encoder) uses temporal patching even for still - # images and requires one num_frames=1 entry per image/tile — mirror the - # SFT-path fix at nemo_rl/data/processors.py:589-594. + # bundled image processors. RADIO uses temporal patching even for still + # images and requires one num_frames=1 entry per image/tile. if "imgs_sizes" in processed and "imgs_sizes" not in multimodal_keys: multimodal_keys.append("imgs_sizes") if "imgs_sizes" in processed and "num_frames" not in processed: @@ -285,6 +298,7 @@ def _attach_multimodal_data_to_user_message( user_message[key] = PackedTensor( value, dim_to_pack=get_dim_to_pack_along(processor, key), + pad_to_max_shape=uses_placeholder and key == "pixel_values", ) @@ -569,9 +583,7 @@ def _postprocess_nemo_gym_to_nemo_rl_result( if processor is not None: images_this_turn = ( - per_turn_images[turn_idx] - if turn_idx < len(per_turn_images) - else [] + per_turn_images[turn_idx] if turn_idx < len(per_turn_images) else [] ) _attach_multimodal_data_to_user_message( user_message, diff --git a/tests/unit/data/datasets/test_mmpr_tiny.py b/tests/unit/data/datasets/test_mmpr_tiny.py index 74822398d7b..4b0c9e4fc28 100644 --- a/tests/unit/data/datasets/test_mmpr_tiny.py +++ b/tests/unit/data/datasets/test_mmpr_tiny.py @@ -194,7 +194,7 @@ def tiny_image_path(tmp_path): ) -def _run_processor(tiny_image_path): +def _run_processor(tiny_image_path, processor=None): """Helper: run vlm_hf_data_processor on an MMPR sample and return (result DatumSpec, stub processor with captured_call_text).""" from nemo_rl.data.interfaces import TaskDataSpec @@ -202,7 +202,7 @@ def _run_processor(tiny_image_path): task_data_spec = TaskDataSpec(task_name="mmpr-tiny") task_data_spec.prompt = _TEST_PROMPT_TEMPLATE - processor = _make_stub_nemotron_processor() + processor = processor or _make_stub_nemotron_processor() sample = { "images": [tiny_image_path], "question": _RAW_QUESTION, @@ -238,6 +238,20 @@ def test_processor_produces_valid_datum_spec(self, tiny_image_path): assert result["task_name"] == "mmpr-tiny" user_message = result["message_log"][0] assert torch.equal(user_message["num_frames"].as_tensor(), torch.tensor([1])) + assert user_message["pixel_values"].pad_to_max_shape is True + assert user_message["pixel_values"].as_tensor().dtype == torch.float32 + + def test_conversation_preprocessor_is_preserved(self, tiny_image_path): + processor = _make_stub_nemotron_processor() + processor.conversation_preprocessor = MagicMock( + return_value={"role": "user", "content": "preprocessed"} + ) + + result, _ = _run_processor(tiny_image_path, processor=processor) + + processor.conversation_preprocessor.assert_called_once() + assert result["vllm_content"] == "preprocessed" + assert processor.captured_call_text == "preprocessed" def test_historical_tiled_processor_gets_media_metadata(self, tiny_image_path): from nemo_rl.data.interfaces import TaskDataSpec diff --git a/tests/unit/data/test_multimodal_image_encoding.py b/tests/unit/data/test_multimodal_image_encoding.py new file mode 100644 index 00000000000..e36c64135e0 --- /dev/null +++ b/tests/unit/data/test_multimodal_image_encoding.py @@ -0,0 +1,100 @@ +# 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. + +from PIL import Image + +from nemo_rl.data.multimodal_utils import ( + encode_images_in_examples, + image_to_data_url, + resolve_to_image, +) + + +def _example(*content_parts: dict) -> dict: + return { + "responses_create_params": { + "input": [{"role": "user", "content": list(content_parts)}] + } + } + + +def _write_png(tmp_path, name: str, size: tuple[int, int]) -> str: + path = tmp_path / name + Image.new("RGB", size, color=(10, 20, 30)).save(path, format="PNG") + return str(path) + + +def test_image_to_data_url_round_trips_through_resolve_to_image(): + url = image_to_data_url(Image.new("RGB", (4, 3))) + assert url.startswith("data:image/png;base64,") + assert resolve_to_image(url).size == (4, 3) + + +def test_resolve_to_image_accepts_file_scheme(tmp_path): + path = _write_png(tmp_path, "img.png", (5, 6)) + assert resolve_to_image(f"file://{path}").size == (5, 6) + assert resolve_to_image(path).size == (5, 6) + + +def test_encode_images_encodes_local_paths_and_file_urls(tmp_path): + plain = _write_png(tmp_path, "plain.png", (2, 2)) + file_url = "file://" + _write_png(tmp_path, "scheme.png", (3, 3)) + + examples = [ + _example( + {"type": "input_image", "image_url": plain}, + {"type": "input_image", "image_url": {"url": file_url}}, + {"type": "input_text", "text": "describe"}, + ) + ] + encode_images_in_examples(examples) + + parts = examples[0]["responses_create_params"]["input"][0]["content"] + assert parts[0]["image_url"].startswith("data:image/png;base64,") + assert parts[1]["image_url"].startswith("data:image/png;base64,") + assert resolve_to_image(parts[0]["image_url"]).size == (2, 2) + assert resolve_to_image(parts[1]["image_url"]).size == (3, 3) + # Non-image parts are untouched. + assert parts[2] == {"type": "input_text", "text": "describe"} + + +def test_encode_images_passes_through_http_and_data_urls(): + data_url = image_to_data_url(Image.new("RGB", (2, 2))) + examples = [ + _example( + {"type": "input_image", "image_url": "https://example.com/cat.png"}, + {"type": "input_image", "image_url": "http://example.com/dog.png"}, + {"type": "input_image", "image_url": data_url}, + ) + ] + encode_images_in_examples(examples) + + parts = examples[0]["responses_create_params"]["input"][0]["content"] + assert parts[0]["image_url"] == "https://example.com/cat.png" + assert parts[1]["image_url"] == "http://example.com/dog.png" + assert parts[2]["image_url"] == data_url + + +def test_encode_images_is_a_noop_for_text_only_examples(): + examples = [_example({"type": "input_text", "text": "no images here"})] + before = [ + dict(part) + for part in examples[0]["responses_create_params"]["input"][0]["content"] + ] + assert encode_images_in_examples(examples) is examples + assert examples[0]["responses_create_params"]["input"][0]["content"] == before + + # Missing/oddly-shaped payloads must not raise. + assert encode_images_in_examples([{}, {"responses_create_params": {}}]) is not None + assert encode_images_in_examples([{"responses_create_params": {"input": "nope"}}]) diff --git a/tests/unit/data/test_multimodal_processor_adapter.py b/tests/unit/data/test_multimodal_processor_adapter.py deleted file mode 100644 index 5e9ccbd1c96..00000000000 --- a/tests/unit/data/test_multimodal_processor_adapter.py +++ /dev/null @@ -1,177 +0,0 @@ -# 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. - -import pytest -import torch -from PIL import Image - -from nemo_rl.data.multimodal_utils import ( - PackedTensor, - extract_multimodal_model_inputs, - process_multimodal_chat, - register_multimodal_processor_adapter, -) - - -def _messages() -> list[dict]: - return [ - { - "role": "user", - "content": [ - {"type": "image", "image": Image.new("RGB", (2, 2))}, - {"type": "text", "text": "describe"}, - ], - } - ] - - -class _ImageProcessor: - model_input_names = ["pixel_values", "image_grid_thw"] - - -class _Tokenizer: - model_input_names = ["input_ids", "attention_mask"] - - -class StandardProcessor: - image_processor = _ImageProcessor() - tokenizer = _Tokenizer() - model_input_names = [ - "input_ids", - "attention_mask", - "pixel_values", - "image_grid_thw", - ] - - def __init__(self): - self.tokenized_messages = None - - def apply_chat_template( - self, messages, *, tokenize, add_generation_prompt, **kwargs - ): - assert add_generation_prompt - if not tokenize: - return "rendered" - self.tokenized_messages = messages - return { - "input_ids": torch.tensor([[1, 2, 3]]), - "pixel_values": torch.ones(2, 3, 2, 2), - "image_grid_thw": torch.tensor([[1, 2, 2]]), - "token_type_ids": torch.tensor([[0, 1, 1]]), - "mm_token_type_ids": torch.tensor([[0, 1, 1]]), - } - - -def test_standard_hf_adapter_and_model_input_extraction(): - processor = StandardProcessor() - formatted, processed = process_multimodal_chat( - processor, _messages(), add_generation_prompt=True - ) - model_inputs = extract_multimodal_model_inputs(processor, processed) - - assert formatted == "rendered" - assert processor.tokenized_messages[0]["content"][0]["type"] == "image" - assert isinstance( - processor.tokenized_messages[0]["content"][0]["image"], Image.Image - ) - assert isinstance(model_inputs["pixel_values"], PackedTensor) - assert isinstance(model_inputs["image_grid_thw"], PackedTensor) - assert model_inputs["token_type_ids"].tolist() == [0, 1, 1] - assert model_inputs["mm_token_type_ids"].tolist() == [0, 1, 1] - - -def test_smolvlm_inputs_pack_along_dimension_one(): - class SmolVLMProcessor(StandardProcessor): - pass - - processor = SmolVLMProcessor() - _, processed = process_multimodal_chat( - processor, _messages(), add_generation_prompt=True - ) - model_inputs = extract_multimodal_model_inputs(processor, processed) - assert model_inputs["pixel_values"].dim_to_pack == 1 - assert model_inputs["image_grid_thw"].dim_to_pack == 1 - - -def test_registered_nemotron_placeholder_adapter(): - class NemotronNanoVLV2Processor(StandardProcessor): - image_token = "" - - def __init__(self): - self.call = None - - def apply_chat_template(self, messages, **kwargs): - assert messages[0]["content"] == "\ndescribe" - return messages[0]["content"] - - def __call__(self, *, text, images, return_tensors): - self.call = (text, images, return_tensors) - return { - "input_ids": torch.tensor([[1, 2, 3]]), - "pixel_values": torch.ones(1, 3, 2, 2), - } - - processor = NemotronNanoVLV2Processor() - formatted, _ = process_multimodal_chat( - processor, _messages(), add_generation_prompt=True - ) - assert formatted == "\ndescribe" - assert processor.call[0] == formatted - assert len(processor.call[1]) == 1 - - -def test_custom_processor_adapter_registration(): - class CustomProcessor(StandardProcessor): - pass - - class CustomAdapter: - def process(self, processor, messages, *, add_generation_prompt): - assert add_generation_prompt - return "custom", { - "input_ids": torch.tensor([[4, 5]]), - "pixel_values": torch.ones(1, 3, 2, 2), - } - - register_multimodal_processor_adapter("CustomProcessor", CustomAdapter()) - formatted, processed = process_multimodal_chat( - CustomProcessor(), _messages(), add_generation_prompt=True - ) - assert formatted == "custom" - assert processed["input_ids"].tolist() == [[4, 5]] - - -def test_images_without_visual_model_inputs_fail_loudly(): - class MissingVisualProcessor(StandardProcessor): - def apply_chat_template( - self, messages, *, tokenize, add_generation_prompt, **kwargs - ): - if not tokenize: - return "rendered" - return {"input_ids": torch.tensor([[1, 2, 3]])} - - with pytest.raises(ValueError, match="returned no visual model inputs"): - process_multimodal_chat( - MissingVisualProcessor(), _messages(), add_generation_prompt=True - ) - - -@pytest.mark.parametrize("key", ["token_type_ids", "mm_token_type_ids"]) -def test_malformed_sequence_auxiliary_length_fails_loudly(key): - processor = StandardProcessor() - processed = { - "input_ids": torch.tensor([[1, 2, 3]]), - key: torch.tensor([[0, 1]]), - } - with pytest.raises(ValueError, match=f"{key!r} has length 2"): - extract_multimodal_model_inputs(processor, processed) diff --git a/tests/unit/environments/test_nemo_gym_mm_utils.py b/tests/unit/environments/test_nemo_gym_mm_utils.py new file mode 100644 index 00000000000..626e4bb1355 --- /dev/null +++ b/tests/unit/environments/test_nemo_gym_mm_utils.py @@ -0,0 +1,99 @@ +from PIL import Image + +from nemo_rl.data.multimodal_utils import image_to_data_url +from nemo_rl.environments.nemo_gym import ( + _extract_input_images_from_message, + _index_per_turn_images, +) + + +def _image(size: tuple[int, int]) -> str: + """Return a data URL for a solid RGB image of the given size.""" + return image_to_data_url(Image.new("RGB", size)) + + +def _user(*data_urls: str) -> dict: + return { + "role": "user", + "content": [{"type": "input_image", "image_url": url} for url in data_urls], + } + + +def _assistant(token_ids: list[int]) -> dict: + return {"role": "assistant", "generation_token_ids": token_ids} + + +def test_extract_input_images_handles_flat_and_dict_image_url(): + item = { + "role": "user", + "content": [ + {"type": "input_image", "image_url": _image((2, 2))}, + {"type": "input_image", "image_url": {"url": _image((3, 3))}}, + {"type": "input_text", "text": "ignore me"}, + ], + } + images = _extract_input_images_from_message(item) + assert [img.size for img in images] == [(2, 2), (3, 3)] + + +def test_extract_input_images_returns_empty_for_string_content(): + assert _extract_input_images_from_message({"role": "user", "content": "hi"}) == [] + assert _extract_input_images_from_message({"role": "user"}) == [] + + +def test_index_per_turn_images_bins_seed_and_intermediate_images(): + seed_obs = [_user(_image((2, 2)))] + output = [ + _assistant([1, 2]), + _user(_image((3, 3)), _image((4, 4))), + _assistant([3, 4]), + ] + per_turn = _index_per_turn_images(seed_obs, output) + + assert len(per_turn) == 2 + assert [img.size for img in per_turn[0]] == [(2, 2)] + assert [img.size for img in per_turn[1]] == [(3, 3), (4, 4)] + + +def test_index_per_turn_images_text_only_rollout_yields_empty_buckets(): + output = [ + {"role": "user", "content": "solve this"}, + _assistant([1, 2]), + {"role": "user", "content": "and this"}, + _assistant([3, 4]), + ] + assert _index_per_turn_images([], output) == [[], []] + + +def test_index_per_turn_images_assigns_tool_result_image_to_next_turn(): + """A tool-result image contributes to the following assistant turn.""" + output = [ + _user(_image((2, 2))), + _assistant([1, 2]), + {"type": "function_call_output", "output": _image((5, 5))}, + _assistant([3, 4]), + ] + per_turn = _index_per_turn_images([], output) + + assert len(per_turn) == 2 + assert [img.size for img in per_turn[0]] == [(2, 2)] + assert [img.size for img in per_turn[1]] == [(5, 5)] + + +def test_index_per_turn_images_aligns_with_postprocess_skip_of_empty_generations(): + """Turns skipped by the postprocess loop must not consume an image bucket. + + ``_postprocess_nemo_gym_to_nemo_rl_result`` skips output items whose + ``generation_token_ids`` is present but empty, so the bucket list must skip + them too or every later turn is attached to the wrong images. + """ + output = [ + _user(_image((2, 2))), + _assistant([]), # all-EOS generation, skipped by the postprocess loop + _user(_image((6, 6))), + _assistant([7, 8]), + ] + per_turn = _index_per_turn_images([], output) + + assert len(per_turn) == 1 + assert [img.size for img in per_turn[0]] == [(2, 2), (6, 6)] From a50f7720c92af7d4d6547f72ebddbe8a09172e1c Mon Sep 17 00:00:00 2001 From: rohitrango Date: Wed, 5 Aug 2026 09:48:43 -0700 Subject: [PATCH 08/27] chore: clean up multimodal processor plumbing Signed-off-by: rohitrango --- nemo_rl/algorithms/grpo.py | 1 - nemo_rl/data/processors.py | 9 +++------ 2 files changed, 3 insertions(+), 7 deletions(-) diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index 2f2ba76e9fe..bc841990bac 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -4510,7 +4510,6 @@ def async_grpo_train( input_lengths, repeated_batch, master_config.policy, - master_config, ) train_data.to("cpu") diff --git a/nemo_rl/data/processors.py b/nemo_rl/data/processors.py index 453f2f9c55f..adad65b4d5e 100644 --- a/nemo_rl/data/processors.py +++ b/nemo_rl/data/processors.py @@ -30,11 +30,6 @@ VLMMessageLogType, ) from nemo_rl.data.llm_message_utils import get_formatted_message_log -from nemo_rl.data.multimodal_utils import ( - get_dim_to_pack_along, - get_multimodal_keys_from_processor, - uses_image_placeholder, -) TokenizerType = PreTrainedTokenizerBase @@ -465,6 +460,9 @@ 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, + get_multimodal_keys_from_processor, + uses_image_placeholder, get_multimodal_default_settings_from_processor, resolve_to_image, ) @@ -560,7 +558,6 @@ def vlm_hf_data_processor( # so we must use processor(text=..., images=...) directly. uses_placeholder = uses_image_placeholder(processor) - message: dict if uses_placeholder and images: # Convert content list to placeholder text format image_token = getattr(processor, "image_token", "") From f6300f861b1902d62b2b740aa2c5c35c5a6ffab4 Mon Sep 17 00:00:00 2001 From: rohitrango Date: Tue, 4 Aug 2026 13:54:42 -0700 Subject: [PATCH 09/27] change per-turn images to get results from tool-calls, etc (anything except assistant) Signed-off-by: rohitrango (cherry picked from commit 85a24b80b107c9fbdd9a6e899a70aab29c196b2c) Signed-off-by: rohitrango --- nemo_rl/environments/nemo_gym.py | 21 ++++++++++++++----- .../environments/test_nemo_gym_mm_utils.py | 12 +++++------ 2 files changed, 22 insertions(+), 11 deletions(-) diff --git a/nemo_rl/environments/nemo_gym.py b/nemo_rl/environments/nemo_gym.py index 6b35359467f..2489a941e38 100644 --- a/nemo_rl/environments/nemo_gym.py +++ b/nemo_rl/environments/nemo_gym.py @@ -195,8 +195,18 @@ def _detect_invalid_tool_call_and_malformed_thinking( def _extract_input_images_from_message(item: dict) -> list[Image.Image]: - """Pull PIL images out of a Responses-API user-role item's content list.""" + """Pull PIL images out of a non-assistant Responses-API item. + + Handles both content-list items (user / tool messages carrying + ``input_image``/``image``/``image_url`` parts) and ``function_call_output`` + items whose ``output`` field is an image data URL. + """ images: list[Image.Image] = [] + if item.get("type") == "function_call_output": + src = item.get("output") + if isinstance(src, str): + images.append(resolve_to_image(src)) + return images content = item.get("content") or [] if not isinstance(content, list): return images @@ -217,10 +227,11 @@ def _extract_input_images_from_message(item: dict) -> list[Image.Image]: def _index_per_turn_images(output: list[dict]) -> list[list[Image.Image]]: - """Bin server-returned user images by the assistant turn that saw them. + """Bin server-returned images by the assistant turn that saw them. - Walks the Responses-API items in order, accumulating images from user-role - items into a pending list, and flushing them into a per-turn bucket each + Walks the Responses-API items in order, accumulating images from every + non-assistant item (user turns, tool messages, ``function_call_output``, + etc.) into a pending list, and flushing them into a per-turn bucket each time a trainable assistant item (one carrying ``generation_token_ids``) is reached. The returned list has one entry per trainable assistant turn, aligned with the postprocess loop's ``turn_idx``. @@ -228,7 +239,7 @@ def _index_per_turn_images(output: list[dict]) -> list[list[Image.Image]]: per_turn: list[list[Image.Image]] = [] pending: list[Image.Image] = [] for item in output: - if item.get("role") == "user": + if item.get("role") != "assistant": pending.extend(_extract_input_images_from_message(item)) elif item.get( "generation_token_ids" diff --git a/tests/unit/environments/test_nemo_gym_mm_utils.py b/tests/unit/environments/test_nemo_gym_mm_utils.py index 626e4bb1355..3b452e6db1c 100644 --- a/tests/unit/environments/test_nemo_gym_mm_utils.py +++ b/tests/unit/environments/test_nemo_gym_mm_utils.py @@ -41,14 +41,14 @@ def test_extract_input_images_returns_empty_for_string_content(): assert _extract_input_images_from_message({"role": "user"}) == [] -def test_index_per_turn_images_bins_seed_and_intermediate_images(): - seed_obs = [_user(_image((2, 2)))] +def test_index_per_turn_images_bins_images(): output = [ + _user(_image((2, 2))), _assistant([1, 2]), _user(_image((3, 3)), _image((4, 4))), _assistant([3, 4]), ] - per_turn = _index_per_turn_images(seed_obs, output) + per_turn = _index_per_turn_images(output) assert len(per_turn) == 2 assert [img.size for img in per_turn[0]] == [(2, 2)] @@ -62,7 +62,7 @@ def test_index_per_turn_images_text_only_rollout_yields_empty_buckets(): {"role": "user", "content": "and this"}, _assistant([3, 4]), ] - assert _index_per_turn_images([], output) == [[], []] + assert _index_per_turn_images(output) == [[], []] def test_index_per_turn_images_assigns_tool_result_image_to_next_turn(): @@ -73,7 +73,7 @@ def test_index_per_turn_images_assigns_tool_result_image_to_next_turn(): {"type": "function_call_output", "output": _image((5, 5))}, _assistant([3, 4]), ] - per_turn = _index_per_turn_images([], output) + per_turn = _index_per_turn_images(output) assert len(per_turn) == 2 assert [img.size for img in per_turn[0]] == [(2, 2)] @@ -93,7 +93,7 @@ def test_index_per_turn_images_aligns_with_postprocess_skip_of_empty_generations _user(_image((6, 6))), _assistant([7, 8]), ] - per_turn = _index_per_turn_images([], output) + per_turn = _index_per_turn_images(output) assert len(per_turn) == 1 assert [img.size for img in per_turn[0]] == [(2, 2), (6, 6)] From 629a1d42ad10564ec9c0a7b3154bd87a6c3bc104 Mon Sep 17 00:00:00 2001 From: rohitrango Date: Tue, 4 Aug 2026 14:43:03 -0700 Subject: [PATCH 10/27] fix(nemo_gym): flush per-turn image bucket on any trainable item MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `_index_per_turn_images` gated its bucket flush on `role == "assistant"`, but `_postprocess_nemo_gym_to_nemo_rl_result` treats every item carrying truthy `generation_token_ids` as a trainable turn — including reasoning-only responses and `function_call` items whose role is not `"assistant"`. The mismatch left the batched flatten path with a `PackedTensor` for normal assistant turns and a missing entry for reasoning/tool-call turns, crashing `PackedTensor.flattened_concat` on async multimodal GRPO runs. Gate the flush on `generation_token_ids` directly so the per-turn image list stays aligned with the postprocess loop's `turn_idx`, and add regression tests for the reasoning-only and function_call cases. Also drop the unused `processor` kwarg from the async GRPO call site. Signed-off-by: rohitrango (cherry picked from commit 4c2537b5c37eaec3af8cc907c9f1483af8936de4) Signed-off-by: rohitrango --- examples/nemo_gym/run_grpo_nemo_gym.py | 1 - nemo_rl/environments/nemo_gym.py | 28 ++++++------ .../environments/test_nemo_gym_mm_utils.py | 43 +++++++++++++++++++ 3 files changed, 59 insertions(+), 13 deletions(-) diff --git a/examples/nemo_gym/run_grpo_nemo_gym.py b/examples/nemo_gym/run_grpo_nemo_gym.py index 3798cc44f1d..4d83b9c129f 100644 --- a/examples/nemo_gym/run_grpo_nemo_gym.py +++ b/examples/nemo_gym/run_grpo_nemo_gym.py @@ -315,7 +315,6 @@ def main() -> None: max_trajectory_age_steps=config.grpo.async_grpo.max_trajectory_age_steps, teacher_worker_groups=teacher_worker_groups, alias_to_group_alias=alias_to_group_alias, - processor=processor, ) else: print("🚀 Running synchronous GRPO training") diff --git a/nemo_rl/environments/nemo_gym.py b/nemo_rl/environments/nemo_gym.py index 2489a941e38..2f93eb45693 100644 --- a/nemo_rl/environments/nemo_gym.py +++ b/nemo_rl/environments/nemo_gym.py @@ -227,25 +227,29 @@ def _extract_input_images_from_message(item: dict) -> list[Image.Image]: def _index_per_turn_images(output: list[dict]) -> list[list[Image.Image]]: - """Bin server-returned images by the assistant turn that saw them. - - Walks the Responses-API items in order, accumulating images from every - non-assistant item (user turns, tool messages, ``function_call_output``, - etc.) into a pending list, and flushing them into a per-turn bucket each - time a trainable assistant item (one carrying ``generation_token_ids``) is - reached. The returned list has one entry per trainable assistant turn, - aligned with the postprocess loop's ``turn_idx``. + """Bin server-returned images by the trainable turn that saw them. + + Walks the Responses-API items in order and flushes ``pending`` into a + per-turn bucket each time it hits an item carrying truthy + ``generation_token_ids`` — matching the exact gate that + ``_postprocess_nemo_gym_to_nemo_rl_result`` uses to decide which items + become trainable turns. Every other item (user turns, tool messages, + ``function_call_output``, non-trainable reasoning) contributes its images + to ``pending`` for the next trainable turn. This ensures the returned list + has one entry per trainable turn, aligned with the postprocess loop's + ``turn_idx`` even when the trainable item's role is not ``assistant`` + (e.g. a reasoning-only response, or a ``function_call``). """ per_turn: list[list[Image.Image]] = [] pending: list[Image.Image] = [] for item in output: - if item.get("role") != "assistant": - pending.extend(_extract_input_images_from_message(item)) - elif item.get( + if item.get( "generation_token_ids" - ): # if the generation token ids are empty, skip appending to bucket (`verifiers_agent` and `hermes_agent` can return empty generation token ids list) + ): # trainable turn; empty generation_token_ids is skipped by the postprocess loop and must not consume a bucket per_turn.append(pending) pending = [] + elif item.get("role") != "assistant": + pending.extend(_extract_input_images_from_message(item)) return per_turn diff --git a/tests/unit/environments/test_nemo_gym_mm_utils.py b/tests/unit/environments/test_nemo_gym_mm_utils.py index 3b452e6db1c..7e77a624d6f 100644 --- a/tests/unit/environments/test_nemo_gym_mm_utils.py +++ b/tests/unit/environments/test_nemo_gym_mm_utils.py @@ -97,3 +97,46 @@ def test_index_per_turn_images_aligns_with_postprocess_skip_of_empty_generations assert len(per_turn) == 1 assert [img.size for img in per_turn[0]] == [(2, 2), (6, 6)] + + +def test_index_per_turn_images_flushes_on_non_assistant_trainable_item(): + """Trainable items whose role is not ``assistant`` (reasoning-only responses, + function_call items) still carry ``generation_token_ids`` and are treated as + turns by the postprocess loop. The per-turn image bucket must flush for them + too, or the batched flatten path will see a ``PackedTensor`` for turns + where the model produced a normal assistant message and a missing key for + turns where it produced only reasoning — crashing + ``PackedTensor.flattened_concat`` on the None entry. + """ + reasoning_only = {"type": "reasoning", "generation_token_ids": [9, 10]} + output = [ + _user(_image((2, 2))), + reasoning_only, + ] + per_turn = _index_per_turn_images(output) + + assert len(per_turn) == 1 + assert [img.size for img in per_turn[0]] == [(2, 2)] + + +def test_index_per_turn_images_flushes_on_function_call_trainable_item(): + """Same as the reasoning-only case, but for tool-calling turns where the + model call's last output item is a ``function_call`` (no ``role`` field).""" + function_call = { + "type": "function_call", + "name": "tool", + "arguments": "{}", + "call_id": "c1", + "generation_token_ids": [11, 12], + } + output = [ + _user(_image((2, 2))), + function_call, + {"type": "function_call_output", "output": _image((5, 5)), "call_id": "c1"}, + _assistant([13, 14]), + ] + per_turn = _index_per_turn_images(output) + + assert len(per_turn) == 2 + assert [img.size for img in per_turn[0]] == [(2, 2)] + assert [img.size for img in per_turn[1]] == [(5, 5)] From 171506f4ca0c3b85730e4786cf898c83986a78ff Mon Sep 17 00:00:00 2001 From: rohitrango Date: Tue, 4 Aug 2026 15:10:15 -0700 Subject: [PATCH 11/27] docs: add Google-style docstrings to image encoding helpers Add Args/Returns sections to image_to_data_url and encode_images_in_examples in nemo_rl/data/multimodal_utils.py. Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: rohitrango (cherry picked from commit e9b89bc8806b538c8855a59ea27e4c8c238537dd) Signed-off-by: rohitrango --- nemo_rl/data/multimodal_utils.py | 37 +++++++++++++++++++++++++++----- 1 file changed, 32 insertions(+), 5 deletions(-) diff --git a/nemo_rl/data/multimodal_utils.py b/nemo_rl/data/multimodal_utils.py index 4cdf7a5fbca..368608f5909 100644 --- a/nemo_rl/data/multimodal_utils.py +++ b/nemo_rl/data/multimodal_utils.py @@ -404,7 +404,18 @@ def resolve_to_image(image_path_or_image: str | Image.Image) -> Image.Image: def image_to_data_url(image: Image.Image, fmt: str = "PNG") -> str: - """Encode a PIL Image as a base64 data URL.""" + """Encode a PIL Image as a base64 ``data:`` URL. + + Args: + image: PIL image to encode. + fmt: PIL image format used for serialization (e.g. ``"PNG"``, ``"JPEG"``). + The value is also lowercased and embedded in the MIME type of the + returned URL. + + Returns: + A ``data:image/;base64,`` URL suitable for embedding in + an OpenAI Responses ``input_image`` content part. + """ buf = BytesIO() image.save(buf, format=fmt) encoded = base64.b64encode(buf.getvalue()).decode("utf-8") @@ -412,11 +423,27 @@ def image_to_data_url(image: Image.Image, fmt: str = "PNG") -> str: def encode_images_in_examples(nemo_gym_examples: list[dict]) -> list[dict]: - """Walk examples and replace local image paths with base64 data URLs. + """Replace local image paths in NeMo Gym examples with base64 data URLs. - Operates in-place on each example's responses_create_params.input[].content[] - items of type 'input_image'. HTTP(S) and data URLs are preserved; local - paths, including file:// URLs, are encoded as data URLs. + Walks each example's ``responses_create_params.input[].content[]`` items + and rewrites any ``input_image`` part whose ``image_url`` is a local path + (or ``file://`` URL) into a base64 ``data:`` URL via + :func:`image_to_data_url`. Parts whose URL already starts with ``http://``, + ``https://``, or ``data:`` are left untouched. Malformed items (non-dict + entries, missing/empty URLs, non-list ``input``/``content``) are skipped + without raising. + + The examples are mutated in place; the same list is also returned for + convenience so callers can chain the call. + + Args: + nemo_gym_examples: List of NeMo Gym example dicts. Each example is + expected to contain a ``responses_create_params`` mapping with an + ``input`` list of Responses API messages. + + Returns: + The same ``nemo_gym_examples`` list, with local image references + rewritten to base64 data URLs in place. """ for example in nemo_gym_examples: input_items = example.get("responses_create_params", {}).get("input", []) From a104ae648f4539c14dad5d6578cfb95c503169e8 Mon Sep 17 00:00:00 2001 From: rohitrango Date: Tue, 4 Aug 2026 15:43:49 -0700 Subject: [PATCH 12/27] change non-default config option Signed-off-by: rohitrango (cherry picked from commit d60004f69be9318b90b1db99f70a56c0eb949148) Signed-off-by: rohitrango --- nemo_rl/environments/nemo_gym.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/nemo_rl/environments/nemo_gym.py b/nemo_rl/environments/nemo_gym.py index 2f93eb45693..e3eda985529 100644 --- a/nemo_rl/environments/nemo_gym.py +++ b/nemo_rl/environments/nemo_gym.py @@ -788,7 +788,7 @@ def setup_nemo_gym_config(config, tokenizer) -> None: # For VLM runs, plumb the tokenizer config into the gym env config so the # NemoGym actor can reconstruct the processor inside itself (needed for # multi-turn multimodal postprocessing). - if config.policy.get("is_vlm", False): + if config.policy.get("is_vlm"): env_cfg = config.env.setdefault("nemo_gym", {}) env_cfg.setdefault("tokenizer_config", dict(config.policy["tokenizer"])) From d3e02b482afa3d7c24b1dab89261def3cc21a3ef Mon Sep 17 00:00:00 2001 From: rohitrango Date: Wed, 5 Aug 2026 09:54:48 -0700 Subject: [PATCH 13/27] build: preserve main dependency configuration Signed-off-by: rohitrango --- pyproject.toml | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 33e071b2d85..9958048b578 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -171,17 +171,17 @@ mcore = [ # sudo dpkg -i cuda-keyring_1.1-1_all.deb # sudo apt-get update # sudo apt-get install cudnn-cuda-13 - # This dependency also needs to be compatible with the spec in Megatron-Bridge/pyproject.toml. - # It is specified here since we don't directly use Megatron-Bridge/pyproject.toml, but a proxy setup.py+pyproject.toml combo - # outside to allow "optionally" installing the megatron path. It's simpler to deal with transformer-engine here in the NeMo RL pyproject.toml - "transformer-engine[pytorch,core_cu13] @ git+https://github.com/NVIDIA/TransformerEngine.git@release_v2.15", - "megatron-core", - "megatron-bridge", - "megatron-energon[av-decode]~=7.0", - # Must match Megatron-Bridge main's transformers requirement (still includes GLM 5.1 support). - "transformers>=5.8.1,<5.9.0", - "nvidia-modelopt[torch]; sys_platform != 'darwin'", - "onnxscript", + + # megatron-bridge is installed straight from the submodule's own pyproject.toml, so its + # dependencies (transformers, transformer-engine, megatron-core[dev,mlm] -> nvidia-modelopt, + # onnxscript, flash-linear-attention, ...) flow transitively and pick up upstream changes on + # every submodule bump + `uv lock`. Do not mirror them here. transformer-engine's exact + # version is still pinned globally via [tool.uv] override-dependencies. + # [te]/[ssm] activate megatron-bridge's transformer-engine and mamba-ssm/causal-conv1d extras. + # megatron-core is intentionally NOT listed: it flows from megatron-bridge's own + # dependency megatron-core[dev,mlm] and its [tool.uv.sources] path mapping, so Megatron-LM + # pyproject changes propagate on bump without any mirror here. + "megatron-bridge[te,ssm]", # Flash-attn version should be selected to satisfy both TE + vLLM requirements (xformers in particular) # https://github.com/NVIDIA/TransformerEngine/blob/v2.3/transformer_engine/pytorch/attention/dot_product_attention/utils.py#L108 # https://github.com/facebookresearch/xformers/blob/8354497deb2c04c67fbb2e2ad911e86530da0e90/xformers/ops/fmha/flash.py#L76 From 428f5ab4c85b07995f3604e8417d4cedee969dc6 Mon Sep 17 00:00:00 2001 From: rohitrango Date: Wed, 5 Aug 2026 10:05:53 -0700 Subject: [PATCH 14/27] (chore): add copyright notice to test Signed-off-by: rohitrango --- tests/unit/environments/test_nemo_gym_mm_utils.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/tests/unit/environments/test_nemo_gym_mm_utils.py b/tests/unit/environments/test_nemo_gym_mm_utils.py index 7e77a624d6f..71ae0c9c25c 100644 --- a/tests/unit/environments/test_nemo_gym_mm_utils.py +++ b/tests/unit/environments/test_nemo_gym_mm_utils.py @@ -1,3 +1,17 @@ +# 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. + from PIL import Image from nemo_rl.data.multimodal_utils import image_to_data_url From 33b70ec26ad5d541da8579c31fa205c4d814ffde Mon Sep 17 00:00:00 2001 From: rohitrango Date: Wed, 5 Aug 2026 10:13:02 -0700 Subject: [PATCH 15/27] (chore): undo vllm chat request change Signed-off-by: rohitrango --- .../generation/vllm/vllm_worker_async.py | 19 ------------------- 1 file changed, 19 deletions(-) diff --git a/nemo_rl/models/generation/vllm/vllm_worker_async.py b/nemo_rl/models/generation/vllm/vllm_worker_async.py index fb993f5de8a..201d0790044 100644 --- a/nemo_rl/models/generation/vllm/vllm_worker_async.py +++ b/nemo_rl/models/generation/vllm/vllm_worker_async.py @@ -682,14 +682,6 @@ async def create_chat_completion( assert request.temperature == generation_config["temperature"] assert request.top_p == generation_config["top_p"] - # Merge recipe-level chat_template_kwargs into the request. Client- - # provided keys win so a caller can still override per request. - if default_chat_template_kwargs: - request.chat_template_kwargs = { - **default_chat_template_kwargs, - **(request.chat_template_kwargs or {}), - } - try: generator = await openai_serving_chat.create_chat_completion( request, raw_request @@ -756,17 +748,6 @@ class NeMoRLServingTokenization(ServingTokenization): @app.post("/tokenize") async def tokenize(request: NeMoRLTokenizeRequest, raw_request: Request): - # Chat-mode tokenize also renders the chat template — inject the - # same default kwargs so /tokenize and /v1/chat/completions produce - # identical prompt tokens under multi-turn. - if default_chat_template_kwargs and hasattr( - request, "chat_template_kwargs" - ): - request.chat_template_kwargs = { - **default_chat_template_kwargs, - **(request.chat_template_kwargs or {}), - } - generator = await openai_serving_tokenization.create_tokenize( request, raw_request ) From ef3f1630e6202b0fe8c9eb911579bf1314ca9682 Mon Sep 17 00:00:00 2001 From: rohitrango Date: Wed, 5 Aug 2026 13:53:30 -0700 Subject: [PATCH 16/27] fix: address NeMo Gym multimodal image indexing issues Ignore text-only function call outputs during image extraction and seed the first trainable turn with images from the initial input messages. Signed-off-by: rohitrango --- nemo_rl/environments/nemo_gym.py | 46 +++++++++++++++++-- .../environments/test_nemo_gym_mm_utils.py | 22 +++++++++ 2 files changed, 63 insertions(+), 5 deletions(-) diff --git a/nemo_rl/environments/nemo_gym.py b/nemo_rl/environments/nemo_gym.py index e3eda985529..9fc7dd7e4bb 100644 --- a/nemo_rl/environments/nemo_gym.py +++ b/nemo_rl/environments/nemo_gym.py @@ -194,17 +194,33 @@ def _detect_invalid_tool_call_and_malformed_thinking( ######################################## +_IMAGE_SRC_PREFIXES = ("data:image/", "http://", "https://", "file://") + + +def _looks_like_image_src(src: str) -> bool: + """True when ``src`` plausibly points at an image the loader can open. + + Guards against tool responses (e.g. ``{"x": 0.65, "y": 0.83}`` from a + click tool) that are strings but not image URLs. Without this, the + indexer forwards the JSON payload to ``resolve_to_image`` → PIL.open, + which treats it as a filesystem path and raises ``FileNotFoundError``. + """ + return src.startswith(_IMAGE_SRC_PREFIXES) + + def _extract_input_images_from_message(item: dict) -> list[Image.Image]: """Pull PIL images out of a non-assistant Responses-API item. Handles both content-list items (user / tool messages carrying ``input_image``/``image``/``image_url`` parts) and ``function_call_output`` - items whose ``output`` field is an image data URL. + items whose ``output`` field is an image data URL. Tool outputs that are + non-image strings (e.g. structured JSON returned by tools like + ``click(x, y)``) contribute zero images to the bucket. """ images: list[Image.Image] = [] if item.get("type") == "function_call_output": src = item.get("output") - if isinstance(src, str): + if isinstance(src, str) and _looks_like_image_src(src): images.append(resolve_to_image(src)) return images content = item.get("content") or [] @@ -226,7 +242,10 @@ def _extract_input_images_from_message(item: dict) -> list[Image.Image]: return images -def _index_per_turn_images(output: list[dict]) -> list[list[Image.Image]]: +def _index_per_turn_images( + output: list[dict], + input_messages: list[dict] | None = None, +) -> list[list[Image.Image]]: """Bin server-returned images by the trainable turn that saw them. Walks the Responses-API items in order and flushes ``pending`` into a @@ -239,9 +258,19 @@ def _index_per_turn_images(output: list[dict]) -> list[list[Image.Image]]: has one entry per trainable turn, aligned with the postprocess loop's ``turn_idx`` even when the trainable item's role is not ``assistant`` (e.g. a reasoning-only response, or a ``function_call``). + + ``input_messages`` is the initial ``responses_create_params.input`` list — + images there (e.g. a single-shot user prompt for tool-based envs like + circle-click) are consumed by the first trainable turn's tokenized prompt + and must land in the first bucket. Agents like ``gym_v_agent`` that keep + ``input`` empty and inject observations as ``function_call_output`` items + are unaffected — the seed is a no-op when ``input_messages`` is empty. """ per_turn: list[list[Image.Image]] = [] pending: list[Image.Image] = [] + for item in input_messages or (): + if isinstance(item, dict) and item.get("role") != "assistant": + pending.extend(_extract_input_images_from_message(item)) for item in output: if item.get( "generation_token_ids" @@ -515,8 +544,15 @@ def _postprocess_nemo_gym_to_nemo_rl_result( ) processor = getattr(self, "_processor", None) - per_turn_images = _index_per_turn_images( - nemo_gym_result["response"]["output"], + per_turn_images = ( + _index_per_turn_images( + nemo_gym_result["response"]["output"], + input_messages=nemo_gym_result.get( + "responses_create_params", {} + ).get("input"), + ) + if processor is not None + else [] ) turn_idx = 0 diff --git a/tests/unit/environments/test_nemo_gym_mm_utils.py b/tests/unit/environments/test_nemo_gym_mm_utils.py index 71ae0c9c25c..3d8a4186e17 100644 --- a/tests/unit/environments/test_nemo_gym_mm_utils.py +++ b/tests/unit/environments/test_nemo_gym_mm_utils.py @@ -55,6 +55,18 @@ def test_extract_input_images_returns_empty_for_string_content(): assert _extract_input_images_from_message({"role": "user"}) == [] +def test_extract_input_images_ignores_text_function_call_output(): + item = { + "type": "function_call_output", + "call_id": "c1", + "output": '{"ok": true}', + } + assert _extract_input_images_from_message(item) == [] + + item["output"] = "Tool failed to create result.png" + assert _extract_input_images_from_message(item) == [] + + def test_index_per_turn_images_bins_images(): output = [ _user(_image((2, 2))), @@ -69,6 +81,16 @@ def test_index_per_turn_images_bins_images(): assert [img.size for img in per_turn[1]] == [(3, 3), (4, 4)] +def test_index_per_turn_images_seeds_first_turn_from_input_messages(): + input_messages = [_user(_image((2, 2)))] + output = [_assistant([1, 2])] + + per_turn = _index_per_turn_images(output, input_messages=input_messages) + + assert len(per_turn) == 1 + assert [img.size for img in per_turn[0]] == [(2, 2)] + + def test_index_per_turn_images_text_only_rollout_yields_empty_buckets(): output = [ {"role": "user", "content": "solve this"}, From 1c5bfbb8ef11350a629c9cd9578f54d4ca40db84 Mon Sep 17 00:00:00 2001 From: rohitrango Date: Wed, 5 Aug 2026 14:31:41 -0700 Subject: [PATCH 17/27] lint fixes Signed-off-by: rohitrango --- nemo_rl/algorithms/grpo.py | 1 + nemo_rl/data/processors.py | 4 ++-- nemo_rl/environments/nemo_gym.py | 12 +++++++----- nemo_rl/models/megatron/setup.py | 5 ++--- 4 files changed, 12 insertions(+), 10 deletions(-) diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index bc841990bac..29e5037b9a7 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -1982,6 +1982,7 @@ def _preserve_router_replay_routed_experts( if router_replay_enabled(policy_config) and "routed_experts" in flat_messages: target["routed_experts"] = flat_messages["routed_experts"] + def _build_async_grpo_train_data( flat_messages: BatchedDataDict, input_lengths: torch.Tensor, diff --git a/nemo_rl/data/processors.py b/nemo_rl/data/processors.py index adad65b4d5e..a6e8ef14727 100644 --- a/nemo_rl/data/processors.py +++ b/nemo_rl/data/processors.py @@ -461,10 +461,10 @@ def vlm_hf_data_processor( from nemo_rl.data.multimodal_utils import ( PackedTensor, get_dim_to_pack_along, - get_multimodal_keys_from_processor, - uses_image_placeholder, get_multimodal_default_settings_from_processor, + get_multimodal_keys_from_processor, resolve_to_image, + uses_image_placeholder, ) # depending on the task, format the data differently diff --git a/nemo_rl/environments/nemo_gym.py b/nemo_rl/environments/nemo_gym.py index 9fc7dd7e4bb..c2f475c1c96 100644 --- a/nemo_rl/environments/nemo_gym.py +++ b/nemo_rl/environments/nemo_gym.py @@ -26,7 +26,6 @@ from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy from transformers import PreTrainedTokenizerBase -from nemo_rl.distributed.ray_actor_environment_registry import get_actor_python_env from nemo_rl.data.multimodal_utils import ( PackedTensor, encode_images_in_examples, @@ -35,7 +34,7 @@ resolve_to_image, uses_image_placeholder, ) - +from nemo_rl.distributed.ray_actor_environment_registry import get_actor_python_env from nemo_rl.distributed.virtual_cluster import ( DEFAULT_GYM_PORT_RANGE_HIGH, DEFAULT_GYM_PORT_RANGE_LOW, @@ -194,6 +193,9 @@ def _detect_invalid_tool_call_and_malformed_thinking( ######################################## +# WARNING: A function-call output beginning with HTTP(S) is accepted here and +# passed to ``resolve_to_image``, which performs an outbound request during +# postprocessing even when the tool result is not actually an image. _IMAGE_SRC_PREFIXES = ("data:image/", "http://", "https://", "file://") @@ -547,9 +549,9 @@ def _postprocess_nemo_gym_to_nemo_rl_result( per_turn_images = ( _index_per_turn_images( nemo_gym_result["response"]["output"], - input_messages=nemo_gym_result.get( - "responses_create_params", {} - ).get("input"), + input_messages=nemo_gym_result.get("responses_create_params", {}).get( + "input" + ), ) if processor is not None else [] diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index 85df8fae8ad..13dc97f5b5f 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -1410,9 +1410,8 @@ def freeze_moe_router(megatron_model): # `.llava_model.language_model`; unwrap that layer first so the # generic `.language_model.decoder.layers` walk below finds the # MoE router. - if ( - getattr(model_module, "llava_model", None) is not None - and hasattr(model_module.llava_model, "language_model") + if getattr(model_module, "llava_model", None) is not None and hasattr( + model_module.llava_model, "language_model" ): model_module = model_module.llava_model if hasattr(model_module, "language_model"): From 062e999129a125e7b65c92fcd5e684846dfe76b9 Mon Sep 17 00:00:00 2001 From: rohitrango Date: Wed, 5 Aug 2026 14:44:06 -0700 Subject: [PATCH 18/27] allow mixed (multimodal, text) batches from nemo-gym batch rollouts Signed-off-by: rohitrango --- nemo_rl/data/llm_message_utils.py | 25 +++++++++++++-- tests/unit/data/test_llm_message_utils.py | 37 +++++++++++++++++++++++ 2 files changed, 59 insertions(+), 3 deletions(-) diff --git a/nemo_rl/data/llm_message_utils.py b/nemo_rl/data/llm_message_utils.py index d3e1a4df903..4c120fd670d 100644 --- a/nemo_rl/data/llm_message_utils.py +++ b/nemo_rl/data/llm_message_utils.py @@ -363,9 +363,28 @@ def batched_message_log_to_flat_message( result = BatchedDataDict() for key in all_keys: values = [seq.get(key) for seq in sequenced_lists] - # if the values are PackedTensors, create a new PackedTensor from the list of values - if values and isinstance(values[0], PackedTensor): - result[key] = PackedTensor.flattened_concat(values) + packed_template = next( + (value for value in values if isinstance(value, PackedTensor)), None + ) + if packed_template is not None: + if any( + value is not None and not isinstance(value, PackedTensor) + for value in values + ): + raise TypeError( + f"Expected PackedTensor or None for {key=}, " + f"got {[type(value).__name__ for value in values]}" + ) + filled_packed_values = cast( + list[PackedTensor], + [ + PackedTensor.empty_like(packed_template) + if value is None + else value + for value in values + ], + ) + result[key] = PackedTensor.flattened_concat(filled_packed_values) continue if not values or not isinstance(values[0], Tensor): result[key] = values diff --git a/tests/unit/data/test_llm_message_utils.py b/tests/unit/data/test_llm_message_utils.py index b39f1175934..ee85ddabf24 100644 --- a/tests/unit/data/test_llm_message_utils.py +++ b/tests/unit/data/test_llm_message_utils.py @@ -772,6 +772,43 @@ def test_batched_message_log_to_flat_message_with_packed_images() -> None: assert torch.equal(input_lengths, torch.tensor([4, 5], dtype=torch.int32)) +@pytest.mark.parametrize("image_first", [True, False]) +def test_batched_message_log_to_flat_message_with_image_free_sample( + image_first: bool, +) -> None: + from nemo_rl.data.multimodal_utils import PackedTensor + + image = torch.randn(1, 3, 4, 4) + image_log: LLMMessageLogType = [ + { + "role": "user", + "token_ids": torch.tensor([1, 2]), + "pixel_values": PackedTensor(image, dim_to_pack=0), + } + ] + image_free_log: LLMMessageLogType = [ + {"role": "user", "token_ids": torch.tensor([3, 4])} + ] + batch_logs = ( + [image_log, image_free_log] + if image_first + else [image_free_log, image_log] + ) + + batched, _ = batched_message_log_to_flat_message(batch_logs) + + pixel_values = batched["pixel_values"] + assert isinstance(pixel_values, PackedTensor) + assert len(pixel_values) == 2 + expected = [image, None] if image_first else [None, image] + for actual, expected_value in zip(pixel_values.tensors, expected): + if expected_value is None: + assert actual is None + else: + assert torch.equal(actual, expected_value) + assert "pixel_values" in batched.get_multimodal_dict() + + @pytest.mark.hf_gated def test_get_formatted_message_log_multimodal_prompt_formatting() -> None: processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-3B-Instruct") From 0edbe98f7aa387b2bac46cea6f0a7d280d1abf51 Mon Sep 17 00:00:00 2001 From: rohitrango Date: Wed, 5 Aug 2026 17:08:05 -0700 Subject: [PATCH 19/27] chore: apply ruff format to llm_message_utils Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: rohitrango --- nemo_rl/data/llm_message_utils.py | 4 +--- tests/unit/data/test_llm_message_utils.py | 4 +--- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/nemo_rl/data/llm_message_utils.py b/nemo_rl/data/llm_message_utils.py index 4c120fd670d..c3e6bf76586 100644 --- a/nemo_rl/data/llm_message_utils.py +++ b/nemo_rl/data/llm_message_utils.py @@ -378,9 +378,7 @@ def batched_message_log_to_flat_message( filled_packed_values = cast( list[PackedTensor], [ - PackedTensor.empty_like(packed_template) - if value is None - else value + PackedTensor.empty_like(packed_template) if value is None else value for value in values ], ) diff --git a/tests/unit/data/test_llm_message_utils.py b/tests/unit/data/test_llm_message_utils.py index ee85ddabf24..113fd9ce0b9 100644 --- a/tests/unit/data/test_llm_message_utils.py +++ b/tests/unit/data/test_llm_message_utils.py @@ -790,9 +790,7 @@ def test_batched_message_log_to_flat_message_with_image_free_sample( {"role": "user", "token_ids": torch.tensor([3, 4])} ] batch_logs = ( - [image_log, image_free_log] - if image_first - else [image_free_log, image_log] + [image_log, image_free_log] if image_first else [image_free_log, image_log] ) batched, _ = batched_message_log_to_flat_message(batch_logs) From a70eacc9340115ffb9876832cf1330b8663badb9 Mon Sep 17 00:00:00 2001 From: rohitrango Date: Wed, 5 Aug 2026 17:54:18 -0700 Subject: [PATCH 20/27] feat(recipes): add Nemotron-Omni 30B circle-click 2n8g VLM-GRPO recipe Non-colocated 2n8g layout (vLLM TP=8 on node 1, Megatron TP=2/EP=8/CP=2 on node 2) for the single-turn Circle-Click NeMo-Gym environment. Model points at the HF repo (nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16); train/eval data_path use /path/to/{train,eval}_dataset.jsonl placeholders so users wire in their own manifests. Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: rohitrango --- ...-circle-click-2n8g-megatron-tp2ep8.v1.yaml | 441 ++++++++++++++++++ 1 file changed, 441 insertions(+) create mode 100644 examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-click-2n8g-megatron-tp2ep8.v1.yaml diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-click-2n8g-megatron-tp2ep8.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-click-2n8g-megatron-tp2ep8.v1.yaml new file mode 100644 index 00000000000..eb99582fe68 --- /dev/null +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-click-2n8g-megatron-tp2ep8.v1.yaml @@ -0,0 +1,441 @@ +# VLM-GRPO recipe: Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16 on the +# Circle-Click NeMo-Gym environment. +# +# Circle-Click is a single-turn function-call environment: the agent is shown +# an image and must issue exactly one `click(x, y)` tool call that lands +# inside the target-colored circle. Verification (hit / miss + reward) +# happens in the resources server. +# +# Shape (2 nodes x 8 GPUs, non-colocated): +# Node 1: vLLM generation (TP=8) +# Node 2: Megatron policy training (TP=2, EP=8, PP=1, CP=2) +# +# Agent harness: `simple_agent` under the `circle_click_simple_agent` group +# (defined in +# 3rdparty/Gym-workspace/Gym/resources_servers/circle_click/configs/circle_click.yaml). +# max_steps=1 enforces the single-tool-call rollout the resources server is +# designed to verify. +# +# Training / validation manifests must contain rows with +# `agent_ref.name: circle_click_simple_agent` (asserted in +# nemo_rl/environments/nemo_gym.py) and Responses-API-formatted input + +# tools bound to the `click` function. The stock example bundled with the +# resources server (resources_servers/circle_click/data/example.jsonl) +# carries the message schema and tool schema but omits `agent_ref`; add it +# before wiring the row into a training manifest. +# +# Run with: +# HF_HOME=${PWD}/.cache/huggingface \ +# TRANSFORMERS_CACHE=${PWD}/.cache/huggingface/hub \ +# RAY_TMPDIR=/tmp/ray \ +# uv run examples/nemo_gym/run_grpo_nemo_gym.py \ +# --config examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-click-2n8g-megatron-tp2ep8.v1.yaml +# +# Env vars (must be set — same set as `experiments/run_gymv_smoke.sh`): +# HF_HOME: persistent HF cache under the mounted worktree. Nemotron-Omni +# uses `trust_remote_code=True`; the driver writes its dynamic +# `modeling_*.py` under `$HF_HOME/modules/transformers_modules//` +# and adds that dir to `sys.path`. Ray workers receive pickled objects +# referencing `transformers_modules..NanoOmni…` and must resolve +# them against the SAME `$HF_HOME`. Without this, workers die at +# actor-init with `ModuleNotFoundError: No module named +# 'transformers_modules'`. +# TRANSFORMERS_CACHE: paired with HF_HOME; controls the model-weights +# download location. +# RAY_TMPDIR: forced under /tmp because Ray's AF_UNIX socket path is +# capped at 107 bytes on Linux; Lustre-rooted `$PWD/tmp` overruns it. +# +# Do not pass `--extra mcore --extra vllm` — pyproject.toml declares them as +# mutually exclusive conflicts. Backend worker venvs are synced with their +# own extras at runtime by NRL's venv manager. + +grpo: + num_prompts_per_step: 1 # smoke: minimal loop + num_generations_per_prompt: 16 # gbs = 1 * 16 = 16 + num_val_generations_per_prompt: 1 + # Circle-Click is a single tool-call environment; horizon_cap is 1 and the + # simple_agent breaks after one function_call → verify. + max_rollout_turns: 1 + max_num_epochs: 1 + max_num_steps: 100 + normalize_rewards: true + use_leave_one_out_baseline: true + val_period: 500 + val_at_start: false + val_at_end: false + overlong_filtering: true + advantage_clip_low: null + advantage_clip_high: null + max_val_samples: null + val_batch_size: null + seed: 42 + async_grpo: + enabled: false # sync + max_trajectory_age_steps: 1 + + batch_multiplier: 1 + use_dynamic_sampling: False + reward_shaping: + enabled: False + reward_scaling: + enabled: False + + seq_logprob_error_threshold: 2 + invalid_tool_call_advantage: null + malformed_thinking_advantage: null + +loss_fn: + reference_policy_kl_penalty: 0 + reference_policy_kl_type: k3 + kl_input_clamp_value: null + kl_output_clamp_value: null + ratio_clip_min: 0.2 + ratio_clip_max: 0.28 + ratio_clip_c: null + use_on_policy_kl_approximation: True + use_importance_sampling_correction: True + sequence_level_importance_ratios: False + token_level_loss: True + truncated_importance_sampling_ratio: null + +checkpointing: + enabled: false + checkpoint_dir: "results/grpo-nemotron-omni-30ba3b-gymv-circle-click" + metric_name: "val:total_reward/mean" + higher_is_better: true + keep_top_k: 1000000 + save_period: 10 + checkpoint_must_save_by: "00:03:40:00" + save_optimizer: true + +policy: + model_name: "nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16" + is_vlm: true # required by the multimodal grpo entrypoint + tokenizer: + name: ${policy.model_name} + # Nano-Omni-Reasoning uses a native block; keep it open and + # preserve prior reasoning across turns (single-turn here, but harmless). + chat_template_kwargs: + enable_thinking: true + truncate_history_thinking: false + train_global_batch_size: ${mul:${grpo.num_prompts_per_step}, ${grpo.num_generations_per_prompt}} + train_micro_batch_size: 1 + generation_batch_size: 64 # HF backend only + logprob_batch_size: 1 + max_total_sequence_length: 8192 + precision: "bfloat16" + logprob_chunk_size: 2048 + + dtensor_cfg: + _v2: true + enabled: false + cpu_offload: False + sequence_parallel: false + activation_checkpointing: false + tensor_parallel_size: 1 + context_parallel_size: 1 + custom_parallel_plan: null + + megatron_cfg: + enabled: true + checkpoint: + async_save: true + empty_unused_memory_level: 1 + activation_checkpointing: true + bias_activation_fusion: False + # 8-GPU policy node: TP=2, EP=8, PP=1, CP=2. + tensor_model_parallel_size: 2 + expert_tensor_parallel_size: 1 + expert_model_parallel_size: 8 + pipeline_model_parallel_size: 1 + num_layers_in_first_pipeline_stage: null + num_layers_in_last_pipeline_stage: null + context_parallel_size: 2 + pipeline_dtype: ${policy.precision} + sequence_parallel: true # only pays off with tp + cp + freeze_moe_router: true + moe_router_dtype: "fp32" + moe_router_load_balancing_type: "none" + moe_router_bias_update_rate: 1e-3 + moe_permute_fusion: true + moe_enable_deepep: false + moe_token_dispatcher_type: "alltoall" + moe_aux_loss_coeff: 0.0 + moe_router_enable_expert_bias: true + apply_rope_fusion: True + defer_fp32_logits: True + track_moe_metrics: True + moe_per_layer_logging: True + moe_shared_expert_overlap: false + gradient_accumulation_fusion: false + use_fused_weighted_squared_relu: false + # RADIO CPE eval mode: keeps vision-tower positional embeddings in eval + # mode during rollout/train (required for the frozen vision path to + # produce stable features). + radio_force_cpe_eval_mode: true + # Empty the CUDA cache before the vLLM refit broadcast so the packed + # staging tensor doesn't OOM against Adam m/v that materialize at the + # end of iter 1. + clear_memory_caches_before_refit: true + + optimizer: + optimizer: "adam" + lr: 3e-6 + min_lr: 3e-6 + weight_decay: 0.0 + bf16: true + fp16: false + params_dtype: "float32" + adam_beta1: 0.9 + adam_beta2: 0.999 + adam_eps: 1e-8 + sgd_momentum: 0.9 + clip_grad: ${policy.max_grad_norm} + use_distributed_optimizer: true + use_precision_aware_optimizer: true + # Offload Adam m/v to CPU during logprob + refit. The 30B policy on + # 8 GPUs with TP=2/EP=8/CP=2 pushes the Megatron worker to ~68 GiB of + # live PyTorch tensors even with activation_checkpointing on; adding + # the packed staging tensor for broadcast_weights_for_collective at + # step 2 OOMs. Offloading Adam frees ~30 GiB per rank at some perf + # cost, which is the right tradeoff for a 2n8g smoke. + optimizer_cpu_offload: true + optimizer_offload_fraction: 1.0 + + scheduler: + start_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} + end_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} + weight_decay_incr_style: "constant" + lr_decay_style: "constant" + lr_decay_iters: null + # Flat LR throughout: skip warmup. Combined with lr_decay_style="constant" + # and min_lr == lr above, this pins the LR at 3e-6 from step 0 onward. + lr_warmup_iters: 0 + lr_warmup_init: 3.0e-7 + + distributed_data_parallel_config: + grad_reduce_in_fp32: false + # Nemotron-Omni carries a sound_encoder / sound_projection tower that + # never sees an audio input in gym-v. Those params live in the DDP + # bucket with requires_grad=True but never receive a backward hook, so + # MCore's async grad-reduce path trips the golden-count assertion at + # param_and_grad_buffer.py:272 on iter 2's zero_grad_buffer(). Disable + # the overlap paths until the sound tower is actually frozen via a + # nemo-rl pre_wrap_hook (mirror of freeze_moe_router). This is a real + # perf cost — remove once the freeze path is wired. + overlap_grad_reduce: false + overlap_param_gather: false + average_in_collective: false + use_custom_fsdp: false + data_parallel_sharding_strategy: "optim_grads_params" + + env_vars: null + + draft: + enabled: false + model_name: null + loss_weight: 0.1 + num_layers: null + aux_layer_indices: null + + dynamic_batching: + enabled: False + train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}} + sequence_length_round: 64 + + sequence_packing: + enabled: True + train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}} + algorithm: "modified_first_fit_decreasing" + sequence_length_round: 64 + + # Base recipe ties this to dtensor_cfg.tensor_parallel_size; we're on + # megatron, so pin to a value that divides the megatron TP layout. + make_sequence_length_divisible_by: 32 + max_grad_norm: 1.0 + + optimizer: null # remove default FSDP optimizer + scheduler: null + + offload_optimizer_for_logprob: False + + generation: + port_range_low: 3000 + port_range_high: 4999 + backend: "vllm" + max_new_tokens: ${policy.max_total_sequence_length} # single-turn; whole budget available + temperature: 1.0 + top_p: 1.0 + top_k: null + stop_token_ids: null + stop_strings: null + # Nano-Omni-Reasoning bad_words: only the vision-side text tokens are + # safe (in-vocab). The audio-delimiter tokens (, + # , ) sit past the LM's text-logit width and cause + # vLLM out-of-bounds writes in bad_words masking. Even the vision tags + # currently diverge sampling vs. the reference adlr-super-v3-omni-nemorl + # branch — leave the list empty for parity. + bad_words: [] + mcore_generation_config: + max_model_len: ${policy.max_total_sequence_length} + transformer_impl: "inference_optimized" + cuda_graph_impl: "local" + inference_cuda_graph_scope: "block" + activation_checkpointing: false + mamba_inference_ssm_states_dtype: "float32" + inference_moe_token_dispatcher_type: "nccl" + inference_grouped_gemm_backend: "vllm" + moe_router_num_groups: null + moe_router_group_topk: null + pipeline_model_parallel_size: 1 + expert_tensor_parallel_size: 1 + expert_model_parallel_size: 8 + sequence_parallel: True + context_parallel_size: 1 + tensor_model_parallel_size: 2 + buffer_size_gb: 20 + num_cuda_graphs: -1 + block_size_tokens: 256 + use_cuda_graphs_for_non_decode_steps: true + enable_chunked_prefill: true + max_tokens: ${policy.max_total_sequence_length} + kv_cache_management_mode: "persist" + materialize_only_last_token_logits: true + num_speculative_tokens: 0 + refit_backend: "nvshmem" + async_engine: true + expose_http_server: true + enable_prefix_caching: true + parsers: + - deepseek-r1-reasoning + - qwen3-coder-tool + vllm_cfg: + async_engine: true + kv_cache_dtype: auto + precision: ${policy.precision} + # Full 8-GPU vLLM node. + tensor_parallel_size: 8 + pipeline_parallel_size: 1 + expert_parallel_size: 1 + # 0.6 leaves ~30 GiB headroom for the weight-refit broadcast from + # Megatron. 0.8 (the smoke default) OOMs at step-2's + # broadcast_weights_for_collective on this recipe because the packed + # staging tensor + policy-side Adam m/v both materialize just before + # the vLLM copy, and clear_memory_caches_before_refit only reclaims + # cached fragments, not live tensors. + gpu_memory_utilization: 0.6 + max_model_len: ${policy.max_total_sequence_length} + enforce_eager: true # skip CUDA-graph capture + use_deep_gemm: False + num_last_layers_in_bf16: 0 + num_first_layers_in_bf16: 0 + # Nano-Omni + prefix caching crashes vLLM's mm cache (mm_hash miss). + # Belt-and-suspenders: also disabled via mm_processor_cache_gb=0 below. + enable_prefix_caching: false + # VLMs need the tokenizer initialized (multimodal grpo entrypoint asserts). + skip_tokenizer_init: false + expose_http_server: true + reasoning_parser_plugin: nemo_rl/models/generation/vllm/reasoning_parsers/nano_v3_reasoning_parser.py + # Nano-Omni chat template expects string content, not the OpenAI + # {type: text, text: ...} list. Keeps the NeMo-Gym prompt-token prefix + # invariant across multi-turn rollouts (single-turn here, still needed + # for cache correctness). + http_server_serving_chat_kwargs: + enable_auto_tools: true + tool_parser: qwen3_coder + reasoning_parser: nano_v3 + chat_template_content_format: string + default_chat_template_kwargs: + enable_thinking: true + truncate_history_thinking: false + + vllm_kwargs: + # Circle-Click is single-turn: one image in the prompt, no + # per-turn observation echoes. + limit_mm_per_prompt: {"image": 1} + # Disable vLLM's mm processor cache — the mm-cache desync crash on + # Nano-Omni is the load-bearing fix. Must sit in vllm_kwargs, not + # vllm_cfg. + mm_processor_cache_gb: 0 + max_num_batched_tokens: 16384 + # Nano-Omni's mamba backbone needs SSM cache in fp32 (accuracy). + mamba_ssm_cache_dtype: "float32" + compilation_config: + backend: eager + + colocated: + # Non-colocated: one dedicated vLLM node, the other node runs Megatron + # training. Total cluster nodes = 2 (see `cluster.num_nodes` below); + # generation takes `resources.num_nodes` of those, training gets the + # remainder. See nemo_rl/algorithms/grpo.py:setup for the split. + enabled: false + resources: + gpus_per_node: 8 + num_nodes: 1 + +data: + max_input_seq_length: null + shuffle: False + num_workers: 1 + use_multiple_dataloader: false + + train: + data_path: /path/to/train_dataset.jsonl + validation: + data_path: /path/to/eval_dataset.jsonl + default: + dataset_name: NemoGymDataset + env_name: "nemo_gym" + prompt_file: null + system_prompt_file: null + processor: "nemo_gym_data_processor" + +env: + should_use_nemo_gym: true + # true: skip expensive train_data_step*.jsonl; false: write full jsonl. + should_log_nemo_gym_responses: true + nemo_gym: + is_trajectory_collection: false + port_range_low: 5000 + port_range_high: 5999 + # Load the Circle-Click resources server config. The + # vllm_model_for_training entry is required for the policy-model side of + # every gym-v run. + config_paths: + - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml + - resources_servers/circle_click/configs/circle_click.yaml + # Agent-side horizon. Outer group name matches the training manifest's + # `agent_ref.name`; inner `simple_agent` key matches the entry in + # circle_click.yaml. Keep max_steps == grpo.max_rollout_turns. + circle_click_simple_agent: + responses_api_agents: + simple_agent: + max_steps: 1 + +logger: + log_dir: "logs/grpo-nemotron-omni-30ba3b-gymv-circle-click" + num_val_samples_to_print: 0 + wandb_enabled: true + tensorboard_enabled: true + mlflow_enabled: false + monitor_gpus: false + swanlab_enabled: false + wandb: + project: "grpo-nemotron-omni-gymv" + name: "grpo-nemotron-omni-30ba3b-gymv-circle-click" + log_nemo_gym_full_result_tables: false + tensorboard: {} + mlflow: + experiment_name: "grpo-nemotron-omni-gymv" + run_name: "grpo-nemotron-omni-30ba3b-gymv-circle-click" + gpu_monitoring: + collection_interval: 10 + flush_interval: 10 + +cluster: + gpus_per_node: 8 + num_nodes: 2 + master_port_range_low: 1400 + master_port_range_high: 1999 From 68a538b844dd7f30389a018236ab8ec472e583e4 Mon Sep 17 00:00:00 2001 From: rohitrango Date: Wed, 5 Aug 2026 18:06:27 -0700 Subject: [PATCH 21/27] feat(nemo_gym): assert placeholder-style processor at actor init The multimodal postprocessing path in _attach_multimodal_data_to_user_message assumes a placeholder-style processor: it reconstructs imgs_sizes / num_frames and builds the pixel_values PackedTensor with pad_to_max_shape=True. A non-placeholder VLM (e.g. Qwen2-VL / LLaVA-style) would silently produce wrong multimodal tensors instead of erroring. Fail loud at NemoGym.__init__ so the misconfiguration is caught at actor construction, well before any rollout. Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: rohitrango --- nemo_rl/environments/nemo_gym.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/nemo_rl/environments/nemo_gym.py b/nemo_rl/environments/nemo_gym.py index c2f475c1c96..2c9e0e77310 100644 --- a/nemo_rl/environments/nemo_gym.py +++ b/nemo_rl/environments/nemo_gym.py @@ -362,6 +362,16 @@ def __init__(self, cfg: NemoGymConfig): from nemo_rl.algorithms.utils import get_tokenizer self._processor = get_tokenizer(tokenizer_config, get_processor=True) + # _attach_multimodal_data_to_user_message assumes a placeholder-style + # processor (imgs_sizes / num_frames reconstruction + pad_to_max_shape + # PackedTensor build). A non-placeholder VLM would silently produce + # wrong multimodal tensors — fail at actor construction instead. + assert uses_image_placeholder(self._processor), ( + "NemoGym multimodal path assumes a placeholder-style processor " + "(see _PLACEHOLDER_STYLE_PROCESSOR_NAMES in nemo_rl/data/multimodal_utils.py); " + f"got {type(self._processor).__name__}. Update " + "_attach_multimodal_data_to_user_message before enabling." + ) def _spinup(self) -> None: """Start the NeMo-Gym head server and rollout collection helper. From b9e1a27e67e4c4e75adc19361ea6865b2ab88b05 Mon Sep 17 00:00:00 2001 From: rohitrango Date: Wed, 5 Aug 2026 18:12:53 -0700 Subject: [PATCH 22/27] chore(recipes): inherit from vlm_grpo_3B_megatron exemplar in circle-click recipe Adds the required defaults: ../../vlm_grpo_3B_megatron.yaml key so the recipe passes the configs-minimize-check pre-commit hook (all recipes under examples/configs/recipes/**/*.yaml must inherit from an exemplar). Drops the standalone header preamble; provenance now lives in the adjacent clevr/mmpr sibling recipes and the commit history. Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: rohitrango --- ...-circle-click-2n8g-megatron-tp2ep8.v1.yaml | 51 +------------------ 1 file changed, 1 insertion(+), 50 deletions(-) diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-click-2n8g-megatron-tp2ep8.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-click-2n8g-megatron-tp2ep8.v1.yaml index eb99582fe68..8b838e8d449 100644 --- a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-click-2n8g-megatron-tp2ep8.v1.yaml +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-click-2n8g-megatron-tp2ep8.v1.yaml @@ -1,53 +1,4 @@ -# VLM-GRPO recipe: Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16 on the -# Circle-Click NeMo-Gym environment. -# -# Circle-Click is a single-turn function-call environment: the agent is shown -# an image and must issue exactly one `click(x, y)` tool call that lands -# inside the target-colored circle. Verification (hit / miss + reward) -# happens in the resources server. -# -# Shape (2 nodes x 8 GPUs, non-colocated): -# Node 1: vLLM generation (TP=8) -# Node 2: Megatron policy training (TP=2, EP=8, PP=1, CP=2) -# -# Agent harness: `simple_agent` under the `circle_click_simple_agent` group -# (defined in -# 3rdparty/Gym-workspace/Gym/resources_servers/circle_click/configs/circle_click.yaml). -# max_steps=1 enforces the single-tool-call rollout the resources server is -# designed to verify. -# -# Training / validation manifests must contain rows with -# `agent_ref.name: circle_click_simple_agent` (asserted in -# nemo_rl/environments/nemo_gym.py) and Responses-API-formatted input + -# tools bound to the `click` function. The stock example bundled with the -# resources server (resources_servers/circle_click/data/example.jsonl) -# carries the message schema and tool schema but omits `agent_ref`; add it -# before wiring the row into a training manifest. -# -# Run with: -# HF_HOME=${PWD}/.cache/huggingface \ -# TRANSFORMERS_CACHE=${PWD}/.cache/huggingface/hub \ -# RAY_TMPDIR=/tmp/ray \ -# uv run examples/nemo_gym/run_grpo_nemo_gym.py \ -# --config examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-click-2n8g-megatron-tp2ep8.v1.yaml -# -# Env vars (must be set — same set as `experiments/run_gymv_smoke.sh`): -# HF_HOME: persistent HF cache under the mounted worktree. Nemotron-Omni -# uses `trust_remote_code=True`; the driver writes its dynamic -# `modeling_*.py` under `$HF_HOME/modules/transformers_modules//` -# and adds that dir to `sys.path`. Ray workers receive pickled objects -# referencing `transformers_modules..NanoOmni…` and must resolve -# them against the SAME `$HF_HOME`. Without this, workers die at -# actor-init with `ModuleNotFoundError: No module named -# 'transformers_modules'`. -# TRANSFORMERS_CACHE: paired with HF_HOME; controls the model-weights -# download location. -# RAY_TMPDIR: forced under /tmp because Ray's AF_UNIX socket path is -# capped at 107 bytes on Linux; Lustre-rooted `$PWD/tmp` overruns it. -# -# Do not pass `--extra mcore --extra vllm` — pyproject.toml declares them as -# mutually exclusive conflicts. Backend worker venvs are synced with their -# own extras at runtime by NRL's venv manager. +defaults: ../../vlm_grpo_3B_megatron.yaml grpo: num_prompts_per_step: 1 # smoke: minimal loop From 4bb7e20df7d94e2e8b588681fdf825b861a1b04b Mon Sep 17 00:00:00 2001 From: rohitrango Date: Wed, 5 Aug 2026 19:45:09 -0700 Subject: [PATCH 23/27] minimized config Signed-off-by: rohitrango --- ...-circle-click-2n8g-megatron-tp2ep8.v1.yaml | 318 +++--------------- 1 file changed, 43 insertions(+), 275 deletions(-) diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-click-2n8g-megatron-tp2ep8.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-click-2n8g-megatron-tp2ep8.v1.yaml index 8b838e8d449..3695064845c 100644 --- a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-click-2n8g-megatron-tp2ep8.v1.yaml +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-click-2n8g-megatron-tp2ep8.v1.yaml @@ -1,298 +1,105 @@ defaults: ../../vlm_grpo_3B_megatron.yaml - grpo: - num_prompts_per_step: 1 # smoke: minimal loop - num_generations_per_prompt: 16 # gbs = 1 * 16 = 16 + num_prompts_per_step: 1 num_val_generations_per_prompt: 1 - # Circle-Click is a single tool-call environment; horizon_cap is 1 and the - # simple_agent breaks after one function_call → verify. - max_rollout_turns: 1 - max_num_epochs: 1 max_num_steps: 100 - normalize_rewards: true - use_leave_one_out_baseline: true val_period: 500 - val_at_start: false - val_at_end: false overlong_filtering: true - advantage_clip_low: null - advantage_clip_high: null max_val_samples: null val_batch_size: null - seed: 42 - async_grpo: - enabled: false # sync - max_trajectory_age_steps: 1 - - batch_multiplier: 1 - use_dynamic_sampling: False - reward_shaping: - enabled: False - reward_scaling: - enabled: False - seq_logprob_error_threshold: 2 - invalid_tool_call_advantage: null - malformed_thinking_advantage: null - loss_fn: reference_policy_kl_penalty: 0 - reference_policy_kl_type: k3 kl_input_clamp_value: null kl_output_clamp_value: null - ratio_clip_min: 0.2 ratio_clip_max: 0.28 - ratio_clip_c: null - use_on_policy_kl_approximation: True - use_importance_sampling_correction: True - sequence_level_importance_ratios: False - token_level_loss: True - truncated_importance_sampling_ratio: null - + use_on_policy_kl_approximation: true + use_importance_sampling_correction: true checkpointing: enabled: false - checkpoint_dir: "results/grpo-nemotron-omni-30ba3b-gymv-circle-click" - metric_name: "val:total_reward/mean" - higher_is_better: true + checkpoint_dir: results/grpo-nemotron-omni-30ba3b-gymv-circle-click + metric_name: val:total_reward/mean keep_top_k: 1000000 - save_period: 10 - checkpoint_must_save_by: "00:03:40:00" - save_optimizer: true - + checkpoint_must_save_by: 00:03:40:00 policy: - model_name: "nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16" - is_vlm: true # required by the multimodal grpo entrypoint + model_name: nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16 + is_vlm: true tokenizer: - name: ${policy.model_name} - # Nano-Omni-Reasoning uses a native block; keep it open and - # preserve prior reasoning across turns (single-turn here, but harmless). chat_template_kwargs: enable_thinking: true truncate_history_thinking: false train_global_batch_size: ${mul:${grpo.num_prompts_per_step}, ${grpo.num_generations_per_prompt}} - train_micro_batch_size: 1 - generation_batch_size: 64 # HF backend only + generation_batch_size: 64 logprob_batch_size: 1 max_total_sequence_length: 8192 - precision: "bfloat16" logprob_chunk_size: 2048 - - dtensor_cfg: - _v2: true - enabled: false - cpu_offload: False - sequence_parallel: false - activation_checkpointing: false - tensor_parallel_size: 1 - context_parallel_size: 1 - custom_parallel_plan: null - megatron_cfg: - enabled: true - checkpoint: - async_save: true - empty_unused_memory_level: 1 activation_checkpointing: true - bias_activation_fusion: False - # 8-GPU policy node: TP=2, EP=8, PP=1, CP=2. + bias_activation_fusion: false tensor_model_parallel_size: 2 - expert_tensor_parallel_size: 1 expert_model_parallel_size: 8 - pipeline_model_parallel_size: 1 - num_layers_in_first_pipeline_stage: null - num_layers_in_last_pipeline_stage: null context_parallel_size: 2 - pipeline_dtype: ${policy.precision} - sequence_parallel: true # only pays off with tp + cp - freeze_moe_router: true - moe_router_dtype: "fp32" - moe_router_load_balancing_type: "none" - moe_router_bias_update_rate: 1e-3 - moe_permute_fusion: true - moe_enable_deepep: false - moe_token_dispatcher_type: "alltoall" + sequence_parallel: true + moe_router_dtype: fp32 + moe_router_bias_update_rate: 0.001 moe_aux_loss_coeff: 0.0 moe_router_enable_expert_bias: true - apply_rope_fusion: True - defer_fp32_logits: True - track_moe_metrics: True - moe_per_layer_logging: True - moe_shared_expert_overlap: false - gradient_accumulation_fusion: false - use_fused_weighted_squared_relu: false - # RADIO CPE eval mode: keeps vision-tower positional embeddings in eval - # mode during rollout/train (required for the frozen vision path to - # produce stable features). + defer_fp32_logits: true + track_moe_metrics: true + moe_per_layer_logging: true radio_force_cpe_eval_mode: true - # Empty the CUDA cache before the vLLM refit broadcast so the packed - # staging tensor doesn't OOM against Adam m/v that materialize at the - # end of iter 1. clear_memory_caches_before_refit: true - optimizer: - optimizer: "adam" - lr: 3e-6 - min_lr: 3e-6 + lr: 3.0e-06 + min_lr: 3.0e-06 weight_decay: 0.0 - bf16: true - fp16: false - params_dtype: "float32" - adam_beta1: 0.9 - adam_beta2: 0.999 - adam_eps: 1e-8 - sgd_momentum: 0.9 - clip_grad: ${policy.max_grad_norm} - use_distributed_optimizer: true - use_precision_aware_optimizer: true - # Offload Adam m/v to CPU during logprob + refit. The 30B policy on - # 8 GPUs with TP=2/EP=8/CP=2 pushes the Megatron worker to ~68 GiB of - # live PyTorch tensors even with activation_checkpointing on; adding - # the packed staging tensor for broadcast_weights_for_collective at - # step 2 OOMs. Offloading Adam frees ~30 GiB per rank at some perf - # cost, which is the right tradeoff for a 2n8g smoke. optimizer_cpu_offload: true optimizer_offload_fraction: 1.0 - scheduler: - start_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} - end_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay} - weight_decay_incr_style: "constant" - lr_decay_style: "constant" lr_decay_iters: null - # Flat LR throughout: skip warmup. Combined with lr_decay_style="constant" - # and min_lr == lr above, this pins the LR at 3e-6 from step 0 onward. lr_warmup_iters: 0 - lr_warmup_init: 3.0e-7 - + lr_warmup_init: 3.0e-07 distributed_data_parallel_config: - grad_reduce_in_fp32: false - # Nemotron-Omni carries a sound_encoder / sound_projection tower that - # never sees an audio input in gym-v. Those params live in the DDP - # bucket with requires_grad=True but never receive a backward hook, so - # MCore's async grad-reduce path trips the golden-count assertion at - # param_and_grad_buffer.py:272 on iter 2's zero_grad_buffer(). Disable - # the overlap paths until the sound tower is actually frozen via a - # nemo-rl pre_wrap_hook (mirror of freeze_moe_router). This is a real - # perf cost — remove once the freeze path is wired. - overlap_grad_reduce: false overlap_param_gather: false average_in_collective: false - use_custom_fsdp: false - data_parallel_sharding_strategy: "optim_grads_params" - - env_vars: null - - draft: - enabled: false - model_name: null - loss_weight: 0.1 - num_layers: null - aux_layer_indices: null - - dynamic_batching: - enabled: False - train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}} - logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}} - sequence_length_round: 64 - sequence_packing: - enabled: True - train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}} - logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}} - algorithm: "modified_first_fit_decreasing" - sequence_length_round: 64 - - # Base recipe ties this to dtensor_cfg.tensor_parallel_size; we're on - # megatron, so pin to a value that divides the megatron TP layout. + enabled: true make_sequence_length_divisible_by: 32 - max_grad_norm: 1.0 - - optimizer: null # remove default FSDP optimizer + optimizer: null scheduler: null - - offload_optimizer_for_logprob: False - generation: - port_range_low: 3000 - port_range_high: 4999 - backend: "vllm" - max_new_tokens: ${policy.max_total_sequence_length} # single-turn; whole budget available - temperature: 1.0 - top_p: 1.0 - top_k: null - stop_token_ids: null - stop_strings: null - # Nano-Omni-Reasoning bad_words: only the vision-side text tokens are - # safe (in-vocab). The audio-delimiter tokens (, - # , ) sit past the LM's text-logit width and cause - # vLLM out-of-bounds writes in bad_words masking. Even the vision tags - # currently diverge sampling vs. the reference adlr-super-v3-omni-nemorl - # branch — leave the list empty for parity. + max_new_tokens: ${policy.max_total_sequence_length} bad_words: [] mcore_generation_config: - max_model_len: ${policy.max_total_sequence_length} - transformer_impl: "inference_optimized" - cuda_graph_impl: "local" - inference_cuda_graph_scope: "block" + transformer_impl: inference_optimized activation_checkpointing: false - mamba_inference_ssm_states_dtype: "float32" - inference_moe_token_dispatcher_type: "nccl" - inference_grouped_gemm_backend: "vllm" + mamba_inference_ssm_states_dtype: float32 + inference_moe_token_dispatcher_type: nccl + inference_grouped_gemm_backend: vllm moe_router_num_groups: null moe_router_group_topk: null pipeline_model_parallel_size: 1 expert_tensor_parallel_size: 1 expert_model_parallel_size: 8 - sequence_parallel: True + sequence_parallel: true context_parallel_size: 1 tensor_model_parallel_size: 2 buffer_size_gb: 20 num_cuda_graphs: -1 - block_size_tokens: 256 - use_cuda_graphs_for_non_decode_steps: true - enable_chunked_prefill: true max_tokens: ${policy.max_total_sequence_length} - kv_cache_management_mode: "persist" - materialize_only_last_token_logits: true - num_speculative_tokens: 0 - refit_backend: "nvshmem" async_engine: true expose_http_server: true enable_prefix_caching: true parsers: - - deepseek-r1-reasoning - - qwen3-coder-tool + - deepseek-r1-reasoning + - qwen3-coder-tool vllm_cfg: async_engine: true - kv_cache_dtype: auto - precision: ${policy.precision} - # Full 8-GPU vLLM node. tensor_parallel_size: 8 - pipeline_parallel_size: 1 - expert_parallel_size: 1 - # 0.6 leaves ~30 GiB headroom for the weight-refit broadcast from - # Megatron. 0.8 (the smoke default) OOMs at step-2's - # broadcast_weights_for_collective on this recipe because the packed - # staging tensor + policy-side Adam m/v both materialize just before - # the vLLM copy, and clear_memory_caches_before_refit only reclaims - # cached fragments, not live tensors. - gpu_memory_utilization: 0.6 - max_model_len: ${policy.max_total_sequence_length} - enforce_eager: true # skip CUDA-graph capture - use_deep_gemm: False - num_last_layers_in_bf16: 0 - num_first_layers_in_bf16: 0 - # Nano-Omni + prefix caching crashes vLLM's mm cache (mm_hash miss). - # Belt-and-suspenders: also disabled via mm_processor_cache_gb=0 below. + enforce_eager: true enable_prefix_caching: false - # VLMs need the tokenizer initialized (multimodal grpo entrypoint asserts). - skip_tokenizer_init: false expose_http_server: true reasoning_parser_plugin: nemo_rl/models/generation/vllm/reasoning_parsers/nano_v3_reasoning_parser.py - # Nano-Omni chat template expects string content, not the OpenAI - # {type: text, text: ...} list. Keeps the NeMo-Gym prompt-token prefix - # invariant across multi-turn rollouts (single-turn here, still needed - # for cache correctness). http_server_serving_chat_kwargs: enable_auto_tools: true tool_parser: qwen3_coder @@ -301,92 +108,53 @@ policy: default_chat_template_kwargs: enable_thinking: true truncate_history_thinking: false - vllm_kwargs: - # Circle-Click is single-turn: one image in the prompt, no - # per-turn observation echoes. - limit_mm_per_prompt: {"image": 1} - # Disable vLLM's mm processor cache — the mm-cache desync crash on - # Nano-Omni is the load-bearing fix. Must sit in vllm_kwargs, not - # vllm_cfg. - mm_processor_cache_gb: 0 + limit_mm_per_prompt: + image: 1 max_num_batched_tokens: 16384 - # Nano-Omni's mamba backbone needs SSM cache in fp32 (accuracy). - mamba_ssm_cache_dtype: "float32" + mamba_ssm_cache_dtype: float32 compilation_config: backend: eager - colocated: - # Non-colocated: one dedicated vLLM node, the other node runs Megatron - # training. Total cluster nodes = 2 (see `cluster.num_nodes` below); - # generation takes `resources.num_nodes` of those, training gets the - # remainder. See nemo_rl/algorithms/grpo.py:setup for the split. enabled: false resources: gpus_per_node: 8 num_nodes: 1 - data: max_input_seq_length: null - shuffle: False - num_workers: 1 - use_multiple_dataloader: false - + shuffle: false train: data_path: /path/to/train_dataset.jsonl validation: data_path: /path/to/eval_dataset.jsonl default: dataset_name: NemoGymDataset - env_name: "nemo_gym" + env_name: nemo_gym prompt_file: null - system_prompt_file: null - processor: "nemo_gym_data_processor" - + processor: nemo_gym_data_processor env: should_use_nemo_gym: true - # true: skip expensive train_data_step*.jsonl; false: write full jsonl. should_log_nemo_gym_responses: true nemo_gym: is_trajectory_collection: false port_range_low: 5000 port_range_high: 5999 - # Load the Circle-Click resources server config. The - # vllm_model_for_training entry is required for the policy-model side of - # every gym-v run. config_paths: - - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml - - resources_servers/circle_click/configs/circle_click.yaml - # Agent-side horizon. Outer group name matches the training manifest's - # `agent_ref.name`; inner `simple_agent` key matches the entry in - # circle_click.yaml. Keep max_steps == grpo.max_rollout_turns. + - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml + - resources_servers/circle_click/configs/circle_click.yaml circle_click_simple_agent: responses_api_agents: simple_agent: max_steps: 1 - logger: - log_dir: "logs/grpo-nemotron-omni-30ba3b-gymv-circle-click" - num_val_samples_to_print: 0 + log_dir: logs/grpo-nemotron-omni-30ba3b-gymv-circle-click wandb_enabled: true - tensorboard_enabled: true - mlflow_enabled: false - monitor_gpus: false - swanlab_enabled: false wandb: - project: "grpo-nemotron-omni-gymv" - name: "grpo-nemotron-omni-30ba3b-gymv-circle-click" - log_nemo_gym_full_result_tables: false - tensorboard: {} + project: grpo-nemotron-omni-gymv + name: grpo-nemotron-omni-30ba3b-gymv-circle-click mlflow: - experiment_name: "grpo-nemotron-omni-gymv" - run_name: "grpo-nemotron-omni-30ba3b-gymv-circle-click" - gpu_monitoring: - collection_interval: 10 - flush_interval: 10 - + experiment_name: grpo-nemotron-omni-gymv + run_name: grpo-nemotron-omni-30ba3b-gymv-circle-click cluster: gpus_per_node: 8 num_nodes: 2 - master_port_range_low: 1400 - master_port_range_high: 1999 From aa5ea0b198203a11b7cb1347789e12cfc5583b8f Mon Sep 17 00:00:00 2001 From: Yi-Fu Wu Date: Thu, 6 Aug 2026 00:06:07 -0700 Subject: [PATCH 24/27] test(vlm): add circle-click gym driver script, disabled for now The circle-click recipe YAML landed without the driver script and suite entry that tests/unit/test_recipes_and_test_suites.py requires, so test_all_recipe_yamls_accounted_for_in_test_suites failed on a 229 vs 228 count mismatch. Add the missing driver. circle_click is a NeMo-Gym env, so it runs through run_grpo_nemo_gym.py rather than run_vlm_grpo.py, and the script regenerates its data via the resources server's generate_data.py (the committed example.jsonl has 5 rows and no agent_ref) with disjoint train/eval seeds. List it in disabled.txt rather than nightly.txt for now: the recipe has not been run end to end, so its reward threshold is an unvalidated smoke bound. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Yi-Fu Wu --- tests/test_suites/disabled.txt | 8 +++ ...3b-circle-click-2n8g-megatron-tp2ep8.v1.sh | 70 +++++++++++++++++++ 2 files changed, 78 insertions(+) create mode 100755 tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-click-2n8g-megatron-tp2ep8.v1.sh diff --git a/tests/test_suites/disabled.txt b/tests/test_suites/disabled.txt index fe460c2185f..14c0c1e28df 100644 --- a/tests/test_suites/disabled.txt +++ b/tests/test_suites/disabled.txt @@ -5,3 +5,11 @@ # grpo-qwen3.5-35ba3b-2n8g-megatron-ep16tp2cp2 run hangs the same way on main, # so this is the pre-existing Qwen3.5 + Megatron + EP hang. tests/test_suites/vlm/vlm_grpo-qwen3.5-35ba3b-geo3k-2n8g-megatron-ep16.sh + +# First multimodal NeMo-Gym recipe (circle_click at the pinned Gym). Disabled on +# landing for two reasons: the recipe has not been run end to end yet, so its +# reward threshold is an unvalidated smoke bound; and nightly.txt is at 3755 of +# its 3800 GPU-hour budget, which this run's 32 GPU-hours would leave only 13 to +# spare. Move to nightly.txt once a real run confirms it converges and the +# budget has room. +tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-click-2n8g-megatron-tp2ep8.v1.sh diff --git a/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-click-2n8g-megatron-tp2ep8.v1.sh b/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-click-2n8g-megatron-tp2ep8.v1.sh new file mode 100755 index 00000000000..f7c4214e7cb --- /dev/null +++ b/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-click-2n8g-megatron-tp2ep8.v1.sh @@ -0,0 +1,70 @@ +#!/bin/bash +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +source $SCRIPT_DIR/common.env + +# ===== BEGIN CONFIG ===== +NUM_NODES=2 +GPUS_PER_NODE=8 +STEPS_PER_RUN=10 +MAX_STEPS=10 +NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) # Round up +# 30B MoE across 2 nodes, plus nemo_gym head-server startup and vLLM warmup on top +# of the 10 steps; 120 min leaves margin for teardown + metric dump. +NUM_MINUTES=120 +# ===== END CONFIG ===== + +exit_if_max_steps_reached + +cd $PROJECT_ROOT + +# circle_click generates its own data (no HF download). Regenerate rather than reuse the +# 5-row example.jsonl so the run never trains on stale committed data, and so train/eval +# are disjoint (distinct --seed-offset). +DATA_DIR=$EXP_DIR/data +mkdir -p $DATA_DIR +GYM_DIR=3rdparty/Gym-workspace/Gym +RAW_TRAIN=$DATA_DIR/circle_click_train_raw.jsonl +RAW_VALIDATION=$DATA_DIR/circle_click_validation_raw.jsonl +( cd $GYM_DIR && uv run python resources_servers/circle_click/generate_data.py \ + --n 512 --seed-offset 0 --out $PROJECT_ROOT/$RAW_TRAIN ) +( cd $GYM_DIR && uv run python resources_servers/circle_click/generate_data.py \ + --n 32 --seed-offset 100000 --out $PROJECT_ROOT/$RAW_VALIDATION ) + +# Attach `agent_ref` so rollouts are routed to the env's agent. The name must match the +# group registered in resources_servers/circle_click/configs/circle_click.yaml. +TRAIN_PATH=$DATA_DIR/circle_click_train.jsonl +VALIDATION_PATH=$DATA_DIR/circle_click_validation.jsonl +jq -c '. + {agent_ref: {name: "circle_click_simple_agent"}}' $RAW_TRAIN > $TRAIN_PATH +jq -c '. + {agent_ref: {name: "circle_click_simple_agent"}}' $RAW_VALIDATION > $VALIDATION_PATH + +# Run the experiment via the gym entrypoint (circle_click is a NeMo-Gym env, so this +# recipe runs through run_grpo_nemo_gym.py rather than run_vlm_grpo.py). +uv run examples/nemo_gym/run_grpo_nemo_gym.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 \ + data.train.data_path=$TRAIN_PATH \ + data.validation.data_path=$VALIDATION_PATH \ + $@ \ + 2>&1 | tee $RUN_LOG + +# Convert tensorboard logs to json +uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS + +# Only run metrics if the target step is reached +if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | map(tonumber) | max' $JSON_METRICS) -ge $MAX_STEPS ]]; then + # Smoke-level threshold: this recipe has not been run end to end yet, so assert only + # that the multimodal gym path produces non-zero reward. Tighten once real runs land. + uv run tests/check_metrics.py $JSON_METRICS \ + 'max(data["train/reward"]) > 0.0' + + # Clean up checkpoint directory after successful run to save space. + rm -rf "$CKPT_DIR" +fi From 91959244d3ba9c048805dbb00a6e56d72a8e9f07 Mon Sep 17 00:00:00 2001 From: Ali Roshan Ghias Date: Thu, 6 Aug 2026 06:16:18 -0700 Subject: [PATCH 25/27] feat: deduplicate multimodal GRPO payloads (cherry picked from commit c3b8dacb3320214e8d5404fb4f31de01a63df2f8) Signed-off-by: Ali Roshan Ghias --- docs/design-docs/multimodal-deduplication.md | 193 ++++++ docs/index.md | 1 + ...-circle-click-2n8g-megatron-tp2ep8.v1.yaml | 2 + ...ni-30ba3b-clevr-1n8g-automodel-ep8.v1.yaml | 2 + ...-30ba3b-clevr-1n8g-megatron-tp8ep8.v1.yaml | 3 + ...mni-30ba3b-mmpr-4n8g-automodel-ep8.v1.yaml | 2 + ...-30ba3b-mmpr-4n8g-megatron-tp8ep16.v1.yaml | 3 +- examples/configs/vlm_grpo_3B.yaml | 2 + examples/nemo_gym/run_grpo_nemo_gym.py | 2 + examples/run_vlm_grpo.py | 1 + nemo_rl/algorithms/async_utils/interfaces.py | 14 + .../algorithms/async_utils/replay_buffer.py | 46 +- .../async_utils/trajectory_collector.py | 55 +- nemo_rl/algorithms/grpo.py | 282 +++++++- nemo_rl/data/llm_message_utils.py | 79 +-- nemo_rl/data/multimodal_utils.py | 492 +++++++++++++- nemo_rl/distributed/batched_data_dict.py | 183 ++++- nemo_rl/environments/nemo_gym.py | 181 +++-- nemo_rl/experience/rollouts.py | 281 +++++++- nemo_rl/experience/sync_rollout_actor.py | 3 + nemo_rl/models/generation/interfaces.py | 4 +- nemo_rl/models/generation/vllm/utils.py | 51 +- .../models/generation/vllm/vllm_generation.py | 28 +- nemo_rl/models/policy/lm_policy.py | 26 + nemo_rl/utils/checkpoint.py | 3 + nemo_rl/utils/multimodal_payload_metrics.py | 403 +++++++++++ pyrefly.toml | 1 + tests/unit/algorithms/test_async_utils.py | 182 ++++- tests/unit/algorithms/test_grpo.py | 164 ++++- .../test_grpo_router_replay_async.py | 22 + tests/unit/data/test_llm_message_utils.py | 50 +- tests/unit/data/test_multimodal_dict.py | 141 +++- .../distributed/test_batched_data_dict.py | 300 ++++++++- tests/unit/environments/test_nemo_gym.py | 300 ++++++++- .../test_nemo_gym_router_replay.py | 6 +- tests/unit/experience/test_rollouts.py | 624 +++++++++++++++++- .../models/automodel/test_automodel_data.py | 25 +- .../unit/models/generation/test_vllm_utils.py | 9 +- .../megatron/test_nemotron_omni_model.py | 150 +++++ tests/unit/test_config_validation.py | 36 +- .../utils/test_multimodal_payload_metrics.py | 304 +++++++++ 41 files changed, 4396 insertions(+), 260 deletions(-) create mode 100644 docs/design-docs/multimodal-deduplication.md create mode 100644 nemo_rl/utils/multimodal_payload_metrics.py create mode 100644 tests/unit/utils/test_multimodal_payload_metrics.py diff --git a/docs/design-docs/multimodal-deduplication.md b/docs/design-docs/multimodal-deduplication.md new file mode 100644 index 00000000000..fb4301dbbfc --- /dev/null +++ b/docs/design-docs/multimodal-deduplication.md @@ -0,0 +1,193 @@ +# Multimodal payload deduplication + +Multimodal GRPO repeats each prompt for multiple generations. Without +deduplication, the driver, Ray object store, and replay buffer may store the +same image, video, or audio payload once per logical generation even though +the model must eventually receive the same logical batch. + +The deduplication feature keeps one physical copy of verified-equivalent media +and preserves a logical row-to-segment mapping until the common policy-worker +materialization boundary. It does not reduce the logical model batch or media +encoder compute. + +## Configuration and supported transport + +The schema defaults both options to disabled when a recipe omits them: + +```yaml +grpo: + deduplicate_multimodal_data: false + debug_payload_metrics: false +``` + +The shared VLM roots declare these disabled defaults. Qualified Nemotron Omni +recipes opt in to deduplication so they continuously exercise that path; other +VLM and text-only recipes remain opt-in. Payload metrics stay disabled because +protocol-5 sizing is diagnostic work, and qualification runs enable it +explicitly. + +`deduplicate_multimodal_data=true` is supported for vLLM generation with the +legacy Ray argument transport (`data_plane.enabled=false`). Configuration +validation fails early for other generation backends and for the TransferQueue +data plane rather than silently running an unqualified path. + +The implementation covers both synchronous and asynchronous GRPO, including +NeMo Gym image rollouts, DAPO cache concatenation, replay push/sample, replay +checkpoints, validation, log-probability calculation, policy training, and +KV-cache calibration. + +## Representation and safety contract + +`PackedTensor` has two representations: + +- Legacy values store one physical tensor entry per logical row. +- Deduplicated values store physical media segments plus CSR-style + `row_offsets` and `segment_indices`. + +Each physical segment receives opaque provenance when deduplication is first +enabled. Provenance survives deepcopy, pickle, replay, slicing, and sharding. +Concatenation re-interns segments only when their provenance matches. + +Prompt or problem identity is never evidence that media is equal. It may +narrow a search, but two trajectories from one prompt can receive different +media from a tool or environment. Those segments receive different provenance +and stay distinct. + +The following invariants apply: + +- Logical row count and order do not change. +- `PackedTensor.as_tensor()` reconstructs logical segment order before dynamic + shape padding and model-owned sequence packing. +- Dtype, device, packing dimension, and `pad_to_max_shape` are preserved. +- Mutable row, conversation, tool, and Gym request containers are copied. + Only explicit immutable media leaves are shared. +- Model-visible labels, masks, rewards, advantages, token-type tensors, and + log probabilities are never deduplicated. +- Coupled inputs such as `pixel_values`/`image_grid_thw`, + `pixel_values`/`imgs_sizes`, `pixel_values`/`num_frames`, and + `pixel_values_videos`/`video_grid_thw` must have compatible ordered per-row + logical segment counts before materialization. + +Operations that change logical rows—slice, filter, reorder, chunk, dynamic +batching, sequence packing, ordinary sharding, and replay concatenation—remap +the CSR indices. Each final worker shard re-interns matching provenance +locally. Copies required on different Ray workers are not incorrectly counted +as avoidable. + +## Where physical copies are removed + +There are four related boundaries: + +1. Prompt repeat copies row containers but shares immutable media leaves. +2. Conversation flattening merges ordered media segments without first + materializing one full tensor per repeated row. +3. Replay/DAPO concatenation represents a missing media key as an explicit + empty logical row when the dedup flag is enabled. +4. Final dynamic/sequence-packed worker shards re-intern provenance after + sharding has scattered prompt groups. + +For vLLM generation, native `vllm_content`, `vllm_images`, `vllm_videos`, and +`vllm_audios` remain the generation representation. Redundant policy-ready +`PackedTensor` media is omitted from a generation call only when every active +row has one of the image, video, or audio side channels that the vLLM formatter +actually consumes. Raw content alone and unconsumed path metadata are never +used to justify suppression. Policy media remains attached to the trajectory +for later training. + +Image, video, and audio leaves use the same ownership and payload-measurement +rules in the generic non-Gym representation. Processor-produced model inputs +and metadata use the generic `PackedTensor` path; native vLLM side channels use +explicit key and typed content recognition. The implementation does not use +model-name branches. + +## Gym and replay + +NeMo Gym may return trajectory-specific images. Initial images are processed +once on the unrepeated prompt only when deduplication is enabled, then +reattached by user-turn ordinal after the Gym Ray call. New Gym images are left +untouched. This avoids resending the original large policy payload through Gym +and avoids substituting images across generations. NeMo Gym audio and video +lineage are not implemented or qualified by this change; those media remain +supported only by the generic non-Gym path. + +Sparse text-only and multimodal trajectory groups are normalized only at the +GRPO replay/DAPO call sites. The default `BatchedDataDict.from_batches()` +missing-key behavior remains strict for other algorithms. + +Replay checkpoints store the compact representation directly. Save and restore +run inside the replay actor, so a buffer-sized state dict is not copied through +the long-lived driver frame. `checkpointing.save_replay_buffer=false` can skip +that checkpoint entirely for especially large runs; resume then regenerates +trajectories. The mapping is self-describing. Async replay assembly detects +restored compact tensors even when the current flag is off, so sparse flag-on +checkpoints can be resumed flag-off; legacy checkpoints also remain readable +after enabling the flag. When legacy and compact batches later concatenate, +legacy segments receive fresh provenance and restored compact segments remain +shared. + +## Policy backends and context parallelism + +All policy backends use the existing common materialization call, +`get_multimodal_dict(as_tensors=True)`. + +| Policy path | Dedup data-path status | +| --- | --- | +| AutoModel/DTensor, CP=1 | Supported by the common worker boundary; exact worker materialization is unit tested, and Qwen2.5-VL G=16 is exercised end to end. | +| Megatron, CP=1 | Supported by the common worker boundary and model-owned packing. | +| Megatron, CP=2 | Supported by the data representation; two-rank Nemotron model-ingress, logprob, loss, and gradient parity is tested. | +| AutoModel VLM, CP>1 | Not claimed; the upstream VLM worker currently rejects this topology independently of deduplication. | +| Megatron, CP>2 | Shared data operations are topology-independent, but model-level qualification is not yet claimed. | + +Model-family support follows processor and backend support. Qwen, Nemotron, and +other VLMs do not need dedicated dedup branches, but a family is only +end-to-end qualified when its maintained recipe and checkpoint have been run +with the feature enabled. This change is exercised with the maintained +Nemotron NeMo Gym image recipe and a Qwen2.5-VL native-rollout recipe on the +Megatron policy backend. Qwen2.5-VL also exercises the AutoModel CP=1 data path +at G=16 with the expected payload reduction. AutoModel TMPE is not used as +correctness evidence for that qualification, and exact trajectory identity +after independently updating and refitting two policies is not claimed. +Concrete job IDs, W&B runs, parity results, and transport measurements belong +in the pull-request validation report so they remain tied to the exact code +revision that was run. + +Gemma is excluded from this change's evaluation set. Its vLLM generation and +AutoModel policy paths have a pre-existing token-logprob mismatch with +deduplication both disabled and enabled, so that model cannot provide a valid +deduplication correctness signal. Resolving the Gemma backend mismatch is +separate work. Audio/video model runs remain unqualified; their shared data +primitives are covered by focused tests. + +## Payload metrics + +`debug_payload_metrics=true` emits stable lines beginning with +`▶ [PAYLOAD]`. Metrics include: + +- physical and logical media bytes; +- physical and logical media segment counts; +- estimated saved bytes and physical-to-logical ratio; +- cloudpickle protocol-5 frame plus out-of-band buffer size; +- total and maximum serialized size for exact unique final DP-shard arguments. + +The measured Python object is the exact object passed at that Ray boundary. +The protocol-5 serialized size is a serialization proxy, not a claim that it +is an exact Ray object-store allocation. Object-store qualification should +pair it with tracked object IDs or an isolated matched cluster delta. Totals +count each DP-shard object once; they do not multiply bytes for TP/CP replicas +that consume the same Ray object or reference. + +When `debug_payload_metrics=false`, call sites return before walking media or +serializing payloads. + +## Qualification expectations + +Correctness tests compare dedup-on materialization against the legacy +representation exactly. Dynamic resolution, pack dimensions 0 and 1, +multi-turn divergence, missing media keys, replay checkpoints, native +image/video/audio leaves, dynamic batching, sequence packing, and shard-local +re-interning require focused coverage. + +End-to-end qualification should record the exact commit, container, model +revision, recipe, overrides, hardware, and Ray version. Pre-shard savings are +compared with `1 - unique_physical/logical_occurrences`; worker savings use the +same formula independently on each final shard. diff --git a/docs/index.md b/docs/index.md index 951f6631716..89f26abd9d2 100644 --- a/docs/index.md +++ b/docs/index.md @@ -358,6 +358,7 @@ design-docs/training-backends.md design-docs/sequence-packing-and-dynamic-batching.md design-docs/env-vars.md design-docs/nemo-gym-integration.md +design-docs/multimodal-deduplication.md design-docs/modelopt-real-quant-architecture.md design-docs/nccl-reshard-refit.md ``` diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-click-2n8g-megatron-tp2ep8.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-click-2n8g-megatron-tp2ep8.v1.yaml index 3695064845c..3ccf06abe9f 100644 --- a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-click-2n8g-megatron-tp2ep8.v1.yaml +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-click-2n8g-megatron-tp2ep8.v1.yaml @@ -1,5 +1,7 @@ defaults: ../../vlm_grpo_3B_megatron.yaml grpo: + deduplicate_multimodal_data: true + debug_payload_metrics: false num_prompts_per_step: 1 num_val_generations_per_prompt: 1 max_num_steps: 100 diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-automodel-ep8.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-automodel-ep8.v1.yaml index c2ed8cbdb39..381bb50a128 100644 --- a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-automodel-ep8.v1.yaml +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-automodel-ep8.v1.yaml @@ -1,5 +1,7 @@ defaults: ../../vlm_grpo_3B.yaml grpo: + deduplicate_multimodal_data: true + debug_payload_metrics: false num_prompts_per_step: 32 val_at_start: true checkpointing: diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8.v1.yaml index b738cf17a13..be7e172f22b 100644 --- a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8.v1.yaml +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8.v1.yaml @@ -1,4 +1,7 @@ defaults: ../../vlm_grpo_3B_megatron.yaml +grpo: + deduplicate_multimodal_data: true + debug_payload_metrics: false loss_fn: reference_policy_kl_penalty: 0.0 checkpointing: diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-automodel-ep8.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-automodel-ep8.v1.yaml index 3a0f328a36a..dc234140a4f 100644 --- a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-automodel-ep8.v1.yaml +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-automodel-ep8.v1.yaml @@ -1,5 +1,7 @@ defaults: ../../vlm_grpo_3B.yaml grpo: + deduplicate_multimodal_data: true + debug_payload_metrics: false num_prompts_per_step: 32 overlong_filtering: true seq_logprob_error_threshold: 2 diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-megatron-tp8ep16.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-megatron-tp8ep16.v1.yaml index 844917c1fc2..115a932e10e 100644 --- a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-megatron-tp8ep16.v1.yaml +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-megatron-tp8ep16.v1.yaml @@ -3,7 +3,8 @@ grpo: num_prompts_per_step: 512 overlong_filtering: true zero_variance_prompt_filtering: false - deduplicate_multimodal_data: false + deduplicate_multimodal_data: true + debug_payload_metrics: false loss_fn: ratio_clip_max: 0.28 use_on_policy_kl_approximation: true diff --git a/examples/configs/vlm_grpo_3B.yaml b/examples/configs/vlm_grpo_3B.yaml index 916a4a9c4f0..19d82e5c845 100644 --- a/examples/configs/vlm_grpo_3B.yaml +++ b/examples/configs/vlm_grpo_3B.yaml @@ -3,6 +3,8 @@ defaults: "grpo_math_1B.yaml" grpo: + deduplicate_multimodal_data: false + debug_payload_metrics: false num_prompts_per_step: 8 reward_shaping: overlong_buffer_length: 512 diff --git a/examples/nemo_gym/run_grpo_nemo_gym.py b/examples/nemo_gym/run_grpo_nemo_gym.py index 4d83b9c129f..cfa6fa464ba 100644 --- a/examples/nemo_gym/run_grpo_nemo_gym.py +++ b/examples/nemo_gym/run_grpo_nemo_gym.py @@ -315,6 +315,7 @@ def main() -> None: max_trajectory_age_steps=config.grpo.async_grpo.max_trajectory_age_steps, teacher_worker_groups=teacher_worker_groups, alias_to_group_alias=alias_to_group_alias, + processor=processor, ) else: print("🚀 Running synchronous GRPO training") @@ -333,6 +334,7 @@ def main() -> None: checkpointer, grpo_state, master_config, + processor=processor, ) diff --git a/examples/run_vlm_grpo.py b/examples/run_vlm_grpo.py index e2b17e43baf..a0030ca20ad 100644 --- a/examples/run_vlm_grpo.py +++ b/examples/run_vlm_grpo.py @@ -146,6 +146,7 @@ def main() -> None: checkpointer, grpo_state, master_config, + processor=processor, ) diff --git a/nemo_rl/algorithms/async_utils/interfaces.py b/nemo_rl/algorithms/async_utils/interfaces.py index 892f3293a0f..824f718b6e7 100644 --- a/nemo_rl/algorithms/async_utils/interfaces.py +++ b/nemo_rl/algorithms/async_utils/interfaces.py @@ -73,6 +73,20 @@ def load_state_dict( """Restore state produced by ``state_dict``.""" ... + def save_to_path(self, path: str) -> int: + """Serialize state directly from the replay actor.""" + ... + + def load_from_path( + self, + path: str, + num_prompts_per_step: int | None = None, + current_training_step: int | None = None, + max_age_steps: int | None = None, + ) -> dict[str, int]: + """Restore state directly in the replay actor.""" + ... + def get_trajectories_needed( self, target_step: int, diff --git a/nemo_rl/algorithms/async_utils/replay_buffer.py b/nemo_rl/algorithms/async_utils/replay_buffer.py index 8b182b349ae..7d408d7b4fe 100644 --- a/nemo_rl/algorithms/async_utils/replay_buffer.py +++ b/nemo_rl/algorithms/async_utils/replay_buffer.py @@ -13,6 +13,7 @@ # limitations under the License. import asyncio +import gc import statistics import threading as _threading import uuid @@ -21,11 +22,16 @@ from typing import Any, Iterable, Optional import ray +import torch from nemo_rl.algorithms.async_utils.interfaces import ReplayBufferProtocol from nemo_rl.data_plane import KVBatchMeta from nemo_rl.data_plane.schema import ROUTED_EXPERTS_FIELD -from nemo_rl.experience.interfaces import PromptGroupRecord +from nemo_rl.experience.interfaces import ( + NEMO_GYM_TASK_INDEX_KEY, + NEXT_NEMO_GYM_TASK_INDEX_KEY, + PromptGroupRecord, +) from nemo_rl.experience.payload import pack_payload, record_to_train_batch from nemo_rl.utils.r3_trace import trace_rollout_payload @@ -340,6 +346,44 @@ def state_dict(self) -> dict[str, Any]: "max_size": self.max_size, } + def save_to_path(self, path: str) -> int: + """Serialize inside the actor without materializing the buffer on the driver.""" + state = self.state_dict() + torch.save(state, path) + num_trajectories = len(state["trajectories"]) + del state + gc.collect() + return num_trajectories + + def load_from_path( + self, + path: str, + num_prompts_per_step: int | None = None, + current_training_step: int | None = None, + max_age_steps: int | None = None, + ) -> dict[str, int]: + """Restore inside the actor and return only compact coordination metadata.""" + state = torch.load(path, weights_only=False) + saved_task_indices = [ + int(trajectory[NEMO_GYM_TASK_INDEX_KEY]) + for trajectory in state.get("trajectories", []) + if trajectory.get(NEMO_GYM_TASK_INDEX_KEY) is not None + ] + next_task_index = max(saved_task_indices, default=-1) + 1 + num_trajectories = len(state["trajectories"]) + self.load_state_dict( + state, + num_prompts_per_step=num_prompts_per_step, + current_training_step=current_training_step, + max_age_steps=max_age_steps, + ) + del state + gc.collect() + return { + "num_trajectories": num_trajectories, + NEXT_NEMO_GYM_TASK_INDEX_KEY: next_task_index, + } + def load_state_dict( self, state: dict[str, Any], diff --git a/nemo_rl/algorithms/async_utils/trajectory_collector.py b/nemo_rl/algorithms/async_utils/trajectory_collector.py index ca57644afa8..f87673b4c57 100644 --- a/nemo_rl/algorithms/async_utils/trajectory_collector.py +++ b/nemo_rl/algorithms/async_utils/trajectory_collector.py @@ -38,10 +38,16 @@ ) from nemo_rl.experience.rollouts import ( RolloutGroupResult, + attach_initial_nemo_gym_image_payloads, run_async_multi_turn_rollout_groups, ) from nemo_rl.models.generation.interfaces import GenerationConfig, GenerationInterface from nemo_rl.utils.logger import should_log_nemo_gym_full_result_tables +from nemo_rl.utils.multimodal_payload_metrics import ( + collect_multimodal_payload_metrics, + drain_multimodal_payload_metrics, + print_multimodal_payload_metrics, +) from nemo_rl.utils.timer import ThreadSafeTimer TokenizerType = PreTrainedTokenizerBase @@ -66,6 +72,7 @@ def __init__( alias_to_group_alias: Optional[dict[str, str]] = None, on_policy_distillation_cfg: Optional[dict[str, Any]] = None, next_nemo_gym_task_index: int = 0, + processor: Any = None, ): self.policy_generation = policy_generation self.tokenizer = tokenizer @@ -75,6 +82,7 @@ def __init__( self.teacher_worker_groups = teacher_worker_groups or {} self.alias_to_group_alias = alias_to_group_alias or {} self.on_policy_distillation_cfg = on_policy_distillation_cfg or {} + self.processor = processor self._has_distillation_teachers = bool(self.teacher_worker_groups) self._teacher_seq_pad_multiple = teacher_seq_pad_multiple( self.teacher_worker_groups, @@ -428,7 +436,23 @@ def _process_batch(self, batch: BatchedDataDict[DatumSpec]) -> None: rollout_batch = batch.slice(0, num_prompts_to_generate) if use_nemo_gym: self._stamp_nemo_gym_task_indices(rollout_batch) - repeated_batch = rollout_batch.repeat_interleave(num_generations) + if self.master_config.grpo.deduplicate_multimodal_data: + attach_initial_nemo_gym_image_payloads( + rollout_batch, self.processor + ) + repeated_batch = rollout_batch.repeat_interleave( + num_generations, + share_immutable_media=( + self.master_config.grpo.deduplicate_multimodal_data + ), + ) + print_multimodal_payload_metrics( + collect_multimodal_payload_metrics( + repeated_batch, + "prompt_repeat_async", + enabled=self.master_config.grpo.debug_payload_metrics, + ) + ) def _run_rollout_batch() -> None: asyncio.run( @@ -605,6 +629,15 @@ def get_efficiency_metrics(self) -> dict[str, float]: self._efficiency_timer.get_timing_metrics(reduction_op="sum"), ) + async def drain_payload_metrics(self) -> dict[str, int | float]: + """Close one drain-to-drain collector/Gym telemetry interval. + + Rollout collection is concurrent with training, so the interval is not + claimed to own the sampled training batch. Call-normalized metrics make + intervals comparable even when their background transfer counts differ. + """ + return drain_multimodal_payload_metrics() + def get_rollouts_state(self) -> dict[str, int]: """Get collector-side rollout state for checkpointing.""" return {NEXT_NEMO_GYM_TASK_INDEX_KEY: self._next_nemo_gym_task_index} @@ -777,6 +810,10 @@ async def _iter_rollout_groups( mask_env_flagged_samples=should_mask_flagged_samples( self.master_config.env ), + deduplicate_multimodal_data=( + self.master_config.grpo.deduplicate_multimodal_data + ), + debug_payload_metrics=self.master_config.grpo.debug_payload_metrics, ): task_index = rollout_result.task_index if task_index is None: @@ -801,6 +838,9 @@ async def _iter_rollout_groups( num_generations=num_generations, max_rollout_turns=self.master_config.grpo.max_rollout_turns, greedy=False, + deduplicate_multimodal_data=( + self.master_config.grpo.deduplicate_multimodal_data + ), ): yield rollout_result @@ -922,11 +962,22 @@ async def _enqueue_rollout_group( } if rollout_result.task_index is not None: trajectory_group[NEMO_GYM_TASK_INDEX_KEY] = rollout_result.task_index - backoff_delay = 0.01 backoff_started_at: float | None = None try: while self.running: + # Every retry is a distinct Ray submission of the full payload. + print_multimodal_payload_metrics( + collect_multimodal_payload_metrics( + ( + trajectory_group, + generation_weight_version, + target_weight_version, + ), + "replay_push", + enabled=self.master_config.grpo.debug_payload_metrics, + ) + ) status = await self.replay_buffer.add.remote( trajectory_group, generation_weight_version, diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index 678dbd3c0c5..6835854633e 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -69,6 +69,7 @@ batched_message_log_to_flat_message, get_keys_from_message_log, ) +from nemo_rl.data.multimodal_utils import PackedTensor from nemo_rl.data.utils import extract_necessary_env_names, load_dataloader_state from nemo_rl.data_plane.interfaces import DataPlaneConfig from nemo_rl.distributed.batched_data_dict import BatchedDataDict @@ -88,6 +89,7 @@ ) from nemo_rl.experience.rollouts import ( EffortLevelsConfig, + attach_initial_nemo_gym_image_payloads, backfill_missing_routed_experts, get_nemo_gym_thinking_tags, run_async_multi_turn_rollout, @@ -124,6 +126,12 @@ should_log_nemo_gym_full_result_tables, ) from nemo_rl.utils.memory_tracker import MemoryTracker +from nemo_rl.utils.multimodal_payload_metrics import ( + collect_multimodal_payload_metrics, + drain_multimodal_payload_metrics, + merge_multimodal_payload_metrics, + print_multimodal_payload_metrics, +) from nemo_rl.utils.nsys import maybe_gpu_profile_step from nemo_rl.utils.timer import TimeoutChecker, Timer from nemo_rl.utils.venvs import create_local_venv_on_each_node @@ -157,6 +165,29 @@ def _get_next_nemo_gym_task_index( return next_task_index +def _save_async_replay_buffer_checkpoint( + replay_buffer: Any, + checkpoint_path: str, + checkpointing_config: CheckpointingConfig, +) -> int | None: + """Checkpoint replay state inside its actor, or skip it when configured.""" + if not checkpointing_config.get("save_replay_buffer", True): + print( + "⏭️ Skipping replay buffer checkpoint " + "(checkpointing.save_replay_buffer=false)" + ) + return None + + print("📦 Saving replay buffer state...") + num_buffered_trajectories = ray.get( + replay_buffer.save_to_path.remote( + os.path.join(checkpoint_path, "replay_buffer.pt") + ) + ) + print(f"✅ Saved replay buffer with {num_buffered_trajectories} trajectories") + return num_buffered_trajectories + + class RewardScalingConfig(BaseModel, extra="allow"): """Configure linear reward scaling with clamping. @@ -282,6 +313,11 @@ class GRPOConfig(BaseModel, extra="allow"): malformed_thinking_advantage: float | None = None # Advantage estimator configuration (grpo or reinforce_plus_plus) adv_estimator: AdvEstimatorConfig = Field(default_factory=AdvEstimatorConfig) + # Share and compact immutable image/video/audio payload segments across + # logical GRPO rows. Prompt identity is never used as proof of equality. + deduplicate_multimodal_data: bool = False + # Emit exact-boundary and logical-vs-physical payload metrics. + debug_payload_metrics: bool = False @dataclass @@ -343,6 +379,24 @@ class MasterConfig(BaseModel, extra="allow"): # =============================================================================== +def _validate_multimodal_dedup_capability(master_config: MasterConfig) -> None: + """Reject configurations whose media transfer path is not qualified.""" + if not master_config.grpo.deduplicate_multimodal_data: + return + + generation_config = master_config.policy["generation"] + if generation_config.get("backend") != "vllm": + raise NotImplementedError( + "grpo.deduplicate_multimodal_data=true is currently qualified " + "only with policy.generation.backend=vllm." + ) + if (master_config.data_plane or {}).get("enabled", False): + raise NotImplementedError( + "grpo.deduplicate_multimodal_data=true is currently supported " + "only when data_plane.enabled=false." + ) + + def setup( master_config: MasterConfig, tokenizer: TokenizerType, @@ -397,6 +451,7 @@ def setup( ) if generation_config["backend"] == "vllm": normalize_vllm_refit_config(cast(VllmConfig, generation_config)) + _validate_multimodal_dedup_capability(master_config) # Set seed for all random number generators set_seed(grpo_config.seed) @@ -938,6 +993,7 @@ def _spinup_nemo_gym(base_urls, model_name): # vllm model loading prefers clean environment, initialize policy_generation before policy in colocated mode backend = generation_config["backend"] generation_config["model_name"] = policy_config["model_name"] # Needed for vLLM + generation_config["debug_payload_metrics"] = grpo_config.debug_payload_metrics remote_transport = None remote_synchronizer_cls = None remote_baseline_init_refs: list[Any] = [] @@ -1006,6 +1062,8 @@ def init_policy(): init_optimizer=True, init_reference_model=init_reference_model, ) + # Keep custom policy_factory call signatures backward compatible. + p.debug_payload_metrics = grpo_config.debug_payload_metrics if remote_transport is not None: assert remote_synchronizer_cls is not None remote_baseline_init_refs.extend( @@ -1611,7 +1669,10 @@ def dynamic_sampling( filtered_repeated_batch if batch_cache is None else BatchedDataDict.from_batches( - [batch_cache, filtered_repeated_batch] + [batch_cache, filtered_repeated_batch], + allow_missing_packed_tensors=( + master_config.grpo.deduplicate_multimodal_data + ), ) ) filtered_repeated_batch = batch_cache @@ -1984,6 +2045,19 @@ def _preserve_router_replay_routed_experts( target["routed_experts"] = flat_messages["routed_experts"] +def _should_normalize_sparse_replay_media( + batches: list[BatchedDataDict], + *, + deduplicate_multimodal_data: bool, +) -> bool: + """Keep sparse compact checkpoints readable across a flag transition.""" + return deduplicate_multimodal_data or any( + isinstance(value, PackedTensor) and value.deduplication_enabled + for batch in batches + for value in batch.values() + ) + + def _build_async_grpo_train_data( flat_messages: BatchedDataDict, input_lengths: torch.Tensor, @@ -2610,6 +2684,7 @@ def grpo_train( checkpointer: CheckpointManager, grpo_save_state: GRPOSaveState, master_config: MasterConfig, + processor: Optional[AutoProcessor] = None, ) -> None: """Run GRPO training algorithm.""" timer = Timer(context={"worker": "driver"}) @@ -2685,10 +2760,19 @@ def grpo_train( step=0, master_config=master_config, logger=logger, + processor=processor, ) policy_generation.finish_generation() logger.log_metrics(val_metrics, current_step, prefix="validation") logger.log_metrics(validation_timings, current_step, prefix="timing/validation") + if master_config.grpo.debug_payload_metrics: + validation_payload_metrics = drain_multimodal_payload_metrics() + if validation_payload_metrics: + logger.log_metrics( + validation_payload_metrics, + current_step, + prefix="validation", + ) stop_message = _validation_early_stop_message( val_metrics, stop_at_validation_threshold, @@ -2745,10 +2829,25 @@ def grpo_train( # Prepare batch print("▶ Preparing batch...", flush=True) with timer.time("data_processing"): + if ( + master_config.grpo.deduplicate_multimodal_data + and _should_use_nemo_gym(master_config) + ): + attach_initial_nemo_gym_image_payloads(batch, processor) # Repeat batch items repeated_batch: BatchedDataDict[DatumSpec] = ( batch.repeat_interleave( - master_config.grpo.num_generations_per_prompt + master_config.grpo.num_generations_per_prompt, + share_immutable_media=( + master_config.grpo.deduplicate_multimodal_data + ), + ) + ) + print_multimodal_payload_metrics( + collect_multimodal_payload_metrics( + repeated_batch, + "prompt_repeat", + enabled=master_config.grpo.debug_payload_metrics, ) ) # Convert LLMMessageLogType to FlatMessagesType for generation @@ -2852,6 +2951,12 @@ def grpo_train( mask_env_flagged_samples=should_mask_flagged_samples( master_config.env ), + deduplicate_multimodal_data=( + master_config.grpo.deduplicate_multimodal_data + ), + debug_payload_metrics=( + master_config.grpo.debug_payload_metrics + ), ) input_ids = nemo_gym_rollout_result.input_ids repeated_batch = nemo_gym_rollout_result.final_batch @@ -2873,6 +2978,9 @@ def grpo_train( ], max_rollout_turns=master_config.grpo.max_rollout_turns, greedy=False, + deduplicate_multimodal_data=( + master_config.grpo.deduplicate_multimodal_data + ), ) else: repeated_batch, rollout_metrics = run_multi_turn_rollout( @@ -2885,6 +2993,9 @@ def grpo_train( ], max_rollout_turns=master_config.grpo.max_rollout_turns, greedy=False, + deduplicate_multimodal_data=( + master_config.grpo.deduplicate_multimodal_data + ), ) policy_generation.finish_generation() # Collect generation logger metrics for performance reporting after each generation step @@ -3055,6 +3166,13 @@ def grpo_train( as_tensors=False ) train_data.update(extra_multimodal_data) + print_multimodal_payload_metrics( + collect_multimodal_payload_metrics( + train_data, + "rollout_to_policy", + enabled=master_config.grpo.debug_payload_metrics, + ) + ) # Router replay (R3) on the legacy data_plane.enabled=false # driver path: routed_experts already rides flat_messages # (attached to message_log during rollout, then batched into @@ -3227,12 +3345,19 @@ def grpo_train( ) early_stop_message: Optional[str] = None - # Run validation if it's a validation step or last step with val_at_end - if ( + should_run_validation = ( val_period > 0 and (total_steps + 1) >= val_start_at and (total_steps + 1) % val_period == 0 - ) or (val_at_end and is_last_step): + ) or (val_at_end and is_last_step) + + # Keep training and validation traffic in separate metric intervals. + payload_metrics: dict[str, int | float] = {} + if master_config.grpo.debug_payload_metrics: + payload_metrics = drain_multimodal_payload_metrics() + + # Run validation if it's a validation step or last step with val_at_end + if should_run_validation: memory_tracker.snapshot_start_of_stage("Validation", dir()) if NEED_REFIT and POLICY_GENERATION_STALE: refit_metrics = refit_policy_generation( @@ -3255,6 +3380,7 @@ def grpo_train( step=total_steps + 1, master_config=master_config, logger=logger, + processor=processor, ) policy_generation.finish_generation() logger.log_metrics( @@ -3263,6 +3389,14 @@ def grpo_train( logger.log_metrics( val_metrics, total_steps + 1, prefix="validation" ) + if master_config.grpo.debug_payload_metrics: + validation_payload_metrics = drain_multimodal_payload_metrics() + if validation_payload_metrics: + logger.log_metrics( + validation_payload_metrics, + total_steps + 1, + prefix="validation", + ) early_stop_message = _validation_early_stop_message( val_metrics, stop_at_validation_threshold, @@ -3603,6 +3737,9 @@ def grpo_train( train_results, metrics, timing_metrics, master_config ) + if payload_metrics: + logger.log_metrics(payload_metrics, total_steps + 1, prefix="") + if refit_metrics: logger.log_metrics(refit_metrics, total_steps + 1, prefix="refit") logger.log_metrics(metrics, total_steps + 1, prefix="train") @@ -3673,6 +3810,7 @@ def validate( step: int, master_config: MasterConfig, logger: Optional[Logger] = None, + processor: Optional[AutoProcessor] = None, ) -> tuple[dict[str, Any], dict[str, Any]]: """Run validation on the validation dataset.""" if val_dataloader is None: @@ -3702,6 +3840,8 @@ def validate( # Use async rollouts when enabled by config/backend defaults. # We cascade NeMo-Gym first since NeMo-Gym also uses async rollouts. if _should_use_nemo_gym(master_config): + if master_config.grpo.deduplicate_multimodal_data: + attach_initial_nemo_gym_image_payloads(val_batch, processor) generation_config = master_config.policy["generation"] nemo_gym_rollout_result = run_nemo_gym_rollout_sync( policy_generation=policy_generation, @@ -3722,6 +3862,10 @@ def validate( mask_env_flagged_samples=should_mask_flagged_samples( master_config.env ), + deduplicate_multimodal_data=( + master_config.grpo.deduplicate_multimodal_data + ), + debug_payload_metrics=master_config.grpo.debug_payload_metrics, ) val_batch = nemo_gym_rollout_result.final_batch gen_metrics = nemo_gym_rollout_result.rollout_metrics @@ -3735,6 +3879,9 @@ def validate( max_seq_len=master_config.policy["max_total_sequence_length"], max_rollout_turns=master_config.grpo.max_rollout_turns, greedy=False, + deduplicate_multimodal_data=( + master_config.grpo.deduplicate_multimodal_data + ), ) else: val_batch, gen_metrics = run_multi_turn_rollout( @@ -3745,6 +3892,9 @@ def validate( max_seq_len=master_config.policy["max_total_sequence_length"], max_rollout_turns=master_config.grpo.max_rollout_turns, greedy=False, + deduplicate_multimodal_data=( + master_config.grpo.deduplicate_multimodal_data + ), ) total_rewards.extend(val_batch["total_reward"].tolist()) @@ -3883,6 +4033,7 @@ def async_grpo_train( max_trajectory_age_steps: int = 1, teacher_worker_groups: Optional[dict[str, Any]] = None, alias_to_group_alias: Optional[dict[str, str]] = None, + processor: Optional[AutoProcessor] = None, ) -> None: """Run asynchronous GRPO training with replay buffer. @@ -3900,6 +4051,8 @@ def async_grpo_train( grpo_save_state: Training state master_config: Master configuration max_trajectory_age_steps: Maximum age (in training steps) for trajectories to be used in training + processor: Optional multimodal processor used to attach compact policy + media to NeMo Gym prompt rows. """ # Ensure we are running with a compatible async generation backend. # Async GRPO (with in-flight weight updates) supports vLLM, Megatron, and TRT-LLM; @@ -4026,18 +4179,15 @@ def async_grpo_train( ) last_checkpoint_path = checkpointer.get_latest_checkpoint_path() - replay_buffer_state = None + replay_buffer_restore_metadata: dict[str, int] | None = None rollouts_state = None if last_checkpoint_path is not None: replay_buffer_path = os.path.join(last_checkpoint_path, "replay_buffer.pt") if os.path.exists(replay_buffer_path): print(f"📦 Restoring replay buffer from checkpoint: {replay_buffer_path}") - # weights_only=False: trajectories are pickled BatchedDataDict/dicts, - # not plain tensors. The checkpoint is a trusted same-job artifact. - replay_buffer_state = torch.load(replay_buffer_path, weights_only=False) - ray.get( - replay_buffer.load_state_dict.remote( - replay_buffer_state, + replay_buffer_restore_metadata = ray.get( + replay_buffer.load_from_path.remote( + replay_buffer_path, num_prompts_per_step=num_prompts_per_step, current_training_step=step, max_age_steps=max_trajectory_age_steps, @@ -4057,8 +4207,13 @@ def async_grpo_train( next_nemo_gym_task_index = _get_next_nemo_gym_task_index( rollouts_state=rollouts_state, - replay_buffer_state=replay_buffer_state, + replay_buffer_state=None, ) + if replay_buffer_restore_metadata is not None: + next_nemo_gym_task_index = max( + next_nemo_gym_task_index, + replay_buffer_restore_metadata[NEXT_NEMO_GYM_TASK_INDEX_KEY], + ) _tc_py_exec = get_actor_python_env( "nemo_rl.algorithms.async_utils.AsyncTrajectoryCollector" @@ -4096,6 +4251,7 @@ def async_grpo_train( alias_to_group_alias=alias_to_group_alias, on_policy_distillation_cfg=opd_module._opd_cfg(master_config), next_nemo_gym_task_index=next_nemo_gym_task_index, + processor=processor, ) # Start trajectory collection in background @@ -4145,7 +4301,7 @@ def async_grpo_train( if val_at_start and step == 0: print("\n🔍 Running initial validation...") # Pause trajectory collection during initial validation - trajectory_collector.pause.remote() + ray.get(trajectory_collector.pause.remote()) initial_val_metrics: Optional[dict[str, Any]] = None try: @@ -4157,11 +4313,20 @@ def async_grpo_train( step=0, master_config=master_config, logger=logger, + processor=processor, ) initial_val_metrics = val_metrics policy_generation.finish_generation() logger.log_metrics(val_metrics, step, prefix="validation") logger.log_metrics(validation_timings, step, prefix="timing/validation") + if master_config.grpo.debug_payload_metrics: + validation_payload_metrics = drain_multimodal_payload_metrics() + if validation_payload_metrics: + logger.log_metrics( + validation_payload_metrics, + step, + prefix="validation", + ) print("✅ Initial validation completed successfully") except Exception as e: print(f"❌ Initial validation failed: {e}") @@ -4317,6 +4482,14 @@ def async_grpo_train( max_age_steps=max_trajectory_age_steps, ) ) + if sample_result is not None: + print_multimodal_payload_metrics( + collect_multimodal_payload_metrics( + sample_result, + "replay_sample", + enabled=master_config.grpo.debug_payload_metrics, + ) + ) if ( sample_result is None @@ -4398,7 +4571,16 @@ def async_grpo_train( # Concatenate per-prompt groups into a single training batch per_prompt_batches = [t["batch"] for t in trajectories] - repeated_batch = BatchedDataDict.from_batches(per_prompt_batches) + normalize_sparse_media = _should_normalize_sparse_replay_media( + per_prompt_batches, + deduplicate_multimodal_data=( + master_config.grpo.deduplicate_multimodal_data + ), + ) + repeated_batch = BatchedDataDict.from_batches( + per_prompt_batches, + allow_missing_packed_tensors=normalize_sparse_media, + ) # Teacher logprobs are stored in batch dict by collection-time # computation and padded by from_batches. Extract here. @@ -4513,6 +4695,13 @@ def async_grpo_train( repeated_batch, master_config.policy, ) + print_multimodal_payload_metrics( + collect_multimodal_payload_metrics( + train_data, + "rollout_to_policy_async", + enabled=master_config.grpo.debug_payload_metrics, + ) + ) train_data.to("cpu") # Training phase (same as sync version) @@ -4687,17 +4876,30 @@ def async_grpo_train( # Validation val_metrics, validation_timings = None, None is_last_step = step + 1 == master_config.grpo.max_num_steps - - # Run validation if it's a validation step or last step with val_at_end - if ( + should_run_validation = ( val_period > 0 and (step + 1) >= val_start_at and (step + 1) % val_period == 0 - ) or (val_at_end and is_last_step): - with timer.time("idle/validation"): - # Pause trajectory collection during validation to reduce memory pressure - trajectory_collector.pause.remote() + ) or (val_at_end and is_last_step) + + payload_metrics: dict[str, int | float] = {} + if should_run_validation: + # Stop new dispatch before separating the training and + # validation payload-metric intervals. + ray.get(trajectory_collector.pause.remote()) + if master_config.grpo.debug_payload_metrics: + payload_metrics = merge_multimodal_payload_metrics( + [ + drain_multimodal_payload_metrics(), + ray.get( + trajectory_collector.drain_payload_metrics.remote() + ), + ] + ) + # Run validation if it's a validation step or last step with val_at_end + if should_run_validation: + with timer.time("idle/validation"): if NEED_REFIT and POLICY_GENERATION_STALE: refit_metrics = refit_policy_generation( policy, @@ -4715,12 +4917,23 @@ def async_grpo_train( step=step + 1, master_config=master_config, logger=logger, + processor=processor, ) policy_generation.finish_generation() logger.log_metrics( validation_timings, step + 1, prefix="timing/validation" ) logger.log_metrics(val_metrics, step + 1, prefix="validation") + if master_config.grpo.debug_payload_metrics: + validation_payload_metrics = ( + drain_multimodal_payload_metrics() + ) + if validation_payload_metrics: + logger.log_metrics( + validation_payload_metrics, + step + 1, + prefix="validation", + ) early_stop_message = _validation_early_stop_message( val_metrics, stop_at_validation_threshold, @@ -4731,8 +4944,6 @@ def async_grpo_train( print(early_stop_message, flush=True) # Explicit GPU memory cleanup after validation in async mode - import gc - gc.collect() torch.cuda.empty_cache() @@ -4905,15 +5116,10 @@ def async_grpo_train( actual_dataloader_state, os.path.join(checkpoint_path, "train_dataloader.pt"), ) - print("📦 Saving replay buffer state...") - replay_buffer_state = ray.get(replay_buffer.state_dict.remote()) - torch.save( - replay_buffer_state, - os.path.join(checkpoint_path, "replay_buffer.pt"), - ) - print( - "✅ Saved replay buffer with " - f"{len(replay_buffer_state['trajectories'])} trajectories" + _save_async_replay_buffer_checkpoint( + replay_buffer, + checkpoint_path, + master_config.checkpointing, ) rollouts_state = ray.get( trajectory_collector.get_rollouts_state.remote() @@ -5057,6 +5263,16 @@ def async_grpo_train( merged_efficiency, total_wall_time, step + 1 ) + if master_config.grpo.debug_payload_metrics and not should_run_validation: + payload_metrics = merge_multimodal_payload_metrics( + [ + drain_multimodal_payload_metrics(), + ray.get(trajectory_collector.drain_payload_metrics.remote()), + ] + ) + if payload_metrics: + logger.log_metrics(payload_metrics, step + 1, prefix="") + if refit_metrics: logger.log_metrics(refit_metrics, step + 1, prefix="refit") logger.log_metrics(performance_metrics, step + 1, prefix="performance") diff --git a/nemo_rl/data/llm_message_utils.py b/nemo_rl/data/llm_message_utils.py index c3e6bf76586..d61ff4b7bda 100644 --- a/nemo_rl/data/llm_message_utils.py +++ b/nemo_rl/data/llm_message_utils.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# 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. @@ -26,9 +26,8 @@ ) 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, load_media_from_message, ) from nemo_rl.distributed.batched_data_dict import BatchedDataDict @@ -115,9 +114,22 @@ def message_log_to_flat_messages( f"tensors for {key=} must have same number of dimensions: {[t.shape for t in result[key]]}" ) from e raise - elif result[key] and isinstance(result[key][0], PackedTensor): + packed_values = [ + value for value in result[key] if isinstance(value, PackedTensor) + ] + if packed_values: + invalid_values = [ + value + for value in result[key] + if value is not None and not isinstance(value, PackedTensor) + ] + if invalid_values: + raise TypeError( + f"Packed multimodal key {key!r} also contains non-packed " + f"values: {[type(value).__name__ for value in invalid_values]}" + ) try: - concat[key] = PackedTensor.concat(result[key]) + concat[key] = PackedTensor.merge_segments(packed_values) except Exception as e: raise RuntimeError( f"Error concatenating packed multimodal data for {key=}" @@ -363,26 +375,31 @@ def batched_message_log_to_flat_message( result = BatchedDataDict() for key in all_keys: values = [seq.get(key) for seq in sequenced_lists] - packed_template = next( - (value for value in values if isinstance(value, PackedTensor)), None - ) - if packed_template is not None: - if any( - value is not None and not isinstance(value, PackedTensor) + packed_values = [value for value in values if isinstance(value, PackedTensor)] + # Preserve one logical row for conversations missing this media key. + # Async replay may concatenate text-only and multimodal prompt groups in + # either order, so the first row cannot determine the value type. + if packed_values: + invalid_values = [ + value for value in values - ): + if value is not None and not isinstance(value, PackedTensor) + ] + if invalid_values: raise TypeError( - f"Expected PackedTensor or None for {key=}, " - f"got {[type(value).__name__ for value in values]}" + f"Packed multimodal key {key!r} also contains non-packed " + f"values: {[type(value).__name__ for value in invalid_values]}" ) - filled_packed_values = cast( - list[PackedTensor], - [ - PackedTensor.empty_like(packed_template) if value is None else value - for value in values - ], - ) - result[key] = PackedTensor.flattened_concat(filled_packed_values) + template = packed_values[0] + aligned_values = [ + ( + value + if isinstance(value, PackedTensor) + else PackedTensor.empty_rows_like(template, 1) + ) + for value in values + ] + result[key] = PackedTensor.flattened_concat(aligned_values) continue if not values or not isinstance(values[0], Tensor): result[key] = values @@ -477,7 +494,6 @@ def get_formatted_message_log( list[dict[str, str]], message_log ) # we just use the str:str parts here - multimodal_keys = get_multimodal_keys_from_processor(tokenizer) multimodal_load_kwargs = get_multimodal_default_settings_from_processor(tokenizer) def _format_content_helper( @@ -642,20 +658,9 @@ def _format_content_helper( ) new_message["token_ids"] = processed_chunk["input_ids"][0] - # add all vlm keys to the message - for key in multimodal_keys: - if key in processed_chunk: - # token_type_ids and mm_token_type_ids are sequence-length tensors - # (one label per token), not visual patch tensors. They must be - # stored as plain tensors and padded like input_ids rather than - # packed as multimodal data. This mirrors processors.py behavior. - if key in ("token_type_ids", "mm_token_type_ids"): - new_message[key] = processed_chunk[key][0] - else: - new_message[key] = PackedTensor( - processed_chunk[key], - dim_to_pack=get_dim_to_pack_along(tokenizer, key), - ) + new_message.update( + extract_multimodal_model_inputs(tokenizer, dict(processed_chunk)) + ) if len(new_message["token_ids"]) == 0: # if there is an empty message, the empty `token_ids` tensor ends up being in fp32, diff --git a/nemo_rl/data/multimodal_utils.py b/nemo_rl/data/multimodal_utils.py index 368608f5909..3528e1be86c 100644 --- a/nemo_rl/data/multimodal_utils.py +++ b/nemo_rl/data/multimodal_utils.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# 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. @@ -16,7 +16,9 @@ import inspect import logging import re +import uuid from collections import defaultdict +from copy import deepcopy from io import BytesIO from typing import Any, Optional, Union @@ -28,6 +30,12 @@ from transformers.audio_utils import load_audio from transformers.video_utils import load_video +VLLM_MULTIMODAL_DATA_KEYS = frozenset({"vllm_images", "vllm_videos", "vllm_audios"}) +NATIVE_MULTIMODAL_KEYS = frozenset({"vllm_content", *VLLM_MULTIMODAL_DATA_KEYS}) +MULTIMODAL_CONTENT_TYPES = frozenset( + {"input_image", "image", "image_url", "video", "audio"} +) + # List of allowed placeholder strings for different media types in the dataset string # e.g. "This is an example of " MEDIA_TAGS = { @@ -88,14 +96,21 @@ def uses_image_placeholder(processor: Any) -> bool: class PackedTensor: - """Wrapper around a list of torch tensors and a dimension along which to pack the tensors. + """A logical batch of rows backed by packable tensor segments. - This class is used to wrap a list of tensors along with a `dim_to_pack` parameter. - It can be used for data that can be packed along different dimensions (such as multimodal data). + The default representation is intentionally the legacy one: every entry in + ``tensors`` is one logical row and no deduplication metadata is allocated. + ``enable_deduplication`` adds stable provenance to the physical segments. + Operations that combine or slice dedup-enabled values then use a CSR-like + logical-row mapping: - `dim_to_pack` is used to specify the dimension along which to pack the tensors. + - ``_row_offsets`` partitions the flattened logical segment references. + - ``_segment_indices`` maps each logical segment reference to ``tensors``. + - ``_segment_provenance`` is stable across deepcopy/pickle and is the only + evidence used to re-intern physical segments. - The list of tensors can be returned as a single packed tensor by calling `as_tensor` which will concatenate the tensors along the `dim_to_pack` dimension. + Prompt identity is deliberately absent: belonging to the same prompt group + makes media a candidate for sharing, but never proves media equality. """ def __init__( @@ -104,6 +119,9 @@ def __init__( dim_to_pack: int, *, pad_to_max_shape: bool = False, + _row_offsets: Optional[list[int]] = None, + _segment_indices: Optional[list[int]] = None, + _segment_provenance: Optional[list[bytes]] = None, ) -> None: """Wrap per-item tensors for concatenation along ``dim_to_pack``. @@ -119,9 +137,10 @@ def __init__( if isinstance(tensors, torch.Tensor): self.tensors: list[Optional[torch.Tensor]] = [tensors] elif isinstance(tensors, list): - assert len(tensors) > 0, ( - "Input tensors to PackedTensor must be a non-empty list" - ) + if not tensors and _row_offsets is None: + raise AssertionError( + "Input tensors to PackedTensor must be a non-empty list" + ) self.tensors: list[Optional[torch.Tensor]] = tensors else: raise ValueError( @@ -129,6 +148,163 @@ def __init__( ) self.dim_to_pack = dim_to_pack self.pad_to_max_shape = pad_to_max_shape + if (_row_offsets is None) != (_segment_indices is None): + raise ValueError( + "_row_offsets and _segment_indices must either both be set or both be None" + ) + if _row_offsets is not None: + if not _row_offsets or _row_offsets[0] != 0: + raise ValueError("_row_offsets must start with 0") + if any( + current > following + for current, following in zip(_row_offsets, _row_offsets[1:]) + ): + raise ValueError("_row_offsets must be non-decreasing") + assert _segment_indices is not None + if _row_offsets[-1] != len(_segment_indices): + raise ValueError( + "_row_offsets must end at the number of logical segment references" + ) + if _segment_indices and ( + min(_segment_indices) < 0 or max(_segment_indices) >= len(self.tensors) + ): + raise ValueError( + "_segment_indices cannot reference an out-of-range physical segment" + ) + if _segment_provenance is not None and len(_segment_provenance) != len( + self.tensors + ): + raise ValueError( + "_segment_provenance must have one entry per physical segment" + ) + self._row_offsets = _row_offsets + self._segment_indices = _segment_indices + self._segment_provenance = _segment_provenance + + def __setstate__(self, state: dict[str, Any]) -> None: + """Restore both current and pre-deduplication pickled instances.""" + self.__dict__.update(state) + self.__dict__.setdefault("_row_offsets", None) + self.__dict__.setdefault("_segment_indices", None) + self.__dict__.setdefault("_segment_provenance", None) + + @property + def deduplication_enabled(self) -> bool: + """Whether this value carries stable physical-segment provenance.""" + return self._segment_provenance is not None + + @property + def logical_segment_count(self) -> int: + """Number of segment occurrences after logical expansion.""" + if self._segment_indices is not None: + return len(self._segment_indices) + return len(self.tensors) + + def logical_segment_counts_by_row(self) -> list[int]: + """Return the number of non-empty media segments in each logical row.""" + if self._row_offsets is None: + return [int(tensor is not None) for tensor in self.tensors] + assert self._segment_indices is not None + return [ + sum( + self.tensors[physical_index] is not None + for physical_index in self._segment_indices[ + self._row_offsets[row] : self._row_offsets[row + 1] + ] + ) + for row in range(len(self)) + ] + + def iter_logical_segments(self): + """Yield physical tensor segments in logical row/segment order.""" + if self._segment_indices is None: + yield from self.tensors + return + for physical_index in self._segment_indices: + yield self.tensors[physical_index] + + def enable_deduplication(self) -> "PackedTensor": + """Assign stable provenance lazily without changing logical contents.""" + if self._segment_provenance is None: + self._segment_provenance = [ + uuid.uuid4().bytes for _ in range(len(self.tensors)) + ] + return self + + def repeat_interleave(self, num_repeats: int) -> "PackedTensor": + """Repeat logical rows while retaining one copy of each physical segment.""" + if not self.deduplication_enabled: + raise ValueError( + "PackedTensor repeat_interleave requires deduplication to be enabled" + ) + if num_repeats < 0: + raise ValueError("num_repeats must be non-negative") + if num_repeats == 0: + return PackedTensor( + [], + self.dim_to_pack, + pad_to_max_shape=self.pad_to_max_shape, + _row_offsets=[0], + _segment_indices=[], + _segment_provenance=[], + ) + segment_indices = [] + row_offsets = [0] + for row in range(len(self)): + row_segments = self._row_segment_indices(row) + for _ in range(num_repeats): + segment_indices.extend(row_segments) + row_offsets.append(len(segment_indices)) + return PackedTensor( + list(self.tensors), + self.dim_to_pack, + pad_to_max_shape=self.pad_to_max_shape, + _row_offsets=row_offsets, + _segment_indices=segment_indices, + _segment_provenance=list(self._segment_provenance or []), + ) + + def _row_segment_indices(self, row: int) -> list[int]: + if self._row_offsets is None: + return [row] + assert self._segment_indices is not None + return self._segment_indices[ + self._row_offsets[row] : self._row_offsets[row + 1] + ] + + def __deepcopy__(self, memo: dict[int, Any]) -> "PackedTensor": + """Share immutable media segments only for an explicitly enabled value.""" + if self._row_offsets is None and not self.deduplication_enabled: + copied = PackedTensor( + [deepcopy(item, memo) for item in self.tensors], + self.dim_to_pack, + pad_to_max_shape=self.pad_to_max_shape, + ) + else: + copied = PackedTensor( + ( + list(self.tensors) + if self.deduplication_enabled + else [deepcopy(item, memo) for item in self.tensors] + ), + self.dim_to_pack, + pad_to_max_shape=self.pad_to_max_shape, + _row_offsets=( + list(self._row_offsets) if self._row_offsets is not None else None + ), + _segment_indices=( + list(self._segment_indices) + if self._segment_indices is not None + else None + ), + _segment_provenance=( + list(self._segment_provenance) + if self._segment_provenance is not None + else None + ), + ) + memo[id(self)] = copied + return copied def as_tensor( self, device: Optional[torch.device] = None @@ -138,7 +314,10 @@ def as_tensor( for i, item in enumerate(self.tensors): if item is not None: self.tensors[i] = item.to(device) - non_none_tensors = [t for t in self.tensors if t is not None] + tensors = self.tensors + if self._segment_indices is not None: + tensors = [self.tensors[index] for index in self._segment_indices] + non_none_tensors = [t for t in tensors if t is not None] if len(non_none_tensors) == 0: return None @@ -188,7 +367,8 @@ def pad_to_batch_shape(tensor: torch.Tensor) -> torch.Tensor: return torch.cat(non_none_tensors, dim=self.dim_to_pack).to(device) def __len__(self) -> int: - # this is the number of tensors in this data wrapper + if self._row_offsets is not None: + return len(self._row_offsets) - 1 return len(self.tensors) def to(self, device: str | torch.device) -> "PackedTensor": @@ -199,18 +379,73 @@ def to(self, device: str | torch.device) -> "PackedTensor": def slice(self, indices: Union[list[int], torch.Tensor]) -> "PackedTensor": idx = indices.tolist() if isinstance(indices, torch.Tensor) else indices - tensors = [self.tensors[i] for i in idx] + if not self.deduplication_enabled and self._row_offsets is None: + tensors = [self.tensors[i] for i in idx] + return PackedTensor( + tensors, + self.dim_to_pack, + pad_to_max_shape=self.pad_to_max_shape, + ) + + physical_remap: dict[int, int] = {} + tensors: list[Optional[torch.Tensor]] = [] + provenances: list[bytes] = [] + segment_indices: list[int] = [] + row_offsets = [0] + for row in idx: + if row < 0: + row += len(self) + if not 0 <= row < len(self): + raise IndexError(f"PackedTensor row index {row} is out of range") + for physical_index in self._row_segment_indices(row): + if physical_index not in physical_remap: + physical_remap[physical_index] = len(tensors) + tensors.append(self.tensors[physical_index]) + if self._segment_provenance is not None: + provenances.append(self._segment_provenance[physical_index]) + segment_indices.append(physical_remap[physical_index]) + row_offsets.append(len(segment_indices)) return PackedTensor( tensors, self.dim_to_pack, pad_to_max_shape=self.pad_to_max_shape, + _row_offsets=row_offsets, + _segment_indices=segment_indices, + _segment_provenance=( + provenances if self._segment_provenance is not None else None + ), ) @classmethod def empty_like(cls, other: "PackedTensor") -> "PackedTensor": - """Return a new PackedTensor with same length and dim_to_pack as `other`, with all entries None.""" + """Return empty logical rows matching ``other``.""" + return cls.empty_rows_like(other, len(other)) + + @classmethod + def empty_rows_like(cls, other: "PackedTensor", num_rows: int) -> "PackedTensor": + """Return ``num_rows`` logical rows containing no media segments.""" + if num_rows < 0: + raise ValueError("num_rows must be non-negative") + if other.deduplication_enabled or other._row_offsets is not None: + return cls( + [], + other.dim_to_pack, + pad_to_max_shape=other.pad_to_max_shape, + _row_offsets=[0] * (num_rows + 1), + _segment_indices=[], + _segment_provenance=[], + ) + if num_rows == 0: + return cls( + [], + other.dim_to_pack, + pad_to_max_shape=other.pad_to_max_shape, + _row_offsets=[0], + _segment_indices=[], + _segment_provenance=None, + ) return cls( - [None] * len(other.tensors), + [None] * num_rows, other.dim_to_pack, pad_to_max_shape=other.pad_to_max_shape, ) @@ -245,7 +480,53 @@ def concat(cls, from_packed_tensors: list["PackedTensor"]) -> "PackedTensor": assert len(set(pad_to_max_shapes)) == 1, ( "All packed tensors must have the same pad_to_max_shape setting" ) - # concatenate the tensors + if any( + packed_tensor.deduplication_enabled + or packed_tensor._row_offsets is not None + for packed_tensor in from_packed_tensors + ): + tensors: list[Optional[torch.Tensor]] = [] + provenances: list[bytes] = [] + provenance_to_physical: dict[bytes, int] = {} + segment_indices: list[int] = [] + row_offsets = [0] + + for packed_tensor in from_packed_tensors: + physical_remap: dict[int, int] = {} + for physical_index, tensor in enumerate(packed_tensor.tensors): + provenance = ( + packed_tensor._segment_provenance[physical_index] + if packed_tensor._segment_provenance is not None + else None + ) + if provenance is not None and provenance in provenance_to_physical: + new_index = provenance_to_physical[provenance] + else: + new_index = len(tensors) + tensors.append(tensor) + if provenance is None: + provenance = uuid.uuid4().bytes + provenances.append(provenance) + provenance_to_physical[provenance] = new_index + physical_remap[physical_index] = new_index + + for row in range(len(packed_tensor)): + segment_indices.extend( + physical_remap[index] + for index in packed_tensor._row_segment_indices(row) + ) + row_offsets.append(len(segment_indices)) + + return cls( + tensors, + dim_to_packs[0], + pad_to_max_shape=pad_to_max_shapes[0], + _row_offsets=row_offsets, + _segment_indices=segment_indices, + _segment_provenance=provenances, + ) + + # Legacy flag-off behavior: concatenate the tensors without metadata. tensors = [] for packed_tensor in from_packed_tensors: tensors.extend(packed_tensor.tensors) @@ -256,6 +537,29 @@ def concat(cls, from_packed_tensors: list["PackedTensor"]) -> "PackedTensor": pad_to_max_shape=pad_to_max_shapes[0], ) + @classmethod + def merge_segments( + cls, from_packed_tensors: list["PackedTensor"] + ) -> "PackedTensor": + """Merge message-turn values into one logical conversation row.""" + if not any( + packed_tensor.deduplication_enabled + or packed_tensor._row_offsets is not None + for packed_tensor in from_packed_tensors + ): + return cls.concat(from_packed_tensors) + + concatenated = cls.concat(from_packed_tensors) + assert concatenated._segment_indices is not None + return cls( + concatenated.tensors, + concatenated.dim_to_pack, + pad_to_max_shape=concatenated.pad_to_max_shape, + _row_offsets=[0, len(concatenated._segment_indices)], + _segment_indices=concatenated._segment_indices, + _segment_provenance=concatenated._segment_provenance, + ) + @classmethod def flattened_concat( cls, from_packed_tensors: list["PackedTensor"] @@ -290,6 +594,12 @@ def flattened_concat( assert len(set(pad_to_max_shapes)) == 1, ( "All packed tensors must have the same pad_to_max_shape setting" ) + if any( + packed_tensor.deduplication_enabled + or packed_tensor._row_offsets is not None + for packed_tensor in from_packed_tensors + ): + return cls.concat(from_packed_tensors) tensors = [p.as_tensor() for p in from_packed_tensors] return cls( tensors, @@ -374,6 +684,95 @@ def get_dim_to_pack_along(processor, key: str) -> int: return 0 +def get_pad_to_max_shape(processor: Any, key: str) -> bool: + """Return whether a processor input must pad non-packing dimensions.""" + return uses_image_placeholder(processor) and key == "pixel_values" + + +def extract_multimodal_model_inputs( + processor: Any, processed: dict[str, Any] +) -> dict[str, PackedTensor | torch.Tensor]: + """Extract packed media inputs and sequence-aligned auxiliary tensors.""" + processed = dict(processed) + if ( + uses_image_placeholder(processor) + and "pixel_values" in processed + and "imgs_sizes" not in processed + and processed["pixel_values"].ndim == 4 + ): + pixel_values = processed["pixel_values"] + num_tiles, _, height, width = pixel_values.shape + processed["imgs_sizes"] = torch.tensor( + [[height, width]] * num_tiles, + dtype=torch.long, + ) + if "imgs_sizes" in processed and "num_frames" not in processed: + processed["num_frames"] = torch.ones( + len(processed["imgs_sizes"]), + dtype=torch.long, + ) + + input_ids = processed.get("input_ids") + if input_ids is None: + raise ValueError("Processor output is missing input_ids.") + if not isinstance(input_ids, torch.Tensor) or input_ids.ndim not in (1, 2): + raise ValueError( + "Processor input_ids must be a one- or two-dimensional torch.Tensor." + ) + if input_ids.ndim == 2 and input_ids.shape[0] != 1: + raise ValueError( + "Multimodal chat processing expects a single conversation, got " + f"input_ids shape {tuple(input_ids.shape)}." + ) + sequence_length = input_ids.shape[-1] + + extracted: dict[str, PackedTensor | torch.Tensor] = {} + multimodal_keys = list(get_multimodal_keys_from_processor(processor)) + for key in ("imgs_sizes", "num_frames"): + if key in processed and key not in multimodal_keys: + multimodal_keys.append(key) + for key in multimodal_keys: + if key not in processed: + continue + value = processed[key] + if not isinstance(value, torch.Tensor): + raise ValueError( + f"Processor model input {key!r} must be a torch.Tensor, got " + f"{type(value).__name__}." + ) + if key == "imgs_sizes": + value = value.to(dtype=torch.int32) + extracted[key] = PackedTensor( + value, + dim_to_pack=get_dim_to_pack_along(processor, key), + pad_to_max_shape=get_pad_to_max_shape(processor, key), + ) + + for key in ("token_type_ids", "mm_token_type_ids"): + if key not in processed: + continue + value = processed[key] + if not isinstance(value, torch.Tensor) or value.ndim not in (1, 2): + raise ValueError( + f"Processor sequence input {key!r} must be a one- or " + "two-dimensional torch.Tensor." + ) + if value.ndim == 2: + if value.shape[0] != 1: + raise ValueError( + f"Processor sequence input {key!r} must contain one " + f"conversation, got shape {tuple(value.shape)}." + ) + value = value[0] + if len(value) != sequence_length: + raise ValueError( + f"Processor sequence input {key!r} has length {len(value)}, " + f"but input_ids has length {sequence_length}." + ) + extracted[key] = value + return extracted + + def resolve_to_image(image_path_or_image: str | Image.Image) -> Image.Image: """Resolve the image path to a PIL.Image object. @@ -422,6 +821,69 @@ def image_to_data_url(image: Image.Image, fmt: str = "PNG") -> str: return f"data:image/{fmt.lower()};base64,{encoded}" +def extract_input_image_sources_from_responses_messages( + messages: Any, +) -> list[str | Image.Image]: + """Extract image sources from Responses-API messages in encounter order.""" + if not isinstance(messages, list): + return [] + + sources: list[str | Image.Image] = [] + for message in messages: + if not isinstance(message, dict): + continue + content = message.get("content") or [] + if not isinstance(content, list): + continue + for part in content: + if not isinstance(part, dict): + continue + if part.get("type") not in ("input_image", "image", "image_url"): + continue + source = part.get("image") or part.get("image_url") or part.get("url") + if isinstance(source, dict): + source = source.get("url") + if isinstance(source, (str, Image.Image)): + sources.append(source) + return sources + + +def extract_input_images_from_responses_messages( + messages: Any, +) -> list[Image.Image]: + """Load images from Responses-API input messages in encounter order.""" + return [ + resolve_to_image(source) + for source in extract_input_image_sources_from_responses_messages(messages) + ] + + +def attach_image_model_inputs_to_message( + message: dict[str, Any], + *, + images: list[Image.Image], + processor: Any, +) -> None: + """Attach processor-owned image tensors without replacing rollout tokens.""" + if not images or processor is None: + return + + image_token = getattr(processor, "image_token", "") + processed = processor( + text=image_token * len(images), + images=images, + return_tensors="pt", + ) + model_inputs = extract_multimodal_model_inputs(processor, dict(processed)) + message.update( + { + key: value + for key, value in model_inputs.items() + if isinstance(value, PackedTensor) + } + ) + + def encode_images_in_examples(nemo_gym_examples: list[dict]) -> list[dict]: """Replace local image paths in NeMo Gym examples with base64 data URLs. diff --git a/nemo_rl/distributed/batched_data_dict.py b/nemo_rl/distributed/batched_data_dict.py index eccdd09ad94..b7a5af0a934 100644 --- a/nemo_rl/distributed/batched_data_dict.py +++ b/nemo_rl/distributed/batched_data_dict.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# 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. @@ -32,6 +32,8 @@ from typing_extensions import Self from nemo_rl.data.multimodal_utils import ( + MULTIMODAL_CONTENT_TYPES, + NATIVE_MULTIMODAL_KEYS, PackedTensor, ) from nemo_rl.data.packing import get_packer @@ -42,6 +44,50 @@ DictT = TypeVar("DictT", bound=Mapping[str, Any]) +_COUPLED_MULTIMODAL_KEYS = ( + ("pixel_values", "image_grid_thw"), + ("pixel_values", "imgs_sizes"), + ("pixel_values", "num_frames"), + ("pixel_values_videos", "video_grid_thw"), +) + + +def _prepare_multimodal_sharing( + value: Any, + *, + media_context: bool = False, +) -> dict[int, Any]: + """Enable PackedTensor provenance and return deepcopy memo entries. + + PackedTensor is an explicit multimodal type. Raw native-vLLM payloads are + shared only under named media keys or typed content parts. Containers are + still deep-copied so rollout rows may diverge safely. + """ + shared_leaves: dict[int, Any] = {} + + def visit(item: Any, in_media_context: bool = False) -> None: + if isinstance(item, PackedTensor): + item.enable_deduplication() + return + if isinstance(item, dict): + content_type = item.get("type") + typed_media = content_type in MULTIMODAL_CONTENT_TYPES + for key, child in item.items(): + visit( + child, + in_media_context or typed_media or key in NATIVE_MULTIMODAL_KEYS, + ) + return + if isinstance(item, (list, tuple)): + for child in item: + visit(child, in_media_context) + return + if in_media_context: + shared_leaves[id(item)] = item + + visit(value, media_context) + return shared_leaves + class SequencePackingArgs(TypedDict): """Configuration settings for sequence packing. @@ -94,6 +140,24 @@ def get_multimodal_dict( self, as_tensors: bool = False, device: Optional[torch.device] = None ) -> dict[str, Any]: """Return a regular dict of tensors or packed multimodal data items.""" + if as_tensors: + for value_key, metadata_key in _COUPLED_MULTIMODAL_KEYS: + value = self.data.get(value_key) + metadata = self.data.get(metadata_key) + if not isinstance(value, PackedTensor) or not isinstance( + metadata, PackedTensor + ): + continue + value_counts = value.logical_segment_counts_by_row() + metadata_counts = metadata.logical_segment_counts_by_row() + if value_counts != metadata_counts: + raise ValueError( + "Coupled multimodal keys must have the same ordered " + f"per-row segment counts, but {value_key!r} has " + f"{value_counts} and {metadata_key!r} has " + f"{metadata_counts}." + ) + multimodal_dict = {} for k, v in self.data.items(): if isinstance(v, PackedTensor): @@ -108,6 +172,8 @@ def from_batches( cls: Type[Self], batches: Sequence[Mapping[Any, Any]], pad_value_dict: Optional[dict[str, int | float]] = None, + *, + allow_missing_packed_tensors: bool = False, ) -> Self: """Given a list of batches, stack the tensors/lists within and put them in a single dictionary. @@ -116,6 +182,9 @@ def from_batches( Args: batches (list[Dict]): A list of dictionaries, each containing a batch of data. pad_value_dict (Optional[dict[str, int]]): An optional dict mapping keys to non-default(0) padding values. + allow_missing_packed_tensors: Represent missing ``PackedTensor`` + media keys as empty logical rows. This is opt-in so ordinary + flag-off concatenation retains its strict key checks. Returns: BatchedDataDict: A new BatchedDataDict containing the stacked data. @@ -128,12 +197,26 @@ def from_batches( def batch_size(item: Mapping[Any, Any]) -> int: if not item: return 0 - value = next(iter(item.values())) - if isinstance(value, PackedTensor): - return len(value) - if isinstance(value, torch.Tensor): - return value.shape[0] - return len(value) + + if not allow_missing_packed_tensors: + # Preserve the legacy shared primitive exactly unless sparse + # PackedTensor normalization was explicitly requested. + return len(next(iter(item.values()))) + + sizes = set() + for value in item.values(): + if isinstance(value, PackedTensor): + sizes.add(len(value)) + elif isinstance(value, torch.Tensor): + sizes.add(value.shape[0]) + else: + sizes.add(len(value)) + if len(sizes) != 1: + raise ValueError( + "Source batch has inconsistent logical row counts: " + f"{sorted(sizes)}." + ) + return next(iter(sizes)) keys = sorted({key for item in batches for key in item}) for k in keys: @@ -143,12 +226,47 @@ def batch_size(item: Mapping[Any, Any]) -> int: if k not in item and batch_size(item) ] if missing_nonempty_batches: - raise KeyError( - f"Key {k!r} is missing from non-empty batches " - f"{missing_nonempty_batches}." - ) + present_values = [item[k] for item in batches if k in item] + if not ( + allow_missing_packed_tensors + and present_values + and all(isinstance(value, PackedTensor) for value in present_values) + ): + raise KeyError( + f"Key {k!r} is missing from non-empty batches " + f"{missing_nonempty_batches}." + ) - list_of_tensors = [item[k] for item in batches if k in item] + template = present_values[0] + assert isinstance(template, PackedTensor) + list_of_tensors = [ + ( + item[k] + if k in item + else PackedTensor.empty_rows_like(template, batch_size(item)) + ) + for item in batches + if k in item or batch_size(item) + ] + else: + list_of_tensors = [item[k] for item in batches if k in item] + + if allow_missing_packed_tensors and isinstance( + list_of_tensors[0], PackedTensor + ): + source_batches = [ + item for item in batches if k in item or batch_size(item) + ] + for batch_index, (item, packed_tensor) in enumerate( + zip(source_batches, list_of_tensors) + ): + expected_rows = batch_size(item) + if len(packed_tensor) != expected_rows: + raise ValueError( + f"PackedTensor key {k!r} has {len(packed_tensor)} " + f"logical rows in source batch {batch_index}, " + f"expected {expected_rows}." + ) if isinstance(list_of_tensors[0], list): tensor_or_list: list[Any] | torch.Tensor = [ @@ -622,7 +740,7 @@ def _get_padded_seqlen(seqlen: int) -> int: aggregated_shards[shard_idx][k] = ( PackedTensor.concat(packed_slices) if packed_slices - else PackedTensor.empty_like(v) + else PackedTensor.empty_rows_like(v, 0) ) else: shard_values = [] @@ -767,12 +885,21 @@ def slice(self, start: int, end: int) -> "SlicedDataDict": sliced_batch[k] = self.data[k][start:end] return sliced_batch - def repeat_interleave(self, num_repeats: int) -> Self: + def repeat_interleave( + self, + num_repeats: int, + *, + share_immutable_media: bool = False, + ) -> Self: """Repeats the batch num_repeats times. For each element in the batch, repeat each value num_repeats times. i.e: {"key": torch.tensor([1, 2, 3]), "other_key": [1, 2, 3]} -> {"key": torch.tensor([1, 1, 2, 2, 3, 3]), "other_key": [1, 1, 2, 2, 3, 3]} + + When ``share_immutable_media`` is enabled, only explicit multimodal + leaves share storage. Every surrounding row/message/content container + remains independent. """ repeated_batch: Self = type(self)() for k, v in self.data.items(): @@ -780,14 +907,30 @@ def repeat_interleave(self, num_repeats: int) -> Self: # For tensors, use repeat_interleave to repeat each element repeated_batch[k] = v.repeat_interleave(num_repeats, dim=0) elif isinstance(v, PackedTensor): - raise NotImplementedError( - "PackedTensor does not currently support repeat_interleave" + if not share_immutable_media: + raise NotImplementedError( + "PackedTensor does not currently support repeat_interleave " + "unless share_immutable_media is enabled" + ) + repeated_batch[k] = v.enable_deduplication().repeat_interleave( + num_repeats ) else: # For lists or other sequences, use a list comprehension to repeat each element - repeated_batch[k] = [ - deepcopy(item) for item in v for _ in range(num_repeats) - ] + repeated_items = [] + for item in v: + shared_leaves = ( + _prepare_multimodal_sharing( + item, + media_context=k in NATIVE_MULTIMODAL_KEYS, + ) + if share_immutable_media + else {} + ) + repeated_items.extend( + deepcopy(item, dict(shared_leaves)) for _ in range(num_repeats) + ) + repeated_batch[k] = repeated_items return repeated_batch def truncate_tensors(self, dim: int, truncated_len: int): @@ -861,9 +1004,9 @@ def size(self) -> int: """Get the batch size of the batch.""" # Get the first key and use its size as the batch size # This assumes all keys have the same batch size - key = next(iter(self.data)) if not self.data: return 0 + key = next(iter(self.data)) if not torch.is_tensor(self.data[key]): return len(self.data[key]) return self.data[key].shape[0] # type: ignore # it's a tensor here diff --git a/nemo_rl/environments/nemo_gym.py b/nemo_rl/environments/nemo_gym.py index 375d3b2d41d..458ebc3b72e 100644 --- a/nemo_rl/environments/nemo_gym.py +++ b/nemo_rl/environments/nemo_gym.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# 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. @@ -17,6 +17,7 @@ import sys from collections import Counter from collections.abc import AsyncGenerator +from copy import deepcopy from pathlib import Path from typing import Any, Dict, List, NotRequired, Optional, TypedDict @@ -27,10 +28,9 @@ from transformers import PreTrainedTokenizerBase from nemo_rl.data.multimodal_utils import ( - PackedTensor, + attach_image_model_inputs_to_message, encode_images_in_examples, - get_dim_to_pack_along, - get_multimodal_keys_from_processor, + extract_input_image_sources_from_responses_messages, resolve_to_image, uses_image_placeholder, ) @@ -284,6 +284,48 @@ def _index_per_turn_images( return per_turn +def _image_sources_equal(left: Any, right: Any) -> bool: + return ( + left == right + if isinstance(left, str) and isinstance(right, str) + else left is right + ) + + +def _without_initial_image_sources( + messages: Any, initial_sources: list[Any] +) -> tuple[Any, bool]: + """Copy Responses messages and remove one ordered copy of initial images.""" + if not isinstance(messages, list): + return messages, False + + filtered = deepcopy(messages) + remaining_sources = list(initial_sources) + for message in filtered: + if not isinstance(message, dict): + continue + content = message.get("content") + if not isinstance(content, list): + continue + + filtered_content = [] + for part in content: + part_sources = extract_input_image_sources_from_responses_messages( + [{"content": [part]}] + ) + if ( + remaining_sources + and len(part_sources) == 1 + and _image_sources_equal(part_sources[0], remaining_sources[0]) + ): + remaining_sources.pop(0) + continue + filtered_content.append(part) + message["content"] = filtered_content + + return filtered, not remaining_sources + + def _attach_multimodal_data_to_user_message( user_message: dict, *, @@ -300,52 +342,11 @@ def _attach_multimodal_data_to_user_message( already contains expanded ``...*N...`` regions, and the processor would try to re-expand every embedded ````. """ - if not images or processor is None: - return - image_token = getattr(processor, "image_token", "") - processed = processor( - text=image_token * len(images), + attach_image_model_inputs_to_message( + user_message, images=images, - return_tensors="pt", + processor=processor, ) - uses_placeholder = uses_image_placeholder(processor) - multimodal_keys = list(get_multimodal_keys_from_processor(processor)) - # Historical checkpoints may emit dynamic image tiles without imgs_sizes. - # Mirror the media-metadata handling in vlm_hf_data_processor. - if ( - uses_placeholder - and "pixel_values" in processed - and "imgs_sizes" not in processed - and processed["pixel_values"].ndim == 4 - ): - pixel_values = processed["pixel_values"] - num_tiles, _, height, width = pixel_values.shape - processed["imgs_sizes"] = torch.tensor( - [[height, width]] * num_tiles, dtype=torch.long - ) - - # imgs_sizes / num_frames are not always declared in model_input_names by - # bundled image processors. RADIO uses temporal patching even for still - # images and requires one num_frames=1 entry per image/tile. - if "imgs_sizes" in processed and "imgs_sizes" not in multimodal_keys: - multimodal_keys.append("imgs_sizes") - if "imgs_sizes" in processed and "num_frames" not in processed: - processed["num_frames"] = torch.ones( - len(processed["imgs_sizes"]), dtype=torch.long - ) - if "num_frames" in processed and "num_frames" not in multimodal_keys: - multimodal_keys.append("num_frames") - for key in multimodal_keys: - if key not in processed: - continue - value = processed[key] - if key == "imgs_sizes": - value = value.to(dtype=torch.int32) - user_message[key] = PackedTensor( - value, - dim_to_pack=get_dim_to_pack_along(processor, key), - pad_to_max_shape=uses_placeholder and key == "pixel_values", - ) @ray.remote(max_restarts=-1, max_task_retries=-1) # pragma: no cover @@ -473,6 +474,7 @@ async def run_rollouts( nemo_gym_examples: list[dict], tokenizer: PreTrainedTokenizerBase, timer_prefix: str, + deduplicate_multimodal_data: bool = False, ) -> AsyncGenerator[tuple[int, dict, dict | None], None]: """Stream postprocessed rollouts as NeMo-Gym tasks complete.""" if not nemo_gym_examples: @@ -511,7 +513,10 @@ async def run_rollouts( with timer.time(label=f"{timer_prefix}/postprocess_results"): nemo_rl_result = self._postprocess_nemo_gym_to_nemo_rl_result( - nemo_gym_result, tokenizer + nemo_gym_row, + nemo_gym_result, + tokenizer, + include_initial_multimodal_data=not deduplicate_multimodal_data, ) if _has_nan_generation_logprobs(nemo_rl_result): raise RuntimeError("Generation logprobs contain NaN") @@ -548,20 +553,73 @@ async def run_rollouts( def _postprocess_nemo_gym_to_nemo_rl_result( self, + nemo_gym_row: dict, nemo_gym_result: dict, tokenizer: PreTrainedTokenizerBase, + *, + include_initial_multimodal_data: bool = True, ) -> dict: assert isinstance(nemo_gym_result, dict), ( f"Hit a non-successful response when querying NeMo Gym for rollouts: {nemo_gym_result}" ) processor = getattr(self, "_processor", None) + response = nemo_gym_result["response"] + result_input = nemo_gym_result["responses_create_params"].get("input", []) + request_input = nemo_gym_row.get("responses_create_params", {}).get("input") + raw_input = ( + request_input + if isinstance(request_input, list) and request_input + else result_input + ) + initial_input = response.get("agent_input") + if not isinstance(initial_input, list) or not initial_input: + initial_input = raw_input + + seed_obs = response.get("seed_obs") + media_messages = ( + seed_obs if isinstance(seed_obs, list) and seed_obs else initial_input + ) + raw_initial_sources = extract_input_image_sources_from_responses_messages( + raw_input + ) + agent_initial_sources = extract_input_image_sources_from_responses_messages( + initial_input + ) + returned_media_sources = extract_input_image_sources_from_responses_messages( + media_messages + ) + initial_media_matches_raw_input = ( + bool(raw_initial_sources) + and len(agent_initial_sources) == len(raw_initial_sources) + and all( + _image_sources_equal(agent_source, raw_source) + for agent_source, raw_source in zip( + agent_initial_sources, raw_initial_sources + ) + ) + ) + returned_media_matches_raw_input = len(returned_media_sources) == len( + raw_initial_sources + ) and all( + _image_sources_equal(returned_source, raw_source) + for returned_source, raw_source in zip( + returned_media_sources, raw_initial_sources + ) + ) + initial_multimodal_data_omitted = ( + not include_initial_multimodal_data + and initial_media_matches_raw_input + and returned_media_matches_raw_input + ) + if initial_multimodal_data_omitted: + media_messages, _ = _without_initial_image_sources( + media_messages, raw_initial_sources + ) per_turn_images = ( _index_per_turn_images( - nemo_gym_result["response"]["output"], - input_messages=nemo_gym_result.get("responses_create_params", {}).get( - "input" - ), + response["output"], + input_messages=media_messages, ) if processor is not None else [] @@ -736,11 +794,26 @@ def _postprocess_nemo_gym_to_nemo_rl_result( f" → If (2): inspect why no assistant content was produced for this rollout." ) - return { + if initial_multimodal_data_omitted: + for container, key in ( + (nemo_gym_result["responses_create_params"], "input"), + (response, "agent_input"), + (response, "seed_obs"), + ): + if key in container: + container[key], _ = _without_initial_image_sources( + container[key], raw_initial_sources + ) + nemo_gym_result["_nemo_rl_initial_media_omitted"] = True + + result = { "message_log": nemo_rl_message_log, "input_message_log": nemo_rl_message_log[:1], "full_result": nemo_gym_result, } + if not include_initial_multimodal_data: + result["_initial_multimodal_data_omitted"] = initial_multimodal_data_omitted + return result def shutdown(self) -> None: self.rh.shutdown() diff --git a/nemo_rl/experience/rollouts.py b/nemo_rl/experience/rollouts.py index 1a6145160b8..60d150fa799 100644 --- a/nemo_rl/experience/rollouts.py +++ b/nemo_rl/experience/rollouts.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# 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. @@ -36,11 +36,19 @@ DatumSpec, FlatMessagesType, LLMMessageLogType, + VLMMessageLogType, ) from nemo_rl.data.llm_message_utils import ( batched_message_log_to_flat_message, get_keys_from_message_log, ) +from nemo_rl.data.multimodal_utils import ( + NATIVE_MULTIMODAL_KEYS, + VLLM_MULTIMODAL_DATA_KEYS, + PackedTensor, + attach_image_model_inputs_to_message, + extract_input_images_from_responses_messages, +) from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.environments.interfaces import ( EnvironmentInterface, @@ -56,11 +64,147 @@ GenerationInterface, GenerationOutputSpec, ) +from nemo_rl.utils.multimodal_payload_metrics import ( + collect_multimodal_payload_metrics, + print_multimodal_payload_metrics, +) from nemo_rl.utils.timer import Timer TokenizerType = PreTrainedTokenizerBase +def _set_untyped_message_field(message: Any, key: str, value: Any) -> None: + """Set a restored extension field outside the typed core message schema.""" + message[key] = value + + +def attach_initial_nemo_gym_image_payloads( + batch: BatchedDataDict[DatumSpec], + processor: Any, +) -> None: + """Attach initial Gym image tensors once, before prompt repeat. + + The NeMo Gym dataset deliberately carries only the Responses request in + ``extra_env_info``. Dedup-enabled GRPO calls this helper on the unrepeated + prompt batch, allowing ``repeat_interleave(..., share_immutable_media=True)`` + to retain one physical processor output per prompt. Flag-off runs never call + this helper. + """ + for message_log, extra_env_info in zip( + batch["message_log"], batch["extra_env_info"] + ): + if extra_env_info is None or not isinstance(extra_env_info, dict): + continue + initial_messages = extra_env_info.get("responses_create_params", {}).get( + "input", [] + ) + images = extract_input_images_from_responses_messages(initial_messages) + if not images: + continue + if processor is None or getattr(processor, "image_processor", None) is None: + raise ValueError( + "NeMo Gym image deduplication requires the multimodal processor " + "to be passed to GRPO." + ) + user_message = next( + (message for message in message_log if message.get("role") == "user"), + None, + ) + if user_message is None: + raise ValueError("NeMo Gym image prompt has no user message to attach to.") + if isinstance(user_message.get("pixel_values"), PackedTensor): + continue + attach_image_model_inputs_to_message( + user_message, + images=images, + processor=processor, + ) + + +def _add_multimodal_generation_payload( + generation_input_data: BatchedDataDict[GenerationDatumSpec], + flat_messages: BatchedDataDict[FlatMessagesType], + active_batch: BatchedDataDict[DatumSpec], + policy_generation: GenerationInterface, + *, + deduplicate_multimodal_data: bool, +) -> None: + """Attach one policy-ready or native-vLLM media representation. + + The compact policy representation remains in ``message_log`` for later + logprob/training construction. When every active row has a native vLLM + prompt, sending that representation as well is redundant. + """ + generation_config = getattr(policy_generation, "cfg", {}) + native_content = active_batch.get("vllm_content") + + def row_has_formatter_consumed_media(row_index: int) -> bool: + for key in VLLM_MULTIMODAL_DATA_KEYS: + rows = active_batch.get(key) + if rows is None or row_index >= len(rows): + continue + value = rows[row_index] + if value is None: + continue + if isinstance(value, (list, tuple, dict, str, bytes)): + if len(value) > 0: + return True + else: + return True + return False + + use_native_vllm_only = ( + deduplicate_multimodal_data + and generation_config.get("backend") == "vllm" + and native_content is not None + and all( + row_has_formatter_consumed_media(row_index) + for row_index in range(len(native_content)) + ) + ) + if not use_native_vllm_only: + generation_input_data.update( + flat_messages.get_multimodal_dict(as_tensors=False) + ) + + for key in NATIVE_MULTIMODAL_KEYS: + if key in active_batch: + generation_input_data[key] = active_batch[key] + + +def _reattach_original_multimodal_payloads( + results: list[dict[str, Any]], + original_message_logs: list[LLMMessageLogType | VLMMessageLogType], +) -> None: + """Restore exact prompt media omitted by a remote Gym rollout. + + User turns are matched by their ordinal position. Only explicit + ``PackedTensor`` values and named native-generation media are restored, so + arbitrary non-text metadata is never misclassified as media. Newly returned + Gym media is left untouched unless it occupies the corresponding original + prompt key. + """ + for result, original_log in zip(results, original_message_logs): + if not result.pop("_initial_multimodal_data_omitted", False): + continue + original_user_messages = [ + message for message in original_log if message.get("role") == "user" + ] + for log_key in ("input_message_log", "message_log"): + target_log = result.get(log_key) + if not target_log: + continue + target_user_messages = [ + message for message in target_log if message.get("role") == "user" + ] + for original, target in zip(original_user_messages, target_user_messages): + for key, value in original.items(): + if isinstance(value, PackedTensor): + _set_untyped_message_field(target, key, value) + elif key in NATIVE_MULTIMODAL_KEYS: + _set_untyped_message_field(target, key, value) + + def _add_r3_fallback_metrics( gen_metrics: dict[str, float | int], generation_outputs: BatchedDataDict, @@ -670,6 +814,7 @@ def run_multi_turn_rollout( max_seq_len: int, max_rollout_turns: int = 999999, greedy: bool = False, + deduplicate_multimodal_data: bool = False, ) -> tuple[BatchedDataDict[DatumSpec], dict[str, Any]]: """Runs a multi-turn rollout loop, interacting with the environment. @@ -681,6 +826,9 @@ def run_multi_turn_rollout( max_rollout_turns: Maximum number of agent-environment interaction turns. max_seq_len: Maximum sequence length allowed. greedy: Whether to use greedy decoding. + deduplicate_multimodal_data: Send only native media through the vLLM + generation boundary while retaining compact policy media for + logprob and training. Returns: Tuple containing: @@ -719,6 +867,8 @@ def run_multi_turn_rollout( # Convert LLMMessageLogType to FlatMessagesType for generation active_batch = current_batch.select_indices(active_indices) + if turn > 0 and "vllm_content" in active_batch: + active_batch["vllm_content"] = [None] * len(active_indices) active_stop_strings = [current_stop_strings[i] for i in active_indices.tolist()] active_flat_messages: BatchedDataDict[FlatMessagesType] @@ -740,19 +890,13 @@ def run_multi_turn_rollout( "stop_strings": active_stop_strings, } ) - # add the multimodal data to the generation input data - multimodal_data = active_flat_messages.get_multimodal_dict(as_tensors=False) - generation_input_data.update(multimodal_data) - - # keep message log for generation - if "vllm_content" in active_batch: - generation_input_data["vllm_content"] = active_batch["vllm_content"] - if "vllm_images" in active_batch: - generation_input_data["vllm_images"] = active_batch["vllm_images"] - if "vllm_videos" in active_batch: - generation_input_data["vllm_videos"] = active_batch["vllm_videos"] - if "vllm_audios" in active_batch: - generation_input_data["vllm_audios"] = active_batch["vllm_audios"] + _add_multimodal_generation_payload( + generation_input_data, + active_flat_messages, + active_batch, + policy_generation, + deduplicate_multimodal_data=deduplicate_multimodal_data, + ) # generate_responses updates active_batch["message_log"] in-place active_batch, generated_ids, gen_metrics = generate_responses( @@ -919,6 +1063,9 @@ async def async_generate_response_for_sample_turn( tokenizer: TokenizerType, max_seq_len: int, greedy: bool = False, + *, + sample_multimodal_data: dict[str, Any] | None = None, + deduplicate_multimodal_data: bool = False, ) -> tuple[list[dict], torch.Tensor, torch.Tensor, dict[str, float]]: """Generate a response for a single sample's turn using async generation. @@ -929,6 +1076,9 @@ async def async_generate_response_for_sample_turn( tokenizer: Tokenizer to use max_seq_len: Maximum sequence length greedy: Whether to use greedy decoding + sample_multimodal_data: Native vLLM media fields for this sample. + deduplicate_multimodal_data: Avoid sending both native and policy-ready + media through the async generation boundary. Returns: Tuple of (updated_message_log, generated_tokens, input_lengths, generation_metrics) @@ -960,6 +1110,15 @@ async def async_generate_response_for_sample_turn( "stop_strings": [sample_stop_strings], } ) + for key, value in (sample_multimodal_data or {}).items(): + dummy_batch[key] = [value] + _add_multimodal_generation_payload( + generation_input_data, + flat_messages, + dummy_batch, + policy_generation, + deduplicate_multimodal_data=deduplicate_multimodal_data, + ) # Generate response using the async version updated_batch, generated_ids, gen_metrics = await generate_responses_async( @@ -988,6 +1147,7 @@ async def run_sample_multi_turn_rollout( max_seq_len: int, max_rollout_turns: int = 999999, greedy: bool = False, + deduplicate_multimodal_data: bool = False, ) -> tuple[dict, dict[str, Any]]: """Run a multi-turn rollout for a single sample. @@ -1003,6 +1163,8 @@ async def run_sample_multi_turn_rollout( max_seq_len: Maximum sequence length max_rollout_turns: Maximum number of turns greedy: Whether to use greedy decoding + deduplicate_multimodal_data: Avoid redundant media at generation + boundaries while preserving compact policy media in the trajectory. Returns: Tuple of (final_sample_state, sample_metrics) @@ -1012,6 +1174,11 @@ async def run_sample_multi_turn_rollout( current_extra_env_info = copy.deepcopy(initial_sample_state["extra_env_info"]) current_stop_strings = initial_sample_state.get("stop_strings", None) task_name = initial_sample_state["task_name"] + sample_multimodal_data = { + key: initial_sample_state[key] + for key in NATIVE_MULTIMODAL_KEYS + if key in initial_sample_state + } # Sample-level metrics total_reward = 0.0 @@ -1040,6 +1207,11 @@ async def run_sample_multi_turn_rollout( # Generate response for this sample using async generation try: + turn_multimodal_data = sample_multimodal_data + if turn > 0 and "vllm_content" in sample_multimodal_data: + turn_multimodal_data = dict(sample_multimodal_data) + turn_multimodal_data["vllm_content"] = None + ( updated_message_log, generated_tokens, @@ -1052,6 +1224,8 @@ async def run_sample_multi_turn_rollout( tokenizer, max_seq_len, greedy=greedy, + sample_multimodal_data=turn_multimodal_data, + deduplicate_multimodal_data=deduplicate_multimodal_data, ) current_message_log = updated_message_log @@ -1288,21 +1462,24 @@ async def _run_multi_turn_rollout_async( max_seq_len: int, max_rollout_turns: int = 999999, greedy: bool = False, + deduplicate_multimodal_data: bool = False, ) -> tuple[BatchedDataDict[DatumSpec], list[dict[str, Any]]]: """Run one native rollout batch and retain metrics at sample granularity.""" batch_size = len(input_batch["message_log"]) sample_initial_states = [] for i in range(batch_size): - sample_initial_states.append( - { - "message_log": input_batch["message_log"][i], - "extra_env_info": input_batch["extra_env_info"][i], - "task_name": input_batch["task_name"][i], - "stop_strings": input_batch.get("stop_strings", [None] * batch_size)[i], - "idx": input_batch.get("idx", list(range(batch_size)))[i], - } - ) + sample_state = { + "message_log": input_batch["message_log"][i], + "extra_env_info": input_batch["extra_env_info"][i], + "task_name": input_batch["task_name"][i], + "stop_strings": input_batch.get("stop_strings", [None] * batch_size)[i], + "idx": input_batch.get("idx", list(range(batch_size)))[i], + } + for key in NATIVE_MULTIMODAL_KEYS: + if key in input_batch: + sample_state[key] = input_batch[key][i] + sample_initial_states.append(sample_state) async def run_single_sample_with_error_handling(i, sample_state): try: @@ -1315,6 +1492,7 @@ async def run_single_sample_with_error_handling(i, sample_state): max_seq_len=max_seq_len, max_rollout_turns=max_rollout_turns, greedy=greedy, + deduplicate_multimodal_data=deduplicate_multimodal_data, ) except Exception as error: raise RuntimeError(f"Error in sample {i} rollout: {error}") from error @@ -1381,6 +1559,7 @@ def run_async_multi_turn_rollout( max_seq_len: int, max_rollout_turns: int = 999999, greedy: bool = False, + deduplicate_multimodal_data: bool = False, ) -> tuple[BatchedDataDict[DatumSpec], dict[str, Any]]: """Run a complete native rollout batch from a synchronous call site. @@ -1413,6 +1592,7 @@ def run_async_multi_turn_rollout( max_seq_len=max_seq_len, max_rollout_turns=max_rollout_turns, greedy=greedy, + deduplicate_multimodal_data=deduplicate_multimodal_data, ) ) return final_batch, _aggregate_multi_turn_rollout_metrics(sample_metrics) @@ -1427,6 +1607,7 @@ async def run_async_multi_turn_rollout_groups( num_generations: int, max_rollout_turns: int = 999999, greedy: bool = False, + deduplicate_multimodal_data: bool = False, ) -> AsyncGenerator[RolloutGroupResult, None]: """Run one native batch, then yield prompt groups with group-local metrics. @@ -1469,6 +1650,7 @@ async def run_async_multi_turn_rollout_groups( max_seq_len=max_seq_len, max_rollout_turns=max_rollout_turns, greedy=greedy, + deduplicate_multimodal_data=deduplicate_multimodal_data, ) for group_index, start in enumerate(range(0, final_batch.size, num_generations)): end = start + num_generations @@ -2059,6 +2241,8 @@ async def run_async_nemo_gym_rollout( thinking_tags: list[str] | tuple[str, ...] | None = None, mask_env_flagged_samples: bool = True, returns_entire_batch: bool = False, + deduplicate_multimodal_data: bool = False, + debug_payload_metrics: bool = False, ) -> AsyncGenerator[NemoGymRolloutResult, None]: """Stream complete NeMo-Gym prompt groups in group-completion order. @@ -2089,6 +2273,10 @@ async def run_async_nemo_gym_rollout( returns_entire_batch: Whether to treat the input as one potentially heterogeneous group. This requires ``num_generations`` to equal the batch size and is used by synchronous callers. + deduplicate_multimodal_data: Omit initial policy-ready media from the + remote Gym return and restore the exact original payload locally. + debug_payload_metrics: Emit logical, physical, and serialized media + payload metrics at the Gym Ray boundary. Yields: ``NemoGymRolloutResult`` objects in prompt-group completion order. Rows @@ -2172,9 +2360,22 @@ async def run_async_nemo_gym_rollout( actor_timing_metrics: dict[str, Any] = {} nemo_gym_environment = task_to_env["nemo_gym"] with timer.time(run_rollouts_timer_label): + ray_arguments = ( + nemo_gym_rows, + tokenizer, + timer_prefix, + deduplicate_multimodal_data, + ) + print_multimodal_payload_metrics( + collect_multimodal_payload_metrics( + ray_arguments, + "nemo_gym_request", + enabled=debug_payload_metrics, + ) + ) rollout_gen = nemo_gym_environment.run_rollouts.options( num_returns="streaming" - ).remote(nemo_gym_rows, tokenizer, timer_prefix) + ).remote(*ray_arguments) rollout_iterator = rollout_gen.__aiter__() while True: @@ -2188,6 +2389,17 @@ async def run_async_nemo_gym_rollout( stream_finished = True else: rowidx, result, timing_metrics = await future + # Measure the received streaming Ray value in the caller. In + # async training this runs in the collector actor; validation + # runs in the driver, so the two phases cannot share a metric + # accumulator even when they share the NeMo-Gym actor. + print_multimodal_payload_metrics( + collect_multimodal_payload_metrics( + (rowidx, result, timing_metrics), + "nemo_gym_return", + enabled=debug_payload_metrics, + ) + ) if not stream_finished: if timing_metrics is not None: @@ -2196,16 +2408,22 @@ async def run_async_nemo_gym_rollout( _tensorize_nemo_gym_result(result) completed_group = accumulator.add(rowidx, result) if completed_group is not None: + group_input_batch = input_batch.slice( + completed_group.group_index * num_generations, + (completed_group.group_index + 1) * num_generations, + ) + if deduplicate_multimodal_data: + _reattach_original_multimodal_payloads( + completed_group.results, + group_input_batch["message_log"], + ) rollout_result = _postprocess_single_nemo_gym_group( nemo_gym_rows=completed_group.rows, results=completed_group.results, timer=timer, timer_prefix=timer_prefix, policy_generation=policy_generation, - input_batch=input_batch.slice( - completed_group.group_index * num_generations, - (completed_group.group_index + 1) * num_generations, - ), + input_batch=group_input_batch, tokenizer=tokenizer, log_full_result_tables=log_full_result_tables, effort_config=effort_config, @@ -2249,6 +2467,8 @@ def run_nemo_gym_rollout_sync( reward_penalty_config: dict[str, Any] | BaseModel | None = None, thinking_tags: list[str] | tuple[str, ...] | None = None, mask_env_flagged_samples: bool = True, + deduplicate_multimodal_data: bool = False, + debug_payload_metrics: bool = False, ) -> NemoGymRolloutResult: """Run and return one complete NeMo-Gym batch synchronously. @@ -2273,6 +2493,9 @@ def run_nemo_gym_rollout_sync( thinking_tags: Optional opening and closing tags used by thinking penalties. mask_env_flagged_samples: Whether to carry env-driven ``mask_sample`` flags in the rollout batch for loss masking. + deduplicate_multimodal_data: Omit initial policy-ready media from the + remote Gym return and restore it from the input batch. + debug_payload_metrics: Emit exact Gym Ray-boundary media payload metrics. Returns: The fully postprocessed NeMo-Gym rollout batch in input-row order. @@ -2304,6 +2527,8 @@ async def _consume_rollout() -> NemoGymRolloutResult: thinking_tags=thinking_tags, mask_env_flagged_samples=mask_env_flagged_samples, returns_entire_batch=True, + deduplicate_multimodal_data=deduplicate_multimodal_data, + debug_payload_metrics=debug_payload_metrics, ): pass if rollout_result is None: diff --git a/nemo_rl/experience/sync_rollout_actor.py b/nemo_rl/experience/sync_rollout_actor.py index a29ceb9696a..1bc5c6c391e 100644 --- a/nemo_rl/experience/sync_rollout_actor.py +++ b/nemo_rl/experience/sync_rollout_actor.py @@ -263,6 +263,8 @@ def rollout_to_tq( else None, reward_penalty_config=cfg.reward_penalties, thinking_tags=get_nemo_gym_thinking_tags(cfg.env), + deduplicate_multimodal_data=cfg.grpo.deduplicate_multimodal_data, + debug_payload_metrics=cfg.grpo.debug_payload_metrics, ) final_batch, rollout_metrics = r.final_batch, r.rollout_metrics else: @@ -275,6 +277,7 @@ def rollout_to_tq( **common, max_seq_len=cfg.policy["max_total_sequence_length"], max_rollout_turns=cfg.grpo.max_rollout_turns, + deduplicate_multimodal_data=cfg.grpo.deduplicate_multimodal_data, ) fb = final_batch.to("cpu") del final_batch diff --git a/nemo_rl/models/generation/interfaces.py b/nemo_rl/models/generation/interfaces.py index a757ad3a8ff..41548a4b8cb 100644 --- a/nemo_rl/models/generation/interfaces.py +++ b/nemo_rl/models/generation/interfaces.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# 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. @@ -213,6 +213,8 @@ class GenerationConfig(TypedDict): _pad_token_id: NotRequired[int] # MTP draft weights arrive via refit if the trainer trains the MTP layer. _mtp_weights_from_refit: NotRequired[bool] + # Internal debug-only measurement of exact Ray generation arguments. + debug_payload_metrics: NotRequired[bool] class GenerationDatumSpec(TypedDict): diff --git a/nemo_rl/models/generation/vllm/utils.py b/nemo_rl/models/generation/vllm/utils.py index b8ec3dee050..e4f9a78ab43 100644 --- a/nemo_rl/models/generation/vllm/utils.py +++ b/nemo_rl/models/generation/vllm/utils.py @@ -91,39 +91,38 @@ def _get_regular_prompt(index: int): token_ids = valid_ids.tolist() return {"prompt_token_ids": token_ids} - # Check if this is VLM generation by looking for message_log with images - # Support for videos/audio/etc. can be added here - # if 'message_log' in data and any('images' in msg for msg in data['message_log']): + def _get_multi_modal_data(index: int) -> dict[str, Any]: + multi_modal_data = {} + images = data.get("vllm_images", None) + if images is not None and len(images[index]) > 0: + multi_modal_data["image"] = ( + images[index][0] if len(images[index]) == 1 else images[index] + ) + audios = data.get("vllm_audios", None) + if audios is not None and len(audios[index]) > 0: + multi_modal_data["audio"] = ( + audios[index][0] if len(audios[index]) == 1 else audios[index] + ) + videos = data.get("vllm_videos", None) + if videos is not None and len(videos[index]) > 0: + multi_modal_data["video"] = ( + videos[index][0] if len(videos[index]) == 1 else videos[index] + ) + return multi_modal_data + + # Native image, audio, and video side channels share this formatter path. if "vllm_content" in data: # VLM generation using content and multi_modal_data for i in range(start_idx, end_idx): msg = data["vllm_content"][i] - # if msg is None, this conversation had no multimodal content, fallback to regular prompt - if msg is None: - prompts.append(_get_regular_prompt(i)) - continue - # init prompt dict - prompt_dict = {"prompt": msg} - # collect multi_modal_data from images, audios, and videos - multi_modal_data = {} - images = data.get("vllm_images", None) - if images is not None and len(images[i]) > 0: - multi_modal_data["image"] = ( - images[i][0] if len(images[i]) == 1 else images[i] - ) - audios = data.get("vllm_audios", None) - if audios is not None and len(audios[i]) > 0: - multi_modal_data["audio"] = ( - audios[i][0] if len(audios[i]) == 1 else audios[i] - ) - videos = data.get("vllm_videos", None) - if videos is not None and len(videos[i]) > 0: - multi_modal_data["video"] = ( - videos[i][0] if len(videos[i]) == 1 else videos[i] - ) + multi_modal_data = _get_multi_modal_data(i) if not multi_modal_data: prompts.append(_get_regular_prompt(i)) continue + # Raw processor content is valid only for the initial turn. Later + # turns use the updated pre-tokenized conversation plus the same + # native media, preventing vLLM from regenerating the stale prompt. + prompt_dict = {"prompt": msg} if msg is not None else _get_regular_prompt(i) prompt_dict["multi_modal_data"] = multi_modal_data prompts.append(prompt_dict) else: diff --git a/nemo_rl/models/generation/vllm/vllm_generation.py b/nemo_rl/models/generation/vllm/vllm_generation.py index 6d0d8d37143..da0e10de45e 100644 --- a/nemo_rl/models/generation/vllm/vllm_generation.py +++ b/nemo_rl/models/generation/vllm/vllm_generation.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# 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. @@ -43,6 +43,11 @@ compute_spec_decode_metrics, resolve_generation_worker_cls, ) +from nemo_rl.utils.multimodal_payload_metrics import ( + collect_multimodal_payload_metrics, + collect_sharded_multimodal_payload_metrics, + print_multimodal_payload_metrics, +) from nemo_rl.weight_sync.interfaces import WeightSynchronizer logger = logging.getLogger(__name__) @@ -638,6 +643,13 @@ def generate( sharded_data: list[SlicedDataDict] = data.shard_by_batch_size( dp_size, allow_uneven_shards=True ) + print_multimodal_payload_metrics( + collect_sharded_multimodal_payload_metrics( + sharded_data, + "vllm_generation", + enabled=self.cfg.get("debug_payload_metrics", False), + ) + ) future_bundle = self.worker_group.run_all_workers_sharded_data( "generate", data=sharded_data, @@ -689,6 +701,13 @@ def generate_text( sharded_data: list[SlicedDataDict] = data.shard_by_batch_size( dp_size, allow_uneven_shards=True ) + print_multimodal_payload_metrics( + collect_sharded_multimodal_payload_metrics( + sharded_data, + "vllm_text_generation", + enabled=self.cfg.get("debug_payload_metrics", False), + ) + ) future_bundle = self.worker_group.run_all_workers_sharded_data( "generate_text", data=sharded_data, @@ -758,6 +777,13 @@ async def _async_generate_base( leader_worker_idx = self.worker_group.get_dp_leader_worker_idx( self.current_generate_dp_shard_idx ) + print_multimodal_payload_metrics( + collect_multimodal_payload_metrics( + data, + "vllm_generation_async", + enabled=self.cfg.get("debug_payload_metrics", False), + ) + ) # Run the async method on the selected leader worker worker_gen_proxy = self.worker_group.run_single_worker_single_data( diff --git a/nemo_rl/models/policy/lm_policy.py b/nemo_rl/models/policy/lm_policy.py index 6a184ded3e9..ef207318437 100644 --- a/nemo_rl/models/policy/lm_policy.py +++ b/nemo_rl/models/policy/lm_policy.py @@ -56,6 +56,10 @@ get_default_hf_config, get_theoretical_tflops, ) +from nemo_rl.utils.multimodal_payload_metrics import ( + collect_sharded_multimodal_payload_metrics, + print_multimodal_payload_metrics, +) from nemo_rl.utils.timer import Timer PathLike = Union[str, "os.PathLike[Any]"] @@ -98,7 +102,9 @@ def __init__( processor: Optional[AutoProcessor] = None, worker_extension_cls_fqn: Optional[str] = None, skip_weight_load: bool = False, + debug_payload_metrics: bool = False, ): + self.debug_payload_metrics = debug_payload_metrics if weights_path: weights_path = os.path.abspath(weights_path) if optimizer_path: @@ -523,6 +529,22 @@ def _shard_for_train( ) return sharded_data + def _report_sharded_payload( + self, + sharded_data: list["SlicedDataDict"], + boundary: str, + ) -> None: + """Measure the exact unique per-DP-shard Ray arguments.""" + if not self.debug_payload_metrics: + return + print_multimodal_payload_metrics( + collect_sharded_multimodal_payload_metrics( + sharded_data, + boundary, + enabled=True, + ) + ) + def get_logprobs( self, data: BatchedDataDict[GenerationDatumSpec], @@ -537,6 +559,7 @@ def get_logprobs( """ with timer.time("get_logprobs/shard_data") if timer else nullcontext(): sharded_data, unsorted_data_indices = self._shard_for_logprob(data) + self._report_sharded_payload(sharded_data, "policy_get_logprobs") with ( timer.time("get_logprobs/submit_logprob_futures") @@ -585,6 +608,7 @@ def get_reference_policy_logprobs( else nullcontext() ): sharded_data, unsorted_data_indices = self._shard_for_logprob(data) + self._report_sharded_payload(sharded_data, "policy_get_reference_logprobs") with ( timer.time( @@ -751,6 +775,7 @@ def train( # Shard and replicate the batch with timer.time("policy_training/sharding_data") if timer else nullcontext(): sharded_data = self._shard_for_train(data, batch_size) + self._report_sharded_payload(sharded_data, "policy_train") if self.flops_tracker is not None: self.flops_tracker.reset() @@ -995,6 +1020,7 @@ def calibrate_qkv_fp8_scales( dp_size, batch_size=None, ) + self._report_sharded_payload(sharded_data, "policy_kv_calibration") futures = self.worker_group.run_all_workers_sharded_data( "calibrate_qkv_fp8_scales", diff --git a/nemo_rl/utils/checkpoint.py b/nemo_rl/utils/checkpoint.py index 8c8de3129d9..4fbbd0ae230 100644 --- a/nemo_rl/utils/checkpoint.py +++ b/nemo_rl/utils/checkpoint.py @@ -126,6 +126,9 @@ class CheckpointingConfig(TypedDict): keep_top_k: NotRequired[int] ft_keep_latest_k: NotRequired[int | None] ft_save_period: NotRequired[int] + # Async GRPO only. Disable to regenerate replay trajectories after resume + # instead of serializing a potentially very large buffer. + save_replay_buffer: NotRequired[bool] checkpoint_must_save_by: NotRequired[str | None] pretrained_checkpoint: NotRequired[PretrainedCheckpointConfig] save_optimizer: NotRequired[bool] # Default: True diff --git a/nemo_rl/utils/multimodal_payload_metrics.py b/nemo_rl/utils/multimodal_payload_metrics.py new file mode 100644 index 00000000000..49c5b82d779 --- /dev/null +++ b/nemo_rl/utils/multimodal_payload_metrics.py @@ -0,0 +1,403 @@ +# 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. + +from collections.abc import Mapping, Sequence +from threading import Lock +from typing import Any + +import numpy as np +import ray.cloudpickle as cloudpickle +import torch +from PIL import Image + +from nemo_rl.data.multimodal_utils import ( + MULTIMODAL_CONTENT_TYPES, + NATIVE_MULTIMODAL_KEYS, + PackedTensor, +) + + +_PENDING_PAYLOAD_METRICS: dict[str, int | float] = {} +_PENDING_PAYLOAD_METRICS_LOCK = Lock() + + +def _tensor_nbytes(value: torch.Tensor | None) -> int: + if value is None: + return 0 + return value.numel() * value.element_size() + + +def _value_nbytes(value: Any, seen: set[int] | None = None) -> int: + """Estimate data bytes, optionally counting shared leaves only once.""" + if value is None: + return 0 + if isinstance(value, (str, bytes, bytearray, memoryview)): + object_id = id(value) + if seen is not None: + if object_id in seen: + return 0 + seen.add(object_id) + if isinstance(value, str): + return len(value.encode("utf-8")) + return len(value) + if torch.is_tensor(value): + object_id = id(value) + if seen is not None: + if object_id in seen: + return 0 + seen.add(object_id) + return _tensor_nbytes(value) + if isinstance(value, np.ndarray): + object_id = id(value) + if seen is not None: + if object_id in seen: + return 0 + seen.add(object_id) + return value.nbytes + if isinstance(value, Image.Image): + object_id = id(value) + if seen is not None: + if object_id in seen: + return 0 + seen.add(object_id) + return len(value.tobytes()) + if isinstance(value, Mapping): + return sum(_value_nbytes(item, seen) for item in value.values()) + if isinstance(value, (list, tuple)): + return sum(_value_nbytes(item, seen) for item in value) + return 0 + + +def _value_segment_count(value: Any, seen: set[int] | None = None) -> int: + """Count media leaves, optionally counting shared objects only once.""" + if value is None: + return 0 + if isinstance( + value, + (str, bytes, bytearray, memoryview, torch.Tensor, np.ndarray, Image.Image), + ): + object_id = id(value) + if seen is not None: + if object_id in seen: + return 0 + seen.add(object_id) + return 1 + if isinstance(value, Mapping): + return sum(_value_segment_count(item, seen) for item in value.values()) + if isinstance(value, (list, tuple)): + return sum(_value_segment_count(item, seen) for item in value) + return 0 + + +def _typed_content_media_nbytes( + value: Any, + seen: set[int] | None = None, +) -> int: + """Count media embedded in typed vLLM content without counting prompt text.""" + if isinstance(value, Mapping): + if value.get("type") in MULTIMODAL_CONTENT_TYPES: + return sum( + _value_nbytes(item, seen) + for key, item in value.items() + if key != "type" + ) + return sum(_typed_content_media_nbytes(item, seen) for item in value.values()) + if isinstance(value, (list, tuple)): + return sum(_typed_content_media_nbytes(item, seen) for item in value) + return 0 + + +def _typed_content_media_segment_count( + value: Any, + seen: set[int] | None = None, +) -> int: + """Count media leaves embedded in typed content without counting text.""" + if isinstance(value, Mapping): + if value.get("type") in MULTIMODAL_CONTENT_TYPES: + return sum( + _value_segment_count(item, seen) + for key, item in value.items() + if key != "type" + ) + return sum( + _typed_content_media_segment_count(item, seen) for item in value.values() + ) + if isinstance(value, (list, tuple)): + return sum(_typed_content_media_segment_count(item, seen) for item in value) + return 0 + + +def protocol5_serialized_nbytes(value: Any) -> int: + """Return cloudpickle protocol-5 frame plus out-of-band buffer bytes.""" + buffers = [] + frame = cloudpickle.dumps(value, protocol=5, buffer_callback=buffers.append) + buffer_bytes = 0 + for buffer in buffers: + raw = buffer.raw() if hasattr(buffer, "raw") else memoryview(buffer) + buffer_bytes += raw.nbytes + return len(frame) + buffer_bytes + + +def collect_multimodal_payload_metrics( + data: Any, + boundary: str, + *, + enabled: bool, +) -> dict[str, int | float]: + """Measure one exact Ray argument without scanning when disabled.""" + if not enabled: + return {} + + totals = { + "physical_media_bytes": 0, + "logical_media_bytes": 0, + "physical_segments": 0, + "logical_segments": 0, + } + per_key: dict[str, int] = {} + seen_native_leaves: set[int] = set() + seen_native_segments: set[int] = set() + seen_packed_leaves: set[int] = set() + seen_packed_leaves_by_key: dict[str, set[int]] = {} + + def visit(key: str, value: Any) -> None: + if isinstance(value, PackedTensor): + key_seen = seen_packed_leaves_by_key.setdefault(key, set()) + key_physical_segments = 0 + key_physical = 0 + for item in value.tensors: + if item is None: + continue + object_id = id(item) + if object_id not in key_seen: + key_seen.add(object_id) + key_physical_segments += 1 + if object_id not in seen_packed_leaves: + seen_packed_leaves.add(object_id) + key_physical += _tensor_nbytes(item) + + logical_items = [ + item for item in value.iter_logical_segments() if item is not None + ] + key_logical = sum(_tensor_nbytes(item) for item in logical_items) + totals["physical_media_bytes"] += key_physical + totals["logical_media_bytes"] += key_logical + totals["physical_segments"] += key_physical_segments + totals["logical_segments"] += len(logical_items) + physical_key = f"payload_counts/{boundary}/{key}/physical_segments" + logical_key = f"payload_counts/{boundary}/{key}/logical_segments" + per_key[physical_key] = per_key.get(physical_key, 0) + key_physical_segments + per_key[logical_key] = per_key.get(logical_key, 0) + len(logical_items) + elif key == "vllm_content": + totals["physical_media_bytes"] += _typed_content_media_nbytes( + value, seen_native_leaves + ) + totals["logical_media_bytes"] += _typed_content_media_nbytes(value) + totals["physical_segments"] += _typed_content_media_segment_count( + value, seen_native_segments + ) + totals["logical_segments"] += _typed_content_media_segment_count(value) + elif key in NATIVE_MULTIMODAL_KEYS: + totals["physical_media_bytes"] += _value_nbytes(value, seen_native_leaves) + totals["logical_media_bytes"] += _value_nbytes(value) + totals["physical_segments"] += _value_segment_count( + value, seen_native_segments + ) + totals["logical_segments"] += _value_segment_count(value) + elif ( + isinstance(value, Mapping) and value.get("type") in MULTIMODAL_CONTENT_TYPES + ): + totals["physical_media_bytes"] += _typed_content_media_nbytes( + value, seen_native_leaves + ) + totals["logical_media_bytes"] += _typed_content_media_nbytes(value) + totals["physical_segments"] += _typed_content_media_segment_count( + value, seen_native_segments + ) + totals["logical_segments"] += _typed_content_media_segment_count(value) + elif isinstance(value, Mapping): + for nested_key, nested_value in value.items(): + visit(str(nested_key), nested_value) + elif isinstance(value, (list, tuple)): + for nested_value in value: + visit(key, nested_value) + + if isinstance(data, Mapping): + for key, value in data.items(): + visit(str(key), value) + else: + visit("root", data) + + physical_media_bytes = totals["physical_media_bytes"] + logical_media_bytes = totals["logical_media_bytes"] + saved_bytes = max(logical_media_bytes - physical_media_bytes, 0) + ratio = ( + float(physical_media_bytes) / float(logical_media_bytes) + if logical_media_bytes + else 1.0 + ) + return { + f"payload_bytes/{boundary}/serialized": protocol5_serialized_nbytes(data), + f"payload_bytes/{boundary}/physical_media": physical_media_bytes, + f"payload_bytes/{boundary}/logical_media": logical_media_bytes, + f"payload_bytes/{boundary}/estimated_saved": saved_bytes, + f"payload_counts/{boundary}/physical_segments": totals["physical_segments"], + f"payload_counts/{boundary}/logical_segments": totals["logical_segments"], + f"payload_counts/{boundary}/calls": 1, + f"payload_ratio/{boundary}/physical_to_logical": ratio, + **per_key, + } + + +def collect_sharded_multimodal_payload_metrics( + shards: Sequence[Mapping[str, Any]], + boundary: str, + *, + enabled: bool, +) -> dict[str, int | float]: + """Aggregate metrics over the exact unique per-DP-shard Ray arguments.""" + if not enabled: + return {} + + per_shard = [ + collect_multimodal_payload_metrics( + shard, + f"{boundary}/shard_{index}", + enabled=True, + ) + for index, shard in enumerate(shards) + ] + serialized = [ + int(metrics[f"payload_bytes/{boundary}/shard_{index}/serialized"]) + for index, metrics in enumerate(per_shard) + ] + physical = [ + int(metrics[f"payload_bytes/{boundary}/shard_{index}/physical_media"]) + for index, metrics in enumerate(per_shard) + ] + logical = [ + int(metrics[f"payload_bytes/{boundary}/shard_{index}/logical_media"]) + for index, metrics in enumerate(per_shard) + ] + physical_segments = [ + int(metrics[f"payload_counts/{boundary}/shard_{index}/physical_segments"]) + for index, metrics in enumerate(per_shard) + ] + logical_segments = [ + int(metrics[f"payload_counts/{boundary}/shard_{index}/logical_segments"]) + for index, metrics in enumerate(per_shard) + ] + total_logical_bytes = sum(logical) + return { + f"payload_bytes/{boundary}/serialized_total": sum(serialized), + f"payload_bytes/{boundary}/serialized_max_shard": max(serialized, default=0), + f"payload_bytes/{boundary}/physical_media_total": sum(physical), + f"payload_bytes/{boundary}/logical_media_total": total_logical_bytes, + f"payload_bytes/{boundary}/estimated_saved_total": max( + total_logical_bytes - sum(physical), 0 + ), + f"payload_counts/{boundary}/physical_segments_total": sum(physical_segments), + f"payload_counts/{boundary}/logical_segments_total": sum(logical_segments), + f"payload_counts/{boundary}/shards": len(shards), + f"payload_counts/{boundary}/calls": len(shards), + f"payload_ratio/{boundary}/physical_to_logical": ( + float(sum(physical)) / float(total_logical_bytes) + if total_logical_bytes + else 1.0 + ), + } + + +def merge_multimodal_payload_metrics( + metric_sets: Sequence[Mapping[str, int | float]], +) -> dict[str, int | float]: + """Aggregate payload measurements collected during one logging interval. + + Byte and segment counts are summed because repeated calls represent distinct + Ray transfers. Per-call maxima and shard counts retain their maximum value, + and physical-to-logical ratios are recomputed from the aggregated byte + totals instead of averaging ratios. + """ + merged: dict[str, int | float] = {} + ratio_keys: set[str] = set() + boundaries: set[str] = set() + for metrics in metric_sets: + for key, value in metrics.items(): + if key.startswith("payload_ratio/") and key.endswith( + "/physical_to_logical" + ): + ratio_keys.add(key) + continue + if key.startswith("payload_bytes/") and key.endswith( + "/serialized_mean_per_call" + ): + continue + if key.startswith("payload_counts/") and key.endswith("/calls"): + boundaries.add( + key.removeprefix("payload_counts/").removesuffix("/calls") + ) + if key.endswith("/serialized_max_shard") or key.endswith("/shards"): + merged[key] = max(merged.get(key, 0), value) + else: + merged[key] = merged.get(key, 0) + value + + for ratio_key in ratio_keys: + boundary = ratio_key.removeprefix("payload_ratio/").removesuffix( + "/physical_to_logical" + ) + physical_key = f"payload_bytes/{boundary}/physical_media" + logical_key = f"payload_bytes/{boundary}/logical_media" + if physical_key not in merged and logical_key not in merged: + physical_key += "_total" + logical_key += "_total" + physical = merged.get(physical_key, 0) + logical = merged.get(logical_key, 0) + merged[ratio_key] = float(physical) / float(logical) if logical else 1.0 + + for boundary in boundaries: + calls = merged[f"payload_counts/{boundary}/calls"] + serialized_key = f"payload_bytes/{boundary}/serialized" + if serialized_key not in merged: + serialized_key += "_total" + if serialized_key in merged: + merged[f"payload_bytes/{boundary}/serialized_mean_per_call"] = ( + float(merged[serialized_key]) / float(calls) if calls else 0.0 + ) + + return merged + + +def drain_multimodal_payload_metrics() -> dict[str, int | float]: + """Return and clear payload measurements recorded in this process.""" + with _PENDING_PAYLOAD_METRICS_LOCK: + pending = dict(_PENDING_PAYLOAD_METRICS) + _PENDING_PAYLOAD_METRICS.clear() + return merge_multimodal_payload_metrics([pending]) + + +def print_multimodal_payload_metrics( + metrics: Mapping[str, int | float], +) -> None: + """Record metrics for the logger and print a stable, scrapeable line.""" + if not metrics: + return + with _PENDING_PAYLOAD_METRICS_LOCK: + merged = merge_multimodal_payload_metrics([_PENDING_PAYLOAD_METRICS, metrics]) + _PENDING_PAYLOAD_METRICS.clear() + _PENDING_PAYLOAD_METRICS.update(merged) + values = [] + for key, value in sorted(metrics.items()): + rendered = f"{value:.6f}" if isinstance(value, float) else str(value) + values.append(f"{key}={rendered}") + print("▶ [PAYLOAD] " + ", ".join(values), flush=True) diff --git a/pyrefly.toml b/pyrefly.toml index 5f55643e09b..494c28a227f 100644 --- a/pyrefly.toml +++ b/pyrefly.toml @@ -213,6 +213,7 @@ project-includes = [ "nemo_rl/utils/config.py", "nemo_rl/utils/fastokens.py", "nemo_rl/utils/grad_norm.py", + "nemo_rl/utils/multimodal_payload_metrics.py", "nemo_rl/utils/native_checkpoint.py", "nemo_rl/utils/nsys.py", "nemo_rl/utils/nvml.py", diff --git a/tests/unit/algorithms/test_async_utils.py b/tests/unit/algorithms/test_async_utils.py index 1e274504fbd..b1c68fa2e17 100644 --- a/tests/unit/algorithms/test_async_utils.py +++ b/tests/unit/algorithms/test_async_utils.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# 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. @@ -42,10 +42,12 @@ GRPOConfig, MasterConfig, _get_next_nemo_gym_task_index, + _should_normalize_sparse_replay_media, add_grpo_token_loss_masks_and_generation_logprobs, extract_initial_prompt_messages, ) from nemo_rl.data.interfaces import DatumSpec, LLMMessageLogType +from nemo_rl.data.multimodal_utils import PackedTensor from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.environments.interfaces import ( EnvironmentInterface, @@ -356,6 +358,67 @@ def test_local_load_state_dict_validates_checkpoint_shape(self): } ) + def test_local_actor_side_checkpoint_preserves_compact_media_and_resume_metadata( + self, tmp_path + ): + checkpoint_path = tmp_path / "replay_buffer.pt" + compact_media = ( + PackedTensor(torch.tensor([[1.0, 2.0]]), dim_to_pack=0) + .enable_deduplication() + .repeat_interleave(2) + ) + source = ReplayBufferImpl(max_size=10) + assert ( + source.add( + { + "batch": {"pixel_values": compact_media}, + "rollout_metrics": {}, + "_ng_task_index": 7, + }, + weight_version=4, + target_weight_version=5, + ) + == "success" + ) + assert ( + source.add( + { + "batch": {"data": "stale"}, + "rollout_metrics": {}, + "_ng_task_index": 41, + }, + weight_version=0, + target_weight_version=5, + ) + == "success" + ) + + assert source.save_to_path(str(checkpoint_path)) == 2 + + restored = ReplayBufferImpl(max_size=10) + metadata = restored.load_from_path( + str(checkpoint_path), + num_prompts_per_step=1, + current_training_step=5, + max_age_steps=1, + ) + + # Metadata accounts for every saved task index, including trajectories + # discarded during resume cleanup, so an index is never reused. + assert metadata == { + "num_trajectories": 2, + "next_ng_task_index": 42, + } + assert restored.size() == 1 + restored_state = restored.state_dict() + restored_media = restored_state["trajectories"][0]["batch"]["pixel_values"] + assert len(restored_media) == 2 + assert len(restored_media.tensors) == 1 + torch.testing.assert_close( + restored_media.as_tensor(), + torch.tensor([[1.0, 2.0], [1.0, 2.0]]), + ) + class TestReplayBuffer: """Test cases for ReplayBuffer.""" @@ -1006,39 +1069,109 @@ def test_replay_buffer_remove_incomplete_resets_watermark_before_first_remaining ray.kill(buffer) - def test_replay_buffer_checkpoint_with_torch_save(self): - """Test that state_dict can be saved and loaded with torch.save/load.""" + def test_replay_buffer_checkpoint_with_torch_save(self, tmp_path): + """Actor-side compact replay checkpoint survives a config flag flip.""" buffer1 = ReplayBuffer.remote(max_size=10) trajectory = { "batch": { "token_ids": torch.tensor([1, 2, 3]), "rewards": torch.tensor([0.5]), + "pixel_values": PackedTensor(torch.tensor([[1.0, 2.0]]), dim_to_pack=0) + .enable_deduplication() + .repeat_interleave(2), }, "rollout_metrics": {"reward": 1.0, "length": 10}, "timestamp": 12345.0, + "_ng_task_index": 11, } ray.get( buffer1.add.remote(trajectory, weight_version=5, target_weight_version=6) ) - state = ray.get(buffer1.state_dict.remote()) - with tempfile.NamedTemporaryFile(suffix=".pt", delete=False) as f: - torch.save(state, f.name) - checkpoint_path = f.name + checkpoint_path = tmp_path / "replay_buffer.pt" + assert ray.get(buffer1.save_to_path.remote(str(checkpoint_path))) == 1 ray.kill(buffer1) - loaded_state = torch.load(checkpoint_path, weights_only=False) buffer2 = ReplayBuffer.remote(max_size=10) - ray.get(buffer2.load_state_dict.remote(loaded_state)) + restore_metadata = ray.get(buffer2.load_from_path.remote(str(checkpoint_path))) + assert restore_metadata == { + "num_trajectories": 1, + "next_ng_task_index": 12, + } assert ray.get(buffer2.size.remote()) == 1 debug_info = ray.get(buffer2.get_debug_info.remote()) assert debug_info["trajectory_versions"] == [5] assert debug_info["target_weight_versions"] == [6] + restored_state = ray.get(buffer2.state_dict.remote()) + restored_media = restored_state["trajectories"][0]["batch"]["pixel_values"] + assert restored_media.deduplication_enabled + assert len(restored_media) == 2 + assert len(restored_media.tensors) == 1 + torch.testing.assert_close( + restored_media.as_tensor(), + torch.tensor([[1.0, 2.0], [1.0, 2.0]]), + ) + + restored_sparse_batches = [ + BatchedDataDict( + { + "token_ids": torch.tensor([[1]]), + "pixel_values": restored_media.slice([0]), + } + ), + BatchedDataDict({"token_ids": torch.tensor([[2]])}), + ] + assert _should_normalize_sparse_replay_media( + restored_sparse_batches, + deduplicate_multimodal_data=False, + ) + restored_sparse = BatchedDataDict.from_batches( + restored_sparse_batches, + allow_missing_packed_tensors=True, + ) + assert len(restored_sparse["pixel_values"]) == 2 + assert restored_sparse["pixel_values"].logical_segment_count == 1 + + # The representation is self-describing: after a flag-on checkpoint is + # restored by a flag-off run, newly collected legacy media can be + # concatenated without expanding the restored physical segment. + legacy_media = PackedTensor(torch.tensor([[3.0, 4.0]]), dim_to_pack=0) + mixed_after_flag_off = BatchedDataDict.from_batches( + [ + {"pixel_values": restored_media}, + {"pixel_values": legacy_media}, + ], + allow_missing_packed_tensors=True, + ) + assert len(mixed_after_flag_off["pixel_values"].tensors) == 2 + torch.testing.assert_close( + mixed_after_flag_off["pixel_values"].as_tensor(), + torch.tensor([[1.0, 2.0], [1.0, 2.0], [3.0, 4.0]]), + ) + + # The inverse transition is valid too: legacy checkpoint media is + # assigned fresh provenance when combined with new compact media. + new_compact_media = ( + PackedTensor(torch.tensor([[5.0, 6.0]]), dim_to_pack=0) + .enable_deduplication() + .repeat_interleave(2) + ) + mixed_after_flag_on = BatchedDataDict.from_batches( + [ + {"pixel_values": legacy_media}, + {"pixel_values": new_compact_media}, + ], + allow_missing_packed_tensors=True, + ) + assert len(mixed_after_flag_on["pixel_values"].tensors) == 2 + torch.testing.assert_close( + mixed_after_flag_on["pixel_values"].as_tensor(), + torch.tensor([[3.0, 4.0], [5.0, 6.0], [5.0, 6.0]]), + ) - os.unlink(checkpoint_path) ray.kill(buffer2) def test_resume_deadlock_precondition_detectable(self): @@ -1159,6 +1292,31 @@ def test_collection_loop_marks_data_exhausted_on_natural_completion(self): assert status["errored"] is False assert status["running"] is False + @pytest.mark.asyncio + async def test_drain_payload_metrics_returns_collector_interval(self, monkeypatch): + collector = self.create_local_collector() + collector.master_config.grpo.debug_payload_metrics = True + monkeypatch.setattr( + "nemo_rl.algorithms.async_utils.trajectory_collector." + "drain_multimodal_payload_metrics", + lambda: { + "payload_bytes/nemo_gym_return/serialized": 180, + "payload_bytes/nemo_gym_return/serialized_mean_per_call": 90, + "payload_bytes/nemo_gym_return/physical_media": 30, + "payload_bytes/nemo_gym_return/logical_media": 150, + "payload_counts/nemo_gym_return/calls": 2, + "payload_ratio/nemo_gym_return/physical_to_logical": 0.2, + }, + ) + + metrics = await collector.drain_payload_metrics() + + assert metrics["payload_counts/nemo_gym_return/calls"] == 2 + assert metrics["payload_bytes/nemo_gym_return/serialized_mean_per_call"] == 90 + assert metrics["payload_bytes/nemo_gym_return/physical_media"] == 30 + assert metrics["payload_bytes/nemo_gym_return/logical_media"] == 150 + assert metrics["payload_ratio/nemo_gym_return/physical_to_logical"] == 0.2 + def test_collection_loop_marks_errored_on_crash(self): """A crash sets errored (not data_exhausted) so driver guards fail fast.""" collector = self.create_local_collector() @@ -1423,7 +1581,8 @@ class FakeBatch: def slice(self, start, end): return self - def repeat_interleave(self, repeats): + def repeat_interleave(self, repeats, *, share_immutable_media=False): + assert not share_immutable_media return self class FailingThread: @@ -1487,6 +1646,7 @@ def is_alive(self): target_weight = 7 collector = self.create_local_collector(replay_buffer=FakeReplayBuffer()) + collector.master_config.grpo.deduplicate_multimodal_data = True collector.running = True def reserve_target(generation_weight_version): diff --git a/tests/unit/algorithms/test_grpo.py b/tests/unit/algorithms/test_grpo.py index dee6cdb260b..c2bc53abe93 100644 --- a/tests/unit/algorithms/test_grpo.py +++ b/tests/unit/algorithms/test_grpo.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# 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. @@ -42,7 +42,9 @@ _raise_if_reward_penalties_enabled_without_nemo_gym, _resolve_logprob_skip_flags, _resolve_message_level_advantage_penalties, + _save_async_replay_buffer_checkpoint, _should_use_async_rollouts, + _validate_multimodal_dedup_capability, _validate_use_kl_in_reward_compat, aggregate_rollout_metrics, async_grpo_train, @@ -60,6 +62,7 @@ ) from nemo_rl.algorithms.utils import calculate_baseline_and_std_per_prompt from nemo_rl.data.interfaces import DatumSpec, LLMMessageLogType +from nemo_rl.data.multimodal_utils import PackedTensor from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.environments.interfaces import ( EnvironmentInterface, @@ -82,6 +85,36 @@ def _mock_policy_generation() -> MagicMock: return policy_generation +@pytest.mark.parametrize( + ("checkpointing_config", "expected_count"), + [ + ({}, 7), + ({"save_replay_buffer": True}, 7), + ({"save_replay_buffer": False}, None), + ], +) +def test_save_async_replay_buffer_checkpoint_gate( + tmp_path, checkpointing_config, expected_count +): + replay_buffer = MagicMock() + replay_buffer.save_to_path.remote.return_value = 7 + + with patch("nemo_rl.algorithms.grpo.ray.get", side_effect=lambda value: value): + count = _save_async_replay_buffer_checkpoint( + replay_buffer, + str(tmp_path), + checkpointing_config, + ) + + assert count == expected_count + if expected_count is None: + replay_buffer.save_to_path.remote.assert_not_called() + else: + replay_buffer.save_to_path.remote.assert_called_once_with( + str(tmp_path / "replay_buffer.pt") + ) + + @patch("nemo_rl.algorithms.grpo.ray") def test_refit_policy_generation_forwards_kv_scales_on_colocated_ipc( mock_ray: MagicMock, @@ -710,6 +743,25 @@ def test_raise_if_message_level_advantage_penalties_enabled_raises_when_set( _raise_if_message_level_advantage_penalties_enabled(master_config) +def test_multimodal_dedup_rejects_unqualified_transfer_paths( + mock_grpo_components, +): + master_config = mock_grpo_components["master_config"] + master_config.grpo.deduplicate_multimodal_data = True + master_config.policy["generation"]["backend"] = "sglang" + + with pytest.raises(NotImplementedError, match="backend=vllm"): + _validate_multimodal_dedup_capability(master_config) + + master_config.policy["generation"]["backend"] = "vllm" + master_config.data_plane = {"enabled": True} + with pytest.raises(NotImplementedError, match="data_plane.enabled=false"): + _validate_multimodal_dedup_capability(master_config) + + master_config.data_plane = {"enabled": False} + _validate_multimodal_dedup_capability(master_config) + + def test_grpo_sync_seq_logprob_error_helper_accepts_dict_result(monkeypatch): from nemo_rl.algorithms import grpo_sync as grpo_sync_mod @@ -852,6 +904,25 @@ def _load_state_dict(state, *args, **kwargs): mock.remote = MagicMock(side_effect=_load_state_dict) return mock + @property + def save_to_path(self): + """Return a mock that checkpoints state without a driver-sized return.""" + mock = MagicMock() + mock.remote = MagicMock(return_value=self._size) + return mock + + @property + def load_from_path(self): + """Return compact restore metadata.""" + mock = MagicMock() + mock.remote = MagicMock( + return_value={ + "num_trajectories": self._size, + "next_ng_task_index": 0, + } + ) + return mock + @property def get_trajectories_needed(self): """Return a mock that reports how many prompt groups are still needed.""" @@ -1651,6 +1722,63 @@ def test_dapo_dynamic_sampling_batch_caching(mock_grpo_components): assert batch_cache is not None +def test_dapo_cache_aligns_deduplicated_media_with_text_only_batch( + mock_grpo_components, +): + def make_batch(prompt: str, *, with_media: bool) -> BatchedDataDict: + message_logs = [ + [ + {"role": "user", "content": prompt}, + {"role": "assistant", "content": f"response_{i}"}, + ] + for i in range(3) + ] + batch = create_mock_batch(3, ["math"] * 3, message_logs) + batch["total_reward"] = torch.tensor([1.0, 0.0, 0.5]) + if with_media: + media = PackedTensor( + torch.tensor([[1.0]]), dim_to_pack=0 + ).enable_deduplication() + batch["pixel_values"] = media.repeat_interleave(3) + return batch + + master_config = mock_grpo_components["master_config"] + master_config.grpo.use_dynamic_sampling = True + master_config.grpo.num_prompts_per_step = 2 + master_config.grpo.num_generations_per_prompt = 3 + master_config.grpo.dynamic_sampling_max_gen_batches = 5 + master_config.grpo.deduplicate_multimodal_data = True + std = torch.tensor([0.4, 0.4, 0.4]) + baseline = torch.tensor([0.5, 0.5, 0.5]) + + _, complete, cache, _ = dynamic_sampling( + make_batch("visual", with_media=True), + std, + baseline, + dynamic_sampling_num_gen_batches=1, + master_config=master_config, + timer=Timer(), + ) + assert not complete + assert cache is not None + + result, complete, _, _ = dynamic_sampling( + make_batch("text", with_media=False), + std, + baseline, + dynamic_sampling_num_gen_batches=2, + master_config=master_config, + timer=Timer(), + batch_cache=cache, + ) + + assert complete + assert result.size == 6 + assert len(result["pixel_values"]) == 6 + assert len(result["pixel_values"].tensors) == 1 + assert result["pixel_values"].slice([3, 4, 5]).as_tensor() is None + + def test_dapo_dynamic_sampling_disabled(mock_grpo_components): """Test that when dynamic sampling is disabled, all prompts are kept regardless of std.""" batch_size = 6 @@ -1942,7 +2070,7 @@ def test_noncolocated_opd_teacher_must_fit_on_one_cluster_node( "initial_skip_flag", [None, False], ) -def test_setup_auto_enables_skip_reference_policy_logprobs_when_kl_penalty_zero( +def test_setup_auto_enables_skip_reference_logprobs_with_legacy_policy_factory( monkeypatch, mock_grpo_components, initial_skip_flag ): from nemo_rl.algorithms import grpo as grpo_mod @@ -1997,6 +2125,29 @@ def prepare_refit_info(self): def set_rollout_num_gpus_per_engine(self, _num_gpus_per_engine): pass + def legacy_policy_factory( + *, + cluster, + config, + tokenizer, + processor, + weights_path, + optimizer_path, + init_optimizer, + init_reference_model, + ): + del ( + cluster, + config, + tokenizer, + processor, + weights_path, + optimizer_path, + init_optimizer, + init_reference_model, + ) + return DummyPolicy() + class DummySGLangGeneration: num_gpus_per_engine = 1 @@ -2018,7 +2169,6 @@ def init_collective(self, *_args, **_kwargs): ) monkeypatch.setattr(grpo_mod, "StatefulDataLoader", DummyLoader) monkeypatch.setattr(grpo_mod, "RayVirtualCluster", DummyCluster) - monkeypatch.setattr(grpo_mod, "Policy", lambda *_args, **_kwargs: DummyPolicy()) monkeypatch.setattr( grpo_mod, "SGLangGeneration", @@ -2059,7 +2209,13 @@ def init_collective(self, *_args, **_kwargs): dataset = MagicMock() dataset.__len__ = MagicMock(return_value=1) - grpo_mod.setup(master_config, tokenizer, dataset, None) + grpo_mod.setup( + master_config, + tokenizer, + dataset, + None, + policy_factory=legacy_policy_factory, + ) assert master_config.grpo.skip_reference_policy_logprobs_calculation is True diff --git a/tests/unit/algorithms/test_grpo_router_replay_async.py b/tests/unit/algorithms/test_grpo_router_replay_async.py index 52aa8369ac5..a831d64731a 100644 --- a/tests/unit/algorithms/test_grpo_router_replay_async.py +++ b/tests/unit/algorithms/test_grpo_router_replay_async.py @@ -109,6 +109,28 @@ def test_build_async_grpo_train_data_preserves_routed_experts_for_r3( assert "routed_experts" not in train_data +def test_build_async_grpo_train_data_accepts_all_text_vlm_replay_batch(): + flat_messages = BatchedDataDict( + { + "token_ids": torch.tensor([[1, 2, 3]]), + "generation_logprobs": torch.zeros(1, 3), + "token_loss_mask": torch.tensor([[0, 1, 1]]), + } + ) + input_lengths = torch.tensor([3]) + repeated_batch = BatchedDataDict({"loss_multiplier": torch.tensor([1.0])}) + + train_data = _build_async_grpo_train_data( + flat_messages, + input_lengths, + repeated_batch, + {**_make_async_master_config().policy, "is_vlm": True}, + ) + + assert train_data["input_ids"].tolist() == [[1, 2, 3]] + assert train_data.get_multimodal_dict(as_tensors=False) == {} + + def test_async_grpo_r3_data_plane_directs_to_single_controller(): master_config = _make_async_master_config(data_plane={"enabled": True}) diff --git a/tests/unit/data/test_llm_message_utils.py b/tests/unit/data/test_llm_message_utils.py index 113fd9ce0b9..8a8e0ab93ca 100644 --- a/tests/unit/data/test_llm_message_utils.py +++ b/tests/unit/data/test_llm_message_utils.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# 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. @@ -32,6 +32,7 @@ get_keys_from_message_log, message_log_to_flat_messages, ) +from nemo_rl.data.multimodal_utils import PackedTensor @pytest.fixture @@ -922,6 +923,53 @@ def test_get_formatted_message_log_debug_off_by_default( assert "DEBUG: Individual message turns" not in captured.out +@pytest.mark.parametrize("visual_first", [False, True]) +def test_batched_flatten_aligns_nested_sparse_multimodal_rows(visual_first): + media = PackedTensor(torch.ones(1, 3, 2, 2), dim_to_pack=0).enable_deduplication() + visual_log = [ + { + "role": "user", + "content": "look", + "token_ids": torch.tensor([1]), + "pixel_values": media, + }, + { + "role": "assistant", + "content": "seen", + "token_ids": torch.tensor([2]), + }, + ] + text_log = [ + { + "role": "user", + "content": "text only", + "token_ids": torch.tensor([3]), + }, + { + "role": "assistant", + "content": "answer", + "token_ids": torch.tensor([4]), + }, + ] + message_logs = [visual_log, text_log] if visual_first else [text_log, visual_log] + + flat, _ = batched_message_log_to_flat_message( + message_logs, pad_value_dict={"token_ids": 0} + ) + packed = flat["pixel_values"] + + assert isinstance(packed, PackedTensor) + assert len(packed) == 2 + assert packed.logical_segment_counts_by_row() == ( + [1, 0] if visual_first else [0, 1] + ) + assert len(packed.tensors) == 1 + torch.testing.assert_close( + flat.get_multimodal_dict(as_tensors=True)["pixel_values"], + torch.ones(1, 3, 2, 2), + ) + + def test_get_formatted_message_log_debug_enabled( raw_chat_message_log: LLMMessageLogType, capsys, diff --git a/tests/unit/data/test_multimodal_dict.py b/tests/unit/data/test_multimodal_dict.py index 23b7cdacdbd..b97bbf9d204 100644 --- a/tests/unit/data/test_multimodal_dict.py +++ b/tests/unit/data/test_multimodal_dict.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# 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. @@ -11,7 +11,10 @@ # 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. +from copy import deepcopy + import pytest +import ray.cloudpickle as cloudpickle import torch from nemo_rl.data.llm_message_utils import batched_message_log_to_flat_message @@ -480,3 +483,139 @@ def test_slice_preserves_pad_to_max_shape_flag(): assert sliced.pad_to_max_shape is True assert sliced.as_tensor().shape == (2, 3, 4, 4) + + +def test_packedtensor_dedup_uses_provenance_not_prompt_position(): + """Only segments descended from the same physical media are compacted.""" + shared = PackedTensor(torch.tensor([[1.0]]), dim_to_pack=0) + shared.enable_deduplication() + shared_copy = deepcopy(shared) + same_prompt_but_different_media = PackedTensor(torch.tensor([[1.0]]), dim_to_pack=0) + same_prompt_but_different_media.enable_deduplication() + + packed = PackedTensor.concat([shared, shared_copy, same_prompt_but_different_media]) + + assert len(packed) == 3 + assert packed.logical_segment_count == 3 + assert len(packed.tensors) == 2 + torch.testing.assert_close(packed.as_tensor(), torch.tensor([[1.0], [1.0], [1.0]])) + + +def test_packedtensor_multiturn_csr_preserves_shared_seed_and_unique_media(): + """Diverged rows retain one seed segment plus their own later segment.""" + seed = PackedTensor(torch.tensor([[1.0]]), dim_to_pack=0) + seed.enable_deduplication() + row_1 = PackedTensor.merge_segments( + [deepcopy(seed), PackedTensor(torch.tensor([[2.0]]), dim_to_pack=0)] + ) + row_2 = PackedTensor.merge_segments( + [deepcopy(seed), PackedTensor(torch.tensor([[3.0]]), dim_to_pack=0)] + ) + + packed = PackedTensor.flattened_concat([row_1, row_2]) + + assert len(packed) == 2 + assert packed.logical_segment_count == 4 + assert len(packed.tensors) == 3 + torch.testing.assert_close( + packed.as_tensor(), torch.tensor([[1.0], [2.0], [1.0], [3.0]]) + ) + + second_row = packed.slice([1]) + assert len(second_row) == 1 + assert len(second_row.tensors) == 2 + torch.testing.assert_close(second_row.as_tensor(), torch.tensor([[1.0], [3.0]])) + + +def test_packedtensor_dedup_expands_before_dynamic_shape_padding(): + """Logical order is restored before non-packing dimensions are padded.""" + first = PackedTensor( + torch.ones(1, 1, 2), + dim_to_pack=0, + pad_to_max_shape=True, + ).enable_deduplication() + second = PackedTensor( + 2 * torch.ones(1, 2, 1), + dim_to_pack=0, + pad_to_max_shape=True, + ).enable_deduplication() + + packed = PackedTensor.concat([first, deepcopy(first), second]) + materialized = packed.as_tensor() + + assert materialized.shape == (3, 2, 2) + torch.testing.assert_close(materialized[0], materialized[1]) + torch.testing.assert_close(materialized[2, :, 0], 2 * torch.ones(2)) + + +def test_packedtensor_dedup_dim_one_slice_empty_and_cloudpickle_roundtrip(): + first = torch.tensor([[1.0], [2.0]]) + second = torch.tensor([[3.0, 4.0], [5.0, 6.0]]) + packed = PackedTensor( + [first, second], + dim_to_pack=1, + ).enable_deduplication() + repeated = packed.repeat_interleave(2) + + assert len(repeated) == 4 + assert len(repeated.tensors) == 2 + torch.testing.assert_close( + repeated.as_tensor(), + torch.cat([first, first, second, second], dim=1), + ) + + selected = repeated.slice([3, 0, -1]) + assert len(selected) == 3 + assert len(selected.tensors) == 2 + torch.testing.assert_close( + selected.as_tensor(), + torch.cat([second, first, second], dim=1), + ) + + restored = cloudpickle.loads(cloudpickle.dumps(selected, protocol=5)) + assert restored.deduplication_enabled + assert len(restored) == 3 + assert len(restored.tensors) == 2 + torch.testing.assert_close(restored.as_tensor(), selected.as_tensor()) + + empty = packed.repeat_interleave(0) + assert len(empty) == 0 + assert empty.logical_segment_count == 0 + assert empty.as_tensor() is None + + +def test_packedtensor_unpickles_pre_deduplication_state(): + tensor = torch.tensor([[1.0], [2.0]]) + legacy = PackedTensor.__new__(PackedTensor) + legacy.__dict__ = { + "tensors": [tensor], + "dim_to_pack": 0, + "pad_to_max_shape": False, + } + + restored = cloudpickle.loads(cloudpickle.dumps(legacy, protocol=5)) + + assert not restored.deduplication_enabled + assert len(restored) == 1 + assert restored.logical_segment_count == 1 + torch.testing.assert_close(restored.as_tensor(), tensor) + restored.enable_deduplication() + assert restored.deduplication_enabled + + +def test_packedtensor_empty_legacy_rows_survive_copy_pickle_and_slice(): + legacy = PackedTensor(torch.tensor([[1.0]]), dim_to_pack=0) + empty = PackedTensor.empty_rows_like(legacy, 0) + + assert len(empty) == 0 + assert not empty.deduplication_enabled + assert empty.as_tensor() is None + + copied = deepcopy(empty) + restored = cloudpickle.loads(cloudpickle.dumps(empty, protocol=5)) + sliced = empty.slice([]) + for value in (copied, restored, sliced): + assert len(value) == 0 + assert value.logical_segment_count == 0 + assert not value.deduplication_enabled + assert value.as_tensor() is None diff --git a/tests/unit/distributed/test_batched_data_dict.py b/tests/unit/distributed/test_batched_data_dict.py index 46b35f69a5f..a9823f0d56c 100644 --- a/tests/unit/distributed/test_batched_data_dict.py +++ b/tests/unit/distributed/test_batched_data_dict.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# 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. @@ -11,6 +11,7 @@ # 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. +import numpy as np import pytest import torch @@ -47,6 +48,27 @@ def test_shard_by_batch_size_basic(): assert torch.equal(sharded[1]["other_tensor"], torch.tensor([12, 13, 16, 17])) +def test_from_batches_flag_off_keeps_legacy_first_field_batch_size(): + class MetadataWithoutBatchLength: + def __len__(self): + raise AssertionError("flag-off batch sizing must not scan metadata") + + batches = [ + { + "tokens": torch.tensor([[1], [2]]), + "metadata": MetadataWithoutBatchLength(), + }, + { + "tokens": torch.tensor([[3]]), + "metadata": MetadataWithoutBatchLength(), + "extra": torch.tensor([1]), + }, + ] + + with pytest.raises(KeyError, match="'extra'"): + BatchedDataDict.from_batches(batches) + + def test_shard_by_batch_size_list_data(): """Test shard_by_batch_size with list data.""" # Create a sample batch with list data @@ -625,6 +647,113 @@ def test_shard_by_batch_size_with_packed_multimodal(): assert tuple(shards[1]["pixel_values"].as_tensor().shape) == (6, 3, 8, 8) +def test_repeat_interleave_shares_only_flagged_multimodal_segments(): + media = PackedTensor(torch.ones(1, 3, 2, 2), dim_to_pack=0) + batch = BatchedDataDict( + { + "message_log": [ + [ + { + "role": "user", + "content": "look", + "token_ids": torch.tensor([1, 2]), + "pixel_values": media, + } + ] + ] + } + ) + + flag_off = batch.repeat_interleave(2) + off_first = flag_off["message_log"][0][0]["pixel_values"] + off_second = flag_off["message_log"][1][0]["pixel_values"] + assert not off_first.deduplication_enabled + assert off_first.tensors[0] is not off_second.tensors[0] + + flag_on = batch.repeat_interleave(2, share_immutable_media=True) + on_first = flag_on["message_log"][0][0]["pixel_values"] + on_second = flag_on["message_log"][1][0]["pixel_values"] + assert on_first.deduplication_enabled + assert on_first is not on_second + assert on_first.tensors[0] is on_second.tensors[0] + assert flag_on["message_log"][0] is not flag_on["message_log"][1] + + +def test_repeat_interleave_shares_native_image_video_and_audio_leaves(): + image = torch.ones(1, 2) + video = np.ones((2, 2), dtype=np.float32) + audio = np.ones(16, dtype=np.float32) + batch = BatchedDataDict( + { + "vllm_images": [[image]], + "vllm_videos": [[video]], + "vllm_audios": [[(audio, 16_000)]], + } + ) + + flag_off = batch.repeat_interleave(2) + assert flag_off["vllm_images"][0][0] is not flag_off["vllm_images"][1][0] + assert flag_off["vllm_videos"][0][0] is not flag_off["vllm_videos"][1][0] + assert flag_off["vllm_audios"][0][0][0] is not flag_off["vllm_audios"][1][0][0] + + flag_on = batch.repeat_interleave(2, share_immutable_media=True) + assert flag_on["vllm_images"][0] is not flag_on["vllm_images"][1] + assert flag_on["vllm_images"][0][0] is flag_on["vllm_images"][1][0] + assert flag_on["vllm_videos"][0][0] is flag_on["vllm_videos"][1][0] + assert flag_on["vllm_audios"][0][0][0] is flag_on["vllm_audios"][1][0][0] + + +def test_shards_reintern_shared_segments_locally(): + media = PackedTensor(torch.ones(1, 2), dim_to_pack=0).enable_deduplication() + repeated_media = media.repeat_interleave(4) + batch = BatchedDataDict( + { + "input_ids": torch.arange(8).reshape(4, 2), + "input_lengths": torch.tensor([2, 2, 2, 2]), + "pixel_values": repeated_media, + } + ) + + shards = batch.shard_by_batch_size(shards=2) + + assert [len(shard["pixel_values"]) for shard in shards] == [2, 2] + assert [len(shard["pixel_values"].tensors) for shard in shards] == [1, 1] + + +def test_sequence_packing_reinterns_shared_segments_per_shard_for_cp_padding(): + media = PackedTensor(torch.ones(1, 2), dim_to_pack=0) + repeated_media = media.enable_deduplication().repeat_interleave(8) + sequence_lengths = torch.tensor([5, 6, 7, 8, 9, 10, 11, 12]) + batch = BatchedDataDict( + { + "input_ids": torch.arange(8 * 12).reshape(8, 12), + "input_lengths": sequence_lengths, + "pixel_values": repeated_media, + } + ) + sequence_packing_args = SequencePackingArgs( + max_tokens_per_microbatch=24, + input_key="input_ids", + input_lengths_key="input_lengths", + algorithm="modified_first_fit_decreasing", + # CP=2 requires sequences to be divisible by 2 * CP. + sequence_length_pad_multiple=4, + ) + + shards, _ = batch.shard_by_batch_size( + shards=2, + sequence_packing_args=sequence_packing_args, + ) + + assert sum(len(shard["pixel_values"]) for shard in shards) == 8 + assert [len(shard["pixel_values"].tensors) for shard in shards] == [1, 1] + for shard in shards: + torch.testing.assert_close( + shard["pixel_values"].as_tensor(), + torch.ones(len(shard["pixel_values"]), 2), + ) + + def test_shard_by_batch_size_allow_uneven_empty_shards_preserve_all_keys(): """Empty trailing shards should preserve all keys with empty values.""" batch = BatchedDataDict( @@ -643,6 +772,7 @@ def test_shard_by_batch_size_allow_uneven_empty_shards_preserve_all_keys(): # Empty trailing shards should preserve all keys and use empty values. for empty_shard in shards[2:]: + assert empty_shard.size == 0 for key, original_value in batch.items(): assert key in empty_shard shard_value = empty_shard[key] @@ -650,6 +780,7 @@ def test_shard_by_batch_size_allow_uneven_empty_shards_preserve_all_keys(): assert shard_value.shape[0] == 0 elif isinstance(original_value, PackedTensor): assert isinstance(shard_value, PackedTensor) + assert len(shard_value) == 0 assert shard_value.as_tensor() is None else: assert shard_value == [] @@ -848,6 +979,173 @@ def test_from_batches_keeps_keys_missing_from_empty_mapping(): assert torch.equal(stacked["routed_experts"], routed_experts) +def test_from_batches_can_align_optional_deduplicated_media_keys(): + shared_pixels = PackedTensor( + torch.tensor([[1.0]]), dim_to_pack=0 + ).enable_deduplication() + pixel_rows = shared_pixels.repeat_interleave(2) + distinct_image_sizes = PackedTensor( + [torch.tensor([[10, 20]]), torch.tensor([[30, 40]])], + dim_to_pack=0, + ).enable_deduplication() + audio_rows = PackedTensor( + torch.tensor([[5.0]]), dim_to_pack=0 + ).enable_deduplication() + + visual_batch = BatchedDataDict( + { + "input_ids": torch.tensor([[1, 2], [3, 4]]), + "pixel_values": pixel_rows, + "imgs_sizes": distinct_image_sizes, + } + ) + audio_batch = BatchedDataDict( + { + "input_ids": torch.tensor([[5, 6]]), + "audio_values": audio_rows, + } + ) + + stacked = BatchedDataDict.from_batches( + [visual_batch, audio_batch], + allow_missing_packed_tensors=True, + ) + + assert stacked.size == 3 + assert { + key: len(stacked[key]) for key in ("pixel_values", "imgs_sizes", "audio_values") + } == { + "pixel_values": 3, + "imgs_sizes": 3, + "audio_values": 3, + } + assert len(stacked["pixel_values"].tensors) == 1 + assert len(stacked["imgs_sizes"].tensors) == 2 + assert stacked["pixel_values"].slice([2]).as_tensor() is None + assert stacked["imgs_sizes"].slice([2]).as_tensor() is None + assert stacked["audio_values"].slice([0, 1]).as_tensor() is None + torch.testing.assert_close( + stacked["audio_values"].slice([2]).as_tensor(), + torch.tensor([[5.0]]), + ) + + +def test_from_batches_optional_media_rejects_cross_key_row_misalignment(): + batch = BatchedDataDict( + { + "pixel_values": PackedTensor( + torch.tensor([[1.0]]), dim_to_pack=0 + ).enable_deduplication(), + "input_ids": torch.tensor([[1, 2], [3, 4]]), + } + ) + + with pytest.raises(ValueError, match="inconsistent logical row counts"): + BatchedDataDict.from_batches( + [batch], + allow_missing_packed_tensors=True, + ) + + +def test_model_materialization_validates_coupled_media_segment_order(): + pixels = PackedTensor( + [torch.tensor([[1.0]]), torch.tensor([[2.0]])], + dim_to_pack=0, + ).enable_deduplication() + image_sizes = PackedTensor( + [torch.tensor([[10, 20]]), torch.tensor([[30, 40]])], + dim_to_pack=0, + ).enable_deduplication() + valid = BatchedDataDict( + { + "pixel_values": pixels, + "imgs_sizes": image_sizes, + } + ) + + materialized = valid.get_multimodal_dict(as_tensors=True) + torch.testing.assert_close( + materialized["pixel_values"], torch.tensor([[1.0], [2.0]]) + ) + torch.testing.assert_close( + materialized["imgs_sizes"], + torch.tensor([[10, 20], [30, 40]]), + ) + + first_row_only = PackedTensor.merge_segments( + [ + PackedTensor( + torch.tensor([[10, 20]]), dim_to_pack=0 + ).enable_deduplication(), + PackedTensor( + torch.tensor([[30, 40]]), dim_to_pack=0 + ).enable_deduplication(), + ] + ) + missing_second_row = PackedTensor.concat( + [first_row_only, PackedTensor.empty_rows_like(first_row_only, 1)] + ) + invalid = BatchedDataDict( + { + "pixel_values": pixels, + "imgs_sizes": missing_second_row, + } + ) + + with pytest.raises(ValueError, match="ordered per-row segment counts"): + invalid.get_multimodal_dict(as_tensors=True) + + +def test_size_supports_packed_tensor_as_first_key_and_empty_batches(): + media = PackedTensor( + [torch.tensor([[1.0]]), torch.tensor([[2.0]])], + dim_to_pack=0, + ) + batch = BatchedDataDict( + { + "pixel_values": media, + "input_ids": torch.tensor([[1, 2], [3, 4]]), + } + ) + + assert batch.size == 2 + assert BatchedDataDict().size == 0 + + +def test_deduplicated_media_survives_chunk_reorder_and_select_indices(): + first = PackedTensor(torch.tensor([[1.0]]), dim_to_pack=0).enable_deduplication() + second = PackedTensor(torch.tensor([[2.0]]), dim_to_pack=0).enable_deduplication() + media = PackedTensor.concat( + [first.repeat_interleave(2), second.repeat_interleave(2)] + ) + batch = BatchedDataDict( + { + "pixel_values": media, + "input_ids": torch.arange(8).reshape(4, 2), + } + ) + + first_chunk = batch.chunk(rank=0, chunks=2) + assert len(first_chunk["pixel_values"].tensors) == 1 + torch.testing.assert_close( + first_chunk["pixel_values"].as_tensor(), + torch.tensor([[1.0], [1.0]]), + ) + + batch.reorder_data([3, 2, 1, 0]) + torch.testing.assert_close( + batch["pixel_values"].as_tensor(), + torch.tensor([[2.0], [2.0], [1.0], [1.0]]), + ) + + selected = batch.select_indices([0, 3]) + assert len(selected["pixel_values"].tensors) == 2 + torch.testing.assert_close( + selected["pixel_values"].as_tensor(), + torch.tensor([[2.0], [1.0]]), + ) + + @pytest.mark.parametrize("pad_to_multiple_of", [1, 32, 64, 256]) def test_sequence_packing_microbatch_boundaries(pad_to_multiple_of): """Test that microbatch boundaries are correctly maintained across chunks with random sequences.""" diff --git a/tests/unit/environments/test_nemo_gym.py b/tests/unit/environments/test_nemo_gym.py index a4ca5759e98..b32f9d7b947 100644 --- a/tests/unit/environments/test_nemo_gym.py +++ b/tests/unit/environments/test_nemo_gym.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# 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. @@ -20,9 +20,11 @@ import ray import requests import torch +from PIL import Image from yaml import safe_load from nemo_rl.algorithms.grpo import MasterConfig +from nemo_rl.data.multimodal_utils import PackedTensor, image_to_data_url from nemo_rl.distributed.ray_actor_environment_registry import ( get_actor_python_env, ) @@ -34,6 +36,7 @@ setup_nemo_gym_config, validate_reward_components_match_scalar, ) +from nemo_rl.experience.rollouts import _reattach_original_multimodal_payloads from nemo_rl.models.generation.vllm import VllmGeneration # cluster and tokenizer are fixture imports @@ -293,7 +296,7 @@ class _MockSelf: result = ( NemoGym.__ray_metadata__.modified_class._postprocess_nemo_gym_to_nemo_rl_result( - _MockSelf(), nemo_gym_result, tokenizer + _MockSelf(), {}, nemo_gym_result, tokenizer ) ) @@ -311,6 +314,295 @@ class _MockSelf: assert nemo_gym_result["response"]["output"][1]["generation_str"] == "6 7" +@pytest.mark.parametrize("include_initial_multimodal_data", [False, True]) +def test_nemo_gym_dedup_redacts_initial_images_from_actor_return( + include_initial_multimodal_data, +): + data_url = image_to_data_url(Image.new("RGB", (2, 2), color="red")) + initial_input = [ + { + "role": "user", + "content": [ + {"type": "input_text", "text": "count"}, + {"type": "input_image", "image_url": data_url}, + ], + } + ] + nemo_gym_result = { + "response": { + "agent_input": deepcopy(initial_input), + "seed_obs": deepcopy(initial_input), + "output": [ + { + "prompt_token_ids": [1, 2], + "generation_token_ids": [3], + "generation_log_probs": [-0.1], + } + ], + }, + "responses_create_params": {"input": deepcopy(initial_input)}, + "reward": 1.0, + } + + class _Tokenizer: + def batch_decode(self, batch): + return ["decoded"] * len(batch) + + class _MockSelf: + cfg = {} + _processor = None + + result = ( + NemoGym.__ray_metadata__.modified_class._postprocess_nemo_gym_to_nemo_rl_result( + _MockSelf(), + {}, + nemo_gym_result, + _Tokenizer(), + include_initial_multimodal_data=include_initial_multimodal_data, + ) + ) + + if include_initial_multimodal_data: + assert "_initial_multimodal_data_omitted" not in result + assert "_nemo_rl_initial_media_omitted" not in result["full_result"] + assert data_url in json.dumps(result["full_result"]) + else: + assert result["_initial_multimodal_data_omitted"] is True + assert result["full_result"]["_nemo_rl_initial_media_omitted"] is True + assert data_url not in json.dumps(result["full_result"]) + assert result["full_result"]["responses_create_params"]["input"][0][ + "content" + ] == [{"type": "input_text", "text": "count"}] + + +def test_nemo_gym_dedup_omits_actor_initial_tensor_and_preserves_later_media(): + initial_url = image_to_data_url(Image.new("RGB", (1, 1), color=(1, 0, 0))) + tool_url = image_to_data_url(Image.new("RGB", (1, 1), color=(2, 0, 0))) + initial_input = [ + { + "role": "user", + "content": [ + {"type": "input_text", "text": "inspect"}, + {"type": "input_image", "image_url": initial_url}, + ], + } + ] + template = { + "response": { + "agent_input": deepcopy(initial_input), + "seed_obs": deepcopy(initial_input), + "output": [ + { + "prompt_token_ids": [1], + "generation_token_ids": [2], + "generation_log_probs": [-0.1], + }, + { + "role": "user", + "content": [ + {"type": "input_image", "image_url": tool_url}, + ], + }, + { + "prompt_token_ids": [1, 2, 3], + "generation_token_ids": [4], + "generation_log_probs": [-0.2], + }, + ], + }, + "responses_create_params": {"input": deepcopy(initial_input)}, + "reward": 1.0, + } + + class _Tokenizer: + def batch_decode(self, batch): + return ["decoded"] * len(batch) + + class _ImageProcessor: + model_input_names = ["pixel_values"] + + class _TextTokenizer: + model_input_names = ["input_ids"] + + class _Processor: + image_token = "" + image_processor = _ImageProcessor() + tokenizer = _TextTokenizer() + model_input_names = ["input_ids", "pixel_values"] + + def __call__(self, *, text, images, return_tensors): + assert text == "" * len(images) + assert return_tensors == "pt" + red_values = [image.getpixel((0, 0))[0] for image in images] + return { + "input_ids": torch.tensor([[1]]), + "pixel_values": torch.tensor(red_values, dtype=torch.float32).view( + -1, 1 + ), + } + + class _MockSelf: + cfg = {} + _processor = _Processor() + + postprocess = ( + NemoGym.__ray_metadata__.modified_class._postprocess_nemo_gym_to_nemo_rl_result + ) + flag_off = postprocess( + _MockSelf(), + {}, + deepcopy(template), + _Tokenizer(), + include_initial_multimodal_data=True, + ) + flag_on = postprocess( + _MockSelf(), + {}, + deepcopy(template), + _Tokenizer(), + include_initial_multimodal_data=False, + ) + + off_users = [ + message for message in flag_off["message_log"] if message["role"] == "user" + ] + on_users = [ + message for message in flag_on["message_log"] if message["role"] == "user" + ] + assert off_users[0]["pixel_values"].as_tensor().item() == 1 + assert off_users[1]["pixel_values"].as_tensor().item() == 2 + assert "pixel_values" not in on_users[0] + assert on_users[1]["pixel_values"].as_tensor().item() == 2 + assert initial_url not in json.dumps(flag_on["full_result"]) + assert tool_url in json.dumps(flag_on["full_result"]) + + original_media = PackedTensor(torch.tensor([[99.0]]), dim_to_pack=0) + _reattach_original_multimodal_payloads( + [flag_on], + [[{"role": "user", "content": "", "pixel_values": original_media}]], + ) + on_users = [ + message for message in flag_on["message_log"] if message["role"] == "user" + ] + assert on_users[0]["pixel_values"] is original_media + assert on_users[1]["pixel_values"].as_tensor().item() == 2 + + +@pytest.mark.parametrize( + ("seed_mode", "expected_pixel_values"), + [ + ("text_only", None), + ("initial_plus_additional", [1.0, 2.0]), + ], +) +def test_nemo_gym_dedup_keeps_authoritative_changed_seed_media( + seed_mode, expected_pixel_values +): + initial_url = image_to_data_url(Image.new("RGB", (1, 1), color=(1, 0, 0))) + additional_url = image_to_data_url(Image.new("RGB", (1, 1), color=(2, 0, 0))) + initial_input = [ + { + "role": "user", + "content": [ + {"type": "input_text", "text": "inspect"}, + {"type": "input_image", "image_url": initial_url}, + ], + } + ] + if seed_mode == "text_only": + seed_obs = [ + { + "role": "user", + "content": [{"type": "input_text", "text": "text only"}], + } + ] + else: + seed_obs = deepcopy(initial_input) + seed_obs[0]["content"].append( + {"type": "input_image", "image_url": additional_url} + ) + + nemo_gym_result = { + "response": { + "agent_input": deepcopy(initial_input), + "seed_obs": seed_obs, + "output": [ + { + "prompt_token_ids": [1], + "generation_token_ids": [2], + "generation_log_probs": [-0.1], + } + ], + }, + "responses_create_params": {"input": deepcopy(initial_input)}, + "reward": 1.0, + } + + class _Tokenizer: + def batch_decode(self, batch): + return ["decoded"] * len(batch) + + class _ImageProcessor: + model_input_names = ["pixel_values"] + + class _TextTokenizer: + model_input_names = ["input_ids"] + + class _Processor: + image_token = "" + image_processor = _ImageProcessor() + tokenizer = _TextTokenizer() + model_input_names = ["input_ids", "pixel_values"] + + def __call__(self, *, text, images, return_tensors): + assert text == "" * len(images) + assert return_tensors == "pt" + red_values = [image.getpixel((0, 0))[0] for image in images] + return { + "input_ids": torch.tensor([[1]]), + "pixel_values": torch.tensor(red_values, dtype=torch.float32).view( + -1, 1 + ), + } + + class _MockSelf: + cfg = {} + _processor = _Processor() + + result = ( + NemoGym.__ray_metadata__.modified_class._postprocess_nemo_gym_to_nemo_rl_result( + _MockSelf(), + {}, + nemo_gym_result, + _Tokenizer(), + include_initial_multimodal_data=False, + ) + ) + + assert result["_initial_multimodal_data_omitted"] is False + user_message = next( + message for message in result["message_log"] if message["role"] == "user" + ) + if expected_pixel_values is None: + assert "pixel_values" not in user_message + else: + assert user_message["pixel_values"].as_tensor().flatten().tolist() == ( + expected_pixel_values + ) + + original_media = PackedTensor(torch.tensor([[99.0]]), dim_to_pack=0) + _reattach_original_multimodal_payloads( + [result], + [[{"role": "user", "content": "", "pixel_values": original_media}]], + ) + if expected_pixel_values is None: + assert "pixel_values" not in user_message + else: + assert user_message["pixel_values"].as_tensor().flatten().tolist() == ( + expected_pixel_values + ) + + def test_nemo_gym_postprocess_no_generation_data_raises(): """When no output item carries generation data, the postprocess should raise a ValueError that reports the prompt length and the response.output item types.""" @@ -335,7 +627,7 @@ class _MockSelf: with pytest.raises(ValueError) as excinfo: NemoGym.__ray_metadata__.modified_class._postprocess_nemo_gym_to_nemo_rl_result( - _MockSelf(), nemo_gym_result, _Tokenizer() + _MockSelf(), {}, nemo_gym_result, _Tokenizer() ) msg = str(excinfo.value) @@ -364,7 +656,7 @@ class _MockSelf: with pytest.raises(ValueError) as excinfo: NemoGym.__ray_metadata__.modified_class._postprocess_nemo_gym_to_nemo_rl_result( - _MockSelf(), nemo_gym_result, _Tokenizer() + _MockSelf(), {}, nemo_gym_result, _Tokenizer() ) msg = str(excinfo.value) diff --git a/tests/unit/environments/test_nemo_gym_router_replay.py b/tests/unit/environments/test_nemo_gym_router_replay.py index 15383d8af5b..f60f629b702 100644 --- a/tests/unit/environments/test_nemo_gym_router_replay.py +++ b/tests/unit/environments/test_nemo_gym_router_replay.py @@ -56,7 +56,7 @@ class _MockSelf: result = ( NemoGym.__ray_metadata__.modified_class._postprocess_nemo_gym_to_nemo_rl_result( - _MockSelf(), nemo_gym_result, _Tokenizer() + _MockSelf(), {}, nemo_gym_result, _Tokenizer() ) ) @@ -90,7 +90,7 @@ class _MockSelf: with pytest.raises(ValueError, match="requires NeMo Gym output items"): NemoGym.__ray_metadata__.modified_class._postprocess_nemo_gym_to_nemo_rl_result( - _MockSelf(), nemo_gym_result, _Tokenizer() + _MockSelf(), {}, nemo_gym_result, _Tokenizer() ) @@ -116,7 +116,7 @@ class _MockSelf: result = ( NemoGym.__ray_metadata__.modified_class._postprocess_nemo_gym_to_nemo_rl_result( - _MockSelf(), nemo_gym_result, _Tokenizer() + _MockSelf(), {}, nemo_gym_result, _Tokenizer() ) ) diff --git a/tests/unit/experience/test_rollouts.py b/tests/unit/experience/test_rollouts.py index 038384a5958..2332278d84f 100644 --- a/tests/unit/experience/test_rollouts.py +++ b/tests/unit/experience/test_rollouts.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# 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. @@ -29,9 +29,11 @@ from nemo_rl.data.datasets.response_datasets import NemoGymDataset from nemo_rl.data.interfaces import DatumSpec from nemo_rl.data.llm_message_utils import batched_message_log_to_flat_message +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.distributed.virtual_cluster import RayVirtualCluster +from nemo_rl.environments.interfaces import EnvironmentReturn from nemo_rl.environments.games.sliding_puzzle import ( SlidingPuzzleConfig, SlidingPuzzleEnv, @@ -41,12 +43,16 @@ from nemo_rl.experience.metric_utils import calculate_single_metric, pct from nemo_rl.experience.rollout_manager import AsyncNemoGymRolloutImpl from nemo_rl.experience.rollouts import ( + _add_multimodal_generation_payload, + _reattach_original_multimodal_payloads, + async_generate_response_for_sample_turn, generate_responses_async, run_async_multi_turn_rollout, run_async_multi_turn_rollout_groups, run_async_nemo_gym_rollout, run_multi_turn_rollout, run_nemo_gym_rollout_sync, + run_sample_multi_turn_rollout, ) from nemo_rl.models.generation import configure_generation_config from nemo_rl.models.generation.vllm import VllmConfig, VllmGeneration @@ -67,6 +73,241 @@ _MultiStepCalculatorLogic, ) + +def test_reattach_original_multimodal_payloads_is_media_only_and_turn_aligned(): + first_image = PackedTensor(torch.tensor([[1.0]]), dim_to_pack=0) + second_image = PackedTensor(torch.tensor([[2.0]]), dim_to_pack=0) + original_logs = [ + [ + { + "role": "user", + "content": "first", + "pixel_values": first_image, + "request_metadata": {"must_not": "reattach"}, + }, + {"role": "assistant", "content": "answer"}, + { + "role": "user", + "content": "second", + "pixel_values": second_image, + "vllm_videos": ["video.mp4"], + }, + ] + ] + results = [ + { + "_initial_multimodal_data_omitted": True, + "input_message_log": [ + {"role": "user", "content": "first"}, + {"role": "user", "content": "second"}, + ], + "message_log": [ + {"role": "system", "content": "system"}, + { + "role": "user", + "content": "first", + "pixel_values": "remote placeholder", + }, + {"role": "assistant", "content": "answer"}, + {"role": "user", "content": "second"}, + ], + } + ] + + _reattach_original_multimodal_payloads(results, original_logs) + + for log_key in ("input_message_log", "message_log"): + user_messages = [ + message for message in results[0][log_key] if message["role"] == "user" + ] + assert user_messages[0]["pixel_values"] is first_image + assert user_messages[1]["pixel_values"] is second_image + assert user_messages[1]["vllm_videos"] == ["video.mp4"] + assert "request_metadata" not in user_messages[0] + + +@pytest.mark.parametrize("omission_marker", [False, None]) +def test_reattach_keeps_authoritative_changed_gym_media(omission_marker): + original_media = PackedTensor(torch.tensor([[1.0]]), dim_to_pack=0) + effective_media = PackedTensor(torch.tensor([[2.0]]), dim_to_pack=0) + result = { + "message_log": [ + { + "role": "user", + "content": "", + "pixel_values": effective_media, + } + ], + } + if omission_marker is not None: + result["_initial_multimodal_data_omitted"] = omission_marker + results = [result] + original_logs = [ + [ + { + "role": "user", + "content": "", + "pixel_values": original_media, + } + ] + ] + + _reattach_original_multimodal_payloads(results, original_logs) + + assert results[0]["message_log"][0]["pixel_values"] is effective_media + assert "_initial_multimodal_data_omitted" not in results[0] + + +def test_nemo_gym_initial_media_stays_compact_through_replay_and_policy_flatten(): + generations = 16 + initial_media = PackedTensor(torch.ones(1, 3, 2, 2), dim_to_pack=0) + prompt_batch = BatchedDataDict( + { + "message_log": [ + [ + { + "role": "user", + "content": "", + "token_ids": torch.tensor([], dtype=torch.long), + "pixel_values": initial_media, + } + ] + ] + } + ) + repeated = prompt_batch.repeat_interleave(generations, share_immutable_media=True) + results = [ + { + "_initial_multimodal_data_omitted": True, + "input_message_log": [ + { + "role": "user", + "content": "", + "token_ids": torch.tensor([1]), + } + ], + "message_log": [ + { + "role": "user", + "content": "", + "token_ids": torch.tensor([1]), + }, + { + "role": "assistant", + "content": "", + "token_ids": torch.tensor([2]), + }, + ], + } + for _ in range(generations) + ] + + _reattach_original_multimodal_payloads(results, repeated["message_log"]) + replay_batch = BatchedDataDict( + {"message_log": [result["message_log"] for result in results]} + ) + flat, _ = batched_message_log_to_flat_message(replay_batch["message_log"]) + media = flat["pixel_values"] + + assert media.deduplication_enabled + assert len(media) == generations + assert media.logical_segment_count == generations + assert len(media.tensors) == 1 + assert media.as_tensor().shape == (generations, 3, 2, 2) + + +def test_dedup_generation_sends_only_native_vllm_media(): + class _Generation: + cfg = {"backend": "vllm"} + + pixel_values = PackedTensor(torch.tensor([[1.0]]), dim_to_pack=0) + flat_messages = BatchedDataDict( + { + "token_ids": torch.tensor([[1, 2]]), + "pixel_values": pixel_values, + } + ) + active_batch = BatchedDataDict( + { + "vllm_content": [" describe"], + "vllm_images": [[torch.ones(1, 2)]], + } + ) + compact_generation_input = BatchedDataDict() + + _add_multimodal_generation_payload( + compact_generation_input, + flat_messages, + active_batch, + _Generation(), + deduplicate_multimodal_data=True, + ) + + assert "pixel_values" not in compact_generation_input + assert compact_generation_input["vllm_content"] == [" describe"] + assert compact_generation_input["vllm_images"] is active_batch["vllm_images"] + + later_turn_batch = BatchedDataDict( + { + "vllm_content": [None], + "vllm_images": active_batch["vllm_images"], + } + ) + later_turn_generation_input = BatchedDataDict() + _add_multimodal_generation_payload( + later_turn_generation_input, + flat_messages, + later_turn_batch, + _Generation(), + deduplicate_multimodal_data=True, + ) + assert "pixel_values" not in later_turn_generation_input + assert later_turn_generation_input["vllm_content"] == [None] + assert later_turn_generation_input["vllm_images"] is active_batch["vllm_images"] + + legacy_generation_input = BatchedDataDict() + _add_multimodal_generation_payload( + legacy_generation_input, + flat_messages, + active_batch, + _Generation(), + deduplicate_multimodal_data=False, + ) + assert legacy_generation_input["pixel_values"] is pixel_values + + +def test_dedup_generation_keeps_policy_media_for_unconsumed_native_metadata(): + class _Generation: + cfg = {"backend": "vllm"} + + input_features = PackedTensor(torch.tensor([[1.0]]), dim_to_pack=0) + flat_messages = BatchedDataDict( + { + "token_ids": torch.tensor([[1, 2]]), + "input_features": input_features, + } + ) + active_batch = BatchedDataDict( + { + "vllm_content": [[{"type": "audio", "audio": "/tmp/unconsumed-audio.wav"}]], + "vllm_audio_paths": [["/tmp/unconsumed-audio.wav"]], + } + ) + generation_input = BatchedDataDict() + + _add_multimodal_generation_payload( + generation_input, + flat_messages, + active_batch, + _Generation(), + deduplicate_multimodal_data=True, + ) + + assert generation_input["input_features"] is input_features + assert generation_input["vllm_content"] is active_batch["vllm_content"] + assert "vllm_audio_paths" not in generation_input + + MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct" @@ -141,6 +382,12 @@ def test_median_like_p50(self): class _DummyTokenizer: pad_token_id = 0 + def __call__(self, text, return_tensors=True, add_special_tokens=False): + class _Tokens: + input_ids = torch.tensor([[7]], dtype=torch.int64) + + return _Tokens() + def batch_decode(self, generated_ids, skip_special_tokens=True): return ["ok" for _ in generated_ids] @@ -167,6 +414,297 @@ async def generate_async(self, data, greedy=False): ) +class _CapturingAsyncVllmGeneration: + cfg = {"backend": "vllm", "vllm_cfg": {"async_engine": True}} + + def __init__(self): + self.generation_input = None + + async def generate_async(self, data, greedy=False): + self.generation_input = data + input_length = int(data["input_lengths"][0]) + output_ids = torch.cat( + (data["input_ids"][0, :input_length], torch.tensor([9])) + ).unsqueeze(0) + yield ( + 0, + BatchedDataDict( + { + "output_ids": output_ids, + "logprobs": torch.zeros_like(output_ids, dtype=torch.float32), + "generation_lengths": torch.tensor([1], dtype=torch.long), + "unpadded_sequence_lengths": torch.tensor( + [input_length + 1], dtype=torch.long + ), + "truncated": torch.tensor([False], dtype=torch.bool), + } + ), + ) + + +class _CapturingSyncVllmGeneration: + cfg = {"backend": "vllm"} + + def __init__(self): + self.calls = [] + + def generate(self, data, greedy=False): + self.calls.append( + { + "input_ids": data["input_ids"].clone(), + "input_lengths": data["input_lengths"].clone(), + "vllm_content": list(data["vllm_content"]), + "vllm_images": data["vllm_images"], + } + ) + input_lengths = data["input_lengths"].to(dtype=torch.long) + output_ids = torch.zeros( + (len(input_lengths), int(input_lengths.max().item()) + 1), + dtype=torch.long, + ) + for row, input_length in enumerate(input_lengths.tolist()): + output_ids[row, :input_length] = data["input_ids"][row, :input_length] + output_ids[row, input_length] = 9 + return BatchedDataDict( + { + "output_ids": output_ids, + "logprobs": torch.zeros_like(output_ids, dtype=torch.float32), + "generation_lengths": torch.ones(len(input_lengths), dtype=torch.long), + "unpadded_sequence_lengths": input_lengths + 1, + "truncated": torch.zeros(len(input_lengths), dtype=torch.bool), + } + ) + + +@pytest.mark.parametrize("deduplicate_multimodal_data", [False, True]) +def test_sync_vlm_multiturn_drops_stale_native_content( + monkeypatch, deduplicate_multimodal_data +): + generation = _CapturingSyncVllmGeneration() + image = torch.ones(3, 2, 2) + policy_media = PackedTensor(torch.ones(1, 3, 2, 2), dim_to_pack=0) + reward_calls = [] + + def fake_rewards(batch, task_to_env): + reward_calls.append(None) + return EnvironmentReturn( + observations=[{"role": "user", "content": "next"}], + metadata=[None], + next_stop_strings=[None], + rewards=torch.tensor([0.0]), + terminateds=torch.tensor([len(reward_calls) >= 2]), + answers=[None], + ) + + monkeypatch.setattr( + "nemo_rl.experience.rollouts.calculate_rewards", + fake_rewards, + ) + + run_multi_turn_rollout( + policy_generation=generation, + input_batch=BatchedDataDict( + { + "message_log": [ + [ + { + "role": "user", + "content": "", + "token_ids": torch.tensor([1]), + "pixel_values": policy_media, + } + ] + ], + "extra_env_info": [None], + "task_name": ["vlm"], + "stop_strings": [None], + "idx": [0], + "vllm_content": [" initial prompt"], + "vllm_images": [[image]], + } + ), + tokenizer=_DummyTokenizer(), + task_to_env={}, + max_seq_len=32, + max_rollout_turns=2, + deduplicate_multimodal_data=deduplicate_multimodal_data, + ) + + assert len(generation.calls) == 2 + assert generation.calls[0]["vllm_content"] == [" initial prompt"] + assert generation.calls[1]["vllm_content"] == [None] + assert generation.calls[0]["vllm_images"][0][0] is image + assert generation.calls[1]["vllm_images"][0][0] is image + assert generation.calls[0]["input_ids"][0, :1].tolist() == [1] + assert generation.calls[1]["input_ids"][0, :3].tolist() == [1, 9, 7] + + +def test_async_vlm_generation_receives_exact_compact_native_media_payload(): + generation = _CapturingAsyncVllmGeneration() + policy_media = PackedTensor(torch.ones(1, 3, 2, 2), dim_to_pack=0) + image = torch.ones(3, 2, 2) + audio = torch.ones(16) + video = torch.ones(2, 3, 2, 2) + message_log = [ + { + "role": "user", + "content": "", + "token_ids": torch.tensor([1]), + "pixel_values": policy_media, + } + ] + + asyncio.run( + async_generate_response_for_sample_turn( + generation, + message_log, + None, + _DummyTokenizer(), + max_seq_len=32, + sample_multimodal_data={ + "vllm_content": "