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

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

技術挑戰

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

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

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

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

現代高效能網路

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

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

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

高效能網路設計哲學

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

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

了解硬體拓樸

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

放大

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

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

TCP/IP 和 RDMA 傳輸的資料路徑展示了其方法上的根本差異：

放大

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

1. 從 GPU 記憶體到使用者空間中的應用程式緩衝區
2. 從使用者空間緩衝區到核心緩衝區
3. 從核心緩衝區到網路卡
4. 接收端會發生相反的過程

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

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

1. 網路卡直接從 GPU 記憶體讀取
2. 資料直接傳輸到遠端網路卡
3. 遠端網路卡直接寫入目的地 GPU 記憶體
4. 應用程式只需要檢查使用者空間中的完成佇列 (CQ) 即可知道傳輸何時完成

有了 RDMA 和適當的硬體配對，在兩個 GPU 之間傳輸資料只需穿越本地 PCIe 交換器和網路：

放大

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

放大

使用 libfabric 建構

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

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

邁向顛峰效能之旅

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

1. 使用 SEND/RECV 實作基本的單向訊息傳輸
2. 已擴展至雙向通訊
3. 新增用於直接 GPU-GPU 傳輸的 GPUDirect RDMA WRITE
4. 擴展以處理多個並行傳輸
5. 引入作業排隊以提高強韌性
6. 在單張網路卡上達到 97.4% 的頻寬利用率

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

1. 作業排隊：我們維護應用程式層級的佇列，而不是直接將作業提交給網路卡。這提供了對抗網路壅塞的強韌性，並簡化了編程模型。
2. 網路暖機：預先建立連線可改善啟動效能。
3. 多執行緒：為每個 GPU 的網路作業配置專用執行緒。
4. CPU 核心綁定：將執行緒綁定到特定的 CPU 核心，以避免 NUMA 效應和快取未命中。
5. 狀態分片：減少執行緒之間的競爭。
6. 作業批次處理：同時提交多個作業。
7. 延遲作業張貼：作業首先在應用程式中排隊。輪詢完成佇列後，我們嘗試將暫掛的作業提交給網路卡，確保有效利用網路資源。
8. 具備 NUMA 感知的資源配置：在正確的 NUMA 節點上配置 libfabric 資源（例如完成佇列），以將記憶體存取延遲降至最低。

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

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

結論

建構高效能網路系統需要同時理解硬體架構和系統設計原則。雖然像 NCCL 這樣的程式庫為常見模式提供了優秀的解決方案，但有時針對特定需求，自訂解決方案是必要的。

我們的歷程證明，只要密切關注系統架構、硬體拓樸和各種最佳化技術，就有可能達到接近理論值的網路效能。關鍵不僅在於理解個別元件，還在於它們如何互動以形成一個完整的系統。

此歷程的完整技術深入解析（包含實作細節和程式碼範例）可在我們的開源儲存庫中取得。

引用

@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}
}