此权重传输引擎是vLLM原生的。任何基于Ray的训练器只需一个WeightSource迭代器即可采用它。此处工作出色！https://vllm.ai/blog/2026-08-22-rdt-weight-transfer…

使用 Ray Direct Transport (RDT) 在 vLLM 中进行大规模分片权重传输

来源：https://vllm.ai/blog/2026-08-22-rdt-weight-transfer

引言

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

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

• vLLM 原生的分片权重传输引擎，适用于一系列模型——稠密模型、使用融合或独立检查点的 MoE 模型，以及量化模型，并利用 vLLM 的原生 RL API (https://vllm.ai/blog/2026-05-28-native-rl-apis)。
• 为 RL 框架提供简单的采用 API，框架只需描述其权重布局，引擎则负责整个传输过程。
• 优化的实现，将预处理与传输重叠进行，使得数据收集、传输和后处理相互重叠。
• 容错的 rollout 演示，展示了使用 NIXL 的 RDT 的容错特性。

我们能够在 48 个 8xH100 节点（32 个训练节点，16 个推理节点）上，在 7.53 秒内完成 Kimi K2 模型在 BF16 格式下的分片权重传输。该实现在 vLLM (https://docs.vllm.ai/en/latest/training/weight_transfer/sharded_rdt/) 中提供，并在 SkyRL (https://github.com/NovaSky-AI/SkyRL/tree/main/examples/train/megatron/sharded_rdt) 中有一个端到端示例。

概述：基于广播的权重传输与使用 RDT (NIXL 后端) 的分片权重传输对比。使用 NCCL 时，训练器 rank 0 与所有推理 rank 形成集合通信组，并通过广播传输完整权重。使用分片权重传输引擎时，我们利用所有训练器 rank 参与传输，并且仅发送所需的分片。传输进一步优化以避免跨 PP rank 收集权重，并跳过收集专家层。 概述：基于广播的权重传输与使用 RDT (NIXL 后端) 的分片权重传输对比。使用 NCCL 时，训练器 rank 0 与所有推理 rank 形成集合通信组，并通过广播传输完整权重。使用分片权重传输引擎时，我们利用所有训练器 rank 参与传输，并且仅发送所需的分片。传输进一步优化以避免跨 PP rank 收集权重，并跳过收集专家层。

背景

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

1. 每个工作器接收整个模型：在 TP8 下，工作器保留每个权重的 1/8，并丢弃其余部分。对于像 Kimi K2 这样的大型 MoE 模型（通常采用宽专家并行部署）情况更糟，其中每层的完整参数可能仍然非常大（数十 GB），损害峰值内存和传输速度。
2. 广播是集合操作：NCCL 要求所有 rank 同步参与，这在动态场景中可能成问题。在大规模下，可能会有落后的 rank 阻塞集合操作，甚至发生副本故障。

虽然之前有大规模分片权重传输的工作（1 (https://www.lmsys.org/blog/2026-04-29-p2p-update/)，2 (https://research.perplexity.ai/articles/weight-transfer-for-rl-post-training-in-under-2-seconds)），但我们的主要关注点是通用性，体现在两个维度：

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

vLLM 中的权重加载

权重之旅

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

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

操作 1-5 在 vLLM 的权重加载器中进行，通过逐层重载 (https://docs.vllm.ai/en/latest/training/layerwise/)。逐层重载有助于确保权重更新在保持内存使用受限的同时保留 CUDA 图。

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

自定义权重加载行为

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

获取该信息的明显方式是计算：读取并行配置，计算出哪个张量的哪个带属于哪个 rank，并相应发送。然而，要执行的具体操作可能因层和模型而异。两个例子是：

1. 分组查询注意力下的 QKV 融合：三个张量 q_proj、k_proj、v_proj 融合成一个张量。在 GQA 下，KV 头的数量可能少于 TP 数量——因此两个工作器可能拉取不同的 Q 张量，但相同的 K 和 V 张量。这与标准 MHA 模型不同，后者的 TP 分片对于 Q、K 和 V 张量是一致的。
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 (https://www.ray.io/) 来协调训练，其中训练和推理 rank 通常作为单独的 Ray actor 进行管理。为了开发我们的分片权重传输引擎，我们利用了 Ray Direct Transport (https://docs.ray.io/en/latest/ray-core/api/direct-transport.html) (RDT)，一个允许 Ray actor 之间直接 GPU-GPU 通信的 Ray API。RDT 允许 Ray actor 方法返回 GPU 张量而无需将它们复制出 GPU。调用者接收一个 ObjectRef (https://docs.ray.io/en/latest/ray-core/objects.html)，当调用者读取它时，字节通过可插拔的传输层（NIXL、NCCL、Gloo）移动。在我们的情况下，我们选择 NIXL 后端以进行灵活的 P2P 通信，允许将自定义权重传输到每个消费者/推理 rank。NIXL 还提供了我们长时间训练运行所需的容错特性。由于 RDT 使用 NIXL 实现基于拉取的传输，我们实现了一个基于拉取的权重传输引擎，其中推理 rank 将从一个或多个映射的训练器 rank 拉取它们所需的分片张量。完整流程如下：

初始化时

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

初始化：训练器 rank all-gather 所有权元数据。Rank-0 将所有权 + 传输元数据传输到推理 rank。推理 rank 运行记录张量空运行以构建分片计划。所有 rank 分配并注册其 RDT 缓冲区以进行权重传输。 初始化：训练器 rank all-gather 所有权元数据。Rank-0 将所有权 + 传输元数据传输到推理 rank。推理 rank 运行记录张量空运行以构建分片计划。所有 rank 分配并注册其 RDT 缓冲区以进行权重传输。

权重同步期间

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

注意力层的权重同步：在一个训练器和一个推理 rank 上权重同步期间的操作概述。显示了注意力层的 Q、K 和 V 张量的权重传输。 注意力层的权重同步：在一个训练器和一个推理 rank 上权重同步期间的操作概述。显示了注意力层的 Q、K 和 V 张量的权重传输。 MoE 层的权重同步：在一个训练器和一个推理 rank 上权重同步期间的操作概述。显示了专家的权重传输。 MoE 层的权重同步：在一个训练器和一个推理 rank 上权重同步期间的操作概述。显示了专家的权重传输。

性能优化

我们记录了构建引擎的历程，并重点介绍了在训练器上的几个重要性能优化。

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

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

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

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

1. **收集包含数千个小型集合操作。**MoE 检查点为每个专家单独命名。Qwen3-235B 有 94 层 × 128 专家 × 几个投影 - 大约 37,000 个张量，其中大部分很小。一次收集它们会导致相当大的开销。
2. **每个 rank 收集所有内容。**在每个训练器 rank 上重建完整张量会导致大量的冗余内存使用。

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

V2 - 优化的迭代器：PP 本地，EP 本地

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

• PP 本地收集。一层的 all-gather 仅在同一流水线阶段的 rank 之间运行。
• EP 本地传输。