From 4c01498dd4ba27204da187a29397cc2e266420aa Mon Sep 17 00:00:00 2001 From: Cory Ye Date: Tue, 18 Aug 2026 10:35:58 -0700 Subject: [PATCH 01/15] Support Megatron-LLM inference backend for multimodal RL. Signed-off-by: Cory Ye --- ...cle-count-1n4g-megatron_generation.v1.yaml | 150 +++++ ...a3b-clevr-1n4g-megatron_generation.v1.yaml | 104 ++++ ...a3b-clevr-8n4g-megatron_generation.v1.yaml | 114 ++++ nemo_rl/data/collate_fn.py | 14 +- nemo_rl/data/multimodal_utils.py | 150 +++-- nemo_rl/data/processors.py | 6 +- nemo_rl/environments/nemo_gym.py | 80 +-- nemo_rl/evals/eval.py | 30 +- nemo_rl/models/generation/megatron/config.py | 11 + .../megatron/megatron_generation.py | 8 +- .../generation/megatron/megatron_worker.py | 304 ++++++++-- nemo_rl/models/generation/megatron/utils.py | 60 ++ nemo_rl/models/megatron/setup.py | 22 +- nemo_rl/models/policy/__init__.py | 4 + .../policy/workers/megatron_policy_worker.py | 2 + scripts/add_nemo_gym_agent_ref.py | 74 +++ scripts/prepare_nemotron_omni_vstat.py | 144 +++++ ...emotron_omni_circle_count_megatron_1n4g.sh | 387 +++++++++++++ .../run_nemotron_omni_clevr_megatron_1n4g.sh | 324 +++++++++++ .../run_nemotron_omni_vstat_megatron_1n4g.sh | 245 +++++++++ ...ubmit_nemotron_omni_clevr_megatron_8n4g.sh | 422 ++++++++++++++ ...ubmit_nemotron_omni_vstat_megatron_8n4g.sh | 519 ++++++++++++++++++ tests/unit/data/datasets/test_mmpr_tiny.py | 17 +- tests/unit/data/test_collate_fn.py | 22 +- tools/install_audio_deps.sh | 42 +- 25 files changed, 3098 insertions(+), 157 deletions(-) create mode 100644 examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-count-1n4g-megatron_generation.v1.yaml create mode 100644 examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.yaml create mode 100644 examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron_generation.v1.yaml create mode 100644 scripts/add_nemo_gym_agent_ref.py create mode 100644 scripts/prepare_nemotron_omni_vstat.py create mode 100755 scripts/run_nemotron_omni_circle_count_megatron_1n4g.sh create mode 100755 scripts/run_nemotron_omni_clevr_megatron_1n4g.sh create mode 100755 scripts/run_nemotron_omni_vstat_megatron_1n4g.sh create mode 100755 scripts/submit_nemotron_omni_clevr_megatron_8n4g.sh create mode 100755 scripts/submit_nemotron_omni_vstat_megatron_8n4g.sh 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..9b1aaaa5508 --- /dev/null +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-count-1n4g-megatron_generation.v1.yaml @@ -0,0 +1,150 @@ +defaults: ../../vlm_grpo_3B_megatron.yaml + +# 1-node × 4-GPU GB200 recipe: Nemotron Omni on NeMo-Gym's circle_count image +# environment. Each row is a synthetic PNG scene carried inline as a base64 +# data URL plus a "how many circles" question; the Gym resources server +# rewards an exact match on \boxed{count}. Default layout is non-colocated +# 2 train / 2 gen with Megatron generation, async GRPO, age>1, and in-flight +# weight updates. +# +# `data` and `env` use _override_ because the parent config wires the CLEVR +# HF dataset and its reward functions at the split level, and split keys win +# over `data.default` (nemo_rl/data/datasets/utils.py:263). +grpo: + num_prompts_per_step: 2 + num_generations_per_prompt: 8 + max_num_steps: 4 + val_period: 0 + val_at_start: false + val_at_end: false + use_dynamic_sampling: false + # The Gym entrypoint refuses a preset max_val_samples and derives both of + # these from the validation dataset length instead. + max_val_samples: null + val_batch_size: null + reward_shaping: + enabled: false + reward_scaling: + enabled: false + 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 + train_micro_batch_size: 1 + logprob_batch_size: 1 + max_total_sequence_length: 8192 + offload_optimizer_for_logprob: false + sequence_packing: + enabled: true + megatron_cfg: + env_vars: + TORCH_CUDA_ARCH_LIST: "10.0" + tensor_model_parallel_size: 2 + expert_model_parallel_size: 2 + expert_tensor_parallel_size: 1 + context_parallel_size: 1 + sequence_parallel: true + bias_activation_fusion: false + activation_checkpointing: true + generation: + backend: megatron + max_new_tokens: 1024 + # Gym clears stop_strings/stop_token_ids in setup_nemo_gym_config, and the + # Megatron engine takes null here (the vLLM path bans the media tokens + # explicitly instead). + bad_words: null + mcore_generation_config: + # Required for Gym: the environment drives generation over the engine's + # OpenAI-compatible HTTP server. + expose_http_server: true + parsers: [] + max_model_len: ${policy.max_total_sequence_length} + cuda_graph_impl: local + inference_cuda_graph_scope: block + buffer_size_gb: 8 + num_cuda_graphs: -1 + block_size_tokens: 256 + use_cuda_graphs_for_non_decode_steps: true + enable_chunked_prefill: true + enable_prefix_caching: false + max_tokens: ${policy.max_total_sequence_length} + kv_cache_management_mode: persist + materialize_only_last_token_logits: true + num_speculative_tokens: 0 + refit_backend: nccl + 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 + vllm_cfg: + async_engine: true + expose_http_server: true + skip_tokenizer_init: false + colocated: + enabled: false + resources: + gpus_per_node: 2 + num_nodes: 1 + +data: + _override_: true + max_input_seq_length: null + shuffle: false + num_workers: 1 + use_multiple_dataloader: false + # The launcher overrides these with absolute paths; the shipped example file + # holds 5 rows, so regenerate a larger set with + # resources_servers/circle_count/generate_data.py for anything but a smoke. + 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 + system_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: + wandb_enabled: false + tensorboard_enabled: false + wandb: + project: "nemo-rl-omni" + name: "nemotron-omni-circle-count-1n4g" + +cluster: + gpus_per_node: 4 + num_nodes: 1 diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.yaml new file mode 100644 index 00000000000..34bbd2a92c7 --- /dev/null +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.yaml @@ -0,0 +1,104 @@ +defaults: ../../vlm_grpo_3B_megatron.yaml + +# 1-node × 4-GPU GB200 recipe: Nemotron Omni CLEVR with Megatron generation. +# Default layout is non-colocated 2 train / 2 gen with async GRPO, age>1, and +# in-flight weight updates. +grpo: + num_prompts_per_step: 2 + num_generations_per_prompt: 8 + max_num_steps: 4 + val_period: 0 + val_at_start: false + val_at_end: false + use_dynamic_sampling: false + reward_shaping: + enabled: false + reward_scaling: + enabled: false + 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-1n4g + +policy: + model_name: nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16 + is_vlm: true + train_global_batch_size: 16 + train_micro_batch_size: 1 + logprob_batch_size: 1 + max_total_sequence_length: 2048 + offload_optimizer_for_logprob: false + sequence_packing: + enabled: true + megatron_cfg: + env_vars: + TORCH_CUDA_ARCH_LIST: "10.0" + tensor_model_parallel_size: 2 + expert_model_parallel_size: 2 + expert_tensor_parallel_size: 1 + context_parallel_size: 1 + sequence_parallel: true + bias_activation_fusion: false + activation_checkpointing: true + generation: + backend: megatron + max_new_tokens: 512 + stop_strings: + - + - + - + - + - + - + mcore_generation_config: + expose_http_server: false + parsers: [] + max_model_len: ${policy.max_total_sequence_length} + cuda_graph_impl: local + inference_cuda_graph_scope: block + buffer_size_gb: 8 + num_cuda_graphs: -1 + block_size_tokens: 256 + use_cuda_graphs_for_non_decode_steps: true + enable_chunked_prefill: true + enable_prefix_caching: false + max_tokens: ${policy.max_total_sequence_length} + kv_cache_management_mode: persist + materialize_only_last_token_logits: true + num_speculative_tokens: 0 + refit_backend: gloo + 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 + colocated: + enabled: false + resources: + gpus_per_node: 2 + num_nodes: 1 + +data: + default: + prompt_file: examples/prompts/clevr_cogent_cot_nemotron_omni.txt + +logger: + wandb_enabled: false + tensorboard_enabled: false + wandb: + project: "nemo-rl-omni" + name: "nemotron-omni-clevr-megatron-1n4g" + +cluster: + gpus_per_node: 4 + num_nodes: 1 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..73ffb6b771f --- /dev/null +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron_generation.v1.yaml @@ -0,0 +1,114 @@ +defaults: ../../vlm_grpo_3B_megatron.yaml + +# 8-node × 4-GPU GB200 recipe: Nemotron Omni CLEVR with Megatron generation. +# Default layout is non-colocated 2 train / 6 gen nodes with async GRPO, age>1, +# and in-flight weight updates (gen-heavy split for long rollouts). +grpo: + # Inference DP=3 (24 GPUs / TP8); two prompt groups per DP replica → 6. + num_prompts_per_step: 6 + num_generations_per_prompt: 8 + max_num_steps: 1000000 + val_period: 10 + val_at_start: false + val_at_end: false + # Training DP=1 (8 GPUs / TP8). + val_batch_size: 64 + max_val_samples: 64 + use_dynamic_sampling: false + reward_shaping: + enabled: false + reward_scaling: + enabled: false + 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}} + train_micro_batch_size: 1 + logprob_batch_size: 1 + max_total_sequence_length: 4096 + offload_optimizer_for_logprob: false + sequence_packing: + enabled: true + megatron_cfg: + env_vars: + TORCH_CUDA_ARCH_LIST: "10.0" + tensor_model_parallel_size: 8 + expert_model_parallel_size: 8 + expert_tensor_parallel_size: 1 + context_parallel_size: 1 + sequence_parallel: true + bias_activation_fusion: false + activation_checkpointing: true + optimizer: + optimizer_cpu_offload: false + optimizer_offload_fraction: 0.0 + exp_avg_dtype: bfloat16 + exp_avg_sq_dtype: bfloat16 + store_param_remainders: true + generation: + backend: megatron + max_new_tokens: 2048 + stop_strings: + - + - + - + - + - + - + mcore_generation_config: + expose_http_server: false + parsers: [] + max_model_len: ${policy.max_total_sequence_length} + cuda_graph_impl: local + inference_cuda_graph_scope: block + buffer_size_gb: 8 + num_cuda_graphs: -1 + block_size_tokens: 256 + use_cuda_graphs_for_non_decode_steps: true + enable_chunked_prefill: true + enable_prefix_caching: false + max_tokens: ${policy.max_total_sequence_length} + kv_cache_management_mode: persist + materialize_only_last_token_logits: true + num_speculative_tokens: 0 + refit_backend: gloo + transformer_impl: transformer_engine + activation_checkpointing: false + tensor_model_parallel_size: 8 + expert_model_parallel_size: 8 + expert_tensor_parallel_size: 1 + sequence_parallel: true + moe_pad_experts_for_cuda_graph_inference: true + colocated: + enabled: false + resources: + gpus_per_node: 4 + num_nodes: 6 + +data: + default: + prompt_file: examples/prompts/clevr_cogent_cot_nemotron_omni.txt + +logger: + wandb_enabled: false + 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/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..e4a670b3570 100644 --- a/nemo_rl/data/multimodal_utils.py +++ b/nemo_rl/data/multimodal_utils.py @@ -22,6 +22,7 @@ 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 @@ -927,6 +928,29 @@ def extract_input_image_sources_from_responses_messages( return sources +def extract_input_video_sources_from_responses_messages(messages: Any) -> list[Any]: + """Extract video sources from Responses-API messages in encounter order.""" + if not isinstance(messages, list): + return [] + + sources: list[Any] = [] + for message in messages: + if not isinstance(message, dict): + continue + content = message.get("content") or [] + if not isinstance(content, list): + continue + for part in content: + if not isinstance(part, dict) or part.get("type") not in VIDEO_CONTENT_TYPES: + continue + source = part.get("video") or part.get("video_url") or part.get("url") + if isinstance(source, dict): + source = source.get("url") + if source is not None: + sources.append(source) + return sources + + def extract_input_images_from_responses_messages( messages: Any, ) -> list[Image.Image]: @@ -1042,30 +1066,45 @@ 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. +_VIDEO_EXT_TO_MIME = { + ".mp4": "mp4", + ".m4v": "mp4", + ".mov": "quicktime", + ".webm": "webm", + ".mkv": "x-matroska", + ".avi": "x-msvideo", +} - 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. +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 - The examples are mutated in place; the same list is also returned for - convenience so callers can chain the call. + 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." + ) - 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. + 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}" - 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]]] = {} + +def normalize_media_in_examples(nemo_gym_examples: list[dict]) -> list[dict]: + """Canonicalize Gym media parts and inline each unique local source.""" + image_targets_by_source: dict[str, list[dict]] = {} for example in nemo_gym_examples: input_items = example.get("responses_create_params", {}).get("input", []) @@ -1078,27 +1117,54 @@ def encode_images_in_examples(nemo_gym_examples: list[dict]) -> list[dict]: if not isinstance(content, list): continue for part in content: - if ( - not isinstance(part, dict) - or part.get("type") not in IMAGE_CONTENT_TYPES - ): + if not isinstance(part, dict): 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: + 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 - if url.startswith(("http://", "https://", "data:")): + + 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 - targets_by_source.setdefault(url, []).append((part, media_key)) + 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: + image_targets_by_source.setdefault(url, []).append(part) + else: + url = video_path_to_data_url(url) + + for key in source_keys: + if key != canonical_key: + part.pop(key, None) + part["type"] = canonical_type + part[canonical_key] = url + if is_image and nested_detail is not None: + part.setdefault("detail", nested_detail) - sources = list(targets_by_source) + sources = list(image_targets_by_source) if sources: with ThreadPoolExecutor( max_workers=NEMO_GYM_IMAGE_ENCODE_MAX_WORKERS @@ -1113,13 +1179,17 @@ def encode_images_in_examples(nemo_gym_examples: list[dict]) -> list[dict]: # 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(): + for source, targets in image_targets_by_source.items(): data_url = encoded_by_source[source] - for part, media_key in targets: - part[media_key] = data_url + for part in targets: + part["image_url"] = data_url return nemo_gym_examples +# Backward-compatible alias for existing callers. +encode_images_in_examples = normalize_media_in_examples + + def get_media_from_message(message: dict[str, Any]) -> dict[str, list[Any]]: """Get all media from a message log item.""" # Handle None or missing content (e.g., assistant messages with only tool_calls) diff --git a/nemo_rl/data/processors.py b/nemo_rl/data/processors.py index 58f09c95856..b52405b10dc 100644 --- a/nemo_rl/data/processors.py +++ b/nemo_rl/data/processors.py @@ -675,9 +675,11 @@ def vlm_hf_data_processor( chat_message[key] = PackedTensor.empty_like(value) 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 have already expanded media positions in + # input_ids. Let vLLM consume those IDs while retaining main's native + # media side channels; other processors keep the formatted-string path. vllm_kwargs = { - "vllm_content": string_formatted_dialog, + "vllm_content": None if uses_placeholder else string_formatted_dialog, "vllm_images": images, "vllm_audios": audios, "vllm_videos": videos, diff --git a/nemo_rl/environments/nemo_gym.py b/nemo_rl/environments/nemo_gym.py index e702f462042..761fc2ea856 100644 --- a/nemo_rl/environments/nemo_gym.py +++ b/nemo_rl/environments/nemo_gym.py @@ -29,8 +29,9 @@ from nemo_rl.data.multimodal_utils import ( attach_image_model_inputs_to_message, - encode_images_in_examples, extract_input_image_sources_from_responses_messages, + extract_input_video_sources_from_responses_messages, + normalize_media_in_examples, resolve_to_image, uses_image_placeholder, ) @@ -414,10 +415,37 @@ def _image_sources_equal(left: Any, right: Any) -> bool: ) -def _without_initial_image_sources( +def _input_media_sources(messages: Any) -> list[tuple[str, Any]]: + if not isinstance(messages, list): + return [] + + sources: list[tuple[str, Any]] = [] + for message in messages: + if not isinstance(message, dict): + continue + content = message.get("content") + if not isinstance(content, list): + continue + for part in content: + part_message = [{"content": [part]}] + images = extract_input_image_sources_from_responses_messages(part_message) + videos = extract_input_video_sources_from_responses_messages(part_message) + sources.extend(("image", source) for source in images) + sources.extend(("video", source) for source in videos) + return sources + + +def _media_sources_equal( + left: tuple[str, Any], + right: tuple[str, Any], +) -> bool: + return left[0] == right[0] and _image_sources_equal(left[1], right[1]) + + +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.""" + """Copy Responses messages and remove one ordered copy of initial media.""" if not isinstance(messages, list): return messages, False @@ -432,13 +460,11 @@ def _without_initial_image_sources( filtered_content = [] for part in content: - part_sources = extract_input_image_sources_from_responses_messages( - [{"content": [part]}] - ) + part_sources = _input_media_sources([{"content": [part]}]) if ( remaining_sources and len(part_sources) == 1 - and _image_sources_equal(part_sources[0], remaining_sources[0]) + and _media_sources_equal(part_sources[0], remaining_sources[0]) ): remaining_sources.pop(0) continue @@ -675,14 +701,7 @@ 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_media_in_examples(nemo_gym_examples) timer.start("_run_rollouts_total") nemo_gym_result_iterator = self.rch.run_examples( @@ -778,20 +797,14 @@ 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_input - ) - agent_initial_sources = extract_input_image_sources_from_responses_messages( - initial_input - ) - returned_media_sources = extract_input_image_sources_from_responses_messages( - media_messages - ) + raw_initial_sources = _input_media_sources(raw_input) + agent_initial_sources = _input_media_sources(initial_input) + returned_media_sources = _input_media_sources(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 +813,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 +824,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 = ( @@ -1003,7 +1016,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 +1125,14 @@ 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": + generation_config["vllm_cfg"]["async_engine"] = True + generation_config["vllm_cfg"]["expose_http_server"] = True + elif backend == "megatron": + 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/evals/eval.py b/nemo_rl/evals/eval.py index 670ed625d7e..1f572d589d9 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,28 @@ async def _run_env_eval_impl( multi_modal_data["video"] = ( videos[i][0] if len(videos[i]) == 1 else videos[i] ) - if multi_modal_data: - prompt_dict["multi_modal_data"] = multi_modal_data + + if multi_modal_data: + vllm_content = batch["vllm_content"][i] + if vllm_content is not None: + prompt_dict = {"prompt": vllm_content} + prompt_display = vllm_content + else: + 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} + prompt_display = "\n".join( + str(message["content"]) for message in message_log + ) + prompt_dict["multi_modal_data"] = multi_modal_data prompts.append(prompt_dict) - prompts_for_display.append(vllm_content) + 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/models/generation/megatron/config.py b/nemo_rl/models/generation/megatron/config.py index d79c4d5ba81..ee6d0d27a72 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] + logprobs_mode: NotRequired[Literal["raw_logprobs", "processed_logprobs"]] refit_backend: Literal["gloo", "nccl", "nvshmem"] num_speculative_tokens: int @@ -55,6 +58,12 @@ class MCoreGenerationSpecificArgs(TypedDict): mamba_inference_ssm_states_dtype: NotRequired[str] mamba_inference_conv_states_dtype: NotRequired[str] + # Raw media preprocessing used by the OpenAI-compatible HTTP endpoint. + video_num_frames: NotRequired[int] + video_temporal_patch_size: NotRequired[int] + video_target_num_patches: NotRequired[int] + multimodal_prompt_config: NotRequired[dict] + # 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 +79,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..d84248e6c6d 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 @@ -231,13 +232,10 @@ def __init__( "policy; when colocated, pass it to the training policy instead." ) - # `self.cfg` exposes the `generation` that matches the `GenerationInterface` contract. - # `self._policy_config` keeps a reference to the full PolicyConfig. - self._policy_config = config + # Dedicated inference may modify its policy config. + self._policy_config = config if policy is not None else deepcopy(config) self.cfg: MCoreGenerationConfig = config["generation"] - # Populated after the first prepare_for_generation (which starts the HTTP server). self.dp_openai_server_base_urls: list[Optional[str]] = [] - # Installed by setup via create_weight_synchronizer. self.weight_synchronizer: Optional["WeightSynchronizer"] = None if policy is not None: diff --git a/nemo_rl/models/generation/megatron/megatron_worker.py b/nemo_rl/models/generation/megatron/megatron_worker.py index 8c9ee1f7c61..969ebef4041 100644 --- a/nemo_rl/models/generation/megatron/megatron_worker.py +++ b/nemo_rl/models/generation/megatron/megatron_worker.py @@ -14,19 +14,23 @@ import asyncio import gc +import importlib import os import threading import time import warnings -from typing import AsyncGenerator, Optional +from dataclasses import replace +from typing import Any, AsyncGenerator, Optional import requests import torch from megatron.core.inference.config import ( + AsyncScheduleMode, InferenceConfig, KVCacheManagementMode, MambaInferenceStateConfig, PrefixCachingCoordinatorPolicy, + VideoProcessingConfig, ) from megatron.core.inference.engines.dynamic_engine import EngineState from megatron.core.inference.sampling_params import SamplingParams @@ -60,9 +64,13 @@ verify_right_padding, ) from nemo_rl.models.generation.megatron.utils import ( + build_image_preprocessing_config, log_gpu_memory, resolve_torch_dtype, ) +from nemo_rl.models.generation.vllm.video_utils import ( + _CACHED_VIDEO_FRAME_MANIFEST_MAGIC, +) from nemo_rl.models.megatron.memory_saver import ( HAVE_TORCH_MEMORY_SAVER, pause_inference_weights, @@ -112,6 +120,60 @@ def _init_inference_engine_state(self) -> None: self._inference_loop = None self._inference_thread = None + def _get_megatron_inference_wrapper_cls(self): + """Resolve the configured Megatron inference wrapper, if any.""" + class_path = self.cfg["generation"]["mcore_generation_config"].get( + "megatron_inference_wrapper" + ) + if class_path is None: + return None + module_name, _, class_name = class_path.rpartition(".") + if not module_name: + raise ValueError( + "megatron_inference_wrapper must be a fully qualified class name." + ) + return getattr(importlib.import_module(module_name), class_name) + + @staticmethod + def _wrapper_supports_modality(inference_wrapper_cls, modality: str) -> bool: + 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 = getattr(self, "processor", None) + if processor is None: + raise ValueError( + "Megatron multimodal generation requires the policy processor." + ) + return build_image_preprocessing_config( + processor.image_processor, + dynamic_resolution=bool( + generation_config.get("image_dynamic_resolution", False) + ), + ) + def _setup_colocated_cuda_graph_managers(self) -> None: """Create inference CUDA-graph managers for shared-model colocated generation. @@ -223,8 +285,12 @@ 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 + ) + 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 +311,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 +338,31 @@ 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 = None + temporal_patch_size = mcore_generation_config.get("video_temporal_patch_size") + if image_preprocessing_config is not None and temporal_patch_size is not None: + video_image_config = image_preprocessing_config + target_num_patches = mcore_generation_config.get("video_target_num_patches") + if target_num_patches is not None: + video_image_config = replace( + video_image_config, + dynamic_resolution_max_patches=int(target_num_patches), + ) + video_preprocessing_config = VideoProcessingConfig( + image_config=video_image_config, + num_frames=int(mcore_generation_config["video_num_frames"]), + temporal_patch_size=int(temporal_patch_size), + frame_manifest_magic=_CACHED_VIDEO_FRAME_MANIFEST_MAGIC, + ) + inference_config = InferenceConfig( block_size_tokens=block_size_tokens, buffer_size_gb=buffer_size_gb, @@ -286,10 +374,16 @@ def _initialize_inference_engine(self, mcore_generation_config: dict) -> None: use_cuda_graphs_for_non_decode_steps=use_cuda_graphs_for_non_decode_steps, use_flashinfer_fused_rope=use_flashinfer_fused_rope, sampling_backend="flashinfer", + async_sched_mode=AsyncScheduleMode( + mcore_generation_config.get("async_sched_mode", "legacy") + ), use_synchronous_zmq_collectives=True, 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"], + vision_embedding_cache_max_bytes=int( + mcore_generation_config.get("vision_embedding_cache_max_bytes", 0) + ), prefix_caching_coordinator_policy=PrefixCachingCoordinatorPolicy( "first_prefix_block" ), @@ -301,21 +395,34 @@ def _initialize_inference_engine(self, mcore_generation_config: dict) -> None: ), logging_step_interval=logging_step_interval, num_speculative_tokens=num_speculative_tokens, - logprobs_mode=mcore_generation_config["logprobs_mode"], + logprobs_mode=mcore_generation_config.get( + "logprobs_mode", "raw_logprobs" + ), 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[ + 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 + model_config, inference_config ) + if media_model is None: + self.inference_wrapped_model = GPTInferenceWrapper( + inference_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( + media_model, self.inference_context + ) text_generation_controller = TextGenerationController( inference_wrapped_model=self.inference_wrapped_model, tokenizer=self.megatron_tokenizer, @@ -408,6 +515,7 @@ def run_loop(): def _setup_openai_api_server(self) -> str: """Start the OpenAI-compatible HTTP server on this worker.""" + from megatron.core.inference.config import MultimodalPromptConfig from megatron.core.inference.text_generation_server.dynamic_text_gen_server.text_generation_server import ( start_text_gen_server, ) @@ -424,6 +532,9 @@ def _setup_openai_api_server(self) -> str: else: server_port = _get_free_port_local() + prompt_config = self.cfg["generation"]["mcore_generation_config"].get( + "multimodal_prompt_config" + ) start_text_gen_server( coordinator_addr=self.coordinator_addr, tokenizer=self.megatron_tokenizer, @@ -432,6 +543,11 @@ def _setup_openai_api_server(self) -> str: parsers=self.cfg["generation"]["mcore_generation_config"]["parsers"], verbose=False, sock=reserved_socket, + multimodal_prompt_config=( + MultimodalPromptConfig.from_dict(prompt_config) + if prompt_config + else None + ), ) base_url = f"http://{ip}:{server_port}/v1" @@ -606,7 +722,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 +746,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 +762,98 @@ 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.""" + 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." + ) + + 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(self, data, index: int): + """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 = self._sample_vision_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" + inference_wrapper_cls = self._get_megatron_inference_wrapper_cls() + if not self._wrapper_supports_modality(inference_wrapper_cls, 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 _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 +866,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 +902,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), @@ -754,7 +987,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 +996,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 +1027,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 +1050,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 +1062,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..0d83f9f7529 100644 --- a/nemo_rl/models/generation/megatron/utils.py +++ b/nemo_rl/models/generation/megatron/utils.py @@ -12,10 +12,70 @@ # See the License for the specific language governing permissions and # limitations under the License. +from typing import Any + import torch +from megatron.core.inference.config import ImageProcessingConfig from megatron.core.inference.utils import device_memory_summary +def build_image_preprocessing_config( + image_processor: Any, *, dynamic_resolution: bool +) -> ImageProcessingConfig: + """Translate an HF image processor to an MCore config.""" + + 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") + + 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 + ] + if missing: + 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), + dynamic_resolution=dynamic_resolution, + 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 resolve_torch_dtype(val): """Convert a value to `torch.dtype`.""" if isinstance(val, torch.dtype): diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index d0f57de802a..295d931065d 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -260,6 +260,26 @@ 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.""" + from megatron.bridge.utils.activation_map import str_to_dtype + + resolved = dict(optimizer_cfg) + for key in _OPTIMIZER_DTYPE_KEYS: + value = resolved.get(key) + if isinstance(value, str): + resolved[key] = str_to_dtype(value) + return resolved + def destroy_parallel_state(): """Safely destroy parallel state and reset async call tracking. @@ -1501,7 +1521,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/scripts/add_nemo_gym_agent_ref.py b/scripts/add_nemo_gym_agent_ref.py new file mode 100644 index 00000000000..95d81dd8f10 --- /dev/null +++ b/scripts/add_nemo_gym_agent_ref.py @@ -0,0 +1,74 @@ +#!/usr/bin/env python3 +# 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. + +"""Stamp an ``agent_ref`` onto NeMo-Gym JSONL rows that are missing one. + +NeMo-RL routes every row to a Gym agent by reading ``agent_ref.name`` +(``nemo_rl/environments/nemo_gym.py``), so a row without it fails the rollout +batch with ``KeyError: 'agent_ref'``. Some Gym generators omit the field even +though the ``example.jsonl`` they ship beside declares it — for instance +``resources_servers/circle_count/generate_data.py``. + +Rewrites files in place and leaves rows that already declare an ``agent_ref`` +untouched, so it is safe to re-run. + +Usage: + python scripts/add_nemo_gym_agent_ref.py \\ + --agent-name circle_count_simple_agent \\ + path/to/train.jsonl path/to/val.jsonl +""" + +import argparse +import json +from pathlib import Path + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "paths", type=Path, nargs="+", help="JSONL files to rewrite in place." + ) + parser.add_argument( + "--agent-name", + required=True, + help=( + "Agent key from the Gym resources server config, e.g. " + "circle_count_simple_agent." + ), + ) + parser.add_argument( + "--agent-type", + default="responses_api_agents", + help="Gym agent type. Only responses_api_agents exists today.", + ) + args = parser.parse_args() + + agent_ref = {"type": args.agent_type, "name": args.agent_name} + + for path in args.paths: + rows = [ + json.loads(line) for line in path.read_text().splitlines() if line.strip() + ] + stamped = 0 + for row in rows: + if "agent_ref" not in row: + row["agent_ref"] = agent_ref + stamped += 1 + path.write_text("".join(json.dumps(row) + "\n" for row in rows)) + print(f"{path}: stamped {stamped} of {len(rows)} row(s)") + + +if __name__ == "__main__": + main() diff --git a/scripts/prepare_nemotron_omni_vstat.py b/scripts/prepare_nemotron_omni_vstat.py new file mode 100644 index 00000000000..e67b7079e2d --- /dev/null +++ b/scripts/prepare_nemotron_omni_vstat.py @@ -0,0 +1,144 @@ +#!/usr/bin/env python3 +"""Download and prepare a small VSTAT split for Nemotron Omni video GRPO.""" + +import argparse +import json +import subprocess +import sys +import zipfile +from pathlib import Path + +import pyarrow.parquet as pq +from huggingface_hub import hf_hub_download + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument( + "--output-dir", + type=Path, + required=True, + help="Directory for extracted media and source JSONL files.", + ) + parser.add_argument( + "--repo-id", + default="ShushengYang/VSTAT", + help="Hugging Face dataset repository.", + ) + parser.add_argument( + "--num-rows", + type=int, + default=8, + help="Maximum number of MCQ examples to prepare.", + ) + return parser.parse_args() + + +def write_jsonl(path: Path, rows: list[dict]) -> None: + with path.open("w", encoding="utf-8") as stream: + for row in rows: + stream.write(json.dumps(row) + "\n") + + +def main() -> None: + args = parse_args() + if args.num_rows < 3: + raise ValueError("--num-rows must be at least 3.") + + root = args.output_dir.expanduser().resolve() + root.mkdir(parents=True, exist_ok=True) + + parquet_path = hf_hub_download( + repo_id=args.repo_id, + repo_type="dataset", + filename="test.parquet", + ) + archive_path = hf_hub_download( + repo_id=args.repo_id, + repo_type="dataset", + filename="videos.zip", + ) + + media_root = root / "media" + marker = media_root / ".extracted" + if not marker.exists(): + media_root.mkdir(parents=True, exist_ok=True) + with zipfile.ZipFile(archive_path) as archive: + archive.extractall(media_root) + marker.touch() + + video_files = { + path.name: path.resolve() for path in media_root.rglob("*.mp4") + } + converted = [] + for row in pq.read_table(parquet_path).to_pylist(): + if str(row.get("answer_type", "")).lower() != "mcq": + continue + + relative_video = str(row["video"]) + video_path = media_root / relative_video + if not video_path.exists(): + video_path = video_files.get(Path(relative_video).name) + if video_path is None or not video_path.exists(): + continue + + choices = [str(choice) for choice in row.get("choices") or []] + if not choices: + continue + letters = "ABCDEFGHIJKLMNOPQRSTUVWXYZ" + options = "\n".join( + f"{letters[index]}. {choice}" + for index, choice in enumerate(choices) + ) + converted.append( + { + "prompt": ( + "Answer the multiple-choice question using the video. " + "Return the final answer as a boxed letter.\n" + f"Question: {row['question']}\n{options}" + ), + "video": str(video_path), + "answer": str(row["answer"]).strip().upper(), + "verifier": "multiple-choice", + } + ) + if len(converted) >= args.num_rows: + break + + if len(converted) < 3: + raise RuntimeError( + f"Only resolved {len(converted)} VSTAT MCQ rows with local videos." + ) + + split = max(2, len(converted) - 2) + train_source = root / "train-source.jsonl" + val_source = root / "val-source.jsonl" + write_jsonl(train_source, converted[:split]) + write_jsonl(val_source, converted[split:]) + + repo_root = Path(__file__).resolve().parents[1] + converter = repo_root / "examples/nemo_gym/prepare_video_dataset.py" + for source, output in ( + (train_source, root / "train-gym.jsonl"), + (val_source, root / "val-gym.jsonl"), + ): + subprocess.run( + [ + sys.executable, + str(converter), + "convert", + "--input", + str(source), + "--output", + str(output), + "--require-video", + ], + cwd=repo_root, + check=True, + ) + + print(f"Prepared {len(converted)} VSTAT examples under {root}") + + +if __name__ == "__main__": + main() diff --git a/scripts/run_nemotron_omni_circle_count_megatron_1n4g.sh b/scripts/run_nemotron_omni_circle_count_megatron_1n4g.sh new file mode 100755 index 00000000000..78ec54898e7 --- /dev/null +++ b/scripts/run_nemotron_omni_circle_count_megatron_1n4g.sh @@ -0,0 +1,387 @@ +#!/usr/bin/env bash +set -euo pipefail + +# 1-node GB200 / 4-GPU smoke for Nemotron Omni on NeMo-Gym's circle_count env. +# circle_count is an IMAGE environment: every row carries one synthetic PNG as a +# base64 data URL plus a "how many circles" question, and the Gym +# resources server rewards an exact match on \boxed{count}. +# +# Unlike the CLEVR launcher this runs the Gym entrypoint +# (examples/nemo_gym/run_grpo_nemo_gym.py): Gym owns the prompts, the reward and +# the generation loop, talking to the policy over the engine's HTTP server. +# +# GENERATION_BACKEND=megatron|vllm (default megatron) +# COLOCATED=true|false (default false) shares every GPU between the trainer and +# the engine instead of splitting them. Colocated async GRPO is only +# supported for Megatron generation. +# ASYNC_GRPO=true|false (default true) selects async or sync GRPO. +# GENERATE_ROWS=N regenerates a larger dataset instead of using the 5-row +# example.jsonl that ships with Gym. +NEMORL="${NEMORL:-/opt/nemo-rl}" +WORKSPACE_ROOT="${WORKSPACE_ROOT:-${NEMORL}/workspace}" +GYM="${GYM:-${NEMORL}/3rdparty/Gym-workspace/Gym}" +MODEL_NAME="${MODEL_NAME:-nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16}" +CONFIG="${CONFIG:-examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-count-1n4g-megatron_generation.v1.yaml}" +GENERATION_BACKEND="${GENERATION_BACKEND:-megatron}" +COLOCATED="${COLOCATED:-false}" +ASYNC_GRPO="${ASYNC_GRPO:-true}" + +cd "${NEMORL}" + +if [[ "${GENERATION_BACKEND}" != "megatron" && "${GENERATION_BACKEND}" != "vllm" ]]; then + echo "GENERATION_BACKEND must be megatron or vllm (got ${GENERATION_BACKEND})." >&2 + exit 1 +fi +if [[ "${COLOCATED}" != "true" && "${COLOCATED}" != "false" ]]; then + echo "COLOCATED must be true or false (got ${COLOCATED})." >&2 + exit 1 +fi +if [[ "${ASYNC_GRPO}" != "true" && "${ASYNC_GRPO}" != "false" ]]; then + echo "ASYNC_GRPO must be true or false (got ${ASYNC_GRPO})." >&2 + exit 1 +fi +if [[ "${COLOCATED}" == "true" && "${ASYNC_GRPO}" == "true" && "${GENERATION_BACKEND}" != "megatron" ]]; then + echo "Colocated async GRPO requires GENERATION_BACKEND=megatron." >&2 + exit 1 +fi +if [[ ! -d "${GYM}" ]]; then + echo "NeMo-Gym checkout is missing: ${GYM}" >&2 + exit 1 +fi + +GPUS_PER_NODE="${GPUS_PER_NODE:-$(nvidia-smi --query-gpu=index --format=csv,noheader | wc -l)}" +if (( GPUS_PER_NODE < 4 )); then + echo "This launcher requires at least four visible GPUs (got ${GPUS_PER_NODE})." >&2 + exit 1 +fi + +if [[ "${COLOCATED}" == "true" ]]; then + # The engine shares the trainer's GPUs and its HTTP server comes up on the + # policy workers, so both worlds span every GPU. + TRAIN_WORLD_SIZE="${GPUS_PER_NODE}" + INFERENCE_WORLD_SIZE="${GPUS_PER_NODE}" + GEN_GPUS_PER_NODE="${GPUS_PER_NODE}" + NUM_GEN_NODES=1 + COLOCATED_ENABLED=true +else + GEN_GPUS_PER_NODE="${GEN_GPUS_PER_NODE:-$((GPUS_PER_NODE / 2))}" + if (( GEN_GPUS_PER_NODE <= 0 || GEN_GPUS_PER_NODE >= GPUS_PER_NODE )); then + echo "Non-colocated mode requires a strict train/inference GPU split." >&2 + exit 1 + fi + TRAIN_WORLD_SIZE=$((GPUS_PER_NODE - GEN_GPUS_PER_NODE)) + INFERENCE_WORLD_SIZE="${GEN_GPUS_PER_NODE}" + NUM_GEN_NODES=1 + COLOCATED_ENABLED=false +fi + +POLICY_TP="${POLICY_TP:-${TRAIN_WORLD_SIZE}}" +INFER_TP="${INFER_TP:-${INFERENCE_WORLD_SIZE}}" +POLICY_CP="${POLICY_CP:-1}" +if [[ "${GENERATION_BACKEND}" == "megatron" && "${POLICY_CP}" != "1" ]]; then + echo "Megatron dynamic inference requires POLICY_CP=1." >&2 + exit 1 +fi +if (( TRAIN_WORLD_SIZE % (POLICY_TP * POLICY_CP) != 0 )); then + echo "Training world size must be divisible by POLICY_TP * POLICY_CP." >&2 + exit 1 +fi +if (( INFERENCE_WORLD_SIZE % INFER_TP != 0 )); then + echo "Inference world size must be divisible by INFER_TP." >&2 + exit 1 +fi + +# With ETP=1, world_size must be divisible by EP (not by TP*EP). +largest_ep() { + local world="$1" + local ep=8 + while (( ep > world || world % ep != 0 )); do + ep=$((ep / 2)) + done + printf '%d' "${ep}" +} + +POLICY_EP="${POLICY_EP:-$(largest_ep "${TRAIN_WORLD_SIZE}")}" +INFER_EP="${INFER_EP:-$(largest_ep "${INFERENCE_WORLD_SIZE}")}" +if (( TRAIN_WORLD_SIZE % POLICY_EP != 0 )); then + echo "Training world size must be divisible by POLICY_EP (ETP=1)." >&2 + exit 1 +fi +if [[ "${GENERATION_BACKEND}" == "vllm" ]]; then + if [[ "${INFER_EP}" != "${INFER_TP}" ]]; then + echo "Forcing INFER_EP=${INFER_TP} for vLLM (was ${INFER_EP}) for async-safe Ray DP." >&2 + fi + INFER_EP="${INFER_TP}" +elif (( INFERENCE_WORLD_SIZE % INFER_EP != 0 )); then + echo "Inference world size must be divisible by INFER_EP (ETP=1)." >&2 + exit 1 +fi +if [[ "${COLOCATED_ENABLED}" == "true" ]]; then + # Colocated Megatron generation reuses the trained model in place, so the + # engine cannot reshard: its parallel layout has to match the policy's. + if (( INFER_TP != POLICY_TP || INFER_EP != POLICY_EP )); then + echo "Colocated mode requires INFER_TP/INFER_EP to match POLICY_TP/POLICY_EP" \ + "(got TP ${INFER_TP} vs ${POLICY_TP}, EP ${INFER_EP} vs ${POLICY_EP})." >&2 + exit 1 + fi +fi + +CACHE_ROOT="${CACHE_ROOT:-${WORKSPACE_ROOT}/cache/nemo-rl-omni}" +export HF_HOME="${HF_HOME:-${CACHE_ROOT}/huggingface}" +export HUGGINGFACE_HUB_CACHE="${HUGGINGFACE_HUB_CACHE:-${HF_HOME}/hub}" +export HF_HUB_CACHE="${HF_HUB_CACHE:-${HUGGINGFACE_HUB_CACHE}}" +export HF_DATASETS_CACHE="${HF_DATASETS_CACHE:-${HF_HOME}/datasets}" +export HF_MODULES_CACHE="${HF_MODULES_CACHE:-${HF_HOME}/modules}" +export TRANSFORMERS_CACHE="${TRANSFORMERS_CACHE:-${HF_HOME}/transformers}" +export NRL_MEGATRON_CHECKPOINT_DIR="${NRL_MEGATRON_CHECKPOINT_DIR:-${CACHE_ROOT}/megatron-checkpoints}" +export XDG_CACHE_HOME="${XDG_CACHE_HOME:-${CACHE_ROOT}/xdg}" +export TORCH_HOME="${TORCH_HOME:-${CACHE_ROOT}/torch}" +export TRITON_CACHE_DIR="${TRITON_CACHE_DIR:-${CACHE_ROOT}/triton}" + +BRIDGE="${NEMORL}/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge" +export PYTHONPATH="${NEMORL}:${GYM}:${BRIDGE}/src:${BRIDGE}/3rdparty/Megatron-LM${PYTHONPATH:+:${PYTHONPATH}}" +export RAY_ENABLE_UV_RUN_RUNTIME_ENV=0 +export NRL_VENVS_TRUST_EXISTING="${NRL_VENVS_TRUST_EXISTING:-1}" +export NRL_FORCE_REBUILD_VENVS="${NRL_FORCE_REBUILD_VENVS:-false}" +export NEMO_RL_VENV_DIR="${NEMO_RL_VENV_DIR:-/opt/ray_venvs}" +# The Gym env servers run in their own venvs; keep them beside the NeMo-RL ones +# so a warm container does not rebuild them on every launch. +export NEMO_GYM_VENV_DIR="${NEMO_GYM_VENV_DIR:-${NEMO_RL_VENV_DIR}}" +export CUDA_DEVICE_MAX_CONNECTIONS="${CUDA_DEVICE_MAX_CONNECTIONS:-1}" +export NVTE_FWD_LAYERNORM_SM_MARGIN="${NVTE_FWD_LAYERNORM_SM_MARGIN:-16}" +export NVTE_BWD_LAYERNORM_SM_MARGIN="${NVTE_BWD_LAYERNORM_SM_MARGIN:-16}" +export FLASHINFER_DISABLE_VERSION_CHECK="${FLASHINFER_DISABLE_VERSION_CHECK:-1}" +export TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST:-10.0}" +export NCCL_DEBUG="${NCCL_DEBUG:-WARN}" + +MAX_STEPS="${MAX_STEPS:-5}" +MAX_SEQUENCE_LENGTH="${MAX_SEQUENCE_LENGTH:-4096}" +MAX_NEW_TOKENS="${MAX_NEW_TOKENS:-512}" +NUM_PROMPTS="${NUM_PROMPTS:-2}" +NUM_GENERATIONS="${NUM_GENERATIONS:-8}" +TRAIN_GBS="${TRAIN_GBS:-$((NUM_PROMPTS * NUM_GENERATIONS))}" +EXPECTED_TRAIN_GBS=$((NUM_PROMPTS * NUM_GENERATIONS)) +if (( TRAIN_GBS != EXPECTED_TRAIN_GBS )); then + echo "TRAIN_GBS (${TRAIN_GBS}) must equal NUM_PROMPTS * NUM_GENERATIONS (${EXPECTED_TRAIN_GBS})." >&2 + exit 1 +fi +REFIT_BACKEND="${REFIT_BACKEND:-nccl}" +JOB_NAME="${JOB_NAME:-nemotron-omni-circle-count-${GENERATION_BACKEND}-1n4g}" +EXP_NAME="${EXP_NAME:-${JOB_NAME}}" +PRECISION_RECIPE="${PRECISION_RECIPE:-bf16}" +WANDB_ENABLED="${WANDB_ENABLED:-false}" +WANDB_PROJ="${WANDB_PROJ:-mllm-rl-dev}" +WANDB_GROUP="${WANDB_GROUP:-adlr}" +WANDB_NAME="${WANDB_NAME:-${EXP_NAME}-${PRECISION_RECIPE}-internal-repo}" +RESULTS_DIR="${RESULTS_DIR:-${WORKSPACE_ROOT}/results/nemo-rl-omni/${JOB_NAME}}" +CHECKPOINTING_ENABLED="${CHECKPOINTING_ENABLED:-false}" +# Host OOM on GB200 when optimizer CPU offload is enabled for this model size. +OPTIMIZER_CPU_OFFLOAD="${OPTIMIZER_CPU_OFFLOAD:-false}" +OFFLOAD_OPTIMIZER_FOR_LOGPROB="${OFFLOAD_OPTIMIZER_FOR_LOGPROB:-false}" +if [[ "${OPTIMIZER_CPU_OFFLOAD}" == "true" ]]; then + OPTIMIZER_OFFLOAD_FRACTION="${OPTIMIZER_OFFLOAD_FRACTION:-1.0}" +else + OPTIMIZER_OFFLOAD_FRACTION="${OPTIMIZER_OFFLOAD_FRACTION:-0.0}" +fi +# Lower-precision Adam moments for the 2-GPU train half (HBM-tight on 1n4g). +EXP_AVG_DTYPE="${EXP_AVG_DTYPE:-bfloat16}" +EXP_AVG_SQ_DTYPE="${EXP_AVG_SQ_DTYPE:-bfloat16}" +STORE_PARAM_REMAINDERS="${STORE_PARAM_REMAINDERS:-true}" +BUFFER_SIZE_GB="${BUFFER_SIZE_GB:-8}" +MAX_TRAJECTORY_AGE_STEPS="${MAX_TRAJECTORY_AGE_STEPS:-2}" +IN_FLIGHT_WEIGHT_UPDATES="${IN_FLIGHT_WEIGHT_UPDATES:-true}" +# Keep CUDA graphs disabled and use chunked prefill by default. +MEGATRON_ENABLE_CHUNKED_PREFILL="${MEGATRON_ENABLE_CHUNKED_PREFILL:-true}" +MEGATRON_CUDA_GRAPH_IMPL="${MEGATRON_CUDA_GRAPH_IMPL:-local}" # local none +MEGATRON_CUDA_GRAPH_SCOPE="${MEGATRON_CUDA_GRAPH_SCOPE:-block}" # block layer none +MEGATRON_NUM_CUDA_GRAPHS="${MEGATRON_NUM_CUDA_GRAPHS:--1}" +MEGATRON_USE_CUDA_GRAPHS_FOR_NON_DECODE="${MEGATRON_USE_CUDA_GRAPHS_FOR_NON_DECODE:-false}" +MOE_PAD_EXPERTS_FOR_CG="${MOE_PAD_EXPERTS_FOR_CG:-true}" + +# Dataset. Gym ships a 5-row example; GENERATE_ROWS=N synthesizes a larger set +# (N train rows plus a disjoint 25%-sized validation split) under DATA_ROOT. +DATA_ROOT="${DATA_ROOT:-${WORKSPACE_ROOT}/datasets/circle-count-smoke}" +GENERATE_ROWS="${GENERATE_ROWS:-64}" +GYM_CIRCLE_COUNT="${GYM}/resources_servers/circle_count" +TRAIN_JSONL="${TRAIN_JSONL:-${GYM_CIRCLE_COUNT}/data/example.jsonl}" +VAL_JSONL="${VAL_JSONL:-${TRAIN_JSONL}}" +if (( GENERATE_ROWS > 0 )); then + VAL_ROWS=$(( GENERATE_ROWS / 4 )) + (( VAL_ROWS > 0 )) || VAL_ROWS=1 + mkdir -p "${DATA_ROOT}" + TRAIN_JSONL="${DATA_ROOT}/train.jsonl" + VAL_JSONL="${DATA_ROOT}/val.jsonl" + echo "Generating ${GENERATE_ROWS} train / ${VAL_ROWS} val circle_count rows under ${DATA_ROOT}" + uv run --no-sync python "${GYM_CIRCLE_COUNT}/generate_data.py" \ + --n "${GENERATE_ROWS}" --out "${TRAIN_JSONL}" --seed-offset 0 + # Offset the seeds so validation scenes are disjoint from training ones. + uv run --no-sync python "${GYM_CIRCLE_COUNT}/generate_data.py" \ + --n "${VAL_ROWS}" --out "${VAL_JSONL}" --seed-offset "${GENERATE_ROWS}" + # generate_data.py omits the agent_ref that NeMo-RL reads off every row to + # route it to a Gym agent, unlike the example.jsonl it ships beside. The name + # is the agent key in resources_servers/circle_count/configs/circle_count.yaml. + uv run --no-sync python scripts/add_nemo_gym_agent_ref.py \ + --agent-name circle_count_simple_agent \ + "${TRAIN_JSONL}" "${VAL_JSONL}" +fi + +# NSYS: ENABLE_NSYS=true NRL_NSYS_PROFILE_STEP_RANGE=1:4 +ENABLE_NSYS="${ENABLE_NSYS:-false}" +NSYS_ENV=() +if [[ "${ENABLE_NSYS}" == "true" ]]; then + NRL_NSYS_WORKER_PATTERNS="${NRL_NSYS_WORKER_PATTERNS:-*policy*,*megatron*}" + NRL_NSYS_PROFILE_STEP_RANGE="${NRL_NSYS_PROFILE_STEP_RANGE:-1:4}" + LD_LIBRARY_PATH="/usr/local/cuda/targets/aarch64-linux/lib:/usr/local/cuda/targets/x86_64-linux/lib:/usr/local/cuda/lib64:/usr/local/cuda/lib:/usr/local/nvidia/lib64:/usr/local/nvidia/lib:/usr/lib/aarch64-linux-gnu:/usr/lib/x86_64-linux-gnu:${LD_LIBRARY_PATH:-}" + NRL_NSYS_EXTRA_OPTIONS="${NRL_NSYS_EXTRA_OPTIONS:-{\"o\":\"/opt/nemo-rl/workspace/nsys/%p\",\"cpuctxsw\":\"none\",\"force-overwrite\":\"true\"}}" + NSYS_ENV=( + "NRL_NSYS_WORKER_PATTERNS=${NRL_NSYS_WORKER_PATTERNS}" + "NRL_NSYS_PROFILE_STEP_RANGE=${NRL_NSYS_PROFILE_STEP_RANGE}" + "LD_LIBRARY_PATH=${LD_LIBRARY_PATH}" + "NRL_NSYS_EXTRA_OPTIONS=${NRL_NSYS_EXTRA_OPTIONS}" + ) + mkdir -p /opt/nemo-rl/workspace/nsys +fi + +mkdir -p \ + "${HF_HUB_CACHE}" \ + "${HF_DATASETS_CACHE}" \ + "${HF_MODULES_CACHE}" \ + "${TRANSFORMERS_CACHE}" \ + "${NRL_MEGATRON_CHECKPOINT_DIR}" \ + "${XDG_CACHE_HOME}" \ + "${TORCH_HOME}" \ + "${TRITON_CACHE_DIR}" \ + "${RESULTS_DIR}" + +if [[ ! -f "${CONFIG}" ]]; then + echo "Config is missing under ${NEMORL}: ${CONFIG}" >&2 + exit 1 +fi +for jsonl in "${TRAIN_JSONL}" "${VAL_JSONL}"; do + if [[ ! -s "${jsonl}" ]]; then + echo "Gym dataset is missing or empty: ${jsonl}" >&2 + echo "Generate one with: GENERATE_ROWS=64 $0" >&2 + exit 1 + fi +done + +COMMON_OVERRIDES=( + cluster.num_nodes=1 + cluster.gpus_per_node="${GPUS_PER_NODE}" + policy.model_name="${MODEL_NAME}" + policy.tokenizer.name="${MODEL_NAME}" + policy.is_vlm=true + policy.megatron_cfg.tensor_model_parallel_size="${POLICY_TP}" + policy.megatron_cfg.expert_model_parallel_size="${POLICY_EP}" + policy.megatron_cfg.expert_tensor_parallel_size=1 + policy.megatron_cfg.context_parallel_size="${POLICY_CP}" + policy.megatron_cfg.optimizer.optimizer_cpu_offload="${OPTIMIZER_CPU_OFFLOAD}" + policy.megatron_cfg.optimizer.optimizer_offload_fraction="${OPTIMIZER_OFFLOAD_FRACTION}" + policy.offload_optimizer_for_logprob="${OFFLOAD_OPTIMIZER_FOR_LOGPROB}" +) +if [[ -n "${EXP_AVG_DTYPE}" ]]; then + COMMON_OVERRIDES+=("++policy.megatron_cfg.optimizer.exp_avg_dtype=${EXP_AVG_DTYPE}") +fi +if [[ -n "${EXP_AVG_SQ_DTYPE}" ]]; then + COMMON_OVERRIDES+=("++policy.megatron_cfg.optimizer.exp_avg_sq_dtype=${EXP_AVG_SQ_DTYPE}") +fi +if [[ -n "${STORE_PARAM_REMAINDERS}" ]]; then + COMMON_OVERRIDES+=("++policy.megatron_cfg.optimizer.store_param_remainders=${STORE_PARAM_REMAINDERS}") +fi +COMMON_OVERRIDES+=( + policy.generation.backend="${GENERATION_BACKEND}" + policy.generation.colocated.enabled="${COLOCATED_ENABLED}" + policy.generation.colocated.resources.num_nodes="${NUM_GEN_NODES}" + policy.generation.colocated.resources.gpus_per_node="${GEN_GPUS_PER_NODE}" + policy.max_total_sequence_length="${MAX_SEQUENCE_LENGTH}" + policy.generation.max_new_tokens="${MAX_NEW_TOKENS}" + data.train.data_path="${TRAIN_JSONL}" + data.validation.data_path="${VAL_JSONL}" + grpo.async_grpo.enabled="${ASYNC_GRPO}" + grpo.async_grpo.max_trajectory_age_steps="${MAX_TRAJECTORY_AGE_STEPS}" + grpo.async_grpo.in_flight_weight_updates="${IN_FLIGHT_WEIGHT_UPDATES}" + loss_fn.use_importance_sampling_correction=true + grpo.num_prompts_per_step="${NUM_PROMPTS}" + grpo.num_generations_per_prompt="${NUM_GENERATIONS}" + grpo.val_period=0 + grpo.val_at_start=false + grpo.val_at_end=false + policy.train_global_batch_size="${TRAIN_GBS}" + grpo.max_num_steps="${MAX_STEPS}" + checkpointing.enabled="${CHECKPOINTING_ENABLED}" + checkpointing.checkpoint_dir="${RESULTS_DIR}" + logger.log_dir="${RESULTS_DIR}" + logger.wandb_enabled="${WANDB_ENABLED}" + logger.tensorboard_enabled=false + logger.wandb.name="${WANDB_NAME}" + logger.wandb.project="${WANDB_PROJ}" + +logger.wandb.entity="${WANDB_GROUP}" +) + +GEN_OVERRIDES=() +if [[ "${GENERATION_BACKEND}" == "megatron" ]]; then + GEN_OVERRIDES=( + policy.generation.mcore_generation_config.tensor_model_parallel_size="${INFER_TP}" + policy.generation.mcore_generation_config.expert_model_parallel_size="${INFER_EP}" + policy.generation.mcore_generation_config.expert_tensor_parallel_size=1 + ++policy.generation.mcore_generation_config.context_parallel_size="${POLICY_CP}" + policy.generation.mcore_generation_config.enable_chunked_prefill="${MEGATRON_ENABLE_CHUNKED_PREFILL}" + ++policy.generation.mcore_generation_config.async_sched_mode=async + policy.generation.mcore_generation_config.cuda_graph_impl="${MEGATRON_CUDA_GRAPH_IMPL}" + policy.generation.mcore_generation_config.inference_cuda_graph_scope="${MEGATRON_CUDA_GRAPH_SCOPE}" + policy.generation.mcore_generation_config.num_cuda_graphs="${MEGATRON_NUM_CUDA_GRAPHS}" + policy.generation.mcore_generation_config.use_cuda_graphs_for_non_decode_steps="${MEGATRON_USE_CUDA_GRAPHS_FOR_NON_DECODE}" + policy.generation.mcore_generation_config.refit_backend="${REFIT_BACKEND}" + policy.generation.mcore_generation_config.buffer_size_gb="${BUFFER_SIZE_GB}" + policy.generation.mcore_generation_config.moe_pad_experts_for_cuda_graph_inference="${MOE_PAD_EXPERTS_FOR_CG}" + policy.generation.mcore_generation_config.max_model_len="${MAX_SEQUENCE_LENGTH}" + policy.generation.mcore_generation_config.max_tokens="${MAX_SEQUENCE_LENGTH}" + ) +else + # Refit packs weights into chunks sized at NRL_REFIT_BUFFER_MEMORY_RATIO * + # total HBM; exported globally on purpose because producer and consumer must + # agree on the chunk boundaries. + export NRL_REFIT_BUFFER_MEMORY_RATIO="${NRL_REFIT_BUFFER_MEMORY_RATIO:-0.005}" + VLLM_GPU_MEMORY_UTILIZATION="${VLLM_GPU_MEMORY_UTILIZATION:-0.5}" + VLLM_ENFORCE_EAGER="${VLLM_ENFORCE_EAGER:-true}" + VLLM_MAX_NUM_BATCHED_TOKENS="${VLLM_MAX_NUM_BATCHED_TOKENS:-${MAX_SEQUENCE_LENGTH}}" + GEN_OVERRIDES=( + policy.generation.vllm_cfg.async_engine="${ASYNC_GRPO}" + policy.generation.vllm_cfg.skip_tokenizer_init=false + policy.generation.vllm_cfg.tensor_parallel_size="${INFER_TP}" + policy.generation.vllm_cfg.pipeline_parallel_size=1 + policy.generation.vllm_cfg.expert_parallel_size="${INFER_EP}" + policy.generation.vllm_cfg.max_model_len="${MAX_SEQUENCE_LENGTH}" + ++policy.generation.vllm_cfg.cap_max_tokens_to_context=true + policy.generation.vllm_cfg.gpu_memory_utilization="${VLLM_GPU_MEMORY_UTILIZATION}" + policy.generation.vllm_cfg.enforce_eager="${VLLM_ENFORCE_EAGER}" + ++policy.generation.vllm_cfg.enable_prefix_caching=false + policy.generation.vllm_cfg.logprobs_mode=raw_logprobs + ++policy.generation.vllm_kwargs.limit_mm_per_prompt.image=1 + ++policy.generation.vllm_kwargs.max_num_batched_tokens="${VLLM_MAX_NUM_BATCHED_TOKENS}" + ++policy.generation.vllm_kwargs.mamba_ssm_cache_dtype=float32 + ++policy.generation.vllm_kwargs.skip_mm_profiling=true + ++policy.generation.vllm_kwargs.kernel_config.enable_flashinfer_autotune=false + ++policy.generation.vllm_kwargs.kernel_config.moe_backend=triton + # vLLM bans these tokens rather than treating them as stop strings, which + # Gym clears anyway. + '++policy.generation.bad_words=["","","","","",""]' + ) +fi + +echo "Launching ${JOB_NAME}: ${GPUS_PER_NODE} visible GPU(s)" +echo " env: NeMo-Gym circle_count (images, base64 data URLs) via ${GYM}" +echo " generation backend: ${GENERATION_BACKEND} colocated=${COLOCATED_ENABLED} async=${ASYNC_GRPO}" +echo " async: max_trajectory_age=${MAX_TRAJECTORY_AGE_STEPS} in_flight_weight_updates=${IN_FLIGHT_WEIGHT_UPDATES}" +echo " training world size: ${TRAIN_WORLD_SIZE} (TP=${POLICY_TP}, EP=${POLICY_EP}, ETP=1)" +echo " inference world size: ${INFERENCE_WORLD_SIZE} (TP=${INFER_TP}, EP=${INFER_EP})" +echo " data: train=${TRAIN_JSONL} ($(wc -l < "${TRAIN_JSONL}") rows) val=${VAL_JSONL} ($(wc -l < "${VAL_JSONL}") rows)" +echo " seq/new_tokens: ${MAX_SEQUENCE_LENGTH}/${MAX_NEW_TOKENS}" +echo " optimizer moments: exp_avg=${EXP_AVG_DTYPE:-} exp_avg_sq=${EXP_AVG_SQ_DTYPE:-} store_param_remainders=${STORE_PARAM_REMAINDERS:-}" +echo " W&B: ${WANDB_GROUP}/${WANDB_PROJ}/${WANDB_NAME} (enabled=${WANDB_ENABLED})" + +exec env "${NSYS_ENV[@]}" uv run --no-sync python examples/nemo_gym/run_grpo_nemo_gym.py \ + --config "${CONFIG}" \ + "${COMMON_OVERRIDES[@]}" \ + "${GEN_OVERRIDES[@]}" \ + "$@" diff --git a/scripts/run_nemotron_omni_clevr_megatron_1n4g.sh b/scripts/run_nemotron_omni_clevr_megatron_1n4g.sh new file mode 100755 index 00000000000..bd7b257cec8 --- /dev/null +++ b/scripts/run_nemotron_omni_clevr_megatron_1n4g.sh @@ -0,0 +1,324 @@ +#!/usr/bin/env bash +set -euo pipefail + +# 1-node GB200 / 4-GPU smoke for Nemotron Omni CLEVR. +# Default: Megatron generation, non-colocated 2/2, async GRPO, age>1, in-flight. +# +# GENERATION_BACKEND=megatron|vllm (default megatron) +# - megatron: MCore generation; colocated async supported (NVIDIA-NeMo/RL#2884) +# - vllm: async multimodal via NVIDIA-NeMo/RL#3414; colocated async not supported +# COLOCATED=true ASYNC_GRPO=true requires GENERATION_BACKEND=megatron +# COLOCATED=true ASYNC_GRPO=false runs sync colocated GRPO +# Optimizer moments default to bf16 on this 1n4g split; set EXP_AVG_DTYPE=float32 +# EXP_AVG_SQ_DTYPE=float32 STORE_PARAM_REMAINDERS=false for full precision. +NEMORL="${NEMORL:-/opt/nemo-rl}" +WORKSPACE_ROOT="${WORKSPACE_ROOT:-${NEMORL}/workspace}" +MODEL_NAME="${MODEL_NAME:-nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16}" +CONFIG="${CONFIG:-examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.yaml}" +GENERATION_BACKEND="${GENERATION_BACKEND:-megatron}" +COLOCATED="${COLOCATED:-false}" +ASYNC_GRPO="${ASYNC_GRPO:-true}" + +cd "${NEMORL}" + +if [[ "${GENERATION_BACKEND}" != "megatron" && "${GENERATION_BACKEND}" != "vllm" ]]; then + echo "GENERATION_BACKEND must be megatron or vllm (got ${GENERATION_BACKEND})." >&2 + exit 1 +fi +if [[ "${COLOCATED}" == "true" && "${ASYNC_GRPO}" == "true" && "${GENERATION_BACKEND}" != "megatron" ]]; then + echo "Colocated async GRPO requires GENERATION_BACKEND=megatron." >&2 + exit 1 +fi + +GPUS_PER_NODE="${GPUS_PER_NODE:-$(nvidia-smi --query-gpu=index --format=csv,noheader | wc -l)}" +if (( GPUS_PER_NODE < 4 )); then + echo "This launcher requires at least four visible GPUs (got ${GPUS_PER_NODE})." >&2 + exit 1 +fi + +if [[ "${COLOCATED}" == "true" ]]; then + TRAIN_WORLD_SIZE="${GPUS_PER_NODE}" + INFERENCE_WORLD_SIZE="${GPUS_PER_NODE}" + GEN_GPUS_PER_NODE="${GPUS_PER_NODE}" + NUM_GEN_NODES=1 + COLOCATED_ENABLED=true +else + GEN_GPUS_PER_NODE="${GEN_GPUS_PER_NODE:-$((GPUS_PER_NODE / 2))}" + if (( GEN_GPUS_PER_NODE <= 0 || GEN_GPUS_PER_NODE >= GPUS_PER_NODE )); then + echo "Non-colocated mode requires a strict train/inference GPU split." >&2 + exit 1 + fi + TRAIN_WORLD_SIZE=$((GPUS_PER_NODE - GEN_GPUS_PER_NODE)) + INFERENCE_WORLD_SIZE="${GEN_GPUS_PER_NODE}" + NUM_GEN_NODES=1 + COLOCATED_ENABLED=false +fi + +POLICY_TP="${POLICY_TP:-${TRAIN_WORLD_SIZE}}" +INFER_TP="${INFER_TP:-${INFERENCE_WORLD_SIZE}}" +POLICY_CP="${POLICY_CP:-1}" +if [[ "${GENERATION_BACKEND}" == "megatron" && "${POLICY_CP}" != "1" ]]; then + echo "Megatron dynamic inference requires POLICY_CP=1." >&2 + exit 1 +fi +if (( TRAIN_WORLD_SIZE % (POLICY_TP * POLICY_CP) != 0 )); then + echo "Training world size must be divisible by POLICY_TP * POLICY_CP." >&2 + exit 1 +fi +if (( INFERENCE_WORLD_SIZE % INFER_TP != 0 )); then + echo "Inference world size must be divisible by INFER_TP." >&2 + exit 1 +fi + +# With ETP=1, world_size must be divisible by EP (not by TP*EP). +largest_ep() { + local world="$1" + local ep=8 + while (( ep > world || world % ep != 0 )); do + ep=$((ep / 2)) + done + printf '%d' "${ep}" +} + +POLICY_EP="${POLICY_EP:-$(largest_ep "${TRAIN_WORLD_SIZE}")}" +INFER_EP="${INFER_EP:-$(largest_ep "${INFERENCE_WORLD_SIZE}")}" +if (( TRAIN_WORLD_SIZE % POLICY_EP != 0 )); then + echo "Training world size must be divisible by POLICY_EP (ETP=1)." >&2 + exit 1 +fi +if [[ "${GENERATION_BACKEND}" == "vllm" ]]; then + if [[ "${INFER_EP}" != "${INFER_TP}" ]]; then + echo "Forcing INFER_EP=${INFER_TP} for vLLM (was ${INFER_EP}) for async-safe Ray DP." >&2 + fi + INFER_EP="${INFER_TP}" +elif (( INFERENCE_WORLD_SIZE % INFER_EP != 0 )); then + echo "Inference world size must be divisible by INFER_EP (ETP=1)." >&2 + exit 1 +fi + +CACHE_ROOT="${CACHE_ROOT:-${WORKSPACE_ROOT}/cache/nemo-rl-omni}" +export HF_HOME="${HF_HOME:-${CACHE_ROOT}/huggingface}" +export HUGGINGFACE_HUB_CACHE="${HUGGINGFACE_HUB_CACHE:-${HF_HOME}/hub}" +export HF_HUB_CACHE="${HF_HUB_CACHE:-${HUGGINGFACE_HUB_CACHE}}" +export HF_DATASETS_CACHE="${HF_DATASETS_CACHE:-${HF_HOME}/datasets}" +export HF_MODULES_CACHE="${HF_MODULES_CACHE:-${HF_HOME}/modules}" +export TRANSFORMERS_CACHE="${TRANSFORMERS_CACHE:-${HF_HOME}/transformers}" +export NRL_MEGATRON_CHECKPOINT_DIR="${NRL_MEGATRON_CHECKPOINT_DIR:-${CACHE_ROOT}/megatron-checkpoints}" +export XDG_CACHE_HOME="${XDG_CACHE_HOME:-${CACHE_ROOT}/xdg}" +export TORCH_HOME="${TORCH_HOME:-${CACHE_ROOT}/torch}" +export TRITON_CACHE_DIR="${TRITON_CACHE_DIR:-${CACHE_ROOT}/triton}" + +BRIDGE="${NEMORL}/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge" +export PYTHONPATH="${NEMORL}:${BRIDGE}/src:${BRIDGE}/3rdparty/Megatron-LM${PYTHONPATH:+:${PYTHONPATH}}" +export RAY_ENABLE_UV_RUN_RUNTIME_ENV=0 +export NRL_VENVS_TRUST_EXISTING="${NRL_VENVS_TRUST_EXISTING:-1}" +export NRL_FORCE_REBUILD_VENVS="${NRL_FORCE_REBUILD_VENVS:-false}" +export NEMO_RL_VENV_DIR="${NEMO_RL_VENV_DIR:-/opt/ray_venvs}" +export CUDA_DEVICE_MAX_CONNECTIONS="${CUDA_DEVICE_MAX_CONNECTIONS:-1}" +export NVTE_FWD_LAYERNORM_SM_MARGIN="${NVTE_FWD_LAYERNORM_SM_MARGIN:-16}" +export NVTE_BWD_LAYERNORM_SM_MARGIN="${NVTE_BWD_LAYERNORM_SM_MARGIN:-16}" +export FLASHINFER_DISABLE_VERSION_CHECK="${FLASHINFER_DISABLE_VERSION_CHECK:-1}" +export TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST:-10.0}" +export NCCL_DEBUG="${NCCL_DEBUG:-WARN}" + +MAX_STEPS="${MAX_STEPS:-5}" +MAX_SEQUENCE_LENGTH="${MAX_SEQUENCE_LENGTH:-2048}" +MAX_NEW_TOKENS="${MAX_NEW_TOKENS:-512}" +NUM_PROMPTS="${NUM_PROMPTS:-2}" +NUM_GENERATIONS="${NUM_GENERATIONS:-8}" +TRAIN_GBS="${TRAIN_GBS:-$((NUM_PROMPTS * NUM_GENERATIONS))}" +EXPECTED_TRAIN_GBS=$((NUM_PROMPTS * NUM_GENERATIONS)) +if (( TRAIN_GBS != EXPECTED_TRAIN_GBS )); then + echo "TRAIN_GBS (${TRAIN_GBS}) must equal NUM_PROMPTS * NUM_GENERATIONS (${EXPECTED_TRAIN_GBS})." >&2 + exit 1 +fi +REFIT_BACKEND="${REFIT_BACKEND:-nccl}" +JOB_NAME="${JOB_NAME:-nemotron-omni-clevr-${GENERATION_BACKEND}-1n4g}" +EXP_NAME="${EXP_NAME:-${JOB_NAME}}" +PRECISION_RECIPE="${PRECISION_RECIPE:-bf16}" +WANDB_ENABLED="${WANDB_ENABLED:-false}" +WANDB_PROJ="${WANDB_PROJ:-mllm-rl-dev}" +WANDB_GROUP="${WANDB_GROUP:-adlr}" +WANDB_NAME="${WANDB_NAME:-${EXP_NAME}-${PRECISION_RECIPE}-internal-repo}" +RESULTS_DIR="${RESULTS_DIR:-${WORKSPACE_ROOT}/results/nemo-rl-omni/${JOB_NAME}}" +CHECKPOINTING_ENABLED="${CHECKPOINTING_ENABLED:-false}" +# Host OOM on GB200 when optimizer CPU offload is enabled for this model size. +OPTIMIZER_CPU_OFFLOAD="${OPTIMIZER_CPU_OFFLOAD:-false}" +OFFLOAD_OPTIMIZER_FOR_LOGPROB="${OFFLOAD_OPTIMIZER_FOR_LOGPROB:-false}" +if [[ "${OPTIMIZER_CPU_OFFLOAD}" == "true" ]]; then + OPTIMIZER_OFFLOAD_FRACTION="${OPTIMIZER_OFFLOAD_FRACTION:-1.0}" +else + OPTIMIZER_OFFLOAD_FRACTION="${OPTIMIZER_OFFLOAD_FRACTION:-0.0}" +fi +# Lower-precision Adam moments for the 2-GPU train half (HBM-tight on 1n4g). +# Override with EXP_AVG_DTYPE=float32 EXP_AVG_SQ_DTYPE=float32 STORE_PARAM_REMAINDERS=false +# for full-precision optimizer state, or set any to empty to skip the override. +EXP_AVG_DTYPE="${EXP_AVG_DTYPE:-bfloat16}" +EXP_AVG_SQ_DTYPE="${EXP_AVG_SQ_DTYPE:-bfloat16}" +STORE_PARAM_REMAINDERS="${STORE_PARAM_REMAINDERS:-true}" +BUFFER_SIZE_GB="${BUFFER_SIZE_GB:-8}" +MAX_TRAJECTORY_AGE_STEPS="${MAX_TRAJECTORY_AGE_STEPS:-2}" +IN_FLIGHT_WEIGHT_UPDATES="${IN_FLIGHT_WEIGHT_UPDATES:-true}" +# Use decode-only CUDA graphs with chunked prefill by default. +MEGATRON_ENABLE_CHUNKED_PREFILL="${MEGATRON_ENABLE_CHUNKED_PREFILL:-true}" +MEGATRON_CUDA_GRAPH_IMPL="${MEGATRON_CUDA_GRAPH_IMPL:-local}" # none local +MEGATRON_CUDA_GRAPH_SCOPE="${MEGATRON_CUDA_GRAPH_SCOPE:-block}" # layer block none +MEGATRON_NUM_CUDA_GRAPHS="${MEGATRON_NUM_CUDA_GRAPHS:--1}" # -1 +MEGATRON_USE_CUDA_GRAPHS_FOR_NON_DECODE="${MEGATRON_USE_CUDA_GRAPHS_FOR_NON_DECODE:-false}" +MOE_PAD_EXPERTS_FOR_CG="${MOE_PAD_EXPERTS_FOR_CG:-true}" + +# NSYS: ENABLE_NSYS=true NRL_NSYS_PROFILE_STEP_RANGE=1:4 +ENABLE_NSYS="${ENABLE_NSYS:-false}" +NSYS_ENV=() +if [[ "${ENABLE_NSYS}" == "true" ]]; then + NRL_NSYS_WORKER_PATTERNS="${NRL_NSYS_WORKER_PATTERNS:-*policy*,*megatron*}" + NRL_NSYS_PROFILE_STEP_RANGE="${NRL_NSYS_PROFILE_STEP_RANGE:-1:4}" + LD_LIBRARY_PATH="/usr/local/cuda/targets/aarch64-linux/lib:/usr/local/cuda/targets/x86_64-linux/lib:/usr/local/cuda/lib64:/usr/local/cuda/lib:/usr/local/nvidia/lib64:/usr/local/nvidia/lib:/usr/lib/aarch64-linux-gnu:/usr/lib/x86_64-linux-gnu:${LD_LIBRARY_PATH:-}" + NRL_NSYS_EXTRA_OPTIONS="${NRL_NSYS_EXTRA_OPTIONS:-{\"o\":\"/opt/nemo-rl/workspace/nsys/%p\",\"cpuctxsw\":\"none\",\"force-overwrite\":\"true\"}}" + NSYS_ENV=( + "NRL_NSYS_WORKER_PATTERNS=${NRL_NSYS_WORKER_PATTERNS}" + "NRL_NSYS_PROFILE_STEP_RANGE=${NRL_NSYS_PROFILE_STEP_RANGE}" + "LD_LIBRARY_PATH=${LD_LIBRARY_PATH}" + "NRL_NSYS_EXTRA_OPTIONS=${NRL_NSYS_EXTRA_OPTIONS}" + ) + mkdir -p /opt/nemo-rl/workspace/nsys +fi + +mkdir -p \ + "${HF_HUB_CACHE}" \ + "${HF_DATASETS_CACHE}" \ + "${HF_MODULES_CACHE}" \ + "${TRANSFORMERS_CACHE}" \ + "${NRL_MEGATRON_CHECKPOINT_DIR}" \ + "${XDG_CACHE_HOME}" \ + "${TORCH_HOME}" \ + "${TRITON_CACHE_DIR}" \ + "${RESULTS_DIR}" + +if [[ ! -f "${CONFIG}" ]]; then + echo "Config is missing under ${NEMORL}: ${CONFIG}" >&2 + exit 1 +fi + +COMMON_OVERRIDES=( + cluster.num_nodes=1 + cluster.gpus_per_node="${GPUS_PER_NODE}" + policy.model_name="${MODEL_NAME}" + policy.tokenizer.name="${MODEL_NAME}" + policy.is_vlm=true + policy.megatron_cfg.tensor_model_parallel_size="${POLICY_TP}" + policy.megatron_cfg.expert_model_parallel_size="${POLICY_EP}" + policy.megatron_cfg.expert_tensor_parallel_size=1 + policy.megatron_cfg.context_parallel_size="${POLICY_CP}" + policy.megatron_cfg.optimizer.optimizer_cpu_offload="${OPTIMIZER_CPU_OFFLOAD}" + policy.megatron_cfg.optimizer.optimizer_offload_fraction="${OPTIMIZER_OFFLOAD_FRACTION}" + policy.offload_optimizer_for_logprob="${OFFLOAD_OPTIMIZER_FOR_LOGPROB}" +) +if [[ -n "${EXP_AVG_DTYPE}" ]]; then + COMMON_OVERRIDES+=("++policy.megatron_cfg.optimizer.exp_avg_dtype=${EXP_AVG_DTYPE}") +fi +if [[ -n "${EXP_AVG_SQ_DTYPE}" ]]; then + COMMON_OVERRIDES+=("++policy.megatron_cfg.optimizer.exp_avg_sq_dtype=${EXP_AVG_SQ_DTYPE}") +fi +if [[ -n "${STORE_PARAM_REMAINDERS}" ]]; then + COMMON_OVERRIDES+=("++policy.megatron_cfg.optimizer.store_param_remainders=${STORE_PARAM_REMAINDERS}") +fi +COMMON_OVERRIDES+=( + policy.generation.backend="${GENERATION_BACKEND}" + policy.generation.colocated.enabled="${COLOCATED_ENABLED}" + policy.generation.colocated.resources.num_nodes="${NUM_GEN_NODES}" + policy.generation.colocated.resources.gpus_per_node="${GEN_GPUS_PER_NODE}" + policy.max_total_sequence_length="${MAX_SEQUENCE_LENGTH}" + policy.generation.max_new_tokens="${MAX_NEW_TOKENS}" + grpo.async_grpo.enabled="${ASYNC_GRPO}" + grpo.async_grpo.max_trajectory_age_steps="${MAX_TRAJECTORY_AGE_STEPS}" + grpo.async_grpo.in_flight_weight_updates="${IN_FLIGHT_WEIGHT_UPDATES}" + loss_fn.use_importance_sampling_correction=true + grpo.num_prompts_per_step="${NUM_PROMPTS}" + grpo.num_generations_per_prompt="${NUM_GENERATIONS}" + grpo.val_period=0 + grpo.val_at_start=false + grpo.val_at_end=false + policy.train_global_batch_size="${TRAIN_GBS}" + grpo.max_num_steps="${MAX_STEPS}" + checkpointing.enabled="${CHECKPOINTING_ENABLED}" + checkpointing.checkpoint_dir="${RESULTS_DIR}" + logger.log_dir="${RESULTS_DIR}" + logger.wandb_enabled="${WANDB_ENABLED}" + logger.tensorboard_enabled=false + logger.wandb.name="${WANDB_NAME}" + logger.wandb.project="${WANDB_PROJ}" + +logger.wandb.entity="${WANDB_GROUP}" +) + +GEN_OVERRIDES=() +if [[ "${GENERATION_BACKEND}" == "megatron" ]]; then + GEN_OVERRIDES=( + ++policy.generation.stop_strings=null + ++policy.generation.bad_words=null + policy.generation.mcore_generation_config.tensor_model_parallel_size="${INFER_TP}" + policy.generation.mcore_generation_config.expert_model_parallel_size="${INFER_EP}" + policy.generation.mcore_generation_config.expert_tensor_parallel_size=1 + ++policy.generation.mcore_generation_config.context_parallel_size="${POLICY_CP}" + ++policy.generation.mcore_generation_config.mamba_inference_ssm_states_dtype=float32 + ++policy.generation.mcore_generation_config.mamba_inference_conv_states_dtype=float32 + ++policy.generation.mcore_generation_config.logprobs_mode=raw_logprobs + policy.generation.mcore_generation_config.enable_chunked_prefill="${MEGATRON_ENABLE_CHUNKED_PREFILL}" + ++policy.generation.mcore_generation_config.async_sched_mode=async + policy.generation.mcore_generation_config.cuda_graph_impl="${MEGATRON_CUDA_GRAPH_IMPL}" + policy.generation.mcore_generation_config.inference_cuda_graph_scope="${MEGATRON_CUDA_GRAPH_SCOPE}" + policy.generation.mcore_generation_config.num_cuda_graphs="${MEGATRON_NUM_CUDA_GRAPHS}" + policy.generation.mcore_generation_config.use_cuda_graphs_for_non_decode_steps="${MEGATRON_USE_CUDA_GRAPHS_FOR_NON_DECODE}" + policy.generation.mcore_generation_config.refit_backend="${REFIT_BACKEND}" + policy.generation.mcore_generation_config.buffer_size_gb="${BUFFER_SIZE_GB}" + policy.generation.mcore_generation_config.moe_pad_experts_for_cuda_graph_inference="${MOE_PAD_EXPERTS_FOR_CG}" + policy.generation.mcore_generation_config.max_model_len="${MAX_SEQUENCE_LENGTH}" + policy.generation.mcore_generation_config.max_tokens="${MAX_SEQUENCE_LENGTH}" + ) +else + # To fit vLLM generation on 1 node. Refit packs weights into chunks sized at + # NRL_REFIT_BUFFER_MEMORY_RATIO * total HBM; exported globally on purpose + # because producer and consumer must agree on the chunk boundaries. + export NRL_REFIT_BUFFER_MEMORY_RATIO="${NRL_REFIT_BUFFER_MEMORY_RATIO:-0.005}" + # Omni vLLM A/B path. Use EP=TP so async GRPO can use Ray DP. + VLLM_GPU_MEMORY_UTILIZATION="${VLLM_GPU_MEMORY_UTILIZATION:-0.5}" + VLLM_ENFORCE_EAGER="${VLLM_ENFORCE_EAGER:-true}" + # Per-step forward-pass token budget; mirrors mcore_generation_config.max_tokens. + VLLM_MAX_NUM_BATCHED_TOKENS="${VLLM_MAX_NUM_BATCHED_TOKENS:-${MAX_SEQUENCE_LENGTH}}" + GEN_OVERRIDES=( + ++policy.generation.stop_strings=null + ++policy.generation.bad_words=null + policy.generation.vllm_cfg.async_engine="${ASYNC_GRPO}" + policy.generation.vllm_cfg.skip_tokenizer_init=false + policy.generation.vllm_cfg.tensor_parallel_size="${INFER_TP}" + policy.generation.vllm_cfg.pipeline_parallel_size=1 + policy.generation.vllm_cfg.expert_parallel_size="${INFER_EP}" + policy.generation.vllm_cfg.max_model_len="${MAX_SEQUENCE_LENGTH}" + ++policy.generation.vllm_cfg.cap_max_tokens_to_context=true + policy.generation.vllm_cfg.gpu_memory_utilization="${VLLM_GPU_MEMORY_UTILIZATION}" + policy.generation.vllm_cfg.enforce_eager="${VLLM_ENFORCE_EAGER}" + ++policy.generation.vllm_cfg.enable_prefix_caching=false + policy.generation.vllm_cfg.logprobs_mode=raw_logprobs + ++policy.generation.vllm_kwargs.limit_mm_per_prompt.image=2 + ++policy.generation.vllm_kwargs.max_num_batched_tokens="${VLLM_MAX_NUM_BATCHED_TOKENS}" + ++policy.generation.vllm_kwargs.mamba_ssm_cache_dtype=float32 + ++policy.generation.vllm_kwargs.skip_mm_profiling=true + ++policy.generation.vllm_kwargs.kernel_config.enable_flashinfer_autotune=false + ++policy.generation.vllm_kwargs.kernel_config.moe_backend=triton + ) +fi + +echo "Launching ${JOB_NAME}: ${GPUS_PER_NODE} visible GPU(s)" +echo " generation backend: ${GENERATION_BACKEND} colocated=${COLOCATED_ENABLED} async=${ASYNC_GRPO}" +echo " training world size: ${TRAIN_WORLD_SIZE} (TP=${POLICY_TP}, EP=${POLICY_EP}, ETP=1)" +echo " inference world size: ${INFERENCE_WORLD_SIZE} (TP=${INFER_TP}, EP=${INFER_EP})" +echo " async: max_trajectory_age=${MAX_TRAJECTORY_AGE_STEPS} in_flight_weight_updates=${IN_FLIGHT_WEIGHT_UPDATES}" +echo " seq/new_tokens: ${MAX_SEQUENCE_LENGTH}/${MAX_NEW_TOKENS}" +echo " optimizer moments: exp_avg=${EXP_AVG_DTYPE:-} exp_avg_sq=${EXP_AVG_SQ_DTYPE:-} store_param_remainders=${STORE_PARAM_REMAINDERS:-}" +echo " W&B: ${WANDB_GROUP}/${WANDB_PROJ}/${WANDB_NAME} (enabled=${WANDB_ENABLED})" + +exec env "${NSYS_ENV[@]}" uv run --no-sync python examples/run_vlm_grpo.py \ + --config "${CONFIG}" \ + "${COMMON_OVERRIDES[@]}" \ + "${GEN_OVERRIDES[@]}" \ + "$@" diff --git a/scripts/run_nemotron_omni_vstat_megatron_1n4g.sh b/scripts/run_nemotron_omni_vstat_megatron_1n4g.sh new file mode 100755 index 00000000000..96914af0e00 --- /dev/null +++ b/scripts/run_nemotron_omni_vstat_megatron_1n4g.sh @@ -0,0 +1,245 @@ +#!/usr/bin/env bash +set -euo pipefail + +# One-node / four-GPU Nemotron Omni video-GRPO smoke test. +# VSTAT is recommended here because Hugging Face hosts both the MCQ annotations +# and the actual MP4 assets; Video-MME commonly requires a separate video fetch. +# +# ASYNC_GRPO=false selects synchronous GRPO. +# MEGATRON_CUDA_GRAPH_IMPL=local enables CUDA graphs; block scope and MoE +# expert padding are selected automatically unless explicitly overridden. + +NEMORL="${NEMORL:-/opt/nemo-rl}" +MODEL_NAME="${MODEL_NAME:-nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16}" +WORKSPACE_ROOT="${WORKSPACE_ROOT:-${NEMORL}/workspace}" +DATA_ROOT="${DATA_ROOT:-${WORKSPACE_ROOT}/datasets/vstat-smoke}" +HF_DATASET="${HF_DATASET:-ShushengYang/VSTAT}" +NUM_DATA_ROWS="${NUM_DATA_ROWS:-8}" +PREPARE_VSTAT="${PREPARE_VSTAT:-false}" +NUM_FRAMES="${NUM_FRAMES:-8}" +TEMPORAL_PATCH_SIZE="${TEMPORAL_PATCH_SIZE:-2}" +VIDEO_TARGET_PATCHES="${VIDEO_TARGET_PATCHES:-256}" +MAX_SEQUENCE_LENGTH="${MAX_SEQUENCE_LENGTH:-4096}" +MAX_NEW_TOKENS="${MAX_NEW_TOKENS:-1024}" +MIN_GENERATION_TOKENS="${MIN_GENERATION_TOKENS:-2000}" +VISION_EMBEDDING_CACHE_MAX_BYTES="${VISION_EMBEDDING_CACHE_MAX_BYTES:-536870912}" +ENABLE_THINKING="${ENABLE_THINKING:-true}" +MAX_STEPS="${MAX_STEPS:-4}" +NUM_PROMPTS_PER_STEP="${NUM_PROMPTS_PER_STEP:-2}" +NUM_GENERATIONS_PER_PROMPT="${NUM_GENERATIONS_PER_PROMPT:-4}" +TRAIN_GBS="${TRAIN_GBS:-$((NUM_PROMPTS_PER_STEP * NUM_GENERATIONS_PER_PROMPT))}" +GEN_GPUS="${GEN_GPUS:-2}" +INFER_EP="${INFER_EP:-${GEN_GPUS}}" +REFIT_BACKEND="${REFIT_BACKEND:-nccl}" +ASYNC_GRPO="${ASYNC_GRPO:-true}" +MAX_TRAJECTORY_AGE_STEPS="${MAX_TRAJECTORY_AGE_STEPS:-2}" +IN_FLIGHT_WEIGHT_UPDATES="${IN_FLIGHT_WEIGHT_UPDATES:-true}" +MEGATRON_CUDA_GRAPH_IMPL="${MEGATRON_CUDA_GRAPH_IMPL:-local}" +if [[ "${MEGATRON_CUDA_GRAPH_IMPL}" == "none" ]]; then + DEFAULT_CUDA_GRAPH_SCOPE=none +else + DEFAULT_CUDA_GRAPH_SCOPE=block +fi +if [[ "${MEGATRON_CUDA_GRAPH_IMPL}" != "none" && "${INFER_EP}" -gt 1 ]]; then + DEFAULT_MOE_PAD_EXPERTS_FOR_CG=true +else + DEFAULT_MOE_PAD_EXPERTS_FOR_CG=false +fi +MEGATRON_CUDA_GRAPH_SCOPE="${MEGATRON_CUDA_GRAPH_SCOPE:-${DEFAULT_CUDA_GRAPH_SCOPE}}" +MEGATRON_NUM_CUDA_GRAPHS="${MEGATRON_NUM_CUDA_GRAPHS:--1}" +MEGATRON_USE_CUDA_GRAPHS_FOR_NON_DECODE="${MEGATRON_USE_CUDA_GRAPHS_FOR_NON_DECODE:-false}" +MOE_PAD_EXPERTS_FOR_CG="${MOE_PAD_EXPERTS_FOR_CG:-${DEFAULT_MOE_PAD_EXPERTS_FOR_CG}}" +MEGATRON_ASYNC_SCHED_MODE="${MEGATRON_ASYNC_SCHED_MODE:-async}" +# Lower-precision Adam moments for the 2-GPU train half (HBM-tight on 1n4g). +# Override with EXP_AVG_DTYPE=float32 EXP_AVG_SQ_DTYPE=float32 STORE_PARAM_REMAINDERS=false +# for full-precision optimizer state. +EXP_AVG_DTYPE="${EXP_AVG_DTYPE:-bfloat16}" +EXP_AVG_SQ_DTYPE="${EXP_AVG_SQ_DTYPE:-bfloat16}" +STORE_PARAM_REMAINDERS="${STORE_PARAM_REMAINDERS:-true}" + +cd "${NEMORL}" + +GPUS_PER_NODE="${GPUS_PER_NODE:-$(nvidia-smi --query-gpu=index --format=csv,noheader | wc -l)}" +if (( GPUS_PER_NODE != 4 )); then + echo "This smoke launcher expects exactly four visible GPUs (got ${GPUS_PER_NODE})." >&2 + exit 1 +fi +if (( GEN_GPUS <= 0 || GEN_GPUS >= GPUS_PER_NODE )); then + echo "GEN_GPUS must leave at least one GPU for training." >&2 + exit 1 +fi +if (( INFER_EP <= 0 || GEN_GPUS % INFER_EP != 0 )); then + echo "INFER_EP must be a positive divisor of GEN_GPUS (got ${INFER_EP})." >&2 + exit 1 +fi +if [[ "${ASYNC_GRPO}" != "true" && "${ASYNC_GRPO}" != "false" ]]; then + echo "ASYNC_GRPO must be true or false (got ${ASYNC_GRPO})." >&2 + exit 1 +fi +if [[ "${MEGATRON_CUDA_GRAPH_IMPL}" != "none" && + "${MEGATRON_CUDA_GRAPH_SCOPE}" == "none" ]]; then + echo "CUDA graphs require a non-none MEGATRON_CUDA_GRAPH_SCOPE." >&2 + exit 1 +fi +if [[ "${MEGATRON_CUDA_GRAPH_IMPL}" != "none" && "${INFER_EP}" -gt 1 && + "${MOE_PAD_EXPERTS_FOR_CG}" != "true" ]]; then + echo "CUDA graphs with expert parallelism require MOE_PAD_EXPERTS_FOR_CG=true." >&2 + exit 1 +fi +TRAIN_GPUS=$((GPUS_PER_NODE - GEN_GPUS)) + +CACHE_ROOT="${CACHE_ROOT:-${WORKSPACE_ROOT}/cache/nemo-rl-omni}" +export HF_HOME="${HF_HOME:-${CACHE_ROOT}/huggingface}" +export HF_HUB_CACHE="${HF_HUB_CACHE:-${HF_HOME}/hub}" +export HF_DATASETS_CACHE="${HF_DATASETS_CACHE:-${HF_HOME}/datasets}" +export NRL_MEGATRON_CHECKPOINT_DIR="${NRL_MEGATRON_CHECKPOINT_DIR:-${CACHE_ROOT}/megatron-checkpoints}" +export NEMO_RL_VENV_DIR="${NEMO_RL_VENV_DIR:-/opt/ray_venvs}" +export NEMO_GYM_VENV_DIR="${NEMO_GYM_VENV_DIR:-/opt/ray_venvs}" +export NEMO_GYM_EXTRA_ROOTS="${NEMO_GYM_EXTRA_ROOTS:-${NEMORL}/3rdparty/Gym-workspace/Gym}" +export NRL_VIDEO_BACKEND="${NRL_VIDEO_BACKEND:-torchcodec}" +export NRL_VIDEO_SAMPLING_STYLE="${NRL_VIDEO_SAMPLING_STYLE:-nemotron_vl}" +export NRL_VIDEO_TEMPORAL_PATCH_SIZE="${TEMPORAL_PATCH_SIZE}" +export RAY_ENABLE_UV_RUN_RUNTIME_ENV=0 +export NRL_VENVS_TRUST_EXISTING="${NRL_VENVS_TRUST_EXISTING:-1}" +export CUDA_DEVICE_MAX_CONNECTIONS="${CUDA_DEVICE_MAX_CONNECTIONS:-1}" +export FLASHINFER_DISABLE_VERSION_CHECK="${FLASHINFER_DISABLE_VERSION_CHECK:-1}" + +AUDIO_DEPS_SCRIPT="${NEMORL}/tools/install_audio_deps.sh" +MEGATRON_WORKER_PYTHON="${RAY_MEGATRON_PYTHON:-${NEMO_RL_VENV_DIR}/nemo_rl.models.policy.workers.megatron_policy_worker.MegatronPolicyWorker/bin/python}" +NEED_AUDIO_VIDEO_DEPS=false +if ! python -c "import torchcodec" >/dev/null 2>&1; then + NEED_AUDIO_VIDEO_DEPS=true +fi +if [[ ! -x "${MEGATRON_WORKER_PYTHON}" ]] || + ! "${MEGATRON_WORKER_PYTHON}" -c "import av" >/dev/null 2>&1; then + NEED_AUDIO_VIDEO_DEPS=true +fi +if [[ "${NEED_AUDIO_VIDEO_DEPS}" == "true" ]]; then + if [[ ! -f "${AUDIO_DEPS_SCRIPT}" ]]; then + echo "Audio/video dependency installer is missing: ${AUDIO_DEPS_SCRIPT}" >&2 + exit 1 + fi + echo "Installing missing audio/video dependencies" + RAY_MEGATRON_PYTHON="${MEGATRON_WORKER_PYTHON}" bash "${AUDIO_DEPS_SCRIPT}" +fi + +# NSYS: ENABLE_NSYS=true NRL_NSYS_PROFILE_STEP_RANGE=1:4 +ENABLE_NSYS="${ENABLE_NSYS:-false}" +NSYS_ENV=() +if [[ "${ENABLE_NSYS}" == "true" ]]; then + NRL_NSYS_WORKER_PATTERNS="${NRL_NSYS_WORKER_PATTERNS:-*policy*,*megatron*}" + NRL_NSYS_PROFILE_STEP_RANGE="${NRL_NSYS_PROFILE_STEP_RANGE:-1:4}" + LD_LIBRARY_PATH="/usr/local/cuda/targets/aarch64-linux/lib:/usr/local/cuda/targets/x86_64-linux/lib:/usr/local/cuda/lib64:/usr/local/cuda/lib:/usr/local/nvidia/lib64:/usr/local/nvidia/lib:/usr/lib/aarch64-linux-gnu:/usr/lib/x86_64-linux-gnu:${LD_LIBRARY_PATH:-}" + NRL_NSYS_EXTRA_OPTIONS="${NRL_NSYS_EXTRA_OPTIONS:-{\"o\":\"${WORKSPACE_ROOT}/nsys/%p\",\"cpuctxsw\":\"none\",\"force-overwrite\":\"true\"}}" + NSYS_ENV=( + "NRL_NSYS_WORKER_PATTERNS=${NRL_NSYS_WORKER_PATTERNS}" + "NRL_NSYS_PROFILE_STEP_RANGE=${NRL_NSYS_PROFILE_STEP_RANGE}" + "LD_LIBRARY_PATH=${LD_LIBRARY_PATH}" + "NRL_NSYS_EXTRA_OPTIONS=${NRL_NSYS_EXTRA_OPTIONS}" + ) + mkdir -p "${WORKSPACE_ROOT}/nsys" +fi + +BRIDGE="${NEMORL}/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge" +export PYTHONPATH="${NEMORL}:${NEMO_GYM_EXTRA_ROOTS}:${BRIDGE}/src:${BRIDGE}/3rdparty/Megatron-LM${PYTHONPATH:+:${PYTHONPATH}}" + +mkdir -p "${DATA_ROOT}" "${HF_HOME}" "${NRL_MEGATRON_CHECKPOINT_DIR}" +export NEMO_RL_VIDEO_TRAIN_JSONL="${DATA_ROOT}/train-gym.jsonl" +export NEMO_RL_VIDEO_VAL_JSONL="${DATA_ROOT}/val-gym.jsonl" +export NEMO_RL_VIDEO_MEDIA_ROOT="${DATA_ROOT}/media" + +VSTAT_PREPARE_SCRIPT="${NEMORL}/scripts/prepare_nemotron_omni_vstat.py" +if [[ ! -f "${VSTAT_PREPARE_SCRIPT}" ]]; then + echo "VSTAT preparation script is missing: ${VSTAT_PREPARE_SCRIPT}" >&2 + exit 1 +fi +if [[ "${PREPARE_VSTAT}" == "true" || + ! -s "${NEMO_RL_VIDEO_TRAIN_JSONL}" || + ! -s "${NEMO_RL_VIDEO_VAL_JSONL}" ]]; then + echo "Preparing VSTAT under ${DATA_ROOT}" + uv run --no-sync python "${VSTAT_PREPARE_SCRIPT}" \ + --output-dir "${DATA_ROOT}" \ + --repo-id "${HF_DATASET}" \ + --num-rows "${NUM_DATA_ROWS}" +fi + +RESULTS_DIR="${RESULTS_DIR:-${WORKSPACE_ROOT}/results/nemo-rl-omni/nemotron-omni-vstat-megatron-1n4g}" +mkdir -p "${RESULTS_DIR}" + +echo "VSTAT: root=${DATA_ROOT} repo=${HF_DATASET} rows=${NUM_DATA_ROWS} prepare=${PREPARE_VSTAT}" +echo " async_grpo=${ASYNC_GRPO} scheduler=${MEGATRON_ASYNC_SCHED_MODE}" +echo " generation: TP=${GEN_GPUS} EP=${INFER_EP}" +echo " cuda_graph_impl=${MEGATRON_CUDA_GRAPH_IMPL} scope=${MEGATRON_CUDA_GRAPH_SCOPE} non_decode=${MEGATRON_USE_CUDA_GRAPHS_FOR_NON_DECODE}" +exec env "${NSYS_ENV[@]}" 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 \ + policy.model_name="${MODEL_NAME}" \ + policy.tokenizer.name="${MODEL_NAME}" \ + policy.tokenizer.chat_template_kwargs.enable_thinking="${ENABLE_THINKING}" \ + policy.is_vlm=true \ + policy.generation.backend=megatron \ + ++policy.generation.bad_words=null \ + policy.generation.colocated.enabled=false \ + policy.generation.colocated.resources.num_nodes=1 \ + policy.generation.colocated.resources.gpus_per_node="${GEN_GPUS}" \ + policy.megatron_cfg.tensor_model_parallel_size="${TRAIN_GPUS}" \ + policy.megatron_cfg.expert_model_parallel_size="${TRAIN_GPUS}" \ + policy.megatron_cfg.expert_tensor_parallel_size=1 \ + policy.megatron_cfg.context_parallel_size=1 \ + policy.megatron_cfg.sequence_parallel=true \ + policy.megatron_cfg.bias_activation_fusion=false \ + policy.megatron_cfg.optimizer.optimizer_cpu_offload=false \ + policy.megatron_cfg.optimizer.optimizer_offload_fraction=0.0 \ + ++policy.megatron_cfg.optimizer.exp_avg_dtype="${EXP_AVG_DTYPE}" \ + ++policy.megatron_cfg.optimizer.exp_avg_sq_dtype="${EXP_AVG_SQ_DTYPE}" \ + ++policy.megatron_cfg.optimizer.store_param_remainders="${STORE_PARAM_REMAINDERS}" \ + policy.generation.mcore_generation_config.tensor_model_parallel_size="${GEN_GPUS}" \ + policy.generation.mcore_generation_config.expert_model_parallel_size="${INFER_EP}" \ + policy.generation.mcore_generation_config.expert_tensor_parallel_size=1 \ + ++policy.generation.mcore_generation_config.context_parallel_size=1 \ + policy.generation.mcore_generation_config.transformer_impl=transformer_engine \ + policy.generation.mcore_generation_config.sequence_parallel=true \ + policy.generation.mcore_generation_config.refit_backend="${REFIT_BACKEND}" \ + policy.generation.mcore_generation_config.buffer_size_gb=8 \ + policy.generation.mcore_generation_config.cuda_graph_impl="${MEGATRON_CUDA_GRAPH_IMPL}" \ + policy.generation.mcore_generation_config.inference_cuda_graph_scope="${MEGATRON_CUDA_GRAPH_SCOPE}" \ + policy.generation.mcore_generation_config.num_cuda_graphs="${MEGATRON_NUM_CUDA_GRAPHS}" \ + ++policy.generation.mcore_generation_config.moe_pad_experts_for_cuda_graph_inference="${MOE_PAD_EXPERTS_FOR_CG}" \ + policy.generation.mcore_generation_config.use_cuda_graphs_for_non_decode_steps="${MEGATRON_USE_CUDA_GRAPHS_FOR_NON_DECODE}" \ + policy.generation.mcore_generation_config.max_model_len="${MAX_SEQUENCE_LENGTH}" \ + policy.generation.mcore_generation_config.max_tokens="${MAX_SEQUENCE_LENGTH}" \ + policy.generation.mcore_generation_config.enable_chunked_prefill=true \ + ++policy.generation.mcore_generation_config.async_sched_mode="${MEGATRON_ASYNC_SCHED_MODE}" \ + policy.generation.mcore_generation_config.enable_prefix_caching=false \ + ++policy.generation.mcore_generation_config.vision_embedding_cache_max_bytes="${VISION_EMBEDDING_CACHE_MAX_BYTES}" \ + policy.generation.mcore_generation_config.parsers=[nemotron-v3-reasoning,qwen3-coder-tool] \ + ++policy.generation.mcore_generation_config.video_num_frames="${NUM_FRAMES}" \ + ++policy.generation.mcore_generation_config.video_temporal_patch_size="${TEMPORAL_PATCH_SIZE}" \ + ++policy.generation.mcore_generation_config.video_target_num_patches="${VIDEO_TARGET_PATCHES}" \ + policy.max_total_sequence_length="${MAX_SEQUENCE_LENGTH}" \ + policy.generation.max_new_tokens="${MAX_NEW_TOKENS}" \ + +data.default.num_frames="${NUM_FRAMES}" \ + +data.default.video_sampling_style=nemotron_vl \ + +data.default.video_temporal_patch_size="${TEMPORAL_PATCH_SIZE}" \ + +data.default.min_generation_tokens="${MIN_GENERATION_TOKENS}" \ + data.default.video_target_num_patches="${VIDEO_TARGET_PATCHES}" \ + data.train.data_path="${NEMO_RL_VIDEO_TRAIN_JSONL}" \ + data.validation.data_path="${NEMO_RL_VIDEO_VAL_JSONL}" \ + ++env.nemo_gym.policy_model.responses_api_models.vllm_model.chat_template_kwargs.enable_thinking="${ENABLE_THINKING}" \ + grpo.deduplicate_multimodal_data=false \ + grpo.async_grpo.enabled="${ASYNC_GRPO}" \ + grpo.async_grpo.max_trajectory_age_steps="${MAX_TRAJECTORY_AGE_STEPS}" \ + grpo.async_grpo.in_flight_weight_updates="${IN_FLIGHT_WEIGHT_UPDATES}" \ + loss_fn.use_importance_sampling_correction=true \ + grpo.num_prompts_per_step="${NUM_PROMPTS_PER_STEP}" \ + grpo.num_generations_per_prompt="${NUM_GENERATIONS_PER_PROMPT}" \ + policy.train_global_batch_size="${TRAIN_GBS}" \ + grpo.max_num_steps="${MAX_STEPS}" \ + grpo.val_period=0 \ + grpo.val_at_start=false \ + grpo.val_at_end=false \ + cluster.num_nodes=1 \ + cluster.gpus_per_node="${GPUS_PER_NODE}" \ + checkpointing.enabled=false \ + logger.log_dir="${RESULTS_DIR}" \ + logger.wandb_enabled=false \ + "$@" diff --git a/scripts/submit_nemotron_omni_clevr_megatron_8n4g.sh b/scripts/submit_nemotron_omni_clevr_megatron_8n4g.sh new file mode 100755 index 00000000000..ea370f90e29 --- /dev/null +++ b/scripts/submit_nemotron_omni_clevr_megatron_8n4g.sh @@ -0,0 +1,422 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Submit the eight-node Omni CLEVR Megatron GRPO recipe through NeMo-RL's +# Ray/Slurm launcher. +# +# Default: non-colocated gen-heavy (6 gen + 2 train). Async colocated is +# supported on this branch (NVIDIA-NeMo/RL#2884); set COLOCATED=true to share +# all GPUs for train+gen. ASYNC_GRPO=true is required for colocated async. +# +# GENERATION_BACKEND=vllm flips inference to Omni vLLM for A/B vs Megatron. +# The A/B defaults disable stop/bad-word filters; native multimodal inputs are +# supported by both synchronous and non-colocated asynchronous rollouts. + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +NEMORL="${NEMORL:-$(cd "${SCRIPT_DIR}/.." && pwd)}" +CONTAINER_NEMORL="${CONTAINER_NEMORL:-/opt/nemo-rl}" +WORKSPACE_ROOT="${WORKSPACE_ROOT:-${NEMORL}/workspace}" + +CONTAINER="${CONTAINER:-/lustre/fsw/portfolios/coreai/users/cye/enroot/nemo_rl_0807.sqsh}" +MODEL_NAME="${MODEL_NAME:-nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16}" +CONFIG="${CONFIG:-examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron_generation.v1.yaml}" +GENERATION_BACKEND="${GENERATION_BACKEND:-megatron}" +COLOCATED="${COLOCATED:-false}" +ASYNC_GRPO="${ASYNC_GRPO:-true}" + +if [[ "${GENERATION_BACKEND}" != "megatron" && "${GENERATION_BACKEND}" != "vllm" ]]; then + echo "GENERATION_BACKEND must be megatron or vllm (got ${GENERATION_BACKEND})." >&2 + exit 1 +fi +if [[ "${COLOCATED}" == "true" && "${ASYNC_GRPO}" == "true" && "${GENERATION_BACKEND}" != "megatron" ]]; then + echo "Colocated async GRPO requires GENERATION_BACKEND=megatron." >&2 + exit 1 +fi + +CACHE_ROOT="${CACHE_ROOT:-${WORKSPACE_ROOT}/cache/nemo-rl-omni}" +export HF_HOME="${HF_HOME:-${CACHE_ROOT}/huggingface}" +export HUGGINGFACE_HUB_CACHE="${HUGGINGFACE_HUB_CACHE:-${HF_HOME}/hub}" +export HF_HUB_CACHE="${HF_HUB_CACHE:-${HUGGINGFACE_HUB_CACHE}}" +export HF_DATASETS_CACHE="${HF_DATASETS_CACHE:-${HF_HOME}/datasets}" +export HF_MODULES_CACHE="${HF_MODULES_CACHE:-${HF_HOME}/modules}" +export TRANSFORMERS_CACHE="${TRANSFORMERS_CACHE:-${HF_HOME}/transformers}" +export NRL_MEGATRON_CHECKPOINT_DIR="${NRL_MEGATRON_CHECKPOINT_DIR:-${CACHE_ROOT}/megatron-checkpoints}" +export XDG_CACHE_HOME="${XDG_CACHE_HOME:-${CACHE_ROOT}/xdg}" +export TORCH_HOME="${TORCH_HOME:-${CACHE_ROOT}/torch}" +export TRITON_CACHE_DIR="${TRITON_CACHE_DIR:-${CACHE_ROOT}/triton}" + +NUM_NODES="${NUM_NODES:-8}" +GPUS_PER_NODE="${GPUS_PER_NODE:-4}" + +if [[ "${COLOCATED}" == "true" ]]; then + # Shared train+gen on every GPU (NVIDIA-NeMo/RL#2884). + TRAIN_WORLD_SIZE=$((NUM_NODES * GPUS_PER_NODE)) + INFERENCE_WORLD_SIZE="${TRAIN_WORLD_SIZE}" + GEN_GPUS_PER_NODE="${GPUS_PER_NODE}" + NUM_GEN_NODES="${NUM_NODES}" + COLOCATED_ENABLED=true +else + NUM_GEN_NODES="${NUM_GEN_NODES:-6}" + if (( NUM_NODES == 1 )); then + if (( NUM_GEN_NODES != 1 )); then + echo "One-node non-colocated mode requires NUM_GEN_NODES=1." >&2 + exit 1 + fi + GEN_GPUS_PER_NODE="${GEN_GPUS_PER_NODE:-$((GPUS_PER_NODE / 2))}" + if (( GPUS_PER_NODE < 4 || GEN_GPUS_PER_NODE <= 0 || GEN_GPUS_PER_NODE >= GPUS_PER_NODE )); then + echo "One-node non-colocated mode requires at least four GPUs and a strict train/inference split." >&2 + exit 1 + fi + TRAIN_WORLD_SIZE=$((GPUS_PER_NODE - GEN_GPUS_PER_NODE)) + else + GEN_GPUS_PER_NODE="${GEN_GPUS_PER_NODE:-${GPUS_PER_NODE}}" + if (( NUM_GEN_NODES <= 0 || NUM_GEN_NODES >= NUM_NODES )); then + echo "Multi-node non-colocated mode requires 0 < NUM_GEN_NODES < NUM_NODES." >&2 + exit 1 + fi + if (( GEN_GPUS_PER_NODE != GPUS_PER_NODE )); then + echo "Multi-node non-colocated inference must reserve complete GPU nodes." >&2 + exit 1 + fi + TRAIN_WORLD_SIZE=$(((NUM_NODES - NUM_GEN_NODES) * GPUS_PER_NODE)) + fi + INFERENCE_WORLD_SIZE=$((NUM_GEN_NODES * GEN_GPUS_PER_NODE)) + COLOCATED_ENABLED=false +fi + +# Prefer TP=8; fall back by powers of two for smaller worlds. +DEFAULT_POLICY_TP=8 +while (( DEFAULT_POLICY_TP > TRAIN_WORLD_SIZE || TRAIN_WORLD_SIZE % DEFAULT_POLICY_TP != 0 )); do + DEFAULT_POLICY_TP=$((DEFAULT_POLICY_TP / 2)) +done +DEFAULT_INFER_TP=8 +while (( DEFAULT_INFER_TP > INFERENCE_WORLD_SIZE || INFERENCE_WORLD_SIZE % DEFAULT_INFER_TP != 0 )); do + DEFAULT_INFER_TP=$((DEFAULT_INFER_TP / 2)) +done +POLICY_TP="${POLICY_TP:-${DEFAULT_POLICY_TP}}" +INFER_TP="${INFER_TP:-${DEFAULT_INFER_TP}}" +POLICY_CP="${POLICY_CP:-1}" + +if [[ "${GENERATION_BACKEND}" == "megatron" && "${POLICY_CP}" != "1" ]]; then + echo "Megatron dynamic inference requires POLICY_CP=1." >&2 + exit 1 +fi +if (( TRAIN_WORLD_SIZE % (POLICY_TP * POLICY_CP) != 0 )); then + echo "Training world size must be divisible by POLICY_TP * POLICY_CP." >&2 + exit 1 +fi +if (( INFERENCE_WORLD_SIZE % INFER_TP != 0 )); then + echo "Inference world size must be divisible by INFER_TP." >&2 + exit 1 +fi +TRAIN_DP_SIZE=$((TRAIN_WORLD_SIZE / (POLICY_TP * POLICY_CP))) +INFERENCE_DP_SIZE=$((INFERENCE_WORLD_SIZE / INFER_TP)) + +# With ETP=1, world_size must be divisible by EP (not by TP*EP). +largest_ep() { + local world="$1" + local ep=8 + while (( ep > world || world % ep != 0 )); do + ep=$((ep / 2)) + done + printf '%d' "${ep}" +} + +POLICY_EP="${POLICY_EP:-$(largest_ep "${TRAIN_WORLD_SIZE}")}" +INFER_EP="${INFER_EP:-$(largest_ep "${INFERENCE_WORLD_SIZE}")}" +if (( TRAIN_WORLD_SIZE % POLICY_EP != 0 )); then + echo "Training world size must be divisible by POLICY_EP (ETP=1)." >&2 + exit 1 +fi +if [[ "${GENERATION_BACKEND}" == "vllm" ]]; then + # vLLM EP = DP_internal * TP. Async GRPO cannot use internal DP, so EP must + # equal TP; Ray DP then matches Megatron's world/TP replica count. + if [[ "${INFER_EP}" != "${INFER_TP}" ]]; then + echo "Forcing INFER_EP=${INFER_TP} for vLLM (was ${INFER_EP}) for async-safe Ray DP." >&2 + fi + INFER_EP="${INFER_TP}" +elif (( INFERENCE_WORLD_SIZE % INFER_EP != 0 )); then + echo "Inference world size must be divisible by INFER_EP (ETP=1)." >&2 + exit 1 +fi + +MAX_STEPS="${MAX_STEPS:-1000000}" +MAX_SEQUENCE_LENGTH="${MAX_SEQUENCE_LENGTH:-4096}" +MAX_NEW_TOKENS="${MAX_NEW_TOKENS:-2048}" +# Two prompt groups per inference DP replica (non-colocated default: 6 gen +# nodes → DP=3 → prompts=6; colocated 8n4g → DP=4 → prompts=8). +NUM_PROMPTS="${NUM_PROMPTS:-$((INFERENCE_DP_SIZE * 2))}" +NUM_GENERATIONS="${NUM_GENERATIONS:-8}" +TRAIN_GBS="${TRAIN_GBS:-$((NUM_PROMPTS * NUM_GENERATIONS))}" +EXPECTED_TRAIN_GBS=$((NUM_PROMPTS * NUM_GENERATIONS)) +VAL_GBS="${VAL_GBS:-64}" +VAL_SIZE="${VAL_SIZE:-64}" +if (( TRAIN_GBS != EXPECTED_TRAIN_GBS )); then + echo "TRAIN_GBS (${TRAIN_GBS}) must equal NUM_PROMPTS * NUM_GENERATIONS (${EXPECTED_TRAIN_GBS})." >&2 + exit 1 +fi +if (( TRAIN_GBS % TRAIN_DP_SIZE != 0 )); then + echo "TRAIN_GBS (${TRAIN_GBS}) must be divisible by training DP size (${TRAIN_DP_SIZE})." >&2 + exit 1 +fi +if (( VAL_GBS % TRAIN_DP_SIZE != 0 )); then + echo "VAL_GBS (${VAL_GBS}) must be divisible by training DP size (${TRAIN_DP_SIZE})." >&2 + exit 1 +fi +REFIT_BACKEND="${REFIT_BACKEND:-nccl}" +OPTIMIZER_CPU_OFFLOAD="${OPTIMIZER_CPU_OFFLOAD:-false}" +BUFFER_SIZE_GB="${BUFFER_SIZE_GB:-8}" +OFFLOAD_OPTIMIZER_FOR_LOGPROB="${OFFLOAD_OPTIMIZER_FOR_LOGPROB:-false}" +if [[ "${OPTIMIZER_CPU_OFFLOAD}" == "true" ]]; then + OPTIMIZER_OFFLOAD_FRACTION="${OPTIMIZER_OFFLOAD_FRACTION:-1.0}" +else + OPTIMIZER_OFFLOAD_FRACTION="${OPTIMIZER_OFFLOAD_FRACTION:-0.0}" +fi +# Lower-precision Adam moments; set to empty to skip the override. +EXP_AVG_DTYPE="${EXP_AVG_DTYPE:-bfloat16}" +EXP_AVG_SQ_DTYPE="${EXP_AVG_SQ_DTYPE:-bfloat16}" +STORE_PARAM_REMAINDERS="${STORE_PARAM_REMAINDERS:-true}" +# Megatron defaults for stable inference/training parity. +MEGATRON_ENABLE_CHUNKED_PREFILL="${MEGATRON_ENABLE_CHUNKED_PREFILL:-true}" +MEGATRON_CUDA_GRAPH_IMPL="${MEGATRON_CUDA_GRAPH_IMPL:-local}" +MEGATRON_CUDA_GRAPH_SCOPE="${MEGATRON_CUDA_GRAPH_SCOPE:-block}" +MEGATRON_NUM_CUDA_GRAPHS="${MEGATRON_NUM_CUDA_GRAPHS:--1}" +MEGATRON_USE_CUDA_GRAPHS_FOR_NON_DECODE="${MEGATRON_USE_CUDA_GRAPHS_FOR_NON_DECODE:-false}" +MOE_PAD_EXPERTS_FOR_CG="${MOE_PAD_EXPERTS_FOR_CG:-true}" +EXTRA_OVERRIDES="${EXTRA_OVERRIDES:-}" +MAX_TRAJECTORY_AGE_STEPS="${MAX_TRAJECTORY_AGE_STEPS:-2}" +IN_FLIGHT_WEIGHT_UPDATES="${IN_FLIGHT_WEIGHT_UPDATES:-true}" + +JOB_NAME="${JOB_NAME:-nemotron-omni-clevr-${GENERATION_BACKEND}-8n4g}" +EXP_NAME="${EXP_NAME:-${JOB_NAME}}" +PRECISION_RECIPE="${PRECISION_RECIPE:-bf16}" +WANDB_ENABLED="${WANDB_ENABLED:-true}" +WANDB_PROJ="${WANDB_PROJ:-mllm-rl-dev}" +WANDB_GROUP="${WANDB_GROUP:-adlr}" +WANDB_NAME="${WANDB_NAME:-${EXP_NAME}-${PRECISION_RECIPE}-internal-repo}" +RESULTS_DIR="${RESULTS_DIR:-${WORKSPACE_ROOT}/results/nemo-rl-omni/${JOB_NAME}}" +CHECKPOINTING_ENABLED="${CHECKPOINTING_ENABLED:-false}" +SLURM_LOG_DIR="${SLURM_LOG_DIR:-${RESULTS_DIR}/slurm}" + +# OCI-HSG Slurm defaults. Override any of these for a reservation or a +# different account/partition. +SBATCH_ACCOUNT="${SBATCH_ACCOUNT:-nemotron_sw_post}" # coreai_dlalgo_mcore nemotron_sw_post +SBATCH_PARTITION="${SBATCH_PARTITION:-batch_long}" +SBATCH_QOS="${SBATCH_QOS:-}" +SBATCH_TIME="${SBATCH_TIME:-04:00:00}" +SBATCH_RESERVATION="${SBATCH_RESERVATION:-}" +SBATCH_SEGMENT="${NUM_NODES}" + +mkdir -p \ + "${HF_HUB_CACHE}" \ + "${HF_DATASETS_CACHE}" \ + "${HF_MODULES_CACHE}" \ + "${TRANSFORMERS_CACHE}" \ + "${NRL_MEGATRON_CHECKPOINT_DIR}" \ + "${XDG_CACHE_HOME}" \ + "${TORCH_HOME}" \ + "${TRITON_CACHE_DIR}" \ + "${SLURM_LOG_DIR}" + +if [[ ! -f "${CONTAINER}" ]]; then + echo "Container image does not exist: ${CONTAINER}" >&2 + exit 1 +fi +if [[ ! -f "${NEMORL}/ray.sub" || ! -f "${NEMORL}/${CONFIG}" ]]; then + echo "NeMo-RL launcher or config is missing under: ${NEMORL}" >&2 + exit 1 +fi + +OPTIMIZER_DTYPE_OVERRIDES="" +if [[ -n "${EXP_AVG_DTYPE}" ]]; then + OPTIMIZER_DTYPE_OVERRIDES+=" ++policy.megatron_cfg.optimizer.exp_avg_dtype=${EXP_AVG_DTYPE}" +fi +if [[ -n "${EXP_AVG_SQ_DTYPE}" ]]; then + OPTIMIZER_DTYPE_OVERRIDES+=" ++policy.megatron_cfg.optimizer.exp_avg_sq_dtype=${EXP_AVG_SQ_DTYPE}" +fi +if [[ -n "${STORE_PARAM_REMAINDERS}" ]]; then + OPTIMIZER_DTYPE_OVERRIDES+=" ++policy.megatron_cfg.optimizer.store_param_remainders=${STORE_PARAM_REMAINDERS}" +fi + +REFIT_ENV_EXPORTS="" +REFIT_BUFFER_MEMORY_RATIO="" + +if [[ "${GENERATION_BACKEND}" == "megatron" ]]; then + GEN_OVERRIDES="\ +++policy.generation.stop_strings=null \ +++policy.generation.bad_words=null \ +policy.generation.mcore_generation_config.tensor_model_parallel_size=${INFER_TP} \ +policy.generation.mcore_generation_config.expert_model_parallel_size=${INFER_EP} \ +policy.generation.mcore_generation_config.expert_tensor_parallel_size=1 \ +++policy.generation.mcore_generation_config.context_parallel_size=${POLICY_CP} \ +++policy.generation.mcore_generation_config.mamba_inference_ssm_states_dtype=float32 \ +++policy.generation.mcore_generation_config.mamba_inference_conv_states_dtype=float32 \ +++policy.generation.mcore_generation_config.logprobs_mode=raw_logprobs \ +policy.generation.mcore_generation_config.enable_chunked_prefill=${MEGATRON_ENABLE_CHUNKED_PREFILL} \ +++policy.generation.mcore_generation_config.async_sched_mode=async \ +policy.generation.mcore_generation_config.cuda_graph_impl=${MEGATRON_CUDA_GRAPH_IMPL} \ +policy.generation.mcore_generation_config.inference_cuda_graph_scope=${MEGATRON_CUDA_GRAPH_SCOPE} \ +policy.generation.mcore_generation_config.num_cuda_graphs=${MEGATRON_NUM_CUDA_GRAPHS} \ +policy.generation.mcore_generation_config.use_cuda_graphs_for_non_decode_steps=${MEGATRON_USE_CUDA_GRAPHS_FOR_NON_DECODE} \ +policy.generation.mcore_generation_config.refit_backend=${REFIT_BACKEND} \ +policy.generation.mcore_generation_config.buffer_size_gb=${BUFFER_SIZE_GB} \ +policy.generation.mcore_generation_config.moe_pad_experts_for_cuda_graph_inference=${MOE_PAD_EXPERTS_FOR_CG} \ +policy.generation.mcore_generation_config.max_model_len=${MAX_SEQUENCE_LENGTH} \ +policy.generation.mcore_generation_config.max_tokens=${MAX_SEQUENCE_LENGTH}" +else + # Non-colocated refit NCCL-broadcasts weights in packed chunks sized at + # NRL_REFIT_BUFFER_MEMORY_RATIO * total HBM (the 0.02 default is 3.7GiB on + # GB200), which the ~170GiB-deep trainer cannot spare. Exported globally on + # purpose: producer and consumer both derive chunk boundaries from it and must + # agree. + REFIT_BUFFER_MEMORY_RATIO="${NRL_REFIT_BUFFER_MEMORY_RATIO:-0.005}" + REFIT_ENV_EXPORTS="export NRL_REFIT_BUFFER_MEMORY_RATIO=${REFIT_BUFFER_MEMORY_RATIO}" + + # Omni vLLM A/B path. + VLLM_GPU_MEMORY_UTILIZATION="${VLLM_GPU_MEMORY_UTILIZATION:-0.5}" + VLLM_ENFORCE_EAGER="${VLLM_ENFORCE_EAGER:-true}" + # Per-step forward-pass token budget; mirrors mcore_generation_config.max_tokens. + VLLM_MAX_NUM_BATCHED_TOKENS="${VLLM_MAX_NUM_BATCHED_TOKENS:-${MAX_SEQUENCE_LENGTH}}" + GEN_OVERRIDES="\ +++policy.generation.stop_strings=null \ +++policy.generation.bad_words=null \ +policy.generation.vllm_cfg.async_engine=${ASYNC_GRPO} \ +policy.generation.vllm_cfg.skip_tokenizer_init=false \ +policy.generation.vllm_cfg.tensor_parallel_size=${INFER_TP} \ +policy.generation.vllm_cfg.pipeline_parallel_size=1 \ +policy.generation.vllm_cfg.expert_parallel_size=${INFER_EP} \ +policy.generation.vllm_cfg.max_model_len=${MAX_SEQUENCE_LENGTH} \ +++policy.generation.vllm_cfg.cap_max_tokens_to_context=true \ +policy.generation.vllm_cfg.gpu_memory_utilization=${VLLM_GPU_MEMORY_UTILIZATION} \ +policy.generation.vllm_cfg.enforce_eager=${VLLM_ENFORCE_EAGER} \ +++policy.generation.vllm_cfg.enable_prefix_caching=false \ +policy.generation.vllm_cfg.logprobs_mode=raw_logprobs \ +++policy.generation.vllm_kwargs.limit_mm_per_prompt.image=2 \ +++policy.generation.vllm_kwargs.max_num_batched_tokens=${VLLM_MAX_NUM_BATCHED_TOKENS} \ +++policy.generation.vllm_kwargs.mamba_ssm_cache_dtype=float32 \ +++policy.generation.vllm_kwargs.skip_mm_profiling=true \ +++policy.generation.vllm_kwargs.kernel_config.enable_flashinfer_autotune=false \ +++policy.generation.vllm_kwargs.kernel_config.moe_backend=triton" +fi + +export NUM_NODES GPUS_PER_NODE CONTAINER +export RAY_ENABLE_UV_RUN_RUNTIME_ENV=0 +export NRL_VENVS_TRUST_EXISTING="${NRL_VENVS_TRUST_EXISTING:-1}" +export NRL_FORCE_REBUILD_VENVS="${NRL_FORCE_REBUILD_VENVS:-false}" +export NEMO_RL_VENV_DIR="${NEMO_RL_VENV_DIR:-/opt/ray_venvs}" +export CUDA_DEVICE_MAX_CONNECTIONS="${CUDA_DEVICE_MAX_CONNECTIONS:-1}" +export FLASHINFER_DISABLE_VERSION_CHECK="${FLASHINFER_DISABLE_VERSION_CHECK:-1}" +export NVTE_FWD_LAYERNORM_SM_MARGIN="${NVTE_FWD_LAYERNORM_SM_MARGIN:-16}" +export NVTE_BWD_LAYERNORM_SM_MARGIN="${NVTE_BWD_LAYERNORM_SM_MARGIN:-16}" +export NCCL_DEBUG="${NCCL_DEBUG:-WARN}" + +BRIDGE="${CONTAINER_NEMORL}/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge" +export COMMAND="\ +set -euo pipefail +NRL_SLURM_JOB_ID=\$(basename \"\$(dirname \"\$0\")\") +NRL_SLURM_JOB_ID=\${NRL_SLURM_JOB_ID%%-*} +cd ${CONTAINER_NEMORL} +export HF_HOME=${HF_HOME} +export HUGGINGFACE_HUB_CACHE=${HUGGINGFACE_HUB_CACHE} +export HF_HUB_CACHE=${HF_HUB_CACHE} +export HF_DATASETS_CACHE=${HF_DATASETS_CACHE} +export HF_MODULES_CACHE=${HF_MODULES_CACHE} +export TRANSFORMERS_CACHE=${TRANSFORMERS_CACHE} +export NRL_MEGATRON_CHECKPOINT_DIR=${NRL_MEGATRON_CHECKPOINT_DIR} +export XDG_CACHE_HOME=${XDG_CACHE_HOME} +export TORCH_HOME=${TORCH_HOME} +export TRITON_CACHE_DIR=${TRITON_CACHE_DIR} +${REFIT_ENV_EXPORTS} +export PYTHONPATH=${CONTAINER_NEMORL}:${BRIDGE}/src:${BRIDGE}/3rdparty/Megatron-LM\${PYTHONPATH:+:\$PYTHONPATH} +uv run --no-sync python examples/run_vlm_grpo.py --config ${CONFIG} \ +cluster.num_nodes=${NUM_NODES} \ +cluster.gpus_per_node=${GPUS_PER_NODE} \ +policy.model_name=${MODEL_NAME} \ +policy.tokenizer.name=${MODEL_NAME} \ +policy.is_vlm=true \ +policy.megatron_cfg.tensor_model_parallel_size=${POLICY_TP} \ +policy.megatron_cfg.expert_model_parallel_size=${POLICY_EP} \ +policy.megatron_cfg.expert_tensor_parallel_size=1 \ +policy.megatron_cfg.context_parallel_size=${POLICY_CP} \ +policy.megatron_cfg.optimizer.optimizer_cpu_offload=${OPTIMIZER_CPU_OFFLOAD} \ +policy.megatron_cfg.optimizer.optimizer_offload_fraction=${OPTIMIZER_OFFLOAD_FRACTION} \ +policy.offload_optimizer_for_logprob=${OFFLOAD_OPTIMIZER_FOR_LOGPROB} \ +${OPTIMIZER_DTYPE_OVERRIDES} \ +policy.generation.backend=${GENERATION_BACKEND} \ +policy.generation.colocated.enabled=${COLOCATED_ENABLED} \ +policy.generation.colocated.resources.num_nodes=${NUM_GEN_NODES} \ +policy.generation.colocated.resources.gpus_per_node=${GEN_GPUS_PER_NODE} \ +${GEN_OVERRIDES} \ +policy.max_total_sequence_length=${MAX_SEQUENCE_LENGTH} \ +policy.generation.max_new_tokens=${MAX_NEW_TOKENS} \ +grpo.async_grpo.enabled=${ASYNC_GRPO} \ +grpo.async_grpo.max_trajectory_age_steps=${MAX_TRAJECTORY_AGE_STEPS} \ +grpo.async_grpo.in_flight_weight_updates=${IN_FLIGHT_WEIGHT_UPDATES} \ +loss_fn.use_importance_sampling_correction=true \ +grpo.num_prompts_per_step=${NUM_PROMPTS} \ +grpo.num_generations_per_prompt=${NUM_GENERATIONS} \ +grpo.val_batch_size=${VAL_GBS} \ +grpo.max_val_samples=${VAL_SIZE} \ +policy.train_global_batch_size=${TRAIN_GBS} \ +grpo.max_num_steps=${MAX_STEPS} \ +checkpointing.enabled=${CHECKPOINTING_ENABLED} \ +checkpointing.checkpoint_dir=${RESULTS_DIR} \ +logger.log_dir=${RESULTS_DIR} \ +logger.wandb_enabled=${WANDB_ENABLED} \ +logger.wandb.name=${WANDB_NAME}-\${NRL_SLURM_JOB_ID} \ +logger.wandb.project=${WANDB_PROJ} \ ++logger.wandb.entity=${WANDB_GROUP} \ +${EXTRA_OVERRIDES}" + +echo "Submitting ${JOB_NAME}: ${NUM_NODES} node(s), ${GPUS_PER_NODE} GPU(s)/node" +if [[ "${COLOCATED_ENABLED}" == "true" ]]; then + echo " layout: colocated async on all ${NUM_NODES} node(s)" +else + echo " split: ${NUM_GEN_NODES} gen node(s) / $((NUM_NODES - NUM_GEN_NODES)) train node(s)" +fi +echo " generation backend: ${GENERATION_BACKEND} colocated=${COLOCATED_ENABLED} async=${ASYNC_GRPO}" +echo " training world size: ${TRAIN_WORLD_SIZE} (TP=${POLICY_TP}, EP=${POLICY_EP}, DP=${TRAIN_DP_SIZE})" +if [[ "${GENERATION_BACKEND}" == "megatron" ]]; then + echo " inference world size: ${INFERENCE_WORLD_SIZE} (TP=${INFER_TP}, EP=${INFER_EP}, DP=${INFERENCE_DP_SIZE})" + echo " refit backend: ${REFIT_BACKEND}" +else + echo " inference world size: ${INFERENCE_WORLD_SIZE} (TP=${INFER_TP}, EP=${INFER_EP}, DP=${INFERENCE_DP_SIZE})" + echo " vLLM: mem_util=${VLLM_GPU_MEMORY_UTILIZATION:-0.5} eager=${VLLM_ENFORCE_EAGER:-true} max_batched_tokens=${VLLM_MAX_NUM_BATCHED_TOKENS:-${MAX_SEQUENCE_LENGTH}}" + if [[ -n "${REFIT_ENV_EXPORTS}" ]]; then + echo " refit: buffer_memory_ratio=${REFIT_BUFFER_MEMORY_RATIO}" + fi +fi +echo " seq/new_tokens: ${MAX_SEQUENCE_LENGTH}/${MAX_NEW_TOKENS}" +echo " prompts/generations/train_gbs: ${NUM_PROMPTS}/${NUM_GENERATIONS}/${TRAIN_GBS}" +echo " async: max_trajectory_age=${MAX_TRAJECTORY_AGE_STEPS} in_flight_weight_updates=${IN_FLIGHT_WEIGHT_UPDATES}" +echo " val_gbs/val_size: ${VAL_GBS}/${VAL_SIZE}" +echo " optimizer moments: exp_avg=${EXP_AVG_DTYPE:-} exp_avg_sq=${EXP_AVG_SQ_DTYPE:-} store_param_remainders=${STORE_PARAM_REMAINDERS:-}" +echo " W&B: ${WANDB_GROUP}/${WANDB_PROJ}/${WANDB_NAME}- (enabled=${WANDB_ENABLED})" +echo " Hugging Face cache: ${HF_HUB_CACHE}" +echo " MCore checkpoint cache: ${NRL_MEGATRON_CHECKPOINT_DIR}" + +SBATCH_ARGS=( + --nodes="${NUM_NODES}" + --account="${SBATCH_ACCOUNT}" + --partition="${SBATCH_PARTITION}" + --job-name="${JOB_NAME}" + --time="${SBATCH_TIME}" + --output="${SLURM_LOG_DIR}/%j.out" + --error="${SLURM_LOG_DIR}/%j.out" + --gres="gpu:${GPUS_PER_NODE}" + --comment='{"OccupiedIdleGPUsJobReaper":{"exemptIdleTimeMins":"90","reason":"data_loading","description":"Async GRPO RL training: training GPUs idle during rollout collection (~30min) and validation each step"}}' + --exclusive + --mem=0 + --dependency=singleton + --segment="${SBATCH_SEGMENT}" +) +if [[ -n "${SBATCH_QOS}" ]]; then + SBATCH_ARGS+=(--qos="${SBATCH_QOS}") +fi +if [[ -n "${SBATCH_RESERVATION}" ]]; then + SBATCH_ARGS+=(--reservation="${SBATCH_RESERVATION}") +fi + +BASE_LOG_DIR="${SLURM_LOG_DIR}" \ +MOUNTS="${MOUNTS:-/lustre:/lustre},${NEMORL}:${CONTAINER_NEMORL}" \ +sbatch "${SBATCH_ARGS[@]}" "${NEMORL}/ray.sub" diff --git a/scripts/submit_nemotron_omni_vstat_megatron_8n4g.sh b/scripts/submit_nemotron_omni_vstat_megatron_8n4g.sh new file mode 100755 index 00000000000..721b78a873b --- /dev/null +++ b/scripts/submit_nemotron_omni_vstat_megatron_8n4g.sh @@ -0,0 +1,519 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Submit the eight-node Nemotron Omni VSTAT video-GRPO recipe through NeMo-RL's +# Ray/Slurm launcher. The default non-colocated layout reserves six nodes for +# generation and two for training. GENERATION_BACKEND=vllm enables the Omni +# vLLM A/B path; COLOCATED=true shares every GPU between training and generation. + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +NEMORL="${NEMORL:-$(cd "${SCRIPT_DIR}/.." && pwd)}" +CONTAINER_NEMORL="${CONTAINER_NEMORL:-/opt/nemo-rl}" +WORKSPACE_ROOT="${WORKSPACE_ROOT:-${NEMORL}/workspace}" + +CONTAINER="${CONTAINER:-/lustre/fsw/portfolios/coreai/users/cye/enroot/nemo_rl_0807.sqsh}" +MODEL_NAME="${MODEL_NAME:-nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16}" +CONFIG="${CONFIG:-examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-16n8g-megatron-tp4ep4-async-gym-video.v1.yaml}" +ENTRYPOINT="${ENTRYPOINT:-examples/nemo_gym/run_grpo_nemo_gym.py}" +GENERATION_BACKEND="${GENERATION_BACKEND:-megatron}" +COLOCATED="${COLOCATED:-false}" +ASYNC_GRPO="${ASYNC_GRPO:-true}" + +if [[ "${GENERATION_BACKEND}" != "megatron" && "${GENERATION_BACKEND}" != "vllm" ]]; then + echo "GENERATION_BACKEND must be megatron or vllm (got ${GENERATION_BACKEND})." >&2 + exit 1 +fi +if [[ "${COLOCATED}" == "true" && "${ASYNC_GRPO}" == "true" && "${GENERATION_BACKEND}" != "megatron" ]]; then + echo "Colocated async GRPO requires GENERATION_BACKEND=megatron." >&2 + exit 1 +fi + +CACHE_ROOT="${CACHE_ROOT:-${WORKSPACE_ROOT}/cache/nemo-rl-omni}" +export HF_HOME="${HF_HOME:-${CACHE_ROOT}/huggingface}" +export HUGGINGFACE_HUB_CACHE="${HUGGINGFACE_HUB_CACHE:-${HF_HOME}/hub}" +export HF_HUB_CACHE="${HF_HUB_CACHE:-${HUGGINGFACE_HUB_CACHE}}" +export HF_DATASETS_CACHE="${HF_DATASETS_CACHE:-${HF_HOME}/datasets}" +export HF_MODULES_CACHE="${HF_MODULES_CACHE:-${HF_HOME}/modules}" +export TRANSFORMERS_CACHE="${TRANSFORMERS_CACHE:-${HF_HOME}/transformers}" +export NRL_MEGATRON_CHECKPOINT_DIR="${NRL_MEGATRON_CHECKPOINT_DIR:-${CACHE_ROOT}/megatron-checkpoints}" +export XDG_CACHE_HOME="${XDG_CACHE_HOME:-${CACHE_ROOT}/xdg}" +export TORCH_HOME="${TORCH_HOME:-${CACHE_ROOT}/torch}" +export TRITON_CACHE_DIR="${TRITON_CACHE_DIR:-${CACHE_ROOT}/triton}" + +NUM_NODES="${NUM_NODES:-8}" +GPUS_PER_NODE="${GPUS_PER_NODE:-4}" + +if [[ "${COLOCATED}" == "true" ]]; then + TRAIN_WORLD_SIZE=$((NUM_NODES * GPUS_PER_NODE)) + INFERENCE_WORLD_SIZE="${TRAIN_WORLD_SIZE}" + GEN_GPUS_PER_NODE="${GPUS_PER_NODE}" + NUM_GEN_NODES="${NUM_NODES}" + COLOCATED_ENABLED=true +else + NUM_GEN_NODES="${NUM_GEN_NODES:-6}" + if (( NUM_NODES == 1 )); then + if (( NUM_GEN_NODES != 1 )); then + echo "One-node non-colocated mode requires NUM_GEN_NODES=1." >&2 + exit 1 + fi + GEN_GPUS_PER_NODE="${GEN_GPUS_PER_NODE:-$((GPUS_PER_NODE / 2))}" + if (( GPUS_PER_NODE < 4 || GEN_GPUS_PER_NODE <= 0 || GEN_GPUS_PER_NODE >= GPUS_PER_NODE )); then + echo "One-node non-colocated mode requires at least four GPUs and a strict train/inference split." >&2 + exit 1 + fi + TRAIN_WORLD_SIZE=$((GPUS_PER_NODE - GEN_GPUS_PER_NODE)) + else + GEN_GPUS_PER_NODE="${GEN_GPUS_PER_NODE:-${GPUS_PER_NODE}}" + if (( NUM_GEN_NODES <= 0 || NUM_GEN_NODES >= NUM_NODES )); then + echo "Multi-node non-colocated mode requires 0 < NUM_GEN_NODES < NUM_NODES." >&2 + exit 1 + fi + if (( GEN_GPUS_PER_NODE != GPUS_PER_NODE )); then + echo "Multi-node non-colocated inference must reserve complete GPU nodes." >&2 + exit 1 + fi + TRAIN_WORLD_SIZE=$(((NUM_NODES - NUM_GEN_NODES) * GPUS_PER_NODE)) + fi + INFERENCE_WORLD_SIZE=$((NUM_GEN_NODES * GEN_GPUS_PER_NODE)) + COLOCATED_ENABLED=false +fi + +# Prefer TP=8 and fall back by powers of two for smaller worlds. +DEFAULT_POLICY_TP=8 +while (( DEFAULT_POLICY_TP > TRAIN_WORLD_SIZE || TRAIN_WORLD_SIZE % DEFAULT_POLICY_TP != 0 )); do + DEFAULT_POLICY_TP=$((DEFAULT_POLICY_TP / 2)) +done +DEFAULT_INFER_TP=8 +while (( DEFAULT_INFER_TP > INFERENCE_WORLD_SIZE || INFERENCE_WORLD_SIZE % DEFAULT_INFER_TP != 0 )); do + DEFAULT_INFER_TP=$((DEFAULT_INFER_TP / 2)) +done +POLICY_TP="${POLICY_TP:-${DEFAULT_POLICY_TP}}" +INFER_TP="${INFER_TP:-${DEFAULT_INFER_TP}}" +POLICY_CP="${POLICY_CP:-1}" + +if [[ "${GENERATION_BACKEND}" == "megatron" && "${POLICY_CP}" != "1" ]]; then + echo "Megatron dynamic inference requires POLICY_CP=1." >&2 + exit 1 +fi +if (( TRAIN_WORLD_SIZE % (POLICY_TP * POLICY_CP) != 0 )); then + echo "Training world size must be divisible by POLICY_TP * POLICY_CP." >&2 + exit 1 +fi +if (( INFERENCE_WORLD_SIZE % INFER_TP != 0 )); then + echo "Inference world size must be divisible by INFER_TP." >&2 + exit 1 +fi +TRAIN_DP_SIZE=$((TRAIN_WORLD_SIZE / (POLICY_TP * POLICY_CP))) +INFERENCE_DP_SIZE=$((INFERENCE_WORLD_SIZE / INFER_TP)) + +# With ETP=1, world_size must be divisible by EP (not by TP*EP). +largest_ep() { + local world="$1" + local ep=8 + while (( ep > world || world % ep != 0 )); do + ep=$((ep / 2)) + done + printf '%d' "${ep}" +} + +POLICY_EP="${POLICY_EP:-$(largest_ep "${TRAIN_WORLD_SIZE}")}" +INFER_EP="${INFER_EP:-$(largest_ep "${INFERENCE_WORLD_SIZE}")}" +if (( TRAIN_WORLD_SIZE % POLICY_EP != 0 )); then + echo "Training world size must be divisible by POLICY_EP (ETP=1)." >&2 + exit 1 +fi +if [[ "${GENERATION_BACKEND}" == "vllm" ]]; then + # Async vLLM cannot use internal DP. EP=TP leaves replica-level DP to Ray. + if [[ "${INFER_EP}" != "${INFER_TP}" ]]; then + echo "Forcing INFER_EP=${INFER_TP} for vLLM (was ${INFER_EP}) for async-safe Ray DP." >&2 + fi + INFER_EP="${INFER_TP}" +elif (( INFERENCE_WORLD_SIZE % INFER_EP != 0 )); then + echo "Inference world size must be divisible by INFER_EP (ETP=1)." >&2 + exit 1 +fi +if [[ "${COLOCATED_ENABLED}" == "true" && "${GENERATION_BACKEND}" == "megatron" ]] && + (( POLICY_TP != INFER_TP || POLICY_EP != INFER_EP )); then + echo "Colocated Megatron requires matching policy/inference TP and EP topology." >&2 + exit 1 +fi + +MAX_STEPS="${MAX_STEPS:-1000000}" +MAX_SEQUENCE_LENGTH="${MAX_SEQUENCE_LENGTH:-8192}" +MAX_NEW_TOKENS="${MAX_NEW_TOKENS:-2048}" +MIN_GENERATION_TOKENS="${MIN_GENERATION_TOKENS:-2000}" +VISION_EMBEDDING_CACHE_MAX_BYTES="${VISION_EMBEDDING_CACHE_MAX_BYTES:-536870912}" +ENABLE_THINKING="${ENABLE_THINKING:-true}" +NUM_FRAMES="${NUM_FRAMES:-16}" +TEMPORAL_PATCH_SIZE="${TEMPORAL_PATCH_SIZE:-2}" +# Match the checkpoint-native vLLM video processor. Policy preprocessing and +# Megatron inference consume this same value below. +VIDEO_TARGET_PATCHES="${VIDEO_TARGET_PATCHES:-1024}" +NUM_DATA_ROWS="${NUM_DATA_ROWS:-256}" +NUM_PROMPTS="${NUM_PROMPTS:-$((INFERENCE_DP_SIZE * 2))}" +NUM_GENERATIONS="${NUM_GENERATIONS:-8}" +TRAIN_GBS="${TRAIN_GBS:-$((NUM_PROMPTS * NUM_GENERATIONS))}" +EXPECTED_TRAIN_GBS=$((NUM_PROMPTS * NUM_GENERATIONS)) +if (( TRAIN_GBS != EXPECTED_TRAIN_GBS )); then + echo "TRAIN_GBS (${TRAIN_GBS}) must equal NUM_PROMPTS * NUM_GENERATIONS (${EXPECTED_TRAIN_GBS})." >&2 + exit 1 +fi +if (( TRAIN_GBS % TRAIN_DP_SIZE != 0 )); then + echo "TRAIN_GBS (${TRAIN_GBS}) must be divisible by training DP size (${TRAIN_DP_SIZE})." >&2 + exit 1 +fi + +DATA_ROOT="${DATA_ROOT:-${CONTAINER_NEMORL}/workspace/datasets/vstat-8n4g}" +HF_DATASET="${HF_DATASET:-ShushengYang/VSTAT}" +PREPARE_VSTAT="${PREPARE_VSTAT:-false}" +export NEMO_RL_VIDEO_TRAIN_JSONL="${NEMO_RL_VIDEO_TRAIN_JSONL:-${DATA_ROOT}/train-gym.jsonl}" +export NEMO_RL_VIDEO_VAL_JSONL="${NEMO_RL_VIDEO_VAL_JSONL:-${DATA_ROOT}/val-gym.jsonl}" +export NEMO_RL_VIDEO_MEDIA_ROOT="${NEMO_RL_VIDEO_MEDIA_ROOT:-${DATA_ROOT}/media}" + +REFIT_BACKEND="${REFIT_BACKEND:-nccl}" +OPTIMIZER_CPU_OFFLOAD="${OPTIMIZER_CPU_OFFLOAD:-false}" +BUFFER_SIZE_GB="${BUFFER_SIZE_GB:-8}" +OFFLOAD_OPTIMIZER_FOR_LOGPROB="${OFFLOAD_OPTIMIZER_FOR_LOGPROB:-false}" +if [[ "${OPTIMIZER_CPU_OFFLOAD}" == "true" ]]; then + OPTIMIZER_OFFLOAD_FRACTION="${OPTIMIZER_OFFLOAD_FRACTION:-1.0}" +else + OPTIMIZER_OFFLOAD_FRACTION="${OPTIMIZER_OFFLOAD_FRACTION:-0.0}" +fi +USE_PRECISION_AWARE_OPTIMIZER="${USE_PRECISION_AWARE_OPTIMIZER:-true}" +EXP_AVG_DTYPE="${EXP_AVG_DTYPE:-bfloat16}" +EXP_AVG_SQ_DTYPE="${EXP_AVG_SQ_DTYPE:-bfloat16}" +STORE_PARAM_REMAINDERS="${STORE_PARAM_REMAINDERS:-true}" +MEGATRON_ENABLE_CHUNKED_PREFILL="${MEGATRON_ENABLE_CHUNKED_PREFILL:-true}" +MEGATRON_CUDA_GRAPH_IMPL="${MEGATRON_CUDA_GRAPH_IMPL:-local}" +MEGATRON_CUDA_GRAPH_SCOPE="${MEGATRON_CUDA_GRAPH_SCOPE:-block}" +MEGATRON_NUM_CUDA_GRAPHS="${MEGATRON_NUM_CUDA_GRAPHS:--1}" +MEGATRON_USE_CUDA_GRAPHS_FOR_NON_DECODE="${MEGATRON_USE_CUDA_GRAPHS_FOR_NON_DECODE:-false}" +MOE_PAD_EXPERTS_FOR_CG="${MOE_PAD_EXPERTS_FOR_CG:-true}" +MAX_TRAJECTORY_AGE_STEPS="${MAX_TRAJECTORY_AGE_STEPS:-2}" +IN_FLIGHT_WEIGHT_UPDATES="${IN_FLIGHT_WEIGHT_UPDATES:-true}" +EXTRA_OVERRIDES="${EXTRA_OVERRIDES:-}" + +JOB_NAME="${JOB_NAME:-nemotron-omni-vstat-${GENERATION_BACKEND}-8n4g}" +EXP_NAME="${EXP_NAME:-${JOB_NAME}}" +PRECISION_RECIPE="${PRECISION_RECIPE:-bf16}" +WANDB_ENABLED="${WANDB_ENABLED:-true}" +WANDB_PROJ="${WANDB_PROJ:-mllm-rl-dev}" +WANDB_GROUP="${WANDB_GROUP:-adlr}" +WANDB_NAME="${WANDB_NAME:-${EXP_NAME}-${PRECISION_RECIPE}-internal-repo}" +RESULTS_DIR="${RESULTS_DIR:-${WORKSPACE_ROOT}/results/nemo-rl-omni/${JOB_NAME}}" +CHECKPOINTING_ENABLED="${CHECKPOINTING_ENABLED:-false}" +SLURM_LOG_DIR="${SLURM_LOG_DIR:-${RESULTS_DIR}/slurm}" + +SBATCH_ACCOUNT="${SBATCH_ACCOUNT:-nemotron_sw_post}" +SBATCH_PARTITION="${SBATCH_PARTITION:-batch_long}" +SBATCH_QOS="${SBATCH_QOS:-}" +SBATCH_TIME="${SBATCH_TIME:-04:00:00}" +SBATCH_RESERVATION="${SBATCH_RESERVATION:-}" +SBATCH_SEGMENT="${NUM_NODES}" + +mkdir -p \ + "${HF_HUB_CACHE}" \ + "${HF_DATASETS_CACHE}" \ + "${HF_MODULES_CACHE}" \ + "${TRANSFORMERS_CACHE}" \ + "${NRL_MEGATRON_CHECKPOINT_DIR}" \ + "${XDG_CACHE_HOME}" \ + "${TORCH_HOME}" \ + "${TRITON_CACHE_DIR}" \ + "${SLURM_LOG_DIR}" + +if [[ ! -f "${CONTAINER}" ]]; then + echo "Container image does not exist: ${CONTAINER}" >&2 + exit 1 +fi +if [[ ! -f "${NEMORL}/ray.sub" || ! -f "${NEMORL}/${CONFIG}" || + ! -f "${NEMORL}/${ENTRYPOINT}" || ! -f "${NEMORL}/scripts/prepare_nemotron_omni_vstat.py" ]]; then + echo "NeMo-RL launcher, config, entrypoint, or VSTAT preparation script is missing under: ${NEMORL}" >&2 + exit 1 +fi + +OPTIMIZER_DTYPE_OVERRIDES="" +if [[ -n "${EXP_AVG_DTYPE}" ]]; then + OPTIMIZER_DTYPE_OVERRIDES+=" ++policy.megatron_cfg.optimizer.exp_avg_dtype=${EXP_AVG_DTYPE}" +fi +if [[ -n "${EXP_AVG_SQ_DTYPE}" ]]; then + OPTIMIZER_DTYPE_OVERRIDES+=" ++policy.megatron_cfg.optimizer.exp_avg_sq_dtype=${EXP_AVG_SQ_DTYPE}" +fi +if [[ -n "${STORE_PARAM_REMAINDERS}" ]]; then + OPTIMIZER_DTYPE_OVERRIDES+=" ++policy.megatron_cfg.optimizer.store_param_remainders=${STORE_PARAM_REMAINDERS}" +fi + +REFIT_ENV_EXPORTS="" +REFIT_BUFFER_MEMORY_RATIO="" +if [[ "${GENERATION_BACKEND}" == "megatron" ]]; then + GEN_OVERRIDES="\ +++policy.generation.stop_strings=null \ +++policy.generation.bad_words=null \ +policy.generation.mcore_generation_config.tensor_model_parallel_size=${INFER_TP} \ +policy.generation.mcore_generation_config.expert_model_parallel_size=${INFER_EP} \ +policy.generation.mcore_generation_config.expert_tensor_parallel_size=1 \ +++policy.generation.mcore_generation_config.context_parallel_size=${POLICY_CP} \ +++policy.generation.mcore_generation_config.mamba_inference_ssm_states_dtype=float32 \ +++policy.generation.mcore_generation_config.mamba_inference_conv_states_dtype=float32 \ +++policy.generation.mcore_generation_config.logprobs_mode=raw_logprobs \ +policy.generation.mcore_generation_config.parsers=[nemotron-v3-reasoning,qwen3-coder-tool] \ +policy.generation.mcore_generation_config.transformer_impl=transformer_engine \ +policy.generation.mcore_generation_config.sequence_parallel=true \ +policy.generation.mcore_generation_config.enable_chunked_prefill=${MEGATRON_ENABLE_CHUNKED_PREFILL} \ +++policy.generation.mcore_generation_config.async_sched_mode=async \ +policy.generation.mcore_generation_config.enable_prefix_caching=false \ +++policy.generation.mcore_generation_config.vision_embedding_cache_max_bytes=${VISION_EMBEDDING_CACHE_MAX_BYTES} \ +policy.generation.mcore_generation_config.cuda_graph_impl=${MEGATRON_CUDA_GRAPH_IMPL} \ +policy.generation.mcore_generation_config.inference_cuda_graph_scope=${MEGATRON_CUDA_GRAPH_SCOPE} \ +policy.generation.mcore_generation_config.num_cuda_graphs=${MEGATRON_NUM_CUDA_GRAPHS} \ +policy.generation.mcore_generation_config.use_cuda_graphs_for_non_decode_steps=${MEGATRON_USE_CUDA_GRAPHS_FOR_NON_DECODE} \ +++policy.generation.mcore_generation_config.moe_pad_experts_for_cuda_graph_inference=${MOE_PAD_EXPERTS_FOR_CG} \ +policy.generation.mcore_generation_config.refit_backend=${REFIT_BACKEND} \ +policy.generation.mcore_generation_config.buffer_size_gb=${BUFFER_SIZE_GB} \ +policy.generation.mcore_generation_config.max_model_len=${MAX_SEQUENCE_LENGTH} \ +policy.generation.mcore_generation_config.max_tokens=${MAX_SEQUENCE_LENGTH} \ +++policy.generation.mcore_generation_config.video_num_frames=${NUM_FRAMES} \ +++policy.generation.mcore_generation_config.video_temporal_patch_size=${TEMPORAL_PATCH_SIZE} \ +++policy.generation.mcore_generation_config.video_target_num_patches=${VIDEO_TARGET_PATCHES}" +else + REFIT_BUFFER_MEMORY_RATIO="${NRL_REFIT_BUFFER_MEMORY_RATIO:-0.005}" + REFIT_ENV_EXPORTS="export NRL_REFIT_BUFFER_MEMORY_RATIO=${REFIT_BUFFER_MEMORY_RATIO}" + VLLM_GPU_MEMORY_UTILIZATION="${VLLM_GPU_MEMORY_UTILIZATION:-0.5}" + VLLM_ENFORCE_EAGER="${VLLM_ENFORCE_EAGER:-true}" + VLLM_MAX_NUM_SEQS="${VLLM_MAX_NUM_SEQS:-4}" + VLLM_MAX_NUM_BATCHED_TOKENS="${VLLM_MAX_NUM_BATCHED_TOKENS:-${MAX_SEQUENCE_LENGTH}}" + GEN_OVERRIDES="\ +++policy.generation.stop_strings=null \ +++policy.generation.bad_words=null \ +policy.generation.vllm_cfg.async_engine=${ASYNC_GRPO} \ +policy.generation.vllm_cfg.skip_tokenizer_init=false \ +policy.generation.vllm_cfg.tensor_parallel_size=${INFER_TP} \ +policy.generation.vllm_cfg.pipeline_parallel_size=1 \ +policy.generation.vllm_cfg.expert_parallel_size=${INFER_EP} \ +policy.generation.vllm_cfg.max_model_len=${MAX_SEQUENCE_LENGTH} \ +++policy.generation.vllm_cfg.cap_max_tokens_to_context=true \ +policy.generation.vllm_cfg.gpu_memory_utilization=${VLLM_GPU_MEMORY_UTILIZATION} \ +policy.generation.vllm_cfg.enforce_eager=${VLLM_ENFORCE_EAGER} \ +policy.generation.vllm_cfg.enable_prefix_caching=false \ +policy.generation.vllm_cfg.logprobs_mode=raw_logprobs \ +policy.generation.vllm_cfg.reset_mm_cache_after_refit=true \ +policy.generation.vllm_cfg.env_vars.NRL_VIDEO_BACKEND=torchcodec \ +policy.generation.vllm_cfg.env_vars.NRL_VIDEO_SAMPLING_STYLE=nemotron_vl \ +policy.generation.vllm_cfg.env_vars.NRL_VIDEO_TEMPORAL_PATCH_SIZE=${TEMPORAL_PATCH_SIZE} \ +policy.generation.vllm_cfg.env_vars.VLLM_VIDEO_LOADER_BACKEND=nemotron_vl \ +policy.generation.vllm_kwargs.allowed_local_media_path=${NEMO_RL_VIDEO_MEDIA_ROOT} \ +policy.generation.vllm_kwargs.mm_processor_cache_gb=0 \ +policy.generation.vllm_kwargs.max_num_seqs=${VLLM_MAX_NUM_SEQS} \ +policy.generation.vllm_kwargs.limit_mm_per_prompt.video.count=1 \ +policy.generation.vllm_kwargs.limit_mm_per_prompt.video.num_frames=${NUM_FRAMES} \ +++policy.generation.vllm_kwargs.limit_mm_per_prompt.image=${NUM_FRAMES} \ +policy.generation.vllm_kwargs.max_num_batched_tokens=${VLLM_MAX_NUM_BATCHED_TOKENS} \ +policy.generation.vllm_kwargs.enable_chunked_prefill=false \ +policy.generation.vllm_kwargs.disable_custom_all_reduce=true \ +policy.generation.vllm_kwargs.attention_backend=FLASH_ATTN \ +policy.generation.vllm_kwargs.attention_config.use_trtllm_attention=false \ +++policy.generation.vllm_kwargs.mamba_ssm_cache_dtype=float32 \ +++policy.generation.vllm_kwargs.skip_mm_profiling=false \ +++policy.generation.vllm_kwargs.kernel_config.enable_flashinfer_autotune=false \ +++policy.generation.vllm_kwargs.kernel_config.moe_backend=triton" +fi + +export NUM_NODES GPUS_PER_NODE CONTAINER +export RAY_ENABLE_UV_RUN_RUNTIME_ENV=0 +export NRL_VENVS_TRUST_EXISTING="${NRL_VENVS_TRUST_EXISTING:-1}" +export NRL_FORCE_REBUILD_VENVS="${NRL_FORCE_REBUILD_VENVS:-false}" +export NEMO_RL_VENV_DIR="${NEMO_RL_VENV_DIR:-/opt/ray_venvs}" +export NEMO_GYM_VENV_DIR="${NEMO_GYM_VENV_DIR:-/opt/ray_venvs}" +export NEMO_GYM_EXTRA_ROOTS="${NEMO_GYM_EXTRA_ROOTS:-${CONTAINER_NEMORL}/3rdparty/Gym-workspace/Gym}" +export NRL_VIDEO_BACKEND="${NRL_VIDEO_BACKEND:-torchcodec}" +export NRL_VIDEO_SAMPLING_STYLE="${NRL_VIDEO_SAMPLING_STYLE:-nemotron_vl}" +export NRL_VIDEO_TEMPORAL_PATCH_SIZE="${TEMPORAL_PATCH_SIZE}" +export VLLM_VIDEO_LOADER_BACKEND="${VLLM_VIDEO_LOADER_BACKEND:-nemotron_vl}" +export TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST:-10.0}" +export CUDA_DEVICE_MAX_CONNECTIONS="${CUDA_DEVICE_MAX_CONNECTIONS:-1}" +export FLASHINFER_DISABLE_VERSION_CHECK="${FLASHINFER_DISABLE_VERSION_CHECK:-1}" +export NVTE_FWD_LAYERNORM_SM_MARGIN="${NVTE_FWD_LAYERNORM_SM_MARGIN:-16}" +export NVTE_BWD_LAYERNORM_SM_MARGIN="${NVTE_BWD_LAYERNORM_SM_MARGIN:-16}" +export NCCL_DEBUG="${NCCL_DEBUG:-WARN}" + +ENABLE_NSYS="${ENABLE_NSYS:-false}" +if [[ "${ENABLE_NSYS}" == "true" ]]; then + export NRL_NSYS_WORKER_PATTERNS="${NRL_NSYS_WORKER_PATTERNS:-*policy*,*megatron*,*vllm*}" + export NRL_NSYS_PROFILE_STEP_RANGE="${NRL_NSYS_PROFILE_STEP_RANGE:-1:4}" + export LD_LIBRARY_PATH="/usr/local/cuda/targets/aarch64-linux/lib:/usr/local/cuda/targets/x86_64-linux/lib:/usr/local/cuda/lib64:/usr/local/cuda/lib:/usr/local/nvidia/lib64:/usr/local/nvidia/lib:/usr/lib/aarch64-linux-gnu:/usr/lib/x86_64-linux-gnu:${LD_LIBRARY_PATH:-}" + export NRL_NSYS_EXTRA_OPTIONS="${NRL_NSYS_EXTRA_OPTIONS:-{\"o\":\"${CONTAINER_NEMORL}/workspace/nsys/%p\",\"cpuctxsw\":\"none\",\"force-overwrite\":\"true\"}}" + mkdir -p "${WORKSPACE_ROOT}/nsys" +fi + +BRIDGE="${CONTAINER_NEMORL}/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge" +export PYTHONPATH="${CONTAINER_NEMORL}:${NEMO_GYM_EXTRA_ROOTS}:${BRIDGE}/src:${BRIDGE}/3rdparty/Megatron-LM${PYTHONPATH:+:${PYTHONPATH}}" +export SETUP_COMMAND="\ +set -euo pipefail +cd ${CONTAINER_NEMORL} +bash tools/install_audio_deps.sh +if [[ ${GENERATION_BACKEND} == vllm ]]; then + for vllm_python in ${NEMO_RL_VENV_DIR}/nemo_rl.models.generation.vllm.*Vllm*Worker/bin/python; do + [[ -x \"\${vllm_python}\" ]] || continue + \"\${vllm_python}\" -m pip install --no-cache-dir \ + --index-url https://download.pytorch.org/whl/cu130 \ + --extra-index-url https://pypi.org/simple \ + torchcodec==0.11.1 + done +fi" + +export COMMAND="\ +set -euo pipefail +NRL_SLURM_JOB_ID=\$(basename \"\$(dirname \"\$0\")\") +NRL_SLURM_JOB_ID=\${NRL_SLURM_JOB_ID%%-*} +cd ${CONTAINER_NEMORL} +export HF_HOME=${HF_HOME} +export HUGGINGFACE_HUB_CACHE=${HUGGINGFACE_HUB_CACHE} +export HF_HUB_CACHE=${HF_HUB_CACHE} +export HF_DATASETS_CACHE=${HF_DATASETS_CACHE} +export HF_MODULES_CACHE=${HF_MODULES_CACHE} +export TRANSFORMERS_CACHE=${TRANSFORMERS_CACHE} +export NRL_MEGATRON_CHECKPOINT_DIR=${NRL_MEGATRON_CHECKPOINT_DIR} +export XDG_CACHE_HOME=${XDG_CACHE_HOME} +export TORCH_HOME=${TORCH_HOME} +export TRITON_CACHE_DIR=${TRITON_CACHE_DIR} +export NEMO_RL_VENV_DIR=${NEMO_RL_VENV_DIR} +export NEMO_GYM_VENV_DIR=${NEMO_GYM_VENV_DIR} +export NEMO_GYM_EXTRA_ROOTS=${NEMO_GYM_EXTRA_ROOTS} +export NRL_VIDEO_BACKEND=${NRL_VIDEO_BACKEND} +export NRL_VIDEO_SAMPLING_STYLE=${NRL_VIDEO_SAMPLING_STYLE} +export NRL_VIDEO_TEMPORAL_PATCH_SIZE=${NRL_VIDEO_TEMPORAL_PATCH_SIZE} +export VLLM_VIDEO_LOADER_BACKEND=${VLLM_VIDEO_LOADER_BACKEND} +export TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST} +export NEMO_RL_VIDEO_TRAIN_JSONL=${NEMO_RL_VIDEO_TRAIN_JSONL} +export NEMO_RL_VIDEO_VAL_JSONL=${NEMO_RL_VIDEO_VAL_JSONL} +export NEMO_RL_VIDEO_MEDIA_ROOT=${NEMO_RL_VIDEO_MEDIA_ROOT} +${REFIT_ENV_EXPORTS} +export PYTHONPATH=${CONTAINER_NEMORL}:${NEMO_GYM_EXTRA_ROOTS}:${BRIDGE}/src:${BRIDGE}/3rdparty/Megatron-LM\${PYTHONPATH:+:\$PYTHONPATH} +mkdir -p ${DATA_ROOT} +if [[ ${PREPARE_VSTAT} == true || ! -s \${NEMO_RL_VIDEO_TRAIN_JSONL} || ! -s \${NEMO_RL_VIDEO_VAL_JSONL} ]]; then + echo \"Preparing VSTAT inside the mounted container path: ${DATA_ROOT}\" + uv run --no-sync python scripts/prepare_nemotron_omni_vstat.py --output-dir ${DATA_ROOT} --repo-id ${HF_DATASET} --num-rows ${NUM_DATA_ROWS} +fi +uv run --no-sync python ${ENTRYPOINT} --config ${CONFIG} \ +cluster.num_nodes=${NUM_NODES} \ +cluster.gpus_per_node=${GPUS_PER_NODE} \ +policy.model_name=${MODEL_NAME} \ +policy.tokenizer.name=${MODEL_NAME} \ +policy.tokenizer.chat_template_kwargs.enable_thinking=${ENABLE_THINKING} \ +policy.is_vlm=true \ +policy.megatron_cfg.env_vars.TORCH_CUDA_ARCH_LIST=\"'${TORCH_CUDA_ARCH_LIST}'\" \ +policy.megatron_cfg.freeze_vision_model=false \ +policy.megatron_cfg.freeze_vision_projection=false \ +policy.megatron_cfg.freeze_moe_router=false \ +policy.megatron_cfg.mtp_num_layers=0 \ +policy.megatron_cfg.mtp_use_repeated_layer=true \ +policy.megatron_cfg.mtp_detach_heads=true \ +policy.megatron_cfg.mtp_loss_scaling_factor=0.0 \ +policy.megatron_cfg.tensor_model_parallel_size=${POLICY_TP} \ +policy.megatron_cfg.pipeline_model_parallel_size=1 \ +policy.megatron_cfg.expert_model_parallel_size=${POLICY_EP} \ +policy.megatron_cfg.expert_tensor_parallel_size=1 \ +policy.megatron_cfg.context_parallel_size=${POLICY_CP} \ +policy.megatron_cfg.sequence_parallel=true \ +policy.megatron_cfg.bias_activation_fusion=false \ +policy.megatron_cfg.moe_shared_expert_overlap=false \ +policy.megatron_cfg.radio_force_cpe_eval_mode=true \ +policy.megatron_cfg.clear_memory_caches_before_refit=true \ +policy.megatron_cfg.distributed_data_parallel_config.overlap_grad_reduce=false \ +policy.megatron_cfg.distributed_data_parallel_config.overlap_param_gather=false \ +policy.megatron_cfg.optimizer.params_dtype=float32 \ +policy.megatron_cfg.optimizer.use_precision_aware_optimizer=${USE_PRECISION_AWARE_OPTIMIZER} \ +policy.megatron_cfg.optimizer.optimizer_cpu_offload=${OPTIMIZER_CPU_OFFLOAD} \ +policy.megatron_cfg.optimizer.optimizer_offload_fraction=${OPTIMIZER_OFFLOAD_FRACTION} \ +policy.offload_optimizer_for_logprob=${OFFLOAD_OPTIMIZER_FOR_LOGPROB} \ +${OPTIMIZER_DTYPE_OVERRIDES} \ +policy.generation.backend=${GENERATION_BACKEND} \ +policy.generation.colocated.enabled=${COLOCATED_ENABLED} \ +policy.generation.colocated.resources.num_nodes=${NUM_GEN_NODES} \ +policy.generation.colocated.resources.gpus_per_node=${GEN_GPUS_PER_NODE} \ +${GEN_OVERRIDES} \ +policy.max_total_sequence_length=${MAX_SEQUENCE_LENGTH} \ +policy.generation.max_new_tokens=${MAX_NEW_TOKENS} \ +data.max_input_seq_length=${MAX_SEQUENCE_LENGTH} \ +data.num_workers=0 \ ++data.default.num_frames=${NUM_FRAMES} \ ++data.default.video_sampling_style=nemotron_vl \ ++data.default.video_temporal_patch_size=${TEMPORAL_PATCH_SIZE} \ ++data.default.min_generation_tokens=${MIN_GENERATION_TOKENS} \ +data.default.video_target_num_patches=${VIDEO_TARGET_PATCHES} \ +data.default.video_maintain_aspect_ratio=true \ +data.train.data_path=${NEMO_RL_VIDEO_TRAIN_JSONL} \ +data.validation.data_path=${NEMO_RL_VIDEO_VAL_JSONL} \ +++env.nemo_gym.policy_model.responses_api_models.vllm_model.chat_template_kwargs.enable_thinking=${ENABLE_THINKING} \ +grpo.deduplicate_multimodal_data=false \ +grpo.async_grpo.enabled=${ASYNC_GRPO} \ +grpo.async_grpo.max_trajectory_age_steps=${MAX_TRAJECTORY_AGE_STEPS} \ +grpo.async_grpo.in_flight_weight_updates=${IN_FLIGHT_WEIGHT_UPDATES} \ +loss_fn.use_importance_sampling_correction=true \ +grpo.num_prompts_per_step=${NUM_PROMPTS} \ +grpo.num_generations_per_prompt=${NUM_GENERATIONS} \ +grpo.val_num_generations_per_prompt=1 \ +policy.train_global_batch_size=${TRAIN_GBS} \ +grpo.max_num_steps=${MAX_STEPS} \ +checkpointing.enabled=${CHECKPOINTING_ENABLED} \ +checkpointing.checkpoint_dir=${RESULTS_DIR} \ +logger.log_dir=${RESULTS_DIR} \ +logger.wandb_enabled=${WANDB_ENABLED} \ +logger.wandb.name=${WANDB_NAME}-\${NRL_SLURM_JOB_ID} \ +logger.wandb.project=${WANDB_PROJ} \ ++logger.wandb.entity=${WANDB_GROUP} \ +${EXTRA_OVERRIDES}" + +echo "Submitting ${JOB_NAME}: ${NUM_NODES} node(s), ${GPUS_PER_NODE} GPU(s)/node" +if [[ "${COLOCATED_ENABLED}" == "true" ]]; then + echo " layout: colocated async on all ${NUM_NODES} node(s)" +else + echo " split: ${NUM_GEN_NODES} gen node(s) / $((NUM_NODES - NUM_GEN_NODES)) train node(s)" +fi +echo " generation backend: ${GENERATION_BACKEND} colocated=${COLOCATED_ENABLED} async=${ASYNC_GRPO}" +echo " training world size: ${TRAIN_WORLD_SIZE} (TP=${POLICY_TP}, EP=${POLICY_EP}, DP=${TRAIN_DP_SIZE})" +echo " inference world size: ${INFERENCE_WORLD_SIZE} (TP=${INFER_TP}, EP=${INFER_EP}, DP=${INFERENCE_DP_SIZE})" +if [[ "${GENERATION_BACKEND}" == "megatron" ]]; then + echo " Megatron: refit=${REFIT_BACKEND} chunked_prefill=${MEGATRON_ENABLE_CHUNKED_PREFILL} CUDA graphs=${MEGATRON_CUDA_GRAPH_IMPL}/${MEGATRON_CUDA_GRAPH_SCOPE}" +else + echo " vLLM: mem_util=${VLLM_GPU_MEMORY_UTILIZATION} eager=${VLLM_ENFORCE_EAGER} max_seqs=${VLLM_MAX_NUM_SEQS} max_batched_tokens=${VLLM_MAX_NUM_BATCHED_TOKENS}" + echo " refit buffer memory ratio: ${REFIT_BUFFER_MEMORY_RATIO}" +fi +echo " seq/new_tokens: ${MAX_SEQUENCE_LENGTH}/${MAX_NEW_TOKENS}" +echo " thinking enabled: ${ENABLE_THINKING}" +echo " prompts/generations/train_gbs: ${NUM_PROMPTS}/${NUM_GENERATIONS}/${TRAIN_GBS}" +echo " async: max_trajectory_age=${MAX_TRAJECTORY_AGE_STEPS} in_flight_weight_updates=${IN_FLIGHT_WEIGHT_UPDATES}" +echo " video: frames=${NUM_FRAMES} temporal_patch=${TEMPORAL_PATCH_SIZE} target_patches=${VIDEO_TARGET_PATCHES}" +echo " VSTAT: root=${DATA_ROOT} repo=${HF_DATASET} rows=${NUM_DATA_ROWS} prepare=${PREPARE_VSTAT} (head only; also runs if JSONL is missing)" +echo " datasets: train=${NEMO_RL_VIDEO_TRAIN_JSONL} val=${NEMO_RL_VIDEO_VAL_JSONL} media=${NEMO_RL_VIDEO_MEDIA_ROOT}" +echo " dependencies: targeted audio/video dependency setup runs on every node before Ray" +echo " NSYS: enabled=${ENABLE_NSYS}${NRL_NSYS_PROFILE_STEP_RANGE:+ step_range=${NRL_NSYS_PROFILE_STEP_RANGE}}" +echo " W&B: ${WANDB_GROUP}/${WANDB_PROJ}/${WANDB_NAME}- (enabled=${WANDB_ENABLED})" +echo " Hugging Face cache: ${HF_HUB_CACHE}" +echo " MCore checkpoint cache: ${NRL_MEGATRON_CHECKPOINT_DIR}" + +SBATCH_ARGS=( + --nodes="${NUM_NODES}" + --account="${SBATCH_ACCOUNT}" + --partition="${SBATCH_PARTITION}" + --job-name="${JOB_NAME}" + --time="${SBATCH_TIME}" + --output="${SLURM_LOG_DIR}/%j.out" + --error="${SLURM_LOG_DIR}/%j.out" + --gres="gpu:${GPUS_PER_NODE}" + --comment='{"OccupiedIdleGPUsJobReaper":{"exemptIdleTimeMins":"90","reason":"data_loading","description":"Async GRPO RL training: training GPUs idle during rollout collection (~30min) and validation each step"}}' + --exclusive + --mem=0 + --dependency=singleton + --segment="${SBATCH_SEGMENT}" +) +if [[ -n "${SBATCH_QOS}" ]]; then + SBATCH_ARGS+=(--qos="${SBATCH_QOS}") +fi +if [[ -n "${SBATCH_RESERVATION}" ]]; then + SBATCH_ARGS+=(--reservation="${SBATCH_RESERVATION}") +fi + +BASE_LOG_DIR="${SLURM_LOG_DIR}" \ +MOUNTS="${MOUNTS:-/lustre:/lustre},${NEMORL}:${CONTAINER_NEMORL}" \ +sbatch "${SBATCH_ARGS[@]}" "${NEMORL}/ray.sub" diff --git a/tests/unit/data/datasets/test_mmpr_tiny.py b/tests/unit/data/datasets/test_mmpr_tiny.py index 4b0c9e4fc28..3365f8e8572 100644 --- a/tests/unit/data/datasets/test_mmpr_tiny.py +++ b/tests/unit/data/datasets/test_mmpr_tiny.py @@ -250,7 +250,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): @@ -286,13 +286,13 @@ 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"] + 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 +310,14 @@ 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 + # Placeholder-style processors send their 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..7cf80429350 100755 --- a/tests/unit/data/test_collate_fn.py +++ b/tests/unit/data/test_collate_fn.py @@ -16,7 +16,7 @@ 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 +149,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/tools/install_audio_deps.sh b/tools/install_audio_deps.sh index 312cf5f3325..9cea2af9739 100755 --- a/tools/install_audio_deps.sh +++ b/tools/install_audio_deps.sh @@ -5,27 +5,33 @@ # # 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 and must be installed into +# the isolated Megatron policy worker environment that imports it. +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] Force-reinstalling PyAV in the Megatron worker environment..." + "$RAY_MEGATRON_PYTHON" -m pip install --no-cache-dir --force-reinstall av +else + echo "[audio-deps] Megatron worker environment not found; skipping PyAV: $RAY_MEGATRON_PYTHON" +fi echo "[audio-deps] Done." From 9d876ecdb0911dbfcb79a56d26e7dd3bc6a441e1 Mon Sep 17 00:00:00 2001 From: Cory Ye Date: Fri, 21 Aug 2026 19:07:49 -0700 Subject: [PATCH 02/15] Post-rebase housekeeping. Signed-off-by: Cory Ye --- ...8g-megatron-tp4ep4-async-gym-video.v1.yaml | 9 +++ ...-circle-click-2n8g-megatron-tp2ep8.v1.yaml | 8 +++ ...cle-count-1n4g-megatron_generation.v1.yaml | 8 +++ ...a3b-clevr-1n4g-megatron_generation.v1.yaml | 8 +++ ...a3b-clevr-8n4g-megatron_generation.v1.yaml | 8 +++ nemo_rl/data/multimodal_utils.py | 50 ++++++--------- nemo_rl/data/processors.py | 5 +- nemo_rl/environments/nemo_gym.py | 64 ++++++++----------- nemo_rl/environments/nemo_gym_video.py | 35 ---------- nemo_rl/evals/eval.py | 1 + nemo_rl/experience/rollouts.py | 12 +++- nemo_rl/models/generation/megatron/config.py | 2 + .../megatron/megatron_generation.py | 6 +- .../generation/megatron/megatron_worker.py | 28 ++------ nemo_rl/models/generation/megatron/utils.py | 29 ++++++++- nemo_rl/models/generation/vllm/video_utils.py | 17 ++--- tests/unit/data/datasets/test_mmpr_tiny.py | 5 +- .../data/test_multimodal_image_encoding.py | 18 +++--- tests/unit/environments/test_nemo_gym.py | 16 ++--- .../generation/test_vllm_video_utils.py | 6 +- 20 files changed, 173 insertions(+), 162 deletions(-) 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..f0db0455927 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,15 @@ policy: make_sequence_length_divisible_by: 32 generation: bad_words: [] + mcore_generation_config: + image_dynamic_resolution: true + megatron_inference_wrapper: megatron.core.inference.model_inference_wrappers.multimodal.nemotron_omni_inference_wrapper.NemotronOmniInferenceWrapper + multimodal_prompt_config: + image_spec: &media_prompt_spec + model_token: + prefix: + suffix: + video_spec: *media_prompt_spec 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..85a9b1cc61b 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,14 @@ policy: max_tokens: ${policy.max_total_sequence_length} expose_http_server: true enable_prefix_caching: true + image_dynamic_resolution: true + megatron_inference_wrapper: megatron.core.inference.model_inference_wrappers.multimodal.nemotron_omni_inference_wrapper.NemotronOmniInferenceWrapper + multimodal_prompt_config: + image_spec: &media_prompt_spec + model_token: + prefix: + suffix: + video_spec: *media_prompt_spec 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 index 9b1aaaa5508..53aa16dfdb3 100644 --- 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 @@ -92,6 +92,14 @@ policy: expert_tensor_parallel_size: 1 sequence_parallel: true moe_pad_experts_for_cuda_graph_inference: true + image_dynamic_resolution: true + megatron_inference_wrapper: megatron.core.inference.model_inference_wrappers.multimodal.nemotron_omni_inference_wrapper.NemotronOmniInferenceWrapper + multimodal_prompt_config: + image_spec: &media_prompt_spec + model_token: + prefix: + suffix: + video_spec: *media_prompt_spec vllm_cfg: async_engine: true expose_http_server: true diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.yaml index 34bbd2a92c7..571e6a9a654 100644 --- a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.yaml +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.yaml @@ -82,6 +82,14 @@ policy: expert_tensor_parallel_size: 1 sequence_parallel: true moe_pad_experts_for_cuda_graph_inference: true + image_dynamic_resolution: true + megatron_inference_wrapper: megatron.core.inference.model_inference_wrappers.multimodal.nemotron_omni_inference_wrapper.NemotronOmniInferenceWrapper + multimodal_prompt_config: + image_spec: &media_prompt_spec + model_token: + prefix: + suffix: + video_spec: *media_prompt_spec colocated: enabled: false resources: 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 index 73ffb6b771f..d12d447e87f 100644 --- 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 @@ -92,6 +92,14 @@ policy: expert_tensor_parallel_size: 1 sequence_parallel: true moe_pad_experts_for_cuda_graph_inference: true + image_dynamic_resolution: true + megatron_inference_wrapper: megatron.core.inference.model_inference_wrappers.multimodal.nemotron_omni_inference_wrapper.NemotronOmniInferenceWrapper + multimodal_prompt_config: + image_spec: &media_prompt_spec + model_token: + prefix: + suffix: + video_spec: *media_prompt_spec colocated: enabled: false resources: diff --git a/nemo_rl/data/multimodal_utils.py b/nemo_rl/data/multimodal_utils.py index e4a670b3570..83922c71a68 100644 --- a/nemo_rl/data/multimodal_utils.py +++ b/nemo_rl/data/multimodal_utils.py @@ -901,14 +901,14 @@ def _encode_single_image_source(source: str) -> str: return data_url -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 @@ -918,36 +918,19 @@ 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"): - continue - source = part.get("image") or part.get("image_url") or part.get("url") - if isinstance(source, dict): - source = source.get("url") - if isinstance(source, (str, Image.Image)): - sources.append(source) - return sources - - -def extract_input_video_sources_from_responses_messages(messages: Any) -> list[Any]: - """Extract video sources from Responses-API messages in encounter order.""" - if not isinstance(messages, list): - return [] - - sources: list[Any] = [] - for message in messages: - if not isinstance(message, dict): - continue - content = message.get("content") or [] - if not isinstance(content, list): - continue - for part in content: - if not isinstance(part, dict) or part.get("type") not in VIDEO_CONTENT_TYPES: + 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("video") or part.get("video_url") or part.get("url") if isinstance(source, dict): - source = source.get("url") + source = source.get("url") or source.get("path") if source is not None: - sources.append(source) + sources.append((media_type, source)) return sources @@ -957,7 +940,10 @@ def extract_input_images_from_responses_messages( """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) + for media_type, source in extract_input_media_sources_from_responses_messages( + messages + ) + if media_type == "image" ] diff --git a/nemo_rl/data/processors.py b/nemo_rl/data/processors.py index b52405b10dc..a37c3c1c6b4 100644 --- a/nemo_rl/data/processors.py +++ b/nemo_rl/data/processors.py @@ -675,9 +675,8 @@ def vlm_hf_data_processor( chat_message[key] = PackedTensor.empty_like(value) loss_multiplier = 0.0 else: - # Placeholder-style processors have already expanded media positions in - # input_ids. Let vLLM consume those IDs while retaining main's native - # media side channels; other processors keep the formatted-string path. + # 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": None if uses_placeholder else string_formatted_dialog, "vllm_images": images, diff --git a/nemo_rl/environments/nemo_gym.py b/nemo_rl/environments/nemo_gym.py index 761fc2ea856..3d080c2c98c 100644 --- a/nemo_rl/environments/nemo_gym.py +++ b/nemo_rl/environments/nemo_gym.py @@ -29,8 +29,7 @@ from nemo_rl.data.multimodal_utils import ( attach_image_model_inputs_to_message, - extract_input_image_sources_from_responses_messages, - extract_input_video_sources_from_responses_messages, + extract_input_media_sources_from_responses_messages, normalize_media_in_examples, resolve_to_image, uses_image_placeholder, @@ -407,45 +406,25 @@ def _index_per_turn_images( return per_turn -def _image_sources_equal(left: Any, right: Any) -> bool: - return ( - left == right - if isinstance(left, str) and isinstance(right, str) - else left is right - ) - - -def _input_media_sources(messages: Any) -> list[tuple[str, Any]]: - if not isinstance(messages, list): - return [] - - sources: list[tuple[str, Any]] = [] - for message in messages: - if not isinstance(message, dict): - continue - content = message.get("content") - if not isinstance(content, list): - continue - for part in content: - part_message = [{"content": [part]}] - images = extract_input_image_sources_from_responses_messages(part_message) - videos = extract_input_video_sources_from_responses_messages(part_message) - sources.extend(("image", source) for source in images) - sources.extend(("video", source) for source in videos) - return sources - - def _media_sources_equal( left: tuple[str, Any], right: tuple[str, Any], ) -> bool: - return left[0] == right[0] and _image_sources_equal(left[1], right[1]) + """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 _without_initial_media_sources( messages: Any, initial_sources: list[Any] ) -> tuple[Any, bool]: - """Copy Responses messages and remove one ordered copy of initial media.""" + """Copy Responses messages and remove one ordered copy of initial images and videos.""" if not isinstance(messages, list): return messages, False @@ -460,7 +439,9 @@ def _without_initial_media_sources( filtered_content = [] for part in content: - part_sources = _input_media_sources([{"content": [part]}]) + part_sources = extract_input_media_sources_from_responses_messages( + [{"content": [part]}] + ) if ( remaining_sources and len(part_sources) == 1 @@ -701,6 +682,9 @@ async def run_rollouts( timer = Timer() counts_left = Counter(row["agent_ref"]["name"] for row in 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") @@ -797,9 +781,15 @@ 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 = _input_media_sources(raw_input) - agent_initial_sources = _input_media_sources(initial_input) - returned_media_sources = _input_media_sources(media_messages) + raw_initial_sources = extract_input_media_sources_from_responses_messages( + raw_input + ) + agent_initial_sources = extract_input_media_sources_from_responses_messages( + initial_input + ) + 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) @@ -1127,9 +1117,11 @@ def setup_nemo_gym_config(config, tokenizer) -> None: 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}.") diff --git a/nemo_rl/environments/nemo_gym_video.py b/nemo_rl/environments/nemo_gym_video.py index 83d3f8bd9ea..d91a98f99a4 100644 --- a/nemo_rl/environments/nemo_gym_video.py +++ b/nemo_rl/environments/nemo_gym_video.py @@ -96,41 +96,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: diff --git a/nemo_rl/evals/eval.py b/nemo_rl/evals/eval.py index 1f572d589d9..17c335e900f 100644 --- a/nemo_rl/evals/eval.py +++ b/nemo_rl/evals/eval.py @@ -355,6 +355,7 @@ async def _run_env_eval_impl( prompt_dict = {"prompt": vllm_content} prompt_display = vllm_content else: + # Placeholder-style processors pass prompt_token_ids with media. prompt_token_ids = [] for message in message_log: token_ids = message["token_ids"] 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 ee6d0d27a72..7b77068b6a9 100644 --- a/nemo_rl/models/generation/megatron/config.py +++ b/nemo_rl/models/generation/megatron/config.py @@ -59,10 +59,12 @@ class MCoreGenerationSpecificArgs(TypedDict): mamba_inference_conv_states_dtype: NotRequired[str] # Raw media preprocessing used by the OpenAI-compatible HTTP endpoint. + image_dynamic_resolution: NotRequired[bool] video_num_frames: NotRequired[int] video_temporal_patch_size: NotRequired[int] video_target_num_patches: NotRequired[int] multimodal_prompt_config: NotRequired[dict] + megatron_inference_wrapper: NotRequired[str] # KV cache lifecycle across suspend/resume: # - "persist": cache stays allocated; CUDA graphs remain valid (default) diff --git a/nemo_rl/models/generation/megatron/megatron_generation.py b/nemo_rl/models/generation/megatron/megatron_generation.py index d84248e6c6d..78e0919c994 100644 --- a/nemo_rl/models/generation/megatron/megatron_generation.py +++ b/nemo_rl/models/generation/megatron/megatron_generation.py @@ -232,10 +232,14 @@ def __init__( "policy; when colocated, pass it to the training policy instead." ) - # Dedicated inference may modify its policy config. + # `self.cfg` exposes the `generation` that matches the `GenerationInterface` contract. + # `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 if policy is not None else deepcopy(config) self.cfg: MCoreGenerationConfig = config["generation"] + # Populated after the first prepare_for_generation (which starts the HTTP server). self.dp_openai_server_base_urls: list[Optional[str]] = [] + # Installed by setup via create_weight_synchronizer. self.weight_synchronizer: Optional["WeightSynchronizer"] = None if policy is not None: diff --git a/nemo_rl/models/generation/megatron/megatron_worker.py b/nemo_rl/models/generation/megatron/megatron_worker.py index 969ebef4041..f250415d61a 100644 --- a/nemo_rl/models/generation/megatron/megatron_worker.py +++ b/nemo_rl/models/generation/megatron/megatron_worker.py @@ -19,7 +19,6 @@ import threading import time import warnings -from dataclasses import replace from typing import Any, AsyncGenerator, Optional import requests @@ -30,7 +29,6 @@ KVCacheManagementMode, MambaInferenceStateConfig, PrefixCachingCoordinatorPolicy, - VideoProcessingConfig, ) from megatron.core.inference.engines.dynamic_engine import EngineState from megatron.core.inference.sampling_params import SamplingParams @@ -57,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, @@ -65,12 +64,10 @@ ) from nemo_rl.models.generation.megatron.utils import ( build_image_preprocessing_config, + build_video_preprocessing_config, log_gpu_memory, resolve_torch_dtype, ) -from nemo_rl.models.generation.vllm.video_utils import ( - _CACHED_VIDEO_FRAME_MANIFEST_MAGIC, -) from nemo_rl.models.megatron.memory_saver import ( HAVE_TORCH_MEMORY_SAVER, pause_inference_weights, @@ -346,22 +343,11 @@ def _initialize_inference_engine(self, mcore_generation_config: dict) -> None: image_preprocessing_config = self._build_image_preprocessing_config( mcore_generation_config ) - video_preprocessing_config = None - temporal_patch_size = mcore_generation_config.get("video_temporal_patch_size") - if image_preprocessing_config is not None and temporal_patch_size is not None: - video_image_config = image_preprocessing_config - target_num_patches = mcore_generation_config.get("video_target_num_patches") - if target_num_patches is not None: - video_image_config = replace( - video_image_config, - dynamic_resolution_max_patches=int(target_num_patches), - ) - video_preprocessing_config = VideoProcessingConfig( - image_config=video_image_config, - num_frames=int(mcore_generation_config["video_num_frames"]), - temporal_patch_size=int(temporal_patch_size), - frame_manifest_magic=_CACHED_VIDEO_FRAME_MANIFEST_MAGIC, - ) + video_preprocessing_config = build_video_preprocessing_config( + image_preprocessing_config, + mcore_generation_config, + frame_manifest_magic=CACHED_VIDEO_FRAME_MANIFEST_MAGIC, + ) inference_config = InferenceConfig( block_size_tokens=block_size_tokens, diff --git a/nemo_rl/models/generation/megatron/utils.py b/nemo_rl/models/generation/megatron/utils.py index 0d83f9f7529..f2a4b926adc 100644 --- a/nemo_rl/models/generation/megatron/utils.py +++ b/nemo_rl/models/generation/megatron/utils.py @@ -12,10 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dataclasses import replace from typing import Any import torch -from megatron.core.inference.config import ImageProcessingConfig +from megatron.core.inference.config import ImageProcessingConfig, VideoProcessingConfig from megatron.core.inference.utils import device_memory_summary @@ -76,6 +77,32 @@ def read(*names: str) -> Any: ) +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.""" + temporal_patch_size = generation_config.get("video_temporal_patch_size") + if image_config is None or temporal_patch_size is None: + return None + + 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(generation_config["video_num_frames"]), + temporal_patch_size=int(temporal_patch_size), + frame_manifest_magic=frame_manifest_magic, + ) + + def resolve_torch_dtype(val): """Convert a value to `torch.dtype`.""" if isinstance(val, torch.dtype): diff --git a/nemo_rl/models/generation/vllm/video_utils.py b/nemo_rl/models/generation/vllm/video_utils.py index c573a95d462..90b5637585f 100644 --- a/nemo_rl/models/generation/vllm/video_utils.py +++ b/nemo_rl/models/generation/vllm/video_utils.py @@ -24,15 +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( num_frames: int, *, @@ -197,11 +198,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 +211,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/tests/unit/data/datasets/test_mmpr_tiny.py b/tests/unit/data/datasets/test_mmpr_tiny.py index 3365f8e8572..bd536cfd02f 100644 --- a/tests/unit/data/datasets/test_mmpr_tiny.py +++ b/tests/unit/data/datasets/test_mmpr_tiny.py @@ -285,7 +285,7 @@ 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) + result, processor = _run_processor(tiny_image_path) processed_text = processor.captured_call_text # Positive: literal \boxed{} must survive prompt formatting @@ -310,7 +310,8 @@ 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. - # Placeholder-style processors send their expanded token IDs to vLLM. + # 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 diff --git a/tests/unit/data/test_multimodal_image_encoding.py b/tests/unit/data/test_multimodal_image_encoding.py index fda52a35ca6..9b150bcefb2 100644 --- a/tests/unit/data/test_multimodal_image_encoding.py +++ b/tests/unit/data/test_multimodal_image_encoding.py @@ -19,8 +19,8 @@ import nemo_rl.data.multimodal_utils as multimodal_utils from nemo_rl.data.multimodal_utils import ( - encode_images_in_examples, image_to_data_url, + normalize_media_in_examples, resolve_to_image, ) @@ -59,7 +59,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 +70,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 +81,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 +90,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" @@ -192,15 +192,15 @@ def fail_to_encode(_): 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/environments/test_nemo_gym.py b/tests/unit/environments/test_nemo_gym.py index 8f5a8e510ec..a5c7808a392 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, + normalize_media_in_examples, ) from nemo_rl.data.utils import setup_response_data from nemo_rl.distributed.ray_actor_environment_registry import ( @@ -50,7 +51,6 @@ _inject_vllm_mm_processor_kwargs, _metadata_extra_body, nemo_gym_example_to_video_datum_spec, - normalize_video_urls_in_examples, ) from nemo_rl.environments.nemotron_utils import ( _expand_nemotron_video_placeholders, @@ -184,7 +184,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 +205,12 @@ 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_extract_static_video_message_rejects_multiple_videos(tmp_path): 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": { From 990df4898ffed7382327a7dba4c01406a2f1d1cf Mon Sep 17 00:00:00 2001 From: Cory Ye Date: Fri, 21 Aug 2026 20:17:22 -0700 Subject: [PATCH 03/15] Pass pre-expanded media tokens to MLLM. Signed-off-by: Cory Ye --- nemo_rl/data/multimodal_utils.py | 131 ++----- nemo_rl/data/processors.py | 2 +- nemo_rl/environments/nemo_gym.py | 190 +---------- ...mo_gym_video.py => nemo_gym_multimodal.py} | 323 +++++++++++++----- nemo_rl/environments/nemo_gym_request.py | 91 +++++ nemo_rl/models/generation/megatron/config.py | 2 + .../generation/megatron/megatron_worker.py | 11 +- .../data/test_multimodal_image_encoding.py | 4 +- tests/unit/environments/test_nemo_gym.py | 18 +- .../test_nemo_gym_image_placeholders.py | 16 +- .../environments/test_nemo_gym_mm_utils.py | 2 +- .../generation/test_megatron_generation.py | 62 ++++ 12 files changed, 455 insertions(+), 397 deletions(-) rename nemo_rl/environments/{nemo_gym_video.py => nemo_gym_multimodal.py} (69%) create mode 100644 nemo_rl/environments/nemo_gym_request.py diff --git a/nemo_rl/data/multimodal_utils.py b/nemo_rl/data/multimodal_utils.py index 83922c71a68..c5c6ab95604 100644 --- a/nemo_rl/data/multimodal_utils.py +++ b/nemo_rl/data/multimodal_utils.py @@ -19,7 +19,6 @@ 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 @@ -41,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 " @@ -891,14 +889,15 @@ 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_media_sources_from_responses_messages( @@ -934,17 +933,19 @@ def extract_input_media_sources_from_responses_messages( 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 media_type, source in extract_input_media_sources_from_responses_messages( - messages - ) - if media_type == "image" - ] +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( @@ -1088,94 +1089,6 @@ def video_path_to_data_url(video_path: str) -> str: return f"data:video/{mime};base64,{encoded}" -def normalize_media_in_examples(nemo_gym_examples: list[dict]) -> list[dict]: - """Canonicalize Gym media parts and inline each unique local source.""" - image_targets_by_source: dict[str, list[dict]] = {} - - 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: - image_targets_by_source.setdefault(url, []).append(part) - else: - url = video_path_to_data_url(url) - - for key in source_keys: - if key != canonical_key: - part.pop(key, None) - part["type"] = canonical_type - part[canonical_key] = url - if is_image and nested_detail is not None: - part.setdefault("detail", nested_detail) - - sources = list(image_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, - ) - ) - - # 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 image_targets_by_source.items(): - data_url = encoded_by_source[source] - for part in targets: - part["image_url"] = data_url - return nemo_gym_examples - - -# Backward-compatible alias for existing callers. -encode_images_in_examples = normalize_media_in_examples - - def get_media_from_message(message: dict[str, Any]) -> dict[str, list[Any]]: """Get all media from a message log item.""" # Handle None or missing content (e.g., assistant messages with only tool_calls) diff --git a/nemo_rl/data/processors.py b/nemo_rl/data/processors.py index a37c3c1c6b4..e4ea9f198d0 100644 --- a/nemo_rl/data/processors.py +++ b/nemo_rl/data/processors.py @@ -800,7 +800,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/environments/nemo_gym.py b/nemo_rl/environments/nemo_gym.py index 3d080c2c98c..80307cb56df 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, extract_input_media_sources_from_responses_messages, - normalize_media_in_examples, - resolve_to_image, + media_sources_equal, uses_image_placeholder, ) from nemo_rl.distributed.ray_actor_environment_registry import get_actor_python_env @@ -42,6 +39,11 @@ _get_node_ip_local, ) from nemo_rl.environments.interfaces import EnvironmentInterface +from nemo_rl.environments.nemo_gym_multimodal import ( + _index_per_turn_images, + _without_initial_media_sources, + normalize_media_in_examples, +) from nemo_rl.experience.failures import ( GymTransportError, RolloutDataFailure, @@ -282,28 +284,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,40 +302,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. @@ -366,120 +312,6 @@ def _is_trainable_output_item(item: dict) -> bool: 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 _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 _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 - - -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.""" @@ -507,7 +339,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. @@ -515,7 +347,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: @@ -794,7 +626,7 @@ def _postprocess_nemo_gym_to_nemo_rl_result( bool(raw_initial_sources) and len(agent_initial_sources) == len(raw_initial_sources) and all( - _media_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 ) @@ -803,7 +635,7 @@ def _postprocess_nemo_gym_to_nemo_rl_result( returned_media_matches_raw_input = len(returned_media_sources) == len( raw_initial_sources ) and all( - _media_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 ) @@ -911,7 +743,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, diff --git a/nemo_rl/environments/nemo_gym_video.py b/nemo_rl/environments/nemo_gym_multimodal.py similarity index 69% rename from nemo_rl/environments/nemo_gym_video.py rename to nemo_rl/environments/nemo_gym_multimodal.py index d91a98f99a4..f81b4021297 100644 --- a/nemo_rl/environments/nemo_gym_video.py +++ b/nemo_rl/environments/nemo_gym_multimodal.py @@ -15,6 +15,8 @@ import copy import json import os +from concurrent.futures import ThreadPoolExecutor +from copy import deepcopy from pathlib import Path from typing import Any, TypeVar, cast from urllib.parse import unquote, urlparse @@ -29,8 +31,18 @@ VIDEO_CONTENT_TYPES, PackedTensor, extract_multimodal_model_inputs, + extract_input_media_sources_from_responses_messages, 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 +55,220 @@ 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.""" + image = resolve_to_image(source) + try: + return image_to_data_url(image) + finally: + image.close() + + +def normalize_media_in_examples(nemo_gym_examples: list[dict]) -> list[dict]: + """Replace local media paths in NeMo Gym examples with data URLs.""" + image_targets_by_source: dict[str, list[dict]] = {} + 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: + image_targets_by_source.setdefault(url, []).append(part) + else: + url = video_path_to_data_url(url) + + for key in source_keys: + if key != canonical_key: + part.pop(key, None) + part["type"] = canonical_type + part[canonical_key] = url + if is_image and nested_detail is not None: + part.setdefault("detail", nested_detail) + + sources = list(image_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, + ) + ) + for source, targets in image_targets_by_source.items(): + data_url = encoded_by_source[source] + for part in targets: + part["image_url"] = data_url + 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 _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 _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 +290,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) @@ -133,7 +348,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) @@ -142,7 +359,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." @@ -214,80 +433,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], @@ -352,7 +497,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..44465b05d7b --- /dev/null +++ b/nemo_rl/environments/nemo_gym_request.py @@ -0,0 +1,91 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import copy +import json +from typing import Any + + +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 diff --git a/nemo_rl/models/generation/megatron/config.py b/nemo_rl/models/generation/megatron/config.py index 7b77068b6a9..41d8a91cb46 100644 --- a/nemo_rl/models/generation/megatron/config.py +++ b/nemo_rl/models/generation/megatron/config.py @@ -44,6 +44,8 @@ class MCoreGenerationSpecificArgs(TypedDict): # - 'block': graphs are owned at the enclosing block (TransformerBlock / HybridBlock). # Only meaningful when cuda_graph_impl='local'. inference_cuda_graph_scope: NotRequired[str] + # Required for EP>1 + local CUDA graphs. + moe_pad_experts_for_cuda_graph_inference: NotRequired[bool] materialize_only_last_token_logits: bool enable_chunked_prefill: bool diff --git a/nemo_rl/models/generation/megatron/megatron_worker.py b/nemo_rl/models/generation/megatron/megatron_worker.py index f250415d61a..f35974f2489 100644 --- a/nemo_rl/models/generation/megatron/megatron_worker.py +++ b/nemo_rl/models/generation/megatron/megatron_worker.py @@ -805,6 +805,12 @@ def _build_prompt_and_multimodal_data(self, data, index: int): length = int(data["input_lengths"][index].item()) prompt = data["input_ids"][index, :length].tolist() imgs, imgs_sizes, num_frames = self._sample_vision_tensors(data, index) + media_cache_keys = data.get("media_cache_key") + media_cache_key = ( + media_cache_keys[index] if media_cache_keys is not None else None + ) + if media_cache_key is not None and not isinstance(media_cache_key, str): + raise TypeError("media_cache_key entries must be strings or None.") if imgs is None: return prompt, None @@ -831,10 +837,13 @@ def _build_prompt_and_multimodal_data(self, data, index: int): } else: modality_data = {"imgs": imgs, "imgs_sizes": imgs_sizes} - return prompt, { + multi_modal_data: dict[str, Any] = { modality: modality_data, "media_tokens_preexpanded": True, } + if media_cache_key is not None: + multi_modal_data["media_cache_key"] = media_cache_key + return prompt, multi_modal_data def _prepare_data_for_generation( self, data: BatchedDataDict[GenerationDatumSpec], greedy: bool = False diff --git a/tests/unit/data/test_multimodal_image_encoding.py b/tests/unit/data/test_multimodal_image_encoding.py index 9b150bcefb2..9706a13dd82 100644 --- a/tests/unit/data/test_multimodal_image_encoding.py +++ b/tests/unit/data/test_multimodal_image_encoding.py @@ -20,9 +20,11 @@ import nemo_rl.data.multimodal_utils as multimodal_utils from nemo_rl.data.multimodal_utils import ( image_to_data_url, - normalize_media_in_examples, resolve_to_image, ) +from nemo_rl.environments.nemo_gym_multimodal import ( + normalize_media_in_examples, +) def _example(*content_parts: dict) -> dict: diff --git a/tests/unit/environments/test_nemo_gym.py b/tests/unit/environments/test_nemo_gym.py index a5c7808a392..32736d803c2 100644 --- a/tests/unit/environments/test_nemo_gym.py +++ b/tests/unit/environments/test_nemo_gym.py @@ -32,7 +32,6 @@ MULTIMODAL_CONTENT_TYPES, PackedTensor, image_to_data_url, - normalize_media_in_examples, ) from nemo_rl.data.utils import setup_response_data from nemo_rl.distributed.ray_actor_environment_registry import ( @@ -46,12 +45,13 @@ 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_media_in_examples, ) +from nemo_rl.environments.nemo_gym_request import _metadata_extra_body from nemo_rl.environments.nemotron_utils import ( _expand_nemotron_video_placeholders, _flatten_nemotron_video_frame_messages, @@ -378,7 +378,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}, @@ -511,7 +511,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( @@ -572,7 +572,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"), ) @@ -659,7 +659,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}, @@ -817,11 +817,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), 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..cdc4e19f67e 100644 --- a/tests/unit/environments/test_nemo_gym_mm_utils.py +++ b/tests/unit/environments/test_nemo_gym_mm_utils.py @@ -15,7 +15,7 @@ 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.environments.nemo_gym_multimodal import ( _extract_input_images_from_message, _index_per_turn_images, ) diff --git a/tests/unit/models/generation/test_megatron_generation.py b/tests/unit/models/generation/test_megatron_generation.py index 95bce59081d..c45a959a10f 100644 --- a/tests/unit/models/generation/test_megatron_generation.py +++ b/tests/unit/models/generation/test_megatron_generation.py @@ -28,6 +28,7 @@ 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.policy import PolicyConfig from nemo_rl.models.policy.lm_policy import Policy from nemo_rl.weight_sync.megatron_weight_synchronizer import ( @@ -37,6 +38,67 @@ model_name = "Qwen/Qwen3-0.6B" + +def test_direct_megatron_media_request_preserves_preexpanded_prompt(): + worker = object.__new__(MegatronGenerationMixin) + worker.cfg = {"generation": {"mcore_generation_config": {}}} + + def sample_vision_tensors(data, index): + return torch.ones(1, 2, 4), torch.tensor([[2, 2]]), None + + def get_wrapper_cls(): + return object + + def wrapper_supports_modality(wrapper_cls, modality): + return modality == "image" + + worker._sample_vision_tensors = sample_vision_tensors + worker._get_megatron_inference_wrapper_cls = get_wrapper_cls + worker._wrapper_supports_modality = wrapper_supports_modality + data = { + "input_ids": torch.tensor([[10, 99, 99, 20, 0]]), + "input_lengths": torch.tensor([4]), + } + + prompt, multi_modal_data = worker._build_prompt_and_multimodal_data(data, 0) + + assert prompt == [10, 99, 99, 20] + assert multi_modal_data["media_tokens_preexpanded"] is True + assert "image" in multi_modal_data + + +def test_direct_megatron_video_request_marks_preexpanded_prompt(): + worker = object.__new__(MegatronGenerationMixin) + worker.cfg = {"generation": {"mcore_generation_config": {}}} + + def sample_vision_tensors(data, index): + return ( + torch.ones(1, 4, 4), + torch.tensor([[2, 2], [2, 2], [2, 2], [2, 2]]), + torch.tensor([4]), + ) + + def get_wrapper_cls(): + return object + + def wrapper_supports_modality(wrapper_cls, modality): + return modality == "video" + + worker._sample_vision_tensors = sample_vision_tensors + worker._get_megatron_inference_wrapper_cls = get_wrapper_cls + worker._wrapper_supports_modality = wrapper_supports_modality + data = { + "input_ids": torch.tensor([[10, 99, 99, 20]]), + "input_lengths": torch.tensor([4]), + } + + prompt, multi_modal_data = worker._build_prompt_and_multimodal_data(data, 0) + + assert prompt == [10, 99, 99, 20] + assert multi_modal_data["media_tokens_preexpanded"] is True + assert "video" in multi_modal_data + + basic_megatron_test_config: PolicyConfig = { "model_name": model_name, "tokenizer": {"name": model_name}, From b7b19f7ac6a495e480680c5e5d2a5c838d04c3e6 Mon Sep 17 00:00:00 2001 From: Cory Ye Date: Fri, 21 Aug 2026 21:44:20 -0700 Subject: [PATCH 04/15] Remove scripts from branch. Signed-off-by: Cory Ye --- scripts/add_nemo_gym_agent_ref.py | 74 --- scripts/prepare_nemotron_omni_vstat.py | 144 ----- ...emotron_omni_circle_count_megatron_1n4g.sh | 387 ------------- .../run_nemotron_omni_clevr_megatron_1n4g.sh | 324 ----------- .../run_nemotron_omni_vstat_megatron_1n4g.sh | 245 --------- ...ubmit_nemotron_omni_clevr_megatron_8n4g.sh | 422 -------------- ...ubmit_nemotron_omni_vstat_megatron_8n4g.sh | 519 ------------------ 7 files changed, 2115 deletions(-) delete mode 100644 scripts/add_nemo_gym_agent_ref.py delete mode 100644 scripts/prepare_nemotron_omni_vstat.py delete mode 100755 scripts/run_nemotron_omni_circle_count_megatron_1n4g.sh delete mode 100755 scripts/run_nemotron_omni_clevr_megatron_1n4g.sh delete mode 100755 scripts/run_nemotron_omni_vstat_megatron_1n4g.sh delete mode 100755 scripts/submit_nemotron_omni_clevr_megatron_8n4g.sh delete mode 100755 scripts/submit_nemotron_omni_vstat_megatron_8n4g.sh diff --git a/scripts/add_nemo_gym_agent_ref.py b/scripts/add_nemo_gym_agent_ref.py deleted file mode 100644 index 95d81dd8f10..00000000000 --- a/scripts/add_nemo_gym_agent_ref.py +++ /dev/null @@ -1,74 +0,0 @@ -#!/usr/bin/env python3 -# 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. - -"""Stamp an ``agent_ref`` onto NeMo-Gym JSONL rows that are missing one. - -NeMo-RL routes every row to a Gym agent by reading ``agent_ref.name`` -(``nemo_rl/environments/nemo_gym.py``), so a row without it fails the rollout -batch with ``KeyError: 'agent_ref'``. Some Gym generators omit the field even -though the ``example.jsonl`` they ship beside declares it — for instance -``resources_servers/circle_count/generate_data.py``. - -Rewrites files in place and leaves rows that already declare an ``agent_ref`` -untouched, so it is safe to re-run. - -Usage: - python scripts/add_nemo_gym_agent_ref.py \\ - --agent-name circle_count_simple_agent \\ - path/to/train.jsonl path/to/val.jsonl -""" - -import argparse -import json -from pathlib import Path - - -def main() -> None: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument( - "paths", type=Path, nargs="+", help="JSONL files to rewrite in place." - ) - parser.add_argument( - "--agent-name", - required=True, - help=( - "Agent key from the Gym resources server config, e.g. " - "circle_count_simple_agent." - ), - ) - parser.add_argument( - "--agent-type", - default="responses_api_agents", - help="Gym agent type. Only responses_api_agents exists today.", - ) - args = parser.parse_args() - - agent_ref = {"type": args.agent_type, "name": args.agent_name} - - for path in args.paths: - rows = [ - json.loads(line) for line in path.read_text().splitlines() if line.strip() - ] - stamped = 0 - for row in rows: - if "agent_ref" not in row: - row["agent_ref"] = agent_ref - stamped += 1 - path.write_text("".join(json.dumps(row) + "\n" for row in rows)) - print(f"{path}: stamped {stamped} of {len(rows)} row(s)") - - -if __name__ == "__main__": - main() diff --git a/scripts/prepare_nemotron_omni_vstat.py b/scripts/prepare_nemotron_omni_vstat.py deleted file mode 100644 index e67b7079e2d..00000000000 --- a/scripts/prepare_nemotron_omni_vstat.py +++ /dev/null @@ -1,144 +0,0 @@ -#!/usr/bin/env python3 -"""Download and prepare a small VSTAT split for Nemotron Omni video GRPO.""" - -import argparse -import json -import subprocess -import sys -import zipfile -from pathlib import Path - -import pyarrow.parquet as pq -from huggingface_hub import hf_hub_download - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser() - parser.add_argument( - "--output-dir", - type=Path, - required=True, - help="Directory for extracted media and source JSONL files.", - ) - parser.add_argument( - "--repo-id", - default="ShushengYang/VSTAT", - help="Hugging Face dataset repository.", - ) - parser.add_argument( - "--num-rows", - type=int, - default=8, - help="Maximum number of MCQ examples to prepare.", - ) - return parser.parse_args() - - -def write_jsonl(path: Path, rows: list[dict]) -> None: - with path.open("w", encoding="utf-8") as stream: - for row in rows: - stream.write(json.dumps(row) + "\n") - - -def main() -> None: - args = parse_args() - if args.num_rows < 3: - raise ValueError("--num-rows must be at least 3.") - - root = args.output_dir.expanduser().resolve() - root.mkdir(parents=True, exist_ok=True) - - parquet_path = hf_hub_download( - repo_id=args.repo_id, - repo_type="dataset", - filename="test.parquet", - ) - archive_path = hf_hub_download( - repo_id=args.repo_id, - repo_type="dataset", - filename="videos.zip", - ) - - media_root = root / "media" - marker = media_root / ".extracted" - if not marker.exists(): - media_root.mkdir(parents=True, exist_ok=True) - with zipfile.ZipFile(archive_path) as archive: - archive.extractall(media_root) - marker.touch() - - video_files = { - path.name: path.resolve() for path in media_root.rglob("*.mp4") - } - converted = [] - for row in pq.read_table(parquet_path).to_pylist(): - if str(row.get("answer_type", "")).lower() != "mcq": - continue - - relative_video = str(row["video"]) - video_path = media_root / relative_video - if not video_path.exists(): - video_path = video_files.get(Path(relative_video).name) - if video_path is None or not video_path.exists(): - continue - - choices = [str(choice) for choice in row.get("choices") or []] - if not choices: - continue - letters = "ABCDEFGHIJKLMNOPQRSTUVWXYZ" - options = "\n".join( - f"{letters[index]}. {choice}" - for index, choice in enumerate(choices) - ) - converted.append( - { - "prompt": ( - "Answer the multiple-choice question using the video. " - "Return the final answer as a boxed letter.\n" - f"Question: {row['question']}\n{options}" - ), - "video": str(video_path), - "answer": str(row["answer"]).strip().upper(), - "verifier": "multiple-choice", - } - ) - if len(converted) >= args.num_rows: - break - - if len(converted) < 3: - raise RuntimeError( - f"Only resolved {len(converted)} VSTAT MCQ rows with local videos." - ) - - split = max(2, len(converted) - 2) - train_source = root / "train-source.jsonl" - val_source = root / "val-source.jsonl" - write_jsonl(train_source, converted[:split]) - write_jsonl(val_source, converted[split:]) - - repo_root = Path(__file__).resolve().parents[1] - converter = repo_root / "examples/nemo_gym/prepare_video_dataset.py" - for source, output in ( - (train_source, root / "train-gym.jsonl"), - (val_source, root / "val-gym.jsonl"), - ): - subprocess.run( - [ - sys.executable, - str(converter), - "convert", - "--input", - str(source), - "--output", - str(output), - "--require-video", - ], - cwd=repo_root, - check=True, - ) - - print(f"Prepared {len(converted)} VSTAT examples under {root}") - - -if __name__ == "__main__": - main() diff --git a/scripts/run_nemotron_omni_circle_count_megatron_1n4g.sh b/scripts/run_nemotron_omni_circle_count_megatron_1n4g.sh deleted file mode 100755 index 78ec54898e7..00000000000 --- a/scripts/run_nemotron_omni_circle_count_megatron_1n4g.sh +++ /dev/null @@ -1,387 +0,0 @@ -#!/usr/bin/env bash -set -euo pipefail - -# 1-node GB200 / 4-GPU smoke for Nemotron Omni on NeMo-Gym's circle_count env. -# circle_count is an IMAGE environment: every row carries one synthetic PNG as a -# base64 data URL plus a "how many circles" question, and the Gym -# resources server rewards an exact match on \boxed{count}. -# -# Unlike the CLEVR launcher this runs the Gym entrypoint -# (examples/nemo_gym/run_grpo_nemo_gym.py): Gym owns the prompts, the reward and -# the generation loop, talking to the policy over the engine's HTTP server. -# -# GENERATION_BACKEND=megatron|vllm (default megatron) -# COLOCATED=true|false (default false) shares every GPU between the trainer and -# the engine instead of splitting them. Colocated async GRPO is only -# supported for Megatron generation. -# ASYNC_GRPO=true|false (default true) selects async or sync GRPO. -# GENERATE_ROWS=N regenerates a larger dataset instead of using the 5-row -# example.jsonl that ships with Gym. -NEMORL="${NEMORL:-/opt/nemo-rl}" -WORKSPACE_ROOT="${WORKSPACE_ROOT:-${NEMORL}/workspace}" -GYM="${GYM:-${NEMORL}/3rdparty/Gym-workspace/Gym}" -MODEL_NAME="${MODEL_NAME:-nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16}" -CONFIG="${CONFIG:-examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-count-1n4g-megatron_generation.v1.yaml}" -GENERATION_BACKEND="${GENERATION_BACKEND:-megatron}" -COLOCATED="${COLOCATED:-false}" -ASYNC_GRPO="${ASYNC_GRPO:-true}" - -cd "${NEMORL}" - -if [[ "${GENERATION_BACKEND}" != "megatron" && "${GENERATION_BACKEND}" != "vllm" ]]; then - echo "GENERATION_BACKEND must be megatron or vllm (got ${GENERATION_BACKEND})." >&2 - exit 1 -fi -if [[ "${COLOCATED}" != "true" && "${COLOCATED}" != "false" ]]; then - echo "COLOCATED must be true or false (got ${COLOCATED})." >&2 - exit 1 -fi -if [[ "${ASYNC_GRPO}" != "true" && "${ASYNC_GRPO}" != "false" ]]; then - echo "ASYNC_GRPO must be true or false (got ${ASYNC_GRPO})." >&2 - exit 1 -fi -if [[ "${COLOCATED}" == "true" && "${ASYNC_GRPO}" == "true" && "${GENERATION_BACKEND}" != "megatron" ]]; then - echo "Colocated async GRPO requires GENERATION_BACKEND=megatron." >&2 - exit 1 -fi -if [[ ! -d "${GYM}" ]]; then - echo "NeMo-Gym checkout is missing: ${GYM}" >&2 - exit 1 -fi - -GPUS_PER_NODE="${GPUS_PER_NODE:-$(nvidia-smi --query-gpu=index --format=csv,noheader | wc -l)}" -if (( GPUS_PER_NODE < 4 )); then - echo "This launcher requires at least four visible GPUs (got ${GPUS_PER_NODE})." >&2 - exit 1 -fi - -if [[ "${COLOCATED}" == "true" ]]; then - # The engine shares the trainer's GPUs and its HTTP server comes up on the - # policy workers, so both worlds span every GPU. - TRAIN_WORLD_SIZE="${GPUS_PER_NODE}" - INFERENCE_WORLD_SIZE="${GPUS_PER_NODE}" - GEN_GPUS_PER_NODE="${GPUS_PER_NODE}" - NUM_GEN_NODES=1 - COLOCATED_ENABLED=true -else - GEN_GPUS_PER_NODE="${GEN_GPUS_PER_NODE:-$((GPUS_PER_NODE / 2))}" - if (( GEN_GPUS_PER_NODE <= 0 || GEN_GPUS_PER_NODE >= GPUS_PER_NODE )); then - echo "Non-colocated mode requires a strict train/inference GPU split." >&2 - exit 1 - fi - TRAIN_WORLD_SIZE=$((GPUS_PER_NODE - GEN_GPUS_PER_NODE)) - INFERENCE_WORLD_SIZE="${GEN_GPUS_PER_NODE}" - NUM_GEN_NODES=1 - COLOCATED_ENABLED=false -fi - -POLICY_TP="${POLICY_TP:-${TRAIN_WORLD_SIZE}}" -INFER_TP="${INFER_TP:-${INFERENCE_WORLD_SIZE}}" -POLICY_CP="${POLICY_CP:-1}" -if [[ "${GENERATION_BACKEND}" == "megatron" && "${POLICY_CP}" != "1" ]]; then - echo "Megatron dynamic inference requires POLICY_CP=1." >&2 - exit 1 -fi -if (( TRAIN_WORLD_SIZE % (POLICY_TP * POLICY_CP) != 0 )); then - echo "Training world size must be divisible by POLICY_TP * POLICY_CP." >&2 - exit 1 -fi -if (( INFERENCE_WORLD_SIZE % INFER_TP != 0 )); then - echo "Inference world size must be divisible by INFER_TP." >&2 - exit 1 -fi - -# With ETP=1, world_size must be divisible by EP (not by TP*EP). -largest_ep() { - local world="$1" - local ep=8 - while (( ep > world || world % ep != 0 )); do - ep=$((ep / 2)) - done - printf '%d' "${ep}" -} - -POLICY_EP="${POLICY_EP:-$(largest_ep "${TRAIN_WORLD_SIZE}")}" -INFER_EP="${INFER_EP:-$(largest_ep "${INFERENCE_WORLD_SIZE}")}" -if (( TRAIN_WORLD_SIZE % POLICY_EP != 0 )); then - echo "Training world size must be divisible by POLICY_EP (ETP=1)." >&2 - exit 1 -fi -if [[ "${GENERATION_BACKEND}" == "vllm" ]]; then - if [[ "${INFER_EP}" != "${INFER_TP}" ]]; then - echo "Forcing INFER_EP=${INFER_TP} for vLLM (was ${INFER_EP}) for async-safe Ray DP." >&2 - fi - INFER_EP="${INFER_TP}" -elif (( INFERENCE_WORLD_SIZE % INFER_EP != 0 )); then - echo "Inference world size must be divisible by INFER_EP (ETP=1)." >&2 - exit 1 -fi -if [[ "${COLOCATED_ENABLED}" == "true" ]]; then - # Colocated Megatron generation reuses the trained model in place, so the - # engine cannot reshard: its parallel layout has to match the policy's. - if (( INFER_TP != POLICY_TP || INFER_EP != POLICY_EP )); then - echo "Colocated mode requires INFER_TP/INFER_EP to match POLICY_TP/POLICY_EP" \ - "(got TP ${INFER_TP} vs ${POLICY_TP}, EP ${INFER_EP} vs ${POLICY_EP})." >&2 - exit 1 - fi -fi - -CACHE_ROOT="${CACHE_ROOT:-${WORKSPACE_ROOT}/cache/nemo-rl-omni}" -export HF_HOME="${HF_HOME:-${CACHE_ROOT}/huggingface}" -export HUGGINGFACE_HUB_CACHE="${HUGGINGFACE_HUB_CACHE:-${HF_HOME}/hub}" -export HF_HUB_CACHE="${HF_HUB_CACHE:-${HUGGINGFACE_HUB_CACHE}}" -export HF_DATASETS_CACHE="${HF_DATASETS_CACHE:-${HF_HOME}/datasets}" -export HF_MODULES_CACHE="${HF_MODULES_CACHE:-${HF_HOME}/modules}" -export TRANSFORMERS_CACHE="${TRANSFORMERS_CACHE:-${HF_HOME}/transformers}" -export NRL_MEGATRON_CHECKPOINT_DIR="${NRL_MEGATRON_CHECKPOINT_DIR:-${CACHE_ROOT}/megatron-checkpoints}" -export XDG_CACHE_HOME="${XDG_CACHE_HOME:-${CACHE_ROOT}/xdg}" -export TORCH_HOME="${TORCH_HOME:-${CACHE_ROOT}/torch}" -export TRITON_CACHE_DIR="${TRITON_CACHE_DIR:-${CACHE_ROOT}/triton}" - -BRIDGE="${NEMORL}/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge" -export PYTHONPATH="${NEMORL}:${GYM}:${BRIDGE}/src:${BRIDGE}/3rdparty/Megatron-LM${PYTHONPATH:+:${PYTHONPATH}}" -export RAY_ENABLE_UV_RUN_RUNTIME_ENV=0 -export NRL_VENVS_TRUST_EXISTING="${NRL_VENVS_TRUST_EXISTING:-1}" -export NRL_FORCE_REBUILD_VENVS="${NRL_FORCE_REBUILD_VENVS:-false}" -export NEMO_RL_VENV_DIR="${NEMO_RL_VENV_DIR:-/opt/ray_venvs}" -# The Gym env servers run in their own venvs; keep them beside the NeMo-RL ones -# so a warm container does not rebuild them on every launch. -export NEMO_GYM_VENV_DIR="${NEMO_GYM_VENV_DIR:-${NEMO_RL_VENV_DIR}}" -export CUDA_DEVICE_MAX_CONNECTIONS="${CUDA_DEVICE_MAX_CONNECTIONS:-1}" -export NVTE_FWD_LAYERNORM_SM_MARGIN="${NVTE_FWD_LAYERNORM_SM_MARGIN:-16}" -export NVTE_BWD_LAYERNORM_SM_MARGIN="${NVTE_BWD_LAYERNORM_SM_MARGIN:-16}" -export FLASHINFER_DISABLE_VERSION_CHECK="${FLASHINFER_DISABLE_VERSION_CHECK:-1}" -export TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST:-10.0}" -export NCCL_DEBUG="${NCCL_DEBUG:-WARN}" - -MAX_STEPS="${MAX_STEPS:-5}" -MAX_SEQUENCE_LENGTH="${MAX_SEQUENCE_LENGTH:-4096}" -MAX_NEW_TOKENS="${MAX_NEW_TOKENS:-512}" -NUM_PROMPTS="${NUM_PROMPTS:-2}" -NUM_GENERATIONS="${NUM_GENERATIONS:-8}" -TRAIN_GBS="${TRAIN_GBS:-$((NUM_PROMPTS * NUM_GENERATIONS))}" -EXPECTED_TRAIN_GBS=$((NUM_PROMPTS * NUM_GENERATIONS)) -if (( TRAIN_GBS != EXPECTED_TRAIN_GBS )); then - echo "TRAIN_GBS (${TRAIN_GBS}) must equal NUM_PROMPTS * NUM_GENERATIONS (${EXPECTED_TRAIN_GBS})." >&2 - exit 1 -fi -REFIT_BACKEND="${REFIT_BACKEND:-nccl}" -JOB_NAME="${JOB_NAME:-nemotron-omni-circle-count-${GENERATION_BACKEND}-1n4g}" -EXP_NAME="${EXP_NAME:-${JOB_NAME}}" -PRECISION_RECIPE="${PRECISION_RECIPE:-bf16}" -WANDB_ENABLED="${WANDB_ENABLED:-false}" -WANDB_PROJ="${WANDB_PROJ:-mllm-rl-dev}" -WANDB_GROUP="${WANDB_GROUP:-adlr}" -WANDB_NAME="${WANDB_NAME:-${EXP_NAME}-${PRECISION_RECIPE}-internal-repo}" -RESULTS_DIR="${RESULTS_DIR:-${WORKSPACE_ROOT}/results/nemo-rl-omni/${JOB_NAME}}" -CHECKPOINTING_ENABLED="${CHECKPOINTING_ENABLED:-false}" -# Host OOM on GB200 when optimizer CPU offload is enabled for this model size. -OPTIMIZER_CPU_OFFLOAD="${OPTIMIZER_CPU_OFFLOAD:-false}" -OFFLOAD_OPTIMIZER_FOR_LOGPROB="${OFFLOAD_OPTIMIZER_FOR_LOGPROB:-false}" -if [[ "${OPTIMIZER_CPU_OFFLOAD}" == "true" ]]; then - OPTIMIZER_OFFLOAD_FRACTION="${OPTIMIZER_OFFLOAD_FRACTION:-1.0}" -else - OPTIMIZER_OFFLOAD_FRACTION="${OPTIMIZER_OFFLOAD_FRACTION:-0.0}" -fi -# Lower-precision Adam moments for the 2-GPU train half (HBM-tight on 1n4g). -EXP_AVG_DTYPE="${EXP_AVG_DTYPE:-bfloat16}" -EXP_AVG_SQ_DTYPE="${EXP_AVG_SQ_DTYPE:-bfloat16}" -STORE_PARAM_REMAINDERS="${STORE_PARAM_REMAINDERS:-true}" -BUFFER_SIZE_GB="${BUFFER_SIZE_GB:-8}" -MAX_TRAJECTORY_AGE_STEPS="${MAX_TRAJECTORY_AGE_STEPS:-2}" -IN_FLIGHT_WEIGHT_UPDATES="${IN_FLIGHT_WEIGHT_UPDATES:-true}" -# Keep CUDA graphs disabled and use chunked prefill by default. -MEGATRON_ENABLE_CHUNKED_PREFILL="${MEGATRON_ENABLE_CHUNKED_PREFILL:-true}" -MEGATRON_CUDA_GRAPH_IMPL="${MEGATRON_CUDA_GRAPH_IMPL:-local}" # local none -MEGATRON_CUDA_GRAPH_SCOPE="${MEGATRON_CUDA_GRAPH_SCOPE:-block}" # block layer none -MEGATRON_NUM_CUDA_GRAPHS="${MEGATRON_NUM_CUDA_GRAPHS:--1}" -MEGATRON_USE_CUDA_GRAPHS_FOR_NON_DECODE="${MEGATRON_USE_CUDA_GRAPHS_FOR_NON_DECODE:-false}" -MOE_PAD_EXPERTS_FOR_CG="${MOE_PAD_EXPERTS_FOR_CG:-true}" - -# Dataset. Gym ships a 5-row example; GENERATE_ROWS=N synthesizes a larger set -# (N train rows plus a disjoint 25%-sized validation split) under DATA_ROOT. -DATA_ROOT="${DATA_ROOT:-${WORKSPACE_ROOT}/datasets/circle-count-smoke}" -GENERATE_ROWS="${GENERATE_ROWS:-64}" -GYM_CIRCLE_COUNT="${GYM}/resources_servers/circle_count" -TRAIN_JSONL="${TRAIN_JSONL:-${GYM_CIRCLE_COUNT}/data/example.jsonl}" -VAL_JSONL="${VAL_JSONL:-${TRAIN_JSONL}}" -if (( GENERATE_ROWS > 0 )); then - VAL_ROWS=$(( GENERATE_ROWS / 4 )) - (( VAL_ROWS > 0 )) || VAL_ROWS=1 - mkdir -p "${DATA_ROOT}" - TRAIN_JSONL="${DATA_ROOT}/train.jsonl" - VAL_JSONL="${DATA_ROOT}/val.jsonl" - echo "Generating ${GENERATE_ROWS} train / ${VAL_ROWS} val circle_count rows under ${DATA_ROOT}" - uv run --no-sync python "${GYM_CIRCLE_COUNT}/generate_data.py" \ - --n "${GENERATE_ROWS}" --out "${TRAIN_JSONL}" --seed-offset 0 - # Offset the seeds so validation scenes are disjoint from training ones. - uv run --no-sync python "${GYM_CIRCLE_COUNT}/generate_data.py" \ - --n "${VAL_ROWS}" --out "${VAL_JSONL}" --seed-offset "${GENERATE_ROWS}" - # generate_data.py omits the agent_ref that NeMo-RL reads off every row to - # route it to a Gym agent, unlike the example.jsonl it ships beside. The name - # is the agent key in resources_servers/circle_count/configs/circle_count.yaml. - uv run --no-sync python scripts/add_nemo_gym_agent_ref.py \ - --agent-name circle_count_simple_agent \ - "${TRAIN_JSONL}" "${VAL_JSONL}" -fi - -# NSYS: ENABLE_NSYS=true NRL_NSYS_PROFILE_STEP_RANGE=1:4 -ENABLE_NSYS="${ENABLE_NSYS:-false}" -NSYS_ENV=() -if [[ "${ENABLE_NSYS}" == "true" ]]; then - NRL_NSYS_WORKER_PATTERNS="${NRL_NSYS_WORKER_PATTERNS:-*policy*,*megatron*}" - NRL_NSYS_PROFILE_STEP_RANGE="${NRL_NSYS_PROFILE_STEP_RANGE:-1:4}" - LD_LIBRARY_PATH="/usr/local/cuda/targets/aarch64-linux/lib:/usr/local/cuda/targets/x86_64-linux/lib:/usr/local/cuda/lib64:/usr/local/cuda/lib:/usr/local/nvidia/lib64:/usr/local/nvidia/lib:/usr/lib/aarch64-linux-gnu:/usr/lib/x86_64-linux-gnu:${LD_LIBRARY_PATH:-}" - NRL_NSYS_EXTRA_OPTIONS="${NRL_NSYS_EXTRA_OPTIONS:-{\"o\":\"/opt/nemo-rl/workspace/nsys/%p\",\"cpuctxsw\":\"none\",\"force-overwrite\":\"true\"}}" - NSYS_ENV=( - "NRL_NSYS_WORKER_PATTERNS=${NRL_NSYS_WORKER_PATTERNS}" - "NRL_NSYS_PROFILE_STEP_RANGE=${NRL_NSYS_PROFILE_STEP_RANGE}" - "LD_LIBRARY_PATH=${LD_LIBRARY_PATH}" - "NRL_NSYS_EXTRA_OPTIONS=${NRL_NSYS_EXTRA_OPTIONS}" - ) - mkdir -p /opt/nemo-rl/workspace/nsys -fi - -mkdir -p \ - "${HF_HUB_CACHE}" \ - "${HF_DATASETS_CACHE}" \ - "${HF_MODULES_CACHE}" \ - "${TRANSFORMERS_CACHE}" \ - "${NRL_MEGATRON_CHECKPOINT_DIR}" \ - "${XDG_CACHE_HOME}" \ - "${TORCH_HOME}" \ - "${TRITON_CACHE_DIR}" \ - "${RESULTS_DIR}" - -if [[ ! -f "${CONFIG}" ]]; then - echo "Config is missing under ${NEMORL}: ${CONFIG}" >&2 - exit 1 -fi -for jsonl in "${TRAIN_JSONL}" "${VAL_JSONL}"; do - if [[ ! -s "${jsonl}" ]]; then - echo "Gym dataset is missing or empty: ${jsonl}" >&2 - echo "Generate one with: GENERATE_ROWS=64 $0" >&2 - exit 1 - fi -done - -COMMON_OVERRIDES=( - cluster.num_nodes=1 - cluster.gpus_per_node="${GPUS_PER_NODE}" - policy.model_name="${MODEL_NAME}" - policy.tokenizer.name="${MODEL_NAME}" - policy.is_vlm=true - policy.megatron_cfg.tensor_model_parallel_size="${POLICY_TP}" - policy.megatron_cfg.expert_model_parallel_size="${POLICY_EP}" - policy.megatron_cfg.expert_tensor_parallel_size=1 - policy.megatron_cfg.context_parallel_size="${POLICY_CP}" - policy.megatron_cfg.optimizer.optimizer_cpu_offload="${OPTIMIZER_CPU_OFFLOAD}" - policy.megatron_cfg.optimizer.optimizer_offload_fraction="${OPTIMIZER_OFFLOAD_FRACTION}" - policy.offload_optimizer_for_logprob="${OFFLOAD_OPTIMIZER_FOR_LOGPROB}" -) -if [[ -n "${EXP_AVG_DTYPE}" ]]; then - COMMON_OVERRIDES+=("++policy.megatron_cfg.optimizer.exp_avg_dtype=${EXP_AVG_DTYPE}") -fi -if [[ -n "${EXP_AVG_SQ_DTYPE}" ]]; then - COMMON_OVERRIDES+=("++policy.megatron_cfg.optimizer.exp_avg_sq_dtype=${EXP_AVG_SQ_DTYPE}") -fi -if [[ -n "${STORE_PARAM_REMAINDERS}" ]]; then - COMMON_OVERRIDES+=("++policy.megatron_cfg.optimizer.store_param_remainders=${STORE_PARAM_REMAINDERS}") -fi -COMMON_OVERRIDES+=( - policy.generation.backend="${GENERATION_BACKEND}" - policy.generation.colocated.enabled="${COLOCATED_ENABLED}" - policy.generation.colocated.resources.num_nodes="${NUM_GEN_NODES}" - policy.generation.colocated.resources.gpus_per_node="${GEN_GPUS_PER_NODE}" - policy.max_total_sequence_length="${MAX_SEQUENCE_LENGTH}" - policy.generation.max_new_tokens="${MAX_NEW_TOKENS}" - data.train.data_path="${TRAIN_JSONL}" - data.validation.data_path="${VAL_JSONL}" - grpo.async_grpo.enabled="${ASYNC_GRPO}" - grpo.async_grpo.max_trajectory_age_steps="${MAX_TRAJECTORY_AGE_STEPS}" - grpo.async_grpo.in_flight_weight_updates="${IN_FLIGHT_WEIGHT_UPDATES}" - loss_fn.use_importance_sampling_correction=true - grpo.num_prompts_per_step="${NUM_PROMPTS}" - grpo.num_generations_per_prompt="${NUM_GENERATIONS}" - grpo.val_period=0 - grpo.val_at_start=false - grpo.val_at_end=false - policy.train_global_batch_size="${TRAIN_GBS}" - grpo.max_num_steps="${MAX_STEPS}" - checkpointing.enabled="${CHECKPOINTING_ENABLED}" - checkpointing.checkpoint_dir="${RESULTS_DIR}" - logger.log_dir="${RESULTS_DIR}" - logger.wandb_enabled="${WANDB_ENABLED}" - logger.tensorboard_enabled=false - logger.wandb.name="${WANDB_NAME}" - logger.wandb.project="${WANDB_PROJ}" - +logger.wandb.entity="${WANDB_GROUP}" -) - -GEN_OVERRIDES=() -if [[ "${GENERATION_BACKEND}" == "megatron" ]]; then - GEN_OVERRIDES=( - policy.generation.mcore_generation_config.tensor_model_parallel_size="${INFER_TP}" - policy.generation.mcore_generation_config.expert_model_parallel_size="${INFER_EP}" - policy.generation.mcore_generation_config.expert_tensor_parallel_size=1 - ++policy.generation.mcore_generation_config.context_parallel_size="${POLICY_CP}" - policy.generation.mcore_generation_config.enable_chunked_prefill="${MEGATRON_ENABLE_CHUNKED_PREFILL}" - ++policy.generation.mcore_generation_config.async_sched_mode=async - policy.generation.mcore_generation_config.cuda_graph_impl="${MEGATRON_CUDA_GRAPH_IMPL}" - policy.generation.mcore_generation_config.inference_cuda_graph_scope="${MEGATRON_CUDA_GRAPH_SCOPE}" - policy.generation.mcore_generation_config.num_cuda_graphs="${MEGATRON_NUM_CUDA_GRAPHS}" - policy.generation.mcore_generation_config.use_cuda_graphs_for_non_decode_steps="${MEGATRON_USE_CUDA_GRAPHS_FOR_NON_DECODE}" - policy.generation.mcore_generation_config.refit_backend="${REFIT_BACKEND}" - policy.generation.mcore_generation_config.buffer_size_gb="${BUFFER_SIZE_GB}" - policy.generation.mcore_generation_config.moe_pad_experts_for_cuda_graph_inference="${MOE_PAD_EXPERTS_FOR_CG}" - policy.generation.mcore_generation_config.max_model_len="${MAX_SEQUENCE_LENGTH}" - policy.generation.mcore_generation_config.max_tokens="${MAX_SEQUENCE_LENGTH}" - ) -else - # Refit packs weights into chunks sized at NRL_REFIT_BUFFER_MEMORY_RATIO * - # total HBM; exported globally on purpose because producer and consumer must - # agree on the chunk boundaries. - export NRL_REFIT_BUFFER_MEMORY_RATIO="${NRL_REFIT_BUFFER_MEMORY_RATIO:-0.005}" - VLLM_GPU_MEMORY_UTILIZATION="${VLLM_GPU_MEMORY_UTILIZATION:-0.5}" - VLLM_ENFORCE_EAGER="${VLLM_ENFORCE_EAGER:-true}" - VLLM_MAX_NUM_BATCHED_TOKENS="${VLLM_MAX_NUM_BATCHED_TOKENS:-${MAX_SEQUENCE_LENGTH}}" - GEN_OVERRIDES=( - policy.generation.vllm_cfg.async_engine="${ASYNC_GRPO}" - policy.generation.vllm_cfg.skip_tokenizer_init=false - policy.generation.vllm_cfg.tensor_parallel_size="${INFER_TP}" - policy.generation.vllm_cfg.pipeline_parallel_size=1 - policy.generation.vllm_cfg.expert_parallel_size="${INFER_EP}" - policy.generation.vllm_cfg.max_model_len="${MAX_SEQUENCE_LENGTH}" - ++policy.generation.vllm_cfg.cap_max_tokens_to_context=true - policy.generation.vllm_cfg.gpu_memory_utilization="${VLLM_GPU_MEMORY_UTILIZATION}" - policy.generation.vllm_cfg.enforce_eager="${VLLM_ENFORCE_EAGER}" - ++policy.generation.vllm_cfg.enable_prefix_caching=false - policy.generation.vllm_cfg.logprobs_mode=raw_logprobs - ++policy.generation.vllm_kwargs.limit_mm_per_prompt.image=1 - ++policy.generation.vllm_kwargs.max_num_batched_tokens="${VLLM_MAX_NUM_BATCHED_TOKENS}" - ++policy.generation.vllm_kwargs.mamba_ssm_cache_dtype=float32 - ++policy.generation.vllm_kwargs.skip_mm_profiling=true - ++policy.generation.vllm_kwargs.kernel_config.enable_flashinfer_autotune=false - ++policy.generation.vllm_kwargs.kernel_config.moe_backend=triton - # vLLM bans these tokens rather than treating them as stop strings, which - # Gym clears anyway. - '++policy.generation.bad_words=["","","","","",""]' - ) -fi - -echo "Launching ${JOB_NAME}: ${GPUS_PER_NODE} visible GPU(s)" -echo " env: NeMo-Gym circle_count (images, base64 data URLs) via ${GYM}" -echo " generation backend: ${GENERATION_BACKEND} colocated=${COLOCATED_ENABLED} async=${ASYNC_GRPO}" -echo " async: max_trajectory_age=${MAX_TRAJECTORY_AGE_STEPS} in_flight_weight_updates=${IN_FLIGHT_WEIGHT_UPDATES}" -echo " training world size: ${TRAIN_WORLD_SIZE} (TP=${POLICY_TP}, EP=${POLICY_EP}, ETP=1)" -echo " inference world size: ${INFERENCE_WORLD_SIZE} (TP=${INFER_TP}, EP=${INFER_EP})" -echo " data: train=${TRAIN_JSONL} ($(wc -l < "${TRAIN_JSONL}") rows) val=${VAL_JSONL} ($(wc -l < "${VAL_JSONL}") rows)" -echo " seq/new_tokens: ${MAX_SEQUENCE_LENGTH}/${MAX_NEW_TOKENS}" -echo " optimizer moments: exp_avg=${EXP_AVG_DTYPE:-} exp_avg_sq=${EXP_AVG_SQ_DTYPE:-} store_param_remainders=${STORE_PARAM_REMAINDERS:-}" -echo " W&B: ${WANDB_GROUP}/${WANDB_PROJ}/${WANDB_NAME} (enabled=${WANDB_ENABLED})" - -exec env "${NSYS_ENV[@]}" uv run --no-sync python examples/nemo_gym/run_grpo_nemo_gym.py \ - --config "${CONFIG}" \ - "${COMMON_OVERRIDES[@]}" \ - "${GEN_OVERRIDES[@]}" \ - "$@" diff --git a/scripts/run_nemotron_omni_clevr_megatron_1n4g.sh b/scripts/run_nemotron_omni_clevr_megatron_1n4g.sh deleted file mode 100755 index bd7b257cec8..00000000000 --- a/scripts/run_nemotron_omni_clevr_megatron_1n4g.sh +++ /dev/null @@ -1,324 +0,0 @@ -#!/usr/bin/env bash -set -euo pipefail - -# 1-node GB200 / 4-GPU smoke for Nemotron Omni CLEVR. -# Default: Megatron generation, non-colocated 2/2, async GRPO, age>1, in-flight. -# -# GENERATION_BACKEND=megatron|vllm (default megatron) -# - megatron: MCore generation; colocated async supported (NVIDIA-NeMo/RL#2884) -# - vllm: async multimodal via NVIDIA-NeMo/RL#3414; colocated async not supported -# COLOCATED=true ASYNC_GRPO=true requires GENERATION_BACKEND=megatron -# COLOCATED=true ASYNC_GRPO=false runs sync colocated GRPO -# Optimizer moments default to bf16 on this 1n4g split; set EXP_AVG_DTYPE=float32 -# EXP_AVG_SQ_DTYPE=float32 STORE_PARAM_REMAINDERS=false for full precision. -NEMORL="${NEMORL:-/opt/nemo-rl}" -WORKSPACE_ROOT="${WORKSPACE_ROOT:-${NEMORL}/workspace}" -MODEL_NAME="${MODEL_NAME:-nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16}" -CONFIG="${CONFIG:-examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.yaml}" -GENERATION_BACKEND="${GENERATION_BACKEND:-megatron}" -COLOCATED="${COLOCATED:-false}" -ASYNC_GRPO="${ASYNC_GRPO:-true}" - -cd "${NEMORL}" - -if [[ "${GENERATION_BACKEND}" != "megatron" && "${GENERATION_BACKEND}" != "vllm" ]]; then - echo "GENERATION_BACKEND must be megatron or vllm (got ${GENERATION_BACKEND})." >&2 - exit 1 -fi -if [[ "${COLOCATED}" == "true" && "${ASYNC_GRPO}" == "true" && "${GENERATION_BACKEND}" != "megatron" ]]; then - echo "Colocated async GRPO requires GENERATION_BACKEND=megatron." >&2 - exit 1 -fi - -GPUS_PER_NODE="${GPUS_PER_NODE:-$(nvidia-smi --query-gpu=index --format=csv,noheader | wc -l)}" -if (( GPUS_PER_NODE < 4 )); then - echo "This launcher requires at least four visible GPUs (got ${GPUS_PER_NODE})." >&2 - exit 1 -fi - -if [[ "${COLOCATED}" == "true" ]]; then - TRAIN_WORLD_SIZE="${GPUS_PER_NODE}" - INFERENCE_WORLD_SIZE="${GPUS_PER_NODE}" - GEN_GPUS_PER_NODE="${GPUS_PER_NODE}" - NUM_GEN_NODES=1 - COLOCATED_ENABLED=true -else - GEN_GPUS_PER_NODE="${GEN_GPUS_PER_NODE:-$((GPUS_PER_NODE / 2))}" - if (( GEN_GPUS_PER_NODE <= 0 || GEN_GPUS_PER_NODE >= GPUS_PER_NODE )); then - echo "Non-colocated mode requires a strict train/inference GPU split." >&2 - exit 1 - fi - TRAIN_WORLD_SIZE=$((GPUS_PER_NODE - GEN_GPUS_PER_NODE)) - INFERENCE_WORLD_SIZE="${GEN_GPUS_PER_NODE}" - NUM_GEN_NODES=1 - COLOCATED_ENABLED=false -fi - -POLICY_TP="${POLICY_TP:-${TRAIN_WORLD_SIZE}}" -INFER_TP="${INFER_TP:-${INFERENCE_WORLD_SIZE}}" -POLICY_CP="${POLICY_CP:-1}" -if [[ "${GENERATION_BACKEND}" == "megatron" && "${POLICY_CP}" != "1" ]]; then - echo "Megatron dynamic inference requires POLICY_CP=1." >&2 - exit 1 -fi -if (( TRAIN_WORLD_SIZE % (POLICY_TP * POLICY_CP) != 0 )); then - echo "Training world size must be divisible by POLICY_TP * POLICY_CP." >&2 - exit 1 -fi -if (( INFERENCE_WORLD_SIZE % INFER_TP != 0 )); then - echo "Inference world size must be divisible by INFER_TP." >&2 - exit 1 -fi - -# With ETP=1, world_size must be divisible by EP (not by TP*EP). -largest_ep() { - local world="$1" - local ep=8 - while (( ep > world || world % ep != 0 )); do - ep=$((ep / 2)) - done - printf '%d' "${ep}" -} - -POLICY_EP="${POLICY_EP:-$(largest_ep "${TRAIN_WORLD_SIZE}")}" -INFER_EP="${INFER_EP:-$(largest_ep "${INFERENCE_WORLD_SIZE}")}" -if (( TRAIN_WORLD_SIZE % POLICY_EP != 0 )); then - echo "Training world size must be divisible by POLICY_EP (ETP=1)." >&2 - exit 1 -fi -if [[ "${GENERATION_BACKEND}" == "vllm" ]]; then - if [[ "${INFER_EP}" != "${INFER_TP}" ]]; then - echo "Forcing INFER_EP=${INFER_TP} for vLLM (was ${INFER_EP}) for async-safe Ray DP." >&2 - fi - INFER_EP="${INFER_TP}" -elif (( INFERENCE_WORLD_SIZE % INFER_EP != 0 )); then - echo "Inference world size must be divisible by INFER_EP (ETP=1)." >&2 - exit 1 -fi - -CACHE_ROOT="${CACHE_ROOT:-${WORKSPACE_ROOT}/cache/nemo-rl-omni}" -export HF_HOME="${HF_HOME:-${CACHE_ROOT}/huggingface}" -export HUGGINGFACE_HUB_CACHE="${HUGGINGFACE_HUB_CACHE:-${HF_HOME}/hub}" -export HF_HUB_CACHE="${HF_HUB_CACHE:-${HUGGINGFACE_HUB_CACHE}}" -export HF_DATASETS_CACHE="${HF_DATASETS_CACHE:-${HF_HOME}/datasets}" -export HF_MODULES_CACHE="${HF_MODULES_CACHE:-${HF_HOME}/modules}" -export TRANSFORMERS_CACHE="${TRANSFORMERS_CACHE:-${HF_HOME}/transformers}" -export NRL_MEGATRON_CHECKPOINT_DIR="${NRL_MEGATRON_CHECKPOINT_DIR:-${CACHE_ROOT}/megatron-checkpoints}" -export XDG_CACHE_HOME="${XDG_CACHE_HOME:-${CACHE_ROOT}/xdg}" -export TORCH_HOME="${TORCH_HOME:-${CACHE_ROOT}/torch}" -export TRITON_CACHE_DIR="${TRITON_CACHE_DIR:-${CACHE_ROOT}/triton}" - -BRIDGE="${NEMORL}/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge" -export PYTHONPATH="${NEMORL}:${BRIDGE}/src:${BRIDGE}/3rdparty/Megatron-LM${PYTHONPATH:+:${PYTHONPATH}}" -export RAY_ENABLE_UV_RUN_RUNTIME_ENV=0 -export NRL_VENVS_TRUST_EXISTING="${NRL_VENVS_TRUST_EXISTING:-1}" -export NRL_FORCE_REBUILD_VENVS="${NRL_FORCE_REBUILD_VENVS:-false}" -export NEMO_RL_VENV_DIR="${NEMO_RL_VENV_DIR:-/opt/ray_venvs}" -export CUDA_DEVICE_MAX_CONNECTIONS="${CUDA_DEVICE_MAX_CONNECTIONS:-1}" -export NVTE_FWD_LAYERNORM_SM_MARGIN="${NVTE_FWD_LAYERNORM_SM_MARGIN:-16}" -export NVTE_BWD_LAYERNORM_SM_MARGIN="${NVTE_BWD_LAYERNORM_SM_MARGIN:-16}" -export FLASHINFER_DISABLE_VERSION_CHECK="${FLASHINFER_DISABLE_VERSION_CHECK:-1}" -export TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST:-10.0}" -export NCCL_DEBUG="${NCCL_DEBUG:-WARN}" - -MAX_STEPS="${MAX_STEPS:-5}" -MAX_SEQUENCE_LENGTH="${MAX_SEQUENCE_LENGTH:-2048}" -MAX_NEW_TOKENS="${MAX_NEW_TOKENS:-512}" -NUM_PROMPTS="${NUM_PROMPTS:-2}" -NUM_GENERATIONS="${NUM_GENERATIONS:-8}" -TRAIN_GBS="${TRAIN_GBS:-$((NUM_PROMPTS * NUM_GENERATIONS))}" -EXPECTED_TRAIN_GBS=$((NUM_PROMPTS * NUM_GENERATIONS)) -if (( TRAIN_GBS != EXPECTED_TRAIN_GBS )); then - echo "TRAIN_GBS (${TRAIN_GBS}) must equal NUM_PROMPTS * NUM_GENERATIONS (${EXPECTED_TRAIN_GBS})." >&2 - exit 1 -fi -REFIT_BACKEND="${REFIT_BACKEND:-nccl}" -JOB_NAME="${JOB_NAME:-nemotron-omni-clevr-${GENERATION_BACKEND}-1n4g}" -EXP_NAME="${EXP_NAME:-${JOB_NAME}}" -PRECISION_RECIPE="${PRECISION_RECIPE:-bf16}" -WANDB_ENABLED="${WANDB_ENABLED:-false}" -WANDB_PROJ="${WANDB_PROJ:-mllm-rl-dev}" -WANDB_GROUP="${WANDB_GROUP:-adlr}" -WANDB_NAME="${WANDB_NAME:-${EXP_NAME}-${PRECISION_RECIPE}-internal-repo}" -RESULTS_DIR="${RESULTS_DIR:-${WORKSPACE_ROOT}/results/nemo-rl-omni/${JOB_NAME}}" -CHECKPOINTING_ENABLED="${CHECKPOINTING_ENABLED:-false}" -# Host OOM on GB200 when optimizer CPU offload is enabled for this model size. -OPTIMIZER_CPU_OFFLOAD="${OPTIMIZER_CPU_OFFLOAD:-false}" -OFFLOAD_OPTIMIZER_FOR_LOGPROB="${OFFLOAD_OPTIMIZER_FOR_LOGPROB:-false}" -if [[ "${OPTIMIZER_CPU_OFFLOAD}" == "true" ]]; then - OPTIMIZER_OFFLOAD_FRACTION="${OPTIMIZER_OFFLOAD_FRACTION:-1.0}" -else - OPTIMIZER_OFFLOAD_FRACTION="${OPTIMIZER_OFFLOAD_FRACTION:-0.0}" -fi -# Lower-precision Adam moments for the 2-GPU train half (HBM-tight on 1n4g). -# Override with EXP_AVG_DTYPE=float32 EXP_AVG_SQ_DTYPE=float32 STORE_PARAM_REMAINDERS=false -# for full-precision optimizer state, or set any to empty to skip the override. -EXP_AVG_DTYPE="${EXP_AVG_DTYPE:-bfloat16}" -EXP_AVG_SQ_DTYPE="${EXP_AVG_SQ_DTYPE:-bfloat16}" -STORE_PARAM_REMAINDERS="${STORE_PARAM_REMAINDERS:-true}" -BUFFER_SIZE_GB="${BUFFER_SIZE_GB:-8}" -MAX_TRAJECTORY_AGE_STEPS="${MAX_TRAJECTORY_AGE_STEPS:-2}" -IN_FLIGHT_WEIGHT_UPDATES="${IN_FLIGHT_WEIGHT_UPDATES:-true}" -# Use decode-only CUDA graphs with chunked prefill by default. -MEGATRON_ENABLE_CHUNKED_PREFILL="${MEGATRON_ENABLE_CHUNKED_PREFILL:-true}" -MEGATRON_CUDA_GRAPH_IMPL="${MEGATRON_CUDA_GRAPH_IMPL:-local}" # none local -MEGATRON_CUDA_GRAPH_SCOPE="${MEGATRON_CUDA_GRAPH_SCOPE:-block}" # layer block none -MEGATRON_NUM_CUDA_GRAPHS="${MEGATRON_NUM_CUDA_GRAPHS:--1}" # -1 -MEGATRON_USE_CUDA_GRAPHS_FOR_NON_DECODE="${MEGATRON_USE_CUDA_GRAPHS_FOR_NON_DECODE:-false}" -MOE_PAD_EXPERTS_FOR_CG="${MOE_PAD_EXPERTS_FOR_CG:-true}" - -# NSYS: ENABLE_NSYS=true NRL_NSYS_PROFILE_STEP_RANGE=1:4 -ENABLE_NSYS="${ENABLE_NSYS:-false}" -NSYS_ENV=() -if [[ "${ENABLE_NSYS}" == "true" ]]; then - NRL_NSYS_WORKER_PATTERNS="${NRL_NSYS_WORKER_PATTERNS:-*policy*,*megatron*}" - NRL_NSYS_PROFILE_STEP_RANGE="${NRL_NSYS_PROFILE_STEP_RANGE:-1:4}" - LD_LIBRARY_PATH="/usr/local/cuda/targets/aarch64-linux/lib:/usr/local/cuda/targets/x86_64-linux/lib:/usr/local/cuda/lib64:/usr/local/cuda/lib:/usr/local/nvidia/lib64:/usr/local/nvidia/lib:/usr/lib/aarch64-linux-gnu:/usr/lib/x86_64-linux-gnu:${LD_LIBRARY_PATH:-}" - NRL_NSYS_EXTRA_OPTIONS="${NRL_NSYS_EXTRA_OPTIONS:-{\"o\":\"/opt/nemo-rl/workspace/nsys/%p\",\"cpuctxsw\":\"none\",\"force-overwrite\":\"true\"}}" - NSYS_ENV=( - "NRL_NSYS_WORKER_PATTERNS=${NRL_NSYS_WORKER_PATTERNS}" - "NRL_NSYS_PROFILE_STEP_RANGE=${NRL_NSYS_PROFILE_STEP_RANGE}" - "LD_LIBRARY_PATH=${LD_LIBRARY_PATH}" - "NRL_NSYS_EXTRA_OPTIONS=${NRL_NSYS_EXTRA_OPTIONS}" - ) - mkdir -p /opt/nemo-rl/workspace/nsys -fi - -mkdir -p \ - "${HF_HUB_CACHE}" \ - "${HF_DATASETS_CACHE}" \ - "${HF_MODULES_CACHE}" \ - "${TRANSFORMERS_CACHE}" \ - "${NRL_MEGATRON_CHECKPOINT_DIR}" \ - "${XDG_CACHE_HOME}" \ - "${TORCH_HOME}" \ - "${TRITON_CACHE_DIR}" \ - "${RESULTS_DIR}" - -if [[ ! -f "${CONFIG}" ]]; then - echo "Config is missing under ${NEMORL}: ${CONFIG}" >&2 - exit 1 -fi - -COMMON_OVERRIDES=( - cluster.num_nodes=1 - cluster.gpus_per_node="${GPUS_PER_NODE}" - policy.model_name="${MODEL_NAME}" - policy.tokenizer.name="${MODEL_NAME}" - policy.is_vlm=true - policy.megatron_cfg.tensor_model_parallel_size="${POLICY_TP}" - policy.megatron_cfg.expert_model_parallel_size="${POLICY_EP}" - policy.megatron_cfg.expert_tensor_parallel_size=1 - policy.megatron_cfg.context_parallel_size="${POLICY_CP}" - policy.megatron_cfg.optimizer.optimizer_cpu_offload="${OPTIMIZER_CPU_OFFLOAD}" - policy.megatron_cfg.optimizer.optimizer_offload_fraction="${OPTIMIZER_OFFLOAD_FRACTION}" - policy.offload_optimizer_for_logprob="${OFFLOAD_OPTIMIZER_FOR_LOGPROB}" -) -if [[ -n "${EXP_AVG_DTYPE}" ]]; then - COMMON_OVERRIDES+=("++policy.megatron_cfg.optimizer.exp_avg_dtype=${EXP_AVG_DTYPE}") -fi -if [[ -n "${EXP_AVG_SQ_DTYPE}" ]]; then - COMMON_OVERRIDES+=("++policy.megatron_cfg.optimizer.exp_avg_sq_dtype=${EXP_AVG_SQ_DTYPE}") -fi -if [[ -n "${STORE_PARAM_REMAINDERS}" ]]; then - COMMON_OVERRIDES+=("++policy.megatron_cfg.optimizer.store_param_remainders=${STORE_PARAM_REMAINDERS}") -fi -COMMON_OVERRIDES+=( - policy.generation.backend="${GENERATION_BACKEND}" - policy.generation.colocated.enabled="${COLOCATED_ENABLED}" - policy.generation.colocated.resources.num_nodes="${NUM_GEN_NODES}" - policy.generation.colocated.resources.gpus_per_node="${GEN_GPUS_PER_NODE}" - policy.max_total_sequence_length="${MAX_SEQUENCE_LENGTH}" - policy.generation.max_new_tokens="${MAX_NEW_TOKENS}" - grpo.async_grpo.enabled="${ASYNC_GRPO}" - grpo.async_grpo.max_trajectory_age_steps="${MAX_TRAJECTORY_AGE_STEPS}" - grpo.async_grpo.in_flight_weight_updates="${IN_FLIGHT_WEIGHT_UPDATES}" - loss_fn.use_importance_sampling_correction=true - grpo.num_prompts_per_step="${NUM_PROMPTS}" - grpo.num_generations_per_prompt="${NUM_GENERATIONS}" - grpo.val_period=0 - grpo.val_at_start=false - grpo.val_at_end=false - policy.train_global_batch_size="${TRAIN_GBS}" - grpo.max_num_steps="${MAX_STEPS}" - checkpointing.enabled="${CHECKPOINTING_ENABLED}" - checkpointing.checkpoint_dir="${RESULTS_DIR}" - logger.log_dir="${RESULTS_DIR}" - logger.wandb_enabled="${WANDB_ENABLED}" - logger.tensorboard_enabled=false - logger.wandb.name="${WANDB_NAME}" - logger.wandb.project="${WANDB_PROJ}" - +logger.wandb.entity="${WANDB_GROUP}" -) - -GEN_OVERRIDES=() -if [[ "${GENERATION_BACKEND}" == "megatron" ]]; then - GEN_OVERRIDES=( - ++policy.generation.stop_strings=null - ++policy.generation.bad_words=null - policy.generation.mcore_generation_config.tensor_model_parallel_size="${INFER_TP}" - policy.generation.mcore_generation_config.expert_model_parallel_size="${INFER_EP}" - policy.generation.mcore_generation_config.expert_tensor_parallel_size=1 - ++policy.generation.mcore_generation_config.context_parallel_size="${POLICY_CP}" - ++policy.generation.mcore_generation_config.mamba_inference_ssm_states_dtype=float32 - ++policy.generation.mcore_generation_config.mamba_inference_conv_states_dtype=float32 - ++policy.generation.mcore_generation_config.logprobs_mode=raw_logprobs - policy.generation.mcore_generation_config.enable_chunked_prefill="${MEGATRON_ENABLE_CHUNKED_PREFILL}" - ++policy.generation.mcore_generation_config.async_sched_mode=async - policy.generation.mcore_generation_config.cuda_graph_impl="${MEGATRON_CUDA_GRAPH_IMPL}" - policy.generation.mcore_generation_config.inference_cuda_graph_scope="${MEGATRON_CUDA_GRAPH_SCOPE}" - policy.generation.mcore_generation_config.num_cuda_graphs="${MEGATRON_NUM_CUDA_GRAPHS}" - policy.generation.mcore_generation_config.use_cuda_graphs_for_non_decode_steps="${MEGATRON_USE_CUDA_GRAPHS_FOR_NON_DECODE}" - policy.generation.mcore_generation_config.refit_backend="${REFIT_BACKEND}" - policy.generation.mcore_generation_config.buffer_size_gb="${BUFFER_SIZE_GB}" - policy.generation.mcore_generation_config.moe_pad_experts_for_cuda_graph_inference="${MOE_PAD_EXPERTS_FOR_CG}" - policy.generation.mcore_generation_config.max_model_len="${MAX_SEQUENCE_LENGTH}" - policy.generation.mcore_generation_config.max_tokens="${MAX_SEQUENCE_LENGTH}" - ) -else - # To fit vLLM generation on 1 node. Refit packs weights into chunks sized at - # NRL_REFIT_BUFFER_MEMORY_RATIO * total HBM; exported globally on purpose - # because producer and consumer must agree on the chunk boundaries. - export NRL_REFIT_BUFFER_MEMORY_RATIO="${NRL_REFIT_BUFFER_MEMORY_RATIO:-0.005}" - # Omni vLLM A/B path. Use EP=TP so async GRPO can use Ray DP. - VLLM_GPU_MEMORY_UTILIZATION="${VLLM_GPU_MEMORY_UTILIZATION:-0.5}" - VLLM_ENFORCE_EAGER="${VLLM_ENFORCE_EAGER:-true}" - # Per-step forward-pass token budget; mirrors mcore_generation_config.max_tokens. - VLLM_MAX_NUM_BATCHED_TOKENS="${VLLM_MAX_NUM_BATCHED_TOKENS:-${MAX_SEQUENCE_LENGTH}}" - GEN_OVERRIDES=( - ++policy.generation.stop_strings=null - ++policy.generation.bad_words=null - policy.generation.vllm_cfg.async_engine="${ASYNC_GRPO}" - policy.generation.vllm_cfg.skip_tokenizer_init=false - policy.generation.vllm_cfg.tensor_parallel_size="${INFER_TP}" - policy.generation.vllm_cfg.pipeline_parallel_size=1 - policy.generation.vllm_cfg.expert_parallel_size="${INFER_EP}" - policy.generation.vllm_cfg.max_model_len="${MAX_SEQUENCE_LENGTH}" - ++policy.generation.vllm_cfg.cap_max_tokens_to_context=true - policy.generation.vllm_cfg.gpu_memory_utilization="${VLLM_GPU_MEMORY_UTILIZATION}" - policy.generation.vllm_cfg.enforce_eager="${VLLM_ENFORCE_EAGER}" - ++policy.generation.vllm_cfg.enable_prefix_caching=false - policy.generation.vllm_cfg.logprobs_mode=raw_logprobs - ++policy.generation.vllm_kwargs.limit_mm_per_prompt.image=2 - ++policy.generation.vllm_kwargs.max_num_batched_tokens="${VLLM_MAX_NUM_BATCHED_TOKENS}" - ++policy.generation.vllm_kwargs.mamba_ssm_cache_dtype=float32 - ++policy.generation.vllm_kwargs.skip_mm_profiling=true - ++policy.generation.vllm_kwargs.kernel_config.enable_flashinfer_autotune=false - ++policy.generation.vllm_kwargs.kernel_config.moe_backend=triton - ) -fi - -echo "Launching ${JOB_NAME}: ${GPUS_PER_NODE} visible GPU(s)" -echo " generation backend: ${GENERATION_BACKEND} colocated=${COLOCATED_ENABLED} async=${ASYNC_GRPO}" -echo " training world size: ${TRAIN_WORLD_SIZE} (TP=${POLICY_TP}, EP=${POLICY_EP}, ETP=1)" -echo " inference world size: ${INFERENCE_WORLD_SIZE} (TP=${INFER_TP}, EP=${INFER_EP})" -echo " async: max_trajectory_age=${MAX_TRAJECTORY_AGE_STEPS} in_flight_weight_updates=${IN_FLIGHT_WEIGHT_UPDATES}" -echo " seq/new_tokens: ${MAX_SEQUENCE_LENGTH}/${MAX_NEW_TOKENS}" -echo " optimizer moments: exp_avg=${EXP_AVG_DTYPE:-} exp_avg_sq=${EXP_AVG_SQ_DTYPE:-} store_param_remainders=${STORE_PARAM_REMAINDERS:-}" -echo " W&B: ${WANDB_GROUP}/${WANDB_PROJ}/${WANDB_NAME} (enabled=${WANDB_ENABLED})" - -exec env "${NSYS_ENV[@]}" uv run --no-sync python examples/run_vlm_grpo.py \ - --config "${CONFIG}" \ - "${COMMON_OVERRIDES[@]}" \ - "${GEN_OVERRIDES[@]}" \ - "$@" diff --git a/scripts/run_nemotron_omni_vstat_megatron_1n4g.sh b/scripts/run_nemotron_omni_vstat_megatron_1n4g.sh deleted file mode 100755 index 96914af0e00..00000000000 --- a/scripts/run_nemotron_omni_vstat_megatron_1n4g.sh +++ /dev/null @@ -1,245 +0,0 @@ -#!/usr/bin/env bash -set -euo pipefail - -# One-node / four-GPU Nemotron Omni video-GRPO smoke test. -# VSTAT is recommended here because Hugging Face hosts both the MCQ annotations -# and the actual MP4 assets; Video-MME commonly requires a separate video fetch. -# -# ASYNC_GRPO=false selects synchronous GRPO. -# MEGATRON_CUDA_GRAPH_IMPL=local enables CUDA graphs; block scope and MoE -# expert padding are selected automatically unless explicitly overridden. - -NEMORL="${NEMORL:-/opt/nemo-rl}" -MODEL_NAME="${MODEL_NAME:-nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16}" -WORKSPACE_ROOT="${WORKSPACE_ROOT:-${NEMORL}/workspace}" -DATA_ROOT="${DATA_ROOT:-${WORKSPACE_ROOT}/datasets/vstat-smoke}" -HF_DATASET="${HF_DATASET:-ShushengYang/VSTAT}" -NUM_DATA_ROWS="${NUM_DATA_ROWS:-8}" -PREPARE_VSTAT="${PREPARE_VSTAT:-false}" -NUM_FRAMES="${NUM_FRAMES:-8}" -TEMPORAL_PATCH_SIZE="${TEMPORAL_PATCH_SIZE:-2}" -VIDEO_TARGET_PATCHES="${VIDEO_TARGET_PATCHES:-256}" -MAX_SEQUENCE_LENGTH="${MAX_SEQUENCE_LENGTH:-4096}" -MAX_NEW_TOKENS="${MAX_NEW_TOKENS:-1024}" -MIN_GENERATION_TOKENS="${MIN_GENERATION_TOKENS:-2000}" -VISION_EMBEDDING_CACHE_MAX_BYTES="${VISION_EMBEDDING_CACHE_MAX_BYTES:-536870912}" -ENABLE_THINKING="${ENABLE_THINKING:-true}" -MAX_STEPS="${MAX_STEPS:-4}" -NUM_PROMPTS_PER_STEP="${NUM_PROMPTS_PER_STEP:-2}" -NUM_GENERATIONS_PER_PROMPT="${NUM_GENERATIONS_PER_PROMPT:-4}" -TRAIN_GBS="${TRAIN_GBS:-$((NUM_PROMPTS_PER_STEP * NUM_GENERATIONS_PER_PROMPT))}" -GEN_GPUS="${GEN_GPUS:-2}" -INFER_EP="${INFER_EP:-${GEN_GPUS}}" -REFIT_BACKEND="${REFIT_BACKEND:-nccl}" -ASYNC_GRPO="${ASYNC_GRPO:-true}" -MAX_TRAJECTORY_AGE_STEPS="${MAX_TRAJECTORY_AGE_STEPS:-2}" -IN_FLIGHT_WEIGHT_UPDATES="${IN_FLIGHT_WEIGHT_UPDATES:-true}" -MEGATRON_CUDA_GRAPH_IMPL="${MEGATRON_CUDA_GRAPH_IMPL:-local}" -if [[ "${MEGATRON_CUDA_GRAPH_IMPL}" == "none" ]]; then - DEFAULT_CUDA_GRAPH_SCOPE=none -else - DEFAULT_CUDA_GRAPH_SCOPE=block -fi -if [[ "${MEGATRON_CUDA_GRAPH_IMPL}" != "none" && "${INFER_EP}" -gt 1 ]]; then - DEFAULT_MOE_PAD_EXPERTS_FOR_CG=true -else - DEFAULT_MOE_PAD_EXPERTS_FOR_CG=false -fi -MEGATRON_CUDA_GRAPH_SCOPE="${MEGATRON_CUDA_GRAPH_SCOPE:-${DEFAULT_CUDA_GRAPH_SCOPE}}" -MEGATRON_NUM_CUDA_GRAPHS="${MEGATRON_NUM_CUDA_GRAPHS:--1}" -MEGATRON_USE_CUDA_GRAPHS_FOR_NON_DECODE="${MEGATRON_USE_CUDA_GRAPHS_FOR_NON_DECODE:-false}" -MOE_PAD_EXPERTS_FOR_CG="${MOE_PAD_EXPERTS_FOR_CG:-${DEFAULT_MOE_PAD_EXPERTS_FOR_CG}}" -MEGATRON_ASYNC_SCHED_MODE="${MEGATRON_ASYNC_SCHED_MODE:-async}" -# Lower-precision Adam moments for the 2-GPU train half (HBM-tight on 1n4g). -# Override with EXP_AVG_DTYPE=float32 EXP_AVG_SQ_DTYPE=float32 STORE_PARAM_REMAINDERS=false -# for full-precision optimizer state. -EXP_AVG_DTYPE="${EXP_AVG_DTYPE:-bfloat16}" -EXP_AVG_SQ_DTYPE="${EXP_AVG_SQ_DTYPE:-bfloat16}" -STORE_PARAM_REMAINDERS="${STORE_PARAM_REMAINDERS:-true}" - -cd "${NEMORL}" - -GPUS_PER_NODE="${GPUS_PER_NODE:-$(nvidia-smi --query-gpu=index --format=csv,noheader | wc -l)}" -if (( GPUS_PER_NODE != 4 )); then - echo "This smoke launcher expects exactly four visible GPUs (got ${GPUS_PER_NODE})." >&2 - exit 1 -fi -if (( GEN_GPUS <= 0 || GEN_GPUS >= GPUS_PER_NODE )); then - echo "GEN_GPUS must leave at least one GPU for training." >&2 - exit 1 -fi -if (( INFER_EP <= 0 || GEN_GPUS % INFER_EP != 0 )); then - echo "INFER_EP must be a positive divisor of GEN_GPUS (got ${INFER_EP})." >&2 - exit 1 -fi -if [[ "${ASYNC_GRPO}" != "true" && "${ASYNC_GRPO}" != "false" ]]; then - echo "ASYNC_GRPO must be true or false (got ${ASYNC_GRPO})." >&2 - exit 1 -fi -if [[ "${MEGATRON_CUDA_GRAPH_IMPL}" != "none" && - "${MEGATRON_CUDA_GRAPH_SCOPE}" == "none" ]]; then - echo "CUDA graphs require a non-none MEGATRON_CUDA_GRAPH_SCOPE." >&2 - exit 1 -fi -if [[ "${MEGATRON_CUDA_GRAPH_IMPL}" != "none" && "${INFER_EP}" -gt 1 && - "${MOE_PAD_EXPERTS_FOR_CG}" != "true" ]]; then - echo "CUDA graphs with expert parallelism require MOE_PAD_EXPERTS_FOR_CG=true." >&2 - exit 1 -fi -TRAIN_GPUS=$((GPUS_PER_NODE - GEN_GPUS)) - -CACHE_ROOT="${CACHE_ROOT:-${WORKSPACE_ROOT}/cache/nemo-rl-omni}" -export HF_HOME="${HF_HOME:-${CACHE_ROOT}/huggingface}" -export HF_HUB_CACHE="${HF_HUB_CACHE:-${HF_HOME}/hub}" -export HF_DATASETS_CACHE="${HF_DATASETS_CACHE:-${HF_HOME}/datasets}" -export NRL_MEGATRON_CHECKPOINT_DIR="${NRL_MEGATRON_CHECKPOINT_DIR:-${CACHE_ROOT}/megatron-checkpoints}" -export NEMO_RL_VENV_DIR="${NEMO_RL_VENV_DIR:-/opt/ray_venvs}" -export NEMO_GYM_VENV_DIR="${NEMO_GYM_VENV_DIR:-/opt/ray_venvs}" -export NEMO_GYM_EXTRA_ROOTS="${NEMO_GYM_EXTRA_ROOTS:-${NEMORL}/3rdparty/Gym-workspace/Gym}" -export NRL_VIDEO_BACKEND="${NRL_VIDEO_BACKEND:-torchcodec}" -export NRL_VIDEO_SAMPLING_STYLE="${NRL_VIDEO_SAMPLING_STYLE:-nemotron_vl}" -export NRL_VIDEO_TEMPORAL_PATCH_SIZE="${TEMPORAL_PATCH_SIZE}" -export RAY_ENABLE_UV_RUN_RUNTIME_ENV=0 -export NRL_VENVS_TRUST_EXISTING="${NRL_VENVS_TRUST_EXISTING:-1}" -export CUDA_DEVICE_MAX_CONNECTIONS="${CUDA_DEVICE_MAX_CONNECTIONS:-1}" -export FLASHINFER_DISABLE_VERSION_CHECK="${FLASHINFER_DISABLE_VERSION_CHECK:-1}" - -AUDIO_DEPS_SCRIPT="${NEMORL}/tools/install_audio_deps.sh" -MEGATRON_WORKER_PYTHON="${RAY_MEGATRON_PYTHON:-${NEMO_RL_VENV_DIR}/nemo_rl.models.policy.workers.megatron_policy_worker.MegatronPolicyWorker/bin/python}" -NEED_AUDIO_VIDEO_DEPS=false -if ! python -c "import torchcodec" >/dev/null 2>&1; then - NEED_AUDIO_VIDEO_DEPS=true -fi -if [[ ! -x "${MEGATRON_WORKER_PYTHON}" ]] || - ! "${MEGATRON_WORKER_PYTHON}" -c "import av" >/dev/null 2>&1; then - NEED_AUDIO_VIDEO_DEPS=true -fi -if [[ "${NEED_AUDIO_VIDEO_DEPS}" == "true" ]]; then - if [[ ! -f "${AUDIO_DEPS_SCRIPT}" ]]; then - echo "Audio/video dependency installer is missing: ${AUDIO_DEPS_SCRIPT}" >&2 - exit 1 - fi - echo "Installing missing audio/video dependencies" - RAY_MEGATRON_PYTHON="${MEGATRON_WORKER_PYTHON}" bash "${AUDIO_DEPS_SCRIPT}" -fi - -# NSYS: ENABLE_NSYS=true NRL_NSYS_PROFILE_STEP_RANGE=1:4 -ENABLE_NSYS="${ENABLE_NSYS:-false}" -NSYS_ENV=() -if [[ "${ENABLE_NSYS}" == "true" ]]; then - NRL_NSYS_WORKER_PATTERNS="${NRL_NSYS_WORKER_PATTERNS:-*policy*,*megatron*}" - NRL_NSYS_PROFILE_STEP_RANGE="${NRL_NSYS_PROFILE_STEP_RANGE:-1:4}" - LD_LIBRARY_PATH="/usr/local/cuda/targets/aarch64-linux/lib:/usr/local/cuda/targets/x86_64-linux/lib:/usr/local/cuda/lib64:/usr/local/cuda/lib:/usr/local/nvidia/lib64:/usr/local/nvidia/lib:/usr/lib/aarch64-linux-gnu:/usr/lib/x86_64-linux-gnu:${LD_LIBRARY_PATH:-}" - NRL_NSYS_EXTRA_OPTIONS="${NRL_NSYS_EXTRA_OPTIONS:-{\"o\":\"${WORKSPACE_ROOT}/nsys/%p\",\"cpuctxsw\":\"none\",\"force-overwrite\":\"true\"}}" - NSYS_ENV=( - "NRL_NSYS_WORKER_PATTERNS=${NRL_NSYS_WORKER_PATTERNS}" - "NRL_NSYS_PROFILE_STEP_RANGE=${NRL_NSYS_PROFILE_STEP_RANGE}" - "LD_LIBRARY_PATH=${LD_LIBRARY_PATH}" - "NRL_NSYS_EXTRA_OPTIONS=${NRL_NSYS_EXTRA_OPTIONS}" - ) - mkdir -p "${WORKSPACE_ROOT}/nsys" -fi - -BRIDGE="${NEMORL}/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge" -export PYTHONPATH="${NEMORL}:${NEMO_GYM_EXTRA_ROOTS}:${BRIDGE}/src:${BRIDGE}/3rdparty/Megatron-LM${PYTHONPATH:+:${PYTHONPATH}}" - -mkdir -p "${DATA_ROOT}" "${HF_HOME}" "${NRL_MEGATRON_CHECKPOINT_DIR}" -export NEMO_RL_VIDEO_TRAIN_JSONL="${DATA_ROOT}/train-gym.jsonl" -export NEMO_RL_VIDEO_VAL_JSONL="${DATA_ROOT}/val-gym.jsonl" -export NEMO_RL_VIDEO_MEDIA_ROOT="${DATA_ROOT}/media" - -VSTAT_PREPARE_SCRIPT="${NEMORL}/scripts/prepare_nemotron_omni_vstat.py" -if [[ ! -f "${VSTAT_PREPARE_SCRIPT}" ]]; then - echo "VSTAT preparation script is missing: ${VSTAT_PREPARE_SCRIPT}" >&2 - exit 1 -fi -if [[ "${PREPARE_VSTAT}" == "true" || - ! -s "${NEMO_RL_VIDEO_TRAIN_JSONL}" || - ! -s "${NEMO_RL_VIDEO_VAL_JSONL}" ]]; then - echo "Preparing VSTAT under ${DATA_ROOT}" - uv run --no-sync python "${VSTAT_PREPARE_SCRIPT}" \ - --output-dir "${DATA_ROOT}" \ - --repo-id "${HF_DATASET}" \ - --num-rows "${NUM_DATA_ROWS}" -fi - -RESULTS_DIR="${RESULTS_DIR:-${WORKSPACE_ROOT}/results/nemo-rl-omni/nemotron-omni-vstat-megatron-1n4g}" -mkdir -p "${RESULTS_DIR}" - -echo "VSTAT: root=${DATA_ROOT} repo=${HF_DATASET} rows=${NUM_DATA_ROWS} prepare=${PREPARE_VSTAT}" -echo " async_grpo=${ASYNC_GRPO} scheduler=${MEGATRON_ASYNC_SCHED_MODE}" -echo " generation: TP=${GEN_GPUS} EP=${INFER_EP}" -echo " cuda_graph_impl=${MEGATRON_CUDA_GRAPH_IMPL} scope=${MEGATRON_CUDA_GRAPH_SCOPE} non_decode=${MEGATRON_USE_CUDA_GRAPHS_FOR_NON_DECODE}" -exec env "${NSYS_ENV[@]}" 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 \ - policy.model_name="${MODEL_NAME}" \ - policy.tokenizer.name="${MODEL_NAME}" \ - policy.tokenizer.chat_template_kwargs.enable_thinking="${ENABLE_THINKING}" \ - policy.is_vlm=true \ - policy.generation.backend=megatron \ - ++policy.generation.bad_words=null \ - policy.generation.colocated.enabled=false \ - policy.generation.colocated.resources.num_nodes=1 \ - policy.generation.colocated.resources.gpus_per_node="${GEN_GPUS}" \ - policy.megatron_cfg.tensor_model_parallel_size="${TRAIN_GPUS}" \ - policy.megatron_cfg.expert_model_parallel_size="${TRAIN_GPUS}" \ - policy.megatron_cfg.expert_tensor_parallel_size=1 \ - policy.megatron_cfg.context_parallel_size=1 \ - policy.megatron_cfg.sequence_parallel=true \ - policy.megatron_cfg.bias_activation_fusion=false \ - policy.megatron_cfg.optimizer.optimizer_cpu_offload=false \ - policy.megatron_cfg.optimizer.optimizer_offload_fraction=0.0 \ - ++policy.megatron_cfg.optimizer.exp_avg_dtype="${EXP_AVG_DTYPE}" \ - ++policy.megatron_cfg.optimizer.exp_avg_sq_dtype="${EXP_AVG_SQ_DTYPE}" \ - ++policy.megatron_cfg.optimizer.store_param_remainders="${STORE_PARAM_REMAINDERS}" \ - policy.generation.mcore_generation_config.tensor_model_parallel_size="${GEN_GPUS}" \ - policy.generation.mcore_generation_config.expert_model_parallel_size="${INFER_EP}" \ - policy.generation.mcore_generation_config.expert_tensor_parallel_size=1 \ - ++policy.generation.mcore_generation_config.context_parallel_size=1 \ - policy.generation.mcore_generation_config.transformer_impl=transformer_engine \ - policy.generation.mcore_generation_config.sequence_parallel=true \ - policy.generation.mcore_generation_config.refit_backend="${REFIT_BACKEND}" \ - policy.generation.mcore_generation_config.buffer_size_gb=8 \ - policy.generation.mcore_generation_config.cuda_graph_impl="${MEGATRON_CUDA_GRAPH_IMPL}" \ - policy.generation.mcore_generation_config.inference_cuda_graph_scope="${MEGATRON_CUDA_GRAPH_SCOPE}" \ - policy.generation.mcore_generation_config.num_cuda_graphs="${MEGATRON_NUM_CUDA_GRAPHS}" \ - ++policy.generation.mcore_generation_config.moe_pad_experts_for_cuda_graph_inference="${MOE_PAD_EXPERTS_FOR_CG}" \ - policy.generation.mcore_generation_config.use_cuda_graphs_for_non_decode_steps="${MEGATRON_USE_CUDA_GRAPHS_FOR_NON_DECODE}" \ - policy.generation.mcore_generation_config.max_model_len="${MAX_SEQUENCE_LENGTH}" \ - policy.generation.mcore_generation_config.max_tokens="${MAX_SEQUENCE_LENGTH}" \ - policy.generation.mcore_generation_config.enable_chunked_prefill=true \ - ++policy.generation.mcore_generation_config.async_sched_mode="${MEGATRON_ASYNC_SCHED_MODE}" \ - policy.generation.mcore_generation_config.enable_prefix_caching=false \ - ++policy.generation.mcore_generation_config.vision_embedding_cache_max_bytes="${VISION_EMBEDDING_CACHE_MAX_BYTES}" \ - policy.generation.mcore_generation_config.parsers=[nemotron-v3-reasoning,qwen3-coder-tool] \ - ++policy.generation.mcore_generation_config.video_num_frames="${NUM_FRAMES}" \ - ++policy.generation.mcore_generation_config.video_temporal_patch_size="${TEMPORAL_PATCH_SIZE}" \ - ++policy.generation.mcore_generation_config.video_target_num_patches="${VIDEO_TARGET_PATCHES}" \ - policy.max_total_sequence_length="${MAX_SEQUENCE_LENGTH}" \ - policy.generation.max_new_tokens="${MAX_NEW_TOKENS}" \ - +data.default.num_frames="${NUM_FRAMES}" \ - +data.default.video_sampling_style=nemotron_vl \ - +data.default.video_temporal_patch_size="${TEMPORAL_PATCH_SIZE}" \ - +data.default.min_generation_tokens="${MIN_GENERATION_TOKENS}" \ - data.default.video_target_num_patches="${VIDEO_TARGET_PATCHES}" \ - data.train.data_path="${NEMO_RL_VIDEO_TRAIN_JSONL}" \ - data.validation.data_path="${NEMO_RL_VIDEO_VAL_JSONL}" \ - ++env.nemo_gym.policy_model.responses_api_models.vllm_model.chat_template_kwargs.enable_thinking="${ENABLE_THINKING}" \ - grpo.deduplicate_multimodal_data=false \ - grpo.async_grpo.enabled="${ASYNC_GRPO}" \ - grpo.async_grpo.max_trajectory_age_steps="${MAX_TRAJECTORY_AGE_STEPS}" \ - grpo.async_grpo.in_flight_weight_updates="${IN_FLIGHT_WEIGHT_UPDATES}" \ - loss_fn.use_importance_sampling_correction=true \ - grpo.num_prompts_per_step="${NUM_PROMPTS_PER_STEP}" \ - grpo.num_generations_per_prompt="${NUM_GENERATIONS_PER_PROMPT}" \ - policy.train_global_batch_size="${TRAIN_GBS}" \ - grpo.max_num_steps="${MAX_STEPS}" \ - grpo.val_period=0 \ - grpo.val_at_start=false \ - grpo.val_at_end=false \ - cluster.num_nodes=1 \ - cluster.gpus_per_node="${GPUS_PER_NODE}" \ - checkpointing.enabled=false \ - logger.log_dir="${RESULTS_DIR}" \ - logger.wandb_enabled=false \ - "$@" diff --git a/scripts/submit_nemotron_omni_clevr_megatron_8n4g.sh b/scripts/submit_nemotron_omni_clevr_megatron_8n4g.sh deleted file mode 100755 index ea370f90e29..00000000000 --- a/scripts/submit_nemotron_omni_clevr_megatron_8n4g.sh +++ /dev/null @@ -1,422 +0,0 @@ -#!/usr/bin/env bash -set -euo pipefail - -# Submit the eight-node Omni CLEVR Megatron GRPO recipe through NeMo-RL's -# Ray/Slurm launcher. -# -# Default: non-colocated gen-heavy (6 gen + 2 train). Async colocated is -# supported on this branch (NVIDIA-NeMo/RL#2884); set COLOCATED=true to share -# all GPUs for train+gen. ASYNC_GRPO=true is required for colocated async. -# -# GENERATION_BACKEND=vllm flips inference to Omni vLLM for A/B vs Megatron. -# The A/B defaults disable stop/bad-word filters; native multimodal inputs are -# supported by both synchronous and non-colocated asynchronous rollouts. - -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -NEMORL="${NEMORL:-$(cd "${SCRIPT_DIR}/.." && pwd)}" -CONTAINER_NEMORL="${CONTAINER_NEMORL:-/opt/nemo-rl}" -WORKSPACE_ROOT="${WORKSPACE_ROOT:-${NEMORL}/workspace}" - -CONTAINER="${CONTAINER:-/lustre/fsw/portfolios/coreai/users/cye/enroot/nemo_rl_0807.sqsh}" -MODEL_NAME="${MODEL_NAME:-nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16}" -CONFIG="${CONFIG:-examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron_generation.v1.yaml}" -GENERATION_BACKEND="${GENERATION_BACKEND:-megatron}" -COLOCATED="${COLOCATED:-false}" -ASYNC_GRPO="${ASYNC_GRPO:-true}" - -if [[ "${GENERATION_BACKEND}" != "megatron" && "${GENERATION_BACKEND}" != "vllm" ]]; then - echo "GENERATION_BACKEND must be megatron or vllm (got ${GENERATION_BACKEND})." >&2 - exit 1 -fi -if [[ "${COLOCATED}" == "true" && "${ASYNC_GRPO}" == "true" && "${GENERATION_BACKEND}" != "megatron" ]]; then - echo "Colocated async GRPO requires GENERATION_BACKEND=megatron." >&2 - exit 1 -fi - -CACHE_ROOT="${CACHE_ROOT:-${WORKSPACE_ROOT}/cache/nemo-rl-omni}" -export HF_HOME="${HF_HOME:-${CACHE_ROOT}/huggingface}" -export HUGGINGFACE_HUB_CACHE="${HUGGINGFACE_HUB_CACHE:-${HF_HOME}/hub}" -export HF_HUB_CACHE="${HF_HUB_CACHE:-${HUGGINGFACE_HUB_CACHE}}" -export HF_DATASETS_CACHE="${HF_DATASETS_CACHE:-${HF_HOME}/datasets}" -export HF_MODULES_CACHE="${HF_MODULES_CACHE:-${HF_HOME}/modules}" -export TRANSFORMERS_CACHE="${TRANSFORMERS_CACHE:-${HF_HOME}/transformers}" -export NRL_MEGATRON_CHECKPOINT_DIR="${NRL_MEGATRON_CHECKPOINT_DIR:-${CACHE_ROOT}/megatron-checkpoints}" -export XDG_CACHE_HOME="${XDG_CACHE_HOME:-${CACHE_ROOT}/xdg}" -export TORCH_HOME="${TORCH_HOME:-${CACHE_ROOT}/torch}" -export TRITON_CACHE_DIR="${TRITON_CACHE_DIR:-${CACHE_ROOT}/triton}" - -NUM_NODES="${NUM_NODES:-8}" -GPUS_PER_NODE="${GPUS_PER_NODE:-4}" - -if [[ "${COLOCATED}" == "true" ]]; then - # Shared train+gen on every GPU (NVIDIA-NeMo/RL#2884). - TRAIN_WORLD_SIZE=$((NUM_NODES * GPUS_PER_NODE)) - INFERENCE_WORLD_SIZE="${TRAIN_WORLD_SIZE}" - GEN_GPUS_PER_NODE="${GPUS_PER_NODE}" - NUM_GEN_NODES="${NUM_NODES}" - COLOCATED_ENABLED=true -else - NUM_GEN_NODES="${NUM_GEN_NODES:-6}" - if (( NUM_NODES == 1 )); then - if (( NUM_GEN_NODES != 1 )); then - echo "One-node non-colocated mode requires NUM_GEN_NODES=1." >&2 - exit 1 - fi - GEN_GPUS_PER_NODE="${GEN_GPUS_PER_NODE:-$((GPUS_PER_NODE / 2))}" - if (( GPUS_PER_NODE < 4 || GEN_GPUS_PER_NODE <= 0 || GEN_GPUS_PER_NODE >= GPUS_PER_NODE )); then - echo "One-node non-colocated mode requires at least four GPUs and a strict train/inference split." >&2 - exit 1 - fi - TRAIN_WORLD_SIZE=$((GPUS_PER_NODE - GEN_GPUS_PER_NODE)) - else - GEN_GPUS_PER_NODE="${GEN_GPUS_PER_NODE:-${GPUS_PER_NODE}}" - if (( NUM_GEN_NODES <= 0 || NUM_GEN_NODES >= NUM_NODES )); then - echo "Multi-node non-colocated mode requires 0 < NUM_GEN_NODES < NUM_NODES." >&2 - exit 1 - fi - if (( GEN_GPUS_PER_NODE != GPUS_PER_NODE )); then - echo "Multi-node non-colocated inference must reserve complete GPU nodes." >&2 - exit 1 - fi - TRAIN_WORLD_SIZE=$(((NUM_NODES - NUM_GEN_NODES) * GPUS_PER_NODE)) - fi - INFERENCE_WORLD_SIZE=$((NUM_GEN_NODES * GEN_GPUS_PER_NODE)) - COLOCATED_ENABLED=false -fi - -# Prefer TP=8; fall back by powers of two for smaller worlds. -DEFAULT_POLICY_TP=8 -while (( DEFAULT_POLICY_TP > TRAIN_WORLD_SIZE || TRAIN_WORLD_SIZE % DEFAULT_POLICY_TP != 0 )); do - DEFAULT_POLICY_TP=$((DEFAULT_POLICY_TP / 2)) -done -DEFAULT_INFER_TP=8 -while (( DEFAULT_INFER_TP > INFERENCE_WORLD_SIZE || INFERENCE_WORLD_SIZE % DEFAULT_INFER_TP != 0 )); do - DEFAULT_INFER_TP=$((DEFAULT_INFER_TP / 2)) -done -POLICY_TP="${POLICY_TP:-${DEFAULT_POLICY_TP}}" -INFER_TP="${INFER_TP:-${DEFAULT_INFER_TP}}" -POLICY_CP="${POLICY_CP:-1}" - -if [[ "${GENERATION_BACKEND}" == "megatron" && "${POLICY_CP}" != "1" ]]; then - echo "Megatron dynamic inference requires POLICY_CP=1." >&2 - exit 1 -fi -if (( TRAIN_WORLD_SIZE % (POLICY_TP * POLICY_CP) != 0 )); then - echo "Training world size must be divisible by POLICY_TP * POLICY_CP." >&2 - exit 1 -fi -if (( INFERENCE_WORLD_SIZE % INFER_TP != 0 )); then - echo "Inference world size must be divisible by INFER_TP." >&2 - exit 1 -fi -TRAIN_DP_SIZE=$((TRAIN_WORLD_SIZE / (POLICY_TP * POLICY_CP))) -INFERENCE_DP_SIZE=$((INFERENCE_WORLD_SIZE / INFER_TP)) - -# With ETP=1, world_size must be divisible by EP (not by TP*EP). -largest_ep() { - local world="$1" - local ep=8 - while (( ep > world || world % ep != 0 )); do - ep=$((ep / 2)) - done - printf '%d' "${ep}" -} - -POLICY_EP="${POLICY_EP:-$(largest_ep "${TRAIN_WORLD_SIZE}")}" -INFER_EP="${INFER_EP:-$(largest_ep "${INFERENCE_WORLD_SIZE}")}" -if (( TRAIN_WORLD_SIZE % POLICY_EP != 0 )); then - echo "Training world size must be divisible by POLICY_EP (ETP=1)." >&2 - exit 1 -fi -if [[ "${GENERATION_BACKEND}" == "vllm" ]]; then - # vLLM EP = DP_internal * TP. Async GRPO cannot use internal DP, so EP must - # equal TP; Ray DP then matches Megatron's world/TP replica count. - if [[ "${INFER_EP}" != "${INFER_TP}" ]]; then - echo "Forcing INFER_EP=${INFER_TP} for vLLM (was ${INFER_EP}) for async-safe Ray DP." >&2 - fi - INFER_EP="${INFER_TP}" -elif (( INFERENCE_WORLD_SIZE % INFER_EP != 0 )); then - echo "Inference world size must be divisible by INFER_EP (ETP=1)." >&2 - exit 1 -fi - -MAX_STEPS="${MAX_STEPS:-1000000}" -MAX_SEQUENCE_LENGTH="${MAX_SEQUENCE_LENGTH:-4096}" -MAX_NEW_TOKENS="${MAX_NEW_TOKENS:-2048}" -# Two prompt groups per inference DP replica (non-colocated default: 6 gen -# nodes → DP=3 → prompts=6; colocated 8n4g → DP=4 → prompts=8). -NUM_PROMPTS="${NUM_PROMPTS:-$((INFERENCE_DP_SIZE * 2))}" -NUM_GENERATIONS="${NUM_GENERATIONS:-8}" -TRAIN_GBS="${TRAIN_GBS:-$((NUM_PROMPTS * NUM_GENERATIONS))}" -EXPECTED_TRAIN_GBS=$((NUM_PROMPTS * NUM_GENERATIONS)) -VAL_GBS="${VAL_GBS:-64}" -VAL_SIZE="${VAL_SIZE:-64}" -if (( TRAIN_GBS != EXPECTED_TRAIN_GBS )); then - echo "TRAIN_GBS (${TRAIN_GBS}) must equal NUM_PROMPTS * NUM_GENERATIONS (${EXPECTED_TRAIN_GBS})." >&2 - exit 1 -fi -if (( TRAIN_GBS % TRAIN_DP_SIZE != 0 )); then - echo "TRAIN_GBS (${TRAIN_GBS}) must be divisible by training DP size (${TRAIN_DP_SIZE})." >&2 - exit 1 -fi -if (( VAL_GBS % TRAIN_DP_SIZE != 0 )); then - echo "VAL_GBS (${VAL_GBS}) must be divisible by training DP size (${TRAIN_DP_SIZE})." >&2 - exit 1 -fi -REFIT_BACKEND="${REFIT_BACKEND:-nccl}" -OPTIMIZER_CPU_OFFLOAD="${OPTIMIZER_CPU_OFFLOAD:-false}" -BUFFER_SIZE_GB="${BUFFER_SIZE_GB:-8}" -OFFLOAD_OPTIMIZER_FOR_LOGPROB="${OFFLOAD_OPTIMIZER_FOR_LOGPROB:-false}" -if [[ "${OPTIMIZER_CPU_OFFLOAD}" == "true" ]]; then - OPTIMIZER_OFFLOAD_FRACTION="${OPTIMIZER_OFFLOAD_FRACTION:-1.0}" -else - OPTIMIZER_OFFLOAD_FRACTION="${OPTIMIZER_OFFLOAD_FRACTION:-0.0}" -fi -# Lower-precision Adam moments; set to empty to skip the override. -EXP_AVG_DTYPE="${EXP_AVG_DTYPE:-bfloat16}" -EXP_AVG_SQ_DTYPE="${EXP_AVG_SQ_DTYPE:-bfloat16}" -STORE_PARAM_REMAINDERS="${STORE_PARAM_REMAINDERS:-true}" -# Megatron defaults for stable inference/training parity. -MEGATRON_ENABLE_CHUNKED_PREFILL="${MEGATRON_ENABLE_CHUNKED_PREFILL:-true}" -MEGATRON_CUDA_GRAPH_IMPL="${MEGATRON_CUDA_GRAPH_IMPL:-local}" -MEGATRON_CUDA_GRAPH_SCOPE="${MEGATRON_CUDA_GRAPH_SCOPE:-block}" -MEGATRON_NUM_CUDA_GRAPHS="${MEGATRON_NUM_CUDA_GRAPHS:--1}" -MEGATRON_USE_CUDA_GRAPHS_FOR_NON_DECODE="${MEGATRON_USE_CUDA_GRAPHS_FOR_NON_DECODE:-false}" -MOE_PAD_EXPERTS_FOR_CG="${MOE_PAD_EXPERTS_FOR_CG:-true}" -EXTRA_OVERRIDES="${EXTRA_OVERRIDES:-}" -MAX_TRAJECTORY_AGE_STEPS="${MAX_TRAJECTORY_AGE_STEPS:-2}" -IN_FLIGHT_WEIGHT_UPDATES="${IN_FLIGHT_WEIGHT_UPDATES:-true}" - -JOB_NAME="${JOB_NAME:-nemotron-omni-clevr-${GENERATION_BACKEND}-8n4g}" -EXP_NAME="${EXP_NAME:-${JOB_NAME}}" -PRECISION_RECIPE="${PRECISION_RECIPE:-bf16}" -WANDB_ENABLED="${WANDB_ENABLED:-true}" -WANDB_PROJ="${WANDB_PROJ:-mllm-rl-dev}" -WANDB_GROUP="${WANDB_GROUP:-adlr}" -WANDB_NAME="${WANDB_NAME:-${EXP_NAME}-${PRECISION_RECIPE}-internal-repo}" -RESULTS_DIR="${RESULTS_DIR:-${WORKSPACE_ROOT}/results/nemo-rl-omni/${JOB_NAME}}" -CHECKPOINTING_ENABLED="${CHECKPOINTING_ENABLED:-false}" -SLURM_LOG_DIR="${SLURM_LOG_DIR:-${RESULTS_DIR}/slurm}" - -# OCI-HSG Slurm defaults. Override any of these for a reservation or a -# different account/partition. -SBATCH_ACCOUNT="${SBATCH_ACCOUNT:-nemotron_sw_post}" # coreai_dlalgo_mcore nemotron_sw_post -SBATCH_PARTITION="${SBATCH_PARTITION:-batch_long}" -SBATCH_QOS="${SBATCH_QOS:-}" -SBATCH_TIME="${SBATCH_TIME:-04:00:00}" -SBATCH_RESERVATION="${SBATCH_RESERVATION:-}" -SBATCH_SEGMENT="${NUM_NODES}" - -mkdir -p \ - "${HF_HUB_CACHE}" \ - "${HF_DATASETS_CACHE}" \ - "${HF_MODULES_CACHE}" \ - "${TRANSFORMERS_CACHE}" \ - "${NRL_MEGATRON_CHECKPOINT_DIR}" \ - "${XDG_CACHE_HOME}" \ - "${TORCH_HOME}" \ - "${TRITON_CACHE_DIR}" \ - "${SLURM_LOG_DIR}" - -if [[ ! -f "${CONTAINER}" ]]; then - echo "Container image does not exist: ${CONTAINER}" >&2 - exit 1 -fi -if [[ ! -f "${NEMORL}/ray.sub" || ! -f "${NEMORL}/${CONFIG}" ]]; then - echo "NeMo-RL launcher or config is missing under: ${NEMORL}" >&2 - exit 1 -fi - -OPTIMIZER_DTYPE_OVERRIDES="" -if [[ -n "${EXP_AVG_DTYPE}" ]]; then - OPTIMIZER_DTYPE_OVERRIDES+=" ++policy.megatron_cfg.optimizer.exp_avg_dtype=${EXP_AVG_DTYPE}" -fi -if [[ -n "${EXP_AVG_SQ_DTYPE}" ]]; then - OPTIMIZER_DTYPE_OVERRIDES+=" ++policy.megatron_cfg.optimizer.exp_avg_sq_dtype=${EXP_AVG_SQ_DTYPE}" -fi -if [[ -n "${STORE_PARAM_REMAINDERS}" ]]; then - OPTIMIZER_DTYPE_OVERRIDES+=" ++policy.megatron_cfg.optimizer.store_param_remainders=${STORE_PARAM_REMAINDERS}" -fi - -REFIT_ENV_EXPORTS="" -REFIT_BUFFER_MEMORY_RATIO="" - -if [[ "${GENERATION_BACKEND}" == "megatron" ]]; then - GEN_OVERRIDES="\ -++policy.generation.stop_strings=null \ -++policy.generation.bad_words=null \ -policy.generation.mcore_generation_config.tensor_model_parallel_size=${INFER_TP} \ -policy.generation.mcore_generation_config.expert_model_parallel_size=${INFER_EP} \ -policy.generation.mcore_generation_config.expert_tensor_parallel_size=1 \ -++policy.generation.mcore_generation_config.context_parallel_size=${POLICY_CP} \ -++policy.generation.mcore_generation_config.mamba_inference_ssm_states_dtype=float32 \ -++policy.generation.mcore_generation_config.mamba_inference_conv_states_dtype=float32 \ -++policy.generation.mcore_generation_config.logprobs_mode=raw_logprobs \ -policy.generation.mcore_generation_config.enable_chunked_prefill=${MEGATRON_ENABLE_CHUNKED_PREFILL} \ -++policy.generation.mcore_generation_config.async_sched_mode=async \ -policy.generation.mcore_generation_config.cuda_graph_impl=${MEGATRON_CUDA_GRAPH_IMPL} \ -policy.generation.mcore_generation_config.inference_cuda_graph_scope=${MEGATRON_CUDA_GRAPH_SCOPE} \ -policy.generation.mcore_generation_config.num_cuda_graphs=${MEGATRON_NUM_CUDA_GRAPHS} \ -policy.generation.mcore_generation_config.use_cuda_graphs_for_non_decode_steps=${MEGATRON_USE_CUDA_GRAPHS_FOR_NON_DECODE} \ -policy.generation.mcore_generation_config.refit_backend=${REFIT_BACKEND} \ -policy.generation.mcore_generation_config.buffer_size_gb=${BUFFER_SIZE_GB} \ -policy.generation.mcore_generation_config.moe_pad_experts_for_cuda_graph_inference=${MOE_PAD_EXPERTS_FOR_CG} \ -policy.generation.mcore_generation_config.max_model_len=${MAX_SEQUENCE_LENGTH} \ -policy.generation.mcore_generation_config.max_tokens=${MAX_SEQUENCE_LENGTH}" -else - # Non-colocated refit NCCL-broadcasts weights in packed chunks sized at - # NRL_REFIT_BUFFER_MEMORY_RATIO * total HBM (the 0.02 default is 3.7GiB on - # GB200), which the ~170GiB-deep trainer cannot spare. Exported globally on - # purpose: producer and consumer both derive chunk boundaries from it and must - # agree. - REFIT_BUFFER_MEMORY_RATIO="${NRL_REFIT_BUFFER_MEMORY_RATIO:-0.005}" - REFIT_ENV_EXPORTS="export NRL_REFIT_BUFFER_MEMORY_RATIO=${REFIT_BUFFER_MEMORY_RATIO}" - - # Omni vLLM A/B path. - VLLM_GPU_MEMORY_UTILIZATION="${VLLM_GPU_MEMORY_UTILIZATION:-0.5}" - VLLM_ENFORCE_EAGER="${VLLM_ENFORCE_EAGER:-true}" - # Per-step forward-pass token budget; mirrors mcore_generation_config.max_tokens. - VLLM_MAX_NUM_BATCHED_TOKENS="${VLLM_MAX_NUM_BATCHED_TOKENS:-${MAX_SEQUENCE_LENGTH}}" - GEN_OVERRIDES="\ -++policy.generation.stop_strings=null \ -++policy.generation.bad_words=null \ -policy.generation.vllm_cfg.async_engine=${ASYNC_GRPO} \ -policy.generation.vllm_cfg.skip_tokenizer_init=false \ -policy.generation.vllm_cfg.tensor_parallel_size=${INFER_TP} \ -policy.generation.vllm_cfg.pipeline_parallel_size=1 \ -policy.generation.vllm_cfg.expert_parallel_size=${INFER_EP} \ -policy.generation.vllm_cfg.max_model_len=${MAX_SEQUENCE_LENGTH} \ -++policy.generation.vllm_cfg.cap_max_tokens_to_context=true \ -policy.generation.vllm_cfg.gpu_memory_utilization=${VLLM_GPU_MEMORY_UTILIZATION} \ -policy.generation.vllm_cfg.enforce_eager=${VLLM_ENFORCE_EAGER} \ -++policy.generation.vllm_cfg.enable_prefix_caching=false \ -policy.generation.vllm_cfg.logprobs_mode=raw_logprobs \ -++policy.generation.vllm_kwargs.limit_mm_per_prompt.image=2 \ -++policy.generation.vllm_kwargs.max_num_batched_tokens=${VLLM_MAX_NUM_BATCHED_TOKENS} \ -++policy.generation.vllm_kwargs.mamba_ssm_cache_dtype=float32 \ -++policy.generation.vllm_kwargs.skip_mm_profiling=true \ -++policy.generation.vllm_kwargs.kernel_config.enable_flashinfer_autotune=false \ -++policy.generation.vllm_kwargs.kernel_config.moe_backend=triton" -fi - -export NUM_NODES GPUS_PER_NODE CONTAINER -export RAY_ENABLE_UV_RUN_RUNTIME_ENV=0 -export NRL_VENVS_TRUST_EXISTING="${NRL_VENVS_TRUST_EXISTING:-1}" -export NRL_FORCE_REBUILD_VENVS="${NRL_FORCE_REBUILD_VENVS:-false}" -export NEMO_RL_VENV_DIR="${NEMO_RL_VENV_DIR:-/opt/ray_venvs}" -export CUDA_DEVICE_MAX_CONNECTIONS="${CUDA_DEVICE_MAX_CONNECTIONS:-1}" -export FLASHINFER_DISABLE_VERSION_CHECK="${FLASHINFER_DISABLE_VERSION_CHECK:-1}" -export NVTE_FWD_LAYERNORM_SM_MARGIN="${NVTE_FWD_LAYERNORM_SM_MARGIN:-16}" -export NVTE_BWD_LAYERNORM_SM_MARGIN="${NVTE_BWD_LAYERNORM_SM_MARGIN:-16}" -export NCCL_DEBUG="${NCCL_DEBUG:-WARN}" - -BRIDGE="${CONTAINER_NEMORL}/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge" -export COMMAND="\ -set -euo pipefail -NRL_SLURM_JOB_ID=\$(basename \"\$(dirname \"\$0\")\") -NRL_SLURM_JOB_ID=\${NRL_SLURM_JOB_ID%%-*} -cd ${CONTAINER_NEMORL} -export HF_HOME=${HF_HOME} -export HUGGINGFACE_HUB_CACHE=${HUGGINGFACE_HUB_CACHE} -export HF_HUB_CACHE=${HF_HUB_CACHE} -export HF_DATASETS_CACHE=${HF_DATASETS_CACHE} -export HF_MODULES_CACHE=${HF_MODULES_CACHE} -export TRANSFORMERS_CACHE=${TRANSFORMERS_CACHE} -export NRL_MEGATRON_CHECKPOINT_DIR=${NRL_MEGATRON_CHECKPOINT_DIR} -export XDG_CACHE_HOME=${XDG_CACHE_HOME} -export TORCH_HOME=${TORCH_HOME} -export TRITON_CACHE_DIR=${TRITON_CACHE_DIR} -${REFIT_ENV_EXPORTS} -export PYTHONPATH=${CONTAINER_NEMORL}:${BRIDGE}/src:${BRIDGE}/3rdparty/Megatron-LM\${PYTHONPATH:+:\$PYTHONPATH} -uv run --no-sync python examples/run_vlm_grpo.py --config ${CONFIG} \ -cluster.num_nodes=${NUM_NODES} \ -cluster.gpus_per_node=${GPUS_PER_NODE} \ -policy.model_name=${MODEL_NAME} \ -policy.tokenizer.name=${MODEL_NAME} \ -policy.is_vlm=true \ -policy.megatron_cfg.tensor_model_parallel_size=${POLICY_TP} \ -policy.megatron_cfg.expert_model_parallel_size=${POLICY_EP} \ -policy.megatron_cfg.expert_tensor_parallel_size=1 \ -policy.megatron_cfg.context_parallel_size=${POLICY_CP} \ -policy.megatron_cfg.optimizer.optimizer_cpu_offload=${OPTIMIZER_CPU_OFFLOAD} \ -policy.megatron_cfg.optimizer.optimizer_offload_fraction=${OPTIMIZER_OFFLOAD_FRACTION} \ -policy.offload_optimizer_for_logprob=${OFFLOAD_OPTIMIZER_FOR_LOGPROB} \ -${OPTIMIZER_DTYPE_OVERRIDES} \ -policy.generation.backend=${GENERATION_BACKEND} \ -policy.generation.colocated.enabled=${COLOCATED_ENABLED} \ -policy.generation.colocated.resources.num_nodes=${NUM_GEN_NODES} \ -policy.generation.colocated.resources.gpus_per_node=${GEN_GPUS_PER_NODE} \ -${GEN_OVERRIDES} \ -policy.max_total_sequence_length=${MAX_SEQUENCE_LENGTH} \ -policy.generation.max_new_tokens=${MAX_NEW_TOKENS} \ -grpo.async_grpo.enabled=${ASYNC_GRPO} \ -grpo.async_grpo.max_trajectory_age_steps=${MAX_TRAJECTORY_AGE_STEPS} \ -grpo.async_grpo.in_flight_weight_updates=${IN_FLIGHT_WEIGHT_UPDATES} \ -loss_fn.use_importance_sampling_correction=true \ -grpo.num_prompts_per_step=${NUM_PROMPTS} \ -grpo.num_generations_per_prompt=${NUM_GENERATIONS} \ -grpo.val_batch_size=${VAL_GBS} \ -grpo.max_val_samples=${VAL_SIZE} \ -policy.train_global_batch_size=${TRAIN_GBS} \ -grpo.max_num_steps=${MAX_STEPS} \ -checkpointing.enabled=${CHECKPOINTING_ENABLED} \ -checkpointing.checkpoint_dir=${RESULTS_DIR} \ -logger.log_dir=${RESULTS_DIR} \ -logger.wandb_enabled=${WANDB_ENABLED} \ -logger.wandb.name=${WANDB_NAME}-\${NRL_SLURM_JOB_ID} \ -logger.wandb.project=${WANDB_PROJ} \ -+logger.wandb.entity=${WANDB_GROUP} \ -${EXTRA_OVERRIDES}" - -echo "Submitting ${JOB_NAME}: ${NUM_NODES} node(s), ${GPUS_PER_NODE} GPU(s)/node" -if [[ "${COLOCATED_ENABLED}" == "true" ]]; then - echo " layout: colocated async on all ${NUM_NODES} node(s)" -else - echo " split: ${NUM_GEN_NODES} gen node(s) / $((NUM_NODES - NUM_GEN_NODES)) train node(s)" -fi -echo " generation backend: ${GENERATION_BACKEND} colocated=${COLOCATED_ENABLED} async=${ASYNC_GRPO}" -echo " training world size: ${TRAIN_WORLD_SIZE} (TP=${POLICY_TP}, EP=${POLICY_EP}, DP=${TRAIN_DP_SIZE})" -if [[ "${GENERATION_BACKEND}" == "megatron" ]]; then - echo " inference world size: ${INFERENCE_WORLD_SIZE} (TP=${INFER_TP}, EP=${INFER_EP}, DP=${INFERENCE_DP_SIZE})" - echo " refit backend: ${REFIT_BACKEND}" -else - echo " inference world size: ${INFERENCE_WORLD_SIZE} (TP=${INFER_TP}, EP=${INFER_EP}, DP=${INFERENCE_DP_SIZE})" - echo " vLLM: mem_util=${VLLM_GPU_MEMORY_UTILIZATION:-0.5} eager=${VLLM_ENFORCE_EAGER:-true} max_batched_tokens=${VLLM_MAX_NUM_BATCHED_TOKENS:-${MAX_SEQUENCE_LENGTH}}" - if [[ -n "${REFIT_ENV_EXPORTS}" ]]; then - echo " refit: buffer_memory_ratio=${REFIT_BUFFER_MEMORY_RATIO}" - fi -fi -echo " seq/new_tokens: ${MAX_SEQUENCE_LENGTH}/${MAX_NEW_TOKENS}" -echo " prompts/generations/train_gbs: ${NUM_PROMPTS}/${NUM_GENERATIONS}/${TRAIN_GBS}" -echo " async: max_trajectory_age=${MAX_TRAJECTORY_AGE_STEPS} in_flight_weight_updates=${IN_FLIGHT_WEIGHT_UPDATES}" -echo " val_gbs/val_size: ${VAL_GBS}/${VAL_SIZE}" -echo " optimizer moments: exp_avg=${EXP_AVG_DTYPE:-} exp_avg_sq=${EXP_AVG_SQ_DTYPE:-} store_param_remainders=${STORE_PARAM_REMAINDERS:-}" -echo " W&B: ${WANDB_GROUP}/${WANDB_PROJ}/${WANDB_NAME}- (enabled=${WANDB_ENABLED})" -echo " Hugging Face cache: ${HF_HUB_CACHE}" -echo " MCore checkpoint cache: ${NRL_MEGATRON_CHECKPOINT_DIR}" - -SBATCH_ARGS=( - --nodes="${NUM_NODES}" - --account="${SBATCH_ACCOUNT}" - --partition="${SBATCH_PARTITION}" - --job-name="${JOB_NAME}" - --time="${SBATCH_TIME}" - --output="${SLURM_LOG_DIR}/%j.out" - --error="${SLURM_LOG_DIR}/%j.out" - --gres="gpu:${GPUS_PER_NODE}" - --comment='{"OccupiedIdleGPUsJobReaper":{"exemptIdleTimeMins":"90","reason":"data_loading","description":"Async GRPO RL training: training GPUs idle during rollout collection (~30min) and validation each step"}}' - --exclusive - --mem=0 - --dependency=singleton - --segment="${SBATCH_SEGMENT}" -) -if [[ -n "${SBATCH_QOS}" ]]; then - SBATCH_ARGS+=(--qos="${SBATCH_QOS}") -fi -if [[ -n "${SBATCH_RESERVATION}" ]]; then - SBATCH_ARGS+=(--reservation="${SBATCH_RESERVATION}") -fi - -BASE_LOG_DIR="${SLURM_LOG_DIR}" \ -MOUNTS="${MOUNTS:-/lustre:/lustre},${NEMORL}:${CONTAINER_NEMORL}" \ -sbatch "${SBATCH_ARGS[@]}" "${NEMORL}/ray.sub" diff --git a/scripts/submit_nemotron_omni_vstat_megatron_8n4g.sh b/scripts/submit_nemotron_omni_vstat_megatron_8n4g.sh deleted file mode 100755 index 721b78a873b..00000000000 --- a/scripts/submit_nemotron_omni_vstat_megatron_8n4g.sh +++ /dev/null @@ -1,519 +0,0 @@ -#!/usr/bin/env bash -set -euo pipefail - -# Submit the eight-node Nemotron Omni VSTAT video-GRPO recipe through NeMo-RL's -# Ray/Slurm launcher. The default non-colocated layout reserves six nodes for -# generation and two for training. GENERATION_BACKEND=vllm enables the Omni -# vLLM A/B path; COLOCATED=true shares every GPU between training and generation. - -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -NEMORL="${NEMORL:-$(cd "${SCRIPT_DIR}/.." && pwd)}" -CONTAINER_NEMORL="${CONTAINER_NEMORL:-/opt/nemo-rl}" -WORKSPACE_ROOT="${WORKSPACE_ROOT:-${NEMORL}/workspace}" - -CONTAINER="${CONTAINER:-/lustre/fsw/portfolios/coreai/users/cye/enroot/nemo_rl_0807.sqsh}" -MODEL_NAME="${MODEL_NAME:-nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16}" -CONFIG="${CONFIG:-examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-16n8g-megatron-tp4ep4-async-gym-video.v1.yaml}" -ENTRYPOINT="${ENTRYPOINT:-examples/nemo_gym/run_grpo_nemo_gym.py}" -GENERATION_BACKEND="${GENERATION_BACKEND:-megatron}" -COLOCATED="${COLOCATED:-false}" -ASYNC_GRPO="${ASYNC_GRPO:-true}" - -if [[ "${GENERATION_BACKEND}" != "megatron" && "${GENERATION_BACKEND}" != "vllm" ]]; then - echo "GENERATION_BACKEND must be megatron or vllm (got ${GENERATION_BACKEND})." >&2 - exit 1 -fi -if [[ "${COLOCATED}" == "true" && "${ASYNC_GRPO}" == "true" && "${GENERATION_BACKEND}" != "megatron" ]]; then - echo "Colocated async GRPO requires GENERATION_BACKEND=megatron." >&2 - exit 1 -fi - -CACHE_ROOT="${CACHE_ROOT:-${WORKSPACE_ROOT}/cache/nemo-rl-omni}" -export HF_HOME="${HF_HOME:-${CACHE_ROOT}/huggingface}" -export HUGGINGFACE_HUB_CACHE="${HUGGINGFACE_HUB_CACHE:-${HF_HOME}/hub}" -export HF_HUB_CACHE="${HF_HUB_CACHE:-${HUGGINGFACE_HUB_CACHE}}" -export HF_DATASETS_CACHE="${HF_DATASETS_CACHE:-${HF_HOME}/datasets}" -export HF_MODULES_CACHE="${HF_MODULES_CACHE:-${HF_HOME}/modules}" -export TRANSFORMERS_CACHE="${TRANSFORMERS_CACHE:-${HF_HOME}/transformers}" -export NRL_MEGATRON_CHECKPOINT_DIR="${NRL_MEGATRON_CHECKPOINT_DIR:-${CACHE_ROOT}/megatron-checkpoints}" -export XDG_CACHE_HOME="${XDG_CACHE_HOME:-${CACHE_ROOT}/xdg}" -export TORCH_HOME="${TORCH_HOME:-${CACHE_ROOT}/torch}" -export TRITON_CACHE_DIR="${TRITON_CACHE_DIR:-${CACHE_ROOT}/triton}" - -NUM_NODES="${NUM_NODES:-8}" -GPUS_PER_NODE="${GPUS_PER_NODE:-4}" - -if [[ "${COLOCATED}" == "true" ]]; then - TRAIN_WORLD_SIZE=$((NUM_NODES * GPUS_PER_NODE)) - INFERENCE_WORLD_SIZE="${TRAIN_WORLD_SIZE}" - GEN_GPUS_PER_NODE="${GPUS_PER_NODE}" - NUM_GEN_NODES="${NUM_NODES}" - COLOCATED_ENABLED=true -else - NUM_GEN_NODES="${NUM_GEN_NODES:-6}" - if (( NUM_NODES == 1 )); then - if (( NUM_GEN_NODES != 1 )); then - echo "One-node non-colocated mode requires NUM_GEN_NODES=1." >&2 - exit 1 - fi - GEN_GPUS_PER_NODE="${GEN_GPUS_PER_NODE:-$((GPUS_PER_NODE / 2))}" - if (( GPUS_PER_NODE < 4 || GEN_GPUS_PER_NODE <= 0 || GEN_GPUS_PER_NODE >= GPUS_PER_NODE )); then - echo "One-node non-colocated mode requires at least four GPUs and a strict train/inference split." >&2 - exit 1 - fi - TRAIN_WORLD_SIZE=$((GPUS_PER_NODE - GEN_GPUS_PER_NODE)) - else - GEN_GPUS_PER_NODE="${GEN_GPUS_PER_NODE:-${GPUS_PER_NODE}}" - if (( NUM_GEN_NODES <= 0 || NUM_GEN_NODES >= NUM_NODES )); then - echo "Multi-node non-colocated mode requires 0 < NUM_GEN_NODES < NUM_NODES." >&2 - exit 1 - fi - if (( GEN_GPUS_PER_NODE != GPUS_PER_NODE )); then - echo "Multi-node non-colocated inference must reserve complete GPU nodes." >&2 - exit 1 - fi - TRAIN_WORLD_SIZE=$(((NUM_NODES - NUM_GEN_NODES) * GPUS_PER_NODE)) - fi - INFERENCE_WORLD_SIZE=$((NUM_GEN_NODES * GEN_GPUS_PER_NODE)) - COLOCATED_ENABLED=false -fi - -# Prefer TP=8 and fall back by powers of two for smaller worlds. -DEFAULT_POLICY_TP=8 -while (( DEFAULT_POLICY_TP > TRAIN_WORLD_SIZE || TRAIN_WORLD_SIZE % DEFAULT_POLICY_TP != 0 )); do - DEFAULT_POLICY_TP=$((DEFAULT_POLICY_TP / 2)) -done -DEFAULT_INFER_TP=8 -while (( DEFAULT_INFER_TP > INFERENCE_WORLD_SIZE || INFERENCE_WORLD_SIZE % DEFAULT_INFER_TP != 0 )); do - DEFAULT_INFER_TP=$((DEFAULT_INFER_TP / 2)) -done -POLICY_TP="${POLICY_TP:-${DEFAULT_POLICY_TP}}" -INFER_TP="${INFER_TP:-${DEFAULT_INFER_TP}}" -POLICY_CP="${POLICY_CP:-1}" - -if [[ "${GENERATION_BACKEND}" == "megatron" && "${POLICY_CP}" != "1" ]]; then - echo "Megatron dynamic inference requires POLICY_CP=1." >&2 - exit 1 -fi -if (( TRAIN_WORLD_SIZE % (POLICY_TP * POLICY_CP) != 0 )); then - echo "Training world size must be divisible by POLICY_TP * POLICY_CP." >&2 - exit 1 -fi -if (( INFERENCE_WORLD_SIZE % INFER_TP != 0 )); then - echo "Inference world size must be divisible by INFER_TP." >&2 - exit 1 -fi -TRAIN_DP_SIZE=$((TRAIN_WORLD_SIZE / (POLICY_TP * POLICY_CP))) -INFERENCE_DP_SIZE=$((INFERENCE_WORLD_SIZE / INFER_TP)) - -# With ETP=1, world_size must be divisible by EP (not by TP*EP). -largest_ep() { - local world="$1" - local ep=8 - while (( ep > world || world % ep != 0 )); do - ep=$((ep / 2)) - done - printf '%d' "${ep}" -} - -POLICY_EP="${POLICY_EP:-$(largest_ep "${TRAIN_WORLD_SIZE}")}" -INFER_EP="${INFER_EP:-$(largest_ep "${INFERENCE_WORLD_SIZE}")}" -if (( TRAIN_WORLD_SIZE % POLICY_EP != 0 )); then - echo "Training world size must be divisible by POLICY_EP (ETP=1)." >&2 - exit 1 -fi -if [[ "${GENERATION_BACKEND}" == "vllm" ]]; then - # Async vLLM cannot use internal DP. EP=TP leaves replica-level DP to Ray. - if [[ "${INFER_EP}" != "${INFER_TP}" ]]; then - echo "Forcing INFER_EP=${INFER_TP} for vLLM (was ${INFER_EP}) for async-safe Ray DP." >&2 - fi - INFER_EP="${INFER_TP}" -elif (( INFERENCE_WORLD_SIZE % INFER_EP != 0 )); then - echo "Inference world size must be divisible by INFER_EP (ETP=1)." >&2 - exit 1 -fi -if [[ "${COLOCATED_ENABLED}" == "true" && "${GENERATION_BACKEND}" == "megatron" ]] && - (( POLICY_TP != INFER_TP || POLICY_EP != INFER_EP )); then - echo "Colocated Megatron requires matching policy/inference TP and EP topology." >&2 - exit 1 -fi - -MAX_STEPS="${MAX_STEPS:-1000000}" -MAX_SEQUENCE_LENGTH="${MAX_SEQUENCE_LENGTH:-8192}" -MAX_NEW_TOKENS="${MAX_NEW_TOKENS:-2048}" -MIN_GENERATION_TOKENS="${MIN_GENERATION_TOKENS:-2000}" -VISION_EMBEDDING_CACHE_MAX_BYTES="${VISION_EMBEDDING_CACHE_MAX_BYTES:-536870912}" -ENABLE_THINKING="${ENABLE_THINKING:-true}" -NUM_FRAMES="${NUM_FRAMES:-16}" -TEMPORAL_PATCH_SIZE="${TEMPORAL_PATCH_SIZE:-2}" -# Match the checkpoint-native vLLM video processor. Policy preprocessing and -# Megatron inference consume this same value below. -VIDEO_TARGET_PATCHES="${VIDEO_TARGET_PATCHES:-1024}" -NUM_DATA_ROWS="${NUM_DATA_ROWS:-256}" -NUM_PROMPTS="${NUM_PROMPTS:-$((INFERENCE_DP_SIZE * 2))}" -NUM_GENERATIONS="${NUM_GENERATIONS:-8}" -TRAIN_GBS="${TRAIN_GBS:-$((NUM_PROMPTS * NUM_GENERATIONS))}" -EXPECTED_TRAIN_GBS=$((NUM_PROMPTS * NUM_GENERATIONS)) -if (( TRAIN_GBS != EXPECTED_TRAIN_GBS )); then - echo "TRAIN_GBS (${TRAIN_GBS}) must equal NUM_PROMPTS * NUM_GENERATIONS (${EXPECTED_TRAIN_GBS})." >&2 - exit 1 -fi -if (( TRAIN_GBS % TRAIN_DP_SIZE != 0 )); then - echo "TRAIN_GBS (${TRAIN_GBS}) must be divisible by training DP size (${TRAIN_DP_SIZE})." >&2 - exit 1 -fi - -DATA_ROOT="${DATA_ROOT:-${CONTAINER_NEMORL}/workspace/datasets/vstat-8n4g}" -HF_DATASET="${HF_DATASET:-ShushengYang/VSTAT}" -PREPARE_VSTAT="${PREPARE_VSTAT:-false}" -export NEMO_RL_VIDEO_TRAIN_JSONL="${NEMO_RL_VIDEO_TRAIN_JSONL:-${DATA_ROOT}/train-gym.jsonl}" -export NEMO_RL_VIDEO_VAL_JSONL="${NEMO_RL_VIDEO_VAL_JSONL:-${DATA_ROOT}/val-gym.jsonl}" -export NEMO_RL_VIDEO_MEDIA_ROOT="${NEMO_RL_VIDEO_MEDIA_ROOT:-${DATA_ROOT}/media}" - -REFIT_BACKEND="${REFIT_BACKEND:-nccl}" -OPTIMIZER_CPU_OFFLOAD="${OPTIMIZER_CPU_OFFLOAD:-false}" -BUFFER_SIZE_GB="${BUFFER_SIZE_GB:-8}" -OFFLOAD_OPTIMIZER_FOR_LOGPROB="${OFFLOAD_OPTIMIZER_FOR_LOGPROB:-false}" -if [[ "${OPTIMIZER_CPU_OFFLOAD}" == "true" ]]; then - OPTIMIZER_OFFLOAD_FRACTION="${OPTIMIZER_OFFLOAD_FRACTION:-1.0}" -else - OPTIMIZER_OFFLOAD_FRACTION="${OPTIMIZER_OFFLOAD_FRACTION:-0.0}" -fi -USE_PRECISION_AWARE_OPTIMIZER="${USE_PRECISION_AWARE_OPTIMIZER:-true}" -EXP_AVG_DTYPE="${EXP_AVG_DTYPE:-bfloat16}" -EXP_AVG_SQ_DTYPE="${EXP_AVG_SQ_DTYPE:-bfloat16}" -STORE_PARAM_REMAINDERS="${STORE_PARAM_REMAINDERS:-true}" -MEGATRON_ENABLE_CHUNKED_PREFILL="${MEGATRON_ENABLE_CHUNKED_PREFILL:-true}" -MEGATRON_CUDA_GRAPH_IMPL="${MEGATRON_CUDA_GRAPH_IMPL:-local}" -MEGATRON_CUDA_GRAPH_SCOPE="${MEGATRON_CUDA_GRAPH_SCOPE:-block}" -MEGATRON_NUM_CUDA_GRAPHS="${MEGATRON_NUM_CUDA_GRAPHS:--1}" -MEGATRON_USE_CUDA_GRAPHS_FOR_NON_DECODE="${MEGATRON_USE_CUDA_GRAPHS_FOR_NON_DECODE:-false}" -MOE_PAD_EXPERTS_FOR_CG="${MOE_PAD_EXPERTS_FOR_CG:-true}" -MAX_TRAJECTORY_AGE_STEPS="${MAX_TRAJECTORY_AGE_STEPS:-2}" -IN_FLIGHT_WEIGHT_UPDATES="${IN_FLIGHT_WEIGHT_UPDATES:-true}" -EXTRA_OVERRIDES="${EXTRA_OVERRIDES:-}" - -JOB_NAME="${JOB_NAME:-nemotron-omni-vstat-${GENERATION_BACKEND}-8n4g}" -EXP_NAME="${EXP_NAME:-${JOB_NAME}}" -PRECISION_RECIPE="${PRECISION_RECIPE:-bf16}" -WANDB_ENABLED="${WANDB_ENABLED:-true}" -WANDB_PROJ="${WANDB_PROJ:-mllm-rl-dev}" -WANDB_GROUP="${WANDB_GROUP:-adlr}" -WANDB_NAME="${WANDB_NAME:-${EXP_NAME}-${PRECISION_RECIPE}-internal-repo}" -RESULTS_DIR="${RESULTS_DIR:-${WORKSPACE_ROOT}/results/nemo-rl-omni/${JOB_NAME}}" -CHECKPOINTING_ENABLED="${CHECKPOINTING_ENABLED:-false}" -SLURM_LOG_DIR="${SLURM_LOG_DIR:-${RESULTS_DIR}/slurm}" - -SBATCH_ACCOUNT="${SBATCH_ACCOUNT:-nemotron_sw_post}" -SBATCH_PARTITION="${SBATCH_PARTITION:-batch_long}" -SBATCH_QOS="${SBATCH_QOS:-}" -SBATCH_TIME="${SBATCH_TIME:-04:00:00}" -SBATCH_RESERVATION="${SBATCH_RESERVATION:-}" -SBATCH_SEGMENT="${NUM_NODES}" - -mkdir -p \ - "${HF_HUB_CACHE}" \ - "${HF_DATASETS_CACHE}" \ - "${HF_MODULES_CACHE}" \ - "${TRANSFORMERS_CACHE}" \ - "${NRL_MEGATRON_CHECKPOINT_DIR}" \ - "${XDG_CACHE_HOME}" \ - "${TORCH_HOME}" \ - "${TRITON_CACHE_DIR}" \ - "${SLURM_LOG_DIR}" - -if [[ ! -f "${CONTAINER}" ]]; then - echo "Container image does not exist: ${CONTAINER}" >&2 - exit 1 -fi -if [[ ! -f "${NEMORL}/ray.sub" || ! -f "${NEMORL}/${CONFIG}" || - ! -f "${NEMORL}/${ENTRYPOINT}" || ! -f "${NEMORL}/scripts/prepare_nemotron_omni_vstat.py" ]]; then - echo "NeMo-RL launcher, config, entrypoint, or VSTAT preparation script is missing under: ${NEMORL}" >&2 - exit 1 -fi - -OPTIMIZER_DTYPE_OVERRIDES="" -if [[ -n "${EXP_AVG_DTYPE}" ]]; then - OPTIMIZER_DTYPE_OVERRIDES+=" ++policy.megatron_cfg.optimizer.exp_avg_dtype=${EXP_AVG_DTYPE}" -fi -if [[ -n "${EXP_AVG_SQ_DTYPE}" ]]; then - OPTIMIZER_DTYPE_OVERRIDES+=" ++policy.megatron_cfg.optimizer.exp_avg_sq_dtype=${EXP_AVG_SQ_DTYPE}" -fi -if [[ -n "${STORE_PARAM_REMAINDERS}" ]]; then - OPTIMIZER_DTYPE_OVERRIDES+=" ++policy.megatron_cfg.optimizer.store_param_remainders=${STORE_PARAM_REMAINDERS}" -fi - -REFIT_ENV_EXPORTS="" -REFIT_BUFFER_MEMORY_RATIO="" -if [[ "${GENERATION_BACKEND}" == "megatron" ]]; then - GEN_OVERRIDES="\ -++policy.generation.stop_strings=null \ -++policy.generation.bad_words=null \ -policy.generation.mcore_generation_config.tensor_model_parallel_size=${INFER_TP} \ -policy.generation.mcore_generation_config.expert_model_parallel_size=${INFER_EP} \ -policy.generation.mcore_generation_config.expert_tensor_parallel_size=1 \ -++policy.generation.mcore_generation_config.context_parallel_size=${POLICY_CP} \ -++policy.generation.mcore_generation_config.mamba_inference_ssm_states_dtype=float32 \ -++policy.generation.mcore_generation_config.mamba_inference_conv_states_dtype=float32 \ -++policy.generation.mcore_generation_config.logprobs_mode=raw_logprobs \ -policy.generation.mcore_generation_config.parsers=[nemotron-v3-reasoning,qwen3-coder-tool] \ -policy.generation.mcore_generation_config.transformer_impl=transformer_engine \ -policy.generation.mcore_generation_config.sequence_parallel=true \ -policy.generation.mcore_generation_config.enable_chunked_prefill=${MEGATRON_ENABLE_CHUNKED_PREFILL} \ -++policy.generation.mcore_generation_config.async_sched_mode=async \ -policy.generation.mcore_generation_config.enable_prefix_caching=false \ -++policy.generation.mcore_generation_config.vision_embedding_cache_max_bytes=${VISION_EMBEDDING_CACHE_MAX_BYTES} \ -policy.generation.mcore_generation_config.cuda_graph_impl=${MEGATRON_CUDA_GRAPH_IMPL} \ -policy.generation.mcore_generation_config.inference_cuda_graph_scope=${MEGATRON_CUDA_GRAPH_SCOPE} \ -policy.generation.mcore_generation_config.num_cuda_graphs=${MEGATRON_NUM_CUDA_GRAPHS} \ -policy.generation.mcore_generation_config.use_cuda_graphs_for_non_decode_steps=${MEGATRON_USE_CUDA_GRAPHS_FOR_NON_DECODE} \ -++policy.generation.mcore_generation_config.moe_pad_experts_for_cuda_graph_inference=${MOE_PAD_EXPERTS_FOR_CG} \ -policy.generation.mcore_generation_config.refit_backend=${REFIT_BACKEND} \ -policy.generation.mcore_generation_config.buffer_size_gb=${BUFFER_SIZE_GB} \ -policy.generation.mcore_generation_config.max_model_len=${MAX_SEQUENCE_LENGTH} \ -policy.generation.mcore_generation_config.max_tokens=${MAX_SEQUENCE_LENGTH} \ -++policy.generation.mcore_generation_config.video_num_frames=${NUM_FRAMES} \ -++policy.generation.mcore_generation_config.video_temporal_patch_size=${TEMPORAL_PATCH_SIZE} \ -++policy.generation.mcore_generation_config.video_target_num_patches=${VIDEO_TARGET_PATCHES}" -else - REFIT_BUFFER_MEMORY_RATIO="${NRL_REFIT_BUFFER_MEMORY_RATIO:-0.005}" - REFIT_ENV_EXPORTS="export NRL_REFIT_BUFFER_MEMORY_RATIO=${REFIT_BUFFER_MEMORY_RATIO}" - VLLM_GPU_MEMORY_UTILIZATION="${VLLM_GPU_MEMORY_UTILIZATION:-0.5}" - VLLM_ENFORCE_EAGER="${VLLM_ENFORCE_EAGER:-true}" - VLLM_MAX_NUM_SEQS="${VLLM_MAX_NUM_SEQS:-4}" - VLLM_MAX_NUM_BATCHED_TOKENS="${VLLM_MAX_NUM_BATCHED_TOKENS:-${MAX_SEQUENCE_LENGTH}}" - GEN_OVERRIDES="\ -++policy.generation.stop_strings=null \ -++policy.generation.bad_words=null \ -policy.generation.vllm_cfg.async_engine=${ASYNC_GRPO} \ -policy.generation.vllm_cfg.skip_tokenizer_init=false \ -policy.generation.vllm_cfg.tensor_parallel_size=${INFER_TP} \ -policy.generation.vllm_cfg.pipeline_parallel_size=1 \ -policy.generation.vllm_cfg.expert_parallel_size=${INFER_EP} \ -policy.generation.vllm_cfg.max_model_len=${MAX_SEQUENCE_LENGTH} \ -++policy.generation.vllm_cfg.cap_max_tokens_to_context=true \ -policy.generation.vllm_cfg.gpu_memory_utilization=${VLLM_GPU_MEMORY_UTILIZATION} \ -policy.generation.vllm_cfg.enforce_eager=${VLLM_ENFORCE_EAGER} \ -policy.generation.vllm_cfg.enable_prefix_caching=false \ -policy.generation.vllm_cfg.logprobs_mode=raw_logprobs \ -policy.generation.vllm_cfg.reset_mm_cache_after_refit=true \ -policy.generation.vllm_cfg.env_vars.NRL_VIDEO_BACKEND=torchcodec \ -policy.generation.vllm_cfg.env_vars.NRL_VIDEO_SAMPLING_STYLE=nemotron_vl \ -policy.generation.vllm_cfg.env_vars.NRL_VIDEO_TEMPORAL_PATCH_SIZE=${TEMPORAL_PATCH_SIZE} \ -policy.generation.vllm_cfg.env_vars.VLLM_VIDEO_LOADER_BACKEND=nemotron_vl \ -policy.generation.vllm_kwargs.allowed_local_media_path=${NEMO_RL_VIDEO_MEDIA_ROOT} \ -policy.generation.vllm_kwargs.mm_processor_cache_gb=0 \ -policy.generation.vllm_kwargs.max_num_seqs=${VLLM_MAX_NUM_SEQS} \ -policy.generation.vllm_kwargs.limit_mm_per_prompt.video.count=1 \ -policy.generation.vllm_kwargs.limit_mm_per_prompt.video.num_frames=${NUM_FRAMES} \ -++policy.generation.vllm_kwargs.limit_mm_per_prompt.image=${NUM_FRAMES} \ -policy.generation.vllm_kwargs.max_num_batched_tokens=${VLLM_MAX_NUM_BATCHED_TOKENS} \ -policy.generation.vllm_kwargs.enable_chunked_prefill=false \ -policy.generation.vllm_kwargs.disable_custom_all_reduce=true \ -policy.generation.vllm_kwargs.attention_backend=FLASH_ATTN \ -policy.generation.vllm_kwargs.attention_config.use_trtllm_attention=false \ -++policy.generation.vllm_kwargs.mamba_ssm_cache_dtype=float32 \ -++policy.generation.vllm_kwargs.skip_mm_profiling=false \ -++policy.generation.vllm_kwargs.kernel_config.enable_flashinfer_autotune=false \ -++policy.generation.vllm_kwargs.kernel_config.moe_backend=triton" -fi - -export NUM_NODES GPUS_PER_NODE CONTAINER -export RAY_ENABLE_UV_RUN_RUNTIME_ENV=0 -export NRL_VENVS_TRUST_EXISTING="${NRL_VENVS_TRUST_EXISTING:-1}" -export NRL_FORCE_REBUILD_VENVS="${NRL_FORCE_REBUILD_VENVS:-false}" -export NEMO_RL_VENV_DIR="${NEMO_RL_VENV_DIR:-/opt/ray_venvs}" -export NEMO_GYM_VENV_DIR="${NEMO_GYM_VENV_DIR:-/opt/ray_venvs}" -export NEMO_GYM_EXTRA_ROOTS="${NEMO_GYM_EXTRA_ROOTS:-${CONTAINER_NEMORL}/3rdparty/Gym-workspace/Gym}" -export NRL_VIDEO_BACKEND="${NRL_VIDEO_BACKEND:-torchcodec}" -export NRL_VIDEO_SAMPLING_STYLE="${NRL_VIDEO_SAMPLING_STYLE:-nemotron_vl}" -export NRL_VIDEO_TEMPORAL_PATCH_SIZE="${TEMPORAL_PATCH_SIZE}" -export VLLM_VIDEO_LOADER_BACKEND="${VLLM_VIDEO_LOADER_BACKEND:-nemotron_vl}" -export TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST:-10.0}" -export CUDA_DEVICE_MAX_CONNECTIONS="${CUDA_DEVICE_MAX_CONNECTIONS:-1}" -export FLASHINFER_DISABLE_VERSION_CHECK="${FLASHINFER_DISABLE_VERSION_CHECK:-1}" -export NVTE_FWD_LAYERNORM_SM_MARGIN="${NVTE_FWD_LAYERNORM_SM_MARGIN:-16}" -export NVTE_BWD_LAYERNORM_SM_MARGIN="${NVTE_BWD_LAYERNORM_SM_MARGIN:-16}" -export NCCL_DEBUG="${NCCL_DEBUG:-WARN}" - -ENABLE_NSYS="${ENABLE_NSYS:-false}" -if [[ "${ENABLE_NSYS}" == "true" ]]; then - export NRL_NSYS_WORKER_PATTERNS="${NRL_NSYS_WORKER_PATTERNS:-*policy*,*megatron*,*vllm*}" - export NRL_NSYS_PROFILE_STEP_RANGE="${NRL_NSYS_PROFILE_STEP_RANGE:-1:4}" - export LD_LIBRARY_PATH="/usr/local/cuda/targets/aarch64-linux/lib:/usr/local/cuda/targets/x86_64-linux/lib:/usr/local/cuda/lib64:/usr/local/cuda/lib:/usr/local/nvidia/lib64:/usr/local/nvidia/lib:/usr/lib/aarch64-linux-gnu:/usr/lib/x86_64-linux-gnu:${LD_LIBRARY_PATH:-}" - export NRL_NSYS_EXTRA_OPTIONS="${NRL_NSYS_EXTRA_OPTIONS:-{\"o\":\"${CONTAINER_NEMORL}/workspace/nsys/%p\",\"cpuctxsw\":\"none\",\"force-overwrite\":\"true\"}}" - mkdir -p "${WORKSPACE_ROOT}/nsys" -fi - -BRIDGE="${CONTAINER_NEMORL}/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge" -export PYTHONPATH="${CONTAINER_NEMORL}:${NEMO_GYM_EXTRA_ROOTS}:${BRIDGE}/src:${BRIDGE}/3rdparty/Megatron-LM${PYTHONPATH:+:${PYTHONPATH}}" -export SETUP_COMMAND="\ -set -euo pipefail -cd ${CONTAINER_NEMORL} -bash tools/install_audio_deps.sh -if [[ ${GENERATION_BACKEND} == vllm ]]; then - for vllm_python in ${NEMO_RL_VENV_DIR}/nemo_rl.models.generation.vllm.*Vllm*Worker/bin/python; do - [[ -x \"\${vllm_python}\" ]] || continue - \"\${vllm_python}\" -m pip install --no-cache-dir \ - --index-url https://download.pytorch.org/whl/cu130 \ - --extra-index-url https://pypi.org/simple \ - torchcodec==0.11.1 - done -fi" - -export COMMAND="\ -set -euo pipefail -NRL_SLURM_JOB_ID=\$(basename \"\$(dirname \"\$0\")\") -NRL_SLURM_JOB_ID=\${NRL_SLURM_JOB_ID%%-*} -cd ${CONTAINER_NEMORL} -export HF_HOME=${HF_HOME} -export HUGGINGFACE_HUB_CACHE=${HUGGINGFACE_HUB_CACHE} -export HF_HUB_CACHE=${HF_HUB_CACHE} -export HF_DATASETS_CACHE=${HF_DATASETS_CACHE} -export HF_MODULES_CACHE=${HF_MODULES_CACHE} -export TRANSFORMERS_CACHE=${TRANSFORMERS_CACHE} -export NRL_MEGATRON_CHECKPOINT_DIR=${NRL_MEGATRON_CHECKPOINT_DIR} -export XDG_CACHE_HOME=${XDG_CACHE_HOME} -export TORCH_HOME=${TORCH_HOME} -export TRITON_CACHE_DIR=${TRITON_CACHE_DIR} -export NEMO_RL_VENV_DIR=${NEMO_RL_VENV_DIR} -export NEMO_GYM_VENV_DIR=${NEMO_GYM_VENV_DIR} -export NEMO_GYM_EXTRA_ROOTS=${NEMO_GYM_EXTRA_ROOTS} -export NRL_VIDEO_BACKEND=${NRL_VIDEO_BACKEND} -export NRL_VIDEO_SAMPLING_STYLE=${NRL_VIDEO_SAMPLING_STYLE} -export NRL_VIDEO_TEMPORAL_PATCH_SIZE=${NRL_VIDEO_TEMPORAL_PATCH_SIZE} -export VLLM_VIDEO_LOADER_BACKEND=${VLLM_VIDEO_LOADER_BACKEND} -export TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST} -export NEMO_RL_VIDEO_TRAIN_JSONL=${NEMO_RL_VIDEO_TRAIN_JSONL} -export NEMO_RL_VIDEO_VAL_JSONL=${NEMO_RL_VIDEO_VAL_JSONL} -export NEMO_RL_VIDEO_MEDIA_ROOT=${NEMO_RL_VIDEO_MEDIA_ROOT} -${REFIT_ENV_EXPORTS} -export PYTHONPATH=${CONTAINER_NEMORL}:${NEMO_GYM_EXTRA_ROOTS}:${BRIDGE}/src:${BRIDGE}/3rdparty/Megatron-LM\${PYTHONPATH:+:\$PYTHONPATH} -mkdir -p ${DATA_ROOT} -if [[ ${PREPARE_VSTAT} == true || ! -s \${NEMO_RL_VIDEO_TRAIN_JSONL} || ! -s \${NEMO_RL_VIDEO_VAL_JSONL} ]]; then - echo \"Preparing VSTAT inside the mounted container path: ${DATA_ROOT}\" - uv run --no-sync python scripts/prepare_nemotron_omni_vstat.py --output-dir ${DATA_ROOT} --repo-id ${HF_DATASET} --num-rows ${NUM_DATA_ROWS} -fi -uv run --no-sync python ${ENTRYPOINT} --config ${CONFIG} \ -cluster.num_nodes=${NUM_NODES} \ -cluster.gpus_per_node=${GPUS_PER_NODE} \ -policy.model_name=${MODEL_NAME} \ -policy.tokenizer.name=${MODEL_NAME} \ -policy.tokenizer.chat_template_kwargs.enable_thinking=${ENABLE_THINKING} \ -policy.is_vlm=true \ -policy.megatron_cfg.env_vars.TORCH_CUDA_ARCH_LIST=\"'${TORCH_CUDA_ARCH_LIST}'\" \ -policy.megatron_cfg.freeze_vision_model=false \ -policy.megatron_cfg.freeze_vision_projection=false \ -policy.megatron_cfg.freeze_moe_router=false \ -policy.megatron_cfg.mtp_num_layers=0 \ -policy.megatron_cfg.mtp_use_repeated_layer=true \ -policy.megatron_cfg.mtp_detach_heads=true \ -policy.megatron_cfg.mtp_loss_scaling_factor=0.0 \ -policy.megatron_cfg.tensor_model_parallel_size=${POLICY_TP} \ -policy.megatron_cfg.pipeline_model_parallel_size=1 \ -policy.megatron_cfg.expert_model_parallel_size=${POLICY_EP} \ -policy.megatron_cfg.expert_tensor_parallel_size=1 \ -policy.megatron_cfg.context_parallel_size=${POLICY_CP} \ -policy.megatron_cfg.sequence_parallel=true \ -policy.megatron_cfg.bias_activation_fusion=false \ -policy.megatron_cfg.moe_shared_expert_overlap=false \ -policy.megatron_cfg.radio_force_cpe_eval_mode=true \ -policy.megatron_cfg.clear_memory_caches_before_refit=true \ -policy.megatron_cfg.distributed_data_parallel_config.overlap_grad_reduce=false \ -policy.megatron_cfg.distributed_data_parallel_config.overlap_param_gather=false \ -policy.megatron_cfg.optimizer.params_dtype=float32 \ -policy.megatron_cfg.optimizer.use_precision_aware_optimizer=${USE_PRECISION_AWARE_OPTIMIZER} \ -policy.megatron_cfg.optimizer.optimizer_cpu_offload=${OPTIMIZER_CPU_OFFLOAD} \ -policy.megatron_cfg.optimizer.optimizer_offload_fraction=${OPTIMIZER_OFFLOAD_FRACTION} \ -policy.offload_optimizer_for_logprob=${OFFLOAD_OPTIMIZER_FOR_LOGPROB} \ -${OPTIMIZER_DTYPE_OVERRIDES} \ -policy.generation.backend=${GENERATION_BACKEND} \ -policy.generation.colocated.enabled=${COLOCATED_ENABLED} \ -policy.generation.colocated.resources.num_nodes=${NUM_GEN_NODES} \ -policy.generation.colocated.resources.gpus_per_node=${GEN_GPUS_PER_NODE} \ -${GEN_OVERRIDES} \ -policy.max_total_sequence_length=${MAX_SEQUENCE_LENGTH} \ -policy.generation.max_new_tokens=${MAX_NEW_TOKENS} \ -data.max_input_seq_length=${MAX_SEQUENCE_LENGTH} \ -data.num_workers=0 \ -+data.default.num_frames=${NUM_FRAMES} \ -+data.default.video_sampling_style=nemotron_vl \ -+data.default.video_temporal_patch_size=${TEMPORAL_PATCH_SIZE} \ -+data.default.min_generation_tokens=${MIN_GENERATION_TOKENS} \ -data.default.video_target_num_patches=${VIDEO_TARGET_PATCHES} \ -data.default.video_maintain_aspect_ratio=true \ -data.train.data_path=${NEMO_RL_VIDEO_TRAIN_JSONL} \ -data.validation.data_path=${NEMO_RL_VIDEO_VAL_JSONL} \ -++env.nemo_gym.policy_model.responses_api_models.vllm_model.chat_template_kwargs.enable_thinking=${ENABLE_THINKING} \ -grpo.deduplicate_multimodal_data=false \ -grpo.async_grpo.enabled=${ASYNC_GRPO} \ -grpo.async_grpo.max_trajectory_age_steps=${MAX_TRAJECTORY_AGE_STEPS} \ -grpo.async_grpo.in_flight_weight_updates=${IN_FLIGHT_WEIGHT_UPDATES} \ -loss_fn.use_importance_sampling_correction=true \ -grpo.num_prompts_per_step=${NUM_PROMPTS} \ -grpo.num_generations_per_prompt=${NUM_GENERATIONS} \ -grpo.val_num_generations_per_prompt=1 \ -policy.train_global_batch_size=${TRAIN_GBS} \ -grpo.max_num_steps=${MAX_STEPS} \ -checkpointing.enabled=${CHECKPOINTING_ENABLED} \ -checkpointing.checkpoint_dir=${RESULTS_DIR} \ -logger.log_dir=${RESULTS_DIR} \ -logger.wandb_enabled=${WANDB_ENABLED} \ -logger.wandb.name=${WANDB_NAME}-\${NRL_SLURM_JOB_ID} \ -logger.wandb.project=${WANDB_PROJ} \ -+logger.wandb.entity=${WANDB_GROUP} \ -${EXTRA_OVERRIDES}" - -echo "Submitting ${JOB_NAME}: ${NUM_NODES} node(s), ${GPUS_PER_NODE} GPU(s)/node" -if [[ "${COLOCATED_ENABLED}" == "true" ]]; then - echo " layout: colocated async on all ${NUM_NODES} node(s)" -else - echo " split: ${NUM_GEN_NODES} gen node(s) / $((NUM_NODES - NUM_GEN_NODES)) train node(s)" -fi -echo " generation backend: ${GENERATION_BACKEND} colocated=${COLOCATED_ENABLED} async=${ASYNC_GRPO}" -echo " training world size: ${TRAIN_WORLD_SIZE} (TP=${POLICY_TP}, EP=${POLICY_EP}, DP=${TRAIN_DP_SIZE})" -echo " inference world size: ${INFERENCE_WORLD_SIZE} (TP=${INFER_TP}, EP=${INFER_EP}, DP=${INFERENCE_DP_SIZE})" -if [[ "${GENERATION_BACKEND}" == "megatron" ]]; then - echo " Megatron: refit=${REFIT_BACKEND} chunked_prefill=${MEGATRON_ENABLE_CHUNKED_PREFILL} CUDA graphs=${MEGATRON_CUDA_GRAPH_IMPL}/${MEGATRON_CUDA_GRAPH_SCOPE}" -else - echo " vLLM: mem_util=${VLLM_GPU_MEMORY_UTILIZATION} eager=${VLLM_ENFORCE_EAGER} max_seqs=${VLLM_MAX_NUM_SEQS} max_batched_tokens=${VLLM_MAX_NUM_BATCHED_TOKENS}" - echo " refit buffer memory ratio: ${REFIT_BUFFER_MEMORY_RATIO}" -fi -echo " seq/new_tokens: ${MAX_SEQUENCE_LENGTH}/${MAX_NEW_TOKENS}" -echo " thinking enabled: ${ENABLE_THINKING}" -echo " prompts/generations/train_gbs: ${NUM_PROMPTS}/${NUM_GENERATIONS}/${TRAIN_GBS}" -echo " async: max_trajectory_age=${MAX_TRAJECTORY_AGE_STEPS} in_flight_weight_updates=${IN_FLIGHT_WEIGHT_UPDATES}" -echo " video: frames=${NUM_FRAMES} temporal_patch=${TEMPORAL_PATCH_SIZE} target_patches=${VIDEO_TARGET_PATCHES}" -echo " VSTAT: root=${DATA_ROOT} repo=${HF_DATASET} rows=${NUM_DATA_ROWS} prepare=${PREPARE_VSTAT} (head only; also runs if JSONL is missing)" -echo " datasets: train=${NEMO_RL_VIDEO_TRAIN_JSONL} val=${NEMO_RL_VIDEO_VAL_JSONL} media=${NEMO_RL_VIDEO_MEDIA_ROOT}" -echo " dependencies: targeted audio/video dependency setup runs on every node before Ray" -echo " NSYS: enabled=${ENABLE_NSYS}${NRL_NSYS_PROFILE_STEP_RANGE:+ step_range=${NRL_NSYS_PROFILE_STEP_RANGE}}" -echo " W&B: ${WANDB_GROUP}/${WANDB_PROJ}/${WANDB_NAME}- (enabled=${WANDB_ENABLED})" -echo " Hugging Face cache: ${HF_HUB_CACHE}" -echo " MCore checkpoint cache: ${NRL_MEGATRON_CHECKPOINT_DIR}" - -SBATCH_ARGS=( - --nodes="${NUM_NODES}" - --account="${SBATCH_ACCOUNT}" - --partition="${SBATCH_PARTITION}" - --job-name="${JOB_NAME}" - --time="${SBATCH_TIME}" - --output="${SLURM_LOG_DIR}/%j.out" - --error="${SLURM_LOG_DIR}/%j.out" - --gres="gpu:${GPUS_PER_NODE}" - --comment='{"OccupiedIdleGPUsJobReaper":{"exemptIdleTimeMins":"90","reason":"data_loading","description":"Async GRPO RL training: training GPUs idle during rollout collection (~30min) and validation each step"}}' - --exclusive - --mem=0 - --dependency=singleton - --segment="${SBATCH_SEGMENT}" -) -if [[ -n "${SBATCH_QOS}" ]]; then - SBATCH_ARGS+=(--qos="${SBATCH_QOS}") -fi -if [[ -n "${SBATCH_RESERVATION}" ]]; then - SBATCH_ARGS+=(--reservation="${SBATCH_RESERVATION}") -fi - -BASE_LOG_DIR="${SLURM_LOG_DIR}" \ -MOUNTS="${MOUNTS:-/lustre:/lustre},${NEMORL}:${CONTAINER_NEMORL}" \ -sbatch "${SBATCH_ARGS[@]}" "${NEMORL}/ray.sub" From 8cf6154f318fff07f9cf810fe97a83f6c4d54ee8 Mon Sep 17 00:00:00 2001 From: Cory Ye Date: Mon, 24 Aug 2026 11:06:14 -0700 Subject: [PATCH 05/15] Deprecate media_cache_key. Signed-off-by: Cory Ye --- nemo_rl/models/generation/megatron/config.py | 1 + .../models/generation/megatron/megatron_worker.py | 14 ++++---------- nemo_rl/models/generation/megatron/utils.py | 6 +++++- 3 files changed, 10 insertions(+), 11 deletions(-) diff --git a/nemo_rl/models/generation/megatron/config.py b/nemo_rl/models/generation/megatron/config.py index 41d8a91cb46..0927024ae03 100644 --- a/nemo_rl/models/generation/megatron/config.py +++ b/nemo_rl/models/generation/megatron/config.py @@ -65,6 +65,7 @@ class MCoreGenerationSpecificArgs(TypedDict): video_num_frames: NotRequired[int] video_temporal_patch_size: NotRequired[int] video_target_num_patches: NotRequired[int] + video_maintain_aspect_ratio: NotRequired[bool] multimodal_prompt_config: NotRequired[dict] megatron_inference_wrapper: NotRequired[str] diff --git a/nemo_rl/models/generation/megatron/megatron_worker.py b/nemo_rl/models/generation/megatron/megatron_worker.py index f35974f2489..807114da359 100644 --- a/nemo_rl/models/generation/megatron/megatron_worker.py +++ b/nemo_rl/models/generation/megatron/megatron_worker.py @@ -347,6 +347,9 @@ def _initialize_inference_engine(self, mcore_generation_config: dict) -> None: image_preprocessing_config, mcore_generation_config, frame_manifest_magic=CACHED_VIDEO_FRAME_MANIFEST_MAGIC, + video_maintain_aspect_ratio=bool( + mcore_generation_config.get("video_maintain_aspect_ratio", True) + ), ) inference_config = InferenceConfig( @@ -805,12 +808,6 @@ def _build_prompt_and_multimodal_data(self, data, index: int): length = int(data["input_lengths"][index].item()) prompt = data["input_ids"][index, :length].tolist() imgs, imgs_sizes, num_frames = self._sample_vision_tensors(data, index) - media_cache_keys = data.get("media_cache_key") - media_cache_key = ( - media_cache_keys[index] if media_cache_keys is not None else None - ) - if media_cache_key is not None and not isinstance(media_cache_key, str): - raise TypeError("media_cache_key entries must be strings or None.") if imgs is None: return prompt, None @@ -837,13 +834,10 @@ def _build_prompt_and_multimodal_data(self, data, index: int): } else: modality_data = {"imgs": imgs, "imgs_sizes": imgs_sizes} - multi_modal_data: dict[str, Any] = { + return prompt, { modality: modality_data, "media_tokens_preexpanded": True, } - if media_cache_key is not None: - multi_modal_data["media_cache_key"] = media_cache_key - return prompt, multi_modal_data def _prepare_data_for_generation( self, data: BatchedDataDict[GenerationDatumSpec], greedy: bool = False diff --git a/nemo_rl/models/generation/megatron/utils.py b/nemo_rl/models/generation/megatron/utils.py index f2a4b926adc..cd7d6886a8c 100644 --- a/nemo_rl/models/generation/megatron/utils.py +++ b/nemo_rl/models/generation/megatron/utils.py @@ -21,7 +21,9 @@ def build_image_preprocessing_config( - image_processor: Any, *, dynamic_resolution: bool + image_processor: Any, + *, + dynamic_resolution: bool, ) -> ImageProcessingConfig: """Translate an HF image processor to an MCore config.""" @@ -82,6 +84,7 @@ def build_video_preprocessing_config( generation_config: dict[str, Any], *, frame_manifest_magic: bytes, + video_maintain_aspect_ratio: bool = True, ) -> VideoProcessingConfig | None: """Build video preprocessing when explicitly enabled by generation config.""" temporal_patch_size = generation_config.get("video_temporal_patch_size") @@ -99,6 +102,7 @@ def build_video_preprocessing_config( image_config=image_config, num_frames=int(generation_config["video_num_frames"]), temporal_patch_size=int(temporal_patch_size), + video_maintain_aspect_ratio=video_maintain_aspect_ratio, frame_manifest_magic=frame_manifest_magic, ) From 907dfe0c15b1f5d551b5fdc5bfe5cdf66246b769 Mon Sep 17 00:00:00 2001 From: Cory Ye Date: Tue, 25 Aug 2026 14:15:43 -0700 Subject: [PATCH 06/15] Preserve bounded image normalization after rebase. Keep main's deduplicated, atomic media encoding behavior in the multimodal module where the branch moved Gym normalization. Signed-off-by: Cory Ye --- nemo_rl/environments/nemo_gym_multimodal.py | 50 +++++++++++++------ .../data/test_multimodal_image_encoding.py | 38 +++++++------- 2 files changed, 57 insertions(+), 31 deletions(-) diff --git a/nemo_rl/environments/nemo_gym_multimodal.py b/nemo_rl/environments/nemo_gym_multimodal.py index f81b4021297..e88f9fa44ce 100644 --- a/nemo_rl/environments/nemo_gym_multimodal.py +++ b/nemo_rl/environments/nemo_gym_multimodal.py @@ -69,7 +69,10 @@ def _encode_single_image_source(source: str) -> str: def normalize_media_in_examples(nemo_gym_examples: list[dict]) -> list[dict]: """Replace local media paths in NeMo Gym examples with data URLs.""" - image_targets_by_source: dict[str, list[dict]] = {} + local_image_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): @@ -116,19 +119,24 @@ def normalize_media_in_examples(nemo_gym_examples: list[dict]) -> list[dict]: raise ValueError(f"{part_type} requires a non-empty media URL") if not url.startswith(("http://", "https://", "data:")): if is_image: - image_targets_by_source.setdefault(url, []).append(part) + local_image_sources.setdefault(url, None) else: url = video_path_to_data_url(url) - for key in source_keys: - if key != canonical_key: - part.pop(key, None) - part["type"] = canonical_type - part[canonical_key] = url - if is_image and nested_detail is not None: - part.setdefault("detail", nested_detail) + pending_mutations.append( + ( + part, + source_keys, + canonical_type, + canonical_key, + is_image, + nested_detail, + url, + ) + ) - sources = list(image_targets_by_source) + sources = list(local_image_sources) + encoded_by_source: dict[str, str] = {} if sources: with ThreadPoolExecutor( max_workers=_NEMO_GYM_IMAGE_ENCODE_MAX_WORKERS @@ -140,10 +148,24 @@ def normalize_media_in_examples(nemo_gym_examples: list[dict]) -> list[dict]: strict=True, ) ) - for source, targets in image_targets_by_source.items(): - data_url = encoded_by_source[source] - for part in targets: - part["image_url"] = data_url + + # 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 + part[canonical_key] = encoded_by_source.get(url, url) + if is_image and nested_detail is not None: + part.setdefault("detail", nested_detail) return nemo_gym_examples diff --git a/tests/unit/data/test_multimodal_image_encoding.py b/tests/unit/data/test_multimodal_image_encoding.py index 9706a13dd82..2cc2cff7109 100644 --- a/tests/unit/data/test_multimodal_image_encoding.py +++ b/tests/unit/data/test_multimodal_image_encoding.py @@ -17,7 +17,7 @@ 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 ( image_to_data_url, resolve_to_image, @@ -106,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 = [ @@ -119,31 +119,35 @@ 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[1]["image_url"] == expected[first] + assert parts[2]["image_url"] == expected[second] def test_encode_images_does_not_partially_mutate_on_error(tmp_path): @@ -157,7 +161,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 @@ -166,15 +170,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 @@ -182,15 +186,15 @@ 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 From 122aa6b7cce7996d67fd9da470296ad3ced0948a Mon Sep 17 00:00:00 2001 From: Cory Ye Date: Tue, 25 Aug 2026 14:19:10 -0700 Subject: [PATCH 07/15] Align multimodal engine state with main lifecycle. Use the full media model for the inference context while retaining the nested language model for hybrid state and lifecycle operations. Signed-off-by: Cory Ye --- .../generation/megatron/megatron_worker.py | 23 ++++++++++--------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/nemo_rl/models/generation/megatron/megatron_worker.py b/nemo_rl/models/generation/megatron/megatron_worker.py index 807114da359..cddc603b32d 100644 --- a/nemo_rl/models/generation/megatron/megatron_worker.py +++ b/nemo_rl/models/generation/megatron/megatron_worker.py @@ -286,6 +286,7 @@ def _initialize_inference_engine(self, mcore_generation_config: dict) -> None: 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 @@ -385,7 +386,7 @@ def _initialize_inference_engine(self, mcore_generation_config: dict) -> None: logging_step_interval=logging_step_interval, num_speculative_tokens=num_speculative_tokens, logprobs_mode=mcore_generation_config.get( - "logprobs_mode", "raw_logprobs" + "logprobs_mode", "processed_logprobs" ), max_requests=max_requests, image_preprocessing_config=image_preprocessing_config, @@ -398,11 +399,11 @@ def _initialize_inference_engine(self, mcore_generation_config: dict) -> None: ] self.inference_context = DynamicInferenceContext( - model_config, inference_config + engine_model.config, inference_config ) if media_model is None: self.inference_wrapped_model = GPTInferenceWrapper( - inference_model, self.inference_context + engine_model, self.inference_context ) else: if inference_wrapper_cls is None: @@ -410,7 +411,7 @@ def _initialize_inference_engine(self, mcore_generation_config: dict) -> None: "Multimodal inference requires megatron_inference_wrapper." ) self.inference_wrapped_model = inference_wrapper_cls( - media_model, self.inference_context + engine_model, self.inference_context ) text_generation_controller = TextGenerationController( inference_wrapped_model=self.inference_wrapped_model, @@ -600,7 +601,8 @@ 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, _ = self._inference_model_and_media_parts() + lang_module = unwrap_model(inference_model) if self.is_generation_colocated: if self._inference_engine_initialized and not self._inference_engine_asleep: @@ -657,10 +659,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 @@ -673,13 +671,16 @@ 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, _ = 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) lang_module.eval() rotary_module = getattr(lang_module, "rotary_pos_emb", None) From a81fd717212c3a895d07e3f79d20eaa6eed272a0 Mon Sep 17 00:00:00 2001 From: Cory Ye Date: Tue, 25 Aug 2026 15:24:03 -0700 Subject: [PATCH 08/15] More rebase fixes. Signed-off-by: Cory Ye --- ...8g-megatron-tp4ep4-async-gym-video.v1.yaml | 1 + ...-circle-click-2n8g-megatron-tp2ep8.v1.yaml | 1 + ...cle-count-1n4g-megatron_generation.v1.yaml | 1 + ...a3b-clevr-1n4g-megatron_generation.v1.yaml | 1 + ...a3b-clevr-8n4g-megatron_generation.v1.yaml | 1 + nemo_rl/data/multimodal_utils.py | 2 ++ .../generation/megatron/megatron_worker.py | 19 ++++++++++++++----- 7 files changed, 21 insertions(+), 5 deletions(-) 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 f0db0455927..ad2e62218da 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 @@ -50,6 +50,7 @@ policy: 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 multimodal_prompt_config: image_spec: &media_prompt_spec 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 85a9b1cc61b..bfdcef5403f 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 @@ -91,6 +91,7 @@ policy: 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 multimodal_prompt_config: image_spec: &media_prompt_spec 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 index 53aa16dfdb3..d3ef612c7bc 100644 --- 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 @@ -93,6 +93,7 @@ policy: 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 multimodal_prompt_config: image_spec: &media_prompt_spec diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.yaml index 571e6a9a654..aaa06035b5f 100644 --- a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.yaml +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.yaml @@ -83,6 +83,7 @@ policy: 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 multimodal_prompt_config: image_spec: &media_prompt_spec 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 index d12d447e87f..da104b0bc29 100644 --- 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 @@ -93,6 +93,7 @@ policy: 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 multimodal_prompt_config: image_spec: &media_prompt_spec diff --git a/nemo_rl/data/multimodal_utils.py b/nemo_rl/data/multimodal_utils.py index c5c6ab95604..abbd7e59db0 100644 --- a/nemo_rl/data/multimodal_utils.py +++ b/nemo_rl/data/multimodal_utils.py @@ -50,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"], diff --git a/nemo_rl/models/generation/megatron/megatron_worker.py b/nemo_rl/models/generation/megatron/megatron_worker.py index cddc603b32d..58dbc9b79e2 100644 --- a/nemo_rl/models/generation/megatron/megatron_worker.py +++ b/nemo_rl/models/generation/megatron/megatron_worker.py @@ -394,7 +394,7 @@ def _initialize_inference_engine(self, mcore_generation_config: dict) -> None: ) if "inference_cuda_graph_scope" in mcore_generation_config: - model_config.inference_cuda_graph_scope = InferenceCudaGraphScope[ + engine_model.config.inference_cuda_graph_scope = InferenceCudaGraphScope[ mcore_generation_config["inference_cuda_graph_scope"] ] @@ -601,8 +601,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") - inference_model, _ = self._inference_model_and_media_parts() + 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: @@ -611,7 +614,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) @@ -676,11 +680,14 @@ def prepare_for_generation(self, tags=None, **kwargs) -> None: if self.inference_model is not None: self._reshard_into_inference_model() - inference_model, _ = self._inference_model_and_media_parts() + 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) @@ -691,7 +698,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 From abef468b7aeb3fd6f8c68f59d6ad9e5580ad8f38 Mon Sep 17 00:00:00 2001 From: Cory Ye Date: Tue, 25 Aug 2026 17:01:53 -0700 Subject: [PATCH 09/15] Retreat the model prompt config back to the wrapper. Users can just create a new wrapper for a new config. Signed-off-by: Cory Ye --- ...3b-16n8g-megatron-tp4ep4-async-gym-video.v1.yaml | 6 ------ ...30ba3b-circle-click-2n8g-megatron-tp2ep8.v1.yaml | 6 ------ ...3b-circle-count-1n4g-megatron_generation.v1.yaml | 6 ------ ...ni-30ba3b-clevr-1n4g-megatron_generation.v1.yaml | 6 ------ ...ni-30ba3b-clevr-8n4g-megatron_generation.v1.yaml | 6 ------ nemo_rl/models/generation/megatron/config.py | 2 +- .../models/generation/megatron/megatron_worker.py | 13 ++++--------- nemo_rl/models/generation/vllm/config.py | 2 +- tests/unit/environments/test_nemo_gym.py | 2 +- .../generation/test_megatron_generation_parse.py | 1 + 10 files changed, 8 insertions(+), 42 deletions(-) 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 ad2e62218da..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 @@ -52,12 +52,6 @@ policy: image_dynamic_resolution: true logprobs_mode: raw_logprobs megatron_inference_wrapper: megatron.core.inference.model_inference_wrappers.multimodal.nemotron_omni_inference_wrapper.NemotronOmniInferenceWrapper - multimodal_prompt_config: - image_spec: &media_prompt_spec - model_token: - prefix: - suffix: - video_spec: *media_prompt_spec 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 bfdcef5403f..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 @@ -93,12 +93,6 @@ policy: image_dynamic_resolution: true logprobs_mode: raw_logprobs megatron_inference_wrapper: megatron.core.inference.model_inference_wrappers.multimodal.nemotron_omni_inference_wrapper.NemotronOmniInferenceWrapper - multimodal_prompt_config: - image_spec: &media_prompt_spec - model_token: - prefix: - suffix: - video_spec: *media_prompt_spec 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 index d3ef612c7bc..25aa3cc8e09 100644 --- 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 @@ -95,12 +95,6 @@ policy: image_dynamic_resolution: true logprobs_mode: raw_logprobs megatron_inference_wrapper: megatron.core.inference.model_inference_wrappers.multimodal.nemotron_omni_inference_wrapper.NemotronOmniInferenceWrapper - multimodal_prompt_config: - image_spec: &media_prompt_spec - model_token: - prefix: - suffix: - video_spec: *media_prompt_spec vllm_cfg: async_engine: true expose_http_server: true diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.yaml index aaa06035b5f..6db937e21ef 100644 --- a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.yaml +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.yaml @@ -85,12 +85,6 @@ policy: image_dynamic_resolution: true logprobs_mode: raw_logprobs megatron_inference_wrapper: megatron.core.inference.model_inference_wrappers.multimodal.nemotron_omni_inference_wrapper.NemotronOmniInferenceWrapper - multimodal_prompt_config: - image_spec: &media_prompt_spec - model_token: - prefix: - suffix: - video_spec: *media_prompt_spec colocated: enabled: false resources: 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 index da104b0bc29..0948bdc6d7c 100644 --- 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 @@ -95,12 +95,6 @@ policy: image_dynamic_resolution: true logprobs_mode: raw_logprobs megatron_inference_wrapper: megatron.core.inference.model_inference_wrappers.multimodal.nemotron_omni_inference_wrapper.NemotronOmniInferenceWrapper - multimodal_prompt_config: - image_spec: &media_prompt_spec - model_token: - prefix: - suffix: - video_spec: *media_prompt_spec colocated: enabled: false resources: diff --git a/nemo_rl/models/generation/megatron/config.py b/nemo_rl/models/generation/megatron/config.py index 0927024ae03..6cc9352df1d 100644 --- a/nemo_rl/models/generation/megatron/config.py +++ b/nemo_rl/models/generation/megatron/config.py @@ -52,6 +52,7 @@ class MCoreGenerationSpecificArgs(TypedDict): enable_prefix_caching: bool async_sched_mode: NotRequired[Literal["legacy", "async"]] vision_embedding_cache_max_bytes: NotRequired[int] + allow_stale_multimodal_embeddings: NotRequired[bool] logprobs_mode: NotRequired[Literal["raw_logprobs", "processed_logprobs"]] refit_backend: Literal["gloo", "nccl", "nvshmem"] @@ -66,7 +67,6 @@ class MCoreGenerationSpecificArgs(TypedDict): video_temporal_patch_size: NotRequired[int] video_target_num_patches: NotRequired[int] video_maintain_aspect_ratio: NotRequired[bool] - multimodal_prompt_config: NotRequired[dict] megatron_inference_wrapper: NotRequired[str] # KV cache lifecycle across suspend/resume: diff --git a/nemo_rl/models/generation/megatron/megatron_worker.py b/nemo_rl/models/generation/megatron/megatron_worker.py index 58dbc9b79e2..459019cf76e 100644 --- a/nemo_rl/models/generation/megatron/megatron_worker.py +++ b/nemo_rl/models/generation/megatron/megatron_worker.py @@ -374,6 +374,9 @@ def _initialize_inference_engine(self, mcore_generation_config: dict) -> None: vision_embedding_cache_max_bytes=int( mcore_generation_config.get("vision_embedding_cache_max_bytes", 0) ), + allow_stale_multimodal_embeddings=bool( + mcore_generation_config.get("allow_stale_multimodal_embeddings", False) + ), prefix_caching_coordinator_policy=PrefixCachingCoordinatorPolicy( "first_prefix_block" ), @@ -505,7 +508,6 @@ def run_loop(): def _setup_openai_api_server(self) -> str: """Start the OpenAI-compatible HTTP server on this worker.""" - from megatron.core.inference.config import MultimodalPromptConfig from megatron.core.inference.text_generation_server.dynamic_text_gen_server.text_generation_server import ( start_text_gen_server, ) @@ -522,9 +524,6 @@ def _setup_openai_api_server(self) -> str: else: server_port = _get_free_port_local() - prompt_config = self.cfg["generation"]["mcore_generation_config"].get( - "multimodal_prompt_config" - ) start_text_gen_server( coordinator_addr=self.coordinator_addr, tokenizer=self.megatron_tokenizer, @@ -533,11 +532,7 @@ def _setup_openai_api_server(self) -> str: parsers=self.cfg["generation"]["mcore_generation_config"]["parsers"], verbose=False, sock=reserved_socket, - multimodal_prompt_config=( - MultimodalPromptConfig.from_dict(prompt_config) - if prompt_config - else None - ), + multimodal_prompt_config=self.inference_wrapped_model.multimodal_prompt_config, ) base_url = f"http://{ip}:{server_port}/v1" 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/tests/unit/environments/test_nemo_gym.py b/tests/unit/environments/test_nemo_gym.py index 32736d803c2..43485c809b2 100644 --- a/tests/unit/environments/test_nemo_gym.py +++ b/tests/unit/environments/test_nemo_gym.py @@ -1478,7 +1478,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(): diff --git a/tests/unit/models/generation/test_megatron_generation_parse.py b/tests/unit/models/generation/test_megatron_generation_parse.py index 6db700ef237..2a9a8655c98 100644 --- a/tests/unit/models/generation/test_megatron_generation_parse.py +++ b/tests/unit/models/generation/test_megatron_generation_parse.py @@ -231,6 +231,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 From 3024ea037bf17c49081ffba2d3b1d5f85f3b8839 Mon Sep 17 00:00:00 2001 From: Cory Ye Date: Wed, 26 Aug 2026 14:03:20 -0700 Subject: [PATCH 10/15] Add tests and lint. Signed-off-by: Cory Ye --- ...cle-count-1n4g-megatron_generation.v1.yaml | 64 +------- ...a3b-clevr-1n4g-megatron_generation.v1.yaml | 54 +----- ...a3b-clevr-8n4g-megatron_generation.v1.yaml | 58 +------ examples/nemo_gym/prepare_video_dataset.py | 2 +- nemo_rl/algorithms/grpo.py | 48 +----- nemo_rl/data/multimodal_utils.py | 1 + nemo_rl/data/processors.py | 2 + nemo_rl/distributed/worker_groups.py | 18 +- nemo_rl/environments/nemo_gym_multimodal.py | 17 +- .../generation/megatron/megatron_worker.py | 4 +- nemo_rl/models/generation/megatron/utils.py | 30 ++-- nemo_rl/models/generation/vllm/video_utils.py | 2 + nemo_rl/utils/venvs.py | 41 ++++- pyrefly.toml | 3 +- .../L1_Functional_Tests_Megatron_Omni.sh | 38 +++++ .../nemotron_omni_clevr_megatron_1n4g.sh | 145 +++++++++++++++++ .../nemotron_omni_gym_video_megatron_1n4g.sh | 154 ++++++++++++++++++ tests/unit/algorithms/test_grpo.py | 11 +- tests/unit/data/test_collate_fn.py | 6 +- .../data/test_multimodal_image_encoding.py | 8 +- tests/unit/distributed/test_worker_groups.py | 27 ++- tests/unit/environments/test_nemo_gym.py | 107 ++++++++++++ .../generation/test_megatron_generation.py | 101 ++++++++++++ tests/unit/test_prepare_video_dataset.py | 2 +- tests/unit/utils/test_venvs.py | 30 +++- 25 files changed, 725 insertions(+), 248 deletions(-) create mode 100755 tests/functional/L1_Functional_Tests_Megatron_Omni.sh create mode 100755 tests/functional/nemotron_omni_clevr_megatron_1n4g.sh create mode 100755 tests/functional/nemotron_omni_gym_video_megatron_1n4g.sh 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 index 25aa3cc8e09..ebc3c231a19 100644 --- 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 @@ -1,90 +1,45 @@ defaults: ../../vlm_grpo_3B_megatron.yaml - -# 1-node × 4-GPU GB200 recipe: Nemotron Omni on NeMo-Gym's circle_count image -# environment. Each row is a synthetic PNG scene carried inline as a base64 -# data URL plus a "how many circles" question; the Gym resources server -# rewards an exact match on \boxed{count}. Default layout is non-colocated -# 2 train / 2 gen with Megatron generation, async GRPO, age>1, and in-flight -# weight updates. -# -# `data` and `env` use _override_ because the parent config wires the CLEVR -# HF dataset and its reward functions at the split level, and split keys win -# over `data.default` (nemo_rl/data/datasets/utils.py:263). grpo: num_prompts_per_step: 2 num_generations_per_prompt: 8 max_num_steps: 4 val_period: 0 - val_at_start: false - val_at_end: false - use_dynamic_sampling: false - # The Gym entrypoint refuses a preset max_val_samples and derives both of - # these from the validation dataset length instead. max_val_samples: null val_batch_size: null - reward_shaping: - enabled: false - reward_scaling: - enabled: false 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 - train_micro_batch_size: 1 logprob_batch_size: 1 max_total_sequence_length: 8192 - offload_optimizer_for_logprob: false sequence_packing: enabled: true megatron_cfg: env_vars: - TORCH_CUDA_ARCH_LIST: "10.0" + TORCH_CUDA_ARCH_LIST: '10.0' tensor_model_parallel_size: 2 expert_model_parallel_size: 2 - expert_tensor_parallel_size: 1 - context_parallel_size: 1 sequence_parallel: true bias_activation_fusion: false activation_checkpointing: true generation: backend: megatron - max_new_tokens: 1024 - # Gym clears stop_strings/stop_token_ids in setup_nemo_gym_config, and the - # Megatron engine takes null here (the vLLM path bans the media tokens - # explicitly instead). bad_words: null mcore_generation_config: - # Required for Gym: the environment drives generation over the engine's - # OpenAI-compatible HTTP server. expose_http_server: true - parsers: [] - max_model_len: ${policy.max_total_sequence_length} - cuda_graph_impl: local - inference_cuda_graph_scope: block buffer_size_gb: 8 num_cuda_graphs: -1 - block_size_tokens: 256 - use_cuda_graphs_for_non_decode_steps: true - enable_chunked_prefill: true - enable_prefix_caching: false max_tokens: ${policy.max_total_sequence_length} - kv_cache_management_mode: persist - materialize_only_last_token_logits: true - num_speculative_tokens: 0 - refit_backend: nccl transformer_impl: transformer_engine activation_checkpointing: false tensor_model_parallel_size: 2 @@ -98,22 +53,15 @@ policy: vllm_cfg: async_engine: true expose_http_server: true - skip_tokenizer_init: false colocated: enabled: false resources: gpus_per_node: 2 num_nodes: 1 - data: _override_: true max_input_seq_length: null shuffle: false - num_workers: 1 - use_multiple_dataloader: false - # The launcher overrides these with absolute paths; the shipped example file - # holds 5 rows, so regenerate a larger set with - # resources_servers/circle_count/generate_data.py for anything but a smoke. train: data_path: 3rdparty/Gym-workspace/Gym/resources_servers/circle_count/data/example.jsonl validation: @@ -122,9 +70,7 @@ data: dataset_name: NemoGymDataset env_name: nemo_gym prompt_file: null - system_prompt_file: null processor: nemo_gym_data_processor - env: _override_: true should_use_nemo_gym: true @@ -140,14 +86,10 @@ env: responses_api_agents: simple_agent: max_steps: 1 - logger: - wandb_enabled: false tensorboard_enabled: false wandb: - project: "nemo-rl-omni" - name: "nemotron-omni-circle-count-1n4g" - + project: nemo-rl-omni + name: nemotron-omni-circle-count-1n4g cluster: gpus_per_node: 4 - num_nodes: 1 diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.yaml index 6db937e21ef..0abac605b1c 100644 --- a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.yaml +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.yaml @@ -1,50 +1,31 @@ defaults: ../../vlm_grpo_3B_megatron.yaml - -# 1-node × 4-GPU GB200 recipe: Nemotron Omni CLEVR with Megatron generation. -# Default layout is non-colocated 2 train / 2 gen with async GRPO, age>1, and -# in-flight weight updates. grpo: num_prompts_per_step: 2 num_generations_per_prompt: 8 max_num_steps: 4 val_period: 0 - val_at_start: false - val_at_end: false - use_dynamic_sampling: false - reward_shaping: - enabled: false - reward_scaling: - enabled: false 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-1n4g - policy: model_name: nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16 is_vlm: true train_global_batch_size: 16 - train_micro_batch_size: 1 logprob_batch_size: 1 - max_total_sequence_length: 2048 - offload_optimizer_for_logprob: false sequence_packing: enabled: true megatron_cfg: env_vars: - TORCH_CUDA_ARCH_LIST: "10.0" + TORCH_CUDA_ARCH_LIST: '10.0' tensor_model_parallel_size: 2 expert_model_parallel_size: 2 - expert_tensor_parallel_size: 1 - context_parallel_size: 1 sequence_parallel: true bias_activation_fusion: false activation_checkpointing: true @@ -52,28 +33,16 @@ policy: backend: megatron max_new_tokens: 512 stop_strings: - - - - - - - - - - - - + - + - + - + - + - + - mcore_generation_config: - expose_http_server: false - parsers: [] - max_model_len: ${policy.max_total_sequence_length} - cuda_graph_impl: local - inference_cuda_graph_scope: block buffer_size_gb: 8 num_cuda_graphs: -1 - block_size_tokens: 256 - use_cuda_graphs_for_non_decode_steps: true - enable_chunked_prefill: true - enable_prefix_caching: false max_tokens: ${policy.max_total_sequence_length} - kv_cache_management_mode: persist - materialize_only_last_token_logits: true - num_speculative_tokens: 0 refit_backend: gloo transformer_impl: transformer_engine activation_checkpointing: false @@ -90,18 +59,13 @@ policy: resources: gpus_per_node: 2 num_nodes: 1 - data: default: prompt_file: examples/prompts/clevr_cogent_cot_nemotron_omni.txt - logger: - wandb_enabled: false tensorboard_enabled: false wandb: - project: "nemo-rl-omni" - name: "nemotron-omni-clevr-megatron-1n4g" - + project: nemo-rl-omni + name: nemotron-omni-clevr-megatron-1n4g cluster: gpus_per_node: 4 - num_nodes: 1 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 index 0948bdc6d7c..ba46f1079ef 100644 --- 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 @@ -1,60 +1,36 @@ defaults: ../../vlm_grpo_3B_megatron.yaml - -# 8-node × 4-GPU GB200 recipe: Nemotron Omni CLEVR with Megatron generation. -# Default layout is non-colocated 2 train / 6 gen nodes with async GRPO, age>1, -# and in-flight weight updates (gen-heavy split for long rollouts). grpo: - # Inference DP=3 (24 GPUs / TP8); two prompt groups per DP replica → 6. num_prompts_per_step: 6 num_generations_per_prompt: 8 - max_num_steps: 1000000 - val_period: 10 - val_at_start: false - val_at_end: false - # Training DP=1 (8 GPUs / TP8). val_batch_size: 64 max_val_samples: 64 - use_dynamic_sampling: false - reward_shaping: - enabled: false - reward_scaling: - enabled: false 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}} - train_micro_batch_size: 1 logprob_batch_size: 1 max_total_sequence_length: 4096 - offload_optimizer_for_logprob: false sequence_packing: enabled: true megatron_cfg: env_vars: - TORCH_CUDA_ARCH_LIST: "10.0" + TORCH_CUDA_ARCH_LIST: '10.0' tensor_model_parallel_size: 8 expert_model_parallel_size: 8 - expert_tensor_parallel_size: 1 - context_parallel_size: 1 sequence_parallel: true bias_activation_fusion: false activation_checkpointing: true optimizer: - optimizer_cpu_offload: false - optimizer_offload_fraction: 0.0 exp_avg_dtype: bfloat16 exp_avg_sq_dtype: bfloat16 store_param_remainders: true @@ -62,28 +38,16 @@ policy: backend: megatron max_new_tokens: 2048 stop_strings: - - - - - - - - - - - - + - + - + - + - + - + - mcore_generation_config: - expose_http_server: false - parsers: [] - max_model_len: ${policy.max_total_sequence_length} - cuda_graph_impl: local - inference_cuda_graph_scope: block buffer_size_gb: 8 num_cuda_graphs: -1 - block_size_tokens: 256 - use_cuda_graphs_for_non_decode_steps: true - enable_chunked_prefill: true - enable_prefix_caching: false max_tokens: ${policy.max_total_sequence_length} - kv_cache_management_mode: persist - materialize_only_last_token_logits: true - num_speculative_tokens: 0 refit_backend: gloo transformer_impl: transformer_engine activation_checkpointing: false @@ -100,18 +64,14 @@ policy: resources: gpus_per_node: 4 num_nodes: 6 - data: default: prompt_file: examples/prompts/clevr_cogent_cot_nemotron_omni.txt - logger: - wandb_enabled: false tensorboard_enabled: false wandb: - project: "nemo-rl-omni" - name: "nemotron-omni-clevr-megatron-8n4g" - + 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/multimodal_utils.py b/nemo_rl/data/multimodal_utils.py index abbd7e59db0..8aaa99915d9 100644 --- a/nemo_rl/data/multimodal_utils.py +++ b/nemo_rl/data/multimodal_utils.py @@ -1064,6 +1064,7 @@ def attach_image_model_inputs_to_message( ".avi": "x-msvideo", } + 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:"): diff --git a/nemo_rl/data/processors.py b/nemo_rl/data/processors.py index e4ea9f198d0..084fed12f6d 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']}") 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_multimodal.py b/nemo_rl/environments/nemo_gym_multimodal.py index e88f9fa44ce..fd3baa2e33d 100644 --- a/nemo_rl/environments/nemo_gym_multimodal.py +++ b/nemo_rl/environments/nemo_gym_multimodal.py @@ -30,8 +30,8 @@ IMAGE_CONTENT_TYPES, VIDEO_CONTENT_TYPES, PackedTensor, - extract_multimodal_model_inputs, 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, @@ -60,11 +60,8 @@ def _encode_single_image_source(source: str) -> str: """Resolve, encode, and close one local image source.""" - image = resolve_to_image(source) - try: + with resolve_to_image(source) as image: return image_to_data_url(image) - finally: - image.close() def normalize_media_in_examples(nemo_gym_examples: list[dict]) -> list[dict]: @@ -101,7 +98,11 @@ def normalize_media_in_examples(nemo_gym_examples: list[dict]) -> list[dict]: 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: + if ( + not present_keys + and part_type == "input_image" + and "file_id" in part + ): continue if len(present_keys) != 1: raise ValueError( @@ -109,7 +110,9 @@ def normalize_media_in_examples(nemo_gym_examples: list[dict]) -> list[dict]: ) source = part[present_keys[0]] - nested_detail = source.get("detail") if isinstance(source, dict) else None + nested_detail = ( + source.get("detail") if isinstance(source, dict) else None + ) url = ( source.get("url") or source.get("path", "") if isinstance(source, dict) diff --git a/nemo_rl/models/generation/megatron/megatron_worker.py b/nemo_rl/models/generation/megatron/megatron_worker.py index 459019cf76e..c150dc3b85a 100644 --- a/nemo_rl/models/generation/megatron/megatron_worker.py +++ b/nemo_rl/models/generation/megatron/megatron_worker.py @@ -789,9 +789,7 @@ def _sample_vision_tensors(self, data, index: int): 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 + 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 diff --git a/nemo_rl/models/generation/megatron/utils.py b/nemo_rl/models/generation/megatron/utils.py index cd7d6886a8c..922654dbf08 100644 --- a/nemo_rl/models/generation/megatron/utils.py +++ b/nemo_rl/models/generation/megatron/utils.py @@ -42,18 +42,24 @@ def read(*names: str) -> Any: pixel_mean = read("norm_mean", "image_mean") pixel_std = read("norm_std", "image_std") - 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 - ] - if missing: + 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 " diff --git a/nemo_rl/models/generation/vllm/video_utils.py b/nemo_rl/models/generation/vllm/video_utils.py index 90b5637585f..88a3059926e 100644 --- a/nemo_rl/models/generation/vllm/video_utils.py +++ b/nemo_rl/models/generation/vllm/video_utils.py @@ -34,6 +34,8 @@ _TORCHCODEC_END_OF_STREAM_ERROR = ( "Requested next frame while there are no more frames left to decode." ) + + def _round_video_frame_count( num_frames: int, *, 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/functional/L1_Functional_Tests_Megatron_Omni.sh b/tests/functional/L1_Functional_Tests_Megatron_Omni.sh new file mode 100755 index 00000000000..1d95ae917ed --- /dev/null +++ b/tests/functional/L1_Functional_Tests_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 < 4 )); then + echo "SKIP: Nemotron Omni functional tests require at least four GPUs" + exit 0 +fi + +# The recipes are intentionally 1n4g even when the CI runner exposes eight GPUs. +export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3}" + +time uv run --no-sync bash ./tests/functional/nemotron_omni_clevr_megatron_1n4g.sh +time uv run --no-sync bash ./tests/functional/nemotron_omni_gym_video_megatron_1n4g.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_1n4g.sh b/tests/functional/nemotron_omni_clevr_megatron_1n4g.sh new file mode 100755 index 00000000000..49a02dc97e1 --- /dev/null +++ b/tests/functional/nemotron_omni_clevr_megatron_1n4g.sh @@ -0,0 +1,145 @@ +#!/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 < 4 )); then + echo "SKIP: Omni CLEVR Megatron smoke requires at least four 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_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. +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 + +uv run --no-sync python examples/run_vlm_grpo.py \ + --config examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.yaml \ + cluster.num_nodes=1 \ + cluster.gpus_per_node=4 \ + policy.megatron_cfg.env_vars.TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST}" \ + policy.megatron_cfg.tensor_model_parallel_size=4 \ + policy.megatron_cfg.expert_model_parallel_size=4 \ + 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.optimizer.optimizer_cpu_offload=false \ + policy.megatron_cfg.optimizer.optimizer_offload_fraction=0.0 \ + ++policy.megatron_cfg.optimizer.exp_avg_dtype=bfloat16 \ + ++policy.megatron_cfg.optimizer.exp_avg_sq_dtype=bfloat16 \ + ++policy.megatron_cfg.optimizer.store_param_remainders=true \ + policy.generation.backend=megatron \ + policy.generation.colocated.enabled=true \ + policy.generation.colocated.resources.num_nodes=1 \ + policy.generation.colocated.resources.gpus_per_node=4 \ + policy.generation.max_new_tokens=128 \ + policy.generation.mcore_generation_config.tensor_model_parallel_size=4 \ + policy.generation.mcore_generation_config.expert_model_parallel_size=4 \ + 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.cuda_graph_impl="${MEGATRON_CUDA_GRAPH_IMPL}" \ + policy.generation.mcore_generation_config.inference_cuda_graph_scope=block \ + policy.generation.mcore_generation_config.num_cuda_graphs=-1 \ + 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 \ + 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/loss"]) < 1e6' \ + 'min(data["train/loss"]) > -1e6' diff --git a/tests/functional/nemotron_omni_gym_video_megatron_1n4g.sh b/tests/functional/nemotron_omni_gym_video_megatron_1n4g.sh new file mode 100755 index 00000000000..798d3b735fa --- /dev/null +++ b/tests/functional/nemotron_omni_gym_video_megatron_1n4g.sh @@ -0,0 +1,154 @@ +#!/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 < 4 )); then + echo "SKIP: Omni Gym-video Megatron smoke requires at least four 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}" +MOE_PAD_EXPERTS_FOR_CG=false + +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}" + +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=4 \ + 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.optimizer.optimizer_cpu_offload=false \ + policy.megatron_cfg.optimizer.optimizer_offload_fraction=0.0 \ + ++policy.megatron_cfg.optimizer.exp_avg_dtype=bfloat16 \ + ++policy.megatron_cfg.optimizer.exp_avg_sq_dtype=bfloat16 \ + ++policy.megatron_cfg.optimizer.store_param_remainders=true \ + policy.generation.backend=megatron \ + ++policy.generation.bad_words=null \ + policy.generation.colocated.enabled=false \ + 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=8 \ + policy.generation.mcore_generation_config.cuda_graph_impl=none \ + policy.generation.mcore_generation_config.inference_cuda_graph_scope=none \ + policy.generation.mcore_generation_config.num_cuda_graphs=0 \ + 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=4096 \ + policy.generation.mcore_generation_config.max_tokens=4096 \ + ++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=4096 \ + +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}" + +RECORDED_STEP=$(jq -r \ + 'if has("train/loss") then (."train/loss" | keys | map(tonumber) | max // 0) else 0 end' \ + "${JSON_METRICS}") +if (( RECORDED_STEP < 1 )); then + echo "[ERROR] Expected at least one completed Gym-video training step" + exit 1 +fi + +uv run --no-sync tests/check_metrics.py "${JSON_METRICS}" \ + 'max(data["train/loss"]) < 1e6' \ + 'min(data["train/loss"]) > -1e6' 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/test_collate_fn.py b/tests/unit/data/test_collate_fn.py index 7cf80429350..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 eval_collate_fn, preference_collate_fn, rl_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 diff --git a/tests/unit/data/test_multimodal_image_encoding.py b/tests/unit/data/test_multimodal_image_encoding.py index 2cc2cff7109..d020fa9e2e0 100644 --- a/tests/unit/data/test_multimodal_image_encoding.py +++ b/tests/unit/data/test_multimodal_image_encoding.py @@ -133,9 +133,7 @@ def tracking_executor(*args, **kwargs): return original_executor(*args, **kwargs) monkeypatch.setattr(nemo_gym_multimodal, "resolve_to_image", tracking_resolve) - monkeypatch.setattr( - nemo_gym_multimodal, "ThreadPoolExecutor", tracking_executor - ) + monkeypatch.setattr(nemo_gym_multimodal, "ThreadPoolExecutor", tracking_executor) normalize_media_in_examples(examples) @@ -208,5 +206,7 @@ def test_normalize_media_is_a_noop_for_text_only_examples(): assert examples[0]["responses_create_params"]["input"][0]["content"] == before # Missing/oddly-shaped payloads must not raise. - assert normalize_media_in_examples([{}, {"responses_create_params": {}}]) is not None + 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 43485c809b2..b8a401d6cb4 100644 --- a/tests/unit/environments/test_nemo_gym.py +++ b/tests/unit/environments/test_nemo_gym.py @@ -1528,6 +1528,113 @@ 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 + + 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/models/generation/test_megatron_generation.py b/tests/unit/models/generation/test_megatron_generation.py index c45a959a10f..52ff3177a4d 100644 --- a/tests/unit/models/generation/test_megatron_generation.py +++ b/tests/unit/models/generation/test_megatron_generation.py @@ -22,6 +22,7 @@ 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 @@ -99,6 +100,106 @@ def wrapper_supports_modality(wrapper_cls, modality): assert "video" in multi_modal_data +@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}, 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..5a5de9a6673 100644 --- a/tests/unit/utils/test_venvs.py +++ b/tests/unit/utils/test_venvs.py @@ -16,7 +16,10 @@ 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, +) from tests.unit.conftest import TEST_ASSETS_DIR @@ -48,3 +51,28 @@ 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 From 10b95211ffbc0d4b884adee1ec8dda161c3fd2da Mon Sep 17 00:00:00 2001 From: Cory Ye Date: Thu, 27 Aug 2026 23:13:54 -0700 Subject: [PATCH 11/15] Address Claude and Fable reviews Signed-off-by: Cory Ye --- ...cle-count-1n4g-megatron_generation.v1.yaml | 1 + ...a3b-clevr-1n4g-megatron_generation.v1.yaml | 1 + nemo_rl/data/multimodal_utils.py | 4 +- nemo_rl/environments/nemo_gym.py | 11 +- nemo_rl/environments/nemo_gym_multimodal.py | 33 ++++-- nemo_rl/evals/eval.py | 43 ++++---- nemo_rl/models/generation/megatron/config.py | 23 ++-- .../generation/megatron/megatron_worker.py | 100 +++++++++++++----- nemo_rl/models/generation/megatron/utils.py | 35 ++++-- tests/unit/environments/test_nemo_gym.py | 29 +++++ .../environments/test_nemo_gym_mm_utils.py | 70 +++++++++++- .../generation/test_megatron_generation.py | 3 + tests/unit/utils/test_venvs.py | 24 +++++ tools/install_audio_deps.sh | 27 +++-- 14 files changed, 321 insertions(+), 83 deletions(-) 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 index ebc3c231a19..aca1e6015a6 100644 --- 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 @@ -62,6 +62,7 @@ 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: diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.yaml index 0abac605b1c..773e057a1da 100644 --- a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.yaml +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.yaml @@ -60,6 +60,7 @@ policy: gpus_per_node: 2 num_nodes: 1 data: + num_workers: 0 default: prompt_file: examples/prompts/clevr_cogent_cot_nemotron_omni.txt logger: diff --git a/nemo_rl/data/multimodal_utils.py b/nemo_rl/data/multimodal_utils.py index 8aaa99915d9..67878e5278a 100644 --- a/nemo_rl/data/multimodal_utils.py +++ b/nemo_rl/data/multimodal_utils.py @@ -930,7 +930,9 @@ def extract_input_media_sources_from_responses_messages( continue if isinstance(source, dict): source = source.get("url") or source.get("path") - if source is not None: + # 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((media_type, source)) return sources diff --git a/nemo_rl/environments/nemo_gym.py b/nemo_rl/environments/nemo_gym.py index 80307cb56df..0f68069db92 100644 --- a/nemo_rl/environments/nemo_gym.py +++ b/nemo_rl/environments/nemo_gym.py @@ -41,6 +41,7 @@ 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, ) @@ -302,16 +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 _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")) - - @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.""" diff --git a/nemo_rl/environments/nemo_gym_multimodal.py b/nemo_rl/environments/nemo_gym_multimodal.py index fd3baa2e33d..3e65b88ba69 100644 --- a/nemo_rl/environments/nemo_gym_multimodal.py +++ b/nemo_rl/environments/nemo_gym_multimodal.py @@ -16,6 +16,7 @@ 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 @@ -60,13 +61,16 @@ def _encode_single_image_source(source: str) -> str: """Resolve, encode, and close one local image source.""" - with resolve_to_image(source) as image: + # `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] ] = [] @@ -124,7 +128,7 @@ def normalize_media_in_examples(nemo_gym_examples: list[dict]) -> list[dict]: if is_image: local_image_sources.setdefault(url, None) else: - url = video_path_to_data_url(url) + local_video_sources.setdefault(url, None) pending_mutations.append( ( @@ -152,6 +156,14 @@ def normalize_media_in_examples(nemo_gym_examples: list[dict]) -> list[dict]: ) ) + # 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, @@ -166,7 +178,8 @@ def normalize_media_in_examples(nemo_gym_examples: list[dict]) -> list[dict]: if key != canonical_key: part.pop(key, None) part["type"] = canonical_type - part[canonical_key] = encoded_by_source.get(url, url) + 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 @@ -220,6 +233,16 @@ def _extract_input_images_from_message(item: dict) -> list[Image.Image]: 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, @@ -250,9 +273,7 @@ def _index_per_turn_images( 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 + if _is_trainable_output_item(item): per_turn.append(pending) pending = [] elif item.get("role") != "assistant": diff --git a/nemo_rl/evals/eval.py b/nemo_rl/evals/eval.py index 17c335e900f..1ea9f766f4e 100644 --- a/nemo_rl/evals/eval.py +++ b/nemo_rl/evals/eval.py @@ -349,26 +349,31 @@ async def _run_env_eval_impl( videos[i][0] if len(videos[i]) == 1 else videos[i] ) - if multi_modal_data: - vllm_content = batch["vllm_content"][i] - if vllm_content is not None: - prompt_dict = {"prompt": vllm_content} - prompt_display = vllm_content - else: - # 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} - prompt_display = "\n".join( - str(message["content"]) for message in message_log + 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(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["multi_modal_data"] = multi_modal_data + 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: diff --git a/nemo_rl/models/generation/megatron/config.py b/nemo_rl/models/generation/megatron/config.py index 6cc9352df1d..7a7afc585ca 100644 --- a/nemo_rl/models/generation/megatron/config.py +++ b/nemo_rl/models/generation/megatron/config.py @@ -44,8 +44,6 @@ class MCoreGenerationSpecificArgs(TypedDict): # - 'block': graphs are owned at the enclosing block (TransformerBlock / HybridBlock). # Only meaningful when cuda_graph_impl='local'. inference_cuda_graph_scope: NotRequired[str] - # Required for EP>1 + local CUDA graphs. - moe_pad_experts_for_cuda_graph_inference: NotRequired[bool] materialize_only_last_token_logits: bool enable_chunked_prefill: bool @@ -53,7 +51,6 @@ class MCoreGenerationSpecificArgs(TypedDict): async_sched_mode: NotRequired[Literal["legacy", "async"]] vision_embedding_cache_max_bytes: NotRequired[int] allow_stale_multimodal_embeddings: NotRequired[bool] - logprobs_mode: NotRequired[Literal["raw_logprobs", "processed_logprobs"]] refit_backend: Literal["gloo", "nccl", "nvshmem"] num_speculative_tokens: int @@ -61,12 +58,24 @@ class MCoreGenerationSpecificArgs(TypedDict): mamba_inference_ssm_states_dtype: NotRequired[str] mamba_inference_conv_states_dtype: NotRequired[str] - # Raw media preprocessing used by the OpenAI-compatible HTTP endpoint. + # Raw media preprocessing. Each key is read only on the path that needs it; + # when omitted, MCore's own dataclass default applies (see + # ImageProcessingConfig / VideoProcessingConfig). + # `image_dynamic_resolution` requires an inference wrapper that supports + # images. The `video_*` keys are read only when `video_temporal_patch_size` + # is set, and `video_num_frames` is then required. image_dynamic_resolution: NotRequired[bool] - video_num_frames: NotRequired[int] - video_temporal_patch_size: NotRequired[int] - video_target_num_patches: NotRequired[int] + 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: diff --git a/nemo_rl/models/generation/megatron/megatron_worker.py b/nemo_rl/models/generation/megatron/megatron_worker.py index c150dc3b85a..0dc14e6d89d 100644 --- a/nemo_rl/models/generation/megatron/megatron_worker.py +++ b/nemo_rl/models/generation/megatron/megatron_worker.py @@ -117,22 +117,42 @@ def _init_inference_engine_state(self) -> None: self._inference_loop = None self._inference_thread = None - def _get_megatron_inference_wrapper_cls(self): - """Resolve the configured Megatron inference wrapper, if any.""" + 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." + "megatron_inference_wrapper must be a fully qualified class name, " + f"got {class_path!r}." ) - return getattr(importlib.import_module(module_name), class_name) + 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, modality: str) -> bool: + 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) @@ -164,11 +184,16 @@ def _build_image_preprocessing_config(self, generation_config: dict[str, Any]): 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"] + ) return build_image_preprocessing_config( processor.image_processor, - dynamic_resolution=bool( - generation_config.get("image_dynamic_resolution", False) - ), + **image_kwargs, ) def _setup_colocated_cuda_graph_managers(self) -> None: @@ -348,11 +373,24 @@ def _initialize_inference_engine(self, mcore_generation_config: dict) -> None: image_preprocessing_config, mcore_generation_config, frame_manifest_magic=CACHED_VIDEO_FRAME_MANIFEST_MAGIC, - video_maintain_aspect_ratio=bool( - mcore_generation_config.get("video_maintain_aspect_ratio", True) - ), ) + # 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, @@ -364,19 +402,11 @@ def _initialize_inference_engine(self, mcore_generation_config: dict) -> None: use_cuda_graphs_for_non_decode_steps=use_cuda_graphs_for_non_decode_steps, use_flashinfer_fused_rope=use_flashinfer_fused_rope, sampling_backend="flashinfer", - async_sched_mode=AsyncScheduleMode( - mcore_generation_config.get("async_sched_mode", "legacy") - ), use_synchronous_zmq_collectives=True, 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"], - vision_embedding_cache_max_bytes=int( - mcore_generation_config.get("vision_embedding_cache_max_bytes", 0) - ), - allow_stale_multimodal_embeddings=bool( - mcore_generation_config.get("allow_stale_multimodal_embeddings", False) - ), + **inference_overrides, prefix_caching_coordinator_policy=PrefixCachingCoordinatorPolicy( "first_prefix_block" ), @@ -388,9 +418,7 @@ def _initialize_inference_engine(self, mcore_generation_config: dict) -> None: ), logging_step_interval=logging_step_interval, num_speculative_tokens=num_speculative_tokens, - logprobs_mode=mcore_generation_config.get( - "logprobs_mode", "processed_logprobs" - ), + logprobs_mode=mcore_generation_config["logprobs_mode"], max_requests=max_requests, image_preprocessing_config=image_preprocessing_config, video_preprocessing_config=video_preprocessing_config, @@ -786,6 +814,25 @@ def _sample_vision_tensors(self, data, index: int): "per-sample PackedTensor value." ) + # `.tensors` is the *physical* segment list. It only lines up with the + # logical row index while media are un-deduplicated (`_row_offsets is + # None`); once `repeat_interleave(..., share_immutable_media=True)` packs + # N rows onto fewer segments, row `index` would silently read another + # row's image. `_validate_multimodal_dedup_capability` rejects + # `deduplicate_multimodal_data` for non-vLLM backends, so this is + # currently unreachable — assert it here, where the assumption is used. + 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 = ( @@ -932,9 +979,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) diff --git a/nemo_rl/models/generation/megatron/utils.py b/nemo_rl/models/generation/megatron/utils.py index 922654dbf08..87ea00c4391 100644 --- a/nemo_rl/models/generation/megatron/utils.py +++ b/nemo_rl/models/generation/megatron/utils.py @@ -23,9 +23,15 @@ def build_image_preprocessing_config( image_processor: Any, *, - dynamic_resolution: bool, + dynamic_resolution: bool | None = None, ) -> ImageProcessingConfig: - """Translate an HF image processor to an MCore config.""" + """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. + """ def read(*names: str) -> Any: for name in names: @@ -74,7 +80,11 @@ def read(*names: str) -> Any: return ImageProcessingConfig( patch_dim=int(patch_dim), - dynamic_resolution=dynamic_resolution, + **( + {} + if dynamic_resolution is None + else {"dynamic_resolution": dynamic_resolution} + ), use_tiling=False, pixel_shuffle=merge_size > 1, spatial_merge_size=merge_size, @@ -90,13 +100,20 @@ def build_video_preprocessing_config( generation_config: dict[str, Any], *, frame_manifest_magic: bytes, - video_maintain_aspect_ratio: bool = True, ) -> VideoProcessingConfig | None: """Build video preprocessing when explicitly enabled by generation config.""" temporal_patch_size = generation_config.get("video_temporal_patch_size") if image_config is None or temporal_patch_size is None: return None + # Read past the early return so text-only and image-only configs never need + # video keys; omit when absent so MCore's own default applies. + video_kwargs: dict[str, Any] = {} + 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( @@ -104,12 +121,18 @@ def build_video_preprocessing_config( dynamic_resolution_max_patches=int(target_num_patches), ) + video_num_frames = generation_config.get("video_num_frames") + if video_num_frames is None: + raise ValueError( + "video_num_frames must be set when video_temporal_patch_size is set." + ) + return VideoProcessingConfig( image_config=image_config, - num_frames=int(generation_config["video_num_frames"]), + num_frames=int(video_num_frames), temporal_patch_size=int(temporal_patch_size), - video_maintain_aspect_ratio=video_maintain_aspect_ratio, frame_manifest_magic=frame_manifest_magic, + **video_kwargs, ) diff --git a/tests/unit/environments/test_nemo_gym.py b/tests/unit/environments/test_nemo_gym.py index b8a401d6cb4..9775072fd79 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 ( @@ -213,6 +214,31 @@ def test_gym_local_video_path_is_inlined_as_data_url(tmp_path): 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): first = tmp_path / "first.mp4" second = tmp_path / "second.mp4" @@ -1606,6 +1632,9 @@ class _MockSelf: 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 diff --git a/tests/unit/environments/test_nemo_gym_mm_utils.py b/tests/unit/environments/test_nemo_gym_mm_utils.py index cdc4e19f67e..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.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 52ff3177a4d..e8a7c9a660b 100644 --- a/tests/unit/models/generation/test_megatron_generation.py +++ b/tests/unit/models/generation/test_megatron_generation.py @@ -40,6 +40,7 @@ model_name = "Qwen/Qwen3-0.6B" +@pytest.mark.mcore def test_direct_megatron_media_request_preserves_preexpanded_prompt(): worker = object.__new__(MegatronGenerationMixin) worker.cfg = {"generation": {"mcore_generation_config": {}}} @@ -68,6 +69,7 @@ def wrapper_supports_modality(wrapper_cls, modality): assert "image" in multi_modal_data +@pytest.mark.mcore def test_direct_megatron_video_request_marks_preexpanded_prompt(): worker = object.__new__(MegatronGenerationMixin) worker.cfg = {"generation": {"mcore_generation_config": {}}} @@ -100,6 +102,7 @@ def wrapper_supports_modality(wrapper_cls, modality): assert "video" in multi_modal_data +@pytest.mark.mcore @pytest.mark.parametrize( ("modality", "num_frames"), [("image", torch.tensor([1])), ("video", torch.tensor([4]))], diff --git a/tests/unit/utils/test_venvs.py b/tests/unit/utils/test_venvs.py index 5a5de9a6673..2635689385b 100644 --- a/tests/unit/utils/test_venvs.py +++ b/tests/unit/utils/test_venvs.py @@ -19,6 +19,7 @@ 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 @@ -76,3 +77,26 @@ def test_add_hf_modules_cache_does_not_duplicate_pythonpath_entry(): ) 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 9cea2af9739..9c89ffd43d2 100755 --- a/tools/install_audio_deps.sh +++ b/tools/install_audio_deps.sh @@ -24,14 +24,27 @@ if ! python -c "import torchcodec" 2>/dev/null; then "torchcodec==0.11.1" fi -# PyAV is intentionally absent from the base image and must be installed into -# the isolated Megatron policy worker environment that imports it. +# 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] Force-reinstalling PyAV in the Megatron worker environment..." - "$RAY_MEGATRON_PYTHON" -m pip install --no-cache-dir --force-reinstall av -else - echo "[audio-deps] Megatron worker environment not found; skipping PyAV: $RAY_MEGATRON_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." From 21fa8e29c18f4440729a52f6274a56afbb100e67 Mon Sep 17 00:00:00 2001 From: Cory Ye Date: Sun, 30 Aug 2026 16:10:19 -0700 Subject: [PATCH 12/15] Address more review feedback. Signed-off-by: Cory Ye --- docs/design-docs/generation.md | 36 ++++ nemo_rl/data/processors.py | 4 +- nemo_rl/environments/nemo_gym_request.py | 64 +++++++ nemo_rl/models/generation/megatron/config.py | 9 +- .../megatron/megatron_generation.py | 8 +- .../generation/megatron/megatron_worker.py | 115 ++---------- nemo_rl/models/generation/megatron/utils.py | 127 +++++++++++-- ...1_Functional_Tests_GB200_Megatron_Omni.sh} | 12 +- ...h => nemotron_omni_clevr_megatron_1n2g.sh} | 19 +- ... nemotron_omni_gym_video_megatron_1n2g.sh} | 8 +- tests/test_suites/disabled.txt | 7 + ...ircle-count-1n4g-megatron_generation.v1.sh | 34 ++++ ...0ba3b-clevr-1n4g-megatron_generation.v1.sh | 34 ++++ ...0ba3b-clevr-8n4g-megatron_generation.v1.sh | 34 ++++ tests/unit/data/datasets/test_mmpr_tiny.py | 30 ++- .../data/test_multimodal_image_encoding.py | 90 ++++++++- tests/unit/environments/test_nemo_gym.py | 54 +++++- .../generation/test_megatron_generation.py | 78 +++++--- .../test_megatron_generation_utils.py | 176 ++++++++++++++++++ 19 files changed, 767 insertions(+), 172 deletions(-) rename tests/functional/{L1_Functional_Tests_Megatron_Omni.sh => L1_Functional_Tests_GB200_Megatron_Omni.sh} (83%) rename tests/functional/{nemotron_omni_clevr_megatron_1n4g.sh => nemotron_omni_clevr_megatron_1n2g.sh} (91%) rename tests/functional/{nemotron_omni_gym_video_megatron_1n4g.sh => nemotron_omni_gym_video_megatron_1n2g.sh} (98%) create mode 100755 tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-count-1n4g-megatron_generation.v1.sh create mode 100755 tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.sh create mode 100755 tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron_generation.v1.sh create mode 100644 tests/unit/models/generation/test_megatron_generation_utils.py diff --git a/docs/design-docs/generation.md b/docs/design-docs/generation.md index 6dec60fb1b9..8cf4161d217 100644 --- a/docs/design-docs/generation.md +++ b/docs/design-docs/generation.md @@ -143,6 +143,42 @@ 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. +- `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/nemo_rl/data/processors.py b/nemo_rl/data/processors.py index 084fed12f6d..ef8744afb7b 100644 --- a/nemo_rl/data/processors.py +++ b/nemo_rl/data/processors.py @@ -680,7 +680,9 @@ def vlm_hf_data_processor( # 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": None if uses_placeholder else string_formatted_dialog, + "vllm_content": ( + None if uses_placeholder and images else string_formatted_dialog + ), "vllm_images": images, "vllm_audios": audios, "vllm_videos": videos, diff --git a/nemo_rl/environments/nemo_gym_request.py b/nemo_rl/environments/nemo_gym_request.py index 44465b05d7b..46c44b689ec 100644 --- a/nemo_rl/environments/nemo_gym_request.py +++ b/nemo_rl/environments/nemo_gym_request.py @@ -12,12 +12,31 @@ # 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): @@ -34,6 +53,22 @@ def _json_mapping(value: Any, *, field_name: str) -> dict[str, Any]: 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") @@ -51,6 +86,22 @@ def _metadata_extra_body(nemo_gym_example: dict[str, Any]) -> dict[str, Any]: 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") @@ -82,6 +133,19 @@ def _chat_template_kwargs_for_processor( 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): diff --git a/nemo_rl/models/generation/megatron/config.py b/nemo_rl/models/generation/megatron/config.py index 7a7afc585ca..ec0111fe395 100644 --- a/nemo_rl/models/generation/megatron/config.py +++ b/nemo_rl/models/generation/megatron/config.py @@ -58,12 +58,9 @@ class MCoreGenerationSpecificArgs(TypedDict): mamba_inference_ssm_states_dtype: NotRequired[str] mamba_inference_conv_states_dtype: NotRequired[str] - # Raw media preprocessing. Each key is read only on the path that needs it; - # when omitted, MCore's own dataclass default applies (see - # ImageProcessingConfig / VideoProcessingConfig). - # `image_dynamic_resolution` requires an inference wrapper that supports - # images. The `video_*` keys are read only when `video_temporal_patch_size` - # is set, and `video_num_frames` is then required. + # Raw media preprocessing corresponding with Megatron's + # ImageProcessingConfig / VideoProcessingConfig. + # `video_num_frames` is required for video. image_dynamic_resolution: NotRequired[bool] video_num_frames: NotRequired[int] # Frames sampled per video. video_temporal_patch_size: NotRequired[int] # Frames per temporal patch. diff --git a/nemo_rl/models/generation/megatron/megatron_generation.py b/nemo_rl/models/generation/megatron/megatron_generation.py index 78e0919c994..53294bce379 100644 --- a/nemo_rl/models/generation/megatron/megatron_generation.py +++ b/nemo_rl/models/generation/megatron/megatron_generation.py @@ -211,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. @@ -235,7 +235,7 @@ def __init__( # `self.cfg` exposes the `generation` that matches the `GenerationInterface` contract. # `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 if policy is not None else deepcopy(config) + self._policy_config = config self.cfg: MCoreGenerationConfig = config["generation"] # Populated after the first prepare_for_generation (which starts the HTTP server). self.dp_openai_server_base_urls: list[Optional[str]] = [] @@ -254,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 0dc14e6d89d..b49795708b4 100644 --- a/nemo_rl/models/generation/megatron/megatron_worker.py +++ b/nemo_rl/models/generation/megatron/megatron_worker.py @@ -64,9 +64,11 @@ ) 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, @@ -86,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 @@ -179,7 +183,7 @@ def _build_image_preprocessing_config(self, generation_config: dict[str, Any]): inference_wrapper_cls = self._get_megatron_inference_wrapper_cls() if not self._wrapper_supports_modality(inference_wrapper_cls, "image"): return None - processor = getattr(self, "processor", None) + processor = self.processor if processor is None: raise ValueError( "Megatron multimodal generation requires the policy processor." @@ -786,108 +790,19 @@ def _merge_stop_strings( def _sample_vision_tensors(self, 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. It only lines up with the - # logical row index while media are un-deduplicated (`_row_offsets is - # None`); once `repeat_interleave(..., share_immutable_media=True)` packs - # N rows onto fewer segments, row `index` would silently read another - # row's image. `_validate_multimodal_dedup_capability` rejects - # `deduplicate_multimodal_data` for non-vLLM backends, so this is - # currently unreachable — assert it here, where the assumption is used. - 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 + 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.""" - length = int(data["input_lengths"][index].item()) - prompt = data["input_ids"][index, :length].tolist() - imgs, imgs_sizes, num_frames = self._sample_vision_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" - inference_wrapper_cls = self._get_megatron_inference_wrapper_cls() - if not self._wrapper_supports_modality(inference_wrapper_cls, 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, - } + 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 diff --git a/nemo_rl/models/generation/megatron/utils.py b/nemo_rl/models/generation/megatron/utils.py index 87ea00c4391..302b3eab48f 100644 --- a/nemo_rl/models/generation/megatron/utils.py +++ b/nemo_rl/models/generation/megatron/utils.py @@ -13,13 +13,120 @@ # limitations under the License. from dataclasses import replace -from typing import Any +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, *, @@ -102,13 +209,16 @@ def build_video_preprocessing_config( frame_manifest_magic: bytes, ) -> VideoProcessingConfig | None: """Build video preprocessing when explicitly enabled by generation config.""" - temporal_patch_size = generation_config.get("video_temporal_patch_size") - if image_config is None or temporal_patch_size is None: + video_num_frames = generation_config.get("video_num_frames") + if image_config is None or video_num_frames is None: return None - # Read past the early return so text-only and image-only configs never need - # video keys; omit when absent so MCore's own default applies. + # 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"] @@ -121,16 +231,9 @@ def build_video_preprocessing_config( dynamic_resolution_max_patches=int(target_num_patches), ) - video_num_frames = generation_config.get("video_num_frames") - if video_num_frames is None: - raise ValueError( - "video_num_frames must be set when video_temporal_patch_size is set." - ) - return VideoProcessingConfig( image_config=image_config, num_frames=int(video_num_frames), - temporal_patch_size=int(temporal_patch_size), frame_manifest_magic=frame_manifest_magic, **video_kwargs, ) diff --git a/tests/functional/L1_Functional_Tests_Megatron_Omni.sh b/tests/functional/L1_Functional_Tests_GB200_Megatron_Omni.sh similarity index 83% rename from tests/functional/L1_Functional_Tests_Megatron_Omni.sh rename to tests/functional/L1_Functional_Tests_GB200_Megatron_Omni.sh index 1d95ae917ed..c8f916ef053 100755 --- a/tests/functional/L1_Functional_Tests_Megatron_Omni.sh +++ b/tests/functional/L1_Functional_Tests_GB200_Megatron_Omni.sh @@ -21,16 +21,16 @@ PROJECT_ROOT=$(realpath "${SCRIPT_DIR}/../..") cd "${PROJECT_ROOT}" GPU_COUNT=$(nvidia-smi --query-gpu=index --format=csv,noheader | wc -l) -if (( GPU_COUNT < 4 )); then - echo "SKIP: Nemotron Omni functional tests require at least four GPUs" +if (( GPU_COUNT < 2 )); then + echo "SKIP: Nemotron Omni functional tests require at least two GB200 GPUs" exit 0 fi -# The recipes are intentionally 1n4g even when the CI runner exposes eight GPUs. -export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3}" +# 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_1n4g.sh -time uv run --no-sync bash ./tests/functional/nemotron_omni_gym_video_megatron_1n4g.sh +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 diff --git a/tests/functional/nemotron_omni_clevr_megatron_1n4g.sh b/tests/functional/nemotron_omni_clevr_megatron_1n2g.sh similarity index 91% rename from tests/functional/nemotron_omni_clevr_megatron_1n4g.sh rename to tests/functional/nemotron_omni_clevr_megatron_1n2g.sh index 49a02dc97e1..a1156ce9436 100755 --- a/tests/functional/nemotron_omni_clevr_megatron_1n4g.sh +++ b/tests/functional/nemotron_omni_clevr_megatron_1n2g.sh @@ -12,8 +12,8 @@ if [[ -z "${HF_TOKEN:-}" ]]; then fi GPU_COUNT=$(nvidia-smi --query-gpu=index --format=csv,noheader | wc -l) -if (( GPU_COUNT < 4 )); then - echo "SKIP: Omni CLEVR Megatron smoke requires at least four visible GPUs" +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) @@ -81,10 +81,10 @@ PY uv run --no-sync python examples/run_vlm_grpo.py \ --config examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.yaml \ cluster.num_nodes=1 \ - cluster.gpus_per_node=4 \ + 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=4 \ - policy.megatron_cfg.expert_model_parallel_size=4 \ + 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 \ @@ -97,16 +97,17 @@ uv run --no-sync python examples/run_vlm_grpo.py \ policy.generation.backend=megatron \ policy.generation.colocated.enabled=true \ policy.generation.colocated.resources.num_nodes=1 \ - policy.generation.colocated.resources.gpus_per_node=4 \ + policy.generation.colocated.resources.gpus_per_node=2 \ policy.generation.max_new_tokens=128 \ - policy.generation.mcore_generation_config.tensor_model_parallel_size=4 \ - policy.generation.mcore_generation_config.expert_model_parallel_size=4 \ + 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=8 \ policy.generation.mcore_generation_config.cuda_graph_impl="${MEGATRON_CUDA_GRAPH_IMPL}" \ policy.generation.mcore_generation_config.inference_cuda_graph_scope=block \ policy.generation.mcore_generation_config.num_cuda_graphs=-1 \ @@ -114,6 +115,8 @@ uv run --no-sync python examples/run_vlm_grpo.py \ 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=4096 \ + policy.generation.mcore_generation_config.max_tokens=4096 \ data.train.dataset_name=ResponseDataset \ ++data.train.data_path="${TRAIN_PATH}" \ data.train.split=train \ diff --git a/tests/functional/nemotron_omni_gym_video_megatron_1n4g.sh b/tests/functional/nemotron_omni_gym_video_megatron_1n2g.sh similarity index 98% rename from tests/functional/nemotron_omni_gym_video_megatron_1n4g.sh rename to tests/functional/nemotron_omni_gym_video_megatron_1n2g.sh index 798d3b735fa..d4d04f733ce 100755 --- a/tests/functional/nemotron_omni_gym_video_megatron_1n4g.sh +++ b/tests/functional/nemotron_omni_gym_video_megatron_1n2g.sh @@ -12,8 +12,8 @@ if [[ -z "${HF_TOKEN:-}" ]]; then fi GPU_COUNT=$(nvidia-smi --query-gpu=index --format=csv,noheader | wc -l) -if (( GPU_COUNT < 4 )); then - echo "SKIP: Omni Gym-video Megatron smoke requires at least four GPUs" +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) @@ -69,7 +69,7 @@ uv run --no-sync examples/nemo_gym/prepare_video_dataset.py convert \ 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=4 \ + 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 \ @@ -85,7 +85,7 @@ uv run --no-sync python examples/nemo_gym/run_grpo_nemo_gym.py \ ++policy.megatron_cfg.optimizer.store_param_remainders=true \ policy.generation.backend=megatron \ ++policy.generation.bad_words=null \ - policy.generation.colocated.enabled=false \ + 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 \ diff --git a/tests/test_suites/disabled.txt b/tests/test_suites/disabled.txt index c06a9f253b9..bea832b11be 100644 --- a/tests/test_suites/disabled.txt +++ b/tests/test_suites/disabled.txt @@ -27,6 +27,13 @@ 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 these multimodal Megatron generation +# functional tests, add golden convergence metrics, and move them to the +# appropriate recurring suite once its resource budget permits. +tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-count-1n4g-megatron_generation.v1.sh +tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.sh +tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-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/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-1n4g-megatron_generation.v1.sh b/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.sh new file mode 100755 index 00000000000..f8d2e75cdf8 --- /dev/null +++ b/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-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..bb6baa7455a --- /dev/null +++ b/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-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=8 +GPUS_PER_NODE=4 +STEPS_PER_RUN=10 +MAX_STEPS=10 +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/unit/data/datasets/test_mmpr_tiny.py b/tests/unit/data/datasets/test_mmpr_tiny.py index bd536cfd02f..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( diff --git a/tests/unit/data/test_multimodal_image_encoding.py b/tests/unit/data/test_multimodal_image_encoding.py index d020fa9e2e0..4d7f9b23322 100644 --- a/tests/unit/data/test_multimodal_image_encoding.py +++ b/tests/unit/data/test_multimodal_image_encoding.py @@ -143,9 +143,93 @@ def tracking_executor(*args, **kwargs): ] for example in examples: parts = example["responses_create_params"]["input"][0]["content"] - assert parts[0]["image_url"] == expected[first] - assert parts[1]["image_url"] == expected[first] - assert parts[2]["image_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): diff --git a/tests/unit/environments/test_nemo_gym.py b/tests/unit/environments/test_nemo_gym.py index 9775072fd79..7d1fe348579 100644 --- a/tests/unit/environments/test_nemo_gym.py +++ b/tests/unit/environments/test_nemo_gym.py @@ -52,7 +52,11 @@ nemo_gym_example_to_video_datum_spec, normalize_media_in_examples, ) -from nemo_rl.environments.nemo_gym_request import _metadata_extra_body +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, _flatten_nemotron_video_frame_messages, @@ -343,6 +347,54 @@ 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": { diff --git a/tests/unit/models/generation/test_megatron_generation.py b/tests/unit/models/generation/test_megatron_generation.py index e8a7c9a660b..3e94c8a2a0a 100644 --- a/tests/unit/models/generation/test_megatron_generation.py +++ b/tests/unit/models/generation/test_megatron_generation.py @@ -30,6 +30,9 @@ 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 ( @@ -41,28 +44,33 @@ @pytest.mark.mcore -def test_direct_megatron_media_request_preserves_preexpanded_prompt(): +def test_multimodal_preprocessing_requires_policy_processor(): + class _ImageWrapper: + supports_image = True + worker = object.__new__(MegatronGenerationMixin) - worker.cfg = {"generation": {"mcore_generation_config": {}}} + worker._get_megatron_inference_wrapper_cls = lambda: _ImageWrapper - def sample_vision_tensors(data, index): - return torch.ones(1, 2, 4), torch.tensor([[2, 2]]), None + with pytest.raises(ValueError, match="requires the policy processor"): + worker._build_image_preprocessing_config({}) - def get_wrapper_cls(): - return object - def wrapper_supports_modality(wrapper_cls, modality): - return modality == "image" +@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 - worker._sample_vision_tensors = sample_vision_tensors - worker._get_megatron_inference_wrapper_cls = get_wrapper_cls - worker._wrapper_supports_modality = wrapper_supports_modality data = { "input_ids": torch.tensor([[10, 99, 99, 20, 0]]), "input_lengths": torch.tensor([4]), } - prompt, multi_modal_data = worker._build_prompt_and_multimodal_data(data, 0) + 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 @@ -70,32 +78,44 @@ def wrapper_supports_modality(wrapper_cls, modality): @pytest.mark.mcore -def test_direct_megatron_video_request_marks_preexpanded_prompt(): - worker = object.__new__(MegatronGenerationMixin) - worker.cfg = {"generation": {"mcore_generation_config": {}}} +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}" + ), + ) - def sample_vision_tensors(data, index): + 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]), ) - def get_wrapper_cls(): - return object - - def wrapper_supports_modality(wrapper_cls, modality): - return modality == "video" - - worker._sample_vision_tensors = sample_vision_tensors - worker._get_megatron_inference_wrapper_cls = get_wrapper_cls - worker._wrapper_supports_modality = wrapper_supports_modality data = { "input_ids": torch.tensor([[10, 99, 99, 20]]), "input_lengths": torch.tensor([4]), } - prompt, multi_modal_data = worker._build_prompt_and_multimodal_data(data, 0) + 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 @@ -785,6 +805,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_utils.py b/tests/unit/models/generation/test_megatron_generation_utils.py new file mode 100644 index 00000000000..5ac71d30cef --- /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, + ) + + assert config.patch_dim == 16 + assert config.dynamic_resolution is True + 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 From 1461a71bc799e24a793c7461ae7633c7156dac32 Mon Sep 17 00:00:00 2001 From: Cory Ye Date: Sun, 30 Aug 2026 21:02:32 -0700 Subject: [PATCH 13/15] Add video_model_type to RL mcore gen config, and lint. Signed-off-by: Cory Ye --- docs/design-docs/generation.md | 1 + nemo_rl/models/generation/megatron/config.py | 1 + .../generation/megatron/megatron_worker.py | 4 +++ nemo_rl/models/generation/megatron/utils.py | 8 ++++++ .../data/datasets/test_audiomcq_dataset.py | 25 +++++++++++-------- tests/unit/environments/test_nemo_gym.py | 4 +-- .../generation/test_megatron_generation.py | 22 ++++++++++++++++ .../test_megatron_generation_utils.py | 6 ++--- 8 files changed, 54 insertions(+), 17 deletions(-) diff --git a/docs/design-docs/generation.md b/docs/design-docs/generation.md index 8cf4161d217..87c40084a58 100644 --- a/docs/design-docs/generation.md +++ b/docs/design-docs/generation.md @@ -170,6 +170,7 @@ data: ``` - `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. diff --git a/nemo_rl/models/generation/megatron/config.py b/nemo_rl/models/generation/megatron/config.py index ec0111fe395..98d0733a80e 100644 --- a/nemo_rl/models/generation/megatron/config.py +++ b/nemo_rl/models/generation/megatron/config.py @@ -61,6 +61,7 @@ class MCoreGenerationSpecificArgs(TypedDict): # 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. diff --git a/nemo_rl/models/generation/megatron/megatron_worker.py b/nemo_rl/models/generation/megatron/megatron_worker.py index b49795708b4..737c89664ec 100644 --- a/nemo_rl/models/generation/megatron/megatron_worker.py +++ b/nemo_rl/models/generation/megatron/megatron_worker.py @@ -195,6 +195,10 @@ def _build_image_preprocessing_config(self, generation_config: dict[str, Any]): 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, diff --git a/nemo_rl/models/generation/megatron/utils.py b/nemo_rl/models/generation/megatron/utils.py index 302b3eab48f..6e40693e2d5 100644 --- a/nemo_rl/models/generation/megatron/utils.py +++ b/nemo_rl/models/generation/megatron/utils.py @@ -131,6 +131,7 @@ 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. @@ -138,6 +139,8 @@ def build_image_preprocessing_config( 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: @@ -192,6 +195,11 @@ def read(*names: str) -> Any: 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, 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/environments/test_nemo_gym.py b/tests/unit/environments/test_nemo_gym.py index 7d1fe348579..eefd569b72a 100644 --- a/tests/unit/environments/test_nemo_gym.py +++ b/tests/unit/environments/test_nemo_gym.py @@ -375,9 +375,7 @@ def test_chat_template_kwargs_for_processor_defaults_to_empty(): def test_chat_template_kwargs_for_processor_rejects_invalid_json(): example = { - "responses_create_params": { - "metadata": {"chat_template_kwargs": "not-json"} - } + "responses_create_params": {"metadata": {"chat_template_kwargs": "not-json"}} } with pytest.raises(ValueError, match="chat_template_kwargs"): diff --git a/tests/unit/models/generation/test_megatron_generation.py b/tests/unit/models/generation/test_megatron_generation.py index 3e94c8a2a0a..f996690deef 100644 --- a/tests/unit/models/generation/test_megatron_generation.py +++ b/tests/unit/models/generation/test_megatron_generation.py @@ -55,6 +55,28 @@ class _ImageWrapper: 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): diff --git a/tests/unit/models/generation/test_megatron_generation_utils.py b/tests/unit/models/generation/test_megatron_generation_utils.py index 5ac71d30cef..f53c5ae0d0d 100644 --- a/tests/unit/models/generation/test_megatron_generation_utils.py +++ b/tests/unit/models/generation/test_megatron_generation_utils.py @@ -43,10 +43,12 @@ def test_build_image_config_handles_dict_patch_and_downsample_ratio(): 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 @@ -63,9 +65,7 @@ def test_build_image_config_handles_dict_patch_and_downsample_ratio(): ({}, 1), ], ) -def test_build_image_config_merge_size_fallbacks( - merge_fields, expected_merge_size -): +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 From 85fc69d03e59652aec1ce6ee48d8d521b0fdca9b Mon Sep 17 00:00:00 2001 From: Cory Ye Date: Tue, 1 Sep 2026 00:51:43 -0700 Subject: [PATCH 14/15] Add short convergence test with CLEVR. Signed-off-by: Cory Ye --- ...a3b-clevr-1n4g-megatron_generation.v1.yaml | 72 ------------------- ...a3b-clevr-8n4g-megatron_generation.v1.yaml | 24 +++---- tests/check_metrics.py | 8 +++ .../nemotron_omni_clevr_megatron_1n2g.sh | 2 +- tests/test_suites/disabled.txt | 7 +- tests/test_suites/nightly_gb200.txt | 1 + ...0ba3b-clevr-1n4g-megatron_generation.v1.sh | 34 --------- ...0ba3b-clevr-8n4g-megatron_generation.v1.sh | 23 ++++-- tests/unit/test_check_metrics.py | 30 +++++++- 9 files changed, 71 insertions(+), 130 deletions(-) delete mode 100644 examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.yaml delete mode 100755 tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.sh diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.yaml deleted file mode 100644 index 773e057a1da..00000000000 --- a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.yaml +++ /dev/null @@ -1,72 +0,0 @@ -defaults: ../../vlm_grpo_3B_megatron.yaml -grpo: - num_prompts_per_step: 2 - num_generations_per_prompt: 8 - max_num_steps: 4 - val_period: 0 - 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-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 - 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 - max_new_tokens: 512 - stop_strings: - - - - - - - - - - - - - mcore_generation_config: - buffer_size_gb: 8 - num_cuda_graphs: -1 - max_tokens: ${policy.max_total_sequence_length} - refit_backend: gloo - 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 - colocated: - enabled: false - resources: - gpus_per_node: 2 - num_nodes: 1 -data: - num_workers: 0 - 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-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 index ba46f1079ef..0f741856edf 100644 --- 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 @@ -1,9 +1,10 @@ defaults: ../../vlm_grpo_3B_megatron.yaml grpo: - num_prompts_per_step: 6 + num_prompts_per_step: 12 num_generations_per_prompt: 8 + val_at_start: true + val_at_end: true val_batch_size: 64 - max_val_samples: 64 async_grpo: enabled: true max_trajectory_age_steps: 2 @@ -37,25 +38,24 @@ policy: generation: backend: megatron max_new_tokens: 2048 - stop_strings: - - - - - - - - - - - - + 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} - refit_backend: gloo - transformer_impl: transformer_engine + 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_pad_experts_for_cuda_graph_inference: 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 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/nemotron_omni_clevr_megatron_1n2g.sh b/tests/functional/nemotron_omni_clevr_megatron_1n2g.sh index a1156ce9436..a3a606684ae 100755 --- a/tests/functional/nemotron_omni_clevr_megatron_1n2g.sh +++ b/tests/functional/nemotron_omni_clevr_megatron_1n2g.sh @@ -79,7 +79,7 @@ for path, count in ((os.environ["TRAIN_PATH"], 64), (os.environ["VAL_PATH"], 2)) PY uv run --no-sync python examples/run_vlm_grpo.py \ - --config examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.yaml \ + --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}" \ diff --git a/tests/test_suites/disabled.txt b/tests/test_suites/disabled.txt index bea832b11be..44e2cf0e5af 100644 --- a/tests/test_suites/disabled.txt +++ b/tests/test_suites/disabled.txt @@ -27,12 +27,9 @@ 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 these multimodal Megatron generation -# functional tests, add golden convergence metrics, and move them to the -# appropriate recurring suite once its resource budget permits. +# 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 -tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.sh -tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-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 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-clevr-1n4g-megatron_generation.v1.sh b/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.sh deleted file mode 100755 index f8d2e75cdf8..00000000000 --- a/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.sh +++ /dev/null @@ -1,34 +0,0 @@ -#!/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 index bb6baa7455a..bc36a391470 100755 --- 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 @@ -2,14 +2,15 @@ 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. +# 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=10 -MAX_STEPS=10 +STEPS_PER_RUN=50 +MAX_STEPS=50 NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN )) NUM_MINUTES=120 # ===== END CONFIG ===== @@ -20,15 +21,27 @@ 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=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/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.""" From 95b0a899f76ed9c91de3fa75df532accd786ccea Mon Sep 17 00:00:00 2001 From: Cory Ye Date: Wed, 2 Sep 2026 00:35:03 -0700 Subject: [PATCH 15/15] Fix functional test and port reservation deadlock. Signed-off-by: Cory Ye --- nemo_rl/distributed/held_port.py | 35 +++++++++++-- nemo_rl/models/megatron/setup.py | 24 +++++++-- .../nemotron_omni_clevr_megatron_1n2g.sh | 35 ++++++++++--- .../nemotron_omni_gym_video_megatron_1n2g.sh | 51 +++++++++++-------- .../test_megatron_generation_parse.py | 16 +++++- .../models/megatron/test_megatron_setup.py | 17 +++++++ 6 files changed, 141 insertions(+), 37 deletions(-) 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/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index 295d931065d..afc14e8dcf0 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -270,14 +270,30 @@ def _sync_distrib_opt(distrib_opt): def _resolve_optimizer_dtype_kwargs(optimizer_cfg: dict[str, Any]) -> dict[str, Any]: - """Resolve optimizer dtype strings.""" - from megatron.bridge.utils.activation_map import str_to_dtype - + """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): - resolved[key] = str_to_dtype(value) + 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 diff --git a/tests/functional/nemotron_omni_clevr_megatron_1n2g.sh b/tests/functional/nemotron_omni_clevr_megatron_1n2g.sh index a3a606684ae..0a61b7b71c6 100755 --- a/tests/functional/nemotron_omni_clevr_megatron_1n2g.sh +++ b/tests/functional/nemotron_omni_clevr_megatron_1n2g.sh @@ -20,6 +20,13 @@ DETECTED_CUDA_ARCH=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader -i 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 @@ -43,6 +50,9 @@ 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 @@ -78,6 +88,9 @@ for path, count in ((os.environ["TRAIN_PATH"], 64), (os.environ["VAL_PATH"], 2)) 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 \ @@ -89,11 +102,16 @@ uv run --no-sync python examples/run_vlm_grpo.py \ 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=true \ + ++policy.megatron_cfg.optimizer.store_param_remainders=false \ policy.generation.backend=megatron \ policy.generation.colocated.enabled=true \ policy.generation.colocated.resources.num_nodes=1 \ @@ -107,16 +125,17 @@ uv run --no-sync python examples/run_vlm_grpo.py \ 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=8 \ + 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=block \ - policy.generation.mcore_generation_config.num_cuda_graphs=-1 \ + 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=4096 \ - policy.generation.mcore_generation_config.max_tokens=4096 \ + 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 \ @@ -144,5 +163,5 @@ uv run --no-sync python examples/run_vlm_grpo.py \ 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/loss"]) < 1e6' \ - 'min(data["train/loss"]) > -1e6' + '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 index d4d04f733ce..3fd6aad2675 100755 --- a/tests/functional/nemotron_omni_gym_video_megatron_1n2g.sh +++ b/tests/functional/nemotron_omni_gym_video_megatron_1n2g.sh @@ -19,7 +19,20 @@ 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}" -MOE_PAD_EXPERTS_FOR_CG=false +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}" @@ -66,6 +79,9 @@ 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 \ @@ -78,11 +94,15 @@ uv run --no-sync python examples/nemo_gym/run_grpo_nemo_gym.py \ 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=true \ + ++policy.megatron_cfg.optimizer.store_param_remainders=false \ policy.generation.backend=megatron \ ++policy.generation.bad_words=null \ policy.generation.colocated.enabled=true \ @@ -98,21 +118,21 @@ uv run --no-sync python examples/nemo_gym/run_grpo_nemo_gym.py \ 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=8 \ - policy.generation.mcore_generation_config.cuda_graph_impl=none \ - policy.generation.mcore_generation_config.inference_cuda_graph_scope=none \ - policy.generation.mcore_generation_config.num_cuda_graphs=0 \ + 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=4096 \ - policy.generation.mcore_generation_config.max_tokens=4096 \ + 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=4096 \ + 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 \ @@ -140,15 +160,6 @@ uv run --no-sync python examples/nemo_gym/run_grpo_nemo_gym.py \ "$@" 2>&1 | tee "${RUN_LOG}" uv run --no-sync tests/json_dump_tb_logs.py "${LOG_DIR}" --output_path "${JSON_METRICS}" - -RECORDED_STEP=$(jq -r \ - 'if has("train/loss") then (."train/loss" | keys | map(tonumber) | max // 0) else 0 end' \ - "${JSON_METRICS}") -if (( RECORDED_STEP < 1 )); then - echo "[ERROR] Expected at least one completed Gym-video training step" - exit 1 -fi - uv run --no-sync tests/check_metrics.py "${JSON_METRICS}" \ - 'max(data["train/loss"]) < 1e6' \ - 'min(data["train/loss"]) > -1e6' + 'max(data["train/gen_kl_error"]) < 0.05' \ + 'all_finite(data["train/reward"])' diff --git a/tests/unit/models/generation/test_megatron_generation_parse.py b/tests/unit/models/generation/test_megatron_generation_parse.py index 2a9a8655c98..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( 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."""