显存省了GPU却吃不满？LLM训练中显存优化与吞吐提升的冲突与解法

1. 显存省下来了GPU 为什么还是吃不满显存占用从 22GB 压到 9GBnvidia-smi里显存那一栏确实绿油油一片很健康结果 GPU 利用率死活卡在 30% 上下风扇都不怎么转训练一个 epoch 的时间跟没优化前一模一样。这个场景我猜不少人都遇到过尤其是刚上手 LLM 训练、兴冲冲搞完梯度检查点、LoRA、8-bit 优化器之后满心以为能起飞结果发现显存是省了速度一点没变。这篇就专门聊这个坑。核心关键词是LLM、GPU、显存、PyTorch、CUDA但真正要解决的不是怎么省显存而是显存省下来之后为什么算力没被用起来。适合已经能跑通训练脚本、做过一轮显存优化、但发现吞吐没提升的中阶玩家。如果你还在纠结torch.cuda.is_available()返回 False那这篇可能稍微超前了一点但看看也无妨迟早会遇到。先把结论摆前面显存和算力是两条独立的流水线省显存的操作很多时候是在给算力流水线添堵。梯度检查点用时间换空间8-bit 优化器增加了量化/反量化开销LoRA 让每个 step 的计算量变小但 kernel 启动次数没变这些都会让 GPU 的 SM流多处理器闲着。GPU 吃不满本质上是喂给它的活不够连续、不够大块而不是显存不够它施展。下面我按自己踩坑的顺序把这件事拆成四块讲先讲清楚显存和利用率为什么是两回事再讲几个最常见的省显存反而拖慢的操作然后是真正能把 GPU 喂饱的实操手段最后是我自己整理的一份排查清单。2. 显存和 GPU 利用率压根不是一回事2.1 显存是仓库利用率是流水线开工率打个比方。显存相当于工厂的仓库GPU 的计算核心SM相当于流水线工人。你把仓库从 22 平米缩到 9 平米确实能放下更多货了但工人干活的速度取决于传送带上有没有连续的活送过来跟仓库多大没直接关系。仓库小一点工人照样能满负荷仓库再大传送带断断续续工人照样摸鱼。nvidia-smi里的 GPU-Util 这一栏官方定义是过去采样周期内至少有一个 kernel 在执行的 SM 占比。注意关键词至少有一个 kernel 在执行。这意味着哪怕一个 SM 上只跑了一个线程这一瞬间它也算忙。所以 GPU-Util 高不代表算力用满了它只代表没完全闲着。真正反映算力利用的是 SM occupancy、Tensor Core 活跃度、显存带宽利用率这些更细的指标。我见过太多人盯着 GPU-Util 90% 就以为万事大吉结果torch.profiler一开发现大部分时间花在 kernel launch 的间隙和等待数据上真正的矩阵乘只占了一小部分。所以第一步认知要摆正GPU-Util 是个粗指标别拿它当唯一真理。2.2 一个 step 里GPU 到底在等什么一个训练 step 拆开看GPU 的时间大致花在这几块数据搬运从 CPU 内存把 batch 拷到显存H2D如果num_workers不够或者数据预处理太重GPU 就在这干等。前向计算一堆 kernel 排队执行矩阵乘、激活、归一化。反向计算通常是前向的两倍计算量。优化器更新Adam 这类优化器要读写大量状态显存带宽压力大。kernel launch 间隙每个 kernel 启动都有几微秒的固定开销kernel 越小越多间隙占比越高。省显存的操作往往是在放大 kernel launch 间隙或者增加额外的搬运/量化开销。比如梯度检查点前向的时候不存中间激活反向的时候重新算一遍等于把前向算了两遍计算量上去了但显存下来了。如果你的瓶颈本来就在计算这一下就更慢了。2.3 显存优化和吞吐优化目标函数不一样这里有个很关键的认知显存优化和吞吐优化是两个不同的目标函数甚至经常互相冲突。优化手段显存效果吞吐效果原因梯度检查点大幅降低通常下降前向重算计算量增加8-bit 优化器降低可能下降量化/反量化开销LoRA大幅降低小 batch 下下降计算量小kernel 启动占比高增大 batch上升上升计算密度提高混合精度降低上升Tensor Core 加速FlashAttention降低上升减少显存读写看这张表就明白了真正既省显存又提吞吐的只有混合精度和 FlashAttention 这类算法级优化其余大部分是拿吞吐换显存。所以当你做完一轮显存优化发现 GPU 吃不满先别急着怀疑代码先想想你用的手段是不是本来就牺牲了吞吐。3. 几个省了显存却拖慢 GPU的典型操作3.1 梯度检查点用两倍前向换一半激活梯度检查点gradient checkpointing的原理是前向的时候只存少数几个检查点的激活反向传播需要中间激活时从最近的检查点重新前向算一遍。显存能降 50% 到 70%代价是计算量增加约 30% 到 50%。问题在于很多人开了检查点之后batch size 没跟着调大。比如原来 batch8 显存爆了开检查点后 batch8 能跑了但计算量多了 40%GPU 利用率反而从 70% 掉到 50%。正确的做法是开了检查点就把省下来的显存换成更大的 batch让每个 step 的计算密度提上去把重算的开销摊薄。我自己的经验是开检查点后 batch 至少要翻倍否则纯亏。如果显存省下来但 batch 加不上去比如被序列长度卡住那检查点就是负优化不如换个思路。3.2 8-bit 优化器省的是状态加的是量化开销Adam 优化器要为每个参数存两份状态一阶矩、二阶矩7B 模型 fp32 状态下光优化器状态就 56GB。bitsandbytes 的 8-bit Adam 把这部分压到 14GB确实香。但每次更新都要做量化、反量化这些操作本身要读写显存、启动 kernel。在小模型或者小 batch 下优化器更新占总时间的比例不低8-bit 带来的额外开销可能吃掉收益。实测下来7B 以上模型、batch 足够大时8-bit 优化器的吞吐损失通常在 5% 以内可以接受但 1B 以下的小模型损失可能到 15%这时候就得掂量了。提示8-bit 优化器省的是优化器状态显存不省激活显存。如果你的瓶颈是激活长序列、大 batch它帮不上忙。3.3 LoRA计算量太小kernel 启动成了瓶颈LoRA 只训练低秩矩阵参数量可能只有全量的 1%显存省得飞起。但正因为计算量小每个 step 的 kernel 执行时间很短kernel launch 的固定开销占比就上去了。GPU 大部分时间在等 CPU 发指令而不是在算。这种情况在单卡小模型上特别明显。解决办法是增大 batch、增大 rank、或者用torch.compile把多个小 kernel 融合成大的。我试过在 LoRA 训练里把 batch 从 4 提到 32GPU 利用率直接从 25% 拉到 75%效果立竿见影。3.4 频繁的 CPU-GPU 同步隐形的性能杀手这个坑很隐蔽。代码里但凡出现.item()、.cpu()、print(loss)、if loss threshold这类操作都会强制同步让 GPU 停下来等 CPU。一个 step 里同步几次GPU 利用率就废了。我见过最离谱的一段代码在训练循环里每个 step 都loss.item()打印还顺手torch.cuda.memory_allocated()查显存。这两个操作加起来GPU 利用率从 85% 掉到 40%。把打印改成每 100 步一次利用率立刻回来。4. 把 GPU 真正喂饱的实操手段4.1 先定位瓶颈用 profiler 而不是猜在动手优化之前先搞清楚瓶颈在哪。PyTorch 自带的 profiler 是最靠谱的工具import torch from torch.profiler import profile, ProfilerActivity with profile( activities[ProfilerActivity.CPU, ProfilerActivity.CUDA], scheduletorch.profiler.schedule(wait1, warmup1, active3, repeat1), on_trace_readytorch.profiler.tensorboard_trace_handler(./log), record_shapesTrue, profile_memoryTrue, ) as prof: for step, batch in enumerate(loader): if step 6: break train_step(batch) prof.step()跑完之后看 TensorBoard 的 trace重点看三件事GPU 上有没有大段空白说明在等数据或等 CPU、kernel 是不是又小又碎说明 launch 开销大、显存带宽利用率高不高说明是不是 memory-bound。如果 GPU 时间线上大片空白问题在数据加载或 CPU 侧如果 kernel 密集但每个都很短问题在 kernel launch如果 kernel 很长但 SM occupancy 低问题在并行度不够。4.2 数据加载让 GPU 永远有活干数据加载是最常见的瓶颈尤其在磁盘慢、预处理重的时候。几个关键参数DataLoader(num_workersN)N 一般设成 CPU 核心数的 1/4 到 1/2太多反而抢资源。我一般从 4 开始试。pin_memoryTrue让数据放在锁页内存里H2D 拷贝能走 DMA快很多。persistent_workersTrue避免每个 epoch 重建 worker 进程。prefetch_factor2或更高提前预取几个 batch。如果预处理特别重比如图像增强、tokenize考虑离线预处理把结果存成二进制文件训练时直接读。我做过一个项目把在线 tokenize 改成离线数据加载时间从每 step 200ms 降到 20msGPU 利用率从 45% 拉到 90%。4.3 增大 batch最直接的提利用率手段batch 越大每个 step 的计算量越大kernel launch 开销被摊薄GPU 越容易吃满。显存省下来之后第一件事就是把 batch 加到显存快满为止。具体怎么加先设一个保守值跑几步看显存峰值然后按 1.5 倍递增直到接近显存上限留 10% 余量给碎片。如果 batch 加到某个值后 OOM就回退一档。配合梯度累积可以在显存受限时模拟大 batchaccum_steps 4 for i, batch in enumerate(loader): loss model(batch) / accum_steps loss.backward() if (i 1) % accum_steps 0: optimizer.step() optimizer.zero_grad()注意梯度累积不减少 kernel launch 次数它只是让优化器更新频率降低。所以它对显存友好但对 kernel launch 瓶颈帮助有限。真正要提利用率还是得靠单次 batch 变大。4.4 torch.compile把碎 kernel 融合起来PyTorch 2.0 引入的torch.compile能把多个小算子融合成一个大 kernel直接减少 launch 次数。对 LoRA 这种计算量小的场景效果特别明显model torch.compile(model, modemax-autotune)max-autotune模式会花几分钟编译但跑起来之后吞吐通常能提 20% 到 50%。我实测过一个 LoRA 微调任务编译前 GPU 利用率 30%编译后 65%训练时间砍掉近一半。注意torch.compile首次编译有开销短任务可能不划算。另外它跟某些自定义 kernel 可能不兼容遇到报错先退回modedefault试试。4.5 混合精度和 FlashAttention算法级的双赢这两个是少数既省显存又提吞吐的手段。混合精度bf16/fp16让 Tensor Core 跑起来矩阵乘速度翻倍显存减半。FlashAttention 重写了注意力计算减少显存读写长序列下提速明显。from torch.cuda.amp import autocast, GradScaler scaler GradScaler() with autocast(dtypetorch.bfloat16): output model(input) loss criterion(output, target) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()bf16 比 fp16 更稳不容易溢出A100/H100 上优先用 bf16。消费级卡如果只支持 fp16记得配 GradScaler。5. 排查清单和常见问题速查5.1 GPU 吃不满的排查顺序遇到 GPU 利用率低我一般按这个顺序查看数据加载profiler 里 GPU 时间线有没有大段空白有就是数据瓶颈调num_workers、pin_memory、prefetch_factor。看同步操作代码里有没有.item()、.cpu()、频繁 print有就删掉或降频。看 batch size显存还有余量吗有就加大 batch。看 kernel 粒度kernel 是不是又小又碎是就上torch.compile。看并行度SM occupancy 低不低低就检查模型并行策略、序列长度。看优化器8-bit 优化器在小模型上是不是负优化是就换回 fp32。5.2 常见问题速查表现象可能原因解决方向GPU-Util 低但显存满显存瓶颈batch 太大减 batch 或开检查点GPU-Util 低且显存空数据或 CPU 瓶颈查 DataLoader、同步操作GPU-Util 高但训练慢kernel 效率低上混合精度、FlashAttention开了检查点更慢batch 没跟着加加大 batch 摊薄重算LoRA 训练利用率低kernel 太碎torch.compile、加大 batch多卡利用率不均负载不均衡检查数据分片、梯度同步5.3 几个我踩过的坑坑一以为显存省了就能加 batch结果序列长度卡死。长序列训练里激活显存跟序列长度平方相关检查点省下来的显存可能只够把 batch 从 2 加到 3收益有限。这种情况不如考虑序列并行或者 FlashAttention。坑二num_workers设太大反而慢。worker 进程之间抢 CPU、抢内存带宽设成 CPU 核心数反而拖慢。我一般设 4 到 8配合persistent_workers。坑三torch.compile编译时间被忽略。小任务跑 100 步就结束编译花 3 分钟纯亏。短任务别开或者用modereduce-overhead减少编译时间。坑四多卡训练里梯度同步成了瓶颈。DDP 每个 step 都要 all-reduce 梯度如果模型大、卡间带宽低GPU 就在等通信。可以开gradient_as_bucket_viewTrue或者用梯度累积减少同步频率。6. 我个人的一点经验显存优化和吞吐优化本质上是两个方向的事。省显存是为了能跑起来提吞吐是为了跑得快前者是及格线后者才是性能。很多人做完显存优化就以为大功告成其实真正的调优才刚开始。我自己的习惯是每做一次显存优化就重新测一遍吞吐并且尝试把省下来的显存换成更大的 batch。如果 batch 加不上去就说明这个优化手段在当前场景下是负收益得换思路。这个循环做几轮基本就能把 GPU 喂到 80% 以上。最后分享一个小技巧训练脚本里加一个简单的吞吐监控每 100 步打印一次samples/sec和 GPU 利用率跑起来一眼就能看出优化有没有效果。比事后翻 profiler 方便多了。import time start time.time() for step, batch in enumerate(loader): train_step(batch) if step % 100 0 and step 0: elapsed time.time() - start throughput 100 * batch_size / elapsed print(fstep {step}, throughput {throughput:.1f} samples/s) start time.time()这个监控加上nvidia-smi -l 1在旁边开着优化效果一目了然。踩过的坑多了就会明白一个道理GPU 吃不满八成不是显存的锅而是你没把活连续地喂给它。