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

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

第1節：介紹

以自然語言處理為代表的深度學習革命正處在類似於美蘇60年代“太空競賽”的境地，導致模型的大小正以爆炸性的速度增長。顯然，大模型帶來了極大的好處，比如OpenAI的超大模型GPT-3已經可以在輕微的人類引導下去生成一些有可讀性的散文，效果明顯是之前的模型無法達到的。同時OpenAI研究院也觀察到，世界最佳模型所消耗的計算資源每3.5個月就要翻一番。這比當時的摩爾定律誇張多了。之前摩爾定律表明每18月單位面積的電晶體數量翻一番。由摩爾定律導致的安迪-比爾定律直接地預示了計算機產業幾十年的輝煌。

所以，為了滿足這些大模型的快速訓練需求，我們需要更大的超級計算機。這場深度學習革命引發的超算革命給GPU帶來了鉅額的投資：過去幾年，英偉達公司的營收見證過50%的年增長率。英偉達公司的股票也從2012年11月的11美元增長到了2020年11月582美元(53倍)。

2015年，谷歌的TPU神經網路專用加速器橫空出世。谷歌宣稱，每個TPU晶片比之前的神經網路晶片在功耗比，低延遲效率，以及最高效能上提升了10倍(Jouppi等人的論文[2],2017)。在之後的兩年，谷歌基於二代TPU的256晶片超計算機就展示出了完美的並行效率(Jouppi等人的論文[3],2020)。之後谷歌基於3代TPU將超級計算機的規模做到了1024晶片的規模(Kumar等人的論文[4],2019)。英偉達和其它GPU供應商也在建造類似規模的AI超級計算機。微軟和OpenAI甚至要打造一個一萬個GPU的超級計算機(Langston的文章[5],2020)。這場類似於“美蘇星球大戰”的比拼將會用精準的模型去試圖逼近通用人工智慧。這一點是否能實現還有待觀察。但是，無疑的是，這場比拼中的各個巨頭對硬體投入是雄心勃勃的。

美蘇太空競賽可以用類似於“繞地飛行”或“登陸月球”等一些里程碑事件去標記成就。然而，這次AI超算競爭卻在開始的時候沒有具體衡量指標。所以，谷歌、英偉達、英特爾、阿里巴巴、華為等工業界巨頭意識到這個問題後提出了MLPerf標準(MLPerf.org)。特別地，MLPerf訓練標準吸引了眾多有高效能計算能力的實體去用超級計算機以最短的時間完成神經網路的訓練(Mattson等人的論文[6],2019)。

事實證明，在過去的兩年，MLPerf標準對深度學習社群產生了積極而深遠的影響。因為重新整理MLPerf標準的努力和創新產生了新的系統最佳化技術，更實用的程式碼庫，更高效的編譯器，以及最合適的應用層程式碼。

常規的TPU超級計算機只有1024個晶片。我們在上次的MLPerfv0.6版本比賽中只用了常規的TPU超級計算機。為了探索MLPerf模型平行計算的極限，我們用第三代TPU晶片組裝了一個4096晶片的多組超級計算機(詳見圖1)。兩個TPU超級計算機在網格結構的X方向上被光纖連線起來(詳見圖2)。這些組間的光纖連結要比普通的組內光纖連結長。我們用了128×32的二維網狀結構將所有的晶片連線起來用於訓練MLPerf的模型。因為第三代TPU晶片的路由表只能容納1024個地址，所以我們用了一種稀疏的路由策略。在這種策略下，每個晶片只能看到與它同一行或同一列的晶片。這個策略足以使all-reduce操作達到峰值的通訊傳輸效率。

圖1：常規TPU超級計算與多組TPU超級計算機[1]

圖2：4組TPU超級計算機的配置，跨組連結在橫向連線了兩個TPU超級計算機[1]

第2節：多種程式設計框架

儘管TPU的主要前端程式設計框架一直是TensorFlow(Abadi等人的論文[7],2016)，TPU硬體和XLA編譯器其實是通用的工具：它們是可以支援其它的程式設計框架。因此，在這篇論文中，我們選擇給TensorFlow和JAX都做評測。JAX(Frostig等人的論文[8],2018)是一種面向研究的基於XLA的數值計算程式設計框架。這兩種程式設計模型都需要額外的軟體工程在多組超級計算機上獲得高效的擴充套件性。但是它們最終可以獲得相似的評測結果。

如圖3所示，由於TensorFlow和JAX在架構上的不同，它們在大規模運算上也有效能差異。首先，它們有不同分階段方法。TensorFlow在Python中嵌入了一種表示式的動態即時程式語言，然後用XLA去即時編譯圖的子集(TensorFlow的圖可以分佈在TPU加速器上和CPU上)。相比之下，JAX少了一個階段：它是一個在Python中嵌入了即時編譯的XLA程式的程式設計框架。其中，即時編譯的XLA程式用於在加速器上的靜態編譯效能以及加速器網路上的並行性。JAX用於動態以及加速計算。所以，TensorFlow比JAX多了一個編譯階段，我們用多執行緒去加速了這個多出來的編譯階段。同時，JAX也需要我們更加小心地去處理Python的瓶頸。比如，把一些類似於資料讀取的干擾性任務移出主執行緒。

圖3：TensorFlow和JAX程式設計模型在第三代TPU上使用的示意圖[1]

其次，TensorFlow和JAX可以支援不同的分散式程式設計模型。JAX採用了一種多客戶端的方式去做分散式程式設計，使得超級計算機中每個主機分別有一份同樣的JAX程式碼(包括Python直譯器)。JAX程式只在兩個地方進行通訊：(1)在初始階段去設定TPU的網格連線。(2)在模型訓練過程中在網路上進行XLA編譯好的all-reduce操作。與之相反的是，TensorFlow採用單客戶端的方式在TPU上進行程式設計。這種方式會給一個Python程序全域性視角並允許它控制整個分散式系統。這個Python程序執行在超級計算機的一個主機上。這個主機將TensorFlow的圖均分後透過RPC在網路上發給剩下的主機執行。

TensorFlow和JAX在實用性和效能特點上也有不同之處。儘管TensorFlow的單客戶端分散式系統能夠讓使用者程式碼直接控制整體的執行負載，JAX的多客戶端方法使得程式碼可以直接控制各個獨立的計算單元。JAX在各個主機上獨立地呼叫XLA編譯器，這依賴於確定性編譯技術去避免不同主機程式的不相容性。然而，TensorFlow只編譯一次並且把二進位制檔案分給所有主機節點。TensorFlow的這種多裝置圖其實能引起Amdahl定律瓶頸。原因是因為客服端程序在圖構造和最佳化時間上的開銷是跟節點數成正比的。然後，除了TPU拓撲網格連結的初始化之外，JAX的設定時間是不隨著節點數增加而增加的。所以，JAX在這一點上看似有更好的可擴充套件性。

作者簡介

尤洋，壁仞科技顧問，是新加坡國立大學計算機系的助理教授(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名UC Berkeley EECS2020博士畢業生中選2人)。更多資訊請檢視他的研究小組主頁(https://ai.comp.nus.edu.sg/)。

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

參考文獻

[1] Sameer Kumar James Bradbury Cliff Young Yu Emma Wang Anselm Levskaya Blake Hechtman Dehao Chen HyoukJoong Lee Mehmet Deveci Naveen Kumar Pankaj Kanwar Shibo Wang Skye Wanderman-Milne Steve Lacy Tao Wang Tayo OguntebiYazhou Zu Yuanzhong Xu Andy Swing EXPLORING THE LIMITS OF CONCURRENCY IN ML TRAINING ON GOOGLE TPUS

[2] Jouppi, N. P., Young, C., Patil, N., Patterson, D., Agrawal, G., Bajwa, R., Bates, S., Bhatia, S., Boden, N., Borchers, A., et al. In-datacenter performance analysis of a tensor processing unit. In Proceedings of the 44th Annual International Symposium on Computer Architecture, pp. 1–12, 2017.

[3] Jouppi, N. P., Yoon, D. H., Kurian, G., Li, S., Patil, N., Laudon, J., Young, C., and Patterson, D. A domainspecific supercomputer for training deep neural networks. Communications of the ACM, 63(7):67–78, 2020.

[4] Kumar, S., Bitorff, V., Chen, D., Chou, C., Hechtman, B., Lee, H., Kumar, N., Mattson, P., Wang, S., Wang, T., et al. Scale MLPerf-0.6 models on Google TPU-v3 pods. arXiv preprint arXiv:1909.09756, 2019.

[5] Langston, J. Microsoft announces new supercomputer, lays out vision for future AI work. MicroSoft Blog, 2020.

[6] 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.

[7] Abadi, M., Barham, P., Chen, J., Chen, Z., Davis, A., Dean, J., Devin, M., Ghemawat, S., Irving, G., Isard, M., et al. TensorFlow: A system for large-scale machine learning. In 12th {USENIX} symposium on operating systems design and implementation ({OSDI} 16), pp. 265–283, 2016.

[8] Frostig, R., Johnson, M. J., and Leary, C. Compiling machine learning programs via high-level tracing. Systems for Machine Learning, 2018.

往期推薦

1、圖結構: 分析神經網路的新視角-1

2、圖結構: 分析神經網路的新視角-2

3、在壁仞科技研究院，我們思考的問題