vLLM 发布基于 Ray Direct Transport 的大规模分片权重传输引擎

引言

在在线强化学习（RL）设置中，模型权重必须定期同步，以确保 rollout 是基于最近的权重版本生成的。随着开源模型持续扩展到万亿级以上的参数量，高效权重传输对于限制内存消耗和传输时间变得至关重要。

在这篇博客中，我们详细介绍了 vLLM 中利用 Ray Direct Transport (RDT) 实现的分片权重传输。我们的贡献如下：

• vLLM 中的原生分片权重传输引擎，适用于多种模型——稠密模型、使用融合或按专家检查点的 MoE，以及量化模型，并利用了 vLLM 中的原生 RL API。
• 一个供 RL 框架采用的简单 API，框架只需描述其权重的布局方式，引擎便负责整个传输过程。
• 一种将预处理与传输重叠的优化实现，使得 gather、传输和后处理彼此重叠。
• 一个容错 rollout 演示，展示了 RDT 与 NIXL 的容错特性。

我们能够在 48 个 8xH100 节点（32 个节点用于训练器，16 个用于推理）上，以 BF16 精度在 7.53 秒内完成 Kimi K2 模型的分片权重传输。该实现已在 vLLM 中提供，并附有 SkyRL 中的端到端示例。

背景

标准的权重同步是 NCCL 广播。训练器将每个参数 all-gather 成 HuggingFace 格式，并将其广播给每个推理 worker。对于中等规模的模型来说这没问题，但随着模型增长，它有以下缺点：

1. 每个 worker 接收整个模型：在 TP8 下，一个 worker 保留每个权重的 ⅛ 并丢弃其余部分。对于像 Kimi K2 这样的大型 MoE 模型（通常部署在宽专家并行下）来说更糟，因为每层的完整参数仍然可能相当大（数十 GB），从而影响峰值内存和传输速度。
2. 广播是一种集合操作：NCCL 要求所有 rank 同步参与，这在动态场景中可能存在问题。在大规模下，可能会出现拖后腿的 rank 导致集合操作停滞，甚至副本故障。

虽然之前已有大规模分片权重传输的工作（1、2），但我们的主要关注点是在两个维度上的通用性：

• 跨模型和布局，与 vLLM 支持的几乎任何模型兼容。
• 跨 RL 框架，允许其他 RL 框架采用优化后的权重传输实现。

vLLM 中的权重加载

权重的旅程

当 HuggingFace 格式的新权重张量到达 vLLM worker 时，它必须经过以下操作：

1. 融合：权重分区被融合，例如注意力层中的 Q、K 和 V 张量
2. 重排：根据原始权重的格式，权重可以被转置或重塑
3. 拆分/选择：融合后的张量可以被分块，或者可以选择参数子集（例如专家并行）
4. 分片：权重可被切片以用于张量并行
5. 复制到缓冲区：权重被复制回按层分配的缓冲区（“逐层缓冲区”）。这些逐层缓冲区是 vLLM 在权重加载期间分配的暂存缓冲区。
6. 处理：权重可选地进行量化，并伴随一些内核特定的操作，如填充、跨步等。
7. 复制：最终处理后的权重被复制到已分配的 GPU 内存中

操作 1-5 在 vLLM 的权重加载器中通过逐层重载进行。逐层重载有助于确保权重更新保留 CUDA 图，同时保持内存使用受限。

理想情况下，在权重传输期间，我们从训练器传输最终处理后的权重（步骤 6 之后），并直接写入实时权重的存储中。然而，为了支持步骤 6 中广泛的后处理操作，我们专注于以 BF16 格式传输分片但未处理的权重（即步骤 4 之后），并让引擎处理其余部分。这使我们还能支持 vLLM 中的不同量化方案。

自定义权重加载行为

将步骤 1-4 移至训练器意味着训练器必须知道，对于每个工作进程和每个权重，该工作进程最终会保留哪些字节。

获取该信息的显而易见的方法是计算：读取并行配置，确定哪个张量的哪个部分属于哪个 rank，并相应地发送。然而，要执行的确切操作可能因层以及模型而异。两个例子是：

1. 分组查询注意力下的 QKV 融合：三个张量（q_proj、k_proj、v_proj）被融合为一个张量。在 GQA 下，KV 头数可能少于 TP rank 数——因此两个工作进程可以拉取不同的 Q 张量，但 K 和 V 张量相同。这与标准 MHA 模型不同，后者中 Q、K 和 V 张量的 TP 分片是一致的。
2. Llama-4 的融合专家：在 Llama-4 中，HuggingFace 格式的专家张量被转置，拆分为 gate_proj 和 up_proj，vLLM 工作进程的专家从中选择。

这两个例子说明了权重加载器可能具有的多样化操作集。鉴于 vLLM 支持的各种架构，在训练器上实现步骤 1-4 将涉及针对每个模型和层的定制操作。避免这种情况的唯一方法是在运行时记录给定配置下步骤 1-4 的确切操作集。

解决方案：“记录张量”试运行

为了支持上述自定义权重加载行为，我们的解决方案如下：在引擎初始化时，我们将一个“记录张量”——一个报告正确形状和数据类型但不拥有数据的张量子类——交给 vLLM 的加载器。每个变换——视图、窄化、转置、重塑等——都会被追加到一个操作链中。当加载器复制到参数时，我们记录它从何处复制以及复制到了何处。我们在权重同步期间利用这个操作序列（“分片计划”）将训练器上的完整张量转换为 vLLM 工作进程所需的分片张量。

由于该计划来自 vLLM 自身的加载器，因此对于这些加载器在不同层和模型上的任何操作，它都是构造正确的。

因此，我们在训练器上执行步骤 1-4，并以 BF16 格式将分片权重传输到每个 vLLM rank。在接收到分片权重后，我们执行剩余的步骤 5-7 以更新每个 rank 上的实时权重。

使用 RDT 的分片权重传输引擎

大多数流行的 RL 框架，如 verl、SkyRL、Slime、NemoRL 等，都使用 Ray 来编排训练，训练和推理 rank 通常作为独立的 Ray actor 进行管理。为了开发我们的分片权重传输引擎，我们因此使用了 Ray Direct Transport（RDT），这是一个允许 Ray actor 之间直接进行 GPU-GPU 通信的 Ray API。RDT 允许 Ray actor 方法返回 GPU 张量，而无需将其从 GPU 上复制下来。调用方会收到一个 ObjectRef，当调用方读取它时，字节会通过可插拔的传输方式（NIXL、NCCL、Gloo）移动。在我们的场景中，我们选择了 NIXL 后端以实现灵活的 P2P 通信，从而允许将自定义权重传输到每个消费者/推理 rank。NIXL 还提供了我们长时间训练运行所需的容错特性。由于 RDT 使用 NIXL 实现了基于拉取的传输，我们实现了一个基于拉取的权重传输引擎，其中推理 rank 会从一个或多个映射的训练器 rank 拉取它们所需的分片张量。完整流程如下：

初始化时

1. 训练器收集所有权元数据：训练器报告每个参数的元数据——名称、dtype 和完整形状——以及训练器布局：每个 rank 上存在哪些层（流水线并行）和哪些权重名称（例如，专家并行下的一部分专家参数）。训练器 rank 对这份所有权元数据进行 all-gather。
2. Rank 0 将传输元数据发送给推理 worker：Rank 0 发送参数和所有权元数据，以及 RDT 传输所需的训练器 Ray actor 名称。
3. 每个 vLLM worker 记录其分片计划：每个 vLLM rank 将执行上述的 recording-tensor 试运行，为每个参数创建由操作链组成的分片计划。
4. 每个 vLLM worker 构建源训练器 rank 的映射：利用传输元数据，每个 vLLM worker 构建源训练器 rank（持有它所需的参数）的映射以及要运行的分片计划。当多个训练器 rank 持有给定参数时，vLLM worker 会以负载均衡的方式选择一个训练器 rank。vLLM worker 将分散到给定参数的可用生产者之间，并且来自不同副本的相同 worker rank 会从同一生产者拉取，以减少内存开销并缩短传输时间。
5. 双方分配并注册各自的 RDT 缓冲区。消费者的目标缓冲区和生产者的源缓冲区会一次性分配，并预先向 NIXL 注册。

权重同步期间

1. 每个训练器 rank 一次收集一个权重组。一个权重组对应一个 transformer 块（注意力 + MoE 层）。我们一次 all-gather 一层，以最小化内存开销。可选地，我们可以选择利用权重局部性，只收集特定的张量。在我们的集成中，我们仅在 TP 维度上进行 all-gather。我们不跨 PP 阶段进行 all-gather，也避免在 EP 下于训练器 rank 上收集专家。对于 EP 下的分布式专家，我们只需在初始化阶段将每个推理 rank 映射到持有所需专家的相关训练 rank。
2. Worker 拉取分片权重。 每个推理 worker 遍历其记录的计划，并向对应的 trainer actor 请求下一批切片。trainer actor 针对收集到的权重重放记录的操作，并将结果连续打包到其注册的 RDT 缓冲区中。随后 worker 通过 RDMA 从该存储读取到自己的缓冲区。
3. Worker 在后台运行 process + copy。 后台线程将每个切片从 worker 侧的 RDT 缓冲区复制到逐层缓冲区，然后 vLLM 引擎运行 process + copy，以获取内核就绪格式的最终权重。
4. Worker 释放权重组。 在完成某个权重组的最后一个切片后，每个 vLLM worker 向拥有该组的 trainer rank 发出信号。一旦所有 vLLM worker 都发出信号，trainer 就会丢弃该组已收集的张量，并可自由收集下一个。
5. trainer 关闭同步，一旦没有在途操作，worker 便完成逐层重新加载。

性能优化

我们记录了构建该引擎的历程，并重点介绍 trainer 上的一些重要性能优化。

为此，我们将使用 SkyRL 中结合 Megatron 和 vLLM 对 Qwen3-235B-A22B 进行权重同步的小规模设置。训练在 4 个 8×H100 节点上进行——两个 trainer 节点和两个推理节点，Megatron 并行度为 TP4/PP2/EP8/ETP1，vLLM 以 DP16/EP16 提供服务，以匹配宽 EP 服务设置。报告的权重同步数字是端到端延迟，包括 all-gather 权重提取，在多次权重同步中取平均，排除第一次冷启动迭代。

作为基线，SkyRL 中的 NCCL 广播实现在相同设置下耗时 64.72 秒。下面，我们重点介绍分片权重传输引擎不同版本的性能，聚焦于我们在 trainer 上如何收集、迭代和传输模型参数。其他一切保持与之前描述相同——trainer 到推理 rank 的映射、记录张量试运行等。

V1 - 一个简单的迭代器（跨所有维度收集）

在这种情况下，我们使用一个简单的迭代器，逐个遍历模型参数，并跨所有维度（TP、PP 和 EP）收集每个参数，并以 HuggingFace 格式生成完整张量。这种方法有两个缺点：

1. 收集过程包含数千个微小的集合通信。 MoE 检查点将每个专家单独命名。Qwen3-235B 有 94 层 × 128 个专家 × 若干投影——大约 37,000 个张量，其中大多数很小。一次收集一个会带来相当大的开销。
2. 每个 rank 都收集所有内容。 在每个 trainer rank 上重建完整张量会导致大量冗余内存使用。

使用这种方法，在上述 Qwen3-235B-A22B 设置下，端到端权重同步时间为 25.02 秒。

V2 - 优化的迭代器：PP 局部、EP 局部

在这种情况下，我们解决了 V1 的两个主要缺点，并按如下方式更改迭代器：

• PP 局部收集。 某一层的 all-gather 仅在同一流水线阶段内的 rank 之间运行。
• EP 本地传输。 专家完全不聚集在一起。训练器不是重新组装 MoE 层中的所有专家，而是让各个 rank 声明哪个 rank 持有哪个专家，推理 rank 则从相应的 rank 拉取。

这些优化对于像 Kimi K2 这样更大的模型尤为重要，不仅是为了节省传输时间，也是为了节省内存：Kimi K2 的完整 MoE 层采用 BF16 格式时约为 30GB。在权重同步期间为每个 GPU 分配如此大的缓冲区很容易导致 OOM。

通过上述优化，端到端权重传输时间从 25.02 秒降至 5.61 秒。请注意，还有一些额外的优化（如元数据缓存）对传输时间影响较小。更多细节见此处。

V3 - 流水线执行

在 V2 中，同步仍然按顺序运行多个操作：all-gather、重放操作和传输。这三个阶段使用不同的资源，可以流水线化。

• 训练器：按权重组聚集。 权重作为一个解码器块聚集。这使得块成为聚集、传输和释放的单位。
• 训练器：重叠聚集和拉取。 训练器在推理 rank 仍在拉取组 N 时聚集组 N+1。
• 训练器：重叠重放和传输。 当一个块的 RDMA 正在落地时，生产者打包并对下一个块运行重放操作。类似地，在推理侧，可以在将张量从当前 RDT 块复制到逐层缓冲区的同时，并行接收下一个块。
• 推理：后台处理： 将权重从 RDT 缓冲区复制到 vLLM 引擎分配的逐层缓冲区后，我们将 Process + Copy 操作（步骤 6 和 7）调度为在后台运行。现在 RDT 缓冲区可用于接收下一层的权重。

通过额外的流水线化，权重同步延迟从 5.61 秒降至 3.49 秒。

最终结果：48 个节点上的 Kimi K2

NIXL 团队在 48 个 8×H100 节点上使用 Kimi K2 验证了权重同步。

训练器设置：Megatron 使用 TP8/PP8/EP32/ETP1
推理设置：vLLM 使用 TP32/EP32

我们进一步估计了最佳理论权重传输时间。该设置的绝对光速（SoL）将是通过网络发送权重的传输时间。训练器占用 32 个节点，每个推理副本占用 4 个节点。PP 大小为 8 时，每个由 4 个节点组成的 PP 组需要向 4 个副本发送约 2TB/8 = 0.25TB 的权重，因此每个组需要从 4 个节点发送约 1 TB 的权重。类似地，每个由 4 个节点组成的推理副本需要接收 2TB 的权重。因此，我们可以通过关注一个推理副本来估计光速。

要传输的字节数 = 2TB 权重
总带宽：400*4 GB/s = 1600 GB/s（使用 InfiniBand）

因此，绝对 SoL 约为 1.25s。然而，由于 vLLM 中的逐层重载逻辑，目前我们只能通过 trainer PP 组进行串行传输。每一层在 GPU 内存上分配了单独的缓冲区，从 PP 组并行传输很容易导致 OOM。因此，对于传输的合理预期 SoL，我们应该转而关注发送端。聚焦于一个 PP 组的传输时间，我们得到每个 PP 组约 0.625s。在 trainer PP 大小为 8 的情况下，此设置下的预期 SoL 为 0.625*8 = 5s。在 7.53s 时，测得的权重同步时间在此设置下预期 SoL 传输时间的约 1.5 倍以内。

rollout 的容错能力

使用 NIXL 的主要优势之一是能够处理故障。使用广播集合通信时，如果组中的某个 rank 发生故障，整个集合通信可能会失败，并且集合通信组需要重新初始化。

为了突出 RDT 的优势，我们展示了 SkyRL 中推理引擎故障的一个场景。当推理引擎发生故障时，运行会继续，但处于降级状态：路由器将流量路由到剩余的推理引擎。trainer rank 在下一次权重同步时仅与存活的引擎通信。副本恢复后，它会在下一个权重同步边界重新加入，接收更新后的权重，并继续处理请求。

与 SkyRL 集成

我们基于 RDT 的权重传输引擎已集成到 SkyRL 中。要使用它，您只需使用以下覆盖配置：

generator.inference_engine.weight_sync_backend=sharded_rdt \
trainer.placement.colocate_all=false

对于其他要采用该引擎的 RL 框架，在 trainer 侧要实现的主要接口是一个 WeightSource 迭代器。

class WeightSource(ABC):
    def metadata(self) -> list[ParamMeta]: ...        # names, dtypes, full shapes — no transfer
    def __iter__(self): ...                           # yield (name, materialized tensor)
 
    # Optional, for sharded trainers — declare what THIS rank holds:
    def held_names(self) -> "Collection[str] | None": ... # which params are yielded?

可选方法 held_names 允许 trainer 精确定义特定 rank 持有哪些参数，从而实现 V2 中的优化。

局限性与后续工作

使用 RDT 的分片权重传输引擎仍处于早期阶段。一些局限性包括：

• Loader 必须保持在可记录的操作范围内。例如，在加载期间检查真实值的 loader 会在初始化时失败
• RDT 目标缓冲区位于 vLLM 的 gpu_memory_utilization 预算之外，必须在选择该比例之前确定大小。
• 当前实现与 vLLM 中的 EPLB 不兼容。
• 权重传输目前在 trainer PP 组之间是串行的，以避免逐层重载导致的 OOM。可以将传输跨 PP 组并行到不同副本以避免此问题。
• 我们目前使用 RDT 进行 GPU -> GPU 传输。对使用 RDT 进行远程 GPU -> CPU 传输的支持最近已添加。我们可以利用远程 GPU -> CPU 传输来避免在推理 rank 的 GPU 内存上分配额外的 RDT 缓冲区。此外，我们被迫同步来自同一 worker 跨多个副本的拉取，以避免在 GPU 上为每个副本分配单独缓冲区的额外开销。如果我们简单地将模型副本存储在 CPU 内存中，也可以避免这种情况。

致谢

这项工作是与 NIXL 团队合作完成的，他们推动了在 Kimi K2 上的大规模验证，并提供了许多有用的技巧来提升权重传输性能。

感谢 Josh Lee 和 Stephanie Wang 对 RDT 的指导，以及 vLLM 团队（尤其是 Ao Shen）提供的宝贵审阅意见。