干过大模型训练的人都清楚“显存不够”往往是压垮训练任务的第一根稻草。不管是7B还是13B模型开工前如果不先算清楚显存预算等跑到一半OOM再回头调batch size、重开任务损失的不只是时间还有心态。这篇是这个系列的第三篇专门讲两件事怎么手算模型训练要吃多少显存以及混合精度训练里那些经常被忽略的技术细节。适合正要准备全参数微调、在选卡或者规划序列长度的朋友读完你至少能回答两个问题这条训练管线需要几张80G卡混合精度到底帮我省了哪部分显存显存估计这件事很多人习惯直接抄别人给的“7B大概要多少G”但别人的batch size、序列长度、是否开重计算都和你不一样抄来必翻车。不如花五分钟把手里的模型配置代进公式得到一个误差可接受的量级。混合精度则更微妙它确实能把训练跑得更快、中间张量更小但如果你以为用了AMP就万事大吉结果显存没降多少、loss还时不时NaN那大概率是没理解FP16/BF16的底牌。这篇就把账一五一十算给你看再附上实操排坑清单。1. 显存都去哪儿了训练时的四类开销1.1 模型参数本身最直白的一笔模型参数有多少个就有多少权重数值占显存。但同样一组数字用FP32存和用FP16存体积差一半。FP32每个参数占4字节FP16/BF16每个参数占2字节。7B参数模型如果全用FP32存储光权重就是 7e9 × 4B ≈ 28GBFP16则约14GB。所以“这个模型多少个B”只是起点真正的显存还要看存储精度和训练方式。1.2 梯度反向传播的“路程记录”训练必须做反向传播每个参数都会对应一个梯度用于指导参数往哪个方向更新。梯度张量的大小和参数一样大因此参数如果是FP32梯度FP32就是每参数4字节如果参数用FP16存储梯度通常也会被算成FP16每参数2字节。这里多说一句在PyTorch标准的torch.cuda.amp里如果你的模型参数保持FP32反向传出来的梯度也会是FP32显存占用并不会因为这个模式而下降。很多新手在这里有误解以为用了AMP参数内存就减半其实不是。1.3 优化器状态真正的大头不是参数以最常用的Adam/AdamW为例它要给每个参数额外保存两个状态一阶动量momentum和二阶动量variance这两个都按FP32存储每个状态4字节。另外如果我们把模型参数转成FP16训练为了不让更新过程因为精度太低而发散还需要单独保留一份FP32的“主权重”副本又是4字节。算下来混合精度训练中每个参数的固定开销可拆成FP16参数副本2字节FP16梯度2字节FP32主权重4字节Adam一阶动量4字节Adam二阶动量4字节合计16字节/参数看到没即使模型FP16存储更新时还得背一份FP32主权重和两个FP32状态每个参数依然要16字节。这一点常被营销话术掩盖混合精度训练并没有让Adam优化器下的“固定显存”减半它真正省下的是中间激活张量以及通过降低计算精度换取的算力优势。用纯FP32训练参数4字节梯度4字节两个状态8字节也是16字节。两者打平。不同优化器和精度组合的每参数固定字节数可以直接查下面这张表训练配置参数梯度优化器状态每参数总字节FP32 Adam4B4B8B16BFP16/FP32主权重 Adam2B2B12B16BFP16无主权重 Adam2B2B8B12BBF16 AdamW ZeRO offload2B2B状态可被切分/卸载不定最后一行意味着通过ZeRO等策略可以把优化器状态切到多卡甚至卸载到CPU这才能突破单卡固定开销的瓶颈。1.4 激活值与中间张量藏得最深的“刺客”前向传播时每一层都会产生中间结果比如Transformer里QKV投影输出、注意力分数、MLP中间层输出等。这些结果在反向传播计算梯度时还要再用一次所以不能算完就扔必须暂存在显存里。它们的大小和batch size、序列长度、注意力头数强相关而和参数量关系不大。很多模型看着参数不多却因为序列长、batch大显存直接爆掉罪魁祸首往往是激活值。这一块也是我们可以手动估算的核心部分下一节专门展开。除了上述四类CUDA context、cuDNN workspace、PyTorch缓存分配器的预留区也会占用一小部分显存通常几十MB到几GB不等但不会成为决定性因素估算时可以忽略实测时再通过监控去看。2. 显存估计用一张纸算出你的卡够不够2.1 固定开销先把参数规模装进去固定开销只依赖参数量P和训练配置和batch size、序列长度无关。计算公式非常简单[ 固定显存(GB) P \times bytes_per_param / 1024^3 ]P是参数个数bytes_per_param按上一节表格取。例如一个7B参数模型用FP16FP32主权重Adam训练每参数16字节固定开销就是[ 7 \times 10^9 \times 16 / 1024^3 \approx 112GB ]也就是说只看参数、梯度、优化器状态就已经需要1.5张80G卡了。还没算激活值和batch。这也是为什么很多人说全参数微调7B模型单卡80G根本放不下除非用LoRA、DeepSpeed ZeRO、CPU offload等手段把它拆走。2.2 激活值粗略估算一个能上手算的经验式严格精确计算激活显存需要逐层追踪张量生命周期工程上太累。这里给一个足够用于“拍脑袋”的近似式。假设我们要训练一个标准Transformer模型配置为层数L隐藏维H注意力头数Abatch size B序列长度S使用FP16/BF16存储中间张量每个Transformer层需要暂存的激活量大约为[ 2 \times (19 \times B \times S \times H 3 \times B \times A \times S^2) ]单位是字节。解释一下两部分来源前半部分包括层的输入、QKV投影结果、注意力上下文输出、MLP两层中间结果、残差连接需要保留的输入等大约折合19份B×S×H的二维矩阵后半部分是注意力分数矩阵的三种形态QK^T的结果、softmax之后的概率、dropout掩码或概率版每个都是B×A×S×S的大方块。因为按2字节精度的FP16存储所以前面乘2。这个式子不追求精确但能反映出两个关键趋势序列长度S影响极大因为注意力矩阵是S²增长batch size和隐藏维是线性增长。如果开启了激活检查点gradient checkpointing代价是反向传播时重算一部分前向激活显存可以大幅降低估算时大约只需要保留“每层输入激活”和一层重算所需的中间量整体量级大约是一个layer的激活量可以按上式“乘以1”而不是“乘以L”来近似。2.3 一个7B模型的计算实例我们用常见配置7B模型L32H4096A32序列长度S2048。先算batch size1时每一层激活[ 19 \times 1 \times 2048 \times 4096 159,383,552 ] [ 3 \times 1 \times 32 \times 2048^2 402,653,184 ]两部分相加约5.62亿个元素乘2字节得到约1.12GB。这是单层激活32层都不开检查点则约36GB。固定开销112GB 激活36GB ≈ 148GB两张80G卡能放下但比较紧张。如果不开激活检查点再把batch size提到4激活变成144GB总显存约256GB两张80G就不够了至少需要4张而且这只是压着内存线跑实际情况还要看PyTorch缓存碎片建议再多留余量。配置固定开销激活开销(估算)总显存7BB1S2048无检查点112GB36GB~148GB7BB1S2048开检查点112GB约4GB~116GB7BB4S2048无检查点112GB144GB~256GB13BB1S2048开检查点约208GB约5GB~213GB13B的固定开销13e9×16B≈208GB所以要跑13B全参训练4卡80G是起步5卡更稳。这也是为什么现在大家更青睐LoRALoRA只训练少量新增参数主模型冻结后可省掉优化器状态和梯度的巨额开销大大降低门槛。2.4 从估算到代码用memory stats验证纸面估算完成后一定要让训练代码“自报家门”。PyTorch的CUDA caching allocator提供的数字比nvidia-smi更接近真实逻辑使用量因为显卡驱动显示的显存可能包含预留而未用的部分。在训练循环里打印import torch def print_gpu_memory(): allocated torch.cuda.memory_allocated() / 1024**3 reserved torch.cuda.memory_reserved() / 1024**3 print(fallocated: {allocated:.2f}GB | reserved: {reserved:.2f}GB)每次step结束时调用一次观察allocated的峰值和reserved的差值。如果allocated离80G很远就开始OOM多半是缓存碎片问题。这时可以设置环境变量export PYTORCH_CUDA_ALLOC_CONFexpandable_segments:True让它利用可扩展内存段通常能显著缓解碎片导致的假OOM。另一种方案是max_split_size_mb调小但实际体验上expandable_segments更有效。3. 混合精度训练它做什么不做什么3.1 FP16和BF16一个怕下溢一个怕精度细FP16用1位符号、5位指数、10位尾数能表示的数值范围很小最大约65504最小的normal值约6e-5。它的问题不是“存不下大数”而是“存不了特别小的数”一旦某个中间结果小于约1e-5再往下就变成0这叫下溢。梯度在深层网络中经常会出现非常小的数值正是FP16的噩梦。BF16则用了和FP32相同的8位指数最大值也是FP32同量级范围宽阔得多几乎不会出现下溢但它只有7位尾数精度比FP16更粗。打个比方FP16像一把刻度更细但量程有限的尺子BF16像一把量程很广但刻度很粗的尺子。大模型训练里梯度数值范围的最主要矛盾往往是小数下溢所以新一代GPU比如A100、H100、RTX 30系以后上BF16往往比FP16更稳不需要额外做loss scaling也能跑。反过来早期的V100、T4这些没有BF16硬件加速的卡就只能用FP16 loss scaling。3.2 Loss Scaling为什么损失放大一倍不是玄学既然FP16存不了太小的数那操作办法很朴素既然梯度在进入FP16存储前变得太小那就先把loss乘一个大因子例如1024。因为链式法则反向传播中所有梯度都会跟着“放大同样倍数”原本可能小于FP16下限的值被搬进可表示范围。等真正更新参数前再把梯度除以这个因子恢复真实梯度。这个“先乘后除”的过程就是loss scaling。手动调一个固定scale也行但不保险因为不同训练阶段的梯度量级差很多。于是就有了动态策略每隔一段时间如果梯度没出现inf/NaN就把scale往大调一旦发现梯度中有inf/NaN说明放大过头了就回退半分并跳过这次优化器更新。PyTorch的GradScaler封装的就是这套逻辑。BF16下通常不需要它因为指数范围足够大但如果训练早期loss突然发散也可以临时罩一层来兜底代价不大。3.3 两种混合精度“流派”torch AMP与完整半精度这是最容易被混在一起说的地方必须讲清楚。第一种是PyTorch原生AMP典型写法是torch.autocastGradScaler。它做的事是在算子层面动态选择精度比如矩阵乘法、卷积这类计算密集且适合TensorCore的操作自动把输入转成FP16去算LayerNorm、Softmax这类对精度敏感的操作仍保持FP32。模型参数本身不会整体转成FP16优化器更新依然在FP32参数上做因此不需要额外维护FP32主权重。它的价值是代码改动极小、速度提升明显、中间张量也是FP16从而省内存。但注意它并没有省下“参数梯度优化器状态”那部分固定内存。第二种是“完整半精度训练”典型代表是DeepSpeed和部分手动实现。先调用model.half()或者model.to(dtypetorch.bfloat16)把模型参数整体转成低精度存储再额外保留一份FP32主权重给优化器用。这时训练管线里的权重和中间激活大部分都是低精度整体显存占用会更接近我们第一节表里的“FP16参数2B”那一行。它的固定开销依然是每参数16字节但因为参数本体是2字节而不是4字节通信和某些矩阵乘法的内存带宽压力会小一些。所以一句话总结真正省显存的是“削减优化器状态/激活”而不是把参数从FP32换成FP16这件事本身。如果你的目标是塞进一张卡最有效的手段是LoRA、ZeRO、激活检查点而不是单纯打开AMP。3.4 什么时候用FP16什么时候用BF16我自己的选择标准很简单如果卡是A100/H100/4090/3090这类支持BF16的优先BF16因为省事不需要跟loss scale斗智斗勇对学习率也不那么敏感。如果只有V100/T4/2080Ti这类不支持BF16加速或支持的生态不完善的卡就用FP16 GradScaler。如果模型本身有非常多的Softmax、LayerNorm、注意力打分这类对精度敏感的结构BF16偶尔会在小模型上出现精度损失可以先跑10个step对比FP16和BF16的loss曲线再决定用哪个。如果发现GradScaler的scale一路降到16甚至更低而梯度还在出现inf那大概率不是scale的问题而是梯度爆炸。先看学习率再看梯度裁剪最后看数据里是不是混入了异常样本。3.5 混合精度训练的关键流程无论用哪种流派优化器更新前都有一个必须完成的顺序前向时启用autocastloss得到FP16/BF16计算下的输出。将loss乘上当前scale再backward梯度被放大。反向后先unscale梯度即除以scale并趁机检查是否有inf/NaN。如果梯度正常做梯度裁剪然后调用optimizer.step()更新参数。最后通过scaler.update()调整下次的scale。这里特别容易踩的坑是梯度裁剪的时机。很多人习惯直接调用torch.nn.utils.clip_grad_norm_但如果你用GradScaler梯度此时还是放大后的状态必须先scaler.unscale_(optimizer)再裁剪。顺序错了裁剪阈值等于形同虚设。4. 实战配置、问题速查与显存兜底方案4.1 PyTorch混合精度训练循环最小可跑代码给一段简洁但完整的最小示例。以FP16 GradScaler为例import torch import torch.nn as nn model MyTransformer() optimizer torch.optim.AdamW(model.parameters(), lr1e-5) scaler torch.amp.GradScaler(cuda) # PyTorch 2.x 新写法 # 旧写法torch.cuda.amp.GradScaler() for batch in dataloader: optimizer.zero_grad() with torch.autocast(device_typecuda, dtypetorch.float16): loss model(batch) scaler.scale(loss).backward() # 关键先unscale再clip scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) scaler.step(optimizer) scaler.update()如果换成BF16可以用with torch.autocast(device_typecuda, dtypetorch.bfloat16): loss model(batch)BF16通常不需要GradScaler但unscale_和step的代码结构仍然建议保留因为万一训练中途出现NaN想启用scale时改起来只差一行。另外注意scaler.unscale_在同一个optimizer上只能调用一次重复调用会报错如果你没有调用scaler.unscale_clip_grad_norm_之前也要保证梯度已经是真实值。把上面的流程当成标准模板存下来能少踩很多坑。4.2 显存不够时的“救命”优先级如果真的OOM了按下面这个顺序调整性价比从高到低缩小batch size到1。这是最直接的激活显存压缩器。batch缩小后用梯度累积来补偿batch size下降带来的不稳定但梯度累积不省显存它只是等价放大batch。打开激活检查点。PyTorch自带接口非常方便model.gradient_checkpointing_enable()代价是大约20%-30%的训练速度损失但激活内存能从几十GB压到几GB划算得不能再划算。 3. 使用ZeRO或CPU offload。ZeRO阶段1把优化器状态切分到多卡阶段2把梯度也切分阶段3把参数也切分。如果只有单卡可以用DeepSpeed的offload_optimizer把优化器状态放到CPU内存。 4. 改用LoRA/QLoRA。这是“显存实在不够”后的终极方案只训练极小一部分低秩适配器主模型冻结或量化7B模型能压到一张消费级显卡。记住梯度累积不省显存能不写就别为了“显存”去写它真正需要它时它的意义是大batch不是省内存。4.3 常见问题速查表现象可能原因处理建议训练到一半OOM激活缓存积累/峰值过高减小batch、开激活检查点、扩容缓存配置OOM但memory_allocated并不高CUDA缓存碎片设置expandable_segments:True或max_split_size_mbloss变成NaN/Inf学习率过大、梯度爆炸降低lr、增加梯度裁剪、检查数据中NaNGradScaler的scale持续下降梯度确实有inf先查lr再查模型结构是否数值不稳定不要只盯loss scaling用了AMP后显存没降参数仍是FP32AMP只省中间张量改用完整半精度方案或用LoRA训练速度反而更慢频繁cast/小算子太多/CPU瓶颈增大batch减少张量转精度次数检查GPU利用率BF16训练收敛明显变差尾数精度不足换FP16scaler或提高模型宽度而不是深度4.4 容易被忽略的“收尾”细节训练结束保存模型时如果用的是FP16 FP32主权重方案不要只保存model.state_dict()否则下次加载的可能是FP16权重直接拿去推理精度可能打折扣。正确做法是保存训练用的主权重或者在保存前用主权重覆盖模型并转回所需精度。PyTorch AMP路线下模型参数本来就在FP32没这个问题但完整半精度方案必须注意。另外我在实际训练中习惯每100步打印一次当前的scale、loss和最大梯度范数。如果scale一直在涨说明训练稳定如果scale掉到16以下且频繁触发跳过step那基本可以断定是学习率太大或者数据出了问题。这比看loss曲线更早暴露危险。最后分享一个经验显存估算这件事我踩过最大的坑是过分相信网上的显存占用截图。别人说7B全参数微调只要70G我照着配结果batch size、seq length一不一样都不知道最后浪费了半天调试。后来我学乖了先把固定开销用每参数16字节秒算出来再根据激活公式估算一个数量级然后直接跑一个batch1、seq短的小实验用memory_allocated验证误差控制在10GB以内后才去正式调参。这套流程虽然不精致但从来不会让我在设备规划上翻车。混合精度也一样它救不了显存规划错误它的价值是让TensorCore干活让中间张量减半让训练速度变快。真正想要把70B模型塞进单卡还得靠LoRA、ZeRO和重计算这些“开源节流”的组合拳。希望这篇能让你下次新开训练任务时第一件事不再是焦虑而是冷静算一遍这张卡到底装不装得下。
大模型训练显存估算与混合精度调优实战
干过大模型训练的人都清楚“显存不够”往往是压垮训练任务的第一根稻草。不管是7B还是13B模型开工前如果不先算清楚显存预算等跑到一半OOM再回头调batch size、重开任务损失的不只是时间还有心态。这篇是这个系列的第三篇专门讲两件事怎么手算模型训练要吃多少显存以及混合精度训练里那些经常被忽略的技术细节。适合正要准备全参数微调、在选卡或者规划序列长度的朋友读完你至少能回答两个问题这条训练管线需要几张80G卡混合精度到底帮我省了哪部分显存显存估计这件事很多人习惯直接抄别人给的“7B大概要多少G”但别人的batch size、序列长度、是否开重计算都和你不一样抄来必翻车。不如花五分钟把手里的模型配置代进公式得到一个误差可接受的量级。混合精度则更微妙它确实能把训练跑得更快、中间张量更小但如果你以为用了AMP就万事大吉结果显存没降多少、loss还时不时NaN那大概率是没理解FP16/BF16的底牌。这篇就把账一五一十算给你看再附上实操排坑清单。1. 显存都去哪儿了训练时的四类开销1.1 模型参数本身最直白的一笔模型参数有多少个就有多少权重数值占显存。但同样一组数字用FP32存和用FP16存体积差一半。FP32每个参数占4字节FP16/BF16每个参数占2字节。7B参数模型如果全用FP32存储光权重就是 7e9 × 4B ≈ 28GBFP16则约14GB。所以“这个模型多少个B”只是起点真正的显存还要看存储精度和训练方式。1.2 梯度反向传播的“路程记录”训练必须做反向传播每个参数都会对应一个梯度用于指导参数往哪个方向更新。梯度张量的大小和参数一样大因此参数如果是FP32梯度FP32就是每参数4字节如果参数用FP16存储梯度通常也会被算成FP16每参数2字节。这里多说一句在PyTorch标准的torch.cuda.amp里如果你的模型参数保持FP32反向传出来的梯度也会是FP32显存占用并不会因为这个模式而下降。很多新手在这里有误解以为用了AMP参数内存就减半其实不是。1.3 优化器状态真正的大头不是参数以最常用的Adam/AdamW为例它要给每个参数额外保存两个状态一阶动量momentum和二阶动量variance这两个都按FP32存储每个状态4字节。另外如果我们把模型参数转成FP16训练为了不让更新过程因为精度太低而发散还需要单独保留一份FP32的“主权重”副本又是4字节。算下来混合精度训练中每个参数的固定开销可拆成FP16参数副本2字节FP16梯度2字节FP32主权重4字节Adam一阶动量4字节Adam二阶动量4字节合计16字节/参数看到没即使模型FP16存储更新时还得背一份FP32主权重和两个FP32状态每个参数依然要16字节。这一点常被营销话术掩盖混合精度训练并没有让Adam优化器下的“固定显存”减半它真正省下的是中间激活张量以及通过降低计算精度换取的算力优势。用纯FP32训练参数4字节梯度4字节两个状态8字节也是16字节。两者打平。不同优化器和精度组合的每参数固定字节数可以直接查下面这张表训练配置参数梯度优化器状态每参数总字节FP32 Adam4B4B8B16BFP16/FP32主权重 Adam2B2B12B16BFP16无主权重 Adam2B2B8B12BBF16 AdamW ZeRO offload2B2B状态可被切分/卸载不定最后一行意味着通过ZeRO等策略可以把优化器状态切到多卡甚至卸载到CPU这才能突破单卡固定开销的瓶颈。1.4 激活值与中间张量藏得最深的“刺客”前向传播时每一层都会产生中间结果比如Transformer里QKV投影输出、注意力分数、MLP中间层输出等。这些结果在反向传播计算梯度时还要再用一次所以不能算完就扔必须暂存在显存里。它们的大小和batch size、序列长度、注意力头数强相关而和参数量关系不大。很多模型看着参数不多却因为序列长、batch大显存直接爆掉罪魁祸首往往是激活值。这一块也是我们可以手动估算的核心部分下一节专门展开。除了上述四类CUDA context、cuDNN workspace、PyTorch缓存分配器的预留区也会占用一小部分显存通常几十MB到几GB不等但不会成为决定性因素估算时可以忽略实测时再通过监控去看。2. 显存估计用一张纸算出你的卡够不够2.1 固定开销先把参数规模装进去固定开销只依赖参数量P和训练配置和batch size、序列长度无关。计算公式非常简单[ 固定显存(GB) P \times bytes_per_param / 1024^3 ]P是参数个数bytes_per_param按上一节表格取。例如一个7B参数模型用FP16FP32主权重Adam训练每参数16字节固定开销就是[ 7 \times 10^9 \times 16 / 1024^3 \approx 112GB ]也就是说只看参数、梯度、优化器状态就已经需要1.5张80G卡了。还没算激活值和batch。这也是为什么很多人说全参数微调7B模型单卡80G根本放不下除非用LoRA、DeepSpeed ZeRO、CPU offload等手段把它拆走。2.2 激活值粗略估算一个能上手算的经验式严格精确计算激活显存需要逐层追踪张量生命周期工程上太累。这里给一个足够用于“拍脑袋”的近似式。假设我们要训练一个标准Transformer模型配置为层数L隐藏维H注意力头数Abatch size B序列长度S使用FP16/BF16存储中间张量每个Transformer层需要暂存的激活量大约为[ 2 \times (19 \times B \times S \times H 3 \times B \times A \times S^2) ]单位是字节。解释一下两部分来源前半部分包括层的输入、QKV投影结果、注意力上下文输出、MLP两层中间结果、残差连接需要保留的输入等大约折合19份B×S×H的二维矩阵后半部分是注意力分数矩阵的三种形态QK^T的结果、softmax之后的概率、dropout掩码或概率版每个都是B×A×S×S的大方块。因为按2字节精度的FP16存储所以前面乘2。这个式子不追求精确但能反映出两个关键趋势序列长度S影响极大因为注意力矩阵是S²增长batch size和隐藏维是线性增长。如果开启了激活检查点gradient checkpointing代价是反向传播时重算一部分前向激活显存可以大幅降低估算时大约只需要保留“每层输入激活”和一层重算所需的中间量整体量级大约是一个layer的激活量可以按上式“乘以1”而不是“乘以L”来近似。2.3 一个7B模型的计算实例我们用常见配置7B模型L32H4096A32序列长度S2048。先算batch size1时每一层激活[ 19 \times 1 \times 2048 \times 4096 159,383,552 ] [ 3 \times 1 \times 32 \times 2048^2 402,653,184 ]两部分相加约5.62亿个元素乘2字节得到约1.12GB。这是单层激活32层都不开检查点则约36GB。固定开销112GB 激活36GB ≈ 148GB两张80G卡能放下但比较紧张。如果不开激活检查点再把batch size提到4激活变成144GB总显存约256GB两张80G就不够了至少需要4张而且这只是压着内存线跑实际情况还要看PyTorch缓存碎片建议再多留余量。配置固定开销激活开销(估算)总显存7BB1S2048无检查点112GB36GB~148GB7BB1S2048开检查点112GB约4GB~116GB7BB4S2048无检查点112GB144GB~256GB13BB1S2048开检查点约208GB约5GB~213GB13B的固定开销13e9×16B≈208GB所以要跑13B全参训练4卡80G是起步5卡更稳。这也是为什么现在大家更青睐LoRALoRA只训练少量新增参数主模型冻结后可省掉优化器状态和梯度的巨额开销大大降低门槛。2.4 从估算到代码用memory stats验证纸面估算完成后一定要让训练代码“自报家门”。PyTorch的CUDA caching allocator提供的数字比nvidia-smi更接近真实逻辑使用量因为显卡驱动显示的显存可能包含预留而未用的部分。在训练循环里打印import torch def print_gpu_memory(): allocated torch.cuda.memory_allocated() / 1024**3 reserved torch.cuda.memory_reserved() / 1024**3 print(fallocated: {allocated:.2f}GB | reserved: {reserved:.2f}GB)每次step结束时调用一次观察allocated的峰值和reserved的差值。如果allocated离80G很远就开始OOM多半是缓存碎片问题。这时可以设置环境变量export PYTORCH_CUDA_ALLOC_CONFexpandable_segments:True让它利用可扩展内存段通常能显著缓解碎片导致的假OOM。另一种方案是max_split_size_mb调小但实际体验上expandable_segments更有效。3. 混合精度训练它做什么不做什么3.1 FP16和BF16一个怕下溢一个怕精度细FP16用1位符号、5位指数、10位尾数能表示的数值范围很小最大约65504最小的normal值约6e-5。它的问题不是“存不下大数”而是“存不了特别小的数”一旦某个中间结果小于约1e-5再往下就变成0这叫下溢。梯度在深层网络中经常会出现非常小的数值正是FP16的噩梦。BF16则用了和FP32相同的8位指数最大值也是FP32同量级范围宽阔得多几乎不会出现下溢但它只有7位尾数精度比FP16更粗。打个比方FP16像一把刻度更细但量程有限的尺子BF16像一把量程很广但刻度很粗的尺子。大模型训练里梯度数值范围的最主要矛盾往往是小数下溢所以新一代GPU比如A100、H100、RTX 30系以后上BF16往往比FP16更稳不需要额外做loss scaling也能跑。反过来早期的V100、T4这些没有BF16硬件加速的卡就只能用FP16 loss scaling。3.2 Loss Scaling为什么损失放大一倍不是玄学既然FP16存不了太小的数那操作办法很朴素既然梯度在进入FP16存储前变得太小那就先把loss乘一个大因子例如1024。因为链式法则反向传播中所有梯度都会跟着“放大同样倍数”原本可能小于FP16下限的值被搬进可表示范围。等真正更新参数前再把梯度除以这个因子恢复真实梯度。这个“先乘后除”的过程就是loss scaling。手动调一个固定scale也行但不保险因为不同训练阶段的梯度量级差很多。于是就有了动态策略每隔一段时间如果梯度没出现inf/NaN就把scale往大调一旦发现梯度中有inf/NaN说明放大过头了就回退半分并跳过这次优化器更新。PyTorch的GradScaler封装的就是这套逻辑。BF16下通常不需要它因为指数范围足够大但如果训练早期loss突然发散也可以临时罩一层来兜底代价不大。3.3 两种混合精度“流派”torch AMP与完整半精度这是最容易被混在一起说的地方必须讲清楚。第一种是PyTorch原生AMP典型写法是torch.autocastGradScaler。它做的事是在算子层面动态选择精度比如矩阵乘法、卷积这类计算密集且适合TensorCore的操作自动把输入转成FP16去算LayerNorm、Softmax这类对精度敏感的操作仍保持FP32。模型参数本身不会整体转成FP16优化器更新依然在FP32参数上做因此不需要额外维护FP32主权重。它的价值是代码改动极小、速度提升明显、中间张量也是FP16从而省内存。但注意它并没有省下“参数梯度优化器状态”那部分固定内存。第二种是“完整半精度训练”典型代表是DeepSpeed和部分手动实现。先调用model.half()或者model.to(dtypetorch.bfloat16)把模型参数整体转成低精度存储再额外保留一份FP32主权重给优化器用。这时训练管线里的权重和中间激活大部分都是低精度整体显存占用会更接近我们第一节表里的“FP16参数2B”那一行。它的固定开销依然是每参数16字节但因为参数本体是2字节而不是4字节通信和某些矩阵乘法的内存带宽压力会小一些。所以一句话总结真正省显存的是“削减优化器状态/激活”而不是把参数从FP32换成FP16这件事本身。如果你的目标是塞进一张卡最有效的手段是LoRA、ZeRO、激活检查点而不是单纯打开AMP。3.4 什么时候用FP16什么时候用BF16我自己的选择标准很简单如果卡是A100/H100/4090/3090这类支持BF16的优先BF16因为省事不需要跟loss scale斗智斗勇对学习率也不那么敏感。如果只有V100/T4/2080Ti这类不支持BF16加速或支持的生态不完善的卡就用FP16 GradScaler。如果模型本身有非常多的Softmax、LayerNorm、注意力打分这类对精度敏感的结构BF16偶尔会在小模型上出现精度损失可以先跑10个step对比FP16和BF16的loss曲线再决定用哪个。如果发现GradScaler的scale一路降到16甚至更低而梯度还在出现inf那大概率不是scale的问题而是梯度爆炸。先看学习率再看梯度裁剪最后看数据里是不是混入了异常样本。3.5 混合精度训练的关键流程无论用哪种流派优化器更新前都有一个必须完成的顺序前向时启用autocastloss得到FP16/BF16计算下的输出。将loss乘上当前scale再backward梯度被放大。反向后先unscale梯度即除以scale并趁机检查是否有inf/NaN。如果梯度正常做梯度裁剪然后调用optimizer.step()更新参数。最后通过scaler.update()调整下次的scale。这里特别容易踩的坑是梯度裁剪的时机。很多人习惯直接调用torch.nn.utils.clip_grad_norm_但如果你用GradScaler梯度此时还是放大后的状态必须先scaler.unscale_(optimizer)再裁剪。顺序错了裁剪阈值等于形同虚设。4. 实战配置、问题速查与显存兜底方案4.1 PyTorch混合精度训练循环最小可跑代码给一段简洁但完整的最小示例。以FP16 GradScaler为例import torch import torch.nn as nn model MyTransformer() optimizer torch.optim.AdamW(model.parameters(), lr1e-5) scaler torch.amp.GradScaler(cuda) # PyTorch 2.x 新写法 # 旧写法torch.cuda.amp.GradScaler() for batch in dataloader: optimizer.zero_grad() with torch.autocast(device_typecuda, dtypetorch.float16): loss model(batch) scaler.scale(loss).backward() # 关键先unscale再clip scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) scaler.step(optimizer) scaler.update()如果换成BF16可以用with torch.autocast(device_typecuda, dtypetorch.bfloat16): loss model(batch)BF16通常不需要GradScaler但unscale_和step的代码结构仍然建议保留因为万一训练中途出现NaN想启用scale时改起来只差一行。另外注意scaler.unscale_在同一个optimizer上只能调用一次重复调用会报错如果你没有调用scaler.unscale_clip_grad_norm_之前也要保证梯度已经是真实值。把上面的流程当成标准模板存下来能少踩很多坑。4.2 显存不够时的“救命”优先级如果真的OOM了按下面这个顺序调整性价比从高到低缩小batch size到1。这是最直接的激活显存压缩器。batch缩小后用梯度累积来补偿batch size下降带来的不稳定但梯度累积不省显存它只是等价放大batch。打开激活检查点。PyTorch自带接口非常方便model.gradient_checkpointing_enable()代价是大约20%-30%的训练速度损失但激活内存能从几十GB压到几GB划算得不能再划算。 3. 使用ZeRO或CPU offload。ZeRO阶段1把优化器状态切分到多卡阶段2把梯度也切分阶段3把参数也切分。如果只有单卡可以用DeepSpeed的offload_optimizer把优化器状态放到CPU内存。 4. 改用LoRA/QLoRA。这是“显存实在不够”后的终极方案只训练极小一部分低秩适配器主模型冻结或量化7B模型能压到一张消费级显卡。记住梯度累积不省显存能不写就别为了“显存”去写它真正需要它时它的意义是大batch不是省内存。4.3 常见问题速查表现象可能原因处理建议训练到一半OOM激活缓存积累/峰值过高减小batch、开激活检查点、扩容缓存配置OOM但memory_allocated并不高CUDA缓存碎片设置expandable_segments:True或max_split_size_mbloss变成NaN/Inf学习率过大、梯度爆炸降低lr、增加梯度裁剪、检查数据中NaNGradScaler的scale持续下降梯度确实有inf先查lr再查模型结构是否数值不稳定不要只盯loss scaling用了AMP后显存没降参数仍是FP32AMP只省中间张量改用完整半精度方案或用LoRA训练速度反而更慢频繁cast/小算子太多/CPU瓶颈增大batch减少张量转精度次数检查GPU利用率BF16训练收敛明显变差尾数精度不足换FP16scaler或提高模型宽度而不是深度4.4 容易被忽略的“收尾”细节训练结束保存模型时如果用的是FP16 FP32主权重方案不要只保存model.state_dict()否则下次加载的可能是FP16权重直接拿去推理精度可能打折扣。正确做法是保存训练用的主权重或者在保存前用主权重覆盖模型并转回所需精度。PyTorch AMP路线下模型参数本来就在FP32没这个问题但完整半精度方案必须注意。另外我在实际训练中习惯每100步打印一次当前的scale、loss和最大梯度范数。如果scale一直在涨说明训练稳定如果scale掉到16以下且频繁触发跳过step那基本可以断定是学习率太大或者数据出了问题。这比看loss曲线更早暴露危险。最后分享一个经验显存估算这件事我踩过最大的坑是过分相信网上的显存占用截图。别人说7B全参数微调只要70G我照着配结果batch size、seq length一不一样都不知道最后浪费了半天调试。后来我学乖了先把固定开销用每参数16字节秒算出来再根据激活公式估算一个数量级然后直接跑一个batch1、seq短的小实验用memory_allocated验证误差控制在10GB以内后才去正式调参。这套流程虽然不精致但从来不会让我在设备规划上翻车。混合精度也一样它救不了显存规划错误它的价值是让TensorCore干活让中间张量减半让训练速度变快。真正想要把70B模型塞进单卡还得靠LoRA、ZeRO和重计算这些“开源节流”的组合拳。希望这篇能让你下次新开训练任务时第一件事不再是焦虑而是冷静算一遍这张卡到底装不装得下。
