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