Appendix G 用于 RL 训练的引擎切换流水线

本文为论文《Kimi K2：面向智能体能力的开源大模型》「中文全文翻译」板块的对应章节，与「论文原文」板块同一章节一一对应，层级保持一致。

(a) 理论上的完美三阶段流水线权重更新

(b) 一个受 PCIE 限制的三阶段流水线

(c) 固定的两阶段流水线

图 13：RL 权重更新流水线

检查点引擎在每个 GPU 上管理三个等大小的设备缓冲区：一个用于加载卸载模型参数的 H2D 缓冲区，以及两个用于 GPU 间广播的 IPC 缓冲区。 这些 IPC 缓冲区与推理引擎共享，使其能够直接访问相同的物理内存。 这三个缓冲区使我们能够将三个步骤安排在一个流水线中。

理论上的三阶段流水线。

如图 13(a) 所示，引入了一个三阶段流水线。 （1）H2D：最新权重的一个分片被异步复制到 H2D 缓冲区。 （2）Broadcast：复制完成后，该分片被复制到一个 IPC 缓冲区并广播到所有设备。 （3）Reload：推理引擎同时从另一个 IPC 缓冲区加载参数。

因 PCIe 饱和而产生的两阶段流水线。

在 NVIDIA H800 集群上，并发的 H2D 和广播会使共享的 PCIe 总线饱和，从而将三个阶段压缩为一个顺序过程（图 13(b)）。 因此我们采用了一个更简单的两阶段方案（图 13(c)）：（1）所有设备执行一次单一的同步 H2D 传输。（2）广播和加载并行进行。

两阶段流水线将受多个同步 H2D 复制操作的限制。但在大规模设备上，模型会被拆分为较小的分片，整个参数集在一次传输中即可装入 H2D 缓冲区，此时开销将消失。

通过重叠 H2D、Broadcast 和 Reload 权重，我们能够获得高带宽，以将权重从训练引擎重分片到所有推理引擎。