Hugging Face 在 TRL 中实现增量权重同步,单步传输从 1.2 GB 降到 20-35 MB
Shipping a Trillion Parameters With a Hub Bucket: Delta Weight Sync in TRL
Hugging Face 在 TRL 中落地增量权重同步(Delta Weight Sync):只把 RL 优化器步之间真正变化的 bf16 权重编码成稀疏 safetensors 文件,上传到 Hugging Face Bucket,再由 vLLM 拉取应用。
TL;DR, because you have models to train and we respect that:
- 异步 RL 有一个肮脏的秘密:每一步,训练器都必须把整个模型发送给推理引擎。对于 bf16 的 7B 模型,那是 14 GB。对于前沿的 1T 模型检查点,那大约是 1 TB。每一步都是如此。
- 事实证明你并不需要这样做。在两个连续的 RL 优化器步骤之间,大约 99% 的 bf16 权重是逐位相同的(最坏情况下也从未低于 98%)。实际的增量非常小。
- 我们提交了一个 TRL PR,它仅将变化的元素编码为稀疏 safetensors 文件,上传到Hugging Face Bucket,并通知 vLLM 去获取它。在 Qwen3-0.6B 上,每步负载从 1.2 GB 降至 20 到 35 MB。
- 锦上添花的是:我们运行了一次完整的分离式训练,其中训练器在一台机器上,vLLM 运行在一个 Hugging Face Space 中,Wordle 环境运行在另一个 Space 中,权重通过单个 Hub bucket 流动。没有共享集群,没有 RDMA,没有 VPN。
异步 RL 现在便宜多了。请继续阅读。
1. 一太字节问题
如果你读过我们之前关于异步 RL 训练全景的文章,你已经知道结论了。每个异步 RL 库,无论它如何拼写“actor model”,或者它的 NCCL 后端被涂成什么颜色,最终都会绊倒在同一个根源上:权重同步。
推理引擎说的是第 N 步的策略。训练器刚刚完成第 N+1 步。新权重必须从一边传到另一边,否则推理引擎会开始无可救药地偏离策略。无论你运行的是同步还是异步,这都位于关键路径上:阻塞式传输是浪费的 GPU 空闲计算,这些 GPU 没有在生成 token。使用稀疏增量路径,你可以将空闲时间压缩到几秒钟,而且训练器甚至不必等待推理引擎就绪:它只需发布“权重就绪”,并在优化器步骤完成的那一刻将权重上传到共享 bucket,而推理引擎则按自己的时间获取。
Fireworks 在他们的文章前沿 RL 比你想象的更便宜中给出了一个非常令人难忘的数字:对于 fp8(他们的设置)下的前沿 1T 参数检查点,完整快照是 1024 GiB,而传统观点认为每次更新 rollout 集群时都必须传输这么多。这种数字会让人开始绘制包含超级集群、RDMA 网络和专用跨区域链路的图表。他们测得的相邻检查点之间的平均增量是 20.3 GiB,即完整模型的 1.98%,并且“连续检查点之间超过 98% 的 bf16 格式权重保持逐位等价”。
Cursor 的Composer 2 报告讲述了一个类似的故事。他们在不同区域运行训练和推理,并用一个共享 S3 bucket(他们的原话)将它们连接起来,训练器每个训练步骤都将压缩的权重差异上传到该 bucket。每个集群独立地从共享增量链下载并重建,“不需要与训练集群直接连接”。双方从不直接通信参数。bucket 就是通信线路。
两篇论文在三点上达成一致,我们想慢慢重复一遍,因为本文的其余部分本质上是一份忠实的开源翻译:
- 在相邻的两个 RL 步骤之间,大多数权重实际上并没有改变。
- 如果你只发送发生变化的部分,你的带宽费用会降低大约两个数量级。
- 如果你把这些微小的差异通过共享对象存储进行路由,你就不再需要让训练器和推理集群位于同一个数据中心。
唯一缺少的是一个你可以pip install的故事版本。所以我们写了一个。
2. 为什么 bf16 RL 权重几乎总是稀疏的
在我们连接任何东西之前,值得理解为什么整个游戏甚至是可以获胜的。“98% 的权重不变”这一说法听起来可疑地像是那种在演示中有效但在实际中崩溃的数字。事实并非如此。它源于 bf16 算术在 RL 使用的学习率下的工作方式。
一个 bf16 数字有 7 个尾数位。在两个连续的 2 的幂之间,恰好有 27=1282^7 = 128 个可表示的值,因此围绕 ∣w∣|w| 的相邻 bf16 数字之间的间距大约为 ∣w∣⋅2−7|w| \cdot 2^{-7}。当更新低于该间距的一半时,即当 ∣Δw∣<∣w∣/256|\Delta w| < |w|/256 时,更新会被 bf16 转换吸收。这就是 PULSE 在其图 3 中绘制的“bf16 可见性阈值”。
现在看看 Adam 做了什么。在例如 3×10−63 \times 10^{-6} 的 RL 学习率下,对单个权重的更新为: Δw=−η⋅m^v^+ϵ\Delta w = -\eta \cdot \frac{\hat{m}}{\sqrt{\hat{v}} + \epsilon}
归一化步长 m^/(v^+ϵ)\hat{m}/(\sqrt{\hat{v}}+\epsilon) 大致为一阶,因此 ∣Δw∣≈η≈3×10−6|\Delta w| \approx \eta \approx 3 \times 10^{-6}。对于大多数权重,∣w∣|w| 大约在 10−210^{-2} 到 10−110^{-1} 之间(PULSE 报告代表性 LLM 权重的中位数为 0.019)。在该量级下的阈值 ∣w∣/256|w|/256 大约为 4×10−54 \times 10^{-5} 到 4×10−44 \times 10^{-4},这大于更新量。
换句话说:优化器在低语,而 bf16 听不见。更新被舍入吸收,ww 的字节表示没有改变,从推理引擎的角度来看,这个权重没有移动。将其乘以几亿个参数,你就能免费得到 >99% 的稀疏性数字,且零近似。
这正是 PULSE 论文(Mihai & Belilovsky, 2026)中形式化论证的观点。他们定义了两个阈值。吸收边界 10η10\eta 是 Adam 更新的保守最坏情况,而有效边界 η\eta 则是你实际所处的区间。bf16 可见性阈值为 ∣w∣/256|w|/256。只要更新量低于可见性阈值,它就会被吸收,bf16 字节不会改变。他们的图 3 将这两个边界与一组具有代表性的 LLM 权重云图进行了对比,结论明确无误:在 η=3×10−6\eta = 3 \times 10^{-6} 时,吸收边界本身已经低于模型中几乎每个权重的可见性阈值。他们在 Qwen2.5(0.5B/1.5B/7B)、Llama-3.2-3B 和 Gemma-3-4B 上进行了实证测量,并一致发现在 400 个训练步骤中,平均每步稀疏度约为 99%,标准差为 0.2% 至 0.4%。最坏情况下的步骤仍高于 98%。因此,<1% 发生变化并非侥幸测量;这是算术所保证的结果。
我们不必通过解析方式预测这一点(事实上,我们曾尝试从 Adam 的 mm 和 vv 统计量预测变化掩码,但召回率惨淡,只有 30%,稍后详述)。我们只需要观察哪些字节发生了翻转。那就是每个参数的一个微小布尔张量,在优化器步骤前后计算得出。
3. HF Buckets 与架构
这就是故事的第二部分登场之处,也是本文从 Fireworks/Cursor 的转述转变为 Hugging Face 相关内容的地方。
3.1 什么是 Bucket?
Bucket 是 Hub 上的一种仓库类型,专为高频对象存储而设计。没有提交仪式,没有 PR 工作流,没有 LFS 的怪癖。你添加文件、列出文件、下载文件。Python 接口只有两个函数:
from huggingface_hub import batch_bucket_files, download_bucket_files
# Trainer side
batch_bucket_files("my-org/wordle-deltas", add=[(buffer, "deltas/step_000042.safetensors")])
# Inference side
download_bucket_files("my-org/wordle-deltas", files=[("deltas/step_000042.safetensors", local_path)])
就是这样。两次函数调用,你的权重就已上路。
在底层,bucket 由 Xet 提供支持,这是 Hub 的内容定义分块存储层。Xet 会查看你上传的每个文件,根据其实际内容(而非固定偏移量)将其切分为块,并与 bucket 中已有的所有内容进行去重。实际效果——在此语境下令人欣喜——是即使我们懒得编写稀疏编码,每一步都只上传完整锚点,Xet 仍然只会传输发生变化的块。稀疏编码 + Xet 堆叠:我们只为移动的部分付费,而且只付一次。
这是 Fireworks 和 Cursor 都采用的“共享 S3 bucket”的开源等价物,区别在于存储层已经了解内容哈希,你现有的 HF token 已经拥有权限,并且它能与堆栈的其余部分(Spaces、数据集、模型)原生组合。
3.2 三个盒子
完整架构恰好有三个盒子加一个共享基底:
- 训练器。 随你放在哪里。一块 GPU、八块 GPU、一台通过 USB 连接 H100 的笔记本电脑,我们不会评判。它拥有模型权重,运行优化器,发出稀疏增量。
- HF Bucket。 单个仓库,两个前缀:
anchors/用于偶尔的完整快照,deltas/用于其间的稀疏补丁。这是双方唯一达成一致的地方。 - vLLM rollout 服务器。 可以放在任何地方,关键是不一定非要和训练器在一起。它从 bucket 拉取数据,应用增量,然后提供 rollout 服务。
- 环境。 以通常的方式挂在 rollout 服务器上(HTTP、函数调用,或者你的环境使用的任何方式)。
需要内化的特性,也就是 Cursor 论文极力推销、并且在这里完全成立的那一点:训练器和 rollout 服务器从不互相通信交换权重。它们交换一个包含 {"repo_id": ..., "filename": ...} 的小型 POST 请求,这就是整个控制平面。实际的字节传输发生在每一侧与 bucket 之间,并行进行,没有共享的网络结构。
为什么这在实践中很重要:
- rollout 服务器可以位于另一个区域、另一个云,或者位于 Hugging Face Space 内 NAT 之后。它不在乎。
- N 个推理副本可以从同一个 bucket 拉取相同的增量,而 Xet 会在所有副本之间对字节进行去重。
- 训练器永远不需要知道存在多少个推理副本、它们在哪里,或者其中一个是否刚刚崩溃。
训练器写入。副本读取。Hub 负责管道连接。
4. 协议
现在我们可以打开引擎盖了。协议有四个部分:一种传输格式、一种 bucket 布局、一个 30 行的 vLLM 扩展,以及一个训练器侧的变化检测器。老实说,代码比听起来要少。
4.1 以 Safetensors 作为传输格式
我们选择 safetensors 作为磁盘和传输格式。它已经是 Hub 上的规范检查点格式,每个合理的框架都能读取它,而且头部携带任意字符串元数据。那个元数据字段就是我们隐藏协议的地方。
bucket 中有两种文件。
锚点看起来像普通检查点:每个参数一个张量,完整的 bf16 权重,每 NN 次同步写入一次(我们默认 N=10N=10)。
anchors/step_000010.safetensors
├── model.layers.0.self_attn.q_proj.weight (bf16, full)
├── model.layers.0.self_attn.k_proj.weight (bf16, full)
└── ...
metadata:
sparse=False, model_version=10, sparsity=0.0
增量才是有趣的部分。对于每个实际发生变化的参数,我们存储两个条目:一个扁平 int32 张量表示元素索引,以及一个 bf16 张量表示这些索引处的值。
deltas/step_000011.safetensors
├── model.layers.0.self_attn.q_proj.weight.indices (int32, [num_changed])
├── model.layers.0.self_attn.q_proj.weight.values (bf16, [num_changed])
├── model.layers.0.mlp.gate_proj.weight.indices
├── model.layers.0.mlp.gate_proj.weight.values
└── ...
metadata:
sparse=True, model_version=11, sparsity=0.9938, changed_params=[...]
这种选择带来几个不错的结果:
- 增量是一个文件。你可以在 Python 中用
safe_open(...)打开它,并检查其中的每个张量。没有专有帧格式,没有长度前缀,没有版本握手。 - 元数据是自描述的。接收方读取
sparse=True/False并分支处理。没有单独的清单。 - 在推理侧通过 mmap 实现零拷贝,当你每隔几秒就要做一次时,这一点很重要。
节奏很直接:每第 N 步一个锚点,其间是增量。两者最终都放在同一个 bucket 中,分别位于 anchors/ 和 deltas/ 前缀下。每个新的推理副本只需要抓取最近的锚点,然后重放此后的增量。
4.2 训练器侧:来自优化器钩子的布尔掩码
训练器需要知道哪些 bf16 元素实际发生了翻转。我们用一个很小的 BF16ChangeDetector 来实现这一点,它在优化器上注册 pre-step 和 post-step 钩子:
class BF16ChangeDetector:
def __init__(self, model, optimizer):
self._pre_step_bf16: dict[str, torch.Tensor] = {}
self._validated_masks: dict[str, torch.Tensor] = {}
optimizer.register_step_pre_hook(self._pre_step_hook)
optimizer.register_step_post_hook(self._post_step_hook)
def _pre_step_hook(self, opt, args, kwargs):
for p in self._params:
self._pre_step_bf16[name_of(p)] = p.detach().to(torch.bfloat16).cpu().clone()
def _post_step_hook(self, opt, args, kwargs):
for p in self._params:
self._validated_masks[name_of(p)] = (
p.detach().to(torch.bfloat16).cpu() != self._pre_step_bf16[name_of(p)]
)
PR 中的实际代码有更多细节(通过 data_ptr() 将优化器参数对象与模型参数匹配,因为 Accelerate 会将它们包装为不同的 Python 对象),但核心思路可以画在餐巾纸上:快照、步进、差分。
这是基准事实。我们尝试了更优雅的路径,即直接从 Adam 的 mm 和 vv 统计量预测掩码,使用 bf16 ULP 阈值。这在原理上是可行的。但在实践中,召回率约为 30%,这意味着我们会发布一个缺失了三分之二实际更新的增量。Adam 的归一化足够混乱,以至于解析阈值并不紧致。所以我们直接比较字节。这需要在训练器侧对模型做一次 bf16 CPU 快照,我们愿意付出这个代价。
新的 _sync_weight 流程的四个阶段是:
- 推理持续运行时上传。 训练器将掩码元素编码到 safetensors 缓冲区,并推送到 bucket。在整个过程中,vLLM 仍在愉快地服务旧策略。
- 暂停 vLLM。 一次简短的 HTTP 调用,数百毫秒。
- 通知
/update_weights。 发送 bucket 坐标。vLLM 下载、应用、返回。 - 恢复。 vLLM 重新上线。
日志行说明了整个过程:
Delta: 1234567/200000000 elements changed (sparsity=99.38%)
[delta_engine] uploaded user/wordle-deltas/deltas/step_000042.safetensors (27.4 MB, ...)
Weight sync: done. Total 9.4s (inference paused 1.1s)
关键的是括号里的数字。推理暂停了 1.1 秒。其余 9.4 秒用于上传,而这发生在 rollout 服务器仍在生成 token 的时候。使用 NCCL 时,我们支付的全部同步时间都算作暂停时间。而在这里,我们将其作为后台时间来支付。
4.3 vLLM 侧:30 行扩展
vLLM 为此提供了一个干净的抽象,称为 WeightTransferEngine。我们实现了一个 DeltaWeightTransferEngine,其 receive_weights 方法在精神上大致是:
def receive_weights(self, update_info, load_weights):
download_bucket_files(update_info.repo_id, files=[(update_info.filename, local_path)])
with safe_open(local_path, framework="pt", device="cpu") as f:
meta = PatchMetadata.from_metadata_dict(f.metadata())
if not meta.sparse:
# Anchor: feed every tensor and snapshot for future deltas
for name in f.keys():
tensor = f.get_tensor(name)
self._bf16_snapshot[name] = tensor.clone()
load_weights([(name, tensor)])
else:
# Delta: apply (indices, values) to snapshot, hand full tensor to vLLM
for name in json.loads(meta.changed_params):
indices = f.get_tensor(f"{name}.indices").long()
values = f.get_tensor(f"{name}.values")
snap = self._bf16_snapshot[name].flatten()
snap[indices] = values
self._bf16_snapshot[name] = snap.reshape(self._bf16_snapshot[name].shape)
load_weights([(name, self._bf16_snapshot[name])])
我们通过 vLLM 的 --worker-extension-cls 标志注册它,这意味着无需 fork vLLM。你将 TRL 安装到与 vLLM 相同的镜像中,将 CLI 指向我们的类,就完成了。
值得一提的是:vLLM 自身正在努力原生支持稀疏权重传输,vllm-project/vllm#40096。它直接在 WeightTransferEngine 基类上添加了 receive_sparse_weights() 和 trainer_send_sparse_weights(),补丁编码为 (indices, values) 并通过 index_copy_() 原地应用,完全消除了 GPU/CPU 验证往返。该 PR 报告称,对于 Qwen3-1.7B 上的稀疏补丁,传输为 0.16 MB 耗时 0.40 ms,而完整稠密发送为 942 MB 耗时 192 ms。
我们在推理侧实现中有一个诚实的注意事项:我们保留模型的 CPU bf16 快照,以便从稀疏 (indices, values) 补丁重建完整张量,因为 vLLM 目前的 load_weights 期望完整张量。一旦 #40096(或其后续)落地并暴露原地稀疏 load_weights 路径,我们就可以直接在 GPU 上应用索引并丢弃快照!
5. 在 Spaces 上真实部署
这是我们引以为豪的部分。到目前为止我们描述的一切都可以在你的笔记本电脑上运行,但通过 Hub bucket 路由权重的意义在于,训练器和 rollout 服务器不必位于彼此附近。因此,我们使用三台机器进行了完全解耦的训练,它们之间没有共享网络:
- 一台带有一个 GPU 的机器运行训练器。
- 一个 Hugging Face Space(Docker SDK,L4 GPU)运行带有我们扩展类的 vLLM。
- 第二个 Hugging Face Space(CPU)运行 Wordle 环境服务器,支持 256 个并发会话。
- 中间是一个 Hub 存储桶。
设置这个真的只需要几个 hf CLI 调用。vLLM Space 的 Dockerfile 基本上就是上游 vLLM 镜像加上 pip install trl@... 再加上入口点:
FROM vllm/vllm-openai:latest
RUN pip install "trl @ git+https://github.com/huggingface/trl.git@delta-weight-sync"
ENV VLLM_SERVER_DEV_MODE=1
EXPOSE 7860
ENTRYPOINT ["vllm", "serve", "Qwen/Qwen3-1.7B", \
"--host", "0.0.0.0", "--port", "7860", \
"--worker-extension-cls", "trl.experimental.async_grpo.delta_engine.DeltaWorkerExtension", \
"--weight-transfer-config", "{\"backend\":\"nccl\"}", \
"--max-model-len", "32768", \
"--gpu-memory-utilization", "0.8"]
将其部署为一个 Space:
hf repos create $USER/vllm-wordle-inference \
--type space --space-sdk docker --flavor l4x1 \
--secrets HF_TOKEN=$HF_TOKEN
hf upload $USER/vllm-wordle-inference examples/scripts/openenv/vllm_space/ --type space
然后从地球上任何能进行 HTTPS 通信的地方启动训练:
python examples/scripts/openenv/async_wordle.py \
--vllm-server-url https://$USER-vllm-wordle-inference.hf.space \
--env-url https://openenv-wordle.hf.space \
--delta-sync-repo-id $USER/wordle-deltas \
--model Qwen/Qwen3-1.7B
训练器从不开放端口。Space 从未看到训练器的 IP。Wordle 环境也不知道它们的存在。它们都通过 Hub 通信。训练在即时 EOS 健全性检查上收敛,然后在真实的 Wordle 推演上收敛:奖励上升,增量负载保持在 20 到 35 MB 范围内,每次同步的推理暂停窗口保持在 1 秒左右。完整的运行日志链接在配套 PR 中。
6. 那么这究竟解锁了什么?
有几件事,我们认为它们很重要。
无需集群的异步 RL 训练。如果你有一块 GPU 和一个 Hugging Face 账户,现在就可以进行真正的解耦训练。你的训练器在 GPU 上;你的推演集群在 Spaces 中;你的环境在另一个 Space 中;权重通过存储桶传输。这过去需要共置设置(带来所有吞吐量妥协)或具有共享网络的真实集群。现在不再需要了。
多副本推理,免费获得。启动两个 vLLM Spaces,或者十个。它们都从同一个存储桶拉取。Xet 对存储进行内容寻址,因此连续的锚点在静止时共享块(这可以防止你的存储桶膨胀),而 Hub 的边缘缓存使得重复下载同一文件变得廉价。想要一个全球分布的推演集群?现在这只是一个小型 DevOps 练习,而不是一个研究项目。
一种可以用现有工具调试的线格式。增量是一个 safetensors 文件。你可以从笔记本中 safe_open 它,列出其键,检查索引,自己计算稀疏性。我们在不透明的 NCCL 流上花了足够多的时间在 tcpdump 中,因此对此表示赞赏。
通往前沿规模的路径。20 到 35 MB 的数字是针对 Qwen3-0.6B 的。有趣的问题是,一旦你调高旋钮,曲线会是什么样子。让我们做个粗略计算。
以 Llama-3.1-405B 为例。在 bf16 中,磁盘上为 810 GB。PULSE 在 RL 学习率下测得每步平均稀疏度约为 99%,因此实际增量约为参数的 1%。他们部署测量的编码在 7B 模型上达到 108 MB,这是 PULSE 报告的 ~130 倍减少。线性扩展到 405B,增量大约为 每步 6 GB。
这在挂钟时间上能带来什么?NCCL 在集群内部很快,当然。假设一个慷慨的 100 GB/s 聚合广播带宽(多节点、RDMA 等)。完整同步是每步 810 GB / 100 GB/s ≈ 8 seconds 的推理暂停。使用增量路径,训练器在生成继续运行的同时在后台将 6 GB 流式传输到存储桶,而推演服务器的实际暂停窗口只是应用步骤,在这个规模上只需几秒钟。因此,即使在我们离开集群之前,增量也将可见暂停减少了 4 倍,并将线上的字节减少了约 130 倍。
现在离开集群。NCCL 跨云完全无法工作。一旦你想要一个 rollout 集群在 us-east,另一个在 eu-west,也许还有一个在 Hugging Face Space 中,基于 bucket 的路径就是唯一的路径。在 1 GB/s 的可用互联网带宽下,一次完整的广播需要 13 分钟;而增量只需 6 秒。
对于 Fireworks 框架中的 1 TB 级模型,他们自己测量的数字显示 20.3 GiB 增量对比 1024 GiB 完整快照,减少了约 50 倍。PULSE 更紧凑的稀疏编码会进一步推动这一点(外推每个增量约 15 GB,接近约 65 倍)。无论哪种方式,你都处于这样一种状态:通过普通对象存储传输权重不再是一种 hack,而是开始成为唯一合理的架构。
7. 我们还有什么待解决
我们并不假装这已经完成。以下是诚实的清单。
- 两个 CPU bf16 快照,多了一个。 训练器保留一个(用于变化检测器),rollout 服务器保留一个(用于为 vLLM 的
load_weights重建完整张量)。第一个我们暂时无法摆脱,直到有人找到一种紧凑的解析掩码,这比看起来要难。当 vLLM 获得稀疏load_weightsAPI 时,第二个就会消失。PR 即将到来。 - 固定的锚点节奏。 我们目前每 NN 步转储一个完整锚点。自适应策略(“当累积漂移超过 X 时锚定”)会在长时间运行中降低锚点成本。
- 多节点 FSDP2 训练器。
BF16ChangeDetector是围绕每个进程的优化器钩子构建的。它应该能干净地推广到 FSDP2,但我们还没有在多节点规模上测量过。PR 中有一个TODO上面有我们的名字。 - 接入优化器。 我们尝试仅从 (m,v)(m, v) 预测掩码,但召回率很低,这意味着解析 bf16 阈值做的事情比教科书公式所暗示的更微妙。我们很想听到任何破解了这个问题的人的意见。
- 与线上压缩叠加。 稀疏 safetensors 和逐块 gzip 是正交的。我们还没有尝试将它们结合起来。尽管我们并不期望巨大的压缩收益。
8. 试一试
- PR:huggingface/trl#5417。分支是
delta-weight-sync。 - 完整的 Wordle 示例:
examples/scripts/openenv/async_wordle.py。 - Spaces Dockerfiles:
examples/scripts/openenv/vllm_space/和examples/scripts/openenv/wordle_space/。 - 背景阅读:我们的 异步 RL 景观文章、Fireworks 1 TB 文章、Cursor Composer 2 报告。
来源:Hugging Face Blog · huggingface.co