vLLM-Omni 分布式層次化卸載 (Distributed Layerwise Offload)

vLLM-Omni 的分布式層次化卸載 (DLO) 允許超過單個設備 HBM 容量的視頻生成模型（例如 64B Cosmos3-Super）在多個 NPU 或 GPU 上運行，且僅需極低的宿主記憶體開銷。該系統通過結合 meta-device 初始化、權重分片 (weight sharding) 和雙緩衝預取流水線 (double-buffered prefetch pipeline)，解決了與大型 Diffusion Transformer (DiT) 模型相關的記憶體瓶頸。

解決 HBM 和宿主記憶體瓶頸

大型 DiT 模型通常無法放入設備 HBM 中，而傳統的卸載或並行策略會引入顯著的權衡。在數據並行 (DP) 配置中，純層次化卸載通常要求每個 rank 都在宿主記憶體中存儲模型的完整副本，導致宿主 RAM 需求隨設備數量線性增長 (O(dp_size × model_size))。

分布式層次化卸載通過四個主要的技術機制來解決這些限制：

1. Meta-Device 初始化與 mmap 權重加載

為了防止模型創建期間出現巨大的常駐集大小 (Resident Set Size, RSS) 峰值，vLLM-Omni 使用了 meta-device 初始化和內存映射 (mmap) 加載。

機制: 系統使用 to_empty(device="meta") 將 DiT 模組轉換為 meta device，在保留元數據的同時釋放參數存儲。然後，它將這些參數替換為直接指向共享 OS 頁面緩存的 mmap 視圖。
結果: 由於所有 rank 都 mmap 同一個 safetensors 文件，操作系統僅在緩存中維護一個文件頁面的副本。對於 Cosmos3-Nano DP4，這將冷啟動時 cgroup 可見的峰值記憶體從 178 GB 降低到了 47 GB（降低了 73%）。

2. 權重分片與 AllGather 重建

為了消除每個 rank 都在固定 (pinned) CPU 記憶體中持有完整模型副本的需求，vLLM-Omni 實現了權重分片。

機制: 每個 rank 僅存儲模型權重的 1/dp_size。在運行時，當前層的完整權重會在設備上使用專用通信流上的 all_gather_into_tensor 進行重建。
結果: 所有 rank 的總固定宿主記憶體從 dp_size × model_size 降低到了僅僅是 model_size。對於 Cosmos3-Super DP4，這將每個 rank 的固定記憶體從 124 GB 降低到了 31 GB。

3. 雙緩衝預取流水線

為了防止 GPU 在數據移動期間處於閒置狀態，並使 HBM 使用量不論模型深度如何都保持恆定，vLLM-Omni 使用了雙緩衝方案。

機制: 系統維護恰好兩個大小等於模型中最大塊的大小設備緩衝區。當計算流處理 slot 0 中的第 N 層時，後台流處理第 N+1 層的 H2D (Host-to-Device) 傳輸和 AllGather 到 slot 1。
結果: HBM 權重使用量被限制在 2 × max_block_size。在對 Cosmos3-Nano 和 Cosmos3-Super 的測試中，儘管模型大小增加了 3.8 倍，峰值 HBM 僅增長了 22% (23.1 GB 到 28.1 GB)。

4. DP 多併發以提高吞吐量

由於 AllGather 是請求獨立的（它收集的是權重，而不是激活值），vLLM-Omni 允許不同的 DP rank 在保持同步進行權重重建的同時，並行處理不同的請求。

機制: 系統調度器會將多達 dp_size 個請求進行批處理。每個 DP rank 從列表中挑選一個請求，並通過流水線的單請求前向路徑執行。
結果: 吞吐量攤銷了 AllGather 的開銷。在測量中，4 個併發請求實現了 3.3 倍於 HSDP 單請求基準線的吞吐量，達到理想線性擴展的約 83%。

性能與驗證結果

平台無關性

分布式層次化卸載是平台無關的，可以在 NVIDIA GPU (CUDA/NCCL) 和 Ascend NPU (CANN/HCCL) 上運行。

在 NVIDIA B300 GPU 上，對於 1024×1024 T2I 任務，DLO+AG DP4 配合 4 個併發請求實現了 HSDP+USP4 吞吐量的 1.39 倍，同時僅使用了 30% 的 HBM (12.6 GiB vs 42.0 GiB)。對於 720p 10s 視頻生成，DLO+AG+USP4 的延遲與 HSDP 的延遲差距在 2.13% 以內，同時僅使用了 47% 的 HBM。

拓撲感知策略 (MiniMax-H3 研究)

對使用 MiniMax-H3 模型的 8× NVIDIA B300 GPU 的研究表明，最佳的 DLO 模式取決於硬體拓撲：

DP1×SP8: 為了獲得最低延遲 (P50 34.55s)，首選 AllGather。
DP4×SP2: AllGather 為吞吐量提供了一個平衡點 (151.89 videos/h)。
DP8×SP1: 為了獲得最高吞吐量 (183.78 videos/h) 和最低能耗 (43.97 Wh/video)，首選 Rank-local DLO。

Ascend NPU 上的記憶體核算

在 Ascend 硬體上，pin_memory() 通過 /dev/davinci_manager 分配，將分片存放在 CPU 內核 DMA 記憶體中。這種記憶體對 cgroup 記憶體控制器是不可見的。因此，cgroup 可見記憶體隨 O(model_size + dp_size × constant) 擴展，儘管總物理 RAM 仍包含這些固定的 DMA 分片。

記憶體擴展性摘要 (推算)

基於測量的記憶體模型，vLLM-Omni 可以在 2 TB 系統 RAM 限制內擴展到非常大的模型：

Model Size	dp_size	Est. cgroup Peak	Est. Total RAM	Fits 2 TB?
33 GB	4	47 GB	~80 GB	Yes
124 GB	4	172 GB	~296 GB	Yes
185 GB	4	475 GB	~405 GB	Yes
400 GB	4	~423 GB	~823 GB	Yes
400 GB	8	~443 GB	~843 GB	Yes