Интеграция с vLLM

В этом документе описан процесс использования vLLM вместе с TRL для ускорения генерации в онлайн-методах, таких как GRPO и Online DPO. Сначала мы кратко (tl;dr) расскажем, как использовать vLLM с TRL, а затем перейдём к подробному описанию того, как это работает под капотом.

TRL в настоящее время поддерживает только версии vLLM от 0.23.0 до 0.33.x. Убедитесь, что у вас установлена версия из этого диапазона, чтобы избежать проблем с совместимостью.

Следующие тренера в настоящее время поддерживают генерацию с vLLM:

• GRPOTrainer
• RLOOTrainer
• experimental.online_dpo.OnlineDPOTrainer

🚀 Как использовать vLLM вместе с TRL для ускорения обучения?

💡 Примечание: Ресурсы, необходимые для этого конкретного примера: один узел с 8 GPU.

При использовании vLLM с TRL сервер 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 с использованием tensor parallelism. Затем запускаем скрипт обучения на других GPU (например, GPU 4-7), передав use_vllm=True в аргументах обучения следующим образом:

Пример простого скрипта 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()

А команда train на отдельных GPU от сервера:

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

Зачем использовать vLLM?

Онлайн-методы генерируют завершения во время обучения, и создание их с помощью собственного метода generate модели является узким местом. vLLM предоставляет эти завершения значительно быстрее благодаря таким техникам, как PagedAttention.

Как TRL использует сервер 🔍

Тренер запрашивает завершения по OpenAI-совместимому эндпоинту /v1/completions, передавая идентификаторы токенов промпта. Мультимодальные промпты обрабатываются иначе: сервер самостоятельно обрабатывает изображения, а полученные признаки связываются с теми же идентификаторами токенов на /inference/v1/generate, поскольку ни один OpenAI-совместимый эндпоинт не принимает идентификаторы токенов и изображения одновременно.

Сервер только генерирует. После каждого шага оптимизатора тренер потоково передаёт в него обновлённые веса по 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, не поддерживается.

vLLM с бэкендом Transformers

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 во время обучения: режим colocation (по умолчанию) и режим сервера.

Режим 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",
)

Ознакомьтесь с документацией используемого вами тренера, чтобы узнать подробности о применении vLLM и параметрах.

Чтобы снизить потребление памяти GPU при запуске vLLM, рассмотрите возможность включения режима сна vLLM.