强化学习训练后权重传输，不到 2 秒

我们近期实现了 Kimi-K2（1T 参数）1.3 秒的跨机器参数更新，将权重从 256 个训练 GPU（BF16）传输至 128 个推理 GPU（FP8）。

在异步强化学习微调中，训练和推理在不同的 GPU 上运行。在每个训练步骤之后，必须将新权重推送到推理节点。对于万亿参数模型，许多现有框架需要数秒甚至数分钟的时间。

通过利用 RDMA 点对点通信，我们能够在不改变推理引擎的情况下实现极速权重传输，并使代码更易于编写和维护。

我们的解决方案基于 RDMA WRITE，这是一种单向原语，源端直接写入目标端的 GPU 内存。

python

def rdma_write(src_ptr, dst_ptr, size, src_mr, dst_mr):
    # Write from local [src_ptr, src_ptr+size) to remote [dst_ptr, dst_ptr+size).
    # src_mr and dst_mr contains Memory Region metadata of both sides.
    ...

目标端甚至无需收到传输通知。这使得我们能够由训练节点驱动实现低延迟、高吞吐、零拷贝的传输，且无需在推理节点上执行任何控制逻辑。

元数据收集 – 控制器从所有训练和推理 GPU 收集参数元数据。

调度计算 – 控制器计算静态权重传输调度，映射哪个训练 GPU 向哪个推理 GPU 发送什么参数，以及发送顺序。

调度分发 – 控制器将调度发送给所有训练 GPU。

执行 – 在每个训练步骤之后，控制器向训练 GPU 发出开始传输的信号。

在定义了高层工作流之后，关键挑战在于如何以万亿参数规模高效执行权重传输。在此，我们详细说明执行路径。

训练中的参数根据 FSDP 配置进行分布。使用 full_tensor()，DeviceMesh 中的所有 GPU 均可重建完整参数，因此所有 GPU 都可以作为权重传输的源端。

多个不相交的 DeviceMesh 构成一个 网格组。由于同一组中的 DeviceMesh 是不相交的，它们的传输不会相互干扰，可以完全并行运行。我们在网格组之间插入全局屏障以强制执行顺序。

权重传输流水线重叠了主机到设备拷贝、张量组装、融合与量化，以及跨 MeshGroups 的 RDMA 传输。

我们将每个参数张量的传输视为一个任务。由于权重传输过程利用了多种硬件资源，我们将权重传输任务拆分为不同的流水线阶段，这些阶段在时间上重叠：

主机到设备内存复制（memcpy） — 如果 FSDP 将权重卸载到 CPU

参数准备 — 使用 full_tensor() 重建完整权重，应用投影融合，必要时进行量化。

RDMA 传输 — 零拷贝写入远程推理 GPU 内存

全局屏障 — 在所有 full_tensor() 调用完成后，通过以太网使用 GLOO 在网格组之间进行同步。

在实现中，我们为每个流水线阶段维护一个任务 FIFO 队列。每当队列头的任务完成该阶段时，它就会被移动到下一个阶段队列的尾部。

full_tensor() 及其他 GPU 操作会引入额外的 GPU 内存使用。为避免内存溢出错误，仅当当前正在进行的任务占用的临时 GPU 内存小于可配置的水位线时，我们才会启动任务执行。

几项设计选择使我们的系统运行速度显著加快，且比常见的开源解决方案更易于维护。

一种常见的模式是通过 rank-0 GPU 汇聚所有参数：在训练 rank-0 上收集，发送到推理 rank-0，然后再次分发。这很快会成为瓶颈，受限于单块 GPU 的 PCIe 带宽和网卡（例如 400 Gbps ≈ 50 GB/s）。

相比之下，我们的点对点设置允许每个训练 GPU 直接发送到每个推理 GPU，从而使整个网络结构达到饱和，而非仅仅占用单个链路。

基于 RPC 的传输使用推理侧缓冲区和 worker 调用，而单边传输则会直接写入权重，无需推理 worker 参与。

一些系统依赖于为每个张量调用推理引擎的 update_weight() 方法。这意味着对推理代码的侵入式更改，以及来自 RPC、序列化和控制平面协调的开销。

利用 RDMA WRITE 原语，我们在推理 GPU 内存中静默更新权重，无需额外拷贝。不涉及控制平面消息，也不涉及 CPU 控制逻辑。无需修改推理引擎。

权重传输过程可以利用四种硬件资源：（1）主机-设备数据移动（2）用于投影融合和量化的 GPU 计算（3）用于数据平面的 RDMA 网络（4）用于控制平面的以太网。

我们的设计将权重传输任务拆分为流水线阶段，允许在不同硬件资源之间轻松实现重叠。

一些实现在每个训练步骤都会重新计算传输调度，重复收集元数据并分发指令。这增加了不必要的控制平面延迟。

我们的调度在初始化时计算一次。每次训练迭代只需重放该计划：控制器发出“执行”信号，GPU 按照其预分配的路由运行。执行过程可预测且轻量。

很容易将整个权重更新过程纠缠在一个单体函数中：收集元数据、名称匹配、节点内汇聚、投影融合、量化、子切片通信域、跨节点网络传输。这不仅难以正确编程，而且更难优化。

在工程实践中，我们将这些步骤分离为独立的组件。每个组件都可以被单独进行单元测试、分析和优化。

快速、可靠的权重传输是大规模强化学习微调的关键构建模块。通过结合 RDMA WRITE 原语、静态传输调度和流水线执行，我们将 Kimi-K2 上的万亿参数更新缩减至仅 1.3 秒。该方法逻辑简单、易于维护，且避免了传统设计的瓶颈。

BibTeX

@online{perplexity2025WeightTransferForRlPostTrainingInUnder2Seconds,
  author = {{Perplexity} and Lequn Chen and Erin McKnight and Vivian Wang},
  title = {{Weight transfer for RL post-training in under 2 seconds}},
  date = {2025-09-24},
  url = {https://www.perplexity.ai/hub/blog/weight-transfer-for-rl-post-training-in-under-2-seconds}
}

RDMA WRITE：单向传输

高层工作流

权重传输执行

设备网格（DeviceMesh）与网格组

任务流水线

GPU 内存使用控制

为何它快速且简单

点对点通信

单向数据传输

流水线

静态调度

清晰的分离

结论

引用