nemo_automodel.components.models.deepseek_v41.packing

One packed text layout for direct forwards and contiguous context parallelism.

Module Contents

Classes

Name	Description
`PackedSequenceLayout`	Compression-aligned packed rows and reversible token coordinates.

Functions

Name	Description
`packed_layout`	Validate packed spans and align document starts for compression groups.

API

class nemo_automodel.components.models.deepseek_v41.packing.PackedSequenceLayout(
    source_indices: torch.Tensor,
    input_positions: torch.Tensor,
    position_ids: torch.Tensor,
    sequence_ids: torch.Tensor
)

Dataclass

Compression-aligned packed rows and reversible token coordinates.

input_positions

Tensor

position_ids

Tensor

sequence_ids

Tensor

source_indices

Tensor

nemo_automodel.components.models.deepseek_v41.packing.PackedSequenceLayout.pack(
    values: torch.Tensor,
    fill: int | float = 0
) -> torch.Tensor

Insert alignment padding without changing real token order.

Parameters:

values

torch.Tensor

Tensor of shape [batch, original_sequence, …], with arbitrary trailing axes.

fill

int | floatDefaults to 0

Value for padding slots.

Returns: torch.Tensor

Independent tensor [batch, padded_sequence, …], preserving gradients of real tokens.

nemo_automodel.components.models.deepseek_v41.packing.PackedSequenceLayout.restore(
    values: torch.Tensor
) -> torch.Tensor

Restore caller coordinates, returning zeros at original padding slots.

Parameters:

values

torch.Tensor

Tensor of shape [batch, padded_sequence, …], with arbitrary trailing axes.

Returns: torch.Tensor

Independent tensor [batch, original_sequence, …] retaining real-token gradients.

nemo_automodel.components.models.deepseek_v41.packing.packed_layout(
    seq_lens: torch.Tensor,
    seq_lens_padded: torch.Tensor | None,
    input_shape: tuple[int, int],
    alignment: int,
    minimum_length: int = 0
) -> nemo_automodel.components.models.deepseek_v41.packing.PackedSequenceLayout

Validate packed spans and align document starts for compression groups.

Parameters:

seq_lens

torch.Tensor

Integer real document lengths [batch, documents] or [documents] for batch one. Zero and -1000 entries denote absent documents.

seq_lens_padded

torch.Tensor | None

Optional integer physical span lengths with the same shape. These include existing per-document padding; omitted means documents are directly concatenated.

input_shape

tuple[int, int]

Original [batch, sequence] dimensions.

alignment

int

Positive multiple for each document’s physical span.

minimum_length

intDefaults to 0

Minimum aligned row length, used to retain a fixed pack budget.

Returns: PackedSequenceLayout

Layout whose tensor fields are documented by PackedSequenceLayout, on seq_lens.device.