TPU のタイプとトポロジ
TPU の構成
TPU を作成するときに、TPU 構成を指定する必要があります。TensorCore または TPU チップの観点から TPU の構成を指定できます。チップの観点から TPU の構成を指定する必要があるのは、大規模なモデルをトレーニングしてパフォーマンスを調整する高度なシナリオだけです。
TensorCore ベースの構成
v2 ~ v4 TPU では、TensorCore ベースの構成を使用できます。TPU バージョンと使用する TensorCore の数を使用して、TensorCore ベースの構成を記述します。たとえば、v4-128
は v4 TPU と 128 個の TensorCore の構成を指定します。gcloud compute tpus tpu-vm create コマンドの --accelerator-type
フラグを使用して、TensorCore ベースの構成を指定します。
$ gcloud compute tpus tpu-vm create tpu-name --zone=zone --accelerator-type=v3-128 --version=tpu-vm-tf-2.12.0
チップベースの TPU 構成
v4 TPU 以降の TPU バージョンでは、TPU チップベースの構成を使用できます。チップベースの構成は、TPU バージョンとトポロジ(TPU チップの物理配列)で記述します。「v4
」などの文字列で TPU バージョンを指定します。チップを 3 次元で配置する方法を記述する 3 タプルを使用してトポロジを指定します。チップベースの構成には、gcloud compute tpus tpu-vm create コマンドの --version
と --topology
フラグを指定します。たとえば、次のコマンドでは TPU バージョンとして TPU v4 を、トポロジとして 2x2x4
(それぞれ 2 つの TensorCore を含む 16 チップ)を指定しています。
$ gcloud compute tpus tpu-vm create tpu-name --zone=zone --type=v4 --topology=2x2x4 --version=tpu-vm-tf-2.12.0
Cloud TPU のバージョンとトポロジの詳細については、システム アーキテクチャをご覧ください。
v4 トポロジの表現は AxBxC
です(ここで、A <=B <=C、A、B、C はすべて 4 以下、または 4 の整数倍)。値 A、B、C は、3 つのディメンションのそれぞれのチップ数です。2A=B=C または 2A=2B=C のトポロジでも、すべての通信用に最適化されたトポロジ バリアントがあります(例: 4×4×8、8×8x16、12×12×24)。
v4 トポロジの詳細については、TPU v4 の構成をご覧ください。
トポロジのバリアント
特定のチップ数に関連付けられたデフォルトのトポロジは、キューブに最も近いものです(トポロジの形状を参照)。この形状は、データの並列 ML トレーニングに最適です。他のトポロジ(モデルとデータの並列処理、シミュレーションの空間パーティショニングなど)は、複数の種類の並列処理のワークロードで役立ちます。これらのワークロードは、トポロジが使用された並列処理に一致する場合に最も効果的です。たとえば、X 次元に 4 方向モデル並列処理、Y ディメンションと Z ディメンションに 256 方向データ並列処理を配置すると、4x16x16 トポロジと一致します。
並列処理のディメンションが複数あるモデルは、並列処理ディメンションが TPU トポロジ ディメンションにマッピングされている場合に最もパフォーマンスを発揮します。これらは通常、データとモデル並列大規模言語モデル(LLM)です。たとえば、トポロジ 8x16x16 の TPU v4 Pod スライスの場合、TPU トポロジのディメンションは 8、16、16 です。8 方向または 16 方向のモデル並列処理(物理 TPU トポロジディメンションの 1 つにマッピング)を使用したほうがパフォーマンスが向上します。4 方向モデル並列処理は、このトポロジには最適ではありません。TPU トポロジのディメンションと整合していませんが、同じ数のチップで 4x16x32 のトポロジを使用すると最適です。
小規模な v4 トポロジ
Cloud TPU は、次の 64 チップよりも小さい TPU v4 スライス(4x4x4 キューブ)をサポートしています。これらの小規模な v4 トポロジは、TensorCore ベースの名前(v4-32 など)またはトポロジ(2x2x4 など)を使用して作成できます。
名前(TensorCore のカウントに基づく) | チップ数 | トポロジ |
v4-8 | 4 | 2x2x1 |
v4-16 | 8 | 2x2x2 |
v4-32 | 16 | 2x2x4 |
v4-64 | 32 | 2x4x4 |
大規模な v4 トポロジ
TPU v4 スライスは 64 チップ単位で利用でき、3 つのディメンションすべてで 4 の倍数の形状になっています。ディメンションも昇順でなければなりません。次の表にいくつかの例を示します。これらのトポロジのいくつかは、トポロジ API を使用してのみ起動できる「カスタム」トポロジです。これは、一般的に使用される名前付きトポロジと同じチップ数を持つためです。 。
名前(TensorCore のカウントに基づく) | チップ数 | トポロジ |
v4-128 | 64 | 4x4x4 |
v4-256 | 128 | 4x4x8 |
v4-512 | 256 | 4x8x8 |
Topology API をご覧ください | 256 | 4x4x16 |
v4-1024 | 512 | 8x8x8 |
v4-1536 | 768 | 8x8x12 |
v4-2048 | 1024 | 8x8x16 |
Topology API をご覧ください | 1024 | 4x16x16 |
v4-4096 | 2048 | 8x16x16 |
… | … | … |
トポロジ API
カスタム トポロジを使用して Cloud TPU Pod スライスを作成するには、gcloud
TPU API を次のように使用します。
$ gcloud compute tpus tpu-vm create tpu-name \ --zone=us-central2-b \ --subnetwork=tpusubnet \ --type=v4 \ --topology=4x4x16 \ --version=runtime-version
トポロジの形状
次の図は、3 つのトポロジトポロジと、関連する topology
フラグを示しています。
TPU タイプの互換性
TPU タイプを、同じ TensorCore 数を持つ別の TPU タイプ(たとえば、v3-128
や v4-128
)に変更し、コードを変更せずにトレーニング スクリプトを実行できます。ただし、チップ数または TensorCore 数の増減を実施した TPU タイプに変更する場合は、かなりの調整と最適化を行う必要があります。詳細については、TPU Pod でのトレーニングをご覧ください。
次のステップ
- TPU アーキテクチャの詳細については、システム アーキテクチャ ページをご覧ください。
- TPU の用途を参照して、Cloud TPU に最適なモデルの種類について確認する。