Hugging Face 的 LOGO
返回文章列表

使用 PPO 的 RLHF 的 N 個實現細節

釋出於 2023 年 10 月 24 日
在 GitHub 上更新
贊
65
Shengyi Costa Huang's avatar
黃聖毅 (Shengyi Costa Huang)
vwxyzjn
Tianlin Liu's avatar
劉天林
tianlinliu0121
訪客
Leandro von Werra's avatar
Leandro von Werra
lvwerra

RLHF/ChatGPT 最近是一個熱門研究課題。為了深入研究 RLHF，本部落格文章嘗試重現 OpenAI 2019 年的原始 RLHF 程式碼庫 openai/lm-human-preferences。儘管是“tensorflow-1.x-ness”，OpenAI 的原始程式碼庫經過了充分評估和基準測試，使其成為研究 RLHF 實現工程細節的良好起點。

我們旨在

  1. 重現 OAI 在風格任務中的結果，並匹配 openai/lm-human-preferences 的學習曲線。
  2. 提供一份實現細節清單，類似於 近端策略最佳化 (PPO) 的 37 個實現細節；無需痛苦地除錯強化學習的精神。
  3. 提供一個易於閱讀且最小的 RLHF 參考實現；

這項工作僅用於教育/學習目的。對於需要更多功能的高階使用者，例如使用 PEFT 執行更大的模型，huggingface/trl 將是一個不錯的選擇。

  • 在匹配學習曲線中，我們展示了我們的主要貢獻：建立一個能夠重現 OAI 在風格任務中結果的程式碼庫，並與 openai/lm-human-preferences 的學習曲線非常接近。
  • 然後，我們將深入探討與重現 OAI 工作相關的實現細節。在通用實現細節中，我們討論了基本細節，例如獎勵/值的生成方式以及響應的生成方式。在獎勵模型實現細節中，我們討論了獎勵歸一化等細節。在策略訓練實現細節中，我們討論了拒絕取樣和獎勵“白化”等細節。
  • 接下來，我們將檢查在獎勵標籤由 gpt2-large 生成的情況下，訓練不同基礎模型（例如，gpt2-xl、falcon-1b）的效果。
  • 最後，我們總結了我們的工作，並討論了其侷限性。

以下是重要連結

匹配學習曲線

我們的主要貢獻是重現 OAI 在風格任務（例如情感和描述性）中的結果。如下圖所示，我們的程式碼庫（橙色曲線）可以生成與 OAI 程式碼庫（藍色曲線）幾乎相同的學習曲線。

Untitled

關於執行 openai/lm-human-preferences 的注意事項

為了進行直接比較，我們運行了 openai/lm-human-preferences 中的原始 RLHF 程式碼，它將提供有價值的指標來幫助驗證和診斷我們的復現。我們能夠設定原始的 TensorFlow 1.x 程式碼，但這需要一個超特定的設定：

  • OAI 的資料集部分損壞/丟失（因此我們用類似的 HF 資料集替換它們，這可能會或可能不會導致效能差異）
  • 它不能在 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 工作相關的實現細節。在本節中，我們將討論基本細節，例如獎勵/值的生成方式以及響應的生成方式。這些細節沒有特定的順序：

  1. 獎勵模型和策略的值頭將 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` 的獎勵。

  2. 使用特殊的填充 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]])}
"""

  3. 相應地調整填充 token 的位置索引

    1. 在計算 logits 時，OAI 的程式碼透過正確地遮蔽填充標記來工作。這是透過找出與填充標記對應的標記索引 (lm_human_preferences/language/model.py#L296-L297)，然後相應地調整它們的位置索引 (lm_human_preferences/language/model.py#L320) 來實現的。

    2. 例如，如果 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)

    3. 關於 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>)
"""

    4. 關於 HF Transformers — generate 期間的 position_ids：在生成期間，我們不應該傳入 position_ids，因為 transformers 中已經調整了 position_ids（參見 huggingface/transformers#/7552）。

    通常，我們在 transformers 中幾乎從不傳入 position_ids。所有的掩碼和移位邏輯都已在例如 generate 函式中實現（需要永久程式碼連結）。

  4. 響應生成取樣固定長度的響應，不帶填充。

    1. 在生成響應時，OAI 使用 top_k=0, top_p=1.0，並且只對詞彙表進行分類取樣 (lm_human_preferences/language/sample.py#L43)，程式碼會持續取樣直到生成固定長度的響應 (lm_human_preferences/policy.py#L103)。值得注意的是，即使遇到 EOS（序列結束）token，它也會繼續取樣。

    2. 關於 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]])
"""

    3. 請注意，在較新的程式碼庫 https://github.com/openai/summarize-from-feedback 中，OAI 在遇到 EOS 標記時確實停止取樣 (summarize_from_feedback/utils/experiment_helpers.py#L19)。然而，在這項工作中，我們旨在進行 1:1 的復現，因此我們保持了即使遇到 EOS 標記也可以繼續取樣的設定。

  5. 獎勵模型和策略訓練的學習率退火。

    1. 正如 Ziegler 等人（2019）所建議的，獎勵模型只訓練一個 epoch，以避免過度擬合有限的人類標註資料（例如，descriptiveness 任務只有大約 5000 個標籤）。在這個單個 epoch 期間，學習率會退火到零 (lm_human_preferences/train_reward.py#L249)。
    2. 與獎勵模型訓練類似，學習率被退火到零 (lm_human_preferences/train_policy.py#L172-L173)。

  6. 對不同程序使用不同的種子

    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)。

獎勵模型實現細節

在本節中，我們將討論獎勵模型特定的實現細節。我們將討論獎勵歸一化和層初始化等細節。這些細節沒有特定的順序：

  1. 獎勵模型僅在最後一個 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)。然而，在這項工作中，我們只堅持原始設定。
  2. 獎勵頭層初始化
    1. 獎勵頭的權重根據 N(0,1/(dmodel +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 頁有一個筆誤，稱分佈為 N(0,1/(dmodel +1)) \mathcal{N}\left(0,1 /\left(d_{\text {model }}+1\right)\right) ，沒有平方根）。
    2. 獎勵頭的偏差設定為 0 (lm_human_preferences/language/model.py#L254)。
  3. 獎勵模型歸一化（前後）
    1. 在論文中，Ziegler 等人 (2019) 提到：“為了保持獎勵模型在訓練過程中的尺度一致，我們對其進行歸一化，使其對於 x∼D,y∼ρ(⋅∣x) 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. 我們用 μD \mu_{\mathcal{D}} 表示經驗均值，σD \sigma_{\mathcal{D}} 表示經驗標準差，gg 表示 reward_gain，bb 表示 reward_bias，μT=0 \mu_{\mathcal{T}} = 0 目標均值，σT=1 \sigma_{\mathcal{T}}=1 目標標準差。那麼我們有以下公式。g∗N(μD,σD)+b=N(g∗μD,g∗σD)+b=N(g∗μD+b,g∗σD)=N(μT,σT)g=σTσDb=μT−g∗μD\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∼ρ(⋅∣x) y \sim \rho(·|x) 來自預訓練語言模型 ρ\rho 。模型 ρ\rho 被固定為參考，並且在獎勵學習中不進行更新 (lm_human_preferences/train_reward.py#L286C1-L286C31)。

策略訓練實現細節

在本節中，我們將深入探討層初始化、資料後處理和 dropout 設定等細節。我們還將探討拒絕取樣、獎勵“白化”和自適應 KL 等技術。這些細節沒有特定的順序：

  1. 透過取樣溫度對 logits 進行縮放。

    1. 在計算響應的對數機率時，模型首先輸出響應中 tokens 的 logits，然後將 logits 除以取樣溫度 (lm_human_preferences/policy.py#L121)。即，logits /= self.temperature
    2. 在非正式測試中，我們發現如果沒有這種縮放，KL 會比預期上升得更快，並且效能會下降。

  2. 值頭層初始化

    1. 價值頭的權重根據 N(0,0)\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)。

  3. 選擇以句點開頭和結尾的查詢文字

    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。

  4. 停用 Dropout

    1. Ziegler 等人 (2019) 建議：“我們不使用 dropout 進行策略訓練。” 這也在程式碼中完成 (lm_human_preferences/policy.py#L48)。

  5. 拒絕取樣

    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 中的一些例子：

        Samples extracted from our reproduction [https://wandb.ai/openrlbenchmark/lm_human_preference_details/runs/djf8yymv/logs](https://wandb.ai/openrlbenchmark/lm_human_preference_details/runs/djf8yymv/logs?workspace=user-costa-huang). Notice the 1st and 3rd example has too many tokens after the period, so its score was replaced by -1.

        從我們的復現中提取的樣本 https://wandb.ai/openrlbenchmark/lm_human_preference_details/runs/djf8yymv/logs。請注意，第一個和第三個示例在句號後有太多 token，因此其分數被替換為 -1。

  6. 折扣因子 = 1

    1. 折扣引數 γ\gamma 設定為 1 (lm_human_preferences/train_policy.py#L56)，這意味著未來獎勵與即時獎勵具有相同的權重。

  7. 訓練迴圈術語：PPO 中的批次和微批次

    1. 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.]

  8. 每個 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]。

  9. 每個迷你批次的獎勵和優勢白化，可選地進行均值平移

    1. 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

    1. 在每個迷你批次中，OAI 會白化獎勵 whiten(rewards, shift_mean=False)，不進行均值平移 (lm_human_preferences/train_policy.py#L325)，並白化優勢 whiten(advantages)，進行均值平移 (lm_human_preferences/train_policy.py#L338)。

    2. 最佳化注意： 如果迷你批次數量為一個（本復現中就是這種情況），我們只需要對獎勵進行一次白化，並計算和白化優勢一次，因為它們的值不會改變。

    3. 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)

  10. 裁剪值函式

    1. 如原始 PPO 中所做 (baselines/ppo2/model.py#L68-L75)，值函式被裁剪 (lm_human_preferences/train_policy.py#L343-L348)，方式與策略目標類似。

  11. 自適應 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 節之前的公式）的梯度更新規則如下：

pytorch adam :θt=θt−1−α⋅m^t/(v^t+ε)=θt−1−α[mt/(1−β1t)]⏟=m^t/[vt/(1−β2t)⏟=v^t+ε]=θt−1−α[mt/(1−β1t)]1−β2tvt+ε1−β2t\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}

tensorflow adam:θt=θt−1−αtmt/(vt+ε^)=θt−1−[α1−β2t/(1−β1t)]⏟=αtmt/(vt+ε^)=θt−1−α[mt/(1−β1t)]1−β2tvt+ε^\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實現之間的區別在於它們的歸一化項，即 ε1−β2t\color{green}{\varepsilon \sqrt{1-\beta_2^t}} 和 ε^\color{green}{\hat{\varepsilon}}。如果我們將 ε^=ε1−β2t\hat{\varepsilon} =\varepsilon \sqrt{1-\beta_2^t} 時，這兩個版本是等效的。然而，在PyTorch和TensorFlow的API中，我們只能透過 `eps` 引數設定 ε\varepsilon (PyTorch) 和 ε^\hat{\varepsilon} (TensorFlow)，這導致了它們的更新等式存在差異。如果我們將 ε\varepsilon 和 ε^\hat{\varepsilon} 設定為相同的值，比如 1e-5，會怎樣？那麼對於TensorFlow Adam，歸一化項 ε^=1e-5\hat{\varepsilon} = \text{1e-5} 只是一個常數。但對於PyTorch Adam，歸一化項 ε1−β2t{\varepsilon \sqrt{1-\beta_2^t}} 隨著時間而變化。重要的是，在時間步長 tt 較小的時候，ε1−β2t{\varepsilon \sqrt{1-\beta_2^t}} 會遠小於1e-5，隨著時間步長的增加，它會逐漸接近1e-5。下圖比較了這兩個歸一化項隨時間步長的變化。

    norma_const_comparison.png

  • 上圖顯示，如果我們在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` 大幅增加。

adam_gpt2.png

adam_gpt2_xl.png

侷限性

值得注意的是，本工作沒有嘗試重現 CNN DM 或 TL;DR 中的摘要工作。這是因為我們發現訓練既耗時又脆弱。

我們所進行的特定訓練執行顯示出較差的 GPU 利用率（約 30%），因此完成一次訓練執行需要近 4 天時間，這非常昂貴（只有 AWS 提供 p3dn.24xlarge，每小時費用為 31.212 美元）

此外，訓練過程也很脆弱。雖然獎勵有所提高，但我們發現很難重現 Ziegler 等人（2019）報告的“智慧複製器”行為。下面是一些示例輸出——顯然，代理在某種程度上過擬合了。更多完整的日誌請參見 https://wandb.ai/openrlbenchmark/lm-human-preferences/runs/1ab47rqi/logs。

tldr1.png

tldr2.png

結論

在這項工作中，我們深入研究了 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

作者：  vwxyzjn 2024年6月12日 • 99

使用開放 LLM 的 Constitutional AI

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

社群

blopen

N-ice ;)

vijetadeshpande

我對白化函式感到困惑。

它不應該是這樣的嗎？

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

但是這裡的函式有，

whitened = (values - mean) * torch.rsqrt(var + 1e-8)
·
vwxyzjn
 文章作者

請注意它是 rsqrt：https://pytorch.org/docs/stable/generated/torch.rsqrt.html

註冊 或 登入 評論

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