在 RLHF 中重新引入强化学习 (RL)

发布于 2024 年 6 月 12 日

在 GitHub 上更新

黄圣毅 (Shengyi Costa Huang)

vwxyzjn

Arash Ahmadian

ArashAhmadian

访客

我们很高兴在 TRL 中推出 RLOO（REINFORCE Leave One-Out）训练器。作为 PPO 的替代方案，RLOO 是一种新的在线 RLHF 训练算法，旨在更易于访问和实现。特别地，**RLOO 需要更少的 GPU 内存，并且收敛时间更短。**如下图所示：

🤑RLOO 的 vRAM 使用量比 PPO **减少约 50-70%**，具体取决于模型大小。
🚀RLOO 在 1B 模型上比 PPO **快 2 倍**，在 6.9B 模型上比 PPO **快 3 倍**。
🔥RLOO 在响应胜率方面（由 GPT4 评判）**与 PPO 具有竞争力**，并且始终优于 DPO 等流行的离线方法。

通过 RLOO，我们将强化学习重新引入 RLHF，使社区能够更轻松地探索在线 RL 方法。这令人兴奋，因为越来越多的研究表明在线 RL 比 DPO 等离线方法更有效（https://arxiv.org/abs/2402.04792，https://arxiv.org/abs/2405.08448）。

alt_text

这篇博客文章将解释 RLOO 训练器背后的动机、其工作原理以及如何在 TRL 中使用它。

动机

PPO 是一种有效的在线 RLHF 训练算法，用于训练 GPT-4 等最先进的模型。然而，由于其高昂的 GPU 内存要求，PPO 在实际使用中可能非常具有挑战性。特别是，PPO 需要将 4 份模型加载到内存中：1）策略模型，2）参考策略模型，3）奖励模型，以及 4）价值模型，如下图所示。PPO 还有许多微妙的实现细节，可能很难做到正确（Engstrom 等人；2020，Huang 等人 2022）。

在 Cohere 的一篇新论文中，Ahmadian 等人 (2024) 回顾了 RLHF 训练的基础，并提出了一种更优雅的方法，称为 RLOO，一种新的在线训练算法。RLOO 只需要将 3 份模型加载到内存中：1）策略模型，2）参考策略模型，以及 3）奖励模型，如上图所示。

重要的是，RLOO 需要更少的内存，这意味着它更容易

在没有 OOM（内存不足错误）的情况下运行
能够加载更大的批次大小
运行更高效、更快。

此外，RLOO 将整个完成令牌建模为单个动作，如下图所示。在下一节中，我们将结合代码片段深入探讨更多细节。

RLOO 的工作原理

RLOO 和 PPO 都有几个共享步骤

策略模型将生成一些完成令牌，并获得当前策略和参考策略下的每个令牌的对数概率。
然后，我们计算每个令牌的 KL 惩罚，即当前策略和参考策略下对数概率之间的差异。
然后，我们从奖励模型中获取整个完成的得分。

从这里开始，常规 PPO 和 RLOO 的方法有所不同。RLOO 有几个关键思想。首先，它将**整个模型完成**视为一个单一动作，而常规 PPO 将**每个完成令牌**视为单个动作。通常，只有 EOS 令牌获得真实奖励，这非常稀疏。常规 PPO 会将奖励归因于 EOS 令牌，而 RLOO 会将该 EOS 奖励归因于整个完成，如下所示。

from torch import Tensor
response = Tensor([4., 5., 6.])
per_token_logprobs = Tensor([-12.3, -8.3, -2.3])
reference_per_token_logprobs = Tensor([-11.3, -8.4, -2.0])
kl = per_token_logprobs - reference_per_token_logprobs
score_from_rm = 1.0
print(f"{kl=}")  # kl=tensor([-1.0000,  0.1000, -0.3000])
per_token_reward = kl.clone()
per_token_reward[-1] += score_from_rm  # assume last token is the EOS token
print(f"{per_token_reward=}")  # per_token_reward=tensor([-1.0000,  0.1000,  0.7000])
print(f"{score_from_rm=}")  # score_from_rm=1.0
print("#### Modeling each token as an action")
for action, reward in zip(response, per_token_reward):
    print(f"{action=}, {reward=}")
# action=tensor(4.), reward=tensor(-1.)
# action=tensor(5.), reward=tensor(0.1000)
# action=tensor(6.), reward=tensor(0.7000)
print("#### Modeling the entire response as an action")
entire_generation_reward = per_token_reward.sum()
print(f"action='entire completion', reward={entire_generation_reward}")
# action='entire completion', reward=-0.2000 (-1 + 0.1 + 0.7)

其次，RLOO 使用 REINFORCE 损失，它基本上将（奖励 - 基线）乘以动作的对数概率。在这里，我们强调每个令牌的 REINFORCE 损失和整个完成的 REINFORCE 损失之间的差异。请注意，对于 PPO 的损失，我们还需要根据价值模型使用广义优势估计 (GAE) 额外计算优势。

from torch import Tensor
response = Tensor([4., 5., 6.])
per_token_logprobs = Tensor([-12.3, -8.3, -2.3])
reference_per_token_logprobs = Tensor([-11.3, -8.4, -2.0])
kl = per_token_logprobs - reference_per_token_logprobs
score_from_rm = 1.0
print(f"{kl=}")  # kl=tensor([-1.0000,  0.1000, -0.3000])
per_token_reward = kl.clone()
per_token_reward[-1] += score_from_rm  # assume last token is the EOS token
print(f"{per_token_reward=}")  # per_token_reward=tensor([-1.0000,  0.1000,  0.7000])
print(f"{score_from_rm=}")  # score_from_rm=1.0
print("#### Modeling each token as an action")
for action, reward in zip(response, per_token_reward):
    print(f"{action=}, {reward=}")
# action=tensor(4.), reward=tensor(-1.)
# action=tensor(5.), reward=tensor(0.1000)
# action=tensor(6.), reward=tensor(0.7000)
print("#### Modeling the entire response as an action")
entire_generation_reward = per_token_reward.sum()
print(f"action='entire completion', reward={entire_generation_reward}")
# action='entire completion', reward=-0.2000 (-1 + 0.1 + 0.7)
baseline = Tensor([0.2, 0.3, 0.4])  # dummy baseline
print("#### Modeling each token as an action")
advantage = per_token_reward - baseline
per_token_reinforce_loss = per_token_logprobs * advantage
print(f"{advantage=}")  # advantage=tensor([-1.2000, -0.2000,  0.3000])
print(f"{per_token_reinforce_loss=}")  # per_token_reinforce_loss=tensor([14.7600,  1.6600, -0.6900])
print(f"{per_token_reinforce_loss.mean()=}")  # per_token_reinforce_loss.mean()=tensor(5.2433)

print("#### Modeling the entire response as an action")
advantage = entire_generation_reward - baseline.sum()
reinforce_loss = per_token_logprobs.sum() * advantage
print(f"{advantage=}")  # advantage=tensor(-1.1000)
print(f"{reinforce_loss=}")  # reinforce_loss=tensor(25.1900)

第三，RLOO 巧妙地计算基线。请注意，我们上面使用了一个虚拟基线。实际上，RLOO 使用批次中所有其他样本的奖励作为基线。下面是一个案例，我们有 3 个提示，每个提示有 4 个完成。我们通过平均同一提示下所有其他完成的奖励来计算每个完成的基线。

import torch
local_batch_size = 3
rloo_k = 4

rlhf_reward = torch.tensor([
    1, 2, 3, # first rlhf reward for three prompts
    2, 3, 4, # second rlhf reward for three prompts
    5, 6, 7, # third rlhf reward for three prompts
    8, 9, 10, # fourth rlhf reward for three prompts
]).float() # here we have 3 prompts which have 4 completions each

# slow impl
baseline = (rlhf_reward.sum(0) - rlhf_reward) / (rloo_k - 1)
advantages = torch.zeros_like(rlhf_reward)
for i in range(0, len(advantages), local_batch_size):
    other_response_rlhf_rewards = []
    for j in range(0, len(advantages), local_batch_size):
        if i != j:
            other_response_rlhf_rewards.append(rlhf_reward[j : j + local_batch_size])
    advantages[i : i + local_batch_size] = rlhf_reward[i : i + local_batch_size] - torch.stack(
        other_response_rlhf_rewards
    ).mean(0)
assert (1 - (2 + 5 + 8) / 3 - advantages[0].item()) < 1e-6
assert (6 - (3 + 2 + 9) / 3 - advantages[7].item()) < 1e-6

# vectorized impl
rlhf_reward = rlhf_reward.reshape(rloo_k, local_batch_size)
baseline = (rlhf_reward.sum(0) - rlhf_reward) / (rloo_k - 1)
vec_advantages = rlhf_reward - baseline
torch.testing.assert_close(vec_advantages.flatten(), advantages)

在此特别感谢 Arash Ahmadian，他提供了上述优势计算的矢量化实现。

开始使用 TRL 中的 RLOO

要开始使用 RLOO，您可以通过 pip install --upgrade trl 安装最新版本的 TRL 并导入 RLOOTrainer。下面是一个简短的片段，展示了一些高级 API 用法。欢迎查看文档

from transformers import (
    AutoModelForCausalLM,
    AutoModelForSequenceClassification,
    AutoTokenizer,
)

from trl.trainer.rloo_trainer import RLOOConfig, RLOOTrainer
from trl.trainer.utils import SIMPLE_QUERY_CHAT_TEMPLATE


base_model_name = "EleutherAI/pythia-1b-deduped"
tokenizer = AutoTokenizer.from_pretrained(base_model_name, padding_side="left")
tokenizer.add_special_tokens({"pad_token": "[PAD]"})
if tokenizer.chat_template is None:
    tokenizer.chat_template = SIMPLE_QUERY_CHAT_TEMPLATE
reward_model = AutoModelForSequenceClassification.from_pretrained(base_model_name, num_labels=1)
ref_policy = AutoModelForCausalLM.from_pretrained(base_model_name)
policy = AutoModelForCausalLM.from_pretrained(base_model_name)

train_dataset = ...  # make sure to have columns "input_ids"
eval_dataset = ...

trainer = RLOOTrainer(
    config=RLOOConfig(
        per_device_train_batch_size=1,
        gradient_accumulation_steps=64,
        total_episodes=30000,
    ),
    tokenizer=tokenizer,
    policy=policy,
    ref_policy=ref_policy,
    reward_model=reward_model,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
)
trainer.train()

这是一个跟踪权重和偏差实验的示例：https://wandb.ai/huggingface/trl/runs/dd2o3g35

在编码 RLOO 和 PPOv2 实现时，我们强调使其更容易提高模型开发的透明度。特别是，我们增强了文档，其中包含对记录指标的解释以及关于阅读和调试这些指标的入门指南。例如，我们建议在训练期间密切关注目标/rlhf_reward，这是 RLHF 训练的最终目标。

为了帮助可视化训练进度，我们定期从模型中记录一些样本完成。这是一个完成示例。在 Weights and Biases 的一个跟踪运行示例中（https://wandb.ai/huggingface/trl/runs/dd2o3g35），它看起来如下所示，允许您查看模型在训练不同阶段的响应。默认情况下，我们在训练期间生成 --num_sample_generations 10，但您可以自定义生成数量。

我们如何在 TRL 中实现 RLOO 训练器

我们基于新的实验性 PPOv2Trainer 实现了 RLOO 训练器，该训练器本身基于 https://arxiv.org/abs/2403.17031。有趣的是，我们实现的 RLOO 训练器仍然使用 PPO 损失。这是因为 REINFORCE 的损失是 PPO 的一个特例（https://arxiv.org/abs/2205.09123）。请注意，尽管对数概率在 REINFORCE 损失中明确出现，但它也隐式存在于 PPO 损失中。眼见为实，让我们用一个简单的例子来演示这一点。

import torch.nn.functional as F
from torch import LongTensor, Tensor, gather, no_grad

action = LongTensor([1])
advantage = Tensor([1.0])
logits = Tensor([[1.0, 2.0, 1.0, 1.0]])
logits.requires_grad = True
all_logprob = F.log_softmax(logits, dim=-1)
with no_grad():
    old_logprob = gather(all_logprob, 1, action.unsqueeze(-1)).squeeze(-1)
logprob = gather(all_logprob, 1, action.unsqueeze(-1)).squeeze(-1)
ratio = (logprob - old_logprob).exp()
ppo_loss = (ratio * advantage).mean() # [πθ(at | st) / πθ_old(at | st) * At]
# when the πθ and πθ_old are the same, the ratio is 1, and PPO's clipping has no effect
ppo_loss.backward()
print(f"{logits.grad=}")  # tensor([[-0.1749,  0.5246, -0.1749, -0.1749]])
logits2 = Tensor([[1.0, 2.0, 1.0, 1.0]])
logits2.requires_grad = True
all_logprob2 = F.log_softmax(logits2, dim=-1)
logprob2 = gather(all_logprob2, 1, action.unsqueeze(-1)).squeeze(-1)
reinforce_loss = logprob2 * advantage  # [log πθ(at | st) * At]
reinforce_loss.mean().backward()
print(f"{logits2.grad=}")  # tensor([[-0.1749,  0.5246, -0.1749, -0.1749]])

实验

为了验证 RLOO 实现是否有效，我们对 Pythia 1B 和 6.9B 模型进行了实验，并在此发布了训练好的检查点：

https://huggingface.co/collections/vwxyzjn/rloo-ppov2-tl-dr-summarize-checkpoints-66679a3bfd95ddf66c97420d

我们直接从 Huang et al., 2024 获取 SFT / RM 模型。为了评估，我们使用 vLLM 加载检查点，并使用 GPT4 作为判断模型来评估生成的 TL;DR 与参考 TL;DR。我们还查看了 GPU 内存使用情况和运行时，如博客文章开头的图所示。要重现我们的工作，请随时查看我们文档中的命令：

主要结果如下：

**🚀高性能 RLOO 检查点：**6.9B 检查点使用 GPT4 作为判断模型获得了 78.7% (k=2) 的偏好率，甚至超过了原始论文中报告的最佳性能 77.9% (k=4) 和 74.2% (k=2)。这是一个好迹象，表明我们的 RLOO 训练正在按预期工作。
- RLOO 1B 检查点的胜率为 40.1%，而 SFT 检查点的胜率为 21.3%。这是一个好迹象，表明 RLOO 训练正在按预期工作。
🤑**更少的 GPU 内存和更快的运行速度**：RLOO 训练使用更少的内存，运行速度更快，使其成为在线 RL 训练的非常有用的算法。

数值稳定性：阴暗面

尽管 RLOO 在性能和计算效率方面具有优势，但我们想强调一些数值问题。具体来说，在生成过程中获得的响应对数概率与在 bf16 下训练前向传播过程中获得的对数概率在数值上略有不同。这给 PPO 和 RLOO 都带来了问题，但对 RLOO 来说更糟，原因如下所述。

例如，假设我们为两个序列生成 10 个令牌。在 fp32 精度下，输出如下所示，其中 ratio = (forward_logprob - generation_logprob).exp()，这是 PPO 用于裁剪的。在第一个 epoch 和第一个 minibatch 下，比率应该完全相同，因为模型还没有进行任何更新。

generation_logprob=tensor([[    -0.1527,     -0.2258,     -3.5535,     -3.4805,     -0.0519,
             -2.3097,     -2.0275,     -0.4597,     -0.1687,     -0.0000],
        [    -0.1527,     -0.2258,     -5.2855,     -0.1686,     -8.4760,
             -4.3118,     -1.0368,     -0.8274,     -1.6342,     -2.6128]],
       device='cuda:0')
forward_logprob=tensor([[-0.1527, -0.2258, -3.5535, -3.4805, -0.0519, -2.3097, -2.0275, -0.4597,
         -0.1687],
        [-0.1527, -0.2258, -5.2855, -0.1686, -8.4760, -4.3118, -1.0368, -0.8274,
         -1.6342]], device='cuda:0', grad_fn=<SqueezeBackward1>)
ratio=tensor([[1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000],
        [1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000]],
       device='cuda:0', grad_fn=<ExpBackward0>)
ratio.mean()=0.9999998211860657
ratio.std()=6.592738373001339e-06
ratio.max()=1.0000133514404297
ratio.min()=0.9999887943267822

然而，在 bf16 下，我们得到

generation_logprob=tensor([[    -0.1426,     -0.1904,     -3.5938,     -3.4688,     -0.0618,
             -2.3906,     -2.0781,     -0.4375,     -0.1562,     -0.0000],
        [    -0.1426,     -0.1904,     -5.2812,     -0.1641,     -8.5625,
             -4.2812,     -1.0078,     -0.8398,     -1.5781,     -2.5781]],
       device='cuda:0', dtype=torch.bfloat16)
forward_logprob=tensor([[-0.1445, -0.1670, -3.5938, -3.5156, -0.0554, -2.2969, -1.9688, -0.5273,
         -0.1953],
        [-0.1445, -0.1670, -5.2812, -0.1533, -8.5625, -4.3125, -1.0000, -0.7852,
         -1.6641]], device='cuda:0', dtype=torch.bfloat16,
       grad_fn=<SqueezeBackward1>)
ratio=tensor([[1.0000, 0.9766, 1.0000, 1.0469, 0.9922, 0.9102, 0.8945, 1.0938, 1.0391],
        [1.0000, 0.9766, 1.0000, 0.9883, 1.0000, 1.0312, 0.9922, 0.9453, 1.0859]],
       device='cuda:0', dtype=torch.bfloat16, grad_fn=<ExpBackward0>)
ratio.mean()=1.0
ratio.std()=0.051025390625
ratio.max()=1.09375
ratio.min()=0.89453125

在 fp16 下，我们得到

generation_logprob=tensor([[    -0.1486,     -0.2212,     -3.5586,     -3.4688,     -0.0526,
             -2.3105,     -2.0254,     -0.4629,     -0.1677,     -0.0000],
        [    -0.1486,     -0.2212,     -5.2852,     -0.1681,     -8.4844,
             -4.3008,     -1.0322,     -0.8286,     -1.6348,     -2.6074]],
       device='cuda:0', dtype=torch.float16)
forward_logprob=tensor([[-0.1486, -0.2212, -3.5586, -3.4805, -0.0529, -2.3066, -2.0332, -0.4629,
         -0.1676],
        [-0.1486, -0.2212, -5.2852, -0.1682, -8.4766, -4.3008, -1.0322, -0.8281,
         -1.6299]], device='cuda:0', dtype=torch.float16,
       grad_fn=<SqueezeBackward1>)
ratio=tensor([[1.0000, 1.0000, 1.0000, 1.0117, 1.0000, 0.9961, 1.0078, 1.0000, 1.0000],
        [1.0000, 1.0000, 1.0000, 1.0000, 0.9922, 1.0000, 1.0000, 0.9995, 0.9951]],
       device='cuda:0', dtype=torch.float16, grad_fn=<ExpBackward0>)
ratio.mean()=1.0
ratio.std()=0.00418853759765625
ratio.max()=1.01171875
ratio.min()=0.9921875

请注意，bf16 的比率由于某种原因非常不稳定。当比率变大时，PPO 的裁剪系数 = 0.2 会生效，**使**比率大于 1.2 或小于 0.8 的令牌的梯度**归零**。对于 RLOO，这个问题更极端，因为我们看到的是 (forward_logprob.sum(1) - generation_logprob.sum(1)).exp() = [ 1.0625, 12.1875]，这意味着整个第二个序列的梯度都被归零了。

在实践中，我们注意到 PPO 会使大约 3% 的批次数据的梯度归零，而 RLOO 会使大约 20-40% 的批次数据归零。理论上，当不使用小批次时，RLOO 应该使 0% 的批次数据归零。重要的是，我们观察到，一旦我们增加生成新批次之前的梯度步数（通过 num_ppo_epochs 和 num_mini_batches），RLOO 的裁剪比率并没有显著变化；这提供了经验证据，表明裁剪比率确实是由于 bf16 的数值问题，而不是像论文中那样行为和最新策略显着不同。

要继续阅读有关最新问题更新，请随时查看 https://github.com/huggingface/transformers/issues/31267。

结论

TRL 中引入 RLOO（REINFORCE Leave One-Out）训练器是强化学习人类反馈（RLHF）在线训练中一个令人兴奋的算法，它为 PPO 提供了一种更易于访问且高效的替代方案。通过减少 GPU 内存使用和简化训练过程，RLOO 能够实现更大的批处理量和更快的训练时间。我们的实验表明，RLOO 在响应胜率方面与 PPO 具有竞争力，并且优于 DPO 检查点，使其成为有效在线 RLHF 的强大工具。探索我们的文档以开始使用！

致谢与感谢

感谢 Lewis Tunstall、Sara Hooker、Omar Sanseviero 和 Leandro Von Werra 对本博客文章提供的有益反馈。

更多博客文章

使用开放 LLM 的 Constitutional AI

作者： 2024年2月1日 • 15

使用直接偏好优化方法对 LLM 进行偏好调优

作者： 2024年1月18日 • 69

社区

通过拖放到文本输入框、粘贴或点击此处上传图片、音频和视频。

点击或粘贴此处以上传图片

· 注册或登录以评论