免费获取学习方案
ARTICLE DETAIL

资讯详情

深耕编程基础知识与建站技术分享的一线实战洞察。

大模型时代国产AI芯片的胜负手:全栈协同

大模型时代国产AI芯片的胜负手:全栈协同 大模型规模还在膨胀AI芯片的迭代已经明显跟不上了。但凡在国产AI芯片上跑过大模型训练和推理的人对“全栈协同”这四个字的体会都会特别深单卡算力指标再好看一到千卡集群就原形毕露算子库少一个融合实现端到端吞吐直接腰斩通信拓扑没调好AllReduce 跑得比硬盘还慢。这篇文章我想从算力账、芯片现状、全栈协同的层次、实操路径和踩坑经验几个角度把“为什么说国内AI芯片的胜负手是全栈协同”这件事讲透。适合算法工程师、芯片软件栈开发者和做AI基础设施选型的人参考。1. 大模型膨胀为什么最先扛不住的是芯片1.1 算力需求不是线性涨是指数涨先算一笔最基本的账。训练一个大模型需要的浮点运算量大约等于 6×N×D其中 N 是模型参数量D 是训练数据量token 数。这个公式是业界公认的粗算口径前向传播和反向传播各占一半左右再算上激活重算之类的损耗实际会再上浮 20% 到 50%。从 GPT-3 的 175B 参数训练到 Chinchilla 时代讲究“数据与参数按比例放大”模型规模一路从 10B 涨到 100B、500B甚至还有 1T 以上的实验模型。同样训练 1T token70B 模型的算力需求是 7B 模型的 10 倍而 1T 参数模型就是 100 倍。单卡算力哪怕每年提升 50%也根本追不上这种规模膨胀的速度。所以大模型时代对芯片的第一重打击是“算力缺口”从单卡级别扩大到了集群级别。以前一张卡能训完 ResNet现在要训一个 70B 模型哪怕用 BF16 混合精度把显存压到极限也至少需要几十张卡并联。于是芯片不再是一个单点硬件而是一个集群网络里的一环。1.2 显存和带宽才是真正的天花板很多人以为只要 FLOPS 够高芯片就够用。但大模型跑起来之后最先爆掉的往往是显存。以 7B 模型为例BF16 权重约 14GB训练时反向传播需要的梯度又要 14GBAdamW 优化器状态包括主权重、动量和方差按每参数 12 字节粗算又是 84GB 左右。这几项加起来光训练一个 7B 模型就要上百 GB 显存这还没算激活值、中间变量和通信缓冲。70B 模型在单卡上根本不可能放下必须走张量并行、流水并行、ZeRO 这类显存卸载策略。推理侧同样紧张。KV Cache 与序列长度和 batch 大小直接成正比一个 70B 模型假设 8K 上下文、32 层、64 头、每头 128 维BF16 精度下每个 token 的 KV 大小约 2×32×64×128×2 字节算下来约 1MB。看起来很小但 batch 到了 64、序列长度到了 32KKV Cache 就可能轻松超过 100GB比权重本身还占地方。带宽问题更隐蔽。大模型推理是典型的“访存密集型”场景生成式解码每个 token 都要把全部权重读一遍。70B 模型一次前向就要读 140GB 数据如果 HBM 带宽只有 3.3TB/s每秒最多只能生成 23 个 token 左右这还是不计算 attention 计算时间的理论上限。所以芯片厂商在堆 FLOPS 的同时更应该在 HBM 带宽、片上 SRAM 容量和内存带宽利用率上下功夫。算力峰值再高数据喂不进去一切白搭。1.3 集群规模上台后通信和调度的账不好算当模型从单机多卡变成千卡、万卡集群真正的瓶颈就开始从“单芯片”转移到“网络”。分布式训练每轮迭代都要做梯度同步典型做法是 Ring AllReduce。每做一次全规约需要传输大约 2 倍模型字节数70B 模型用 BF16 训练每次迭代光梯度同步就要传 280GB 数据。如果卡间互联带宽只有 200Gbps光通信时间就让人无法接受。通信占比直接决定了扩展效率1000 张卡看起来是单卡的 1000 倍算力但如果通信设计不好扩展效率可能只有 50% 甚至更低。而且大集群里慢节点问题非常突出几千张卡里总有一张散热异常、PCIe 降速、网络重传整个训练任务的速度会被这张“拖后腿”的卡拉低。有没有弹性容错、能不能动态剔除慢节点、Checkpoint 恢复要多久这些都属于芯片之外、但在大模型时代同样决定成败的问题。所以说大模型膨胀给芯片“上强度”的方式已经不再只是“单卡 FLOPS 多少个 T”而是从显存、带宽、互联、调度、容错整个系统一起被推到了极限。2. 国产AI芯片在赢什么又在赌什么2.1 硬件架构的追赶矩阵单元、存储墙与Chiplet国产AI芯片在硬件层面的策略普遍是“扬长避短”在制程工艺不占优势的情况下通过架构创新把有效算力堆上去。最常见的手段有三类。第一是加大矩阵计算单元的面积占比。大模型核心计算是 GEMM也就是矩阵乘法如果芯片能把更多晶体管留给矩阵单元少放一些通用控制逻辑单位面积内的峰值算力自然更高。NPU、脉动阵列、二维阵列架构都属于这个思路。第二是提升片上存储容量和带宽。HBM 不够用的情况下把 SRAM 做大、做多让数据尽量留在片上。因为 GEMM 的访存密集程度极高如果每次计算都要去 HBM 拿数据再大的 FLOPS 也会被访存拖死。第三是 Chiplet 和先进封装。把计算die、存储 die、IO die 通过先进封装组合到一颗芯片里可以有效避免单芯片面积过大带来的良率和功耗问题。多个 chiplet 拼在一起成本可控扩展性也更灵活。这些做法的共同点是单看峰值算力国产芯片已经能和国际主流产品有得一拼真到了跑模型的时候差距才慢慢暴露出来。2.2 软件栈的鸿沟算子、编译器和调试工具硬件展示的跑分再高都是厂商用“精心调过的样例”跑出来的。真实场景下你会遇到什么第一算子缺失。模型里有特殊的激活函数、自定义 attention 算子的变体、某层网络写成很不常见的形式结果编译不过去或者直接 fallback 到极其低效的通用实现。跑一个 benchmark 只剩 30% 的利用率原因很可能只是一个算子实现不完整。第二编译器不够聪明。英伟达生态里的对标方案往往在 CUDA 层面就有大量算子模板和并行策略而国产芯片的编译器在面对动态 shape、复杂数据流和融合逻辑时经常生成效率低下的指令序列。同样的模型在别人家芯片上能自动融合的算子在国产芯片上可能被拆成几十次内核启动。第三调试工具的成熟度差一截。出了 NaN、显存越界、Slow Node你需要在 profiling 工具里快速定位到具体算子、具体卡、具体通信时序。如果 profiling 工具集成度不高采样 overhead 又大排查一个分布式训练问题可能要花几天时间。这也是为什么我一直认为国产AI芯片现在的竞争重点不在纸面跑分而在“别让开发者骂娘”。2.3 兼容还是原生绕开生态还是重构生态面对成熟的加速计算生态国产芯片通常有两条路线。一是兼容路线在指令翻译层或 API 层兼容主流加速接口。这样做的好处是迁移成本低已有模型代码能跑起来。但风险在于翻译层有性能损耗且永远是在追别人定义的新特性一旦对手更新架构翻译层又得加班加点适配。二是原生路线围绕自家芯片重新造一套软件栈、算子库、编译器和调试工具。性能上限高但生态积累很慢开发者习惯迁移是个大问题。现实没有完美答案大多数厂商是“兼容先行原生随后”。但不管是哪条路线最终能走多远取决于有没有做“全栈协同”的意愿和执行力。芯片、算子库、编译器、框架、训练策略、集群调度彼此之间是深度耦合的。如果只做一个“能用”的兼容层而不去围绕芯片特点优化模型算法和框架策略那国产芯片永远只能吃别人定义好的红利很难形成自己的优势。3. 全栈协同到底在“协同”什么3.1 芯片与算子协同GEMM、融合与数据布局算子是软件栈和硬件之间最短的连接线。全栈协同的第一层就是算子层面的协同。以最核心的矩阵乘法 GEMM 为例。芯片的矩阵单元尺寸决定了分块策略假设矩阵单元一次能算 16×16 的 tile那么你就得把大矩阵切成 16 的整数倍小块每个小块的计算结果要留在片上累加再配合寄存器级的双缓冲隐藏内存访问延迟。如果数据在内存里的布局是 row-major 但芯片加载器偏好 column-major你不做 padding 或 transpose访存效率立刻下降。再举一个更实际的例子大模型推理里的注意力计算通常把 Query 和 Key 算完点积后要经过 softmax 再乘 Value如果每一步都写回 HBM白白多出好几轮访存。FlashAttention 这类融合算子就是把整个 attention 块放在片上流式处理用分段计算加在线 softmax 避免完整矩阵的上上下下。国产芯片要实现同等效率不能只靠抄 FlashAttention 的思想还得根据自家 SRAM 大小、寄存器数量和矩阵单元的尺寸去做重新切分。所以第一层协同的本质就是让算子实现照着芯片的微架构特点去做融合、分块和调度。这要求芯片架构师、算子库工程师和编译器团队从一开始就坐在一起而不是芯片流片完再想办法适配。3.2 并行训练框架协同从DP到TP/PP/EP的组合拳第二层协同在训练框架。大模型卡多了之后光靠数据并行显然不行每个 GPU 复制一份完整模型显存放不下。所以分布式训练要把模型“切”开常用的切法包括数据并行DP每张卡持有完整模型只切数据批次。实现简单但模型太大时不适用。张量并行TP把一层的大矩阵按行/列切开分布在多张卡上层内通信频繁需要高速互联。流水并行PP按层切分像工厂流水线一样把不同层放到不同卡上通信相对较少但存在流水气泡。专家并行EPMoE 模型里每个专家放到不同卡上token 按路由分散通信模式比较特殊。全栈协同的第二层就是把并行策略和底层通信原语、显存分配方式深度绑定。举一个真实例子如果芯片的高速互联只支持特定消息大小效率最高那么切分模型时就要尽量把切分粒度对齐到这个高效区间。如果框架自动把 Tensor 切成 1MB 的碎片传输而硬件最适合 4MB 的消息通信性能就白丢一半。框架层的协同还涉及显存管理。ZeRO 阶段一到三本质是用通信换显存把优化器状态、梯度、参数依次做分片。可是在国产芯片上这套机制是否真的能吃到红利取决于通信库能不能处理小消息的聚合、内存分配器会不会产生大量碎片、以及多卡同步是否高效。框架不感知芯片就谈不上协同。3.3 集群调度协同拓扑、故障与弹性大模型训练跑在集群上芯片之间的物理拓扑直接决定了通信效率。常见的 8 卡机内互联、跨机 RDMA 组网都有固定的 bandwidth 差异。全栈协同要做到“通信感知调度”把通信量大的并行组比如 TP 组分配到物理拓扑上紧邻的卡比如同一交换机下、同一机柜内。把通信量小的并行维度比如 PP可以适当跨机甚至利用机间带宽。调度器要能感知每个节点的实时健康状态自动把慢节点上的任务迁移走。故障恢复同样不能忽视。大模型训练跑一个月中途必然有卡或节点出问题。如果 Checkpoint 频率太低一次故障损失十几个小时频率太高写入开销又会压掉训练吞吐。全栈协同的集群层解决方案通常是做分层 Checkpoint模型权重、优化器状态、RNG 状态分不同优先级处理重要状态用异步方式持续写入普通状态周期性落盘。3.4 模型算法协同让大模型适应芯片很多人做模型加速的思路是把现有模型原封不动搬到新芯片上跑不动就猜是芯片不行。但其实很多时候算法侧只要稍微让步就能释放硬件潜力。MoE 模型天然适合多芯片并行。把不同的专家分到不同卡上推理时只有少数专家被激活但要解决的是 token 路由带来的通信热点如果芯片的集合通信能支持稀疏聚合模式效果会好很多。投机采样也很有代表性。推理时用一个小模型先生成若干候选 token再让大模型验证。小模型的结构要是恰好适合芯片上的小 batch 高效执行而大模型的大 batch 验证又能吃满算力整个系统的 token 吞吐可以翻倍。设计投机采样器的 draft 模型时就要考虑目标芯片的并行特性和算子开销。量化则是深度绑定硬件的典型。不同芯片对低比特计算的支持方式完全不同有的支持 INT8 的向量指令有的可以做 FP8 的矩阵乘累加有的对 INT4 权重做了专门的 load-and-dequant 流水。量化方案必须照着芯片的指令集和数值范围去选否则精度和性能很难兼顾。这也是我反复强调“全栈协同”的原因越到后期模型和芯片之间就不再是“适配关系”而是“互为设计”的关系。4. 在国产芯片上加速大模型实操路径与踩坑实录4.1 部署的第一步不是调参而是先确认“算力边界”我见过太多人拿到国产加速卡第一件事就是跑模型、看精度然后抱怨“为什么这么慢”。正确的做法是先做一次无损的硬件基准测试。先把 NVIDIA 生态里常用的矩阵乘法 benchmark 移植过来分别测不同 shape 下的 FLOPS再测一遍 HBM 读写带宽、卡间通信带宽和延迟。不要用厂商默认的“验证跑分”用例而是用你真实模型里最常见的 shape。我一般会在部署前先写一个小脚本跑三件事用真实 GEMM shape 测试单卡矩阵乘峰值效率看看能到理论峰值的百分之多少。测一下 8 卡 AllReduce 在不同消息大小下的带宽曲线。测一下多卡之间的点对点延迟特别是不同机柜之间的延迟。这样做完你就能在脑海里画出一张“哪些环节还有潜力哪些环节已经到上限”的图。后续模型调优时心里有底。4.2 算子缺失和数据格式不对怎么办真实模型跑到国产芯片上最容易出现的报错就是某某算子当前后端不支持。这时候有三个解决路径优先级从高到低第一尽量调整模型代码的结构。比如把自定义 attention 函数改成框架内置的 SDPA 接口如果底层已经有融合实现直接调用效率最高。我见过不少模型的 attention 是作者自己手写的结构上和标准实现有一些小差异但其实可以通过调整 mask 和 dropout 的写法完全等价地替换成标准接口。第二用图改写绕过。某些算子组合起来其实可以被替换。比如 layer_norm residual 两个线性层如果芯片的算子库里有融合的 FFN 算子就直接替换整个模块如果编译器支持自定义 pattern 匹配可以做自动图改写。第三实在绕不开的算子自己实现。如果芯片提供低门槛的编程接口可以直接写自定义 kernel如果没有就只能用上层语言去组合更基础的算子性能会打折但至少能跑通。这个优先级我在项目里反复强调不要一上来就写自定义 kernel大概率是在错误的地方花大力气。数据格式是另一个隐形坑。国产加速卡对 Tensor 内存布局经常有偏好比如某些 NCHW 和 NHWC 之间切换成本极高。很多模型代码里会随手调用 permute、transpose、view 之类的操作看起来只是改变一个 shape实际会让底层存储布局变化引发多余的拷贝。调优的时候优先检查模型里有没有频繁的 transpose尽量把布局在设计阶段就固定下来。4.3 分布式训练跑不起来的排查思路分布式训练跑不起来是最磨人心态的。常见问题我整理了一个排查顺序先看通信库版本和集群网络配置是否一致。多机训练时不同节点的 RDMA、TCP 配置不一致就会导致连接超时。再看卡间拓扑是否被框架正确识别。很多国产芯片支持多路互联但如果框架默认按 PCIe Switch 分组就会把本来应该走高速互联的卡分到慢速通信域性能暴跌。然后看是不是 AllReduce 消息大小触发了通信库的不同算法路径。消息太大时ring 算法优势明显消息太小时tree 算法可能更高效。通信库如果不能自己动态选择就要业务侧手动设置。最后看是否存在“活锁”等待同步的卡在同一个通信域里因为某张卡显存分配失败而一直阻塞其他卡全部空转。这种情况往往不会直接报错只会表现为“训练卡住”。排查的时候我习惯用一个小技巧先把多卡训练降成 2 卡把模型缩到极小逐层加复杂度。这样能快速把问题从“模型太大”和“基础设施坏了”两类原因里分离出来。4.4 性能从“能跑”到“跑满”量化、批处理与KV Cache调优跑通只是开始让模型把芯片用满才是关键。推理侧的调优第一看 batch 和 SPLIT 策略。大模型推理如果 batch 太小访存密集特性会让芯片利用率很低batch 加大之后计算密度上升但 KV Cache 也会暴涨。我一般会先用测带宽的小工具摸清芯片在什么 batch 下达到性能平台再决定服务端要不要做 dynamic batching 和 continuous batching。千万别拿单请求性能去推整体吞吐那会严重低估真实负载。第二看量化。国产芯片对低比特的支持差别很大有的可以用 FP16 加 AWQ 这种权重压缩方案有的更适合跑 INT8 动态量化。实操时要注意不能只看权重能不能塞进显存还要看反量化算子是不是 fused 到了矩阵计算里。如果权重每次加载都要额外解压吞吐会被吃掉一大截。训练侧的调优除了前面讲的并行策略还要注意梯度累积步数和通信重叠。有的框架支持把反向传播和梯度同步做重叠在计算当前层梯度时异步通信前几层的梯度。如果芯片驱动和通信库支持多流这个优化能让端到端训练时间缩短 15% 到 30%。5. 全栈协同的团队与优先级经验5.1 做平台的人要和做模型的人坐在一起全栈协同真正落地的时候会碰到一个普遍的组织难题芯片厂商的软件团队、框架开发团队、算法业务团队往往分属不同部门目标也不一致。芯片团队关心如何把硅片跑满载框架团队关心兼容多少模型算法团队关心指标跌没跌。如果不打通很容易出现典型的“优化断层”芯片团队把某个 benchmark 调得很好看但框架层不感知业务层也没人去消费这个红利。我在实践中见过两种有效做法。一种是让算法工程师进入芯片测评组的日常流程用真实模型跑分倒逼芯片算子设计另一种是建立联合攻关小组专门针对某几个关键大模型做全链路优化每一层优化都要有对应负责人最终用端到端收益作为统一口径。5.2 优先级排序算力、易用性、生态、成本四个维度预算和人力总是有限的全栈协同工作和起手就想把一切都做完美。我习惯用四个维度判断优先级算力缺口当前模型在目标芯片上距离理论峰值差多少如果算子融合能解决 30% 的差距就先做算子融合。易用性缺口文档、调试工具、示例代码是否完备如果工程师连 profiling 都做不了后续一切优化都是空谈。生态缺口主流大模型能不能开箱即用如果跑一个模型要改上百行代码这个芯片永远不会被真正用起来。成本缺口综合采购成本、机房功耗、运维投入和人工改造成本是不是比现有方案划算这几个维度互相影响。比如生态缺口再大如果算力优势明显还是值得投入反过来模型跑起来虽然快但算子文档一塌糊涂也会劝退大部分用户。真正有经验的团队会把“算力缺口”和“易用性缺口”放在最前面因为这两个直接决定技术能不能成立。5.3 别怕给模型“动刀子”在国产芯片上做优化最大的心态障碍就是“模型是别人训练的不能改”。但实操中很多优化要动模型结构的一点点“皮毛”收益却很大。举个最常见的例子把输出层的 logits 计算从密集 softmax 改成 chunked softmax可以显著减少中间矩阵的显存占用对芯片而言显存压力小了batch 就可以加大端到端吞吐自然提升。还有 attention 里的因果 mask如果芯片支持 2D 掩码的稀疏计算那么模型就可以少一点无效计算但这需要模型把 mask 的声明方式改一下对齐到芯片支持的模式。我不建议为了堆跑分去改模型语义这不负责任。但如果是等价变换或者对下游影响可测量的轻微调整那该改就改。芯片特性和模型设计的双向适配才是全栈协同的真正含义。6. 常见问题速查国产芯片上跑大模型我遇到的典型故障这一节我把几年来在国产加速卡上遇到的高频问题整理成表格供大家排查时直接照抄。现象可能原因排查思路解决经验模型加载后显存报错权重精度和 KV Cache 预留分配冲突先用最小 batch 逐个开关开启二分定位哪项配置溢出关掉框架自动内存池手动设置显存分片上限单卡利用率低只有 30%算子未融合频繁内核启动用 profiling 工具查内核启动时间占比查找是否有对应融合算子或临时降低小算子调用频率多卡扩展效率只有 50%通信拓扑没对齐到物理互联打印通信域里的卡编号顺序重排卡顺序让 TP 组落在同一高速互联域内AllReduce 偶尔超时网络拥塞或网卡中断观察重传率和网络中断计数开启通信库的 QoS或换到专用的 high-priority 通道训练 loss 和参考实现不一致混合精度策略或算子数值差异对比逐步打开 FP32 与 BF16 下的 loss 曲线手动设置缩放因子必要时关闭某些算子的低比特分支推理 batch 打上去后延迟剧增KV Cache 访问模式导致缓存不友好检查是否开启 PagedAttention 一类方案调整 page 大小和预取策略或改用 chunked prefill训练中途某卡突然变慢散热降频或 PCIe 链路重协商看系统日志和卡温升曲线通过调度器重启任务并剔除慢节点模型量化后精度掉太多没有做量化感知训练或校准集不够检查校准集是否覆盖长尾数据用少量训练数据重新做 AWQ/RTN 校准算子 fallback 到了 CPU国产芯片后端没有对应算子实现查编译日志里的 fallback 标记用图改写或替换成等价算子组合框架与驱动版本不匹配底层运行时 API 变了看启动日志里的版本校验信息统一版本特别注意驱动、通信库、框架三者配套关系这张表每个问题我都实际遇到过其中最让我印象深刻的是通信拓扑问题。有次我们在测试国产芯片的多机训练8 机 64 卡扩展效率只有 55%一开始都怀疑是互联带宽不行。后来用工具把卡间物理拓扑打出来才发现框架按默认顺序把 TP 通信组分配到了不同机柜的卡上等于让高频通信走了远路。重新排序后同一组 TP 落在同一机柜内扩展效率直接拉到 72%。还有一个很实用的经验性能调优不要先怀疑芯片先怀疑自己的使用姿势。这句话说起来像废话但掉坑掉得多了就会发现百分之七八十的性能问题都出在软件栈配置、数据布局、通信拓扑和算子实现选择这些地方。真正“芯片算力不够”的场景反而少见。这不是替国产芯片说好话而是我踩过的坑确实大多在软件层。我个人这两年做国产芯片适配最大的体会是全栈协同不是锦上添花而是决定一款国产 AI 芯片能不能真正打进大模型时代的入场券。芯片硬件只在流片之前被讨论而软件栈、工具链、框架适配和模型协同则要在产品生命周期里被反复打磨。谁能在“全栈”这条链路上跑得更顺谁就有资格站在牌桌上。最后再分享一个小建议如果你所在团队正在评估国产 AI 芯片别只让采购和 CEO 看 PPT也不要用英伟达的 benchmark 脚本直接跑。挑一个你们真实业务里最常用的大模型指定一个真实延迟目标让芯片厂商的软件团队陪你把模型从训练到推理完整跑一遍。谁家的软件栈能在两周内给你提交一份端到端性能报告并且问题清单越来越少那家才是有全栈协同能力的选手。
返回列表