vLLM 集成

本文档将指导您如何将 vLLM 与 TRL 结合使用，以在 GRPO 和 Online DPO 等在线方法中实现更快的生成速度。我们首先总结了如何将 vLLM 与 TRL 结合使用的“太长不看版”，然后将深入探讨其底层工作原理。让我们开始吧！🔥

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

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

vLLM 服务器和 TRL 训练器必须使用不同的 CUDA 设备以避免冲突。

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

pip install "trl[vllm]"

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

CUDA_VISIBLE_DEVICES=0,1,2,3 trl vllm-serve --model Qwen/Qwen2.5-7B --tensor-parallel-size 2 --data-parallel-size 2

服务器运行后，您可以使用它来生成用于训练的补全内容。在下面的示例中，我们使用 GRPOTrainer 来训练模型，并使用 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

dataset = load_dataset("trl-lib/tldr", split="train")

# Dummy reward function: count the number of unique characters in the completions
def reward_num_unique_chars(completions, **kwargs):
    return [len(set(c)) for c in completions]

training_args = GRPOConfig(
    output_dir="my_test",
    use_vllm=True,
    bf16=True,
    gradient_checkpointing=True,
)

trainer = GRPOTrainer(
    model="Qwen/Qwen2.5-7B",
    args=training_args,
    reward_funcs=reward_num_unique_chars,
    train_dataset=dataset,
)

trainer.train()

以及在与服务器不同的 GPU 上的训练命令

CUDA_VISIBLE_DEVICES=4,5,6,7 accelerate launch train.py

🎬 回顾：为什么我们需要在在线方法中使用 vLLM？

像 GRPO 或 Online DPO 这样的在线方法要求模型在训练期间生成补全内容，这些内容随后用于计算奖励信号。然而，生成过程可能非常耗时，特别是对于大型或推理模型。在默认设置（没有 vLLM）中，补全内容是使用（未包装的）模型的 `generate` 方法生成的。这种方法很快成为一个主要瓶颈——生成速度慢且效率低下，尤其是在处理大批量或大型模型时。因此，训练时间显著增加，整体效率下降。为了解决这个问题，我们转向 vLLM，它能实现更快、更具可扩展性的生成，帮助消除在线方法中的这一瓶颈。

🤔 vLLM 如何解决生成缓慢的问题？

如果你做过自回归解码器训练，你就会知道所有输入到 LLM 的词元都会生成它们的注意力键（key）和值（value）张量，这些张量保存在 GPU 内存中，以便后续生成基于它们的词元。这些缓存的键和值张量通常被称为 KV 缓存。然而，存储 KV 缓存会占用大量内存，所以 vLLM 使用一种名为 PagedAttention 的技术来解决这个问题。PagedAttention 受到操作系统虚拟内存概念的启发，将连续的键和值存储在非连续的内存空间中，这样效率更高。其具体细节超出了本文档的范围，但简而言之，它允许模型以更高效的方式存储键和值，从而减少内存占用并加速生成过程。如果你感兴趣，请务必查看 vLLM PagedAttention 以获取更多详细信息。

🤔 当你运行 `trl vllm-serve --model <model_name>` 时，究竟发生了什么？

当你运行例如：

CUDA_VISIBLE_DEVICES=0,1,2,3 trl vllm-serve --model Qwen/Qwen2.5-7B --tensor-parallel-size 1 --data-parallel-size 4

会发生以下情况

vllm

vLLM 首先会生成多个工作进程（workers）来并行处理传入的请求。工作进程的数量由 --tensor-parallel-size 和 --data-parallel-size 的值相乘决定。在本例中，它会生成 4 个工作进程（1 × 4）。每个工作进程独立运行，并处理一部分传入的请求——这些请求基本上是发送到服务器用于生成的提示（prompts）。一个关键点是，这 4 个工作进程是并行运行的，每个进程负责处理总传入负载的一部分。
一旦传入的请求（提示）被分配到各个工作进程，模型就开始生成补全内容。在内部，模型的权重会根据 --tensor-parallel-size 参数被分割到多个 GPU 上——这就是张量并行的处理方式。同时，数据并行（由 --data-parallel-size 控制）确保不同的请求集在工作进程之间独立处理。简而言之：张量并行将模型分割到多个 GPU 上，而数据并行将一批请求分割到不同的模型副本上。
尽管 GPU 独立并并行处理请求，它们仍然需要相互通信。请记住，每个 GPU 只处理传入提示的一部分（例如，有 4 个 GPU 和 8 个提示，使用 --data-parallel-size=4 时，每个 GPU 处理 2 个提示）。这种 GPU 间的通信由 NVIDIA 的 NCCL 库高效管理。通信主要确保每个 GPU 获得其正确部分的传入请求——这种通信是轻量级的，不会干扰生成本身。另外，每个提示要生成的补全数量由 GRPO 配置中的 num_generations 设置控制。例如，如果你设置 num_generations=2（如上图所示），每个提示将有两个补全。因此，对于 8 个提示和 num_generations=2，你总共会得到 16 个补全——无论 GPU 数量或并行设置如何。

🥸 运行服务器时底层发生的更多细节

vLLM 服务器通过运行以下命令启动：trl vllm-serve --model Qwen/Qwen2.5-7B。
服务器运行后，它会根据客户端（训练器）的请求，使用 `vllm_client.generate`（此处）生成补全内容。
然后，客户端（训练器）从服务器请求这些补全内容。
这些补全内容用于计算奖励信号。
根据奖励信号和模型的输出，计算损失并执行反向传播以更新模型的权重。
注意：服务器只负责生成补全内容——它不训练模型。因此，模型的权重不会在服务器上更新。反向传播完成后，客户端使用 vllm_client.update_named_param(name, param.data) 将更新后的权重发送到服务器。

使用 vLLM 时，请确保为训练和生成分配的 GPU 是分开的，以避免 NCCL 通信冲突。如果你不设置 `CUDA_VISIBLE_DEVICES` 环境变量，训练脚本将默认使用所有可用的 GPU，这可能会导致设备冲突。从 TRL v0.19.1 之后的下一个版本开始，代码会自动检测并防止使用相同的设备，并在 vllm 服务器进程中引发错误。

RuntimeError: Attempting to use the same CUDA device for multiple distinct roles/ranks within the same communicator. 
Ensure that trainer is using different devices than vLLM server.

例如，如果你想使用 GPU 4-7 进行训练，而服务器在 GPU 0-3 上运行，请设置

CUDA_VISIBLE_DEVICES=4,5,6,7 accelerate launch train.py

🍷 vLLM 的更多自定义选项？

您可以通过传递额外的参数来自定义服务器配置。

$ trl vllm-serve --help
usage: trl vllm-serve [-h] --model MODEL [--revision REVISION] [--tensor_parallel_size TENSOR_PARALLEL_SIZE]
                      [--data_parallel_size DATA_PARALLEL_SIZE] [--host HOST] [--port PORT]
                      [--gpu_memory_utilization GPU_MEMORY_UTILIZATION] [--dtype DTYPE] [--max_model_len MAX_MODEL_LEN]
                      [--enable_prefix_caching ENABLE_PREFIX_CACHING] [--enforce_eager ENFORCE_EAGER] [--log_level LOG_LEVEL]

options:
  -h, --help            Show this help message and exit
  --model MODEL         Model name or path to load the model from. (default: None)
  --revision REVISION   Revision to use for the model. If not specified, the default branch will be used. (default: None)
  --tensor_parallel_size TENSOR_PARALLEL_SIZE, --tensor-parallel-size TENSOR_PARALLEL_SIZE
                        Number of tensor parallel workers to use. (default: 1)
  --data_parallel_size DATA_PARALLEL_SIZE, --data-parallel-size DATA_PARALLEL_SIZE
                        Number of data parallel workers to use. (default: 1)
  --host HOST           Host address to run the server on. (default: 0.0.0.0)
  --port PORT           Port to run the server on. (default: 8000)
  --gpu_memory_utilization GPU_MEMORY_UTILIZATION, --gpu-memory-utilization GPU_MEMORY_UTILIZATION
                        Ratio (between 0 and 1) of GPU memory to reserve for the model weights, activations, and KV cache on the device
                        dedicated to generation powered by vLLM. Higher values will increase the KV cache size and thus improve the
                        model's throughput. However, if the value is too high, it may cause out-of-memory (OOM) errors during
                        initialization. (default: 0.9)
  --dtype DTYPE         Data type to use for vLLM generation. If set to 'auto', the data type will be automatically determined based on
                        the model configuration. Find the supported values in the vLLM documentation. (default: auto)
  --max_model_len MAX_MODEL_LEN, --max-model-len MAX_MODEL_LEN
                        If set, the `max_model_len` to use for vLLM. This can be useful when running with reduced
                        `vllm_gpu_memory_utilization`, leading to a reduced KV cache size. If not set, vLLM will use the model context
                        size, which might be much larger than the KV cache, leading to inefficiencies. (default: None)
  --enable_prefix_caching ENABLE_PREFIX_CACHING, --enable-prefix-caching ENABLE_PREFIX_CACHING
                        Whether to enable prefix caching in vLLM. If set to `True`, ensure that the model and the hardware support this
                        feature. (default: None)
  --enforce_eager ENFORCE_EAGER, --enforce-eager ENFORCE_EAGER
                        Whether to enforce eager execution. If set to `True`, we will disable CUDA graph and always execute the model
                        in eager mode. If `False` (default behavior), we will use CUDA graph and eager execution in hybrid. (default:
                        None)
  --log_level LOG_LEVEL, --log-level LOG_LEVEL
                        Log level for uvicorn. Possible choices: 'critical', 'error', 'warning', 'info', 'debug', 'trace'. (default:
                        info)

🥳 好了，现在服务器已经运行，我们如何用它来生成补全内容呢？

运行训练脚本，并在训练参数中传递 use_vllm=True

from trl import GRPOConfig

training_args = GRPOConfig(..., use_vllm=True)

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

首先，最重要的是，始终记住最佳设置取决于

模型大小
您拥有的 GPU 数量
GPU 内存大小
您使用的批处理大小
您发送到服务器的请求数量（提示）
您使用的 `max_model_len`（这是模型可以处理的输入序列的最大长度，也称为上下文窗口大小）
您为每个请求生成的补全数量 (`num_generations`)

考虑到这些因素，我们对 Qwen 模型家族（3B、7B、14B、32B）使用 8 个 H100 GPU 进行的实验表明

对于中等大小的模型（3B-14B）和适度的上下文窗口（max_len < 8k），将全部容量用于数据并行可以获得更好的吞吐量。设置 (tp=1, dp=8) 效果最佳。
对于较大的模型（32B）和较长的上下文窗口（max_len > 8k），较小的 DP 大小结合一些模型侧并行表现更好。例如，对于具有较大上下文窗口的 32B 模型，(tp=2, dp=4) 是一个很好的设置。

vLLM 与 Transformers 后端

vLLM 现在支持 transformers 后端用于模型实现。只需在配置的 `vllm_model_impl` 中或通过参数解析器传入 `transformers`，即可设置使用 transformers 后端。这对于 LLM 和 VLM 都适用。请看下面的示例，您可以在此处获取更多信息。

CUDA_DEVICE_ORDER=PCI_BUS_ID CUDA_VISIBLE_DEVICES=0 trl vllm-serve --model Qwen/Qwen
2.5-VL-3B-Instruct --tensor-parallel-size 1 --port 8000 --enforce_eager --vllm_model_impl transformers

< > 在 GitHub 上更新