nemo_automodel.components.models.common.packing

Model-derived packing contracts and native metadata normalization.

HF flash attention consumes the public varlen kwargs built by dataset collaters. Native FA4 and recurrent consumers use separate unpadded document metadata. Neither path modifies Transformers globals.

Module Contents

Classes

Name	Description
`AttentionBackendSelection`	Typed attention selection carried by a built custom model.
`DocumentAttentionMaskConsumer`	Model that composes document eligibility with its own attention policy.
`NativeFA4Consumer`	Model whose attention layers invoke Automodel’s native FA4 callable.
`PackedAttentionBackendValidator`	Model with additional constraints beyond a backend’s generic contract.
`PackedMaskConsumer`	Model that owns masking and consumes compact document IDs.
`PackedMetadataEntryOwner`	Model whose attention policy needs metadata at a specific forward entry.
`PackingCapabilities`	Model-derived loader requests and native attention setup information.
`PackingMetadataConsumer`	Consumer of AutoModel’s gathered-token indices and cumulative lengths.
`_HFAttentionConsumer`	Model whose active attention path dispatches through its HF config.

Functions

Name	Description
`_allows_ordinary_causal_attention`	Whether each [B,S] attention-ID map is one eligible prefix plus padding.
`_flatten_packed_metadata_at_model_entry`	Flatten packed metadata once after pipeline microbatch splitting.
`_install_native_fa4_metadata_hook`	Install the once-per-forward native FA4 metadata normalizer.
`_model_attn_implementation`	Return the packing-relevant attention backend an already-built model runs with.
`_prepare_flash_metadata_before_input_transfer`	Prepare host offsets before FSDP moves the document map to CUDA.
`_prepare_packed_attention_at_model_entry`	Materialize the selected document segmentation after PP splitting.
`_strip_unused_packed_metadata_at_model_entry`	Keep dataset-only NEAT offsets away from a legacy THD consumer.
`_validate_model_consumes_packed_contract`	Raise if `model`’s `forward` cannot receive the typed packing contract.
`configure_packing`	Configure the model consumer and return its dataset packing contract.
`configure_packing_for_models`	Validate and install the adapters for batched packed execution.
`flatten_packed_sequence_metadata`	Normalize batch-major packed metadata to a single flat token stream.
`get_model_attn_implementation`	Return the NEAT/BSHD attention implementation used by a built model.
`get_packing_capabilities`	Describe loader requests without configuring or validating packed execution.
`get_packing_capabilities_for_models`	Combine loader return requests across local model or pipeline stages.
`is_indexed_packed_mask`	Return `True` iff `attention_mask` is an Automodel-style indexed packing mask.
`validate_flash_packing_support`	Fail loudly when flash neat packing cannot be consumed by the backend or model.
`validate_thd_packing_support`	Reject THD inputs before loader construction if a local stage cannot consume them.

Data

PackedMaskType

_FLASH_ATTN_IMPLEMENTATIONS

logger

API

class nemo_automodel.components.models.common.packing.AttentionBackendSelection()

Protocol

Typed attention selection carried by a built custom model.

attn

AttentionBackend

Selected native attention backend.

class nemo_automodel.components.models.common.packing.DocumentAttentionMaskConsumer()

Protocol

Model that composes document eligibility with its own attention policy.

nemo_automodel.components.models.common.packing.DocumentAttentionMaskConsumer.prepare_document_attention_mask(
    attention_mask: torch.Tensor | dict[str, torch.Tensor] | None,
    mm_token_type_ids: torch.Tensor | None,
    document_ids: torch.Tensor | None
) -> torch.Tensor | dict[str, torch.Tensor] | None

Compose selected document attention with a model-owned policy.

Parameters:

attention_mask

torch.Tensor | dict[str, torch.Tensor] | None

Integer selected IDs [batch, sequence], a prepared mask [batch, 1, sequence, sequence], a mapping of such masks, or None for native masking.

mm_token_type_ids

torch.Tensor | None

Optional modality IDs [batch, sequence].

document_ids

torch.Tensor | None

Optional original membership [batch, sequence], used only by this model policy rather than attention/MTP segmentation.

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

A prepared mask or mapping of masks [batch, 1, sequence, sequence],

class nemo_automodel.components.models.common.packing.NativeFA4Consumer()

Protocol

Model whose attention layers invoke Automodel’s native FA4 callable.

_uses_native_fa4

bool

class nemo_automodel.components.models.common.packing.PackedAttentionBackendValidator()

Protocol

Model with additional constraints beyond a backend’s generic contract.

nemo_automodel.components.models.common.packing.PackedAttentionBackendValidator.validate_packed_attention_backend(
    attn_implementation: str
) -> None

Reject a backend that cannot preserve this model’s packed policy.

class nemo_automodel.components.models.common.packing.PackedMaskConsumer()

Protocol

Model that owns masking and consumes compact document IDs.

packed_mask_type

PackedMaskType

class nemo_automodel.components.models.common.packing.PackedMetadataEntryOwner()

Protocol

Model whose attention policy needs metadata at a specific forward entry.

packed_metadata_entry_module

Module

Return the called module that owns packed attention preparation.

class nemo_automodel.components.models.common.packing.PackingCapabilities(
    packed_mask_type: nemo_automodel.components.models.common.packing.PackedMaskType,
    return_packed_sequence_metadata: bool = False,
    uses_native_fa4: bool = False,
    return_document_ids: bool = False,
    max_sequence_length: int | None = None
)

Dataclass

Model-derived loader requests and native attention setup information.

packed_mask_type and return_packed_sequence_metadata apply only to batched packed inputs. THD derives kernel metadata from lengths instead. return_document_ids requests original membership in either format. uses_native_fa4 is model setup information, not an emitted batch field. max_sequence_length is the configured decoder input-token limit per packed document, before batch flattening; it applies to either output format.

max_sequence_length

int | None = None

packed_mask_type

PackedMaskType

return_document_ids

bool = False

return_packed_sequence_metadata

bool = False

uses_native_fa4

bool = False

nemo_automodel.components.models.common.packing.PackingCapabilities.prepare_attention_inputs(
    document_ids: torch.Tensor
) -> dict[str, torch.Tensor | int]

Convert native collator document IDs using the selected model interface.

Parameters:

document_ids

torch.Tensor

Integer tensor [batch, sequence], with positive IDs for document membership and zero for padding. CPU collation stays on CPU; model-entry preparation independently handles transfer.

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

An attention mask [batch, 1, sequence, sequence], compact document

class nemo_automodel.components.models.common.packing.PackingMetadataConsumer()

Protocol

Consumer of AutoModel’s gathered-token indices and cumulative lengths.

packed_token_indices selects non-padding tokens; cu_seqlens and max_seqlen describe that gathered stream, not physical padded storage. Qwen3.5, Kimi Linear/K3 use these boundaries for recurrent-state processing; AutoModel’s FA4 path also needs the unpadded layout. This is independent of the data loader or packing algorithm. The batch-major bundle is normalized by flatten_packed_sequence_metadata after PP splitting, where needed.

requires_packed_sequence_metadata

bool

class nemo_automodel.components.models.common.packing._HFAttentionConsumer()

Protocol

Model whose active attention path dispatches through its HF config.

_uses_hf_attention

bool

Whether HF’s dispatch key controls attention rather than backend.attn.

nemo_automodel.components.models.common.packing._allows_ordinary_causal_attention(
    document_ids: torch.Tensor
) -> bool

Whether each [B,S] attention-ID map is one eligible prefix plus padding.

Any positive ID is valid, independently in each batch document. Prefix or interior exclusions and distinct subdocuments require an explicit mask. Trailing padding cannot affect earlier causal outputs; its loss is ignored.

nemo_automodel.components.models.common.packing._flatten_packed_metadata_at_model_entry(
    _module: torch.nn.Module,
    args: tuple[typing.Any, ...],
    kwargs: dict[str, typing.Any]
) -> tuple[tuple[typing.Any, ...], dict[str, typing.Any]] | None

Flatten packed metadata once after pipeline microbatch splitting.

Parameters:

_module

torch.nn.Module

Native FA4 model receiving the forward call.

args

tuple[Any, ...]

Positional model inputs. Tensor entries retain their model-specific layouts and are returned unchanged.

kwargs

dict[str, Any]

Model inputs containing packed_token_indices of shape [batch, sequence] or [tokens] and cu_seqlens of shape [batch, max_documents + 1] or [documents + 1].

Returns: tuple[tuple[Any, ...], dict[str, Any]] | None

Updated positional and keyword inputs whose packed token indices have

Raises:

ValueError: If pre-packed THD is requested, only one metadata tensor is supplied, or either layout is inconsistent with the current microbatch.

nemo_automodel.components.models.common.packing._install_native_fa4_metadata_hook(
    model: torch.nn.Module
) -> None

Install the once-per-forward native FA4 metadata normalizer.

nemo_automodel.components.models.common.packing._model_attn_implementation(
    model: torch.nn.Module
) -> str | None

Return the packing-relevant attention backend an already-built model runs with.

model.config._attn_implementation is a Transformers dispatch key, whose vocabulary is wider than the mask layouts packing knows about: when flash attention is requested but only the kernels package provides it, Transformers records a kernels-hub id instead of the mainline name. Those ids are mapped back so a model genuinely running varlen flash attention is packed as such. Other live string keys are returned unchanged and therefore use packing’s conservative block-causal default.

nemo_automodel.components.models.common.packing._prepare_flash_metadata_before_input_transfer(
    module: torch.nn.Module,
    args: tuple[typing.Any, ...],
    kwargs: dict[str, typing.Any],
    attn_implementation: str
) -> tuple[tuple[typing.Any, ...], dict[str, typing.Any]] | None

Prepare host offsets before FSDP moves the document map to CUDA.

PP has already selected the microbatch. Keep text Flash kwargs bundled until the decoder hook consumes them: VLM vision towers have their own boundaries and must not receive the text offsets through forwarded keyword arguments.

Parameters:

module

torch.nn.Module

Outer model or pipeline stage, before its FSDP input-transfer hook.

args

tuple[Any, ...]

Token IDs [batch, sequence] or hidden states [batch, sequence, hidden].

kwargs

dict[str, Any]

Forward inputs with CPU attention document IDs [batch, sequence].

attn_implementation

str

Configured FlashAttention backend.

Returns: tuple[tuple[Any, ...], dict[str, Any]] | None

Prepared inputs with device-local document IDs and privately bundled

nemo_automodel.components.models.common.packing._prepare_packed_attention_at_model_entry(
    module: torch.nn.Module,
    args: tuple[typing.Any, ...],
    kwargs: dict[str, typing.Any],
    packed_mask_type: nemo_automodel.components.models.common.packing.PackedMaskType,
    attn_implementation: str
) -> tuple[tuple[typing.Any, ...], dict[str, typing.Any]] | None

Materialize the selected document segmentation after PP splitting.

Parameters:

module

torch.nn.Module

Consumer owning the attention implementation.

args

tuple[Any, ...]

Forward inputs, including token IDs [batch, sequence] or hidden states [batch, sequence, hidden].

kwargs

dict[str, Any]

Forward inputs with integer _packed_seq_ids [batch, sequence] describing effective attention and MTP boundaries, zero at excluded slots. CPU IDs allow host Flash offset construction before transfer. Optional document_ids [batch, sequence] retain original membership for a model-owned policy. Positions [batch, sequence] or [position_axes, batch, sequence] are passed through unchanged.

packed_mask_type

PackedMaskType

Required attention input representation.

attn_implementation

str

Configured backend, used for mask dtype.

Returns: tuple[tuple[Any, ...], dict[str, Any]] | None

Updated args/kwargs retaining all input layouts. Selected IDs move to

nemo_automodel.components.models.common.packing._strip_unused_packed_metadata_at_model_entry(
    _module: torch.nn.Module,
    args: tuple[typing.Any, ...],
    kwargs: dict[str, typing.Any]
) -> tuple[tuple[typing.Any, ...], dict[str, typing.Any]] | None

Keep dataset-only NEAT offsets away from a legacy THD consumer.

Parameters:

_module

torch.nn.Module

Model that does not require explicit NEAT metadata.

args

tuple[Any, ...]

Positional model inputs, returned unchanged in their original layouts.

kwargs

dict[str, Any]

Model inputs. packed_token_indices of shape [batch, sequence] or [tokens] identifies NEAT metadata; its companion cu_seqlens has shape [batch, max_documents + 1] or [documents + 1].

Returns: tuple[tuple[Any, ...], dict[str, Any]] | None

Inputs without NEAT token indices, cumulative lengths, and integer

nemo_automodel.components.models.common.packing._validate_model_consumes_packed_contract(
    model: torch.nn.Module | None
) -> None

Raise if model’s forward cannot receive the typed packing contract.

A model consumes the contract when its forward accepts **kwargs (HF models thread FlashAttentionKwargs this way), names _packed_seq_ids (the custom-model document map), or names all four varlen kwargs explicitly. A forward whose signature cannot be introspected is rejected.

Parameters:

model

torch.nn.Module | None

Already-built model, or a DDP wrapper exposing .module.

nemo_automodel.components.models.common.packing.configure_packing(
    attn_implementation: str,
    model: torch.nn.Module | None = None
) -> nemo_automodel.components.models.common.packing.PackingCapabilities

Configure the model consumer and return its dataset packing contract.

HF attention accepts public masks/varlen kwargs or a CPU batch-major attention document map. Maps are materialized at the model-owned entry point (HF decoder by default) without modifying Transformers globals. Flash offsets are prepared before FSDP input transfer and exposed only to the text consumer. Native FA4 receives batch-major metadata flattened at model entry after microbatch splitting. Other consumers strip unused native metadata so it cannot accidentally select their THD path.

Parameters:

attn_implementation

str

Attention implementation used by the built model.

model

torch.nn.Module | NoneDefaults to None

Optional built model declaring additional packed-metadata needs.

Returns: PackingCapabilities

Structural packed-data requirements for the dataset collater.

Raises:

ValueError: If native FA4 packing lacks a declared compatible consumer.
RuntimeError: If HF flash attention cannot consume the public varlen contract.

nemo_automodel.components.models.common.packing.configure_packing_for_models(
    models: collections.abc.Sequence[torch.nn.Module]
) -> nemo_automodel.components.models.common.packing.PackingCapabilities

Validate and install the adapters for batched packed execution.

Always obtain loader requests with get_packing_capabilities_for_models. Call this setup only for packed batched input. THD validates its length-based interface separately and must not install these mask/Flash/gather hooks.

Parameters:

models

Sequence[torch.nn.Module]

Nonempty sequence of built local stages consuming the same batch.

Returns: PackingCapabilities

The same loader requests as the side-effect-free model query.

Raises:

ValueError: If stages require incompatible batched mask layouts.

nemo_automodel.components.models.common.packing.flatten_packed_sequence_metadata(
    packed_token_indices: torch.Tensor,
    cu_seqlens: torch.Tensor,
    batch_size: int,
    sequence_length: int
) -> tuple[torch.Tensor, torch.Tensor]

Normalize batch-major packed metadata to a single flat token stream.

Parameters:

packed_token_indices

torch.Tensor

Tensor of shape [batch, sequence] containing row-local token indices with -1 padding, or a pre-flattened tensor of shape [tokens] containing indices into [batch, sequence].

cu_seqlens

torch.Tensor

Tensor of shape [batch, max_documents + 1] containing row-local cumulative lengths with -1 padding, or a pre-flattened tensor of shape [documents + 1].

batch_size

int

Batch dimension of the padded token layout.

sequence_length

int

Sequence dimension of the padded token layout.

Returns: torch.Tensor

Flat token indices of shape [tokens] into the padded [batch, sequence]

Raises:

ValueError: If the layouts, boundaries, or represented token counts are inconsistent.

nemo_automodel.components.models.common.packing.get_model_attn_implementation(
    model: torch.nn.Module
) -> str

Return the NEAT/BSHD attention implementation used by a built model.

Custom models expose a typed backend.attn. Hugging Face models record their resolved dispatch key on model.config; this reflects preload and fallback decisions that are intentionally absent from the recipe config. Models mixing native and HF attention declare which dispatch handles BSHD. A separate pre-packed THD path does not determine the NEAT mask layout.

Parameters:

model

torch.nn.Module

Constructed model, optionally wrapped in DDP.

Returns: str

The active attention backend or HF dispatch key, normalized for packing.

Raises:

TypeError: If model is not a torch.nn.Module.

nemo_automodel.components.models.common.packing.get_packing_capabilities(
    attn_implementation: str,
    model: torch.nn.Module | None = None
) -> nemo_automodel.components.models.common.packing.PackingCapabilities

Describe loader requests without configuring or validating packed execution.

This cheap query also serves unpacked and THD loaders. It does not install hooks or invoke model packed-backend validators. Batched packing support is checked separately by configure_packing; THD has its own validation.

Parameters:

attn_implementation

str

Resolved attention implementation.

model

torch.nn.Module | NoneDefaults to None

Optional built consumer declaring its metadata requirements.

Returns: PackingCapabilities

Loader return requests plus model-owned native FA4 setup information.

nemo_automodel.components.models.common.packing.get_packing_capabilities_for_models(
    models: collections.abc.Sequence[torch.nn.Module]
) -> nemo_automodel.components.models.common.packing.PackingCapabilities

Combine loader return requests across local model or pipeline stages.

This query does not configure models or validate an execution format. Batched packed execution separately checks that stage mask layouts agree.

Parameters:

models

Sequence[torch.nn.Module]

Nonempty sequence of built local stages consuming the same batch.

Returns: PackingCapabilities

The first stage’s mask interface, the union of optional return requests,

Raises:

ValueError: If no model stages were supplied.

nemo_automodel.components.models.common.packing.is_indexed_packed_mask(
    attention_mask: torch.Tensor | None
) -> bool

Return True iff attention_mask is an Automodel-style indexed packing mask.

The Automodel neat_packed_vlm_collater (and the LLM equivalent) encode packed-sample boundaries by marking document i (1-based) with the integer i and using 0 for padding (e.g. [1, 1, 1, 2, 2, 3, 3, 0, 0]). Any value greater than 1 is therefore a sufficient signal that two or more documents are packed into the same row. A standard 0/1 attention mask never has values > 1.

nemo_automodel.components.models.common.packing.validate_flash_packing_support(
    attn_implementation: str = 'sdpa',
    model: torch.nn.Module | None = None
) -> None

Fail loudly when flash neat packing cannot be consumed by the backend or model.

Neat packing under flash attention relies on the varlen FlashAttention contract: the collater emits FlashAttentionKwargs (cu_seq_lens_q/ cu_seq_lens_k/max_length_q/max_length_k) which HuggingFace threads into flash_attn_varlen_func. This function verifies, before training starts, that (1) the installed Transformers build exposes that contract and (2) the model’s forward can actually receive it — either as **kwargs, explicit varlen parameters, or the custom-model _packed_seq_ids map. Any gap raises here instead of silently dropping the cumulative lengths and enabling cross-document attention.

Non-flash backends (sdpa / eager) use a 4D block-causal mask built by the collater and need no validation.

Parameters:

attn_implementation

strDefaults to 'sdpa'

The attention implementation the model runs with.

model

torch.nn.Module | NoneDefaults to None

Optional already-built model (or DDP wrapper) to check for the ability to receive the typed packing contract. Skipped when None or when the model’s forward signature cannot be introspected.

Raises:

RuntimeError: If a flash backend is requested but the installed Transformers _flash_attention_forward does not accept the public varlen kwargs, or the model’s forward can consume neither the varlen kwargs nor _packed_seq_ids.

nemo_automodel.components.models.common.packing.validate_thd_packing_support(
    models: collections.abc.Sequence[torch.nn.Module]
) -> None

Reject THD inputs before loader construction if a local stage cannot consume them.

THD uses the model’s seq_lens path rather than the mask/Flash conversion installed by configure_packing_for_models. Validate the existing model/backend capability without selecting a different backend or disabling packing.

nemo_automodel.components.models.common.packing.PackedMaskType = Literal['block_causal', 'document_ids', 'flash_varlen']

nemo_automodel.components.models.common.packing._FLASH_ATTN_IMPLEMENTATIONS = ('flash_attention_2', 'flash_attention_3', 'flash_attention_4')

nemo_automodel.components.models.common.packing.logger = logging.getLogger(__name__)