Hugging Face TRL：共置 vLLM 以实现高效 GRPO 训练

Hugging Face 已在 TRL 库中引入对 vLLM 共置的支持，使训练和推理能够共享同一批 GPU。这消除了服务器模式推理的“ping-pong”低效——在生成期间训练 GPU 处于空闲状态，而在训练期间推理 GPU 也处于空闲状态，从而提升吞吐量并降低硬件需求。

服务器模式推理的低效

在 TRL v0.18.0 之前，vLLM 仅在服务器模式下受支持。在此配置中，vLLM 作为独立进程运行在专用 GPU 上，并通过 HTTP 与训练脚本通信。这种架构在诸如 Group Relative Policy Optimization（GRPO）等在线学习算法中会产生显著瓶颈，因为生成是持续进行的：

  • GPU 利用率低下：在 vLLM 服务器生成完成时，训练 GPU 处于空闲状态。
  • 资源浪费：必须专门为推理配备专用 GPU，导致成本上升。
  • 通信开销：依赖 REST API 调用会引入网络延迟。

共置设计与实现

共置使 vLLM 能够与训练代码在同一分布式进程组中并行运行。训练器不再使用外部服务器，而是通过外部启动器在进程内启动 vLLM。

关键技术能力

  • 统一执行：训练和推理任务轮流使用同一 GPU 资源，减少空闲时间。
  • 直接内存通信：通过移除 HTTP 调用，vLLM 通过原生 Python 调用与训练循环通信。
  • 分布式兼容性：系统兼容 torchrun、张量并行（Tensor Parallelism，TP）和数据并行（Data Parallelism，DP），采用单程序多数据（Single Program, Multiple Data，SPMD）执行模式。
  • 简化部署：vLLM 直接在训练作业内部受控，省去单独的服务器脚本。

配置

用户可以通过在 GRPOConfig 中设置 vllm_mode="colocate" 来启用此功能：

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

性能基准

对比服务器模式（8 块 GPU 中有 1 块专用于 vLLM）和共置模式（全部 8 块 GPU 用于训练）的实验显示出一致的吞吐量提升。为保证公平，服务器模式的吞吐量按 8/7 的系数进行了归一化。

模型特定结果

  • 1.5B 模型：在最大批量大小时实现最高 1.43 倍加速。然而，提升张量并行（TP）会因通信开销而降低性能。
  • 7B 模型：在不同批量大小下实现最高 1.35 倍加速，提升 TP 时最高可达 1.73 倍加速，表明更大的模型从分片中受益更多。
  • 72B 模型（Qwen2.5-Math-72B）：共置设置比普通 TRL 快约 1.26 倍，即使使用的 GPU 少了 4 块。

扩展到大模型（72B+）

训练如 Qwen2.5-Math-72B 这样的大模型需要高级内存管理，以防止训练与生成之间的资源争用。

vLLM 睡眠模式

Hugging Face 将 vLLM 的 sleep() API 集成到 GRPO 循环中。Level 2 睡眠用于将模型权重和 KV 缓存全部从 GPU 卸载，最大化训练步骤的空闲内存，避免内存争用。

DeepSpeed 与 Accelerate 优化

为保持稳定性和效率，采用了以下优化：

  • DeepSpeed ZeRO Stage 3：在 GPU 之间划分模型权重、梯度和优化器状态。
  • CPU 卸载：将优化器状态移至 CPU（"offload_optimizer": {"device": "cpu"}），释放 GPU 内存。
  • 内存管理：启用 contiguous_gradients 以减少内存碎片，并使用 overlap_comm 加速训练。
  • Accelerate 集成：用于多 GPU/多节点启动和数据并行。

挑战与经验教训

实现共置 vLLM 的过程中暴露出若干技术难点：

  • vLLM ≥ 0.8.0 的确定性：外部启动器下的张量并行出现 bug，原因是需要显式设置随机种子。
  • 缓冲区恢复：Level 2 睡眠中的 bug（Issue #16564）需要修复，以在唤醒后显式恢复模型缓冲区（例如 BatchNorm 的运行均值/方差）。
  • 关闭稳定性：已知问题是 vLLM 睡眠在退出训练进程时会导致段错误（Issue #16993），但不影响训练结果。

关于模型质量的结论

在 Math500 和 AIME24 上的基准测试证实，共置不会削弱模型性能。共置和普通设置的奖励曲线几乎相同，且共置训练的模型与普通训练的模型表现持平，验证了效率提升并未以准确性为代价。

Sources