diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-click-2n8g-megatron-tp2ep8.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-click-2n8g-megatron-tp2ep8.v1.yaml index 3695064845c..116fa267a7e 100644 --- a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-click-2n8g-megatron-tp2ep8.v1.yaml +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-click-2n8g-megatron-tp2ep8.v1.yaml @@ -1,5 +1,6 @@ defaults: ../../vlm_grpo_3B_megatron.yaml grpo: + deduplicate_multimodal_data: true num_prompts_per_step: 1 num_val_generations_per_prompt: 1 max_num_steps: 100 diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-automodel-ep8.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-automodel-ep8.v1.yaml index c2ed8cbdb39..c94d84349f6 100644 --- a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-automodel-ep8.v1.yaml +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-automodel-ep8.v1.yaml @@ -1,5 +1,6 @@ defaults: ../../vlm_grpo_3B.yaml grpo: + deduplicate_multimodal_data: true num_prompts_per_step: 32 val_at_start: true checkpointing: diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8.v1.yaml index b738cf17a13..6771b5bbc80 100644 --- a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8.v1.yaml +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8.v1.yaml @@ -1,4 +1,6 @@ defaults: ../../vlm_grpo_3B_megatron.yaml +grpo: + deduplicate_multimodal_data: true loss_fn: reference_policy_kl_penalty: 0.0 checkpointing: diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-automodel-ep8.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-automodel-ep8.v1.yaml index 3a0f328a36a..54cf70c9981 100644 --- a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-automodel-ep8.v1.yaml +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-automodel-ep8.v1.yaml @@ -1,5 +1,6 @@ defaults: ../../vlm_grpo_3B.yaml grpo: + deduplicate_multimodal_data: true num_prompts_per_step: 32 overlong_filtering: true seq_logprob_error_threshold: 2 diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-megatron-tp8ep16.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-megatron-tp8ep16.v1.yaml index 844917c1fc2..d84e512c877 100644 --- a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-megatron-tp8ep16.v1.yaml +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-megatron-tp8ep16.v1.yaml @@ -3,7 +3,7 @@ grpo: num_prompts_per_step: 512 overlong_filtering: true zero_variance_prompt_filtering: false - deduplicate_multimodal_data: false + deduplicate_multimodal_data: true loss_fn: ratio_clip_max: 0.28 use_on_policy_kl_approximation: true diff --git a/examples/configs/vlm_grpo_3B.yaml b/examples/configs/vlm_grpo_3B.yaml index 916a4a9c4f0..19d82e5c845 100644 --- a/examples/configs/vlm_grpo_3B.yaml +++ b/examples/configs/vlm_grpo_3B.yaml @@ -3,6 +3,8 @@ defaults: "grpo_math_1B.yaml" grpo: + deduplicate_multimodal_data: false + debug_payload_metrics: false num_prompts_per_step: 8 reward_shaping: overlong_buffer_length: 512 diff --git a/examples/nemo_gym/run_grpo_nemo_gym.py b/examples/nemo_gym/run_grpo_nemo_gym.py index 4d83b9c129f..cfa6fa464ba 100644 --- a/examples/nemo_gym/run_grpo_nemo_gym.py +++ b/examples/nemo_gym/run_grpo_nemo_gym.py @@ -315,6 +315,7 @@ def main() -> None: max_trajectory_age_steps=config.grpo.async_grpo.max_trajectory_age_steps, teacher_worker_groups=teacher_worker_groups, alias_to_group_alias=alias_to_group_alias, + processor=processor, ) else: print("šŸš€ Running synchronous GRPO training") @@ -333,6 +334,7 @@ def main() -> None: checkpointer, grpo_state, master_config, + processor=processor, ) diff --git a/examples/run_vlm_grpo.py b/examples/run_vlm_grpo.py index e2b17e43baf..a0030ca20ad 100644 --- a/examples/run_vlm_grpo.py +++ b/examples/run_vlm_grpo.py @@ -146,6 +146,7 @@ def main() -> None: checkpointer, grpo_state, master_config, + processor=processor, ) diff --git a/nemo_rl/algorithms/async_utils/interfaces.py b/nemo_rl/algorithms/async_utils/interfaces.py index 892f3293a0f..824f718b6e7 100644 --- a/nemo_rl/algorithms/async_utils/interfaces.py +++ b/nemo_rl/algorithms/async_utils/interfaces.py @@ -73,6 +73,20 @@ def load_state_dict( """Restore state produced by ``state_dict``.""" ... + def save_to_path(self, path: str) -> int: + """Serialize state directly from the replay actor.""" + ... + + def load_from_path( + self, + path: str, + num_prompts_per_step: int | None = None, + current_training_step: int | None = None, + max_age_steps: int | None = None, + ) -> dict[str, int]: + """Restore state directly in the replay actor.""" + ... + def get_trajectories_needed( self, target_step: int, diff --git a/nemo_rl/algorithms/async_utils/replay_buffer.py b/nemo_rl/algorithms/async_utils/replay_buffer.py index 8b182b349ae..7d408d7b4fe 100644 --- a/nemo_rl/algorithms/async_utils/replay_buffer.py +++ b/nemo_rl/algorithms/async_utils/replay_buffer.py @@ -13,6 +13,7 @@ # limitations under the License. import asyncio +import gc import statistics import threading as _threading import uuid @@ -21,11 +22,16 @@ from typing import Any, Iterable, Optional import ray +import torch from nemo_rl.algorithms.async_utils.interfaces import ReplayBufferProtocol from nemo_rl.data_plane import KVBatchMeta from nemo_rl.data_plane.schema import ROUTED_EXPERTS_FIELD -from nemo_rl.experience.interfaces import PromptGroupRecord +from nemo_rl.experience.interfaces import ( + NEMO_GYM_TASK_INDEX_KEY, + NEXT_NEMO_GYM_TASK_INDEX_KEY, + PromptGroupRecord, +) from nemo_rl.experience.payload import pack_payload, record_to_train_batch from nemo_rl.utils.r3_trace import trace_rollout_payload @@ -340,6 +346,44 @@ def state_dict(self) -> dict[str, Any]: "max_size": self.max_size, } + def save_to_path(self, path: str) -> int: + """Serialize inside the actor without materializing the buffer on the driver.""" + state = self.state_dict() + torch.save(state, path) + num_trajectories = len(state["trajectories"]) + del state + gc.collect() + return num_trajectories + + def load_from_path( + self, + path: str, + num_prompts_per_step: int | None = None, + current_training_step: int | None = None, + max_age_steps: int | None = None, + ) -> dict[str, int]: + """Restore inside the actor and return only compact coordination metadata.""" + state = torch.load(path, weights_only=False) + saved_task_indices = [ + int(trajectory[NEMO_GYM_TASK_INDEX_KEY]) + for trajectory in state.get("trajectories", []) + if trajectory.get(NEMO_GYM_TASK_INDEX_KEY) is not None + ] + next_task_index = max(saved_task_indices, default=-1) + 1 + num_trajectories = len(state["trajectories"]) + self.load_state_dict( + state, + num_prompts_per_step=num_prompts_per_step, + current_training_step=current_training_step, + max_age_steps=max_age_steps, + ) + del state + gc.collect() + return { + "num_trajectories": num_trajectories, + NEXT_NEMO_GYM_TASK_INDEX_KEY: next_task_index, + } + def load_state_dict( self, state: dict[str, Any], diff --git a/nemo_rl/algorithms/async_utils/trajectory_collector.py b/nemo_rl/algorithms/async_utils/trajectory_collector.py index ca57644afa8..f87673b4c57 100644 --- a/nemo_rl/algorithms/async_utils/trajectory_collector.py +++ b/nemo_rl/algorithms/async_utils/trajectory_collector.py @@ -38,10 +38,16 @@ ) from nemo_rl.experience.rollouts import ( RolloutGroupResult, + attach_initial_nemo_gym_image_payloads, run_async_multi_turn_rollout_groups, ) from nemo_rl.models.generation.interfaces import GenerationConfig, GenerationInterface from nemo_rl.utils.logger import should_log_nemo_gym_full_result_tables +from nemo_rl.utils.multimodal_payload_metrics import ( + collect_multimodal_payload_metrics, + drain_multimodal_payload_metrics, + print_multimodal_payload_metrics, +) from nemo_rl.utils.timer import ThreadSafeTimer TokenizerType = PreTrainedTokenizerBase @@ -66,6 +72,7 @@ def __init__( alias_to_group_alias: Optional[dict[str, str]] = None, on_policy_distillation_cfg: Optional[dict[str, Any]] = None, next_nemo_gym_task_index: int = 0, + processor: Any = None, ): self.policy_generation = policy_generation self.tokenizer = tokenizer @@ -75,6 +82,7 @@ def __init__( self.teacher_worker_groups = teacher_worker_groups or {} self.alias_to_group_alias = alias_to_group_alias or {} self.on_policy_distillation_cfg = on_policy_distillation_cfg or {} + self.processor = processor self._has_distillation_teachers = bool(self.teacher_worker_groups) self._teacher_seq_pad_multiple = teacher_seq_pad_multiple( self.teacher_worker_groups, @@ -428,7 +436,23 @@ def _process_batch(self, batch: BatchedDataDict[DatumSpec]) -> None: rollout_batch = batch.slice(0, num_prompts_to_generate) if use_nemo_gym: self._stamp_nemo_gym_task_indices(rollout_batch) - repeated_batch = rollout_batch.repeat_interleave(num_generations) + if self.master_config.grpo.deduplicate_multimodal_data: + attach_initial_nemo_gym_image_payloads( + rollout_batch, self.processor + ) + repeated_batch = rollout_batch.repeat_interleave( + num_generations, + share_immutable_media=( + self.master_config.grpo.deduplicate_multimodal_data + ), + ) + print_multimodal_payload_metrics( + collect_multimodal_payload_metrics( + repeated_batch, + "prompt_repeat_async", + enabled=self.master_config.grpo.debug_payload_metrics, + ) + ) def _run_rollout_batch() -> None: asyncio.run( @@ -605,6 +629,15 @@ def get_efficiency_metrics(self) -> dict[str, float]: self._efficiency_timer.get_timing_metrics(reduction_op="sum"), ) + async def drain_payload_metrics(self) -> dict[str, int | float]: + """Close one drain-to-drain collector/Gym telemetry interval. + + Rollout collection is concurrent with training, so the interval is not + claimed to own the sampled training batch. Call-normalized metrics make + intervals comparable even when their background transfer counts differ. + """ + return drain_multimodal_payload_metrics() + def get_rollouts_state(self) -> dict[str, int]: """Get collector-side rollout state for checkpointing.""" return {NEXT_NEMO_GYM_TASK_INDEX_KEY: self._next_nemo_gym_task_index} @@ -777,6 +810,10 @@ async def _iter_rollout_groups( mask_env_flagged_samples=should_mask_flagged_samples( self.master_config.env ), + deduplicate_multimodal_data=( + self.master_config.grpo.deduplicate_multimodal_data + ), + debug_payload_metrics=self.master_config.grpo.debug_payload_metrics, ): task_index = rollout_result.task_index if task_index is None: @@ -801,6 +838,9 @@ async def _iter_rollout_groups( num_generations=num_generations, max_rollout_turns=self.master_config.grpo.max_rollout_turns, greedy=False, + deduplicate_multimodal_data=( + self.master_config.grpo.deduplicate_multimodal_data + ), ): yield rollout_result @@ -922,11 +962,22 @@ async def _enqueue_rollout_group( } if rollout_result.task_index is not None: trajectory_group[NEMO_GYM_TASK_INDEX_KEY] = rollout_result.task_index - backoff_delay = 0.01 backoff_started_at: float | None = None try: while self.running: + # Every retry is a distinct Ray submission of the full payload. + print_multimodal_payload_metrics( + collect_multimodal_payload_metrics( + ( + trajectory_group, + generation_weight_version, + target_weight_version, + ), + "replay_push", + enabled=self.master_config.grpo.debug_payload_metrics, + ) + ) status = await self.replay_buffer.add.remote( trajectory_group, generation_weight_version, diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index 6d612e50b0d..e1d04189d1c 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -69,6 +69,7 @@ batched_message_log_to_flat_message, get_keys_from_message_log, ) +from nemo_rl.data.multimodal_utils import PackedTensor from nemo_rl.data.utils import extract_necessary_env_names, load_dataloader_state from nemo_rl.data_plane.interfaces import DataPlaneConfig from nemo_rl.distributed.batched_data_dict import BatchedDataDict @@ -88,6 +89,7 @@ ) from nemo_rl.experience.rollouts import ( EffortLevelsConfig, + attach_initial_nemo_gym_image_payloads, backfill_missing_routed_experts, get_nemo_gym_thinking_tags, run_async_multi_turn_rollout, @@ -125,6 +127,12 @@ should_log_nemo_gym_full_result_tables, ) from nemo_rl.utils.memory_tracker import MemoryTracker +from nemo_rl.utils.multimodal_payload_metrics import ( + collect_multimodal_payload_metrics, + drain_multimodal_payload_metrics, + merge_multimodal_payload_metrics, + print_multimodal_payload_metrics, +) from nemo_rl.utils.nsys import maybe_gpu_profile_step from nemo_rl.utils.timer import TimeoutChecker, Timer from nemo_rl.utils.venvs import create_local_venv_on_each_node @@ -158,6 +166,29 @@ def _get_next_nemo_gym_task_index( return next_task_index +def _save_async_replay_buffer_checkpoint( + replay_buffer: Any, + checkpoint_path: str, + checkpointing_config: CheckpointingConfig, +) -> int | None: + """Checkpoint replay state inside its actor, or skip it when configured.""" + if not checkpointing_config.get("save_replay_buffer", True): + print( + "ā­ļø Skipping replay buffer checkpoint " + "(checkpointing.save_replay_buffer=false)" + ) + return None + + print("šŸ“¦ Saving replay buffer state...") + num_buffered_trajectories = ray.get( + replay_buffer.save_to_path.remote( + os.path.join(checkpoint_path, "replay_buffer.pt") + ) + ) + print(f"āœ… Saved replay buffer with {num_buffered_trajectories} trajectories") + return num_buffered_trajectories + + class RewardScalingConfig(BaseModel, extra="allow"): """Configure linear reward scaling with clamping. @@ -289,6 +320,11 @@ class GRPOConfig(BaseModel, extra="allow"): malformed_thinking_advantage: float | None = None # Advantage estimator configuration (grpo or reinforce_plus_plus) adv_estimator: AdvEstimatorConfig = Field(default_factory=AdvEstimatorConfig) + # Share and compact immutable image/video/audio payload segments across + # logical GRPO rows. Prompt identity is never used as proof of equality. + deduplicate_multimodal_data: bool = False + # Emit exact-boundary and logical-vs-physical payload metrics. + debug_payload_metrics: bool = False @dataclass @@ -350,6 +386,24 @@ class MasterConfig(BaseModel, extra="allow"): # =============================================================================== +def _validate_multimodal_dedup_capability(master_config: MasterConfig) -> None: + """Reject configurations whose media transfer path is not qualified.""" + if not master_config.grpo.deduplicate_multimodal_data: + return + + generation_config = master_config.policy["generation"] + if generation_config.get("backend") != "vllm": + raise NotImplementedError( + "grpo.deduplicate_multimodal_data=true is currently qualified " + "only with policy.generation.backend=vllm." + ) + if (master_config.data_plane or {}).get("enabled", False): + raise NotImplementedError( + "grpo.deduplicate_multimodal_data=true is currently supported " + "only when data_plane.enabled=false." + ) + + def setup( master_config: MasterConfig, tokenizer: TokenizerType, @@ -404,6 +458,7 @@ def setup( ) if generation_config["backend"] == "vllm": normalize_vllm_refit_config(cast(VllmConfig, generation_config)) + _validate_multimodal_dedup_capability(master_config) # Validation-only sampling is honored only on the NeMo-Gym vLLM rollout # path; everywhere else validation must sample exactly like training. @@ -997,6 +1052,7 @@ def _spinup_nemo_gym(base_urls, model_name): # vllm model loading prefers clean environment, initialize policy_generation before policy in colocated mode backend = generation_config["backend"] generation_config["model_name"] = policy_config["model_name"] # Needed for vLLM + generation_config["debug_payload_metrics"] = grpo_config.debug_payload_metrics remote_transport = None remote_synchronizer_cls = None remote_baseline_init_refs: list[Any] = [] @@ -1065,6 +1121,8 @@ def init_policy(): init_optimizer=True, init_reference_model=init_reference_model, ) + # Keep custom policy_factory call signatures backward compatible. + p.debug_payload_metrics = grpo_config.debug_payload_metrics if remote_transport is not None: assert remote_synchronizer_cls is not None remote_baseline_init_refs.extend( @@ -1670,7 +1728,10 @@ def dynamic_sampling( filtered_repeated_batch if batch_cache is None else BatchedDataDict.from_batches( - [batch_cache, filtered_repeated_batch] + [batch_cache, filtered_repeated_batch], + allow_missing_packed_tensors=( + master_config.grpo.deduplicate_multimodal_data + ), ) ) filtered_repeated_batch = batch_cache @@ -2043,6 +2104,19 @@ def _preserve_router_replay_routed_experts( target["routed_experts"] = flat_messages["routed_experts"] +def _should_normalize_sparse_replay_media( + batches: list[BatchedDataDict], + *, + deduplicate_multimodal_data: bool, +) -> bool: + """Keep sparse compact checkpoints readable across a flag transition.""" + return deduplicate_multimodal_data or any( + isinstance(value, PackedTensor) and value.deduplication_enabled + for batch in batches + for value in batch.values() + ) + + def _build_async_grpo_train_data( flat_messages: BatchedDataDict, input_lengths: torch.Tensor, @@ -2061,7 +2135,9 @@ def _build_async_grpo_train_data( ) _preserve_router_replay_routed_experts(train_data, flat_messages, policy_config) # update multimodal data unconditionally - extra_multimodal_data = flat_messages.get_multimodal_dict(as_tensors=False) + extra_multimodal_data = flat_messages.get_multimodal_dict( + as_tensors=False, pixel_dtype=torch.bfloat16 + ) train_data.update(extra_multimodal_data) return train_data @@ -2669,6 +2745,7 @@ def grpo_train( checkpointer: CheckpointManager, grpo_save_state: GRPOSaveState, master_config: MasterConfig, + processor: Optional[AutoProcessor] = None, ) -> None: """Run GRPO training algorithm.""" timer = Timer(context={"worker": "driver"}) @@ -2744,10 +2821,19 @@ def grpo_train( step=0, master_config=master_config, logger=logger, + processor=processor, ) policy_generation.finish_generation() logger.log_metrics(val_metrics, current_step, prefix="validation") logger.log_metrics(validation_timings, current_step, prefix="timing/validation") + if master_config.grpo.debug_payload_metrics: + validation_payload_metrics = drain_multimodal_payload_metrics() + if validation_payload_metrics: + logger.log_metrics( + validation_payload_metrics, + current_step, + prefix="validation", + ) stop_message = _validation_early_stop_message( val_metrics, stop_at_validation_threshold, @@ -2804,10 +2890,25 @@ def grpo_train( # Prepare batch print("ā–¶ Preparing batch...", flush=True) with timer.time("data_processing"): + if ( + master_config.grpo.deduplicate_multimodal_data + and _should_use_nemo_gym(master_config) + ): + attach_initial_nemo_gym_image_payloads(batch, processor) # Repeat batch items repeated_batch: BatchedDataDict[DatumSpec] = ( batch.repeat_interleave( - master_config.grpo.num_generations_per_prompt + master_config.grpo.num_generations_per_prompt, + share_immutable_media=( + master_config.grpo.deduplicate_multimodal_data + ), + ) + ) + print_multimodal_payload_metrics( + collect_multimodal_payload_metrics( + repeated_batch, + "prompt_repeat", + enabled=master_config.grpo.debug_payload_metrics, ) ) # Convert LLMMessageLogType to FlatMessagesType for generation @@ -2849,7 +2950,9 @@ def grpo_train( } ) calibration_data.update( - calib_flat.get_multimodal_dict(as_tensors=False) + calib_flat.get_multimodal_dict( + as_tensors=False, pixel_dtype=torch.bfloat16 + ) ) calibration_data.to("cpu") kv_scales_cache = policy.calibrate_qkv_fp8_scales( @@ -2911,6 +3014,12 @@ def grpo_train( mask_env_flagged_samples=should_mask_flagged_samples( master_config.env ), + deduplicate_multimodal_data=( + master_config.grpo.deduplicate_multimodal_data + ), + debug_payload_metrics=( + master_config.grpo.debug_payload_metrics + ), ) input_ids = nemo_gym_rollout_result.input_ids repeated_batch = nemo_gym_rollout_result.final_batch @@ -2932,6 +3041,9 @@ def grpo_train( ], max_rollout_turns=master_config.grpo.max_rollout_turns, greedy=False, + deduplicate_multimodal_data=( + master_config.grpo.deduplicate_multimodal_data + ), ) else: repeated_batch, rollout_metrics = run_multi_turn_rollout( @@ -2944,6 +3056,9 @@ def grpo_train( ], max_rollout_turns=master_config.grpo.max_rollout_turns, greedy=False, + deduplicate_multimodal_data=( + master_config.grpo.deduplicate_multimodal_data + ), ) policy_generation.finish_generation() # Collect generation logger metrics for performance reporting after each generation step @@ -3111,9 +3226,16 @@ def grpo_train( # this will be mini-batched inside the policy, so maintain the packed multimodal structure # This is also used to populate part of the downstream logprob calculation data extra_multimodal_data = flat_messages.get_multimodal_dict( - as_tensors=False + as_tensors=False, pixel_dtype=torch.bfloat16 ) train_data.update(extra_multimodal_data) + print_multimodal_payload_metrics( + collect_multimodal_payload_metrics( + train_data, + "rollout_to_policy", + enabled=master_config.grpo.debug_payload_metrics, + ) + ) # Router replay (R3) on the legacy data_plane.enabled=false # driver path: routed_experts already rides flat_messages # (attached to message_log during rollout, then batched into @@ -3286,12 +3408,19 @@ def grpo_train( ) early_stop_message: Optional[str] = None - # Run validation if it's a validation step or last step with val_at_end - if ( + should_run_validation = ( val_period > 0 and (total_steps + 1) >= val_start_at and (total_steps + 1) % val_period == 0 - ) or (val_at_end and is_last_step): + ) or (val_at_end and is_last_step) + + # Keep training and validation traffic in separate metric intervals. + payload_metrics: dict[str, int | float] = {} + if master_config.grpo.debug_payload_metrics: + payload_metrics = drain_multimodal_payload_metrics() + + # Run validation if it's a validation step or last step with val_at_end + if should_run_validation: memory_tracker.snapshot_start_of_stage("Validation", dir()) if NEED_REFIT and POLICY_GENERATION_STALE: refit_metrics = refit_policy_generation( @@ -3314,6 +3443,7 @@ def grpo_train( step=total_steps + 1, master_config=master_config, logger=logger, + processor=processor, ) policy_generation.finish_generation() logger.log_metrics( @@ -3322,6 +3452,14 @@ def grpo_train( logger.log_metrics( val_metrics, total_steps + 1, prefix="validation" ) + if master_config.grpo.debug_payload_metrics: + validation_payload_metrics = drain_multimodal_payload_metrics() + if validation_payload_metrics: + logger.log_metrics( + validation_payload_metrics, + total_steps + 1, + prefix="validation", + ) early_stop_message = _validation_early_stop_message( val_metrics, stop_at_validation_threshold, @@ -3662,6 +3800,9 @@ def grpo_train( train_results, metrics, timing_metrics, master_config ) + if payload_metrics: + logger.log_metrics(payload_metrics, total_steps + 1, prefix="") + if refit_metrics: logger.log_metrics(refit_metrics, total_steps + 1, prefix="refit") logger.log_metrics(metrics, total_steps + 1, prefix="train") @@ -3732,6 +3873,7 @@ def validate( step: int, master_config: MasterConfig, logger: Optional[Logger] = None, + processor: Optional[AutoProcessor] = None, ) -> tuple[dict[str, Any], dict[str, Any]]: """Run validation on the validation dataset.""" if val_dataloader is None: @@ -3768,6 +3910,8 @@ def validate( # Use async rollouts when enabled by config/backend defaults. # We cascade NeMo-Gym first since NeMo-Gym also uses async rollouts. if _should_use_nemo_gym(master_config): + if master_config.grpo.deduplicate_multimodal_data: + attach_initial_nemo_gym_image_payloads(val_batch, processor) generation_config = master_config.policy["generation"] # Validation-only sampling (e.g. near-greedy validation); # defaults to the train profile via the exemplar YAML @@ -3797,6 +3941,10 @@ def validate( mask_env_flagged_samples=should_mask_flagged_samples( master_config.env ), + deduplicate_multimodal_data=( + master_config.grpo.deduplicate_multimodal_data + ), + debug_payload_metrics=master_config.grpo.debug_payload_metrics, ) val_batch = nemo_gym_rollout_result.final_batch gen_metrics = nemo_gym_rollout_result.rollout_metrics @@ -3810,6 +3958,9 @@ def validate( max_seq_len=master_config.policy["max_total_sequence_length"], max_rollout_turns=master_config.grpo.max_rollout_turns, greedy=False, + deduplicate_multimodal_data=( + master_config.grpo.deduplicate_multimodal_data + ), ) else: val_batch, gen_metrics = run_multi_turn_rollout( @@ -3820,6 +3971,9 @@ def validate( max_seq_len=master_config.policy["max_total_sequence_length"], max_rollout_turns=master_config.grpo.max_rollout_turns, greedy=False, + deduplicate_multimodal_data=( + master_config.grpo.deduplicate_multimodal_data + ), ) total_rewards.extend(val_batch["total_reward"].tolist()) @@ -3975,6 +4129,7 @@ def async_grpo_train( max_trajectory_age_steps: int = 1, teacher_worker_groups: Optional[dict[str, Any]] = None, alias_to_group_alias: Optional[dict[str, str]] = None, + processor: Optional[AutoProcessor] = None, ) -> None: """Run asynchronous GRPO training with replay buffer. @@ -3992,6 +4147,8 @@ def async_grpo_train( grpo_save_state: Training state master_config: Master configuration max_trajectory_age_steps: Maximum age (in training steps) for trajectories to be used in training + processor: Optional multimodal processor used to attach compact policy + media to NeMo Gym prompt rows. """ # Ensure we are running with a compatible async generation backend. # Async GRPO (with in-flight weight updates) supports vLLM, Megatron, and TRT-LLM; @@ -4118,18 +4275,15 @@ def async_grpo_train( ) last_checkpoint_path = checkpointer.get_latest_checkpoint_path() - replay_buffer_state = None + replay_buffer_restore_metadata: dict[str, int] | None = None rollouts_state = None if last_checkpoint_path is not None: replay_buffer_path = os.path.join(last_checkpoint_path, "replay_buffer.pt") if os.path.exists(replay_buffer_path): print(f"šŸ“¦ Restoring replay buffer from checkpoint: {replay_buffer_path}") - # weights_only=False: trajectories are pickled BatchedDataDict/dicts, - # not plain tensors. The checkpoint is a trusted same-job artifact. - replay_buffer_state = torch.load(replay_buffer_path, weights_only=False) - ray.get( - replay_buffer.load_state_dict.remote( - replay_buffer_state, + replay_buffer_restore_metadata = ray.get( + replay_buffer.load_from_path.remote( + replay_buffer_path, num_prompts_per_step=num_prompts_per_step, current_training_step=step, max_age_steps=max_trajectory_age_steps, @@ -4149,8 +4303,13 @@ def async_grpo_train( next_nemo_gym_task_index = _get_next_nemo_gym_task_index( rollouts_state=rollouts_state, - replay_buffer_state=replay_buffer_state, + replay_buffer_state=None, ) + if replay_buffer_restore_metadata is not None: + next_nemo_gym_task_index = max( + next_nemo_gym_task_index, + replay_buffer_restore_metadata[NEXT_NEMO_GYM_TASK_INDEX_KEY], + ) _tc_py_exec = get_actor_python_env( "nemo_rl.algorithms.async_utils.AsyncTrajectoryCollector" @@ -4188,6 +4347,7 @@ def async_grpo_train( alias_to_group_alias=alias_to_group_alias, on_policy_distillation_cfg=opd_module._opd_cfg(master_config), next_nemo_gym_task_index=next_nemo_gym_task_index, + processor=processor, ) # Start trajectory collection in background @@ -4237,7 +4397,7 @@ def async_grpo_train( if val_at_start and step == 0: print("\nšŸ” Running initial validation...") # Pause trajectory collection during initial validation - trajectory_collector.pause.remote() + ray.get(trajectory_collector.pause.remote()) initial_val_metrics: Optional[dict[str, Any]] = None try: @@ -4249,11 +4409,20 @@ def async_grpo_train( step=0, master_config=master_config, logger=logger, + processor=processor, ) initial_val_metrics = val_metrics policy_generation.finish_generation() logger.log_metrics(val_metrics, step, prefix="validation") logger.log_metrics(validation_timings, step, prefix="timing/validation") + if master_config.grpo.debug_payload_metrics: + validation_payload_metrics = drain_multimodal_payload_metrics() + if validation_payload_metrics: + logger.log_metrics( + validation_payload_metrics, + step, + prefix="validation", + ) print("āœ… Initial validation completed successfully") except Exception as e: print(f"āŒ Initial validation failed: {e}") @@ -4409,6 +4578,14 @@ def async_grpo_train( max_age_steps=max_trajectory_age_steps, ) ) + if sample_result is not None: + print_multimodal_payload_metrics( + collect_multimodal_payload_metrics( + sample_result, + "replay_sample", + enabled=master_config.grpo.debug_payload_metrics, + ) + ) if ( sample_result is None @@ -4490,7 +4667,16 @@ def async_grpo_train( # Concatenate per-prompt groups into a single training batch per_prompt_batches = [t["batch"] for t in trajectories] - repeated_batch = BatchedDataDict.from_batches(per_prompt_batches) + normalize_sparse_media = _should_normalize_sparse_replay_media( + per_prompt_batches, + deduplicate_multimodal_data=( + master_config.grpo.deduplicate_multimodal_data + ), + ) + repeated_batch = BatchedDataDict.from_batches( + per_prompt_batches, + allow_missing_packed_tensors=normalize_sparse_media, + ) # Teacher logprobs are stored in batch dict by collection-time # computation and padded by from_batches. Extract here. @@ -4605,6 +4791,13 @@ def async_grpo_train( repeated_batch, master_config.policy, ) + print_multimodal_payload_metrics( + collect_multimodal_payload_metrics( + train_data, + "rollout_to_policy_async", + enabled=master_config.grpo.debug_payload_metrics, + ) + ) train_data.to("cpu") # Training phase (same as sync version) @@ -4779,17 +4972,30 @@ def async_grpo_train( # Validation val_metrics, validation_timings = None, None is_last_step = step + 1 == master_config.grpo.max_num_steps - - # Run validation if it's a validation step or last step with val_at_end - if ( + should_run_validation = ( val_period > 0 and (step + 1) >= val_start_at and (step + 1) % val_period == 0 - ) or (val_at_end and is_last_step): - with timer.time("idle/validation"): - # Pause trajectory collection during validation to reduce memory pressure - trajectory_collector.pause.remote() + ) or (val_at_end and is_last_step) + + payload_metrics: dict[str, int | float] = {} + if should_run_validation: + # Stop new dispatch before separating the training and + # validation payload-metric intervals. + ray.get(trajectory_collector.pause.remote()) + if master_config.grpo.debug_payload_metrics: + payload_metrics = merge_multimodal_payload_metrics( + [ + drain_multimodal_payload_metrics(), + ray.get( + trajectory_collector.drain_payload_metrics.remote() + ), + ] + ) + # Run validation if it's a validation step or last step with val_at_end + if should_run_validation: + with timer.time("idle/validation"): if NEED_REFIT and POLICY_GENERATION_STALE: refit_metrics = refit_policy_generation( policy, @@ -4807,12 +5013,23 @@ def async_grpo_train( step=step + 1, master_config=master_config, logger=logger, + processor=processor, ) policy_generation.finish_generation() logger.log_metrics( validation_timings, step + 1, prefix="timing/validation" ) logger.log_metrics(val_metrics, step + 1, prefix="validation") + if master_config.grpo.debug_payload_metrics: + validation_payload_metrics = ( + drain_multimodal_payload_metrics() + ) + if validation_payload_metrics: + logger.log_metrics( + validation_payload_metrics, + step + 1, + prefix="validation", + ) early_stop_message = _validation_early_stop_message( val_metrics, stop_at_validation_threshold, @@ -4823,8 +5040,6 @@ def async_grpo_train( print(early_stop_message, flush=True) # Explicit GPU memory cleanup after validation in async mode - import gc - gc.collect() torch.cuda.empty_cache() @@ -4997,15 +5212,10 @@ def async_grpo_train( actual_dataloader_state, os.path.join(checkpoint_path, "train_dataloader.pt"), ) - print("šŸ“¦ Saving replay buffer state...") - replay_buffer_state = ray.get(replay_buffer.state_dict.remote()) - torch.save( - replay_buffer_state, - os.path.join(checkpoint_path, "replay_buffer.pt"), - ) - print( - "āœ… Saved replay buffer with " - f"{len(replay_buffer_state['trajectories'])} trajectories" + _save_async_replay_buffer_checkpoint( + replay_buffer, + checkpoint_path, + master_config.checkpointing, ) rollouts_state = ray.get( trajectory_collector.get_rollouts_state.remote() @@ -5149,6 +5359,16 @@ def async_grpo_train( merged_efficiency, total_wall_time, step + 1 ) + if master_config.grpo.debug_payload_metrics and not should_run_validation: + payload_metrics = merge_multimodal_payload_metrics( + [ + drain_multimodal_payload_metrics(), + ray.get(trajectory_collector.drain_payload_metrics.remote()), + ] + ) + if payload_metrics: + logger.log_metrics(payload_metrics, step + 1, prefix="") + if refit_metrics: logger.log_metrics(refit_metrics, step + 1, prefix="refit") logger.log_metrics(performance_metrics, step + 1, prefix="performance") diff --git a/nemo_rl/data/llm_message_utils.py b/nemo_rl/data/llm_message_utils.py index c3e6bf76586..d61ff4b7bda 100644 --- a/nemo_rl/data/llm_message_utils.py +++ b/nemo_rl/data/llm_message_utils.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -26,9 +26,8 @@ ) from nemo_rl.data.multimodal_utils import ( PackedTensor, - get_dim_to_pack_along, + extract_multimodal_model_inputs, get_multimodal_default_settings_from_processor, - get_multimodal_keys_from_processor, load_media_from_message, ) from nemo_rl.distributed.batched_data_dict import BatchedDataDict @@ -115,9 +114,22 @@ def message_log_to_flat_messages( f"tensors for {key=} must have same number of dimensions: {[t.shape for t in result[key]]}" ) from e raise - elif result[key] and isinstance(result[key][0], PackedTensor): + packed_values = [ + value for value in result[key] if isinstance(value, PackedTensor) + ] + if packed_values: + invalid_values = [ + value + for value in result[key] + if value is not None and not isinstance(value, PackedTensor) + ] + if invalid_values: + raise TypeError( + f"Packed multimodal key {key!r} also contains non-packed " + f"values: {[type(value).__name__ for value in invalid_values]}" + ) try: - concat[key] = PackedTensor.concat(result[key]) + concat[key] = PackedTensor.merge_segments(packed_values) except Exception as e: raise RuntimeError( f"Error concatenating packed multimodal data for {key=}" @@ -363,26 +375,31 @@ def batched_message_log_to_flat_message( result = BatchedDataDict() for key in all_keys: values = [seq.get(key) for seq in sequenced_lists] - packed_template = next( - (value for value in values if isinstance(value, PackedTensor)), None - ) - if packed_template is not None: - if any( - value is not None and not isinstance(value, PackedTensor) + packed_values = [value for value in values if isinstance(value, PackedTensor)] + # Preserve one logical row for conversations missing this media key. + # Async replay may concatenate text-only and multimodal prompt groups in + # either order, so the first row cannot determine the value type. + if packed_values: + invalid_values = [ + value for value in values - ): + if value is not None and not isinstance(value, PackedTensor) + ] + if invalid_values: raise TypeError( - f"Expected PackedTensor or None for {key=}, " - f"got {[type(value).__name__ for value in values]}" + f"Packed multimodal key {key!r} also contains non-packed " + f"values: {[type(value).__name__ for value in invalid_values]}" ) - filled_packed_values = cast( - list[PackedTensor], - [ - PackedTensor.empty_like(packed_template) if value is None else value - for value in values - ], - ) - result[key] = PackedTensor.flattened_concat(filled_packed_values) + template = packed_values[0] + aligned_values = [ + ( + value + if isinstance(value, PackedTensor) + else PackedTensor.empty_rows_like(template, 1) + ) + for value in values + ] + result[key] = PackedTensor.flattened_concat(aligned_values) continue if not values or not isinstance(values[0], Tensor): result[key] = values @@ -477,7 +494,6 @@ def get_formatted_message_log( list[dict[str, str]], message_log ) # we just use the str:str parts here - multimodal_keys = get_multimodal_keys_from_processor(tokenizer) multimodal_load_kwargs = get_multimodal_default_settings_from_processor(tokenizer) def _format_content_helper( @@ -642,20 +658,9 @@ def _format_content_helper( ) new_message["token_ids"] = processed_chunk["input_ids"][0] - # add all vlm keys to the message - for key in multimodal_keys: - if key in processed_chunk: - # token_type_ids and mm_token_type_ids are sequence-length tensors - # (one label per token), not visual patch tensors. They must be - # stored as plain tensors and padded like input_ids rather than - # packed as multimodal data. This mirrors processors.py behavior. - if key in ("token_type_ids", "mm_token_type_ids"): - new_message[key] = processed_chunk[key][0] - else: - new_message[key] = PackedTensor( - processed_chunk[key], - dim_to_pack=get_dim_to_pack_along(tokenizer, key), - ) + new_message.update( + extract_multimodal_model_inputs(tokenizer, dict(processed_chunk)) + ) if len(new_message["token_ids"]) == 0: # if there is an empty message, the empty `token_ids` tensor ends up being in fp32, diff --git a/nemo_rl/data/multimodal_utils.py b/nemo_rl/data/multimodal_utils.py index 368608f5909..49146e9bc8c 100644 --- a/nemo_rl/data/multimodal_utils.py +++ b/nemo_rl/data/multimodal_utils.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -16,7 +16,9 @@ import inspect import logging import re +import uuid from collections import defaultdict +from copy import deepcopy from io import BytesIO from typing import Any, Optional, Union @@ -28,6 +30,12 @@ from transformers.audio_utils import load_audio from transformers.video_utils import load_video +VLLM_MULTIMODAL_DATA_KEYS = frozenset({"vllm_images", "vllm_videos", "vllm_audios"}) +NATIVE_MULTIMODAL_KEYS = frozenset({"vllm_content", *VLLM_MULTIMODAL_DATA_KEYS}) +MULTIMODAL_CONTENT_TYPES = frozenset( + {"input_image", "image", "image_url", "video", "audio"} +) + # List of allowed placeholder strings for different media types in the dataset string # e.g. "This is an example of " MEDIA_TAGS = { @@ -88,14 +96,21 @@ def uses_image_placeholder(processor: Any) -> bool: class PackedTensor: - """Wrapper around a list of torch tensors and a dimension along which to pack the tensors. + """A logical batch of rows backed by packable tensor segments. - This class is used to wrap a list of tensors along with a `dim_to_pack` parameter. - It can be used for data that can be packed along different dimensions (such as multimodal data). + The default representation is intentionally the legacy one: every entry in + ``tensors`` is one logical row and no deduplication metadata is allocated. + ``enable_deduplication`` adds stable provenance to the physical segments. + Operations that combine or slice dedup-enabled values then use a CSR-like + logical-row mapping: - `dim_to_pack` is used to specify the dimension along which to pack the tensors. + - ``_row_offsets`` partitions the flattened logical segment references. + - ``_segment_indices`` maps each logical segment reference to ``tensors``. + - ``_segment_provenance`` is stable across deepcopy/pickle and is the only + evidence used to re-intern physical segments. - The list of tensors can be returned as a single packed tensor by calling `as_tensor` which will concatenate the tensors along the `dim_to_pack` dimension. + Prompt identity is deliberately absent: belonging to the same prompt group + makes media a candidate for sharing, but never proves media equality. """ def __init__( @@ -104,6 +119,9 @@ def __init__( dim_to_pack: int, *, pad_to_max_shape: bool = False, + _row_offsets: Optional[list[int]] = None, + _segment_indices: Optional[list[int]] = None, + _segment_provenance: Optional[list[bytes]] = None, ) -> None: """Wrap per-item tensors for concatenation along ``dim_to_pack``. @@ -119,9 +137,10 @@ def __init__( if isinstance(tensors, torch.Tensor): self.tensors: list[Optional[torch.Tensor]] = [tensors] elif isinstance(tensors, list): - assert len(tensors) > 0, ( - "Input tensors to PackedTensor must be a non-empty list" - ) + if not tensors and _row_offsets is None: + raise AssertionError( + "Input tensors to PackedTensor must be a non-empty list" + ) self.tensors: list[Optional[torch.Tensor]] = tensors else: raise ValueError( @@ -129,6 +148,163 @@ def __init__( ) self.dim_to_pack = dim_to_pack self.pad_to_max_shape = pad_to_max_shape + if (_row_offsets is None) != (_segment_indices is None): + raise ValueError( + "_row_offsets and _segment_indices must either both be set or both be None" + ) + if _row_offsets is not None: + if not _row_offsets or _row_offsets[0] != 0: + raise ValueError("_row_offsets must start with 0") + if any( + current > following + for current, following in zip(_row_offsets, _row_offsets[1:]) + ): + raise ValueError("_row_offsets must be non-decreasing") + assert _segment_indices is not None + if _row_offsets[-1] != len(_segment_indices): + raise ValueError( + "_row_offsets must end at the number of logical segment references" + ) + if _segment_indices and ( + min(_segment_indices) < 0 or max(_segment_indices) >= len(self.tensors) + ): + raise ValueError( + "_segment_indices cannot reference an out-of-range physical segment" + ) + if _segment_provenance is not None and len(_segment_provenance) != len( + self.tensors + ): + raise ValueError( + "_segment_provenance must have one entry per physical segment" + ) + self._row_offsets = _row_offsets + self._segment_indices = _segment_indices + self._segment_provenance = _segment_provenance + + def __setstate__(self, state: dict[str, Any]) -> None: + """Restore both current and pre-deduplication pickled instances.""" + self.__dict__.update(state) + self.__dict__.setdefault("_row_offsets", None) + self.__dict__.setdefault("_segment_indices", None) + self.__dict__.setdefault("_segment_provenance", None) + + @property + def deduplication_enabled(self) -> bool: + """Whether this value carries stable physical-segment provenance.""" + return self._segment_provenance is not None + + @property + def logical_segment_count(self) -> int: + """Number of segment occurrences after logical expansion.""" + if self._segment_indices is not None: + return len(self._segment_indices) + return len(self.tensors) + + def logical_segment_counts_by_row(self) -> list[int]: + """Return the number of non-empty media segments in each logical row.""" + if self._row_offsets is None: + return [int(tensor is not None) for tensor in self.tensors] + assert self._segment_indices is not None + return [ + sum( + self.tensors[physical_index] is not None + for physical_index in self._segment_indices[ + self._row_offsets[row] : self._row_offsets[row + 1] + ] + ) + for row in range(len(self)) + ] + + def iter_logical_segments(self): + """Yield physical tensor segments in logical row/segment order.""" + if self._segment_indices is None: + yield from self.tensors + return + for physical_index in self._segment_indices: + yield self.tensors[physical_index] + + def enable_deduplication(self) -> "PackedTensor": + """Assign stable provenance lazily without changing logical contents.""" + if self._segment_provenance is None: + self._segment_provenance = [ + uuid.uuid4().bytes for _ in range(len(self.tensors)) + ] + return self + + def repeat_interleave(self, num_repeats: int) -> "PackedTensor": + """Repeat logical rows while retaining one copy of each physical segment.""" + if not self.deduplication_enabled: + raise ValueError( + "PackedTensor repeat_interleave requires deduplication to be enabled" + ) + if num_repeats < 0: + raise ValueError("num_repeats must be non-negative") + if num_repeats == 0: + return PackedTensor( + [], + self.dim_to_pack, + pad_to_max_shape=self.pad_to_max_shape, + _row_offsets=[0], + _segment_indices=[], + _segment_provenance=[], + ) + segment_indices = [] + row_offsets = [0] + for row in range(len(self)): + row_segments = self._row_segment_indices(row) + for _ in range(num_repeats): + segment_indices.extend(row_segments) + row_offsets.append(len(segment_indices)) + return PackedTensor( + list(self.tensors), + self.dim_to_pack, + pad_to_max_shape=self.pad_to_max_shape, + _row_offsets=row_offsets, + _segment_indices=segment_indices, + _segment_provenance=list(self._segment_provenance or []), + ) + + def _row_segment_indices(self, row: int) -> list[int]: + if self._row_offsets is None: + return [row] + assert self._segment_indices is not None + return self._segment_indices[ + self._row_offsets[row] : self._row_offsets[row + 1] + ] + + def __deepcopy__(self, memo: dict[int, Any]) -> "PackedTensor": + """Share immutable media segments only for an explicitly enabled value.""" + if self._row_offsets is None and not self.deduplication_enabled: + copied = PackedTensor( + [deepcopy(item, memo) for item in self.tensors], + self.dim_to_pack, + pad_to_max_shape=self.pad_to_max_shape, + ) + else: + copied = PackedTensor( + ( + list(self.tensors) + if self.deduplication_enabled + else [deepcopy(item, memo) for item in self.tensors] + ), + self.dim_to_pack, + pad_to_max_shape=self.pad_to_max_shape, + _row_offsets=( + list(self._row_offsets) if self._row_offsets is not None else None + ), + _segment_indices=( + list(self._segment_indices) + if self._segment_indices is not None + else None + ), + _segment_provenance=( + list(self._segment_provenance) + if self._segment_provenance is not None + else None + ), + ) + memo[id(self)] = copied + return copied def as_tensor( self, device: Optional[torch.device] = None @@ -138,7 +314,10 @@ def as_tensor( for i, item in enumerate(self.tensors): if item is not None: self.tensors[i] = item.to(device) - non_none_tensors = [t for t in self.tensors if t is not None] + tensors = self.tensors + if self._segment_indices is not None: + tensors = [self.tensors[index] for index in self._segment_indices] + non_none_tensors = [t for t in tensors if t is not None] if len(non_none_tensors) == 0: return None @@ -188,7 +367,8 @@ def pad_to_batch_shape(tensor: torch.Tensor) -> torch.Tensor: return torch.cat(non_none_tensors, dim=self.dim_to_pack).to(device) def __len__(self) -> int: - # this is the number of tensors in this data wrapper + if self._row_offsets is not None: + return len(self._row_offsets) - 1 return len(self.tensors) def to(self, device: str | torch.device) -> "PackedTensor": @@ -197,20 +377,107 @@ def to(self, device: str | torch.device) -> "PackedTensor": ] return self + def to_dtype(self, dtype: torch.dtype) -> "PackedTensor": + """Return a dtype-converted value without expanding logical segments. + + Dtype conversion creates new physical tensor values, so deduplicated + inputs receive new provenance. The logical row-to-segment mapping is + preserved exactly. + """ + if all(item is None or item.dtype == dtype for item in self.tensors): + return self + + return PackedTensor( + [ + item.to(dtype=dtype) if item is not None else None + for item in self.tensors + ], + self.dim_to_pack, + pad_to_max_shape=self.pad_to_max_shape, + _row_offsets=( + list(self._row_offsets) if self._row_offsets is not None else None + ), + _segment_indices=( + list(self._segment_indices) + if self._segment_indices is not None + else None + ), + _segment_provenance=( + [uuid.uuid4().bytes for _ in self.tensors] + if self._segment_provenance is not None + else None + ), + ) + def slice(self, indices: Union[list[int], torch.Tensor]) -> "PackedTensor": idx = indices.tolist() if isinstance(indices, torch.Tensor) else indices - tensors = [self.tensors[i] for i in idx] + if not self.deduplication_enabled and self._row_offsets is None: + tensors = [self.tensors[i] for i in idx] + return PackedTensor( + tensors, + self.dim_to_pack, + pad_to_max_shape=self.pad_to_max_shape, + ) + + physical_remap: dict[int, int] = {} + tensors: list[Optional[torch.Tensor]] = [] + provenances: list[bytes] = [] + segment_indices: list[int] = [] + row_offsets = [0] + for row in idx: + if row < 0: + row += len(self) + if not 0 <= row < len(self): + raise IndexError(f"PackedTensor row index {row} is out of range") + for physical_index in self._row_segment_indices(row): + if physical_index not in physical_remap: + physical_remap[physical_index] = len(tensors) + tensors.append(self.tensors[physical_index]) + if self._segment_provenance is not None: + provenances.append(self._segment_provenance[physical_index]) + segment_indices.append(physical_remap[physical_index]) + row_offsets.append(len(segment_indices)) return PackedTensor( tensors, self.dim_to_pack, pad_to_max_shape=self.pad_to_max_shape, + _row_offsets=row_offsets, + _segment_indices=segment_indices, + _segment_provenance=( + provenances if self._segment_provenance is not None else None + ), ) @classmethod def empty_like(cls, other: "PackedTensor") -> "PackedTensor": - """Return a new PackedTensor with same length and dim_to_pack as `other`, with all entries None.""" + """Return empty logical rows matching ``other``.""" + return cls.empty_rows_like(other, len(other)) + + @classmethod + def empty_rows_like(cls, other: "PackedTensor", num_rows: int) -> "PackedTensor": + """Return ``num_rows`` logical rows containing no media segments.""" + if num_rows < 0: + raise ValueError("num_rows must be non-negative") + if other.deduplication_enabled or other._row_offsets is not None: + return cls( + [], + other.dim_to_pack, + pad_to_max_shape=other.pad_to_max_shape, + _row_offsets=[0] * (num_rows + 1), + _segment_indices=[], + _segment_provenance=[], + ) + if num_rows == 0: + return cls( + [], + other.dim_to_pack, + pad_to_max_shape=other.pad_to_max_shape, + _row_offsets=[0], + _segment_indices=[], + _segment_provenance=None, + ) return cls( - [None] * len(other.tensors), + [None] * num_rows, other.dim_to_pack, pad_to_max_shape=other.pad_to_max_shape, ) @@ -245,7 +512,53 @@ def concat(cls, from_packed_tensors: list["PackedTensor"]) -> "PackedTensor": assert len(set(pad_to_max_shapes)) == 1, ( "All packed tensors must have the same pad_to_max_shape setting" ) - # concatenate the tensors + if any( + packed_tensor.deduplication_enabled + or packed_tensor._row_offsets is not None + for packed_tensor in from_packed_tensors + ): + tensors: list[Optional[torch.Tensor]] = [] + provenances: list[bytes] = [] + provenance_to_physical: dict[bytes, int] = {} + segment_indices: list[int] = [] + row_offsets = [0] + + for packed_tensor in from_packed_tensors: + physical_remap: dict[int, int] = {} + for physical_index, tensor in enumerate(packed_tensor.tensors): + provenance = ( + packed_tensor._segment_provenance[physical_index] + if packed_tensor._segment_provenance is not None + else None + ) + if provenance is not None and provenance in provenance_to_physical: + new_index = provenance_to_physical[provenance] + else: + new_index = len(tensors) + tensors.append(tensor) + if provenance is None: + provenance = uuid.uuid4().bytes + provenances.append(provenance) + provenance_to_physical[provenance] = new_index + physical_remap[physical_index] = new_index + + for row in range(len(packed_tensor)): + segment_indices.extend( + physical_remap[index] + for index in packed_tensor._row_segment_indices(row) + ) + row_offsets.append(len(segment_indices)) + + return cls( + tensors, + dim_to_packs[0], + pad_to_max_shape=pad_to_max_shapes[0], + _row_offsets=row_offsets, + _segment_indices=segment_indices, + _segment_provenance=provenances, + ) + + # Legacy flag-off behavior: concatenate the tensors without metadata. tensors = [] for packed_tensor in from_packed_tensors: tensors.extend(packed_tensor.tensors) @@ -256,6 +569,29 @@ def concat(cls, from_packed_tensors: list["PackedTensor"]) -> "PackedTensor": pad_to_max_shape=pad_to_max_shapes[0], ) + @classmethod + def merge_segments( + cls, from_packed_tensors: list["PackedTensor"] + ) -> "PackedTensor": + """Merge message-turn values into one logical conversation row.""" + if not any( + packed_tensor.deduplication_enabled + or packed_tensor._row_offsets is not None + for packed_tensor in from_packed_tensors + ): + return cls.concat(from_packed_tensors) + + concatenated = cls.concat(from_packed_tensors) + assert concatenated._segment_indices is not None + return cls( + concatenated.tensors, + concatenated.dim_to_pack, + pad_to_max_shape=concatenated.pad_to_max_shape, + _row_offsets=[0, len(concatenated._segment_indices)], + _segment_indices=concatenated._segment_indices, + _segment_provenance=concatenated._segment_provenance, + ) + @classmethod def flattened_concat( cls, from_packed_tensors: list["PackedTensor"] @@ -290,6 +626,12 @@ def flattened_concat( assert len(set(pad_to_max_shapes)) == 1, ( "All packed tensors must have the same pad_to_max_shape setting" ) + if any( + packed_tensor.deduplication_enabled + or packed_tensor._row_offsets is not None + for packed_tensor in from_packed_tensors + ): + return cls.concat(from_packed_tensors) tensors = [p.as_tensor() for p in from_packed_tensors] return cls( tensors, @@ -374,6 +716,95 @@ def get_dim_to_pack_along(processor, key: str) -> int: return 0 +def get_pad_to_max_shape(processor: Any, key: str) -> bool: + """Return whether a processor input must pad non-packing dimensions.""" + return uses_image_placeholder(processor) and key == "pixel_values" + + +def extract_multimodal_model_inputs( + processor: Any, processed: dict[str, Any] +) -> dict[str, PackedTensor | torch.Tensor]: + """Extract packed media inputs and sequence-aligned auxiliary tensors.""" + processed = dict(processed) + if ( + uses_image_placeholder(processor) + and "pixel_values" in processed + and "imgs_sizes" not in processed + and processed["pixel_values"].ndim == 4 + ): + pixel_values = processed["pixel_values"] + num_tiles, _, height, width = pixel_values.shape + processed["imgs_sizes"] = torch.tensor( + [[height, width]] * num_tiles, + dtype=torch.long, + ) + if "imgs_sizes" in processed and "num_frames" not in processed: + processed["num_frames"] = torch.ones( + len(processed["imgs_sizes"]), + dtype=torch.long, + ) + + input_ids = processed.get("input_ids") + if input_ids is None: + raise ValueError("Processor output is missing input_ids.") + if not isinstance(input_ids, torch.Tensor) or input_ids.ndim not in (1, 2): + raise ValueError( + "Processor input_ids must be a one- or two-dimensional torch.Tensor." + ) + if input_ids.ndim == 2 and input_ids.shape[0] != 1: + raise ValueError( + "Multimodal chat processing expects a single conversation, got " + f"input_ids shape {tuple(input_ids.shape)}." + ) + sequence_length = input_ids.shape[-1] + + extracted: dict[str, PackedTensor | torch.Tensor] = {} + multimodal_keys = list(get_multimodal_keys_from_processor(processor)) + for key in ("imgs_sizes", "num_frames"): + if key in processed and key not in multimodal_keys: + multimodal_keys.append(key) + for key in multimodal_keys: + if key not in processed: + continue + value = processed[key] + if not isinstance(value, torch.Tensor): + raise ValueError( + f"Processor model input {key!r} must be a torch.Tensor, got " + f"{type(value).__name__}." + ) + if key == "imgs_sizes": + value = value.to(dtype=torch.int32) + extracted[key] = PackedTensor( + value, + dim_to_pack=get_dim_to_pack_along(processor, key), + pad_to_max_shape=get_pad_to_max_shape(processor, key), + ) + + for key in ("token_type_ids", "mm_token_type_ids"): + if key not in processed: + continue + value = processed[key] + if not isinstance(value, torch.Tensor) or value.ndim not in (1, 2): + raise ValueError( + f"Processor sequence input {key!r} must be a one- or " + "two-dimensional torch.Tensor." + ) + if value.ndim == 2: + if value.shape[0] != 1: + raise ValueError( + f"Processor sequence input {key!r} must contain one " + f"conversation, got shape {tuple(value.shape)}." + ) + value = value[0] + if len(value) != sequence_length: + raise ValueError( + f"Processor sequence input {key!r} has length {len(value)}, " + f"but input_ids has length {sequence_length}." + ) + extracted[key] = value + return extracted + + def resolve_to_image(image_path_or_image: str | Image.Image) -> Image.Image: """Resolve the image path to a PIL.Image object. @@ -422,6 +853,69 @@ def image_to_data_url(image: Image.Image, fmt: str = "PNG") -> str: return f"data:image/{fmt.lower()};base64,{encoded}" +def extract_input_image_sources_from_responses_messages( + messages: Any, +) -> list[str | Image.Image]: + """Extract image sources from Responses-API messages in encounter order.""" + if not isinstance(messages, list): + return [] + + sources: list[str | Image.Image] = [] + for message in messages: + if not isinstance(message, dict): + continue + content = message.get("content") or [] + if not isinstance(content, list): + continue + for part in content: + if not isinstance(part, dict): + continue + if part.get("type") not in ("input_image", "image", "image_url"): + continue + source = part.get("image") or part.get("image_url") or part.get("url") + if isinstance(source, dict): + source = source.get("url") + if isinstance(source, (str, Image.Image)): + sources.append(source) + return sources + + +def extract_input_images_from_responses_messages( + messages: Any, +) -> list[Image.Image]: + """Load images from Responses-API input messages in encounter order.""" + return [ + resolve_to_image(source) + for source in extract_input_image_sources_from_responses_messages(messages) + ] + + +def attach_image_model_inputs_to_message( + message: dict[str, Any], + *, + images: list[Image.Image], + processor: Any, +) -> None: + """Attach processor-owned image tensors without replacing rollout tokens.""" + if not images or processor is None: + return + + image_token = getattr(processor, "image_token", "") + processed = processor( + text=image_token * len(images), + images=images, + return_tensors="pt", + ) + model_inputs = extract_multimodal_model_inputs(processor, dict(processed)) + message.update( + { + key: value + for key, value in model_inputs.items() + if isinstance(value, PackedTensor) + } + ) + + def encode_images_in_examples(nemo_gym_examples: list[dict]) -> list[dict]: """Replace local image paths in NeMo Gym examples with base64 data URLs. diff --git a/nemo_rl/distributed/batched_data_dict.py b/nemo_rl/distributed/batched_data_dict.py index eccdd09ad94..f270ef46fdc 100644 --- a/nemo_rl/distributed/batched_data_dict.py +++ b/nemo_rl/distributed/batched_data_dict.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -32,6 +32,8 @@ from typing_extensions import Self from nemo_rl.data.multimodal_utils import ( + MULTIMODAL_CONTENT_TYPES, + NATIVE_MULTIMODAL_KEYS, PackedTensor, ) from nemo_rl.data.packing import get_packer @@ -42,6 +44,50 @@ DictT = TypeVar("DictT", bound=Mapping[str, Any]) +_COUPLED_MULTIMODAL_KEYS = ( + ("pixel_values", "image_grid_thw"), + ("pixel_values", "imgs_sizes"), + ("pixel_values", "num_frames"), + ("pixel_values_videos", "video_grid_thw"), +) + + +def _prepare_multimodal_sharing( + value: Any, + *, + media_context: bool = False, +) -> dict[int, Any]: + """Enable PackedTensor provenance and return deepcopy memo entries. + + PackedTensor is an explicit multimodal type. Raw native-vLLM payloads are + shared only under named media keys or typed content parts. Containers are + still deep-copied so rollout rows may diverge safely. + """ + shared_leaves: dict[int, Any] = {} + + def visit(item: Any, in_media_context: bool = False) -> None: + if isinstance(item, PackedTensor): + item.enable_deduplication() + return + if isinstance(item, dict): + content_type = item.get("type") + typed_media = content_type in MULTIMODAL_CONTENT_TYPES + for key, child in item.items(): + visit( + child, + in_media_context or typed_media or key in NATIVE_MULTIMODAL_KEYS, + ) + return + if isinstance(item, (list, tuple)): + for child in item: + visit(child, in_media_context) + return + if in_media_context: + shared_leaves[id(item)] = item + + visit(value, media_context) + return shared_leaves + class SequencePackingArgs(TypedDict): """Configuration settings for sequence packing. @@ -77,6 +123,8 @@ class DynamicBatchingArgs(TypedDict): class BatchedDataDict(UserDict, Generic[DictT]): + _PIXEL_DTYPE_CAST_KEYS = frozenset({"pixel_values"}) + # keys that are model specific, but not part of the PackedTensor ADDITIONAL_OPTIONAL_KEY_TENSORS = [ "token_type_ids", # specific to gemma3 that tells where the image tokens are in the sequence, not required for llm-only inference/training @@ -91,12 +139,39 @@ def __init__(self, *args, **kwargs): self.elem_counts_per_gb = None def get_multimodal_dict( - self, as_tensors: bool = False, device: Optional[torch.device] = None + self, + as_tensors: bool = False, + device: Optional[torch.device] = None, + pixel_dtype: Optional[torch.dtype] = None, ) -> dict[str, Any]: - """Return a regular dict of tensors or packed multimodal data items.""" + """Return a regular dict of tensors or packed multimodal data items. + + ``pixel_dtype`` converts pixel tensors without materializing repeated + logical segments. This is used to reduce policy-bound Ray payloads. + """ + if as_tensors: + for value_key, metadata_key in _COUPLED_MULTIMODAL_KEYS: + value = self.data.get(value_key) + metadata = self.data.get(metadata_key) + if not isinstance(value, PackedTensor) or not isinstance( + metadata, PackedTensor + ): + continue + value_counts = value.logical_segment_counts_by_row() + metadata_counts = metadata.logical_segment_counts_by_row() + if value_counts != metadata_counts: + raise ValueError( + "Coupled multimodal keys must have the same ordered " + f"per-row segment counts, but {value_key!r} has " + f"{value_counts} and {metadata_key!r} has " + f"{metadata_counts}." + ) + multimodal_dict = {} for k, v in self.data.items(): if isinstance(v, PackedTensor): + if pixel_dtype is not None and k in self._PIXEL_DTYPE_CAST_KEYS: + v = v.to_dtype(pixel_dtype) multimodal_dict[k] = v.as_tensor(device=device) if as_tensors else v elif k in self.ADDITIONAL_OPTIONAL_KEY_TENSORS: multimodal_dict[k] = v @@ -108,6 +183,8 @@ def from_batches( cls: Type[Self], batches: Sequence[Mapping[Any, Any]], pad_value_dict: Optional[dict[str, int | float]] = None, + *, + allow_missing_packed_tensors: bool = False, ) -> Self: """Given a list of batches, stack the tensors/lists within and put them in a single dictionary. @@ -116,6 +193,9 @@ def from_batches( Args: batches (list[Dict]): A list of dictionaries, each containing a batch of data. pad_value_dict (Optional[dict[str, int]]): An optional dict mapping keys to non-default(0) padding values. + allow_missing_packed_tensors: Represent missing ``PackedTensor`` + media keys as empty logical rows. This is opt-in so ordinary + flag-off concatenation retains its strict key checks. Returns: BatchedDataDict: A new BatchedDataDict containing the stacked data. @@ -128,12 +208,26 @@ def from_batches( def batch_size(item: Mapping[Any, Any]) -> int: if not item: return 0 - value = next(iter(item.values())) - if isinstance(value, PackedTensor): - return len(value) - if isinstance(value, torch.Tensor): - return value.shape[0] - return len(value) + + if not allow_missing_packed_tensors: + # Preserve the legacy shared primitive exactly unless sparse + # PackedTensor normalization was explicitly requested. + return len(next(iter(item.values()))) + + sizes = set() + for value in item.values(): + if isinstance(value, PackedTensor): + sizes.add(len(value)) + elif isinstance(value, torch.Tensor): + sizes.add(value.shape[0]) + else: + sizes.add(len(value)) + if len(sizes) != 1: + raise ValueError( + "Source batch has inconsistent logical row counts: " + f"{sorted(sizes)}." + ) + return next(iter(sizes)) keys = sorted({key for item in batches for key in item}) for k in keys: @@ -143,12 +237,47 @@ def batch_size(item: Mapping[Any, Any]) -> int: if k not in item and batch_size(item) ] if missing_nonempty_batches: - raise KeyError( - f"Key {k!r} is missing from non-empty batches " - f"{missing_nonempty_batches}." - ) + present_values = [item[k] for item in batches if k in item] + if not ( + allow_missing_packed_tensors + and present_values + and all(isinstance(value, PackedTensor) for value in present_values) + ): + raise KeyError( + f"Key {k!r} is missing from non-empty batches " + f"{missing_nonempty_batches}." + ) + + template = present_values[0] + assert isinstance(template, PackedTensor) + list_of_tensors = [ + ( + item[k] + if k in item + else PackedTensor.empty_rows_like(template, batch_size(item)) + ) + for item in batches + if k in item or batch_size(item) + ] + else: + list_of_tensors = [item[k] for item in batches if k in item] - list_of_tensors = [item[k] for item in batches if k in item] + if allow_missing_packed_tensors and isinstance( + list_of_tensors[0], PackedTensor + ): + source_batches = [ + item for item in batches if k in item or batch_size(item) + ] + for batch_index, (item, packed_tensor) in enumerate( + zip(source_batches, list_of_tensors) + ): + expected_rows = batch_size(item) + if len(packed_tensor) != expected_rows: + raise ValueError( + f"PackedTensor key {k!r} has {len(packed_tensor)} " + f"logical rows in source batch {batch_index}, " + f"expected {expected_rows}." + ) if isinstance(list_of_tensors[0], list): tensor_or_list: list[Any] | torch.Tensor = [ @@ -622,7 +751,7 @@ def _get_padded_seqlen(seqlen: int) -> int: aggregated_shards[shard_idx][k] = ( PackedTensor.concat(packed_slices) if packed_slices - else PackedTensor.empty_like(v) + else PackedTensor.empty_rows_like(v, 0) ) else: shard_values = [] @@ -767,12 +896,21 @@ def slice(self, start: int, end: int) -> "SlicedDataDict": sliced_batch[k] = self.data[k][start:end] return sliced_batch - def repeat_interleave(self, num_repeats: int) -> Self: + def repeat_interleave( + self, + num_repeats: int, + *, + share_immutable_media: bool = False, + ) -> Self: """Repeats the batch num_repeats times. For each element in the batch, repeat each value num_repeats times. i.e: {"key": torch.tensor([1, 2, 3]), "other_key": [1, 2, 3]} -> {"key": torch.tensor([1, 1, 2, 2, 3, 3]), "other_key": [1, 1, 2, 2, 3, 3]} + + When ``share_immutable_media`` is enabled, only explicit multimodal + leaves share storage. Every surrounding row/message/content container + remains independent. """ repeated_batch: Self = type(self)() for k, v in self.data.items(): @@ -780,14 +918,30 @@ def repeat_interleave(self, num_repeats: int) -> Self: # For tensors, use repeat_interleave to repeat each element repeated_batch[k] = v.repeat_interleave(num_repeats, dim=0) elif isinstance(v, PackedTensor): - raise NotImplementedError( - "PackedTensor does not currently support repeat_interleave" + if not share_immutable_media: + raise NotImplementedError( + "PackedTensor does not currently support repeat_interleave " + "unless share_immutable_media is enabled" + ) + repeated_batch[k] = v.enable_deduplication().repeat_interleave( + num_repeats ) else: # For lists or other sequences, use a list comprehension to repeat each element - repeated_batch[k] = [ - deepcopy(item) for item in v for _ in range(num_repeats) - ] + repeated_items = [] + for item in v: + shared_leaves = ( + _prepare_multimodal_sharing( + item, + media_context=k in NATIVE_MULTIMODAL_KEYS, + ) + if share_immutable_media + else {} + ) + repeated_items.extend( + deepcopy(item, dict(shared_leaves)) for _ in range(num_repeats) + ) + repeated_batch[k] = repeated_items return repeated_batch def truncate_tensors(self, dim: int, truncated_len: int): @@ -861,9 +1015,9 @@ def size(self) -> int: """Get the batch size of the batch.""" # Get the first key and use its size as the batch size # This assumes all keys have the same batch size - key = next(iter(self.data)) if not self.data: return 0 + key = next(iter(self.data)) if not torch.is_tensor(self.data[key]): return len(self.data[key]) return self.data[key].shape[0] # type: ignore # it's a tensor here diff --git a/nemo_rl/environments/nemo_gym.py b/nemo_rl/environments/nemo_gym.py index 375d3b2d41d..458ebc3b72e 100644 --- a/nemo_rl/environments/nemo_gym.py +++ b/nemo_rl/environments/nemo_gym.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -17,6 +17,7 @@ import sys from collections import Counter from collections.abc import AsyncGenerator +from copy import deepcopy from pathlib import Path from typing import Any, Dict, List, NotRequired, Optional, TypedDict @@ -27,10 +28,9 @@ from transformers import PreTrainedTokenizerBase from nemo_rl.data.multimodal_utils import ( - PackedTensor, + attach_image_model_inputs_to_message, encode_images_in_examples, - get_dim_to_pack_along, - get_multimodal_keys_from_processor, + extract_input_image_sources_from_responses_messages, resolve_to_image, uses_image_placeholder, ) @@ -284,6 +284,48 @@ def _index_per_turn_images( return per_turn +def _image_sources_equal(left: Any, right: Any) -> bool: + return ( + left == right + if isinstance(left, str) and isinstance(right, str) + else left is right + ) + + +def _without_initial_image_sources( + messages: Any, initial_sources: list[Any] +) -> tuple[Any, bool]: + """Copy Responses messages and remove one ordered copy of initial images.""" + if not isinstance(messages, list): + return messages, False + + filtered = deepcopy(messages) + remaining_sources = list(initial_sources) + for message in filtered: + if not isinstance(message, dict): + continue + content = message.get("content") + if not isinstance(content, list): + continue + + filtered_content = [] + for part in content: + part_sources = extract_input_image_sources_from_responses_messages( + [{"content": [part]}] + ) + if ( + remaining_sources + and len(part_sources) == 1 + and _image_sources_equal(part_sources[0], remaining_sources[0]) + ): + remaining_sources.pop(0) + continue + filtered_content.append(part) + message["content"] = filtered_content + + return filtered, not remaining_sources + + def _attach_multimodal_data_to_user_message( user_message: dict, *, @@ -300,52 +342,11 @@ def _attach_multimodal_data_to_user_message( already contains expanded ``...*N...`` regions, and the processor would try to re-expand every embedded ````. """ - if not images or processor is None: - return - image_token = getattr(processor, "image_token", "") - processed = processor( - text=image_token * len(images), + attach_image_model_inputs_to_message( + user_message, images=images, - return_tensors="pt", + processor=processor, ) - uses_placeholder = uses_image_placeholder(processor) - multimodal_keys = list(get_multimodal_keys_from_processor(processor)) - # Historical checkpoints may emit dynamic image tiles without imgs_sizes. - # Mirror the media-metadata handling in vlm_hf_data_processor. - if ( - uses_placeholder - and "pixel_values" in processed - and "imgs_sizes" not in processed - and processed["pixel_values"].ndim == 4 - ): - pixel_values = processed["pixel_values"] - num_tiles, _, height, width = pixel_values.shape - processed["imgs_sizes"] = torch.tensor( - [[height, width]] * num_tiles, dtype=torch.long - ) - - # imgs_sizes / num_frames are not always declared in model_input_names by - # bundled image processors. RADIO uses temporal patching even for still - # images and requires one num_frames=1 entry per image/tile. - if "imgs_sizes" in processed and "imgs_sizes" not in multimodal_keys: - multimodal_keys.append("imgs_sizes") - if "imgs_sizes" in processed and "num_frames" not in processed: - processed["num_frames"] = torch.ones( - len(processed["imgs_sizes"]), dtype=torch.long - ) - if "num_frames" in processed and "num_frames" not in multimodal_keys: - multimodal_keys.append("num_frames") - for key in multimodal_keys: - if key not in processed: - continue - value = processed[key] - if key == "imgs_sizes": - value = value.to(dtype=torch.int32) - user_message[key] = PackedTensor( - value, - dim_to_pack=get_dim_to_pack_along(processor, key), - pad_to_max_shape=uses_placeholder and key == "pixel_values", - ) @ray.remote(max_restarts=-1, max_task_retries=-1) # pragma: no cover @@ -473,6 +474,7 @@ async def run_rollouts( nemo_gym_examples: list[dict], tokenizer: PreTrainedTokenizerBase, timer_prefix: str, + deduplicate_multimodal_data: bool = False, ) -> AsyncGenerator[tuple[int, dict, dict | None], None]: """Stream postprocessed rollouts as NeMo-Gym tasks complete.""" if not nemo_gym_examples: @@ -511,7 +513,10 @@ async def run_rollouts( with timer.time(label=f"{timer_prefix}/postprocess_results"): nemo_rl_result = self._postprocess_nemo_gym_to_nemo_rl_result( - nemo_gym_result, tokenizer + nemo_gym_row, + nemo_gym_result, + tokenizer, + include_initial_multimodal_data=not deduplicate_multimodal_data, ) if _has_nan_generation_logprobs(nemo_rl_result): raise RuntimeError("Generation logprobs contain NaN") @@ -548,20 +553,73 @@ async def run_rollouts( def _postprocess_nemo_gym_to_nemo_rl_result( self, + nemo_gym_row: dict, nemo_gym_result: dict, tokenizer: PreTrainedTokenizerBase, + *, + include_initial_multimodal_data: bool = True, ) -> dict: assert isinstance(nemo_gym_result, dict), ( f"Hit a non-successful response when querying NeMo Gym for rollouts: {nemo_gym_result}" ) processor = getattr(self, "_processor", None) + response = nemo_gym_result["response"] + result_input = nemo_gym_result["responses_create_params"].get("input", []) + request_input = nemo_gym_row.get("responses_create_params", {}).get("input") + raw_input = ( + request_input + if isinstance(request_input, list) and request_input + else result_input + ) + initial_input = response.get("agent_input") + if not isinstance(initial_input, list) or not initial_input: + initial_input = raw_input + + seed_obs = response.get("seed_obs") + media_messages = ( + seed_obs if isinstance(seed_obs, list) and seed_obs else initial_input + ) + raw_initial_sources = extract_input_image_sources_from_responses_messages( + raw_input + ) + agent_initial_sources = extract_input_image_sources_from_responses_messages( + initial_input + ) + returned_media_sources = extract_input_image_sources_from_responses_messages( + media_messages + ) + initial_media_matches_raw_input = ( + bool(raw_initial_sources) + and len(agent_initial_sources) == len(raw_initial_sources) + and all( + _image_sources_equal(agent_source, raw_source) + for agent_source, raw_source in zip( + agent_initial_sources, raw_initial_sources + ) + ) + ) + returned_media_matches_raw_input = len(returned_media_sources) == len( + raw_initial_sources + ) and all( + _image_sources_equal(returned_source, raw_source) + for returned_source, raw_source in zip( + returned_media_sources, raw_initial_sources + ) + ) + initial_multimodal_data_omitted = ( + not include_initial_multimodal_data + and initial_media_matches_raw_input + and returned_media_matches_raw_input + ) + if initial_multimodal_data_omitted: + media_messages, _ = _without_initial_image_sources( + media_messages, raw_initial_sources + ) per_turn_images = ( _index_per_turn_images( - nemo_gym_result["response"]["output"], - input_messages=nemo_gym_result.get("responses_create_params", {}).get( - "input" - ), + response["output"], + input_messages=media_messages, ) if processor is not None else [] @@ -736,11 +794,26 @@ def _postprocess_nemo_gym_to_nemo_rl_result( f" ā†’ If (2): inspect why no assistant content was produced for this rollout." ) - return { + if initial_multimodal_data_omitted: + for container, key in ( + (nemo_gym_result["responses_create_params"], "input"), + (response, "agent_input"), + (response, "seed_obs"), + ): + if key in container: + container[key], _ = _without_initial_image_sources( + container[key], raw_initial_sources + ) + nemo_gym_result["_nemo_rl_initial_media_omitted"] = True + + result = { "message_log": nemo_rl_message_log, "input_message_log": nemo_rl_message_log[:1], "full_result": nemo_gym_result, } + if not include_initial_multimodal_data: + result["_initial_multimodal_data_omitted"] = initial_multimodal_data_omitted + return result def shutdown(self) -> None: self.rh.shutdown() diff --git a/nemo_rl/experience/rollouts.py b/nemo_rl/experience/rollouts.py index 84ff46b13d3..c30ae496667 100644 --- a/nemo_rl/experience/rollouts.py +++ b/nemo_rl/experience/rollouts.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -36,11 +36,19 @@ DatumSpec, FlatMessagesType, LLMMessageLogType, + VLMMessageLogType, ) from nemo_rl.data.llm_message_utils import ( batched_message_log_to_flat_message, get_keys_from_message_log, ) +from nemo_rl.data.multimodal_utils import ( + NATIVE_MULTIMODAL_KEYS, + VLLM_MULTIMODAL_DATA_KEYS, + PackedTensor, + attach_image_model_inputs_to_message, + extract_input_images_from_responses_messages, +) from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.environments.interfaces import ( EnvironmentInterface, @@ -57,11 +65,147 @@ GenerationOutputSpec, GenerationSamplingParams, ) +from nemo_rl.utils.multimodal_payload_metrics import ( + collect_multimodal_payload_metrics, + print_multimodal_payload_metrics, +) from nemo_rl.utils.timer import Timer TokenizerType = PreTrainedTokenizerBase +def _set_untyped_message_field(message: Any, key: str, value: Any) -> None: + """Set a restored extension field outside the typed core message schema.""" + message[key] = value + + +def attach_initial_nemo_gym_image_payloads( + batch: BatchedDataDict[DatumSpec], + processor: Any, +) -> None: + """Attach initial Gym image tensors once, before prompt repeat. + + The NeMo Gym dataset deliberately carries only the Responses request in + ``extra_env_info``. Dedup-enabled GRPO calls this helper on the unrepeated + prompt batch, allowing ``repeat_interleave(..., share_immutable_media=True)`` + to retain one physical processor output per prompt. Flag-off runs never call + this helper. + """ + for message_log, extra_env_info in zip( + batch["message_log"], batch["extra_env_info"] + ): + if extra_env_info is None or not isinstance(extra_env_info, dict): + continue + initial_messages = extra_env_info.get("responses_create_params", {}).get( + "input", [] + ) + images = extract_input_images_from_responses_messages(initial_messages) + if not images: + continue + if processor is None or getattr(processor, "image_processor", None) is None: + raise ValueError( + "NeMo Gym image deduplication requires the multimodal processor " + "to be passed to GRPO." + ) + user_message = next( + (message for message in message_log if message.get("role") == "user"), + None, + ) + if user_message is None: + raise ValueError("NeMo Gym image prompt has no user message to attach to.") + if isinstance(user_message.get("pixel_values"), PackedTensor): + continue + attach_image_model_inputs_to_message( + user_message, + images=images, + processor=processor, + ) + + +def _add_multimodal_generation_payload( + generation_input_data: BatchedDataDict[GenerationDatumSpec], + flat_messages: BatchedDataDict[FlatMessagesType], + active_batch: BatchedDataDict[DatumSpec], + policy_generation: GenerationInterface, + *, + deduplicate_multimodal_data: bool, +) -> None: + """Attach one policy-ready or native-vLLM media representation. + + The compact policy representation remains in ``message_log`` for later + logprob/training construction. When every active row has a native vLLM + prompt, sending that representation as well is redundant. + """ + generation_config = getattr(policy_generation, "cfg", {}) + native_content = active_batch.get("vllm_content") + + def row_has_formatter_consumed_media(row_index: int) -> bool: + for key in VLLM_MULTIMODAL_DATA_KEYS: + rows = active_batch.get(key) + if rows is None or row_index >= len(rows): + continue + value = rows[row_index] + if value is None: + continue + if isinstance(value, (list, tuple, dict, str, bytes)): + if len(value) > 0: + return True + else: + return True + return False + + use_native_vllm_only = ( + deduplicate_multimodal_data + and generation_config.get("backend") == "vllm" + and native_content is not None + and all( + row_has_formatter_consumed_media(row_index) + for row_index in range(len(native_content)) + ) + ) + if not use_native_vllm_only: + generation_input_data.update( + flat_messages.get_multimodal_dict(as_tensors=False) + ) + + for key in NATIVE_MULTIMODAL_KEYS: + if key in active_batch: + generation_input_data[key] = active_batch[key] + + +def _reattach_original_multimodal_payloads( + results: list[dict[str, Any]], + original_message_logs: list[LLMMessageLogType | VLMMessageLogType], +) -> None: + """Restore exact prompt media omitted by a remote Gym rollout. + + User turns are matched by their ordinal position. Only explicit + ``PackedTensor`` values and named native-generation media are restored, so + arbitrary non-text metadata is never misclassified as media. Newly returned + Gym media is left untouched unless it occupies the corresponding original + prompt key. + """ + for result, original_log in zip(results, original_message_logs): + if not result.pop("_initial_multimodal_data_omitted", False): + continue + original_user_messages = [ + message for message in original_log if message.get("role") == "user" + ] + for log_key in ("input_message_log", "message_log"): + target_log = result.get(log_key) + if not target_log: + continue + target_user_messages = [ + message for message in target_log if message.get("role") == "user" + ] + for original, target in zip(original_user_messages, target_user_messages): + for key, value in original.items(): + if isinstance(value, PackedTensor): + _set_untyped_message_field(target, key, value) + elif key in NATIVE_MULTIMODAL_KEYS: + _set_untyped_message_field(target, key, value) + + def _add_r3_fallback_metrics( gen_metrics: dict[str, float | int], generation_outputs: BatchedDataDict, @@ -671,6 +815,7 @@ def run_multi_turn_rollout( max_seq_len: int, max_rollout_turns: int = 999999, greedy: bool = False, + deduplicate_multimodal_data: bool = False, ) -> tuple[BatchedDataDict[DatumSpec], dict[str, Any]]: """Runs a multi-turn rollout loop, interacting with the environment. @@ -682,6 +827,9 @@ def run_multi_turn_rollout( max_rollout_turns: Maximum number of agent-environment interaction turns. max_seq_len: Maximum sequence length allowed. greedy: Whether to use greedy decoding. + deduplicate_multimodal_data: Send only native media through the vLLM + generation boundary while retaining compact policy media for + logprob and training. Returns: Tuple containing: @@ -720,6 +868,8 @@ def run_multi_turn_rollout( # Convert LLMMessageLogType to FlatMessagesType for generation active_batch = current_batch.select_indices(active_indices) + if turn > 0 and "vllm_content" in active_batch: + active_batch["vllm_content"] = [None] * len(active_indices) active_stop_strings = [current_stop_strings[i] for i in active_indices.tolist()] active_flat_messages: BatchedDataDict[FlatMessagesType] @@ -741,19 +891,13 @@ def run_multi_turn_rollout( "stop_strings": active_stop_strings, } ) - # add the multimodal data to the generation input data - multimodal_data = active_flat_messages.get_multimodal_dict(as_tensors=False) - generation_input_data.update(multimodal_data) - - # keep message log for generation - if "vllm_content" in active_batch: - generation_input_data["vllm_content"] = active_batch["vllm_content"] - if "vllm_images" in active_batch: - generation_input_data["vllm_images"] = active_batch["vllm_images"] - if "vllm_videos" in active_batch: - generation_input_data["vllm_videos"] = active_batch["vllm_videos"] - if "vllm_audios" in active_batch: - generation_input_data["vllm_audios"] = active_batch["vllm_audios"] + _add_multimodal_generation_payload( + generation_input_data, + active_flat_messages, + active_batch, + policy_generation, + deduplicate_multimodal_data=deduplicate_multimodal_data, + ) # generate_responses updates active_batch["message_log"] in-place active_batch, generated_ids, gen_metrics = generate_responses( @@ -920,6 +1064,9 @@ async def async_generate_response_for_sample_turn( tokenizer: TokenizerType, max_seq_len: int, greedy: bool = False, + *, + sample_multimodal_data: dict[str, Any] | None = None, + deduplicate_multimodal_data: bool = False, ) -> tuple[list[dict], torch.Tensor, torch.Tensor, dict[str, float]]: """Generate a response for a single sample's turn using async generation. @@ -930,6 +1077,9 @@ async def async_generate_response_for_sample_turn( tokenizer: Tokenizer to use max_seq_len: Maximum sequence length greedy: Whether to use greedy decoding + sample_multimodal_data: Native vLLM media fields for this sample. + deduplicate_multimodal_data: Avoid sending both native and policy-ready + media through the async generation boundary. Returns: Tuple of (updated_message_log, generated_tokens, input_lengths, generation_metrics) @@ -961,6 +1111,15 @@ async def async_generate_response_for_sample_turn( "stop_strings": [sample_stop_strings], } ) + for key, value in (sample_multimodal_data or {}).items(): + dummy_batch[key] = [value] + _add_multimodal_generation_payload( + generation_input_data, + flat_messages, + dummy_batch, + policy_generation, + deduplicate_multimodal_data=deduplicate_multimodal_data, + ) # Generate response using the async version updated_batch, generated_ids, gen_metrics = await generate_responses_async( @@ -989,6 +1148,7 @@ async def run_sample_multi_turn_rollout( max_seq_len: int, max_rollout_turns: int = 999999, greedy: bool = False, + deduplicate_multimodal_data: bool = False, ) -> tuple[dict, dict[str, Any]]: """Run a multi-turn rollout for a single sample. @@ -1004,6 +1164,8 @@ async def run_sample_multi_turn_rollout( max_seq_len: Maximum sequence length max_rollout_turns: Maximum number of turns greedy: Whether to use greedy decoding + deduplicate_multimodal_data: Avoid redundant media at generation + boundaries while preserving compact policy media in the trajectory. Returns: Tuple of (final_sample_state, sample_metrics) @@ -1013,6 +1175,11 @@ async def run_sample_multi_turn_rollout( current_extra_env_info = copy.deepcopy(initial_sample_state["extra_env_info"]) current_stop_strings = initial_sample_state.get("stop_strings", None) task_name = initial_sample_state["task_name"] + sample_multimodal_data = { + key: initial_sample_state[key] + for key in NATIVE_MULTIMODAL_KEYS + if key in initial_sample_state + } # Sample-level metrics total_reward = 0.0 @@ -1041,6 +1208,11 @@ async def run_sample_multi_turn_rollout( # Generate response for this sample using async generation try: + turn_multimodal_data = sample_multimodal_data + if turn > 0 and "vllm_content" in sample_multimodal_data: + turn_multimodal_data = dict(sample_multimodal_data) + turn_multimodal_data["vllm_content"] = None + ( updated_message_log, generated_tokens, @@ -1053,6 +1225,8 @@ async def run_sample_multi_turn_rollout( tokenizer, max_seq_len, greedy=greedy, + sample_multimodal_data=turn_multimodal_data, + deduplicate_multimodal_data=deduplicate_multimodal_data, ) current_message_log = updated_message_log @@ -1289,21 +1463,24 @@ async def _run_multi_turn_rollout_async( max_seq_len: int, max_rollout_turns: int = 999999, greedy: bool = False, + deduplicate_multimodal_data: bool = False, ) -> tuple[BatchedDataDict[DatumSpec], list[dict[str, Any]]]: """Run one native rollout batch and retain metrics at sample granularity.""" batch_size = len(input_batch["message_log"]) sample_initial_states = [] for i in range(batch_size): - sample_initial_states.append( - { - "message_log": input_batch["message_log"][i], - "extra_env_info": input_batch["extra_env_info"][i], - "task_name": input_batch["task_name"][i], - "stop_strings": input_batch.get("stop_strings", [None] * batch_size)[i], - "idx": input_batch.get("idx", list(range(batch_size)))[i], - } - ) + sample_state = { + "message_log": input_batch["message_log"][i], + "extra_env_info": input_batch["extra_env_info"][i], + "task_name": input_batch["task_name"][i], + "stop_strings": input_batch.get("stop_strings", [None] * batch_size)[i], + "idx": input_batch.get("idx", list(range(batch_size)))[i], + } + for key in NATIVE_MULTIMODAL_KEYS: + if key in input_batch: + sample_state[key] = input_batch[key][i] + sample_initial_states.append(sample_state) async def run_single_sample_with_error_handling(i, sample_state): try: @@ -1316,6 +1493,7 @@ async def run_single_sample_with_error_handling(i, sample_state): max_seq_len=max_seq_len, max_rollout_turns=max_rollout_turns, greedy=greedy, + deduplicate_multimodal_data=deduplicate_multimodal_data, ) except Exception as error: raise RuntimeError(f"Error in sample {i} rollout: {error}") from error @@ -1382,6 +1560,7 @@ def run_async_multi_turn_rollout( max_seq_len: int, max_rollout_turns: int = 999999, greedy: bool = False, + deduplicate_multimodal_data: bool = False, ) -> tuple[BatchedDataDict[DatumSpec], dict[str, Any]]: """Run a complete native rollout batch from a synchronous call site. @@ -1414,6 +1593,7 @@ def run_async_multi_turn_rollout( max_seq_len=max_seq_len, max_rollout_turns=max_rollout_turns, greedy=greedy, + deduplicate_multimodal_data=deduplicate_multimodal_data, ) ) return final_batch, _aggregate_multi_turn_rollout_metrics(sample_metrics) @@ -1428,6 +1608,7 @@ async def run_async_multi_turn_rollout_groups( num_generations: int, max_rollout_turns: int = 999999, greedy: bool = False, + deduplicate_multimodal_data: bool = False, ) -> AsyncGenerator[RolloutGroupResult, None]: """Run one native batch, then yield prompt groups with group-local metrics. @@ -1470,6 +1651,7 @@ async def run_async_multi_turn_rollout_groups( max_seq_len=max_seq_len, max_rollout_turns=max_rollout_turns, greedy=greedy, + deduplicate_multimodal_data=deduplicate_multimodal_data, ) for group_index, start in enumerate(range(0, final_batch.size, num_generations)): end = start + num_generations @@ -2063,6 +2245,8 @@ async def run_async_nemo_gym_rollout( mask_env_flagged_samples: bool = True, returns_entire_batch: bool = False, sampling_params: Optional[GenerationSamplingParams] = None, + deduplicate_multimodal_data: bool = False, + debug_payload_metrics: bool = False, ) -> AsyncGenerator[NemoGymRolloutResult, None]: """Stream complete NeMo-Gym prompt groups in group-completion order. @@ -2096,6 +2280,10 @@ async def run_async_nemo_gym_rollout( sampling_params: Sampling profile stamped onto every NeMo-Gym row. ``None`` uses the train profile from ``generation_config``; validation passes its own profile explicitly. + deduplicate_multimodal_data: Omit initial policy-ready media from the + remote Gym return and restore the exact original payload locally. + debug_payload_metrics: Emit logical, physical, and serialized media + payload metrics at the Gym Ray boundary. Yields: ``NemoGymRolloutResult`` objects in prompt-group completion order. Rows @@ -2183,9 +2371,22 @@ async def run_async_nemo_gym_rollout( actor_timing_metrics: dict[str, Any] = {} nemo_gym_environment = task_to_env["nemo_gym"] with timer.time(run_rollouts_timer_label): + ray_arguments = ( + nemo_gym_rows, + tokenizer, + timer_prefix, + deduplicate_multimodal_data, + ) + print_multimodal_payload_metrics( + collect_multimodal_payload_metrics( + ray_arguments, + "nemo_gym_request", + enabled=debug_payload_metrics, + ) + ) rollout_gen = nemo_gym_environment.run_rollouts.options( num_returns="streaming" - ).remote(nemo_gym_rows, tokenizer, timer_prefix) + ).remote(*ray_arguments) rollout_iterator = rollout_gen.__aiter__() while True: @@ -2199,6 +2400,17 @@ async def run_async_nemo_gym_rollout( stream_finished = True else: rowidx, result, timing_metrics = await future + # Measure the received streaming Ray value in the caller. In + # async training this runs in the collector actor; validation + # runs in the driver, so the two phases cannot share a metric + # accumulator even when they share the NeMo-Gym actor. + print_multimodal_payload_metrics( + collect_multimodal_payload_metrics( + (rowidx, result, timing_metrics), + "nemo_gym_return", + enabled=debug_payload_metrics, + ) + ) if not stream_finished: if timing_metrics is not None: @@ -2207,16 +2419,22 @@ async def run_async_nemo_gym_rollout( _tensorize_nemo_gym_result(result) completed_group = accumulator.add(rowidx, result) if completed_group is not None: + group_input_batch = input_batch.slice( + completed_group.group_index * num_generations, + (completed_group.group_index + 1) * num_generations, + ) + if deduplicate_multimodal_data: + _reattach_original_multimodal_payloads( + completed_group.results, + group_input_batch["message_log"], + ) rollout_result = _postprocess_single_nemo_gym_group( nemo_gym_rows=completed_group.rows, results=completed_group.results, timer=timer, timer_prefix=timer_prefix, policy_generation=policy_generation, - input_batch=input_batch.slice( - completed_group.group_index * num_generations, - (completed_group.group_index + 1) * num_generations, - ), + input_batch=group_input_batch, tokenizer=tokenizer, log_full_result_tables=log_full_result_tables, effort_config=effort_config, @@ -2261,6 +2479,8 @@ def run_nemo_gym_rollout_sync( thinking_tags: list[str] | tuple[str, ...] | None = None, sampling_params: Optional[GenerationSamplingParams] = None, mask_env_flagged_samples: bool = True, + deduplicate_multimodal_data: bool = False, + debug_payload_metrics: bool = False, ) -> NemoGymRolloutResult: """Run and return one complete NeMo-Gym batch synchronously. @@ -2288,6 +2508,9 @@ def run_nemo_gym_rollout_sync( validation passes its own profile explicitly. mask_env_flagged_samples: Whether to carry env-driven ``mask_sample`` flags in the rollout batch for loss masking. + deduplicate_multimodal_data: Omit initial policy-ready media from the + remote Gym return and restore it from the input batch. + debug_payload_metrics: Emit exact Gym Ray-boundary media payload metrics. Returns: The fully postprocessed NeMo-Gym rollout batch in input-row order. @@ -2320,6 +2543,8 @@ async def _consume_rollout() -> NemoGymRolloutResult: mask_env_flagged_samples=mask_env_flagged_samples, returns_entire_batch=True, sampling_params=sampling_params, + deduplicate_multimodal_data=deduplicate_multimodal_data, + debug_payload_metrics=debug_payload_metrics, ): pass if rollout_result is None: diff --git a/nemo_rl/experience/sync_rollout_actor.py b/nemo_rl/experience/sync_rollout_actor.py index a29ceb9696a..1bc5c6c391e 100644 --- a/nemo_rl/experience/sync_rollout_actor.py +++ b/nemo_rl/experience/sync_rollout_actor.py @@ -263,6 +263,8 @@ def rollout_to_tq( else None, reward_penalty_config=cfg.reward_penalties, thinking_tags=get_nemo_gym_thinking_tags(cfg.env), + deduplicate_multimodal_data=cfg.grpo.deduplicate_multimodal_data, + debug_payload_metrics=cfg.grpo.debug_payload_metrics, ) final_batch, rollout_metrics = r.final_batch, r.rollout_metrics else: @@ -275,6 +277,7 @@ def rollout_to_tq( **common, max_seq_len=cfg.policy["max_total_sequence_length"], max_rollout_turns=cfg.grpo.max_rollout_turns, + deduplicate_multimodal_data=cfg.grpo.deduplicate_multimodal_data, ) fb = final_batch.to("cpu") del final_batch diff --git a/nemo_rl/models/generation/interfaces.py b/nemo_rl/models/generation/interfaces.py index 791657c394f..8a84f0ddd5a 100644 --- a/nemo_rl/models/generation/interfaces.py +++ b/nemo_rl/models/generation/interfaces.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -221,6 +221,8 @@ class GenerationConfig(TypedDict): _pad_token_id: NotRequired[int] # MTP draft weights arrive via refit if the trainer trains the MTP layer. _mtp_weights_from_refit: NotRequired[bool] + # Internal debug-only measurement of exact Ray generation arguments. + debug_payload_metrics: NotRequired[bool] @dataclass diff --git a/nemo_rl/models/generation/vllm/utils.py b/nemo_rl/models/generation/vllm/utils.py index b8ec3dee050..e4f9a78ab43 100644 --- a/nemo_rl/models/generation/vllm/utils.py +++ b/nemo_rl/models/generation/vllm/utils.py @@ -91,39 +91,38 @@ def _get_regular_prompt(index: int): token_ids = valid_ids.tolist() return {"prompt_token_ids": token_ids} - # Check if this is VLM generation by looking for message_log with images - # Support for videos/audio/etc. can be added here - # if 'message_log' in data and any('images' in msg for msg in data['message_log']): + def _get_multi_modal_data(index: int) -> dict[str, Any]: + multi_modal_data = {} + images = data.get("vllm_images", None) + if images is not None and len(images[index]) > 0: + multi_modal_data["image"] = ( + images[index][0] if len(images[index]) == 1 else images[index] + ) + audios = data.get("vllm_audios", None) + if audios is not None and len(audios[index]) > 0: + multi_modal_data["audio"] = ( + audios[index][0] if len(audios[index]) == 1 else audios[index] + ) + videos = data.get("vllm_videos", None) + if videos is not None and len(videos[index]) > 0: + multi_modal_data["video"] = ( + videos[index][0] if len(videos[index]) == 1 else videos[index] + ) + return multi_modal_data + + # Native image, audio, and video side channels share this formatter path. if "vllm_content" in data: # VLM generation using content and multi_modal_data for i in range(start_idx, end_idx): msg = data["vllm_content"][i] - # if msg is None, this conversation had no multimodal content, fallback to regular prompt - if msg is None: - prompts.append(_get_regular_prompt(i)) - continue - # init prompt dict - prompt_dict = {"prompt": msg} - # collect multi_modal_data from images, audios, and videos - multi_modal_data = {} - images = data.get("vllm_images", None) - if images is not None and len(images[i]) > 0: - multi_modal_data["image"] = ( - images[i][0] if len(images[i]) == 1 else images[i] - ) - audios = data.get("vllm_audios", None) - if audios is not None and len(audios[i]) > 0: - multi_modal_data["audio"] = ( - audios[i][0] if len(audios[i]) == 1 else audios[i] - ) - videos = data.get("vllm_videos", None) - if videos is not None and len(videos[i]) > 0: - multi_modal_data["video"] = ( - videos[i][0] if len(videos[i]) == 1 else videos[i] - ) + multi_modal_data = _get_multi_modal_data(i) if not multi_modal_data: prompts.append(_get_regular_prompt(i)) continue + # Raw processor content is valid only for the initial turn. Later + # turns use the updated pre-tokenized conversation plus the same + # native media, preventing vLLM from regenerating the stale prompt. + prompt_dict = {"prompt": msg} if msg is not None else _get_regular_prompt(i) prompt_dict["multi_modal_data"] = multi_modal_data prompts.append(prompt_dict) else: diff --git a/nemo_rl/models/generation/vllm/vllm_generation.py b/nemo_rl/models/generation/vllm/vllm_generation.py index 6d0d8d37143..da0e10de45e 100644 --- a/nemo_rl/models/generation/vllm/vllm_generation.py +++ b/nemo_rl/models/generation/vllm/vllm_generation.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -43,6 +43,11 @@ compute_spec_decode_metrics, resolve_generation_worker_cls, ) +from nemo_rl.utils.multimodal_payload_metrics import ( + collect_multimodal_payload_metrics, + collect_sharded_multimodal_payload_metrics, + print_multimodal_payload_metrics, +) from nemo_rl.weight_sync.interfaces import WeightSynchronizer logger = logging.getLogger(__name__) @@ -638,6 +643,13 @@ def generate( sharded_data: list[SlicedDataDict] = data.shard_by_batch_size( dp_size, allow_uneven_shards=True ) + print_multimodal_payload_metrics( + collect_sharded_multimodal_payload_metrics( + sharded_data, + "vllm_generation", + enabled=self.cfg.get("debug_payload_metrics", False), + ) + ) future_bundle = self.worker_group.run_all_workers_sharded_data( "generate", data=sharded_data, @@ -689,6 +701,13 @@ def generate_text( sharded_data: list[SlicedDataDict] = data.shard_by_batch_size( dp_size, allow_uneven_shards=True ) + print_multimodal_payload_metrics( + collect_sharded_multimodal_payload_metrics( + sharded_data, + "vllm_text_generation", + enabled=self.cfg.get("debug_payload_metrics", False), + ) + ) future_bundle = self.worker_group.run_all_workers_sharded_data( "generate_text", data=sharded_data, @@ -758,6 +777,13 @@ async def _async_generate_base( leader_worker_idx = self.worker_group.get_dp_leader_worker_idx( self.current_generate_dp_shard_idx ) + print_multimodal_payload_metrics( + collect_multimodal_payload_metrics( + data, + "vllm_generation_async", + enabled=self.cfg.get("debug_payload_metrics", False), + ) + ) # Run the async method on the selected leader worker worker_gen_proxy = self.worker_group.run_single_worker_single_data( diff --git a/nemo_rl/models/policy/lm_policy.py b/nemo_rl/models/policy/lm_policy.py index 6a184ded3e9..ef207318437 100644 --- a/nemo_rl/models/policy/lm_policy.py +++ b/nemo_rl/models/policy/lm_policy.py @@ -56,6 +56,10 @@ get_default_hf_config, get_theoretical_tflops, ) +from nemo_rl.utils.multimodal_payload_metrics import ( + collect_sharded_multimodal_payload_metrics, + print_multimodal_payload_metrics, +) from nemo_rl.utils.timer import Timer PathLike = Union[str, "os.PathLike[Any]"] @@ -98,7 +102,9 @@ def __init__( processor: Optional[AutoProcessor] = None, worker_extension_cls_fqn: Optional[str] = None, skip_weight_load: bool = False, + debug_payload_metrics: bool = False, ): + self.debug_payload_metrics = debug_payload_metrics if weights_path: weights_path = os.path.abspath(weights_path) if optimizer_path: @@ -523,6 +529,22 @@ def _shard_for_train( ) return sharded_data + def _report_sharded_payload( + self, + sharded_data: list["SlicedDataDict"], + boundary: str, + ) -> None: + """Measure the exact unique per-DP-shard Ray arguments.""" + if not self.debug_payload_metrics: + return + print_multimodal_payload_metrics( + collect_sharded_multimodal_payload_metrics( + sharded_data, + boundary, + enabled=True, + ) + ) + def get_logprobs( self, data: BatchedDataDict[GenerationDatumSpec], @@ -537,6 +559,7 @@ def get_logprobs( """ with timer.time("get_logprobs/shard_data") if timer else nullcontext(): sharded_data, unsorted_data_indices = self._shard_for_logprob(data) + self._report_sharded_payload(sharded_data, "policy_get_logprobs") with ( timer.time("get_logprobs/submit_logprob_futures") @@ -585,6 +608,7 @@ def get_reference_policy_logprobs( else nullcontext() ): sharded_data, unsorted_data_indices = self._shard_for_logprob(data) + self._report_sharded_payload(sharded_data, "policy_get_reference_logprobs") with ( timer.time( @@ -751,6 +775,7 @@ def train( # Shard and replicate the batch with timer.time("policy_training/sharding_data") if timer else nullcontext(): sharded_data = self._shard_for_train(data, batch_size) + self._report_sharded_payload(sharded_data, "policy_train") if self.flops_tracker is not None: self.flops_tracker.reset() @@ -995,6 +1020,7 @@ def calibrate_qkv_fp8_scales( dp_size, batch_size=None, ) + self._report_sharded_payload(sharded_data, "policy_kv_calibration") futures = self.worker_group.run_all_workers_sharded_data( "calibrate_qkv_fp8_scales", diff --git a/nemo_rl/utils/checkpoint.py b/nemo_rl/utils/checkpoint.py index 8c8de3129d9..4fbbd0ae230 100644 --- a/nemo_rl/utils/checkpoint.py +++ b/nemo_rl/utils/checkpoint.py @@ -126,6 +126,9 @@ class CheckpointingConfig(TypedDict): keep_top_k: NotRequired[int] ft_keep_latest_k: NotRequired[int | None] ft_save_period: NotRequired[int] + # Async GRPO only. Disable to regenerate replay trajectories after resume + # instead of serializing a potentially very large buffer. + save_replay_buffer: NotRequired[bool] checkpoint_must_save_by: NotRequired[str | None] pretrained_checkpoint: NotRequired[PretrainedCheckpointConfig] save_optimizer: NotRequired[bool] # Default: True diff --git a/nemo_rl/utils/multimodal_payload_metrics.py b/nemo_rl/utils/multimodal_payload_metrics.py new file mode 100644 index 00000000000..49c5b82d779 --- /dev/null +++ b/nemo_rl/utils/multimodal_payload_metrics.py @@ -0,0 +1,403 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from collections.abc import Mapping, Sequence +from threading import Lock +from typing import Any + +import numpy as np +import ray.cloudpickle as cloudpickle +import torch +from PIL import Image + +from nemo_rl.data.multimodal_utils import ( + MULTIMODAL_CONTENT_TYPES, + NATIVE_MULTIMODAL_KEYS, + PackedTensor, +) + + +_PENDING_PAYLOAD_METRICS: dict[str, int | float] = {} +_PENDING_PAYLOAD_METRICS_LOCK = Lock() + + +def _tensor_nbytes(value: torch.Tensor | None) -> int: + if value is None: + return 0 + return value.numel() * value.element_size() + + +def _value_nbytes(value: Any, seen: set[int] | None = None) -> int: + """Estimate data bytes, optionally counting shared leaves only once.""" + if value is None: + return 0 + if isinstance(value, (str, bytes, bytearray, memoryview)): + object_id = id(value) + if seen is not None: + if object_id in seen: + return 0 + seen.add(object_id) + if isinstance(value, str): + return len(value.encode("utf-8")) + return len(value) + if torch.is_tensor(value): + object_id = id(value) + if seen is not None: + if object_id in seen: + return 0 + seen.add(object_id) + return _tensor_nbytes(value) + if isinstance(value, np.ndarray): + object_id = id(value) + if seen is not None: + if object_id in seen: + return 0 + seen.add(object_id) + return value.nbytes + if isinstance(value, Image.Image): + object_id = id(value) + if seen is not None: + if object_id in seen: + return 0 + seen.add(object_id) + return len(value.tobytes()) + if isinstance(value, Mapping): + return sum(_value_nbytes(item, seen) for item in value.values()) + if isinstance(value, (list, tuple)): + return sum(_value_nbytes(item, seen) for item in value) + return 0 + + +def _value_segment_count(value: Any, seen: set[int] | None = None) -> int: + """Count media leaves, optionally counting shared objects only once.""" + if value is None: + return 0 + if isinstance( + value, + (str, bytes, bytearray, memoryview, torch.Tensor, np.ndarray, Image.Image), + ): + object_id = id(value) + if seen is not None: + if object_id in seen: + return 0 + seen.add(object_id) + return 1 + if isinstance(value, Mapping): + return sum(_value_segment_count(item, seen) for item in value.values()) + if isinstance(value, (list, tuple)): + return sum(_value_segment_count(item, seen) for item in value) + return 0 + + +def _typed_content_media_nbytes( + value: Any, + seen: set[int] | None = None, +) -> int: + """Count media embedded in typed vLLM content without counting prompt text.""" + if isinstance(value, Mapping): + if value.get("type") in MULTIMODAL_CONTENT_TYPES: + return sum( + _value_nbytes(item, seen) + for key, item in value.items() + if key != "type" + ) + return sum(_typed_content_media_nbytes(item, seen) for item in value.values()) + if isinstance(value, (list, tuple)): + return sum(_typed_content_media_nbytes(item, seen) for item in value) + return 0 + + +def _typed_content_media_segment_count( + value: Any, + seen: set[int] | None = None, +) -> int: + """Count media leaves embedded in typed content without counting text.""" + if isinstance(value, Mapping): + if value.get("type") in MULTIMODAL_CONTENT_TYPES: + return sum( + _value_segment_count(item, seen) + for key, item in value.items() + if key != "type" + ) + return sum( + _typed_content_media_segment_count(item, seen) for item in value.values() + ) + if isinstance(value, (list, tuple)): + return sum(_typed_content_media_segment_count(item, seen) for item in value) + return 0 + + +def protocol5_serialized_nbytes(value: Any) -> int: + """Return cloudpickle protocol-5 frame plus out-of-band buffer bytes.""" + buffers = [] + frame = cloudpickle.dumps(value, protocol=5, buffer_callback=buffers.append) + buffer_bytes = 0 + for buffer in buffers: + raw = buffer.raw() if hasattr(buffer, "raw") else memoryview(buffer) + buffer_bytes += raw.nbytes + return len(frame) + buffer_bytes + + +def collect_multimodal_payload_metrics( + data: Any, + boundary: str, + *, + enabled: bool, +) -> dict[str, int | float]: + """Measure one exact Ray argument without scanning when disabled.""" + if not enabled: + return {} + + totals = { + "physical_media_bytes": 0, + "logical_media_bytes": 0, + "physical_segments": 0, + "logical_segments": 0, + } + per_key: dict[str, int] = {} + seen_native_leaves: set[int] = set() + seen_native_segments: set[int] = set() + seen_packed_leaves: set[int] = set() + seen_packed_leaves_by_key: dict[str, set[int]] = {} + + def visit(key: str, value: Any) -> None: + if isinstance(value, PackedTensor): + key_seen = seen_packed_leaves_by_key.setdefault(key, set()) + key_physical_segments = 0 + key_physical = 0 + for item in value.tensors: + if item is None: + continue + object_id = id(item) + if object_id not in key_seen: + key_seen.add(object_id) + key_physical_segments += 1 + if object_id not in seen_packed_leaves: + seen_packed_leaves.add(object_id) + key_physical += _tensor_nbytes(item) + + logical_items = [ + item for item in value.iter_logical_segments() if item is not None + ] + key_logical = sum(_tensor_nbytes(item) for item in logical_items) + totals["physical_media_bytes"] += key_physical + totals["logical_media_bytes"] += key_logical + totals["physical_segments"] += key_physical_segments + totals["logical_segments"] += len(logical_items) + physical_key = f"payload_counts/{boundary}/{key}/physical_segments" + logical_key = f"payload_counts/{boundary}/{key}/logical_segments" + per_key[physical_key] = per_key.get(physical_key, 0) + key_physical_segments + per_key[logical_key] = per_key.get(logical_key, 0) + len(logical_items) + elif key == "vllm_content": + totals["physical_media_bytes"] += _typed_content_media_nbytes( + value, seen_native_leaves + ) + totals["logical_media_bytes"] += _typed_content_media_nbytes(value) + totals["physical_segments"] += _typed_content_media_segment_count( + value, seen_native_segments + ) + totals["logical_segments"] += _typed_content_media_segment_count(value) + elif key in NATIVE_MULTIMODAL_KEYS: + totals["physical_media_bytes"] += _value_nbytes(value, seen_native_leaves) + totals["logical_media_bytes"] += _value_nbytes(value) + totals["physical_segments"] += _value_segment_count( + value, seen_native_segments + ) + totals["logical_segments"] += _value_segment_count(value) + elif ( + isinstance(value, Mapping) and value.get("type") in MULTIMODAL_CONTENT_TYPES + ): + totals["physical_media_bytes"] += _typed_content_media_nbytes( + value, seen_native_leaves + ) + totals["logical_media_bytes"] += _typed_content_media_nbytes(value) + totals["physical_segments"] += _typed_content_media_segment_count( + value, seen_native_segments + ) + totals["logical_segments"] += _typed_content_media_segment_count(value) + elif isinstance(value, Mapping): + for nested_key, nested_value in value.items(): + visit(str(nested_key), nested_value) + elif isinstance(value, (list, tuple)): + for nested_value in value: + visit(key, nested_value) + + if isinstance(data, Mapping): + for key, value in data.items(): + visit(str(key), value) + else: + visit("root", data) + + physical_media_bytes = totals["physical_media_bytes"] + logical_media_bytes = totals["logical_media_bytes"] + saved_bytes = max(logical_media_bytes - physical_media_bytes, 0) + ratio = ( + float(physical_media_bytes) / float(logical_media_bytes) + if logical_media_bytes + else 1.0 + ) + return { + f"payload_bytes/{boundary}/serialized": protocol5_serialized_nbytes(data), + f"payload_bytes/{boundary}/physical_media": physical_media_bytes, + f"payload_bytes/{boundary}/logical_media": logical_media_bytes, + f"payload_bytes/{boundary}/estimated_saved": saved_bytes, + f"payload_counts/{boundary}/physical_segments": totals["physical_segments"], + f"payload_counts/{boundary}/logical_segments": totals["logical_segments"], + f"payload_counts/{boundary}/calls": 1, + f"payload_ratio/{boundary}/physical_to_logical": ratio, + **per_key, + } + + +def collect_sharded_multimodal_payload_metrics( + shards: Sequence[Mapping[str, Any]], + boundary: str, + *, + enabled: bool, +) -> dict[str, int | float]: + """Aggregate metrics over the exact unique per-DP-shard Ray arguments.""" + if not enabled: + return {} + + per_shard = [ + collect_multimodal_payload_metrics( + shard, + f"{boundary}/shard_{index}", + enabled=True, + ) + for index, shard in enumerate(shards) + ] + serialized = [ + int(metrics[f"payload_bytes/{boundary}/shard_{index}/serialized"]) + for index, metrics in enumerate(per_shard) + ] + physical = [ + int(metrics[f"payload_bytes/{boundary}/shard_{index}/physical_media"]) + for index, metrics in enumerate(per_shard) + ] + logical = [ + int(metrics[f"payload_bytes/{boundary}/shard_{index}/logical_media"]) + for index, metrics in enumerate(per_shard) + ] + physical_segments = [ + int(metrics[f"payload_counts/{boundary}/shard_{index}/physical_segments"]) + for index, metrics in enumerate(per_shard) + ] + logical_segments = [ + int(metrics[f"payload_counts/{boundary}/shard_{index}/logical_segments"]) + for index, metrics in enumerate(per_shard) + ] + total_logical_bytes = sum(logical) + return { + f"payload_bytes/{boundary}/serialized_total": sum(serialized), + f"payload_bytes/{boundary}/serialized_max_shard": max(serialized, default=0), + f"payload_bytes/{boundary}/physical_media_total": sum(physical), + f"payload_bytes/{boundary}/logical_media_total": total_logical_bytes, + f"payload_bytes/{boundary}/estimated_saved_total": max( + total_logical_bytes - sum(physical), 0 + ), + f"payload_counts/{boundary}/physical_segments_total": sum(physical_segments), + f"payload_counts/{boundary}/logical_segments_total": sum(logical_segments), + f"payload_counts/{boundary}/shards": len(shards), + f"payload_counts/{boundary}/calls": len(shards), + f"payload_ratio/{boundary}/physical_to_logical": ( + float(sum(physical)) / float(total_logical_bytes) + if total_logical_bytes + else 1.0 + ), + } + + +def merge_multimodal_payload_metrics( + metric_sets: Sequence[Mapping[str, int | float]], +) -> dict[str, int | float]: + """Aggregate payload measurements collected during one logging interval. + + Byte and segment counts are summed because repeated calls represent distinct + Ray transfers. Per-call maxima and shard counts retain their maximum value, + and physical-to-logical ratios are recomputed from the aggregated byte + totals instead of averaging ratios. + """ + merged: dict[str, int | float] = {} + ratio_keys: set[str] = set() + boundaries: set[str] = set() + for metrics in metric_sets: + for key, value in metrics.items(): + if key.startswith("payload_ratio/") and key.endswith( + "/physical_to_logical" + ): + ratio_keys.add(key) + continue + if key.startswith("payload_bytes/") and key.endswith( + "/serialized_mean_per_call" + ): + continue + if key.startswith("payload_counts/") and key.endswith("/calls"): + boundaries.add( + key.removeprefix("payload_counts/").removesuffix("/calls") + ) + if key.endswith("/serialized_max_shard") or key.endswith("/shards"): + merged[key] = max(merged.get(key, 0), value) + else: + merged[key] = merged.get(key, 0) + value + + for ratio_key in ratio_keys: + boundary = ratio_key.removeprefix("payload_ratio/").removesuffix( + "/physical_to_logical" + ) + physical_key = f"payload_bytes/{boundary}/physical_media" + logical_key = f"payload_bytes/{boundary}/logical_media" + if physical_key not in merged and logical_key not in merged: + physical_key += "_total" + logical_key += "_total" + physical = merged.get(physical_key, 0) + logical = merged.get(logical_key, 0) + merged[ratio_key] = float(physical) / float(logical) if logical else 1.0 + + for boundary in boundaries: + calls = merged[f"payload_counts/{boundary}/calls"] + serialized_key = f"payload_bytes/{boundary}/serialized" + if serialized_key not in merged: + serialized_key += "_total" + if serialized_key in merged: + merged[f"payload_bytes/{boundary}/serialized_mean_per_call"] = ( + float(merged[serialized_key]) / float(calls) if calls else 0.0 + ) + + return merged + + +def drain_multimodal_payload_metrics() -> dict[str, int | float]: + """Return and clear payload measurements recorded in this process.""" + with _PENDING_PAYLOAD_METRICS_LOCK: + pending = dict(_PENDING_PAYLOAD_METRICS) + _PENDING_PAYLOAD_METRICS.clear() + return merge_multimodal_payload_metrics([pending]) + + +def print_multimodal_payload_metrics( + metrics: Mapping[str, int | float], +) -> None: + """Record metrics for the logger and print a stable, scrapeable line.""" + if not metrics: + return + with _PENDING_PAYLOAD_METRICS_LOCK: + merged = merge_multimodal_payload_metrics([_PENDING_PAYLOAD_METRICS, metrics]) + _PENDING_PAYLOAD_METRICS.clear() + _PENDING_PAYLOAD_METRICS.update(merged) + values = [] + for key, value in sorted(metrics.items()): + rendered = f"{value:.6f}" if isinstance(value, float) else str(value) + values.append(f"{key}={rendered}") + print("ā–¶ [PAYLOAD] " + ", ".join(values), flush=True) diff --git a/pyrefly.toml b/pyrefly.toml index 5f55643e09b..494c28a227f 100644 --- a/pyrefly.toml +++ b/pyrefly.toml @@ -213,6 +213,7 @@ project-includes = [ "nemo_rl/utils/config.py", "nemo_rl/utils/fastokens.py", "nemo_rl/utils/grad_norm.py", + "nemo_rl/utils/multimodal_payload_metrics.py", "nemo_rl/utils/native_checkpoint.py", "nemo_rl/utils/nsys.py", "nemo_rl/utils/nvml.py", diff --git a/tests/unit/algorithms/test_async_utils.py b/tests/unit/algorithms/test_async_utils.py index 1e274504fbd..b1c68fa2e17 100644 --- a/tests/unit/algorithms/test_async_utils.py +++ b/tests/unit/algorithms/test_async_utils.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -42,10 +42,12 @@ GRPOConfig, MasterConfig, _get_next_nemo_gym_task_index, + _should_normalize_sparse_replay_media, add_grpo_token_loss_masks_and_generation_logprobs, extract_initial_prompt_messages, ) from nemo_rl.data.interfaces import DatumSpec, LLMMessageLogType +from nemo_rl.data.multimodal_utils import PackedTensor from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.environments.interfaces import ( EnvironmentInterface, @@ -356,6 +358,67 @@ def test_local_load_state_dict_validates_checkpoint_shape(self): } ) + def test_local_actor_side_checkpoint_preserves_compact_media_and_resume_metadata( + self, tmp_path + ): + checkpoint_path = tmp_path / "replay_buffer.pt" + compact_media = ( + PackedTensor(torch.tensor([[1.0, 2.0]]), dim_to_pack=0) + .enable_deduplication() + .repeat_interleave(2) + ) + source = ReplayBufferImpl(max_size=10) + assert ( + source.add( + { + "batch": {"pixel_values": compact_media}, + "rollout_metrics": {}, + "_ng_task_index": 7, + }, + weight_version=4, + target_weight_version=5, + ) + == "success" + ) + assert ( + source.add( + { + "batch": {"data": "stale"}, + "rollout_metrics": {}, + "_ng_task_index": 41, + }, + weight_version=0, + target_weight_version=5, + ) + == "success" + ) + + assert source.save_to_path(str(checkpoint_path)) == 2 + + restored = ReplayBufferImpl(max_size=10) + metadata = restored.load_from_path( + str(checkpoint_path), + num_prompts_per_step=1, + current_training_step=5, + max_age_steps=1, + ) + + # Metadata accounts for every saved task index, including trajectories + # discarded during resume cleanup, so an index is never reused. + assert metadata == { + "num_trajectories": 2, + "next_ng_task_index": 42, + } + assert restored.size() == 1 + restored_state = restored.state_dict() + restored_media = restored_state["trajectories"][0]["batch"]["pixel_values"] + assert len(restored_media) == 2 + assert len(restored_media.tensors) == 1 + torch.testing.assert_close( + restored_media.as_tensor(), + torch.tensor([[1.0, 2.0], [1.0, 2.0]]), + ) + class TestReplayBuffer: """Test cases for ReplayBuffer.""" @@ -1006,39 +1069,109 @@ def test_replay_buffer_remove_incomplete_resets_watermark_before_first_remaining ray.kill(buffer) - def test_replay_buffer_checkpoint_with_torch_save(self): - """Test that state_dict can be saved and loaded with torch.save/load.""" + def test_replay_buffer_checkpoint_with_torch_save(self, tmp_path): + """Actor-side compact replay checkpoint survives a config flag flip.""" buffer1 = ReplayBuffer.remote(max_size=10) trajectory = { "batch": { "token_ids": torch.tensor([1, 2, 3]), "rewards": torch.tensor([0.5]), + "pixel_values": PackedTensor(torch.tensor([[1.0, 2.0]]), dim_to_pack=0) + .enable_deduplication() + .repeat_interleave(2), }, "rollout_metrics": {"reward": 1.0, "length": 10}, "timestamp": 12345.0, + "_ng_task_index": 11, } ray.get( buffer1.add.remote(trajectory, weight_version=5, target_weight_version=6) ) - state = ray.get(buffer1.state_dict.remote()) - with tempfile.NamedTemporaryFile(suffix=".pt", delete=False) as f: - torch.save(state, f.name) - checkpoint_path = f.name + checkpoint_path = tmp_path / "replay_buffer.pt" + assert ray.get(buffer1.save_to_path.remote(str(checkpoint_path))) == 1 ray.kill(buffer1) - loaded_state = torch.load(checkpoint_path, weights_only=False) buffer2 = ReplayBuffer.remote(max_size=10) - ray.get(buffer2.load_state_dict.remote(loaded_state)) + restore_metadata = ray.get(buffer2.load_from_path.remote(str(checkpoint_path))) + assert restore_metadata == { + "num_trajectories": 1, + "next_ng_task_index": 12, + } assert ray.get(buffer2.size.remote()) == 1 debug_info = ray.get(buffer2.get_debug_info.remote()) assert debug_info["trajectory_versions"] == [5] assert debug_info["target_weight_versions"] == [6] + restored_state = ray.get(buffer2.state_dict.remote()) + restored_media = restored_state["trajectories"][0]["batch"]["pixel_values"] + assert restored_media.deduplication_enabled + assert len(restored_media) == 2 + assert len(restored_media.tensors) == 1 + torch.testing.assert_close( + restored_media.as_tensor(), + torch.tensor([[1.0, 2.0], [1.0, 2.0]]), + ) + + restored_sparse_batches = [ + BatchedDataDict( + { + "token_ids": torch.tensor([[1]]), + "pixel_values": restored_media.slice([0]), + } + ), + BatchedDataDict({"token_ids": torch.tensor([[2]])}), + ] + assert _should_normalize_sparse_replay_media( + restored_sparse_batches, + deduplicate_multimodal_data=False, + ) + restored_sparse = BatchedDataDict.from_batches( + restored_sparse_batches, + allow_missing_packed_tensors=True, + ) + assert len(restored_sparse["pixel_values"]) == 2 + assert restored_sparse["pixel_values"].logical_segment_count == 1 + + # The representation is self-describing: after a flag-on checkpoint is + # restored by a flag-off run, newly collected legacy media can be + # concatenated without expanding the restored physical segment. + legacy_media = PackedTensor(torch.tensor([[3.0, 4.0]]), dim_to_pack=0) + mixed_after_flag_off = BatchedDataDict.from_batches( + [ + {"pixel_values": restored_media}, + {"pixel_values": legacy_media}, + ], + allow_missing_packed_tensors=True, + ) + assert len(mixed_after_flag_off["pixel_values"].tensors) == 2 + torch.testing.assert_close( + mixed_after_flag_off["pixel_values"].as_tensor(), + torch.tensor([[1.0, 2.0], [1.0, 2.0], [3.0, 4.0]]), + ) + + # The inverse transition is valid too: legacy checkpoint media is + # assigned fresh provenance when combined with new compact media. + new_compact_media = ( + PackedTensor(torch.tensor([[5.0, 6.0]]), dim_to_pack=0) + .enable_deduplication() + .repeat_interleave(2) + ) + mixed_after_flag_on = BatchedDataDict.from_batches( + [ + {"pixel_values": legacy_media}, + {"pixel_values": new_compact_media}, + ], + allow_missing_packed_tensors=True, + ) + assert len(mixed_after_flag_on["pixel_values"].tensors) == 2 + torch.testing.assert_close( + mixed_after_flag_on["pixel_values"].as_tensor(), + torch.tensor([[3.0, 4.0], [5.0, 6.0], [5.0, 6.0]]), + ) - os.unlink(checkpoint_path) ray.kill(buffer2) def test_resume_deadlock_precondition_detectable(self): @@ -1159,6 +1292,31 @@ def test_collection_loop_marks_data_exhausted_on_natural_completion(self): assert status["errored"] is False assert status["running"] is False + @pytest.mark.asyncio + async def test_drain_payload_metrics_returns_collector_interval(self, monkeypatch): + collector = self.create_local_collector() + collector.master_config.grpo.debug_payload_metrics = True + monkeypatch.setattr( + "nemo_rl.algorithms.async_utils.trajectory_collector." + "drain_multimodal_payload_metrics", + lambda: { + "payload_bytes/nemo_gym_return/serialized": 180, + "payload_bytes/nemo_gym_return/serialized_mean_per_call": 90, + "payload_bytes/nemo_gym_return/physical_media": 30, + "payload_bytes/nemo_gym_return/logical_media": 150, + "payload_counts/nemo_gym_return/calls": 2, + "payload_ratio/nemo_gym_return/physical_to_logical": 0.2, + }, + ) + + metrics = await collector.drain_payload_metrics() + + assert metrics["payload_counts/nemo_gym_return/calls"] == 2 + assert metrics["payload_bytes/nemo_gym_return/serialized_mean_per_call"] == 90 + assert metrics["payload_bytes/nemo_gym_return/physical_media"] == 30 + assert metrics["payload_bytes/nemo_gym_return/logical_media"] == 150 + assert metrics["payload_ratio/nemo_gym_return/physical_to_logical"] == 0.2 + def test_collection_loop_marks_errored_on_crash(self): """A crash sets errored (not data_exhausted) so driver guards fail fast.""" collector = self.create_local_collector() @@ -1423,7 +1581,8 @@ class FakeBatch: def slice(self, start, end): return self - def repeat_interleave(self, repeats): + def repeat_interleave(self, repeats, *, share_immutable_media=False): + assert not share_immutable_media return self class FailingThread: @@ -1487,6 +1646,7 @@ def is_alive(self): target_weight = 7 collector = self.create_local_collector(replay_buffer=FakeReplayBuffer()) + collector.master_config.grpo.deduplicate_multimodal_data = True collector.running = True def reserve_target(generation_weight_version): diff --git a/tests/unit/algorithms/test_grpo.py b/tests/unit/algorithms/test_grpo.py index 479e5479919..498f742a42c 100644 --- a/tests/unit/algorithms/test_grpo.py +++ b/tests/unit/algorithms/test_grpo.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -42,7 +42,9 @@ _raise_if_reward_penalties_enabled_without_nemo_gym, _resolve_logprob_skip_flags, _resolve_message_level_advantage_penalties, + _save_async_replay_buffer_checkpoint, _should_use_async_rollouts, + _validate_multimodal_dedup_capability, _validate_use_kl_in_reward_compat, aggregate_rollout_metrics, async_grpo_train, @@ -61,6 +63,7 @@ ) from nemo_rl.algorithms.utils import calculate_baseline_and_std_per_prompt from nemo_rl.data.interfaces import DatumSpec, LLMMessageLogType +from nemo_rl.data.multimodal_utils import PackedTensor from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.environments.interfaces import ( EnvironmentInterface, @@ -83,6 +86,36 @@ def _mock_policy_generation() -> MagicMock: return policy_generation +@pytest.mark.parametrize( + ("checkpointing_config", "expected_count"), + [ + ({}, 7), + ({"save_replay_buffer": True}, 7), + ({"save_replay_buffer": False}, None), + ], +) +def test_save_async_replay_buffer_checkpoint_gate( + tmp_path, checkpointing_config, expected_count +): + replay_buffer = MagicMock() + replay_buffer.save_to_path.remote.return_value = 7 + + with patch("nemo_rl.algorithms.grpo.ray.get", side_effect=lambda value: value): + count = _save_async_replay_buffer_checkpoint( + replay_buffer, + str(tmp_path), + checkpointing_config, + ) + + assert count == expected_count + if expected_count is None: + replay_buffer.save_to_path.remote.assert_not_called() + else: + replay_buffer.save_to_path.remote.assert_called_once_with( + str(tmp_path / "replay_buffer.pt") + ) + + @patch("nemo_rl.algorithms.grpo.ray") def test_refit_policy_generation_forwards_kv_scales_on_colocated_ipc( mock_ray: MagicMock, @@ -715,6 +748,25 @@ def test_raise_if_message_level_advantage_penalties_enabled_raises_when_set( _raise_if_message_level_advantage_penalties_enabled(master_config) +def test_multimodal_dedup_rejects_unqualified_transfer_paths( + mock_grpo_components, +): + master_config = mock_grpo_components["master_config"] + master_config.grpo.deduplicate_multimodal_data = True + master_config.policy["generation"]["backend"] = "sglang" + + with pytest.raises(NotImplementedError, match="backend=vllm"): + _validate_multimodal_dedup_capability(master_config) + + master_config.policy["generation"]["backend"] = "vllm" + master_config.data_plane = {"enabled": True} + with pytest.raises(NotImplementedError, match="data_plane.enabled=false"): + _validate_multimodal_dedup_capability(master_config) + + master_config.data_plane = {"enabled": False} + _validate_multimodal_dedup_capability(master_config) + + def test_grpo_sync_seq_logprob_error_helper_accepts_dict_result(monkeypatch): from nemo_rl.algorithms import grpo_sync as grpo_sync_mod @@ -857,6 +909,25 @@ def _load_state_dict(state, *args, **kwargs): mock.remote = MagicMock(side_effect=_load_state_dict) return mock + @property + def save_to_path(self): + """Return a mock that checkpoints state without a driver-sized return.""" + mock = MagicMock() + mock.remote = MagicMock(return_value=self._size) + return mock + + @property + def load_from_path(self): + """Return compact restore metadata.""" + mock = MagicMock() + mock.remote = MagicMock( + return_value={ + "num_trajectories": self._size, + "next_ng_task_index": 0, + } + ) + return mock + @property def get_trajectories_needed(self): """Return a mock that reports how many prompt groups are still needed.""" @@ -1656,6 +1727,63 @@ def test_dapo_dynamic_sampling_batch_caching(mock_grpo_components): assert batch_cache is not None +def test_dapo_cache_aligns_deduplicated_media_with_text_only_batch( + mock_grpo_components, +): + def make_batch(prompt: str, *, with_media: bool) -> BatchedDataDict: + message_logs = [ + [ + {"role": "user", "content": prompt}, + {"role": "assistant", "content": f"response_{i}"}, + ] + for i in range(3) + ] + batch = create_mock_batch(3, ["math"] * 3, message_logs) + batch["total_reward"] = torch.tensor([1.0, 0.0, 0.5]) + if with_media: + media = PackedTensor( + torch.tensor([[1.0]]), dim_to_pack=0 + ).enable_deduplication() + batch["pixel_values"] = media.repeat_interleave(3) + return batch + + master_config = mock_grpo_components["master_config"] + master_config.grpo.use_dynamic_sampling = True + master_config.grpo.num_prompts_per_step = 2 + master_config.grpo.num_generations_per_prompt = 3 + master_config.grpo.dynamic_sampling_max_gen_batches = 5 + master_config.grpo.deduplicate_multimodal_data = True + std = torch.tensor([0.4, 0.4, 0.4]) + baseline = torch.tensor([0.5, 0.5, 0.5]) + + _, complete, cache, _ = dynamic_sampling( + make_batch("visual", with_media=True), + std, + baseline, + dynamic_sampling_num_gen_batches=1, + master_config=master_config, + timer=Timer(), + ) + assert not complete + assert cache is not None + + result, complete, _, _ = dynamic_sampling( + make_batch("text", with_media=False), + std, + baseline, + dynamic_sampling_num_gen_batches=2, + master_config=master_config, + timer=Timer(), + batch_cache=cache, + ) + + assert complete + assert result.size == 6 + assert len(result["pixel_values"]) == 6 + assert len(result["pixel_values"].tensors) == 1 + assert result["pixel_values"].slice([3, 4, 5]).as_tensor() is None + + def test_dapo_dynamic_sampling_disabled(mock_grpo_components): """Test that when dynamic sampling is disabled, all prompts are kept regardless of std.""" batch_size = 6 @@ -1947,7 +2075,7 @@ def test_noncolocated_opd_teacher_must_fit_on_one_cluster_node( "initial_skip_flag", [None, False], ) -def test_setup_auto_enables_skip_reference_policy_logprobs_when_kl_penalty_zero( +def test_setup_auto_enables_skip_reference_logprobs_with_legacy_policy_factory( monkeypatch, mock_grpo_components, initial_skip_flag ): from nemo_rl.algorithms import grpo as grpo_mod @@ -2002,6 +2130,29 @@ def prepare_refit_info(self): def set_rollout_num_gpus_per_engine(self, _num_gpus_per_engine): pass + def legacy_policy_factory( + *, + cluster, + config, + tokenizer, + processor, + weights_path, + optimizer_path, + init_optimizer, + init_reference_model, + ): + del ( + cluster, + config, + tokenizer, + processor, + weights_path, + optimizer_path, + init_optimizer, + init_reference_model, + ) + return DummyPolicy() + class DummySGLangGeneration: num_gpus_per_engine = 1 @@ -2023,7 +2174,6 @@ def init_collective(self, *_args, **_kwargs): ) monkeypatch.setattr(grpo_mod, "StatefulDataLoader", DummyLoader) monkeypatch.setattr(grpo_mod, "RayVirtualCluster", DummyCluster) - monkeypatch.setattr(grpo_mod, "Policy", lambda *_args, **_kwargs: DummyPolicy()) monkeypatch.setattr( grpo_mod, "SGLangGeneration", @@ -2064,7 +2214,13 @@ def init_collective(self, *_args, **_kwargs): dataset = MagicMock() dataset.__len__ = MagicMock(return_value=1) - grpo_mod.setup(master_config, tokenizer, dataset, None) + grpo_mod.setup( + master_config, + tokenizer, + dataset, + None, + policy_factory=legacy_policy_factory, + ) assert master_config.grpo.skip_reference_policy_logprobs_calculation is True diff --git a/tests/unit/algorithms/test_grpo_router_replay_async.py b/tests/unit/algorithms/test_grpo_router_replay_async.py index 52aa8369ac5..824ddd41e72 100644 --- a/tests/unit/algorithms/test_grpo_router_replay_async.py +++ b/tests/unit/algorithms/test_grpo_router_replay_async.py @@ -24,6 +24,7 @@ _initial_grpo_save_state, async_grpo_train, ) +from nemo_rl.data.multimodal_utils import PackedTensor from nemo_rl.distributed.batched_data_dict import BatchedDataDict @@ -109,6 +110,58 @@ def test_build_async_grpo_train_data_preserves_routed_experts_for_r3( assert "routed_experts" not in train_data +def test_build_async_grpo_train_data_accepts_all_text_vlm_replay_batch(): + flat_messages = BatchedDataDict( + { + "token_ids": torch.tensor([[1, 2, 3]]), + "generation_logprobs": torch.zeros(1, 3), + "token_loss_mask": torch.tensor([[0, 1, 1]]), + } + ) + input_lengths = torch.tensor([3]) + repeated_batch = BatchedDataDict({"loss_multiplier": torch.tensor([1.0])}) + + train_data = _build_async_grpo_train_data( + flat_messages, + input_lengths, + repeated_batch, + {**_make_async_master_config().policy, "is_vlm": True}, + ) + + assert train_data["input_ids"].tolist() == [[1, 2, 3]] + assert train_data.get_multimodal_dict(as_tensors=False) == {} + + +def test_build_async_grpo_train_data_precasts_pixels_without_expanding_dedup(): + pixels = PackedTensor( + [torch.randn(2, 3, 8, 8, dtype=torch.float32)], dim_to_pack=0 + ).enable_deduplication() + pixels = pixels.repeat_interleave(4) + flat_messages = BatchedDataDict( + { + "token_ids": torch.tensor([[1, 2, 3]] * 4), + "generation_logprobs": torch.zeros(4, 3), + "token_loss_mask": torch.tensor([[0, 1, 1]] * 4), + "pixel_values": pixels, + } + ) + + train_data = _build_async_grpo_train_data( + flat_messages, + torch.tensor([3] * 4), + BatchedDataDict({"loss_multiplier": torch.ones(4)}), + {"router_replay": {"enabled": False}}, + ) + + cast_pixels = train_data["pixel_values"] + assert isinstance(cast_pixels, PackedTensor) + assert len(cast_pixels) == 4 + assert len(cast_pixels.tensors) == 1 + assert cast_pixels.logical_segment_count == 4 + assert cast_pixels.tensors[0].dtype == torch.bfloat16 + assert pixels.tensors[0].dtype == torch.float32 + + def test_async_grpo_r3_data_plane_directs_to_single_controller(): master_config = _make_async_master_config(data_plane={"enabled": True}) diff --git a/tests/unit/data/test_llm_message_utils.py b/tests/unit/data/test_llm_message_utils.py index 113fd9ce0b9..8a8e0ab93ca 100644 --- a/tests/unit/data/test_llm_message_utils.py +++ b/tests/unit/data/test_llm_message_utils.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -32,6 +32,7 @@ get_keys_from_message_log, message_log_to_flat_messages, ) +from nemo_rl.data.multimodal_utils import PackedTensor @pytest.fixture @@ -922,6 +923,53 @@ def test_get_formatted_message_log_debug_off_by_default( assert "DEBUG: Individual message turns" not in captured.out +@pytest.mark.parametrize("visual_first", [False, True]) +def test_batched_flatten_aligns_nested_sparse_multimodal_rows(visual_first): + media = PackedTensor(torch.ones(1, 3, 2, 2), dim_to_pack=0).enable_deduplication() + visual_log = [ + { + "role": "user", + "content": "look", + "token_ids": torch.tensor([1]), + "pixel_values": media, + }, + { + "role": "assistant", + "content": "seen", + "token_ids": torch.tensor([2]), + }, + ] + text_log = [ + { + "role": "user", + "content": "text only", + "token_ids": torch.tensor([3]), + }, + { + "role": "assistant", + "content": "answer", + "token_ids": torch.tensor([4]), + }, + ] + message_logs = [visual_log, text_log] if visual_first else [text_log, visual_log] + + flat, _ = batched_message_log_to_flat_message( + message_logs, pad_value_dict={"token_ids": 0} + ) + packed = flat["pixel_values"] + + assert isinstance(packed, PackedTensor) + assert len(packed) == 2 + assert packed.logical_segment_counts_by_row() == ( + [1, 0] if visual_first else [0, 1] + ) + assert len(packed.tensors) == 1 + torch.testing.assert_close( + flat.get_multimodal_dict(as_tensors=True)["pixel_values"], + torch.ones(1, 3, 2, 2), + ) + + def test_get_formatted_message_log_debug_enabled( raw_chat_message_log: LLMMessageLogType, capsys, diff --git a/tests/unit/data/test_multimodal_dict.py b/tests/unit/data/test_multimodal_dict.py index 23b7cdacdbd..b97bbf9d204 100644 --- a/tests/unit/data/test_multimodal_dict.py +++ b/tests/unit/data/test_multimodal_dict.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -11,7 +11,10 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +from copy import deepcopy + import pytest +import ray.cloudpickle as cloudpickle import torch from nemo_rl.data.llm_message_utils import batched_message_log_to_flat_message @@ -480,3 +483,139 @@ def test_slice_preserves_pad_to_max_shape_flag(): assert sliced.pad_to_max_shape is True assert sliced.as_tensor().shape == (2, 3, 4, 4) + + +def test_packedtensor_dedup_uses_provenance_not_prompt_position(): + """Only segments descended from the same physical media are compacted.""" + shared = PackedTensor(torch.tensor([[1.0]]), dim_to_pack=0) + shared.enable_deduplication() + shared_copy = deepcopy(shared) + same_prompt_but_different_media = PackedTensor(torch.tensor([[1.0]]), dim_to_pack=0) + same_prompt_but_different_media.enable_deduplication() + + packed = PackedTensor.concat([shared, shared_copy, same_prompt_but_different_media]) + + assert len(packed) == 3 + assert packed.logical_segment_count == 3 + assert len(packed.tensors) == 2 + torch.testing.assert_close(packed.as_tensor(), torch.tensor([[1.0], [1.0], [1.0]])) + + +def test_packedtensor_multiturn_csr_preserves_shared_seed_and_unique_media(): + """Diverged rows retain one seed segment plus their own later segment.""" + seed = PackedTensor(torch.tensor([[1.0]]), dim_to_pack=0) + seed.enable_deduplication() + row_1 = PackedTensor.merge_segments( + [deepcopy(seed), PackedTensor(torch.tensor([[2.0]]), dim_to_pack=0)] + ) + row_2 = PackedTensor.merge_segments( + [deepcopy(seed), PackedTensor(torch.tensor([[3.0]]), dim_to_pack=0)] + ) + + packed = PackedTensor.flattened_concat([row_1, row_2]) + + assert len(packed) == 2 + assert packed.logical_segment_count == 4 + assert len(packed.tensors) == 3 + torch.testing.assert_close( + packed.as_tensor(), torch.tensor([[1.0], [2.0], [1.0], [3.0]]) + ) + + second_row = packed.slice([1]) + assert len(second_row) == 1 + assert len(second_row.tensors) == 2 + torch.testing.assert_close(second_row.as_tensor(), torch.tensor([[1.0], [3.0]])) + + +def test_packedtensor_dedup_expands_before_dynamic_shape_padding(): + """Logical order is restored before non-packing dimensions are padded.""" + first = PackedTensor( + torch.ones(1, 1, 2), + dim_to_pack=0, + pad_to_max_shape=True, + ).enable_deduplication() + second = PackedTensor( + 2 * torch.ones(1, 2, 1), + dim_to_pack=0, + pad_to_max_shape=True, + ).enable_deduplication() + + packed = PackedTensor.concat([first, deepcopy(first), second]) + materialized = packed.as_tensor() + + assert materialized.shape == (3, 2, 2) + torch.testing.assert_close(materialized[0], materialized[1]) + torch.testing.assert_close(materialized[2, :, 0], 2 * torch.ones(2)) + + +def test_packedtensor_dedup_dim_one_slice_empty_and_cloudpickle_roundtrip(): + first = torch.tensor([[1.0], [2.0]]) + second = torch.tensor([[3.0, 4.0], [5.0, 6.0]]) + packed = PackedTensor( + [first, second], + dim_to_pack=1, + ).enable_deduplication() + repeated = packed.repeat_interleave(2) + + assert len(repeated) == 4 + assert len(repeated.tensors) == 2 + torch.testing.assert_close( + repeated.as_tensor(), + torch.cat([first, first, second, second], dim=1), + ) + + selected = repeated.slice([3, 0, -1]) + assert len(selected) == 3 + assert len(selected.tensors) == 2 + torch.testing.assert_close( + selected.as_tensor(), + torch.cat([second, first, second], dim=1), + ) + + restored = cloudpickle.loads(cloudpickle.dumps(selected, protocol=5)) + assert restored.deduplication_enabled + assert len(restored) == 3 + assert len(restored.tensors) == 2 + torch.testing.assert_close(restored.as_tensor(), selected.as_tensor()) + + empty = packed.repeat_interleave(0) + assert len(empty) == 0 + assert empty.logical_segment_count == 0 + assert empty.as_tensor() is None + + +def test_packedtensor_unpickles_pre_deduplication_state(): + tensor = torch.tensor([[1.0], [2.0]]) + legacy = PackedTensor.__new__(PackedTensor) + legacy.__dict__ = { + "tensors": [tensor], + "dim_to_pack": 0, + "pad_to_max_shape": False, + } + + restored = cloudpickle.loads(cloudpickle.dumps(legacy, protocol=5)) + + assert not restored.deduplication_enabled + assert len(restored) == 1 + assert restored.logical_segment_count == 1 + torch.testing.assert_close(restored.as_tensor(), tensor) + restored.enable_deduplication() + assert restored.deduplication_enabled + + +def test_packedtensor_empty_legacy_rows_survive_copy_pickle_and_slice(): + legacy = PackedTensor(torch.tensor([[1.0]]), dim_to_pack=0) + empty = PackedTensor.empty_rows_like(legacy, 0) + + assert len(empty) == 0 + assert not empty.deduplication_enabled + assert empty.as_tensor() is None + + copied = deepcopy(empty) + restored = cloudpickle.loads(cloudpickle.dumps(empty, protocol=5)) + sliced = empty.slice([]) + for value in (copied, restored, sliced): + assert len(value) == 0 + assert value.logical_segment_count == 0 + assert not value.deduplication_enabled + assert value.as_tensor() is None diff --git a/tests/unit/distributed/test_batched_data_dict.py b/tests/unit/distributed/test_batched_data_dict.py index 46b35f69a5f..7a478027cf5 100644 --- a/tests/unit/distributed/test_batched_data_dict.py +++ b/tests/unit/distributed/test_batched_data_dict.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -11,6 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import numpy as np import pytest import torch @@ -47,6 +48,27 @@ def test_shard_by_batch_size_basic(): assert torch.equal(sharded[1]["other_tensor"], torch.tensor([12, 13, 16, 17])) +def test_from_batches_flag_off_keeps_legacy_first_field_batch_size(): + class MetadataWithoutBatchLength: + def __len__(self): + raise AssertionError("flag-off batch sizing must not scan metadata") + + batches = [ + { + "tokens": torch.tensor([[1], [2]]), + "metadata": MetadataWithoutBatchLength(), + }, + { + "tokens": torch.tensor([[3]]), + "metadata": MetadataWithoutBatchLength(), + "extra": torch.tensor([1]), + }, + ] + + with pytest.raises(KeyError, match="'extra'"): + BatchedDataDict.from_batches(batches) + + def test_shard_by_batch_size_list_data(): """Test shard_by_batch_size with list data.""" # Create a sample batch with list data @@ -625,6 +647,113 @@ def test_shard_by_batch_size_with_packed_multimodal(): assert tuple(shards[1]["pixel_values"].as_tensor().shape) == (6, 3, 8, 8) +def test_repeat_interleave_shares_only_flagged_multimodal_segments(): + media = PackedTensor(torch.ones(1, 3, 2, 2), dim_to_pack=0) + batch = BatchedDataDict( + { + "message_log": [ + [ + { + "role": "user", + "content": "look", + "token_ids": torch.tensor([1, 2]), + "pixel_values": media, + } + ] + ] + } + ) + + flag_off = batch.repeat_interleave(2) + off_first = flag_off["message_log"][0][0]["pixel_values"] + off_second = flag_off["message_log"][1][0]["pixel_values"] + assert not off_first.deduplication_enabled + assert off_first.tensors[0] is not off_second.tensors[0] + + flag_on = batch.repeat_interleave(2, share_immutable_media=True) + on_first = flag_on["message_log"][0][0]["pixel_values"] + on_second = flag_on["message_log"][1][0]["pixel_values"] + assert on_first.deduplication_enabled + assert on_first is not on_second + assert on_first.tensors[0] is on_second.tensors[0] + assert flag_on["message_log"][0] is not flag_on["message_log"][1] + + +def test_repeat_interleave_shares_native_image_video_and_audio_leaves(): + image = torch.ones(1, 2) + video = np.ones((2, 2), dtype=np.float32) + audio = np.ones(16, dtype=np.float32) + batch = BatchedDataDict( + { + "vllm_images": [[image]], + "vllm_videos": [[video]], + "vllm_audios": [[(audio, 16_000)]], + } + ) + + flag_off = batch.repeat_interleave(2) + assert flag_off["vllm_images"][0][0] is not flag_off["vllm_images"][1][0] + assert flag_off["vllm_videos"][0][0] is not flag_off["vllm_videos"][1][0] + assert flag_off["vllm_audios"][0][0][0] is not flag_off["vllm_audios"][1][0][0] + + flag_on = batch.repeat_interleave(2, share_immutable_media=True) + assert flag_on["vllm_images"][0] is not flag_on["vllm_images"][1] + assert flag_on["vllm_images"][0][0] is flag_on["vllm_images"][1][0] + assert flag_on["vllm_videos"][0][0] is flag_on["vllm_videos"][1][0] + assert flag_on["vllm_audios"][0][0][0] is flag_on["vllm_audios"][1][0][0] + + +def test_shards_reintern_shared_segments_locally(): + media = PackedTensor(torch.ones(1, 2), dim_to_pack=0).enable_deduplication() + repeated_media = media.repeat_interleave(4) + batch = BatchedDataDict( + { + "input_ids": torch.arange(8).reshape(4, 2), + "input_lengths": torch.tensor([2, 2, 2, 2]), + "pixel_values": repeated_media, + } + ) + + shards = batch.shard_by_batch_size(shards=2) + + assert [len(shard["pixel_values"]) for shard in shards] == [2, 2] + assert [len(shard["pixel_values"].tensors) for shard in shards] == [1, 1] + + +def test_sequence_packing_reinterns_shared_segments_per_shard_for_cp_padding(): + media = PackedTensor(torch.ones(1, 2), dim_to_pack=0) + repeated_media = media.enable_deduplication().repeat_interleave(8) + sequence_lengths = torch.tensor([5, 6, 7, 8, 9, 10, 11, 12]) + batch = BatchedDataDict( + { + "input_ids": torch.arange(8 * 12).reshape(8, 12), + "input_lengths": sequence_lengths, + "pixel_values": repeated_media, + } + ) + sequence_packing_args = SequencePackingArgs( + max_tokens_per_microbatch=24, + input_key="input_ids", + input_lengths_key="input_lengths", + algorithm="modified_first_fit_decreasing", + # CP=2 requires sequences to be divisible by 2 * CP. + sequence_length_pad_multiple=4, + ) + + shards, _ = batch.shard_by_batch_size( + shards=2, + sequence_packing_args=sequence_packing_args, + ) + + assert sum(len(shard["pixel_values"]) for shard in shards) == 8 + assert [len(shard["pixel_values"].tensors) for shard in shards] == [1, 1] + for shard in shards: + torch.testing.assert_close( + shard["pixel_values"].as_tensor(), + torch.ones(len(shard["pixel_values"]), 2), + ) + + def test_shard_by_batch_size_allow_uneven_empty_shards_preserve_all_keys(): """Empty trailing shards should preserve all keys with empty values.""" batch = BatchedDataDict( @@ -643,6 +772,7 @@ def test_shard_by_batch_size_allow_uneven_empty_shards_preserve_all_keys(): # Empty trailing shards should preserve all keys and use empty values. for empty_shard in shards[2:]: + assert empty_shard.size == 0 for key, original_value in batch.items(): assert key in empty_shard shard_value = empty_shard[key] @@ -650,6 +780,7 @@ def test_shard_by_batch_size_allow_uneven_empty_shards_preserve_all_keys(): assert shard_value.shape[0] == 0 elif isinstance(original_value, PackedTensor): assert isinstance(shard_value, PackedTensor) + assert len(shard_value) == 0 assert shard_value.as_tensor() is None else: assert shard_value == [] @@ -695,6 +826,39 @@ def test_get_multimodal_dict_mixed_content_and_device_move(): ].device.type == ("cuda" if torch.cuda.is_available() else "cpu") +def test_get_multimodal_dict_casts_only_pixels_without_materializing_dedup(): + pixels = PackedTensor( + [torch.randn(2, 3, 8, 8, dtype=torch.float32)], dim_to_pack=0 + ).enable_deduplication() + pixels = pixels.repeat_interleave(4) + image_sizes = PackedTensor( + [torch.tensor([[8, 8]], dtype=torch.int64)], dim_to_pack=0 + ).enable_deduplication() + image_sizes = image_sizes.repeat_interleave(4) + original_provenance = list(pixels._segment_provenance) + batch = BatchedDataDict({"pixel_values": pixels, "imgs_sizes": image_sizes}) + + multimodal = batch.get_multimodal_dict(as_tensors=False, pixel_dtype=torch.bfloat16) + cast_pixels = multimodal["pixel_values"] + + assert isinstance(cast_pixels, PackedTensor) + assert len(cast_pixels) == 4 + assert len(cast_pixels.tensors) == 1 + assert cast_pixels.logical_segment_count == 4 + assert cast_pixels.tensors[0].dtype == torch.bfloat16 + assert pixels.tensors[0].dtype == torch.float32 + assert cast_pixels._row_offsets == pixels._row_offsets + assert cast_pixels._segment_indices == pixels._segment_indices + assert cast_pixels._segment_provenance != original_provenance + assert multimodal["imgs_sizes"] is image_sizes + + materialized = batch.get_multimodal_dict( + as_tensors=True, pixel_dtype=torch.bfloat16 + ) + assert materialized["pixel_values"].dtype == torch.bfloat16 + assert materialized["pixel_values"].shape[0] == 8 + + def test_from_batches_pads_3d_tensors_along_sequence_dim(): """from_batches should pad 3D tensors along the sequence dimension before stacking.""" @@ -848,6 +1012,173 @@ def test_from_batches_keeps_keys_missing_from_empty_mapping(): assert torch.equal(stacked["routed_experts"], routed_experts) +def test_from_batches_can_align_optional_deduplicated_media_keys(): + shared_pixels = PackedTensor( + torch.tensor([[1.0]]), dim_to_pack=0 + ).enable_deduplication() + pixel_rows = shared_pixels.repeat_interleave(2) + distinct_image_sizes = PackedTensor( + [torch.tensor([[10, 20]]), torch.tensor([[30, 40]])], + dim_to_pack=0, + ).enable_deduplication() + audio_rows = PackedTensor( + torch.tensor([[5.0]]), dim_to_pack=0 + ).enable_deduplication() + + visual_batch = BatchedDataDict( + { + "input_ids": torch.tensor([[1, 2], [3, 4]]), + "pixel_values": pixel_rows, + "imgs_sizes": distinct_image_sizes, + } + ) + audio_batch = BatchedDataDict( + { + "input_ids": torch.tensor([[5, 6]]), + "audio_values": audio_rows, + } + ) + + stacked = BatchedDataDict.from_batches( + [visual_batch, audio_batch], + allow_missing_packed_tensors=True, + ) + + assert stacked.size == 3 + assert { + key: len(stacked[key]) for key in ("pixel_values", "imgs_sizes", "audio_values") + } == { + "pixel_values": 3, + "imgs_sizes": 3, + "audio_values": 3, + } + assert len(stacked["pixel_values"].tensors) == 1 + assert len(stacked["imgs_sizes"].tensors) == 2 + assert stacked["pixel_values"].slice([2]).as_tensor() is None + assert stacked["imgs_sizes"].slice([2]).as_tensor() is None + assert stacked["audio_values"].slice([0, 1]).as_tensor() is None + torch.testing.assert_close( + stacked["audio_values"].slice([2]).as_tensor(), + torch.tensor([[5.0]]), + ) + + +def test_from_batches_optional_media_rejects_cross_key_row_misalignment(): + batch = BatchedDataDict( + { + "pixel_values": PackedTensor( + torch.tensor([[1.0]]), dim_to_pack=0 + ).enable_deduplication(), + "input_ids": torch.tensor([[1, 2], [3, 4]]), + } + ) + + with pytest.raises(ValueError, match="inconsistent logical row counts"): + BatchedDataDict.from_batches( + [batch], + allow_missing_packed_tensors=True, + ) + + +def test_model_materialization_validates_coupled_media_segment_order(): + pixels = PackedTensor( + [torch.tensor([[1.0]]), torch.tensor([[2.0]])], + dim_to_pack=0, + ).enable_deduplication() + image_sizes = PackedTensor( + [torch.tensor([[10, 20]]), torch.tensor([[30, 40]])], + dim_to_pack=0, + ).enable_deduplication() + valid = BatchedDataDict( + { + "pixel_values": pixels, + "imgs_sizes": image_sizes, + } + ) + + materialized = valid.get_multimodal_dict(as_tensors=True) + torch.testing.assert_close( + materialized["pixel_values"], torch.tensor([[1.0], [2.0]]) + ) + torch.testing.assert_close( + materialized["imgs_sizes"], + torch.tensor([[10, 20], [30, 40]]), + ) + + first_row_only = PackedTensor.merge_segments( + [ + PackedTensor( + torch.tensor([[10, 20]]), dim_to_pack=0 + ).enable_deduplication(), + PackedTensor( + torch.tensor([[30, 40]]), dim_to_pack=0 + ).enable_deduplication(), + ] + ) + missing_second_row = PackedTensor.concat( + [first_row_only, PackedTensor.empty_rows_like(first_row_only, 1)] + ) + invalid = BatchedDataDict( + { + "pixel_values": pixels, + "imgs_sizes": missing_second_row, + } + ) + + with pytest.raises(ValueError, match="ordered per-row segment counts"): + invalid.get_multimodal_dict(as_tensors=True) + + +def test_size_supports_packed_tensor_as_first_key_and_empty_batches(): + media = PackedTensor( + [torch.tensor([[1.0]]), torch.tensor([[2.0]])], + dim_to_pack=0, + ) + batch = BatchedDataDict( + { + "pixel_values": media, + "input_ids": torch.tensor([[1, 2], [3, 4]]), + } + ) + + assert batch.size == 2 + assert BatchedDataDict().size == 0 + + +def test_deduplicated_media_survives_chunk_reorder_and_select_indices(): + first = PackedTensor(torch.tensor([[1.0]]), dim_to_pack=0).enable_deduplication() + second = PackedTensor(torch.tensor([[2.0]]), dim_to_pack=0).enable_deduplication() + media = PackedTensor.concat( + [first.repeat_interleave(2), second.repeat_interleave(2)] + ) + batch = BatchedDataDict( + { + "pixel_values": media, + "input_ids": torch.arange(8).reshape(4, 2), + } + ) + + first_chunk = batch.chunk(rank=0, chunks=2) + assert len(first_chunk["pixel_values"].tensors) == 1 + torch.testing.assert_close( + first_chunk["pixel_values"].as_tensor(), + torch.tensor([[1.0], [1.0]]), + ) + + batch.reorder_data([3, 2, 1, 0]) + torch.testing.assert_close( + batch["pixel_values"].as_tensor(), + torch.tensor([[2.0], [2.0], [1.0], [1.0]]), + ) + + selected = batch.select_indices([0, 3]) + assert len(selected["pixel_values"].tensors) == 2 + torch.testing.assert_close( + selected["pixel_values"].as_tensor(), + torch.tensor([[2.0], [1.0]]), + ) + + @pytest.mark.parametrize("pad_to_multiple_of", [1, 32, 64, 256]) def test_sequence_packing_microbatch_boundaries(pad_to_multiple_of): """Test that microbatch boundaries are correctly maintained across chunks with random sequences.""" diff --git a/tests/unit/environments/test_nemo_gym.py b/tests/unit/environments/test_nemo_gym.py index a4ca5759e98..b32f9d7b947 100644 --- a/tests/unit/environments/test_nemo_gym.py +++ b/tests/unit/environments/test_nemo_gym.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -20,9 +20,11 @@ import ray import requests import torch +from PIL import Image from yaml import safe_load from nemo_rl.algorithms.grpo import MasterConfig +from nemo_rl.data.multimodal_utils import PackedTensor, image_to_data_url from nemo_rl.distributed.ray_actor_environment_registry import ( get_actor_python_env, ) @@ -34,6 +36,7 @@ setup_nemo_gym_config, validate_reward_components_match_scalar, ) +from nemo_rl.experience.rollouts import _reattach_original_multimodal_payloads from nemo_rl.models.generation.vllm import VllmGeneration # cluster and tokenizer are fixture imports @@ -293,7 +296,7 @@ class _MockSelf: result = ( NemoGym.__ray_metadata__.modified_class._postprocess_nemo_gym_to_nemo_rl_result( - _MockSelf(), nemo_gym_result, tokenizer + _MockSelf(), {}, nemo_gym_result, tokenizer ) ) @@ -311,6 +314,295 @@ class _MockSelf: assert nemo_gym_result["response"]["output"][1]["generation_str"] == "6 7" +@pytest.mark.parametrize("include_initial_multimodal_data", [False, True]) +def test_nemo_gym_dedup_redacts_initial_images_from_actor_return( + include_initial_multimodal_data, +): + data_url = image_to_data_url(Image.new("RGB", (2, 2), color="red")) + initial_input = [ + { + "role": "user", + "content": [ + {"type": "input_text", "text": "count"}, + {"type": "input_image", "image_url": data_url}, + ], + } + ] + nemo_gym_result = { + "response": { + "agent_input": deepcopy(initial_input), + "seed_obs": deepcopy(initial_input), + "output": [ + { + "prompt_token_ids": [1, 2], + "generation_token_ids": [3], + "generation_log_probs": [-0.1], + } + ], + }, + "responses_create_params": {"input": deepcopy(initial_input)}, + "reward": 1.0, + } + + class _Tokenizer: + def batch_decode(self, batch): + return ["decoded"] * len(batch) + + class _MockSelf: + cfg = {} + _processor = None + + result = ( + NemoGym.__ray_metadata__.modified_class._postprocess_nemo_gym_to_nemo_rl_result( + _MockSelf(), + {}, + nemo_gym_result, + _Tokenizer(), + include_initial_multimodal_data=include_initial_multimodal_data, + ) + ) + + if include_initial_multimodal_data: + assert "_initial_multimodal_data_omitted" not in result + assert "_nemo_rl_initial_media_omitted" not in result["full_result"] + assert data_url in json.dumps(result["full_result"]) + else: + assert result["_initial_multimodal_data_omitted"] is True + assert result["full_result"]["_nemo_rl_initial_media_omitted"] is True + assert data_url not in json.dumps(result["full_result"]) + assert result["full_result"]["responses_create_params"]["input"][0][ + "content" + ] == [{"type": "input_text", "text": "count"}] + + +def test_nemo_gym_dedup_omits_actor_initial_tensor_and_preserves_later_media(): + initial_url = image_to_data_url(Image.new("RGB", (1, 1), color=(1, 0, 0))) + tool_url = image_to_data_url(Image.new("RGB", (1, 1), color=(2, 0, 0))) + initial_input = [ + { + "role": "user", + "content": [ + {"type": "input_text", "text": "inspect"}, + {"type": "input_image", "image_url": initial_url}, + ], + } + ] + template = { + "response": { + "agent_input": deepcopy(initial_input), + "seed_obs": deepcopy(initial_input), + "output": [ + { + "prompt_token_ids": [1], + "generation_token_ids": [2], + "generation_log_probs": [-0.1], + }, + { + "role": "user", + "content": [ + {"type": "input_image", "image_url": tool_url}, + ], + }, + { + "prompt_token_ids": [1, 2, 3], + "generation_token_ids": [4], + "generation_log_probs": [-0.2], + }, + ], + }, + "responses_create_params": {"input": deepcopy(initial_input)}, + "reward": 1.0, + } + + class _Tokenizer: + def batch_decode(self, batch): + return ["decoded"] * len(batch) + + class _ImageProcessor: + model_input_names = ["pixel_values"] + + class _TextTokenizer: + model_input_names = ["input_ids"] + + class _Processor: + image_token = "" + image_processor = _ImageProcessor() + tokenizer = _TextTokenizer() + model_input_names = ["input_ids", "pixel_values"] + + def __call__(self, *, text, images, return_tensors): + assert text == "" * len(images) + assert return_tensors == "pt" + red_values = [image.getpixel((0, 0))[0] for image in images] + return { + "input_ids": torch.tensor([[1]]), + "pixel_values": torch.tensor(red_values, dtype=torch.float32).view( + -1, 1 + ), + } + + class _MockSelf: + cfg = {} + _processor = _Processor() + + postprocess = ( + NemoGym.__ray_metadata__.modified_class._postprocess_nemo_gym_to_nemo_rl_result + ) + flag_off = postprocess( + _MockSelf(), + {}, + deepcopy(template), + _Tokenizer(), + include_initial_multimodal_data=True, + ) + flag_on = postprocess( + _MockSelf(), + {}, + deepcopy(template), + _Tokenizer(), + include_initial_multimodal_data=False, + ) + + off_users = [ + message for message in flag_off["message_log"] if message["role"] == "user" + ] + on_users = [ + message for message in flag_on["message_log"] if message["role"] == "user" + ] + assert off_users[0]["pixel_values"].as_tensor().item() == 1 + assert off_users[1]["pixel_values"].as_tensor().item() == 2 + assert "pixel_values" not in on_users[0] + assert on_users[1]["pixel_values"].as_tensor().item() == 2 + assert initial_url not in json.dumps(flag_on["full_result"]) + assert tool_url in json.dumps(flag_on["full_result"]) + + original_media = PackedTensor(torch.tensor([[99.0]]), dim_to_pack=0) + _reattach_original_multimodal_payloads( + [flag_on], + [[{"role": "user", "content": "", "pixel_values": original_media}]], + ) + on_users = [ + message for message in flag_on["message_log"] if message["role"] == "user" + ] + assert on_users[0]["pixel_values"] is original_media + assert on_users[1]["pixel_values"].as_tensor().item() == 2 + + +@pytest.mark.parametrize( + ("seed_mode", "expected_pixel_values"), + [ + ("text_only", None), + ("initial_plus_additional", [1.0, 2.0]), + ], +) +def test_nemo_gym_dedup_keeps_authoritative_changed_seed_media( + seed_mode, expected_pixel_values +): + initial_url = image_to_data_url(Image.new("RGB", (1, 1), color=(1, 0, 0))) + additional_url = image_to_data_url(Image.new("RGB", (1, 1), color=(2, 0, 0))) + initial_input = [ + { + "role": "user", + "content": [ + {"type": "input_text", "text": "inspect"}, + {"type": "input_image", "image_url": initial_url}, + ], + } + ] + if seed_mode == "text_only": + seed_obs = [ + { + "role": "user", + "content": [{"type": "input_text", "text": "text only"}], + } + ] + else: + seed_obs = deepcopy(initial_input) + seed_obs[0]["content"].append( + {"type": "input_image", "image_url": additional_url} + ) + + nemo_gym_result = { + "response": { + "agent_input": deepcopy(initial_input), + "seed_obs": seed_obs, + "output": [ + { + "prompt_token_ids": [1], + "generation_token_ids": [2], + "generation_log_probs": [-0.1], + } + ], + }, + "responses_create_params": {"input": deepcopy(initial_input)}, + "reward": 1.0, + } + + class _Tokenizer: + def batch_decode(self, batch): + return ["decoded"] * len(batch) + + class _ImageProcessor: + model_input_names = ["pixel_values"] + + class _TextTokenizer: + model_input_names = ["input_ids"] + + class _Processor: + image_token = "" + image_processor = _ImageProcessor() + tokenizer = _TextTokenizer() + model_input_names = ["input_ids", "pixel_values"] + + def __call__(self, *, text, images, return_tensors): + assert text == "" * len(images) + assert return_tensors == "pt" + red_values = [image.getpixel((0, 0))[0] for image in images] + return { + "input_ids": torch.tensor([[1]]), + "pixel_values": torch.tensor(red_values, dtype=torch.float32).view( + -1, 1 + ), + } + + class _MockSelf: + cfg = {} + _processor = _Processor() + + result = ( + NemoGym.__ray_metadata__.modified_class._postprocess_nemo_gym_to_nemo_rl_result( + _MockSelf(), + {}, + nemo_gym_result, + _Tokenizer(), + include_initial_multimodal_data=False, + ) + ) + + assert result["_initial_multimodal_data_omitted"] is False + user_message = next( + message for message in result["message_log"] if message["role"] == "user" + ) + if expected_pixel_values is None: + assert "pixel_values" not in user_message + else: + assert user_message["pixel_values"].as_tensor().flatten().tolist() == ( + expected_pixel_values + ) + + original_media = PackedTensor(torch.tensor([[99.0]]), dim_to_pack=0) + _reattach_original_multimodal_payloads( + [result], + [[{"role": "user", "content": "", "pixel_values": original_media}]], + ) + if expected_pixel_values is None: + assert "pixel_values" not in user_message + else: + assert user_message["pixel_values"].as_tensor().flatten().tolist() == ( + expected_pixel_values + ) + + def test_nemo_gym_postprocess_no_generation_data_raises(): """When no output item carries generation data, the postprocess should raise a ValueError that reports the prompt length and the response.output item types.""" @@ -335,7 +627,7 @@ class _MockSelf: with pytest.raises(ValueError) as excinfo: NemoGym.__ray_metadata__.modified_class._postprocess_nemo_gym_to_nemo_rl_result( - _MockSelf(), nemo_gym_result, _Tokenizer() + _MockSelf(), {}, nemo_gym_result, _Tokenizer() ) msg = str(excinfo.value) @@ -364,7 +656,7 @@ class _MockSelf: with pytest.raises(ValueError) as excinfo: NemoGym.__ray_metadata__.modified_class._postprocess_nemo_gym_to_nemo_rl_result( - _MockSelf(), nemo_gym_result, _Tokenizer() + _MockSelf(), {}, nemo_gym_result, _Tokenizer() ) msg = str(excinfo.value) diff --git a/tests/unit/environments/test_nemo_gym_router_replay.py b/tests/unit/environments/test_nemo_gym_router_replay.py index 15383d8af5b..f60f629b702 100644 --- a/tests/unit/environments/test_nemo_gym_router_replay.py +++ b/tests/unit/environments/test_nemo_gym_router_replay.py @@ -56,7 +56,7 @@ class _MockSelf: result = ( NemoGym.__ray_metadata__.modified_class._postprocess_nemo_gym_to_nemo_rl_result( - _MockSelf(), nemo_gym_result, _Tokenizer() + _MockSelf(), {}, nemo_gym_result, _Tokenizer() ) ) @@ -90,7 +90,7 @@ class _MockSelf: with pytest.raises(ValueError, match="requires NeMo Gym output items"): NemoGym.__ray_metadata__.modified_class._postprocess_nemo_gym_to_nemo_rl_result( - _MockSelf(), nemo_gym_result, _Tokenizer() + _MockSelf(), {}, nemo_gym_result, _Tokenizer() ) @@ -116,7 +116,7 @@ class _MockSelf: result = ( NemoGym.__ray_metadata__.modified_class._postprocess_nemo_gym_to_nemo_rl_result( - _MockSelf(), nemo_gym_result, _Tokenizer() + _MockSelf(), {}, nemo_gym_result, _Tokenizer() ) ) diff --git a/tests/unit/experience/test_rollouts.py b/tests/unit/experience/test_rollouts.py index 40676612602..0321601313e 100644 --- a/tests/unit/experience/test_rollouts.py +++ b/tests/unit/experience/test_rollouts.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -29,9 +29,11 @@ from nemo_rl.data.datasets.response_datasets import NemoGymDataset from nemo_rl.data.interfaces import DatumSpec from nemo_rl.data.llm_message_utils import batched_message_log_to_flat_message +from nemo_rl.data.multimodal_utils import PackedTensor from nemo_rl.data.processors import nemo_gym_data_processor from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.distributed.virtual_cluster import RayVirtualCluster +from nemo_rl.environments.interfaces import EnvironmentReturn from nemo_rl.environments.games.sliding_puzzle import ( SlidingPuzzleConfig, SlidingPuzzleEnv, @@ -41,12 +43,16 @@ from nemo_rl.experience.metric_utils import calculate_single_metric, pct from nemo_rl.experience.rollout_manager import AsyncNemoGymRolloutImpl from nemo_rl.experience.rollouts import ( + _add_multimodal_generation_payload, + _reattach_original_multimodal_payloads, + async_generate_response_for_sample_turn, generate_responses_async, run_async_multi_turn_rollout, run_async_multi_turn_rollout_groups, run_async_nemo_gym_rollout, run_multi_turn_rollout, run_nemo_gym_rollout_sync, + run_sample_multi_turn_rollout, ) from nemo_rl.models.generation import configure_generation_config from nemo_rl.models.generation.vllm import VllmConfig, VllmGeneration @@ -67,6 +73,241 @@ _MultiStepCalculatorLogic, ) + +def test_reattach_original_multimodal_payloads_is_media_only_and_turn_aligned(): + first_image = PackedTensor(torch.tensor([[1.0]]), dim_to_pack=0) + second_image = PackedTensor(torch.tensor([[2.0]]), dim_to_pack=0) + original_logs = [ + [ + { + "role": "user", + "content": "first", + "pixel_values": first_image, + "request_metadata": {"must_not": "reattach"}, + }, + {"role": "assistant", "content": "answer"}, + { + "role": "user", + "content": "second", + "pixel_values": second_image, + "vllm_videos": ["video.mp4"], + }, + ] + ] + results = [ + { + "_initial_multimodal_data_omitted": True, + "input_message_log": [ + {"role": "user", "content": "first"}, + {"role": "user", "content": "second"}, + ], + "message_log": [ + {"role": "system", "content": "system"}, + { + "role": "user", + "content": "first", + "pixel_values": "remote placeholder", + }, + {"role": "assistant", "content": "answer"}, + {"role": "user", "content": "second"}, + ], + } + ] + + _reattach_original_multimodal_payloads(results, original_logs) + + for log_key in ("input_message_log", "message_log"): + user_messages = [ + message for message in results[0][log_key] if message["role"] == "user" + ] + assert user_messages[0]["pixel_values"] is first_image + assert user_messages[1]["pixel_values"] is second_image + assert user_messages[1]["vllm_videos"] == ["video.mp4"] + assert "request_metadata" not in user_messages[0] + + +@pytest.mark.parametrize("omission_marker", [False, None]) +def test_reattach_keeps_authoritative_changed_gym_media(omission_marker): + original_media = PackedTensor(torch.tensor([[1.0]]), dim_to_pack=0) + effective_media = PackedTensor(torch.tensor([[2.0]]), dim_to_pack=0) + result = { + "message_log": [ + { + "role": "user", + "content": "", + "pixel_values": effective_media, + } + ], + } + if omission_marker is not None: + result["_initial_multimodal_data_omitted"] = omission_marker + results = [result] + original_logs = [ + [ + { + "role": "user", + "content": "", + "pixel_values": original_media, + } + ] + ] + + _reattach_original_multimodal_payloads(results, original_logs) + + assert results[0]["message_log"][0]["pixel_values"] is effective_media + assert "_initial_multimodal_data_omitted" not in results[0] + + +def test_nemo_gym_initial_media_stays_compact_through_replay_and_policy_flatten(): + generations = 16 + initial_media = PackedTensor(torch.ones(1, 3, 2, 2), dim_to_pack=0) + prompt_batch = BatchedDataDict( + { + "message_log": [ + [ + { + "role": "user", + "content": "", + "token_ids": torch.tensor([], dtype=torch.long), + "pixel_values": initial_media, + } + ] + ] + } + ) + repeated = prompt_batch.repeat_interleave(generations, share_immutable_media=True) + results = [ + { + "_initial_multimodal_data_omitted": True, + "input_message_log": [ + { + "role": "user", + "content": "", + "token_ids": torch.tensor([1]), + } + ], + "message_log": [ + { + "role": "user", + "content": "", + "token_ids": torch.tensor([1]), + }, + { + "role": "assistant", + "content": "", + "token_ids": torch.tensor([2]), + }, + ], + } + for _ in range(generations) + ] + + _reattach_original_multimodal_payloads(results, repeated["message_log"]) + replay_batch = BatchedDataDict( + {"message_log": [result["message_log"] for result in results]} + ) + flat, _ = batched_message_log_to_flat_message(replay_batch["message_log"]) + media = flat["pixel_values"] + + assert media.deduplication_enabled + assert len(media) == generations + assert media.logical_segment_count == generations + assert len(media.tensors) == 1 + assert media.as_tensor().shape == (generations, 3, 2, 2) + + +def test_dedup_generation_sends_only_native_vllm_media(): + class _Generation: + cfg = {"backend": "vllm"} + + pixel_values = PackedTensor(torch.tensor([[1.0]]), dim_to_pack=0) + flat_messages = BatchedDataDict( + { + "token_ids": torch.tensor([[1, 2]]), + "pixel_values": pixel_values, + } + ) + active_batch = BatchedDataDict( + { + "vllm_content": [" describe"], + "vllm_images": [[torch.ones(1, 2)]], + } + ) + compact_generation_input = BatchedDataDict() + + _add_multimodal_generation_payload( + compact_generation_input, + flat_messages, + active_batch, + _Generation(), + deduplicate_multimodal_data=True, + ) + + assert "pixel_values" not in compact_generation_input + assert compact_generation_input["vllm_content"] == [" describe"] + assert compact_generation_input["vllm_images"] is active_batch["vllm_images"] + + later_turn_batch = BatchedDataDict( + { + "vllm_content": [None], + "vllm_images": active_batch["vllm_images"], + } + ) + later_turn_generation_input = BatchedDataDict() + _add_multimodal_generation_payload( + later_turn_generation_input, + flat_messages, + later_turn_batch, + _Generation(), + deduplicate_multimodal_data=True, + ) + assert "pixel_values" not in later_turn_generation_input + assert later_turn_generation_input["vllm_content"] == [None] + assert later_turn_generation_input["vllm_images"] is active_batch["vllm_images"] + + legacy_generation_input = BatchedDataDict() + _add_multimodal_generation_payload( + legacy_generation_input, + flat_messages, + active_batch, + _Generation(), + deduplicate_multimodal_data=False, + ) + assert legacy_generation_input["pixel_values"] is pixel_values + + +def test_dedup_generation_keeps_policy_media_for_unconsumed_native_metadata(): + class _Generation: + cfg = {"backend": "vllm"} + + input_features = PackedTensor(torch.tensor([[1.0]]), dim_to_pack=0) + flat_messages = BatchedDataDict( + { + "token_ids": torch.tensor([[1, 2]]), + "input_features": input_features, + } + ) + active_batch = BatchedDataDict( + { + "vllm_content": [[{"type": "audio", "audio": "/tmp/unconsumed-audio.wav"}]], + "vllm_audio_paths": [["/tmp/unconsumed-audio.wav"]], + } + ) + generation_input = BatchedDataDict() + + _add_multimodal_generation_payload( + generation_input, + flat_messages, + active_batch, + _Generation(), + deduplicate_multimodal_data=True, + ) + + assert generation_input["input_features"] is input_features + assert generation_input["vllm_content"] is active_batch["vllm_content"] + assert "vllm_audio_paths" not in generation_input + + MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct" @@ -141,6 +382,12 @@ def test_median_like_p50(self): class _DummyTokenizer: pad_token_id = 0 + def __call__(self, text, return_tensors=True, add_special_tokens=False): + class _Tokens: + input_ids = torch.tensor([[7]], dtype=torch.int64) + + return _Tokens() + def batch_decode(self, generated_ids, skip_special_tokens=True): return ["ok" for _ in generated_ids] @@ -167,6 +414,297 @@ async def generate_async(self, data, greedy=False): ) +class _CapturingAsyncVllmGeneration: + cfg = {"backend": "vllm", "vllm_cfg": {"async_engine": True}} + + def __init__(self): + self.generation_input = None + + async def generate_async(self, data, greedy=False): + self.generation_input = data + input_length = int(data["input_lengths"][0]) + output_ids = torch.cat( + (data["input_ids"][0, :input_length], torch.tensor([9])) + ).unsqueeze(0) + yield ( + 0, + BatchedDataDict( + { + "output_ids": output_ids, + "logprobs": torch.zeros_like(output_ids, dtype=torch.float32), + "generation_lengths": torch.tensor([1], dtype=torch.long), + "unpadded_sequence_lengths": torch.tensor( + [input_length + 1], dtype=torch.long + ), + "truncated": torch.tensor([False], dtype=torch.bool), + } + ), + ) + + +class _CapturingSyncVllmGeneration: + cfg = {"backend": "vllm"} + + def __init__(self): + self.calls = [] + + def generate(self, data, greedy=False): + self.calls.append( + { + "input_ids": data["input_ids"].clone(), + "input_lengths": data["input_lengths"].clone(), + "vllm_content": list(data["vllm_content"]), + "vllm_images": data["vllm_images"], + } + ) + input_lengths = data["input_lengths"].to(dtype=torch.long) + output_ids = torch.zeros( + (len(input_lengths), int(input_lengths.max().item()) + 1), + dtype=torch.long, + ) + for row, input_length in enumerate(input_lengths.tolist()): + output_ids[row, :input_length] = data["input_ids"][row, :input_length] + output_ids[row, input_length] = 9 + return BatchedDataDict( + { + "output_ids": output_ids, + "logprobs": torch.zeros_like(output_ids, dtype=torch.float32), + "generation_lengths": torch.ones(len(input_lengths), dtype=torch.long), + "unpadded_sequence_lengths": input_lengths + 1, + "truncated": torch.zeros(len(input_lengths), dtype=torch.bool), + } + ) + + +@pytest.mark.parametrize("deduplicate_multimodal_data", [False, True]) +def test_sync_vlm_multiturn_drops_stale_native_content( + monkeypatch, deduplicate_multimodal_data +): + generation = _CapturingSyncVllmGeneration() + image = torch.ones(3, 2, 2) + policy_media = PackedTensor(torch.ones(1, 3, 2, 2), dim_to_pack=0) + reward_calls = [] + + def fake_rewards(batch, task_to_env): + reward_calls.append(None) + return EnvironmentReturn( + observations=[{"role": "user", "content": "next"}], + metadata=[None], + next_stop_strings=[None], + rewards=torch.tensor([0.0]), + terminateds=torch.tensor([len(reward_calls) >= 2]), + answers=[None], + ) + + monkeypatch.setattr( + "nemo_rl.experience.rollouts.calculate_rewards", + fake_rewards, + ) + + run_multi_turn_rollout( + policy_generation=generation, + input_batch=BatchedDataDict( + { + "message_log": [ + [ + { + "role": "user", + "content": "", + "token_ids": torch.tensor([1]), + "pixel_values": policy_media, + } + ] + ], + "extra_env_info": [None], + "task_name": ["vlm"], + "stop_strings": [None], + "idx": [0], + "vllm_content": [" initial prompt"], + "vllm_images": [[image]], + } + ), + tokenizer=_DummyTokenizer(), + task_to_env={}, + max_seq_len=32, + max_rollout_turns=2, + deduplicate_multimodal_data=deduplicate_multimodal_data, + ) + + assert len(generation.calls) == 2 + assert generation.calls[0]["vllm_content"] == [" initial prompt"] + assert generation.calls[1]["vllm_content"] == [None] + assert generation.calls[0]["vllm_images"][0][0] is image + assert generation.calls[1]["vllm_images"][0][0] is image + assert generation.calls[0]["input_ids"][0, :1].tolist() == [1] + assert generation.calls[1]["input_ids"][0, :3].tolist() == [1, 9, 7] + + +def test_async_vlm_generation_receives_exact_compact_native_media_payload(): + generation = _CapturingAsyncVllmGeneration() + policy_media = PackedTensor(torch.ones(1, 3, 2, 2), dim_to_pack=0) + image = torch.ones(3, 2, 2) + audio = torch.ones(16) + video = torch.ones(2, 3, 2, 2) + message_log = [ + { + "role": "user", + "content": "", + "token_ids": torch.tensor([1]), + "pixel_values": policy_media, + } + ] + + asyncio.run( + async_generate_response_for_sample_turn( + generation, + message_log, + None, + _DummyTokenizer(), + max_seq_len=32, + sample_multimodal_data={ + "vllm_content": "