Skip to navigation
ReferenceFull Library ReferenceNemo AutomodelNemo AutomodelComponentsDistributedPipeliningnemo_automodel.components.distributed.pipelining.media

nemo_automodel.components.distributed.pipelining.media

View as Markdown

Consume explicitly staged pipeline media in microbatch order.

Module Contents

Classes

NameDescription
VlmMediaChunkOne microbatch’s views of staged image tensors.

Functions

NameDescription
next_vlm_media_chunkAdvance 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
)
Dataclass

One microbatch’s views of staged image tensors.

Pixels retain their processor layout, with images or patches on axis zero. Grids have shape [images, coordinates]. Empty media chunks become None; nonempty tensors alias the staged inputs unchanged.

image_grid
Tensor | None
pixel_values
Tensor | None
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 chunk lists. Pixel tensors retain their image/patch-major layouts; metadata is [images, coordinates].

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.