近日,阿里云機(jī)器學(xué)習(xí)平臺PAI主導(dǎo)的論文《Whale: Efficient Giant Model Training over Heterogeneous GPUs》,高效大模型訓(xùn)練框架Whale入選USENIX ATC’22。Whale通過對不同并行化策略進(jìn)行統(tǒng)一抽象、封裝,在一套分布式訓(xùn)練框架中支持多種并行策略,并進(jìn)行顯存、計(jì)算、通信等全方位的優(yōu)化,來提供易用、高效的分布式訓(xùn)練框架。
USENIX Annul Technical Conference (USENIX ATC),是計(jì)算機(jī)系統(tǒng)領(lǐng)域國際頂級學(xué)術(shù)會議 (CCF-A),自1992年舉辦第一屆USENIX ATC會議以來,至今已成功舉辦30多屆,在學(xué)術(shù)和工業(yè)界都有巨大的影響力。USENIX ATC2022將于2022年7月11日召開。此次入選意味著阿里云機(jī)器學(xué)習(xí)平臺PAI自研的深度學(xué)習(xí)分布式模型訓(xùn)練系統(tǒng)達(dá)到了全球業(yè)界先進(jìn)水平,獲得了國際學(xué)者的認(rèn)可,展現(xiàn)了中國機(jī)器學(xué)習(xí)系統(tǒng)技術(shù)創(chuàng)新在國際上的競爭力。
Whale是阿里云機(jī)器學(xué)習(xí)PAI平臺自研的高效、通用、硬件感知的大模型分布式訓(xùn)練框架,現(xiàn)已開源,開源后的名稱是EPL(Easy Parallel Library)。Whale通過統(tǒng)一的策略抽象來實(shí)現(xiàn)各種分布式策略的表達(dá),并通過硬件感知和自動(dòng)計(jì)算圖改寫及優(yōu)化完成高效的分布式模型實(shí)現(xiàn)。

隨著模型參數(shù)規(guī)模飛速增長,模型效果也在逐步提高,但同時(shí)也為訓(xùn)練框架帶來更大的挑戰(zhàn)。當(dāng)前的業(yè)界分布式訓(xùn)練框架只支持少量的并行策略,缺乏一個(gè)統(tǒng)一的抽象來支持所有的并行策略及其混合策略。同時(shí)實(shí)現(xiàn)復(fù)雜的并行策略需要大量的模型代碼改造和對底層系統(tǒng)的理解,大大增加了用戶的使用難度。由于集群中異構(gòu)GPU計(jì)算能力和顯存的差異,靜態(tài)的并行化策略無法充分利用異構(gòu)資源實(shí)現(xiàn)高效訓(xùn)練。
針對這些問題,Whale抽象并定義了兩個(gè)分布式原語(replicate和split) 。用戶可以通過在模型上添加幾行原語標(biāo)記,即來表達(dá)和轉(zhuǎn)換各種并行化策略及其組合,極大降低了分布式框架的使用門檻。Whale runtime將用戶的標(biāo)記信息融合到計(jì)算圖中,自動(dòng)完成模型的并行化優(yōu)化。同時(shí)Whale提供了基于硬件感知的自動(dòng)化分布式并行策略,優(yōu)化在異構(gòu)GPU集群上分布式訓(xùn)練性能。Whale的設(shè)計(jì)很好地平衡了模型用戶的干預(yù)和系統(tǒng)優(yōu)化機(jī)會,讓每一個(gè)算法工程師都能輕松高效訓(xùn)練分布式大模型任務(wù)。借助Whale,阿里達(dá)摩院M6模型使用512張V100 GPU在10天內(nèi)即可完成10萬億多模態(tài)預(yù)訓(xùn)練模型。
Whale起源于阿里內(nèi)部業(yè)務(wù),在阿里內(nèi)部已經(jīng)支持圖像、推薦、語音、視頻、自然語言、多模態(tài)等業(yè)務(wù)場景。并且能無感支持模型規(guī)模橫向擴(kuò)展,最大完成了10萬億規(guī)模的M6模型訓(xùn)練。同時(shí)Whale已經(jīng)集成在阿里云機(jī)器學(xué)習(xí)平臺PAI中,大家可以在阿里云PAI平臺上使用Whale。機(jī)器學(xué)習(xí)平臺PAI是面向開發(fā)者和企業(yè)的AI工程化平臺,提供了覆蓋數(shù)據(jù)準(zhǔn)備、模型開發(fā)、模型訓(xùn)練、模型部署的全流程服務(wù)。
Whale(EPL)開源地址(https://github.com/alibaba/EasyParallelLibrary),歡迎大家來試用和給出建議。
論文名稱:Whale: Efficient Giant Model Training over Heterogeneous GPUs
論文作者:賈賢艷,江樂,王昂,肖文聰,石子驥,張杰,李昕元,陳浪石,李永,鄭禎,劉小勇,林偉
開源鏈接:https://github.com/alibaba/easyparallellibrary
論文鏈接:https://www.usenix.org/conference/atc22/presentation/jia-xianyan