Large-Scale Sharded Weight Transfer With Ray Direct Transport (RDT) In vLLM

发表时间: 2026-08 · Blog post by Aaron Hao, Sumanth Hegde, Gal Meirom, Istvan Haller, Kourosh Hakhamaneshi, Gavin Parnaby, Moein Khazraee, Omri Kahalon (blog.vllm.ai)

原文: https://blog.vllm.ai/blog/2026-08-22-rdt-weight-transfer

作者/机构：Aaron Hao, Sumanth Hegde, Gal Meirom, Istvan Haller, Kourosh Hakhamaneshi, Gavin Parnaby, Moein Khazraee, Omri Kahalon

1. 主要贡献

在在线强化学习（RL）设置中，必须定期同步模型权重，以确保生成的rollout来自于最新的权重版本。随着开源模型的参数量扩展到万亿级别，高效的权重传输对于限制内存消耗和缩短传输时间变得至关重要。本文在vLLM中详细介绍了一种利用Ray直接传输（Ray Direct Transport, RDT）的分片权重传输实现。本文的核心贡献如下：

• 构建了vLLM原生的分片权重传输引擎：该引擎利用vLLM中的原生RL API，能够跨多种模型工作，包括稠密模型、具有融合或独立专家检查点的MoE模型以及量化模型。
• 为RL框架提供了简单的API接入：RL框架只需描述其权重的布局方式，传输引擎即可接管并完成整个传输过程。
• 实现了预处理与传输重叠的优化机制：通过优化实现，使得权重的收集（gather）、传输（transfer）和后处理（post-processing）能够相互重叠执行。
• 展示了容错的rollout能力：通过结合NIXL后端的RDT，演示了系统在长时间训练中具备的容错特性。

利用该引擎，本文在48个8xH100节点（32个训练节点，16个推理节点）上，仅用$7.53$秒就完成了Kimi K2模型在BF16精度下的分片权重传输。

2. 背景知识与设计原则

标准权重同步的局限性：标准的权重同步通常是一个NCCL广播过程。训练器（trainer）将每个参数全收集（all-gather）成HuggingFace格式，然后将其广播给每一个推理工作节点（worker）。对于中等规模的模型，这种方式是可行的，但随着模型规模的增长，它暴露出了明显的缺点。

大规模模型面临的内存与传输瓶颈：首先，每个工作节点都会接收到整个模型。例如在TP8（张量并行度为8）的情况下，一个工作节点只保留每个权重的$1/8$，并丢弃其余部分。对于像Kimi K2这样的大型MoE模型（通常在宽专家并行下部署），这种做法的劣势更为严重，因为每层的完整参数量仍然可能非常大（达到数十GB），这极大地损害了峰值内存以及传输速度。

集合通信的同步开销与脆弱性：其次，广播是一种集合通信（collective）操作。NCCL要求所有rank（进程）同步参与，这在动态场景中可能会产生问题。在大规模集群中，可能会出现落后的节点（straggler ranks）从而导致整个集合通信停滞，甚至可能发生副本故障（replica failures）。

本文设计的通用性原则：虽然之前已有关于大规模分片权重传输的研究，但本文的主要关注点在于两个维度上的通用性（generality）：一是跨模型和布局的通用性，即兼容vLLM支持的几乎所有模型；二是跨RL框架的通用性，允许其他RL框架轻松采用这种优化的权重传输实现。

3. 方法细节

3.1 vLLM中的权重加载机制

权重在vLLM中的生命周期：当一个新的HuggingFace格式权重张量到达vLLM worker时，必须经历以下七个操作步骤：1. 融合（Fuse）：合并权重分区，例如注意力层中的$Q$、$K$和$V$张量；2. 重新布局（Relayout）：根据原始权重格式对权重进行转置或重塑；3. 拆分/选择（Split/select）：对融合后的张量进行分块，或选择部分参数（例如专家并行）；4. 分片（Shard）：为了张量并行（TP）对权重进行切片；5. 拷贝至缓冲区（Copy Into Buffer）：将权重拷贝回为每层分配的缓冲区（“逐层缓冲区”），这些是vLLM在加载权重时分配的暂存缓冲区；6. 处理（Process）：可选择对权重进行量化，并执行一些特定于内核的操作，如填充（padding）、步幅（striding）等；7. 拷贝（Copy）：将最终处理好的权重拷贝到已分配的GPU内存中。

逐层重载机制：上述的步骤1至步骤5发生在vLLM的权重加载器中，通过逐层重载（layerwise reloading）来实现。逐层重载有助于确保权重更新能够保留CUDA图，同时保持内存使用量在有界范围内。

传输边界的确定：在理想情况下，在权重传输期间，我们从训练器传输最终处理好的权重（即步骤6之后的状态），并直接写入活动权重的存储空间中。然而，为了支持步骤6中广泛的后处理操作（如不同的量化方案），本文将重点放在传输BF16格式的、已分片但未处理的权重上（即步骤4之后的状态），并让引擎处理剩余的步骤。这使得系统能够兼容vLLM提供的各种量化方案。

自定义权重加载行为的挑战：将步骤1至步骤4转移到训练器端，意味着训练器必须知道，对于每一个工作节点和每一个权重，该节点最终将保留哪些确切的字节。

静态计算的不可行性：最直观的方法是通过计算得出结论：读取并行配置，计算出哪个张量的哪个区段属于哪个rank，然后据此发送。然而，要执行的确切操作可能会因层和模型的不同而有所差异。本文列举了两个例子：第一个是分组查询注意力（GQA）下的QKV融合。三个张量（q_proj、k_proj、v_proj）被融合为一个张量。在GQA下，KV头的数量可能少于TP rank的数量，因此两个工作节点可能会拉取不同的$Q$张量，但拉取完全相同的$K$和$V$张量。这与标准MHA（多头注意力）模型不同，后者$Q$、$K$和$V$张量的TP分片是一致的。第二个例子是Llama-4的融合专家。在Llama-4中，HuggingFace格式的专家张量被转置，拆分为gate_proj和up_proj，然后从中选择vLLM工作节点所需的专家。

动态记录解决方案：这两个例子说明了权重加载器可能具有的多种不同操作。由于vLLM支持各种架构，在训练器上为每个模型和每层实现步骤1至步骤4将涉及定制化的操作。避免这种定制化的唯一方法是在运行时记录给定配置下步骤1至步骤4的确切操作集。

“记录张量”空跑（Dry Run）机制：为了支持上述自定义权重加载行为，本文的解决方案如下：在引擎初始化时，我们向vLLM的加载器传递一个“记录张量（recording tensor）”——这是一个张量子类，它报告正确的形状和数据类型，但不包含任何实际数据。每一个变换操作（如view、narrow、transpose、reshape等）都会被追加到一个操作链中。当加载器将数据拷贝到参数中时，系统会记录它从哪里拷贝以及拷贝到了哪里。本文利用这个操作序列（即“分片计划”，sharding plan）在权重同步期间将训练器上的完整张量转换为vLLM工作节点所需的分片张量。

正确性保障与执行流程：因为这个计划直接来自于vLLM自身的加载器，所以无论这些加载器在不同层和模型中执行什么操作，该计划在构造上都是绝对正确的。因此，本文在训练器上执行步骤1至步骤4，并将BF16格式的分片权重传输给每个vLLM rank。在接收到分片权重后，我们在每个rank上执行剩余的步骤5至步骤7以更新活动的权重。

3.2 基于RDT的分片权重传输引擎

通信后端的选择：大多数流行的RL框架（如verl、SkyRL、Slime、NemoRL等）都使用Ray来编排训练，其中训练和推理rank通常作为独立的Ray actor进行管理。为了开发分片权重传输引擎，本文利用了Ray直接传输（Ray Direct Transport, RDT），这是一个允许Ray actor之间进行直接GPU到GPU通信的Ray API。RDT允许Ray actor方法返回GPU张量而无需将其拷贝出GPU。调用者会收到一个ObjectRef，当调用者读取它时，字节会通过可插拔的传输后端（如NIXL、NCCL、Gloo）进行移动。

基于NIXL的拉取式传输架构：在本文的案例中，选择NIXL后端以实现灵活的P2P通信，允许将自定义权重传输给每个消费者/推理rank。NIXL还提供了长训练周期所需的容错属性。由于RDT通过NIXL实现了基于拉取（pull-based）的传输，本文实现了一个基于拉取的权重传输引擎，推理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注册。

权重同步阶段的收集与拉取：1. 每个训练器rank一次收集一个权重组：一个权重组对应于一个Transformer块（注意力层 + MoE层）。系统一次全收集一层以最小化内存开销。作为可选项，系统可以选择利用权重局部性，仅收集特定的张量。在本文的集成中，系统仅跨TP进行收集，不跨PP阶段收集，并且也避免在EP（专家并行）下收集训练器rank上的专家。对于EP下的分布式专家，系统只需在初始化阶段将每个推理rank映射到具有所需专家的相关训练rank即可。2. 工作节点拉取分片权重：每个推理工作节点遍历其记录的计划，并向相应的训练器actor请求下一批切片。训练器actor对收集到的权重重放记录的操作，并将结果连续打包到其注册的RDT缓冲区中。然后，工作节点通过RDMA从该存储空间读取数据到自己的缓冲区中。

权重同步阶段的后台处理与清理：3. 工作节点在后台运行处理+拷贝操作：一个后台线程将每个切片从工作节点侧的RDT缓冲区拷贝到逐层缓冲区中，然后vLLM引擎运行处理+拷贝操作，以获得内核就绪格式的最终权重。4. 工作节点释放权重组：在处理完一个权重组的最后一个切片后，每个vLLM工作节点向拥有该权重的训练器rank发送信号。一旦每个vLLM工作节点都发送了信号，训练器就会丢弃该组收集到的张量，并可以自由地收集下一个张量。5. 训练器关闭同步：一旦没有正在传输的数据，并且工作节点完成了逐层重载，训练器就会关闭同步。

4. 实验环境

• 测试模型：小规模测试采用Qwen3-235B-A22B；大规模测试采用Kimi K2；容错测试采用Qwen3-32B。

• 硬件配置：

  • 小规模测试：4个8×H100节点（2个训练节点，2个推理节点）。
  • 大规模测试：48个8×H100节点（32个用于训练，16个用于推理）。网络连接采用InfiniBand（聚合带宽支持约1600 GB/s）。

• 软件配置：SkyRL、Megatron、vLLM、Ray（使用RDT与NIXL后端）。

• 并行策略参数：
  • Qwen3-235B测试：Megatron端采用TP4/PP2/EP8/ETP1，vLLM端采用DP16/EP16（匹配宽EP服务设置）。
  • Kimi K2测试：Megatron端采用TP8/PP8/EP32/ETP1，vLLM端采用TP32/EP32。

5. 实验结果

5.1 性能优化过程（基于Qwen3-235B-A22B）

本文记录了构建引擎的过程，并重点介绍了训练器上的几项重要性能优化。基线方案（SkyRL中的NCCL广播实现）在上述设置下的端到端权重同步时间为$64.72$秒。

• V1 - 简单迭代器（跨所有维度收集）：此版本使用一个简单的迭代器逐个遍历模型参数，并跨所有维度（TP、PP和EP）收集每个参数，生成HuggingFace格式的完整张量。结果：端到端权重同步时间为$25.02$秒。分析：其缺点在于1）收集操作涉及数以千计的微小集合通信（MoE检查点单独命名每个专家，Qwen3-235B有约$37,000$个张量，导致巨大开销）；2）每个rank都在重建完整张量，导致大量冗余内存使用。
• V2 - 优化迭代器（PP局部、EP局部）：此版本针对V1的缺点进行了修改。1）PP局部收集：一层的全收集仅在同一流水线阶段的rank之间运行。2）EP局部传输：根本不收集专家。训练器rank声明哪个rank持有哪个专家，推理rank直接从适当的rank中拉取。结果：端到端传输时间从$25.02$秒降至$5.61$秒。分析：这些优化对于像Kimi K2这样的大型模型尤为重要，不仅节省了传输时间，还节省了内存（Kimi K2一个完整的BF16 MoE层约为30GB，同步期间分配如此大的缓冲区极易导致OOM）。
• V3 - 流水线执行：在V2中，同步仍然顺序运行多个操作（全收集、重放操作和传输）。V3将这些使用不同资源的阶段进行了流水线化。1）训练器按权重组收集（以一个解码器块为单位）。2）训练器重叠收集和拉取：在推理rank拉取第$N$组时，训练器收集第$N+1$组。3）训练器重叠重放和传输：当一个块的RDMA数据落地时，生产者打包并在下一个块上运行重放操作。推理端也类似，在将张量从RDT缓冲区拷贝到逐层缓冲区时，并行接收下一个块。4）推理端后台处理：拷贝到vLLM分配的逐层缓冲区后，在后台调度处理+拷贝操作，此时RDT缓冲区可用于接收下一层的权重。结果：权重同步延迟进一步降至$3.49$秒。
  
  

5.2 最终结果（基于Kimi K2，48个节点）

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

• 实验内容：在32个训练节点和16个推理节点间同步Kimi K2权重。
• 实验结果：每次同步传输的字节数为$7.9$ TB，权重同步时间仅为$7.53$秒，实现的聚合带宽为$1,049$ GB/s。
• 分析结论：通过评估理论极限（SoL），在PP大小为8的设置下，考虑到vLLM逐层重载逻辑的限制（为避免OOM，必须在训练器PP组上串行传输），预期的SoL传输时间约为$5$秒。实测的$7.53$秒同步时间在预期SoL传输时间的$1.5$倍以内，证明了该引擎的高效性。

5.3 Rollout的容错性演示

• 实验内容：在4个8xH100节点上使用Qwen3-32B模型进行Text2SQL任务训练，部署4个推理副本。在第20步和第40步通过杀死一个推理引擎来模拟故障，并在几步后将其恢复上线。
• 实验结果：当推理引擎发生故障时，运行在降级状态下继续（路由器将流量路由到剩余的推理引擎）。训练器rank在下一次权重同步期间仅与活动引擎通信。副本恢复后，它会在下一个权重同步边界重新加入，接收更新的权重，并继续服务请求。
• 分析结论：证明了使用NIXL后端的RDT具备强大的容错能力。在广播集合通信中，单一rank故障会导致整个集合通信失败并需要重新初始化，而RDT+NIXL方案下训练照常进行，收敛性不受影响。
  

6. 结论

本文成功在vLLM中实现了一个原生的分片权重传输引擎，通过结合Ray Direct Transport (RDT) 和 NIXL 后端，解决了超大规模模型（如万亿参数的MoE模型）在RL训练中权重同步的内存与时间瓶颈问题。该方案不仅提供了跨模型架构和RL框架的通用性，还通过流水线重叠执行和局部化收集技术实现了接近理论极限的传输速度（在48节点上同步Kimi K2仅需7.53秒），同时赋予了系统在长周期训练中至关重要的容错能力。

局限性与未来展望：
该引擎目前仍处于早期阶段，存在以下局限性：1）加载器必须保持在可记录操作的范围内（例如，在加载期间检查真实值的加载器会在初始化时失败）；2）RDT目标缓冲区位于vLLM的gpu_memory_utilization预算之外，必须在选择该比例之前确定其大小；3）当前实现与vLLM中的EPLB不兼容；4）权重传输目前在训练器PP组之间是串行的以避免OOM，未来可以将其并行化到不同的副本以解决此问题；5）目前使用的是GPU到GPU的传输，最近Ray添加了对远程GPU到CPU传输的支持。未来可以利用这一特性避免在推理rank的GPU内存上分配额外的RDT缓冲区。此外，目前被迫在多个副本之间同步拉取同一工作节点的数据以避免额外的GPU缓冲区开销，如果简单地将模型副本存储在CPU内存上，这种限制也可以被消除。

7. 补充细节

7.1 与SkyRL的集成

配置覆盖机制：本文基于RDT的权重传输引擎已经集成到SkyRL中。要使用它，用户只需在配置中简单地使用以下参数覆盖：

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

WeightSource接口要求：为了让其他RL框架能够采用该引擎，在训练器端需要实现的主要接口是一个WeightSource迭代器。具体代码结构如下：

class WeightSource(ABC):
    def metadata(self) -> list[ParamMeta]: ...        # 名称，数据类型，完整形状 — 不涉及传输
    def __iter__(self): ...                           # 产出 (名称, 实体化张量)
 
    # 可选，针对分片训练器 — 声明“这个”rank持有内容：
    def held_names(self) -> "Collection[str] | None": ... # 产出哪些参数？

可选的held_names方法允许训练器准确定义特定rank持有哪些参数，从而启用V2版本中的优化（如EP局部传输）。