讓AI簡單且強大:深度學習引擎OneFlow技術實踐

機器之心發表於2019-04-25

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

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

讓AI簡單且強大:深度學習引擎OneFlow技術實踐

提綱:

  • 研發 OneFlow 的動機

  • OneFlow 技術突破

  • 總結

01 研發 OneFlow 的動機

軟體 OneFlow 簡介

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

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

案例:

  • 2015 Microsoft Resnet

  • 2016 Baidu Deep Speech 2

  • 2017 Google NMT

讓AI簡單且強大:深度學習引擎OneFlow技術實踐

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

讓AI簡單且強大:深度學習引擎OneFlow技術實踐

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

讓AI簡單且強大:深度學習引擎OneFlow技術實踐

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

讓AI簡單且強大:深度學習引擎OneFlow技術實踐

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

讓AI簡單且強大:深度學習引擎OneFlow技術實踐

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

基於純硬體的解決思路

讓AI簡單且強大:深度學習引擎OneFlow技術實踐

案例:

  • 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 的架構。

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

讓AI簡單且強大:深度學習引擎OneFlow技術實踐

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

理念:縱向擴充套件與橫向擴充套件

1.縱向擴充套件

讓AI簡單且強大:深度學習引擎OneFlow技術實踐

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

橫向擴充套件

讓AI簡單且強大:深度學習引擎OneFlow技術實踐

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

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

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

通訊密集,延遲敏感

讓AI簡單且強大:深度學習引擎OneFlow技術實踐

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

02OneFlow 技術突破

讓AI簡單且強大:深度學習引擎OneFlow技術實踐

基於靜態排程的流式計算引擎

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

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

產品對比

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

讓AI簡單且強大:深度學習引擎OneFlow技術實踐

人有我優,用資料並行加速 CNN 訓練

讓AI簡單且強大:深度學習引擎OneFlow技術實踐

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

讓AI簡單且強大:深度學習引擎OneFlow技術實踐

阿姆達爾定律

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

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

讓AI簡單且強大:深度學習引擎OneFlow技術實踐

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 模型。

讓AI簡單且強大:深度學習引擎OneFlow技術實踐

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

03 總結

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

相關文章