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