TPU ポッドでのトレーニング

概要

TPU は、TPU Pod にスケールアウトされるように設計されています。TPU Pod は、専用の高速ネットワーク インターフェースによって接続された TPU デバイスのコレクションです。1 つの TPU Pod に最大 2,048 個の TPU コアを設定して、複数の TPU で処理負荷を分散できます。各 TPU ボードは、データの読み込みや前処理などのために、高性能の CPU ベースのホストマシンに接続されます。多数の TPU を最大限に活用するには、いくつかのトレーニング タスク パラメータを調整する必要があります。このドキュメントでは、共通の問題、モデルに加える必要がある変更、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 バケットを使用していることを確認してください。

データ ストレージに NFS を使用する

TPU Pod を作成すると、TPU ノードごとに個別の VM が作成されます。デフォルトでは、各 TPU VM に異なるユーザー ID(UID)が割り当てられます。これにより、複数のノードから同じ NFS ディレクトリにアクセスしようとすると問題が発生します。同じディレクトリでは、ノードが異なるとオーナーも異なるため、標準の Linux 権限はノード間で適用されません。たとえば、あるノードのプロセスは、別のノードによって作成されたログ ディレクトリに書き込むことができません。

この問題を回避するには、OS Login を有効にします。OS Login では、特定の VM インスタンスまたはプロジェクトを構成できます。詳細については、OS Login の設定をご覧ください。

TPU Pod での開発に向けたワークフローのベスト プラクティス

新しい TPU ワークロードを開発する場合は、通常、最小の TPU で開発を開始し、より大きな TPU サイズを徐々に反復していくのが最適です。最小の TPU(v2-8 や v3-8 など)から始めます。

  • ワークロードの機能をテストする
  • パフォーマンス ツールを使用してパフォーマンスをテストして検証する

ワークロードが機能し、パフォーマンスの目標を達成したら、v2-32 または v3-32 にスケールアップします。目的の TPU サイズに達するまで、TPU のサイズを徐々に増やして反復処理ながら、スケーラビリティ(機能とパフォーマンス)を検証します。