diff --git a/examples/models/nemotron/nemotron_3_omni/README.md b/examples/models/nemotron/nemotron_3_omni/README.md
index a66d6b9e6e..6ae501ecc9 100644
--- a/examples/models/nemotron/nemotron_3_omni/README.md
+++ b/examples/models/nemotron/nemotron_3_omni/README.md
@@ -247,21 +247,27 @@ mode:
image's native H×W and produces a variable per-image token count
(`temporal_patch_dim=1`, no temporal fusion). Images are supported on **both** the
HF (`DirectHFSFTDatasetConfig`) path and the Energon path.
-- Video inputs → temporal video embedder (non-dynamic resolution): Frames are resized onto a fixed 512×512 canvas and fused in consecutive
- pairs (`temporal_patch_dim=2`, `separate_video_embedder=True`), so every
- video contributes a constant number of tokens per frame-pair. Videos are supported
- on the **Energon path only**.
+- Video inputs → temporal video embedder: Frames are resized with the public
+ processor's aspect-preserving video policy and fused in consecutive pairs
+ (`temporal_patch_dim=2`, `separate_video_embedder=True`). The resulting token
+ count is derived independently for each frame pair. Videos are supported on
+ the **Energon path only**; the fixed 512×512 policy remains available as an
+ explicit compatibility mode.
-Set the four flags below as a matched column — a mismatched set produces incorrect model's expected input.
+Set the fields below as a matched column; a mismatched set produces input that
+does not match the model configuration.
-| Field | Dynamic resolution (images, variable H×W) | Temporal video (videos, fused pairs, 512²) |
-|----------------------------------------------------|-------------------------------------------|--------------------------------------------|
-| `dataset.task_encoder.use_temporal_video_embedder` | `False` | `True` |
-| `model.temporal_patchrecipe_dim` | `1` | `2` |
-| `model.separate_video_embedder` | `False` | `True` |
-| `model.temporal_ckpt_compat` | `False` | `True` |
+| Field | Dynamic resolution (images, variable H×W) | Temporal video (videos, fused pairs) |
+|----------------------------------------------------|-------------------------------------------|--------------------------------------|
+| `dataset.task_encoder.use_temporal_video_embedder` | `False` | `True` |
+| `dataset.task_encoder.temporal_video_resize_mode` | N/A | `"processor"` |
+| `model.temporal_patch_dim` | `1` | `2` |
+| `model.separate_video_embedder` | `False` | `True` |
+| `model.temporal_ckpt_compat` | `False` | `True` |
-Note:`dataset.task_encoder.use_temporal_video_embedder` only applies to the Energon data path.
+Note: `dataset.task_encoder.use_temporal_video_embedder` and
+`dataset.task_encoder.temporal_video_resize_mode` apply only to the Energon data
+path.
### Image-Text — CORD-V2
@@ -296,13 +302,13 @@ embedder: frames are fused in pairs (`temporal_patch_dim=2`,
`separate_video_embedder=True`) and audio is fed through the Parakeet
encoder. Recipe base: `nemotron_omni_valor32k_*_config`.
-The public processor uses aspect-preserving dynamic sizes for video inference.
-Pinned MCore does not yet support ragged non-square temporal tubelets, so the
-Bridge training path uses an antialiased bicubic 512×512 compatibility canvas.
-Temporal mode uses that canvas for every visual item in the batch, including
-standalone images. Non-temporal image recipes use the public dynamic-resolution
-sizes. Supporting the public non-square video layout requires MCore to
-pixel-shuffle each temporal chunk with its own spatial grid.
+The shipped VALOR32K Energon recipes set
+`dataset.task_encoder.temporal_video_resize_mode="processor"`. Bridge therefore
+uses the public processor's aspect-preserving video grids and expands each
+tubelet placeholder to the number of vision features that grid produces. Frames
+inside one tubelet must share a grid, while different tubelets may have different
+grids and token counts. Use `"fixed_512"` only when reproducing checkpoints or
+runs prepared with the previous square compatibility policy.
Prepare the Energon shards once. For the full walkthrough, see
[`tutorials/data/valor32k-avqa/data-preparation.md`](../../../../tutorials/data/valor32k-avqa/data-preparation.md).
diff --git a/src/megatron/bridge/data/builders/energon.py b/src/megatron/bridge/data/builders/energon.py
index eedebc262b..428a9b4b8e 100644
--- a/src/megatron/bridge/data/builders/energon.py
+++ b/src/megatron/bridge/data/builders/energon.py
@@ -102,8 +102,12 @@ class NemotronOmniEnergonTaskEncoderConfig:
``visual_keys`` is retained for configuration compatibility, but Omni owns
its visual input contract and supports only ``("pixel_values",)``.
+ ``temporal_video_resize_mode="processor"`` applies the public
+ aspect-preserving video grid and exact per-tubelet token counts; the
+ ``"fixed_512"`` default preserves the previous square policy.
``collapse_image_tokens=True`` selects the deprecated LLaVA compatibility
- path; the default ``False`` selects the canonical expanded-sequence path.
+ path and cannot be combined with processor-driven temporal sizing. The
+ default ``False`` selects the canonical expanded-sequence path.
"""
hf_processor_path: str
@@ -115,6 +119,7 @@ class NemotronOmniEnergonTaskEncoderConfig:
video_nframes: int
use_temporal_video_embedder: bool
patch_dim: int
+ temporal_video_resize_mode: Literal["fixed_512", "processor"] = "fixed_512"
collapse_image_tokens: bool = False
trust_remote_code: bool | None = None
@@ -128,6 +133,12 @@ def validate(self) -> None:
raise ValueError(f"{field_name} must be greater than 0.")
if self.video_fps <= 0:
raise ValueError("video_fps must be greater than 0.")
+ if self.temporal_video_resize_mode not in ("fixed_512", "processor"):
+ raise ValueError("temporal_video_resize_mode must be either 'fixed_512' or 'processor'.")
+ if self.collapse_image_tokens and self.temporal_video_resize_mode == "processor":
+ raise ValueError(
+ "temporal_video_resize_mode='processor' requires the canonical expanded-sequence contract."
+ )
if not self.visual_keys or tuple(self.visual_keys) != ("pixel_values",):
raise ValueError("Nemotron Omni visual_keys must be exactly ('pixel_values',).")
@@ -307,6 +318,7 @@ def build_energon_task_encoder(config: EnergonDatasetConfig) -> Any:
video_nframes=task_config.video_nframes,
use_temporal_video_embedder=task_config.use_temporal_video_embedder,
patch_dim=task_config.patch_dim,
+ temporal_video_resize_mode=task_config.temporal_video_resize_mode,
collapse_image_tokens=task_config.collapse_image_tokens,
pad_to_max_length=config.pad_to_max_length,
pad_to_multiple_of=config.pad_to_multiple_of,
diff --git a/src/megatron/bridge/data/energon/nemotron_omni_task_encoder.py b/src/megatron/bridge/data/energon/nemotron_omni_task_encoder.py
index 6e4aae65a1..cdc320a0b3 100644
--- a/src/megatron/bridge/data/energon/nemotron_omni_task_encoder.py
+++ b/src/megatron/bridge/data/energon/nemotron_omni_task_encoder.py
@@ -16,7 +16,7 @@
from collections.abc import Iterator, MutableMapping
from dataclasses import dataclass, fields
-from typing import Any, Mapping, Sequence
+from typing import Any, Literal, Mapping, Sequence
import torch
@@ -128,7 +128,8 @@ class NemotronOmniTaskEncoder(HFTaskEncoder):
assistant masking, modality-token expansion, padding, and in-batch packing
are performed by the canonical expanded-sequence collator for both
Direct-HF and Energon datasets. ``collapse_image_tokens=True`` selects the
- deprecated LLaVA compatibility contract.
+ deprecated LLaVA compatibility contract. Processor-driven temporal video
+ resizing is an Energon option and requires the canonical contract.
"""
def __init__(
@@ -143,6 +144,7 @@ def __init__(
video_nframes: int = 8,
use_temporal_video_embedder: bool = False,
patch_dim: int = 16,
+ temporal_video_resize_mode: Literal["fixed_512", "processor"] = "fixed_512",
pad_to_max_length: bool = False,
pad_to_multiple_of: int = 128,
enable_in_batch_packing: bool = False,
@@ -168,6 +170,7 @@ def __init__(
self.video_nframes = video_nframes
self.use_temporal_video_embedder = use_temporal_video_embedder
self.patch_dim = patch_dim
+ self.temporal_video_resize_mode = temporal_video_resize_mode
self.collapse_image_tokens = collapse_image_tokens
def collate_fn(self, examples: list[dict[str, Any]]) -> dict[str, Any]:
@@ -188,6 +191,7 @@ def collate_fn(self, examples: list[dict[str, Any]]) -> dict[str, Any]:
video_nframes=self.video_nframes,
use_temporal_video_embedder=self.use_temporal_video_embedder,
patch_dim=self.patch_dim,
+ temporal_video_resize_mode=self.temporal_video_resize_mode,
)
def batch(self, samples: list[HFEnergonSample]) -> NemotronOmniTaskBatch:
diff --git a/src/megatron/bridge/models/nemotron_omni/data/collate_fn.py b/src/megatron/bridge/models/nemotron_omni/data/collate_fn.py
index 2038c516a8..b90fcdca13 100644
--- a/src/megatron/bridge/models/nemotron_omni/data/collate_fn.py
+++ b/src/megatron/bridge/models/nemotron_omni/data/collate_fn.py
@@ -20,7 +20,7 @@
import tempfile
import warnings
from collections.abc import Mapping, Sequence
-from typing import Any
+from typing import Any, Literal
import numpy as np
import torch
@@ -38,7 +38,9 @@
from megatron.bridge.models.nemotron_omni.nemotron_omni_utils import (
COMPACT_IMAGE_PLACEHOLDER,
patchify_temporal_frame,
+ processor_patchify_temporal_frames,
temporal_model_frames,
+ temporal_tubelet_feature_counts,
)
from megatron.bridge.training.utils.visual_inputs import GenericVisualInputs
@@ -237,10 +239,16 @@ def _prepare_temporal_rows(
video_fps: float,
video_nframes: int,
patch_dim: int,
+ temporal_video_resize_mode: Literal["fixed_512", "processor"],
) -> tuple[dict[str, Any], list[dict[str, Any]], torch.Tensor]:
"""Build row-local temporal prompts and one packed all-frame vision tensor."""
if temporal_patch_size < 1:
raise ValueError("temporal_patch_size must be at least 1.")
+ if temporal_video_resize_mode not in ("fixed_512", "processor"):
+ raise ValueError(
+ "temporal_video_resize_mode must be either 'fixed_512' or 'processor', "
+ f"got {temporal_video_resize_mode!r}."
+ )
frame_height = frame_width = VISION_FRAME_SIZE
token_rows: list[torch.Tensor] = []
mask_examples: list[dict[str, Any]] = []
@@ -294,11 +302,20 @@ def _prepare_temporal_rows(
row_placeholder_count += 1
text_parts.append("\n".join(video_lines))
model_frames = temporal_model_frames(frames, temporal_patch_size)
- all_patches.extend(
- _patchify_frame(frame, height=frame_height, width=frame_width, patch_dim=patch_dim)
- for frame in model_frames
- )
- all_sizes.extend([[frame_height, frame_width]] * len(model_frames))
+ if temporal_video_resize_mode == "processor":
+ packed_frames, frame_sizes = processor_patchify_temporal_frames(
+ model_frames,
+ image_processor=processor.image_processor,
+ patch_dim=patch_dim,
+ )
+ all_patches.append(packed_frames.squeeze(0))
+ all_sizes.extend(frame_sizes.tolist())
+ else:
+ all_patches.extend(
+ _patchify_frame(frame, height=frame_height, width=frame_width, patch_dim=patch_dim)
+ for frame in model_frames
+ )
+ all_sizes.extend([[frame_height, frame_width]] * len(model_frames))
all_num_frames.append(len(model_frames))
elif isinstance(item, Mapping) and item.get("type") == "text":
text_parts.append(str(item.get("text", "")))
@@ -851,12 +868,17 @@ def nemotron_omni_collate_fn(
video_nframes: int = 8,
use_temporal_video_embedder: bool = False,
patch_dim: int = 16,
+ temporal_video_resize_mode: Literal["fixed_512", "processor"] = "fixed_512",
collapse_image_tokens: bool = False,
) -> dict[str, Any]:
"""Build one model-ready Omni batch from either HF or Energon examples.
The canonical :class:`NemotronOmniModel` consumes the processor-expanded
token sequence, with one image placeholder for every projected feature.
+ ``temporal_video_resize_mode="processor"`` uses the image processor's
+ aspect-preserving video grid and derives each tubelet's placeholder count
+ from its returned size metadata. The default ``"fixed_512"`` preserves the
+ prior square policy and Direct-HF behavior.
Use :func:`nemotron_omni_llava_collate_fn` for the legacy LLaVA
collapse/expand contract.
"""
@@ -868,6 +890,11 @@ def nemotron_omni_collate_fn(
FutureWarning,
stacklevel=2,
)
+ if use_temporal_video_embedder and temporal_video_resize_mode == "processor":
+ raise ValueError(
+ "Processor-driven temporal video sizing is supported only by the canonical expanded-sequence "
+ "NemotronOmniModel; the deprecated LLaVA collapse/expand contract requires fixed_512."
+ )
_validate_nemotron_omni_visual_keys(visual_keys)
del start_of_response_token, min_pixels, max_pixels
if not examples:
@@ -891,6 +918,7 @@ def nemotron_omni_collate_fn(
video_fps=video_fps,
video_nframes=video_nframes,
patch_dim=patch_dim,
+ temporal_video_resize_mode=temporal_video_resize_mode,
)
use_per_image_token_counts = False
else:
@@ -936,13 +964,18 @@ def nemotron_omni_collate_fn(
adjusted, loss_mask = _adjust_image_placeholders(batch, loss_mask, processor, num_tiles)
batch["input_ids"] = adjusted["input_ids"]
batch["attention_mask"] = adjusted["attention_mask"]
- elif use_temporal_video_embedder and num_tiles is not None:
- tokens_per_tubelet = _pixel_shuffled_token_count(
- height=VISION_FRAME_SIZE,
- width=VISION_FRAME_SIZE,
+ elif use_temporal_video_embedder and num_tiles is not None and num_tiles.numel() > 0:
+ replacement_counts = temporal_tubelet_feature_counts(
+ batch["imgs_sizes"],
+ batch["num_frames"],
+ temporal_patch_size=temporal_patch_size,
patch_dim=patch_dim,
)
- replacement_counts = torch.full_like(num_tiles, tokens_per_tubelet)
+ if replacement_counts.numel() != num_tiles.numel():
+ raise ValueError(
+ "Temporal vision metadata produced "
+ f"{replacement_counts.numel()} feature counts for {num_tiles.numel()} compact placeholders."
+ )
adjusted, loss_mask = _adjust_image_placeholders(
batch,
loss_mask,
diff --git a/src/megatron/bridge/models/nemotron_omni/nemotron_omni_utils.py b/src/megatron/bridge/models/nemotron_omni/nemotron_omni_utils.py
index 3d7fc045f9..670977fcff 100644
--- a/src/megatron/bridge/models/nemotron_omni/nemotron_omni_utils.py
+++ b/src/megatron/bridge/models/nemotron_omni/nemotron_omni_utils.py
@@ -27,12 +27,12 @@
def patchify_temporal_frame(frame: Any, *, height: int, width: int, patch_dim: int) -> torch.Tensor:
- """Resize and normalize one frame for MCore's square temporal RADIO path.
+ """Resize and normalize one frame for the fixed temporal RADIO policy.
- The public HF processor preserves aspect ratio, but pinned MCore requires
- temporal tubelets to share one square spatial grid. This helper is shared
- by training collation and inference so both paths use the same antialiased
- bicubic interpolation, RADIO normalization, and patch layout.
+ This compatibility helper intentionally places every frame on the supplied
+ canvas. It is shared by fixed-policy training collation and inference so
+ both paths use the same antialiased bicubic interpolation, RADIO
+ normalization, and patch layout.
Args:
frame: PIL-compatible image with ``convert("RGB")`` support.
@@ -89,6 +89,185 @@ def temporal_model_frames(frames: Sequence[_FrameT], temporal_patch_size: int) -
return model_frames
+def processor_patchify_temporal_frames(
+ frames: Sequence[Any],
+ *,
+ image_processor: Any,
+ patch_dim: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+ """Apply the public processor's video resize policy and pack normalized frames.
+
+ The Nemotron Omni remote image processor selects a patch-grid size using a
+ separate aspect-preserving video policy. It exposes that policy through the
+ same temporary ``_is_video_mode`` switch used by the public multimodal
+ processor. Reusing the processor output here keeps Bridge pixels, grid
+ metadata, and placeholder counts aligned with HF/vLLM preprocessing.
+
+ Args:
+ frames: Frames belonging to one source video, in model order.
+ image_processor: Nemotron Omni dynamic-resolution image processor.
+ patch_dim: Vision patch edge length.
+
+ Returns:
+ A pair containing packed patches with shape
+ ``[1, total_patches, 3 * patch_dim**2]`` and one ``(height, width)``
+ metadata row per frame.
+
+ Raises:
+ ValueError: If the processor contract or returned frame metadata is
+ incompatible with RADIO's packed temporal input.
+ """
+ if not frames:
+ raise ValueError("Processor-driven temporal preprocessing requires at least one frame.")
+ if patch_dim <= 0:
+ raise ValueError("patch_dim must be greater than 0.")
+ if not hasattr(image_processor, "_is_video_mode"):
+ raise ValueError(
+ "Processor-driven temporal preprocessing requires a Nemotron Omni image processor "
+ "with the public '_is_video_mode' video contract."
+ )
+
+ previous_video_mode = image_processor._is_video_mode
+ image_processor._is_video_mode = True
+ try:
+ processed = image_processor(images=list(frames), return_tensors=None)
+ finally:
+ image_processor._is_video_mode = previous_video_mode
+
+ pixel_values = processed.get("pixel_values")
+ if isinstance(pixel_values, list):
+ frame_tensors = [torch.as_tensor(value) for value in pixel_values]
+ elif torch.is_tensor(pixel_values) and pixel_values.ndim == 4:
+ frame_tensors = list(pixel_values.unbind(0))
+ elif torch.is_tensor(pixel_values) and pixel_values.ndim == 3:
+ frame_tensors = [pixel_values]
+ else:
+ shape = getattr(pixel_values, "shape", None)
+ raise ValueError(f"Video image processor returned unsupported pixel_values shape {shape}.")
+
+ raw_imgs_sizes = processed.get("imgs_sizes")
+ if raw_imgs_sizes is None:
+ raise ValueError("Video image processor must return imgs_sizes metadata.")
+ imgs_sizes = torch.as_tensor(raw_imgs_sizes, dtype=torch.long)
+ if imgs_sizes.ndim != 2 or imgs_sizes.shape != (len(frame_tensors), 2):
+ raise ValueError(
+ "Video image processor must return one (height, width) row per frame; "
+ f"got {tuple(imgs_sizes.shape)} for {len(frame_tensors)} frames."
+ )
+ if len(frame_tensors) != len(frames):
+ raise ValueError(f"Video image processor returned {len(frame_tensors)} tensors for {len(frames)} frames.")
+
+ reported_tokens = processed.get("num_tokens")
+ if reported_tokens is None or len(reported_tokens) != len(frame_tensors):
+ count = None if reported_tokens is None else len(reported_tokens)
+ raise ValueError(
+ "Video image processor must return one num_tokens entry per frame; "
+ f"got {count} entries for {len(frame_tensors)} frames."
+ )
+
+ patches = []
+ for frame_index, (frame, size, reported_count) in enumerate(
+ zip(frame_tensors, imgs_sizes.tolist(), reported_tokens, strict=True)
+ ):
+ if frame.ndim != 3:
+ raise ValueError(f"Processed video frame {frame_index} must have shape [3,H,W], got {tuple(frame.shape)}.")
+ channels, height, width = frame.shape
+ expected_height, expected_width = (int(value) for value in size)
+ if channels != 3 or (height, width) != (expected_height, expected_width):
+ raise ValueError(
+ f"Processed video frame {frame_index} shape {tuple(frame.shape)} does not match "
+ f"imgs_sizes row {(expected_height, expected_width)}."
+ )
+ if height % patch_dim or width % patch_dim:
+ raise ValueError(f"Video frame {height}x{width} is not divisible by patch_dim={patch_dim}.")
+ patch_rows, patch_cols = height // patch_dim, width // patch_dim
+ if patch_rows % 2 or patch_cols % 2:
+ raise ValueError(f"Video patch grid {patch_rows}x{patch_cols} is not divisible by the 2x2 pixel shuffle.")
+ expected_count = (patch_rows * patch_cols) // 4
+ if int(reported_count) != expected_count:
+ raise ValueError(
+ f"Video image processor reported {int(reported_count)} tokens for frame {frame_index}, "
+ f"but grid {patch_rows}x{patch_cols} produces {expected_count}."
+ )
+ patches.append(
+ frame.reshape(channels, patch_rows, patch_dim, patch_cols, patch_dim)
+ .permute(1, 3, 0, 2, 4)
+ .reshape(patch_rows * patch_cols, channels * patch_dim * patch_dim)
+ .contiguous()
+ )
+
+ return torch.cat(patches, dim=0).unsqueeze(0).contiguous(), imgs_sizes
+
+
+def temporal_tubelet_feature_counts(
+ imgs_sizes: torch.Tensor,
+ num_frames: torch.Tensor,
+ *,
+ temporal_patch_size: int,
+ patch_dim: int,
+ pixel_shuffle_factor: int = 2,
+) -> torch.Tensor:
+ """Compute projected RADIO feature counts in temporal tubelet order.
+
+ Args:
+ imgs_sizes: One ``(height, width)`` row per ungrouped input frame.
+ num_frames: Number of frame rows owned by each image or video item.
+ temporal_patch_size: Frames fused into one video tubelet.
+ patch_dim: Vision patch edge length.
+ pixel_shuffle_factor: Spatial reduction factor per dimension.
+
+ Returns:
+ One feature count per image or temporal tubelet, aligned with compact
+ ``
`` wrappers.
+
+ Raises:
+ ValueError: If metadata is inconsistent, a patch grid cannot be
+ shuffled, or frames inside one tubelet use different grids.
+ """
+ if imgs_sizes.ndim != 2 or imgs_sizes.shape[1] != 2:
+ raise ValueError(f"imgs_sizes must have shape [N, 2], got {tuple(imgs_sizes.shape)}.")
+ if temporal_patch_size <= 0:
+ raise ValueError("temporal_patch_size must be greater than 0.")
+ if patch_dim <= 0:
+ raise ValueError("patch_dim must be greater than 0.")
+ if pixel_shuffle_factor <= 0:
+ raise ValueError("pixel_shuffle_factor must be greater than 0.")
+
+ frame_counts = [int(count) for count in num_frames.reshape(-1).tolist()]
+ if not frame_counts or any(count <= 0 for count in frame_counts):
+ raise ValueError("num_frames must contain positive entries.")
+ if sum(frame_counts) != imgs_sizes.shape[0]:
+ raise ValueError(
+ f"num_frames accounts for {sum(frame_counts)} frames but imgs_sizes has {imgs_sizes.shape[0]} rows."
+ )
+
+ sizes = [(int(height), int(width)) for height, width in imgs_sizes.tolist()]
+ feature_counts = []
+ frame_offset = 0
+ for media_index, frame_count in enumerate(frame_counts):
+ group_width = 1 if frame_count == 1 else temporal_patch_size
+ for group_start in range(0, frame_count, group_width):
+ group = sizes[frame_offset + group_start : frame_offset + min(group_start + group_width, frame_count)]
+ if any(size != group[0] for size in group[1:]):
+ raise ValueError(
+ f"Temporal tubelet {len(feature_counts)} in media item {media_index} has inconsistent "
+ f"frame sizes {group}."
+ )
+ height, width = group[0]
+ if height <= 0 or width <= 0 or height % patch_dim or width % patch_dim:
+ raise ValueError(f"Frame size {height}x{width} is not divisible by patch_dim={patch_dim}.")
+ patch_rows, patch_cols = height // patch_dim, width // patch_dim
+ if patch_rows % pixel_shuffle_factor or patch_cols % pixel_shuffle_factor:
+ raise ValueError(
+ f"Patch grid {patch_rows}x{patch_cols} is not divisible by the "
+ f"{pixel_shuffle_factor}x{pixel_shuffle_factor} pixel shuffle."
+ )
+ feature_counts.append((patch_rows * patch_cols) // (pixel_shuffle_factor**2))
+ frame_offset += frame_count
+
+ return torch.tensor(feature_counts, dtype=torch.int, device=imgs_sizes.device)
+
+
def inference_num_image_tiles(
imgs_sizes: torch.Tensor,
*,
diff --git a/src/megatron/bridge/recipes/nemotron_omni/h100/nemotron_omni.py b/src/megatron/bridge/recipes/nemotron_omni/h100/nemotron_omni.py
index b0d00b58a8..75064ce1e6 100644
--- a/src/megatron/bridge/recipes/nemotron_omni/h100/nemotron_omni.py
+++ b/src/megatron/bridge/recipes/nemotron_omni/h100/nemotron_omni.py
@@ -54,6 +54,7 @@ def _make_nemotron_omni_energon_dataset(micro_batch_size: int) -> EnergonDataset
video_nframes=8,
use_temporal_video_embedder=True,
patch_dim=16,
+ temporal_video_resize_mode="processor",
trust_remote_code=True,
),
enable_in_batch_packing=False,
diff --git a/tests/unit_tests/data/builders/test_energon_builder.py b/tests/unit_tests/data/builders/test_energon_builder.py
index e7b494a7aa..ebd37a8d03 100644
--- a/tests/unit_tests/data/builders/test_energon_builder.py
+++ b/tests/unit_tests/data/builders/test_energon_builder.py
@@ -283,6 +283,7 @@ def test_nemotron_factory_preserves_omni_settings(monkeypatch: pytest.MonkeyPatc
video_nframes=8,
use_temporal_video_embedder=True,
patch_dim=16,
+ temporal_video_resize_mode="processor",
),
enable_in_batch_packing=True,
)
@@ -303,6 +304,7 @@ def test_nemotron_factory_preserves_omni_settings(monkeypatch: pytest.MonkeyPatc
assert encoder_cls.call_args.kwargs["processor"] is processor
assert encoder_cls.call_args.kwargs["max_audio_duration"] == 10.0
assert encoder_cls.call_args.kwargs["use_temporal_video_embedder"] is True
+ assert encoder_cls.call_args.kwargs["temporal_video_resize_mode"] == "processor"
assert encoder_cls.call_args.kwargs["collapse_image_tokens"] is False
assert encoder_cls.call_args.kwargs["enable_in_batch_packing"] is True
@@ -324,6 +326,35 @@ def test_nemotron_config_rejects_unsupported_visual_keys():
config.validate()
+@pytest.mark.parametrize(
+ ("overrides", "match"),
+ [
+ ({"temporal_video_resize_mode": "unknown"}, "temporal_video_resize_mode"),
+ (
+ {"temporal_video_resize_mode": "processor", "collapse_image_tokens": True},
+ "canonical expanded-sequence",
+ ),
+ ],
+)
+def test_nemotron_config_rejects_invalid_temporal_video_resize_settings(overrides, match):
+ kwargs = {
+ "hf_processor_path": "nvidia/model",
+ "max_audio_duration": 10.0,
+ "num_mel_bins": 128,
+ "visual_keys": ("pixel_values",),
+ "temporal_patch_size": 2,
+ "video_fps": 1.0,
+ "video_nframes": 8,
+ "use_temporal_video_embedder": True,
+ "patch_dim": 16,
+ **overrides,
+ }
+ config = NemotronOmniEnergonTaskEncoderConfig(**kwargs)
+
+ with pytest.raises(ValueError, match=match):
+ config.validate()
+
+
def test_builder_honors_requested_splits_and_reuses_runtime_encoder(monkeypatch: pytest.MonkeyPatch):
config = _qwen_config(
micro_batch_size=1,
diff --git a/tests/unit_tests/data/collators/test_model_collators.py b/tests/unit_tests/data/collators/test_model_collators.py
index d97ab55d4e..63fe3a205c 100644
--- a/tests/unit_tests/data/collators/test_model_collators.py
+++ b/tests/unit_tests/data/collators/test_model_collators.py
@@ -2634,6 +2634,8 @@ def test_nemotron_omni_expanded_collate_emits_one_placeholder_per_temporal_featu
"input_ids": input_ids,
"attention_mask": torch.ones_like(input_ids),
"visual_inputs": GenericVisualInputs(pixel_values=torch.ones(1, 1, 768)),
+ "imgs_sizes": torch.tensor([[512, 512], [512, 512]]),
+ "num_frames": torch.tensor([2]),
}
examples = [{"conversation": [{"role": "user", "content": "one tubelet"}]}]
monkeypatch.setattr(
diff --git a/tests/unit_tests/data/energon/test_nemotron_omni_task_encoder.py b/tests/unit_tests/data/energon/test_nemotron_omni_task_encoder.py
index cea2f8daa6..63478df5dd 100644
--- a/tests/unit_tests/data/energon/test_nemotron_omni_task_encoder.py
+++ b/tests/unit_tests/data/energon/test_nemotron_omni_task_encoder.py
@@ -108,6 +108,33 @@ def __call__(self, **kwargs):
return output
+class _AspectVideoImageProcessor:
+ max_num_patches = 13312
+ _is_video_mode = False
+
+ def __init__(self):
+ self.calls = []
+
+ def __call__(self, *, images, return_tensors):
+ assert self._is_video_mode is True
+ assert return_tensors is None
+ self.calls.append(images)
+ sizes = []
+ pixels = []
+ token_counts = []
+ for image in images:
+ if image.width * 9 == image.height * 16:
+ height, width = 384, 672
+ elif image.width == image.height * 2:
+ height, width = 352, 704
+ else:
+ raise AssertionError(f"Unexpected test aspect ratio {image.size}")
+ sizes.append((height, width))
+ pixels.append(torch.ones(3, height, width))
+ token_counts.append((height // 16) * (width // 16) // 4)
+ return {"pixel_values": pixels, "imgs_sizes": sizes, "num_tokens": token_counts}
+
+
def _sample(conversation, *, key="sample", imgs=None, videos=None, audio=None):
return ChatMLSample(
**sample_metadata_kwargs(key=key, restore_key=(), subflavors={}),
@@ -404,6 +431,127 @@ def test_energon_temporal_video_defaults_to_expanded_contract(monkeypatch):
assert batch.num_frames.tolist() == [2]
+@pytest.mark.parametrize("resize_mode", ["fixed_512", "processor"])
+def test_energon_temporal_mode_all_text_batch_has_no_visual_metadata(monkeypatch, resize_mode):
+ monkeypatch.setattr(omni_collate, "build_assistant_loss_mask", _mask_all_tokens)
+ processor = _Processor([[1, 21, PAD_AND_END_ID]])
+ encoder = NemotronOmniTaskEncoder(
+ processor=processor,
+ seq_length=16,
+ use_temporal_video_embedder=True,
+ temporal_video_resize_mode=resize_mode,
+ pad_to_multiple_of=1,
+ )
+ encoded = encoder.encode_sample(_sample([{"role": "user", "content": "Describe the clip."}]))
+
+ batch = encoder.batch([encoded])
+
+ assert batch.visual_inputs is None
+ assert batch.imgs_sizes is None
+ assert batch.num_frames is None
+ assert batch.num_image_tiles is None
+
+
+def test_energon_processor_temporal_video_uses_per_tubelet_feature_counts(monkeypatch):
+ from PIL import Image
+
+ monkeypatch.setattr(omni_collate, "build_assistant_loss_mask", _mask_all_tokens)
+ rows = [
+ [1, IMG_START_ID, IMAGE_TOKEN_ID, IMG_END_ID, 21, PAD_AND_END_ID],
+ [2, IMG_START_ID, IMAGE_TOKEN_ID, IMG_END_ID, 31, PAD_AND_END_ID],
+ ]
+ processor = _Processor(rows)
+ processor.image_processor = _AspectVideoImageProcessor()
+ encoder = NemotronOmniTaskEncoder(
+ processor=processor,
+ seq_length=512,
+ temporal_patch_size=2,
+ use_temporal_video_embedder=True,
+ patch_dim=16,
+ temporal_video_resize_mode="processor",
+ pad_to_multiple_of=1,
+ )
+ source_samples = [
+ _sample(
+ [{"role": "user", "content": [{"type": "video"}]}],
+ key="wide-16-9",
+ videos=[[Image.new("RGB", (160, 90)), Image.new("RGB", (160, 90))]],
+ ),
+ _sample(
+ [{"role": "user", "content": [{"type": "video"}]}],
+ key="wide-2-1",
+ videos=[[Image.new("RGB", (200, 100)), Image.new("RGB", (200, 100))]],
+ ),
+ ]
+
+ batch = encoder.batch([encoder.encode_sample(sample) for sample in source_samples])
+
+ assert [(row == IMAGE_TOKEN_ID).sum().item() for row in batch.input_ids] == [252, 242]
+ assert batch.attention_mask.sum(dim=1).tolist() == [257, 247]
+ assert batch.imgs_sizes.tolist() == [[384, 672], [384, 672], [352, 704], [352, 704]]
+ assert batch.num_frames.tolist() == [2, 2]
+ assert batch.num_image_tiles.tolist() == [1, 1]
+ assert batch.visual_inputs.pixel_values.shape == (1, 3952, 768)
+ assert len(processor.image_processor.calls) == 2
+
+
+def test_energon_processor_temporal_video_packing_preserves_ragged_boundaries(monkeypatch):
+ from PIL import Image
+
+ monkeypatch.setattr(omni_collate, "build_assistant_loss_mask", _mask_all_tokens)
+ rows = [
+ [1, IMG_START_ID, IMAGE_TOKEN_ID, IMG_END_ID, 21, PAD_AND_END_ID],
+ [2, IMG_START_ID, IMAGE_TOKEN_ID, IMG_END_ID, 31, PAD_AND_END_ID],
+ ]
+ processor = _Processor(rows)
+ processor.image_processor = _AspectVideoImageProcessor()
+ encoder = NemotronOmniTaskEncoder(
+ processor=processor,
+ seq_length=512,
+ temporal_patch_size=2,
+ use_temporal_video_embedder=True,
+ patch_dim=16,
+ temporal_video_resize_mode="processor",
+ enable_in_batch_packing=True,
+ in_batch_packing_pad_to_multiple_of=8,
+ pad_to_multiple_of=1,
+ )
+ samples = [
+ _sample(
+ [{"role": "user", "content": [{"type": "video"}]}],
+ key="wide-16-9",
+ videos=[[Image.new("RGB", (160, 90)), Image.new("RGB", (160, 90))]],
+ ),
+ _sample(
+ [{"role": "user", "content": [{"type": "video"}]}],
+ key="wide-2-1",
+ videos=[[Image.new("RGB", (200, 100)), Image.new("RGB", (200, 100))]],
+ ),
+ ]
+
+ batch = encoder.batch([encoder.encode_sample(sample) for sample in samples])
+
+ assert batch.input_ids.shape == (1, 512)
+ assert int((batch.input_ids == IMAGE_TOKEN_ID).sum().item()) == 494
+ assert batch.cu_seqlens_q.tolist() == [0, 257, 504]
+ assert batch.cu_seqlens_q_padded.tolist() == [0, 264, 512]
+ assert batch.total_tokens == 512
+ assert batch.padding_mask.sum().item() == 8
+
+
+def test_energon_processor_temporal_video_rejects_legacy_contract():
+ encoder = NemotronOmniTaskEncoder(
+ processor=_Processor([[1]]),
+ use_temporal_video_embedder=True,
+ temporal_video_resize_mode="processor",
+ collapse_image_tokens=True,
+ pad_to_multiple_of=1,
+ )
+
+ with pytest.raises(ValueError, match="canonical expanded-sequence"):
+ encoder.collate_fn([{"conversation": [{"role": "user", "content": "text"}]}])
+
+
def test_energon_multiple_raw_video_bytes_keep_placeholder_order():
raw_videos = [b"first-mp4", b"second-mp4"]
encoder = NemotronOmniTaskEncoder(processor=_Processor([[1]]), pad_to_multiple_of=1)
diff --git a/tests/unit_tests/models/nemotron_omni/test_nemotron_omni_model.py b/tests/unit_tests/models/nemotron_omni/test_nemotron_omni_model.py
index aba8b7fba0..6e90d5e24d 100644
--- a/tests/unit_tests/models/nemotron_omni/test_nemotron_omni_model.py
+++ b/tests/unit_tests/models/nemotron_omni/test_nemotron_omni_model.py
@@ -826,6 +826,32 @@ def test_real_radio_multiframe_video_forward(single_rank_model_parallel):
assert torch.isfinite(output).all()
+@pytest.mark.run_only_on("GPU")
+def test_real_radio_ragged_rectangular_multiframe_video_forward(single_rank_model_parallel):
+ del single_rank_model_parallel
+ provider = _TinyOmniProvider()
+ provider.finalize()
+ model = provider.provide().cuda().eval()
+ input_ids = torch.tensor([[7, 18, 18, 18, 18, 18, 9, 10]], device="cuda")
+ patch_features = 3 * model.patch_dim**2
+
+ with torch.no_grad():
+ output = model(
+ input_ids=input_ids,
+ attention_mask=torch.ones_like(input_ids, dtype=torch.bool),
+ pixel_values=torch.randn(1, 40, patch_features, device="cuda"),
+ imgs_sizes=torch.tensor(
+ [[32, 64], [32, 64], [32, 96], [32, 96]],
+ dtype=torch.int32,
+ device="cuda",
+ ),
+ num_frames=torch.tensor([2, 2], dtype=torch.int32, device="cuda"),
+ )
+
+ assert output.shape == (1, 8, 128)
+ assert torch.isfinite(output).all()
+
+
@pytest.mark.run_only_on("GPU")
def test_packed_mamba_resets_state_between_samples(single_rank_model_parallel):
del single_rank_model_parallel
diff --git a/tests/unit_tests/models/nemotron_omni/test_nemotron_omni_utils.py b/tests/unit_tests/models/nemotron_omni/test_nemotron_omni_utils.py
index eccc9779a1..814eb9634a 100644
--- a/tests/unit_tests/models/nemotron_omni/test_nemotron_omni_utils.py
+++ b/tests/unit_tests/models/nemotron_omni/test_nemotron_omni_utils.py
@@ -23,8 +23,10 @@
inference_expanded_image_token_counts,
inference_merged_sequence_length,
inference_num_image_tiles,
+ processor_patchify_temporal_frames,
select_inference_next_token,
temporal_model_frames,
+ temporal_tubelet_feature_counts,
valid_audio_feature_lengths,
)
from megatron.bridge.models.nemotron_vl.nemotron_vl_utils import adjust_image_tokens
@@ -96,6 +98,68 @@ def test_temporal_model_frames_rejects_invalid_patch_size():
temporal_model_frames([object()], 0)
+def test_processor_patchify_temporal_frames_preserves_processor_pixels_and_sizes():
+ class _VideoImageProcessor:
+ _is_video_mode = False
+
+ def __call__(self, *, images, return_tensors):
+ assert self._is_video_mode is True
+ assert len(images) == 2
+ assert return_tensors is None
+ return {
+ "pixel_values": torch.arange(2 * 3 * 32 * 64, dtype=torch.float32).reshape(2, 3, 32, 64),
+ "imgs_sizes": [(32, 64), (32, 64)],
+ "num_tokens": [2, 2],
+ }
+
+ image_processor = _VideoImageProcessor()
+
+ patches, imgs_sizes = processor_patchify_temporal_frames(
+ [object(), object()],
+ image_processor=image_processor,
+ patch_dim=16,
+ )
+
+ assert patches.shape == (1, 16, 768)
+ assert imgs_sizes.tolist() == [[32, 64], [32, 64]]
+ assert image_processor._is_video_mode is False
+ first_patch = patches[0, 0].reshape(3, 16, 16)
+ expected = torch.arange(3 * 32 * 64, dtype=torch.float32).reshape(3, 32, 64)[:, :16, :16]
+ assert torch.equal(first_patch, expected)
+
+
+def test_temporal_tubelet_feature_counts_uses_each_post_resize_grid():
+ counts = temporal_tubelet_feature_counts(
+ torch.tensor([[384, 672], [384, 672], [352, 704], [352, 704]]),
+ torch.tensor([2, 2]),
+ temporal_patch_size=2,
+ patch_dim=16,
+ )
+
+ assert counts.tolist() == [252, 242]
+
+
+def test_temporal_tubelet_feature_counts_preserves_incomplete_final_group():
+ counts = temporal_tubelet_feature_counts(
+ torch.tensor([[384, 672], [384, 672], [384, 672]]),
+ torch.tensor([3]),
+ temporal_patch_size=2,
+ patch_dim=16,
+ )
+
+ assert counts.tolist() == [252, 252]
+
+
+def test_temporal_tubelet_feature_counts_rejects_mixed_grids_within_tubelet():
+ with pytest.raises(ValueError, match="inconsistent frame sizes"):
+ temporal_tubelet_feature_counts(
+ torch.tensor([[384, 672], [352, 704]]),
+ torch.tensor([2]),
+ temporal_patch_size=2,
+ patch_dim=16,
+ )
+
+
def test_inference_num_image_tiles_uses_post_shuffle_dynamic_image_counts():
imgs_sizes = torch.tensor([[512, 512], [512, 256], [256, 256]])
diff --git a/tests/unit_tests/recipes/nemotron_omni/test_nemotron_omni_recipes.py b/tests/unit_tests/recipes/nemotron_omni/test_nemotron_omni_recipes.py
index 1e3645c6b2..b00c63e29d 100644
--- a/tests/unit_tests/recipes/nemotron_omni/test_nemotron_omni_recipes.py
+++ b/tests/unit_tests/recipes/nemotron_omni/test_nemotron_omni_recipes.py
@@ -220,6 +220,7 @@ def test_valor32k_sft_recipe_uses_temporal_omni_task_encoder_config(fake_process
assert cfg.dataset.task_encoder.num_mel_bins == 128
assert cfg.dataset.task_encoder.use_temporal_video_embedder is True
assert cfg.dataset.task_encoder.patch_dim == 16
+ assert cfg.dataset.task_encoder.temporal_video_resize_mode == "processor"
assert cfg.dataset.task_encoder.collapse_image_tokens is False
assert cfg.model.temporal_patch_dim == 2
assert cfg.model.separate_video_embedder is True
@@ -236,6 +237,7 @@ def test_valor32k_peft_recipe_configures_lora_and_freezing(fake_processor):
assert isinstance(cfg.dataset, EnergonDatasetConfig)
assert isinstance(cfg.dataset.task_encoder, NemotronOmniEnergonTaskEncoderConfig)
assert cfg.dataset.task_encoder.use_temporal_video_embedder is True
+ assert cfg.dataset.task_encoder.temporal_video_resize_mode == "processor"
assert cfg.peft is not None
assert cfg.peft.target_modules == [
"linear_qkv",