谷歌如何打造世界上最快的AI超級計算機系統？（下）

最近谷歌公司發表了一篇轟動人工智慧系統界的論文[1]，介紹了他們如何基於MLPerf標準去重新整理深度學習訓練速度的世界紀錄。比如，23秒完成BERT訓練！28秒完成ImageNet訓練！文章洋洋灑灑，從系統，演算法，編譯器，應用角度全方位地闡述了谷歌如何打造基於他們自己晶片的AI超級計算機以及深度學習系統。文章一共19位作者，包含了谷歌深度學習系統團隊的一些專家。本文簡要討論一下谷歌公司的這篇文章，分為上下兩部分。如有問題，可以透過郵箱[email protected]聯絡筆者(新加坡國立大學高效能人工智慧實驗室主任、上海壁仞公司兼職顧問尤洋)。筆者曾在UC Berkeley讀博期間在谷歌公司總部的谷歌大腦團隊實習4次。

前文回顧：谷歌如何打造世界上最快的AI超級計算機系統？（上）

第3節：擴充套件技術

第3.1節：模型並行

對於有些模型而言，僅僅用資料並行是不夠的，所以我們加上模型並行在多組TPU超級計算機上獲得更高的併發性。我們用XLA的SPMD(單程式多資料)劃分工具基於輕量級的標註去自動平分模型圖。在影像分割模型中(例如SSD和Mask-RCNN)，我們透過標註輸入圖片去實現空間分割。SPMD劃分工具能自動沿著空間維度去做平行計算。這些模型有相對大的空間維度(比如Mask-RCNN的空間維度是8000×1333，SSD的空間維度是300×300)。SPMD劃分工具透過插入邊界交換通訊操作基於空間分割計算去計算下一步所需要的響應資料。SSD和Mask-RCNN都能進行8個核上的空間劃分區以實現最大化的併發度。我們又做了兩點：通訊最佳化以及透過XLA編譯器的SPMD方法(Lepikhin等人的論文[2],2020)去消除Amdahl瓶頸。這兩點可以提高空間劃分的併發度。例如，雖然Mask-RCNN能用最大batch size是256，但是我們卻可以將Mask-RCNN的訓練擴充套件到1024個加速器核上。

在像MLPerf的transformer基準這樣的語言模型中，空間維度是很小的。所以，我們依照Shazeer等人在2018年提出的方法[3]去劃分特徵維度。我們實現的時候用SPMD劃分工具去做標註。在這個方法裡，模型的引數和響應被劃分在TPU超級計算機的一個片上(超級計算機在網格視角上被劃分成了很多片)。在前向傳播中，每個片上的核完成區域性矩陣的矩陣乘法計算。然後我們透過all-reduce操作把片上所有核上的區域性響應資料整合起來用於模型下一層的計算。反向傳播中有類似的區域性矩陣的矩陣乘法以及之後的all-reduce去整合響應資料和梯度資料。正如模型的引數被劃分那樣，梯度被劃分到不同核以及TPU超級計算機的每個片上，之後區域性資料被整合起來。我們將在第3.3節中介紹多組超級計算機上最佳化梯度整合的技術。

第3.2節：引數更新劃分

在傳統的資料並行中，模型引數被複制並在每次訓練迴圈結束後被最佳化器更新。然而，當每個核的批次數不夠大的時候，計算或許會變成一個瓶頸。例如，我們以MLPerf的BERT訓練為例，在512個第三代TPU晶片上，LAMB最佳化器的引數更新時間可以佔到整個迴圈時間的18%。Xu等人在2020年提出了引數更新劃分技術[4]，這種分散式計算技術首先執行一個reduce-scatter操作，然後使得每個加速器有整合梯度的一部分。這樣每個加速器就可以算出相應的被更新的區域性引數。在下一步，每個被更新的區域性引數被全域性廣播到各個加速器，這樣使得每個加速器上都有被更新的全域性引數。為了獲得更高的加速比，我們同時用資料並行和模型並行去處理引數更新劃分。在影像分割模型中，引數是被複制的，這種情況下引數更新劃分類似於資料並行。然後，當引數被分佈後到不同的核之後，我們就執行多個併發的引數更新劃分。

第3.3節：最佳化全域性加法通訊操作

根據Mattson等人在2019年發表的論文[5]我們知道，在MLPerf中，梯度加法對實現高的強可擴充套件性效率非常重要。為了在多組TPU超級計算機上最佳化梯度加法通訊操作，我們在網格的Y方向上充分利用torus wrap連結。我們用一個雙向的環連結在網格的Y方向上去執行reduce-scatter操作，這樣的輸出是沿Y方向上全域性累加梯度的一部分。接下來，我們在X方向上執行reduce-scatter操作。之後我們輸入區域性梯度去進行引數更新。被更新的引數會先後在網格的X方向和網格的Y方向上被廣播出去。值得注意的是，在資料並行中，網格X方向上的傳輸負載是網格Y方向上傳輸負載的32倍。

在MLPerf的transformer基準中，我們透過將模型引數劃分到4個鄰近的TPU核上去執行分散式矩陣矩陣乘法操作。這些核被放置在X方向上的一條線上。在神經網路訓練的前向傳播中，all-reduce操作被執行在X方向上鄰居的短的環形網路連結上。然而，沿著X方向的梯度加法會跳過模型並行的鄰居(因為它們劃分了梯度)。圖4展示了在Transformer基準中不同的環狀加法操作。在BERT和Transformer模型訓練中，我們也用了bfloat16浮點數運算格式去進一步減少梯度加法的開銷。

圖4：這個圖片展示了在4個晶片上基於16×8網格的模型並行。這裡展示了3個不同的環狀reduction操作：(1)黑色的環狀reduction用於模型並行的前向操作；(2)紅色的環用於梯度塊在Y方向的reduce scatter操作；(3)藍色虛線表示在模型同類中梯度reduction操作(這裡只展示了模型變好為0的情況)[1]。

第3.4節：評估指標的分散式計算

深度神經網路訓練和測試在TPU加速器中執行在一個緊密的迴圈中。訓練的結果可以更新處在TPU加速器上HBM(高速頻寬記憶體)中的模型引數。根據MLPerf的規則，更新的引數隨後每隔幾個epoch就用於預測從而得到MLPerf的輸出指標。在一些benchmark中，當預測的批次數大於預測資料集的樣本數時，我們會在預測資料集後邊加上一些仿造資料(僅僅為了便於處理資料，不影響實驗結果)。在TensorFlow實現中，預測輸出的張量資料被用來在TPU主機上計算預測指標(比如ResNet-50基準的Top-1精度)。然而，在JAX實現中，預測指標的計算卻是全分散式的，我們透過全域性加法操作得到最終的結果。

第3.5節：輸入資料的流水線最佳化

一個擴充套件ResNet-50的重要挑戰就是主機輸入資料流水線上的負載不均衡。由於多組超級計算機的規模太大了，所以一些主機的輸入資料流水線有很高的JPEG資料解壓開銷。我們的解決方案是把解壓後的圖片存在主機記憶體中，這樣主機輸入流水線只用執行3個操作：(1)隨機裁剪;(2) 隨機翻轉；(3)按MLPerf的要求進行圖片正則化使得他們有常量的平均值和方差。這會極大地增加主機輸入資料流水線的吞吐率，使得它能創造一個大的預取緩衝區(prefetch buffer)。所以，主機資料流水線預處理一個大的輸入圖片的同時，它能向TPU輸送預取緩衝區的圖片。這樣，我們就消除了多組TPU超級計算機輸入資料流水線上的負載不均衡。

第4節：模型最佳化

第4.1節：BERT

用維基百科資料集去訓練BERT(Devlin等人的論文[6]，2018)是MLPerf0.7版本中新加入的基準。這是一個用於語言理解的基於雙向transformer結構的預訓練任務。在擴充套件BERT的時候，我們用LAMB最佳化器(尤洋等人的論文[7]，2019)可以很有效地把BERT的batch size擴充套件到很大。這樣，我們可以在4096個晶片上做有效地資料並行。這裡筆者需要指出一點，筆者在ICLR2020會議上發表論文中提出LAMB最佳化器，當時把BERT的訓練時間從3天縮短到了76分鐘。本文中，谷歌也用LAMB最佳化器，為何訓練時間只有24秒？原因有兩點：一方面，谷歌確實比筆者用了更大規模的超級計算機；另一方面，更重要的是，MLPerf收斂標準跟筆者的Devlin標準是不一定的，所以MLPerf標準的計算代價本身就比Devlin標準低。

在大系統上擴充套件BERT需要兩方面的數值最佳化：每步所需要的時間以及收斂需要的步數。

除了第3節提到的最佳化技巧之外，在模型層面，我們在減少BERT每步所需的時間上做了以下最佳化：減輕各體系結構單元(比如記憶體頻寬，向量單元，以及暫存器)的壓力使得TPU上的矩陣計算單元有最小的流水線開銷。為了降低記憶體頻寬的開銷，我們對模型響應和梯度加法用了bfloat16浮點數運算格式，這極大地減少了每步運算所需的時間並且沒有對模型的收斂產生負面影響。為了減低向量單元的壓力，我們利用標量乘法和矩陣乘法的可變性把標量乘法和標量除法移到矩陣矩陣乘法維度小的那一邊。為了降低暫存器溢位的頻率，我們把諸如層正則化變數的一些小變數合併成一個大的TensorFlow張量。這樣會大大減少儲存在暫存器中的變數地址數目，因此可以減少每步訓練的時間。

為了減少收斂所需的步數，我們優化了超引數以及在輸入資料流水線中做了打亂資料順序的操作。首先，我們依賴LAMB演算法的擴充套件性用谷歌的Vizier(Golovin等人的論文[8]，2017)自動調參工具去調了大批次情況下的超引數。這樣允許我們最大化資料並行，使得我們相比於模型並行用更短的時間達到相同的精度。其次，我們優化了輸入資料流水線中打亂資料順序的方法從而能保證大系統的收斂性。這一點對保證最佳化器的隨機性很重要，因為有上千個主機的大系統往往分給每個主機的資料並不多。

第4.2節：ResNet-50

何愷明博士提出的ResNet-50模型是機器學習系統評測中最常用的模型之一。MLPerf用一千類ImageNet資料集訓練ResNet-50模型作為圖片分類的基準。特別地，MLPerf用了Goyal等人在2017年提出的1.5版本的ResNet-50。這個1.5版本的ResNet-50模型跟最初何愷明博士提出ResNet-50模型有輕微的差異。為了能夠把ResNet-50基準擴充套件到多組三代TPU超級計算機上，我們應用了一系列最佳化技術：(1)分散式評測；(2)分散式批標準化；(3)引數更新劃分；(4)梯度加法最佳化。MLPerf的0.7版本用了尤洋(筆者)等人提出的LARS最佳化器去動態擴充套件學習率，從而能讓ResNet-50增大批次數去最大化資料並行度。在調節了動量的超引數後，我們能夠用65536的批次數在88個epoch內完成訓練。

第5節：效能分析

表格1展示了TensorFlow(1.x版本)和JAX在多組TPU超級計算機上完成MLPerf各個基準模型所需要的端到端訓練時間。表格1也展示了我們在0.7版MLPerf的提交相對於0.6版MLPerf提交的加速比。基於這篇文章所使用的技術，我們在半分鐘內完成了4個重要模型的訓練。

表格1：TensorFlow(1.x版本)和JAX在多組TPU超級計算機上完成MLPerf各個基準模型所需要的端到端訓練時間。這裡也展示了我們在0.7版MLPerf的提交相對於0.6版MLPerf提交的加速比。

表格2比較了TensorFlow和JAX在有4096個晶片和2048個晶片的大規模分散式系統上初始化時間的差異。TensorFlow的初始化時間介於498秒到1040秒之間。然後JAX的初始化時間要低很多，只是介於122秒到294秒之間。這樣因為它們在分散式程式設計模型的區別造成的。它們之間的權衡，效能比較，以及可用性分析詳見第2節的描述。

表格2：TensorFlow和JAX在大規模分散式系統上的初始化時間對比[1]

從圖5中我們可以觀察到ResNet-50模型訓練加速比與TPU晶片數量的關係。意料之中的是，ResNet-50取得的加速比接近最理想的擴充套件性。這是因為用ImageNet去訓練ResNet-50的計算通訊的比值(處理一張圖片的浮點運算次數與梯度中引數的個數之間的比值)非常高。

圖5：ResNet-50訓練加速比與晶片數目的關係[1]

為了進一步研究Amdahl定律在大系統上的瓶頸，我們在圖6中展示了ImageNet/ResNet-50訓練中每步所需要的計算時間以及每步所需要的all-reduce通訊時間。在圖6中，藍色區域顯示的是計算時間，紅色區域顯示的通訊時間。這兩個區域加起來就是每步訓練在裝置上需要的總時間。

圖6：ResNet-50訓練計算時間和all-reduce通訊時間與晶片數目的關係[1]

我們在圖7中展示了BERT訓練的加速比與TPU晶片數量的關係。從16晶片到4096晶片，BERT訓練展示了極高的可擴充套件性。我們在圖8中展示BERT訓練中每步的計算時間和每步的all-reduce通訊時間。藍色區域顯示的是計算時間，紅色區域顯示的通訊時間。這兩個區域加起來就是每步訓練在裝置上需要的總時間。

圖7：BERT訓練加速比與晶片數目的關係[1]

圖8：BERT訓練計算時間和all-reduce通訊時間與晶片數目的關係[1]

由於SSD，MaskR-CNN，以及Transformer三個模型的訓練速度在有限的batch size的情況下依賴於模型並行。所以，我們在圖9中展示了在一個主機上，模型並行產生的加速比與TPU核數之間的關係。

圖9：基於TensorFlow實現，在8個TPU核上(相當於一個主機)，模型並行產生的加速比與TPU核數之間的關係[1]

當前，用於建造人工智慧超級計算機的晶片主要包括第三代TPU晶片，英偉達Volta V100 GPU晶片，以及英偉達Ampere A100GPU晶片。人們分別用這三種晶片建造了人工智慧超級計算機。我們在圖10中比較了不同晶片建造的超級計算機對應的MLPerf基準模型訓練的端到端時間。為了給大家一個公平的比較，我們在圖11中展示了不同晶片建造的超級計算機對應的加速比。這裡的加速比指的是相對於自身16個晶片的加速比，反應的是可擴充套件性。

圖10：不同晶片建造的超級計算機對應的MLPerf基準模型訓練的端到端時間[1]

圖11：不同晶片建造的超級計算機對應的加速比。這裡的加速比指的是相對於自身16個晶片的加速比，反應的是可擴充套件性[1]

作者簡介

尤洋，壁仞科技顧問，是新加坡國立大學計算機系的助理教授(tenure-track)。他從加州大學伯克利分校計算機係獲得了博士學位。他的導師是美國科學院院士，工程院院士，ACM/IEEEfellow，伯克利前計算機系主任以及EECS學院院長James Demmel教授。尤洋的研究興趣包括高效能計算，並行演算法，以及機器學習。他當前的研究重點是大規模深度學習訓練演算法的分散式最佳化。他曾創造ImageNet以及BERT訓練速度的世界紀錄，並被ScienceDaily，TheNextWeb，i-programmer等幾十家媒體廣泛報道。他設計的演算法被廣泛應用於谷歌，微軟，英特爾，英偉達等科技巨頭。尤洋近三年以第一作者身份在NIPS，ICLR，Supercomputing，IPDPS，ICS等國際重要會議或期刊上發表論文十餘篇。他曾以第一作者身份獲得了國際並行與分散式處理大會(IPDPS)的最佳論文(0.8%獲獎率)和國際並行處理大會(ICPP)的最佳論文(0.3%獲獎率)。尤洋曾獲清華大學優秀畢業生，北京市優秀畢業生，國家獎學金，以及當時清華大學計算機係數額最高的西貝爾獎學金。他還在2017年獲得美國計算機協會(ACM)官網上唯一頒給在讀博士生的ACM-IEEECS George Michael Memorial HPC Fellowships。他也獲得了頒發給伯克利優秀畢業生的LotfiA. Zadeh Prize。尤洋還被伯克利提名為ACM Doctoral Dissertation Award候選人(81名UCBerkeley EECS2020博士畢業生中選2人)。更多資訊請檢視他的研究小組主頁(https://ai.comp.nus.edu.sg/)。

參考文獻

[1]Sameer Kumar James Bradbury Cliff Young Yu Emma Wang AnselmLevskaya Blake Hechtman Dehao Chen HyoukJoong Lee Mehmet Deveci Naveen Kumar Pankaj Kanwar Shibo Wang Skye Wanderman-Milne SteveLacy Tao Wang Tayo Oguntebi Yazhou Zu Yuanzhong Xu Andy SwingEXPLORING THE LIMITS OF CONCURRENCY IN ML TRAINING ON GOOGLE TPUS

[2]Lepikhin, D., Lee, H., Xu, Y., Chen, D., Firat, O., Huang, Y.,Krikun,M., Shazeer, N., and Chen, Z. Gshard: Scaling giant models withconditional computation and automatic sharding. arXiv preprintarXiv:2006.16668, 2020.

[3]Shazeer, N., Cheng, Y., Parmar, N., Tran, D., Vaswani,A.,Koanantakool, P., Hawkins, P., Lee, H., Hong, M., Young, C., etal. Mesh-TensorFlow: Deep learning for supercomputers. In Advances inNeural Information Processing Systems, pp. 10414–10423, 2018.

[4]Xu, Y., Lee, H., Chen, D., Choi, H., Hechtman, B., and Wang, S.Automatic cross-replica sharding of weight update in data-paralleltraining. arXiv preprint arXiv:2004.13336, 2020.

[5]Mattson, P., Cheng, C., Coleman, C., Diamos, G., Micikevicius, P.,Patterson, D., Tang, H., Wei, G.-Y., Bailis, P., Bittorf, V., et al.MLPerf training benchmark. arXiv preprint arXiv:1910.01500, 2019.

[6]Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K. Bert:Pre-training of deep bidirectional transformers for languageunderstanding. arXiv preprint arXiv:1810.04805, 2018.

[7]You, Y., Li, J., Reddi, S., Hseu, J., Kumar, S., Bhojanapalli, S.,Song, X. d., Demmel, J., Keutzer, K., and Hsieh, C.-J. Large batchoptimization for deep learning: Training BERT in 76 minutes. arXivpreprint arXiv:1904.00962, 2019.

[8]Golovin, D., Solnik, B., Moitra, S., Kochanski, G., Karro, J.,andSculley, D. Google Vizier: A service for black-box optimization. InProceedings of the 23rd ACM SIGKDD international conference onknowledge discovery and data mining, pp. 1487–1495, 2017.