从算子优化到AI Infra：大模型训练性能的隐形战场

训练群里最动人的一刻往往是算法同学和资源管理员都沉默的时候。Loss已经开始收敛可吞吐就是上不去加了卡也白加。这时候一个平时不常露面的名字会突然冒出来贴一张性能分析的时间线淡淡地说一句“瓶颈不在算力在通信。”这个人的岗位多半写的是AI Infra。大模型这两年把“算法”捧得很高但真正下场跑过训练的人都知道一个万亿参数模型能稳定地训练出来背后靠的绝不是某一次灵光一闪的网络结构改动。算子库、分布式策略、推理引擎、资源调度、数据管道这些名字听起来像“工程杂活”的东西恰恰决定了模型能不能训起来、跑多快、省多少钱。这篇文章想聊的就是大模型背后这群“看不见的人”到底在忙什么以及他们每天面对的那些问题是怎么从一个个小小的“算子”长成整个AI Infra体系的。1. 从“算子”说起大模型的最小积木到底是什么1.1 一个矩阵乘法背后发生了什么很多人第一次接触“算子”这个词是在PyTorch里。torch.matmul、torch.nn.LayerNorm敲一行代码就调完感觉像个黑盒。但算子Operator的本质是一次对张量Tensor的数学运算的最小封装输入几个张量经过某种计算输出几个张量。矩阵乘、卷积、归一化、激活函数、Softmax全算算子。真正有意思的是算子落在GPU上会经历什么。假设你要算一个大的矩阵乘数据先要从显存HBM搬到芯片的寄存器或共享内存里算完部分和再写回显存。搬运一块数据可能需要几十个周期而计算本身可能只需要几个周期。也就是说性能瓶颈经常不在“算”而在“搬”。所以一个成熟的算子实现要做的事情远多于“按公式算出来”把数据切成适合线程块处理的Tile保证访存局部性用向量化加载指令一次读入更多数据让连续线程访问连续地址避免内存冲突在共享内存做数据复用减少重复从显存读数据这些细节听起来琐碎但这才是GPU上算子优化的真实工作。用过CUDA写算子的人都知道一个简单的向量加法跑得慢往往不是因为你不会用加号而是因为内存访问模式不理想。换一种线程映射方式可能快两倍。同类事情做得多了你就会理解为什么那些做算子库的人桌上永远摆着一本“性能优化白皮书”。1.2 手写算子为什么让人又爱又恨我在不同团队见过很多刚接触算子开发的同学第一反应是“这有什么难的照着公式翻译成循环就行”。真写一版跑起来之后往往会被性能数据泼一盆冷水手写的版本比框架自带的慢5倍以上。问题几乎都出在几个经典地方。首先是线程束发散。GPU以32个线程为一束warp执行指令如果同一个warp里的线程走了不同分支那么每个分支都会被串行执行一遍代价直接翻倍。其次是共享内存的bank冲突。共享内存被分成多个bank当同一warp的多个线程访问同一个bank的不同地址时硬件会进行多次分发性能直线下降。再有就是最容易被忽略的启动开销。一个kernel的启动大约需要几微秒如果计算本身只有几微秒那启动开销就占了半壁江山。这也是为什么算子融合如此重要。我有个习惯不管写什么算子先写一个最简单的朴素版本跑通逻辑。然后对着性能分析报告一步步优化而不是直接一上来就写那种高度精巧的实现。因为精巧版本的每一行代码都包含了多层假设一旦出问题排查成本会高到你怀疑人生。先跑对的版本再跑快的版本这个顺序能帮你省掉大量调试时间。1.3 算子融合把两趟活并成一趟融合Fusion大概是算子领域最出名的思路了做法是把多个连续算子合并成一个kernel减少反复读写显存的次数。比如矩阵乘法后面通常跟一个偏置加法和一个激活函数如果不融合数据要写回显存再读出来多走好几趟。融合之后中间结果只存在于寄存器或共享内存里直接进入下一步计算访问代价小了一个数量级。这里有个更极致的例子。Attention计算里传统实现要把QK^T的结果写回显存Softmax之后再读出来再和V相乘中间矩阵的大小跟序列长度的平方成正比。序列一长读写量就非常可怕。后来出现了一批“Flash”风格的融合算子把QK^T、Softmax、PV三个阶段全部在片上完成只在最外层按块遍历整个IO量从O(n²)降到了O(n)。这不是靠改硬件做到的纯粹是软件层面的算子融合设计。理解融合的价值可以类比成做饭以前每处理完一道工序就把半成品端回仓库下一道工序再取出来。融合之后你直接在灶台边完成了所有步骤只把最后的成品端上桌。省下来的“端盘子”时间在GPU上就是省了几百上千次显存往返。这也是为什么我们今天看到的主流训练框架都会在计算图层面自动寻找可以融合的相邻算子。2. AI Infra的真实疆域四张网织出大模型的底座2.1 训练框架把千卡集群调度成“一台机器”当模型只有十几亿参数时单张卡勉强还能跑虽然慢但至少不需要太复杂的技术。可到了千亿参数规模显存放不下算力不够用数据太多跑不完问题突然从“算法”变成了“系统”。这时候AI Infra的第一大块工作就出现了分布式训练。它要回答的核心问题是——把一个模型拆到一千张卡上如何保证它们协同得像一台机器一样。并行策略各有各的切法数据并行DP每张卡都有一份完整模型各自处理不同的数据批次定期同步梯度张量并行TP把一层里的矩阵按维度切开分散到多张卡上共同计算流水并行PP按层把网络切成多段每张卡负责其中一段数据像流水线一样流过序列并行SP在序列长度这个维度上进行切分解决长序列的显存压力这几种策略不是互斥的通常要组合使用。比如常见的3D并行就是把数据并行、张量并行、流水并行结合起来。每种并行对应不同的通信模式对集群拓扑的敏感程度也不一样。张量并行每个Transformer层都要做多次AllReduce对节点内通信带宽要求极高流水并行只在小批量边界通信对跨节点带宽相对温和。调并行配置本质上就是在模型大小、显存上限、通信带宽和可扩展性之间找平衡点。我见过不止一个团队模型结构没变损失函数没动只是把并行方案从“数据并行Tensor并行”改成“换成流水并行为主、张量并行为辅”同样的千卡集群吞吐直接涨了近一倍。所以说训练框架的优化空间往往比调几个超参数来得更猛烈。2.2 推理引擎把几十GB的模型塞进一张卡训练解决的是“能不能训出来”推理解决的是“能不能用得起”。模型训练完之后总要上线服务而大模型推理和传统推理差别巨大。不是说模型部署上去就能跑关键是代价。最突出的一点是KV Cache。生成式模型是逐token输出的每生成一个字都要把之前所有token的Key和Value缓存起来复用。序列一长缓存大小就迅速膨胀。有人统计过一个百亿参数模型在长上下文场景下KV Cache占用的显存比模型权重还多。推理引擎的很多工作其实都在跟这个缓存较劲。主流推理引擎通常要解决四件事。第一是显存管理KV Cache怎么划分、怎么复用、怎么换出避免碎片化浪费第二是并发策略连续批处理continuous batching让不同请求在同一个模型副本上交错执行利用率大幅提升第三是量化压缩把权重从FP16压到INT8甚至更低精度用精度换显存同时尽量让损失降到极小第四是张量并行推理在一张卡放不下时把模型切到多张卡上但推理的通信时延又比训练更敏感需要更精细的调度。很多人以为部署模型很简单把模型文件拷到GPU上起个HTTP服务就完了。真上线处理百万级请求时你才会发现需要关心显存会不会爆、长请求和短请求怎么插队、并发太高会不会互相拖慢。推理引擎里的每一块优化最后都会直接体现为用户看到的首字时延和生成速度。AI Infra的人在这个环节干的活本质上是在“模型效果不变”的约束下把资源成本打到最低。2.3 资源调度显存、队列和故障恢复的平衡术再往下是资源调度层。如果说训练框架关心“一张模型怎么在千卡上跑”那么调度平台关心的问题是“几百个任务怎么在一大片GPU上排队”。真实的数据中心里不可能只跑一个任务。不同团队提交的微调任务、预训练任务、推理任务挤在一起这时候调度器要决定谁先跑、跑多少个节点、显存怎么隔离、任务挂了怎么自动恢复。看似平平无奇做起来却到处是决策要不要允许抢占低优先级任务能否让位给高优先级任务显存和计算怎么绑定调度只算显存内存碎片会导致计算闲置只算计算会导致显存爆掉一个任务崩溃后应该立刻重启还是等待人工介入多个小任务挤在节点上任务之间会不会互相抢占带宽导致性能震荡这些问题不像写算子那样有明确的对错更多是权衡。但我见过太多项目卡是够的任务就是跑不起来原因往往不是GPU坏了而是调度策略太要么太保守要么太激进。调度层的稳定性直接决定了整个集群的实际利用率。在这个层面AI Infra更像是在“经营一家酒店”有人负责房间分配有人负责打扫清洁任务回收与资源释放有人负责处理客人投诉故障恢复。每一项看起来不轰轰烈烈但任何一项掉链子前台服务质量立刻崩掉。2.4 数据管道GPU最怕的“等菜”最后一个经常被忽略的大块是数据管道。训练一个模型不光要有GPU还要有数据源源不断地送进来。很多人一开始以为这就是“写个dataloader”等真上了多机多卡才发现问题多得很。首先原始数据往往保存在分布式文件系统里读取路径很长。每个epoch都要重新采集、解码、增强这些操作如果落在CPU上而GPU算得又很快那么GPU就会有大量时间在空转等待数据。你打开nvidia-smi看GPU利用率90%但算的全部是上一批数据新数据还没准备好——这种情况比比皆是。其次不同样本的长度差异会导致计算图不断变化。如果直接按原始长度塞进训练短样本和长样本在同一个batch里短的会浪费大量padding位。于是数据管道要把样本按长度分桶再动态padding让同一个batch内部尽可能齐整。这个优化有时能把有效吞吐提升20%以上。做数据管道的人常常自嘲是“数据包工头”但这项工作恰恰是整个训练性能的地基。GPU利用率再高如果喂进去的数据分布不对模型训练效果也会出问题。先保证数据“不断供”再优化数据“质量”这两句话是每个训练团队都应该挂在墙上的。3. 这群人一天到底在忙什么一次性能攻坚的完整切片3.1 起手式让profiling数据先开口AI Infra工程师接到最多的问题是“我的训练好慢帮我看看为什么”新手常犯的错误是上来就猜——是不是数据增强太慢了是不是梯度爆炸了猜的代价很高因为理由千千万真正的原因往往只有一个。有经验的工程师第一件事永远是做剖析profiling。用性能分析工具抓出一条训练步里每个阶段的耗时占比数据加载用了多少前向计算用了多少反向用了多少梯度同步用了多少。一张时间线图出来瓶颈往往一目了然。有一个很实用的概念叫Roofline模型把硬件能力当作“天花板”横轴是计算强度每字节访存能做多少次浮点运算纵轴是性能。如果任务落在“带宽受限”区域再怎么优化kernel计算也没用应该去减少显存访问如果落在“算力受限”区域那才值得去抠算子内部的计算效率。用这个视角看问题能避免很多无用功。我记得有一次替某开发团队定位训练慢的问题他们怀疑是自己的自定义算子写崩了。结果profile之后发现自定义算子只占总时长4%真正的大头是分布式通信——跨节点梯度同步占了37%。这个结果让所有人都很意外。所以我的习惯很固定先让数据开口再判断方向绝不凭感觉拍脑袋。3.2 定位为什么千卡集群跑不出理论算力找到方向之后更复杂的排查才刚开始。那次通信占37%的问题从方向到具体原因还隔着好几层我们逐步拉网第一层用通信测试工具给各节点间的AllReduce打点发现部分“节点对”的带宽明显低于理论值。奇怪的是这些节点网卡型号完全一样理论上应该没有任何区别。第二层检查网络拓扑。训练任务的节点分配是调度器随机给的有些任务的两个节点之间隔着好几层交换机有些则在同一个机柜内。跨交换机的通信走得远、冲突多带宽自然上不去。问题由此变成为什么集群没有把通信密集的并行策略约束在同一机柜内第三层检查驱动和固件版本。部分节点的网卡固件版本落后导致某些硬件特性没有开启全链路带宽一路打折。这种问题最隐蔽因为从逻辑上看一切正常只有压测才能暴露。最终修复不是改代码而是改分配策略让张量并行这种通信密集的组必须落在同一个物理机柜内同时升级了网卡固件。两周后再看整体训练吞吐提升了47%。这次经历让我意识到分布式训练的性能问题一半藏在软件里另一半藏在物理环境里。环境信息和拓扑信息如果不纳入考虑再好的并行策略也发挥不出来。3.3 解决通信重叠与计算图优化精确找到瓶颈之后优化的手段往往是多管齐下。通信占大头那就把通信和计算重叠起来有大量小的通信消息那就把消息打包成大块再发计算和通信互相等待那就调整分桶策略让梯度同步的节奏隐藏在反向计算间隙。这个思路放到实际场景里是这样操作的。PyTorch的分布式训练通常把梯度按bucket分组每个bucket算完就立即发起通信而不是等所有梯度都算完。这样通信和计算可以有一部分并行GPU在跑反向传播的同时网卡也在搬运已经算好的梯度。听起来简单但bucket大小的设置很有讲究设置太小通信消息碎、消息数量多网络开销反而高设置太大计算的尾部会长时间等待最后一个bucket的通信完成。这个参数的调优典型地属于“不试几次你根本不知道最优曲线在哪”的问题。另一个常用手段是通信算子融合。把一个集合通信操作里的多个不同张量打包成一个大的发送缓冲区让网卡一次发送更多数据能显著提升小消息场景的吞吐。这个优化的效果在千卡规模下尤其明显因为消息数量会随卡数增长而单条消息大小可能不变。做完这些优化之后一定要重新做性能剖析二次验证。很多时候第一次优化引入新瓶颈只是因为指标看起来好了实际上更差。我一般会设定目标基线只有通过profile确认整体时间线已经达到预期才允许提交上线。3.4 收尾性能基线应该写进文档性能攻坚的最后一步往往不是改代码而是把经验固化下来。我们团队有个规矩每次定位到瓶颈并优化完毕都会倒推几个问题最初的性能基线是多少优化后是多少瓶颈判断依据是什么下次遇到类似问题第一步该看哪里把这些写进文档的最大价值不是让别人照着做而是避免同一个坑被反复踩。算子优化、通信优化、调度优化每一项都依赖当时的硬件拓扑和软件版本。三个月后再有人报类似的问题翻出文档先确认环境是否一致再按旧结论排查能省掉大量重复定位的时间。我自己有个很深的感受AI Infra的工作成果大多看不见但它的积累实际上非常依赖文档化和长期复盘。一次成功的优化如果只落在口头上过几个月就归零如果变成可复现的基线和方法论它就是团队资产。4. 为什么“算子”会演化成“AI Infra”复杂度飙升的逻辑4.1 单卡时代算子只是本地函数往回看几年模型还在单卡能跑的阶段时所谓“AI底层”确实就是算子。写一个算子保证它算得快、算得准就能满足绝大多数需求。那时候大家讨论的是某个kernel的指令级优化是共享内存用多少字节是能不能减少一次寄存器溢出。整个优化空间是局部的也是可见的——你改一行代码性能曲线立刻发生变化。那时很多做算法的人也顺手写写算子感觉“底层优化不过如此”。单卡场景下的系统复杂度被掩盖了因为一切都在一块芯片里完成不需要考虑跨设备通信、资源竞争、任务调度。大家能专心研究模型结构本身。4.2 多卡时代通信成了第二维度当模型大到单卡放不下分布式成为必须情况就变了。同一份代码从单卡搬到多卡突然多了一个维度通信。每一步训练前向计算在各张卡上是并行的但梯度同步需要把所有卡的结果汇总。通信的开销完全取决于网络拓扑、消息大小和并行策略而这三个因素在单卡时代根本不存在。以最经典的AllReduce为例。朴素做法是选一张卡收集所有梯度算完再分发给所有人但这样的话中心节点会成为瓶颈。于是有了环形AllReduce把数据分成N块每张卡只从相邻卡拿一块同时把自己手里的另一块传给下一个节点所有卡并行收发。这样一来总通信时间不再随卡数线性增长而是基本只跟每块数据的大小有关。这个设计简直像是接力传纸条的艺术。但通信维度的加入意味着同样的模型同样的卡数你用不同的并行切法性能可能是天壤之别。这也直接催生了一个新岗位AI Infra工程师他们要专门处理“怎么拆、怎么放、怎么通信”这些算法同学通常不太关心的系统问题。4.3 集群时代故障才是常态规模继续往上走更极端的事情开始发生。一个训练任务占用几千张卡哪怕每张卡的稳定运行率高达99.9%几千张卡同时运行一天至少有一张卡出故障的概率几乎是必然的。也就是说大规模训练里“容错”不再是一个加分项而是基本生存条件。于是Infra团队要做的不只是让训练跑得快更要让训练“跑了很久之后突然挂掉还能接得上”。这就需要周期性的模型参数checkpoint落盘需要任务调度器自动感知节点故障并把任务迁移到健康节点需要分布式存储保证checkpoint能快速写入和恢复。一旦这些底层能力缺失再贵的集群也只是堆废铁。这部分工作观众基本看不到因为一切顺利时它就是透明的。可一旦某一天训练任务跑了三周在第22天的时候一个节点掉线如果没有好的容错恢复机制前21天算力全部报废。你立刻就能理解为什么Infra工程师总把“checkpoint写得好”挂在嘴边——这是用一次次的深夜故障换来的教训。4.4 巨大的价值锚点提速30%意味着什么很多人不理解AI Infra的价值因为在模型开发流程里它是一个纯粹的成本中心不产生新模型能力不提高模型效果还总是要钱买卡、买带宽、买存储。但换个角度算一笔账就清楚了一个千卡规模的训练任务每跑一天的电费加折旧都是几十万量级。如果Infra能把训练效率从50%提到80%相当于同样的任务可以缩短将近一半时间省下的钱远超Infra团队的工资预算。更关键的是训练周期缩短意味着迭代速度加快。别人一周试两个版本的模型你可以试四个版本在效果竞赛中这种隐性优势会直接转化为模型能力优势。所以越大的团队越舍得投入Infra他们心里清楚每一笔看起来花在“工程”上的钱最终都会以更低的成本、更快的速度从产品里赚回来。理解了这个逻辑你会明白从“算子”到“AI Infra”的演变不是某个团队的选择而是规模倒逼的必然。当你的世界从一张卡变成一万张卡原本的“本地函数”问题会逐渐长大成“分布式系统”“容错调度”“存储IO”等一整张问题网。它是模型长大的影子也是行业GPU化之后避不开的重心。5. 想入行AI Infra基本功、案例和避坑心得5.1 五门必修课不算冷门的冷门经常有人私信问我转行做AI Infra需要准备什么。我的答案通常很直接跟市面上流行的“AI速成课”完全不同。以下五样是绕不开的基本功C与计算机体系结构。CPU缓存、内存模型、指令流水线不理解清楚看GPU优化文档会像看天书GPU编程。CUDA或者至少一种GPU编程语言要会写简单的kernel并理解线程块、共享内存、寄存器这些核心概念分布式系统基础。包括一致性、容错、调度、网络通信不需要精通全部但起码要知道AllReduce为什么是对的深度学习框架的源码阅读能力。不要求看懂每一行但要能跟踪一个算子的完整调用链路Linux性能和网络工具的使用。会看带宽、时延、IO占用能自己动手写压测脚本这些技能看起来又硬又杂但正因为门槛高行业里真正愿意系统学的人并不多。反而给了认真准备的人机会。我的建议是不要一上来就追最炫的论文先去训练框架里找一个简单的算子从源码层面彻底弄懂它是怎么执行的。读完一个算子的完整链路你对整个体系的感知会强过看十篇综述。5.2 一个小案例把归一化层从38us干到20us说一个我亲自做过的小优化用来展示这个方法论的实用程度。当时有个模型的热点函数是某个归一化算子单次耗时38微秒不算特别长但它在一个训练步里被调用了上千次累加起来相当可观。第一版实现很朴素就是两趟遍历第一趟算均值和方差第二趟做归一化。这在逻辑上没有问题但对显存带宽很不友好因为数据要被读完写一遍再读一遍。于是第一件优化就是改成单趟遍历使用Welford算法在线更新均值和方差把数据访问量减少了一半。第二件是调整内存访问模式。原来的实现是每个线程处理一个标量元素访问不连续导致带宽利用率偏低。改成每个线程连续读取多个元素并用向量化加载指令一次性搬入多个数据点吞吐立刻提升了30%左右。第三件是消除共享内存的bank冲突。重排了线程到数据的映射保证同一束线程访问共享内存时不会撞到同一个bank。到这一步耗时已经压到22微秒。再做了一轮寄存器级别的细节调整最终稳定在20微秒。整个过程没有引入什么高深算法就是三步减少访存、向量化访问、消除冲突。每一处优化都用profile数据验证过没有一步是靠猜的。5.3 避坑三则先测再调、环境一致、写个checkpoint最后分享三条我觉得新人最容易踩的坑都是我亲眼见过或者自己踩过的。第一先测再调永远不要凭感觉优化。性能问题最反直觉的一点是很多你以为的瓶颈根本不是瓶颈。我见过有人花一周优化一个数据加载函数结果profile之后发现它只占2%的时间真正的瓶颈是GPU利用率不足。记住一个原则在没有profile数据之前你没有资格谈优化。第二环境和复现的一致性非常重要。同一个模型在单机环境跑和在多机环境下跑性能特征完全不一样。驱动版本、网卡固件、交换机配置、拓扑布局任何一点不同都可能让结果出现数倍的差异。所有性能记录都要连带环境信息一起记录否则你无法判断优化结论是否还能复用。第三checkpoint不是可选功能。很多刚做分布式训练的人觉得隔几步保存一下权重太浪费时间和磁盘于是把保存间隔拉得很长。直到某天训练挂掉发现好几天的算力全白费了才会明白checkpoint是整个训练系统的安全网。定期保存、异地冗余、能快速恢复这才是一个训练任务真正“可交付”的标准。我自己在项目里养成了一个习惯从早期做算子优化到后来做集群级调优一直坚持每次改动前后都做完整的性能基线对比并把环境信息写进表格。这个习惯曾经在好几周后帮我快速定位到一次性能回退。别嫌麻烦它会救你于水火。最后再分享一点个人体会。AI Infra这份工作最大的特点是“做成了没人夸做差了全怪你”。因为一切顺利时整个系统就是透明的一旦出问题所有人都会来找你。但恰恰是这份默默无闻让我觉得它很有价值——你不需要出现在聚光灯下只需要让那些“不可能跑起来”的训练真正跑起来让那些“慢得没法用”的模型变得能服务用户。说实话每次看到一条本来要跑两周的训练任务因为自己调了并行方案变成五天完成那种踏实的成就感比发一篇论文还来得更实在。