nemo_automodel.components.datasets.vlm.pp_media

Module Contents

Functions

Name	Description
`_select_image_grid`	-
`_will_run_forward_metadata_inference`	Return whether the next schedule step will probe stage 0 with a real forward.
`chunk_step3_media`	Chunk image tensors with per-sample patch counts for PP microbatches.
`chunk_vlm_media`	Split VLM pixel values and media metadata into PP microbatch chunks.
`prepare_vlm_media_for_pp`	Move VLM media tensors into pre-chunked PP media storage on the batch.
`stage_vlm_media_for_pp`	Attach dataloader-prepared VLM media chunks to PP stage 0 for one schedule call.
`wrap_vlm_collate_for_pp`	Wrap a VLM collate function so it prepares media tensors for PP.

Data

VLM_PP_MEDIA_KEY

_VLM_MEDIA_KEYS

API

nemo_automodel.components.datasets.vlm.pp_media._select_image_grid(
    image_grid_hws: torch.Tensor | None,
    image_grid_thw: torch.Tensor | None,
    image_sizes: torch.Tensor | None,
    image_position_ids: torch.Tensor | None
) -> torch.Tensor | None

nemo_automodel.components.datasets.vlm.pp_media._will_run_forward_metadata_inference(
    pp: typing.Any
) -> bool

Return whether the next schedule step will probe stage 0 with a real forward.

Both the legacy _configure_outputs_meta API and the newer populated _user_meta API provide analytical stage metadata, so neither needs a forward probe. An unpopulated _user_meta still uses dynamic inference.

nemo_automodel.components.datasets.vlm.pp_media.chunk_step3_media(
    pixel_values: torch.Tensor,
    batch_size: int,
    n_microbatches: int,
    num_patches: torch.Tensor | None = None,
    patch_pixel_values: torch.Tensor | None = None,
    patch_newline_mask: torch.Tensor | None = None
) -> dict[str, list[torch.Tensor]]

Chunk image tensors with per-sample patch counts for PP microbatches.

Step3 processors emit one full image per sample in pixel_values and a flat list of optional crop patches in patch_pixel_values. num_patches maps samples to the flat patch tensor. Processors may also emit pixel_values itself as a flat patch tensor using the same mapping.

Parameters:

pixel_values

torch.Tensor

Either [batch, ...] (one image tensor per sample) or a row-flattened [total_patches, ...] tensor whose patches are concatenated along axis 0 in sample order, where total_patches == sum(num_patches).

batch_size

int

Number of samples in the batch.

n_microbatches

int

Number of PP microbatches to split the batch into.

num_patches

torch.Tensor | NoneDefaults to None

Tensor of shape [batch]; entry i is the patch count for sample i. Defaults to all-zeros when the processor emits no crop patches.

patch_pixel_values

torch.Tensor | NoneDefaults to None

Optional row-flattened crop patches of shape [total_patches, ...] indexed by num_patches.

patch_newline_mask

torch.Tensor | NoneDefaults to None

Optional tensor of shape [total_patches] marking patch-row newline positions, indexed by num_patches.

Returns: dict[str, list[torch.Tensor]]

dict[str, list[torch.Tensor]]: Per-microbatch slices keyed by

nemo_automodel.components.datasets.vlm.pp_media.chunk_vlm_media(
    pixel_values: torch.Tensor | list[torch.Tensor],
    image_grid: torch.Tensor,
    batch_size: int,
    n_microbatches: int,
    n_images_per_sample: torch.Tensor | None = None
) -> tuple[list[torch.Tensor | list[torch.Tensor]], list[torch.Tensor]]

Split VLM pixel values and media metadata into PP microbatch chunks.

Explicit image counts take precedence over a coincidental equality between the number of images and documents. Text-only documents own empty chunks.

Parameters:

pixel_values

torch.Tensor | list[torch.Tensor]

Image-major pixels [images, channels, height, width], padded patches [images, max_patches, features], flat patches [total_patches, features] or [total_patches, channels, patch_height, patch_width], or a list of [channels, height, width] tensors with possibly different image sizes.

image_grid

torch.Tensor

Per-image metadata [images, coordinates]. For flat patches, the product of each grid gives the corresponding patch count.

batch_size

int

Number of batch documents owning the media.

n_microbatches

int

Number of document-axis chunks used by the PP schedule.

n_images_per_sample

torch.Tensor | NoneDefaults to None

Optional image counts [batch]. Without counts, exactly one image per document is assumed.

Returns: list[torch.Tensor | list[torch.Tensor]]

Pixel and metadata lists in microbatch order, retaining each input’s

nemo_automodel.components.datasets.vlm.pp_media.prepare_vlm_media_for_pp(
    batch: collections.abc.MutableMapping[str, typing.Any],
    batch_size: int,
    n_microbatches: int
) -> collections.abc.MutableMapping[str, typing.Any]

Move VLM media tensors into pre-chunked PP media storage on the batch.

This is intended to run from VLM collate/dataloader code when PP is enabled. The returned batch no longer carries raw media tensors that PyTorch PP would chunk by row incorrectly; instead it carries VLM_PP_MEDIA_KEY with per-microbatch media chunks.

nemo_automodel.components.datasets.vlm.pp_media.stage_vlm_media_for_pp(
    pp: typing.Any,
    model_parts: list[torch.nn.Module],
    batch: collections.abc.MutableMapping[str, typing.Any]
)

Attach dataloader-prepared VLM media chunks to PP stage 0 for one schedule call.

nemo_automodel.components.datasets.vlm.pp_media.wrap_vlm_collate_for_pp(
    collate_fn: collections.abc.Callable[[Any], collections.abc.MutableMapping[str, typing.Any]],
    n_microbatches: int
) -> collections.abc.Callable[[Any], collections.abc.MutableMapping[str, typing.Any]]

Wrap a VLM collate function so it prepares media tensors for PP.

nemo_automodel.components.datasets.vlm.pp_media.VLM_PP_MEDIA_KEY = '_vlm_pp_media_chunks'

nemo_automodel.components.datasets.vlm.pp_media._VLM_MEDIA_KEYS = ('pixel_values', 'patch_pixel_values', 'num_patches', 'patch_newline_mask', 'ima...