微軟開源深度學習優化庫DeepSpeed 可訓練1000億參數的模型
人工智能的最新趨勢是,更大的自然語言模型可以提供更好的準確性,但是由於成本、時間和代碼集成的障礙,較大的模型難以訓練。微軟日前開源了一個深度學習優化庫DeepSpeed,通過提高規模、速度、可用性並降低成本,可以在當前一代的GPU集群上訓練具有超過1000億個參數的深度學習模型,極大促進大型模型的訓練。
同時,與最新技術相比,其係統性能可以提高5 倍以上。
根據微軟的介紹,DeepSpeed庫中有一個名為ZeRO(零冗餘優化器,Zero Redundancy Optimizer)的組件,這是一種新的並行優化器,它可以大大減少模型和數據並行所需的資源,同時可以大量增加可訓練的參數數量。研究人員利用這些突破創建了圖靈自然語言生成模型(Turing-NLG),這是最大的公開語言模型,參數為170億。
ZeRO 作為DeepSpeed 的一部分,是一種用於大規模分佈式深度學習的新內存優化技術,它可以在當前的GPU 集群上訓練具有1000 億個參數的深度學習模型,其吞吐量是當前最佳系統的3 到5 倍。它還為訓練具有數万億個參數的模型提供了一條清晰的思路。
ZeRO 具有三個主要的優化階段,分別對應於優化器狀態、梯度和參數分區。
ZeRO 克服了數據並行和模型並行的局限性,同時實現兩者的優點,它通過跨數據並行進程將模型狀態劃分為上圖所示的參數、梯度和優化器狀態分區,而不是複制它們,從而消除了數據並行進程之間的內存冗餘。在訓練期間使用動態通信規劃(dynamic communication schedule),在分佈式設備之間共享必要的狀態,以保持數據並行的計算粒度和通信量。
目前實施了ZeRO 的第一階段,即優化器狀態分區(簡稱ZeRO-OS),具有支持1000 億參數模型的強大能力,此階段與DeepSpeed 一起發布。
DeepSpeed 與 PyTorch 兼容,DeepSpeed API 是在PyTorch 上進行的輕量級封裝,這意味著開發者可以使用 PyTorch 中的一切,而無需學習新平台。此外,DeepSpeed 管理著所有樣板化的SOTA 訓練技術,例如分佈式訓練、混合精度、梯度累積和檢查點,開發者可以專注於模型開發。同時,開發者僅需對PyTorch 模型進行幾行代碼的更改,就可以利用DeepSpeed 獨特的效率和效益優勢來提高速度和規模。
DeepSpeed 在以下四個方面都表現出色:
- 規模:目前最先進的大型模型,例如OpenAI GPT-2、NVIDIA Megatron-LM和Google T5,分別具有15億、83億和110億個參數,而DeepSpeed的ZeRO第一階段提供系統支持,以運行多達1000億個參數的模型,這是比當前最先進的模型大10倍。未來計劃增加對ZeRO第二和第三階段的支持,從而提供高達2000億個乃至數万億個參數的模型的能力。
- 速度:在各種硬件上,目前觀察到的吞吐量比當前最先進技術高出5倍。例如,為了在GPT系列工作負載上訓練大型模型,DeepSpeed將基於ZeRO的數據並行與NVIDIA Megatron-LM模型並行相結合,在具有低帶寬互連的NVIDIA GPU集群上(沒有NVIDIA NVLink或Infiniband),與僅對具有15億參數的標準GPT-2模型使用Megatron-LM相比,DeepSpeed將吞吐量提高了3.75倍。在具有高帶寬互連的NVIDIA DGX-2集群上,對於20至800億個參數的模型,速度要快3到5倍。這些吞吐量的提高來自DeepSpeed更高的內存效率以及使用較低程度的模型並行和較大的批處理量來擬合這些模型的能力。
- 成本:提高吞吐量意味著大大降低訓練成本,例如,要訓練具有200億個參數的模型,DeepSpeed需要的資源是原來的3/4。
- 易用性:只需更改幾行代碼即可使PyTorch模型使用DeepSpeed和ZeRO。與當前的模型並行庫相比,DeepSpeed不需要重新設計代碼或重構模型,它也沒有對模型尺寸、批處理大小或任何其它訓練參數加以限制。對於參數多達60億的模型,可以方便地使用由ZeRO提供的數據並行能力,而無需模型並行。而相比之下,對於參數超過13億的模型,標準數據並行將耗盡內存。ZeRO第二和第三階段將進一步增加僅通過數據並行即可訓練的模型大小。此外,DeepSpeed支持ZeRO支持的數據並行與模型並行的靈活組合。
更具體的介紹查看微軟的博客: