Async GRPO with LoRA across HF Jobs: a bucket, a proxy, and no NCCL
Async GRPO with LoRA across HF Jobs: a bucket, a proxy, and no NCCL
跨 Hugging Face Jobs 的异步 GRPO 与 LoRA：一个存储桶、一个代理，无需 NCCL
TL;DR AsyncGRPOTrainer now supports training a LoRA adapter and syncing 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.
简而言之：AsyncGRPOTrainer 现在支持训练 LoRA 适配器，并仅将该适配器同步到 vLLM（TRL v1.14）。由于 rank-1 适配器仅有几兆字节，它可以通过挂载在每个 Job 中的存储桶（Storage Bucket）进行传输，而无需使用 NCCL。训练器和 vLLM 副本作为独立的 Hugging Face Jobs 在不同的机器上运行。副本前端的一个小型代理负责添加认证头，将每个 rollout 路由到已持有其 KV 前缀的副本，并将适配器加载指令广播给所有副本。AsyncGRPO 指标揭示了瓶颈所在。在 500 步的测试中，五次运行将相同的配方执行时间从 3 小时 27 分钟缩短至 53 分钟。
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 支持最近通过 PR #7017 进入了 TRL 的 AsyncGRPOTrainer，并随 TRL v1.14 发布。异步训练器现在可以训练适配器而非完整模型，并且仅将 LoRA 适配器同步到 vLLM。本文介绍了一个基于此构建的实际项目，其中训练和推理不再共享同一台机器。
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.
LoRA 训练特别适合强化学习（RL），正如 Thinking Machines 的博客《LoRA Without Regret》所展示的那样。他们证明了即使在 rank 为 1 的情况下，LoRA 在策略梯度强化学习中也能媲美全参数微调。这源于优势函数（advantage function）在每个 episode 中仅提供 ~O(1) 比特的信息，因此从总信息量的角度来看，每一步需要学习的内容并不多。一个 rank-1 的适配器有足够的容量来吸收这些信息。
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.
LoRA 训练在系统层面也有显著影响。对于一个 1.5B 的模型，rank-1 适配器仅有几兆字节，而完整模型约为 3 GB。我们无需在每次更新后将完整策略发送给推理工作节点，只需发送适配器即可。vLLM 还可以同时加载多个适配器。旧的 rollout 使用它们开始时的策略完成，而新的 rollout 则使用最新的策略。
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).
TRL 的 AsyncGRPOTrainer 已经实现了训练与生成的解耦。训练器和 vLLM 可以在不同的机器上以各自的速度运行。在单节点或集群环境中，如果两个进程共享文件系统或可以组成 NCCL 组，这很容易实现。我们想要的是在 Hugging Face Jobs 上运行相同的设置。本质上，一个 HF Job 就是运行在单个虚拟机上的一个容器。这意味着一个 Job 目前无法跨多个节点（至少目前如此）来同时容纳训练器和一组 vLLM 服务器（每个节点最多限制为 8xH200）。
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.
AsyncGRPOTrainer 正是为这种规模而构建的，因此问题变成了：如果我们放弃“训练器和推理服务器必须共享节点”的要求，我们能走多远？好吧，如果是全权重同步，答案是“走不远”。每次更新都必须在机器之间传输数 GB 的数据，这正是密集集群中 NCCL 的作用，但 Jobs 无法跨节点通信。没有共享的本地磁盘，显然也没有共享的 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.
有了 LoRA，同步仅需几兆字节。至于文件系统部分，HF Jobs 提供了由存储桶支持的卷！这些存储桶可以作为 FUSE 文件系统挂载到每个 Job 中，足以作为节点间的共享文件系统。完全不需要在 Jobs 之间建立网络路径。
The setup ended up being quite small: a trainer Job running AsyncGRPOTrainer with 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.
最终的设置非常精简：一个运行 AsyncGRPOTrainer（配合 LoRA 和 FSDP，稍后详述）的训练器 Job，两个 vLLM Job（每个服务于基础模型加上训练器最新发布的适配器），一个挂载在所有三个 Job 中相同路径下的存储桶（这是适配器从训练器传输到服务器的方式），以及一个代理服务器。我们将深入探讨为什么需要代理，但简单来说，我们需要一个代理将每个 rollout 路由到最可能持有其 KV 缓存的副本，并将每个适配器更新广播给所有 vLLM 副本。
The architecture: leveraging Hugging Face Jobs and Storage Buckets 🪣
架构：利用 Hugging Face Jobs 和存储桶 🪣
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.
AsyncGRPOTrainer 中新的仅适配器同步路径工作原理如下：训练器不向 vLLM 发送张量。每隔几个优化器步骤，它将适配器保存到 <output_dir>/.vllm_lora/trl-policy-v{N}，通过原子重命名发布该目录，然后将其路径发送到 vLLM 的 /v1/load_lora_adapter 端点。vLLM 从磁盘加载文件，因此 rollout 工作节点随后可以请求 model="trl-policy-v{N}"。这就是 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:
该端点接收的是路径而非张量，因此训练器和服务器需要共享文件系统。在 Slurm 集群上，这是网络文件系统。在 Jobs 上，正如我们之前提到的，通过将存储桶作为卷挂载到每个 Job 的相同路径下，我们实现了同样的效果。在底层，它使用 hf-mount，将存储桶作为 POSIX 文件系统暴露在容器内：
# every Job gets the same bucket at the same absolute path
# 每个 Job 在相同的绝对路径下获得同一个存储桶
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.
TRL 或 vLLM 无需为此进行任何更改。训练器写入 /lora/<run>/.vllm_lora/，服务器从同一路径读取。POST 请求中发送的路径在每个容器内都是有效的。
The three Jobs and the bucket. TRL talks to the proxy over localhost, the proxy talks to the replicas over HTTPS, and the adapter directory travels through the bucket mount. 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.
三个 Jobs 和存储桶。TRL 通过 localhost 与代理通信，代理通过 HTTPS 与副本通信，适配器目录通过存储桶挂载进行传输。请注意，我们还将检查点和最终适配器存储在存储桶中。HF Jobs 是临时的，但被抢占的训练器可以恢复训练，因为最终适配器始终持久化在存储桶中，不会在 Job 停止时丢失。
The three Jobs
三个 Jobs
The vLLM replicas
vLLM 副本
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.
每个副本使用一个 GPU 和标准的 vllm/vllm-openai 镜像。我们只需要启用运行时 LoRA 加载并预留足够的适配器槽位。适配器槽位的数量取决于 max_staleness。在 AsyncGRPOTrainer 中，每次权重同步都会使策略版本增加 1，而 max_staleness 表示 rollout 样本可以落后于当前策略的版本数。