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