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

基本概念

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

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

RL概念圖

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

PS：這里要注意區(qū)分價(jià)值和獎(jiǎng)勵(lì)：價(jià)值是未來累計(jì)獎(jiǎng)勵(lì)的期望。獎(jiǎng)勵(lì)是我們做出該動(dòng)作后立即獲取的收益。

RLHF過程

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

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

Step2：Training Reward Model

數(shù)據(jù)

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

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

模型結(jié)構(gòu)

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

Reward Model

訓(xùn)練目標(biāo)

訓(xùn)練Reward Model是一個(gè)排序任務(wù)，針對(duì)query，輸入chosen和rejected response，訓(xùn)練目標(biāo)盡可能的使得chosen和rejected的差值更大，損失函數(shù)為：

以上就是第二步Training Reward Model的全部過程，基于rank loss訓(xùn)練了一個(gè)打分模型。在第三步強(qiáng)化學(xué)習(xí)中，reward模型將扮演環(huán)境的角色，針對(duì)模型預(yù)測(cè)的字符給出獎(jiǎng)勵(lì)分?jǐn)?shù)。

Step3：RLHF

整體結(jié)構(gòu)

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

DeepSpeed-Chat RLHF

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

• Actor Model：由SFT之后的模型初始化而來。作為策略（policy）模型，用于接收上文，做出動(dòng)作，預(yù)測(cè)下一個(gè)字符。學(xué)習(xí)完畢之后，我們最終使用的就是這個(gè)模型。
• Reference Model：和Actor Model同樣初始化自SFT Model，訓(xùn)練過程中凍結(jié)參數(shù)，用于和Actor Model做對(duì)比，保證模型不要偏離原始SFT Model太多。
• Reward Model：作為環(huán)境（env），訓(xùn)練過程中凍結(jié)參數(shù)，針對(duì)每一個(gè)狀態(tài)，給出獎(jiǎng)勵(lì)分?jǐn)?shù)。
• Critic Model：由Reward Model初始化而來，用于近似價(jià)值函數(shù)，輸入為狀態(tài)s，估計(jì)當(dāng)前狀態(tài)的價(jià)值V。

訓(xùn)練過程

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

首先是make_experience，首先在訓(xùn)練數(shù)據(jù)中抽取一部分query，然后Actor Model生成答案。然后我們依據(jù)這條答案獲取我們所需要的經(jīng)驗(yàn)：

• actor_logits:由Actor Model產(chǎn)生，包含對(duì)答案所有詞的概率分布。
• reference_logits：由Reference Model產(chǎn)生，包含對(duì)答案所有詞語的概率分布，用于和actor logits進(jìn)行對(duì)比，防止actor model偏離SFT Model太遠(yuǎn)。
• reward_score: 由Reward Model產(chǎn)生，為當(dāng)前句子狀態(tài)下，立即獲取的收益分?jǐn)?shù)。
• values：由Critic Model產(chǎn)生，估計(jì)當(dāng)前句子狀態(tài)下，到完成生成，可以獲取的回報(bào)。

整體流程如下：

make experience

然后在learn的時(shí)候，通過所產(chǎn)生的經(jīng)驗(yàn)進(jìn)行學(xué)習(xí)。我們通過Actor Model與Critic Model近似策略函數(shù)和價(jià)值函數(shù)，整體流程如下：

learn

關(guān)于learn這部分，詳細(xì)介紹下Critic Model訓(xùn)練和Actor Model訓(xùn)練過程。

Critic Loss

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

Actor Loss

舉個(gè)例子，如果我們模型有「上，下，左，右」四個(gè)動(dòng)作，分別有累計(jì)獎(jiǎng)勵(lì)「10，20，30，40」，我們做出任意動(dòng)作，都會(huì)獲取正向的累計(jì)獎(jiǎng)勵(lì)，因此模型也會(huì)向這個(gè)動(dòng)作更新參數(shù)。而實(shí)際上，我們累計(jì)獎(jiǎng)勵(lì)的平均期望為25，對(duì)于動(dòng)作「上，下」我們都應(yīng)該受到懲罰。

以上就是第三步的核心內(nèi)容，RL過程整體分為兩步，make experience和learn。我們首先采樣數(shù)據(jù)，然后生成結(jié)果，Reward Model給出環(huán)境的獎(jiǎng)勵(lì)，以及Critic Model對(duì)結(jié)果進(jìn)行評(píng)判。之后我們依據(jù)所獲取的經(jīng)驗(yàn)來對(duì)模型進(jìn)行更新。

DeepSpeed-Chat實(shí)踐與踩坑

Step2 Training Reward Model

這個(gè)步驟基本就是全部按照DeepSpeed-Chat代碼來了，使用cohere-zh的數(shù)據(jù)，大約構(gòu)造了2w條chosen-rejected pair用于訓(xùn)練。最后訓(xùn)練訓(xùn)練了一個(gè)epoch，在驗(yàn)證集上準(zhǔn)確率到了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

踩坑&解決方案

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

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

不過目前DeepSpeed-Chat也沒有解決，需要關(guān)閉Hybrid Engine進(jìn)行訓(xùn)練。

• DeepSpeed-Chat還有一個(gè)很嚴(yán)重的問題就是，在make experience的時(shí)候，強(qiáng)制Actor Model生成到最大長(zhǎng)度（設(shè)置max_length=min_length=max_min_length），這樣子導(dǎo)致模型生成偏差很大。對(duì)于一個(gè)簡(jiǎn)單的問題，模型可能本來生成簡(jiǎn)單的一句話就可以完美回答了，但是卻必須強(qiáng)制生成到最大長(zhǎng)度，這樣訓(xùn)練的模型和我們實(shí)際用起來的模型是有區(qū)別的。

對(duì)于這個(gè)問題，可以通過修改generate中的參數(shù)eos_token_id來解決。設(shè)置一個(gè)虛假的結(jié)束符，然后模型就可能生成正常的結(jié)束符。然后我們?cè)跇?gòu)建attention_mask來遮蔽掉結(jié)束符后面的回答。例如：
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)]

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

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

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

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

Trick

trick方面主要參考了The 37 Implementation Details of Proximal Policy Optimization以及影響PPO算法性能的10個(gè)關(guān)鍵技巧（附PPO算法簡(jiǎn)潔Pytorch實(shí)現(xiàn)）。對(duì)于部分trick，進(jìn)行了嘗試。

• Normalization of Advantages

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

• Overall Loss and Entropy Bonus

• 為了提高算法的探索能力，在actor的loss中增加一項(xiàng)策略熵，并乘以一個(gè)系數(shù)entropy_coef，使得在優(yōu)化actor_loss的同時(shí)，讓策略的熵盡可能大。一般我們?cè)O(shè)置entropy_coef=0.01

• 總體損失變?yōu)椋簂oss = policy_loss - entropy * entropy_coefficient + value_loss * value_coefficient

• 沒有指標(biāo)的量化，從我個(gè)人看結(jié)果而言，感覺沒有太多提升

• Reward Scale

• 對(duì)reward進(jìn)行縮放，將reward除以標(biāo)準(zhǔn)差

• 從訓(xùn)練log來看，對(duì)穩(wěn)定critic loss效果很好，畢竟將reward 進(jìn)行縮放之后，降低了return的估計(jì)區(qū)間。

• 沒有指標(biāo)的量化，從我個(gè)人看結(jié)果而言，提升很大，推薦使用。

關(guān)于其他的一些trick，如學(xué)習(xí)率衰減、梯度裁剪、Value Clipping等，本身框架就包含了，就不進(jìn)行特別說明了。當(dāng)然，這些trick在不同的數(shù)據(jù)或者模型上都會(huì)有不同的效果，需要自己進(jìn)行探索。

通過以上這些方法，最后也順利的完成了訓(xùn)練，希望對(duì)大家能有所幫助，也希望大家能分享自己有用的經(jīng)驗(yàn)與理解，關(guān)于RL自己了解的也很少，還需要更多的學(xué)習(xí)。