AWS Sagemaker Hyperpod 上的高效能 GPU 記憶體傳輸

現代深度學習基礎設施通常需要在機器之間的 GPU 之間傳輸大量資料。在 Perplexity，我們遇到了一項獨特的技術挑戰：如何以最大可能的速度在機器之間高效地傳輸非連續的 GPU 記憶體區域。我們的目標平台 AWS p5 執行個體透過 32 張網路卡提供了高達 3200 Gbps 的強大網路頻寬。本文分享了我們建構客製化高效能網路解決方案的歷程，該解決方案達到了此理論頻寬的 97.1%。

我們的使用案例提出了幾項關鍵的技術需求：

遠端 GPU 之間非連續記憶體區塊的高頻寬傳輸
在不中斷持續作業的情況下，動態新增或移除 Kubernetes 部署節點的能力
支援點對點通訊模式

雖然 NVIDIA 的 NCCL 函式庫是分散式深度學習的事實標準，但它並不適合我們的使用案例：

NCCL 擅長集合通訊，但需要建立一個靜態的「世界」(world)，這意味著在調整參與節點時需要重新啟動整個叢集。
NCCL 的同步通訊模型增加了我們非同步工作負載的複雜性
我們希望直接控制記憶體傳輸模式以進行最佳化
建構我們自己的解決方案提供了寶貴的學習機會

為了理解我們的解決方案，讓我們首先探討現代高效能網路與傳統網路有何不同。

我們日常使用的大多數網路都依賴 TCP/IP 協定，其中應用程式透過作業系統核心使用 Socket 與網路卡進行通訊。然而，高效能網路使用 RDMA（遠端直接記憶體存取）——這是一種完全不同的硬體與軟體堆疊，可在不涉及 CPU 的情況下實現機器之間的直接記憶體存取。

AWS 提供了 Elastic Fabric Adapter (EFA)，這是一種實作 Amazon 客製化協定 Scalable Reliable Datagram (SRD) 的自訂網路介面。與傳統 TCP/IP 網路不同（資料必須在使用者空間、核心空間與網路緩衝區之間多次複製），支援 RDMA 的 EFA 允許在 GPU 記憶體與網路卡之間直接傳輸資料，完全繞過 CPU。

建構高效能網路系統需要重新思考幾個基本假設：

緩衝區所有權：與傳統 Socket（由核心管理網路緩衝區並需要在使用者空間與核心空間之間複製資料）不同，RDMA 要求應用程式自行管理緩衝區。當應用程式發起網路作業時，它會將緩衝區所有權轉移給網路卡，直到作業完成，從而消除了資料複製的需求。
記憶體註冊：應用程式必須向作業系統核心註冊記憶體區域。核心會設定虛擬位址對應，允許 CPU、GPU 與網路卡都能理解相同的虛擬位址。此註冊為一次性操作，能夠啟用後續的零複製資料傳輸。
控制平面與資料平面：高效能網路將操作分為兩類：
1. 控制平面操作（如連線設定與記憶體註冊）會經過核心以確保安全性
2. 資料平面操作（實際資料傳輸）繞過核心以獲得最高效能
傳輸前的接收：在沒有核心管理緩衝區的情況下，應用程式必須預先發布接收操作，指定傳入資料應放置的位置。這與應用程式可隨時接收資料的 Socket 模型有根本上的轉變。
基於輪詢的完成機制：應用程式直接輪詢硬體完成佇列，而不是透過 epoll 等機制等待核心通知。這消除了系統呼叫的開銷，並允許對已完成的操作立即做出反應。
硬體拓撲感知：理解並針對硬體拓撲進行最佳化，對於達到最大效能至關重要。

AWS p5 執行個體具有複雜的內部架構。如下所示，每個執行個體包含兩個形成兩個 NUMA 節點的 CPU 插槽，每個 NUMA 節點連接到四個 PCIe 交換器：

在每個 PCIe 交換器下，我們發現：

四張 100 Gbps EFA 網路卡
一個 NVIDIA H100 GPU
一個 NVMe SSD

TCP/IP 與 RDMA 傳輸的資料路徑展示了它們在方法上的根本差異：

GPU、使用者空間應用程式、Linux 核心和網路介面之間的 TCP 與 GPUDirect RDMA 資料傳輸路徑。

使用 TCP/IP（左側）時，資料必須多次複製：

從 GPU 記憶體到使用者空間中的應用程式緩衝區
從使用者空間緩衝區到核心緩衝區
從核心緩衝區到網路卡
接收端也會發生反向處理過程

每次複製操作都會消耗 CPU 週期與記憶體頻寬。應用程式還必須針對每個網路操作在使用者空間與核心空間之間進行內容切換。

相比之下，RDMA（右側）提供了真正的零複製資料傳輸：

網路卡直接讀取 GPU 記憶體
資料直接傳輸至遠端網路卡
遠端網路卡直接寫入目標 GPU 記憶體
應用程式只需檢查使用者空間中的完成佇列 (CQ) 即可得知傳輸何時完成

透過 RDMA 與正確的硬體配對，在兩個 GPU 之間傳輸資料僅需遍歷本機 PCIe 交換器與網路：

透過 PCIe Gen5 x16 與四條 100Gbps EFA 連線相連的成對本地與遠端 GPU。

相比之下，TCP/IP 傳輸必須透過主記憶體多次複製資料，導致嚴重的 PCIe 匯流排擁塞：

主記憶體、CPU、四個 PCIe 交換器與四個 GPU 之間因 TCP 傳輸造成的系統匯流排壅塞。

我們使用了 libfabric——一個為 Fabric 服務提供通用介面的框架。我們的實作使用了兩種類型的 RDMA 操作：

用於攜帶記憶體區域相關元資料（例如偏移量與大小）控制訊息的雙向 RDMA (SEND/RECV)。
用於實際資料傳輸的單向 RDMA WRITE，其中每個 WRITE 操作處理一個連續的記憶體區塊。

我們的開發進展經歷了幾個階段：

使用 SEND/RECV 實作了基本的單向訊息傳輸
擴展至雙向通訊
增加了用於直接 GPU-GPU 傳輸的 GPUDirect RDMA WRITE
擴展以處理多個併發傳輸
引入操作佇列以提高穩健性
在單張網路卡上實現了 97.4% 的頻寬利用率

當擴展至 32 張網路卡時，我們實作了幾項關鍵的最佳化：

操作佇列：我們維持一個應用程式層級的佇列，而不是直接將操作提交給網路卡。這提供了針對網路擁塞的穩健性並簡化了程式設計模型。
網路暖機：預先建立連線改善了啟動效能。
多討論串：為每個 GPU 的網路操作提供專用討論串。
CPU 核心綁定：將討論串綁定到特定的 CPU 核心，以避免 NUMA 效應與快取未命中。
狀態分片：減少討論串之間的爭用。
操作批次處理：將多個操作合併提交。
延遲操作提交：操作首先在應用程式中排隊。輪詢完成佇列後，我們嘗試將待處理的操作提交給網路卡，確保網路資源的高效使用。
NUMA 感知資源配置：在正確的 NUMA 節點上配置 libfabric 資源（如完成佇列），以最小化記憶體存取延遲。

透過這些最佳化，我們最終達到了 3,108 Gbps 的效能——即理論最大頻寬的 97.1%。

下方的影片展示了我們命令列程式的實際運作情況。它將非連續的 GPU 記憶體區塊傳輸到遠端節點，達到了 3108.283 Gbps 的傳輸速度——證明了網路基礎設施近乎理論值的頻寬利用率：

建構高效能網路系統需要同時理解硬體架構與系統設計原則。雖然像 NCCL 這樣的函式庫為常見模式提供了出色的解決方案，但有時針對特定需求必須採用客製化解決方案。

我們的旅程證明了透過細心關注系統架構、硬體拓撲以及各種最佳化技術，實現近乎理論值的網路效能是可能的。關鍵不僅在於理解各個元件，還在於它們如何交互作用以形成一個完整的系統。

此旅程的完整技術深度解析，包括實作細節與程式碼範例，可在我們的開源儲存庫中取得。

BibTeX

@online{perplexity2025HighPerformanceGpuMemoryTransferOnAwsSagemakerHyperpod,
  author = {{Perplexity} and Lequn Chen and Phi Hoang},
  title = {{High-performance GPU memory transfer on AWS Sagemaker Hyperpod}},
  date = {2025-02-10},
  url = {https://www.perplexity.ai/hub/blog/high-performance-gpu-memory-transfer-on-aws-sagemaker-hyperpod}
}

技術挑戰

現代高效能網路

高效能網路設計哲學

理解硬體拓撲

使用 libfabric 建構

邁向巔峰效能之旅

結論

引用