AI硬件05-手把手拆解训练芯片：五大核心模块全解析，训练芯片五大模块 vs 你家的 CPU：架构差异在哪

本文为《AI 硬件体系深度调研》系列第 05 篇。上一篇我们给芯片篇画了总纲AI 芯片是算力体系的最小核心单元分训练三高、推理三低两大阵营。本篇开始往里钻把一颗训练芯片拆开看它内部五大核心模块是如何分工、如何配合的。黄金 100 字开头你是否好奇一颗训练芯片里到底塞了些什么才让它这么能算芯片手册里的一堆模块名词看得人头大却不知彼此怎么配合。本文按数据流顺序把训练芯片五大核心模块逐个拆开讲明白。一、张量计算核心阵列算力的发动机训练芯片的心脏是一块由数千个专用计算单元排成的二维阵列。它的唯一任务就是做矩阵乘加Multiply-Accumulate简称 MAC。这一块是整颗芯片里最值钱、也最烧电的部分。1.1 先搞懂张量到底在算什么张量听起来高大上其实你可以把它理解成多维数组。一维是向量二维是矩阵三维以上统称张量。深度神经网络的每一层本质上都是张量之间的运算——其中最核心、最耗时的就是矩阵乘法。为什么矩阵乘法这么重要我们看一个最简单的全连接层输入是一个向量 x权重是一个矩阵 W输出就是 W 乘以 x。一次矩阵乘法就是成千上万次乘再加。模型越大矩阵越大乘加次数就越多。千亿参数的大模型一次前向传播要做的乘加运算是天文数字。具体到硬件层面一次乘再加MAC就是先做一次乘法、再做一次加法。张量核心里的每个计算单元每个时钟周期就能完成一次这样的 MAC。几千个单元同时开工一个周期就是几千次 MAC这就是并行算力的来源。 一句话AI 计算 90% 以上的时间都在做同一件事——矩阵乘加。所以谁能把矩阵乘加做得最快谁就是算力王者。1.2 为什么是二维阵列而不是一堆核心你可能会想CPU 里也有乘法单元为什么训练芯片要专门排一个二维阵列答案在并行度和数据复用。矩阵乘法有一个黄金特性同一个数据可以被反复使用。算 A×B 时A 的每一行要和 B 的每一列相乘一个数据会被用到几十次、几百次。二维阵列的妙处就是让数据横向流、纵向流权重从一侧流入激活值从另一侧流入每个交叉点上的计算单元同时做乘加。这样几千个单元可以同时工作而不是像 CPU 那样一个接一个地算。这里的关键是数据复用一个权重数据进入阵列后会被同一行的所有单元同时使用一个激活值进入阵列后会被同一列的所有单元同时使用。一次搬运多个单元共享数据搬运的开销就被摊薄了。反观 CPU同样的数据要反复从寄存器里取搬运开销高得多。打个比方CPU 是一个老师批改作业一本本慢慢批张量阵列是几百个老师同时改每人负责一道题。速度差就在这里。这个比喻还能再延伸一层CPU 的老师不但改得慢还得自己翻作业本从内存取数张量阵列的老师们则排成整齐的方阵作业本自动流转到每个人面前数据流驱动。一边是人找活一边是活找人效率自然天差地别。1.3 卷积也是矩阵乘加可能有读者会问卷积神经网络CNN里用的是卷积不是矩阵乘法啊其实卷积可以通过数据重排im2col转换成矩阵乘法。把滑动的卷积窗口展开成矩阵的行把卷积核展开成矩阵的列卷积运算就变成了一个规整的矩阵乘加。所以无论 CNN 还是 Transformer底层通通收敛到矩阵乘加——这也是为什么训练芯片只要把矩阵乘加做到极致就能通吃各种模型。im2col 的本质是把滑动窗口这种带位置变化的运算摊平成一个没有位置概念的矩阵乘法。代价是数据会有冗余相邻窗口有重叠但换来的是运算模式的统一——芯片只需要做好矩阵乘加这一件事就能覆盖所有层的计算。⚠️ 避坑警告不要以为张量核心和GPU 的 CUDA Core是一回事。张量核心是专门做矩阵乘加的专用单元一个张量核心一个时钟周期能完成一个小矩阵的乘加而普通 CUDA Core 是通用标量单元。两者在芯片里是并存的各司其职。1.4 阵列规模决定算力上限张量阵列里计算单元的数量直接决定芯片的峰值算力。单元越多、频率越高算力越强。旗舰训练芯片的张量阵列通常由数千个单元组成每个单元每周期完成一次乘加再乘上频率和阵列规模就得到了我们熟悉的2000 TFLOPS级别的恐怖数字。但这个发动机有个前提它必须被喂饱数据。阵列再快如果数据送不上来单元就只能空转。于是芯片设计者要解决的下一个问题就是——怎么把数据以最快的速度送到阵列嘴边。这就是下一节要讲的片上内存。这里引入一个关键概念算力与带宽的比例匹配。阵列算力翻一倍数据供给能力也得翻一倍否则新增的算力就是摆设。训练芯片设计里算力/带宽比是一个必须精打细算的指标失衡了就会一头是瓶颈。二、三级片上内存把数据放在手边张量阵列算得飞快但它有一个致命需求每个时钟周期都要有新数据送进来。如果数据要跑到片外的显存里去取来回一趟的延迟够阵列空转几百个周期。所以芯片里必须有贴近阵列的、极高速的缓存。2.1 三级缓存的分工训练芯片的片上内存通常分成三级每一级都有明确的职责层级名称作用特点1激活值缓存Activation Buffer存放当前层的输入/输出激活值频繁读写带宽极高2权重缓存Weight Buffer存放当前层要用的权重只读为主反复使用3SRAM 缓存通用高速缓存缓冲中间结果灵活调度容量稍大这三级的共同特点容量都很小通常几十 KB 到几 MB但带宽极高且离计算阵列极近。给一个数量级感受三级片上缓存加起来可能只有几 MB而片外 HBM 动辄80GB差了四个数量级。但片上缓存的访问延迟只有纳秒级带宽却是 HBM 的几倍到十几倍。容量和速度永远是反比芯片设计就是在中间找平衡。2.2 为什么必须分级你可能会问为什么不直接把所有数据都塞进最快的缓存答案是成本与面积。最快的 SRAM单位面积的成本和功耗都远高于片外显存。芯片的硅面积就那么大全堆 SRAM 就放不下计算单元了。所以只能分级——把最常用、最紧急的数据放在最快的地方把大块数据放在稍慢但更大的地方。打个比方计算阵列是灶台三级缓存是灶台边的调料架片外 HBM 是十米外的冷库。炒菜时常用的油盐酱醋放调料架缓存大量的菜放冷库HBM需要时再往灶台搬。这样既不用把整个冷库搬到灶台边面积不允许又能保证炒菜不停火。这个分级策略其实是计算机体系结构里一条经典原则的极致应用——数据访问越频繁、越紧急就放得离计算越近。训练芯片把它做到了三层甚至更多层本质都是同一个思路让最快的存储伺候最忙的计算。2.3 激活值和权重为什么要分开缓存这里有个容易被忽略的细节激活值和权重的读写模式完全不同。权重在一次推理或训练的一层里基本是只读、反复用的——同一批权重会被成千上万个输入数据反复乘。所以权重缓存可以设计成读多写少甚至支持广播。激活值则是频繁读、频繁写的——每一层的输出就是下一层的输入数据不停地产生、不停地被消费。所以激活值缓存要支持高频读写。正因为两者模式不同把它们分成两个独立的缓存Weight Buffer 和 Activation Buffer可以各自优化避免互相干扰。这是芯片设计师从数据流里抠出来的效率。举一个广播的例子同一批权重要同时喂给阵列里的几千个单元如果权重缓存放不下、要去 HBM 反复取那带宽早就被拖垮了。把权重独立缓存并支持广播一次取出来就能同时送到所有单元这就是分开缓存的实际价值。 一句话三级片上缓存本质是用面积换速度——把离灶台最近的黄金位置留给最常用的数据。2.4 缓存虽小却是性能关键别看三级缓存容量加起来可能才几 MB它直接决定了张量阵列能不能持续满负荷运转。缓存命中率高阵列就满血缓存命中率低阵列就空转等数据。这也是为什么下一篇要讲的脉动阵列如此重要——脉动阵列的设计目标之一就是让数据在阵列里流动起来最大程度复用缓存里的数据减少与片外的交互。缓存和阵列的配合是训练芯片性能的生死线。换句话说训练芯片真正的瓶颈往往不在算而在喂——谁能把数据喂得更快、更准谁就能把峰值算力兑现得更高。这个判断会贯穿后面每一篇硬件分析。三、HBM 控制器打通片外显存片上缓存虽快但容量太小装不下大模型的权重。真正的大块数据存在片外的HBM高带宽内存里。而连接片上缓存和片外 HBM的就是HBM 控制器。3.1 HBM 是什么HBM全称 High Bandwidth Memory是一种堆叠式的高带宽内存。它把多层 DRAM 芯片像三明治一样垂直堆叠起来通过极宽的数据通路与主芯片相连。普通显卡用的 GDDR 内存带宽一般在几百 GB/s而训练芯片用的 HBM带宽动辄几个 TB/s是前者的好几倍。上一篇我们讲过训练芯片三高之一就是超高内存带宽这个带宽的实现靠的就是 HBM。为什么 HBM 能做到这么高带宽核心是堆叠 超宽总线。普通 GDDR 是平铺在电路板上的走线长、位宽有限HBM 则把多层 DRAM 垂直堆叠紧贴着主芯片封装用上千条短走线直连把距离和位宽这两个带宽的天花板同时捅破。3.2 HBM 控制器的职责有了 HBM 还不够还得有人指挥数据进出。HBM 控制器就是芯片里专门负责调度显存与片上缓存之间数据交换的模块。它的工作包括发起读写根据计算需求把权重、激活值从 HBM 搬到片上缓存调度顺序决定先搬谁、后搬谁尽量让计算不中断批量传输把零散的小请求合并成大的连续传输提高效率。这四点里“批量传输尤其重要。HBM 对连续大块读写效率最高对零散小请求效率极低。控制器把无数小请求合并成大传输才能把 HBM 的带宽真正吃满。这也是为什么数据排布”layout在训练里如此关键——排布得好访问就是连续的排布得差访问就是碎片化的。打个比方HBM 是大仓库片上缓存是车间的小货架HBM 控制器就是那个仓库管理员。他的本事在于——提前预判车间接下来要什么料在车间用完之前就把下一批料从仓库送到货架上让车间永不停工。3.3 预取是隐藏延迟的关键HBM 再快从片外到片上还是有一段物理距离单次访问的延迟依然存在。HBM 控制器的核心技巧就是预取Prefetch在计算还没用完当前数据之前就提前把下一步要用的数据搬进缓存。这样当张量阵列真正需要新数据时数据已经在缓存里等着了片外访问的延迟就被隐藏在了计算时间里。这就像餐厅里服务员在你吃完这道菜之前就已经把下一道菜端到了传菜口——你几乎感觉不到等待。反过来如果预取判断错了——提前搬来的数据用不上真正要用的数据反而没搬来——阵列就会断粮只能空转等待。预取的命中率直接决定训练芯片实际能发挥出几成算力。这也是为什么大模型训练时框架层的数据加载策略和芯片层的预取策略同样重要。 效率技巧评估训练芯片时不要只盯着峰值算力还要看内存带宽和预取能力。带宽不够算力再高也是饿着肚子跑。四、多芯片互联模块为多卡训练而生上一篇我们提到训练芯片三高的第三高是多芯片高速协同。这一节我们就看这颗芯片里专门为多卡协同而生的互联模块。4.1 为什么单卡不够大模型的参数量早就超过了一颗芯片能装下的范围。以千亿参数模型为例即使用 BF16 精度权重就有200GB 以上远超单卡显存通常 80GB。所以训练必须把模型切分到几十张、几百张甚至几千张卡上让它们协同工作。切分的方式有很多种按层切流水并行、按张量切张量并行、按数据切数据并行。但无论哪种切法卡与卡之间都必须通信——因为模型是一个整体切开的每一块都依赖其他块的结果。协同训练时卡与卡之间要频繁交换数据——梯度要同步、激活值要传递、权重要广播。如果卡与卡之间的通信慢整批训练就会被最慢的那张卡拖住。4.2 互联模块的作用多芯片互联模块就是芯片上专门负责芯片间直接数据交互的接口。它的作用可以概括为三点点对点直连让两颗芯片之间可以直接传输数据不必绕道 CPU 或系统总线高速带宽卡间互联带宽动辄几百 GB/s 甚至上 TB/s尽量追平片内通信低延迟让跨卡的数据交换几乎感觉不到跨机的延迟。在大规模集群里互联还涉及拓扑结构——几百上千张卡怎么连是全互联、还是分组、还是环状。不同的拓扑决定了通信效率和成本之间的权衡。互联模块提供的是单芯片级的直连能力而拓扑则是系统层面怎么把这种能力组织起来。打个比方多卡训练是一群人一起算一道大题互联模块就是他们之间的对讲机。对讲机越清晰、越同步协作效率越高对讲机一卡壳大家就得停下来等。4.3 芯片间直接交互有多重要这里的关键词是直接。如果芯片之间不能直接通信就必须通过 CPU、PCIe、网络层层中转每一层中转都是一次延迟和带宽损耗。训练芯片把互联模块直接集成在芯片内部让芯片 A 的数据可以点对点直达芯片 B中间不绕任何弯路。这就是为什么大规模训练集群里卡间通信能做到接近片内通信的速度——因为互联模块把卡与卡的距离压缩到了芯片与芯片的距离。不过要诚实地说卡间通信再快也快不过片内。所以大规模训练的理想状态是把通信和计算重叠起来一边算、一边传让通信的时间藏在计算里。这需要互联模块、调度器和上层框架三方配合是训练系统优化的核心战场之一。 一句话互联模块是训练芯片从单打独斗走向千卡军团的桥梁。没有它大模型训练无从谈起。五、全局控制单元芯片的大脑前面四个模块分别负责算张量阵列、“缓存”三级片上内存、“取数”HBM 控制器、“协同”互联模块。但谁来决定什么时候算什么、数据从哪里来、结果往哪里去答案是全局控制单元——芯片的大脑。5.1 解析框架指令训练框架如 PyTorch、TensorFlow会下发一条条计算指令比如执行这个矩阵乘法“把结果写回缓存”。全局控制单元负责解析这些指令把它们翻译成芯片内部各个模块能理解的信号。现代训练框架的指令往往已经是高度抽象的计算图节点而不是最底层的硬件信号。从图到电路之间隔着好几层翻译先拆成算子再拆成微操作最后才是控制信号。全局控制单元就站在最底层负责把这最后一步翻译做对、做快。这一步看似简单实则关键框架的指令是高层语言芯片的模块是底层硬件控制单元就是中间的翻译官。翻译得准不准、快不快直接影响整颗芯片的利用率。5.2 拆分调度到各核心一条矩阵乘法指令往往要拆分成无数个小任务分配到阵列里成千上万个计算单元上。全局控制单元负责这个拆分和调度——哪个单元算矩阵的哪一块、哪些数据先送、哪些后送都由它统一安排。打个比方全局控制单元是工厂的调度中心张量阵列是车间里的机器控制单元负责把订单拆成工序、派发给每台机器。调度得好机器满负荷调度得差机器空转等料。调度的粒度也很讲究拆得太粗单元之间负载不均快的等慢的拆得太细调度本身的开销又太大。一个好的调度器要在负载均衡和调度开销之间找到最优解这是芯片设计里极吃经验的环节。5.3 同步时序训练芯片是高度并行的——几千个计算单元同时干活数据在缓存、显存、阵列之间不停流动。如果各模块的节奏对不上就会出现算得太快、数据没到或数据到了、还没轮到算的错乱。全局控制单元的第三个职责就是同步各模块的时序让取数—缓存—计算—写回这条流水线严丝合缝地运转。这也是为什么训练芯片的运行非常依赖精确的时钟和调度——一个时序上的失误就可能让整颗芯片的效率腰斩。可以想象一条流水线取数、缓存、计算、写回四个工位同时开工。任何一个工位节奏快了或慢了都会让后面的工位要么等、要么堵。全局控制单元就像流水线的节拍器让每个工位踩着同一个节拍整体吞吐才能最大化。 一句话总结张量阵列是肌肉三级缓存是近处的补给HBM 控制器是远端的后勤互联模块是对外的协作全局控制单元是大脑。五大模块缺一不可共同组成一台精密的算力机器。配图图 1训练芯片五大模块结构图flowchart LR CTRL[全局控制单元br/解析指令/调度/同步] -- ARRAY[张量计算核心阵列br/数千专用单元·矩阵乘加] CTRL -- CACHE[三级片上内存br/SRAM / Weight Buffer / Activation Buffer] CACHE -- ARRAY HBMCTRL[HBM 控制器] -- CACHE HBMCTRL -- HBM[片外 HBM 高带宽内存] LINK[多芯片互联模块] -- ARRAY LINK -- OTHER[其他训练芯片]五大模块各司其职控制单元是大脑阵列是肌肉缓存是近补给HBM 是远后勤互联是对外协作。图 2数据在三级片上内存与 HBM 之间的流动示意图sequenceDiagram participant HBM as 片外HBM participant CTRL as HBM控制器 participant CACHE as 三级片上缓存 participant ARRAY as 张量计算阵列 CTRL-HBM: 预取权重/激活值 HBM--CTRL: 返回数据 CTRL-CACHE: 写入缓存 CACHE-ARRAY: 按需供给数据 ARRAY-CACHE: 写回中间结果 CACHE-CTRL: 溢出回写 CTRL-HBM: 写回显存数据从 HBM 经控制器预取到缓存再以极低延迟供给阵列算完的中间结果写回缓存最终溢出回写 HBM。预取的目标就是让阵列永远有数据可算。看到这里你会发现训练芯片的五大模块其实都围绕同一个目标转让张量阵列一刻不停地做矩阵乘加。缓存是为了喂饱它HBM 控制器是为了搬得动它互联是为了放大它控制单元是为了调度它。理解了这条主线训练芯片就不再是一堆神秘名词的堆砌。写在最后一颗训练芯片拆开来看就是这五大模块张量计算核心阵列负责算三级片上内存负责喂HBM 控制器负责搬多芯片互联负责连全局控制单元负责指挥。它们环环相扣把海量的矩阵乘法变成一场有条不紊的流水线。【思考题】三级片上缓存容量都很小却要喂饱庞大的阵列调度策略出错会怎样欢迎评论。提示想想缓存未命中时阵列会不会空转以及预取失败对整体吞吐的影响。【系列文章预告】下一篇讲训练芯片的灵魂——脉动阵列为什么它能把数据搬运开销降一个数量级。标签训练芯片架构、张量核心、片上内存、HBM、多芯片互联、AI芯片、硬件架构张量计算核心阵列数千专用单元二维阵列专做矩阵乘加/卷积、分层高带宽片上内存SRAM 缓存 权重缓存 Weight Buffer 激活值缓存 Activation Buffer 三级、HBM 控制器对接外部 HBM调度显存与片上缓存大规模交互、多芯片互联模块支持大规模多卡协同芯片间直接数据交互、全局控制单元解析框架指令拆分调度到各核心同步时序。