feat(data): add PackedTensor preprocessing modes

nemo_rl/data/processors.py

@@ -389,6 +389,7 @@ def vlm_preference_preprocessor(
     placeholder_style_processors = {
         "NemotronNanoVLV2Processor",
         "NemotronH_Nano_Omni_Reasoning_V3Processor",
+        "NemotronH_Omni_Reasoning_V3Processor",
     }
     message_processor = (
         _NemotronOmniPreferenceProcessorProxy(processor)
@@ -404,25 +405,33 @@ def _format_branch(completion: dict[str, Any]) -> VLMMessageLogType:
             task_data_spec,
         )
 
-        # Mirror the canonical Nemotron Omni metadata contract. Dynamic-resolution
-        # image batches may differ spatially across rows, while imgs_sizes
-        # preserves the true crop consumed by model-owned patchification.
+        # Mirror the canonical Nemotron Omni metadata. Record native image sizes
+        # before patchification removes the spatial dimensions.
         for raw_message in message_log:
             message = cast(Any, raw_message)
             pixel_values = message.get("pixel_values")
             if not isinstance(pixel_values, PackedTensor):
                 continue
-            pixel_values.pad_to_max_shape = True
-            pixels = pixel_values.as_tensor()
-            if pixels is not None and pixels.ndim == 4 and "imgs_sizes" not in message:
-                num_images, _, height, width = pixels.shape
+            if "imgs_sizes" not in message:
+                image_sizes: list[list[int]] = []
+                for pixels in pixel_values.iter_logical_segments():
+                    if pixels is None:
+                        continue
+                    if pixels.ndim != 4:
+                        raise ValueError(
+                            "Nemotron Omni pixel values must be [N, C, H, W] "
+                            f"before patchification, got {tuple(pixels.shape)}"
+                        )
+                    image_sizes.extend(
+                        [[int(pixels.shape[-2]), int(pixels.shape[-1])]]
+                        * int(pixels.shape[0])
+                    )
                 message["imgs_sizes"] = PackedTensor(
-                    torch.tensor(
-                        [[height, width]] * num_images,
-                        dtype=torch.long,
-                    ),
+                    torch.tensor(image_sizes, dtype=torch.long),
                     dim_to_pack=0,
                 )
+            pixel_values.preprocess_mode = "patchify"
+            pixel_values.preprocess_kwargs = {"patch_dim": 16}
             imgs_sizes = message.get("imgs_sizes")
             if isinstance(imgs_sizes, PackedTensor) and "num_frames" not in message:
                 sizes = imgs_sizes.as_tensor()

nemo_rl/data_plane/worker_mixin.py

@@ -116,7 +116,8 @@ def _broadcast_batched_data_dict(
                                 "empty_packed",
                                 len(v),
                                 v.dim_to_pack,
-                                v.pad_to_max_shape,
+                                v.preprocess_mode,
+                                v.preprocess_kwargs,
                             )
                         )
                         continue
@@ -130,7 +131,8 @@ def _broadcast_batched_data_dict(
                             str(values.device),
                             nested.offsets().tolist(),
                             shapes,
-                            v.pad_to_max_shape,
+                            v.preprocess_mode,
+                            v.preprocess_kwargs,
                         )
                     )
                 elif (
@@ -206,7 +208,14 @@ def _broadcast_batched_data_dict(
             ):
                 out[key] = tensor.to(src_device)
         elif kind == "packed_wire":
-            dtype_str, src_device, offsets, shapes, pad_to_max_shape = entry[2:]
+            (
+                dtype_str,
+                src_device,
+                offsets,
+                shapes,
+                preprocess_mode,
+                preprocess_kwargs,
+            ) = entry[2:]
             if is_leader:
                 flat = leader_flat[key].to(bcast_device)
             else:
@@ -224,17 +233,21 @@ def _broadcast_batched_data_dict(
                 if torch.device(src_device).type != torch.device(bcast_device).type:
                     nested = nested.to(src_device)
                 out[key] = PackedTensor.from_wire(
-                    nested, shapes, pad_to_max_shape=pad_to_max_shape
+                    nested,
+                    shapes,
+                    preprocess_mode=preprocess_mode,
+                    preprocess_kwargs=preprocess_kwargs,
                 )
         elif kind == "empty_packed":
             # Structural only: no payload, so followers rebuild from the
             # geometry and land on the leader's key set.
-            n_rows, dim_to_pack, pad_to_max_shape = entry[2:]
+            n_rows, dim_to_pack, preprocess_mode, preprocess_kwargs = entry[2:]
             if not is_leader:
                 out[key] = PackedTensor(
                     [None] * n_rows,
                     dim_to_pack,
-                    pad_to_max_shape=pad_to_max_shape,
+                    preprocess_mode=preprocess_mode,
+                    preprocess_kwargs=preprocess_kwargs,
                 )
         else:
             if not is_leader:

tests/unit/data/datasets/test_mmpr_tiny.py

@@ -241,7 +241,8 @@ def test_processor_produces_valid_datum_spec(self, tiny_image_path):
         assert result["task_name"] == "mmpr-tiny"
         user_message = result["message_log"][0]
         assert torch.equal(user_message["num_frames"].as_tensor(), torch.tensor([1]))
-        assert user_message["pixel_values"].pad_to_max_shape is True
+        assert user_message["pixel_values"].preprocess_mode == "patchify"
+        assert user_message["pixel_values"].preprocess_kwargs == {"patch_dim": 16}
         assert user_message["pixel_values"].as_tensor().dtype == torch.float32
 
     def test_text_only_row_preserves_formatted_vllm_content(self):