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..3695064845c --- /dev/null +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-click-2n8g-megatron-tp2ep8.v1.yaml @@ -0,0 +1,160 @@ +defaults: ../../vlm_grpo_3B_megatron.yaml +grpo: + num_prompts_per_step: 1 + num_val_generations_per_prompt: 1 + max_num_steps: 100 + val_period: 500 + overlong_filtering: true + max_val_samples: null + val_batch_size: null + seq_logprob_error_threshold: 2 +loss_fn: + reference_policy_kl_penalty: 0 + kl_input_clamp_value: null + kl_output_clamp_value: null + ratio_clip_max: 0.28 + 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 + keep_top_k: 1000000 + checkpoint_must_save_by: 00:03:40:00 +policy: + model_name: nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16 + is_vlm: true + tokenizer: + 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}} + generation_batch_size: 64 + logprob_batch_size: 1 + max_total_sequence_length: 8192 + logprob_chunk_size: 2048 + megatron_cfg: + activation_checkpointing: true + bias_activation_fusion: false + tensor_model_parallel_size: 2 + expert_model_parallel_size: 8 + context_parallel_size: 2 + 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 + defer_fp32_logits: true + track_moe_metrics: true + moe_per_layer_logging: true + radio_force_cpe_eval_mode: true + clear_memory_caches_before_refit: true + optimizer: + lr: 3.0e-06 + min_lr: 3.0e-06 + weight_decay: 0.0 + optimizer_cpu_offload: true + optimizer_offload_fraction: 1.0 + scheduler: + lr_decay_iters: null + lr_warmup_iters: 0 + lr_warmup_init: 3.0e-07 + distributed_data_parallel_config: + overlap_param_gather: false + average_in_collective: false + sequence_packing: + enabled: true + make_sequence_length_divisible_by: 32 + optimizer: null + scheduler: null + generation: + max_new_tokens: ${policy.max_total_sequence_length} + bad_words: [] + mcore_generation_config: + transformer_impl: inference_optimized + 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 + max_tokens: ${policy.max_total_sequence_length} + async_engine: true + expose_http_server: true + enable_prefix_caching: true + parsers: + - deepseek-r1-reasoning + - qwen3-coder-tool + vllm_cfg: + async_engine: true + tensor_parallel_size: 8 + enforce_eager: true + enable_prefix_caching: false + expose_http_server: true + reasoning_parser_plugin: nemo_rl/models/generation/vllm/reasoning_parsers/nano_v3_reasoning_parser.py + 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: + limit_mm_per_prompt: + image: 1 + max_num_batched_tokens: 16384 + mamba_ssm_cache_dtype: float32 + compilation_config: + backend: eager + colocated: + enabled: false + resources: + gpus_per_node: 8 + num_nodes: 1 +data: + max_input_seq_length: null + 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 + prompt_file: null + processor: nemo_gym_data_processor +env: + should_use_nemo_gym: true + should_log_nemo_gym_responses: true + nemo_gym: + is_trajectory_collection: false + port_range_low: 5000 + port_range_high: 5999 + config_paths: + - 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 + wandb_enabled: true + wandb: + 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 +cluster: + gpus_per_node: 8 + num_nodes: 2 diff --git a/examples/nemo_gym/grpo_nanov3.yaml b/examples/nemo_gym/grpo_nanov3.yaml index 103d470f610..3c371660e41 100644 --- a/examples/nemo_gym/grpo_nanov3.yaml +++ b/examples/nemo_gym/grpo_nanov3.yaml @@ -240,7 +240,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/run_grpo_nemo_gym.py b/examples/nemo_gym/run_grpo_nemo_gym.py index 89081a6f948..4d83b9c129f 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") diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index 278abeab7da..cc85990b55d 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -2042,6 +2042,9 @@ 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) + train_data.update(extra_multimodal_data) return train_data @@ -3988,6 +3991,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): diff --git a/nemo_rl/data/llm_message_utils.py b/nemo_rl/data/llm_message_utils.py index d3e1a4df903..c3e6bf76586 100644 --- a/nemo_rl/data/llm_message_utils.py +++ b/nemo_rl/data/llm_message_utils.py @@ -363,9 +363,26 @@ 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/nemo_rl/data/multimodal_utils.py b/nemo_rl/data/multimodal_utils.py index 470618d1e74..368608f5909 100644 --- a/nemo_rl/data/multimodal_utils.py +++ b/nemo_rl/data/multimodal_utils.py @@ -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 @@ -67,6 +74,19 @@ logger = logging.getLogger(__name__) +def uses_image_placeholder(processor: Any) -> bool: + """Return whether a processor requires explicit image placeholders. + + Args: + processor: Multimodal processor to classify. + + Returns: + Whether the processor expands image placeholders through ``__call__`` + rather than tokenized ``apply_chat_template``. + """ + return type(processor).__name__ in _PLACEHOLDER_STYLE_PROCESSOR_NAMES + + class PackedTensor: """Wrapper around a list of torch tensors and a dimension along which to pack the tensors. @@ -376,11 +396,79 @@ 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. + + 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") + return f"data:image/{fmt.lower()};base64,{encoded}" + + +def encode_images_in_examples(nemo_gym_examples: list[dict]) -> list[dict]: + """Replace local image paths in NeMo Gym examples with base64 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", []) + 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..a6e8ef14727 100644 --- a/nemo_rl/data/processors.py +++ b/nemo_rl/data/processors.py @@ -464,6 +464,7 @@ def vlm_hf_data_processor( 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 @@ -555,13 +556,9 @@ def vlm_hf_data_processor( # 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 + uses_placeholder = uses_image_placeholder(processor) - if _uses_image_placeholder and images: + if uses_placeholder and images: # Convert content list to placeholder text format image_token = getattr(processor, "image_token", "") text_parts = [] @@ -582,27 +579,25 @@ def vlm_hf_data_processor( 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, add_generation_prompt=True, ) - # this is the id-tokenized and image processed conversation template for the policy - if _uses_image_placeholder and images: + 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: dict = processor( + message = processor( text=string_formatted_dialog, images=images, return_tensors="pt", ) else: - message: dict = processor.apply_chat_template( + message = processor.apply_chat_template( [user_message_for_tokenize], tokenize=True, add_generation_prompt=True, @@ -620,7 +615,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 @@ -646,7 +641,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 @@ -789,7 +784,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 1b3daba4612..375d3b2d41d 100644 --- a/nemo_rl/environments/nemo_gym.py +++ b/nemo_rl/environments/nemo_gym.py @@ -22,9 +22,18 @@ import ray import torch +from PIL import Image from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy 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.ray_actor_environment_registry import get_actor_python_env from nemo_rl.distributed.virtual_cluster import ( DEFAULT_GYM_PORT_RANGE_HIGH, @@ -33,6 +42,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 @@ -112,6 +122,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[TokenizerConfig] + ] # For processor reconstruction inside the actor def _detect_invalid_tool_call_and_malformed_thinking( @@ -174,12 +188,190 @@ def _detect_invalid_tool_call_and_malformed_thinking( return is_invalid_tool_call, has_malformed_thinking +######################################## +# Multimodal helpers +######################################## + + +# 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://") + + +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. 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) and _looks_like_image_src(src): + images.append(resolve_to_image(src)) + return images + 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( + 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 + 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``). + + ``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" + ): # 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 + + +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", + ) + 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 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) + # _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. @@ -293,6 +485,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 +547,27 @@ 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_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"]["output"], + input_messages=nemo_gym_result.get("responses_create_params", {}).get( + "input" + ), + ) + if processor is not None + else [] + ) + turn_idx = 0 + nemo_rl_message_log = [] seen_token_ids: List[int] = [] batch_decode_items = [] @@ -378,6 +590,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") @@ -437,6 +650,16 @@ 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. @@ -471,6 +694,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( @@ -616,6 +840,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"): + 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], @@ -653,6 +884,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. @@ -668,6 +900,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/megatron/setup.py b/nemo_rl/models/megatron/setup.py index 2776acffbef..e8faea4f9cf 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -1407,6 +1407,14 @@ 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/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 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_llm_message_utils.py b/tests/unit/data/test_llm_message_utils.py index b39f1175934..113fd9ce0b9 100644 --- a/tests/unit/data/test_llm_message_utils.py +++ b/tests/unit/data/test_llm_message_utils.py @@ -772,6 +772,41 @@ 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") 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/environments/test_nemo_gym_mm_utils.py b/tests/unit/environments/test_nemo_gym_mm_utils.py new file mode 100644 index 00000000000..3d8a4186e17 --- /dev/null +++ b/tests/unit/environments/test_nemo_gym_mm_utils.py @@ -0,0 +1,178 @@ +# 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 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_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))), + _assistant([1, 2]), + _user(_image((3, 3)), _image((4, 4))), + _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]] == [(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"}, + _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)] + + +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)]