Autograd y el loop de optimización, con nombre y apellido

1. El DAG dinámico: por qué PyTorch permite control de flujo

torch.autograd mantiene, mientras se ejecuta el forward pass, un grafo acíclico dirigido (DAG) de objetos Function — las hojas son los tensores de entrada, la raíz es el tensor de salida (típicamente el loss). Cada operación queda registrada como un nodo con su función de derivada correspondiente, accesible vía .grad_fn:

z = torch.matmul(x, w) + b
loss = F.binary_cross_entropy_with_logits(z, y)

print(z.grad_fn)     # <AddBackward0 object>
print(loss.grad_fn)  # <BinaryCrossEntropyWithLogitsBackward0 object>

El detalle que la documentación oficial marca como clave: este grafo se reconstruye desde cero en cada forward pass. No hay un grafo estático definido de antemano (como en TensorFlow 1.x o en Theano) — cada llamada a .backward() consume el grafo actual, y el próximo forward pass arma uno nuevo. Esto es exactamente lo que permite usar if, loops y control de flujo arbitrario de Python dentro de un modelo: la forma, el tamaño y las operaciones pueden cambiar en cada iteración sin ningún problema, porque no hay un grafo fijo que reconciliar.

"Consumir el grafo" tiene un significado preciso: para calcular derivadas, cada nodo guardó los tensores intermedios que necesita (la entrada de un matmul, la salida de una sigmoid) — eso es la memoria de activaciones, la razón por la que entrenar cuesta mucho más memoria que inferir. backward() libera esos tensores a medida que los usa. Por eso un segundo backward() sobre el mismo grafo da RuntimeError: Trying to backward through the graph a second time (or directly access saved tensors after they have already been freed). Si de verdad necesitás dos pasadas (dos losses que comparten un forward, o derivadas de orden superior), la primera se llama con retain_graph=True:

out = model(x)
loss_a, loss_b = crit_a(out, y), crit_b(out, y)
loss_a.backward(retain_graph=True)   # conserva las activaciones
loss_b.backward()                    # segunda pasada sobre el mismo grafo
# ...aunque casi siempre es más simple y barato: (loss_a + loss_b).backward()

torch.autograd.grad es la interfaz funcional de lo mismo: en vez de escribir en .grad de cada parámetro, devuelve los gradientes respecto de los tensores que le pidas. Es lo que se usa para penalizar gradientes (gradient penalty de WGAN-GP), para derivadas segundas (create_graph=True deja el resultado dentro del grafo) y en métodos de interpretabilidad que derivan la salida respecto de la entrada:

x = torch.randn(4, requires_grad=True)
y = (x ** 3).sum()
(dy_dx,) = torch.autograd.grad(y, x, create_graph=True)     # 3x²; x.grad sigue en None
(d2y_dx2,) = torch.autograd.grad(dy_dx.sum(), x)              # 6x

2. requires_grad, no_grad, inference_mode y detach

Solo los tensores marcados con requires_grad=True acumulan historial de gradiente — y solo se puede leer .grad en los leaf nodes del grafo (los tensores de entrada, no los resultados intermedios; para un intermedio hay que pedirlo antes con .retain_grad()). Para inferencia, donde no hace falta ningún gradiente, dos formas equivalentes de desactivarlo:

with torch.no_grad():
    z = torch.matmul(x, w) + b   # z.requires_grad == False acá dentro

z_det = z.detach()                # mismo efecto, sobre un tensor puntual

Dos razones prácticas para desactivarlo: marcar partes del modelo como parámetros congelados (exactamente lo que pasa con el modelo base cuando entrenás solo un adapter LoRA), y acelerar el forward pass cuando no hace falta backward — construir y mantener el grafo tiene un costo real de memoria y cómputo que no vale la pena pagar en pura inferencia.

Son cuatro herramientas parecidas con alcances distintos, y conviene tener claro cuál va en cada caso:

with torch.inference_mode():
    logits = model(x)
logits.requires_grad_()            # RuntimeError: Inference tensors cannot be saved for backward...
loss = (logits * w).sum(); loss.backward()   # ...y tampoco se pueden meter en un grafo nuevo

Regla corta: inference_mode para servir, no_grad para evaluar mientras entrenás. Los dos son ortogonales a model.eval() — ver §5.

3. Los gradientes se acumulan — por diseño, no por bug

Cada llamada a .backward() suma el gradiente calculado al valor que ya había en .grad, en vez de reemplazarlo. Es intencional (útil para acumular gradientes de varios mini-batches antes de un solo step, una técnica real para simular batches más grandes de los que entran en memoria), pero significa que sin limpiar explícitamente entre steps, los gradientes de iteraciones distintas se mezclan:

optimizer.zero_grad()    # 1. limpiar gradientes de la iteración anterior
loss.backward()          # 2. calcular gradientes nuevos, sumándolos a .grad
optimizer.step()         # 3. ajustar parámetros usando esos gradientes

Estos tres pasos, en este orden exacto (o con zero_grad() al final, algunas variantes lo ubican ahí), son el corazón de cualquier training loop de PyTorch que vayas a ver — desde un nn.Linear de dos capas hasta el fine-tuning de un LLM de 70B parámetros con LoRA.

Acumulación de gradientes, en código

La "técnica real" merece verse escrita, porque es la forma en que entrenás un batch efectivo de 32 en una GPU donde solo entran 8. La clave es dividir la loss por la cantidad de micro-batches: la loss promedio de un batch de 32 es el promedio de las losses de cuatro batches de 8, y como el gradiente es lineal, sumar cuatro gradientes de loss/4 da exactamente el gradiente del batch grande (con la salvedad de BatchNorm, que ve estadísticas de 8 y no de 32):

accum = 4                                        # batch efectivo = batch_size * accum
optimizer.zero_grad()
for i, (X, y) in enumerate(dataloader):
    loss = loss_fn(model(X), y) / accum          # dividir: así la suma de 4 backward == 1 batch grande
    loss.backward()                              # suma sobre .grad, no reemplaza
    if (i + 1) % accum == 0:
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        optimizer.step()
        optimizer.zero_grad()

clip_grad_norm_ va siempre en el mismo lugar: después del último backward() y antes de step(). Reescala todos los gradientes juntos para que su norma global no supere max_norm (1,0 es el valor habitual en transformers), y devuelve la norma que había antes de recortar — un número que vale la pena loggear, porque un pico repentino ahí es el primer aviso de que el entrenamiento se va a desestabilizar. Es la misma gradient_accumulation_steps + max_grad_norm del Trainer de HF y de los TrainingArguments que usaste en el Capítulo 3.

4. El loop de optimización completo

Con loss, optimizer y los tres pasos de arriba, el loop completo de la documentación oficial:

def train_loop(dataloader, model, loss_fn, optimizer):
    model.train()
    for batch, (X, y) in enumerate(dataloader):
        pred = model(X)
        loss = loss_fn(pred, y)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

def test_loop(dataloader, model, loss_fn):
    model.eval()
    test_loss, correct = 0, 0
    with torch.no_grad():
        for X, y in dataloader:
            pred = model(X)
            test_loss += loss_fn(pred, y).item()
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()

nn.CrossEntropyLoss combina internamente LogSoftmax + NLLLoss (Negative Log Likelihood) — por eso el modelo devuelve logits crudos, sin softmax aplicado a mano, y es la loss la que se encarga de normalizarlos. Entre los optimizadores disponibles (torch.optim), SGD es el más simple; Adam y RMSProp suelen converger más rápido en la práctica para la mayoría de arquitecturas modernas.

Fijate en el .item() de test_loop: no es cosmético. loss_fn(pred, y) es un tensor con grad_fn (aun bajo no_grad es un tensor, y fuera de no_grad arrastra el grafo entero); .item() lo convierte a un float de Python y suelta todo. Un running_loss += loss sin .item() en el loop de entrenamiento guarda una referencia al grafo de cada iteración: la memoria crece batch a batch hasta el OOM, y el mensaje de error no menciona el logging por ningún lado.

5. model.train() vs model.eval(): un interruptor, no cosmética

La documentación oficial insiste en un detalle que es fácil pasar por alto: llamar model.train() antes de entrenar y model.eval() antes de evaluar no es una convención cosmética. Capas como BatchNorm y Dropout se comportan de forma distinta en cada modo — BatchNorm usa estadísticas del batch actual en training pero un promedio acumulado en eval; Dropout apaga neuronas al azar en training pero las deja todas activas en eval. Olvidarse de este switch produce resultados de inferencia inconsistentes y difíciles de diagnosticar, porque el modelo "funciona", solo que con números ligeramente distintos cada vez de una forma que no tiene una causa obvia hasta que se revisa este detalle puntual.

model.eval() y torch.no_grad() se confunden porque siempre aparecen juntos, pero son dos cosas independientes que hay que hacer las dos:

Y el olvido simétrico: después de evaluar, model.train() de vuelta. El Trainer de HF y el Learner de fastai lo hacen por vos; en un loop escrito a mano, la evaluación de mitad de epoch que deja el modelo en eval es un bug clásico — el entrenamiento sigue, sin Dropout y con BatchNorm congelada, y nadie lo nota hasta que las curvas se ven raras.

6. Lo que el tutorial no cubre: precisión mixta y schedulers

Precisión mixta: autocast + (a veces) GradScaler

Entrenar todo en fp32 desperdicia la mitad del ancho de banda y los tensor cores. La receta estándar es mixed precision: los parámetros y el optimizador quedan en fp32 (los master weights), pero el forward y el backward corren en 16 bits donde es seguro. torch.autocast decide por operación — matmuls y convoluciones en bf16/fp16, reducciones, softmax y normas en fp32 — sin que toques el modelo:

scaler = torch.amp.GradScaler("cuda")                       # solo hace falta con float16

for X, y in dataloader:
    with torch.autocast(device_type="cuda", dtype=torch.float16):
        loss = loss_fn(model(X), y)                       # forward en 16 bits donde es seguro
    scaler.scale(loss).backward()                         # backward sobre la loss escalada
    scaler.unscale_(optimizer)                            # des-escalar ANTES de clipear
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    scaler.step(optimizer)                                # salta el step si hubo inf/nan
    scaler.update()                                       # ajusta el factor de escala
    optimizer.zero_grad()

GradScaler existe por el rango chico de fp16 (ver la tabla de dtypes de la página 1): multiplica la loss por un factor grande para que los gradientes pequeños no se vayan a cero, los divide de vuelta antes del step, y si detecta inf/nan salta ese step y baja el factor. Con bfloat16 no hace falta: tiene el mismo rango que fp32, así que el loop queda en dos líneas — y es la razón por la que bf16 es el default en cualquier GPU Ampere o más nueva:

with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
    loss = loss_fn(model(X), y)
loss.backward()                                           # sin scaler
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
optimizer.step(); optimizer.zero_grad()

Esto es lo que activan bf16=True / fp16=True en TrainingArguments. Distinto de cargar el modelo directamente en bf16 (dtype=torch.bfloat16): ahí no hay master weights en fp32, las actualizaciones chicas del optimizador se pierden por redondeo, y es aceptable para inferencia o para LoRA (donde los pesos que se actualizan son pocos y se pueden mantener en fp32), pero no para full fine-tuning serio.

Schedulers: el learning rate no es una constante

Un scheduler cambia el learning rate del optimizador a lo largo del entrenamiento. Dos que cubren casi todo: OneCycleLR sube el lr linealmente desde un valor bajo hasta max_lr y después baja (por defecto con coseno) hasta casi cero — es el fit_one_cycle de fastai, y se llama por batch. CosineAnnealingLR baja el lr siguiendo medio coseno desde el valor inicial hasta eta_min a lo largo de T_max steps — combinado con un warmup lineal es la curva de casi todos los LLMs. La regla de ubicación: scheduler.step() va después de optimizer.step(), una vez por cada step real del optimizador (con acumulación de gradientes, adentro del mismo if):

steps_total = epochs * len(dataloader) // accum
scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=3e-4, total_steps=steps_total)
# o: CosineAnnealingLR(optimizer, T_max=steps_total, eta_min=1e-5)

for i, (X, y) in enumerate(dataloader):
    (loss_fn(model(X), y) / accum).backward()
    if (i + 1) % accum == 0:
        optimizer.step()
        scheduler.step()                  # después de optimizer.step(), nunca antes
        optimizer.zero_grad()
        lr_actual = scheduler.get_last_lr()[0]   # para loggear

Si invertís el orden, PyTorch avisa (Detected call of lr_scheduler.step() before optimizer.step()): el primer step usaría el lr equivocado. Y si construís OneCycleLR con un total_steps menor que los steps reales, el scheduler lanza una excepción al pasarse — mejor que fallar en silencio.

7. Errores comunes

Casi todos son de orden: una línea que falta o está en el lugar equivocado del loop.

8. Resumen

1. Autograd reconstruye el grafo de cómputo en cada forward pass — dinámico, no estático, lo que permite control de flujo arbitrario de Python dentro del modelo. backward() consume el grafo (libera las activaciones); una segunda pasada requiere retain_graph=True, y torch.autograd.grad es la versión funcional que no escribe en .grad.
2. no_grad()/detach() desactivan el tracking de gradientes — para congelar parámetros o acelerar inferencia pura; inference_mode() es la versión más estricta y rápida para servir.
3. Los gradientes se acumulan por diseño; zero_grad() es obligatorio en cada iteración salvo que quieras acumular a propósito — y entonces dividís la loss por el número de micro-batches y hacés el step cada N.
4. El loop de optimización son pasos fijos en orden fijo: backward(), clip_grad_norm_, step(), scheduler.step(), zero_grad() — el mismo patrón en cualquier escala de modelo. .item() en todo lo que loggees.
5. model.train()/model.eval() cambian el comportamiento real de BatchNorm y Dropout — no es un detalle estético, es funcional, y es independiente de no_grad: en evaluación van los dos.
6. Precisión mixta: autocast corre el forward en 16 bits con master weights en fp32; fp16 necesita GradScaler, bf16 no. AdamW es el optimizador por defecto de los transformers; OneCycle y coseno con warmup, las curvas de lr habituales.

9. Ponete a prueba

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

10. Para practicar

import torch, torch.nn as nn
torch.manual_seed(0)
model = nn.Sequential(nn.Linear(10, 32), nn.ReLU(), nn.Linear(32, 3))
loss_fn = nn.CrossEntropyLoss()
X, y = torch.randn(8, 10), torch.randint(0, 3, (8,))

# (a) un solo batch de 8
model.zero_grad()
loss_fn(model(X), y).backward()
g_full = [p.grad.clone() for p in model.parameters()]

# (b) dos micro-batches de 4, loss dividida por 2
model.zero_grad()
for Xi, yi in ((X[:4], y[:4]), (X[4:], y[4:])):
    (loss_fn(model(Xi), yi) / 2).backward()
g_accum = [p.grad.clone() for p in model.parameters()]

print(max((a - b).abs().max().item() for a, b in zip(g_full, g_accum)))

# (c) sin dividir: ¿cuánto más grande queda?
model.zero_grad()
for Xi, yi in ((X[:4], y[:4]), (X[4:], y[4:])):
    loss_fn(model(Xi), yi).backward()
print(model[0].weight.grad.norm() / g_full[0].norm())

loss = loss_fn(model(X), y)
loss.backward()
loss.backward()                              # (a)

with torch.inference_mode():
    out = model(X)
loss_fn(out, y).backward()                   # (b)

total = 0
for _ in range(5):
    total += loss_fn(model(X), y)
print(total.grad_fn, total.item())           # (c): hay grad_fn → el grafo de las 5 iteraciones sigue vivo

import torch, torch.nn as nn
lin = nn.Linear(64, 64)
x = torch.randn(8, 64)
with torch.autocast(device_type="cpu", dtype=torch.bfloat16):
    h = lin(x)
    s = h.float().softmax(-1)
    loss = h.sum()
print("activación:", h.dtype, "| parámetro:", lin.weight.dtype)
loss.backward()
print("gradiente:", lin.weight.grad.dtype)

11. Fuentes y para seguir

• Automatic Differentiation with torch.autograd — Learn the Basics — el tutorial original de §1 y §2, con la explicación del DAG y el ejemplo de Jacobiano.
• Optimizing Model Parameters — Learn the Basics — el loop completo de §4 con hiperparámetros, epochs y la salida de entrenamiento de FashionMNIST.
• Autograd mechanics — la nota oficial: cómo se guardan los tensores intermedios, no_grad vs inference_mode, in-place y contadores de versión. Leela cuando un error de autograd no te cierre.
• Automatic Mixed Precision examples — las recetas oficiales de autocast + GradScaler, incluidas las variantes con acumulación de gradientes y clipping (el orden de unscale_ viene de acá).
• Mixed Precision Training (Micikevicius et al., 2017, arXiv:1710.03740) — el paper que introdujo master weights en fp32 + loss scaling; explica con histogramas de gradientes por qué fp16 necesita el scaler.
• Decoupled Weight Decay Regularization (Loshchilov & Hutter, 2017, arXiv:1711.05101) — el paper de AdamW: por qué el weight decay "clásico" de Adam no hace lo que uno cree y cómo desacoplarlo.