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 上更新