*本文系SDNLAB編譯自Juniper技術(shù)專家兼高級(jí)工程總監(jiān)Sharada Yeluri博客

在本文中，Juniper技術(shù)專家兼高級(jí)工程總監(jiān)Sharada Yeluri深入研究了 AI/ML 訓(xùn)練/推理中使用的集體操作，并討論如何將其中一些功能卸載到網(wǎng)絡(luò)交換機(jī)上，以減少擁塞并提高結(jié)構(gòu)的性能。最后，Sharada Yeluri 以Juniper 的 Trio 架構(gòu)為例，展示了可編程交換機(jī)如何無(wú)縫卸載集體操作。

什么是“集體操作”？

在由眾多通過(guò)互聯(lián)結(jié)構(gòu)相連的處理節(jié)點(diǎn)（如GPU）構(gòu)成的并行/分布式計(jì)算系統(tǒng)中，“集體操作”是指涉及一組處理節(jié)點(diǎn)間通信的一系列操作，用于執(zhí)行協(xié)調(diào)性的任務(wù)。這些任務(wù)可能包括將數(shù)據(jù)從一個(gè)節(jié)點(diǎn)分發(fā)給所有節(jié)點(diǎn)、將所有節(jié)點(diǎn)的數(shù)據(jù)收集到一個(gè)節(jié)點(diǎn)、聚合所有節(jié)點(diǎn)之間的數(shù)據(jù)等。

深度學(xué)習(xí)框架支持庫(kù)，可實(shí)現(xiàn) GPU 組之間的集體通信。Nvidia的集體通信庫(kù)（NCCL）針對(duì)其GPU架構(gòu)高效地實(shí)現(xiàn)了集體操作。當(dāng)一個(gè)模型在一組GPU之間進(jìn)行分區(qū)時(shí)，NCCL負(fù)責(zé)管理它們之間的所有通信。

下面是常用的集體操作：

Reduce：從所有節(jié)點(diǎn)聚合數(shù)據(jù)（求和或平均）并將結(jié)果發(fā)送給其中一個(gè)節(jié)點(diǎn)。

AllReduce：聚合所有節(jié)點(diǎn)的數(shù)據(jù)，并將結(jié)果發(fā)送回所有節(jié)點(diǎn)。

AllReduce

ReduceScatter：聚合所有節(jié)點(diǎn)的數(shù)據(jù)，并將結(jié)果（每個(gè)節(jié)點(diǎn)獲得結(jié)果的唯一子集）分發(fā)到所有節(jié)點(diǎn)。

ReduceScatter

廣播：將數(shù)據(jù)從一個(gè)節(jié)點(diǎn)發(fā)送到組中的所有其他節(jié)點(diǎn)。

廣播

AllGather：收集數(shù)據(jù)的不同部分，并將其分發(fā)給所有節(jié)點(diǎn)。

Scatter：將一個(gè)節(jié)點(diǎn)的不同值分發(fā)到所有節(jié)點(diǎn)。

AlltoAll：將所有節(jié)點(diǎn)的數(shù)據(jù)分發(fā)到所有節(jié)點(diǎn)。

有些集體操作可以通過(guò)使用集體操作的組合來(lái)實(shí)現(xiàn)。例如，AllReduce可以實(shí)現(xiàn)為先執(zhí)行ReduceScatter操作，然后再執(zhí)行AllGather操作。

AI/ML 框架中的集體操作

NCCL 實(shí)現(xiàn)了多種集體操作算法，包括環(huán)形（Ring）、樹(shù)形（Tree）以及雙二叉樹(shù)（Double Binary Trees）等，根據(jù)數(shù)據(jù)大小、涉及的 GPU 數(shù)量和網(wǎng)絡(luò)拓?fù)渥詣?dòng)選擇最為高效的算法。這些算法旨在優(yōu)化集體操作中的數(shù)據(jù)傳輸。

例如，在使用AllReduce集體操作進(jìn)行梯度聚合時(shí)，梯度可以按照環(huán)形模式從一個(gè) GPU 發(fā)送到另一個(gè) GPU，其中每個(gè)GPU都會(huì)將其從上一個(gè)GPU接收到的梯度與其本地計(jì)算出的梯度進(jìn)行聚合，然后再將結(jié)果發(fā)送給下一個(gè)GPU。這個(gè)過(guò)程很慢，因?yàn)樘荻染酆鲜前错樞蛲瓿傻模罱K的結(jié)果也會(huì)按順序在環(huán)形拓?fù)渲袀骰氐剿?GPU。

在AllReduce集體操作中，GPU也可以排列為二叉樹(shù)結(jié)構(gòu)。在這個(gè)結(jié)構(gòu)中，每個(gè)葉節(jié)點(diǎn)將其存儲(chǔ)的所有參數(shù)的梯度發(fā)送到其父節(jié)點(diǎn)，父節(jié)點(diǎn)則會(huì)將接收到的梯度與來(lái)自其兄弟葉節(jié)點(diǎn)的對(duì)應(yīng)梯度進(jìn)行求和。此過(guò)程以遞歸方式繼續(xù)，直到所有梯度在樹(shù)的根節(jié)點(diǎn)處聚合。在根節(jié)點(diǎn)擁有所有梯度的總和后，必須將聚合梯度發(fā)送回該樹(shù)中的所有節(jié)點(diǎn)，以更新其模型參數(shù)的本地副本。根節(jié)點(diǎn)首先將聚合梯度發(fā)送給其子節(jié)點(diǎn)，子節(jié)點(diǎn)再依次傳遞給它們的子節(jié)點(diǎn)，如此遞歸進(jìn)行，直到所有節(jié)點(diǎn)都收到了更新后的梯度。

下圖展示了7個(gè)GPU以二叉樹(shù)的形式排列以進(jìn)行梯度聚合。假設(shè)這些 GPU 是連接到不同葉交換機(jī)和主干交換機(jī)的大型網(wǎng)絡(luò)拓?fù)涞囊徊糠帧Ｔ搱D還顯示了 GPU 之間梯度聚合的流量模式。在網(wǎng)絡(luò)中，這些交換機(jī)是被動(dòng)設(shè)備，負(fù)責(zé)轉(zhuǎn)發(fā)GPU之間的通信結(jié)果。

梯度聚合流量（僅顯示 4 個(gè) GPU）

集體操作的卸載

如果網(wǎng)絡(luò)交換機(jī)可以幫助卸載部分或全部的集體操作，結(jié)果會(huì)怎樣？在這個(gè)例子中，每個(gè)交換機(jī)可以對(duì)從屬于AllReduce集體操作的 GPU 接收到的梯度進(jìn)行部分求和，并將結(jié)果傳遞給下一層級(jí)的交換機(jī)。最終的結(jié)果將廣播給該集體內(nèi)的所有GPU。

上圖展示了網(wǎng)內(nèi)聚合如何加速梯度聚合過(guò)程：通過(guò)減少數(shù)據(jù)跳轉(zhuǎn)次數(shù)（從而降低延遲）和減輕網(wǎng)絡(luò)流量。

這樣做可以減少此集體操作的延遲和網(wǎng)絡(luò)擁塞。它還可以在訓(xùn)練期間卸載 GPU 計(jì)算資源以專注于更重要的任務(wù)。同樣，AllGather和Broadcast這兩種集體操作也能從在網(wǎng)絡(luò)交換機(jī)中卸載任務(wù)中獲益。

Nvidia 在其 InfiniBand 和 NVLink 交換機(jī)中通過(guò) SHARP（可擴(kuò)展層次聚合和歸約協(xié)議）支持此功能。SHARP的主要目標(biāo)是在網(wǎng)絡(luò)內(nèi)部直接卸載并加速?gòu)?fù)雜的集體操作，減少需要在網(wǎng)絡(luò)上傳輸?shù)臄?shù)據(jù)量，從而降低整體的通信時(shí)間。SHARP是一個(gè)專有的協(xié)議，僅與Nvidia的InfiniBand/ NVLink交換機(jī)兼容。它在Nvidia的 AI 堆棧中無(wú)縫運(yùn)作。

深度學(xué)習(xí)框架，如TensorFlow和PyTorch，通常依賴于MPI（消息傳遞接口）或NCCL來(lái)進(jìn)行集體操作。Nvidia 確保這些框架能夠利用SHARP來(lái)獲得網(wǎng)絡(luò)內(nèi)計(jì)算的性能優(yōu)勢(shì)。然而，Nvidia 并未公開(kāi)發(fā)布SHARPv3集體操作前后性能改進(jìn)的具體對(duì)比結(jié)果。不過(guò)，一些關(guān)于早期SHARP版本的文章顯示，使用SHARP后，訓(xùn)練性能得到了17%-20%的提升。

以太網(wǎng)結(jié)構(gòu)和在網(wǎng)集體操作

在使用 ROCEv2 進(jìn)行 GPU間通信的以太網(wǎng)結(jié)構(gòu)，目前尚未有任何針對(duì)在網(wǎng)集體操作 (INC) 的開(kāi)放標(biāo)準(zhǔn)。超以太網(wǎng)聯(lián)盟（UEC）正在開(kāi)發(fā)一種新的傳輸協(xié)議 (UET) 及其配套的INC。一旦標(biāo)準(zhǔn)成熟，這項(xiàng)技術(shù)有可能被交換機(jī)和AI框架采納。然而，從任何標(biāo)準(zhǔn)發(fā)布到其原生部署在網(wǎng)絡(luò)設(shè)備硬件中，通常會(huì)有3-4年的滯后期。并且，要對(duì)硬件進(jìn)行這些功能的優(yōu)化，往往需要經(jīng)歷幾代產(chǎn)品的發(fā)展周期。

在數(shù)據(jù)平面交換機(jī)中執(zhí)行INC需要硬件對(duì)各種操作的支持。

解析新協(xié)議/數(shù)據(jù)包格式，以確定數(shù)據(jù)包是否是集體操作的一部分。

能夠深入查看數(shù)據(jù)包并提取所有有效載荷進(jìn)行處理。

能夠?qū)τ行лd荷進(jìn)行操作，并跨多個(gè)數(shù)據(jù)包累積結(jié)果。這需要支持各種浮點(diǎn)/整數(shù)格式的算術(shù)運(yùn)算，而以太網(wǎng)交換機(jī)通常不具備這些功能。

能夠從內(nèi)存中讀取累積的結(jié)果，創(chuàng)建新的數(shù)據(jù)包并將其發(fā)送出去。

能夠處理?yè)砣榧喜僮髁髁烤S護(hù)QoS，特別是在芯片中發(fā)生多個(gè)并行集合操作時(shí)。

能夠從網(wǎng)絡(luò)或鏈路錯(cuò)誤中恢復(fù)，并在集體操作因網(wǎng)絡(luò)故障、鏈路問(wèn)題或落后的GPU而導(dǎo)致失敗時(shí)，通知終端主機(jī)。

有幾種方法可以實(shí)現(xiàn)上述的集體操作。

在交換機(jī)硬件中直接實(shí)現(xiàn)原生支持始終是最高效的方式。然而，在交換機(jī)制造商考慮添加此類(lèi)功能之前，相關(guān)標(biāo)準(zhǔn)和技術(shù)規(guī)范需要進(jìn)一步成熟和完善。若這些交換機(jī)旨在滿足廣泛的應(yīng)用需求，則額外的芯片面積和成本可能會(huì)成為其在AI/ML應(yīng)用之外推廣的阻礙。

另一種方法是將集體操作卸載到連接到交換機(jī)的協(xié)處理器上。這些協(xié)處理器可以包含 CPU 核心，通過(guò)在其上運(yùn)行的軟件來(lái)實(shí)現(xiàn)集體操作。或者，它們可以是具有集體處理原生功能的FPGA/ASIC。協(xié)處理器通常只處理一部分WAN帶寬，以保持較低的成本。這種做法可以讓以太網(wǎng)交換機(jī)保持輕量化和能效，同時(shí)使數(shù)據(jù)中心能夠選擇性地在某些交換機(jī)上卸載集體操作。

這些交換機(jī)能靈活地解析新報(bào)頭和有效載荷。此外，由于幾乎不需要硬件變更，可編程交換機(jī)能夠迅速實(shí)現(xiàn)新標(biāo)準(zhǔn)，并具備支持多種集體操作協(xié)議的能力。然而，這些交換機(jī)的帶寬往往較低（通常為常規(guī)交換機(jī)的5到10倍），因此對(duì)于構(gòu)建大型網(wǎng)絡(luò)結(jié)構(gòu)來(lái)說(shuō)，成本效益不高。盡管如此，這些交換機(jī)可以作為主高帶寬交換機(jī)的協(xié)處理器，協(xié)助卸載集體操作和其他處理任務(wù)。

在后文中，將以Juniper 的 Trio 架構(gòu)（用于 MX 系列路由器/交換機(jī)）為例，解釋如何在具有靈活數(shù)據(jù)包處理引擎的可編程交換機(jī)中實(shí)現(xiàn)集體操作。這部分內(nèi)容主要基于MIT研究人員與瞻博網(wǎng)絡(luò)合作發(fā)表的SIGCOMM論文。

使用 Trio 進(jìn)行在網(wǎng)計(jì)算

Trio 概述

Trio6擁有1.6Tbps的帶寬。其查找子系統(tǒng)包含多個(gè)包處理引擎 (PPE) 和加速器，用于執(zhí)行哈希/記錄、鎖定和過(guò)濾器等特殊功能。通過(guò)使用高帶寬交叉開(kāi)關(guān)(cross bar)，PPE 可與加速器和內(nèi)存子系統(tǒng)進(jìn)行通信。

每個(gè) PPE 都是一個(gè) 256 位寬的 VLIW 微碼引擎。它具有 8-stage執(zhí)行pipeline，可以同時(shí)支持多個(gè)線程（20+）。PPE遵循桶式 pipeline架構(gòu)，這意味著線程在任何時(shí)候只能處于單個(gè) pipeline階段。這種設(shè)計(jì)簡(jiǎn)化了旁路邏輯，創(chuàng)造出一個(gè)在面積和功耗上更優(yōu)的優(yōu)化方案。PPE 針對(duì)更高的總吞吐量進(jìn)行了優(yōu)化。每個(gè)微碼指令可以控制多個(gè)算術(shù)邏輯單元（ALU）、支持多種操作數(shù)選擇方式、結(jié)果路由以及復(fù)雜的多路分支，從而實(shí)現(xiàn)了豐富的包處理功能。

當(dāng)數(shù)據(jù)包到達(dá)時(shí)，前大約 200B 會(huì)被發(fā)送到一個(gè)空閑的 PPE 線程進(jìn)行處理。如果需要更深入地檢查數(shù)據(jù)包，PPE 線程還可以從數(shù)據(jù)包尾部讀取額外字節(jié)。此功能對(duì)于集體操作至關(guān)重要，因?yàn)椴僮鲾?shù)（梯度、參數(shù)）往往超出前 200 個(gè)字節(jié)并占據(jù)了整個(gè)數(shù)據(jù)包有效負(fù)載。

哈希引擎（加速器）包含多個(gè)哈希表，PPE 可以在其中插入、查找或刪除條目。哈希記錄存儲(chǔ)在 DataMemory 中，可以由 PPE 修改。

內(nèi)存子系統(tǒng)包含用于包處理結(jié)構(gòu)的DataMemory和用于延遲帶寬緩沖的包緩沖區(qū)。DataMemory前面有一個(gè)大型片上內(nèi)存。片上內(nèi)存被劃分為兩部分：一部分作為DataMemory地址空間的擴(kuò)展，另一部分作為DataMemory訪問(wèn)的大緩存。這兩部分之間的劃分是可變的。

數(shù)據(jù)包處理涉及許多讀取-修改-寫(xiě)入操作。有時(shí)，多個(gè) PPE 可以訪問(wèn)同一位置以更新其內(nèi)容。一種簡(jiǎn)單的方法是賦予每個(gè)線程對(duì)內(nèi)存位置的完全所有權(quán)，直到其讀取-修改-寫(xiě)入操作完成。但是，這樣做效率低下，可能會(huì)大大降低性能。在 Trio 中，PPE 將讀取-修改-寫(xiě)入操作卸載到內(nèi)存子系統(tǒng)。內(nèi)存子系統(tǒng)包含多個(gè)讀取-修改-寫(xiě)入引擎（每個(gè)引擎處理特定的地址子集），這些引擎可以在每個(gè)周期處理這些請(qǐng)求（8字節(jié)）。當(dāng)多個(gè)請(qǐng)求到達(dá)特定引擎的某個(gè)內(nèi)存位置時(shí)，引擎會(huì)按順序處理這些請(qǐng)求，從而確保更新的一致性。

該架構(gòu)還允許通過(guò)讀取數(shù)據(jù)內(nèi)存的內(nèi)容、使用 PPE 處理附加適當(dāng)?shù)膱?bào)頭、將新數(shù)據(jù)包寫(xiě)入數(shù)據(jù)包緩沖區(qū)，并將其排隊(duì)發(fā)送，從而創(chuàng)建新報(bào)文。因此，Trio架構(gòu)具備了實(shí)現(xiàn)網(wǎng)絡(luò)內(nèi)集體操作所需的所有關(guān)鍵組件。

Trio PFE

這種架構(gòu)的優(yōu)點(diǎn)在于，交換機(jī)中的任何處理都不是硬編碼的。它具有靈活性，可以支持和解析任何新協(xié)議（UEC或自定義協(xié)議）。

網(wǎng)內(nèi)聚合流程

本節(jié)介紹了論文中使用的網(wǎng)內(nèi)聚合流程。

首先，需要定義集體數(shù)據(jù)包的報(bào)頭（通常在 UDP 層下）和數(shù)據(jù)包格式。在 SIGCOMM 論文中，作者創(chuàng)建了一個(gè)名為 Trio-ML 的自定義報(bào)頭。識(shí)別聚合線程、源/目標(biāo)GPU、梯度塊以及塊中梯度數(shù)量的足夠信息。ML 框架通常允許插件來(lái)支持自定義通信協(xié)議。

在訓(xùn)練開(kāi)始前，控制平面在作業(yè)配置時(shí)在哈希表中創(chuàng)建作業(yè)記錄，并一直持續(xù)到作業(yè)完成。這些記錄包含了參與作業(yè)的所有源（GPU）的信息，正在被聚合的區(qū)塊數(shù)量（梯度的子集），以及如何創(chuàng)建響應(yīng)數(shù)據(jù)包的信息（包括數(shù)據(jù)包的目的地等）。

當(dāng)作業(yè)/塊的第一個(gè)數(shù)據(jù)包到達(dá)時(shí)，PPE線程會(huì)在哈希表中創(chuàng)建一個(gè)塊記錄。這個(gè)記錄追蹤集體操作的狀態(tài)（在本例是梯度聚合），包括記錄哪些源（源GPU）尚未交付梯度，以及指向DataMemory中聚合緩沖區(qū)的指針。一旦塊聚合完成，這個(gè)塊記錄就會(huì)被移除。

當(dāng)數(shù)據(jù)包到達(dá) PPE 線程時(shí)，它會(huì)解析報(bào)頭并在哈希表中查找作業(yè)/塊 ID。如果塊記錄不存在，它會(huì)創(chuàng)建記錄并在 DataMemory 中為該塊分配一個(gè)聚合緩沖區(qū)。如果塊記錄已經(jīng)存在，則執(zhí)行讀取-修改-寫(xiě)入以將當(dāng)前數(shù)據(jù)包的梯度聚合到聚合緩沖區(qū)。如果這是最后一個(gè)需要聚合梯度的源（塊記錄指示所有源的狀態(tài)），則生成響應(yīng)數(shù)據(jù)包，將其寫(xiě)入數(shù)據(jù)包緩沖存儲(chǔ)器，并將其排入排隊(duì)子系統(tǒng)。之后清理塊記錄。

訓(xùn)練結(jié)束后，系統(tǒng)（在控制平面）可以清理作業(yè)記錄以釋放分配的空間。

網(wǎng)絡(luò)內(nèi)聚合流程（飲用自 SIGCOMM 2022 Juniper/MIT 論文）

在一個(gè)小型設(shè)置（使用六臺(tái)A100 GPU服務(wù)器和圖像識(shí)別中使用的小型DNN模型）中的結(jié)果顯示，在訓(xùn)練過(guò)程中性能提高了1.8倍。這是一個(gè)概念驗(yàn)證（POC）設(shè)計(jì)。

雖然使用數(shù)據(jù)包處理引擎（或網(wǎng)絡(luò)處理器）的網(wǎng)絡(luò)交換機(jī)在交換機(jī)內(nèi)部實(shí)現(xiàn)集體操作方面提供了極大的靈活性，但這些交換機(jī)并不像使用固定pipeline數(shù)據(jù)包處理的淺/深緩沖交換機(jī)那樣具有高端口密度。此外，由于數(shù)據(jù)包處理預(yù)算有限，在網(wǎng)絡(luò)集體操作中花費(fèi)的時(shí)間越多，數(shù)據(jù)包處理性能就越低。

解決此問(wèn)題的一種方法是將這些可編程交換機(jī)作為協(xié)處理器，連接到常規(guī)的高帶寬交換機(jī)上，并使用這些交換機(jī)的部分 WAN 帶寬與協(xié)處理器進(jìn)行通信，以處理集體操作。

在構(gòu)建支持可編程交換機(jī)進(jìn)行網(wǎng)內(nèi)集體操作的架構(gòu)時(shí)，沒(méi)有一種萬(wàn)能的解決方案。這取決于集群大小、正在訓(xùn)練的模型類(lèi)型以及要卸載的集體操作。

總結(jié)

在本文中，作者解釋了深度學(xué)習(xí)/通用人工智能訓(xùn)練中使用的集體操作，以及網(wǎng)絡(luò)設(shè)備如何幫助卸載這些操作。在網(wǎng)計(jì)算的爭(zhēng)論已經(jīng)存在了十多年，但尚未在業(yè)界引起足夠的關(guān)注。

然而，隨著最近大語(yǔ)言模型和通用人工智能模型訓(xùn)練工作負(fù)載的激增，以及 GPU 的稀缺性和高成本，任何在網(wǎng)絡(luò)中進(jìn)行的卸載都將直接轉(zhuǎn)化為公有云和數(shù)據(jù)中心顯著的成本優(yōu)化與效能提升。Nvidia 已經(jīng)在其 InfiniBand 和 NVLink 交換機(jī)中原生支持此功能。隨著 UEC 聯(lián)盟致力于標(biāo)準(zhǔn)化 INC，這些操作進(jìn)入高帶寬以太網(wǎng)交換機(jī)只是時(shí)間問(wèn)題。與此同時(shí)，可編程網(wǎng)絡(luò)設(shè)備可能會(huì)單獨(dú)使用或作為協(xié)處理器來(lái)卸載訓(xùn)練工作負(fù)載中的集體操作。

總體而言，隨著從公有云到邊緣計(jì)算等各個(gè)領(lǐng)域的 AI/ML 工作負(fù)載呈指數(shù)級(jí)增長(zhǎng)，交換機(jī)/路由器將開(kāi)始發(fā)揮重要作用，不僅可以盡可能快地傳輸數(shù)據(jù)，還可以卸載某些操作以提高性能并降低成本。