Integración de vLLM

Este documento te guiará a través del proceso de usar vLLM con TRL para una generación más rápida en métodos en línea como GRPO y Online DPO. Primero resumiremos un resumen breve (tl;dr) sobre cómo usar vLLM con TRL, y luego entraremos en los detalles de cómo funciona por debajo.

TRL actualmente solo admite versiones de vLLM desde 0.23.0 hasta 0.33.x. Asegúrate de tener instalada una versión dentro de este rango para evitar problemas de compatibilidad.

Los siguientes entrenadores admiten actualmente la generación con vLLM:

• GRPOTrainer
• RLOOTrainer
• experimental.online_dpo.OnlineDPOTrainer

🚀 ¿Cómo puedo usar vLLM con TRL para acelerar el entrenamiento?

💡 Nota: Recursos necesarios para este ejemplo específico: un nodo único con 8 GPUs.

Al usar vLLM con TRL, el servidor de vLLM y el trainer deben ejecutarse en dispositivos CUDA separados para evitar conflictos. Para obtener orientación sobre cómo configurarlo correctamente, consulte Modos de uso de vLLM durante el entrenamiento.

Primero, instala vLLM usando el siguiente comando:

pip install "trl[vllm]"

Luego, ejecuta el servidor en GPUs específicas (p. ej., GPUs 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

Una vez que el servidor esté en ejecución, puede usarse para generar completions para el entrenamiento. En el ejemplo siguiente, se utilizan los distintos trainers admitidos con el servidor vLLM para la generación. Los argumentos --tensor-parallel-size y --data-parallel-size controlan cómo se fragmentan el modelo y los datos entre las GPUs.

En este ejemplo, dividimos un modelo en 4 GPUs mediante paralelismo de tensores. A continuación, ejecutamos el script de entrenamiento en otras GPUs (p. ej., GPUs 4-7) pasando use_vllm=True en los argumentos de entrenamiento de la siguiente manera:

Ejemplo de un simple script 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()

Y el comando de entrenamiento en GPUs separadas del servidor:

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

¿Por qué usar vLLM?

Los métodos en línea generan completiones durante el entrenamiento, y generarlas con el propio generate del modelo es el cuello de botella. vLLM sirve esas completiones mucho más rápido, gracias a técnicas como PagedAttention.

Cómo TRL utiliza el servidor 🔍

El trainer solicita completions en el endpoint /v1/completions compatible con OpenAI, enviando los IDs de tokens del prompt. Los prompts multimodales siguen una ruta diferente: el servidor procesa las imágenes por su cuenta, y las características resultantes se emparejan con los mismos IDs de tokens en /inference/v1/generate, ya que ningún endpoint compatible con OpenAI acepta IDs de tokens e imágenes a la vez.

El servidor solo genera. Después de cada paso del optimizador, el entrenador transmite los pesos actualizados a través de NCCL, anunciándolos con /start_weight_update y /update_weights y confirmándolos con /finish_weight_update.

Uso avanzado

🍷 ¿Más opciones de personalización con vLLM?

Puedes personalizar la configuración del servidor pasando cualquier argumento de vllm serve, por ejemplo --tensor-parallel-size, --data-parallel-size, --max-model-len, --enable-prefix-caching, --enforce-eager, --kv-cache-dtype o --trust-remote-code. Ejecuta vllm serve --help para ver la lista completa.

Solo se requieren los siguientes elementos para TRL:

trl vllm-serve está en desuso: ahora solo construye este comando y ejecuta el servidor de vLLM. Imprime el comando exacto de vllm serve que ejecuta, para que puedas copiarlo y eliminar el envoltorio.

💆🏻‍♀️ ¿Cuál es la mejor configuración distribuida?

Escala la generación con --tensor-parallel-size. El paralelismo de datos ya no es útil para modelos densos: desde vLLM PR #30739 (publicado en 0.14.0), no se admite la escalabilidad de paralelismo de datos sin conexión para modelos que no son MoE.

vLLM con backend de Transformers

vLLM puede utilizar el backend de Transformers para las implementaciones de modelos, que funciona tanto para LLMs como para VLMs. Para habilitarlo, establece vllm_model_impl="transformers" en tu configuración o pásalo mediante el argumento de línea de comandos.

Para más detalles, consulta vLLM Transformers Backend.

Ejemplo:

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

Modos de uso de vLLM durante el entrenamiento

TRL admite dos modos para integrar vLLM durante el entrenamiento: modo colocate (predeterminado) y modo server.

Modo Colocate

En el modo colocate, vLLM se ejecuta dentro del proceso del entrenador y comparte la memoria de GPU con el modelo de entrenamiento. Esto evita iniciar un servidor separado y puede mejorar la utilización de la GPU, pero puede provocar una contención de memoria en las GPUs de entrenamiento. Este es el modo predeterminado.

Ejemplo de configuración:

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
)

Modo servidor

En modo servidor, vLLM se ejecuta como un proceso separado en GPUs dedicadas y se comunica con el entrenador a través de HTTP. Esta configuración es ideal si tienes GPUs dedicadas a la inferencia.

Ejemplo de configuración:

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",
)

Consulte la documentación del trainer que está utilizando para obtener detalles específicos sobre el uso de vLLM y los parámetros.

Para reducir el uso de memoria de GPU al ejecutar vLLM, considere activar el modo de suspensión de vLLM.