Google TPUv2/v3架构及其设计演化

Google TPUv2/v3架构及其设计演化

  • 发布时间:2020-09-24 10:25
  • 访问量:

【概要描述】TPUv2是Google在TPUv1的Systolic Array基础上使用VLIW指令架构进行的大规模架构更新,TPUv3则在TPUv2基础上的PPA优化。

Google TPUv2/v3架构及其设计演化

【概要描述】TPUv2是Google在TPUv1的Systolic Array基础上使用VLIW指令架构进行的大规模架构更新,TPUv3则在TPUv2基础上的PPA优化。

  • 分类:研究院
  • 作者:壁仞科技研究院
  • 来源:
  • 发布时间:2020-09-24 10:25
  • 访问量:
详情

摘要

TPUv2是Google在TPUv1的Systolic Array基础上使用VLIW指令架构进行的大规模架构更新,TPUv3则在TPUv2基础上的PPA优化。本文基于公开信息讨论Google TPU的架构演进,供大家参考,如有错误疏漏还请批评指正。

 

Table 1TPU v1-v4 chip架构指标[2]


Table 1 TPU v1-v4 chip架构指标[2]

注:向量单元数量为推导数值,TPUv4数值源自MLperf training 0.7[3]

 

TPUv1到TPUv2的设计思路和架构演化

 

TPUv1对于google来说是其进入芯片设计领域的第一步试水的作品,主打inference。而TPUv2的战略意义在于补齐了google在training方面的能力,在这之前training一直是GPU的天下,而其它架构例如CPU/DSP的架构和GPU相比各方面都无法抗衡,TPU在推出之后,google在training领域也能和GPU竞争了。

关于ML training的一些挑战,google认为主要有以下几个方面:

1 计算:训练需要反向计算,包括转置,batch-norm统计,这些是前向所没有的需求。

2 存储:训练需要存储前向的activation,用于计算反向的gradient。所以需要更大的存储空间。

3 精度:训练需要更大的动态范围,所以int8已经不够用了,必须上浮点。

4 编程性:新型神经网络的研发非常快,需要支持不同的调参优化器,需要有能力构造新型的神经网络。

5 并行性:对于训练而言,如何将多个device组成一个超大系统成为设计目标之一,多卡和多节点的支持变成硬性需求。

Google在文献2中介绍了一下的自己团队,看上去人可能并不是很多,并不像GPU/手机芯片这种动则几百人的前端团队,用中国的话说,google的team 非常精干。但是,我认为,也正是这种人员配置,使其只能选择一定的复杂度的架构和schedule,是google选择偏向VLIW架构的原因之一。

中间google花了详细的几页PPT介绍了一下自己的TPUv1到TPUv2的演化路线。

 

Figure 1TPUv1 架构图[1]

 

Figure 2TPUv2架构图[1]

 

我提取精要,总结如下:

  • Dedicate的activation storage和accumulators设计被取消。
  • 增加了vector memory。
  • activation pipe换成一个vector unit,其中替代了acc和activation pipe处理的功能。
  • matrix multiply unit不再直接访存,而是通过vector unit进行数据的读取和写回。
  • 将DDR替换乘HBM,latency更小,带宽更大。
  • 同时为了支持多卡多节点互联,设计了interconnect模块。

这里的改动基本上都是google针对training的特点而进行的架构升级,这里提一下第4点,这样的好处我认为有两点,第一,也就是op fuse,矩阵乘卷积操作后不用写回memory就能进行element wise操作。第二,DDR3之间和matrix multiply unit进行load store 操作变成了matrix multiply unit和vector memory的操作,这样效率上可能会有一定损失,但是访存灵活度提高。可以实现更灵活的数据复用策略。估计这里可能也参考了GPU的tensor core的一些访存设计。

在我看来,TPUv2实现了从fixed pipeline结构加速器到可编程处理器结构的思路转变,同样的存储单元和运算单元,变成了程序显式可编程单元,而不是简单的配置寄存器,硬件调度,这样可编程性和灵活度大大提高,而这些正是训练时所需要的。

 

TPUv2模块详细介绍

 

TPUv2的ISA 是VLIW形式, 一个指令bundle 322 bit,里面包括2个scalar slot,4个vector slot (2 for load/store) 2个matrix slot(push,pop)和1个misc slot,也就是一条指令里面包含9个指令slot, 指令包里面有6个立即数用于计算或者寻址。

对于VLIW而言,硬件设计相对简单,但是指令之间fetch和issue是紧耦合的,一旦有一个slot的指令有dependency,导致流水线hazard,那么整个指令包将被stall住,因此complier编译设计的时候十分小心,不然slot之间互相牵制,产生很多空nop指令slot,导致流水线气泡,performance惨不忍睹。

TPUv2整个chip结构里面包括core x 2、pcie x 2、interconnect router、link x 4、HBM x 2。

 

Figure 3TPUv2 chip结构[1]

 

其中Core里面主要分为4块:

·标量单元

·向量单元

·矩阵乘法单元

·转置和重排单元

标量单元相当简洁清楚,标量单元负责从instruction buffer中fetch整个VLIW指令bundle,然后issue到ALU或者mem,进行load/store/mov/alu/dma操作,这里标量配有2路ALU,32x32b register file,4kx32b的memory,这里还有sync flags模块用于进行一些外部的同步控制操作。另一方面标量部分还控制着往vector和matrix unit进行issue指令的过程。

 

Figure 4 TPUv2标量单元[1]

 

向量处理部分主要有vector memory,vector reg, vector alu, vector memory的大小为32k x 32b=128KB,register file是32个entry,32bit,一共8个lane,两路vector ALU,本人猜测这里有一个是激活函数,一个进行乘加mac,这里我们可以知道访存的宽度是32x128=256b,一共有128组这样的vector unit,所以一个core的memory是16MB, mac数是1024。而matrix unit的访存操作都通过vector reg进行,注意这里从矩阵输出的数据还有一些绿色的storage被标出用于缓冲数据。

Google提到,TPUv2的向量单元在设计之初由于batch-norm的加入导致vector算力成为瓶颈,因此在TPUv2中大幅强化了vector能力[2]。

 

Figure 5 Tpuv2向量单元[1]

 

矩阵单元脉动阵列是TPUv1留给TPUv2的核心财富之一。包括128x128的脉动阵列,bf16格式的乘法,fp32的累加单元,其中左矩阵和结果是流动的,右矩阵是stationary在矩阵中的。且只有右矩阵能进行转置操作,这里我们从反向算法角度,可以推测出,在计算正向和反向activation gradient的时候,右矩阵应该是输入weight,而在计算反向weight gradient的时候,input activation应该在右矩阵。Google详细说明了其选择128x128这种脉动阵列原因,红线表示每个操作数会进行多少次乘加操作,蓝线表示矩阵乘法的mac利用率。对于design来说,1个256x256单元和16个64x64单元和4个128x128单元所用的mac总的数量是相同的,都是65536个mac。

 

Figure 6Tpuv2矩阵单元尺寸设计[1]

 

首先是operations per operand ratio这条红线,对于256x256的脉动阵列,对应的逻辑操作是矩阵乘法m256xn256xk256, 每个operand in A会在算256个结果元素时用到,同理,64x64,每个operand in A会被64个结果元素用到,在数据复用性上来说,前者是后者的4倍,但是要找到一个应用有很大的数据复用性并不容易。另外一条蓝线是mac利用率,对于尺寸更大的矩阵乘来说,某些应用就是一些小的矩阵乘法,其算力利用率不高。而数据复用性和利用率对于mac array来说其实是两个相矛盾的指标,因此2者权衡之下,Google选择了4个128x128的matrix multiply单元。

另外,同样mac数量的矩阵,面积占比不一样,大矩阵乘法器运算单元占比大,而小的矩阵单元,IO和control面积会占比较大。

这里强调一下,TPUv1的矩阵单元只能进行int8操作,TPUv2/TPUv3的矩阵单元只能进行bf16操作,并使用bf16进行inference运算,这种单一的数据格式使得整个芯片功耗和面积上面TPU对NV tensor core都具有一定优势,但是在inference应用的TOPS/W会不如Nvidia T4。

另外,TPU中还提供了矩阵的transpose,reduction,permutation等操作,能够实现不同lane之间数据的互换重排。

Memory系统是芯片重要设计部分,对于访存操作, load/store对象都是SRAM scratchpads,同时Google提到了striding over vector,能从一块地址连续的数据中取出某些数据段,这些数据的首地址按stride跳跃的,数据的搬移通过异步DMA操作进行,软件需要参与管理。但是好处是系统的访存调度是可预测的,对应performance分析很容易得到准确的数值。Google 在load store操作中通过sync flags来进行stall或者indicate操作完成,防止out of order造成的程序语义错误。

 

Figure 7 Tpuv2 2D torus结构[1]

 

 

Figure 8 TPUv2 chip互联结构[1]

 

Google的TPU board 互联设计中包含一个router,和2个core相连,同时将2个HBM也连入router,router上面有4个link,每个link 400Gbps,这些link和其它TPU board相链接,对应2D torus结构中虚线框中的8个输入输出。2D torus的ring的特点在于每次reduce操作都在充分利用所有core的算力,过程是先是在水平方向做scatter-reduce,然后在垂直方向做all reduce,最后在水平方向做all-gather进行数据复制。从软件看来,DMA操作就像访问HBM,用不同chip id来区分不同的core。

 

Figure 9TPUv2 floorplan[2]

 

TPUv2的Floorplan的设计看起来中规中矩,值得注意的是,这里vector unit/memory分别到了两个physical partition。当中间的matrix 单元在读写数据时,其实它面对的是一个物理上分布式的存储结构,而在程序看来是一块完整的L1 memory,这里和含光的分布式L1设计有类似的地方。

 

TPUv3 架构演化

TPUv3 架构其实就是在TPUv2上面做的PPA优化,具体包括:

l矩阵乘法单元x2

lCore频率提升30%

l提升30% HBM带宽

l提升2倍的HBM存储容量

l提升4倍的节点能力

l提升30%的link带宽

 

Figure 10 TPUv3 架构更新[1]

 

实际应用部署落地

 

任何处理器,比起benchmark,实际应用产品的经验反馈更能体现现实情况,这些反馈对于设计的迭代也是非常非常重要的一环。Google天生就有大量的应用需求、数据和模型,在这一点上有很多一般的芯片公司所不具备的实战能力。

这里Google主要列出了6大应用,3种网络类型

这里Google主要列出了6大应用,3种网络类型

 

Figure 11 TPU实际应用部署[1]

 

总结

 

Google TPUv1可以说是Google试水芯片的作品,其脉动阵列在5年后看依然是经典设计,很多大学将其放进教科书,TPUv2完美的继承了脉动阵列的设计,而可编程性的VLIW指令集的加入让TPUv2重新define了整个架构,这些转变其实是翻天覆地的变化,从硬件设计、compiler、driver这些角度来说,相当于重新设计了一款芯片。而TPUv3在TPUv2的基础上PPA方便做了进一步提升,可以说TPUv3使其架构真正稳定下来。要知道google在现在依然采购在Nvidia T4等产品,说明Google现在还不能完全替代GPU的解决方案!“your beautiful DSA can fail if best-practice algorithms change”[2],随着算法的快速演进,这让我们更加期待现在测试中的TPUv4中会引入哪些新feature。是引入SIMT?还是引入新的内存模型?新的任务分配调度模型?让我们拭目以待!

 

参考资料

 

[1] Thomas Norrie, Nishant Patil, Doe Hyun Yoon, George Kurian, Sheng Li, James Laudon, Cliff Young, Norman P. Jouppi, and David Patterson, "Google’s Training Chips Revealed: TPUv2 and TPUv3", HOTCHIPS 2020

[2] NORMAN P. JOUPPI, DOE HYUN YOON, GEORGE KURIAN, SHENG LI, NISHANT PATIL, JAMES LAUDON, CLIFF YOUNG, AND DAVID PATTERSON, "A Domain-Specific Supercomputer for Training Deep Neural Networks", communications of THE ACM JULY 2020 VOL.64 NO.7

[3] MLPerf Training v0.7 Results https://mlperf.org/training-results-0-7/

 

近期文章

流体力学与物理导引神经网络

流体力学与物理导引神经网络

现代科学中众多复杂的关键应用依赖着对流体运动的精确预测。然而,由于流体力学方程数值求解的复杂性,基于传统数值方法的高分辨率长周期数值模拟运算量大且难以保证数值稳定性。本文将深入介绍PINN方法在流体计算中的应用并分析其与传统数值方法的差别。
2021-10-25
物理导引神经网络方法分析

物理导引神经网络方法分析

随着GPU能力的提升,支撑深度学习的软硬件生态得到了快速发展。通过深度学习来解决科学计算问题成了一种趋势,其中用深度学习来求解偏微分方程的方法也逐渐兴起。尤其引人注意的是一种称为物理导引神经网络的方法,其为科学计算领域注入了新的活力。
2021-10-18

联系我们

这是描述信息

招贤纳士

公众号二维码
这是描述信息

 版权所有   ©   上海壁仞科技有限公司   |    沪ICP备19047354号