當下，RLHF/ChatGPT 已經變成了一個非常流行的話題。我們正在致力于更多有關 RLHF 的研究，這篇博客嘗試復現 OpenAI 在 2019 年開源的原始 RLHF 代碼庫，其倉庫位置位于 openai/lm-human-preferences。盡管它具有 “tensorflow-1.x” 的特性，但 OpenAI 的原始代碼庫評估和基準測試非常完善，使其成為研究 RLHF 實現工程細節的好地方。

我們的目標是:

1. 復現 OAI 在風格化任務中的結果，并匹配 openai/lm-human-preferences 的學習曲線。
2. 提供一個實現細節的清單，類似于 近端優化策略的 37 個實施細節 (The 37 Implementation Details of Proximal Policy Optimization) 和 沒有痛苦折磨的調試 RL (Debugging RL, Without the Agonizing Pain) 的風格;
3. 提供一個易于閱讀且簡潔的 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 的說明

為了直觀比較，我們運行了原始的 RLHF 代碼，其倉庫位置位于 openai/lm-human-preferences，它將提供寶貴的指標，以幫助驗證和診斷我們的復現。我們能夠設置原始的 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 上運行，因為它沒有實現梯度累積。相反，它使用一個大的 BS (批量大小)，并在 8 個 GPU 上分割 batch (批量)，僅在 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 是 BS (批量大小)，T 是序列長度，而 1 代表獎勵頭的輸出結構的維度為 1 (lm_human_preferences/rewards.py#L105-L107, lm_human_preferences/policy.py#L111)。
   3. T 意味著每個 token 都有與其和前文關聯的獎勵。例如，eyes token 將有一個與他在想某事，但他的眼神很難讀懂 相對應的獎勵。

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 的 transformers — 填充 token 的注解。 根據 (transformers#2630#issuecomment-578159876)，在 GPT 和 GPT-2 的預訓練期間沒有使用填充 token; 因此，transformer 的 gpt2 模型與其分詞器沒有關聯的官方填充 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 的代碼通過適當地屏蔽填充 token 來工作。這是通過找出與填充 token 相對應的 token 索引來實現的 (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 的注解。 我們可以通過 1) 左填充和 2) 傳入適當的 position_ids ，使用 Hugging Face 的 transformer 復制精確的 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 ——在 生成 過程中的 position_ids 的注解: 在生成過程中，我們不應傳入 position_ids ，因為在 transformers 中， position_ids 已經以某種方式被調整了。當我在生成過程中也傳入 position_ids 時，性能會災難性地惡化。

   通常情況下，我們幾乎從不在 transformers 中傳遞 position_ids 。所有的遮蔽 (masking) 和移位 (shifting) logic 已經實現，例如，在 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. TODO: why does generation collapse with this?
          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 中，當遇到 EOS token 時，OAI 確實會停止采樣 (summarize_from_feedback/utils/experiment_helpers.py#L19)。然而，在這項工作中，我們的目標是進行 1:1 的復刻，所以我們調整了設置，即使遇到 eos_token 也可以繼續采樣。

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

   1. 正如 Ziegler 等人 (2019) 建議的，獎勵模型只訓練一個 epcho，以避免過度擬合有限量的人類注釋數據 (例如，descriptiveness 任務只有大約 5000 個標簽)。在這個單一的 epcho 中，學習率會退火至零 (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. 注: 我認為數據集的洗牌 (shuffling) 存在一個錯誤——由于某種原因，數據集是使用相同的種子進行洗牌的 (lm_human_preferences/lm_tasks.py#L94-L97)。

獎勵模型實現細節

在本節中，我們討論了獎勵模型特定的實現細節。我們討論了諸如獎勵歸一化和層初始化等細節。以下是這些細節，不按特定順序排列:

1. 獎勵模型只輸出最后一個 token 的值。

   1. 請注意，在對 query 和 response 的連接進行前向傳遞后獲得的獎勵將具有形狀 (B, T, 1) ，其中 B 是 BS(批量大小)，T 是序列長度 (始終相同; 在 OAI 的設置中，它是 query_length + response_length = 64 + 24 = 88 ，用于風格任務，參見 launch.py#L9-L11)，1 是獎勵頭其維度為 1。對于 RLHF (Reinforcement Learning from Human Feedback，通過人類反饋進行強化學習) 的目的，原始代碼庫提取最后一個 token 的獎勵 (lm_human_preferences/rewards.py#L132)，因此獎勵將只具有形狀 (B, 1) 。
   2. 請注意，在較新的代碼庫 openai/summarize-from-feedback 中，OAI 在遇到 EOS token 時停止采樣 (summarize_from_feedback/utils/experiment_helpers.py#L19)。在提取獎勵時，它將確定 last_response_index ，即 EOS token 之前的索引 (#L11-L13)，并在該索引處提取獎勵 (summarize_from_feedback/reward_model.py#L59)。但在此工作中，我們只是堅持原始設置。

2. 獎勵頭層初始化

   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. 獎勵頭的 bias (偏置) 設為 0 (lm_human_preferences/language/model.py#L254)。

3. 獎勵模型的前后歸一化

   1. 在論文中，Ziegler 等人 (2019) 提到“為了保持訓練過程中獎勵模型的規模一致，我們將其歸一化，使其在 $ x \sim \mathcal{D}, y \sim \rho(·|x) $ 的情況下，均值為 0，方差為 1”。為了執行歸一化過程，代碼首先創建了 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。以下是這些細節，排列不分先后:

1. 通過采樣溫度來縮放 logits

   1. 在計算響應的對數概率時，模型首先輸出響應中 token 的 logits，然后用采樣溫度除以這些 logits (lm_human_preferences/policy.py#L121)。即 logits /= self.temperature
   2. 在一個非正式的測試中，我們發現如果不進行此縮放，KL 散度會比預期更快地上升，性能會下降。

2. 價值頭層的初始化

   1. 價值頭的權重是根據 \(\mathcal{N}(0,0)\) 進行初始化的 (lm_human_preferences/language/model.py#L368、lm_human_preferences/language/model.py#L251-L252)。
   2. 獎勵頭的 bias (偏置) 設置為 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，我們發現沒有必要提取以句號開始和結束的句子，因為數據集已經是這樣預處理過的 (例如，"usually , he would be tearing around the living room , playing with his toys." ) 為此，我們為 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 個 token 之間有一個句號，然后在那個句號處截斷 (這是‘句子結束’的粗略近似。我們選擇它是因為它很容易集成到 RL 循環中，即使是粗略的近似也足以使人類評估任務變得稍微容易一些)。在 RL 微調期間，我們對沒有這樣的句號的延續給予固定獎勵 -1。”
   2. 具體來說，通過以下步驟實現此目的:

      1. token 截斷: 我們想要在第一個出現在響應的 truncate_after 位置之后的 truncate_token 處截斷 (lm_human_preferences/train_policy.py#L378)。

         1. 代碼注釋: “中心示例: 將截斷 token 后的所有 token 替換為填充 token”

      2. 在截斷響應上運行獎勵模型: 在 token 截斷過程將響應截斷后，代碼然后在 截斷的響應 上運行獎勵模型。

      3. 拒絕采樣: 如果在第 16 和 24 個 token 之間沒有句號，那么將響應的分數替換為固定的低值 (例如 -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。請注意，第 1 和第 3 個示例在句號后有太多 token，因此其分數被替換為 -1。

6. 折現因子 (discount factor) = 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 來幫助處理梯度累積的情況。在每個時期，它都會洗牌批次索引。

      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. 基于每個標記的 KL 懲罰

   • 代碼為獎勵添加了每個標記的 KL 懲罰 (lm_human_preferences/train_policy.py#L150-L153)，以阻止策略與原始策略差異過大。

   • 以 “usually, he would” 為例，它被標記化為 [23073, 11, 339, 561] 。假設我們使用 [23073] 作為查詢，[11, 339, 561] 作為響應。然后在默認的 gpt2 參數下，響應標記將具有參考策略的對數概率 logprobs=[-3.3213, -4.9980, -3.8690] 。

     • 在第一個 PPO 更新時期和小批次更新時，激活策略將具有相同的對數概率new_logprobs=[-3.3213, -4.9980, -3.8690] 。因此，每個標記的 KL 懲罰將為 kl = new_logprobs - logprobs = [0., 0., 0.] 。
     • 但是，在第一個梯度反向傳播后，我們可能會得到 new_logprob=[3.3213, -4.9980, -3.8690] ，因此每個標記的 KL 懲罰變為 kl = new_logprobs - logprobs = [-0.3315, -0.0426, 0.6351] 。
     • 隨后，non_score_reward = beta * kl ，其中 beta 是 KL 懲罰系數 \(beta\)，它被添加到從獎勵模型獲得的 score 中，以創建用于訓練的 rewards 。score 僅在每個回合 ( episode ) 結束時給出，可能類似于 [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
      

   2. 在每個小批次中，OAI 使用 whiten(rewards, shift_mean=False) 對獎勵進行白化，不對均值進行平移處理 (lm_human_preferences/train_policy.py#L325)，并使用平移后的均值對優勢進行白化 whiten(advantages) (lm_human_preferences/train_policy.py#L338)。

   3. 優化注意事項: 如果小批次的數量為一 (在此復現中是這種情況)，我們只需要對獎勵進行白化、計算并對優勢進行一次白化，因為它們的值不會改變。

   4. TensorFlow vs PyTorch 注意事項: tf.moments 與 torch.var 的不同行為: 由于方差計算方式不同，Torch 和 TensorFlow 中的白化行為不同:

      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 和 Ba，2014) 的符號表示，我們可以得到 PyTorch Adam (Kingma 和 Ba 論文的算法 1) 和 TensorFlow 風格 Adam (Kingma 和 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}}\)。如果我們設置 \(\hat{\varepsilon} = \varepsilon \sqrt{1-\beta_2^t}\)，則這兩個版本是等價的。然而，在 PyTorch 和 TensorFlow 的 API 中，我們只能通過 eps 參數設置 \(\varepsilon\) (PyTorch) 和 \(\hat{\varepsilon}\) (TensorFlow)，從而導致它們的更新方程存在差異。如果我們將 \(\varepsilon\) 和 \(\hat{\varepsilon}\) 都設置為相同的值，比如 1e-5 會發生什么？那么對于 TensorFlow Adam，歸一化項 \(\hat{\varepsilon} = \text{1e-5}\) 就是一個常數。但對于 PyTorch Adam，歸一化項 \({\varepsilon \sqrt{1-\beta_2^t}}\) 隨著時間的推移而變化。重要的是，當時間步 \(t\) 較小時，該項 \({\varepsilon \sqrt{1-\beta_2^t}}\) 明顯小于 1e-5，隨著時間步增加，逐漸接近 1e-5。下面的圖表比較了這兩個歸一化項隨著時間步的變化情況:

• 上圖顯示，如果我們在 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 優化器的前兩個訓練周期的指標作為基準。以下是一些關鍵指標:

	OAI 的 TF1 Adam	PyTorch 的 Adam	我們自定義的類似 TensorFlow 風格的 Adam
policy/approxkl	0.00037167023	0.0023672834504395723	0.000374998344341293
policy/clipfrac	0.0045572915	0.02018229104578495	0.0052083334885537624
ratio_mean	1.0051285	1.0105520486831665	1.0044583082199097
ratio_var	0.0007716546	0.005374275613576174	0.0007942612282931805
ratio_max	1.227216	1.8121057748794556	1.250215768814087
ratio_min	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 是經過兩個訓練周期后，初始策略和當前策略之間的標記對數概率差異。具有更大的 logprob_diff_var 意味著對數概率變化的幅度比 OAI 的 TF1 Adam 大。
  • PyTorch 的 Adam 呈現更極端的最大和最小比率。這里的 ratio = torch.exp(logprobs_diff) 。具有 ratio_max=1.8121057748794556 意味著對于某些標記，在當前策略下抽取該標記的概率要比 OAI 的 TF1 Adam 高 1.8 倍，而后者僅為 1.2 倍。
  • 更大的 policy/approxkl 和 policy/clipfrac。由于激進的更新，比率被剪切的次數 多 4.4 倍，近似的 KL 散度大 6 倍。
  • 這種激進的更新可能會導致進一步的問題。例如，PyTorch 的 Adam 中的logprob_diff_mean 要大 1.7 倍，這將對下一個獎勵計算中的 KL 懲罰產生 1.7 倍大的影響; 這可能會被累積。實際上，這可能與著名的 KL 散度問題有關—— KL 懲罰遠大于它應該的值，模型可能會更多地關注它并進行更多優化，從而導致負的 KL 散度。

• 更大的模型受到更多影響。我們進行了一些實驗，比較了 PyTorch 的 Adam (代號 pt_adam ) 和我們自定義的類似 TensorFlow 風格的 Adam (代號 tf_adam ) 在 gpt2 和 gpt2-xl 上的性能。我們發現在 gpt2 下性能大致相似; 但是在 gpt2-xl 下，我們觀察到了更激進的更新，這意味著更大的模型受到了更多的影響。

  • 當在 gpt2-xl 中初始策略更新更為激進時，訓練動態會受到影響。例如，我們發現使用 pt_adam 時，sentiment 的 objective/kl 和 objective/scores 峰值要大得多， 在其中一個隨機種子中，最大的 KL 值達到了 17.5 ，這表明了不希望的過度優化。
  • 此外，由于 KL 更大，許多其他訓練指標也受到影響。例如，我們觀察到更大的 clipfrac (ratio 被 PPO 的目標裁剪系數 0.2 裁剪的時間比例) 和 approxkl 。

局限性

注意到這項工作沒有嘗試復現 CNN DM 中的摘要工作。這是因為我們發現訓練耗時且不穩定。

我們的特定訓練運行顯示 GPU 利用率較低 (約 30%)，因此一個訓練運行需要近 4 天的時間，這非常昂貴 (只有 AWS 銷售 p3dn.24xlarge，每小時費用為 31.212 美元)。

此外，訓練也很不穩定。雖然獎勵值上升，但我們發現難以復現 Ziegler 等人 (2019 年) 報告的“智能復制”行為。以下是一些樣本輸出 — 顯然，智能體出現了某種程度的過擬合。請查看 https://wandb.ai/openrlbenchmark/lm-human-preferences/runs/1ab47rqi/logs 以獲取更完整的日志。

總結

在這項工作中，我們深入研究了 OpenAI 的原始 RLHF (Reinforcement Learning from Human Feedback) 代碼庫，并編制了其實施細節的列表。我們還創建了一個最小的基礎版本，當數據集和超參數受控制時，可以復現與 OpenAI 原始 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},
}

---

英文原文: https://hf.co/blog/the_n_implementation_details_of_rlhf_with_ppo

原文作者: Shengyi Costa Huang, Tianlin Liu, Leandro von We

譯者: innovation64

審校/排版: zhongdongy (阿東)

總結

以上是生活随笔為你收集整理的使用 PPO 算法进行 RLHF 的 N 步实现细节的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。