使用 PPO 的 RLHF 的 N 个实现细节

发布于 2023 年 10 月 24 日

在 GitHub 上更新

黄圣毅 (Shengyi Costa Huang)

访客

RLHF/ChatGPT 最近是一个热门研究课题。为了深入研究 RLHF，本博客文章尝试重现 OpenAI 2019 年的原始 RLHF 代码库 openai/lm-human-preferences。尽管是“tensorflow-1.x-ness”，OpenAI 的原始代码库经过了充分评估和基准测试，使其成为研究 RLHF 实现工程细节的良好起点。

我们旨在

重现 OAI 在风格任务中的结果，并匹配 openai/lm-human-preferences 的学习曲线。
提供一份实现细节清单，类似于 近端策略优化 (PPO) 的 37 个实现细节；无需痛苦地调试强化学习的精神。
提供一个易于阅读且最小的 RLHF 参考实现；

这项工作仅用于教育/学习目的。对于需要更多功能的高级用户，例如使用 PEFT 运行更大的模型，huggingface/trl 将是一个不错的选择。

在匹配学习曲线中，我们展示了我们的主要贡献：创建一个能够重现 OAI 在风格任务中结果的代码库，并与 openai/lm-human-preferences 的学习曲线非常接近。
然后，我们将深入探讨与重现 OAI 工作相关的实现细节。在通用实现细节中，我们讨论了基本细节，例如奖励/值的生成方式以及响应的生成方式。在奖励模型实现细节中，我们讨论了奖励归一化等细节。在策略训练实现细节中，我们讨论了拒绝采样和奖励“白化”等细节。
- 在PyTorch Adam 优化器在 RLHF 方面的数值问题中，我们强调了 TensorFlow 和 PyTorch 之间 Adam 实现的一个非常有趣的差异，这导致模型训练中的激进更新。
接下来，我们将检查在奖励标签由 gpt2-large 生成的情况下，训练不同基础模型（例如，gpt2-xl、falcon-1b）的效果。
最后，我们总结了我们的工作，并讨论了其局限性。

以下是重要链接

💾 我们的复现代码库 https://github.com/vwxyzjn/lm-human-preference-details
🤗 RLHF 模型比较演示： https://huggingface.co/spaces/lm-human-preference-details/rlhf-demo
🐝 所有 w&b 训练日志 https://wandb.ai/openrlbenchmark/lm_human_preference_details

匹配学习曲线

我们的主要贡献是重现 OAI 在风格任务（例如情感和描述性）中的结果。如下图所示，我们的代码库（橙色曲线）可以生成与 OAI 代码库（蓝色曲线）几乎相同的学习曲线。

关于运行 openai/lm-human-preferences 的注意事项

为了进行直接比较，我们运行了 openai/lm-human-preferences 中的原始 RLHF 代码，它将提供有价值的指标来帮助验证和诊断我们的复现。我们能够设置原始的 TensorFlow 1.x 代码，但这需要一个超特定的设置：

OAI 的数据集部分损坏/丢失（因此我们用类似的 HF 数据集替换它们，这可能会或可能不会导致性能差异）
- 具体来说，它的图书数据集在 OpenAI 的 GCP - Azure 迁移期间丢失了 (https://github.com/openai/lm-human-preferences/issues/17#issuecomment-1044051496)。我用 Hugging Face 的 bookcorpus 数据集替换了图书数据集，这原则上就是 OAI 使用的。
它不能在 1 个 V100 上运行，因为它没有实现梯度累积。相反，它使用大批量大小并将批量分散到 8 个 GPU 上，并且在仅 1 个 GPU 上会 OOM。
它不能在 8 个 A100 上运行，因为它使用 TensorFlow 1.x，这与 Cuda 8+ 不兼容
它不能在 8 个 V100 (16GB) 上运行，因为它会 OOM
它只能在 8 个 V100 (32GB) 上运行，这只能由 AWS 作为 p3dn.24xlarge 实例提供。

通用实现细节

现在我们将深入探讨与重现 OAI 工作相关的实现细节。在本节中，我们将讨论基本细节，例如奖励/值的生成方式以及响应的生成方式。这些细节没有特定的顺序：

奖励模型和策略的值头将 query 和 response 的连接作为输入
1. 奖励模型和策略的值头*不*只关注响应。相反，它将 query 和 response 连接在一起作为 query_response (lm_human_preferences/rewards.py#L105-L107)。
2. 因此，例如，如果 query = "他安静了一分钟，眼睛无法解读"，并且 response = "他看着左手，那只手握着他伸出的手臂。"，那么奖励模型和策略的值会对 query_response = "他安静了一分钟，眼睛无法解读。他看着左手，那只手握着他伸出的手臂。" 进行前向传播，并生成形状为 (B, T, 1) 的奖励和值，其中 B 是批量大小，T 是序列长度，1 是奖励头维度 (lm_human_preferences/rewards.py#L105-L107, lm_human_preferences/policy.py#L111)。
3. T 表示每个 token 都与其之前的上下文关联着一个奖励。例如，`eyes` token 将对应于 `he was quiet for a minute, his eyes` 的奖励。
使用特殊的填充 token 进行填充并截断输入。
1. OAI 为查询 query_length 设置了固定的输入长度；它会用 pad_token 填充过短的序列 (lm_human_preferences/language/datasets.py#L66-L67)，并截断过长的序列 (lm_human_preferences/language/datasets.py#L57)。有关概念的通用介绍，请参阅此处。在填充输入时，OAI 使用了一个词汇表之外的 token (lm_human_preferences/language/encodings.py#L56)。
  1. 关于 HF 的 transformer——填充 token。根据 (transformers#2630#issuecomment-578159876)，GPT 和 GPT-2 的预训练过程中未使用填充 token；因此，transformer 的 gpt2 模型没有与其 tokenizer 关联的官方填充 token。常见的做法是设置 tokenizer.pad_token = tokenizer.eos_token，但在本工作中，我们将区分这两个特殊 token 以匹配 OAI 的原始设置，因此我们将使用 tokenizer.add_special_tokens({"pad_token": "[PAD]"})。
  请注意，没有填充 token 是解码器模型的默认设置，因为它们在预训练期间使用“打包”进行训练，这意味着许多序列被连接并由 EOS token 分隔，并且这些序列中总是具有最大长度的块在预训练期间被送入模型。
2. 将所有内容整合在一起，这是一个示例：
```
import transformers
tokenizer = transformers.AutoTokenizer.from_pretrained("gpt2", padding_side="right")
tokenizer.add_special_tokens({"pad_token": "[PAD]"})
query_length = 5
texts = [
    "usually, he would",
    "she thought about it",
]    
tokens = []
for text in texts:
    tokens.append(tokenizer.encode(text)[:query_length])

print("tokens", tokens)
inputs = tokenizer.pad(
    {"input_ids": tokens},
    padding="max_length",
    max_length=query_length,
    return_tensors="pt",
    return_attention_mask=True,
)
print("inputs", inputs)

"""prints are
tokens [[23073, 11, 339, 561], [7091, 1807, 546, 340]]
inputs {'input_ids': tensor([[23073,    11,   339,   561, 50257],
        [ 7091,  1807,   546,   340, 50257]]), 'attention_mask': tensor([[1, 1, 1, 1, 0],
        [1, 1, 1, 1, 0]])}
"""
```

相应地调整填充 token 的位置索引

在计算 logits 时，OAI 的代码通过正确地屏蔽填充标记来工作。这是通过找出与填充标记对应的标记索引 (lm_human_preferences/language/model.py#L296-L297)，然后相应地调整它们的位置索引 (lm_human_preferences/language/model.py#L320) 来实现的。

例如，如果 query=[23073, 50259, 50259] 和 response=[11, 339, 561]，其中 (50259 是 OAI 的填充 token)，它会创建位置索引为 [[0 1 1 1 2 3]] 和 logits 如下。请注意，对应于填充 token 的 logits 保持不变！这就是我们应该在重现中实现的效果。

all_logits [[[ -35.28693   -34.2875    -38.16074  ...  -41.595802  -41.082108
    -35.36577 ]
  [ -35.28693   -34.2875    -38.16074  ...  -41.595802  -41.082108
    -35.36577 ]
  [ -35.28693   -34.2875    -38.16074  ...  -41.595802  -41.082108
    -35.36577 ]
  [-111.303955 -110.94471  -112.90624  ... -113.13064  -113.7788
   -109.17345 ]
  [-111.51512  -109.61077  -114.90231  ... -118.43514  -111.56671
   -112.12478 ]
  [-122.69775  -121.84468  -128.27417  ... -132.28055  -130.39604
   -125.707756]]] (1, 6, 50257)

关于 HF Transformers 的注意事项 — position_ids 和 padding_side。我们可以使用 Hugging Face 的 Transformer，通过 1) 左填充和 2) 传入适当的 position_ids 来精确复现 logits。

import torch
import transformers
tokenizer = transformers.AutoTokenizer.from_pretrained("gpt2", padding_side="right")
tokenizer.add_special_tokens({"pad_token": "[PAD]"})
pad_id = tokenizer.pad_token_id
query = torch.tensor([
    [pad_id, pad_id, 23073],
])
response = torch.tensor([
    [11, 339, 561],
])
temperature = 1.0

query = torch.tensor(query)
response = torch.tensor(response).long()
context_length = query.shape[1]
query_response = torch.cat((query, response), 1)
pretrained_model = transformers.AutoModelForCausalLM.from_pretrained("gpt2")
def forward(policy, query_responses, tokenizer):
    attention_mask = query_responses != tokenizer.pad_token_id
    position_ids = attention_mask.cumsum(1) - attention_mask.long()  # exclusive cumsum
    input_ids = query_responses.clone()
    input_ids[~attention_mask] = 0
    return policy(
        input_ids=input_ids,
        attention_mask=attention_mask,
        position_ids=position_ids,
        return_dict=True,
        output_hidden_states=True,
    )
output = forward(pretrained_model, query_response, tokenizer)
logits = output.logits
logits /= temperature
print(logits)

"""
tensor([[[ -26.9395,  -26.4709,  -30.0456,  ...,  -33.2208,  -33.2884,
           -27.4360],
         [ -27.1677,  -26.7330,  -30.2386,  ...,  -33.6813,  -33.6931,
           -27.5928],
         [ -35.2869,  -34.2875,  -38.1608,  ...,  -41.5958,  -41.0821,
           -35.3658],
         [-111.3040, -110.9447, -112.9062,  ..., -113.1306, -113.7788,
          -109.1734],
         [-111.5152, -109.6108, -114.9024,  ..., -118.4352, -111.5668,
          -112.1248],
         [-122.6978, -121.8447, -128.2742,  ..., -132.2805, -130.3961,
          -125.7078]]], grad_fn=<DivBackward0>)
"""

关于 HF Transformers — generate 期间的 position_ids：在生成期间，我们不应该传入 position_ids，因为 transformers 中已经调整了 position_ids（参见 huggingface/transformers#/7552）。

通常，我们在 transformers 中几乎从不传入 position_ids。所有的掩码和移位逻辑都已在例如 generate 函数中实现（需要永久代码链接）。

响应生成采样固定长度的响应，不带填充。

在生成响应时，OAI 使用 top_k=0, top_p=1.0，并且只对词汇表进行分类采样 (lm_human_preferences/language/sample.py#L43)，代码会持续采样直到生成固定长度的响应 (lm_human_preferences/policy.py#L103)。值得注意的是，即使遇到 EOS（序列结束）token，它也会继续采样。

关于 HF Transformers 的注意事项 — 采样可能在 eos_token 处停止： 在 transformers 中，生成可能会在 eos_token 处停止 (src/transformers/generation/utils.py#L2248-L2256)，这与 OAI 的设置不同。为了保持设置一致，我们需要将 pretrained_model.generation_config.eos_token_id = None, pretrained_model.generation_config.pad_token_id = None。请注意，transformers.GenerationConfig(eos_token_id=None, pad_token_id=None, ...) 无效，因为 pretrained_model.generation_config 会覆盖并设置 eos_token。

import torch
import transformers
tokenizer = transformers.AutoTokenizer.from_pretrained("gpt2", padding_side="right")
tokenizer.add_special_tokens({"pad_token": "[PAD]"})
pad_id = tokenizer.pad_token_id
query = torch.tensor([
    [pad_id, pad_id, 23073],
])
response = torch.tensor([
    [11, 339, 561],
])
response_length = 4
temperature = 0.7
pretrained_model = transformers.AutoModelForCausalLM.from_pretrained("gpt2")
pretrained_model.generation_config.eos_token_id = None # disable `pad_token_id` and `eos_token_id` because we just want to
pretrained_model.generation_config.pad_token_id = None  # generate tokens without truncation / padding
generation_config = transformers.GenerationConfig(
    max_new_tokens=response_length,
    min_new_tokens=response_length,
    temperature=temperature,
    top_k=0.0,
    top_p=1.0,
    do_sample=True,
)
context_length = query.shape[1]
attention_mask = query != tokenizer.pad_token_id
input_ids = query.clone()
input_ids[~attention_mask] = 0  # set padding tokens to 0
output = pretrained_model.generate(
    input_ids=input_ids,
    attention_mask=attention_mask,
    # position_ids=attention_mask.cumsum(1) - attention_mask.long(), # generation collapsed if this was turned on.
    generation_config=generation_config,
    return_dict_in_generate=True,
)
print(output.sequences)

"""
tensor([[    0,     0, 23073, 16851,    11,   475,   991]])
"""

请注意，在较新的代码库 https://github.com/openai/summarize-from-feedback 中，OAI 在遇到 EOS 标记时确实停止采样 (summarize_from_feedback/utils/experiment_helpers.py#L19)。然而，在这项工作中，我们旨在进行 1:1 的复现，因此我们保持了即使遇到 EOS 标记也可以继续采样的设置。

奖励模型和策略训练的学习率退火。
1. 正如 Ziegler 等人（2019）所建议的，奖励模型只训练一个 epoch，以避免过度拟合有限的人类标注数据（例如，descriptiveness 任务只有大约 5000 个标签）。在这个单个 epoch 期间，学习率会退火到零 (lm_human_preferences/train_reward.py#L249)。
2. 与奖励模型训练类似，学习率被退火到零 (lm_human_preferences/train_policy.py#L172-L173)。
对不同进程使用不同的种子
1. 在生成 8 个 GPU 进程以进行数据并行时，OAI 为每个进程设置了不同的随机种子 (lm_human_preferences/utils/core.py#L108-L111)。在实现方面，这是通过 local_seed = args.seed + process_rank * 100003 完成的。例如，该种子将使模型产生不同的响应并获得不同的分数。
  1. 注意：我发现数据集洗牌有一个 bug——数据集由于某种原因使用相同的种子进行洗牌 (lm_human_preferences/lm_tasks.py#L94-L97)。

奖励模型实现细节

在本节中，我们将讨论奖励模型特定的实现细节。我们将讨论奖励归一化和层初始化等细节。这些细节没有特定的顺序：

奖励模型仅在最后一个 token 处输出值。
1. 请注意，在对 query 和 response 的拼接进行前向传播后获得的奖励将具有形状 (B, T, 1)，其中 B 是批量大小，T 是序列长度（始终相同；在 OAI 的风格任务设置中为 query_length + response_length = 64 + 24 = 88，参见 launch.py#L9-L11），1 是奖励头维度为 1。出于 RLHF 的目的，原始代码库提取了最后一个 token 的奖励 (lm_human_preferences/rewards.py#L132)，因此奖励将只有形状 (B, 1)。
2. 请注意，在较新的代码库 openai/summarize-from-feedback 中，OAI 在遇到 EOS 标记时停止采样 (summarize_from_feedback/utils/experiment_helpers.py#L19)。在提取奖励时，它会识别 last_response_index，即 EOS 标记之前的索引 (#L11-L13)，并提取该索引处的奖励 (summarize_from_feedback/reward_model.py#L59)。然而，在这项工作中，我们只坚持原始设置。
奖励头层初始化
1. 奖励头的权重根据 $\mathcal{N}\left(0,1 /\left(\sqrt{d_{\text {model }}+1}\right)\right)$ 初始化 (lm_human_preferences/language/model.py#L368, lm_human_preferences/language/model.py#L251-L252)。这与 Stiennon 等人 (2020) 的设置一致 (summarize_from-feedback/query_response_model.py#L106-L107)（附注：Stiennon 等人 (2020) 在第 17 页有一个笔误，称分布为 $\mathcal{N}\left(0,1 /\left(d_{\text {model }}+1\right)\right)$ ，没有平方根）。
2. 奖励头的偏差设置为 0 (lm_human_preferences/language/model.py#L254)。
奖励模型归一化（前后）
1. 在论文中，Ziegler 等人 (2019) 提到：“为了保持奖励模型在训练过程中的尺度一致，我们对其进行归一化，使其对于 $x \sim \mathcal{D}, y \sim \rho(·|x)$ 。”为了执行归一化过程，代码首先创建 reward_gain 和 reward_bias，这样奖励可以通过 reward = reward * reward_gain + reward_bias 计算 (lm_human_preferences/rewards.py#L50-L51)。
2. 在执行归一化过程时，代码首先设置 reward_gain=1, reward_bias=0 (lm_human_preferences/train_reward.py#L211)，然后从目标数据集（例如 bookcorpus, tldr, cnndm）收集采样查询、完整响应和评估奖励。然后它获取评估奖励的经验平均值和标准差 (lm_human_preferences/train_reward.py#L162-L167)，并尝试计算 reward_gain 和 reward_bias 应该是什么。
3. 我们用 $\mu_{\mathcal{D}}$ 表示经验均值， $\sigma_{\mathcal{D}}$ 表示经验标准差， $g$ 表示 reward_gain， $b$ 表示 reward_bias， $\mu_{\mathcal{T}} = 0$ 目标均值， $\sigma_{\mathcal{T}}=1$ 目标标准差。那么我们有以下公式。 $\begin{aligned}g*\mathcal{N}(\mu_{\mathcal{D}}, \sigma_{\mathcal{D}}) + b &= \mathcal{N}(g*\mu_{\mathcal{D}}, g*\sigma_{\mathcal{D}}) + b\\&= \mathcal{N}(g*\mu_{\mathcal{D}} + b, g*\sigma_{\mathcal{D}}) \\&= \mathcal{N}(\mu_{\mathcal{T}}, \sigma_{\mathcal{T}}) \\g &= \frac{\sigma_{\mathcal{T}}}{\sigma_{\mathcal{D}}} \\b &= \mu_{\mathcal{T}} - g*\mu_{\mathcal{D}}\end{aligned}$
4. 归一化过程在奖励模型训练之前和之后都进行应用 (lm_human_preferences/train_reward.py#L232-L234, lm_human_preferences/train_reward.py#L252-L254)。
5. 请注意，我们为了归一化而生成的响应 $y \sim \rho(·|x)$ 来自预训练语言模型 $\rho$ 。模型 $\rho$ 被固定为参考，并且在奖励学习中不进行更新 (lm_human_preferences/train_reward.py#L286C1-L286C31)。

策略训练实现细节

在本节中，我们将深入探讨层初始化、数据后处理和 dropout 设置等细节。我们还将探讨拒绝采样、奖励“白化”和自适应 KL 等技术。这些细节没有特定的顺序：

通过采样温度对 logits 进行缩放。
1. 在计算响应的对数概率时，模型首先输出响应中 tokens 的 logits，然后将 logits 除以采样温度 (lm_human_preferences/policy.py#L121)。即，logits /= self.temperature
2. 在非正式测试中，我们发现如果没有这种缩放，KL 会比预期上升得更快，并且性能会下降。
值头层初始化
1. 价值头的权重根据 $\mathcal{N}\left(0,0\right)$ 初始化 (lm_human_preferences/language/model.py#L368, lm_human_preferences/language/model.py#L251-L252)。这是
2. 奖励头的偏差设置为 0 (lm_human_preferences/language/model.py#L254)。
选择以句点开头和结尾的查询文本
1. 这是数据预处理的一部分；
  1. 尝试只选择 start_text="." 之后的文本 (lm_human_preferences/language/datasets.py#L51)
  2. 尝试选择紧邻 end_text="." 之前的文本 (lm_human_preferences/language/datasets.py#L61)
  3. 然后填充文本 (lm_human_preferences/language/datasets.py#L66-L67)
2. 当运行 openai/lm-human-preferences 时，OAI 的数据集部分损坏/丢失 (openai/lm-human-preferences/issues/17#issuecomment-104405149)，所以我们不得不将其替换为类似的 HF 数据集，这可能会或可能不会导致性能差异）
3. 对于图书数据集，我们使用了 https://huggingface.co/datasets/bookcorpus，我们发现没有必要提取以句号开头和结尾的句子，因为该数据集已经以这种方式预处理过（例如，“通常，他会在客厅里跑来跑去，玩着他的玩具。”）。为此，我们为 sentiment 和 descriptiveness 任务设置了 start_text=None, end_text=None。
禁用 Dropout
1. Ziegler 等人 (2019) 建议：“我们不使用 dropout 进行策略训练。” 这也在代码中完成 (lm_human_preferences/policy.py#L48)。
拒绝采样
1. Ziegler 等人 (2019) 建议：“我们使用拒绝采样来确保在标记 16 和 24 之间有一个句号，然后在此句号处截断（这是对‘句子结束’的粗略近似。我们选择它是因为它易于集成到强化学习循环中，即使是粗略的近似也足以达到使人类评估任务更容易的目的）。在强化学习微调过程中，我们通过给予没有此类句号的延续固定奖励 -1 来惩罚它们。”
2. 具体来说，这是通过以下步骤实现的：
  1. Token 截断：我们希望在响应中位于位置 truncate_after 或之后首次出现的 truncate_token 处进行截断 (lm_human_preferences/train_policy.py#L378)
    1. 代码注释：“核心示例：将 truncate_token 后所有 token 替换为 padding_token”
  2. 在截断响应上运行奖励模型：在响应被 token 截断过程截断后，代码接着在截断响应上运行奖励模型。
  3. 拒绝采样：如果 token 16 和 24 之间没有句点，则将响应的分数替换为固定的低值（例如 -1）(lm_human_preferences/train_policy.py#L384, lm_human_preferences/train_policy.py#L384-L402)
    1. 代码注释：“核心示例：确保样本包含 truncate_token"
    2. 代码注释：“只对通过该函数的人类响应进行查询”
  4. 举例说明在 descriptiveness 中的一些例子：
    
    从我们的复现中提取的样本 https://wandb.ai/openrlbenchmark/lm_human_preference_details/runs/djf8yymv/logs。请注意，第一个和第三个示例在句号后有太多 token，因此其分数被替换为 -1。
折扣因子 = 1
1. 折扣参数 $\gamma$ 设置为 1 (lm_human_preferences/train_policy.py#L56)，这意味着未来奖励与即时奖励具有相同的权重。

训练循环术语：PPO 中的批次和微批次

OAI 使用以下训练循环 (lm_human_preferences/train_policy.py#L184-L192)。注意：我们额外添加了 micro_batch_size 以帮助处理梯度累积的情况。在每个 epoch 中，它会打乱批次索引。

import numpy as np
batch_size = 8
nminibatches = 2
gradient_accumulation_steps = 2
mini_batch_size = batch_size // nminibatches
micro_batch_size = mini_batch_size // gradient_accumulation_steps
data = np.arange(batch_size).astype(np.float32)
print("data:", data)
print("batch_size:", batch_size)
print("mini_batch_size:", mini_batch_size)
print("micro_batch_size:", micro_batch_size)
for epoch in range(4):
    batch_inds = np.random.permutation(batch_size)
    print("epoch:", epoch, "batch_inds:", batch_inds)
    for mini_batch_start in range(0, batch_size, mini_batch_size):
        mini_batch_end = mini_batch_start + mini_batch_size
        mini_batch_inds = batch_inds[mini_batch_start:mini_batch_end]
        
        # `optimizer.zero_grad()` set optimizer to zero for gradient accumulation
        for micro_batch_start in range(0, mini_batch_size, micro_batch_size):
            micro_batch_end = micro_batch_start + micro_batch_size 
            micro_batch_inds = mini_batch_inds[micro_batch_start:micro_batch_end]
            print("____⏩ a forward pass on", data[micro_batch_inds])
        # `optimizer.step()`
        print("⏪ a backward pass on", data[mini_batch_inds])

# data: [0. 1. 2. 3. 4. 5. 6. 7.]
# batch_size: 8
# mini_batch_size: 4
# micro_batch_size: 2
# epoch: 0 batch_inds: [6 4 0 7 3 5 1 2]
# ____⏩ a forward pass on [6. 4.]
# ____⏩ a forward pass on [0. 7.]
# ⏪ a backward pass on [6. 4. 0. 7.]
# ____⏩ a forward pass on [3. 5.]
# ____⏩ a forward pass on [1. 2.]
# ⏪ a backward pass on [3. 5. 1. 2.]
# epoch: 1 batch_inds: [6 7 3 2 0 4 5 1]
# ____⏩ a forward pass on [6. 7.]
# ____⏩ a forward pass on [3. 2.]
# ⏪ a backward pass on [6. 7. 3. 2.]
# ____⏩ a forward pass on [0. 4.]
# ____⏩ a forward pass on [5. 1.]
# ⏪ a backward pass on [0. 4. 5. 1.]
# epoch: 2 batch_inds: [1 4 5 6 0 7 3 2]
# ____⏩ a forward pass on [1. 4.]
# ____⏩ a forward pass on [5. 6.]
# ⏪ a backward pass on [1. 4. 5. 6.]
# ____⏩ a forward pass on [0. 7.]
# ____⏩ a forward pass on [3. 2.]
# ⏪ a backward pass on [0. 7. 3. 2.]
# epoch: 3 batch_inds: [7 2 4 1 3 0 6 5]
# ____⏩ a forward pass on [7. 2.]
# ____⏩ a forward pass on [4. 1.]
# ⏪ a backward pass on [7. 2. 4. 1.]
# ____⏩ a forward pass on [3. 0.]
# ____⏩ a forward pass on [6. 5.]
# ⏪ a backward pass on [3. 0. 6. 5.]

每个 token 的 KL 惩罚
- 代码添加了每个 token 的 KL 惩罚 (lm_human_preferences/train_policy.py#L150-L153) 到奖励中，以阻止策略与原始策略差异过大。
- 以 "usually, he would" 为例，它被 token 化为 [23073, 11, 339, 561]。假设我们将 [23073] 用作查询，[11, 339, 561] 用作响应。那么在默认的 gpt2 参数下，响应 token 的参考策略的对数概率为 logprobs=[-3.3213, -4.9980, -3.8690]。
  - 在第一次 PPO 更新 epoch 和 minibatch 更新期间，活跃策略将具有相同的对数概率 new_logprobs=[-3.3213, -4.9980, -3.8690]。因此，每个 token 的 KL 惩罚将是 kl = new_logprobs - logprobs = [0., 0., 0.,]。
  - 然而，在第一次梯度反向传播之后，我们可能会得到 new_logprob=[-3.6528, -5.0406, -3.2339]，因此每个 token 的 KL 惩罚变为 kl = new_logprobs - logprobs = [-0.3315, -0.0426, 0.6351]
  - 然后 non_score_reward = beta * kl，其中 beta 是 KL 惩罚系数 $\beta$ ，它被添加到从奖励模型获得的 score 中，以创建用于训练的 rewards。score 仅在回合结束时给出；它可能看起来像 [0.4,]，我们有 rewards = [beta * -0.3315, beta * -0.0426, beta * 0.6351 + 0.4]。

每个迷你批次的奖励和优势白化，可选地进行均值平移

OAI 实现了一个 whiten 函数，其工作方式如下，通过减去均值然后除以标准差来归一化 values。可选地，whiten 可以通过 shift_mean=True 将白化后的 values 的均值移回。

def whiten(values, shift_mean=True):
    mean, var = torch.mean(values), torch.var(values, unbiased=False)
    whitened = (values - mean) * torch.rsqrt(var + 1e-8)
    if not shift_mean:
        whitened += mean
    return whitened

在每个迷你批次中，OAI 会白化奖励 whiten(rewards, shift_mean=False)，不进行均值平移 (lm_human_preferences/train_policy.py#L325)，并白化优势 whiten(advantages)，进行均值平移 (lm_human_preferences/train_policy.py#L338)。
优化注意： 如果迷你批次数量为一个（本复现中就是这种情况），我们只需要对奖励进行一次白化，并计算和白化优势一次，因为它们的值不会改变。

TensorFlow 与 PyTorch 注： tf.moments 与 torch.var 的行为不同：由于方差计算方式不同，torch 和 tf 中白化的行为也不同。

import numpy as np
import tensorflow as tf
import torch

def whiten_tf(values, shift_mean=True):
    mean, var = tf.nn.moments(values, axes=list(range(values.shape.rank)))
    mean = tf.Print(mean, [mean], 'mean', summarize=100)
    var = tf.Print(var, [var], 'var', summarize=100)
    whitened = (values - mean) * tf.rsqrt(var + 1e-8)
    if not shift_mean:
        whitened += mean
    return whitened

def whiten_pt(values, shift_mean=True, unbiased=True):
    mean, var = torch.mean(values), torch.var(values, unbiased=unbiased)
    print("mean", mean)
    print("var", var)
    whitened = (values - mean) * torch.rsqrt(var + 1e-8)
    if not shift_mean:
        whitened += mean
    return whitened

rewards = np.array([
    [1.2, 1.3, 1.4],
    [1.5, 1.6, 1.7],
    [1.8, 1.9, 2.0],
])

with tf.Session() as sess:
    print(sess.run(whiten_tf(tf.constant(rewards, dtype=tf.float32), shift_mean=False)))
    print(whiten_pt(torch.tensor(rewards), shift_mean=False, unbiased=True))
    print(whiten_pt(torch.tensor(rewards), shift_mean=False, unbiased=False))

mean[1.5999999]
var[0.0666666627]
[[0.05080712 0.4381051  0.8254035 ]
 [1.2127019  1.6000004  1.9872988 ]
 [2.3745968  2.7618952  3.1491938 ]]
mean tensor(1.6000, dtype=torch.float64)
var tensor(0.0750, dtype=torch.float64)
tensor([[0.1394, 0.5046, 0.8697],
        [1.2349, 1.6000, 1.9651],
        [2.3303, 2.6954, 3.0606]], dtype=torch.float64)
mean tensor(1.6000, dtype=torch.float64)
var tensor(0.0667, dtype=torch.float64)
tensor([[0.0508, 0.4381, 0.8254],
        [1.2127, 1.6000, 1.9873],
        [2.3746, 2.7619, 3.1492]], dtype=torch.float64)

裁剪值函数
1. 如原始 PPO 中所做 (baselines/ppo2/model.py#L68-L75)，值函数被裁剪 (lm_human_preferences/train_policy.py#L343-L348)，方式与策略目标类似。
自适应 KL
- KL 散度惩罚系数 $\beta$ 根据当前策略和先前策略之间的 KL 散度进行自适应调整。如果 KL 散度超出预定义的目标范围，则调整惩罚系数以使其更接近目标范围 (lm_human_preferences/train_policy.py#L115-L124)。其实现如下：
```
class AdaptiveKLController:
    def __init__(self, init_kl_coef, hparams):
        self.value = init_kl_coef
        self.hparams = hparams

    def update(self, current, n_steps):
        target = self.hparams.target
        proportional_error = np.clip(current / target - 1, -0.2, 0.2)
        mult = 1 + proportional_error * n_steps / self.hparams.horizon
        self.value *= mult
```
- 对于本工作中检查的 sentiment 和 descriptiveness 任务，我们有 init_kl_coef=0.15, hparams.target=6, hparams.horizon=10000。

PyTorch Adam 优化器在 RLHF 方面的数值问题

这个实现细节非常有趣，值得专门一节来讨论。
PyTorch Adam 优化器 (torch.optim.Adam.html) 的实现与 TensorFlow 的 Adam 优化器（TF1 Adam 在 tensorflow/v1.15.2/adam.py，TF2 Adam 在 keras/adam.py#L26-L220）不同。特别是，PyTorch 遵循 Kingma 和 Ba 的 Adam 论文中的算法 1 (arxiv/1412.6980)，而 TensorFlow 使用该论文 2.1 节之前的公式，并且其此处所指的 epsilon 是论文中的 epsilon hat。在伪代码比较中，我们有以下内容：

### pytorch adam implementation:
bias_correction1 = 1 - beta1 ** step
bias_correction2 = 1 - beta2 ** step
step_size = lr / bias_correction1
bias_correction2_sqrt = _dispatch_sqrt(bias_correction2)
denom = (exp_avg_sq.sqrt() / bias_correction2_sqrt).add_(eps)
param.addcdiv_(exp_avg, denom, value=-step_size)

### tensorflow adam implementation:
lr_t = lr * _dispatch_sqrt((1 - beta2 ** step)) / (1 - beta1 ** step)
denom = exp_avg_sq.sqrt().add_(eps)
param.addcdiv_(exp_avg, denom, value=-lr_t)

让我们比较一下 PyTorch 风格和 TensorFlow 风格 Adam 的更新方程。遵循 Adam 论文 (Kingma and Ba, 2014) 的表示法，PyTorch Adam（Kingma and Ba 论文中的算法 1）和 TensorFlow 风格 Adam（Kingma and Ba 论文中 2.1 节之前的公式）的梯度更新规则如下：

$\begin{aligned}\text{pytorch adam :}\quad \theta_t & =\theta_{t-1}-\alpha \cdot \hat{m}_t /\left(\sqrt{\hat{v}_t}+\varepsilon\right) \\& =\theta_{t-1}- \alpha \underbrace{\left[m_t /\left(1-\beta_1^t\right)\right]}_{=\hat{m}_t} /\left[\sqrt{\underbrace{v_t /\left(1-\beta_2^t\right)}_{=\hat{v}_t} }+\varepsilon\right]\\& =\theta_{t-1}- \alpha\left[m_t /\left(1-\beta_1^t\right)\right]\frac{\sqrt{1-\beta_2^t}}{\sqrt{v_t}+\color{green}{\varepsilon \sqrt{1-\beta_2^t}}}\end{aligned}$

$\begin{aligned}\text{tensorflow adam:}\quad \theta_t & =\theta_{t-1}-\alpha_t m_t /\left(\sqrt{v_t}+\hat{\varepsilon}\right) \\& =\theta_{t-1}-\underbrace{\left[\alpha \sqrt{1-\beta_2^t} /\left(1-\beta_1^t\right)\right]}_{=\alpha_t} m_t /\left(\sqrt{v_t}+\hat{\varepsilon}\right) \\& =\theta_{t-1}- \alpha\left[m_t /\left(1-\beta_1^t\right)\right] \frac{\sqrt{1-\beta_2^t}}{\sqrt{v_t}+\color{green}{\hat{\varepsilon}}} \end{aligned}$

上述等式强调了PyTorch和TensorFlow实现之间的区别在于它们的归一化项，即 $\color{green}{\varepsilon \sqrt{1-\beta_2^t}}$ 和 $\color{green}{\hat{\varepsilon}}$
上图显示，如果我们在PyTorch Adam和TensorFlow Adam中设置相同的 `eps`，那么PyTorch Adam在训练初期使用的归一化项要比TensorFlow Adam小得多。换句话说，PyTorch Adam在训练早期会进行更激进的梯度更新。我们的实验支持这一发现，如下文所示。

这如何影响可复现性和性能？为了对齐设置，我们记录了 https://github.com/openai/lm-human-preferences 中的原始查询、响应和奖励，并将它们保存到 https://huggingface.co/datasets/vwxyzjn/lm-human-preferences-debug/tree/main。我还记录了使用TF1的 `AdamOptimizer` 优化器进行前两个epoch训练的指标作为基准。以下是一些关键指标：

	OAI 的 TF1 Adam	PyTorch 的 Adam	我们自定义的 TensorFlow 风格 Adam
策略/近似 KL 散度	0.00037167023	0.0023672834504395723	0.000374998344341293
策略/剪裁比例	0.0045572915	0.02018229104578495	0.0052083334885537624
比率均值	1.0051285	1.0105520486831665	1.0044583082199097
比率方差	0.0007716546	0.005374275613576174	0.0007942612282931805
比率最大值	1.227216	1.8121057748794556	1.250215768814087
比率最小值	0.7400441	0.4011387825012207	0.7299948930740356
logprob_diff_mean	0.0047487603	0.008101251907646656	0.004073789343237877
logprob_diff_var	0.0007207897	0.004668936599045992	0.0007334011606872082
logprob_diff_max	0.20474821	0.594489574432373	0.22331619262695312
logprob_diff_min	-0.30104542	-0.9134478569030762	-0.31471776962280273

由于某种原因，PyTorch 的 `Adam` 会产生更激进的更新。以下是一些证据：
- PyTorch 的 `Adam` 的 `logprob_diff_var` 高出 6 倍。这里的 `logprobs_diff = new_logprobs - logprobs` 是指在两个 epoch 的训练后，初始策略和当前策略之间 token 对数概率的差异。`logprob_diff_var` 更大意味着对数概率变化的幅度比 OAI 的 TF1 Adam 更大。
- PyTorch 的 `Adam` 呈现出更极端的比率最大值和最小值。这里的 `ratio = torch.exp(logprobs_diff)`。`ratio_max=1.8121057748794556` 意味着对于某些 token，在当前策略下采样该 token 的概率是 1.8 倍，而 OAI 的 TF1 Adam 仅为 1.2 倍。
- `policy/approxkl` `policy/clipfrac` 更大。 由于激进的更新，比率被裁剪的频率高出 4.4 倍，近似 KL 散度大出 6 倍。
- 这种激进的更新很可能会导致进一步的问题。例如，PyTorch 的 `Adam` 的 `logprob_diff_mean` 大 1.7 倍，这将在下一次奖励计算中对应 1.7 倍大的 KL 惩罚；这可能会被放大。事实上，这可能与著名的 KL 散度问题有关——KL 惩罚远大于其应有的值，模型可能会更多地关注和优化它，从而导致负的 KL 散度。
更大的模型受影响更大。我们对 PyTorch 的 `Adam`（代号 `pt_adam`）和我们自定义的 TensorFlow 风格（代号 `tf_adam`）在 `gpt2` 和 `gpt2-xl` 上进行了比较实验。我们发现，在 `gpt2` 下，性能大致相似；然而，在 `gpt2-xl` 上，我们观察到更激进的更新，这意味着更大的模型受此问题的影响更大。
- 当 `gpt2-xl` 中初始策略更新更激进时，训练动态会受到影响。例如，我们看到 `pt_adam` 的 `objective/kl` 和 `objective/scores` 出现更大的尖峰，尤其是在 `sentiment` 中——在其中一个随机种子中，**最大的 KL 甚至高达 17.5**，这表明存在不必要的过度优化。
- 此外，由于 KL 值较大，许多其他训练指标也受到影响。例如，我们看到 `clipfrac`（比率被 PPO 目标剪裁系数 0.2 剪裁的时间分数）和 `approxkl` 大幅增加。

局限性

值得注意的是，本工作没有尝试重现 CNN DM 或 TL;DR 中的摘要工作。这是因为我们发现训练既耗时又脆弱。

我们所进行的特定训练运行显示出较差的 GPU 利用率（约 30%），因此完成一次训练运行需要近 4 天时间，这非常昂贵（只有 AWS 提供 p3dn.24xlarge，每小时费用为 31.212 美元）

此外，训练过程也很脆弱。虽然奖励有所提高，但我们发现很难重现 Ziegler 等人（2019）报告的“智能复制器”行为。下面是一些示例输出——显然，代理在某种程度上过拟合了。更多完整的日志请参见 https://wandb.ai/openrlbenchmark/lm-human-preferences/runs/1ab47rqi/logs。

结论

在这项工作中，我们深入研究了 OAI 的原始 RLHF 代码库，并整理了其实现细节列表。我们还创建了一个最小基础，当数据集和超参数受控时，它能重现 OAI 原始 RLHF 代码库的相同学习曲线。此外，我们还发现了一些令人惊讶的实现细节，例如 Adam 优化器的设置导致了早期 RLHF 训练中的激进更新。

致谢

这项工作得到了 Hugging Face 的 Big Science 集群 🤗 的支持。我们还要感谢 @lewtun 和 @natolambert 的有益讨论。

Bibtex

@article{Huang2023implementation,
  author = {Huang, Shengyi and Liu, Tianlin and von Werra, Leandro},
  title = {The N Implementation Details of RLHF with PPO},
  journal = {Hugging Face Blog},
  year = {2023},
  note = {https://huggingface.co/blog/the_n_implementation_details_of_rlhf_with_ppo},
}

更多博客文章

让 RL 回归 RLHF

作者： 2024年6月12日 • 99

使用开放 LLM 的 Constitutional AI

作者： 2024年2月1日 • 15

社区

blopen

1月19日

N-ice ;)

vijetadeshpande

2月28日

我对白化函数感到困惑。

它不应该是这样的吗？

whitened = (values - mean) / torch.rsqrt(var + 1e-8)

但是这里的函数有，

whitened = (values - mean) * torch.rsqrt(var + 1e-8)

vwxyzjn

文章作者 2月28日

请注意它是 rsqrt：https://pytorch.ac.cn/docs/stable/generated/torch.rsqrt.html

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

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

· 注册或登录评论