nemo_automodel.components.datasets.zephon.attention

Stateless mapping from Zephon tensors to AutoModel’s packed input contract.

Module Contents

Classes

Name	Description
`TrainingBatch`	Collated CPU tensors; AutoModel aliases are applied only when yielding.

Functions

Name	Description
`adapt_model_inputs`	Prepare attention membership from Zephon’s collated document IDs.
`enable_xdoc_attention`	Join attention and position ranges inside each packed batch document.
`require_supported_packing_contract`	Validate model input interfaces once, during loader construction.

API

class nemo_automodel.components.datasets.zephon.attention.TrainingBatch

Bases: typing.TypedDict

Collated CPU tensors; AutoModel aliases are applied only when yielding.

Token fields have shape [batch, sequence]. document_ids preserves subdocument membership (zero for padding); attention_document_ids also excludes processor-ineligible tokens and joins subdocuments when xdoc is on. Model positions are [batch, sequence] or [axes, batch, sequence] for mRoPE.

_packed_seq_ids

NotRequired[Tensor]

attention_document_ids

NotRequired[Tensor]

attention_mask

NotRequired[Tensor]

cu_seqlens

NotRequired[Tensor]

document_ids

NotRequired[Tensor]

input_ids

Tensor

labels

Tensor

max_seqlen

NotRequired[int]

packed_token_indices

NotRequired[Tensor]

padding_mask

Tensor

position_ids

Tensor

qkv_format

NotRequired[str]

seq_lens

NotRequired[Tensor]

seq_lens_padded

NotRequired[Tensor]

nemo_automodel.components.datasets.zephon.attention.adapt_model_inputs(
    batch: nemo_automodel.components.datasets.zephon.attention.TrainingBatch,
    attention_eligibility: torch.Tensor | None = None
) -> nemo_automodel.components.datasets.zephon.attention.TrainingBatch

Prepare attention membership from Zephon’s collated document IDs.

This follows to_training so IDs already match input tokens after shifting and padding. Zephon owns structural membership; processor eligibility and the consumer’s attention policy belong to this integration. AutoModel field names are applied later, when the loader yields the completed batch.

Parameters:

batch

TrainingBatch

Owned mapping with int64 document_ids [batch, sequence] from Zephon. Token/label/padding tensors have the same shape; model positions may additionally have a leading RoPE-axes dimension. Mutates mapping entries only; never modifies an existing tensor in place.

attention_eligibility

torch.Tensor | NoneDefaults to None

Optional bool [batch, sequence], supplied by the processor to exclude synthetic image tokens independently of padding.

Returns: TrainingBatch

The same mapping with attention_document_ids [batch, sequence]. It

nemo_automodel.components.datasets.zephon.attention.enable_xdoc_attention(
    batch: nemo_automodel.components.datasets.zephon.attention.TrainingBatch,
    packing_contract: nemo_automodel.components.datasets.packing.PackedSequenceContract
) -> nemo_automodel.components.datasets.zephon.attention.TrainingBatch

Join attention and position ranges inside each packed batch document.

Parameters:

batch

TrainingBatch

Mapped training tensors as documented by adapt_model_inputs; document_ids and attention_document_ids have shape [batch, sequence]; zero in the latter marks excluded positions. Model position_ids has shape [batch, sequence] or [position_axes, batch, sequence] for multi-axis RoPE. Subdocument coordinates are nonnegative and local to each subdocument. The owned mapping is updated in place, but no existing tensor is mutated.

packing_contract

PackedSequenceContract

Model-owned packed attention representation.

Returns: TrainingBatch

The same mapping with translated position IDs and joined attention IDs

nemo_automodel.components.datasets.zephon.attention.require_supported_packing_contract(
    contract: nemo_automodel.components.datasets.packing.PackedSequenceContract
) -> None

Validate model input interfaces once, during loader construction.

The text and image config builders call this before sources/workers start; the contract is fixed for that loader and is not rechecked for each batch.

return_packed_sequence_metadata specifically requests AutoModel’s gathered-token bundle: packed_token_indices [batch, sequence] (-1 at padding), cu_seqlens [batch, boundaries] (-1 at unused entries), and max_seqlen. These describe the stream after gathering non-padding tokens. AutoModel’s FA4 and recurrent/hybrid models use it for kernel layouts and recurrent-state resets. Zephon’s cu_seqlens instead describes physical offsets including padding. The final batch map computes the gathered-token offsets with AutoModel’s build_packed_sequence_metadata on the selected attention document IDs. Xdoc therefore joins both attention and recurrent-state boundaries; MTP uses those same selected boundaries. Keeping this bundle batch-major lets the model flatten it after PP microbatch splitting.