TPU 硬核科普详解

微信视频号：sph0RgSyDYV47z6

快手号：4874645212

抖音号：dy0so323fq2w

小红书号：95619019828

本文是由 Henry Ko（github: henryhmko）发布于其个人技术播客的一篇英文播客，深入介绍和分析了谷歌 TPU 芯片的设计理念。「科技写作侠 Patrick」和 Claude 一起对这篇博客完成了中文编译，并且添加了字数接近原文的概念注释，帮助所有背景和年龄的读者都能够阅读和理解这篇文章，旨在让中文读者能够更深入了解 GPU 之外最重要的 AI 计算芯片。

以下为原文：

我最近在大量使用TPU，看到它们与GPU相比有着如此不同的设计理念，这很有趣。

注释：TPU (Tensor Processing Unit) 是Google开发的专门用于AI计算的芯片，类似于GPU，但专门为AI模型的训练和推理优化。GPU原本是为图形处理设计的，后来被广泛用于AI计算。

TPU的主要优势在于其可扩展性。这是通过硬件端（如能效和模块化）和软件端（如XLA编译器）的协同设计实现的。

注释：可扩展性指的是能够通过增加更多芯片来提升整体计算能力。硬件端指芯片本身的设计，软件端指配套的程序和工具。XLA编译器是Google开发的一种能够优化AI计算的程序。

01 背景

简单介绍一下TPU，它是Google的ASIC芯片，专注于两个因素：极高的矩阵乘法（matmul）吞吐量和能效。

注释：ASIC (Application-Specific Integrated Circuit) 是专用集成电路，即为特定用途专门设计的芯片。矩阵乘法是AI计算中最核心的数学运算，就像加法对于计算器一样重要。吞吐量指单位时间内能处理的计算量，能效指每消耗一单位电能能完成多少计算。

TPU的起源可以追溯到2006年的Google，当时他们首次评估是否应该采用GPU、FPGA或定制ASIC。那时只有少数应用需要专用硬件，他们认为这些需求可以通过引入大型数据中心的多余CPU算力来满足。但这在2013年发生了变化，当时Google的语音搜索功能开始运行在神经网络上，内部预测显示如果该功能普及，他们将需要更多的算力。

注释：FPGA (Field-Programmable Gate Array) 是可编程门阵列，是一种可以重新配置的芯片。CPU是通用处理器，适合各种计算但在AI计算上效率较低。神经网络是模拟人脑工作方式的AI算法。算力指计算能力，就像马力对于汽车一样。

时至今日，TPU为Google的大部分AI服务提供动力。当然，这包括Gemini或Veo的训练和推理，也包括部署他们的推荐模型（DLRM）。

注释：Gemini是Google的AI聊天机器人，Veo是Google的AI视频生成工具。训练指教会AI模型如何工作，推理指使用已训练好的模型来处理实际任务。推荐模型是用来推荐内容的AI系统，比如YouTube推荐视频。

让我们深入了解TPU的内部结构，从底层开始。

02 TPU单芯片级别

我将重点关注TPUv4的架构图，但这个布局或多或少适用于最新一代的TPU（如TPUv6p "Trillium"；TPUv7 "Ironwood"的详细信息在2025年6月撰写时尚未公布）。

注释：TPUv4、TPUv6p、TPUv7是TPU的不同版本，就像iPhone 14、iPhone 15一样，数字越大通常性能越好。架构图是芯片内部结构的示意图，就像建筑物的平面图。

以下是单个TPUv4芯片的布局：

 

 

添加图片注释，不超过 140 字（可选）

TPU单芯片 + TensorCore

在每个芯片中，有两个TPU TensorCore，负责我们的计算。（注：推理专用TPU只有一个TensorCore）。两个TensorCore共享内存单元：CMEM（128MiB）和HBM（32GiB）。

注释：TensorCore是芯片内部负责AI计算的核心部件，类似于大脑中的神经元。推理专用指只用来运行已训练好的AI模型，不用来训练新模型。内存是临时存储数据的地方，MiB和GiB是内存容量单位，1GiB约等于1000MiB。

在每个TensorCore内部，有我们的计算单元和较小的内存缓冲区：

1. 矩阵乘法单元（MXU）

• 这是TensorCore的关键组件，是一个128x128的脉动阵列（systolic array）
• 我们将在下面详细介绍脉动阵列

注释：矩阵乘法单元专门用来进行矩阵乘法运算。128x128意味着它可以同时处理128行×128列的数据，这是一个相当大的计算阵列。

2. 向量单元（VPU）

• 通用逐元素操作（如ReLU、点乘加/乘、归约）

注释：向量单元处理一些简单的数学运算。ReLU是一种常用的AI算法中的数学函数，点乘是基本数学运算，归约是将多个数值合并成一个的操作。

3. 向量内存（VMEM；32MiB）

• 内存缓冲区。HBM中的数据在TensorCore进行任何计算之前先复制到VMEM中

注释：这就像是工作台上的临时储物盒，数据从大仓库（HBM）搬到工作台附近的小盒子（VMEM）里，方便快速取用。

4. 标量单元 + 标量内存（SMEM；10MiB）

• 告诉VPU和MXU要做什么
• 管理控制流、标量操作和内存地址生成

注释：标量单元就像是指挥官，告诉其他部件该做什么。控制流是程序执行的顺序，标量是单个数值（相对于向量的多个数值），内存地址是数据在内存中的位置标识。

如果你来自NVIDIA GPU的背景，有一些初始观察可能会让你感到困惑：

1. TPU上的片上内存单元（CMEM、VMEM、SMEM）比GPU上的L1、L2缓存大得多
2. TPU上的HBM也比GPU上的HBM小得多
3. 负责计算的"核心"似乎少了很多

注释：NVIDIA是制造GPU的知名公司。L1、L2缓存是GPU内部的小型高速内存。这里的对比说明TPU和GPU有不同的设计策略：TPU使用较少但更大的内存单元，而GPU使用很多小的处理核心。

这与GPU完全相反，GPU有较小的L1、L2缓存（H100分别为256KB和50MB），较大的HBM（H100为80GB），以及数万个核心。

注释：H100是NVIDIA的高端GPU型号。KB、MB、GB是存储容量单位，1GB=1000MB=1000000KB。数万个核心意味着GPU内部有成千上万个小的计算单元。

在我们深入之前，回想一下TPU就像GPU一样能够实现极高的吞吐量。TPU v5p每芯片可达到500 TFLOPs/sec，使用8960芯片的完整pod可实现约4.45 ExaFLOPs/sec。最新的"Ironwood"TPUv7据说每pod（9216芯片）可达到42.5 ExaFLOPS/sec。

注释：TFLOPs和ExaFLOPs是计算速度单位，类似于汽车的马力。1 ExaFLOPs = 1000000 TFLOPs，这些数字表示每秒能进行多少万亿次浮点运算。Pod是多个芯片组成的集群，就像多台电脑连接成一个超级计算机。

要理解TPU如何实现这一点，我们需要了解它们的设计理念。

03 TPU设计理念

TPU通过依靠两大支柱和一个关键假设来实现惊人的吞吐量和能效：脉动阵列+流水线、提前编译（AoT），以及假设大多数操作可以以很好映射到脉动阵列的方式表达。幸运的是，在我们现代深度学习时代，矩阵乘法占据了大部分计算，这很适合脉动阵列。

注释：设计理念是指导芯片设计的核心思想。脉动阵列是一种特殊的计算架构，流水线是一种提高效率的技术。提前编译意味着在程序运行前就把它翻译成芯片能理解的指令。深度学习是现代AI的主要技术。

TPU设计选择 #1：脉动阵列 + 流水线

问：什么是脉动阵列？

脉动阵列是一种硬件设计架构，由相互连接的处理元素（PE）网格组成。每个PE执行小型计算（如乘法累加）并将结果传递给相邻的PE。

注释：想象一个工厂的流水线，每个工人（PE）负责一个简单操作，然后把半成品传给下一个工人。脉动阵列就像这样的二维网格，数据像心跳一样有节奏地在网格中流动，因此叫"脉动"。

 

 

添加图片注释，不超过 140 字（可选）

这种设计的好处是，一旦数据传入脉动阵列，就不需要额外的控制逻辑来处理数据。此外，当给定足够大的脉动阵列时，除了输入和输出之外没有内存读写。

注释：控制逻辑是告诉芯片如何处理数据的指令系统。脉动阵列的优势是一旦启动就能自动运行，不需要复杂的指挥，就像训练有素的军队方阵一样。

由于其严格的组织结构，脉动阵列只能处理具有固定数据流模式的操作，但幸运的是矩阵乘法和卷积完全适合这种模式。

注释：固定数据流模式意味着数据按照预定的、规律的方式流动。卷积是AI中常用的另一种数学运算，特别用于图像处理。这就像某些机器只能生产特定类型的产品，但恰好这些产品是最需要的。

此外，流水线有明显的机会可以将计算与数据移动重叠。下面是TPU上流水线点操作的图表。

注释：流水线技术让不同的操作同时进行，就像汽车生产线上，当第二辆车开始组装时，第一辆车同时在进行油漆工作。这大大提高了效率。

 

 

添加图片注释，不超过 140 字（可选）

附：脉动阵列的缺点 - 稀疏性

你可以看到脉动阵列喜欢密集矩阵（即当每个PE几乎每个周期都处于活动状态时）。然而，缺点是对于相同大小的稀疏矩阵没有性能改进：仍需要执行相同数量的周期，即使对于零值元素，PE仍在工作

注释：密集矩阵是大部分位置都有有效数据的矩阵，稀疏矩阵是大部分位置都是零的矩阵。这就像让所有工人都同时工作效率最高，但如果有些工人没有实际工作要做（处理零值），他们还是要空转，浪费了资源。

如果深度学习社区倾向于更多不规则稀疏性（如MoE），处理脉动阵列的这种系统性限制将变得更重要。

注释：MoE (Mixture of Experts) 是一种AI模型架构，只激活部分神经网络，产生稀疏的计算模式。不规则稀疏性是指零值分布不规律，这对脉动阵列来说更难处理。

TPU设计选择 #2：提前编译（AoT）+ 减少对缓存的依赖

本节解释TPU如何通过TPU + XLA编译器的硬件-软件协同设计来避免缓存，从而实现高能效。

注释：缓存是芯片内部的高速小容量存储器，用来临时保存常用数据。硬件-软件协同设计是指硬件和软件一起设计，相互配合以达到最佳性能。

首先，回想一下传统缓存是为了处理不可预测的内存访问模式而设计的。一个应用程序的程序可能与其他应用程序的程序有着截然不同的内存访问模式。本质上，缓存允许硬件灵活并适应广泛的应用程序。这是GPU非常灵活的硬件（注：相比TPU）的一个重要原因。

注释：内存访问模式是指程序访问数据的顺序和方式。传统芯片使用缓存来应对各种不同的程序需求，就像图书馆把热门书籍放在容易拿到的地方，但不知道每个读者会要什么书。

然而，缓存访问（以及一般的内存访问）消耗大量能源。下面是芯片上操作的粗略能源成本估算（45nm，0.9V）。这里的关键要点是内存访问和控制占用了我们大部分能源，而算术运算消耗的能源要少得多。

注释：45nm和0.9V是芯片制造的技术参数，数字越小通常表示技术越先进。这个比较显示，移动数据比计算数据消耗更多能源，就像运输货物比加工货物更耗油。

 

 

添加图片注释，不超过 140 字（可选）

但如果你的应用程序非常具体，其计算/内存访问模式高度可预测呢？

作为一个极端例子，如果我们的编译器能够提前找出所有必需的内存访问，那么我们的硬件只需要一个作为缓冲区的暂存存储器，而不需要缓存。

注释：这就像如果你提前知道今天要做的所有事情，就可以一次性准备好所有需要的工具，而不需要一个复杂的储物系统来应对各种可能的需求。

这正是TPU理念的目标，也是TPU与XLA编译器协同设计以实现这一目标的确切原因。XLA编译器通过提前分析计算图来生成优化程序。

注释：计算图是AI模型计算过程的可视化表示，显示数据如何在不同操作间流动。提前分析就像提前规划旅行路线，知道每一步要去哪里。

问：但JAX也能很好地与TPU配合工作，但它们使用@jit？

JAX+XLA在TPU上处于JIT和AOT的混合空间，因此会产生混淆。当我们第一次在JAX中调用jitted函数时，JAX会跟踪它以创建静态计算图。这会传递给XLA编译器，在那里它被转换为TPU的完全静态二进制文件。正是在最后的转换阶段进行TPU特定的优化（如最小化内存访问）以适应TPU。

注释：JAX是Google开发的AI编程框架。JIT (Just-In-Time) 编译是运行时编译，AOT (Ahead-Of-Time) 是提前编译。@jit是JAX中的一个装饰器，用来标记需要优化的函数。静态意味着在运行前就确定，二进制文件是机器能直接执行的程序。

但有一个注意事项：如果jitted函数使用不同的输入形状运行，它们必须重新编译和缓存。这就是为什么JAX在任何动态填充或具有依赖于输入的不同长度for循环的层上表现不佳。

注释：输入形状是指数据的尺寸大小。动态填充是指根据实际数据调整尺寸。这就像为不同尺寸的衣服需要不同的生产线设置一样，频繁改变会降低效率。

当然，这种方法听起来很好，但也有不便的缺点。缺乏灵活性，这种对编译器的严重依赖是一把双刃剑。

但为什么Google仍然追求这种设计理念？

TPU和能效（TPUv4）

之前的能源使用图表对TPU来说不是准确的表示，所以这里是TPUv4的能源使用分解。注意TPUv4是7nm的，45nm在这里是为了比较。

注释：7nm和45nm是芯片制造工艺，数字越小表示技术越先进，能效越高。这就像从粗糙的手工制作发展到精密的机器制造。

 

 

添加图片注释，不超过 140 字（可选）

 

 

添加图片注释，不超过 140 字（可选）

左侧的条形图向我们展示了数值的可视化，但需要注意的是现代芯片使用HBM3，比这里显示的DDR3/4 DRAM消耗的能源少得多。尽管如此，这表明内存操作消耗的能源要多几个数量级。

注释：HBM3、DDR3/4是不同类型的内存技术，HBM3更先进但更昂贵。数量级是指10倍、100倍、1000倍这样的差别。这说明移动数据比计算数据要耗能得多。

这与现代扩展定律有很好的联系：我们非常乐意增加FLOPS来换取减少的内存操作。因此减少内存操作有双重优化效益，因为它们不仅使程序快速，而且消耗的能源也少得多。

注释：扩展定律是指导AI发展的经验规律。FLOPS是每秒浮点运算次数，衡量计算能力。这里的思路是：宁愿多做计算也要少搬数据，因为搬数据更耗能耗时。

04 TPU多芯片级别

让我们向上一级，看看TPU在多芯片设置中如何工作。

注释：多芯片设置是指把多个芯片连接起来协同工作，就像把多台电脑连接成一个更强大的系统。

Tray级别（又称"Board"；4芯片）

单个TPU tray由4个TPU芯片或8个TensorCore（简称"cores"）组成。每个tray都有自己的CPU主机（注：对于推理TPU，一个主机访问2个tray，因为每个芯片只有1个core）。

注释：Tray就像一个托盘，装着几个芯片。CPU主机是控制芯片的普通处理器，就像乐队指挥指挥多个乐器演奏员一样。

 

 

添加图片注释，不超过 140 字（可选）

主机⇔芯片连接是PCIe，但芯片⇔芯片连接是Inter-Core Interconnect（ICI），具有更高的带宽。

注释：PCIe是一种常见的计算机内部连接标准，ICI是Google专门为TPU设计的连接方式。带宽是指单位时间内能传输的数据量，就像道路的宽度决定了车流量。

但ICI连接扩展到多个tray。为此，我们需要移至Rack级别。

Rack级别（4x4x4芯片）

TPU特别令人兴奋的部分在于其可扩展性，我们从Rack级别开始看到这一点。

注释：Rack是机架，就像服务器机房里的大柜子，可以装很多设备。4x4x4意味着在三维空间中排列成4×4×4=64个芯片。

TPU机架由连接在4x4x4 3D环面中的64个TPU组成。如果你看过Google关于TPU的宣传材料如下，这是8个TPU机架的图像。

注释：3D环面是一种三维网络拓扑结构，每个节点都与相邻节点连接，边缘的节点还会连接到对面，形成环状结构。这样设计可以减少任意两点间的最大距离。

 

 

添加图片注释，不超过 140 字（可选）

但在我们深入机架之前，我们需要澄清一些令人困惑的术语：机架 vs pod vs slice。

问："TPU Rack" vs "TPU Pod" vs "TPU Slice"有什么区别？

不同的Google资源使用它们略有不同，有时"TPU Pods"和"TPU Slices"可以互换使用。但在本文中，我们将坚持使用Google TPU论文和GCP TPU文档中使用的定义。

1. TPU Rack：

• 包含64个芯片的物理单元。也称为"cube"

1. TPU Pod：

• 可以通过ICI和光纤连接的TPU的最大单元
• 也称为"Superpod"或"Full pod"。例如，TPUv4的TPU Pod将包含4096个芯片或64个TPU Rack

1. TPU Slice：

• 介于4个芯片和Superpod大小之间的任何TPU配置

注释：这些术语定义了不同规模的TPU系统。Rack像是一栋楼，Pod像是一个小区，Slice像是小区中任意大小的一片区域。理解这些概念有助于了解Google如何组织和管理大规模的AI计算资源。

关键区别是TPU Rack和TPU Pod是物理测量单位，而TPU Slice是抽象单位。当然，设置TPU Slice有重要的物理属性，但我们现在将其抽象化。

现在，我们将使用物理测量单位：TPU Rack和TPU Pod。这是因为看到TPU系统如何物理连接在一起可以让我们更好地理解TPU设计理念。

现在回到TPU机架（针对TPUv4）：

单个TPU Rack由通过ICI和光电路交换（OCS）连接在一起的64个芯片组成。本质上，我们将多个tray连接在一起来模拟64个芯片的系统。这种将较小部分组合成超级计算机的主题在后面会继续出现。

注释：光电路交换（OCS）是一种使用光信号进行高速数据传输的技术。这种层级组合的思路类似于积木，小积木组成大积木，大积木组成更大的结构。

下面是单个TPUv4 Rack的图表。这是一个4x4x4 3D环面，其中每个节点是一个芯片，蓝色箭头是ICI，而面上的线是OCS。

 

 

添加图片注释，不超过 140 字（可选）

这个图表引发了几个问题。为什么OCS只用于面？换句话说，使用OCS有什么好处？有3个主要好处，我们将在下面介绍其他两个。

OCS的好处 #1：环绕

通过环绕实现节点间更快通信

OCS还为给定的TPU配置充当环绕。这将两个节点之间最坏情况的跳数从N-1跳减少到每轴(N-1)/2跳，因为每个轴都变成了环（1D环面）。

注释：跳数是指数据从一个芯片传到另一个芯片需要经过多少个中间芯片。环绕就像在地球上，从中国到美国可以向东飞也可以向西飞，选择更短的路径。这大大减少了通信延迟。

这种效应在我们进一步扩展时变得更重要，因为减少芯片到芯片的通信延迟对高并行化至关重要。

注释：并行化是指多个芯片同时工作来解决同一个问题。就像合唱团，如果成员之间不能快速协调，就无法演奏复杂的乐曲。

附：并非所有TPU都有3D环面拓扑

注：较老的TPU代（如TPUv2、v3）和推理TPU（如TPUv5e、TPUv6e）有2D环面拓扑，而不是像下面这样的3D环面。然而，TPUv7"Ironwood"似乎是3D环面，尽管它被宣传为推理芯片（注：我只是从他们的宣传材料推测）。

注释：2D环面是平面的环状连接，3D环面是立体的环状连接。就像从平面地图发展到立体地球仪，3D拓扑能提供更好的连接性能。

 

 

添加图片注释，不超过 140 字（可选）

Full Pod级别（又称"Superpod"；TPUv4的4096芯片）

就像我们将多个芯片连接在一起制成TPU Rack一样，我们可以连接多个机架制成一个大型Superpod。

Superpod还指TPU可以达到的最大相互连接芯片配置（仅使用ICI和OCS）。这根据代数而变化，但对于TPUv4是4096个芯片（即64个4x4x4芯片的Rack）。对于最新的TPUv7"Ironwood"，这是9216个芯片。

注释：4096和9216是非常大的数字，相当于几千到近万个芯片协同工作。这就像从一个小工厂发展到一个巨大的工业园区。

下图显示了TPUv4的一个Superpod。

 

 

添加图片注释，不超过 140 字（可选）

注意每个立方体（即TPU Rack）如何通过OCS相互连接。这也允许我们在pod中取TPU切片。

使用OCS的TPU Slice

我们可以请求pod内的TPU子集，这些就是TPU Slice。但即使你想要N个芯片，也有多种拓扑可供选择。

例如，假设你想要总共512个芯片。你可以要求立方体（8x8x8）、雪茄形状（4x4x32）或矩形（4x8x16）。选择slice的拓扑本身就是一个超参数。

注释：拓扑是指芯片的连接方式和排列形状。不同的形状适合不同的AI任务，就像不同的工具适合不同的工作。超参数是需要人工设定的参数。

你选择的拓扑将影响节点之间的通信带宽。这直接影响不同并行方法的性能。

例如，立方体（如8x8x8）对于全对全通信（如数据并行或张量并行）是首选，因为它具有最高的二分带宽。然而，雪茄形状（如4x4x32）对于流水线并行更好，因为它可以与顺序层更快地通信（假设一层适合4x4芯片的子切片）。

注释：全对全通信是指每个芯片都需要与所有其他芯片通信。数据并行是把数据分给不同芯片，张量并行是把模型分给不同芯片，流水线并行是把模型的不同层分给不同芯片按顺序执行。二分带宽是衡量网络连接性能的指标。

 

 

添加图片注释，不超过 140 字（可选）

当然，最优拓扑将取决于模型，找到这个拓扑本身也是一项工作。TPUv4论文还测量了这一点，以显示拓扑变化如何加速吞吐量（注：我不确定第一行指的是哪种LLM架构，因为没有指定）。

注释：LLM (Large Language Model) 是大型语言模型，如ChatGPT、Claude等。不同的模型架构需要不同的芯片排列方式来达到最佳性能，这需要大量实验来确定。

 

 

添加图片注释，不超过 140 字（可选）

我们介绍了TPU Slice，但有一个重要特性有助于TPU的高运行稳定性。

这些slice由于OCS不必是连续的机架。这是我们之前没有提到的使用OCS的第二个好处——可能也是最大的好处。OCS的好处 #2：（可重新配置的）非连续多节点Slice

注意这与硬连线多个节点来模拟非连续slice不同。由于OCS是交换机而不是硬连线，节点间的物理线路要少得多，因此允许更高的可扩展性（即更大的TPU Pod大小）。

注释：硬连线是指用物理电缆直接连接，交换机是可以灵活改变连接路径的设备。这就像传统电话需要接线员手动连接线路，而现代电话可以自动选择最佳路径。

这允许灵活的节点配置以进行扩展。例如，假设我们想在单个pod上运行三个作业。虽然简单调度不允许这样做，但OCS连接允许我们抽象化节点的位置，将整个pod视为"节点包"。

注释：这意味着不同的AI任务可以使用pod中不相邻的芯片，就像在停车场中，你可以使用任意位置的停车位，而不必使用连续的位置。这大大提高了资源利用率。

 

 

添加图片注释，不超过 140 字（可选）

这增加了pod利用率，在节点故障的情况下可能更容易维护。Google将此描述为"死节点的爆炸半径小"。然而，我不确定当只有某些节点必须关闭时，其液体冷却将如何受到影响。

注释：爆炸半径是指一个故障影响范围的大小。液体冷却是用液体来为芯片散热的技术，类似于汽车的水冷系统。

最后，这种灵活OCS的一个有趣扩展：我们还可以改变TPU Slice的拓扑，比如从常规环面到扭曲环面。

OCS的好处 #3：扭曲TPU拓扑

我们之前看到如何通过为固定数量的芯片改变(x,y,z)维度来实现不同的TPU Slice拓扑。但这次，我们将使用固定的(x,y,z)维度，而是改变它们如何连接在一起以实现不同的拓扑。

一个显著的例子是从雪茄形状的常规环面转到扭曲雪茄环面，如下所示。

注释：扭曲拓扑是指改变连接方式，就像把一张纸扭转后再连接两端，形成莫比乌斯环一样。这种设计可以缩短某些节点之间的通信距离。

 

 

添加图片注释，不超过 140 字（可选）

扭曲环面情况允许跨扭曲2D平面的芯片之间更快的通信。这对于加速全对全通信特别有用。

让我们深入了解一个具体场景，其中这将有所帮助。

使用扭曲环面加速训练

理论上，扭曲环面将为张量并行（TP）带来最大的好处，因为每层有多个all-gather和reduce-scatter操作。它可能为数据并行（DP）带来适度的好处，因为每个训练步骤也有一个allreduce，但这会不太频繁。

注释：all-gather和reduce-scatter是分布式计算中的通信模式，all-gather是收集所有数据，reduce-scatter是分散处理结果。allreduce是将所有节点的结果合并后再分发。这些操作需要大量的芯片间通信。

想象我们正在训练标准的仅解码器transformer，我们想采用大量并行来加速训练。我们将看到下面两个场景。

注释：仅解码器transformer是现代大型语言模型（如GPT、Claude）使用的主要架构。解码器负责生成文本，而编码器负责理解输入文本。

场景 #1：4x4x16拓扑（TP + PP；总共256芯片）

我们的z轴将是我们的流水线并行（PP）维度，我们的2D TP维度将是4x4。本质上，假设每层k位于z=k，每层在16个芯片上分片。如果没有明确绘制，假设标准OCS连接（即最近邻）。

注释：这意味着AI模型的不同层放在不同的z位置（如第1层在z=1，第2层在z=2），每一层的计算分布在4×4=16个芯片上。最近邻连接是指每个芯片只与相邻的芯片直接连接。

 

 

添加图片注释，不超过 140 字（可选）

我们将在每个z=k处扭曲2D环面，这使得每个TP层中芯片之间的通信更快。沿着我们的PP维度扭曲是不必要的，因为它们主要依赖点对点通信。

注释：点对点通信是指两个特定芯片之间的直接通信，不需要经过其他芯片。这就像流水线上相邻工位之间的直接传递。

注：实际上，当芯片数量大于4x4时，扭曲环面带来好处。我们使用4x4只是为了可视化目的。

场景 #2：16x4x16拓扑（DP + TP + PP；总共1024芯片）

作为扩展，我们将向之前的场景添加DP维度4。这意味着沿x轴有4个场景#1模型。

注释：这相当于同时训练4个相同的模型副本，每个副本使用不同的数据批次，最后合并结果。这是扩大训练规模的常用方法。

 

 

添加图片注释，不超过 140 字（可选）

注意扭曲环面限于每个DP模型的每个TP维度（即给定k=1…16的每个z=k的4x4 2D平面）。DP维度只有环绕，使每行成为大小为16的水平环。

你可能注意到有8x8x16的替代拓扑（即2x2 DP维度），但这变得更复杂，因为我们混合了DP和TP维度。具体来说，不清楚我们如何为y轴构建OCS环绕，同时为每个TP维度适应扭曲环面。

注释：这说明在设计大规模系统时，不同的优化技术之间可能会产生冲突，需要仔细平衡各种因素。

Multi-Pod级别（又称"Multislice"；TPUv4的4096+芯片）

 

 

添加图片注释，不超过 140 字（可选）

TPU层次结构的最终级别是Multi-pod级别。这是你可以将多个pod视为一台大机器的地方。但是，pod之间的通信是通过数据中心网络（DCN）完成的，其带宽比ICI低。

注释：数据中心网络是连接不同服务器机架的网络基础设施，就像连接不同建筑物的道路系统。由于距离更远，速度会比建筑物内部的连接慢一些。

 

 

添加图片注释，不超过 140 字（可选）

这就是PaLM的训练方式。它花了56天在6144个TPUv4（2个pod）上训练。下面你可以看到6个pod上的TPU作业分配：绿色是PaLM，红色是无分配，其余是其他作业。注意每个方块是4x4x4 TPU立方体。

注释：PaLM是Google开发的大型语言模型。56天的训练时间和6144个芯片显示了训练现代AI模型需要的巨大计算资源。

 

 

添加图片注释，不超过 140 字（可选）

实现这一点本身就很困难，但使这更令人印象深刻的是对开发者体验的关注。具体来说，它关于"我们如何尽可能抽象化模型扩展的系统/硬件部分？"这个问题。

注释：开发者体验是指程序员使用这些工具的便利程度。抽象化意味着隐藏复杂的技术细节，让开发者可以专注于AI模型本身而不是底层硬件。

Google的答案是让XLA编译器负责协调大规模芯片间通信。通过研究人员给出的正确标志（即DP、FSDP、TP的并行维度、slice数量等），XLA编译器为手头的TPU拓扑插入正确的分层集合（Xu et al, 2021: GSPMD）。目标是尽可能少的代码更改来实现大规模训练。

注释：FSDP (Fully Sharded Data Parallel) 是一种更高级的数据并行技术。分层集合是指在不同层级组织通信操作。这就像有一个智能助手，你只需要告诉它要做什么，它会自动处理所有复杂的协调工作。

例如，这里是Google博客中跨多个slice的all-reduce操作的分解。

 

 

添加图片注释，不超过 140 字（可选）

这表明XLA编译器处理slice之间和slice内的通信集合。

对于一个具体例子，可能有如下的TPU拓扑用于训练模型。激活通信通过ICI在slice内发生，而梯度通信将通过DCN跨slice发生（即跨DCN DP dim）。

注释：激活是神经网络中层与层之间传递的数据，梯度是训练过程中用来更新模型参数的信息。这种设计让高频的激活通信在高速网络内进行，低频的梯度通信通过较慢但足够的网络进行。

 

 

添加图片注释，不超过 140 字（可选）

05 将图表置于现实生活的角度

当你有实际硬件照片时，我发现将图表置于透视中很有帮助。下面是一个总结。

注释：看到真实的硬件照片有助于理解这些抽象概念在现实中的样子，就像看到真实的工厂照片有助于理解生产流程图。

如果你看过Google关于TPU的宣传材料图片，你可能遇到过下面这个图像。

 

 

添加图片注释，不超过 140 字（可选）

这是8个TPU Pod，其中每个单元是我们上面看到的4x4x4 3D环面。pod中的每行有2个tray，意味着每行有8个TPU芯片。

这里是单个TPUv4 tray：

 

 

添加图片注释，不超过 140 字（可选）

注意图表简化为只有一个PCIe端口，但在实际tray上有4个PCIe端口（在左侧）——每个TPU一个。

注释：PCIe端口是连接芯片与主机系统的接口，就像USB接口连接外设与电脑一样。

这里是单个芯片：

 

 

添加图片注释，不超过 140 字（可选）

中心部分是ASIC，周围的4个块是HBM堆栈。这是我们看到的TPU v4，所以它内部有2个TensorCore，因此总共有4个HBM堆栈。

注释：ASIC是实际进行计算的芯片核心，HBM堆栈是高速内存模块。4个HBM堆栈为2个TensorCore提供足够的内存带宽。

我找不到TPUv4的芯片平面图，所以这里是TPUv4i的，它很相似，只是它只有1个TensorCore，因为它是推理芯片。

注释：芯片平面图显示了芯片内部各个组件的布局，就像建筑物的平面图显示房间布局一样。

 

 

添加图片注释，不超过 140 字（可选）

注意CMEM在TPUv4i的布局中占用相当大的空间。

注释：CMEM占用大空间说明TPU设计中大容量片上内存的重要性，这与GPU设计形成对比。

 

 

微信视频号：sph0RgSyDYV47z6

快手号：4874645212

抖音号：dy0so323fq2w

小红书号：95619019828

参考文献链接

参考文献

[1] Google博客：TPU Multi-Slice Training [2] Xu, et al. "GSPMD: General and Scalable Parallelization for ML Computation Graphs" [3] Jouppi et al. "Ten Lessons From Three Generations Shaped Google's TPUv4i" [4] How to Scale Your Model - TPUs [5] Domain Specific Architectures for AI Inference - TPUs [6] HotChips 2023: TPUv4 [7] Google Cloud Docs: TPUv4 [8] Jouppi et al. "In-Datacenter Performance Analysis of a Tensor Processing Unit" -- TPU起源论文 [9] Jouppi et al. "TPU v4"-- TPUv4论文 [10] PaLM训练视频 [11] HotChips 2021: "Challenges in large scale training of Giant Transformers on Google TPU machines" [12] HotChips 2020: "Exploring Limits of ML Training on Google TPUs" [13] Google博客：Ironwood [14] HotChips 2019: "Cloud TPU: Codesigning Architecture and Infrastructure" [15] ETH Zurich的计算架构讲座28：脉动阵列架构 [16] Patterson演讲："A Decade of Machine Learning Accelerators: Lessons Learned and Carbon Footprint" [17] Camara et al. "Twisted Torus Topologies for Enhanced Interconnection Networks." [18] Horowitz文章："Computing's Energy Problem(and what we can do about it)"

所有人都能看懂的 TPU 硬核科普