AI Infra全景学习笔记：从GPU算力到分布式训练与推理

做个AI Infra的学习整理我一直觉得这活儿特别容易让人陷入误区。一搜AI Infra跳出来全是Kubernetes、GPU驱动、PyTorch分布式、Ray、vLLM……词都认识串起来却糊成一片。我自己就是从单机训练一路摸到多机多卡踩了不少坑之后才慢慢把这张地图拼完整的。这篇笔记不追求面面俱到而是想帮你先建立一个清晰的知识坐标系再往里填细节。适合刚接触大模型工程、想从只跑模型进阶到懂系统的工程师也适合需要跟算法同学对焦、不想被术语唬住的运维和平台开发。我先把这套东西拆成六个层面来记全景分层、计算层、网络与存储、平台调度、推理服务、以及学习路径。每一层我都会讲清楚它解决什么问题、关键概念有哪些、实际干活的时候该盯哪些指标。1. 先从全景图说起AI Infra 的四层结构与分工边界最开始我犯过一个错误想一口气把整个AI Infra学完结果看了两天资料直接劝退。后来才明白这东西不是一个独立学科而是分布式系统、高性能计算、编译原理、操作系统、网络工程在跑大模型这件事上的交汇。如果硬要画一张全景图我会把AI Infra分成这么几个圈层硬件层GPU、HBM显存、NVLink、InfiniBand/RoCE网卡、NVMe存储、液冷/风冷机柜。系统层驱动与CUDA运行时、容器与镜像、Kubernetes调度、GPU共享与隔离MIG、时间片。框架与中间件层PyTorch/DeepSpeed/Megatron-LM、分布式通信库NCCL/RCCL、并行策略数据并行、张量并行、流水线并行、任务编排Ray、Volcano、Internal Scheduler。服务层模型推理引擎vLLM、TGI、TensorRT-LLM、网关与负载均衡、KV Cache管理与量化、可观测性Metrics/Tracing/Logging。这四个圈层层层依赖又各有各的痛点。硬件层决定你能跑多大的模型系统层决定混合编排的弹性框架层决定单/多卡训练的效率服务层决定线上业务的延迟和成本。我在整理笔记时发现很多人在某个圈层扎得很深却缺乏上下游视角。比如搞推理的只知道调vLLM参数却不理解底层显存带宽对prefill和decode的差异意味着什么搞平台的只关注K8s调度却不知道一个模型checkpoint有多大、加载要多长时间。所以我特别建议无论你现在在哪个岗位先把四个圈层都过一遍哪怕只是知道这东西是干什么的、卡在哪儿也很有价值。从分工边界上看AI Infra工程师要做的事情其实就三件把算力调度好、把数据喂得快、把模型跑得稳。听起来朴素干起来全是细活。接下来我们从上往下说先说最直接影响训练性能和成本的计算层。2. 计算层的关键账本GPU 选型、集群规模与算力利用率我先说一个反直觉的结论AI Infra最容易忽视的成本黑洞不是GPU买贵了而是算力利用率长期低于50%你还在以为机器不够用。我见过不少团队一报训练需求就说要加卡结果一查监控GPU利用率只有三成多大量的时间花在了数据加载等待、通信同步、以及由于并行策略配置不当导致的空闲气泡bubble。所以学习AI Infra的第一步不是研究哪块卡跑分最高而是学会算账。2.1 选卡标准账面算力 vs 有效算力我们常拿FP16 TFLOPS来看显卡纸面性能但真正训练大模型时瓶颈往往不在纯计算而在显存容量和通信带宽。比如一块主流卡账面算力看起来很高可一旦模型权重加优化器状态加上KV Cache总量超过显存就得做梯度重计算activation checkpointing或者张量并行切分这两个操作都会拉低有效利用率。我自己的选型逻辑是三步估单个模型副本的显存占用。以7B参数为例FP16权重大约14GBAdam优化器状态再加28GB左右梯度再来14GB总共大概56GB起步。这还没算激活值所以你就能理解为什么很多团队做微调宁可上LoRA也不全参。算通信与计算的配比。单机八卡靠NVLink内部互联机间靠网卡如果模型切分的维度导致每步迭代都要跨机通信那网卡带宽就会成为硬瓶颈。看生态兼容性。CUDA/FlashAttention/FlashMLA这些底层算子是不是对这套硬件做了充分优化会让训练工程省很多事。2.2 算力利用率的几个关键指标我每次给团队做训练性能排查必看四个数指标含义合理区间参考训练吞吐Tokens/Second/GPU每条GPU每秒处理的token数量取决于模型规模和序列长度看趋势不看绝对值GPU计算利用率SM利用率流处理器有没有真正在算训练集中且无数据饥饿时应在80%以上显存带宽利用率实际读写带宽占理论值的比例短序列场景尤其重要低说明kernel没喂饱通信等待占比一步迭代里等NCCL的时间占比大规模并行时要尽量压到20%以下这里有个经验如果GPU利用率低先查数据加载管线再查通信最后再怀疑自己的并行策略。顺序反过来很容易浪费时间。数据管线的问题通常症状明显——你会在nvidia-smi里看到GPU利用率像锯齿一样波动而CPU和磁盘IO却跑满了。2.3 集群规模别一上来就奔千卡很多人研究AI Infra张口闭口千卡集群。但我的建议是从64卡规模理解分布式训练的所有核心问题已经完全够了。千卡集群遇到的大部分问题网卡拥塞、checkpoint写爆存储、调度爆炸在64卡的规模上已经有苗头。先在小规模上把数据加载、容错、监控这些基本功练扎实再往大做才是正路。这一层的笔记核心是一句话计算层的核心账本是单卡效率和集群效率的乘积不是堆卡数量。3. 网络与存储模型并行时代的隐形瓶颈如果说GPU是AI Infra的发动机那网络和存储就是供油系统。Model参数越来越大以后单机放不下模型成了常态于是模型并行成了绕不开的话题。而模型并行的一个重要代价就是每一轮迭代都要在卡与卡、机与机之间同步梯度。你把训练任务从单机八卡扩到跨机只要网络跟不上跑出来的加速比会让你怀疑人生。3.1 RDMA为什么普通以太网喂不饱GPU常规TCP协议栈在传输小包数据时有大量内核中断和内存拷贝开销GPU之间的梯度同步动辄几十GB级吞吐普通以太网根本扛不住。所以现代AI集群普遍用RDMA远程直接内存访问常见的两种实现是InfiniBand和RoCERDMA over Converged Ethernet。前者贵但省心后者是把RDMA塞进以太网相对省钱但需要开PFC流控并配置ECN避免拥塞丢包。我学习时用一个类比来理解普通网络像是快递员送货每件都要扫码登记RDMA像是货车直接把仓库门对门对接省掉了中间环节。搞懂这个你就明白为什么网卡叫HCA/RDNS不叫普通网卡也明白为什么交换机不能随便买——它必须支持无损网络。3.2 网络拓扑与通信模式训练集群的典型拓扑是两级/三级Fat-Tree或叶脊结构。机架内带宽最强跨机架次之跨POD最弱。分布式框架的通信模式主要有集合通信All-Reduce梯度聚合最常用的操作。怎么避免跨机通信频繁发生直接决定了扩展效率。点到点通信P2P流水线并行的stage间传activation和梯度。稀疏通信MoE模型的专家分发通信模式更加随机对网络压力更大。我的笔记里画了一条关系链并行策略决定通信模式通信模式决定网络需求网络需求决定硬件选型。所以讨论网络之前先想明白你的并行策略是几路张量并行、几路流水线并行、几路数据并行。3.3 存储checkpoint 才是被低估的大佬很多人都把注意力放在网络却忽略了存储层。大模型训练每隔一段时间就要落盘checkpoint动辄几百GB甚至上TB。像一个大模型checkpoint以TB级计往普通NFS写一次可能要几分钟如果训练崩溃需要回滚加载checkpoint的几分钟甚至更久会直接变成纯GPU空闲时间。存储这块我的建议是分层设计热路径训练节点本地NVMe做数据缓存与临时输出。温路径高性能并行文件系统承载checkpoint和数据集要重点看聚合带宽。冷路径对象存储存放历史实验记录和归档数据。另外还有一个常被提的优化方向是异步checkpoint把模型权重复制到CPU内存再由后台线程慢慢写盘这样主训练线程不阻塞。但异步checkpoint有个坑它要求CPU内存够大而且写盘期间如果节点崩溃恢复点可能不一致需要框架层做好屏障处理。我建议你在笔记里单独列一个容错小节把自动恢复、快照一致性、掉线节点替换这三类问题放一起学。4. 训练平台资源调度、容错与实验管理有了硬件、网络、存储你还需要一个能把任务落到这堆资源上的大脑。训练平台这件事市面上的方案大致可以分三档直接用Kubernetes 自定义调度用Argo/Volcano这类批处理框架或者直接用云厂商的托管平台。我学习的时候发现大部分自建平台的复杂度都出在三个地方配额管理、抢占与容错、实验追踪。4.1 调度策略从按资源申请到按拓扑感知K8s原生的调度器对AI训练任务并不友好。它只知道你申请了几个GPU却不知道你这几个GPU最好放在同一台机器上因为NVLink。多机任务还希望机架间距离尽量近否则NCCL通信会因为跨交换机跳数变多而变慢。这就是拓扑感知调度Topology-Aware Scheduling也叫Gang Scheduling或coscheduling要解决的问题。我把调度这块理解为订酒店你带了8个人希望住同一层楼且房间相邻但酒店系统只按你有几张床来安排结果把你拆到不同楼层你串门聊天光走路就累得不行。同理没有拓扑感知的任务调度通信开销分分钟吃掉并行收益。实际工程中一个任务通常有最小并行单元比如8卡一个Pod调度器按这个单元整体分配避免半路只拿到几块卡、任务卡死等待。这里有个工程细节如果你用K8s建议看看volcano或者koordinator这类增强调度器它们在队列、抢占、拓扑亲和这些维度上比默认调度器完整很多——但底层的kube-scheduler扩展机制也得懂出了问题你要能定位到是调度策略拒绝还是资源不足拒绝。4.2 容错断点续训的真正难点分布式训练有一个很讨厌的特性——一份训练任务的任一节点挂了整个任务就废了。框架虽然实现了All-Reduce但通信协议有一方断连同步就会超时。所以平台层面必须做节点健康检查、任务重启、checkpoint自动回滚。断点续训的难点在于重启之后的恢复粒度。如果你每隔N步存一次checkpoint那么崩溃后就回到N步之前这部分算力就浪费了。优化方向是更频繁地存小checkpoint或者分层checkpoint——优化器状态和模型参数分开存频率可以不同。实践里有些团队还会做Preemption处理就是节点要被收回之前预留缓冲时间让任务自己保存状态类似云平台的Spot实例机制。这些机制学起来并不难难在你想清楚恢复时间目标RTO是多少以及允许丢失多少步训练进度RPO。4.3 实验追踪被忽略的科研基础设施很多笔记都不提实验追踪但我认为它是AI Infra不可或缺的一环。训练大模型是烧钱的实验如果连哪份代码、哪个超参、哪份数据集产出了哪个指标都追踪不了你根本没法复盘。早期的MLflow、现阶段的Weights Biases、以及开源的TensorBoard核心能力都是同一件事把实验过程固化为可复现的数据。Infra工程师要做的是在平台层面默认接入这套追踪让训练任务启动时自动记录代码commit、环境镜像、数据集版本、关键指标曲线。我自己甚至会建议你重点看一个方向数据集版本管理。很多人只盯模型版本但训练效果变了往往是因为数据悄悄变了。这块虽然不直接属于传统Infra的范畴但在AI Infra的语境里它跟存储和CI/CD强相关是值得单独记一章的内容。5. 推理侧的基础设施从批处理到高并发服务的架构迁移前四层讲的主要是训练这个场景但AI Infra的第二大块是推理。推理Infra跟训练有一个本质区别训练是吞吐优先推理是延迟、吞吐、成本三者同时算账。同一个模型你怎么部署、怎么切分、怎么优化的收益可能比模型效果本身还大。5.1 推理引擎的演化逻辑从朴素部署到KV Cache复用最开始大家推理大模型就是一个模型实例接一个请求扛不住并发就横向多副本。但这套方式对Transformer系模型有巨大浪费——每个请求都要重新计算之前token的key和value。后来业界普遍引入PagedAttention和KV Cache管理才让推理吞吐上了一个大台阶。我的理解是这样的KV Cache像是临时便签记录已经看过的上下文。朴素方案里每个请求自己带一份全新便签重复内容反复重算vLLM的PagedAttention则像给便签分页管理内存按需分配、还能在多个请求之间共享公共前缀。这两者的吞吐差距在长上下文场景下可能有好几倍。5.2 Prefill 与 Decode 的算力错配大模型推理有个很有意思的现象同一个GPU在处理输入prefill时算力吃满像训练在处理输出decode时却因为逐token生成变成显存带宽受限。如果你用一个引擎同时服务流式对话请求你会发现GPU计算利用率可能很波动一会儿99%一会儿20%不到。工程上的解法是把prefill和decode解耦如分离式推理架构、动态batching、以及用Continous Batching机制让不同请求穿插利用GPU的间隙。笔记里我建议对比三个引擎维度支持的量化格式、是否支持张量并行多卡推理、对动态shape批处理的支持能力。推理侧还有一个躲不开的话题是量化。从FP16到INT8再到INT4显存占用一路下降吞吐一路上升但效果损失也一路爬升。我的建议是不必追求极端低bit先看你的场景对困惑度Perplexity或者下游任务指标是否敏感再决定要不要往下压。工程上量化带来的显存节省有可能让你把一批请求塞进更少卡里但如果你要做多卡张量并行推理要特别留意通信和量化粒度的配合。5.3 服务治理网关、限流、灰度与可观测性模型推理服务本质上是高并发微服务所以常规服务的网关、限流、灰度发布这些概念照样适用。但多了一层模型版本要和配置一起做灰度而且推理服务的负载曲线不是匀速的可能突然来一波热点请求。如果你直接把所有请求都打到同一个推理引擎很有可能会把KV Cache打爆、造成OOM。我个人的实操建议是三层防护接入层限流按用户维度或按并发维度把突发流量挡在外面避免拖垮整批实例。队列与弹性伸缩对延迟不敏感的场景用队列把请求削峰填谷对实时场景通过HPA/Prometheus指标自动扩缩容。这里要重点观察in-flight请求数和GPU显存占用单纯看QPS太粗了。实例隔离把不同的模型、不同租户拆到不同实例组避免相互影响。可观测性这块除了常规的延迟、错误率、QPS我建议把 GPU利用率、显存占用、KV Cache命中率单独画成面板。这几个指标的联动能帮你快速判断瓶颈是算力不够、显存不够还是并发策略没生效。6. 学习路径与知识地图我整理的最后一份干货既然这篇笔记叫全景学习笔记最后我想把学习路径也交代一下分享我是怎么一步步把这张图啃下来的省得你再走弯路。6.1 阶段划分从跑通到理解到扩展我自己把学习AI Infra分为四个阶段每个阶段都有明确的输出物阶段目标典型输出建议用时第一阶段跑通在多卡上成功训练一个中小模型能解释DDP与DeepSpeed Zero的配置差异1~2周第二阶段理解掌握并行策略的原理和通信开销能画出一张并行策略-通信模式-硬件需求映射表1个月第三阶段优化定位并解决训练性能瓶颈能给出GPU利用率低下的排查报告持续迭代第四阶段架构设计或运维一个中型训练平台能独立配置调度、容错、监控、推理服务一套链路3个月以上我不建议一上来就啃论文。先做实验再做性能剖析最后才去看论文里的优化公式。理论只有在你能联系到实际痛点时才会真正留下来。6.2 我踩过的最深的坑分享几个我真实踩过、且常规文档里很少写的坑也许对你排查问题有直接帮助坑一NCCL超时的原因往往不在网络而在存储或CPU。有一次训练任务报通信超时我查了几天网络配置最后发现是checkpoint写入把存储IO占满连带影响了网络的元数据操作。排查通信问题一定要先排除邻居干扰。坑二容器里忘了配置IPC共享内存。PyTorch的DataLoader多进程要用共享内存在K8s Pod里如果shmSize太小你会看到奇怪的数据加载停顿。坑三监控指标采样间隔太粗。默认15秒采样一次根本看不到训练迭代里算-等-算的微观波动。排障时把采样间隔调到1秒以内很多问题图景立刻清晰。坑四推理压测只在短序列流量下做。真实业务的序列长度可能是几百到几千不等长短混合下的显存分配和批处理效率差异巨大只测固定长度会严重高估系统容量。6.3 最后的知识地图速查如果只能保留一页笔记我会留下这张速查清单作为日常对照训练任务慢先查数据加载再查GPU利用率再查通信比例最后看并行策略。显存不够先考虑激活重计算再考虑ZeRO优化器状态切分再考虑张量并行最后才考虑降batch size。推理延迟高先看prefill和decode的占比再看KV Cache配置再看批量大小最后看模型量化。任务老崩溃先看节点健康再看checkpoint恢复逻辑再看网络波动最后看是否是集群资源争抢。集群利用率低先看调度器是否支持拓扑感知再看排队策略最后看是否有大量被占但没跑满的时间片。这张速查可能不求解百发百中但胜在能让你快速定位大方向。我到现在做性能排查仍然先过一遍这条思路链再去翻具体日志和指标。这个领域更新太快今天写的内容可能过半年就有新名字、新框架冒出来。但底层的思想方法不会变理解每一层在解决什么问题、上下层如何互相制约你就能在任何一个新工具出现时快速把它安放进自己的知识地图里。希望这份笔记能帮你少走一些我走过的弯路。