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 mas 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 cuantizacion. Esta leccion recorre un pipeline de entrenamiento completo con cada parametro explicado.

Configuracion 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 evaluacion:

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 hiperparametro mas impactante. Para QLoRA con rango 16:

  • Comienza con 2e-4 para datasets pequenos (menos de 1,000 ejemplos)
  • Usa 1e-4 para datasets mas grandes (mas de 5,000 ejemplos)
  • Si la perdida de entrenamiento oscila salvajemente, reducela

Tamano de batch y acumulacion de gradientes: El tamano de batch efectivo es per_device_batch_size * gradient_accumulation_steps * num_gpus. Tamanos de batch efectivos mas grandes producen entrenamiento mas estable pero requieren mas memoria. Comienza con un batch efectivo de 16.

Epocas: Para datasets pequenos (200-500 ejemplos), usa 3-5 epocas. Para datasets mas grandes (5,000+), 1-2 epocas a menudo es suficiente. Observa la perdida de validacion 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: Alineacion por Preferencia

Direct Preference Optimization (DPO) entrena al modelo con pares de respuestas preferidas y rechazadas. Se usa despues de SFT para alinear aun mas 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()

Parametro beta de DPO: Controla cuanto se desvia el modelo de la politica de referencia. Beta mas bajo (0.05) significa mas desviacion, beta mas alto (0.5) mantiene el modelo mas 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 Practicos

  • Siempre monitorea la memoria GPU durante los primeros pasos de entrenamiento. Usa nvidia-smi o torch.cuda.memory_summary().
  • Comienza con un subconjunto pequeno 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 numeros de perdida solos no te dicen si el modelo esta aprendiendo el comportamiento correcto.

En la proxima leccion, cubriremos Unsloth — un framework de entrenamiento alternativo que ofrece mejoras de velocidad de 2x y reduccion de memoria del 60% sobre el entrenamiento vanilla de Hugging Face.