Intégration de vLLM

Ce document vous guide à travers le processus d’utilisation de vLLM avec TRL pour une génération plus rapide dans les méthodes en ligne telles que GRPO et Online DPO. Nous commençons par un résumé (tl;dr) sur la façon d’utiliser vLLM avec TRL, puis nous détaillons le fonctionnement interne.

TRL prend actuellement en charge uniquement les versions de vLLM comprises entre 0.23.0 et 0.33.x. Veuillez vous assurer d'avoir installé une version comprise dans cette plage afin d'éviter tout problème de compatibilité.

Les formateurs suivants prennent actuellement en charge la génération avec vLLM :

• GRPOTrainer
• RLOOTrainer
• experimental.online_dpo.OnlineDPOTrainer

🚀 Comment utiliser vLLM avec TRL pour accélérer l'entraînement ?

💡 Note : Ressources requises pour cet exemple spécifique : un nœud unique avec 8 GPU.

Lors de l'utilisation de vLLM avec TRL, le serveur vLLM et le formateur doivent s'exécuter sur des appareils CUDA distincts afin d'éviter les conflits. Pour des conseils sur la configuration appropriée, consultez Modes d'utilisation de vLLM pendant l'entraînement.

Installez d'abord vLLM à l'aide de la commande suivante :

pip install "trl[vllm]"

Puis exécutez le serveur sur des GPU spécifiques (par exemple, les 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

Une fois le serveur en cours d’exécution, vous pouvez l’utiliser pour générer des complétions pour l’entraînement. Dans l’exemple ci-dessous, nous utilisons les différents formateurs pris en charge en utilisant le serveur vLLM pour la génération. Les arguments --tensor-parallel-size et --data-parallel-size contrôlent la manière dont le modèle et les données sont répartis sur les GPU.

Dans cet exemple, nous fragmentons un modèle sur 4 GPU à l'aide du parallélisme de tenseurs. Ensuite, exécutez le script d'entraînement sur des GPU différents (par exemple, les GPU 4-7) en passant use_vllm=True dans les arguments d'entraînement comme suit :

Exemple d'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()

Et la commande train sur des GPU distincts du serveur :

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

Pourquoi utiliser vLLM ?

Les méthodes en ligne génèrent des complétions pendant l’entraînement, et leur génération via generate du modèle constitue le goulot d’étranglement. vLLM sert ces complétions beaucoup plus rapidement, grâce à des techniques telles que PagedAttention.

Comment TRL utilise le serveur 🔍

Le formateur demande des complétions sur le point de terminaison compatible OpenAI /v1/completions en envoyant les identifiants de jetons du prompt. Les prompts multimodaux suivent un autre chemin : le serveur traite les images lui-même, et les caractéristiques résultantes sont associées aux mêmes identifiants de jetons sur /inference/v1/generate, car aucun point de terminaison compatible OpenAI n’accepte à la fois des identifiants de jetons et des images.

Le serveur ne fait que générer. Après chaque étape de l’optimiseur, le formateur y transmet les poids mis à jour via NCCL, en les annonçant avec /start_weight_update et /update_weights et en les validant avec /finish_weight_update.

Utilisation avancée

🍷 Plus d'options de personnalisation avec vLLM ?

Vous pouvez personnaliser la configuration du serveur en transmettant tout argument de vllm serve, par exemple --tensor-parallel-size, --data-parallel-size, --max-model-len, --enable-prefix-caching, --enforce-eager, --kv-cache-dtype ou --trust-remote-code. Exécutez vllm serve --help pour obtenir la liste complète.

Seuls les éléments suivants sont requis par TRL :

trl vllm-serve est obsolète : il ne fait désormais que construire cette commande et exécuter le serveur vLLM. Il affiche la commande vllm serve exacte qu’il exécute, afin que vous puissiez la copier et supprimer l’outil d’encapsulation.

💆🏻‍♀️ Quelle est la meilleure configuration distribuée ?

Échellez la génération avec --tensor-parallel-size. Le parallélisme de données n’est plus utile pour les modèles denses : depuis vLLM PR #30739 (publié dans 0.14.0), le passage à l’échelle hors ligne par parallélisme de données pour les modèles non-MoE n’est plus pris en charge.

vLLM avec le backend Transformers

vLLM peut utiliser le backend Transformers pour les implémentations de modèles, ce qui fonctionne aussi bien pour les LLM que pour les VLM. Pour l’activer, définissez vllm_model_impl="transformers" dans votre configuration ou passez-le via l’argument en ligne de commande.

Pour plus de détails, consultez vLLM Transformers Backend.

Exemple :

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

Modes d'utilisation de vLLM pendant l'entraînement

TRL prend en charge deux modes pour intégrer vLLM pendant l’entraînement : le mode colocate (par défaut) et le mode serveur.

Mode Colocate

En mode colocate, vLLM s’exécute dans le processus d’entraînement et partage la mémoire GPU avec le modèle d’entraînement. Cela évite de démarrer un serveur séparé et peut améliorer l’utilisation des GPU, mais peut entraîner une concurrence mémoire sur les GPU d’entraînement. C’est le mode par défaut.

Exemple de configuration :

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
)

Mode serveur

En mode serveur, vLLM s'exécute en tant que processus distinct sur des GPU dédiés et communique avec l'entraîneur via HTTP. Cette configuration est idéale si vous disposez de GPU dédiés à l'inférence.

Exemple de configuration :

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

Consultez la documentation de l'entraîneur que vous utilisez pour obtenir des détails spécifiques sur l'utilisation de vLLM et les paramètres.

Pour réduire l'utilisation de la mémoire GPU lors de l'exécution de vLLM, envisagez d'activer le mode veille de vLLM.