大模型强化学习对齐演进:从 PPO、DPO 到 GRPO 核心数学原理与显存优化实践
在大语言模型(LLM)的训练范式中,经历了预训练(Pre-training)与监督微调(SFT)之后,模型虽然掌握了语言结构与任务指令,但往往存在“对齐税(Alignment Tax)”、长逻辑链推理幻觉以及过度迎合用户等问题。基于人类反馈的强化学习(RLHF)是引导模型价值观对齐、激发思维链(Chain-of-Thought, CoT)深度推理能力的核心途径。
然而,传统的强化学习方案对显存与分布式工程架构的要求极高。本文系统推导从经典 PPO、离线 DPO 到现代推理模型广泛采用的 GRPO(Group Relative Policy Optimization)的技术演进脉络,拆解各算法底层的数学推导与显存消耗模型,并给出生产级工程优化实践。
一、 经典 RLHF 范式:PPO 的四模型显存困境
在基于 Proximal Policy Optimization(PPO)的经典 RLHF 架构中,系统必须在 GPU 集群上同时维护并协调 4 个同等或相似规模的模型:
┌───────────────────────┐
│ Prompt (输入) │
└──────────┬────────────┘
│
┌────────────────┴────────────────┐
▼ ▼
┌──────────────────┐ ┌──────────────────┐
│ Actor (策略模型) │ │ Reference (基准) │ (参数冻结,防跑偏)
│ 生成回答并更新 │ │ 计算 KL 散度 │
└────────┬─────────┘ └────────┬─────────┘
│ 输出 Sequence │
├─────────────────────────────────┘
▼
┌──────────────────┐ ┌──────────────────┐
│ Critic (价值) │ │ Reward (奖励) │ (参数冻结,提供标量奖励)
│ 拟合状态价值 V │ │ 预测最终打分 │
└──────────────────┘ └──────────────────┘
1. 显存开销与通信瓶颈
- Actor 模型:参与前向与反向传播,保存权重、梯度与 AdamW 优化器状态(以 70B 为例,全参训练需占用约 112 GB 显存,未计激活值)。
- Critic 模型(Value Network):与 Actor 同等参数规模,专门输出标量 用于估计状态价值。同样需要参与反向传播,维护一套独立的优化器状态。
- Reference 与 Reward 模型:虽然参数冻结不产生梯度,但依然各自占据 140 GB 左右的高带宽显存用于前向推理。
工程痛点:四模型交织引发极高的流水线气泡(Pipeline Bubble)与巨额通信开销。若 Actor 与 Critic 规模相当,训练成本往往达到纯 SFT 的 3~4 倍以上,极易引发显存爆炸。
二、 离线简化路线:DPO 的数学代换与局限
为了摆脱显存中笨重的 Critic 与 Reward 模型,直接偏好优化(Direct Preference Optimization, DPO)通过巧妙的数学变换,直接利用策略模型本身表示奖励函数。
1. 理论推导核心
Bradley-Terry 偏好模型中,给定偏好回答 与劣质回答 ,人类偏好概率为:
在标准带 KL 正则项的强化学习目标下:
通过拉格朗日乘子法求解,最优策略 与隐式真实奖励 存在严格的闭式解关系:
将上述 的表示直接代入 Bradley-Terry 偏好模型中,配分函数 在差分中被直接抵消。由此推导出无需显式 Reward 模型的 DPO 损失函数:
2. DPO 的工程权衡
- 优势:彻底移除了 Critic 和独立的 Reward 模型,训练流程完全等价于成对数据的有监督交叉熵,显存开销与 SFT 持平。
- 致命缺陷:属于离线强化学习(Offline RL)。模型无法在当前生成空间中进行主动“探索(Exploration)”,只能在静态标注数据集内拟合。在涉及长链路数学推导、代码编译等需要在线试错的强推理场景中,DPO 极易出现概率漂移(OOD)与模式坍塌。
三、 在线探索新范式:GRPO(群体相对策略优化)
为了兼顾 PPO 的“在线动态试错”优势,同时彻底摒弃 Critic 模型的高昂显存负担,GRPO(Group Relative Policy Optimization) 应运而生。
1. 核心架构重构:基于组采样的无 Critic 估计
GRPO 的核心创新在于:完全移除 Critic 模型,改为通过“组内相对对比”直接估计优势(Advantage)。
执行时序:
- 组采样(Group Rollout):对于输入 Prompt ,利用当前的旧策略模型 并发生成一组( 个)候选回答 。
- 多维打分(Reward Assignment):
- 规则打分(Rule-based):对于数学或代码任务,直接使用单元测试通过率、编译器报错、答案正则表达式进行 0/1 判定;
- 轻量模型打分:对于自由文本,引入轻量 Reward 模型或格式评分函数。
- 得到每个候选的标量奖励序列 。
- 组内标准化计算优势(Advantage):
不依赖价值网络预测状态基准线,直接使用该组输出的样本均值和标准差对奖励进行归一化:
2. GRPO 优化目标函数
GRPO 的代理优化目标结合了 PPO 的 Clipped Surrogate Objective 与显式 KL 散度惩罚:
- 大幅节省显存:省去了整个庞大的 Critic 模型结构及其反向传播图,集群显存直接节约近 50%。
- 天然基准对齐:若某批次采样的全部回答得分均高,组内标准化会强制模型区分“优秀”与“极致优秀”;若普遍得分低,能精准奖励相对偏差较小的回答,具备天然的防方差漂移能力。
四、 GRPO 优势计算与目标损失实现(PyTorch 核心范式)
以下代码演示在批处理场景下,GRPO 计算组内标准化优势与 Clipped 策略梯度的核心实现:
import torch
import torch.nn.functional as F
def compute_grpo_advantage(rewards: torch.Tensor, group_size: int, eps: float = 1e-8) -> torch.Tensor:
"""
计算组内相对优势 (Group-Relative Advantage)
:param rewards: 扁平化奖励张量 [Batch_Size * Group_Size]
:param group_size: 每个 Prompt 生成的样本数 (G)
:return: 归一化后的优势张量 [Batch_Size * Group_Size]
"""
# 转换为 [Batch_Size, Group_Size]
rewards_grouped = rewards.view(-1, group_size)
# 计算组内均值与标准差 (沿组维度)
mean = rewards_grouped.mean(dim=1, keepdim=True)
std = rewards_grouped.std(dim=1, keepdim=True)
# 优势归一化
advantages = (rewards_grouped - mean) / (std + eps)
# 展平回一维供损失计算
return advantages.view(-1)
def compute_grpo_loss(
log_probs: torch.Tensor, # 当前模型 log(pi_theta(y|x)) [N, Seq_Len]
old_log_probs: torch.Tensor, # Rollout 时旧模型 log(pi_old(y|x)) [N, Seq_Len]
ref_log_probs: torch.Tensor, # Reference 冻结模型 log(pi_ref(y|x)) [N, Seq_Len]
advantages: torch.Tensor, # 组优势 [N]
attention_mask: torch.Tensor, # Token 掩码 [N, Seq_Len]
clip_eps: float = 0.2,
beta: float = 0.04
) -> torch.Tensor:
"""
计算单个批次下的 GRPO 联合损失
"""
# 1. 序列级别的比率计算 (Importance Sampling Ratio)
# 对 Token 维度的对数概率求和,得到整句对数概率
seq_log_prob = (log_probs * attention_mask).sum(dim=-1)
seq_old_log_prob = (old_log_probs * attention_mask).sum(dim=-1)
ratio = torch.exp(seq_log_prob - seq_old_log_prob)
# 2. PPO-style 截断策略损失
surr1 = ratio * advantages
surr2 = torch.clamp(ratio, 1.0 - clip_eps, 1.0 + clip_eps) * advantages
policy_loss = -torch.min(surr1, surr2).mean()
# 3. 无偏 KL 散度估计 (基于 Schulman 逼近公式: (r - 1) - log(r))
# 保证 KL 恒为非负,避免数值震荡
log_ratio_ref = ref_log_probs - log_probs
kl_penalty = (torch.exp(log_ratio_ref) - 1.0 - log_ratio_ref) * attention_mask
kl_loss = beta * (kl_penalty.sum(dim=-1)).mean()
total_loss = policy_loss + kl_loss
return total_loss
五、 生产环境工程调优与避坑指南
1. 组大小(Group Size, )的调谐陷阱
- 方差与开销的平衡:若 过小(如 ),组内样本标准差极易失真,微小的随机波动会被无限制放大,导致策略更新剧烈震荡;若 过大(如 ),虽然优势估计极准,但 Rollout 阶段的推理开销会占满训练耗时的 80% 以上。
- 工程推荐:对于数学与代码推理任务,生产环境通常将 设定在 6 到 16 之间,在样本多样性与集群吞吐间取得最优收益。
2. 奖励黑客(Reward Hacking)与长度作弊防御
在强化学习中,随着训练步数深入,模型会敏锐地挖掘环境规则的漏洞以套取高分:
- 长度偏见:模型发现“只要输出非常长的推理过程(即使中间包含无意义的反复自问自答),就能获得更高奖励”。
- 规避策略:
- 硬性 Token 惩罚:在 Reward 函数中注入动态步幅惩罚项:超出基础阈值后的 Token 施加微弱的负线性惩罚;
- 严格格式校验(Format Reward):必须输出严格的 XML 标签(如
<think>...</think><answer>...</answer>),一旦发现标签未闭合、结构破损或重复复读,直接判为 0 分。
3. 全局 Reward 坍塌与零方差保护
在数学题难度极高或极低的批次中,组内 个候选回答可能会全部答错(全 0 分)或全部答对(全 1 分):
- 此时组内样本标准差 。如果不加干预,直接归一化会导致除以接近 0 的数字,产生巨大的浮点梯度冲击;
- 容错机制:当组内样本得分全同()时,直接将该组所有样本的 Advantage 强制归零,该 Prompt 在本轮迭代中不贡献任何策略梯度,保护权重不被无序更新破坏。
- 点赞
- 收藏
- 关注作者
评论(0)