nemo_automodel.components.datasets.llm.packed_sequence

Module Contents

Functions

Name	Description
`_convert_to_tensors`	Converts a pack into tensors. Pack comes in as a dict of lists and is converted to tensors.
`_fill_labels_with_cross_entropy_ignore_idx`	-
`_materialize_row`	Return `dataset[i]` for the row’s `__index__` (a `datasets.Dataset.map` worker).
`_pad_pack`	Pad an aligned pack to its fixed token budget.
`_should_stop_packing`	If max packs is set, stop packing when we reach that number.
`_split_and_add_pack`	Finalize complete documents and carry the last document into the next pack.
`_tensorize_and_pad_pack`	Tensorize and pad one pack using the layout documented by _pad_pack.
`build_block_causal_additive_mask`	Build a `[B, 1, T, T]` additive block-causal mask directly on `device`.
`create_block_causal_mask`	Creates causal mask block for specified lengths.
`pack_dataset`	Pack the dataset to defined length.
`packed_block_causal_mask`	Create a 2D block causal document mask for a batch of packed sequences.
`tokenize_dataset_parallel`	Pre-tokenize a map-style dataset in parallel before packing.

Data

CROSS_ENTROPY_IGNORE_IDX

PACK_TYPE

logger

API

nemo_automodel.components.datasets.llm.packed_sequence._convert_to_tensors(
    pack: nemo_automodel.components.datasets.llm.packed_sequence.PACK_TYPE
) -> nemo_automodel.components.datasets.llm.packed_sequence.PACK_TYPE

Converts a pack into tensors. Pack comes in as a dict of lists and is converted to tensors.

nemo_automodel.components.datasets.llm.packed_sequence._fill_labels_with_cross_entropy_ignore_idx(
    labels: list[int],
    loss_mask: list[int]
) -> list[int]

nemo_automodel.components.datasets.llm.packed_sequence._materialize_row(
    index_row: dict[str, int],
    dataset: torch.utils.data.Dataset
) -> dict

Return dataset[i] for the row’s __index__ (a datasets.Dataset.map worker).

The row is produced by the dataset’s own __getitem__, so any per-index behavior (e.g. a sample-skipping fallback) is preserved exactly rather than re-implemented.

Parameters:

index_row

dict[str, int]

A single mapped row holding the integer "__index__".

dataset

torch.utils.data.Dataset

The map-style dataset to index.

Returns: dict

The tokenized sample dict dataset[index_row["__index__"]].

nemo_automodel.components.datasets.llm.packed_sequence._pad_pack(
    pack: nemo_automodel.components.datasets.llm.packed_sequence.PACK_TYPE,
    padding_idx: int,
    packed_sequence_size: int,
    cross_entropy_ignore_idx: int = CROSS_ENTROPY_IGNORE_IDX,
    cp_size: int = 1,
    pad_to_multiple_of: int = 1
) -> nemo_automodel.components.datasets.llm.packed_sequence.PACK_TYPE

Pad an aligned pack to its fixed token budget.

Parameters:

pack

PACK_TYPE

Token tensors input_ids, labels and position_ids [tokens], plus real document lengths seq_lens [documents].

padding_idx

int

Input padding token.

packed_sequence_size

int

Physical token budget, including alignment padding.

cross_entropy_ignore_idx

intDefaults to CROSS_ENTROPY_IGNORE_IDX

Label value excluded from loss.

cp_size

intDefaults to 1

CP geometry used by generic THD packing.

pad_to_multiple_of

intDefaults to 1

Additional per-document alignment independent of CP.

Returns: PACK_TYPE

Tensor mapping with token fields [packed_sequence_size], real seq_lens

nemo_automodel.components.datasets.llm.packed_sequence._should_stop_packing(
    max_packs: int,
    packs: list[nemo_automodel.components.datasets.llm.packed_sequence.PACK_TYPE]
) -> bool

If max packs is set, stop packing when we reach that number.

nemo_automodel.components.datasets.llm.packed_sequence._split_and_add_pack(
    current_pack: nemo_automodel.components.datasets.llm.packed_sequence.PACK_TYPE,
    packs: list[nemo_automodel.components.datasets.llm.packed_sequence.PACK_TYPE],
    previous_sample_boundary: int,
    packed_sequence_size: int,
    padding_idx: int,
    cross_entropy_ignore_idx = CROSS_ENTROPY_IGNORE_IDX,
    cp_size: int = 1,
    pad_to_multiple_of: int = 1
) -> nemo_automodel.components.datasets.llm.packed_sequence.PACK_TYPE

Finalize complete documents and carry the last document into the next pack.

Parameters:

current_pack

PACK_TYPE

Lists of tokens [tokens] and real lengths [documents].

packs

list[PACK_TYPE]

Output list receiving one tensorized, padded pack.

previous_sample_boundary

int

Physical token offset before the last document.

packed_sequence_size

int

Fixed physical token budget, including alignment.

padding_idx

int

Input padding token.

cross_entropy_ignore_idx

Defaults to CROSS_ENTROPY_IGNORE_IDX

Padding label.

cp_size

intDefaults to 1

Generic THD context-parallel size.

pad_to_multiple_of

intDefaults to 1

Additional document alignment.

Returns: PACK_TYPE

List-valued token fields and lengths for the remaining document, preserving order.

nemo_automodel.components.datasets.llm.packed_sequence._tensorize_and_pad_pack(
    pack: nemo_automodel.components.datasets.llm.packed_sequence.PACK_TYPE,
    padding_idx: int,
    packed_sequence_size: int,
    cross_entropy_ignore_idx: int = CROSS_ENTROPY_IGNORE_IDX,
    cp_size: int = 1,
    pad_to_multiple_of: int = 1
) -> nemo_automodel.components.datasets.llm.packed_sequence.PACK_TYPE

Tensorize and pad one pack using the layout documented by _pad_pack.

Parameters:

pack

PACK_TYPE

Lists of token fields [tokens] and document lengths [documents].

padding_idx

int

Input padding token.

packed_sequence_size

int

Fixed physical token budget.

cross_entropy_ignore_idx

intDefaults to CROSS_ENTROPY_IGNORE_IDX

Padding label.

cp_size

intDefaults to 1

Generic THD context-parallel size.

pad_to_multiple_of

intDefaults to 1

Additional document alignment.

Returns: PACK_TYPE

Tensor mapping with token fields [packed_sequence_size] and lengths [documents].

nemo_automodel.components.datasets.llm.packed_sequence.build_block_causal_additive_mask(
    seq_lens: torch.Tensor,
    seq_length: int,
    dtype: torch.dtype,
    device: torch.device
) -> torch.Tensor

Build a [B, 1, T, T] additive block-causal mask directly on device.

In-document causal attention is allowed (0); cross-document and padding positions are finfo(dtype).min. seq_lens is the [B, max_docs] 0-padded per-document length tensor; each row’s non-zero entries sum to seq_length (trailing pad folded into the final document).

Raises ValueError if seq_lens is not 2D or any row does not sum to seq_length — a malformed row would otherwise silently misbucket the boundary cumsum (e.g. isolating the final token into a phantom document) and train the draft on wrong-context supervision with no error. This mirrors the fail-loud sum check the FlashAttention varlen path already performs in _seq_lens_to_cu_seqlens.

nemo_automodel.components.datasets.llm.packed_sequence.create_block_causal_mask(
    seq_lens: list[torch.Tensor]
) -> torch.Tensor

Creates causal mask block for specified lengths.

In particular, given a batch tensor of seq lens defining the lengths of samples in each pack, Construct a 2D block causal mask for each pack in the batch. For example, if a single sample’s seq_lens is [3, 2, 1], the mask would be:: mask = [ [1, 0, 0, 0, 0, 0], [1, 1, 0, 0, 0, 0], [1, 1, 1, 0, 0, 0], [0, 0, 0, 1, 0, 0], [0, 0, 0, 1, 1, 0], [0, 0, 0, 0, 0, 1], ]

Parameters:

seq_lens

List[torch.Tensor]

Sequence lengths of samples in each pack in the batch, shape (batch_size, n), where n is the max number of sequences in a pack and can vary across packs.

Returns: torch.Tensor

Block causal mask of shape (batch_size, packed_sequence_size, packed_sequence_size).

nemo_automodel.components.datasets.llm.packed_sequence.pack_dataset(
    dataset: torch.utils.data.Dataset | datasets.Dataset | datasets.DatasetDict | list[dict[str, list[int]]],
    split: str | None,
    packed_sequence_size: int,
    max_packs: int | None = None,
    padding_idx: int = 0,
    drop_long_samples: bool = True,
    cp_size: int = 1,
    pad_to_multiple_of: int = 1
) -> datasets.Dataset

Pack the dataset to defined length.

In particular, it will iterate through the dataset. Use a buffer to hold samples until packed_sequence_size, then append the buffer to packs as a single “packed” sample. Continue until max_packs or end of dataset.

Parameters:

dataset

torch.utils.data.Dataset | Dataset | DatasetDict | list[dict[str, list[int]]]

Actual dataset (can be ‘train’, ‘val’ or ‘test’)

split

str

Whether the dataset is ‘train’, ‘val’ or ‘test’

packed_sequence_size

int

Number of tokens in a pack

max_packs

intDefaults to None

Maximum number of packs. Default: None

drop_long_samples

boolDefaults to True

If True, drop samples that are longer than packed_sequence_size.

cp_size

intDefaults to 1

Context parallel size. When > 1, each sequence will be padded to be divisible by 2*cp_size for context parallel processing. Default: 1 (no CP).

pad_to_multiple_of

intDefaults to 1

Align each document’s physical span to this many tokens, including this padding in the pack budget. Real lengths and shifted labels are preserved. Default: 1. This alignment is independent of context-parallel size.

Returns: Dataset

Arrow dataset of fixed-length input_ids, already-shifted labels and position_ids

nemo_automodel.components.datasets.llm.packed_sequence.packed_block_causal_mask(
    seq_lens: list[torch.Tensor]
)

Create a 2D block causal document mask for a batch of packed sequences.

Parameters:

seq_lens

List[torch.Tensor]

Sequence lengths of samples in each pack in the batch, shape (batch_size, n), where n is the max number of sequences in a pack and can vary across packs.

Returns:

BlockMask or Tensor if torch version < 2.5.0.

nemo_automodel.components.datasets.llm.packed_sequence.tokenize_dataset_parallel(
    dataset: torch.utils.data.Dataset,
    num_proc: int
) -> datasets.Dataset

Pre-tokenize a map-style dataset in parallel before packing.

Iterating a lazily-tokenizing dataset (e.g. ChatDataset) tokenizes each sample serially in a single process, which dominates the cost of preparing a packed dataset. This helper front-loads that work by calling the dataset’s own __getitem__ over range(len(dataset)) with num_proc worker processes and returns a HuggingFace Dataset of the tokenized rows in original order. The result can be handed to :func:pack_dataset (or neat_pack_dataset) unchanged.

Caveats (the caller in build_dataloader gates on the first three):

dataset must be indexable (__len__ + __getitem__) and its __getitem__ must be a deterministic pure function of the index. Rows are produced in num_proc separate worker processes, so a __getitem__ that consults process-global state (RNG-driven augmentation, counters) would diverge from serial iteration.
Does not honor max_packs: the whole dataset is tokenized up front, whereas serial packing stops early once max_packs is reached.
dataset is pickled into every worker via fn_kwargs. This is cheap for the Arrow-backed datasets used here (self.dataset pickles as an mmap reference), but an in-memory (list-backed) dataset is copied once per worker, so size num_proc for the node’s memory as well as its CPUs.
Runs independently on each data-parallel rank, like the serial packing pass it replaces; dp_local_ranks * num_proc should not oversubscribe the node’s CPUs. Caching is disabled so every run re-tokenizes fresh.

Parameters:

dataset

torch.utils.data.Dataset

A map-style dataset exposing __len__ and __getitem__ whose items are tokenized dicts (input_ids, labels and, optionally, loss_mask), each a list[int].

num_proc

int

Number of worker processes used for tokenization. Must be > 1 for this helper to be worthwhile; the caller gates on that.

Returns: Dataset

A HuggingFace Dataset of the tokenized rows in original order.

nemo_automodel.components.datasets.llm.packed_sequence.CROSS_ENTROPY_IGNORE_IDX = -100

nemo_automodel.components.datasets.llm.packed_sequence.PACK_TYPE = dict[str, torch.Tensor | list[int]]

nemo_automodel.components.datasets.llm.packed_sequence.logger = logging.getLogger(__name__)