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

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

什么是“集體操作”？

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

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

下面是常用的集體操作：

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

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

AllReduce

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

ReduceScatter

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

廣播

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

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

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

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

AI/ML 框架中的集體操作

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

例如，在使用AllReduce集體操作進(jìn)行梯度聚合時，梯度可以按照環(huán)形模式從一個 GPU 發(fā)送到另一個 GPU，其中每個GPU都會將其從上一個GPU接收到的梯度與其本地計算出的梯度進(jìn)行聚合，然后再將結(jié)果發(fā)送給下一個GPU。這個過程很慢，因為梯度聚合是按順序完成的，最終的結(jié)果也會按順序在環(huán)形拓?fù)渲袀骰氐剿?GPU。

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

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

梯度聚合流量（僅顯示 4 個 GPU）

集體操作的卸載

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

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

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

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

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

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

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

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

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

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

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

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

能夠處理擁塞并為集合操作流量維護(hù)QoS，特別是在芯片中發(fā)生多個并行集合操作時。

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

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

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

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

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

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

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

Trio 概述

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

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

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

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

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

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

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

Trio PFE

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

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

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

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

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

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

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

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

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

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

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

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

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

總結(jié)

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

然而，隨著最近大語言模型和通用人工智能模型訓(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ī)只是時間問題。與此同時，可編程網(wǎng)絡(luò)設(shè)備可能會單獨使用或作為協(xié)處理器來卸載訓(xùn)練工作負(fù)載中的集體操作。

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