大模型训练的显存经常不够用，大家第一反应是开梯度检查、切混合精度，很少有人系统算过优化器状态到底吃了多少显存。Adam优化器每个参数要维护两个FP32状态（一阶动量m和二阶动量v），一个参数就是8字节。7B模型全参微调，光优化器状态就是56GB，一张卡塞得下但再加激活值、参数和梯度就爆了。所以把优化器状态挪到CPU，也就是标题里说的CPU Offload，是目前最实用的解法定式。这篇内容就把我实际做过的三种CPU Offload方案、配置参数和踩坑记录从头到尾聊一遍。

1. 先算一笔显存账：为什么优化器状态最该挪走

1.1 一个参数的完整显存账单

很多人以为混合精度训练能省一半显存，其实只对了一半。以Adam + 混合精度为例，一个参数在训练过程中要在显存里维护的可不止一份数据：

混合精度下GPU上常驻的往往是FP16参数副本，但Adam的m和v无论如何都必须是FP32，不然更新精度崩掉。所以一个参数保守要占用2+4+4+4=14字节，用FP32则直接到16字节。

拿最常见的13B模型举例：参数副本26GB，梯度按FP16算26GB，主权重52GB，优化器状态104GB。就算不统计激活值，也已经超过200GB。单卡A100 80GB在这种量级面前完全没戏。即便缩到7B，Adam状态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。这两个参数控制梯度聚合和参数收集的桶大小，单位是字节。桶越大，单次通信的数据量越大，吞吐越高，但峰值显存也会高。我试过把两个值调成5e8，7B模型在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和offload，DeepSpeed会在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。整个流程拆开就三步：

1. 构造Optimizer时，把state['exp_avg']和state['exp_avg_sq']创建成torch.zeros_like(p, device='cpu')。
2. 每次step时，用non_blocking=True把它们转到GPU。
3. 更新完成后用copy_()把GPU上的新状态覆盖回CPU，并释放GPU侧临时张量。

为什么是异步拷贝？因为PCIe传输不是瞬间完成的，non_blocking可以让拷贝不阻塞主线程，等后面真正要算子运算时再同步。基础版我先把逻辑写清楚，进阶版再讲怎么把传输藏起来。

3.2 核心代码实现（基础版）

下面这个类是一个可运行的Adam CPU Offload版本，只处理单机单卡、FP32参数的场景，用于演示核心机制。

import torch class CPUOffloadAdam(torch.optim.Optimizer): def __init__(self, params, lr=1e-4, betas=(0.9, 0.999), eps=1e-8): defaults = dict(lr=lr, betas=betas, eps=eps) 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, device="cpu") state["exp_avg_sq"] = torch.zeros_like(p.data, device="cpu") state["step"] = 0 @torch.no_grad() def step(self, closure=None): 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_blocking=True) exp_avg_sq = state["exp_avg_sq"].to(p.device, non_blocking=True) # 2. GPU上完成更新 bias_correction1 = 1.0 - beta1 ** state["step"] bias_correction2 = 1.0 - beta2 ** state["step"] exp_avg.mul_(beta1).add_(grad, alpha=1.0 - beta1) exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1.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_blocking=True) state["exp_avg_sq"].copy_(exp_avg_sq, non_blocking=True) # 释放GPU临时张量 del exp_avg, exp_avg_sq return loss

几个细节要注意：

• grad必须转成FP32再参与动量更新，否则在混合精度下梯度是FP16，精度损失在长期训练里会累积成收敛问题。
• 每次step后一定要copy回CPU，否则下一次step在CPU上读到的还是旧状态，训练等于白跑。
• copy_用non_blocking=True是有条件的，目标张量是CPU张量时并没有实际意义，但加上无害；真正有意义的是源张量是GPU张量、目标是固定页内存时。这里为了演示统一加了，生产环境可以按精度要求去掉。

3.3 让传输和计算重叠的进阶写法

基础版有一个明显问题：每个参数的state拷贝都要等PCIe传输几百微秒，整体step时间会拖长。更好的做法是引入独立CUDA Stream，将传输和一部分计算排到不同的流里，让它们重叠执行。

def step_advanced(self, closure=None): 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_blocking=True) state["exp_avg_sq"].to(p.device, non_blocking=True) 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.0，CPU是高端至强，batch size按每卡1条样本，FP16混合精度微调。记录了三组配置：

看到差距没有？框架版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模型，优化器状态就要56GB，A100 80GB单卡先占掉70%容量，明显该搬。
• 你更在意能跑起来，而不是收敛速度。离线微调任务、非交互式任务，多花点时间无所谓。
• 训练脚本里batch size已经压到最小值，显存还是差一口气，而且瓶颈恰好是优化器状态而不是激活值。

反过来，如果模型在纯GPU模式下已经跑得稳稳当当，就别折腾offload。还有，如果你用的是PCIE 3.0的老机器，offload带来的带宽瓶颈会被放大，实测PCIE3比PCIE4的step时间多30%以上；不是万不得已，优先加显存或缩小模型规模。

5. 实操中避坑：典型问题与排查实录

5.1 常见问题速查表

5.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好不少。这个思路跟工程里“分布式的关键是把正确的数据放在正确的地方”一个道理。