本文內容節選自由msup主辦的第七屆TOP100summit,北京一流科技有限公司首席科學家袁進輝(老師木)分享的《讓AI簡單且強大:深度學習引擎OneFlow背後的技術實踐》實錄。

北京一流科技有限公司將自動編排並行模式、靜態調度、流式執行等創新性技術相融合,構建成一套自動支持數據並行、模型並行及流水並行等多種模式的分散式深度學習框架,降低了分散式訓練門檻、極大的提高了硬體使用率。該框架已經成功幫助眾多頭部互聯網公司及人工智慧企業提升了大模型訓練效率,節約了硬體運營和使用成本,達到了降本增效的效果。一流科技是一家為企業客戶提供面向大規模大計算大模型等深度學習框架的人工智慧領域科技創新公司。

分享者袁進輝是北京一流科技有限公司創始人,任首席科學家。2008年7月在清華大學計算機系獲得工學博士學位,獲得清華大學優秀博士學位論文獎。2013年加入微軟亞洲研究院從事大規模機器學習平臺的研發工作。2014年發明瞭當時世界上最快的主題模型訓練演算法和系統LightLDA,只用數十臺伺服器即可完成以前數千臺伺服器才能實現的大規模主題模型,該技術成功應用於微軟在線廣告系統,被當時主管研究的全球副總裁周以真稱為「年度最好成果」。2015年至2016年底,專註於搭建基於異構集羣的深度學習平臺,項目榮獲微軟亞洲研究院院長特別獎 (top 1%)。2017年創立北京一流科技有限公司,致力於打造分散式深度學習平臺的事實工業標準。

編者按:2018年11月30日-12月3日,第七屆全球軟體案例研究峯會在北京國家會議中心盛大開幕,現場解讀2018年「壹佰案例榜單」。本文為北京一流科技有限公司首席科學家袁進輝(老師木)分享的《讓AI簡單且強大:深度學習引擎OneFlow背後的技術實踐》案例實錄。

提綱:

  • 研發OneFlow的動機
  • OneFlow技術突破
  • 總結

01研發OneFlow的動機

軟體OneFlow簡介

業界有人工智慧浪潮的三駕馬車之說,即數據、演算法、算力。具體到算力,業界更多關注的是硬體,譬如GPU,甚至是TPU之類的AI專用晶元。但是,人們發現,有了更快的加速器之後,制約大規模分散式訓練算力的瓶頸是軟體。怎麼幫助數據科學家和研究員們更輕鬆的把各種演算法在底層硬體上跑起來,而且充分釋放底層硬體的潛力,正是軟體框架需要解決的問題。目前,已有的開源深度學習框架對數據並行場景解決的比較好,但遇到模型越來越大的場景就沒有好辦法。用戶或者束手無策,或者只能付出極大成本基於開源框架做深度定製開發來滿足需求。OneFlow團隊的目標是研發一個通用框架自動解決這些問題,讓那些沒有框架研發能力的團隊也能夠享受分散式GPU集羣帶來的效率,這是我們歷時兩年多研發一套全新深度學習框架的初衷。

背後的動機:計算力是深度學習發展的最重要的推動力。

案例:2015 Microsoft Resnet

2016 Baidu Deep Speech 2

2017 Google NMT

2015年微軟研究院發明的ResNet需要的計算量是7乘以10的18次方次計算(ExaFlops)。當然,可以推算一下用一顆24核的CPU來計算,大概需要多久能完成這些計算,也可以推算用幾千個核心的GPU來算需要多長時間。可能是需要幾個月或幾個星期的時間。除了計算量在增長,模型大小也在增長,ResNet 這種CNN模型通常是幾千萬參數,需要幾百兆位元組的存儲空間,百度研發的Deep Speech模型到了三億參數的規模,然後Google的機器翻譯模型NMT,已經到了幾十億參數,整個模型在一塊GPU上已經放不下了。這種情況,數據並行無濟於事,需要模型並行或流水並行來解決大模型的分散式訓練問題。很不幸,目前還沒有開源框架支持這些需求,一些大公司通過內部定製的系統來支持這種需求。

今年上半年Facebook發布了一個研究結果,收集35億張弱標註圖片,使用幾百塊GPU,經過接近一個月的時間,訓練了一個用於圖片分類的卷積神經網路模型,它能做到什麼效果呢?能提高6個百分點的準確率。這是非常了不得的成績,演算法基本上沒什麼變化,僅僅是通過採用更多的數據和計算就能把top-1的準確率提高了這麼多。要知道,對於一個商業價值很高的場景,提高0.5個百分點可能是一個團隊一年的KPI。

九月份Google發表了BigGAN模型,研究人員通過提高圖片的解析度來訓練更大的GAN模型,CNN中間的activation和反向gradient會非常多,當然計算量也會大的非常多,基於TPU集羣來完成訓練。這個手段同樣獲得了比以前的GAN模型好的多的效果。

上個月,Google又發表了BERT模型,相當於一種大的多的transformer模型,在16個TPU上訓練了4天,然後基於這個語言模型作為主幹網路去解決各種常見的自然語言處理任務,發現在各任務上全面超越了以前的方法。很不幸,目前還沒有出現在GPU集羣上從零開始訓練BERT-Large模型的辦法,如果想在自己的業務裏應用BERT,只能去下載Google預訓練好的模型,然後做少量微調來使用。這倒不是資源不足的問題,即使是已經搭建了大規模的GPU集羣的客戶也無能為力,有錢也解決不了。

深度學習經過這幾年的爆髮式發展,特別引人注目的演算法層面的創新越來越少了,今年比較吸引眼球的進步都來自於計算力,也就是人們常說的「大力出奇蹟」的方式。怎麼才能讓更多的企業用戶能享受到算力提升的紅利,幫助演算法科學家完成更多的KPI, 這是OneFlow非常關心的問題。常言道,工欲善其事必先利其器,框架在深度學習研究和落地的過程中就扮演了「工具」的角色,好的工具能大大加速人工智慧研發的效率,甚至可能成為行業競爭的決勝法寶。從BigGAN和BERT等例子也可以看出來,當一家公司掌握了其他人不掌握的工具時,就可以引領演算法研究的潮流,反過來,當一家公司的基礎設施跟不上的時候,也就沒辦法做前沿探索,即使是做研究也只能跟在Google後面,因此稱深度學習框架是人工智慧制高點的戰略武器一點不為過。

基於純硬體的解決思路

案例:Nvidia DGX-2

IBM Power9 Server

英偉達通過銷售GPU成為這一波AI計算力紅利的最大受益者,英偉達除了把單個設備做的越來越快,也做了伺服器架構方面的諸多創新,出品了一系列超級計算盒子,每個盒子裡面可以集成8個或者是16個計算力非常強的GPU(譬如DGX-1是P100,今年推出的DGX-2是V100),更特別的是,這些GPU之間使用了非常高速的互聯,能夠實現GPU之間點對點150GB以上的傳輸帶寬,比常見的PCIe帶寬要高一個數量級。這種設計使得DGX伺服器能夠使得16塊GPU一起工作時幾乎像一個單體晶元那樣輸出超強算力。

當然還有比DGX更特別的伺服器,比如說IBM出的Power9 Server,它的獨特之處在於他的CPU使用了不同於Intel x86 CPU的架構,而且支持CPU和GPU之間NV Link互連,意味著CPU和GPU之間的數據傳輸也能夠做到150GB以上的帶寬。目前世界排名第一的超級計算機Summit就使用了類似Power9 Server的架構。

基於這麼強的硬體就能解決計算力的問題嗎?

IBM和Nvidia一起搭建了世界上最強的超級計算機Summit,一共用了2萬多塊V100 GPU,還使用了最先進的互聯技術(NVLink, Infiniband),要說最強的硬體,除了TPU Cluster,應該沒有更好的了,這是不是就夠了呢?IBM首席科學家在今年的ASPLOS(計算機體系結構頂級會議)上做了一個特邀報告,主題是「只有很強的硬體,沒有很好的軟體還是不能解決擴展問題」。現在國內擁有幾千塊GPU乃至上萬塊GPU的頭部公司不在少數,但基於開源框架能訓練BERT-Large模型嗎?不行,這也印證了軟體框架瓶頸的問題:購買了很多的硬體,但用不起來,或者說不能很好的用起來。

理念:縱向擴展與橫向擴展

縱向擴展

縱向擴展是通過把單個設備或者是單個機器做的越來越強,或通過編譯器優化的手段讓作業在一個設備上或者是一個機器內部把硬體性能發揮到極致來滿足現在日益增長的計算需求。硬體從多核架構CPU發展到眾核架構GPU,GPU從P100到V100, 為了追求更高的效率,甚至研發FPGA或ASIC晶元來獲得更高算力。當前最知名的AI晶元是Google的TPU,寒武紀,華為,阿里,百度等本土公司也在研發AI晶元。AI晶元的主要問題是有物理限制(譬如製程,功耗,同步時鐘等等約束),人們不能生產出計算力任意大的晶元。也有人把這個現象稱為硅基擴展瓶頸(Silicon Scaling)。除了提高單個晶元的吞吐率,英偉達的DGX也是縱向擴展的例子,DGX通過在一個機器內部高速互聯手段實現晶元之間點對點極高的傳輸帶寬,從而使得多晶元間協作起來更加高效。

橫向擴展

如果縱向擴展仍不能滿足需求,人們繼續把多臺伺服器通過高速乙太網或Infiniband連接起來組成集羣來實現更高算力。如果能投入多少硬體資源,就得到多少計算力,那計算力瓶頸就迎刃而解了。理想很豐滿,現實很骨感。一方面,晶元間互聯帶寬要比片內數據訪問帶寬低一到兩個數量級,在晶元間搬運數據成為瓶頸,另一方面,編寫在多晶元上高效運行的軟體非常挑戰,以深度學習為例,神經網路的結構不同,效率最高的並行方式(邏輯任務向物理計算單元的映射)也不同。在集羣上實現線性加速比縱向擴展更有想像空間,但實現難度更大。一個理想的橫向擴展方案,不管底層實際使用了多少鬆散耦合在一起的晶元,在上層用戶眼裡就像在一個專門為當前任務打造的巨大單體晶元一樣,編程簡單而且任務運行時能把底層每一個獨立的晶元都利用充分。要實現這個目的,必須依靠軟體框架。

邏輯任務到物理拓撲之間的最優映射複雜多變

給定一個特定的神經網路模型和一批計算資源,從任務到設備之間的映射有多種方式,但不同的映射方案運行效率不同。哪種方案最優既取決於作業本身的特性,也取決於底層硬體的拓撲。神經網路由很多局部計算(通常稱為kernel)搭建組成,每一個局部計算是採用數據並行,還是模型並行取決於這個局部任務的計算傳輸比。現在業界討論比較多的卷積運算參數量很小,但中間結果量大,所以最划算的方法是對數據進行切分,不同的設備處理不同的數據,在設備之間偶爾進行參數同步,這種數據並行方法基本上是一個已經被解決的問題。還有一些運算,中間計算結果相對於參數量更少,就適合模型並行。還有一些網路參數量很大或中間計算結果都很大,可能採用流水並行(也就是接力的形式)是最優的。模型並行和流水並行中通信的數據路由要比數據並行複雜,同時,怎麼重疊計算和傳輸從而提高設備利用率也非常挑戰,現有開源框架對這些更複雜的並行模式的支持還比較初級。

通信密集,延遲敏感

左圖展示了一個常見的大數據處理引擎的架構,集羣中的計算資源一般分成用於中心調度的Master節點和用於處理數據的Worker節點。Master節點以有向無環圖(DAG)的方式管理整個作業的進度,同時監控所有Worker的資源使用狀況,在合適的時機把一個子任務(Task)分配給某個Worker去做,某個Worker在完成一個子任務之後,會向Master節點彙報,等待Master分配新的任務。在傳統大數據處理中,Worker執行一個子任務的時間量級一般在幾十秒鐘或數分鐘。其它開銷,諸如發生在Master節點那裡的排隊開銷,Master和Worker之間對話的時間開銷,以及數據傳輸開銷都是數十毫秒,相對於Worker的工作時間可以被忽略。但是深度學習訓練的負載與此不同,深度學習訓練更接近流式計算,一方面是因為隨機梯度下降演算法採用的小批次訓練,計算粒度小,另一方面是因為底層硬體吞吐率可能是CPU的數十倍,計算太快。直接後果就是,數據處理時間越來越短,每個子任務可能幾百毫秒就完成了,在這種情況下,之前可忽略的那種幾十乃至幾百毫秒的開銷就非常顯著了,如果不能通過技術手段把這些開銷消除或掩蓋掉,整個系統的性能就非常低。

02OneFlow技術突破

基於靜態調度的流式計算引擎

為了對任意作業和資源都達到類似巨大單體專用晶元的效果,OneFlow首創了靜態調度(左圖)和流式執行(右圖)架構的深度學習框架。靜態調度是什麼思路呢?它來自於計算機體系結構。我們熟知的CPU晶元中真正做算術運算的器件只佔很小比例的面積,大部分面積在做亂序執行,流水線和緩衝區的管理。學界和工業界很久以前就開始探索怎麼讓晶元的有效面積儘可能多的做算術運算,靜態調度的思路應運而生,基本上是把流水管理,指令排布之類的工作從硬體轉移至編譯器。這樣硬體複雜度就可以大幅降低,當然編譯器複雜度肯定會提高很多。有一個叫VLIW(超長指令集架構)的指令集就採用了這種思路。OneFlow的靜態調度體現在兩方面,首先,編譯器自動解決從邏輯任務到硬體資源的映射,包括數據並行,模型並行,流水並行的設備分配以及數據路由方案,大大降低了分散式編程的複雜度,用戶只需要關心任務的邏輯結構以及本次任務可使用的硬體資源,而不用去編程實現數據在硬體資源中的流動機制;其次,靜態調度把所有能在正式運行之前得到的調度策略,資源管理策略等問題都在編譯階段解決,運行時就不需要在線求解最優的調度方案,從而大大降低運行時開銷。

經過靜態編譯,每個設備負責運行的子任務是預先可知的,每個子任務的上下游依賴也預先可知,在運行任務時,就不再需要中心調度器,只需要支持上下游任務之間局部的握手信號即可,即生產者向消費者發送的請求以及消費者向生產者發送的確認,整個系統以全鏈路非同步的方式運行。這個思路也來自於晶元設計領域一種叫非同步電路的技術。OneFlow另一個區別於其它深度學習框架的特色是把數據搬運看成一等公民,在靜態分析階段就把磁碟IO,主存和設備之間數據搬運,節點間數據搬運看作和計算同等重要的任務,在代價分析和調度策略裏作為一等公民進行顯式建模,從而得到重疊傳輸和計算的最優方案。與此相對,已有開源框架把數據搬運當成二等公民處理,編譯期的注意力主要集中在計算的優化上。熟悉軟體定義網路(SDN)技術的朋友可以發現,OneFlow編譯器相當於網路的控制平面,用於獲取數據計算和轉發策略,運行時相當於網路的數據平面,執行體依照控制層面的策略去轉發和處理數據。

產品對比

OneFlow歷經兩年的研發,2018年10月份才推出1.0版本,還是一個很年輕的系統,目前正在客戶的生產環境裡面試用和迭代。實事求是的講,我們在模型的豐富程度,易用性,多語言支持等方面還有比較大的提升空間,目前是落後於其它框架的。但是,OneFlow在企業級大規模應用上是稱得上遙遙領先的:(1)分散式最容易使用,用戶在寫程序的時候是感受不到多機和單機的區別的;(2)OneFlow支持數據並行,模型並行和流水並行,而其它框架只支持最容易支持的數據並行;(3)OneFlow在分散式訓練時的擴展能力,加速比是最優秀的。這些特點也正是OneFlow作為企業級深度學習框架,比已有開源深度學習框架優秀之處。

人有我優,用數據並行加速CNN訓練

卷積神經網路(CNN)作為最容易解決的一個問題,是大家最喜歡拿來做基準測試的應用。在過去一年,很多公司用數據並行方法,已經可以用數千塊GPU做到幾分鐘就在ImageNet數據集上訓練好ResNet模型。如果發現TensorFlow的參數伺服器不給力,上層使用Horovod,底層使用Nvidia NCCL已經可以做到很漂亮的結果。需要注意的是,以前社區有一個認識是TensorFlow並行做的不好,速度比其它框架慢,實際上今天已經不是這樣了,TensorFlow團隊的benchmark項目(github.com/tensorflow/b)針對CNN做了很多優化,做數據並行已經是開源框架裏最優秀之一了。我們使用完全一樣的演算法和硬體(V100 GPU, 100Gbps RDMA網路),和TensorFlow benchmark對比會發現,無論是基於單機多卡,還是多機多卡都是比TensorFlow快。上圖左邊是OneFlow,右邊是TensorFlow,除了AlexNet遇到硬體瓶頸,OneFlow都能做到線性加速,TensorFlow在單機多卡和多機多卡上與OneFlow還是有一定的差距。

阿姆達爾定律

上面的評測結果中,在32卡時,OneFlow仍是線性加速,當卡數增加到一定程度,譬如幾百或者是上千時遲早會遇到天花板。並行效率不同的系統,只是遇到天花板時間早晚的問題,這是阿姆達爾定律所揭示的規律。比如說上圖綠色曲線表示一個並行度(parallel portion)為95%的任務,什麼時候遇到天花板呢?可以計算出來,加速到20倍的時候就到了天花板了,後面投入再多的資源進去它也不可能再加速了。假設系統的並行度不隨卡數變化,在卡數少時,大部分系統還是比較接近線性的,各個系統之間差別很小,但當卡數增多時,系統遲早會遇到天花板,即使增加再多的GPU也不會進一步提升吞吐率。這表明,在卡數比較少時實現線性加速比不一定能在卡很多時還能實現線性加速,但在卡數較少時就實現不了線性加速,在卡數更多時肯定距離線性加速更遠。由此可見,把系統的運行時開銷優化到極致,對於大規模集羣訓練效率是至關重要的。

人無我有,分散式訓練BERT-Large模型

BERT-Large是谷歌最近推出的一個學習語言模型的大型神經網路,基本上在常見的自然語言處理任務上都顯著超越了以前的方法。BERT-large有24層,整個模型大概1.3G,每一層中間結果都蠻大的,如果不做內存優化,對於32GB顯存的V100,一次也就處理八九個句子。雖然BERT是個大殺器,但客戶想基於自己語料重新訓練一個BERT-Large模型,卻不可能。谷歌在TPU Cluster上用16個TPU訓練BERT-Large需要4天時間。沒有TPU的用戶,只能使用GPU,最主要的是,現在還沒有開源的分散式解決方案,谷歌放出來TensorFlow代碼只支持單GPU卡,如果用戶做一些定製去支持分散式,很遺憾,加速比也很不理想。如左上角圖所示,即使是在有NVLink互聯的單機八卡伺服器上,TensorFlow也只能實現四五倍的加速,按這種加速比去推算一下,即使是使用幾十塊V100也是需要一個月以上的時間。在Google BERT論文發表後不久,我們團隊就基於OneFlow實現了和TensorFlow準確率一樣的BERT,在單機八卡伺服器上數據並行接近線性加速,在8機64卡的配置下,也能跑到50倍以上的加速比。這還不是線性加速比,我們正在做一些優化工作,不久以後對於BERT-Large在多機多卡也能實現線性加速比。OneFlow現在的實現在單精度條件下只需要8天就能訓練出來BERT-Large模型,如果加上半精度支持,時間會再縮短一半,只需要三四天。需要指出的是,Google BERT的詞典只有4萬個單詞,當詞表達到幾十萬或上百萬級別時,embedding層就無法用數據並行計算了,必須做模型並行,而後續的層次可以繼續使用數據並行,也就是混合併行,OneFlow可以很方便的支持起來。最近,我們已經開始為幾家頭部互聯網公司提供BERT訓練服務,在客戶自己的數據集上訓練BERT-Large模型。

以訓練安防領域的大規模人臉識別模型為例,當人臉類別達到百萬級時,最後的全連接層必須使用模型並行,要解決這個問題,用戶就不得不深度hack 已有開源框架,此時會面臨易用性和高效性的難題。詞嵌入和廣告/推薦系統領域也存在許多大模型的問題,模型容量可達幾十GB甚至幾百GB乃至TB,也只有少數頭部企業不計研發成本才能做一些定製開發來支持這些需求。OneFlow 可以很方便高效的支持這些需求,大大節省用戶成本,幫助用戶完成以前搞不定的事情。

05總結

一路走來,我們深切體會了do right things, do things right 如此重要。在諸多方向裏,我們經過反覆論證,認為這個領域最關鍵也最難的問題是橫向擴展,從公司成立之初,就立下解決這個業界公認難題的目標。不同於其它框架的技術路線,OneFlow以軟硬協同設計為指導思想,從晶元設計領域借鑒了大量有益的思路,在純軟體層面解決了橫向擴展難題。我們深信現在OneFlow的技術路線是解決深度學習橫向擴展難題的必由之路,在我們走通這條路徑之後,很高興看到技術社區其它團隊已經開始沿著這個方向進發。創新和創造是OneFlow決勝的法寶,僅僅follow已有框架走過的路是不可能實現超越的,唯有創新纔有機會。最後,我們深感真正有價值的事都是長跑,除了技術因素,情懷和堅持也必不可少,seeing is believing, believing is seeing。

歡迎關注OneFlow公眾號


推薦閱讀:
相關文章