★ FEATURED ARTICLE 大模型训练的显存经常不够用大家第一反应是开梯度检查、切混合精度很少有人系统算过优化器状态到底吃了多少显存。Adam优化器每个参数要维护两个FP32状态一阶动量m和二阶动量v一个参数就是8字节。7B模型全参微调光优化器状态就是56GB一张卡塞得下但再加激活值、参数和梯度就爆了。所以把优化器状态挪到CPU也就是标题里说的CPU Offload是目前最实用的解法定式。这篇内容就把我实际做过的三种CPU Offload方案、配置参数和踩坑记录从头到尾聊一遍。1. 先算一笔显存账为什么优化器状态最该挪走1.1 一个参数的完整显存账单很多人以为混合精度训练能省一半显存其实只对了一半。以Adam 混合精度为例一个参数在训练过程中要在显存里维护的可不止一份数据数据精度每参数字节数FP16/FP32参数副本FP16/FP322/4梯度FP16/FP322/4优化器一阶动量 mFP324优化器二阶动量 vFP324主权重Master WeightFP324混合精度下GPU上常驻的往往是FP16参数副本但Adam的m和v无论如何都必须是FP32不然更新精度崩掉。所以一个参数保守要占用244414字节用FP32则直接到16字节。拿最常见的13B模型举例参数副本26GB梯度按FP16算26GB主权重52GB优化器状态104GB。就算不统计激活值也已经超过200GB。单卡A100 80GB在这种量级面前完全没戏。即便缩到7BAdam状态56GB也是绕不过去的坎。1.2 为什么优先搬优化器状态而不是梯度显存不够时梯度也可以放CPU但实际工程里几乎都先搬优化器状态原因是梯度在反向传播中要反复使用。反向传播是一层层回传的每一层的梯度要参与再往前的梯度计算如果梯度在CPU每层都要跨PCIe拷贝回来反向过程会被打断速度惨不忍睹。而优化器状态的使用频率低很多它只在step()更新那一下被读和写而且可以整批异步预取。打个比方梯度是流水线上的核心工具随时要用手边拿优化器状态是仓库里的备件每周用一次放在远一点的仓库并不会让产线停工。另外优化器状态搬走后主权重FP32也可以顺带放CPU因为主权重的更新跟着优化器走。DeepSpeed的ZeRO-Offload就是这么设计的GPU只保留FP16参数副本和临时梯度FP32主权重、m、v全放CPU。显存直接省掉一整块大项。2. DeepSpeed ZeRO-Offload 的标准配置法2.1 一份可直接套用的ds_config如果你不想自己撸底层逻辑DeepSpeed的ZeRO Stage 2 CPU Offload是目前最成熟的方案。它做的事有两层ZeRO Stage 2把优化器状态按数据并行度切分每个rank只持有一部分Offload再把每个rank持有的那部分状态放到CPU。显存立刻降一大截。我最近在微调一个7B模型时用的配置长这样{ train_batch_size: 8, gradient_accumulation_steps: 1, fp16: { enabled: true, auto_cast: true, loss_scale: 0, initial_scale_power: 16, loss_scale_window: 1000 }, zero_optimization: { stage: 2, offload_optimizer: { device: cpu, pin_memory: true, ratio: 1.0 }, allgather_bucket_size: 50000000, reduce_bucket_size: 50000000, overlap_comm: true, contiguous_gradients: true }, optimizer: { type: AdamW, params: { lr: 1e-5, betas: [0.9, 0.999], eps: 1e-8, weight_decay: 0.01 } }, scheduler: { type: WarmupDecayLR, params: { warmup_max_lr: 1e-5, warmup_num_steps: 200, total_num_steps: 5000 } }, communication_data_type: fp16 }训练入口不用改太多无非是把model、optimizer、scheduler丢给DeepSpeed engine然后循环里换成engine.backward()和engine.step()。2.2 关键参数的作用与调节空间上面这份配置里最容易踩坑的是allgather_bucket_size和reduce_bucket_size。这两个参数控制梯度聚合和参数收集的桶大小单位是字节。桶越大单次通信的数据量越大吞吐越高但峰值显存也会高。我试过把两个值调成5e87B模型在8卡A800上单卡显存能多占2GB左右换来的却是step时间提升10%。如果显存本来就紧按5e7用更安全通信慢一点但不会莫名OOM。pin_memory这一项值得开。开启后生成的CPU张量会固定在物理内存页里不会被换页host到device的拷贝可以直接走DMA省一次中间缓冲传输效率明显更高。代价是这些页面不会被系统回收主机内存占用看着会偏大。如果你的服务器CPU内存很紧张可以关掉它但传输耗时大概会多20%到30%。ratio这个参数通常保持1.0。它表示优化器状态放进CPU的比例。如果你机器CPU内存不够又想要一点折中可以设成0.5但状态被切得七零八落训练收敛很受影响我不建议这么干。另外offload_optimizer开起来后fp16开关也必须配套。如果同时开着fp16和offloadDeepSpeed会在CPU维护FP32主权重在GPU维护FP16参数副本更新流程是GPU参数字典取回FP16参数 - 传给CPU端优化器 - CPU更新主权重 - 转成FP16 - 拷贝回GPU。这个过程冗长但显存最省。老版本DeepSpeed在fp16关闭的情况下也能offload只是主权重和参数都成了FP32显存优势没那么明显。3. 不引框架手写一个CPU优化器卸载器有人不想上DeepSpeed觉得引入一堆依赖太笨重。这个想法我完全理解。手动实现CPU Offload其实也不算复杂尤其在单卡或者单机多卡、状态不需要切分的情况下。3.1 核心思想与整体流程手写版的思路一句话就能说明白优化器state里的m和v张量在创建时就不放在GPU而是放在CPU每次step()先把它们从CPU异步拷贝到GPU等计算更新完再把更新后的值写回CPU。整个流程拆开就三步构造Optimizer时把state[exp_avg]和state[exp_avg_sq]创建成torch.zeros_like(p, devicecpu)。每次step时用non_blockingTrue把它们转到GPU。更新完成后用copy_()把GPU上的新状态覆盖回CPU并释放GPU侧临时张量。为什么是异步拷贝因为PCIe传输不是瞬间完成的non_blocking可以让拷贝不阻塞主线程等后面真正要算子运算时再同步。基础版我先把逻辑写清楚进阶版再讲怎么把传输藏起来。3.2 核心代码实现基础版下面这个类是一个可运行的Adam CPU Offload版本只处理单机单卡、FP32参数的场景用于演示核心机制。import torch class CPUOffloadAdam(torch.optim.Optimizer): def __init__(self, params, lr1e-4, betas(0.9, 0.999), eps1e-8): defaults dict(lrlr, betasbetas, epseps) super().__init__(params, defaults) for group in self.param_groups: for p in group[params]: state self.state[p] # 状态直接建在CPU上 state[exp_avg] torch.zeros_like(p.data, devicecpu) state[exp_avg_sq] torch.zeros_like(p.data, devicecpu) state[step] 0 torch.no_grad() def step(self, closureNone): loss None if closure is not None: with torch.enable_grad(): loss closure() for group in self.param_groups: beta1, beta2 group[betas] eps group[eps] lr group[lr] for p in group[params]: if p.grad is None: continue state self.state[p] state[step] 1 grad p.grad.float() # 1. CPU - GPU 拷贝 exp_avg state[exp_avg].to(p.device, non_blockingTrue) exp_avg_sq state[exp_avg_sq].to(p.device, non_blockingTrue) # 2. GPU上完成更新 bias_correction1 1.0 - beta1 ** state[step] bias_correction2 1.0 - beta2 ** state[step] exp_avg.mul_(beta1).add_(grad, alpha1.0 - beta1) exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value1.0 - beta2) denom (exp_avg_sq.sqrt() / bias_correction2 ** 0.5).add_(eps) step_size lr / bias_correction1 p.data.addcdiv_(exp_avg, denom, value-step_size) # 3. 更新后的状态写回CPU state[exp_avg].copy_(exp_avg, non_blockingTrue) state[exp_avg_sq].copy_(exp_avg_sq, non_blockingTrue) # 释放GPU临时张量 del exp_avg, exp_avg_sq return loss几个细节要注意grad必须转成FP32再参与动量更新否则在混合精度下梯度是FP16精度损失在长期训练里会累积成收敛问题。每次step后一定要copy回CPU否则下一次step在CPU上读到的还是旧状态训练等于白跑。copy_用non_blockingTrue是有条件的目标张量是CPU张量时并没有实际意义但加上无害真正有意义的是源张量是GPU张量、目标是固定页内存时。这里为了演示统一加了生产环境可以按精度要求去掉。3.3 让传输和计算重叠的进阶写法基础版有一个明显问题每个参数的state拷贝都要等PCIe传输几百微秒整体step时间会拖长。更好的做法是引入独立CUDA Stream将传输和一部分计算排到不同的流里让它们重叠执行。def step_advanced(self, closureNone): transfer_stream torch.cuda.Stream() transfer_stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(transfer_stream): for group in self.param_groups: for p in group[params]: if p.grad is None: continue state self.state[p] state[exp_avg].to(p.device, non_blockingTrue) state[exp_avg_sq].to(p.device, non_blockingTrue) torch.cuda.current_stream().wait_stream(transfer_stream) for group in self.param_groups: for p in group[params]: if p.grad is None: continue # 此刻状态已经就位正常完成更新 state self.state[p] exp_avg state[exp_avg].to(p.device) exp_avg_sq state[exp_avg_sq].to(p.device) # ... 更新逻辑同上 ... state[exp_avg].copy_(exp_avg) state[exp_avg_sq].copy_(exp_avg_sq)进阶版需要注意transfer_stream里做的是异步拷贝当前流在wait_stream之后才会读到数据所以更新阶段用同步.to(p.device)就够。如果觉得手写麻烦其实还有一个折中方案用PyTorch Optimizer标准写法只在radix大的层比如超大Embedding上单独做CPU Offload其他层保持GPU。实测某些热门Lang大头把Embedding和部分大MLP的优化器状态offload出去显存能降20%step时间只多5%左右。这种局部offload经常比全量offload更平衡。4. 什么时候该用、什么时候别用性能与容量权衡4.1 一次实测对比数据我在一台双路服务器上做过一组对比测试模型是7B数据并行度8卡PCIE 3.0CPU是高端至强batch size按每卡1条样本FP16混合精度微调。记录了三组配置配置单卡峰值显存平均step时间注纯GPUZeRO-2m/v在GPU31.2GB2.1s接近显存上限ZeRO-2 CPU Offloadpin_memorytrue17.4GB3.3s显存降14GB手写全量CPU Offload单卡模拟16.8GB5.6s传输出血看到差距没有框架版step时间增加约57%手写版直接翻了2.7倍。核心原因是框架版有overlap_comm传输和通信有重叠而且梯度、参数用桶式聚合传输是连续的而手写版逐参数拷贝PCIe利用率被拉低了。如果你想手写务必把所有参数的状态拼成一块大tensor再整体拷贝别一个参数一个参数来。显存下降的数字非常稳定7B模型省出来13到14GB足够把batch size从1提到2或者给激活值留缓冲。但代价是训练速度变慢所以CPU Offload本质上是拿带宽换容量的买卖。4.2 判断CPU Offload是否适合你的三句话我的判断标准很简单遇到下面任一情况offload优化器状态就是划算的模型参数量*8字节已经超过可用显存的25%以上。比如7B模型优化器状态就要56GBA100 80GB单卡先占掉70%容量明显该搬。你更在意能跑起来而不是收敛速度。离线微调任务、非交互式任务多花点时间无所谓。训练脚本里batch size已经压到最小值显存还是差一口气而且瓶颈恰好是优化器状态而不是激活值。反过来如果模型在纯GPU模式下已经跑得稳稳当当就别折腾offload。还有如果你用的是PCIE 3.0的老机器offload带来的带宽瓶颈会被放大实测PCIE3比PCIE4的step时间多30%以上不是万不得已优先加显存或缩小模型规模。5. 实操中避坑典型问题与排查实录5.1 常见问题速查表现象原因解决方法开启offload后反而OOM没有配合ZeRO状态切分单卡仍在重复持有全量状态把zero stage设为2或者手写时确认每个rank只保留自己的分片step时间暴涨3倍以上逐参数走PCIe传输没有批量化合并状态为连续张量启用overlap_commpin_memoryCPU内存占用异常高pin_memory缓冲区过大或未释放调小allgather_bucket_size关pin_memory监控主机RSSloss变成NaNCPU端主权重更新后没有同步回GPU参数检查fp16参数副本是否在step后重新拷贝核对state张量设备多卡显存不均衡ZeRO切分颗粒大某些层参数多调整桶大小检查模型是否分层不均衡单独offload大Embedding训练一开始就卡死用了DDP又想混DeepSpeed engine两者不能混用DDP和ZeRO只能选一个用DeepSpeed的引擎统一forward/backward5.2 我的排障思路和小技巧第一件事永远是看通没通传输。训练卡死的时候先检查是不是DDP和ZeRO混在一起了。有一次一个同学接手我的配置为了托管数据并行在脚本外面又包了一层DDP结果step直接死等。记住DeepSpeed的engine本身就已经做了数据并行和梯度同步外面再包DDP会锁死。第二件事是看CPU端内存碎片。手写版里我犯过一个低级错误每次step都新建CPU张量旧的不释放主机内存越涨越高最后被系统杀掉。正确做法是在__init__里创建一次状态step里只做copy_回写别反复torch.zeros。CPU内存分配是真有碎片的频繁申请释放会碎到找不出连续大块。第三件事是验证state到底在哪个设备。排查NaN时我习惯在step开头加两行打印print(self.state[p][exp_avg].device) print(p.device, p.grad.device)看设备不匹配就立刻定位。很多NaN问题根本不是数值溢出而是FP32状态还在CPU上却被当成GPU张量做了就地更新PyTorch自动隐式搬运精度和语义全乱掉了。还有一个实际技巧训练过程中监控主机内存和PCIe带宽。Linux下用nvidia-smi dmon看设备侧利用率pidstat -p pid -r看RSS。如果是PCIe带宽持续跑满而GPU计算利用率低说明offload太多了把部分小层放回GPU反而更快。6. 最后的个人体会CPU Offload优化器状态这件事折腾过一次之后再看别的显存优化手段会通很多。它的本质是重新规划数据流哪些张量必须待在计算单元旁边哪些可以放到远程仓库按需调用。优化器状态恰好是那个可以放远程仓库的典型因为更新频率低、数据量大、结构规则。DeepSpeed帮你做的是把一次粗糙的手工搬运用流式打包、通信重叠、异步预取这些手段打磨得更高效。如果让我给刚上手的人一个建议先用DeepSpeed配置跑通观察step时间和显存曲线再去手写别上来就撸。手写能帮你理解原理但框架的优化细节确实经过千锤百炼很多坑他们早就填平了。等你用框架跑通一轮训练再试图手写也不迟。最后分享一个我后来一直在用的小技巧不要把offload当成二选一的全量开关。把模型里那些参数量巨大但又不太影响训练的层比如大词表Embedding单独offload其他层保持GPU经常能兼顾速度和容量效果比全量offload好不少。这个思路跟工程里“分布式的关键是把正确的数据放在正确的地方”一个道理。
阅读完成 · 觉得有帮助？
