世界第一块AI训练芯片(Google TPUv2)的原理 #AI芯片 #TPU系列原创

叭叭叭，Hello，大家好，想成功先发疯，不顾一切向前冲，我是人到中年冲不动的Z米，虽然天天负能量喊着冲不动啊，但实际上呢比谁都减，加油奥利给，今天呢，我们来到了谷歌TPU第二代的芯片的第二代呢是谷歌第一款训练卡是我们今天呢去看看训练卡跟推尼卡有什么不一样，谷歌推出的第一款训练卡呢，对整个AI芯片领域呢，能有哪些划时代的变化，在整个AI芯片系片里面呢？我们现在来到了国外AI芯片里面，特别是在AI谷歌TPU系列这一个系列里面，真正的去看看谷歌TPU芯片的架构剖析，我们现在呢已经到了第三个小节，TPU第一款训练卡，去看看TPU第二代TPU2有什么不一样哦，朱明老师你好啊，我想问一下。

呃，为什么你要分享这么多节谷歌的TP的课程呢？实际上呢，不管谷歌的训练卡和推理卡也好，他的每一代呢，周米老师都觉得它是一个跨时代的一个意义。首先呢，我们回顾一下谷歌的t bo1里面呢，就首创用了T8的数据格式来做AI的推理，这个呢是在幺三年的时候已经做好规划了，确实非常时的。而在谷歌的TPU2里面呢，用了BF16，现在我们经常用来训练大模型，解决数据不收敛，数据跑的问题，这个呢是在2016年17年就引入的一个架构，但是对于国内的其他厂商呢，可能到了2022年到2023年才开始慢慢引入的这些谷歌在七八年前已经引入的概念，所以我们说谷歌的TPU非常具有划时代的意义，虽然不一定谷歌的TPU它自己很成功，但是它对整个AI芯片的业界呢，带来很多不一样。

的创新点，现在我们整体看看谷歌TPU2的整个产品形态，那这一块呢，就是它的TPU2里面的整一块板卡，这里面呢就镶嵌的四块TPU2的核心训练的芯片，每一块散热卡下面呢，都有一款简单的小小的TPU2的芯片，四块连接在一起呢，形成我们整一块板卡。现在我们看一下具体的谷歌TPV1到V2呢，它最重要的改变呢，就是从推理到训练，大家都知道呢，在训练过程当中呢，会比推理的过程复杂很多，而国内的AI芯片厂商呢，大部分都聚焦于推理，因为推理确实芯片做起来呢，还是相对来说容易很多，整个软件在内会薄很多，而训练呢，它确实会复杂很多很多，我太难了，那训练有哪些难点呢？这里面总米呢总结了五点，那第一点就是训练的并行化费非常的难，因为每个推理的任务都是独立的，我们在整个AI芯片的集群呢。

可以在推理任务里面横向的去扩展，但是训练呢，我们简单的一个小模型和4.50呢，就需要迭代百万次，那这个时候呢，我们要调动跨集群的并行的资源，所以说训练的并行难度呢是很复杂的。第二个呢，就是更复杂的计算，因为呢我们不仅仅去算正向传播或者一个正向的推理A该框架呢，还要做自动微分计算权重参数和输入数据的导数，包括我们的数据格式呢，都要求更高的精度。第三个呢，就是需要更大的内存或者更大的显存呢，因为在推理过程当中呢，我们可以非踌速的去复用我们的内存或者存，但是在训练的过程当中呢，我们会产生反向传播引起的临时变量，还有我们优化器的一些临时的变量和状态，在整体的过程当中呢，会比我们简单的训练，在整体对内存的需求呢，可能会膨胀八到九倍，所以说它需要更大的内存，那第四个呢，就是更加具有可编程性，因为训练的算法和模型呢，在快速。

度的眼镜会比推理呢更加复杂，更加多变，所以这一块呢，大家AI芯片要演进的非常的快，那最后一个呢，就是高精度的数据格式，在推理的过程当中，我们之前都说了可以他用INT8的整形做推理，但是在训练过程当中呢，我们的梯度性是是非常非常的小的，我们要捕捉非常微小的一些数据的变化，于是呢就引入了BF16啊PF16啊P12等缓和进度的计算，所以说训练会比推理难很多很多，而谷歌TPU的第二代呢，就做的非常的巧妙，它将TPU第一代的架构呢做了一些稍微的改动，就变成了谷歌TPU第二代用来做训练的，而针对T第二代的架构的修改呢，就是我们训练跟推理的差别，我们今天呢，就会重点的去看看训练跟推理的差别，做了哪些修改点。我们现在。

正式交给命运的齿轮，看看谷歌TPU第二代跟谷歌TPU地带有什么具体的差异，首先呢，最大的差异中，你觉得是呢，是引入了BF16，因为现在训练大模型真的是BF16非常的重要，没有了BF16做大模型的预训练呢，很容易跑飞，不过呢，这是体外话，我们真正的来到了谷歌第二代的一些真正的改变。第一个呢，就是谷歌第一代呢，其实有两个存储区啊，一个叫做accumul，一个是accation storage，也就是右边这个图的这两个黄色方框，在推理的场景当中呢，设置了一些专用的存储模块呢，实际上对我们的计算呢，会对我们的存储呢，是非常有帮助的，因为它比较多in specific嘛，专门针对我们的领域来的，在训练过程当中呢，为了提升可编程性呢，于是呢，就把刚才的actation排line呢跟ACCUMULATOR2个单元先换换位置，换完位置之后呢，再把accumulator这个模块呢跟actation story这两个模块呢和。

并起来，变成更加类似于传统EECPU里面的Le cash，于是呢就变成一个VE memory，我们的可编程的缓存，有了这个内容之后呢，下个我们针对activation pipeline呢进行一个改变，Activation pipeline呢，一开始在谷歌P地带呢，是专门针对激活的函数进行处理的，也就是我们卷积之后有个BA normalization，然后在些激活，这个激活特殊的ARU的计算，这个特殊的alu计算呢，已经没有办法去满足我们的训练的场景了，于是呢，把它变成一个vect的unit，可以专门针对类似于激活还有soft ma这一类vect的计算呢，专门在进行计算和处理。接在呢，我们来到了第三个模块，也就是mxu，第三个重要的改变呢，就是mxu，原来MU呢是跟VE memory相连接的，现在呢，用了新的方式之后，谷歌TV2呢把MSU跟VE呢进行一个连接。

所有的数据的出口，所有的计算呢，都是通过vect unit进行分发，这个时候呢，MU脉冲阵列单元的计算，特别是针对矩阵的计算呢，就成为了向量单元vu的一个斜处理器，这种结构呢，对编译器和编程人员来说呢更友好，我们只关心VE unit的控制，然后呢，通过VE unit来去下发到不同的计算单元上面。到了第四个改变点呢，就是DDR3一开始呢，D代的时候使用DDR3的内存，是用把所有的推理的场景当中呢，每一次需要用到的一些权重的直接加载进来去计算，但是呢，在训练的过程当中呢，我们会在中间产生非常大量的变量和权重，那这个时候呢，DDR3的回血速度是非常慢的，于是呢，在训练的过程当中呢，我们把DDR3跟we memory放在一起，然后再做一个改变，把DDR3呢换成HBM，使得我们。

整体的读写的速度快了20倍，哎呀，这一期视频又不小心讲多了好多的废话，我们现在呢，才来到了T2的具体的核心，谷歌T2的核心呢，其实有非常多啊，有呢有呢有M呢，还有一个transport，一共有四个unit单元，看到有这么多个单元呢，不知道国内的AI芯片厂商的设计者，还有国内AI芯片的从业者，或者国内AI系统的从业子呢，有没有一些似同相似的感觉，现在是不是很多练卡，特别是D都是采用这种方式呢？这也是你觉得谷歌的TPU的很多的计算的方式或者他的思想呢，都被大家借鉴过来了，百因必有果，你的报应就是我，那现在呢，我们看一下它的指令呢，具体从哪里获取过来的，首先呢，指令不在是通过我们的CPU去获取的。而是通过和发。

是单元处于出超强的指令集，并且将具体的指令呢发送到VP啊，也就是vect process unit，然后再给其他企处理器呢进行控制执行的，这里面所有最重要的应口呢，就是我们的VEUN，具体VEUN展开呢就是左边的红色的这个大框框，整体的VE呢有一个controller，通过control呢来去下发我们的具体的指令给到不同的鞋处理器上面，VE自己呢也可以做一些简单的a Lu的运算里呢，自己是一个sin的结构，里面呢有非常多个a Lu，而Lu里面呢又分为非常多个a Lu的和具体的参数规格呢，大家也可以慢慢的去细品细品。接着我们来到了第二个内容，第二个计算单元就是unit unit的用的比较少，更多的是计算一些简单的标量的运算，那标量的运算呢，其实在AI的运算过程当中呢，它不。

说非常重要，但是呢，没有它肯定是不行的，现在呢，我们来到之前在谷歌TP第一代里面很重要的一核心就是MSU脉动阵列，每个脉动增列呢，就变成128乘以EMAIL8，对比起V1版本的256乘256呢，少了一半哦，数据的计算格式呢，支持混合精度，支持BF16跟FP322种格式，最后还有一个另点呢，就是对我们的矩阵的特殊的计算进行提供的，例如矩阵的转制的矩阵的规约置换矩阵，这些呢都是在我们训练的过程当中，特别是反向的传播过程当中呢，会经常遇到的，诶慢子慢子周明老师你好啊，我有个问题哦，刚才你不是说到V2呢变成128乘以128吗？为什么谷歌不再用256乘256，而是用128乘以二八小了一。

呢，小新你问的这个问题呢，非常的及时哦，我们看看右边的这个图案，如果螨虫帧粒的大小呢，是256乘256呢，它整体的计算的峰值可以做的非常非常的高，但是呢，利用率可能会变得非常的低，因为不是所有的卷积都能够吃得下这么大的一个矩阵的计算的，而放到64乘64呢，虽然发现呢利用率多了，但是呢性能却下降了，于是呢，谷歌就采用128乘以二八权衡芯片的电路面积，还有计算的峰值和算力的利用率，专门的采用了这个核心的位宽，终于一口气讲完这么多，来到了最后一个内容，就是谷歌TPU的内存和互联的方式，不知道大家觉得十分钟比较合适，还是像明老师这样呢，把每一节课的内容都详细的展开，或者尽可能的详细的展开呢，果家。

我更好的意见呢，欢迎弹幕，那这里面呢，我们看看谷歌T里面的一个具体的互联的方式，在H和memory里面呢，增加了一个connected这么这么一个模块，用来方便TPU去组成一个pod的超级计算机的形态的句子，模块呢，就是右边的这个框框，Interconnected water，在训练的场景当中呢，因为需要有大量的数据不断的去未经我们的神经网络不断轮训的迭代循环去训练，所以单个核心呢是不够用的，我们需要用到并行的计算，我们需要用到集群的能力，值得注意的就是TPU核心跟TPU核心之间的互联带宽，而不是卡跟卡之间，也不是机跟机之间的一个互联带宽，现在呢，我们看一下核之间的互联，核金的互联，最核心的模块呢，就是我们刚才讲到的右边的这个红色的conct vote，这个模块能够实现的2D的面连接，也就是2D的to，我们将会在下节内容里面呢，在谷歌。

UV3里面详细的去展开这个to us的架构，还有它的互联的整个集群的内容，刚才我们看到的所有的内容呢，都是左边的这个模块，实际上呢，谷歌TPU的一个芯片里面呢，有两个核心，也就是有一个CALL1和一个二组成一块小小的芯片，整个芯片里面呢，是通过interconineed water呢去把两个核心连接在一起，最后呢，变成简单的架构图呢，就像这样所示，那这个图里面呢，线越宽代表它的带宽越高，所以说t PU Co跟HBN之间的带宽是最高的，TPU扩跟TPU扩之间的带宽呢是其次的，而TPU整一个新力在利用还新力之间的带宽呢，会低成，这也是英伟达的GPU呢，去建立多级的缓存，通过减少英伟达GPU的食延呢，从而提升GPU的计算的吞吐。很不容易，很不容易，我们终于来到了TPU。

V2最后的总结和思考啦，看到这里面的人呢，可能就并不多了，里面我们刚才呢有两个核心，上面是一个核心，下面是一个核心，中间呢通过interconineed vote呢把它们连接起来，这里面呢，最大电路呢还是MXU2个MSU，还有里面的memory，我们尝试做一个总结TV2呢，主要是提供了两个两个核心，具体呢，将激活函数，也就是我们的改称可编程性更高的向量单元VE unit，那下面几个都是VE unit带来的好处，使用了VE unit呢代替了之前的双缓存，而MSU呢，作为向量单元的些处理器，提升我们的可编程性，方便我们的AI编译器呢做一个整体的调度和控制，最后还增加了SCUN呢，还有还有其他UN，那在内存的互联方面呢，就使用了HBN代替了DDR3，而且还使得HBM和VM之间呢增加了一个互联的模块，使得tu跟t puu间呢提升整个。

互联的带宽，那整体的区别呢，我们看看架构图啊，原来control的四级流水呢，现在变成右边的这个图，更多的是通过VEUN呢进行一个控制调度、分发，这个整个就是我们TPV跟VR最大的区别，也是从推理到训练最大的改变。今天的内容呢，就这么多，谢谢各位喽，拜了个拜，卷的不行了，卷的不行了，AI系统的全站知识都会分享在这里，欢迎打开和吐槽，给我一键三连，给我一键三连啊，谢谢各位，拜了个拜。