> 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

Flash Attention packing support via monkey-patching.

When `attn_implementation="flash_attention_2"` and neat packing is enabled,
the collater produces an **indexed** attention mask `[B, S]` where each
position contains the 1-based document index (0 = padding).  For example::

\[1, 1, 2, 2, 2, 0]   # 2 tokens in doc 1, 3 in doc 2, 1 padding

To make HuggingFace's flash attention path use `flash_attn_varlen_func`
with per-document `cu_seqlens`, we monkey-patch two functions:

1. `transformers.modeling_flash_attention_utils._get_unpad_data` — extracts
   per-document sequence lengths from the indexed mask and builds cu\_seqlens.
2. `transformers.models.qwen3_vl.modeling_qwen3_vl.create_causal_mask` —
   returns the 2D indexed mask as-is, bypassing 4D mask creation.

This is the same approach used by LlamaFactory.

## Module Contents

### Functions

| Name                                                                                                                  | Description                                                                    |
| --------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------ |
| [`_passthrough_create_causal_mask`](#nemo_automodel-components-models-common-packing-_passthrough_create_causal_mask) | Replacement for `create_causal_mask` that passes through packed masks.         |
| [`configure_packing`](#nemo_automodel-components-models-common-packing-configure_packing)                             | Apply monkey-patches for packed-sequence training with flash\_attention\_2.    |
| [`get_attn_implementation`](#nemo_automodel-components-models-common-packing-get_attn_implementation)                 | Determine the attention backend from model config.                             |
| [`get_seqlens_in_batch`](#nemo_automodel-components-models-common-packing-get_seqlens_in_batch)                       | Extract per-document sequence lengths from an indexed attention mask.          |
| [`get_unpad_data`](#nemo_automodel-components-models-common-packing-get_unpad_data)                                   | Prepare indices and cu\_seqlens for `flash_attn_varlen_func`.                  |
| [`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. |

### Data

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

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

### API

```python
nemo_automodel.components.models.common.packing._passthrough_create_causal_mask(
    config = None,
    input_embeds = None,
    inputs_embeds = None,
    attention_mask = None,
    cache_position = None,
    past_key_values = None,
    position_ids = None,
    kwargs = {}
)
```

Replacement for `create_causal_mask` that passes through packed masks.

FA2 handles masking internally, so always pass through.  For non-FA2
backends, pass through packed masks but delegate normal 2D masks to HF.

```python
nemo_automodel.components.models.common.packing.configure_packing(
    attn_implementation: str = 'sdpa'
) -> None
```

Apply monkey-patches for packed-sequence training with flash\_attention\_2.

Only patches when `attn_implementation == "flash_attention_2"`.

**Parameters:**

The attention implementation used by the model.

```python
nemo_automodel.components.models.common.packing.get_attn_implementation(
    cfg_model
)
```

Determine the attention backend from model config.

Custom models store it in `backend.attn`; HF models use `attn_implementation`.

```python
nemo_automodel.components.models.common.packing.get_seqlens_in_batch(
    attention_mask: torch.Tensor
) -> torch.Tensor
```

Extract per-document sequence lengths from an indexed attention mask.

Example::

\>>> get\_seqlens\_in\_batch(torch.tensor(\[\[1, 1, 2, 2, 2, 0],
...                                    \[1, 2, 2, 3, 3, 3]]))
tensor(\[2, 3, 1, 2, 3])

**Parameters:**

`[B, S]` integer tensor where each position contains
the 1-based document index (0 = padding).

**Returns:** `torch.Tensor`

1D tensor of all individual document lengths across the batch.

```python
nemo_automodel.components.models.common.packing.get_unpad_data(
    attention_mask: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor, int]
```

Prepare indices and cu\_seqlens for `flash_attn_varlen_func`.

This is a drop-in replacement for
`transformers.modeling_flash_attention_utils._get_unpad_data`
that handles **indexed** attention masks (values 1, 2, 3, …) instead of
binary (0/1) masks.  Each unique non-zero value is treated as a separate
document, so `flash_attn_varlen_func` applies causal attention
*within* each document without cross-document attention.

Example::

\>>> get\_unpad\_data(torch.tensor(\[\[1, 1, 2, 2, 2, 0],
...                              \[1, 2, 2, 3, 3, 3]]))
(tensor(\[0, 1, 2, 3, 4, 6, 7, 8, 9, 10, 11]),
tensor(\[ 0,  2,  5,  6,  8, 11], dtype=torch.int32),
3\)

**Returns:** `torch.Tensor`

Indices of non-padding tokens from the flattened sequence.

```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._PACKING_PATCH_MODULES = ['transformers.models.qwen2.modeling_qwen2', 'transformers.models.qwen2_5_vl.mod...
```

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