返回 文章 build CMS 文章

Async GRPO 与 LoRA 跨 HF Jobs:用存储桶和代理替代 NCCL

用几兆字节的 LoRA 适配器替代千兆字节的全权重同步,在 HF Jobs 上实现无 NCCL 的异步 GRPO 训练。

Async GRPOLoRAHF JobsvLLM
成长分 / 100 74 综合收获、行动、留存与影响

Async GRPO 与 LoRA 跨 HF Jobs:用存储桶和代理替代 NCCL
为什么值得读了解如何在不共享节点、没有 NCCL 的情况下,利用 HF Jobs 和 Storage Bucket 运行异步 GRPO 训练。

学习仅适配器同步、带版本号的适配器命名以及前缀缓存感知路由如何保证训练与推理策略一致。

关键洞察
  1. LoRA 秩为 1 的适配器仅几兆字节,可通过 Storage Bucket 挂载传输,避免跨节点 NCCL 通信。
  2. vLLM 运行时 LoRA 加载端点接收路径而非张量,因此训练器和服务器只需共享文件系统。
  3. 带版本号的适配器名称防止前缀缓存跨策略版本错误匹配,确保 ratio 稳定在 1.000。
转成行动

深入阅读

正文与原文对照

原文保真覆盖:全文原文字符:33121

文本生成 • 2B • 已更新 • 403k • 1.58k

TL;DR

AsyncGRPOTrainer

现在可以训练一个 LoRA 适配器,并仅将该适配器同步到 vLLM(TRL v1.14)。- 一个秩为 1 的适配器只有几兆字节,因此它可以通过挂载在每个 Job 中的 Storage Bucket 传输,而不是通过 NCCL。训练器和 vLLM 副本作为独立的 Hugging Face Jobs 运行在不同的机器上。

  • 副本前面的一个小型代理添加认证头,将每个 rollout 路由到已经持有其 KV 前缀的副本,并将适配器加载广播到每个副本。
  • AsyncGRPO 指标显示了瓶颈所在。五次运行使用相同的配方,500 步从 3 小时 27 分钟缩短到 53 分钟。

LoRA 支持最近在 TRL 的 AsyncGRPOTrainer 中落地,LoRA 训练特别适合 RL,正如 Thinking Machines 的博客 LoRA Without Regret 所示。他们表明,LoRA 可以在策略梯度 RL 中匹配完全微调,即使秩为 1。这源于优势函数每个 episode 只提供 ~O(1)

比特的信息,因此从总信息比特的角度来看,每一步没有太多可学习的。一个秩为 1 的适配器有足够的容量来吸收它。

LoRA 训练还有一个系统层面的影响。一个 1.5B 模型的秩为 1 的适配器只有几兆字节,而完整模型大约 3 GB。我们不必在每次更新后将完整的策略发送给推理工作器,只需发送适配器即可。vLLM 还可以同时加载多个适配器。旧的 rollout 使用它们开始时的策略完成,而新的 rollout 使用最新的策略。

TRL 的 AsyncGRPOTrainer

已经将训练和生成分离。训练器和 vLLM 可以在不同的机器上以各自的速度运行。这在单节点或集群设置中很容易,因为两个进程共享文件系统或可以组成 NCCL 组。

我们想要的是使用 Hugging Face Jobs 运行相同的设置。本质上,一个 HF Job 是在一个 VM 上运行的一个容器。这意味着一个 Job 无法生成多个节点(至少目前如此)来容纳一个训练器和一组 vLLM 服务器(每个节点最多限制为 8xH200)。AsyncGRPOTrainer

正是为这种规模构建的,所以问题变成了:如果我们放弃训练器和推理服务器共享节点的要求,我们能走多远?

嗯,如果使用全权重同步,答案将是“走不远”。每次更新都必须在机器之间移动千兆字节的数据,这在密集集群中正是 NCCL 的用途,但 Jobs 无法节点通信。没有共享的本地磁盘,显然也没有共享的 localhost

。使用 LoRA,一次同步只有几兆字节。对于文件系统部分,HF Jobs 提供了由 Storage Buckets 支持的卷!这些桶可以在每个 Job 中作为 FUSE 文件系统挂载,足以作为节点之间的共享文件系统。完全不需要 Job 之间的网络路径。

设置最终变得相当小:

AsyncGRPOTrainer

使用 LoRA(以及 FSDP,稍后详述),AsyncGRPOTrainer 中新的仅适配器同步路径

工作方式如下。训练器不向 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 文件系统暴露在容器内:

# 每个 Job 在相同的绝对路径下获得相同的 bucket
hf jobs run ... -v hf://buckets/aminediroHF/asyncgrpo-lora-buckets:/lora ...

TRL 或 vLLM 无需为此做任何改动。训练器写入 /lora/<run>/.vllm_lora/

,服务器从同一路径读取。POST 请求中发送的路径在每个容器内都已有效。

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

每个副本使用一个 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 在途的策略。

# --expose 8000 可通过 https://<job_id>--8000.hf.jobs 访问
# -v ...:/lora:ro 只读:服务器仅读取适配器
# VLLM_ALLOW_RUNTIME_LORA_UPDATING=1 启用 /v1/load_lora_adapter
# VLLM_SERVER_DEV_MODE=1 启用 /pause、/resume、/server_info(TRL 需要全部三个)
# --max-loras 6 max_staleness=4 -> 4+2 个适配器槽位
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 可能从某一个策略版本获取其前缀,而从下一个版本获取其解码。训练器将无从察觉,而这会表现为 ratio

偏离 1。带版本的名称使这种情况不可能发生:一个名称始终意味着同一组权重,而缓存的前缀永远不可能匹配更新的版本。

我们选择了 sail/Sanity-Test-R1D-1.5B,即来自

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 个补全,最多生成 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", # on the bucket: adapters, checkpoints and the final adapter all land here
vllm_server_base_url="http://localhost:8000", # the proxy, not a vLLM Job; TRL never sees the Jobs URLs
max_staleness=4,
weight_sync_steps=4, # publish an adapter every 4 optimizer steps
save_strategy="steps", save_steps=50, # checkpoints go to the same bucket -> resume after preemption
...
)
trainer = AsyncGRPOTrainer(
model="Qwen/Qwen2.5-Math-1.5B",
args=config,
peft_config=LoraConfig(r=1, lora_alpha=2, target_modules="all-linear"), # plain LoRA vLLM can serve as-is
...
)

初始化期间,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 秩,因此其他秩会继续以新策略名称提供基础模型。在 Jobs 上,这个问题根本不会出现,因为每个副本都是独立的机器。因此,数据并行必须上移一层,放在能将适配器加载扇出到每个副本的东西中。

因此,我们在训练器 Job 上的 127.0.0.1:8000

运行一个小型代理,并将 TRL 指向它,就像它是一个单一的 vLLM 服务器一样。除了添加头之外,代理在功能上做两件事:

快速提醒一下为什么这很重要。生成补全有两个阶段,工作负载特征非常不同:

那些键和值就是 KV 缓存。由于注意力是因果的,一个 token 的 KV 只依赖于它之前的 token,而不依赖于它之后的 token。因此,共享前缀的两个请求共享该前缀的 KV,而已经将其缓存在缓存中的副本可以 完全跳过预填充的那部分。现在的整个游戏就是找到那个副本,这样请求就可以受益于落在已经见过其前缀的副本上。

vLLM 将其 前缀 KV 缓存 存储在 16 个 token 的块中。由于 GRPO,rollout worker 发送 G

个具有相同提示的请求(在我们的例子中 G=8

)。如果它们都到达同一个副本,第一个请求计算预填充,接下来的七个重用它。使用轮询路由,一半会到达没有缓存前缀的副本,那四个请求将重做预填充工作并浪费宝贵的 GPU 计算。

我们的路由器 的工作是跟踪哪个副本见过哪个 块哈希。一个重要细节是哈希是链式的,因此块 3 的哈希代表块 1、2 和 3,而不仅仅是块 3。这反映了因果注意力:块 3 的 KV 只有在块 1 和 2 也相同的情况下才有效。我们还用适配器名称作为链的种子,因为 KV 缓存也依赖于生成它的适配器:为策略 v3 缓存的前缀对策略 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 由两个副本共同拥有,并且已经有多个后继,h2h8 仅由 A 拥有,每个都有一个单一后继,而 h2'h6' 仅由 B 拥有。

在实践中,一次运行中的每个提示都以相同的 token 开头。在这里,它是聊天模板和系统提示,它们构成了所有 1,460 个问题的前 23 个 token。在代理设置中,它将是工具描述,而在多轮环境中,它将是共享的对话历史。这些块在几秒钟内就存在于每个副本的缓存中,因此对它们进行匹配并不能告诉我们特定提示位于何处。

如果一个块被每个副本都服务过,或者它有多个后继,那么它就是公共的。公共块在路由过程中被忽略,因为它们不能识别特定的提示。更多内容见下文!

5. 选择一个副本。 路由器计算每个副本上有多少个前导块匹配,并移除公共前缀。剩下的就是特定于此提示的块数。然后:

以下是应用于四个请求的规则。从副本 A 和 B 都有 3 个请求在途的状态开始,并且只有模板块 h1 在两个副本上已知。

h2h8 由 A 拥有。A 现在有 4 个请求在途。h2'h6' 由 B 拥有。h2h8 也由 B 拥有。现在每个副本都服务过这些块,因此问题 0 也变得公共,其后续的 rollout 仅由负载决定。6. 重用预填充。 请求 2 是执行所有这些操作的原因。它重用了请求 1 计算的预填充:块 1 到 8 已经在 A 的 KV 缓存中,因此 A 直接跳到解码完成。如果它去了 B,B 将再次预填充所有 135 个 token,而 A 的缓存则闲置不用。请求 3 展示了为什么需要 common 规则。如果没有它,共享的聊天模板会让每个新提示看起来都像缓存命中。请求 4 保持负载有界。节省一次预填充不值得让一个副本远远落后。

def choose(self, upstreams, model, prompt):
hashes = self.block_hashes(model, prompt) # 对 16 个 token 的块进行链式 blake2b,以 `model` 作为种子
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
...将 `pick` 记录为每个块以及每个块的后继块的所有者...
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) # this replica's mount has not seen the directory yet
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

,在 all-linear

上使用 LoRA r=1

,每步 128 个完成,每个提示 8 次 rollout。它运行 500 步,每 50 步保存一个检查点。训练器使用一个 h200x2

作业,两个 vLLM 副本各使用一个 h200

作业。运行全部三个每小时花费约 20 美元。

每次同步,训练器时钟,126 次同步 之前 现在(p50)
整个同步 30.8 秒 8.5 秒(最小 6.6,最大 9.2)
其中:暂停两个副本 0.3 秒 0.3 秒
适配器全收集并保存到桶 0.6 秒 1.1 秒
两个副本接受适配器 ~29 秒 ~7 秒

所有 252 次适配器加载均成功 🎉:126 次同步乘以 2 个副本。六次在第二次尝试时成功,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 秒
前向 + 反向 21.9 秒
等待 rollout 0.02 秒
rollout 队列占用 476 / 512
训练器 MFU 3.9%

rollout 队列始终保持满状态,工作进程大部分时间被背压阻塞。第二个副本在此配置下基本无用。我们将在后文看到,我们如何通过一系列运行在训练和生成之间转移瓶颈,使运行速度提升了 3.9 倍。

图 1. trackio 运行 r1-dp2。面板:奖励及其 20 步滚动均值和 50 步块均值,以及 0.99 至 1.01 轴上的比率。奖励在 500 步内从 0.15 攀升至 0.44;比率全程保持在 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 队列空间的时间。它们完全相反,不应同时很高。结合队列大小,它们告诉我们哪一侧慢。

我们进行了五次实验。每次实验都从上次运行仪表板中可见的问题开始。除非另有说明,模型、配方和三作业布局保持不变。以下名称是 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/backpressure_s

。评分随后停止排空自己的输入队列,因此生成也无法移交下一组,这就是 rollout/score_block_s

两者是同一个卡顿,首先在评分阶段看到,然后传播回生成阶段。诊断很简单:队列满且 rollout 等待为零、背压高,意味着训练器太慢。队列空且 rollout 等待上升、没有背压,意味着生成太慢。比较 perf/mfu_wall_clock

perf/mfu_fwd_bwd

也显示了训练器 GPU 有多少时间花在等待而不是训练上。

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 工作器每组因背压阻塞 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 模型,这完全是延迟受限的。

r1-dp2-tb16k

:打包微批次

修复方法不是改变批次大小。我们保持每个优化器步 128 个完成,只改变它们在 GPU 上的布局方式:不是每个微批次一个序列,而是将许多序列密集地打包到每一行中。训练器通过 token 预算批处理 支持这一点。当 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 跃升至 25k tokens/s,尽管我们在 vLLM 侧没有做任何改变。队列不再持续满,因此副本终于可以运行。这就是为什么我们不喜欢孤立地优化流水线阶段。我们需要小心评估整个系统,因为慢速阶段可能掩盖其之前所有部分的真实性能。

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 秒升至 0.5 秒:瓶颈转移到了生成。

使用 gradient_checkpointing=False

时,前向和反向降至 4.6 秒,几乎正好少了一次前向的时间,MFU 达到 23%。队列现在降至 71,rollout 等待从 0.04 秒升至 0.6 秒。训练器消耗样本的速度快于两个副本生成样本的速度。我们成功地将瓶颈转移到了生成。

这暴露了另外两个成本。每四步进行一次 7.6 秒的权重同步现在占用了 25% 的挂钟时间。当每步耗时 23 秒时,它只占 8%。此外,反向传播仍然比前向传播慢 2.5 倍。在基础权重冻结的情况下,每步大约有 2 秒看起来不像正常的模型计算。

r1-dp3-tb16k-nockpt

:三个副本,以及一个意外

由于生成现在太慢,我们添加了第三个副本。我们还将代理中的适配器重试间隔从 2 秒减少到 0.5 秒,并禁用了 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 秒降至 5.8 秒;生成和前向+反向没有变化;两次运行中 rollout/inflight 读数均为 128,这是第三个副本遇到的上限。

权重同步从 7.6 秒降至 5.8 秒,因此更短的重试有所帮助。前向和反向保持在 4.6 秒,这排除了重新分片的影响。生成从 25k 仅提升到 26k tokens/s。第三个副本基本上没有起作用。

原因就在 rollout/inflight

中:每次运行都是 128。代理显示这些请求在三个副本上分为 44 + 43 + 41。max_inflight_tasks

限制的是整个 rollout worker 的并发数,而不是每个副本的并发数。一个 1.5B 模型在 H200 上处理 43 个和 130 个并发序列时,每个 token 的成本几乎相同。将 128 个请求分配到三个 GPU 上,其吞吐量与分配到两个 GPU 上几乎相同。

所以 vLLM 不是限制因素。我们自己客户端侧的常量才是。我们之前设置得比较保守,因为我们不知道数百个长 HTTPS 请求通过公共 Jobs 代理会表现如何。此时,已有 130,000 次 rollout 完成通过了它,没有出现任何传输错误。

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 秒降至 4.8 秒,而奖励曲线彼此重合;最后一次运行的队列重新填充至约 690/768,其陈旧度稳定在 2。运行 2 到 4 在仪表板回答了问题后提前停止。

当有 384 个请求在途时,每个副本获得 128 个。队列迅速填充至约 690/768 并保持在那里。背压恢复到 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
挂钟时间
perf/step_s , p50
22.9 秒
perf/fwd_bwd_s , p50
21.9 秒
perf/mfu_fwd_bwd
3.9 %
batch/samples_per_step
128
训练的样本数
perf/weight_sync_s , p50
8.5 秒
sample/staleness_mean
1.5
奖励,前 20 步 → 后 20 步

图 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 # ~15 min, three Jobs, cancels the servers when done
MAX_STEPS=500 ./run_all.sh --wait # run 1: the reference batch shape, ~3.5 h
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 # run 5: same recipe, ~55 min

AsyncGRPOTrainer

AsyncGRPOTrainer

采用仅适配器 vLLM 同步,发布于 TRL v1.14。hf-mount

hf-mount-repro

sail/Sanity-Test-R1D-1.5B

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}

}


在可视化仪表板中显示跟踪的 I/O 统计信息