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 でのトレーニングをご覧ください。

次のステップ