做 LLM 训练调优这几年我最大的感觉是很多人一碰到显存不足就急着改代码、调参数甚至直接换大卡却很少有人先做一件事——把显存占用“测”清楚。这篇是“大模型显存优化篇”的 task3核心就两个字测量加上一个决策。具体说就是在训练侧把参数、梯度、优化器状态、激活值和临时缓冲区这五块显存的真实占用量化出来然后根据测量结果去排优先级、定方案该开 activation checkpointing 还是该换 LoRA该调微批大小还是该上 FSDP。这篇文章适合正在做大模型微调、全参数训练或者在 80G 卡上为 OOM 折腾的人。我把自己的测量脚本、决策步骤和踩过的坑都整理出来了可以直接照着做。1. 训练侧显存到底花在哪先搞清楚五块大头1.1 显存组成不只是模型参数很多人以为显存只装模型权重这是最大的误区。训练一个 7B 模型和推理一个 7B 模型显存需求完全不同。推理时你老老实实做 forward主要就是权重 KV cache训练时你不仅要跑 forward还要把每一层的中间输出存下来拿去做反向传播所以显存里同时住着五类东西模型参数、梯度、优化器状态、激活值中间输出、以及 PyTorch 缓存分配器临时占用的 buffer。我习惯把它们分成“静态大头”和“动态大头”。静态大头包括参数、梯度和优化器状态这三块的数量级基本可以靠参数量直接算出来动态大头主要是激活值它会随着 micro-batch size、序列长度、层数的变化剧烈波动。临时 buffer 一般占比小但碎片化时可能虚高到十几 GB。所以当你看到 nvidia-smi 里占用 70Gtorch.cuda.memory_allocated() 却只报 50G 时别慌那多出来的部分很可能就是 reserved 缓存没还回去。1.2 不同优化器对显存的影响优化器状态是训练侧最容易被低估的部分。用 AdamW 全量更新 7B 模型时每个参数要有两份 fp32 的额外状态一阶动量 m 和二阶动量 v每个状态 4 字节合起来就是 8 字节每参数。再加上模型参数本身 2 字节bf16/fp16梯度 2 字节算下来一个 7B 模型全参数训练的理论最下占用大概是7B × (2 2 8) 84GB。注意这只是静态部分还没算激活值也就是说单张 80G 卡根本放不下。看到这个数字你就明白为什么大模型微调实战中大家普遍优先选择 LoRA、QLoRA或者 FSDP/ZeRO 做分片。因为优化器状态是大头而 LoRA 只训练新增的少量低秩矩阵冻结原模型参数待更新的参数量可能只有 0.1% 到 1%优化器状态自然就小到可以忽略。换个角度如果你坚持全参数微调那要么上多卡并行把静态部分摊到多张卡上要么用 Adafactor 这类省状态的优化器Adafactor 把 Adam 的两个状态合并成按行和按列的二次矩统计显存占用接近 SGD。1.3 训练前先算理论基线我每次做训练侧测量前的第一步不是抓脚本跑数据而是在纸上先算一条理论基线。公式很简单模型参数显存 参数量 × 参数精度字节数。bf16 算 2fp32 算 4。梯度显存 参数量 × 2通常与参数同精度。优化器状态显存 参数量 × 状态字节数。AdamW fp32 是 8SGD 是 4Adafactor 约 2。激活值显存 ≈ micro-batch size × 序列长度 × hidden size × 层数 × 若干常数。这个常数一般取 8 到 16取决于是否用重计算、attention 实现等在线性层和 attention 层之间有差异。比如 7B 模型用 bf16 全量 AdamW 微调基线就是上面说的 84GB 静态占用。用 bf16 LoRA 微调静态部分约 14GB参数 0.1GBLoRA 参数和优化器状态剩下的预算全都可以分给激活值。有了这个基数你再去跑测量脚本看到某个数字不对心里就有谱了。2. 动手测量显存工具、Hook 与实验设计2.1 用 PyTorch 自带工具快速定位峰值说实话PyTorch 自带的显存接口已经够用不需要一开始就上 nvidia-smi 或 CUDA Events。我最常用的是这一组组合torch.cuda.reset_peak_memory_stats(device)把峰值的记录清空一般在训练循环开始前调用。torch.cuda.memory_allocated(device)当前实际已分配的张量内存单位字节。torch.cuda.max_memory_allocated(device)自上次 reset 以来的峰值分配内存。torch.cuda.memory_reserved(device)缓存分配器实际向显卡申请的内存也就是说它包括了上面那些没还给 CUDACachingAllocator 的部分。实际操作中我会在 step 的前后各打印一次这三个值再结合 nvidia-smi 的占用对比。如果memory_allocated远小于 nvidia-smi 里的进程占用说明大部分显存是被 PyTorch 缓存吃掉了这类情况通常不是训练张量本身不够而是峰值申请后缓存不释放后续内存碎片化后无法高效复用。这个坑我在第 4 部分会展开讲。新版 PyTorch 里torch.cuda.memory_summary()仍然可用但更推荐直接读torch.cuda.memory_stats()返回的字典里面的allocated_bytes.all.peak、reserved_bytes.all.peak这些字段更方便程序化处理。还有一个技巧在跑大 batch 出现 OOM 后把PYTORCH_CUDA_ALLOC_CONF设成max_split_size_mb512或expandable_segments:True有时候没改代码就把 OOM 解决了。2.2 逐层记录激活值把显存“大头”抓出来静态部分的显存可以用公式算出来但激活值这种和输入、batch、序列长都强相关的动态量最好用 Hook 实测。我的做法是给模型的每个子模块注册一个 forward hook记录模块输出张量的numel()乘以element_size()这样就能得到每层 forward 输出的大小。实际操作时注意三点第一不要统计model.eval()下torch.no_grad()的 profiler那是推理场景训练侧要在正常的 forward/backward 环境里测。第二不需要每步都测用一个固定的 dummy input 跑一个前向把各层输出内存快照存下来基本就能看出激活值分布。第三attention 类模块的输出可能有 tuple需要判断类型再取torch.Tensor。我还会用torch.profiler.execution_graph或简单打印torch.cuda.memory_snapshot()的堆栈信息去定位是哪个 op 分配了最大 block。这个方法一开始看起来繁琐但对于找出“为什么同一个 batch 大小别人的显存比你少 10G”这类问题是唯一的解题路径。2.3 控制变量怎么设计一次有效的测量实验测量很容易被数据分布、batch 大小、并行策略甚至随机数顺序干扰所以一定要控制变量。我的实验设计模板是这样的固定全局 batch size比如 32固定训练步数比如 3 步只改变 micro-batch size。固定 micro-batch size固定序列长度只改变模型并行设置比如关闭/开启 activation checkpointing。固定其他条件单独改变混合精度策略bf16 vs fp16 vs fp32。记录每步的allocated峰值、reserved峰值、耗时以及是否 OOM。有个关键点训练时显存峰值通常出现在第一步 backprop 结束时因为这时候计算图还没释放优化器状态也全部建立起来了。所以不要拿单步 forward 的峰值下结论一定要跑到第 2-3 步让 PyTorch 的惰性初始化全部触发。我实测过第一次 forward 后allocated可能只有最终峰值的一半优化器参数的二值 mask、分片索引等都是延迟到具体更新时才分配。2.4 一个可复用的最小测量脚本下面这段代码是我在项目里常用的最小测量框架它能在 5 分钟内拿到训练侧显存分布import torch import torch.nn as nn def measure_peak_memory(model, data_loader, devicecuda:0, steps3): model model.to(device) optimizer torch.optim.AdamW(model.parameters(), lr1e-5) # 把记录清零 torch.cuda.reset_peak_memory_stats(device) for i, batch in enumerate(data_loader): if i steps: break input_ids batch[input_ids].to(device) labels batch[labels].to(device) outputs model(input_idsinput_ids, labelslabels) loss outputs.loss optimizer.zero_grad() loss.backward() optimizer.step() print(fstep {i}: allocated{torch.cuda.memory_allocated(device)/2**30:.2f} GB, fpeak{torch.cuda.max_memory_allocated(device)/2**30:.2f} GB, freserved{torch.cuda.memory_reserved(device)/2**30:.2f} GB) # 详细汇总 print(torch.cuda.memory_summary(devicedevice, abbreviatedTrue))如果要测逐层激活值在训练循环前插入 hook具体参考上一小节。注意必须在loss.backward()之后再取峰值因为 backward 过程中还会创建一些临时矩阵这些矩阵虽然不跨 step 存活但会让max_memory_allocated跳得非常高。3. 显存预算决策把 80G 卡花在刀刃上3.1 预算框架先定上限再列大头测量做完下一步就是做显存预算。我喜欢把这个问题当成“家庭月度开销”先确认一个月的总盘子GPU 总显存再给固定支出参数、梯度、优化器状态、弹性支出激活值、临时 buffer、机动储蓄预留 headroom各自分配比例。总盘子的上限一般取 85% 到 90% 的实际卡显存。比如 80G 的 A100/H100我建议预算上限定在 72G 左右剩下 8G 留给 CUDA context、cuDNN workspace、通信缓冲区以及突发峰值。如果把它压到 75G 以上训练过程只要稍微遇到较长序列或通信量波动就有概率随机 OOM。接下来把测量得到的数据填入预算表。典型 7B 模型 LoRA 微调场景可能是这样的模型参数冻结bf1614GBLoRA 可训练参数 优化器状态不到 0.5GB激活值seq_len2048micro_batch4开启重计算6GB临时 buffer 和通信缓冲2GB合计22.5GB远低于 72G 预算所以还能继续加大 batch 或序列长度。同样是 7B 模型全参数 AdamW 微调就完全不同了。静态部分 84GB 直接超预算所以决策必须是“换状态”要么用 FSDP/ZeRO-3 把静态部分分片到 2 张卡每卡分到约 42GB要么用 LoRA要么把优化器状态 offload 到 CPU 内存每卡静态降至 16GB 左右但会牺牲一部分吞吐。这个取舍就是“预算决策”的本质显存不够时不是在压缩 batch 和序列长度上死磕而是优先处理最胖的那块静态占用。3.2 针对不同训练模式的具体决策我在实际项目里总结出了一个简化决策表适合大多数 7B-70B 级模型训练模式静态占用情况推荐预算策略全量微调 AdamW极高约 12 字节/参数优先 FSDP/ZeRO-3 或多卡 TP/PP若单卡则建议 offload全量微调 Adafactor偏高约 6 字节/参数单卡可尝试配合激活重计算LoRA 微调bf16低约 2 字节/参数 适配器预留大头给激活值可放心加长序列QLoRA 4bit 量化极低约 0.5-0.8 字节/参数适合消费级显卡或超长上下文决策逻辑永远是先看静态部分有没有爆再看激活部分有没有超。静态爆了你做 activation checkpointing 救不了因为激活值只是动态部分静态该占的还是占。激活超了才有必要开重计算、降 micro-batch、减少序列长度。很多新手一 OOM 就开重计算开完发现没改善就是因为爆的是静态部分。另外并行策略也会改变预算。DDP 会给每个 rank 复制一份全量模型和梯度Adam 优化器状态也是全量所以静态占用不降FSDP/ZeRO 才真正把静态占用分摊到多卡。这意味着你在 DDP 模式下把 batch 调得再小7B 全量微调照样放不进 80G 卡。这一点在团队里很少有人明说但它是预算决策的分水岭。3.3 从测量到决策的最小流程我建议把整个“测量-决策”流程固定成四个步骤方便每个人照着做跑一次带reset_peak_memory_stats的最小训练拿到 allocated 峰值和 reserved 峰值。用公式或表格拆出静态占用参数梯度优化器状态剩余部分大致归因到激活值。对比总预算比如 72G。若静态超预算先改训练模式或并行策略不要动 batch若静态在预算内但总占用超再考虑压缩动态部分。每次只改一个变量比如先开 activation checkpointing再调 micro-batch size记录峰值变化最后确定一手不爆、性能能接受的配置。第 4 步要特别注意开 activation checkpointing 会让前向多算一遍token 吞吐会下降但显存峰值可能从 70G 降到 40G。这时候要判断值不值。我个人经验是如果显存峰值在预算上限附近先开 checkpointing如果还有 20G 以上的余量优先加 batch因为模型带宽利用率往往比省显存更影响训练效率。3.4 7B 模型在 A100-80G 上的预算示例直接给一个可以复用的算例。假设 7B 模型bf16 权重LoRA 微调目标 seq_len4096A100-80G 单卡总预算 72G。静态模型权重 14GLoRA 可训练层通常只占全部参数的 0.5%约 35M 参数AdamW 状态约 0.28G合下来算 14.5G。激活值7B 模型 hidden size 一般是 409632 层多头注意。粗略估算 1 张卡 micro-batch1 且关闭重计算时激活值可到 8-12Gmicro-batch2 会翻一倍。这里按开启 activation checkpointing 后的 6-8G 估。通信/缓冲2-3G。合计约 23-26G剩余预算充足因此可以把 micro-batch 提到 4 甚至 8或者把序列长度推到 8192而不用动 FSDP。如果是全量 AdamW 微调静态 84G 超出总预算所以只能 FSDP 双卡。分片后每卡静态约 42G 激活 8G 50G还在 72G 预算内。这时候再考虑要不要开 checkpointing就属于锦上添花而不是救火了。4. 真实案例一次 OOM 排查与预算复盘4.1 现场与报错解析之前我帮人排查过一起 7B 全量微调 OOM 问题。卡是 80G A100代码用的是 DDP AdamW bf16micro-batch 设为 16seq_len 2048。训练到第一步就报CUDA out of memory错误提示只有一行没有堆栈。很多人到这一步直接去搜“如何增大 PyTorch 显存”或者把 batch 降到 8但那样解决不了本质问题。我先让训练脚本在 OOM 前打印torch.cuda.memory_summary(abbreviatedTrue)同时用PYTORCH_CUDA_ALLOC_CONFexpandable_segments:True再跑一次。结果发现allocated峰值约 72Greserved却到了 80G也就是说实际分配张量只差一点就顶到物理上限。进一步分析72G 里静态部分参数梯度Adam 状态约 84G按公式算确实静态就 84G这说明 DDP 下每卡都是全量物理上注定放不进 80G 卡。所以问题不在 batch 大小而在并行策略选错了。4.2 测量发现真正瓶颈用代码把静态部分和动态部分分离。7B 参数占了 14G bf16梯度也占 14GAdamW 优化器状态 56G合计 84G。看到这个数字dummy 都能判断出哪怕 micro-batch1用 DDP 全量 AdamW 也照样 OOM。而 72G 这个测量值之所以小于 84G是因为 PyTorch 在 OOM 前还没完成第一轮优化器 stepAdam 状态还没来得及全部建立所以实际测量到的就不是最终峰值。这给我们一个教训训练侧测量不能只测崩溃前的快照要在心里给优化器状态留出“延迟初始化”的余量。我当时让代码直接构造一个 Adam 优化器并对所有参数 group 做一次optimizer.step()随机梯度也算强行触发状态初始化再测峰值这样才拿到真正的稳态占用。这一步对预算决策极其重要否则你会低估静态占用误判成激活值膨胀。4.3 最终决策与优化效果测量完成后决策变得很清晰放弃 DDP切到 FSDP并开启优化器状态 offload。因为静态 84G 即便 FSDP 分到 2 卡每卡约 42G加上激活 10G基本能塞进 80G。最后实际配置是2 卡 FSDPCPU offload优化器状态activation checkpointing开启micro-batch 从 16 降到 8全局 batch 保持不变。这样每卡峰值降到 30G 左右显存余量充足吞吐只损失不到 20%。这个案例中如果只把 micro-batch 降到 1 或者换 Adafactor也能救但效率不一样。Adafactor 静态占用大幅减少但收敛稳定性通常比 Adam 差FSDP 保持 AdamW 的大模型微调实战惯例收敛质量不降。所以决策时要综合考虑训练质量、吞吐和改造成本不要只看显存数字。4.4 reserved vs allocated显存缓存引发的虚高上面案例里还有一个细节reserved比allocated高不少。PyTorch 默认的 CUDACachingAllocator 会把释放的小块内存留在自己手中不立刻还给 GPU导致nvidia-smi看起来占用很高但程序内部allocated却不高。这种情况容易造成两种误判一种是以为自己的模型很吃显存其实是缓存碎片另一种是跑完大 batch 后立刻换小 batch结果依然 OOM因为缓存里可复用的块不够。遇到这类问题我的排查顺序是查看memory_reserved和memory_allocated的差值。看memory_stats里的num_free_blocks和active.Allocated.current。尝试设置PYTORCH_CUDA_ALLOC_CONFmax_split_size_mb:128或expandable_segments:True强制缓存分配器更激进地合并碎片。如果只是单次训练到结尾允许显存不归还直接忽略这个差值。如果你在实现长序列微调强烈建议用分段式数据加载器而不是一次性把超长文本塞进一个 batch这会直接减少峰值块的大小碎片化明显缓解。5. 常见问题速查表与实操心得5.1 常见问题速查表现象可能原因排查/解决方向训练第一步就 OOM静态占用超预算参数梯度优化器状态检查并行策略改用 FSDP/ZeRO 或 LoRA前几步正常十几步后 OOM优化器状态延迟初始化或中间计算图未释放多跑几步再测峰值打印每步max_memory_allocated显存峰值不高但整体卡顿缓存碎片化reserved虚高调PYTORCH_CUDA_ALLOC_CONF或减小 max_split_size_mb开启 checkpointing 后仍 OOM问题的瓶颈是静态部分而非激活值先算静态占用是否超预算若超预算checkpointing 救不了loss 正常但显存缓慢增长可能代码里意外地保留了 loss 或 logits 的引用用weakref或检查打印张量是否被保存到列表多卡训练时各卡 OOM 不一致数据不均匀或通信缓冲差异确认数据 sampler 是否给到相同 batch检查梯度同步时机5.2 几条省显存的实操心得第一尽量在 dataloader 里做 padding 而不是在训练 loop 里做。很多开源代码会先把 batch 中所有序列 pad 到最长序列再交给模型这样浪费的显存是“最长减最短”的部分。如果你控制好 bucket 策略比如按长度分组再 padding可以平白省下 20%-30% 的激活显存。第二关于 bf16 和 fp16不要只看占用。bf16 和 fp16 的张量都是 2 字节每参数但 fp16 在反向容易溢出需要 master weights 和 gradient scaling有的库还会额外汇出 optimizer 状态实际上显存开销反而更高。所以我做训练侧测量时凡是涉及到混合精度都会把“额外状态”算进预算而不是只看模型的 dtype。第三小心 norm 层和 attention mask。LayerNorm 全量重算的开销小但某些实现会把 weight 变成 fp32如果模型很大且每层都有 norm累加起来就是好几 GB。建议在测量后检查一下模型里 fp32 参数的比例如果超过 5%考虑消融测试是否对质量有实际影响。第四通信库的缓冲也需要预算。比如 NCCL 的临时 buffer 默认可能占用几百 MB 到几 GB在多卡训练时会在每张卡上重复分配。如果你用 FSDP 后显存依然吃紧可以尝试设置NCCL_BUFFSIZE或NCCL_P2P_LEVEL但不要随意调参最好先确认瓶颈是通信缓冲还是激活值。5.3 个人体会与后续扩展我在实际项目中做过的所有显存优化几乎没有一次是靠“玄学调参”成功的全部是先测量、再决策、后调参。测量侧的关键不是写多炫的脚本而是把静态和动态分开把 allocated 和 reserved 分开把真实 OOM 和缓存虚高分开。这三点想清楚显存预算决策就完成了大半。后面如果继续做 task4我大概率会往“训练吞吐与显存联合优化”的方向走显存只是资源的一部分显存和算力、带宽之间存在权衡好的决策应该在保证不 OOM 的前提下让总吞吐最高。比如同一个 batch 配比下是开大 batch 降 checkpointing还是小 batch 开 checkpointing最终要以 tokens/秒为准。这套方法论也可以反向用到推理侧的 KV cache 预算里。最后分享一个我每次都会用的小习惯训练脚本里永远在train_step的入口保留一行torch.cuda.reset_peak_memory_stats()出口打印max_memory_allocated。这样每次调参都能拿到一个可对比的显存曲线而不是靠目测 nvidia-smi。有了这个基线后续做任何优化心里都会特别踏实。