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 训练框架集成的粘合层,处理权重传输、模型补丁等关键功能。