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`	Pads a pack to `packed_sequence_size`.
`_should_stop_packing`	If max packs is set, stop packing when we reach that number.
`_split_and_add_pack`	Splits the current pack at the boundary, processes it, adds it to `packs`.
`_tensorize_and_pad_pack`	converts to tensors, pads a pack and returns it.
`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
) -> nemo_automodel.components.datasets.llm.packed_sequence.PACK_TYPE

Pads a pack to packed_sequence_size.

seq_lens contains original lengths. seq_lens_padded applies CP padding (if cp_size > 1) and pack-level padding.

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
) -> nemo_automodel.components.datasets.llm.packed_sequence.PACK_TYPE

Splits the current pack at the boundary, processes it, adds it to packs.

…and returns the start of the next pack.

TODO(@akoumparouli): refactor.

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

converts to tensors, pads a pack and returns it.

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,
    split,
    packed_sequence_size,
    max_packs = None,
    padding_idx = 0,
    drop_long_samples = True,
    cp_size = 1
)

Pack the dataset to defined length.

In particulat, 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

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

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