TPU 内部架构解析:从芯片设计到超大规模集群
微信视频号:sph0RgSyDYV47z6
快手号:4874645212
抖音号:dy0so323fq2w
小红书号:95619019828
B站1:UID:3546863642871878
B站2:UID: 3546955410049087
本文深入解析了谷歌张量处理单元(TPU)的硬件架构与设计哲学。内容涵盖从单个TPU芯片的组成元件(包括矩阵乘法单元、向量处理单元及内存层次结构)到多芯片互联的拓扑结构(如2D/3D环面连接与光路交换技术)。文章详细探讨了TPU如何通过脉动阵列实现极高矩阵运算吞吐量,借助提前编译优化能效,并利用可重构互联拓扑支持从数芯片到数千芯片的灵活扩展。同时,分析了不同并行策略(数据并行、流水并行、张量并行)在各类TPU切片拓扑上的性能影响,并辅以实际部署案例(如PaLM训练)说明其大规模应用。
-
1. 背景介绍
-
2. TPU单芯片架构
-
3. TPU设计哲学
-
• 脉动阵列与流水线
-
• 提前编译与缓存策略
-
4. TPU多芯片架构
-
• 板级互联(4芯片)
-
• 机架级互联(64芯片)
-
• 全Pod级互联(4096芯片)
-
• 多Pod级互联
-
5. 实际系统对应关系
-
6. 参考文献
1 背景介绍
TPU(Tensor Processing Unit)是谷歌专为机器学习工作负载设计的专用集成电路(ASIC)。其核心设计目标聚焦于两个关键因素:极高的矩阵乘法吞吐量和优异的能效表现。
TPU的起源可追溯至2006年,当时谷歌正在评估使用GPU、FPGA还是定制ASIC来满足计算需求。早期仅少数应用需要专用硬件,可通过数据中心过剩的CPU资源满足。这一情况在2013年发生转变,谷歌语音搜索功能开始基于神经网络,内部预测表明若该功能普及,计算需求将大幅增长。
现今,TPU已支撑谷歌大多数AI服务,包括Gemini和Veo模型的训练与推理,以及推荐模型(如DLRM)的部署。
2 TPU单芯片架构
以下分析以TPUv4为主,但其架构布局也适用于最新代次TPU(如2025年6月尚未发布详细信息的TPUv7 "Ironwood")。
单个TPUv4芯片布局如下:
TPU单芯片与TensorCore
每个芯片包含两个TPU TensorCore,负责计算任务(注:专用于推理的TPU仅有一个TensorCore)。两个TensorCore共享内存单元:CMEM(128MiB)和HBM(32GiB)。
每个TensorCore内部包含计算单元及小型内存缓冲区:
-
1. 矩阵乘法单元(MXU)
-
• 这是TensorCore的核心组件,是一个128x128的脉动阵列。
-
• 脉动阵列的细节将在后文详述。
-
2. 向量单元(VPU)
-
• 处理通用逐元素操作(如ReLU、逐点加/乘、归约运算)。
-
3. 向量内存(VMEM; 32MiB)
-
• 内存缓冲区。数据从HBM复制到VMEM后,TensorCore才能执行计算。
-
4. 标量单元与标量内存(SMEM; 10MiB)
-
• 控制VPU和MXU的操作。
-
• 管理控制流、标量运算及内存地址生成。
对于熟悉NVIDIA GPU的读者,以下观察可能有所不同:
-
1. TPU上的片上内存单元(CMEM、VMEM、SMEM)远大于GPU的L1、L2缓存(H100的L1和L2缓存分别为256KB和50MB)。
-
2. TPU的HBM容量远小于GPU的HBM(H100为80GB)。
-
3. 负责计算的核心数量明显更少。
这与GPU设计形成对比:GPU具有较小的L1、L2缓存、较大的HBM以及数万个核心。
需要指出的是,TPU同样能够实现极高的吞吐量。TPU v5p每芯片可达到500 TFLOPs/s,一个包含8960芯片的完整Pod可提供约4.45 ExaFLOPs/s。最新的"Ironwood" TPUv7据称每个Pod(9216芯片)可达到42.5 ExaFLOPS/s。
理解TPU如何实现这一性能,需要深入其设计哲学。
3 TPU设计哲学
TPU通过两大支柱及一个关键假设实现卓越吞吐量和能效:脉动阵列与流水线技术、提前(AoT)编译,以及假设大多数操作可以映射到脉动阵列。在现代深度学习领域,矩阵乘法构成了计算的主体,恰好适合脉动阵列处理。
3.1 脉动阵列与流水线
脉动阵列是一种硬件设计架构,由互联的处理元件(PE)网格组成。每个PE执行小型计算(如乘加运算)并将结果传递给相邻PE。
脉动阵列示意图
这种设计的好处在于,一旦数据进入脉动阵列,无需额外的控制逻辑来指导数据处理。并且,当脉动阵列足够大时,除了输入和输出外,没有其他内存读写操作。
由于其 rigid 的组织结构,脉动阵列只能处理具有固定数据流模式的操作,但矩阵乘法和卷积完美契合这种模式。
此外,流水线技术可以清晰地重叠计算与数据移动。下图展示了TPU上流水线逐点操作的示意图。
流水线逐点操作(源自"How to Scale Your Model" [4]
脉动阵列的局限性:稀疏性 脉动阵列擅长处理稠密矩阵(即每个PE几乎每个周期都处于活跃状态)。然而,其缺点在于对相同大小的稀疏矩阵无法带来性能提升:即使对于零值元素,PE仍需工作,执行相同数量的周期。
应对脉动阵列的这一系统性限制将变得愈发重要,特别是在深度学习社区更青睐不规则稀疏性(如混合专家模型MoE)的背景下。
3.2 提前编译与减少对缓存的依赖
本节阐述TPU如何通过硬件-软件协同设计(TPU + XLA编译器)避免使用缓存,从而实现高能效。
传统缓存设计用于处理不可预测的内存访问模式。不同应用程序的程序可能具有截然不同的内存访问模式。本质上,缓存使硬件能够灵活适应广泛的应用范围。这是GPU(相对于TPU而言)非常灵活的一个重要原因。
然而,缓存访问(以及一般的内存访问)消耗大量能量。下图显示了芯片上操作的能量成本粗略估计(45nm, 0.9V;[18])。关键要点是内存访问和控制消耗了大部分能量,而算术运算消耗的能量显著较少。
不同操作的能量成本比较
但如果应用程序非常特定,其计算/内存访问模式高度可预测呢?
作为一个极端例子,如果编译器可以提前确定所有需要的内存访问,那么硬件可以仅依靠暂存器内存作为缓冲区,而无需缓存。
这正是TPU设计哲学所追求的,也是TPU与XLA编译器协同设计以实现这一目标的原因。XLA编译器通过提前分析计算图来生成优化程序。
JAX与TPU的协作 JAX+XLA在TPU上处于JIT(即时编译)和AOT(提前编译)的混合空间,这有时会造成困惑。当首次调用JAX中的jitted函数时,JAX会对其进行跟踪以创建静态计算图。该图被传递给XLA编译器,转换为TPU专用的完全静态二进制文件。正是在最后的转换阶段,进行了TPU特定的优化(例如最小化内存访问)以定制流程。
但需要注意:如果使用不同输入形状运行,jitted函数需要重新编译和缓存。这就是JAX在处理任何动态填充或长度依赖于输入的不同长度的循环层时表现不佳的原因。
当然,这种方法听起来很好,但也存在不便之处,即缺乏灵活性,并且严重依赖编译器是一把双刃剑。
TPU与能效(TPUv4) 之前的能量使用图并不能准确代表TPU,以下是TPUv4的能量使用细分。请注意,TPUv4采用7nm工艺,图中包含45nm数据用于比较([3], [16])。
TPUv4能量细分
每次操作能量(TPUv4, 7 nm)
TPUv4能量细分
左侧条形图可视化地展示了数值,但需要注意现代芯片使用HBM3,其能耗远低于此处的DDR3/4 DRAM。尽管如此,该图仍表明内存操作消耗的能量高出几个数量级。
这与现代扩展定律很好地联系起来:我们非常乐意增加FLOPS以换取减少内存操作。因此,减少内存操作具有双重优化好处,因为它们不仅使程序快速,而且消耗的能量也少得多。
4 TPU多芯片架构
4.1 板级(又称“板卡”;4芯片)
单个TPU板包含4个TPU芯片或8个TensorCore(简称“核心”)。每个板有自己的CPU主机(注:对于推理TPU,一个主机访问2个板,因为每个芯片只有1个核心)。
主机⇔芯片连接通过PCIe,而芯片⇔芯片连接通过核心间互联(ICI),后者具有更高带宽。
但ICI连接可以进一步扩展到多个板卡。为此,需要上升到机架级别。
添加图片注释,不超过 140 字(可选)
4.2 机架级(4x4x4芯片)
TPU尤其令人兴奋的部分在于其可扩展性,这从机架级别开始显现。
一个TPU机架由64个TPU组成,以4x4x4 3D环面方式连接。如果您见过谷歌如下的TPU宣传材料,这张图片展示的是8个TPU机架。
8个TPU机架(TPUv4)
在深入探讨机架之前,需要澄清一些容易混淆的术语:机架(Rack)、Pod和切片(Slice)。
机架、Pod与切片的区别 不同的谷歌资料来源对它们的使用略有不同,有时“TPU Pod”和“TPU Slice”可以互换使用。但在本文中,我们将遵循谷歌TPU论文和GCP TPU文档中使用的定义([3], [7], [9])。
-
1. TPU机架 (TPU Rack):
-
• 包含64个芯片的物理单元。也称为“立方体”(cube)。
-
2. TPU Pod:
-
• 通过ICI和光纤连接的最大TPU单元。
-
• 也称为“超级Pod”(Superpod)或“完整Pod”。例如,TPUv4的一个TPU Pod由4096个芯片或64个TPU机架组成。
-
3. TPU切片 (TPU Slice):
-
• 介于4个芯片和超级Pod大小之间的任何TPU配置。
关键区别在于TPU机架和TPU Pod是物理度量单位,而TPU切片是一个抽象单位。当然,设置TPU切片有重要的物理属性,但我们现在暂时抽象掉这一点。
目前,我们将使用物理度量单位:TPU机架和TPU Pod。了解TPU系统如何物理连接在一起,可以让我们更好地理解TPU的设计理念。
回到TPU机架(TPUv4): 一个单独的TPU机架由64个通过ICI和光路交换(OCS)连接在一起的芯片组成。本质上,我们将多个板卡连接在一起以模拟一个64芯片的系统。这种将小部件组合成超级计算机的主题在后面会继续出现。
下图是一个单独的TPUv4机架示意图。它是一个4x4x4 3D环面,每个节点是一个芯片,蓝色箭头是ICI,面上的线是OCS(根据[7]重绘)。
带有OCS的单个TPU机架
然而,此图引出了几个问题。为什么OCS只用于面? 换句话说,使用OCS的好处是什么? 有三大好处,我们将在下面介绍另外两个。
OCS的好处 #1:环绕连接 通过环绕实现节点间更快的通信。 OCS也充当给定TPU配置的环绕连接。这将两个节点之间最坏情况的跳数从每轴N-1跳减少到(N-1)/2跳,因为每个轴变成一个环(1D环面)。 随着规模扩大,这种效应变得更加重要,因为减少芯片间通信延迟对于高并行化至关重要。
并非所有TPU都具有3D环面拓扑 注意:较旧的TPU代次(如TPUv2, v3)和推理TPU(如TPUv5e, TPUv6e)具有2D环面拓扑,而不是如下所示的3D环面。然而,TPUv7 "Ironwood" 似乎是3D环面,尽管它被宣传为推理芯片(注:此判断仅基于其宣传材料)。
2D环面拓扑
4.3 全Pod级(又称“超级Pod”;TPUv4为4096芯片)
就像我们连接多个芯片组成一个TPU机架一样,我们可以连接多个机架组成一个大型超级Pod。 超级Pod也指的是TPU可以达到的互连芯片(仅使用ICI和OCS)的最大配置。接下来还有多Pod级别,但这必须通过较慢的互连实现,我们将在之后介绍。 这因代次而异,但对于TPUv4是4096个芯片(即64个机架,每个4x4x4芯片)。对于最新的TPUv7 "Ironwood",这是9216个芯片。 下图显示了TPUv4的一个超级Pod。
TPUv4的超级Pod(64个机架)
注意每个立方体(即一个TPU机架)如何通过OCS相互连接。这也允许我们在一个Pod中获取TPU切片。
使用OCS的TPU切片 我们可以请求Pod内的TPU子集,这些就是TPU切片。但即使你想要N个芯片,也有多种拓扑可以选择。 例如,假设你需要总共512个芯片。你可以请求一个立方体(8x8x8)、一个雪茄形状(4x4x32)或一个矩形(4x8x16)。选择切片的拓扑本身就是一个超参数。 你选择的拓扑将影响节点间的通信带宽。这直接影响了不同并行方法的性能。 例如,立方体(例如8x8x8)更适合全通信(all-to-all communications),例如数据并行或张量并行,因为它具有最高的二分带宽。然而,雪茄形状(例如4x4x32)更适合流水线并行,因为它可以更快地与顺序层通信(假设一层适合一个4x4芯片的子切片)。
示例TPU拓扑
当然,最佳拓扑取决于模型,找到它本身也是一项工作。TPUv4论文[9]也测量了这一点,以显示拓扑变化如何加速吞吐量(注:我不确定第一行指的是哪种LLM架构,因为它未指定)。
不同拓扑下的吞吐量改进
我们介绍了TPU切片,但有一个重要特性有助于TPU的高运行稳定性。 由于OCS的存在,这些切片不必是连续的机架。这是使用OCS的第二个好处——也可能是最大的——我们之前没有提到的。
OCS的好处 #2:(可重构的)非连续多节点切片 请注意,这与硬连接多个节点以模拟非连续切片不同。由于OCS是一个交换机而不是硬连接的,节点之间的物理线路要少得多,从而允许更高的可扩展性(即更大的TPU Pod大小)。 这允许大规模灵活的节点配置。例如,假设我们想在一个Pod上运行三个作业。虽然朴素的调度不允许这样做,但OCS连接允许我们抽象掉节点的位置,并将整个Pod视为仅仅一个“节点袋”(根据[6]重绘)。
单个作业可以将Pod中的TPU机架视为“节点袋”
这提高了Pod利用率,并且在节点发生故障时可能更易于维护。谷歌将其描述为“死节点爆炸半径小”。但是,我不确定当只有某些节点必须关闭时,其液体冷却会如何受到影响。 最后,这种灵活OCS有一个有趣的扩展:我们也可以改变TPU切片的拓扑,例如从常规环面到扭曲环面。
OCS的好处 #3:扭曲的TPU拓扑 我们之前看到如何通过改变固定数量芯片的(x,y,z)维度来实现不同的TPU切片拓扑。这次,然而,我们将处理固定的(x,y,z)维度,但改为改变它们的连接方式以实现不同的拓扑。 一个显著的例子是从雪茄形状的常规环面变为扭曲雪茄环面,如下所示。
常规环面与扭曲环面(来源:TPUv4论文 [9])
扭曲环面情况允许跨扭曲2D平面的芯片之间更快地通信。这对于加速全通信特别有用。 让我们更深入一点,想象一个这将有所帮助的具体场景。
使用扭曲环面加速训练 理论上,扭曲环面将为张量并行(TP)带来最大好处,因为每层有多个全收集(all-gather)和规约散点(reduce-scatter)操作。它也可能为数据并行(DP)带来中等好处,因为每个训练步骤也有一个全归约(all-reduce),但这将不那么频繁。 假设我们正在训练一个标准的仅解码器Transformer,并且我们希望采用大量并行来加速训练。我们将在下面看到两种场景。
场景 #1: 4x4x16 拓扑 (TP + PP; 总共 256 芯片) 我们的 z 轴将是我们的流水线并行 (PP) 维度,我们的 2D TP 维度将是 4x4。本质上,假设每层 k 位于 z=k,并且每层被分片到 16 个芯片上。假设没有明确绘制的情况下使用标准 OCS 连接(即最近邻)。
具有 TP + PP 的 4x4x16 拓扑
我们将在每个 z=k 处扭曲 2D 环面,这使得每个 TP 层内的芯片通信更快。沿着我们的 PP 维度扭曲是不必要的,因为它们主要依赖点对点通信。 注意: 实际上,当芯片数量大于 4x4 时,扭曲环面会带来好处。我们使用 4x4 仅用于可视化目的。
具有 DP + TP + PP 的 16x4x16 拓扑
注意扭曲环面仅限于每个 DP 模型的每个 TP 维度(即,对于每个 z=k,给定 k=1…16,一个 4x4 2D 平面)。DP 维度只有一个环绕,因此每一行变成一个大小为 16 的水平环。 你可能已经注意到存在另一种 8x8x16 的拓扑(即 2x2 DP 维度),但这变得更加复杂,因为我们将 DP 和 TP 维度混合在一起。具体来说,不清楚我们应该如何为 y 轴构建 OCS 环绕,同时为每个 TP 维度容纳扭曲环面。
4.4 多Pod级(又称“多切片”;TPUv4为4096+芯片)
TPU层次结构的最终级别是多Pod级别。在这里,您可以将多个Pod视为一台大型机器。然而,Pod之间的通信通过数据中心网络(DCN)进行,其带宽低于ICI。
通过DCN连接的两个Pod [1]
该图显示了如何配置多Pod训练。 PaLM就是这样训练的。它花费了56天时间在6144个TPUv4(2个Pod)上进行训练。下面您可以看到跨越6个Pod的TPU作业分配:绿色用于PaLM,红色表示未分配,其余是其他作业。注意每个方块是一个4x4x4的TPU立方体。
训练PaLM的TPU Pod利用率 [6]
实现这一点本身就很困难,但更令人印象深刻的是对开发者体验的关注。具体来说,重点是关注“我们如何能尽可能抽象化模型扩展的系统/HW部分?”这个问题。 谷歌的答案是让XLA编译器负责在大规模芯片之间协调通信。通过研究人员提供正确的标志(即DP、FSDP、TP的并行维度、切片数量等),XLA编译器为手头的TPU拓扑插入正确的分层集合操作(Xu et al, 2021: GSPMD [2])。目标是使大规模训练在尽可能少的代码更改下进行。 例如,以下是谷歌博客[1]中跨多个切片的全归约操作分解。
跨Pod的XLA全归约
这表明XLA编译器负责处理切片之间和切片内的通信集合操作。 举一个具体例子,训练一个模型的TPU拓扑可能如下所示。激活通信通过ICI在切片内进行,而梯度通信将通过DCN跨切片进行(即跨DCN DP维度)[1]。
多Pod通信示意图
5 实际系统对应关系
我发现当有硬件的实际照片时,将图表放入背景中会很有帮助。下面是一个总结。 如果你见过谷歌TPU宣传材料的图片,你可能遇到过下面这张图片。
8个TPU机架(TPUv4)
这是8个TPU Pod,其中每个单元是我们上面看到的4x4x4 3D环面。一个Pod中的每行有2个板卡,意味着每行有8个TPU芯片。 这是一个单独的TPUv4板卡:
TPUv4板卡
请注意,该图简化为只有一个PCIe端口,但在实际板卡上,有4个PCIe端口(在左侧)——每个TPU一个。 下面是一个单独的芯片:
TPUv4芯片,中心为ASIC + 4个HBM堆栈
中心部分是ASIC,周围的4个块是HBM堆栈。这是我们看到的TPU v4,所以内部有2个TensorCores,因此总共有4个HBM堆栈。
我找不到TPUv4的芯片布局图,所以这里有一个TPUv4i的布局图,它类似,只是因为它是一个推理芯片,所以只有1个TensorCore [3]。
TPUv4i芯片布局
注意CMEM在TPUv4i的布局上占据了相当多的空间。
6 致谢
感谢Google TPU Research Cloud (TRC) 提供的TPU支持!
微信视频号:sph0RgSyDYV47z6
快手号:4874645212
抖音号:dy0so323fq2w
小红书号:95619019828
B站1:UID:3546863642871878
B站2:UID: 3546955410049087
参考文献链接
7 参考文献
[1] Google Blog: TPU Multi-Slice Training[1] [2] Xu, et al. "GSPMD: General and Scalable Parallelizaton for ML Computation Graphs"[2] [3] Jouppi et al. "Ten Lessons From Three Generations Shaped Google's TPUv4i"[3] [4] How to Scale Your Model - TPUs[4] [5] Domain Specific Architectures for AI Inference - TPUs[5] [6] HotChips 2023: TPUv4[6] [7] Google Cloud Docs: TPUv4[7] [8] Jouppi et al. "In-Datacenter Performance Analysis of a Tensor Processing Unit" -- TPU origins paper[8] [9] Jouppi et al. "TPU v4"-- TPUv4 paper[9] [10] PaLM training video[10] [11] HotChips 2021: "Challenges in large scale training of Giant Transformers on Google TPU machines"[11] [12] HotChips 2020: "Exploring Limits of ML Training on Google TPUs"[12] [13] Google Blog: Ironwood[13] [14] HotChips 2019: "Cloud TPU: Codesigning Architecture and Infrastructure"[14] [15] ETH Zurich's Comp Arch Lecture 28: Systolic Array Architectures[15] [16] Patterson presentation: "A Decade of Machine Learning Accelerators: Lessons Learned and Carbon Footprint"[16] [17] Camara et al. "Twisted Torus Topologies for Enhanced Interconnection Networks."[17] [18] Horowitz article: "Computing's Energy Problem(and what we can do about it)"[18]
本文翻译自TPU Deep Dive[19]
引用链接
[1] Google Blog: TPU Multi-Slice Training: https://cloud.google.com/blog/products/compute/using-cloud-tpu-multislice-to-scale-ai-workloads [2] Xu, et al. "GSPMD: General and Scalable Parallelizaton for ML Computation Graphs": https://arxiv.org/pdf/2105.04663 [3] Jouppi et al. "Ten Lessons From Three Generations Shaped Google's TPUv4i": https://gwern.net/doc/ai/scaling/hardware/2021-jouppi.pdf [4] How to Scale Your Model - TPUs: https://jax-ml.github.io/scaling-book/tpus/ [5] Domain Specific Architectures for AI Inference - TPUs: https://fleetwood.dev/posts/domain-specific-architectures#google-tpu [6] HotChips 2023: TPUv4: https://hc2023.hotchips.org/assets/program/conference/day2/ML%20training/HC2023.Session5.ML_Training.Google.Norm_Jouppi.Andy_Swing.Final_2023-08-25.pdf [7] Google Cloud Docs: TPUv4: https://cloud.google.com/tpu/docs/v4 [8] Jouppi et al. "In-Datacenter Performance Analysis of a Tensor Processing Unit" -- TPU origins paper: https://arxiv.org/abs/1704.04760 [9] Jouppi et al. "TPU v4"-- TPUv4 paper: https://arxiv.org/abs/2304.01433 [10] PaLM training video: https://www.youtube.com/watch?v=0yPFBxkOKRY [11] HotChips 2021: "Challenges in large scale training of Giant Transformers on Google TPU machines": https://hc33.hotchips.org/assets/program/tutorials/HC2021.Google.Sameer%20Kumar.pdf [12] HotChips 2020: "Exploring Limits of ML Training on Google TPUs": https://hc32.hotchips.org/assets/program/tutorials/HC2020.Google.SameerKumarDehaoChen.v02.pdf [13] Google Blog: Ironwood: https://blog.google/products/google-cloud/ironwood-tpu-age-of-inference/ [14] HotChips 2019: "Cloud TPU: Codesigning Architecture and Infrastructure": https://old.hotchips.org/hc31/HC31_T3_Cloud_TPU_Codesign.pdf [15] ETH Zurich's Comp Arch Lecture 28: Systolic Array Architectures: https://www.youtube.com/watch?v=XkgtANeDrm8 [16] Patterson presentation: "A Decade of Machine Learning Accelerators: Lessons Learned and Carbon Footprint": https://www.cs.ucla.edu/wp-content/uploads/cs/PATTERSON-10-Lessons-4-TPU-gens-CO2e-45-minutes.pdf [17] Camara et al. "Twisted Torus Topologies for Enhanced Interconnection Networks.": https://personales.unican.es/vallejoe/Publications/C%C3%A1mara%20-%20TPDS'10%20-%20Twisted%20Torus%20Topologies%20for%20Enhanced%20Interconnection%20Networks.pdf [18] Horowitz article: "Computing's Energy Problem(and what we can do about it)": https://gwern.net/doc/cs/hardware/2014-horowitz-2.pdf [19] TPU Deep Dive: https://henryhmko.github.io/posts/tpu/tpu.html
人工智能芯片与自动驾驶

浙公网安备 33010602011771号