使用 vLLM 加速 RLHF，来自 OpenRLHF 的最佳实践
随着训练具备推理能力的大语言模型（LLM）需求的增长，基于人类反馈的强化学习（RLHF）已成为核心技术。然而，传统的 RLHF 流水线——尤其是使用近端策略优化（PPO）的流水线——往往受到巨大的计算开销阻碍。对于擅长复杂推理任务的模型（如 OpenAI-o1 和 DeepSeek-R1）而言，这一挑战尤为突出。在这些模型中，生成长思维链（CoT）输出的时间可能占总训练时间的 90%。这些模型必须产生详细的、分步骤的推理过程，可能跨越数千个 token，使得推理阶段比训练阶段更加耗时。作为领先的推理框架，vLLM 为生成 RLHF 样本和更新模型权重提供了一个用户友好的接口。
OpenRLHF 的设计
为了在 RLHF 框架的性能和易用性之间取得平衡，OpenRLHF 被设计为一个高性能且用户友好的解决方案，集成了 Ray、vLLM、零冗余优化器（ZeRO-3）和自动张量并行（AutoTP）等关键技术。
Ray 是 OpenRLHF 分布式架构的支柱。凭借强大的调度和编排功能，Ray 可以高效管理复杂的数据流和计算，包括跨多个节点分发基于规则的奖励模型。
带有 Ray Executor 和 AutoTP 的 vLLM 在加速推理方面发挥着核心作用。通过内置对 Ray Executors 的支持以及与 HuggingFace Transformers 的集成，它能够通过 AutoTP 实现高效的权重更新，从而带来高吞吐量和内存高效的 LLM 生成。
带有 HuggingFace Transformers 的 ZeRO-3 是一种来自 DeepSpeed 的内存优化方法，使 OpenRLHF 能够训练大模型，而无需使用 Megatron 等重量级框架。这种与 HuggingFace 的无缝集成，允许轻松加载和微调预训练模型。
Ray、vLLM、ZeRO-3 和 HuggingFace Transformers 共同构成了一个尖端且精简的 RLHF 训练加速解决方案。该架构还影响了诸如 veRL 等其他框架，它们采用了类似的范式来实现可扩展且高效的 RLHF 训练。OpenRLHF 也是第一个基于 Ray、vLLM 和 ZeRO-3 开发的开源 RLHF 框架，已被 Google、字节跳动、阿里巴巴、美团、Berkeley Starling 团队等使用。
使用 vLLM Ray Executor 实现 RLHF 加速
OpenRLHF 和 vLLM 提供了一套简洁高效的 API，简化了 RLHF 流水线内的交互。通过实现自定义的 WorkerExtension 类，用户可以处理训练和推理组件之间的权重同步。环境变量 VLLM_RAY_PER_WORKER_GPUS 和 VLLM_RAY_BUNDLE_INDICES 允许对每个 Worker 进行细粒度的 GPU 资源分配，从而实现多个组件共享一个 GPU 组的混合引擎配置。
# rlhf_utils.py
class ColocateWorkerExtension:
"""
Extension class for vLLM workers to handle weight synchronization.
This class ensures compatibility with both vLLM V0 and V1.
"""
def report_device_id(self) -> str:
"""Report the unique device ID for this worker"""
from vllm.platforms import current_platform
self.device_uuid = current_platform.get_device_uuid(self.device.index)
return self.device_uuid
def update_weights_from_ipc_handles(self, ipc_handles):
"""Update model weights using IPC handles"""
handles = ipc_handles[self.device_uuid]
device_id = self.device.index
weights = []
for name, handle in handles.items():
func, args = handle
list_args = list(args)
list_args[6] = device_id # Update device ID for current process
tensor = func(*list_args)
weights.append((name, tensor))
self.model_runner.model.load_weights(weights=weights)
torch.cuda.synchronize()
# main.py
class MyLLM(LLM):
"""
Custom LLM class to handle GPU resource allocation and bundle indices.
This ensures proper GPU utilization and placement group management.
"""
def __init__(self, *args, bundle_indices: list, **kwargs):
# Prevent Ray from manipulating CUDA_VISIBLE_DEVICES at the top level
os.environ.pop("CUDA_VISIBLE_DEVICES", None)
# Configure GPU utilization per worker
os.environ["VLLM_RAY_PER_WORKER_GPUS"] = "0.4"
os.environ["VLLM_RAY_BUNDLE_INDICES"] = ",".join(map(str, bundle_indices))
super().__init__(*args, **kwargs)
# Create Ray's placement group for GPU allocation
pg = placement_group([{"GPU": 1, "CPU": 0}] * 4)
ray.get(pg.ready())
# Create inference engines
inference_engines = []
for bundle_indices in [[0, 1], [2, 3]]:
llm = ray.remote(
num_gpus=0,
scheduling_strategy=PlacementGroupSchedulingStrategy(
placement_group=pg
)
)(MyLLM).remote(
model="facebook/opt-125m",
tensor_parallel_size=2,
distributed_executor_backend="ray",
gpu_memory_utilization=0.4,
worker_extension_cls="rlhf_utils.ColocateWorkerExtension",
bundle_indices=bundle_indices
)
inference_engines.append(llm)完整的 RLHF 示例详细介绍了如何初始化指定 GPU 数量的 Ray，创建放置组来管理资源，并定义训练 Actor 和推理引擎。训练 Actor 管理模型初始化和权重更新，而推理引擎则通过 vLLM 提供模型服务。权重同步使用 CUDA IPC 或 NCCL 执行，确保在整个 RLHF 流水线中保持一致性和高效性。
致谢
我们要向 vLLM 的贡献者们表示诚挚的感谢，包括 Kaichao You、Cody Yu、Rui Qiao 等人，没有他们的帮助，OpenRLHF 与 vLLM 的集成将无法实现。来自 vLLM 团队的 Kaichao You 领导了这次 RLHF 集成工作。
OpenRLHF 项目是第一个基于 Ray 和 vLLM 的开源 RLHF 框架。我们要感谢 Jian Hu、Songlin Jiang、Zilin Zhu、Xibin Wu 等人，他们对 OpenRLHF 项目的 Ray、vLLM Wrapper 和混合引擎组件做出了重大贡献。Jian Hu 领导了该开发工作。