人手一個ChatGPT的夢想,就要實現(xiàn)了?
剛剛,微軟開源了一個可以在模型訓(xùn)練中加入完整RLHF流程的系統(tǒng)框架——DeepSpeed Chat。
也就是說,各種規(guī)模的高質(zhì)量類ChatGPT模型,現(xiàn)在都唾手可得了!
項目地址:https://Github.com/microsoft/DeepSpeed
一鍵解鎖千億級ChatGPT,輕松省錢15倍
眾所周知,由于OpenAI太不Open,開源社區(qū)為了讓更多人能用上類ChatGPT模型,相繼推出了LLaMa、Alpaca、Vicuna、Databricks-Dolly等模型。
但由于缺乏一個支持端到端的RLHF規(guī)模化系統(tǒng),目前類ChatGPT模型的訓(xùn)練仍然十分困難。而DeepSpeed Chat的出現(xiàn),正好補全了這個「bug」。
更亮的是,DeepSpeed Chat把成本大大地打了下來。
此前,昂貴的多GPU設(shè)置超出了許多研究者的能力范圍,并且,即使能訪問多GPU集群,現(xiàn)有的方法也無力負擔(dān)數(shù)千億參數(shù)ChatGPT模型的訓(xùn)練。
現(xiàn)在,只要花1620美元,就可以通過混合引擎DeepSpeed-HE,在2.1天內(nèi)訓(xùn)練一個OPT-66B模型。
而如果使用多節(jié)點、多GPU系統(tǒng),DeepSpeed-HE可以花320美元,在1.25小時內(nèi)訓(xùn)練一個OPT-13B模型,花5120美元,就能在不到一天的時間內(nèi)訓(xùn)練一個OPT-175B模型。
前Meta AI專家Elvis激動轉(zhuǎn)發(fā),稱這是一件大事,并表示好奇DeepSpeed Chat和ColossalChat相比起來如何。
下面,我們就來看看效果如何。
經(jīng)過DeepSpeed-Chat的訓(xùn)練,13億參數(shù)版「ChatGPT」在問答環(huán)節(jié)上的表現(xiàn)非常亮眼。不僅能get到問題的上下文關(guān)系,而且給出的答案也有模有樣。
在多輪對話中,這個13億參數(shù)版「ChatGPT」所展示出的性能,也完全超越了這個規(guī)模的固有印象。
一段代碼,生成你的第一個ChatGPT
當(dāng)然,在體驗之前,還需要把環(huán)境配置一下:
一杯咖啡,訓(xùn)完13億參數(shù)版ChatGPT
如果你只有大約1-2小時的咖啡或午餐休息時間,也可以嘗試使用DeepSpeed-Chat訓(xùn)練一個「小玩具」。
團隊特地準備了一個針對1.3B模型的訓(xùn)練示例,可在消費級GPU上進行測試。最棒的是,當(dāng)你從午休回來時,一切都已準備就緒。
配備48GB顯存的消費級NVIDIA A6000 GPU:
一個GPU Node,半天搞定130億參數(shù)
如果你只有半天的時間,以及一臺服務(wù)器節(jié)點,則可以通過預(yù)訓(xùn)練的OPT-13B作為actor模型,OPT-350M作為reward模型,來生成一個130億參數(shù)的類ChatGPT模型。
單DGX節(jié)點,搭載了8個NVIDIA A100-40G GPU:
超省錢云方案,訓(xùn)練660億參數(shù)模型
如果你可以使用多節(jié)點集群或云資源,并希望訓(xùn)練一個更大、更高質(zhì)量的模型。那么只需基于下面這行代碼,輸入你想要的模型大小(如66B)和GPU數(shù)量(如64)。
8個DGX節(jié)點,每個節(jié)點配備8個NVIDIA A100-80G GPU:
具體來說,針對不同規(guī)模的模型和硬件配置,DeepSpeed-RLHF系統(tǒng)所需的時間和成本如下:
DeepSpeed Chat是個啥?
DeepSpeed Chat是一種通用系統(tǒng)框架,能夠?qū)崿F(xiàn)類似ChatGPT模型的端到端RLHF訓(xùn)練,從而幫助我們生成自己的高質(zhì)量類ChatGPT模型。
DeepSpeed Chat具有以下三大核心功能:
1. 簡化ChatGPT類型模型的訓(xùn)練和強化推理體驗
開發(fā)者只需一個腳本,就能實現(xiàn)多個訓(xùn)練步驟,并且在完成后還可以利用推理API進行對話式交互測試。
2. DeepSpeed-RLHF模塊
DeepSpeed-RLHF復(fù)刻了InstructGPT論文中的訓(xùn)練模式,并提供了數(shù)據(jù)抽象和混合功能,支持開發(fā)者使用多個不同來源的數(shù)據(jù)源進行訓(xùn)練。
3. DeepSpeed-RLHF系統(tǒng)
團隊將DeepSpeed的訓(xùn)練(training engine)和推理能力(inference engine) 整合成了一個統(tǒng)一的混合引擎(DeepSpeed Hybrid Engine or DeepSpeed-HE)中,用于RLHF訓(xùn)練。由于,DeepSpeed-HE能夠無縫地在推理和訓(xùn)練模式之間切換,因此可以利用來自DeepSpeed-Inference的各種優(yōu)化。
DeepSpeed-RLHF系統(tǒng)在大規(guī)模訓(xùn)練中具有無與倫比的效率,使復(fù)雜的RLHF訓(xùn)練變得快速、經(jīng)濟并且易于大規(guī)模推廣:
①高效且經(jīng)濟:
DeepSpeed-HE比現(xiàn)有系統(tǒng)快15倍以上,使RLHF訓(xùn)練快速且經(jīng)濟實惠。
例如,DeepSpeed-HE在Azure云上只需9小時即可訓(xùn)練一個OPT-13B模型,只需18小時即可訓(xùn)練一個OPT-30B模型。這兩種訓(xùn)練分別花費不到300美元和600美元。
②卓越的擴展性:
DeepSpeed-HE能夠支持訓(xùn)練擁有數(shù)千億參數(shù)的模型,并在多節(jié)點多GPU系統(tǒng)上展現(xiàn)出卓越的擴展性。
因此,即使是一個擁有130億參數(shù)的模型,也只需1.25小時就能完成訓(xùn)練。而對于擁有1750 億參數(shù)的模型,使用DeepSpeed-HE進行訓(xùn)練也只需不到一天的時間。
③實現(xiàn)RLHF訓(xùn)練的普及化:
僅憑單個GPU,DeepSpeed-HE就能支持訓(xùn)練超過130億參數(shù)的模型。這使得那些無法使用多GPU系統(tǒng)的數(shù)據(jù)科學(xué)家和研究者不僅能夠輕松創(chuàng)建輕量級的RLHF模型,還能創(chuàng)建大型且功能強大的模型,以應(yīng)對不同的使用場景。
完整的RLHF訓(xùn)練流程
為了提供無縫的訓(xùn)練體驗,研究者遵循InstructGPT,并在DeepSpeed-Chat中包含了一個完整的端到端訓(xùn)練流程。
DeepSpeed-Chat的RLHF訓(xùn)練流程圖示,包含了一些可選擇的功能
流程包括三個主要步驟:
第 1 步:
監(jiān)督微調(diào) (SFT),使用精選的人類回答來微調(diào)預(yù)訓(xùn)練的語言模型,以應(yīng)對各種查詢。
第 2 步:
獎勵模型微調(diào),用一個包含人類對同一查詢的多個答案打分的數(shù)據(jù)集,來訓(xùn)練一個獨立的(通常比SFT小)獎勵模型(RW)。
第 3 步:
RLHF訓(xùn)練,在這一步,SFT模型通過使用近似策略優(yōu)化(PPO)算法,從RW模型的獎勵反饋進一步微調(diào)。
在步驟3中,研究者還提供了兩個附加功能,來幫助提高模型質(zhì)量:
- 指數(shù)移動平均線(EMA)的收集,可以選擇一個基于EMA的檢查點,進行最終評估。
- 混合訓(xùn)練,將預(yù)訓(xùn)練目標(即下一個詞預(yù)測)與 PPO 目標混合,以防止在公共基準(如SQuAD2.0)上的性能回歸。
EMA和混合訓(xùn)練這兩個訓(xùn)練特征,常常被其他的開源框架所忽略,因為它們并不會妨礙訓(xùn)練的進行。
然而,根據(jù)InstructGPT,EMA檢查點往往比傳統(tǒng)的最終訓(xùn)練模型提供更好的響應(yīng)質(zhì)量,而混合訓(xùn)練可以幫助模型保持訓(xùn)練前的基準解決能力。
因此,研究者為用戶提供了這些功能,讓他們可以充分獲得InstructGPT中描述的訓(xùn)練經(jīng)驗。
而除了與InstructGPT論文高度一致外,研究者還提供了功能,讓開發(fā)者使用多種數(shù)據(jù)資源,訓(xùn)練自己的RLHF模型:
數(shù)據(jù)抽象和混合能力:
DeepSpeed-Chat配備了(1)抽象數(shù)據(jù)集層,以統(tǒng)一不同數(shù)據(jù)集的格式;以及(2)數(shù)據(jù)拆分/混合功能,從而使多個數(shù)據(jù)集被適當(dāng)?shù)鼗旌希缓笤?個訓(xùn)練階段進行分割。
DeepSpeed混合引擎
指示引導(dǎo)的RLHF管道的第1步和第2步,類似于大模型的常規(guī)微調(diào),它們由基于ZeRO的優(yōu)化和DeepSpeed訓(xùn)練中靈活的并行策略組合,來實現(xiàn)規(guī)模和速度。
而管道的第3步是在性能影響方面最復(fù)雜的部分。
每個迭代都需要有效地處理兩個階段:a)推理階段,用于token/經(jīng)驗的生成,產(chǎn)生訓(xùn)練的輸入;b)訓(xùn)練階段,更新actor和reward模型的權(quán)重,以及它們之間的互動和調(diào)度。
它引入了兩個主要的困難:(1)內(nèi)存成本,因為在整個第三階段需要運行多個SFT和RW模型;(2)生成回答階段的速度較慢,如果沒有正確加速,將顯著拖慢整個第三階段。
此外,研究者在第三階段增加的兩個重要功能——指數(shù)移動平均(EMA)收集和混合訓(xùn)練,將產(chǎn)生額外的內(nèi)存和訓(xùn)練成本。
為了應(yīng)對這些挑戰(zhàn),研究者將DeepSpeed訓(xùn)練和推理的全部系統(tǒng)能力組成了一個統(tǒng)一的基礎(chǔ)設(shè)施,即混合引擎(Hybrid Engine)。
它利用原有的DeepSpeed引擎進行快速訓(xùn)練模式,同時毫不費力地應(yīng)用DeepSpeed推理引擎進行生成/評估模式,為第三階段的RLHF訓(xùn)練提供了一個更快的訓(xùn)練系統(tǒng)。
如下圖所示,DeepSpeed訓(xùn)練和推理引擎之間的過渡是無縫的:通過為actor模型啟用典型的eval和train模式,在運行推理和訓(xùn)練流程時,DeepSpeed選擇了不同的優(yōu)化,以更快地運行模型,并提高整個系統(tǒng)的吞吐量。
用于加速RLHF流程中最耗時部分的DeepSpeed混合引擎設(shè)計
在RLHF訓(xùn)練的經(jīng)驗生成階段的推理執(zhí)行過程中,DeepSpeed混合引擎使用輕量級的內(nèi)存管理系統(tǒng),來處理KV緩存和中間結(jié)果,同時使用高度優(yōu)化的推理CUDA核和張量并行計算,與現(xiàn)有方案相比,實現(xiàn)了吞吐量(每秒token數(shù))的大幅提升。
在訓(xùn)練期間,混合引擎啟用了內(nèi)存優(yōu)化技術(shù),如DeepSpeed的ZeRO系列技術(shù)和低階自適應(yīng)(LoRA)。
而研究者設(shè)計和實現(xiàn)這些系統(tǒng)優(yōu)化的方式是,讓它們彼此兼容,并可以組合在一起,在統(tǒng)一的混合引擎下提供最高的訓(xùn)練效率。
混合引擎可以在訓(xùn)練和推理中無縫地改變模型分區(qū),以支持基于張量并行的推理,和基于ZeRO的訓(xùn)練分片機制。
它還可以重新配置內(nèi)存系統(tǒng),以便在每一種模式中最大限度地提高內(nèi)存可用性。
這就避免了內(nèi)存分配瓶頸,能夠支持大的batch size,讓性能大大提升。
總之,混合引擎推動了現(xiàn)代RLHF訓(xùn)練的邊界,為RLHF工作負載提供了無與倫比的規(guī)模和系統(tǒng)效率。
效果評估
與Colossal-AI或HuggingFace-DDP等現(xiàn)有系統(tǒng)相比,DeepSpeed-Chat具有超過一個數(shù)量級的吞吐量,能夠在相同的延遲預(yù)算下訓(xùn)練更大的演員模型或以更低的成本訓(xùn)練相似大小的模型。
例如,在單個GPU上,DeepSpeed使RLHF訓(xùn)練的吞吐量提高了10倍以上。雖然CAI-Coati和HF-DDP都可以運行1.3B的模型,但DeepSpeed可以在相同的硬件上運行6.5B模型,直接高出5倍。
在單個節(jié)點的多個GPU上,DeepSpeed-Chat在系統(tǒng)吞吐量方面比CAI-Coati提速6-19倍,HF-DDP提速1.4-10.5倍。
團隊表示,DeepSpeed-Chat能夠獲得如此優(yōu)異的結(jié)果,關(guān)鍵原因之一便是混合引擎在生成階段提供的加速。
參考資料:
https://github.com/microsoft/DeepSpeed