ZeRO 具有三個主要的優(yōu)化階段,分別對應于優(yōu)化器狀態(tài)、梯度和參數分區(qū)。
ZeRO 克服了數據并行和模型并行的局限性,同時實現兩者的優(yōu)點,它通過跨數據并行進程將模型狀態(tài)劃分為上圖所示的參數、梯度和優(yōu)化器狀態(tài)分區(qū),而不是復制它們,從而消除了數據并行進程之間的內存冗余。在訓練期間使用動態(tài)通信規(guī)劃(dynamic communication schedule),在分布式設備之間共享必要的狀態(tài),以保持數據并行的計算粒度和通信量。
目前實施了 ZeRO 的第一階段,即優(yōu)化器狀態(tài)分區(qū)(簡稱 ZeRO-OS),具有支持 1000 億參數模型的強大能力,此階段與 DeepSpeed 一起發(fā)布。
DeepSpeed 與 PyTorch 兼容,DeepSpeed API 是在 PyTorch 上進行的輕量級封裝,這意味著開發(fā)者可以使用 PyTorch 中的一切,而無需學習新平臺。此外,DeepSpeed 管理著所有樣板化的 SOTA 訓練技術,例如分布式訓練、混合精度、梯度累積和檢查點,開發(fā)者可以專注于模型開發(fā)。同時,開發(fā)者僅需對 PyTorch 模型進行幾行代碼的更改,就可以利用 DeepSpeed 獨特的效率和效益優(yōu)勢來提高速度和規(guī)模。
DeepSpeed 在以下四個方面都表現出色:
規(guī)模:目前最先進的大型模型,例如 OpenAI GPT-2、NVIDIA Megatron-LM 和 Google T5,分別具有 15 億、83 億和 110 億個參數,而 DeepSpeed 的 ZeRO 第一階段提供系統(tǒng)支持,以運行多達 1000 億個參數的模型,這是比當前最先進的模型大 10 倍。未來計劃增加對 ZeRO 第二和第三階段的支持,從而提供高達 2000 億個乃至數萬億個參數的模型的能力。