nemo_automodel.components.distributed.pipelining.media

Consume explicitly staged pipeline media in microbatch order.

Module Contents

Classes

Name	Description
`VlmMediaChunk`	One microbatch’s views of staged image and video tensors.

Functions

Name	Description
`next_vlm_media_chunk`	Advance staged media once, including for a text-only microbatch.

API

class nemo_automodel.components.distributed.pipelining.media.VlmMediaChunk(
    pixel_values: torch.Tensor | None,
    image_grid: torch.Tensor | None,
    pixel_values_videos: torch.Tensor | None,
    video_grid: torch.Tensor | None,
    patch_pixel_values: torch.Tensor | None = None,
    num_patches: torch.Tensor | None = None,
    patch_newline_mask: torch.Tensor | None = None
)

Dataclass

One microbatch’s views of staged image and video tensors.

Absent or empty staged tensors become None; every nonempty field aliases its staged tensor unchanged.

image_grid

Tensor | None

Per-image metadata with images on axis 0: patch grids [images, 2] (height, width) or [images, 3] (time, height, width), original sizes [images, 2], or per-patch positions [images, max_patches, 2].

num_patches

Tensor | None = None

Tensor of shape [batch] with each sample’s crop-patch count (Step3 stages one image per sample); it indexes patch_pixel_values and patch_newline_mask.

patch_newline_mask

Tensor | None = None

Tensor of shape [patches] marking patch-row newline positions.

patch_pixel_values

Tensor | None = None

Crop patches of shape [patches, channels, height, width], concatenated in image order.

pixel_values

Tensor | None

Image pixels in processor layout with images or patches on axis 0: complete images [images, channels, height, width], padded patches [images, max_patches, features], flat patch rows [patches, patch_features], or patch-major [patches, channels, patch_height, patch_width].

pixel_values_videos

Tensor | None

Video pixels with patches on axis 0, e.g. flat patch rows [patches, patch_features] in video order.

video_grid

Tensor | None

Tensor of shape [videos, 3] storing each video’s (time, height, width) patch grid.

nemo_automodel.components.distributed.pipelining.media.next_vlm_media_chunk(
    module: torch.nn.Module
) -> nemo_automodel.components.distributed.pipelining.media.VlmMediaChunk | None

Advance staged media once, including for a text-only microbatch.

Parameters:

module

torch.nn.Module

Pipeline stage carrying image/video chunk lists. Pixel tensors retain their image/patch-major layouts; metadata is [media, coordinates]. Optional patch pixels retain [patches, channels, height, width], patch counts [images], and newline masks [patches] for patch-based consumers.

Returns: VlmMediaChunk | None

The next chunk with original nonempty tensor views, or None when no media

Raises:

ValueError: If a media or grid list is exhausted before the schedule.