TPU v5e

このドキュメントでは、Cloud TPU v5e のアーキテクチャとサポートされている構成について説明します。

TPU v5e では、単一ホストとマルチホストのトレーニングと単一ホストの推論がサポートされています。マルチホストの推論は Sax を使用することでサポートされます。詳細については、大規模言語モデルのサービス提供をご覧ください。

システム アーキテクチャ

各 v5e チップには 1 つの TensorCore が含まれています。各 TensorCore には、4 つのマトリックス乗算ユニット(MXU)、1 つのベクトル ユニット、1 つのスカラー ユニットがあります。

次の図は、TPU v5 チップを示しています。

v5e チップの図

次の表に、v5e の主なチップ仕様とその値を示します。

主なチップ仕様 v5e の値
チップあたりのピーク コンピューティング(bf16) 197 TFLOPs
HBM2 の容量と帯域幅 16 GB, 819 GBps
インターチップ相互接続 BW 1600 Gbps

次の表に、v5e の Pod の仕様とその値を示します。

Pod の主な仕様 v5e の値
TPU Pod のサイズ 256 チップ
相互接続トポロジ 2D トーラス
Pod あたりのピーク コンピューティング 100 PetaOps(Int8)
Pod あたりの all-reduce 帯域幅 51.2 TB/s
Pod あたりの二分割帯域幅 1.6 TB/s
Pod あたりのデータセンターのネットワーク帯域幅 6.4 Tbps

構成

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 形式の文字列です。ここで、AB は各方向のチップ数です。

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 チップが同じ Non Uniform Memory Access(NUMA)ノードを共有します。

8 チップ v5e TPU VM の場合、NUMA パーティション内では CPU-TPU 通信の方が効率的です。たとえば、次の図では、CPU0-Chip0 通信は CPU0-Chip4 通信よりも高速です。

NUMA ノードの通信

サービス提供用の Cloud TPU v5e タイプ

単一ホストのサービス提供は、最大 8 個の v5e チップでサポートされています。サポートされている構成は、1x1 スライス、2x2 スライス、2x4 スライスです。各スライスには、それぞれ 1、4、8 個のチップが搭載されています。

サービス提供をサポートする TPU v5e 構成: 1x1、2x2、2x4。

サービス提供ジョブに 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 がプリエンプトされる可能性があります。