作者簡介

周棟，北京郵電大學(xué) 未來網(wǎng)絡(luò)理論與應(yīng)用實(shí)驗(yàn)室博士生，研究方向?yàn)槲磥砭W(wǎng)絡(luò)體系架構(gòu)、可編程網(wǎng)絡(luò)。郵箱: hardworkzd@163.com

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

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

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

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

?

什么是分布式機(jī)器學(xué)習(xí)？

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

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

| 數(shù)據(jù)并行（左）和模型并行（右）示意圖

數(shù)據(jù)并行是指訓(xùn)練集群中的機(jī)器使用不同的數(shù)據(jù)進(jìn)行并行計(jì)算，數(shù)據(jù)的劃分可以是直接劃分也可以是隨機(jī)采樣。

模型并行則是將不同的模型參數(shù)劃分到不同的機(jī)器上，根據(jù)模型結(jié)構(gòu)的依賴關(guān)系進(jìn)行計(jì)算和通信。

?

在分布式機(jī)器學(xué)習(xí)中，網(wǎng)絡(luò)能做什么？

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

網(wǎng)絡(luò)通信在分布式機(jī)器學(xué)習(xí)中的作用具體涉及哪些方面呢？又該如何去優(yōu)化？

首先是邏輯通信拓?fù)?/strong>，在AI模型訓(xùn)練的每輪迭代中，各個(gè)計(jì)算節(jié)點(diǎn)完成自己的計(jì)算任務(wù)后，需要通過某種邏輯通信拓?fù)溥M(jìn)行結(jié)果的聚合，完成一輪次的學(xué)習(xí)。根據(jù)計(jì)算節(jié)點(diǎn)的功能邏輯是否相似，可以劃分為兩類：

除了通信拓?fù)洌?strong>通信步調(diào)也是制約分布式機(jī)器學(xué)習(xí)整體性能的關(guān)鍵。通信步調(diào)關(guān)乎每輪迭代后模型全局更新的方式，分為異步和同步兩種基本方式：

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

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

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

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

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

| All-Reduce過程示意圖

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

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

|?網(wǎng)內(nèi)聚合原理示意圖

?

結(jié)語

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

編輯：黃飛

?