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硬件加速。这些技术共同构成了现代分布式训练的基础。