# TPUでの学習

[TPU (Tensor Processing Unit)](https://cloud.google.com/tpu/docs/intro-to-tpu)は、モデルを効率的にトレーニングするために特化して設計されたハードウェアの一種です。AccelerateはTPUでのトレーニングをサポートしていますが、グラフのコンパイルなど、いくつか注意すべき点があります。本チュートリアルではコンパイルについて簡単に説明し、詳細については[AccelerateによるTPUでのトレーニング](../concept_guides/training_tpu)ガイドを参照してください。

## コンパイル

TPUは、フォワードパス、バックワードパス、オプティマイザーステップなどのトレーニングステップ内のすべての操作のグラフを作成します。そのため、このグラフの構築とコンパイルに時間がかかるため、最初のトレーニングステップは常に時間がかかります。しかし、コンパイルが完了するとキャッシュされ、それ以降のステップははるかに高速になります。

ポイントは、コードを再度コンパイルしないことです。そうしないと、トレーニングが非常に遅くなります。つまり、すべての操作が完全に同一でなければなりません：

* バッチ内のすべてのテンソルは同じ長さである必要があります（NLPタスクでは動的パディングを使用しないなど）
* コードは静的である必要があります（例えば、LSTMのように、入力によって長さが異なるforループを含むレイヤーは使用できません）

## 重みの共有

一般的な言語モデルの設計では、embedding層とsoftmax層の重みを共有（weight tying）します。ただし、モデルをTPUに移動した場合（自分で移動した場合でも、[prepare()](/docs/accelerate/main/ja/package_reference/accelerator#accelerate.Accelerator.prepare)メソッドに渡した場合でも）、weight tyingが解除されるため、重みを再度共有し直す必要があります。

TPU用のスクリプトに特殊な動作（ウェイトのタイイングなど）を追加するには、まず`distributed_type`を`DistributedType.TPU`に設定してください。その後、[tie_weights](https://huggingface.co/docs/transformers/main/en/main_classes/model#transformers.PreTrainedModel.tie_weights)メソッドを使用してウェイトをタイイングできます。

```py
if accelerator.distributed_type == DistributedType.TPU:
    model.tie_weights()
```