diff --git a/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge b/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge index d352aceda8e..b11414c71b1 160000 --- a/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge +++ b/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge @@ -1 +1 @@ -Subproject commit d352aceda8ed4136f1db787bcf449c1b210a2438 +Subproject commit b11414c71b15e54d333eb49346ed199f20fa9021 diff --git a/docs/design-docs/generation.md b/docs/design-docs/generation.md index 6dec60fb1b9..87c40084a58 100644 --- a/docs/design-docs/generation.md +++ b/docs/design-docs/generation.md @@ -143,6 +143,43 @@ The `mcore_generation_config` section controls Megatron Core inference engine be - **num_cuda_graphs**: Number of CUDA graphs to pre-allocate for different batch sizes. More graphs can improve performance by avoiding runtime graph capture, but consume more memory. - **max_tokens**: Maximum total number of tokens (across all requests) that can be processed simultaneously. This limits the maximum batch size and sequence length combinations. Increasing this might throw OOM depending on vocab size and buffer size allocated. +### Multimodal Megatron Generation + +Megatron inference supports image and video inputs in NeMo-RL. Enable multimodal processing with `policy.is_vlm: true`, use the `megatron` generation backend, and provide a `megatron_inference_wrapper`. The wrapper must subclass `megatron.core.inference.model_inference_wrappers.abstract_model_inference_wrapper.AbstractModelInferenceWrapper` in Megatron-Core and declare `supports_ = True` for each supported modality. + +```yaml +policy: + is_vlm: true + generation: + backend: megatron + mcore_generation_config: + megatron_inference_wrapper: megatron.core.inference.model_inference_wrappers.multimodal.nemotron_omni_inference_wrapper.NemotronOmniInferenceWrapper + image_dynamic_resolution: true + video_num_frames: 16 + video_temporal_patch_size: 2 + video_target_num_patches: 2048 + video_maintain_aspect_ratio: true + vision_embedding_cache_max_bytes: 0 + allow_stale_multimodal_embeddings: false +data: + default: + num_frames: 16 + video_temporal_patch_size: 2 + video_target_num_patches: 2048 + video_maintain_aspect_ratio: true +``` + +- `image_dynamic_resolution` preserves variable image shapes instead of forcing one fixed resolution; for example, a wide image uses a wider patch grid than a square image. +- `vision_model_type` optionally selects the MCore vision encoder type used by image and video preprocessing. Set it to the encoder expected by the inference wrapper; when omitted, MCore uses its default (`radio`). +- `num_frames` controls uniform video-frame sampling. Use `video_num_frames` for the corresponding MCore key. +- `video_temporal_patch_size` groups sampled frames into temporal tubelets; for example, size `2` turns 16 frames into 8 temporal groups. +- `video_target_num_patches` sets `num_patches_per_frame = patch_height * patch_width <= video_target_num_patches`, which produces `num_patches_per_frame * num_frames / video_temporal_patch_size` total video patches prior to spatial merging (i.e. further grouped / concatenated into MxM patch blocks) that are provided to the vision encoder. +- `video_maintain_aspect_ratio=true` keeps `patch_width / patch_height ~= source_width / source_height`; `false` uses `patch_width = patch_height ~= sqrt(video_target_num_patches)` (for example, `sqrt(256) = 16`). +- `vision_embedding_cache_max_bytes` limits GPU memory used to reuse vision embeddings for repeated media; `0` disables the cache, while `1073741824` permits up to 1 GiB. +- `allow_stale_multimodal_embeddings` controls whether cached embeddings survive model-weight changes. Keep it `false` for RL refits; use `true` only when weights remain fixed. +- `expose_http_server` should be `true` for NeMo Gym. + +Keep the video preprocessing values identical in `data.default` and `mcore_generation_config` to avoid disparity between the training policy and inference generation. ## Usage Examples diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-16n8g-megatron-tp4ep4-async-gym-video.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-16n8g-megatron-tp4ep4-async-gym-video.v1.yaml index 2647cfc4d8a..20799328e5d 100644 --- a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-16n8g-megatron-tp4ep4-async-gym-video.v1.yaml +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-16n8g-megatron-tp4ep4-async-gym-video.v1.yaml @@ -48,6 +48,10 @@ policy: make_sequence_length_divisible_by: 32 generation: bad_words: [] + mcore_generation_config: + image_dynamic_resolution: true + logprobs_mode: raw_logprobs + megatron_inference_wrapper: megatron.core.inference.model_inference_wrappers.multimodal.nemotron_omni_inference_wrapper.NemotronOmniInferenceWrapper vllm_cfg: reset_encoder_cache_after_weight_update: false video: diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-click-2n8g-megatron-tp2ep8.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-click-2n8g-megatron-tp2ep8.v1.yaml index 38ff8f89fe6..2f198cb56c6 100644 --- a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-click-2n8g-megatron-tp2ep8.v1.yaml +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-click-2n8g-megatron-tp2ep8.v1.yaml @@ -90,6 +90,9 @@ policy: max_tokens: ${policy.max_total_sequence_length} expose_http_server: true enable_prefix_caching: true + image_dynamic_resolution: true + logprobs_mode: raw_logprobs + megatron_inference_wrapper: megatron.core.inference.model_inference_wrappers.multimodal.nemotron_omni_inference_wrapper.NemotronOmniInferenceWrapper parsers: - deepseek-r1-reasoning - qwen3-coder-tool diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-count-1n4g-megatron_generation.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-count-1n4g-megatron_generation.v1.yaml new file mode 100644 index 00000000000..aca1e6015a6 --- /dev/null +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-count-1n4g-megatron_generation.v1.yaml @@ -0,0 +1,96 @@ +defaults: ../../vlm_grpo_3B_megatron.yaml +grpo: + num_prompts_per_step: 2 + num_generations_per_prompt: 8 + max_num_steps: 4 + val_period: 0 + max_val_samples: null + val_batch_size: null + async_grpo: + enabled: true + max_trajectory_age_steps: 2 + in_flight_weight_updates: true +loss_fn: + reference_policy_kl_penalty: 0.0 + use_importance_sampling_correction: true +checkpointing: + enabled: false + checkpoint_dir: results/nemo-rl-omni/nemotron-omni-circle-count-1n4g +policy: + model_name: nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16 + is_vlm: true + train_global_batch_size: 16 + logprob_batch_size: 1 + max_total_sequence_length: 8192 + sequence_packing: + enabled: true + megatron_cfg: + env_vars: + TORCH_CUDA_ARCH_LIST: '10.0' + tensor_model_parallel_size: 2 + expert_model_parallel_size: 2 + sequence_parallel: true + bias_activation_fusion: false + activation_checkpointing: true + generation: + backend: megatron + bad_words: null + mcore_generation_config: + expose_http_server: true + buffer_size_gb: 8 + num_cuda_graphs: -1 + max_tokens: ${policy.max_total_sequence_length} + transformer_impl: transformer_engine + activation_checkpointing: false + tensor_model_parallel_size: 2 + expert_model_parallel_size: 2 + expert_tensor_parallel_size: 1 + sequence_parallel: true + moe_pad_experts_for_cuda_graph_inference: true + image_dynamic_resolution: true + logprobs_mode: raw_logprobs + megatron_inference_wrapper: megatron.core.inference.model_inference_wrappers.multimodal.nemotron_omni_inference_wrapper.NemotronOmniInferenceWrapper + vllm_cfg: + async_engine: true + expose_http_server: true + colocated: + enabled: false + resources: + gpus_per_node: 2 + num_nodes: 1 +data: + _override_: true + max_input_seq_length: null + shuffle: false + num_workers: 0 + train: + data_path: 3rdparty/Gym-workspace/Gym/resources_servers/circle_count/data/example.jsonl + validation: + data_path: 3rdparty/Gym-workspace/Gym/resources_servers/circle_count/data/example.jsonl + default: + dataset_name: NemoGymDataset + env_name: nemo_gym + prompt_file: null + processor: nemo_gym_data_processor +env: + _override_: true + should_use_nemo_gym: true + should_log_nemo_gym_responses: true + nemo_gym: + is_trajectory_collection: false + port_range_low: 5000 + port_range_high: 5999 + config_paths: + - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml + - resources_servers/circle_count/configs/circle_count.yaml + circle_count_simple_agent: + responses_api_agents: + simple_agent: + max_steps: 1 +logger: + tensorboard_enabled: false + wandb: + project: nemo-rl-omni + name: nemotron-omni-circle-count-1n4g +cluster: + gpus_per_node: 4 diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron-single-controller-async.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron-single-controller-async.v1.yaml new file mode 100644 index 00000000000..ae9150e691d --- /dev/null +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron-single-controller-async.v1.yaml @@ -0,0 +1,51 @@ +# NeMo-RL v2 SingleController overlay for one-node, non-colocated Omni CLEVR. +defaults: ./vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.yaml + +grpo: + async_grpo: null + val_period: 0 + val_at_start: false + val_at_end: false + +data_plane: + enabled: true + impl: transfer_queue + backend: simple + claim_meta_poll_interval_s: 0.5 + simple: + num_storage_units: 2 + +async_rl: + sampler: + name: in_order + max_lookahead_versions: 1 + recompute_kv_cache_after_weight_updates: false + min_groups_for_streaming_train: ${grpo.num_prompts_per_step} + max_inflight_prompts: ${mul:${grpo.num_prompts_per_step}, 2} + max_buffered_rollouts: ${mul:${grpo.num_prompts_per_step}, 2} + diagnostics: true + +policy: + generation: + backend: megatron + colocated: + enabled: false + resources: + num_nodes: 1 + gpus_per_node: 2 + mcore_generation_config: + transformer_impl: inference_optimized + moe_router_dtype: fp32 + moe_pad_experts_for_cuda_graph_inference: false + cuda_graph_impl: local + inference_cuda_graph_scope: block + num_cuda_graphs: -1 + use_cuda_graphs_for_non_decode_steps: false + enable_chunked_prefill: true + async_sched_mode: async + kv_cache_management_mode: persist + refit_backend: nccl + +cluster: + num_nodes: 1 + gpus_per_node: 4 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..773e057a1da --- /dev/null +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.yaml @@ -0,0 +1,72 @@ +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-single-controller-async.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-single-controller-async.v1.yaml new file mode 100644 index 00000000000..be31246f1b4 --- /dev/null +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-single-controller-async.v1.yaml @@ -0,0 +1,51 @@ +# NeMo-RL v2 SingleController overlay for eight-node, non-colocated Omni CLEVR. +defaults: ./vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron_generation.v1.yaml + +grpo: + async_grpo: null + val_period: 0 + val_at_start: false + val_at_end: false + +data_plane: + enabled: true + impl: transfer_queue + backend: simple + claim_meta_poll_interval_s: 0.5 + simple: + num_storage_units: 16 + +async_rl: + sampler: + name: in_order + max_lookahead_versions: 1 + recompute_kv_cache_after_weight_updates: false + min_groups_for_streaming_train: ${grpo.num_prompts_per_step} + max_inflight_prompts: ${mul:${grpo.num_prompts_per_step}, 2} + max_buffered_rollouts: ${mul:${grpo.num_prompts_per_step}, 2} + diagnostics: false + +policy: + generation: + backend: megatron + colocated: + enabled: false + resources: + num_nodes: 6 + gpus_per_node: 4 + mcore_generation_config: + transformer_impl: inference_optimized + moe_router_dtype: fp32 + moe_pad_experts_for_cuda_graph_inference: false + cuda_graph_impl: local + inference_cuda_graph_scope: block + num_cuda_graphs: -1 + use_cuda_graphs_for_non_decode_steps: false + enable_chunked_prefill: true + async_sched_mode: async + kv_cache_management_mode: persist + refit_backend: nccl + +cluster: + num_nodes: 8 + gpus_per_node: 4 diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron_generation.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron_generation.v1.yaml new file mode 100644 index 00000000000..ba46f1079ef --- /dev/null +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron_generation.v1.yaml @@ -0,0 +1,77 @@ +defaults: ../../vlm_grpo_3B_megatron.yaml +grpo: + num_prompts_per_step: 6 + num_generations_per_prompt: 8 + val_batch_size: 64 + max_val_samples: 64 + async_grpo: + enabled: true + max_trajectory_age_steps: 2 + in_flight_weight_updates: true +loss_fn: + reference_policy_kl_penalty: 0.0 + use_importance_sampling_correction: true +checkpointing: + enabled: false + checkpoint_dir: results/nemo-rl-omni/nemotron-omni-clevr-megatron-8n4g +policy: + model_name: nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16 + is_vlm: true + train_global_batch_size: ${mul:${grpo.num_prompts_per_step}, ${grpo.num_generations_per_prompt}} + logprob_batch_size: 1 + max_total_sequence_length: 4096 + sequence_packing: + enabled: true + megatron_cfg: + env_vars: + TORCH_CUDA_ARCH_LIST: '10.0' + tensor_model_parallel_size: 8 + expert_model_parallel_size: 8 + sequence_parallel: true + bias_activation_fusion: false + activation_checkpointing: true + optimizer: + exp_avg_dtype: bfloat16 + exp_avg_sq_dtype: bfloat16 + store_param_remainders: true + generation: + backend: megatron + max_new_tokens: 2048 + 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: 8 + expert_model_parallel_size: 8 + 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: 4 + num_nodes: 6 +data: + default: + prompt_file: examples/prompts/clevr_cogent_cot_nemotron_omni.txt +logger: + tensorboard_enabled: false + wandb: + project: nemo-rl-omni + name: nemotron-omni-clevr-megatron-8n4g +cluster: + gpus_per_node: 4 + num_nodes: 8 diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-1n4g-megatron-single-controller-async.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-1n4g-megatron-single-controller-async.v1.yaml new file mode 100644 index 00000000000..4eca9a9e08f --- /dev/null +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-1n4g-megatron-single-controller-async.v1.yaml @@ -0,0 +1,59 @@ +# NeMo-RL v2 SingleController overlay for one-node, non-colocated Omni VSTAT. +defaults: ./vlm_grpo-nemotron-omni-30ba3b-16n8g-megatron-tp4ep4-async-gym-video.v1.yaml + +grpo: + async_grpo: null + val_period: 0 + val_at_start: false + val_at_end: false + deduplicate_multimodal_data: false + overlong_filtering: false + +data_plane: + enabled: true + impl: transfer_queue + backend: simple + claim_meta_poll_interval_s: 0.5 + simple: + num_storage_units: 2 + +async_rl: + sampler: + name: in_order + max_lookahead_versions: 1 + recompute_kv_cache_after_weight_updates: false + min_groups_for_streaming_train: ${grpo.num_prompts_per_step} + max_inflight_prompts: ${mul:${grpo.num_prompts_per_step}, 2} + max_buffered_rollouts: ${mul:${grpo.num_prompts_per_step}, 2} + diagnostics: true + +policy: + megatron_cfg: + tensor_model_parallel_size: 2 + expert_model_parallel_size: 2 + generation: + backend: megatron + colocated: + enabled: false + resources: + num_nodes: 1 + gpus_per_node: 2 + mcore_generation_config: + transformer_impl: inference_optimized + tensor_model_parallel_size: 2 + expert_model_parallel_size: 2 + moe_router_dtype: fp32 + moe_pad_experts_for_cuda_graph_inference: false + cuda_graph_impl: local + inference_cuda_graph_scope: block + num_cuda_graphs: -1 + use_cuda_graphs_for_non_decode_steps: false + enable_chunked_prefill: true + enable_prefix_caching: false + async_sched_mode: async + kv_cache_management_mode: persist + refit_backend: nccl + +cluster: + num_nodes: 1 + gpus_per_node: 4 diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-8n4g-megatron-single-controller-async.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-8n4g-megatron-single-controller-async.v1.yaml new file mode 100644 index 00000000000..6214e10aa2f --- /dev/null +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-8n4g-megatron-single-controller-async.v1.yaml @@ -0,0 +1,54 @@ +# NeMo-RL v2 SingleController overlay for eight-node, non-colocated Omni VSTAT. +defaults: ./vlm_grpo-nemotron-omni-30ba3b-16n8g-megatron-tp4ep4-async-gym-video.v1.yaml + +grpo: + async_grpo: null + val_period: 0 + val_at_start: false + val_at_end: false + deduplicate_multimodal_data: false + overlong_filtering: false + +data_plane: + enabled: true + impl: transfer_queue + backend: simple + claim_meta_poll_interval_s: 0.5 + simple: + num_storage_units: 16 + +async_rl: + sampler: + name: in_order + max_lookahead_versions: 1 + recompute_kv_cache_after_weight_updates: false + min_groups_for_streaming_train: ${grpo.num_prompts_per_step} + max_inflight_prompts: ${mul:${grpo.num_prompts_per_step}, 2} + max_buffered_rollouts: ${mul:${grpo.num_prompts_per_step}, 2} + diagnostics: false + +policy: + generation: + backend: megatron + colocated: + enabled: false + resources: + num_nodes: 6 + gpus_per_node: 4 + mcore_generation_config: + transformer_impl: inference_optimized + moe_router_dtype: fp32 + moe_pad_experts_for_cuda_graph_inference: false + cuda_graph_impl: local + inference_cuda_graph_scope: block + num_cuda_graphs: -1 + use_cuda_graphs_for_non_decode_steps: false + enable_chunked_prefill: true + enable_prefix_caching: false + async_sched_mode: async + kv_cache_management_mode: persist + refit_backend: nccl + +cluster: + num_nodes: 8 + gpus_per_node: 4 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/examples/run_grpo_single_controller.py b/examples/run_grpo_single_controller.py index 3a151016a2a..8456f51e7d1 100644 --- a/examples/run_grpo_single_controller.py +++ b/examples/run_grpo_single_controller.py @@ -126,7 +126,12 @@ def main() -> None: maybe_configure_data_plane_env(config.data_plane) init_ray() - tokenizer = get_tokenizer(config.policy["tokenizer"]) + processor = None + if config.policy.get("is_vlm", False): + processor = get_tokenizer(config.policy["tokenizer"], get_processor=True) + tokenizer = processor.tokenizer + else: + tokenizer = get_tokenizer(config.policy["tokenizer"]) assert config.policy["generation"] is not None, ( "A generation config is required for SC-driven async GRPO" ) @@ -144,7 +149,9 @@ def main() -> None: if bool(config.env.get("should_use_nemo_gym")): setup_nemo_gym_config(config, tokenizer) - actor_args, setup_timing_metrics = setup_single_controller(config, tokenizer) + actor_args, setup_timing_metrics = setup_single_controller( + config, tokenizer, processor=processor + ) print("🚀 Launching SingleControllerActor") sc = SingleControllerActor.remote( 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/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py index f6a6eba9f76..c23cf7c64ce 100644 --- a/nemo_rl/algorithms/single_controller_utils/setup.py +++ b/nemo_rl/algorithms/single_controller_utils/setup.py @@ -69,6 +69,10 @@ ) from nemo_rl.algorithms.utils import set_seed from nemo_rl.data.collate_fn import rl_collate_fn +from nemo_rl.data.multimodal_utils import ( + UNDECLARED_MULTIMODAL_MODEL_INPUTS, + get_multimodal_keys_from_processor, +) from nemo_rl.data.utils import load_dataloader_state, setup_response_data from nemo_rl.data_plane import ( DATA_PLANE_CHECKPOINT_SCHEMA_VERSION, @@ -79,6 +83,7 @@ from nemo_rl.data_plane.schema import ( SC_ROLLOUT_SCHEMA_FIELDS, fields_with_optional_routed_experts, + packed_tensor_wire_fields, ) from nemo_rl.distributed.virtual_cluster import ( RayVirtualCluster, @@ -1013,6 +1018,8 @@ def setup_single_controller( # ========================== # TODO: add validate dataset wiring. use_nemo_gym = should_use_nemo_gym(master_config) + data_tokenizer = processor if processor is not None else tokenizer + is_vlm = processor is not None if use_nemo_gym and generation_config["backend"] not in ("vllm", "megatron"): raise NotImplementedError( "SC NeMo-Gym integration currently supports the vllm and megatron backends only; got " @@ -1023,13 +1030,18 @@ def setup_single_controller( if use_nemo_gym: # NeMo-Gym creates the env actor outside setup_response_data; we wire # it in after generation is up (it needs the OpenAI server URLs). - response_data = setup_response_data(tokenizer, data_config, env_configs=None) + response_data = setup_response_data( + data_tokenizer, data_config, env_configs=None, is_vlm=is_vlm + ) assert len(response_data) == 2 dataset, _val_dataset = response_data env_handles: dict[str, EnvironmentInterface] = {} else: response_data = setup_response_data( - tokenizer, data_config, env_configs=master_config.env + data_tokenizer, + data_config, + env_configs=master_config.env, + is_vlm=is_vlm, ) assert len(response_data) == 4 dataset, _val_dataset, env_handles, _val_env_handles = response_data @@ -1093,6 +1105,7 @@ def setup_single_controller( megatron_reserved_url = None megatron_port_holder = None reserved_http_server_port = None + weight_synchronizer: Optional[WeightSynchronizer] = None if megatron_backend: generation_config["model_name"] = master_config.policy["model_name"] @@ -1252,7 +1265,6 @@ def _build_generation_then_trainer( build_tasks["trainer"] = _build_trainer_and_value # Submit build tasks and get results - weight_synchronizer: Optional[WeightSynchronizer] = None try: with ThreadPoolExecutor(max_workers=len(build_tasks)) as executor: submitted = {k: executor.submit(fn) for k, fn in build_tasks.items()} @@ -1279,7 +1291,9 @@ def _build_generation_then_trainer( train_cluster=train_cluster, inference_cluster=inference_cluster, refit_buffer_size_gb=policy_config.get("refit_buffer_size_gb"), + refit_timeout_s=master_config.async_rl.generation_fleet_health.refit_timeout_s, ) + generation.weight_synchronizer = weight_synchronizer weight_synchronizer.init_communicator() setup_timing_metrics.collective_init_time_s = time.perf_counter() - t0 t0 = time.perf_counter() @@ -1357,12 +1371,24 @@ def _build_generation_then_trainer( # SingleController reuses one partition for the run. Warm every known # tensor field before rollout, policy, and teacher writers become # concurrent; TransferQueue otherwise registers field names lazily. + partition_fields = fields_with_optional_routed_experts( + SC_ROLLOUT_SCHEMA_FIELDS, + enabled=router_replay_enabled(policy_config), + ) + if processor is not None: + partition_fields.extend( + field + for field in packed_tensor_wire_fields( + [ + *get_multimodal_keys_from_processor(processor), + *UNDECLARED_MULTIMODAL_MODEL_INPUTS, + ] + ) + if field not in partition_fields + ) dp_client.register_partition( partition_id=partition_id, - fields=fields_with_optional_routed_experts( - SC_ROLLOUT_SCHEMA_FIELDS, - enabled=router_replay_enabled(policy_config), - ), + fields=partition_fields, num_samples=( master_config.async_rl.max_buffered_rollouts * algo_cfg.num_generations_per_prompt @@ -1383,6 +1409,7 @@ def _build_generation_then_trainer( refit_buffer_size_gb=policy_config.get("refit_buffer_size_gb"), refit_timeout_s=master_config.async_rl.generation_fleet_health.refit_timeout_s, ) + generation.weight_synchronizer = weight_synchronizer weight_synchronizer.init_communicator() setup_timing_metrics.collective_init_time_s = time.perf_counter() - t0 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..2844e95979b 100644 --- a/nemo_rl/data/multimodal_utils.py +++ b/nemo_rl/data/multimodal_utils.py @@ -19,9 +19,9 @@ import uuid from collections import defaultdict from collections.abc import Sequence -from concurrent.futures import ThreadPoolExecutor from copy import deepcopy from io import BytesIO +from pathlib import Path from typing import Any, Optional, Union import requests @@ -40,7 +40,6 @@ MULTIMODAL_CONTENT_TYPES = frozenset( {*IMAGE_CONTENT_TYPES, *VIDEO_CONTENT_TYPES, *AUDIO_CONTENT_TYPES} ) -NEMO_GYM_IMAGE_ENCODE_MAX_WORKERS = 8 # List of allowed placeholder strings for different media types in the dataset string # e.g. "This is an example of " @@ -51,6 +50,8 @@ "video-audio": "", } MEDIA_TAGS_REVERSED = {v: k for k, v in MEDIA_TAGS.items()} +CACHED_VIDEO_FRAME_MANIFEST_MAGIC = b"NEMO_RL_CACHED_VIDEO_FRAMES_V1\n" +CACHED_VIDEO_FRAME_MANIFEST_MIME = "video/x-nemo-rl-cached-frames" DEFAULT_MEDIA_EXTENSIONS = { "image": ["png", "jpeg", "jpg", "img"], @@ -669,6 +670,16 @@ def flattened_concat( ) +# Model inputs some remote-code processors omit from ``model_input_names`` even +# though their forward requires them. Keep extraction and TQ schema warmup aligned. +UNDECLARED_MULTIMODAL_MODEL_INPUTS = ( + "imgs_sizes", + "num_frames", + "pixel_values_flat", + "image_num_patches", +) + + def get_multimodal_keys_from_processor(processor) -> list[str]: """Get keys of the multimodal data that can be used as model inputs. @@ -792,12 +803,7 @@ def extract_multimodal_model_inputs( # TODO(rohitrango): Let ProcessorInterface declare model-specific media inputs. # Some remote-code processors omit these inputs from model_input_names even # though their model forward requires them. - for key in ( - "imgs_sizes", - "num_frames", - "pixel_values_flat", - "image_num_patches", - ): + for key in UNDECLARED_MULTIMODAL_MODEL_INPUTS: if key in processed and key not in multimodal_keys: multimodal_keys.append(key) for key in multimodal_keys: @@ -890,24 +896,25 @@ def image_to_data_url(image: Image.Image, fmt: str = "PNG") -> str: return f"data:image/{fmt.lower()};base64,{encoded}" -def _encode_single_image_source(source: str) -> str: - """Resolve and encode one image source.""" - image = resolve_to_image(source) - try: - data_url = image_to_data_url(image) - finally: - image.close() - return data_url +def get_responses_content_part_url(part: dict[str, Any], *keys: str) -> str: + """Return a string media source from a Responses/Chat content part.""" + for key in keys: + value = part.get(key) + if isinstance(value, dict): + value = value.get("url") or value.get("path") + if isinstance(value, str) and value: + return value + return "" -def extract_input_image_sources_from_responses_messages( +def extract_input_media_sources_from_responses_messages( messages: Any, -) -> list[str | Image.Image]: - """Extract image sources from Responses-API messages in encounter order.""" +) -> list[tuple[str, Any]]: + """Extract tagged image and video sources in encounter order.""" if not isinstance(messages, list): return [] - sources: list[str | Image.Image] = [] + sources: list[tuple[str, Any]] = [] for message in messages: if not isinstance(message, dict): continue @@ -917,24 +924,37 @@ def extract_input_image_sources_from_responses_messages( for part in content: if not isinstance(part, dict): continue - if part.get("type") not in ("input_image", "image", "image_url"): + part_type = part.get("type") + if part_type in IMAGE_CONTENT_TYPES: + media_type = "image" + source = part.get("image") or part.get("image_url") or part.get("url") + elif part_type in VIDEO_CONTENT_TYPES: + media_type = "video" + source = part.get("video") or part.get("video_url") or part.get("url") + else: continue - source = part.get("image") or part.get("image_url") or part.get("url") if isinstance(source, dict): - source = source.get("url") + source = source.get("url") or source.get("path") + # Skip non-str/non-Image sources: callers hand these straight to + # `resolve_to_image`, which would raise on e.g. an int `image_url`. if isinstance(source, (str, Image.Image)): - sources.append(source) + sources.append((media_type, source)) return sources -def extract_input_images_from_responses_messages( - messages: Any, -) -> list[Image.Image]: - """Load images from Responses-API input messages in encounter order.""" - return [ - resolve_to_image(source) - for source in extract_input_image_sources_from_responses_messages(messages) - ] +def media_sources_equal( + left: tuple[str, Any], + right: tuple[str, Any], +) -> bool: + """Compare tagged media by string value or object identity.""" + if left[0] != right[0]: + return False + left_source, right_source = left[1], right[1] + return ( + left_source == right_source + if isinstance(left_source, str) and isinstance(right_source, str) + else left_source is right_source + ) def _materialize_ragged_pixel_values( @@ -1042,82 +1062,41 @@ def attach_image_model_inputs_to_message( ) -def encode_images_in_examples(nemo_gym_examples: list[dict]) -> list[dict]: - """Replace local image paths in NeMo Gym examples with base64 data URLs. - - Walks each example's ``responses_create_params.input[].content[]`` items, - collects local image references, encodes each unique source once using a - bounded thread pool, and rewrites every corresponding image part with the - resulting base64 ``data:`` URL. Parts whose URL already starts with - ``http://``, ``https://``, or ``data:`` are left untouched. Malformed items - (non-dict entries, missing/empty URLs, non-list ``input``/``content``) are - skipped without raising. - - The examples are mutated in place; the same list is also returned for - convenience so callers can chain the call. +_VIDEO_EXT_TO_MIME = { + ".mp4": "mp4", + ".m4v": "mp4", + ".mov": "quicktime", + ".webm": "webm", + ".mkv": "x-matroska", + ".avi": "x-msvideo", +} - Args: - nemo_gym_examples: List of NeMo Gym example dicts. Each example is - expected to contain a ``responses_create_params`` mapping with an - ``input`` list of Responses API messages. - Returns: - The same ``nemo_gym_examples`` list, with local image references - rewritten to base64 data URLs in place. - """ - targets_by_source: dict[str, list[tuple[dict, str]]] = {} +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 - for example in nemo_gym_examples: - input_items = example.get("responses_create_params", {}).get("input", []) - if not isinstance(input_items, list): - continue - for item in input_items: - if not isinstance(item, dict): - continue - content = item.get("content", []) - if not isinstance(content, list): - continue - for part in content: - if ( - not isinstance(part, dict) - or part.get("type") not in IMAGE_CONTENT_TYPES - ): - continue - media_key = next( - (key for key in ("image_url", "image", "url") if key in part), - None, - ) - if media_key is None: - continue - url = part.get(media_key) - if isinstance(url, dict): - url = url.get("url") or url.get("path") or "" - if not isinstance(url, str) or not url: - continue - if url.startswith(("http://", "https://", "data:")): - continue - targets_by_source.setdefault(url, []).append((part, media_key)) - - sources = list(targets_by_source) - if sources: - with ThreadPoolExecutor( - max_workers=NEMO_GYM_IMAGE_ENCODE_MAX_WORKERS - ) as executor: - encoded_by_source = dict( - zip( - sources, - executor.map(_encode_single_image_source, sources), - strict=True, - ) - ) + resolved = ( + video_path.removeprefix("file://") + if video_path.startswith("file://") + else str(Path(video_path).expanduser().resolve()) + ) + path = Path(resolved) + if not path.is_file(): + raise FileNotFoundError( + f"Video path resolved to {resolved!r}, which does not exist." + ) - # Keep payload mutation on the caller thread after worker-owned images - # have been closed and every unique source has been encoded. - for source, targets in targets_by_source.items(): - data_url = encoded_by_source[source] - for part, media_key in targets: - part[media_key] = data_url - return nemo_gym_examples + ext = path.suffix.lower() + mime = _VIDEO_EXT_TO_MIME.get(ext) + if mime is None: + raise ValueError( + f"Unsupported video extension {ext!r} for {resolved!r}. " + f"Supported: {sorted(_VIDEO_EXT_TO_MIME)}." + ) + encoded = base64.b64encode(path.read_bytes()).decode("ascii") + return f"data:video/{mime};base64,{encoded}" def get_media_from_message(message: dict[str, Any]) -> dict[str, list[Any]]: diff --git a/nemo_rl/data/processors.py b/nemo_rl/data/processors.py index 58f09c95856..ef8744afb7b 100644 --- a/nemo_rl/data/processors.py +++ b/nemo_rl/data/processors.py @@ -490,6 +490,8 @@ def vlm_hf_data_processor( pass # Daily-Omni data is already formatted by DailyOmniDataset.format_data elif datum_dict["task_name"] in ("intent-train", "intent-bench"): pass # IntentDataset.format_data already produces the message structure + elif "messages" in datum_dict: + pass # Generic ResponseDataset data can already use the message structure else: raise ValueError(f"No data processor for task {datum_dict['task_name']}") @@ -676,8 +678,11 @@ def vlm_hf_data_processor( loss_multiplier = 0.0 else: # get the prompt content! (use this for vllm-backend that needs formatted dialog and list of images/audios) for the entire conversation + # Placeholder-style processors set vllm_content to None so vLLM uses expanded input_ids. vllm_kwargs = { - "vllm_content": string_formatted_dialog, + "vllm_content": ( + None if uses_placeholder and images else string_formatted_dialog + ), "vllm_images": images, "vllm_audios": audios, "vllm_videos": videos, @@ -799,7 +804,7 @@ def nemo_gym_data_processor( "Gym video data requires a multimodal processor with " "apply_chat_template and tokenizer attributes" ) - from nemo_rl.environments.nemo_gym_video import ( + from nemo_rl.environments.nemo_gym_multimodal import ( nemo_gym_example_to_video_datum_spec, ) diff --git a/nemo_rl/data_plane/codec.py b/nemo_rl/data_plane/codec.py index 919bf96a13d..045eed067c5 100644 --- a/nemo_rl/data_plane/codec.py +++ b/nemo_rl/data_plane/codec.py @@ -28,6 +28,8 @@ :func:`response_from_nested` to extract the response slice from a (prompt+response) nested tensor. +* Multimodal ``PackedTensor`` fields ride as row-jagged tensor payloads plus + compact reconstruction metadata. * Non-tensor object fields ride as ``NonTensorStack`` / ``NonTensorData`` leaves (TQ-native passthrough). :func:`materialize` decodes them back to ``np.ndarray(dtype=object)`` for the trainer. @@ -41,7 +43,7 @@ import torch from tensordict import TensorDict, TensorDictBase -from nemo_rl.data_plane.schema import Layout +from nemo_rl.data_plane.schema import Layout, PACKED_TENSOR_META_PREFIX if TYPE_CHECKING: # Type-only import. At runtime, BatchedDataDict is loaded lazily @@ -130,8 +132,134 @@ def unwrap_wire_stripped_payload(item: Any) -> Any: return item +def _pack_packed_tensor_field( + key: str, + value: Any, +) -> dict[str, torch.Tensor]: + """Encode a PackedTensor as a row-jagged tensor plus reconstruction metadata.""" + from nemo_rl.data.multimodal_utils import PackedTensor + + if not isinstance(value, PackedTensor): + raise TypeError(f"{key!r} is not a PackedTensor") + + logical_rows: list[torch.Tensor | None] = [] + for row_idx in range(len(value)): + row = value.slice([row_idx]).as_tensor() + if row is None: + logical_rows.append(None) + continue + pack_dim = value.dim_to_pack + if pack_dim < 0: + pack_dim += row.dim() + if not 0 <= pack_dim < row.dim(): + raise IndexError( + f"PackedTensor field {key!r} has dim_to_pack={value.dim_to_pack} " + f"for rank-{row.dim()} row" + ) + logical_rows.append(row.movedim(pack_dim, 0).detach()) + + nonempty = [row for row in logical_rows if row is not None] + if nonempty: + ranks = {row.dim() for row in nonempty} + dtypes = {row.dtype for row in nonempty} + devices = {row.device for row in nonempty} + if len(ranks) != 1 or len(dtypes) != 1 or len(devices) != 1: + raise ValueError( + f"PackedTensor field {key!r} must have one rank, dtype, and device " + "across logical rows" + ) + rank = nonempty[0].dim() + trailing_shape = tuple( + max(row.shape[dim] for row in nonempty) for dim in range(1, rank) + ) + if not value.pad_to_max_shape and any( + tuple(row.shape[1:]) != trailing_shape for row in nonempty + ): + raise ValueError( + f"PackedTensor field {key!r} has mismatched non-packing dimensions " + "without pad_to_max_shape" + ) + + canonical_rows: list[torch.Tensor] = [] + for row in logical_rows: + if row is None: + canonical_rows.append( + nonempty[0].new_empty((0, *trailing_shape)) + ) + continue + if tuple(row.shape[1:]) != trailing_shape: + padding: list[int] = [] + for dim in reversed(range(row.dim())): + padding.extend( + (0, 0 if dim == 0 else trailing_shape[dim - 1] - row.shape[dim]) + ) + row = torch.nn.functional.pad(row, padding) + canonical_rows.append(row.contiguous()) + payload = stack_or_nest(canonical_rows) + else: + payload = torch.empty((len(value), 0)) + + lengths = torch.tensor( + [0 if row is None else row.shape[0] for row in logical_rows], + dtype=torch.long, + ) + metadata = torch.stack( + ( + lengths, + torch.full_like(lengths, value.dim_to_pack), + torch.full_like(lengths, int(value.pad_to_max_shape)), + ), + dim=1, + ) + return { + key: payload, + f"{PACKED_TENSOR_META_PREFIX}{key}": metadata, + } + + +def _unpack_packed_tensor_field( + payload: torch.Tensor, + metadata: torch.Tensor, +) -> Any: + """Reconstruct a PackedTensor encoded by :func:`_pack_packed_tensor_field`.""" + from nemo_rl.data.multimodal_utils import PackedTensor + + if metadata.dim() != 2 or metadata.shape[1] != 3: + raise ValueError( + "PackedTensor metadata must have shape [batch, 3], got " + f"{tuple(metadata.shape)}" + ) + lengths = metadata[:, 0].tolist() + dims = metadata[:, 1].tolist() + pad_flags = metadata[:, 2].tolist() + if len(set(dims)) != 1 or len(set(pad_flags)) != 1: + raise ValueError("PackedTensor reconstruction settings must be constant per field") + + wire_rows = list(payload.unbind()) if payload.is_nested else list(payload.unbind(0)) + if len(wire_rows) != len(lengths): + raise ValueError( + f"PackedTensor payload has {len(wire_rows)} rows but metadata has " + f"{len(lengths)}" + ) + + dim_to_pack = int(dims[0]) if dims else 0 + rows: list[torch.Tensor | None] = [] + for row, length in zip(wire_rows, lengths, strict=True): + length = int(length) + if length == 0: + rows.append(None) + continue + canonical = row[:length] + rows.append(canonical.movedim(0, dim_to_pack).contiguous()) + return PackedTensor( + rows, + dim_to_pack=dim_to_pack, + pad_to_max_shape=bool(pad_flags[0]) if pad_flags else False, + ) + + def pack_jagged_fields( - fields: "dict[str, torch.Tensor | np.ndarray]", + fields: "dict[str, Any]", *, lengths: torch.Tensor | None, token_aligned_fields: set[str] | frozenset[str] | None = None, @@ -147,8 +275,8 @@ def pack_jagged_fields( of truth for both :func:`kv_first_write` and :func:`write_columns`. Args: - fields: Column name → tensor or object array. Other value types - raise ``TypeError``. + fields: Column name → tensor, PackedTensor, or object array. Other + value types raise ``TypeError``. lengths: Per-row valid lengths used by :func:`pack_per_token_field`. ``None`` disables jagged conversion entirely. token_aligned_fields: Field names known to be per-token. These use @@ -161,9 +289,17 @@ def pack_jagged_fields( """ n = int(lengths.shape[0]) if lengths is not None else 0 token_aligned_fields = token_aligned_fields or frozenset() + from nemo_rl.data.multimodal_utils import PackedTensor + packed: dict[str, Any] = {} for k, v in fields.items(): - if isinstance(v, np.ndarray) and v.dtype == object: + if isinstance(v, PackedTensor): + if len(v) != n: + raise ValueError( + f"PackedTensor field {k!r} has {len(v)} rows, expected {n}" + ) + packed.update(_pack_packed_tensor_field(k, v)) + elif isinstance(v, np.ndarray) and v.dtype == object: # tensordict==0.12.2 wire bug: a NonTensorStack stored as a # TensorDict leaf returns as a LinkedList on parent # __getitem__, losing identity. ndarray(dtype=object) @@ -177,7 +313,7 @@ def pack_jagged_fields( else: raise TypeError( f"pack_jagged_fields: unsupported value type for {k!r}: {type(v)}. " - "Use torch.Tensor or np.ndarray(dtype=object)." + "Use torch.Tensor, PackedTensor, or np.ndarray(dtype=object)." ) return TensorDict(packed, batch_size=[n]) @@ -290,8 +426,31 @@ def materialize( pads = pad_value_dict or {} out: dict[str, Any] = {} + available_keys = set(td.keys(include_nested=False)) + for key in available_keys: + if not key.startswith(PACKED_TENSOR_META_PREFIX): + continue + payload_key = key.removeprefix(PACKED_TENSOR_META_PREFIX) + if payload_key not in available_keys: + raise ValueError( + f"PackedTensor metadata field {key!r} has no payload field " + f"{payload_key!r}" + ) # pyrefly: inference cycle on tensordict.items() loop var. for key, val in td.items(include_nested=False): # type: ignore[bad-assignment] + if key.startswith(PACKED_TENSOR_META_PREFIX): + continue + packed_meta_key = f"{PACKED_TENSOR_META_PREFIX}{key}" + if packed_meta_key in available_keys: + metadata = td[packed_meta_key] + if not isinstance(val, torch.Tensor) or not isinstance( + metadata, torch.Tensor + ): + raise TypeError( + f"PackedTensor wire fields for {key!r} must both be tensors" + ) + out[key] = _unpack_packed_tensor_field(val, metadata) + continue if isinstance(val, NonTensorStack): # ``np.asarray(list, dtype=object)`` would probe each item's # ``__iter__`` to detect a nested array. A wire-stripped TD @@ -328,6 +487,7 @@ def materialize( if ( pad_to_seqlen > 0 and isinstance(padded, torch.Tensor) + and not padded.is_nested and padded.dim() >= 2 and padded.shape[1] < pad_to_seqlen ): diff --git a/nemo_rl/data_plane/column_io.py b/nemo_rl/data_plane/column_io.py index 9bbdb5616c0..fa04685e70c 100644 --- a/nemo_rl/data_plane/column_io.py +++ b/nemo_rl/data_plane/column_io.py @@ -53,6 +53,8 @@ "token_mask", "sample_mask", "routed_experts", + "token_type_ids", + "mm_token_type_ids", } ) @@ -109,7 +111,7 @@ def read_columns( def write_columns( dp_client: DataPlaneClient, meta: KVBatchMeta, - fields: "dict[str, torch.Tensor | np.ndarray]", + fields: "dict[str, Any]", ) -> None: """``put_samples(meta.sample_ids, fields=...)``. @@ -190,10 +192,13 @@ def kv_first_write( f"kv_first_write: tags ({len(tags)}) must match batch size ({n})" ) lengths = final_batch_cpu["input_lengths"] - fields: dict[str, torch.Tensor | np.ndarray] = { + from nemo_rl.data.multimodal_utils import PackedTensor + + fields: dict[str, Any] = { k: v for k, v in final_batch_cpu.items() if isinstance(v, torch.Tensor) + or isinstance(v, PackedTensor) or (isinstance(v, np.ndarray) and v.dtype == object) } td = pack_jagged_fields( diff --git a/nemo_rl/data_plane/interfaces.py b/nemo_rl/data_plane/interfaces.py index fcff2e398f4..8031df22dd9 100644 --- a/nemo_rl/data_plane/interfaces.py +++ b/nemo_rl/data_plane/interfaces.py @@ -296,7 +296,7 @@ def slice(self, start: int, stop: int) -> "KVBatchMeta": ) def concat(self, *others: "KVBatchMeta") -> "KVBatchMeta": - """Append ``others`` to ``self``. All metas must share ``partition_id``.""" + """Append ``others`` and union their fields in first-seen order.""" if any(o.partition_id != self.partition_id for o in others): raise ValueError("KVBatchMeta.concat: partition_ids must match") all_m = (self, *others) @@ -309,9 +309,14 @@ def concat(self, *others: "KVBatchMeta") -> "KVBatchMeta": ) all_have_tags = all(m.tags is not None for m in all_m) tags = [t for m in all_m for t in (m.tags or [])] if all_have_tags else None - return self._replace( + merged_fields = list( + dict.fromkeys(field for meta in all_m for field in (meta.fields or [])) + ) + result = self._replace( sample_ids=sample_ids, sequence_lengths=seq_lens, tags=tags ) + result.fields = merged_fields or None + return result def drop(self, indices: "Sequence[int]") -> "KVBatchMeta | None": """Complement of :meth:`subset`. Returns ``None`` when all rows are dropped.""" diff --git a/nemo_rl/data_plane/schema.py b/nemo_rl/data_plane/schema.py index 9c1f9d1e2ef..67ea1341c77 100644 --- a/nemo_rl/data_plane/schema.py +++ b/nemo_rl/data_plane/schema.py @@ -18,6 +18,11 @@ # Materialization layout for `codec.materialize` / `read_columns` / worker fetch. Layout = Literal["padded", "jagged"] +# Companion column used to carry the row lengths and reconstruction settings for +# a PackedTensor field. The payload itself keeps the original model-input name. +PACKED_TENSOR_META_PREFIX = "__nrl_packed_tensor_meta__" +MULTIMODAL_AUXILIARY_FIELDS = ("token_type_ids", "mm_token_type_ids") + # Per-shard packing metadata keys in `KVBatchMeta.extra_info`. MICRO_BATCH_INDICES = "micro_batch_indices" MICRO_BATCH_LENGTHS = "micro_batch_lengths" @@ -127,3 +132,46 @@ def fields_with_optional_routed_experts( if enabled and ROUTED_EXPERTS_FIELD not in out: out.append(ROUTED_EXPERTS_FIELD) return out + + +def fields_with_packed_tensor_payload( + fields: Sequence[str], + available_fields: Sequence[str] | None, +) -> list[str]: + """Include PackedTensor payload/metadata columns present in a TQ record. + + Use this when narrowing an existing record for a consumer that can use + media, such as the policy path. For partition schema registration, use + :func:`packed_tensor_wire_fields` instead. Text-only consumers should not + request media columns. + """ + out = list(fields) + if not available_fields: + return out + available = set(available_fields) + for field in MULTIMODAL_AUXILIARY_FIELDS: + if field in available and field not in out: + out.append(field) + for field in available_fields: + if not field.startswith(PACKED_TENSOR_META_PREFIX): + continue + payload_field = field.removeprefix(PACKED_TENSOR_META_PREFIX) + if payload_field not in available: + continue + if payload_field not in out: + out.append(payload_field) + if field not in out: + out.append(field) + return out + + +def packed_tensor_wire_fields(field_names: Sequence[str]) -> list[str]: + """Return payload and companion-column names to pre-register for model inputs.""" + out: list[str] = [] + for field in field_names: + if field not in out: + out.append(field) + meta_field = f"{PACKED_TENSOR_META_PREFIX}{field}" + if meta_field not in out: + out.append(meta_field) + return out diff --git a/nemo_rl/data_plane/worker_mixin.py b/nemo_rl/data_plane/worker_mixin.py index e545ae48370..2762cec2b6b 100644 --- a/nemo_rl/data_plane/worker_mixin.py +++ b/nemo_rl/data_plane/worker_mixin.py @@ -35,6 +35,7 @@ FetchPolicy = Literal["auto", "independent", "leader_broadcast"] from nemo_rl.data.llm_message_utils import attach_message_log_view +from nemo_rl.data.multimodal_utils import PackedTensor from nemo_rl.data_plane.schema import ( ELEM_COUNTS_PER_GB, GLOBAL_FORWARD_PAD_SEQLEN, @@ -80,6 +81,28 @@ def _broadcast_batched_data_dict( descriptor.append( (k, "tensor", str(v.dtype), tuple(v.shape), str(v.device)) ) + elif isinstance(v, PackedTensor): + descriptor.append( + ( + k, + "packed_tensor", + v.dim_to_pack, + v.pad_to_max_shape, + v._row_offsets, + v._segment_indices, + v._segment_provenance, + [ + None + if tensor is None + else ( + str(tensor.dtype), + tuple(tensor.shape), + str(tensor.device), + ) + for tensor in v.tensors + ], + ) + ) else: descriptor.append((k, "raw", v)) payload: list[Any] = [descriptor] @@ -114,6 +137,54 @@ def _broadcast_batched_data_dict( and torch.device(src_device).type != torch.device(bcast_device).type ): out[key] = tensor.to(src_device) + elif kind == "packed_tensor": + ( + _, + _, + dim_to_pack, + pad_to_max_shape, + row_offsets, + segment_indices, + segment_provenance, + tensor_descriptors, + ) = entry + if is_leader: + packed_value = out[key] + tensors = packed_value.tensors + else: + tensors = [ + None + if tensor_descriptor is None + else torch.empty( + tensor_descriptor[1], + dtype=getattr(torch, tensor_descriptor[0].split(".")[-1]), + device=bcast_device, + ) + for tensor_descriptor in tensor_descriptors + ] + for index, tensor_descriptor in enumerate(tensor_descriptors): + if tensor_descriptor is None: + continue + tensor = tensors[index] + assert tensor is not None + src_device = tensor_descriptor[2] + if tensor.device.type != torch.device(bcast_device).type: + tensor = tensor.to(bcast_device) + torch.distributed.broadcast(tensor, src=src, group=group) + if ( + not is_leader + and torch.device(src_device).type != torch.device(bcast_device).type + ): + tensor = tensor.to(src_device) + tensors[index] = tensor + out[key] = PackedTensor( + tensors, + dim_to_pack=dim_to_pack, + pad_to_max_shape=pad_to_max_shape, + _row_offsets=row_offsets, + _segment_indices=segment_indices, + _segment_provenance=segment_provenance, + ) else: if not is_leader: out[key] = entry[2] @@ -135,12 +206,6 @@ def setup_data_plane(self, cfg: DataPlaneConfig) -> None: Called once by the driver after worker construction. Idempotent. """ - if getattr(self, "model_slices_context_parallel_inputs", False): - raise NotImplementedError( - "TransferQueue/SingleController does not yet support models that " - "insert media before context-parallel input selection. Use the " - "synchronous NeMo-RL policy path for Nemotron Omni." - ) if self._dp_client is not None: return from nemo_rl.data_plane import build_data_plane_client diff --git a/nemo_rl/distributed/worker_groups.py b/nemo_rl/distributed/worker_groups.py index 7d88794f549..3cfd2d405ca 100644 --- a/nemo_rl/distributed/worker_groups.py +++ b/nemo_rl/distributed/worker_groups.py @@ -32,10 +32,22 @@ from nemo_rl.distributed.virtual_cluster import RayVirtualCluster from nemo_rl.distributed.worker_group_utils import recursive_merge_options from nemo_rl.utils.venvs import ( + add_hf_modules_cache_to_pythonpath, create_local_venv_on_each_node, ) +def _get_initializer_env_vars(env_vars: dict[str, str]) -> dict[str, str]: + """Build the environment needed to unpickle worker constructor arguments.""" + initializer_env_vars = { + key: env_vars[key] if key in env_vars else os.environ[key] + for key in ("HF_HOME", "HF_MODULES_CACHE", "PYTHONPATH") + if key in env_vars or key in os.environ + } + + return add_hf_modules_cache_to_pythonpath(initializer_env_vars) + + @dataclass class MultiWorkerFuture: """Container for Ray futures with associated worker information.""" @@ -502,11 +514,7 @@ def _create_workers_from_bundle_indices( # import-related variables that trust_remote_code classes need to # resolve their generated modules have to travel with it. unique_pg_indices = sorted({pg_idx for pg_idx, _ in bundle_indices_list}) - initializer_env_vars = { - key: env_vars[key] - for key in ("HF_HOME", "HF_MODULES_CACHE", "PYTHONPATH") - if key in env_vars - } + initializer_env_vars = _get_initializer_env_vars(env_vars) initializer_runtime_env = {} if py_executable != sys.executable: initializer_runtime_env["py_executable"] = py_executable diff --git a/nemo_rl/environments/nemo_gym.py b/nemo_rl/environments/nemo_gym.py index e702f462042..0f68069db92 100644 --- a/nemo_rl/environments/nemo_gym.py +++ b/nemo_rl/environments/nemo_gym.py @@ -17,21 +17,18 @@ import sys from collections import Counter from collections.abc import AsyncGenerator, Mapping -from copy import deepcopy from pathlib import Path from typing import Any, Dict, List, NotRequired, Optional, Protocol, TypedDict import ray import torch -from PIL import Image from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy from transformers import PreTrainedTokenizerBase from nemo_rl.data.multimodal_utils import ( attach_image_model_inputs_to_message, - encode_images_in_examples, - extract_input_image_sources_from_responses_messages, - resolve_to_image, + extract_input_media_sources_from_responses_messages, + media_sources_equal, uses_image_placeholder, ) from nemo_rl.distributed.ray_actor_environment_registry import get_actor_python_env @@ -42,6 +39,12 @@ _get_node_ip_local, ) from nemo_rl.environments.interfaces import EnvironmentInterface +from nemo_rl.environments.nemo_gym_multimodal import ( + _index_per_turn_images, + _is_trainable_output_item, + _without_initial_media_sources, + normalize_media_in_examples, +) from nemo_rl.experience.failures import ( GymTransportError, RolloutDataFailure, @@ -282,28 +285,6 @@ def _detect_invalid_tool_call_and_malformed_thinking( return is_invalid_tool_call, has_malformed_thinking -######################################## -# Multimodal helpers -######################################## - - -# WARNING: A function-call output beginning with HTTP(S) is accepted here and -# passed to ``resolve_to_image``, which performs an outbound request during -# postprocessing even when the tool result is not actually an image. -_IMAGE_SRC_PREFIXES = ("data:image/", "http://", "https://", "file://") - - -def _looks_like_image_src(src: str) -> bool: - """True when ``src`` plausibly points at an image the loader can open. - - Guards against tool responses (e.g. ``{"x": 0.65, "y": 0.83}`` from a - click tool) that are strings but not image URLs. Without this, the - indexer forwards the JSON payload to ``resolve_to_image`` → PIL.open, - which treats it as a filesystem path and raises ``FileNotFoundError``. - """ - return src.startswith(_IMAGE_SRC_PREFIXES) - - def get_pad_dynamic_image_shapes(env_config: Mapping[str, Any]) -> bool: """Return nemo_gym's pad_dynamic_image_shapes from an env config, or False. @@ -322,157 +303,6 @@ def get_pad_dynamic_image_shapes(env_config: Mapping[str, Any]) -> bool: return bool(nemo_gym_config.get("pad_dynamic_image_shapes")) -def _extract_input_images_from_message(item: dict) -> list[Image.Image]: - """Pull PIL images out of a non-assistant Responses-API item. - - Handles both content-list items (user / tool messages carrying - ``input_image``/``image``/``image_url`` parts) and ``function_call_output`` - items whose ``output`` field is an image data URL. Tool outputs that are - non-image strings (e.g. structured JSON returned by tools like - ``click(x, y)``) contribute zero images to the bucket. - """ - images: list[Image.Image] = [] - if item.get("type") == "function_call_output": - src = item.get("output") - if isinstance(src, str) and _looks_like_image_src(src): - images.append(resolve_to_image(src)) - return images - content = item.get("content") or [] - if not isinstance(content, list): - return images - for part in content: - if not isinstance(part, dict): - continue - if part.get("type") not in ("input_image", "image", "image_url"): - continue - src = part.get("image") or part.get("image_url") or part.get("url") - if src is None: - continue - if isinstance(src, dict): - src = src.get("url") - if src is None: - continue - images.append(resolve_to_image(src)) - return images - - -def _is_trainable_output_item(item: dict) -> bool: - """Report whether an output item becomes a trainable assistant turn. - - The postprocess loop skips items whose ``generation_token_ids`` is missing - *or* empty, so per-turn image binning has to use the same predicate or the - two walks disagree and every later turn gets the wrong images. - """ - return bool(item.get("generation_token_ids")) - - -def _index_per_turn_images( - output: list[dict], - input_messages: list[dict] | None = None, -) -> list[list[Image.Image]]: - """Bin server-returned images by the trainable turn that saw them. - - Walks the Responses-API items in order and flushes ``pending`` into a - per-turn bucket each time it hits an item carrying truthy - ``generation_token_ids`` — matching the exact gate that - ``_postprocess_nemo_gym_to_nemo_rl_result`` uses to decide which items - become trainable turns. Every other item (user turns, tool messages, - ``function_call_output``, non-trainable reasoning) contributes its images - to ``pending`` for the next trainable turn. This ensures the returned list - has one entry per trainable turn, aligned with the postprocess loop's - ``turn_idx`` even when the trainable item's role is not ``assistant`` - (e.g. a reasoning-only response, or a ``function_call``). - - ``input_messages`` is the initial ``responses_create_params.input`` list — - images there (e.g. a single-shot user prompt for tool-based envs like - circle-click) are consumed by the first trainable turn's tokenized prompt - and must land in the first bucket. Agents like ``gym_v_agent`` that keep - ``input`` empty and inject observations as ``function_call_output`` items - are unaffected — the seed is a no-op when ``input_messages`` is empty. - """ - per_turn: list[list[Image.Image]] = [] - pending: list[Image.Image] = [] - for item in input_messages or (): - if isinstance(item, dict) and item.get("role") != "assistant": - pending.extend(_extract_input_images_from_message(item)) - for item in output: - if item.get( - "generation_token_ids" - ): # trainable turn; empty generation_token_ids is skipped by the postprocess loop and must not consume a bucket - per_turn.append(pending) - pending = [] - elif item.get("role") != "assistant": - pending.extend(_extract_input_images_from_message(item)) - return per_turn - - -def _image_sources_equal(left: Any, right: Any) -> bool: - return ( - left == right - if isinstance(left, str) and isinstance(right, str) - else left is right - ) - - -def _without_initial_image_sources( - messages: Any, initial_sources: list[Any] -) -> tuple[Any, bool]: - """Copy Responses messages and remove one ordered copy of initial images.""" - if not isinstance(messages, list): - return messages, False - - filtered = deepcopy(messages) - remaining_sources = list(initial_sources) - for message in filtered: - if not isinstance(message, dict): - continue - content = message.get("content") - if not isinstance(content, list): - continue - - filtered_content = [] - for part in content: - part_sources = extract_input_image_sources_from_responses_messages( - [{"content": [part]}] - ) - if ( - remaining_sources - and len(part_sources) == 1 - and _image_sources_equal(part_sources[0], remaining_sources[0]) - ): - remaining_sources.pop(0) - continue - filtered_content.append(part) - message["content"] = filtered_content - - return filtered, not remaining_sources - - -def _attach_multimodal_data_to_user_message( - user_message: dict, - *, - images: list[Image.Image], - processor: Any, - pad_dynamic_image_shapes: bool = False, -) -> None: - """Attach per-turn multimodal tensors to ``user_message``. - - The processor is only invoked to extract multimodal tensors (pixel_values, - imgs_sizes, num_patches, etc.); its text output is discarded — vLLM's - tokens remain the trajectory. We therefore feed it the minimal placeholder - text it needs to count image regions: one ``processor.image_token`` per - image. Passing the vLLM-decoded text does not work because that text - already contains expanded ``...*N...`` regions, and the - processor would try to re-expand every embedded ````. - """ - attach_image_model_inputs_to_message( - user_message, - images=images, - processor=processor, - pad_dynamic_image_shapes=pad_dynamic_image_shapes, - ) - - @ray.remote(max_restarts=-1, max_task_retries=-1) # pragma: no cover class NemoGym(EnvironmentInterface): """This environment class isn't really used for training. It's really meant as an integration wrapper around NeMo-Gym that hooks into the existing NeMo RL resource management via ray. So there is still one source of truth for resource management in NeMo RL.""" @@ -500,7 +330,7 @@ def __init__(self, cfg: NemoGymConfig): from nemo_rl.algorithms.utils import get_tokenizer self._processor = get_tokenizer(tokenizer_config, get_processor=True) - # _attach_multimodal_data_to_user_message assumes a placeholder-style + # attach_image_model_inputs_to_message assumes a placeholder-style # processor (imgs_sizes / num_frames reconstruction + pad_to_max_shape # PackedTensor build). A non-placeholder VLM would silently produce # wrong multimodal tensors — fail at actor construction instead. @@ -508,7 +338,7 @@ def __init__(self, cfg: NemoGymConfig): "NemoGym multimodal path assumes a placeholder-style processor " "(see _PLACEHOLDER_STYLE_PROCESSOR_NAMES in nemo_rl/data/multimodal_utils.py); " f"got {type(self._processor).__name__}. Update " - "_attach_multimodal_data_to_user_message before enabling." + "attach_image_model_inputs_to_message before enabling." ) def _require_spinup(self) -> None: @@ -675,14 +505,10 @@ async def run_rollouts( timer = Timer() counts_left = Counter(row["agent_ref"]["name"] for row in nemo_gym_examples) - from nemo_rl.environments.nemo_gym_video import ( - normalize_video_urls_in_examples, - ) - - # Normalize local media before shipping requests to vLLM. Both helpers - # are no-ops for text-only rows and already-qualified URLs. - normalize_video_urls_in_examples(nemo_gym_examples) - encode_images_in_examples(nemo_gym_examples) + # Normalize local media before shipping requests to vLLM. Helper is a no-op + # for text-only rows and already-qualified URLs. + # Megatron's HTTP backend consumes the same normalized Responses payload. + normalize_media_in_examples(nemo_gym_examples) timer.start("_run_rollouts_total") nemo_gym_result_iterator = self.rch.run_examples( @@ -778,20 +604,20 @@ def _postprocess_nemo_gym_to_nemo_rl_result( media_messages = ( seed_obs if isinstance(seed_obs, list) and seed_obs else initial_input ) - raw_initial_sources = extract_input_image_sources_from_responses_messages( + raw_initial_sources = extract_input_media_sources_from_responses_messages( raw_input ) - agent_initial_sources = extract_input_image_sources_from_responses_messages( + agent_initial_sources = extract_input_media_sources_from_responses_messages( initial_input ) - returned_media_sources = extract_input_image_sources_from_responses_messages( + returned_media_sources = extract_input_media_sources_from_responses_messages( media_messages ) initial_media_matches_raw_input = ( bool(raw_initial_sources) and len(agent_initial_sources) == len(raw_initial_sources) and all( - _image_sources_equal(agent_source, raw_source) + media_sources_equal(agent_source, raw_source) for agent_source, raw_source in zip( agent_initial_sources, raw_initial_sources ) @@ -800,7 +626,7 @@ def _postprocess_nemo_gym_to_nemo_rl_result( returned_media_matches_raw_input = len(returned_media_sources) == len( raw_initial_sources ) and all( - _image_sources_equal(returned_source, raw_source) + media_sources_equal(returned_source, raw_source) for returned_source, raw_source in zip( returned_media_sources, raw_initial_sources ) @@ -811,7 +637,7 @@ def _postprocess_nemo_gym_to_nemo_rl_result( and returned_media_matches_raw_input ) if initial_multimodal_data_omitted: - media_messages, _ = _without_initial_image_sources( + media_messages, _ = _without_initial_media_sources( media_messages, raw_initial_sources ) per_turn_images = ( @@ -908,7 +734,7 @@ def _postprocess_nemo_gym_to_nemo_rl_result( images_this_turn = ( per_turn_images[turn_idx] if turn_idx < len(per_turn_images) else [] ) - _attach_multimodal_data_to_user_message( + attach_image_model_inputs_to_message( user_message, images=images_this_turn, processor=processor, @@ -1003,7 +829,7 @@ def _postprocess_nemo_gym_to_nemo_rl_result( (response, "seed_obs"), ): if key in container: - container[key], _ = _without_initial_image_sources( + container[key], _ = _without_initial_media_sources( container[key], raw_initial_sources ) @@ -1112,9 +938,16 @@ def validate_reward_components_match_scalar(nemo_gym_results: List[dict]) -> Non def setup_nemo_gym_config(config, tokenizer) -> None: generation_config = config.policy["generation"] - # Enable the http server. Requires both async engine and the expose_http_server flag - generation_config["vllm_cfg"]["async_engine"] = True - generation_config["vllm_cfg"]["expose_http_server"] = True + backend = generation_config.get("backend") + if backend == "vllm": + # Enable the http server. Requires both async engine and the expose_http_server flag + generation_config["vllm_cfg"]["async_engine"] = True + generation_config["vllm_cfg"]["expose_http_server"] = True + elif backend == "megatron": + # Enable the http server for Gym dispatch over the Megatron generation backend. + generation_config["mcore_generation_config"]["expose_http_server"] = True + else: + raise ValueError(f"NeMo Gym does not support generation backend {backend!r}.") # Stop strings or token ids are not supported generation_config["stop_strings"] = None diff --git a/nemo_rl/environments/nemo_gym_video.py b/nemo_rl/environments/nemo_gym_multimodal.py similarity index 67% rename from nemo_rl/environments/nemo_gym_video.py rename to nemo_rl/environments/nemo_gym_multimodal.py index 83d3f8bd9ea..3e65b88ba69 100644 --- a/nemo_rl/environments/nemo_gym_video.py +++ b/nemo_rl/environments/nemo_gym_multimodal.py @@ -15,6 +15,9 @@ import copy import json import os +from concurrent.futures import ThreadPoolExecutor +from contextlib import closing +from copy import deepcopy from pathlib import Path from typing import Any, TypeVar, cast from urllib.parse import unquote, urlparse @@ -28,9 +31,19 @@ IMAGE_CONTENT_TYPES, VIDEO_CONTENT_TYPES, PackedTensor, + extract_input_media_sources_from_responses_messages, extract_multimodal_model_inputs, get_dim_to_pack_along, + get_responses_content_part_url, + image_to_data_url, + media_sources_equal, resolve_to_image, + video_path_to_data_url, +) +from nemo_rl.environments.nemo_gym_request import ( + _chat_template_kwargs_for_processor, + _deep_merge_dict, + _json_mapping, ) from nemo_rl.environments.nemotron_utils import ( NEMOTRON_VIDEO_PROCESSOR_NAMES, @@ -43,6 +56,265 @@ load_video_frames_with_metadata, ) +_NEMO_GYM_IMAGE_ENCODE_MAX_WORKERS = 8 + + +def _encode_single_image_source(source: str) -> str: + """Resolve, encode, and close one local image source.""" + # `closing` (not a bare `with`): PIL's Image.__exit__ is a no-op, so only an + # explicit close() releases the buffer. + with closing(resolve_to_image(source)) as image: + return image_to_data_url(image) + + +def normalize_media_in_examples(nemo_gym_examples: list[dict]) -> list[dict]: + """Replace local media paths in NeMo Gym examples with data URLs.""" + local_image_sources: dict[str, None] = {} + local_video_sources: dict[str, None] = {} + pending_mutations: list[ + tuple[dict, tuple[str, str, str], str, str, bool, Any, str] + ] = [] + for example in nemo_gym_examples: + input_items = example.get("responses_create_params", {}).get("input", []) + if not isinstance(input_items, list): + continue + for item in input_items: + if not isinstance(item, dict): + continue + content = item.get("content", []) + if not isinstance(content, list): + continue + for part in content: + if not isinstance(part, dict): + continue + part_type = part.get("type") + if part_type in IMAGE_CONTENT_TYPES: + source_keys = ("image_url", "image", "url") + canonical_type = "input_image" + canonical_key = "image_url" + is_image = True + elif part_type in VIDEO_CONTENT_TYPES: + source_keys = ("video_url", "video", "url") + canonical_type = "input_video" + canonical_key = "video_url" + is_image = False + else: + continue + + present_keys = [key for key in source_keys if key in part] + if ( + not present_keys + and part_type == "input_image" + and "file_id" in part + ): + continue + if len(present_keys) != 1: + raise ValueError( + f"{part_type} requires exactly one of {source_keys}" + ) + + source = part[present_keys[0]] + nested_detail = ( + source.get("detail") if isinstance(source, dict) else None + ) + url = ( + source.get("url") or source.get("path", "") + if isinstance(source, dict) + else source + ) + if not isinstance(url, str) or not url: + raise ValueError(f"{part_type} requires a non-empty media URL") + if not url.startswith(("http://", "https://", "data:")): + if is_image: + local_image_sources.setdefault(url, None) + else: + local_video_sources.setdefault(url, None) + + pending_mutations.append( + ( + part, + source_keys, + canonical_type, + canonical_key, + is_image, + nested_detail, + url, + ) + ) + + sources = list(local_image_sources) + encoded_by_source: dict[str, str] = {} + if sources: + with ThreadPoolExecutor( + max_workers=_NEMO_GYM_IMAGE_ENCODE_MAX_WORKERS + ) as executor: + encoded_by_source = dict( + zip( + sources, + executor.map(_encode_single_image_source, sources), + strict=True, + ) + ) + + # Encode each unique video once. A video shared by G generations then points + # every part at the same string, instead of G separate base64 copies of the + # same file. Kept sequential: these payloads are large enough that encoding + # several at once would spike driver memory. + encoded_video_by_source: dict[str, str] = { + source: video_path_to_data_url(source) for source in local_video_sources + } + + # Apply mutations only after every local source was encoded successfully. + for ( + part, + source_keys, + canonical_type, + canonical_key, + is_image, + nested_detail, + url, + ) in pending_mutations: + for key in source_keys: + if key != canonical_key: + part.pop(key, None) + part["type"] = canonical_type + encoded = encoded_by_source if is_image else encoded_video_by_source + part[canonical_key] = encoded.get(url, url) + if is_image and nested_detail is not None: + part.setdefault("detail", nested_detail) + return nemo_gym_examples + + +# WARNING: A function-call output beginning with HTTP(S) is accepted here and +# passed to ``resolve_to_image``, which performs an outbound request during +# postprocessing even when the tool result is not actually an image. +_IMAGE_SRC_PREFIXES = ("data:image/", "http://", "https://", "file://") + + +def _looks_like_image_src(src: str) -> bool: + """True when ``src`` plausibly points at an image the loader can open. + + Guards against tool responses (e.g. ``{"x": 0.65, "y": 0.83}`` from a + click tool) that are strings but not image URLs. Without this, the + indexer forwards the JSON payload to ``resolve_to_image`` → PIL.open, + which treats it as a filesystem path and raises ``FileNotFoundError``. + """ + return src.startswith(_IMAGE_SRC_PREFIXES) + + +def _extract_input_images_from_message(item: dict) -> list[Image.Image]: + """Pull PIL images out of a non-assistant Responses-API item. + + Handles both content-list items (user / tool messages carrying + ``input_image``/``image``/``image_url`` parts) and ``function_call_output`` + items whose ``output`` field is an image data URL. Tool outputs that are + non-image strings (e.g. structured JSON returned by tools like + ``click(x, y)``) contribute zero images to the bucket. + """ + images: list[Image.Image] = [] + if item.get("type") == "function_call_output": + src = item.get("output") + if isinstance(src, str) and _looks_like_image_src(src): + images.append(resolve_to_image(src)) + return images + content = item.get("content") or [] + if not isinstance(content, list): + return images + for part in content: + if not isinstance(part, dict): + continue + if part.get("type") not in ("input_image", "image", "image_url"): + continue + src = part.get("image") or part.get("image_url") or part.get("url") + if isinstance(src, dict): + src = src.get("url") + if src is not None: + images.append(resolve_to_image(src)) + return images + + +def _is_trainable_output_item(item: dict) -> bool: + """Report whether an output item becomes a trainable assistant turn. + + The postprocess loop skips items whose ``generation_token_ids`` is missing + *or* empty, so per-turn image binning has to use the same predicate or the + two walks disagree and every later turn gets the wrong images. + """ + return bool(item.get("generation_token_ids")) + + +def _index_per_turn_images( + output: list[dict], + input_messages: list[dict] | None = None, +) -> list[list[Image.Image]]: + """Bin server-returned images by the trainable turn that saw them. + + Walks the Responses-API items in order and flushes ``pending`` into a + per-turn bucket each time it hits an item carrying truthy + ``generation_token_ids`` — matching the exact gate that + ``_postprocess_nemo_gym_to_nemo_rl_result`` uses to decide which items + become trainable turns. Every other item (user turns, tool messages, + ``function_call_output``, non-trainable reasoning) contributes its images + to ``pending`` for the next trainable turn. This ensures the returned list + has one entry per trainable turn, aligned with the postprocess loop's + ``turn_idx`` even when the trainable item's role is not ``assistant`` + (e.g. a reasoning-only response, or a ``function_call``). + + ``input_messages`` is the initial ``responses_create_params.input`` list — + images there (e.g. a single-shot user prompt for tool-based envs like + circle-click) are consumed by the first trainable turn's tokenized prompt + and must land in the first bucket. Agents like ``gym_v_agent`` that keep + ``input`` empty and inject observations as ``function_call_output`` items + are unaffected — the seed is a no-op when ``input_messages`` is empty. + """ + per_turn: list[list[Image.Image]] = [] + pending: list[Image.Image] = [] + for item in input_messages or (): + if isinstance(item, dict) and item.get("role") != "assistant": + pending.extend(_extract_input_images_from_message(item)) + for item in output: + if _is_trainable_output_item(item): + per_turn.append(pending) + pending = [] + elif item.get("role") != "assistant": + pending.extend(_extract_input_images_from_message(item)) + return per_turn + + +def _without_initial_media_sources( + messages: Any, initial_sources: list[Any] +) -> tuple[Any, bool]: + """Copy Responses messages and remove one ordered copy of initial images and videos.""" + if not isinstance(messages, list): + return messages, False + + filtered = deepcopy(messages) + remaining_sources = list(initial_sources) + for message in filtered: + if not isinstance(message, dict): + continue + content = message.get("content") + if not isinstance(content, list): + continue + + filtered_content = [] + for part in content: + part_sources = extract_input_media_sources_from_responses_messages( + [{"content": [part]}] + ) + if ( + remaining_sources + and len(part_sources) == 1 + and media_sources_equal(part_sources[0], remaining_sources[0]) + ): + remaining_sources.pop(0) + continue + filtered_content.append(part) + message["content"] = filtered_content + + return filtered, not remaining_sources + + _VideoConfigValue = TypeVar("_VideoConfigValue") _LOCAL_VIDEO_METADATA_KEYS = frozenset( { @@ -64,17 +336,6 @@ def _require_video_config_value( return value -def _get_content_part_url(part: dict[str, Any], *keys: str) -> str: - """Return a string media source from a Responses/Chat content part.""" - for key in keys: - value = part.get(key) - if isinstance(value, dict): - value = value.get("url") or value.get("path") - if isinstance(value, str) and value: - return value - return "" - - def _resolve_local_video_path(source: str) -> str: """Resolve a local video source and reject unsupported remote schemes.""" parsed = urlparse(source) @@ -96,41 +357,6 @@ def _resolve_local_video_path(source: str) -> str: return str(path.resolve()) -def normalize_video_urls_in_examples(examples: list[dict[str, Any]]) -> None: - """Convert bare local video paths to file URLs before Gym dispatch.""" - for example in examples: - input_items = example.get("responses_create_params", {}).get("input", []) - if not isinstance(input_items, list): - continue - for item in input_items: - if not isinstance(item, dict): - continue - content = item.get("content", []) - if not isinstance(content, list): - continue - for part in content: - if ( - not isinstance(part, dict) - or part.get("type") not in VIDEO_CONTENT_TYPES - ): - continue - media_key = next( - (key for key in ("video_url", "video", "url") if key in part), - None, - ) - if media_key is None: - continue - source = _get_content_part_url(part, media_key) - if not source or urlparse(source).scheme: - continue - normalized = Path(_resolve_local_video_path(source)).as_uri() - original = part[media_key] - if isinstance(original, dict): - original["url"] = normalized - else: - part[media_key] = normalized - - def _extract_static_video_messages( nemo_gym_example: dict[str, Any], ) -> tuple[list[dict[str, Any]], str | None] | None: @@ -168,7 +394,9 @@ def _extract_static_video_messages( if part_type == "input_text": hf_content.append({"type": "text", "text": part["text"]}) elif part_type in VIDEO_CONTENT_TYPES: - source = _get_content_part_url(part, "video_url", "video", "url") + source = get_responses_content_part_url( + part, "video_url", "video", "url" + ) if not source: raise ValueError(f"{part_type} requires a non-empty video URL") video_sources.append(source) @@ -177,7 +405,9 @@ def _extract_static_video_messages( if not part.get("_is_video_frame"): has_still_images = True continue - source = _get_content_part_url(part, "image_url", "image", "url") + source = get_responses_content_part_url( + part, "image_url", "image", "url" + ) if not source: raise ValueError( "Cached Gym video frames require a non-empty image URL." @@ -249,80 +479,6 @@ def _extract_static_video_messages( return hf_messages, _resolve_local_video_path(video_sources[0]) -def _json_mapping(value: Any, *, field_name: str) -> dict[str, Any]: - if isinstance(value, dict): - return copy.deepcopy(value) - if not isinstance(value, str): - raise TypeError(f"{field_name} must be a JSON object string or a dict") - if not value.strip(): - raise ValueError(f"{field_name} must not be empty") - try: - decoded = json.loads(value) - except json.JSONDecodeError as exc: - raise ValueError(f"{field_name} must contain valid JSON") from exc - if not isinstance(decoded, dict): - raise TypeError(f"{field_name} JSON must decode to an object") - return decoded - - -def _metadata_extra_body(nemo_gym_example: dict[str, Any]) -> dict[str, Any]: - params = nemo_gym_example.get("responses_create_params", {}) - if not isinstance(params, dict): - raise TypeError("responses_create_params must be a dict") - metadata = params.get("metadata", {}) - if not isinstance(metadata, dict): - raise TypeError("responses_create_params.metadata must be a dict") - if "extra_body" not in metadata: - return {} - return _json_mapping( - metadata["extra_body"], - field_name="responses_create_params.metadata.extra_body", - ) - - -def _chat_template_kwargs_for_processor( - nemo_gym_example: dict[str, Any], -) -> dict[str, Any]: - params = nemo_gym_example.get("responses_create_params", {}) - if not isinstance(params, dict): - raise TypeError("responses_create_params must be a dict") - metadata = params.get("metadata", {}) - if not isinstance(metadata, dict): - raise TypeError("responses_create_params.metadata must be a dict") - - extra_body = _metadata_extra_body(nemo_gym_example) - processor_kwargs: dict[str, Any] = {} - raw_chat_template_kwargs = metadata.get( - "chat_template_kwargs", extra_body.get("chat_template_kwargs") - ) - chat_template_kwargs = ( - _json_mapping( - raw_chat_template_kwargs, - field_name="responses_create_params.metadata.chat_template_kwargs", - ) - if raw_chat_template_kwargs is not None - else {} - ) - if chat_template_kwargs: - processor_kwargs["chat_template_kwargs"] = chat_template_kwargs - enable_thinking = chat_template_kwargs.get( - "enable_thinking", extra_body.get("enable_thinking") - ) - if enable_thinking is not None: - processor_kwargs["enable_thinking"] = enable_thinking - return processor_kwargs - - -def _deep_merge_dict(base: dict[str, Any], update: dict[str, Any]) -> dict[str, Any]: - merged = copy.deepcopy(base) - for key, value in update.items(): - if isinstance(value, dict) and isinstance(merged.get(key), dict): - merged[key] = _deep_merge_dict(merged[key], value) - else: - merged[key] = copy.deepcopy(value) - return merged - - def _inject_vllm_mm_processor_kwargs( nemo_gym_example: dict[str, Any], mm_processor_kwargs: dict[str, Any], @@ -387,7 +543,9 @@ def _replace_cached_video_frames_with_native_video( for part in content: if not isinstance(part, dict) or not part.get("_is_video_frame"): continue - frame_path = _get_content_part_url(part, "image_url", "image", "url") + frame_path = get_responses_content_part_url( + part, "image_url", "image", "url" + ) if not frame_path: raise ValueError( "Cached Gym video frames require a non-empty image URL." diff --git a/nemo_rl/environments/nemo_gym_request.py b/nemo_rl/environments/nemo_gym_request.py new file mode 100644 index 00000000000..46c44b689ec --- /dev/null +++ b/nemo_rl/environments/nemo_gym_request.py @@ -0,0 +1,155 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Helpers for reading and merging NeMo Gym request metadata.""" + +import copy +import json +from typing import Any + + +def _json_mapping(value: Any, *, field_name: str) -> dict[str, Any]: + """Return a copied dict from a mapping or JSON object string. + + Example: + ``_json_mapping('{"enabled": true}', field_name="options")`` returns + ``{"enabled": True}``. + + Args: + value: Dict or JSON object string. + field_name: Field name used in errors. + + Returns: + A new dictionary. + + Raises: + TypeError: If the value is not a dict or JSON object string. + ValueError: If the string is empty or invalid JSON. + """ + if isinstance(value, dict): + return copy.deepcopy(value) + if not isinstance(value, str): + raise TypeError(f"{field_name} must be a JSON object string or a dict") + if not value.strip(): + raise ValueError(f"{field_name} must not be empty") + try: + decoded = json.loads(value) + except json.JSONDecodeError as exc: + raise ValueError(f"{field_name} must contain valid JSON") from exc + if not isinstance(decoded, dict): + raise TypeError(f"{field_name} JSON must decode to an object") + return decoded + + +def _metadata_extra_body(nemo_gym_example: dict[str, Any]) -> dict[str, Any]: + """Read ``metadata.extra_body`` as a dict. + + Example: + An example with ``metadata.extra_body='{"seed": 1}'`` returns + ``{"seed": 1}``. + + Args: + nemo_gym_example: Example containing Responses API parameters. + + Returns: + Parsed ``extra_body``, or an empty dict when absent. + + Raises: + TypeError: If request parameters or metadata are not dictionaries. + ValueError: If ``extra_body`` contains invalid JSON. + """ + params = nemo_gym_example.get("responses_create_params", {}) + if not isinstance(params, dict): + raise TypeError("responses_create_params must be a dict") + metadata = params.get("metadata", {}) + if not isinstance(metadata, dict): + raise TypeError("responses_create_params.metadata must be a dict") + if "extra_body" not in metadata: + return {} + return _json_mapping( + metadata["extra_body"], + field_name="responses_create_params.metadata.extra_body", + ) + + +def _chat_template_kwargs_for_processor( + nemo_gym_example: dict[str, Any], +) -> dict[str, Any]: + """Build processor kwargs from NeMo Gym chat-template metadata. + + Example: + ``{"chat_template_kwargs": {"enable_thinking": False}}`` becomes the + processor kwarg with the same name and value. + + Args: + nemo_gym_example: Example containing Responses API parameters. + + Returns: + Keyword arguments for the processor's chat template. + + Raises: + TypeError: If request metadata has an unsupported type. + ValueError: If a JSON metadata value is empty or invalid. + """ + params = nemo_gym_example.get("responses_create_params", {}) + if not isinstance(params, dict): + raise TypeError("responses_create_params must be a dict") + metadata = params.get("metadata", {}) + if not isinstance(metadata, dict): + raise TypeError("responses_create_params.metadata must be a dict") + + extra_body = _metadata_extra_body(nemo_gym_example) + processor_kwargs: dict[str, Any] = {} + raw_chat_template_kwargs = metadata.get( + "chat_template_kwargs", extra_body.get("chat_template_kwargs") + ) + chat_template_kwargs = ( + _json_mapping( + raw_chat_template_kwargs, + field_name="responses_create_params.metadata.chat_template_kwargs", + ) + if raw_chat_template_kwargs is not None + else {} + ) + if chat_template_kwargs: + processor_kwargs["chat_template_kwargs"] = chat_template_kwargs + enable_thinking = chat_template_kwargs.get( + "enable_thinking", extra_body.get("enable_thinking") + ) + if enable_thinking is not None: + processor_kwargs["enable_thinking"] = enable_thinking + return processor_kwargs + + +def _deep_merge_dict(base: dict[str, Any], update: dict[str, Any]) -> dict[str, Any]: + """Recursively merge two dictionaries without modifying either input. + + Example: + Merging ``{"a": {"b": 1}}`` with ``{"a": {"c": 2}}`` returns + ``{"a": {"b": 1, "c": 2}}``. + + Args: + base: Initial mapping. + update: Values to merge into ``base``. + + Returns: + A recursively merged deep copy. + """ + merged = copy.deepcopy(base) + for key, value in update.items(): + if isinstance(value, dict) and isinstance(merged.get(key), dict): + merged[key] = _deep_merge_dict(merged[key], value) + else: + merged[key] = copy.deepcopy(value) + return merged diff --git a/nemo_rl/evals/eval.py b/nemo_rl/evals/eval.py index 670ed625d7e..1ea9f766f4e 100644 --- a/nemo_rl/evals/eval.py +++ b/nemo_rl/evals/eval.py @@ -331,10 +331,8 @@ async def _run_env_eval_impl( prompts = [] prompts_for_display = [] for i, message_log in enumerate(batch["message_log"]): - if is_multimodal and batch["vllm_content"][i] is not None: - vllm_content = batch["vllm_content"][i] - prompt_dict = {"prompt": vllm_content} - multi_modal_data = {} + multi_modal_data = {} + if is_multimodal: audios = batch.get("vllm_audios", None) if audios is not None and len(audios[i]) > 0: multi_modal_data["audio"] = ( @@ -350,10 +348,34 @@ async def _run_env_eval_impl( multi_modal_data["video"] = ( videos[i][0] if len(videos[i]) == 1 else videos[i] ) + + vllm_content = batch["vllm_content"][i] if is_multimodal else None + if vllm_content is not None: + prompt_dict = {"prompt": vllm_content} + prompt_display = vllm_content if multi_modal_data: prompt_dict["multi_modal_data"] = multi_modal_data prompts.append(prompt_dict) - prompts_for_display.append(vllm_content) + prompts_for_display.append(prompt_display) + elif multi_modal_data: + # Placeholder-style processors pass prompt_token_ids with media. + prompt_token_ids = [] + for message in message_log: + token_ids = message["token_ids"] + prompt_token_ids.extend( + token_ids.tolist() + if isinstance(token_ids, torch.Tensor) + else token_ids + ) + prompt_dict = { + "prompt_token_ids": prompt_token_ids, + "multi_modal_data": multi_modal_data, + } + prompt_display = "\n".join( + str(message["content"]) for message in message_log + ) + prompts.append(prompt_dict) + prompts_for_display.append(prompt_display) else: # Text-only fallback: use raw prompt strings (vLLM will tokenize them). # Note: utils.py's format_prompt_for_vllm_generation uses pre-tokenized diff --git a/nemo_rl/experience/payload.py b/nemo_rl/experience/payload.py index e10a9026fa0..417a827ab85 100644 --- a/nemo_rl/experience/payload.py +++ b/nemo_rl/experience/payload.py @@ -68,7 +68,7 @@ def record_to_train_batch( Returns: BatchedDataDict with input_ids, input_lengths, generation_logprobs, token_mask, sample_mask, prompt_ids_for_adv, total_reward, violation counts, and optional - routed_experts. + routed_experts, plus multimodal model inputs preserved as PackedTensor values. """ # Lazy imports: grpo and llm_message_utils transitively pull # experience.rollouts, so importing at module top risks a cycle. @@ -121,6 +121,7 @@ def record_to_train_batch( } if ROUTED_EXPERTS_FIELD in flat: train_data[ROUTED_EXPERTS_FIELD] = flat[ROUTED_EXPERTS_FIELD] + train_data.update(flat.get_multimodal_dict(as_tensors=False)) return BatchedDataDict[Any](train_data) @@ -145,10 +146,13 @@ def pack_payload( """ lengths = train_batch["input_lengths"] n = int(lengths.shape[0]) - tensor_fields: dict[str, torch.Tensor | np.ndarray] = { + from nemo_rl.data.multimodal_utils import PackedTensor + + tensor_fields: dict[str, Any] = { k: v for k, v in train_batch.items() if isinstance(v, torch.Tensor) + or isinstance(v, PackedTensor) or (isinstance(v, np.ndarray) and v.dtype == object) } fields_td = pack_jagged_fields( diff --git a/nemo_rl/experience/rollout_manager.py b/nemo_rl/experience/rollout_manager.py index 504f146c5f0..cf07e03c27a 100644 --- a/nemo_rl/experience/rollout_manager.py +++ b/nemo_rl/experience/rollout_manager.py @@ -30,6 +30,7 @@ TQReplayBuffer, ) from nemo_rl.data.interfaces import DatumSpec, LLMMessageLogType +from nemo_rl.data.llm_message_utils import batched_message_log_to_flat_message from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.environments.interfaces import EnvironmentInterface from nemo_rl.experience.failures import ( @@ -616,9 +617,14 @@ async def _generate_response( Returns: Tuple of (assistant_message, input_lengths, gen_metrics) """ - # Prepare generation input - input_ids = torch.cat([m["token_ids"] for m in message_log]).unsqueeze(0) - input_lengths = torch.tensor([input_ids.shape[1]], dtype=torch.int32) + # Flatten both tokens and model-ready multimodal inputs. Building this + # from token_ids alone leaves expanded media placeholders in the prompt + # without the pixel tensors Megatron needs to project. + flat_messages, input_lengths = batched_message_log_to_flat_message( + [message_log], + pad_value_dict={"token_ids": self._tokenizer.pad_token_id}, + ) + input_ids = flat_messages["token_ids"] generation_input_data = BatchedDataDict[GenerationDatumSpec]( { "input_ids": input_ids, @@ -626,6 +632,9 @@ async def _generate_response( "stop_strings": [stop_strings], } ) + generation_input_data.update( + flat_messages.get_multimodal_dict(as_tensors=False) + ) # Generate response # TODO: update generate_async to return a single item directly diff --git a/nemo_rl/experience/rollouts.py b/nemo_rl/experience/rollouts.py index 0ca333aa242..4ba0f1fdf5f 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.distributed.batched_data_dict import BatchedDataDict from nemo_rl.environments.interfaces import ( @@ -125,7 +126,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/experience/sync_rollout_actor.py b/nemo_rl/experience/sync_rollout_actor.py index 9f7f7e9ea03..28dda440d69 100644 --- a/nemo_rl/experience/sync_rollout_actor.py +++ b/nemo_rl/experience/sync_rollout_actor.py @@ -43,6 +43,7 @@ import ray import torch +from nemo_rl.data.multimodal_utils import PackedTensor from nemo_rl.data_plane.column_io import kv_first_write from nemo_rl.data_plane.interfaces import KVBatchMeta from nemo_rl.data_plane.schema import ROUTED_EXPERTS_FIELD @@ -316,7 +317,7 @@ def rollout_to_tq( if ROUTED_EXPERTS_FIELD in flat: bulk_batch[ROUTED_EXPERTS_FIELD] = flat[ROUTED_EXPERTS_FIELD] for k, v in flat.get_multimodal_dict(as_tensors=False).items(): - if isinstance(v, torch.Tensor): + if isinstance(v, (torch.Tensor, PackedTensor)): bulk_batch[k] = v # ``content`` (raw assistant text per sample) — rides TQ as a # NonTensorStack so the driver can fetch it back at jsonl time diff --git a/nemo_rl/models/generation/megatron/config.py b/nemo_rl/models/generation/megatron/config.py index d79c4d5ba81..98d0733a80e 100644 --- a/nemo_rl/models/generation/megatron/config.py +++ b/nemo_rl/models/generation/megatron/config.py @@ -48,6 +48,9 @@ class MCoreGenerationSpecificArgs(TypedDict): materialize_only_last_token_logits: bool enable_chunked_prefill: bool enable_prefix_caching: bool + async_sched_mode: NotRequired[Literal["legacy", "async"]] + vision_embedding_cache_max_bytes: NotRequired[int] + allow_stale_multimodal_embeddings: NotRequired[bool] refit_backend: Literal["gloo", "nccl", "nvshmem"] num_speculative_tokens: int @@ -55,6 +58,24 @@ class MCoreGenerationSpecificArgs(TypedDict): mamba_inference_ssm_states_dtype: NotRequired[str] mamba_inference_conv_states_dtype: NotRequired[str] + # Raw media preprocessing corresponding with Megatron's + # ImageProcessingConfig / VideoProcessingConfig. + # `video_num_frames` is required for video. + vision_model_type: NotRequired[str] + image_dynamic_resolution: NotRequired[bool] + video_num_frames: NotRequired[int] # Frames sampled per video. + video_temporal_patch_size: NotRequired[int] # Frames per temporal patch. + video_target_num_patches: NotRequired[int] # Overrides the image max-patch budget. + video_maintain_aspect_ratio: NotRequired[bool] + + # Fully-qualified class path of the MCore inference wrapper, e.g. + # "megatron.core.inference.model_inference_wrappers.multimodal. + # nemotron_omni_inference_wrapper.NemotronOmniInferenceWrapper". + # Resolved by `_get_megatron_inference_wrapper_cls`; its `supports_*` + # attributes gate which modalities are preprocessed. Not media preprocessing + # itself, and used on the direct generate path as well as the HTTP endpoint. + megatron_inference_wrapper: NotRequired[str] + # KV cache lifecycle across suspend/resume: # - "persist": cache stays allocated; CUDA graphs remain valid (default) # - "offload": cache is moved off-GPU between iterations @@ -70,6 +91,8 @@ class MCoreGenerationSpecificArgs(TypedDict): # FP8/MXFP8 for the dedicated (non-colocated) inference model; # merged into its `megatron_cfg` by `merged_inference_megatron_cfg`. fp8_cfg: NotRequired[Fp8Config] + # Merged into megatron_cfg for gen workers; required for EP>1 + local CUDA graphs. + moe_pad_experts_for_cuda_graph_inference: NotRequired[bool] class MCoreGenerationConfig(GenerationConfig): diff --git a/nemo_rl/models/generation/megatron/megatron_generation.py b/nemo_rl/models/generation/megatron/megatron_generation.py index 7298ec95687..53294bce379 100644 --- a/nemo_rl/models/generation/megatron/megatron_generation.py +++ b/nemo_rl/models/generation/megatron/megatron_generation.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +from copy import deepcopy from typing import TYPE_CHECKING, Any, AsyncGenerator, Optional, cast import ray @@ -210,8 +211,8 @@ def __init__( Args: config: PolicyConfig for the Megatron model. tokenizer: The tokenizer for the model. - cluster: Cluster to deploy a dedicated inference Policy on. - policy: Existing training Policy to reuse for generation. + cluster: Cluster for a dedicated, non-colocated inference Policy. + policy: Existing training Policy reused for colocated generation. name_prefix: Prefix for naming the worker group (non-colocated only). processor: Optional processor for VLMs (non-colocated only). skip_weight_load: Do not load the weights from the checkpoint; refit will do it. @@ -232,7 +233,8 @@ def __init__( ) # `self.cfg` exposes the `generation` that matches the `GenerationInterface` contract. - # `self._policy_config` keeps a reference to the full PolicyConfig. + # `self._policy_config` keeps a reference to the full PolicyConfig. Dedicated + # inference receives a copy because worker setup may modify it. self._policy_config = config self.cfg: MCoreGenerationConfig = config["generation"] # Populated after the first prepare_for_generation (which starts the HTTP server). @@ -252,7 +254,7 @@ def __init__( # Stand up a dedicated inference-only policy. self._owns_policy = True self._policy_config = { - **config, + **deepcopy(config), "megatron_cfg": self.effective_megatron_cfg(config), } # Reserve GPUs before Policy workers grab them, to prevent disjoint NVLS domains. diff --git a/nemo_rl/models/generation/megatron/megatron_worker.py b/nemo_rl/models/generation/megatron/megatron_worker.py index 8c9ee1f7c61..737c89664ec 100644 --- a/nemo_rl/models/generation/megatron/megatron_worker.py +++ b/nemo_rl/models/generation/megatron/megatron_worker.py @@ -14,15 +14,17 @@ import asyncio import gc +import importlib import os import threading import time import warnings -from typing import AsyncGenerator, Optional +from typing import Any, AsyncGenerator, Optional import requests import torch from megatron.core.inference.config import ( + AsyncScheduleMode, InferenceConfig, KVCacheManagementMode, MambaInferenceStateConfig, @@ -53,6 +55,7 @@ ) from megatron.core.utils import unwrap_model +from nemo_rl.data.multimodal_utils import CACHED_VIDEO_FRAME_MANIFEST_MAGIC from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.models.generation.interfaces import ( GenerationDatumSpec, @@ -60,8 +63,12 @@ verify_right_padding, ) from nemo_rl.models.generation.megatron.utils import ( + build_image_preprocessing_config, + build_prompt_and_multimodal_data, + build_video_preprocessing_config, log_gpu_memory, resolve_torch_dtype, + sample_vision_tensors, ) from nemo_rl.models.megatron.memory_saver import ( HAVE_TORCH_MEMORY_SAVER, @@ -81,12 +88,14 @@ class MegatronGenerationMixin: - rank: global rank (used for logging). - tokenizer: HF tokenizer. - megatron_tokenizer: tokenizer for inference. + - processor: optional multimodal processor. - is_generation_colocated: Whether colocated or distributed. - _reserved_http_server_socket: driver-reserved server socket, or None. """ # Colocated-reshard hosts assign the dedicated inference-layout model here # (see MegatronPolicyWorkerImpl._build_colocated_inference_model). + processor: Optional[Any] = None inference_model = None _colocated_reshard_plan = None @@ -112,6 +121,89 @@ def _init_inference_engine_state(self) -> None: self._inference_loop = None self._inference_thread = None + def _get_megatron_inference_wrapper_cls(self) -> Optional[type]: + """Resolve the configured Megatron inference wrapper, if any. + + Returns: + The wrapper class, or None when no wrapper is configured. + """ + class_path = self.cfg["generation"]["mcore_generation_config"].get( + "megatron_inference_wrapper" + ) + if class_path is None: + return None + # Resolved once per worker: this is called per sample during generation. + cached = getattr(self, "_megatron_inference_wrapper_cls", None) + if cached is not None: + return cached + module_name, _, class_name = class_path.rpartition(".") + if not module_name: + raise ValueError( + "megatron_inference_wrapper must be a fully qualified class name, " + f"got {class_path!r}." + ) + try: + wrapper_cls = getattr(importlib.import_module(module_name), class_name) + except (ImportError, AttributeError) as e: + raise ValueError( + f"Could not resolve megatron_inference_wrapper {class_path!r} " + f"(from policy.generation.mcore_generation_config): {e}" + ) from e + self._megatron_inference_wrapper_cls = wrapper_cls + return wrapper_cls + + @staticmethod + def _wrapper_supports_modality( + inference_wrapper_cls: Optional[type], modality: str + ) -> bool: + """Whether the configured inference wrapper advertises `modality` support.""" + return bool( + inference_wrapper_cls is not None + and getattr(inference_wrapper_cls, f"supports_{modality}", False) + ) + + def _inference_model_and_media_parts(self, inference_wrapper_cls=None): + """Return the language model and its optional multimodal parent.""" + model = unwrap_model(self._gen_model()) + if isinstance(model, (list, tuple)): + if len(model) != 1: + raise NotImplementedError("Virtual pipeline models are not supported.") + model = model[0] + if inference_wrapper_cls is None: + inference_wrapper_cls = self._get_megatron_inference_wrapper_cls() + if not any( + self._wrapper_supports_modality(inference_wrapper_cls, modality) + for modality in ("image", "video", "audio") + ): + return model, None + return model.language_model, model + + def _build_image_preprocessing_config(self, generation_config: dict[str, Any]): + """Build raw-image preprocessing settings.""" + inference_wrapper_cls = self._get_megatron_inference_wrapper_cls() + if not self._wrapper_supports_modality(inference_wrapper_cls, "image"): + return None + processor = self.processor + if processor is None: + raise ValueError( + "Megatron multimodal generation requires the policy processor." + ) + # Omit absent keys entirely so MCore's own dataclass defaults apply; + # passing None would override them (these fields are `bool`, not `Optional`). + image_kwargs: dict[str, Any] = {} + if "image_dynamic_resolution" in generation_config: + image_kwargs["dynamic_resolution"] = bool( + generation_config["image_dynamic_resolution"] + ) + if "vision_model_type" in generation_config: + image_kwargs["vision_model_type"] = str( + generation_config["vision_model_type"] + ) + return build_image_preprocessing_config( + processor.image_processor, + **image_kwargs, + ) + def _setup_colocated_cuda_graph_managers(self) -> None: """Create inference CUDA-graph managers for shared-model colocated generation. @@ -223,8 +315,13 @@ def _initialize_inference_engine(self, mcore_generation_config: dict) -> None: ) from megatron.core.utils import get_attr_wrapped_model - gen_model = self._gen_model() - pg_collection = get_attr_wrapped_model(gen_model, "pg_collection") + inference_wrapper_cls = self._get_megatron_inference_wrapper_cls() + inference_model, media_model = self._inference_model_and_media_parts( + inference_wrapper_cls + ) + engine_model = media_model if media_model is not None else self._gen_model() + pg_collection = get_attr_wrapped_model(self._gen_model(), "pg_collection") + model_config = inference_model.config buffer_size_gb = mcore_generation_config["buffer_size_gb"] num_cuda_graphs = mcore_generation_config["num_cuda_graphs"] @@ -245,7 +342,9 @@ def _initialize_inference_engine(self, mcore_generation_config: dict) -> None: num_speculative_tokens = mcore_generation_config["num_speculative_tokens"] max_requests = mcore_generation_config.get("max_requests") - mamba_inference_state_config = MambaInferenceStateConfig.from_model(gen_model) + mamba_inference_state_config = MambaInferenceStateConfig.from_model( + inference_model + ) is_hybrid_model = mamba_inference_state_config is not None if is_hybrid_model: if ( @@ -270,11 +369,36 @@ def _initialize_inference_engine(self, mcore_generation_config: dict) -> None: logging_step_interval = 0 # flashinfer's fused-RoPE kernel only dispatches fp16/bf16 q/k. - use_flashinfer_fused_rope = gen_model.config.params_dtype in ( + use_flashinfer_fused_rope = model_config.params_dtype in ( torch.float16, torch.bfloat16, ) + image_preprocessing_config = self._build_image_preprocessing_config( + mcore_generation_config + ) + video_preprocessing_config = build_video_preprocessing_config( + image_preprocessing_config, + mcore_generation_config, + frame_manifest_magic=CACHED_VIDEO_FRAME_MANIFEST_MAGIC, + ) + + # Only forward keys the config actually sets, so MCore's InferenceConfig + # defaults stay the single source of truth for the ones it omits. + inference_overrides: dict[str, Any] = {} + if "async_sched_mode" in mcore_generation_config: + inference_overrides["async_sched_mode"] = AsyncScheduleMode( + mcore_generation_config["async_sched_mode"] + ) + if "vision_embedding_cache_max_bytes" in mcore_generation_config: + inference_overrides["vision_embedding_cache_max_bytes"] = int( + mcore_generation_config["vision_embedding_cache_max_bytes"] + ) + if "allow_stale_multimodal_embeddings" in mcore_generation_config: + inference_overrides["allow_stale_multimodal_embeddings"] = bool( + mcore_generation_config["allow_stale_multimodal_embeddings"] + ) + inference_config = InferenceConfig( block_size_tokens=block_size_tokens, buffer_size_gb=buffer_size_gb, @@ -290,6 +414,7 @@ def _initialize_inference_engine(self, mcore_generation_config: dict) -> None: materialize_only_last_token_logits=materialize_only_last_token_logits, enable_chunked_prefill=enable_chunked_prefill, enable_prefix_caching=mcore_generation_config["enable_prefix_caching"], + **inference_overrides, prefix_caching_coordinator_policy=PrefixCachingCoordinatorPolicy( "first_prefix_block" ), @@ -303,19 +428,30 @@ def _initialize_inference_engine(self, mcore_generation_config: dict) -> None: num_speculative_tokens=num_speculative_tokens, logprobs_mode=mcore_generation_config["logprobs_mode"], max_requests=max_requests, + image_preprocessing_config=image_preprocessing_config, + video_preprocessing_config=video_preprocessing_config, ) if "inference_cuda_graph_scope" in mcore_generation_config: - gen_model.config.inference_cuda_graph_scope = InferenceCudaGraphScope[ + engine_model.config.inference_cuda_graph_scope = InferenceCudaGraphScope[ mcore_generation_config["inference_cuda_graph_scope"] ] self.inference_context = DynamicInferenceContext( - gen_model.config, inference_config - ) - self.inference_wrapped_model = GPTInferenceWrapper( - gen_model, self.inference_context + engine_model.config, inference_config ) + if media_model is None: + self.inference_wrapped_model = GPTInferenceWrapper( + engine_model, self.inference_context + ) + else: + if inference_wrapper_cls is None: + raise ValueError( + "Multimodal inference requires megatron_inference_wrapper." + ) + self.inference_wrapped_model = inference_wrapper_cls( + engine_model, self.inference_context + ) text_generation_controller = TextGenerationController( inference_wrapped_model=self.inference_wrapped_model, tokenizer=self.megatron_tokenizer, @@ -432,6 +568,7 @@ def _setup_openai_api_server(self) -> str: parsers=self.cfg["generation"]["mcore_generation_config"]["parsers"], verbose=False, sock=reserved_socket, + multimodal_prompt_config=self.inference_wrapped_model.multimodal_prompt_config, ) base_url = f"http://{ip}:{server_port}/v1" @@ -495,7 +632,11 @@ def finish_generation(self, *, release_gpu: bool = True) -> None: print(f"[Rank {self.rank}] finishing generation", flush=True) log_gpu_memory("finish_generation START") - lang_module = unwrap_model(self._gen_model()) + inference_model, media_model = self._inference_model_and_media_parts() + lang_module = unwrap_model(inference_model) + graph_module = unwrap_model( + media_model if media_model is not None else inference_model + ) if self.is_generation_colocated: if self._inference_engine_initialized and not self._inference_engine_asleep: @@ -504,7 +645,8 @@ def finish_generation(self, *, release_gpu: bool = True) -> None: "cuda_graph_impl" ] if cuda_graph_impl != "none": - toggle_cuda_graphs(lang_module, set_to="none") + # Restore the same full model tree cached during worker setup. + toggle_cuda_graphs(graph_module, set_to="none") # Need to turn off padding before training. # Gains nightly MoE coverage once #2884 and #3570 merge. set_decode_expert_padding(lang_module, set_to=False) @@ -552,10 +694,6 @@ def prepare_for_generation(self, tags=None, **kwargs) -> None: if self._colocated_reshard_plan is not None: self._build_colocated_inference_model(self.cfg) - gen_model = self._gen_model() - # `flash_decode` selects Megatron Inference's deprecated static-batching decode path, - # which would cause an assertion error if taken. - gen_model.config.flash_decode = False if self.is_generation_colocated and self.inference_model is None: self.model = self.move_model( self.model, "cuda", move_params=True, move_grads=False @@ -568,13 +706,19 @@ def prepare_for_generation(self, tags=None, **kwargs) -> None: and self._forward_pre_hook_enabled() ): self._disable_forward_pre_hook_until_next_train_step(param_sync=True) - gen_model = self.model # Colocated reshard (hosts without a dedicated inference model skip it). if self.inference_model is not None: self._reshard_into_inference_model() - lang_module = unwrap_model(gen_model) + inference_model, media_model = self._inference_model_and_media_parts() + # `flash_decode` selects Megatron Inference's deprecated static-batching decode path, + # which would cause an assertion error if taken. + inference_model.config.flash_decode = False + lang_module = unwrap_model(inference_model) + graph_module = unwrap_model( + media_model if media_model is not None else inference_model + ) lang_module.eval() rotary_module = getattr(lang_module, "rotary_pos_emb", None) @@ -585,7 +729,9 @@ def prepare_for_generation(self, tags=None, **kwargs) -> None: cuda_graph_impl = mcore_generation_config["cuda_graph_impl"] if cuda_graph_impl != "none": - toggle_cuda_graphs(lang_module, set_to=cuda_graph_impl) + # Use the same root object as _setup_colocated_cuda_graph_managers; + # CUDA-graph manager caches are keyed by model identity. + toggle_cuda_graphs(graph_module, set_to=cuda_graph_impl) # tags=["weights"] means we are inside refit_policy_generation between # suspend_for_refit and the weight transfer — the engine was intentionally @@ -606,7 +752,11 @@ def report_dp_openai_server_base_url(self) -> Optional[str]: return self.base_url def _build_sampling_params( - self, greedy: bool, stop_words: Optional[list[str]] + self, + greedy: bool, + stop_words: Optional[list[str]], + *, + return_prompt_tokens: bool = False, ) -> SamplingParams: """Build mcore SamplingParams for a single request.""" top_k_cfg = self.cfg["generation"]["top_k"] @@ -626,6 +776,7 @@ def _build_sampling_params( num_tokens_to_generate=self.cfg["generation"]["max_new_tokens"], termination_id=self.megatron_tokenizer.eod, stop_words=stop_words, + return_prompt_tokens=return_prompt_tokens, ) def _merge_stop_strings( @@ -641,10 +792,26 @@ def _merge_stop_strings( stop_set.update(sample_ss) return list(stop_set) if stop_set else None + def _sample_vision_tensors(self, data, index: int): + """Return one sample's vision tensors from RL PackedTensors.""" + return sample_vision_tensors(data, index) + + def _build_prompt_and_multimodal_data(self, data, index: int): + """Build one pre-expanded token prompt and optional MCore media payload.""" + return build_prompt_and_multimodal_data( + data, + index, + sample_tensors=self._sample_vision_tensors, + supports_modality=lambda modality: self._wrapper_supports_modality( + self._get_megatron_inference_wrapper_cls(), + modality, + ), + ) + def _prepare_data_for_generation( self, data: BatchedDataDict[GenerationDatumSpec], greedy: bool = False - ) -> tuple[torch.Tensor, torch.Tensor, list[SamplingParams]]: - """Build the prompt tensors and a per-request SamplingParams for each sample.""" + ) -> tuple[list[list[int]], list[Optional[Any]], list[SamplingParams]]: + """Build prompts, optional media payloads, and sampling parameters.""" if data is not None: assert isinstance(data, BatchedDataDict), ( f"data must be a BatchedDataDict, got type: {type(data)}" @@ -657,21 +824,29 @@ def _prepare_data_for_generation( f"Input to Megatron Generation worker is not properly right-padded: {error_msg}" ) - prompt_tokens_tensor = data["input_ids"].cuda() - prompt_lengths_tensor = data["input_lengths"] - batch_stop_strings = data.get("stop_strings", []) + prompts: list[list[int]] = [] + multi_modal_data_list: list[Optional[Any]] = [] sampling_params = [] - for i in range(prompt_tokens_tensor.size(0)): + for i in range(data.size): + prompt, multi_modal_data = self._build_prompt_and_multimodal_data(data, i) sample_stop_strings = ( batch_stop_strings[i] if i < len(batch_stop_strings) else None ) stop_words = self._merge_stop_strings( [sample_stop_strings] if sample_stop_strings else None ) - sampling_params.append(self._build_sampling_params(greedy, stop_words)) + prompts.append(prompt) + multi_modal_data_list.append(multi_modal_data) + sampling_params.append( + self._build_sampling_params( + greedy, + stop_words, + return_prompt_tokens=multi_modal_data is not None, + ) + ) - return prompt_tokens_tensor, prompt_lengths_tensor, sampling_params + return prompts, multi_modal_data_list, sampling_params def _parse_result_to_batched_data_dict( self, @@ -685,6 +860,22 @@ def _parse_result_to_batched_data_dict( max_gen_seq_len = max(len(x.generated_tokens) for x in result) padded_input_length = input_ids.size(1) + expected_prompt_lengths = [int(length) for length in input_lengths.tolist()] + inference_prompt_lengths = [ + len(x.prompt_tokens) + if getattr(x, "prompt_tokens", None) is not None + else expected_prompt_lengths[i] + for i, x in enumerate(result) + ] + if any(getattr(x, "prompt_tokens", None) is not None for x in result): + if inference_prompt_lengths != expected_prompt_lengths: + raise RuntimeError( + "Megatron inference prompt lengths do not match the training " + "processor input lengths: " + f"inference={inference_prompt_lengths}, " + f"training={expected_prompt_lengths}." + ) + max_seq_len = padded_input_length + max_gen_seq_len output_ids_padded = torch.full( (batch_size, max_seq_len), @@ -707,9 +898,10 @@ def _parse_result_to_batched_data_dict( ) for i in range(batch_size): # Take the prompt from the request we submitted rather than from the - # engine's reply: mcore only echoes prompt_tokens back when - # SamplingParams.return_prompt_tokens is set, and asking for them would - # ship the whole prompt over ZMQ for data we already hold. + # engine's reply. Multimodal requests do set + # SamplingParams.return_prompt_tokens, but only so the echoed tokens + # can be length-checked below; the padded output is still built from + # the prompt we already hold rather than shipped back over ZMQ. prompt_len = input_lengths[i].item() generated_tokens = result[i].generated_tokens seq_len = prompt_len + len(generated_tokens) @@ -754,7 +946,7 @@ def generate( - generation_lengths: Lengths of each response - unpadded_sequence_lengths: Lengths of each input + generated sequence """ - prompt_tokens_tensor, prompt_lengths_tensor, sampling_params = ( + prompts, multi_modal_data_list, sampling_params = ( self._prepare_data_for_generation(data, greedy) ) if self._inference_loop is None: @@ -763,8 +955,8 @@ def generate( ) future = asyncio.run_coroutine_threadsafe( self._generate_with_persistent_engine( - prompt_tokens_tensor, - prompt_lengths_tensor, + prompts, + multi_modal_data_list, sampling_params, ), self._inference_loop, @@ -794,13 +986,13 @@ async def _generate_single_item( index: int, ) -> tuple[int, BatchedDataDict[GenerationOutputSpec]]: datum = data.get_batch(index, 1) - prompt_tokens_tensor, prompt_lengths_tensor, sampling_params = ( + prompts, multi_modal_data_list, sampling_params = ( self._prepare_data_for_generation(datum, greedy) ) future = asyncio.run_coroutine_threadsafe( self._generate_with_persistent_engine( - prompt_tokens_tensor, - prompt_lengths_tensor, + prompts, + multi_modal_data_list, sampling_params, ), self._inference_loop, @@ -817,8 +1009,8 @@ async def _generate_single_item( async def _generate_with_persistent_engine( self, - prompt_tokens_tensor: torch.Tensor, - prompt_lengths_tensor: torch.Tensor, + prompts: list[list[int]], + multi_modal_data_list: list[Optional[Any]], sampling_params: list[SamplingParams], ) -> list: """Submit requests through the persistent inference client (rank 0 only).""" @@ -829,17 +1021,18 @@ async def _generate_with_persistent_engine( "Only rank 0 creates a client to communicate with the coordinator" ) - print( - f"[Rank {dist_rank}] Submitting {prompt_tokens_tensor.size(0)} requests to coordinator" - ) + print(f"[Rank {dist_rank}] Submitting {len(prompts)} requests to coordinator") futures = [] - for prompt_tokens, prompt_len, request_sampling_params in zip( - prompt_tokens_tensor, prompt_lengths_tensor, sampling_params, strict=True + for prompt, multi_modal_data, request_sampling_params in zip( + prompts, multi_modal_data_list, sampling_params, strict=True ): - prompt = prompt_tokens[: prompt_len.item()].tolist() futures.append( - self.inference_client.add_request(prompt, request_sampling_params) + self.inference_client.add_request( + prompt, + request_sampling_params, + multi_modal_data=multi_modal_data, + ) ) results: list[DynamicInferenceRequest] = await asyncio.gather(*futures) diff --git a/nemo_rl/models/generation/megatron/utils.py b/nemo_rl/models/generation/megatron/utils.py index 339d57e46cc..6e40693e2d5 100644 --- a/nemo_rl/models/generation/megatron/utils.py +++ b/nemo_rl/models/generation/megatron/utils.py @@ -12,10 +12,241 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dataclasses import replace +from typing import Any, Callable + import torch +from megatron.core.inference.config import ImageProcessingConfig, VideoProcessingConfig from megatron.core.inference.utils import device_memory_summary +def sample_vision_tensors(data, index: int): + """Return one sample's vision tensors from RL PackedTensors.""" + from nemo_rl.data.multimodal_utils import PackedTensor + + pixel_values = data.get("pixel_values") + imgs_sizes = data.get("imgs_sizes") + packed_num_frames = data.get("num_frames") + if pixel_values is None and imgs_sizes is None: + if packed_num_frames is not None: + raise ValueError("num_frames was provided without vision tensors.") + return None, None, None + if pixel_values is None or imgs_sizes is None: + raise ValueError( + "Megatron image generation requires both pixel_values and imgs_sizes." + ) + if not isinstance(pixel_values, PackedTensor) or not isinstance( + imgs_sizes, PackedTensor + ): + raise TypeError( + "Megatron image generation expects pixel_values and imgs_sizes " + "as per-sample PackedTensor values." + ) + if packed_num_frames is not None and not isinstance( + packed_num_frames, PackedTensor + ): + raise TypeError( + "Megatron video generation expects num_frames as a " + "per-sample PackedTensor value." + ) + + # `.tensors` is the physical segment list and only matches logical row + # indices while media are not deduplicated. + for name, packed in ( + ("pixel_values", pixel_values), + ("imgs_sizes", imgs_sizes), + ("num_frames", packed_num_frames), + ): + if packed is not None and packed._row_offsets is not None: + raise ValueError( + f"Megatron generation cannot index deduplicated {name}; " + "set deduplicate_multimodal_data=false (it is only supported " + "for the vLLM backend)." + ) + + imgs = pixel_values.tensors[index] + sizes = imgs_sizes.tensors[index] + num_frames = ( + packed_num_frames.tensors[index] if packed_num_frames is not None else None + ) + if imgs is None and sizes is None: + return None, None, None + if imgs is None or sizes is None: + raise ValueError( + "Megatron image generation requires matching per-sample " + "pixel_values and imgs_sizes." + ) + if imgs.ndim == 3: + imgs = imgs.unsqueeze(0) + if sizes.ndim == 1: + sizes = sizes.unsqueeze(0) + if num_frames is not None: + num_frames = num_frames.to(dtype=torch.int32).reshape(-1) + return imgs, sizes, num_frames + + +def build_prompt_and_multimodal_data( + data, + index: int, + *, + supports_modality: Callable[[str], bool], + sample_tensors: Callable = sample_vision_tensors, +): + """Build one pre-expanded token prompt and optional MCore media payload.""" + length = int(data["input_lengths"][index].item()) + prompt = data["input_ids"][index, :length].tolist() + imgs, imgs_sizes, num_frames = sample_tensors(data, index) + if imgs is None: + return prompt, None + + assert imgs_sizes is not None + is_video = num_frames is not None and bool(torch.any(num_frames > 1).item()) + modality = "video" if is_video else "image" + if not supports_modality(modality): + raise ValueError( + f"The configured megatron_inference_wrapper does not support " + f"{modality} inputs." + ) + if is_video: + if int(num_frames.sum().item()) != int(imgs_sizes.shape[0]): + raise ValueError( + "Video num_frames must partition imgs_sizes exactly: " + f"sum(num_frames)={int(num_frames.sum().item())}, " + f"imgs_sizes={imgs_sizes.shape[0]}." + ) + modality_data = { + "imgs": imgs, + "imgs_sizes": imgs_sizes, + "num_frames": num_frames, + } + else: + modality_data = {"imgs": imgs, "imgs_sizes": imgs_sizes} + return prompt, { + modality: modality_data, + "media_tokens_preexpanded": True, + } + + +def build_image_preprocessing_config( + image_processor: Any, + *, + dynamic_resolution: bool | None = None, + vision_model_type: str | None = None, +) -> ImageProcessingConfig: + """Translate an HF image processor to an MCore config. + + Args: + image_processor: HF image processor to read patch/normalization fields from. + dynamic_resolution: Override for `ImageProcessingConfig.dynamic_resolution`. + `None` leaves MCore's own default in place. + vision_model_type: Override for `ImageProcessingConfig.vision_model_type`. + `None` leaves MCore's own default in place. + """ + + def read(*names: str) -> Any: + for name in names: + value = getattr(image_processor, name, None) + if value is not None: + return value + return None + + patch_dim = read("patch_size", "patch_dim") + if isinstance(patch_dim, dict): + patch_dim = patch_dim.get("height", patch_dim.get("width")) + min_patches = read("min_num_patches") + max_patches = read("max_num_patches") + pixel_mean = read("norm_mean", "image_mean") + pixel_std = read("norm_std", "image_std") + + if ( + patch_dim is None + or min_patches is None + or max_patches is None + or pixel_mean is None + or pixel_std is None + ): + missing = [ + name + for name, value in ( + ("patch_size", patch_dim), + ("min_num_patches", min_patches), + ("max_num_patches", max_patches), + ("norm_mean", pixel_mean), + ("norm_std", pixel_std), + ) + if value is None + ] + raise ValueError( + f"{type(image_processor).__name__} does not expose {', '.join(missing)}, " + "so MCore cannot preprocess raw images the way this model's data " + "pipeline does." + ) + + downsample_ratio = read("downsample_ratio") + if downsample_ratio is not None: + merge_size = int(round(1.0 / float(downsample_ratio))) + else: + merge_size = int(read("merge_size", "spatial_merge_size") or 1) + + return ImageProcessingConfig( + patch_dim=int(patch_dim), + **( + {} + if dynamic_resolution is None + else {"dynamic_resolution": dynamic_resolution} + ), + **( + {} + if vision_model_type is None + else {"vision_model_type": vision_model_type} + ), + use_tiling=False, + pixel_shuffle=merge_size > 1, + spatial_merge_size=merge_size, + dynamic_resolution_min_patches=int(min_patches), + dynamic_resolution_max_patches=int(max_patches), + pixel_mean=[float(value) for value in pixel_mean], + pixel_std=[float(value) for value in pixel_std], + ) + + +def build_video_preprocessing_config( + image_config: ImageProcessingConfig | None, + generation_config: dict[str, Any], + *, + frame_manifest_magic: bytes, +) -> VideoProcessingConfig | None: + """Build video preprocessing when explicitly enabled by generation config.""" + video_num_frames = generation_config.get("video_num_frames") + if image_config is None or video_num_frames is None: + return None + + # Video configs. + video_kwargs: dict[str, Any] = {} + if "video_temporal_patch_size" in generation_config: + video_kwargs["temporal_patch_size"] = int( + generation_config["video_temporal_patch_size"] + ) + if "video_maintain_aspect_ratio" in generation_config: + video_kwargs["video_maintain_aspect_ratio"] = bool( + generation_config["video_maintain_aspect_ratio"] + ) + + target_num_patches = generation_config.get("video_target_num_patches") + if target_num_patches is not None: + image_config = replace( + image_config, + dynamic_resolution_max_patches=int(target_num_patches), + ) + + return VideoProcessingConfig( + image_config=image_config, + num_frames=int(video_num_frames), + frame_manifest_magic=frame_manifest_magic, + **video_kwargs, + ) + + def resolve_torch_dtype(val): """Convert a value to `torch.dtype`.""" if isinstance(val, torch.dtype): diff --git a/nemo_rl/models/generation/vllm/config.py b/nemo_rl/models/generation/vllm/config.py index dabf01d0b1e..c4d69be527b 100644 --- a/nemo_rl/models/generation/vllm/config.py +++ b/nemo_rl/models/generation/vllm/config.py @@ -289,7 +289,7 @@ def normalize_vllm_refit_config(config: VllmConfig) -> VllmRefitConfig | None: raise ValueError( "vllm_cfg.reset_encoder_cache_after_weight_update is not supported " f"with refit_transport={transport!r}: this transport's refit path " - "does not reset the multimodal encoder cache, so stale vision " + "does not reset the multimodal encoder cache, so stale multimodal " "embeddings would silently survive weight updates. Supported " "transports: null (collective/IPC) and 'nccl_reshard'." ) diff --git a/nemo_rl/models/generation/vllm/video_utils.py b/nemo_rl/models/generation/vllm/video_utils.py index c573a95d462..88a3059926e 100644 --- a/nemo_rl/models/generation/vllm/video_utils.py +++ b/nemo_rl/models/generation/vllm/video_utils.py @@ -24,13 +24,16 @@ import torch from PIL import Image +from nemo_rl.data.multimodal_utils import ( + CACHED_VIDEO_FRAME_MANIFEST_MAGIC, + CACHED_VIDEO_FRAME_MANIFEST_MIME, +) + VideoSamplingStyle = Literal["nemotron_vl"] _TORCHCODEC_END_OF_STREAM_ERROR = ( "Requested next frame while there are no more frames left to decode." ) -_CACHED_VIDEO_FRAME_MANIFEST_MAGIC = b"NEMO_RL_CACHED_VIDEO_FRAMES_V1\n" -_CACHED_VIDEO_FRAME_MANIFEST_MIME = "video/x-nemo-rl-cached-frames" def _round_video_frame_count( @@ -197,11 +200,11 @@ def build_cached_video_frame_data_url( "frame_paths": resolved_frames, "metadata": build_cached_video_frame_metadata(len(resolved_frames)), } - payload = _CACHED_VIDEO_FRAME_MANIFEST_MAGIC + json.dumps( + payload = CACHED_VIDEO_FRAME_MANIFEST_MAGIC + json.dumps( manifest, separators=(",", ":") ).encode("utf-8") encoded = base64.b64encode(payload).decode("ascii") - return f"data:{_CACHED_VIDEO_FRAME_MANIFEST_MIME};base64,{encoded}" + return f"data:{CACHED_VIDEO_FRAME_MANIFEST_MIME};base64,{encoded}" def _load_cached_video_frame_manifest( @@ -210,11 +213,11 @@ def _load_cached_video_frame_manifest( num_frames: int, ) -> tuple[np.ndarray, dict[str, Any]] | None: """Load an internal cached-frame manifest passed through vLLM VideoMediaIO.""" - if not data.startswith(_CACHED_VIDEO_FRAME_MANIFEST_MAGIC): + if not data.startswith(CACHED_VIDEO_FRAME_MANIFEST_MAGIC): return None try: - manifest = json.loads(data[len(_CACHED_VIDEO_FRAME_MANIFEST_MAGIC) :]) + manifest = json.loads(data[len(CACHED_VIDEO_FRAME_MANIFEST_MAGIC) :]) except (UnicodeDecodeError, json.JSONDecodeError) as exc: raise ValueError("Invalid cached Gym video frame manifest.") from exc if not isinstance(manifest, dict): diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index d0f57de802a..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 5c2063cd67b..aae2280e5b3 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/tq_policy.py b/nemo_rl/models/policy/tq_policy.py index 9f68526a937..d39bf036cd8 100644 --- a/nemo_rl/models/policy/tq_policy.py +++ b/nemo_rl/models/policy/tq_policy.py @@ -45,6 +45,7 @@ DP_TRAIN_FIELDS, LP_SEED_FIELDS, fields_with_optional_routed_experts, + fields_with_packed_tensor_payload, ) from nemo_rl.models.policy.lm_policy import Policy from nemo_rl.utils.flops_tracker import get_theoretical_tflops @@ -139,7 +140,9 @@ def load_data_plane_checkpoint(self, checkpoint_dir: str | Path) -> dict[str, An def shutdown(self) -> bool: # type: ignore[override] """Close the TQ client before shutting down the worker group.""" try: - self.dp_client.close() + dp_client = getattr(self, "dp_client", None) + if dp_client is not None: + dp_client.close() except Exception as e: warnings.warn(f"Error closing data-plane client: {e}") return super().shutdown() @@ -216,9 +219,9 @@ def _logprob_dispatch( ) -> None: """Shared body of get_logprobs_from_meta / get_reference_policy_logprobs_from_meta. - Logprob workers need only LP_SEED_FIELDS — narrow the meta's - field list so ``_fetch`` doesn't pull rollout-only payload (e.g. - multimodal). The same shape is used for both prev_lp and ref_lp. + Logprob workers need LP_SEED_FIELDS plus any multimodal model inputs + encoded in the rollout payload. The same shape is used for both prev_lp + and ref_lp. Workers compute the per-token tensor and commit it to TQ via the leader-rank ``_write_back_result_field``; the Ray return is always None, so this dispatcher just waits for completion. @@ -226,9 +229,12 @@ def _logprob_dispatch( spa, dba = self._packing_args("logprob_mb_tokens") lp_meta = self._isolated_meta( meta, - fields=fields_with_optional_routed_experts( - LP_SEED_FIELDS, - enabled=self._router_replay_enabled and include_router_replay, + fields=fields_with_packed_tensor_payload( + fields_with_optional_routed_experts( + LP_SEED_FIELDS, + enabled=self._router_replay_enabled and include_router_replay, + ), + meta.fields, ), task_name=task_name, ) @@ -333,8 +339,11 @@ def train_from_meta( # skipped this step (e.g. ``prev_logprobs`` under force_on_policy_ratio). train_meta = self._isolated_meta( meta, - fields=fields_with_optional_routed_experts( - train_fields, enabled=self._router_replay_enabled + fields=fields_with_packed_tensor_payload( + fields_with_optional_routed_experts( + train_fields, enabled=self._router_replay_enabled + ), + meta.fields, ), task_name="train", ) @@ -456,8 +465,11 @@ def train_microbatches_from_meta( spa, dba = self._packing_args("train_mb_tokens") train_meta = self._isolated_meta( meta, - fields=fields_with_optional_routed_experts( - DP_TRAIN_FIELDS, enabled=self._router_replay_enabled + fields=fields_with_packed_tensor_payload( + fields_with_optional_routed_experts( + DP_TRAIN_FIELDS, enabled=self._router_replay_enabled + ), + meta.fields, ), task_name="train", ) diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index 904a1359b5f..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( @@ -3341,10 +3343,13 @@ def prepare_for_training(self, *args, **kwargs): def finish_inference(self) -> None: """Offload model params to CPU after inference. Only used in PPO.""" + # MambaMixer.eval() recomputes and caches a state transition decay, + # -torch.exp(self.A_log.float()). Set the model in inference mode + # before offloading the model parameters (including self.A_log). + self.model.eval() self.model = self.move_model( self.model, "cpu", move_params=True, move_grads=False ) - self.model.eval() gc.collect() torch.cuda.empty_cache() @@ -3490,10 +3495,13 @@ def offload_after_refit(self): and self.inference_model is None and self._colocated_reshard_plan is None ) + # MambaMixer.eval() recomputes and caches a state transition decay, + # -torch.exp(self.A_log.float()). Set the model in inference mode + # before offloading the model parameters (including self.A_log). + self.model.eval() self.model = self.move_model( self.model, "cpu", move_params=not keep_params_for_generation ) - self.model.eval() torch.randn(1).cuda() # wake up torch allocator self.offload_before_refit() # rerun the old offload function diff --git a/nemo_rl/models/value/tq_value.py b/nemo_rl/models/value/tq_value.py index 062df22f12a..c01a8eea44f 100644 --- a/nemo_rl/models/value/tq_value.py +++ b/nemo_rl/models/value/tq_value.py @@ -110,6 +110,8 @@ def get_values_from_meta( timer: Optional timer for nested get_values measurements. """ spa, dba = self._packing_args("logprob_mb_tokens") + # The critic is text-only (built with is_vlm=False), so media columns + # are deliberately not fetched here. value_meta = self._isolated_meta( meta, fields=list(VALUE_SEED_FIELDS), @@ -161,6 +163,8 @@ def train_from_meta( micro_batch_size = mbs or self.cfg["train_micro_batch_size"] spa, dba = self._packing_args("train_mb_tokens") + # The critic is text-only (built with is_vlm=False), so media columns + # are deliberately not fetched here. train_meta = self._isolated_meta( meta, fields=list(DP_VALUE_TRAIN_FIELDS), 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/pyproject.toml b/pyproject.toml index 6b20966223d..9ebf4d40df0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -456,7 +456,7 @@ constraint-dependencies = [ # whose [tool.uv.sources] pins them to git. uv requires such URL deps to also appear as a # direct requirement or constraint of the root project. Keep these URLs/revs in sync with # Megatron-LM's [tool.uv.sources] (uv errors loudly on mismatch after a submodule bump). - "emerging-optimizers @ git+https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git@v0.2.0", + "emerging-optimizers @ git+https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git@v0.3.0", "fast-hadamard-transform @ git+https://github.com/Dao-AILab/fast-hadamard-transform.git@f134af63deb2df17e1171a9ec1ea4a7d8604d5ca", # Same story for megatron-fsdp, which nemo-automodel (r0.6.0) resolves from a # Megatron-LM fork via its own [tool.uv.sources] instead of PyPI. Keep this URL/rev diff --git a/pyrefly.toml b/pyrefly.toml index 5c93a2b0a46..ce9c2163780 100644 --- a/pyrefly.toml +++ b/pyrefly.toml @@ -145,7 +145,8 @@ project-includes = [ "nemo_rl/environments/interfaces.py", "nemo_rl/environments/math_environment.py", "nemo_rl/environments/metrics.py", - "nemo_rl/environments/nemo_gym_video.py", + "nemo_rl/environments/nemo_gym_multimodal.py", + "nemo_rl/environments/nemo_gym_request.py", "nemo_rl/environments/nemotron_utils.py", "nemo_rl/environments/rewards.py", "nemo_rl/environments/utils.py", diff --git a/tests/check_metrics.py b/tests/check_metrics.py index ac774435d0a..ff18d649488 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 every recorded value is finite.""" + vals = [float(v) for v in value.values()] + return bool(vals) and builtins.all(math.isfinite(v) for v in vals) + + def ratio_above(value, threshold): """Return the ratio of values that are >= threshold. @@ -148,6 +155,7 @@ def evaluate_check(data: dict, check: str) -> tuple[bool, str, object]: "max": max, "mean": mean, "median": median, + "all_finite": all_finite, "ratio_above": ratio_above, } diff --git a/tests/functional/L1_Functional_Tests_GB200_Megatron_Omni.sh b/tests/functional/L1_Functional_Tests_GB200_Megatron_Omni.sh new file mode 100755 index 00000000000..c8f916ef053 --- /dev/null +++ b/tests/functional/L1_Functional_Tests_GB200_Megatron_Omni.sh @@ -0,0 +1,38 @@ +#!/bin/bash +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +set -xeuo pipefail + +SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &> /dev/null && pwd) +PROJECT_ROOT=$(realpath "${SCRIPT_DIR}/../..") + +cd "${PROJECT_ROOT}" + +GPU_COUNT=$(nvidia-smi --query-gpu=index --format=csv,noheader | wc -l) +if (( GPU_COUNT < 2 )); then + echo "SKIP: Nemotron Omni functional tests require at least two GB200 GPUs" + exit 0 +fi + +# Both tests colocate TP2/EP2 training and generation on two GB200 GPUs. +export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1}" + +time uv run --no-sync bash ./tests/functional/nemotron_omni_clevr_megatron_1n2g.sh +time uv run --no-sync bash ./tests/functional/nemotron_omni_gym_video_megatron_1n2g.sh + +cd "${PROJECT_ROOT}/tests" +if compgen -G ".coverage*" > /dev/null; then + coverage combine .coverage* +fi diff --git a/tests/functional/nemotron_omni_clevr_megatron_1n2g.sh b/tests/functional/nemotron_omni_clevr_megatron_1n2g.sh new file mode 100755 index 00000000000..a1156ce9436 --- /dev/null +++ b/tests/functional/nemotron_omni_clevr_megatron_1n2g.sh @@ -0,0 +1,148 @@ +#!/usr/bin/env bash +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +set -euo pipefail + +SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd) +PROJECT_ROOT=$(realpath "${SCRIPT_DIR}/../..") + +if [[ -z "${HF_TOKEN:-}" ]]; then + echo "SKIP: HF_TOKEN is required for the Omni checkpoint" + exit 0 +fi + +GPU_COUNT=$(nvidia-smi --query-gpu=index --format=csv,noheader | wc -l) +if (( GPU_COUNT < 2 )); then + echo "SKIP: Omni CLEVR Megatron smoke requires at least two visible GPUs" + exit 0 +fi +DETECTED_CUDA_ARCH=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader -i 0) +export TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST:-${DETECTED_CUDA_ARCH}}" +MEGATRON_TRANSFORMER_IMPL="${MEGATRON_TRANSFORMER_IMPL:-inference_optimized}" +MEGATRON_CUDA_GRAPH_IMPL="${MEGATRON_CUDA_GRAPH_IMPL:-local}" +if [[ "${MEGATRON_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=2 \ + policy.megatron_cfg.env_vars.TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST}" \ + policy.megatron_cfg.tensor_model_parallel_size=2 \ + policy.megatron_cfg.expert_model_parallel_size=2 \ + policy.megatron_cfg.expert_tensor_parallel_size=1 \ + policy.megatron_cfg.context_parallel_size=1 \ + policy.megatron_cfg.sequence_parallel=true \ + policy.megatron_cfg.activation_checkpointing=true \ + policy.megatron_cfg.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=2 \ + policy.generation.max_new_tokens=128 \ + policy.generation.mcore_generation_config.tensor_model_parallel_size=2 \ + policy.generation.mcore_generation_config.expert_model_parallel_size=2 \ + policy.generation.mcore_generation_config.expert_tensor_parallel_size=1 \ + ++policy.generation.mcore_generation_config.context_parallel_size=1 \ + ++policy.generation.mcore_generation_config.moe_router_dtype=fp32 \ + policy.generation.mcore_generation_config.transformer_impl="${MEGATRON_TRANSFORMER_IMPL}" \ + policy.generation.mcore_generation_config.sequence_parallel=true \ + policy.generation.mcore_generation_config.refit_backend=nccl \ + policy.generation.mcore_generation_config.buffer_size_gb=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 \ + 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 \ + 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_1n2g.sh b/tests/functional/nemotron_omni_gym_video_megatron_1n2g.sh new file mode 100755 index 00000000000..d4d04f733ce --- /dev/null +++ b/tests/functional/nemotron_omni_gym_video_megatron_1n2g.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 < 2 )); then + echo "SKIP: Omni Gym-video Megatron smoke requires at least two GPUs" + exit 0 +fi +DETECTED_CUDA_ARCH=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader -i 0) +export TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST:-${DETECTED_CUDA_ARCH}}" +MEGATRON_TRANSFORMER_IMPL="${MEGATRON_TRANSFORMER_IMPL:-inference_optimized}" +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=2 \ + policy.megatron_cfg.env_vars.TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST}" \ + policy.megatron_cfg.tensor_model_parallel_size=2 \ + policy.megatron_cfg.pipeline_model_parallel_size=1 \ + policy.megatron_cfg.expert_model_parallel_size=2 \ + policy.megatron_cfg.expert_tensor_parallel_size=1 \ + policy.megatron_cfg.context_parallel_size=1 \ + policy.megatron_cfg.sequence_parallel=true \ + policy.megatron_cfg.activation_checkpointing=true \ + policy.megatron_cfg.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=true \ + policy.generation.colocated.resources.num_nodes=1 \ + policy.generation.colocated.resources.gpus_per_node=2 \ + policy.generation.max_new_tokens=128 \ + policy.generation.mcore_generation_config.expose_http_server=true \ + policy.generation.mcore_generation_config.tensor_model_parallel_size=2 \ + policy.generation.mcore_generation_config.expert_model_parallel_size=2 \ + policy.generation.mcore_generation_config.expert_tensor_parallel_size=1 \ + ++policy.generation.mcore_generation_config.context_parallel_size=1 \ + ++policy.generation.mcore_generation_config.moe_router_dtype=fp32 \ + policy.generation.mcore_generation_config.transformer_impl="${MEGATRON_TRANSFORMER_IMPL}" \ + policy.generation.mcore_generation_config.sequence_parallel=true \ + policy.generation.mcore_generation_config.refit_backend=nccl \ + policy.generation.mcore_generation_config.buffer_size_gb=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/test_suites/disabled.txt b/tests/test_suites/disabled.txt index c06a9f253b9..05183d8ad38 100644 --- a/tests/test_suites/disabled.txt +++ b/tests/test_suites/disabled.txt @@ -27,6 +27,19 @@ 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 + +# TODO(@cspades): Multimodal SingleController/TransferQueue recipes awaiting prepared +# fixtures and validation. The 1n4g CLEVR sibling remains enabled as an L1 smoke test. +tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-single-controller-async.v1.sh +tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-1n4g-megatron-single-controller-async.v1.sh +tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-8n4g-megatron-single-controller-async.v1.sh + # Nemotron Super Omni: 16-node topology, and the checkpoint and multimodal Gym # blend are too large to ship with the repo, so these are invoked manually via # examples/nemo_gym/nemotron-3-super-omni/super_omni_launch.sh rather than run diff --git a/tests/test_suites/nightly.txt b/tests/test_suites/nightly.txt index 19edb80e753..f4ddbf7b90e 100644 --- a/tests/test_suites/nightly.txt +++ b/tests/test_suites/nightly.txt @@ -54,6 +54,9 @@ tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-automodel-ep8.v1.s tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8.v1.sh tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-megatron-tp8ep16.v1.sh +# L1 multimodal SingleController/TransferQueue smoke coverage +tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron-single-controller-async.v1.sh + # Functional Qwen3.5-35B VLM GRPO run # The AutoModel variant is re-enabled with the vLLM 0.25.1 bump (no longer hits # https://github.com/vllm-project/vllm/issues/36237). The Megatron variant still 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-single-controller-async.v1.sh b/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron-single-controller-async.v1.sh new file mode 100755 index 00000000000..52e9c9ab9cc --- /dev/null +++ b/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron-single-controller-async.v1.sh @@ -0,0 +1,64 @@ +#!/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. + +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +source "$SCRIPT_DIR/common.env" + +# Two steps exercise CP=1 multimodal generation, TQ policy training, weight +# refit, and a post-refit rollout without treating this smoke test as a +# convergence run. CP>1 + multimodal + TQ remains untested. +# ===== BEGIN CONFIG ===== +NUM_NODES=1 +GPUS_PER_NODE=4 +STEPS_PER_RUN=2 +MAX_STEPS=2 +NUM_RUNS=1 +NUM_MINUTES=120 +# ===== END CONFIG ===== + +exit_if_max_steps_reached + +cd "$PROJECT_ROOT" + +uv run examples/run_grpo_single_controller.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" + +RECORDED_STEP=$(jq -r \ + 'if has("train/loss") then (."train/loss" | keys | map(tonumber) | max // 0) else 0 end' \ + "$JSON_METRICS") +if [[ "$RECORDED_STEP" -lt 1 ]]; then + echo "[ERROR] Expected at least one completed training step" + exit 1 +fi + +uv run tests/check_metrics.py "$JSON_METRICS" \ + 'all_finite(data["train/token_mult_prob_error"])' \ + 'max(data["train/loss"]) < 1000000.0' \ + 'min(data["train/loss"]) > -1000000.0' + +rm -rf "$CKPT_DIR" 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-single-controller-async.v1.sh b/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-single-controller-async.v1.sh new file mode 100755 index 00000000000..48ba6451ce1 --- /dev/null +++ b/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-single-controller-async.v1.sh @@ -0,0 +1,46 @@ +#!/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. + +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +source "$SCRIPT_DIR/common.env" + +# ===== BEGIN CONFIG ===== +NUM_NODES=8 +GPUS_PER_NODE=4 +STEPS_PER_RUN=2 +MAX_STEPS=2 +NUM_RUNS=1 +NUM_MINUTES=120 +# ===== END CONFIG ===== + +exit_if_max_steps_reached + +cd "$PROJECT_ROOT" + +uv run examples/run_grpo_single_controller.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/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-1n4g-megatron-single-controller-async.v1.sh b/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-1n4g-megatron-single-controller-async.v1.sh new file mode 100755 index 00000000000..65703c42e74 --- /dev/null +++ b/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-1n4g-megatron-single-controller-async.v1.sh @@ -0,0 +1,48 @@ +#!/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. + +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +source "$SCRIPT_DIR/common.env" + +# TODO(@cspades): Prepare a self-contained VSTAT fixture and validate this +# driver before moving it from disabled.txt into a recurring suite. +# ===== BEGIN CONFIG ===== +NUM_NODES=1 +GPUS_PER_NODE=4 +STEPS_PER_RUN=2 +MAX_STEPS=2 +NUM_RUNS=1 +NUM_MINUTES=120 +# ===== END CONFIG ===== + +exit_if_max_steps_reached + +cd "$PROJECT_ROOT" + +uv run examples/run_grpo_single_controller.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-vstat-8n4g-megatron-single-controller-async.v1.sh b/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-8n4g-megatron-single-controller-async.v1.sh new file mode 100755 index 00000000000..3f65e783dbc --- /dev/null +++ b/tests/test_suites/vlm/vlm_grpo-nemotron-omni-30ba3b-vstat-8n4g-megatron-single-controller-async.v1.sh @@ -0,0 +1,48 @@ +#!/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. + +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) +source "$SCRIPT_DIR/common.env" + +# TODO(@cspades): Prepare a self-contained VSTAT fixture and validate this +# driver before moving it from disabled.txt into a recurring suite. +# ===== BEGIN CONFIG ===== +NUM_NODES=8 +GPUS_PER_NODE=4 +STEPS_PER_RUN=2 +MAX_STEPS=2 +NUM_RUNS=1 +NUM_MINUTES=120 +# ===== END CONFIG ===== + +exit_if_max_steps_reached + +cd "$PROJECT_ROOT" + +uv run examples/run_grpo_single_controller.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/algorithms/test_grpo.py b/tests/unit/algorithms/test_grpo.py index f2f6f875fe4..3fca2187900 100644 --- a/tests/unit/algorithms/test_grpo.py +++ b/tests/unit/algorithms/test_grpo.py @@ -1236,16 +1236,11 @@ def mock_async_grpo_infrastructure( checkpoint_cut_ordinal=checkpoint_cut_ordinal, ) - # Patch venv creation + # Patch actor runtime environment creation stack.enter_context( patch( - "nemo_rl.algorithms.grpo.create_local_venv_on_each_node", - return_value="/fake/venv", - ) - ) - stack.enter_context( - patch( - "nemo_rl.algorithms.grpo.get_actor_python_env", return_value="/fake/python" + "nemo_rl.algorithms.grpo.make_actor_runtime_env", + return_value={"py_executable": "/fake/python", "env_vars": {}}, ) ) diff --git a/tests/unit/data/datasets/test_audiomcq_dataset.py b/tests/unit/data/datasets/test_audiomcq_dataset.py index 5168dd20012..b95cb9b0886 100644 --- a/tests/unit/data/datasets/test_audiomcq_dataset.py +++ b/tests/unit/data/datasets/test_audiomcq_dataset.py @@ -427,22 +427,25 @@ def test_vlm_hf_data_processor_returns_audiomcq_datum_spec( assert result["extra_env_info"]["ground_truth"] == rows[0]["answer"] assert result["extra_env_info"]["choices"] == rows[0]["choices"] - def test_dispatcher_rejects_unknown_task_name(self): + def test_dispatcher_accepts_generic_preformatted_messages(self): from nemo_rl.data.interfaces import TaskDataSpec from nemo_rl.data.processors import vlm_hf_data_processor - bogus_datum = { - "task_name": "definitely-not-a-task", + generic_datum = { + "task_name": "custom-vlm-task", "messages": [ {"role": "user", "content": [{"type": "text", "text": "hi"}]}, {"role": "assistant", "content": "hello"}, ], } - with pytest.raises(ValueError, match="No data processor for task"): - vlm_hf_data_processor( - datum_dict=bogus_datum, - task_data_spec=TaskDataSpec(task_name="definitely-not-a-task"), - processor=_FakeProcessor(), - max_seq_length=4096, - idx=0, - ) + result = vlm_hf_data_processor( + datum_dict=generic_datum, + task_data_spec=TaskDataSpec(task_name="custom-vlm-task"), + processor=_FakeProcessor(), + max_seq_length=4096, + idx=0, + ) + + assert result["task_name"] == "custom-vlm-task" + assert result["extra_env_info"]["ground_truth"] == "hello" + assert result["vllm_content"] == "fake" diff --git a/tests/unit/data/datasets/test_mmpr_tiny.py b/tests/unit/data/datasets/test_mmpr_tiny.py index 4b0c9e4fc28..7ec4c34a538 100644 --- a/tests/unit/data/datasets/test_mmpr_tiny.py +++ b/tests/unit/data/datasets/test_mmpr_tiny.py @@ -160,7 +160,10 @@ def apply_chat_template(self, messages, **kwargs): for item in content: if isinstance(item, dict) and "text" in item: parts.append(item["text"]) - return " ".join(parts) + formatted_text = " ".join(parts) + if kwargs.get("tokenize"): + return {"input_ids": fake_input_ids} + return formatted_text def __call__(self, text=None, images=None, **kwargs): self.captured_call_text = text @@ -241,6 +244,31 @@ def test_processor_produces_valid_datum_spec(self, tiny_image_path): assert user_message["pixel_values"].pad_to_max_shape is True assert user_message["pixel_values"].as_tensor().dtype == torch.float32 + def test_text_only_row_preserves_formatted_vllm_content(self): + from nemo_rl.data.interfaces import TaskDataSpec + from nemo_rl.data.processors import vlm_hf_data_processor + + processor = _make_stub_nemotron_processor() + task_data_spec = TaskDataSpec(task_name="text-only") + task_data_spec.prompt = "Answer: {}" + + result = vlm_hf_data_processor( + datum_dict={ + "messages": [ + {"role": "user", "content": "What is 2 + 2?"}, + {"role": "assistant", "content": "4"}, + ], + "task_name": "text-only", + }, + task_data_spec=task_data_spec, + processor=processor, + max_seq_length=8192, + idx=0, + ) + + assert result["vllm_content"] == "Answer: What is 2 + 2?" + assert result["vllm_images"] == [] + def test_conversation_preprocessor_is_preserved(self, tiny_image_path): processor = _make_stub_nemotron_processor() processor.conversation_preprocessor = MagicMock( @@ -250,7 +278,7 @@ def test_conversation_preprocessor_is_preserved(self, tiny_image_path): result, _ = _run_processor(tiny_image_path, processor=processor) processor.conversation_preprocessor.assert_called_once() - assert result["vllm_content"] == "preprocessed" + assert result["vllm_content"] is None assert processor.captured_call_text == "preprocessed" def test_historical_tiled_processor_gets_media_metadata(self, tiny_image_path): @@ -285,14 +313,14 @@ def test_historical_tiled_processor_gets_media_metadata(self, tiny_image_path): def test_prompted_text_contains_boxed_literal_and_no_raw_dataset_string( self, tiny_image_path ): - result, _ = _run_processor(tiny_image_path) - vllm_content = result["vllm_content"] + result, processor = _run_processor(tiny_image_path) + processed_text = processor.captured_call_text # Positive: literal \boxed{} must survive prompt formatting - assert "\\boxed{}" in vllm_content + assert "\\boxed{}" in processed_text # Negative: the raw dataset string (with prefix) must NOT leak through - assert _RAW_QUESTION not in vllm_content + assert _RAW_QUESTION not in processed_text def test_placeholder_conversion_exact_string(self, tiny_image_path): """Verify the exact tokenizer input for the placeholder-style processor path. @@ -310,15 +338,15 @@ def test_placeholder_conversion_exact_string(self, tiny_image_path): # The stub's apply_chat_template joins message parts with spaces, # so the captured text passed to __call__ is the chat-templated string. - # Verify the vllm_content (which is apply_chat_template output) matches. - vllm_content = result["vllm_content"] - assert vllm_content == expected_tokenizer_input + # Verify the apply_chat_template output through captured_call_text below; + # placeholder-style processors send expanded token IDs to vLLM. + assert result["vllm_content"] is None # Verify exactly one token in the final output - assert vllm_content.count("") == 1 + assert processor.captured_call_text.count("") == 1 # Verify the question text is present - assert _CLEAN_QUESTION in vllm_content + assert _CLEAN_QUESTION in processor.captured_call_text # Verify the captured __call__ text also matches # (processor.__call__ receives the apply_chat_template output) diff --git a/tests/unit/data/test_collate_fn.py b/tests/unit/data/test_collate_fn.py index e531f4cb5e7..194b798f8e6 100755 --- a/tests/unit/data/test_collate_fn.py +++ b/tests/unit/data/test_collate_fn.py @@ -16,7 +16,11 @@ import torch -from nemo_rl.data.collate_fn import preference_collate_fn +from nemo_rl.data.collate_fn import ( + eval_collate_fn, + preference_collate_fn, + rl_collate_fn, +) from nemo_rl.data.interfaces import DatumSpec from nemo_rl.distributed.batched_data_dict import BatchedDataDict @@ -149,3 +153,23 @@ def test_preference_collate_fn(): assert torch.equal( train_data["input_ids"][1][3:5], torch.tensor([8, 9]) ) # assistant + + +def test_collate_preserves_native_media_when_vllm_content_is_none(): + image = object() + datum = DatumSpec( + message_log=[], + length=1, + loss_multiplier=1.0, + extra_env_info={}, + idx=0, + task_name="vlm", + vllm_content=None, + vllm_images=[image], + vllm_audios=[], + vllm_videos=[], + ) + + for batch in (rl_collate_fn([datum]), eval_collate_fn([datum])): + assert batch["vllm_content"] == [None] + assert batch["vllm_images"] == [[image]] diff --git a/tests/unit/data/test_multimodal_image_encoding.py b/tests/unit/data/test_multimodal_image_encoding.py index fda52a35ca6..4d7f9b23322 100644 --- a/tests/unit/data/test_multimodal_image_encoding.py +++ b/tests/unit/data/test_multimodal_image_encoding.py @@ -17,12 +17,14 @@ import pytest from PIL import Image -import nemo_rl.data.multimodal_utils as multimodal_utils +import nemo_rl.environments.nemo_gym_multimodal as nemo_gym_multimodal from nemo_rl.data.multimodal_utils import ( - encode_images_in_examples, image_to_data_url, resolve_to_image, ) +from nemo_rl.environments.nemo_gym_multimodal import ( + normalize_media_in_examples, +) def _example(*content_parts: dict) -> dict: @@ -59,7 +61,7 @@ def test_resolve_to_image_accepts_file_scheme(tmp_path): assert resolve_to_image(path).size == (5, 6) -def test_encode_images_encodes_local_paths_and_file_urls(tmp_path): +def test_normalize_media_encodes_local_image_paths_and_file_urls(tmp_path): plain = _write_png(tmp_path, "plain.png", (2, 2)) file_url = "file://" + _write_png(tmp_path, "scheme.png", (3, 3)) @@ -70,7 +72,7 @@ def test_encode_images_encodes_local_paths_and_file_urls(tmp_path): {"type": "input_text", "text": "describe"}, ) ] - encode_images_in_examples(examples) + normalize_media_in_examples(examples) parts = examples[0]["responses_create_params"]["input"][0]["content"] assert parts[0]["image_url"].startswith("data:image/png;base64,") @@ -81,7 +83,7 @@ def test_encode_images_encodes_local_paths_and_file_urls(tmp_path): assert parts[2] == {"type": "input_text", "text": "describe"} -def test_encode_images_passes_through_http_and_data_urls(): +def test_normalize_media_passes_through_http_and_data_urls(): data_url = image_to_data_url(Image.new("RGB", (2, 2))) examples = [ _example( @@ -90,7 +92,7 @@ def test_encode_images_passes_through_http_and_data_urls(): {"type": "input_image", "image_url": data_url}, ) ] - encode_images_in_examples(examples) + normalize_media_in_examples(examples) parts = examples[0]["responses_create_params"]["input"][0]["content"] assert parts[0]["image_url"] == "https://example.com/cat.png" @@ -104,7 +106,7 @@ def test_encode_images_deduplicates_sources_and_uses_a_bounded_thread_pool( first = _write_png(tmp_path, "first.png", (2, 3)) second = _write_png(tmp_path, "second.png", (4, 5)) expected = { - source: multimodal_utils._encode_single_image_source(source) + source: nemo_gym_multimodal._encode_single_image_source(source) for source in (first, second) } examples = [ @@ -117,31 +119,117 @@ def test_encode_images_deduplicates_sources_and_uses_a_bounded_thread_pool( ] resolve_calls = [] - original_resolve = multimodal_utils.resolve_to_image + original_resolve = nemo_gym_multimodal.resolve_to_image def tracking_resolve(source): resolve_calls.append(source) return original_resolve(source) observed_max_workers = [] - original_executor = multimodal_utils.ThreadPoolExecutor + original_executor = nemo_gym_multimodal.ThreadPoolExecutor def tracking_executor(*args, **kwargs): observed_max_workers.append(kwargs.get("max_workers")) return original_executor(*args, **kwargs) - monkeypatch.setattr(multimodal_utils, "resolve_to_image", tracking_resolve) - monkeypatch.setattr(multimodal_utils, "ThreadPoolExecutor", tracking_executor) + monkeypatch.setattr(nemo_gym_multimodal, "resolve_to_image", tracking_resolve) + monkeypatch.setattr(nemo_gym_multimodal, "ThreadPoolExecutor", tracking_executor) - encode_images_in_examples(examples) + normalize_media_in_examples(examples) assert Counter(resolve_calls) == Counter({first: 1, second: 1}) - assert observed_max_workers == [multimodal_utils.NEMO_GYM_IMAGE_ENCODE_MAX_WORKERS] + assert observed_max_workers == [ + nemo_gym_multimodal._NEMO_GYM_IMAGE_ENCODE_MAX_WORKERS + ] for example in examples: parts = example["responses_create_params"]["input"][0]["content"] - assert parts[0]["image_url"] == expected[first] - assert parts[1]["image"] == expected[first] - assert parts[2]["url"] == expected[second] + assert parts[0] == { + "type": "input_image", + "image_url": expected[first], + } + assert parts[1] == { + "type": "input_image", + "image_url": expected[first], + } + assert parts[2] == { + "type": "input_image", + "image_url": expected[second], + } + + +@pytest.mark.parametrize( + "part", + [ + {"type": "image", "image": "a.png", "url": "b.png"}, + {"type": "image"}, + ], +) +def test_normalize_media_requires_exactly_one_source_key(part): + with pytest.raises(ValueError, match="requires exactly one"): + normalize_media_in_examples([_example(part)]) + + +@pytest.mark.parametrize( + "source", + ["", {"url": ""}], +) +def test_normalize_media_rejects_empty_urls(source): + with pytest.raises(ValueError, match="requires a non-empty media URL"): + normalize_media_in_examples( + [_example({"type": "input_image", "image_url": source})] + ) + + +def test_normalize_media_preserves_input_image_file_id(): + part = {"type": "input_image", "file_id": "file-123", "detail": "high"} + examples = [_example(part)] + + normalize_media_in_examples(examples) + + assert examples[0]["responses_create_params"]["input"][0]["content"][0] == { + "type": "input_image", + "file_id": "file-123", + "detail": "high", + } + + +def test_normalize_media_promotes_nested_image_detail(): + data_url = image_to_data_url(Image.new("RGB", (2, 2))) + examples = [ + _example( + { + "type": "image", + "image": {"url": data_url, "detail": "high"}, + } + ) + ] + + normalize_media_in_examples(examples) + + assert examples[0]["responses_create_params"]["input"][0]["content"][0] == { + "type": "input_image", + "image_url": data_url, + "detail": "high", + } + + +def test_normalize_media_canonicalizes_local_video(tmp_path, monkeypatch): + video_path = tmp_path / "clip.mp4" + video_path.write_bytes(b"video") + encoded_url = "data:video/mp4;base64,dmlkZW8=" + monkeypatch.setattr( + nemo_gym_multimodal, + "video_path_to_data_url", + lambda source: encoded_url, + ) + examples = [_example({"type": "video", "url": str(video_path)})] + + normalize_media_in_examples(examples) + + assert examples[0]["responses_create_params"]["input"][0]["content"][0] == { + "type": "input_video", + "video_url": encoded_url, + } def test_encode_images_does_not_partially_mutate_on_error(tmp_path): @@ -155,7 +243,7 @@ def test_encode_images_does_not_partially_mutate_on_error(tmp_path): ] with pytest.raises(FileNotFoundError): - encode_images_in_examples(examples) + normalize_media_in_examples(examples) parts = examples[0]["responses_create_params"]["input"][0]["content"] assert parts[0]["image_url"] == existing @@ -164,15 +252,15 @@ def test_encode_images_does_not_partially_mutate_on_error(tmp_path): def test_encode_single_image_source_closes_image_on_success(monkeypatch): image = _CloseTrackingImage() - monkeypatch.setattr(multimodal_utils, "resolve_to_image", lambda _: image) + monkeypatch.setattr(nemo_gym_multimodal, "resolve_to_image", lambda _: image) monkeypatch.setattr( - multimodal_utils, + nemo_gym_multimodal, "image_to_data_url", lambda _: "data:image/png;base64,AA", ) assert ( - multimodal_utils._encode_single_image_source("image.png") + nemo_gym_multimodal._encode_single_image_source("image.png") == "data:image/png;base64,AA" ) assert image.closed @@ -180,27 +268,29 @@ def test_encode_single_image_source_closes_image_on_success(monkeypatch): def test_encode_single_image_source_closes_image_on_error(monkeypatch): image = _CloseTrackingImage() - monkeypatch.setattr(multimodal_utils, "resolve_to_image", lambda _: image) + monkeypatch.setattr(nemo_gym_multimodal, "resolve_to_image", lambda _: image) def fail_to_encode(_): raise RuntimeError("encoding failed") - monkeypatch.setattr(multimodal_utils, "image_to_data_url", fail_to_encode) + monkeypatch.setattr(nemo_gym_multimodal, "image_to_data_url", fail_to_encode) with pytest.raises(RuntimeError, match="encoding failed"): - multimodal_utils._encode_single_image_source("image.png") + nemo_gym_multimodal._encode_single_image_source("image.png") assert image.closed -def test_encode_images_is_a_noop_for_text_only_examples(): +def test_normalize_media_is_a_noop_for_text_only_examples(): examples = [_example({"type": "input_text", "text": "no images here"})] before = [ dict(part) for part in examples[0]["responses_create_params"]["input"][0]["content"] ] - assert encode_images_in_examples(examples) is examples + assert normalize_media_in_examples(examples) is examples assert examples[0]["responses_create_params"]["input"][0]["content"] == before # Missing/oddly-shaped payloads must not raise. - assert encode_images_in_examples([{}, {"responses_create_params": {}}]) is not None - assert encode_images_in_examples([{"responses_create_params": {"input": "nope"}}]) + assert ( + normalize_media_in_examples([{}, {"responses_create_params": {}}]) is not None + ) + assert normalize_media_in_examples([{"responses_create_params": {"input": "nope"}}]) diff --git a/tests/unit/data_plane/test_codec_jagged.py b/tests/unit/data_plane/test_codec_jagged.py index 5dfc57839fc..58743155154 100644 --- a/tests/unit/data_plane/test_codec_jagged.py +++ b/tests/unit/data_plane/test_codec_jagged.py @@ -24,12 +24,14 @@ import torch from tensordict import TensorDict +from nemo_rl.data.multimodal_utils import PackedTensor from nemo_rl.data_plane.codec import ( materialize, pack_jagged_fields, response_from_nested, to_nested_by_length, ) +from nemo_rl.data_plane.schema import PACKED_TENSOR_META_PREFIX from ._rollout_shapes import make_rollout_batch @@ -135,6 +137,20 @@ def test_materialize_jagged_layout_passes_nested_through() -> None: assert bdd["x"].is_nested +def test_materialize_jagged_layout_ignores_forward_pad_target() -> None: + """A token-length target must not be applied to raw nested leaves.""" + padded, lens = _padded([[1, 2], [3, 4, 5]], pad=0) + nested = to_nested_by_length(padded, lens) + td = TensorDict({"x": nested}, batch_size=[2]) + + bdd = materialize(td, layout="jagged", pad_to_seqlen=8) + + assert bdd["x"].is_nested + rows = list(bdd["x"].unbind()) + assert torch.equal(rows[0], torch.tensor([1, 2])) + assert torch.equal(rows[1], torch.tensor([3, 4, 5])) + + def test_materialize_default_pad_value_is_zero() -> None: """No pad_value_dict → fields pad with 0.""" padded, lens = _padded([[1, 2, 3], [4]], pad=0) @@ -257,3 +273,140 @@ def test_pack_jagged_fields_forced_per_token_field_drops_extra_padding() -> None assert torch.equal(out["advantages"][1], advantages[1, :5]) assert torch.equal(out["advantages"][0, 3:], torch.zeros(2)) assert torch.equal(out["extra_2d"], extra) + + +# ── PackedTensor wire invariants ────────────────────────────────────── + + +def test_packed_tensor_wire_round_trip_preserves_arbitrary_pack_dim() -> None: + """Field names are dynamic; metadata preserves rows, dtype, and packing dim.""" + rows = [ + torch.arange(6, dtype=torch.float32).reshape(2, 3), + None, + torch.arange(4, dtype=torch.float32).reshape(2, 2) + 10, + ] + packed = PackedTensor(rows, dim_to_pack=1) + + wire = pack_jagged_fields( + {"pixel_values_flat": packed}, + lengths=torch.tensor([4, 2, 3]), + ) + + metadata_key = f"{PACKED_TENSOR_META_PREFIX}pixel_values_flat" + assert wire["pixel_values_flat"].is_nested + assert torch.equal( + wire[metadata_key], + torch.tensor([[3, 1, 0], [0, 1, 0], [2, 1, 0]]), + ) + + restored = materialize(wire)["pixel_values_flat"] + assert isinstance(restored, PackedTensor) + assert restored.dim_to_pack == 1 + assert restored.pad_to_max_shape is False + assert restored.as_tensor().dtype == torch.float32 + for row_idx, expected in enumerate(rows): + actual = restored.slice([row_idx]).as_tensor() + if expected is None: + assert actual is None + else: + assert torch.equal(actual, expected) + + +def test_packed_tensor_wire_round_trip_preserves_mixed_media_fields() -> None: + """Independent media fields keep their own dtype and reconstruction metadata.""" + pixels = PackedTensor( + [ + torch.arange(24, dtype=torch.float32).reshape(2, 3, 4), + torch.arange(8, dtype=torch.float32).reshape(1, 2, 4), + ], + dim_to_pack=0, + pad_to_max_shape=True, + ) + image_sizes = PackedTensor( + [torch.tensor([[12, 16]]), torch.tensor([[8, 8]])], + dim_to_pack=0, + ) + num_frames = PackedTensor( + [torch.tensor([1], dtype=torch.int32), torch.tensor([2], dtype=torch.int32)], + dim_to_pack=0, + ) + + wire = pack_jagged_fields( + { + "pixel_values": pixels, + "imgs_sizes": image_sizes, + "num_frames": num_frames, + }, + lengths=torch.tensor([4, 3]), + ) + restored = materialize(wire) + + restored_pixels = restored["pixel_values"] + assert isinstance(restored_pixels, PackedTensor) + assert restored_pixels.pad_to_max_shape is True + assert restored_pixels.as_tensor().dtype == torch.float32 + assert restored_pixels.slice([0]).as_tensor().shape == (2, 3, 4) + assert restored_pixels.slice([1]).as_tensor().shape == (1, 3, 4) + assert torch.equal( + restored_pixels.slice([1]).as_tensor()[:, :2], + pixels.slice([1]).as_tensor(), + ) + assert torch.count_nonzero(restored_pixels.slice([1]).as_tensor()[:, 2:]) == 0 + + assert restored["imgs_sizes"].as_tensor().dtype == torch.int64 + assert torch.equal(restored["imgs_sizes"].as_tensor(), image_sizes.as_tensor()) + assert restored["num_frames"].as_tensor().dtype == torch.int32 + assert torch.equal(restored["num_frames"].as_tensor(), num_frames.as_tensor()) + + +def test_materialize_rejects_orphan_packed_tensor_metadata() -> None: + metadata_key = f"{PACKED_TENSOR_META_PREFIX}pixel_values" + td = TensorDict( + {metadata_key: torch.tensor([[1, 0, 0], [1, 0, 0]])}, + batch_size=[2], + ) + + with pytest.raises(ValueError, match="has no payload field"): + materialize(td) + + +def test_materialize_rejects_malformed_packed_tensor_metadata() -> None: + packed = PackedTensor( + [torch.ones(1, 2), torch.ones(2, 2)], + dim_to_pack=0, + ) + wire = pack_jagged_fields({"pixel_values": packed}, lengths=torch.tensor([2, 3])) + metadata_key = f"{PACKED_TENSOR_META_PREFIX}pixel_values" + malformed = TensorDict( + { + "pixel_values": wire["pixel_values"], + metadata_key: torch.ones((2, 2), dtype=torch.long), + }, + batch_size=[2], + ) + + with pytest.raises(ValueError, match=r"shape \[batch, 3\]"): + materialize(malformed) + + +def test_materialized_media_survives_token_microbatch_truncation() -> None: + """Media is reconstructed before token truncation and is not sequence-padded.""" + packed = PackedTensor( + [torch.arange(6, dtype=torch.float32).reshape(2, 3)], + dim_to_pack=0, + ) + wire = pack_jagged_fields( + { + "input_ids": torch.tensor([[1, 2, 3, 0]]), + "pixel_values": packed, + }, + lengths=torch.tensor([3]), + token_aligned_fields=frozenset({"input_ids"}), + ) + + restored = materialize(wire, layout="padded", pad_to_seqlen=8) + restored.truncate_tensors(dim=1, truncated_len=3) + + assert restored["input_ids"].shape == (1, 3) + assert isinstance(restored["pixel_values"], PackedTensor) + assert torch.equal(restored["pixel_values"].as_tensor(), packed.as_tensor()) diff --git a/tests/unit/data_plane/test_kvbatchmeta.py b/tests/unit/data_plane/test_kvbatchmeta.py index a8dc3bc822d..1fbabd6faa9 100644 --- a/tests/unit/data_plane/test_kvbatchmeta.py +++ b/tests/unit/data_plane/test_kvbatchmeta.py @@ -247,6 +247,34 @@ def test_tags_none_when_either_side_missing_in_concat(): assert with_tags.concat(without).tags is None +def test_concat_unions_payload_fields_in_first_seen_order(): + text = KVBatchMeta( + partition_id="p", + task_name="train", + sample_ids=["a"], + fields=["input_ids", "input_lengths"], + ) + multimodal = KVBatchMeta( + partition_id="p", + task_name="train", + sample_ids=["b"], + fields=[ + "input_ids", + "pixel_values", + "__nrl_packed_tensor_meta__pixel_values", + ], + ) + + joined = text.concat(multimodal) + + assert joined.fields == [ + "input_ids", + "input_lengths", + "pixel_values", + "__nrl_packed_tensor_meta__pixel_values", + ] + + # ── Realistic tags from the rollout-shapes helper ── diff --git a/tests/unit/data_plane/test_leader_broadcast.py b/tests/unit/data_plane/test_leader_broadcast.py index 5a74f438c42..eeec85cd392 100644 --- a/tests/unit/data_plane/test_leader_broadcast.py +++ b/tests/unit/data_plane/test_leader_broadcast.py @@ -25,6 +25,7 @@ import torch.distributed as dist import torch.multiprocessing as mp +from nemo_rl.data.multimodal_utils import PackedTensor from nemo_rl.data_plane.worker_mixin import _broadcast_batched_data_dict from nemo_rl.distributed.batched_data_dict import BatchedDataDict @@ -45,6 +46,14 @@ def _worker(rank: int, world_size: int, tmp_init_file: str, q): { "input_ids": torch.arange(12, dtype=torch.long).reshape(3, 4), "input_lengths": torch.tensor([4, 3, 2], dtype=torch.int32), + "pixel_values": PackedTensor( + [ + torch.ones(2, 4), + None, + torch.full((1, 4), 3.0), + ], + dim_to_pack=0, + ), "scalar_meta": "step_42", } ) @@ -61,6 +70,12 @@ def _worker(rank: int, world_size: int, tmp_init_file: str, q): assert torch.equal( out["input_lengths"], torch.tensor([4, 3, 2], dtype=torch.int32) ) + assert isinstance(out["pixel_values"], PackedTensor) + assert out["pixel_values"].logical_segment_counts_by_row() == [1, 0, 1] + assert torch.equal( + out["pixel_values"].as_tensor(), + torch.cat((torch.ones(2, 4), torch.full((1, 4), 3.0))), + ) assert out["scalar_meta"] == "step_42" q.put((rank, "ok")) except Exception as e: # pragma: no cover — surface failures to parent diff --git a/tests/unit/data_plane/test_preshard_extras.py b/tests/unit/data_plane/test_preshard_extras.py index 0c5b9e0d62f..9b5e43029ad 100644 --- a/tests/unit/data_plane/test_preshard_extras.py +++ b/tests/unit/data_plane/test_preshard_extras.py @@ -29,11 +29,12 @@ import torch +from nemo_rl.data.multimodal_utils import PackedTensor from nemo_rl.data_plane import KVBatchMeta from nemo_rl.data_plane.adapters.noop import NoOpDataPlaneClient from nemo_rl.data_plane.column_io import kv_first_write, read_columns from nemo_rl.data_plane.preshard import shard_meta_for_dp -from nemo_rl.data_plane.schema import DP_TRAIN_FIELDS +from nemo_rl.data_plane.schema import DP_TRAIN_FIELDS, packed_tensor_wire_fields from nemo_rl.distributed.batched_data_dict import BatchedDataDict from ._rollout_shapes import ( @@ -94,6 +95,51 @@ def test_kv_first_write_carries_multimodal_extras(): assert fetched["pixel_values"].shape == (4, 3, 4, 4) +def test_kv_first_write_round_trips_mixed_packed_tensor_rows(): + """Packed media crosses the client boundary with text-only rows intact.""" + client = NoOpDataPlaneClient() + client.register_partition( + partition_id="train", + fields=[ + "input_ids", + "input_lengths", + "token_mask", + "sample_mask", + "generation_logprobs", + *packed_tensor_wire_fields(["pixel_values"]), + ], + num_samples=3, + consumer_tasks=["train"], + ) + fb = _final_batch(3) + fb["pixel_values"] = PackedTensor( + [ + torch.arange(8, dtype=torch.float32).reshape(2, 4), + None, + torch.arange(4, dtype=torch.float32).reshape(1, 4) + 20, + ], + dim_to_pack=0, + ) + + meta = kv_first_write( + fb, + sample_ids=["u0", "u1", "u2"], + dp_client=client, + partition_id="train", + ) + out = read_columns( + client, + meta, + select_fields=packed_tensor_wire_fields(["pixel_values"]), + ) + + restored = out["pixel_values"] + assert isinstance(restored, PackedTensor) + assert restored.logical_segment_counts_by_row() == [1, 0, 1] + assert restored.as_tensor().dtype == torch.float32 + assert torch.equal(restored.as_tensor(), fb["pixel_values"].as_tensor()) + + def test_kv_first_write_keys_match_uids_x_ngen(): """Keys round-trip: caller mints ``f"{uid}_g{i}"``, helper preserves them in ``meta.sample_ids`` byte-for-byte.""" diff --git a/tests/unit/distributed/test_worker_groups.py b/tests/unit/distributed/test_worker_groups.py index e1a2363792e..8ceaef0a2be 100644 --- a/tests/unit/distributed/test_worker_groups.py +++ b/tests/unit/distributed/test_worker_groups.py @@ -25,7 +25,32 @@ PY_EXECUTABLES, ) from nemo_rl.distributed.virtual_cluster import RayVirtualCluster -from nemo_rl.distributed.worker_groups import RayWorkerBuilder, RayWorkerGroup +from nemo_rl.distributed.worker_groups import ( + RayWorkerBuilder, + RayWorkerGroup, + _get_initializer_env_vars, +) + + +def test_initializer_env_adds_hf_modules_cache_to_pythonpath(monkeypatch): + monkeypatch.setenv("PYTHONPATH", "/driver/pythonpath") + + result = _get_initializer_env_vars({"HF_MODULES_CACHE": "/hf/modules"}) + + assert result["HF_MODULES_CACHE"] == "/hf/modules" + assert result["PYTHONPATH"].split(os.pathsep) == [ + "/hf/modules", + "/driver/pythonpath", + ] + + +def test_initializer_env_does_not_duplicate_hf_modules_cache(monkeypatch): + monkeypatch.setenv("HF_MODULES_CACHE", "/hf/modules") + monkeypatch.setenv("PYTHONPATH", f"/project{os.pathsep}/hf/modules") + + result = _get_initializer_env_vars({}) + + assert result["PYTHONPATH"] == f"/project{os.pathsep}/hf/modules" @ray.remote diff --git a/tests/unit/environments/test_nemo_gym.py b/tests/unit/environments/test_nemo_gym.py index 8f5a8e510ec..eefd569b72a 100644 --- a/tests/unit/environments/test_nemo_gym.py +++ b/tests/unit/environments/test_nemo_gym.py @@ -32,6 +32,7 @@ MULTIMODAL_CONTENT_TYPES, PackedTensor, image_to_data_url, + video_path_to_data_url, ) from nemo_rl.data.utils import setup_response_data from nemo_rl.distributed.ray_actor_environment_registry import ( @@ -45,12 +46,16 @@ setup_nemo_gym_config, validate_reward_components_match_scalar, ) -from nemo_rl.environments.nemo_gym_video import ( +from nemo_rl.environments.nemo_gym_multimodal import ( _extract_static_video_messages, _inject_vllm_mm_processor_kwargs, - _metadata_extra_body, nemo_gym_example_to_video_datum_spec, - normalize_video_urls_in_examples, + normalize_media_in_examples, +) +from nemo_rl.environments.nemo_gym_request import ( + _chat_template_kwargs_for_processor, + _deep_merge_dict, + _metadata_extra_body, ) from nemo_rl.environments.nemotron_utils import ( _expand_nemotron_video_placeholders, @@ -184,7 +189,7 @@ def test_extract_static_video_message_ignores_still_image_only_row(): assert _extract_static_video_messages(example) is None -def test_gym_local_video_path_is_normalized_to_file_url(tmp_path): +def test_gym_local_video_path_is_inlined_as_data_url(tmp_path): video_path = tmp_path / "clip with spaces.mp4" video_path.write_bytes(b"test") examples = [ @@ -205,14 +210,37 @@ def test_gym_local_video_path_is_normalized_to_file_url(tmp_path): } ] - normalize_video_urls_in_examples(examples) + normalize_media_in_examples(examples) - assert ( - examples[0]["responses_create_params"]["input"][0]["content"][0]["video_url"][ - "url" - ] - == video_path.resolve().as_uri() - ) + video_url = examples[0]["responses_create_params"]["input"][0]["content"][0][ + "video_url" + ] + assert video_url.startswith("data:video/mp4;base64,") + + +def test_video_path_to_data_url_rejects_unsupported_and_missing_paths(tmp_path): + """Bad local video sources must fail loudly rather than inline garbage.""" + unsupported = tmp_path / "clip.gif" + unsupported.write_bytes(b"test") + with pytest.raises(ValueError, match="Unsupported video extension"): + video_path_to_data_url(str(unsupported)) + + with pytest.raises(FileNotFoundError, match="does not exist"): + video_path_to_data_url(str(tmp_path / "missing.mp4")) + + +def test_video_path_to_data_url_passes_through_data_urls_and_accepts_file_scheme( + tmp_path, +): + already_inlined = "data:video/mp4;base64,dG95" + assert video_path_to_data_url(already_inlined) == already_inlined + + video_path = tmp_path / "clip.mp4" + video_path.write_bytes(b"toy-video") + from_plain = video_path_to_data_url(str(video_path)) + from_scheme = video_path_to_data_url(f"file://{video_path}") + assert from_plain.startswith("data:video/mp4;base64,") + assert from_plain == from_scheme def test_extract_static_video_message_rejects_multiple_videos(tmp_path): @@ -319,6 +347,52 @@ def test_video_metadata_rejects_invalid_extra_body(extra_body): _metadata_extra_body(example) +@pytest.mark.parametrize( + "chat_template_kwargs", + [ + {"enable_thinking": False}, + '{"enable_thinking": false}', + ], +) +def test_chat_template_kwargs_for_processor_accepts_mapping_or_json( + chat_template_kwargs, +): + example = { + "responses_create_params": { + "metadata": {"chat_template_kwargs": chat_template_kwargs} + } + } + + assert _chat_template_kwargs_for_processor(example) == { + "chat_template_kwargs": {"enable_thinking": False}, + "enable_thinking": False, + } + + +def test_chat_template_kwargs_for_processor_defaults_to_empty(): + assert _chat_template_kwargs_for_processor({}) == {} + + +def test_chat_template_kwargs_for_processor_rejects_invalid_json(): + example = { + "responses_create_params": {"metadata": {"chat_template_kwargs": "not-json"}} + } + + with pytest.raises(ValueError, match="chat_template_kwargs"): + _chat_template_kwargs_for_processor(example) + + +def test_deep_merge_dict_merges_nested_values_without_mutating_inputs(): + base = {"nested": {"left": 1}, "unchanged": [1]} + update = {"nested": {"right": 2}, "unchanged": [2]} + + merged = _deep_merge_dict(base, update) + + assert merged == {"nested": {"left": 1, "right": 2}, "unchanged": [2]} + assert base == {"nested": {"left": 1}, "unchanged": [1]} + assert update == {"nested": {"right": 2}, "unchanged": [2]} + + def test_video_metadata_canonicalizes_mapping_extra_body_to_json_string(): example = { "responses_create_params": { @@ -380,7 +454,7 @@ def test_video_datum_uses_temporal_processor_contract(monkeypatch, tmp_path): frames = np.zeros((4, 8, 8, 3), dtype=np.uint8) monkeypatch.setattr( - "nemo_rl.environments.nemo_gym_video.load_video_frames_with_metadata", + "nemo_rl.environments.nemo_gym_multimodal.load_video_frames_with_metadata", lambda *args, **kwargs: ( frames, {"frames_indices": [0, 3, 6, 9], "fps": 3.0}, @@ -513,7 +587,7 @@ def fake_video_processor( } monkeypatch.setattr( - "nemo_rl.environments.nemo_gym_video.nemo_gym_example_to_video_datum_spec", + "nemo_rl.environments.nemo_gym_multimodal.nemo_gym_example_to_video_datum_spec", fake_video_processor, ) processor = SimpleNamespace( @@ -574,7 +648,7 @@ def test_video_datum_uses_cached_frames_without_decoding_video(monkeypatch, tmp_ } } monkeypatch.setattr( - "nemo_rl.environments.nemo_gym_video._video_to_image_content", + "nemo_rl.environments.nemo_gym_multimodal._video_to_image_content", lambda *args, **kwargs: pytest.fail("cached frames must not decode the video"), ) @@ -661,7 +735,7 @@ def test_nemotron_video_datum_uses_dynamic_tubelet_inputs(monkeypatch, tmp_path) } frames = np.zeros((4, 8, 16, 3), dtype=np.uint8) monkeypatch.setattr( - "nemo_rl.environments.nemo_gym_video.load_video_frames_with_metadata", + "nemo_rl.environments.nemo_gym_multimodal.load_video_frames_with_metadata", lambda *args, **kwargs: ( frames, {"frames_indices": [0, 3, 6, 9], "fps": 3.0}, @@ -819,11 +893,11 @@ def fake_manifest_builder(paths): return "data:video/x-nemo-rl-cached-frames;base64,dGVzdA==" monkeypatch.setattr( - "nemo_rl.environments.nemo_gym_video.build_cached_video_frame_data_url", + "nemo_rl.environments.nemo_gym_multimodal.build_cached_video_frame_data_url", fake_manifest_builder, ) monkeypatch.setattr( - "nemo_rl.environments.nemo_gym_video.process_nemotron_video_frames", + "nemo_rl.environments.nemo_gym_multimodal.process_nemotron_video_frames", lambda *args, **kwargs: { "input_ids": torch.tensor([[7, 18, 18, 9]]), "pixel_values": torch.ones(4, 3, 8, 8), @@ -1480,7 +1554,7 @@ class _RolloutCollectionHelper: def run_examples(self, examples, head_server_config): del head_server_config content = examples[0]["responses_create_params"]["input"][0]["content"] - assert content[0]["video_url"] == video_path.resolve().as_uri() + assert content[0]["video_url"].startswith("data:video/mp4;base64,") assert content[1]["image_url"].startswith("data:image/png;base64,") async def _completed_result(): @@ -1530,6 +1604,116 @@ def _postprocess_nemo_gym_to_nemo_rl_result( asyncio.run(_run()) +@pytest.mark.parametrize("modality", ["image", "video"]) +def test_nemo_gym_megatron_multimodal_response_round_trip(tmp_path, modality): + """Round-trip normalized media and a mocked Megatron HTTP response through Gym.""" + + async def _run(): + if modality == "image": + media_path = tmp_path / "clevr.png" + Image.new("RGB", (2, 2), color="red").save(media_path) + media_part = {"type": "input_image", "image_url": str(media_path)} + expected_prefix = "data:image/png;base64," + else: + media_path = tmp_path / "vstat.mp4" + media_path.write_bytes(b"toy-video") + media_part = {"type": "input_video", "video_url": str(media_path)} + expected_prefix = "data:video/mp4;base64," + + row = { + "_rowidx": 3, + "agent_ref": {"name": "mock-megatron-agent"}, + "responses_create_params": { + "input": [ + { + "role": "user", + "content": [ + media_part, + {"type": "input_text", "text": "What is shown?"}, + ], + } + ] + }, + } + + class _Tokenizer: + def batch_decode(self, batches): + return [" ".join(map(str, token_ids)) for token_ids in batches] + + class _RolloutCollectionHelper: + def run_examples(self, examples, head_server_config): + assert head_server_config.backend == "megatron" + dispatched_row = examples[0] + dispatched_part = dispatched_row["responses_create_params"]["input"][0][ + "content" + ][0] + media_url = dispatched_part[f"{modality}_url"] + assert media_url.startswith(expected_prefix) + + mocked_result = { + "responses_create_params": { + "input": deepcopy( + dispatched_row["responses_create_params"]["input"] + ) + }, + "response": { + "agent_input": deepcopy( + dispatched_row["responses_create_params"]["input"] + ), + "output": [ + { + "type": "message", + "prompt_token_ids": [10, 99, 20], + "generation_token_ids": [71, 72], + "generation_log_probs": [-0.25, -0.5], + } + ], + }, + } + + async def _completed_result(): + return dispatched_row, mocked_result + + return [_completed_result()] + + class _MockSelf: + cfg = {} + rch = _RolloutCollectionHelper() + head_server_config = SimpleNamespace(backend="megatron") + _tokenizer = _Tokenizer() + _processor = None + # Bind the real postprocess: the assertions below are about its + # message_log output, not about run_rollouts' dispatch alone. + _postprocess_nemo_gym_to_nemo_rl_result = NemoGym.__ray_metadata__.modified_class._postprocess_nemo_gym_to_nemo_rl_result + + def _require_spinup(self): + pass + + streamed = [] + async for item in NemoGym.__ray_metadata__.modified_class.run_rollouts( + _MockSelf(), [row], "test" + ): + streamed.append(item) + + row_index, result, _metrics = streamed[0] + assert row_index == 3 + assert [message["role"] for message in result["message_log"]] == [ + "user", + "assistant", + ] + assert result["message_log"][0]["token_ids"].tolist() == [10, 99, 20] + assert result["message_log"][1]["token_ids"].tolist() == [71, 72] + assert result["message_log"][1]["generation_logprobs"].tolist() == [ + -0.25, + -0.5, + ] + assert result["full_result"]["response"]["output"][0]["generation_str"] == ( + "71 72" + ) + + asyncio.run(_run()) + + def test_nemo_gym_postprocess_no_generation_data_raises(): """When no output item carries generation data, the postprocess should raise a ValueError that reports the prompt length and the response.output item types.""" diff --git a/tests/unit/environments/test_nemo_gym_image_placeholders.py b/tests/unit/environments/test_nemo_gym_image_placeholders.py index 73359f9d4c4..5a05ddf366d 100644 --- a/tests/unit/environments/test_nemo_gym_image_placeholders.py +++ b/tests/unit/environments/test_nemo_gym_image_placeholders.py @@ -16,10 +16,10 @@ import torch from PIL import Image -from nemo_rl.environments.nemo_gym import _attach_multimodal_data_to_user_message +from nemo_rl.data.multimodal_utils import attach_image_model_inputs_to_message # -------------------------------------------------------------------------- -# ragged pixel_values path in _attach_multimodal_data_to_user_message +# ragged pixel_values path in attach_image_model_inputs_to_message # -------------------------------------------------------------------------- @@ -72,7 +72,7 @@ def test_ragged_output_requested_only_for_multi_image_turns(): """The ragged switch needs both the flag and more than one image.""" for count, flag, expected in [(2, True, None), (1, True, "pt"), (2, False, "pt")]: processor = NemotronNanoVLV2Processor(torch.zeros(count, 3, 4, 4)) - _attach_multimodal_data_to_user_message( + attach_image_model_inputs_to_message( {}, images=_images(count), processor=processor, @@ -87,7 +87,7 @@ def test_ragged_pixel_values_are_padded_to_one_tensor(): """Heterogeneous CHW tensors become a single padded tensor for the message.""" processor = _ragged((3, 2, 4), (3, 6, 4)) user_message: dict = {} - _attach_multimodal_data_to_user_message( + attach_image_model_inputs_to_message( user_message, images=_images(2), processor=processor, @@ -103,7 +103,7 @@ def test_ragged_pixel_values_are_padded_to_one_tensor(): def test_ragged_pixel_values_reject_non_chw_entries(): processor = _ragged((3, 2, 4), (2, 4)) with pytest.raises(ValueError, match="one CHW tensor per image"): - _attach_multimodal_data_to_user_message( + attach_image_model_inputs_to_message( {}, images=_images(2), processor=processor, @@ -114,7 +114,7 @@ def test_ragged_pixel_values_reject_non_chw_entries(): def test_ragged_pixel_values_reject_mixed_channel_counts(): processor = _ragged((3, 2, 4), (1, 2, 4)) with pytest.raises(ValueError, match="same channel count"): - _attach_multimodal_data_to_user_message( + attach_image_model_inputs_to_message( {}, images=_images(2), processor=processor, @@ -124,10 +124,10 @@ def test_ragged_pixel_values_reject_mixed_channel_counts(): def test_attach_is_a_noop_without_images_or_processor(): user_message: dict = {} - _attach_multimodal_data_to_user_message( + attach_image_model_inputs_to_message( user_message, images=[], processor=NemotronNanoVLV2Processor(None) ) - _attach_multimodal_data_to_user_message( + attach_image_model_inputs_to_message( user_message, images=_images(1), processor=None ) assert user_message == {} diff --git a/tests/unit/environments/test_nemo_gym_mm_utils.py b/tests/unit/environments/test_nemo_gym_mm_utils.py index 3d8a4186e17..6f2a785b8ce 100644 --- a/tests/unit/environments/test_nemo_gym_mm_utils.py +++ b/tests/unit/environments/test_nemo_gym_mm_utils.py @@ -14,10 +14,14 @@ from PIL import Image -from nemo_rl.data.multimodal_utils import image_to_data_url -from nemo_rl.environments.nemo_gym import ( +from nemo_rl.data.multimodal_utils import ( + extract_input_media_sources_from_responses_messages, + image_to_data_url, +) +from nemo_rl.environments.nemo_gym_multimodal import ( _extract_input_images_from_message, _index_per_turn_images, + _without_initial_media_sources, ) @@ -176,3 +180,67 @@ def test_index_per_turn_images_flushes_on_function_call_trainable_item(): assert len(per_turn) == 2 assert [img.size for img in per_turn[0]] == [(2, 2)] assert [img.size for img in per_turn[1]] == [(5, 5)] + + +def test_without_initial_media_sources_strips_videos_and_images_in_order(): + """Video parts must be de-duplicated alongside images, in encounter order.""" + image_url = _image((2, 2)) + video_url = "data:video/mp4;base64,dG95" + messages = [ + { + "role": "user", + "content": [ + {"type": "input_video", "video_url": video_url}, + {"type": "input_image", "image_url": image_url}, + {"type": "input_text", "text": "What is shown?"}, + ], + } + ] + initial_sources = extract_input_media_sources_from_responses_messages(messages) + assert initial_sources == [("video", video_url), ("image", image_url)] + + filtered, fully_consumed = _without_initial_media_sources(messages, initial_sources) + + assert fully_consumed is True + assert filtered[0]["content"] == [{"type": "input_text", "text": "What is shown?"}] + # The caller's messages must not be mutated in place. + assert len(messages[0]["content"]) == 3 + + +def test_without_initial_media_sources_keeps_media_the_agent_added(): + """Only the ordered prefix of initial sources is removed; extras survive.""" + initial_image = _image((2, 2)) + agent_image = _image((4, 4)) + messages = [ + { + "role": "user", + "content": [{"type": "input_image", "image_url": initial_image}], + }, + { + "role": "user", + "content": [{"type": "input_image", "image_url": agent_image}], + }, + ] + + filtered, fully_consumed = _without_initial_media_sources( + messages, [("image", initial_image)] + ) + + assert fully_consumed is True + assert filtered[0]["content"] == [] + assert filtered[1]["content"] == [{"type": "input_image", "image_url": agent_image}] + + +def test_without_initial_media_sources_reports_unconsumed_sources(): + """A source that never appears leaves the consumed flag False.""" + filtered, fully_consumed = _without_initial_media_sources( + [{"role": "user", "content": [{"type": "input_text", "text": "hi"}]}], + [("image", "data:image/png;base64,AA")], + ) + + assert fully_consumed is False + assert filtered[0]["content"] == [{"type": "input_text", "text": "hi"}] + + +def test_without_initial_media_sources_passes_through_non_list_messages(): + assert _without_initial_media_sources("not-a-list", []) == ("not-a-list", False) diff --git a/tests/unit/experience/test_payload.py b/tests/unit/experience/test_payload.py index 5029d3ebb63..f506e42fb80 100644 --- a/tests/unit/experience/test_payload.py +++ b/tests/unit/experience/test_payload.py @@ -16,6 +16,9 @@ import torch +from nemo_rl.data.multimodal_utils import PackedTensor +from nemo_rl.data_plane.codec import materialize +from nemo_rl.data_plane.schema import PACKED_TENSOR_META_PREFIX from nemo_rl.experience.interfaces import Completion, PromptGroupRecord from nemo_rl.experience.payload import pack_payload, record_to_train_batch @@ -162,6 +165,37 @@ def test_record_to_train_batch_omits_routed_experts_when_absent() -> None: assert "routed_experts" not in fields +def test_multimodal_packed_tensor_round_trips_through_tq_payload() -> None: + completions = [ + _completion(route_start=10, reward=1.0, with_routes=False), + _completion(route_start=30, reward=2.0, with_routes=False), + ] + media = torch.arange(8, dtype=torch.float32).reshape(2, 4) + completions[0].message_log[0]["pixel_values"] = PackedTensor(media, dim_to_pack=0) + + train_batch = record_to_train_batch( + _record(completions), + pad_value_dict={"token_ids": 0, "input_ids": 0}, + ) + assert isinstance(train_batch["pixel_values"], PackedTensor) + + _, fields, _ = pack_payload( + train_batch, + weight_version=3, + group_id="group", + prompt_idx=17, + ) + assert "pixel_values" in fields + assert f"{PACKED_TENSOR_META_PREFIX}pixel_values" in fields + + restored = materialize(fields) + restored_media = restored["pixel_values"] + assert isinstance(restored_media, PackedTensor) + assert len(restored_media) == 2 + assert restored_media.logical_segment_counts_by_row() == [1, 0] + assert torch.equal(restored_media.as_tensor(), media) + + def _failed_completion() -> Completion: """A trajectory whose first generation raised: prompt only, no routes.""" return Completion( diff --git a/tests/unit/experience/test_rollout_manager.py b/tests/unit/experience/test_rollout_manager.py index 9553910460b..7997897b00d 100644 --- a/tests/unit/experience/test_rollout_manager.py +++ b/tests/unit/experience/test_rollout_manager.py @@ -29,6 +29,7 @@ import tempfile import uuid from copy import deepcopy +from types import SimpleNamespace import pytest import torch @@ -40,11 +41,13 @@ from nemo_rl.data.collate_fn import rl_collate_fn from nemo_rl.data.datasets.response_datasets import NemoGymDataset from nemo_rl.data.interfaces import DatumSpec +from nemo_rl.data.multimodal_utils import PackedTensor from nemo_rl.data.processors import nemo_gym_data_processor from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.experience.interfaces import Completion, PromptGroupRecord from nemo_rl.experience.rollout_manager import ( AsyncNemoGymRolloutImpl, + AsyncRolloutImpl, RolloutManager, RolloutOutcome, RolloutRetryPolicy, @@ -86,6 +89,55 @@ async def apply(): return _run(apply()) +def test_generate_response_forwards_message_log_media_to_generation() -> None: + captured: dict[str, BatchedDataDict] = {} + + class _Generation: + async def generate_async(self, data): + captured["data"] = data + input_len = int(data["input_lengths"][0]) + yield 0, BatchedDataDict( + { + "output_ids": torch.cat( + (data["input_ids"], torch.tensor([[42]])), dim=1 + ), + "unpadded_sequence_lengths": torch.tensor([input_len + 1]), + "logprobs": torch.zeros(1, input_len + 1), + } + ) + + manager = object.__new__(AsyncRolloutImpl) + manager._policy_generation = _Generation() + manager._tokenizer = SimpleNamespace( + pad_token_id=0, + decode=lambda *_args, **_kwargs: "answer", + ) + manager._timeouts = SimpleNamespace(generation_s=10.0) + pixel_values = PackedTensor(torch.ones(2, 3, 4, 4), dim_to_pack=0) + imgs_sizes = PackedTensor(torch.tensor([[4, 4], [4, 4]]), dim_to_pack=0) + message_log = [ + { + "role": "user", + "content": "image", + "token_ids": torch.tensor([1, 2, 3]), + "pixel_values": pixel_values, + "imgs_sizes": imgs_sizes, + } + ] + + _run(manager._generate_response(message_log, None)) + + generation_data = captured["data"] + assert isinstance(generation_data["pixel_values"], PackedTensor) + assert isinstance(generation_data["imgs_sizes"], PackedTensor) + assert torch.equal( + generation_data["pixel_values"].as_tensor(), pixel_values.as_tensor() + ) + assert torch.equal( + generation_data["imgs_sizes"].as_tensor(), imgs_sizes.as_tensor() + ) + + class _FakeBuffer: """Minimal TQReplayBuffer stand-in that records reserve/commit calls.""" diff --git a/tests/unit/models/generation/test_megatron_generation.py b/tests/unit/models/generation/test_megatron_generation.py index 95bce59081d..f996690deef 100644 --- a/tests/unit/models/generation/test_megatron_generation.py +++ b/tests/unit/models/generation/test_megatron_generation.py @@ -22,12 +22,17 @@ from nemo_rl.algorithms.grpo import refit_policy_generation from nemo_rl.algorithms.utils import get_tokenizer +from nemo_rl.data.multimodal_utils import PackedTensor from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.distributed.virtual_cluster import RayVirtualCluster from nemo_rl.models.generation.megatron import MegatronGeneration, megatron_generation from nemo_rl.models.generation.megatron.config import ( dedicated_inference_megatron_cfg, ) +from nemo_rl.models.generation.megatron.megatron_worker import MegatronGenerationMixin +from nemo_rl.models.generation.megatron.utils import ( + build_prompt_and_multimodal_data, +) from nemo_rl.models.policy import PolicyConfig from nemo_rl.models.policy.lm_policy import Policy from nemo_rl.weight_sync.megatron_weight_synchronizer import ( @@ -37,6 +42,209 @@ model_name = "Qwen/Qwen3-0.6B" + +@pytest.mark.mcore +def test_multimodal_preprocessing_requires_policy_processor(): + class _ImageWrapper: + supports_image = True + + worker = object.__new__(MegatronGenerationMixin) + worker._get_megatron_inference_wrapper_cls = lambda: _ImageWrapper + + with pytest.raises(ValueError, match="requires the policy processor"): + worker._build_image_preprocessing_config({}) + + +@pytest.mark.mcore +def test_multimodal_preprocessing_forwards_vision_model_type(): + class _ImageWrapper: + supports_image = True + + worker = object.__new__(MegatronGenerationMixin) + worker._get_megatron_inference_wrapper_cls = lambda: _ImageWrapper + worker.processor = SimpleNamespace( + image_processor=SimpleNamespace( + patch_size=14, + min_num_patches=1, + max_num_patches=32, + norm_mean=[0.1, 0.2, 0.3], + norm_std=[0.4, 0.5, 0.6], + ) + ) + + config = worker._build_image_preprocessing_config({"vision_model_type": "qwen-vl"}) + + assert config.vision_model_type == "qwen-vl" + + +@pytest.mark.mcore +def test_direct_megatron_media_request_preserves_preexpanded_prompt(): + def fake_sample_vision_tensors(data, index): + return torch.ones(1, 2, 4), torch.tensor([[2, 2]]), None + + data = { + "input_ids": torch.tensor([[10, 99, 99, 20, 0]]), + "input_lengths": torch.tensor([4]), + } + + prompt, multi_modal_data = build_prompt_and_multimodal_data( + data, + 0, + sample_tensors=fake_sample_vision_tensors, + supports_modality=lambda modality: modality == "image", + ) + + assert prompt == [10, 99, 99, 20] + assert multi_modal_data["media_tokens_preexpanded"] is True + assert "image" in multi_modal_data + + +@pytest.mark.mcore +def test_text_only_request_does_not_resolve_multimodal_capabilities(): + data = { + "input_ids": torch.tensor([[10, 20, 0]]), + "input_lengths": torch.tensor([2]), + } + + prompt, multi_modal_data = build_prompt_and_multimodal_data( + data, + 0, + supports_modality=lambda modality: pytest.fail( + f"unexpected capability lookup for {modality}" + ), + ) + + assert prompt == [10, 20] + assert multi_modal_data is None + + +@pytest.mark.mcore +def test_direct_megatron_video_request_marks_preexpanded_prompt(): + def fake_sample_vision_tensors(data, index): + return ( + torch.ones(1, 4, 4), + torch.tensor([[2, 2], [2, 2], [2, 2], [2, 2]]), + torch.tensor([4]), + ) + + data = { + "input_ids": torch.tensor([[10, 99, 99, 20]]), + "input_lengths": torch.tensor([4]), + } + + prompt, multi_modal_data = build_prompt_and_multimodal_data( + data, + 0, + sample_tensors=fake_sample_vision_tensors, + supports_modality=lambda modality: modality == "video", + ) + + assert prompt == [10, 99, 99, 20] + assert multi_modal_data["media_tokens_preexpanded"] is True + assert "video" in multi_modal_data + + +@pytest.mark.mcore +@pytest.mark.parametrize( + ("modality", "num_frames"), + [("image", torch.tensor([1])), ("video", torch.tensor([4]))], +) +def test_direct_megatron_multimodal_generate_round_trip( + monkeypatch, modality, num_frames +): + """Exercise RL request construction and response packing around a mocked MCore LLM.""" + + class _MultimodalWrapper: + supports_text = True + supports_image = True + supports_video = True + supports_audio = False + + worker = object.__new__(MegatronGenerationMixin) + worker.cfg = { + "generation": { + "temperature": 1.0, + "top_k": None, + "top_p": 1.0, + "max_new_tokens": 2, + "stop_strings": None, + "mcore_generation_config": {}, + } + } + worker.tokenizer = SimpleNamespace(pad_token_id=0) + worker.megatron_tokenizer = SimpleNamespace(eod=2) + worker._inference_loop = object() + worker._get_megatron_inference_wrapper_cls = lambda: _MultimodalWrapper + + frame_count = int(num_frames.sum()) + pixels = torch.arange(frame_count * 12, dtype=torch.float32).reshape( + frame_count, 3, 2, 2 + ) + sizes = torch.tensor([[2, 2]] * frame_count) + data = BatchedDataDict( + { + "input_ids": torch.tensor([[10, 99, 99, 20]]), + "input_lengths": torch.tensor([4]), + "pixel_values": PackedTensor([pixels], dim_to_pack=0), + "imgs_sizes": PackedTensor([sizes], dim_to_pack=0), + "num_frames": PackedTensor([num_frames], dim_to_pack=0), + } + ) + + captured = {} + mocked_call = object() + + def mock_generate(prompts, multi_modal_data, sampling_params): + captured.update( + prompts=prompts, + multi_modal_data=multi_modal_data, + sampling_params=sampling_params, + ) + return mocked_call + + replies = [ + SimpleNamespace( + prompt_tokens=torch.tensor([10, 99, 99, 20]), + generated_tokens=[71, 72], + generated_log_probs=[-0.25, -0.5], + ) + ] + worker._generate_with_persistent_engine = mock_generate + + class _CompletedFuture: + def result(self): + return replies + + def mock_run_coroutine_threadsafe(call, loop): + assert call is mocked_call + assert loop is worker._inference_loop + return _CompletedFuture() + + monkeypatch.setattr( + "nemo_rl.models.generation.megatron.megatron_worker.asyncio.run_coroutine_threadsafe", + mock_run_coroutine_threadsafe, + ) + + output = worker.generate(data=data) + + assert captured["prompts"] == [[10, 99, 99, 20]] + media = captured["multi_modal_data"][0] + assert media["media_tokens_preexpanded"] is True + assert set(media) == {modality, "media_tokens_preexpanded"} + assert torch.equal(media[modality]["imgs"], pixels) + assert torch.equal(media[modality]["imgs_sizes"], sizes) + if modality == "video": + assert torch.equal(media["video"]["num_frames"], num_frames.to(torch.int32)) + else: + assert "num_frames" not in media["image"] + assert captured["sampling_params"][0].return_prompt_tokens is True + + assert output["output_ids"][0].tolist() == [10, 99, 99, 20, 71, 72] + assert output["logprobs"][0].tolist() == [0.0, 0.0, 0.0, 0.0, -0.25, -0.5] + assert output["generation_lengths"].tolist() == [2] + assert output["unpadded_sequence_lengths"].tolist() == [6] + + basic_megatron_test_config: PolicyConfig = { "model_name": model_name, "tokenizer": {"name": model_name}, @@ -619,6 +827,12 @@ def test_megatron_generation_non_colocated_refit( cluster=generation_cluster, skip_weight_load=skip_weight_load, ) + assert mg._policy_config is not config + assert mg._policy_config["generation"] is not config["generation"] + assert ( + mg._policy_config["generation"]["mcore_generation_config"] + is not config["generation"]["mcore_generation_config"] + ) # Wire the refit collective the way grpo.setup does: through the # weight synchronizer, which refit_policy_generation dispatches to. diff --git a/tests/unit/models/generation/test_megatron_generation_parse.py b/tests/unit/models/generation/test_megatron_generation_parse.py index 6db700ef237..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 diff --git a/tests/unit/models/generation/test_megatron_generation_utils.py b/tests/unit/models/generation/test_megatron_generation_utils.py new file mode 100644 index 00000000000..f53c5ae0d0d --- /dev/null +++ b/tests/unit/models/generation/test_megatron_generation_utils.py @@ -0,0 +1,176 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from types import SimpleNamespace + +import pytest + +from nemo_rl.models.generation.megatron.utils import ( + build_image_preprocessing_config, + build_video_preprocessing_config, +) + +pytestmark = pytest.mark.mcore + + +def _image_processor(**overrides): + fields = { + "patch_size": 14, + "min_num_patches": 1, + "max_num_patches": 32, + "norm_mean": [0.1, 0.2, 0.3], + "norm_std": [0.4, 0.5, 0.6], + } + fields.update(overrides) + return SimpleNamespace(**fields) + + +def test_build_image_config_handles_dict_patch_and_downsample_ratio(): + config = build_image_preprocessing_config( + _image_processor( + patch_size={"height": 16, "width": 16}, + downsample_ratio=0.5, + ), + dynamic_resolution=True, + vision_model_type="qwen-vl", + ) + + assert config.patch_dim == 16 + assert config.dynamic_resolution is True + assert config.vision_model_type == "qwen-vl" + assert config.pixel_shuffle is True + assert config.spatial_merge_size == 2 + assert config.dynamic_resolution_min_patches == 1 + assert config.dynamic_resolution_max_patches == 32 + assert config.pixel_mean == [0.1, 0.2, 0.3] + assert config.pixel_std == [0.4, 0.5, 0.6] + + +@pytest.mark.parametrize( + ("merge_fields", "expected_merge_size"), + [ + ({"merge_size": 4}, 4), + ({"spatial_merge_size": 3}, 3), + ({}, 1), + ], +) +def test_build_image_config_merge_size_fallbacks(merge_fields, expected_merge_size): + config = build_image_preprocessing_config(_image_processor(**merge_fields)) + + assert config.spatial_merge_size == expected_merge_size + assert config.pixel_shuffle is (expected_merge_size > 1) + + +def test_build_image_config_accepts_alternate_field_names(): + config = build_image_preprocessing_config( + SimpleNamespace( + patch_dim=12, + min_num_patches=2, + max_num_patches=24, + image_mean=(0.1, 0.2, 0.3), + image_std=(0.7, 0.8, 0.9), + ) + ) + + assert config.patch_dim == 12 + assert config.dynamic_resolution_min_patches == 2 + assert config.dynamic_resolution_max_patches == 24 + assert config.pixel_mean == [0.1, 0.2, 0.3] + assert config.pixel_std == [0.7, 0.8, 0.9] + + +def test_build_image_config_error_names_all_missing_fields(): + with pytest.raises(ValueError) as exc_info: + build_image_preprocessing_config(SimpleNamespace()) + + for field in ( + "patch_size", + "min_num_patches", + "max_num_patches", + "norm_mean", + "norm_std", + ): + assert field in str(exc_info.value) + + +def test_build_video_config_returns_none_when_disabled(): + image_config = build_image_preprocessing_config(_image_processor()) + + assert ( + build_video_preprocessing_config( + None, + {"video_temporal_patch_size": 2, "video_num_frames": 8}, + frame_manifest_magic=b"manifest", + ) + is None + ) + assert ( + build_video_preprocessing_config( + image_config, + {}, + frame_manifest_magic=b"manifest", + ) + is None + ) + + +def test_build_video_config_is_not_enabled_by_temporal_patch_size_alone(): + image_config = build_image_preprocessing_config(_image_processor()) + + assert ( + build_video_preprocessing_config( + image_config, + {"video_temporal_patch_size": 2}, + frame_manifest_magic=b"manifest", + ) + is None + ) + + +def test_build_video_config_uses_default_temporal_patch_size(): + image_config = build_image_preprocessing_config(_image_processor()) + + video_config = build_video_preprocessing_config( + image_config, + {"video_num_frames": 8}, + frame_manifest_magic=b"manifest", + ) + + assert video_config is not None + assert video_config.num_frames == 8 + assert video_config.temporal_patch_size == 1 + + +def test_build_video_config_overrides_patch_budget_without_mutating_image_config(): + image_config = build_image_preprocessing_config(_image_processor()) + + video_config = build_video_preprocessing_config( + image_config, + { + "video_num_frames": 8, + "video_temporal_patch_size": 2, + "video_target_num_patches": 64, + "video_maintain_aspect_ratio": False, + }, + frame_manifest_magic=b"manifest", + ) + + assert video_config is not None + assert video_config.image_config is not image_config + assert image_config.dynamic_resolution_max_patches == 32 + assert video_config.image_config.dynamic_resolution_max_patches == 64 + assert video_config.num_frames == 8 + assert video_config.temporal_patch_size == 2 + assert video_config.frame_manifest_magic == b"manifest" + assert video_config.video_maintain_aspect_ratio is False diff --git a/tests/unit/models/generation/test_vllm_video_utils.py b/tests/unit/models/generation/test_vllm_video_utils.py index 468778f9677..dc691b720e1 100644 --- a/tests/unit/models/generation/test_vllm_video_utils.py +++ b/tests/unit/models/generation/test_vllm_video_utils.py @@ -158,7 +158,7 @@ def __init__(self, *_args, **_kwargs): Image.fromarray(frame).save(frame_path) expected_frames.append(frame) frame_paths.append(str(frame_path)) - payload = utils._CACHED_VIDEO_FRAME_MANIFEST_MAGIC + json.dumps( + payload = utils.CACHED_VIDEO_FRAME_MANIFEST_MAGIC + json.dumps( { "frame_paths": frame_paths, "metadata": { @@ -189,7 +189,7 @@ def test_cached_video_data_url_requires_no_driver_decoder(monkeypatch, tmp_path) _, encoded = data_url.split(",", 1) payload = base64.b64decode(encoded) - manifest = json.loads(payload[len(utils._CACHED_VIDEO_FRAME_MANIFEST_MAGIC) :]) + manifest = json.loads(payload[len(utils.CACHED_VIDEO_FRAME_MANIFEST_MAGIC) :]) assert manifest["frame_paths"] == frame_paths assert manifest["metadata"]["frames_indices"] == [0, 1, 2, 3] assert manifest["metadata"]["fps"] == 1.0 @@ -201,7 +201,7 @@ def test_cached_video_manifest_does_not_import_torchcodec(monkeypatch, tmp_path) monkeypatch.setenv("NEMO_RL_VIDEO_MEDIA_ROOT", str(tmp_path)) frame_path = tmp_path / "frame.png" Image.new("RGB", (2, 2)).save(frame_path) - payload = utils._CACHED_VIDEO_FRAME_MANIFEST_MAGIC + json.dumps( + payload = utils.CACHED_VIDEO_FRAME_MANIFEST_MAGIC + json.dumps( { "frame_paths": [str(frame_path)], "metadata": { diff --git a/tests/unit/models/policy/test_megatron_worker.py b/tests/unit/models/policy/test_megatron_worker.py index 1755cbac3f3..9ee11b7912b 100644 --- a/tests/unit/models/policy/test_megatron_worker.py +++ b/tests/unit/models/policy/test_megatron_worker.py @@ -18,7 +18,7 @@ from pathlib import Path from types import SimpleNamespace from typing import Optional -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import numpy as np import pytest @@ -149,18 +149,23 @@ class ModelSlicesContextParallelInputs: assert not _model_slices_context_parallel_inputs(object()) -def test_model_cp_slicing_rejects_transfer_queue_setup(): +def test_setup_data_plane_builds_client(): from nemo_rl.models.policy.workers.megatron_policy_worker import ( MegatronPolicyWorkerImpl, ) worker = object.__new__(MegatronPolicyWorkerImpl) - worker.model_slices_context_parallel_inputs = True + worker._dp_client = None + cfg = MagicMock() + client = MagicMock() - with pytest.raises( - NotImplementedError, match="TransferQueue/SingleController does not yet support" - ): - worker.setup_data_plane(MagicMock()) + with patch( + "nemo_rl.data_plane.build_data_plane_client", return_value=client + ) as build_client: + worker.setup_data_plane(cfg) + + build_client.assert_called_once_with(cfg, bootstrap=False) + assert worker._dp_client is client def test_refit_size_estimate_preserves_integral_buffer_dtype(): @@ -328,6 +333,7 @@ def test_megatron_offload_after_refit_finalizes_before_model_move( move_kwargs = [] worker = object.__new__(MegatronPolicyWorkerImpl) worker.model = _FakeTrainableModel() + worker.model.eval = lambda: events.append("eval") worker.cfg = ( {"generation": {"backend": generation_backend}} if generation_backend else {} ) @@ -360,9 +366,33 @@ def cuda(self): assert events[0] == "finalize_async_save" assert events.index("finalize_async_save") < events.index("move_model") + assert events.index("eval") < events.index("move_model") assert move_kwargs[0]["move_params"] is expect_move_params +def test_megatron_finish_inference_evals_before_model_offload(monkeypatch): + """Mamba decode caches must refresh before CUDA parameter storage is released.""" + from nemo_rl.models.policy.workers.megatron_policy_worker import ( + MegatronPolicyWorkerImpl, + ) + + events = [] + move_kwargs = [] + worker = object.__new__(MegatronPolicyWorkerImpl) + worker.model = _FakeTrainableModel() + worker.model.eval = lambda: events.append("eval") + worker.move_model = lambda model, device, **kwargs: ( + events.append("move_model") or move_kwargs.append(kwargs) or model + ) + + monkeypatch.setattr(torch.cuda, "empty_cache", lambda: None) + + MegatronPolicyWorkerImpl.finish_inference(worker) + + assert events == ["eval", "move_model"] + assert move_kwargs == [{"move_params": True, "move_grads": False}] + + def test_megatron_save_checkpoint_onloads_model_before_save(monkeypatch): """Params offloaded by colocated generation must be onloaded before the save walks them.""" import nemo_rl.models.policy.workers.megatron_policy_worker as worker_module diff --git a/tests/unit/models/policy/test_split_api_wrappers.py b/tests/unit/models/policy/test_split_api_wrappers.py index 34ccd1c1a9c..2e96294ca22 100644 --- a/tests/unit/models/policy/test_split_api_wrappers.py +++ b/tests/unit/models/policy/test_split_api_wrappers.py @@ -32,7 +32,11 @@ from unittest.mock import MagicMock, patch from nemo_rl.data_plane import KVBatchMeta -from nemo_rl.data_plane.schema import DP_TRAIN_FIELDS, ROUTED_EXPERTS_FIELD +from nemo_rl.data_plane.schema import ( + DP_TRAIN_FIELDS, + PACKED_TENSOR_META_PREFIX, + ROUTED_EXPERTS_FIELD, +) from nemo_rl.data_plane.worker_mixin import TQWorkerMixin from nemo_rl.models.policy.tq_policy import TQPolicy @@ -181,6 +185,53 @@ def test_train_microbatches_requests_routed_experts_for_router_replay(self): train_meta = mock_shard.call_args.args[0] assert train_meta.fields == [*DP_TRAIN_FIELDS, ROUTED_EXPERTS_FIELD] + def test_policy_dispatches_request_the_same_multimodal_columns(self): + """Prev/ref logprob and both train APIs must fetch one identical media set.""" + p, wg = _make_tq_policy() + media_fields = [ + "token_type_ids", + "pixel_values", + f"{PACKED_TENSOR_META_PREFIX}pixel_values", + "imgs_sizes", + f"{PACKED_TENSOR_META_PREFIX}imgs_sizes", + ] + meta = KVBatchMeta( + partition_id="train", + task_name="train", + sample_ids=["s0", "s1"], + fields=[*DP_TRAIN_FIELDS, *media_fields], + ) + + dispatched_media: list[list[str]] = [] + calls = [ + lambda: p.get_logprobs_from_meta(meta), + lambda: p.get_reference_policy_logprobs_from_meta(meta), + lambda: p.train_from_meta(meta, loss_fn="LF"), + lambda: p.train_microbatches_from_meta(meta), + ] + for call in calls: + with ( + patch.object(TQPolicy, "_stamp_pad_seqlen"), + patch.object(TQPolicy, "_packing_args", return_value=(None, None)), + patch( + "nemo_rl.models.policy.tq_policy.shard_meta_for_dp", + return_value=([meta, meta], None), + ) as mock_shard, + patch( + "nemo_rl.models.policy.tq_policy._aggregate_train_results", + return_value={}, + ), + ): + wg.get_all_worker_results.reset_mock() + call() + + dispatch_meta = mock_shard.call_args.args[0] + dispatched_media.append( + [field for field in dispatch_meta.fields if field in media_fields] + ) + + assert dispatched_media == [media_fields] * len(calls) + def test_finish_dedupes_replica_twins(self): """TP/CP twins return identical metric copies; aggregating without the is_replica_leader filter inflates every per-token metric.""" diff --git a/tests/unit/single_controller/test_entrypoint.py b/tests/unit/single_controller/test_entrypoint.py index c4436936db8..c3b931e037f 100644 --- a/tests/unit/single_controller/test_entrypoint.py +++ b/tests/unit/single_controller/test_entrypoint.py @@ -89,7 +89,7 @@ def main_context(monkeypatch: pytest.MonkeyPatch) -> SimpleNamespace: monkeypatch.setattr( run_grpo_single_controller, "setup_single_controller", - lambda *_args: (actor_args, SetupTimingMetrics()), + lambda *_args, **_kwargs: (actor_args, SetupTimingMetrics()), ) monkeypatch.setattr( run_grpo_single_controller.SingleControllerActor, @@ -170,3 +170,30 @@ def test_main_configures_generation_for_trained_mtp( assert ( main_context.config.policy["generation"] is main_context.configured_generation ) + + +def test_main_passes_processor_for_vlm( + main_context: SimpleNamespace, + monkeypatch: pytest.MonkeyPatch, +) -> None: + processor = SimpleNamespace(tokenizer="vlm-tokenizer") + get_tokenizer = MagicMock(return_value=processor) + setup_single_controller = MagicMock( + return_value=(main_context.actor_args, SetupTimingMetrics()) + ) + main_context.config.policy["is_vlm"] = True + monkeypatch.setattr(run_grpo_single_controller, "get_tokenizer", get_tokenizer) + monkeypatch.setattr( + run_grpo_single_controller, + "setup_single_controller", + setup_single_controller, + ) + + run_grpo_single_controller.main() + + get_tokenizer.assert_called_once_with( + main_context.config.policy["tokenizer"], get_processor=True + ) + setup_single_controller.assert_called_once_with( + main_context.config, "vlm-tokenizer", processor=processor + ) diff --git a/tests/unit/single_controller/test_setup.py b/tests/unit/single_controller/test_setup.py index d5bc887a3ad..2a79be85358 100644 --- a/tests/unit/single_controller/test_setup.py +++ b/tests/unit/single_controller/test_setup.py @@ -923,6 +923,34 @@ def test_env_handles_sourced_from_setup_response_data(self, patched_factories): assert call_kwargs["env_configs"] == {"math": math_env_cfg} assert actor_args.env_handles is patched_factories["env_handles"] + def test_vlm_processor_used_for_data_and_environment_setup( + self, patched_factories + ): + mc = _make_master_config(env={"clevr-cogent": {"some": "value"}}) + tokenizer = MagicMock(pad_token_id=0) + processor = MagicMock(tokenizer=tokenizer) + processor.model_input_names = ["input_ids", "pixel_values", "image_grid_thw"] + + actor_args, _ = setup_single_controller(mc, tokenizer, processor=processor) + + call_args, call_kwargs = patched_factories["setup_response_data"].call_args + assert call_args[0] is processor + assert call_kwargs["env_configs"] == { + "clevr-cogent": {"some": "value"} + } + assert call_kwargs["is_vlm"] is True + warmup_fields = actor_args.dp_client.register_partition.call_args.kwargs["fields"] + assert "pixel_values" in warmup_fields + assert "__nrl_packed_tensor_meta__pixel_values" in warmup_fields + for field in ( + "imgs_sizes", + "num_frames", + "pixel_values_flat", + "image_num_patches", + ): + assert field in warmup_fields + assert f"__nrl_packed_tensor_meta__{field}" in warmup_fields + def test_weight_sync_factory_args(self, patched_factories): """create_weight_synchronizer receives policy / generation / topology.""" mc = _make_master_config(colocated=False, backend="vllm") @@ -1358,13 +1386,15 @@ def _spinup_gym(**_): mock_megatron.return_value.finish_generation.assert_called_once_with() if gym: # Gym spins up on the reserved URL, before the served-address - # cross-check — so the mismatch leg sees it too. + # cross-check — so the mismatch leg sees it too. The initial refit + # must happen during that wait because it starts Megatron's server. _, spinup_kwargs = mock_spinup.call_args assert spinup_kwargs["base_urls"] == [reserved_url] # The initial refit ran in setup, against the collective brought up # there; the served-address check reads the URLs it populated. weight_sync.init_communicator.assert_called_once_with() weight_sync.sync_weights.assert_called_once_with() + assert mock_megatron.return_value.weight_synchronizer is weight_sync else: mock_spinup.assert_not_called() # Native: the actor's startup sync performs the initial refit. diff --git a/tests/unit/test_check_metrics.py b/tests/unit/test_check_metrics.py index 73cb08f4690..6f417cb15eb 100644 --- a/tests/unit/test_check_metrics.py +++ b/tests/unit/test_check_metrics.py @@ -21,7 +21,15 @@ tests_dir = Path(__file__).parent.parent sys.path.insert(0, str(tests_dir)) -from check_metrics import evaluate_check, max, mean, median, min, ratio_above +from check_metrics import ( + all_finite, + evaluate_check, + max, + mean, + median, + min, + ratio_above, +) class TestMeanFunction: @@ -160,6 +168,18 @@ def test_max_with_string_values(self): assert result == 8.8 +class TestAllFiniteFunction: + def test_requires_at_least_one_value(self): + assert all_finite({}) is False + + @pytest.mark.parametrize("invalid", [float("nan"), float("inf"), float("-inf")]) + def test_rejects_non_finite_values(self, invalid): + assert all_finite({"1": 1.0, "2": invalid}) is False + + def test_accepts_finite_numeric_strings(self): + assert all_finite({"1": "1.0", "2": "-2.5"}) is True + + class TestRatioAboveFunction: """Test the ratio_above function.""" @@ -306,6 +326,14 @@ def test_evaluate_check_with_ratio_above(self): assert passed is False assert value == 0.4 + def test_evaluate_check_with_all_finite(self): + passed, _, value = evaluate_check( + {"loss": {"1": 1.0, "2": 0.5}}, + 'all_finite(data["loss"])', + ) + assert passed is True + assert value is True + class TestRealWorldScenarios: """Test scenarios that match real-world usage patterns.""" diff --git a/tests/unit/test_prepare_video_dataset.py b/tests/unit/test_prepare_video_dataset.py index 9b58c856e5e..cc68186841e 100644 --- a/tests/unit/test_prepare_video_dataset.py +++ b/tests/unit/test_prepare_video_dataset.py @@ -67,7 +67,7 @@ def test_converter_skips_missing_local_videos_when_requested( content = converted_rows[0]["responses_create_params"]["input"][0]["content"] assert content[0]["video_url"] == str(existing_video.resolve()) metadata = converted_rows[0]["responses_create_params"]["metadata"] - assert metadata["chat_template_kwargs"] == {"enable_thinking": True} + assert json.loads(metadata["chat_template_kwargs"]) == {"enable_thinking": True} assert "extra_body" not in metadata assert "extraction_mode" not in converted_rows[0] assert "Skipped 1 non-video or duplicate rows" in capsys.readouterr().out diff --git a/tests/unit/utils/test_venvs.py b/tests/unit/utils/test_venvs.py index 9e54541bb06..2635689385b 100644 --- a/tests/unit/utils/test_venvs.py +++ b/tests/unit/utils/test_venvs.py @@ -16,7 +16,11 @@ from tempfile import TemporaryDirectory from unittest.mock import patch -from nemo_rl.utils.venvs import create_local_venv +from nemo_rl.utils.venvs import ( + add_hf_modules_cache_to_pythonpath, + create_local_venv, + make_actor_runtime_env, +) from tests.unit.conftest import TEST_ASSETS_DIR @@ -48,3 +52,51 @@ def test_create_local_venv(): # Verify the command executed successfully (return code 0) assert result.returncode == 0, f"Failed to import sphinx: {result.stderr}" assert "Sphinx package is installed" in result.stdout + + +def test_add_hf_modules_cache_to_pythonpath(): + result = add_hf_modules_cache_to_pythonpath( + { + "HF_MODULES_CACHE": "/hf/modules", + "PYTHONPATH": f"/project{os.pathsep}/other", + } + ) + + assert result["PYTHONPATH"].split(os.pathsep) == [ + "/hf/modules", + "/project", + "/other", + ] + + +def test_add_hf_modules_cache_does_not_duplicate_pythonpath_entry(): + pythonpath = f"/project{os.pathsep}/hf/modules" + + result = add_hf_modules_cache_to_pythonpath( + {"HF_MODULES_CACHE": "/hf/modules", "PYTHONPATH": pythonpath} + ) + + assert result["PYTHONPATH"] == pythonpath + + +def test_make_actor_runtime_env_builds_local_venv_for_uv_python_executable(): + """Mirrors the inline venv-creation logic that used to live in grpo.py.""" + with ( + patch( + "nemo_rl.distributed.ray_actor_environment_registry.get_actor_python_env", + return_value="uv run --group vllm", + ) as mock_get_env, + patch( + "nemo_rl.utils.venvs.create_local_venv_on_each_node", + return_value="/fake/venv/bin/python", + ) as mock_create_venv, + ): + runtime_env = make_actor_runtime_env("some.module.SomeActor") + + mock_get_env.assert_called_once_with("some.module.SomeActor") + mock_create_venv.assert_called_once_with( + "uv run --group vllm", "some.module.SomeActor" + ) + assert runtime_env["py_executable"] == "/fake/venv/bin/python" + assert runtime_env["env_vars"]["VIRTUAL_ENV"] == "/fake/venv" + assert runtime_env["env_vars"]["UV_PROJECT_ENVIRONMENT"] == "/fake/venv" diff --git a/tools/install_audio_deps.sh b/tools/install_audio_deps.sh index 312cf5f3325..9c89ffd43d2 100755 --- a/tools/install_audio_deps.sh +++ b/tools/install_audio_deps.sh @@ -5,27 +5,46 @@ # # bash tools/install_audio_deps.sh # -# Safe to call multiple times — exits immediately if already installed. +# Safe to call multiple times. set -euo pipefail -# Fast exit: if torchcodec imports cleanly it already has FFmpeg available. -if python -c "import torchcodec" 2>/dev/null; then - echo "[audio-deps] Already installed and functional, skipping." - exit 0 -fi +if ! python -c "import torchcodec" 2>/dev/null; then + # Install system FFmpeg — torchcodec dlopens libavcodec.so.* at runtime. + echo "[audio-deps] Installing system FFmpeg..." + apt-get update && apt-get install -y --no-install-recommends ffmpeg -# Install system FFmpeg — torchcodec dlopens libavcodec.so.* at runtime. -echo "[audio-deps] Installing system FFmpeg..." -apt-get update && apt-get install -y --no-install-recommends ffmpeg + # torchaudio 2.11+ routes torchaudio.load through torchcodec, so both are needed. + # --no-config prevents the project's [tool.uv] overrides from interfering. + echo "[audio-deps] Installing torchaudio==2.11.0 and torchcodec..." + uv pip install --no-config \ + --index-url https://download.pytorch.org/whl/cu130 \ + --extra-index-url https://pypi.org/simple \ + --reinstall-package torchaudio \ + "torchaudio==2.11.0" \ + "torchcodec==0.11.1" +fi -# torchaudio 2.11+ routes torchaudio.load through torchcodec, so both are needed. -# --no-config prevents the project's [tool.uv] overrides from interfering. -echo "[audio-deps] Installing torchaudio==2.11.0 and torchcodec..." -uv pip install --no-config \ - --index-url https://download.pytorch.org/whl/cu130 \ - --extra-index-url https://pypi.org/simple \ - --reinstall-package torchaudio \ - "torchaudio==2.11.0" \ - "torchcodec>=0.3.0" +# PyAV is intentionally absent from the base image (pyproject excludes it via +# `av; sys_platform == 'never'` because it bundles CVE-carrying codec libs), so it +# must be installed after the fact into the isolated Megatron policy worker +# environment that imports it. `--no-config` bypasses that exclusion; the version +# floor is therefore restated here to keep pyproject's CVE-2026-40962 constraint. +# +# The worker venv is created lazily at worker start, so run this AFTER the +# Megatron worker has been created at least once (or point RAY_MEGATRON_PYTHON at +# an existing venv). +RAY_MEGATRON_PYTHON="${RAY_MEGATRON_PYTHON:-/opt/ray_venvs/nemo_rl.models.policy.workers.megatron_policy_worker.MegatronPolicyWorker/bin/python}" +if [[ ! -x "$RAY_MEGATRON_PYTHON" ]]; then + echo "[audio-deps] ERROR: Megatron worker environment not found: $RAY_MEGATRON_PYTHON" >&2 + echo "[audio-deps] It is created on first worker start. Run this script after that," >&2 + echo "[audio-deps] or set RAY_MEGATRON_PYTHON to an existing worker interpreter." >&2 + exit 1 +fi +if ! "$RAY_MEGATRON_PYTHON" -c "import av" 2>/dev/null; then + # `uv pip install --python` targets the venv directly; these venvs are built + # by `uv venv` without `--seed`, so they have no pip to invoke. + echo "[audio-deps] Installing PyAV in the Megatron worker environment..." + uv pip install --no-config --python "$RAY_MEGATRON_PYTHON" "av>=17.1.0" +fi echo "[audio-deps] Done." diff --git a/uv.lock b/uv.lock index f0d07eb780b..e8fbe88ccc4 100644 --- a/uv.lock +++ b/uv.lock @@ -88,7 +88,7 @@ constraints = [ { name = "cryptography", specifier = ">=48.0.1" }, { name = "diffusers", specifier = ">=0.38.0" }, { name = "dulwich", specifier = ">=1.2.5" }, - { name = "emerging-optimizers", git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git?rev=v0.2.0" }, + { name = "emerging-optimizers", git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git?rev=v0.3.0" }, { name = "fast-hadamard-transform", git = "https://github.com/Dao-AILab/fast-hadamard-transform.git?rev=f134af63deb2df17e1171a9ec1ea4a7d8604d5ca" }, { name = "megatron-fsdp", git = "https://github.com/yuhezhang-ai/Megatron-LM.git?subdirectory=megatron%2Fcore%2Fdistributed%2Ffsdp%2Fsrc&rev=455389c480af6b3acdca74c7830c68b3274eb083" }, { name = "onnx", specifier = ">=1.21.0rc4" }, @@ -549,6 +549,37 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/7e/16/fbe8e1e185a45042f7cd3a282def5bb8d95bb69ab9e9ef6a5368aa17e426/audioread-3.1.0-py3-none-any.whl", hash = "sha256:b30d1df6c5d3de5dcef0fb0e256f6ea17bdcf5f979408df0297d8a408e2971b4", size = 23143, upload-time = "2025-10-26T19:44:12.016Z" }, ] +[[package]] +name = "awkward" +version = "2.13.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "awkward-cpp" }, + { name = "fsspec" }, + { name = "numpy" }, + { name = "packaging" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/43/97/2728ab879ed3edaecfbd8e9daa7d88a2f89a7ea584e413c4d8971fc0414d/awkward-2.13.0.tar.gz", hash = "sha256:36f127573295e1ddf65b551bf071b803ad4af21bd14519f4c4dd34953cb3efc2", size = 6479149, upload-time = "2026-08-14T16:52:43.873Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/25/34/b8eece9a8d024defc08418ea883361b3d7e3300392719ee27535e93885a5/awkward-2.13.0-py3-none-any.whl", hash = "sha256:ff40879e7179a6f14f4c4ee5f5297e3dcd027b99b4c64188bdb7dccdf625e982", size = 990112, upload-time = "2026-08-14T16:52:42.018Z" }, +] + +[[package]] +name = "awkward-cpp" +version = "56" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/43/92/6325292ccc2be3ede227fd0b4bd7b23bc7597790f31227171bec4009d2df/awkward_cpp-56.tar.gz", hash = "sha256:cd3635fab926c6630c0a85d92b59a15851c4f5a881ea749984069027634f8f52", size = 1501723, upload-time = "2026-08-14T15:30:10.748Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e5/8f/d824319c6576839b00c14dae4857faaa8c2ac9f1c1e4b76dc003b18b8232/awkward_cpp-56-cp313-cp313-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5991ab389a365defe7f60a6314cf8658733ad4737f1ce528d39508c20d776ea9", size = 626456, upload-time = "2026-08-14T15:29:26.14Z" }, + { url = "https://files.pythonhosted.org/packages/65/cc/dc7f4e98b538da691a28f6c352a6a8ea6c842f04b64db5257086996bd1dd/awkward_cpp-56-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f68511c490e47d009e9ba1076f91870a9abee3bc82ac285a1965fa57aea6f104", size = 698747, upload-time = "2026-08-14T15:29:28.101Z" }, + { url = "https://files.pythonhosted.org/packages/2b/04/65153eade3f380535ddc0bd70a03c36c4e9eddd2eed1adc76999359a37d7/awkward_cpp-56-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:26b0c3b6452352a9d1788a106d4153305330cc13d5e10756d5e50712dd526d53", size = 1626987, upload-time = "2026-08-14T15:29:30.258Z" }, + { url = "https://files.pythonhosted.org/packages/3e/57/282e4981754bf1a4996b13fc11174eb6bcacd5653d0f338d410ec59a9c4b/awkward_cpp-56-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:00023f08d2ed82e2b8e0ba2e4637fa62b8164163f692db44e7ebd27d207f1955", size = 1745804, upload-time = "2026-08-14T15:29:31.817Z" }, + { url = "https://files.pythonhosted.org/packages/ce/42/6b5243f804d4214a3885998e82a34d421e6e69a830651d195cf49ccca759/awkward_cpp-56-cp313-cp313-pyemscripten_2025_0_wasm32.whl", hash = "sha256:1417c538d468374e86ad7f5e0b49f458e001c9e843afac3ae866fa7832433b75", size = 234427, upload-time = "2026-08-14T15:29:33.315Z" }, +] + [[package]] name = "awscrt" version = "0.36.0" @@ -1580,8 +1611,8 @@ wheels = [ [[package]] name = "emerging-optimizers" -version = "0.2.0" -source = { git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git?rev=v0.2.0#1effa026ff096b7fa1063ca2fba19d98be6e6cdf" } +version = "0.3.0" +source = { git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git?rev=v0.3.0#b309e2f01cda75dc96a6dc1a2355a7b3b64b5e16" } dependencies = [ { name = "absl-py" }, { name = "torch" }, @@ -2252,6 +2283,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b3/bb/d71d6da82763528c2c2ed6b59a9d6142c6595545a4c448e2085d155e88c2/gguf-0.19.0-py3-none-any.whl", hash = "sha256:70bcd10edfe697fb2dad6e40af2234b9d8ece9a41a99761405121ebda1c3c1cd", size = 118475, upload-time = "2026-05-06T13:04:02.588Z" }, ] +[[package]] +name = "gigatoken" +version = "0.10.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "awkward" }, + { name = "numpy" }, + { name = "typer" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/33/8a/fa097b404650a9eaea59de80bc8c33ad8ccb6b4a17ff6f33f213ec091057/gigatoken-0.10.0.tar.gz", hash = "sha256:562fd4284eacdebd8a8043ce21ddcdacb61210a036e6e700cd13bdca2b2af7ac", size = 1330282, upload-time = "2026-07-25T19:15:42.934Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/35/6f/d393a7735cbf775f612da4d0674f7ecc8900e0be4989fd46cbe23f650965/gigatoken-0.10.0-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:46f726b77e4132914f64ed42c7b81d82bdc40e9302138fcaca0b394eed62fc47", size = 4486995, upload-time = "2026-07-25T19:15:35.573Z" }, + { url = "https://files.pythonhosted.org/packages/e7/32/a0b3af6db26b3704224d3ad856ebadfdb6ed5d50210ab16ab9ad27b30673/gigatoken-0.10.0-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:2f3481d0c067beacf1c0a7640d956edbc8973014ba8cfa3e18bfafb02f78e53b", size = 5336884, upload-time = "2026-07-25T19:15:37.245Z" }, +] + [[package]] name = "gitdb" version = "4.0.12" @@ -3482,7 +3528,6 @@ dev = [ { name = "emerging-optimizers" }, { name = "fast-hadamard-transform" }, { name = "fastapi", extra = ["standard"], marker = "extra == 'extra-7-nemo-rl-mcore' or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-modelopt' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nemo-gym' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nvrx' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-8-nemo-gym-vllm') or (extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-8-nemo-gym-vllm') or (extra == 'extra-7-nemo-rl-vllm' and extra == 'extra-8-nemo-gym-vllm') or (extra == 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra == 'extra-8-nemo-gym-vllm')" }, - { name = "flash-linear-attention" }, { name = "flashinfer-python", version = "0.6.8.post1", source = { registry = "https://pypi.org/simple" } }, { name = "hypercorn" }, { name = "megatron-energon", marker = "extra == 'extra-7-nemo-rl-mcore' or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-modelopt' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nemo-gym' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nvrx' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-8-nemo-gym-vllm') or (extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-8-nemo-gym-vllm') or (extra == 'extra-7-nemo-rl-vllm' and extra == 'extra-8-nemo-gym-vllm') or (extra == 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra == 'extra-8-nemo-gym-vllm')" }, @@ -3495,8 +3540,10 @@ dev = [ { name = "openai", extra = ["aiohttp"], marker = "extra == 'extra-7-nemo-rl-mcore' or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-modelopt' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nemo-gym' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nvrx' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-8-nemo-gym-vllm') or (extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-8-nemo-gym-vllm') or (extra == 'extra-7-nemo-rl-vllm' and extra == 'extra-8-nemo-gym-vllm') or (extra == 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra == 'extra-8-nemo-gym-vllm')" }, { name = "opentelemetry-api" }, { name = "orjson" }, + { name = "pillow" }, { name = "quart" }, { name = "tensorstore" }, + { name = "torch-memory-saver", version = "0.0.10b1", source = { git = "https://github.com/fzyzcjy/torch_memory_saver.git?rev=9bc9a442e6d108c7b7903def199896a005143aaf#9bc9a442e6d108c7b7903def199896a005143aaf" } }, { name = "tqdm" }, { name = "wget" }, { name = "zstandard" }, @@ -3504,6 +3551,7 @@ dev = [ mlm = [ { name = "accelerate" }, { name = "flask-restful" }, + { name = "gigatoken" }, { name = "omegaconf" }, { name = "sentencepiece" }, { name = "tiktoken" }, @@ -3522,13 +3570,15 @@ requires-dist = [ { name = "causal-conv1d", marker = "extra == 'ssm'", specifier = "~=1.5" }, { name = "datasets", marker = "extra == 'dev'" }, { name = "einops", marker = "extra == 'dev'", specifier = "~=0.8" }, - { name = "emerging-optimizers", marker = "extra == 'dev'", git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git?rev=v0.2.0" }, + { name = "emerging-optimizers", marker = "extra == 'dev'", git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git?rev=v0.3.0" }, { name = "fast-hadamard-transform", marker = "extra == 'dev'", git = "https://github.com/Dao-AILab/fast-hadamard-transform.git?rev=f134af63deb2df17e1171a9ec1ea4a7d8604d5ca" }, { name = "fastapi", marker = "extra == 'dev'", specifier = "~=0.50" }, - { name = "flash-linear-attention", marker = "extra == 'dev'", specifier = "==0.5.1" }, + { name = "flash-linear-attention", marker = "extra == 'ssm'", specifier = "==0.5.1" }, { name = "flashinfer-python", marker = "extra == 'dev'", specifier = ">=0.5.0,<0.7.0" }, { name = "flask-restful", marker = "extra == 'mlm'" }, { name = "flask-restful", marker = "extra == 'training'" }, + { name = "gigatoken", marker = "extra == 'mlm'" }, + { name = "gigatoken", marker = "extra == 'training'" }, { name = "hypercorn", marker = "extra == 'dev'" }, { name = "mamba-ssm", marker = "extra == 'ssm'", git = "https://github.com/state-spaces/mamba.git?rev=0048fbf2e7b2f214dcbe703ea3dec2b9647595e1" }, { name = "megatron-energon", extras = ["av-decode"], marker = "extra == 'dev'", specifier = "~=7.0" }, @@ -3546,6 +3596,7 @@ requires-dist = [ { name = "opentelemetry-api", marker = "extra == 'dev'", specifier = ">=1.33.1,<2" }, { name = "orjson", marker = "extra == 'dev'" }, { name = "packaging", specifier = ">=24.2" }, + { name = "pillow", marker = "extra == 'dev'" }, { name = "quart", marker = "extra == 'dev'" }, { name = "sentencepiece", marker = "extra == 'mlm'" }, { name = "sentencepiece", marker = "extra == 'training'" }, @@ -3553,6 +3604,7 @@ requires-dist = [ { name = "tiktoken", marker = "extra == 'mlm'" }, { name = "tiktoken", marker = "extra == 'training'" }, { name = "torch", specifier = ">=2.6.0" }, + { name = "torch-memory-saver", marker = "extra == 'dev'", git = "https://github.com/fzyzcjy/torch_memory_saver.git?rev=9bc9a442e6d108c7b7903def199896a005143aaf" }, { name = "tqdm", marker = "extra == 'dev'" }, { name = "transformer-engine", extras = ["core-cu13", "pytorch"], marker = "extra == 'te'", git = "https://github.com/NVIDIA/TransformerEngine.git?rev=4329ff84bfbdaa778a33cba02a15fb0807c64689" }, { name = "transformers", marker = "extra == 'mlm'" }, @@ -3598,7 +3650,7 @@ linting = [ ] no-pypi-wheels = [ { name = "deep-gemm", git = "https://github.com/deepseek-ai/DeepGEMM.git?rev=714dd1a4a980f7937a74343d19a8eba4fe321480" }, - { name = "emerging-optimizers", git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git?rev=v0.2.0" }, + { name = "emerging-optimizers", git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git?rev=v0.3.0" }, { name = "flash-mla", git = "https://github.com/deepseek-ai/FlashMLA?rev=nv_dev" }, ] test = [ @@ -4449,7 +4501,7 @@ mcore = [ { name = "nvidia-cutlass-dsl", version = "4.5.2", source = { registry = "https://pypi.org/simple" }, extra = ["cu13"], marker = "extra == 'extra-7-nemo-rl-mcore' or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-modelopt' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nemo-gym' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-nvrx' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-8-nemo-gym-vllm') or (extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-8-nemo-gym-vllm') or (extra == 'extra-7-nemo-rl-vllm' and extra == 'extra-8-nemo-gym-vllm') or (extra == 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra == 'extra-8-nemo-gym-vllm')" }, { name = "nvshmem4py-cu13" }, { name = "tilelang", version = "0.1.9", source = { registry = "https://pypi.org/simple" } }, - { name = "torch-memory-saver" }, + { name = "torch-memory-saver", version = "0.0.10b1", source = { git = "https://github.com/fzyzcjy/torch_memory_saver.git?rev=9bc9a442e6d108c7b7903def199896a005143aaf#9bc9a442e6d108c7b7903def199896a005143aaf" } }, { name = "transformers", version = "5.12.1", source = { registry = "https://pypi.org/simple" } }, ] modelopt = [ @@ -7125,7 +7177,7 @@ dependencies = [ { name = "timm" }, { name = "tokenspeed-mla", version = "0.1.8", source = { registry = "https://pypi.org/simple" } }, { name = "torch" }, - { name = "torch-memory-saver" }, + { name = "torch-memory-saver", version = "0.0.9.post1", source = { registry = "https://pypi.org/simple" } }, { name = "torchao", version = "0.17.0", source = { registry = "https://pypi.org/simple" } }, { name = "torchaudio" }, { name = "torchcodec", version = "0.11.1", source = { registry = "https://pypi.org/simple" }, marker = "platform_machine == 'x86_64'" }, @@ -8221,12 +8273,25 @@ wheels = [ name = "torch-memory-saver" version = "0.0.9.post1" source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "platform_machine == 'x86_64' and sys_platform == 'linux'", + "platform_machine == 'aarch64' and sys_platform == 'linux'", +] sdist = { url = "https://files.pythonhosted.org/packages/81/fd/42aad783d433fd69dc108b1b2ee5860fcf33e20e5440b899bc004ff97d70/torch_memory_saver-0.0.9.post1.tar.gz", hash = "sha256:25fd4b691ed3242c3a18b2bef0dbe9de84d2e7068b96a37686a923d55c274f43", size = 15209, upload-time = "2026-05-02T05:30:54.649Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/69/76/9a88e20f0be461af58165cc87663076a1fdfa53893f8371442d2e4ccabd1/torch_memory_saver-0.0.9.post1-cp39-abi3-manylinux2014_aarch64.whl", hash = "sha256:2c3505ba2bd7aa9ac75392417c675e8be686029da7c5f35037c59e2e8a3dfece", size = 1008682, upload-time = "2026-05-02T05:30:53.267Z" }, { url = "https://files.pythonhosted.org/packages/49/67/6789f9048b836615d07e419fd2f253b83fb9b0f86f34d03ab75ac1f07049/torch_memory_saver-0.0.9.post1-cp39-abi3-manylinux2014_x86_64.whl", hash = "sha256:b4e560fcd88a0641efb21ac530bb79a6e78fc0b31c9a654fa3b5eb7c3629e555", size = 1015832, upload-time = "2026-05-02T05:30:51.317Z" }, ] +[[package]] +name = "torch-memory-saver" +version = "0.0.10b1" +source = { git = "https://github.com/fzyzcjy/torch_memory_saver.git?rev=9bc9a442e6d108c7b7903def199896a005143aaf#9bc9a442e6d108c7b7903def199896a005143aaf" } +resolution-markers = [ + "platform_machine == 'x86_64' and sys_platform == 'linux'", + "platform_machine == 'aarch64' and sys_platform == 'linux'", +] + [[package]] name = "torchao" version = "0.15.0"