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