2024-05-07 11:31:24
Las técnicas de optimización de memoria como la cuantización y el ajuste fino distribuido están democratizando los grandes modelos de lenguaje, permitiendo su implementación en entornos con recursos limitados y facilitando su entrenamiento y ajuste fino.
En el mundo de los modelos de lenguaje de gran escala (LLM), el consumo de memoria es un desafío crucial. Estos modelos gigantescos, con miles de millones de parámetros, requieren una enorme cantidad de memoria para su inferencia y ajuste fino. Afortunadamente, las técnicas de optimización de memoria como la cuantización y el entrenamiento distribuido están permitiendo que estos LLM sean más accesibles y escalables.
La cuantización es una técnica poderosa que reduce drásticamente los requisitos de memoria al representar los parámetros del modelo con una menor precisión numérica. Por ejemplo, con la cuantización de 8 bits, la memoria necesaria para el modelo OctoCode de 15 mil millones de parámetros se reduce de 31 GB a solo 15 GB. ¡Y con la cuantización de 4 bits, la huella de memoria se reduce aún más a 9.5 GB! Esto permite ejecutar este modelo masivo en GPUs de consumo como la RTX 3090 con 24 GB de VRAM.
Sin embargo, es importante tener en cuenta que una cuantización más agresiva puede degradar la precisión del modelo. Existe un equilibrio entre el ahorro de memoria y la precisión que los usuarios deben evaluar según sus necesidades.
La cuantización es una técnica poderosa que puede permitir la implementación de LLM en entornos con recursos limitados como instancias en la nube, dispositivos periféricos o incluso teléfonos móviles al reducir drásticamente la huella de memoria.
Ajuste Fino Distribuido para Modelos Masivos
Mientras que la cuantización es clave para la inferencia eficiente, el ajuste fino de LLM requiere técnicas adicionales como el paralelismo de tensores y el paralelismo de modelos para gestionar los requisitos de memoria durante el entrenamiento.
Durante el ajuste fino, el consumo máximo de memoria suele ser de 3 a 4 veces mayor que en la inferencia debido a los requisitos adicionales de memoria para gradientes, estados del optimizador y activaciones almacenadas.
- Gradients
- Estados del optimizador
- Activaciones almacenadas del pase hacia adelante para la retropropagación
Una estimación conservadora es que el ajuste fino de un LLM con X mil millones de parámetros requiere alrededor de 4 * (2X) = 8X GB de VRAM en precisión bfloat16.
Por ejemplo, el ajuste fino del modelo LLaMA de 7B parámetros requeriría aproximadamente 7 * 8 = 56GB de VRAM por GPU en precisión bfloat16. Esto supera la capacidad de memoria de las GPUs actuales, lo que hace necesarias técnicas de ajuste fino distribuido.
Técnicas de Ajuste Fino Distribuido
Se han propuesto varios métodos de ajuste fino distribuido para superar las limitaciones de memoria de la GPU:
- Paralelismo de Datos: Replica el modelo en múltiples GPUs y distribuye los lotes de entrenamiento, reduciendo el tiempo de entrenamiento pero no el requisito de memoria máximo.
- ZeRO Stage 3: Una forma avanzada de paralelismo de datos que divide los parámetros, gradientes y estados del optimizador entre las GPUs, reduciendo la memoria en comparación con el paralelismo de datos clásico.
- Paralelismo de Tensores: Divide los parámetros del modelo en filas o columnas y los distribuye entre las GPUs, con cada GPU operando en un conjunto particionado de parámetros, gradientes y estados del optimizador.
- Paralelismo de Pipelines: Divide las capas del modelo entre diferentes GPUs/trabajadores, con cada dispositivo ejecutando un subconjunto de las capas y pasando las activaciones entre ellos.
Estimar el uso de memoria para estos métodos distribuidos no es trivial, ya que la distribución de parámetros, gradientes, activaciones y estados del optimizador varía entre técnicas. Además, diferentes componentes como el cuerpo del transformador y la cabeza de modelado de lenguaje pueden mostrar diferentes comportamientos de asignación de memoria.
La Solución LLMem
Recientemente, los investigadores propusieron LLMem, una solución que estima con precisión el consumo de memoria de la GPU al aplicar métodos de ajuste fino distribuido a LLMs en múltiples GPUs.
LLMem considera factores como la recombinación de parámetros antes de la computación, la recolección de salidas en el pase hacia atrás y las diferentes estrategias de asignación de memoria para el cuerpo del transformador y la cabeza de modelado de lenguaje.
Los resultados experimentales muestran que LLMem puede estimar el uso máximo de memoria de la GPU para el ajuste fino de LLM en una sola GPU con tasas de error de hasta 1,6%, superando la tasa de error promedio de DNNMem, líder en el sector, del 42,6%. Al aplicar métodos de ajuste fino distribuido a LLMs con más de mil millones de parámetros en múltiples GPUs, LLMem logra una impresionante tasa de error promedio del 3,0%
