Asynchronous GRPO with LoRA adapters on Hugging Face Jobs: vLLM replicas, bucket sync and proxy routing

TRL (PR #7017, shipped in TRL v1.14) added LoRA adapter support to AsyncGRPOTrainer so the trainer can train a LoRA adapter and sync only that adapter to vLLM. Because a rank‑1 adapter is a few megabytes, adapters can be shared between separate Hugging Face Jobs via a Storage Bucket mount rather than moving gigabytes between machines. A small proxy running with the trainer adds the required authorization header, routes rollouts by KV‑prefix affinity to replicas that already hold the prefix, and broadcasts adapter loads to all replicas.

Why LoRA fits RL workloads

LoRA (low‑rank adapters) is especially suitable for policy‑gradient reinforcement learning: Thinking Machines’ “LoRA Without Regret” shows rank‑1 LoRA can match full fine‑tuning for such RL. From an information point of view an episode’s advantage provides only ~O(1) bits per episode, so a low‑capacity rank‑1 adapter can absorb the training signal.

There is an important systems consequence: a rank‑1 adapter for a 1.5B model is a few megabytes vs ~3 GB for the full model. vLLM can keep several adapters loaded simultaneously so old rollouts can finish with the policy they started with, while new rollouts use the latest policy.

Architecture on Hugging Face Jobs

Hugging Face Jobs run a single container on a single VM; there is no shared local disk or NCCL across Jobs. The solution is to mount a Storage Bucket (via hf-mount) as a POSIX filesystem at the same absolute path in every Job. The trainer saves adapter directories like <output_dir>/.vllm_lora/trl-policy-v{N}, renames atomically, and POSTs the path to vLLM’s /v1/load_lora_adapter. vLLM reads the files from the shared path.

The deployment used in the experiments contains:

• a trainer Job running AsyncGRPOTrainer with LoRA (and FSDP in some runs),
• two vLLM Jobs, each serving the base model and whatever adapters the trainer last published,
• a Storage Bucket mounted at the same path in all three Jobs,
• a small proxy on the trainer Job that adds Authorization headers, routes completion requests to the replica most likely to already have the prefix in cache, and broadcasts state‑changing calls (adapter load/unload, pause/resume) to every replica.

No code changes were required in TRL or vLLM to use this filesystem‑based adapter sync path.

vLLM configuration and adapter slots

The experiments pinned vLLM to v0.27.1 (the runtime LoRA flags and endpoints used are that version’s). Because AsyncGRPOTrainer configured max_staleness=4, each vLLM replica must serve the current policy plus the four previous versions simultaneously. That requires max_loras=6 on the server (4 + 2 for the swap transient). If fewer slots are configured, vLLM would evict policies still in flight at every sync.

Versioned adapter names were used (trl-policy-v{N}) to avoid cache mismatches: the KV cache is keyed by adapter name, and reusing a single name could allow a cached prefix computed under an older policy to be reused against a newer decode, corrupting the policy matching ratio.

Dataset, model and hyperparameters

• TRL v1.14 (PR #7017) provided AsyncGRPOTrainer LoRA support.
• vLLM image: vllm/vllm-openai:v0.27.1.
• Model: Qwen/Qwen2.5-Math-1.5B.
• Dataset: sail/Sanity-Test-R1D-1.5B (from Qi et al., 2025). The dataset contains 1,460 MATH problems filtered for 20–80% baseline success rate.
• LoRA recipe: rank r=1, lora_alpha=2, learning rate 4e-5, 8 samples per prompt, 128 completions per step, max 3,000 generated tokens, context length 4,096.

The proxy and routing by KV prefix affinity

The proxy is necessary because:

1. Exposed HF Job ports require an Authorization: Bearer <HF token> header on every request; the proxy adds it so TRL doesn’t need to manage HF Job URLs.
2. We want multiple GPUs producing generations. On single‑node vLLM DP modes, adapter-only sync is unsupported because /v1/load_lora_adapter reaches only the answering DP rank. With separate Jobs each replica is a full machine, so the proxy must fan‑out adapter loads to all replicas.

Functionally the proxy acts as a single server to TRL at http://localhost:8000 while doing two things: route per‑completion requests to the best replica so the G (here G=8) rollout requests reuse a prefilled KV cache, and broadcast state‑changing requests to all replicas.

Routing algorithm (overview):

• Split the prompt into 16‑token blocks (vLLM’s block size) and compute chained hashes seeded by the adapter name. Because hashes are chained, h_k encodes blocks 1..k and thus reflects causal attention.
• The router records which replicas served which block hashes and which successors followed each hash (cap successors at two).
• A block is considered common if every replica has served it or it has multiple successors; common blocks are ignored in affinity checks because they identify no single replica.
• For each replica compute the number of leading specific (non-common) blocks matched. Then:
  • If one replica has specific blocks and it is not more than cfg.imbalance (8) requests ahead of the least‑loaded replica, route there (affinity hit).
  • If it is more than that, send to the least‑loaded replica (spill).
  • If no replica has specific blocks, send to the least‑loaded (unmatched, round‑robin ties).

This logic maximizes reuse of prefill work: the first request that reaches a replica computes the prefill (attentions / KV cache), subsequent rollout requests with the same prefix reuse it.

Broadcasting adapter loads and transactional semantics

Adapter loads are broadcasted to all replicas and handled as all‑or‑nothing: the proxy retries replicas that report "No adapter found for <path>" (delay in bucket mount visibility), and if any other replica fails the proxy unloads the adapter from replicas that accepted it so a policy name never exists only on a subset.

From TRL’s perspective the proxy reports itself as a single data_parallel_size=1 server so AsyncGRPOTrainer selects adapter‑only sync mode.

Full run results and metric‑driven tuning

All timing metrics reported come from the trainer logs and the trackio dashboard. Experiment setup: trainer Job h200x2, each vLLM replica on one h200 Job. Running all three Jobs costs ~20 USD/hour.

Key aggregate results (500 steps, same recipe):

• Run 1 (r1-dp2): baseline with default batching — p50 step 22.9 s, forward+backward 21.9 s, MFU 3.9%, 500 steps took 3 h 27 min. Rollout queue stayed nearly full (476/512), so trainer was the bottleneck.
• Run 2 (r1-dp2-tb16k): token_budget=16384, gradient_accumulation_steps=6 — samples_per_row rose from 1.0 to ~12.7, microbatches fell 64→6, step time fell 23 s → 5.9 s, MFU rose 3.9% → 19%, generation throughput rose 4.6k → 25k tok/s because the queue was no longer constantly full.
• Run 3 (r1-dp2-tb16k-nockpt): gradient_checkpointing=False — forward+backward dropped to 4.6 s, MFU 23%, the bottleneck shifted to generation (rollout wait increased).
• Run 4 (r1-dp3-tb16k-nockpt): added a third replica — weight sync improved (7.6 s → 5.8 s) but generated token/s barely changed (25k → 26k). Diagnosis: client-side max_inflight_tasks=128 limited concurrency; splitting 128 across three GPUs gave similar per‑GPU concurrency as splitting across two.
• Run 5 (r1-dp3-inflight384): raised max_inflight_tasks=384 and queue_maxsize=768 — each replica got ~128 in flight, queue ~690/768, step p50 4.8 s, forward+backward 4.6 s, mean staleness 2.0, reward curves similar. Wall-clock for 500 steps fell from 3 h 27 min to 53 min (3.9× speedup). Samples trained increased from 64,000 → 84,078.

Routing statistics after the full run (64,728 rollouts): routed [31928, 32800], affinity 54,712, spilled 820, unmatched 9,196. Affinity hits were 84.5%, unmatched 14.2% (close to theoretical minimum of 12.5% for G=8), spills 1.3%.

All 252 adapter loads (126 syncs × 2 replicas) succeeded; six loads succeeded on the second attempt and 246 on the third, consistent with transient bucket mount delays.

Conclusion and reproducibility

Adapter‑only sync for AsyncGRPOTrainer makes it practical to separate training and inference across Hugging Face Jobs by using Storage Buckets and a small proxy. Detailed AsyncGRPO metrics allow diagnosing whether the bottleneck is training or generation and guide effective changes: token packing, disabling gradient checkpointing when feasible, increasing the client-side in‑flight cap, and lowering adapter retry intervals improved throughput and GPU utilization, yielding a 3.9× faster run for the same 500 steps and comparable reward curves.

To reproduce: clone https://github.com/AmineDiro/hfjobs-lora-buckets, authenticate with HF, and use the provided run_all.sh with the environment variables shown in the repository. A short smoke run (~20 steps) takes ~15 minutes and cancels servers after completion; the full 500-step optimized run is reported at ~52–55 minutes in the experiments above.

References

• TRL AsyncGRPOTrainer and logged metrics; PR #7017 (LoRA support, adapter-only vLLM sync, TRL v1.14).
• vLLM v0.27.1 runtime LoRA flags and endpoints.
• Thinking Machines Lab, "LoRA Without Regret" (September 2025).
• Penghui Qi et al., "Defeating the Training-Inference Mismatch via FP16" (arXiv:2510.26788, 2025) — source of the Sanity dataset and LoRA recipe.
• Reproduction repo: aminediro/hfjobs-lora-buckets (run_all.sh and helper scripts).