1. 项目概述为什么设备端DNN训练需要“特制”的硬件加速器深度神经网络DNN如今已经渗透到我们生活的方方面面从手机相册的人脸识别到智能音箱的语音交互背后都是训练好的模型在默默工作。但不知道你有没有想过为什么这些智能功能总是需要联网更新或者只能在云端进行“学习”核心原因在于传统的DNN训练是一个极其“烧算力”和“吃内存”的过程动辄需要成百上千张GPU卡在数据中心跑上好几天这显然不是我们手边的手机、智能摄像头或可穿戴设备能承受的。然而真正的智能应该是个性化和实时进化的。想象一下你的手机摄影助手能根据你独特的审美偏好比如特别喜欢某种色调进行微调或者一个家庭陪伴机器人能记住每个家庭成员的习惯并做出不同反应。这些场景都需要模型在设备端、利用本地产生的私有数据进行持续学习即“设备端训练”而不是把所有隐私数据都上传到云端。这就引出了我们今天要深入探讨的核心问题如何在算力有限、电池供电的边缘设备上高效地完成DNN训练这个艰巨任务答案就是为训练“特制”的硬件加速器。你可能会说市面上不是已经有大量AI推理芯片了吗没错但推理和训练是两码事。打个比方推理就像开一辆已经调校好的赛车在固定赛道上跑圈而训练则是工程师在车间里反复拆解、测试、调整这辆赛车的每一个部件。后者对“车间”计算硬件的要求要复杂和苛刻得多。这篇分享我就结合自己多年在AI芯片设计一线的经验为你拆解设备端DNN训练加速器的核心技术挑战与设计思路这不仅仅是学术综述更是工程实践中必须直面的硬骨头。2. 训练与推理的本质差异硬件设计挑战的根源要设计高效的训练加速器首先必须透彻理解训练过程本身为何如此“难搞”。DNN训练本质是一个迭代优化过程每次迭代包含三个核心步骤前向传播输入数据经过网络各层得到预测输出。这步和推理类似。反向传播计算预测输出与真实标签之间的误差损失并将这个误差从网络输出层逐层反向传递至输入层。这一步需要计算每一层参数的梯度。权重更新根据反向传播计算出的梯度使用优化器如SGD、Adam来更新网络中的权重参数以期在下一次迭代中降低误差。这三个步骤循环往复直到模型性能收敛。正是这个循环过程带来了与推理截然不同的三大硬件设计挑战。2.1 数据流冲突一个数据三种“吃法”这是训练硬件设计遇到的第一个拦路虎。在推理时数据权重和输入特征的流动方向是固定的、单一的。但在训练中同一份数据在不同的步骤中被访问和使用的模式完全不同。以最常用的卷积层为例前向传播需要按输入通道顺序读取权重和输入特征图进行乘累加运算。反向传播需要读取转置后的权重输出通道和输入通道维度互换与上一层传来的梯度来计算本层的输入梯度。权重更新需要将本层的输入特征图与输出梯度进行外积运算来计算权重的梯度。问题来了为了最大化数据复用、提升计算效率硬件加速器通常采用某种固定的数据存储布局。例如为了优化前向传播我们可能把权重张量在内存中按照[输出通道 输入通道 高度 宽度]的顺序连续存储。但到了反向传播阶段我们需要按输入通道维度去连续访问权重而此时的存储布局导致访问变成了跨步访问——就像你想连续读一本书的每一页第一行但书却被装订成每次翻页都跳转到随机行效率极低。这种跨步访问会严重降低内存的有效带宽增加访存延迟和能耗。硬件设计必须在支持灵活多变的数据流和保持高硬件利用率之间找到平衡。实操心得在早期架构定义阶段必须用真实的主流网络如ResNet、MobileNet对三个步骤进行详细的访存模式分析。不能只看理论计算量FLOPS访存瓶颈往往是能效的真正杀手。我们曾经在一个早期设计中忽略了这一点导致反向传播阶段的实际性能只有前向传播的30%。2.2 内存墙中间特征的“存储海啸”推理时每一层的输出特征图在计算完成后就可以丢弃因为后续不再需要。因此内存中只需要同时保留当前层的输入和输出特征即可峰值内存占用相对可控。训练时情况截然不同。为了进行反向传播和权重更新前向传播过程中每一层产生的中间特征图都必须被保存下来直到对应的反向传播计算完成。对于一个批大小为N的训练过程这相当于需要同时存储网络所有层在N个样本上的中间结果。对于深度网络这会产生巨大的内存占用我们称之为“激活内存”问题。此外现代网络普遍使用的批归一化层在训练时也需要计算并保存整个批次的均值和方差进一步增加了内存开销。我们的实测数据显示在训练ResNet-50时中间特征的内存访问量占总访存量的70%以上远超权重参数本身。因此训练加速器的设计必须将中间特征的压缩与高效管理置于最高优先级而不是像推理芯片那样主要关注权重压缩。2.3 计算特性动态范围与稀疏性的挑战宽动态范围的梯度在反向传播中梯度值可能分布在一个非常广的范围内从极小的值到很大的值。推理中常用的低比特定点数格式如INT8因其表示范围有限很容易导致梯度下溢变成0或溢出变成最大值从而使训练失败或精度大幅下降。因此训练硬件通常需要支持浮点数格式如FP16、BF16来保证动态范围但浮点运算单元的能效比远低于定点单元。稀疏性的差异推理可以利用权重稀疏性通过剪枝获得大量0值来跳过计算节省功耗。但在训练中权重是持续更新的无法保持静态的稀疏模式。虽然前向传播中ReLU激活函数会产生激活稀疏性很多0值但在计算量巨大的反向传播步骤中输入的梯度往往不是稀疏的这使得推理芯片中高效的“零值跳过”逻辑在训练中收益有限。因此训练硬件的计算单元设计需要在支持必要数值精度的前提下探索新的稀疏性利用方式和定制化的低位宽格式以逼近推理硬件的能效水平。3. 核心优化技术一驾驭多变的数据流面对三种步骤的不同数据流硬件架构师主要有两种武器阵列重排和可配置数据通路。3.1 阵列重排以空间换时间的预处理思路很直接既然数据在内存中的布局不适合当前计算步骤那我就在数据加载到计算核心之前先把它在片上缓存或寄存器中重新排列好。具体实现通常由DMA控制器或专用置换网络完成。例如在开始反向传播计算前将一块权重数据从“输出通道优先”的布局转置为“输入通道优先”的布局然后再送入处理单元阵列进行计算。优点实现相对简单对计算核心的微架构改动小。软件可控灵活性高。缺点与坑点额外的存储与搬运开销重排需要额外的缓冲区并且数据被搬运了两次加载到缓冲区重排后再送入PE增加了功耗和延迟。不适合大数据块对于很大的张量重排操作本身的延迟可能无法被计算时间掩盖成为性能瓶颈。与压缩技术冲突如果数据为了节省带宽已经过压缩如稀疏编码那么重排前需要先解压重排后再压缩开销巨大往往得不偿失。设计权衡在我们的芯片设计中我们只对权重张量的小数据块进行重排。因为权重在训练中相对稳定且块大小可控。而对于巨大的、已压缩的中间特征张量我们则放弃了重排方案转而采用下面这种更根本的方法。3.2 可配置数据通路让硬件“变形”以适应数据这是一种更优雅但也更复杂的解决方案。其核心思想是保持数据在内存中的存储布局不变但让处理单元阵列内部的数据流网络能够动态重构以适应不同步骤的计算模式。架构举例一个典型的可配置数据通路PE阵列可能包含可切换的累加路径在前向传播时PE将部分和沿“输出通道”维度累加在反向传播时则切换到沿“输入通道”维度累加。灵活的网络互连通过可配置的片上网络让数据能从不同的方向如行、列、输出站、权重站流入和流出PE实现输出站、权重站、输入站等多种数据流模式。异构核心甚至可以为前向、反向、权重更新分别设计微架构略有侧重的专用计算核心通过任务调度器来分配任务。优点从根本上避免了数据重排的搬运开销和延迟。能始终保持较高的处理单元利用率。与数据压缩技术兼容性更好。挑战硬件设计复杂度高互连网络和控制器设计是关键。可能增加芯片面积和布线拥塞。需要编译器或驱动程序的深度配合以正确配置每个层、每个步骤的数据流。一个真实案例我们参考过一篇论文中的设计它采用了一种“2D环面互连”结构。PE之间通过可配置的链路连接在运行时可以根据指令将阵列重构成最适合当前计算步骤的拓扑比如为矩阵乘优化或为卷积优化从而在三种训练步骤中都达到了超过85%的硬件利用率。4. 核心优化技术二向内存墙发起总攻——压缩策略如前所述中间特征是内存消耗的大户。幸运的是这些特征数据具有一些可被压缩的特性稀疏性ReLU等激活函数会产生大量零值。局部性零值或相似值往往成片出现。容错性对数值精度有一定容忍度适当的量化不会导致训练发散。针对这些特性压缩技术主要分为三类4.1 稀疏压缩直接干掉零值这是最直观的压缩方法。不存储零值只存储非零值及其位置信息。常用格式零值压缩用一个与原始张量同尺寸的位掩码来标记每个元素是否为零。简单但索引开销大。游程编码存储非零值以及连续零值的个数。对于零值聚集的情况压缩率高。压缩稀疏行/列经典稀疏矩阵格式存储非零值及行/列索引。适合不规则稀疏模式。硬件支持关键稀疏压缩的收益必须在硬件层面实现才能转化为能效。这需要稀疏感知的加载单元能直接读取压缩格式的数据流。计算跳过逻辑PE阵列需要能够根据索引信息跳过与零值相关的乘加运算。这涉及到不规则工作负载的调度和负载均衡问题设计不好反而会降低利用率。4.2 基于概率的压缩高频值用短码分析张量中数值的分布对出现概率高的值用较短的比特位编码对罕见值用长码。这类似于哈夫曼编码的思想。实际应用例如有研究观察到中间特征的指数部分如果采用类浮点格式集中在少数几个值上。于是他们将最常见的3种指数值用2比特编码其他值仍用原始比特表示整体上减少了存储空间。注意事项这种方法需要在线或离线的统计分析来确定码表并且编解码过程会引入额外的计算开销。适用于那些相对稳定、数据分布可预测的张量。4.3 低位宽量化在精度与效率间走钢丝这是提升能效的“大杀器”。将FP32或FP16的数据用更低的位宽如FP8、INT8甚至混合精度来表示。训练友好的低位宽格式设计扩展动态范围梯度需要宽动态范围。因此自定义的8位浮点格式如E5M25位指数2位尾数比传统的INT8更适合训练因为它能通过指数位覆盖更大的数值范围。混合精度训练这是目前的主流实践。权重用高精度如FP16/BF16存储和更新确保收敛性前向和反向传播中的激活和梯度计算使用低精度如FP8提升计算速度。硬件需要支持不同精度格式之间的高效转换。细粒度混合精度更进一步甚至可以在一个张量内部对数值较大的重要元素采用较高精度如FP16对数值较小、影响不大的元素采用较低精度如FP8实现动态的精度分配。踩坑记录我们最初尝试全FP8训练时在几个复杂数据集上出现了收敛困难。后来分析发现某些层在训练初期梯度非常小FP8的尾数精度不足导致更新停滞。最终方案是采用“权重FP16激活FP8”的混合精度并在优化器状态保持FP32这才在保证收敛的同时获得了显著的能效提升。5. 核心优化技术三计算单元的能效革命内存访问省下来的能量最终要靠高效的计算单元来转化为算力。训练硬件的计算优化围绕两大主题利用稀疏性和定制精度硬件。5.1 超越输入零跳过输出零预测与权重剪枝既然训练中权重的稀疏性难以利用反向传播的输入也不稀疏那还有别的零可以跳过吗有那就是输出零。核心洞察在反向传播经过ReLU层时其梯度取决于前向传播时该神经元的输入是否大于0。如果前向时输入为负那么反向传播的梯度就是0。这个信息在前向传播结束时就已经确定了硬件实现在前向传播计算ReLU时额外生成一个“梯度掩码”位图记录哪些位置将会产生零梯度。在反向传播到达该层时硬件直接读取这个掩码。对于掩码指示为零输出的位置跳过所有计算其梯度所需的乘加运算。这可以跳过大量不必要的计算尤其是在深层网络。更激进的策略——训练中剪枝一些研究尝试在训练过程中就动态地识别并剪枝掉不重要的权重连接将其置零。一旦某个权重被判定为可剪枝在后续的所有训练步骤中与该权重相关的计算和梯度更新都可以被永久跳过。这需要算法和硬件的紧密协同硬件需要支持动态的稀疏模式。5.2 定制化处理单元为混合精度而生支持混合精度和自定义格式不能仅仅靠软件模拟必须在硬件层面设计高效的算术逻辑单元。设计挑战与选择浮点单元共享设计一个既能处理FP16又能处理FP8或2个FP8的融合乘加单元。通常可以共享面积最大的尾数乘法器和加法器但需要两套指数处理和对齐逻辑。关键在于平衡面积、功耗和灵活性。定点动态缩放采用定点格式但为每个张量甚至每个通道动态地维护一个缩放因子。计算时使用低比特定点数但通过缩放因子来恢复动态范围。硬件需要支持高效的缩放、饱和及舍入逻辑。近似计算对于训练中容错性更高的部分如某些梯度累加甚至可以引入近似加法器或乘法器以进一步降低功耗。一个具体设计我们研究过一款芯片其PE内部包含一个可配置的FMA。在FP16模式下它作为一个标准的半精度浮点乘加器工作。在FP8模式下它可以被拆分成两个独立的FP8乘加器将计算吞吐量翻倍。虽然控制逻辑复杂了些但相比部署两个独立的FP16单元面积和能效优势非常明显。6. 技术协同与权衡以LNPU芯片设计为例单独应用上述任何一项技术都能带来收益但真正的艺术在于如何让它们协同工作而不是相互掣肘。这里以论文中提到的LNPU芯片为例看看这些技术是如何被整合的。LNPU的核心设计哲学是联合优化内存访问与计算。它提出了“细粒度混合精度”编码方案将一个中间特征张量划分为两部分一部分是数值较大、需要较高精度的元素用FP16存储另一部分是数值较小、可以用FP8覆盖的元素。对这两部分数据分别进行游程编码压缩消除零值。压缩后的数据流可以直接送入PE阵列无需解压。PE阵列支持FP8/FP16混合精度计算和零值跳过。这个方案的巧妙之处在于压缩与计算联动RLE压缩直接减少了内存访问量而压缩后的数据格式非零值块正好适合PE的零跳过逻辑避免了先解压再判断的 overhead。精度与效率平衡FGMP根据数据分布动态分配精度在保证训练收敛的前提下最大化使用了高效的FP8计算。然而协同也意味着约束数据流优化的限制由于数据被RLE压缩其地址变得不规则使得阵列重排操作变得极其困难且低效。因此LNPU放弃了对压缩特征张量进行重排只对较小的、未压缩的权重张量块进行重排。微架构的固定为了支持基于输入的零跳过PE阵列中输入数据的数据通路被固定。这在一定程度上限制了为权重更新步骤优化数据通路的灵活性需要通过反转权重和输出的数据通路来弥补。这个例子深刻地说明芯片设计是一个系统工程没有银弹。每一项优化技术的引入都需要评估其对其他模块的影响在全局收益和局部代价之间做出权衡。7. 常见问题与设计陷阱实录在实际流片和调试过程中我们遇到了不少教科书上不会写的“坑”。问题一模拟精度完美上板训练发散现象在软件模拟器和RTL仿真中采用自定义8位格式的训练loss曲线正常下降。但芯片实测时训练几个epoch后loss突然爆炸。排查最终定位到是舍入模式不一致。软件模拟使用标准的“最近偶数”舍入而为了硬件面积优化我们最初在乘法器后使用了简单的“截断”舍入。在训练初期梯度很小时这种系统性偏差经过数万次迭代后被放大导致发散。解决硬件上改为支持“随机舍入”或“最近偶数舍入”虽然增加了一点面积但保证了数值稳定性。问题二稀疏加速效果不及预期现象设计了复杂的零跳过逻辑但实际运行ResNet训练时整体加速比远低于理论值。分析使用性能剖析工具发现虽然跳过了大量零值计算但负载不均衡导致严重。某些PE因为处理的数据块稀疏度高早早完工而其他PE还在忙碌造成整体计算阵列等待。解决引入了动态负载均衡机制。在数据调度层将压缩后的非零数据块进行统计并尽可能均匀地分配到各个PE上。同时在PE内部采用细粒度的任务队列进一步平滑计算波动。问题三片上缓存容量与带宽的博弈现象为了减少外部内存访问我们增大了用于存储中间特征的片上SRAM容量。但芯片功耗反而上升且最高频率下降。分析大容量SRAM带来了巨大的静态功耗和访问延迟。更糟糕的是由于SRAM端口数量有限在同时进行前向存特征和反向取特征时出现了端口争用产生了停顿。解决采用了分级存储策略。使用一小块高速、多端口的寄存器堆来缓存当前计算最活跃的数据块用较大但较慢的SRAM作为下一级缓存。并通过数据预取和访问调度算法尽可能隐藏SRAM的访问延迟。最终找到了一个面积、功耗和性能的最佳平衡点。问题四编译器与硬件的协同调试现象硬件测得的性能总是比手工优化的内核代码低30%以上。排查问题出在数据布局的隐式转换上。编译器为了通用性有时会在不同内存格式间插入隐式的转换操作例如NHWC转NCHW这些操作在源码中不可见但消耗了额外的周期和带宽。解决与编译器团队深度合作定义了一套清晰的硬件原语接口和内存布局约束。迫使算法开发者在明确知晓硬件偏好布局的情况下进行编程虽然增加了使用门槛但彻底释放了硬件性能。8. 未来展望与设计建议设备端DNN训练加速器仍是一个快速发展的领域。结合我们的经验我认为以下几个方向值得深入关注算法-硬件协同设计的深化未来的突破将更依赖于算法和硬件的共同创新。例如开发对硬件更友好的新型优化器、激活函数或网络结构能够天然产生更高的稀疏性或更规整的数据访问模式。存算一体技术的探索将部分计算移到内存内部可以彻底颠覆“内存墙”问题。虽然目前存算一体主要用于推理但其在训练中用于权重更新、梯度累加等内存密集型操作潜力巨大。动态自适应硬件训练的不同阶段初期、中期、末期对数值精度、稀疏性的需求是不同的。能够实时监测训练状态如梯度分布、loss变化并动态调整硬件精度、电压频率的“自适应加速器”可能成为下一代芯片的标配。标准化与生态建设目前各家训练加速器的指令集、数据格式、编程模型各异给开发者带来了巨大负担。推动设备端训练硬件接口和格式的标准化将是产业规模化应用的关键。对于想要踏入这个领域的设计师我的建议是永远从真实的工作负载出发。不要只盯着峰值算力这个数字。搭建一个从算法、编译器、架构到电路的全栈评估平台用完整的训练任务而不仅是几个算子来驱动你的设计决策。在资源受限的边缘设备上实现高效的训练是一场在算力、内存、功耗和精度之间进行的精妙平衡艺术每一个晶体管都要用在刀刃上。
