Training on TPU with TensorFlow

What is a TPU?

TPUはTensor Processing Unit（テンソル処理ユニット）の略です。これらはGoogleが設計したハードウェアで、ニューラルネットワーク内のテンソル計算を大幅に高速化するために使用されます。これはGPUのようなものです。ネットワークのトレーニングと推論の両方に使用できます。一般的にはGoogleのクラウドサービスを介してアクセスされますが、Google ColabとKaggle Kernelsを通じても無料で小規模のTPUに直接アクセスできます。

🤗 TransformersのすべてのTensorFlowモデルはKerasモデルですので、この文書のほとんどの方法は一般的にKerasモデル用のTPUトレーニングに適用できます！ただし、TransformersとDatasetsのHuggingFaceエコシステム（hug-o-system？）に固有のポイントもいくつかあり、それについては適用するときにそれを示します。

What kinds of TPU are available?

新しいユーザーは、さまざまなTPUとそのアクセス方法に関する幅広い情報によく混乱します。理解するための最初の重要な違いは、TPUノードとTPU VMの違いです。

TPUノードを使用すると、事実上リモートのTPUに間接的にアクセスします。別個のVMが必要で、ネットワークとデータパイプラインを初期化し、それらをリモートノードに転送します。Google ColabでTPUを使用すると、TPUノードスタイルでアクセスしています。

TPUノードを使用すると、それに慣れていない人々にはかなり予期しない動作が発生することがあります！特に、TPUはPythonコードを実行しているマシンと物理的に異なるシステムに配置されているため、データはローカルマシンにローカルで格納されているデータパイプラインが完全に失敗します。代わりに、データはGoogle Cloud Storageに格納する必要があります。ここでデータパイプラインはリモートのTPUノードで実行されている場合でも、データにアクセスできます。

TPUにアクセスする第二の方法は、TPU VMを介してです。TPU VMを使用する場合、TPUが接続されているマシンに直接接続します。これはGPU VMでトレーニングを行うのと同様です。TPU VMは一般的にデータパイプラインに関しては特に作業がしやすく、上記のすべての警告はTPU VMには適用されません！

これは主観的な文書ですので、こちらの意見です：可能な限りTPUノードの使用を避けてください。 TPU VMよりも混乱しやすく、デバッグが難しいです。将来的にはサポートされなくなる可能性もあります - Googleの最新のTPUであるTPUv4は、TPU VMとしてのみアクセスできるため、TPUノードは将来的には「レガシー」のアクセス方法になる可能性が高いです。ただし、無料でTPUにアクセスできるのはColabとKaggle Kernelsの場合があります。その場合、どうしても使用しなければならない場合の取り扱い方法を説明しようとします！詳細はTPUの例のノートブックで詳細な説明を確認してください。

What sizes of TPU are available?

単一のTPU（v2-8/v3-8/v4-8）は8つのレプリカを実行します。TPUは数百から数千のレプリカを同時に実行できるポッドに存在します。単一のTPUよりも多くのTPUを使用するが、ポッド全体ではない場合（たとえばv3-32）、TPUフリートはポッドスライスとして参照されます。

Colabを介して無料のTPUにアクセスする場合、通常は単一のv2-8 TPUが提供されます。

I keep hearing about this XLA thing. What’s XLA, and how does it relate to TPUs?

XLAは、TensorFlowとJAXの両方で使用される最適化コンパイラです。JAXでは唯一のコンパイラであり、TensorFlowではオプションですが（しかしTPUでは必須です！）、Kerasモデルをトレーニングする際にmodel.compile()に引数jit_compile=Trueを渡すことで最も簡単に有効にできます。エラーが発生せず、パフォーマンスが良好であれば、それはTPUに移行する準備が整った良い兆候です！

TPU上でのデバッグは一般的にCPU/GPUよりも少し難しいため、TPUで試す前にまずCPU/GPUでXLAを使用してコードを実行することをお勧めします。もちろん、長時間トレーニングする必要はありません。モデルとデータパイプラインが期待通りに動作するかを確認するための数ステップだけです。

How do I make my model XLA compatible?

多くの場合、コードはすでにXLA互換かもしれません！ただし、XLAでは動作する通常のTensorFlowでも動作しないいくつかの要素があります。以下に、3つの主要なルールにまとめています：

XLA Rule #1: Your code cannot have “data-dependent conditionals”

これは、任意のifステートメントがtf.Tensor内の値に依存していない必要があることを意味します。例えば、次のコードブロックはXLAでコンパイルできません！

if tf.reduce_sum(tensor) > 10:
    tensor = tensor / 2.0

これは最初は非常に制限的に思えるかもしれませんが、ほとんどのニューラルネットコードはこれを行う必要はありません。通常、この制約を回避するためにtf.condを使用するか（ドキュメントはこちらを参照）、条件を削除して代わりに指示変数を使用したりすることができます。次のように：

sum_over_10 = tf.cast(tf.reduce_sum(tensor) > 10, tf.float32)
tensor = tensor / (1.0 + sum_over_10)

このコードは、上記のコードとまったく同じ効果を持っていますが、条件を回避することで、XLAで問題なくコンパイルできることを確認します！

XLA Rule #2: Your code cannot have “data-dependent shapes”

これは、コード内のすべての tf.Tensor オブジェクトの形状が、その値に依存しないことを意味します。たとえば、tf.unique 関数はXLAでコンパイルできないので、このルールに違反します。なぜなら、これは入力 Tensor の一意の値の各インスタンスを含む tensor を返すためです。この出力の形状は、入力 Tensor の重複具合によって異なるため、XLAはそれを処理しないことになります！

一般的に、ほとんどのニューラルネットワークコードはデフォルトでルール＃2に従います。ただし、いくつかの一般的なケースでは問題が発生することがあります。非常に一般的なケースの1つは、ラベルマスキングを使用する場合です。ラベルを無視して損失を計算する場所を示すために、ラベルを負の値に設定する方法です。NumPyまたはPyTorchのラベルマスキングをサポートする損失関数を見ると、次のようなブールインデックスを使用したコードがよく見られます：

label_mask = labels >= 0
masked_outputs = outputs[label_mask]
masked_labels = labels[label_mask]
loss = compute_loss(masked_outputs, masked_labels)
mean_loss = torch.mean(loss)

このコードはNumPyやPyTorchでは完全に機能しますが、XLAでは動作しません！なぜなら、masked_outputsとmasked_labelsの形状はマスクされた位置の数に依存するため、これはデータ依存の形状になります。ただし、ルール＃1と同様に、このコードを書き直して、データ依存の形状なしでまったく同じ出力を生成できることがあります。

label_mask = tf.cast(labels >= 0, tf.float32)
loss = compute_loss(outputs, labels)
loss = loss * label_mask  # Set negative label positions to 0
mean_loss = tf.reduce_sum(loss) / tf.reduce_sum(label_mask)

ここでは、データ依存の形状を避けるために、各位置で損失を計算してから、平均を計算する際に分子と分母の両方でマスクされた位置をゼロ化する方法を紹介します。これにより、最初のアプローチとまったく同じ結果が得られますが、XLA互換性を維持します。注意点として、ルール＃1と同じトリックを使用します - tf.boolをtf.float32に変換して指標変数として使用します。これは非常に便利なトリックですので、自分のコードをXLAに変換する必要がある場合には覚えておいてください！

XLA Rule #3: XLA will need to recompile your model for every different input shape it sees

これは重要なルールです。これはつまり、入力形状が非常に変動的な場合、XLA はモデルを何度も再コンパイルする必要があるため、大きなパフォーマンスの問題が発生する可能性があるということです。これは NLP モデルで一般的に発生し、トークナイズ後の入力テキストの長さが異なる場合があります。他のモダリティでは、静的な形状が一般的であり、このルールはほとんど問題になりません。

ルール＃3を回避する方法は何でしょうか？鍵は「パディング」です - すべての入力を同じ長さにパディングし、次に「attention_mask」を使用することで、可変形状と同じ結果を得ることができますが、XLA の問題は発生しません。ただし、過度のパディングも深刻な遅延を引き起こす可能性があります - データセット全体で最大の長さにすべてのサンプルをパディングすると、多くの計算とメモリを無駄にする可能性があります！

この問題には完璧な解決策はありませんが、いくつかのトリックを試すことができます。非常に便利なトリックの1つは、バッチのサンプルを32または64トークンの倍数までパディングすることです。これにより、トークン数がわずかに増加するだけで、すべての入力形状が32または64の倍数である必要があるため、一意の入力形状の数が大幅に減少します。一意の入力形状が少ないと、XLA の再コンパイルが少なくなります！

How do I actually train my model on TPU?

一度トレーニングが XLA 互換性があることを確認し、（TPU Node/Colab を使用する場合は）データセットが適切に準備されている場合、TPU 上で実行することは驚くほど簡単です！コードを変更する必要があるのは、いくつかの行を追加して TPU を初期化し、モデルとデータセットが TPUStrategy スコープ内で作成されるようにすることだけです。これを実際に見るには、TPU のサンプルノートブックをご覧ください！

Summary

ここでは多くの情報が提供されましたので、TPU でモデルをトレーニングする際に以下のチェックリストを使用できます：

• コードが XLA の三つのルールに従っていることを確認します。
• CPU/GPU で jit_compile=True を使用してモデルをコンパイルし、XLA でトレーニングできることを確認します。
• データセットをメモリに読み込むか、TPU 互換のデータセット読み込みアプローチを使用します（ノートブックを参照）。
• コードを Colab（アクセラレータを「TPU」に設定）または Google Cloud の TPU VM に移行します。
• TPU 初期化コードを追加します（ノートブックを参照）。
• TPUStrategy を作成し、データセットの読み込みとモデルの作成が strategy.scope() 内で行われることを確認します（ノートブックを参照）。
• TPU に移行する際に jit_compile=True を外すのを忘れないでください！
• 🙏🙏🙏🥺🥺🥺
• model.fit() を呼び出します。
• おめでとうございます！