ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

AI芯片软硬件协同设计:脉动阵列与FP8量化实践

AI芯片软硬件协同设计:脉动阵列与FP8量化实践 1. AI芯片软硬件协同设计的核心逻辑1.1 为什么软硬件必须一起设计做AI芯片这行的人都有一个共识硬件堆算力容易让算力真正跑满难。我见过太多团队芯片流片回来理论峰值算力标得漂亮实际跑主流模型连30%的利用率都摸不到。问题出在哪出在软硬件各做各的。AI芯片和通用CPU最大的区别在于它的计算模式高度特化。矩阵乘法、卷积、注意力机制这些操作有固定的数据流模式硬件如果按照这些模式去定制数据通路效率能比通用架构高出一到两个数量级。但代价是一旦硬件定型能高效支持的算子集合就固定了。这时候如果软件栈不能精准地把模型映射到硬件上芯片就是一块昂贵的硅片。所以软硬件协同设计的本质是在硬件设计阶段就考虑软件怎么用在软件优化阶段理解硬件的瓶颈在哪。具体到工程实践这意味着几个关键决策要一起做数值格式的选择直接决定乘法器面积和功耗也决定量化工具链的设计数据复用策略决定片上缓存层次和带宽需求也决定编译器的调度算法指令集架构决定编程模型也决定算子库怎么写这三个决策环环相扣任何一个单独优化都可能让另外两个变成瓶颈。1.2 从模型到芯片的完整映射链路一个典型的AI芯片软件栈从上层到下层大致是这样的PyTorch / ONNX 模型 ↓ 图优化与算子融合编译器前端 ↓ 量化与数值格式转换 ↓ 算子映射与调度编译器后端 ↓ 指令生成与内存分配 ↓ 硬件执行每一层都在做一件事把上一层的抽象逐步降低直到变成硬件能执行的信号。这个链路里最容易出问题的地方是量化与算子映射的衔接。量化把FP32的权重和激活值压成FP8或INT8但不同硬件对FP8的支持方式不一样——有的支持E4M3和E5M2两种格式有的只支持一种有的甚至对累加器精度有特殊要求。如果编译器不知道这些细节生成的代码要么跑不对要么跑不快。我个人的经验是做AI芯片的软件栈量化工具链和编译器后端必须是同一拨人设计或者至少要有非常紧密的沟通。否则量化方案选了一个硬件不擅长的格式后面怎么调都别扭。1.3 当前主流技术路线的取舍市面上AI芯片的软硬件设计大致分三个流派第一派通用矩阵加速路线。代表是各类GPGPU和脉动阵列架构。硬件提供大规模的MAC阵列软件通过编译器把各种算子拆解成矩阵乘法。优点是通用性好新算子适配快缺点是数据搬运开销大能效比受限于内存带宽。第二派数据流架构路线。代表是各种TPU-like设计。硬件按照固定的数据流模式如权重 stationary、输出 stationary设计软件需要把模型严格映射到这种模式上。优点是能效比极高缺点是灵活性差不支持的算子只能回退到CPU。第三派可重构架构路线。硬件提供可配置的计算单元和互连网络软件在运行时动态配置。优点是兼顾灵活性和效率缺点是编译复杂度高工具链开发难度大。选哪条路线取决于目标场景。如果是云端训练通用性优先第一派更合适如果是边缘推理能效比优先第二派更有优势如果是需要支持多种模型且对能效有要求第三派值得考虑但工具链投入会很大。2. 脉动阵列的基本原理与设计要点2.1 脉动阵列到底在解决什么问题脉动阵列这个概念最早是1978年H.T. Kung提出的当时是为了解决VLSI时代计算单元和内存之间带宽不匹配的问题。放到今天的AI芯片语境下它解决的是同一个问题计算单元太快内存太慢数据搬运成了瓶颈。传统CPU做矩阵乘法是取一个数、算一次、存回去内存访问次数和计算次数是同一个量级。脉动阵列的思路是让数据在计算单元之间流动每个数据被取进来之后参与多次计算再离开。这样内存访问次数可以降到计算次数的几分之一甚至几十分之一。用一个生活化的类比传统计算像是一个厨师做菜每做一道菜都要去仓库拿一次食材。脉动阵列像是一条流水线食材从一端进入经过多个工位每个工位加工一下再传给下一个最后从另一端出来。仓库只需要在流水线入口送一次食材出口收一次成品。2.2 权重固定型脉动阵列的工作机制以最经典的权重固定Weight Stationary脉动阵列为例假设我们要计算矩阵乘法 C A × B其中A是M×K的激活矩阵B是K×N的权重矩阵。硬件是一个M×N的PEProcessing Element阵列。每个PE里存一个权重值第(i,j)个PE存的是B的第i行第j列的元素。计算过程是这样的激活值从左侧流入第i行的激活值a[i][k]依次进入第i行的第一个PE每个PE收到激活值后与本地权重相乘结果累加到本地累加器激活值继续向右流动进入下一个PE经过K个周期后每个PE的累加器里就是C的一个元素这个过程中权重只在初始化时加载一次之后一直留在PE里。激活值从左边流入横向流动。部分和留在PE内部不需要搬来搬去。注意权重固定型脉动阵列适合权重矩阵较大的场景因为权重加载的开销可以被大量激活值的计算摊薄。如果权重矩阵很小加载开销占比就高了。2.3 输出固定型与行固定型的对比除了权重固定还有两种常见的脉动阵列数据流输出固定Output Stationary每个PE负责计算输出矩阵的一个元素激活值和权重都从外部流入在PE内部相乘累加。这种方式的优点是输出不需要在PE之间传递适合输出矩阵较大的场景。缺点是激活值和权重都需要广播对互连网络的压力较大。行固定Row Stationary这是Google TPU采用的方式结合了权重固定和输出固定的特点。每一行PE负责计算输出的一行权重在行内固定激活值在行内流动部分和在行间传递。这种方式在卷积神经网络上表现很好因为卷积可以展开成矩阵乘法且权重复用率高。三种数据流的对比如下数据流类型权重加载次数激活值加载次数部分和传递适用场景权重固定1次N次无权重矩阵大激活矩阵小输出固定M次K次无输出矩阵大行固定1次N次行间传递卷积为主权重复用高实际芯片设计中很少有纯用一种数据流的通常是混合方案。比如卷积层用行固定全连接层用权重固定通过配置寄存器切换模式。2.4 脉动阵列的尺寸选择与利用率计算脉动阵列的尺寸M×N是一个关键设计参数。尺寸越大峰值算力越高但利用率可能越低。原因很简单如果矩阵维度小于阵列尺寸部分PE就会闲置。假设阵列是128×128处理一个64×64的矩阵乘法。实际用到的PE只有64×644096个而总共有16384个PE利用率只有25%。如果处理的是256×256的矩阵利用率可以到100%。利用率计算公式利用率 (实际计算量) / (阵列峰值算力 × 时间)对于M×N阵列处理K维矩阵乘法如果M、N、K都大于等于阵列尺寸利用率接近100%。如果某个维度小于阵列尺寸利用率就会下降。实际模型中矩阵维度是变化的。Transformer的注意力矩阵可能是序列长度×序列长度序列长度从几十到几千不等。所以脉动阵列的尺寸选择要折中太大则小矩阵利用率低太小则大矩阵算得慢。我个人的经验是128×128到256×256是比较甜点的区间。再大小模型利用率掉得厉害再小大模型算力不够。3. 数值格式的选择与FP8实践3.1 为什么数值格式是AI芯片的核心设计决策数值格式决定了三件事计算精度、硬件面积、功耗。这三者互相制约。用FP32做推理精度绰绰有余但乘法器和加法器的面积是FP16的4倍左右功耗也是4倍左右。用INT8面积和功耗都小但精度损失可能影响模型准确率。FP8是这两者之间的折中面积和功耗接近INT8但动态范围比INT8大得多对异常值的容忍度更好。FP8有两种主流格式E4M34位指数3位尾数。动态范围较小但精度较高。适合前向传播的激活值和权重。E5M25位指数2位尾数。动态范围较大但精度较低。适合梯度计算因为梯度值分布范围广。注意E4M3和E5M2的选择不是随意的。E4M3能表示的最大值是448E5M2能表示的最大值是57344。如果激活值中有超过448的值用E4M3就会溢出必须用E5M2或者做缩放。3.2 FP8量化对模型精度的影响FP8量化不是简单地把FP32截断成FP8。直接截断会导致精度大幅下降因为FP8的尾数只有2到3位舍入误差很大。实际做法是缩放舍入统计张量的最大值或分位数计算缩放因子把张量映射到FP8的可表示范围对缩放后的值做舍入缩放因子的选择很关键。如果按最大值缩放异常值会拉低整体精度如果按分位数缩放超出范围的值会被截断。实践中常用的是动态缩放每个batch或每个通道单独计算缩放因子。我实测下来FP8量化对大多数视觉模型的影响在1%以内对语言模型的影响稍大在1%到3%之间。但如果量化方案设计得不好掉点可能超过10%。3.3 FP8与GGUF量化格式的关系GGUF是llama.cpp推出的一种模型文件格式里面包含了多种量化方案。常见的GGUF量化类型有Q4_0、Q4_K、Q5_K、Q8_0等数字代表位数字母代表具体的量化策略。FP8和GGUF量化是不同层面的东西。FP8是一种数值格式GGUF是一种文件格式和量化方案集合。GGUF里也可以包含FP8量化的模型但更常见的是INT4、INT5、INT8等整数量化。如果要把FP8模型转成GGUF格式需要注意几点GGUF的量化工具链主要支持整数量化FP8支持有限FP8的缩放因子需要额外存储GGUF格式对元数据的支持需要确认转换后的模型在推理时的反量化开销需要考虑提示如果你手头有FP8量化的模型想用GGUF生态的工具跑最稳妥的做法是先反量化回FP16再用GGUF的工具重新量化。直接转换可能会遇到格式不兼容的问题。3.4 FP8在训练和推理中的不同策略FP8在训练和推理中的使用方式差别很大。推理场景权重和激活值都用FP8累加器用FP16或FP32。缩放因子可以离线计算好固定在模型文件里。推理时不需要动态调整硬件实现简单。训练场景前向传播用FP8反向传播的梯度用FP8或FP16权重更新用FP32。缩放因子需要动态调整因为训练过程中激活值的分布会变化。这就需要硬件支持动态缩放实现复杂度高不少。目前FP8训练的主流方案是混合精度大部分计算用FP8关键部分如LayerNorm、Softmax用FP16或FP32。这样既能享受FP8的速度又能保持训练的稳定性。4. 软硬件协同的实操要点4.1 从模型到硬件的映射流程把一个训练好的模型部署到AI芯片上大致要经过这些步骤模型导出从训练框架导出为ONNX或类似格式图优化常量折叠、算子融合、死代码消除量化把FP32权重和激活值转成FP8或INT8算子映射把优化后的算子映射到硬件支持的指令内存分配为每个张量分配片上或片外内存指令调度生成指令序列安排计算和数据搬运的顺序代码生成输出硬件可执行的二进制这个流程里算子映射和指令调度是最容易出问题的环节。因为硬件支持的算子集合有限遇到不支持的算子就要拆解或回退。拆解的方式直接影响性能回退到CPU则可能成为瓶颈。4.2 算子融合的收益与风险算子融合是编译器优化的常用手段。比如把ConvBNReLU融合成一个算子可以减少内存访问次数提升性能。融合的收益可以用一个简单模型估算假设三个算子单独执行每个算子需要读一次输入、写一次输出总共6次内存访问。融合后只需要读一次输入、写一次输出内存访问降到2次。如果内存访问是瓶颈性能提升可能接近3倍。但融合也有风险融合后的算子可能超出硬件资源限制如寄存器数量、片上缓存大小融合可能改变数值精度如BN的均值和方差在融合后需要重新计算融合可能影响调试因为中间结果不可见我个人的做法是先融合收益大的算子保留必要的中间结果用于调试。等模型稳定后再做更激进的融合。4.3 内存层次的设计与数据复用AI芯片的内存层次通常包括寄存器文件最快容量最小每个PE私有片上缓存较快容量中等多个PE共享片外内存最慢容量最大全局共享数据复用的核心思想是让数据在最快的存储层次里被多次使用。比如权重加载到寄存器后参与多次乘加运算激活值加载到片上缓存后被多个PE读取。数据复用率可以用算术强度Arithmetic Intensity来衡量算术强度 计算操作数 / 内存访问字节数算术强度越高说明每次内存访问对应的计算越多对内存带宽的要求越低。脉动阵列的设计目标就是提高算术强度。以一个128×128的脉动阵列为例处理128×128×128的矩阵乘法计算量128×128×128 2,097,152次乘加内存访问读入128×128的激活矩阵和128×128的权重矩阵共2×128×128×4字节 131,072字节算术强度2,097,152 / 131,072 16次乘加/字节如果内存带宽是100GB/s理论算力就是100×16 1600 GOPS。这个估算可以帮助判断内存带宽是否足够。4.4 编译器的调度策略编译器后端要做的事情简单说就是在满足资源约束的前提下最小化执行时间。资源约束包括片上缓存大小、寄存器数量、指令队列深度等。执行时间取决于计算时间和数据搬运时间的重叠程度。一个好的调度策略应该做到计算和数据搬运重叠在计算当前数据块的同时预取下一个数据块双缓冲用两块缓存交替一块用于计算一块用于加载流水线填充让脉动阵列的流水线尽量满这些策略说起来简单实现起来要考虑很多细节。比如双缓冲的切换时机太早切换会导致数据还没算完就被覆盖太晚切换会导致计算单元等待数据。我踩过的一个坑是预取距离设得太远导致片上缓存不够用。预取距离是指提前多少个周期开始加载下一个数据块。设得太小数据来不及加载设得太大缓存里堆了太多数据挤占了其他用途的空间。后来我们做了一个自适应预取距离的机制根据缓存压力动态调整效果好了很多。5. 常见问题与排查技巧5.1 精度掉点问题的排查思路FP8量化后模型精度掉点是最常见的问题。排查思路可以按这个顺序来第一步确认掉点是否在可接受范围。不同任务对精度的敏感度不同。分类任务掉1%可能可以接受检测任务掉1%可能就不能接受了。第二步定位掉点发生在哪一层。逐层量化看哪一层的输出误差最大。通常是那些激活值分布范围大的层比如注意力层的Softmax之前。第三步检查缩放因子。缩放因子是按最大值算的还是按分位数算的分位数选了多少如果按最大值算异常值会拉低整体精度如果分位数选得太低截断误差会大。第四步考虑混合精度。对精度敏感的层保留FP16或FP32其他层用FP8。混合精度的粒度可以是层级别也可以是通道级别。提示排查精度问题时建议保留一个FP32的参考实现逐层对比输出。这样能快速定位问题层。5.2 性能不达预期的常见原因芯片跑起来但性能不达预期可能的原因有现象可能原因排查方法算力利用率低矩阵维度小于阵列尺寸检查模型各层的矩阵维度算力利用率低内存带宽不足计算算术强度对比带宽算力利用率低算子融合不够查看生成的指令序列延迟高数据搬运未重叠检查是否有双缓冲延迟高流水线填充不足检查流水线深度和启动开销功耗高数据搬运过多检查数据复用率我遇到最多的情况是矩阵维度不匹配导致的利用率低。比如阵列是128×128但模型里有很多64×64的矩阵乘法利用率直接掉到25%。解决办法要么是调整阵列尺寸要么是把多个小矩阵拼成一个大矩阵。5.3 软硬件接口的调试技巧软硬件接口是最容易出问题的地方因为两边的人用的语言和工具都不一样。硬件工程师看波形软件工程师看日志。接口出问题时两边经常互相甩锅。我的经验是在接口上定义一套完整的断言和计数器。断言用于检查协议合规性比如握手信号是否成对出现、数据是否在有效期内稳定。计数器用于统计接口的吞吐量、延迟、错误次数。这套机制在流片前就能发现大部分接口问题。流片后如果还有问题也可以通过计数器快速定位是硬件的问题还是软件的问题。5.4 模型适配的避坑清单最后整理一份模型适配的避坑清单都是实际踩过的坑动态shape支持很多AI芯片对动态shape支持不好模型导出时尽量固定shape或者把动态维度放在最外层算子版本兼容ONNX的算子版本更新很快导出时注意目标硬件支持的算子版本量化校准集校准集要覆盖实际推理时的数据分布否则缩放因子会偏内存对齐很多硬件要求内存地址对齐分配内存时注意对齐要求边界情况padding、stride、dilation这些参数在边界情况下的行为要和硬件确认数值溢出FP8的表示范围有限中间结果可能溢出需要插入缩放或截断这些坑每一个都可能导致模型跑不对或跑不快而且排查起来都很费时间。提前了解可以省很多事。我个人在实际操作中的体会是AI芯片的软硬件设计没有银弹。脉动阵列、FP8、算子融合这些技术都是工具用得好不好取决于对具体场景的理解。同样一个模型在不同的硬件上最优的映射方案可能完全不同。所以做这个方向既要懂硬件原理又要懂模型特性还要有足够的耐心去调优。踩坑是常态但每踩一个坑对系统的理解就深一层。
返回列表