make_vllm_worker¶
- class torchrl.modules.llm.make_vllm_worker(*, model_name: str, devices: list[torch.device | int] | None = None, num_devices: int | None = None, make_ray_worker: bool = True, enforce_eager: bool = False, **kwargs)[源代码]¶
创建一个支持张量并行化的 vLLM 推理引擎。
- 参数:
model_name (str) – 传递给 vLLM.LLM 的模型名称。
devices (list[torch.device | int], optional) – 要使用的设备列表。与 num_devices 互斥。
num_devices (int, optional) – 要使用的设备数量。与 devices 互斥。
make_ray_worker (bool, optional) – 是否创建 Ray actor。默认为 True。
enforce_eager (bool, optional) – 是否强制执行即时执行。默认为 False。
**kwargs – 传递给 vLLM.LLM.__init__ 的其他参数。
- 返回:
返回 Ray worker 包装器或本地 LLM 包装器，两者都实现了 RLvLLMEngine。
- 返回类型:
RayLLMWorker | LocalLLMWrapper
示例
>>> # Create a 2-GPU tensor parallel worker with Ray >>> worker = make_vllm_worker("Qwen/Qwen2.5-3B", num_devices=2) >>> # Create a local LLM instance on GPU 1 >>> llm = make_vllm_worker("Qwen/Qwen2.5-3B", devices=[1], make_ray_worker=False)
