GPU 加速的权重传送（Weight Teleportation，约 17 分钟阅读）

# GPU 加速的权重瞬移 Source: https://humansand.ai/blog/gpu-accelerated-weight-teleportation ## 引言\# (https://humansand.ai/blog/gpu-accelerated-weight-teleportation#introduction) 大规模 RL 与大规模预训练或大规模推理都不同。一个关键区别是，训练权重与采样器权重理想情况下应尽可能保持一致，这就要求训练器向采样器发送权重更新。由此产生了一个瓶颈：要么使用高速集群内网络发送权重更新，要么接受尝试传输数 TB 模型权重所带来的往往不可忽视的延迟，更不用说还要阻塞等待采样器了。我们引入了一种更快的权重更新范式，在大多数情况下，无论采样器与训练器是否位于同一集群内，都能获得优于集群内网络的速度。 该设计从被广泛采用的 delta 权重同步思想出发¹ Miles v0\.1: Production\-Level Post\-Training, §4\.3 Disk\-Delta Updates (https://arxiv.org/html/2609.08368v1#S4.SS3)\.。我们共同设计了一整条 Blackwell 原生的压缩/解压流水线，结合了以下技术： 1. 大规模并行的 GPU 压缩，具备分布式所有权与容错能力。 2. 精心设计的双重压缩算法，配合异构的 CPU\-GPU 解压。 3. Blackwell 解压引擎（DE）与硬件原生解压算法，将 DE 与普通 GPU 计算完全重叠。 4. 训练器与采样器、CPU 与 GPU 之间极小的内存缓冲开销。 结果，在我们的 Ampersand 生产级 RL 配置中，采样器暂停时间降到 0\.23 s 以下，端到端权重同步时间降到 4 s 以下，充分压满了 Blackwell 解压引擎流水线，使解压 ASIC 成为采样器暂停时间的瓶颈。无论从端到端权重同步时间还是采样器暂停时间来看，该实现都足够快，以至于对大多数真实工作负载而言，集群内 P2P RDMA 反而更慢。 在本文余下部分，我们使用以下内容作为一致的工作负载代理。模型是 NVFP4 GLM\-5\.3，其规范权重张量共 445 GB。RL 训练信号是稀疏的，每次训练步骤只有极小一部分权重被更新，因此我们修改检查点的 0\.085% 作为一次稳态更新的代理。训练器由 8 个节点、64 块 B300 GPU 组成。每个 rollout 接收器在一个 8 GPU B300 节点上运行两个 TP4 NVFP4 GLM\-5\.3 引擎。为清晰起见，图中使用了更少的 GPU 和模型层数。 ## CPU 磁盘 delta 权重同步\# (https://humansand.ai/blog/gpu-accelerated-weight-teleportation#cpu-disk-delta-weight-sync) 要解释我们的设计演进，我们需要理解 delta 权重同步的基础，以及为什么朴素的实现并不足够。 ### Delta 更新代数\# (https://humansand.ai/blog/gpu-accelerated-weight-teleportation#delta-update-algebra) `` # trainer mask = weights_old ⊕ weights_new # rollout weights_updated = weights_old ⊕ mask = weights_new ``` delta 权重同步的核心是按位 `XOR`（`⊕`）操作。在训练步骤之前，rollout 引擎与训练器持有相同的权重。训练器完成一个训练步骤，从 `weights_old` 前进到 `weights_new`。训练器对 `weights_old` 与 `weights_new` 做按位 `XOR`，得到一个与权重形状完全相同的掩码。掩码中的 `1` 表示该位置的权重位发生了翻转，`0` 表示未改变。rollout 引擎需要 `weights_new` 才能基于新策略为下一步 RL 采样。为了得到它们，引擎只需翻转相同的权重位即可。因此，我们将掩码传给 rollout 引擎，并在其 `weights_old` 与掩码之间做按位 `XOR`。由于权重被视为一长串连续的比特，其具体数据类型并不重要，也不涉及浮点运算或舍入。 但掩码与权重大小相同，因此直接传输掩码相比传输完整权重并不能节省任何东西。真正让 delta 同步方法有吸引力的是下面这个洞察。在 RL 中，训练信号是稀疏的，每次训练步骤只有极小一部分权重被更新。这意味着掩码中绝大部分是 `0`，只有极小一部分是 `1`。无损压缩算法随后可以将这些掩码压缩成极小的传输负载。 ### 作为交接介质的模型检查点\# (https://humansand.ai/blog/gpu-accelerated-weight-teleportation#model-checkpoint-as-a-handoff) 现在我们来看 RL 循环中一个朴素但清晰的 delta 权重同步设计。训练器与 rollout 引擎需要同步它们的检查点，所以让我们用模型检查点作为交接介质。训练器执行一个训练步骤，计算掩码，将其压缩成一个小负载，然后把压缩后的掩码发送给 rollout 引擎。rollout 引擎加载并解压掩码，修补检查点以匹配新的训练器权重，然后从该检查点重新加载权重。由于检查点就是交接介质，rollout 引擎可以复用与服务器启动相同的代码路径。我们无需处理 rollout 引擎如何重新分片或映射权重。 Miles¹ Miles v0\.1: Production\-Level Post\-Training, §4\.3 Disk\-Delta Updates (https://arxiv.org/html/2609.08368v1#S4.SS3)\. 中现有的磁盘 delta 实现在很大程度上遵循这一设计： 从上图的测量来看，压缩后的负载仅为 0\.876 GB，符合我们的预期。然而，这一收益已被其他显著的开销完全抵消： - 训练器压缩带来额外延迟。 - 检查点修补需要一次磁盘往返。 - rollout 引擎需要从磁盘重新加载整个检查点。 ## GPU 加速的（解）压缩\# (https://humansand.ai/blog/gpu-accelerated-weight-teleportation#gpu-accelerated-de-compression) RL 技术栈已经运行在分布式 GPU 集群之上，因此直观的下一步是让这条流水线 GPU 加速，利用其强大的并行计算能力和高带宽 HBM。 ### 借助 Blackwell 原生 DE 构建强基线\# (https://humansand.ai/blog/gpu-accelerated-weight-teleportation#a-strong-baseline-with-blackwell-native-de) Blackwell 具有一个固定功能的解压引擎（DE）² Decompression Engine FAQ (https://docs.nvidia.com/cuda/nvcomp/decompression_engine_faq.html)\.，能极其高效地执行某些受支持的解压算法。为了使我们的基线足够强，我们直接将 GPU 加速流水线与 GPU 原生（解）压缩算法以及 Blackwell 原生专用硬件进行协同设计。 下表比较了 GPU 上的（解）压缩算法³ NVIDIA nvCOMP Quick Decision Matrix (https://docs.nvidia.com/cuda/nvcomp/algorithm_selection.html#quick-decision-matrix)\, 定性评级。： *表 1。GPU 编解码器的权衡与解压引擎支持情况。* 如上文 Fig\. 4 所示，主要开销位于接收端，包括解压、检查点修补和引擎重载。因此，我们优先优化接收端的解压性能，从 LZ4/Snappy 开始构建快速的 GPU 编解码流水线。 LZ4 与 Snappy 来自同一算法家族⁴ LZ4 与 Snappy 都是面向字节的 LZ77 编解码器：两者都编码字面量与反向拷贝，没有熵编码阶段。参见 Snappy format (https://github.com/google/snappy/blob/main/format_description.txt) 与 LZ4 block format (https://github.com/lz4/lz4/blob/dev/doc/lz4_Block_format.md)\.，它们在拷贝命令的编码方式上有一个显著差异。单个 Snappy 拷贝命令最多覆盖 64 字节，而 LZ4 可以用更少的拷贝命令直接编码更长的区域。由于 RL 更新掩码大多由长段的零填充，LZ4 往往比 Snappy 将最终负载压缩得更紧凑。请注意，具体结果可能因实现而异。我们的实验结果支持这一直觉。在同一个代理掩码上，nvCOMP LZ4 比 Snappy 压缩得更紧凑。LZ4 基线产生 5\.36 GB 的完整发布量。因此，我们为该设计选择 LZ4 作为编解码器。 既然我们已决定使用 GPU 进行解压，解压后的掩码已经驻留在 GPU HBM 中，为检查点修补再付出一次磁盘往返并不合理。因此，我们实现了 Triton kernel，直接在 GPU HBM 中原地执行按位 `XOR`。模型层与解压掩码之间的分块（tiling）和地址映射可能相当复杂，因此我们在 CPU 上准备这些元数据，让 Triton kernel 从预先准备好的计划中读取。结果，整个模型的 `XOR` 应用阶段可以由单次 Triton kernel 启动极其高效地完成。 在训练器一侧，我们利用分布式训练已有的全部 GPU 并行压缩掩码。分布式训练本身就意味着分布式的张量所有权，但它仍可能重复保存某些模型权重，尤其是在使用上下文并行（CP）和数据并行（DP）时。对于压缩，我们重新映射权重所有权，使所有 GPU rank 恰好整体覆盖一次完整的模型检查点，每个 rank 只拥有其为分布式训练所持权重的一个子集。这带来三大优势： 1. 通过在现有训练 GPU 上并行化，压缩速度大幅提升。 2. 权重聚合流量大幅减少，因为我们只需聚合 rank 本地的压缩负载，而非原始模型权重。 3. 在所有现有 rank 上，持有旧版本权重所带来的内存压力大幅缓解。 将设计整合起来：所有训练器 GPU 并发压缩，rollout GPU 直接使用 LZ4 通过 DE 解压，单次 Triton kernel 启动即可更新 HBM 中的所有模型权重。所有权重同步的网络流量都携带压缩负载，磁盘检查点完全退出热路径。这已经是一个极其强大的基线：对于相同的代理负载，训练器侧的准备与压缩耗时 0\.834 s，接收端的工作耗时 6\.041 s。 ### 与低精度 RL 协同\# (https://humansand.ai/blog/gpu-accelerated-weight-teleportation#working-with-low-precision-rl) 由于压缩、解压和按位 `XOR` 都把模型权重视为连续的比特缓冲区，delta 权重同步对低精度 RL 也能无损工作。 低精度权重甚至可以吸收小的高精度更新，例如低位尾数位的变化，从而使 delta 负载更小。 一个有趣的细节是如何与高性能 MoE kernel 协同，比如我们在*Faster\-than\-light NVFP4 decode—with precision* (https://humansand.ai/blog/faster-than-light-nvfp4-decode) 中描述的那一个。为了实现激进的融合，高性能 kernel 可能会将逻辑权重重排为不同的物理布局。参见下面将 gate 与 up 权重按 64 为一组配对的示例。早先以模型检查点作为交接避开了这种复杂性，但当我们想要原地更新权重时，就必须解决它。 由于解压后的掩码与模型权重形状相同，我们可以以完全相同的方式重新映射掩码，使它们再次与重新映射后的模型权重逐点对齐。这为模型层与解压掩码之间的分块和地址映射增加了另一层复杂性，但所有这些元数据都在 CPU 上准备。执行 `XOR` 应用的 Triton kernel 只需读取准备好的计划。 ### 剩余的权衡\# (https://humansand.ai/blog/gpu-accelerated-weight-teleportation#the-remaining-trade-offs) 我们已经拥有一个极其强大的基线：对于相同的代理负载，训练器侧的准备与压缩耗时 0\.834 s，接收端的工作耗时 6\.041 s。然而，仍存在两个关键缺陷： 1. **负载大小。**在相同代理工作负载上，nvCOMP LZ4 负载大小为 5\.36 GB，而 CPU Zstd 基线仅为 0\.876 GB，因此 LZ4 传输的字节数是其 6\.1 倍。对 GPU 友好的压缩算法并不必然针对高压缩比进行优化。根据表 1，没有任何单一压缩算法能同时具备高压缩比、快速解压和 DE 支持。 2. **GPU 缓冲区大小。**GPU 解压也意味着需要 GPU 输入和输出缓冲区。为了在 HBM 中容纳大型解压掩码，我们必须驱逐 KV cache。首先，这使内存生命周期管理复杂化。更重要的是，驱逐 KV cache 实际上使得将 rollout 解码与解压重叠变得不可能，因为 KV cache 已经不存在了。 ## 从加速器到瞬移器\# (https://humansand.ai/blog/gpu-accelerated-weight-teleportation#from-accelerator-to-teleporter) 我们现在展示如何将 GPU 从加速器变成瞬移器（teleporter），把权重同步导致的 rollout 暂停降至 0\.23 s 以下。 ### 双重压缩与异构解压\# (https://humansand.ai/blog/gpu-accelerated-weight-teleportation#double-compression-with-heterogeneous-decompression) 压缩后的负载仍可能包含重复，尤其是使用对 GPU 友好、未必针对压缩比优化的算法时。因此我们考虑在 LZ4 或 Snappy 等对 GPU 友好的编解码器之上再加一轮 Zstd 压缩： *表 2。外层 Zstd 压缩前后的负载大小。* 从经验上看，Snappy \+ Zstd 相比 LZ4 基线带来了有意义的负载缩减。然而，在 LZ4 上再加 Zstd 带来的改善较小。一种可能的解释是，Snappy 更简单的拷贝命令更容易被 Zstd 压缩，尽管 Snappy 在单次遍历中编码长段零序列的效率不如 LZ4。 双重（解）压缩为接收端带来的设计机会不止于减小负载大小。由于外层 Zstd 阶段处理的是已经压缩过的小负载，CPU 可以高效地完成这一解压，而无需占用 GPU 计算或 HBM。在 CPU 接收 delta 更新负载、准备元数据、执行外层解压并将得到的 Snappy 负载放入锁页 CPU RAM 缓冲区以供 GPU 之后消费的同时，rollout 解码可以正常继续。 DE 可以直接消费锁页 CPU RAM 并输出到 GPU HBM⁵ CUDA Driver API`cuMemCreate` (https://docs.nvidia.com/cuda/cuda-driver-api/cuda_driver_api/group__CUDA__VA.html#cuMemCreate__CUmemGenericAllocationHandleP.s.CUmemAllocationPropCP.unsigned-l-l)\, 使用`CU\_MEM\_LOCATION\_TYPE\_HOST\_NUMA`和`CU\_MEM\_CREATE\_USAGE\_HW\_DECOMPRESS` (https://docs.nvidia.com/cuda/cuda-driver-api/cuda_driver_api/group__CUDA__TYPES.html#c.CU_MEM_CREATE_USAGE_HW_DECOMPRESS)于`allocFlags\.usage`。普通的锁页分配并不足够；参见 NVIDIA 的DE memory requirements (https://docs.nvidia.com/cuda/nvcomp/decompression_engine_faq.html#does-de-work-with-host-memory)\.，这与该 CPU\-GPU 解压流水线无缝衔接。这也消除了在 HBM 中设置输入暂存缓冲区的需求，略微减轻了 GPU 内存压力。我们还从经验上发现，这种方法达到了比我们自己的环形缓冲区实现更高的主机\-设备传输带宽。 ### 重叠 DE 与 SM\# (https://humansand.ai/blog/gpu-accelerated-weight-teleportation#overlapping-de-and-sm) 为了充分压满 DE，我们使用了最后一项特性。由于 DE 是用于解压的固定功能硬件模块，其工作可以与流式多处理器（SM）上的 GPU 计算重叠。对我们的用例而言，这意味着将 DE 解压与 SM 上的 Triton 按位 `XOR` kernel 重叠。我们设计了一个双缓冲的软件流水线，其中 DE 解压到一个槽位，而