Hugging Face 的 LOGO
返回文章列表

在 RLHF 中重新引入強化學習 (RL)

釋出於 2024 年 6 月 12 日
在 GitHub 上更新
贊
99
Shengyi Costa Huang's avatar
黃聖毅 (Shengyi Costa Huang)
vwxyzjn
Arash Ahmadian's avatar
Arash Ahmadian
ArashAhmadian
訪客

我們很高興在 TRL 中推出 RLOO（REINFORCE Leave One-Out）訓練器。作為 PPO 的替代方案，RLOO 是一種新的線上 RLHF 訓練演算法，旨在更易於訪問和實現。特別地，**RLOO 需要更少的 GPU 記憶體，並且收斂時間更短。**如下圖所示：

  1. 🤑RLOO 的 vRAM 使用量比 PPO **減少約 50-70%**，具體取決於模型大小。
  2. 🚀RLOO 在 1B 模型上比 PPO **快 2 倍**，在 6.9B 模型上比 PPO **快 3 倍**。
  3. 🔥RLOO 在響應勝率方面（由 GPT4 評判）**與 PPO 具有競爭力**，並且始終優於 DPO 等流行的離線方法。

透過 RLOO，我們將強化學習重新引入 RLHF，使社群能夠更輕鬆地探索線上 RL 方法。這令人興奮，因為越來越多的研究表明線上 RL 比 DPO 等離線方法更有效（https://arxiv.org/abs/2402.04792，https://arxiv.org/abs/2405.08448）。

alt_text

alt_text

這篇部落格文章將解釋 RLOO 訓練器背後的動機、其工作原理以及如何在 TRL 中使用它。

動機

PPO 是一種有效的線上 RLHF 訓練演算法，用於訓練 GPT-4 等最先進的模型。然而，由於其高昂的 GPU 記憶體要求，PPO 在實際使用中可能非常具有挑戰性。特別是，PPO 需要將 4 份模型載入到記憶體中：1）策略模型，2）參考策略模型，3）獎勵模型，以及 4）價值模型，如下圖所示。PPO 還有許多微妙的實現細節，可能很難做到正確（Engstrom 等人；2020，Huang 等人 2022）。

alt_text

在 Cohere 的一篇新論文中，Ahmadian 等人 (2024) 回顧了 RLHF 訓練的基礎，並提出了一種更優雅的方法，稱為 RLOO，一種新的線上訓練演算法。RLOO 只需要將 3 份模型載入到記憶體中：1）策略模型，2）參考策略模型，以及 3）獎勵模型，如上圖所示。

重要的是，RLOO 需要更少的記憶體，這意味著它更容易

  1. 在沒有 OOM（記憶體不足錯誤）的情況下執行
  2. 能夠載入更大的批次大小
  3. 執行更高效、更快。

此外，RLOO 將整個完成令牌建模為單個動作，如下圖所示。在下一節中，我們將結合程式碼片段深入探討更多細節。

alt_text

RLOO 的工作原理

RLOO 和 PPO 都有幾個共享步驟

  1. 策略模型將生成一些完成令牌，並獲得當前策略和參考策略下的每個令牌的對數機率。

  2. 然後，我們計算每個令牌的 KL 懲罰，即當前策略和參考策略下對數機率之間的差異。

  3. 然後，我們從獎勵模型中獲取整個完成的得分。

從這裡開始，常規 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

alt_text

在編碼 RLOO 和 PPOv2 實現時，我們強調使其更容易提高模型開發的透明度。特別是，我們增強了文件，其中包含對記錄指標的解釋以及關於閱讀和除錯這些指標的入門指南。例如，我們建議在訓練期間密切關注目標/rlhf_reward，這是 RLHF 訓練的最終目標。

alt_text

alt_text

為了幫助視覺化訓練進度，我們定期從模型中記錄一些樣本完成。這是一個完成示例。在 Weights and Biases 的一個跟蹤執行示例中（https://wandb.ai/huggingface/trl/runs/dd2o3g35），它看起來如下所示，允許您檢視模型在訓練不同階段的響應。預設情況下，我們在訓練期間生成 --num_sample_generations 10，但您可以自定義生成數量。

alt_text

我們如何在 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 模型進行了實驗，並在此釋出了訓練好的檢查點：

我們直接從 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

作者：  vwxyzjn 2024年2月1日 • 15

使用直接偏好最佳化方法對 LLM 進行偏好調優

作者：  kashif 2024年1月18日 • 69

社群

註冊或 登入以評論

贊
99
© . This site is unofficial and not affiliated with Hugging Face, Inc.