diff --git a/docs/design-docs/generation.md b/docs/design-docs/generation.md index 6dec60fb1b9..87c40084a58 100644 --- a/docs/design-docs/generation.md +++ b/docs/design-docs/generation.md @@ -143,6 +143,43 @@ The `mcore_generation_config` section controls Megatron Core inference engine be - **num_cuda_graphs**: Number of CUDA graphs to pre-allocate for different batch sizes. More graphs can improve performance by avoiding runtime graph capture, but consume more memory. - **max_tokens**: Maximum total number of tokens (across all requests) that can be processed simultaneously. This limits the maximum batch size and sequence length combinations. Increasing this might throw OOM depending on vocab size and buffer size allocated. +### Multimodal Megatron Generation + +Megatron inference supports image and video inputs in NeMo-RL. Enable multimodal processing with `policy.is_vlm: true`, use the `megatron` generation backend, and provide a `megatron_inference_wrapper`. The wrapper must subclass `megatron.core.inference.model_inference_wrappers.abstract_model_inference_wrapper.AbstractModelInferenceWrapper` in Megatron-Core and declare `supports_ = True` for each supported modality. + +```yaml +policy: + is_vlm: true + generation: + backend: megatron + mcore_generation_config: + megatron_inference_wrapper: megatron.core.inference.model_inference_wrappers.multimodal.nemotron_omni_inference_wrapper.NemotronOmniInferenceWrapper + image_dynamic_resolution: true + video_num_frames: 16 + video_temporal_patch_size: 2 + video_target_num_patches: 2048 + video_maintain_aspect_ratio: true + vision_embedding_cache_max_bytes: 0 + allow_stale_multimodal_embeddings: false +data: + default: + num_frames: 16 + video_temporal_patch_size: 2 + video_target_num_patches: 2048 + video_maintain_aspect_ratio: true +``` + +- `image_dynamic_resolution` preserves variable image shapes instead of forcing one fixed resolution; for example, a wide image uses a wider patch grid than a square image. +- `vision_model_type` optionally selects the MCore vision encoder type used by image and video preprocessing. Set it to the encoder expected by the inference wrapper; when omitted, MCore uses its default (`radio`). +- `num_frames` controls uniform video-frame sampling. Use `video_num_frames` for the corresponding MCore key. +- `video_temporal_patch_size` groups sampled frames into temporal tubelets; for example, size `2` turns 16 frames into 8 temporal groups. +- `video_target_num_patches` sets `num_patches_per_frame = patch_height * patch_width <= video_target_num_patches`, which produces `num_patches_per_frame * num_frames / video_temporal_patch_size` total video patches prior to spatial merging (i.e. further grouped / concatenated into MxM patch blocks) that are provided to the vision encoder. +- `video_maintain_aspect_ratio=true` keeps `patch_width / patch_height ~= source_width / source_height`; `false` uses `patch_width = patch_height ~= sqrt(video_target_num_patches)` (for example, `sqrt(256) = 16`). +- `vision_embedding_cache_max_bytes` limits GPU memory used to reuse vision embeddings for repeated media; `0` disables the cache, while `1073741824` permits up to 1 GiB. +- `allow_stale_multimodal_embeddings` controls whether cached embeddings survive model-weight changes. Keep it `false` for RL refits; use `true` only when weights remain fixed. +- `expose_http_server` should be `true` for NeMo Gym. + +Keep the video preprocessing values identical in `data.default` and `mcore_generation_config` to avoid disparity between the training policy and inference generation. ## Usage Examples diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-16n8g-megatron-tp4ep4-async-gym-video.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-16n8g-megatron-tp4ep4-async-gym-video.v1.yaml index 2647cfc4d8a..20799328e5d 100644 --- a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-16n8g-megatron-tp4ep4-async-gym-video.v1.yaml +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-16n8g-megatron-tp4ep4-async-gym-video.v1.yaml @@ -48,6 +48,10 @@ policy: make_sequence_length_divisible_by: 32 generation: bad_words: [] + mcore_generation_config: + image_dynamic_resolution: true + logprobs_mode: raw_logprobs + megatron_inference_wrapper: megatron.core.inference.model_inference_wrappers.multimodal.nemotron_omni_inference_wrapper.NemotronOmniInferenceWrapper vllm_cfg: reset_encoder_cache_after_weight_update: false video: 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 38ff8f89fe6..2f198cb56c6 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 @@ -90,6 +90,9 @@ policy: max_tokens: ${policy.max_total_sequence_length} expose_http_server: true enable_prefix_caching: true + image_dynamic_resolution: true + logprobs_mode: raw_logprobs + megatron_inference_wrapper: megatron.core.inference.model_inference_wrappers.multimodal.nemotron_omni_inference_wrapper.NemotronOmniInferenceWrapper parsers: - deepseek-r1-reasoning - qwen3-coder-tool diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-count-1n4g-megatron_generation.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-count-1n4g-megatron_generation.v1.yaml new file mode 100644 index 00000000000..aca1e6015a6 --- /dev/null +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-count-1n4g-megatron_generation.v1.yaml @@ -0,0 +1,96 @@ +defaults: ../../vlm_grpo_3B_megatron.yaml +grpo: + num_prompts_per_step: 2 + num_generations_per_prompt: 8 + max_num_steps: 4 + val_period: 0 + max_val_samples: null + val_batch_size: null + async_grpo: + enabled: true + max_trajectory_age_steps: 2 + in_flight_weight_updates: true +loss_fn: + reference_policy_kl_penalty: 0.0 + use_importance_sampling_correction: true +checkpointing: + enabled: false + checkpoint_dir: results/nemo-rl-omni/nemotron-omni-circle-count-1n4g +policy: + model_name: nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16 + is_vlm: true + train_global_batch_size: 16 + logprob_batch_size: 1 + max_total_sequence_length: 8192 + sequence_packing: + enabled: true + megatron_cfg: + env_vars: + TORCH_CUDA_ARCH_LIST: '10.0' + tensor_model_parallel_size: 2 + expert_model_parallel_size: 2 + sequence_parallel: true + bias_activation_fusion: false + activation_checkpointing: true + generation: + backend: megatron + bad_words: null + mcore_generation_config: + expose_http_server: true + buffer_size_gb: 8 + num_cuda_graphs: -1 + max_tokens: ${policy.max_total_sequence_length} + transformer_impl: transformer_engine + activation_checkpointing: false + tensor_model_parallel_size: 2 + expert_model_parallel_size: 2 + expert_tensor_parallel_size: 1 + sequence_parallel: true + moe_pad_experts_for_cuda_graph_inference: true + image_dynamic_resolution: true + logprobs_mode: raw_logprobs + megatron_inference_wrapper: megatron.core.inference.model_inference_wrappers.multimodal.nemotron_omni_inference_wrapper.NemotronOmniInferenceWrapper + vllm_cfg: + async_engine: true + expose_http_server: true + colocated: + enabled: false + resources: + gpus_per_node: 2 + num_nodes: 1 +data: + _override_: true + max_input_seq_length: null + shuffle: false + num_workers: 0 + train: + data_path: 3rdparty/Gym-workspace/Gym/resources_servers/circle_count/data/example.jsonl + validation: + data_path: 3rdparty/Gym-workspace/Gym/resources_servers/circle_count/data/example.jsonl + default: + dataset_name: NemoGymDataset + env_name: nemo_gym + prompt_file: null + processor: nemo_gym_data_processor +env: + _override_: true + 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_count/configs/circle_count.yaml + circle_count_simple_agent: + responses_api_agents: + simple_agent: + max_steps: 1 +logger: + tensorboard_enabled: false + wandb: + project: nemo-rl-omni + name: nemotron-omni-circle-count-1n4g +cluster: + gpus_per_node: 4 diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron_generation.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron_generation.v1.yaml new file mode 100644 index 00000000000..0f741856edf --- /dev/null +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron_generation.v1.yaml @@ -0,0 +1,77 @@ +defaults: ../../vlm_grpo_3B_megatron.yaml +grpo: + num_prompts_per_step: 12 + num_generations_per_prompt: 8 + val_at_start: true + val_at_end: true + val_batch_size: 64 + async_grpo: + enabled: true + max_trajectory_age_steps: 2 + in_flight_weight_updates: true +loss_fn: + reference_policy_kl_penalty: 0.0 + use_importance_sampling_correction: true +checkpointing: + enabled: false + checkpoint_dir: results/nemo-rl-omni/nemotron-omni-clevr-megatron-8n4g +policy: + model_name: nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16 + is_vlm: true + train_global_batch_size: ${mul:${grpo.num_prompts_per_step}, ${grpo.num_generations_per_prompt}} + logprob_batch_size: 1 + max_total_sequence_length: 4096 + sequence_packing: + enabled: true + megatron_cfg: + env_vars: + TORCH_CUDA_ARCH_LIST: '10.0' + tensor_model_parallel_size: 8 + expert_model_parallel_size: 8 + sequence_parallel: true + bias_activation_fusion: false + activation_checkpointing: true + optimizer: + exp_avg_dtype: bfloat16 + exp_avg_sq_dtype: bfloat16 + store_param_remainders: true + generation: + backend: megatron + max_new_tokens: 2048 + bad_words: null + mcore_generation_config: + buffer_size_gb: 8 + async_sched_mode: async + num_cuda_graphs: -1 + use_cuda_graphs_for_non_decode_steps: false + max_tokens: ${policy.max_total_sequence_length} + transformer_impl: inference_optimized + activation_checkpointing: false + tensor_model_parallel_size: 8 + expert_model_parallel_size: 8 + expert_tensor_parallel_size: 1 + context_parallel_size: 1 + sequence_parallel: true + moe_router_dtype: fp32 + moe_pad_experts_for_cuda_graph_inference: false + mamba_inference_ssm_states_dtype: float32 + mamba_inference_conv_states_dtype: float32 + image_dynamic_resolution: true + logprobs_mode: raw_logprobs + megatron_inference_wrapper: megatron.core.inference.model_inference_wrappers.multimodal.nemotron_omni_inference_wrapper.NemotronOmniInferenceWrapper + colocated: + enabled: false + resources: + gpus_per_node: 4 + num_nodes: 6 +data: + default: + prompt_file: examples/prompts/clevr_cogent_cot_nemotron_omni.txt +logger: + tensorboard_enabled: false + wandb: + project: nemo-rl-omni + name: nemotron-omni-clevr-megatron-8n4g +cluster: + gpus_per_node: 4 + num_nodes: 8 diff --git a/examples/nemo_gym/prepare_video_dataset.py b/examples/nemo_gym/prepare_video_dataset.py index 7d18a3e0685..3976c03ae22 100644 --- a/examples/nemo_gym/prepare_video_dataset.py +++ b/examples/nemo_gym/prepare_video_dataset.py @@ -256,7 +256,7 @@ def convert(args: argparse.Namespace) -> None: row["responses_create_params"] = { "input": input_messages, "metadata": { - "chat_template_kwargs": {"enable_thinking": True}, + "chat_template_kwargs": json.dumps({"enable_thinking": True}), }, } raw_answer = source_row.get("answer") diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index 66ec0b5bc5e..68c4c903c87 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -77,7 +77,6 @@ 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 -from nemo_rl.distributed.ray_actor_environment_registry import get_actor_python_env from nemo_rl.distributed.virtual_cluster import ( TOPO_RANK_UNKNOWN, ClusterConfig, @@ -158,7 +157,7 @@ ) 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 +from nemo_rl.utils.venvs import make_actor_runtime_env from nemo_rl.weight_sync.checkpoint_engine_config import ( checkpoint_engine_refit_config, ) @@ -4553,28 +4552,9 @@ def async_grpo_train( print(f" - train_global_batch_size: {train_gbs}") print(f" - min_trajectories_needed: {min_trajectories_needed} (async mode)") - _replay_py_exec = get_actor_python_env( + _replay_runtime_env = make_actor_runtime_env( "nemo_rl.algorithms.async_utils.ReplayBuffer" ) - if _replay_py_exec.startswith("uv"): - # Lazily build a dedicated venv across all Ray nodes on-demand. - _replay_py_exec = create_local_venv_on_each_node( - _replay_py_exec, - "nemo_rl.algorithms.async_utils.ReplayBuffer", - ) - - _replay_py_venv = os.path.dirname( - os.path.dirname(_replay_py_exec) - ) # to remove the "bin/python" suffix - - _replay_runtime_env = { - "py_executable": _replay_py_exec, - "env_vars": { - **os.environ, - "VIRTUAL_ENV": _replay_py_venv, - "UV_PROJECT_ENVIRONMENT": _replay_py_venv, - }, - } # Calculate optimal buffer size based on generation limits to prevent length bias # Each weight version generates exactly num_prompts_per_step trajectories @@ -4681,29 +4661,13 @@ def async_grpo_train( set(trained_task_indices) if frontier_restore else set() ) - _tc_py_exec = get_actor_python_env( - "nemo_rl.algorithms.async_utils.AsyncTrajectoryCollector" - ) - if _tc_py_exec.startswith("uv"): - _tc_py_exec = create_local_venv_on_each_node( - _tc_py_exec, - "nemo_rl.algorithms.async_utils.AsyncTrajectoryCollector", - ) - - _tc_py_venv = os.path.dirname( - os.path.dirname(_tc_py_exec) - ) # to remove the "bin/python" suffix - - _tc_runtime_env = { - "py_executable": _tc_py_exec, - "env_vars": { - **os.environ, - "VIRTUAL_ENV": _tc_py_venv, - "UV_PROJECT_ENVIRONMENT": _tc_py_venv, + _tc_runtime_env = make_actor_runtime_env( + "nemo_rl.algorithms.async_utils.AsyncTrajectoryCollector", + extra_env_vars={ # Names this actor's spans the way RayWorkerGroup names its groups'. "NRL_WORKER_GROUP": "trajectory_collector", }, - } + ) # Captured inside rl.grpo.job, so the collector's spans join this run's # trace instead of starting their own roots. Empty unless the job group is diff --git a/nemo_rl/data/collate_fn.py b/nemo_rl/data/collate_fn.py index 86f91b247ee..bcd5fbbf753 100644 --- a/nemo_rl/data/collate_fn.py +++ b/nemo_rl/data/collate_fn.py @@ -45,11 +45,10 @@ def rl_collate_fn(data_batch: list[DatumSpec]) -> BatchedDataDict[Any]: # Extract stop_strings if present stop_strings = [datum.get("stop_strings", None) for datum in data_batch] - # check if any of the data batch has vllm content and images + # Presence of the key selects vLLM's native-media path. Placeholder-style + # processors intentionally set the content to None so vLLM uses input_ids. extra_args = {} - if any( - [datum_spec.get("vllm_content", None) is not None for datum_spec in data_batch] - ): + if any("vllm_content" in datum_spec for datum_spec in data_batch): vllm_content = [ datum_spec.get("vllm_content", None) for datum_spec in data_batch ] @@ -119,11 +118,10 @@ def eval_collate_fn(data_batch: list[DatumSpec]) -> BatchedDataDict[Any]: idx = [datum_spec["idx"] for datum_spec in data_batch] task_names = [datum_spec.get("task_name", None) for datum_spec in data_batch] - # Check if any of the data batch has vllm content (multimodal data) + # Preserve native media when placeholder-style processors intentionally + # set vllm_content to None in favor of their expanded input_ids. extra_args = {} - if any( - datum_spec.get("vllm_content", None) is not None for datum_spec in data_batch - ): + if any("vllm_content" in datum_spec for datum_spec in data_batch): extra_args["vllm_content"] = [ datum_spec.get("vllm_content", None) for datum_spec in data_batch ] diff --git a/nemo_rl/data/multimodal_utils.py b/nemo_rl/data/multimodal_utils.py index 2a746a36097..67878e5278a 100644 --- a/nemo_rl/data/multimodal_utils.py +++ b/nemo_rl/data/multimodal_utils.py @@ -19,9 +19,9 @@ import uuid from collections import defaultdict from collections.abc import Sequence -from concurrent.futures import ThreadPoolExecutor from copy import deepcopy from io import BytesIO +from pathlib import Path from typing import Any, Optional, Union import requests @@ -40,7 +40,6 @@ MULTIMODAL_CONTENT_TYPES = frozenset( {*IMAGE_CONTENT_TYPES, *VIDEO_CONTENT_TYPES, *AUDIO_CONTENT_TYPES} ) -NEMO_GYM_IMAGE_ENCODE_MAX_WORKERS = 8 # List of allowed placeholder strings for different media types in the dataset string # e.g. "This is an example of " @@ -51,6 +50,8 @@ "video-audio": "", } MEDIA_TAGS_REVERSED = {v: k for k, v in MEDIA_TAGS.items()} +CACHED_VIDEO_FRAME_MANIFEST_MAGIC = b"NEMO_RL_CACHED_VIDEO_FRAMES_V1\n" +CACHED_VIDEO_FRAME_MANIFEST_MIME = "video/x-nemo-rl-cached-frames" DEFAULT_MEDIA_EXTENSIONS = { "image": ["png", "jpeg", "jpg", "img"], @@ -890,24 +891,25 @@ def image_to_data_url(image: Image.Image, fmt: str = "PNG") -> str: return f"data:image/{fmt.lower()};base64,{encoded}" -def _encode_single_image_source(source: str) -> str: - """Resolve and encode one image source.""" - image = resolve_to_image(source) - try: - data_url = image_to_data_url(image) - finally: - image.close() - return data_url +def get_responses_content_part_url(part: dict[str, Any], *keys: str) -> str: + """Return a string media source from a Responses/Chat content part.""" + for key in keys: + value = part.get(key) + if isinstance(value, dict): + value = value.get("url") or value.get("path") + if isinstance(value, str) and value: + return value + return "" -def extract_input_image_sources_from_responses_messages( +def extract_input_media_sources_from_responses_messages( messages: Any, -) -> list[str | Image.Image]: - """Extract image sources from Responses-API messages in encounter order.""" +) -> list[tuple[str, Any]]: + """Extract tagged image and video sources in encounter order.""" if not isinstance(messages, list): return [] - sources: list[str | Image.Image] = [] + sources: list[tuple[str, Any]] = [] for message in messages: if not isinstance(message, dict): continue @@ -917,24 +919,37 @@ def extract_input_image_sources_from_responses_messages( for part in content: if not isinstance(part, dict): continue - if part.get("type") not in ("input_image", "image", "image_url"): + part_type = part.get("type") + if part_type in IMAGE_CONTENT_TYPES: + media_type = "image" + source = part.get("image") or part.get("image_url") or part.get("url") + elif part_type in VIDEO_CONTENT_TYPES: + media_type = "video" + source = part.get("video") or part.get("video_url") or part.get("url") + else: continue - source = part.get("image") or part.get("image_url") or part.get("url") if isinstance(source, dict): - source = source.get("url") + source = source.get("url") or source.get("path") + # Skip non-str/non-Image sources: callers hand these straight to + # `resolve_to_image`, which would raise on e.g. an int `image_url`. if isinstance(source, (str, Image.Image)): - sources.append(source) + sources.append((media_type, 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 media_sources_equal( + left: tuple[str, Any], + right: tuple[str, Any], +) -> bool: + """Compare tagged media by string value or object identity.""" + if left[0] != right[0]: + return False + left_source, right_source = left[1], right[1] + return ( + left_source == right_source + if isinstance(left_source, str) and isinstance(right_source, str) + else left_source is right_source + ) def _materialize_ragged_pixel_values( @@ -1042,82 +1057,41 @@ def attach_image_model_inputs_to_message( ) -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, - collects local image references, encodes each unique source once using a - bounded thread pool, and rewrites every corresponding image part with the - resulting base64 ``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. +_VIDEO_EXT_TO_MIME = { + ".mp4": "mp4", + ".m4v": "mp4", + ".mov": "quicktime", + ".webm": "webm", + ".mkv": "x-matroska", + ".avi": "x-msvideo", +} - 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. +def video_path_to_data_url(video_path: str) -> str: + """Inline a local or ``file://`` video as a base64 data URL.""" + if video_path.startswith("data:"): + return video_path - Returns: - The same ``nemo_gym_examples`` list, with local image references - rewritten to base64 data URLs in place. - """ - targets_by_source: dict[str, list[tuple[dict, str]]] = {} - - 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") not in IMAGE_CONTENT_TYPES - ): - continue - media_key = next( - (key for key in ("image_url", "image", "url") if key in part), - None, - ) - if media_key is None: - continue - url = part.get(media_key) - if isinstance(url, dict): - url = url.get("url") or url.get("path") or "" - if not isinstance(url, str) or not url: - continue - if url.startswith(("http://", "https://", "data:")): - continue - targets_by_source.setdefault(url, []).append((part, media_key)) - - sources = list(targets_by_source) - if sources: - with ThreadPoolExecutor( - max_workers=NEMO_GYM_IMAGE_ENCODE_MAX_WORKERS - ) as executor: - encoded_by_source = dict( - zip( - sources, - executor.map(_encode_single_image_source, sources), - strict=True, - ) - ) + resolved = ( + video_path.removeprefix("file://") + if video_path.startswith("file://") + else str(Path(video_path).expanduser().resolve()) + ) + path = Path(resolved) + if not path.is_file(): + raise FileNotFoundError( + f"Video path resolved to {resolved!r}, which does not exist." + ) - # Keep payload mutation on the caller thread after worker-owned images - # have been closed and every unique source has been encoded. - for source, targets in targets_by_source.items(): - data_url = encoded_by_source[source] - for part, media_key in targets: - part[media_key] = data_url - return nemo_gym_examples + ext = path.suffix.lower() + mime = _VIDEO_EXT_TO_MIME.get(ext) + if mime is None: + raise ValueError( + f"Unsupported video extension {ext!r} for {resolved!r}. " + f"Supported: {sorted(_VIDEO_EXT_TO_MIME)}." + ) + encoded = base64.b64encode(path.read_bytes()).decode("ascii") + return f"data:video/{mime};base64,{encoded}" def get_media_from_message(message: dict[str, Any]) -> dict[str, list[Any]]: diff --git a/nemo_rl/data/processors.py b/nemo_rl/data/processors.py index 58f09c95856..ef8744afb7b 100644 --- a/nemo_rl/data/processors.py +++ b/nemo_rl/data/processors.py @@ -490,6 +490,8 @@ def vlm_hf_data_processor( pass # Daily-Omni data is already formatted by DailyOmniDataset.format_data elif datum_dict["task_name"] in ("intent-train", "intent-bench"): pass # IntentDataset.format_data already produces the message structure + elif "messages" in datum_dict: + pass # Generic ResponseDataset data can already use the message structure else: raise ValueError(f"No data processor for task {datum_dict['task_name']}") @@ -676,8 +678,11 @@ def vlm_hf_data_processor( loss_multiplier = 0.0 else: # get the prompt content! (use this for vllm-backend that needs formatted dialog and list of images/audios) for the entire conversation + # Placeholder-style processors set vllm_content to None so vLLM uses expanded input_ids. vllm_kwargs = { - "vllm_content": string_formatted_dialog, + "vllm_content": ( + None if uses_placeholder and images else string_formatted_dialog + ), "vllm_images": images, "vllm_audios": audios, "vllm_videos": videos, @@ -799,7 +804,7 @@ def nemo_gym_data_processor( "Gym video data requires a multimodal processor with " "apply_chat_template and tokenizer attributes" ) - from nemo_rl.environments.nemo_gym_video import ( + from nemo_rl.environments.nemo_gym_multimodal import ( nemo_gym_example_to_video_datum_spec, ) diff --git a/nemo_rl/distributed/held_port.py b/nemo_rl/distributed/held_port.py index 510a40af714..18d68be06b6 100644 --- a/nemo_rl/distributed/held_port.py +++ b/nemo_rl/distributed/held_port.py @@ -18,6 +18,8 @@ from nemo_rl.distributed.virtual_cluster import _get_node_ip_local +_HANDOFF_RELEASED = b"\x01" + def _held_port_uds_name(port: int) -> str: """Abstract-namespace Unix socket where a HeldPortReservation serves its fd.""" @@ -34,19 +36,35 @@ def receive_held_socket(port: int) -> socket.socket: The live listening socket, duplicated into this process. """ client = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + fds: list[int] = [] + received_socket: socket.socket try: client.connect(_held_port_uds_name(port)) _, fds, _, _ = socket.recv_fds(client, 1024, 1) + if not fds: + raise RuntimeError(f"Port holder for port {port} sent no file descriptor.") + # The receiving process now owns a duplicate of the reservation socket, + # but callers such as MCore close it before binding their own + # SO_REUSEPORT listeners. Wait until the holder has closed its original + # descriptor so those listeners cannot race the old, non-reusable + # reservation socket and fail with EADDRINUSE. + released = client.recv(1) + if released != _HANDOFF_RELEASED: + socket.close(fds.pop()) + raise RuntimeError( + f"Port holder for port {port} did not confirm releasing its socket." + ) + received_socket = socket.socket(fileno=fds.pop()) except OSError as e: + for fd in fds: + socket.close(fd) raise RuntimeError( f"Could not receive the reserved server socket for port {port}: " "the port holder on this node is gone, so the pre-published URL would be unreachable." ) from e finally: client.close() - if not fds: - raise RuntimeError(f"Port holder for port {port} sent no file descriptor.") - return socket.socket(fileno=fds[0]) + return received_socket class HeldPortReservation: @@ -58,6 +76,11 @@ class HeldPortReservation: def __init__(self) -> None: self._sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + # MCore starts multiple frontend replicas that each bind this port with + # SO_REUSEPORT. Make the reservation socket part of the same reuse group + # so their binds remain valid while the handed-off fd is being closed. + self._sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + self._sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1) self._sock.bind(("", 0)) self._sock.listen(128) self._port = self._sock.getsockname()[1] @@ -75,10 +98,14 @@ def _serve_fd_once(self) -> None: conn, _ = self._uds.accept() try: socket.send_fds(conn, [b"s"], [self._sock.fileno()]) + # The receiver holds a duplicate fd, so the port remains reserved. + # Close this copy before acknowledging the handoff; the receiver may + # immediately close its copy and rebind the port with SO_REUSEPORT. + self._sock.close() + conn.sendall(_HANDOFF_RELEASED) finally: conn.close() self._uds.close() - # The receiver holds a duplicate fd; the local one is done. self._sock.close() diff --git a/nemo_rl/distributed/worker_groups.py b/nemo_rl/distributed/worker_groups.py index 7d88794f549..3cfd2d405ca 100644 --- a/nemo_rl/distributed/worker_groups.py +++ b/nemo_rl/distributed/worker_groups.py @@ -32,10 +32,22 @@ from nemo_rl.distributed.virtual_cluster import RayVirtualCluster from nemo_rl.distributed.worker_group_utils import recursive_merge_options from nemo_rl.utils.venvs import ( + add_hf_modules_cache_to_pythonpath, create_local_venv_on_each_node, ) +def _get_initializer_env_vars(env_vars: dict[str, str]) -> dict[str, str]: + """Build the environment needed to unpickle worker constructor arguments.""" + initializer_env_vars = { + key: env_vars[key] if key in env_vars else os.environ[key] + for key in ("HF_HOME", "HF_MODULES_CACHE", "PYTHONPATH") + if key in env_vars or key in os.environ + } + + return add_hf_modules_cache_to_pythonpath(initializer_env_vars) + + @dataclass class MultiWorkerFuture: """Container for Ray futures with associated worker information.""" @@ -502,11 +514,7 @@ def _create_workers_from_bundle_indices( # import-related variables that trust_remote_code classes need to # resolve their generated modules have to travel with it. unique_pg_indices = sorted({pg_idx for pg_idx, _ in bundle_indices_list}) - initializer_env_vars = { - key: env_vars[key] - for key in ("HF_HOME", "HF_MODULES_CACHE", "PYTHONPATH") - if key in env_vars - } + initializer_env_vars = _get_initializer_env_vars(env_vars) initializer_runtime_env = {} if py_executable != sys.executable: initializer_runtime_env["py_executable"] = py_executable diff --git a/nemo_rl/environments/nemo_gym.py b/nemo_rl/environments/nemo_gym.py index e702f462042..0f68069db92 100644 --- a/nemo_rl/environments/nemo_gym.py +++ b/nemo_rl/environments/nemo_gym.py @@ -17,21 +17,18 @@ import sys from collections import Counter from collections.abc import AsyncGenerator, Mapping -from copy import deepcopy from pathlib import Path from typing import Any, Dict, List, NotRequired, Optional, Protocol, TypedDict 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 ( attach_image_model_inputs_to_message, - encode_images_in_examples, - extract_input_image_sources_from_responses_messages, - resolve_to_image, + extract_input_media_sources_from_responses_messages, + media_sources_equal, uses_image_placeholder, ) from nemo_rl.distributed.ray_actor_environment_registry import get_actor_python_env @@ -42,6 +39,12 @@ _get_node_ip_local, ) from nemo_rl.environments.interfaces import EnvironmentInterface +from nemo_rl.environments.nemo_gym_multimodal import ( + _index_per_turn_images, + _is_trainable_output_item, + _without_initial_media_sources, + normalize_media_in_examples, +) from nemo_rl.experience.failures import ( GymTransportError, RolloutDataFailure, @@ -282,28 +285,6 @@ 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 get_pad_dynamic_image_shapes(env_config: Mapping[str, Any]) -> bool: """Return nemo_gym's pad_dynamic_image_shapes from an env config, or False. @@ -322,157 +303,6 @@ def get_pad_dynamic_image_shapes(env_config: Mapping[str, Any]) -> bool: return bool(nemo_gym_config.get("pad_dynamic_image_shapes")) -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 _is_trainable_output_item(item: dict) -> bool: - """Report whether an output item becomes a trainable assistant turn. - - The postprocess loop skips items whose ``generation_token_ids`` is missing - *or* empty, so per-turn image binning has to use the same predicate or the - two walks disagree and every later turn gets the wrong images. - """ - return bool(item.get("generation_token_ids")) - - -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 _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, - *, - images: list[Image.Image], - processor: Any, - pad_dynamic_image_shapes: bool = False, -) -> 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 ````. - """ - attach_image_model_inputs_to_message( - user_message, - images=images, - processor=processor, - pad_dynamic_image_shapes=pad_dynamic_image_shapes, - ) - - @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.""" @@ -500,7 +330,7 @@ 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 + # attach_image_model_inputs_to_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. @@ -508,7 +338,7 @@ def __init__(self, cfg: NemoGymConfig): "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." + "attach_image_model_inputs_to_message before enabling." ) def _require_spinup(self) -> None: @@ -675,14 +505,10 @@ async def run_rollouts( timer = Timer() counts_left = Counter(row["agent_ref"]["name"] for row in nemo_gym_examples) - from nemo_rl.environments.nemo_gym_video import ( - normalize_video_urls_in_examples, - ) - - # Normalize local media before shipping requests to vLLM. Both helpers - # are no-ops for text-only rows and already-qualified URLs. - normalize_video_urls_in_examples(nemo_gym_examples) - encode_images_in_examples(nemo_gym_examples) + # Normalize local media before shipping requests to vLLM. Helper is a no-op + # for text-only rows and already-qualified URLs. + # Megatron's HTTP backend consumes the same normalized Responses payload. + normalize_media_in_examples(nemo_gym_examples) timer.start("_run_rollouts_total") nemo_gym_result_iterator = self.rch.run_examples( @@ -778,20 +604,20 @@ def _postprocess_nemo_gym_to_nemo_rl_result( 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_initial_sources = extract_input_media_sources_from_responses_messages( raw_input ) - agent_initial_sources = extract_input_image_sources_from_responses_messages( + agent_initial_sources = extract_input_media_sources_from_responses_messages( initial_input ) - returned_media_sources = extract_input_image_sources_from_responses_messages( + returned_media_sources = extract_input_media_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) + media_sources_equal(agent_source, raw_source) for agent_source, raw_source in zip( agent_initial_sources, raw_initial_sources ) @@ -800,7 +626,7 @@ def _postprocess_nemo_gym_to_nemo_rl_result( returned_media_matches_raw_input = len(returned_media_sources) == len( raw_initial_sources ) and all( - _image_sources_equal(returned_source, raw_source) + media_sources_equal(returned_source, raw_source) for returned_source, raw_source in zip( returned_media_sources, raw_initial_sources ) @@ -811,7 +637,7 @@ def _postprocess_nemo_gym_to_nemo_rl_result( and returned_media_matches_raw_input ) if initial_multimodal_data_omitted: - media_messages, _ = _without_initial_image_sources( + media_messages, _ = _without_initial_media_sources( media_messages, raw_initial_sources ) per_turn_images = ( @@ -908,7 +734,7 @@ def _postprocess_nemo_gym_to_nemo_rl_result( images_this_turn = ( per_turn_images[turn_idx] if turn_idx < len(per_turn_images) else [] ) - _attach_multimodal_data_to_user_message( + attach_image_model_inputs_to_message( user_message, images=images_this_turn, processor=processor, @@ -1003,7 +829,7 @@ def _postprocess_nemo_gym_to_nemo_rl_result( (response, "seed_obs"), ): if key in container: - container[key], _ = _without_initial_image_sources( + container[key], _ = _without_initial_media_sources( container[key], raw_initial_sources ) @@ -1112,9 +938,16 @@ def validate_reward_components_match_scalar(nemo_gym_results: List[dict]) -> Non def setup_nemo_gym_config(config, tokenizer) -> None: generation_config = config.policy["generation"] - # Enable the http server. Requires both async engine and the expose_http_server flag - generation_config["vllm_cfg"]["async_engine"] = True - generation_config["vllm_cfg"]["expose_http_server"] = True + backend = generation_config.get("backend") + if backend == "vllm": + # Enable the http server. Requires both async engine and the expose_http_server flag + generation_config["vllm_cfg"]["async_engine"] = True + generation_config["vllm_cfg"]["expose_http_server"] = True + elif backend == "megatron": + # Enable the http server for Gym dispatch over the Megatron generation backend. + generation_config["mcore_generation_config"]["expose_http_server"] = True + else: + raise ValueError(f"NeMo Gym does not support generation backend {backend!r}.") # Stop strings or token ids are not supported generation_config["stop_strings"] = None diff --git a/nemo_rl/environments/nemo_gym_video.py b/nemo_rl/environments/nemo_gym_multimodal.py similarity index 67% rename from nemo_rl/environments/nemo_gym_video.py rename to nemo_rl/environments/nemo_gym_multimodal.py index 83d3f8bd9ea..3e65b88ba69 100644 --- a/nemo_rl/environments/nemo_gym_video.py +++ b/nemo_rl/environments/nemo_gym_multimodal.py @@ -15,6 +15,9 @@ import copy import json import os +from concurrent.futures import ThreadPoolExecutor +from contextlib import closing +from copy import deepcopy from pathlib import Path from typing import Any, TypeVar, cast from urllib.parse import unquote, urlparse @@ -28,9 +31,19 @@ IMAGE_CONTENT_TYPES, VIDEO_CONTENT_TYPES, PackedTensor, + extract_input_media_sources_from_responses_messages, extract_multimodal_model_inputs, get_dim_to_pack_along, + get_responses_content_part_url, + image_to_data_url, + media_sources_equal, resolve_to_image, + video_path_to_data_url, +) +from nemo_rl.environments.nemo_gym_request import ( + _chat_template_kwargs_for_processor, + _deep_merge_dict, + _json_mapping, ) from nemo_rl.environments.nemotron_utils import ( NEMOTRON_VIDEO_PROCESSOR_NAMES, @@ -43,6 +56,265 @@ load_video_frames_with_metadata, ) +_NEMO_GYM_IMAGE_ENCODE_MAX_WORKERS = 8 + + +def _encode_single_image_source(source: str) -> str: + """Resolve, encode, and close one local image source.""" + # `closing` (not a bare `with`): PIL's Image.__exit__ is a no-op, so only an + # explicit close() releases the buffer. + with closing(resolve_to_image(source)) as image: + return image_to_data_url(image) + + +def normalize_media_in_examples(nemo_gym_examples: list[dict]) -> list[dict]: + """Replace local media paths in NeMo Gym examples with data URLs.""" + local_image_sources: dict[str, None] = {} + local_video_sources: dict[str, None] = {} + pending_mutations: list[ + tuple[dict, tuple[str, str, str], str, str, bool, Any, str] + ] = [] + 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): + continue + part_type = part.get("type") + if part_type in IMAGE_CONTENT_TYPES: + source_keys = ("image_url", "image", "url") + canonical_type = "input_image" + canonical_key = "image_url" + is_image = True + elif part_type in VIDEO_CONTENT_TYPES: + source_keys = ("video_url", "video", "url") + canonical_type = "input_video" + canonical_key = "video_url" + is_image = False + else: + continue + + present_keys = [key for key in source_keys if key in part] + if ( + not present_keys + and part_type == "input_image" + and "file_id" in part + ): + continue + if len(present_keys) != 1: + raise ValueError( + f"{part_type} requires exactly one of {source_keys}" + ) + + source = part[present_keys[0]] + nested_detail = ( + source.get("detail") if isinstance(source, dict) else None + ) + url = ( + source.get("url") or source.get("path", "") + if isinstance(source, dict) + else source + ) + if not isinstance(url, str) or not url: + raise ValueError(f"{part_type} requires a non-empty media URL") + if not url.startswith(("http://", "https://", "data:")): + if is_image: + local_image_sources.setdefault(url, None) + else: + local_video_sources.setdefault(url, None) + + pending_mutations.append( + ( + part, + source_keys, + canonical_type, + canonical_key, + is_image, + nested_detail, + url, + ) + ) + + sources = list(local_image_sources) + encoded_by_source: dict[str, str] = {} + if sources: + with ThreadPoolExecutor( + max_workers=_NEMO_GYM_IMAGE_ENCODE_MAX_WORKERS + ) as executor: + encoded_by_source = dict( + zip( + sources, + executor.map(_encode_single_image_source, sources), + strict=True, + ) + ) + + # Encode each unique video once. A video shared by G generations then points + # every part at the same string, instead of G separate base64 copies of the + # same file. Kept sequential: these payloads are large enough that encoding + # several at once would spike driver memory. + encoded_video_by_source: dict[str, str] = { + source: video_path_to_data_url(source) for source in local_video_sources + } + + # Apply mutations only after every local source was encoded successfully. + for ( + part, + source_keys, + canonical_type, + canonical_key, + is_image, + nested_detail, + url, + ) in pending_mutations: + for key in source_keys: + if key != canonical_key: + part.pop(key, None) + part["type"] = canonical_type + encoded = encoded_by_source if is_image else encoded_video_by_source + part[canonical_key] = encoded.get(url, url) + if is_image and nested_detail is not None: + part.setdefault("detail", nested_detail) + return nemo_gym_examples + + +# 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 isinstance(src, dict): + src = src.get("url") + if src is not None: + images.append(resolve_to_image(src)) + return images + + +def _is_trainable_output_item(item: dict) -> bool: + """Report whether an output item becomes a trainable assistant turn. + + The postprocess loop skips items whose ``generation_token_ids`` is missing + *or* empty, so per-turn image binning has to use the same predicate or the + two walks disagree and every later turn gets the wrong images. + """ + return bool(item.get("generation_token_ids")) + + +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 _is_trainable_output_item(item): + per_turn.append(pending) + pending = [] + elif item.get("role") != "assistant": + pending.extend(_extract_input_images_from_message(item)) + return per_turn + + +def _without_initial_media_sources( + messages: Any, initial_sources: list[Any] +) -> tuple[Any, bool]: + """Copy Responses messages and remove one ordered copy of initial images and videos.""" + 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_media_sources_from_responses_messages( + [{"content": [part]}] + ) + if ( + remaining_sources + and len(part_sources) == 1 + and media_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 + + _VideoConfigValue = TypeVar("_VideoConfigValue") _LOCAL_VIDEO_METADATA_KEYS = frozenset( { @@ -64,17 +336,6 @@ def _require_video_config_value( return value -def _get_content_part_url(part: dict[str, Any], *keys: str) -> str: - """Return a string media source from a Responses/Chat content part.""" - for key in keys: - value = part.get(key) - if isinstance(value, dict): - value = value.get("url") or value.get("path") - if isinstance(value, str) and value: - return value - return "" - - def _resolve_local_video_path(source: str) -> str: """Resolve a local video source and reject unsupported remote schemes.""" parsed = urlparse(source) @@ -96,41 +357,6 @@ def _resolve_local_video_path(source: str) -> str: return str(path.resolve()) -def normalize_video_urls_in_examples(examples: list[dict[str, Any]]) -> None: - """Convert bare local video paths to file URLs before Gym dispatch.""" - for example in 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") not in VIDEO_CONTENT_TYPES - ): - continue - media_key = next( - (key for key in ("video_url", "video", "url") if key in part), - None, - ) - if media_key is None: - continue - source = _get_content_part_url(part, media_key) - if not source or urlparse(source).scheme: - continue - normalized = Path(_resolve_local_video_path(source)).as_uri() - original = part[media_key] - if isinstance(original, dict): - original["url"] = normalized - else: - part[media_key] = normalized - - def _extract_static_video_messages( nemo_gym_example: dict[str, Any], ) -> tuple[list[dict[str, Any]], str | None] | None: @@ -168,7 +394,9 @@ def _extract_static_video_messages( if part_type == "input_text": hf_content.append({"type": "text", "text": part["text"]}) elif part_type in VIDEO_CONTENT_TYPES: - source = _get_content_part_url(part, "video_url", "video", "url") + source = get_responses_content_part_url( + part, "video_url", "video", "url" + ) if not source: raise ValueError(f"{part_type} requires a non-empty video URL") video_sources.append(source) @@ -177,7 +405,9 @@ def _extract_static_video_messages( if not part.get("_is_video_frame"): has_still_images = True continue - source = _get_content_part_url(part, "image_url", "image", "url") + source = get_responses_content_part_url( + part, "image_url", "image", "url" + ) if not source: raise ValueError( "Cached Gym video frames require a non-empty image URL." @@ -249,80 +479,6 @@ def _extract_static_video_messages( return hf_messages, _resolve_local_video_path(video_sources[0]) -def _json_mapping(value: Any, *, field_name: str) -> dict[str, Any]: - if isinstance(value, dict): - return copy.deepcopy(value) - if not isinstance(value, str): - raise TypeError(f"{field_name} must be a JSON object string or a dict") - if not value.strip(): - raise ValueError(f"{field_name} must not be empty") - try: - decoded = json.loads(value) - except json.JSONDecodeError as exc: - raise ValueError(f"{field_name} must contain valid JSON") from exc - if not isinstance(decoded, dict): - raise TypeError(f"{field_name} JSON must decode to an object") - return decoded - - -def _metadata_extra_body(nemo_gym_example: dict[str, Any]) -> dict[str, Any]: - params = nemo_gym_example.get("responses_create_params", {}) - if not isinstance(params, dict): - raise TypeError("responses_create_params must be a dict") - metadata = params.get("metadata", {}) - if not isinstance(metadata, dict): - raise TypeError("responses_create_params.metadata must be a dict") - if "extra_body" not in metadata: - return {} - return _json_mapping( - metadata["extra_body"], - field_name="responses_create_params.metadata.extra_body", - ) - - -def _chat_template_kwargs_for_processor( - nemo_gym_example: dict[str, Any], -) -> dict[str, Any]: - params = nemo_gym_example.get("responses_create_params", {}) - if not isinstance(params, dict): - raise TypeError("responses_create_params must be a dict") - metadata = params.get("metadata", {}) - if not isinstance(metadata, dict): - raise TypeError("responses_create_params.metadata must be a dict") - - extra_body = _metadata_extra_body(nemo_gym_example) - processor_kwargs: dict[str, Any] = {} - raw_chat_template_kwargs = metadata.get( - "chat_template_kwargs", extra_body.get("chat_template_kwargs") - ) - chat_template_kwargs = ( - _json_mapping( - raw_chat_template_kwargs, - field_name="responses_create_params.metadata.chat_template_kwargs", - ) - if raw_chat_template_kwargs is not None - else {} - ) - if chat_template_kwargs: - processor_kwargs["chat_template_kwargs"] = chat_template_kwargs - enable_thinking = chat_template_kwargs.get( - "enable_thinking", extra_body.get("enable_thinking") - ) - if enable_thinking is not None: - processor_kwargs["enable_thinking"] = enable_thinking - return processor_kwargs - - -def _deep_merge_dict(base: dict[str, Any], update: dict[str, Any]) -> dict[str, Any]: - merged = copy.deepcopy(base) - for key, value in update.items(): - if isinstance(value, dict) and isinstance(merged.get(key), dict): - merged[key] = _deep_merge_dict(merged[key], value) - else: - merged[key] = copy.deepcopy(value) - return merged - - def _inject_vllm_mm_processor_kwargs( nemo_gym_example: dict[str, Any], mm_processor_kwargs: dict[str, Any], @@ -387,7 +543,9 @@ def _replace_cached_video_frames_with_native_video( for part in content: if not isinstance(part, dict) or not part.get("_is_video_frame"): continue - frame_path = _get_content_part_url(part, "image_url", "image", "url") + frame_path = get_responses_content_part_url( + part, "image_url", "image", "url" + ) if not frame_path: raise ValueError( "Cached Gym video frames require a non-empty image URL." diff --git a/nemo_rl/environments/nemo_gym_request.py b/nemo_rl/environments/nemo_gym_request.py new file mode 100644 index 00000000000..46c44b689ec --- /dev/null +++ b/nemo_rl/environments/nemo_gym_request.py @@ -0,0 +1,155 @@ +# 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. + +"""Helpers for reading and merging NeMo Gym request metadata.""" + +import copy +import json +from typing import Any + + +def _json_mapping(value: Any, *, field_name: str) -> dict[str, Any]: + """Return a copied dict from a mapping or JSON object string. + + Example: + ``_json_mapping('{"enabled": true}', field_name="options")`` returns + ``{"enabled": True}``. + + Args: + value: Dict or JSON object string. + field_name: Field name used in errors. + + Returns: + A new dictionary. + + Raises: + TypeError: If the value is not a dict or JSON object string. + ValueError: If the string is empty or invalid JSON. + """ + if isinstance(value, dict): + return copy.deepcopy(value) + if not isinstance(value, str): + raise TypeError(f"{field_name} must be a JSON object string or a dict") + if not value.strip(): + raise ValueError(f"{field_name} must not be empty") + try: + decoded = json.loads(value) + except json.JSONDecodeError as exc: + raise ValueError(f"{field_name} must contain valid JSON") from exc + if not isinstance(decoded, dict): + raise TypeError(f"{field_name} JSON must decode to an object") + return decoded + + +def _metadata_extra_body(nemo_gym_example: dict[str, Any]) -> dict[str, Any]: + """Read ``metadata.extra_body`` as a dict. + + Example: + An example with ``metadata.extra_body='{"seed": 1}'`` returns + ``{"seed": 1}``. + + Args: + nemo_gym_example: Example containing Responses API parameters. + + Returns: + Parsed ``extra_body``, or an empty dict when absent. + + Raises: + TypeError: If request parameters or metadata are not dictionaries. + ValueError: If ``extra_body`` contains invalid JSON. + """ + params = nemo_gym_example.get("responses_create_params", {}) + if not isinstance(params, dict): + raise TypeError("responses_create_params must be a dict") + metadata = params.get("metadata", {}) + if not isinstance(metadata, dict): + raise TypeError("responses_create_params.metadata must be a dict") + if "extra_body" not in metadata: + return {} + return _json_mapping( + metadata["extra_body"], + field_name="responses_create_params.metadata.extra_body", + ) + + +def _chat_template_kwargs_for_processor( + nemo_gym_example: dict[str, Any], +) -> dict[str, Any]: + """Build processor kwargs from NeMo Gym chat-template metadata. + + Example: + ``{"chat_template_kwargs": {"enable_thinking": False}}`` becomes the + processor kwarg with the same name and value. + + Args: + nemo_gym_example: Example containing Responses API parameters. + + Returns: + Keyword arguments for the processor's chat template. + + Raises: + TypeError: If request metadata has an unsupported type. + ValueError: If a JSON metadata value is empty or invalid. + """ + params = nemo_gym_example.get("responses_create_params", {}) + if not isinstance(params, dict): + raise TypeError("responses_create_params must be a dict") + metadata = params.get("metadata", {}) + if not isinstance(metadata, dict): + raise TypeError("responses_create_params.metadata must be a dict") + + extra_body = _metadata_extra_body(nemo_gym_example) + processor_kwargs: dict[str, Any] = {} + raw_chat_template_kwargs = metadata.get( + "chat_template_kwargs", extra_body.get("chat_template_kwargs") + ) + chat_template_kwargs = ( + _json_mapping( + raw_chat_template_kwargs, + field_name="responses_create_params.metadata.chat_template_kwargs", + ) + if raw_chat_template_kwargs is not None + else {} + ) + if chat_template_kwargs: + processor_kwargs["chat_template_kwargs"] = chat_template_kwargs + enable_thinking = chat_template_kwargs.get( + "enable_thinking", extra_body.get("enable_thinking") + ) + if enable_thinking is not None: + processor_kwargs["enable_thinking"] = enable_thinking + return processor_kwargs + + +def _deep_merge_dict(base: dict[str, Any], update: dict[str, Any]) -> dict[str, Any]: + """Recursively merge two dictionaries without modifying either input. + + Example: + Merging ``{"a": {"b": 1}}`` with ``{"a": {"c": 2}}`` returns + ``{"a": {"b": 1, "c": 2}}``. + + Args: + base: Initial mapping. + update: Values to merge into ``base``. + + Returns: + A recursively merged deep copy. + """ + merged = copy.deepcopy(base) + for key, value in update.items(): + if isinstance(value, dict) and isinstance(merged.get(key), dict): + merged[key] = _deep_merge_dict(merged[key], value) + else: + merged[key] = copy.deepcopy(value) + return merged diff --git a/nemo_rl/evals/eval.py b/nemo_rl/evals/eval.py index 670ed625d7e..1ea9f766f4e 100644 --- a/nemo_rl/evals/eval.py +++ b/nemo_rl/evals/eval.py @@ -331,10 +331,8 @@ async def _run_env_eval_impl( prompts = [] prompts_for_display = [] for i, message_log in enumerate(batch["message_log"]): - if is_multimodal and batch["vllm_content"][i] is not None: - vllm_content = batch["vllm_content"][i] - prompt_dict = {"prompt": vllm_content} - multi_modal_data = {} + multi_modal_data = {} + if is_multimodal: audios = batch.get("vllm_audios", None) if audios is not None and len(audios[i]) > 0: multi_modal_data["audio"] = ( @@ -350,10 +348,34 @@ async def _run_env_eval_impl( multi_modal_data["video"] = ( videos[i][0] if len(videos[i]) == 1 else videos[i] ) + + vllm_content = batch["vllm_content"][i] if is_multimodal else None + if vllm_content is not None: + prompt_dict = {"prompt": vllm_content} + prompt_display = vllm_content if multi_modal_data: prompt_dict["multi_modal_data"] = multi_modal_data prompts.append(prompt_dict) - prompts_for_display.append(vllm_content) + prompts_for_display.append(prompt_display) + elif multi_modal_data: + # Placeholder-style processors pass prompt_token_ids with media. + prompt_token_ids = [] + for message in message_log: + token_ids = message["token_ids"] + prompt_token_ids.extend( + token_ids.tolist() + if isinstance(token_ids, torch.Tensor) + else token_ids + ) + prompt_dict = { + "prompt_token_ids": prompt_token_ids, + "multi_modal_data": multi_modal_data, + } + prompt_display = "\n".join( + str(message["content"]) for message in message_log + ) + prompts.append(prompt_dict) + prompts_for_display.append(prompt_display) else: # Text-only fallback: use raw prompt strings (vLLM will tokenize them). # Note: utils.py's format_prompt_for_vllm_generation uses pre-tokenized diff --git a/nemo_rl/experience/rollouts.py b/nemo_rl/experience/rollouts.py index 898bf8e55d8..46209cbadd0 100644 --- a/nemo_rl/experience/rollouts.py +++ b/nemo_rl/experience/rollouts.py @@ -47,7 +47,8 @@ VLLM_MULTIMODAL_DATA_KEYS, PackedTensor, attach_image_model_inputs_to_message, - extract_input_images_from_responses_messages, + extract_input_media_sources_from_responses_messages, + resolve_to_image, ) from nemo_rl.data_plane.schema import MASK_SAMPLE from nemo_rl.distributed.batched_data_dict import BatchedDataDict @@ -126,7 +127,14 @@ def attach_initial_nemo_gym_image_payloads( initial_messages = extra_env_info.get("responses_create_params", {}).get( "input", [] ) - images = extract_input_images_from_responses_messages(initial_messages) + # Load images from Responses-API input messages in encounter order. + images = [ + resolve_to_image(source) + for media_type, source in extract_input_media_sources_from_responses_messages( + initial_messages + ) + if media_type == "image" + ] if not images: continue if processor is None or getattr(processor, "image_processor", None) is None: diff --git a/nemo_rl/models/generation/megatron/config.py b/nemo_rl/models/generation/megatron/config.py index d79c4d5ba81..98d0733a80e 100644 --- a/nemo_rl/models/generation/megatron/config.py +++ b/nemo_rl/models/generation/megatron/config.py @@ -48,6 +48,9 @@ class MCoreGenerationSpecificArgs(TypedDict): materialize_only_last_token_logits: bool enable_chunked_prefill: bool enable_prefix_caching: bool + async_sched_mode: NotRequired[Literal["legacy", "async"]] + vision_embedding_cache_max_bytes: NotRequired[int] + allow_stale_multimodal_embeddings: NotRequired[bool] refit_backend: Literal["gloo", "nccl", "nvshmem"] num_speculative_tokens: int @@ -55,6 +58,24 @@ class MCoreGenerationSpecificArgs(TypedDict): mamba_inference_ssm_states_dtype: NotRequired[str] mamba_inference_conv_states_dtype: NotRequired[str] + # Raw media preprocessing corresponding with Megatron's + # ImageProcessingConfig / VideoProcessingConfig. + # `video_num_frames` is required for video. + vision_model_type: NotRequired[str] + image_dynamic_resolution: NotRequired[bool] + video_num_frames: NotRequired[int] # Frames sampled per video. + video_temporal_patch_size: NotRequired[int] # Frames per temporal patch. + video_target_num_patches: NotRequired[int] # Overrides the image max-patch budget. + video_maintain_aspect_ratio: NotRequired[bool] + + # Fully-qualified class path of the MCore inference wrapper, e.g. + # "megatron.core.inference.model_inference_wrappers.multimodal. + # nemotron_omni_inference_wrapper.NemotronOmniInferenceWrapper". + # Resolved by `_get_megatron_inference_wrapper_cls`; its `supports_*` + # attributes gate which modalities are preprocessed. Not media preprocessing + # itself, and used on the direct generate path as well as the HTTP endpoint. + megatron_inference_wrapper: NotRequired[str] + # KV cache lifecycle across suspend/resume: # - "persist": cache stays allocated; CUDA graphs remain valid (default) # - "offload": cache is moved off-GPU between iterations @@ -70,6 +91,8 @@ class MCoreGenerationSpecificArgs(TypedDict): # FP8/MXFP8 for the dedicated (non-colocated) inference model; # merged into its `megatron_cfg` by `merged_inference_megatron_cfg`. fp8_cfg: NotRequired[Fp8Config] + # Merged into megatron_cfg for gen workers; required for EP>1 + local CUDA graphs. + moe_pad_experts_for_cuda_graph_inference: NotRequired[bool] class MCoreGenerationConfig(GenerationConfig): diff --git a/nemo_rl/models/generation/megatron/megatron_generation.py b/nemo_rl/models/generation/megatron/megatron_generation.py index 7298ec95687..53294bce379 100644 --- a/nemo_rl/models/generation/megatron/megatron_generation.py +++ b/nemo_rl/models/generation/megatron/megatron_generation.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +from copy import deepcopy from typing import TYPE_CHECKING, Any, AsyncGenerator, Optional, cast import ray @@ -210,8 +211,8 @@ def __init__( Args: config: PolicyConfig for the Megatron model. tokenizer: The tokenizer for the model. - cluster: Cluster to deploy a dedicated inference Policy on. - policy: Existing training Policy to reuse for generation. + cluster: Cluster for a dedicated, non-colocated inference Policy. + policy: Existing training Policy reused for colocated generation. name_prefix: Prefix for naming the worker group (non-colocated only). processor: Optional processor for VLMs (non-colocated only). skip_weight_load: Do not load the weights from the checkpoint; refit will do it. @@ -232,7 +233,8 @@ def __init__( ) # `self.cfg` exposes the `generation` that matches the `GenerationInterface` contract. - # `self._policy_config` keeps a reference to the full PolicyConfig. + # `self._policy_config` keeps a reference to the full PolicyConfig. Dedicated + # inference receives a copy because worker setup may modify it. self._policy_config = config self.cfg: MCoreGenerationConfig = config["generation"] # Populated after the first prepare_for_generation (which starts the HTTP server). @@ -252,7 +254,7 @@ def __init__( # Stand up a dedicated inference-only policy. self._owns_policy = True self._policy_config = { - **config, + **deepcopy(config), "megatron_cfg": self.effective_megatron_cfg(config), } # Reserve GPUs before Policy workers grab them, to prevent disjoint NVLS domains. diff --git a/nemo_rl/models/generation/megatron/megatron_worker.py b/nemo_rl/models/generation/megatron/megatron_worker.py index 8c9ee1f7c61..737c89664ec 100644 --- a/nemo_rl/models/generation/megatron/megatron_worker.py +++ b/nemo_rl/models/generation/megatron/megatron_worker.py @@ -14,15 +14,17 @@ import asyncio import gc +import importlib import os import threading import time import warnings -from typing import AsyncGenerator, Optional +from typing import Any, AsyncGenerator, Optional import requests import torch from megatron.core.inference.config import ( + AsyncScheduleMode, InferenceConfig, KVCacheManagementMode, MambaInferenceStateConfig, @@ -53,6 +55,7 @@ ) from megatron.core.utils import unwrap_model +from nemo_rl.data.multimodal_utils import CACHED_VIDEO_FRAME_MANIFEST_MAGIC from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.models.generation.interfaces import ( GenerationDatumSpec, @@ -60,8 +63,12 @@ verify_right_padding, ) from nemo_rl.models.generation.megatron.utils import ( + build_image_preprocessing_config, + build_prompt_and_multimodal_data, + build_video_preprocessing_config, log_gpu_memory, resolve_torch_dtype, + sample_vision_tensors, ) from nemo_rl.models.megatron.memory_saver import ( HAVE_TORCH_MEMORY_SAVER, @@ -81,12 +88,14 @@ class MegatronGenerationMixin: - rank: global rank (used for logging). - tokenizer: HF tokenizer. - megatron_tokenizer: tokenizer for inference. + - processor: optional multimodal processor. - is_generation_colocated: Whether colocated or distributed. - _reserved_http_server_socket: driver-reserved server socket, or None. """ # Colocated-reshard hosts assign the dedicated inference-layout model here # (see MegatronPolicyWorkerImpl._build_colocated_inference_model). + processor: Optional[Any] = None inference_model = None _colocated_reshard_plan = None @@ -112,6 +121,89 @@ def _init_inference_engine_state(self) -> None: self._inference_loop = None self._inference_thread = None + def _get_megatron_inference_wrapper_cls(self) -> Optional[type]: + """Resolve the configured Megatron inference wrapper, if any. + + Returns: + The wrapper class, or None when no wrapper is configured. + """ + class_path = self.cfg["generation"]["mcore_generation_config"].get( + "megatron_inference_wrapper" + ) + if class_path is None: + return None + # Resolved once per worker: this is called per sample during generation. + cached = getattr(self, "_megatron_inference_wrapper_cls", None) + if cached is not None: + return cached + module_name, _, class_name = class_path.rpartition(".") + if not module_name: + raise ValueError( + "megatron_inference_wrapper must be a fully qualified class name, " + f"got {class_path!r}." + ) + try: + wrapper_cls = getattr(importlib.import_module(module_name), class_name) + except (ImportError, AttributeError) as e: + raise ValueError( + f"Could not resolve megatron_inference_wrapper {class_path!r} " + f"(from policy.generation.mcore_generation_config): {e}" + ) from e + self._megatron_inference_wrapper_cls = wrapper_cls + return wrapper_cls + + @staticmethod + def _wrapper_supports_modality( + inference_wrapper_cls: Optional[type], modality: str + ) -> bool: + """Whether the configured inference wrapper advertises `modality` support.""" + return bool( + inference_wrapper_cls is not None + and getattr(inference_wrapper_cls, f"supports_{modality}", False) + ) + + def _inference_model_and_media_parts(self, inference_wrapper_cls=None): + """Return the language model and its optional multimodal parent.""" + model = unwrap_model(self._gen_model()) + if isinstance(model, (list, tuple)): + if len(model) != 1: + raise NotImplementedError("Virtual pipeline models are not supported.") + model = model[0] + if inference_wrapper_cls is None: + inference_wrapper_cls = self._get_megatron_inference_wrapper_cls() + if not any( + self._wrapper_supports_modality(inference_wrapper_cls, modality) + for modality in ("image", "video", "audio") + ): + return model, None + return model.language_model, model + + def _build_image_preprocessing_config(self, generation_config: dict[str, Any]): + """Build raw-image preprocessing settings.""" + inference_wrapper_cls = self._get_megatron_inference_wrapper_cls() + if not self._wrapper_supports_modality(inference_wrapper_cls, "image"): + return None + processor = self.processor + if processor is None: + raise ValueError( + "Megatron multimodal generation requires the policy processor." + ) + # Omit absent keys entirely so MCore's own dataclass defaults apply; + # passing None would override them (these fields are `bool`, not `Optional`). + image_kwargs: dict[str, Any] = {} + if "image_dynamic_resolution" in generation_config: + image_kwargs["dynamic_resolution"] = bool( + generation_config["image_dynamic_resolution"] + ) + if "vision_model_type" in generation_config: + image_kwargs["vision_model_type"] = str( + generation_config["vision_model_type"] + ) + return build_image_preprocessing_config( + processor.image_processor, + **image_kwargs, + ) + def _setup_colocated_cuda_graph_managers(self) -> None: """Create inference CUDA-graph managers for shared-model colocated generation. @@ -223,8 +315,13 @@ def _initialize_inference_engine(self, mcore_generation_config: dict) -> None: ) from megatron.core.utils import get_attr_wrapped_model - gen_model = self._gen_model() - pg_collection = get_attr_wrapped_model(gen_model, "pg_collection") + inference_wrapper_cls = self._get_megatron_inference_wrapper_cls() + inference_model, media_model = self._inference_model_and_media_parts( + inference_wrapper_cls + ) + engine_model = media_model if media_model is not None else self._gen_model() + pg_collection = get_attr_wrapped_model(self._gen_model(), "pg_collection") + model_config = inference_model.config buffer_size_gb = mcore_generation_config["buffer_size_gb"] num_cuda_graphs = mcore_generation_config["num_cuda_graphs"] @@ -245,7 +342,9 @@ def _initialize_inference_engine(self, mcore_generation_config: dict) -> None: num_speculative_tokens = mcore_generation_config["num_speculative_tokens"] max_requests = mcore_generation_config.get("max_requests") - mamba_inference_state_config = MambaInferenceStateConfig.from_model(gen_model) + mamba_inference_state_config = MambaInferenceStateConfig.from_model( + inference_model + ) is_hybrid_model = mamba_inference_state_config is not None if is_hybrid_model: if ( @@ -270,11 +369,36 @@ def _initialize_inference_engine(self, mcore_generation_config: dict) -> None: logging_step_interval = 0 # flashinfer's fused-RoPE kernel only dispatches fp16/bf16 q/k. - use_flashinfer_fused_rope = gen_model.config.params_dtype in ( + use_flashinfer_fused_rope = model_config.params_dtype in ( torch.float16, torch.bfloat16, ) + image_preprocessing_config = self._build_image_preprocessing_config( + mcore_generation_config + ) + video_preprocessing_config = build_video_preprocessing_config( + image_preprocessing_config, + mcore_generation_config, + frame_manifest_magic=CACHED_VIDEO_FRAME_MANIFEST_MAGIC, + ) + + # Only forward keys the config actually sets, so MCore's InferenceConfig + # defaults stay the single source of truth for the ones it omits. + inference_overrides: dict[str, Any] = {} + if "async_sched_mode" in mcore_generation_config: + inference_overrides["async_sched_mode"] = AsyncScheduleMode( + mcore_generation_config["async_sched_mode"] + ) + if "vision_embedding_cache_max_bytes" in mcore_generation_config: + inference_overrides["vision_embedding_cache_max_bytes"] = int( + mcore_generation_config["vision_embedding_cache_max_bytes"] + ) + if "allow_stale_multimodal_embeddings" in mcore_generation_config: + inference_overrides["allow_stale_multimodal_embeddings"] = bool( + mcore_generation_config["allow_stale_multimodal_embeddings"] + ) + inference_config = InferenceConfig( block_size_tokens=block_size_tokens, buffer_size_gb=buffer_size_gb, @@ -290,6 +414,7 @@ def _initialize_inference_engine(self, mcore_generation_config: dict) -> None: materialize_only_last_token_logits=materialize_only_last_token_logits, enable_chunked_prefill=enable_chunked_prefill, enable_prefix_caching=mcore_generation_config["enable_prefix_caching"], + **inference_overrides, prefix_caching_coordinator_policy=PrefixCachingCoordinatorPolicy( "first_prefix_block" ), @@ -303,19 +428,30 @@ def _initialize_inference_engine(self, mcore_generation_config: dict) -> None: num_speculative_tokens=num_speculative_tokens, logprobs_mode=mcore_generation_config["logprobs_mode"], max_requests=max_requests, + image_preprocessing_config=image_preprocessing_config, + video_preprocessing_config=video_preprocessing_config, ) if "inference_cuda_graph_scope" in mcore_generation_config: - gen_model.config.inference_cuda_graph_scope = InferenceCudaGraphScope[ + engine_model.config.inference_cuda_graph_scope = InferenceCudaGraphScope[ mcore_generation_config["inference_cuda_graph_scope"] ] self.inference_context = DynamicInferenceContext( - gen_model.config, inference_config - ) - self.inference_wrapped_model = GPTInferenceWrapper( - gen_model, self.inference_context + engine_model.config, inference_config ) + if media_model is None: + self.inference_wrapped_model = GPTInferenceWrapper( + engine_model, self.inference_context + ) + else: + if inference_wrapper_cls is None: + raise ValueError( + "Multimodal inference requires megatron_inference_wrapper." + ) + self.inference_wrapped_model = inference_wrapper_cls( + engine_model, self.inference_context + ) text_generation_controller = TextGenerationController( inference_wrapped_model=self.inference_wrapped_model, tokenizer=self.megatron_tokenizer, @@ -432,6 +568,7 @@ def _setup_openai_api_server(self) -> str: parsers=self.cfg["generation"]["mcore_generation_config"]["parsers"], verbose=False, sock=reserved_socket, + multimodal_prompt_config=self.inference_wrapped_model.multimodal_prompt_config, ) base_url = f"http://{ip}:{server_port}/v1" @@ -495,7 +632,11 @@ def finish_generation(self, *, release_gpu: bool = True) -> None: print(f"[Rank {self.rank}] finishing generation", flush=True) log_gpu_memory("finish_generation START") - lang_module = unwrap_model(self._gen_model()) + inference_model, media_model = self._inference_model_and_media_parts() + lang_module = unwrap_model(inference_model) + graph_module = unwrap_model( + media_model if media_model is not None else inference_model + ) if self.is_generation_colocated: if self._inference_engine_initialized and not self._inference_engine_asleep: @@ -504,7 +645,8 @@ def finish_generation(self, *, release_gpu: bool = True) -> None: "cuda_graph_impl" ] if cuda_graph_impl != "none": - toggle_cuda_graphs(lang_module, set_to="none") + # Restore the same full model tree cached during worker setup. + toggle_cuda_graphs(graph_module, set_to="none") # Need to turn off padding before training. # Gains nightly MoE coverage once #2884 and #3570 merge. set_decode_expert_padding(lang_module, set_to=False) @@ -552,10 +694,6 @@ def prepare_for_generation(self, tags=None, **kwargs) -> None: if self._colocated_reshard_plan is not None: self._build_colocated_inference_model(self.cfg) - gen_model = self._gen_model() - # `flash_decode` selects Megatron Inference's deprecated static-batching decode path, - # which would cause an assertion error if taken. - gen_model.config.flash_decode = False if self.is_generation_colocated and self.inference_model is None: self.model = self.move_model( self.model, "cuda", move_params=True, move_grads=False @@ -568,13 +706,19 @@ def prepare_for_generation(self, tags=None, **kwargs) -> None: and self._forward_pre_hook_enabled() ): self._disable_forward_pre_hook_until_next_train_step(param_sync=True) - gen_model = self.model # Colocated reshard (hosts without a dedicated inference model skip it). if self.inference_model is not None: self._reshard_into_inference_model() - lang_module = unwrap_model(gen_model) + inference_model, media_model = self._inference_model_and_media_parts() + # `flash_decode` selects Megatron Inference's deprecated static-batching decode path, + # which would cause an assertion error if taken. + inference_model.config.flash_decode = False + lang_module = unwrap_model(inference_model) + graph_module = unwrap_model( + media_model if media_model is not None else inference_model + ) lang_module.eval() rotary_module = getattr(lang_module, "rotary_pos_emb", None) @@ -585,7 +729,9 @@ def prepare_for_generation(self, tags=None, **kwargs) -> None: cuda_graph_impl = mcore_generation_config["cuda_graph_impl"] if cuda_graph_impl != "none": - toggle_cuda_graphs(lang_module, set_to=cuda_graph_impl) + # Use the same root object as _setup_colocated_cuda_graph_managers; + # CUDA-graph manager caches are keyed by model identity. + toggle_cuda_graphs(graph_module, set_to=cuda_graph_impl) # tags=["weights"] means we are inside refit_policy_generation between # suspend_for_refit and the weight transfer — the engine was intentionally @@ -606,7 +752,11 @@ def report_dp_openai_server_base_url(self) -> Optional[str]: return self.base_url def _build_sampling_params( - self, greedy: bool, stop_words: Optional[list[str]] + self, + greedy: bool, + stop_words: Optional[list[str]], + *, + return_prompt_tokens: bool = False, ) -> SamplingParams: """Build mcore SamplingParams for a single request.""" top_k_cfg = self.cfg["generation"]["top_k"] @@ -626,6 +776,7 @@ def _build_sampling_params( num_tokens_to_generate=self.cfg["generation"]["max_new_tokens"], termination_id=self.megatron_tokenizer.eod, stop_words=stop_words, + return_prompt_tokens=return_prompt_tokens, ) def _merge_stop_strings( @@ -641,10 +792,26 @@ def _merge_stop_strings( stop_set.update(sample_ss) return list(stop_set) if stop_set else None + def _sample_vision_tensors(self, data, index: int): + """Return one sample's vision tensors from RL PackedTensors.""" + return sample_vision_tensors(data, index) + + def _build_prompt_and_multimodal_data(self, data, index: int): + """Build one pre-expanded token prompt and optional MCore media payload.""" + return build_prompt_and_multimodal_data( + data, + index, + sample_tensors=self._sample_vision_tensors, + supports_modality=lambda modality: self._wrapper_supports_modality( + self._get_megatron_inference_wrapper_cls(), + modality, + ), + ) + def _prepare_data_for_generation( self, data: BatchedDataDict[GenerationDatumSpec], greedy: bool = False - ) -> tuple[torch.Tensor, torch.Tensor, list[SamplingParams]]: - """Build the prompt tensors and a per-request SamplingParams for each sample.""" + ) -> tuple[list[list[int]], list[Optional[Any]], list[SamplingParams]]: + """Build prompts, optional media payloads, and sampling parameters.""" if data is not None: assert isinstance(data, BatchedDataDict), ( f"data must be a BatchedDataDict, got type: {type(data)}" @@ -657,21 +824,29 @@ def _prepare_data_for_generation( f"Input to Megatron Generation worker is not properly right-padded: {error_msg}" ) - prompt_tokens_tensor = data["input_ids"].cuda() - prompt_lengths_tensor = data["input_lengths"] - batch_stop_strings = data.get("stop_strings", []) + prompts: list[list[int]] = [] + multi_modal_data_list: list[Optional[Any]] = [] sampling_params = [] - for i in range(prompt_tokens_tensor.size(0)): + for i in range(data.size): + prompt, multi_modal_data = self._build_prompt_and_multimodal_data(data, i) sample_stop_strings = ( batch_stop_strings[i] if i < len(batch_stop_strings) else None ) stop_words = self._merge_stop_strings( [sample_stop_strings] if sample_stop_strings else None ) - sampling_params.append(self._build_sampling_params(greedy, stop_words)) + prompts.append(prompt) + multi_modal_data_list.append(multi_modal_data) + sampling_params.append( + self._build_sampling_params( + greedy, + stop_words, + return_prompt_tokens=multi_modal_data is not None, + ) + ) - return prompt_tokens_tensor, prompt_lengths_tensor, sampling_params + return prompts, multi_modal_data_list, sampling_params def _parse_result_to_batched_data_dict( self, @@ -685,6 +860,22 @@ def _parse_result_to_batched_data_dict( max_gen_seq_len = max(len(x.generated_tokens) for x in result) padded_input_length = input_ids.size(1) + expected_prompt_lengths = [int(length) for length in input_lengths.tolist()] + inference_prompt_lengths = [ + len(x.prompt_tokens) + if getattr(x, "prompt_tokens", None) is not None + else expected_prompt_lengths[i] + for i, x in enumerate(result) + ] + if any(getattr(x, "prompt_tokens", None) is not None for x in result): + if inference_prompt_lengths != expected_prompt_lengths: + raise RuntimeError( + "Megatron inference prompt lengths do not match the training " + "processor input lengths: " + f"inference={inference_prompt_lengths}, " + f"training={expected_prompt_lengths}." + ) + max_seq_len = padded_input_length + max_gen_seq_len output_ids_padded = torch.full( (batch_size, max_seq_len), @@ -707,9 +898,10 @@ def _parse_result_to_batched_data_dict( ) for i in range(batch_size): # Take the prompt from the request we submitted rather than from the - # engine's reply: mcore only echoes prompt_tokens back when - # SamplingParams.return_prompt_tokens is set, and asking for them would - # ship the whole prompt over ZMQ for data we already hold. + # engine's reply. Multimodal requests do set + # SamplingParams.return_prompt_tokens, but only so the echoed tokens + # can be length-checked below; the padded output is still built from + # the prompt we already hold rather than shipped back over ZMQ. prompt_len = input_lengths[i].item() generated_tokens = result[i].generated_tokens seq_len = prompt_len + len(generated_tokens) @@ -754,7 +946,7 @@ def generate( - generation_lengths: Lengths of each response - unpadded_sequence_lengths: Lengths of each input + generated sequence """ - prompt_tokens_tensor, prompt_lengths_tensor, sampling_params = ( + prompts, multi_modal_data_list, sampling_params = ( self._prepare_data_for_generation(data, greedy) ) if self._inference_loop is None: @@ -763,8 +955,8 @@ def generate( ) future = asyncio.run_coroutine_threadsafe( self._generate_with_persistent_engine( - prompt_tokens_tensor, - prompt_lengths_tensor, + prompts, + multi_modal_data_list, sampling_params, ), self._inference_loop, @@ -794,13 +986,13 @@ async def _generate_single_item( index: int, ) -> tuple[int, BatchedDataDict[GenerationOutputSpec]]: datum = data.get_batch(index, 1) - prompt_tokens_tensor, prompt_lengths_tensor, sampling_params = ( + prompts, multi_modal_data_list, sampling_params = ( self._prepare_data_for_generation(datum, greedy) ) future = asyncio.run_coroutine_threadsafe( self._generate_with_persistent_engine( - prompt_tokens_tensor, - prompt_lengths_tensor, + prompts, + multi_modal_data_list, sampling_params, ), self._inference_loop, @@ -817,8 +1009,8 @@ async def _generate_single_item( async def _generate_with_persistent_engine( self, - prompt_tokens_tensor: torch.Tensor, - prompt_lengths_tensor: torch.Tensor, + prompts: list[list[int]], + multi_modal_data_list: list[Optional[Any]], sampling_params: list[SamplingParams], ) -> list: """Submit requests through the persistent inference client (rank 0 only).""" @@ -829,17 +1021,18 @@ async def _generate_with_persistent_engine( "Only rank 0 creates a client to communicate with the coordinator" ) - print( - f"[Rank {dist_rank}] Submitting {prompt_tokens_tensor.size(0)} requests to coordinator" - ) + print(f"[Rank {dist_rank}] Submitting {len(prompts)} requests to coordinator") futures = [] - for prompt_tokens, prompt_len, request_sampling_params in zip( - prompt_tokens_tensor, prompt_lengths_tensor, sampling_params, strict=True + for prompt, multi_modal_data, request_sampling_params in zip( + prompts, multi_modal_data_list, sampling_params, strict=True ): - prompt = prompt_tokens[: prompt_len.item()].tolist() futures.append( - self.inference_client.add_request(prompt, request_sampling_params) + self.inference_client.add_request( + prompt, + request_sampling_params, + multi_modal_data=multi_modal_data, + ) ) results: list[DynamicInferenceRequest] = await asyncio.gather(*futures) diff --git a/nemo_rl/models/generation/megatron/utils.py b/nemo_rl/models/generation/megatron/utils.py index 339d57e46cc..6e40693e2d5 100644 --- a/nemo_rl/models/generation/megatron/utils.py +++ b/nemo_rl/models/generation/megatron/utils.py @@ -12,10 +12,241 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dataclasses import replace +from typing import Any, Callable + import torch +from megatron.core.inference.config import ImageProcessingConfig, VideoProcessingConfig from megatron.core.inference.utils import device_memory_summary +def sample_vision_tensors(data, index: int): + """Return one sample's vision tensors from RL PackedTensors.""" + from nemo_rl.data.multimodal_utils import PackedTensor + + pixel_values = data.get("pixel_values") + imgs_sizes = data.get("imgs_sizes") + packed_num_frames = data.get("num_frames") + if pixel_values is None and imgs_sizes is None: + if packed_num_frames is not None: + raise ValueError("num_frames was provided without vision tensors.") + return None, None, None + if pixel_values is None or imgs_sizes is None: + raise ValueError( + "Megatron image generation requires both pixel_values and imgs_sizes." + ) + if not isinstance(pixel_values, PackedTensor) or not isinstance( + imgs_sizes, PackedTensor + ): + raise TypeError( + "Megatron image generation expects pixel_values and imgs_sizes " + "as per-sample PackedTensor values." + ) + if packed_num_frames is not None and not isinstance( + packed_num_frames, PackedTensor + ): + raise TypeError( + "Megatron video generation expects num_frames as a " + "per-sample PackedTensor value." + ) + + # `.tensors` is the physical segment list and only matches logical row + # indices while media are not deduplicated. + for name, packed in ( + ("pixel_values", pixel_values), + ("imgs_sizes", imgs_sizes), + ("num_frames", packed_num_frames), + ): + if packed is not None and packed._row_offsets is not None: + raise ValueError( + f"Megatron generation cannot index deduplicated {name}; " + "set deduplicate_multimodal_data=false (it is only supported " + "for the vLLM backend)." + ) + + imgs = pixel_values.tensors[index] + sizes = imgs_sizes.tensors[index] + num_frames = ( + packed_num_frames.tensors[index] if packed_num_frames is not None else None + ) + if imgs is None and sizes is None: + return None, None, None + if imgs is None or sizes is None: + raise ValueError( + "Megatron image generation requires matching per-sample " + "pixel_values and imgs_sizes." + ) + if imgs.ndim == 3: + imgs = imgs.unsqueeze(0) + if sizes.ndim == 1: + sizes = sizes.unsqueeze(0) + if num_frames is not None: + num_frames = num_frames.to(dtype=torch.int32).reshape(-1) + return imgs, sizes, num_frames + + +def build_prompt_and_multimodal_data( + data, + index: int, + *, + supports_modality: Callable[[str], bool], + sample_tensors: Callable = sample_vision_tensors, +): + """Build one pre-expanded token prompt and optional MCore media payload.""" + length = int(data["input_lengths"][index].item()) + prompt = data["input_ids"][index, :length].tolist() + imgs, imgs_sizes, num_frames = sample_tensors(data, index) + if imgs is None: + return prompt, None + + assert imgs_sizes is not None + is_video = num_frames is not None and bool(torch.any(num_frames > 1).item()) + modality = "video" if is_video else "image" + if not supports_modality(modality): + raise ValueError( + f"The configured megatron_inference_wrapper does not support " + f"{modality} inputs." + ) + if is_video: + if int(num_frames.sum().item()) != int(imgs_sizes.shape[0]): + raise ValueError( + "Video num_frames must partition imgs_sizes exactly: " + f"sum(num_frames)={int(num_frames.sum().item())}, " + f"imgs_sizes={imgs_sizes.shape[0]}." + ) + modality_data = { + "imgs": imgs, + "imgs_sizes": imgs_sizes, + "num_frames": num_frames, + } + else: + modality_data = {"imgs": imgs, "imgs_sizes": imgs_sizes} + return prompt, { + modality: modality_data, + "media_tokens_preexpanded": True, + } + + +def build_image_preprocessing_config( + image_processor: Any, + *, + dynamic_resolution: bool | None = None, + vision_model_type: str | None = None, +) -> ImageProcessingConfig: + """Translate an HF image processor to an MCore config. + + Args: + image_processor: HF image processor to read patch/normalization fields from. + dynamic_resolution: Override for `ImageProcessingConfig.dynamic_resolution`. + `None` leaves MCore's own default in place. + vision_model_type: Override for `ImageProcessingConfig.vision_model_type`. + `None` leaves MCore's own default in place. + """ + + def read(*names: str) -> Any: + for name in names: + value = getattr(image_processor, name, None) + if value is not None: + return value + return None + + patch_dim = read("patch_size", "patch_dim") + if isinstance(patch_dim, dict): + patch_dim = patch_dim.get("height", patch_dim.get("width")) + min_patches = read("min_num_patches") + max_patches = read("max_num_patches") + pixel_mean = read("norm_mean", "image_mean") + pixel_std = read("norm_std", "image_std") + + if ( + patch_dim is None + or min_patches is None + or max_patches is None + or pixel_mean is None + or pixel_std is None + ): + missing = [ + name + for name, value in ( + ("patch_size", patch_dim), + ("min_num_patches", min_patches), + ("max_num_patches", max_patches), + ("norm_mean", pixel_mean), + ("norm_std", pixel_std), + ) + if value is None + ] + raise ValueError( + f"{type(image_processor).__name__} does not expose {', '.join(missing)}, " + "so MCore cannot preprocess raw images the way this model's data " + "pipeline does." + ) + + downsample_ratio = read("downsample_ratio") + if downsample_ratio is not None: + merge_size = int(round(1.0 / float(downsample_ratio))) + else: + merge_size = int(read("merge_size", "spatial_merge_size") or 1) + + return ImageProcessingConfig( + patch_dim=int(patch_dim), + **( + {} + if dynamic_resolution is None + else {"dynamic_resolution": dynamic_resolution} + ), + **( + {} + if vision_model_type is None + else {"vision_model_type": vision_model_type} + ), + use_tiling=False, + pixel_shuffle=merge_size > 1, + spatial_merge_size=merge_size, + dynamic_resolution_min_patches=int(min_patches), + dynamic_resolution_max_patches=int(max_patches), + pixel_mean=[float(value) for value in pixel_mean], + pixel_std=[float(value) for value in pixel_std], + ) + + +def build_video_preprocessing_config( + image_config: ImageProcessingConfig | None, + generation_config: dict[str, Any], + *, + frame_manifest_magic: bytes, +) -> VideoProcessingConfig | None: + """Build video preprocessing when explicitly enabled by generation config.""" + video_num_frames = generation_config.get("video_num_frames") + if image_config is None or video_num_frames is None: + return None + + # Video configs. + video_kwargs: dict[str, Any] = {} + if "video_temporal_patch_size" in generation_config: + video_kwargs["temporal_patch_size"] = int( + generation_config["video_temporal_patch_size"] + ) + if "video_maintain_aspect_ratio" in generation_config: + video_kwargs["video_maintain_aspect_ratio"] = bool( + generation_config["video_maintain_aspect_ratio"] + ) + + target_num_patches = generation_config.get("video_target_num_patches") + if target_num_patches is not None: + image_config = replace( + image_config, + dynamic_resolution_max_patches=int(target_num_patches), + ) + + return VideoProcessingConfig( + image_config=image_config, + num_frames=int(video_num_frames), + frame_manifest_magic=frame_manifest_magic, + **video_kwargs, + ) + + def resolve_torch_dtype(val): """Convert a value to `torch.dtype`.""" if isinstance(val, torch.dtype): diff --git a/nemo_rl/models/generation/vllm/config.py b/nemo_rl/models/generation/vllm/config.py index dabf01d0b1e..c4d69be527b 100644 --- a/nemo_rl/models/generation/vllm/config.py +++ b/nemo_rl/models/generation/vllm/config.py @@ -289,7 +289,7 @@ def normalize_vllm_refit_config(config: VllmConfig) -> VllmRefitConfig | None: raise ValueError( "vllm_cfg.reset_encoder_cache_after_weight_update is not supported " f"with refit_transport={transport!r}: this transport's refit path " - "does not reset the multimodal encoder cache, so stale vision " + "does not reset the multimodal encoder cache, so stale multimodal " "embeddings would silently survive weight updates. Supported " "transports: null (collective/IPC) and 'nccl_reshard'." ) diff --git a/nemo_rl/models/generation/vllm/video_utils.py b/nemo_rl/models/generation/vllm/video_utils.py index c573a95d462..88a3059926e 100644 --- a/nemo_rl/models/generation/vllm/video_utils.py +++ b/nemo_rl/models/generation/vllm/video_utils.py @@ -24,13 +24,16 @@ import torch from PIL import Image +from nemo_rl.data.multimodal_utils import ( + CACHED_VIDEO_FRAME_MANIFEST_MAGIC, + CACHED_VIDEO_FRAME_MANIFEST_MIME, +) + VideoSamplingStyle = Literal["nemotron_vl"] _TORCHCODEC_END_OF_STREAM_ERROR = ( "Requested next frame while there are no more frames left to decode." ) -_CACHED_VIDEO_FRAME_MANIFEST_MAGIC = b"NEMO_RL_CACHED_VIDEO_FRAMES_V1\n" -_CACHED_VIDEO_FRAME_MANIFEST_MIME = "video/x-nemo-rl-cached-frames" def _round_video_frame_count( @@ -197,11 +200,11 @@ def build_cached_video_frame_data_url( "frame_paths": resolved_frames, "metadata": build_cached_video_frame_metadata(len(resolved_frames)), } - payload = _CACHED_VIDEO_FRAME_MANIFEST_MAGIC + json.dumps( + payload = CACHED_VIDEO_FRAME_MANIFEST_MAGIC + json.dumps( manifest, separators=(",", ":") ).encode("utf-8") encoded = base64.b64encode(payload).decode("ascii") - return f"data:{_CACHED_VIDEO_FRAME_MANIFEST_MIME};base64,{encoded}" + return f"data:{CACHED_VIDEO_FRAME_MANIFEST_MIME};base64,{encoded}" def _load_cached_video_frame_manifest( @@ -210,11 +213,11 @@ def _load_cached_video_frame_manifest( num_frames: int, ) -> tuple[np.ndarray, dict[str, Any]] | None: """Load an internal cached-frame manifest passed through vLLM VideoMediaIO.""" - if not data.startswith(_CACHED_VIDEO_FRAME_MANIFEST_MAGIC): + if not data.startswith(CACHED_VIDEO_FRAME_MANIFEST_MAGIC): return None try: - manifest = json.loads(data[len(_CACHED_VIDEO_FRAME_MANIFEST_MAGIC) :]) + manifest = json.loads(data[len(CACHED_VIDEO_FRAME_MANIFEST_MAGIC) :]) except (UnicodeDecodeError, json.JSONDecodeError) as exc: raise ValueError("Invalid cached Gym video frame manifest.") from exc if not isinstance(manifest, dict): diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index d0f57de802a..afc14e8dcf0 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -260,6 +260,42 @@ def _sync_distrib_opt(distrib_opt): TokenizerType = TypeVar("TokenizerType", bound=PreTrainedTokenizerBase) +_OPTIMIZER_DTYPE_KEYS = ( + "params_dtype", + "main_grads_dtype", + "main_params_dtype", + "exp_avg_dtype", + "exp_avg_sq_dtype", +) + + +def _resolve_optimizer_dtype_kwargs(optimizer_cfg: dict[str, Any]) -> dict[str, Any]: + """Resolve optimizer dtype strings, including TE's uint8-backed FP8 moments.""" + resolved = dict(optimizer_cfg) + dtype_aliases = { + "fp32": torch.float32, + "float32": torch.float32, + "fp16": torch.float16, + "float16": torch.float16, + "bf16": torch.bfloat16, + "bfloat16": torch.bfloat16, + "fp8": torch.uint8, + "uint8": torch.uint8, + } + for key in _OPTIMIZER_DTYPE_KEYS: + value = resolved.get(key) + if isinstance(value, str): + normalized = value.lower().removeprefix("torch.") + try: + resolved[key] = dtype_aliases[normalized] + except KeyError as e: + raise ValueError( + f"Unsupported optimizer dtype {value!r} for {key}. " + "Supported Transformer Engine FusedAdam dtype aliases: " + f"{', '.join(dtype_aliases)}" + ) from e + return resolved + def destroy_parallel_state(): """Safely destroy parallel state and reset async call tracking. @@ -1501,7 +1537,7 @@ def _create_megatron_config( "overlap_param_gather" ] optimizer_kwargs = { - **config["megatron_cfg"]["optimizer"], + **_resolve_optimizer_dtype_kwargs(config["megatron_cfg"]["optimizer"]), "overlap_param_gather": overlap_param_gather, "reuse_grad_buf_for_mxfp8_param_ag": reuse_grad_buf_for_mxfp8_param_ag, } diff --git a/nemo_rl/models/policy/__init__.py b/nemo_rl/models/policy/__init__.py index fed800dcf99..300995b7697 100644 --- a/nemo_rl/models/policy/__init__.py +++ b/nemo_rl/models/policy/__init__.py @@ -265,6 +265,10 @@ class MegatronOptimizerConfig(TypedDict): optimizer_offload_fraction: float # overlap optimizer state transfers with CPU optimizer updates overlap_cpu_optimizer_d2h_h2d: NotRequired[bool] + # Precision-aware Adam moment / remainder dtypes (YAML strings resolved in setup). + exp_avg_dtype: NotRequired[str] + exp_avg_sq_dtype: NotRequired[str] + store_param_remainders: NotRequired[bool] class MegatronSchedulerConfig(TypedDict): diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index 90ed5149cee..3135f88721c 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -430,6 +430,7 @@ def __init__( optimizer_path: Optional[str] = None, init_optimizer: bool = True, init_reference_model: bool = True, + processor: Optional[Any] = None, *, worker_sharding_annotations: NamedSharding, skip_weight_load: bool = False, @@ -523,6 +524,7 @@ def __init__( self.tokenizer = tokenizer if self.tokenizer.pad_token is None: self.tokenizer.pad_token = self.tokenizer.eos_token + self.processor = processor # Step 3: Setup model configuration runtime_config = validate_and_set_config( diff --git a/nemo_rl/utils/venvs.py b/nemo_rl/utils/venvs.py index fa5b8acb738..65f6d8b0930 100644 --- a/nemo_rl/utils/venvs.py +++ b/nemo_rl/utils/venvs.py @@ -30,6 +30,25 @@ logger = logging.getLogger(__name__) +def add_hf_modules_cache_to_pythonpath(env_vars: dict[str, str]) -> dict[str, str]: + """Make Hugging Face ``trust_remote_code`` modules importable by Ray actors.""" + result = env_vars.copy() + modules_cache = result.get("HF_MODULES_CACHE") + if modules_cache is None: + try: + from transformers.utils import HF_MODULES_CACHE + except ImportError: + return result + modules_cache = HF_MODULES_CACHE + result["HF_MODULES_CACHE"] = modules_cache + + pythonpath = result.get("PYTHONPATH", "") + path_entries = pythonpath.split(os.pathsep) if pythonpath else [] + if modules_cache not in path_entries: + result["PYTHONPATH"] = os.pathsep.join([modules_cache, *path_entries]) + return result + + @lru_cache(maxsize=None) def create_local_venv( py_executable: str, venv_name: str, force_rebuild: bool = False @@ -193,13 +212,18 @@ def create_local_venv_on_each_node(py_executable: str, venv_name: str): return paths[0] -def make_actor_runtime_env(actor_class_fqn: str) -> dict: +def make_actor_runtime_env( + actor_class_fqn: str, + *, + extra_env_vars: dict[str, str] | None = None, +) -> dict: """Build a Ray ``runtime_env`` for one of our registered actors. Resolves the actor's tier-specific py_executable via the registry, materializes a per-node venv when uv-managed, and packages it with ``VIRTUAL_ENV`` / ``UV_PROJECT_ENVIRONMENT`` env vars so workers see - the same interpreter as the driver. + the same interpreter as the driver. Additional actor-specific environment + variables can be supplied via ``extra_env_vars``. Used by ReplayBuffer, AsyncTrajectoryCollector, and SyncRolloutActor — three actors that need the VLLM tier's venv on every node. Also @@ -215,11 +239,16 @@ def make_actor_runtime_env(actor_class_fqn: str) -> dict: if py_exec.startswith("uv"): py_exec = create_local_venv_on_each_node(py_exec, actor_class_fqn) venv = os.path.dirname(os.path.dirname(py_exec)) # strip bin/python - return { - "py_executable": py_exec, - "env_vars": { + env_vars = add_hf_modules_cache_to_pythonpath( + { **os.environ, "VIRTUAL_ENV": venv, "UV_PROJECT_ENVIRONMENT": venv, - }, + } + ) + if extra_env_vars: + env_vars.update(extra_env_vars) + return { + "py_executable": py_exec, + "env_vars": env_vars, } diff --git a/pyrefly.toml b/pyrefly.toml index 5c93a2b0a46..ce9c2163780 100644 --- a/pyrefly.toml +++ b/pyrefly.toml @@ -145,7 +145,8 @@ project-includes = [ "nemo_rl/environments/interfaces.py", "nemo_rl/environments/math_environment.py", "nemo_rl/environments/metrics.py", - "nemo_rl/environments/nemo_gym_video.py", + "nemo_rl/environments/nemo_gym_multimodal.py", + "nemo_rl/environments/nemo_gym_request.py", "nemo_rl/environments/nemotron_utils.py", "nemo_rl/environments/rewards.py", "nemo_rl/environments/utils.py", diff --git a/tests/check_metrics.py b/tests/check_metrics.py index ac774435d0a..547397df400 100755 --- a/tests/check_metrics.py +++ b/tests/check_metrics.py @@ -20,6 +20,7 @@ import argparse import builtins import json +import math import statistics import sys @@ -38,6 +39,12 @@ def max(value): return builtins.max(float(v) for v in value.values()) +def all_finite(value): + """Return whether a metric is present and all recorded values are finite.""" + values = [float(v) for v in value.values()] + return bool(values) and builtins.all(math.isfinite(value) for value in values) + + def ratio_above(value, threshold): """Return the ratio of values that are >= threshold. @@ -148,6 +155,7 @@ def evaluate_check(data: dict, check: str) -> tuple[bool, str, object]: "max": max, "mean": mean, "median": median, + "all_finite": all_finite, "ratio_above": ratio_above, } diff --git a/tests/functional/L1_Functional_Tests_GB200_Megatron_Omni.sh b/tests/functional/L1_Functional_Tests_GB200_Megatron_Omni.sh new file mode 100755 index 00000000000..c8f916ef053 --- /dev/null +++ b/tests/functional/L1_Functional_Tests_GB200_Megatron_Omni.sh @@ -0,0 +1,38 @@ +#!/bin/bash +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +set -xeuo pipefail + +SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &> /dev/null && pwd) +PROJECT_ROOT=$(realpath "${SCRIPT_DIR}/../..") + +cd "${PROJECT_ROOT}" + +GPU_COUNT=$(nvidia-smi --query-gpu=index --format=csv,noheader | wc -l) +if (( GPU_COUNT < 2 )); then + echo "SKIP: Nemotron Omni functional tests require at least two GB200 GPUs" + exit 0 +fi + +# Both tests colocate TP2/EP2 training and generation on two GB200 GPUs. +export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1}" + +time uv run --no-sync bash ./tests/functional/nemotron_omni_clevr_megatron_1n2g.sh +time uv run --no-sync bash ./tests/functional/nemotron_omni_gym_video_megatron_1n2g.sh + +cd "${PROJECT_ROOT}/tests" +if compgen -G ".coverage*" > /dev/null; then + coverage combine .coverage* +fi diff --git a/tests/functional/nemotron_omni_clevr_megatron_1n2g.sh b/tests/functional/nemotron_omni_clevr_megatron_1n2g.sh new file mode 100755 index 00000000000..0a61b7b71c6 --- /dev/null +++ b/tests/functional/nemotron_omni_clevr_megatron_1n2g.sh @@ -0,0 +1,167 @@ +#!/usr/bin/env bash +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +set -euo pipefail + +SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd) +PROJECT_ROOT=$(realpath "${SCRIPT_DIR}/../..") + +if [[ -z "${HF_TOKEN:-}" ]]; then + echo "SKIP: HF_TOKEN is required for the Omni checkpoint" + exit 0 +fi + +GPU_COUNT=$(nvidia-smi --query-gpu=index --format=csv,noheader | wc -l) +if (( GPU_COUNT < 2 )); then + echo "SKIP: Omni CLEVR Megatron smoke requires at least two visible GPUs" + exit 0 +fi +DETECTED_CUDA_ARCH=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader -i 0) +export TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST:-${DETECTED_CUDA_ARCH}}" +MEGATRON_TRANSFORMER_IMPL="${MEGATRON_TRANSFORMER_IMPL:-inference_optimized}" +MEGATRON_CUDA_GRAPH_IMPL="${MEGATRON_CUDA_GRAPH_IMPL:-local}" +if [[ "${MEGATRON_CUDA_GRAPH_IMPL}" == "local" ]]; then + INFERENCE_CUDA_GRAPH_SCOPE=block + NUM_CUDA_GRAPHS=-1 +else + INFERENCE_CUDA_GRAPH_SCOPE=none + NUM_CUDA_GRAPHS=0 +fi +if [[ "${MEGATRON_TRANSFORMER_IMPL}" != "inference_optimized" && + "${MEGATRON_CUDA_GRAPH_IMPL}" == "local" ]]; then + MOE_PAD_EXPERTS_FOR_CG=true +else + MOE_PAD_EXPERTS_FOR_CG=false +fi + +EXP_NAME=$(basename "$0" .sh) +EXP_DIR="${SCRIPT_DIR}/${EXP_NAME}" +LOG_DIR="${EXP_DIR}/logs" +DATA_ROOT="${EXP_DIR}/data" +TRAIN_PATH="${DATA_ROOT}/train.jsonl" +VAL_PATH="${DATA_ROOT}/val.jsonl" +JSON_METRICS="${EXP_DIR}/metrics.json" +RUN_LOG="${EXP_DIR}/run.log" +rm -rf "${EXP_DIR}" +mkdir -p "${LOG_DIR}" "${DATA_ROOT}" + +cd "${PROJECT_ROOT}" +export PYTHONPATH="${PROJECT_ROOT}:${PYTHONPATH:-}" + +# Use a tiny local image dataset. Downloading the full 70K CLEVR training split +# adds several minutes to a one-step smoke and does not improve E2E coverage. +# TODO(@cspades): Replace Omni 30B with a smaller pretrained model. +# For now, just use this as a partially-trainable functional test +# for inference and multimodal RL. +TRAIN_PATH="${TRAIN_PATH}" VAL_PATH="${VAL_PATH}" uv run --no-sync python - <<'PY' +import base64 +import io +import json +import os + +from PIL import Image + +buffer = io.BytesIO() +Image.new("RGB", (224, 224), color="red").save(buffer, format="PNG") +image_url = "data:image/png;base64," + base64.b64encode(buffer.getvalue()).decode() + +def sample(index: int) -> dict: + return { + "messages": [ + { + "role": "user", + "content": [ + {"type": "image", "image": image_url}, + { + "type": "text", + "text": f"Sample {index}: What color is the image?", + }, + ], + }, + {"role": "assistant", "content": "red"}, + ] + } + +for path, count in ((os.environ["TRAIN_PATH"], 64), (os.environ["VAL_PATH"], 2)): + with open(path, "w") as output: + for index in range(count): + output.write(json.dumps(sample(index)) + "\n") +PY + +# TODO(@cspades): Replace Omni 30B with a smaller pretrained model. +# For now, just use this as a partially-trainable functional test +# (frozen decoder trunk) for inference and multimodal RL. +uv run --no-sync python examples/run_vlm_grpo.py \ + --config examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron_generation.v1.yaml \ + cluster.num_nodes=1 \ + cluster.gpus_per_node=2 \ + policy.megatron_cfg.env_vars.TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST}" \ + policy.megatron_cfg.tensor_model_parallel_size=2 \ + policy.megatron_cfg.expert_model_parallel_size=2 \ + policy.megatron_cfg.expert_tensor_parallel_size=1 \ + policy.megatron_cfg.context_parallel_size=1 \ + policy.megatron_cfg.sequence_parallel=true \ + policy.megatron_cfg.activation_checkpointing=true \ + ++policy.megatron_cfg.freeze_config.freeze_language_model=true \ + +policy.megatron_cfg.bias_dropout_fusion=false \ + policy.megatron_cfg.optimizer.optimizer_cpu_offload=false \ + policy.megatron_cfg.optimizer.optimizer_offload_fraction=0.0 \ + ++policy.megatron_cfg.optimizer.params_dtype=bfloat16 \ + ++policy.megatron_cfg.optimizer.main_grads_dtype=bfloat16 \ + ++policy.megatron_cfg.optimizer.main_params_dtype=float16 \ + ++policy.megatron_cfg.optimizer.exp_avg_dtype=bfloat16 \ + ++policy.megatron_cfg.optimizer.exp_avg_sq_dtype=bfloat16 \ + ++policy.megatron_cfg.optimizer.store_param_remainders=false \ + policy.generation.backend=megatron \ + policy.generation.colocated.enabled=true \ + policy.generation.colocated.resources.num_nodes=1 \ + policy.generation.colocated.resources.gpus_per_node=2 \ + policy.generation.max_new_tokens=128 \ + policy.generation.mcore_generation_config.tensor_model_parallel_size=2 \ + policy.generation.mcore_generation_config.expert_model_parallel_size=2 \ + policy.generation.mcore_generation_config.expert_tensor_parallel_size=1 \ + ++policy.generation.mcore_generation_config.context_parallel_size=1 \ + ++policy.generation.mcore_generation_config.moe_router_dtype=fp32 \ + policy.generation.mcore_generation_config.transformer_impl="${MEGATRON_TRANSFORMER_IMPL}" \ + policy.generation.mcore_generation_config.sequence_parallel=true \ + policy.generation.mcore_generation_config.refit_backend=nccl \ + policy.generation.mcore_generation_config.buffer_size_gb=2 \ + policy.generation.mcore_generation_config.cuda_graph_impl="${MEGATRON_CUDA_GRAPH_IMPL}" \ + policy.generation.mcore_generation_config.inference_cuda_graph_scope="${INFERENCE_CUDA_GRAPH_SCOPE}" \ + policy.generation.mcore_generation_config.num_cuda_graphs="${NUM_CUDA_GRAPHS}" \ + policy.generation.mcore_generation_config.use_cuda_graphs_for_non_decode_steps=false \ + policy.generation.mcore_generation_config.moe_pad_experts_for_cuda_graph_inference="${MOE_PAD_EXPERTS_FOR_CG}" \ + policy.generation.mcore_generation_config.enable_chunked_prefill=true \ + ++policy.generation.mcore_generation_config.async_sched_mode=async \ + policy.generation.mcore_generation_config.max_model_len=1024 \ + policy.generation.mcore_generation_config.max_tokens=1024 \ + policy.max_total_sequence_length=1024 \ + data.train.dataset_name=ResponseDataset \ + ++data.train.data_path="${TRAIN_PATH}" \ + data.train.split=train \ + data.validation.dataset_name=ResponseDataset \ + ++data.validation.data_path="${VAL_PATH}" \ + data.validation.split=train \ + data.num_workers=0 \ + grpo.async_grpo.enabled=true \ + grpo.async_grpo.max_trajectory_age_steps=2 \ + grpo.async_grpo.in_flight_weight_updates=true \ + grpo.num_prompts_per_step=1 \ + grpo.num_generations_per_prompt=2 \ + grpo.max_num_steps=1 \ + grpo.val_period=0 \ + grpo.val_at_start=false \ + grpo.val_at_end=false \ + policy.train_global_batch_size=2 \ + policy.train_micro_batch_size=1 \ + logger.tensorboard_enabled=true \ + logger.log_dir="${LOG_DIR}" \ + logger.wandb_enabled=false \ + logger.monitor_gpus=false \ + checkpointing.enabled=false \ + "$@" 2>&1 | tee "${RUN_LOG}" + +uv run --no-sync tests/json_dump_tb_logs.py "${LOG_DIR}" --output_path "${JSON_METRICS}" +uv run --no-sync tests/check_metrics.py "${JSON_METRICS}" \ + 'max(data["train/gen_kl_error"]) < 0.05' \ + 'all_finite(data["train/reward"])' diff --git a/tests/functional/nemotron_omni_gym_video_megatron_1n2g.sh b/tests/functional/nemotron_omni_gym_video_megatron_1n2g.sh new file mode 100755 index 00000000000..3fd6aad2675 --- /dev/null +++ b/tests/functional/nemotron_omni_gym_video_megatron_1n2g.sh @@ -0,0 +1,165 @@ +#!/usr/bin/env bash +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +set -euo pipefail + +SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd) +PROJECT_ROOT=$(realpath "${SCRIPT_DIR}/../..") + +if [[ -z "${HF_TOKEN:-}" ]]; then + echo "SKIP: HF_TOKEN is required for the Omni checkpoint" + exit 0 +fi + +GPU_COUNT=$(nvidia-smi --query-gpu=index --format=csv,noheader | wc -l) +if (( GPU_COUNT < 2 )); then + echo "SKIP: Omni Gym-video Megatron smoke requires at least two GPUs" + exit 0 +fi +DETECTED_CUDA_ARCH=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader -i 0) +export TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST:-${DETECTED_CUDA_ARCH}}" +MEGATRON_TRANSFORMER_IMPL="${MEGATRON_TRANSFORMER_IMPL:-inference_optimized}" +MEGATRON_CUDA_GRAPH_IMPL="${MEGATRON_CUDA_GRAPH_IMPL:-local}" +if [[ "${MEGATRON_CUDA_GRAPH_IMPL}" == "local" ]]; then + INFERENCE_CUDA_GRAPH_SCOPE=block + NUM_CUDA_GRAPHS=-1 +else + INFERENCE_CUDA_GRAPH_SCOPE=none + NUM_CUDA_GRAPHS=0 +fi +if [[ "${MEGATRON_TRANSFORMER_IMPL}" != "inference_optimized" && + "${MEGATRON_CUDA_GRAPH_IMPL}" == "local" ]]; then + MOE_PAD_EXPERTS_FOR_CG=true +else + MOE_PAD_EXPERTS_FOR_CG=false +fi + +EXP_NAME=$(basename "$0" .sh) +EXP_DIR="${SCRIPT_DIR}/${EXP_NAME}" +LOG_DIR="${EXP_DIR}/logs" +DATA_ROOT="${EXP_DIR}/data" +VIDEO_PATH="${DATA_ROOT}/red.mp4" +RAW_TRAIN_PATH="${DATA_ROOT}/train-raw.jsonl" +RAW_VAL_PATH="${DATA_ROOT}/val-raw.jsonl" +TRAIN_PATH="${DATA_ROOT}/train-gym.jsonl" +VAL_PATH="${DATA_ROOT}/val-gym.jsonl" +JSON_METRICS="${EXP_DIR}/metrics.json" +RUN_LOG="${EXP_DIR}/run.log" +rm -rf "${EXP_DIR}" +mkdir -p "${LOG_DIR}" "${DATA_ROOT}" + +cd "${PROJECT_ROOT}" +export PYTHONPATH="${PROJECT_ROOT}:${PYTHONPATH:-}" +export NRL_VIDEO_BACKEND=torchcodec +export NRL_VIDEO_SAMPLING_STYLE=nemotron_vl +export NRL_VIDEO_TEMPORAL_PATCH_SIZE=2 + +bash tools/install_audio_deps.sh +ffmpeg -hide_banner -loglevel error -y \ + -f lavfi -i color=c=red:s=224x224:r=8:d=2 \ + -c:v libx264 -pix_fmt yuv420p "${VIDEO_PATH}" + +for sample_id in $(seq 1 64); do + jq -nc \ + --arg prompt "Sample ${sample_id}: What color fills the video? A. Red B. Blue" \ + --arg video "${VIDEO_PATH}" \ + '{prompt: $prompt, video: $video, answer: "A", verifier: "mcqa"}' +done > "${RAW_TRAIN_PATH}" +for sample_id in $(seq 1 2); do + jq -nc \ + --arg prompt "Validation ${sample_id}: What color fills the video? A. Red B. Blue" \ + --arg video "${VIDEO_PATH}" \ + '{prompt: $prompt, video: $video, answer: "A", verifier: "mcqa"}' +done > "${RAW_VAL_PATH}" + +uv run --no-sync examples/nemo_gym/prepare_video_dataset.py convert \ + --input "${RAW_TRAIN_PATH}" \ + --output "${TRAIN_PATH}" +uv run --no-sync examples/nemo_gym/prepare_video_dataset.py convert \ + --input "${RAW_VAL_PATH}" \ + --output "${VAL_PATH}" + +# TODO(@cspades): Replace Omni 30B with a smaller pretrained model. +# For now, just use this as a partially-trainable functional test +# (frozen decoder trunk) for inference and multimodal RL. +uv run --no-sync python examples/nemo_gym/run_grpo_nemo_gym.py \ + --config examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-16n8g-megatron-tp4ep4-async-gym-video.v1.yaml \ + cluster.num_nodes=1 \ + cluster.gpus_per_node=2 \ + policy.megatron_cfg.env_vars.TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST}" \ + policy.megatron_cfg.tensor_model_parallel_size=2 \ + policy.megatron_cfg.pipeline_model_parallel_size=1 \ + policy.megatron_cfg.expert_model_parallel_size=2 \ + policy.megatron_cfg.expert_tensor_parallel_size=1 \ + policy.megatron_cfg.context_parallel_size=1 \ + policy.megatron_cfg.sequence_parallel=true \ + policy.megatron_cfg.activation_checkpointing=true \ + ++policy.megatron_cfg.freeze_config.freeze_language_model=true \ + policy.megatron_cfg.optimizer.optimizer_cpu_offload=false \ + policy.megatron_cfg.optimizer.optimizer_offload_fraction=0.0 \ + ++policy.megatron_cfg.optimizer.params_dtype=bfloat16 \ + ++policy.megatron_cfg.optimizer.main_grads_dtype=bfloat16 \ + ++policy.megatron_cfg.optimizer.main_params_dtype=float16 \ + ++policy.megatron_cfg.optimizer.exp_avg_dtype=bfloat16 \ + ++policy.megatron_cfg.optimizer.exp_avg_sq_dtype=bfloat16 \ + ++policy.megatron_cfg.optimizer.store_param_remainders=false \ + policy.generation.backend=megatron \ + ++policy.generation.bad_words=null \ + policy.generation.colocated.enabled=true \ + policy.generation.colocated.resources.num_nodes=1 \ + policy.generation.colocated.resources.gpus_per_node=2 \ + policy.generation.max_new_tokens=128 \ + policy.generation.mcore_generation_config.expose_http_server=true \ + policy.generation.mcore_generation_config.tensor_model_parallel_size=2 \ + policy.generation.mcore_generation_config.expert_model_parallel_size=2 \ + policy.generation.mcore_generation_config.expert_tensor_parallel_size=1 \ + ++policy.generation.mcore_generation_config.context_parallel_size=1 \ + ++policy.generation.mcore_generation_config.moe_router_dtype=fp32 \ + policy.generation.mcore_generation_config.transformer_impl="${MEGATRON_TRANSFORMER_IMPL}" \ + policy.generation.mcore_generation_config.sequence_parallel=true \ + policy.generation.mcore_generation_config.refit_backend=nccl \ + policy.generation.mcore_generation_config.buffer_size_gb=2 \ + policy.generation.mcore_generation_config.cuda_graph_impl="${MEGATRON_CUDA_GRAPH_IMPL}" \ + policy.generation.mcore_generation_config.inference_cuda_graph_scope="${INFERENCE_CUDA_GRAPH_SCOPE}" \ + policy.generation.mcore_generation_config.num_cuda_graphs="${NUM_CUDA_GRAPHS}" \ + policy.generation.mcore_generation_config.use_cuda_graphs_for_non_decode_steps=false \ + ++policy.generation.mcore_generation_config.moe_pad_experts_for_cuda_graph_inference="${MOE_PAD_EXPERTS_FOR_CG}" \ + policy.generation.mcore_generation_config.enable_chunked_prefill=true \ + ++policy.generation.mcore_generation_config.async_sched_mode=async \ + policy.generation.mcore_generation_config.enable_prefix_caching=true \ + policy.generation.mcore_generation_config.max_model_len=1024 \ + policy.generation.mcore_generation_config.max_tokens=1024 \ + ++policy.generation.mcore_generation_config.video_num_frames=8 \ + ++policy.generation.mcore_generation_config.video_temporal_patch_size=2 \ + ++policy.generation.mcore_generation_config.video_target_num_patches=256 \ + policy.max_total_sequence_length=1024 \ + +data.default.num_frames=8 \ + +data.default.video_sampling_style=nemotron_vl \ + +data.default.video_temporal_patch_size=2 \ + +data.default.min_generation_tokens=128 \ + data.default.video_target_num_patches=256 \ + data.train.data_path="${TRAIN_PATH}" \ + data.validation.data_path="${VAL_PATH}" \ + grpo.deduplicate_multimodal_data=false \ + grpo.async_grpo.enabled=true \ + grpo.async_grpo.max_trajectory_age_steps=2 \ + grpo.async_grpo.in_flight_weight_updates=true \ + grpo.num_prompts_per_step=1 \ + grpo.num_generations_per_prompt=2 \ + grpo.max_num_steps=1 \ + grpo.val_period=0 \ + grpo.val_at_start=false \ + grpo.val_at_end=false \ + policy.train_global_batch_size=2 \ + policy.train_micro_batch_size=1 \ + logger.tensorboard_enabled=true \ + logger.log_dir="${LOG_DIR}" \ + logger.wandb_enabled=false \ + logger.monitor_gpus=false \ + checkpointing.enabled=false \ + "$@" 2>&1 | tee "${RUN_LOG}" + +uv run --no-sync tests/json_dump_tb_logs.py "${LOG_DIR}" --output_path "${JSON_METRICS}" +uv run --no-sync tests/check_metrics.py "${JSON_METRICS}" \ + 'max(data["train/gen_kl_error"]) < 0.05' \ + 'all_finite(data["train/reward"])' diff --git a/tests/test_suites/disabled.txt b/tests/test_suites/disabled.txt index c06a9f253b9..44e2cf0e5af 100644 --- a/tests/test_suites/disabled.txt +++ b/tests/test_suites/disabled.txt @@ -27,6 +27,10 @@ tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-2n8g-megatron-tp4ep4-gym-vid # budget has room. tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-click-2n8g-megatron-tp2ep8.v1.sh +# TODO(@cspades): Run and validate this multimodal Megatron generation +# functional test before moving it to a recurring suite. +tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-count-1n4g-megatron_generation.v1.sh + # Nemotron Super Omni: 16-node topology, and the checkpoint and multimodal Gym # blend are too large to ship with the repo, so these are invoked manually via # examples/nemo_gym/nemotron-3-super-omni/super_omni_launch.sh rather than run diff --git a/tests/test_suites/nightly_gb200.txt b/tests/test_suites/nightly_gb200.txt index c6ac7ae01a7..a1850e81ace 100644 --- a/tests/test_suites/nightly_gb200.txt +++ b/tests/test_suites/nightly_gb200.txt @@ -32,6 +32,7 @@ tests/test_suites/llm/grpo-nanov3-30BA3B-4n4g-megatron_colocated_reshard.sh # Functional VLM run tests/test_suites/vlm/vlm_grpo-qwen2.5-vl-3b-instruct-clevr-1n4g-dtensor2tp1.v1.sh tests/test_suites/vlm/vlm_grpo-qwen2.5-vl-3b-instruct-clevr-1n4g-megatrontp1.v1.sh +tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron_generation.v1.sh # Deepscaler (short tests) tests/test_suites/llm/grpo-deepscaler-1.5b-1n4g-8K.sh diff --git a/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-count-1n4g-megatron_generation.v1.sh b/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-count-1n4g-megatron_generation.v1.sh new file mode 100755 index 00000000000..f8d2e75cdf8 --- /dev/null +++ b/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-count-1n4g-megatron_generation.v1.sh @@ -0,0 +1,34 @@ +#!/bin/bash +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +source $SCRIPT_DIR/common.env + +# TODO(@cspades): Run and validate this functional test, then add golden +# convergence metrics before enabling it in a recurring suite. + +# ===== BEGIN CONFIG ===== +NUM_NODES=1 +GPUS_PER_NODE=4 +STEPS_PER_RUN=4 +MAX_STEPS=4 +NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) +NUM_MINUTES=120 +# ===== END CONFIG ===== + +exit_if_max_steps_reached + +cd $PROJECT_ROOT +uv run examples/run_vlm_grpo.py \ + --config $CONFIG_PATH \ + grpo.max_num_steps=$MAX_STEPS \ + logger.log_dir=$LOG_DIR \ + logger.wandb_enabled=True \ + logger.wandb.project=nemo-rl \ + logger.wandb.name=$EXP_NAME \ + logger.monitor_gpus=True \ + logger.tensorboard_enabled=True \ + checkpointing.enabled=True \ + checkpointing.checkpoint_dir=$CKPT_DIR \ + $@ \ + 2>&1 | tee $RUN_LOG + +uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS diff --git a/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron_generation.v1.sh b/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron_generation.v1.sh new file mode 100755 index 00000000000..bc36a391470 --- /dev/null +++ b/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron_generation.v1.sh @@ -0,0 +1,47 @@ +#!/bin/bash +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +source $SCRIPT_DIR/common.env + +# Compare the pretrained baseline against 50 policy updates. Require active, +# numerically healthy training and a modest validation improvement; CLEVR +# reward remains too noisy over 50 steps to require a monotonic reward trend. + +# ===== BEGIN CONFIG ===== +NUM_NODES=8 +GPUS_PER_NODE=4 +STEPS_PER_RUN=50 +MAX_STEPS=50 +NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) +NUM_MINUTES=120 +# ===== END CONFIG ===== + +exit_if_max_steps_reached + +cd $PROJECT_ROOT +uv run examples/run_vlm_grpo.py \ + --config $CONFIG_PATH \ + grpo.max_num_steps=$MAX_STEPS \ + policy.megatron_cfg.scheduler.lr_warmup_iters=10 \ + logger.log_dir=$LOG_DIR \ + logger.wandb_enabled=True \ + logger.wandb.project=nemo-rl \ + logger.wandb.name=$EXP_NAME \ + logger.monitor_gpus=True \ + logger.tensorboard_enabled=True \ + checkpointing.enabled=False \ + checkpointing.checkpoint_dir=$CKPT_DIR \ + $@ \ + 2>&1 | tee $RUN_LOG + +uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS + +uv run tests/check_metrics.py "$JSON_METRICS" \ + 'all_finite(data["train/loss"])' \ + 'all_finite(data["train/grad_norm"])' \ + 'min(data["train/grad_norm"]) > 0' \ + 'all_finite(data["train/token_mult_prob_error"])' \ + 'mean(data["train/reward"], range_start=-10) > 0.6' \ + '"0" in data["validation/accuracy"]' \ + '"50" in data["validation/accuracy"]' \ + 'data["validation/accuracy"]["50"] > 0.6' \ + 'data["validation/accuracy"]["50"] > data["validation/accuracy"]["0"] + 0.01' diff --git a/tests/unit/algorithms/test_grpo.py b/tests/unit/algorithms/test_grpo.py index f2f6f875fe4..3fca2187900 100644 --- a/tests/unit/algorithms/test_grpo.py +++ b/tests/unit/algorithms/test_grpo.py @@ -1236,16 +1236,11 @@ def mock_async_grpo_infrastructure( checkpoint_cut_ordinal=checkpoint_cut_ordinal, ) - # Patch venv creation + # Patch actor runtime environment creation stack.enter_context( patch( - "nemo_rl.algorithms.grpo.create_local_venv_on_each_node", - return_value="/fake/venv", - ) - ) - stack.enter_context( - patch( - "nemo_rl.algorithms.grpo.get_actor_python_env", return_value="/fake/python" + "nemo_rl.algorithms.grpo.make_actor_runtime_env", + return_value={"py_executable": "/fake/python", "env_vars": {}}, ) ) diff --git a/tests/unit/data/datasets/test_audiomcq_dataset.py b/tests/unit/data/datasets/test_audiomcq_dataset.py index 5168dd20012..b95cb9b0886 100644 --- a/tests/unit/data/datasets/test_audiomcq_dataset.py +++ b/tests/unit/data/datasets/test_audiomcq_dataset.py @@ -427,22 +427,25 @@ def test_vlm_hf_data_processor_returns_audiomcq_datum_spec( assert result["extra_env_info"]["ground_truth"] == rows[0]["answer"] assert result["extra_env_info"]["choices"] == rows[0]["choices"] - def test_dispatcher_rejects_unknown_task_name(self): + def test_dispatcher_accepts_generic_preformatted_messages(self): from nemo_rl.data.interfaces import TaskDataSpec from nemo_rl.data.processors import vlm_hf_data_processor - bogus_datum = { - "task_name": "definitely-not-a-task", + generic_datum = { + "task_name": "custom-vlm-task", "messages": [ {"role": "user", "content": [{"type": "text", "text": "hi"}]}, {"role": "assistant", "content": "hello"}, ], } - with pytest.raises(ValueError, match="No data processor for task"): - vlm_hf_data_processor( - datum_dict=bogus_datum, - task_data_spec=TaskDataSpec(task_name="definitely-not-a-task"), - processor=_FakeProcessor(), - max_seq_length=4096, - idx=0, - ) + result = vlm_hf_data_processor( + datum_dict=generic_datum, + task_data_spec=TaskDataSpec(task_name="custom-vlm-task"), + processor=_FakeProcessor(), + max_seq_length=4096, + idx=0, + ) + + assert result["task_name"] == "custom-vlm-task" + assert result["extra_env_info"]["ground_truth"] == "hello" + assert result["vllm_content"] == "fake" diff --git a/tests/unit/data/datasets/test_mmpr_tiny.py b/tests/unit/data/datasets/test_mmpr_tiny.py index 4b0c9e4fc28..7ec4c34a538 100644 --- a/tests/unit/data/datasets/test_mmpr_tiny.py +++ b/tests/unit/data/datasets/test_mmpr_tiny.py @@ -160,7 +160,10 @@ def apply_chat_template(self, messages, **kwargs): for item in content: if isinstance(item, dict) and "text" in item: parts.append(item["text"]) - return " ".join(parts) + formatted_text = " ".join(parts) + if kwargs.get("tokenize"): + return {"input_ids": fake_input_ids} + return formatted_text def __call__(self, text=None, images=None, **kwargs): self.captured_call_text = text @@ -241,6 +244,31 @@ def test_processor_produces_valid_datum_spec(self, tiny_image_path): assert user_message["pixel_values"].pad_to_max_shape is True assert user_message["pixel_values"].as_tensor().dtype == torch.float32 + def test_text_only_row_preserves_formatted_vllm_content(self): + from nemo_rl.data.interfaces import TaskDataSpec + from nemo_rl.data.processors import vlm_hf_data_processor + + processor = _make_stub_nemotron_processor() + task_data_spec = TaskDataSpec(task_name="text-only") + task_data_spec.prompt = "Answer: {}" + + result = vlm_hf_data_processor( + datum_dict={ + "messages": [ + {"role": "user", "content": "What is 2 + 2?"}, + {"role": "assistant", "content": "4"}, + ], + "task_name": "text-only", + }, + task_data_spec=task_data_spec, + processor=processor, + max_seq_length=8192, + idx=0, + ) + + assert result["vllm_content"] == "Answer: What is 2 + 2?" + assert result["vllm_images"] == [] + def test_conversation_preprocessor_is_preserved(self, tiny_image_path): processor = _make_stub_nemotron_processor() processor.conversation_preprocessor = MagicMock( @@ -250,7 +278,7 @@ def test_conversation_preprocessor_is_preserved(self, tiny_image_path): result, _ = _run_processor(tiny_image_path, processor=processor) processor.conversation_preprocessor.assert_called_once() - assert result["vllm_content"] == "preprocessed" + assert result["vllm_content"] is None assert processor.captured_call_text == "preprocessed" def test_historical_tiled_processor_gets_media_metadata(self, tiny_image_path): @@ -285,14 +313,14 @@ def test_historical_tiled_processor_gets_media_metadata(self, tiny_image_path): def test_prompted_text_contains_boxed_literal_and_no_raw_dataset_string( self, tiny_image_path ): - result, _ = _run_processor(tiny_image_path) - vllm_content = result["vllm_content"] + result, processor = _run_processor(tiny_image_path) + processed_text = processor.captured_call_text # Positive: literal \boxed{} must survive prompt formatting - assert "\\boxed{}" in vllm_content + assert "\\boxed{}" in processed_text # Negative: the raw dataset string (with prefix) must NOT leak through - assert _RAW_QUESTION not in vllm_content + assert _RAW_QUESTION not in processed_text def test_placeholder_conversion_exact_string(self, tiny_image_path): """Verify the exact tokenizer input for the placeholder-style processor path. @@ -310,15 +338,15 @@ def test_placeholder_conversion_exact_string(self, tiny_image_path): # The stub's apply_chat_template joins message parts with spaces, # so the captured text passed to __call__ is the chat-templated string. - # Verify the vllm_content (which is apply_chat_template output) matches. - vllm_content = result["vllm_content"] - assert vllm_content == expected_tokenizer_input + # Verify the apply_chat_template output through captured_call_text below; + # placeholder-style processors send expanded token IDs to vLLM. + assert result["vllm_content"] is None # Verify exactly one token in the final output - assert vllm_content.count("") == 1 + assert processor.captured_call_text.count("") == 1 # Verify the question text is present - assert _CLEAN_QUESTION in vllm_content + assert _CLEAN_QUESTION in processor.captured_call_text # Verify the captured __call__ text also matches # (processor.__call__ receives the apply_chat_template output) diff --git a/tests/unit/data/test_collate_fn.py b/tests/unit/data/test_collate_fn.py index e531f4cb5e7..194b798f8e6 100755 --- a/tests/unit/data/test_collate_fn.py +++ b/tests/unit/data/test_collate_fn.py @@ -16,7 +16,11 @@ import torch -from nemo_rl.data.collate_fn import preference_collate_fn +from nemo_rl.data.collate_fn import ( + eval_collate_fn, + preference_collate_fn, + rl_collate_fn, +) from nemo_rl.data.interfaces import DatumSpec from nemo_rl.distributed.batched_data_dict import BatchedDataDict @@ -149,3 +153,23 @@ def test_preference_collate_fn(): assert torch.equal( train_data["input_ids"][1][3:5], torch.tensor([8, 9]) ) # assistant + + +def test_collate_preserves_native_media_when_vllm_content_is_none(): + image = object() + datum = DatumSpec( + message_log=[], + length=1, + loss_multiplier=1.0, + extra_env_info={}, + idx=0, + task_name="vlm", + vllm_content=None, + vllm_images=[image], + vllm_audios=[], + vllm_videos=[], + ) + + for batch in (rl_collate_fn([datum]), eval_collate_fn([datum])): + assert batch["vllm_content"] == [None] + assert batch["vllm_images"] == [[image]] diff --git a/tests/unit/data/test_multimodal_image_encoding.py b/tests/unit/data/test_multimodal_image_encoding.py index fda52a35ca6..4d7f9b23322 100644 --- a/tests/unit/data/test_multimodal_image_encoding.py +++ b/tests/unit/data/test_multimodal_image_encoding.py @@ -17,12 +17,14 @@ import pytest from PIL import Image -import nemo_rl.data.multimodal_utils as multimodal_utils +import nemo_rl.environments.nemo_gym_multimodal as nemo_gym_multimodal from nemo_rl.data.multimodal_utils import ( - encode_images_in_examples, image_to_data_url, resolve_to_image, ) +from nemo_rl.environments.nemo_gym_multimodal import ( + normalize_media_in_examples, +) def _example(*content_parts: dict) -> dict: @@ -59,7 +61,7 @@ def test_resolve_to_image_accepts_file_scheme(tmp_path): assert resolve_to_image(path).size == (5, 6) -def test_encode_images_encodes_local_paths_and_file_urls(tmp_path): +def test_normalize_media_encodes_local_image_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)) @@ -70,7 +72,7 @@ def test_encode_images_encodes_local_paths_and_file_urls(tmp_path): {"type": "input_text", "text": "describe"}, ) ] - encode_images_in_examples(examples) + normalize_media_in_examples(examples) parts = examples[0]["responses_create_params"]["input"][0]["content"] assert parts[0]["image_url"].startswith("data:image/png;base64,") @@ -81,7 +83,7 @@ def test_encode_images_encodes_local_paths_and_file_urls(tmp_path): assert parts[2] == {"type": "input_text", "text": "describe"} -def test_encode_images_passes_through_http_and_data_urls(): +def test_normalize_media_passes_through_http_and_data_urls(): data_url = image_to_data_url(Image.new("RGB", (2, 2))) examples = [ _example( @@ -90,7 +92,7 @@ def test_encode_images_passes_through_http_and_data_urls(): {"type": "input_image", "image_url": data_url}, ) ] - encode_images_in_examples(examples) + normalize_media_in_examples(examples) parts = examples[0]["responses_create_params"]["input"][0]["content"] assert parts[0]["image_url"] == "https://example.com/cat.png" @@ -104,7 +106,7 @@ def test_encode_images_deduplicates_sources_and_uses_a_bounded_thread_pool( first = _write_png(tmp_path, "first.png", (2, 3)) second = _write_png(tmp_path, "second.png", (4, 5)) expected = { - source: multimodal_utils._encode_single_image_source(source) + source: nemo_gym_multimodal._encode_single_image_source(source) for source in (first, second) } examples = [ @@ -117,31 +119,117 @@ def test_encode_images_deduplicates_sources_and_uses_a_bounded_thread_pool( ] resolve_calls = [] - original_resolve = multimodal_utils.resolve_to_image + original_resolve = nemo_gym_multimodal.resolve_to_image def tracking_resolve(source): resolve_calls.append(source) return original_resolve(source) observed_max_workers = [] - original_executor = multimodal_utils.ThreadPoolExecutor + original_executor = nemo_gym_multimodal.ThreadPoolExecutor def tracking_executor(*args, **kwargs): observed_max_workers.append(kwargs.get("max_workers")) return original_executor(*args, **kwargs) - monkeypatch.setattr(multimodal_utils, "resolve_to_image", tracking_resolve) - monkeypatch.setattr(multimodal_utils, "ThreadPoolExecutor", tracking_executor) + monkeypatch.setattr(nemo_gym_multimodal, "resolve_to_image", tracking_resolve) + monkeypatch.setattr(nemo_gym_multimodal, "ThreadPoolExecutor", tracking_executor) - encode_images_in_examples(examples) + normalize_media_in_examples(examples) assert Counter(resolve_calls) == Counter({first: 1, second: 1}) - assert observed_max_workers == [multimodal_utils.NEMO_GYM_IMAGE_ENCODE_MAX_WORKERS] + assert observed_max_workers == [ + nemo_gym_multimodal._NEMO_GYM_IMAGE_ENCODE_MAX_WORKERS + ] for example in examples: parts = example["responses_create_params"]["input"][0]["content"] - assert parts[0]["image_url"] == expected[first] - assert parts[1]["image"] == expected[first] - assert parts[2]["url"] == expected[second] + assert parts[0] == { + "type": "input_image", + "image_url": expected[first], + } + assert parts[1] == { + "type": "input_image", + "image_url": expected[first], + } + assert parts[2] == { + "type": "input_image", + "image_url": expected[second], + } + + +@pytest.mark.parametrize( + "part", + [ + {"type": "image", "image": "a.png", "url": "b.png"}, + {"type": "image"}, + ], +) +def test_normalize_media_requires_exactly_one_source_key(part): + with pytest.raises(ValueError, match="requires exactly one"): + normalize_media_in_examples([_example(part)]) + + +@pytest.mark.parametrize( + "source", + ["", {"url": ""}], +) +def test_normalize_media_rejects_empty_urls(source): + with pytest.raises(ValueError, match="requires a non-empty media URL"): + normalize_media_in_examples( + [_example({"type": "input_image", "image_url": source})] + ) + + +def test_normalize_media_preserves_input_image_file_id(): + part = {"type": "input_image", "file_id": "file-123", "detail": "high"} + examples = [_example(part)] + + normalize_media_in_examples(examples) + + assert examples[0]["responses_create_params"]["input"][0]["content"][0] == { + "type": "input_image", + "file_id": "file-123", + "detail": "high", + } + + +def test_normalize_media_promotes_nested_image_detail(): + data_url = image_to_data_url(Image.new("RGB", (2, 2))) + examples = [ + _example( + { + "type": "image", + "image": {"url": data_url, "detail": "high"}, + } + ) + ] + + normalize_media_in_examples(examples) + + assert examples[0]["responses_create_params"]["input"][0]["content"][0] == { + "type": "input_image", + "image_url": data_url, + "detail": "high", + } + + +def test_normalize_media_canonicalizes_local_video(tmp_path, monkeypatch): + video_path = tmp_path / "clip.mp4" + video_path.write_bytes(b"video") + encoded_url = "data:video/mp4;base64,dmlkZW8=" + monkeypatch.setattr( + nemo_gym_multimodal, + "video_path_to_data_url", + lambda source: encoded_url, + ) + examples = [_example({"type": "video", "url": str(video_path)})] + + normalize_media_in_examples(examples) + + assert examples[0]["responses_create_params"]["input"][0]["content"][0] == { + "type": "input_video", + "video_url": encoded_url, + } def test_encode_images_does_not_partially_mutate_on_error(tmp_path): @@ -155,7 +243,7 @@ def test_encode_images_does_not_partially_mutate_on_error(tmp_path): ] with pytest.raises(FileNotFoundError): - encode_images_in_examples(examples) + normalize_media_in_examples(examples) parts = examples[0]["responses_create_params"]["input"][0]["content"] assert parts[0]["image_url"] == existing @@ -164,15 +252,15 @@ def test_encode_images_does_not_partially_mutate_on_error(tmp_path): def test_encode_single_image_source_closes_image_on_success(monkeypatch): image = _CloseTrackingImage() - monkeypatch.setattr(multimodal_utils, "resolve_to_image", lambda _: image) + monkeypatch.setattr(nemo_gym_multimodal, "resolve_to_image", lambda _: image) monkeypatch.setattr( - multimodal_utils, + nemo_gym_multimodal, "image_to_data_url", lambda _: "data:image/png;base64,AA", ) assert ( - multimodal_utils._encode_single_image_source("image.png") + nemo_gym_multimodal._encode_single_image_source("image.png") == "data:image/png;base64,AA" ) assert image.closed @@ -180,27 +268,29 @@ def test_encode_single_image_source_closes_image_on_success(monkeypatch): def test_encode_single_image_source_closes_image_on_error(monkeypatch): image = _CloseTrackingImage() - monkeypatch.setattr(multimodal_utils, "resolve_to_image", lambda _: image) + monkeypatch.setattr(nemo_gym_multimodal, "resolve_to_image", lambda _: image) def fail_to_encode(_): raise RuntimeError("encoding failed") - monkeypatch.setattr(multimodal_utils, "image_to_data_url", fail_to_encode) + monkeypatch.setattr(nemo_gym_multimodal, "image_to_data_url", fail_to_encode) with pytest.raises(RuntimeError, match="encoding failed"): - multimodal_utils._encode_single_image_source("image.png") + nemo_gym_multimodal._encode_single_image_source("image.png") assert image.closed -def test_encode_images_is_a_noop_for_text_only_examples(): +def test_normalize_media_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 normalize_media_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"}}]) + assert ( + normalize_media_in_examples([{}, {"responses_create_params": {}}]) is not None + ) + assert normalize_media_in_examples([{"responses_create_params": {"input": "nope"}}]) diff --git a/tests/unit/distributed/test_worker_groups.py b/tests/unit/distributed/test_worker_groups.py index e1a2363792e..8ceaef0a2be 100644 --- a/tests/unit/distributed/test_worker_groups.py +++ b/tests/unit/distributed/test_worker_groups.py @@ -25,7 +25,32 @@ PY_EXECUTABLES, ) from nemo_rl.distributed.virtual_cluster import RayVirtualCluster -from nemo_rl.distributed.worker_groups import RayWorkerBuilder, RayWorkerGroup +from nemo_rl.distributed.worker_groups import ( + RayWorkerBuilder, + RayWorkerGroup, + _get_initializer_env_vars, +) + + +def test_initializer_env_adds_hf_modules_cache_to_pythonpath(monkeypatch): + monkeypatch.setenv("PYTHONPATH", "/driver/pythonpath") + + result = _get_initializer_env_vars({"HF_MODULES_CACHE": "/hf/modules"}) + + assert result["HF_MODULES_CACHE"] == "/hf/modules" + assert result["PYTHONPATH"].split(os.pathsep) == [ + "/hf/modules", + "/driver/pythonpath", + ] + + +def test_initializer_env_does_not_duplicate_hf_modules_cache(monkeypatch): + monkeypatch.setenv("HF_MODULES_CACHE", "/hf/modules") + monkeypatch.setenv("PYTHONPATH", f"/project{os.pathsep}/hf/modules") + + result = _get_initializer_env_vars({}) + + assert result["PYTHONPATH"] == f"/project{os.pathsep}/hf/modules" @ray.remote diff --git a/tests/unit/environments/test_nemo_gym.py b/tests/unit/environments/test_nemo_gym.py index 8f5a8e510ec..eefd569b72a 100644 --- a/tests/unit/environments/test_nemo_gym.py +++ b/tests/unit/environments/test_nemo_gym.py @@ -32,6 +32,7 @@ MULTIMODAL_CONTENT_TYPES, PackedTensor, image_to_data_url, + video_path_to_data_url, ) from nemo_rl.data.utils import setup_response_data from nemo_rl.distributed.ray_actor_environment_registry import ( @@ -45,12 +46,16 @@ setup_nemo_gym_config, validate_reward_components_match_scalar, ) -from nemo_rl.environments.nemo_gym_video import ( +from nemo_rl.environments.nemo_gym_multimodal import ( _extract_static_video_messages, _inject_vllm_mm_processor_kwargs, - _metadata_extra_body, nemo_gym_example_to_video_datum_spec, - normalize_video_urls_in_examples, + normalize_media_in_examples, +) +from nemo_rl.environments.nemo_gym_request import ( + _chat_template_kwargs_for_processor, + _deep_merge_dict, + _metadata_extra_body, ) from nemo_rl.environments.nemotron_utils import ( _expand_nemotron_video_placeholders, @@ -184,7 +189,7 @@ def test_extract_static_video_message_ignores_still_image_only_row(): assert _extract_static_video_messages(example) is None -def test_gym_local_video_path_is_normalized_to_file_url(tmp_path): +def test_gym_local_video_path_is_inlined_as_data_url(tmp_path): video_path = tmp_path / "clip with spaces.mp4" video_path.write_bytes(b"test") examples = [ @@ -205,14 +210,37 @@ def test_gym_local_video_path_is_normalized_to_file_url(tmp_path): } ] - normalize_video_urls_in_examples(examples) + normalize_media_in_examples(examples) - assert ( - examples[0]["responses_create_params"]["input"][0]["content"][0]["video_url"][ - "url" - ] - == video_path.resolve().as_uri() - ) + video_url = examples[0]["responses_create_params"]["input"][0]["content"][0][ + "video_url" + ] + assert video_url.startswith("data:video/mp4;base64,") + + +def test_video_path_to_data_url_rejects_unsupported_and_missing_paths(tmp_path): + """Bad local video sources must fail loudly rather than inline garbage.""" + unsupported = tmp_path / "clip.gif" + unsupported.write_bytes(b"test") + with pytest.raises(ValueError, match="Unsupported video extension"): + video_path_to_data_url(str(unsupported)) + + with pytest.raises(FileNotFoundError, match="does not exist"): + video_path_to_data_url(str(tmp_path / "missing.mp4")) + + +def test_video_path_to_data_url_passes_through_data_urls_and_accepts_file_scheme( + tmp_path, +): + already_inlined = "data:video/mp4;base64,dG95" + assert video_path_to_data_url(already_inlined) == already_inlined + + video_path = tmp_path / "clip.mp4" + video_path.write_bytes(b"toy-video") + from_plain = video_path_to_data_url(str(video_path)) + from_scheme = video_path_to_data_url(f"file://{video_path}") + assert from_plain.startswith("data:video/mp4;base64,") + assert from_plain == from_scheme def test_extract_static_video_message_rejects_multiple_videos(tmp_path): @@ -319,6 +347,52 @@ def test_video_metadata_rejects_invalid_extra_body(extra_body): _metadata_extra_body(example) +@pytest.mark.parametrize( + "chat_template_kwargs", + [ + {"enable_thinking": False}, + '{"enable_thinking": false}', + ], +) +def test_chat_template_kwargs_for_processor_accepts_mapping_or_json( + chat_template_kwargs, +): + example = { + "responses_create_params": { + "metadata": {"chat_template_kwargs": chat_template_kwargs} + } + } + + assert _chat_template_kwargs_for_processor(example) == { + "chat_template_kwargs": {"enable_thinking": False}, + "enable_thinking": False, + } + + +def test_chat_template_kwargs_for_processor_defaults_to_empty(): + assert _chat_template_kwargs_for_processor({}) == {} + + +def test_chat_template_kwargs_for_processor_rejects_invalid_json(): + example = { + "responses_create_params": {"metadata": {"chat_template_kwargs": "not-json"}} + } + + with pytest.raises(ValueError, match="chat_template_kwargs"): + _chat_template_kwargs_for_processor(example) + + +def test_deep_merge_dict_merges_nested_values_without_mutating_inputs(): + base = {"nested": {"left": 1}, "unchanged": [1]} + update = {"nested": {"right": 2}, "unchanged": [2]} + + merged = _deep_merge_dict(base, update) + + assert merged == {"nested": {"left": 1, "right": 2}, "unchanged": [2]} + assert base == {"nested": {"left": 1}, "unchanged": [1]} + assert update == {"nested": {"right": 2}, "unchanged": [2]} + + def test_video_metadata_canonicalizes_mapping_extra_body_to_json_string(): example = { "responses_create_params": { @@ -380,7 +454,7 @@ def test_video_datum_uses_temporal_processor_contract(monkeypatch, tmp_path): frames = np.zeros((4, 8, 8, 3), dtype=np.uint8) monkeypatch.setattr( - "nemo_rl.environments.nemo_gym_video.load_video_frames_with_metadata", + "nemo_rl.environments.nemo_gym_multimodal.load_video_frames_with_metadata", lambda *args, **kwargs: ( frames, {"frames_indices": [0, 3, 6, 9], "fps": 3.0}, @@ -513,7 +587,7 @@ def fake_video_processor( } monkeypatch.setattr( - "nemo_rl.environments.nemo_gym_video.nemo_gym_example_to_video_datum_spec", + "nemo_rl.environments.nemo_gym_multimodal.nemo_gym_example_to_video_datum_spec", fake_video_processor, ) processor = SimpleNamespace( @@ -574,7 +648,7 @@ def test_video_datum_uses_cached_frames_without_decoding_video(monkeypatch, tmp_ } } monkeypatch.setattr( - "nemo_rl.environments.nemo_gym_video._video_to_image_content", + "nemo_rl.environments.nemo_gym_multimodal._video_to_image_content", lambda *args, **kwargs: pytest.fail("cached frames must not decode the video"), ) @@ -661,7 +735,7 @@ def test_nemotron_video_datum_uses_dynamic_tubelet_inputs(monkeypatch, tmp_path) } frames = np.zeros((4, 8, 16, 3), dtype=np.uint8) monkeypatch.setattr( - "nemo_rl.environments.nemo_gym_video.load_video_frames_with_metadata", + "nemo_rl.environments.nemo_gym_multimodal.load_video_frames_with_metadata", lambda *args, **kwargs: ( frames, {"frames_indices": [0, 3, 6, 9], "fps": 3.0}, @@ -819,11 +893,11 @@ def fake_manifest_builder(paths): return "data:video/x-nemo-rl-cached-frames;base64,dGVzdA==" monkeypatch.setattr( - "nemo_rl.environments.nemo_gym_video.build_cached_video_frame_data_url", + "nemo_rl.environments.nemo_gym_multimodal.build_cached_video_frame_data_url", fake_manifest_builder, ) monkeypatch.setattr( - "nemo_rl.environments.nemo_gym_video.process_nemotron_video_frames", + "nemo_rl.environments.nemo_gym_multimodal.process_nemotron_video_frames", lambda *args, **kwargs: { "input_ids": torch.tensor([[7, 18, 18, 9]]), "pixel_values": torch.ones(4, 3, 8, 8), @@ -1480,7 +1554,7 @@ class _RolloutCollectionHelper: def run_examples(self, examples, head_server_config): del head_server_config content = examples[0]["responses_create_params"]["input"][0]["content"] - assert content[0]["video_url"] == video_path.resolve().as_uri() + assert content[0]["video_url"].startswith("data:video/mp4;base64,") assert content[1]["image_url"].startswith("data:image/png;base64,") async def _completed_result(): @@ -1530,6 +1604,116 @@ def _postprocess_nemo_gym_to_nemo_rl_result( asyncio.run(_run()) +@pytest.mark.parametrize("modality", ["image", "video"]) +def test_nemo_gym_megatron_multimodal_response_round_trip(tmp_path, modality): + """Round-trip normalized media and a mocked Megatron HTTP response through Gym.""" + + async def _run(): + if modality == "image": + media_path = tmp_path / "clevr.png" + Image.new("RGB", (2, 2), color="red").save(media_path) + media_part = {"type": "input_image", "image_url": str(media_path)} + expected_prefix = "data:image/png;base64," + else: + media_path = tmp_path / "vstat.mp4" + media_path.write_bytes(b"toy-video") + media_part = {"type": "input_video", "video_url": str(media_path)} + expected_prefix = "data:video/mp4;base64," + + row = { + "_rowidx": 3, + "agent_ref": {"name": "mock-megatron-agent"}, + "responses_create_params": { + "input": [ + { + "role": "user", + "content": [ + media_part, + {"type": "input_text", "text": "What is shown?"}, + ], + } + ] + }, + } + + class _Tokenizer: + def batch_decode(self, batches): + return [" ".join(map(str, token_ids)) for token_ids in batches] + + class _RolloutCollectionHelper: + def run_examples(self, examples, head_server_config): + assert head_server_config.backend == "megatron" + dispatched_row = examples[0] + dispatched_part = dispatched_row["responses_create_params"]["input"][0][ + "content" + ][0] + media_url = dispatched_part[f"{modality}_url"] + assert media_url.startswith(expected_prefix) + + mocked_result = { + "responses_create_params": { + "input": deepcopy( + dispatched_row["responses_create_params"]["input"] + ) + }, + "response": { + "agent_input": deepcopy( + dispatched_row["responses_create_params"]["input"] + ), + "output": [ + { + "type": "message", + "prompt_token_ids": [10, 99, 20], + "generation_token_ids": [71, 72], + "generation_log_probs": [-0.25, -0.5], + } + ], + }, + } + + async def _completed_result(): + return dispatched_row, mocked_result + + return [_completed_result()] + + class _MockSelf: + cfg = {} + rch = _RolloutCollectionHelper() + head_server_config = SimpleNamespace(backend="megatron") + _tokenizer = _Tokenizer() + _processor = None + # Bind the real postprocess: the assertions below are about its + # message_log output, not about run_rollouts' dispatch alone. + _postprocess_nemo_gym_to_nemo_rl_result = NemoGym.__ray_metadata__.modified_class._postprocess_nemo_gym_to_nemo_rl_result + + def _require_spinup(self): + pass + + streamed = [] + async for item in NemoGym.__ray_metadata__.modified_class.run_rollouts( + _MockSelf(), [row], "test" + ): + streamed.append(item) + + row_index, result, _metrics = streamed[0] + assert row_index == 3 + assert [message["role"] for message in result["message_log"]] == [ + "user", + "assistant", + ] + assert result["message_log"][0]["token_ids"].tolist() == [10, 99, 20] + assert result["message_log"][1]["token_ids"].tolist() == [71, 72] + assert result["message_log"][1]["generation_logprobs"].tolist() == [ + -0.25, + -0.5, + ] + assert result["full_result"]["response"]["output"][0]["generation_str"] == ( + "71 72" + ) + + asyncio.run(_run()) + + 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.""" diff --git a/tests/unit/environments/test_nemo_gym_image_placeholders.py b/tests/unit/environments/test_nemo_gym_image_placeholders.py index 73359f9d4c4..5a05ddf366d 100644 --- a/tests/unit/environments/test_nemo_gym_image_placeholders.py +++ b/tests/unit/environments/test_nemo_gym_image_placeholders.py @@ -16,10 +16,10 @@ import torch from PIL import Image -from nemo_rl.environments.nemo_gym import _attach_multimodal_data_to_user_message +from nemo_rl.data.multimodal_utils import attach_image_model_inputs_to_message # -------------------------------------------------------------------------- -# ragged pixel_values path in _attach_multimodal_data_to_user_message +# ragged pixel_values path in attach_image_model_inputs_to_message # -------------------------------------------------------------------------- @@ -72,7 +72,7 @@ def test_ragged_output_requested_only_for_multi_image_turns(): """The ragged switch needs both the flag and more than one image.""" for count, flag, expected in [(2, True, None), (1, True, "pt"), (2, False, "pt")]: processor = NemotronNanoVLV2Processor(torch.zeros(count, 3, 4, 4)) - _attach_multimodal_data_to_user_message( + attach_image_model_inputs_to_message( {}, images=_images(count), processor=processor, @@ -87,7 +87,7 @@ def test_ragged_pixel_values_are_padded_to_one_tensor(): """Heterogeneous CHW tensors become a single padded tensor for the message.""" processor = _ragged((3, 2, 4), (3, 6, 4)) user_message: dict = {} - _attach_multimodal_data_to_user_message( + attach_image_model_inputs_to_message( user_message, images=_images(2), processor=processor, @@ -103,7 +103,7 @@ def test_ragged_pixel_values_are_padded_to_one_tensor(): def test_ragged_pixel_values_reject_non_chw_entries(): processor = _ragged((3, 2, 4), (2, 4)) with pytest.raises(ValueError, match="one CHW tensor per image"): - _attach_multimodal_data_to_user_message( + attach_image_model_inputs_to_message( {}, images=_images(2), processor=processor, @@ -114,7 +114,7 @@ def test_ragged_pixel_values_reject_non_chw_entries(): def test_ragged_pixel_values_reject_mixed_channel_counts(): processor = _ragged((3, 2, 4), (1, 2, 4)) with pytest.raises(ValueError, match="same channel count"): - _attach_multimodal_data_to_user_message( + attach_image_model_inputs_to_message( {}, images=_images(2), processor=processor, @@ -124,10 +124,10 @@ def test_ragged_pixel_values_reject_mixed_channel_counts(): def test_attach_is_a_noop_without_images_or_processor(): user_message: dict = {} - _attach_multimodal_data_to_user_message( + attach_image_model_inputs_to_message( user_message, images=[], processor=NemotronNanoVLV2Processor(None) ) - _attach_multimodal_data_to_user_message( + attach_image_model_inputs_to_message( user_message, images=_images(1), processor=None ) assert user_message == {} diff --git a/tests/unit/environments/test_nemo_gym_mm_utils.py b/tests/unit/environments/test_nemo_gym_mm_utils.py index 3d8a4186e17..6f2a785b8ce 100644 --- a/tests/unit/environments/test_nemo_gym_mm_utils.py +++ b/tests/unit/environments/test_nemo_gym_mm_utils.py @@ -14,10 +14,14 @@ from PIL import Image -from nemo_rl.data.multimodal_utils import image_to_data_url -from nemo_rl.environments.nemo_gym import ( +from nemo_rl.data.multimodal_utils import ( + extract_input_media_sources_from_responses_messages, + image_to_data_url, +) +from nemo_rl.environments.nemo_gym_multimodal import ( _extract_input_images_from_message, _index_per_turn_images, + _without_initial_media_sources, ) @@ -176,3 +180,67 @@ def test_index_per_turn_images_flushes_on_function_call_trainable_item(): 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_without_initial_media_sources_strips_videos_and_images_in_order(): + """Video parts must be de-duplicated alongside images, in encounter order.""" + image_url = _image((2, 2)) + video_url = "data:video/mp4;base64,dG95" + messages = [ + { + "role": "user", + "content": [ + {"type": "input_video", "video_url": video_url}, + {"type": "input_image", "image_url": image_url}, + {"type": "input_text", "text": "What is shown?"}, + ], + } + ] + initial_sources = extract_input_media_sources_from_responses_messages(messages) + assert initial_sources == [("video", video_url), ("image", image_url)] + + filtered, fully_consumed = _without_initial_media_sources(messages, initial_sources) + + assert fully_consumed is True + assert filtered[0]["content"] == [{"type": "input_text", "text": "What is shown?"}] + # The caller's messages must not be mutated in place. + assert len(messages[0]["content"]) == 3 + + +def test_without_initial_media_sources_keeps_media_the_agent_added(): + """Only the ordered prefix of initial sources is removed; extras survive.""" + initial_image = _image((2, 2)) + agent_image = _image((4, 4)) + messages = [ + { + "role": "user", + "content": [{"type": "input_image", "image_url": initial_image}], + }, + { + "role": "user", + "content": [{"type": "input_image", "image_url": agent_image}], + }, + ] + + filtered, fully_consumed = _without_initial_media_sources( + messages, [("image", initial_image)] + ) + + assert fully_consumed is True + assert filtered[0]["content"] == [] + assert filtered[1]["content"] == [{"type": "input_image", "image_url": agent_image}] + + +def test_without_initial_media_sources_reports_unconsumed_sources(): + """A source that never appears leaves the consumed flag False.""" + filtered, fully_consumed = _without_initial_media_sources( + [{"role": "user", "content": [{"type": "input_text", "text": "hi"}]}], + [("image", "data:image/png;base64,AA")], + ) + + assert fully_consumed is False + assert filtered[0]["content"] == [{"type": "input_text", "text": "hi"}] + + +def test_without_initial_media_sources_passes_through_non_list_messages(): + assert _without_initial_media_sources("not-a-list", []) == ("not-a-list", False) diff --git a/tests/unit/models/generation/test_megatron_generation.py b/tests/unit/models/generation/test_megatron_generation.py index 95bce59081d..f996690deef 100644 --- a/tests/unit/models/generation/test_megatron_generation.py +++ b/tests/unit/models/generation/test_megatron_generation.py @@ -22,12 +22,17 @@ from nemo_rl.algorithms.grpo import refit_policy_generation from nemo_rl.algorithms.utils import get_tokenizer +from nemo_rl.data.multimodal_utils import PackedTensor from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.distributed.virtual_cluster import RayVirtualCluster from nemo_rl.models.generation.megatron import MegatronGeneration, megatron_generation from nemo_rl.models.generation.megatron.config import ( dedicated_inference_megatron_cfg, ) +from nemo_rl.models.generation.megatron.megatron_worker import MegatronGenerationMixin +from nemo_rl.models.generation.megatron.utils import ( + build_prompt_and_multimodal_data, +) from nemo_rl.models.policy import PolicyConfig from nemo_rl.models.policy.lm_policy import Policy from nemo_rl.weight_sync.megatron_weight_synchronizer import ( @@ -37,6 +42,209 @@ model_name = "Qwen/Qwen3-0.6B" + +@pytest.mark.mcore +def test_multimodal_preprocessing_requires_policy_processor(): + class _ImageWrapper: + supports_image = True + + worker = object.__new__(MegatronGenerationMixin) + worker._get_megatron_inference_wrapper_cls = lambda: _ImageWrapper + + with pytest.raises(ValueError, match="requires the policy processor"): + worker._build_image_preprocessing_config({}) + + +@pytest.mark.mcore +def test_multimodal_preprocessing_forwards_vision_model_type(): + class _ImageWrapper: + supports_image = True + + worker = object.__new__(MegatronGenerationMixin) + worker._get_megatron_inference_wrapper_cls = lambda: _ImageWrapper + worker.processor = SimpleNamespace( + image_processor=SimpleNamespace( + patch_size=14, + min_num_patches=1, + max_num_patches=32, + norm_mean=[0.1, 0.2, 0.3], + norm_std=[0.4, 0.5, 0.6], + ) + ) + + config = worker._build_image_preprocessing_config({"vision_model_type": "qwen-vl"}) + + assert config.vision_model_type == "qwen-vl" + + +@pytest.mark.mcore +def test_direct_megatron_media_request_preserves_preexpanded_prompt(): + def fake_sample_vision_tensors(data, index): + return torch.ones(1, 2, 4), torch.tensor([[2, 2]]), None + + data = { + "input_ids": torch.tensor([[10, 99, 99, 20, 0]]), + "input_lengths": torch.tensor([4]), + } + + prompt, multi_modal_data = build_prompt_and_multimodal_data( + data, + 0, + sample_tensors=fake_sample_vision_tensors, + supports_modality=lambda modality: modality == "image", + ) + + assert prompt == [10, 99, 99, 20] + assert multi_modal_data["media_tokens_preexpanded"] is True + assert "image" in multi_modal_data + + +@pytest.mark.mcore +def test_text_only_request_does_not_resolve_multimodal_capabilities(): + data = { + "input_ids": torch.tensor([[10, 20, 0]]), + "input_lengths": torch.tensor([2]), + } + + prompt, multi_modal_data = build_prompt_and_multimodal_data( + data, + 0, + supports_modality=lambda modality: pytest.fail( + f"unexpected capability lookup for {modality}" + ), + ) + + assert prompt == [10, 20] + assert multi_modal_data is None + + +@pytest.mark.mcore +def test_direct_megatron_video_request_marks_preexpanded_prompt(): + def fake_sample_vision_tensors(data, index): + return ( + torch.ones(1, 4, 4), + torch.tensor([[2, 2], [2, 2], [2, 2], [2, 2]]), + torch.tensor([4]), + ) + + data = { + "input_ids": torch.tensor([[10, 99, 99, 20]]), + "input_lengths": torch.tensor([4]), + } + + prompt, multi_modal_data = build_prompt_and_multimodal_data( + data, + 0, + sample_tensors=fake_sample_vision_tensors, + supports_modality=lambda modality: modality == "video", + ) + + assert prompt == [10, 99, 99, 20] + assert multi_modal_data["media_tokens_preexpanded"] is True + assert "video" in multi_modal_data + + +@pytest.mark.mcore +@pytest.mark.parametrize( + ("modality", "num_frames"), + [("image", torch.tensor([1])), ("video", torch.tensor([4]))], +) +def test_direct_megatron_multimodal_generate_round_trip( + monkeypatch, modality, num_frames +): + """Exercise RL request construction and response packing around a mocked MCore LLM.""" + + class _MultimodalWrapper: + supports_text = True + supports_image = True + supports_video = True + supports_audio = False + + worker = object.__new__(MegatronGenerationMixin) + worker.cfg = { + "generation": { + "temperature": 1.0, + "top_k": None, + "top_p": 1.0, + "max_new_tokens": 2, + "stop_strings": None, + "mcore_generation_config": {}, + } + } + worker.tokenizer = SimpleNamespace(pad_token_id=0) + worker.megatron_tokenizer = SimpleNamespace(eod=2) + worker._inference_loop = object() + worker._get_megatron_inference_wrapper_cls = lambda: _MultimodalWrapper + + frame_count = int(num_frames.sum()) + pixels = torch.arange(frame_count * 12, dtype=torch.float32).reshape( + frame_count, 3, 2, 2 + ) + sizes = torch.tensor([[2, 2]] * frame_count) + data = BatchedDataDict( + { + "input_ids": torch.tensor([[10, 99, 99, 20]]), + "input_lengths": torch.tensor([4]), + "pixel_values": PackedTensor([pixels], dim_to_pack=0), + "imgs_sizes": PackedTensor([sizes], dim_to_pack=0), + "num_frames": PackedTensor([num_frames], dim_to_pack=0), + } + ) + + captured = {} + mocked_call = object() + + def mock_generate(prompts, multi_modal_data, sampling_params): + captured.update( + prompts=prompts, + multi_modal_data=multi_modal_data, + sampling_params=sampling_params, + ) + return mocked_call + + replies = [ + SimpleNamespace( + prompt_tokens=torch.tensor([10, 99, 99, 20]), + generated_tokens=[71, 72], + generated_log_probs=[-0.25, -0.5], + ) + ] + worker._generate_with_persistent_engine = mock_generate + + class _CompletedFuture: + def result(self): + return replies + + def mock_run_coroutine_threadsafe(call, loop): + assert call is mocked_call + assert loop is worker._inference_loop + return _CompletedFuture() + + monkeypatch.setattr( + "nemo_rl.models.generation.megatron.megatron_worker.asyncio.run_coroutine_threadsafe", + mock_run_coroutine_threadsafe, + ) + + output = worker.generate(data=data) + + assert captured["prompts"] == [[10, 99, 99, 20]] + media = captured["multi_modal_data"][0] + assert media["media_tokens_preexpanded"] is True + assert set(media) == {modality, "media_tokens_preexpanded"} + assert torch.equal(media[modality]["imgs"], pixels) + assert torch.equal(media[modality]["imgs_sizes"], sizes) + if modality == "video": + assert torch.equal(media["video"]["num_frames"], num_frames.to(torch.int32)) + else: + assert "num_frames" not in media["image"] + assert captured["sampling_params"][0].return_prompt_tokens is True + + assert output["output_ids"][0].tolist() == [10, 99, 99, 20, 71, 72] + assert output["logprobs"][0].tolist() == [0.0, 0.0, 0.0, 0.0, -0.25, -0.5] + assert output["generation_lengths"].tolist() == [2] + assert output["unpadded_sequence_lengths"].tolist() == [6] + + basic_megatron_test_config: PolicyConfig = { "model_name": model_name, "tokenizer": {"name": model_name}, @@ -619,6 +827,12 @@ def test_megatron_generation_non_colocated_refit( cluster=generation_cluster, skip_weight_load=skip_weight_load, ) + assert mg._policy_config is not config + assert mg._policy_config["generation"] is not config["generation"] + assert ( + mg._policy_config["generation"]["mcore_generation_config"] + is not config["generation"]["mcore_generation_config"] + ) # Wire the refit collective the way grpo.setup does: through the # weight synchronizer, which refit_policy_generation dispatches to. diff --git a/tests/unit/models/generation/test_megatron_generation_parse.py b/tests/unit/models/generation/test_megatron_generation_parse.py index 6db700ef237..062f9b2edb6 100644 --- a/tests/unit/models/generation/test_megatron_generation_parse.py +++ b/tests/unit/models/generation/test_megatron_generation_parse.py @@ -197,13 +197,27 @@ def test_http_server_port_reservation(monkeypatch): pass # Worker-side adoption: the same live socket, duplicated across the - # process boundary; still the same port, still accepting. + # process boundary. The holder confirms that its original descriptor is + # closed before this returns, preventing MCore's SO_REUSEPORT listeners + # from racing the old non-reusable socket. reserved = receive_held_socket(port) try: + assert holder._sock.fileno() == -1 assert reserved.getsockname()[1] == port with socket.create_connection(("127.0.0.1", port), timeout=5): pass + # MCore closes the handed-off fd and gives every frontend replica its + # own SO_REUSEPORT listener. Verify that such a listener can join the + # reservation's reuse group even before this duplicate is closed. + replica = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + try: + replica.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + replica.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1) + replica.bind(("0.0.0.0", port)) + finally: + replica.close() + # Server start with the network and MLM server stubbed out. started = {} monkeypatch.setattr( @@ -231,6 +245,7 @@ def test_http_server_port_reservation(monkeypatch): rank=0, cfg={"generation": {"mcore_generation_config": {"parsers": []}}}, _reserved_http_server_socket=reserved_socket, + inference_wrapped_model=SimpleNamespace(multimodal_prompt_config=None), ) base_url = MegatronGenerationMixin._setup_openai_api_server(worker) assert started["sock"] is reserved_socket diff --git a/tests/unit/models/generation/test_megatron_generation_utils.py b/tests/unit/models/generation/test_megatron_generation_utils.py new file mode 100644 index 00000000000..f53c5ae0d0d --- /dev/null +++ b/tests/unit/models/generation/test_megatron_generation_utils.py @@ -0,0 +1,176 @@ +# 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 types import SimpleNamespace + +import pytest + +from nemo_rl.models.generation.megatron.utils import ( + build_image_preprocessing_config, + build_video_preprocessing_config, +) + +pytestmark = pytest.mark.mcore + + +def _image_processor(**overrides): + fields = { + "patch_size": 14, + "min_num_patches": 1, + "max_num_patches": 32, + "norm_mean": [0.1, 0.2, 0.3], + "norm_std": [0.4, 0.5, 0.6], + } + fields.update(overrides) + return SimpleNamespace(**fields) + + +def test_build_image_config_handles_dict_patch_and_downsample_ratio(): + config = build_image_preprocessing_config( + _image_processor( + patch_size={"height": 16, "width": 16}, + downsample_ratio=0.5, + ), + dynamic_resolution=True, + vision_model_type="qwen-vl", + ) + + assert config.patch_dim == 16 + assert config.dynamic_resolution is True + assert config.vision_model_type == "qwen-vl" + assert config.pixel_shuffle is True + assert config.spatial_merge_size == 2 + assert config.dynamic_resolution_min_patches == 1 + assert config.dynamic_resolution_max_patches == 32 + assert config.pixel_mean == [0.1, 0.2, 0.3] + assert config.pixel_std == [0.4, 0.5, 0.6] + + +@pytest.mark.parametrize( + ("merge_fields", "expected_merge_size"), + [ + ({"merge_size": 4}, 4), + ({"spatial_merge_size": 3}, 3), + ({}, 1), + ], +) +def test_build_image_config_merge_size_fallbacks(merge_fields, expected_merge_size): + config = build_image_preprocessing_config(_image_processor(**merge_fields)) + + assert config.spatial_merge_size == expected_merge_size + assert config.pixel_shuffle is (expected_merge_size > 1) + + +def test_build_image_config_accepts_alternate_field_names(): + config = build_image_preprocessing_config( + SimpleNamespace( + patch_dim=12, + min_num_patches=2, + max_num_patches=24, + image_mean=(0.1, 0.2, 0.3), + image_std=(0.7, 0.8, 0.9), + ) + ) + + assert config.patch_dim == 12 + assert config.dynamic_resolution_min_patches == 2 + assert config.dynamic_resolution_max_patches == 24 + assert config.pixel_mean == [0.1, 0.2, 0.3] + assert config.pixel_std == [0.7, 0.8, 0.9] + + +def test_build_image_config_error_names_all_missing_fields(): + with pytest.raises(ValueError) as exc_info: + build_image_preprocessing_config(SimpleNamespace()) + + for field in ( + "patch_size", + "min_num_patches", + "max_num_patches", + "norm_mean", + "norm_std", + ): + assert field in str(exc_info.value) + + +def test_build_video_config_returns_none_when_disabled(): + image_config = build_image_preprocessing_config(_image_processor()) + + assert ( + build_video_preprocessing_config( + None, + {"video_temporal_patch_size": 2, "video_num_frames": 8}, + frame_manifest_magic=b"manifest", + ) + is None + ) + assert ( + build_video_preprocessing_config( + image_config, + {}, + frame_manifest_magic=b"manifest", + ) + is None + ) + + +def test_build_video_config_is_not_enabled_by_temporal_patch_size_alone(): + image_config = build_image_preprocessing_config(_image_processor()) + + assert ( + build_video_preprocessing_config( + image_config, + {"video_temporal_patch_size": 2}, + frame_manifest_magic=b"manifest", + ) + is None + ) + + +def test_build_video_config_uses_default_temporal_patch_size(): + image_config = build_image_preprocessing_config(_image_processor()) + + video_config = build_video_preprocessing_config( + image_config, + {"video_num_frames": 8}, + frame_manifest_magic=b"manifest", + ) + + assert video_config is not None + assert video_config.num_frames == 8 + assert video_config.temporal_patch_size == 1 + + +def test_build_video_config_overrides_patch_budget_without_mutating_image_config(): + image_config = build_image_preprocessing_config(_image_processor()) + + video_config = build_video_preprocessing_config( + image_config, + { + "video_num_frames": 8, + "video_temporal_patch_size": 2, + "video_target_num_patches": 64, + "video_maintain_aspect_ratio": False, + }, + frame_manifest_magic=b"manifest", + ) + + assert video_config is not None + assert video_config.image_config is not image_config + assert image_config.dynamic_resolution_max_patches == 32 + assert video_config.image_config.dynamic_resolution_max_patches == 64 + assert video_config.num_frames == 8 + assert video_config.temporal_patch_size == 2 + assert video_config.frame_manifest_magic == b"manifest" + assert video_config.video_maintain_aspect_ratio is False diff --git a/tests/unit/models/generation/test_vllm_video_utils.py b/tests/unit/models/generation/test_vllm_video_utils.py index 468778f9677..dc691b720e1 100644 --- a/tests/unit/models/generation/test_vllm_video_utils.py +++ b/tests/unit/models/generation/test_vllm_video_utils.py @@ -158,7 +158,7 @@ def __init__(self, *_args, **_kwargs): Image.fromarray(frame).save(frame_path) expected_frames.append(frame) frame_paths.append(str(frame_path)) - payload = utils._CACHED_VIDEO_FRAME_MANIFEST_MAGIC + json.dumps( + payload = utils.CACHED_VIDEO_FRAME_MANIFEST_MAGIC + json.dumps( { "frame_paths": frame_paths, "metadata": { @@ -189,7 +189,7 @@ def test_cached_video_data_url_requires_no_driver_decoder(monkeypatch, tmp_path) _, encoded = data_url.split(",", 1) payload = base64.b64decode(encoded) - manifest = json.loads(payload[len(utils._CACHED_VIDEO_FRAME_MANIFEST_MAGIC) :]) + manifest = json.loads(payload[len(utils.CACHED_VIDEO_FRAME_MANIFEST_MAGIC) :]) assert manifest["frame_paths"] == frame_paths assert manifest["metadata"]["frames_indices"] == [0, 1, 2, 3] assert manifest["metadata"]["fps"] == 1.0 @@ -201,7 +201,7 @@ def test_cached_video_manifest_does_not_import_torchcodec(monkeypatch, tmp_path) monkeypatch.setenv("NEMO_RL_VIDEO_MEDIA_ROOT", str(tmp_path)) frame_path = tmp_path / "frame.png" Image.new("RGB", (2, 2)).save(frame_path) - payload = utils._CACHED_VIDEO_FRAME_MANIFEST_MAGIC + json.dumps( + payload = utils.CACHED_VIDEO_FRAME_MANIFEST_MAGIC + json.dumps( { "frame_paths": [str(frame_path)], "metadata": { diff --git a/tests/unit/models/megatron/test_megatron_setup.py b/tests/unit/models/megatron/test_megatron_setup.py index 2cbd02361eb..22994ccd10d 100644 --- a/tests/unit/models/megatron/test_megatron_setup.py +++ b/tests/unit/models/megatron/test_megatron_setup.py @@ -57,6 +57,23 @@ def finalize(self) -> None: self.finalized = True +@pytest.mark.mcore +def test_resolve_optimizer_fp8_moment_dtypes(): + from nemo_rl.models.megatron.setup import _resolve_optimizer_dtype_kwargs + + resolved = _resolve_optimizer_dtype_kwargs( + { + "main_params_dtype": "float16", + "exp_avg_dtype": "fp8", + "exp_avg_sq_dtype": "torch.uint8", + } + ) + + assert resolved["main_params_dtype"] is torch.float16 + assert resolved["exp_avg_dtype"] is torch.uint8 + assert resolved["exp_avg_sq_dtype"] is torch.uint8 + + @pytest.mark.mcore class TestValidateModelPaths: """Tests for validate_model_paths function.""" diff --git a/tests/unit/test_check_metrics.py b/tests/unit/test_check_metrics.py index 73cb08f4690..6f417cb15eb 100644 --- a/tests/unit/test_check_metrics.py +++ b/tests/unit/test_check_metrics.py @@ -21,7 +21,15 @@ tests_dir = Path(__file__).parent.parent sys.path.insert(0, str(tests_dir)) -from check_metrics import evaluate_check, max, mean, median, min, ratio_above +from check_metrics import ( + all_finite, + evaluate_check, + max, + mean, + median, + min, + ratio_above, +) class TestMeanFunction: @@ -160,6 +168,18 @@ def test_max_with_string_values(self): assert result == 8.8 +class TestAllFiniteFunction: + def test_requires_at_least_one_value(self): + assert all_finite({}) is False + + @pytest.mark.parametrize("invalid", [float("nan"), float("inf"), float("-inf")]) + def test_rejects_non_finite_values(self, invalid): + assert all_finite({"1": 1.0, "2": invalid}) is False + + def test_accepts_finite_numeric_strings(self): + assert all_finite({"1": "1.0", "2": "-2.5"}) is True + + class TestRatioAboveFunction: """Test the ratio_above function.""" @@ -306,6 +326,14 @@ def test_evaluate_check_with_ratio_above(self): assert passed is False assert value == 0.4 + def test_evaluate_check_with_all_finite(self): + passed, _, value = evaluate_check( + {"loss": {"1": 1.0, "2": 0.5}}, + 'all_finite(data["loss"])', + ) + assert passed is True + assert value is True + class TestRealWorldScenarios: """Test scenarios that match real-world usage patterns.""" diff --git a/tests/unit/test_prepare_video_dataset.py b/tests/unit/test_prepare_video_dataset.py index 9b58c856e5e..cc68186841e 100644 --- a/tests/unit/test_prepare_video_dataset.py +++ b/tests/unit/test_prepare_video_dataset.py @@ -67,7 +67,7 @@ def test_converter_skips_missing_local_videos_when_requested( content = converted_rows[0]["responses_create_params"]["input"][0]["content"] assert content[0]["video_url"] == str(existing_video.resolve()) metadata = converted_rows[0]["responses_create_params"]["metadata"] - assert metadata["chat_template_kwargs"] == {"enable_thinking": True} + assert json.loads(metadata["chat_template_kwargs"]) == {"enable_thinking": True} assert "extra_body" not in metadata assert "extraction_mode" not in converted_rows[0] assert "Skipped 1 non-video or duplicate rows" in capsys.readouterr().out diff --git a/tests/unit/utils/test_venvs.py b/tests/unit/utils/test_venvs.py index 9e54541bb06..2635689385b 100644 --- a/tests/unit/utils/test_venvs.py +++ b/tests/unit/utils/test_venvs.py @@ -16,7 +16,11 @@ from tempfile import TemporaryDirectory from unittest.mock import patch -from nemo_rl.utils.venvs import create_local_venv +from nemo_rl.utils.venvs import ( + add_hf_modules_cache_to_pythonpath, + create_local_venv, + make_actor_runtime_env, +) from tests.unit.conftest import TEST_ASSETS_DIR @@ -48,3 +52,51 @@ def test_create_local_venv(): # Verify the command executed successfully (return code 0) assert result.returncode == 0, f"Failed to import sphinx: {result.stderr}" assert "Sphinx package is installed" in result.stdout + + +def test_add_hf_modules_cache_to_pythonpath(): + result = add_hf_modules_cache_to_pythonpath( + { + "HF_MODULES_CACHE": "/hf/modules", + "PYTHONPATH": f"/project{os.pathsep}/other", + } + ) + + assert result["PYTHONPATH"].split(os.pathsep) == [ + "/hf/modules", + "/project", + "/other", + ] + + +def test_add_hf_modules_cache_does_not_duplicate_pythonpath_entry(): + pythonpath = f"/project{os.pathsep}/hf/modules" + + result = add_hf_modules_cache_to_pythonpath( + {"HF_MODULES_CACHE": "/hf/modules", "PYTHONPATH": pythonpath} + ) + + assert result["PYTHONPATH"] == pythonpath + + +def test_make_actor_runtime_env_builds_local_venv_for_uv_python_executable(): + """Mirrors the inline venv-creation logic that used to live in grpo.py.""" + with ( + patch( + "nemo_rl.distributed.ray_actor_environment_registry.get_actor_python_env", + return_value="uv run --group vllm", + ) as mock_get_env, + patch( + "nemo_rl.utils.venvs.create_local_venv_on_each_node", + return_value="/fake/venv/bin/python", + ) as mock_create_venv, + ): + runtime_env = make_actor_runtime_env("some.module.SomeActor") + + mock_get_env.assert_called_once_with("some.module.SomeActor") + mock_create_venv.assert_called_once_with( + "uv run --group vllm", "some.module.SomeActor" + ) + assert runtime_env["py_executable"] == "/fake/venv/bin/python" + assert runtime_env["env_vars"]["VIRTUAL_ENV"] == "/fake/venv" + assert runtime_env["env_vars"]["UV_PROJECT_ENVIRONMENT"] == "/fake/venv" diff --git a/tools/install_audio_deps.sh b/tools/install_audio_deps.sh index 312cf5f3325..9c89ffd43d2 100755 --- a/tools/install_audio_deps.sh +++ b/tools/install_audio_deps.sh @@ -5,27 +5,46 @@ # # bash tools/install_audio_deps.sh # -# Safe to call multiple times — exits immediately if already installed. +# Safe to call multiple times. set -euo pipefail -# Fast exit: if torchcodec imports cleanly it already has FFmpeg available. -if python -c "import torchcodec" 2>/dev/null; then - echo "[audio-deps] Already installed and functional, skipping." - exit 0 -fi +if ! python -c "import torchcodec" 2>/dev/null; then + # Install system FFmpeg — torchcodec dlopens libavcodec.so.* at runtime. + echo "[audio-deps] Installing system FFmpeg..." + apt-get update && apt-get install -y --no-install-recommends ffmpeg -# Install system FFmpeg — torchcodec dlopens libavcodec.so.* at runtime. -echo "[audio-deps] Installing system FFmpeg..." -apt-get update && apt-get install -y --no-install-recommends ffmpeg + # torchaudio 2.11+ routes torchaudio.load through torchcodec, so both are needed. + # --no-config prevents the project's [tool.uv] overrides from interfering. + echo "[audio-deps] Installing torchaudio==2.11.0 and torchcodec..." + uv pip install --no-config \ + --index-url https://download.pytorch.org/whl/cu130 \ + --extra-index-url https://pypi.org/simple \ + --reinstall-package torchaudio \ + "torchaudio==2.11.0" \ + "torchcodec==0.11.1" +fi -# torchaudio 2.11+ routes torchaudio.load through torchcodec, so both are needed. -# --no-config prevents the project's [tool.uv] overrides from interfering. -echo "[audio-deps] Installing torchaudio==2.11.0 and torchcodec..." -uv pip install --no-config \ - --index-url https://download.pytorch.org/whl/cu130 \ - --extra-index-url https://pypi.org/simple \ - --reinstall-package torchaudio \ - "torchaudio==2.11.0" \ - "torchcodec>=0.3.0" +# PyAV is intentionally absent from the base image (pyproject excludes it via +# `av; sys_platform == 'never'` because it bundles CVE-carrying codec libs), so it +# must be installed after the fact into the isolated Megatron policy worker +# environment that imports it. `--no-config` bypasses that exclusion; the version +# floor is therefore restated here to keep pyproject's CVE-2026-40962 constraint. +# +# The worker venv is created lazily at worker start, so run this AFTER the +# Megatron worker has been created at least once (or point RAY_MEGATRON_PYTHON at +# an existing venv). +RAY_MEGATRON_PYTHON="${RAY_MEGATRON_PYTHON:-/opt/ray_venvs/nemo_rl.models.policy.workers.megatron_policy_worker.MegatronPolicyWorker/bin/python}" +if [[ ! -x "$RAY_MEGATRON_PYTHON" ]]; then + echo "[audio-deps] ERROR: Megatron worker environment not found: $RAY_MEGATRON_PYTHON" >&2 + echo "[audio-deps] It is created on first worker start. Run this script after that," >&2 + echo "[audio-deps] or set RAY_MEGATRON_PYTHON to an existing worker interpreter." >&2 + exit 1 +fi +if ! "$RAY_MEGATRON_PYTHON" -c "import av" 2>/dev/null; then + # `uv pip install --python` targets the venv directly; these venvs are built + # by `uv venv` without `--seed`, so they have no pip to invoke. + echo "[audio-deps] Installing PyAV in the Megatron worker environment..." + uv pip install --no-config --python "$RAY_MEGATRON_PYTHON" "av>=17.1.0" +fi echo "[audio-deps] Done."