TPU の構成
v5p TPU の構成
TPU v5p Pod は、再構成可能な高速リンクで相互接続された 8,960 個のチップで構成されます。TPU v5p' の柔軟なネットワーキングでは、さまざまな方法で同じサイズのスライスにチップを接続できます。gcloud compute tpus tpu-vm create コマンドを使用して TPU スライスを作成する場合は、AcceleratorType パラメータまたは AcceleratorConfig パラメータを使用してそのタイプとシェイプを指定します。
次の表では、v5p でサポートされている最も一般的なシングル スライスの形状と、1 キューブより大きいほとんどの（全部ではない）キューブの形状を示します。最大の v5p 形状は 16x16x24（6,144 チップ、96 キューブ）です。
| スライス形状 | VM サイズ | コア数 | チップ数 | マシン数 | キューブ数 | Twisted のサポート |
| 2x2x1 | フルホスト | 8 | 4 | 1 | なし | なし |
| 2x2x2 | フルホスト | 16 | 8 | 2 | なし | なし |
| 2x4x4 | フルホスト | 64 | 32 | 8 | なし | なし |
| 4x4x4 | フルホスト | 128 | 64 | 16 | 1 | なし |
| 4x4x8 | フルホスト | 256 | 128 | 32 | 2 | ○ |
| 4x8x8 | フルホスト | 512 | 256 | 64 | 4 | ○ |
| 8x8x8 | フルホスト | 1024 | 512 | 128 | 8 | なし |
| 8x8x16 | フルホスト | 2048 | 1024 | 256 | 16 | ○ |
| 8x16x16 | フルホスト | 4,096 | 2048 | 512 | 32 | ○ |
| 16x16x16 | フルホスト | 8,192 | 4,096 | 1024 | 64 | なし |
| 16x16x24 | フルホスト | 12,288 | 6,144 | 1,536 | 96 | なし |
シングル スライス トレーニングは最大 6,144 チップに対応します。マルチスライスを使用すると、18,432 チップまで拡張できます。 マルチスライスの詳細については、Cloud TPU マルチスライスの概要をご覧ください。
AcceleratorType パラメータの使用
TPU リソースを割り当てる場合は、--accelerator-type 引数を使用してスライス内の TensorCore 数を指定します。--accelerator-type は書式付きの文字列「v$VERSION_NUMBERv-$CORES_COUNT」です。たとえば、v5p-32 は、32 個の TensorCore（16 チップ）がある v5p TPU スライスを指定します。
v5p トレーニング ジョブに TPU をプロビジョニングするには、CLI または TPU API 作成リクエストで次のいずれかのアクセラレータ タイプを使用します。
- v5p-8
- v5p-16
- v5p-32
- v5p-64
- v5p-128（1 フルキューブ/ラック）
- v5p-256（2 キューブ）
- v5p-512
- v5p-1024～v5p-12288
AcceleratorConfig パラメータの使用
v5p 以降の Cloud TPU バージョンでは、AcceleratorConfig が Cloud TPU v4 の場合とほぼ同じ方法で使用されます。違いは、TPU タイプを --type=v4 として指定する代わりに、使用している TPU バージョン（例: v5p リリースの場合は --type=v5p）として指定することです。
Cloud TPU の ICI の復元性
ICI の復元性により、キューブ間で TPU を接続する光リンクと光回路スイッチ（OCS）のフォールト トレランスが向上します。（キューブ内の ICI 接続では、影響を受けない銅リンクが使用されます）。ICI 復元力により、ICI 接続は OCS 障害と光 ICI 障害を回避して経路設定されます。その結果、TPU スライスのスケジューリング可用性が改善されますが、ICI のパフォーマンスが一時的に低下するというトレードオフがあります。
Cloud TPU v4 と同様に、1 キューブ以上の v5p スライスでは、ICI の復元力がデフォルトで有効になります。
- v5p-128（アクセラレータ タイプを指定する場合）
- 4x4x4（アクセラレータ構成を指定する場合）
VM、ホスト、スライスのプロパティ
| プロパティ | TPU の値 |
| v5p チップの数 | 4 |
| vCPU の数 | 208（NUMA バインディングを使用して、クロス NUMA パフォーマンスのペナルティを回避する場合は半分のみ使用可能） |
| RAM（GB） | 448（NUMA バインディングを使用して、クロス NUMA パフォーマンスのペナルティを回避する場合は半分のみ使用可能） |
| NUMA ノードの数 | 2 |
| NIC のスループット（Gbps） | 200 |
Pod 内の TensorCore 数、チップ数、ホスト/VM 数、キューブ間の関係:
| コア数 | チップ | ホスト/VM | キューブ | |
|---|---|---|---|---|
| ホスト | 8 | 4 | 1 | |
| キューブ（別名ラック） | 128 | 64 | 16 | 1 |
| サポートされている最大のスライス | 12,288 | 6,144 | 1,536 | 96 |
| v5p フル Pod | 17,920 | 8,960 | 2,240 | 140 |
TPU v5e の構成
Cloud TPU v5e は、トレーニングと推論（サービス提供）を組み合わせたプロダクトです。トレーニング環境と推論環境を区別するには、TPU API で AcceleratorType フラグまたは AcceleratorConfig フラグを使用するか、GKE ノードプールを作成するときに --machine-type フラグを使用します。
トレーニング ジョブはスループットと可用性に対して最適化されていますが、サービス提供ジョブはレイテンシに対して最適化されています。そのため、サービス提供用にプロビジョニングされた TPU でのトレーニング ジョブの可用性が低下する可能性があります。同様に、トレーニング用にプロビジョニングされた TPU で実行されるサービス提供ジョブでも、レイテンシが高くなる可能性があります。
AcceleratorType を使用して、使用する TensorCore の数を指定します。
AcceleratorType は、gcloud CLI または Google Cloud コンソールを使用して TPU を作成するときに指定します。AcceleratorType に指定する値は、v$VERSION_NUMBER-$CHIP_COUNT の形式の文字列です。
AcceleratorConfig を使用して、使用する TensorCore の数を指定することもできます。ただし、TPU v5e にはカスタム 2D トポロジ バリアントがないため、AcceleratorConfig の使用と AcceleratorType の使用に違いはありません。
AcceleratorConfig を使用して TPU v5e を構成するには、--version フラグと --topology フラグを使用します。--version を、使用する TPU バージョンに設定し、--topology をスライスの TPU チップの物理配置に設定します。AcceleratorConfig に指定する値は AxB 形式の文字列です。ここで、A と B は各方向のチップ数です。
v5e では、次の 2D スライスの形状がサポートされています。
| トポロジ | TPU チップ数 | ホストの数 |
| 1×1 | 1 | 1/8 |
| 2x2 | 4 | 1/2 |
| 2x4 | 8 | 1 |
| 4x4 | 16 | 2 |
| 4x8 | 32 | 4 |
| 8x8 | 64 | 8 |
| 8x16 | 128 | 16 |
| 16x16 | 256 | 32 |
v5e TPU スライスの各 TPU VM には 1 つ、4 つ、または 8 つのチップが含まれています。4 チップ以下の小さなスライスでは、すべての TPU チップが同じ非均一メモリアクセス（NUMA）ノードを共有します。
8 チップ v5e TPU VM の場合、NUMA パーティション内では CPU-TPU 通信の方が効率的です。たとえば、次の図では、CPU0-Chip0 通信は CPU0-Chip4 通信よりも高速です。
サービス提供用の Cloud TPU v5e タイプ
単一ホストの提供は、最大 8 つの v5e チップでサポートされます。次の構成がサポートされています（1x1、2x2、2x4 スライス）。各スライスには、それぞれ 1、4、8 チップが搭載されています。
サービス提供ジョブに TPU をプロビジョニングするには、CLI または API TPU 作成リクエストで次のいずれかのアクセラレータ タイプを使用します。
| AcceleratorType（TPU API） | マシンタイプ（GKE API） |
|---|---|
v5litepod-1 |
ct5lp-hightpu-1t |
v5litepod-4 |
ct5lp-hightpu-4t |
v5litepod-8 |
ct5lp-hightpu-8t |
8 を超える v5e チップ（マルチホスト サービングとも呼ばれる）でのサービス提供は、Sax を使用することでサポートされます。詳細については、大規模言語モデルの提供をご覧ください。
トレーニング用の Cloud TPU v5e タイプ
トレーニングは最大 256 個のチップに対応しています。
v5e トレーニング ジョブに TPU をプロビジョニングするには、CLI または API TPU 作成リクエストで次のいずれかのアクセラレータ タイプを使用します。
| AcceleratorType（TPU API） | マシンタイプ（GKE API） | トポロジ |
|---|---|---|
v5litepod-16 |
ct5lp-hightpu-4t |
4x4 |
v5litepod-32 |
ct5lp-hightpu-4t |
4x8 |
v5litepod-64 |
ct5lp-hightpu-4t |
8x8 |
v5litepod-128 |
ct5lp-hightpu-4t |
8x16 |
v5litepod-256 |
ct5lp-hightpu-4t |
16x16 |
v5e TPU VM タイプの比較:
| VM のタイプ | n2d-48-24-v5lite-tpu | n2d-192-112-v5lite-tpu | n2d-384-224-v5lite-tpu |
| # of v5e chips | 1 | 4 | 8 |
| vCPU の数 | 24 | 112 | 224 |
| RAM（GB） | 48 | 192 | 384 |
| NUMA ノードの数 | 1 | 1 | 2 |
| 適用先 | v5litepod-1 | v5litepod-4 | v5litepod-8 |
| 中断 | 高 | 中 | 低 |
より多くのチップを必要とするワークロードのために、スケジューラはチップ数の少ない VM をプリエンプトすることがあります。したがって、8 チップの VM は 1 チップ VM と 4 チップ VM をプリエンプトする可能性が高くなります。
TPU v4 の構成
TPU v4 Pod は、再構成可能な高速リンクで相互接続された 4,096 のチップで構成されます。TPU v4 の柔軟なネットワーキングでは、さまざまな方法で同じサイズの Pod スライスにチップを接続できます。TPU Pod スライスを作成する場合は、TPU のバージョンと必要な TPU リソースの数を指定します。TPU v4 Pod スライスを作成する場合は、AcceleratorType と AcceleratorConfig のいずれかでタイプとサイズを指定できます。
AcceleratorType の使用
トポロジを指定していない場合は、AcceleratorType を使用します。AcceleratorType を使用して v4 TPU を構成するには、TPU Pod スライスの作成時に --accelerator-type フラグを使用します。--accelerator-type を、TPU のバージョンと使用する TensorCore の数を含む文字列に設定します。たとえば、32 個の TensorCore を使用して v4 Pod スライスを作成するには、--accelerator-type=v4-32 を使用します。
次のコマンドは、--accelerator-type フラグを使用して 512 の TensorCore を搭載した v4 TPU Pod スライスを作成します。
$ gcloud compute tpus tpu-vm create tpu-name
--zone=zone
--accelerator-type=v4-512
--version=tpu-vm-tf-2.16.1-pod-pjrt
TPU バージョン（v4）の後の数字は、TensorCore の数を指定しています。v4 TPU には 2 つの TensorCore があるため、TPU チップの数は 512/2 = 256 になります。
AcceleratorConfig の使用
TPU スライスの物理トポロジをカスタマイズする場合は、AcceleratorConfig を使用します。これは通常、256 チップを超える Pod スライスを使用したパフォーマンス調整に必要です。
AcceleratorConfig を使用して v4 TPU を構成するには、--version フラグと --topology フラグを使用します。--version を、使用する TPU バージョンに設定し、--topology を Pod スライスの TPU チップの物理配置に設定します。
3 タプル AxBxC を使用して TPU トポロジを指定します。ここで、A<=B<=C で、A、B、C はすべて 4 以下または 4 の整数倍です。値 A、B、C は、3 つのディメンションのチップ数です。たとえば、16 チップの v4 Pod スライスを作成するには、--version=v4 と --topology=2x2x4 を設定します。
次のコマンドは、128 個の TPU チップを 4x4x8 のアレイに配置した v4 TPU Pod スライスを作成します。
$ gcloud compute tpus tpu-vm create tpu-name
--zone=zone
--type=v4
--topology=4x4x8
--version=tpu-vm-tf-2.16.1-pod-pjrt
2A=B=C または 2A=2B=C のトポロジでも、すべての通信用に最適化されたトポロジ バリアントがあります（例: 4×4×8、8×8x16、12×12×24）。これらをツイスト トリトポロジといいます。
次の図は、一般的な TPU v4 トポロジを示しています。
より大きな Pod スライスを、1 つまたは複数の 4x4x4 の「キューブ」チップから構築できます。
ツイスト トリ トポロジ
一部の v4 3D トーラス スライスには、ツイスト トポロジ トポロジとして知られているものを使用するオプションがあります。たとえば、2 つの v4 キューブは 4x4x8 スライスまたは 4x4x8_twisted として配置できます。ツイスト トポロジでは、二分帯域幅が大幅に高くなります。全二分帯域幅を増やすとグローバル通信パターンを使用するワークロードで役立ちます。ツイスト トポロジでは、ほとんどのモデルのパフォーマンスが向上します。ただし、大規模な TPU 埋め込みワークロードのほうがメリットがあります。
データ並列処理を唯一の並列処理戦略として使用するワークロードでは、ツイスト トポロジのパフォーマンスがわずかに向上する可能性があります。LLM の場合、ツイスト トポロジを使用したパフォーマンスは、並列処理のタイプ（DP、MP など）によって異なる場合があります。モデルに最適なパフォーマンスを判断するには、ツイスト トポロジを使用して LLM をトレーニングすることをおすすめします。FSDP MaxText モデルの一部の実験では、ツイスト トポロジを使用して MFU が 1～2 回改善されました。
ツイスト トポロジの主な利点は、非対称トポロジ（たとえば、4×4×8）を密接に関連する対称トポロジに変換できることです。対称トポロジには次のような多くの利点があります。
- 負荷分散の改善
- 全二分帯域幅
- パケット ルートが短い
これらの利点は、最終的に、多くのグローバル通信パターンのパフォーマンス向上につながります。
TPU ソフトウェアは、各ディメンションのサイズが最小のディメンションのサイズ（例: 4x4x8、4×8×8、12x12x24）に等しいか、2 倍であるスライスのツイスト トリをサポートします。
例として、スライス内の（X、Y）座標にラベルが付けられた TPU を備えた 4×2 トーラス トポロジについて考えてみます。
このトポロジグラフのエッジは、わかりやすくするために無向エッジとして示しています。実際には、各エッジは TPU 間の双方向接続です。図のとおり、このグリッドの一方の側と反対側の端とを囲むエッジを、ラップアラウンド エッジと呼んでいます。
このトポロジをツイストすると、完全に対称な 4×2 ツイスト トーラス トポロジになります。
この図と前の図で変更されたのは、Y ラップアラウンド エッジのみです。同じ X 座標で別の TPU に接続する代わりに、座標 X+2 mod 4 で TPU に接続するようにシフトされています。
同じ考え方を、さまざまなディメンション サイズとさまざまな数のディメンションに一般化できます。結果として得られるネットワークは、各ディメンションが最小ディメンションのサイズと等しいか、2 倍である限り、対称になります。
次の表は、サポートされているツイスト トポロジと、ツイストされていないトポロジを使用した場合の二分帯域幅の理論的な増加を示しています。
| チップ トポロジ | 二分帯域幅の理論上の増加 とツイストされていないトポロジ |
|---|---|
| 4×4×8_twisted | 最大 70% |
| 8x8x16_twisted | |
| 12×12×24_twisted | |
| 4×8×8_twisted | ～40% |
| 8×16×16_twisted |
TPU v4 トポロジのバリアント
同じ数のチップを含む一部のトポロジは、さまざまな方法で配置できます。たとえば、512 チップ（1,024 Tensor コア）の TPU Pod スライスは、4x4x32、4x8x16、8x8x8 のトポロジを使用して構成できます。2,048 チップ（4,096 TensorCores）の TPU Pod スライスは、さらに多くの 4x4x128、4x8x64、4x16x32、8x16x16 のトポロジ オプションを提供します。2,048 チップ（4,096 TensorCores）の TPU Pod スライスは、さらに多くの 4x4x128、4x8x64、4x16x32、8x16x16 のトポロジ オプションを提供します。
特定のチップ数に関連付けられるデフォルトのトポロジは、キューブに最も近いものです（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 のトポロジを使用すると最適です。
TPU v4 構成は、64 チップより小さいトポロジ（小規模トポロジ）と 64 チップを超えるトポロジ（大規模トポロジ）の 2 つのグループで構成されています。
小規模な v4 トポロジ
Cloud TPU は、次の 64 チップよりも小さい TPU v4 Pod スライス（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 Pod スライスは 64 チップ単位で利用でき、3 つのディメンションすべてで 4 の倍数の形状になっています。ディメンションも昇順にする必要があります。次の表に、いくつかの例を示します。これらのトポロジの一部は、--type フラグと --topology フラグを使用してのみ作成できる「カスタム」トポロジです。これは、チップを配置する方法は複数あるためです。
次のコマンドは、512 個の TPU チップを 8x8x8 のアレイに配置した v4 TPU Pod スライスを作成します。
$ gcloud compute tpus tpu-vm create tpu-name
--zone=zone
--type=v4
--topology=8x8x8
--version=tpu-vm-tf-2.16.1-pod-pjrt
--accelerator-type を使用して、同じ数の TensorCore を持つ v4 TPU Pod スライスを作成できます。
$ gcloud compute tpus tpu-vm create tpu-name
--zone=zone
--accelerator-type=v4-1024
--version=tpu-vm-tf-2.16.1-pod-pjrt
| 名前（TensorCore のカウントに基づく） | チップ数 | トポロジ |
| v4-128 | 64 | 4x4x4 |
| v4-256 | 128 | 4x4x8 |
| v4-512 | 256 | 4x8x8 |
なし - --type フラグと --topology フラグを使用する必要があります。 |
256 | 4x4x16 |
| v4-1024 | 512 | 8x8x8 |
| v4-1536 | 768 | 8x8x12 |
| v4-2048 | 1024 | 8x8x16 |
なし - --type フラグと --topology フラグを使用する必要があります。 |
1024 | 4x16x16 |
| v4-4096 | 2048 | 8x16x16 |
| … | … | … |
TPU v3 の構成
TPU v3 Pod は、高速リンクで相互接続された 1,024 個のチップで構成されます。TPU v3 デバイスまたは Pod スライスを作成するには、gcloud compute tpus tpu-vm コマンドに --accelerator-type フラグを使用します。アクセラレータ タイプを指定するには、TPU のバージョンと TPU コアの数を指定します。単一の v3 TPU の場合は、--accelerator-type=v3-8 を使用します。128 個の TensorCore を使用する v3 Pod スライスの場合は、--accelerator-type=v3-128 を使用します。
次のコマンドは、128 個の TensorCore を使用して v2 TPU Pod スライスを作成する方法を示しています。
$ gcloud compute tpus tpu-vm create tpu-name
--zone=zone
--accelerator-type=v3-128
--version=tpu-vm-tf-2.16.1-pjrt
次の表に、サポートされている v2 TPU タイプを示します。
| TPU バージョン | サポートの終了 |
|---|---|
| v3-8 | （終了日未定） |
| v3-32 | （終了日未定） |
| v3-128 | （終了日未定） |
| v3-256 | （終了日未定） |
| v3-512 | （終了日未定） |
| v3-1024 | （終了日未定） |
| v3-2048 | （終了日未定） |
TPU の管理の詳細については、TPU の管理をご覧ください。Cloud TPU のさまざまなバージョンの詳細については、システム アーキテクチャをご覧ください。
TPU v2 の構成
TPU v2 Pod は、再構成可能な高速リンクで相互接続された 512 のチップで構成されます。TPU v2 Pod スライスを作成するには、gcloud compute tpus tpu-vm コマンドに --accelerator-type フラグを使用します。アクセラレータ タイプを指定するには、TPU のバージョンと TPU コアの数を指定します。単一の v2 TPU の場合は、--accelerator-type=v2-8 を使用します。128 個の TensorCore を使用する v3 Pod スライスの場合は、--accelerator-type=v2-128 を使用します。
次のコマンドは、128 個の TensorCore を使用して v2 TPU Pod スライスを作成する方法を示しています。
$ gcloud compute tpus tpu-vm create tpu-name
--zone=zone
--accelerator-type=v2-128
--version=tpu-vm-tf-2.16.1-pjrt
TPU の管理の詳細については、TPU の管理をご覧ください。Cloud TPU のさまざまなバージョンの詳細については、システム アーキテクチャをご覧ください。
次の表に、サポートされている v2 TPU タイプを示します。
| TPU バージョン | サポートの終了 |
|---|---|
| v2-8 | （終了日未定） |
| v2-32 | （終了日未定） |
| v2-128 | （終了日未定） |
| v2-256 | （終了日未定） |
| v2-512 | （終了日未定） |
TPU タイプの互換性
TPU タイプを、同じ TensorCore 数を持つ別の TPU タイプ（たとえば、v3-128 や v4-128）に変更し、コードを変更せずにトレーニング スクリプトを実行できます。ただし、TensorCores コア数を増減させた TPU タイプに変更する場合は、かなりの調整と最適化を行う必要があります。詳細については、TPU Pod でのトレーニングをご覧ください。
TPU VM ソフトウェア バージョン
このセクションでは、TPU VM アーキテクチャで TPU に対して使用する TPU ソフトウェア バージョンについて説明します。 TPU ノードのアーキテクチャについては、TPU ノードのソフトウェア バージョンをご覧ください。
TPU ソフトウェア バージョンは、TensorFlow、PyTorch、JAX のフレームワークで使用できます。
TensorFlow
モデルの書き込まれた TensorFlow のバージョンと一致する TPU ソフトウェア バージョンを使用します。
TensorFlow 2.15.0 以降では、ストリーム エグゼキュータ（SE）ランタイムまたは PJRT ランタイムも指定する必要があります。たとえば、PJRT ランタイムで TensorFlow 2.16.1 を使用している場合は、tpu-vm-tf-2.16.1-pjrt TPU ソフトウェア バージョンを使用します。TensorFlow 2.15.0 より前のバージョンは、ストリーム エグゼキュータのみをサポートしています。PJRT の詳細については、TensorFlow PJRT のサポートをご覧ください。
現在サポートされている TensorFlow TPU VM ソフトウェア バージョンは次のとおりです。
- tpu-vm-tf-2.16.1-pjrt
- tpu-vm-tf-2.16.1-se
- tpu-vm-tf-2.15.0-pjrt
- tpu-vm-tf-2.15.0-se
- tpu-vm-tf-2.14.1
- tpu-vm-tf-2.14.0
- tpu-vm-tf-2.13.1
- tpu-vm-tf-2.13.0
- tpu-vm-tf-2.12.1
- tpu-vm-tf-2.12.0
- tpu-vm-tf-2.11.1
- tpu-vm-tf-2.11.0
- tpu-vm-tf-2.10.1
- tpu-vm-tf-2.10.0
- tpu-vm-tf-2.9.3
- tpu-vm-tf-2.9.1
- tpu-vm-tf-2.8.4
- tpu-vm-tf-2.8.3
- tpu-vm-tf-2.8.0
- tpu-vm-tf-2.7.4
- tpu-vm-tf-2.7.3
TensorFlow パッチ バージョンの詳細については、サポートされている TensorFlow パッチ バージョンをご覧ください。
TensorFlow PJRT のサポート
TensorFlow 2.15.0 以降では、TPU で TensorFlow に PJRT インターフェースを使用できます。PJRT は自動デバイス メモリ デフラグメンテーション機能を備えており、ハードウェアとフレームワークのインテグレーションを簡素化します。PJRT の詳細については、Google オープンソース ブログの PJRT: ML ハードウェアとフレームワークのインテグレーションの簡素化をご覧ください。
TPU v2、v3、v4 のすべての機能が PJRT ランタイムに移行されたわけではありません。次の表に、PJRT またはストリーム エグゼキュータでサポートされている機能を示します。
| アクセラレータ | 特徴 | PJRT でサポートされている | ストリーム エグゼキュータでサポートされている |
|---|---|---|---|
| TPU v2~v4 | 高密度なコンピューティング（TPU 埋め込み API なし） | ○ | ○ |
| TPU v2~v4 | 高密度なコンピューティング API + TPU 埋め込み API | いいえ | ○ |
| TPU v2~v4 | ソフトデバイス配置で tf.summary/tf.print |
いいえ | ○ |
| TPU v5e | 高密度なコンピューティング（TPU 埋め込み API なし） | ○ | いいえ |
| TPU v5e | TPU 埋め込み API | なし - TPU v5e では TPU 埋め込み API がサポートされません | なし |
| TPU v5p | 高密度なコンピューティング（TPU 埋め込み API なし） | ○ | × |
| TPU v5p | TPU 埋め込み API | ○ | × |
TensorFlow バージョン 2.10.0 以前を使用した TPU v4
TensorFlow で TPU v4 を使用してモデルをトレーニングする場合、TensorFlow バージョン 2.10.0 以前では、次の表に示す v4 固有のバージョンを使用します。使用している TensorFlow のバージョンが表に表示されていない場合は、TensorFlow セクションのガイダンスに従ってください。
| TensorFlow のバージョン | TPU ソフトウェア バージョン |
|---|---|
| 2.10.0 | tpu-vm-tf-2.10.0-v4、tpu-vm-tf-2.10.0-pod-v4 |
| 2.9.3 | tpu-vm-tf-2.9.3-v4、tpu-vm-tf-2.9.3-pod-v4 |
| 2.9.2 | tpu-vm-tf-2.9.2-v4、tpu-vm-tf-2.9.2-pod-v4 |
| 2.9.1 | tpu-vm-tf-2.9.1-v4、tpu-vm-tf-2.9.1-pod-v4 |
Libtpu のバージョン
TPU VM は、TensorFlow を使用して作成され、対応する Libtpu ライブラリがプリインストールされます。独自の VM イメージを作成する場合は、次の TensorFlow TPU ソフトウェア バージョンと対応する libtpu バージョンを指定します。
| TensorFlow のバージョン | libtpu.so のバージョン |
|---|---|
| 2.16.1 | 1.10.1 |
| 2.15.0 | 1.9.0 |
| 2.14.1 | 1.8.1 |
| 2.14.0 | 1.8.0 |
| 2.13.1 | 1.7.1 |
| 2.13.0 | 1.7.0 |
| 2.12.1 | 1.6.1 |
| 2.12.0 | 1.6.0 |
| 2.11.1 | 1.5.1 |
| 2.11.0 | 1.5.0 |
| 2.10.1 | 1.4.1 |
| 2.10.0 | 1.4.0 |
| 2.9.3 | 1.3.2 |
| 2.9.1 | 1.3.0 |
| 2.8.3 | 1.2.3 |
| 2.8.* | 1.2.0 |
| 2.7.3 | 1.1.2 |
PyTorch
モデルの書き込まれた PyTorch のバージョンと一致する TPU ソフトウェア バージョンを使用します。たとえば、PyTorch 1.13 と TPU v2 または v3 を使用している場合は、tpu-vm-pt-1.13 TPU ソフトウェア バージョンを使用します。TPU v4 を使用している場合は、tpu-vm-v4-pt-1.13 TPU ソフトウェア バージョンを使用します。TPU Pod に同じ TPU ソフトウェア バージョン（v2-32、v3-128、v4-32 など）が使用されます。現在サポートされている TPU ソフトウェア バージョンは次のとおりです。
TPU v2/v3:
- tpu-vm-pt-2.0（pytorch-2.0）
- tpu-vm-pt-1.13（pytorch-1.13）
- tpu-vm-pt-1.12（pytorch-1.12）
- tpu-vm-pt-1.11（pytorch-1.11）
- tpu-vm-pt-1.10（pytorch-1.10）
- v2-alpha（pytorch-1.8.1）
TPU v4:
- tpu-vm-v4-pt-2.0（pytorch-2.0）
- tpu-vm-v4-pt-1.13 (pytorch-1.13)
TPU v5（v5e と v5p）:
- v2-alpha-tpuv5（pytorch-2.0）
TPU VM を作成すると、TPU VM に最新バージョンの PyTorch がプリインストールされます。PyTorch をインストールすると、libtpu.so の正しいバージョンが自動的にインストールされます。
現在の PyTorch ソフトウェア バージョンを変更するには、PyTorch バージョンの変更をご覧ください。
JAX
TPU VM に JAX を手動でインストールする必要があります。TPU v2 と v3 用の JAX 固有の TPU ソフトウェア（ランタイム）バージョンはありません。それ以降の TPU バージョンでは、次のソフトウェア バージョンを使用します。
- TPU v4: tpu-vm-v4-base
- TPU v5e: v2-alpha-tpuv5
- TPU v5p: v2-alpha-tpuv5
JAX をインストールすると、libtpu.so の正しいバージョンが自動的にインストールされます。
TPU ノードのソフトウェア バージョン
このセクションでは、TPU ノード アーキテクチャで TPU に対して使用する TPU ソフトウェア バージョンについて説明します。TPU VM アーキテクチャについては、TPU VM ソフトウェア バージョンをご覧ください。
TPU ソフトウェア バージョンは、TensorFlow、PyTorch、JAX のフレームワークで使用できます。
TensorFlow
モデルの書き込まれた TensorFlow のバージョンと一致する TPU ソフトウェア バージョンを使用します。たとえば、TensorFlow 2.12.0 を使用している場合は、2.12.0 TPU ソフトウェア バージョンを使用します。TensorFlow 固有の TPU ソフトウェア バージョンは次のとおりです。
- 2.12.1
- 2.12.0
- 2.11.1
- 2.11.0
- 2.10.1
- 2.10.0
- 2.9.3
- 2.9.1
- 2.8.4
- 2.8.2
- 2.7.3
TensorFlow パッチ バージョンの詳細については、サポートされている TensorFlow パッチ バージョンをご覧ください。
TPU ノードを作成すると、最新バージョンの TensorFlow が TPU ノードにプリインストールされます。
PyTorch
モデルの書き込まれた PyTorch のバージョンと一致する TPU ソフトウェア バージョンを使用します。たとえば、PyTorch 1.9 を使用している場合は、pytorch-1.9 ソフトウェア バージョンを使用します。
PyTorch 固有の TPU ソフトウェア バージョンは次のとおりです。
- pytorch-2.0
- pytorch-1.13
- pytorch-1.12
- pytorch-1.11
- pytorch-1.10
- pytorch-1.9
- pytorch-1.8
- pytorch-1.7
pytorch-1.6
pytorch-nightly
TPU ノードを作成すると、TPU ノードに最新バージョンの PyTorch がプリインストールされます。
JAX
TPU VM に JAX を手動でインストールする必要があるため、JAX 固有の TPU ソフトウェア バージョンはプリインストールされていません。TensorFlow に含まれるソフトウェア バージョンはすべて使用できます。
次のステップ
- TPU アーキテクチャの詳細については、システム アーキテクチャ ページをご覧ください。
- TPU の用途を参照して、Cloud TPU に最適なモデルの種類について確認する。