TRL documentation
Acelerar el entrenamiento
Acelerar el entrenamiento
Esta guía cubre varios métodos para acelerar el entrenamiento en TRL. Cada técnica incluye ejemplos mínimos con enlaces a documentación más completa.
vLLM para una generación rápida en métodos en línea
Los métodos online, como GRPO o Online DPO, requieren que el modelo genere completados, un proceso que suele ser lento y puede afectar significativamente el tiempo de entrenamiento. Para acelerar la generación, puedes usar vLLM, una biblioteca que permite una generación rápida mediante, entre otras cosas, PagedAttention. Los entrenadores online de TRL admiten vLLM, lo que mejora considerablemente la velocidad de entrenamiento. Para más detalles, consulta vLLM Integration.
Para usar vLLM, primero instálalo usando:
pip install trl[vllm]
Primero, inicie un servidor de vLLM ejecutando:
VLLM_SERVER_DEV_MODE=1 vllm serve <model_name> \
--weight-transfer-config '{"backend": "nccl"}' \
--logprobs-mode processed_logprobs \
--max-logprobs -1Luego, ejecuta el script de entrenamiento y pasa use_vllm=True en los argumentos de entrenamiento.
from trl.experimental.online_dpo import OnlineDPOConfig
training_args = OnlineDPOConfig(..., use_vllm=True, vllm_mode="server")Primero, inicie un servidor de vLLM ejecutando:
VLLM_SERVER_DEV_MODE=1 vllm serve <model_name> \
--weight-transfer-config '{"backend": "nccl"}' \
--logprobs-mode processed_logprobs \
--max-logprobs -1Luego, ejecuta el script de entrenamiento y pasa use_vllm=True en los argumentos de entrenamiento.
from trl import GRPOConfig
training_args = GRPOConfig(..., use_vllm=True, vllm_mode="server")Puedes personalizar la configuración del servidor pasando argumentos adicionales. Para obtener más información, consulta integración de vLLM.
Al usar vLLM, asegúrese de que las GPUs asignadas para el entrenamiento y la generación sean distintas para evitar conflictos de recursos. Por ejemplo, si planea usar 4 GPUs para el entrenamiento y otras 4 para la generación con vLLM, puede especificar la asignación de GPUs mediante
CUDA_VISIBLE_DEVICES.Establece las GPUs 0-3 para la generación de 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 -1Y las GPU 4-7 para el entrenamiento:
CUDA_VISIBLE_DEVICES=4,5,6,7 torchrun --nproc_per_node 4 train.py
Primero, inicie un servidor de vLLM ejecutando:
VLLM_SERVER_DEV_MODE=1 vllm serve <model_name> \
--weight-transfer-config '{"backend": "nccl"}' \
--logprobs-mode processed_logprobs \
--max-logprobs -1Luego, ejecuta el script de entrenamiento y pasa use_vllm=True en los argumentos de entrenamiento.
from trl import RLOOConfig
training_args = RLOOConfig(..., use_vllm=True, vllm_mode="server")Puedes personalizar la configuración del servidor pasando argumentos adicionales. Para obtener más información, consulta integración de vLLM.
Al usar vLLM, asegúrese de que las GPUs asignadas para el entrenamiento y la generación sean distintas para evitar conflictos de recursos. Por ejemplo, si planea usar 4 GPUs para el entrenamiento y otras 4 para la generación con vLLM, puede especificar la asignación de GPUs mediante
CUDA_VISIBLE_DEVICES.Establece las GPUs 0-3 para la generación de 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 -1Y las GPU 4-7 para el entrenamiento:
CUDA_VISIBLE_DEVICES=4,5,6,7 torchrun --nproc_per_node 4 train.py
Implementaciones de atención optimizadas
TRL admite diversas implementaciones de atención optimizadas que pueden acelerar significativamente el entrenamiento y reducir el uso de memoria. Puede utilizar kernels preoptimizados directamente desde Kernels Hub o un backend de atención construido manualmente.
Puedes usar kernels de atención preoptimizados del Hub sin compilación manual:
from trl import SFTConfig
training_args = SFTConfig(..., model_init_kwargs={"attn_implementation": "kernels-community/flash-attn2"})Otras opciones incluyen kernels-community/vllm-flash-attn3 y kernels-community/paged-attention.
La atención optimizada funciona en todos los entrenadores de TRL. Para más detalles, consulte Integración con Kernels Hub.
Construir manualmente backends de atención optimizados es complejo y consume mucho tiempo. No se recomienda en absoluto a menos que sea estrictamente necesario. En su lugar, considere usar Kernels del Hub, como se describe en la sección anterior.
Si has instalado manualmente un backend de atención optimizado como Flash Attention 2, puedes especificarlo en los argumentos de entrenamiento:
from trl import SFTConfig
training_args = SFTConfig(..., model_init_kwargs={"attn_implementation": "flash_attention_2"})Entrenamiento con precisión mixta
El entrenamiento con precisión mixta usando bf16 o fp16 puede acelerar el entrenamiento y reducir el uso de memoria con un impacto mínimo en la calidad del modelo.
from trl import SFTConfig
training_args = SFTConfig(..., bf16=True) # or fp16=True for older GPUsUse bf16=True para GPUs Ampere (A100, RTX 30xx) o más nuevas, y fp16=True para GPUs más antiguas. El entrenamiento con precisión mixta está soportado en todos los entrenadores de TRL.