feat(grpo): deduplicate multimodal payloads

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

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:

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:

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

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

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

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,
         )
 
 

examples/run_vlm_grpo.py

@@ -146,6 +146,7 @@ def main() -> None:
         checkpointer,
         grpo_state,
         master_config,
+        processor=processor,
     )
 
 

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,

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],

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,