显存优化后GPU利用率仍低?LLM训练瓶颈排查与调优指南
1. 显存省下来了GPU 为什么还是吃不满显存优化这件事做过 LLM 训练的人多少都碰过。梯度检查点、ZeRO、LoRA、量化加载一套组合拳下来nvidia-smi里的显存占用确实从爆红变成了绿色心里那块石头刚落地转头一看 GPU 利用率——30% 到 50% 之间来回晃风扇都不怎么转。显存是省下来了可 GPU 还是吃不满训练速度并没有跟着显存下降而线性提升。这个问题在 LLM Training Lab 系列里被单独拎出来讲说明它不是个例。我自己的经验是显存优化和 GPU 利用率优化是两条不同的战线。显存解决的是“能不能跑”利用率解决的是“跑得快不快”。很多人把显存省下来之后就直接开始训练结果发现 step time 没怎么变甚至因为多了些 offload 操作反而更慢了。这篇文章面向的是已经在跑 LLM 训练、做过显存优化、但发现 GPU 利用率上不去的从业者。我会从计算图、数据管线、通信开销、kernel 启动这几个角度拆开讲把“显存省了但 GPU 吃不满”这件事的根因一条条捋清楚并给出可以直接抄的排查方法和调优参数。读完之后你应该能判断自己的瓶颈到底在哪一层而不是盲目加 batch size 或者换卡。2. 先搞清楚 GPU 利用率到底在衡量什么2.1 利用率数字背后的真实含义nvidia-smi里那个 GPU-Util 百分比很多人把它当成“GPU 有多忙”的指标。严格来说它衡量的是在一个采样周期内GPU 上至少有一个 kernel 在执行的时间占比。注意是时间占比不是计算单元的占用率。这意味着什么呢。假设你的模型有一堆小 kernel每个 kernel 执行 10 微秒但 kernel 之间的启动间隔是 90 微秒那 GPU-Util 可能显示 10% 左右。但实际上 SM 里的 CUDA Core 可能在那 10 微秒里是满负荷的。反过来如果有一个大 kernel 连续跑满GPU-Util 显示 100%但里面可能有一半的 warp 在等内存。所以 GPU-Util 低只说明 GPU 有大量时间处于空闲等待状态不说明计算单元没被用满。要判断真实瓶颈得结合sm__throughput、dram__throughput、smsp__issue_active这些指标一起看。Nsight Systems 和 Nsight Compute 是干这个事的后面会讲怎么用。2.2 显存优化和利用率之间的关系显存优化手段大致分几类重计算gradient checkpointing、参数分片ZeRO/FSDP、低精度FP16/BF16/FP8、offloadCPU/NVMe。这些手段对 GPU 利用率的影响方向是不一样的。重计算是用计算换显存前向多算一遍GPU 计算量增加利用率通常会上升但 step time 会变长。参数分片会引入额外的通信如果通信和计算没有 overlap 好GPU 就会在等通信的时候空转。Offload 更直接数据在 CPU 和 GPU 之间搬来搬去PCIe 带宽成为瓶颈GPU 大量时间在等数据。我见过最常见的情况是用了 FSDP 或者 DeepSpeed ZeRO-3显存确实降下来了但 GPU 利用率一直在 40% 左右。原因就是 all-gather 和 reduce-scatter 的通信没有和计算重叠GPU 在等参数从其他卡传过来。显存优化不是免费的。每省下一块显存通常都要用计算、通信或者时间去买单。关键是搞清楚这笔买卖划不划算。3. 数据管线最容易被忽略的 GPU 空转源头3.1 DataLoader 到底卡在哪GPU 吃不满第一个要排查的往往不是模型本身而是数据管线。我做过一个统计在中小规模训练任务里数据管线导致的 GPU 空转占了利用率问题的四成以上。PyTorch 的 DataLoader 默认num_workers0意思是数据加载在主进程里同步做。这时候 GPU 算完一个 batch要等 CPU 把下一个 batch 的数据读出来、做增强、拼成 tensor、再拷贝到 GPU。这一整套下来GPU 就在那儿干等。你可以用torch.profiler或者简单的计时来验证。在训练循环里分别记录数据加载时间和前反向时间import time import torch for step, batch in enumerate(dataloader): t0 time.perf_counter() data batch.to(device, non_blockingTrue) t1 time.perf_counter() outputs model(data) loss criterion(outputs, targets) loss.backward() optimizer.step() optimizer.zero_grad() t2 time.perf_counter() if step % 50 0: print(fdata load: {(t1-t0)*1000:.2f}ms, compute: {(t2-t1)*1000:.2f}ms)如果 data load 时间接近甚至超过 compute 时间那 GPU 利用率低就是数据管线的问题跟显存优化没关系。3.2 把 DataLoader 调顺的几个关键参数num_workers是最直接的。设成 CPU 物理核心数的 2 到 4 倍通常比较合适但也不是越大越好。worker 太多会导致进程切换开销和内存占用上升反而拖慢。我一般从 4 开始试逐步加到 8 或 16看 step time 有没有继续下降。pin_memoryTrue配合non_blockingTrue能让 H2D 拷贝和计算重叠。pin memory 把数据放在锁页内存里DMA 可以直接从这块内存搬到显存不经过 CPU 中转速度快很多。这个组合基本是标配。prefetch_factor控制每个 worker 预取多少个 batch。默认是 2在数据增强比较重的时候可以调到 4 或 6让 worker 提前把数据准备好。persistent_workersTrue在 epoch 之间保持 worker 进程存活避免每个 epoch 重新 fork 的开销。数据量大的时候这个参数能省不少时间。还有一个容易踩的坑如果用了自定义的 Dataset__getitem__里做了很重的 Python 操作比如 PIL 图片解码、numpy 变换那 worker 再多也救不了。这种情况要么把增强逻辑用 DALI 或者 torchvision 的 v2 transform 重写要么预先把数据转成 webdataset 或者 lmdb 格式。3.3 一个真实的数据管线调优案例之前调过一个 7B 模型的训练任务A100 80G 单卡显存优化做到位了但 GPU-Util 只有 35%。用 profiler 一看每个 step 里数据加载占了 60% 的时间。排查下来是 Dataset 的__getitem__里每次都在做 tokenize。文本数据在训练时实时 tokenizePython 的 tokenizer 又是纯 CPU 操作速度很慢。改成离线预处理把所有样本提前 tokenize 好存成二进制文件训练时直接 mmap 读取step time 从 1.8 秒降到了 0.6 秒GPU-Util 直接拉到 85% 以上。这个案例说明一个事显存优化做完之后一定要先确认数据管线不是瓶颈否则后面所有的计算优化都是白费力气。4. 计算图与 kernel 启动开销4.1 小 kernel 太多GPU 在等 launchLLM 训练里有一类 GPU 利用率问题根源在于 kernel 启动开销。CUDA kernel 的启动不是免费的每次 launch 大概有几微秒的固定开销。如果模型里全是小算子比如 element-wise 的 add、mul、激活函数每个 kernel 执行时间可能只有几微秒那 GPU 大部分时间都在等 CPU 发指令。PyTorch 的 eager 模式是逐个算子执行的每个算子对应一个 kernel launch。一个 Transformer layer 里可能有几十个算子一层一层串下来launch 开销累积起来很可观。解决办法有几个方向。一是用torch.compile做算子融合把多个小算子合并成一个 kernel。PyTorch 2.0 之后这个已经比较成熟了对 LLM 训练场景效果明显。二是用 CUDA Graph 把整个计算图捕获下来一次性提交消除重复的 launch 开销。三是手写融合算子比如把 LayerNorm 和 residual add 融合把 SwiGLU 的多个操作融合成一个 kernel。4.2 torch.compile 在 LLM 训练里的实际效果torch.compile用起来很简单model torch.compile(model, modemax-autotune)max-autotune模式会尝试不同的 kernel 实现和融合策略编译时间比较长但运行时的 kernel 数量会大幅减少。我实测过一个 1.3B 的模型eager 模式下每个 step 有 2000 多个 kernel launch用torch.compile之后降到 300 多个GPU-Util 从 55% 提到 78%。不过torch.compile在 LLM 训练里也有坑。动态 shape 会导致反复编译如果 batch 里序列长度变化很大编译缓存会爆掉。解决办法是 padding 到固定长度或者用dynamicTrue让编译器处理动态 shape但性能会打折扣。另外torch.compile和 FSDP、DeepSpeed 这些分布式框架的兼容性需要确认版本。PyTorch 2.1 之后对 FSDP 的支持好了很多但早期版本会有各种奇怪的报错。4.3 CUDA Graph 的适用场景CUDA Graph 适合计算图固定、shape 固定的场景。LLM 训练如果 padding 到固定长度每个 step 的计算图完全一样就适合用 CUDA Graph。PyTorch 里可以用torch.cuda.graphs来捕获g torch.cuda.CUDAGraph() with torch.cuda.graph(g): static_output model(static_input)但 CUDA Graph 和梯度检查点、动态控制流不太兼容。梯度检查点会在反向时重新执行前向计算图不是静态的。所以用了 gradient checkpointing 的模型CUDA Graph 基本用不了。算子融合和 CUDA Graph 解决的是 launch 开销问题但如果你的瓶颈在数据管线或者通信这两个手段帮助有限。先定位瓶颈再选工具。5. 分布式训练中的通信与计算重叠5.1 FSDP 和 ZeRO 的通信开销从哪来FSDP 和 DeepSpeed ZeRO-3 的核心思路是把参数、梯度、优化器状态分片到各张卡上用的时候再 all-gather 回来。这个 all-gather 操作是通信密集型的如果没和计算重叠GPU 就会在等参数的时候空转。具体来说FSDP 在前向时对每一层做 all-gather拿到完整参数后计算算完就把非本地的参数丢掉。反向时再 all-gather 一次然后 reduce-scatter 梯度。如果 all-gather 是同步的GPU 在通信期间完全空闲。DeepSpeed ZeRO-3 也类似但它有stage3_prefetch_bucket_size和stage3_param_persistence_threshold这些参数可以调让预取提前发生把通信藏在计算后面。5.2 让通信和计算重叠的实操配置FSDP 里有个backward_prefetch参数可以设成BACKWARD_PRE或者BACKWARD_POST。BACKWARD_PRE会在反向计算开始前就预取下一层的参数重叠效果更好但显存占用会高一些。from torch.distributed.fsdp import FullyShardedDataParallel as FSDP from torch.distributed.fsdp import BackwardPrefetch model FSDP( model, backward_prefetchBackwardPrefetch.BACKWARD_PRE, forward_prefetchTrue, limit_all_gathersTrue, )forward_prefetchTrue让前向的 all-gather 也提前发起。limit_all_gathersTrue限制同时进行的 all-gather 数量避免通信队列堆积。DeepSpeed 这边关键是调stage3_prefetch_bucket_size。这个值设得太小预取不够提前通信暴露出来设得太大显存占用上升。一般从5e7开始试根据模型层大小调整。还有一个通用手段是用torch.distributed的ProcessGroupNCCL配置开启NCCL_ASYNC_ERROR_HANDLING和调整NCCL_BUFFSIZE。NCCL 的 ring 算法在小消息时延迟低tree 算法在大消息时带宽高具体选哪个要看 all-gather 的消息大小。5.3 通信瓶颈的排查方法判断是不是通信瓶颈可以看 Nsight Systems 的 timeline。如果看到 NCCL kernel 和计算 kernel 之间有明显的空隙GPU 在 NCCL 期间完全空闲那就是通信没重叠好。另一个简单方法是算通信量。以 FSDP 为例每个 step 的通信量大约是2 * 参数量 * (1 1/world_size)字节all-gather 加 reduce-scatter。如果这个量除以卡间带宽得到的时间和 step time 里计算部分的时间接近那通信就是瓶颈。比如 7B 模型BF16 参数world_size8通信量约2 * 7e9 * 2 * (1 1/8) ≈ 31.5 GB。如果卡间带宽是 200 GB/sNVLink通信时间约 157ms。如果 step time 是 400ms那通信占了接近 40%必须做重叠优化。6. 显存优化手段本身带来的副作用6.1 梯度检查点的计算代价梯度检查点activation checkpointing是省显存最常用的手段但它会让前向计算多做一遍。具体来说正常反向需要保存所有中间激活用了 checkpointing 之后只保存部分层的激活反向时重新计算这些层的中间结果。这个重计算的开销取决于 checkpoint 的粒度。如果每层都 checkpoint前向计算量增加约 33%因为反向本来就要算一遍前向的导数现在多算一遍前向。如果每隔几层 checkpoint开销小一些但省显存的效果也差一些。PyTorch 里可以用torch.utils.checkpoint.checkpoint包住需要 checkpoint 的模块。FSDP 也有activation_checkpointing参数可以按层应用。from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import ( checkpoint_wrapper, CheckpointImpl, apply_activation_checkpointing, ) checkpoint_wrapper_fn lambda m: checkpoint_wrapper( m, checkpoint_implCheckpointImpl.NO_REENTRANT ) apply_activation_checkpointing(model, checkpoint_wrapper_fncheckpoint_wrapper_fn)NO_REENTRANT模式比默认的 reentrant 模式效率高因为它不会在反向时重新触发 autograd 的 hook减少了开销。PyTorch 2.1 之后推荐用这个。6.2 Offload 的带宽陷阱CPU offload 把优化器状态或者参数放到 CPU 内存里用的时候再搬到 GPU。这个操作走 PCIe带宽比显存带宽低一两个数量级。PCIe 4.0 x16 的理论带宽是 32 GB/s实际能到 25 GB/s 左右而 A100 的显存带宽是 2 TB/s差了将近 80 倍。如果 offload 的数据量大GPU 就会大量时间在等 H2D 拷贝。DeepSpeed 的 ZeRO-Offload 和 ZeRO-Infinity 都有这个问题。用的时候要算清楚每个 step 需要 offload 多少数据除以 PCIe 带宽看这个时间占 step time 的比例。我一般建议除非显存实在不够否则优先用参数分片而不是 offload。分片走的是卡间高速互联NVLink 或 InfiniBand带宽比 PCIe 高得多。6.3 低精度训练的精度与速度权衡FP16 和 BF16 能省一半显存同时提升计算吞吐Tensor Core 在低精度下算力更高。但 FP16 有数值范围小的问题容易溢出需要 loss scaling。BF16 范围大基本不用 scaling但精度低一些。FP8 是更新的选项H100 之后的卡支持。FP8 能把显存再降一半计算吞吐再翻倍但对模型精度的影响需要仔细评估。目前 LLM 训练里 FP8 主要用于前向反向还是用 BF16混合精度策略比较复杂。低精度本身不会导致 GPU 利用率下降反而会提升。但如果因为精度问题频繁出现 loss spike需要回滚 step 或者调整 scaling factor那训练效率就会受影响。7. 用 Nsight 工具定位真实瓶颈7.1 Nsight Systems 看 timelineNsight Systems 是系统级的 profiler能看到 CPU、GPU、通信、IO 的完整时间线。用法很简单nsys profile -o profile_output --force-overwrite true python train.py跑完之后用nsys-ui打开.nsys-rep文件看 GPU 那一行。如果 GPU 上有大段空白说明 GPU 在等什么。把鼠标放到空白处看对应的 CPU 在干什么或者 NCCL 是不是在跑。常见的模式有几种。GPU 空白对应 CPU 在 DataLoader 里那是数据瓶颈。GPU 空白对应 NCCL kernel那是通信瓶颈。GPU 上 kernel 很碎间隔很大那是 launch 开销问题。7.2 Nsight Compute 看 kernel 内部Nsight Compute 是 kernel 级的 profiler能看单个 kernel 的 SM 占用率、内存吞吐、warp stall 原因。用法ncu --set full -o kernel_profile python train.py重点看几个指标。sm__throughput.avg.pct_of_peak_sustained_elapsed是 SM 计算单元的利用率如果这个值低说明计算单元没吃满。dram__throughput.avg.pct_of_peak_sustained_elapsed是显存带宽利用率如果这个高但 SM 低说明是内存瓶颈。smsp__issue_active.avg.pct_of_peak_sustained_elapsed是指令发射率低说明 warp 调度有问题。对于 LLM 训练常见的 kernel 瓶颈是 attention 的 softmax 和 LayerNorm这些操作有大量的 reduction 和 element-wise 计算内存访问模式不够优化。FlashAttention 就是专门解决这个问题的把 attention 的多个操作融合成一个 kernel减少显存读写。7.3 PyTorch Profiler 的轻量替代如果不想装 NsightPyTorch 自带的 profiler 也能用from torch.profiler import profile, ProfilerActivity with profile( activities[ProfilerActivity.CPU, ProfilerActivity.CUDA], scheduletorch.profiler.schedule(wait1, warmup1, active3), on_trace_readytorch.profiler.tensorboard_trace_handler(./log), ) as prof: for step, batch in enumerate(dataloader): train_step(batch) prof.step()跑完之后用 TensorBoard 看 trace能看到每个算子的 CPU 和 GPU 时间。虽然粒度不如 Nsight 细但胜在方便适合快速定位。8. 常见问题速查与避坑经验8.1 GPU 利用率问题排查表现象可能原因排查方法解决方向GPU-Util 低GPU 上有大段空白数据管线瓶颈profiler 看 data load 时间增加 num_workerspin_memory预处理数据GPU 空白对应 NCCL通信未重叠Nsight 看 NCCL 和计算的时间关系调 prefetch开 forward_prefetchkernel 碎间隔大launch 开销看 kernel 数量和间隔torch.compileCUDA Graph算子融合SM 利用率低但 DRAM 高内存瓶颈Nsight Compute 看 dram throughput用 FlashAttention优化内存访问step time 波动大动态 shape 或数据不均衡看每个 step 的时间分布padding 到固定长度数据 shuffle显存够但 batch 上不去激活值占用大看显存分配gradient checkpointing序列并行8.2 几个我踩过的坑第一个坑是盲目加 batch size。显存省下来之后第一反应是把 batch size 加大觉得这样 GPU 就能吃满。但实际上如果瓶颈在数据管线加大 batch size 只会让数据加载更慢GPU 利用率反而下降。正确的顺序是先确认数据管线不是瓶颈再考虑加 batch。第二个坑是num_workers设太大。有次设了 32 个 worker结果 CPU 内存爆了而且进程切换开销让数据加载更慢。后来发现 8 个 worker 就够了再多是浪费。第三个坑是 FSDP 的limit_all_gathers没开。默认情况下 FSDP 会同时发起多个 all-gather通信队列堆积反而拖慢。开了limit_all_gathersTrue之后all-gather 按顺序进行和计算的重叠更好。第四个坑是torch.compile和 gradient checkpointing 一起用。早期版本这两个有冲突编译后的图在 checkpoint 重计算时会出错。PyTorch 2.2 之后好了很多但还是建议先分别验证再组合使用。8.3 一个实用的调优顺序显存优化做完之后我一般按这个顺序调 GPU 利用率先用 profiler 确认数据管线不是瓶颈。data load 时间应该小于 compute 时间的 10%。看 Nsight timeline确认 GPU 空白对应的是什么。是数据、通信还是 launch。如果是 launch 开销上torch.compile看 kernel 数量下降多少。如果是通信调 FSDP/DeepSpeed 的 prefetch 参数开 forward_prefetch。如果是内存瓶颈换 FlashAttention优化 LayerNorm 和激活函数。最后再考虑加 batch size 或者换并行策略。这个顺序的核心逻辑是先解决“GPU 在等什么”再解决“GPU 算得慢不慢”。前者是利用率问题后者是吞吐问题两个不要混在一起调。调优最忌讳的是同时改多个参数。改一个测一次记录 step time 和 GPU-Util确认有效再改下一个。否则出了问题都不知道是哪个参数导致的。9. 一些关于硬件和环境的补充热词里提到了 RTX 4060 Laptop、6G 显存、RTX 3060 12G 这些消费级卡跑 LLM 的场景。这些卡跑推理还行跑训练基本不现实显存和算力都不够。如果非要在消费级卡上做实验建议用 LoRA 或者 QLoRA只训练少量参数显存需求能降到 6G 到 8G。但即使是 LoRA 微调GPU 利用率也容易上不去。消费级卡的 PCIe 通道数少数据加载和 offload 的带宽瓶颈更明显。这时候数据管线的优化比计算优化更重要因为卡本身算力就有限数据供不上就是纯浪费。另外WSL 环境下跑 PyTorch 训练CUDA 的版本兼容性要特别注意。WSL 里的 CUDA 驱动是 Windows 驱动透传的版本必须和 Windows 侧的驱动匹配。PyTorch 的 CUDA 版本、cuDNN 版本、驱动版本三者要对上否则会出现各种奇怪的报错。我一般建议在 WSL 里用 conda 管理环境PyTorch 用官方 conda 源安装能省很多版本兼容的麻烦。最后说一个观察GPU 利用率这个问题本质上是个系统工程问题。它涉及数据、计算、通信、存储多个环节任何一个环节拖后腿GPU 就吃不饱。显存优化只是让你能跑起来真正要让 GPU 跑满得把整个 pipeline 都调顺。这也是为什么同样一张卡有人能跑到 90% 利用率有人只能跑到 40%差距不在硬件在对整个训练流程的理解和调优。