TL;DR
AsyncGRPOTrainer现在可以训练一个 LoRA 适配器，并仅将该适配器同步到 vLLM（TRL v1.14）。- 一个 rank-1 适配器只有几兆字节，因此它可以通过挂载在每个 Job 中的 Storage Bucket 传输，而不必经由 NCCL。训练器和 vLLM 副本作为独立的 Hugging Face Jobs 运行在不同的机器上。
- 副本前端的一个小型代理会添加认证头，将每次 rollout 路由到已经持有其 KV 前缀的副本，并将适配器加载广播到每个副本。
- AsyncGRPO 的指标揭示了瓶颈所在。五次运行采用相同的配方，500 步的耗时从 3 小时 27 分钟缩短到 53 分钟。
LoRA 支持最近落地到了 TRL 的AsyncGRPOTrainer随着PR #7017，并随 TRL v1.14 一同发布。异步训练器现在可以训练适配器而非完整模型，并且只将 LoRA 适配器同步到 vLLM。本文介绍了一个基于它构建的真实项目，在该项目中训练和推理不再共享同一台机器。
LoRA 训练特别适合 RL，正如 Thinking Machines 的博客 LoRA Without Regret 所示。他们表明，LoRA 在策略梯度 RL 中可以与全量微调相媲美，即使 rank 为 1 也是如此。这源于一个事实：优势函数每个 episode 只提供 ~O(1) 比特的信息，因此从信息总比特数的角度来看，每一步并没有太多可学习的内容。一个 rank-1 适配器就有足够的容量来吸收它。
LoRA 训练还会带来系统层面的影响。对于一个 1.5B 模型，rank-1 适配器只有几 MB，而完整模型约为 3 GB。我们不必在每次更新后把完整的策略发送给推理 worker，只需发送适配器即可。vLLM 还可以同时加载多个适配器。旧的 rollout 用它们启动时的策略跑完，而新的 rollout 则使用最新的策略。
TRL 的 AsyncGRPOTrainer 已经将训练和生成分离。训练器和 vLLM 可以运行在不同的机器上，并以各自的速率运行。在单节点或集群环境中，只要两个进程共享文件系统或能组成 NCCL 通信组，这就很容易实现。
我们想要的是用 Hugging Face Jobs 运行同样的配置。本质上，一个 HF Job 就是运行在一台 VM 上的一个容器。这意味着一个 Job 无法拉起多个节点（至少目前如此）来同时承载一个训练器和一组 vLLM 服务器（每个节点最多只能用到 8xH200）。AsyncGRPOTrainer 正是为这种规模而构建的，所以问题就变成了：如果我们放弃训练器和推理服务器必须共享同一节点这一要求，我们能走多远？
嗯，如果做全权重同步，答案会是"走不远"。每次更新都必须在机器之间传输数 GB 的数据，这正是 NCCL 在密集集群中的用途，但 Jobs 无法跨节点通信。没有共享的本地磁盘，显然也没有共享的 localhost。而使用 LoRA，一次同步只有几 MB。至于文件系统部分，HF Jobs 提供了由 Storage Buckets 支持的卷！这些 bucket 随后可以作为 FUSE 文件系统挂载到每个 Job 中，足以充当节点之间的共享文件系统。Job 之间完全不需要网络通路。
这套配置最终相当小巧：
- 一个 trainer Job，运行
AsyncGRPOTrainer并使用 LoRA（以及 FSDP，稍后详述）， - 两个 vLLM Job，各自服务基础模型以及 trainer 最近发布的任何适配器，
- 一个 Storage Bucket，以相同路径挂载到这三者中，适配器正是通过这种方式从 trainer 传到服务器，
- 一个 代理服务器。我们会深入探讨为什么需要它，但从高层来看，我们需要一个代理，将每次 rollout 路由到最有可能持有其 KV cache 的副本，并将每次适配器更新广播到所有 vLLM 副本。
架构：利用 Hugging Face Jobs 和 Storage Buckets 🪣
AsyncGRPOTrainer 中新增的仅适配器同步路径是这样工作的。trainer 不向 vLLM 发送张量。每隔几个优化器步骤，它就把适配器保存到 <output_dir>/.vllm_lora/trl-policy-v{N} 下，通过原子重命名发布该目录，然后将其路径发送到 vLLM 的 /v1/load_lora_adapter 端点。vLLM 从磁盘加载这些文件，因此 rollout worker 随后可以请求 model="trl-policy-v{N}"。
这就是 vLLM 中运行时适配器加载的现有工作方式。该端点接收的是路径，而不是张量，因此训练器和服务器需要共享同一个文件系统。在 Slurm 集群上，那就是网络文件系统。在 Jobs 上，我们通过将 Storage Bucket 挂载为卷，使其在每个 Job 中的相同路径下可用，从而实现同样的效果，正如我们之前提到的。在底层，它使用了 hf-mount，将存储桶作为 POSIX 文件系统暴露在容器内部：
hf jobs run ... -v hf://buckets/aminediroHF/asyncgrpo-lora-buckets:/lora ...
TRL 或 vLLM 无需为此做任何改动。训练器写入 /lora/<run>/.vllm_lora/，服务器从同一路径读取。POST 请求中发送的路径在每个容器内部已经是有效的。
请注意，我们还将检查点和最终适配器存储在存储桶中。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 块仍然会匹配，因此 prefill 不会被重做，一个 rollout 可能从某个策略版本获取其前缀，而从下一个版本获取其 decode。训练器将无从分辨，而这会表现为 ratio 偏离 1。带版本的名称使这成为不可能：一个名称始终意味着同一组权重，而缓存的前缀永远无法匹配更新的版本。
数据集选择：Sanity 集
我们选择了 sail/Sanity-Test-R1D-1.5B，即来自 Defeating the Training-Inference Mismatch via FP16（Qi 等人，2025）的数据集。复现代码在 sail-sg/Precision-RL 中。
作者用 DeepSeek-R1-Distill-Qwen-1.5B 为每道 MATH 题生成了 40 个答案。他们保留了成功率在 20% 到 80% 之间的题目，最终得到 1,460 道题。这个数据集非常适合做 RL 验证，因为这些题目对该模型来说既不是已经解出来的，也不是完全没希望的，这意味着模型可以在训练早期获得良好的信号并加以改进。
作为一个稳健的端到端测试，这非常棒：如果某个 vLLM 副本悄悄以适配器名称提供基础模型，我们希望能在几十步内从曲线上看到这一点。此外，这个数据集足够小，不到两小时就能完整跑一遍。
我们还采用了论文 LoRA 脚本中的超参数，见 oat/scripts/lora：Qwen/Qwen2.5-Math-1.5B，LoRA rank 为 1、alpha 为 2，学习率为 4e-5，每个提示词 8 个样本，每步 128 个补全，最多生成 3,000 个 token，上下文为 4,096 个 token。
训练器
训练器使用同一个 vllm/vllm-openai:v0.27.1 镜像，并在其上安装了 TRL。我们当时运行的是 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的情况，会回退到合并权重同步，并给出警告。日志中应包含Adapter-only vLLM sync enabled。
代理
现在进入有趣的部分。我们需要在训练器和 vLLM Jobs 之间加一个代理，原因有两个：
暴露的 Job 端口要求每个请求都带上
Authorization: Bearer <HF token>头。代理正是添加该头的地方，因此 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 cache。由于注意力是因果的，一个 token 的 KV 只取决于它之前的 token，而不取决于它之后的 token。因此，两个共享前缀的请求会共享该前缀的 KV，而一个已经将其缓存在本地的副本可以完全跳过那部分 prefill。现在整个关键就在于找到那个副本，这样请求就能受益于落到已经见过其前缀的副本上。
vLLM 将其 prefix KV cache 以 16 个 token 为一个块进行存储。由于 GRPO，rollout worker 会发送 G 个具有相同提示词的请求（在我们的场景中是 G=8）。如果它们全部到达同一个副本，第一个请求会计算 prefill，接下来的七个请求则复用它。在轮询路由下，其中一半会发往没有缓存该前缀的副本，那四个请求将重新执行 prefill 工作，浪费宝贵的 GPU 算力。
我们的路由器的职责是追踪哪个副本已经见过哪个 block hash。一个重要细节是，这些哈希是链式的，因此第 3 个块的哈希代表的是第 1、2、3 个块，而不仅仅是第 3 个块。这对应了因果注意力：第 3 个块的 KV 只有在第 1 和第 2 个块也相同的情况下才有效。我们还用 adapter 名称作为链的种子，因为 KV cache 也取决于生成它的 adapter：为 policy v3 缓存的前缀对 policy v4 毫无用处！
视频完整演示了选择副本的整个决策过程。以下步骤通过一个真实的 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 的聊天模板开头。它们的第一个块完全相同，但块 2 已经包含了问题文本。从那里开始哈希就不同了。
4. 存储所有者。 对于每个哈希，路由器会记住哪些副本为其提供了服务，以及哪些哈希排在它之后（我们将后继集合上限设为两个，因为我们只需要知道一个块是有一个后续还是多个后续）。经过几个提示词之后，模板块 h1 由两个副本共同拥有，并且已经有多个后继，h2 到 h8 仅由 A 拥有且各自只有一个后继，而 h2' 到 h6' 仅由 B 拥有。
在实践中，一次运行中的每个提示词都以相同的 token 开头。在这里，就是聊天模板和系统提示词，它们构成了全部 1,460 个问题的前 23 个 token。在智能体场景中，它会是工具描述；在多轮环境中，它会是共享的对话历史。这些块在几秒内就存在于每个副本的缓存中，因此对它们进行匹配并不能告诉我们某个特定提示词位于何处。
如果一个块被每个副本都服务过，或者它有不止一个后继，那么它就是公共块。公共块在路由过程中会被忽略，因为它们无法标识某个特定的提示词。详见下文！
5. 选择一个副本。路由器会统计每个副本上有多少个前导块匹配，并移除公共前缀。剩下的就是该提示词特有的块数量。然后：
- 如果一个副本拥有特有块，并且它没有被淹没——也就是说它最多比负载最低的副本多 8 个请求——那么请求就会发往那里。我们称之为一次亲和命中。
- 如果一个副本拥有特有块，但它比负载最低的副本多出 8 个以上的请求，我们就放弃缓存，把请求发送到负载最低的副本。我们称之为一次溢出。
- 如果没有任何副本拥有特有块，这就是一个新提示词。它会发往负载最低的副本，若出现并列则采用轮询。我们称之为未匹配。
下面是将该规则应用于四个请求的过程。从副本 A 和 B 都各有 3 个请求在途的状态开始，并且两个副本上都只有模板块 h1 是已知的。
- 请求 1，问题 0，rollout 1。 两个副本都匹配一个块，即模板，而该块是共有的。因此没有任何特定块在任何地方匹配。该请求未被匹配，两个副本负载相同，轮询将其发送到 A。路由器记录
h2到h8归 A 所有。A 现在有 4 个请求在途。 - 请求 2，问题 0，rollout 2。 相同的提示词。A 匹配全部 8 个块，B 只匹配模板。去掉那一个共有块后，A 有 7 个特定块，B 没有。A 只比 B 领先 1 个请求，远在 8 的限制之内，因此该请求发往 A。这是一次亲和性命中：A 的 KV cache 中已经拥有整个提示词。
- 请求 3，问题 1，rollout 1。 一个新的提示词。两个副本都只匹配模板块，因此没有任何特定块匹配。该请求未被匹配，发往负载最低的副本 B，B 有 3 个在途请求，而 A 有 5 个。路由器记录
h2'到h6'归 B 所有。 - 请求 4，问题 0，第 9 次 rollout。假设此时 A 已有 12 个请求在途，而 B 回到了 3 个。A 仍然持有那 7 个特定块，但它现在领先 B 9 个请求，超过了上限。该请求溢出到 B。B 对问题 0 执行一次预填充，路由器将
h2到h8记录为同样归 B 所有。现在每个副本都已服务过这些块，因此问题 0 也变为公共块，其后续 rollout 仅按负载进行分配。
6. 复用预填充。请求 2 正是做这一切的原因。它复用了请求 1 计算出的预填充：块 1 到 8 已经在 A 的 KV cache 中，因此 A 直接跳到解码补全。如果它去了 B，B 会重新预填充全部 135 个 token，而 A 的缓存则闲置不用。请求 3 说明了为什么需要 common 规则。没有它，共享的聊天模板会让每个新提示词看起来都像是缓存命中。请求 4 则将负载保持在有界范围内。节省一次预填充并不值得让某个副本远远落后。
def choose(self, upstreams, model, prompt):
hashes = self.block_hashes(model, prompt)
matched = self.matched_prefix(hashes)
common = self.common_prefix_len(hashes)
specific = [max(0, m - common) for m in matched]
least = min(u.inflight for u in upstreams)
best = max(range(self.n), key=lambda i: (specific[i], -upstreams[i].inflight))
if specific[best] > 0 and upstreams[best].inflight - least <= self.cfg.imbalance:
pick = best
else:
candidates = [i for i in range(self.n) if upstreams[i].inflight == least]
pick = candidates[self.rr % len(candidates)]
self.rr += 1
...record `pick` as an owner of every block, and each block's successor...
return upstreams[pick]
common 前缀是最麻烦的部分。每个请求都以相同的系统提示词和聊天模板开头。简单的最大前缀匹配会让第一个副本几乎匹配上每一个新提示词。我们改为通过扇出来检测共享前缀：一个有多个不同后继的块是公共块，而一个总是通向同一个后继的块则属于某个特定提示词。只有该公共前缀之后的块才计入亲和性。
广播适配器
代理还需要将适配器加载广播到每个副本。我们将该操作视为全有或全无。每个副本都有自己的存储桶挂载，因此它们不一定会在完全相同的时刻看到新适配器。No adapter found for <path> 错误通常意味着某个存储桶挂载尚未跟上，我们只重试该副本。对于任何其他错误，我们会从已接受该适配器的副本上将其卸载，以确保某个策略名称绝不会只存在于部分副本上。
async def load_one(u):
while True:
status, _, out = await send(u, "POST", "/v1/load_lora_adapter", headers, body)
if status == 200 or "No adapter found" not in out.decode() or time.monotonic() > deadline:
return u, status, out
await asyncio.sleep(cfg.lora_retry_s)
results = await asyncio.gather(*(load_one(u) for u in ups))
if any(st != 200 for _, st, _ in results):
await asyncio.gather(*(send(u, "POST", "/v1/unload_lora_adapter", headers, unload) for u, st, _ in results if st == 200))
return web.Response(status=504 if timed_out else st, text="rolled back on the others")
我们同样以相同方式广播 /pause、/resume 和 /v1/unload_lora_adapter。/health 只有在每个副本都健康时才返回 200。/server_info 和 /v1/models 只需要一个应答。从 TRL 的角度来看，代理就是一个单一的 data_parallel_size=1 服务器，因此它选择仅适配器同步。
我们起初怀疑 Python asyncio 代理是否会成为瓶颈。它并不会（至少在这个规模下不会）。最多有 128 个非流式 JSON 请求在途，而路由只计算几个哈希。一个线程就能轻松处理。一个需要处理更多流量的更精细路由器，可能需要用更快的语言来编写（我们看到你了 🦀）。
完整运行结果
下面的数字来自训练器的记录指标，位于 trackio 上。该次运行使用Qwen/Qwen2.5-Math-1.5B、LoRAr=1在all-linear上，每步 128 个补全，每个提示词 8 次 rollout。它运行 500 步，每 50 步保存一次检查点。训练器使用一个h200x2Job，两个 vLLM 副本各使用一个h200Job。运行这三者每小时约花费 $20。
权重同步
| 每次同步，训练器时钟，126 次同步 | 之前 | 现在（p50） |
|---|---|---|
| 整个同步 | 30.8 s | 8.5 s（最小 6.6，最大 9.2） |
| 其中：暂停两个副本 | 0.3 s | 0.3 s |
| adapter 全收集并保存到 bucket | 0.6 s | 1.1 s |
| 两个副本都接受 adapter | ~29 s | ~7 s |
全部 252 次 adapter 加载均成功 🎉：126 次同步乘以 2 个副本。其中 6 次在第二次尝试时成功，246 次在第三次尝试时成功。
路由
在运行结束时，经过 64,728 次 rollout 之后，代理的计数器显示：
routed [31928, 32800] affinity 54712 spilled 820 unmatched 9196
每个提示词有 8 次 rollout，八个请求中至少有一个必然是冷启动。因此理论最小值是 12.5%。路由器得到 14.2% 未匹配请求、84.5% 亲和性命中和 1.3% 溢出。除非我们开始基于真实测量负载、使用更深入的推理侧指标来观察每个副本的负载，否则这里已经没有多少可优化的空间了。
时间都花在哪里了
第一个配置有一个相当明显的问题：瓶颈在训练器，而不是生成。在 500 步中，我们有：
| 每个优化器步，p50 | |
|---|---|
| 步 | 22.9 s |
| 前向 + 反向 | 21.9 s |
| 等待 rollout | 0.02 s |
| rollout 队列占用率 | 512 中的 476 |
| 训练器 MFU | 3.9 % |
部署队列始终处于满载状态，而工作进程主要被背压阻塞。在这种配置下，第二个副本基本上毫无用处。我们将在后文中看到，我们如何通过一系列运行在训练和生成之间转移瓶颈，使运行速度提升了 3.9 倍。
奖励
图 1. trackio 运行 r1-dp2。面板：reward 及其 20 步滚动均值和 50 步分块均值，以及 ratio 在 0.99 到 1.01 的坐标轴上。奖励在 500 步内从 0.15 攀升至 0.44；ratio 全程保持在 0.9993 到 1.0004 之间。
500 步耗时 3 小时 27 分钟。平均奖励从前 20 步的 0.145 上升到后 20 步的 0.438。对于本次测试更重要的是，ratio 每一步都保持在 1.000！由 vLLM 服务的策略始终与训练器用于给 rollout 打分的策略一致。这在全部 126 次同步中都成立。平均陈旧度为 1.5 个策略版本，最大值为 4。trackio 仪表盘中有完整曲线。
我们有了无可辩驳的证据，证明 LoRA AsyncGRPO 有效！现在让我们通过查看近期详细的 AsyncGRPO 指标，来看看如何改进训练运行。
追逐瓶颈的乒乓
异步 RL 是训练与生成之间的流水线。如果一侧跟不上，让另一侧变快也无济于事。幸运的是，在 AsyncGRPOTrainer 中我们添加了足够多的计时和指标，可以直接看到这一点。
所有有用的指标都记录在已记录指标部分。perf/rollout_wait_s告诉我们训练器等待样本的时间有多长。rollout/backpressure_s告诉我们生成在 rollout 队列中等待空间的时间有多长。两者截然相反，不应同时处于高位。结合队列大小，它们能告诉我们哪一侧是瓶颈。
我们运行了五次实验。每一次都从上一轮运行仪表盘中可见的问题出发。除非另有说明，模型、配方和三 Job 布局保持不变。以下名称是 trackio 运行名称。
阅读仪表盘
我们保持以下四组指标可见：
perf/step_s和perf/fwd_bwd_s：优化器一步需要多长时间，以及其中前向+反向传播占多少。如果一步的耗时几乎与前向+反向传播相当，那么训练器显然受算力限制。perf/rollout_wait_s：训练器在开始一步之前等待样本的时间有多长。接近零意味着生成领先于训练。样本可以立即用于训练。sample/rollout_queue_size对比queue_maxsize：两侧之间的缓冲区。满意味着生成正在被限流；空意味着训练器正在挨饿。rollout/backpressure_s和rollout/score_block_s：rollout worker 因该缓冲区已满而被阻塞了多长时间。该 worker 是一个两阶段流水线：生成阶段将完成的组交给评分阶段，评分阶段将已评分的样本推入 rollout 缓冲区。当缓冲区已满时，评分阶段无法入队并发生阻塞，这就是rollout/backpressure_s。随后评分阶段停止排空自己的输入队列，因此生成阶段也无法移交下一组，这就是rollout/score_block_s。两者是同一个停滞，先在评分阶段被观察到，然后反向传播到生成阶段。
诊断很简单：队列已满且 rollout 等待为零、背压很高，意味着训练器太慢。队列为空且 rollout 等待上升、没有背压，意味着生成太慢。将 perf/mfu_wall_clock 与 perf/mfu_fwd_bwd 进行比较，还能看出训练器 GPU 有多少时间花在等待上而非训练上。
运行 1，r1-dp2：一个跟不上节奏的训练器
图 2. trackio 运行 r1-dp2。面板：perf/step_s、perf/fwd_bwd_s、sample/rollout_queue_size、rollout/backpressure_s。步进时间与前向+反向几乎完全重叠；队列几乎钉在 512 中的 476 附近，背压从未低于每个 rollout 组 11 秒：受训练器制约。
perf/step_s 为 22.9 秒，perf/fwd_bwd_s 为 21.9 秒。前向与反向传播占用了步进时间的 96%。队列始终处于满载状态，训练器等待 rollout 仅需 0.02 秒，而 rollout worker 每组因背压阻塞耗时 15 秒。两个 vLLM 副本的生成速度快于训练器的消费速度。所报告的 4.6k tokens/s 并非它们的实际极限；它们只是无处安放更多输出。
批次指标解释了糟糕的 3.9% MFU。batch/microbatches_per_step 为 64，batch/samples_per_row 为 1.0。每个 rank 处理一条约 1.2k tokens 的序列，每步处理 64 次。这来自参考配方的 per_device_train_batch_size=1。对于 H200 上的 1.5B 模型而言，这完全是延迟受限的。
运行 2，r1-dp2-tb16k：打包微批次
修复方法不是改变批次大小。我们保持每个优化器步进 128 个补全，只改变它们在 GPU 上的排布方式：不再每个微批次放一条序列，而是将多条序列密集地打包进每一行。训练器通过 token-budget batching 支持这一点。借助 token_budget > 0，它将多个样本打包进每个 rank 的一行无填充行中。一个优化器步进处理 gradient_accumulation_steps 行。我们设置 token_budget=16384 和 gradient_accumulation_steps=6。
图 3. trackio 运行 r1-dp2 和 r1-dp2-tb16k 在其前 154 步的叠加对比。面板：batch/samples_per_row、batch/microbatches_per_step、perf/step_s、perf/fwd_bwd_s、perf/mfu_fwd_bwd、rollout/generated_tok_s。打包使每行样本数从 1 增至 13，微批次从 64 降至 6，步进时间从 23 秒降至 5.9 秒，生成速度从 4.2k 提升至 27.5k tok/s，而 vLLM 侧毫无变化。
batch/samples_per_row 从 1.0 提升到约 12.7，微批数量从 64 降到 6。各行填充率达到 95%！前向和反向传播从 21.9 秒降到 5.6 秒，而 MFU 从 3.9% 升到 19%。现在每一步大约训练 150 个样本，因为各行的打包效果比平均长度估计所预测的更好。
生成速度也从 4.6k tokens/s 跃升到 25k tokens/s，尽管我们在 vLLM 侧没有做任何改动。队列不再持续满载，因此副本终于能够运行起来。这就是为什么我们不喜欢孤立地优化流水线各阶段。我们需要谨慎地评估整个系统，因为一个缓慢的阶段可能会掩盖它之前所有环节的真实性能。
第 3 次运行，r1-dp2-tb16k-nockpt：停止重新计算前向
perf/fwd_s 为 1.34 秒，而 perf/fwd_bwd_s 为 5.6 秒。正常的反向传播开销大约是前向的两倍，而在基础权重冻结的情况下应该更接近一倍。3.2 的比值很可疑。
原因是 AsyncGRPOConfig 默认使用 gradient_checkpointing=True。每个微批在反向传播期间都会重新计算其前向。这也解释了为什么一个 16k token 的行在 141 GB 的 H200 上只占用 25 GB！对这个训练器来说这是一种内存优化，但在这个特定场景下我们不需要它：模型足够小，可以在保留激活值用于反向传播的情况下装入显存。
图 4. trackio 运行 r1-dp2-tb16k 和 r1-dp2-tb16k-nockpt 在其前 134 步上的叠加对比。面板：perf/fwd_s、perf/fwd_bwd_s、perf/weight_sync_s、sample/rollout_queue_size、perf/rollout_wait_s、perf/mfu_fwd_bwd。前向+反向下降了一次前向的时间；队列从约 420 降至约 60，rollout 等待从 0.02 s 升至 0.5 s：瓶颈转移到了生成环节。
使用 gradient_checkpointing=False 后，前向和反向降至 4.6 s，几乎正好少了一次前向的时间，MFU 达到 23%。队列现在降至 71，rollout 等待从 0.04 s 升至 0.6 s。训练器消耗样本的速度快于两个副本生成样本的速度。我们成功将瓶颈转移到了生成环节。
这又暴露出两项额外开销。每四步进行一次的 7.6 秒权重同步现在占用了 25% 的墙钟时间。而当每步耗时 23 秒时，这一比例仅为 8%。此外，反向传播仍然比前向传播慢 2.5 倍。在基础权重冻结的情况下，每步约有 2 秒的时间不符合正常的模型计算特征。
运行 4，r1-dp3-tb16k-nockpt：三个副本，以及一个意外发现
由于生成现在太慢了，我们增加了第三个副本。我们还将代理中的适配器重试间隔从 2 s 缩短到 0.5 s，并禁用了 fsdp_reshard_after_forward，以检查 FSDP2 的重新聚合是否导致了反向传播中多出的那 2 秒。
图 5. trackio 运行 r1-dp2-tb16k-nockpt 和 r1-dp3-tb16k-nockpt 在其前 134 步上的叠加对比。各面板：perf/weight_sync_s、rollout/generated_tok_s、rollout/inflight、perf/fwd_bwd_s。同步从 7.6 s 降至 5.8 s；生成和前向+反向没有变化；rollout/inflight 在两次运行中都显示为 128，这正是第三个副本撞上的上限。
权重同步从 7.6 s 降至 5.8 s，因此更短的 retry 起了作用。前向和反向保持在 4.6 s，这排除了重新分片（resharding）的可能。生成从 25k 仅提升到 26k tokens/s。第三个副本基本上什么都没做。
原因就在 rollout/inflight 里：每次运行都是 128。代理显示这些请求在三个副本上拆分为 44 + 43 + 41。max_inflight_tasks 限制的是整个 rollout worker 的并发数，而不是每个副本的并发数。H200 上的一个 1.5B 模型处理 43 个和 130 个并发序列时，每个 token 的成本几乎相同。把 128 个请求分到三块 GPU 上，吞吐量与分到两块 GPU 上几乎一样。
所以 vLLM 并不是瓶颈。我们自己客户端侧的常量才是。我们当初把它设得很保守，因为不知道数百个长 HTTPS 请求经过公共 Jobs 代理会表现如何。而到这时，已有 130,000 次 rollout 完成通过了它，没有出现一次传输错误。
运行 5，r1-dp3-inflight384：解除对在途请求的上限
max_inflight_tasks=384 和 queue_maxsize=768，没有别的。
图 6。全部五次 trackio 运行（r1-dp2、r1-dp2-tb16k、r1-dp2-tb16k-nockpt、r1-dp3-tb16k-nockpt、r1-dp3-inflight384）叠加显示，x 轴以步数为单位。各面板：perf/step_s、reward、sample/rollout_queue_size、sample/staleness_mean。在整个序列中，单步时间从 22.9 s 降至 4.8 s，而奖励曲线始终彼此重合；最后一次运行的队列重新填充至 768 中的约 690，其陈旧度稳定在 2。一旦仪表盘回答了问题，运行 2 到 4 就被提前停止。
在有 384 个请求在途的情况下，每个副本分到 128 个。队列迅速填充到 768 中的约 690 并保持在那里。背压回到 5 秒，rollout 等待降至 0.03 秒。训练再次成为瓶颈！前向和后向传播耗时 4.6 秒，权重同步摊销增加 1.5 秒，单步时间中位数为 4.8 秒。
平均陈旧度从 1.5 个版本升至 2.0 个版本，因为样本在更大的队列中等待更久。这仍低于 max_staleness=4，且 ratio 仍非常接近 1.000。
记分板
| 500 步 | 运行 1 r1-dp2 | 运行 5 r1-dp3-inflight384 |
|---|---|---|
| 挂钟时间 | 3 小时 27 分钟 | 53 分钟 |
perf/step_s，p50 | 22.9 s | 4.8 s |
perf/fwd_bwd_s，p50 | 21.9 s | 4.6 s |
perf/mfu_fwd_bwd | 3.9 % | 23.5 % |
batch/samples_per_step | 128 | 168 |
| 训练样本数 | 64 000 | 84 078 |
perf/weight_sync_s，p50 | 8.5 s | 6.2 s |
sample/staleness_mean | 1.5 | 2.0 |
| 奖励，前 20 → 后 20 步 | 0.145 → 0.438 | 0.145 → 0.416 |
图 7. trackio 运行 r1-dp2 和 r1-dp3-inflight384，奖励相对于自首次优化器步骤以来的挂钟分钟数。相同的配方，相同的 500 步，相同的最终奖励；运行 5 在 52 分钟内达到，而不是 3 小时 26 分钟。
最终运行快了 3.9 倍，并在多 31 % 的样本上进行训练，奖励曲线基本相同。打包、禁用检查点以及提高在途限制带来了差异。在每种情况下，仪表盘都在前十分钟内清楚地显示了这一点。
试一试
git clone https://github.com/AmineDiro/hfjobs-lora-buckets && cd hfjobs-lora-buckets
hf auth login
MAX_STEPS=20 RUN_TAG=smoke ./run_all.sh --wait
MAX_STEPS=500 ./run_all.sh --wait
TOKEN_BUDGET=16384 GRAD_ACCUM=6 GRADIENT_CHECKPOINTING=0 PROXY_LORA_RETRY_S=0.5 \
MAX_INFLIGHT=384 QUEUE_MAXSIZE=768 MAX_STEPS=500 ./run_all.sh --wait
参考文献
- John Schulman 等人，LoRA Without Regret，Thinking Machines Lab，2025 年 9 月。论证 rank-1 LoRA 在策略梯度 RL 中可与全量微调相媲美，以及原因。
- TRL，
AsyncGRPOTrainer及其 记录的指标。 - TRL PR #7017：对
AsyncGRPOTrainer的 PEFT/LoRA 支持，并带有仅适配器 vLLM 同步，已在 TRL v1.14 中发布。 - Hugging Face Jobs 和 Storage Buckets；
hf-mount。 hf-mount-repro：30 秒负缓存卡顿的双脚本复现。- 本文中每次运行的 trackio dashboard。
- Penghui Qi、Zichen Liu、Xiangxin Zhou、Tianyu Pang、Chao Du、Wee Sun Lee、Min Lin，Defeating the Training-Inference Mismatch via FP16，arXiv:2510.26788，2025。Sanity 数据集
sail/Sanity-Test-R1D-1.5B与 LoRA 配方sail-sg/Precision-RL的来源，oat/scripts/lora/bf16_grpo_tis_lora.sh。
@article{qi2025precisionrl,
title={Defeating the Training-Inference Mismatch via FP16},
author={Qi, Penghui and Liu, Zichen and Zhou, Xiangxin and Pang, Tianyu and Du, Chao and Lee, Wee Sun and Lin, Min},
journal={arXiv preprint arXiv:2510.26788},
year={2025}
}
TL;DR
AsyncGRPOTrainercan now train a LoRA adapter and sync only that adapter to vLLM (TRL v1.14).- A rank-1 adapter is a few megabytes, so it can travel through a Storage Bucket mounted in every Job instead of over NCCL. The trainer and the vLLM replicas run as separate Hugging Face Jobs on separate machines.
- A small proxy in front of the replicas adds the auth header, routes each rollout to the replica that already holds its KV prefix, and broadcasts adapter loads to every replica.
- The AsyncGRPO metrics show where the bottleneck sits. Five runs take the same recipe from 3 h 27 min to 53 min for 500 steps.
LoRA support recently landed in TRL's AsyncGRPOTrainer with PR #7017, and ships with TRL v1.14. The asynchronous trainer can now train an adapter instead of the full model, and it syncs only the LoRA adapter to vLLM. This post covers a real-world project built on top of it, where training and inference no longer share a machine.
LoRA training is particularly suited for RL, as shown in Thinking Machines's blog LoRA Without Regret. They show that LoRA can match full fine-tuning for policy-gradient RL, even with rank 1. This stems from the fact that the advantage function only gives ~O(1) bits of information per episode, so there is not that much to learn from each step, from a total-bits-of-information point of view. A rank-1 adapter has enough capacity to absorb it.
There is also a systems consequence of LoRA training. A rank-1 adapter for a 1.5B model is a few megabytes, while the full model is around 3 GB. Instead of sending the full policy to the inference workers after every update, we can just send the adapter. vLLM can also keep several adapters loaded at once. Old rollouts finish with the policy they started with, while new rollouts use the latest one.
TRL's AsyncGRPOTrainer already separates training and generation. The trainer and vLLM can run on different machines and at their own speed. This is easy in a single-node or cluster setting where both processes share a filesystem or can form an NCCL group.
What we want is to run the same setup with Hugging Face Jobs. Essentially, an HF Job is one container running on one VM. This means that one Job cannot spawn multiple nodes (at least for now) to hold a trainer and a fleet of vLLM servers (we are limited to 8xH200 at most per node). The AsyncGRPOTrainer is built for exactly that kind of scale, so the question became: how far can we get if we drop the requirement that the trainer and the inference servers share a node?
Well, with a full-weight sync, the answer would be "not far". Every update would have to move gigabytes between machines, which is what NCCL is for in a dense cluster, but Jobs can't communicate across nodes. There is no shared local disk and obviously no shared localhost. With LoRA, a sync is only a few megabytes. For the filesystem part, HF Jobs provide volumes backed by Storage Buckets! These buckets can then be mounted as a FUSE filesystem in every Job and are enough to work as a shared FS between nodes. No network path between the Jobs is needed at all.
The setup ended up being quite small:
- a trainer Job running
AsyncGRPOTrainerwith LoRA (and FSDP, more on that later), - two vLLM Jobs, each serving the base model plus whatever adapter the trainer last published,
- a Storage Bucket mounted in all three at the same path, which is how the adapter gets from the trainer to the servers,
- a proxy server. We'll dive deeper into why we need one, but at a high level we need a proxy that routes each rollout to the replica most likely to hold its KV cache, and broadcasts every adapter update to all vLLM replicas.
The architecture: leveraging Hugging Face Jobs and Storage Buckets 🪣
The new adapter-only sync path in AsyncGRPOTrainer works like this. The trainer does not send tensors to vLLM. Every few optimizer steps, it saves the adapter under <output_dir>/.vllm_lora/trl-policy-v{N}, publishes the directory with an atomic rename, then sends its path to vLLM's /v1/load_lora_adapter endpoint. vLLM loads the files from disk, so the rollout worker can then request model="trl-policy-v{N}".
This is how runtime adapter loading already works in vLLM. The endpoint takes a path, not tensors, so the trainer and the server are expected to share a filesystem. On a Slurm cluster, that is the network filesystem. On Jobs, we get the same thing by mounting a Storage Bucket as a volume at the same path in every Job, as we mentioned earlier. Under the hood, it uses hf-mount, which exposes the bucket as a POSIX filesystem inside the container:
hf jobs run ... -v hf://buckets/aminediroHF/asyncgrpo-lora-buckets:/lora ...
Nothing in TRL or vLLM had to change for this. The trainer writes to /lora/<run>/.vllm_lora/ and the servers read from the same path. The path sent in the POST request is already valid inside every container.
Note that we also store the checkpoints and the final adapter in the bucket. The HF Jobs are ephemeral, but a preempted trainer can resume training, as the final adapter is always persisted to the bucket and is never lost when the Job stops.
The three Jobs
The vLLM replicas
Each replica uses one GPU and the stock vllm/vllm-openai image. We only need to enable runtime LoRA loading and reserve enough adapter slots.
The number of adapter slots follows from max_staleness. In AsyncGRPOTrainer, every weight sync bumps the policy version by one, and max_staleness is how many versions a rollout sample may lag behind the current policy before the trainer discards it. With max_staleness=4, a sample generated under trl-policy-v3 is still used for training while the trainer is at v7. A rollout that started under v3 must also be able to finish under v3. So at any moment, vLLM has to serve the current policy plus the four before it. That is why the trainer keeps max_staleness + 1 adapter versions registered and unloads anything older. Each sync loads the new version before it unloads the oldest one, which needs one more slot during the swap. That gives --max-loras 6. With only five, vLLM would silently evict a policy that still has rollouts in flight at every sync.
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
We pin vLLM to v0.27.1. vLLM moves fast, and the flags above and the runtime LoRA endpoints are the ones that version exposes, so treat the version as part of the recipe.
There is another possible design where the trainer keeps only the latest adapter and always publishes it under the same name. We did not go that way, because vLLM keys its prefix cache by adapter name. With a single name, KV blocks computed under the previous weights would still match after the swap, so the prefill would not be redone and a rollout could get its prefix from one policy version and its decode from the next. The trainer would have no way to tell, and it would show up as ratio drifting away from 1. Versioned names make this impossible: a name always means one set of weights, and a cached prefix can never match a newer version.
The dataset choice: the Sanity set
We chose sail/Sanity-Test-R1D-1.5B, the dataset from Defeating the Training-Inference Mismatch via FP16 (Qi et al., 2025). The reproduction code is in sail-sg/Precision-RL.
The authors generated 40 answers for each MATH problem with DeepSeek-R1-Distill-Qwen-1.5B. They kept problems with a success rate between 20% and 80%, yielding 1,460 questions. This dataset is really good for RL validation because the questions are neither already solved nor completely hopeless for that model, meaning the model can get a good early signal to train on and improve.
This is awesome as a robust end-to-end test: if one vLLM replica silently serves the base model under an adapter name, we want to see that in the curve within a few dozen steps. Also, this dataset is small enough to cycle through in less than two hours.
We also take the hyperparameters from the paper's LoRA scripts in oat/scripts/lora: Qwen/Qwen2.5-Math-1.5B, LoRA rank 1 with alpha 2, a learning rate of 4e-5, 8 samples per prompt, 128 completions per step, a maximum of 3,000 generated tokens and a 4,096-token context.
The trainer
The trainer uses the same vllm/vllm-openai:v0.27.1 image with TRL installed on top. We ran the PR branch at the time; the same code now ships in TRL v1.14. The training script is a normal AsyncGRPOTrainer script. The only Job-specific values are the output directory and the server 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"),
...
)
During initialization, TRL calls
/server_info. If it finds alora_config, it uses adapter-only sync. Configurations vLLM cannot serve directly, such as DoRA,modules_to_save, or a rank above--max-lora-rank, fall back to merged-weight sync with a warning. The log should containAdapter-only vLLM sync enabled.
The proxy
Now onto the fun stuff. We need a proxy between the trainer and the vLLM Jobs for two reasons:
Exposed Job ports require an
Authorization: Bearer <HF token>header on every request. The proxy is where that header gets added, so TRL does not need to know about it.We want more than one GPU generating. On a single vLLM server, the usual way to get that is
--data-parallel-size > 1, but TRL refuses adapter-only sync in that mode, for a good reason: a call to/v1/load_lora_adapteronly reaches the DP rank that answers it, so the other ranks would keep serving the base model under the new policy name. On Jobs the question does not even arise, since each replica is its own machine. So the data parallelism has to live one level up, in something that fans the adapter load out to every replica.
We therefore run a small proxy at 127.0.0.1:8000 on the trainer Job and point TRL to it as if it were a single vLLM server. Besides adding the header, the proxy does two things functionally:
- It sends each completion request to one replica, chosen so that the eight rollouts of a prompt land where their prefix is already cached (details on this below).
- It broadcasts every state-changing request, such as adapter loads, pause and resume, to all replicas, so that a policy name means the same thing everywhere.
Routing rollouts by KV prefix
A quick reminder of why this matters. Generating a completion has two phases with very different workload profiles:
- The prefill processes the whole prompt at once and computes the attention keys and values for every prompt token.
- The decode phase then produces one token at a time, and each new token attends to the keys and values of all the tokens before it.
Those keys and values are the KV cache. Because attention is causal, the KV of a token depends only on the tokens before it, not on what comes after. Two requests that share a prefix therefore share the KV of that prefix, and a replica that already has it in cache can skip that part of the prefill entirely. The whole game now is to find that replica, so a request can benefit from landing on the replica that has already seen its prefix.
vLLM stores its prefix KV cache in blocks of 16 tokens. Because of GRPO, the rollout worker sends G requests with the same prompt (in our case G=8). If they all reach the same replica, the first request computes the prefill and the next seven reuse it. With round-robin routing, half would go to a replica that doesn't have the prefix cached, and those four requests would redo the prefill work and waste valuable GPU compute.
The job of our router is to track which replica has seen which block hash. One important detail is that the hashes are chained, so the hash of block 3 represents blocks 1, 2 and 3, not just block 3. This mirrors causal attention: the KV of block 3 is only valid if blocks 1 and 2 are the same too. We also seed the chain with the adapter name because the KV cache also depends on the adapter that generated it: a prefix cached for policy v3 is useless for policy v4!
The video walks through the entire decision process for choosing a replica. The steps below go through a real 135-token completion request example (from the Sanity dataset problems):
1. Split the prompt into blocks. The router receives token ids and cuts them into 16-token blocks, just like vLLM. It only hashes complete blocks, so the last 7 tokens are ignored here.
2. Hash the prefix. Each block is hashed with the previous hash, starting from the adapter seed. h3 therefore identifies blocks 1, 2 and 3 in order. Two prompts with the same first k blocks get the same hashes up to hk. Once one block changes, every hash after it changes too. This is why we seed the hashes with the adapter name: the same prompt under trl-policy-v4 starts from another seed and cannot match entries from v3. This is what we want because the old KV blocks were computed with different weights.
3. Compare two prompts. Problem 1 has 103 tokens. Both prompts start with the same 23-token chat template. Their first block is identical, but block 2 already contains the problem text. The hashes differ from there.
4. Store the owners. For every hash, the router remembers which replicas served it and which hashes came after it (we cap the successor set at two because we only need to know whether a block has one continuation or several). After a few prompts, the template block h1 is owned by both replicas and already has several successors, h2 to h8 are owned by A only and each has a single successor, and h2' to h6' are owned by B only.
In practice, every prompt in a run starts with the same tokens. Here, it is the chat template and the system prompt, which amount to the first 23 tokens of all 1,460 problems. In an agent setting, it would be the tool descriptions, and in a multi-turn environment it would be the shared conversation history. These blocks are in every replica's cache within seconds, so matching on them tells us nothing about where a particular prompt lives.
A block is common if every replica has served it, or if it has more than one successor. Common blocks are ignored during routing because they do not identify a particular prompt. More on this below!
5. Pick a replica. The router counts how many leading blocks match on each replica and removes the common prefix. What is left is the number of blocks specific to this prompt. Then:
- If one replica has specific blocks and it is not swamped, meaning it is at most 8 requests ahead of the least-loaded replica, the request goes there. We call this an affinity hit.
- If a replica has specific blocks but it is more than 8 requests ahead, we give up on the cache and send the request to the least-loaded replica. We call this a spill.
- If no replica has specific blocks, this is a new prompt. It goes to the least-loaded replica, round-robin on ties. We call this unmatched.
Here is the rule applied to four requests. Start from a state where replicas A and B both have 3 requests in flight, and only the template block h1 is known on both replicas.
- Request 1, problem 0, rollout 1. Both replicas match one block, the template, and that block is common. So nothing specific matches anywhere. The request is unmatched, both replicas are equally loaded, and round-robin sends it to A. The router records
h2toh8as owned by A. A now has 4 requests in flight. - Request 2, problem 0, rollout 2. Same prompt. A matches all 8 blocks, B matches only the template. After removing the one common block, A has 7 specific blocks and B has none. A is only 1 request ahead of B, well within the limit of 8, so the request goes to A. This is an affinity hit: A already has the whole prompt in its KV cache.
- Request 3, problem 1, rollout 1. A new prompt. Both replicas match only the template block, so nothing specific matches. The request is unmatched and goes to the least-loaded replica, B, which has 3 in flight against A's 5. The router records
h2'toh6'as owned by B. - Request 4, problem 0, rollout 9. Suppose that by now A has 12 requests in flight while B is back to 3. A still has the 7 specific blocks, but it is now 9 requests ahead of B, more than the limit. The request spills to B. B prefills problem 0 once, and the router records
h2toh8as owned by B as well. Every replica has now served those blocks, so problem 0 becomes common too, and its later rollouts are placed by load alone.
6. Reuse the prefill. Request 2 is the reason for doing all this. It reuses the prefill computed by Request 1: blocks 1 to 8 are already in A's KV cache, so A skips straight to decoding the completion. Had it gone to B, B would have prefilled all 135 tokens again while A's cache sat unused. Request 3 shows why the common rule is needed. Without it, the shared chat template would make every new prompt look like a cache hit. Request 4 keeps the load bounded. Saving one prefill is not worth letting a replica fall far behind.
def choose(self, upstreams, model, prompt):
hashes = self.block_hashes(model, prompt)
matched = self.matched_prefix(hashes)
common = self.common_prefix_len(hashes)
specific = [max(0, m - common) for m in matched]
least = min(u.inflight for u in upstreams)
best = max(range(self.n), key=lambda i: (specific[i], -upstreams[i].inflight))
if specific[best] > 0 and upstreams[best].inflight - least <= self.cfg.imbalance:
pick = best
else:
candidates = [i for i in range(self.n) if upstreams[i].inflight == least]
pick = candidates[self.rr % len(candidates)]
self.rr += 1
...record `pick` as an owner of every block, and each block's successor...
return upstreams[pick]
The common prefix is the annoying part. Every request starts with the same system prompt and chat template. A simple longest-prefix match would give the first replica a match for almost every new prompt. We detect the shared prefix through fan-out instead: a block with several different successors is common, while a block that always leads to the same successor belongs to a particular prompt. Only the blocks after that common prefix count as affinity.
Broadcasting the adapter
The proxy also needs to broadcast the adapter loads to every replica. We treat the operation as all-or-nothing. Each replica has its own bucket mount, so they do not necessarily see a new adapter at exactly the same time. A No adapter found for <path> error usually means that one bucket mount has not caught up yet, and we retry only that replica. For any other error, we unload the adapter from the replicas that accepted it so that a policy name never exists on only part of the replicas.
async def load_one(u):
while True:
status, _, out = await send(u, "POST", "/v1/load_lora_adapter", headers, body)
if status == 200 or "No adapter found" not in out.decode() or time.monotonic() > deadline:
return u, status, out
await asyncio.sleep(cfg.lora_retry_s)
results = await asyncio.gather(*(load_one(u) for u in ups))
if any(st != 200 for _, st, _ in results):
await asyncio.gather(*(send(u, "POST", "/v1/unload_lora_adapter", headers, unload) for u, st, _ in results if st == 200))
return web.Response(status=504 if timed_out else st, text="rolled back on the others")
We also broadcast /pause, /resume and /v1/unload_lora_adapter in the same way. /health returns 200 only if every replica is healthy. /server_info and /v1/models only need one answer. From TRL's point of view, the proxy is a single data_parallel_size=1 server, so it selects adapter-only sync.
We initially wondered whether a Python asyncio proxy would become a bottleneck. It does not (at least at this scale). There are at most 128 non-streaming JSON requests in flight, and routing only computes a few hashes. One thread handles this easily. A more refined router that needs to handle more traffic would probably need to be written in a faster language (We see you 🦀).
Full run results
The numbers below come from the trainer's logged metrics on trackio. The run uses Qwen/Qwen2.5-Math-1.5B, LoRA r=1 on all-linear, 128 completions per step and 8 rollouts per prompt. It runs for 500 steps and saves a checkpoint every 50 steps. The trainer uses an h200x2 Job and each of the two vLLM replicas uses one h200 Job. Running all three costs ~$20 per hour.
Weight sync
| per sync, trainer's clock, 126 syncs | before | now (p50) |
|---|---|---|
| whole sync | 30.8 s | 8.5 s (min 6.6, max 9.2) |
| of which: pause both replicas | 0.3 s | 0.3 s |
| adapter all-gather and save to the bucket | 0.6 s | 1.1 s |
| both replicas accept the adapter | ~29 s | ~7 s |
All 252 adapter loads succeeded 🎉: 126 syncs times 2 replicas. Six succeeded on the second attempt and 246 on the third.
Routing
At the end of the run, after 64,728 rollouts, the proxy's counters read:
routed [31928, 32800] affinity 54712 spilled 820 unmatched 9196
With 8 rollouts per prompt, at least one of the eight requests must be cold. The theoretical minimum is therefore 12.5 %. The router gets 14.2 % unmatched requests, 84.5 % affinity hits, and 1.3 % spills. There is not much left to gain here unless we start looking at each replica's load using deeper inference-side metrics based on real measured load.
Where the time goes
The first configuration has a pretty obvious problem: the trainer is the bottleneck, not generation. Over the 500 steps, we have:
| per optimizer step, p50 | |
|---|---|
| step | 22.9 s |
| forward + backward | 21.9 s |
| waiting for rollouts | 0.02 s |
| rollout queue occupancy | 476 of 512 |
| trainer MFU | 3.9 % |
The rollout queue stays full, and the worker is mostly blocked by backpressure. The second replica is basically useless in this configuration. We'll see later in the post how we went through runs that moved the bottleneck between training and generation to make the run 3.9× faster.
Reward
Figure 1. trackio run r1-dp2. Panels: reward with its 20-step rolling mean and 50-step block means, and ratio on a 0.99 to 1.01 axis. Reward climbs from 0.15 to 0.44 over 500 steps; ratio stays between 0.9993 and 1.0004 throughout.
500 steps took 3 h 27 min. Mean reward goes from 0.145 over the first 20 steps to 0.438 over the last 20. More importantly for this test, ratio stays at 1.000 for every step! The policy served by vLLM always matches the one used by the trainer to score the rollout. This held across all 126 syncs. Mean staleness was 1.5 policy versions, against a maximum of 4. The trackio dashboard has the full curves.
We have our undeniable proof that LoRA AsyncGRPO works! Let's now see how we can improve our training runs by looking at the recent detailed AsyncGRPO metrics.
Chasing the bottleneck ping-pong
Async RL is a pipeline between training and generation. Making one side faster does nothing if the other side cannot keep up. Fortunately, in AsyncGRPOTrainer we've added enough timings and metrics to see this directly.
All the useful metrics are documented in the Logged metrics section. perf/rollout_wait_s tells us how long the trainer waits for samples. rollout/backpressure_s tells us how long generation waits for space in the rollout queue. They are diametrically opposed to each other and should not both be high. Together with the queue size, they tell us which side is slow.
We ran five experiments. Each one starts from a problem visible in the previous run's dashboard. Unless mentioned otherwise, the model, recipe and three-Job layout stay the same. The names below are the trackio run names.
Reading the dashboard
We keep these four groups of metrics visible:
perf/step_sandperf/fwd_bwd_s: how long an optimizer step takes, and how much of it is forward+backward. If a step takes nearly as long as the forward+backward, then the trainer is clearly compute-bound.perf/rollout_wait_s: how long the trainer sat waiting for samples before it could start a step. Near zero means generation is ahead of training. Samples are available immediately to be trained on.sample/rollout_queue_sizeagainstqueue_maxsize: the buffer between the two sides. Full means generation is being throttled; empty means the trainer is starving.rollout/backpressure_sandrollout/score_block_s: how long the rollout worker sat blocked because that buffer was full. The worker is a two-stage pipeline: generation hands finished groups to a scoring stage, and scoring pushes scored samples into the rollout buffer. When the buffer is full, scoring cannot enqueue and blocks, which isrollout/backpressure_s. Scoring then stops draining its own input queue, so generation cannot hand over the next group either, which isrollout/score_block_s. Both are the same stall, seen first at the scoring stage and then propagated back to generation.
The diagnosis is simple: a full queue with zero rollout wait and high backpressure means the trainer is too slow. An empty queue with rising rollout wait and no backpressure means generation is too slow. Comparing perf/mfu_wall_clock with perf/mfu_fwd_bwd also shows how much time the trainer GPUs spend waiting instead of training.
Run 1, r1-dp2: a trainer that cannot keep up
Figure 2. trackio run r1-dp2. Panels: perf/step_s, perf/fwd_bwd_s, sample/rollout_queue_size, rollout/backpressure_s. Step time and forward+backward overlap almost completely; the queue sits pinned near 476 of 512 and backpressure never drops below 11 s per rollout group: trainer-bound.
perf/step_s is 22.9 s and perf/fwd_bwd_s is 21.9 s. Forward and backward take 96 % of the step time. The queue stays full and the trainer waits only 0.02 s for rollouts, and the rollout worker spends 15 seconds per group blocked by backpressure. The two vLLM replicas generate faster than the trainer consumes. The reported 4.6k tokens/s is not their actual limit; they simply have nowhere to put more output.
The batch metrics explain the terrible 3.9 % MFU. batch/microbatches_per_step is 64 and batch/samples_per_row is 1.0. Each rank processes one sequence of around 1.2k tokens, 64 times per step. This comes from the reference recipe's per_device_train_batch_size=1. For a 1.5B model on an H200, this is completely latency-bound.
Run 2, r1-dp2-tb16k: pack the microbatch
The fix is not to change the batch size. We keep 128 completions per optimizer step and only change how they are laid out on the GPU: instead of one sequence per microbatch, we pack many sequences densely into each row. The trainer supports this through token-budget batching. With token_budget > 0, it packs several samples into one padding-free row per rank. An optimizer step processes gradient_accumulation_steps rows. We set token_budget=16384 and gradient_accumulation_steps=6.
Figure 3. trackio runs r1-dp2 and r1-dp2-tb16k overlaid over their first 154 steps. Panels: batch/samples_per_row, batch/microbatches_per_step, perf/step_s, perf/fwd_bwd_s, perf/mfu_fwd_bwd, rollout/generated_tok_s. Packing takes samples per row from 1 to 13, microbatches from 64 to 6, step time from 23 s to 5.9 s, and generation from 4.2k to 27.5k tok/s with no change on the vLLM side.
batch/samples_per_row goes from 1.0 to ~12.7 and the number of microbatches drops from 64 to 6. The rows are 95 % full! Forward and backward fall from 21.9 s to 5.6 s, while MFU rises from 3.9 % to 19 %. We now train on around 150 samples per step because the rows pack better than the mean-length estimate predicted.
Generation also jumps from 4.6k to 25k tokens/s, even though we changed nothing on the vLLM side. The queue is no longer constantly full, so the replicas can finally run. This is why we do not like optimizing pipeline stages in isolation. We need to be careful to evaluate the whole system, as a slow stage can hide the real performance of everything before it.
Run 3, r1-dp2-tb16k-nockpt: stop recomputing the forward
perf/fwd_s is 1.34 s while perf/fwd_bwd_s is 5.6 s. A normal backward costs roughly twice the forward, and with frozen base weights it should be closer to once. A ratio of 3.2 is suspicious.
The reason is that AsyncGRPOConfig defaults to gradient_checkpointing=True. Every microbatch recomputes its forward during the backward. This also explains why a 16k-token row only uses 25 GB on a 141 GB H200! It's a memory optimization for this trainer, but we don't need it in this specific case: the model is small enough to fit into VRAM with the activations kept for the backward.
Figure 4. trackio runs r1-dp2-tb16k and r1-dp2-tb16k-nockpt overlaid over their first 134 steps. Panels: perf/fwd_s, perf/fwd_bwd_s, perf/weight_sync_s, sample/rollout_queue_size, perf/rollout_wait_s, perf/mfu_fwd_bwd. Forward+backward drops by one forward; the queue falls from ~420 to ~60, and rollout wait rises from 0.02 s to 0.5 s: the bottleneck shifts to generation.
With gradient_checkpointing=False, forward and backward drop to 4.6 s, almost exactly one forward less, and MFU reaches 23 %. The queue now falls to 71 and rollout wait rises from 0.04 s to 0.6 s. The trainer consumes samples faster than two replicas generate them. We successfully moved the bottleneck to generation.
This exposes two more costs. A 7.6-second weight sync every four steps now takes 25 % of wall-clock time. It was only 8 % when each step took 23 seconds. Also, backward is still 2.5 times slower than forward. With frozen base weights, there are around 2 seconds per step that do not look like normal model math.
Run 4, r1-dp3-tb16k-nockpt: three replicas, and a surprise
Since generation was now too slow, we added a third replica. We also reduced the adapter retry interval from 2 s to 0.5 s in the proxy and disabled fsdp_reshard_after_forward to check whether FSDP2 re-gathers caused the extra 2 seconds in backward.
Figure 5. trackio runs r1-dp2-tb16k-nockpt and r1-dp3-tb16k-nockpt overlaid over their first 134 steps. Panels: perf/weight_sync_s, rollout/generated_tok_s, rollout/inflight, perf/fwd_bwd_s. Sync falls from 7.6 s to 5.8 s; generation and forward+backward do not move; rollout/inflight reads 128 in both runs, which is the cap the third replica ran into.
Weight sync falls from 7.6 s to 5.8 s, so the shorter retry helps. Forward and backward stay at 4.6 s, which rules out resharding. Generation moves from 25k to only 26k tokens/s. The third replica does basically nothing.
The reason was sitting in rollout/inflight: 128 in every run. The proxy shows those requests split as 44 + 43 + 41 across the three replicas. max_inflight_tasks limits concurrency for the whole rollout worker, not per replica. A 1.5B model on an H200 processes 43 and 130 concurrent sequences at almost the same cost per token. Splitting 128 requests over three GPUs gives nearly the same throughput as splitting them over two.
So vLLM was not the limit. Our own client-side constant was. We had set it conservatively because we did not know how hundreds of long HTTPS requests would behave through the public Jobs proxy. At this point, 130,000 rollout completions had crossed it without a single transport error.
Run 5, r1-dp3-inflight384: lift the cap on in-flight requests
max_inflight_tasks=384 and queue_maxsize=768, nothing else.
Figure 6. All five trackio runs (r1-dp2, r1-dp2-tb16k, r1-dp2-tb16k-nockpt, r1-dp3-tb16k-nockpt, r1-dp3-inflight384) overlaid, x-axis in steps. Panels: perf/step_s, reward, sample/rollout_queue_size, sample/staleness_mean. Step time falls from 22.9 s to 4.8 s across the series while the reward curves stay on top of each other; the last run's queue refills to ~690 of 768 and its staleness settles at 2. Runs 2 to 4 were stopped early once the dashboard had answered the question.
With 384 requests in flight, each replica gets 128. The queue quickly fills to around 690 out of 768 and stays there. Backpressure returns to 5 seconds and rollout wait falls to 0.03 seconds. Training is the bottleneck again! Forward and backward take 4.6 seconds, weight sync adds an amortized 1.5 seconds, and median step time is 4.8 seconds.
Mean staleness rises from 1.5 to 2.0 versions because samples wait longer in the larger queue. This is still below max_staleness=4, and ratio remains very close to 1.000.
The scoreboard
| 500 steps | run 1 r1-dp2 | run 5 r1-dp3-inflight384 |
|---|---|---|
| wall clock | 3 h 27 min | 53 min |
perf/step_s, p50 | 22.9 s | 4.8 s |
perf/fwd_bwd_s, p50 | 21.9 s | 4.6 s |
perf/mfu_fwd_bwd | 3.9 % | 23.5 % |
batch/samples_per_step | 128 | 168 |
| samples trained | 64 000 | 84 078 |
perf/weight_sync_s, p50 | 8.5 s | 6.2 s |
sample/staleness_mean | 1.5 | 2.0 |
| reward, first 20 → last 20 steps | 0.145 → 0.438 | 0.145 → 0.416 |
Figure 7. trackio runs r1-dp2 and r1-dp3-inflight384, reward against wall-clock minutes since the first optimizer step. Same recipe, same 500 steps, same final reward; run 5 gets there in 52 minutes instead of 3 h 26 min.
The final run is 3.9× faster and trains on 31 % more samples, with basically the same reward curve. Packing, disabling checkpointing and raising the in-flight limit made the difference. In each case, the dashboard made this clear within the first ten minutes.
Try it
git clone https://github.com/AmineDiro/hfjobs-lora-buckets && cd hfjobs-lora-buckets
hf auth login
MAX_STEPS=20 RUN_TAG=smoke ./run_all.sh --wait
MAX_STEPS=500 ./run_all.sh --wait
TOKEN_BUDGET=16384 GRAD_ACCUM=6 GRADIENT_CHECKPOINTING=0 PROXY_LORA_RETRY_S=0.5 \
MAX_INFLIGHT=384 QUEUE_MAXSIZE=768 MAX_STEPS=500 ./run_all.sh --wait
References
- John Schulman et al., LoRA Without Regret, Thinking Machines Lab, September 2025. The case that rank-1 LoRA matches full fine-tuning for policy-gradient RL, and why.
- TRL,
AsyncGRPOTrainerand its logged metrics. - TRL PR #7017: PEFT/LoRA support for
AsyncGRPOTrainerwith adapter-only vLLM sync, released in TRL v1.14. - Hugging Face Jobs and Storage Buckets;
hf-mount. hf-mount-repro: the two-script reproduction of the 30-second negative-cache stall.- The trackio dashboard for every run in this post.
- Penghui Qi, Zichen Liu, Xiangxin Zhou, Tianyu Pang, Chao Du, Wee Sun Lee, Min Lin, Defeating the Training-Inference Mismatch via FP16, arXiv:2510.26788, 2025. Source of the Sanity dataset
sail/Sanity-Test-R1D-1.5Band of the LoRA recipe,sail-sg/Precision-RL,oat/scripts/lora/bf16_grpo_tis_lora.sh.
@article{qi2025precisionrl,
title={Defeating the Training-Inference Mismatch via FP16},
author={Qi, Penghui and Liu, Zichen and Zhou, Xiangxin and Pang, Tianyu and Du, Chao and Lee, Wee Sun and Lin, Min},
journal={arXiv preprint arXiv:2510.26788},
year={2025}
}