> 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.llm.packed_sequence

## Module Contents

### Functions

| Name                                                                                                                                               | Description                                                                                 |
| -------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------- |
| [`_convert_to_tensors`](#nemo_automodel-components-datasets-llm-packed_sequence-_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`](#nemo_automodel-components-datasets-llm-packed_sequence-_fill_labels_with_cross_entropy_ignore_idx) | -                                                                                           |
| [`_materialize_row`](#nemo_automodel-components-datasets-llm-packed_sequence-_materialize_row)                                                     | Return `dataset[i]` for the row's `__index__` (a `datasets.Dataset.map` worker).            |
| [`_pad_pack`](#nemo_automodel-components-datasets-llm-packed_sequence-_pad_pack)                                                                   | Pads a pack to `packed_sequence_size`.                                                      |
| [`_should_stop_packing`](#nemo_automodel-components-datasets-llm-packed_sequence-_should_stop_packing)                                             | If max packs is set, stop packing when we reach that number.                                |
| [`_split_and_add_pack`](#nemo_automodel-components-datasets-llm-packed_sequence-_split_and_add_pack)                                               | Splits the current pack at the boundary, processes it, adds it to `packs`.                  |
| [`_tensorize_and_pad_pack`](#nemo_automodel-components-datasets-llm-packed_sequence-_tensorize_and_pad_pack)                                       | converts to tensors, pads a pack and returns it.                                            |
| [`build_block_causal_additive_mask`](#nemo_automodel-components-datasets-llm-packed_sequence-build_block_causal_additive_mask)                     | Build a `[B, 1, T, T]` additive block-causal mask directly on `device`.                     |
| [`create_block_causal_mask`](#nemo_automodel-components-datasets-llm-packed_sequence-create_block_causal_mask)                                     | Creates causal mask block for specified lengths.                                            |
| [`pack_dataset`](#nemo_automodel-components-datasets-llm-packed_sequence-pack_dataset)                                                             | Pack the dataset to defined length.                                                         |
| [`packed_block_causal_mask`](#nemo_automodel-components-datasets-llm-packed_sequence-packed_block_causal_mask)                                     | Create a 2D block causal document mask for a batch of packed sequences.                     |
| [`tokenize_dataset_parallel`](#nemo_automodel-components-datasets-llm-packed_sequence-tokenize_dataset_parallel)                                   | Pre-tokenize a map-style dataset in parallel before packing.                                |

### Data

[`CROSS_ENTROPY_IGNORE_IDX`](#nemo_automodel-components-datasets-llm-packed_sequence-CROSS_ENTROPY_IGNORE_IDX)

[`PACK_TYPE`](#nemo_automodel-components-datasets-llm-packed_sequence-PACK_TYPE)

[`logger`](#nemo_automodel-components-datasets-llm-packed_sequence-logger)

### API

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

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

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

A single mapped row holding the integer `"__index__"`.

The map-style dataset to index.

**Returns:** `dict`

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

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

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

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

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

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

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

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

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

Actual dataset (can be 'train', 'val' or 'test')

Whether the dataset is 'train', 'val' or 'test'

Number of tokens in a pack

Maximum number of packs. Default: None

If True, drop samples that are longer than packed\_sequence\_size.

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

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

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.

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

A map-style dataset exposing `__len__` and `__getitem__`
whose items are tokenized dicts (`input_ids`, `labels` and,
optionally, `loss_mask`), each a `list[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.

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

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

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