1. 大模型背后的隐形战场:从算子到AI Infra到底在忙什么
很多人第一次听到"算子"这个词,脑子里浮现的是数学课本里的符号,或者某个高深莫测的公式。但如果你真的接触过大模型训练和推理的底层,就会发现这个词其实特别接地气——它就像是做菜时的"基本动作":切菜、翻炒、调味。每一个动作单独看都很简单,但要把几万个动作编排好,让整道菜在预定时间内端上桌,那就是另一回事了。
我在这行摸爬滚打这些年,最大的感受就是:外界看到的永远是模型又刷新了什么榜单、又支持了多长的上下文,但真正决定这些东西能不能跑起来、跑得快不快、跑得稳不稳的,是背后那群搞算子、搞编译、搞分布式、搞推理优化的人。他们不太出现在聚光灯下,但少了他们,再漂亮的模型架构也只是一堆躺在论文里的公式。
这篇文章想聊的就是这群人到底在忙什么。从最底层的算子优化,到中间的编译与图优化,再到上层的分布式训练和推理服务化,我会把这条链路上每个环节的核心问题、常见做法、踩坑经验都掰开揉碎讲一遍。不管你是刚入行的算法工程师,还是想往AI Infra方向转的后端开发,或者只是对大模型底层好奇的技术爱好者,应该都能从中找到对自己有用的东西。
先给一个整体的认知框架:算子是计算的最小单元,AI Infra是把这些单元组织起来、让它们高效运转的基础设施。两者之间的关系,有点像"砖块"和"建筑公司"——砖块质量决定建筑上限,但建筑公司决定了你能不能按时按质把楼盖起来。下面我按这条链路从上到下、再从下到上走一遍。
2. 算子层:一切性能问题的源头
2.1 算子到底是什么,为什么它这么重要
用最直白的话说,算子就是一段在特定硬件上执行特定数学运算的代码。比如矩阵乘法、卷积、归一化、激活函数,每一个都是一个算子。模型训练和推理的过程,本质上就是把这些算子按一定顺序串起来执行。
那为什么算子这么关键?因为大模型的绝大部分时间都花在少数几个算子上。以Transformer架构为例,一个典型的推理过程里,矩阵乘法和注意力相关的计算能占到总耗时的百分之七八十甚至更高。这意味着什么?意味着你只要把这几个核心算子优化好,整体性能就能有质的提升。反过来,如果某个算子实现得不好,它就会成为整个系统的瓶颈,其他部分再快也没用。
我见过不少团队,模型结构设计得很漂亮,但一上线就发现推理延迟高得离谱,排查半天发现是某个自定义算子在特定shape下触发了低效路径。这种问题在论文里永远不会提到,但在工程实践中几乎天天遇到。
2.2 算子优化的几个核心方向
算子优化不是简单地"把代码写快一点",它涉及多个层面的权衡。我把它归纳为四个方向:
第一个方向是算法层面的优化。同一个数学运算,可以用不同的算法实现,复杂度可能差一个数量级。比如注意力计算,标准实现是O(n²)的复杂度,但通过分块、稀疏化、近似等方法,可以大幅降低实际计算量。这类优化需要你对数学本身有深刻理解,不是单纯写代码能解决的。
第二个方向是内存访问优化。现代GPU的算力其实过剩,真正的瓶颈往往在内存带宽上。一个算子如果频繁地从显存读写数据,那它的实际性能会被带宽卡死。所以好的算子实现会尽量做算子融合,把多个连续操作合并成一个,减少中间结果的落盘。比如把"矩阵乘+偏置加+激活函数"融合成一个算子,中间结果一直在寄存器或共享内存里流转,不写回显存,性能提升非常明显。
第三个方向是并行策略优化。GPU有几千个核心,怎么把计算任务合理地分配到这些核心上,是个技术活。分得太细,调度开销大;分得太粗,负载不均衡。还要考虑warp级别的分歧、bank conflict这些细节。这部分经验性很强,很多调优参数没有理论最优解,只能靠实测。
第四个方向是硬件特性适配。不同代际的GPU,指令集、Tensor Core规格、内存层次结构都不一样。一个在上一代硬件上跑得飞起的算子,换到新一代可能就不行了。所以算子库需要针对不同硬件做专门适配,这也是为什么各大厂商都在维护自己的算子库。
2.3 手写算子 vs 自动生成:怎么选
在实际项目中,算子来源主要有三种:直接用现成库、手写、自动生成。这三者各有适用场景,我做个对比:
| 方式 | 优势 | 劣势 | 适用场景 |
|---|---|---|---|
| 现成库 | 稳定、经过充分测试、维护成本低 | 灵活性差、特殊shape可能不覆盖 | 标准算子、通用场景 |
| 手写 | 可针对特定场景极致优化 | 开发周期长、容易出bug、维护难 | 核心瓶颈算子、特殊硬件 |
| 自动生成 | 开发效率高、可快速覆盖多硬件 | 生成质量参差、调试困难 | 算子种类多、硬件迭代快 |
我的经验是:先用现成库,遇到瓶颈再考虑手写或自动生成。很多团队一上来就想自己写算子,结果花了几个月,性能还不如调优后的现成库。正确的做法是先做profiling,找到真正的瓶颈,再决定要不要投入人力去啃硬骨头。
注意:手写算子最大的坑不是性能,而是正确性。浮点运算的精度问题、边界条件的处理、不同shape下的行为一致性,这些都需要大量测试用例来保证。我见过因为一个边界条件没处理好,导致模型在特定输入下输出完全错误的案例。
2.4 算子开发中那些文档不会告诉你的坑
第一个坑是数值精度。GPU上的浮点运算和CPU不完全一样,尤其是用了Tensor Core之后,累加顺序的变化可能导致结果有细微差异。单次差异很小,但在深层网络里会逐层放大。所以算子开发必须做精度对齐测试,不能只看性能。
第二个坑是动态shape。很多算子在大模型场景下需要处理变长序列,shape是不固定的。如果算子只针对固定shape优化,遇到动态shape就会退化甚至报错。解决办法通常是准备多套kernel,根据实际shape动态选择。
第三个坑是多卡一致性。分布式训练时,同一个算子在不同卡上的行为必须完全一致,否则会导致梯度不同步。这个问题在单卡测试时根本发现不了,一到多卡就暴露。
3. 编译与图优化:把算子串成高效流水线
3.1 为什么需要编译层
有了好的算子,下一步就是怎么把它们组织起来。如果只是简单地按顺序调用,那中间结果的读写、kernel启动的开销会吃掉大量性能。编译层的核心任务就是在保持计算结果不变的前提下,重新组织算子的执行方式,减少开销。
打个比方:算子像是厨房里的各种工具,编译层就是那个负责安排做菜流程的人。同样是切菜、炒菜、装盘,好的流程安排能让厨师少走很多冤枉路,出菜速度自然就快。
3.2 图优化的几个关键手段
算子融合是最常见也最有效的手段。把多个小算子合并成一个大算子,减少kernel启动次数和内存读写。比如常见的"卷积+批归一化+激活"三合一,在很多框架里都是默认开启的。
内存复用是另一个重点。深度学习模型的计算图里,很多中间张量在用完之后就可以释放,后面的张量可以复用这块内存。好的内存规划算法能显著降低峰值显存占用,这对于大模型来说至关重要——显存就是硬约束,省下来就能跑更大的batch或者更长的序列。
常量折叠和死代码消除属于比较基础的优化,但在实际项目里经常能捡到不少性能。比如有些计算在推理阶段是固定的,完全可以提前算好,不用每次推理都重算。
布局转换优化是个容易被忽视的点。不同的算子可能对数据布局有不同的偏好,有的喜欢NCHW,有的喜欢NHWC。如果频繁转换布局,开销会很大。好的编译器会尽量统一布局,或者在必要时做低成本的转换。
3.3 动态图 vs 静态图:实际项目怎么选
这个问题争论了很多年,我的看法是:训练用动态图,推理用静态图,这是目前最务实的方案。
动态图的优势是调试方便、灵活,适合研究阶段快速迭代。你可以像写普通Python代码一样写模型,出问题容易定位。但它的缺点是执行效率低,因为每次都要重新解释执行。
静态图则相反,它在执行前就把整个计算图编译好了,可以做充分的优化,执行效率高。但调试起来麻烦,出错信息往往很难懂。
现在很多框架都在做"动态图写、静态图跑"的混合模式,兼顾两者优点。实际项目中,我建议训练阶段用动态图保证开发效率,部署前再转成静态图做推理优化。
3.4 编译层的实操经验
在实际做图优化的时候,有几个经验值得分享:
第一,优化要有优先级。不是所有优化都值得做。先做profiling,找到热点,集中精力优化那百分之二十的关键路径。我见过团队花大力气优化了一个只占总耗时百分之三的算子,投入产出比极低。
第二,优化要可回退。图优化有时候会引入bug,尤其是涉及数值精度的优化。所以一定要有开关,出问题能快速回退到未优化版本。
第三,要建立性能基线。每次优化前后都要有明确的性能数据对比,不能凭感觉说"好像快了一点"。建立自动化的性能测试流水线,每次代码变更都跑一遍,防止性能回退。
4. 分布式训练:把成千上万张卡拧成一股绳
4.1 分布式训练的核心挑战
单卡跑不动大模型,这是共识。但多卡不是简单地把卡插上去就行,分布式训练要解决的核心问题是:如何让多张卡高效协同,同时保证训练结果的正确性。
这里面的挑战主要有三个:通信开销、负载均衡、容错。通信开销是最直接的——卡越多,需要同步的数据越多,通信时间可能超过计算时间。负载均衡是指不同卡的计算量要尽量均匀,否则快的卡要等慢的卡。容错是指大规模训练中硬件故障是常态,必须有机制保证训练不中断。
4.2 并行策略的选择
分布式训练有几种主流并行策略,我做个对比:
| 策略 | 原理 | 优势 | 劣势 | 适用场景 |
|---|---|---|---|---|
| 数据并行 | 每卡一份完整模型,数据分片 | 实现简单、扩展性好 | 显存占用高 | 模型能放进单卡 |
| 张量并行 | 单层内切分计算 | 降低单卡显存 | 通信频繁 | 单层特别大 |
| 流水并行 | 按层切分到不同卡 | 降低单卡显存 | 有气泡、调度复杂 | 层数特别多 |
| 混合并行 | 多种策略组合 | 兼顾各方面 | 实现复杂 | 超大模型 |
实际项目中,混合并行是标配。比如常见的配置是:张量并行控制在单机内(利用高带宽的机内互联),流水并行跨机,数据并行在最外层。这样能最大化利用不同层级的带宽差异。
4.3 通信优化的实操技巧
通信是分布式训练的性能杀手,优化手段主要有几个:
梯度压缩是最直接的手段。把梯度从FP32压缩到FP16甚至INT8,通信量直接减半甚至更多。但要注意精度损失,通常需要配合误差补偿机制。
通信与计算重叠是另一个关键技巧。在计算的同时进行通信,把通信时间藏起来。这需要框架层面的支持,把计算图和通信图做细粒度的调度。
拓扑感知的通信调度也很重要。不同的卡之间带宽不一样,机内带宽远高于机间。所以通信调度要尽量让高带宽的链路承担更多通信量。
提示:分布式训练调优是个系统工程,不要指望一个参数调优就能解决问题。建议先用小规模集群验证策略,再逐步扩大规模。
4.4 大规模训练中的稳定性问题
大规模训练最怕的就是训练到一半挂了。我总结了几类常见问题:
硬件故障是最常见的,尤其是卡越多,故障概率越高。解决办法是做好checkpoint,并且支持从checkpoint快速恢复。checkpoint的频率要权衡——太频繁影响性能,太稀疏丢的进度多。
数值不稳定也很常见。大模型训练中梯度爆炸、loss突然变成NaN的情况时有发生。通常需要配合梯度裁剪、loss scaling等手段。但更重要的是要有监控,一旦发现异常能及时干预。
通信超时在跨机训练中经常遇到。网络抖动、某张卡负载过高都可能导致超时。需要设置合理的超时时间,并且有重试机制。
5. 推理服务化:让模型真正跑在生产环境
5.1 推理和训练的本质区别
很多人以为推理就是训练的前向传播,没什么区别。实际上,推理场景的约束和训练完全不同:
训练追求的是吞吐量,可以慢慢跑,只要总体效率高就行。推理追求的是延迟和吞吐的平衡,用户等不了太久,同时服务成本要可控。
训练可以接受一定的数值误差,推理对精度要求更严格,因为直接影响用户体验。训练可以批量处理,推理往往要处理单条请求,batch size可能只有1。
这些差异决定了推理优化需要一套完全不同的思路。
5.2 推理优化的核心手段
量化是推理优化最有效的手段之一。把模型权重和激活值从FP16降到INT8甚至INT4,显存占用和计算量都大幅下降。但量化会带来精度损失,需要做校准和微调来补偿。
KV Cache优化是大模型推理特有的问题。自回归生成时,每次都要用到之前所有token的key和value,如果不做缓存,计算量会随序列长度平方增长。但缓存又占显存,所以需要在缓存策略上做权衡。常见的优化包括分页管理、量化缓存、稀疏化等。
批处理调度是提升吞吐的关键。把多个请求攒成一批一起处理,能显著提升GPU利用率。但攒批会增加延迟,所以需要动态调整批大小,在延迟和吞吐之间找平衡。
投机采样是近两年比较火的技术。用一个小的草稿模型先快速生成几个token,再用大模型验证,能显著加速生成过程。这个技术的难点在于草稿模型的选择和验证策略的设计。
5.3 服务化框架的选型考量
推理服务化框架的选择,我建议从几个维度考虑:
性能是首要的,包括延迟、吞吐、显存占用。但要注意,benchmark数据只能参考,实际场景下的表现才是关键。
易用性也很重要。框架的API设计是否合理、文档是否完善、社区是否活跃,这些直接影响开发和维护成本。
扩展性决定了你能不能应对业务增长。是否支持多模型、是否支持动态扩缩容、是否支持自定义算子,这些都要提前考虑。
生态兼容性不能忽视。能不能和现有的监控、日志、调度系统集成,往往比框架本身的性能更重要。
5.4 生产环境中的那些坑
冷启动问题是推理服务常见的痛点。模型加载、显存分配、kernel编译都需要时间,导致第一个请求延迟特别高。解决办法包括预热、模型常驻、懒加载等。
显存碎片是另一个隐蔽的问题。长时间运行后,显存会出现碎片,导致明明有足够显存却分配失败。需要做好显存池化管理,或者定期重启服务。
请求分布不均也很常见。高峰期请求暴涨,低峰期资源闲置。需要做好弹性伸缩,同时要有过载保护机制,避免雪崩。
6. 常见问题与排查技巧实录
6.1 性能问题排查速查表
| 现象 | 可能原因 | 排查方法 | 解决思路 |
|---|---|---|---|
| 训练速度慢 | 通信瓶颈 | 看通信占比 | 优化并行策略、梯度压缩 |
| 推理延迟高 | 算子低效 | profiling热点 | 算子优化、融合 |
| 显存不足 | 中间结果多 | 看显存快照 | 内存复用、梯度检查点 |
| 精度下降 | 数值问题 | 逐层对比 | 精度对齐、混合精度调优 |
| 训练不稳定 | 梯度异常 | 监控梯度范数 | 梯度裁剪、loss scaling |
6.2 几个真实的排查案例
案例一:推理延迟突然翻倍。排查发现是某个算子在特定shape下走了低效路径。原因是输入序列长度刚好超过了某个阈值,触发了不同的kernel。解决办法是调整分块策略,让两个路径的性能更接近。
案例二:多卡训练速度不升反降。排查发现是通信和计算没有重叠,通信时间完全暴露。通过调整并行策略,把张量并行限制在机内,问题解决。
案例三:模型输出偶尔出现乱码。排查发现是KV Cache的边界处理有bug,在特定序列长度下会读到脏数据。这类问题很难复现,需要加详细的日志和断言。
6.3 避坑经验总结
第一,建立完善的监控体系。性能问题往往不是突然出现的,而是逐渐劣化的。有监控才能及早发现。
第二,保持可复现性。每次实验都要记录完整的配置和环境信息,否则出了问题根本没法排查。
第三,不要过早优化。先让系统跑起来,再找瓶颈优化。我见过太多团队在还没跑通的情况下就开始抠性能,结果方向完全错了。
第四,重视测试。算子测试、精度测试、性能测试、压力测试,一个都不能少。尤其是边界条件的测试,很多bug都藏在那里。
7. 这个方向后续可以怎么深入
如果你对AI Infra这个方向感兴趣,我的建议是从一个具体的点切入,不要贪多。比如先把算子优化搞透,理解GPU的执行模型、内存层次、并行机制。有了这个基础,再看编译、分布式、推理,会发现很多概念是相通的。
另一个建议是多动手。这个领域光看书是学不会的,必须实际跑代码、做profiling、调参数。哪怕只是把一个简单的矩阵乘法算子优化到接近硬件峰值,收获也比看十篇论文大。
最后,保持对硬件的关注。AI Infra和硬件结合非常紧密,新一代硬件的特性往往会带来新的优化机会。能第一时间吃透新硬件的人,在这个领域永远有优势。
我个人在实际操作中的体会是,这个方向最大的乐趣在于:你做的每一分优化,都能直接看到效果。延迟降了多少、吞吐提了多少、成本省了多少,都是实实在在的数字。这种正反馈,是很多方向给不了的。