炼丹的人最怕看到的东西排名第一的不是loss突变成NaN而是那句“CUDA out of memory”。尤其是你明明准备好了数据、写完了训练脚本信心满满地启动结果第二次迭代直接炸掉。更难受的是有时候你只是把batch size调小了一点又跑起来了但你根本不知道刚才那一下到底是“谁”把显存吃满了。这个task3要解决的就是这个事训练侧的显存测量以及基于测量结果做出的预算决策。前两篇里我们把显存的基本构成和推理侧优化过了一遍但训练侧的账完全不是一回事——多出来的梯度、优化器状态、激活峰值每一项都可能比模型本身还要占地方。这篇我会直接给出可复用的测量代码、指标解读方式还有一套从“粗算”到“实测验证”的预算决策流程目标读者是正在做微调、准备尝试预训练、以及想在单卡上把显存榨干的各位。1. 训练侧显存开销的五张账单先搞清楚钱花在哪了很多人一OOM就开骂模型太大实际上有相当一部分情况是“优化器状态”和“激活值”这两个隐形大户把显存吃干净的。训练过程和推理最大的区别在于推理只需要把权重放在显存里然后让数据流过一遍就行而训练需要在每个step里保存大量中间状态用来反向传播更新参数。我把这些开销拆成五张账单顺序按照“长期驻留”到“临时浮动”来排。1.1 权重与梯度训练和推理的账本差异权重账单大家最熟悉。模型有多少参数就得有多少显存来装它。FP32下一参数占4字节BF16/FP16下占2字节。一个8B模型BF16精度裸权重就是16GB这还不算别的。梯度是训练侧新增的第一笔开销。反向传播时每个参数都要算出一个梯度形状和参数完全一致所以全参微调时梯度账单一上来就是2字节/参数。要注意的是梯度的生命周期很短它在反向传播过程中逐步生成等到optimizer.step()消费完之后就被释放。听上去很美好但偏偏在“backward结束到step之前”这一小段窗口里权重、梯度、优化器状态是同时存在的这就是峰值最常出现的地方。1.2 优化器状态真正的大头往往是它这账很多人第一次看会惊掉下巴。混合精度训练里AdamW优化器不但要保存一份FP32精度的主权重副本还要保存一阶动量m和二阶动量v全部是FP32。也就是说每个可训练参数要额外吃 4 4 4 12 字节。算一笔账8B模型全参微调光是优化器状态就要 8B × 12B 96GB。加上BF16权重16GB、梯度16GB静态开销就128GB了。所以别再问为什么8B模型在单张80G上也做不了全参微调这不是卡不行是Adam的账本太重。如果换成SGD优化器状态几乎可以忽略但为了训练收敛性大多数人还是得用Adam系优化器。这也是为什么后面预算决策时“换优化器”和“减少可训练参数”往往是性价比最高的两条路。1.3 激活值与临时内存动态开销的变数激活值指的是前向传播过程中每层计算出来的中间张量。反向传播要算梯度就得用到这些中间结果所以它们必须被保留下来。这一项有多大没法像权重那样用一个简单公式算死因为它跟batch size、序列长度、hidden size、层数、Transformer结构里有没有重计算都有关系。有一个可以做量级估计的粗算方式核心激活大约等于 batch × seq_len × hidden_size × layers × 精度字节数再乘一个经验系数2到4倍覆盖Attention矩阵、FFN中间层等。拿8B模型hidden 4096、32层、单条2048 token的序列举例BF16下光核心激活就是 1 × 2048 × 4096 × 32 × 2 ≈ 0.5GB乘上经验系数后可能到1到2GB。但如果batch size提到8、序列到4096这个数会快速涨到十几甚至几十GB直接变成比权重还大的开销。梯度检查点就是针对这一项的牺牲一部分计算时间不保存所有层激活反向时重新算一遍通常能省掉60%-80%的激活显存。1.4 CUDA上下文与缓存分配器的隐藏开销最后这账单很多人都没概念。只要你在PyTorch里第一次用CUDA驱动就会为进程建立CUDA context固定占用几百MB到1GB不等多卡、多进程时还会翻倍。另外PyTorch的缓存分配器会预先向驱动申请一大块显存作为“内存池”再按需切给你。这个池子里的内存在你眼里是“reserved”但不是“allocated”可从nvidia-smi看过去它已经实打实地占在卡上了。这五张账单可以汇总成一张每参数字节数表方便后面粗算。开销项混合精度全参微调BF16Adam8B模型参考值主要减负手段模型权重2B/参数16GB量化、冻结梯度2B/参数16GB冻结层、LoRA优化器状态Adam12B/参数96GB8bit Adam、Adafactor激活值动态数GB到数十GB梯度检查点、降batchCUDA上下文缓存池固定0.5-1GB少建进程、多卡共享2. 把显存测量做到位从打点到全链路分析知道显存花在哪了下一步就是真正把它测出来。训练侧的测量不是看一眼nvidia-smi就完事你需要三种不同的手段配合PyTorch内部统计接口负责看程序视角Torch Profiler负责看算子视角nvidia-smi负责看GPU全局视角。2.1 最直接的统计接口allocated/reserved/peak三件套PyTorch提供了一组专门的内存统计函数这几兄弟要先分清楚torch.cuda.memory_allocated()当前确实被张量对象占用的字节数这是“分配器视角下正在使用的内存”。torch.cuda.memory_reserved()缓存分配器向驱动申请到的总字节数包含碎片和未使用的block。用它来判断“这个进程实际占了多少显存”更准确。torch.cuda.max_memory_allocated()从程序开始或上一次reset以来allocated达到过的峰值。训练里最关心的就是它。我把这三兄弟封装成一个打点函数效果很好def print_mem(tag): gb 1024 ** 3 print( f{tag:24s} | allocated {torch.cuda.memory_allocated()/gb:6.2f}GB | freserved {torch.cuda.memory_reserved()/gb:6.2f}GB | fpeak {torch.cuda.max_memory_allocated()/gb:6.2f}GB )这段代码建议直接写进你的公共训练工具库后面所有实验都用得上。2.2 在训练循环里埋点一条日志看全生命周期测量要的是“分阶段”数据不是光看最终峰值。我会在关键节点各打一次点加载模型后、冻结参数后、创建优化器后、第一次forward后、第一次backward后、第一次optimizer.step后。每个节点的allocated与reserved差值能告诉你当前阶段有多少内存是“已经被张量占用的”有多少是“被分配器捏在手里的”。还有一个非常关键的动作在开始测峰值前先调用torch.cuda.reset_peak_memory_stats()。不然你看到的peak可能是之前某个残留操作留下的不是本次训练循环产生的。样例流程如下# 模型/优化器准备阶段 print_mem(model loaded) torch.cuda.reset_peak_memory_stats() # 训练循环只跑1-2个step for step, batch in enumerate(train_loader): print_mem(before forward) loss model(**batch) print_mem(after forward) loss.backward() print_mem(after backward) optimizer.step() print_mem(after optimizer.step) if step 1: break实测下来你会发现一个规律after backward和after forward这两个点的峰值差距往往非常大。尤其是全参微调时backward的一瞬间激活值和梯度同时存在这个窗口的allocated直接拉满。把日志对齐到线上线下的卡型基本上就能定位到“哪一步把显存掏空了”。2.3 想细到算子级别Torch Profiler的内存事件打点只能看到时间线的大节点如果还想知道到底是哪个算子分配了最多的显存那就得上torch.profiler。注意观察profile_memory参数PyTorch 2.x里可以通过它开启内存事件的记录具体用法from torch.profiler import profile, ProfilerActivity with profile( activities[ProfilerActivity.CUDA], record_shapesTrue, profile_memoryTrue, ) as prof: loss model(**batch) loss.backward() print(prof.key_averages().table(sort_byself_cuda_memory_usage, row_limit20))输出会按“自身分配显存”排序你一眼就能看到哪个操作最贪吃。我在微调Llama系模型时table里排前面的通常是FlashAttention的临时buffer和FFN中间层。这两个就是激活值的大头看到它们之后你对“要不要开Gradient Checkpointing”基本就有答案了。2.4 nvidia-smi外部采样的用法和局限nvidia-smi适合做GPU全局观测但它有两个先天不足必须要知道。第一采样频率不够密默认2秒一次很容易错过那种一闪而过的峰值第二它显示的是进程整体显存不是程序内部分配的数值而且多进程、多卡环境下的per-process统计有时还会不准。所以我的习惯是nvidia-smi 只看“趋势”PyTorch打点拿“精确值”。两个结合既能确认程序内部发生了什么又能确认卡上其他进程到底占了多少。外部采样时可以用一条bash命令做成轻量日志nvidia-smi --query-gpumemory.used --formatcsv -l 13. 测量数据的正确打开方式峰值、碎片与余量拿到一堆打点数据之后怎么解读才是最关键的。初学者容易犯的错是只记了一次max然后照着这个数决定要不要调batch size。这不够至少还要看懂三件事峰值出现在哪个阶段、碎片有没有制造假OOM、留多少余量才安全。3.1 判断峰值出现的时刻为什么第一轮迭代最准训练循环里第一个step会触发非常多的“一次性分配”CUDA context、模型参数加载、优化器状态、缓存分配器的初始池、甚至cuDNN的benchmark扫描。所以第一个step的峰值往往是最真实的“冷启动峰值”后面几个step因为分配器已经开始复用内存峰值多半不会再上涨。这意味着测峰值要在第一个step之前reset第一个step结束后立刻读max不要等到跑了几百个step再去看。如果你是在一个已经训练了一半的进程里做测量还得小心分配器的状态已经被前面的迭代污染这时候测出来的不是干净值。顺带提一句开启梯度检查点之后峰值的位置会移动。因为前向传播不再保存全量激活反向传播时会重新执行一次前向来算激活所以峰值大概率出现在“反向重计算”这段而不是最初forward末尾。3.2 碎片化问题allocated没爆但OOM了一个非常迷惑的场景我用memory_allocated()看明明只用了18GB24GB的卡怎么还OOM答案通常是碎片化。PyTorch缓存分配器从驱动拿到的是一大块连续显存当你反复创建各种形状的张量时这些block会越切越碎等到你需要一个大连续张量比如一个超大激活矩阵时找不到连续空间就报OOM了。判断方法很直接看memory_reserved()。如果reserved已经接近24GB而allocated还不到20GB说明大量内存被切成小块碎片躺在池子里。缓解手段有把PYTORCH_CUDA_ALLOC_CONFmax_split_size_mb:128设小一点减少大块被拆碎的概率。尽量固定batch size和序列长度避免训练中频繁出现形状暴涨的张量。在关键变换比如从embedding输出到attention输入时用contiguous保持一致性减少额外的临时拷贝。实在不行PyTorch 2.1 还可以试试PYTORCH_CUDA_ALLOC_CONFbackend:cudaMallocAsync用异步分配器缓解碎片但某些老卡和旧驱动不一定支持。3.3 余量设计可用显存不等于显卡标称显存标称24GB的卡实际能给你训练用的远没有24GB。驱动和CUDA context一层层吃下来真正可用的显存通常要打个95折如果再开图形界面、浏览器、其他进程还要再往下压。经验做法是按“标称值减1GB再到1.5GB”来算可用容量预算表里再留出约5%-10%的余量给临时分配抖动。打个比方24G卡我就按22.5G可用来规划80G卡按77G左右来规划。如果预算计算出来贴着24G的边我一定会认为它放不下因为训练过程中还有DataLoader的pin memory、通信库的buffer、偶尔的碎片波动这些变量。宁可把batch调小点用梯度累积补也不要赌那一两个GB。4. 预算决策从粗算到最优配置的完整路径测量是手段预算决策才是目的。所谓预算决策就是在训练前用一张纸和几行代码判断这张卡、这个模型、这个配置到底跑不跑得动以及跑不动时先砍哪块。4.1 静态开销粗算参数规模 × 单位字节先算静态开销就是权重、梯度、优化器状态这三种“每参数开销”的总和。训练方案每参数开销字节8B模型静态总额全参微调BF16 AdamW2 2 12 16约128GB全参微调BF16 Adafactor2 2 4 ≈ 8约64GB全参微调BF16 8bit Adam2 2 4 8约64GBLoRA微调BF16冻结主权重基础权重2B/参数 极小优化器开销约16GB起步QLoRA4bit量化主权重基础权重0.5B/参数 极小优化器开销约4GB起步注意全参微调那一行128GB意味着什么单张80G、甚至双卡80G想做传统意义的全参微调都够呛所以业界的8B微调基本默认走LoRA路线。这行粗算直接决定了你等下选择的训练范式。4.2 决策减配顺序先动哪一块性价比最高如果粗算超过了目标卡的可用显存按下面的优先级去砍砍完再重新粗算。这个顺序是我多次实验后总结的“对训练质量伤害从低到高”的排序把精度降到BF16/FP16这是白捡的基本盘。调小batch size用梯度累积补足有效batch训练效果几乎不变显存立降。开梯度检查点gradient checkpointing用大概20%-30%的额外时间换回30%-60%的激活显存。换低内存优化器8bit Adam或Adafactor优化器状态直接砍到三分之一。再用LoRA/QLoRA把可训练参数砍到0.1%-1%彻底摆脱全参微调的重账本。最后才是CPU offload、多卡分片这种会明显拖慢速度的方案。为什么这个顺序靠谱前四步基本不改变模型结构和训练轨迹后面换LoRA会让优化空间改变QLoRA更是引入了量化误差。所以尽量在“不动训练范式”的范围内把显存腾出来实在不行才动模型结构。4.3 两个实例预算表8B微调在24G卡上的完整推演用一张24G卡看两个典型场景你就能直观感受预算决策是怎么落地的。场景A8B模型全参微调。静态开销128GB直接出局。不用往下算了24G卡上全参微调8B的唯一出路是多卡分片或者接受极慢的offload而且哪怕真的跑起来batch size也会小到几乎没有实用意义。这个场景真正合理的方案是“换赛道”。场景B8B模型LoRABF16冻结主权重。预算是这样的主权重16GB、LoRA可训练参数0.2GB左右、优化器状态梯度不足0.5GB、激活值按1条2048 token的序列估4GB左右加上CUDA context 1GB。合计约21.7GB卡在上面的22.5GB可用容量附近。这时候决策就来了要么开梯度检查点把激活值打到1.5GB以下这样总预算降到19GB左右留出足够余量要么序列长度提到4096时一定OOM必须砍batch或开重计算。这就是预算决策的价值——不需要真跑一次才知道行不行。场景C同一模型QLoRA 4bit。主权重4GB、LoRA开销0.2GB、激活值2GB、context 1GB加起来7.5GB左右。16G卡甚至部分10G卡都能跑但代价是量化误差和略慢的forward适合那种“只要跑得动效果别差太多”的场景。4.4 决策验证训练前的烧水测试预算做完不要直接全量训练先用一个“烧水测试”验证拿1个batch跑2到3个step监控峰值和预算表是否吻合。这个过程花不了五分钟但能挡住90%的上线即OOM。测试脚本就复用第二节的埋点代码观察几个关键值第一个step后的max_memory_allocated()是否接近预算峰值memory_reserved()离目标卡可用容量还有多少空间backward阶段是否有异常的临时分配。如果预算表和实测值偏差超过15%说明我漏了某项临时开销就得回去检查是不是有额外的中间张量或者你的框架替你分配了不该分配的buffer。5. 踩坑记录与生产建议5.1 我踩过的显存测量坑这节写几个真实踩过的坑希望能让你绕过去。第一个坑只看allocated不看reserved导致我一度以为模型只占14GB实际上分配器已经预留了20GB。后来做压力测试才发现那些“预留”的内存一旦被碎片卡住照样OOM。现在我的打点函数永远同时输出allocated和reserved。第二个坑峰值统计没有reset。有一回我写了个推理测试然后再跑训练测量结果max一直显示的是推理阶段的峰值让我白调了半天batch size。后来养成了“测谁之前先reset谁”的习惯不光是训练每个阶段测量前都单独reset。第三个坑nvidia-smi的进程显存对多进程训练不友好。我早期用DeepSpeed stage2时nvidia-smi显示的per-process memory和每个进程真实占用对不上把自己误导了很久。后来干脆不纠结于nvidia-smi的数字主看PyTorch内部统计nvidia-smi只用来确认“卡上有没有别的人占显存”。5.2 不同训练框架的显存表现差异用框架训练和用裸PyTorch训练显存表现是不一样的。HuggingFace Trainer会在底层封装数据collator和评估循环评估阶段也会占额外显存所以测峰值时要覆盖评估阶段的forward。DeepSpeed的ZeRO stage2、stage3会把优化器状态甚至梯度分片出去主卡显存大幅下降但CPU内存会上升测量时要注意主机侧的内存是不是也爆了。FSDP类似而且通信buffer会带来少量固定显存开销这些在预算里都应该留一点。不要指望框架官方文档里给一个“显存基准”就照搬因为你的模型、序列长度、batch、数据维度全都不一样。唯一合理的方式是拿自己的脚本完整跑一遍烧水测试。5.3 学完这个task我的日常显存管控习惯我现在做训练实验的习惯是把print_mem封装成一个装饰器挂在我的Trainer里每个step自动输出每接一个新模型、新任务先跑一次烧水测试生成一份“显存预算报表”所有实验的峰值数据统一记到一个表里长期维护。这样后面对比不同优化器、不同序列长度、甚至不同PyTorch版本时我有的是数据说话而不是靠猜。最后分享一个小技巧正式训练长跑之前我会故意把batch size调成预算表上限的1.2倍跑一个step看它OOM。这个“预期内OOM测试”能让你知道真实边界在哪等正式配置跑起来时心里特别踏实。这个方法我推荐给每个被显存问题逼疯过的炼丹师。