利用 Amazon FSx for Lustre 上的 GPUDirect 和 TurboQuant 加速 LLM 模型載入並擴大上下文視窗

在 AWS GPU 例項上迭代部署大型語言模型 (LLM) 時，模型載入到 GPU 高頻寬記憶體 (HBM) 的時間會隨著模型規模的增大而顯著增加，導致冷啟動首次令牌時間 (TTFT) 過長。本文介紹瞭如何透過 Amazon FSx for Lustre 與 NVIDIA GPUDirect Storage (GDS) 的結合，加上巧妙的規劃，從根本改變冷啟動 TTFT 的計算方式，將模型載入時間從分鐘縮短到秒。同時，還介紹了新發布的 TurboQuant KV 快取對增加上下文視窗的巨大作用。

背景：AWS 上的 NVIDIA Blackwell 架構 AWS 近期推出了基於 NVIDIA Blackwell 架構的 Amazon EC2 P6e 和 P6 例項系列。旗艦級 P6e UltraServer 將 72 個 NVIDIA Blackwell GPU 整合到一個 NVLink 域中，提供 130 TB/s 的雙向頻寬、13.4 TB 的 HBM3e 和 360 petaflops 的 FP8 計算能力。這些 UltraServer 通常用於萬億引數級前沿模型的大規模分散式訓練。本文則聚焦於最佳化單個 P6 或 P5en 例項的冷啟動 TTFT。

模型載入瓶頸 傳統 CPU 模型載入將檢查點流經 CPU 記憶體，並依次透過 PCIe 將權重複制到每個 GPU。而分片 GDS 載入則預分檢查點，所有 GPU 直接透過 EFA 從 FSx for Lustre 並行讀取分片到 HBM，完全繞過 CPU。對於 Llama 3.1 405B 模型（約 800 GB BF16 檢查點），傳統載入需要 10–20 分鐘，而 GDS 幾乎可以消除這一瓶頸。長時間載入不僅影響冷啟動延遲，還影響自動擴縮響應速度、故障恢復以及成本效率。

直接路徑：FSx for Lustre 與 GPUDirect Storage Amazon FSx for Lustre 是一種高效能並行檔案系統，與 NVIDIA GDS 結合時，可以建立多條直達 GPU 記憶體的資料路徑。此整合依賴兩個關鍵技術：Amazon Elastic Fabric Adapter (EFA) 和 NVIDIA GPUDirect Storage。EFA 使用 SRD 協議繞過作業系統開銷，P5en 例項具有 16 個 200 Gbps 的 EFA 介面，總聚合網路頻寬達 3,200 Gbps。GDS 允許 DMA 直接從網路介面傳輸到 GPU HBM。測試中使用了 Persistent_2 EFA 檔案系統，每 TiB 提供 1000 MBps 吞吐量，20 個 OST 可提供約 94 GiB/s 的檔案系統吞吐量。

分片並行載入 以 P5en 例項（8 個 H200 GPU，每個 141 GB HBM3e）為例，對於 Llama 3.1 405B 模型（FP8 時約 400 GB），權重無法裝入單個 GPU，需要張量並行。方法分為四個階段：

• 階段 0：配置基礎設施：在同一個 VPC 和可用區中準備啟用 EFA 的 FSx for Lustre 檔案系統和配置了 GDS 的 GPU 例項。使用 AWS CloudFormation 模板和設定指令碼可自動完成。
• 階段 1：預分片和預量化：離線使用 vLLM 將模型分割為 8 個張量並行分片並量化為 FP8，儲存到 FSx for Lustre。檢查點大小從 ~800 GB 降至 ~400 GB。此步驟僅在模型更新時需要重複。
• 階段 2：GDS 並行載入：啟動推理伺服器時，每個 GPU 透過 GDS 直接從 FSx for Lustre 讀取自己的分片，並行載入權重到 HBM。
• 階段 3：推理：模型就緒後立即開始服務推理請求。

TurboQuant KV 快取：TurboQuant 是一種量化技術，以 4 位甚至 2 位精度儲存 KV 快取，在不影響質量的情況下將每個令牌的 KV 快取記憶體佔用減少 4–8 倍，從而支援更長的上下文視窗。

透過結合這些技術，使用者可以在 AWS 上實現接近即時的大型模型冷啟動，並支援更長的上下文。