大模型火成这样身边改行做算法的人越来越多但很多团队真正卡住的地方反而不是模型结构而是“模型压根跑不起来”“训练一挂就是三天”“推理贵到不敢上线”。这时候就得有一群人站出来去折腾那些看不见摸不着的东西算子怎么写、显存怎么省、分布式怎么切、推理怎么加速。这群人就是做 AI Infra 的。我做了几年这个方向正好借这篇东西把“算子”到“AI Infra”这条链路背后到底在忙什么讲清楚也给对这块感兴趣、想入行的朋友一个真实的视角这活不简单但天花板很高而且缺口一直很大。1. 拆开“算子”和“AI Infra”这两个词很多刚接触的人以为“算子”是个高深概念其实没那么玄乎。神经网络里无论是卷积、矩阵乘法、归一化还是注意力机制最终落到 GPU 上执行的都是一个个基础计算单元这些计算单元就叫算子。模型训练就是把海量数据喂给一连串算子一层一层算过去再一层一层把梯度传回来。PyTorch 里写torch.matmul(a, b)一行代码背后就是一个算子。框架帮你封装好了一切所以算法工程师天天写算子调用却不一定知道算子内部发生了什么。AI Infra 就更宽泛了。它不是某个单独组件而是让大模型能够高效训练、稳定部署的整套基础设施数据管道怎么吞吐、训练任务怎么调度、几千张卡怎么协同、推理服务怎么压延迟、存储怎么不拖后腿。如果说算法工程师负责“让模型变聪明”那 AI Infra 工程师负责“让模型跑得动、跑得快、跑得便宜”。算子开发是 AI Infra 里最硬核的环节之一但 AI Infra 远不止算子。1.1 从一个矩阵乘法的加速说起先看一个具体例子矩阵乘法是所有大模型里出现频率最高的算子。一个朴素实现直接三重循环每算一个输出元素就去显存里取数据GPU 的有效算力可能只剩峰值的十分之一。原因在于 GPU 的核心是数量多但单核简单它极其依赖数据复用——同一份数据如果能被多个计算单元反复使用就能把性能拉满如果每次都重新访存计算单元只能干等。业界怎么优化这类算子核心思路无非几条分块Tiling让数据尽量留在片上缓存里向量化一次多处理几个数双缓冲把访存和计算重叠起来。我曾优化过一个某团队的矩阵乘算子只把分块大小从 32 调到 128再配合寄存器重排单测速度就提升了接近 4 倍后来把数据排布改成更适合向量加载的格式又提升了 2 到 3 倍。一来一回单个算子快了 12 倍以上。这就是为什么大模型训练时同样一张卡别人显存能塞 4096 的长度你塞 2048 就爆了——算子地址重排和显存分配策略不同内存足迹能差出数倍。1.2 算子开发怎么就成了“脏活累活”有人说算子开发是“脏活累活”这评价不算冤枉。因为一个好的算子要同时考虑计算效率、访存效率、数值稳定性和硬件适配性。你做出来的算子跑得再快数值稍微差一点训练就出现 NaN前功尽弃。还得分心处理不同架构的差异有的卡擅长稠密矩阵乘有的卡对稀疏操作更友好有的是半精度强、有的是 INT8 强。同一个算子想在各家硬件上都有好表现基本等于每个平台各写一遍优化。但换个角度看这活的价值也是实打实的。谁掌握了算子的性能谁就掌握了训练和推理的成本曲线。一个 FlashAttention 式的算子优化能让长文本训练从“显卡爆显存”变成“顺利跑完”这种成就感是调 prompt 替代不了的。真正想把这块做好除了要懂 GPU 硬件细节、汇编指令还得有扎实的并行计算功底而这些东西都是需要时间磨的。2. AI Infra 到底是什么把幕后的活儿分成四块很多人误以为 AI Infra 只是“买卡、装环境、写脚本”。实际上一个成熟的 AI Infra 团队要做的事可以分成四块拼图训练基础设施、推理优化、资源调度与数据平台。训练基础设施解决的是“让大模型能训练起来”。这里涉及分布式并行策略怎么选、通信拓扑怎么设计、断点续训怎么做得可靠、混合精度怎么配。这些问题不解决上千张卡启动任务时只要有一张卡掉线整个集群就得重来损失以小时计。推理优化解决的是“让模型能用起来”。大模型训练完了不是终点部署上线才是。怎么压首字延迟、怎么提高吞吐、怎么用量化手段把模型塞进更小的显存、怎么用 KV Cache 让长文本生成不重算历史每一个都能单独开一门课。比拼的不是“能不能跑”而是“同样的服务能力谁的硬件成本低一倍”。资源调度负责“让卡不闲着”。一个集群几十个任务排队谁先跑、谁后跑抢占和排队策略怎么定直接影响 GPU 利用率。很多公司的 GPU 利用率提不上去不是卡不够而是调度策略太粗糙。做得好的人能把有效利用率从 30% 抬到 70% 以上这个收益比单纯买新卡大得多。数据平台管的是“喂给模型的东西靠不靠谱”。大模型训练的数据动不动就是几千亿 token清洗、去重、配比、流式读取任何一环出问题模型效果都会受影响。训练跑到一半发现数据里有大量重复垃圾损失的是几十万度的电费和时间成本。2.1 用“开餐厅”比喻 AI Infra把大模型比作开餐厅算法工程师是主厨负责研究新菜AI Infra 团队则是一家店的选址、后厨水电、供应链、排风管道和出餐动线。主厨很重要但没有可靠的水电和供应链再好的菜谱也端不上桌。GPU 就是炉灶数据管道就是菜市场到后厨的运输线分布式框架就是后厨的分工制度。每一环都有看不见的人在做“笨功夫”但正是这些笨功夫决定了餐厅能同时接待多少客人、翻台速度多快、毛利率多高。一个很典型的场景算法团队在单卡上验证模型效果觉得没问题满怀信心拿到集群上训练结果 8 卡效率只有单卡的 3 倍32 卡效率只有单卡的 7 倍。算法工程师往往第一反应是“框架不行”但真相常常出在数据加载线程数太少、CPU 预处理跟不上、或者通信库配置没优化。这就是 AI Infra 的价值把硬件的每分力气都使在刀刃上。2.2 为什么说 AI Infra 决定了大模型的成本线大模型训练的成本惊人这里的成本不仅是钱还是时间。同样一个千亿参数模型A 团队用 2048 张卡跑 30 天B 团队只用了 1024 张卡跑 25 天背后差的就是 Infra 水平。每省下一半卡换来的是上千万的预算节余。而推理阶段更夸张同一个模型有的团队只靠算子融合和显存优化就能把单卡吞吐翻倍。长期来看AI Infra 团队不是“后勤部门”而是实打实的利润中心。3. 这群人到底在忙什么一次大模型训练背后的真实工作流这一节我想完整走一遍“大模型从零到跑起来”的过程让大家直观感受 AI Infra 人员每天面临的算术题、选择题和“玄学”排查。3.1 从单卡到千卡先算一笔显存账拿一个 7B 参数的模型举例假设我们用混合精度训练FP16 存权重和梯度FP32 存优化器状态。第一步先把账算清楚FP16 权重7B x 2 字节 ≈ 14GBFP16 梯度同样是 14GBFP32 优化器状态Adam 需要保存一阶动量、二阶动量和 FP32 的模型副本每项 7B x 4 字节共 84GB光这三项就已经 112GB 了还没算激活值。一张主流训练卡的显存只有 80GB 左右单卡根本装不下。这意味着必须上并行策略而不是买一张更大的卡。激活值也不是小数目输入序列越长、batch 越大、模型越深激活值占用越高。就算模型权重能塞下长序列训练时激活值也可能直接让显存爆掉这时就得考虑激活重计算用计算换显存只存部分中间结果反向传播时重新算一遍或者梯度检查点。这步“显存算术”是所有 AI Infra 工程师的基本功。方案选型之前先估算而不是上了训练任务才发现 OOM这是职业素养的体现。等到规模大了就不只是“一张卡塞不塞得下”的问题而是“多张卡之间怎么分、通信开销能不能接受”的问题。3.2 并行四件套数据、张量、流水线、序列并行单卡放不下自然就得上多卡。多卡并行不是简单“多复制几份模型”而是要分情况选策略并行策略比喻切分维度适合场景主要成本数据并行多人各抄一本书复制模型、切分数据模型能塞进单卡但数据量太大每步需要同步梯度张量并行一章拆给几个人写把单层算子按行/列切开单层参数超出一张卡显存通信量极大需要高速互联流水线并行书稿分章节流水接力把网络按层切成几段模型层数深、张量并行受限流水线存在空泡率序列并行一本书拆页轮流读把长序列切到多张卡超长上下文的注意力计算引入了额外的序列通信实际生产环境绝不是只用一种。一般 7B 到 13B 用数据并行加 ZeRO优化器状态分片就够了65B 以上的大模型往往是张量并行 流水线并行 数据并行三层混合每一层的并行度都要根据集群拓扑、带宽、显存预算仔细调。我在某次压测中就踩过坑某个团队把 8 卡张量并行的 parallel size 从 4 调到 8理论上每卡显存更省了实际却因为卡间通信频繁整卡利用率从 45% 跌到 20%。这说明并行度不是越大越好得结合通信开销来看。一个通用的判断方式先估算每步通信的数据量和通信耗时再算计算耗时让通信时间和计算时间尽量重叠别让卡闲着等数据。3.3 通信才是真正的隐形瓶颈很多人把注意力放在算力上真正做过大规模训练的人会告诉你通信优化做好了你省下的时间比换更好的 GPU 还多。数据并行里最常见的是每步结束时做梯度同步把每张卡的梯度汇聚起来求平均。用的算法叫 AllReduce。朴素做法是中心化聚合所有卡把梯度发给 0 号卡算完再广播回去0 号卡瞬间成为瓶颈。现在普遍用的是 Ring AllReduce把 P 张卡排成一个环通信量从“所有数据过单点”变成“每个节点只传相邻的部分”通信总时间能随卡数增长大幅改善。同样的通信数据量Ring 结构下传输耗时和端到端时延低得多。但通信优化没有银弹。跨机通信走网络机内通信走高速总线两者的带宽差异可能是数量级的。并行切分如果没考虑物理拓扑让跨机通信做了太多活整体性能一定被拉垮。所以我做方案时一定会先把物理拓扑画出来哪几张卡在同一台机器上、哪几个节点在同一机柜里、网络是收敛的还是无收敛的再决定并行组怎么分。3.4 断点续训和容错最无聊但最省钱的活大规模训练动不动跑几个星期任何一个节点宕机、网络闪断、驱动报错任务就可能从头再来。AI Infra 团队花了很多精力在做检查点保存和自动恢复每隔一段时间将模型权重、优化器状态、数据偏移量完整落盘任务挂掉后自动拉起从最近的检查点恢复而不是从头开始。这块工作看着不“性感”但价值极大。一次千卡训练挂掉重启光浪费的电费和时间就够买一台好服务器了。有经验的团队会把检查点频率、保存位置、恢复流程全做成自动化让训练任务像水一样流过去偶尔有涟漪也不会断流。4. 推理优化让大模型从“能跑”到“跑得又快又便宜”训练广受关注但推理才是真正天天烧钱的部分。同一个大模型推理优化做得好不好成本差距可能是 5 倍甚至 10 倍。这块也是 AI Infra 团队最能体现“价值感”的地方。4.1 KV Cache用显存换速度的经典做法Transformer 生成文本时每个 token 都要重新计算前面所有 token 的 Key 和 Value。如果不缓存每生成一个 token 都要重算一整遍成本会随序列长度呈平方级上升。KV Cache 的思路就是把这些历史上算过的 Key/Value 存起来生成新 token 时只增量计算复杂度立刻降了一个量级。但 KV Cache 不是免费的它吃显存。算一笔账假设一个 32 层、隐藏维度 5120 的模型每个 token 需要缓存的 KV 数据大约是 2K 和 V乘 32 层乘 5120 维也就是约 32 万个浮点数。序列长度 4096 时需要 13 亿个浮点数以 FP16 存储就是 2.68GB 左右。用户在验证码风格的长对话场景一多这部分的显存占用会比模型权重还高。所以现在有各种 KV Cache 压缩策略比如分组查询注意力GQA、KV 量化、H2D 缓存分层。选不选、怎么选都是 Infra 工程师在模型上线前就要定好的事。4.2 量化把模型装进显存的最直接手段推理和训练一个巨大差异是推理可以忍受更低的数值精度。一个 175B 参数模型FP16 权重需要 350GB 显存单卡无论如何也塞不下但如果用 4bit 量化权重就只有约 87.5GB两张卡就能放得下。别以为量化就是简单把 FP16 转 INT4。如果量化是均匀的遇到数值分布偏斜的参数就会损失惨重。实践中常用分组量化把一小块权重分成一组每组单独算 scale 和 zero point把误差控制在一定范围内。更聪明的方法是做混合精度敏感层保留高精度非敏感层用低精度。某团队的量化方案上线后单卡吞吐提升了两倍多而评测分数几乎没变。这种“近乎白拿”的收益不做 Infra 的人是很难想象到的。4.3 端到端延迟先在纸上算一遍再动手推理延迟是可以提前估算的。假设一个 7B 模型FP16 权重 14GB单卡显存带宽约为 2TB/s。生成第一个 token 的 prefill 阶段是计算密集型的需要按矩阵乘的 FLOPs 来算但后续的 decode 阶段由于 batch 很小每生成一个 token 都相当于拿整个模型权重做一次矩阵向量乘此时瓶颈几乎不在算力而在带宽。每生成一个 token需要把 14GB 权重从显存读一遍按 2TB/s 算理论延迟下限约为 7ms。这个数字比很多人印象中的“模型很重”要快但它只是下限。一旦 batch 变大虽然单 token 延迟会被分摊GPU 核心利用率上升吞吐提升了每个请求的感受延迟反而可能增加。所以推理服务必须在延迟和吞吐之间做权衡。定服务目标时我会先问清楚是优先首字延迟还是优先并发吞吐这两个方向会导向完全不同的优化工作。4.4 一次推理优化 40% 延迟的实操记录这里分享一个实战案例。某团队的对话模型在单卡上推理平均生成一个 token 要 22ms用户体感偏慢。先做 profiling发现两个明显问题一是算子层面有大量小的矩阵运算串行执行GPU 利用率只有 25%二是 KV Cache 命中率低导致重复计算了不少历史 token。针对性做了三件事把相邻的小算子做算子融合把多个尺寸小的 GEMM 合并成一个大 GEMM减少了内核启动和内存搬运的开销再把 KV Cache 改成按请求动态分配减少显存碎片最后调整了量化策略让层间精度分配更合理。整个改完平均生成延迟压到了 13ms 左右下降超过 40%。没有太多花哨技巧就是一遍遍 profiling、定位、改动、验证。推理优化的本质就是“找出瓶颈并消灭它”而瓶颈往往不在你直觉以为的位置。5. 常见问题与排查技巧实录那些看着像玄学的故障做 AI Infra 最烦的不是难的技术题而是看起来完全没有头绪的诡异故障。这里把我自己遇到的有代表性的问题整理成一份速查表照着排查能省大量时间。现象可能原因排查方向训练 loss 不降数据管道卡死、数据重复、学习率配置错误先看数据加载耗时和 batch 样本分布显存明明够却 OOM显存碎片化、缓存未复用检查分配器策略打印显存块分布GPU 利用率低CPU 预处理慢、通信等待、算子核启动过密做 timeline 分析看 GPU 空闲区间落在哪多卡训练速度不随卡数增长通信开销过大、parallel size 不当对比单卡/多卡耗时算通信占比推理首字延迟超高prefill 阶段计算量爆炸、未做输入长度限制拆开 prefill 和 decode 分别压测偶发训练中断节点过热、网络丢包、驱动问题看系统日志、硬件事件、抓网络数据5.1 训练 loss 不降先检查数据管道某次技术分享时有团队说用某个模型训练了三天loss 纹丝不动。一开始大家都怀疑是模型结构改出了问题排查了一整天没结果。后来我看了一眼数据加载代码发现 DataLoader 里 prefetch factor 被设成了 1GPU 每算一个 batch 都要等 CPU 跑完一轮数据预处理。GPU 利用率不到 10%所谓“训练三天”其实大部分时间在空转。这就是 Infra 人员的价值很多问题看起来像算法玄学实际根子在工程细节上。排查数据管道问题最简单的方法是记录每个 batch 的时间戳看有没有周期性延迟。正常情况数据加载耗时应该远小于计算耗时一旦出现“计算等数据”的情况整个训练效率会断崖式下跌。CPU 侧多开几个预处理进程、用内存映射把数据文件缓存起来、把 tokenize 结果先持久化都是立竿见影的解法。5.2 显存碎片化的坑训练一个长序列模型时会频繁申请、释放大小不一的显存块。如果分配器回收不及时或策略不当显存会碎成一片片小区域明明总量还有余量就是申请不到连续内存直接报 OOM。这种问题最容易在改 batch size 或者序列长度后出现。排查时先在启动脚本里开启显存快照把分配记录保存下来打印分配块的地址分布。如果看到大量小的空闲块散布在已用块之间基本就能确定是碎片化。解决办法包括开启显存预分配、用缓存分配器复用同一尺寸块、以及把动态 shape 改成静态 shape 减少申请次数。小改动收益大。5.3 GPU 利用率忽高忽低多半是在等通信多卡训练时如果集群里各卡步调不一致有的卡在等梯度同步GPU 利用率就会像心电图一样上下跳。要抓这种问题必须做一次端到端的 timeline 分析看每张卡的计算时间、通信等待时间、数据加载时间分别占多少。如果通信等待时间占比居高不下就该考虑重新设计并行切分方案或者调整通信库的参数设置比如通信计算重叠、梯度分桶大小。经验之谈通信优化不要一上来就换算法先量化问题。拿到 profile 数据看到底是拓扑限制了带宽还是消息太小导致传输效率低。很多情况下只是因为梯度同步把大消息拆成了上千个小消息每个小消息都要额外开销把发送缓冲调大、分批合并速度就能提升明显。5.4 profiling 工具才是真正的“照妖镜”接手任何性能问题我的第一动作永远是先跑 profiling而不是凭感觉改代码。一个完整的性能剖析要能回答三个问题时间花在哪、显存用在哪、通信堵在哪。时间占比用 timeline 视图看显存占用用内存快照看通信开销用通信库的统计页面看。拿到数据再动手所有“玄学”都会变成清晰的工程决策。6. 给想入局的人一点实在话总有人问我现在做算法卷成什么样了AI Infra 是不是一个好方向我的回答是这个方向需求旺盛、壁垒高、不容易被替代但它的门槛也摆在那里。它需要一定的体系结构基础需要耐得住性子做 profile、压测、调参需要直面“改了一周代码性能只涨了 5%”的挫败感。如果你动了入局的念头我给你几个判断标准你是否愿意去看一层层底层实现是否享受把性能从 60% 推到 95% 的成就感是否能在面对一个毫无头绪的故障时保持冷静。如果你觉得这些是乐趣而不是折磨那这个方向大概率适合你。我刚入行的时候导师跟我说过一句话Infra 的活就是把“不可能”变成“可能”再把“可能”变成“默认配置”。做了几年之后我深以为然。外面的人看到的是模型参数的跃迁我们看到的是一条条训练曲线、一块块显存碎片、一次次通信握手。如果你也想成为大模型背后那个“看不见的人”不要被复杂的术语吓退从读懂一张 profiling 图开始从把一个算子在真实硬件上跑快一倍开始。有耐心、沉得住气你会发现自己正站在这个行业最重要的一条赛道上。
复杂度详解)
