跳转至

utils.py — vLLM Rollout 工具

文件概述

提供 vLLM 推理引擎的工具类和函数(约 295 行),包括 Worker 扩展类、LoRA 支持、FP8 量化等。

核心类

vLLMColocateWorkerExtension

vLLM Worker 的扩展类,在 vLLM 子进程中运行,提供权重更新和模型补丁功能。

class vLLMColocateWorkerExtension:
    """vLLM worker 子进程中的扩展

    通过 worker_extension_cls 参数传递给 vLLM,
    在每个 TP worker 中实例化。
    """

    def __new__(cls, **kwargs):
        # 1. LoRA 补丁
        VLLMHijack.hijack()
        # 2. FP8 量化补丁
        if os.environ.get("VERL_VLLM_FP8_QUANT_ENABLED"):
            apply_vllm_fp8_patches()
        # 3. QAT 补丁
        if _is_qat_model:
            apply_qat_patches()

    def update_weights_from_ipc(self, peft_config=None, base_sync_done=False, use_shm=False):
        """从 CUDA IPC 接收并加载权重

        使用 BucketedWeightReceiver 分桶接收权重,
        然后调用 model.load_weights() 加载到模型中。
        """
        receiver = BucketedWeightReceiver(
            zmq_handle=self._get_zmq_handle(),
            device=self.device,
            use_shm=use_shm,
        )
        receiver.receive_weights(
            on_bucket_received=lambda weights: self._update_weights(weights, ...)
        )

核心函数

get_device_uuid

def get_device_uuid(device_id: int) -> str:
    """获取 GPU 设备 UUID(用于 IPC 句柄命名)"""
    # CUDA: 使用 vllm.platforms.current_platform
    # NPU: 使用 ASCEND_RT_VISIBLE_DEVICES

build_cli_args_from_config

def build_cli_args_from_config(config: dict) -> list[str]:
    """将配置字典转换为 vLLM CLI 参数列表

    示例:
    {"dtype": "bfloat16", "enable_lora": True}
    → ["--dtype", "bfloat16", "--enable-lora"]
    """

与其他模块的关系

  • vLLMColocateWorkerExtension 被 vLLM 子进程加载
  • 使用 bucketed_weight_transfer.py 接收权重
  • 被 vllm_async_server.py 引用

小结

这些工具是 vLLM 推理引擎与 verl 训练框架集成的粘合层,处理权重传输、模型补丁等关键功能。