nemo_automodel.components.models.deepseek_v41.dspark_target

Packed DSpark target forward for DeepSeek V4.1.

Module Contents

Functions

Name	Description
`aligned_packed_layout`	Lay out packed rows with every document starting on a compression-group boundary.
`dspark_packed_forward`	Run a frozen DeepSeek V4.1 forward on packed rows for DSpark feature capture.

API

nemo_automodel.components.models.deepseek_v41.dspark_target.aligned_packed_layout(
    seq_lens: torch.Tensor,
    input_shape: tuple[int, int],
    alignment: int,
    sync_across_ranks: bool
) -> nemo_automodel.components.models.deepseek_v41.packing.PackedSequenceLayout

Lay out packed rows with every document starting on a compression-group boundary.

Parameters:

seq_lens

torch.Tensor

Integer document lengths of shape [batch, documents]; zero and -1000 entries denote absent documents.

input_shape

tuple[int, int]

Original [batch, sequence] dimensions; the aligned rows are at least sequence tokens wide.

alignment

int

Positive multiple for each document’s physical span. A compressed group must start where the document’s first group would start on its own, so this is the least common multiple of the active compression ratios.

sync_across_ranks

bool

Max-reduce the row length over the default process group. HybridEP needs the same token count from every rank, and rows of the same width need different amounts of alignment padding.

Returns: PackedSequenceLayout

Layout of the aligned rows; its tensors have shape [batch, aligned_sequence].

nemo_automodel.components.models.deepseek_v41.dspark_target.dspark_packed_forward(
    model: torch.nn.Module,
    input_ids: torch.Tensor,
    seq_lens: torch.Tensor
) -> nemo_automodel.components.models.deepseek_v41.packing.PackedSequenceLayout

Run a frozen DeepSeek V4.1 forward on packed rows for DSpark feature capture.

Documents are aligned as in the seq_lens path of the model’s forward, and the prepared packed_seq_ids keep every attention path inside its document. With the HybridEP dispatcher every rank pads to the longest aligned row of all ranks. DSpark reads hooked activations, so only the last position’s logits are projected.

Parameters:

model

nn.Module

DeepseekV41ForCausalLM target.

input_ids

torch.Tensor

Token IDs of shape [batch, sequence] holding packed documents.

seq_lens

torch.Tensor

Document lengths of shape [batch, documents]; each row sums to sequence and zero entries pad the document axis.

Returns: PackedSequenceLayout

Layout whose restore maps a tensor of shape [batch, aligned_sequence, …]