接手一台 8 卡机器你最先看到的通常不是算力而是一堆刺眼的数字。做 AI Infra 这几年训练与调度领域翻来覆去最常遇见的场景不是资源不够而是资源就在那里睡觉。8 张卡 7 张闲剩下那张也不见得跑得多健康这种状态我见过太多。很多人第一反应是代码写得差、框架选得不对、分布式方案太老结果排查到最后发现根本原因往往特别朴素单机调优没做干净。准备认真搞 AI Infra训练与调度这块你绕不开单机调优它就是你还欠着的第一笔账躲不掉的。这篇文章我把这套算账思路完整写出来。从怎么诊断真正的利用率到 DataLoader 为什么会拖后腿再到 DDP 改造和拓扑通信的细节最后给一张可以直接照着做的体检表。适合混合精度训练刚起步、显卡利用率常年低于 50%、以及准备从单卡升到多卡的工程同学参考。1. 先算账8卡7闲的利用率从哪开始查1.1 别被GPU Utilization骗了SM占用才是真实算力水位拿到一台新机器或新任务很多人第一件事是敲nvidia-smi看到 GPU-Util 80%、90% 就觉得跑得挺满。这个数字的迷惑性非常大。GPU-Util 在驱动层面统计的是采样周期内至少有一个 GPU 引擎处于工作状态的时间比例。注意它说的是有引擎在工作不是算力被高效利用了。举个例子如果你的模型在反复等待数据拷贝、执行很小的 kernel、或者在同步点空转GPU 引擎依然会显示为忙但 SM流式多处理器内部可能大量执行的是等待指令和低效访存根本没有多少有效矩阵乘法和卷积在跑。这就像任务管理器里 CPU 100%但大部分时间都花在忙等锁上真正执行的用户态指令少得可怜。看利用率不能只看一个百分比。我一般用三件套做第一轮健康检查工具命令/方式能看什么nvidia-smi默认输出显存占用、功耗、基础 GPU-Utilnvidia-smi dmonnvidia-smi dmon -s pucctm -d 1SM 活动率、功耗、显存读写、PCIe 流量Nsight Systemsnsys profile每个 kernel 耗时、数据拷贝间隙、空闲段分布重点看 SM 百分比和功耗。A100 满载训练时功耗一般维持在接近 TDP 的区间如果 SM 只有 50% 但显存读写已经拉满那大概率是访存密集型或者数据搬运卡住了。SM 不忙但功耗也很低说明 GPU 根本没使上力气这时候问题八成出在上游数据供给或者同步等待。1.2 三组指标定位瓶颈吞吐、功耗、等待时间愿意细看数的人建议别只看实时曲线先花十分钟建一个 baseline。做法很简单固定一个模型配置关闭日志和可视化跑固定的几十个 step统计真正每秒处理多少样本同时用nvidia-smi dmon记录 SM 占用和功耗。把三条数据放一起看实际吞吐比如某种 CV 模型单卡 A100预期能达到约 800 samples/s实测只有 250。吞吐差了三倍光靠看显存占用根本发现不了。功耗水位GPU 功耗显著低于满载值说明计算单元没被喂饱很可能在等数据或等通信。空闲间隙用 Nsight Systems 或 PyTorch Profiler 看 timeline如果 kernel 之间出现大段空白那基本能断定是同步等待或数据未就位。有个反直觉的事实有时候显存占满但 SM 很低未必是显存不够而是数据管道太慢导致 step 间隔太大旧 batch 一直停在显存里没被消费。处理这种问题后面会专门讲。诊断这一步的核心目标就是别被表面的忙骗了要用吞吐和 SM 占用这两个硬指标把瓶颈定位到具体环节。2. DataLoader 是隐形瓶颈数据搬运比计算更拖后腿2.1 数据管线的饥饿‑过饱模型为什么GPU会等饭很多人觉得 GPU 利用率低就该调模型结构或者换分布式框架但我在实际项目里遇到最多的其实是数据没喂上来。用一个备菜的比喻GPU 是动作极快的厨师CPU 和存储系统是切菜配菜的帮厨。帮厨切得慢厨师再快也只能握着锅铲等。训练循环里的基本情况就是dataloader每次产出一个 batch如果产出速度比 GPU 消费速度慢GPU 就进入饥饿状态。数据管线可以拆成四段读取、解码、变换增强、批处理与传输。每一段都可能成为瓶颈而且往往是多重瓶颈叠加。常见典型问题包括几万个零散小图片文件随机读取每次取样本都要做随机磁盘 IO顺序性极差在 CPU 上做复杂的在线增强比如随机裁剪、旋转、色彩抖动叠加增强计算量比模型一个 forward 还高Python 的 GIL 和跨进程数据交互导致的开销在 worker 个数不合理时被无限放大数据从 CPU 拷贝到 GPU 时用了页缓存内存导致传输速度不稳定。定位数据管线是不是罪魁祸首有个土办法很有效。把训练网络和优化器全部注释掉只保留 dataloader 的迭代和batch.to(device)单独测量它每秒钟能吐出多少 batch。如果单独跑能达到很高的吞吐但一进训练循环就卡说明 GPU 在等菜如果单独跑也很慢那就说明数据管线本身已经需要大改。2.2 五个旋钮一次拧到位从num_workers到数据格式DataLoader 调优并不神秘核心就是把 CPU 端的管线能力补齐到能稳定覆盖 GPU 消费速度。我一般按顺序调下面几个参数DataLoader( dataset, batch_size64, num_workers16, # 通常取物理核数的一半左右需要实测 pin_memoryTrue, # 锁页内存加速 CPU-GPU 拷贝 persistent_workersTrue, # 跨 epoch 复用 worker 进程 prefetch_factor8, drop_lastTrue, )逐个解释一下背后的逻辑num_workers并不是越大越好。worker 太少了数据预处理速度跟不上worker 太多了进程切换和内存占用反而把 CPU 资源吃光拖累主进程。常规经验是先设物理核心数的一半然后往上加看训练吞吐什么时候不再增长甚至回落那个点就是最优值。pin_memoryTrue让 host 端数据放在锁页内存这样to(device, non_blockingTrue)时可以异步传输不会阻塞计算。代价是内存占用会上升小内存机器要控制 worker 总数。persistent_workersTrue避免每个 epoch 都重新创建 worker 子进程。每次重建 worker 的隐含开销在数据集大、epoch 多的任务里非常明显。prefetch_factor默认是 2代表每个 worker 预取多少个 batch。调大到 4 或 8能给训练循环多加一层缓冲对数据抖动有明显的平滑作用。除了这些参数数据存储格式也很关键。如果 dataset 还是几万张小图片单线程逐个读文件哪怕num_workers拉到 32 也救不回来。业界的通用做法是把样本合并成大文件顺序读取。比如用 TFRecord、WebDataset或者提前把所有数据打包成内存映射格式。理由很朴素训练时读的是数据分布不是随机 IO 的噪声。顺序读一个大文件和随机读几千个小文件的性能差距可以到十到几十倍。把上面五个旋钮拧到位之后原先那种GPU 疯狂闪烁但 SM 占用上不去的现象一般会缓解大半。我手头一个 ResNet 训练任务单纯把num_workers从 4 提到 16、打开pin_memory和persistent_workers单卡吞吐就涨了 40%。3. 显存、拓扑与通信单机内同样有网络拥堵3.1 数据并行下的显存墙batch size、梯度累积与重计算数据管线理顺后另一个常见吃利用率的地方是显存和通信。单机 8 卡做数据并行每个 GPU 上会放一份完整的模型副本、优化器状态和梯度缓冲区。对于大模型来说显存墙很快就出现一个 batch 稍微开大一点OOM 直接报错开小了单卡算力又喂不饱。面对显存墙先分清能用什么手段降低单卡 batch size 梯度累积等效扩大 batch但不需要一次占用太大显存。注意梯度累积不是免费的累加的步数越多weight update 的频率就越低收敛节奏会变需要配合学习率调整。激活重计算用torch.utils.checkpoint把中间激活不存backward 时重新算。适合大模型和长序列代价是额外计算量一般会明显提高训练时间所以只在显存确实不足时用。混合精度 AMPA100/H100 上收益显著降低显存占用同时提升 Tensor Core 利用率。这个后面单独讲。梯度累积的写法有一个特别容易踩的坑。很多人直接把每个 mini-batch 的 loss 加起来求平均再 backward这没错但如果配合梯度累积要注意把 loss 除以累积步数或者等价地让每个 mini-batch 对梯度的贡献是平均而非求和。否则等效 batch size 变大之后梯度的 scale 会跟着变大学习率也要莫名其妙调。更进一步在 DDP 场景里默认每次backward()都会触发一次梯度 all-reduce。如果做梯度累积中间那些 mini-batch 的同步其实是白白浪费的。PyTorch 提供了model.no_sync()上下文可以包住中间的梯度累积步骤只在最后一次backward()时同步一次梯度。这在通信开销大的机器上是实打实的提升。3.2 拓扑决定通信成本NVLink还是PCIe一眼看清单机多卡环境里通信成本往往被低估。数据并行的每一步训练结束后都要把所有卡的梯度做一次全局同步也就是 all-reduce。通信量只和模型参数大小有关跟 batch size 无关。模型越大每步训练中通信时间占比越高。但通信成本具体有多高又取决于卡与卡之间的物理连接。很多人不关心机器内部拓扑觉得反正都在一台机器里能慢到哪去。实际上 PCIe 和 NVLink 的带宽差距是三到十倍的量级。排查时我总会跑一下nvidia-smi topo -m这个命令用矩阵方式打印出每对 GPU 之间的连接类型。看到NV#NVLink 链路和PIX/PHB走 PCIe是完全不同的待遇8 张卡通过 NVSwitch 全互联任意两张卡之间都有高带宽低延迟做 all-reduce 很舒服如果是两张一组走 NVLink组间走 PCIe 的拓扑通信就会形成明显的层级差异如果整机走 PCIe switch带宽和延迟都差很多大规模数据并行要特别注意通信与计算的重叠。针对拓扑差异能够做的调优有限但也关键。首先是环境变量和通信库设置比如训练中把NCCL_DEBUGINFO打开观察通信频率和耗时如果 NVLInk 或 PCIe 通信本身异常可以检查NCCL_P2P_DISABLE、NCCL_IB_DISABLE这些开关看是否走了退避路径。注意默认情况下不要乱动这些变量它们更多是排查工具。通信优化的另外一个重点是 batch size 与通信频率的平衡。在数据并行中每步都要 all-reducebatch size 越大通信频率越低等效每个 batch 通信一次计算占比越高。所以如果有显存余量尽量扩大单卡 batch size而不是让梯度累积背锅。比起小 batch 累积 大量同步大 batch 少量同步会舒服很多。4. 最小改造路线从单卡脚本到吃满8卡的DDP4.1 为什么是DDP而不是DataParallel梯度同步方式的代差把 DataLoader 和通信链路理清楚之后再来看代码层面的改造。单卡脚本转多卡很多人下意识用DataParallel因为代码改动最小。但如果你追求利用率我强烈建议直接用DistributedDataParallelDDP不要碰 DP。原因在实现机制上DP 是单进程多线程模型副本分布到多卡后每次 forward 需要从所有卡把输出 gather 到主卡去算 lossbackward 后再把梯度 gather 到主卡更新完再广播回每个卡。主卡既要做计算又要做通信汇聚必然成为瓶颈显存占用也比别的卡高。DDP 是每个进程一张卡每个进程独立 forward/backward梯度通过 ring all-reduce 在所有卡之间对等同步。没有明显的中心节点通信天然更均衡扩展性好得多。从代码改动量上看DDP 并不比 DP 复杂。最小改造只需要四件事初始化进程组、用DistributedSampler切分数据、包装模型、用torchrun启动。我贴一段基本可跑的模板import torch import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.data import DataLoader, DistributedSampler def main(): dist.init_process_group(backendnccl) local_rank dist.get_rank() % torch.cuda.device_count() torch.cuda.set_device(local_rank) dataset MyDataset(...) sampler DistributedSampler(dataset) loader DataLoader(dataset, batch_size64, samplersampler, num_workers8, pin_memoryTrue) model MyModel().to(local_rank) model DDP(model, device_ids[local_rank]) for epoch in range(epochs): sampler.set_epoch(epoch) # 每个 epoch 都打乱数据划分 for x, y in loader: x, y x.to(local_rank), y.to(local_rank) loss model(x, y) loss.backward() optimizer.step() optimizer.zero_grad() if __name__ __main__: main()启动命令是torchrun --nproc_per_node8 train.py有几个非常现实的坑日志、断言、模型保存和打印一定要只在rank 0上执行否则 8 个进程会打 8 份日志看起来混乱不说日志打印本身还会阻塞训练循环。DistributedSampler一定要在每轮 epoch 调用set_epoch(epoch)保证数据划分在每个 epoch 重新打乱。还有一个容易踩的不要随手加dist.barrier()。没有同步需求就不加加了很容易因为某个进程提前异常导致所有进程死锁。4.2 DDP最小改造与AMP、梯度累积的联动单卡能用混合精度多卡当然也能但 AMP 和梯度累积、DDP 放在一起时有个细节特别容易出问题GradScaler的调用时机。我先给一段比较稳的写法scaler torch.cuda.amp.GradScaler() accum_steps 4 optimizer.zero_grad() for step, (x, y) in enumerate(loader): x, y x.to(local_rank), y.to(local_rank) with torch.cuda.amp.autocast(): loss model(x, y) / accum_steps # 梯度累积阶段最后一个 mini-batch 之前不触发同步 is_last ((step 1) % accum_steps 0) context model.no_sync() if not is_last else nullcontext() with context: scaler.scale(loss).backward() if is_last: scaler.step(optimizer) scaler.update() optimizer.zero_grad()解释几个关键点loss model(x, y) / accum_steps这一步已经把累积步数考虑进去了梯度累积不会导致梯度 scale 变大。中间步骤包在model.no_sync()里避免每一次 mini-batch backward 都触发 all-reduce只在最后一个 mini-batch 同步通信量直接降到原来的 1/accum_steps。scaler.step(optimizer)和scaler.update()只在实际更新权重的那个 step 出现不要在每个 mini-batch 都调否则动态 loss scaling 的状态会乱。AMP DDP 改造完成之后有一个必须观察的指标每个 step 的耗时是否稳定。不稳定通常意味着某些 rank 的数据供给不均匀这会直接导致其他卡在同步点等它。可以看 loss 曲线在梯度累积切换之后是否正常收敛。如果 loss 出现剧烈震荡优先把学习率降下来而不是立刻怀疑是混合精度的问题。5. 复盘与延伸一张体检表找出下一处浪费5.1 三次误判案例日志打印、等待锁、显存碎片把一套环境调下来你一定会遇到几个让人怀疑人生的时刻。我把碰过的几个典型误判写出来大家遇到类似情况可以直接对照排查。误判一日志和可视化把训练拖死。有一次训练吞吐怎么调都上不去GPU 利用率时不时掉到 30%。用 profiler 抓了一下发现每隔几步就有一段接近一秒钟的空白。后来才发现训练循环里有个同步上传指标到可视化面板的调用每次上传都会阻塞主进程。你以为在调模型其实在调网络 IO。解决方式很简单指标异步上传或者把上传频率从每步一次改成每 N 步一次。tqdm也是一样在终端里每步刷新和在 notebook 里每步刷新开销完全不同。误判二num_workers 开太大反而更慢。有次为了贪省事把num_workers直接拉到了 32结果吞吐不升反降。因为数据增强里大量用了 OpenCV每个 worker 都有一份独立的 OpenCV 线程池CPU 上下文切换开销被放大主进程也被挤到一边。后来按物理核数一半设并减少了部分在线增强改用预处理缓存问题就消失了。别人给的默认值只能当参考一定要在自己的机器上做 benchmark。误判三显存看起来满就以为是显存不足。初期排查容易把显存占用高等同于显存瓶颈。有一次任务显存占用接近 24GB单卡但 SM 占用只有 20%怎么看都像显存不够。后来发现是数据管道的预取不足GPU 每完成一个 step 都要花很长时间等数据旧数据一直驻留在显存里看起来好像很挤实际上根本没用起来。把 DataLoader 的prefetch_factor调大之后显存里的死数据开始流动SM 占用立刻上来了。这两件事表面上是同一个现象成因完全不同排查路径一个在显存一个在数据供给。误判四小文件 IO 导致数据卡顿却一直在调网络结构。这个问题在 CV 场景特别常见。数据集是几万张 JPEG 小图每轮 epoch 都要随机读一遍。即使加了缓存冷启动阶段也慢得离谱。正解是把数据打包成 WebDataset 或 TFRecord把随机小 IO 变成顺序大 IO。换来换去都绕不开存储介质访问模式这个底层原因。5.2 一张利用率体检表与后续调优路线日常处理利用率问题我习惯把症状和排查手段浓缩成一张表。普通团队遇到训练慢先照着这个表走一轮大多数问题都能定位症状检查手段常见根因优先动作GPU SM 低、功耗低、数据管线耗时占比高单独跑 dataloader 测吞吐DataLoader 参数不合理或预处理过重调 num_workers / pin_memory / prefetch数据打包每个 step 间隔有明显空白PyTorch Profiler 看 timeline日志、可视化、同步等待阻塞移除同步 IOrank0 外不打印8 卡一起跑总吞吐不如单卡×8nvidia-smi topo -m NCCL debug通信占比过高或拓扑层级差异扩大 batch、梯度累积配合 no_sync、检查通信路径显存占用高但 SM 低对比显存占用与 SM 时间线数据饥饿导致旧 batch 堆积优化数据预取确认 pin_memory 生效多卡 loss 震荡或收敛异常观察 loss 与学习率等效 batch 变大后学习率未调整学习率线性缩放或降级回小 batch单机调优的合理预期是什么不是让你把利用率从 20% 干到 99%。训练循环里永远有采样、评估、日志、通信这些必要开销就算全优化到位利用率能稳定在 80%85% 已经非常健康。我做过的几个典型任务从调优前的 30%50% 提到 75%85%端到端训练时间缩短了一半以上这个幅度在工程上已经很可观。单机 8 卡调明白了再往上走多机多卡才有意义。多机训练要面对 RDMA 网络、全局拓扑、通信调度等更复杂的问题但如果单机阶段的数据供给、通信路径和日志同步都没收拾干净上多机只是把同样的账放大好几倍再算一遍。AI Infra 这行很多时候拼的不是会多少花哨框架而是能不能先把眼前这台机器每一张卡都喂饱。先把这笔账还清后面路才走得稳。