★ FEATURED ARTICLE 1. 为什么GPU显存成了DNN训练的硬瓶颈做过深度学习训练的人都有一个共同体会模型结构还没跑通显存先炸了。尤其是近些年Transformer架构横扫CV和NLP两个领域之后参数量从几百万飙升到几十亿甚至上千亿一张消费级显卡连模型权重都装不下更别提训练过程中还要缓存激活值、梯度、优化器状态这些额外开销。我刚开始接触大模型微调的时候用一张24GB显存的卡跑一个7B参数的模型batch size只能设成1序列长度还得砍到512训练速度慢得让人抓狂。后来一步步研究显存优化的各种技术才慢慢把利用率提上来。这篇内容就是把我这些年积累的GPU显存优化经验做一次系统梳理。从显存到底被谁吃掉了到混合精度、梯度检查点、ZeRO系列、Offload、FlashAttention这些主流方案各自的原理和适用场景再到实际训练中怎么组合使用这些技术把显存压到最低我都会掰开揉碎讲清楚。不管你是刚入门的学生还是在工业界做模型训练的工程师只要你的工作涉及DNN训练这些内容都能直接拿来参考。先给一个直观的感受一个标准的训练步骤中显存占用大致来自四个部分——模型参数、梯度、优化器状态、激活值。以Adam优化器为例每个参数需要存储FP32的权重副本、FP32的梯度、FP32的一阶动量和二阶动量加起来就是4倍参数量乘以4字节。一个1B参数的模型光优化器状态就要占16GB。再加上激活值显存需求轻松突破30GB。这就是为什么显存优化不是可选项而是必选项。2. 显存占用的四大来源逐一拆解2.1 模型参数与梯度的显存开销模型参数是显存占用的基础部分。假设模型有P个参数用FP32存储那参数本身就占4P字节。梯度跟参数一一对应也是4P字节。如果使用混合精度训练通常会维护一份FP16的参数副本用于前向和反向计算同时保留一份FP32的主权重用于参数更新。这样参数部分就变成了2PFP16 4PFP32 6P字节梯度也类似。举个例子一个7B参数的模型纯FP32训练时参数加梯度就是7B × 8 56GB这还没算优化器状态和激活值。所以现在几乎没有人在FP32下训练大模型混合精度是标配。2.2 优化器状态的显存黑洞优化器状态往往是被忽视的显存大户。不同优化器的状态量差异很大优化器类型每个参数的状态数额外显存FP32SGD00SGD Momentum14PAdam28PAdamW28PAdafactor分解存储约2P~4PAdam系列优化器因为要维护一阶动量和二阶动量每个参数额外需要8字节。7B模型用Adam训练优化器状态就是56GB。这也是为什么ZeRO Stage 1专门用来切分优化器状态——它确实是最大的那块。2.3 激活值被低估的显存杀手激活值是前向传播过程中每一层的输出反向传播时需要用来计算梯度。激活值的显存占用跟batch size、序列长度、隐藏维度都成正比。在Transformer结构中激活值占用往往是参数量的好几倍。具体来说一个L层Transformer每层的激活值包括注意力层的Q、K、V矩阵注意力分数矩阵FFN中间层的输出等。以序列长度S、隐藏维度H、batch size B为例单层注意力分数矩阵就是B × num_heads × S × S当S2048时这个矩阵的大小相当可观。很多人只关注模型参数能不能装下忽略了激活值才是训练时显存溢出的最常见原因。推理时可以用no_grad省掉激活值但训练时省不了。2.4 临时缓冲区与碎片化除了上述三大块CUDA上下文本身要占几百MBcuDNN和cuBLAS的workspace需要预留空间通信操作如AllReduce需要临时缓冲区。更隐蔽的是显存碎片化——反复申请和释放不同大小的显存块会导致虽然有足够的空闲显存总量但没有一块连续的空间能满足分配请求直接OOM。3. 混合精度训练最基础的显存优化手段3.1 FP16与BF16的选择逻辑混合精度训练的核心思路是用低精度格式存储和计算同时保留高精度主权重保证数值稳定性。FP16和BF16是两种主流选择FP161位符号 5位指数 10位尾数。尾数精度高但指数范围窄最大约65504容易出现上溢。训练时通常需要loss scaling来防止梯度下溢。BF161位符号 8位指数 7位尾数。指数范围和FP32一致不需要loss scaling但尾数精度低对某些敏感操作可能有影响。我的经验是如果硬件支持BF16Ampere及以后的架构优先用BF16省心。如果只有FP16那就老老实实配loss scaling用动态loss scaling基本不需要手动调。3.2 混合精度训练的实操配置以PyTorch为例最简洁的方式是使用torch.cuda.ampfrom torch.cuda.amp import autocast, GradScaler scaler GradScaler() for data, target in dataloader: optimizer.zero_grad() with autocast(dtypetorch.bfloat16): output model(data) loss criterion(output, target) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()如果用的是BF16其实可以不用GradScaler直接loss.backward()就行。但用FP16的话GradScaler是必须的。注意混合精度不是万能的。某些操作如softmax、layer norm、loss计算对精度敏感框架会自动把这些操作保持在FP32下执行。如果你手动实现了某些自定义算子记得检查是否需要强制FP32。3.3 混合精度能省多少显存粗略估算参数和激活值从FP32变成FP16/BF16这两部分显存直接减半。但优化器状态通常还是FP32除非用8-bit优化器所以总体显存节省大约在30%~40%左右。对于一个原本需要40GB显存的训练任务混合精度后可能降到25GB左右效果还是很明显的。4. 梯度检查点用计算换显存的经典方案4.1 梯度检查点的原理梯度检查点Gradient Checkpointing也叫激活重计算Activation Recomputation。核心思想是前向传播时只保存部分层的激活值检查点反向传播时从最近的检查点重新做一次前向计算恢复中间激活值再计算梯度。这样做的好处是激活值显存从O(L)降到O(√L)如果每√L层设一个检查点代价是增加一次前向计算训练时间大约增加30%左右。4.2 PyTorch中的实现方式PyTorch提供了torch.utils.checkpointfrom torch.utils.checkpoint import checkpoint class TransformerBlock(nn.Module): def forward(self, x): # 方式一对整个block做checkpoint return checkpoint(self._forward, x, use_reentrantFalse) def _forward(self, x): x x self.attention(self.norm1(x)) x x self.ffn(self.norm2(x)) return xuse_reentrantFalse是新版推荐的方式支持更多场景如输入不需要梯度的情况。4.3 梯度检查点的适用场景与注意事项梯度检查点最适合以下场景模型层数很深激活值占用远大于参数占用显存紧张但计算资源相对充裕训练吞吐量不是首要瓶颈注意事项不是所有层都值得做checkpoint。通常对激活值大的层如注意力层做对激活值小的层如LayerNorm做反而浪费计算。checkpoint的粒度需要调。每个block都做checkpoint会增加太多计算每隔几层做一次更划算。使用checkpoint后dropout等随机操作在重计算时可能产生不一致的结果需要确保随机种子固定。5. ZeRO系列数据并行的显存优化利器5.1 ZeRO的三个阶段ZeROZero Redundancy Optimizer是微软提出的一套数据并行显存优化方案核心思想是消除数据并行中的显存冗余。在标准数据并行中每张卡都保存完整的模型参数、梯度和优化器状态这造成了大量重复。ZeRO分三个阶段逐步切分这些状态阶段切分内容显存节省N卡通信量增加Stage 1优化器状态约4倍少量Stage 2优化器状态 梯度约8倍中等Stage 3优化器状态 梯度 参数约N倍较大以7B模型、8卡训练为例标准数据并行每卡需要约80GB显存ZeRO Stage 1降到约20GBStage 2降到约10GBStage 3降到约2GB每卡只存1/8的参数。5.2 ZeRO Stage 3的通信开销与权衡Stage 3虽然显存节省最明显但通信量也最大。每次前向传播需要从其他卡收集参数反向传播需要再次收集并做ReduceScatter。在慢速网络如普通以太网上通信可能成为瓶颈导致GPU利用率大幅下降。我的经验是如果卡间是NVLink或高速InfiniBandStage 3完全可用如果是普通千兆以太网Stage 2可能是更平衡的选择。另外Stage 3配合梯度检查点使用效果最好因为参数切分后每卡计算量减少重计算的开销相对可接受。5.3 实操配置示例以DeepSpeed为例一个典型的ZeRO Stage 2配置{ train_batch_size: 64, gradient_accumulation_steps: 4, fp16: { enabled: true, loss_scale: 0, initial_scale_power: 16 }, zero_optimization: { stage: 2, allgather_partitions: true, allgather_bucket_size: 5e8, overlap_comm: true, reduce_scatter: true, reduce_bucket_size: 5e8, contiguous_gradients: true } }关键参数说明overlap_comm: true通信和计算重叠隐藏通信延迟allgather_bucket_size和reduce_bucket_size控制通信桶大小太小通信频繁太大占显存contiguous_gradients: true梯度连续存储减少碎片6. Offload技术把显存压力转移到CPU和NVMe6.1 CPU Offload的基本原理Offload的思路很直接GPU显存不够那就把暂时不用的数据放到CPU内存甚至NVMe硬盘上需要时再传回GPU。ZeRO Offload就是把优化器状态和梯度放到CPU内存参数更新在CPU上完成只把更新后的参数传回GPU。这样做的好处是显存需求大幅降低单卡也能训练大模型。代价是PCIe传输带宽远低于GPU显存带宽训练速度会明显下降。实测下来CPU Offload通常会让训练速度降低2~5倍具体取决于PCIe版本和CPU内存带宽。6.2 NVMe Offload的适用场景当CPU内存也不够时可以把优化器状态放到NVMe SSD上。NVMe的顺序读写带宽可以达到3~7GB/s虽然比内存慢但比机械硬盘快得多。NVMe Offload适合以下场景模型极大CPU内存装不下优化器状态训练对速度要求不高但对模型规模有硬性要求有高速NVMe阵列可用注意Offload对PCIe带宽非常敏感。如果用的是PCIe 3.0 x16带宽约16GB/sOffload后训练速度下降会非常明显。PCIe 4.0 x16翻倍到32GB/s情况会好很多。6.3 Offload的配置与调优DeepSpeed中开启CPU Offload{ zero_optimization: { stage: 3, offload_optimizer: { device: cpu, pin_memory: true }, offload_param: { device: cpu, pin_memory: true } } }pin_memory: true使用锁页内存加速CPU和GPU之间的传输。但锁页内存会占用物理内存不能超过系统限制。调优建议只Offload优化器状态不Offload参数通常能取得较好的平衡如果CPU核心数多可以设置numa相关参数优化内存访问监控PCIe带宽利用率如果接近饱和说明Offload是瓶颈7. FlashAttention与算子级显存优化7.1 标准注意力的显存问题标准自注意力的计算过程是Q乘以K的转置得到注意力分数矩阵softmax后再乘以V。这个注意力分数矩阵的大小是B × num_heads × S × S。当S4096时单头单batch的分数矩阵就是16M个元素FP16下占32MB。如果有多头和多层这个占用会迅速累积。7.2 FlashAttention的核心思路FlashAttention通过分块计算tiling和重计算recomputation避免了显式存储完整的注意力分数矩阵。它把Q、K、V分成小块在SRAM中完成注意力计算只把最终结果写回HBM。这样注意力分数矩阵根本不落到显存中显存占用从O(S²)降到O(S)。实际效果在序列长度2048时FlashAttention能减少约30%的激活值显存序列长度越长节省越明显。到8192时节省可以超过50%。7.3 其他算子级优化除了FlashAttention还有一些算子级优化值得关注Memory-Efficient AttentionPyTorch内置的优化注意力实现思路类似FlashAttentionFused Adam把Adam的多个操作融合成一个CUDA kernel减少中间变量的显存分配8-bit Optimizer用8位量化存储优化器状态显存占用降到原来的1/4这些优化通常只需要改几行代码或加一个配置项性价比很高。8. 组合策略如何在实际训练中搭配使用8.1 不同规模模型的推荐配置根据模型规模和硬件条件我整理了一套推荐配置模型规模硬件条件推荐方案1B单卡24GB混合精度 梯度检查点1B~7B单卡80GB混合精度 FlashAttention 梯度检查点7B~13B多卡80GBZeRO Stage 2 混合精度 FlashAttention13B~70B多卡80GBZeRO Stage 3 CPU Offload 混合精度70B多卡 NVMeZeRO Stage 3 NVMe Offload 混合精度8.2 显存优化的优先级排序当显存不够时建议按以下顺序尝试优化开启混合精度几乎无成本省30%~40%使用FlashAttention改几行代码省注意力部分的激活值开启梯度检查点增加30%计算省大量激活值使用ZeRO Stage 1或2多卡场景省优化器状态和梯度使用ZeRO Stage 3省参数但通信开销大开启Offload最后手段速度下降明显8.3 监控与调优工具推荐几个实用的显存监控工具nvidia-smi最基础的显存查看但看不到碎片化情况torch.cuda.memory_summary()PyTorch内置能看到分配和保留的显存py3nvmlPython接口的NVML可以编程监控DeepSpeed的flops profiler能看到每层的显存和计算开销实操心得训练时先用小batch跑几步用torch.cuda.max_memory_allocated()记录峰值显存再逐步增大batch size找到OOM的临界点。这样比盲目试错高效得多。9. 常见问题与排查技巧实录9.1 OOM了但nvidia-smi显示还有空闲显存这是典型的显存碎片化问题。PyTorch的缓存分配器会预留显存块但可能没有足够大的连续块满足新分配请求。解决方法设置PYTORCH_CUDA_ALLOC_CONFexpandable_segments:True让分配器支持可扩展段减少不必要的显存分配和释放尽量复用张量在训练循环开始时调用torch.cuda.empty_cache()清理缓存但不要在每步都调会影响性能9.2 混合精度训练出现loss NaNFP16训练时loss变成NaN通常是梯度上溢或下溢导致的。排查步骤检查loss scaling是否正常工作动态loss scaling的初始值是否合适检查是否有除零操作或log(0)的情况尝试用BF16替代FP16BF16的数值范围更大对敏感层如LayerNorm强制FP329.3 ZeRO Stage 3训练速度骤降Stage 3的通信量很大如果训练速度比Stage 2慢很多检查卡间互联是否是NVLink或InfiniBand普通以太网下Stage 3可能不划算allgather_bucket_size是否设置合理太小会导致通信频繁是否开启了overlap_comm让通信和计算重叠梯度检查点是否开启Stage 3下重计算的开销相对更可接受9.4 梯度检查点导致训练不稳定梯度检查点重计算时如果模型中有随机操作如dropout需要确保重计算时使用相同的随机状态。PyTorch的checkpoint默认会保存和恢复RNG状态但自定义实现时容易忽略。另外checkpoint的层如果包含BatchNorm重计算时统计量可能不一致建议在checkpoint层中用LayerNorm替代。9.5 Offload后CPU内存不足CPU Offload会把优化器状态放到内存如果模型很大CPU内存可能先于GPU显存耗尽。解决方法使用NVMe Offload把部分状态放到SSD减少pin_memory的使用锁页内存会占用物理内存增加CPU内存或使用内存压缩技术10. 一些踩坑后的个人体会显存优化这件事最忌讳的就是一上来就堆技术。我见过不少同行模型才1B参数就把ZeRO Stage 3、CPU Offload、梯度检查点全开上结果训练速度慢得没法看最后发现其实混合精度加FlashAttention就够了。优化的顺序应该是从低成本到高成本先试那些几乎不影响速度的方案实在不够再上重武器。另一个体会是显存监控要贯穿整个训练过程。有时候训练刚开始显存够用跑了几百步之后慢慢涨上去最后OOM这通常是内存泄漏或者某些缓存没清理。养成定期打印显存使用情况的习惯能帮你提前发现问题。最后不同框架的显存优化实现差异很大。PyTorch原生的FSDP、DeepSpeed的ZeRO、Megatron-LM的并行策略各有优劣选型时要结合团队的技术栈和硬件条件。没有银弹只有最适合当前场景的组合。
阅读完成 · 觉得有帮助？
