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

我們在看AIGC大模型、算力集群相關文獻的時候，經常會看到“ All-Reduce ”、“ All-to-All ”這樣的詞組。

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

故事還是要從英偉達說起。

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

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

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

NCCL支持多種編程語言和網絡，在算法層面提供了豐富的外部API，可以大幅提升通信網絡性能，也可以讓GPU之間的通信設計變得更簡單。

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

1、點對點通信（Point to point communication，P2P）

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

點對點大家一看就懂，就是兩個點之間進行通信。一個是Sender，一個是Receiver。

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

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

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

沒錯，All-Reduce和All-to-All，就是其中之二。

這些通信原語是構建復雜通信行為的“原子操作”。現在所有復雜的AI算力集群，內部通信都是基于這些通信原語。它們極大地提升了并行計算的效率和便利性。

接下來，我們就逐個解釋一下，這些通信原語的意思。

▉Broadcast（1對多的廣播）

這個最簡單。當主節(jié)點執(zhí)行Broadcast操作時，數據會從主節(jié)點發(fā)送至其他所有節(jié)點。

Broadcast是一個典型的分發(fā)、散播行為。在分布式機器學習中，Broadcast常用于網絡參數的初始化。

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

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

▉Gather（多對1的收集）

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

▉All-Gather（多對多的收集）

Gather是多個到一個，All-Gather是多個到多個。

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

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

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

經常用Excel表格的童鞋，對這些簡單運算應該不陌生。例如SUM，就是求和。MIN，就是找出最小值。

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

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

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

All-Reduce，這個是我們在文章開頭提到的，AI領域非常常見的一個詞組。

在大模型訓練中，經常會用到數據并行（DP）這個并行方式。里面就有AIl Reduce這個關鍵操作。

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

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

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

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

Reduce-Scatter稍微有點復雜、燒腦。

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

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

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

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

▉All-to-All（多對多的全互連）

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

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

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

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

上面這個圖，大家如果學過工程數學的話，就會發(fā)現，它就是一個矩陣倒置。或者說，是Excel里的行列倒轉。

All-to-All的核心目標是重分布。它不進行聚合運算，而是專注于在不同節(jié)點間重新分布數據塊。

以后小棗君會給大家介紹， All-to-All 操作在大模型訓練中的混合并行策略里至關重要。例如， 當需要從數據并行組切換到模型并行組時，All-to-All可以高效地重組數據。

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

最后還要提一個有趣的結構——環(huán)（Ring）。

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

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

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

Ring All-Reduce算法分為兩個階段：

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

對于第k個worker，這個worker會把第k份數據發(fā)給下一個worker，同時從前一個worker收到第k-1份數據。

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

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

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

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

▉最后的話

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

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

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

敬請關注！