TRL documentation
Accélérer l'entraînement
Accélérer l'entraînement
Ce guide présente diverses méthodes pour accélérer l’entraînement dans TRL. Chaque technique est accompagnée d’exemples minimaux et de liens vers une documentation plus complète.
vLLM pour une génération rapide dans les méthodes en ligne
Les méthodes en ligne telles que GRPO ou Online DPO nécessitent que le modèle génère des complétions, ce qui est souvent un processus lent et peut avoir un impact significatif sur le temps d’entraînement. Pour accélérer la génération, vous pouvez utiliser vLLM, une bibliothèque qui permet une génération rapide grâce, entre autres, à PagedAttention. Les formateurs en ligne de TRL prennent en charge vLLM, améliorant considérablement la vitesse d’entraînement. Pour plus de détails, consultez Intégration vLLM.
Pour utiliser vLLM, commencez par l'installer à l'aide de :
pip install trl[vllm]
Tout d'abord, démarrez un serveur vLLM en exécutant :
VLLM_SERVER_DEV_MODE=1 vllm serve <model_name> \
--weight-transfer-config '{"backend": "nccl"}' \
--logprobs-mode processed_logprobs \
--max-logprobs -1Ensuite, exécutez le script d'entraînement et passez use_vllm=True dans les arguments d'entraînement.
from trl.experimental.online_dpo import OnlineDPOConfig
training_args = OnlineDPOConfig(..., use_vllm=True, vllm_mode="server")Tout d'abord, démarrez un serveur vLLM en exécutant :
VLLM_SERVER_DEV_MODE=1 vllm serve <model_name> \
--weight-transfer-config '{"backend": "nccl"}' \
--logprobs-mode processed_logprobs \
--max-logprobs -1Ensuite, exécutez le script d'entraînement et passez use_vllm=True dans les arguments d'entraînement.
from trl import GRPOConfig
training_args = GRPOConfig(..., use_vllm=True, vllm_mode="server")Vous pouvez personnaliser la configuration du serveur en passant des arguments supplémentaires. Pour plus d’informations, consultez l’intégration vLLM.
Lors de l’utilisation de vLLM, veillez à ce que les GPU affectés à l’entraînement et à la génération soient distincts afin d’éviter tout conflit de ressources. Par exemple, si vous prévoyez d’utiliser 4 GPU pour l’entraînement et 4 autres pour la génération avec vLLM, vous pouvez spécifier l’allocation des GPU à l’aide de
CUDA_VISIBLE_DEVICES.Définir les GPU 0-3 pour la génération vLLM :
CUDA_VISIBLE_DEVICES=0,1,2,3 VLLM_SERVER_DEV_MODE=1 vllm serve <model_name> \ --weight-transfer-config '{"backend": "nccl"}' \ --logprobs-mode processed_logprobs \ --max-logprobs -1Et les GPU 4-7 pour l'entraînement :
CUDA_VISIBLE_DEVICES=4,5,6,7 torchrun --nproc_per_node 4 train.py
Tout d'abord, démarrez un serveur vLLM en exécutant :
VLLM_SERVER_DEV_MODE=1 vllm serve <model_name> \
--weight-transfer-config '{"backend": "nccl"}' \
--logprobs-mode processed_logprobs \
--max-logprobs -1Ensuite, exécutez le script d'entraînement et passez use_vllm=True dans les arguments d'entraînement.
from trl import RLOOConfig
training_args = RLOOConfig(..., use_vllm=True, vllm_mode="server")Vous pouvez personnaliser la configuration du serveur en passant des arguments supplémentaires. Pour plus d’informations, consultez l’intégration vLLM.
Lors de l’utilisation de vLLM, veillez à ce que les GPU affectés à l’entraînement et à la génération soient distincts afin d’éviter tout conflit de ressources. Par exemple, si vous prévoyez d’utiliser 4 GPU pour l’entraînement et 4 autres pour la génération avec vLLM, vous pouvez spécifier l’allocation des GPU à l’aide de
CUDA_VISIBLE_DEVICES.Définir les GPU 0-3 pour la génération vLLM :
CUDA_VISIBLE_DEVICES=0,1,2,3 VLLM_SERVER_DEV_MODE=1 vllm serve <model_name> \ --weight-transfer-config '{"backend": "nccl"}' \ --logprobs-mode processed_logprobs \ --max-logprobs -1Et les GPU 4-7 pour l'entraînement :
CUDA_VISIBLE_DEVICES=4,5,6,7 torchrun --nproc_per_node 4 train.py
Implémentations d'attention optimisées
TRL prend en charge diverses implémentations d'attention optimisées qui peuvent accélérer considérablement l'entraînement tout en réduisant l'utilisation de la mémoire. Vous pouvez utiliser directement des noyaux pré-optimisés depuis le Kernels Hub ou un backend d'attention construit manuellement.
Vous pouvez utiliser des noyaux d'attention pré-optimisés depuis le Hub sans compilation manuelle :
from trl import SFTConfig
training_args = SFTConfig(..., model_init_kwargs={"attn_implementation": "kernels-community/flash-attn2"})D'autres options incluent kernels-community/vllm-flash-attn3 et kernels-community/paged-attention.
L'attention optimisée fonctionne avec tous les formateurs TRL. Pour plus de détails, consultez Intégration Kernels Hub.
Construire manuellement des backends d’attention optimisés est complexe et chronophage. Il n’est jamais recommandé de le faire sauf si c’est absolument nécessaire. Pensez plutôt à utiliser les Kernels du Hub, comme décrit dans la section précédente.
Si vous avez installé manuellement un backend d'attention optimisé comme Flash Attention 2, vous pouvez le spécifier dans les arguments d'entraînement :
from trl import SFTConfig
training_args = SFTConfig(..., model_init_kwargs={"attn_implementation": "flash_attention_2"})Entraînement en précision mixte
L’entraînement en précision mixte utilisant bf16 ou fp16 peut accélérer l’entraînement et réduire l’utilisation de la mémoire avec un impact minimal sur la qualité du modèle.
from trl import SFTConfig
training_args = SFTConfig(..., bf16=True) # or fp16=True for older GPUsUtilisez bf16=True pour les GPU Ampere (A100, RTX 30xx) ou plus récents, et fp16=True pour les GPU plus anciens. L’entraînement en précision mixte est pris en charge par tous les formateurs TRL.