作者簡介

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

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

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

表 GPT系列模型參數和數據規模

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

?

什么是分布式機器學習？

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

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

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

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

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

?

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

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

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

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

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

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

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

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

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

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

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

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

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

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

| All-Reduce過程示意圖

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

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

|?網內聚合原理示意圖

?

結語

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

編輯：黃飛

?