vLLM 使用 Ray Direct Transport 实现大规模分片权重传输

概述

vLLM 实现了一个原生的分片权重传输引擎，利用 Ray Direct Transport (RDT) 来优化在线强化学习 (RL) 设置中训练器 (trainer) 与推理工作节点 (inference worker) 之间的模型权重同步。该系统取代了传统的 NCCL broadcast 方法，后者在万亿参数规模下经常面临内存瓶颈和同步停顿问题，转而采用一种基于拉取 (pull-based) 的分片方法，从而降低了峰值内存使用并提高了传输速度。

基于广播的同步的局限性

标准的权重同步通常依赖于 NCCL broadcast，其中训练器将参数 all-gather 到 HuggingFace 格式并广播给每个推理工作节点。这种方法对于大规模模型带来了两个主要挑战：

内存效率低下：在诸如 Tensor Parallelism 8 (TP8) 的配置中，工作节点接收完整的模型，但仅保留 1/8 的权重，并丢弃其余部分。对于大型混合专家模型 (MoE)，这会产生巨大的峰值内存开销。
集体同步 (Collective Synchronization)：NCCL 要求所有 rank 参与同步。落后者 (straggler) rank 或副本故障可能会导致整个集体操作停顿，这对于动态、大规模环境来说是一个问题。

技术实现：分片权重传输

记录张量试运行 (The Recording Tensor Dry Run)

为了确保与各种模型架构的兼容性（例如 Llama-4 的 fused experts 或各种模型中的 GQA），vLLM 在初始化期间使用“记录张量”试运行。vLLM 的加载器被提供了一个张量子类，该子类报告形状 (shape) 和数据类型 (dtype) 但不包含数据。每一个转换操作——包括 views, narrows, transposes, 和 reshapes——都会被记录为一系列操作链（即“分片计划”）。

该计划允许训练器执行初始布局操作（fusion, relayout, splitting, 和 sharding），并仅传输每个 vLLM rank 所需的特定分片权重（BF16 格式），从而确保该过程在构建时即是正确的。

Ray Direct Transport (RDT) 与 NIXL

该引擎利用带有 NIXL 后端的 Ray Direct Transport (RDT) 来实现 Ray actor 之间的直接 GPU-to-GPU 通信。这种架构实现了一个基于拉取的系统，推理 rank 仅从映射的训练器 rank 中拉取所需的特定分片张量。

初始化流程包含五个步骤：

训练器 rank all-gather 所有权元数据（参数名称、dtypes 和 shapes）。
Rank 0 将此元数据和训练器 Ray actor 名称传输给推理工作节点。
每个 vLLM 工作节点通过记录张量试运行创建其分片计划。
工作节点以负载均衡的方式将自身映射到源训练器 rank。
生产者和消费者预先分配并注册 RDT 缓冲区。

性能优化

vLLM 对引擎进行了三个版本的迭代，以优化 Qwen3-235B-A22B 模型的端到端延迟（TP4/PP2/EP8 训练器到 DP16/EP16 vLLM 服务器）：

V1 (Simple Iterator)：逐个收集所有维度（TP, PP, EP）的参数。这导致了数千个微小的集体操作和冗余的内存使用，同步时间为 25.02s。
V2 (PP/EP-Local)：实现了 PP-local 收集和 EP-local 传输（专家模型不进行收集；推理 rank 直接从持有专家的 rank 中拉取）。这使得同步时间减少到 5.61s。
V3 (Pipelined Execution)：引入了 all-gather、replay 操作和 RDMA 传输的重叠 (overlapping)。通过在解码器块组中收集权重并进行后台处理，同步延迟降低到了 3.49s。

大规模验证：Kimi K2

在 48 个 8xH100 节点（32 个训练器节点，16 个推理节点）上的 Kimi K2 模型验证展示了以下结果：

Metric	Value
Bytes moved per sync	7.9 TB
Weight sync time	7.53s
Aggregate bandwidth	1,049 GB/s

考虑到 vLLM 的逐层重新加载逻辑的约束，这一性能大约是该特定设置下预期“光速” (SoL) 传输时间的 1.5 倍。

容错性与集成

通过使用 NIXL 而不是广播集体操作，该系统本质上对故障具有更强的韧性。如果推理引擎发生故障，路由器会继续将流量导向剩余的引擎，而训练器仅在下次同步时与在线的引擎进行通信。一旦故障副本被恢复，它将在下一个同步边界重新加入，并接收更新的权重，而不会影响整体收敛。

框架集成

该引擎已集成到 SkyRL 中。其他 RL 框架可以通过实现一个 WeightSource 迭代器来采用它，该迭代器提供参数元数据并产生（yields）实化的张量，并带有一个可选的 held_names 方法以启用 EP/PP-local 优化。

当前局限性

加载器约束：加载器必须使用可记录的操作；那些在加载过程中检查实际值的加载器将会失败。
内存预算：RDT 目标缓冲区存在于 vLLM 的 gpu_memory_utilization 预算之外。
兼容性：当前的实现与 vLLM 中的 EPLB 不兼容。
序列化：传输目前在训练器 PP 组之间是串行的，以防止在逐层重新加载期间发生 OOM。