TPUでの学習

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

コンパイル

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

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

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

重みの共有

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

TPU用のスクリプトに特殊な動作（ウェイトのタイイングなど）を追加するには、まずdistributed_typeをDistributedType.TPUに設定してください。その後、tie_weightsメソッドを使用してウェイトをタイイングできます。

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