Saltar al contenido
Lección 6 de 12

Entrenamiento con Hugging Face

5 min read

El Stack de Entrenamiento de Hugging Face

Hugging Face proporciona el ecosistema más maduro y ampliamente usado para fine-tuning de LLMs. El stack consiste en cuatro bibliotecas principales: transformers para carga de modelos, peft para LoRA/QLoRA, trl para entrenamiento (SFTTrainer, DPOTrainer) y bitsandbytes para cuantización. Esta lección recorre un pipeline de entrenamiento completo con cada parámetro explicado.

Configuración y Dependencias

pip install torch transformers peft trl datasets bitsandbytes accelerate

Para soporte de Flash Attention 2 (recomendado para velocidad):

pip install flash-attn --no-build-isolation

Cargando el Modelo Base con QLoRA

import torch
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
)
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training

# Identificador del modelo
model_name = "meta-llama/Llama-3.1-8B-Instruct"

# Configuracion de cuantizacion de 4 bits
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True,
)

# Cargar tokenizador
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"

# Cargar modelo con cuantizacion
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=bnb_config,
    device_map="auto",
    attn_implementation="flash_attention_2",  # Usar Flash Attention
    torch_dtype=torch.bfloat16,
)

# Preparar modelo para entrenamiento QLoRA
model = prepare_model_for_kbit_training(model)
model.config.use_cache = False  # Deshabilitar cache KV durante entrenamiento

Detalles clave:

  • pad_token = eos_token previene errores al agrupar secuencias de diferentes longitudes en batches.
  • padding_side = "right" es requerido para modelos de lenguaje causal.
  • use_cache = False deshabilita el cache KV que entra en conflicto con gradient checkpointing.

Configurando LoRA

from peft import LoraConfig, TaskType

lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    target_modules=[
        "q_proj", "k_proj", "v_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj",
    ],
    lora_dropout=0.05,
    bias="none",
    task_type=TaskType.CAUSAL_LM,
)

# Aplicar LoRA al modelo
model = get_peft_model(model, lora_config)

# Imprimir parametros entrenables
model.print_trainable_parameters()
# Salida: trainable params: 20,971,520 || all params: 8,051,232,768 || trainable%: 0.26%

Cargando el Dataset

from datasets import load_dataset

# Cargar desde archivos JSONL locales
dataset = load_dataset("json", data_files={
    "train": "data/train.jsonl",
    "validation": "data/val.jsonl",
})

# O cargar desde Hugging Face Hub
# dataset = load_dataset("tu-usuario/tu-dataset")

print(f"Train: {len(dataset['train'])} ejemplos")
print(f"Validacion: {len(dataset['validation'])} ejemplos")

SFTTrainer: Supervised Fine-Tuning

El SFTTrainer de trl maneja el bucle de entrenamiento, el formateo de datos y la evaluación:

from trl import SFTTrainer
from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="./output/llama3-finetuned",

    # Duracion del entrenamiento
    num_train_epochs=3,                  # 2-5 epocas es tipico
    max_steps=-1,                        # -1 significa usar num_train_epochs

    # Tamano de batch
    per_device_train_batch_size=2,       # Ajustar segun VRAM
    per_device_eval_batch_size=2,
    gradient_accumulation_steps=8,       # Batch efectivo: 2 * 8 = 16

    # Tasa de aprendizaje
    learning_rate=2e-4,                  # 1e-4 a 2e-4 para QLoRA
    lr_scheduler_type="cosine",          # Decaimiento coseno
    warmup_ratio=0.05,                   # 5% de pasos de calentamiento

    # Optimizacion
    optim="paged_adamw_8bit",            # Optimizador eficiente en memoria
    weight_decay=0.01,
    max_grad_norm=1.0,                   # Recorte de gradientes

    # Precision
    bf16=True,                           # Usar bfloat16

    # Registro
    logging_steps=10,
    logging_first_step=True,

    # Evaluacion
    eval_strategy="steps",
    eval_steps=50,

    # Guardado
    save_strategy="steps",
    save_steps=100,
    save_total_limit=3,                  # Mantener solo 3 mejores checkpoints

    # Otros
    gradient_checkpointing=True,         # Intercambiar computo por memoria
    gradient_checkpointing_kwargs={"use_reentrant": False},
    report_to="tensorboard",             # O "wandb"
    seed=42,
)

trainer = SFTTrainer(
    model=model,
    args=training_args,
    train_dataset=dataset["train"],
    eval_dataset=dataset["validation"],
    processing_class=tokenizer,
    max_seq_length=2048,
    packing=False,                       # True para ejemplos cortos
    dataset_text_field=None,             # Usar formato 'messages'
)

Entendiendo los Argumentos Clave de Entrenamiento

Tasa de aprendizaje: El hiperparámetro más impactante. Para QLoRA con rango 16:

  • Comienza con 2e-4 para datasets pequeños (menos de 1,000 ejemplos)
  • Usa 1e-4 para datasets más grandes (más de 5,000 ejemplos)
  • Si la pérdida de entrenamiento oscila salvajemente, redúcela

Tamaño de batch y acumulación de gradientes: El tamaño de batch efectivo es per_device_batch_size * gradient_accumulation_steps * num_gpus. Tamaños de batch efectivos más grandes producen entrenamiento más estable pero requieren más memoria. Comienza con un batch efectivo de 16.

Épocas: Para datasets pequeños (200-500 ejemplos), usa 3-5 épocas. Para datasets más grandes (5,000+), 1-2 épocas a menudo es suficiente. Observa la pérdida de validación para detectar sobreajuste.

Calentamiento: Incrementa gradualmente la tasa de aprendizaje al inicio del entrenamiento. 5% del total de pasos es un valor seguro por defecto. Ayuda a prevenir inestabilidad temprana.

Ejecutando el Entrenamiento

# Iniciar entrenamiento
trainer.train()

# Guardar el adaptador LoRA final
trainer.save_model("./output/llama3-finetuned/final")
tokenizer.save_pretrained("./output/llama3-finetuned/final")

Entrenamiento DPO: Alineación por Preferencia

Direct Preference Optimization (DPO) entrena al modelo con pares de respuestas preferidas y rechazadas. Se usa después de SFT para alinear aún más el comportamiento del modelo.

from trl import DPOTrainer, DPOConfig

# Cargar tu dataset de preferencias
# Cada ejemplo necesita: prompt, chosen, rejected
dpo_dataset = load_dataset("json", data_files="data/preferences.jsonl")

dpo_config = DPOConfig(
    output_dir="./output/llama3-dpo",
    num_train_epochs=1,                  # DPO tipicamente necesita menos epocas
    per_device_train_batch_size=2,
    gradient_accumulation_steps=4,
    learning_rate=5e-5,                  # LR mas bajo que SFT
    lr_scheduler_type="cosine",
    warmup_ratio=0.1,
    beta=0.1,                            # Parametro de temperatura DPO
    bf16=True,
    optim="paged_adamw_8bit",
    gradient_checkpointing=True,
    logging_steps=10,
    eval_strategy="steps",
    eval_steps=50,
    report_to="tensorboard",
)

dpo_trainer = DPOTrainer(
    model=model,
    args=dpo_config,
    train_dataset=dpo_dataset["train"],
    eval_dataset=dpo_dataset["validation"],
    processing_class=tokenizer,
)

dpo_trainer.train()

Parámetro beta de DPO: Controla cuánto se desvía el modelo de la política de referencia. Beta más bajo (0.05) significa más desviación, beta más alto (0.5) mantiene el modelo más cerca de su comportamiento SFT. Comienza con 0.1.

Probando tu Modelo Fine-Tuned

from peft import PeftModel

# Cargar el modelo fine-tuned para inferencia
base_model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=bnb_config,
    device_map="auto",
    torch_dtype=torch.bfloat16,
)
model = PeftModel.from_pretrained(base_model, "./output/llama3-finetuned/final")

# Generar una respuesta
messages = [
    {"role": "system", "content": "Eres un asistente util."},
    {"role": "user", "content": "Explica que es LoRA en terminos simples."}
]

inputs = tokenizer.apply_chat_template(messages, return_tensors="pt").to(model.device)

outputs = model.generate(
    inputs,
    max_new_tokens=256,
    temperature=0.7,
    do_sample=True,
)

response = tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True)
print(response)

Consejos Prácticos

  • Siempre monitorea la memoria GPU durante los primeros pasos de entrenamiento. Usa nvidia-smi o torch.cuda.memory_summary().
  • Comienza con un subconjunto pequeño de tus datos (50-100 ejemplos) para verificar que el pipeline funciona antes de entrenar con el dataset completo.
  • Guarda checkpoints frecuentemente al inicio del entrenamiento. Si algo sale mal en el paso 500, no quieres reiniciar desde cero.
  • Revisa algunas salidas del modelo en cada checkpoint. Los números de pérdida solos no te dicen si el modelo está aprendiendo el comportamiento correcto.

En la próxima lección, cubriremos Unsloth — un framework de entrenamiento alternativo que ofrece mejoras de velocidad de 2x y reducción de memoria del 60% sobre el entrenamiento vanilla de Hugging Face.