Hugging Face TRL: Co-located vLLM for Efficient GRPO Training

Hugging Face has introduced support for co-locating vLLM within the TRL library, enabling training and inference to share the same GPUs. This eliminates the "ping-pong" inefficiency of server-mode inference, where training GPUs sit idle during generation and inference GPUs sit idle during training, leading to higher throughput and reduced hardware requirements.

The Inefficiency of Server-Mode Inference

Prior to TRL v0.18.0, vLLM was supported only in server mode. In this configuration, vLLM runs as a separate process on dedicated GPUs, communicating with the training script via HTTP. This architecture creates significant bottlenecks in online learning algorithms like Group Relative Policy Optimization (GRPO), where generation occurs constantly:

  • GPU Underutilization: Training GPUs remain idle while the vLLM server generates completions.
  • Resource Waste: Dedicated GPUs must be provisioned solely for inference, increasing costs.
  • Communication Overhead: Reliance on REST API calls introduces networking latency.

Co-located Design and Implementation

Co-location allows vLLM to run alongside the training code within the same distributed process group. Instead of an external server, the trainer launches vLLM in-process using an external launcher.

Key Technical Capabilities

  • Unified Execution: Training and inference tasks take turns using the same GPU resources, reducing idle time.
  • Direct Memory Communication: By removing HTTP calls, vLLM communicates with the training loop via native Python calls.
  • Distributed Compatibility: The system is compatible with torchrun, Tensor Parallelism (TP), and Data Parallelism (DP), utilizing a Single Program, Multiple Data (SPMD) execution pattern.
  • Simplified Deployment: vLLM is controlled directly inside the training job, removing the need for separate server scripts.

Configuration

Users can enable this feature by setting vllm_mode="colocate" in the GRPOConfig:

training_args = GRPOConfig(
    ...,
    use_vllm=True,
    vllm_mode="colocate",
)

Performance Benchmarks

Experiments comparing server mode (where 1 of 8 GPUs is dedicated to vLLM) and co-locate mode (where all 8 GPUs are used for training) show consistent throughput gains. To ensure fairness, server mode throughput was normalized by a factor of 8/7.

Model-Specific Results

  • 1.5B Model: Achieved up to 1.43× speedup at the largest batch sizes. However, increasing Tensor Parallelism (TP) reduced performance due to communication overhead.
  • 7B Model: Achieved up to 1.35× speedup with varying batch sizes and up to 1.73× speedup when increasing TP, indicating that larger models benefit more from sharding.
  • 72B Model (Qwen2.5-Math-72B): The co-locate setup was approximately 1.26× faster than plain TRL, even while using 4 fewer GPUs.

Scaling to Large Models (72B+)

Training models as large as Qwen2.5-Math-72B requires advanced memory management to prevent contention between training and generation.

vLLM Sleep Mode

Hugging Face integrated vLLM’s sleep() API into the GRPO loop. Level 2 sleep is used to unload both model weights and the KV cache entirely from the GPU, maximizing free memory for training steps and avoiding memory contention.

DeepSpeed and Accelerate Optimizations

To maintain stability and efficiency, the following optimizations are employed:

  • DeepSpeed ZeRO Stage 3: Partitions model weights, gradients, and optimizer states across GPUs.
  • CPU Offloading: Optimizer states are moved to the CPU ("offload_optimizer": {"device": "cpu"}) to free GPU memory.
  • Memory Management: contiguous_gradients is enabled to reduce memory fragmentation, and overlap_comm is used to speed up training.
  • Accelerate Integration: Used for multi-GPU/multi-node launching and data parallelism.

Challenges and Lessons Learned

Implementation of co-located vLLM revealed several technical hurdles:

  • Determinism in vLLM ≥ 0.8.0: A bug in Tensor Parallelism with the external launcher was traced to a requirement for explicitly setting the random seed.
  • Buffer Restoration: A bug in Level 2 sleep (Issue #16564) required a fix to explicitly restore model buffers (e.g., BatchNorm running mean/var) after waking from sleep.
  • Shutdown Stability: A known issue exists where vLLM sleep causes a segmentation fault upon exiting the training process (Issue #16993), though this does not affect training results.

Conclusion on Model Quality

Benchmarks on Math500 and AIME24 confirm that co-location does not compromise model performance. Reward curves for co-located and plain setups are nearly identical, and the co-locate-trained models perform on par with plain-trained models, validating that efficiency gains do not come at the cost of accuracy.

Sources