> 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.neat_packing_vlm

Neat packing for VLM (vision-language model) pre-tokenized datasets.

Packing is split into two phases:

1. **Plan** (instant) — scan raw dataset for estimated token lengths,
   run `greedy_knapsack` to assign samples to bins.  No tokenization,
   no media loading.
2. **Materialize** (lazy, in `__getitem__`) — when the DataLoader
   requests pack *k*, load + tokenize + shift + concat the samples
   assigned to bin *k*.  Runs in DataLoader worker processes, fully
   parallel.

This keeps the packing setup O(N) and lightweight, while the expensive
tokenization + media loading is distributed across `num_workers`.

## Module Contents

### Classes

| Name                                                                                                                | Description                                                        |
| ------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------ |
| [`NeatPackConfig`](#nemo_automodel-components-datasets-vlm-neat_packing_vlm-NeatPackConfig)                         | Construction-time configuration for the VLM neat-packing pipeline. |
| [`PackedDatasetWrapper`](#nemo_automodel-components-datasets-vlm-neat_packing_vlm-PackedDatasetWrapper)             | A Dataset that materializes packs lazily in `__getitem__`.         |
| [`PackedDatasetWrapperConfig`](#nemo_automodel-components-datasets-vlm-neat_packing_vlm-PackedDatasetWrapperConfig) | Construction-time configuration for :class:`PackedDatasetWrapper`. |

### Functions

| Name                                                                                                                  | Description                                                                |
| --------------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------- |
| [`_aligned_length`](#nemo_automodel-components-datasets-vlm-neat_packing_vlm-_aligned_length)                         | Round `length` up to `alignment` without changing zero-length samples.     |
| [`_build_packed_vlm_sample`](#nemo_automodel-components-datasets-vlm-neat_packing_vlm-_build_packed_vlm_sample)       | Concatenate shifted VLM samples, including per-document alignment padding. |
| [`_compute_mrope_position_ids`](#nemo_automodel-components-datasets-vlm-neat_packing_vlm-_compute_mrope_position_ids) | Compute mRoPE 3D position IDs for a single sample.                         |
| [`_estimate_image_tokens`](#nemo_automodel-components-datasets-vlm-neat_packing_vlm-_estimate_image_tokens)           | Estimate token count for one image from its `[height, width]` metadata.    |
| [`_estimate_sample_length`](#nemo_automodel-components-datasets-vlm-neat_packing_vlm-_estimate_sample_length)         | Estimate token count from raw conversation without tokenization.           |
| [`_estimate_video_tokens`](#nemo_automodel-components-datasets-vlm-neat_packing_vlm-_estimate_video_tokens)           | Estimate token count for one video from its                                |
| [`_shift_sample`](#nemo_automodel-components-datasets-vlm-neat_packing_vlm-_shift_sample)                             | Apply per-sample autoregressive shift before concatenation.                |
| [`greedy_knapsack_vt_balanced`](#nemo_automodel-components-datasets-vlm-neat_packing_vlm-greedy_knapsack_vt_balanced) | Pack samples with standard FFD, then interleave bins by VT for balance.    |
| [`neat_pack_dataset_vlm`](#nemo_automodel-components-datasets-vlm-neat_packing_vlm-neat_pack_dataset_vlm)             | Create a lazily-packed VLM dataset.                                        |

### Data

[`MEDIA_KEYS`](#nemo_automodel-components-datasets-vlm-neat_packing_vlm-MEDIA_KEYS)

[`logger`](#nemo_automodel-components-datasets-vlm-neat_packing_vlm-logger)

### API

```python
class nemo_automodel.components.datasets.vlm.neat_packing_vlm.NeatPackConfig(
    pack_size: int = 2048,
    drop_long_samples: bool = False,
    max_packs: int | None = None,
    packing_ratio: float = 1.0,
    balance_media_tokens: bool = True,
    collate_max_length: int | None = None,
    attn_implementation: str | None = None,
    packing_format: typing.Literal['neat', 'thd'] = 'neat'
)
```

Dataclass

Construction-time configuration for the VLM neat-packing pipeline.

**`attn_implementation`** `str | None = None`

Optional packed-mask backend override used only with context parallelism.

---

**`balance_media_tokens`** `bool = True`

If `True`, use VT-balanced knapsack to distribute visual tokens evenly.

---

**`collate_max_length`** `int | None = None`

Optional maximum padded length used by the packed collator.

---

**`drop_long_samples`** `bool = False`

If `True`, samples whose estimated length exceeds `pack_size` are dropped.

---

**`max_packs`** `int | None = None`

Optional cap on the total number of packs produced.

---

**`pack_size`** `int = 2048`

Target packed sequence length (after autoregressive shift).

---

**`packing_format`** `Literal['neat', 'thd'] = 'neat'`

Packed collator format. `thd` emits Transformer Engine sequence metadata.

---

**`packing_ratio`** `float = 1.0`

Knapsack bin fill ratio. Values below 1.0 leave headroom for estimation errors.

---

```python
nemo_automodel.components.datasets.vlm.neat_packing_vlm.NeatPackConfig.__post_init__() -> None
```

```python
nemo_automodel.components.datasets.vlm.neat_packing_vlm.NeatPackConfig.build(
    dataset: 'PreTokenizedDatasetWrapper',
    padding_idx: int,
    ds_raw: torch.utils.data.Dataset | collections.abc.Sequence[dict[str, object]] | None = None,
    get_rope_index: typing.Callable[..., object] | None = None,
    processor: 'ProcessorMixin | None' = None,
    cp_size: int = 1
) -> 'PackedDatasetWrapper'
```

Build a neat-packed VLM dataset from this config.

**Parameters:**

**`dataset`** `'PreTokenizedDatasetWrapper'`

`PreTokenizedDatasetWrapper` for per-sample tokenization.

---

**`padding_idx`** `int`

Runtime tokenizer padding token ID.

---

**`ds_raw`** `torch.utils.data.Dataset | Sequence[dict[str, object]] | None` — default: None

Raw conversations dataset for fast length estimation. Falls back to
`len(dataset)` when `None`.

---

**`get_rope_index`** `Callable[..., object] | None` — default: None

Optional `model.get_rope_index` callable for mRoPE support.

---

**`processor`** `'ProcessorMixin | None'` — default: None

Optional HuggingFace processor for accurate media token estimation.

---

**`cp_size`** `int` — default: 1

Runtime context-parallel size. THD packing aligns every
document to `2 * cp_size` when greater than one.

---

```python
class nemo_automodel.components.datasets.vlm.neat_packing_vlm.PackedDatasetWrapper(
    inner_dataset,
    bins: list[list[int]],
    pack_size: int,
    padding_idx: int = 0,
    sequence_alignment: int = 1,
    get_rope_index: typing.Callable | None = None,
    max_retries: int = 10
)
```

**Bases:** `Dataset`

A Dataset that materializes packs lazily in `__getitem__`.

The constructor only stores bin assignments (which sample indices go
into each pack).  The actual tokenization, media loading, shift, and
concatenation happen when a pack is requested — inside DataLoader
worker processes, fully parallel.

**Parameters:**

**`inner_dataset`**

The `PreTokenizedDatasetWrapper` that tokenizes
individual samples.

---

**`bins`** `list[list[int]]`

List of bins from `greedy_knapsack`, where each bin is a
list of sample indices into `inner_dataset`.

---

**`pack_size`** `int`

Target packed sequence length (after shift).

---

**`padding_idx`** `int` — default: 0

Token ID for padding.

---

**`sequence_alignment`** `int` — default: 1

Per-document token alignment included in capacity
checks and materialization.

---

**`get_rope_index`** `Callable | None` — default: None

Optional `model.get_rope_index` for mRoPE.

---

**`max_retries`** `int` — default: 10

Max retries when a sample fails to tokenize.

---

**`has_mrope`** `= get_rope_index is not None`

---

```python
nemo_automodel.components.datasets.vlm.neat_packing_vlm.PackedDatasetWrapper.__getitem__(
    pack_idx: int
) -> dict
```

Materialize one pack: tokenize + shift + concat all samples in the bin.

```python
nemo_automodel.components.datasets.vlm.neat_packing_vlm.PackedDatasetWrapper.__len__()
```

```python
nemo_automodel.components.datasets.vlm.neat_packing_vlm.PackedDatasetWrapper.robust_collate(
    collate_fn
)
```

Wrap collate\_fn with retry logic, delegating to inner dataset.

```python
class nemo_automodel.components.datasets.vlm.neat_packing_vlm.PackedDatasetWrapperConfig(
    pack_size: int = 2048,
    sequence_alignment: int = 1,
    max_retries: int = 10
)
```

Dataclass

Construction-time configuration for :class:`PackedDatasetWrapper`.

**`max_retries`** `int = 10`

Max retries when a sample fails to tokenize during materialization.

---

**`pack_size`** `int = 2048`

Target packed sequence length (after autoregressive shift).

---

**`sequence_alignment`** `int = 1`

Per-document token alignment applied during materialization.

---

```python
nemo_automodel.components.datasets.vlm.neat_packing_vlm.PackedDatasetWrapperConfig.build(
    inner_dataset: 'PreTokenizedDatasetWrapper',
    bins: list[list[int]],
    padding_idx: int,
    get_rope_index: typing.Callable[..., object] | None = None
) -> 'PackedDatasetWrapper'
```

Build a :class:`PackedDatasetWrapper` from this config.

**Parameters:**

**`inner_dataset`** `'PreTokenizedDatasetWrapper'`

The tokenizing dataset (e.g. `PreTokenizedDatasetWrapper`).

---

**`bins`** `list[list[int]]`

Bin assignments from `greedy_knapsack` / `neat_pack_dataset_vlm`.

---

**`padding_idx`** `int`

Runtime tokenizer padding token ID.

---

**`sequence_alignment`**

Per-document token alignment from this config.

---

**`get_rope_index`** `Callable[..., object] | None` — default: None

Optional `model.get_rope_index` callable for mRoPE support.

---

```python
nemo_automodel.components.datasets.vlm.neat_packing_vlm._aligned_length(
    length: int,
    alignment: int
) -> int
```

Round `length` up to `alignment` without changing zero-length samples.

```python
nemo_automodel.components.datasets.vlm.neat_packing_vlm._build_packed_vlm_sample(
    samples: list[dict],
    pack_size: int,
    padding_idx: int,
    has_mrope: bool = False,
    sequence_alignment: int = 1
) -> dict
```

Concatenate shifted VLM samples, including per-document alignment padding.

```python
nemo_automodel.components.datasets.vlm.neat_packing_vlm._compute_mrope_position_ids(
    sample: dict,
    get_rope_index: typing.Callable
) -> torch.Tensor | None
```

Compute mRoPE 3D position IDs for a single sample.

**Parameters:**

**`sample`** `dict`

Pretokenized sample dict holding `input_ids` of shape `[seq]`
(or `[1, seq]`) plus optional `image_grid_thw`/`video_grid_thw`/
`attention_mask`/`mm_token_type_ids` tensors mirroring it.

---

**`get_rope_index`** `Callable`

Bound `model.get_rope_index` callable. When its
signature requires `mm_token_type_ids` and the sample lacks it, the
tensor is rebuilt from the bound model config's image/video token ids.

---

**Returns:** `torch.Tensor | None`

Position ids of shape `[3, seq_len]`, or `None` if not applicable.

```python
nemo_automodel.components.datasets.vlm.neat_packing_vlm._estimate_image_tokens(
    img_meta,
    image_cfg: dict
) -> int
```

Estimate token count for one image from its `[height, width]` metadata.

```python
nemo_automodel.components.datasets.vlm.neat_packing_vlm._estimate_sample_length(
    example: dict,
    image_cfg: dict | None = None,
    video_cfg: dict | None = None,
    return_media_tokens: bool = False
) -> int | tuple[int, int]
```

Estimate token count from raw conversation without tokenization.

Uses pre-computed `_text_tokens` (from `precompute_tokens.py`) when
available, otherwise falls back to `chars // 3`.  Media tokens are
estimated via `smart_resize` when processor configs are provided,
otherwise falls back to 500 per media item.

**Parameters:**

**`return_media_tokens`** `bool` — default: False

If True, return `(total_tokens, media_tokens)`
instead of just `total_tokens`.

---

```python
nemo_automodel.components.datasets.vlm.neat_packing_vlm._estimate_video_tokens(
    vid_meta,
    video_cfg: dict
) -> int
```

Estimate token count for one video from its
`[total_frames, height, width, fps, duration]` metadata.

```python
nemo_automodel.components.datasets.vlm.neat_packing_vlm._shift_sample(
    sample: dict,
    has_mrope: bool = False
) -> dict
```

Apply per-sample autoregressive shift before concatenation.

```python
nemo_automodel.components.datasets.vlm.neat_packing_vlm.greedy_knapsack_vt_balanced(
    lengths: list[int],
    max_length: int,
    visual_tokens: list[int]
) -> list[list[int]]
```

Pack samples with standard FFD, then interleave bins by VT for balance.

Uses the standard greedy knapsack (FFD) for optimal packing efficiency,
then reorders bins so that consecutive packs have similar visual token
counts.  This ensures data-parallel ranks in the same training step
process packs with comparable VIT workload, reducing straggler effects.

**Parameters:**

**`lengths`** `list[int]`

Total token length (text + media) per sample.

---

**`max_length`** `int`

Maximum capacity per pack.

---

**`visual_tokens`** `list[int]`

Number of media tokens per sample.

---

**Returns:** `list[list[int]]`

A list of bins, where each bin is a list of sample indices.

```python
nemo_automodel.components.datasets.vlm.neat_packing_vlm.neat_pack_dataset_vlm(
    dataset,
    pack_size: int,
    padding_idx: int = 0,
    drop_long_samples: bool = False,
    max_packs: int | None = None,
    get_rope_index: typing.Callable | None = None,
    ds_raw = None,
    packing_ratio: float = 1.0,
    processor = None,
    balance_media_tokens: bool = True,
    sequence_alignment: int = 1
) -> nemo_automodel.components.datasets.vlm.neat_packing_vlm.PackedDatasetWrapper
```

Create a lazily-packed VLM dataset.

1. Estimates token lengths from `ds_raw` (no tokenization).
2. Runs knapsack to assign samples to bins.  When
   `balance_media_tokens=True` (default), uses a two-phase
   algorithm that balances visual token counts across packs,
   reducing VIT compute/memory imbalance and straggler effects.
3. Returns a `PackedDatasetWrapper` whose `__getitem__` tokenizes
   and builds packs on-the-fly in DataLoader workers.

**Parameters:**

**`dataset`**

`PreTokenizedDatasetWrapper` for per-sample tokenization.

---

**`pack_size`** `int`

Target packed sequence length (after shift).

---

**`padding_idx`** `int` — default: 0

Token ID for padding.

---

**`drop_long_samples`** `bool` — default: False

Drop samples whose estimated length exceeds
`pack_size`.

---

**`max_packs`** `int | None` — default: None

Optional cap on number of packs.

---

**`get_rope_index`** `Callable | None` — default: None

Optional `model.get_rope_index` for mRoPE.

---

**`ds_raw`** — default: None

Raw dataset (conversations) for fast length estimation.
Falls back to `len(dataset)` if not provided.

---

**`packing_ratio`** `float` — default: 1.0

Fill ratio for knapsack bins (default 1.0).
E.g. `0.9` means knapsack only fills bins to `pack_size * 0.9`,
leaving 10% headroom to absorb estimation errors.  This reduces
overflow drops at `__getitem__` time.  The actual `pack_size`
is still used as the hard limit.

---

**`processor`** — default: None

Optional HuggingFace processor (e.g. `Qwen2VLProcessor`).
Used to extract `image_processor` / `video_processor` configs
for accurate media token estimation via `smart_resize`.

---

**`balance_media_tokens`** `bool` — default: True

If True (default), use VT-balanced knapsack
that distributes visual tokens evenly across packs.  Falls back
to standard knapsack if no media tokens are detected.

---

**`sequence_alignment`** `int` — default: 1

Per-document token alignment included in both
knapsack capacity and lazy materialization.

---

**Returns:** `PackedDatasetWrapper`

A `PackedDatasetWrapper` (torch Dataset).

```python
nemo_automodel.components.datasets.vlm.neat_packing_vlm.MEDIA_KEYS = ('pixel_values', 'image_grid_thw', 'image_position_ids', 'pixel_values_videos', ...
```

```python
nemo_automodel.components.datasets.vlm.neat_packing_vlm.logger = logging.getLogger(__name__)
```