TL;DR

  • •Fireworks 的新 ARCv3 压缩器减小了从训练器发送到生成强化学习(RL)rollout 的机器的 BF16 权重更新大小。
  • •在 1,000 个生产环境 RL 权重更新增量上,ARCv3 生成的载荷比 ARCv2 小了近 50%,并且能无损重建训练器精确的 BF16 权重。
  • •更小的传输量有助于 rollout 集群更贴近当前策略,使得跨区域使用算力而无需一个巨大的同地集群变得更加可行。
  • •ARCv3 已在Fireworks Training API中通过fireworks-delta-compression提供,面向使用自有训练器搭配 Fireworks rollout 的团队。

随着越来越多的团队投身强化学习以训练自己的前沿规模模型,有两件事比以往任何时候都更重要:获取训练算力,以及优化在训练机器之间传输数据的管道。RL 对这些管道的要求格外苛刻。一个前沿模型可能拥有数万亿参数,而训练器需要随着训练的推进,持续不断地将更新后的版本推送到生成 rollout 的机器集群。

当这些管道跟不上时,团队就被推向唯一的选择:一个庞大的同地集群,所有设备都连接在同一物理楼层上。这既昂贵又难以获得,还把较小的参与者挡在门外。在 Fireworks,我们穷尽一切可行路径,避免让这成为唯一的选择。否则,我们最终会得到一个只有少数几家公司才能进入的训练市场。

在之前的一篇文章中,前沿 RL 比你想象的更便宜,我们详细介绍了如何在实践中实现跨区域 RL。核心观察是,在我们考察的 RL 工作负载中,连续检查点之间只有约 2% 的模型 BF16 权重发生了变化。因此,我们不是每次训练器执行一步就把完整的 1 TB 检查点发送到每台 rollout 机器,而是发送一个紧凑的增量——对变化内容的压缩描述——并在另一端重建更新后的模型。正是这一点让训练运行能在三四个区域之间保持同步,而无需集群之间的专用高带宽网络。

这个增量的大小有助于判断它在实际中效果如何。较小的增量有助于 rollout 集群更快更新、更贴近当前策略,并从各区域的可用算力中获取资源。较大的增量则可能意味着训练器开始跑得比 rollout 快,集群逐渐过时,多区域部署相对单一超大集群的优势开始减弱。

How ARCv3 deltas move from trainer to rollouts
ARCv3 增量如何从训练器传输到 rollout。每个训练器分片并行地将其那部分压缩后的权重增量上传到共享对象存储(S3)。随后,Fireworks API 向每个区域发出信号,告知有新更新可用,每个 rollout 区域拉取所需的分片并在本地重建更新后的模型。训练器从不直接与 rollout 集群通信,这正是让全球分布式推理集群能够通过普通网络链路保持同步的原因。

我们发布 ARCv3,这是我们用来构建这些增量的压缩器的最新版本。这些改进适用于 BF16 权重,ARCv3 产生的载荷大小几乎只有上一版的一半,平均增量载荷约为原始 BF16 权重大小的 0.19%,低于此前的 0.36%(越低越好)。该压缩是无损的:rollout 端重建的模型与训练器产出的模型逐比特完全一致。ARCv3 已在 Fireworks Training API 中提供,可通过fireworks-delta-compression供那些使用自有训练器搭配 Fireworks rollout 的团队使用。

深入 ARCv3:并非所有更新都生而平等

这一改进源自 BF16 权重在训练步骤之间变化方式的一个特定属性。

当我们运行生产环境的 RL 训练会话并观察检查点之间实际发生的情况时,我们观察到一种强烈的不对称性。在某一步中约 2% 发生变化的 BF16 权重里,绝大多数只出现了尾数位移,而指数和符号保持不变。触及指数的更新很少见,翻转符号的更新则更为罕见。换句话说,大多数权重更新只是微调,而非跳跃。

ARCv3 直接编码这种不对称性。未变化的权重值会从载荷中省略。仅尾数变化的更新只贡献其尾数位。更罕见的更新——指数发生偏移或符号翻转——会被单独打包,每次更新的变化字段作为一个整体块编码。在 rollout 侧,两个流都会被应用到上一个检查点,并将结果与训练器的原始版本进行校验和比对,以验证重建结果是否匹配。

How ARCv3 encodes a block of weights
ARCv3 如何编码一个权重块。每个 prev/next 对进行 XOR 运算,并按实际发生变化的字段进行分类。未变化的权重值(XOR 为零)会从输出中省略。仅尾数变化的更新(最常见的情况)只贡献其尾数位(已高亮)。更罕见的更新——指数发生偏移或符号翻转——会被打包在一起,变化字段作为一个块编码。输出流的形态由变化的内容决定,而非输入格式的固定宽度。

对于模型中的每个张量,ARCv3 还会并行运行多种通用压缩算法,并保留产生最小输出的那一种。这比预先选定单一算法消耗更多 CPU,但训练机器通常在 GPU 进行实际训练时有空闲的 CPU 周期,所以值得花这些周期。

ARCv3 与该领域的对比

我们在从生产 RL 工作负载中抽取的 1,000 个真实权重更新增量上,将 ARCv3 与三种替代方案进行了基准测试。所有更新均为 BF16,且四种压缩器均无损运行,因此我们比较的是纯载荷大小:每种压缩器描述完全相同的权重更新需要多少字节。

ARCv3 的替代方案包括:

压缩器压缩后 / 原始 BF16 权重比

压缩器

压缩后 / 原始 BF16 权重比

ARCv3

0.19%

PULSESync

0.35%

ARCv2

0.36%

HF TRL(草稿)

0.67%

这里越低越好,ARCv3 在此基准测试中产生的载荷最小。在测试的 BF16 权重张量上,ARCv3 将压缩后的增量缩小到次优方法的大约一半。每个 rollout 区域拉取所需的增量片段,并重建训练器的精确 BF16 权重。这有助于让全球分布的 rollout 集群更贴近当前策略,这正是跨区域 RL 能够大规模运作的关键。

使用 ARCv3

如果你使用 Fireworks Trainer SDK 进行 rollout,则无需做任何事,因为 ARCv3 已经是默认选项。你的 rollout 集群将以更小的传输量运行,而你无需更改任何代码。

如果你自带训练器并使用 Fireworks 进行 rollout,你可以直接通过fireworks-delta-compression包集成 ARCv3。请参阅Fireworks rollouts 文档了解完整的集成流程。

bash

复制

1

pip install fireworks-delta-compression

python

复制

123456789101112131415161718192021

importtorch

fromfireworks_delta_compressionimportdelta_compress,delta_decompress

src=torch.randn(1024,1024,dtype=torch.bfloat16)

dst=src+1e-3*torch.randn_like(src)

compressed_diff,checksum=delta_compress(

src=src,

dst=dst,

compression_format="arc_v3",

)

restored=delta_decompress(

src=src,

compressed_diff=compressed_diff,

checksum=checksum,

compression_format="arc_v3",

)

# 与原始数据逐位相同

asserttorch.equal(restored.view(torch.uint8),dst.view(torch.uint8))

让分布式计算为 RL 服务

RL 权重更新具有结构,压缩器可以利用这一点。将罕见的指数变化与常见的尾数更新分离,并在其上竞速几个通用编解码器,为我们带来了几乎两倍于我们自己先前版本的压缩率,而代价只是在已有空闲 CPU 的机器上增加相对较低的 CPU 开销。

我们之所以推进这一点,是因为我们从增量中削减的每一个字节,都是在训练步骤中无需跨洲传输的字节。这让我们更接近于让跨三个或四个区域集群的训练感觉像是在同一栋楼里训练。它帮助更多团队在想法上竞争,而不是在能购买多少同地计算资源上竞争。

对于构建专用模型的团队来说,回报在于扩展 RL 时拥有更大的灵活性。Fireworks 负责在各 rollout 区域之间进行权重分发与同步。

ARCv3 是我们目前所处的阶段。我们会继续推进。

把你的训练器带到 Fireworks。