TPU與GPU集羣：集體通信的解剖

在2026年，訓練和推理Transformer模型已成為大規模分佈式系統問題。為了在集羣中切分模型，我們依賴數據並行、張量/模型並行、FSDP和專家並行等技術，這些技術底層都建立在少量核心集體操作之上，例如全規約（All-Reduce）、全收集（All-Gather）、規約散播（Reduce-Scatter）和全到全（All-to-All）。理解這些集體通信如何工作，對於推理現代Transformer系統的性能至關重要。

TPU集羣的拓撲結構比GPU集羣更為統一。TPU芯片通過芯片間互連（ICI）直接連接到相鄰芯片，每個芯片有4個或6個最近鄰，分別對應2D環面（如TPU v5e）和3D環面（如TPU v5p）。例如，一個v5e的16×16超級片（Pod）構成一個完整的2D環面，具有環繞連接。當請求更小的片（Slice）時，若尺寸小於16，環繞連接消失，退化為網格。跨Pod通信通過數據中心網絡（DCN）實現，但DCN帶寬遠低於ICI，因此跨Pod通信容易成為瓶頸。

帶寬層次結構從芯片內部HBM到ICI、PCIe、再到DCN逐漸降低。以v5e為例，ICI單向帶寬為45 GB/s，PCIe為32 GB/s，DCN為12 GB/s。這一層次決定了數據移動的成本：越靠近計算核心，速度越快。理解這一層次對於優化集體操作至關重要。例如，在大消息傳輸中，環狀算法可以充分利用帶寬，但若消息較小，每跳約1微秒的延遲會佔據主導，此時樹狀算法因只需log₂步而更優。

以All-Gather為例，在TPU環面上，數據可以沿環的軸向高效傳播。以4×4網格為例，移動一個8 MiB矩陣通過兩條平行路徑，每條路徑6跳，總時間約8微秒（忽略延遲）。當考慮延遲時，若消息粒度約為45 KB（1微秒可傳輸的數據量），延遲因素不可忽略。因此，理解通信是帶寬受限還是延遲受限，對於性能調優至關重要。

此外，All-Reduce可分解為Reduce-Scatter和All-Gather，而All-to-All用於MoE模型的專家並行。GPU集羣則使用NVLink和InfiniBand，拓撲通常為Fat-Tree，節點內集體操作可通過SHARP硬件加速。這些技術共同構成了現代分佈式訓練的基礎。