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:

transformers.modeling_flash_attention_utils._get_unpad_data — extracts per-document sequence lengths from the indexed mask and builds cu_seqlens.
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`	Replacement for `create_causal_mask` that passes through packed masks.
`_patch_preprocess_mask_arguments_for_packing`	Keep indexed packing masks intact for the supported FA2 path.
`configure_packing`	Apply monkey-patches for packed-sequence training with flash_attention_2.
`get_attn_implementation`	Determine the attention backend from model config.
`get_seqlens_in_batch`	Extract per-document sequence lengths from an indexed attention mask.
`get_unpad_data`	Prepare indices and cu_seqlens for `flash_attn_varlen_func`.
`is_indexed_packed_mask`	Return `True` iff `attention_mask` is an Automodel-style indexed packing mask.

Data

_PACKING_PATCH_MODULES

logger

API

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.

nemo_automodel.components.models.common.packing._patch_preprocess_mask_arguments_for_packing() -> None

Keep indexed packing masks intact for the supported FA2 path.

Transformers 5.x preprocesses 2D attention masks before dispatching attention. For flash attention this can coerce integer indexed masks (1, 2, ... per packed document) to bool masks, losing the document boundaries that get_unpad_data needs. Preserve indexed 2D masks for FA2 so the patched flash-attention path can derive per-document cu_seqlens. Validate the private Transformers contract before installing the shim so an incompatible dependency fails instead of silently enabling cross-document attention.

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:

attn_implementation

strDefaults to 'sdpa'

The attention implementation used by the model.

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.

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:

attention_mask

torch.Tensor

[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.

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.

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

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