Ajusta más y entrena más rápido con ZeRO a través de DeepSpeed y FairScale

Ajusta y entrena más rápido con ZeRO usando DeepSpeed y FairScale.

Una entrada de blog invitada del compañero de Hugging Face Stas Bekman

A medida que los modelos de Aprendizaje Automático recientes han crecido mucho más rápido que la cantidad de memoria GPU añadida a las tarjetas recién lanzadas, muchos usuarios no pueden entrenar o incluso cargar algunos de esos enormes modelos en su hardware. Si bien hay un esfuerzo en curso para destilar algunos de esos enormes modelos para que tengan un tamaño más manejable, ese esfuerzo no está produciendo modelos lo suficientemente pequeños lo suficientemente rápido.

En otoño de 2019, Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase y Yuxiong He publicaron un artículo: ZeRO: Optimizaciones de memoria hacia el entrenamiento de modelos de billones de parámetros, que contiene una gran cantidad de ideas nuevas ingeniosas sobre cómo se puede hacer que el hardware haga mucho más de lo que se creía posible anteriormente. Poco tiempo después, se lanzó DeepSpeed y dio al mundo la implementación de código abierto de la mayoría de las ideas de ese artículo (algunas ideas aún están en desarrollo) y, al mismo tiempo, un equipo de Facebook lanzó FairScale, que también implementó algunas de las ideas clave del artículo ZeRO.

Si utiliza el Entrenador de Hugging Face, a partir de transformers v4.2.0 tiene soporte experimental para las características de ZeRO de DeepSpeed y FairScale. Los nuevos argumentos de línea de comandos --sharded_ddp y --deepspeed del Trainer proporcionan la integración de FairScale y DeepSpeed, respectivamente. Aquí está la documentación completa .

Esta entrada de blog describirá cómo puede beneficiarse de ZeRO sin importar si posee solo una GPU o un conjunto completo de ellas.

Hagamos un pequeño ajuste fino con un experimento de tarea de traducción, utilizando un modelo t5-large y el script finetune_trainer.py que puede encontrar en ejemplos/seq2seq en el repositorio de GitHub de transformers.

Tenemos 2x 24GB (Titan RTX) GPUs para realizar las pruebas.

Esto es solo una prueba de concepto, por lo que seguramente se pueden mejorar las cosas aún más, por lo que realizaremos la prueba en una pequeña muestra de 2000 elementos para el entrenamiento y 500 elementos para la evaluación para realizar las comparaciones. La evaluación realiza, por defecto, una búsqueda de haz de tamaño 4, por lo que es más lenta que el entrenamiento con el mismo número de muestras, por eso se utilizaron 4 veces menos elementos de evaluación en estas pruebas.

Aquí están los argumentos clave de línea de comandos de nuestra línea base:

export BS=16
python -m torch.distributed.launch --nproc_per_node=2 ./finetune_trainer.py \
--model_name_or_path t5-large --n_train 2000 --n_val 500 \
--per_device_eval_batch_size $BS --per_device_train_batch_size $BS \
--task translation_en_to_ro [...]

Solo estamos utilizando DistributedDataParallel (DDP) y nada más para mejorar el rendimiento de la línea base. Pude ajustar un tamaño de lote (BS) de 16 antes de tener un error de falta de memoria (OOM).

Tenga en cuenta que, por simplicidad y para facilitar la comprensión, solo he mostrado los argumentos de línea de comandos importantes para esta demostración. Encontrará la línea de comandos completa en esta publicación .

A continuación, vamos a volver a ejecutar la prueba cada vez agregando uno de los siguientes:

  1. --fp16
  2. --sharded_ddp (fairscale)
  3. --sharded_ddp --fp16 (fairscale)
  4. --deepspeed sin desactivar la CPU
  5. --deepspeed con desactivación de la CPU

Dado que la optimización clave aquí es que cada técnica despliega la memoria GPU de manera más eficiente, vamos a intentar aumentar continuamente el tamaño del lote y esperamos que el entrenamiento y la evaluación se completen más rápido (manteniendo las métricas estables o incluso mejorando algunas, pero no nos enfocaremos en eso aquí).

Recuerde que las etapas de entrenamiento y evaluación son muy diferentes entre sí, porque durante el entrenamiento se modifican los pesos del modelo, se calculan los gradientes y se almacenan los estados del optimizador. Durante la evaluación, nada de esto sucede, pero en esta tarea particular de traducción, el modelo intentará buscar la mejor hipótesis, por lo que en realidad tiene que realizar múltiples ejecuciones antes de estar satisfecho. Por eso no es rápido, especialmente cuando un modelo es grande.

Veamos los resultados de estas seis ejecuciones de prueba:

Es fácil ver que tanto FairScale como DeepSpeed proporcionan grandes mejoras sobre la línea base, tanto en el tiempo total de entrenamiento y evaluación, como en el tamaño del lote. DeepSpeed implementa más magia a partir de esta escritura y parece ser el ganador a corto plazo, pero FairScale es más fácil de implementar. Para DeepSpeed, necesitas escribir un archivo de configuración simple y cambiar el lanzador de tu línea de comandos, con FairScale solo necesitas agregar el argumento de línea de comandos --sharded_ddp, por lo que es posible que desees probarlo primero, ya que es el objetivo más fácil de alcanzar.

Siguiendo la regla del 80:20, solo he dedicado unas pocas horas a estas pruebas y no he intentado exprimir cada MB y segundo refinando los argumentos de línea de comandos y la configuración, ya que es bastante obvio en la tabla simple qué es lo que querrías probar a continuación. Cuando te enfrentes a un proyecto real que se ejecutará durante horas y tal vez días, definitivamente dedica más tiempo para asegurarte de utilizar los hiperparámetros más óptimos para realizar tu trabajo más rápido y a un costo mínimo.

Si deseas experimentar con esta prueba tú mismo o quieres conocer más detalles sobre el hardware y software utilizados para ejecutarla, por favor, consulta esta publicación.

Mientras que FairScale nos brinda un impulso solo con varias GPUs, DeepSpeed tiene un regalo incluso para aquellos de nosotros con una sola GPU.

Intentemos lo imposible: entrenemos t5-3b en una tarjeta RTX-3090 de 24 GB.

Primero intentemos ajustar el enorme t5-3b utilizando la configuración normal de una sola GPU:

export BS=1
CUDA_VISIBLE_DEVICES=0 ./finetune_trainer.py \
--model_name_or_path t5-3b --n_train 60 --n_val 10 \
--per_device_eval_batch_size $BS --per_device_train_batch_size $BS \
--task translation_en_to_ro --fp16 [...]

No hay suerte, incluso con BS=1 obtenemos:

RuntimeError: CUDA sin memoria. Se intentó asignar 64.00 MiB (GPU 0; capacidad total de 23.70 GiB;
21.37 GiB ya asignados; 45.69 MiB libres; 22.05 GiB reservados en total por PyTorch)

Ten en cuenta que, como antes, solo estoy mostrando las partes importantes y los argumentos completos de la línea de comandos se pueden encontrar aquí.

Ahora actualiza tus transformers a la versión 4.2.0 o superior, luego instala DeepSpeed:

pip install deepspeed

y volvamos a intentarlo, esta vez agregando DeepSpeed a la línea de comandos:

export BS=20
CUDA_VISIBLE_DEVICES=0 deepspeed --num_gpus=1 ./finetune_trainer.py \
--model_name_or_path t5-3b --n_train 60 --n_val 10 \
--per_device_eval_batch_size $BS --per_device_train_batch_size $BS \
--task translation_en_to_ro --fp16 --deepspeed ds_config_1gpu.json [...]

¡Et voilà! Obtenemos un tamaño de lote de 20 entrenado sin problemas. Probablemente podría aumentarlo aún más. El programa falló con OOM en BS=30.

Aquí están los resultados relevantes:

2021-01-12 19:06:31 | INFO | __main__ |   train_n_objs = 60
2021-01-12 19:06:31 | INFO | __main__ |   train_runtime = 8.8511
2021-01-12 19:06:35 | INFO | __main__ |   val_n_objs = 10
2021-01-12 19:06:35 | INFO | __main__ |   val_runtime = 3.5329

No podemos comparar estos resultados con la línea base, ya que la línea base ni siquiera se inicia y falla de inmediato con OOM.

¡Simplemente asombroso!

Solo utilicé una muestra pequeña, ya que estaba principalmente interesado en poder entrenar y evaluar con este modelo enorme que normalmente no cabría en una GPU de 24 GB.

Si deseas experimentar con esta prueba tú mismo o quieres conocer más detalles sobre el hardware y software utilizados para ejecutarla, por favor, consulta esta publicación.

Dado que transformers solo integra estas soluciones fabulosas y no formó parte de su invención, compartiré los recursos donde puedes descubrir todos los detalles por ti mismo. Pero aquí tienes algunos conocimientos rápidos que pueden ayudarte a comprender cómo ZeRO gestiona estos logros asombrosos.

La característica clave de ZeRO es agregar almacenamiento de datos distribuido al concepto bastante familiar de entrenamiento paralelo de datos.

La computación en cada GPU es exactamente la misma que en el entrenamiento paralelo de datos, pero los parámetros, gradientes y estados del optimizador se almacenan de manera distribuida/particionada en todas las GPUs y se recuperan solo cuando es necesario.

El siguiente diagrama, que proviene de esta publicación del blog, ilustra cómo funciona esto:

El enfoque ingenioso de ZeRO es particionar los parámetros, gradientes y estados del optimizador de manera equitativa en todas las GPUs y dar a cada GPU solo una partición (también llamada fragmento). Esto lleva a cero superposición en el almacenamiento de datos entre las GPUs. En tiempo de ejecución, cada GPU construye los datos de cada capa sobre la marcha pidiendo a las GPUs participantes que envíen la información que le falta.

Esta idea puede ser difícil de entender, y encontrarás mi intento de explicación aquí.

A partir de esta escritura, FairScale y DeepSpeed solo realizan particionamiento (fragmentación) para los estados del optimizador y los gradientes. Se supone que la fragmentación de los parámetros del modelo llegará pronto en DeepSpeed y FairScale.

Otra característica poderosa es ZeRO-Offload (paper). Esta función descarga parte del procesamiento y las necesidades de memoria en la CPU del host, lo que permite ajustar más en la GPU. Viste su impacto dramático en el éxito de ejecutar t5-3b en una GPU de 24GB.

Otro problema del que mucha gente se queja en los foros de PyTorch es la fragmentación de la memoria de la GPU. A menudo, se produce un error de OOM que puede verse así:

RuntimeError: CUDA sin memoria. Se intentó asignar 1.48 GiB (GPU 0; capacidad total de 23.65 GiB;
16.22 GiB ya asignados; 111.12 MiB libres; 22.52 GiB reservados en total por PyTorch)

El programa desea asignar ~1.5GB y la GPU aún tiene unos 6-7GB de memoria no utilizada, pero informa tener solo ~100MB de memoria libre contigua y falla con el error de OOM. Esto ocurre cuando se asignan y desasignan fragmentos de diferentes tamaños una y otra vez, y con el tiempo se crean agujeros que generan fragmentación de memoria, donde hay mucha memoria no utilizada pero no hay fragmentos contiguos del tamaño deseado. En el ejemplo anterior, el programa probablemente podría asignar 100MB de memoria contigua, pero claramente no puede obtener 1.5GB en un solo fragmento.

DeepSpeed aborda este problema gestionando la memoria de la GPU por sí misma y asegurando que las asignaciones de memoria a largo plazo no se mezclen con las de corto plazo y, por lo tanto, haya mucha menos fragmentación. Si bien el artículo no entra en detalles, el código fuente está disponible, por lo que es posible ver cómo DeepSpeed logra esto.

Como ZeRO significa Optimizador sin Redundancia Cero, es fácil ver que cumple con su nombre.

Además del soporte anticipado para la fragmentación de parámetros del modelo en DeepSpeed, ya se han lanzado nuevas características que aún no hemos explorado. Estas incluyen Atención Escasa de DeepSpeed y 1-bit Adam, que se supone que disminuyen el uso de memoria y reducen drásticamente la sobrecarga de comunicación entre GPU, lo que debería llevar a un entrenamiento aún más rápido y admitir modelos aún más grandes.

Confío en que también veremos nuevas aportaciones del equipo de FairScale. Creo que también están trabajando en ZeRO etapa 3.

Aún más emocionante, ZeRO se está integrando en PyTorch.

Si los resultados compartidos en esta publicación del blog te parecen interesantes, por favor procede aquí para obtener detalles sobre cómo usar DeepSpeed y FairScale con el Entrenador de transformers.

Por supuesto, puedes modificar tu propio entrenador para integrar DeepSpeed y FairScale, según las instrucciones de cada proyecto, o puedes “hacer trampa” y ver cómo lo hicimos en el Entrenador de transformers. Si optas por esto último, para orientarte, busca en el código fuente deepspeed y/o sharded_ddp.

La buena noticia es que ZeRO no requiere modificaciones en el modelo. Las únicas modificaciones requeridas están en el código de entrenamiento.

Si encuentras algún problema con la integración de alguno de estos proyectos, por favor, abre un problema en transformers.

Pero si tienes problemas con la instalación, configuración y implementación de DeepSpeed y FairScale, debes consultar a los expertos en sus dominios, por lo tanto, por favor, usa DeepSpeed Issue o FairScale Issue en su lugar.

Mientras no sea necesario entender cómo funcionan ninguno de estos proyectos y simplemente puedas implementarlos a través del transformers Trainer, si deseas comprender los porqués y los cómos, consulta los siguientes recursos.

  • FairScale GitHub

  • DeepSpeed GitHub

  • Artículo: ZeRO: Optimizaciones de memoria hacia el entrenamiento de modelos de billones de parámetros. El artículo es muy interesante, pero es muy conciso.

  • Aquí tienes un buen video que discute el artículo con imágenes

  • Artículo: ZeRO-Offload: Democratizando el entrenamiento de modelos a escala de miles de millones. Recién publicado, este artículo profundiza en los detalles de la función ZeRO Offload.

  • Configuración y tutoriales de DeepSpeed

  • Además del artículo, recomiendo encarecidamente leer las siguientes publicaciones de blog detalladas con diagramas:

    • DeepSpeed: Entrenamiento de modelos a escala extrema para todos
    • ZeRO & DeepSpeed: Nuevas optimizaciones del sistema permiten entrenar modelos con más de 100 mil millones de parámetros
    • Turing-NLG: Un modelo de lenguaje de 17 mil millones de parámetros de Microsoft
  • Ejemplos de DeepSpeed en GitHub

Quedamos bastante asombrados con el increíble nivel de apoyo que recibimos de los equipos de desarrollo de FairScale y DeepSpeed mientras trabajábamos en la integración de esos proyectos en transformers.

En particular, me gustaría agradecer a:

  • Benjamin Lefaudeux @blefaudeux
  • Mandeep Baines @msbaines

del equipo de FairScale y:

  • Jeff Rasley @jeffra
  • Olatunji Ruwase @tjruwase
  • Samyam Rajbhandari @samyam

del equipo de DeepSpeed por su generoso y atento apoyo y resolución rápida de los problemas que hemos encontrado.

Y a HuggingFace por proporcionar acceso al hardware en el que se ejecutaron las pruebas de rendimiento.

Sylvain Gugger @sgugger y Stas Bekman @stas00 trabajaron en la integración de estos proyectos.

We will continue to update Zepes; if you have any questions or suggestions, please contact us!

Share:

Was this article helpful?

93 out of 132 found this helpful

Discover more

Inteligencia Artificial

Eso es gracioso, pero los modelos de IA no entienden la broma.

Un equipo multi-institucional de científicos probó la capacidad de los modelos de inteligencia artificial para entend...

Inteligencia Artificial

Conoce MovieChat un innovador sistema de comprensión de video que integra modelos fundamentales de video y grandes modelos de lenguaje.

Los Modelos de Lenguaje Grande (LLMs, por sus siglas en inglés) han avanzado considerablemente en el sector de Proces...

Inteligencia Artificial

Creando empatía artificial

Los científicos se esfuerzan por añadir la habilidad de entender lo que los demás sienten a la inteligencia artificial.

Inteligencia Artificial

Luma AI lanza Genie un nuevo modelo de IA generativa en 3D que te permite crear objetos en 3D a partir de texto.

En el modelado 3D, crear objetos 3D realistas a menudo ha sido una tarea compleja y que consume mucho tiempo. Las per...

Inteligencia Artificial

Investigadores de UC Berkeley presentan Nerfstudio un marco de trabajo en Python para el desarrollo de Neural Radiance Field (NeRF)

¿Quién no es fan de Iron Man? Se ve realmente genial cuando está trabajando en su laboratorio. Todos los hologramas y...