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 \times model_size))。

分布式层级卸载通过四个主要技术机制解决了这些限制：

1. Meta-Device 初始化与 mmap 权重加载

为了防止在模型创建期间出现巨大的驻留集大小 (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 \times 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 \times 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. la 吞吐量，达到了理想线性扩展的约 83%。

性能与验证结果

平台无关性

分布式层级卸载是平台无关的，可以在 NVIDIA GPU (CUDA/NCCL) 和 Ascend NPU (CANN/HCCL) 上运行。

在 NVIDIA B300 GPU 上，对于 1024\times 1024 T2I 任务，使用 4 个并发请求的 DLO+AG DP4 实现了 HSDP+USP4 吞吐量的 1.39 倍，同时仅使用了 30% 的 HBM (12.6 GiB vs 42.0 GiB)。对于 720p 10s 视频生成，DLO+AG+USP4 的延迟与 HSDP 相比在 2.13% 以内，而 HBM 使用量仅为 47%。

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

使用 MiniMax-H3 模型在 8\times NVIDIA B300 GPU 上的研究表明，最佳 DLO 模式取决于硬件拓扑：

DP1\times SP8：AllGather 模式对于最低延迟 (P50 34.55s) 最为理想。
DP4\times SP2：AllGather 模式为吞吐量提供了平衡点 (151.89 videos/h)。
DP8\times SP1：Rank-local DLO 模式对于最高吞吐量 (183.78 videos/h) 和最低能耗 (43.97 Wh/video) 最为理想。

Ascend NPU 内存统计

在 Ascend 硬件上，pin_memory() 通过 /dev/davinci_manager 进行分配，将分片存放在 CPU 内核 DMA 内存中。这种内存对 cgroup 内存控制器是不可见的。因此，cgroup 可见的内存随 O(model_size + dp_size \times 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	4	~405 GB	Yes
400 GB	4	~423 GB	~823 GB	Yes
400 GB	8	~443 GB	~843 GB	Yes