这几块牌子在机房里躺了大半年老板问起训练进度你打开nvidia-smi8 张卡里 7 张 util 是 0%只有一张在 90% 上下跳。Loss 还在掉模型也没报错但每个 step 慢得像老牛拉破车。干 AI Infra 这几年这种情况见得太多了大家张口就是“集群调度”“弹性扩缩容”“分布式训练框架”可真到了现场绝大多数项目的第一道坎根本不是多机多卡而是单机这关压根没过。我始终觉得训练和调度这个主题从上往下讲容易把人带沟里去。正确的姿势是先从单机调优开始把账算清楚:8 张卡为什么 7 张闲?数据卡在哪个环节?通信有没有在等人?这些问题不解决给你 1000 张卡也是白搭。这篇文章是 AI Infra 系列第五章的上半部分专门聊单机训练环境里的性能调优。我尽量不堆概念全用能直接落地的东西说话:从时间构成拆解、DataLoader 参数调节、显存测算、混合精度、到单机与多机调度之间的边界一条线捋下来。适合正在做模型训练、经常被“GPU 利用率低”“训练慢”困扰的算法工程师和 AI Infra 从业者参考。1. 先认清账本:8 卡 7 闲到底闲在哪1.1 训练一个 step 的时间到底花在哪了排查单机性能问题我习惯先建立一个最基本的时间模型。拿 PyTorch 训练循环举例一个 step 大致可以拆成几段:数据加载:CPU 读磁盘或网络存储做解码、裁剪、归一化、随机增强然后从内存搬到显存;前向计算:GPU 执行若干算子产出 loss;反向传播:GPU 计算梯度这一步和前向一样吃算力;参数更新:优化器 step涉及 CPU/GPU 之间同步也会产生等待;通信同步:DDP 模式下 all-reduce 梯度单机 8 卡也有这一步只是走 NVLink 或 PCIe 比跨机快。很多人的直观理解是:一个训练 step 的主要开销在 GPU 计算。所以 GPU util 高就觉得自己“跑满了”。实际上数据加载、CPU 预处理、内存拷贝、同步锁竞争这些环节非常容易把 GPU 拖到“饿肚子”状态。GPU util 低不代表卡不行而往往是上游供给跟不上。判断标准其实很简单:用一个空的、不做任何数据加载的模型跑同样 batch size 的前向和反向对比这个理論耗时和真实训练的 step 耗时。如果真实耗时是理论耗时的好几倍那么问题大概率不在 GPU 算力上而在数据管线和同步环节。这套方法我给自己起名“空转对比法”排查效率很高。1.2 7 张卡全在等:显存分配与通信等待的连锁反应8 卡只跑 1 张还有一种隐藏原因:显存。很多人开了 DDP但没有意识到每张卡的 batch size、输入尺寸会直接影响显存占用。假设代码里写死了总 batch_size 32然后平均分到 8 卡每卡只有 4 张图模型大一点可能一卡 4 张直接 OOM。如果代码里 fallback 到梯度累积或者自动缩小 batch那就可能出现某些 rank 的 worker 因为数据不均衡早早算完开始空转等待。再一个常见坑是 NCCL 通信等待。DDP 在 backward 之后要做梯度 all-reduce通信量等于模型参数大小乘以梯度字节数。如果模型里有大量小 tensor通信次数会非常多。这时候 8 张卡的计算轮番结束但因为每次同步都受最慢的一张拖累其他卡只能干等。从 util 上看就是一张卡偶尔跳另外几张全闲着。更麻烦的是这种问题在单卡上完全看不出来因为单卡根本没有跨卡同步。这也是为什么我坚持让团队先把单卡性能和单机 8 卡性能梳理清楚再谈多级扩展。2. 单机调优第一步:数据管线是最大的隐形坑2.1 不跑 profile 你根本不知道 DataLoader 有多慢AI Infra 这行流行一句话:“GPU 算不过来是假瓶颈数据喂不上来才是真瓶颈。”我接手过的项目里至少有三分之一所谓“训练慢”的根因最后都落在数据加载和预处理上。最典型的情况是:硬盘是普通 SATA SSD数据集是几万张高分辨率 JPEG每次 epoch 都要重新解码。你想象一下磁盘读了图片CPU 还要做 JPEG 解码、随机裁剪、缩放、归一化、转 tensor一套操作下来几十毫秒。GPU 算一个 step 可能只要 20 毫秒但数据准备要 80 毫秒那 GPU 有一大半时间就是在等。这时候nvidia-smi里恰恰能看到 GPU util 不算低偶尔冲到 50%但 step 耗时就是压不下去。怎么确认?我通常做一次“数据供给压测”:在训练循环里去掉计算只跑 DataLoader 的迭代统计每秒能出多少个 batch。如果这个值大于训练 step 的吞吐需求说明数据供给暂时够;如果低于需求对不起数据管线就是第一嫌疑人。2.2 DataLoader 参数调节的实操经验PyTorch 的 DataLoader 参数看着就那几个但调不好差距巨大。以下是我经手多个实际训练项目后沉淀下来的经验值未必适合所有场景但可以当作起点来用:参数建议值说明num_workersCPU 核心数的一半到四分之三太少CPU 来不及处理;太多进程切换和内存拷贝开销反而增大prefetch_factor2 到 4每个 worker 预取 batch 数能掩盖单次处理抖动pin_memoryTrue把数据锁在页锁定内存加快 CPU 到 GPU 的拷贝persistent_workersTrue避免每个 epoch 重建 worker 进程省去启动开销有一点要特别说:num_workers 不是越大越好。我见过有人直接在 64 核机器上设 64结果内存暴涨、CPU 上下文切换疯狂系统甚至把训练进程给 OOM kill 掉。比较合理的思路是先用psutil.cpu_count()拿到物理核心数再按物理核数的一半起步逐个往上加观察吞吐变化。通常在某个值之后吞吐就不再提升甚至下降那个值就是当前环境的上限。另外自定义 Dataset 的__getitem__里千万别做重活。比如每次都从磁盘读一个超大 JSON 然后解析这个操作放到 worker 里会让整条管线卡死。最理想的做法是:在__init__阶段就把元数据全部加载到内存__getitem__只做索引和轻量变换。实在有重活也应该用缓存机制把它挡在训练循环之外。2.3 存储选型与预取策略:为什么 SSD 也救不了你很多人觉得把机械盘换成 NVMe SSD 就万事大吉。实际没那么简单小文件随机读取在 SSD 上确实快不少但如果数据是几百万张图的小文件即便 NVMe 也会因为 inode 开销、目录遍历等因素产生大量延迟。我在一个 CV 项目里试过几种方案数据量大概 2TB 的图片训练机器是 8 卡 A100。直接读原始图片step 时间在 500ms 左右;如果把所有图片打包成 TFRecord 或 WebDataset 格式用顺序读的方式step 能压到 180ms。差别就在“随机小文件 IO”和“顺序大文件 IO”上前者对存储系统的随机读能力极其严苛后者吃的是纯顺序带宽。另外一个容易忽略的点是 Page Cache。Linux 会把读过的文件缓存在内存里第二次 epoch 去读同一批数据就会快很多。所以如果你发现第一个 epoch 很慢、第二个 epoch 突然变快不代表环境变好了只是 Page Cache 命中了。这种情况建议优先考虑把整个数据集放内存(比如 /dev/shm 或用系统内存缓存),不过这又涉及内存容量实际操作前要掂量一下。3. GPU 侧资源管理:显存、精度与算子3.1 如何精准测算显存占用:先知道你差多少聊完数据再来看看显存。很多时候训练没挂但性能上不去是因为显存接近上限导致 PyTorch 的缓存分配器反复腾挪。之前 OOM 一次就够折腾现在不 OOM但每步都在内存碎片和显存寻址上浪费时间。建议在训练脚本里加上几行监控代码:import torch # 在训练循环结束时打印最大显存占用 print(fGPU 0 max memory allocated: {torch.cuda.max_memory_allocated(0) / 1024**3:.2f} GB) print(fGPU 0 max memory reserved: {torch.cuda.max_memory_reserved(0) / 1024**3:.2f} GB)allocated是实际张量占用的显存reserved是 PyTorch 向驱动预占的缓存。reserved 通常比 allocated 大不少这是正常的因为 PyTorch 会保留显存缓存避免反复申请。但如果 reserved 长期接近显存上限你就该考虑手动调小 batch size或者用torch.cuda.empty_cache()在合适的时机做一次释放。一个可复用的经验公式:先用一个小 batch size 跑一个 step记录max_memory_allocated然后用“显存上限 - 模型参数和优化器状态占用”的余量去估算能撑起多大的 batch。这一步虽然粗糙但比盲目调参强太多。通常优化器状态(Adam 的动量和方差)占用的显存是模型参数的好几倍,这点在算余量时务必算进去。3.2 混合精度训练的开启与 loss scaling混合精度训练是单机调优里性价比最高的手段之一一块支持 Tensor Core 的卡开启 FP16 后训练速度能直接翻倍或更多。但这东西不是简简单单把模型参数转成half()就完事背后还有一套工程细节。PyTorch 里最常见的做法是用torch.cuda.amp:scaler torch.cuda.amp.GradScaler() for data, target in loader: optimizer.zero_grad() with torch.cuda.amp.autocast(): output model(data) loss criterion(output, target) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()需要注意的坑主要在 loss scaling 上。FP16 的表示范围窄梯度值太小会直接下溢成 0所以要用GradScaler放大 loss反向传播后再缩小回正常数值范围。如果出现 loss 变成 NaN 或 inf很多人第一反应是学习率太大但别忘了检查是不是 scale 因子没更新好或者模型里某些算子在 FP16 下数值不稳定。一个简单验证办法:先用 FP32 跑一个 step对比 FP16 下的梯度统计定位数值出问题的算子。torch.cuda.amp.autocast默认会把大多数卷积和矩阵乘算子自动切到 FP16但某些算子(比如涉及大范围归约的)仍会留在 FP32。所以混合精度的收益也取决于模型结构。CNN、Transformer 这类矩阵密集型的模型收益最大而含大量自定义算子的模型可能收益有限。如果开启后速度提升不明显可以用 PyTorch Profiler 看看哪些算子还在用 FP32。3.3 梯度累积、activation checkpointing 与显存置换显存不够的时候梯度累积是一个常用手段。梯度累积的意思是:不更新参数而是连续多个 step 的梯度累加在一起再统一更新。它模拟了大 batch size 的效果但每个 step 的 batch size 可以很小显存压力就小。直接改两个地方:设置累积步数accumulation_steps然后只在step % accumulation_steps 0时optimizer.step()。一个容易踩的坑是 BatchNorm。如果模型里有 BN 层梯度累积时每个小 batch 的均值和方差统计会有偏差极端情况下模型不稳定。解决途径要么换 GroupNorm要么调整累积步数和学习率的比例。我一般会在实验记录里写明:累积 N 步时等效 batch size N × 单步 batch size学习率也要相应调整否则收敛曲线会出问题。activation checkpointing(也叫梯度检查点)是另一种省显存方式。它不保存所有中间激活值而是只存部分反向传播时再重新计算。这个“重计算”会带来额外耗时但显存节省非常可观。对于 Vision Transformer 或大语言模型这类结构开 activation checkpointing 后可以把 batch size 提高到原来的两倍以上。实际使用时建议只对有大量激活值的模块开启而不是整个网络无脑开尽量把重计算开销控制在可接受范围内。4. 单机与多机的边界:为什么调度问题先压在单机上4.1 单机 8 卡跑不满上多机只会更糟很多人提到调度下意识会想 Kubernetes、Slurm、集群排队这些偏平台层的东西。但真正的调度第一步发生在单机内部:一个 step 里的 CPU 线程、数据 worker、GPU 计算、显存拷贝、NCCL 通信这些资源的协调和排队本身就是调度的雏形。你单机 8 卡都跑不满,比如只有一张卡在算、其余在等通信,那把这些卡分布在多台机器上只会放大等待——跨机通信的延迟和带宽都远不如单机 NVLink/PCIe。我见过一个团队模型训练慢第一反应是上多机分布式,结果拓扑图都画好了后来一查发现单机内部 DataLoader 的瓶颈就占了 60% 的耗时。这种属于典型的本末倒置。正确的检查顺序应该是:单卡性能是不是已经最优?单机 8 卡 DDP 的扩展效率有没有跑到线性?如果单机都跑不动别急着碰多机。4.2 定位瓶颈的工具集与实操流程聊几个我日常排查单机训练问题最常用的工具和数据nvidia-smi:看 GPU util、显存、温度、功耗。注意 util 是采样时间片里的活跃程度不是精确的算力占用所以 util 高不代表算子效率高;nvtop:类似系统top,但看 GPU 状态,实时性好;py-spy:类似 Python 版perf可以 dump 训练进程的调用栈,定位卡在哪个 Python 函数;nsys profile/nsys-ui:NVIDIA 的性能分析工具,能看到 kernel 计算时间、数据传输时间、空闲时间,信息量最大。我定位瓶颈的基本流程是:先跑 50 个 step用nvidia-smi dmon观察 GPU util、显存和 PCIe 读写量。如果 util 低但 PCIe 读写很高说明数据拷贝是瓶颈;再用py-spy dump查看进程卡在哪儿。如果很多 worker 都停在recv或join说明 worker 同步和队列等待有问题;最后用nsys profile --tracecuda,nvtx,osrt跑一个 step看 kernel 间隙。如果 kernel 之间的空隙远大于 kernel 本身耗时那一定是数据准备或者 CPU 同步拖了后腿。这套组合拳基本能在半小时内把单机训练瓶颈定位到具体环节。工具不需要全上从nvidia-smi和py-spy开始就够了,实在不够再用nsys。5. 常见问题速查表与避坑清单5.1 训练慢、卡死、OOM 的高频故障对照下面是我这几年来回踩过、也在团队里给新人培训时反复强调的几张排查表按“症状—原因—解法”来列方便直接抄作业。症状常见原因排查方向与解法GPU util 低CPU 占用高DataLoader 预处理太重用py-spy看 worker 栈;把重预处理改为缓存或离线做;调大num_workers第 1 个 epoch 极慢后面变快Page Cache 命中数据缓存到内存考虑将数据集整体放入内存;用 WebDataset/TFRecord 顺序读某些 step 偶发特别慢CPU 被系统进程抢占或磁盘抖动检查 cgroup、strace;稳定环境优先用物理机或专用节点8 卡 DDP 时 util 都低梯度通信频繁小 tensor 过多尝试梯度累积/梯度压缩;检查模型里是否有大量小参数显存 reserved 很高但不 OOMPyTorch 显存缓存策略用empty_cache();降低 batch size;启用 activation checkpointing开启 AMP 后 loss 变成 NaNFP16 下梯度下溢或算子不稳定检查 GradScaler;对敏感算子强制 FP32;调小学习率GPU util 看似 100%但 step 时间不降GPU 被部分低效算子占满用nsys看 kernel 耗时;尝试算子融合、torch.compile这张表不可能覆盖所有情况,但单机训练类问题里,十有七八最后都会落到这几类根因上。5.2 复现性、日志与坏数据集的处理再分享两个跟定位瓶颈不直接相关、但能让整个调试过程舒服很多的经验。第一是复现性。训练性能排查里改动一个参数就要对比前后差异但如果你每次都从头训练或者随机数种子不固定那对比结果根本不可信。建议在训练脚本启动时就固定random.seed、numpy.random.seed和torch.manual_seed必要时也设置torch.backends.cudnn.deterministic True。虽然这会略微降低性能但调试阶段的高可复现比那点算力损失更值。第二是坏样本处理。数据管线卡顿并不总是系统瓶颈,有时候是数据集里混了损坏图片或异常长文本。比如一张几 MB 的超大图每次随机采样到它当个 step 就可能比别的 step 慢 10 倍。我在真实项目里就被这种“慢样本”坑过损失曲线稳定下降但每过一段时间跳出一个尖峰。后来在 Dataset 里加了“样本耗时统计”把所有 5 秒的单样本打印出来才定位到几张异常大图。把这类样本预处理成统一尺寸后训练稳定很多。这里也顺手放一个小建议:训练脚本里一定要把关键信息(step 耗时、显存占用、GPU util、数据加载耗时)自动记录到日志或 wandb 上别只靠人眼盯nvidia-smi。没有历史数据你后面回查“这周训练为什么变慢”时会无从下手。6. 从单机调优说起给调度留个口子回到标题那句话“8 张卡 7 张闲?单机调优才是你欠下的第一笔账。”这句话不是我危言耸听。太多项目把大量的算力预算砸在集群平台、调度平台、多机框架上结果底下的训练任务连单机都跑不顺。调度从来不是空对空的平台概念它落到最底层就是每一块 GPU 的 util、每一条数据链路是不是都被喂饱了。单机调优做扎实之后再去看多机就会发现思路完全不一样。你不会一上来就想“我要不要用 Ray 或者 Kubernetes 加个弹性伸缩”而是会先算一笔账:单机扩展效率是多少?通信占比是多少?在哪个规模开始收益递减?这些问题的起点全是这一章里讲的单机数据。我自己带过不少人发现一个规律:能把单机 8 卡利用率和扩展系数讲明白、测明白的人后面做多级集群调度往往也思路清晰。反过来一上来就撸平台工具的人经常在资源需求和任务特性这块翻车因为对训练任务本身没有体感。如果你现在也被“8 卡 7 闲”这种诡异状态折磨建议先把这篇文章提到的东西逐一排查一遍。把 DataLoader 调好把显存占用摸清楚把 AMP 开起来跑出基线再回头看多机调度的下一步。你的第一笔账还清了后面的事情自然顺了。