TPUとGPUクラスタ：集団通信の解剖

2026年、トランスフォーマーモデルのトレーニングと推論は大規模分散システムの問題となっています。クラスタ全体にモデルをシャーディングするために、データ並列、テンソル/モデル並列、FSDP、エキスパート並列などの手法に依存しています。これらの手法の基盤には、全収集（All-Reduce）、全ギャザー（All-Gather）、リデューススキャッター（Reduce-Scatter）、全対全（All-to-All）といった少数のコア集団操作があります。これらの集団通信がどのように機能するかを理解することは、現代のトランスフォーマーシステムのパフォーマンスを推論する上で重要です。

TPUクラスタのトポロジはGPUクラスタよりも均一です。TPUチップはチップ間相互接続（ICI）を介して隣接チップに直接接続され、各チップは4つまたは6つの最近傍を持ちます（それぞれ2Dトーラス（例：TPU v5e）および3Dトーラス（例：TPU v5p）に対応）。例えば、v5eの16×16スーパーポッドは完全な2Dトーラスを形成し、ラップアラウンド接続があります。より小さなスライスをリクエストすると、サイズが16未満の場合、ラップアラウンド接続が失われ、メッシュに退化します。ポッド間通信はデータセンターネットワーキング（DCN）を介して行われますが、DCNの帯域幅はICIよりもはるかに低いため、ポッド間通信がボトルネックになりがちです。

帯域幅階層は、チップ内のHBMからICI、PCIe、そしてDCNへと低下します。v5eを例にとると、ICIの一方向帯域幅は45 GB/s、PCIeは32 GB/s、DCNは12 GB/sです。この階層がデータ移動のコストを決定します：計算コアに近いほど高速です。この階層を理解することは、集団操作の最適化に不可欠です。例えば、大メッセージ転送ではリングアルゴリズムが帯域幅を最大限活用できますが、メッセージが小さい場合、1ホップあたり約1マイクロ秒のレイテンシが支配的になり、ツリーアルゴリズム（log₂ステップのみ必要）の方が適切になります。

All-Gatherを例にとると、TPUトーラス上ではデータがリングの軸に沿って効率的に伝搬します。4×4メッシュの場合、8 MiBの行列を移動するのに2つの並列パス（各6ホップ）を使用し、総時間は約8マイクロ秒（レイテンシ無視）です。レイテンシを考慮すると、メッセージの粒度が約45 KB（1マイクロ秒で転送可能なデータ量）の場合、レイテンシ要因は無視できなくなります。したがって、通信が帯域幅制限かレイテンシ制限かを理解することが、性能チューニングに重要です。

さらに、All-ReduceはReduce-ScatterとAll-Gatherに分解でき、All-to-AllはMoEモデルのエキスパート並列に使用されます。GPUクラスタではNVLinkとInfiniBandが使用され、トポロジは通常Fat-Treeであり、ノード内の集団操作はSHARPハードウェアで高速化できます。これらの技術が現代の分散トレーニングの基盤を形成しています。