# Google TPU

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

## インストール

[torch_tpu リポジトリ](https://github.com/google-pytorch/torch_tpu)の手順に従って`torch_tpu`をインストールしてください。`torch_tpu`は`torch`をインポートした際に自動的に読み込まれます。PyTorchがTPUチップを認識していることを確認してください。

```py
import torch

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

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

## 推論を実行する

`device_map="auto"` は TPU を検出できません。自動配置は [Accelerate](https://huggingface.co/docs/accelerate/index) が利用可能なデバイスを列挙する仕組みに依存していますが、Accelerate は TPU デバイスを列挙しないため、モデルは CPU にサイレントにロードされます。[from_pretrained()](/docs/transformers/main/ja/main_classes/model#transformers.PreTrainedModel.from_pretrained) に `device_map="tpu"` を、または [Pipeline](/docs/transformers/main/ja/main_classes/pipelines#transformers.Pipeline) に `device="tpu"` を明示的に指定してください。

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

```py
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))
```

```txt
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()](/docs/transformers/main/ja/main_classes/text_generation#transformers.GenerationMixin.generate) は `torch.compile` と `tpu` バックエンドを使用して、フォワードパスを自動的にコンパイルします。`generate` に `disable_compile=True` を渡すことで、この動作を無効にできます。

### テンソル並列

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

```py
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` を起動してください。

```py
# launch.py
import sys

import torch
from torch.distributed.run import main

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

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

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

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

`tp_size` で割り切れない数の key-value heads を持つモデルは、attention レイヤーをシャーディングできず、forward pass で shape エラーが発生します。代わりに、`tp_plan` でルールをオーバーライドして、attention の射影をすべてのチップに複製してください。オーバーライドしないルールは、モデルのデフォルトのプランが維持されます。例えば、[GLM-4-9B-0414](https://huggingface.co/zai-org/GLM-4-9B-0414) は 2 つの key-value heads を持ち、4 つのチップに分割することはできません。

```py
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](https://github.com/huggingface/kernels) ライブラリは v0.17 から TPU をサポートしています。TPU カーネルは JAX と [Pallas](https://docs.jax.dev/en/latest/pallas/index.html) のコードをラップするため、`jax` を `kernels` とともにインストールしてください。

```bash
pip install "transformers[kernels]" jax
```

Hub からカーネルを読み込む方法については、[カーネルの読み込み](./kernel_doc/loading_kernels)を参照してください。

## 制限事項

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