Async GRPO with LoRA across HF Jobs: A Bucket, a Proxy, No NCCL

TRL‘s AsyncGRPOTrainer now supports LoRA and adapter-only vLLM sync, shipping in TRL v1.14 via PR #7017. Instead of sending full weights over NCCL, the trainer publishes a few-megabyte rank-1 adapter to a shared Storage Bucket, and vLLM replicas load it from disk. This lets training and inference run as separate Hugging Face Jobs on separate machines, with no network path between Jobs and no shared localhost.

The architecture is small: one trainer Job running AsyncGRPOTrainer with LoRA and FSDP, two or three vLLM Jobs serving the base model plus the latest adapter, a Storage Bucket mounted as FUSE in every Job at the same path, and a small Python asyncio proxy on the trainer’s localhost. The proxy adds the required HF auth header, routes each rollout to the replica most likely to hold its KV prefix, and broadcasts adapter loads, pause, resume, and unload to every replica. Nothing in TRL or vLLM changed for the filesystem part: the trainer writes an adapter under output_dir/.vllm_lora/trl-policy-v{N}, publishes it with an atomic rename, and posts the path to vLLM‘s /v1/load_lora_adapter endpoint.

Versioned adapter names are deliberate. vLLM keys its prefix cache by adapter name, so reusing one name could let rollouts get a prefix from one policy version and decode under another. With versioned names, a cached prefix can never match a newer policy. Adapter slots follow max_staleness: with max_staleness=4, vLLM must serve the current policy plus the four before it, and one extra slot during the load/unload swap, giving –max-loras 6. The recipe pins vllm/vllm-openai:v0.27.1 and enables VLLM_ALLOW_RUNTIME_LORA_UPDATING and VLLM_SERVER_DEV_MODE.

For validation, the project uses sail/Sanity-Test-R1D-1.5B, the dataset from the FP16 training-inference mismatch paper by Qi et al. (2025), with 1,460 MATH problems whose success rate is between 20% and 80%. Hyperparameters come from that paper’s LoRA scripts: Qwen/Qwen2.5-Math-1.5B, LoRA rank 1 with alpha 2, learning rate 4e-5, 8 samples per prompt, 128 completions per step, 3,000 max generated tokens, and a 4,096-token context.

Routing works by hashing prompts into 16-token blocks with chained blake2b hashes seeded by the adapter name, mirroring vLLM‘s KV block structure and causal attention. The router records which replica served each hash and which hashes followed it. Blocks shared by every replica, detected by fan-out rather than prefix length, are ignored because they don’t identify a prompt. A request goes to the replica with matching specific blocks if it is no more than 8 requests ahead of the least-loaded replica (affinity hit); otherwise it spills to the least-loaded replica, and new prompts are round-robin. Adapter broadcast is all-or-nothing: a “No adapter found” error is retried because bucket mounts lag, and any other failure triggers an unload rollback on replicas that already accepted the adapter.

Five runs show where the real bottleneck sits. The baseline took 3 h 27 min for 500 steps: step time was 22.9 s, forward+backward 21.9 s, MFU 3.9%, because one sequence per microbatch on an H200 was completely latency-bound. Setting token_budget=16384 and gradient_accumulation_steps=6 packed ~12.7 samples per row, cut microbatches from 64 to 6, and dropped step time to 5.9 s while generation rose from 4.6k to 25k tokens/s. Disabling gradient_checkpointing removed one forward pass, bringing forward+backward to 4.6 s but shifting the bottleneck to generation. A third replica barely helped because the client-side max_inflight_tasks=128 capped concurrency for the whole rollout worker, not per replica. Raising it to 384 with queue_maxsize=768 gave each of three replicas 128 requests in flight.

The final configuration is 3.9x faster: 53 min instead of 3 h 27 min, median step 4.8 s, MFU 23.5%, and 84,078 samples trained instead of 64,000, with essentially the same reward curve (0.145 to about 0.416–0.438) and ratio staying near 1.000 across all 126 syncs. All 252 adapter loads succeeded, with six on a second attempt and 246 on a third. Mean staleness rose from 1.5 to 2.0 versions, still below the max of 4. The run costs roughly $20/hour across the trainer and vLLM Jobs, and the full reproduction is available in the hfjobs-lora-buckets repository.

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

View Original