From e6bedc1dd1a074ee8fa1fc8fcb2aee70d2044300 Mon Sep 17 00:00:00 2001
From: rohitrango <rohit.rango@gmail.com>
Date: Thu, 10 Sep 2026 10:21:06 -0700
Subject: [PATCH 1/3] feat(data): add PackedTensor patchify preprocessing

Signed-off-by: rohitrango <rohit.rango@gmail.com>
---
 nemo_rl/data/multimodal_utils.py              | 193 ++++++++++++-----
 nemo_rl/data/processors.py                    |  31 ++-
 nemo_rl/data_plane/worker_mixin.py            |  25 ++-
 tests/unit/data/datasets/test_mmpr_tiny.py    |   3 +-
 tests/unit/data/test_multimodal_dict.py       | 200 ++++++++++++++----
 .../data/test_vlm_preference_processor.py     |   3 +-
 .../unit/data_plane/test_leader_broadcast.py  |   9 +-
 tests/unit/data_plane/test_local_sft.py       |   4 +-
 .../test_nemo_gym_image_placeholders.py       |  15 +-
 .../megatron/test_nemotron_omni_model.py      |   6 +-
 10 files changed, 367 insertions(+), 122 deletions(-)

diff --git a/nemo_rl/data/multimodal_utils.py b/nemo_rl/data/multimodal_utils.py
index 76f1c59a172..28d1a70176c 100644
--- a/nemo_rl/data/multimodal_utils.py
+++ b/nemo_rl/data/multimodal_utils.py
@@ -70,6 +70,7 @@
     {
         "NemotronNanoVLV2Processor",
         "NemotronH_Nano_Omni_Reasoning_V3Processor",
+        "NemotronH_Omni_Reasoning_V3Processor",
     }
 )
 
@@ -166,7 +167,8 @@ def row_shapes_key(field: str) -> str:
 # Keys inside the :func:`row_shapes_key` tag value. A plain dict rather than a
 # record because ``tags`` rides TQ's own serializer.
 ROW_GEOMETRY_SHAPES = "shapes"
-ROW_GEOMETRY_PAD = "pad"
+ROW_GEOMETRY_PREPROCESS_MODE = "preprocess_mode"
+ROW_GEOMETRY_PREPROCESS_KWARGS = "preprocess_kwargs"
 
 
 # Include-list of multimodal fields every forward-running dispatch (logprob
@@ -200,7 +202,7 @@ def multimodal_row_tags(
     contents.
 
     Carries ``shapes`` (per-row, and unrecoverable once ``to_wire`` flattens)
-    and ``pad`` (the field's policy flag). Deliberately *not* a pad target: the
+    and the field's preprocessing settings. Deliberately *not* a pad target: the
     width padding lands at is scratch that the model discards -- mcore crops it
     via ``imgs_sizes`` before patchification, and the AutoModel path rejects
     mixed-resolution batches outright -- so each consumer pads to its own view
@@ -224,11 +226,11 @@ def multimodal_row_tags(
                 "sample_ids, so a disagreement here would pair one sample's "
                 "pixels with another's shapes."
             )
-        pad = value.pad_to_max_shape
         for row, row_shapes in enumerate(shapes):
             tags[row][row_shapes_key(key)] = {
                 ROW_GEOMETRY_SHAPES: row_shapes,
-                ROW_GEOMETRY_PAD: pad,
+                ROW_GEOMETRY_PREPROCESS_MODE: value.preprocess_mode,
+                ROW_GEOMETRY_PREPROCESS_KWARGS: dict(value.preprocess_kwargs),
             }
     # ``None`` rather than ``B`` empty dicts: a text-only run has no packed
     # field, and an all-empty tags list would still be pickled on every
@@ -273,16 +275,80 @@ def reassemble_packed_multimodal(
                 + ". to_wire flattens each row, so without the companion the "
                 "true per-segment shapes are unrecoverable."
             )
-        # Indexed, not ``.get``-with-default: a producer-side rename of either
-        # key must fail here rather than silently restore ``pad=False``, which
-        # changes what ``as_tensor`` hands the vision encoder.
+        # Indexed, not ``.get``-with-default: a producer-side rename of any key
+        # must fail here rather than silently change what ``as_tensor`` hands
+        # the vision encoder.
         fields[key] = PackedTensor.from_wire(
             value,
             [[] if r is None else r[ROW_GEOMETRY_SHAPES] for r in rows],  # type: ignore[union-attr]
-            pad_to_max_shape=bool(present[0][ROW_GEOMETRY_PAD]),
+            preprocess_mode=present[0][ROW_GEOMETRY_PREPROCESS_MODE],
+            preprocess_kwargs=present[0][ROW_GEOMETRY_PREPROCESS_KWARGS],
         )
 
 
+def _patchify_segments(segments: list[torch.Tensor], *, patch_dim: int) -> torch.Tensor:
+    """Cut pixel segments into vision patches and pack them into one sequence.
+
+    Each ``[N, channels, H, W]`` segment is processed at its native resolution
+    into a ``[C_i, P²]`` block, where ``C_i`` is its spatial patch count and
+    ``P²`` is the flattened patch width (``channels * patch_dim**2``). Blocks
+    are packed along dimension zero, then a batch dimension is added to produce
+    ``[1, total_C, P²]``. Already-patchified segments in that final layout are
+    accepted so repeated materialization is safe.
+    """
+    if patch_dim <= 0:
+        raise ValueError(f"patch_dim must be positive, got {patch_dim}")
+
+    flattened: list[torch.Tensor] = []
+    for segment in segments:
+        if segment.ndim == 3:
+            if segment.shape[0] != 1:
+                raise ValueError(
+                    "Pre-patchified segments must be [1, total_C, P²], "
+                    f"got shape {tuple(segment.shape)}"
+                )
+            flattened.append(segment[0])
+            continue
+        if segment.ndim != 4:
+            raise ValueError(
+                "patchify expects [N, C, H, W] pixel segments or "
+                "[1, total_C, P²] pre-patchified segments, got shape "
+                f"{tuple(segment.shape)}"
+            )
+        count, channels, height, width = segment.shape
+        if height % patch_dim or width % patch_dim:
+            raise ValueError(
+                f"Image size {(height, width)} is not divisible by "
+                f"patch_dim={patch_dim}"
+            )
+        rows = height // patch_dim
+        columns = width // patch_dim
+        flattened.append(
+            segment.reshape(count, channels, rows, patch_dim, columns, patch_dim)
+            .permute(0, 2, 4, 1, 3, 5)
+            .reshape(count * rows * columns, channels * patch_dim * patch_dim)
+        )
+
+    widths = {tensor.shape[-1] for tensor in flattened}
+    if len(widths) != 1:
+        raise ValueError(
+            f"patchify produced mismatched P² widths {sorted(widths)}; "
+            "the segments do not share a channel count"
+        )
+    return torch.cat(flattened, dim=0).unsqueeze(0).contiguous()
+
+
+def _shared_preprocess_spec(
+    from_packed_tensors: list["PackedTensor"],
+) -> dict[str, Any]:
+    """Return the preprocessing setting shared by every input."""
+    first = from_packed_tensors[0]._preprocess_spec
+    assert all(
+        packed_tensor._preprocess_spec == first for packed_tensor in from_packed_tensors
+    ), "All packed tensors must have the same preprocess setting"
+    return first
+
+
 class PackedTensor:
     """A logical batch of rows backed by packable tensor segments.
 
@@ -326,7 +392,8 @@ def __init__(
         tensors: Union[torch.Tensor, list[Optional[torch.Tensor]], list[None]],
         dim_to_pack: int,
         *,
-        pad_to_max_shape: bool = False,
+        preprocess_mode: Optional[str] = None,
+        preprocess_kwargs: Optional[dict[str, Any]] = None,
         _row_offsets: Optional[list[int]] = None,
         _segment_indices: Optional[list[int]] = None,
         _segment_provenance: Optional[list[bytes]] = None,
@@ -337,8 +404,10 @@ def __init__(
             tensors: A tensor or list of per-item tensors. List entries may be
                 ``None`` for items without this modality.
             dim_to_pack: Dimension along which ``as_tensor`` concatenates.
-            pad_to_max_shape: Pad every non-packing dimension to its batch-wide
-                maximum before concatenating. All tensors must have the same rank.
+            preprocess_mode: Optional preprocessing applied by ``as_tensor``.
+                Supported values are ``pad_to_max_shape`` and ``patchify``.
+            preprocess_kwargs: Extra arguments for ``preprocess_mode``. Patchify
+                accepts ``patch_dim``.
         """
         assert tensors is not None, "Input tensors to PackedTensor cannot be None"
 
@@ -355,7 +424,13 @@ def __init__(
                 f"Unsupported type for input tensors to PackedTensor: {type(tensors)}"
             )
         self.dim_to_pack = dim_to_pack
-        self.pad_to_max_shape = pad_to_max_shape
+        if preprocess_mode not in (None, "pad_to_max_shape", "patchify"):
+            raise ValueError(
+                f"Unknown preprocess_mode {preprocess_mode!r}; expected None, "
+                "'pad_to_max_shape', or 'patchify'"
+            )
+        self.preprocess_mode = preprocess_mode
+        self.preprocess_kwargs: dict[str, Any] = dict(preprocess_kwargs or {})
         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"
@@ -396,6 +471,14 @@ def __setstate__(self, state: dict[str, Any]) -> None:
         self.__dict__.setdefault("_segment_indices", None)
         self.__dict__.setdefault("_segment_provenance", None)
 
+    @property
+    def _preprocess_spec(self) -> dict[str, Any]:
+        """Return keyword arguments that preserve preprocessing in a copy."""
+        return {
+            "preprocess_mode": self.preprocess_mode,
+            "preprocess_kwargs": self.preprocess_kwargs,
+        }
+
     @property
     def deduplication_enabled(self) -> bool:
         """Whether this value carries stable physical-segment provenance."""
@@ -446,7 +529,7 @@ def __deepcopy__(self, memo: dict[int, Any]) -> "PackedTensor":
             copied = PackedTensor(
                 [deepcopy(item, memo) for item in self.tensors],
                 self.dim_to_pack,
-                pad_to_max_shape=self.pad_to_max_shape,
+                **self._preprocess_spec,
             )
         else:
             copied = PackedTensor(
@@ -456,7 +539,7 @@ def __deepcopy__(self, memo: dict[int, Any]) -> "PackedTensor":
                     else [deepcopy(item, memo) for item in self.tensors]
                 ),
                 self.dim_to_pack,
-                pad_to_max_shape=self.pad_to_max_shape,
+                **self._preprocess_spec,
                 _row_offsets=(
                     list(self._row_offsets) if self._row_offsets is not None else None
                 ),
@@ -489,12 +572,21 @@ def as_tensor(
         if len(non_none_tensors) == 0:
             return None
 
+        if self.preprocess_mode == "patchify":
+            if self.dim_to_pack != 0:
+                raise ValueError(
+                    f"patchify requires dim_to_pack=0, got {self.dim_to_pack}"
+                )
+            return _patchify_segments(non_none_tensors, **self.preprocess_kwargs).to(
+                device
+            )
+
         # Some multimodal processors produce a different shape per prompt,
         # such as dynamic-resolution images, variable-frame videos, or audio
         # feature sequences. Concatenation already permits the packing
         # dimension to vary; when explicitly requested, pad every other
         # dimension to the largest size in the batch.
-        if self.pad_to_max_shape:
+        if self.preprocess_mode == "pad_to_max_shape":
             ranks = {tensor.ndim for tensor in non_none_tensors}
             if len(ranks) != 1:
                 raise ValueError(
@@ -600,7 +692,7 @@ def converted(item: Optional[torch.Tensor]) -> Optional[torch.Tensor]:
                 else list(self.tensors)
             ),
             self.dim_to_pack,
-            pad_to_max_shape=self.pad_to_max_shape,
+            **self._preprocess_spec,
             _row_offsets=(
                 list(self._row_offsets) if self._row_offsets is not None else None
             ),
@@ -627,7 +719,7 @@ def slice(self, indices: Union[list[int], torch.Tensor]) -> "PackedTensor":
             return PackedTensor(
                 tensors,
                 self.dim_to_pack,
-                pad_to_max_shape=self.pad_to_max_shape,
+                **self._preprocess_spec,
             )
 
         physical_remap: dict[int, int] = {}
@@ -651,7 +743,7 @@ def slice(self, indices: Union[list[int], torch.Tensor]) -> "PackedTensor":
         return PackedTensor(
             tensors,
             self.dim_to_pack,
-            pad_to_max_shape=self.pad_to_max_shape,
+            **self._preprocess_spec,
             _row_offsets=row_offsets,
             _segment_indices=segment_indices,
             _segment_provenance=(
@@ -673,7 +765,7 @@ def empty_rows_like(cls, other: "PackedTensor", num_rows: int) -> "PackedTensor"
             return cls(
                 [],
                 other.dim_to_pack,
-                pad_to_max_shape=other.pad_to_max_shape,
+                **other._preprocess_spec,
                 _row_offsets=[0] * (num_rows + 1),
                 _segment_indices=[],
                 _segment_provenance=[],
@@ -682,7 +774,7 @@ def empty_rows_like(cls, other: "PackedTensor", num_rows: int) -> "PackedTensor"
             return cls(
                 [],
                 other.dim_to_pack,
-                pad_to_max_shape=other.pad_to_max_shape,
+                **other._preprocess_spec,
                 _row_offsets=[0],
                 _segment_indices=[],
                 _segment_provenance=None,
@@ -690,7 +782,7 @@ def empty_rows_like(cls, other: "PackedTensor", num_rows: int) -> "PackedTensor"
         return cls(
             [None] * num_rows,
             other.dim_to_pack,
-            pad_to_max_shape=other.pad_to_max_shape,
+            **other._preprocess_spec,
         )
 
     @classmethod
@@ -719,10 +811,7 @@ def concat(cls, from_packed_tensors: list["PackedTensor"]) -> "PackedTensor":
         assert len(set(dim_to_packs)) == 1, (
             "All packed tensors must have the same dim_to_pack"
         )
-        pad_to_max_shapes = [batch.pad_to_max_shape for batch in from_packed_tensors]
-        assert len(set(pad_to_max_shapes)) == 1, (
-            "All packed tensors must have the same pad_to_max_shape setting"
-        )
+        preprocess_spec = _shared_preprocess_spec(from_packed_tensors)
         if any(
             packed_tensor.deduplication_enabled
             or packed_tensor._row_offsets is not None
@@ -763,7 +852,7 @@ def concat(cls, from_packed_tensors: list["PackedTensor"]) -> "PackedTensor":
             return cls(
                 tensors,
                 dim_to_packs[0],
-                pad_to_max_shape=pad_to_max_shapes[0],
+                **preprocess_spec,
                 _row_offsets=row_offsets,
                 _segment_indices=segment_indices,
                 _segment_provenance=provenances,
@@ -777,7 +866,7 @@ def concat(cls, from_packed_tensors: list["PackedTensor"]) -> "PackedTensor":
         return cls(
             tensors,
             dim_to_pack,
-            pad_to_max_shape=pad_to_max_shapes[0],
+            **preprocess_spec,
         )
 
     @classmethod
@@ -801,7 +890,7 @@ def merge_segments(
         return cls(
             concatenated.tensors,
             concatenated.dim_to_pack,
-            pad_to_max_shape=concatenated.pad_to_max_shape,
+            **concatenated._preprocess_spec,
             _row_offsets=[0, len(concatenated._segment_indices)],
             _segment_indices=concatenated._segment_indices,
             _segment_provenance=concatenated._segment_provenance,
@@ -837,10 +926,7 @@ def flattened_concat(
         assert len(set(dim_to_packs)) == 1, (
             "All packed tensors must have the same dim_to_pack"
         )
-        pad_to_max_shapes = [batch.pad_to_max_shape for batch in from_packed_tensors]
-        assert len(set(pad_to_max_shapes)) == 1, (
-            "All packed tensors must have the same pad_to_max_shape setting"
-        )
+        preprocess_spec = _shared_preprocess_spec(from_packed_tensors)
         if any(
             packed_tensor.deduplication_enabled
             or packed_tensor._row_offsets is not None
@@ -855,7 +941,7 @@ def flattened_concat(
         return cls(
             tensors,
             from_packed_tensors[0].dim_to_pack,
-            pad_to_max_shape=pad_to_max_shapes[0],
+            **preprocess_spec,
         )
 
     # ── Wire encoding (data-plane roundtrip) ─────────────────────────
@@ -949,12 +1035,12 @@ def to_wire(
         #     storage, and TQ never falls back to the deprecated strided layout.
         #   * The per-row concat is 1-D, so it cannot raise on segments whose
         #     trailing dims differ -- which is what previously forced
-        #     ``pad_to_max_shape`` to pad *before* the concat.
+        #     preprocessing to pad *before* the concat.
         #
         # The true shapes travel beside the payload (see the returned
         # ``shapes``) because TQ derives ``per_sample_shapes`` from what it is
         # handed: give it flat rows and it records flat lengths. Padding still
-        # happens for ``pad_to_max_shape`` values, but in worker memory at use
+        # happens for values that need it, but in worker memory at use
         # time via :meth:`as_tensor`, not on the wire.
         shapes = self._shapes_of(row_segments)
         # ``reshape(-1)`` on contiguous processor output is a view, so the
@@ -991,7 +1077,8 @@ def from_wire(
         nested: torch.Tensor,
         shapes: list[list[list[int]]],
         *,
-        pad_to_max_shape: bool = False,
+        preprocess_mode: Optional[str] = None,
+        preprocess_kwargs: Optional[dict[str, Any]] = None,
     ) -> Optional["PackedTensor"]:
         """Reconstruct from the value produced by :meth:`to_wire`.
 
@@ -1010,10 +1097,9 @@ def from_wire(
         reconstructs as legacy does: ``as_tensor`` returns ``None`` and
         ``logical_segment_counts_by_row`` reports 0 rather than 1.
 
-        ``pad_to_max_shape`` is restored onto the rebuilt value as a flag, not
+        The preprocessing settings are restored onto the rebuilt value, not
         materialized. Segments come back at their true shapes and stay separate
-        via the CSR row map, so nothing is padded or concatenated here;
-        :meth:`as_tensor` pads at use time.
+        via the CSR row map; :meth:`as_tensor` preprocesses at use time.
 
         Mirrors :meth:`to_wire`; both assume ``dim_to_pack=0``.
         """
@@ -1054,7 +1140,8 @@ def from_wire(
         return cls(
             segments_flat,  # type: ignore[arg-type]
             dim_to_pack=0,
-            pad_to_max_shape=pad_to_max_shape,
+            preprocess_mode=preprocess_mode,
+            preprocess_kwargs=preprocess_kwargs,
             _row_offsets=row_offsets,
             _segment_indices=list(range(len(segments_flat))),
         )
@@ -1072,7 +1159,7 @@ def encode_multimodal_for_wire(
     Payload only. Per-token fields ride rectangular; packed fields ride as one
     flattened ``torch.jagged`` value. The geometry :meth:`PackedTensor.from_wire`
     needs to undo that flattening -- per-row segment shapes plus the
-    ``pad_to_max_shape`` flag -- is minted separately by
+    preprocessing settings -- is minted separately by
     :func:`multimodal_row_tags` and shipped on ``KVBatchMeta.tags``. TQ cannot
     derive it, because it reads ``per_sample_shapes`` off the flattened rows it
     is handed.
@@ -1182,9 +1269,14 @@ 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 get_preprocess(processor: Any, key: str) -> dict[str, Any]:
+    """Return materialization preprocessing for one processor input."""
+    if uses_image_placeholder(processor) and key == "pixel_values":
+        return {
+            "preprocess_mode": "patchify",
+            "preprocess_kwargs": {"patch_dim": 16},
+        }
+    return {"preprocess_mode": None, "preprocess_kwargs": {}}
 
 
 def extract_multimodal_model_inputs(
@@ -1246,7 +1338,7 @@ def extract_multimodal_model_inputs(
         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),
+            **get_preprocess(processor, key),
         )
 
     for key in ("token_type_ids", "mm_token_type_ids"):
@@ -1386,13 +1478,13 @@ def media_sources_equal(
 def _materialize_ragged_pixel_values(
     processed: dict[str, Any], processor: Any
 ) -> dict[str, Any]:
-    """Fold a ragged per-image ``pixel_values`` list into one padded tensor.
+    """Fold a ragged per-image ``pixel_values`` list into one patch sequence.
 
     Processors with dynamic per-image resolution return a list of CHW tensors
     rather than a stacked batch. ``imgs_sizes`` is derived from the *unpadded*
     shapes first, since those exact sizes are what the projector slices with;
-    padding happens afterwards so downstream sees the single tensor its
-    torch.Tensor contract expects.
+    patchification happens afterwards so downstream sees the single tensor its
+    ``torch.Tensor`` contract expects.
     """
     processed = dict(processed)
     pixel_values = processed.get("pixel_values")
@@ -1418,7 +1510,7 @@ def _materialize_ragged_pixel_values(
 def _stack_ragged_pixel_values(
     processed: dict[str, Any], tiles: list[torch.Tensor], processor: Any
 ) -> None:
-    """Derive imgs_sizes from unpadded shapes, then pad into one tensor."""
+    """Derive image sizes, then patchify native-shape tiles into one tensor."""
     if uses_image_placeholder(processor) and "imgs_sizes" not in processed:
         processed["imgs_sizes"] = torch.tensor(
             [[int(item.shape[-2]), int(item.shape[-1])] for item in tiles],
@@ -1427,7 +1519,8 @@ def _stack_ragged_pixel_values(
     stacked = PackedTensor(
         [item.unsqueeze(0) for item in tiles],
         dim_to_pack=0,
-        pad_to_max_shape=True,
+        preprocess_mode="patchify",
+        preprocess_kwargs={"patch_dim": 16},
     ).as_tensor()
     assert stacked is not None
     processed["pixel_values"] = stacked
diff --git a/nemo_rl/data/processors.py b/nemo_rl/data/processors.py
index cf5ed322e69..73f1849c162 100644
--- a/nemo_rl/data/processors.py
+++ b/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()
diff --git a/nemo_rl/data_plane/worker_mixin.py b/nemo_rl/data_plane/worker_mixin.py
index 581f371e1bf..5f99349786b 100644
--- a/nemo_rl/data_plane/worker_mixin.py
+++ b/nemo_rl/data_plane/worker_mixin.py
@@ -117,7 +117,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
@@ -131,7 +132,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 (
@@ -207,7 +209,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:
@@ -225,17 +234,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:
diff --git a/tests/unit/data/datasets/test_mmpr_tiny.py b/tests/unit/data/datasets/test_mmpr_tiny.py
index 257227ab380..e304fbaecb1 100644
--- a/tests/unit/data/datasets/test_mmpr_tiny.py
+++ b/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):
diff --git a/tests/unit/data/test_multimodal_dict.py b/tests/unit/data/test_multimodal_dict.py
index c88883316a2..cf7f71e2d76 100644
--- a/tests/unit/data/test_multimodal_dict.py
+++ b/tests/unit/data/test_multimodal_dict.py
@@ -23,8 +23,10 @@
     PER_TOKEN_MULTIMODAL_FIELDS,
     PackedTensor,
     encode_multimodal_for_wire,
+    get_preprocess,
     multimodal_row_tags,
     reassemble_packed_multimodal,
+    uses_image_placeholder,
 )
 from nemo_rl.distributed.batched_data_dict import (
     BatchedDataDict,
@@ -55,6 +57,28 @@ def test_packed_data_basic():
     assert torch.equal(batch.as_tensor(), expected_tensor)
 
 
+@pytest.mark.parametrize(
+    "processor_name",
+    [
+        "NemotronNanoVLV2Processor",
+        "NemotronH_Nano_Omni_Reasoning_V3Processor",
+        "NemotronH_Omni_Reasoning_V3Processor",
+    ],
+)
+def test_placeholder_processors_use_patchify(processor_name):
+    processor = type(processor_name, (), {})()
+
+    assert uses_image_placeholder(processor)
+    assert get_preprocess(processor, "pixel_values") == {
+        "preprocess_mode": "patchify",
+        "preprocess_kwargs": {"patch_dim": 16},
+    }
+    assert get_preprocess(processor, "imgs_sizes") == {
+        "preprocess_mode": None,
+        "preprocess_kwargs": {},
+    }
+
+
 def test_shard_by_batch_size_with_packed_data():
     """Test shard_by_batch_size with packed multimodal data."""
     # Create sample data
@@ -387,7 +411,7 @@ def test_packedtensor_pads_mixed_dynamic_resolution_images():
     second = 2 * torch.ones(1, 3, 4, 2)
 
     packed = PackedTensor(
-        [first, second], dim_to_pack=0, pad_to_max_shape=True
+        [first, second], dim_to_pack=0, preprocess_mode="pad_to_max_shape"
     ).as_tensor()
 
     assert packed.shape == (2, 3, 4, 4)
@@ -412,7 +436,7 @@ def test_dynamic_resolution_padding_is_cropped_before_radio_patchification():
     padded = PackedTensor(
         [small, large],
         dim_to_pack=0,
-        pad_to_max_shape=True,
+        preprocess_mode="pad_to_max_shape",
     ).as_tensor()
     # Use nonzero garbage so this test cannot pass merely because F.pad uses zero.
     padded[0, :, 32:, :] = 123
@@ -461,7 +485,7 @@ def test_packedtensor_pad_to_max_shape_supports_audio_and_video(
     second = 2 * torch.ones(second_shape)
 
     packed = PackedTensor(
-        [first, second], dim_to_pack=0, pad_to_max_shape=True
+        [first, second], dim_to_pack=0, preprocess_mode="pad_to_max_shape"
     ).as_tensor()
 
     assert packed.shape == expected_shape
@@ -476,7 +500,7 @@ def test_pad_to_max_shape_rejects_mismatched_ranks():
         PackedTensor(
             [torch.ones(1, 3, 4), torch.ones(1, 3)],
             dim_to_pack=0,
-            pad_to_max_shape=True,
+            preprocess_mode="pad_to_max_shape",
         ).as_tensor()
 
 
@@ -485,7 +509,7 @@ def test_pad_to_max_shape_rejects_out_of_range_dim():
         PackedTensor(
             [torch.ones(1, 3, 4), torch.ones(2, 3, 4)],
             dim_to_pack=3,
-            pad_to_max_shape=True,
+            preprocess_mode="pad_to_max_shape",
         ).as_tensor()
 
 
@@ -493,25 +517,128 @@ def test_pad_to_max_shape_supports_negative_pack_dim():
     packed = PackedTensor(
         [torch.ones(2, 3, 1), 2 * torch.ones(4, 3, 1)],
         dim_to_pack=-3,
-        pad_to_max_shape=True,
+        preprocess_mode="pad_to_max_shape",
     ).as_tensor()
 
     assert packed.shape == (6, 3, 1)
 
 
-def test_slice_preserves_pad_to_max_shape_flag():
+def test_slice_preserves_preprocess_spec():
     packed = PackedTensor(
         [torch.ones(1, 3, 2, 4), 2 * torch.ones(1, 3, 4, 2)],
         dim_to_pack=0,
-        pad_to_max_shape=True,
+        preprocess_mode="pad_to_max_shape",
+        preprocess_kwargs={},
     )
 
     sliced = packed.slice([0, 1])
 
-    assert sliced.pad_to_max_shape is True
+    assert sliced.preprocess_mode == "pad_to_max_shape"
+    assert sliced.preprocess_kwargs == {}
     assert sliced.as_tensor().shape == (2, 3, 4, 4)
 
 
+def test_packedtensor_rejects_unknown_preprocess_mode():
+    with pytest.raises(ValueError, match="Unknown preprocess_mode"):
+        PackedTensor(
+            torch.ones(1, 3, 4, 4),
+            dim_to_pack=0,
+            preprocess_mode="jagged",
+        )
+
+
+def test_patchify_packs_mixed_resolutions_without_padding():
+    packed = PackedTensor(
+        [torch.ones(1, 3, 32, 32), 2 * torch.ones(1, 3, 64, 32)],
+        dim_to_pack=0,
+        preprocess_mode="patchify",
+        preprocess_kwargs={"patch_dim": 16},
+    ).as_tensor()
+
+    assert packed.shape == (1, 12, 768)
+    assert torch.all(packed[0, :4] == 1)
+    assert torch.all(packed[0, 4:] == 2)
+
+
+def test_patchify_preserves_pixel_order_within_a_patch():
+    image = torch.arange(3 * 2 * 2, dtype=torch.float32).reshape(1, 3, 2, 2)
+
+    packed = PackedTensor(
+        [image],
+        dim_to_pack=0,
+        preprocess_mode="patchify",
+        preprocess_kwargs={"patch_dim": 2},
+    ).as_tensor()
+
+    assert packed.shape == (1, 1, 12)
+    torch.testing.assert_close(packed[0, 0], image.reshape(12))
+
+
+def test_patchify_accepts_already_patchified_segments():
+    raw = PackedTensor(
+        [torch.ones(1, 3, 32, 32)],
+        dim_to_pack=0,
+        preprocess_mode="patchify",
+        preprocess_kwargs={"patch_dim": 16},
+    )
+
+    once = raw.as_tensor()
+    assert once is not None
+    twice = PackedTensor(
+        [once],
+        dim_to_pack=0,
+        preprocess_mode="patchify",
+        preprocess_kwargs={"patch_dim": 16},
+    ).as_tensor()
+
+    torch.testing.assert_close(once, twice)
+
+
+def test_patchify_survives_flattened_concat():
+    first = PackedTensor(
+        [torch.ones(1, 3, 32, 32)],
+        dim_to_pack=0,
+        preprocess_mode="patchify",
+        preprocess_kwargs={"patch_dim": 16},
+    )
+    second = PackedTensor(
+        [2 * torch.ones(1, 3, 64, 32)],
+        dim_to_pack=0,
+        preprocess_mode="patchify",
+        preprocess_kwargs={"patch_dim": 16},
+    )
+
+    flattened = PackedTensor.flattened_concat([first, second])
+
+    assert len(flattened) == 2
+    assert flattened.as_tensor().shape == (1, 12, 768)
+    torch.testing.assert_close(
+        flattened.as_tensor(), PackedTensor.concat([first, second]).as_tensor()
+    )
+
+
+def test_patchify_rejects_indivisible_image_size():
+    with pytest.raises(ValueError, match="not divisible by patch_dim=16"):
+        PackedTensor(
+            [torch.ones(1, 3, 30, 32)],
+            dim_to_pack=0,
+            preprocess_mode="patchify",
+            preprocess_kwargs={"patch_dim": 16},
+        ).as_tensor()
+
+
+def test_concat_rejects_mixed_preprocess_settings():
+    padded = PackedTensor(
+        torch.ones(1, 3, 4, 4),
+        dim_to_pack=0,
+        preprocess_mode="pad_to_max_shape",
+    )
+    plain = PackedTensor(torch.ones(1, 3, 4, 4), dim_to_pack=0)
+
+    with pytest.raises(AssertionError, match="same preprocess setting"):
+        PackedTensor.concat([padded, plain])
+
+
 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)
@@ -559,12 +686,12 @@ def test_packedtensor_dedup_expands_before_dynamic_shape_padding():
     first = PackedTensor(
         torch.ones(1, 1, 2),
         dim_to_pack=0,
-        pad_to_max_shape=True,
+        preprocess_mode="pad_to_max_shape",
     ).enable_deduplication()
     second = PackedTensor(
         2 * torch.ones(1, 2, 1),
         dim_to_pack=0,
-        pad_to_max_shape=True,
+        preprocess_mode="pad_to_max_shape",
     ).enable_deduplication()
 
     packed = PackedTensor.concat([first, deepcopy(first), second])
@@ -632,25 +759,6 @@ def test_packedtensor_compact_dim_one_slice_empty_and_cloudpickle_roundtrip():
     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 sum(restored.logical_segment_counts_by_row()) == 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)
@@ -708,7 +816,7 @@ def test_to_wire_does_not_pad_segments_before_concat_under_dedup():
     packed = PackedTensor(
         [torch.ones(1, 3, 2, 4), 2 * torch.ones(1, 3, 4, 2)],
         dim_to_pack=0,
-        pad_to_max_shape=True,
+        preprocess_mode="pad_to_max_shape",
         _row_offsets=[0, 2],
         _segment_indices=[0, 1],
     )
@@ -721,7 +829,9 @@ def test_to_wire_does_not_pad_segments_before_concat_under_dedup():
     assert [t.numel() for t in nested.unbind()] == [48]
     assert shapes == [[[1, 3, 2, 4], [1, 3, 4, 2]]]
 
-    restored = PackedTensor.from_wire(nested, shapes, pad_to_max_shape=True).as_tensor()
+    restored = PackedTensor.from_wire(
+        nested, shapes, preprocess_mode="pad_to_max_shape"
+    ).as_tensor()
     assert torch.equal(restored, expected)
 
 
@@ -758,7 +868,9 @@ def test_to_wire_does_not_materialize_pad_to_max_shape():
     # Same rank, different trailing dims — nemotron-omni style tiles.
     first = torch.ones(1, 3, 2, 4)
     second = 2 * torch.ones(2, 3, 4, 2)
-    packed = PackedTensor([first, second], dim_to_pack=0, pad_to_max_shape=True)
+    packed = PackedTensor(
+        [first, second], dim_to_pack=0, preprocess_mode="pad_to_max_shape"
+    )
 
     nested, shapes = packed.to_wire()
     rows = list(nested.unbind())
@@ -769,7 +881,9 @@ def test_to_wire_does_not_materialize_pad_to_max_shape():
 
     # Padding is reapplied on read, reproducing the pre-wire as_tensor().
     assert torch.equal(
-        PackedTensor.from_wire(nested, shapes, pad_to_max_shape=True).as_tensor(),
+        PackedTensor.from_wire(
+            nested, shapes, preprocess_mode="pad_to_max_shape"
+        ).as_tensor(),
         packed.as_tensor(),
     )
 
@@ -836,7 +950,8 @@ def test_encode_multimodal_for_wire_packed_emits_single_nested_entry():
         [[3, 4]],
         [[1, 4]],
     ]
-    assert tags[0]["pixel_values__row_shapes"]["pad"] is False
+    assert tags[0]["pixel_values__row_shapes"]["preprocess_mode"] is None
+    assert tags[0]["pixel_values__row_shapes"]["preprocess_kwargs"] == {}
 
 
 def test_multimodal_row_tags_does_not_encode_the_payload():
@@ -882,13 +997,20 @@ def test_reassemble_packed_multimodal_raises_without_companion():
 
 
 def test_reassemble_packed_multimodal_round_trips_with_companion():
-    packed = PackedTensor([torch.ones(3, 4), torch.ones(1, 4)], dim_to_pack=0)
+    packed = PackedTensor(
+        [torch.ones(1, 3, 32, 32), torch.ones(1, 3, 16, 32)],
+        dim_to_pack=0,
+        preprocess_mode="patchify",
+        preprocess_kwargs={"patch_dim": 16},
+    )
     nested, _ = packed.to_wire()
     tags = multimodal_row_tags({"pixel_values": packed}, len(packed))
 
     fields = {"pixel_values": nested}
     reassemble_packed_multimodal(fields, tags)
 
+    assert fields["pixel_values"].preprocess_mode == "patchify"
+    assert fields["pixel_values"].preprocess_kwargs == {"patch_dim": 16}
     assert torch.equal(fields["pixel_values"].as_tensor(), packed.as_tensor())
 
 
@@ -1034,11 +1156,13 @@ def test_to_wire_carries_mixed_rank_rows():
     load-bearing. Reshaping on read restores the original ranks.
     """
     rows = [torch.ones(1, 3, 2), torch.ones(2, 3)]
-    packed = PackedTensor(list(rows), dim_to_pack=0, pad_to_max_shape=True)
+    packed = PackedTensor(list(rows), dim_to_pack=0, preprocess_mode="pad_to_max_shape")
 
     nested, shapes = packed.to_wire()
     assert [t.numel() for t in nested.unbind()] == [6, 6]
     assert shapes == [[[1, 3, 2]], [[2, 3]]]
 
-    restored = PackedTensor.from_wire(nested, shapes, pad_to_max_shape=True)
+    restored = PackedTensor.from_wire(
+        nested, shapes, preprocess_mode="pad_to_max_shape"
+    )
     assert [tuple(t.shape) for t in restored.tensors] == [(1, 3, 2), (2, 3)]
diff --git a/tests/unit/data/test_vlm_preference_processor.py b/tests/unit/data/test_vlm_preference_processor.py
index 308212df768..3f59df888f4 100644
--- a/tests/unit/data/test_vlm_preference_processor.py
+++ b/tests/unit/data/test_vlm_preference_processor.py
@@ -101,6 +101,7 @@ def test_vlm_preference_processor_adds_nemotron_omni_media_metadata():
         result["message_log_rejected"],
     ):
         message = message_log[0]
-        assert message["pixel_values"].pad_to_max_shape
+        assert message["pixel_values"].preprocess_mode == "patchify"
+        assert message["pixel_values"].preprocess_kwargs == {"patch_dim": 16}
         assert message["imgs_sizes"].as_tensor().tolist() == [[15, 23]]
         assert message["num_frames"].as_tensor().tolist() == [1]
diff --git a/tests/unit/data_plane/test_leader_broadcast.py b/tests/unit/data_plane/test_leader_broadcast.py
index 4523ae93356..c48ecb9bee6 100644
--- a/tests/unit/data_plane/test_leader_broadcast.py
+++ b/tests/unit/data_plane/test_leader_broadcast.py
@@ -91,7 +91,7 @@ def _packed(rows):
     return PackedTensor(
         [r.clone() if r is not None else None for r in rows],
         dim_to_pack=0,
-        pad_to_max_shape=True,
+        preprocess_mode="pad_to_max_shape",
     )
 
 
@@ -148,7 +148,9 @@ def _all_empty_body(rank: int):
             {
                 "input_ids": torch.arange(8, dtype=torch.long).reshape(2, 4),
                 "pixel_values": PackedTensor(
-                    [None, None], dim_to_pack=0, pad_to_max_shape=True
+                    [None, None],
+                    dim_to_pack=0,
+                    preprocess_mode="pad_to_max_shape",
                 ),
             }
         )
@@ -165,7 +167,8 @@ def _all_empty_body(rank: int):
     assert isinstance(packed, PackedTensor), type(packed).__name__
     assert packed.logical_segment_counts_by_row() == [0, 0]
     assert packed.as_tensor() is None
-    assert packed.pad_to_max_shape is True
+    assert packed.preprocess_mode == "pad_to_max_shape"
+    assert packed.preprocess_kwargs == {}
 
 
 def _unsupported_type_body(rank: int):
diff --git a/tests/unit/data_plane/test_local_sft.py b/tests/unit/data_plane/test_local_sft.py
index 8dd2e426739..d5937c4645b 100644
--- a/tests/unit/data_plane/test_local_sft.py
+++ b/tests/unit/data_plane/test_local_sft.py
@@ -47,7 +47,7 @@ def _put_multimodal_batch(
     pixels = PackedTensor(
         [torch.full((1, 2), 1.0), torch.full((2, 2), 2.0)],
         dim_to_pack=0,
-        pad_to_max_shape=True,
+        preprocess_mode="pad_to_max_shape",
     ).enable_deduplication()
     fields = local_batch_to_tensordict(
         {
@@ -80,7 +80,7 @@ def test_local_round_trip_preserves_tensor_and_packed_tensor_fields() -> None:
     pixels = batch["pixel_values"]
     assert isinstance(pixels, PackedTensor)
     assert pixels.dim_to_pack == 0
-    assert pixels.pad_to_max_shape
+    assert pixels.preprocess_mode == "pad_to_max_shape"
     assert pixels.deduplication_enabled
     assert pixels.logical_segment_counts_by_row() == [1, 1]
     assert torch.equal(pixels.tensors[0], torch.full((1, 2), 1.0))
diff --git a/tests/unit/environments/test_nemo_gym_image_placeholders.py b/tests/unit/environments/test_nemo_gym_image_placeholders.py
index 5a05ddf366d..a5bf369acb5 100644
--- a/tests/unit/environments/test_nemo_gym_image_placeholders.py
+++ b/tests/unit/environments/test_nemo_gym_image_placeholders.py
@@ -60,7 +60,9 @@ def __call__(self, *, text, images, return_tensors):
 def _ragged(*shapes: tuple[int, ...]) -> NemotronNanoVLV2Processor:
     return NemotronNanoVLV2Processor(
         [torch.ones(*shape) for shape in shapes],
-        imgs_sizes=torch.tensor([[4, 4]] * len(shapes), dtype=torch.long),
+        imgs_sizes=torch.tensor(
+            [[shape[-2], shape[-1]] for shape in shapes], dtype=torch.long
+        ),
     )
 
 
@@ -83,9 +85,9 @@ def test_ragged_output_requested_only_for_multi_image_turns():
         )
 
 
-def test_ragged_pixel_values_are_padded_to_one_tensor():
-    """Heterogeneous CHW tensors become a single padded tensor for the message."""
-    processor = _ragged((3, 2, 4), (3, 6, 4))
+def test_ragged_pixel_values_are_patchified_to_one_tensor():
+    """Heterogeneous CHW tensors become one packed patch sequence."""
+    processor = _ragged((3, 16, 32), (3, 32, 16))
     user_message: dict = {}
     attach_image_model_inputs_to_message(
         user_message,
@@ -94,10 +96,7 @@ def test_ragged_pixel_values_are_padded_to_one_tensor():
         pad_dynamic_image_shapes=True,
     )
     packed = user_message["pixel_values"].as_tensor()
-    # Two images, padded up to the tallest, channels preserved.
-    assert packed.shape[0] == 2
-    assert packed.shape[-3] == 3
-    assert packed.shape[-2] == 6
+    assert packed.shape == (1, 4, 768)
 
 
 def test_ragged_pixel_values_reject_non_chw_entries():
diff --git a/tests/unit/models/megatron/test_nemotron_omni_model.py b/tests/unit/models/megatron/test_nemotron_omni_model.py
index 67990577c87..8ee53222459 100644
--- a/tests/unit/models/megatron/test_nemotron_omni_model.py
+++ b/tests/unit/models/megatron/test_nemotron_omni_model.py
@@ -246,7 +246,8 @@ def _deduplicated_expanded_fixture(device: torch.device):
             "pixel_values": PackedTensor(
                 [image.clone(), image.clone()],
                 dim_to_pack=0,
-                pad_to_max_shape=True,
+                preprocess_mode="patchify",
+                preprocess_kwargs={"patch_dim": 16},
             ),
             "imgs_sizes": PackedTensor(
                 [image_size.clone(), image_size.clone()],
@@ -257,7 +258,8 @@ def _deduplicated_expanded_fixture(device: torch.device):
     pixel_row = PackedTensor(
         image,
         dim_to_pack=0,
-        pad_to_max_shape=True,
+        preprocess_mode="patchify",
+        preprocess_kwargs={"patch_dim": 16},
     ).enable_deduplication()
     image_size_row = PackedTensor(
         image_size,

From 3597f4e29637adb3c3d8de2f38da392eba3b04a4 Mon Sep 17 00:00:00 2001
From: rohitrango <rohit.rango@gmail.com>
Date: Thu, 17 Sep 2026 10:29:35 -0700
Subject: [PATCH 2/3] fix(multimodal): materialize pixels per backend

Signed-off-by: rohitrango <rohit.rango@gmail.com>
---
 nemo_rl/data/multimodal_utils.py              | 57 ++++++++++++-------
 nemo_rl/data/processors.py                    | 22 +++----
 nemo_rl/distributed/batched_data_dict.py      |  4 +-
 nemo_rl/models/automodel/data.py              |  4 +-
 .../policy/workers/dtensor_policy_worker.py   |  6 +-
 5 files changed, 56 insertions(+), 37 deletions(-)

diff --git a/nemo_rl/data/multimodal_utils.py b/nemo_rl/data/multimodal_utils.py
index 28d1a70176c..e684b21b369 100644
--- a/nemo_rl/data/multimodal_utils.py
+++ b/nemo_rl/data/multimodal_utils.py
@@ -299,12 +299,13 @@ def _patchify_segments(segments: list[torch.Tensor], *, patch_dim: int) -> torch
     if patch_dim <= 0:
         raise ValueError(f"patch_dim must be positive, got {patch_dim}")
 
+    patch_features = 3 * patch_dim**2
     flattened: list[torch.Tensor] = []
     for segment in segments:
         if segment.ndim == 3:
-            if segment.shape[0] != 1:
+            if segment.shape[0] != 1 or segment.shape[-1] != patch_features:
                 raise ValueError(
-                    "Pre-patchified segments must be [1, total_C, P²], "
+                    f"Pre-patchified segments must be [1, total_C, P²] with P²={patch_features}, "
                     f"got shape {tuple(segment.shape)}"
                 )
             flattened.append(segment[0])
@@ -429,6 +430,10 @@ def __init__(
                 f"Unknown preprocess_mode {preprocess_mode!r}; expected None, "
                 "'pad_to_max_shape', or 'patchify'"
             )
+        if preprocess_mode == "patchify" and not (preprocess_kwargs or {}).get(
+            "patch_dim"
+        ):
+            raise ValueError("patchify requires patch_dim")
         self.preprocess_mode = preprocess_mode
         self.preprocess_kwargs: dict[str, Any] = dict(preprocess_kwargs or {})
         if (_row_offsets is None) != (_segment_indices is None):
@@ -558,8 +563,9 @@ def __deepcopy__(self, memo: dict[int, Any]) -> "PackedTensor":
         return copied
 
     def as_tensor(
-        self, device: Optional[torch.device] = None
+        self, device: torch.device | None = None, mode: str | None = None
     ) -> Optional[torch.Tensor]:
+        mode = mode or self.preprocess_mode
         if device is not None:
             # Move only non-None tensors to device, preserve Nones
             for i, item in enumerate(self.tensors):
@@ -572,7 +578,7 @@ def as_tensor(
         if len(non_none_tensors) == 0:
             return None
 
-        if self.preprocess_mode == "patchify":
+        if mode == "patchify":
             if self.dim_to_pack != 0:
                 raise ValueError(
                     f"patchify requires dim_to_pack=0, got {self.dim_to_pack}"
@@ -586,7 +592,7 @@ def as_tensor(
         # feature sequences. Concatenation already permits the packing
         # dimension to vary; when explicitly requested, pad every other
         # dimension to the largest size in the batch.
-        if self.preprocess_mode == "pad_to_max_shape":
+        if mode == "pad_to_max_shape":
             ranks = {tensor.ndim for tensor in non_none_tensors}
             if len(ranks) != 1:
                 raise ValueError(
@@ -1272,9 +1278,11 @@ def get_dim_to_pack_along(processor, key: str) -> int:
 def get_preprocess(processor: Any, key: str) -> dict[str, Any]:
     """Return materialization preprocessing for one processor input."""
     if uses_image_placeholder(processor) and key == "pixel_values":
+        image_processor = getattr(processor, "image_processor", processor)
+        patch_dim = getattr(image_processor, "patch_size", 16)
         return {
             "preprocess_mode": "patchify",
-            "preprocess_kwargs": {"patch_dim": 16},
+            "preprocess_kwargs": {"patch_dim": patch_dim},
         }
     return {"preprocess_mode": None, "preprocess_kwargs": {}}
 
@@ -1287,6 +1295,7 @@ def extract_multimodal_model_inputs(
     if (
         uses_image_placeholder(processor)
         and "pixel_values" in processed
+        and isinstance(processed["pixel_values"], torch.Tensor)
         and "imgs_sizes" not in processed
         and processed["pixel_values"].ndim == 4
     ):
@@ -1301,6 +1310,21 @@ def extract_multimodal_model_inputs(
             len(processed["imgs_sizes"]),
             dtype=torch.long,
         )
+    sizes = processed.get("imgs_sizes")
+    if uses_image_placeholder(processor) and isinstance(sizes, torch.Tensor):
+        pixels = processed["pixel_values"]
+        segments = pixels if isinstance(pixels, list) else [pixels]
+        preprocess = get_preprocess(processor, "pixel_values")
+        patch_dim = preprocess["preprocess_kwargs"]["patch_dim"]
+        pixel_count = sum(
+            segment.shape[1]
+            if segment.ndim == 3
+            else segment.numel() // (3 * patch_dim**2)
+            for segment in segments
+        )
+        size_count = int(torch.prod(sizes // patch_dim, dim=1).sum())
+        if pixel_count != size_count:
+            raise ValueError("pixel_values and imgs_sizes have different patch counts")
 
     input_ids = processed.get("input_ids")
     if input_ids is None:
@@ -1328,7 +1352,7 @@ def extract_multimodal_model_inputs(
         if key not in processed:
             continue
         value = processed[key]
-        if not isinstance(value, torch.Tensor):
+        if not isinstance(value, (torch.Tensor, list)):
             raise ValueError(
                 f"Processor model input {key!r} must be a torch.Tensor, got "
                 f"{type(value).__name__}."
@@ -1478,13 +1502,9 @@ def media_sources_equal(
 def _materialize_ragged_pixel_values(
     processed: dict[str, Any], processor: Any
 ) -> dict[str, Any]:
-    """Fold a ragged per-image ``pixel_values`` list into one patch sequence.
+    """Preserve a ragged per-image ``pixel_values`` list for materialization.
 
-    Processors with dynamic per-image resolution return a list of CHW tensors
-    rather than a stacked batch. ``imgs_sizes`` is derived from the *unpadded*
-    shapes first, since those exact sizes are what the projector slices with;
-    patchification happens afterwards so downstream sees the single tensor its
-    ``torch.Tensor`` contract expects.
+    Tiles remain separate until ``PackedTensor.as_tensor`` materializes them.
     """
     processed = dict(processed)
     pixel_values = processed.get("pixel_values")
@@ -1510,20 +1530,13 @@ def _materialize_ragged_pixel_values(
 def _stack_ragged_pixel_values(
     processed: dict[str, Any], tiles: list[torch.Tensor], processor: Any
 ) -> None:
-    """Derive image sizes, then patchify native-shape tiles into one tensor."""
+    """Derive image sizes and preserve native-shape tiles for patchification."""
     if uses_image_placeholder(processor) and "imgs_sizes" not in processed:
         processed["imgs_sizes"] = torch.tensor(
             [[int(item.shape[-2]), int(item.shape[-1])] for item in tiles],
             dtype=torch.long,
         )
-    stacked = PackedTensor(
-        [item.unsqueeze(0) for item in tiles],
-        dim_to_pack=0,
-        preprocess_mode="patchify",
-        preprocess_kwargs={"patch_dim": 16},
-    ).as_tensor()
-    assert stacked is not None
-    processed["pixel_values"] = stacked
+    processed["pixel_values"] = [item.unsqueeze(0) for item in tiles]
 
 
 def _restore_tensors(processed: dict[str, Any]) -> None:
diff --git a/nemo_rl/data/processors.py b/nemo_rl/data/processors.py
index 73f1849c162..4c1439a3a3f 100644
--- a/nemo_rl/data/processors.py
+++ b/nemo_rl/data/processors.py
@@ -377,7 +377,11 @@ def vlm_preference_preprocessor(
     THD input; the canonical ``NemotronOmniModel`` inserts media embeddings
     before selecting this rank's context-parallel tokens.
     """
-    from nemo_rl.data.multimodal_utils import PackedTensor
+    from nemo_rl.data.multimodal_utils import (
+        PackedTensor,
+        get_preprocess,
+        uses_image_placeholder,
+    )
 
     completions = datum_dict["completions"]
     if len(completions) != 2:
@@ -386,14 +390,9 @@ def vlm_preference_preprocessor(
     if ordered[0]["rank"] == ordered[1]["rank"]:
         raise ValueError("Tied preference ranks are not supported")
 
-    placeholder_style_processors = {
-        "NemotronNanoVLV2Processor",
-        "NemotronH_Nano_Omni_Reasoning_V3Processor",
-        "NemotronH_Omni_Reasoning_V3Processor",
-    }
     message_processor = (
         _NemotronOmniPreferenceProcessorProxy(processor)
-        if type(processor).__name__ in placeholder_style_processors
+        if uses_image_placeholder(processor)
         else processor
     )
 
@@ -407,7 +406,7 @@ def _format_branch(completion: dict[str, Any]) -> VLMMessageLogType:
 
         # Mirror the canonical Nemotron Omni metadata. Record native image sizes
         # before patchification removes the spatial dimensions.
-        for raw_message in message_log:
+        for raw_message in message_log if uses_image_placeholder(processor) else []:
             message = cast(Any, raw_message)
             pixel_values = message.get("pixel_values")
             if not isinstance(pixel_values, PackedTensor):
@@ -430,8 +429,11 @@ def _format_branch(completion: dict[str, Any]) -> VLMMessageLogType:
                     torch.tensor(image_sizes, dtype=torch.long),
                     dim_to_pack=0,
                 )
-            pixel_values.preprocess_mode = "patchify"
-            pixel_values.preprocess_kwargs = {"patch_dim": 16}
+            message["pixel_values"] = PackedTensor(
+                pixel_values.tensors,
+                pixel_values.dim_to_pack,
+                **get_preprocess(processor, "pixel_values"),
+            )
             imgs_sizes = message.get("imgs_sizes")
             if isinstance(imgs_sizes, PackedTensor) and "num_frames" not in message:
                 sizes = imgs_sizes.as_tensor()
diff --git a/nemo_rl/distributed/batched_data_dict.py b/nemo_rl/distributed/batched_data_dict.py
index 29ed579fb20..95454b024ab 100644
--- a/nemo_rl/distributed/batched_data_dict.py
+++ b/nemo_rl/distributed/batched_data_dict.py
@@ -141,6 +141,7 @@ def get_multimodal_dict(
         as_tensors: bool = False,
         device: Optional[torch.device] = None,
         pixel_dtype: Optional[torch.dtype] = None,
+        pixel_preprocess_mode: Optional[str] = None,
     ) -> dict[str, Any]:
         """Return the multimodal fields as a dict.
 
@@ -191,7 +192,8 @@ def get_multimodal_dict(
                 # unwrapping via as_tensor).
                 if pixel_dtype is not None and k in self._PIXEL_DTYPE_CAST_KEYS:
                     v = v.to_dtype(pixel_dtype)
-                result[k] = v.as_tensor(device=device) if as_tensors else v
+                preprocess_mode = pixel_preprocess_mode if k == "pixel_values" else None
+                result[k] = v.as_tensor(device, preprocess_mode) if as_tensors else v
             elif k in PER_TOKEN_MULTIMODAL_FIELDS:
                 # Plain per-token tensor: emit as-is.
                 result[k] = v
diff --git a/nemo_rl/models/automodel/data.py b/nemo_rl/models/automodel/data.py
index dfd5422c222..f0ca4f03c2c 100644
--- a/nemo_rl/models/automodel/data.py
+++ b/nemo_rl/models/automodel/data.py
@@ -298,7 +298,9 @@ def process_microbatch(
         flash_attn_kwargs = {}
 
     # Add vlm kwargs to model call
-    vlm_kwargs = mb.get_multimodal_dict(as_tensors=True, device=input_ids.device)
+    vlm_kwargs = mb.get_multimodal_dict(
+        True, input_ids.device, None, "pad_to_max_shape"
+    )
     if len(vlm_kwargs) > 0:
         # if there are multimodal kwargs, we don't need to add position_ids (computed internally)
         position_ids = None
diff --git a/nemo_rl/models/policy/workers/dtensor_policy_worker.py b/nemo_rl/models/policy/workers/dtensor_policy_worker.py
index 29e12e3b8ad..0537b4a7c25 100644
--- a/nemo_rl/models/policy/workers/dtensor_policy_worker.py
+++ b/nemo_rl/models/policy/workers/dtensor_policy_worker.py
@@ -792,7 +792,7 @@ def train(
 
                         # add vlm kwargs to model call
                         vlm_kwargs = mb.get_multimodal_dict(
-                            as_tensors=True, device=input_ids.device
+                            True, input_ids.device, None, "pad_to_max_shape"
                         )
                         vlm_kwargs = filter_multimodal_kwargs_for_model(
                             self.model, vlm_kwargs
@@ -1096,7 +1096,7 @@ def get_logprobs(
                 input_ids = lp_batch.get("input_ids").cuda()
                 input_lengths = lp_batch.get("input_lengths")
                 vlm_kwargs = lp_batch.get_multimodal_dict(
-                    as_tensors=True, device=input_ids.device
+                    True, input_ids.device, None, "pad_to_max_shape"
                 )
                 vlm_kwargs = filter_multimodal_kwargs_for_model(self.model, vlm_kwargs)
 
@@ -1537,7 +1537,7 @@ def get_topk_logits(
                 input_ids = lp_batch.get("input_ids").cuda()
                 input_lengths = lp_batch.get("input_lengths")
                 vlm_kwargs = lp_batch.get_multimodal_dict(
-                    as_tensors=True, device=input_ids.device
+                    True, input_ids.device, None, "pad_to_max_shape"
                 )
                 vlm_kwargs = filter_multimodal_kwargs_for_model(self.model, vlm_kwargs)
                 batch_size, seq_len = input_ids.shape

From 2fc65963cf505a41d41728dfd9a22e17de5df43d Mon Sep 17 00:00:00 2001
From: rohitrango <rohit.rango@gmail.com>
Date: Thu, 17 Sep 2026 11:44:21 -0700
Subject: [PATCH 3/3] test(multimodal): cover patchify contracts

Signed-off-by: rohitrango <rohit.rango@gmail.com>
---
 nemo_rl/data/multimodal_utils.py              | 15 +++----
 nemo_rl/environments/nemo_gym.py              |  2 +-
 nemo_rl/models/automodel/data.py              |  8 ++--
 tests/unit/data/test_multimodal_dict.py       | 42 +++++++++++++++++++
 .../data/test_vlm_preference_processor.py     | 12 +++---
 5 files changed, 61 insertions(+), 18 deletions(-)

diff --git a/nemo_rl/data/multimodal_utils.py b/nemo_rl/data/multimodal_utils.py
index e684b21b369..038cdf7b24e 100644
--- a/nemo_rl/data/multimodal_utils.py
+++ b/nemo_rl/data/multimodal_utils.py
@@ -203,10 +203,9 @@ def multimodal_row_tags(
 
     Carries ``shapes`` (per-row, and unrecoverable once ``to_wire`` flattens)
     and the field's preprocessing settings. Deliberately *not* a pad target: the
-    width padding lands at is scratch that the model discards -- mcore crops it
-    via ``imgs_sizes`` before patchification, and the AutoModel path rejects
-    mixed-resolution batches outright -- so each consumer pads to its own view
-    and nothing batch-wide has to be agreed across shards.
+    width padding lands at is scratch that the model discards. Mcore consumes
+    pre-patchified pixels without spatial dimensions, while AutoModel pads at
+    materialization, so nothing batch-wide has to be agreed across shards.
     """
     tags: list[dict[str, Any]] = [{} for _ in range(sample_count)]
     for key, value in multimodal.items():
@@ -290,7 +289,7 @@ def _patchify_segments(segments: list[torch.Tensor], *, patch_dim: int) -> torch
     """Cut pixel segments into vision patches and pack them into one sequence.
 
     Each ``[N, channels, H, W]`` segment is processed at its native resolution
-    into a ``[C_i, P²]`` block, where ``C_i`` is its spatial patch count and
+    into a ``[C_i, P²]`` block, where ``C_i`` is ``N * rows * columns`` and
     ``P²`` is the flattened patch width (``channels * patch_dim**2``). Blocks
     are packed along dimension zero, then a batch dimension is added to produce
     ``[1, total_C, P²]``. Already-patchified segments in that final layout are
@@ -407,8 +406,10 @@ def __init__(
             dim_to_pack: Dimension along which ``as_tensor`` concatenates.
             preprocess_mode: Optional preprocessing applied by ``as_tensor``.
                 Supported values are ``pad_to_max_shape`` and ``patchify``.
-            preprocess_kwargs: Extra arguments for ``preprocess_mode``. Patchify
-                accepts ``patch_dim``.
+                Patchify requires ``dim_to_pack=0`` and changes 4-D inputs into
+                a 3-D ``[1, total_patches, P²]`` tensor.
+            preprocess_kwargs: Extra arguments for ``preprocess_mode``.
+                Patchify requires ``patch_dim``.
         """
         assert tensors is not None, "Input tensors to PackedTensor cannot be None"
 
diff --git a/nemo_rl/environments/nemo_gym.py b/nemo_rl/environments/nemo_gym.py
index 6412af9d394..43b78129518 100644
--- a/nemo_rl/environments/nemo_gym.py
+++ b/nemo_rl/environments/nemo_gym.py
@@ -281,7 +281,7 @@ class NemoGymConfig(TypedDict):
     ]  # For processor reconstruction inside the actor
     pad_dynamic_image_shapes: NotRequired[
         bool
-    ]  # Normalize heterogeneous image tensors while retaining exact imgs_sizes
+    ]  # Preserve heterogeneous shapes for native-resolution patchification
     # Ledger-authoritative token capture (token_capture.enabled): the dumped
     # TokenCaptureConfig. Turns on external staging in Gym's policy model
     # server, switches run_rollouts to receipt mode, and assembles receipts
diff --git a/nemo_rl/models/automodel/data.py b/nemo_rl/models/automodel/data.py
index f0ca4f03c2c..b8897835ed1 100644
--- a/nemo_rl/models/automodel/data.py
+++ b/nemo_rl/models/automodel/data.py
@@ -64,12 +64,12 @@ def filter_multimodal_kwargs_for_model(
     accepted_kwargs = _accepted_forward_kwargs(type(model))
     if accepted_kwargs is None:
         return multimodal_kwargs
-    # A forward that cannot consume imgs_sizes also cannot crop the per-image
-    # pad_to_max_shape padding, so mixed-resolution batches would feed padded
-    # pixels to the vision encoder and mismatch the placeholder count. This is
+    # AutoModel materializes pixels with pad_to_max_shape. A forward that cannot
+    # consume imgs_sizes cannot crop that padding, so mixed-resolution batches
+    # would mismatch the placeholder count. This is
     # the AutoModel Nemotron Omni path (nvidia/Nemotron-3-Nano-Omni-30B-A3B-
     # Reasoning-BF16), whose HF forward takes pixel_values but not imgs_sizes,
-    # unlike the mcore NemotronOmniModel which crops via imgs_sizes.
+    # unlike mcore, which consumes native-resolution pre-patchified pixels.
     imgs_sizes = multimodal_kwargs.get("imgs_sizes")
     if (
         imgs_sizes is not None
diff --git a/tests/unit/data/test_multimodal_dict.py b/tests/unit/data/test_multimodal_dict.py
index cf7f71e2d76..869db09fb07 100644
--- a/tests/unit/data/test_multimodal_dict.py
+++ b/tests/unit/data/test_multimodal_dict.py
@@ -574,6 +574,28 @@ def test_patchify_preserves_pixel_order_within_a_patch():
     torch.testing.assert_close(packed[0, 0], image.reshape(12))
 
 
+def test_patchify_orders_patches_row_major():
+    image = torch.arange(16, dtype=torch.float32).reshape(1, 1, 4, 4)
+    packed = PackedTensor(
+        [image],
+        dim_to_pack=0,
+        preprocess_mode="patchify",
+        preprocess_kwargs={"patch_dim": 2},
+    ).as_tensor()
+
+    torch.testing.assert_close(
+        packed[0],
+        torch.tensor(
+            [
+                [0.0, 1.0, 4.0, 5.0],
+                [2.0, 3.0, 6.0, 7.0],
+                [8.0, 9.0, 12.0, 13.0],
+                [10.0, 11.0, 14.0, 15.0],
+            ]
+        ),
+    )
+
+
 def test_patchify_accepts_already_patchified_segments():
     raw = PackedTensor(
         [torch.ones(1, 3, 32, 32)],
@@ -639,6 +661,26 @@ def test_concat_rejects_mixed_preprocess_settings():
         PackedTensor.concat([padded, plain])
 
 
+def test_concat_rejects_different_patch_dims():
+    first = PackedTensor(
+        torch.ones(1, 3, 32, 32),
+        dim_to_pack=0,
+        preprocess_mode="patchify",
+        preprocess_kwargs={"patch_dim": 16},
+    )
+    second = PackedTensor(
+        torch.ones(1, 3, 32, 32),
+        dim_to_pack=0,
+        preprocess_mode="patchify",
+        preprocess_kwargs={"patch_dim": 8},
+    )
+
+    with pytest.raises(AssertionError, match="same preprocess setting"):
+        PackedTensor.concat([first, second])
+    with pytest.raises(AssertionError, match="same preprocess setting"):
+        PackedTensor.flattened_concat([first, second])
+
+
 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)
diff --git a/tests/unit/data/test_vlm_preference_processor.py b/tests/unit/data/test_vlm_preference_processor.py
index 3f59df888f4..c10d6ceb0ed 100644
--- a/tests/unit/data/test_vlm_preference_processor.py
+++ b/tests/unit/data/test_vlm_preference_processor.py
@@ -28,6 +28,7 @@ def get_vocab(self):
 
 class _ImageProcessor:
     model_input_names = ["pixel_values"]
+    patch_size = 16
 
 
 class NemotronH_Nano_Omni_Reasoning_V3Processor:
@@ -55,10 +56,8 @@ def __call__(self, text, images=None, **kwargs):
             self.saw_explicit_image_placeholder = "<image>" in text
             return {
                 "input_ids": torch.tensor([[9, 10]]),
-                "pixel_values": torch.ones(1, 3, 16, 24),
-                # The processor can report an unpadded crop that differs from
-                # the pixel tensor's padded spatial shape.
-                "imgs_sizes": torch.tensor([[15, 23]]),
+                "pixel_values": torch.ones(1, 3, 16, 32),
+                "imgs_sizes": torch.tensor([[16, 32]]),
             }
         return {"input_ids": torch.tensor([[11]])}
 
@@ -70,7 +69,7 @@ def test_vlm_preference_processor_adds_nemotron_omni_media_metadata():
             "context": [
                 {
                     "role": "user",
-                    "content": [{"type": "image", "image": Image.new("RGB", (24, 16))}],
+                    "content": [{"type": "image", "image": Image.new("RGB", (32, 16))}],
                 }
             ],
             "completions": [
@@ -103,5 +102,6 @@ def test_vlm_preference_processor_adds_nemotron_omni_media_metadata():
         message = message_log[0]
         assert message["pixel_values"].preprocess_mode == "patchify"
         assert message["pixel_values"].preprocess_kwargs == {"patch_dim": 16}
-        assert message["imgs_sizes"].as_tensor().tolist() == [[15, 23]]
+        assert message["pixel_values"].as_tensor().shape == (1, 2, 768)
+        assert message["imgs_sizes"].as_tensor().tolist() == [[16, 32]]
         assert message["num_frames"].as_tensor().tolist() == [1]