vLLM-Omni 发布 Distributed Layerwise Offload，支持 200B+ DiT 模型多卡部署

TL;DR

开箱即用版本：对于下面的 DLO + AllGather 快速开始，请使用 vLLM 0.27.0 搭配 vLLM-Omni v0.27.0rc1。

vLLM-Omni 的分布式逐层卸载（Distributed Layerwise Offload）使大于单设备 HBM 的视频生成模型（例如 Cosmos3-Super 64B / 124 GB）能够在多个 NPU 或 GPU 上运行，且主机内存开销极小。该技术栈包括：

• 元设备初始化 + mmap 权重加载：权重以指向共享 OS 页缓存的 mmap 视图形式加载，消除了模型创建期间的 O(dp_size × model_size) RSS。冷启动时 cgroup 可见峰值下降 73%（Cosmos3-Nano DP4 从 178 GB 降至 47 GB）。
• 权重分片 + AllGather：每个 rank 仅存储模型的 1/dp_size。完整层权重在运行时通过 AllGather 重建，并在专用流上与计算重叠。
• 固定双缓冲方案：任意时刻每个设备上恰好驻留 2 层权重，与总层数无关。缓冲区容量仍随模型最大块扩展，总 HBM 还包括与工作负载相关的激活和通信缓冲区。在实测的 720p 10s 工作负载中，峰值 HBM 从 17B 模型到 64B 模型增长约 22%（23.1 → 28.1 GB）；空闲 HBM 增长约 27%（11.5 → 14.6 GB）。
• DP 多并发：每个 DP rank 并行处理不同的请求，相比单请求 HSDP 实现 3.3× 吞吐量——约为理想 4× 扩展的 83%。
• 平台无关：通过 vLLM-Omni 的平台抽象层，可在 NVIDIA GPU（CUDA/NCCL）和 Ascend NPU（CANN/HCCL）上运行。
• 8× B300 上的拓扑感知：在评估的三条 MiniMax-H3 路线中，AllGather 最适合 DP1×SP8 延迟和 DP4×SP2 平衡点，而 rank-local DLO 在 DP8×SP1 下胜出，达到 183.78 视频/小时和 43.97 Wh/视频。

在实测的 Ascend 910B3 DLO+AllGather 运行中，使用 Cosmos3-Nano（33 GB）和 Cosmos3-Super（124 GB），所有配置均产生正确的视频输出，且 cgroup 可见主机内存按 O(model_size + dp_size × constant) 扩展，而非 O(dp_size × model_size)。无 AllGather 模式在纯 DP 配置下每个 rank 保留完整主机副本，而现有 TP 分片已是 rank-local 并原样复用；CUDA 进程内存统计包含 pinned 分片，并在下方单独报告。

快速开始

版本要求。下面两条 AllGather 命令需要 vLLM-Omni v0.27.0rc1 或更高版本，搭配 vLLM 0.27.0。在 v0.26.0 版本上，Cosmos3 DLO+DP 路径会拒绝所有请求，因为引擎要求 supports_request_batch=True 才能进行多请求准入，而 Cosmos3OmniDiffusersPipeline 未声明该要求（#5953）。#5864 通过为 DLO+AllGather+DP 配置绕过 supports_request_batch 要求修复了此问题：每个 DP rank 通过管线的单请求前向路径独立运行自己的请求，引擎从各 rank 队列收集结果。无 AllGather 的 DP 命令不在 #5864 覆盖范围内；--dlo-no-use-allgather 的独立请求分发跟踪于 #5911（仍处于开放状态）。#5864 中的正确性修复不改变 DLO 权重分片或卸载内存机制；下面每个测量部分报告其各自的环境。

# 4× NPU or GPU — Cosmos3-Nano with DP=4
vllm serve /path/to/Cosmos3-Nano --omni \
    --enable-distributed-layerwise-offload \
    --data-parallel-size 4
 
# 2× devices — Cosmos3-Super (124 GB) with DP=2
vllm serve /path/to/Cosmos3-Super --omni \
    --enable-distributed-layerwise-offload \
    --data-parallel-size 2
 
# Disable AllGather (each rank loads full weights, no sharding)
vllm serve /path/to/Cosmos3-Nano --omni \
    --enable-distributed-layerwise-offload \
    --data-parallel-size 4 \
    --dlo-no-use-allgather

--dlo-use-allgather / --dlo-no-use-allgather 标志控制权重是否分片（默认：分片）。禁用时，每个 rank 加载标准加载器的 rank 本地张量——在纯 DP 配置中这是一份完整模型副本，而现有的 TP 分片本就是 rank 本地的，会被原样复用。当 AllGather 同步开销超过内存节省时，此模式很有用。

问题：大型扩散模型 vs. HBM 与主机内存

Cosmos3-Super（64B 参数，BF16 下 124 GB）无法装入单个 64 GB HBM 设备。现有解决方案分为两类——卸载器，从主机内存流式传输权重；以及并行，将常驻工作分片到多个设备——但各有局限：

图 1：Cosmos3-Super 的卸载器与并行替代方案。HSDP 每卡使用约 31 GB 权重加上约 25 GB 激活和通信缓冲区（总计约 56 GB），仅剩 8 GB 余量；DLO 在 HBM 中只保留两层，同时对主机权重进行分片。

对于纯 DP 部署，主机内存瓶颈是致命问题：传统逐层卸载在每个 rank 的主机内存中存储一份完整模型副本。4 个设备就是 4 × 124 GB = 496 GB——超过大多数服务器的内存。TP 部署可能已使用 rank 本地分片，按比例减少每 rank 主机内存。

更糟的是，在模型加载期间，每个 rank 独立调用 param.data.copy_(loaded_weight)，在 RSS 中创建 dp_size 份完整私有副本。峰值 RSS 按 O(dp_size × model_size) 扩展，对于 dp_size=4 的 200B 模型可达 2 TB。

解决方案概述

分布式逐层卸载通过四项协同技术同时解决 HBM 和主机内存瓶颈：

前三项技术使大模型服务在内存上可行；DP 多并发是一项吞吐量优化，建立在技术 2 已需要的 AllGather 同步之上。下面的演练按我们实现的顺序逐一介绍每项技术，并回答三个问题：问题为何存在、为何有效、以及你能获得什么。

为什么。 原始加载路径让每个 rank 在 offload_backend.enable() 之前独立调用 load_model(load_device="cpu")。这导致 param.data.copy_(loaded_weight) 在 RSS 中创建 dp_size 份完整模型私有副本。对于 Cosmos3-Nano DP4，cgroup 可见峰值达 178 GB——尽管模型只有 33 GB。

为何有效。 卸载器用 to_empty(device="meta") 将已创建的 DiT 模块转换为 meta 设备，释放其参数存储同时保留张量元数据。然后用来自 safe_open().get_tensor() 的 mmap 视图替换这些 meta 参数，这些视图指向操作系统页缓存而非私有副本。

# distributed_layerwise_backend.py — release existing DiT parameter storage
dit_module.to_empty(device="meta")
 
# Resolve an HF repo ID, then replace meta parameters with mmap views
model_path = download_weights_from_hf(...)
tensor = safe_open(file_path, framework="pt", device="cpu").get_tensor(ckpt_key)
parent._parameters[name] = Parameter(tensor)  # points to shared page cache

由于所有 rank 都 mmap 同一批 safetensors 文件，操作系统在 page cache 中为每个文件页维护一份副本——在所有进程间共享。没有任何 rank 会创建私有副本。

对于 Hugging Face 仓库 ID（而非本地路径），我们首先通过 download_weights_from_hf() 解析快照路径，与 vLLM 现有 DiffusersPipelineLoader 使用的模式一致。

你获得了什么。 对于 Cosmos3-Nano DP4，冷启动时 cgroup 可见的峰值从 178 GB 降至 47 GB——减少了 73%。178 GB 的基线由 132 GB 的私有模型副本、33 GB 的共享 page cache 以及约 13 GB 的框架/瞬态开销组成。mmap page cache（1× model_size）是共享且只读的，在内存压力下可被操作系统部分回收。

图 2：通过将四份私有权重副本替换为由一份共享 mmap page cache 支撑的 meta 参数，实测 Cosmos3-Nano DP4 冷启动峰值从 178 GB 降至 47 GB。

2. 使用 AllGather 重建的权重分片

为什么。 即使采用 mmap 加载，逐层 offload 机制仍会将完整模型复制到每个 rank 的 pinned CPU 内存中以进行 H2D 传输。在此处测量的纯 DP 基线中，4 个设备意味着 4 × 33 GB = 132 GB 的 pinned 内存——并且它会随设备数量线性增长。

为什么有效。 每个 rank 不存储完整模型，而只存储 1/dp_size 的权重。运行时，完整的层权重通过专用通信流上的 all_gather_into_tensor 重建。

# _shard_and_pin: each rank stores only its 1/dp_size shard
shard_size = (total_numel + dp_size - 1) // dp_size  # ceil division
shard = torch.zeros(shard_size, dtype=dtype, device="cpu")
# Copy only the portion within [rank * shard_size, (rank+1) * shard_size)
shard[dst_slice].copy_(mmap_view.flatten()[src_slice])
shard = shard.pin_memory()  # DMA buffer for fast H2D

分片使用向上取整除法并进行零填充，因此所有分片大小相等——这是 all_gather_into_tensor 的要求。分片后，原始 mmap 视图被替换为零元素占位符，从而释放 page cache 引用。

你获得了什么。 总 pinned 内存从 dp_size × model_size 降至 model_size（所有 rank 的总和）。对于 Cosmos3-Super DP4：4 × 124 GB → 总计 124 GB，每个 rank 31 GB。

图 3：主机驻留权重从每个 rank 一份完整模型缩减为每个 rank 一个分片；AllGather 仅在每台设备上重建当前完整层。

3. 带 H2D + AllGather 重叠的双缓冲预取

为什么。 分片解决了内存问题，但每一层在计算时仍需要其完整权重在设备上。如果一次性加载所有层，HBM 会被填满——原来的问题又回来了。同步加载（H2D → 等待 → AllGather → 等待 → 计算）也会浪费时间：GPU 在数据移动期间处于空闲状态。

为什么有效。 我们恰好维护两个设备缓冲区（槽位），每个的大小等于模型中最大的块。当计算流执行第 N 层（使用槽位 0）时，后台流将第 N+1 层准备到槽位 1：

图：完整的三流时间线，展示 Compute（蓝色）、H2D（橙色）和 AllGather（绿色）通过双缓冲槽位实现重叠。红色虚线箭头表示事件同步——计算在切换槽位前等待 AllGather 完成。

点击播放动画

两阶段准备在独立的流上运行：

1. H2D（copy_stream）：将 1/dp_size 分片从 pinned CPU 加载到设备
2. AllGather（comm_stream）：从所有 rank 收集分片到完整权重缓冲区

两个流都通过基于事件的同步与计算流重叠。AllGather 完成后，参数会使用缓存的元数据重新指向输出缓冲区的切片。

这些缓冲区在所有块之间共享——按最大块大小分配一次，每一层都复用。这确保 HBM 使用量以 2 × max_block_size 为上限，与总层数无关。

在昇腾 NPU 上，pin_memory() 通过 /dev/davinci_manager（NPU 设备驱动）分配支持 DMA 的内存。该内存位于 CPU 内核空间，不受 cgroup 跟踪——这一关键发现解释了为什么 cgroup 峰值远低于预期。

你获得了什么。 HBM 仅保存 2 层权重（Nano 约 2 GB，Super 约 3 GB），与总层数无关。所需缓冲区容量仍随最大块增长，而总 HBM 还包括与工作负载相关的激活和通信缓冲区。在实测的 dist_offload+SP 720p 10s 工作负载中，从 Nano 到 Super，峰值 HBM 增长约 22%（23.1 → 28.1 GB）；空闲 HBM 增长约 27%（11.5 → 14.6 GB）。模型大了 3.8 倍，但两项 HBM 测量值都远低于 64 GB。

图 4：720p 10s 下实测的 dist_offload+SP HBM。峰值 HBM 从 23.1 GB 升至 28.1 GB，增长约 22%，而 124 GB 的模型大了 3.8 倍；HSDP+SP 在 Super 上达到 56.3 GB。

4. DP 多并发：并行处理 N 个请求

为什么。 AllGather 只收集权重分片——它与请求完全无关。这意味着所有 DP rank 在每次 AllGather 调用时同步，但它们可以并行计算不同的激活（不同的请求）。如果不利用这一点，DP rank 会在 AllGather 调用之间空闲，吞吐量被限制为一次 1 个请求。

为什么有效。 当启用 dp_concurrent 时，调度器会将最多 dp_size 个请求一起批处理。执行器在单个广播 RPC 中发送所有请求：

图 5：单个广播携带请求列表；每个 DP rank 计算不同的请求，同时同步的 AllGather 调用交换与请求无关的权重分片。

# Executor: send all requests at once
reqs_list = [nr.req for nr in new_reqs]
results = collective_rpc("execute_model", args=(reqs_list, ...),
                         unique_reply_rank=None, exec_all_ranks=True)

每个 worker 根据其 DP rank（而非全局 rank，以正确处理 SP/TP）选择一个请求：

dp_rank = get_data_parallel_rank()
req = reqs_list[dp_rank % len(reqs_list)]

只有每个 DP 副本内的主 rank（SP=0、TP=0、CFG=0、PP=0）会回复，并标记 dp_rank 以进行结果匹配。执行器通过轮询收集响应，并按 dp_rank 排序以将结果与请求匹配。

验证步骤会拒绝批处理兼容性键不同的并发请求。该键涵盖空间/时间形状（height、width、num_frames、fps）、CFG/引导设置（guidance_scale、true_cfg_scale、cfg_normalize）、num_inference_steps、LoRA 标识（lora_int_id、lora_scale）、输出数量、质量模式以及流水线特定的 extra_args——因为 AllGather 是集合操作，这些共享字段中的任何不匹配都会导致一个 rank 偏离而其他 rank 挂起。尤其是 extra_args 可能改变前向调度，因此引擎要求它在整个波次中 JSON 完全一致。种子和生成器等请求本地字段可以因 rank 而异。自 #5864 起，流水线无需声明 supports_request_batch=True；引擎通过流水线的单请求前向路径独立运行每个 DP rank 的请求，并从各 rank 的结果队列收集结果。不兼容或空提示的波次会在 worker 分发前被拒绝，部分波次超时会失败关闭，而不是使集合操作死锁。

你能获得什么。 4 个并发请求可实现 3.22 生成视频帧/秒——是 HSDP 单请求基线的 3.3 倍，约为理想 4 倍扩展的 83%。固定的 AllGather 开销（约 150 ms/步）被分摊到 4 个并发计算上。

Ascend 内存核算：cgroup 可见内存 vs. 物理 RAM

朴素的分析会预期主机内存中占用 2× model_size：页缓存（1× 模型）+ 分片缓冲区（共 1× 模型）。但在 Ascend NPU 上，pin_memory() 通过 /dev/davinci_manager 分配，将分片放置在 cgroup 内存控制器不可见的 CPU 内核 DMA 内存中。物理 RAM ≈ 页缓存 + 固定分片 + 框架开销；cgroup 看不到固定 DMA 部分，但服务器仍需要那么多物理 RAM。

图 6：Cosmos3-Nano DP2 的 Ascend 内存核算。cgroup 看到共享页缓存和框架 RSS，而通过 /dev/davinci_manager 分配的固定分片驻留在驱动管理的 CPU DMA 内存中，而非 NPU HBM。

通过干净的测量验证（Cosmos3-Nano DP2，全新 cgroup）：

cgroup usage_in_bytes = 49 GB = cache(31) + rss(18)  ← exact match, no extra
cgroup kmem           = 0 GB
davinci_manager RSS   = 0 kB  (in /proc/<pid>/smaps)
NPU HBM per card      = 10 GB  (< 14.5 GB shard → shard NOT in HBM)
Slab                  = 3.3 GB  (too small for 29 GB shard)

这意味着 cgroup 可见内存按 O(model_size + dp_size × 常数) 扩展，而非 O(dp_size × model_size)——但总物理 RAM 是 cgroup 可见内存加上 cgroup 无法看到的固定 DMA 分片。对于 dp_size=4 的 200B 模型：约 423 GB cgroup + 约 400 GB 内核 DMA = 约 823 GB 总物理 RAM（可容纳于 2 TB），而未使用 mmap 时为 2000 GB。

验证结果

所有测试均在 Ascend 910B3（64 GB HBM/卡，2 TB 系统 RAM）、Cosmos3-Nano（33 GB）和 Cosmos3-Super（124 GB）上进行。

正确性

主机内存（cgroup 峰值）

NPU HBM

对于所测量的 dist_offload+SP 720p 10s 工作负载，峰值 HBM 从 Nano 到 Super 增长约 22%（23.1 → 28.1 GB），而空闲 HBM 增长约 27%（11.5 → 14.6 GB）。设备上仅驻留 2 层权重，因此大 3.8 倍的模型仍远低于 64 GB 上限。

性能

这些 Ascend 测量使用 Cosmos3-Nano，832×480，29 帧，35 个去噪步。生成帧/秒 是每墙钟秒产生的聚合输出视频帧数（29 frames × outputs per wave / wave latency），而非视频的播放帧率。

AllGather 开销 = 150 ms/步（72 ms 流切换 + 10 ms HCCL + 68 ms Python 调度），在 Cosmos3-Nano DP4 上测量。通信量随层维度、参与方数量和拓扑而变化。在 4 个并发请求下，这一固定成本被分摊 4 倍。

NVIDIA B300 GPU 结果

为了验证平台无关性，我们在 NVIDIA B300 SXM6 GPU 上运行了相同的 DLO 栈。下面的 Cosmos3 测试使用 Cosmos3-Super BF16（124 GB）、4× NVIDIA B300（物理 GPU 1、5、6、7）、Python 3.12.3、PyTorch 2.11.0+cu130、CUDA 13.0、vLLM 0.23.0 以及 vLLM-Omni 提交 9772bb32（PR #5397 的合并前快照；最终合并后的 head 包含此基准测试中不存在的后续 loader-gating 和 TP/mmap 验证更改）。随后的 MiniMax-H3 小节记录了其自身的 vLLM/vLLM-Omni 版本、enforce_eager=True 标志以及本地流水线补丁；这些细节适用于 MiniMax-H3 研究，不假定用于 Cosmos3 运行。

正确性通过所有策略间字节完全相同的输出哈希进行验证。例如，T2I 种子 42 在 DLO+AG、no-AG、DLO+USP4、legacy layerwise+USP4 和 HSDP+USP4 上产生了相同的 SHA256 6e7d2a8c63b88391...。T2V 832×480×29f 种子 17 在所有策略上产生了相同的 666,029 字节输出（SHA256 c5d38f5d21ca619e...）。

CUDA 进程树 PSS 包括共享页面缓存、固定 CPU 分片和框架内存。Ascend cgroup 测量排除了 /dev/davinci_manager 支持的固定分片，因此 GPU PSS 和 Ascend cgroup 数据不能直接比较。

1024×1024 T2I，50 步

DLO+AG DP4 在 4 个并发请求下达到 HSDP+USP4 吞吐量的 1.39×，而 HBM 使用量仅为后者的 30%（12.6 GiB 对 42.0 GiB）。

832×480 T2V，29 帧，35 步

工作负载延迟和 HBM（35 步，DLO+AG DP4 对 HSDP+USP4）

在 720p 10s（241f）上，DLO+AG+USP4 在 214.53s 内完成——与 HSDP 的 210.05s 相差在 2.13% 以内——输出字节完全相同（SHA256 08cb679322996ea6...），而 HBM 使用量仅为 HSDP 的 47%（24.99 GiB 对 53.73 GiB）。

8× B300 上的 MiniMax-H3：DLO 模式依赖于拓扑结构

另一项由 Shunyang Li 开展的 MiniMax-H3 B300 研究 测试了 DP、SP 和 DLO 执行模式在单台 8× NVIDIA B300 SXM6 AC 节点上的交互方式。与上述 Cosmos3 测量不同，此工作负载同时生成视频和音频：768×1344、124 视频帧、立体声音频、BF16、每个副本 batch size 为 1，请求 50 步（49 次调度器去噪更新）。该研究的 environment.json.txt 报告了 vLLM 0.24.0 和 vLLM-Omni 0.26.0rc2.dev11+g6607f4a7f（源提交 9e73ee1）；运行器设置了 enforce_eager=True（图编译已禁用），并对 pipeline_minimax_h3.py 应用了本地子组广播补丁。这些结果不代表未经修改的发布版本或默认的编译图路径。下面选定的每条 T2VA 路线包含两个引擎生命周期内各 20 个测量波次，每个生命周期前进行一次完整预热。吞吐量为输出数量除以波次时间；能耗对每个输出的八 GPU 板卡功率求和进行积分，且不减去空闲基线；外部 nvidia-smi 采样器以 0.758s 的中位间隔记录内存和功率。

图 7：三条评估路线内测得的服务前沿。增加 DP 以每波延迟换取并发输出容量；在 DP8×SP1 时，首选 DLO 模式从 AllGather 变为 rank-local。

配对的五波模式对比解释了为什么不存在单一的全局 DLO 策略。在 DP1×SP8 下，AllGather 使用 SP 组，将吞吐量提升 129.4%，同时将 P50 延迟降低 56.6%。在 DP4×SP2 下，其吞吐量收益缩小至 2.2%。在 DP8×SP1 下，AllGather 使吞吐量降低 4.1%，P50 延迟增加 3.8%，并将实测每 GPU 峰值从 20.03 提高到 94.03 GiB，因此 rank-local DLO 更受青睐。FL2VA 首帧和 Ref2VA 图像+音频测试保持了相同的延迟-吞吐量排序。

图 8：在三条评估路线中（每条路线 n=5 个测量波次），FL2VA 首帧 I2VA 和 Ref2VA 图像+音频改变了绝对延迟和吞吐量，同时保持了 DP1×SP8 → DP4×SP2 → DP8×SP1 的前沿排序。来源：MiniMax-H3 B300 研究工件。

这些结果是拓扑研究，而非通用的生产声明。DP2×SP4 未被测量；实验覆盖单节点、单一输入集、单一分辨率和帧数，且为形状验证而非感知质量。它使用了源提交 9e73ee1 加上记录的本地子组广播修复，运行时警告所测试的 vLLM-Omni 和 vLLM 版本未与发布版本对齐。存档提供了 PDF、CSV、105 个波次样本、环境哈希、本地差异和基准运行器 以供独立审查。

下表是基于上述实测内存模型的主机容量外推；并未实际运行 200B 级模型，该规模下的最大块大小、HBM 余量、带宽、延迟和输出质量仍未经验证。

致谢

我们感谢 vLLM-Omni 的贡献者，包括 @hsliuustc0106 和 @yuanheng-zhao 提供的详尽代码审查反馈，Shunyang Li（@lishunyang12）完成的 MiniMax-H3 B300 拓扑研究与可复现性工件，以及昇腾 NPU 团队提供的硬件支持。

参考文献

源代码：

• 分布式逐层卸载后端、meta 转换与 mmap 加载：distributed_layerwise_backend.py
• OffloadConfig 与策略选择：base.py
• 多队列执行器：multiproc_executor.py
• DP 多并发 worker：diffusion_worker.py
• 单元测试：test_distributed_layerwise_backend.py

RFC 与 PR：

• RFC：GitHub Issue #5396
• 实现 PR：vllm-omni#5397
• DLO DP 并发请求修复：vllm-omni#5864
• rank-local DLO DP 的独立请求：vllm-omni#5911

模型与基准测试工件：

• Cosmos3-Nano：33 GB safetensors（17B 参数，72 个 block）
• Cosmos3-Super：124 GB safetensors（64B 参数，128 个 block）
• MiniMax-H3：B300 DLO 研究笔记与可复现性工件