nemo_automodel.components.datasets.vlm.neat_packing_vlm

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

Packing is split into two phases:

Plan (instant) — scan raw dataset for estimated token lengths, run greedy_knapsack to assign samples to bins. No tokenization, no media loading.
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`	Construction-time configuration for the VLM neat-packing pipeline.
`PackedDatasetWrapper`	A Dataset that materializes packs lazily in `__getitem__`.
`PackedDatasetWrapperConfig`	Construction-time configuration for :class:`PackedDatasetWrapper`.

Functions

Name	Description
`_build_packed_vlm_sample`	Concatenate multiple shifted VLM samples into one packed sample.
`_compute_mrope_position_ids`	Compute mRoPE 3D position IDs for a single sample.
`_estimate_image_tokens`	Estimate token count for one image from its `[height, width]` metadata.
`_estimate_sample_length`	Estimate token count from raw conversation without tokenization.
`_estimate_video_tokens`	Estimate token count for one video from its
`_shift_sample`	Apply per-sample autoregressive shift before concatenation.
`greedy_knapsack_vt_balanced`	Pack samples with standard FFD, then interleave bins by VT for balance.
`neat_pack_dataset_vlm`	Create a lazily-packed VLM dataset.

Data

MEDIA_KEYS

logger

API

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.

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

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
) -> '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]] | NoneDefaults to None

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

get_rope_index

Callable[..., object] | NoneDefaults to None

Optional model.get_rope_index callable for mRoPE support.

processor

'ProcessorMixin | None'Defaults to None

Optional HuggingFace processor for accurate media token estimation.

class nemo_automodel.components.datasets.vlm.neat_packing_vlm.PackedDatasetWrapper(
    inner_dataset,
    bins: list[list[int]],
    pack_size: int,
    padding_idx: int = 0,
    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

intDefaults to 0

Token ID for padding.

get_rope_index

Callable | NoneDefaults to None

Optional model.get_rope_index for mRoPE.

max_retries

intDefaults to 10

Max retries when a sample fails to tokenize.

has_mrope

= get_rope_index is not None

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.

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

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

Wrap collate_fn with retry logic, delegating to inner dataset.

class nemo_automodel.components.datasets.vlm.neat_packing_vlm.PackedDatasetWrapperConfig(
    pack_size: int = 2048,
    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).

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.

get_rope_index

Callable[..., object] | NoneDefaults to None

Optional model.get_rope_index callable for mRoPE support.

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
) -> dict

Concatenate multiple shifted VLM samples into one packed sample.

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.

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.

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

boolDefaults to False

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

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.

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

Apply per-sample autoregressive shift before concatenation.

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.

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
) -> nemo_automodel.components.datasets.vlm.neat_packing_vlm.PackedDatasetWrapper

Create a lazily-packed VLM dataset.

Estimates token lengths from ds_raw (no tokenization).
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.
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

intDefaults to 0

Token ID for padding.

drop_long_samples

boolDefaults to False

Drop samples whose estimated length exceeds pack_size.

max_packs

int | NoneDefaults to None

Optional cap on number of packs.

get_rope_index

Callable | NoneDefaults to None

Optional model.get_rope_index for mRoPE.

ds_raw

Defaults to None

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

packing_ratio

floatDefaults to 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

Defaults to 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

boolDefaults to 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.

Returns: PackedDatasetWrapper

A PackedDatasetWrapper (torch Dataset).

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

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