> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.components.datasets.vlm.pp_media

## Module Contents

### Functions

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

### Data

[`VLM_PP_MEDIA_KEY`](#nemo_automodel-components-datasets-vlm-pp_media-VLM_PP_MEDIA_KEY)

[`_VLM_MEDIA_KEYS`](#nemo_automodel-components-datasets-vlm-pp_media-_VLM_MEDIA_KEYS)

[`__all__`](#nemo_automodel-components-datasets-vlm-pp_media-__all__)

### API

```python
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
```

```python
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.

```python
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 | None` — default: 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 | None` — default: None

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

---

**`patch_newline_mask`** `torch.Tensor | None` — default: 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

```python
nemo_automodel.components.datasets.vlm.pp_media.chunk_vlm_media(
    pixel_values: 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]]
```

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

Handles four layouts:

1. `[N, C, H, W]` with `N == batch_size` -- one full image per sample.
2. `[N, max_patches, D]` with `N == batch_size` -- padded patches per image.
3. Flat patches `[total_patches, D]` with per-sample media counts from
   `n_images_per_sample`.
4. Flat patches with `n_images == batch_size` -- legacy one-image-per-sample.

```python
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.

```python
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.

```python
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.

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

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

```python
nemo_automodel.components.datasets.vlm.pp_media.__all__ = ['VLM_PP_MEDIA_KEY', 'chunk_vlm_media', 'prepare_vlm_media_for_pp', 'stage_vlm_m...
```