大模型背后的AI Infra：从算子优化到GPU显存管理-北京尧图网络

前阵子有朋友问我你们那群做AI Infra的人一天到晚到底在忙什么我说大部分时间我们不是在琢磨怎么把模型调得更聪明而是在跟GPU显存、网卡带宽、算子调度这些“看不见的东西”较劲。他说他不信说训练大模型不就是写写网络结构、喂喂数据然后等Loss往下掉吗我说你要真这么想那说明你还没见过一个40GB的模型怎么硬塞进只能装24GB显存的卡里也没见过训练跑到第七天突然节点宕机、整个进度归零的崩溃现场。这篇文章我想把这些“看不见的人”的工作讲清楚。从最小的“算子”优化一路拆到大模型训练和推理背后的AI Infra系统聊聊我们平时到底在忙什么、为什么这些活不能省、以及如果你也想往这个方向走哪些东西最值得啃。内容会比较干但都是我踩过坑之后的真实理解应该能让你对“大模型背后”这四个字有个完全不同的认识。1. 从“算子”说起模型跑起来之前有人先跟数学打了个招呼1.1 算子不是算法是GPU的“精确动作”很多做业务开发的同事第一次听到“算子”这个词都会愣一下。算子对应的英文是Operator在深度学习里它指的不是“算法策略”而是计算图上的一个基本运算节点。比如矩阵乘法是一次算子ReLU激活是一次算子LayerNorm是一次算子Softmax是一次算子。你写一行x W看起来只是调用了PyTorch的matmul但这一行背后对应的是一个经过高度调优的CUDA kernel它要在NVIDIA GPU上把输入矩阵切块、调度线程、组织共享内存、压满Tensor Core才能跑出接近理论峰值的性能。这里面的差距可以有多大同样的一个矩阵乘法Pytorch默认调用的通用实现和专门针对形状做特化的实现性能可以差出3到5倍如果你用的是FP16或者FP8格式差距还能拉得更大。很多人在调优的时候发现模型跑得慢第一反应是加卡、换显存更大的GPU但往往真正的问题出在算子层面kernel没有融合、访存模式不友好、或者根本没把Tensor Core用起来。换句话说你在GPU上花的每一分算力最后都要落到一个个算子头上算子不快别的都是空谈。1.2 从裸算子到AI Infra中间隔着三层分工我自己的理解里“大模型背后的人”大致可以分成三层虽然很多公司边界模糊但思考问题的方式完全不同。第一层是算法工程师他们关心的是模型结构怎么设计、数据怎么清洗、Loss怎么收敛。他们工作在“计算图”之上很少关心GPU指令长什么样。第二层是框架与编译工程师他们做的事情介于算子和算法之间比如把PyTorch的nn.Module变成可以高效执行的中间表示做算子融合、图优化、内存规划。第三层才是真正意义上的AI Infra工程师负责把模型放到大规模集群里训练和推理这包括并行策略、显存管理、通信优化、调度系统、容错恢复、推理加速等等。我见过最典型的误解是算法同学觉得自己改个网络结构AI Infra的人就能让它在任何规模的集群上自动跑起来。但现实是并行切分方式要跟着模型结构走显存占用要一层一层算通信拓扑要跟网络结构匹配。你改了模型的某个维度后面整个分布式方案可能都要跟着调。所以这些年“AI Infra”火的背后本质上是因为大模型的规模和复杂度已经超出了“单机调包”能处理的范围它需要一整套系统工程来兜底。2. 一个算子的自我修养为什么一行乘法背后要优化半天2.1 先测再调不要靠感觉优化我不太建议谁一上来就埋头读CUDA代码做算子优化的第一步永远是Profiling也就是量化瓶颈。以前我在某个做推理加速的项目里接手过一个服务情况是单卡吞吐上不去领导怀疑是模型太大、显存不够。我拿Nsight System跑了一遍发现GPU的SM占用率只有30%大部分时间都花在了等待数据从CPU侧拷贝到GPU侧根本就不是显存容量的问题。后来我把数据预处理的后处理部分从Python换成了异步加载再把连续的小kernel合并成两三个融合kernel吞吐直接翻了一倍。这个例子很老套但我每次讲给新人听都有效因为它说明了一个道理算子优化的第一步不是“我要写一个更牛的kernel”而是先搞清楚瓶颈到底在计算、访存、还是通信。用PyTorch的话说就是先看看torch.profiler的输出kernel time、GPU utilization、Memcpy的占比哪个占了头名哪个就是要下手的地方。2.2 Kernel融合把十次数据搬运合并成一次为什么要做kernel融合我经常跟非技术同事打一个比方想象你在厨房做饭每次要拿一种调料都要从冰箱走到灶台再走回去十种调料就要来回十趟。如果一次把所有调料都拿齐菜一下子就做完了。GPU的世界里每一次kernel launch都意味着数据要从显存读进寄存器、算完再写回显存尤其是中间结果还要经过全局内存的往返。你做十个elementwise操作就是十次读写全局内存如果能把它们合并成一个kernel读写次数就省成了两次延迟和带宽压力全部降下来。最典型的例子是LayerNorm它本身有求均值、求方差、归一化、缩放偏移好几个步骤但现代的实现早就把这些步骤合在了一个kernel里面中间结果直接留在寄存器和共享内存里完全不用绕到全局内存。再比如残差连接、Dropout、激活函数这三件事在实践中也经常融合成一个算子。这样的小优化在单层网络里感觉不明显但放到几千层的Transformer里省下来的是实打实的训练时长。2.3 FlashAttention这类“重写”才是算子的更高境界比kernel融合更高一档的算子优化是改变算法的访存复杂度。这里最出名的例子就是Attention的优化。标准Attention的复杂度是序列长度的平方要计算完整的注意力矩阵这个矩阵如果序列长度是4096那就是1600万个数存下来就是一堆显存开销。而FlashAttention的思路是不把完整的注意力矩阵写出来而是对Q、K、V做分块切分让每一个计算块尽量在SRAM即GPU上的高速缓存里完成通过重计算的技巧把对HBM的访问量从O(N²)降到O(N)级别。我第一次看完FlashAttention的论文时最大的震撼是它没有改变数学结果但通过精细的访存调度把一个看似不可能省掉的O(N²)访存开销给打掉了。这类工作做多了以后你会发现算子上限的突破很少靠硬件凭空变快更多是靠“数学重排”加上“硬件调度”的双重配合这也是为什么很多人觉得算子优化更像“在GPU上做算法设计”而不是简单写写内核。3. 训练大模型时显存管理是精打细算的“物流调度”3.1 一张显卡装不下的大模型聊完了算子这个微观层面我把镜头拉高一点看看训练大模型时最头疼的问题显存。你可能听过一个说法训练一个7B参数的模型光参数在FP16下就要14GB显存。但实际训练时你还要放梯度梯度跟参数一样是14GB。然后是优化器状态Adam优化器要为每个参数存一份一阶动量、一份二阶动量加在一起又是28GB。你还没算激活值呢光这三样就已经56GB了一张80GB的卡看着好像能放但真正放上激活值、中间缓冲区、通信缓冲区以后马上就会爆。而且这只是7B放到70B、几百B参数的模型上单卡装下整个模型这件事本身就是天方夜谭。这就引出了AI Infra最核心的一个命题一个放不下的模型怎么拆开放在很多张卡上还能高效地一起干活这里的“拆”是分很多维度的每个维度都有不同的代价这是所有做训练系统的人躲不开的“显式管理的折磨”。3.2 反向传播的元凶激活值才是显存大头很多人以为大模型训练占显存最多的是参数这是错的。我自己实测下来激活值activation才是那个“会呼吸的显存杀手”。前向传播的时候你把每一层的中间输出保存下来是因为反向传播的时候要用它计算梯度。Transformer每层中间要存的东西不仅包括Hidden State还包括Attention矩阵、Softmax的中间结果、LayerNorm归一化之前的输入等等。序列长度是4096、批大小一小每层加起来就是好几个GB层数一多直接指数级膨胀。所以业界最常用的一招叫激活重计算英文叫Activation Checkpointing。原理特别朴素前向传播的时候不再保存每一层的激活值只保存一小部分关键节点反向传播需要某层激活值的时候临时从保存的节点重新算一遍前向再把需要用到的值算出来。它本质上是拿“计算时间”换“显存空间”一般能省掉60%到70%的激活显存代价是训练速度下降10%到20%。在很多场景下这是唯一能跑起来的方案你别嫌它慢没有它模型根本塞不进卡里。3.3 Offload与“换页”策略GPU吃不下CPU来兜底激活重计算还是没解决问题怎么办那就得把一部分状态“借”到CPU内存里这就是Offload。很多做训练的同学一开始很排斥Offload因为CPU和GPU之间的PCIe带宽比NVLink带宽差了一个数量级频繁拷贝会把训练拖慢。但在显存不够用的硬约束下Offload往往是从“跑不起来”到“能跑起来”的差别。我这边的实践经验是把参数扎成块按需加载到GPU上用完了再挪回去本质上就是一个显存层面的“LRU缓存策略”。你可以把它想象成仓库放不下货只好在隔壁借了个便宜仓库但每次取货都要花时间在路上。所以Offload的工程设计核心就是“什么时候搬、搬多大块、怎么尽量少搬”这些都是需要根据你实际模型和卡的数量去算的没有一个放之四海而皆准的参数。4. 把一坨上千亿参数拆开并行策略是系统工程的主战场4.1 四种并行之间的关系切batch、切权重、切层、切序列如果一个模型大到单卡放不下就得走上分布式训练这条路。我经常用一句话概括并行策略的世界把一个大活拆成很多小活分给很多人同时干关键问题是每个人怎么干、谁来传话、传多少话。工程上常见的切法有四种。数据并行最简单相当于每个人都拿着完整模型但看不同的数据批次每算完一步大家把梯度同步一下这就要用All-Reduce通信。张量并行是把一个权重矩阵拦腰切成多块放在多张卡上每张卡算一部分这适合单卡放不下某一层的矩阵运算但代价是每次计算都要做好几轮通信。流水线并行是按层切把整个模型按层分成好几段第一张卡算前几层第二张卡算后几层数据像工厂流水线一样流过去缺点是会有流水线气泡一部分卡在等待。序列并行则是大模型时代比较新的玩法把超长序列的注意力计算拆开让多张卡一起算同一个序列块。实际训练一个千亿级模型几乎都是组合拳数据并行在外面套着里面再叠张量并行中间穿插流水线并行。这个组合顺序非常讲究因为每一层并行都会改变通信流量和显存分布排列顺序不对通信量能差出好几倍。这些东西听起来简单真正搭配起来全是排列组合的学问。4.2 通信是一切痛苦的根源在所有并行策略里最让AI Infra工程师头疼的从来不是计算而是通信。张量并行每做一次前向后向传播都要在Tensor Core上算完局部结果以后用All-Reduce把分块结果合并起来这个All-Reduce本身就要走NVLink把所有卡的中间结果汇总后再广播回去。层数一多通信叠加起来训练效率就上不去。如果你租的集群里几张卡不在同一个物理节点上还得走跨节点的RDMA网络那延迟和带宽问题更明显。我有一次帮一个训练任务排查效率很低的问题GPU利用率看着很高但其实大量空闲时间在等网络同步。后来把张量并行的组网从跨节点改成了单节点内也就是把通信最密集的几张卡尽量放在同一台物理机器里训练速度立竿见影提升了两成。所以做AI Infra的人看一个训练任务第一眼看的是“这组并行策略跟底层网络拓扑匹不匹配”而不是“算力够不够”。4.3 训练跑到第七天GPU挂了怎么优雅地爬起来最后一个绕不开的话题是容错。大模型训练动不动就跑几周甚至几个月中间任何一张卡宕机或者网络闪断都会导致训练中断。如果没有容错机制所有的进度都得从头再来那种感觉我只能说是“花钱买崩溃”。常见的做法是在训练过程中周期性保存checkpoint故障发生后从最近的checkpoint恢复但保存checkpoint本身也有代价太频繁会拖慢训练太稀疏又容易丢太多进度。现代AI Infra在容错上做了很多文章比如异步保存checkpoint、把checkpoint写到分布式文件系统而不是本地盘、甚至设计了环形内存备份来降低恢复时间。还有一个很容易被忽视的细节恢复之后还要把集群里所有节点的随机数状态统一重置否则数据加载顺序和模型初始化的随机性都会乱掉训练出来的结果可能跟中断前对不上。这些坑单靠算法同学是想不到的全部属于AI Infra的活。5. 训练只是上半场推理侧的工作同样满地鸡毛5.1 推理最大的敌人是KV Cache模型训练好只是万里长征走完一半真正把它变成一个低延迟、高吞吐的服务推理系统同样一大堆问题。推理跟训练最大的区别是训练可以拿时间换空间推理不行你要在几百毫秒内给出响应还要尽量卡住成本。推理侧的显存大头从激活值变成了KV Cache。Transformer做生成的时候每生成一个token都要把前文所有token的Key和Value重新参与注意力计算。如果每次都重新算前文计算量会随着序列长度线性增长所以工程上会把前面的Key和Value缓存下来这就是KV Cache。序列长度一长KV Cache也变成显存巨兽一个支持4096上下文长度的模型并发一高显存瞬间就告急。我见过不少推理团队被这个问题折腾得够呛最后的解决思路要么是限制并发、要么是量化缓存、要么是动态批处理来尽量复用已有缓存。5.2 吞吐与延迟之间的取舍量化、投机解码、PD分离推理优化的世界里术语也多但本质上都是在一个跷跷板上找平衡延迟低并发就少并发高延迟就容易炸。业界常见的手段有量化把权重从FP16压到INT8甚至FP8推理服务就能吃更多人有投机解码Speculative Decoding先用一个小模型草拟好几个token再用大模型一次验证效果好的时候能接近两倍加速但这个跟模型本身关系很大不是所有模型都适合。还有一个比较新的趋势叫PD分离把Prefill阶段和Decode阶段分开部署。因为Prefill阶段是计算密集的适合用高算力卡去算整段promptDecode阶段是访存密集的适合用另一批卡专门处理逐步生成。两个阶段混在一起的时候总会有人排队等拆开以后整体吞吐能提升不少。这些方案听着都很酷但每一个都要根据实际业务流量画像去调不是装个包就能生效的。6. 这群“看不见的人”的一天定位卡、盯通信、背锅、拆雷6.1 可观测性差是最要命的问题说了这么多系统细节我想再聊聊这群人日常的工作状态。其实AI Infra日常最多的时间不是写代码而是“定位问题”。训练不收敛是数据问题还是数值精度问题GPU利用率低是数据加载慢还是通信等待久显存OOM是哪一层激活值飙了这些问题要回答底层必须有一整套可观测性系统GPU监控、通信时延监控、内存分配追踪、模型数据流profiling。我自己最大的体会是很多AI Infra团队的工作效率差不是因为算法不行而是因为监控不完善。出了问题全靠猜这是最费时间的。所以我到一个新团队第一件事永远是先把监控补齐每个训练任务至少能实时看到GPU利用率、SM占用率、显存占用曲线、通信耗时占比、数据加载耗时。没有这些所有优化都是瞎忙。6.2 从“训练跑通”到“持续交付”的转变以前做传统深度学习Infra的压力主要在产品上线那几天大模型时代不一样了模型不是训完一次就完事要持续迭代、持续调优每次迭代都要重新评估算力集群的容量。模型越来越大数据越来越多成本压力也跟着上来所以现在还要做容量规划、成本分析、混合调度。很多公司内部已经出现了GPU资源池的说法模型训练占用多少卡、推理服务占用多少卡、临时任务能不能挤上去全部要靠调度系统统一管理。这个调度系统的目标已经不只是“把任务排上”而是要能容忍异构GPU、混合并行策略、优先级抢占、自动扩容。今天早上领导说要跑一个500B模型的测试任务晚上就有另一个团队说要上线一个低延迟推理服务两边抢卡怎么办这个博弈也是AI Infra日常的一部分而且压力不小。6.3 怎么跟“看不到这些事”的人协作最后说一下那群“看不见的人”和其他岗位的沟通问题这个我感触很深。算法同事提需求的时候经常会说“你帮我多上几张卡不就行了”但加卡不只是一个数字变大要考虑数据并行的梯度同步会变慢、张量并行跨节点通信会变慢、显存分布要重新算这些都会影响收益。不加卡吧业务会觉得你不上心加了卡吧效果没达标锅还是你的。我后来总结出一条经验跟不懂系统的人沟通不要讲“通信开销”“并行切分”这些词直接给预期收益和时间成本。比如“加两张卡最多提速10%但要多花一天改配置当前效率瓶颈是数据加载不是算力”这样业务侧一听就明白也不会有不切实际的预期。AI Infra工程师本质上是在技术深度和业务期望之间做翻译这个能力有时候比写代码还重要。做AI Infra这些年我自己最深的体会是大家看到的是一个模型越来越聪明背后其实是无数个俯身在GPU日志、网络拓扑和显存碎片里抠日子的人。他们做的活大多是别人看不见的但缺了他们所谓大模型就只是一堆压箱底的网络参数。如果你也对这套系统感兴趣我建议从三件事入手先把PyTorch的profiler用透知道自己的模型瓶颈在哪再去读一读FlashAttention和vLLM的源码看别人怎么在硬件限制下把效率抠出来最后找个训练任务亲手盯几天监控图把数据加载、通信、计算这三段彩带的时间分布背下来。这几个基本功练完你再回来看大模型背后的系统会发现每个“看不见的地方”都藏着可以动手的活儿。

大模型背后的AI Infra：从算子优化到GPU显存管理

相关推荐

前端工程师如何转型Agent开发：从状态管理到智能编排的实战指南

AI日报全流程拆解：从信息采集到深度内容生产的方法论

Pro/E模型练习100例：建模入门必练的实操路径

更多资讯

用生存分析与变点检测识别数字用户隐性流失

Android DLNA 开发实战：DMR、DMC、DMS 三件套落地与避坑指南

内网离线部署RAG知识库问答机器人：从硬件选型到检索调优实践

Vue 项目中修改 Cesium 工具栏提示文字的三种方法与实操记录

多智能体协同架构：MCP协议与A2A通信实战指南

零基础微信理财通基金实操指南：从绑卡到买卖全流程

今日更新

小波包能量谱原理与MATLAB实现：从频带分解到故障特征提取

基于LSTM的网易云音乐评论情感分析：数据清洗与模型调优实战

Alpha-Beta剪枝与极大极小值算法：博弈树搜索原理与实战优化

本周精选

交通拥堵预测毕设：基于Python+LSTM的车辆流量预测系统全流程解析

MATLAB/Simulink搭建10机39节点电力系统暂态稳定仿真实战

无人机航拍三维重建全流程：从SfM到网格生成的避坑指南

本月热门

我发现了一个新思路：用 Remotion + Claude Code 像写代码一样自动化生成短视频

Windows下 Codex 中 Chrome 和 Computer Use 插件不可用问题排查及解决参考方式：TaoToken 统一 Key 配置与验证

2026 大模型集体涨价：用 Python 做企业 Token 成本测算与选型避坑（附配置）

读完文章，想聊聊您的网站？