到底什么是All-Reduce、All-to-All？

我們?cè)诳碅IGC大模型、算力集群相關(guān)文獻(xiàn)的時(shí)候，經(jīng)常會(huì)看到“ All-Reduce ”、“ All-to-All ”這樣的詞組。

大家知道它們是什么意思嗎？

故事還是要從英偉達(dá)說起。

2014年，為了取代傳統(tǒng)的PCIe協(xié)議，英偉達(dá)推出了全新的NVLINK技術(shù)，用于GPU和GPU之間的點(diǎn)對(duì)點(diǎn)高速互聯(lián)。

后來，NVLINK技術(shù)不斷迭代。 2018年，為了實(shí)現(xiàn)8顆GPU之間的互連，英偉達(dá)又發(fā)布了NVSwitch 1.0。

在搗鼓GPU互連技術(shù)的過程中，英偉達(dá)還發(fā)明了一個(gè)名叫 NCCL（NVIDIA Collective Communications Library）的集合通信庫(kù)。

NCCL支持多種編程語(yǔ)言和網(wǎng)絡(luò)，在算法層面提供了豐富的外部API，可以大幅提升通信網(wǎng)絡(luò)性能，也可以讓GPU之間的通信設(shè)計(jì)變得更簡(jiǎn)單。

在NCCL中，定義了兩種通信模式：

1、點(diǎn)對(duì)點(diǎn)通信（Point to point communication，P2P）

2、集合通信（Collective communication，CC）

點(diǎn)對(duì)點(diǎn)大家一看就懂，就是兩個(gè)點(diǎn)之間進(jìn)行通信。一個(gè)是Sender，一個(gè)是Receiver。

什么是集合通信呢？是指一組（多個(gè)）節(jié)點(diǎn)內(nèi)進(jìn)行通信。在我們傳統(tǒng)通信里，就是點(diǎn)到多點(diǎn)，多點(diǎn)到多點(diǎn)，涉及到組網(wǎng)（網(wǎng)狀、星狀、環(huán)狀、mesh等）那種。

NCCL還定義了一些計(jì)算節(jié)點(diǎn)之間數(shù)據(jù)交換的基本操作模式，并將其命名為——“通信原語(yǔ)（也有寫作“通信元語(yǔ)”）。

這些通信原語(yǔ)包括：Broadcast、Scatter、Gather、All-Gather、Reduce、All-Reduce、Reduce-Scatter、All-to-All等。

沒錯(cuò)，All-Reduce和All-to-All，就是其中之二。

這些通信原語(yǔ)是構(gòu)建復(fù)雜通信行為的“原子操作”。現(xiàn)在所有復(fù)雜的AI算力集群，內(nèi)部通信都是基于這些通信原語(yǔ)。它們極大地提升了并行計(jì)算的效率和便利性。

接下來，我們就逐個(gè)解釋一下，這些通信原語(yǔ)的意思。

▉Broadcast（1對(duì)多的廣播）

這個(gè)最簡(jiǎn)單。當(dāng)主節(jié)點(diǎn)執(zhí)行Broadcast操作時(shí)，數(shù)據(jù)會(huì)從主節(jié)點(diǎn)發(fā)送至其他所有節(jié)點(diǎn)。

Broadcast是一個(gè)典型的分發(fā)、散播行為。在分布式機(jī)器學(xué)習(xí)中，Broadcast常用于網(wǎng)絡(luò)參數(shù)的初始化。

▉Scatter（1對(duì)多的發(fā)散）

Scatter也是一種分發(fā)、散播行為。它也是 將主節(jié)點(diǎn)的數(shù)據(jù)發(fā)送至其他所有節(jié)點(diǎn)。只不過， Broadcast發(fā)送的是完整數(shù)據(jù)，而 Scatter是將數(shù)據(jù)進(jìn)行切割后，再分發(fā)，就像分生日蛋糕。

▉Gather（多對(duì)1的收集）

Gather，是將多個(gè)sender（發(fā)送節(jié)點(diǎn)）上的數(shù)據(jù)收集到單個(gè)節(jié)點(diǎn)上，可以理解為反向的Scatter。

▉All-Gather（多對(duì)多的收集）

Gather是多個(gè)到一個(gè)，All-Gather是多個(gè)到多個(gè)。

All-Gather是將多個(gè)sender（發(fā)送節(jié)點(diǎn)）上的數(shù)據(jù)收集到多個(gè)節(jié)點(diǎn)上。它相當(dāng)于多個(gè)Gather操作。或者說，是一個(gè)Gather操作之后，跟著一個(gè)Broadcast操作。

▉Reduce（多對(duì)1的規(guī)約）

Reduce的英文意思是“減少、降低”。在集合通信里，它表示“規(guī)約”運(yùn)算，是一系列簡(jiǎn)單運(yùn)算操作（包括：SUM、MIN、MAX、PROD、LOR等）的統(tǒng)稱。

經(jīng)常用Excel表格的童鞋，對(duì)這些簡(jiǎn)單運(yùn)算應(yīng)該不陌生。例如SUM，就是求和。MIN，就是找出最小值。

其實(shí)說白了，Reduce就是：輸入多個(gè)數(shù)，執(zhí)行操作后，得到更少的數(shù)（例如1個(gè)數(shù)）。

下面這個(gè)，就是以ReduceSum（求和規(guī)約）為例：

▉All-Reduce（多對(duì)多的規(guī)約）

All-Reduce，這個(gè)是我們?cè)谖恼麻_頭提到的，AI領(lǐng)域非常常見的一個(gè)詞組。

在大模型訓(xùn)練中，經(jīng)常會(huì)用到數(shù)據(jù)并行（DP）這個(gè)并行方式。里面就有AIl Reduce這個(gè)關(guān)鍵操作。

我們以All Reduce Sum（求和）為例：

首先，對(duì)所有節(jié)點(diǎn)進(jìn)行數(shù)據(jù)收集。然后，對(duì)數(shù)據(jù)進(jìn)行求和。再然后，把結(jié)果重新發(fā)回給所有節(jié)點(diǎn)。

在大模型訓(xùn)練中，Server GPU節(jié)點(diǎn)收集的數(shù)據(jù)，就是各個(gè)Worker GPU節(jié)點(diǎn)計(jì)算得出的“梯度”。求和之后再發(fā)回的過程，是“更新梯度”。看不懂沒關(guān)系，以后小棗君會(huì)再介紹。

▉Reduce-Scatter（組合的規(guī)約與發(fā)散）

Reduce-Scatter稍微有點(diǎn)復(fù)雜、燒腦。

它是先歸約（Reduce），再分散（Scatter）。具體來說：

首先，在所有參與計(jì)算的GPU節(jié)點(diǎn)上，對(duì)位于相同位置或索引的數(shù)據(jù)塊執(zhí)行指定的規(guī)約運(yùn)算（例如求和SUM）。

接著，將規(guī)約后的完整結(jié)果按維度切分，并將不同的數(shù)據(jù)塊分發(fā)給各個(gè)節(jié)點(diǎn)。最終，每個(gè)節(jié)點(diǎn)只得到整個(gè)規(guī)約結(jié)果的一部分，而不是全部。

簡(jiǎn)單來說，它先對(duì)所有數(shù)據(jù)進(jìn)行“匯總計(jì)算”，然后再將計(jì)算好的結(jié)果“分散下發(fā)”。

▉All-to-All（多對(duì)多的全互連）

AIl-to-AII也是AI領(lǐng)域出現(xiàn)頻率很高的一個(gè)詞組。它是全交換操作，可以讓每個(gè)節(jié)點(diǎn)都獲取其他節(jié)點(diǎn)的值。

在使用All-to-All時(shí)，每一個(gè)節(jié)點(diǎn)都會(huì)向任意一個(gè)節(jié)點(diǎn)發(fā)送消息，每一個(gè)節(jié)點(diǎn)也都會(huì)接收到任意一個(gè)節(jié)點(diǎn)的消息。每個(gè)節(jié)點(diǎn)的接收緩沖區(qū)和發(fā)送緩沖區(qū)都是一個(gè)分為若干個(gè)數(shù)據(jù)塊的數(shù)組。

All-to-All的具體操作是：將節(jié)點(diǎn)i的發(fā)送緩沖區(qū)中的第j塊數(shù)據(jù)發(fā)送給節(jié)點(diǎn)j。節(jié)點(diǎn)j將接收到的來自節(jié)點(diǎn)i的數(shù)據(jù)塊，放在自身接收緩沖區(qū)的第i塊位置。

All-to-All與All-Gather相比較，區(qū)別在于：All-Gather操作中，不同節(jié)點(diǎn)向某一節(jié)點(diǎn)收集到的數(shù)據(jù)是完全相同的。而在All-to-All中，不同的節(jié)點(diǎn)向某一節(jié)點(diǎn)收集到的數(shù)據(jù)是不同的。在每個(gè)節(jié)點(diǎn)的發(fā)送緩沖區(qū)中，為每個(gè)節(jié)點(diǎn)都單獨(dú)準(zhǔn)備了一塊數(shù)據(jù)。

上面這個(gè)圖，大家如果學(xué)過工程數(shù)學(xué)的話，就會(huì)發(fā)現(xiàn)，它就是一個(gè)矩陣倒置。或者說，是Excel里的行列倒轉(zhuǎn)。

All-to-All的核心目標(biāo)是重分布。它不進(jìn)行聚合運(yùn)算，而是專注于在不同節(jié)點(diǎn)間重新分布數(shù)據(jù)塊。

以后小棗君會(huì)給大家介紹， All-to-All 操作在大模型訓(xùn)練中的混合并行策略里至關(guān)重要。例如， 當(dāng)需要從數(shù)據(jù)并行組切換到模型并行組時(shí)，All-to-All可以高效地重組數(shù)據(jù)。

▉Ring-base collective（基于環(huán)的集合）

最后還要提一個(gè)有趣的結(jié)構(gòu)——環(huán)（Ring）。

Ring-base collective是將所有的通信節(jié)點(diǎn)通過首位相連形成一個(gè)單向環(huán)，數(shù)據(jù)在環(huán)上依次傳輸。

傳輸方式有兩種，一種是一次性傳輸全部，還有一種，是對(duì)數(shù)據(jù)進(jìn)行切割，然后分別發(fā)送。

All-Reduce里有一種Ring All-Reduce（環(huán)形全規(guī)約）算法。它是通過組合Reduce-Scatter和All-Gather兩個(gè)操作來實(shí)現(xiàn)的。

Ring All-Reduce算法分為兩個(gè)階段：

第一階段，將N個(gè)worker分布在一個(gè)環(huán)上，并且把每個(gè)worker的數(shù)據(jù)分成N份。

對(duì)于第k個(gè)worker，這個(gè)worker會(huì)把第k份數(shù)據(jù)發(fā)給下一個(gè)worker，同時(shí)從前一個(gè)worker收到第k-1份數(shù)據(jù)。

然后，第k個(gè)worker會(huì)把收到的第k-1份數(shù)據(jù)和自己的第k-1份數(shù)據(jù)整合，再將整合的數(shù)據(jù)發(fā)送給下一個(gè)worker。

以此循環(huán)N次之后，每一個(gè)worker都會(huì)包含最終整合結(jié)果的一份。

第二階段，每個(gè)worker將整合好的部分發(fā)送給下一個(gè)worker。worker在收到數(shù)據(jù)之后，更新自身數(shù)據(jù)對(duì)應(yīng)的部分即可。

很顯然，這種環(huán)形算法可以解決傳統(tǒng)All-Reduce中Server節(jié)點(diǎn)的能力瓶頸問題。

▉最后的話

好啦，以上就是常見通信原語(yǔ)的具體工作原理。

AI大模型訓(xùn)練推理任務(wù)，是由海量的GPU共同完成的。而這些GPU之間的通信，就是基于上面這些通信原語(yǔ)模型。

下一期，小棗君會(huì)詳細(xì)介紹一下大模型訓(xùn)推任務(wù)中的并行計(jì)算方式，以及這些通信原語(yǔ)究竟是如何運(yùn)用于不同的并行計(jì)算方式中。

敬請(qǐng)關(guān)注！