Fine-tuning: de usar un modelo a entrenarlo vos

0. De usar un modelo a entrenarlo

El ejemplo del capítulo original sigue siendo el correcto para aprender la mecánica: tomar BERT (un encoder-only, Capítulo 1) y ajustarlo para una tarea concreta — detectar si dos frases dicen lo mismo (el dataset MRPC, 5.801 pares de oraciones). Es chico, rápido de entrenar, y expone cada pieza del proceso sin la complejidad de un LLM de 30B parámetros de por medio.

Vas a ver tres formas de llegar al mismo resultado, de más alto nivel a más bajo nivel:

1. Preparar los datos

Con datasets (la librería, no el concepto) cargar un dataset del Hub es una línea:

from datasets import load_dataset

raw_datasets = load_dataset("glue", "mrpc")
# DatasetDict con splits train (3668) / validation (408) / test (1725)

Tokenizar un par de oraciones (no una sola, como en el Capítulo 2) agrega un campo nuevo:

inputs = tokenizer("Esta es la primera oración.", "Esta es la segunda.")
# input_ids, attention_mask, y ahora también token_type_ids:
# [0,0,0,0,0,0,0,0, 1,1,1,1,1,1,1] — 0 = primera oración, 1 = segunda

token_type_ids le dice al modelo dónde termina una oración y empieza la otra — BERT lo necesita porque se preentrenó con un objetivo llamado next sentence prediction (predecir si una oración sigue a otra en el texto original).

Para tokenizar el dataset completo sin cargarlo todo en RAM, Dataset.map(batched=True) aplica la función en lotes, aprovechando que el tokenizer de tokenizers está escrito en Rust y paraleliza internamente:

def tokenize_function(example):
    return tokenizer(example["sentence1"], example["sentence2"], truncation=True)

tokenized_datasets = raw_datasets.map(tokenize_function, batched=True)

Notá que acá no se pasa padding=True. La razón es padding dinámico: rellenar cada lote hasta su propio máximo (no el máximo de todo el dataset) ahorra cómputo real cuando las longitudes varían mucho. DataCollatorWithPadding hace ese relleno al momento de armar cada batch, no antes. Se construye con el tokenizer, porque necesita saber cuál es el pad_token_id y de qué lado rellenar:

from transformers import DataCollatorWithPadding

data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

# Probalo a mano con 8 ejemplos: cada uno mide distinto, el batch sale rectangular
samples = [tokenized_datasets["train"][i] for i in range(8)]
samples = [{k: v for k, v in s.items() if k not in ["idx", "sentence1", "sentence2"]} for s in samples]
batch = data_collator(samples)
{k: v.shape for k, v in batch.items()}
# {'input_ids': torch.Size([8, 67]), 'attention_mask': ..., 'token_type_ids': ..., 'labels': torch.Size([8])}

El 67 es el largo del ejemplo más largo de esos ocho — otro batch va a tener otro ancho. Fijate que las columnas de texto se sacan antes: el collator solo sabe rellenar tensores, y un string lo hace fallar.

2. Fine-tunear con la Trainer API

Falta una pieza antes de entrenar: la función que convierte las predicciones en una métrica. Sin ella, Trainer solo reporta la loss, que sirve para ver si el entrenamiento avanza pero no te dice cuántos pares clasificó bien. Trainer le pasa una tupla (logits, labels) sobre todo el set de evaluación; hay que convertir los logits en una clase con argmax y comparar:

import evaluate
import numpy as np

metric = evaluate.load("glue", "mrpc")   # accuracy + F1, las métricas oficiales de la tarea

def compute_metrics(eval_preds):
    logits, labels = eval_preds
    predictions = np.argmax(logits, axis=-1)   # [n, 2] → [n]: la clase con mayor logit
    return metric.compute(predictions=predictions, references=labels)
# → {'accuracy': 0.85..., 'f1': 0.89...} después de entrenar

Con los datos y la métrica listos, definir el entrenamiento son tres objetos:

from transformers import TrainingArguments, AutoModelForSequenceClassification, Trainer

training_args = TrainingArguments("test-trainer", eval_strategy="epoch")
model = AutoModelForSequenceClassification.from_pretrained(checkpoint, num_labels=2)

trainer = Trainer(
    model, training_args,
    train_dataset=tokenized_datasets["train"],
    eval_dataset=tokenized_datasets["validation"],
    data_collator=data_collator,
    processing_class=tokenizer,
    compute_metrics=compute_metrics,
)
trainer.train()

num_labels=2 importa: BERT nunca vio esta tarea en su preentrenamiento, así que AutoModelForSequenceClassification descarta el head original y le pone uno nuevo, inicializado al azar — el warning que ves al cargar el modelo es exactamente eso, no un error. compute_metrics es la función que convierte logits en accuracy/F1 usando la librería evaluate; sin ella, Trainer solo te va a reportar la loss.

Cuatro flags de TrainingArguments que definen el rendimiento y la viabilidad física en hardware:

3. Un training loop a mano

Todo lo que Trainer hace por vos, en PyTorch puro, es esto — vale la pena leerlo una vez para saber qué se esconde detrás del .train(). Primero la preparación, que es donde Trainer esconde más trabajo: dejar el dataset con solo las columnas que el modelo acepta, armar los DataLoader con el collator, y crear el optimizador y el scheduler:

import torch
from torch.optim import AdamW
from torch.utils.data import DataLoader
from transformers import get_scheduler

# El modelo espera input_ids, attention_mask, token_type_ids y labels — nada más
tokenized_datasets = tokenized_datasets.remove_columns(["sentence1", "sentence2", "idx"])
tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
tokenized_datasets.set_format("torch")

train_dataloader = DataLoader(tokenized_datasets["train"], shuffle=True, batch_size=8, collate_fn=data_collator)
eval_dataloader = DataLoader(tokenized_datasets["validation"], batch_size=8, collate_fn=data_collator)

model = AutoModelForSequenceClassification.from_pretrained(checkpoint, num_labels=2)
optimizer = AdamW(model.parameters(), lr=5e-5)

num_epochs = 3
num_training_steps = num_epochs * len(train_dataloader)   # 3 × 459 = 1377 pasos
lr_scheduler = get_scheduler(
    "linear", optimizer=optimizer, num_warmup_steps=0, num_training_steps=num_training_steps
)   # el lr baja linealmente de 5e-5 a 0 a lo largo de los 1377 pasos

device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
model.to(device)

Y el loop en sí:

model.train()
for epoch in range(num_epochs):
    for batch in train_dataloader:
        batch = {k: v.to(device) for k, v in batch.items()}
        outputs = model(**batch)
        loss = outputs.loss
        loss.backward()

        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()

Cinco pasos, siempre en este orden: forward (calcular la salida y la loss), backward (propagar el gradiente), optimizer.step() (actualizar pesos con ese gradiente), scheduler.step() (ajustar el learning rate), zero_grad() (limpiar el gradiente para el próximo batch — si te olvidás este paso, los gradientes se acumulan silenciosamente entre batches).

La evaluación es el mismo recorrido sin las tres últimas líneas, y con dos cosas que Trainer también hace sin que lo veas: model.eval() apaga el dropout (si no, cada evaluación da un número distinto), y torch.no_grad() evita construir el grafo de gradientes — más rápido y varias veces menos memoria:

metric = evaluate.load("glue", "mrpc")
model.eval()
for batch in eval_dataloader:
    batch = {k: v.to(device) for k, v in batch.items()}
    with torch.no_grad():
        outputs = model(**batch)
    predictions = torch.argmax(outputs.logits, dim=-1)
    metric.add_batch(predictions=predictions, references=batch["labels"])

metric.compute()   # {'accuracy': 0.85..., 'f1': 0.89...}
model.train()      # si vas a seguir entrenando, volvé al modo train

El optimizador estándar universal para transformers es AdamW (Loshchilov & Hutter, 2017). Para entender por qué Adam clásico falla y por qué AdamW es indispensable, hay que observar su formulación matemática exacta:

¿Por qué falla Adam con regularización L2 clásica? En SGD tradicional, minimizar $\mathcal{L}+\frac{1}{2}\lambda \parallel \mathbf{W}{\parallel }^2$ suma $\lambda \mathbf{W}$ al gradiente: ${\tilde{\mathbf{g}}}_t={\mathbf{g}}_t+\lambda {\mathbf{W}}_t$. Si se mete ese gradiente modificado dentro de Adam, el paso efectivo resulta:

El término de penalización queda dividido por $\sqrt{{\hat{\mathbf{v}}}_t}$. Esto produce una catástrofe matemática: los pesos que reciben gradientes grandes y frecuentes tienen un ${\mathbf{v}}_t$ enorme, suprimiendo casi por completo su regularización. Por el contrario, pesos con gradientes diminutos o infrecuentes sufren una contracción desproporcionada. AdamW soluciona esto desacoplando el decay: resta ${\eta }_t\lambda {\mathbf{W}}_t$ directamente, logrando una penalización constante e invariante a la escala de los gradientes para todos los parámetros de la red.

La física de VRAM: ¿por qué fine-tunear consume 16 bytes por parámetro?

En inferencia (Capítulo 2) solo se cargan los pesos del modelo ($2\text{B}$ por parámetro en BF16). En entrenamiento completo con precisión mixta y AdamW, el estado estático de la GPU escala a 16 bytes por parámetro:

¿Por qué es estrictamente obligatoria la copia maestra en FP32? En BF16, la mantisa dispone de únicamente 7 bits, con un épsilon de máquina de ${ϵ}_{\text{BF16}}=2^{-7}\approx 7.81\times {10}^{-3}$. Si el learning rate es $\eta ={10}^{-5}$ y el paso de actualización es $\mathrm{\Delta }\mathbf{W}\sim {10}^{-5}$, al sumarlo a un peso de magnitud $\approx 1.0$ en BF16 ocurre subdesbordamiento aritmético: $1.0+{10}^{-5}=1.0$. La actualización desaparece por completo y el modelo queda congelado en el tiempo. La copia maestra en FP32 (${ϵ}_{\text{FP32}}\approx 1.19\times {10}^{-7}$) acumula estos micro-pasos con fidelidad, y luego proyecta el resultado a BF16 para el siguiente forward pass.

Para correr el mismo loop en varias GPUs sin reescribirlo, Accelerate envuelve los mismos objetos:

from accelerate import Accelerator

accelerator = Accelerator()
train_dl, eval_dl, model, optimizer = accelerator.prepare(train_dataloader, eval_dataloader, model, optimizer)
# reemplazar loss.backward() por accelerator.backward(loss), nada más

4. Leer las curvas de entrenamiento

Un número de loss aislado no te dice nada — la forma de la curva en el tiempo sí. Esta parte no estaba en el curso original de 2022; se agregó porque interpretar mal una curva es de los errores más comunes y más caros (entrenás horas para nada, o peor, publicás un modelo roto).

Un detalle que confunde la primera vez: la curva de accuracy avanza en escalones, no suave como la loss. La loss puede mejorar aunque la predicción siga siendo incorrecta (0.3 → 0.4 en un clasificador binario sigue siendo "0"); la accuracy solo salta cuando la predicción cruza el umbral que la vuelve correcta.

5. Lo que cambió: fine-tuning completo ya no es la norma

Todo lo de arriba ajusta los 110 millones de parámetros de BERT, todos. Para un modelo de ese tamaño está perfecto — es rápido, es simple, y no hay ninguna razón para complicarlo. Para un LLM de 8B, 30B o 70B parámetros, hacer lo mismo significa guardar en memoria el estado del optimizador (Adam guarda dos números extra por parámetro) para cada uno de esos parámetros — varias veces el tamaño del modelo, solo en overhead de entrenamiento.

LoRA (Low-Rank Adaptation) ataca esto desde otro ángulo: en vez de tocar la matriz de pesos original W, la congelás, y entrenás dos matrices mucho más chicas (A y B) cuyo producto se suma a la salida de W. El modelo base nunca cambia; lo único que se entrena y se guarda es ese par de matrices de bajo rango — típicamente menos del 1% de los parámetros originales.

from peft import LoraConfig, get_peft_model

config = LoraConfig(r=16, lora_alpha=32, target_modules=["q_proj", "v_proj"])
model = get_peft_model(base_model, config)   # base_model: Llama-2-7B
model.print_trainable_parameters()
# trainable params: 8,388,608 || all params: 6,746,804,224 || trainable%: 0.1243

El número de parámetros entrenables se puede verificar a mano, y vale la pena hacerlo una vez. Cada matriz W apuntada, de tamaño d_out × d_in, recibe un A de r × d_in y un B de d_out × r: en total r · (d_in + d_out) parámetros nuevos. En Llama-2-7B, q_proj y v_proj son de 4096 × 4096, y hay 32 capas:

16 × (4096 + 4096) = 131.072 parámetros por matriz
× 2 matrices (q_proj, v_proj) × 32 capas = 8.388.608   ✓

1. Inicialización asimétrica: preservación estricta de la función base en $t=0$

Para no destruir las representaciones aprendidas durante el preentrenamiento al iniciar el fine-tuning, LoRA inicializa:

En el paso $t=0$, $\mathrm{\Delta }\mathbf{W}$ es idénticamente la matriz nula: el modelo adaptado produce exactamente las mismas salidas que el modelo base original. Si $\mathbf{B}$ también se inicializara con ruido gaussiano, la suma agregaría perturbaciones aleatorias masivas sobre cada capa que destruirían el conocimiento preentrenado en los primeros pasos.

2. Ley asociativa y eficiencia de FLOPs: ¿por qué no multiplicar $B\cdot A$ primero?

Para un tensor de activaciones $\mathbf{X}\in {\mathbb{ℝ}}^{B\times N\times d_{\text{in}}}$, la multiplicación aprovecha la asociatividad:

• $\mathbf{X}{\mathbf{A}}^{\top }$ proyecta $[B,N,d_{\text{in}}]\times [d_{\text{in}},r]\to [B,N,r]$: consume $2BN{d}_{\text{in}}r$ FLOPs.
• Multiplicar por ${\mathbf{B}}^{\top }$ proyecta $[B,N,r]\times [r,d_{\text{out}}]\to [B,N,d_{\text{out}}]$: consume $2BNr{d}_{\text{out}}$ FLOPs.
• Total de FLOPs en LoRA: $2BNr(d_{\text{in}}+d_{\text{out}})$.

Si por error calculáramos $\mathrm{\Delta }\mathbf{W}=\mathbf{B}\mathbf{A}$ primero, materializaríamos una matriz densa de $[d_{\text{out}},d_{\text{in}}]$, costando $2d_{\text{out}}r{d}_{\text{in}}+2BN{d}_{\text{in}}{d}_{\text{out}}$ FLOPs y destruyendo el ahorro de memoria. La relación de cómputo adicional de LoRA frente a la capa base ($2BN{d}_{\text{in}}{d}_{\text{out}}$) es:

Para dimensiones estándar de Llama ($d_{\text{in}}=d_{\text{out}}=4096,r=16$): $16\times \frac{2}{4096}=\frac{32}{4096}=\frac{1}{128}\approx 0.78\%$. El cómputo del adaptador LoRA agrega menos del 1% de FLOPs al forward pass.

3. El colapso del estado del optimizador en VRAM

Como la matriz base $\mathbf{W}$ está congelada (requires_grad=False), PyTorch autograd no calcula ${\nabla }_{\mathbf{W}}\mathcal{L}$ ni reserva sus gradientes. En consecuencia, AdamW no crea copia maestra en FP32 ni momentos $m_t,v_t$ para $\mathbf{W}$. En un modelo de 8B, esto elimina de golpe 112 GB de VRAM estática (los 14 bytes por parámetro de gradientes + estados de optimizador). Los gradientes solo se propagan a $\mathbf{A}$ y $\mathbf{B}$:

En Llama-3-8B con target_modules="all-linear" y $r=16$, los parámetros entrenables suman apenas $\sim 40\text{M}$ ($\approx 0.5\%$). El estado de optimizador completo para ellos ocupa: $40\times {10}^6\times 16\text{bytes}\approx 640\text{MB}$.

4. Fusión de pesos para inferencia con latencia cero

Al finalizar el entrenamiento, si el modelo base $\mathbf{W}$ está en punto flotante (BF16 o FP16), no es necesario mantener la rama paralela en inferencia:

Mediante model.merge_and_unload(), la corrección de bajo rango se suma directamente en la matriz de pesos del modelo base. El modelo final resultante tiene exactamente la misma arquitectura y tamaño que el original, ejecutándose con cero sobrecarga de latencia y sin kernels especiales en vLLM o TGI.

Los dos números de LoraConfig que más tenés que entender:

• r es el rango: el ancho de A y B. Cuanto más alto, más capacidad tiene el adaptador para representar cambios (y más parámetros pesa, linealmente). Valores habituales: 8 a 64. Un r más alto no siempre mejora: el paper de LoRA muestra que rangos muy bajos (4-8) ya capturan casi todo en muchas tareas.
• lora_alpha es un factor de escala: el producto B·A se multiplica por alpha / r antes de sumarse a la salida de W. Con r=16, lora_alpha=32 la escala es 2. Tener el cociente fijo permite cambiar r sin tener que retocar el learning rate — por eso la convención alpha = 2r es tan común.

Y target_modules: el ejemplo clásico (y el del paper original de LoRA) apunta solo a q_proj y v_proj. El paper de QLoRA midió esto con cuidado y encontró que apuntar todas las capas lineales — las cuatro de atención (q_proj, k_proj, v_proj, o_proj) más las del MLP (gate_proj, up_proj, down_proj en Llama) — era necesario para igualar la calidad del fine-tuning completo; con solo atención se quedaba corto. Hoy es el default razonable (target_modules="all-linear" en PEFT), a cambio de unas pocas veces más parámetros en el adaptador — que siguen siendo una fracción ínfima del modelo.

QLoRA combina las dos ideas: el modelo base se carga cuantizado (originalmente en NF4, un formato de 4 bits distinto al NVFP4 de Blackwell que ya conocés) y congelado, y las matrices LoRA se entrenan encima en precisión más alta — la razón por la que hoy podés fine-tunear un modelo de 70B en una sola GPU de consumo.

Nada de esto invalida las secciones anteriores — Trainer, el training loop manual, y las curvas de aprendizaje funcionan exactamente igual con LoRA. Lo único que cambia es qué parámetros tienen requires_grad=True.

6. Errores comunes

La mayoría se ve en la curva de loss antes que en un traceback — por eso la sección 4 viene antes que esta.

7. Resumen

1. Fine-tuning parte de datos tokenizados en lotes con padding dinámico — nunca rellenes al máximo del dataset entero si podés rellenar al máximo del batch.
2. Trainer resuelve el 95% de los casos en pocas líneas; el training loop manual existe para cuando necesitás control fino, y Accelerate lo escala a multi-GPU casi gratis.
3. Las curvas de loss y accuracy tienen formas reconocibles — overfitting, underfitting, entrenamiento errático — cada una con su propio remedio.
4. AdamW desacopla el weight decay del paso adaptativo de Adam; por eso es el optimizador por defecto para transformers, y por eso lr y weight_decay se ajustan por separado.
5. Ajustar el 100% de los parámetros es la excepción hoy, no la regla: LoRA y QLoRA son el camino por defecto para cualquier LLM que no te entre cómodo en memoria para un fine-tune completo — y ya los viste corriendo en tus propios proyectos. r es capacidad, alpha/r es escala, y apuntar todas las capas lineales es lo que iguala al fine-tune completo.

8. Ponete a prueba

Respondé antes de abrir cada respuesta. Si fallás una, volvé a la sección indicada.

9. Para practicar

import numpy as np, evaluate
from datasets import load_dataset
from transformers import (AutoTokenizer, AutoModelForSequenceClassification,
                          DataCollatorWithPadding, TrainingArguments, Trainer)

ckpt = "distilbert-base-uncased"
raw = load_dataset("glue", "mrpc")
tok = AutoTokenizer.from_pretrained(ckpt)
tokenized = raw.map(lambda ex: tok(ex["sentence1"], ex["sentence2"], truncation=True), batched=True)
metric = evaluate.load("glue", "mrpc")

def compute_metrics(p):
    return metric.compute(predictions=np.argmax(p[0], axis=-1), references=p[1])

model = AutoModelForSequenceClassification.from_pretrained(ckpt, num_labels=2)
args = TrainingArguments("mrpc-mini", eval_strategy="epoch", num_train_epochs=2,
                         per_device_train_batch_size=16, report_to="none")
trainer = Trainer(model, args,
                  train_dataset=tokenized["train"].shuffle(seed=42).select(range(500)),
                  eval_dataset=tokenized["validation"],
                  data_collator=DataCollatorWithPadding(tokenizer=tok),
                  processing_class=tok, compute_metrics=compute_metrics)

print("antes:", trainer.evaluate())
trainer.train()
print("después:", trainer.evaluate())

import torch
from torch.optim import AdamW
from torch.utils.data import DataLoader

ds = tokenized["train"].shuffle(seed=42).select(range(480))
ds = ds.remove_columns(["sentence1", "sentence2", "idx"]).rename_column("label", "labels")
ds.set_format("torch")
dl = DataLoader(ds, batch_size=8, shuffle=True, collate_fn=DataCollatorWithPadding(tokenizer=tok))

def correr(lr, con_zero_grad):
    torch.manual_seed(0)
    m = AutoModelForSequenceClassification.from_pretrained(ckpt, num_labels=2)
    opt = AdamW(m.parameters(), lr=lr)
    m.train()
    for paso, batch in enumerate(dl):
        loss = m(**batch).loss
        loss.backward()
        opt.step()
        if con_zero_grad:
            opt.zero_grad()
        if paso % 10 == 0:
            print(f"lr={lr} zero_grad={con_zero_grad} paso {paso}: {loss.item():.3f}")

correr(5e-5, True)
correr(5e-5, False)
correr(1e-3, True)

import torch
from transformers import AutoModelForCausalLM
from peft import LoraConfig, get_peft_model

base = AutoModelForCausalLM.from_pretrained("google/gemma-4-E2B-it", dtype=torch.bfloat16, device_map="cuda")
cfg = base.config.get_text_config()
print(cfg.num_hidden_layers, cfg.hidden_size, cfg.num_attention_heads, cfg.num_key_value_heads, cfg.num_kv_shared_layers)

for r in (8, 16):
    m = get_peft_model(base, LoraConfig(r=r, lora_alpha=2*r, task_type="CAUSAL_LM",
                                        target_modules=r".*language_model.*\.(q_proj|v_proj)$"))   # solo el modelo de texto
    m.print_trainable_parameters()
    m.unload()   # saca el adaptador para poder aplicar otro sobre el mismo base

10. Fuentes y para seguir

• Hugging Face LLM Course — Chapter 3: Fine-tuning a pretrained model — el capítulo original, con el ejemplo MRPC completo en Trainer, loop manual y Accelerate.
• Transformers docs — Trainer — la referencia de Trainer y TrainingArguments; buscá load_best_model_at_end, max_grad_norm y gradient_accumulation_steps.
• Decoupled Weight Decay Regularization (Loshchilov & Hutter, 2017) — el paper de AdamW; la figura 2 muestra en una imagen por qué el decay acoplado de Adam no funciona como debería.
• LoRA: Low-Rank Adaptation of Large Language Models (Hu et al., 2021) — el paper original; la sección 7 es el estudio de qué rango y qué matrices apuntar.
• QLoRA: Efficient Finetuning of Quantized LLMs (Dettmers et al., 2023) — NF4, doble cuantización, y el resultado de que hay que apuntar todas las capas lineales para igualar al fine-tune completo.
• PEFT — documentación — LoraConfig con todos sus campos, y las guías de cómo combinar, fusionar y servir adaptadores.