利用 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 上實現接近即時的大型模型冷啓動，並支持更長的上下文。