0
  • 聊天消息
  • 系統(tǒng)消息
  • 評論與回復(fù)
登錄后你可以
  • 下載海量資料
  • 學(xué)習(xí)在線課程
  • 觀看技術(shù)視頻
  • 寫文章/發(fā)帖/加入社區(qū)
會員中心
創(chuàng)作中心

完善資料讓更多小伙伴認(rèn)識你,還能領(lǐng)取20積分哦,立即完善>

3天內(nèi)不再提示

圖解大模型RLHF系列之:人人都能看懂的PPO原理與源碼解讀

jf_pmFSk4VX ? 來源:GiantPandaCV ? 2024-01-14 11:19 ? 次閱讀

大家好,最近我又讀了讀RLHF的相關(guān)paper和一些開源實踐,有了一些心得體會,整理成這篇文章。過去在RLHF的初學(xué)階段,有一個問題最直接地困惑著我:

  • 如何在NLP語境下理解強(qiáng)化學(xué)習(xí)的框架?例如,我知道強(qiáng)化學(xué)習(xí)中有Agent、Environment、Reward、State等要素,但是在NLP語境中,它們指什么?語言模型又是如何根據(jù)獎勵做更新的?

為了解答這個問題,我翻閱了很多資料,看了許多的公式推導(dǎo),去研究RLHF的整體框架和loss設(shè)計。雖然吭吭哧哧地入門了,但是這個過程實在痛苦,最主要的原因是:理論的部分太多,直觀的解釋太少。

所以,在寫這篇文章時,我直接從一個RLHF開源項目源碼入手(deepspeed-chat),根據(jù)源碼的實現(xiàn)細(xì)節(jié),給出盡可能豐富的訓(xùn)練流程圖,并對所有的公式給出直觀的解釋。希望可以幫助大家更具象地感受RLHF的訓(xùn)練流程。對于沒有強(qiáng)化學(xué)習(xí)背景的朋友,也可以無痛閱讀本文。關(guān)于RLHF,各家的開源代碼間都會有一些差異,同時也不止PPO一種RLHF方式。感興趣的朋友,也可以讀讀別家的源碼,做一些對比。后續(xù)有時間,這個系列也會對各種RLHF方式進(jìn)行比較。

整體內(nèi)容如下:
【一、強(qiáng)化學(xué)習(xí)概述】
1.1 強(qiáng)化學(xué)習(xí)整體流程
1.2 價值函數(shù)
【二、NLP中的強(qiáng)化學(xué)習(xí)】
【三、RLHF中的四個重要角色】
3.1 Actor Model
3.2 Reference Model
3.3 Critic Model
3.4 Reward Model
【四、RLHF中的loss計算】
4.1 Actor loss
(1) 直觀設(shè)計
(2) 引入優(yōu)勢
(3) 重新設(shè)計獎勵函數(shù)
(4) 重新設(shè)計優(yōu)勢
(5)ppo_epoch: 引入新約束,提升訓(xùn)練效率
(6) Actor loss小結(jié)
【五、Critic loss】
(1) 實際收益優(yōu)化
(2) 預(yù)估收益優(yōu)化

一、強(qiáng)化學(xué)習(xí)概述

1. 強(qiáng)化學(xué)習(xí)整體流程

f6713c9a-b223-11ee-8b88-92fbcf53809c.png
  • 強(qiáng)化學(xué)習(xí)的兩個實體:智能體(Agent)環(huán)境(Environment)
  • 強(qiáng)化學(xué)習(xí)中兩個實體的交互:
    • 狀態(tài)空間S:S即為State,指環(huán)境中所有可能狀態(tài)的集合
    • 動作空間A:A即為Action,指智能體所有可能動作的集合
    • 獎勵R:R即為Reward,指智能體在環(huán)境的某一狀態(tài)下所獲得的獎勵。

以上圖為例,智能體與環(huán)境的交互過程如下:

  • 在時刻,環(huán)境的狀態(tài)為,達(dá)到這一狀態(tài)所獲得的獎勵為
  • 智能體觀測到與,采取相應(yīng)動作
  • 智能體采取后,環(huán)境狀態(tài)變?yōu)椋玫较鄳?yīng)的獎勵

智能體在這個過程中學(xué)習(xí),它的最終目標(biāo)是:找到一個策略,這個策略根據(jù)當(dāng)前觀測到的環(huán)境狀態(tài)和獎勵反饋,來選擇最佳的動作。

1.2 價值函數(shù)

在1.1中,我們談到了獎勵值,它表示環(huán)境進(jìn)入狀態(tài)下的即時獎勵。

但如果只考慮即時獎勵,目光似乎太短淺了:當(dāng)下的狀態(tài)和動作會影響到未來的狀態(tài)和動作,進(jìn)而影響到未來的整體收益。

所以,一種更好的設(shè)計方式是:t時刻狀態(tài)s的總收益 = 身處狀態(tài)s能帶來的即時收益 + 從狀態(tài)s出發(fā)后能帶來的未來收益。寫成表達(dá)式就是:

其中:

  • :時刻的總收益,注意這個收益蘊涵了“即時”和“未來”的概念
  • :時刻的即時收益
  • :時刻的總收益,注意這個收益蘊涵了“即時”和“未來”的概念。而對來說就是“未來”。
  • :折扣因子。它決定了我們在多大程度上考慮將“未來收益”納入“當(dāng)下收益”。

注:在這里,我們不展開討論RL中關(guān)于價值函數(shù)的一系列假設(shè)與推導(dǎo),而是直接給出一個便于理解的簡化結(jié)果,方便沒有RL背景的朋友能傾注更多在“PPO策略具體怎么做”及“對PPO的直覺理解”上。

二、NLP中的強(qiáng)化學(xué)習(xí)

我們在第一部分介紹了通用強(qiáng)化學(xué)習(xí)的流程,那么我們要怎么把這個流程對應(yīng)到NLP任務(wù)中呢?換句話說,NLP任務(wù)中的智能體、環(huán)境、狀態(tài)、動作等等,都是指什么呢?

f68cc596-b223-11ee-8b88-92fbcf53809c.png

回想一下我們對NLP任務(wù)做強(qiáng)化學(xué)習(xí)(RLHF)的目的:我們希望給模型一個prompt,讓模型能生成符合人類喜好的response。再回想一下gpt模型做推理的過程:每個時刻只產(chǎn)生一個token,即token是一個一個蹦出來的,先有上一個token,再有下一個token。

復(fù)習(xí)了這兩點,現(xiàn)在我們可以更好解讀上面這張圖了:

  • 我們先喂給模型一個prompt,期望它能產(chǎn)出符合人類喜好的response

  • 在時刻,模型根據(jù)上文,產(chǎn)出一個token,這個token即對應(yīng)著強(qiáng)化學(xué)習(xí)中的動作,我們記為。因此不難理解,在NLP語境下,強(qiáng)化學(xué)習(xí)任務(wù)的動作空間就對應(yīng)著詞表。

  • 在時刻,模型產(chǎn)出token 對應(yīng)著的即時收益為,總收益為(復(fù)習(xí)一下,蘊含著“即時收益”與“未來收益”兩個內(nèi)容)。這個收益即可以理解為“對人類喜好的衡量”。此刻,模型的狀態(tài)從變?yōu)?/strong>,也就是從“上文”變成“上文 + 新產(chǎn)出的token”

  • 在NLP語境下,智能體是語言模型本身,環(huán)境則對應(yīng)著它產(chǎn)出的語料

這樣,我們就大致解釋了NLP語境下的強(qiáng)化學(xué)習(xí)框架,不過針對上面這張圖,你可能還有以下問題:

1)問題1:圖中的下標(biāo)是不是寫得不太對?例如根據(jù)第一部分的介紹,應(yīng)該對應(yīng)著,應(yīng)該對應(yīng)著,以此類推?

答:你說的對。但這里我們不用太糾結(jié)下標(biāo)的問題,只需要記住在對應(yīng)的response token位置,會產(chǎn)生相應(yīng)的即時獎勵和總收益即可。之所以用圖中這樣的下標(biāo),是更方便我們后續(xù)理解代碼。

(2)問題2:我知道肯定是由語言模型產(chǎn)生的,那么是怎么來的呢,也是語言模型產(chǎn)生的嗎?

答:先直接說結(jié)論,是由我們的語言模型產(chǎn)生的,則分別由另外兩個模型來產(chǎn)生,在后文中我們會細(xì)說。

(3)問題3:語言模型的參數(shù)在什么時候更新?是觀測到一個,就更新一次參數(shù),然后再去產(chǎn)生嗎?

答:當(dāng)然不是。你只看到某個時刻的收益,就急著用它更新模型,這也太莽撞了。我們肯定是要等有足夠的觀測數(shù)據(jù)了(例如等模型把完整的response生成完),再去更新它的參數(shù)。這一點我們也放在后文細(xì)說。

(4)問題4:再談?wù)?/strong>吧,在NLP的語境下我還是不太理解它們。

答:

  • 首先,“收益”的含義是“對人類喜好的衡量”
  • :即時收益,指語言模型當(dāng)下產(chǎn)生token 帶來的收益
  • :實際期望總收益(即時+未來),指對語言模型“當(dāng)下產(chǎn)生token ,一直到整個response生產(chǎn)結(jié)束”后的期收益預(yù)估。因為當(dāng)下語言模型還沒產(chǎn)出后的token,所以我們只是對它之后一系列動作的收益做了估計,因而稱為“期望總收益”。

三、RLHF中的四個重要角色

在本節(jié)中,我們在第二部分的基礎(chǔ)上更進(jìn)一步:更詳細(xì)理清NLP語境下RLHF的運作流程。

我們從第二部分中已經(jīng)知道:生成token 和對應(yīng)收益的并不是一個模型。那么在RLHF中到底有幾個模型?他們是怎么配合做訓(xùn)練的?而我們最終要的是哪個模型?

f69ba4c6-b223-11ee-8b88-92fbcf53809c.png

如上圖,在RLHF-PPO階段,一共有四個主要模型,分別是:

  • Actor Model:演員模型,這就是我們想要訓(xùn)練的目標(biāo)語言模型
  • Critic Model:評論家模型,它的作用是預(yù)估總收益
  • Reward Model:獎勵模型,它的作用是計算即時收益
  • Reference Model:參考模型,它的作用是在RLHF階段給語言模型增加一些“約束”,防止語言模型訓(xùn)歪(朝不受控制的方向更新,效果可能越來越差)

其中:

  • Actor/Critic Model在RLHF階段是需要訓(xùn)練的(圖中給這兩個模型加了粗邊,就是表示這個含義);而Reward/Reference Model參數(shù)凍結(jié)的。
  • Critic/Reward/Reference Model共同組成了一個“獎勵-loss”計算體系(我自己命名的,為了方便理解),我們綜合它們的結(jié)果計算loss,用于更新Actor和Critic Model

我們把這四個部分展開說說。

3.1 Actor Model (演員模型)

正如前文所說,Actor就是我們想要訓(xùn)練的目標(biāo)語言模型。我們一般用SFT階段產(chǎn)出的SFT模型來對它做初始化。

f6ac0bf4-b223-11ee-8b88-92fbcf53809c.png

我們的最終目的是讓Actor模型能產(chǎn)生符合人類喜好的response。所以我們的策略是,先喂給Actor一條prompt(這里假設(shè)batch_size = 1,所以是1條prompt),讓它生成對應(yīng)的response。然后,我們再將“prompt + response"送入我們的“獎勵-loss”計算體系中去算得最后的loss,用于更新actor。

3.2 Reference Model(參考模型)

Reference Model(以下簡稱Ref模型)一般也用SFT階段得到的SFT模型做初始化,在訓(xùn)練過程中,它的參數(shù)是凍結(jié)的。Ref模型的主要作用是防止Actor”訓(xùn)歪”,那么它具體是怎么做到這一點的呢?

f6bcedfc-b223-11ee-8b88-92fbcf53809c.png

“防止模型訓(xùn)歪”換一個更詳細(xì)的解釋是:我們希望訓(xùn)練出來的Actor模型既能達(dá)到符合人類喜好的目的,又盡量讓它和SFT模型不要差異太大。簡言之,我們希望兩個模型的輸出分布盡量相似。那什么指標(biāo)能用來衡量輸出分布的相似度呢?我們自然而然想到了KL散度。

如圖所示:

  • 對Actor模型,我們喂給它一個prompt,它正常輸出對應(yīng)的response。那么response中每一個token肯定有它對應(yīng)的log_prob結(jié)果呀,我們把這樣的結(jié)果記為log_probs

  • 對Ref模型,我們把Actor生成的"prompt + response"喂給它,那么它同樣能給出每個token的log_prob結(jié)果,我們記其為ref_log_probs

  • 那么這兩個模型的輸出分布相似度就可以用ref_log_probs - log_probs來衡量,我們可以從兩個方面來理解這個公式:

    • 從直覺上理解,ref_log_probs越高,說明Ref模型對Actor模型輸出的肯定性越大。即Ref模型也認(rèn)為,對于某個,輸出某個的概率也很高()。這時可以認(rèn)為Actor模型較Ref模型沒有訓(xùn)歪

    • 從KL散度上理解,

      ,這個值越小意味著兩個分布的相似性越高。而這個值越小等價于ref_log_probs - log_probs越大

注:你可能已經(jīng)注意到,按照KL散度的定義,這里寫成log_probs - ref_log_probs更合適一些。這里之所以寫成ref_log_probs - log_probs,是為了方便大家從直覺上了解這個公式。

現(xiàn)在,我們已經(jīng)知道怎么利用Ref模型和KL散度來防止Actor訓(xùn)歪了。KL散度將在后續(xù)被用于loss的計算,我們在后文中會詳細(xì)解釋。

3.3 Critic Model(評論家模型)

Critic Model用于預(yù)測期望總收益,和Actor模型一樣,它需要做參數(shù)更新。實踐中,Critic Model的設(shè)計和初始化方式也有很多種,例如和Actor共享部分參數(shù)、從RW階段的Reward Model初始化而來等等。我們講解時,和deepspeed-chat的實現(xiàn)保持一致:從RW階段的Reward Model初始化而來。

你可能想問:訓(xùn)練Actor模型我能理解,但我還是不明白,為什么要單獨訓(xùn)練一個Critic模型用于預(yù)測收益呢?

這是因為,當(dāng)我們在前文討論總收益(即時 + 未來)時,我們是站在上帝視角的,也就是這個就是客觀存在的、真正的總收益。但是我們在訓(xùn)練模型時,就沒有這個上帝視角加成了,也就是在時刻,我們給不出客觀存在的總收益,我們只能訓(xùn)練一個模型去預(yù)測它。

所以總結(jié)來說,在RLHF中,我們不僅要訓(xùn)練模型生成符合人類喜好的內(nèi)容的能力(Actor),也要提升模型對人類喜好量化判斷的能力(Critic)。這就是Critic模型存在的意義。我們來看看它的大致架構(gòu):

f6ca783c-b223-11ee-8b88-92fbcf53809c.png

deepspeed-chat采用了Reward模型作為它的初始化,所以這里我們也按Reward模型的架構(gòu)來簡單畫畫它。你可以簡單理解成,Reward/Critic模型和Actor模型的架構(gòu)是很相似的(畢竟輸入都一樣),同時,它在最后一層增加了一個Value Head層,該層是個簡單的線形層,用于將原始輸出結(jié)果映射成單一的值。

在圖中,表示Critic模型對時刻及未來(response完成)的收益預(yù)估。

3.4 Reward Model(獎勵模型)

Reward Model用于計算生成token 的即時收益,它就是RW階段所訓(xùn)練的獎勵模型,在RLHF過程中,它的參數(shù)是凍結(jié)的。

你可能想問:為什么Critic模型要參與訓(xùn)練,而同樣是和收益相關(guān)的Reward模型的參數(shù)就可以凍結(jié)呢?

這是因為,Reward模型是站在上帝視角的。這個上帝視角有兩層含義:

  • 第一點,Reward模型是經(jīng)過和“估算收益”相關(guān)的訓(xùn)練的,因此在RLHF階段它可以直接被當(dāng)作一個能產(chǎn)生客觀值的模型。
  • 第二點,Reward模型代表的含義就是“即時收益”,你的token 已經(jīng)產(chǎn)生,因此即時收益自然可以立刻算出。

你還可能想問:我已經(jīng)用Critic預(yù)測出了,而這個包含了“即時”和“未來”的概念,那我還需要代表“即時”的做什么呢?直接用不就好了嗎?

為了解答這個問題,我們先回顧下1.2部分中給出的價值函數(shù):

這個函數(shù)告訴我們,我們當(dāng)前可以用兩個結(jié)果來表示時刻的總收益:

  • 結(jié)果1:Critic模型預(yù)測的

  • 結(jié)果2:Reward模型預(yù)測的和critic模型預(yù)測的

那么哪一個結(jié)果更靠近上帝視角給出的客觀值呢?當(dāng)然是結(jié)果2,因為結(jié)果1全靠預(yù)測,而結(jié)果2中的是事實數(shù)據(jù)。

我們知道Critic模型也是參與參數(shù)更新的,我們可以用MSE(上帝視角的客觀收益-Critic模型預(yù)測的收益)來衡量它的loss。但是上帝視角的客觀收益我們是不知道的,只能用已知事實數(shù)據(jù)去逼近它,所以我們就用來做近似。這就是同時存在的意義

Reward模型和critic模型非常相似,這里我們就只給出架構(gòu)圖,不再做過多的說明。關(guān)于Reward模型的訓(xùn)練過程,后續(xù)有時間也會出個原理和代碼解析。

f6d4d836-b223-11ee-8b88-92fbcf53809c.png

四、RLHF中的loss計算

到目前為止,我們已經(jīng)基本了解了RLHF的訓(xùn)練框架,以及其中的四個重要角色(訓(xùn)練一個RLHF,有4個模型在硬件上跑,可想而知對存儲的壓力)。在本節(jié)中,我們一起來解讀RLHF的loss計算方式。在解讀中,我們會再一次理一遍RLHF的整體訓(xùn)練過程,填補(bǔ)相關(guān)細(xì)節(jié)。在這之后,我們就可以來看代碼解析了。

在第三部分的講解中,我們知道Actor和Critic模型都會做參數(shù)更新,所以我們的loss也分成2個:

  • Actor loss:用于評估Actor是否產(chǎn)生了符合人類喜好的結(jié)果,將作用于Actor的BWD上。

  • Critic loss:用于評估Critic是否正確預(yù)測了人類的喜好,將作用于Critic的BWD上。

我們詳細(xì)來看這兩者。

4.1 Actor loss

(1)直觀設(shè)計

我們先來看一個直觀的loss設(shè)計方式:

  • Actor接收到當(dāng)前上文,產(chǎn)出token()

  • Critic根據(jù),產(chǎn)出對總收益的預(yù)測

  • 那么Actor loss可以設(shè)計為:

求和符號表示我們只考慮response部分所有token的loss,為了表達(dá)簡便,我們先把這個求和符號略去(下文也是同理),也就是說:

我們希望minimize這個actor_loss。

這個設(shè)計的直觀解釋是:

  • 當(dāng)時,意味著Critic對Actor當(dāng)前采取的動作給了正向反饋,因此我們就需要在訓(xùn)練迭代中提高,這樣就能達(dá)到減小loss的作用。
  • 當(dāng)時,意味著Critic對Actor當(dāng)前采取的動作給了負(fù)向反饋,因此我們就需要在訓(xùn)練迭代中降低,這樣就能到達(dá)到減小loss的作用。

一句話總結(jié):這個loss設(shè)計的含義是,對上文而言,如果token產(chǎn)生的收益較高,那就增大它出現(xiàn)的概率,否則降低它出現(xiàn)的概率。

(2)引入優(yōu)勢(Advantage)

在開始講解之前,我們舉個小例子:

假設(shè)在王者中,中路想支援發(fā)育路,這時中路有兩種選擇:1. 走自家野區(qū)。2. 走大龍路。

中路選擇走大龍路,當(dāng)她做出這個決定后,Critic告訴她可以收1個人頭。結(jié)果,此刻對面打野正在自家采靈芝,對面也沒有什么茍草英雄,中路一路直上,最終收割2個人頭。

因為實際收割的人頭比預(yù)期要多1個,中路嘗到了甜頭,所以她增大了“支援上路走大龍路”的概率。

這個多出來的“甜頭”,就叫做“優(yōu)勢”(Advantage)。

對NLP任務(wù)來說,如果Critic對的總收益預(yù)測為,但實際執(zhí)行后的總收益是,我們就定義優(yōu)勢為:

我們用替換掉,則此刻actor_loss變?yōu)椋?/span>

(3)重新設(shè)計

總結(jié)一下,到目前為止,我們的actor_loss形式為:

其中,

同時注意,這個actor_loss應(yīng)該是response的所有token loss的sum或者avg。這里為了表達(dá)方便,我們的公式略去了求和或求平均的符號。

按照這個理解,應(yīng)該表示每個Actor產(chǎn)出token 帶來的即時收益,正如下圖所示(其中表示最后一個時刻):

f6e124e2-b223-11ee-8b88-92fbcf53809c.png

但在deepspeed-chat的RLHF實踐中,對做了另一種設(shè)計:

  • :常量,可以理解成是一個控制比例的縮放因子,在deepspeed-chat中默認(rèn)設(shè)為0.1

  • :這一項你是不是非常眼熟,這就是我們在3.2部分介紹的Actor和Ref模型間的KL散度呀,寫成更容易理解的形式,就是ref_log_probs - log_probs。在3.2中我們說過,為了防止模型訓(xùn)歪,我們需要把這個KL散度加入loss計算中,所以這里我們就在做這件事

基于這些,上面這個對的設(shè)計可理解成:

  • 當(dāng)時,我們更加關(guān)心Actor是否有在Ref的約束下生產(chǎn)token
  • 當(dāng)時,我們不僅關(guān)心Actor是否遵從了Ref的約束,也關(guān)心真正的即時收益

需要注意的是,的設(shè)計并不只有這一種(其實我覺得只取時刻的不太合理,deepspeed在自己的代碼注釋中也有提過,可以嘗試把最后一個時刻的替換成所有token的平均值;或者在我來看時也應(yīng)該納入對的考慮)。

代碼實踐如下:

defcompute_rewards(self,prompts,log_probs,ref_log_probs,reward_score,
action_mask):
"""
 reward_function:計算最終的reward分?jǐn)?shù)
復(fù)習(xí)一下幾個相關(guān)參數(shù)的默認(rèn)值:
self.kl_ctl=0.1
self.clip_reward_value=5

對于batch中的某個prompt來說,它最終的reward分?jǐn)?shù)為:
(1)先計算actor和ref_model的logit相似度:-self.kl_ctl *(log_probs - ref_log_probs)
其實寫成self.kl_ctl*(ref_log_probs-log_probs)更好理解些
這個值越大,說明ref_model對actor生成的結(jié)果的認(rèn)可度越高(即表明rlhf沒有訓(xùn)歪),
沒有訓(xùn)歪的情況下我們也應(yīng)該給模型一些獎勵,這個獎勵就是self.kl_ctl*(ref_log_probs-log_probs)

(2)由于我們只取最后一個token對應(yīng)位置的分?jǐn)?shù)作為reward_score,因此我們只需要:
self.kl_ctl*(ref_log_probs-log_probs)的最后一位+reward_score

(3)同時我們對reward_score也做了大小限制,最大不超過self.clip_reward_value(超過統(tǒng)一給成self.clip_reward_value),
最小不低于-self.clip_reward_value(低于統(tǒng)一給成-self.clip_reward_value)

(4)最后返回的rewards大小為:(batch_size, 各條數(shù)據(jù)的長度),對batch中的每條數(shù)據(jù)來說:
- response的最后一位:self.kl_ctl *(ref_log_probs - log_probs)的最后一位+ reward_score
- response的其余位置:self.kl_ctl *(ref_log_probs - log_probs)

"""

kl_divergence_estimate=-self.kl_ctl*(log_probs-ref_log_probs)
rewards=kl_divergence_estimate
#---------------------------------------------------------------------------------------------------
#response開始的位置
#(因為我們對prompt做過padding處理,因此batch中每個prompt長度一致,也就意味著每個response開始的位置一致)
#(所以這里start是不加s的,只是一個int)
#---------------------------------------------------------------------------------------------------
start=prompts.shape[1]-1
#---------------------------------------------------------------------------------------------------
#response結(jié)束的位置
#(因為一個batch中,每個response的長度不一樣,所以response的結(jié)束位置也不一樣)
#(所以這里end是加s的,ends的尺寸是(batch_size,)
#---------------------------------------------------------------------------------------------------
ends=start+action_mask[:,start:].sum(1)+1
#---------------------------------------------------------------------------------------------------
#對rewards_score做限制
#---------------------------------------------------------------------------------------------------
reward_clip=torch.clamp(reward_score,-self.clip_reward_value,
self.clip_reward_value)
batch_size=log_probs.shape[0]
forjinrange(batch_size):
rewards[j,start:ends[j]][-1]+=reward_clip[j]#

returnrewards

(4)重新設(shè)計優(yōu)勢

好,再總結(jié)一下,目前為止我們的actor_loss為:

其中,

同時,我們對進(jìn)行來改造,使其能夠衡量Actor模型是否遵從了Ref模型的約束。

現(xiàn)在我們把改造焦點放在上,回想一下,既然對于收益而言,分為即時和未來,那么對于優(yōu)勢而言,是不是也能引入對未來優(yōu)勢的考量呢?這樣,我們就可以把改寫成如下形式:

(熟悉強(qiáng)化學(xué)習(xí)的朋友應(yīng)該能一眼看出這是GAE,這里我們不打算做復(fù)雜的介紹,一切都站在直覺的角度理解)

其中,新引入的也是一個常量,可將其理解為權(quán)衡因子,直覺上看它控制了在計算當(dāng)前優(yōu)勢時對未來優(yōu)勢的考量。(從強(qiáng)化學(xué)習(xí)的角度上,它控制了優(yōu)勢估計的方差和偏差)

看到這里,你可能想問:這個代表未來優(yōu)勢的,我要怎么算呢?

注意到,對于最后一個時刻,它的未來收益()和未來優(yōu)勢()都是0,也就是,這是可以直接算出來的。而有了,我們不就能從后往前,通過動態(tài)規(guī)劃的方法,把所有時刻的優(yōu)勢都依次算出來了嗎?

代碼實踐如下(其中返回值中的returns表示實際收益,將被用于計算Critic模型的loss,可以參見4.2,其余細(xì)節(jié)都在代碼注釋中):

defget_advantages_and_returns(self,values,rewards,start):
"""
Adoptedfromhttps://github.com/CarperAI/trlx/blob/main/trlx/models/modeling_ppo.py#L134

沒有引入GAE前的t時刻的優(yōu)勢值:
detal_t=r_t+gamma*V_t+1-V_t
其中:
-r_t表示t時刻的即時收益
-V_t+1表示未來時刻的預(yù)期收益
-r_t+gamma*V_t+1可理解成t時刻的實際預(yù)期收益
-V_t可理解成t時刻的預(yù)估預(yù)期收益(是模型,例如criticmodel自己估算出來的)

引入GAE后的t時刻的優(yōu)勢值:
A_t=delta_t+gamma*lambda*A_t+1
粗暴理解為在t時刻時,不僅考慮當(dāng)下優(yōu)勢,還考慮了未來的優(yōu)勢
為了知道A_t, 我們得知道A_t+1,所以在本算法中采取了從后往前做動態(tài)規(guī)劃求解的方法,也即:
假設(shè)T是最后一個時刻,則有A_T+1=0,所以有:A_T=delta_T
知道了A_T,就可以依次往前倒推,把A_t-1,A_t-2之類都算出來了

引入GAE后t時刻的實際預(yù)期收益
returns_t=A_t+V_t
=delta_t+gamma*lambda*A_t+1+V_t
=r_t+gamma*V_t+1-V_t+gamma*lambda*A_t+1+V_t
=r_t+gamma*(V_t+1+lambda*A_t+1)

注意,這里不管是advantages還是returns,都只算response的部分
"""

#Adoptedfromhttps://github.com/CarperAI/trlx/blob/main/trlx/models/modeling_ppo.py#L134
lastgaelam=0
advantages_reversed=[]
length=rewards.size()[-1]
#注意這里用了reversed,是采取從后往前倒推計算的方式
fortinreversed(range(start,length)):
nextvalues=values[:,t+1]ift1else0.0
delta=rewards[:,t]+self.gamma*nextvalues-values[:,t]
lastgaelam=delta+self.gamma*self.lam*lastgaelam
advantages_reversed.append(lastgaelam)
advantages=torch.stack(advantages_reversed[::-1],dim=1)#優(yōu)勢
returns=advantages+values[:,start:]#實際收益
#values:預(yù)期收益
returnadvantages.detach(),returns

(5) PPO-epoch: 引入新約束

總結(jié)一下,目前為止我們的actor_loss為:

其中,

同時

  • 我們已經(jīng)對進(jìn)行來改造,使其能夠衡量Actor模型是否遵從了Ref模型的約束。
  • 我們已經(jīng)對進(jìn)行改造,使其不僅考慮了當(dāng)前時刻的優(yōu)勢,還考慮了未來的優(yōu)勢

基于這些改造,我們重新理一遍RLHF-PPO的訓(xùn)練過程。

f6f2920e-b223-11ee-8b88-92fbcf53809c.png
  • 第一步,我們準(zhǔn)備一個batch的prompts

  • 第二步,我們將這個batch的prompts喂給Actor模型,讓它生成對應(yīng)的responses

  • 第三步,我們把prompt+responses喂給我們的Critic/Reward/Reference模型,讓它生成用于計算actor/critic loss的數(shù)據(jù),按照強(qiáng)化學(xué)習(xí)的術(shù)語,我們稱這些數(shù)據(jù)為經(jīng)驗(experiences)。critic loss我們將在后文做詳細(xì)講解,目前我們只把目光聚焦到actor loss上

  • 第四步,我們根據(jù)這些經(jīng)驗,實際計算出actor/critic loss,然后更新Actor和Critic模型

這些步驟都很符合直覺,但是細(xì)心的你肯定發(fā)現(xiàn)了,文字描述中的第四步和圖例中的第四步有差異:圖中說,這一個batch的經(jīng)驗值將被用于n次模型更新,這是什么意思呢?

我們知道,在強(qiáng)化學(xué)習(xí)中,收集一個batch的經(jīng)驗是非常耗時的。對應(yīng)到我們RLHF的例子中,收集一次經(jīng)驗,它要等四個模型做完推理才可以,正是因此,一個batch的經(jīng)驗,只用于計算1次loss,更新1次Actor和Critic模型,好像有點太浪費了。

所以,我們自然而然想到,1個batch的經(jīng)驗,能不能用來計算ppo-epochs次loss,更新ppo-epochs次Actor和Critic模型?簡單寫一下偽代碼,我們想要:

#--------------------------------------------------------------
#初始化RLHF中的四個模型
#--------------------------------------------------------------
actor,critic,reward,ref=initialize_models()

#--------------------------------------------------------------
#訓(xùn)練
#--------------------------------------------------------------
#對于每一個batch的數(shù)據(jù)
foriinsteps:
#先收集經(jīng)驗值
exps=generate_experience(prompts,actor,critic,reward,ref)
#一個batch的經(jīng)驗值將被用于計算ppo_epochs次loss,更新ppo_epochs次模型
#這也意味著,當(dāng)你計算一次新loss時,你用的是更新后的模型
forjinppo_epochs:
actor_loss=cal_actor_loss(exps,actor)
critic_loss=cal_critic_loss(exps,critic)

actor.backward(actor_loss)
actor.step()

critc.backward(critic_loss)
critic.step()

而如果我們想讓一個batch的經(jīng)驗值被重復(fù)使用ppo_epochs次,等價于我們想要Actor在這個過程中,模擬和環(huán)境交互ppo_epochs次。舉個例子:

  • 如果1個batch的經(jīng)驗值只使用1次,那么在本次更新完后,Actor就吃新的batch,正常和環(huán)境交互,產(chǎn)出新的經(jīng)驗值

  • 但如果1個batch的經(jīng)驗值被使用ppo_epochs次,在這ppo_epochs中,Actor是不吃任何新數(shù)據(jù),不做任何交互的,所以我們只能讓Actor“模擬”一下和環(huán)境交互的過程,吐出一些新數(shù)據(jù)出來。

那怎么讓Actor模擬呢?很簡單,讓它觀察一下之前的數(shù)據(jù)長什么樣,讓它依葫蘆畫瓢,不就行了嗎?我們假設(shè)最開始吃batch,吐出經(jīng)驗的actor叫,而在偽代碼中,每次做完ppo_epochs而更新的actor叫,那么我們只要盡量保證每次更新后的能模仿最開始的那個,不就行了嗎?

誒!是不是很眼熟!兩個分布,通過什么方法讓它們相近!那當(dāng)然是KL散度!所以,再回到我們的actor_loss上來,它現(xiàn)在就可被改進(jìn)成:

我們再稍作一些改動將log去掉(這個其實不是“稍作改動去掉log”的事,是涉及到PPO中重要性采樣的相關(guān)內(nèi)容,大家有興趣可以參考https://www.cnblogs.com/xingzheai/p/15931681.html):

其中,表示真正吃了batch,產(chǎn)出經(jīng)驗值的Actor;P表示ppo_epochs中實時迭代更新的Actor,它在模仿的行為。所以這個公式從直覺上也可以理解成:在Actor想通過模擬交互的方式,使用一個batch的經(jīng)驗值更新自己時,它需要收到真正吃到batch的那個時刻的Actor的約束,這樣才能在有效利用batch,提升訓(xùn)練速度的基礎(chǔ)上,保持訓(xùn)練的穩(wěn)定。


但是,謹(jǐn)慎的你可能此時又有新的擔(dān)心了:雖然我們在更新Actor的過程中用做了約束,但如的約束能力不夠,比如說還是超出了可接受的范圍,那怎么辦?

很簡單,那就剪裁(clip)它吧!

我們給設(shè)置一個范圍,例如(0.8 ,1.2),也就是如果這個值一旦超過1.2,那就統(tǒng)一變成1.2;一旦小于0.8,那就統(tǒng)一變成0.8。這樣就能保證和的分布相似性在我們的掌控之內(nèi)了。此時actor_loss變?yōu)椋?/span>

這時要注意,如果超過變化范圍,將強(qiáng)制設(shè)定為一個常數(shù)后,就說明這一部分的loss和Actor模型無關(guān)了,而這項本身也與Actor無關(guān)。所以相當(dāng)于,在超過約束范圍時,我們停止對Actor模型進(jìn)行更新。

整體代碼如下:

defactor_loss_fn(self,logprobs,old_logprobs,advantages,mask):
"""
logprobs:實時計算的,response部分的prob(只有這個是隨著actor實時更新而改變的)
 old_logprobs:老策略中,response部分的prob (這個是固定的,不隨actor實時更新而改變)
 advantages:老策略中,response部分每個token對應(yīng)的優(yōu)勢(這個是固定的,不隨actor實時更新而改變)
 mask:老策略中,response部分對應(yīng)的mask情況這個是固定的,不隨actor實時更新而改變)

之所以要引入logprobs計算actor_loss,是因為我們不希望策略每次更新的幅度太大,防止模型訓(xùn)歪

self.cliprange:默認(rèn)值是0.2
"""
##policygradientloss
#-------------------------------------------------------------------------------------
#計算新舊策略間的KL散度
#-------------------------------------------------------------------------------------
log_ratio=(logprobs-old_logprobs)*mask
ratio=torch.exp(log_ratio)
#-------------------------------------------------------------------------------------
#計算原始loss和截斷l(xiāng)oss
#-------------------------------------------------------------------------------------
pg_loss1=-advantages*ratio
pg_loss2=-advantages*torch.clamp(ratio,1.0-self.cliprange,1.0+self.cliprange)
pg_loss=torch.sum(torch.max(pg_loss1,pg_loss2)*mask)/mask.sum()#最后是取每個非mask的responsetoken的平均loss作為最終loss
returnpg_loss

(6)Actor loss小結(jié)

(1)~(5)中我們一步步樹立了actor_loss的改進(jìn)過程,這里我們就做一個總結(jié)吧:

其中:

  • 我們已經(jīng)對進(jìn)行來改造,使其能夠衡量Actor模型是否遵從了Ref模型的約束

  • 我們已經(jīng)對進(jìn)行改造,使其不僅考慮了當(dāng)前時刻的優(yōu)勢,還考慮了未來的優(yōu)勢

  • 我們重復(fù)利用了1個batch的數(shù)據(jù),使本來只能被用來做1次模型更新的它現(xiàn)在能被用來做ppo_epochs次模型更新。我們使用真正吃了batch,產(chǎn)出經(jīng)驗值的那個時刻的Actor分布來約束ppo_epochs中更新的Actor分布

  • 我們考慮了剪裁機(jī)制(clip),在ppo_epochs次更新中,一旦Actor的更新幅度超過我們的控制范圍,則不對它進(jìn)行參數(shù)更新。

4.2 Critic loss

我們知道,1個batch產(chǎn)出的經(jīng)驗值,不僅被用來更新Actor,還被用來更新Critic。對于Critic loss,我們不再像Actor loss一樣給出一個“演變過程”的解讀,我們直接來看它最后的設(shè)計。

首先,在之前的解說中,你可能有這樣一個印象:

  • :Critic對t時刻的總收益的預(yù)估,這個總收益包含即時和未來的概念(預(yù)估收益)

  • :Reward計算出的即時收益,Critic預(yù)測出的及之后時候的收益的折現(xiàn),這是比更接近t時刻真值總收益的一個值(實際收益)

所以,我們的第一想法是:

現(xiàn)在,我們對“實際收益”和“預(yù)估收益”都做一些優(yōu)化。

(1)實際收益優(yōu)化

我們原始的實際收益為,但是當(dāng)我們在actor_loss中引入“優(yōu)勢”的概念時,“優(yōu)勢”中刻畫了更為豐富的實時收益信息,所以,我們將實際收益優(yōu)化為:

(2)預(yù)估收益優(yōu)化

我們原始的預(yù)估收益為。

類比于Actor,Critic模型在ppo_epochs的過程中也是不斷更新的。所以這個可以理解成是,也就是真正吃了batch,參與產(chǎn)出經(jīng)驗的那個時候的Critic產(chǎn)出的收益預(yù)測結(jié)果。

我們同樣想用舊模型去約束新模型,但對于Critic我們采用的約束策略就比較簡單了,我們直接看代碼,從中可以看出,我們用老設(shè)計了了一個變動范圍,然后用這個變動范圍去約束新

#self.cliprange_value是一個常量
#old_values:老critic的預(yù)測結(jié)果
# values:新critic的預(yù)測結(jié)果
values_clipped=torch.clamp(
values,
old_values-self.cliprange_value,
old_values+self.cliprange_value,
)

那么最終我們就取實際收益和預(yù)估收益的MSE做為loss就好,這里注意,計算實際收益時都是老Critic(真正吃了batch的那個)產(chǎn)出的結(jié)果,而預(yù)估收益是隨著ppo_epochs而變動的。

代碼如下:

defcritic_loss_fn(self,values,old_values,returns,mask):
"""
values:實時critic跑出來的預(yù)估預(yù)期收益(是變動的,隨著ppoepoch迭代而改變)
 old_values:老critic跑出來的預(yù)估預(yù)期收益(是固定值)
 returns:實際預(yù)期收益
 mask:response部分的mask

self.cliprange_value=0.2
"""
##valueloss
#用舊的value去約束新的value
values_clipped=torch.clamp(
values,
old_values-self.cliprange_value,
old_values+self.cliprange_value,
)
ifself.compute_fp32_loss:
values=values.float()
values_clipped=values_clipped.float()

#critic模型的loss定義為(預(yù)估預(yù)期收益-實際預(yù)期收益)**2
vf_loss1=(values-returns)**2
vf_loss2=(values_clipped-returns)**2
vf_loss=0.5*torch.sum(
torch.max(vf_loss1,vf_loss2)*mask)/mask.sum()#同樣,最后也是把criticloss平均到每個token上
returnvf_loss

至此,關(guān)于RLHF-PPO訓(xùn)練的核心部分和代碼解讀就講完了,建議大家親自閱讀源碼,如果有硬件條件,可以動手實踐。源碼地址:

https://github.com/microsoft/DeepSpeedExamples/tree/master/applications/DeepSpeed-Chat/training/step3_rlhf_finetuning


聲明:本文內(nèi)容及配圖由入駐作者撰寫或者入駐合作網(wǎng)站授權(quán)轉(zhuǎn)載。文章觀點僅代表作者本人,不代表電子發(fā)燒友網(wǎng)立場。文章及其配圖僅供工程師學(xué)習(xí)之用,如有內(nèi)容侵權(quán)或者其他違規(guī)問題,請聯(lián)系本站處理。 舉報投訴
  • 源碼
    +關(guān)注

    關(guān)注

    8

    文章

    641

    瀏覽量

    29213
  • 強(qiáng)化學(xué)習(xí)

    關(guān)注

    4

    文章

    266

    瀏覽量

    11254
  • 大模型
    +關(guān)注

    關(guān)注

    2

    文章

    2450

    瀏覽量

    2706

原文標(biāo)題:圖解大模型RLHF系列之:人人都能看懂的PPO原理與源碼解讀

文章出處:【微信號:GiantPandaCV,微信公眾號:GiantPandaCV】歡迎添加關(guān)注!文章轉(zhuǎn)載請注明出處。

收藏 人收藏

    評論

    相關(guān)推薦

    運行tas5086gui軟件后要實現(xiàn)低音要加載cfg文件嗎?

    請問運行tas5086gui軟件后要實現(xiàn)低音 要加載cfg文件嗎?能告訴我步驟嗎 ?pdf文件不怎么能看懂。謝謝了?。?
    發(fā)表于 11-06 06:43

    解讀 MEMS 可編程 LVCMOS 振蕩器 SiT1602 系列:精準(zhǔn)頻率的創(chuàng)新

    解讀 MEMS 可編程 LVCMOS 振蕩器 SiT1602 系列:精準(zhǔn)頻率的創(chuàng)新
    的頭像 發(fā)表于 08-09 15:39 ?366次閱讀
    <b class='flag-5'>解讀</b> MEMS 可編程 LVCMOS 振蕩器 SiT1602 <b class='flag-5'>系列</b>:精準(zhǔn)頻率的創(chuàng)新<b class='flag-5'>之</b>選

    解讀 MEMS 可編程 LVCMOS 振蕩器 SiT8008 系列:精準(zhǔn)與靈活的時脈

    解讀 MEMS 可編程 LVCMOS 振蕩器 SiT8008 系列:精準(zhǔn)與靈活的時脈
    的頭像 發(fā)表于 08-09 10:29 ?309次閱讀
    <b class='flag-5'>解讀</b> MEMS 可編程 LVCMOS 振蕩器 SiT8008 <b class='flag-5'>系列</b>:精準(zhǔn)與靈活的時脈<b class='flag-5'>之</b>選

    深度解讀 VCXO VG7050CDN:可變晶體振蕩器的卓越

    深度解讀 VCXO VG7050CDN:可變晶體振蕩器的卓越
    的頭像 發(fā)表于 07-24 10:58 ?361次閱讀

    解讀PyTorch模型訓(xùn)練過程

    PyTorch作為一個開源的機(jī)器學(xué)習(xí)庫,以其動態(tài)計算圖、易于使用的API和強(qiáng)大的靈活性,在深度學(xué)習(xí)領(lǐng)域得到了廣泛的應(yīng)用。本文將深入解讀PyTorch模型訓(xùn)練的全過程,包括數(shù)據(jù)準(zhǔn)備、模型構(gòu)建、訓(xùn)練循環(huán)、評估與保存等關(guān)鍵步驟,并結(jié)合
    的頭像 發(fā)表于 07-03 16:07 ?1072次閱讀

    labview實例源碼控壓取樣系統(tǒng)

    labview源碼,包含報表、曲線、通訊等
    發(fā)表于 06-06 11:23 ?1次下載

    這一次,騰訊要讓人人都能「用好」大模型

    “2024年的大模型百花齊放、各出奇招,而騰訊卻想要走一條自己的路?!苯裉斓拇?b class='flag-5'>模型ToB賽道,仍然是一個非共識的時代——大模型如何落地,給企業(yè)產(chǎn)生實際價值,即使一年過去,也沒有一個標(biāo)準(zhǔn)答案。而廠商
    的頭像 發(fā)表于 05-24 08:04 ?740次閱讀
    這一次,騰訊要讓<b class='flag-5'>人人</b><b class='flag-5'>都能</b>「用好」大<b class='flag-5'>模型</b>

    硬件工程師只要會照著芯片規(guī)格書畫外圍電路就夠了嗎?高級硬件工程師多了這項技能

    我們要關(guān)注的是芯片內(nèi)部的電路拓?fù)?,要想達(dá)到高級硬件工程師也必須能看懂芯片內(nèi)部的電路拓?fù)?,在?guī)格書中也叫做BLOCK DIAGRAM。 今天就介紹兩種穩(wěn)壓IC的內(nèi)部電路拓?fù)洹?/div>
    的頭像 發(fā)表于 05-11 08:46 ?1732次閱讀
    硬件工程師只要會照著芯片規(guī)格書畫外圍電路就夠了嗎?高級硬件工程師多了這項技能

    CubeAI-7.0.0生成的C語言代碼,神經(jīng)網(wǎng)絡(luò)運行函數(shù)是哪個,輸入輸出分別是哪個變量???

    CubeAI-7.0.0生成的C語言代碼,神經(jīng)網(wǎng)絡(luò)運行函數(shù)是哪個,輸入輸出分別是哪個變量???求教大佬,之前5.1.2還能看懂,現(xiàn)在7.0.0看不懂了。
    發(fā)表于 04-01 07:30

    微軟4月1日推出生成式AI安全產(chǎn)品“Securit Copilot”

    腳本反編程:自動解碼惡意軟件,實現(xiàn)無須手工逆向工程,讓每位分析師都能看懂入侵者的操作;分析復(fù)雜命令行腳本,實現(xiàn)自然語言解釋,找到相關(guān)實體并關(guān)聯(lián);
    的頭像 發(fā)表于 03-14 10:28 ?499次閱讀

    聊聊小公司如何做大模型

    通過SFT、DPO、RLHF等技術(shù)訓(xùn)練了領(lǐng)域?qū)懽?b class='flag-5'>模型。實測下來,在該領(lǐng)域?qū)懽魃希瑥?qiáng)于國內(nèi)大多數(shù)的閉源模型。
    的頭像 發(fā)表于 02-21 11:28 ?657次閱讀

    程序中的R地址都是什么意思?怎么樣才能看懂?

    程序中的R地址都是什么意思?怎么樣才能看懂? 在程序中,R地址通常指的是寄存器地址,它是用來存儲和訪問計算機(jī)中的數(shù)據(jù)的硬件部件。寄存器是計算機(jī)中最快的內(nèi)存形式,它位于中央處理器(CPU)內(nèi)部或與
    的頭像 發(fā)表于 02-18 10:49 ?1115次閱讀

    模型Reward Model的trick應(yīng)用技巧

    借助對比學(xué)習(xí)和元學(xué)習(xí)的方法。增加對比學(xué)習(xí)的loss,對比學(xué)習(xí)通過增強(qiáng)模型區(qū)分能力,來增強(qiáng)RM的對好壞的區(qū)分水平。元學(xué)習(xí)則使獎勵模型能夠維持區(qū)分分布外樣本的細(xì)微差異,這種方法可以用于迭代式的RLHF優(yōu)化。
    發(fā)表于 01-25 09:31 ?1025次閱讀
    大<b class='flag-5'>模型</b>Reward Model的trick應(yīng)用技巧

    OneFlow Softmax算子源碼解讀BlockSoftmax

    寫在前面:筆者這段時間工作太忙,身心俱疲,博客停更了一段時間,現(xiàn)在重新?lián)炱饋怼1疚闹饕?b class='flag-5'>解讀 OneFlow 框架的第二種 Softmax 源碼實現(xiàn)細(xì)節(jié),即 block 級別的 Softmax。
    的頭像 發(fā)表于 01-08 09:26 ?716次閱讀
    OneFlow Softmax算子<b class='flag-5'>源碼</b><b class='flag-5'>解讀</b><b class='flag-5'>之</b>BlockSoftmax

    OneFlow Softmax算子源碼解讀WarpSoftmax

    寫在前面:近來筆者偶然間接觸了一個深度學(xué)習(xí)框架 OneFlow,所以這段時間主要在閱讀 OneFlow 框架的 cuda 源碼。官方源碼基于不同場景分三種方式實現(xiàn) Softmax,本文主要介紹其中一種的實現(xiàn)過程,即 Warp 級別 Softmax,適用于矩陣寬度不超過 1
    的頭像 發(fā)表于 01-08 09:24 ?858次閱讀
    OneFlow Softmax算子<b class='flag-5'>源碼</b><b class='flag-5'>解讀</b><b class='flag-5'>之</b>WarpSoftmax