Axolotl - Herramienta de Post-Entrenamiento Potente para Modelos de IA

Resumen del Proyecto

Axolotl es una herramienta de código abierto diseñada específicamente para simplificar el proceso de post-entrenamiento de varios modelos de IA. El post-entrenamiento se refiere a cualquier modificación o entrenamiento adicional realizado en un modelo pre-entrenado, incluyendo el ajuste fino completo del modelo, el ajuste eficiente de parámetros (como LoRA y QLoRA), el ajuste fino supervisado (SFT), el ajuste de instrucciones y las técnicas de alineación. Esta herramienta es compatible con diversas arquitecturas de modelos y configuraciones de entrenamiento, lo que permite a los usuarios comenzar fácilmente a utilizar estas avanzadas técnicas de entrenamiento.

Características Principales

Soporte de Modelos

• Arquitecturas de Modelos Diversas: Soporta el entrenamiento de varios modelos de Hugging Face, incluyendo LLaMA, Pythia, Falcon, MPT, Mistral, Mixtral y otros modelos de lenguaje grandes populares.
• Métodos de Entrenamiento Flexibles: Soporta ajuste fino completo (Full Fine-tuning), LoRA, QLoRA, ReLoRA, GPTQ y otros métodos de entrenamiento.

Gestión de Configuración

• Archivos de Configuración YAML: Utiliza archivos YAML simples que contienen todas las configuraciones necesarias para el preprocesamiento de datos, el entrenamiento/ajuste fino del modelo, la inferencia o la evaluación del modelo.
• Anulación por CLI: Permite anular la configuración en el archivo de configuración a través de parámetros de línea de comandos.
• Configuración Flexible: Permite personalizar varios parámetros de entrenamiento y configuraciones del modelo.

Capacidad de Procesamiento de Datos

• Conjuntos de Datos Multiformato: Soporta la carga de conjuntos de datos locales, de HuggingFace y en la nube (S3, Azure, GCP, OCI).
• Formato Personalizado: Permite utilizar formatos personalizados o importar directamente conjuntos de datos ya tokenizados.
• Preprocesamiento de Conjuntos de Datos: Incorpora potentes funciones de preprocesamiento de datos.

Optimización del Rendimiento

• Técnicas de Optimización Avanzadas: Integra xformers, Flash Attention, Liger Kernel, escalado de cuerda y técnicas de multi-empaquetado.
• Soporte Multi-GPU: Soporta entrenamiento con una sola GPU o con múltiples GPUs a través de FSDP o DeepSpeed.
• Entrenamiento Eficiente: Optimizado para GPUs NVIDIA (Ampere o versiones más recientes, con soporte para bf16 y Flash Attention) y GPUs AMD.

Despliegue y Monitorización

• Listo para la Nube: Proporciona imágenes Docker y paquetes PyPI, que se pueden utilizar en plataformas en la nube y hardware local.
• Registro de Resultados: Soporta el registro de resultados y puntos de control en WandB, MLflow o Comet.
• Soporte de Monitorización: Integra varias herramientas de seguimiento y monitorización de experimentos.

Requisitos Técnicos

Requisitos de Hardware

• GPU NVIDIA (Ampere o versiones más recientes, para bf16 y Flash Attention) o GPU AMD.
• Suficiente memoria de GPU para el entrenamiento del modelo.

Requisitos de Software

• Python 3.11
• PyTorch ≥2.4.1
• Paquetes de dependencia relacionados.

Instalación

Instalación Rápida

pip3 install -U packaging==23.2 setuptools==75.8.0 wheel ninja
pip3 install --no-build-isolation axolotl[flash-attn,deepspeed]

# Descargar archivos de configuración de ejemplo
axolotl fetch examples
axolotl fetch deepspeed_configs  # Opcional

Instalación desde el Código Fuente

git clone https://github.com/axolotl-ai-cloud/axolotl.git
cd axolotl
pip3 install -U packaging setuptools wheel ninja
pip3 install --no-build-isolation -e '.[flash-attn,deepspeed]'

Vía Docker

docker run --gpus '"all"' --rm -it axolotlai/axolotl:main-latest

Uso

Flujo de Trabajo Básico

1. Obtener la Configuración de Ejemplo:

   axolotl fetch examples
   

2. Entrenar el Modelo:

   axolotl train examples/llama-3/lora-1b.yml
   

3. Personalizar la Configuración: Modificar los parámetros en el archivo de configuración YAML según sea necesario.

Estructura del Archivo de Configuración

Axolotl utiliza archivos de configuración YAML para controlar todo el flujo de entrenamiento, incluyendo:

• Selección y parámetros del modelo
• Configuración y preprocesamiento del conjunto de datos
• Hiperparámetros de entrenamiento
• Configuración del optimizador
• Monitorización y registro

Matriz de Compatibilidad de Modelos Soportados

Modelo	fp16/fp32	LoRA	QLoRA	GPTQ	Flash Attn	xformers
LLaMA	✅	✅	✅	✅	✅	✅
Mistral	✅	✅	✅	✅	✅	✅
Mixtral-MoE	✅	✅	✅	❓	❓	❓
Pythia	✅	✅	✅	❌	❌	❌
Falcon	✅	✅	✅	❌	❌	❌
Qwen	✅	✅	✅	❓	❓	❓
Gemma	✅	✅	✅	❓	❓	✅

✅: Soportado ❌: No Soportado ❓: No Probado

Escenarios de Aplicación

Campo de Investigación

• Investigación sobre el ajuste fino de modelos de lenguaje grandes
• Experimentación con métodos de entrenamiento eficientes en parámetros
• Investigación sobre la alineación y seguridad de modelos

Aplicaciones Industriales

• Personalización de modelos a nivel empresarial
• Entrenamiento de modelos específicos de dominio
• Optimización de modelos para funcionalidades de productos

Formación Educativa

• Enseñanza de cursos de IA/ML
• Desarrollo de proyectos prácticos
• Mejora de habilidades técnicas

Ventajas del Proyecto

1. Facilidad de Uso: Controla flujos de entrenamiento complejos a través de archivos de configuración YAML simples.
2. Flexibilidad: Soporta diversas arquitecturas de modelos y métodos de entrenamiento.
3. Rendimiento: Integra las últimas técnicas de optimización, proporcionando una experiencia de entrenamiento eficiente.
4. Escalabilidad: Soporta varias escalas de entrenamiento, desde una sola GPU hasta múltiples nodos.
5. Código Abierto: Licencia Apache 2.0, completamente de código abierto y de uso gratuito.