作者簡介

周棟，北京郵電大學 未來網絡理論與應用實驗室博士生，研究方向為未來網絡體系架構、可編程網絡。郵箱: hardworkzd@163.com

最近ChatGPT可以說是火爆各大圈子，不同于以往的對話AI，ChatGPT能夠更好地理解人類的意圖——生成的內容是根據歷史對話信息生成的，并不只是局限于當前問題；而且能夠完成的事情更加豐富——可以幫助你寫代碼、想文案、改論文、做PPT等等。

如此之大的AI表現力的突破究竟來自哪里？當然，模型算法的不斷優(yōu)化改進是重中之重，但我們也很難一時半會從浩如煙海的算法技術細節(jié)中整理出思路。我們可以從另一個更為直觀的角度來理解這種進步，那就是模型參數的規(guī)模。ChatGPT源自于OpenAI推出的GPT系列模型，僅GPT-1就有上億的參數量，到ChatGPT所使用的GPT-3.5更是膨脹到夸張的萬億參數。這種大模型其實并不是獨一份，騰訊的混元模型參數規(guī)模也是達到萬億級別 [1]，阿里達摩院的M6模型參數更是突破10萬億[2]。

表 GPT系列模型參數和數據規(guī)模

神經網絡越深，參數越多，模型的能力就越強，就能解決更復雜的問題?；诖?，AI大模型層出不窮，但是同時也伴隨著一個問題——大模型+大數據該如何進行訓練？不論是模型還是數據，規(guī)模大到一定程度就無法在單機上存儲，而且單機訓練的時間開銷也是難以接受的，這時就需要在一個集群上進行訓練，這也是大規(guī)模AI訓練的必經之路——分布式機器學習。

?

什么是分布式機器學習？

分布式機器學習使用計算機集群對訓練任務進行并行計算，從而擺脫單機的存儲和計算上的限制，加速訓練任務的完成。從訓練流程上來看，分布式與單機都是迭代式的，但不同的是，分布式在每輪迭代中需要進行任務的劃分、節(jié)點并行計算、結果的聚合等步驟，然后繼續(xù)下一輪迭代。

大規(guī)模AI訓練的主要難題就在于“大”，參數規(guī)模大，數據集大，無法單機存儲，計算時間太長。因此，將任務按照數據或者模型進行劃分，并行執(zhí)行以加速整體訓練，成為兩種基礎的任務劃分方式或者說并行模式。

| 數據并行（左）和模型并行（右）示意圖

數據并行是指訓練集群中的機器使用不同的數據進行并行計算，數據的劃分可以是直接劃分也可以是隨機采樣。

模型并行則是將不同的模型參數劃分到不同的機器上，根據模型結構的依賴關系進行計算和通信。

?

在分布式機器學習中，網絡能做什么？

相比單機的學習，分布式機器學習涉及多個計算節(jié)點的協(xié)作，網絡通信既是差別所在也是最為關鍵。值得一提的是，隨著AI訓練中最開始使用的通用計算CPU被更為專業(yè)快速的GPU、TPU取代，端側的計算能力越來越強，這進一步加劇了網絡通信的瓶頸問題，更隨著大規(guī)模AI的爆火，網絡通信方面的優(yōu)化已經刻不容緩。

網絡通信在分布式機器學習中的作用具體涉及哪些方面呢？又該如何去優(yōu)化？

首先是邏輯通信拓撲，在AI模型訓練的每輪迭代中，各個計算節(jié)點完成自己的計算任務后，需要通過某種邏輯通信拓撲進行結果的聚合，完成一輪次的學習。根據計算節(jié)點的功能邏輯是否相似，可以劃分為兩類：

去中心化：在去中心化的拓撲中，各個計算節(jié)點的工作邏輯大致相同，完成計算后會通過某種分布式機制或協(xié)議完成節(jié)點間計算結果的同步和聚合，拓撲的形狀也是各異的，有環(huán)狀、樹狀、星形等等。典型的例子為Ring Allreduce架構。 ?

中心化：在中心化的拓撲中，計算節(jié)點存在功能角色的差異。一部分是工作節(jié)點，主要負責計算；另一部分是中心節(jié)點（或者服務器節(jié)點），負責模型全局更新聚合、下發(fā)最新參數等管理職責。拓撲一般是中心節(jié)點和工作節(jié)點相連，工作節(jié)點間不相連而互不干擾。典型例子為參數服務器架構（Parameter Server）。

| 去中心化拓撲（左）和中心化拓撲（右）

相較而言，去中心化的拓撲能夠更充分利用帶寬，但存在可擴展性不高的問題。中心化的拓撲擴展性較好，但是隨著規(guī)模上升，中心節(jié)點反而成為通信瓶頸。因此，需要根據模型訓練任務和計算集群的特性對拓撲進行選擇和搭配。

除了通信拓撲，通信步調也是制約分布式機器學習整體性能的關鍵。通信步調關乎每輪迭代后模型全局更新的方式，分為異步和同步兩種基本方式：

同步：每輪迭代中，當所有節(jié)點完成計算后，才能完成一次全局的模型更新。

異步：每輪迭代中，先完成計算的節(jié)點不必等待未完成的節(jié)點，而是上傳模型更新，并直接開始下一輪的迭代。

同步方式的優(yōu)勢在于邏輯簡單、理論性能有保障，但是在節(jié)點性能有較大差異的情況下會導致集群資源閑置，可擴展性不高。異步方式的優(yōu)勢在于能夠更好地利用集群計算資源，擴展性強，但是由于參數不同步會影響模型的收斂速度，理論上性能保障較差。兩種基本方式都有一定的局限性，因此也有一些改進的算法被提出，比如延遲同步并行算法SSP假設一個相對同質化的集群，剛開始采用異步的方式，但是會控制迭代次數的極差不超過預設閾值，如果超過閾值則進行等待或丟棄等處理。

以上是分布式機器學習領域中有關網絡部分的經典問題，而近幾年隨著網絡領域在軟硬件方面的發(fā)展，對于分布式機器學習的優(yōu)化加速，也有一些令人眼前一亮的方向和應用成果。

在減少通信的處理延遲的方面，學術界和產業(yè)界聚焦于高性能通信庫和通信協(xié)議。高性能通信庫的代表是NCCL，它是NVIDIA為了更好地支持多GPU的使用、提高集群通信效率而提出的，包括單機多卡和多機多卡的使用場景。NCCL實現了很多用于GPU之間進行并行協(xié)作計算的基本通信原語（Communication Primitives），并在實現上進行各種優(yōu)化以達到高性能。比如下圖所示的All-Reduce原語，是將多個節(jié)點上的數據進行加和，同步到每個節(jié)點，最終每個節(jié)點的數據都是原來數據的加和。這種操作在分布式機器學習中十分常見：每個節(jié)點的更新也需要總體進行加和平均，然后同步到每個節(jié)點。

| All-Reduce過程示意圖

對于網絡通信協(xié)議，越來越多的理論分析和實踐經驗認識到，主機端網絡通信相關的處理也占了很大一部分時延，于是針對協(xié)議棧處理和協(xié)議本身的優(yōu)化也被應用在分布式機器學習領域。零拷貝、內核旁路、CPU卸載等優(yōu)化思想被廣泛地使用在各種通信技術當中：比如RDMA/RoCE/iWARP可以繞過遠程主機內核直接存取訪問內存，能夠應用于多GPU之間的數據傳輸；協(xié)議棧卸載技術在硬件網卡上實現部分甚至全部的協(xié)議處理過程，解放了CPU資源，并且能夠繞過內核直接與用戶空間的應用交互數據，大大減少了主機側的網絡處理時延。

在減少傳輸的數據量方面，傳統(tǒng)的方法是模型的分解壓縮、參數的低精度量化等與模型和優(yōu)化算法耦合的方法。而近年來一些工作利用新興的可編程網絡設備，在網絡中內置處理函數，處理一些關鍵的通信步驟，能夠從另一個維度減少數據傳輸量，從而降低整體的傳輸時延。比如barefoot、微軟參與提出的SwitchML架構[3]以及清華團隊提出的ATP架構[4]，針對深度神經網絡的分布式訓練，通過在可編程設備上實現模型更新聚合并廣播分發(fā)的步驟，大大減少了整體的數據傳輸量，實現了很好的加速效果。

|?網內聚合原理示意圖

?

結語

當然，在大規(guī)模AI的實際訓練過程中，往往需要綜合考慮多個方面進行系統(tǒng)性的設計，比如在設計通信拓撲時需要考慮優(yōu)化算法的性質、更新方式是異步還是同步，選擇并行模式時要針對模型和數據的特性進行切分等等。以上所列舉的方向和技術方法并沒有面面俱到，只是拋磚引玉。相信隨著大規(guī)模AI的火熱，分布式機器學習技術會得到更進一步的發(fā)展，而網絡作為分布式的關鍵，也將在軟件和硬件等方面迎來進一步的發(fā)展空間。

編輯：黃飛

?