Async GRPO with LoRA Across HF Jobs: Shared Storage, a Proxy, and No NCCL

Async GRPO with LoRA Across HF Jobs: Shared Storage, a Proxy, and No NCCL

Overview

TRL v1.14 adds LoRA support to AsyncGRPOTrainer. The trainer can now update a LoRA adapter and synchronize only that adapter with vLLM, rather than transferring the full model.

A rank-1 adapter for a 1.5B model is only a few megabytes, compared with roughly 3 GB for the complete model. That difference makes it practical to run training and inference as separate Hugging Face Jobs on separate machines. The adapter can be written to a shared Storage Bucket instead of being transferred through NCCL.

The resulting system includes:

• A trainer Job running AsyncGRPOTrainer with LoRA and FSDP.
• Two vLLM Jobs, each serving the base model and the most recently published adapter.
• A Storage Bucket mounted at the same path in all Jobs.
• A proxy that adds authentication, routes requests according to KV-cache locality, and broadcasts adapter updates.

Across five runs, the same 500-step recipe improved from 3 h 27 min to 53 min. The final run was 3.9 times faster while producing a similar reward curve.

Why LoRA Enables Separate Jobs

LoRA is well suited to reinforcement learning. The LoRA Without Regret work from Thinking Machines reports that LoRA can match full fine-tuning for policy-gradient RL, including with rank 1. The advantage function provides approximately O(1) bits of information per episode, so a rank-1 adapter can have enough capacity for the update required at each step.

LoRA also changes the systems problem. A rank-1 adapter for a 1.5B model is a few megabytes, while the full model is about 3 GB. Instead of sending the full policy to inference workers after every update, the trainer can publish only the adapter. vLLM can keep several adapters loaded simultaneously, allowing older rollouts to finish with the policy under which they began while newer rollouts use the latest version.

AsyncGRPOTrainer already separates training from generation. In a cluster or single-node environment, the two processes can share a filesystem or communicate through NCCL. Hugging Face Jobs are different: one Job runs in one container on one VM, and a Job cannot currently span multiple nodes. Jobs also do not provide a shared local disk, shared localhost, or cross-node NCCL communication.

Storage Buckets provide the missing shared filesystem. Each Job mounts the same bucket through hf-mount, exposing it as a POSIX filesystem. The trainer writes the adapter to the bucket, and the vLLM Jobs read it from the same path.

Architecture with Hugging Face Jobs and Storage Buckets

The adapter-only synchronization path works as follows:

1. The trainer saves the adapter under <output_dir>/.vllm_lora/trl-policy-v{N}.
2. It publishes the directory with an atomic rename.
3. It sends the path to vLLM's /v1/load_lora_adapter endpoint.
4. vLLM reads the adapter files from the shared filesystem.
5. Rollout workers request the adapter using model="trl-policy-v{N}".

A bucket is mounted at the same absolute path in every Job:

## Every Job receives the same bucket at the same path
hf jobs run ... -v hf://buckets/aminediroHF/asyncgrpo-lora-buckets:/lora ...

The trainer writes to /lora/<run>/.vllm_lora/, and the servers read from that location. No changes to TRL or vLLM are required for the shared path.

Checkpoints and the final adapter are also stored in the bucket. Because Jobs are ephemeral, this allows a preempted trainer to resume without losing the persisted adapter or checkpoints.

The Three Jobs

vLLM Replicas

Each vLLM replica uses one GPU and the vllm/vllm-openai image. Runtime LoRA loading is enabled, and enough adapter slots are reserved for the allowed staleness.

With max_staleness=4, a rollout generated under trl-policy-v3 can still be used when the trainer has reached v7. vLLM therefore needs to serve the current policy plus the four previous versions. During a swap, it must load the new adapter before unloading the oldest one, requiring one additional slot. This results in:

max_staleness + 2 = 6 adapter slots

With only five slots, vLLM could evict a policy that still had rollouts in flight.

## --expose 8000                       available at https://<job_id>--8000.hf.jobs
## -v ...:/lora:ro                     the server reads adapters only
## VLLM_ALLOW_RUNTIME_LORA_UPDATING=1  enables /v1/load_lora_adapter
## VLLM_SERVER_DEV_MODE=1              enables /pause, /resume, and /server_info
## --max-loras 6                       max_staleness=4 requires 4+2 slots
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

The setup pins vLLM to v0.27.1, because the flags and runtime LoRA endpoints used here are provided by that version.

The adapters use versioned names rather than replacing one adapter under a constant name. vLLM keys its prefix cache by adapter name. Reusing one name could allow KV blocks created with earlier weights to match after an update, causing a rollout to use one policy for prefill and another for decoding. Versioned adapter names prevent that mismatch because each name identifies one set of weights.

Dataset: The Sanity Set

The experiments use sail/Sanity-Test-R1D-1.5B, introduced in Defeating the Training-Inference Mismatch via FP16 by Qi et al. (2025). The reproduction code is available in sail-sg/Precision-RL.

The authors generated 40 answers for each MATH problem with DeepSeek-R1-Distill-Qwen-1.5B. They retained problems with success rates between 20% and 80%, producing 1,460 questions. This makes the dataset suitable for RL validation because the problems are neither already solved nor entirely beyond the model's ability.

The experiment uses the paper's LoRA settings:

• Model: Qwen/Qwen2.5-Math-1.5B
• LoRA rank: 1
• LoRA alpha: 2
• Learning rate: 4e-5
• Samples per prompt: 8
• Completions per step: 128
• Maximum generated tokens: 3,000
• Context length: 4,096 tokens

Trainer Configuration

The trainer uses vllm/vllm-openai:v0.27.1 with TRL installed. The training script is a standard AsyncGRPOTrainer configuration, with the bucket output directory and proxy URL as the Job-specific values.

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

config = AsyncGRPOConfig(
    output_dir="/lora/sanity-lora-r1",
    vllm_server_base_url="http://localhost:8000",
    max_staleness=4,
    weight_sync_steps=4,
    save_strategy="steps",
    save_steps=50,
    ...
)

trainer = AsyncGRPOTrainer(
    model="Qwen/Qwen2.5-Math-1.5B",
    args=config,
    peft_config=LoraConfig(
        r=1,
        lora_alpha=2,
        target_modules="all-linear",
    ),
    ...
)

During initialization, TRL calls /server_info. If the response contains a lora_config, TRL enables adapter-only synchronization. Configurations that vLLM cannot serve directly, including DoRA, modules_to_save, or a rank greater than --max-lora-rank, fall back to merged-weight synchronization with a warning. The log should include Adapter-only vLLM sync enabled.

The Proxy

The proxy is required for two reasons.

First, exposed Job ports require an Authorization: Bearer <HF token> header on every request. The proxy adds this header so that TRL does not need to manage the Job credentials.

Second, generation uses multiple vLLM machines. TRL does not use adapter-only synchronization with --data-parallel-size > 1, because /v1/load_lora_adapter would reach only the responding data-parallel rank. Other ranks could continue serving the base model under the new policy name. With separate Jobs, each replica is an independent machine, so adapter loading must be broadcast at the proxy layer.

The proxy runs at 127.0.0.1:8000 on the trainer Job. TRL treats it as one vLLM server. The proxy:

• Routes completion requests to replicas with the relevant KV prefix already cached.
• Sends adapter loads, pause requests, resume requests, and other state-changing operations to every replica.

Routing Rollouts by KV Prefix

Completion generation has two phases:

• Prefill processes the entire prompt and computes attention keys and values for its tokens.
• Decode generates one token at a time, attending to the cached keys and values.

These keys and values form the KV cache. Because attention is causal, requests sharing a prefix can share the corresponding KV blocks. A replica that already holds the prefix can skip that part of prefill.

vLLM stores prefix KV cache data in blocks of 16 tokens. GRPO sends G requests with the same prompt, with G=8 in this experiment. If all eight requests reach one replica, the first performs the prefill and the other seven reuse it. Round-robin routing would send some requests to a replica without the prefix, repeating the prefill unnecessarily.

The router tracks which replica has seen each block hash. The hashes are chained: the hash for block 3 represents blocks 1, 2, and 3. This reflects causal attention, because the KV for block 3 is valid only when the earlier blocks are identical. The adapter name seeds the chain, ensuring that a prefix cached for trl-policy-v3 cannot match one for trl-policy-v4.

The routing process is:

1. Split the prompt into blocks. Token IDs are divided into 16-token blocks. Only complete blocks are hashed, so a final partial block is ignored.
2. Hash the prefix. Each block is hashed with the previous hash, starting from the adapter name. Two prompts with the same first k blocks share hashes through hk.
3. Compare prompts. Shared chat-template tokens may match, while problem-specific blocks diverge.
4. Record ownership. The router records which replicas served each hash and which hashes followed it.
5. Choose a replica. The router counts matching leading blocks, removes the common prefix, and routes according to cache affinity and load.
6. Reuse the prefill. A request with a matching prompt-specific prefix can skip the cached prefill.

A block is considered common when every replica has served it or when it has multiple successors. Common blocks do not identify a particular prompt and are excluded from affinity decisions. This avoids treating a shared system prompt or chat template as a useful cache hit.

The routing rules are:

• If a replica has prompt-specific blocks and is no more than 8 requests ahead of the least-loaded replica, route there as an affinity hit.
• If it has the relevant blocks but is more than 8 requests ahead, route to the least-loaded replica as a spill.
• If no replica has prompt-specific blocks, route to the least-loaded replica, using round-robin selection for ties. This is an unmatched request.

A simplified version of the decision logic is:

def choose(self, upstreams, model, prompt):
    hashes = self.block_hashes(model, prompt)
    matched = self.matched_prefix(hashes)
    common = self.common_prefix_len(hashes)
    specific = [max(0, m - common) for m in matched]
    least = min(u.inflight for u in upstreams)
    best = max(
        range(self.n),
        key=lambda i: (specific[i], -upstreams[i].inflight),
    )

    if (
        specific[best] > 0
        and upstreams[best].inflight - least <= self.cfg.imbalance
    ):
        pick = best
    else:
        candidates = [
            i for i in range(self.n)
            if upstreams[i].inflight == least
        ]
        pick = candidates[self.rr % len(candidates)]
        self.rr += 1

    # Record pick as an owner of every block and successor.
    return upstreams[pick]

Broadcasting Adapters

The proxy treats adapter loading as an all-or-nothing operation. Each replica has its own bucket mount, so replicas may see a newly published adapter at slightly different times. A No adapter found for <path> response normally means that the mount has not caught up, so only that replica is retried.

For another type of error, the proxy unloads the adapter from replicas that accepted it. This ensures that a policy name does not exist on only part of the fleet.

async def load_one(u):
    while True:
        status, _, out = await send(
            u,
            "POST",
            "/v1/load_lora_adapter",
            headers,
            body,
        )
        if (
            status == 200
            or "No adapter found" not in out.decode()
            or time.monotonic() > deadline
        ):
            return u, status, out
        await asyncio.sleep(cfg.lora_retry_s)

results = await asyncio.gather(*(load_one(u) for u in ups))
if any(st != 200 for _, st, _ in results):
    await asyncio.gather(
        *(
            send(
                u,
                "POST",
                "/v1/unload_lora_adapter",
                headers,
                unload,
            )
            for u, st, _ in results
            if st == 200
        )
    )
    return web.Response(
        status=504 if timed_out else st,
        text="rolled back on the others",
    )

The proxy broadcasts /pause, /resume, and /v1/unload_lora_adapter in the same way. /health returns 200 only when every replica is healthy. /server_info and /v1/models need only one response.

At this scale, a Python asyncio proxy did not become a bottleneck. There were at most 128 non-streaming JSON requests in flight, and routing required only a few hash calculations.

Full-Run Results

The results use Qwen/Qwen2.5-Math-1.5B, LoRA r=1 on all-linear, 128 completions per step, and 8 rollouts per prompt. Each run has 500 steps and saves checkpoints every 50 steps.

The trainer uses an h200x2 Job. Each vLLM replica uses one h200 Job. The three Jobs cost approximately $20 per hour.

Weight Synchronization

Across 126 synchronizations, the measurements were:

All 252 adapter loads succeeded. Six succeeded on the second attempt and 246 succeeded on the third.

Routing

After 64,728 rollouts, the proxy counters were:

routed [31928, 32800]
affinity 54712
spilled 820
unmatched 9196

With eight rollouts per prompt, at least one request in each group must be cold. The theoretical minimum cold-request rate is 12.5%. The router recorded 14.2% unmatched requests, 84.5% affinity hits, and 1.3% spills.

Initial Bottleneck

The first configuration was limited by the trainer rather than generation:

The rollout queue remained full, and the trainer was mostly blocked by backpressure. The second vLLM replica was largely unused because the trainer could not consume samples quickly enough.

Reward and Policy Consistency

The first 500-step run took 3 h 27 min. Mean reward increased from 0.145 over the first 20 steps to 0.438 over the last 20 steps.

The ratio metric remained at 1.000 throughout, staying between 0.9993 and 1.0004. This indicates that vLLM served the same policy used by the trainer to score the rollouts across all 126 synchronizations. Mean staleness was 1.5 policy versions, below the maximum of 4.

Moving the Bottleneck

Async RL forms a pipeline between training and generation. Improving one side does not help if the other side cannot keep up. The most useful metrics are:

• perf/step_s and perf/fwd_bwd_s: optimizer-step duration and forward/backward duration.
• perf/rollout_wait_s: time the trainer waits for samples.
• sample/rollout_queue_size: the buffer between generation and training.
• rollout/backpressure_s and rollout/score_block_s: time the rollout worker is blocked because the queue is full.

A full queue, nearly zero rollout wait, and high backpressure indicate a trainer-bound run. An empty queue, increasing rollout wait, and little backpressure indicate that generation is too slow.

Run 1: r1-dp2

The first run was trainer-bound. perf/step_s was 22.9 seconds and perf/fwd_bwd_s was 21.9 seconds. Forward and backward consumed 96% of the step. The queue stayed near 476 of 512, while the rollout worker spent about 15 seconds per group blocked by backpressure.

The batch configuration used per_device_train_batch_size=1, resulting in 64 microbatches per step and one sample per row. For a 1.5B model on an H200, this was latency-bound and produced only 3.9% MFU.

Run 2: r1-dp2-tb16k

The next run kept 128 completions per optimizer step but packed multiple samples into each GPU row using token-budget batching:

• token_budget=16384
• gradient_accumulation_steps=6

Samples per row increased from 1.0 to approximately 12.7, while microbatches fell from 64 to 6. Forward and backward time decreased from 21.9 seconds to 5.6 seconds, and MFU increased from 3.9% to 19%.

Generation increased from approximately 4.6k to 25k tokens per second without changing the vLLM configuration. The queue was no longer continuously full, allowing the replicas to use their available capacity.

Run 3: r1-dp2-tb16k-nockpt

The packed run still spent 5.6 seconds on forward and backward, while forward alone took 1.34 seconds. The ratio suggested that gradient checkpointing was recomputing the forward pass during backward.

AsyncGRPOConfig defaults to gradient_checkpointing=True. The 1.5B model fit in the 141 GB H200 memory without checkpointing, so the next run set gradient_checkpointing=False.

Forward and backward time fell to 4.6 seconds, and MFU reached 23%. The queue decreased to 71 and rollout wait increased from 0.04 to 0.6 seconds, shifting the bottleneck to generation.

Weight synchronization, which took about 7.6 seconds every four steps, now represented 25% of wall-clock time. Backward remained about 2.5 times slower than forward, leaving approximately two seconds per step unexplained by ordinary model computation with frozen base weights.

Run 4: r1-dp3-tb16k-nockpt

A third vLLM replica was added because generation had become the slower stage. The adapter retry interval was reduced from 2 seconds to 0.5 seconds, and fsdp_reshard_after_forward was disabled to test whether FSDP2 re-gathering caused the extra backward time.

Weight synchronization decreased from 7.6 to 5.8 seconds, confirming that the shorter retry interval helped. Forward and backward remained at 4.6 seconds, ruling out resharding as the cause of the extra time. Generation increased only from 25k to 26k tokens per second.

The limiting factor was rollout/inflight, which remained at 128. The requests were divided approximately 44, 43, and 41 across the three replicas. max_inflight_tasks limits concurrency for the entire rollout worker, not for each replica. As a result, three GPUs received almost the same total workload as two.

Run 5: r1-dp3-inflight384

The final run changed only:

max_inflight_tasks=384
queue_maxsize=768

With 384 requests in flight, each replica received about 128 requests. The queue filled to approximately 690 of 768. Backpressure returned to 5 seconds, and rollout wait fell to 0.03 seconds, making training the bottleneck again.

Forward and backward took 4.6 seconds, weight synchronization added an amortized 1.5 seconds, and median step time was 4.8 seconds. Mean staleness increased from 1.5 to 2.0 versions because samples waited longer in the larger queue. It remained below max_staleness=4, while ratio stayed close to 1.000.

Results Summary

The final run was 3.9 times faster and trained 31% more samples, with a similar reward curve. Token-budget packing, disabling gradient checkpointing, and increasing the in-flight request limit produced the improvement.

Reproduction Commands

git clone https://github.com/AmineDiro/hfjobs-lora-buckets
cd hfjobs-lora-buckets
hf auth login

MAX_STEPS=20 RUN_TAG=smoke ./run_all.sh --wait
MAX_STEPS=500 ./run_all.sh --wait

TOKEN_BUDGET=16384 \
GRAD_ACCUM=6 \
GRADIENT_CHECKPOINTING=0 \
PROXY_LORA_RETRY_S=0.5 \
MAX_INFLIGHT=384 \
QUEUE_MAXSIZE=768 \
MAX_STEPS=500 \
./run_all.sh --wait

References

• John Schulman et al., LoRA Without Regret, Thinking Machines Lab, September 2025.
• TRL AsyncGRPOTrainer and its logged metrics.
• TRL PR #7017, which added PEFT and LoRA support for AsyncGRPOTrainer with adapter-only vLLM synchronization in TRL v1.14.
• Hugging Face Jobs, Storage Buckets, and hf-mount.
• Penghui Qi, Zichen Liu, Xiangxin Zhou, Tianyu Pang, Chao Du, Wee Sun Lee, and Min Lin, Defeating the Training-Inference Mismatch via FP16, arXiv:2510.26788, 2025.
• Dataset: sail/Sanity-Test-R1D-1.5B.
• Reproduction code: sail-sg/Precision-RL.

@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}
}