从算子优化到AI Infra：大模型工程链路的底层密码

先说结论这篇文章我想聊聊这些年在大模型工程链路里真正“搬砖”的岗位——从最底层的算子优化到上层一点的分布式训练框架再到整个AI InfraAI基础设施的搭建与运维。你可能看过不少讲Transformer架构、讲训练技巧的文章但大模型真正能跑起来背后靠的是另一个世界一群不怎么发论文、不怎么做宣传的人日复一日地和显存、带宽、算子死磕。这篇文章适合三类人看一是想入行AI Infra方向但还没想清楚具体做什么的开发者二是已经在做模型训练或推理、想搞明白上下游在忙什么的算法工程师三是对“大模型公司到底在招什么样的人”感到好奇的同学。我会从算子这个最不起眼的单元讲起一路聊到分布式训练和推理部署最后聊一聊这个岗位的真实日常和技能要求。1. 算子大模型的“原子操作”1.1 什么是算子算子这个词翻译自Operator在深度学习里指的就是对一个张量做的数学操作。加法、乘法、矩阵乘、卷积、归一化、Softmax、LayerNorm这些都是算子。任何一个神经网络模型无论多大、多复杂从计算图的角度看都是由成百上千个算子组成的。大模型也不例外。Transformer结构里的核心计算拆到最底层就是矩阵乘、注意力计算里的QKV投影、Softmax、LayerNorm、残差连接。模型的训练和推理过程本质上就是让这些算子在前向传播和反向传播中一遍又一遍地执行。但这里有个非常关键的细节算子不是自动就能跑得快的。同一个矩阵乘用一个朴素的实现去跑和用高度优化的算子库去跑性能差距可能是几十倍甚至上百倍。这是什么概念呢同样的训练任务有人三天跑完有人可能得跑三个月。我举一个生活化的例子算一道三位数乘法你按小学列竖式的方法笔算和一个心算高手用各种技巧直接口算结果是一样的但速度差得很远。算子也是这样数学上都是同一个操作但怎么组织数据访问、怎么利用硬件的并行计算单元、怎么减少中间结果的读写次数这些工程细节决定了最终能跑多快。1.2 计算密集和访存密集优化算子的第一步是搞清楚这个算子到底属于哪一类。从硬件利用的角度算子大致分成两种类型计算密集型和访存密集型。计算密集型的典型代表是矩阵乘。一个形状为[M, K]的矩阵乘以[K, N]的矩阵算一遍需要做M乘N乘K次乘加运算。当维度大起来比如M、N、K都是4096那就是六千多亿次乘加。这种算子的特点是计算量大但每个数据可以被反复复用所以瓶颈在于怎么让GPU的计算单元满负荷运转。访存密集型的典型代表是LayerNorm、Softmax这类操作。它们对每个元素只是做一次或几次运算计算量本身很小但数据要先从显存里读出来算完再写回去。比如对一批数据做均值方差归一化每个元素只做加减乘除几次但数据需要从头到尾读一遍。这种算子的瓶颈在数据的搬运速度而不是计算速度。用生活化的语言说计算密集型像是一个工厂里工人很多、每个人都在做复杂加工的情况瓶颈在于工人的数量和工作效率访存密集型像是流水线上大部分时间都在搬原料真正加工的过程很短瓶颈在于传送带的速度。这两类算子的优化思路完全相反。计算密集型的算子核心优化方向是提高计算效率比如让数据在寄存器里多复用几次、减少片上存储的访问冲突访存密集型的算子核心优化方向是减少数据搬运比如合并读请求、减少中间结果的读写甚至把多个访存密集型算子合并成一个。正是这两类算子的差异衍生出了另一个非常重要的概念算子融合。1.3 为什么算子融合能加速算子融合指的是把多个算子合并成一个算子来执行。听起来简单但效果往往非常惊人。最典型的例子是LayerNorm的融合。标准的LayerNorm实现里有好几个步骤计算均值、计算方差、归一化、缩放和平移。如果不做融合每一步都要把整个张量从显存里读一遍写一遍。假设模型的隐藏层维度是8192每一步读写一次就是16KB数据多几个步骤就是几十KB的无效搬运。融合之后的LayerNorm把多个步骤放在GPU的一次内核启动里完成数据只要从显存读一次、在寄存器或高速缓存里完成所有计算、写回一次就够了。再比如QKV投影的融合。Transformer注意力机制里的Query、Key、Value计算数学上其实是同一个输入矩阵乘以三个不同的权重矩阵。很多框架在初期会把这三个矩阵乘拆成三个独立算子于是数据要被读三次、写三次。而融合的方案是把三个权重矩阵拼成一个大的权重矩阵一次矩阵乘完成三组投影。这省掉的不只是计算更重要的是省掉了一遍又一遍的中间结果读写。还有一个经典案例是把残差连接、归一化、Dropout这些算子融合进注意力或前馈计算内部。Transformer块里有大量的“先加残差再归一化再做线性变换”的结构如果每个操作都单独执行一个Transformer层可能有十几次甚至二十几次的显存读写。而把这些操作融合进相邻的计算单元之后层与层之间的中间结果开销能大幅压缩。不过需要提醒的是算子融合也不是无脑做就有效。融合之后的算子如果逻辑过于复杂独立的核心代码超过一定规模反而可能因为寄存器溢出或者指令缓存命中率下降而变慢。实际工程中融合的边界选在哪里是需要靠profile数据说话不是拍脑袋定的。2. 从单一算子到分布式训练2.1 单卡装不下分布式来拯救算子层面的优化解决的是“单个GPU上跑得快不快”的问题。但今天的大模型训练就算是能用上最新的超大显存GPU也远远装不下一个千亿级参数的模型。我来算一笔账一个百亿参数约10B的模型光是把参数用FP16精度存下来就需要大约20GB显存。这还没完训练的优化器状态比如常用的Adam优化器需要保存参数的一阶动量、二阶动量、以及参数本身的副本这部分开销通常是模型参数量的好几倍。再加上前向传播的中间激活值、反向传播的梯度一个百亿参数的模型训练至少要上百GB的显存才勉强转得动。单卡装不下就得把计算和存储分摊到多张卡上。这就是分布式训练要解决的核心问题。但分布式训练不是简单地把模型切成几块分别放上去就行关键的问题是怎么切以及切完之后各张卡之间怎么同步数据。2.2 数据并行、张量并行、流水线并行目前大模型训练中最常用的并行策略有三类数据并行、张量并行、流水线并行。数据并行是最容易理解的一种每张卡上都放一份完整的模型副本然后各拿一批不同的数据并行计算最后把梯度汇总同步。理论上用N张卡做数据并行训练吞吐量可以接近提高到N倍。但在模型大到单卡放不下的时候这个方法就失效了。张量并行解决的就是模型放不下问题把同一层的权重矩阵按行或按列切开分到多张卡上每张卡只负责维护一部分权重同时计算时通过通信把结果拼接起来。最典型的例子是矩阵乘的列并行把权重矩阵按列分成两块分给两张卡输入数据同时发给两张卡每张卡各自算出一部分输出然后拼起来得到完整的输出。这种方式在每个Transformer层都会引入通信开销所以张量并行通常只在单机多卡的场景下使用因为机内GPU之间的通信带宽高、延迟低。流水线并行是把模型按层切段每一段放在不同的GPU上。GPU 0负责前几层GPU 1负责中间几层GPU 2负责最后几层。数据从前往后流像工厂流水线一样。这种方案的问题是会产生“气泡”——在某一个时间段里前面的卡在算第一段的同时后面的卡可能在等待。为了减少气泡工程上会把一条数据拆成多个微批次交错地灌进流水线尽可能让每张卡都处于忙碌状态。这三类并行策略不是互斥的实际的大规模训练通常是把它们组合起来使用。比如一个训练任务先按张量并行处理单机多卡的计算拆分再跨机器用流水线并行把模型层分段最后对每个分段做数据并行来充分利用集群算力。2.3 通信开销分布式训练的隐形瓶颈部署过分布式训练的人都知道最头疼的不是显存够不够而是通信。分布式训练里数据并行需要同步梯度张量并行需要互相传输计算结果这两件事都要通过GPU之间的通信完成。而通信恰恰是最慢的一环。GPU计算用纳秒到微秒级别衡量而跨节点的网络通信一毫秒到几十毫秒都是正常的。当模型规模大到一定程度通信时间甚至可能超过计算时间。举个例子一个七十亿参数的模型做数据并行训练每张卡上每个训练步骤产生的梯度数据就有几十GB。这些梯度要在所有卡之间做一次全量汇聚即便使用目前服务器集群常见的400Gbps高速网络也需要好几秒的纯传输时间。如果优化器还依赖梯度做同步更新那这个通信就是每个训练步骤都无法绕开的固定开销。所以做分布式训练的时候通信拓扑、通信量与计算量的比例、通信如何与计算重叠这些因素的重要性不亚于模型本身设计。现在主流的分布式训练框架都尽量把通信放到计算的同时进行让GPU在等待通信结果的时候先算不依赖该结果的后续部分从而把通信开销“藏”在计算里。这一层优化做得好不好直接影响几千张卡跑起来到底能有多高的利用率。3. 推理部署模型是造出来了关键是怎么“用”起来3.1 量化和推理加速训练要花很多资源推理其实一样需要大量工程投入。差别的只是方向训练看重吞吐量推理看重延迟和单请求成本。推理环节的第一个关键技术是量化。训练好的模型参数默认是一个FP16或FP32的浮点数表示每个参数占2到4字节的显存。量化把这个精度压缩到INT8甚至更低比如INT4。参数占的空间大幅缩小计算速度也会提升——因为整数计算在GPU上通常比浮点计算吞吐更高。以现在常用的W8A8量化为例权重用8位整数表示激活也就是每层的输入也用8位整数表示。相比FP16模型显存占用量直接减半计算密集的场景下速度提升也很明显。更激进的做法是W4A16权重4位、激活16位显存占用更低但对模型精度影响更大需要做精细的校准和适配。量化最讲究的是怎么选择校准数据、以什么标准决定浮点往整数的映射范围。选得不好模型在基准评测集上可能掉点零点几个点但在真实场景里出现明显的质量下降。这个环节需要反复实验从“能用”到“好用”往往要花好几周时间调参。3.2 推理里的状态管理KV Cache推理过程中的另一个关键优化是KV Cache。Transformer模型在生成每个新词的时候需要重新计算注意力层里所有历史词对应的键和值。如果不做缓存每次生成都从头算一遍全局的注意力那时间复杂度会随着生成长度的增加而平方级上升。KV Cache就是把这些计算过的键和值存下来下一次生成时直接复用。KV Cache听起来简单真正做起来很麻烦。它需要多少显存举一个实际场景一个七十亿参数的模型输入输出上下文长度设定在64K每个序列的KV Cache大小可以达到几百MB甚至更高。并发请求多的时候KV Cache的总占用很快就把显存吃满。这就是为什么现在很多推理框架在做PagedAttention这类优化——把KV Cache像操作系统的内存分页一样管理起来按需分配、动态换入换出。这样每个序列占用的显存不是固定的而是真正需要多少才分配多少。在一个多用户并发的推理服务里这种方案能把可并发数提高好几倍这是线上服务商用化的关键。3.3 连续批处理推理服务的另一个核心概念是连续批处理。我见过的很多刚接触服务部署的同学对推理批处理的认知还停留在“同时来很多请求攒一攒一起算”。这种方式的问题在于在线服务的请求到达时间是随机的如果一次只处理一个请求GPU利用率很低如果攒够一批再开工延迟又会变大。连续批处理的做法是不需要把一批请求凑齐再开算而是任何一个请求到了只要GPU还有空闲的计算资源立刻把请求加入计算并且当某个请求生成完结束立即把计算资源让给新请求。这相当于把“排队上车”变成了“随到随走”GPU的利用率完全被填满。这种方案对调度层提出了很高的要求要能随时中断当前的计算、插入新的序列、再做恢复。推理框架在底层通常需要真正实现GPU计算任务的精细调度这已经非常接近操作系统的进程调度。4. AI Infra比模型本身更值钱的那一层4.1 AI Infra到底包含什么讲清楚了算子优化和分布式训练的底层逻辑就能理解AI Infra这个更上层的概念了。如果说算子是零件分布式训练是装配生产线那AI Infra就是整座工厂的规划与运维体系。它包含了训练框架、推理服务框架、资源调度系统、数据流水线、监控系统、容错机制。很多公司里AI Infra团队负责的不只是让模型能跑起来还包括让它稳定、高效、节省成本地跑起来。举几个实际场景。一个团队要跑一次上百卡规模的训练任务AI Infra的人要做的事情包括帮助训练脚本选择合适的并行策略并预估显存开销、配置好训练框架的参数、为服务搭建自动扩缩容的推理平台、在任务失败时自动重启并恢复训练进度而不是从头再来、监控全链路的性能指标并定位瓶颈。这些工作在项目初期往往被忽视因为一个模型从零训练到能跑出结果算法团队短期内更关心的是模型效果。但等到模型训练开始规模化集群利用率上不去、任务频繁失败、显存OOM不断AI Infra的价值就完全体现出来了。4.2 集群管理从单任务到规模化调度大模型训练和推理的一个显著特征是资源需求和任务时长都很极端。一个训练任务可能要在上百张GPU卡上跑几天甚至几周。这个过程中任何一张卡出现故障任务都可能中断。如果没有重要的冗余和容错机制重来一次的成本是几十万投入打水漂。AI Infra里的集群管理核心职责之一就是做任务调度和容错恢复。在系统层面需要用调度器为不同团队的任务分配GPU资源、跟踪资源占用、实施优先级策略。在训练框架层面需要支持定期把模型参数和优化器状态做快照任务挂了之后不用从头训起而是从最近的快照继续跑。行业里有一句经验训练上千张卡的集群如果没有任何容错设计遇到故障的概率几乎是百分之百。一个星期的训练任务里碰到几次卡故障、网络抖动、驱动问题都很正常。能不能自动恢复、恢复的粒度有多细直接决定了训练项目的效率。另一个趋势是异构资源管理。GPU型号越来越多、显存大小不一、计算能力不同如何把不同批次采购的异构卡统一纳入调度池既不浪费性能又保证任务稳定完成是集群管理长期要解决的问题。4.3 为什么AI Infra会成为热门岗位从就业市场的角度看AI Infra可以说是目前AI领域里需求最旺盛的方向之一。原因不复杂模型开源越做越多架构越来越收敛真正把开源模型用好、用稳定、用得省钱的团队核心竞争力恰恰在基础设施层面。同样的模型别人跑一次推理一块钱你的团队跑一次一块二这就是差距。同样的训练任务别人一个礼拜跑完你的团队要跑两个礼拜这也是差距。在模型质量差不多的前提下成本、速度和稳定性是决定一个AI产品能不能活下去的致命因素而这些全部落在AI Infra团队的肩上。5. 从我的角度聊聊这个岗位的一天和踩过的坑5.1 一个AI Infra工程师的日常我自己做这个方向有几年了日常工作大致是以下几种。一是写和改CUDA算子比如某个模型里新加了一个自定义激活函数需要写一个融合算子让它不要在GPU内存里来来回回搬运数据。二是调分布式训练参数比如训练任务跑得慢了需要分析到底是数据加载卡住了还是通信瓶颈还是某个算子本身就写得差。三是做推理服务优化比如上线一个七B模型在保证延迟不超过50毫秒的前提下把单机并发数从20提到60。这些工作听起来没有一个跟“模型训练”直接相关但没有这些工作大模型训练就不可能规模化。一个模型论文的诞生在业界的工程链路里可能意味着几个月的算子开发、框架调优和服务打磨这些都不太会出现在论文作者栏里但它们是模型真正可用的前提。5.2 一些实操经验和教训分享几个我在实际工作中踩过的坑这些在很多内部文档里不一定写得很详细。第一个坑是低估了通信对训练的影响。我早期参与过一次大规模训练单卡计算利用率高得吓人但整体训练速度就是提不上去。后来用性能分析工具一查才发现每步训练里卡之间的梯度同步时间比计算时间还长。解决办法也很俗套调整并行策略把通信量大的张量并行改成流水线并行让通信彻底错开计算时间训练速度直接翻倍。第二个坑是盲目追求算子融合。有一次我把一个Transformer层的十几个算子全部揉到一个内核里代码写了一百多行结果实测性能不但没变快反而变慢了。原因很反直觉融合太深之后中间结果全部堆在寄存器里寄存器溢出导致大量数据被挤到本地内存性能就塌了。后来我把融合拆回两三个级别效果反而好。这让我意识到优化这种东西必须以profile数据为依据不能凭空想象。第三个坑是KV Cache的占用。早前做推理服务的时候我按一个序列的最大长度预分配了KV Cache显存结果并发一上来显存直接爆掉。后来改成按需分配以页为单位动态分配并发数翻了好几倍。这个经验其实在模型推理优化里已经算老生长谈但每次问下来还是有不少人没真正理解“KV Cache是动态增长的”这句话意味着什么。5.3 想入行需要准备哪些技能如果你对这个方向感兴趣我个人的建议是底层基础一定要打牢。CUDA和并行计算最好能自己动手写几个算子分布式系统的基本概念比如通信、一致性、容错这些也要有扎实的理解。Python只是上层工具真正解决问题的时候C、CUDA、性能分析工具这些平时不太性感的技能才是压舱石。另外一定要学会使用profiler工具。很多人优化性能靠猜哪慢了就瞎改哪里这种习惯非常致命。在AI Infra这个方向使用性能分析工具定位瓶颈是基本功。把profile数据读明白你会知道GPU算力利用率是多少、内存带宽用了多少、通信花了多少时间然后才能做有针对性的优化。6. 写在后面的个人体会最后说点个人的真实感受。AI Infra这个方向在行业内存在感不算高很少有普通用户知道一家AI公司的核心竞争力可能有一大半在基础设施团队。但恰恰是这些人决定了模型能不能规模化、成本能不能降下来、服务能不能稳定跑在线上。模型架构迭代得很快但算子优化、显存管理、分布式通信、容错调度这些底层工程问题会长期存在且越来越重要。说句接地气的话这条路确实辛苦要对计算机体系结构、操作系统、并行计算、分布式系统都有深入的理解工作内容也有相当一部分是枯燥的调参和排查。但几万张卡在手里能高效稳定地跑起来的那种成就感也不是发两篇论文能替代的。如果你对“做点东西让模型真正能跑起来”这件事有执念这个方向至少在未来五到十年里都是一个值得下的注。