Cómo afinar un modelo de lenguaje de 70B en casa

• Open source que combina FSDP (Fully Sharded Data Parallel) y QLoRA (Quantization + Low-Rank Adaptation)
• Permite ajustar eficientemente un modelo de lenguaje de tamaño 70B con una computadora de escritorio común y dos GPU gaming estándar

Contexto

• Aunque no hay una gran diferencia de rendimiento entre el hardware de nivel centro de datos y las GPU gaming instaladas en computadoras de escritorio, la diferencia de precio es enorme
• Las GPU gaming tienen poca memoria, por lo que es difícil entrenar modelos de lenguaje de gran escala
• Answer.AI busca resolver este problema encontrando una forma de entrenar los mejores modelos open source usando GPU gaming económicas

QLoRA: entrenar modelos más grandes en una sola GPU

• QLoRA combina cuantización y LoRA para reducir los pesos del modelo a 4 bits, disminuyendo el uso de memoria de la GPU.
• LoRA evita las limitaciones de memoria al no entrenar todo el modelo de lenguaje de gran escala, sino agregando y entrenando pequeñas matrices adaptadoras.
• QLoRA permite entrenar un modelo de 65B en una tarjeta de 48GB, pero sigue siendo insuficiente para entrenar un modelo de 70B en una GPU gaming de 24GB.

FSDP: ampliar la escala de entrenamiento con varias GPU

• FSDP (Fully Sharded Data Parallel) distribuye los parámetros del modelo entre varias GPU para poder usarlas al mismo tiempo.
• FSDP permite aplicar el rendimiento de DDP (Distributed Data Parallel) a modelos demasiado grandes para caber en una sola GPU.
• Por ejemplo, un modelo no cuantizado de 70B (70 mil millones de parámetros) usa 140GB de RAM (porque cada parámetro se almacena en 16 bits, es decir, 2 bytes), por lo que incluso una tarjeta NVIDIA H100 (¡con un costo de unos $40,000 por tarjeta!) se queda corta. Se necesitan 80GB de RAM
  • Sin embargo, con FSDP se pueden combinar 4 GPU H100 para obtener un total de 320GB de RAM. Y ahorrar $150000

Combinar FSDP y QLoRA

• Answer.AI reduce el tamaño del modelo con QLoRA y luego lo distribuye en varias tarjetas gaming de 24GB con FSDP, permitiendo un entrenamiento eficiente.
• Gracias a esto, ahora es posible entrenar un modelo de 70B con dos GPU gaming comunes.
• A fines de 2023 se discutió la idea de combinar FSDP y QLoRA, y luego se resolvieron los problemas necesarios para hacerlo posible.
• En este proceso se combinaron y utilizaron varias tecnologías y librerías
  • Almacenar los parámetros cuantizados en el mismo tipo de dato que el "Computation Type" del modelo
  • Gradient checkpointing, que en lugar de guardar todos los gradientes almacena activaciones en varios "checkpoints" del modelo completo y luego vuelve a ejecutar el paso forward según sea necesario para recalcular los gradientes
  • CPU offloading, que guarda los pesos en la RAM de la CPU cuando no se usan en lugar de mantenerlos en la GPU, reduciendo de forma importante la memoria de GPU necesaria
  • Flash Attention 2: calcula la atención de manera eficiente usando kernels de Cuda optimizados para memoria
  • Descubrieron HQQ (Half-Quadratic Quantization) y lo integraron con FSDP (50 veces más rápido y preciso que GPTQ)

Cómo usar FSDP/QLoRA

• Para usar FSDP se necesitan dos o más GPU. (En Runpod hay cajas dual 3090 disponibles por $0.6 por hora)
• Instala la última versión de Transformers, PEFT, Bitsandbytes (y HQQ si lo vas a usar), clona el repositorio de Answer.AI y sigue el README para ejecutarlo.
• Para entrenar Alpaca con 2 tarjetas de 24GB:
  python train.py --train_type qlora --dataset alpaca --batch_size 8 --gradient_accumulation_steps 2 --output_dir qlora_output --log_to wandb
• Todavía está en una etapa inicial (alpha/preview)
• También planean publicar resultados de benchmarking en unas semanas

Opinión de GN⁺

• Este artículo será muy interesante para quienes trabajan en IA y machine learning. En especial, destaca porque propone una nueva forma de entrenar eficientemente modelos de lenguaje de gran escala.
• Si esta tecnología logra adoptarse ampliamente, podría reducir de manera importante la barrera de costo en la investigación y desarrollo de IA. Eso también podría abrir una gran oportunidad para laboratorios pequeños o desarrolladores individuales.
• Sin embargo, para que este enfoque se adopte de forma generalizada, la estabilidad y la facilidad de uso serán factores clave. Como todavía está en una etapa temprana, el sistema deberá mejorar mediante pruebas y retroalimentación de la comunidad.
• Al introducir una nueva tecnología, siempre hay que considerar la compatibilidad con los sistemas existentes, el rendimiento y el costo. El beneficio de elegir esta tecnología es mejorar la accesibilidad al entrenamiento de modelos grandes, pero al estar aún en una fase inicial, podría tener desventajas en estabilidad y soporte.

[GN⁺]

Comentarios de Hacker News

• El artículo es excelente, pero es una pena que hubo muchas oportunidades de usar la palabra 'nibble' y se desaprovecharon todas.

  • Expresa que el artículo es interesante, pero lamenta que no se haya aprovechado la palabra 'nibble'.

• Como consumidor, me interesaría ver cuánto tardaría en ajustarse finamente un modelo de 70b con un sistema que usa pesos de 4 bits; por ejemplo, una tabla que muestre tamaño del modelo, tamaño de entrenamiento, costo de entrenamiento y la pérdida de calidad esperada con esta técnica.

  • Muestra interés en el costo final del ajuste fino y pide información detallada, incluyendo tamaño del modelo y costo de entrenamiento.

• Me pregunto si se puede aplicar una codificación de (-1,0,1). Si sale ese modelo, pido a answer.ai que use la misma técnica para ajustar finamente modelos grandes con dos GPU pequeñas, y también una evaluación de costo/beneficio.

  • Plantea preguntas sobre la aplicabilidad técnica de un nuevo método de codificación y sobre el análisis de costo-beneficio.

• Este tipo de investigación es muy útil para permitir que los usuarios individuales compitan en el mercado de los LLM. Es un gran trabajo, en línea con lo que plantea el libro 'De cero a uno'.

  • Evalúa positivamente que la investigación sea útil para usuarios individuales y contribuya a resolver un problema específico.

• No se mencionó que todos los primeros lugares en el desafío de eficiencia de NeurIPS usaron QLoRA; es un buen ejemplo de la importancia y utilidad de QLoRA.

  • Destaca la importancia y utilidad de la tecnología QLoRA, mencionando que fue reconocida en esa competencia.

• Es el enfoque más interesante y constructivo que he visto en los LLM de 'código abierto'. Ojalá la academia también lo adoptara como estándar.

  • Valora muy positivamente este enfoque hacia los LLM de código abierto y espera que también se adopte como estándar en la academia.

• Ojalá el entrenamiento rápido en Metal hubiera llegado antes. Por ahora, Mac Studio es mucho más barato y lo suficientemente rápido para inferencia con modelos grandes, pero casi no existen soluciones de entrenamiento.

  • Expresa expectativas sobre entrenamiento rápido en la plataforma Metal y opina sobre la eficiencia actual de Mac Studio.

• La calidad y claridad de la escritura técnica son excelentes.

  • Elogia la calidad y claridad de la documentación técnica.

• Me pregunto si es posible volver a entrenar los pesos de un LLM (70B) con mis propios datos.

  • Pregunta si es posible reentrenar un LLM con datos personales.

• Esto es un avance revolucionario para quienes ajustan finamente LLM con un presupuesto de hardware limitado.

  • Considera que esta tecnología representa un gran avance para usuarios con presupuestos de hardware restringidos.

• Me pregunto por qué eligieron DeepSpeed en lugar de FSDP.

  • Expresa curiosidad sobre la razón para elegir DeepSpeed en vez de FSDP.

• Me pregunto si usaron alguna herramienta para ajustar el rendimiento una vez que lograron hacer fine-tuning del modelo.

  • Pregunta por las herramientas utilizadas para optimizar el rendimiento.

• Me pregunto si han probado usar sparsificación; siento que eso podría permitir el ajuste fino con menos RAM.

  • Opina sobre el uso de técnicas de sparsificación y su posible capacidad para reducir el uso de RAM.