根據(jù)微軟的介紹,DeepSpeed 庫中有一個(gè)名為 ZeRO(零冗余優(yōu)化器,Zero Redundancy Optimizer)的組件,這是一種新的并行優(yōu)化器,它可以大大減少模型和數(shù)據(jù)并行所需的資源,同時(shí)可以大量增加可訓(xùn)練的參數(shù)數(shù)量。研究人員利用這些突破創(chuàng)建了圖靈自然語言生成模型(Turing-NLG),這是最大的公開語言模型,參數(shù)為 170 億。
ZeRO 作為 DeepSpeed 的一部分,是一種用于大規(guī)模分布式深度學(xué)習(xí)的新內(nèi)存優(yōu)化技術(shù),它可以在當(dāng)前的 GPU 集群上訓(xùn)練具有 1000 億個(gè)參數(shù)的深度學(xué)習(xí)模型,其吞吐量是當(dāng)前最佳系統(tǒng)的 3 到 5 倍。它還為訓(xùn)練具有數(shù)萬億個(gè)參數(shù)的模型提供了一條清晰的思路。
ZeRO 具有三個(gè)主要的優(yōu)化階段,分別對(duì)應(yīng)于優(yōu)化器狀態(tài)、梯度和參數(shù)分區(qū)。
ZeRO 克服了數(shù)據(jù)并行和模型并行的局限性,同時(shí)實(shí)現(xiàn)兩者的優(yōu)點(diǎn),它通過跨數(shù)據(jù)并行進(jìn)程將模型狀態(tài)劃分為上圖所示的參數(shù)、梯度和優(yōu)化器狀態(tài)分區(qū),而不是復(fù)制它們,從而消除了數(shù)據(jù)并行進(jìn)程之間的內(nèi)存冗余。在訓(xùn)練期間使用動(dòng)態(tài)通信規(guī)劃(dynamic communication schedule),在分布式設(shè)備之間共享必要的狀態(tài),以保持?jǐn)?shù)據(jù)并行的計(jì)算粒度和通信量。
目前實(shí)施了 ZeRO 的第一階段,即優(yōu)化器狀態(tài)分區(qū)(簡(jiǎn)稱 ZeRO-OS),具有支持 1000 億參數(shù)模型的強(qiáng)大能力,此階段與 DeepSpeed 一起發(fā)布。
DeepSpeed 與 PyTorch 兼容,DeepSpeed API 是在 PyTorch 上進(jìn)行的輕量級(jí)封裝,這意味著開發(fā)者可以使用 PyTorch 中的一切,而無需學(xué)習(xí)新平臺(tái)。此外,DeepSpeed 管理著所有樣板化的 SOTA 訓(xùn)練技術(shù),例如分布式訓(xùn)練、混合精度、梯度累積和檢查點(diǎn),開發(fā)者可以專注于模型開發(fā)。同時(shí),開發(fā)者僅需對(duì) PyTorch 模型進(jìn)行幾行代碼的更改,就可以利用 DeepSpeed 獨(dú)特的效率和效益優(yōu)勢(shì)來提高速度和規(guī)模。
DeepSpeed 在以下四個(gè)方面都表現(xiàn)出色: