Async GRPO with LoRA across HF Jobs: a bucket, a proxy, and no NCCL
<div class="SVELTE_HYDRATER contents" data-target="RepoCodeCopy" data-props="{}"><div></div></div> <div> <a target="_blank" class="btn mb-5 font-sans text-sm no-underline" href="https://github.com/huggingface/blog/blob/main/asyncgrpo-lora-hfjobs.md">Update on GitHub</a></div> <div class="relative overflow-clip [overflow-clip-margin:3rem] [&>:last-child]:mb-0">
<blockquote>
<p></p><strong>TL;DR</strong><p></p>
<ul>
<li><code>AsyncGRPOTrainer</code> can now train a LoRA adapter and sync only that adapter to vLLM (TRL v1.14).</li>
<li>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.</li>
<li>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.</li>
<li>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.</li>
</ul>
</blockquote>
<p>LoRA support recently landed in TRL's <a href="https://huggingface.co/docs/trl/en/async_grpo_trainer"><code>AsyncGRPOTrainer</code></a> with <a href="https://github.com/huggingface/trl/pull/7017" rel="nofollow">PR #7017</a>, 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.</p>
<p>LoRA training is particularly suited for RL, as shown in Thinking Machines's blog <a href="https://thinkingmachines.ai/blog/lora/" rel="nofollow">LoRA Without Regret</a>. 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 <code>~O(1)</code> 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.</p>
<p>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.</p>
<p>TRL's <code>AsyncGRPOTrainer</code> 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.</p>
<p>What we want is to run the same setup with <a href="https://huggingface.co/docs/huggingface_hub/guides/jobs">Hugging Face Jobs</a>. 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 <code>AsyncGRPOTrainer</code> 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?</p>
<p>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 <em>across</em> nodes. There is no shared local disk and obviously no shared <code>localhost</code>. With LoRA, a sync is only a few megabytes. For the filesystem part, HF Jobs provide volumes backed by <a href="https://huggingface.co/docs/hub/storage-buckets">Storage Buckets</a>! 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.</p>
<p>The setup ended up being quite small:</p>
<ul>
<li>a <strong>trainer Job</strong> running <code>AsyncGRPOTrainer</code> with LoRA (and FSDP, more on that later),</li>
<li><strong>two vLLM Jobs</strong>, each serving the base model plus whatever adapter the trainer last published,</li>
<li>a <strong>Storage Bucket</strong> mounted in all three at the same path, which is how the adapter gets from the trainer to the servers,</li>
<li>a <strong>proxy server</strong>. 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.</li>
</ul>
<h2 class="relative group flex items-baseline">
<a id="the-architecture-leveraging-hugging-face-jobs-and-storage-buckets-🪣" class="block pr-1.5 text-lg md:absolute md:p-1.5 md:opacity-0 md:group-hover:opacity-100 md:right-full" href="https://huggingface.co/blog/asyncgrpo-lora-hfjobs#the-architecture-leveraging-hugging-face-jobs-and-storage-buckets-%F0%9F%AA%A3" rel="nofollow">
<span class="header-link"><svg class="text-gray-500 hover:text-black dark:hover:text-gray-200 w-4" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span>
</a>
<span>
The architecture: leveraging Hugging Face Jobs and Storage Buckets 🪣
</span>
</h2>
<p>The new adapter-only sync path in <code>AsyncGRPOTrainer</code> works like this. The trainer does not send tensors to vLLM. Every few optimizer steps, it saves the adapter under <code><output_dir>/.vllm_lora/trl-policy-v{N}</code>, publishes the directory with an atomic rename, then sends its path to vLLM's <code>/v1/load_lora_adapter</code> endpoint. vLLM loads the files from disk, so the rollout worker can then request <code>model="trl-policy-v{N}"</code>.</p>
<p>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 <a href="https://huggingface.co/docs/hub/storage-buckets">Storage Bucket</a> as a volume at the same path in every Job, as we mentioned earlier. Under the hood, it uses <a href="https://github.com/huggingface/hf-mount" rel="nofollow"><code>hf-mount</code></a>, which exposes the bucket as a POSIX filesystem inside the container:</p>
<pre><code class="language-sh"><span class="hljs-comment"># every Job gets the same bucket at the same absolute path</span>
hf <span class="hljs-built_in">jobs</span> run ... -v hf://buckets/aminediroHF/asyncgrpo-lora-buckets:/lora ...
</code></pre>
<p>Nothing in TRL or vLLM had to change for this. The trainer writes to <code>/lora/<run>/.vllm_lora/</code> and the servers read from the same path. The path sent in the POST request is already valid inside every container.</p>
<figure class="image text-center">
<img src="https://imgproxy.citrusreader.com/-ypDAZUDB87CT_EIoOjxyHzrGl7dSt4Ns7qfg_70iug/raw:1/aHR0cHM6Ly9odWdnaW5nZmFjZS5jby9kYXRhc2V0cy9odWdnaW5nZmFjZS9kb2N1bWVudGF0aW9uLWltYWdlcy9yZXNvbHZlL21haW4vYmxvZy9hc3luY2dycG8tbG9yYS1oZmpvYnMvYXJjaGl0ZWN0dXJlLnBuZw" alt="Async GRPO with LoRA across Hugging Face Jobs. The trainer Job runs AsyncGRPOTrainer and the proxy, two vLLM Jobs serve the base model plus the latest adapter, and a Storage Bucket is mounted at /lora in all three." referrerpolicy="no-referrer">
<figcaption>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.</figcaption>
</figure>
<p>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.</p>
<h2 class="relative group flex items-baseline">
<a id="the-three-jobs" class="block pr-1.5 text-lg md:absolute md:p-1.5 md:opacity-0 md:group-hover:opacity-100 md:right-full" href="https://huggingface.co/blog/asyncgrpo-lora-hfjobs#the-three-jobs" rel="nofollow">
<span class="header-link"><svg class="text-gray-500 hover:text-black dark:hover:text-gray-200 w-4" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span>
</a>
<span>
The three Jobs
</span>
</h2>
<h3 class="relative group flex items-baseline">
<a id="the-vllm-replicas" class="block pr-1.5 text-lg md:absolute md:p-1.5 md:opacity-0 md:group-hover:opacity-100 md:right-full" href="https://huggingface.co/blog/asyncgrpo-lora-hfjobs#the-vllm-replicas" rel="nofollow">
<span class="header-link"><svg class="text-gray-500 hover:text-black dark:hover:text-gray-200 w-4" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span>
</a>
<span>
The vLLM replicas
</span>
</h3>
<p>Each replica uses one GPU and the stock <code>vllm/vllm-openai</code> image. We only need to enable <a href="https://docs.vllm.ai/en/latest/features/lora/?h=lora#serving-lora-adapters" rel="nofollow">runtime LoRA</a> loading and reserve enough adapter slots.</p>
<p>The number of adapter slots follows from <code>max_staleness</code>. In <code>AsyncGRPOTrainer</code>, every weight sync bumps the policy version by one, and <code>max_staleness</code> is how many versions a rollout sample may lag behind the current policy before the trainer discards it. With <code>max_staleness=4</code>, a sample generated under <code>trl-policy-v3</code> is still used for training while the trainer is at <code>v7</code>. A rollout that started under <code>v3</code> must also be able to finish under <code>v3</code>. So at any moment, vLLM has to serve the current policy plus the four before it. That is why the trainer keeps <code>max_staleness + 1</code> 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 <code>--max-loras 6</code>. With only five, vLLM would silently evict a policy that still has rollouts in flight at every sync.</p>
<pre><code class="language-sh"><span class="hljs-comment"># --expose 8000 reachable at https://<job_id>--8000.hf.jobs</span>
<span class="hljs-comment"># -v ...:/lora:ro read-only: the server only reads adapters</span>
<span class="hljs-comment"># VLLM_ALLOW_RUNTIME_LORA_UPDATING=1 enables /v1/load_lora_adapter</span>
<span class="hljs-comment"># VLLM_SERVER_DEV_MODE=1 enables /pause, /resume, /server_info (TRL needs all three)</span>
<span class="hljs-comment"># --max-loras 6 max_staleness=4 -> 4+2 adapter slots</span>
<span class="hljs-keyword">for</span> replica <span class="hljs-keyword">in</span> 1 2; <span class="hljs-keyword">do</span>
hf <span class="hljs-built_in">jobs</span> run --detach --flavor h200 --<span class="hljs-built_in">timeout</span> 8h --secrets HF_TOKEN \
--expose 8000 \
-v <span class="hljs-string">"hf://buckets/<span class="hljs-variable">${BUCKET}</span>:/lora:ro"</span> \
-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
<span class="hljs-keyword">done</span>
</code></pre>
<p>We pin vLLM to <code>v0.27.1</code>. 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.</p>
<p>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 <code>ratio</code> 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.</p>
<h3 class="relative group flex items-baseline">
<a id="the-dataset-choice-the-sanity-set" class="block pr-1.5 text-lg md:absolute md:p-1.5 md:opacity-0 md:group-hover:opacity-100 md:right-full" href="https://huggingface.co/blog/asyncgrpo-lora-hfjobs#the-dataset-choice-the-sanity-set" rel="nofollow">
<span class="header-link"><svg class="text-gray-500 hover:text-black dark:hover:text-gray-200 w-4" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span>
</a>
<span>
The dataset choice: the Sanity set
</span>
</h3>
<p>We chose <a href="https://huggingface.co/datasets/sail/Sanity-Test-R1D-1.5B"><code>sail/Sanity-Test-R1D-1.5B</code></a>, the dataset from <a href="https://huggingface.co/papers/2510.26788">Defeating the Training-Inference Mismatch via FP16</a> (Qi et al., 2025). The reproduction code is in <a href="https://github.com/sail-sg/Precision-RL" rel="nofollow"><code>sail-sg/Precision-RL</code></a>.</p>
<p>The authors generated 40 answers for each MATH problem with <a href="https://huggingface.co/deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B">DeepSeek-R1-Distill-Qwen-1.5B</a>. 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.</p>
<p>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.</p>
<p>We also take the hyperparameters from the paper's LoRA scripts in <a href="https://github.com/sail-sg/Precision-RL/tree/main/oat/scripts/lora" rel="nofollow"><code>oat/scripts/lora</code></a>: <code>Qwen/Qwen2.5-Math-1.5B</code>, 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.</p>
<h3 class="relative group flex items-baseline">
<a id="the-trainer" class="block pr-1.5 text-lg md:absolute md:p-1.5 md:opacity-0 md:group-hover:opacity-100 md:right-full" href="https://huggingface.co/blog/asyncgrpo-lora-hfjobs#the-trainer" rel="nofollow">
<span class="header-link"><svg class="text-gray-500 hover:text-black dark:hover:text-gray-200 w-4" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span>
</a>
<span>
The trainer
</span>
</h3>
<p>The trainer uses the same <code>vllm/vllm-openai:v0.27.1</code> 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 <code>AsyncGRPOTrainer</code> script. The only Job-specific values are the output directory and the server URL.</p>
<pre><code class="language-python"><span class="hljs-keyword">from</span> peft <span class="hljs-keyword">import</span> LoraConfig
<span class="hljs-keyword">from</span> trl.experimental.async_grpo <span class="hljs-keyword">import</span> AsyncGRPOConfig, AsyncGRPOTrainer
config = AsyncGRPOConfig(
output_dir=<span class="hljs-string">"/lora/sanity-lora-r1"</span>, <span class="hljs-comment"># on the bucket: adapters, checkpoints and the final adapter all land here</span>
vllm_server_base_url=<span class="hljs-string">"http://localhost:8000"</span>, <span class="hljs-comment"># the proxy, not a vLLM Job; TRL never sees the Jobs URLs</span>
max_staleness=<span class="hljs-number">4</span>,
weight_sync_steps=<span class="hljs-number">4</span>, <span class="hljs-comment"># publish an adapter every 4 optimizer steps</span>
save_strategy=<span class="hljs-string">"steps"</span>, save_steps=<span class="hljs-number">50</span>, <span class="hljs-comment"># checkpoints go to the same bucket -> resume after preemption</span>
...
)
trainer = AsyncGRPOTrainer(
model=<span class="hljs-string">"Qwen/Qwen2.5-Math-1.5B"</span>,
args=config,
peft_config=LoraConfig(r=<span class="hljs-number">1</span>, lora_alpha=<span class="hljs-number">2</span>, target_modules=<span class="hljs-string">"all-linear"</span>), <span class="hljs-comment"># plain LoRA vLLM can serve as-is</span>
...
)
</code></pre>
<blockquote class="note">
<p>During initialization, TRL calls <code>/server_info</code>. If it finds a <code>lora_config</code>, it uses adapter-only sync. Configurations vLLM cannot serve directly, such as DoRA, <code>modules_to_save</code>, or a rank above <code>--max-lora-rank</code>, fall back to merged-weight sync with a warning. The log should contain <code>Adapter-only vLLM sync enabled</code>.</p>
</blockquote>
<h2 class="relative group flex items-baseline">
<a id="the-proxy" class="block pr-1.5 text-lg md:absolute md:p-1.5 md:opacity-0 md:group-hover:opacity-100 md:right-full" href="https://huggingface.co/blog/asyncgrpo-lora-hfjobs#the-proxy" rel="nofollow">
<span class="header-link"><svg class="text-gray-500 hover:text-black dark:hover:text-gray-200 w-4" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span>
</a>
<span>
The proxy
</span>
</h2>
<p>Now onto the fun stuff. We need a proxy between the trainer and the vLLM Jobs for two reasons:</p>
<ol>
<li><p>Exposed Job ports require an <code>Authorization: Bearer <HF token></code> header on every request. The proxy is where that header gets added, so TRL does not need to know about it.</p>
</li>
<li><p>We want more than one GPU generating. On a single vLLM server, the usual way to get that is <code>--data-parallel-size > 1</code>, but TRL refuses adapter-only sync in that mode, for a good reason: a call to <code>/v1/load_lora_adapter</code> only 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.</p>
</li>
</ol>
<p>We therefore run a small proxy at <code>127.0.0.1:8000</code> 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:</p>
<ul>
<li>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).</li>
<li>It broadcasts every <em>state-changing</em> request, such as <strong>adapter loads, pause and resume</strong>, to all replicas, so that a policy name means the same thing everywhere.</li>
</ul>
<h3 class="relative group flex items-baseline">
<a id="routing-rollouts-by-kv-prefix" class="block pr-1.5 text-lg md:absolute md:p-1.5 md:opacity-0 md:group-hover:opacity-100 md:right-full" href="https://huggingface.co/blog/asyncgrpo-lora-hfjobs#routing-rollouts-by-kv-prefix" rel="nofollow">
<span class="header-link"><svg class="text-gray-500 hover:text-black dark:hover:text-gray-200 w-4" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span>
</a>
<span>
Routing rollouts by KV prefix
</span>
</h3>
<p>A quick reminder of why this matters. Generating a completion has two phases with very different workload profiles:</p>
<ul>
<li>The prefill processes the <strong>whole prompt at once</strong> and computes the attention keys and values for every prompt token.</li>
<li>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.</li>
</ul>
<p>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 <strong>skip that part of the prefill entirely</strong>. 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.</p>
<p>vLLM stores its <a href="https://github.com/vllm-project/vllm/blob/main/vllm/v1/core/kv_cache_utils.py" rel="nofollow">prefix KV cache</a> in blocks of 16 tokens. Because of GRPO, the rollout worker sends <code>G</code> requests with the same prompt (in our case <code>G=8</code>). 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.</p>
<p>The job of <a href="https://github.com/AmineDiro/hfjobs-lora-buckets/blob/main/src/lora_proxy.py" rel="nofollow">our router</a> is to track which replica has seen which <strong>block hash</strong>. 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!</p>
<figure class="image text-center">
<video controls="" autoplay="" loop="" muted="" playsinline="" style="max-width: 100%; margin: auto;">
<source src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/asyncgrpo-lora-hfjobs/kv-prefix-router.mp4" type="video/mp4">
Your browser does not support the video tag.
</video>
<figcaption>The routing decision for two prompts and four requests on two replicas: 16-token blocks, chained hashes, the common prefix, one affinity hit and one spill.</figcaption>
</figure>
<p>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):</p>
<p><strong>1. Split the prompt into blocks.</strong> 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.</p>
<p><strong>2. Hash the prefix.</strong> Each block is hashed with the previous hash, starting from the adapter seed. <code>h3</code> therefore identifies blocks 1, 2 and 3 in order. Two prompts with the same first <code>k</code> blocks get the same hashes up to <code>hk</code>. 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 <code>trl-policy-v4</code> starts from another seed and cannot match entries from <code>v3</code>. This is what we want because the old KV blocks were computed with different weights.</p>
<p><strong>3. Compare two prompts.</strong> 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.</p>
<p><strong>4. Store the owners.</strong> 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 <code>h1</code> is owned by both replicas and already has several successors, <code>h2</code> to <code>h8</code> are owned by A only and each has a single successor, and <code>h2'</code> to <code>h6'</code> are owned by B only.</p>
<p>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.</p>
<p>A block is <em>common</em> 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!</p>
<p><strong>5. Pick a replica.</strong> 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:</p>
<ul>
<li>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 <strong>affinity</strong> hit.</li>
<li>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 <strong>spill</strong>.</li>
<li>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 <strong>unmatched</strong>.</li>
</ul>
<p>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 <code>h1</code> is known on both replicas.</p>
<ul>
<li><strong>Request 1, problem 0, rollout 1.</strong> 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 <code>h2</code> to <code>h8</code> as owned by A. A now has 4 requests in flight.</li>
<li><strong>Request 2, problem 0, rollout 2.</strong> 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.</li>
<li><strong>Request 3, problem 1, rollout 1.</strong> 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 <code>h2'</code> to <code>h6'</code> as owned by B.</li>
<li><strong>Request 4, problem 0, rollout 9.</strong> 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 <code>h2</code> to <code>h8</code> as 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.</li>
</ul>
<p><strong>6. Reuse the prefill.</strong> 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 <code>common</code> 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.</p>
<pre><code class="language-python"><span class="hljs-keyword">def</span> <span class="hljs-title function_">choose</span>(<span class="hljs-params">self, upstreams, model, prompt</span>):
hashes = self.block_hashes(model, prompt) <span class="hljs-comment"># chained blake2b over 16-token blocks, seeded with `model`</span>
matched = self.matched_prefix(hashes) <span class="hljs-comment"># per replica: leading blocks it has served</span>
common = self.common_prefix_len(hashes) <span class="hljs-comment"># leading blocks that identify no prompt (see below)</span>
specific = [<span class="hljs-built_in">max</span>(<span class="hljs-number">0</span>, m - common) <span class="hljs-keyword">for</span> m <span class="hljs-keyword">in</span> matched] <span class="hljs-comment"># what actually distinguishes replicas</span>
least = <span class="hljs-built_in">min</span>(u.inflight <span class="hljs-keyword">for</span> u <span class="hljs-keyword">in</span> upstreams)
best = <span class="hljs-built_in">max</span>(<span class="hljs-built_in">range</span>(self.n), key=<span class="hljs-keyword">lambda</span> i: (specific[i], -upstreams[i].inflight))
<span class="hljs-keyword">if</span> specific[best] > <span class="hljs-number">0</span> <span class="hljs-keyword">and</span> upstreams[best].inflight - least <= self.cfg.imbalance:
pick = best <span class="hljs-comment"># affinity: the replica that has this prompt, and is not swamped</span>
<span class="hljs-keyword">else</span>:
candidates = [i <span class="hljs-keyword">for</span> i <span class="hljs-keyword">in</span> <span class="hljs-built_in">range</span>(self.n) <span class="hljs-keyword">if</span> upstreams[i].inflight == least]
pick = candidates[self.rr % <span class="hljs-built_in">len</span>(candidates)] <span class="hljs-comment"># spill or new prompt: least-loaded, round robin on ties</span>
self.rr += <span class="hljs-number">1</span>
...record `pick` <span class="hljs-keyword">as</span> an owner of every block, <span class="hljs-keyword">and</span> each block<span class="hljs-string">'s successor...</span>
<span class="hljs-string"> return upstreams[pick]</span>
</code></pre>
<p>The <code>common</code> 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.</p>
<h3 class="relative group flex items-baseline">
<a id="broadcasting-the-adapter" class="block pr-1.5 text-lg md:absolute md:p-1.5 md:opacity-0 md:group-hover:opacity-100 md:right-full" href="https://huggingface.co/blog/asyncgrpo-lora-hfjobs#broadcasting-the-adapter" rel="nofollow">
<span class="header-link"><svg class="text-gray-500 hover:text-black dark:hover:text-gray-200 w-4" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span>
</a>
<span>
Broadcasting the adapter
</span>
</h3>
<p>The proxy also needs to broadcast the adapter loads to every replica. We treat the operation as <strong>all-or-nothing</strong>. Each replica has its own bucket mount, so they do not necessarily see a new adapter at exactly the same time. A <code>No adapter found for <path></code> 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.</p>
<pre><code class="language-python"><span class="hljs-keyword">async</span> <span class="hljs-keyword">def</span> <span class="hljs-title function_">load_one</span>(<span class="hljs-params">u</span>):
<span class="hljs-keyword">while</span> <span class="hljs-literal">True</span>:
status, _, out = <span class="hljs-keyword">await</span> send(u, <span class="hljs-string">"POST"</span>, <span class="hljs-string">"/v1/load_lora_adapter"</span>, headers, body)
<span class="hljs-keyword">if</span> status == <span class="hljs-number">200</span> <span class="hljs-keyword">or</span> <span class="hljs-string">"No adapter found"</span> <span class="hljs-keyword">not</span> <span class="hljs-keyword">in</span> out.decode() <span class="hljs-keyword">or</span> time.monotonic() > deadline:
<span class="hljs-keyword">return</span> u, status, out
<span class="hljs-keyword">await</span> asyncio.sleep(cfg.lora_retry_s) <span class="hljs-comment"># this replica's mount has not seen the directory yet</span>
results = <span class="hljs-keyword">await</span> asyncio.gather(*(load_one(u) <span class="hljs-keyword">for</span> u <span class="hljs-keyword">in</span> ups))
<span class="hljs-keyword">if</span> <span class="hljs-built_in">any</span>(st != <span class="hljs-number">200</span> <span class="hljs-keyword">for</span> _, st, _ <span class="hljs-keyword">in</span> results):
<span class="hljs-keyword">await</span> asyncio.gather(*(send(u, <span class="hljs-string">"POST"</span>, <span class="hljs-string">"/v1/unload_lora_adapter"</span>, headers, unload) <span class="hljs-keyword">for</span> u, st, _ <span class="hljs-keyword">in</span> results <span class="hljs-keyword">if</span> st == <span class="hljs-number">200</span>))
<span class="hljs-keyword">return</span> web.Response(status=<span class="hljs-number">504</span> <span class="hljs-keyword">if</span> timed_out <span class="hljs-keyword">else</span> st, text=<span class="hljs-string">"rolled back on the others"</span>)
</code></pre>
<p>We also broadcast <code>/pause</code>, <code>/resume</code> and <code>/v1/unload_lora_adapter</code> in the same way. <code>/health</code> returns 200 only if every replica is healthy. <code>/server_info</code> and <code>/v1/models</code> only need one answer. From TRL's point of view, the proxy is a single <code>data_parallel_size=1</code> server, so it selects adapter-only sync.</p>
<p>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 🦀).</p>
<h2 class="relative group flex items-baseline">
<a id="full-run-results" class="block pr-1.5 text-lg md:absolute md:p-1.5 md:opacity-0 md:group-hover:opacity-100 md:right-full" href="https://huggingface.co/blog/asyncgrpo-lora-hfjobs#full-run-results" rel="nofollow">
<span class="header-link"><svg class="text-gray-500 hover:text-black dark:hover:text-gray-200 w-4" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span>
</a>
<span>
Full run results
</span>
</h2>
<p>The numbers below come from the trainer's <a href="https://huggingface.co/docs/trl/en/async_grpo_trainer#logged-metrics">logged metrics</a> on trackio. The run uses <code>Qwen/Qwen2.5-Math-1.5B</code>, LoRA <code>r=1</code> on <code>all-linear</code>, 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 <code>h200x2</code> Job and each of the two vLLM replicas uses one <code>h200</code> Job. Running all three costs ~$20 per hour.</p>
<h3 class="relative group flex items-baseline">
<a id="weight-sync" class="block pr-1.5 text-lg md:absolute md:p-1.5 md:opacity-0 md:group-hover:opacity-100 md:right-full" href="https://huggingface.co/blog/asyncgrpo-lora-hfjobs#weight-sync" rel="nofollow">
<span class="header-link"><svg class="text-gray-500 hover:text-black dark:hover:text-gray-200 w-4" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span>
</a>
<span>
Weight sync
</span>
</h3>
<div class="max-w-full overflow-auto">
<table>
<thead><tr>
<th>per sync, trainer's clock, 126 syncs</th>
<th>before</th>
<th>now (p50)</th>
</tr>
</thead><tbody><tr>
<td>whole sync</td>
<td>30.8 s</td>
<td><strong>8.5 s</strong> (min 6.6, max 9.2)</td>
</tr>
<tr>
<td>of which: pause both replicas</td>
<td>0.3 s</td>
<td>0.3 s</td>
</tr>
<tr>
<td>adapter all-gather and save to the bucket</td>
<td>0.6 s</td>
<td>1.1 s</td>
</tr>
<tr>
<td>both replicas accept the adapter</td>
<td>~29 s</td>
<td>~7 s</td>
</tr>
</tbody>
</table>
</div>
<p>All 252 adapter loads succeeded 🎉: 126 syncs times 2 replicas. Six succeeded on the second attempt and 246 on the third.</p>
<h3 class="relative group flex items-baseline">
<a id="routing" class="block pr-1.5 text-lg md:absolute md:p-1.5 md:opacity-0 md:group-hover:opacity-100 md:right-full" href="https://huggingface.co/blog/asyncgrpo-lora-hfjobs#routing" rel="nofollow">
<span class="header-link"><svg class="text-gray-500 hover:text-black dark:hover:text-gray-200 w-4" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span>
</a>
<span>
Routing
</span>
</h3>
<p>At the end of the run, after 64,728 rollouts, the proxy's counters read:</p>
<pre><code>routed [31928, 32800] affinity 54712 spilled 820 unmatched 9196
</code></pre>
<p>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.</p>
<h3 class="relative group flex items-baseline">
<a id="where-the-time-goes" class="block pr-1.5 text-lg md:absolute md:p-1.5 md:opacity-0 md:group-hover:opacity-100 md:right-full" href="https://huggingface.co/blog/asyncgrpo-lora-hfjobs#where-the-time-goes" rel="nofollow">
<span class="header-link"><svg class="text-gray-500 hover:text-black dark:hover:text-gray-200 w-4" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span>
</a>
<span>
Where the time goes
</span>
</h3>
<p>The first configuration has a pretty obvious problem: the trainer is the bottleneck, not generation. Over the 500 steps, we have:</p>
<div class="max-w-full overflow-auto">
<table>
<thead><tr>
<th>per optimizer step, p50</th>
<th></th>
</tr>
</thead><tbody><tr>
<td>step</td>
<td>22.9 s</td>
</tr>
<tr>
<td>forward + backward</td>
<td>21.9 s</td>
</tr>
<tr>
<td>waiting for rollouts</td>
<td>0.02 s</td>
</tr>
<tr>
<td>rollout queue occupancy</td>
<td>476 of 512</td>
</tr>
<tr>
<td>trainer MFU</td>
<td>3.9 %</td>
</tr>
</tbody>
</table>
</div>
<p>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.</p>
<h3 class="relative group flex items-baseline">
<a id="reward" class="block pr-1.5 text-lg md:absolute md:p-1.5 md:opacity-0 md:group-hover:opacity-100 md:right-full" href="https://huggingface.co/blog/asyncgrpo-lora-hfjobs#reward" rel="nofollow">
<span class="header-link"><svg class="text-gray-500 hover:text-black dark:hover:text-gray-200 w-4" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span>
</a>
<span>
Reward
</span>
</h3>
<p><a href="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/asyncgrpo-lora-hfjobs/fig0-r1-dp2-reward.png" rel="nofollow"><img src="https://imgproxy.citrusreader.com/qrfP-14bKkzWwlle5Dh3fdWd4Qv7FMPGAmDF0sjMrgE/raw:1/aHR0cHM6Ly9odWdnaW5nZmFjZS5jby9kYXRhc2V0cy9odWdnaW5nZmFjZS9kb2N1bWVudGF0aW9uLWltYWdlcy9yZXNvbHZlL21haW4vYmxvZy9hc3luY2dycG8tbG9yYS1oZmpvYnMvZmlnMC1yMS1kcDItcmV3YXJkLnBuZw" alt="Figure 1" referrerpolicy="no-referrer"></a></p>
<p><em>Figure 1. trackio run <code>r1-dp2</code>. Panels: <code>reward</code> with its 20-step rolling mean and 50-step block means, and <code>ratio</code> on a 0.99 to 1.01 axis. Reward climbs from 0.15 to 0.44 over 500 steps; <code>ratio</code> stays between 0.9993 and 1.0004 throughout.</em></p>
<p>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, <code>ratio</code> 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 <a href="https://huggingface.co/spaces/aminediroHF/async-grpo-lora-buckets">trackio dashboard</a> has the full curves.</p>
<p>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.</p>
<h2 class="relative group flex items-baseline">
<a id="chasing-the-bottleneck-ping-pong" class="block pr-1.5 text-lg md:absolute md:p-1.5 md:opacity-0 md:group-hover:opacity-100 md:right-full" href="https://huggingface.co/blog/asyncgrpo-lora-hfjobs#chasing-the-bottleneck-ping-pong" rel="nofollow">
<span class="header-link"><svg class="text-gray-500 hover:text-black dark:hover:text-gray-200 w-4" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span>
</a>
<span>
Chasing the bottleneck ping-pong
</span>
</h2>
<p>Async RL is a pipeline between training and generation. Making one side faster does nothing if the other side cannot keep up. Fortunately, in <code>AsyncGRPOTrainer</code> we've added enough timings and metrics to see this directly.</p>
<p>All the useful metrics are documented in the <a href="https://huggingface.co/docs/trl/en/async_grpo_trainer#logged-metrics">Logged metrics</a> section. <code>perf/rollout_wait_s</code> tells us how long the trainer waits for samples. <code>rollout/backpressure_s</code> 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.</p>
<p>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.</p>
<h3 class="relative group flex items-baseline">
<a id="reading-the-dashboard" class="block pr-1.5 text-lg md:absolute md:p-1.5 md:opacity-0 md:group-hover:opacity-100 md:right-full" href="https://huggingface.co/blog/asyncgrpo-lora-hfjobs#reading-the-dashboard" rel="nofollow">
<span class="header-link"><svg class="text-gray-500 hover:text-black dark:hover:text-gray-200 w-4" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span>
</a>
<span>
Reading the dashboard
</span>
</h3>
<p>We keep these four groups of metrics visible:</p>
<ul>
<li><strong><code>perf/step_s</code></strong> and <strong><code>perf/fwd_bwd_s</code></strong>: 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.</li>
<li><strong><code>perf/rollout_wait_s</code></strong>: 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.</li>
<li><strong><code>sample/rollout_queue_size</code></strong> against <code>queue_maxsize</code>: the buffer between the two sides. Full means generation is being throttled; empty means the trainer is starving.</li>
<li><strong><code>rollout/backpressure_s</code></strong> and <strong><code>rollout/score_block_s</code></strong>: 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 is <code>rollout/backpressure_s</code>. Scoring then stops draining its own input queue, so generation cannot hand over the next group either, which is <code>rollout/score_block_s</code>. Both are the same stall, seen first at the scoring stage and then propagated back to generation.</li>
</ul>
<p>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 <code>perf/mfu_wall_clock</code> with <code>perf/mfu_fwd_bwd</code> also shows how much time the trainer GPUs spend waiting instead of training.</p>
<h3 class="relative group flex items-baseline">
<a id="run-1-r1-dp2-a-trainer-that-cannot-keep-up" class="block pr-1.5 text-lg md:absolute md:p-1.5 md:opacity-0 md:group-hover:opacity-100 md:right-full" href="https://huggingface.co/blog/asyncgrpo-lora-hfjobs#run-1-r1-dp2-a-trainer-that-cannot-keep-up" rel="nofollow">
<span class="header-link"><svg class="text-gray-500 hover:text-black dark:hover:text-gray-200 w-4" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span>
</a>
<span>
Run 1, <code>r1-dp2</code>: a trainer that cannot keep up
</span>
</h3>
<p><a href="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/asyncgrpo-lora-hfjobs/fig1-r1-dp2-trainer-bound.png" rel="nofollow"><img src="https://imgproxy.citrusreader.com/zudD71X_dUULrQnbhkv4oBF5Q-8I_flmdHj4XtXs0nA/raw:1/aHR0cHM6Ly9odWdnaW5nZmFjZS5jby9kYXRhc2V0cy9odWdnaW5nZmFjZS9kb2N1bWVudGF0aW9uLWltYWdlcy9yZXNvbHZlL21haW4vYmxvZy9hc3luY2dycG8tbG9yYS1oZmpvYnMvZmlnMS1yMS1kcDItdHJhaW5lci1ib3VuZC5wbmc" alt="Figure 2" referrerpolicy="no-referrer"></a></p>
<p><em>Figure 2. trackio run <code>r1-dp2</code>. Panels: <code>perf/step_s</code>, <code>perf/fwd_bwd_s</code>, <code>sample/rollout_queue_size</code>, <code>rollout/backpressure_s</code>. 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.</em></p>
<p><code>perf/step_s</code> is 22.9 s and <code>perf/fwd_bwd_s</code> 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.</p>
<p>The batch metrics explain the terrible 3.9 % MFU. <code>batch/microbatches_per_step</code> is 64 and <code>batch/samples_per_row</code> is 1.0. Each rank processes one sequence of around 1.2k tokens, 64 times per step. This comes from the reference recipe's <code>per_device_train_batch_size=1</code>. For a 1.5B model on an H200, this is completely latency-bound.</p>
<h3 class="relative group flex items-baseline">
<a id="run-2-r1-dp2-tb16k-pack-the-microbatch" class="block pr-1.5 text-lg md:absolute md:p-1.5 md:opacity-0 md:group-hover:opacity-100 md:right-full" href="https://huggingface.co/blog/asyncgrpo-lora-hfjobs#run-2-r1-dp2-tb16k-pack-the-microbatch" rel="nofollow">
<span class="header-link"><svg class="text-gray-500 hover:text-black dark:hover:text-gray-200 w-4" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span>
</a>
<span>
Run 2, <code>r1-dp2-tb16k</code>: pack the microbatch
</span>
</h3>
<p>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 <strong>token-budget batching</strong>. With <code>token_budget > 0</code>, it packs several samples into one padding-free row per rank. An optimizer step processes <code>gradient_accumulation_steps</code> rows. We set <code>token_budget=16384</code> and <code>gradient_accumulation_steps=6</code>.</p>
<p><a href="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/asyncgrpo-lora-hfjobs/fig2-r1-dp2-vs-tb16k-packing.png" rel="nofollow"><img src="https://imgproxy.citrusreader.com/M_aPUT3QkmiTTfRCCU6kl4dHe7FZESGDqKxkAnmFwRM/raw:1/aHR0cHM6Ly9odWdnaW5nZmFjZS5jby9kYXRhc2V0cy9odWdnaW5nZmFjZS9kb2N1bWVudGF0aW9uLWltYWdlcy9yZXNvbHZlL21haW4vYmxvZy9hc3luY2dycG8tbG9yYS1oZmpvYnMvZmlnMi1yMS1kcDItdnMtdGIxNmstcGFja2luZy5wbmc" alt="Figure 3" referrerpolicy="no-referrer"></a></p>
<p><em>Figure 3. trackio runs <code>r1-dp2</code> and <code>r1-dp2-tb16k</code> overlaid over their first 154 steps. Panels: <code>batch/samples_per_row</code>, <code>batch/microbatches_per_step</code>, <code>perf/step_s</code>, <code>perf/fwd_bwd_s</code>, <code>perf/mfu_fwd_bwd</code>, <code>rollout/generated_tok_s</code>. 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.</em></p>
<p><code>batch/samples_per_row</code> 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.</p>
<p>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.</p>
<h3 class="relative group flex items-baseline">
<a id="run-3-r1-dp2-tb16k-nockpt-stop-recomputing-the-forward" class="block pr-1.5 text-lg md:absolute md:p-1.5 md:opacity-0 md:group-hover:opacity-100 md:right-full" href="https://huggingface.co/blog/asyncgrpo-lora-hfjobs#run-3-r1-dp2-tb16k-nockpt-stop-recomputing-the-forward" rel="nofollow">
<span class="header-link"><svg class="text-gray-500 hover:text-black dark:hover:text-gray-200 w-4" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span>
</a>
<span>
Run 3, <code>r1-dp2-tb16k-nockpt</code>: stop recomputing the forward
</span>
</h3>
<p><code>perf/fwd_s</code> is 1.34 s while <code>perf/fwd_bwd_s</code> 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.</p>
<p>The reason is that <code>AsyncGRPOConfig</code> defaults to <code>gradient_checkpointing=True</code>. 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.</p>
<p><a href="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/asyncgrpo-lora-hfjobs/fig3-tb16k-vs-nockpt-crossover.png" rel="nofollow"><img src="https://imgproxy.citrusreader.com/L_FeSCY8nlnJ1K8BqKMDbKLVY8klgsWlT9E3OH03LIo/raw:1/aHR0cHM6Ly9odWdnaW5nZmFjZS5jby9kYXRhc2V0cy9odWdnaW5nZmFjZS9kb2N1bWVudGF0aW9uLWltYWdlcy9yZXNvbHZlL21haW4vYmxvZy9hc3luY2dycG8tbG9yYS1oZmpvYnMvZmlnMy10YjE2ay12cy1ub2NrcHQtY3Jvc3NvdmVyLnBuZw" alt="Figure 4" referrerpolicy="no-referrer"></a></p>
<p><em>Figure 4. trackio runs <code>r1-dp2-tb16k</code> and <code>r1-dp2-tb16k-nockpt</code> overlaid over their first 134 steps. Panels: <code>perf/fwd_s</code>, <code>perf/fwd_bwd_s</code>, <code>perf/weight_sync_s</code>, <code>sample/rollout_queue_size</code>, <code>perf/rollout_wait_s</code>, <code>perf/mfu_fwd_bwd</code>. 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.</em></p>
<p>With <code>gradient_checkpointing=False</code>, 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 <em>successfully</em> moved the bottleneck to generation.</p>
<p>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.</p>
<h3 class="relative group flex items-baseline">
<a id="run-4-r1-dp3-tb16k-nockpt-three-replicas-and-a-surprise" class="block pr-1.5 text-lg md:absolute md:p-1.5 md:opacity-0 md:group-hover:opacity-100 md:right-full" href="https://huggingface.co/blog/asyncgrpo-lora-hfjobs#run-4-r1-dp3-tb16k-nockpt-three-replicas-and-a-surprise" rel="nofollow">
<span class="header-link"><svg class="text-gray-500 hover:text-black dark:hover:text-gray-200 w-4" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span>
</a>
<span>
Run 4, <code>r1-dp3-tb16k-nockpt</code>: three replicas, and a surprise
</span>
</h3>
<p>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 <code>fsdp_reshard_after_forward</code> to check whether FSDP2 re-gathers caused the extra 2 seconds in backward.</p>
<p><a href="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/asyncgrpo-lora-hfjobs/fig4-dp2-vs-dp3-inflight-cap.png" rel="nofollow"><img src="https://imgproxy.citrusreader.com/crnUE5Ua1RDFM947v5FmtGBtNliIkVJFdlbUN1W8ScE/raw:1/aHR0cHM6Ly9odWdnaW5nZmFjZS5jby9kYXRhc2V0cy9odWdnaW5nZmFjZS9kb2N1bWVudGF0aW9uLWltYWdlcy9yZXNvbHZlL21haW4vYmxvZy9hc3luY2dycG8tbG9yYS1oZmpvYnMvZmlnNC1kcDItdnMtZHAzLWluZmxpZ2h0LWNhcC5wbmc" alt="Figure 5" referrerpolicy="no-referrer"></a></p>
<p><em>Figure 5. trackio runs <code>r1-dp2-tb16k-nockpt</code> and <code>r1-dp3-tb16k-nockpt</code> overlaid over their first 134 steps. Panels: <code>perf/weight_sync_s</code>, <code>rollout/generated_tok_s</code>, <code>rollout/inflight</code>, <code>perf/fwd_bwd_s</code>. Sync falls from 7.6 s to 5.8 s; generation and forward+backward do not move; <code>rollout/inflight</code> reads 128 in both runs, which is the cap the third replica ran into.</em></p>
<p>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.</p>
<p>The reason was sitting in <code>rollout/inflight</code>: 128 in every run. The proxy shows those requests split as 44 + 43 + 41 across the three replicas. <code>max_inflight_tasks</code> 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.</p>
<p>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.</p>
<h3 class="relative group flex items-baseline">
<a id="run-5-r1-dp3-inflight384-lift-the-cap-on-in-flight-requests" class="block pr-1.5 text-lg md:absolute md:p-1.5 md:opacity-0 md:group-hover:opacity-100 md:right-full" href="https://huggingface.co/blog/asyncgrpo-lora-hfjobs#run-5-r1-dp3-inflight384-lift-the-cap-on-in-flight-requests" rel="nofollow">
<span class="header-link"><svg class="text-gray-500 hover:text-black dark:hover:text-gray-200 w-4" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span>
</a>
<span>
Run 5, <code>r1-dp3-inflight384</code>: lift the cap on in-flight requests
</span>
</h3>
<p><code>max_inflight_tasks=384</code> and <code>queue_maxsize=768</code>, nothing else.</p>
<p><a href="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/asyncgrpo-lora-hfjobs/fig5-all-runs-scoreboard.png" rel="nofollow"><img src="https://imgproxy.citrusreader.com/fVDBdZ-qCJh4O6UJLmv65zaq9l1KIkVJBzWTsIa-Klc/raw:1/aHR0cHM6Ly9odWdnaW5nZmFjZS5jby9kYXRhc2V0cy9odWdnaW5nZmFjZS9kb2N1bWVudGF0aW9uLWltYWdlcy9yZXNvbHZlL21haW4vYmxvZy9hc3luY2dycG8tbG9yYS1oZmpvYnMvZmlnNS1hbGwtcnVucy1zY29yZWJvYXJkLnBuZw" alt="Figure 6" referrerpolicy="no-referrer"></a></p>
<p><em>Figure 6. All five trackio runs (<code>r1-dp2</code>, <code>r1-dp2-tb16k</code>, <code>r1-dp2-tb16k-nockpt</code>, <code>r1-dp3-tb16k-nockpt</code>, <code>r1-dp3-inflight384</code>) overlaid, x-axis in steps. Panels: <code>perf/step_s</code>, <code>reward</code>, <code>sample/rollout_queue_size</code>, <code>sample/staleness_mean</code>. 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.</em></p>
<p>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.</p>
<p>Mean staleness rises from 1.5 to 2.0 versions because samples wait longer in the larger queue. This is still below <code>max_staleness=4</code>, and <code>ratio</code> remains very close to 1.000.</p>
<h3 class="relative group flex items-baseline">
<a id="the-scoreboard" class="block pr-1.5 text-lg md:absolute md:p-1.5 md:opacity-0 md:group-hover:opacity-100 md:right-full" href="https://huggingface.co/blog/asyncgrpo-lora-hfjobs#the-scoreboard" rel="nofollow">
<span class="header-link"><svg class="text-gray-500 hover:text-black dark:hover:text-gray-200 w-4" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span>
</a>
<span>
The scoreboard
</span>
</h3>
<div class="max-w-full overflow-auto">
<table>
<thead><tr>
<th>500 steps</th>
<th>run 1 <code>r1-dp2</code></th>
<th>run 5 <code>r1-dp3-inflight384</code></th>
</tr>
</thead><tbody><tr>
<td>wall clock</td>
<td>3 h 27 min</td>
<td><strong>53 min</strong></td>
</tr>
<tr>
<td><code>perf/step_s</code>, p50</td>
<td>22.9 s</td>
<td>4.8 s</td>
</tr>
<tr>
<td><code>perf/fwd_bwd_s</code>, p50</td>
<td>21.9 s</td>
<td>4.6 s</td>
</tr>
<tr>
<td><code>perf/mfu_fwd_bwd</code></td>
<td>3.9 %</td>
<td>23.5 %</td>
</tr>
<tr>
<td><code>batch/samples_per_step</code></td>
<td>128</td>
<td>168</td>
</tr>
<tr>
<td>samples trained</td>
<td>64 000</td>
<td>84 078</td>
</tr>
<tr>
<td><code>perf/weight_sync_s</code>, p50</td>
<td>8.5 s</td>
<td>6.2 s</td>
</tr>
<tr>
<td><code>sample/staleness_mean</code></td>
<td>1.5</td>
<td>2.0</td>
</tr>
<tr>
<td>reward, first 20 → last 20 steps</td>
<td>0.145 → 0.438</td>
<td>0.145 → 0.416</td>
</tr>
</tbody>
</table>
</div>
<p><a href="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/asyncgrpo-lora-hfjobs/fig6-reward-vs-wallclock.png" rel="nofollow"><img src="https://imgproxy.citrusreader.com/iOObaiFhH8dPjSvrTNoOGtI8P9TqgVnNbksswmzwzFQ/raw:1/aHR0cHM6Ly9odWdnaW5nZmFjZS5jby9kYXRhc2V0cy9odWdnaW5nZmFjZS9kb2N1bWVudGF0aW9uLWltYWdlcy9yZXNvbHZlL21haW4vYmxvZy9hc3luY2dycG8tbG9yYS1oZmpvYnMvZmlnNi1yZXdhcmQtdnMtd2FsbGNsb2NrLnBuZw" alt="Figure 7" referrerpolicy="no-referrer"></a></p>
<p><em>Figure 7. trackio runs <code>r1-dp2</code> and <code>r1-dp3-inflight384</code>, 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.</em></p>
<p>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.</p>
<h2 class="relative group flex items-baseline">
<a id="try-it" class="block pr-1.5 text-lg md:absolute md:p-1.5 md:opacity-0 md:group-hover:opacity-100 md:right-full" href="https://huggingface.co/blog/asyncgrpo-lora-hfjobs#try-it" rel="nofollow">
<span class="header-link"><svg class="text-gray-500 hover:text-black dark:hover:text-gray-200 w-4" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span>
</a>
<span>
Try it
</span>
</h2>
<pre><code class="language-sh">git <span class="hljs-built_in">clone</span> https://github.com/AmineDiro/hfjobs-lora-buckets && <span class="hljs-built_in">cd</span> hfjobs-lora-buckets
hf auth login
MAX_STEPS=20 RUN_TAG=smoke ./run_all.sh --<span class="hljs-built_in">wait</span> <span class="hljs-comment"># ~15 min, three Jobs, cancels the servers when done</span>
MAX_STEPS=500 ./run_all.sh --<span class="hljs-built_in">wait</span> <span class="hljs-comment"># run 1: the reference batch shape, ~3.5 h</span>
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 --<span class="hljs-built_in">wait</span> <span class="hljs-comment"># run 5: same recipe, ~55 min</span>
</code></pre>
<h2 class="relative group flex items-baseline">
<a id="references" class="block pr-1.5 text-lg md:absolute md:p-1.5 md:opacity-0 md:group-hover:opacity-100 md:right-full" href="https://huggingface.co/blog/asyncgrpo-lora-hfjobs#references" rel="nofollow">
<span class="header-link"><svg class="text-gray-500 hover:text-black dark:hover:text-gray-200 w-4" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span>
</a>
<span>
References
</span>
</h2>
<ul>
<li>John Schulman et al., <a href="https://thinkingmachines.ai/blog/lora/" rel="nofollow">LoRA Without Regret</a>, Thinking Machines Lab, September 2025. The case that rank-1 LoRA matches full fine-tuning for policy-gradient RL, and why.</li>
<li>TRL, <a href="https://huggingface.co/docs/trl/en/async_grpo_trainer"><code>AsyncGRPOTrainer</code></a> and its <a href="https://huggingface.co/docs/trl/en/async_grpo_trainer#logged-metrics">logged metrics</a>.</li>
<li>TRL <a href="https://github.com/huggingface/trl/pull/7017" rel="nofollow">PR #7017</a>: PEFT/LoRA support for <code>AsyncGRPOTrainer</code> with adapter-only vLLM sync, released in TRL v1.14.</li>
<li>Hugging Face <a href="https://huggingface.co/docs/huggingface_hub/guides/jobs">Jobs</a> and <a href="https://huggingface.co/docs/hub/storage-buckets">Storage Buckets</a>; <a href="https://github.com/huggingface/hf-mount" rel="nofollow"><code>hf-mount</code></a>.</li>
<li><a href="https://github.com/AmineDiro/hf-mount-repro" rel="nofollow"><code>hf-mount-repro</code></a>: the two-script reproduction of the 30-second negative-cache stall.</li>
<li>The <a href="https://huggingface.co/spaces/aminediroHF/async-grpo-lora-buckets">trackio dashboard</a> for every run in this post.</li>
<li>Penghui Qi, Zichen Liu, Xiangxin Zhou, Tianyu Pang, Chao Du, Wee Sun Lee, Min Lin, <a href="https://huggingface.co/papers/2510.26788">Defeating the Training-Inference Mismatch via FP16</a>, arXiv:2510.26788, 2025. Source of the Sanity dataset <a href="https://huggingface.co/datasets/sail/Sanity-Test-R1D-1.5B"><code>sail/Sanity-Test-R1D-1.5B</code></a> and of the LoRA recipe, <a href="https://github.com/sail-sg/Precision-RL" rel="nofollow"><code>sail-sg/Precision-RL</code></a>, <code>oat/scripts/lora/bf16_grpo_tis_lora.sh</code>.</li>
</ul>
<pre><code class="language-bibtex">@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}
}
</code></pre>
</div>