https://m.toutiao.com/is/AT7tVup/
編輯:Aeneas 好困
人手一個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
當然,在體驗之前,還需要把環境配置一下:
git clone https://github.com/microsoft/DeepSpeed.gitcd DeepSpeedpip install .git clone https://github.com/microsoft/DeepSpeedExamples.gitcd DeepSpeedExamples/applications/DeepSpeed-Chat/pip install -r requirements.txt
如果你只有大約1-2小時的咖啡或午餐休息時間,也可以嘗試使用DeepSpeed-Chat訓練一個「小玩具」。
團隊特地準備了一個針對1.3B模型的訓練示例,可在消費級GPU上進行測試。最棒的是,當你從午休回來時,一切都已準備就緒。
python train.py --actor-model facebook/opt-1.3b --reward-model facebook/opt-350m --num-gpus 1
配備48GB顯存的消費級NVIDIA A6000 GPU:
如果你只有半天的時間,以及一臺服務器節點,則可以通過預訓練的OPT-13B作為actor模型,OPT-350M作為reward模型,來生成一個130億參數的類ChatGPT模型:
python train.py --actor-model facebook/opt-13b --reward-model facebook/opt-350m --num-gpus 8
單DGX節點,搭載了8個NVIDIA A100-40G GPU:
如果你可以使用多節點集群或云資源,并希望訓練一個更大、更高質量的模型。那么只需基于下面這行代碼,輸入你想要的模型大小(如66B)和GPU數量(如64):
python train.py --actor-model facebook/opt-66b --reward-model facebook/opt-350m --num-gpus 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進行訓練也只需不到一天的時間。
僅憑單個GPU,DeepSpeed-HE就能支持訓練超過130億參數的模型。這使得那些無法使用多GPU系統的數據科學家和研究者不僅能夠輕松創建輕量級的RLHF模型,還能創建大型且功能強大的模型,以應對不同的使用場景。
完整的RLHF訓練流程
為了提供無縫的訓練體驗,研究者遵循InstructGPT,并在DeepSpeed-Chat中包含了一個完整的端到端訓練流程。
DeepSpeed-Chat的RLHF訓練流程圖示,包含了一些可選擇的功能
流程包括三個主要步驟:
監督微調 (SFT),使用精選的人類回答來微調預訓練的語言模型,以應對各種查詢。
獎勵模型微調,用一個包含人類對同一查詢的多個答案打分的數據集,來訓練一個獨立的(通常比SFT小)獎勵模型(RW)。
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