【vLLM 学习】Rlhf Colocate

vLLM 是一款专为大语言模型推理加速而设计的框架，实现了 KV 缓存内存几乎零浪费，解决了内存管理瓶颈问题。

更多 vLLM 中文文档及教程可访问 →vllm.hyper.ai/

*在线运行 vLLM 入门教程：零基础分步指南

源码examples/offline_inference/rlhf_colocate.py

# SPDX-License-Identifier: Apache-2.0

"""
一个简单的演示，展示如何将 vLLM 工作进程与训练执行器（training actors）
协同部署在同一 GPU上，适用于类 RLHF 应用。

关键要点：
- 通过正确设置 VLLM_RAY_PER_WORKER_GPUS 和 VLLM_RAY_BUNDLE_INDICES，
  使用 Ray 控制 vLLM 工作进程的部署位置
- 使用 CUDA-IPC 传递张量，因为在同一 GPU 上存在多个进程时 NCCL 无法正常工作
"""
import os

import ray
import torch
from ray.util.placement_group import placement_group
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy

from vllm import LLM


class MyLLM(LLM):

    def __init__(self, *args, bundle_indices: list, **kwargs):

        # 临时方案使脚本能运行
        # 阻止Ray在顶层操作CUDA_VISIBLE_DEVICES
        os.environ.pop("CUDA_VISIBLE_DEVICES", None)

        # 每个工作进程将使用 0.4 个 GPU，这样我们可以在同一 GPU 上调度 2 个实例
        os.environ["VLLM_RAY_PER_WORKER_GPUS"] = "0.4"
        os.environ["VLLM_RAY_BUNDLE_INDICES"] = ",".join(
            map(str, bundle_indices))
        print(f"creating LLM with bundle_indices={bundle_indices}")
        super().__init__(*args, **kwargs)


class RayTrainingActor:

    def __init__(self):

        # ray 将 CUDA_VISIBLE_DEVICES 设置为分配的 GPU
        from transformers import AutoModelForCausalLM
        self.model = AutoModelForCausalLM.from_pretrained("facebook/opt-125m")
        self.model.to("cuda:0")
        for name, p in self.model.named_parameters():
            p.data.zero_()
        torch.cuda.synchronize()
        # get_device_uuid 的参数是
        # 可见设备中 GPU 的索引
        from vllm.platforms import current_platform
        self.device_uuid = current_platform.get_device_uuid(0)

    def report_device_id(self) -> str:
        return self.device_uuid

    def get_weight_ipc_handles(self):
        from torch.multiprocessing.reductions import reduce_tensor
        data = {}
        for name, p in self.model.named_parameters():

            # 训练执行器（training actor）可能只拥有部分权重，
            # 需要从所有执行器进行 all-gather 操作获取完整权重。
            # 出于演示目的，此处我们假设所有训练执行器都拥有完整权重。
            data[name] = reduce_tensor(p.detach())
        return {self.device_uuid: data}


# ray 管理4 GPU
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3"
ray.init()

# 我们需要将 vLLM 实例和训练执行器（training actor）协同部署在同一组 GPU 上
# 具体部署方案如下：
# GPU 0 和 1：训练执行器 0、1 和 vLLM 实例 0（TP=2）
# GPU 2 和 3：训练执行器 2、3 和 vLLM 实例 1（TP=2）

pg = placement_group([{"GPU": 1, "CPU": 0}] * 4)
ray.get(pg.ready())
print(f"placement group has bundles {pg.bundle_specs=}")

training_actors = []
training_actor_device_ids = []
inference_engines = []
inference_engine_device_ids = []

for bundle_index in [0, 1, 2, 3]:
    training_actor = ray.remote(
        num_cpus=0,
        num_gpus=0.4,
        scheduling_strategy=PlacementGroupSchedulingStrategy(
            placement_group=pg,
            placement_group_capture_child_tasks=True,
            placement_group_bundle_index=bundle_index,
        ),
    )(RayTrainingActor).remote()
    training_actors.append(training_actor)

for bundle_index, training_actor in enumerate(training_actors):
    device_id = ray.get(training_actor.report_device_id.remote())
    print(f"training actor {bundle_index} is on {device_id}")
    training_actor_device_ids.append(device_id)

for (i, bundle_indices) in enumerate([[0, 1], [2, 3]]):

    # and cause unexpected behaviors.
    # 重要:创建 vLLM 实例时，我们需要
    # 确保目标 GPU 上没有 GPU 活动，
    # 否则，它们将干扰 vLLM 内存分析，
    # 并引起意外的行为。
    llm = ray.remote(
        num_cpus=0,
        num_gpus=0,
        scheduling_strategy=PlacementGroupSchedulingStrategy(
            placement_group=pg,
            placement_group_capture_child_tasks=True,
        ),
    )(MyLLM).remote(
        model="facebook/opt-125m",
        enforce_eager=True,
        worker_extension_cls="rlhf_utils.ColocateWorkerExtension",
        tensor_parallel_size=2,
        distributed_executor_backend="ray",
        gpu_memory_utilization=0.4,
        bundle_indices=bundle_indices,
    )
    inference_engines.append(llm)
    # don't call any method on the inference engine here,
    # otherwise it will block until the vLLM instance is created.
    # 在此处的推理引擎上不要调用任何方法，
    # 否则，它将锁定直到创建 vLLM 实例。

for i, llm in enumerate(inference_engines):
    inference_engine_device_ids.append(
        ray.get(llm.collective_rpc.remote("report_device_id", args=tuple())))
    print(f"inference engine {i} is on {inference_engine_device_ids[-1]}")

# 检查部署情况
# 前两个训练执行器(training actors)应当
# 与第一个推理引擎(inference engine)部署在同一GPU上
assert training_actor_device_ids[:2] == inference_engine_device_ids[0]

# 最后两个训练执行器(training actors)应当
# 与第二个推理引擎(inference engine)部署在同一GPU上
assert training_actor_device_ids[2:] == inference_engine_device_ids[1]

print("gather all the IPC handles from the training actors")
ipc_handles = {}
for actor in training_actors:
    ipc_handles.update(ray.get(actor.get_weight_ipc_handles.remote()))

print("update the weights of the inference engines")
for llm in inference_engines:
    ray.get(
        llm.collective_rpc.remote("update_weights_from_ipc_handles",
                                  args=(ipc_handles, )))
print("check if the weights are updated")
for llm in inference_engines:
    assert ray.get(
        llm.collective_rpc.remote("check_weights_changed", args=tuple()))