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

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

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

項(xiàng)目地址：https://github.com/microsoft/DeepSpeed

一鍵解鎖千億級(jí)ChatGPT，輕松省錢15倍

眾所周知，由于OpenAI太不Open，開(kāi)源社區(qū)為了讓更多人能用上類ChatGPT模型，相繼推出了LLaMa、Alpaca、Vicuna、Databricks-Dolly等模型。

但由于缺乏一個(gè)支持端到端的RLHF規(guī)模化系統(tǒng)，目前類ChatGPT模型的訓(xùn)練仍然十分困難。而DeepSpeed Chat的出現(xiàn)，正好補(bǔ)全了這個(gè)「bug」。

更亮的是，DeepSpeed Chat把成本大大地打了下來(lái)。

此前，昂貴的多GPU設(shè)置超出了許多研究者的能力范圍，并且，即使能訪問(wèn)多GPU集群，現(xiàn)有的方法也無(wú)力負(fù)擔(dān)數(shù)千億參數(shù)ChatGPT模型的訓(xùn)練。

現(xiàn)在，只要花1620美元，就可以通過(guò)混合引擎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相比起來(lái)如何。

下面，我們就來(lái)看看效果如何。

經(jīng)過(guò)DeepSpeed-Chat的訓(xùn)練，13億參數(shù)版「ChatGPT」在問(wèn)答環(huán)節(jié)上的表現(xiàn)非常亮眼。不僅能get到問(wèn)題的上下文關(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/applications/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)行測(cè)試。最棒的是，當(dāng)你從午休回來(lái)時(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)，則可以通過(guò)預(yù)訓(xùn)練的OPT-13B作為actor模型，OPT-350M作為reward模型，來(lái)生成一個(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：

具體來(lái)說(shuō)，針對(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. 簡(jiǎn)化ChatGPT類型模型的訓(xùn)練和強(qiáng)化推理體驗(yàn)?

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

2. DeepSpeed-RLHF模塊?

DeepSpeed-RLHF復(fù)刻了InstructGPT論文中的訓(xùn)練模式，并提供了數(shù)據(jù)抽象和混合功能，支持開(kāi)發(fā)者使用多個(gè)不同來(lái)源的數(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能夠無(wú)縫地在推理和訓(xùn)練模式之間切換，因此可以利用來(lái)自DeepSpeed-Inference的各種優(yōu)化。

DeepSpeed-RLHF系統(tǒng)在大規(guī)模訓(xùn)練中具有無(wú)與倫比的效率，使復(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)練超過(guò)130億參數(shù)的模型。這使得那些無(wú)法使用多GPU系統(tǒng)的數(shù)據(jù)科學(xué)家和研究者不僅能夠輕松創(chuàng)建輕量級(jí)的RLHF模型，還能創(chuàng)建大型且功能強(qiáng)大的模型，以應(yīng)對(duì)不同的使用場(chǎng)景。

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

為了提供無(wú)縫的訓(xùn)練體驗(yàn)，研究者遵循InstructGPT，并在DeepSpeed-Chat中包含了一個(gè)完整的端到端訓(xùn)練流程。

DeepSpeed-Chat的RLHF訓(xùn)練流程圖示，包含了一些可選擇的功能

流程包括三個(gè)主要步驟：

• 第 1 步：

監(jiān)督微調(diào) (SFT)，使用精選的人類回答來(lái)微調(diào)預(yù)訓(xùn)練的語(yǔ)言模型，以應(yīng)對(duì)各種查詢。

• 第 2 步：

獎(jiǎng)勵(lì)模型微調(diào)，用一個(gè)包含人類對(duì)同一查詢的多個(gè)答案打分的數(shù)據(jù)集，來(lái)訓(xùn)練一個(gè)獨(dú)立的（通常比SFT小）獎(jiǎng)勵(lì)模型（RW）。

• 第 3 步：

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

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

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

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

EMA和混合訓(xùn)練這兩個(gè)訓(xùn)練特征，常常被其他的開(kāi)源框架所忽略，因?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論文高度一致外，研究者還提供了功能，讓開(kāi)發(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ù)鼗旌希缓笤?個(gè)訓(xùn)練階段進(jìn)行分割。

DeepSpeed混合引擎

指示引導(dǎo)的RLHF管道的第1步和第2步，類似于大模型的常規(guī)微調(diào)，它們由基于ZeRO的優(yōu)化和DeepSpeed訓(xùn)練中靈活的并行策略組合，來(lái)實(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）生成回答階段的速度較慢，如果沒(méi)有正確加速，將顯著拖慢整個(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)練和推理引擎之間的過(guò)渡是無(wú)縫的：通過(guò)為actor模型啟用典型的eval和train模式，在運(yùn)行推理和訓(xùn)練流程時(shí)，DeepSpeed選擇了不同的優(yōu)化，以更快地運(yùn)行模型，并提高整個(gè)系統(tǒng)的吞吐量。  

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

在RLHF訓(xùn)練的經(jīng)驗(yàn)生成階段的推理執(zhí)行過(guò)程中，DeepSpeed混合引擎使用輕量級(jí)的內(nèi)存管理系統(tǒng)，來(lái)處理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)練和推理中無(wú)縫地改變模型分區(qū)，以支持基于張量并行的推理，和基于ZeRO的訓(xùn)練分片機(jī)制。

它還可以重新配置內(nèi)存系統(tǒng)，以便在每一種模式中最大限度地提高內(nèi)存可用性。

這就避免了內(nèi)存分配瓶頸，能夠支持大的batch size，讓性能大大提升。

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

效果評(píng)估

與Colossal-AI或HuggingFace-DDP等現(xiàn)有系統(tǒng)相比，DeepSpeed-Chat具有超過(guò)一個(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)鍵原因之一便是混合引擎在生成階段提供的加速。