多GPU通信效率提升4倍，RLHF生成提升2.25倍！DeepSpeed ZeRO++重磅升級
過去半年，由ChatGPT引領(lǐng)的生成式大型語言模型技術(shù)，以其強大的「通用性」徹底顛覆了AI世界，普通人也可以很容易地使用AI工具來進行摘要、靈感創(chuàng)作、輔助編程、多語言翻譯等任務(wù)。
不過，訓(xùn)練這種超大規(guī)模的模型往往需要數(shù)百個、甚至數(shù)千個GPU來存儲和計算數(shù)據(jù)，比如訓(xùn)練5300億參數(shù)的Megatron-Turning NLG就使用了超過4000塊Nvidia A100 GPU
想要高效地利用硬件資源需要設(shè)計復(fù)雜的優(yōu)化系統(tǒng)，將模型劃分割成適合于單個設(shè)備內(nèi)存的片段，然后跨設(shè)備進行高效的并行計算；同時，為了能夠讓深度學(xué)習(xí)社區(qū)更方便地對大型模型進行訓(xùn)練，這些優(yōu)化操作必須易于使用。
2020年2月，微軟開源了深度學(xué)習(xí)訓(xùn)練優(yōu)化庫DeepSpeed，并提供了內(nèi)存優(yōu)化技術(shù)ZeRO（零冗余優(yōu)化器，Zero Redundancy Optimizer），提高了可訓(xùn)練模型的規(guī)模、訓(xùn)練速度、降低了成本、提升了可用性，極大地推動了大模型的訓(xùn)練，已應(yīng)用于TNLG-17 B、Bloom-176 B、MPT-7 B、Jurrasic-1等模型的訓(xùn)練上。
圖片
倉庫鏈接：https://github.com/microsoft/DeepSpeed
但在某些場景下，ZeRO系列技術(shù)會在GPU之間產(chǎn)生更高的數(shù)據(jù)傳輸開銷，降低模型的訓(xùn)練效率，尤其是：
1. 總的GPU數(shù)量較多，導(dǎo)致分配到每個GPU上batch size較少，硬件間通信更頻繁；
2. 在低端（low-end）集群上進行訓(xùn)練時，跨節(jié)點的網(wǎng)絡(luò)帶寬有限，導(dǎo)致GPU間的通信延遲較高。
為了解決這些問題，微軟DeepSpeed團隊發(fā)布了新一代基于ZeRO的通信優(yōu)化策略系統(tǒng)ZeRO++，可以為大型模型訓(xùn)練提供更高的效率，不受batch size限制或跨設(shè)備帶寬限制。
ZeRO++利用量化結(jié)合數(shù)據(jù)和通信重映射，將總通信量減少4倍，并且不會影響模型質(zhì)量。
在數(shù)千個GPU上進行訓(xùn)練，或是在幾百、幾十個GPU上微調(diào)時，在每個GPU分配到的batch size都比較小的情況下，ZeRO++的吞吐量比ZeRO高出2.2倍，可以直接降低訓(xùn)練時間和成本。
在低帶寬集群下，ZeRO++可以使其實現(xiàn)與4倍帶寬集群相似的吞吐量，可以在更廣泛的集群中實現(xiàn)高效的大型模型訓(xùn)練。
ZeRO++還可以加速使用RLHF技術(shù)的ChatGPT類模型的訓(xùn)練。
雖然ZeRO++主要是為訓(xùn)練而設(shè)計的，但通訊開銷在訓(xùn)練和推理過程中都存在，其優(yōu)化也適用于ZeRO-Inference，可以提高訓(xùn)練對話模型中使用的人工反饋強化學(xué)習(xí)（RLHF）等工作負(fù)載的效率。
與原始ZeRO相比，將ZeRO++與DeepSpeed-Chat集成后，可以將RLHF訓(xùn)練的生成速度提高2倍，強化學(xué)習(xí)訓(xùn)練階段提高1.3倍。
ZeRO++
ZeRO是數(shù)據(jù)并行的一種內(nèi)存高效變體（memory efficient variation），其中模型states在所有GPU上進行分區(qū)而非復(fù)制，并在訓(xùn)練期間使用基于收集/廣播（gather/broadcast）的通信集體進行重建，也使得ZeRO能夠高效地利用聚合GPU內(nèi)存和跨所有設(shè)備的計算，同時提供簡單易用的數(shù)據(jù)并行訓(xùn)練。
假設(shè)模型大小為M，則在前向傳遞過程中，ZeRO執(zhí)行所有收集/廣播操作，以便在有需要之前收集完每個模型層的參數(shù)（通信量為M）；
在反向傳遞中，ZeRO對每一層的參數(shù)采用類似的通信模式以計算局部梯度（通信量為M）；
此外，ZeRO在使用reduce或reduce-scatter通信集合計算每個局部梯度之后立即對每個局部梯度進行平均和分區(qū)（通信量為M）。
所以ZeRO總共需要3M的通信量，均勻分布在兩次全收集/廣播和一次reduce/reduce-scatter操作中。
為了減少這些通信開銷，ZeRO++設(shè)計了三組通信優(yōu)化，分別針對上述三個通信集合。
ZeRO (qwZ)的量化權(quán)重通信
為了減少全收集（all-gather）過程中的參數(shù)通信量，研究人員選擇采用權(quán)值量化，在通信之前將每個模型參數(shù)從FP16（兩個字節(jié)）壓縮到INT8（一個字節(jié)）數(shù)據(jù)類型，并在通信之后反量化權(quán)值。
不過只是簡單地對權(quán)重進行量化可能會降低模型訓(xùn)練精度，為了在提速的同時保持訓(xùn)練精度，研究人員采用了基于塊（block-based）的量化，即對模型參數(shù)的每個子集進行獨立的量化。
目前尚不存在用于高性能、基于塊量化的實現(xiàn)，所以研究人員從頭開始實現(xiàn)高度優(yōu)化的量化CUDA內(nèi)核；與基本量化相比，準(zhǔn)確性提高了3倍，速度提高了5倍。
圖片
ZeRO（hpZ）的分層權(quán)重劃分
為了減少反向傳遞期間權(quán)重上的所有g(shù)ather操作的通信開銷，研究人員選擇犧牲GPU內(nèi)存換取通信效率。
具體來說，與在ZeRO中將整個模型權(quán)重分散在所有機器上不同，ZeRO++在每臺機器內(nèi)維護一個完整的模型副本，以更高的內(nèi)存開銷為代價，使得可以用機器內(nèi)全收集/廣播操作來替換昂貴的跨機器全收集/廣播權(quán)重。
因為機器內(nèi)的通信帶寬更高，所以實際性能會得到提升。
圖片
用于ZeRO（qgZ）的量化梯度通信
使用reduce-scatter來降低梯度的通信成本更具挑戰(zhàn)性，因為直接應(yīng)用量化來減少通信量是不可行的，即使將基于塊的權(quán)重量化為低精度，梯度降低也會累積并放大量化誤差。
為了解決這個問題，可以在通信之前只量化梯度，但在任何reduce操作之前將梯度反量化到全精度。
為了高效地實現(xiàn)，研究人員發(fā)明了一個基于all-to-all，新的量化梯度通信范式qgZ，功能上等同于compressed reduce-scatter collective操作。
qgZ旨在解決兩個問題：
1. 克服在INT 4/INT 8低精度中實現(xiàn)reduce-scatter導(dǎo)致的準(zhǔn)確度下降；
2. 避免由長序列的量化和反量化（dequantization）步驟中導(dǎo)致的準(zhǔn)確率下降和等待時間開銷，這兩個步驟是基于ring或tree的reduce scatter傳統(tǒng)方法所必需的。
qgZ是一種全新的基于層次化all-to-all的方法，包括三個主要步驟：
1. 梯度切片重新排序
2. 節(jié)點內(nèi)通信和reduction
3. 節(jié)點間通信和reduction
首先，在任何通信操作執(zhí)行之前，對梯度進行切片并進行張量切片重新排序以保證最終的梯度放置，在通信結(jié)束時在每個GPU上是正確的。
然后量化重排序后的梯度切片，在每個節(jié)點內(nèi)進行all-to-all通信，反量化接收到的梯度切片，并進行局部reductions
再次量化局部reduction的梯度，進行節(jié)點間all-to-all通信，反量化接收到的梯度，并計算最終的高精度梯度降低，得到最終結(jié)果。
使用層次化方法的原因是為了減少跨節(jié)點通信量，更確切地說，如果每個節(jié)點有N個GPU、模型大小為M和量化比為Z，單跳all-to-all會生成M*N/Z個跨節(jié)點流量。
相比之下，分層方法將每個GPU的跨節(jié)點流量從M/Z減少到M/（Z*N），總通信量可以降到M*N/（Z*N）= M/Z
通過重疊節(jié)點內(nèi)和節(jié)點間通信以及融合CUDA內(nèi)核（張量切片重新排序+節(jié)點內(nèi)量化）和（節(jié)點內(nèi)去量化+節(jié)點內(nèi)歸約+節(jié)點間量化）來進一步優(yōu)化qgZ的端到端延遲。
總結(jié)
通過整合上述三個組件，ZeRO++將跨節(jié)點通信量從3M減少到0.75M。
使用qwZ將模型權(quán)重從M減少到0.5M；使用hpZ消除了反向傳播過程中的跨節(jié)點all-gather，將通信從M減少到0；使用qgZ將反向傳遞期間的跨節(jié)點reduce-scatter通信從M降低到0.25M
ZeRO++加速LLM訓(xùn)練
在實驗部分，研究人員展示了在384個Nvidia V100 GPU上使用真實LLM訓(xùn)練場景對ZeRO++的評估結(jié)果。
在高帶寬集群中測試ZeRO++吞吐量的改進，對于不同的模型尺寸和micro batch size，使用4x Infiniband（IB）的400 Gbps跨節(jié)點互連，每個都以100 Gbps運行。
圖片
每個GPU使用1000 token，ZeRO++比ZeRO-3的吞吐量提高了28%到36%。
對于2000 micro batch size，ZeRO++比ZeRO-3實現(xiàn)了24%至29%的吞吐量增益。
在低帶寬集群中，在100 Gbps網(wǎng)絡(luò)等低網(wǎng)絡(luò)環(huán)境中，ZeRO++的性能明顯優(yōu)于ZeRO-3
圖片
與ZeRO-3相比，ZeRO++在端到端吞吐量方面實現(xiàn)了高達(dá)2.2倍的加速，平均而言，ZeRO++比ZeRO-3基線實現(xiàn)了約2倍的加速。
ZeRO++結(jié)合DeepSpeed-Chat進行RLHF訓(xùn)練
RLHF訓(xùn)練背景
ChatGPT類模型的基礎(chǔ)是大型語言模型，然后使用RLHF進行微調(diào)。
RLHF由生成（推理）階段和訓(xùn)練階段組成：在生成階段，actor模型將部分對話作為輸入，并使用一系列前向傳遞來生成回復(fù)。
在訓(xùn)練階段，critic模型按質(zhì)量對生成的回復(fù)進行排序，為actor模型提供強化信號；使用排名對actor模型進行微調(diào)，使其能夠在后續(xù)迭代中生成更準(zhǔn)確和適當(dāng)?shù)幕貜?fù)。
RLHF使用了四個模型（actor, reference, critic, reward），所以需要大量的內(nèi)存，采用低秩自適應(yīng)（LoRA）可以解決RLHF的內(nèi)存壓力。
LoRA凍結(jié)了預(yù)訓(xùn)練的模型權(quán)重，并將可訓(xùn)練的秩分解矩陣注入到Transformer架構(gòu)的每一層中，從而顯著減少了可訓(xùn)練參數(shù)的數(shù)量。
LoRA可以通過減少內(nèi)存使用量、允許更大的batch size來加快RLHF的訓(xùn)練速度，從而大大提高吞吐量。
DeepSpeed-Chat結(jié)合ZeRO++
用LoRA進行RLHF訓(xùn)練是因為大多數(shù)模型權(quán)重都是凍結(jié)的，也意味著ZeRO++可以將這些凍結(jié)的權(quán)重保持在INT 4/8中量化，無需存儲為FP 16并在每次通信操作之前重復(fù)量化。
以這種方式使用ZeRO++進行RLHF訓(xùn)練可以減少內(nèi)存使用和通信量，通過減少通信以及由于減少內(nèi)存使用而使用更大的batch size大小來提高訓(xùn)練吞吐量。
在生成階段，ZeRO++使用hpZ保持每個節(jié)點內(nèi)的所有權(quán)重通信，以利用更高的節(jié)點內(nèi)通信帶寬，減少通信量，進一步提高生成吞吐量。
圖片
研究人員對比了不同尺寸的actor模型的RLHF生成吞吐量，并比較了32個V100 GPU上的30 B和66 B actor模型的ZeRO和ZeRO++，結(jié)果顯示，ZeRO++能夠?qū)崿F(xiàn)比ZeRO高2.25倍的RLHF生成吞吐量。
在16個V100 GPU上的訓(xùn)練階段加速對比中，ZeRO++實現(xiàn)了比ZeRO高1.26倍的吞吐量，主要是由于ZeRO++實現(xiàn)了更低的通信和更大的batch size
參考資料：
