Async GRPO with LoRA across HF Jobs: a bucket, a proxy, and no NCCL
huggingface ·
TRL's AsyncGRPOTrainer has gained LoRA support (PR #7017, shipping in TRL v1.14), allowing it to train an adapter rather than the full model and sync only that adapter to vLLM. The blog describes a project running this setup across separate Hugging Face Jobs, since a Job is one container on one VM and cannot span nodes to hold both a trainer and vLLM replicas. Because Jobs cannot communicate over NCCL and have no shared filesystem, the adapter—only a few megabytes for a rank-1 LoRA versus roughly 3 GB for the full 1.5B model—is passed through a Storage Bucket mounted at the same path in every Job using hf-mount. The setup consists of a trainer Job running AsyncGRPOTrainer with LoRA and FSDP, two vLLM Jobs serving the base model plus the latest published adapter, the shared bucket, and a proxy. The proxy adds the Authorization header required by exposed Job ports, routes each rollout to the replica most likely to hold its KV cache prefix, and broadcasts adapter updates to all replicas. Adapters are published with versioned names under .vllm_lora/trl-policy-v{N}, because vLLM keys its prefix cache by adapter name and a single reused name could let a rollout mix prefix and decode from different policy versions. Training uses the sail/Sanity-Test-R1D-1.5B dataset, containing 1,460 MATH questions with 20–80% success rates for DeepSeek-R1-Distill-Qwen-1.5B, chosen as a fast end-to-end validation set. Five successive runs of the same recipe reduced the time for 500 steps from 3 hours 27 minutes to 53 minutes, with AsyncGRPO metrics used to locate bottlenecks.
huggingface