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

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

什么是“集體操作”？

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

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

下面是常用的集體操作：

Reduce：從所有節點聚合數據（求和或平均）并將結果發送給其中一個節點。

AllReduce：聚合所有節點的數據，并將結果發送回所有節點。

AllReduce

ReduceScatter：聚合所有節點的數據，并將結果（每個節點獲得結果的唯一子集）分發到所有節點。

ReduceScatter

廣播：將數據從一個節點發送到組中的所有其他節點。

廣播

AllGather：收集數據的不同部分，并將其分發給所有節點。

Scatter：將一個節點的不同值分發到所有節點。

AlltoAll：將所有節點的數據分發到所有節點。

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

AI/ML 框架中的集體操作

NCCL 實現了多種集體操作算法，包括環形（Ring）、樹形（Tree）以及雙二叉樹（Double Binary Trees）等，根據數據大小、涉及的 GPU 數量和網絡拓撲自動選擇最為高效的算法。這些算法旨在優化集體操作中的數據傳輸。

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

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

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

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

集體操作的卸載

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

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

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

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

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

以太網結構和在網集體操作

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

在數據平面交換機中執行INC需要硬件對各種操作的支持。

解析新協議/數據包格式，以確定數據包是否是集體操作的一部分。

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

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

能夠從內存中讀取累積的結果，創建新的數據包并將其發送出去。

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

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

有幾種方法可以實現上述的集體操作。

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

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

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

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

使用 Trio 進行在網計算

Trio 概述

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

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

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

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

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

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

該架構還允許通過讀取數據內存的內容、使用 PPE 處理附加適當的報頭、將新數據包寫入數據包緩沖區，并將其排隊發送，從而創建新報文。因此，Trio架構具備了實現網絡內集體操作所需的所有關鍵組件。

Trio PFE

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

網內聚合流程

本節介紹了論文中使用的網內聚合流程。

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

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

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

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

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

網絡內聚合流程（飲用自 SIGCOMM 2022 Juniper/MIT 論文）

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

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

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

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

總結

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

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

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