AI系统-12谷歌跨时代AI芯片TPU

249 阅读26分钟

image.png

AI芯片最有动力又有技术实力公司那就是谷歌了,而且其代表了一个时代的技术创新,本文详细的讲解下TPU,带你走进AI芯片的历史和源头,让你感受下什么才是强大。

TPU基本一年更新一代,目前已经到了7代,详细参考:cloud.google.com/tpu/docs/sy…

为啥谷歌的产品这么超前跨时代,其他公司只有抄,还是抄几年前的产品构架的份?

个人感觉国内的公司对美国公司的技术抄还是抄人家大概5年前的。为什么国外这么厉害,感觉还是谷歌牛人多(高校、创新文化、资本、吸纳全球人才等),还有比较好的创新环境,给一个课题让这些牛人想怎么搞怎么搞,自由的搞。“创意精英”是谷歌管理理念中的另一个重要概念。

《重新定义公司:谷歌是如何运营的》中指出,未来企业的竞争力并非取决于管理者的控制能力,而是取决于能否吸引和留住具有创新能力的创意精英。创意精英不仅具备专业知识,还拥有极高的自主性和创新精神。他们敢于尝试新方法,不惧怕失败,并且能够从失败中总结经验。谷歌通过一种灵活、扁平的企业文化吸引了大量创意精英,并在公司的日常工作中持续赋予他们更多的自由权。谷歌的许多产品创新,正是来自这些创意精英在工作中碰撞出新想法,进而推动产品走向市场的过程。这样灵活的工作环境,让这些人才在谷歌找到满足感,从而在公司内部激发出源源不断的创新力。

本文主要参考《ZOMI酱:AI系统》:chenzomi12.github.io/02Hardware0…

1. TPU发展历史

谷歌作为互联网公司,在互联网时代积累了海量的数据,当处理这些数据的时候就需要算力,最开始2006年利用数据中心的剩余算力就可以。但是随着手机语音识别功能的需求剧增,那么当时谷歌的数据中心需要双倍的算力才能满足日益增长的计算需求,而仅仅依靠传统 CPU 来满足这种需求是非常昂贵的。于是,在这个背景下,谷歌开始了 TPU 的设计。

通常一个芯片的开发需要几年的时间,然而谷歌不愧是谷歌,TPU 从立项到大规模部署只用了 15 个月。TPU 课程的领头人 Norm Jouppi 说到:“芯片设计过程异常迅速,这本身就是一项非凡的成就。令人惊叹的是,我们首批交付的硅片无需进行任何错误修正或掩模的更改。考虑到在整个芯片构建过程中,我们还在同步进行团队的组建,紧接着迅速招募 RTL(寄存器传输级)设计专家,并且急切地补充设计验证团队,整个工作节奏非常紧张。”

从这个过程中就可以看出互联网所追求的速度和试错精神,跟传统的做设备公司是不一样的,先搞出来验证市场再优化,市场接受的产品虽然是豆腐渣工程也可以再找牛人重构,反正有资金支持。另外互联网公司跟资本挂钩的很重,资本也在盯着先进的技术应用,所以有需求和技术那就干。

image.png

  • 1代只搞推理
  • 2代及之后的可以用于训练和推理,28nm改为了16nm
  • 3代在2代基础上增大了性能,TPU v3 的一个显著改进是其 2D torus 互连结构
  • 4代全新架构使用7nm,TPU v4 首次应用了 3D torus 的互联方式,提供了比 2D torus 更高的带宽和更优的性能,能够支持多达 4,096 个 TPU v4 核心,在 TPU v4 POD 中总共提供了 1.1260 exaflops 的 BF16 峰值算力。
  • Pod是互联的超级计算集群

实体是这样: image.png

在终端消费市场,具体就是手机上也是需要搭载TPU的,歌基于 Edge TPU 的框架研发了继任芯片 —— Pixel Neural Core,在 2019 年 10 月发布的 Pixel 4 上首次搭载。

image.png

在移动端集成AI芯片,其实相当的炸裂,首先在手机上应用,但是这直接为智能汽车、智能机器人的实现铺平了道路。

其实在消费市场,如果有黑科技像手机里面集成了NPU/TPU那么就可能是杀手锏,跟不上就是死路一条,但是又有技术壁垒,小的厂商没有资本跟上很可能就面临淘汰。

2. TPU架构

张量处理单元 (TPU) 是 Google 设计的专用集成电路 (ASIC),用于加速机器学习工作负载。Cloud TPU 是一种 Google Cloud 服务,可将 TPU 用作可扩缩资源。

TPU 旨在快速执行矩阵运算,非常适合机器学习工作负载。

CPU 是一种基于冯·诺依曼结构的通用处理器。这意味着 CPU 与软件和内存协同工作,如下所示:

image6.gif CPU 最大的优点是它们的灵活性。您可以在 CPU 上为许多不同类型的应用加载任何类型的软件。例如,您可以使用 CPU 在 PC 上进行文字处理、控制火箭发动机、处理银行交易或通过神经网络对图片进行分类。

对于每次计算,CPU 从内存加载值,对值执行计算,然后将结果存储回内存中。与计算速度相比,内存访问速度较慢,并可能会限制 CPU 的总吞吐量。这通常称为冯·诺依曼瓶颈

GPU 的工作方式

为了提高吞吐量,GPU 在单个处理器中包含数千个算术逻辑单元 (ALU)。现代 GPU 通常包含 2500 - 5000 个 ALU。如此大量的处理器意味着您可以同时执行数千次乘法和加法运算。

image2.gif 这种 GPU 架构非常适合并行处理大量运算(例如神经网络中的矩阵运算)的应用。实际上,在用于深度学习的典型训练工作负载上,GPU 的吞吐量可比 CPU 高出一个数量级。

不过,GPU 仍然是一种通用处理器,必须支持许多不同应用和软件。因此,GPU 与 CPU 存在同样的问题。对于数千个 ALU 中的每一次计算,GPU 都必须访问寄存器或共享内存,以读取运算对象以及存储中间计算结果。

TPU 的工作方式

Google 设计了 Cloud TPU,它们是专门用于神经网络工作负载的矩阵处理器。TPU 不能运行文字处理程序、控制火箭引擎或执行银行交易,但它们可以很快地处理神经网络中使用的大量矩阵运算。

TPU 的主要任务是矩阵处理,这是乘法和累加运算的组合。TPU 包含数千个乘法累加器,这些累加器彼此直接连接以形成大型物理矩阵。这称为脉动阵列架构。Cloud TPU v3 在单个处理器上包含两个 128 x 128 ALU 的脉动阵列。

TPU 主机将数据流式传输到馈入队列中。TPU 从馈入队列加载数据,并将其存储在 HBM 内存中。计算完成后,TPU 会将结果加载到馈出队列中。然后,TPU 主机从馈出队列读取结果并将其存储在主机的内存中。

为了执行矩阵运算,TPU 将 HBM 内存中的参数加载到矩阵乘法单元 (MXU) 中。

image4_5pfb45w.gif

然后,TPU 从内存加载数据。每次执行乘法运算时,系统都会将结果传递给下一个乘法累加器。输出是数据和参数之间所有乘法结果的总和。在矩阵乘法过程中,不需要访问内存。

image1_2pdcvle.gif 因此,TPU 可以在神经网络计算中实现高计算吞吐量。

3. TPU V1

第一代 TPU 主要服务于 8 比特的矩阵计算,由 CPU 通过 PCIe 3.0 总线驱动 CISC 指令。采用 28nm 工艺制造,频率为 700MHz,热设计功耗为 40 瓦。

初代 TPU 主要针对 2015 年左右最火的神经网络进行优化,主要分为以下三类:

  • MLP 多层感知机(MultiLayer Perceptron)
  • CNN 卷积神经网络(Convolutional Neural Network)
  • RNN 递归神经网络(Recurrent Neural Network)& LSTM 长短期记忆(Long Short-Term Memory)

而在这三类中,由于 RNN 和 LSTM 的高复杂度,初代 TPU 只能在前两种模型框架的推理场景进行优化。

TPUV1的三点优化:

  1. 低精度Int8:因为对于一些智能任务例如图像语音识别降低精度对最后的识别结果影响不大。
  2. MXU(Matrix Multiply Unit,矩阵乘法单元)的计算使用脉动阵列:如下图就是取数据-》计算-》存数据,这个组合过程串起来,计算完不用存数据了,直接给下一次计算使用。打个恶心的比方就是狗吃屎,屎还没落地就被狗嘴接住吃了,省事。

image.png

  1. 专用硬件:专门为神经网络推理设计,由于目标的单一化,TPU 上的控制逻辑(control logic)只占了芯片的 2%(远低于 CPU 和 GPU)。同时,由于 TPU 在简化设计和增加芯片内存基础上又缩小了芯片的大小,TPU 的成本控制和良品率也远远优于其他体型更大的芯片。

image.png

简略的框图:

image.png

TPU v1 上几个关键的元素:

  • DDR3 DRAM / Weight FIFI: 第一个讲的是谷歌的 Weight FIFO,位于图像的右上角。在模型推理场景,通过 DDR3-2133 接口,模型权重会被存储在 TPU v1 上的 DDR3 RAM 内。这些权重通过 PCIe 从主机计算机的内存中“预加载”到这些芯片上,然后可以传输到芯片的 Weight FIFO 内存中,以便矩阵乘法单元(MXU)使用。

  • 矩阵乘法单元 MXU(MMU) : MXU 以脉动阵列的形式工作,能够提供 256×256×8 bit 的乘加计算,在每个时钟周期输出 256 个 16 bit 的部分和计算结果。矩阵单元里面包含一个 64 kb 的 Weight Tile 以及一个双缓存单元用于缓存回调。MXU 被谷歌的工程师们描述为“TPU 的心脏”,在本节后面的内容中,我们会更加细致的去剖析这一部分的设计。

  • 累加器 Accumulators: 这是一个能够存储 4 MiB 的 32-bit 数据的累加单元,用来存储 MXU 计算后的结果。这 4 MiB 代表着 4096 个,每个含 256 个元素的 32 位累加器。为什么是 4096 呢?原因是谷歌的工程师注意到每字节的运算次数需要达到峰值性能大约是 1350,继而将其向上舍入到 2048,然后再翻倍让便编译器在运行至峰值性能时能使用双缓冲,这也就是为什么这个累加器被设计为 4 MiB。

  • 控制指令 Control: 每个芯片都需要一个控制模块,而在 TPU 中,整个控制单元采用了四级流水线设计。控制单元的主要任务是接收通过 PCIE 总线或主机(CPU)传递的指令,并将这些指令执行于 TPU。这些指令源自 CPU,而芯片的指令集为 CISC,共包含 12 条指令。TPU 采用 CISC 而非更简单的 WISC 指令集的原因是,谷歌定义的每条指令的平均执行周期为 10-20 个时钟周期,这使得每条指令相对复杂。特别是,TPU 中的各种单元,如 MXU、UB 和 AU,都定义了一些专门为神经网络设计的高级复杂指令:

    • Read_Host_Memory: 从 CPU host 读取数据到 Unified Buffer
    • Read_Weight: 从 Weight DRAM 读取数据到 Weight FIFO
    • Matrix_Multiply/Convolve: 执行乘法或卷积运算
    • Activate: 执行 ReLUSigmoid 等激活计算
    • Write_Host_Memory: 把计算结果数据从 Unified Buffer 输出到 CPU host

通过这些指令,TPU 能够顺序地执行读取、写入、计算和激活操作,从而处理神经网络各层的具体计算需求。

那么在神经网络AI计算盛行的今天,大家还都买GPU去计算,为啥不去买专门适合AI的TPU,答案就是:谷歌不卖。谷歌只卖成品的云计算服务,这样就用不着你自己买卡计算了,直接买更贵的成品,这其实也是傻瓜化,这样谷歌挣更多的钱。也可能芯片并不是谷歌的强项,拿去直接竞争干不过芯片公司。

4. TPU V2

TPU v2 的设计集成了高带宽的存储解决方案,具有 16 GiB 的 HBM,并能够提供最高 600 GB/s 的内存带宽,以及 45 TFLOPS 的浮点运算能力,用于支持更加高效的内存访问、数据操作、和复杂运算。是一个训练卡。

进一步地,谷歌将四个 TPU v2 排列成了性能为 180 TFLOPS 的四芯片模块,并将这样的 64 个模块组成一个一共有 256 片 TPU v2 集成的 TPU v2 Pod,理论峰值计算量达到了恐怖的 11.5 PFLOPS。

训练业务场景的变化:

image.png 推理在训练好的模型结构和参数基础上,一次前向传播得到模型输出过程。相对于训练,推理不涉及梯度和损失优化,因此,协同着神经网络模型的鲁棒性,模型对于数据精度的需求则较低。

相比之下,训练过程是一个比推理复杂很多的过程。通常来说训练过程通过设计合适 AI 模型结构以及损失函数和优化算法,将数据集以 mini-batch 反复进行前向计算并计算损失,反向计算梯度利用优化函数来更新模型,使得损失函数最小从而使得模型收敛。在训练过程中,最关键的步骤就是梯度的计算和反向传播,同时在这个过程中也需要计算工程中常见的优化器、学习率调度器等工具,因此数据计算的精度便变得非常重要,一旦在某个计算节点出了问题,则有可能引发梯度爆炸、梯度消失等问题导致模型无法收敛。

那么细化下来,训练场景到底有哪些难点呢?

  • 更难的数据并行:推理阶段,每个推理任务都是独立的,因此 DSA 芯片集群可以横向拓展。而在训练场景,一个模型就需要迭代百万次,模型所有的参数在每一个迭代中都要进行调整,因此需要协调跨集群进行并行计算。
  • 计算复杂度更高:反向传播需要对模型每一个阶段的每一个权重参数和输入的偏导数进行计算,包括数据格式更高精度的激活值和转置权重矩阵的矩阵乘法计算(WTδ)。
  • 内存需求更大:权证更新的数据来自于前向和反向传播的临时变量,在模型的每一层,临时变量都需要被保留,因而提高了对于现存的负荷,在大模型中临时变量,包括每一层的激活值、优化器的值会将模型的大小膨胀至原始大小的 8~9 倍。
  • 更具可编程性:训练算法和模型的快速变化可能会导致在芯片设计时的设计很快被淘汰果实,因此训练芯片需要有更高的可编程性从而适应日新月异的模型架构更新。
  • 高精度的数据格式:整数 INT8 可以用于推理,但是训练期间需要充分捕捉梯度信息,通常需要 FP16,BF16 和 FP32 等混合精度计算以保证模型能够收敛。

为了解决上面的问题在1代的基础上进行改动:

4.1 TPUV2相对于V1的改动

改动一:Vector Memory

在 TPU v1 中,我们可以看到它有两个存储区域:

  • Accumulator 负责储存矩阵乘积结果
  • Activation Storage 负责储存激活函数输出

image.png 在推理场景中,专门的存储模块对计算和存储非常有帮助,因为它们更加 Domain Specific。但是在训练过程中,为了提升可编程性,如下图所示,TPU v2 交换了 Accumulators 和 Activation Pipeline 这两个独立缓冲区的位置,并将它们合并为 Vector Memory,变得更像传统架构中的 L1 Cache,从而提升了可编程性。

改动二:Vector Unit

image.png

在 TPU v1 中,Activation Pipeline 是专门针对激活函数场景特殊处理的,即卷积之后有一个 Batch Normalization 再接一个特殊的激活函数的 ALU 计算。然而,这种特殊的 ALU 计算无法满足训练场景的需求。因此,如下图所示,曾经的 Activation Pipeline 在 TPU v2 中变成了一个 Vector Unit,用于专门处理一系列的向量激活函数。

改动三:MXU

image.png

在 TPU v1 中,MXU 是和 Vector Memory 相连接的,而在 TPU v2 中,谷歌将 MXU 和 Vector Unit 进行了连接,所有的数据出口和计算都由 Vector Unit 进行分发。如此一来,MXU 所有的矩阵相关计算就成为了 Vector Unit 的协处理器,这对编译器和编程人员更加友好。

改动四:DDR3

image.png

在 TPU v1 中使用 DDR3 内存是为了将推理场景中需要用到的一些权重直接加载进来用于计算。然而在训练过程中会产生许多中间层的变量和权重,这时 DDR3 的回写速度无法满足需求。因此,在 TPU v2 的训练过程中,如下图所示,我们将 DDR3 与 Vector Memory 放在一起,并将 DDR3 的位置换成 HBM,从而加速了 20 倍的读写速度。

4.2 TPU V2架构

那么改进后的TPU core如下:

image.png

标量单元(Scalar Unit)

上图是 TPU 核心的简单图示,我们可以看到标量单元是处理计算的起点,它从指令存储器中取出完整的 VLIW(超长指令集),执行其中的标量操作,并将指令传递给向量和矩阵单元进行后续处理。VLIW 由两个标量槽、四个向量槽(其中两个用于向量加载/存储)、两个矩阵槽(一个推和一个弹出)、一个杂项槽(一个简单的例子是延迟指令)和六个立即数组成。

Alt text

那我们看一下指令到底是从哪里获得的。Core Sequencer 不再从 CPU 获取指令,而是从 Instruction Mem 取出 VLIW 指令,使用 4K 32-bit 的标量内存来执行标量运算,有 32 个 32 位的标量寄存器,而将向量指令送到 VPU。322-bit 宽的 VLIM 可以发送 8 个操作,两个标量,两个向量 ALU,向量 load,向量 store,和一对从矩阵乘法来的排队数据。

这个从内存获取指令的操作比较逆天,这就是存算一体啊。可谓一切为了吞吐量,CPU拖后腿那就不用CPU,直接用自己快速的内存和搭配自己的硬件来解决。这个标量部分应该有自己的运算核(小型定制CPU,例如RISC-V或者其他IP)。

矢量单元(Vector Unit)

如下面右侧的图所示,和 Scalar Unit 相连接的就是 Vector Unit,下图左图是 Vector Unit 其中的一个 Vector Lane(矢量通道),整个 Vector Unit 包含 128 个这样的矢量通道。每个通道包含一个额外的 8 路执行维度,称为子通道。每个子通道配备一个双发射 32 位 ALU,并连接到一个 32 深度的寄存器文件。这种设计允许矢量计算单元在每个时钟周期内同时操作 8 组 128 宽度的矢量,子通道设计提高了矢量与矩阵计算的比例,特别适用于批量归一化操作。

Alt text

矩阵乘法单元(MXU)

MXU 一直是 TPU 的核心,TPU v2 选择使用 128×128 的 MXU 而不是和 TPU v1 一样的 256×256 的 MXU,主要是由于利用率和面积效率的综合考虑。在下面这张图中,谷歌的模拟器表明,四个 128×128 的 MXU 的卷积模型利用率为 37%到 48%,明显高于一个 256×256 的 MXU(22%到 30%),且占用的芯片面积相同,原因在于某些卷积计算天生小于 256×256,所以 MXU 的部分会闲置。这意味着在相同面积下,使用多个较小的 MXU 可以实现更高的计算效率。尽管十六个 64×64 的 MXU 利用率稍高(38%到 52%),但它们需要更多的面积,因为较小的 MXU 面积受 I/O 和控制线的限制,而不是乘法器。因此,256×256 的 MXU 在带宽、面积和利用率之间达到了更好的平衡,使其成为 TPU v2 的最佳选择。

Alt text

除了变得更小,这一代的 MXU 还有以下主要特点:

  • 数值和精度:使用 bfloat16 格式进行乘法运算,该格式具有与 float32 相同的指数范围,但尾数位数较少。累加则使用 32 位浮点数进行。
  • 能效:与 IEEE 16 位浮点数相比,bfloat16 具有大约 1.5 倍的能效优势。
  • 采用和易用性:相比 fp16,bfloat16 对于深度学习更易于使用,因为它不需要进行损失缩放。由于这种易用性以及能效和面积效率,bfloat16 已被业界广为采用,也是深度学习最重要的 building blocks 之一。

转置/规约/置换核心(TRP Unit)

这一部分的核心用于进行 128×128 的矩阵的特殊计算操作(如转置 Transpose、规约 Reduction、置换 Permute),允许矩阵数据重新排布,提升编程易用性。这些步骤都是在训练过程中,特别是反向传播中经常遇到的矩阵相关场景,TPU v2 因此对于这一部分进行了特殊优化。

4.3 TPU V2互联

image.png

Interconnect Router这个模块能够实现 2D 的环面连接(2D Torus),从而组成 Pod 超级计算机。TPUv2 Pod 在提供了传统集群交换机 2.5 倍的双向带宽的同时,省去了 Infiniband 网络卡、Infiniband 交换机以及通过 CPU 主机通信的延迟成本,大大提升了整体计算的效率。

总体的架构图就是下图,下图中的线越宽代表越大的带宽。所以说里面 TPU 核心和 HBM 存储的带宽是最大的,其次就是 TPU 核心、HBM 和 Interconnect Router 的带宽。

Alt text

一般第二代的产品都差不到哪里去

一般第一代产品主要还是市场验证和应急使用,设计时就需要精简,因为复杂了可能芯片流片失败,功能实现异常直接就搞砸了,而且成本巨大,可能都没有研发第二代的机会了。一般第一代能用肯定就不会搞第二代了,那么第二代一定是要把第一代的缺点和新的需求都集成进去,而且会大胆的设计,另外设计时间充足对于架构也有非常大的优化,基本就定型了产品,后续新的产品也基本要基于这个框架搞了,这也是历史的包袱。所以TPU的架构基本第二代就定型了,后续就是堆叠、局部速度优化、算子配置替换等。

5. TPU V3

image.png TPU v3 实际上就是 TPU v2 的增强版。TPU v3 相比 TPU v2 有约 1.35 倍的时钟频率、ICI 带宽和内存带宽,两倍 MXU 数量,峰值性能提高 2.7 倍。除了显眼的蓝色外,相比于 TPU v2,TPU v3 在只增大 10%体积的情况下增加了 MXU 的数量,从 2 个翻倍到了 4 个。同时 TPU v3 时钟频率加快了 30%,进一步加快了计算速度;同时内存带宽扩大了 30%,容量翻倍;此外芯片之间的带宽也扩大了 30%,可连接的节点数是之前的 4 倍。

回到 2017-2018 年左右,Bert 出现之前,很多人是不相信一个模型需要用到一个集群进行训练的,因为当时很多的模型只需要单卡就能进行训练。而实际上的情况是,仅有(从当前的视角看)3 亿参数 Bert 在 4 个 TPUv3 Pod 上训练了整整四天,而当前各家的万亿参数的模型大部分都是用万卡的集群训练数个月的结果。

TPU v2 和 v3 大幅度强化了芯片的互联能力,最主要的核心就是为了搭建这样的超大计算集群。

首先看下TPU V2的集群: image.png

  • A:四个 TPU v2 芯片和散热片
  • B:2 个 BlueLink 25GB/s 电缆接口。其中 BlueLink 是 IBM BlueLink 端口协议,每 Socket 25Gb/s 的带宽,主要是提供 NPU 或是 TPU 之间的网络互联。
  • C:Intel 全路径体系结构(OPA)电缆。其中 OPA 为英特尔 Intel Omni-Path Architecture(OPA)互联架构,与 InfiniBand 相似。
  • D:电路板电源连接器
  • 支持两种网络配置,分别为 10Gbps 以太网和 100Gbps Intel OPA 连接

image.png

  • A 和 D:CPU 机架
  • B 和 C:TPU v2 机架
  • 蓝色框:电源管理系统(UPS)
  • 红色框:电源接口
  • 绿色框:机架式网络交换机和机架式交换机顶部,这部分更多的是网络模块

TPU V3的集群:

image.png

  • A:四个 TPU v2 芯片和液冷散热管;
  • B:2 个 BlueLink 25GB/s 电缆接口
  • C:Intel 全路径体系结构(OPA)电缆
  • D:电路板电源连接器
  • 支持两种网络配置,分别为 10Gbps 以太网和 100Gbps Intel OPA 连接

image.png 从上面的 TPU v3 Pod 的形态中我们就可以看到,相比于 TPU v2 Pod,它的规模大了非常多,有了更多的铜管和电缆,并且在芯片规模上整整大了 4 倍。TPU v3 Pod(1024 chips)

下面是虚拟架构图,整体的架构图也是比较明显的。AI 框架通过 RPC 远程连接到 TPU Host,基于 CPU 去控制 TPU 去实现真正的互联运作执行。

Alt text

  • TPU v2 的技术革新:谷歌的 TPU v2 通过增加核间互连结构(ICI),使得最多 256 个 TPU v2 能够组成一个高效的超级计算机。这种结构支持高效的大规模数据处理,尤其适合神经网络的训练。
  • TPU v3 的性能提升:谷歌进一步扩展其技术,通过组合 1024 个 TPU v3 创建了 TPU POD 超级计算机。该服务器采用水冷系统,功率提升高达 1.6 倍,而模具尺寸仅比 TPU v2 增加 6%。
  • 高效的集群构建:TPU v2 集群利用交换机提供的虚拟电路和无死锁路由功能,加上 ICI 结构,形成了高效的 2D tours。这种配置提供了 15.9T/s 的平分带宽,相比传统的集群组网,省去了集群网卡、交换机的成本,以及与集群 CPU 的通信延迟。

6. TPU V4

TPU v4 通过采用 7 nm 工艺,翻倍的 MXU 数量,以及显著增加的内存容量和带宽,实现了性能的大幅提升。特别引人注目的是 Sparse Core 的引入,它专门针对稀疏计算进行了优化,使得 TPU v4 在处理深度学习中的 Embedding 层时更加高效。此外,3D Torus 互联方式的首次亮相,以及 TPU v4 Pod 的构建,展示了谷歌在大规模并行计算和高效互联方面的突破。Palomar 光路开关芯片的开发,利用 MEMS 技术,进一步降低了系统延迟和功耗,同时减少了网络成本。

image.png

TPU v4 的架构图,每个 TPU v4 都有两个 Tensor Core,每个 Tensor Core 里面有分别六个单元组成,其中有四个是 TPU 的最核心的脉动阵列 MXU,另外还有两个单元,一个是 Scalar Unit,另一个是 Vector Unit,这些都在前面的章节讲过。同时我们还看到两个 HBM 也就是内存模块,为了降低电缆的时延,它们分别被放在了 Tensor Core 的左右两边。

image.png

  • 制造工艺:从前代的 16 nm 提升到了最新的 7 nm 工艺。
  • 硬件增加:TPU v4 上的 MXU 数量相比前代 TPU v3,又翻了一倍,而内存则增加了 9 倍,增加到了 244 MB,虽然 HBM 内存的容量依然是 32 GB,但内存带宽增加到了 1.2 TB/s,相比前代提升了 33%。
  • Sparse Core:TPU v4 的硬件增加了 Sparse Core 的支持用于利好稀疏计算,基于 TPU v4 改良的自有 Transformer 模型架构
  • 3D torus:TPU v4 首次亮相了 3D torus 的互联方式,紧密耦合 4096 个 TPU v4 引擎,使得 TPU v4 POD 总计提供 1.126 Exaflops 的 BF16 峰值算力。