Async GRPO with LoRA across Hugging Face Jobs

Hugging Face has implemented LoRA support in AsyncGRPOTrainer (TRL v1.14), allowing the trainer to sync only a small LoRA adapter to vLLM rather than the full model weights. By leveraging Hugging Face Storage Buckets as a shared filesystem and a custom proxy for request routing, this architecture enables training and inference to run on separate machines (Hugging Face Jobs) without requiring NCCL or a shared local disk.

Architecture: Distributed Sync via Storage Buckets

AsyncGRPOTrainer now supports an adapter-only sync path that eliminates the need for tensors to be sent directly to vLLM. Instead, the trainer saves the LoRA adapter to a specific directory in a Storage Bucket, performs an atomic rename, and notifies vLLM via the /v1/load_lora_adapter endpoint.

Because Hugging Face Jobs can mount Storage Buckets as FUSE filesystems using hf-mount, the trainer and vLLM replicas can share the same absolute path across different VMs. This removes the requirement for the trainer and inference servers to share a physical node or a dense cluster network.

The Job Layout

The system consists of three primary components:

  • Trainer Job: Runs AsyncGRPOTrainer with LoRA and FSDP.
  • vLLM Jobs: Multiple replicas serving the base model and loading the latest adapters from the bucket.
  • Proxy Server: A small asyncio-based proxy running on the trainer Job that handles authentication headers and routes requests to vLLM replicas.

The Proxy: KV-Prefix Routing and Broadcasting

To maximize efficiency across multiple vLLM replicas, a custom proxy is used to manage request distribution and state synchronization.

Routing by KV-Prefix

To avoid redundant prefill computations, the proxy routes requests based on the KV cache prefix. It splits prompts into 16-token blocks and computes chained hashes seeded with the adapter name.

The router tracks which replica has served which block hash. If a request's prompt matches a prefix already cached on a specific replica (and that replica is not overloaded), the request is routed there (an "affinity hit"). This prevents the system from re-computing the prefill for the same prompt across multiple rollouts, which is critical for GRPO where multiple completions are generated for a single prompt.

State Broadcasting

Since each vLLM replica is a separate Job, the proxy ensures consistency by broadcasting state-changing requests—such as adapter loads, pauses, and resumes—to all replicas. This ensures that a specific policy version name refers to the same weights across the entire fleet.

Performance Optimization and Bottleneck Analysis

Using the sail/Sanity-Test-R1D-1.5B dataset and a Qwen/Qwen2.5-Math-1.5B model, Hugging Face conducted five experimental runs to optimize the pipeline. The results demonstrate that the bottleneck in Async RL can shift between training and generation.

Key Optimizations

  1. Token-Budget Batching: Moving from a per-device train batch size of 1 to token-budget batching (e.g., token_budget=16384) increased MFU from 3.9% to 19% by packing multiple sequences into each row, reducing the number of microbatches.
  2. Disabling Gradient Checkpointing: For smaller models (1.5B), disabling gradient checkpointing reduced the forward+backward time by eliminating redundant forward passes, shifting the bottleneck from the trainer to the generation replicas.
  3. Increasing In-Flight Requests: Raising max_inflight_tasks (e.g., to 384) allowed the system to fully utilize multiple vLLM replicas, preventing the client-side concurrency limit from throttling throughput.

Final Results

By combining these optimizations, the total time to complete 500 steps was reduced from 3 hours 27 minutes to 53 minutes (a 3.9x speedup).

Metric Run 1 (Baseline) Run 5 (Optimized)
Wall Clock Time 3 h 27 min 53 min
Median Step Time 22.9 s 4.8 s
MFU (Fwd/Bwd) 3.9% 23.5%
Samples Trained 64,000 84,078
Mean Staleness 1.5 versions 2.0 versions

Technical Implementation Details

  • vLLM Version: Pinned to v0.27.1 for compatibility with runtime LoRA endpoints.
  • Adapter Slots: To support max_staleness=4, vLLM is configured with --max-loras 6 to ensure the current policy and previous versions remain loaded during swaps.
  • Consistency: Versioned adapter names are used to prevent the KV cache from incorrectly matching a prefix generated by an older policy version.

Sources