Skip to content

fix(model): use dynamic temporal video token counts #5590

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
cuichenx wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

fix(model): use dynamic temporal video token counts #5590

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
44 changes: 25 additions & 19 deletions examples/models/nemotron/nemotron_3_omni/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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).
Expand Down
14 changes: 13 additions & 1 deletion src/megatron/bridge/data/builders/energon.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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

Expand All @@ -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',).")

Expand Down Expand Up @@ -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,
Expand Down
8 changes: 6 additions & 2 deletions src/megatron/bridge/data/energon/nemotron_omni_task_encoder.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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__(
Expand All @@ -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,
Expand All @@ -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]:
Expand All @@ -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:
Expand Down
55 changes: 44 additions & 11 deletions src/megatron/bridge/models/nemotron_omni/data/collate_fn.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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

Expand Down Expand Up @@ -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]] = []
Expand Down Expand Up @@ -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", "")))
Expand Down Expand Up @@ -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.
"""
Expand All @@ -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:
Expand All @@ -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:
Expand Down Expand Up @@ -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,
Expand Down
Loading
Loading