最近兩月倒騰了一波RLHF，從ColossalAI到TRLX以及DeepSpeed-Chat，最后基于DeepSpeed-Chat成功訓練上了自己的模型，最后效果也是肉眼可見的提升。對這一部分進行下總結，包括原理，代碼以及踩坑與解決方案。

基本概念

首先還是解釋一下一些概念，從NLP的角度舉一些例子。

首先是RL中的Policy，State，Action。

RL概念圖

接下來介紹Reward，Return，Q，V。

PS：這里要注意區分價值和獎勵：價值是未來累計獎勵的期望。獎勵是我們做出該動作后立即獲取的收益。

RLHF過程

整個過程主要是分為三步：SFT，Training Reward Model，RLHF。

這里主要介紹一下Training Reward Model和RLHF。

Step2：Training Reward Model

數據

首先是這一部分數據的構成，每一條數據由query，chosen response， rejected response構成，chosen response相較于rejected response質量更高。我使用的數據來自Cohere/miracl-zh-queries-22-12，里面對于每條query包含了positive_passages和negative_passages兩個部分，positive部分可以視為chosen，negative部分視為rejected，就可以采樣構建我們訓練Reward Model的數據了：

{
"query": "聯合國總部在哪里？",
"chosen": "聯合國總部大樓（亦稱聯合國大廈）是聯合國總部的所在地，位于美國紐約市曼哈頓東側，屬于國際領土，因此只要是會員國國民持有護照就可以進入，包括與美國無邦交的聯合國會員國。聯合國總部大樓位于紐約市，其西側邊界為第一大道、南側邊界為東42街、北側邊界為東48街、東側邊界為東河，從聯合國總部大樓可以俯瞰東河。此大樓于1949年和1950年間興建，土地購自于當時的紐約房地產家，面積闊達17英畝（約6.87973公頃）。在此之前，洛克斐勒家族有意提供其在紐約州威斯特徹斯特郡洛克菲勒莊園的土地，但因距離曼哈頓遙遠而作罷；之后，納爾遜·洛克菲勒便協助新的土地的購買，其父小約翰·戴維森·洛克菲勒則捐助了850萬美元協助興建大樓。", 
"rejected": "聯合國的15個專門機構（如教科文組織）都沒有設在總部。然而，有一些“自治附屬機構”（如聯合國兒童基金會）的總部設在聯合國總部。"
}

模型結構

Reward Model相較于原始的SFT Model，在后面加上了一個value head，value head是一個Linear，輸入維度為模型的hidden_dim，輸出維度為1，輸出表示模型預測每一字符獲取的得分。DeepSpeed-Chat中使用最后一個字符的得分作為整個response的得分（當然也可以使用整個句子中每個字符的平均分作為整體的得分）。

Reward Model

訓練目標

訓練Reward Model是一個排序任務，針對query，輸入chosen和rejected response，訓練目標盡可能的使得chosen和rejected的差值更大，損失函數為：

以上就是第二步Training Reward Model的全部過程，基于rank loss訓練了一個打分模型。在第三步強化學習中，reward模型將扮演環境的角色，針對模型預測的字符給出獎勵分數。

Step3：RLHF

整體結構

首先來從整體上看一下這部分（這里就只介紹RL部分，PTX就是加上了預訓練任務）：

DeepSpeed-Chat RLHF

RLHF基于A2C方法，這一步包含了四個模型:

• Actor Model：由SFT之后的模型初始化而來。作為策略（policy）模型，用于接收上文，做出動作，預測下一個字符。學習完畢之后，我們最終使用的就是這個模型。
• Reference Model：和Actor Model同樣初始化自SFT Model，訓練過程中凍結參數，用于和Actor Model做對比，保證模型不要偏離原始SFT Model太多。
• Reward Model：作為環境（env），訓練過程中凍結參數，針對每一個狀態，給出獎勵分數。
• Critic Model：由Reward Model初始化而來，用于近似價值函數，輸入為狀態s，估計當前狀態的價值V。

訓練過程

接下來梳理一遍訓練過程。訓練過程整體分為兩步：maker experience和learn。

首先是make_experience，首先在訓練數據中抽取一部分query，然后Actor Model生成答案。然后我們依據這條答案獲取我們所需要的經驗：

• actor_logits:由Actor Model產生，包含對答案所有詞的概率分布。
• reference_logits：由Reference Model產生，包含對答案所有詞語的概率分布，用于和actor logits進行對比，防止actor model偏離SFT Model太遠。
• reward_score: 由Reward Model產生，為當前句子狀態下，立即獲取的收益分數。
• values：由Critic Model產生，估計當前句子狀態下，到完成生成，可以獲取的回報。

整體流程如下：

make experience

然后在learn的時候，通過所產生的經驗進行學習。我們通過Actor Model與Critic Model近似策略函數和價值函數，整體流程如下：

learn

關于learn這部分，詳細介紹下Critic Model訓練和Actor Model訓練過程。

Critic Loss

Critic Model估計當前狀態可以獲取的價值，也就是我們前面所說的V值。模型的輸入為狀態s，也就是當前模型生成的句子，輸出為狀態價值V(s)。

Actor Loss

舉個例子，如果我們模型有「上，下，左，右」四個動作，分別有累計獎勵「10，20，30，40」，我們做出任意動作，都會獲取正向的累計獎勵，因此模型也會向這個動作更新參數。而實際上，我們累計獎勵的平均期望為25，對于動作「上，下」我們都應該受到懲罰。

以上就是第三步的核心內容，RL過程整體分為兩步，make experience和learn。我們首先采樣數據，然后生成結果，Reward Model給出環境的獎勵，以及Critic Model對結果進行評判。之后我們依據所獲取的經驗來對模型進行更新。

DeepSpeed-Chat實踐與踩坑

Step2 Training Reward Model

這個步驟基本就是全部按照DeepSpeed-Chat代碼來了，使用cohere-zh的數據，大約構造了2w條chosen-rejected pair用于訓練。最后訓練訓練了一個epoch，在驗證集上準確率到了0.79左右。

***** Evaluating reward, Epoch 1/1 *****
step: 499/2287, chosen_last_scores (higher is better) : 5.074454307556152,reject_last_scores (lower is better) : 0.5599770545959473, acc (higher is better) : 0.812000036239624
step: 999/2287, chosen_last_scores (higher is better) : 5.084388732910156,reject_last_scores (lower is better) : 0.7938708662986755, acc (higher is better) : 0.7940000295639038
step: 1499/2287, chosen_last_scores (higher is better) : 5.106724262237549,reject_last_scores (lower is better) : 0.7971451878547668, acc (higher is better) : 0.7986666560173035
step: 1999/2287, chosen_last_scores (higher is better) : 5.0183587074279785,reject_last_scores (lower is better) : 0.672178328037262, acc (higher is better) : 0.7955000400543213
chosen_last_scores (higher is better) : 5.028912544250488,reject_last_scores (lower is better) : 0.7077188491821289, acc (higher is better) : 0.7936161160469055

Step3 RLHF

踩坑&解決方案

在這個步驟中，從跑通到收斂還是有不少麻煩，分享一些比較重要的點：

• 訓練過程中發現make experience的時候，model.generate()產生的答案全是重復的無意義字。這里就很奇怪了，最后發現是開啟了DeepSpeed Hybrid Engine所導致的，查看了issue之后，也發現了有類似的問題。

不過目前DeepSpeed-Chat也沒有解決，需要關閉Hybrid Engine進行訓練。

• DeepSpeed-Chat還有一個很嚴重的問題就是，在make experience的時候，強制Actor Model生成到最大長度（設置max_length=min_length=max_min_length），這樣子導致模型生成偏差很大。對于一個簡單的問題，模型可能本來生成簡單的一句話就可以完美回答了，但是卻必須強制生成到最大長度，這樣訓練的模型和我們實際用起來的模型是有區別的。

對于這個問題，可以通過修改generate中的參數eos_token_id來解決。設置一個虛假的結束符，然后模型就可能生成正常的結束符。然后我們在構建attention_mask來遮蔽掉結束符后面的回答。例如：
seq = [0,0,0,0, prompt, answer, eos_token, other_word]
mask = [0,0,0,0,1(prompt),1(answer),1(eos_token),0(other_word)]

通過以上兩步基本就可以跑通流程了，但是訓練過程中還遇到了一個比較大的問題，就是Critic Loss并不收斂，越來越大。

具體的原因是：訓練后期，隨著模型輸出答案越來越好，我們的reward值也會越來越高，導致我們最終累計回報return的區間也會越來越大。而前面說過，我們Critic Loss是一個MSE損失，因此訓練后期，隨著return的估計范圍越來越大，Critic Model就越難估計。在我訓練的過程中，一開始return的范圍是在3-4左右，訓練后期漲到了18-20，因此我們需要想點辦法來約束一下我們的return。

首先的超參的調節，γ參數為折扣回報率，DeepSpeed-Chat中初始設置為1，可以將其調小一些來緩解。其次的話，使用reward scale這一trick幫助也非常大。

通過以上這些步驟，基本上我們正常訓練了，模型最后也能看到一些效果，但是需要取得更好的效果，我們就需要引入一些trick了。

Trick

trick方面主要參考了The 37 Implementation Details of Proximal Policy Optimization以及影響PPO算法性能的10個關鍵技巧（附PPO算法簡潔Pytorch實現）。對于部分trick，進行了嘗試。

• Normalization of Advantages

• 將Advantage進行歸一化：adv=(adv-mean)/std
• 在mini-batch上進行
• 沒有指標的量化，從我個人看結果而言，感覺提升不大

• Overall Loss and Entropy Bonus

• 為了提高算法的探索能力，在actor的loss中增加一項策略熵，并乘以一個系數entropy_coef，使得在優化actor_loss的同時，讓策略的熵盡可能大。一般我們設置entropy_coef=0.01

• 總體損失變為：loss = policy_loss - entropy * entropy_coefficient + value_loss * value_coefficient

• 沒有指標的量化，從我個人看結果而言，感覺沒有太多提升

• Reward Scale

• 對reward進行縮放，將reward除以標準差

• 從訓練log來看，對穩定critic loss效果很好，畢竟將reward 進行縮放之后，降低了return的估計區間。

• 沒有指標的量化，從我個人看結果而言，提升很大，推薦使用。

關于其他的一些trick，如學習率衰減、梯度裁剪、Value Clipping等，本身框架就包含了，就不進行特別說明了。當然，這些trick在不同的數據或者模型上都會有不同的效果，需要自己進行探索。

通過以上這些方法，最后也順利的完成了訓練，希望對大家能有所幫助，也希望大家能分享自己有用的經驗與理解，關于RL自己了解的也很少，還需要更多的學習。