vLLM 集成

本文档将指导您如何使用 vLLM 与 TRL 配合，以在 GRPO 和 Online DPO 等在线方法中实现更快的生成。我们首先简要总结 vLLM 与 TRL 的使用方法，然后深入介绍其底层工作原理。

TRL 目前仅支持 0.23.0 到 0.33.x 版本的 vLLM。请确保您安装的版本在此范围内，以避免兼容性问题。

以下训练器目前支持使用 vLLM 进行生成：

• GRPOTrainer
• RLOOTrainer
• experimental.online_dpo.OnlineDPOTrainer

🚀 如何使用 vLLM 与 TRL 来加速训练？

💡 注意：此特定示例所需的资源：一个配备 8 个 GPU 的节点。

在使用 TRL 搭配 vLLM 时，vLLM 服务器和训练器必须在不同的 CUDA 设备上运行，以避免冲突。 有关如何正确配置此设置的指导，请参阅训练期间使用 vLLM 的模式。

首先，使用以下命令安装 vLLM：

pip install "trl[vllm]"

然后在特定 GPU 上运行服务器（例如 GPU 0-3）：

CUDA_VISIBLE_DEVICES=0,1,2,3 VLLM_SERVER_DEV_MODE=1 vllm serve Qwen/Qwen2.5-7B --tensor-parallel-size 4 \
    --weight-transfer-config '{"backend": "nccl"}' \
    --logprobs-mode processed_logprobs \
    --max-logprobs -1

服务器运行后，即可使用它来生成用于训练的补全。在下方的示例中，我们使用 vLLM 服务器进行生成，并采用了不同的受支持的训练器。--tensor-parallel-size 和 --data-parallel-size 参数用于控制模型和数据在 GPU 上的分片方式。

在此示例中，我们使用张量并行将一个模型分片到 4 个 GPU 上。然后，通过在训练参数中传入 use_vllm=True，在不同的 GPU（例如 GPU 4-7）上运行训练脚本，如下所示：

一个简单的 train.py 脚本示例：

from datasets import load_dataset
from trl import GRPOTrainer, GRPOConfig
from trl.rewards import accuracy_reward

dataset = load_dataset("trl-lib/DeepMath-103K", split="train")

trainer = GRPOTrainer(
    model="Qwen/Qwen2.5-7B",
    args=GRPOConfig(use_vllm=True, vllm_mode="server"),
    reward_funcs=accuracy_reward,
    train_dataset=dataset,
)

trainer.train()

from datasets import load_dataset
from trl.experimental.online_dpo import OnlineDPOConfig, OnlineDPOTrainer
from trl.rewards import accuracy_reward

dataset = load_dataset("trl-lib/DeepMath-103K", split="train")

trainer = OnlineDPOTrainer(
    model="Qwen/Qwen2.5-7B",
    args=OnlineDPOConfig(use_vllm=True, vllm_mode="server"),
    reward_funcs=accuracy_reward,
    train_dataset=dataset,
)

trainer.train()

from datasets import load_dataset
from trl import RLOOTrainer, RLOOConfig
from trl.rewards import accuracy_reward

dataset = load_dataset("trl-lib/DeepMath-103K", split="train")

trainer = RLOOTrainer(
    model="Qwen/Qwen2.5-7B",
    args=RLOOConfig(use_vllm=True, vllm_mode="server"),
    reward_funcs=accuracy_reward,
    train_dataset=dataset,
)

trainer.train()

以及在服务器上独立 GPU 上的 train 命令：

CUDA_VISIBLE_DEVICES=4,5,6,7 torchrun --nproc_per_node 4 train.py

为什么使用 vLLM？

在线方法在训练过程中生成补全，而使用模型自身的 generate 来生成补全往往是性能瓶颈。得益于 PagedAttention 等技术，vLLM 能够以快得多的速度提供这些补全。

TRL 如何使用服务器 🔍

Trainer 会请求 OpenAI 兼容的 /v1/completions 端点以获取补全，并发送提示词（prompt）的 token ID。多模态提示词走不同的路径：服务器自行处理图像，并将生成的特征与相同的 token ID 配对，通过 /inference/v1/generate 发送，因为没有任何 OpenAI 兼容端点能同时接收 token ID 和图像。

服务器仅负责生成。在每个优化器步骤之后，训练器会通过 NCCL 将更新后的权重流式传输到服务器中，并使用 /start_weight_update 和 /update_weights 进行通知，最后通过 /finish_weight_update 提交这些权重。

高级用法

🍷 使用 vLLM 获得更多自定义选项？

你可以通过传入任意 vllm serve 参数来自定义服务器配置，例如 --tensor-parallel-size、--data-parallel-size、--max-model-len、--enable-prefix-caching、--enforce-eager、--kv-cache-dtype 或 --trust-remote-code。运行 vllm serve --help 可查看完整列表。

TRL 仅要求以下内容：

trl vllm-serve 已弃用：它现在仅构建该命令并运行 vLLM 的服务器。它会打印出所执行的精确 vllm serve 命令，因此您可以直接复制该命令并移除包装器。

💆🏻‍♀️ 最佳的分布式配置是什么？

使用 --tensor-parallel-size 扩展生成规模。数据并行对稠密模型不再有帮助：自 vLLM PR #30739（发布于 0.14.0）起，不再支持非 MoE 模型的离线数据并行扩展。

使用 Transformers 后端的 vLLM

vLLM 可以使用 Transformers 后端 作为模型实现，适用于 LLM 和 VLM。 要启用此功能，请在配置中设置 vllm_model_impl="transformers"，或通过命令行参数传入。

有关更多详细信息，请参阅 vLLM Transformers Backend。

示例：

CUDA_DEVICE_ORDER=PCI_BUS_ID CUDA_VISIBLE_DEVICES=0 VLLM_SERVER_DEV_MODE=1 vllm serve Qwen/Qwen2.5-VL-3B-Instruct \
    --tensor-parallel-size 1 --port 8000 --enforce-eager --model-impl transformers \
    --weight-transfer-config '{"backend": "nccl"}' \
    --logprobs-mode processed_logprobs \
    --max-logprobs -1 \
    --enable-scale-out  # vLLM 0.30.0 and later, drop on earlier versions

训练期间使用 vLLM 的模式

TRL 支持在训练期间集成 vLLM 的两种模式：colocate 模式（默认）和 server 模式。

Colocate 模式

在 colocate 模式下，vLLM 在训练器进程内运行，并与训练模型共享 GPU 内存。 这避免了启动单独的服务器，并可能提高 GPU 利用率，但可能导致训练 GPU 上的内存争用。这是默认模式。

示例配置：

from trl import GRPOConfig

training_args = GRPOConfig(
    ...,
    use_vllm=True,  # vllm_mode="colocate" by default
)

from trl.experimental.online_dpo import OnlineDPOConfig

training_args = OnlineDPOConfig(
    ...,
    use_vllm=True,  # vllm_mode="colocate" by default
)

from trl import RLOOConfig

training_args = RLOOConfig(
    ...,
    use_vllm=True,  # vllm_mode="colocate" by default
)

服务器模式

在 服务器模式下，vLLM 作为独立进程运行在专用 GPU 上，并通过 HTTP 与训练器通信。 如果您拥有专用于推理的 GPU，这种配置非常理想。

示例配置：

from trl import GRPOConfig

training_args = GRPOConfig(
    ...,
    use_vllm=True,
    vllm_mode="server",
)

from trl.experimental.online_dpo import OnlineDPOConfig

training_args = OnlineDPOConfig(
    ...,
    use_vllm=True,
    vllm_mode="server",
)

from trl import RLOOConfig

training_args = RLOOConfig(
    ...,
    use_vllm=True,
    vllm_mode="server",
)

请查阅你所使用的 trainer 的文档，以获取有关 vLLM 用法和参数的具体细节。

在运行 vLLM 时，若要降低 GPU 内存占用，请考虑启用 vLLM 睡眠模式。