PPO训练LLM四模型显存占用分析

可视化演示

四模型概览

你走进排练厅，看到四位演员已经就位：主角手里攥着厚厚一沓新台词（Actor 模型），导演助理拿着记分板准备对每个动作打分（Critic 模型），一位“原著复读机”捧着原版剧本随时比对（Reference 模型），还有一位观众代表按着脉搏感受“这届观众买不买账”（Reward 模型）。四者同台，各司其职，这就是 PPO（近端策略优化）训练大语言模型的标准配置。

在 PPO 的 RLHF（基于人类反馈的强化学习）流程里，Actor 是唯一被持续更新的“学生”，它每产出一句回答，Critic 就立刻评估“这句值多少分”，Reward 模型给出奖励信号，Reference 模型则负责计算 KL 散度，防止主角跑调跑得太远。四个模型协同完成“采样—评估—更新—再采样”的循环。

显存占用由参数量、精度和训练状态三因素决定。通常 Actor 和 Reference 是同一个尺寸的底座模型（如 7B），Critic 是较小的价值网络（如 0.4B），Reward 则是更小的奖励模型（如 0.3B）。规模差距直接决定了显存占比：主角和原著复读机是“大块头”，而评论家和观众代表是“轻量级”。

Actor显存

主角演员是全场最“耗空间”的角色。他不仅要站在台上（存储权重），还要在排练中一遍遍对台词（计算梯度），更要把每一次改进心得记成厚厚的排练笔记（优化器状态）。这三样东西，在深度学习里统称“模型状态”。如果我们用 N 表示模型参数量，那么纯 FP32 训练时，仅凭参数本身就要占 4N 字节，梯度同样要占 4N 字节，而 Adam 优化器还要为一阶矩和二阶矩各存一份 4N 字节——加起来就是 16N 字节。

明明生活中我们每天都要做注意力管理，AI 却是在“背”每一个词的笔记——对每个参数，Adam 都要记录“上一次往哪个方向改、改了多少”（一阶矩），以及“这个参数到底多难学”（二阶矩）。这也是思维导图里“Adam 状态占用”最扎心的一点：它比权重本身还要大 2 倍。

$$M_{\text{model}} = M_{\text{weights}} + M_{\text{grads}} + M_{\text{optimizer}} = 4N+4N+8N = 16N \quad(\text{纯FP32})$$

1. 权重占位：参数本身以 4 字节单精度存储，占 4N 字节。

$$M_{\text{weights}} = 4N$$

2. 梯度占位：反向传播要算出每个参数的修正量，同样 4N 字节。

$$M_{\text{grads}} = 4N$$

3. 一阶矩笔记：Adam 记录梯度指数移动平均，即“惯性方向”，再占 4N。

$$M_{\text{first\_moment}} = 4N$$

4. 二阶矩笔记：记录梯度平方的移动平均，即“对每个参数的学习节奏”，又占 4N。

$$M_{\text{second\_moment}} = 4N$$

如果把这个模型换成 7B 参数，16N 就是 112GB——一块 A100-80GB 根本放不下！好在实际训练中我们几乎不会用纯 FP32 硬扛，而是用“轻装版剧本”（BF16 混合精度），把权重和梯度各减半到 2N，这样模型状态就降到 14+14+42=70GB，勉强挤进 80GB 的舞台。正因优化器状态占了大头，显存优化的首要目标就是“压缩排练笔记”。

Critic显存

导演助理（Critic）的任务是给主角每一句台词打分。他也要“揣摩演技”，也就是说 Critic 模型同样需要反向传播更新参数，所以它同样要准备权重、梯度和优化器状态的“三件套”。但好在助理体型小（通常只有 Actor 的 1/10 到 1/20），参数存储本身不算惊人。

然而 Critic 有一个隐藏的大胃口：前向激活值。在计算“这句台词值几分”的时候，每一层的中间结果都要先缓存下来，等反向传播时再取用。激活值大小和批次大小、序列长度、层数都成正比。当输入序列很长时，这部分的量级可能超过权重本身。这也是思维导图“前向激活开销”这一节点的要害：模型小 ≠ 显存一定小，关键在于激活值这座“临时道具山”堆得多高。

$$M_{\text{activation}} = \mathcal{O}(L \times S \times B \times D)$$

1. 层数堆叠：每层 Transformer 都要保存一份中间激活，层数 L 越大，叠加越多。

$$M_{\text{per\_layer}} \propto L$$

2. 序列长度放大：序列越长，注意力矩阵越庞大，激活值与长度 S 成线性甚至平方关系。

$$M_{\text{attn}} \propto S^2$$

3. 批次数倍：批次越大，同时处理的样本越多，显存按 B 倍增长。

$$M_{\text{batch}} \propto B$$

在工程实践中，我们经常会发现 Critic 虽小，但激活值却让它“站得比预想的宽”。应对办法是使用梯度检查点（用时间换空间）：不缓存全部激活，反向传播时重算一次前向。这样激活显存能从 O(L × S × B × D) 降到 O(L × S × B)，非常适合长序列任务。

Ref与Reward

两位“裁判”就好演了。Reference 模型（原著复读机）和 Reward 模型（观众代表）全程只需“品评”，不需“练习”——它们不做反向传播、不更新参数。相比之下，它们只需要权重和前向过程中的激活缓存，不需要梯度，更不需要优化器状态。这一下子就省掉了约 70% 的训练显存开销。

RLHF 的奖励构造依赖一个关键公式：Reward = Reward 模型打分 − β × KL（策略 || Ref），KL 散度正是由 Ref 和 Actor 的输出分布共同算出的。由于 Ref 和 Reward 都只需要前向推理，它们甚至可以用 FP16/BF16 加载并冻结参数，进一步压低显存。唯一要注意的是：它们的前向激活仍然需要保留，用于计算 Reward 和 KL 信号。

$$M_{\text{reward}} = M_{\text{ref}} + M_{\text{re\_model}} \approx 2N_{\text{ref}} + 2N_{\text{rm}}$$

1. Ref 权重：以半精度存储，约 2N（N 为 Ref 参数量）。

$$M_{\text{ref}} \approx 2N_{\text{ref}}$$

2. Reward 权重：同样按半精度存储，约 2N。

$$M_{\text{re\_model}} \approx 2N_{\text{rm}}$$

3. 无梯度项：两者都不计算梯度，所以没有 M_grads 和 M_optimizer。

$$M_{\text{grads}} = M_{\text{optimizer}} = 0$$

在以 7B Actor 为基准的例子里，Ref 约占 14GB（按 BF16），Reward 只有不到 1GB，而同时刻的 Actor 光是优化器状态就扛着 42GB。所以聪明的主播（工程师）往往会想：既然裁判不训练，能否把裁判换一张更小的推理卡？这样就能给主角腾出宝贵的训练舞台。

显存公式

把整个舞台画成一张总平面图：总显存 = 模型状态 + 激活值 + 临时缓冲区。后者包括通信 buffer、CUDA context 等“杂项道具”。理解总公式，是资源规划的第一步。

$$M_{\text{total}} = M_{\text{model}} + M_{\text{activation}} + M_{\text{buffer}}$$

1. 模型状态项：权重 + 梯度 + 优化器状态，是静态显存主体。

$$M_{\text{model}} = M_{\text{weights}} + M_{\text{grads}} + M_{\text{optimizer}}$$

2. 激活值项：随批次 × 序列长度 × 层数动态变化。

$$M_{\text{activation}} = \mathcal{O}(L \times S \times B \times D)$$

3. 缓冲区项：包括 AllReduce 通信缓冲和临时张量，通常预留 10%~20%。

$$M_{\text{buffer}} \approx 0.1 \times M_{\text{total}}$$

拿 7B、BF16、Adam 优化器举例：模型状态 = 14 + 14 + 42 = 70GB；激活值按常见 batch=1、seq=1024 大约 5~15GB；再加上通信缓冲约 10%。总显存 = 70 + 10 + 8 ≈ 88GB，已经超过单卡 80GB。这就是为什么单卡训练 7B 级别的 PPO 非常局促，工程师必须启动混合精度 + ZeRO + 检查点等“减负三板斧”。

混合精度

为了把“16N“压缩成”10N“，道具组想出一个妙招：剧本用大字版（FP32）珍藏在保险柜里（Master 权重），排练时演员只带轻便的小抄（BF16 副本）。这就是混合精度的核心思路：权重、梯度和激活都用 16 位半精度参与计算，但维护一份 FP32 的”权威主权重”用于更新。

$$M_{\text{mixed}} = 2N + 2N + 8N = 12N \approx 10N$$

1. BF16 权重：模型加载半精度副本，占 2N 字节。

$$M_{\text{bf16\_weights}} = 2N$$

2. BF16 梯度：反向传播算出的梯度也以半精度存储，占 2N。

$$M_{\text{bf16\_grads}} = 2N$$

3. FP32 主权重：每个参数仍保留一份 4N 的高精度主副本，保证收敛稳定。

$$M_{\text{master\_weights}} = 4N$$

4. 优化器状态：Adam 的一阶、二阶矩仍建议 FP32 保存，占 8N。

$$M_{\text{optimizer}} = 8N$$

为什么不用纯 FP16？因为梯度太小会“溢出”变成零，就像音量旋钮扭到最低就听不见了。这时梯度缩放器（Grad Scaler）会在反向传播前把梯度放大，更新后再缩回去，确保小梯度也能被听见。BF16 因为保留了 8 位指数位，动态范围更宽，在 LLM 训练中比 FP16 更受欢迎。总之，混合精度相当于把“16 系行李”减少到“10~12 系”，省出的空间立竿见影。

ZeRO策略

当单张舞台实在挤不下，聪明的导演会租用多间排练室，把剧本、道具和笔记分散存放，排练时再“隔空取物”。这就是 ZeRO（零冗余优化器）的思想：把优化器状态、梯度乃至参数切分成碎片，分给 P 张卡各自保管，从而把单卡显存压到原来的 1/P。

$$M_{\text{ZeRO}} = 2N + \frac{12N}{P} \quad(\text{Adam + BF16混合精度})$$

1. 每卡必存权重副本：前向计算时需要完整权重，因此 2N 字节是保底的。

$$M_{\text{weights}} = 2N$$

2. 切分模型状态：梯度 + 优化器状态共 12N 字节被切成 P 片，每卡只拿 1/P。

$$M_{\text{shard}} = \frac{12N}{P}$$

3. P 越大越省：卡数从 1 增至 8，12N/P 从 12N 降到 1.5N，单卡显存趋于 2N 下限。

$$\lim_{P \to \infty} M_{\text{ZeRO}} = 2N$$

在 DeepSpeed 中，ZeRO-1 只切优化器状态，ZeRO-2 再加切梯度，ZeRO-3 连参数都切开（此时权重从“每卡一份”变成“全局仅一份”）。ZeRO-3 省得最多，但通信开销也最大，因为每次前向都要 AllGather 权重。结合思维导图的“道具分仓”比喻：组员越多，每人带的道具越少，但互相借道具的跑来跑去也更频繁。

显存规划

掌握了上述知识，最后一步就是“排班表”：先算总量，再留缓冲，最后设警报。工程上建议至少预留 20%~30% 的余量，因为激活值、临时张量和 CUDA context 常在训练中途突然膨胀。假设你计划用 8×A100-80GB 训练 7B 模型，先用 10~12N 的混合精度公式粗算：模型状态 70GB，再加激活和缓冲约 88GB——单卡不够，但 ZeRO-3 + 梯度检查点能把它压到 20GB 以下，8 卡绰绰有余。

还有几个省钱利器：梯度检查点用重计算换显存，可将激活显存减少约一半以上；Flash Attention 优化注意力计算，让长序列不再“爆心跳”；KV Cache 也要算进推理环节的显存，尤其在 PPO 采样时。规划时建议在训练脚本里加上显存监控，一旦超过阈值就自动减小 batch 或触发 offload。

$$M_{\text{recommended}} = 1.2 \sim 1.3 \times \left( M_{\text{model}} + M_{\text{activation}} + M_{\text{buffer}} \right)$$

1. 安全系数：显存余量取 1.2~1.3，防止临时峰值击穿。

$$k = 1.2 \sim 1.3$$

2. 总显存乘系数：把前文三项相加后再乘 k，得到推荐显存预算。

$$M_{\text{recommended}} = k \times M_{\text{total}}$$

OOM 一旦发生，训练进程会直接崩溃，几小时的心血瞬间归零。所以“先算后训、留足余量、实时监控”是 RLHF 工程师的保命三原则。好消息是，各家框架（DeepSpeed、HF TRL、Megatron）都已内置显存预估工具，团队排卡前跑一遍估算脚本，就能把崩溃风险挡在上线之前。

AI 应用场景

Python 演示

import torchimport torch.nn as nnimport torch.optim as optim  class PolicyNet(nn.Module):    def __init__(self, dim=64):        super().__init__()        self.fc = nn.Sequential(nn.Linear(dim, 128), nn.Tanh(), nn.Linear(128, dim))     def forward(self, x):        return self.fc(x)  class ValueNet(nn.Module):    def __init__(self, dim=64):        super().__init__()        self.fc = nn.Sequential(nn.Linear(dim, 64), nn.Tanh(), nn.Linear(64, 1))     def forward(self, x):        return self.fc(x)  def estimate_gb(model, dtype_bytes=2):    total = sum(p.numel() for p in model.parameters())    return {        "params": total,        "weights_gb": total * dtype_bytes / 1e9,        "optimizer_gb": total * 8 / 1e9,    }  def ppo_step(actor, critic, ref_model, reward_model, states, actions, returns, old_logp, beta=0.05):    logits = actor(states)    logp = nn.functional.log_softmax(logits, dim=-1).gather(1, actions.unsqueeze(-1)).squeeze(-1)    ref_logp = nn.functional.log_softmax(ref_model(states), dim=-1).gather(1, actions.unsqueeze(-1)).squeeze(-1)    rewards = reward_model(states).squeeze(-1).detach()    kl = torch.mean(logp - ref_logp)    adv = rewards + beta * kl    ratio = torch.exp(logp - old_logp)    loss = -torch.mean(torch.clamp(ratio, 0.8, 1.2) * adv.detach()) + 0.1 * torch.mean((returns - critic(states).squeeze(-1)) ** 2)    return loss  def main():    actor = PolicyNet()    critic = ValueNet()    ref = PolicyNet()    reward = ValueNet()    ref.eval()    reward.eval()    for m in (ref, reward):        for p in m.parameters():            p.requires_grad = False    opt = optim.Adam(actor.parameters(), lr=3e-4)    states = torch.randn(32, 64)    actions = torch.randint(0, 64, (32,))    returns = torch.randn(32)    old_logp = torch.randn(32)    for _ in range(20):        opt.zero_grad()        loss = ppo_step(actor, critic, ref, reward, states, actions, returns, old_logp)        loss.backward()        opt.step()    for name, model in (("Actor", actor), ("Critic", critic), ("Ref", ref), ("Reward", reward)):        info = estimate_gb(model)        print(name, "params:", info["params"], "weights(GB):", round(info["weights_gb"], 4), "optimizer(GB):", round(info["optimizer_gb"], 4))  if __name__ == "__main__":    main()

1. class PolicyNet(nn.Module)定义 Actor / Ref 共用的策略网络，模拟生成动作的“主角”。
2. class ValueNet(nn.Module)定义 Critic / Reward 共用的网络，模拟打分和给奖励的“裁判”。
3. def estimate_gb按 BF16 权重（2 字节）+ Adam 优化器（8 字节）估算显存。
4. beta * klKL 散度系数，限制 Actor 偏离 Ref 太远，对应“别把原著改得面目全非”。
5. torch.clamp(ratio, 0.8, 1.2)PPO 的截断（clip）核心，防止策略一步跨太大导致训练崩溃。
6. for m in (ref, reward)冻结 Ref 与 Reward，正是“裁判不训练”的工程实现。

常见误区

面试题