> This page is for version 0.5.0 · 26.06.
> For other versions, use one of these documentation indexes:
> - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt
> - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt
> - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt
> - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt

> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# 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`](#nemo_automodel-components-models-common-packing-AttentionBackendSelection) | Typed attention selection carried by a built custom model.                |
| [`NativeFA4Consumer`](#nemo_automodel-components-models-common-packing-NativeFA4Consumer)                 | Model whose attention layers invoke Automodel's native FA4 callable.      |
| [`PackedMaskConsumer`](#nemo_automodel-components-models-common-packing-PackedMaskConsumer)               | Model that owns masking and consumes compact document IDs.                |
| [`PackingCapabilities`](#nemo_automodel-components-models-common-packing-PackingCapabilities)             | Model-owned requirements for dataset packing and model adaptation.        |
| [`PackingMetadataConsumer`](#nemo_automodel-components-models-common-packing-PackingMetadataConsumer)     | Model that needs dataset-constructed metadata for packed recurrent state. |
| [`_HFAttentionConsumer`](#nemo_automodel-components-models-common-packing-_HFAttentionConsumer)           | Model whose active attention path dispatches through its HF config.       |

### Functions

| Name                                                                                                                                            | Description                                                                     |
| ----------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------- |
| [`_flatten_packed_metadata_at_model_entry`](#nemo_automodel-components-models-common-packing-_flatten_packed_metadata_at_model_entry)           | Flatten packed metadata once after pipeline microbatch splitting.               |
| [`_install_native_fa4_metadata_hook`](#nemo_automodel-components-models-common-packing-_install_native_fa4_metadata_hook)                       | Install the once-per-forward native FA4 metadata normalizer.                    |
| [`_model_attn_implementation`](#nemo_automodel-components-models-common-packing-_model_attn_implementation)                                     | Return the packing-relevant attention backend an already-built model runs with. |
| [`_strip_unused_packed_metadata_at_model_entry`](#nemo_automodel-components-models-common-packing-_strip_unused_packed_metadata_at_model_entry) | Keep dataset-only NEAT offsets away from a legacy THD consumer.                 |
| [`_validate_model_consumes_packed_contract`](#nemo_automodel-components-models-common-packing-_validate_model_consumes_packed_contract)         | Raise if `model`'s `forward` cannot receive the typed packing contract.         |
| [`configure_packing`](#nemo_automodel-components-models-common-packing-configure_packing)                                                       | Configure the model consumer and return its dataset packing contract.           |
| [`configure_packing_for_models`](#nemo_automodel-components-models-common-packing-configure_packing_for_models)                                 | Configure local pipeline stages and combine their dataset requirements.         |
| [`flatten_packed_sequence_metadata`](#nemo_automodel-components-models-common-packing-flatten_packed_sequence_metadata)                         | Normalize batch-major packed metadata to a single flat token stream.            |
| [`get_model_attn_implementation`](#nemo_automodel-components-models-common-packing-get_model_attn_implementation)                               | Return the NEAT/BSHD attention implementation used by a built model.            |
| [`get_packing_capabilities`](#nemo_automodel-components-models-common-packing-get_packing_capabilities)                                         | Map a model attention implementation to semantic packed-data requirements.      |
| [`is_indexed_packed_mask`](#nemo_automodel-components-models-common-packing-is_indexed_packed_mask)                                             | Return `True` iff `attention_mask` is an Automodel-style indexed packing mask.  |
| [`validate_flash_packing_support`](#nemo_automodel-components-models-common-packing-validate_flash_packing_support)                             | Fail loudly when flash neat packing cannot be consumed by the backend or model. |

### Data

[`PackedMaskType`](#nemo_automodel-components-models-common-packing-PackedMaskType)

[`_FLASH_ATTN_IMPLEMENTATIONS`](#nemo_automodel-components-models-common-packing-_FLASH_ATTN_IMPLEMENTATIONS)

[`_PACKED_VARLEN_KWARGS`](#nemo_automodel-components-models-common-packing-_PACKED_VARLEN_KWARGS)

[`logger`](#nemo_automodel-components-models-common-packing-logger)

### API

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

Protocol

Typed attention selection carried by a built custom model.

**`attn`** `AttentionBackend`

Selected native attention backend.

---

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

Protocol

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

**`_uses_native_fa4`** `bool`

---

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

Protocol

Model that owns masking and consumes compact document IDs.

**`packed_mask_type`** `PackedMaskType`

---

```python
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`

---

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

Protocol

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

**`requires_packed_sequence_metadata`** `bool`

---

```python
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.

---

```python
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.

```python
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.

```python
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.

```python
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

```python
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`.

---

```python
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 | None` — default: 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.

```python
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.

```python
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.

```python
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.

```python
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 | None` — default: 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.

```python
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.

```python
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`** `str` — default: 'sdpa'

The attention implementation the model runs with.

---

**`model`** `torch.nn.Module | None` — default: 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`.

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

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

```python
nemo_automodel.components.models.common.packing._PACKED_VARLEN_KWARGS = ('cu_seq_lens_q', 'cu_seq_lens_k', 'max_length_q', 'max_length_k')
```

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