返回资讯中心
外部精选
软件工程
#开发者实时流

新方法实现GPU加速权重传输,采样器暂停低于0.23秒

来源标题:GPU-Accelerated Weight Teleportation

研究团队提出一种名为“权重传送”的快速权重更新方案,在Ampersand生产级强化学习配置下将采样器暂停时间降至0.23秒以内,端到端权重同步低于4秒。该方案基于增量权重同步思路,设计了Blackwell原生压缩/解压流水线,速度在多数场景下超过集群内网络传输。即使采样器与训练器不在同一集群也能适用。

Hacker News 最新ezelikman2026/11/2 · 今天30 分钟阅读

以下正文由 AI 翻译自 Hacker News 最新(原文版权归原站所有),可能存在偏差,重要内容请以原文为准。

大规模RL与大规模预训练、大规模推理都不一样。一个关键区别在于,训练权重和采样器权重最好尽可能保持一致,这就要求训练器把权重更新发送给采样器。这就产生了一个瓶颈:要么使用高速的集群内网络来发送权重更新,要么接受传输数TB模型权重所带来的、往往相当可观的延迟——更不用说还要在采样器上阻塞等待了。我们引入了一种更快的权重更新范式,在大多数情况下能实现优于集群内网络的速度,无论采样器和训练器是否位于同一个集群内。

该设计始于一个被广泛采用的思路,即增量权重同步¹(Miles v0.1: Production-Level Post-Training, §4.3 Disk-Delta Updates)。我们共同设计了一整套Blackwell原生的压缩/解压缩流水线,结合了以下技术:

最终,在我们的Ampersand生产RL配置中,我们把采样器暂停时间降到了0.23秒以下,端到端权重同步时间降到了4秒以下,完全打满了Blackwell解压缩引擎流水线,使解压缩ASIC成为采样器暂停时间的瓶颈。无论在端到端权重同步时间还是采样器暂停时间上,这个实现都如此之快,以至于对于大多数真实工作负载,即便是集群内P2P RDMA也更慢。

在本文余下部分,我们使用以下一致的工作负载代理。模型是NVFP4 GLM-5.3,具有445 GB的规范权重张量。RL训练信号是稀疏的,每个训练步骤中只有极小一部分权重被更新,因此我们以checkpoint的0.085%作为一次稳态更新的代理。训练器有8个节点,配备64块B300 GPU。每个rollout接收端在一个8-GPU B300节点上运行两个TP4 NVFP4 GLM-5.3引擎。为清晰起见,图示中使用了更少的GPU和模型层。

为了解释我们的设计演进过程,我们需要理解增量权重同步的基础,以及为什么朴素实现还不够。

trainer

mask = weights_old ⊕ weights_new

rollout

weights_updated = weights_old ⊕ mask = weights_new

增量权重同步的核心是按位异或(⊕)操作。在一个训练步骤之前,rollout引擎和训练器持有相同的权重。训练器完成一个训练步骤,从weights_old推进到weights_new。训练器对weights_old和weights_new执行按位异或,得到一个与权重形状完全相同的mask。mask中的1表示该位置的权重位发生了翻转,0表示未变化。rollout引擎需要weights_new,以便在下一个RL步骤中从新策略采样。为获得它们,引擎只需翻转相同的权重位。因此,我们把mask传输给rollout引擎,并在它们的weights_old和mask之间执行按位异或。由于权重被视为一长段连续的比特序列,它们的具体数据类型并不重要,也不涉及浮点运算或舍入。

但mask与权重的大小相同,因此直接传输mask相比传输完整权重并无任何节省。真正让增量同步方案有吸引力的是以下洞见。在RL中,训练信号是稀疏的,因此每个训练步骤中只有极小一部分权重被更新。这意味着mask中大部分是0,只有极小一部分是1。然后,无损压缩算法可以把mask压缩成一个很小的传输载荷。

现在我们来看RL循环中一个朴素但清晰的设计。训练器和rollout引擎需要同步它们的checkpoint,因此我们用模型checkpoint作为交接点。训练器执行一个训练步骤,计算mask,将其压缩成一个小载荷,并把压缩后的mask发送给rollout引擎。rollout引擎加载并解压mask,修补checkpoint以匹配新的训练器权重,然后从该checkpoint重新加载权重。由于checkpoint是交接点,rollout引擎复用了与服务器启动相同的代码路径。我们不需要处理rollout引擎如何重新分片或映射权重。

Miles中现有的disk-delta实现¹在很大程度上遵循这一设计:

从上面的测量来看,压缩载荷只有0.876 GB,符合我们的预期。然而,这一收益被其他显著开销完全抵消了:

RL栈已经运行在分布式GPU集群上,因此直观的下一步是让这条流水线GPU加速,利用其强大的并行计算能力和高带宽HBM。

Blackwell有一个固定功能的解压缩引擎(DE)²,能以极高效率执行某些受支持的解压缩算法。为使我们的基线足够强,我们直接将GPU加速流水线与GPU原生(解)压缩算法和Blackwell原生专用硬件进行协同设计。

下表比较了GPU上的(解)压缩算法³:

表1. GPU编解码器权衡与解压缩引擎支持。

如上文图4所示,主要开销在接收端,包括解压缩、checkpoint修补和引擎重载。因此,我们优先考虑接收端解压缩性能,从LZ4/Snappy开始构建快速的GPU编解码器流水线。

(原文较长,此处译出前 6000 字,完整内容请点「打开原文」。)

正文由 探How 从来源站点 RSS 同步,内容未经改写;遇到排版缺失或需要图片、视频时请以原文为准。

打开原文
去社区讨论 →