Ajuste fino de 20B LLMs con RLHF en una GPU de consumo de 24GB

20B LLMs con RLHF ajustados en una GPU de 24GB.

¡Estamos emocionados de lanzar oficialmente la integración de trl con peft para hacer que el ajuste fino del Modelo de Lenguaje Grande (LLM) con Aprendizaje por Reforzamiento sea más accesible para todos! En esta publicación, explicamos por qué esta es una alternativa competitiva a los enfoques de ajuste fino existentes.

Ten en cuenta que peft es una herramienta general que se puede aplicar a muchos casos de uso de ML, pero es particularmente interesante para RLHF, ¡ya que este método requiere mucha memoria!

Si quieres sumergirte directamente en el código, echa un vistazo a los scripts de ejemplo directamente en la página de documentación de TRL.

Introducción

LLMs y RLHF

LLMs combinados con RLHF (Aprendizaje por Reforzamiento con Retroalimentación Humana) parecen ser el enfoque preferido para construir sistemas de IA muy poderosos, como ChatGPT.

El entrenamiento de un modelo de lenguaje con RLHF generalmente implica los siguientes tres pasos:

• Pronóstico de series de tiempo probabilísticas multivariadas con In...
• Jupyter X Hugging Face
• Entrena tu ControlNet con difusores

1- Ajustar finamente un LLM preentrenado en un dominio específico o corpus de instrucciones y demostraciones humanas

2- Recopilar un conjunto de datos anotados por humanos y entrenar un modelo de recompensa

3- Ajustar finamente aún más el LLM del paso 1 con el modelo de recompensa y este conjunto de datos utilizando RL (por ejemplo, PPO)

La elección del LLM base es bastante crucial aquí. En este momento de la escritura, los LLM ajustados a instrucciones son los mejores LLM de código abierto que se pueden usar “tal cual” para muchas tareas. Algunos modelos destacados son: BLOOMZ, Flan-T5, Flan-UL2 y OPT-IML. El inconveniente de estos modelos es su tamaño. Para obtener un modelo decente, necesitas al menos usar modelos a una escala de 10B+ que requerirían hasta 40 GB de memoria GPU en precisión completa, ¡solo para ajustar el modelo en un solo dispositivo GPU sin realizar ningún entrenamiento en absoluto!

¿Qué es TRL?

La biblioteca trl tiene como objetivo facilitar y hacer más flexible el paso de RL para que cualquier persona pueda ajustar finamente su LM usando RL en su conjunto de datos personalizado y configuración de entrenamiento. Entre muchas otras aplicaciones, puedes usar este algoritmo para ajustar finamente un modelo para generar críticas positivas de películas, generar de manera controlada o hacer que el modelo sea menos tóxico.

Usando trl, puedes ejecutar uno de los algoritmos de RL profundo más populares, PPO, de manera distribuida o en un solo dispositivo. Aprovechamos accelerate del ecosistema de Hugging Face para hacer esto posible, para que cualquier usuario pueda escalar los experimentos hasta una escala interesante.

Ajustar finamente un modelo de lenguaje con RL sigue aproximadamente el siguiente protocolo detallado a continuación. Esto requiere tener 2 copias del modelo original; para evitar que el modelo activo se desvíe demasiado de su comportamiento / distribución original, necesitas calcular los logits del modelo de referencia en cada paso de optimización. Esto agrega una restricción difícil al proceso de optimización, ya que siempre necesitas al menos dos copias del modelo por dispositivo GPU. Si el modelo crece en tamaño, se vuelve cada vez más complicado ajustar la configuración en una sola GPU.

En trl, también puedes usar capas compartidas entre los modelos de referencia y activo para evitar copias completas. Un ejemplo concreto de esta característica se muestra en el ejemplo de desintoxicación.

Entrenamiento a gran escala

Entrenar a gran escala puede ser un desafío. El primer desafío es ajustar el modelo y los estados de su optimizador en los dispositivos GPU disponibles. La cantidad de memoria GPU que ocupa un solo parámetro depende de su “precisión” (o más específicamente, dtype). Los dtype más comunes son float32 (32 bits), float16 y bfloat16 (16 bits). Más recientemente, se admiten precisiones “exóticas” de forma nativa para entrenamiento e inferencia (con ciertas condiciones y restricciones), como int8 (8 bits). En resumen, para cargar un modelo en un dispositivo GPU, cada mil millones de parámetros cuesta 4 GB en precisión float32, 2 GB en float16 y 1 GB en int8. Si deseas obtener más información sobre este tema, echa un vistazo a esta publicación de blog que profundiza: https://huggingface.co/blog/hf-bitsandbytes-integration.

Si usas un optimizador AdamW, cada parámetro necesita 8 bytes (por ejemplo, si tu modelo tiene 1B de parámetros, el optimizador completo de AdamW del modelo requeriría 8GB de memoria GPU – fuente).

Se han adoptado muchas técnicas para abordar estos desafíos a gran escala. Los paradigmas más familiares son el Paralelismo de Tuberías, el Paralelismo de Tensores y el Paralelismo de Datos.

Con el paralelismo de datos, el mismo modelo se aloja en paralelo en varias máquinas y cada instancia recibe un lote de datos diferente. Esta es la estrategia de paralelismo más directa, esencialmente replicando el caso de una sola GPU, y ya es compatible con trl. Con el Paralelismo de Tuberías y el Paralelismo de Tensores, el modelo en sí se distribuye en varias máquinas: en el Paralelismo de Tuberías, el modelo se divide por capas, mientras que el Paralelismo de Tensores divide las operaciones de tensores entre las GPU (por ejemplo, multiplicaciones de matrices). Con estas estrategias de Paralelismo de Modelos, es necesario dividir los pesos del modelo entre muchos dispositivos, lo que requiere definir un protocolo de comunicación de las activaciones y gradientes entre los procesos. Esto no es trivial de implementar y puede requerir la adopción de algunos marcos de trabajo como Megatron-DeepSpeed o Nemo. También es importante resaltar otras herramientas que son esenciales para escalar el entrenamiento de modelos de lenguaje como el Punto de Control de Activación Adaptativo y los núcleos fusionados. Se puede encontrar más información sobre los paradigmas de paralelismo aquí.

Por lo tanto, nos hicimos la siguiente pregunta: ¿hasta dónde podemos llegar solo con el paralelismo de datos? ¿Podemos usar herramientas existentes para adaptar procesos de entrenamiento super grandes (incluyendo el modelo activo, el modelo de referencia y los estados del optimizador) en un solo dispositivo? La respuesta parece ser sí. Los ingredientes principales son: adaptadores y multiplicación de matrices de 8 bits. Cubriremos estos temas en las siguientes secciones:

Multiplicación de matrices de 8 bits

La multiplicación de matrices de 8 bits eficiente es un método que se introdujo por primera vez en el artículo LLM.int8() y tiene como objetivo resolver el problema de degradación del rendimiento al cuantizar modelos a gran escala. El método propuesto descompone las multiplicaciones de matrices que se aplican en capas lineales en dos etapas: la parte de estados ocultos atípicos que se va a realizar en float16 y la parte “no atípica” que se realiza en int8.

En pocas palabras, se puede reducir el tamaño de un modelo de precisión completa en 4 (por lo tanto, en 2 para modelos de media precisión) si se utiliza la multiplicación de matrices de 8 bits.

Adaptación de rango bajo y PEFT

En 2021, un artículo llamado LoRA: Low-Rank Adaption of Large Language Models demostró que se puede realizar el ajuste fino de modelos de lenguaje grandes congelando los pesos preentrenados y creando versiones de rango bajo de las matrices de atención de las capas de consulta y valor. Estas matrices de rango bajo tienen muchos menos parámetros que el modelo original, lo que permite el ajuste fino con mucha menos memoria de GPU. Los autores demuestran que el ajuste fino de los adaptadores de rango bajo logró resultados comparables al ajuste fino del modelo preentrenado completo.

Esta técnica permite el ajuste fino de LLMs utilizando una fracción de los requisitos de memoria. Sin embargo, hay algunas desventajas. El pase hacia adelante y hacia atrás es aproximadamente el doble de lento debido a las multiplicaciones de matrices adicionales en las capas de adaptadores.

¿Qué es PEFT?

Parameter-Efficient Fine-Tuning (PEFT) es una biblioteca de Hugging Face, creada para admitir la creación y el ajuste fino de capas de adaptadores en LLMs. peft está integrado de manera transparente con 🤗 Accelerate para modelos a gran escala que aprovechan DeepSpeed y Big Model Inference.

La biblioteca admite muchos modelos de vanguardia y tiene un conjunto extenso de ejemplos, que incluyen:

• Modelado del lenguaje causal
• Generación condicional
• Clasificación de imágenes
• Entrenamiento de 8 bits int8
• Adaptación de rango bajo de modelos Dreambooth
• Segmentación semántica
• Clasificación de secuencias
• Clasificación de tokens

La biblioteca aún está en desarrollo extenso y activo, con muchas características próximas que se anunciarán en los próximos meses.

Ajuste fino de modelos de 20B parámetros con adaptadores de rango bajo

Ahora que se han cubierto los requisitos previos, repasemos todo el proceso paso a paso y expliquemos con figuras cómo se puede ajustar finamente un LLM de 20B parámetros con RL utilizando las herramientas mencionadas anteriormente en una sola GPU de 24GB.

Paso 1: Cargar su modelo activo en precisión de 8 bits

Una reducción de memoria “gratis” de un LLM utilizando transformers es cargar su modelo en precisión de 8 bits utilizando el método descrito en LLM.int8. Esto se puede hacer simplemente agregando el indicador load_in_8bit=True al llamar al método from_pretrained (puede leer más al respecto aquí).

Como se menciona en la sección anterior, un “truco” para calcular la cantidad de memoria GPU que necesitaría para cargar su modelo es pensar en términos de “miles de millones de parámetros”. Dado que un byte necesita 8 bits, necesita 4 GB por cada mil millones de parámetros para un modelo de precisión completa (32 bits = 4 bytes), 2 GB por cada mil millones de parámetros para un modelo de media precisión, y 1 GB por cada mil millones de parámetros para un modelo int8.

Entonces, en primer lugar, simplemente carguemos el modelo activo en 8 bits. ¡Veamos qué necesitamos hacer para el segundo paso!

Paso 2: Agregar adaptadores adicionales entrenables usando peft

El segundo paso es cargar adaptadores dentro del modelo y hacer que estos adaptadores sean entrenables. Esto permite una reducción drástica en la cantidad de pesos entrenables que se necesitan para el modelo activo. Este paso aprovecha la biblioteca peft y se puede realizar con algunas líneas de código. Tenga en cuenta que una vez que los adaptadores están entrenados, puede enviarlos fácilmente al Hub para usarlos más adelante.

Paso 3: Usar el mismo modelo para obtener los logits de referencia y activos

Dado que los adaptadores pueden desactivarse, podemos usar el mismo modelo para obtener los logits de referencia y activos para PPO, ¡sin tener que crear dos copias del mismo modelo! Esto aprovecha una característica en la biblioteca peft, que es el administrador de contexto disable_adapters.

Resumen de los scripts de entrenamiento:

Ahora describiremos cómo entrenamos un modelo gpt-neox de 20 mil millones de parámetros utilizando transformers, peft y trl. El objetivo final de este ejemplo era ajustar finamente un LLM para generar críticas positivas de películas en un entorno con restricciones de memoria. Se podrían aplicar pasos similares para otras tareas, como modelos de diálogo.

En general, hubo tres pasos clave y scripts de entrenamiento:

1. Script – Ajuste fino de un Adaptador de Rango Bajo en un modelo de 8 bits congelado para generación de texto en el conjunto de datos imdb.
2. Script – Fusión de las capas de adaptador en los pesos del modelo base y almacenamiento de estos en el Hub.
3. Script – Ajuste fino de sentimiento de un Adaptador de Rango Bajo para crear críticas positivas.

Probamos estos pasos en una GPU NVIDIA 4090 de 24 GB. Si bien es posible realizar toda la ejecución de entrenamiento en una GPU de 24 GB, las ejecuciones completas de entrenamiento se realizaron en un solo A100 en el clúster de investigación 🤗.

El primer paso en el proceso de entrenamiento fue el ajuste fino del modelo preentrenado. Normalmente esto requeriría varias GPU A100 de 80 GB de gama alta, por lo que optamos por entrenar un adaptador de rango bajo. Tratamos esto como un ajuste fino de modelado de lenguaje causal y entrenamos durante una época de ejemplos del conjunto de datos imdb, que presenta críticas de películas y etiquetas que indican si son de sentimiento positivo o negativo.

Para llevar el modelo adaptado y realizar un ajuste fino adicional con RL, primero necesitábamos combinar los pesos adaptados, esto se logró cargando el modelo preentrenado y el adaptador en punto flotante de 16 bits y resumiendo con matrices de pesos (con la escala adecuada aplicada).

Finalmente, pudimos ajustar finamente otro adaptador de rango bajo, sobre el modelo imdb-finetuned congelado. Utilizamos un clasificador de sentimiento de imdb para proporcionar las recompensas para el algoritmo de RL.

El informe completo de Pesos y Sesgos está disponible para este experimento aquí, si desea consultar más gráficos y generaciones de texto.

Conclusión

Hemos implementado una nueva funcionalidad en trl que permite a los usuarios ajustar finamente modelos de lenguaje grandes utilizando RLHF a un costo razonable aprovechando las bibliotecas peft y bitsandbytes. Hemos demostrado que es posible ajustar finamente gpt-neo-x (¡40 GB en bfloat16!) en una GPU de consumo de 24 GB, y esperamos que esta integración sea ampliamente utilizada por la comunidad para ajustar finamente modelos más grandes utilizando RLHF y compartir artefactos excelentes.

Hemos identificado algunas direcciones interesantes para los próximos pasos para empujar los límites de esta integración.

• ¿Cómo escalará esto en el entorno multi-GPU? Principalmente exploraremos cómo esta integración escalará en relación al número de GPUs, si es posible aplicar la Paralelización de Datos directamente o si requerirá la adopción de alguna nueva característica en alguna de las bibliotecas involucradas.
• ¿Qué herramientas podemos aprovechar para aumentar la velocidad de entrenamiento? Hemos observado que la principal desventaja de esta integración es la velocidad general de entrenamiento. En el futuro estaríamos interesados en explorar las posibles direcciones para hacer el entrenamiento mucho más rápido.

Referencias

• Paradigmas de paralelismo: https://huggingface.co/docs/transformers/v4.17.0/es/parallelism
• Integración de 8 bits en transformers: https://huggingface.co/blog/hf-bitsandbytes-integration
• Documento LLM.int8: https://arxiv.org/abs/2208.07339
• Explicación de la comprobación de gradientes: https://docs.aws.amazon.com/sagemaker/latest/dg/model-parallel-extended-features-pytorch-activation-checkpointing.html