基于近端策略优化（PPO）的大语言模型训练进阶技巧

最新推荐文章于 2025-12-20 19:32:34 发布

原创最新推荐文章于 2025-12-20 19:32:34 发布 · 957 阅读

27 ·

CC 4.0 BY-SA版权

文章标签：

#语言模型 #人工智能 #自然语言处理

部署运行你感兴趣的模型镜像

1. 引言

基于人类反馈的强化学习（Reinforcement Learning from Human Feedback, RLHF）结合近端策略优化（Proximal Policy Optimization, PPO）算法[1]，是微调大语言模型（Large Language Models, LLMs）的有效方法。该方法通过可靠的PPO算法与人类评估者的反馈，提升模型生成回复的质量。然而，采用PPO训练LLMs时面临多重挑战，包括维持训练过程的稳定性，以及实现优于直接偏好优化（Direct Preference Optimization, DPO）[2]的性能表现。为此，OpenRLHF 总结了RLHF与PPO结合的实用训练技巧，旨在帮助研究者更便捷地微调LLMs，同时兼顾训练稳定性与高性能目标。来源：Advanced Tricks for Training Large Language Models with Proximal Policy Optimization

github：OPenRLHF

2 基于PPO训练大语言模型的进阶技巧

OpenRLHF 提出三类基于PPO的训练技术：

1）大语言模型（LLM）专项技巧;

2）PPO算法专项技巧;

3）近期研究中的创新策略。

其中，LLM专项技巧与PPO算法专项技巧已在多个强化学习框架（参见文献[3][4]）中实现应用并验证有效；而近期论文提出的创新策略，其针对特定任务场景的适用性仍有待进一步验证。

2.1 大语言模型（LLM）专项技巧

token 级KL惩罚

在训练过程中，系统会逐 token 计算强化学习（RL）模型与监督微调（SFT）模型生成响应的概率分布之间的KL散度（衡量两个概率分布差异的指标）[11]。

该token级KL散度值将作为惩罚项融入奖励函数的计算。具体而言，每个 token 对应的奖励值按如下方式表示：

$r(s_t, a_t) = \textbf{I}(s_t =[\text{EOS}])r(x,y)-\beta \text{KL}(t) \ \ \ (1)$

$\text{KL}(t) = \log({\pi_{\theta_{\text{old}}}(a_t|s_t)^{\text{RL}}}/{\pi^{\text{SFT}}(a_t|s_t)})\ \ \ (2)$

其中x是 prompt ，y是 response， $\textbf{I}(s_t =[\text{EOS}])$ 是表示t是否为最后一个 token 的恒等函数。

公式1的代码：

'''
r：奖励分数（如 reward model 输出），形状为 [batch]。
kl_coef：KL 散度惩罚系数。
kl：KL 散度（如 logp - logp_ref），形状为 [batch, seq_len]。
action_mask：动作掩码，标记哪些 token 是有效动作。
reward_clip_range：奖励裁剪区间（可选）。
'''
def compute_reward(
    r: Union[torch.Tensor, float],
    kl_coef: float,
    kl: Union[torch.Tensor, list[torch.Tensor]],
    action_mask: Optional[torch.Tensor] = None,
    reward_clip_range: Tuple[float, float] = None,
) -> Union[torch.Tensor, list[torch.Tensor]]:
    if kl_coef <= 0.0:
        kl_coef = 0.0

    if reward_clip_range:
        r = r.clamp(min=reward_clip_range[0], max=reward_clip_range[1])

    kl_reward = -kl_coef * kl
    # The following code is equivalent to:
    #
    # last_reward = torch.zeros_like(kl)
    # for i in range(last_reward.size(0)):
    #     for t in reversed(range(last_reward.size(1))):
    #         if action_mask[i][t] > 0.5:
    #             last_reward[i][t] = r[i]
    #             break
    #
    # 反向查找每一行 action_mask 最后一个为 1 的位置（即最后一个有效动作的索引）
    eos_indices = action_mask.size(1) - 1 - action_mask.long().fliplr().argmax(dim=1, keepdim=True)
    # 在这些位置上将 r 的值赋给 last_reward，其余位置为 0
    last_reward = torch.zeros_like(kl).scatter_(dim=1, index=eos_indices, src=r.unsqueeze(1).to(kl.dtype))

    reward = last_reward + kl_reward

    return reward

广义优势估计（GAE）

GAE [10]，一种TD(λ)回报估计方法，用于在PPO中估计逐 token 奖励。在实践中，我们通常将λ设置为1，将GAE方法转换为蒙特卡罗估计方法。将GAE的λ设置为1，这可以减少在使用基于规则的奖励模型时价值网络引入的偏差。

    @torch.no_grad()
    def get_advantages_and_returns(
        self,
        values: torch.Tensor,
        rewards: torch.Tensor,
        action_mask: torch.Tensor,
        gamma: float,
        lambd: float,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Function that computes advantages and returns from rewards and values.
        Calculated as in the original PPO paper: https://arxiv.org/abs/1707.06347
        Note that rewards may include a KL divergence loss term.

        Advantages looks like this:
        Adv1 =  R1 + γ * λ * R2     + γ^2 * λ^2 * R3       + ...
              - V1 + γ * (1 - λ) V2 + γ^2 * λ * (1 - λ) V3 + ...

        Returns looks like this:
        Ret1 =  R1 + γ * λ * R2     + γ^2 * λ^2 * R3       + ...
                   + γ * (1 - λ) V2 + γ^2 * λ * (1 - λ) V3 + ...

        Input:
        - values: Tensor of shape (batch_size, response_size)
        - rewards: Tensor of shape (batch_size, response_size)

        Output:
        - advantages: Tensor of shape (batch_size, response_size)
        - returns: Tensor of shape (batch_size, response_size)
        """
        lastgaelam = 0
        advantages_reversed = []
        response_length = rewards.size(1)

        # Mask invalid responses
        if action_mask is not None:
            values = action_mask * values
            rewards = action_mask * rewards

        # 反向遍历序列，递推计算 GAE
        '''
        对于每个时间步 t（从后往前递推）：
        delta_t = reward_t + gamma * value_{t+1} - value_t
        advantage_t = delta_t + gamma * lambda * advantage_{t+1}
        reward_t：第 t 步的奖励
        value_t：第 t 步的 value
        gamma：折扣因子
        lambda：GAE 衰减因子
        '''

        for t in reversed(range(response_length)):
            nextvalues = values[:, t + 1] if t < response_length - 1 else 0.0
            delta = rewards[:, t] + gamma * nextvalues - values[:, t]
            lastgaelam = delta + gamma * lambd * lastgaelam
            advantages_reversed.append(lastgaelam)
        advantages = torch.stack(advantages_reversed[::-1], dim=1)
        # 回报
        returns = advantages + values
        return advantages.detach(), returns

添加SFT损失：将额外的有监督逐 token 预测损失与KL散度一起纳入PPO，可以保留SFT模型的现有能力。
token 级损失：当启用样本打包时，OpenRLHF会从训练样本中移除填充内容，确保损失是在token 级别而不是每个样本上计算的，这会提高最终性能。

2.2 PPO算法专项技巧

模型初始化：在使用PPO训练大语言模型（LLM）时，必须初始化两个核心模型：策略模型（actor model）与价值模型（critic model）[6, 7]。具体而言，将策略模型初始化为监督微调（SFT）模型，并将价值模型初始化为奖励模型，能够有效保障PPO训练的高效性。

    # configure model
    # load huggingface model
    actor = Actor(
        args.pretrain,
        use_flash_attention_2=args.flash_attn,
        bf16=args.bf16,
        load_in_4bit=args.load_in_4bit,
        lora_rank=args.lora_rank,
        lora_alpha=args.lora_alpha,
        target_modules=args.target_modules,
        lora_dropout=args.lora_dropout,
        ds_config=strategy.get_ds_train_config(is_actor=True),
    )

    if args.actor_init_on_gpu:
        actor = actor.to(torch.cuda.current_device())

    critic = get_llm_for_sequence_regression(
        args.reward_pretrain,
        "critic",
        normalize_reward=args.normalize_reward,
        use_flash_attention_2=args.flash_attn,
        bf16=args.bf16,
        load_in_4bit=args.load_in_4bit,
        lora_rank=args.lora_rank,
        lora_alpha=args.lora_alpha,
        target_modules=args.target_modules,
        lora_dropout=args.lora_dropout,
        ds_config=strategy.get_ds_train_config(is_actor=False),
        head_prefix=args.head_prefix,
    )

Adam优化器学习率设置：策略模型（actor model）的Adam优化器学习率约为监督微调（SFT）模型的十分之一。例如，在OpenRLHF框架中，SFT模型的Adam学习率设为5×10−6，而策略模型的学习率则为5×10−7。此外，价值模型（critic model）的Adam学习率约为SFT模型的两倍，典型取值为9×10−6。
小批量更新策略：在训练过程中，PPO实现会对规模为N×M的训练数据索引进行随机打乱（其中N表示回放缓冲区容量，M代表每条响应的token长度），随后将其划分为多个小批量（mini-batch）以计算梯度并更新策略网络。(使用DataLoader函数)
价值函数损失裁剪[5]：PPO对价值函数的裁剪处理与PPO的裁剪代理目标[6, 7]机制类似。给定目标价值 Vtarg=收益=优势值+价值估计，PPO通过最小化以下损失函数来训练价值网络：

$Loss_v = \max[(V_{\theta_t} - V_{targ})^2, (\text{clip}(V_{\theta_t}, V_{\theta_{t-1}} - \epsilon, V_{\theta_{t-1}} + \epsilon) - V_{targ})^2] \ \ \ (3)$

class ValueLoss(nn.Module):
    """
    Value Loss for PPO
    """

    def __init__(self, clip_eps: float = None) -> None:
        super().__init__()
        self.clip_eps = clip_eps

    def forward(
        self,
        values: torch.Tensor,
        old_values: torch.Tensor,
        returns: torch.Tensor,
        action_mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        if self.clip_eps is not None:
            values_clipped = old_values + (values - old_values).clamp(-self.clip_eps, self.clip_eps)
            surr1 = (values_clipped - returns) ** 2
            surr2 = (values - returns) ** 2
            loss = torch.max(surr1, surr2)
        else:
            loss = (values - returns) ** 2

        loss = masked_mean(loss, action_mask, dim=-1).mean()
        return 0.5 * loss

奖励归一化与裁剪[5]：Bradley-Terry模型易出现过拟合简单样本而忽略困难样本的情况，导致简单响应对的奖励差距大于困难响应对[6, 7]。这种差异会引发梯度失衡问题。通过奖励归一化与裁剪技术，可均衡所有样本的奖励分布，从而缓解该效应。实际应用中，我们采用Z分数归一化方法：r=(r−μ)/δ，其中μ为奖励训练数据集的均值，δ为标准差。
优势值归一化：在使用均方误差损失训练价值网络时，算法对部分较大数值较为敏感。通过对优势值进行归一化处理，能够有效缓解这些大数值带来的影响[5]。实际应用中，我们同样采用Z分数归一化方法：r=(r−μ)/δ，其中μ为当前批次样本的均值，δ为当前批次样本的标准差。
小规模训练周期：在策略网络训练过程中，必须严格限制训练周期数。该约束条件能够有效防止更新后模型的响应分布与旧模型产生显著偏离，从而避免基于非策略样本（off-policy samples）的训练。具体而言，我们将训练周期数设定为1。

2.3 - 创新策略

训练初期冻结策略模型参数：在策略模型更新策略的过程中，其生成响应的分布会逐渐偏离历史分布。这种偏移会导致优势函数估计失准，并可能引发后续策略更新的误导。因此，通过在训练初期冻结策略模型参数，可确保响应分布保持稳定[6]，从而使价值模型能够基于该固定分布学习到准确的评估结果。
奖励基线方法[7][8]：侯等人[7]研究发现，单个任务内及跨不同任务的奖励分数存在显著波动。因此，若不考虑提示词与任务背景的影响，仅凭高奖励分数无法可靠判定响应质量。为解决这一问题，研究者引入参考基线机制以消除绝对奖励值的波动影响：

$r(x, y) = r(x ,y) - r(x, y_{\text{ref}}) \ \ \ (4)$

其中， $y_{\text{ref}}$ 表示当前策略（即LLM模型）针对提示词 x 所生成的响应。这是一种改进的奖励归一化方法。

PPO训练中的拒绝采样微调：拒绝采样微调方法通过从模型生成K个候选输出，基于奖励模型筛选最优候选响应，并利用该最优样本对模型进行微调。Touvron等人[9]采用拒绝采样微调与近端策略优化（PPO）的序列化组合策略来更新策略网络：在每个训练阶段，首先通过拒绝采样更新模型，再应用PPO算法进行策略优化。

3 参考文献

[1] Schulman J, Wolski F, Dhariwal P, et al. Proximal policy optimization algorithms[J]. arXiv preprint arXiv:1707.06347, 2017.

[2] Rafailov R, Sharma A, Mitchell E, et al. Direct preference optimization: Your language model is secretly a reward model[J]. Advances in Neural Information Processing Systems, 2024, 36.

[3] Reinforcement Learning From Human Feedback — My sample book

[4] Hu J, Wu X, Wang W, et al. OpenRLHF: An Easy-to-use, Scalable and High-performance RLHF Framework[J]. arXiv preprint arXiv:2405.11143, 2024.

[5] Engstrom L, Ilyas A, Santurkar S, et al. Implementation matters in deep policy gradients: A case study on ppo and trpo[J]. arXiv preprint arXiv:2005.12729, 2020.

[6] Cai Z, Cao M, Chen H, et al. Internlm2 technical report[J]. arXiv preprint arXiv:2403.17297, 2024.

[7] Hou Z, Niu Y, Du Z, et al. ChatGLM-RLHF: Practices of Aligning Large Language Models with Human Feedback[J]. arXiv preprint arXiv:2404.00934, 2024.

[8] Shen W, Zhang X, Yao Y, et al. Improving Reinforcement Learning from Human Feedback Using Contrastive Rewards[J]. arXiv preprint arXiv:2403.07708, 2024.

[9] Touvron H, Martin L, Stone K, et al. Llama 2: Open foundation and fine-tuned chat models[J]. arXiv preprint arXiv:2307.09288, 2023.

[10] Schulman J, Moritz P, Levine S, et al. High-dimensional continuous control using generalized advantage estimation[J]. arXiv preprint arXiv:1506.02438, 2015.

[11] Ramamurthy R, Ammanabrolu P, Brantley K, et al. Is reinforcement learning (not) for natural language processing: Benchmarks, baselines, and building blocks for natural language policy optimization[J]. arXiv preprint arXiv:2210.01241, 2022.

您可能感兴趣的与本文相关的镜像