vllm_rollout.py — vLLM 推理适配器

文件概述

vLLM 推理引擎的 ServerAdapter 实现（约 203 行）和 vLLMReplica 副本管理。ServerAdapter 是训练框架与 vLLM 推理服务器之间的桥梁。

核心类

ServerAdapter(BaseRollout)

训练侧的 vLLM 适配器，通过 Ray 远程调用控制 vLLM 服务器。

class ServerAdapter(BaseRollout):
    def __init__(self, config, model_config, device_mesh, replica_rank=-1):
        # ZMQ IPC 句柄，用于权重传输
        self.device_uuid = get_device_uuid(get_device_id())
        self.zmq_handle = f"ipc:///tmp/rl-colocate-zmq-{self.device_uuid}.sock"

    async def resume(self, tags):
        """唤醒 vLLM 服务器（恢复显存）"""
        await self._execute_method("wake_up", kwargs={"tags": tags})

    async def release(self):
        """休眠 vLLM 服务器（释放显存）"""
        await self._execute_method("sleep", kwargs={"level": self.sleep_level})

    async def update_weights(self, weights, global_steps=None, **kwargs):
        """通过 CUDA IPC 或共享内存传输权重到 vLLM worker"""
        # 1. 通知 vLLM worker 准备接收权重
        future = await self._execute_method("update_weights_from_ipc", non_block=True, ...)

        # 2. 通过 ZMQ + CUDA IPC 发送权重
        sender = BucketedWeightSender(zmq_handle=self.zmq_handle, ...)
        await sender.async_send_weights(weights)

        # 3. 等待 vLLM worker 加载完成
        await future

        # 4. 清理前缀缓存
        await self.server_handle.clear_kv_cache.remote()

vLLMReplica(RolloutReplica)

管理 vLLM 服务器的多节点部署。

class vLLMReplica(RolloutReplica):
    async def launch_servers(self):
        """在每个节点启动 vLLM HTTP 服务器

        流程:
        1. 获取每个 worker 的 node_id 和 CUDA_VISIBLE_DEVICES
        2. 在每个节点创建 vLLMHttpServer actor
        3. 启动 HTTP 服务
        4. 获取服务地址
        """

权重传输流程

训练 Worker (Actor)                    vLLM Worker
      │                                    │
      │ 1. update_weights_from_ipc()       │
      ├───────────────────────────────────→│
      │                                    │ (准备接收)
      │ 2. ZMQ + CUDA IPC 发送权重         │
      ├════════════════════════════════════→│
      │    (分桶传输，避免 OOM)              │ (加载权重)
      │                                    │
      │ 3. clear_kv_cache()                │
      ├───────────────────────────────────→│
      │                                    │

与其他模块的关系

• 继承 rollout/base.py 的 BaseRollout
• 使用 bucketed_weight_transfer.py 的 BucketedWeightSender
• 控制 vllm_async_server.py 中的 vLLMHttpServer
• 被 fsdp_workers.py / engine_workers.py 使用

小结

ServerAdapter 是训练框架控制 vLLM 推理引擎的核心接口，实现了权重同步、显存管理等 Hybrid Engine 关键操作。