騰訊機智機器學習平臺由TEG架構平臺部和運營管理部團隊攜手,並和香港浸會大學電腦科學系褚曉文教授團隊深度合作聯袂打造。
背景
2018年6月25日,OpenAI在其Dota2 5v5中取得一定成績後介紹,其在訓練中batch size取100W,而1v1的訓練batch size更是達到800W;訓練時間則是以周計。騰訊內部對遊戲AI一直非常重視,也面臨大batch size收斂精度和低訓練速度慢的問題;目前batch size超過10K則收斂不到基準精度,訓練時間以天計,這對於快速迭代模型來說是遠遠不夠的。
目前業界考驗大batch size收斂能力和大資料集上訓練速度的一個權威基準是如何在ImageNet資料集上,用更大的batch size,在更短的時間內將ResNet-50/AlexNet這兩個典型的網路模型訓練到標準精度;國外多個團隊作了嘗試並取得了進展,比如UC Berkely等高校的團隊可在20分鐘將ResNet-50訓練到基準精度。
研究和解決這個問題,可以積累豐富的大batch size收斂最佳化經驗和大叢集高效能訓練經驗,並將這些經驗應用到解決遊戲AI類實際業務中;這也是我們研究這個問題的初衷。
一 4分鐘內訓練ImageNet
騰訊機智 機器學習平臺團隊,在ImageNet資料集上,4分鐘訓練好AlexNet,6.6分鐘訓練好ResNet-50,創造了AI訓練世界新紀錄。
在這之前,業界最好的水平來自:
① 日本Perferred Network公司Chainer團隊,其15分鐘訓練好ResNet-50;
② UC Berkely等高校的團隊,11分鐘訓練好AlexNet.
二 機器學習領域訓練背景
在AlexNet網路模型出現後的過去幾年中,深度學習有了長足的發展和進步,尤其是在影像、語音、機器翻譯、自然語言處理等領域帶來了跨越式提升。在AlphaGo使用深度學習方法戰勝世界圍棋冠軍李世石之後,大家對人工智慧未來的期望被再一次點燃,人工智慧成為各個領域議論的焦點。但與之相伴的也有很多問題:
• 資料量大:
有些模型的訓練資料動輒上TB,使得多輪訓練時資料讀取成為非常耗時的部分。
• 計算模型複雜:
深度網路的一個特點就是結構越深、越複雜,所表達的特徵就越豐富,在此思想下,最新的網路結構越來越複雜,從AlexNet的8層,VGG-19的19層,ResNet-50的50層,到Inception-ResNet-V2的467層和ResNet-1000的1202層等。
• 引數量大:
深度神經網路由於層次很多,引數量往往很大。ResNet-50有2500萬引數量,AlexNet有6200萬的引數量,而VGG-16引數量則達到1.38億,有的語言模型引數量甚至超過10個億[5]。
• 超引數範圍廣泛:
隨著模型複雜度的提升,模型中可供調節的超引數數量及數值範圍也在增多。例如,在CIFAR-10資料集上訓練的ResNet模型有16個可調的超引數[8],當多數超引數的取值為連續域的情況下,如此少量的超引數仍然可能造成組合爆炸。因此,最近也出現了以谷歌的Vizier為代表的系統,採用最佳化的搜尋及學習演算法為模型自動適配合適的超引數值的集合。
所有上面這些問題,對訓練速度帶來巨大的挑戰和要求。
從2010年以來,每年的ImageNet大規模視覺識別挑戰賽(ILSVRC [1],下文簡稱ImageNet挑戰賽)作為最權威的檢驗影像識別演算法效能的基準,都是機器學習領域的焦點。
隨著全世界研究者的不斷努力,ImageNet的Top-5錯誤率從2010年的28%左右,下降到2012年的15.4%(AlexNet),最終在2017年Top-5錯誤率已經下降到3%左右,遠優於人類5%的水平[2]。
在這個迭代過程其中,兩個典型的網路,AlexNet和ResNet-50具有里程碑的意義。然而,在一個英偉達的 M40 GPU 上用 ResNet-50 訓練 ImageNet 需要 14 天;如果用一個序列程式在單核 CPU 上訓練可能需要幾十年才能完成[3]。因此,如何在更短的時間內在ImageNet上訓練好AlexNet和ResNet-50一直是科研工作者研究的課題。
很多研究團隊都進行了深入嘗試,比如Facebook人工智慧實驗室與應用機器學習團隊可在1小時訓練好ImageNet [4];目前業界最好的水平來自:
① 日本Perferred Network公司Chainer團隊,其15分鐘訓練好ResNet-50; [9]
② UC Berkely等高校的團隊,11分鐘訓練好AlexNet. [3]
機智團隊想在這個問題上做出新貢獻,推動AI行業向前發展,助力AI業務取得成功。
三 訓練速度提升的挑戰
如第二節所述,由於以上四個主要矛盾,深度學習訓練時間常常以小時和天計算,如何提升訓練效率,加快模型訓練迭代效率,成了機智團隊的關注重點。要提升訓練速度,主要面臨挑戰有如下幾個方面:
3.1 大batch size帶來精度損失
為了充分利用大規模叢集算力以達到提升訓練速度的目的,人們不斷的提升batch size大小,這是因為更大的batch size允許我們在擴充套件GPU數量的同時不降低每個GPU的計算負載。
然而,過度增大batch size會帶來明顯的精度損失!這是因為在大batch size(相對於訓練樣本數)情況下,樣本隨機性降低,梯度下降方向趨於穩定,訓練就由SGD向GD趨近,這導致模型更容易收斂於初始點附近的某個區域性最優解,從而抵消了計算力增加帶來的好處。如何既增大batch size,又不降低精度,是機智團隊面臨的首要挑戰。
3.2 多機多卡擴充套件性差
深度訓練通常採用資料並行模式,資料並行模式將樣本分配給不同的GPU進行訓練。相比模型並行,資料並行簡單且可擴充套件,成為目前主流的分散式訓練方式。
分散式訓練資料並行模式下,經典的部署方式是獨立的引數伺服器(Parameter Server)來做訓練過程中梯度的收集、分發和更新工作,每一次迭代所有的GPU都要與PS多次通訊來獲取、更新引數;當節點超過一定數量時,PS的頻寬以及處理能力將成為整個系統的瓶頸。
AI訓練系統和傳統後臺系統之間的一個最主要區別是,傳統後臺系統可以透過增加節點的方式來分擔訪問請求,節點之間沒有強相關的關係;而AI訓練系統在訓練模型時需要參與訓練的所有節點都不斷的與模型引數伺服器交換和更新資料,這無形中相當於對整個系統增加了一把大鎖,對整個系統中單節點的頻寬和處理能力要求非常高,這也是AI訓練系統的特別之處,不能透過簡單的增加節點來提升系統負載能力,還需要解決多節點的擴充套件性問題。
所以如何在架構部署和演算法層面減少對頻寬需求,控制多機擴充套件中引數傳輸對訓練速度的影響,使AI訓練叢集效能可線性擴充套件,是機智團隊面臨的另一項挑戰。
3.3 如何選擇合適的超參
此外,由於超參較多,而每一個超參分佈範圍較廣,使得超參調優的耗時較長,特別是針對ImageNet這種超大資料集的情況。前文提過,CIFAR-10資料集上訓練的ResNet模型就有16個超參。
隨著專案進展,團隊還引入了很多新的關鍵技術,如後面將會提到的LARS演算法、分層同步演算法、梯度融合策略,Batch Norm替換等都會增加模型超引數量,如何在可接受的時間內尋找到較優解,是機智團隊面臨的第三個重大挑戰。
四 訓練速度提升的關鍵技術
機智團隊針對上述挑戰,分別在大batch size訓練,多機擴充套件性,及超參調整方法上取得突破,並應用到ImageNet訓練場景中,能夠在6.6分鐘內完成ResNet-50訓練,4分鐘完成AlexNet訓練——這是迄今為止ImageNet訓練的最高世界紀錄。在這個過程中,機智團隊在吸收業界最佳實踐的同時,深度融合了多項原創性關鍵技術。
4.1 超大batch size 的穩定收斂能力
1)半精度訓練與層次自適應速率縮放(LARS)演算法相結合
為了提升大batch size情況下的可擴充套件性,機智團隊將訓練資料和引數採用半精度浮點數的方式來表示,以減少計算量的同時降低頻寬需求。但半精度浮點數的表示方式不可避免的會降低模型收斂精度。
為了解決精度下降問題,機智團隊引入了層次自適應速率縮放(LARS)演算法。LARS演算法由You et al. (2017)[3]最先提出,該演算法透過對不同的層使用不同的Learning Rate,大幅度提升了大batch size場景下的訓練精度,但實際測試發現,直接將LARS演算法應用於半精度模型訓練造成很大的精度損失,這是由於乘以LARS係數後, 很多引數因半精度數值表示範圍較小而直接歸0。
為此,機智團隊引入了混合精度訓練方法來解決這個問題,透過將半精度引數轉化成單精度,然後再與LARS結合,即半精度訓練,單精度LARS最佳化及引數更新。相應的,在更新引數時使用loss scaling方法成倍擴大loss(並對應減少學習率)避免歸0影響精度。測試結果顯示,這一方法,一方面保證了計算速度,另一方面也取得了很好的收斂效果。
2)模型和引數的改進
我們在32K下復現了You et al. (2017)的測試結果,但batch size擴充套件到64K時,訓練精度未能達到基準準確性。為了提高在64K下的收斂準確性,我們對引數和模型進行了改進:1) 只對weight做正則化。 2)在You et al. (2017)模型的基礎上,進一步改進AlexNet模型。
正則化透過在損失函式後加一項懲罰項,是常用的防止模型過擬合的策略。大多數深度學習框架預設會對所有可學習的引數做正則化,包括weight, bias, BN beta和gamma(batch norm中可學習的引數)。我們發現bias, beta, gamma的引數量相對於weight來說非常小,對於AlexNet模型,bias, beta, gamma引數量總和僅佔總引數量的0.02%。
因此,bias, beta, gamma不會給模型帶來過擬合,如果我們對這些引數進行正則化,增加了計算量,還會讓模型損失一些靈活性。經過實驗驗證,我們發現不對bias, beta, gamma做正則化,模型提高了約1.3%的準確性。
最佳化正則化策略後模型收斂性得到了提升,但是AlexNet還是沒有達到基準準確性。透過對AlexNet訓練引數和輸出特徵的分析,我們發現Pool5的特徵分佈(如下圖示,顏色淺表示資料分佈少,顏色深的表示資料分佈多,總體看資料分佈範圍很廣)隨著迭代步數的增加,方差越來越大,分佈變化很劇烈,這樣會導致學習和收斂變得困難。
這個結果啟發我們在Pool5後插入一個Batch Norm層,用以規範特徵的分佈。如下圖所示,完善AlexNet後,BN5後面輸出的特徵圖分佈更加均勻,64K batch在100 epochs下收斂到58.7%,不損失準確性的同時完成了加速訓練。
3)超參調優
模型超參調優是深度學習中成本最高的部分之一,大模型的訓練過程是以小時或天計算的,特別是訓練資料集較大的情況下。因此需要有更高效的思路和方法來調優超參,機智平臺團隊在這個方面主要採取瞭如下思路:
• 引數步長由粗到細:調優引數值先以較大步長進行劃分,可以減少引數組合數量,當確定大的最優範圍之後再逐漸細化調整,例如在調整學習速率時,採取較大步長測試發現:學習率lr較大時,收斂速度前期快、後期平緩,lr較小時,前期平緩、後期較快,根據這個規律繼續做細微調整,最終得到多個不同區間的最佳學習速率;
• 低精度調參:在低精度訓練過程中,遇到的最大的一個問題就是精度丟失的問題,透過分析相關資料,放大低精度表示邊緣數值,保證引數的有效性是迴歸高精度計算的重要方法;
• 初始化資料的調參:隨著網路層數的增多,由於啟用函式的非線性,初始化引數使得模型變得不容易收斂,可以像VGGNet那樣透過首先訓練一個淺層的網路,再透過淺層網路的引數遞進初始化深層網路引數,也可以根據輸入輸出通道數的範圍來初始化初始值,一般以輸入通道數較為常見;對於全連線網路層則採用高斯分佈即可;對於shortcut的batch norm,引數gamma初始化為零。
以上思路在4分鐘訓練ImageNet專案中提升了調參效率。但調參是個浩繁的工作,後續將由內部正在測試的AutoML系統來進行。
透過以上三個方面,在ImageNet資料集上,機智平臺可將ResNet-50/AlexNet在batch size 為64K時訓練到基準精度!
4.2 超大規模GPU叢集(1024+GPUs)線性擴充套件能力
1)引數更新去中心化
資料並行訓練方式下,每一次迭代都需要做梯度規約,以TensorFlow為代表的經典分散式訓練部署方式中,中心化的引數伺服器(Parameter Server)承擔了梯度的收集、平均和分發工作,這樣的部署方式下PS的訪問頻寬容易成為瓶頸,嚴重影響可擴充套件性,機智團隊最初應對方法是引入HPC領域常用的去中心化的Allreduce方式,然而目前流行的NCCL2或baidu-allreduce中的Allreduce採用的基於環形拓撲的通訊方式,在超大規模GPU叢集場景下資料通訊會有很大的延時開銷。
機智團隊進一步將Allreduce演算法進行了改進,併成功的部署在1024+GPUs的異構叢集中,達到了理想的擴充套件效率。
2)利用分層同步和梯度分段融合最佳化Ring Allreduce
在分散式通訊中,引數傳輸耗時可用如下公式表達:
其中α 代表單節點單次資料傳輸的延時,比如準備資料,傳送資料介面呼叫等;P為資料在節點間傳輸的次數,通常是節點個數的一個倍數,β 為引數傳輸耗時的係數,不同引數傳輸方法,這個係數不同;B為網路頻寬,M為引數總位元組數,(M/B)為單次完整引數傳輸耗時。[6]
由以上公式可以看出,引數M越大,第二項所佔的比重就越大,受第一項的影響就越小,即P對整體時間的影響則越小;引數M越小,則第一項所佔的時間不可忽略,隨著P的增大,對總體時間則影響越大。對於傳輸採取Ring Allreduce演算法來講, 全域性規約操作對頻寬的需求接近於常數,不隨節點數增加而增加,所以β * (M/B)接近為常數,可變因數為 α * P;網路模型越小,傳輸的資料量越小越散,則 α * P 這塊的比重越大,整體的擴充套件性也就越差。
比如,在AlexNet神經網路中,除了兩層引數量較大的全連線層,其餘的BN層和卷積層引數量較少,各層的引數分佈差異很大。在我們的實驗環境中,使用Ring Allreduce傳輸方式,測試不同資料包大小傳輸耗時如下圖所示。從圖中可以看出,Ring Allreduce的時間開銷會隨著GPU個數的增加而顯著增大。
此外,傳輸資料塊太小也不能充分利用頻寬,多個小塊傳輸帶來極大的overhead,如圖所示。可以看到傳送同量資料時,小資料包額外開銷大,不能高效利用頻寬。
經過以上對神經網路每層的引數和資料包傳輸特性的分析,機智團隊得出以下結論:
(1) 隨著叢集節點增多,Ring Allreduce傳輸模式不夠高效。隨著節點增多,傳輸耗時中α* P部分比例會逐步增大。
(2) Ring Allreduce演算法對小Tensor不夠友好。演算法將待規約的資料拆分為N等份(N為節點總數),這導致節點數大幅增加時,Tensor碎片化,通訊網路傳輸大量小資料包,頻寬利用率很低。
針對上述問題,機智團隊提出瞭如下改進方案:
(1) 分層同步與Ring Allreduce有機結合:對叢集內GPU節點進行分組,減少P對整體時間的影響。如前論述,當P的值對系統效能影響較大時,根據具體的叢集網路結構分層,同時跨節點規約使用Ring Allreduce演算法規約, 這一改進有效的減少了每層Ring參與的節點數,降低了傳輸耗時中 α * P 的佔比。如下圖所示。原本需要對16個(即P=16)GPU進行AllReduce,現將16個GPU分為4組,每組4個GPU,首先在組內進行Reduce(4組並行執行,P1=4),然後再以每組的主GPU間進行Allreduce(P2=4),最後在每組內進行Broadcast(P3=4),這樣便大大地減少了P的影響,從而提高Allreduce的效能。
(2) 梯度融合,多次梯度傳輸合併為一次:根據具體模型設定合適的Tensor size閾值,將多次梯度傳輸合併為一次,同時超過閾值大小的Tensor不再參與融合;這樣可以防止Tensor過度碎片化,從而提升了頻寬利用率,降低了傳輸耗時。
(3) GDR技術加速Ring Allreduce:在前述方案的基礎上,將GDR技術應用於跨節點Ring,這減少了主存和視訊記憶體之間的Copy操作,同時為GPU執行規約計算提供了便利;
注:GDR(GPU Direct RDMA)是RDMA技術的GPU版本,可以提供遠端節點間視訊記憶體直接訪問,能大幅降低了CPU工作負載。
在具體到ImageNet訓練問題中,在測試梯度融合時,機智團隊根據對模型各層引數大小分析和實測結果,提出了分段融合策略,將AlexNet和ResNet-50各層分為兩段,各融合為一個Tensor參與Ring Allreduce。經過測試和分析,在1024卡場景下,AlexNet在20層和21層處分段可以達到最好效果;ResNet-50在76層和77層之間分段速度達到最優,如下圖所示。
經分段融合策略後,極大提高了反向計算和傳輸的並行度,提升了訓練速度。目前分段融合已可根據前向計算和反向計算耗時,及傳輸耗時,結合實際硬體配置和網路模型對傳輸效能進行建模,自動實現最優分段策略,自適應地選擇需要合併的引數,以達到系統最佳擴充套件效能。
3)使用Pipeline機制降低IO延遲
GPU引入深度學習之後,模型訓練速度越來越快,最優加速效能不僅依賴於高速的計算硬體,也要求有一個高效的資料輸入管道。
在一輪訓練迭代中,CPU首先從磁碟讀取資料並作預處理,然後將資料載入到計算裝置(GPU)上去。一般實現中,當 CPU 在準備資料時,GPU處於閒置狀態;相反,當GPU在訓練模型時,CPU 處於閒置狀態。因此,總訓練時間是 CPU預處理 和 GPU訓練時間的總和。
機智團隊為解決IO問題,將訓練樣本集部署在由SSD盤組成的儲存系統上,以保證CPU訪問資料的頻寬不受網路限制;同時,更為關鍵的是引入了Pipeline機制,該機制將一次訓練迭代中的資料讀入及處理和模型計算並行起來,在做模型計算的同時做下一輪迭代的資料讀取處理並放入自定義“無鎖”佇列,並透過GPU預取機制提前把處理好的資料從佇列中同步到GPU視訊記憶體中,當做下一輪的模型計算時直接從視訊記憶體讀取資料而不需要再從CPU或磁碟讀取,做到將資料讀取隱藏,IO和計算並行起來。
透過以上三方面,機智平臺在1024卡以上跑ResNet-50,擴充套件性也可以達到~99%,在2048卡上跑ResNet-50,擴充套件性還可以保持在97%!
五 平臺價值
人工智慧越來越多的融入到人們生活當中,涵蓋各行各業,包括衣食住行,交通、個性化產品等;可以服務於人們需求的人工智慧服務在未來會像水和電一樣成為基本需求,機智機器學習平臺正是在這種背景下應運而生。
提升訓練ImageNet的速度,只是機智團隊推動AI發展工作中一小部分;事實上,服務好遊戲AI等AI業務,助力AI團隊在建設AI服務時,聚焦使用者需求,而AI服務背後的模型訓練、最佳化,模型部署和運營觸手可得,才是機智的真正使命。機智機器學習平臺當前主要提供訓練加速能力。
1)訓練加速
快速完成模型訓練意味著可以做更多的模型/演算法嘗試;這不僅是體現平臺能力的一個重要指標,也是增加業務競爭力的一個關鍵所在。如果一個產品模型的訓練時間以週記或以月記,那這個產品也是生命力不旺盛的。以下是機智的兩個典型應用案例:
• X業務的訓練資料是結構化資料且量大,模型比較複雜;一直以來訓練好一個模型需要一天以上,這對於快速迭代模型演算法來說是遠遠不夠的。在應用了機智機器學習平臺後,可以在約10分鐘時間迭代一個模型,極大的加速了訓練迭代效率,為業務的成功奠定了堅實的基礎。
• 計算機視覺,是人工智慧應用的重要領域,已在交通、安防、零售、機器人等各種場景中應用開來。計算機視覺中網路模型的關鍵部件就是CNN,快速的將CNN網路訓練好可以極大的提升產品落地速度。以業界最著名的ImageNet資料集為例,要將ResNet-50模型在其上訓練到75%以上準確率,普通單機單卡需要一週左右時間。在機智機器學習平臺上,這樣的應用只需要6.6分鐘。
六 展望
未來,機智團隊將繼續保障遊戲AI業務的快速迭代,把解決imagenet訓練問題中所積累的加速方案應用到遊戲AI中,將遊戲AI的大batch size收斂問題和訓練速度問題徹底解決,協助業務取得新的突破性成果。此外,機智團隊將在平臺效能,功能上持續完善:
1)平臺效能
機智將結合模型壓縮/剪枝,量化技術,在通訊傳輸中統籌多種all-reduce演算法,做到對不同場景的模型:
• 計算瓶頸的模型多機達到線性擴充套件
• 傳輸瓶頸的模型多機達到90%以上擴充套件效率
2)平臺功能
a) AutoML(自動調參)
機智團隊認為,演算法工程師應當專注於建立新網路,推導新公式,調整超引數的浩繁工作,交給機智平臺幫你自動完成。
b) 一站式的管理服務
機器學習的模型訓練過程是個複雜的系統工整,涉及到對視覺化任務的管理,對各種資源的管理(比如CPU, GPU, FPGA, ASIC等)和排程,對訓練資料和結果資料的管理,高質量的服務體系等,將這一整套流程都打通,並且做到對使用者友好,所見即所得,是演算法工程師驗證想法最基本的需求。機智平臺將提供一站式的管理服務,想你所想,助你成功。
此外,計算機視覺類加速只是起點,未來在功能方面,機智平臺將支援多場景、多模型;結合更廣義的AutoML技術,讓AI技術賦能更廣大的業務,我們的目標是:
為使用者提供訓練,推理,模型託管全流程計算加速服務。
最終,建立從訓練加速到部署上線一站式服務平臺,打造AI服務基礎設施,助力AI業務取得成功。
七 致謝
在研究和解決imagenet訓練專案中,機智團隊小夥伴們通力合作,最終在這個問題上取得了突破。在此要特別感謝團隊內小夥伴——TEG兄弟部門運營管理部,香港浸會大學電腦科學系褚曉文教授團隊;是大家的精誠團結和專業精神,才得以讓我們在這個業界權威基準上取得新的重大突破。
此外還要隆重感謝機智平臺合作伙伴——TEG兄弟部門AI平臺部,網路平臺部,感謝兄弟部門小夥伴們一直以來的支援和信任。期盼在未來的前進道路上,機智團隊仍然能與各位同行,去創造新的好成績。
八 參考資料
[1] Large Scale Visual Recognition Challenge 2017
http://image-net.org/challenges/LSVRC/2017/
[2] http://karpathy.github.io/2014/09/02/what-i-learned-from-competing-against-a-convnet-on-imagenet/
[3] ImageNet Training in Minutes
https://arxiv.org/pdf/1709.05011.pdf
[4] Facebook 1小時訓練ImageNet
https://research.fb.com/publications/ImageNet1kIn1h/
[5] Exploring the Limits of Language Modeling
https://arxiv.org/pdf/1602.02410.pdf
[6]Optimization of Collective Communication Operations in MPICH
http://www.mcs.anl.gov/~thakur/papers/ijhpca-coll.pdf
原allreduce的公式為:t = 2(P − 1)α + (2(P − 1)β + (P −1)γ)M/P,本文根據場景進行了簡化描述。
其中P為程序個數(此次亦表示GPU個數),α 代表單節點單次資料傳輸的延時,比如準備資料,傳送資料介面呼叫等;β為點對點間單個位元組的傳輸時間,通常表示為1/B,B為網路頻寬;γ為兩個引數的單次加法運算時間;M為引數總位元組數。
[7]Google Vizier: A Service for Black-Box Optimization 4.4.2節
https://ai.google/research/pubs/pub46180
[8]PowerAI DDL, https://arxiv.org/pdf/1708.02188.pdf
[9] Extremely Large Minibatch SGD:Training ResNet-50 on ImageNet in 15 Minutes, https://arxiv.org/abs/1711.04325
本專案論文《Highly Scalable Deep Learning Training System with Mixed-Precision: Training ImageNet in Four Minutes》,將首先在arxiv上提交供查閱。