パディングなしの学習

Padding-free training（packingとも呼ばれる）では、各サンプルを固定長にパディングする代わりに、複数のサンプルを単一のシーケンスに連結します。モデルは各サンプルの終了位置を知る必要があり、これにより（線形）attentionがサンプル間でトークンを混在させないようにします。

これらの境界を指定するには、2つの方法があります。

• データコラテーターを使って、事前にそれらを準備してください。
• 実行時にposition_idsから推論します。

推奨されるアプローチは、データコラテーターです。本ガイドでは、その理由とposition_idsパスに関する注意点について説明します。

position_ids から境界を推論することは推奨される方法ではなく、標準的な attention モデルでのみ機能します。Qwen3-Next や Qwen3.5 (Gated DeltaNet) などの Linear-attention モデル、および convolution ベースのモデルは position_ids の境界を無視するため、data collator が必要です。Linear attention と convolution モデルを参照してください。

データコラクターで境界を準備する

境界の kwargs を事前に準備しておくと、上記の問題が解消され、compile を使用する場合としない場合とで同じ挙動になります。

DataCollatorWithFlattening を使用して、各バッチをフラット化し、境界情報を返します。return_flash_attn_kwargs=True を設定することで、collator が実行時に position_ids から推論するのではなく、事前に境界を計算します。Trainer に渡す際は、フラット化されたバッチがすでに境界をエンコードしているため、attention_mask を追加しないでください。マスクはパッキングされたレイアウトと競合するためです。

Padding-free は、標準的なアテンションモデルに対して FlashAttention の実装に依存しています。フラット化されたバッチに必要な可変長パスを公開しているのは、FlashAttention カーネルのみだからです。

kernels ライブラリをインストールします。このライブラリは、ローカルでのビルドを必要とせずに、事前ビルド済みの FlashAttention カーネルを取得します。また、flash-attn がローカルにインストールされていない場合のフォールバックとしても機能します。モデルは attn_implementation="kernels-community/flash-attn2" を指定して読み込みます。

import torch
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer, DataCollatorWithFlattening, Trainer, TrainingArguments

model_id = "meta-llama/Llama-3.2-1B"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    attn_implementation="flash_attention_2",
    device_map="auto",
)

dataset = load_dataset("Salesforce/wikitext", "wikitext-2-raw-v1", split="train")
dataset = dataset.map(
    lambda example: tokenizer(example["text"], truncation=True, max_length=512),
    remove_columns=dataset.column_names,
)

# return_flash_attn_kwargs=True precomputes the sequence boundaries
data_collator = DataCollatorWithFlattening(return_flash_attn_kwargs=True)

trainer = Trainer(
    model=model,
    args=TrainingArguments(output_dir="padding-free-llama"),
    train_dataset=dataset,
    data_collator=data_collator,
)
trainer.train()

position_ids から境界を推論する

FlashAttention は position_ids のみからパディングなしのバッチを検出でき、TRL などの下流フレームワークがこれに依存しているため、後方互換性を保つ目的で残されています。

position_ids に依存することには、2つの問題があります。

• position_idsからのパックシーケンスの検出は、動的かつデータ依存のチェックです。コンパイルなしでも動作しますが、torch.compile下ではグラフの分断を引き起こします。実際のバッチサイズは通常より大きいため、このチェックの実行頻度を抑制するために、現在はbatch_size == 1に限定されています。
• Compiled FlashAttention は、一部の kwargs がプレーンな Python の int であることを強制します。実行時に position_ids からこれらを推論すると、デバイスとホスト間の同期が発生し、古い PyTorch バージョンでは、テンソルから int への変換により追加のグラフブレイクが発生します。

線形アテンションと畳み込みモデル

Gated DeltaNet (GDN)、その他の linear-attention レイヤー、および causal convolution は、設計上、position_ids のみのパスを持ちません。これらのモデルでは、collator を用いてデータを準備することが唯一サポートされている方法です。

GDN、linear-attention、またはcausal convolutionモデルでは、position_idsだけに依存しないでください。データコラクターで、seq_idxを含む境界のkwargsを準備してください。

これらのモデルでは、return_flash_attn_kwargs=True と return_seq_idx=True の両方を設定してください。

from transformers import DataCollatorWithFlattening

data_collator = DataCollatorWithFlattening(
    return_flash_attn_kwargs=True,
    return_seq_idx=True,
)

正確なカーネルパッケージは、モデルの元の実装によって異なります。Qwen3-Next や Qwen3.5 などの Gated DeltaNet モデルは flash-linear-attention を使用し、Bamba などの Mamba ベースのモデルは mamba-ssm を使用します。どちらも畳み込みには causal-conv1d に依存しています。適切なカーネルがない場合、モデルは境界の kwargs を無視し、サンプル間でトークンを混在させる参照実装にフォールバックします。

これらのカーネルの多くは、kernelsライブラリを通じて利用可能であり、互換性のあるビルドを取得できます。flash-linear-attentionは、通常、直接インストールする必要があります。

境界の kwargs が欠落している場合、カーネルはバッチ全体を単一のシーケンスとして静かに扱います。ランタイムチェックはデータ依存の分岐を追加するため torch.compile と競合するため、エラーや警告は発生しません。

次のステップ

• その他のコラレーターについては、データコラレーターのガイドを参照してください。
• DataCollatorWithFlattening の API リファレンスをご覧ください。すべての引数が記載されています。
• ベンチマークとより詳細な解説については、Flash Attention を用いたパッキングによる Hugging Face の学習効率の向上をお読みください。