人手一個ChatGPT！微軟DeepSpeed Chat震撼發(fā)布，一鍵RLHF訓(xùn)練千億級大模型

作者：新智元 2023-04-12 16:16:53

人工智能新聞

微軟開源的DeepSpeed Chat，讓開發(fā)者實(shí)現(xiàn)了人手一個ChatGPT的夢想！

人手一個ChatGPT的夢想，就要實(shí)現(xiàn)了？

剛剛，微軟開源了一個可以在模型訓(xùn)練中加入完整RLHF流程的系統(tǒng)框架——DeepSpeed Chat。

也就是說，各種規(guī)模的高質(zhì)量類ChatGPT模型，現(xiàn)在都唾手可得了！

項(xiàng)目地址：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)，正好補(bǔ)全了這個「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)練一個OPT-66B模型。

而如果使用多節(jié)點(diǎn)、多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)然，在體驗(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/applications/DeepSpeed-Chat/
pip install -r requirements.txt

一杯咖啡，訓(xùn)完13億參數(shù)版ChatGPT

如果你只有大約1-2小時的咖啡或午餐休息時間，也可以嘗試使用DeepSpeed-Chat訓(xùn)練一個「小玩具」。

團(tuán)隊(duì)特地準(zhǔn)備了一個針對1.3B模型的訓(xùn)練示例，可在消費(fèi)級GPU上進(jìn)行測試。最棒的是，當(dāng)你從午休回來時，一切都已準(zhǔn)備就緒。

python train.py --actor-model facebook/opt-1.3b --reward-model facebook/opt-350m --num-gpus 1

配備48GB顯存的消費(fèi)級NVIDIA A6000 GPU：

一個GPU Node，半天搞定130億參數(shù)

如果你只有半天的時間，以及一臺服務(wù)器節(jié)點(diǎn)，則可以通過預(yù)訓(xùn)練的OPT-13B作為actor模型，OPT-350M作為reward模型，來生成一個130億參數(shù)的類ChatGPT模型：

python train.py --actor-model facebook/opt-13b --reward-model facebook/opt-350m --num-gpus 8

單DGX節(jié)點(diǎn)，搭載了8個NVIDIA A100-40G GPU：

超省錢云方案，訓(xùn)練660億參數(shù)模型

如果你可以使用多節(jié)點(diǎn)集群或云資源，并希望訓(xùn)練一個更大、更高質(zhì)量的模型。那么只需基于下面這行代碼，輸入你想要的模型大小（如66B）和GPU數(shù)量（如64）：

python train.py --actor-model facebook/opt-66b --reward-model facebook/opt-350m --num-gpus 64

8個DGX節(jié)點(diǎn)，每個節(jié)點(diǎn)配備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)練和強(qiáng)化推理體驗(yàn)?

開發(fā)者只需一個腳本，就能實(shí)現(xiàn)多個訓(xùn)練步驟，并且在完成后還可以利用推理API進(jìn)行對話式交互測試。

2. DeepSpeed-RLHF模塊?

DeepSpeed-RLHF復(fù)刻了InstructGPT論文中的訓(xùn)練模式，并提供了數(shù)據(jù)抽象和混合功能，支持開發(fā)者使用多個不同來源的數(shù)據(jù)源進(jìn)行訓(xùn)練。

3. DeepSpeed-RLHF系統(tǒng)?

團(tuán)隊(duì)將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)濟(jì)并且易于大規(guī)模推廣：

高效且經(jīng)濟(jì)：

DeepSpeed-HE比現(xiàn)有系統(tǒng)快15倍以上，使RLHF訓(xùn)練快速且經(jīng)濟(jì)實(shí)惠。

例如，DeepSpeed-HE在Azure云上只需9小時即可訓(xùn)練一個OPT-13B模型，只需18小時即可訓(xùn)練一個OPT-30B模型。這兩種訓(xùn)練分別花費(fèi)不到300美元和600美元。

卓越的擴(kuò)展性：

DeepSpeed-HE能夠支持訓(xùn)練擁有數(shù)千億參數(shù)的模型，并在多節(jié)點(diǎn)多GPU系統(tǒng)上展現(xiàn)出卓越的擴(kuò)展性。

因此，即使是一個擁有130億參數(shù)的模型，也只需1.25小時就能完成訓(xùn)練。而對于擁有1750 億參數(shù)的模型，使用DeepSpeed-HE進(jìn)行訓(xùn)練也只需不到一天的時間。

實(shí)現(xiàn)RLHF訓(xùn)練的普及化：

僅憑單個GPU，DeepSpeed-HE就能支持訓(xùn)練超過130億參數(shù)的模型。這使得那些無法使用多GPU系統(tǒng)的數(shù)據(jù)科學(xué)家和研究者不僅能夠輕松創(chuàng)建輕量級的RLHF模型，還能創(chuàng)建大型且功能強(qiáng)大的模型，以應(yīng)對不同的使用場景。

完整的RLHF訓(xùn)練流程

為了提供無縫的訓(xùn)練體驗(yà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)練一個獨(dú)立的（通常比SFT小）獎勵模型（RW）。

第 3 步：

RLHF訓(xùn)練，在這一步，SFT模型通過使用近似策略優(yōu)化（PPO）算法，從RW模型的獎勵反饋進(jìn)一步微調(diào)。

在步驟3中，研究者還提供了兩個附加功能，來幫助提高模型質(zhì)量：

- 指數(shù)移動平均線（EMA）的收集，可以選擇一個基于EMA的檢查點(diǎn)，進(jìn)行最終評估。

- 混合訓(xùn)練，將預(yù)訓(xùn)練目標(biāo)（即下一個詞預(yù)測）與 PPO 目標(biāo)混合，以防止在公共基準(zhǔn)（如SQuAD2.0）上的性能回歸。

EMA和混合訓(xùn)練這兩個訓(xùn)練特征，常常被其他的開源框架所忽略，因?yàn)樗鼈儾⒉粫恋K訓(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ù)拆分/混合功能，從而使多個數(shù)據(jù)集被適當(dāng)?shù)鼗旌希缓笤?個訓(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ù)雜的部分。

每個迭代都需要有效地處理兩個階段：a）推理階段，用于token/經(jīng)驗(yàn)的生成，產(chǎn)生訓(xùn)練的輸入；b）訓(xùn)練階段，更新actor和reward模型的權(quán)重，以及它們之間的互動和調(diào)度。

它引入了兩個主要的困難：（1）內(nèi)存成本，因?yàn)樵谡麄€第三階段需要運(yùn)行多個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引擎進(jìn)行快速訓(xùn)練模式，同時毫不費(fèi)力地應(yīng)用DeepSpeed推理引擎進(jìn)行生成/評估模式，為第三階段的RLHF訓(xùn)練提供了一個更快的訓(xùn)練系統(tǒng)。

如下圖所示，DeepSpeed訓(xùn)練和推理引擎之間的過渡是無縫的：通過為actor模型啟用典型的eval和train模式，在運(yùn)行推理和訓(xùn)練流程時，DeepSpeed選擇了不同的優(yōu)化，以更快地運(yùn)行模型，并提高整個系統(tǒng)的吞吐量。

用于加速RLHF流程中最耗時部分的DeepSpeed混合引擎設(shè)計(jì)

在RLHF訓(xùn)練的經(jīng)驗(yàn)生成階段的推理執(zhí)行過程中，DeepSpeed混合引擎使用輕量級的內(nèi)存管理系統(tǒng)，來處理KV緩存和中間結(jié)果，同時使用高度優(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，讓性能大大提升。

總之，混合引擎推動了現(xiàn)代RLHF訓(xùn)練的邊界，為RLHF工作負(fù)載提供了無與倫比的規(guī)模和系統(tǒng)效率。