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

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

什么是“集體操作”？

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

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

下面是常用的集體操作：

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

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

AllReduce

ReduceScatter：聚合所有節(jié)點的數(shù)據(jù)，并將結果（每個節(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)絡拓撲自動選擇最為高效的算法。這些算法旨在優(yōu)化集體操作中的數(shù)據(jù)傳輸。

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

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

下圖展示了7個GPU以二叉樹的形式排列以進行梯度聚合。假設這些 GPU 是連接到不同葉交換機和主干交換機的大型網(wǎng)絡拓撲的一部分。該圖還顯示了 GPU 之間梯度聚合的流量模式。在網(wǎng)絡中，這些交換機是被動設備，負責轉發(fā)GPU之間的通信結果。

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

集體操作的卸載

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

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

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

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

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

以太網(wǎng)結構和在網(wǎng)集體操作

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

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

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

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

能夠對有效載荷進行操作，并跨多個數(shù)據(jù)包累積結果。這需要支持各種浮點/整數(shù)格式的算術運算，而以太網(wǎng)交換機通常不具備這些功能。

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

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

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

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

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

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

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

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

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

Trio 概述

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

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

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

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

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

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

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

Trio PFE

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

網(wǎng)內聚合流程

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

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

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

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

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

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

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

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

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

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

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

總結

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

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

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