使用 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 高效管理複雜的資料流和計算，包括將基於規則的獎勵模型分佈到多個節點上。

vLLM with Ray Executor and AutoTP 在加速推理方面發揮著核心作用。憑藉對 Ray Executors 的內建支援以及與 HuggingFace Transformers 的整合，它透過 AutoTP 實現高效的權重更新，從而實現高吞吐量和記憶體高效的 LLM 生成。

ZeRO-3 with HuggingFace Transformers，來自 DeepSpeed 的一種記憶體最佳化方法，使 OpenRLHF 能夠訓練大型模型，而無需像 Megatron 那樣重量級的框架。與 HuggingFace 的無縫整合使得預訓練模型的載入和微調變得簡單。

Ray, vLLM, ZeRO-3 和 HuggingFace Transformers 共同打造了一個前沿且精簡的 RLHF 訓練加速解決方案。該架構也影響了其他框架，例如 veRL，它們採用了類似的範例來實現可擴充套件且高效的 RLHF 訓練。OpenRLHF 也是首個基於 Ray, vLLM 和 ZeRO-3 開發的開源 RLHF 框架，已被 Google, 字節跳動, 阿里巴巴, 美團, 伯克利 Starling 團隊等使用。

Ray and vLLM in OpenRLHF

如上圖所示，OpenRLHF 使用 Ray 的 Placement Group API 來靈活排程 RLHF 流水線的元件，包括 vLLM 引擎, Actor, Critic, Reference 和 Reward 模型。儘管單獨表示，這些元件可以共同安置在共享的 Ray placement groups 中，以最大限度地提高資源效率。例如，所有模組可以在混合引擎配置中在同一 GPU 組內執行，或者特定元件——如 Actor 和 Critic——可以組合在一起。所有模組都由一箇中央 Ray Actor 編排，該 Actor 管理整個訓練生命週期。Actor 和 vLLM 引擎之間的權重同步透過高效能通訊方法處理，例如 NVIDIA Collective Communications Library (NCCL) 或混合引擎設定中的 CUDA Inter-Process Communication (IPC) 記憶體傳輸。

使用 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, 建立 placement group 來管理資源, 以及定義訓練 actors 和推理引擎。訓練 actors 管理模型的初始化和權重更新，而推理引擎透過 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 和 Hybrid Engine 元件做出的重大貢獻。由 Jian Hu 負責開發工作。