nemo_automodel.components.datasets.packed_seq

Typed packed-sequence metadata for neat sequence packing.

NeMo AutoModel’s neat packing encodes document boundaries in an indexed map [B, S] where each position holds the 1-based document index it belongs to (0 = padding), e.g. [1, 1, 2, 2, 2, 0]. This module turns that map into typed metadata that drives HuggingFace’s varlen FlashAttention path via transformers.modeling_flash_attention_utils.FlashAttentionKwargs (cu_seq_lens_q/cu_seq_lens_k/max_length_q/max_length_k).

The varlen path reshapes [B, S, ...] to [B * S, ...] without unpadding (see the is_fa_with_varlen_kwargs branch in transformers.modeling_flash_attention_utils._flash_attention_forward). cu_seqlens therefore spans the whole flattened batch, including padding runs, and a boundary is forced at every row start so a document that fills one row cannot merge with the first document of the next row once the batch is flattened.

Module Contents

Classes

Name	Description
`PackedSeqParams`	Packed-sequence metadata derived from a neat-packing document map.

Functions

Name	Description
`packed_seq_params_from_doc_ids`	Build packed-sequence metadata from a neat-packing document map.
`to_flash_attention_kwargs`	Convert packed-sequence params to HuggingFace `FlashAttentionKwargs`.

API

class nemo_automodel.components.datasets.packed_seq.PackedSeqParams(
    cu_seqlens: torch.Tensor,
    max_seqlen: int,
    doc_ids: torch.Tensor
)

Dataclass

Packed-sequence metadata derived from a neat-packing document map.

cu_seqlens

Tensor

Int32 tensor of shape [num_segments + 1]. Cumulative token counts over the flattened [batch * sequence] axis, including trailing padding runs, so cu_seqlens[-1] == batch * sequence. Segment k spans flattened positions [cu_seqlens[k], cu_seqlens[k + 1]).

doc_ids

Tensor

Int tensor of shape [batch, sequence] with the 1-based document index per token (0 = padding). Preserved unchanged for loss functions and context-parallel consumers.

max_seqlen

int

Length of the longest segment, sized for the varlen flash-attention kernel.

nemo_automodel.components.datasets.packed_seq.packed_seq_params_from_doc_ids(
    doc_ids: torch.Tensor
) -> nemo_automodel.components.datasets.packed_seq.PackedSeqParams

Build packed-sequence metadata from a neat-packing document map.

Segments are maximal runs of an identical document index within a row, plus a forced boundary at every row start. Padding runs (index 0) become their own segments so cu_seqlens covers the full flattened [batch * sequence] axis that the varlen flash-attention path consumes without unpadding. Padding tokens therefore attend only among themselves, which is harmless because their labels are ignored by the loss.

Parameters:

doc_ids

torch.Tensor

Int tensor of shape [batch, sequence] where each position holds the 1-based document index it belongs to (0 = padding).

Returns: PackedSeqParams

PackedSeqParams carrying cu_seqlens, max_seqlen, and the original

nemo_automodel.components.datasets.packed_seq.to_flash_attention_kwargs(
    params: nemo_automodel.components.datasets.packed_seq.PackedSeqParams
) -> dict[str, torch.Tensor | int]

Convert packed-sequence params to HuggingFace FlashAttentionKwargs.

The returned mapping matches transformers FlashAttentionKwargs so a packed batch drives flash_attn_varlen_func directly. Query and key share one layout (self-attention), so the q and k entries are equal.

Parameters:

params

PackedSeqParams

Packed-sequence metadata. params.cu_seqlens is an int32 tensor of shape [num_segments + 1].

Returns: dict[str, torch.Tensor | int]

Mapping with cu_seq_lens_q/cu_seq_lens_k (the int32 cu_seqlens