Hugging Face's logo

Transformers documentation

Google TPU

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v5.22.1).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

Google TPU

Google Cloud TPUs は torch_tpu パッケージを介して PyTorch をネイティブで実行し、tpu デバイスを登録します。torch_tpu は、確認した各入力形状ごとに実行するオペレーションをコンパイルするため、入力形状を静的に保つことで再コンパイルを回避できます。

インストール

torch_tpu リポジトリの手順に従ってtorch_tpuをインストールしてください。torch_tpuはtorchをインポートした際に自動的に読み込まれます。PyTorchがTPUチップを認識していることを確認してください。

import torch

print(torch.accelerator.is_available(), torch.accelerator.device_count())

True と TPU チップの数が出力されます。例えば True 4 のように表示されます。

推論を実行する

device_map="auto" は TPU を検出できません。自動配置は Accelerate が利用可能なデバイスを列挙する機能に依存していますが、Accelerate は TPU デバイスを列挙しないため、モデルは CPU にサイレントにロードされます。from_pretrained() に device_map="tpu" を、または Pipeline に device="tpu" を明示的に指定してください。

.to(model.device) を呼び出す代わりに、生成の入力を CPU に渡します。generate() は、forward が消費するテンソルのみを TPU に移動します。各ステップで増加するテンソル（出力シーケンス、attention_mask）は CPU に保持されるため、コンパイルされたグラフは安定した状態を維持します。generate は、入力と同じデバイスで出力を返します。詳細については、生成テンソルを CPU に保持するを参照してください。TPU では、モデルはデフォルトで SDPA を使用します。

from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B", dtype="bfloat16", device_map="tpu"
)

# leave the inputs on CPU
inputs = tokenizer("The French Bread Law states", return_tensors="pt")
output = model.generate(
    **inputs, do_sample=True, max_new_tokens=50, cache_implementation="static"
)
print(tokenizer.decode(output[0], skip_special_tokens=True))
The French Bread Law states that a farmer can sell their harvest of bread at a certain price. The price is determined by the market price and the price of the bread, and the price is set by the farmers' decision-making committee. What are the implications of this law on

静的キャッシュを使用する場合、generate() は torch.compile と tpu バックエンドを使用して、フォワードパスを自動的にコンパイルします。generate に disable_compile=True を渡すことで、この動作を無効にできます。

テンソル並列

torch_tpu は TPU チップごとに 1 プロセスを実行し、各プロセスは自分のチップを tpu:0 として認識します。DistributedConfig をテンソル並列サイズを指定して作成し、from_pretrained() に渡します。Transformers はモデルの読み込み時にシャード化を行います。

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, DistributedConfig

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-8B")
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-8B",
    dtype=torch.bfloat16,
    distributed_config=DistributedConfig(tp_size=4),
)

# leave the inputs on CPU
inputs = tokenizer("The French Bread Law states", return_tensors="pt")
output = model.generate(
    **inputs, do_sample=False, max_new_tokens=50, cache_implementation="static"
)
if torch.distributed.get_rank() == 0:
    print(tokenizer.decode(output[0], skip_special_tokens=True))

torch_tpu は、プロセスの開始前にいくつかの環境変数を設定する必要がありますが、torchrun だけではこれを行いません。torch.tpu.distributed.set_tpu_launch_env() を呼び出し、Python から torchrun を起動してください。

# launch.py
import sys

import torch
from torch.distributed.run import main

torch.tpu.distributed.set_tpu_launch_env()
main(sys.argv[1:])

通常のtorchrun引数をランチャーに渡します。

python launch.py --nproc_per_node=4 generate.py

一部のモデルは、TPUチップ間でシャード化できるまでに適応処理が必要です。

キーバリューヘッド数がtp_sizeで割り切れないモデルは、アテンション層をシャーディングできず、フォワードパスで形状エラーが発生します。代わりに、tp_plan内のルールをオーバーライドして、アテンションの射影をすべてのチップにレプリケーションしてください。オーバーライドしないルールは、モデルのデフォルトプランが維持されます。例えば、GLM-4-9B-0414はキーバリューヘッドが2つであり、4つのチップに分割できません。

model = AutoModelForCausalLM.from_pretrained(
    "zai-org/GLM-4-9B-0414",
    dtype=torch.bfloat16,
    distributed_config=DistributedConfig(
        tp_size=4,
        tp_plan={
            "layers.*.self_attn.q_proj": "replicated_with_grad_allreduce",
            "layers.*.self_attn.k_proj": "replicated_with_grad_allreduce",
            "layers.*.self_attn.v_proj": "replicated_with_grad_allreduce",
            "layers.*.self_attn.o_proj": "replicated_with_grad_allreduce",
        },
    ),
)

Qwen3-0.6B のような重みが共有されているモデルは、tp_size が 1 より大きい場合に読み込みに失敗します。

Hubカーネルを使用する

kernels ライブラリは v0.17 から TPU をサポートしています。TPU カーネルは JAX と Pallas のコードをラップするため、jax を kernels とともにインストールしてください。

pip install "transformers[kernels]" jax

Hub からカーネルを読み込む方法については、カーネルの読み込みを参照してください。

制限事項

  • device_map="auto" を指定すると、モデルは TPU ではなく CPU にロードされます。device_map="tpu" を明示的に指定してください。
  • 動的な形状は再コンパイルをトリガーします。ステップごとに入力の形状が変化すると、生成が低速になります。より高速な生成には、静的なキャッシュを使用してください。
  • 各プロセスは1つのTPUチップにのみアクセスできます。複数のチップで実行するには、チップごとに1つのプロセスを使用してください。
  • TPU での学習および量子化のための Transformers 統合は現在進行中です。そのため、一部の機能がまだ動作しない場合があります。問題が発生した場合は、issue を作成してください。
Update on GitHub

←概要 AWS Trainium→