Loading post…
Hugging FaceAmine Dirhoussi, Quentin Gallouédec, Kashif Rasul, Sergio PaniegoExplainer
Async GRPO with LoRA across HF Jobs: a bucket, a proxy, and no NCCL
Summary
An AsyncGRPO setup uses LoRA adapter-only synchronization to separate vLLM Jobs, rather than moving full model weights over NCCL. A Storage Bucket mounted at the same path lets the trainer publish versioned adapters that replicas load, while a proxy adds authentication, routes rollouts toward KV-cached prefixes, and broadcasts adapter updates. Versioned adapter names prevent cached prefixes from mixing policy versions, and max_staleness=4 requires six LoRA slots during replacement. Five runs identify client-side concurrency, checkpointing, and packing as important throughput factors: raising in-flight requests to 384, disabling checkpointing, and packing improved median step time from 22.9 to 4.8 seconds. At 500 steps, the final run took 53 minutes instead of 3 hours 27 minutes, trained 31% more samples, and had basically the same reward curve.
Context
Hugging Face Jobs run one container on one VM, so the trainer and vLLM fleet cannot share a node, local disk, localhost, or an NCCL group. The design uses three Jobs, a Storage Bucket mounted at the same absolute path, and a proxy connecting the trainer to separate vLLM replicas.
Approach / What changed
The trainer saves each LoRA adapter under a versioned directory in the shared bucket and asks vLLM to load it by path. Two vLLM replicas serve the base model with runtime LoRA loading; the proxy adds the authentication header, routes rollout requests, and broadcasts adapter loads. The setup keeps max_staleness + 1 versions registered and uses an additional slot during swaps, resulting in --max-loras 6 for max_staleness=4.
Takeaways
A rank-1 adapter for a 1.5B model is a few megabytes, compared with around 3 GB for the full model, making bucket-based synchronization practical across separate Jobs.
Versioned adapter names keep vLLM prefix-cache entries tied to one policy version; reusing one name could mix a cached prefix from an old adapter with decoding from a new one.
Increasing max_inflight_tasks to 384, using queue_maxsize=768, disabling checkpointing, and packing improved the 500-step run from 3 hours 27 minutes to 53 minutes.
FAQ
Why does the setup use versioned adapter names instead of replacing one adapter in place?
vLLM keys its prefix cache by adapter name. Reusing one name could let cached blocks from previous weights match after a swap, mixing policy versions; versioned names prevent that.
Why are six LoRA slots needed when max_staleness is 4?
The server must retain the current policy plus four earlier versions so in-flight rollouts can finish. Loading the new version before unloading the oldest requires one additional slot during the swap, making six necessary.
What changes produced the large reduction in training time?
The final run used packing, disabled gradient checkpointing, and raised the in-flight request limit to 384 with a queue size of 768. Median step time fell from 22.9 seconds to 4.8 seconds, while the reward curve remained basically the same.
Source and freshness
- Source
- huggingface.co
- Published
- Added to Yomu
Summaries are generated from the original by a language model. Check the original before relying on it.
Keep reading
Related reading
The Good Stuff
Paid picks we use ourselves. Links carry a referral code — a subscription through one earns us a commission at no extra cost to you.
The best engineering reads, once a week
One useful email. Unsubscribe anytime.