Axolotl Fine-Tuning en Ubuntu Multi-GPU: Guía Paso a Paso para Llama y Mistral

1. Por Qué Axolotl Domina el Ecosistema Open-Source

Entrenar modelos de lenguaje con scripts artesanales de PyTorch implica gestionar manualmente decenas de bibliotecas propensas a conflictos de compatibilidad: tokenizadores, formatos de chat (ChatML, Alpaca, Llama-3), envoltorios de DeepSpeed, hooks de gradient checkpointing y optimizadores paged de 8 bits. Axolotl unifica todo este ecosistema en un motor robusto y probado en batalla por laboratorios de investigación y empresas en todo el mundo.

2. Tutorial Paso a Paso: Instalación en Ubuntu 22.04 / 24.04 LTS

# 1. Crear entorno virtual aislado
python3 -m venv axolotl-env
source axolotl-env/bin/activate

# 2. Instalar PyTorch con soporte CUDA 12.4
pip install --upgrade pip
pip install torch torchvision --extra-index-url https://download.pytorch.org/whl/cu124

# 3. Clonar e instalar Axolotl con dependencias optimizadas
git clone https://github.com/axolotl-ai-cloud/axolotl
cd axolotl
pip install packaging ninja
pip install -e '.[flash-attn,deepspeed]'

3. Configuración Declarativa: config.yaml

A continuación se muestra un archivo de configuración completo para ajustar Llama-3.1-8B-Instruct con LoRA en un servidor con 2 GPUs:

base_model: meta-llama/Meta-Llama-3.1-8B-Instruct
model_type: LlamaForCausalLM
tokenizer_type: AutoTokenizer

load_in_8bit: false
load_in_4bit: false
strict: false

datasets:
  - path: garage-bAInd/Open-Platypus
    type: alpaca

dataset_prepared_path: ./last_run_prepared
val_set_size: 0.05
output_dir: ./lora-llama3-platypus

adapter: lora
lora_r: 32
lora_alpha: 64
lora_dropout: 0.05
lora_target_linear: true

sequence_len: 4096
sample_packing: true
pad_to_sequence_len: true

gradient_accumulation_steps: 2
micro_batch_size: 2
num_epochs: 3
optimizer: adamw_torch
learning_rate: 0.0002
lr_scheduler: cosine

flash_attention: true
gradient_checkpointing: true
fsdp:
  - full_shard
  - auto_wrap
fsdp_config:
  fsdp_limit_all_gathers: true
  fsdp_sync_module_states: true
  fsdp_offload_params: false

warmup_steps: 50
eval_steps: 100
save_steps: 200
logging_steps: 10
weight_decay: 0.0

4. Lanzamiento del Entrenamiento Multi-GPU con Accelerate

# Ejecutar en 2 GPUs con FSDP
accelerate launch -m axolotl.cli.train config.yaml

5. Checklist para Evitar Fallos en Axolotl

• Verifica que sample_packing: true esté activado; puede duplicar la velocidad de entrenamiento al consolidar secuencias heterogéneas.
• Si experimentas picos de VRAM durante la primera época, activa gradient_checkpointing: true.

6. Preprocesamiento Acelerado de Datasets con axolotl.cli.preprocess

En datasets de cientos de miles de conversaciones, tokenizar el texto durante el propio bucle de entrenamiento provoca cuellos de botella severos en la CPU. Axolotl permite pre-computar y empaquetar todas las secuencias en disco antes de encender las GPUs:

# Preprocesar y tokenizar el dataset en paralelo
python3 -m axolotl.cli.preprocess config.yaml

Esto genera un caché binario serializado que las GPUs pueden leer a la velocidad máxima del bus PCIe (más de 5.000 tokens/segundo por tarjeta).

7. Configuración de DeepSpeed ZeRO-3 para Modelos de 70B Parámetros

Si deseas entrenar un modelo que no cabe en la memoria de un solo nodo incluso con LoRA, añade a tu config.yaml la configuración de DeepSpeed ZeRO-3 con offloading a la memoria RAM del sistema:

deepspeed: /app/axolotl/deepspeed_configs/zero3_bf16.json

ZeRO-3 particiona los pesos del modelo, los gradientes y los estados del optimizador a lo largo de todas las tarjetas disponibles, permitiendo entrenar Llama-3-70B en 4 GPUs de 24 GB sin errores de memoria.

8. Diagnóstico de Errores Críticos durante el Entrenamiento con Axolotl

1. 'CUDA error: an illegal memory access was encountered'

Suele producirse por un desajuste entre la versión del kernel de FlashAttention-2 y la versión exacta de PyTorch/CUDA instalada en el sistema anfitrión.

Solución: Reinstala FlashAttention compilando localmente contra los binarios del sistema: pip install flash-attn --no-build-isolation --force-reinstall.

2. Gradientes Explosivos y Pérdida que se Convierte en NaN (Loss = NaN)

Si la tasa de aprendizaje es demasiado alta o la precisión se fijó en float16 en lugar de bfloat16, las multiplicaciones de matrices en capas profundas desbordarán los exponentes flotantes.

Solución: Cambia a bf16: true y fp16: false en tu config.yaml, y activa el recorte de gradientes: clip_grad_norm: 1.0.

9. Automatización de Exportación a Hugging Face

Axolotl permite publicar directamente el adaptador o el modelo fusionado en tu repositorio de Hugging Face al finalizar la última época de entrenamiento configurando:

hub_model_id: tu-organizacion/llama3-especializado-soporte
hub_strategy: end
hub_token: hf_token_secreto_aqui

10. Manifiesto Completo de Despliegue con Docker y Multi-GPU

version: '3.8'

services:
  axolotl-trainer:
    image: winglian/axolotl:main-latest
    container_name: axolotl-job
    ipc: host
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]
    volumes:
      - ./config.yaml:/workspace/axolotl/config.yaml
      - ./dataset:/workspace/axolotl/dataset
      - ./output:/workspace/axolotl/output
    command: accelerate launch -m axolotl.cli.train config.yaml

11. Conclusión y Resumen de Buenas Prácticas

Axolotl es el marco de trabajo industrial definitivo para fine-tuning: flexibilidad declarativa, máxima eficiencia con FlashAttention y FSDP, y soporte probado para modelos de última generación.

12. Monitoreo del Clúster con Weights & Biases (WandB) y Alertas de VRAM

Axolotl soporta registro directo en WandB para visualizar en tiempo real las curvas de pérdida (Train Loss, Validation Loss), la norma de los gradientes (Gradient Norm) y el uso de memoria de cada GPU:

wandb_project: fine-tuning-corporativo
wandb_entity: mi-organizacion
wandb_watch: gradients
wandb_run_id: run_001_llama3_lora

13. Checklist Final de Seguridad y Almacenamiento

• Verifica con df -h que la partición donde se almacenan los checkpoints intermedios disponga de al menos 200 GB libres.
• Habilita el guardado en formato SafeTensors para blindar el repositorio frente a vulnerabilidades de ejecución remota de código inherentes a los archivos pickle tradicionales.

14. Guía de Optimización de Hiperparámetros para Evitar el Olvido Catastrófico

Al ajustar modelos en dominios técnicos específicos (como medicina o derecho), el modelo puede olvidar conocimientos generales de sentido común si el fine-tuning es demasiado agresivo:

• Mezcla de Datos Generales (Data Replay): Incluye entre un 10% y un 15% de datos de instrucciones generales (como Open-Platypus o ShareGPT) mezclados con tu corpus técnico especializado.
• LoRA Rank ($r$) y Alpha ($alpha$): Fija $r=16$ y $alpha=32$ para adaptaciones ligeras de estilo; sube a $r=64$ y $alpha=128$ únicamente si estás enseñando al modelo una sintaxis o lenguaje de programación completamente nuevo.
• Weight Decay: Establece un decaimiento de pesos de 0.01 para regularizar los pesos del adaptador y evitar que se sobrefocalicen en palabras clave infrecuentes.

15. Resumen Ejecutivo del Flujo de Trabajo Industrial

Axolotl representa el estándar de ingeniería para equipos que exigen reproducibilidad, velocidad extrema y control absoluto sobre sus pipelines de fine-tuning multi-GPU.

16. Estrategias de Particionado de Gradientes con DeepSpeed ZeRO-2 vs ZeRO-3

En clústeres multi-GPU con Axolotl, entender la diferencia entre etapas de DeepSpeed ZeRO es vital para balancear velocidad y ahorro de VRAM:

• ZeRO-1: Particiona únicamente los estados del optimizador (ahorra ~4x de memoria de optimizador con sobrecarga de red despreciable).
• ZeRO-2: Particiona los estados del optimizador y los gradientes (ahorra hasta un 70% de VRAM, ideal para LoRA en 2-4 GPUs).
• ZeRO-3: Particiona los pesos estáticos del modelo además de gradientes y optimizadores (permite entrenar modelos de 70B en GPUs de 24 GB, pero incrementa el tráfico de comunicación colectiva en un 50%).

Para fine-tuning con adaptadores LoRA en un solo servidor de 4 GPUs, ZeRO-2 es la configuración óptima al ofrecer la máxima velocidad sin penalización de comunicación inter-nodo.

17. Automatización de Guardado de Checkpoints en Almacenamiento S3 / Cloud

Durante entrenamientos que duran varios días en instancias spot o clústeres bajo demanda, un fallo imprevisto de la máquina puede arruinar horas de cómputo si los pesos solo residen en disco local. Se configura la sincronización periódica:

# Sincronizar checkpoints con AWS S3 o Cloudflare R2 tras cada época
aws s3 sync ./lora-llama3-platypus s3://mi-bucket-checkpoints/run_001/ --exclude "*.tmp"

Axolotl permite reanudar el entrenamiento desde el último checkpoint guardado de forma transparente añadiendo el argumento --resume_from_checkpoint en la llamada de Accelerate.

18. Matriz de Compatibilidad de Hardware para Axolotl

19. Optimización de Memoria con Paged Optimizers de 8 bits en Bitsandbytes

Para aquellos casos donde la memoria VRAM esté al borde del límite físico, Axolotl permite sustituir el optimizador AdamW estándar de 32 bits por paged_adamw_8bit. Esta técnica pagina dinámicamente los estados del optimizador entre la VRAM y la memoria RAM del sistema operativo cuando detecta picos de activación, previniendo cuelgues por CUDA OOM en secuencias largas de más de 4.096 tokens sin ralentizar el paso hacia adelante.