Integración de DeepSpeed

TRL admite el entrenamiento con DeepSpeed, una biblioteca que implementa técnicas avanzadas de optimización del entrenamiento. Estas incluyen la partición del estado del optimizador, el descargado (offloading), la partición de gradientes y más.

DeepSpeed integra el Zero Redundancy Optimizer (ZeRO), que permite escalar el tamaño del modelo en proporción al número de dispositivos manteniendo una alta eficiencia.

Instalación

Para usar DeepSpeed con TRL, instálalo usando el siguiente comando:

pip install deepspeed

Ejecución de scripts de entrenamiento con DeepSpeed

No se requieren modificaciones en su script de entrenamiento. DeepSpeed se habilita a través de los argumentos de entrenamiento, igual que con cualquier Trainer: establezca deepspeed en la ruta de un archivo de configuración de DeepSpeed y ejecute con torchrun.

training_args = SFTConfig(..., deepspeed="<DEEPSPEED_CONFIG_FILE.json>")

torchrun --nproc_per_node 8 train.py

Los scripts que analizan sus argumentos con TrlParser, como los que están detrás de la CLI de trl, los obtienen de la línea de comandos:

trl sft ... --deepspeed <DEEPSPEED_CONFIG_FILE.json>

Proporcionamos archivos de configuración de DeepSpeed listos para usar en el directorio examples/deepspeed_configs. Por ejemplo, para ejecutar el entrenamiento con ZeRO Stage 2:

trl sft ... --deepspeed examples/deepspeed_configs/zero2.json

Los valores establecidos en "auto" en estos archivos (tamaño de lote, acumulación de gradientes, precisión, …) se completan a partir de los argumentos de entrenamiento.

Recursos adicionales

Consulte la documentación de DeepSpeed de 🤗 Transformers para obtener más información sobre la integración de DeepSpeed.