1. 从调包到造轮子为什么需要一张AI Infra知识地图很多人第一次接触AI Infra这个概念是在面试被问到一个千亿参数模型训练时显存到底被什么吃掉了却答不上来的时候。我自己也经历过这个阶段——能跑通PyTorch的Demo能调HuggingFace的接口但一旦模型规模上去、集群节点变多、推理延迟要求变严整个人就懵了。问题不在于不会写代码而在于脑子里没有一张完整的基础设施地图不知道一个请求从发出到返回中间到底经过了哪些层、每层在解决什么问题、瓶颈通常出现在哪里。AI Infra人工智能基础设施这个词听起来很宏大但拆开来看它其实就是回答三个问题算力怎么组织、数据怎么流动、模型怎么高效跑起来。它横跨了硬件层GPU/NPU、互联、存储、系统层驱动、容器、调度、框架层训练框架、推理引擎、以及工程层分布式策略、量化压缩、服务编排。这四层里任何一层没搞明白实际项目里就会踩坑——比如你明明用了A100训练速度却只有别人的三分之一比如推理服务QPS上不去排查半天发现是KV Cache没管好。这篇笔记适合三类人看一是刚转行做AI工程、想系统补齐基础设施知识的开发者二是做算法但经常被工程问题卡住的同学三是准备面试、需要把零散知识点串成体系的人。我不会只讲概念而是按这一层解决什么问题、核心组件有哪些、关键参数怎么算、实际踩过什么坑的逻辑来展开。整篇内容基于我自己的学习和项目实践整理涉及具体数字的地方会给出计算过程涉及选型的地方会说明取舍理由。你可以把它当成一张可以反复查阅的地图而不是一篇读完就忘的科普。2. 算力底座GPU、互联与显存的三本账2.1 显存到底被谁吃掉了一次完整的账目拆解几乎所有AI Infra的入门问题最后都会落到显存上。我见过太多人训练时遇到OOMOut Of Memory第一反应是换个更大的卡但其实显存是被四部分吃掉的搞清楚比例才能对症下药。以训练一个参数量为 Φ 的模型为例假设用混合精度fp16计算、fp32主权重加Adam优化器显存占用大致分四块模型权重fp16存储占 2Φ 字节fp32主权重副本占 4Φ 字节合计 6Φ。梯度fp16梯度占 2Φ 字节。优化器状态Adam需要保存一阶动量m和二阶动量v通常用fp32占 8Φ 字节。激活值这部分和batch size、序列长度、网络结构强相关最难估算往往是大头。把前三项加起来就是 16Φ 字节。也就是说一个10B参数的模型光权重、梯度、优化器状态就要 160GB 显存。一张80GB的卡根本放不下这就是为什么必须上分布式——要么切模型张量并行、流水并行要么切优化器状态ZeRO要么两者结合。激活值这块我单独说因为它最容易被低估。Transformer的激活值和层数 L、序列长度 S、隐藏维度 H、batch size B 大致成正比。实际项目里我遇到过batch size从8调到16就OOM的情况排查下来就是激活值翻倍。解决办法有梯度检查点Gradient Checkpointing用计算换显存能把激活值占用降到原来的平方根级别代价是训练速度慢20%到30%。这个取舍在显存紧张时非常值得。提示估算显存时先算 16Φ 这个下限再根据序列长度和batch size估激活值。如果 16Φ 已经超过单卡容量别犹豫直接上分布式方案硬扛只会浪费时间。2.2 卡间互联为什么你的多卡训练没有线性加速很多人以为多卡训练就是卡越多越快实测下来经常是4卡只有2.5倍加速8卡只有4倍。瓶颈往往不在计算而在卡间通信。这里必须理解两个概念通信带宽和通信拓扑。单机内多卡通信走的是NVLink或PCIe。NVLink的带宽远高于PCIe比如某代NVLink单卡双向带宽能到600GB/s以上而PCIe 4.0 x16只有约32GB/s。差距接近20倍。所以同一台机器内卡间通信走NVLink还是PCIe对张量并行这种通信密集的策略影响巨大。跨机通信走的是网络比如InfiniBand或RoCE。这里的关键参数是带宽和延迟。分布式训练里All-Reduce是最常见的通信原语它的耗时大致和数据量除以带宽成正比同时受延迟影响。当模型很大、梯度很多时通信时间可能超过计算时间这时候就需要通信与计算重叠overlap让梯度在反向传播算出来的同时就开始传输。我踩过的一个坑是以为买了高带宽网卡就万事大吉结果拓扑没规划好跨机通信走了低速路径训练效率直接腰斩。后来才明白集群网络拓扑比如胖树结构和网卡绑定方式都会影响实际带宽。选型时不能只看单卡参数要看整机、整集群的实际通信能力。2.3 算力选型训练卡和推理卡不是一回事新手常犯的错是拿训练卡去做推理或者反过来。这两类负载对硬件的要求完全不同。训练阶段看重的是算力密度、显存容量、互联带宽。因为要存权重、梯度、优化器状态显存越大越好因为要做张量并行和流水并行互联带宽越高越好因为要跑很久稳定性和散热也重要。推理阶段看重的是显存带宽、能效比、批处理能力。推理时权重是只读的不需要存梯度和优化器状态所以显存容量要求相对低但显存带宽直接决定token生成速度。另外推理对延迟敏感需要低延迟的调度和批处理策略。下面这张表是我整理的一个粗略对照帮助理解两类负载的差异维度训练负载推理负载显存容量需求极高权重梯度优化器激活中等主要是权重KV Cache显存带宽需求高极高决定生成速度互联带宽需求极高并行通信中等多卡推理时才有计算精度fp16/bf16为主部分fp32int8/int4量化常见批处理大batch提升吞吐动态batch平衡延迟与吞吐典型瓶颈通信、显存容量显存带宽、调度延迟理解这张表你在做资源规划时就不会盲目堆卡而是根据负载类型匹配硬件。3. 分布式训练切分策略与通信原语的实战逻辑3.1 数据并行、张量并行、流水并行到底怎么选分布式训练的核心思想就一句话单卡放不下或算不快就把工作切到多卡上。但怎么切有讲究。数据并行Data Parallelism是最直观的每张卡存一份完整模型喂不同的数据算完梯度后做All-Reduce求平均。它的优点是实现简单缺点是每张卡都要存完整模型和优化器状态显存利用率低。当模型大到单卡放不下时数据并行就失效了。张量并行Tensor Parallelism是把单个矩阵运算切到多卡上。比如一个大的线性层按列或按行切开每张卡算一部分再通过All-Reduce或All-Gather合并结果。它的优点是能切分模型本身缺点是通信频繁对互联带宽要求极高通常只在单机内用NVLink做。流水并行Pipeline Parallelism是把模型按层切成若干段每张卡负责一段数据像流水线一样依次流过。它的优点是通信量相对小只在段边界传激活值缺点是存在气泡bubble——流水线填充和排空阶段有些卡在空转。为了减少气泡可以用微批micro-batch技术把一个大batch拆成多个小batch依次灌入。实际项目里这三种策略往往是组合使用的。比如一个超大模型可能同时用张量并行机内 流水并行机间 数据并行副本。这个组合就是常说的3D并行。选型的核心原则是通信密集的切分放在高带宽域内通信稀疏的切分放在跨机域。3.2 ZeRO用通信换显存的经典取舍ZeROZero Redundancy Optimizer是微软提出的一套显存优化方案核心思想是消除数据并行中的冗余存储。标准数据并行里每张卡都存了完整的优化器状态、梯度和权重这其实是巨大的浪费。ZeRO分三个阶段逐步消除冗余ZeRO-1只切分优化器状态。显存降到约 4Φ 2Φ 2Φ/ NN为卡数通信量和标准数据并行相当。ZeRO-2再切分梯度。显存进一步降到 2Φ (2Φ12Φ)/N。ZeRO-3连权重也切分。显存降到 16Φ/N但通信量显著增加因为每次前向和反向都要All-Gather权重。这里有个关键取舍ZeRO阶段越高显存越省但通信越多。ZeRO-3虽然能把显存压到最低但通信开销可能让训练变慢。我的经验是如果显存够用优先用ZeRO-1或ZeRO-2只有当单卡实在放不下时才上ZeRO-3并且要配合高带宽网络。注意ZeRO-3的通信模式和参数配置强相关配置不当会导致通信量爆炸。实际使用时要关注通信桶大小bucket size这个参数太小会导致频繁小通信太大则占用额外显存。3.3 通信原语All-Reduce、All-Gather、Reduce-Scatter的区别分布式训练里绕不开三个通信原语很多人分不清它们的区别导致看代码时一头雾水。All-Reduce每张卡都有一份数据先求和Reduce再把结果广播给所有卡All。结果是每张卡都拿到完整的求和结果。数据并行求梯度平均用的就是它。All-Gather每张卡有一份数据分片收集起来拼成完整数据分发给所有卡。ZeRO-3前向传播时收集权重用的就是它。Reduce-Scatter先求和再把结果切分每张卡只拿一部分。ZeRO-2/3反向传播时用的就是它。这三个原语的关系可以用一句话概括All-Reduce Reduce-Scatter All-Gather。理解这个分解就能理解为什么ZeRO-3的通信量是ZeRO-1的1.5倍左右——因为它多了权重的All-Gather。实际调优时通信原语的实现比如Ring All-Reduce、Tree All-Reduce在不同拓扑下性能差异很大。框架通常会根据集群拓扑自动选择但了解原理能帮你在性能异常时定位问题。4. 推理引擎从KV Cache到连续批处理的关键技术4.1 KV Cache自回归生成的显存黑洞大模型推理和训练最大的不同在于自回归生成。生成第 t 个token时需要用到前面所有token的Key和Value。如果每次都重新计算复杂度是 O(n²)完全不可接受。所以推理引擎会把历史token的Key和Value缓存下来这就是KV Cache。KV Cache的显存占用可以估算假设层数 L、隐藏维度 H、序列长度 S、batch size B精度为fp16那么KV Cache大小约为 2 × L × H × S × B × 2 字节。以一个70B模型为例L80H8192S4096B1算下来KV Cache就要约 10GB。batch size一上去KV Cache迅速膨胀成为显存的主要占用者。这就引出了几个关键优化方向PagedAttention把KV Cache按页管理像操作系统管理内存一样减少碎片提升显存利用率。这是某开源推理引擎的核心创新。KV Cache量化把KV Cache从fp16降到int8甚至int4显存直接减半或更多代价是精度略降。前缀共享多个请求如果有相同的前缀比如相同的system prompt可以共享KV Cache避免重复计算。我在实际部署时发现KV Cache管理的好坏直接决定了一个推理服务能扛多少并发。同样的硬件优化前后QPS可能差好几倍。4.2 连续批处理让GPU不再空转传统批处理是攒一批请求一起算算完再攒下一批。问题是每个请求的生成长度不一样短请求算完了长请求还在算短请求占的算力就浪费了。GPU利用率上不去延迟还高。连续批处理Continuous Batching的思路是不等整批算完而是每生成一个token就检查一次谁算完了就把它踢出去同时把新请求加进来。这样GPU始终在处理当前还活着的请求利用率大幅提升。这个机制听起来简单实现起来要考虑很多细节新请求怎么插入、KV Cache怎么动态分配、不同请求的序列长度怎么对齐。这也是为什么推理引擎的调度器是个技术活。实测下来连续批处理相比静态批处理吞吐能提升数倍尤其在高并发场景下。4.3 量化精度与速度的平衡术量化是推理优化的另一大利器。核心思想是用更低的精度表示权重和激活值减少显存占用和计算量。常见的量化方案有训练后量化PTQ模型训练完后直接量化简单快速但精度损失可能较大。量化感知训练QAT训练时就模拟量化误差精度更好但需要重新训练。权重量化 vs 权重激活量化只量化权重实现简单量化激活能进一步加速但对硬件和算子支持要求高。量化位宽上int8是最成熟的int4在近两年也逐步可用。我的经验是int8量化通常精度损失很小可以直接上int4量化要仔细评估某些任务上精度掉得明显需要配合更精细的量化策略比如分组量化、混合精度量化。提示量化不是免费的午餐。上量化前一定要在验证集上对比精度尤其是对精度敏感的任务如代码生成、数学推理量化带来的微小误差可能被放大。5. 工程落地调度、监控与成本控制5.1 集群调度让任务找到合适的卡有了硬件和框架还需要一层调度系统把任务分配到合适的节点上。AI训练任务对调度的要求很特殊需要GPU资源、需要特定的拓扑比如同一台机器内的卡、任务运行时间长、失败后需要断点续训。常见的调度方案有基于Kubernetes的也有专用的调度器。核心要解决的问题包括资源分配如何把GPU、CPU、内存、网络作为一个整体分配避免碎片。拓扑感知如何保证需要高带宽通信的任务被分配到同一台机器或同一网络域内。抢占与优先级高优先级任务能否抢占低优先级任务的资源。故障恢复节点挂了任务能否自动重启并从checkpoint恢复。我见过的一个典型问题是调度器把需要张量并行的任务分散到了不同机器上导致跨机通信成为瓶颈训练效率极低。这就是拓扑感知没做好。选调度器时一定要确认它支持拓扑感知调度。5.2 监控训练崩了怎么快速定位训练任务跑起来只是开始能不能稳定跑完才是关键。监控要覆盖几个层面硬件层GPU利用率、显存占用、温度、功耗、ECC错误。通信层网络带宽、延迟、丢包。框架层训练loss、梯度范数、学习率、吞吐tokens/s。业务层checkpoint保存、任务进度、预计完成时间。我踩过的坑是只监控了loss结果训练突然变慢排查半天发现是某张卡的显存出现了ECC错误导致计算降速。如果早点监控硬件指标就能第一时间发现。所以监控一定要分层硬件指标和框架指标都要看。另一个经验是梯度范数是个非常有用的指标。如果梯度范数突然爆炸说明训练不稳定可能要调小学习率或加梯度裁剪如果梯度范数长期接近零说明模型可能陷入饱和要检查初始化或激活函数。5.3 成本控制算力是钱省着花AI Infra最终要面对的是成本。算力很贵训练和推理都在烧钱。控制成本有几个方向提高利用率通过更好的调度、更大的batch、更优的并行策略让GPU不空转。混合精度和量化用更低的精度换取更高的吞吐。弹性伸缩推理服务根据负载动态扩缩容低峰期释放资源。Spot实例/抢占式资源用更便宜但可能被回收的资源跑容错性高的任务配合checkpoint。模型压缩蒸馏、剪枝、量化减小模型体积和计算量。我的体会是成本优化不是某一个环节的事而是贯穿整个链路的。从选卡、选框架、选并行策略到调度、监控、推理优化每一环都有省钱的空间。关键是先建立度量——知道钱花在哪才能优化。6. 我在这条路上踩过的几个真实坑第一个坑是盲目追求大卡。早期遇到OOM第一反应是换更大的卡结果换了卡还是OOM因为激活值随batch size增长卡再大也扛不住。后来才学会先算账搞清楚显存被什么吃掉再决定是切模型还是减batch还是上梯度检查点。第二个坑是忽视通信拓扑。有次做多机训练怎么调都慢最后发现是网络拓扑没规划好跨机通信走了低速路径。这件事让我明白AI Infra不是单卡的事集群层面的规划同样重要。第三个坑是推理不做压测。上线前觉得模型能跑就行结果一上量就崩。后来才建立完整的压测流程测不同并发下的延迟和吞吐测KV Cache的显存占用测量化的精度损失。压测做扎实了上线才稳。第四个坑是监控只看loss。前面提过硬件故障、通信异常都可能让训练变慢但loss不一定立刻反映出来。分层监控是必须的。这些坑的共同点是问题不在算法而在基础设施。算法决定了模型能做什么基础设施决定了模型能不能高效、稳定、低成本地做出来。这也是为什么AI Infra值得系统学习——它是把算法变成产品的桥梁。7. 给想深入AI Infra的人几条实用建议如果你刚开始学AI Infra我的建议是从显存账目和通信原语入手。这两个是理解一切分布式策略的基础。把 16Φ 这个公式记住把All-Reduce、All-Gather、Reduce-Scatter的区别搞清楚很多看似复杂的方案就变得清晰了。然后动手跑一遍小规模分布式训练。不用真的搞多机单机多卡就够。用数据并行跑一个中等模型观察显存变化、通信开销、加速比。再试试ZeRO的不同阶段感受显存和速度的取舍。这种亲手验证的经验比看十篇文章都管用。推理方向重点研究KV Cache和批处理调度。这是推理引擎性能的核心。可以拿一个开源推理引擎读它的调度器代码理解连续批处理是怎么实现的。再自己压测一下看看不同参数下的延迟和吞吐曲线。最后养成算账的习惯。显存要算通信量要算成本要算。AI Infra的很多决策本质上是资源约束下的优化问题。算清楚了决策就有依据而不是凭感觉。这个领域变化很快新硬件、新框架、新优化技术层出不穷。但底层的原理——算力怎么组织、数据怎么流动、瓶颈在哪里——是相对稳定的。把原理吃透再去看新东西就能快速抓住重点。我自己也是这么一路走过来的从只会调包到能独立设计训练和推理方案靠的就是把这张知识地图一点点补全。希望这篇笔记能帮你少走一些弯路。