Transformers documentation
概要
概要
Transformers は、モデルを高速かつ低コストで利用可能にするための複数の推論最適化技術を提供しています。選択肢には、メモリトラフィックを削減するための代替的なアテンション機構、高速な実行のためのコードコンパイル、スループット向上のための最適化されたカーネルが含まれます。これらの技術を組み合わせることで、最大限のパフォーマンスを発揮できます。
デバイスのセットアップやハードウェア固有の設定については、Hardwareガイドを参照してください。以下の最適化を適用する前に、お使いのハードウェアでどの技術がサポートされているかを確認してください。
メモリと速度は密接に関連していますが、同一ではありません。メモリフットプリントを縮小すると、移動させるデータ量が減るため、モデルは「高速」になります。純粋な速度最適化は必ずしもメモリを削減するわけではなく、場合によっては使用量が増えることもあります。ユースケースとハードウェアに応じて、適切な最適化を選択してください。
以下の表を参照して、最適化手法を選択してください。
| 技術 | 速度 | メモリ |
|---|---|---|
| コンパイル | ✅ | |
| Attention バックエンド | ✅ | ✅ |
| カーネル | ✅ | ✅ |
| 量子化 | ✅ | ✅ |
| キャッシュ | ✅ | ✅ |
| 並列性 | ✅ | |
| 連続バッチ処理 | ✅ |
このガイドでは、Transformersの最適化に関するクイックスタートをご案内します。
コンパイル
torch.compile は Python のオーバーヘッドを削減し、オペレーションを融合し、形状やハードウェアに最適化されたカーネルを生成します。初回の実行でウォームアップが行われ、以降の実行では高速なコンパイル済みパスが使用されます。
torch.compileを自動的にトリガーするには、generate()に固定サイズキャッシュを渡します。
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-0.6B", dtype=torch.float16, device_map="auto")
input = tokenizer("The French Bread Law states", return_tensors="pt").to(model.device)
output = model.generate(**input, do_sample=False, max_new_tokens=20, cache_implementation="static")
tokenizer.batch_decode(output, skip_special_tokens=True)[0]モデルが各ステップで再コンパイルされるのを防ぐため、generate()の外で
torch.compile(model)を呼び出さないでください。
アテンションのバックエンド
代替のアテンションバックエンドはメモリトラフィックを低減します。例えば、FlashAttentionはアテンションの計算をタイル化し、大きな中間テンソルを回避することでメモリフットプリントを削減します。
attn_implementation を from_pretrained() に設定して、最適化されたアテンションバックエンドを読み込みます。
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-0.6B", attn_implementation="flash_attention_2")カーネル
Kernelsは演算を融合してスループットを向上させ、メモリ使用量を削減します。Kernelsライブラリは、Hubから最適化された計算カーネルを、柔軟かつバージョン安全な方法で読み込みます。
FlashAttention-2のパッケージをインストールせずに、最適化されたカーネルを読み込む例を以下に示します。
import torch
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(
"Qwen/Qwen3-0.6B", attn_implementation="kernels-community/flash-attn2"
)量子化
Quantization はすべてのパラメータのサイズを縮小し、メモリ使用量を削減するとともに、より多くの演算を実行できるため処理速度を向上させます。
quantization_config 引数に量子化設定を渡します。from_pretrained() で使用します。各量子化バックエンドには、異なる引数を持つ個別の設定があります。以下の例では、bitsandbytes バックエンドを使用してモデルを4ビットに量子化し、計算のデータ型を設定しています。
import torch
from transformers import AutoModelForCausalLM, BitsAndBytesConfig
bnb_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16)
model = AutoModelForCausalLM.from_pretrained(
"allenai/Olmo-3-7B-Think", quantization_config=bnb_config
)キャッシュ
Cachingは、すべてのトークンに対して再計算するのではなく、過去のキーと値を再利用することで生成を高速化します。過去のキーと値を保存するメモリコストを相殺・低減するため、TransformersはキャッシュをCPUへオフロードすることをサポートしています。GPU上に保持されるのは現在のレイヤーのみです。
generate() の cache_implementation 引数を使用して、キャッシュ戦略を設定します。
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")
model = AutoModelForCausalLM.from_pretrained(
"Qwen/Qwen3-0.6B", attn_implementation="kernels-community/flash-attn2"
)
inputs = tokenizer("The Le Décret Pain states that a baguette must,", return_tensors="pt")
outputs = model.generate(**inputs, do_sample=False, max_new_tokens=50, cache_implementation="offloaded")並列性
Parallelismは、モデルを複数のデバイスに分散して配置し、1つのデバイスに収まらないほど大きなモデルでも高速に実行できるようにする手法です。このアプローチでは、シャード化のオーバーヘッドや結果を同期するための通信が発生するため、メモリ使用量が増加します。
テンソル並列では、モデルのレイヤーを複数のデバイスに分割します。DistributedConfig(tp_size=N)で使用するデバイス数を設定し、from_pretrained()に渡すことで有効にできます。
from transformers import AutoModelForCausalLM, DistributedConfig
distributed_config = DistributedConfig(tp_size=4)
model = AutoModelForCausalLM.from_pretrained(
"meta-llama/Meta-Llama-3-8B-Instruct",
distributed_config=distributed_config,
)
print(model._tp_plan)連続バッチ処理
Continuous batchingは、動的スケジューリングとチャンク化プリフィルによりGPUを常に稼働させ、スループットを最大化します。Servingアプリケーションでは、これを用いて複数の受信リクエストを並行して処理します。
generate_batch() を使用して、連続バッチ処理を有効にします。
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.generation import ContinuousBatchingConfig, GenerationConfig
model = AutoModelForCausalLM.from_pretrained(
"Qwen/Qwen3-0.6B",
attn_implementation="sdpa",
device_map="auto",
torch_dtype=torch.bfloat16,
)
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")
prompts = [
"The Le Décret Pain states that a baguette must",
"Explain gravity in one sentence.",
"Name the capital of France.",
]
inputs = [tokenizer.encode(p) for p in prompts]
generation_config = GenerationConfig(
max_new_tokens=32,
eos_token_id=tokenizer.eos_token_id,
pad_token_id=tokenizer.pad_token_id,
do_sample=False,
)
continuous_batching_config = ContinuousBatchingConfig(max_batch_tokens=512)
outputs = model.generate_batch(
inputs=inputs,
generation_config=generation_config,
continuous_batching_config=continuous_batching_config,
)
for request_id, output in outputs.items():
text = tokenizer.decode(output.generated_tokens, skip_special_tokens=True)
print(f"[{request_id}] {text}")