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.
`NativeFA4Consumer`	Model whose attention layers invoke Automodel’s native FA4 callable.
`PackedMaskConsumer`	Model that owns masking and consumes compact document IDs.
`PackingCapabilities`	Model-owned requirements for dataset packing and model adaptation.
`PackingMetadataConsumer`	Model that needs dataset-constructed metadata for packed recurrent state.
`_HFAttentionConsumer`	Model whose active attention path dispatches through its HF config.

Functions

Name	Description
`_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.
`_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`	Configure local pipeline stages and combine their dataset requirements.
`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`	Map a model attention implementation to semantic packed-data requirements.
`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.

Data

PackedMaskType

_FLASH_ATTN_IMPLEMENTATIONS

_PACKED_VARLEN_KWARGS

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.NativeFA4Consumer()

Protocol

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

_uses_native_fa4

bool

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.PackingCapabilities(
    packed_mask_type: nemo_automodel.components.models.common.packing.PackedMaskType,
    requires_packed_sequence_metadata: bool = False,
    uses_native_fa4: bool = False
)

Dataclass

Model-owned requirements for dataset packing and model adaptation.

packed_mask_type

PackedMaskType

requires_packed_sequence_metadata

bool = False

uses_native_fa4

bool = False

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

Protocol

Model that needs dataset-constructed metadata for packed recurrent state.

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._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._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 flash attention receives public varlen kwargs from the collater, without modifying Transformers globals. 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:

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

Configure local pipeline stages and combine their dataset requirements.

Parameters:

models

Sequence[torch.nn.Module]

Nonempty sequence of built model stages consuming the same batch.

Returns: PackingCapabilities

Packing contract with metadata enabled if any stage needs it.

Raises:

ValueError: If there are no models or their packed mask layouts disagree.

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

Map a model attention implementation to semantic packed-data requirements.

Parameters:

attn_implementation

str

Attention implementation resolved from the built model.

model

torch.nn.Module | NoneDefaults to None

Optional built model, inspected through explicit packed-metadata, mask, and native-FA4 consumer capabilities.

Returns: PackingCapabilities

Structural capabilities consumed by dataset packing. Backend names do not

Raises:

ValueError: If native FA4 packing is requested without a native consumer.

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.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._PACKED_VARLEN_KWARGS = ('cu_seq_lens_q', 'cu_seq_lens_k', 'max_length_q', 'max_length_k')

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