← 返回 2026-09-20 简报

在 HF Jobs 上结合 LoRA 的异步 GRPO:一个 Bucket、一个代理,无需 NCCL

语音播报
摘要
事件:Hugging Face Jobs 结合 TRL v1.14 异步 GRPO 训练 LoRA 适配器,通过挂载存储桶实现跨节点同步,无需 NCCL。 要点:仅同步几兆的 LoRA 权重;五轮测试耗时从3小时27分降至53分钟;利用代理路由并广播适配器加载。 影响:打破单节点限制,让训练与推理分离部署更灵活高效,大幅降低网络开销,加速 RLHF 迭代流程。

TL;DR

AsyncGRPOTrainer 现在可以训练 LoRA 适配器,并将该适配器仅同步到 vLLM(TRL v1.14)。

一个 rank-1 适配器仅有几兆字节大小,因此它可以通过挂载在每个 Job 中的 Storage Bucket 进行传输,而不是通过 NCCL。Trainer 和 vLLM 副本作为独立的 Hugging Face Jobs 运行在独立的机器上。

位于副本前端的一个小型代理负责添加认证头(auth header),将每个 rollout 路由到已持有其 KV prefix 的副本,并向所有副本广播适配器加载操作。

AsyncGRPO 指标显示了瓶颈所在。五次运行将同一配方从 3 小时 27 分钟缩短至 500 步仅需 53 分钟。

LoRA 支持最近通过 PR #7017 合并到了 TRL 的 AsyncGRPOTrainer 中,并随 TRL v1.14 发布。异步 Trainer 现在可以训练适配器而非完整模型,并且仅将 LoRA 适配器同步到 vLLM。本文介绍了一个基于此构建的实际项目,其中训练和推理不再共享同一台机器。

LoRA 训练特别适用于强化学习(RL),正如 Thinking Machines 的博客文章《LoRA Without Regret》所示。他们表明,即使使用 rank 1,LoRA 也能在策略梯度 RL 中媲美全量微调。这源于这样一个事实:优势函数在每个 episode 中仅提供约 O(1) 比特的信息,因此从总信息量的角度来看,每一步需要学习的内容并不多。一个 rank-1 适配器具有足够的容量来吸收这些信息。

LoRA 训练还有一个系统层面的影响。对于 1.5B 模型,rank-1 适配器仅有几兆字节,而完整模型约为 3 GB。我们无需在每个更新后将完整的策略发送给推理工作者,只需发送适配器即可。vLLM 还可以同时加载多个适配器。旧的 rollout 使用其启动时的策略完成,而新的 rollout 则使用最新的策略。

TRL 的 AsyncGRPOTrainer 已经实现了训练和生成的分离。Trainer 和 vLLM 可以运行在不同的机器上,并以各自的速度运行。在单节点或集群环境中,如果两个进程共享文件系统或可以形成 NCCL 组,这很容易实现。

我们想要的是使用 Hugging Face Jobs 运行相同的设置。本质上,HF Job 是在单个 VM 上运行的一个容器。这意味着一个 Job 无法生成多个节点(至少目前如此)来同时容纳 Trainer 和 vLLM 服务器集群(每个节点最多限制为 8xH200)。AsyncGRPOTrainer 正是为这种规模而构建的,因此问题变成了:如果我们放弃 Trainer 和推理服务器必须共享同一节点的要求,我们能走多远?

嗯,如果是全权重同步,答案将是“走不远”。每次更新都必须在机器之间移动吉字节(GB)的数据,这正是 NCCL 在密集集群中用于做的事情,但 Jobs 无法跨节点通信。没有共享本地磁盘,显然也没有共享 localhost 。有了 LoRA,同步仅需几兆字节。对于文件系统部分,HF Jobs 提供由 Storage Buckets 支持的卷!这些存储桶可以挂载为每个 Job 中的 FUSE 文件系统,足以在节点之间充当共享 FS。根本不需要 Jobs 之间的网络路径。

最终设置相当精简:

一个运行带有 LoRA(以及 FSDP,稍后详述)的 AsyncGRPOTrainer 的 Trainer Job,

两个 vLLM Job,每个都服务于基础模型以及 Trainer 最近发布的任意适配器,

在所有三个 Job 中以相同路径挂载的一个 Storage Bucket,适配器正是通过它从 Trainer 传输到服务器的,

一个代理服务器。我们将深入探讨为什么需要它,但大体上,我们需要一个代理,将每个 rollout 路由到最可能持有其 KV 缓存的副本,并将每次适配器更新广播给所有 vLLM 副本。

架构:利用 Hugging Face Jobs 和 Storage Buckets 🪣

AsyncGRPOTrainer 中的新仅适配器同步路径工作原理如下。该训练器不会向 vLLM 发送张量。每隔几个优化器步骤,它会在 /.vllm_lora/trl-policy-v{N} 下保存适配器,通过原子重命名发布该目录,然后将其路径发送至 vLLM 的 /v1/load_lora_adapter 端点。vLLM 从磁盘加载文件,因此 rollout 工作进程随后可以请求 model="trl-policy-v{N}"。

这正是 vLLM 中运行时适配器加载的工作方式。该端点接受路径而非张量,因此训练器和服务器被期望共享一个文件系统。在 Slurm 集群上,这是网络文件系统。在 Jobs 中,我们通过将存储桶挂载为每个 Job 中相同路径下的卷来获得相同的效果,正如我们之前提到的那样。在底层,它使用 hf-mount,将存储桶作为 POSIX 文件系统暴露在容器内:

hf jobs run ... -v hf://buckets/aminediroHF/asyncgrpo-lora-buckets:/lora ...

TRL 或 vLLM 都不需要为此进行更改。训练器写入 /lora//.vllm_lora/,服务器从同一路径读取。POST 请求中发送的路径在每个容器内都是有效的。

三个 Job 和存储桶。TRL 通过 localhost 与代理通信,代理通过 HTTPS 与副本通信,适配器目录通过存储桶挂载传输。

请注意,我们还将检查点和最终适配器存储在存储桶中。HF Jobs 是临时的,但被抢占的训练器可以恢复训练,因为最终适配器始终持久化到存储桶中,当 Job 停止时永远不会丢失。

三个 Job

vLLM 副本

每个副本使用一个 GPU 和标准的 vllm/vllm-openai 镜像。我们只需要启用运行时 LoRA 加载并预留足够的适配器槽位。

适配器槽位的数量由 max_staleness 决定。在 AsyncGRPOTrainer 中,每次权重同步都会使策略版本加一,而 max_staleness 是 rollout 样本可能滞后于当前策略多少个版本,之后训练器才会将其丢弃。使用 max_staleness=4 ,在 trl-policy-v3 下生成的样本仍可用于训练,而此时训练器已处于 v7 。在 v3 下开始的 rollout 也必须能够在 v3 下完成。因此,在任何时刻,vLLM 必须服务当前策略以及之前的四个版本。这就是为什么训练器保留 max_staleness + 1 个适配器版本注册并卸载任何更早的版本的原因。每次同步都会先加载新版本,然后再卸载最旧版本,这在交换期间需要多一个槽位。这给出了 --max-loras 6 。如果只有五个,vLLM 将在每次同步时静默驱逐仍有 rollout 在飞行中的策略。

for replica in 1 2; do
hf jobs run --detach --flavor h200 -- timeout 8h --secrets HF_TOKEN \
--expose 8000 \
-v "hf://buckets/ ${BUCKET} :/lora:ro" \
-e VLLM_ALLOW_RUNTIME_LORA_UPDATING=1 \
-e VLLM_SERVER_DEV_MODE=1 \
-- vllm/vllm-openai:v0.27.1 \
vllm serve Qwen/Qwen2.5-Math-1.5B --host 0.0.0.0 --port 8000 \
--max-model-len 4096 --logprobs-mode processed_logprobs --generation-config vllm \
--enable-lora --max-lora-rank 1 --max-loras 6
done

我们将 vLLM 固定为 v0.27.1 。vLLM 更新迅速,上述标志和运行时 LoRA 端点是该版本所暴露的,因此请将此版本视为配方的一部分。

还有另一种可能的设计,即训练器仅保留最新的适配器并始终使用相同的名称发布它。我们没有选择这种方式,因为 vLLM 根据其适配器名称对前缀缓存进行键控。使用单一名称,在交换后之前权重下计算的 KV 块仍然会匹配,因此预填充不会重新执行,rollout 可能会从其策略版本获取前缀并从下一个版本获取解码。训练器无法察觉这一点,这将表现为比率偏离 1。版本化的名称使这种情况成为不可能:一个名称始终代表一组权重,且缓存的前缀永远无法匹配更新的版本。

数据集选择:Sanity 集

我们选择了 sail/Sanity-Test-R1D-1.5B,该数据集来自《通过 FP16 克服训练与推理不匹配》(Qi 等人,2025)。复现代码位于 sail-sg/Precision-RL 中。

作者使用 DeepSeek-R1-Distill-Qwen-1.5B 为每个 MATH 问题生成了 40 个答案。他们保留了成功率在 20% 到 80% 之间的问题,从而得到 1,460 个问题。这个数据集非常适合 RL 验证,因为这些问题的难度对于该模型来说既不是已经解决也不是完全无望,这意味着模型可以获得良好的早期信号来进行训练并提升性能。

这作为一个稳健的端到端测试非常棒:如果一个 vLLM 副本在适配器名称下静默地服务基础模型,我们希望能在几十步内的曲线中看到这一点。此外,这个数据集足够小,可以在不到两小时内完成一个周期的遍历。

我们还采用了 oat/scripts/lora 中论文 LoRA 脚本里的超参数:Qwen/Qwen2.5-Math-1.5B,LoRA rank 为 1,alpha 为 2,学习率为 4e-5,每个提示词 8 个样本,每步 128 个完成结果,最大生成 token 数为 3,000,上下文长度为 4,096 个 token。

训练器

训练器使用安装了 TRL 的 vllm/vllm-openai:v0.27.1 镜像。我们当时运行了 PR 分支;相同的代码现在包含在 TRL v1.14 中。训练脚本是一个标准的 AsyncGRPOTrainer 脚本。唯一的 Job 特定值是输出目录和服务器 URL。

from peft import LoraConfig
from trl.experimental.async_grpo import AsyncGRPOConfig, AsyncGRPOTrainer

config = AsyncGRPOConfig(
output_dir= "/lora/sanity-lora-r1" ,
vllm_server_base_url= "http://localhost:8000" ,
max_staleness= 4 ,
weight_sync_steps= 4 ,
save_strategy= "steps" , save_steps= 50 ,
...
)
trainer = AsyncGRPOTrainer(
model= "Qwen/Qwen2.5-Math-1.5B" ,
args=config,
peft_config=LoraConfig(r= 1 , lora_alpha= 2 , target_modules= "all-linear" ),
...
)

在初始化期间,TRL 会调用 /server_info 。如果它找到 lora_config ,它将使用仅适配器的同步。vLLM 无法直接服务的配置,如 DoRA、modules_to_save 或超过 --max-lora-rank 的 rank,将回退到合并权重同步并发出警告。日志中应包含 “Adapter-only vLLM sync enabled”(已启用仅适配器的 vLLM 同步)。

代理

现在进入有趣的部分。我们需要在训练器和 vLLM Jobs 之间设置一个代理,原因有二:

暴露的 Job 端口需要在每个请求上添加 Authorization: Bearer 头。代理是添加该头部的地方,因此 TRL 无需了解它。

我们希望不止一个 GPU 进行生成。在单个 vLLM 服务器上,通常的方法是 --data-parallel-size > 1 ,但 TRL 拒绝在这种模式下使用仅适配器的同步,这是有充分理由的:对 /v1/load_lora_adapter 的调用只会到达回答它的 DP rank,因此其他 rank 将继续在新策略名称下服务基础模型。在 Jobs 上,这个问题甚至不存在,因为每个副本都是独立的机器。因此,数据并行必须高出一个层级,存在于某种将适配器负载分发到每个副本的东西中。

因此,我们在训练器 Job 上的 127.0.0.1:8000 运行一个小型代理,并将 TRL 指向它,就像它是一个单一的 vLLM 服务器一样。除了添加头部外,代理在功能上还做两件事:

它将每个完成请求发送给一个副本,选择该副本使得提示词的八个 rollout 落在其前缀已经缓存的位置(细节见下文)。

它将每个改变状态的请求(如适配器加载、暂停和恢复)广播给所有副本,以便策略名称在任何地方都具有相同的含义。

通过 KV 前缀路由 rollout

快速回顾一下为什么这很重要。生成完成结果有两个阶段,具有非常不同的工作负载特征:

prefill 阶段一次性处理整个提示词,并为每个提示词 token 计算注意力键和值。

decode 阶段随后一次生成一个 token,每个新 token 都会关注其之前所有 token 的键和值。

这些键和值构成了 KV 缓存。由于注意力机制是因果的,一个 token 的 KV 仅取决于它之前的 token,而不受后续内容的影响。因此,共享前缀的两个请求也共享该前缀的 KV,而已经将其缓存在副本中的那个副本可以完全跳过预填充(prefill)的那部分工作。现在的关键在于找到那个副本,以便请求能够受益于落在已经见过其前缀的副本上。

vLLM 将其前缀 KV 缓存存储在 16 个 token 的块中。由于 GRPO 的存在,rollout 工作者会发送 G 个具有相同提示的请求(在我们的案例中 G=8)。如果它们都到达同一个副本,第一个请求会计算预填充,而接下来的七个请求则会复用该结果。如果使用轮询路由,一半的请求会去往没有缓存前缀的副本,这四个请求将重新执行预填充工作,从而浪费宝贵的 GPU 算力。

我们路由器的任务是追踪哪个副本见过哪个块哈希。一个重要的细节是,这些哈希是链式连接的,因此第 3 个块的哈希代表的是第 1、2 和 3 个块,而不仅仅是第 3 个块。这反映了因果注意力的特性:只有当第 1 和第 2 个块也相同时,第 3 个块的 KV 才是有效的。我们还将适配器名称作为链的种子,因为 KV 缓存还取决于生成它的适配器:为策略 v3 缓存的前缀对于策略 v4 毫无用处!

您的浏览器不支持视频标签。

两个提示和四个请求在两个副本上的路由决策示意图:16-token 块、链式哈希、公共前缀、一次亲和命中(affinity hit)和一次溢出(spill)。

视频演示了选择副本的整个决策过程。以下步骤通过一个真实的 135-token 补全请求示例(来自 Sanity 数据集的问题)进行说明:

  1. 将提示拆分为块。路由器接收 token ID,并将其切割成 16-token 的块,就像 vLLM 一样。它只对完整的块进行哈希处理,因此这里忽略了最后 7 个 token。

  2. 对前缀进行哈希。每个块都与前一个哈希进行哈希运算,从适配器种子开始。因此 h3 按顺序标识了第 1、2 和 3 个块。两个具有相同前 k 个块的提示会得到相同的哈希值直到 hk。一旦某个块发生变化,其后的所有哈希值也会随之改变。这就是我们使用适配器名称对哈希进行种子的原因:在 trl-policy-v4 下的相同提示会从另一个种子开始,无法与 v3 的条目匹配。这正是我们想要的,因为旧的 KV 块是使用不同的权重计算的。

  3. 比较两个提示。问题 1 有 103 个 token。两个提示都以相同的 23-token 聊天模板开头。它们的第一块是相同的,但第二块已经包含了问题文本。从那里开始哈希值不同。

  4. 存储所有者。对于每个哈希,路由器会记住哪些副本服务了它以及哪些哈希紧随其后(我们将后继集合限制为两个,因为我们只需要知道一个块有一个还是多个延续)。经过几个提示后,模板块 h1 由两个副本共同拥有,并且已经有多个后继;h2 到 h8 仅由 A 拥有,且每个都有一个单一的后继;而 h2' 到 h6' 仅由 B 拥有。

在实践中,运行中的每个提示都以相同的 token 开头。在这里,它是聊天模板和系统提示,相当于所有 1,460 个问题中的前 23 个 token。在智能体(agent)设置中,它将是工具描述;在多轮环境中,它将是共享的对话历史。这些块会在几秒内进入每个副本的缓存,因此基于它们进行匹配无法告诉我们特定提示位于何处。

如果一个块由每个副本都服务过,或者它有多个后继,则该块被视为公共块。在路由期间会忽略公共块,因为它们无法标识特定的提示。下文将对此进行更多说明!

  1. 选择一个副本。路由器计算每个副本上匹配的起始块数量,并移除公共前缀。剩下的就是该提示特有的块数。然后:

如果一个副本包含特定的块,并且它没有过载(即其请求数最多比负载最低的副本多 8 个),则该请求会被发送到该副本。我们称之为亲和命中(affinity hit)。

如果一个副本包含特定的块,但其请求数超过负载最低副本 8 个以上,我们将放弃缓存,并将请求发送给负载最低的副本。我们称之为溢出(spill)。

如果没有任何副本包含特定的块,则说明这是一个新的提示(prompt)。该请求会被发送到负载最低的副本;若存在多个负载相同的副本,则采用轮询方式分配。我们称之为未匹配(unmatched)。

以下是将该规则应用于四个请求的示例。从以下状态开始:

原文链接:https://huggingface.co/blog/asyncgrpo-lora-hfjobs
来源:Hugging Face
以上内容由 AI 自动翻译,仅供参考。
← 返回简报