显存估算与混合精度：大模型训练OOM避坑实战指南

很多朋友做大模型训练第一关不是代码能不能跑而是显存够不够。这个问题听着简单但如果不把显存估计和混合精度训练放在一起算十有八九会在第一个OOM面前懵掉。我最近在A100上配一个13B的全参数训练任务估算表写了三页纸结果第一版配置跑了一百多个step就炸了显存。回头复盘时发现问题不在工具而在对显存构成的理解。这篇整理就是我当时最想看的东西训练一个模型显存到底花在哪几块混合精度训练到底省了什么、没省什么以及真的被OOM卡住时按什么顺序排查最有效。适合准备在单卡或多卡上做全参数训练的同学也适合刚入大模型训练、想搞清楚一张卡到底能塞多大模型的读者。1. 显存到底被谁吃掉了四大占用的测算模型1.1 一张公式表把训练态说清楚大模型训练显存和推理显存完全是两码事。推理只需要权重和一部分临时激活训练却要同时容纳权重、梯度、优化器状态以及反向传播所需的大量中间结果。把这几块分开算才能定位OOM的真正来源。先定义几个符号P 表示模型参数量b 表示单个参数的字节数。FP32 是4字节FP16/BF16 是2字节INT8/FP8 是1字节。全FP32训练时用AdamW优化器显存的权重相关部分如下项目计算说明模型权重4PFP32保存一份梯度4P反向传播累积梯度FP32优化器一阶动量 m4PAdam 需要优化器二阶动量 v4PAdam 需要合计16P不含激活值、临时张量也就是说一个 7B 模型全FP32加AdamW纯状态就要 16 × 7 112 GB。一张 A100 80G 根本装不下更别提反向传播还要吃激活。1.2 拿7B模型实际算一遍再说混合精度。主流做法是 FP16/BF16 前向和反向同时保留一份 FP32 主权重更新时用主权重做累积避免低精度浮点的舍入误差。这样权重相关部分就变成FP32主权重4P低精度权重副本2P梯度4P在反向传播中通常按FP32累积优化器状态8P合计18P注意这里比全FP32多了 2P。这是因为混合精度同时保留了两份权重。7B 就是 126 GB比全FP32反而多了14GB。很多教程一张嘴就说混合精度省显存这里得立刻打个问号。那混合精度到底省在哪答案是计算速度和激活值不是模型状态。这一点理解透了后面所有配置决策都不会跑偏。1.3 为什么说混合精度省显存是个需要拆开看的说法如果只看权重、梯度、优化器这三块混合精度的显存收益几乎为零甚至还有开销。但实际工程中混合精度训练通常能让你把 batch size 开得更大这是因为计算图的中间张量大部分是低精度。autocast 作用下的线性层、注意力分数、MLP 中间结果以 FP16/BF16 存在相比 FP32 直接少一半。激活值存储减半。反向传播依赖前向的中间结果这些结果以低精度保存长序列大batch下节省非常可观。通信量减少。梯度是低精度传输的话多卡训练带宽压力降低。所以混合精度省显存这个结论只有在激活值占比高的场景才严格成立。如果序列短、batch小模型状态是大头混合精度的显存优势不明显。但如果序列长度4096、batch凑到4以上激活值轻松超过20GB混合精度和checkpointing带来的减幅就是决定性的。2. 混合精度训练省下的不是状态是计算和激活2.1 FP16/BF16与FP32的取舍混合精度不是简单地把模型切成float16它背后是浮点数数值分布的取舍。FP16有5位指数、10位尾数。它能表示的最大值约65504精度大约千分之一到万分之一级别。问题是梯度经常小到1e-7甚至1e-9在FP16下直接变成0也就是下溢。梯度一旦为0权重根本不会更新训练看起来loss不动其实就是被精度坑了。BF16则是8位指数、7位尾数。指数范围跟FP32一致能表示的最小非零数约1e-38几乎不会下溢。但尾数只有7位相当于小数点后两位半的十进制精度。所以BF16不会让梯度直接消失但每个数值的相对误差比FP16大。实践中我把这两种精度总结为FP16适合对精度要求中等、动态范围可控的场景必须配梯度缩放BF16适合训练大模型动态范围够用省去了loss scaling绝大多数框架默认都选它。2.2 master weight / loss scaling / 动态scale机制这三个机制是混合精度能稳定训练的根基。master weight 指的是保留一份FP32权重副本。前向传播和反向传播用低精度权重计算但每次优化器更新时是在FP32主权重上完成的。为什么要绕这么一圈因为 Adam 更新是很小的增量如果直接在FP16权重上加这次增量很可能小于FP16的最小分辨率权重纹丝不动。保留主权重相当于给权重穿了一件防弹衣。loss scaling 是为FP16准备的。反向传播得到的梯度除以scale之前是放大的这样梯度就不会因为过小而变成0。PyTorch 的 GradScaler 会在检测到 inf/NaN 时自动调整 scale这就是动态缩放。from torch.cuda.amp import GradScaler scaler GradScaler() for step, batch in enumerate(dataloader): with torch.autocast(device_typecuda, dtypetorch.float16): loss model(batch) scaler.scale(loss).backward() scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) scaler.step(optimizer) scaler.update()如果用 BF16基本不需要 GradScalerautocast 的 dtype 传 torch.bfloat16 就行。这能省不少调参精力也是我推荐BF16的一个重要原因。2.3 在训练循环里如何落地混合精度不是把所有张量转成FP16这么简单。你如果自己手动把 input 和 model 都 .half()大概率训练一两步就 NaN。原因在于LayerNorm、Softmax、损失函数这些计算对精度极其敏感低精度会导致数值漂移。PyTorch 的 autocast 做得比较聪明的地方是维护了一张算子白名单支持低精度的算子用FP16数值敏感的算子自动保留FP32。实战里我建议不要手动到处 .half()直接用 autocast 包裹前向和 loss 计算让框架来决策。with torch.autocast(device_typecuda, dtypetorch.bfloat16): logits model(input_ids) loss loss_fn(logits, labels)这里再提醒一句如果你的模型里自定义了某些 op比如用纯 Python 手写的 attentionautocast 不会自动帮你处理。需要手动把中间计算的张量降到合适的精度或者干脆用 torch.bfloat16 显式转换。这是很多混合精度跑起来很慢或不省显存的隐藏原因。2.4 从13B模型的角度看收益我们用13B、seq4096、batch1这个配置来算一笔账。模型状态这块混合精度 18P 18 × 13 234 GB这是分布到多卡上的。激活值方面按Transformer的粗略估算后文会展开一条4096的序列要吃将近28GB激活。如果不用混合精度激活区是FP32要56GB用了混合精度激活区降到28GB直接省28GB。模型状态虽然多了26GB但激活占的是每卡本地显存模型状态可以靠多卡摊掉两者性质完全不同。所以混合精度对大模型训练的收益主要体现在本地激活和中间张量上。这也是为什么长上下文场景对混合精度需求极其迫切。序列越长激活越大混合精度的优势就越明显。3. 激活显存那个公式最容易被低估的隐形杀手3.1 激活为什么比参数更吃显存很多人第一次看到激活值估算结果都会吓一跳。原因是反向传播需要前向时每一层的输出这些中间结果必须保留到backward否则没法算梯度。Transformer层数越多、序列越长保留的中间张量就越恐怖。业界有个很实用的经验公式来自Korthikanti等人的论文对于Decoder-only的Transformer每个token大约要占 34 × L × H 字节的激活其中L是层数H是隐藏维度。这个公式假设了标准注意力加MLP的结构实测误差大概在10%-20%以内足够做方案评估。这里要特别注意这个值是未开激活重计算的量。开了gradient checkpointing激活占用会大幅下降后面专门说。3.2 7B/13B/70B在seq4096下的预算表我把常见模型规格代入公式算出来更有体感。模型层数 L隐藏维度 H每token激活seq4096单条激活7B324096约4.46MB18.3GB13B405120约6.96MB28.5GB70B808192约22.3MB91.3GB注意这只是 batch1 的数量。如果 batch413B的激活直接来到114GB单张80G卡光激活就爆。很多OOM就是死在这里——模型状态通过了多卡分片塞进去了结果batch里每多一条序列都在疯狂吃激活而你没提前算。3.3 checkpointing、激活partition与计算换显存的取舍gradient checkpointing 不是把所有前向结果都丢掉而是只保留每个Transformer block的输入。反向传播时如果需要某层的中间结果就临时用保存的输入再算一遍前向。这样激活占用通常能减到原来的1/4甚至更低代价是大约30%的额外前向计算。在HuggingFace Transformers里开启非常方便model.gradient_checkpointing_enable()我实际跑过一个13B的LoRA任务开checkpointing之后激活从28GB降到不到6GB训练速度只慢了大约25%。对于显存紧张但不那么追求极限吞吐的场景这笔买卖很划算。再往下走还有 activation offload把激活张量搬到CPU或NVMe换显存但不增加计算。代价是PCIe传输延迟速度损耗比checkpointing更明显。常规训练里我建议优先用checkpointingoffload留给推理而不是训练。4. 从OOM到摸清底细我的显存排查与验证流程4.1 先量后调两条命令看清峰值遇到OOM第一反应不该是狂调batch size而是先量化。PyTorch给了两个非常好用的接口torch.cuda.reset_peak_memory_stats() # 训练几个step print(f峰值显存: {torch.cuda.max_memory_allocated() / 1024**3:.1f} GiB) torch.cuda.memory_summary(abbreviatedTrue)max_memory_allocated 给出的是这个进程在GPU上分配过的最大显存比 nvidia-smi 的进程占用更准确。nvidia-smi 显示的是驱动角度的reserved memory包含CUDA缓存和碎片不代表你真正创建的张量大小。只看 nvidia-smi 往往会高估实际需求。4.2 一次完整的OOM定位过程说一个我最近的真实案例。13B全参数训练4卡A100 40G开BF16混合精度和ZeRO-2。按公式模型状态126GBZeRO-2分片后每卡约63GB这已经超过40G理论上就该预警了。但我当时想省事没仔细算激活就开跑了。结果前100个step还好第120步开始OOM而且每次都在同一个位置崩。后来打了 max_memory_allocated发现峰值38.5GB离40G只差一点点。问题很明确ZeRO-2把梯度和优化器分片在4卡上但每卡仍然持有全量主权重和低精度权重这部分就有42GB光这已经超了激活再压上来必死。我的调整路径是从ZeRO-2切到ZeRO-3让参数也分片模型状态每卡降到约31.5GB。batch size保持4不动开gradient checkpointing激活从28GB降到约7GB。峰值降到36GB左右稳定训练速度损失在可接受范围内。整个过程最强的教训是先用公式估算再上卡验证。如果一开始就算清楚63GB装不进40G也就不用白跑两小时等OOM了。4.3 常见误判与兜底手段排查显存问题要会看现象我整理了几个踩过的坑现象真实原因对策训练一两步就NaN混合精度下LossScale失效或手动half导致敏感算子低精度改用autocast检查GradScaler的scale是否被反复调小特定step稳定OOM数据集中存在超长样本padding后序列长度暴涨按长度分组batch或设max_seq_len截断GPU利用率低但显存不够大量临时张量在autocast外创建膨胀成FP32把整段前向和loss放进autocast上下文cudaMalloc 频繁触发显存碎片化严重设置PYTORCH_CUDA_ALLOC_CONFexpandable_segments:True还有个容易被忽略的坑torch.cuda.empty_cache() 并不会把显存还给你它只是把缓存池清空、归还给驱动但下一次分配又要重新建立缓存。OOM时调用empty_cache通常是自欺欺人该排查的峰值排查别指望清理能解决紧张局面。4.4 各种配置的对显存的影响如果显存只差一点点优先考虑以下三个微调手段而不是粗暴降低batch把优化器换成功率更低的Adafactor或8bit Adam优化器状态从8P降到2P左右13B能省78GB的分布压力。调整allocator策略PYTORCH_CUDA_ALLOC_CONFexpandable_segments:True 在有些场景能减少碎片但不是所有卡都生效自己测了才知道。检查模型是否加载了不必要的fp32副本比如 .half() 之后又 hot swap导致同一份参数在显存里留了两份。打印峰值后对比估算值很容易看出来。5. 多卡训练下的显存账单ZeRO各阶段的取舍5.1 DP到ZeRO-3每一层分片省了什么单卡放不下就要靠多卡把模型状态摊开。PyTorch DDP并不省显存它每张卡都持有全量模型状态只做梯度同步。真正省显存的是ZeRO系列DeepSpeed和FSDP都能实现。把7B混合精度的126GB模型状态拿出来看并行方式每卡维护内容4卡估算每卡模型状态DDP主权重低精度权重梯度优化器状态全部全量126GBZeRO-1优化器状态分片其他全量约84GBZeRO-2优化器状态梯度分片权重全量约63GBZeRO-3参数梯度优化器状态全部分片约31.5GBZeRO-3看起来最省显存但每个step都要把参数从多卡聚合到当前卡上做前向通信量比ZeRO-2高一个量级。所以选型不是越高级越好而是够用就好。5.2 DeepSpeed/FSDP配置要点我最常给到用户的建议单机4卡以下、模型不超过13BZeRO-2基本够模型到了30B以上或者目标是单卡能塞更多才上ZeRO-3。DeepSpeed的配置长这样zero_optimization: stage: 2 offload_optimizer: device: cpu pin_memory: true这里 offload_optimizer 是把优化器状态放到CPU内存能进一步降低GPU占用代价是每次更新都要走PCIe。实测情况下offload优化器会让step时间增加30%到50%如果GPU显存不缺别开。5.3 显存省了通信和时间去哪了ZeRO省显存本质是拿通信换显存前向要all-gather参数反向要reduce-scatter梯度optimizer step前还可能再来一轮通信。如果你把batch size剁到2来迁就ZeRO-3那GPU利用率通常很低因为大量时间卡在等待通信完成。我一直坚持的原则是先用激活手段混合精度、checkpointing把每卡峰值降到80%以下再考虑分布式分片不要一上来就ZeRO-3offload。前者牺牲的是固定比例的计算后者牺牲的是通信延迟通信瓶颈在多卡扩展下远比计算代价难补。6. 实测参考表与我这段时间跑模型的经验6.1 常用规模/卡型/配置参考表最后给一张我常拿来当起点的配置参考表基于BF16混合精度未开激活offload模型规模推荐最小显存方案备注7B1×A100 80G checkpointing全参训练有点紧LoRA更宽松7B4×A100 40G ZeRO-2较稳可开batch813B4×A100 80G ZeRO-2 checkpointing模型状态每卡约58GB激活控制在5GB内13B8×A100 40G ZeRO-3 checkpointing通信压力大但能完整跑全参70B8×A100 80G ZeRO-3 offload optimizer勉强全参通常建议先试LoRA/QLoRA这些数字不是精确到字节的只是我用来做方向判断的基准。每换一个模型我一定会花二十分钟把公式重新填一遍16P或18P的模型状态、34LH × seq × batch的激活、checkpointing的折扣、ZeRO阶段的分摊然后再上卡验证峰值。这半小时能省下后面一整晚的修复时间。6.2 几个我不会再犯的操作习惯不要用梯度累积代替显存优化。batch太小再累积梯度训练曲线会显得特别不稳定而且累积步数变长debug起来很受罪。先把batch和seq调到一个不OOM的稳妥值再根据收敛曲线上调。开混合精度后batch size翻倍不一定掉精度但一定要盯loss曲线。FP16或BF16改变的是数值表示偶尔会在微调任务上出现损失函数变吵的情况这时候优先排除LossScale问题而不是急着把batch缩回去。用显存换时间时要节制。checkpointing能开但别同时把activation offload、CPU offload、ZeRO-3全开了每一步都在PCIe上爬训练速度会惨不忍睹。优先砍计算缓存让给通讯。排查OOM时先把峰值打出来再动配置。我一直觉得算一次显存预算胜过在卡上焦虑一晚上。如果你也经常被OOM打断试试把这张表贴在你工位上下次改配置前先花几分钟填满它你会发现自己省下很多不必要的返工。