Async GRPO with LoRA across HF Jobs: a bucket, a proxy, and no NCCL

What it claims

TRL v1.14's AsyncGRPOTrainer (PR #7017) can train a LoRA adapter and sync only that adapter to vLLM. Because a rank-1 adapter for a 1.5B model is a few MB rather than ~3 GB, it can travel through a Storage Bucket mounted as a FUSE filesystem in every Hugging Face Job. This lets the trainer and the vLLM replicas run on separate machines with no NCCL and no network path between them. The authors state that TRL and vLLM needed no changes for this. All the glue lives in bucket mounts and a small proxy.

How it works

• Adapter publish. Every weight_sync_steps the trainer writes <output_dir>/.vllm_lora/trl-policy-v{N} and publishes it with an atomic rename. It then POSTs the path to vLLM's /v1/load_lora_adapter, which loads the adapter from disk.
• Versioned names are deliberate. vLLM keys its prefix cache by adapter name. Reusing a single name would let a rollout take its prefill from one policy and its decode from the next, which would show up as ratio drifting from 1.
• Slot sizing. You need --max-loras = max_staleness + 2, which is 6 for staleness 4: the current policy, the 4 older ones, and one extra slot during the load-before-unload swap. With 5 slots, vLLM silently evicts a policy that still has rollouts in flight.
• Sync mode selection. TRL checks /server_info. DoRA, modules_to_save, or a rank above --max-lora-rank fall back to merged-weight sync with a warning. Look for Adapter-only vLLM sync enabled in the log.
• Why a proxy, not --data-parallel-size. TRL refuses adapter-only sync with vLLM DP>1, because an adapter load reaches only one DP rank. The proxy runs on the trainer Job at 127.0.0.1:8000, adds the Authorization: Bearer header, and presents itself as one DP=1 server.
• Adapter broadcast is all-or-nothing. A No adapter found error means that replica's mount is lagging, so the proxy retries just that replica. Any other error unloads the adapter from the replicas that accepted it. /pause, /resume and unload are broadcast the same way.
• KV-prefix routing. The proxy computes chained blake2b hashes over 16-token blocks, seeded with the adapter name. A block counts as "common" if every replica has served it or it has more than one successor. Common blocks, such as the shared chat template, are ignored when matching.

flowchart TD
  R[Rollout request] --> H[Hash 16-token blocks, seed = adapter name]
  H --> S{Replica with prompt-specific blocks?}
  S -- no --> U[Unmatched: least-loaded, round-robin]
  S -- yes --> L{≤8 requests ahead of least-loaded?}
  L -- yes --> A[Affinity hit: reuse prefill]
  L -- no --> P[Spill: least-loaded]

Evidence

The setup is Qwen2.5-Math-1.5B with LoRA r=1, alpha 2, all-linear, lr 4e-5, 8 rollouts/prompt, 128 completions/step, 500 steps. The dataset is sail/Sanity-Test-R1D-1.5B (1,460 MATH problems). The trainer runs on an h200x2 Job and each vLLM replica on an h200 Job, at about $20/hour for all three.

• Correctness: All 252 adapter loads succeeded. ratio stayed at 1.000 (range 0.9993–1.0004) across 126 syncs. Reward went from 0.145 to 0.438.
• Routing over 64,728 rollouts: 84.5% affinity hits, 14.2% unmatched (the theoretical floor is 12.5%), 1.3% spills.
• Bottleneck runs: 1. Baseline: 22.9 s/step, 3.9% MFU, 64 microbatches of one sequence each. The run was trainer-bound and took 3h27m. 2. token_budget=16384, gradient_accumulation_steps=6: ~12.7 samples/row, fwd+bwd 21.9 s → 5.6 s, MFU 19%. Generation rose from 4.6k to 25k tok/s with no vLLM change. 3. gradient_checkpointing=False (the default is True): fwd+bwd 4.6 s, MFU 23%. The run became generation-bound, and the 7.6 s weight sync now cost 25% of wall-clock. 4. Third replica: no throughput gain. max_inflight_tasks=128 is a global cap, not per-replica. Cutting the proxy retry interval from 2 s to 0.5 s brought sync from 7.6 s to 5.8 s. 5. max_inflight_tasks=384, queue_maxsize=768: 4.8 s median step and 53 min total. That is 3.9× faster on 31% more samples, with the same reward curve. Mean staleness went from 1.5 to 2.0.
• Diagnosing the bottleneck: A full queue with near-zero perf/rollout_wait_s and high rollout/backpressure_s means the trainer is slow. An emptying queue with rising rollout wait means generation is slow.

How to use it

Pin vllm/vllm-openai:v0.27.1; the flags and endpoints are version-specific. Each replica needs VLLM_ALLOW_RUNTIME_LORA_UPDATING=1, VLLM_SERVER_DEV_MODE=1, --enable-lora --max-lora-rank 1 --max-loras 6, and the same bucket mounted at the same path. Set the trainer's output_dir on the bucket so checkpoints survive preemption.

git clone https://github.com/AmineDiro/hfjobs-lora-buckets && cd hfjobs-lora-buckets
hf auth login
MAX_STEPS=20 RUN_TAG=smoke ./run_all.sh --wait        # ~15 min, three Jobs, cancels the servers when done
MAX_STEPS=500 ./run_all.sh --wait                     # run 1: the reference batch shape, ~3.5 h
TOKEN_BUDGET=16384 GRAD_ACCUM=6 GRADIENT_CHECKPOINTING=0 PROXY_LORA_RETRY_S=0.5 \
  MAX_INFLIGHT=384 QUEUE_MAXSIZE=768 MAX_STEPS=500 ./run_all.sh --wait   # run 5: same recipe, ~55 min

Limits

• Everything was tested at one scale: a 1.5B model, rank 1, and up to 3 replicas. The single-threaded asyncio proxy was fine at ≤128 requests in flight, and the authors expect a higher-traffic router to need a faster language.
• AsyncGRPOTrainer is imported from trl.experimental.
• Backward still runs about 2.5× forward, roughly 2 s/step unexplained. FSDP resharding was ruled out as the cause.
• Bucket mounts lag. Most adapter loads needed 3 attempts (246 of 252). A separate repro, hf-mount-repro, documents a 30-second negative-cache stall.
• Runs 2–4 were stopped early, so only runs 1 and 5 cover the full 500 steps.