TPU ポッドでのトレーニング
概要
TPU は、TPU Pod にスケールアウトされるように設計されています。TPU Pod は、専用の高速ネットワーク インターフェースによって接続された TPU デバイスのコレクションです。TPU Pod を使用すると、処理負荷を複数の TPU に分散できます。各 TPU ボードは、データの読み込みや前処理などのために、高性能の CPU ベースのホストマシンに接続されます。多数の TPU を最大限に活用するには、いくつかのトレーニング タスク パラメータを調整する必要があります。
TPU Pod を使用したトレーニングの設定は、フレームワークごとに異なります。各フレームワークを使用した Pod でのトレーニングの詳細については、次のリンクをご覧ください。
以降のセクションでは、一般的な問題、モデルに加える必要がある変更、Pod の障害を軽減または回避するためのベスト プラクティスについて説明します。
バッチサイズとトレーニング ステップのスケーリング
規模が大きい TPU タイプで線形スケーリングを行うには、コアごとのバッチサイズを同じにします。
たとえば、v2-8 でバッチサイズ 1024 を使用する場合、v2-32 ではバッチサイズ 4096（4 * 1024）を使用します。これは TPU ハードウェアを十分に活用しています。小さいバッチサイズを使用できますが、そうしてもトレーニングを直線的にスケーリングすることはできません。
一部のモデルには、1 つのステップが 1 つのデータバッチの処理に対応する train_steps フラグが含まれています。バッチサイズを増やす場合は、トレーニング ステップの数を減らして、トレーニング サンプルの合計数が変わらないようにします。
たとえば、100 ステップのバッチサイズが 1000 の場合、トレーニング中に 100,000 のサンプルが処理されます。ワーカーが 4 つで、バッチサイズが 4,000 であれば、同じく 100,000 サンプルを処理できるように、ステップ数を 25 に調整する必要があります。モデルで epochs フラグを使用する場合、ステップ数をスケーリングする必要はありません。
バッチサイズが大きいと、モデルの収束動作が変わるため、学習率など、一部のハイパーパラメータの調整も必要になる場合があります。
TPU Pod と同じリージョン内でのリージョン Google Cloud Storage バケットの使用
一般に、TPU トレーニングにおすすめの方法は、常に同じリージョン内のリソースを使用することです。Google Cloud Storage バケットと TPU が同じリージョンに存在する場合に、データ転送率がより高くなるため、TPU Pod を使用する場合は、リソースのリージョンが特に重要です。
データセットとチェックポイントのトレーニングに使用する TPU と同じリージョンで、リージョンの Google Cloud Storage バケットを使用していることを確認してください。
TPU Pod での開発に向けたワークフローのベスト プラクティス
新しい TPU ワークロードを開発する場合は、通常、最小の TPU で開発を開始し、より大きな TPU サイズを徐々に反復していくのが最適です。まず、下位の TPU バージョン（v2-8 や v3-8 など）を使用します。
- ワークロードの機能をテストする
- パフォーマンス ツールを使用してパフォーマンスをテストして検証する
ワークロードが機能し、パフォーマンスの目標を達成したら、v2-32 や v3-32 などの上位の TPU タイプにスケールアップします。目的の TPU サイズに達するまで、TPU のサイズを徐々に増やして反復処理ながら、スケーラビリティ（機能とパフォーマンス）を検証します。