1. 从“调包侠”到系统思维为什么AI Infra值得你花时间啃下来很多人做AI项目起步就是pip install一堆框架然后照着示例代码跑通一个Demo就觉得自己“入门”了。但一旦模型规模上去、并发请求变多、训练任务从单卡变成多机多卡各种诡异的问题就会像潮水一样涌出来——显存莫名其妙爆掉、GPU利用率长期在30%以下徘徊、分布式训练跑着跑着就hang住、推理服务的P99延迟忽高忽低。这些问题的根因往往不在模型代码本身而在于你对底层基础设施的理解深度不够。AI Infra也就是人工智能基础设施说白了就是支撑AI任务从开发到落地全流程的那套底层系统。它涵盖的范围非常广从最底层的GPU/TPU等加速硬件到CUDA、cuDNN这些计算库再到PyTorch、TensorFlow等深度学习框架往上还有分布式训练框架、模型推理引擎、资源调度平台、存储与网络方案以及贯穿始终的监控与性能分析工具链。你可以把它想象成一座冰山——水面之上是大家都能看到的模型结构和算法创新水面之下则是庞大而复杂的工程体系而恰恰是这部分决定了你的AI系统能不能真正跑得快、稳得住、扩得开。这篇学习笔记适合谁看如果你已经会用PyTorch或TensorFlow训练一些中小模型但对“为什么我的GPU利用率上不去”“分布式训练到底怎么切分数据和模型”“推理服务怎么优化延迟和吞吐”这些问题感到模糊那这篇内容就是为你准备的。我不会只给你一堆名词解释而是会从实际工程问题出发把每个知识模块背后的“为什么”讲清楚再配上可操作的排查思路和配置示例。整个知识全景我会拆成几个核心板块来讲每个板块都会告诉你它解决什么问题、关键概念是什么、实际工作中怎么用、容易踩哪些坑。先给你一个全局视角。AI Infra的知识体系大致可以分成五层硬件层、计算库层、框架层、平台层和应用层。硬件层包括GPU、TPU、NPU以及高速互联技术如NVLink、InfiniBand计算库层是CUDA、cuDNN、NCCL这些直接和硬件打交道的库框架层是PyTorch、JAX、MindSpore等平台层涵盖训练平台、推理平台、调度系统、特征存储等应用层则是具体的业务场景落地。每一层都有其核心知识点和常见问题而真正的高手往往是在层与层之间的交界处有深刻理解的人。比如你知道PyTorch的DataLoader的num_workers设置为什么会影响GPU利用率吗这背后涉及的是数据加载、CPU预处理、PCIe传输和GPU计算之间的流水线配合问题。接下来我就按照这个分层逻辑把每一层的核心知识、实操要点和避坑经验逐一展开。2. 硬件与计算库层GPU利用率上不去的根因往往在这里2.1 GPU架构与显存层次为什么你的模型总是OOM很多人遇到显存溢出OOM的第一反应是“模型太大了”然后就去减小batch size。但实际情况往往更复杂。要真正理解显存问题你得先搞清楚GPU的显存层次结构。以主流的NVIDIA GPU为例显存从大到小依次是全局内存HBM/GDDR、L2缓存、共享内存Shared Memory、寄存器文件。全局内存容量最大但访问延迟最高共享内存和寄存器则速度快但容量极小。CUDA编程的核心优化手段之一就是尽可能让数据在共享内存和寄存器中复用减少对全局内存的访问。那为什么训练时显存占用会远超模型参数本身的大小因为显存里不只有模型权重还有优化器状态、梯度、激活值activations、临时缓冲区等。以一个参数量为P的模型为例如果用Adam优化器做混合精度训练显存占用大致可以这样估算模型权重FP16占2P字节梯度FP16占2P字节Adam的动量和方差FP32各占4P字节再加上FP32的权重副本4P字节光这些就是16P字节。如果P是70亿7B那就是112GB已经超过单张80GB显卡的容量了。这还没算激活值——激活值的显存占用和batch size、序列长度、网络深度都成正比在长序列场景下可能比参数本身还大。所以当你遇到OOM时不要只盯着batch size。你可以从这几个方向排查第一检查是否开启了混合精度训练AMP这能直接把权重和激活值的显存占用减半第二考虑使用梯度检查点Gradient Checkpointing用计算换显存把激活值占用从O(n)降到O(sqrt(n))第三使用ZeRO系列优化技术把优化器状态、梯度、参数分片到多张卡上第四检查是否有内存泄漏比如在训练循环里不断累积loss列表却没有释放。我见过不少新手在验证阶段忘了加torch.no_grad()导致计算图不断累积显存一路飙升直到OOM。注意torch.cuda.memory_allocated()和torch.cuda.max_memory_allocated()是你排查显存问题的好朋友前者看当前占用后者看峰值占用。养成在关键节点打印这两个值的习惯能帮你快速定位显存增长的来源。2.2 CUDA与cuDNN版本匹配这件事比你想象的更重要CUDA是NVIDIA推出的通用并行计算平台cuDNN是基于CUDA的深度学习加速库。你安装PyTorch时实际上安装的是预编译好的、链接了特定CUDA和cuDNN版本的二进制包。版本不匹配是新手最容易踩的坑之一。比如你的驱动支持CUDA 12.1但你装的PyTorch是编译自CUDA 11.8的版本那可能也能跑但某些算子会回退到慢速路径甚至直接报错。怎么查版本用nvidia-smi看驱动支持的CUDA版本右上角那个CUDA Version用nvcc --version看当前安装的CUDA Toolkit版本用torch.version.cuda看PyTorch编译时用的CUDA版本。理想情况下这三个版本应该兼容。驱动版本决定了你最高能用哪个CUDA版本而PyTorch的预编译包则决定了你实际用的版本。如果你需要特定版本的CUDA最省心的方式是去PyTorch官网找对应的安装命令用conda或pip安装预编译包而不是自己从源码编译。cuDNN的版本匹配同样关键。cuDNN为卷积、池化、归一化等操作提供了高度优化的实现。不同版本的cuDNN对同一算子的优化程度可能差异很大。如果你发现某个卷积层的速度异常慢可以检查一下cuDNN的版本并尝试设置torch.backends.cudnn.benchmark True让cuDNN在第一次运行时自动搜索最适合当前输入尺寸的算法。不过要注意这个设置适合输入尺寸固定的场景如果输入尺寸变化频繁自动搜索的开销反而会拖慢速度。2.3 NCCL与多卡通信分布式训练的隐形瓶颈当你的训练从单卡扩展到多卡NCCLNVIDIA Collective Communications Library就成了关键角色。它负责实现All-Reduce、Broadcast、All-Gather等集合通信操作是数据并行训练中梯度同步的底层依赖。很多分布式训练“卡住”或“变慢”的问题根因都在NCCL通信上。一个常见的现象是单机多卡训练时GPU利用率很高但扩展到多机多卡后利用率反而下降了。这通常是因为跨机通信的带宽远低于机内通信。机内多卡通过NVLink或PCIe互联带宽可以达到几百GB/s而跨机通信走的是InfiniBand或以太网带宽可能只有几十到一百Gb/s。当梯度同步的数据量很大时通信时间就会成为瓶颈。这时候你需要考虑是否使用了梯度累积来减少通信频率是否使用了通信压缩技术如梯度量化是否合理设置了NCCL的环境变量比如NCCL_IB_DISABLE、NCCL_SOCKET_IFNAME、NCCL_DEBUGINFO来排查通信问题提示设置NCCL_DEBUGINFO后运行你的分布式训练脚本NCCL会打印出详细的通信拓扑和使用的传输方式。如果你看到它走了Socket而不是InfiniBand那说明IB网络配置可能有问题需要检查驱动和网络接口设置。3. 框架与训练层从单卡到千卡训练效率的优化逻辑3.1 PyTorch的DataLoader与数据流水线GPU饿肚子的真相GPU利用率低最常见的原因不是计算本身慢而是数据供给跟不上。你可以把GPU想象成一个超级大厨做一道菜只需要几毫秒但配菜员数据加载每道菜要花几十毫秒去准备食材那大厨大部分时间都在等着。PyTorch的DataLoader就是那个配菜员它的num_workers参数决定了有多少个配菜员同时工作。num_workers设多少合适这取决于你的数据预处理复杂度和CPU核心数。如果预处理很轻比如只是读取图片和归一化设成CPU核心数的一半到三分之二就够了如果预处理很重比如复杂的增强、解码、tokenization可以设得更多甚至超过CPU核心数因为I/O等待时CPU是空闲的。但也不是越大越好每个worker都会占用内存worker太多可能导致内存不足或进程调度开销过大。我的经验是从num_workers4开始逐步增加观察GPU利用率和数据加载时间的变化找到拐点。另外几个关键参数pin_memoryTrue可以把数据加载到锁页内存加速CPU到GPU的传输prefetch_factor控制每个worker预取多少个batchpersistent_workersTrue可以避免每个epoch结束后重新创建worker进程。还有一个容易忽略的点如果你的数据集很小能全部放进显存那直接用TensorDataset把数据搬到GPU上完全绕过DataLoader速度会快很多。3.2 混合精度训练省显存、提速度但要注意数值稳定性混合精度训练AMP是现在训练大模型的标配。它的核心思想是在前向和反向传播中用FP16或BF16进行计算但在参数更新时用FP32保持精度。这样做的好处很明显显存占用减半计算速度提升因为GPU的FP16算力通常是FP32的数倍同时通过损失缩放Loss Scaling技术来避免梯度下溢。PyTorch的torch.cuda.amp用起来很简单用autocast上下文管理器包裹前向传播用GradScaler来缩放损失和更新参数。但这里有几个坑要注意。第一不是所有操作都适合FP16比如softmax、layer norm、loss计算等对数值范围敏感的操作autocast会自动把它们保持在FP32你不需要手动干预。第二如果遇到梯度爆炸或NaN先检查是不是损失缩放的问题可以尝试减小初始缩放因子或使用动态损失缩放。第三BF16比FP16的数值范围更大不容易溢出如果你的GPU支持BF16如A100、H100优先用BF16。3.3 分布式训练策略数据并行、模型并行与流水线并行当单卡放不下模型或单卡训练太慢时就需要分布式训练。常见的策略有三种数据并行Data Parallelism、模型并行Model Parallelism和流水线并行Pipeline Parallelism。数据并行是最常用的每张卡持有完整的模型副本但处理不同的数据batch然后通过All-Reduce同步梯度。PyTorch的DistributedDataParallelDDP是官方推荐的数据并行方案比旧的DataParallelDP效率高很多因为DDP使用多进程而不是多线程避免了GIL的限制而且通信效率更高。使用DDP时每个进程独立加载数据通过DistributedSampler保证各进程看到的数据不重叠。模型并行则是把模型的不同层放到不同的卡上适合单层参数就大到单卡放不下的情况。但模型并行的通信开销大因为每层的前向和反向都需要跨卡传输激活值和梯度。流水线并行是模型并行的一种改进把模型按层切分成多个阶段每个阶段放在不同的卡上并通过微批次micro-batch技术让不同阶段的计算重叠起来减少GPU空闲时间。Megatron-LM和DeepSpeed是这方面比较成熟的框架。实际训练千亿参数模型时往往是三种并行策略的组合也就是所谓的3D并行。比如在节点内用张量并行把单个矩阵运算切分到多卡在节点间用流水线并行同时叠加数据并行。这种组合的配置非常复杂需要根据具体的模型结构、集群拓扑和网络带宽来调优。3.4 梯度累积与检查点小显存也能训大模型梯度累积Gradient Accumulation是一个简单但极其有用的技巧。它的思路是在显存有限的情况下用多个小batch分别计算梯度并累积等累积到一定步数后再统一更新参数。这样等效于用了更大的batch size但显存占用不变。实现上就是在训练循环里每算完一个batch的loss后除以累积步数再反向传播但不立即调用optimizer.step()而是等累积够了再更新并清零梯度。梯度检查点Gradient Checkpointing则是另一种用时间换空间的方法。正常训练时前向传播的所有中间激活值都会被保存下来用于反向传播这占用了大量显存。梯度检查点的做法是只保存部分层的激活值在反向传播需要用到其他层的激活值时重新做一次前向计算。这样显存占用可以从O(n)降到O(sqrt(n))代价是多了大约30%的计算量。PyTorch提供了torch.utils.checkpoint.checkpoint函数可以很方便地对模型中的某些模块启用检查点。4. 推理与部署层让模型真正跑在生产环境里4.1 推理引擎选型ONNX Runtime、TensorRT还是TorchScript训练好的模型要上线服务第一步就是把它转换成适合推理的格式。PyTorch原生支持TorchScript可以把模型序列化并脱离Python解释器运行适合快速部署。ONNX Runtime则是一个跨框架的推理引擎支持PyTorch、TensorFlow等导出的ONNX模型在CPU和GPU上都有不错的优化。TensorRT是NVIDIA的推理优化引擎针对NVIDIA GPU做了极致的算子融合、精度校准和内存优化通常能带来数倍的性能提升但只支持NVIDIA硬件且转换过程相对复杂。怎么选如果你的服务要跑在多种硬件上ONNX Runtime是更通用的选择如果追求极致性能且只用NVIDIA GPUTensorRT是首选如果只是快速验证或模型结构比较特殊TorchScript最省事。实际生产中很多团队会先用ONNX Runtime做快速上线再针对核心模型用TensorRT做深度优化。4.2 动态批处理与连续批处理提升吞吐的关键推理服务的吞吐量很大程度上取决于批处理策略。静态批处理要求凑够固定数量的请求才处理延迟不可控。动态批处理Dynamic Batching则是在一个时间窗口内尽可能多地收集请求凑成一个batch一起推理窗口结束时即使没凑满也立即处理。Triton Inference Server等推理服务器都支持动态批处理配置。更进一步的优化是连续批处理Continuous Batching也叫迭代级批处理。传统批处理中一个batch里的所有请求必须等最长的那个生成完才能释放导致GPU利用率低。连续批处理则是在每个解码步结束后把已经完成的请求移出batch并立即加入新的等待请求让GPU始终处于满负荷状态。这项技术在大语言模型推理中效果显著vLLM和TensorRT-LLM都实现了连续批处理。4.3 量化与蒸馏模型压缩的两条主流路径量化是把模型参数和激活值从高精度如FP32转换为低精度如INT8、INT4的过程能大幅减少显存占用和计算量。常见的量化方法有训练后量化PTQ和量化感知训练QAT。PTQ不需要重新训练只需要一小批校准数据来统计激活值的分布然后确定量化参数简单快捷但精度损失可能较大。QAT则在训练过程中模拟量化误差让模型适应低精度计算精度保持更好但需要重新训练。蒸馏Knowledge Distillation是另一种思路用一个大模型教师模型来指导一个小模型学生模型的训练让学生模型模仿教师模型的输出分布。这样得到的小模型参数量少、推理快但性能接近大模型。在实际部署中量化和蒸馏经常结合使用先用蒸馏得到一个较小的模型再对这个小模型做量化进一步压缩。5. 平台与工具链让AI系统可观测、可调度、可扩展5.1 资源调度与编排Kubernetes在AI场景下的适配Kubernetes已经成为容器编排的事实标准但原生K8s对AI工作负载的支持并不完美。AI训练任务通常需要GPU、需要多机多卡通信、需要共享存储、需要特定的网络拓扑这些需求催生了一批专门为AI场景优化的调度器比如Volcano、KubeFlow、Slurm等。在K8s上跑AI任务几个关键点需要注意第一GPU资源的调度需要安装NVIDIA Device Plugin它会把节点上的GPU暴露为可调度资源第二多机多卡训练需要Pod之间能通过高速网络通信通常需要配置RDMA和InfiniBand第三训练任务往往需要持久化存储来保存checkpoint和数据集需要配置合适的StorageClass和PVC第四训练任务的生命周期管理排队、优先级、抢占、重试需要额外的调度策略支持。5.2 监控与性能分析别等出问题了才去看日志AI系统的监控和传统Web服务有很大不同。除了常规的CPU、内存、网络指标你还需要关注GPU利用率、显存占用、GPU温度、SM流多处理器占用率、显存带宽利用率、NVLink带宽、PCIe带宽等。NVIDIA提供了DCGMData Center GPU Manager来采集这些指标可以集成到Prometheus和Grafana中做可视化。性能分析方面PyTorch Profiler和NVIDIA Nsight Systems是两个最常用的工具。PyTorch Profiler可以告诉你每个算子的耗时、显存占用、调用栈帮你定位模型中的性能瓶颈。Nsight Systems则更底层可以看到CUDA kernel的执行时间、内存拷贝、CPU-GPU同步等细节。我通常的做法是先用PyTorch Profiler做粗粒度分析找到最耗时的几个算子再用Nsight Systems做细粒度分析看这些算子的kernel实现是否有优化空间。提示torch.profiler.profile的schedule参数可以控制采样的时机和频率避免profiling本身对性能造成太大影响。建议在训练稳定后只对少数几个step做profiling。5.3 存储与数据管理别让I/O拖了后腿AI训练的数据量往往很大从几TB到几PB不等。存储系统的吞吐和延迟直接影响数据加载速度进而影响GPU利用率。常见的存储方案有本地NVMe SSD速度快但容量有限、不共享、NFS共享方便但性能一般、分布式文件系统如Lustre、GPFS性能好但运维复杂、对象存储如S3兼容存储容量大但延迟高。选择存储方案时要考虑你的数据访问模式。如果数据集不大且固定可以预先缓存到本地SSD如果需要多机共享且读取频繁分布式文件系统更合适如果数据主要是顺序读取且对延迟不敏感对象存储也能用。另外数据格式也很关键。小文件太多会严重拖慢读取速度建议把数据打包成WebDataset、TFRecord或Parquet等格式减少文件数量提高读取效率。6. 我踩过的那些坑和给你的一些实在建议回顾我学习和实践AI Infra的过程有几个坑印象特别深刻。第一个坑是盲目追求最新版本。有次为了用某个新特性把CUDA、cuDNN、PyTorch全升到了最新版结果发现某个依赖库还没适配整个训练环境崩了花了两天才回滚。从那以后我养成了用Docker镜像固定环境的习惯每个项目一个镜像版本明确记录再也不怕环境漂移。第二个坑是忽视数据加载的优化。早期训练一个图像模型GPU利用率一直在40%左右我以为是模型太复杂后来用PyTorch Profiler一查发现80%的时间花在数据加载上。把num_workers从2调到8加上pin_memoryTrue和persistent_workersTrueGPU利用率直接拉到85%以上。这件事让我明白优化要从瓶颈入手而不是凭感觉猜。第三个坑是分布式训练的调试。第一次跑多机多卡时任务总是跑几分钟就hang住日志里也看不出明显错误。后来设置了NCCL_DEBUGINFO才发现是其中一台机器的IB网卡没配好NCCL回退到了Socket通信而Socket通信在某些网络配置下会超时。修复网络配置后问题解决。这个经历告诉我分布式系统的排查一定要善用底层库提供的调试信息。如果你刚开始系统学习AI Infra我的建议是不要一上来就啃CUDA编程或分布式理论而是从你手头的项目出发遇到什么问题就深入学什么。比如训练慢就去学性能分析显存不够就去学混合精度和梯度检查点要上线服务就去学推理引擎和量化。带着问题学知识才会真正内化。另外多动手搭环境、多跑benchmark、多看官方文档和源码比看一堆二手教程有用得多。AI Infra是一个实践性极强的领域纸上得来终觉浅绝知此事要躬行。