1. 为什么需要一张AI Infra知识全景图搞AI Infra的人大概都有过这种体验面试时被问到“一个千亿参数模型训练时显存怎么分布”脑子里能蹦出七八个知识点但就是串不成一条线工作中遇到训练吞吐上不去排查了半天发现是数据加载的IO瓶颈但之前一直盯着GPU利用率在看。这种“知识点散落一地、连不成网”的状态几乎是每个AI Infra从业者都会经历的阶段。我自己是从传统后端转过来的刚开始接触分布式训练的时候光是Megatron的并行策略就啃了快两周TP、PP、DP、EP这几个概念单独看都懂但放到一个真实的训练任务里它们怎么组合、怎么互相影响、通信开销怎么算完全是懵的。后来踩了不少坑也跟不少同行聊过发现大家的问题都差不多不是缺资料而是缺一张能把所有东西串起来的“全景图”。这份学习笔记就是干这个的。它不打算把每个知识点都讲透——那是教科书的事——而是要把AI Infra这个领域的知识骨架搭出来让你知道每个模块在整体中的位置、它解决什么问题、跟其他模块怎么交互。适合谁看如果你是刚转行做AI Infra的工程师或者做算法想往Infra方向靠又或者面试前需要快速梳理知识体系这份笔记应该能帮上忙。如果你已经是资深Infra工程师也可以拿它当个checklist看看有没有遗漏的角落。2. AI Infra知识体系的整体拆解2.1 从训练到推理的完整链路AI Infra说到底就是一件事让模型能在硬件上高效地跑起来。这个“跑”分两个阶段——训练和推理两者的优化目标完全不同。训练追求的是吞吐量和收敛速度推理追求的是延迟和成本。但它们的底层依赖是共通的计算、存储、通信、调度。我习惯把AI Infra的知识体系分成五层来看。最底层是硬件层包括GPU、NPU、TPU这些加速器以及它们之间的互联方式NVLink、NVSwitch、InfiniBand、RoCE。往上是通信层集合通信库NCCL、HCCL、Gloo和通信原语AllReduce、AllGather、ReduceScatter、AlltoAll。再往上是并行策略层数据并行、张量并行、流水线并行、专家并行以及它们的组合方式。然后是框架与引擎层训练框架Megatron-LM、DeepSpeed、FairScale和推理引擎vLLM、TensorRT-LLM、SGLang。最上面是平台与调度层Kubernetes、Slurm、Ray以及资源调度、容错、弹性伸缩这些。这五层不是孤立的。举个例子你选了什么并行策略直接决定了通信层的压力分布通信层的瓶颈又反过来限制并行策略的选择。再比如推理引擎的PagedAttention机制本质上是在存储层做文章但它对上层调度的影响是巨大的——显存利用率上去了单卡能跑的并发数就多了调度策略就得跟着调。2.2 训练与推理的核心差异很多人会把训练和推理混在一起学结果越学越乱。我的建议是先分开理解再找共通点。训练的核心矛盾是“显存墙”和“通信墙”。显存墙指的是单卡显存放不下整个模型的状态参数、梯度、优化器状态所以需要切分。通信墙指的是切分之后卡间需要同步数据通信开销可能吃掉计算收益。这两个矛盾催生了各种并行策略和显存优化技术ZeRO、梯度检查点、混合精度。推理的核心矛盾是“延迟”和“吞吐”的权衡。延迟要求你尽快返回结果吞吐要求你同时服务尽可能多的请求。这两个目标天然冲突——batch size大了吞吐高但延迟也高。所以推理引擎的核心工作就是在这两者之间找平衡点PagedAttention、Continuous Batching、投机解码这些技术都是围绕这个展开的。但两者也有共通的地方。比如量化技术训练里有FP8训练推理里有INT8/INT4量化底层都是数值精度的取舍。再比如算子优化FlashAttention在训练和推理里都用得上只是侧重点不同。2.3 知识全景的构建逻辑我构建这张全景图的逻辑是“以问题为导向”。不是按技术分类来罗列而是按“你会遇到什么问题”来组织。比如显存不够了你能想到哪些手段切分并行策略、压缩量化、梯度检查点、卸载Offload到CPU/NVMe、复用PagedAttention、显存池。再比如训练慢了排查思路是什么先看GPU利用率如果低就看是数据加载瓶颈还是通信瓶颈如果高就看是不是算子效率问题。这种问题导向的组织方式比按技术分类更贴近实际工作场景。下面这张表是我整理的AI Infra核心模块与对应问题的映射关系可以当作一个速查表来用。核心模块解决的核心问题关键技术典型工具并行策略单卡放不下、算得慢TP/PP/DP/EP/ZeROMegatron-LM、DeepSpeed通信优化卡间同步开销大集合通信、通信重叠、拓扑感知NCCL、HCCL显存优化显存不够用梯度检查点、Offload、量化DeepSpeed、bitsandbytes计算优化算子效率低FlashAttention、融合算子、编译优化Triton、TVM、torch.compile推理引擎延迟高、吞吐低PagedAttention、Continuous BatchingvLLM、TensorRT-LLM调度与编排资源利用率低弹性伸缩、容错、优先级调度Kubernetes、Ray、Slurm存储与IO数据加载慢数据分片、预取、缓存WebDataset、DALI这张表不是让你背的而是让你在遇到问题时能快速定位到相关模块。比如你发现训练时GPU利用率只有30%查表就知道该往“存储与IO”和“通信优化”两个方向去排查。3. 核心模块深度解析与实操要点3.1 并行策略怎么切模型才合理并行策略是AI Infra里最核心也最容易让人晕的部分。我刚开始学的时候把TP、PP、DP的定义背得滚瓜烂熟但一到实际配置就不知道该怎么选。后来总结了一个简单的判断逻辑先看模型能不能单卡放下放不下就考虑TP或PP再看单卡能不能算得过来算不过来就加DP如果模型有MoE结构再考虑EP。数据并行DP是最简单的每张卡持有完整模型副本处理不同的数据batch梯度通过AllReduce同步。优点是实现简单、扩展性好缺点是显存冗余——每张卡都要存一份完整的模型状态。所以DP适合模型能单卡放下、但计算量大的场景。张量并行TP是把单个算子比如矩阵乘切到多张卡上。比如一个[4096, 4096]的权重矩阵可以按列切成两份每张卡存[4096, 2048]。TP的优点是单层内的计算和显存都被切分了缺点是通信频繁——每层前向和反向都需要AllReduce。所以TP通常限制在单机内NVLink带宽高跨机TP通信开销太大。流水线并行PP是把模型按层切分不同卡负责不同的层。比如32层的模型切成4段每张卡负责8层。PP的优点是通信量小只在层边界传激活值缺点是存在流水线气泡——前面的卡在算的时候后面的卡在等。为了减少气泡需要配合微批处理micro-batch和调度策略1F1B、Interleaved 1F1B。专家并行EP是针对MoE模型的把不同的专家分配到不同的卡上。EP的难点在于路由——每个token需要根据门控网络的结果发送到对应的专家卡上这涉及AlltoAll通信。EP的通信模式比TP和PP都复杂调优难度也更大。实际配置时我常用的组合是单机内用TP8NVLink全互联跨机用PP2~4DP根据总卡数除一下。比如64卡训练一个70B模型可以配TP8、PP4、DP2。这个配置不是最优的但通常是个合理的起点。注意TP和PP的切分维度是模型维度DP的切分维度是数据维度。配置时总卡数 TP × PP × DP。如果用了EP总卡数还要乘以EP的倍数但EP通常和DP共享卡。3.2 通信优化别让卡间同步拖后腿通信是分布式训练的隐形杀手。我见过太多案例GPU利用率上不去排查半天发现是通信在等。通信优化的核心思路就两个减少通信量和隐藏通信开销。减少通信量的手段包括梯度压缩比如只传Top-K的梯度、通信量化FP16转INT8、通信调度把多个小通信合并成一个大通信。隐藏通信开销的手段主要是通信重叠——在计算的同时进行通信让通信时间被计算时间覆盖。NCCL是NVIDIA的集合通信库也是目前最常用的。它支持Ring AllReduce、Tree AllReduce、NVLSNVLink SHARP等多种算法会根据拓扑自动选择。但自动选择不一定最优特别是在跨机场景下。我遇到过跨机AllReduce慢的问题后来手动指定了Ring算法并调整了chunk size性能提升了30%。调NCCL有几个关键环境变量NCCL_ALGO指定算法NCCL_PROTO指定协议LL、LL128、SimpleNCCL_BUFFSIZE调整缓冲区大小NCCL_NTHREADS调整线程数。这些参数没有万能配置需要根据实际拓扑和消息大小来调。一般来说小消息用LL128协议延迟低大消息用Simple协议带宽高。实操心得调NCCL之前先用nccl-tests跑一遍AllReduce基准测试了解当前拓扑下的理论带宽和实际带宽差距。如果实际带宽只有理论值的50%以下说明配置有问题。3.3 显存优化把每一MB都用在刀刃上显存是训练里最稀缺的资源。一个70B模型的训练状态参数梯度优化器状态在FP32下需要70×4×3840GB加上激活值轻松超过1TB。单卡80GB的A100/H100根本放不下所以必须做显存优化。显存优化的手段我按“性价比”排个序混合精度训练FP16/BF16是最简单也最有效的直接把显存占用砍半梯度检查点Gradient Checkpointing用计算换显存把激活值从O(n)降到O(√n)代价是增加30%左右的计算量ZeRO把优化器状态、梯度、参数分片到多卡上ZeRO-1只分片优化器状态ZeRO-2加分片梯度ZeRO-3加分片参数Offload把暂时不用的数据放到CPU内存或NVMe上代价是PCIe带宽成为瓶颈。这些手段可以叠加使用。我常用的组合是BF16 梯度检查点 ZeRO-2。这个组合在大多数场景下能平衡显存和速度。如果还不够再上ZeRO-3或Offload。注意ZeRO-3和TP同时使用时需要小心因为两者都在切分参数可能导致通信量爆炸。实践中通常ZeRO-3配合PP使用而不是TP。3.4 推理引擎延迟与吞吐的平衡术推理引擎的优化目标跟训练完全不同。训练可以慢慢跑推理必须快。但“快”有两个维度单请求延迟低TTFT和TPOT小和整体吞吐高每秒处理的token数多。这两个目标往往冲突。vLLM是目前最流行的开源推理引擎它的核心创新是PagedAttention。传统推理引擎为每个请求预分配连续的KV Cache显存导致碎片化严重实际利用率可能只有20%-40%。PagedAttention把KV Cache切成固定大小的block按需分配显存利用率能到90%以上。显存利用率上去了单卡能同时服务的请求数就多了吞吐自然就高了。Continuous Batching是另一个关键优化。传统静态batching要等一个batch里所有请求都完成才能开始下一批导致短请求等长请求。Continuous Batching在每次迭代时动态加入新请求、移除已完成请求GPU利用率大幅提升。投机解码Speculative Decoding是降低延迟的手段。用一个小的草稿模型先猜几个token然后用大模型验证。如果猜对了就跳过这些token的计算相当于用小的计算量换大的计算量。实测下来在代码生成等确定性强的任务上投机解码能带来2-3倍的加速。实操心得vLLM的gpu_memory_utilization参数很关键默认0.9但如果你的模型有较大的激活值峰值可能需要降到0.85甚至0.8。另外max_num_seqs和max_num_batched_tokens需要根据实际请求长度分布来调不是越大越好。3.5 调度与编排让资源利用率最大化调度层是AI Infra里最容易被忽视但影响最大的部分。一个训练集群的GPU利用率很大程度上取决于调度策略。Kubernetes是目前最主流的容器编排平台但原生K8s对AI训练的支持并不好。训练任务通常需要gang scheduling所有Pod同时启动而K8s默认是逐个调度的容易导致死锁。所以通常需要配合Volcano、Kueue等调度器。弹性伸缩是另一个关键能力。训练任务可能因为各种原因中断硬件故障、抢占如果能自动恢复并继续训练能大幅减少人工干预。这需要checkpoint机制和故障检测的配合。注意弹性伸缩不是万能的。如果任务频繁中断checkpoint的开销可能超过收益。实践中通常设置一个阈值比如中断超过3次就暂停任务并告警。4. 实操过程与核心环节实现4.1 从零搭建一个分布式训练环境假设你手头有4台机器每台8张A100 80GB通过InfiniBand互联。目标是训练一个13B参数的模型。下面是我会走的完整流程。第一步环境检查。先确认GPU驱动、CUDA版本、NCCL版本一致。用nvidia-smi topo -m查看拓扑确认NVLink和IB的连接情况。用ibstat确认IB链路状态。这一步很关键我遇到过因为一台机器NCCL版本不一致导致训练随机hang住的情况排查了两天才发现。第二步安装依赖。PyTorch、Megatron-LM或DeepSpeed、Apex可选、FlashAttention。版本兼容性是个大坑建议用NGC容器或者conda锁定版本。我通常用nvcr.io/nvidia/pytorch:23.10-py3这个镜像里面已经配好了大部分依赖。第三步配置并行策略。13B模型在BF16下参数占26GB梯度26GB优化器状态Adam52GB合计104GB。单卡80GB放不下所以需要切分。我选TP4、PP2、DP4总卡数4×2×432刚好4台机器。TP4在单机内NVLinkPP2跨机IBDP4跨机。第四步启动训练。用torchrun或mpirun启动。关键参数包括--tensor-model-parallel-size 4、--pipeline-model-parallel-size 2、--micro-batch-size 4、--global-batch-size 512。global batch size micro batch size × DP × gradient accumulation steps。这里512 4 × 4 × 32所以gradient accumulation steps32。第五步监控与调优。启动后先看GPU利用率如果低于80%就要排查。用nvidia-smi dmon看GPU利用率用nsys或py-spy做profiling。常见瓶颈包括数据加载慢用DALI或WebDataset加速、通信慢调NCCL参数、算子效率低用FlashAttention替换原生Attention。4.2 推理服务的部署与压测推理部署比训练简单但坑也不少。以vLLM部署一个7B模型为例。第一步模型转换。如果模型是HuggingFace格式vLLM可以直接加载。如果是自定义格式需要先转换。注意dtype的选择——FP16通常比BF16快但BF16数值稳定性更好。第二步启动服务。python -m vllm.entrypoints.openai.api_server --model /path/to/model --tensor-parallel-size 2 --gpu-memory-utilization 0.9。tensor-parallel-size根据模型大小和卡数来定7B模型单卡就能跑但为了吞吐可以TP2。第三步压测。用vllm.benchmarks.benchmark_serving或locust做压测。关键指标是TTFT首token延迟、TPOT每token延迟、吞吐tokens/s。压测时要模拟真实的请求长度分布不要只用固定长度。第四步调优。如果TTFT高可能是prefill阶段慢可以调大max_num_batched_tokens。如果TPOT高可能是decode阶段慢可以调大max_num_seqs。如果吞吐低可能是显存不够导致batch size上不去可以降低gpu_memory_utilization或启用量化。实操心得vLLM的--enable-prefix-caching对多轮对话场景效果很好能缓存系统prompt的KV Cache减少重复计算。但要注意如果系统prompt经常变缓存命中率会很低反而增加开销。4.3 性能瓶颈的定位与解决性能瓶颈的定位有一套通用的方法论。我通常按这个顺序排查先看GPU利用率再看显存占用再看通信开销最后看算子效率。GPU利用率低70%通常意味着有等待。等待可能来自数据加载CPU瓶颈、通信网络瓶颈、或者同步锁竞争。用nvidia-smi dmon看利用率波动如果周期性波动可能是数据加载如果持续低可能是通信。显存占用高但利用率低可能是batch size太小。这时候要检查是不是显存碎片化严重或者KV Cache分配不合理。通信开销大用nccl-tests确认实际带宽。如果实际带宽远低于理论带宽检查拓扑配置和NCCL参数。算子效率低用nsys做profiling看哪个kernel耗时最长。常见的低效算子包括LayerNorm、Softmax、Attention。用融合算子或FlashAttention替换。5. 常见问题与排查技巧实录5.1 训练类问题速查问题现象可能原因排查方法解决方案训练loss不下降学习率太大、数据有问题、初始化不对检查数据标签、打印梯度范数调小学习率、检查数据管道训练随机hang住NCCL版本不一致、网络抖动、死锁检查各节点NCCL版本、IB链路状态统一版本、增加超时重试GPU利用率周期性波动数据加载瓶颈用py-spy看CPU占用用DALI/WebDataset加速、增加预取显存OOMbatch size太大、激活值峰值高用torch.cuda.memory_summary()减小batch size、开梯度检查点通信时间占比高并行策略不合理、NCCL配置差用nsys看通信kernel耗时调并行策略、调NCCL参数5.2 推理类问题速查问题现象可能原因排查方法解决方案TTFT高prefill阶段慢、排队严重看prefill和decode的时间占比调大max_num_batched_tokens、加卡TPOT高decode阶段慢、KV Cache碎片化看KV Cache利用率调大max_num_seqs、用PagedAttention吞吐低batch size上不去、显存不够看显存占用和batch size量化、调gpu_memory_utilization输出质量差量化精度损失、采样参数不对对比FP16和量化后的输出调量化配置、调temperature/top_p服务不稳定显存泄漏、请求超时监控显存变化、看日志定期重启、加超时机制5.3 独家避坑技巧技巧一训练前先跑一个mini batch。用1/100的数据跑一遍完整流程确认没有shape错误、没有OOM、没有通信问题。这个习惯帮我省了无数时间。技巧二checkpoint要存优化器状态。只存模型参数不存优化器状态恢复训练后loss会跳变。我见过有人为了省空间只存参数结果恢复后训练效果差很多。技巧三推理服务要设max_model_len。不设的话vLLM会按模型的最大长度预分配KV Cache显存浪费严重。根据实际业务需求设一个合理的值比如4096或8192。技巧四NCCL超时要设长一点。默认的NCCL超时是30分钟但在大规模训练中一次AllReduce可能超过这个时间。设成NCCL_TIMEOUT180030分钟或更长。技巧五用wandb或tensorboard记录一切。loss、学习率、梯度范数、GPU利用率、显存占用、通信时间。出问题的时候这些日志是唯一的线索。6. 知识全景的持续更新与扩展AI Infra这个领域变化太快了。去年还在讨论ZeRO-3和TP的组合今年MoE和EP就成了热点。去年vLLM还是唯一选择今年SGLang和TensorRT-LLM就追上来了。所以这张全景图不是静态的需要持续更新。我自己的做法是每季度做一次知识梳理看看有没有新的技术点需要加进来。关注几个核心来源顶会论文NeurIPS、ICML、OSDI、NSDI、开源项目更新Megatron-LM、DeepSpeed、vLLM、以及一线工程师的博客和分享。另外知识全景的价值在于“连接”而不是“罗列”。我见过很多人收集了一堆资料但都是孤立的点。真正的理解是知道TP和PP为什么不能随便组合、知道PagedAttention为什么能提升吞吐、知道NCCL的Ring和Tree算法分别适合什么场景。这种“知道为什么”的能力才是AI Infra工程师的核心竞争力。最后分享一个我自己的学习习惯每学一个新知识点就试着把它挂到已有的知识树上。比如学PagedAttention的时候我会想它跟操作系统的虚拟内存分页有什么相似之处跟ZeRO的显存分片有什么不同。这种类比和对比能让知识记得更牢也更容易在需要的时候调用出来。