English · Español

02 — El motor de autograd¶

El motor de autograd de PyTorch es exactamente lo que construimos en Fases 7–8: forward registra un grafo, backward lo recorre en reversa. Esta página formaliza la captura, los nodos grad_fn, las hojas, y muestra cómo torch.library.custom_op registra un backward para una operación nueva. El ejemplo corriente es linear(x, W, b) con el LM head del grammar MiniGPT.

Esta página es el motor de autograd, hecho explícito. Tras ella podrás recorrer la cadena grad_fn para cualquier forward, derivar el backward a mano y verificar tu derivación contra .backward() de PyTorch hasta el acuerdo numérico a fp32.

El modelo de dos líneas¶

forward:  every op on a requires_grad tensor records a grad_fn node into the graph.
backward: .backward() walks the graph in reverse-topological order, calling each grad_fn's backward formula.

Ya está. La complejidad reside en (a) qué ops registran, (b) cuál es la fórmula de backward de cada una, © cómo se almacena el grafo, (d) cuándo se libera. La Fase 25 cubre cada una.

El forward: captura del grafo¶

x = torch.randn(2, 64, requires_grad=True)
W = torch.randn(600, 64, requires_grad=True)
b = torch.randn(600, requires_grad=True)

y = torch.nn.functional.linear(x, W, b)   # y.grad_fn = <AddmmBackward0>
loss = y.sum()                             # loss.grad_fn = <SumBackward0>

El motor de autograd registra dos nodos — AddmmBackward0 y SumBackward0 — enlazados por una arista. El grafo en este punto:

x ─┐
W ─┤
b ─┴→ AddmmBackward0 → y → SumBackward0 → loss

Cada nodo contiene:

Tensores guardados necesarios para su fórmula de backward. Para Addmm: guarda x y W (necesarios para el gradiente del matmul).
Aristas a nodos padres (o a hojas). Cada arista sabe qué salida del padre alimenta qué entrada de este nodo.
Puntero a la función de backward — la implementación en C++ de la fórmula del gradiente.

Las hojas (x, W, b) tienen grad_fn = None (no fueron producidas por una op) y is_leaf = True. Son las salidas de .backward(): los gradientes se acumulan en .grad sobre las hojas.

El backward: recorrido en reversa¶

.backward():

Empieza con loss (un escalar por defecto; si no, pasas gradient=ones_like(loss)).
Llama a SumBackward0.backward(grad=1.0) → devuelve dy = ones_like(y).
Llama a AddmmBackward0.backward(grad=dy) → devuelve tres gradientes (uno por entrada):
dx = dy @ W (forma (2, 64))
dW = dy.T @ x (forma (600, 64))
db = dy.sum(dim=0) (forma (600,))
Cada gradiente se suma al .grad de la hoja correspondiente.

loss.backward()
print(x.grad.shape, W.grad.shape, b.grad.shape)   # (2, 64), (600, 64), (600,)

El recorrido es topológico inverso. Los ciclos están prohibidos (autograd lanza error si detecta uno — raro pero posible con hooks).

Derivar `AddmmBackward0` a mano¶

El forward: \(y = b + x W^T\) (con broadcasting sobre \(b\)).

Para una pérdida escalar \(L = \sum y\):

\[\frac{\partial L}{\partial y_{ij}} = 1 \quad \text{(from sum)}\]

Regla de la cadena:

\[\frac{\partial L}{\partial x_{ik}} = \sum_j \frac{\partial L}{\partial y_{ij}} \cdot \frac{\partial y_{ij}}{\partial x_{ik}} = \sum_j 1 \cdot W_{jk} = \sum_j W_{jk}\]

En forma matricial: \(\nabla_x L = (\nabla_y L) W\), forma (2, 600) @ (600, 64) = (2, 64). ✓

\[\frac{\partial L}{\partial W_{jk}} = \sum_i \frac{\partial L}{\partial y_{ij}} \cdot \frac{\partial y_{ij}}{\partial W_{jk}} = \sum_i 1 \cdot x_{ik} = \sum_i x_{ik}\]

En forma matricial: \(\nabla_W L = (\nabla_y L)^T x\), forma (600, 2) @ (2, 64) = (600, 64). ✓

\[\frac{\partial L}{\partial b_j} = \sum_i \frac{\partial L}{\partial y_{ij}} \cdot \frac{\partial y_{ij}}{\partial b_j} = \sum_i 1 \cdot 1 = 2\]

En forma matricial: \(\nabla_b L = (\nabla_y L).\text{sum}(\text{dim}=0)\), forma (600,). Igual al tamaño de batch 2 (cada posición de salida sumada sobre la dim de batch).

Verifica en PyTorch:

loss.backward()
assert torch.allclose(x.grad, torch.ones_like(y) @ W)
assert torch.allclose(W.grad, torch.ones_like(y).T @ x)
assert torch.allclose(b.grad, torch.ones_like(y).sum(dim=0))

Si esos pasan — y lo hacen, a 1e-7 en fp32 — has replicado a mano la fórmula de AddmmBackward0. El laboratorio 01 te hace hacer este ejercicio.

Éste es el contenido entero del motor de autograd: captura de grafo en forward, recorrido en reversa en backward, cada nodo conociendo su derivada. La Fase 7 implementó esto para escalares; la Fase 8 para tensores; la versión de PyTorch es la misma idea a escala.

Tensores guardados y memoria¶

Cada grad_fn guarda los tensores que necesita para el backward. AddmmBackward0 guarda x y W (no b — su gradiente no depende de b). Los tensores guardados aumentan la memoria pico durante el entrenamiento (se mantienen vivos hasta que se ejecuta el backward).

Optimizaciones:

torch.utils.checkpoint: recomputar los tensores guardados en lugar de almacenarlos. Cambia cómputo por memoria.
torch.no_grad(): saltarse la construcción del grafo por completo. Usado en inferencia.
.detach(): produce un tensor nuevo con requires_grad=False, rompiendo el grafo en ese punto.

El laboratorio 01 mide la memoria pico con y sin torch.no_grad() para un forward a través del grammar MiniGPT.

¿Cuándo se libera el grafo?¶

Tras completarse .backward() — por defecto. Los tensores guardados se liberan. Si necesitas llamar a .backward() dos veces sobre el mismo grafo, usa retain_graph=True.

Olvidar retain_graph=True cuando hace falta es un mensaje de error común: "Trying to backward through the graph a second time". El motor de autograd libera con avidez los tensores guardados para ahorrar memoria.

Autograd personalizado: `torch.library.custom_op`¶

La API moderna (torch 2.1+) para registrar una op personalizada con autograd:

import torch
from torch import Tensor

@torch.library.custom_op("mylib::softmax_triton", mutates_args=())
def softmax_triton(x: Tensor) -> Tensor:
    # Implementation (calls into Triton kernel; Phase 24's softmax).
    return triton_softmax_impl(x)

# Shape-inference for torch.compile / FakeTensor:
@softmax_triton.register_fake
def _(x):
    return torch.empty_like(x)

# Backward formula:
def softmax_triton_backward(ctx, grad_output):
    y = ctx.saved_tensors[0]
    # Softmax backward: dy = y * (dL/dy - sum(y * dL/dy, dim=-1, keepdim=True))
    return y * (grad_output - (y * grad_output).sum(dim=-1, keepdim=True))

def softmax_triton_setup_context(ctx, inputs, output):
    ctx.save_for_backward(output)

softmax_triton.register_autograd(
    softmax_triton_backward,
    setup_context=softmax_triton_setup_context,
)

Ahora torch.ops.mylib.softmax_triton(x) se comporta como una op nativa de PyTorch:

El dispatcher la encuentra.
Autograd registra SoftmaxTritonBackward en el grafo durante el forward.
.backward() invoca la fórmula registrada.
torch.compile puede trazarla (gracias a la inferencia de formas de register_fake).
gradcheck valida la fórmula de backward numéricamente.

El laboratorio 02 recorre este registro exacto.

`torch.autograd.gradcheck`¶

La herramienta de PyTorch para verificar una implementación de backward:

from torch.autograd import gradcheck
x = torch.randn(4, 8, dtype=torch.float64, requires_grad=True)
gradcheck(torch.ops.mylib.softmax_triton, (x,), eps=1e-6, atol=1e-4)

gradcheck estima numéricamente el gradiente (mediante diferencias finitas) y compara con el backward analítico. Si no coinciden, la fórmula de backward está mal. Usa entradas en fp64 (fp32 tiene demasiado ruido para verificaciones por diferencias finitas).

Éste es el primer test que ejecutas tras registrar un backward personalizado. Si gradcheck falla, tu fórmula de backward tiene un bug; debugéalo antes de integrarlo en un bucle de entrenamiento real.

Errores comunes de autograd¶

Error	Causa	Solución
`RuntimeError: grad can be implicitly created only for scalar outputs`	Se llamó a `.backward()` sobre una salida no escalar	Pasa `gradient=ones_like(y)` o llama a `.sum().backward()`
`RuntimeError: Trying to backward through the graph a second time`	Llamar a `.backward()` dos veces	`retain_graph=True` en la primera llamada
`RuntimeError: ... is at version N; expected version M`	Op in-place sobre un tensor guardado	Evita ops in-place, o `.clone()` antes
`grad_fn=None` en una no-hoja	El tensor se creó dentro de `torch.no_grad()` o con `.detach()`	Recrea con el tracking de grad activado
Falla `gradcheck`	Fórmula de backward errónea, o tensores guardados erróneos	Re-deriva en papel; verifica el contexto guardado

Lo que deberías ahora ser capaz de hacer¶

Recorrer la cadena grad_fn del forward de cualquier modelo.
Derivar la fórmula de backward para cualquier composición de linear, relu, softmax, cross_entropy.
Usar torch.library.custom_op para registrar una nueva op con backward.
Usar gradcheck para validar numéricamente el backward.
Predecir si la memoria pico de un modelo está dominada por parámetros, activaciones o tensores guardados.

Lo que esta página NO cubre¶

torch.autograd.Function (la API antigua). Mencionada; el laboratorio 02 usa la API moderna custom_op exclusivamente.
__torch_dispatch__ para interceptación de autograd. De nicho; sólo relevante si estás construyendo un framework paralelo sobre PyTorch.
Gradientes de segundo orden (create_graph=True). Usado en meta-learning; fuera del alcance del currículo.
Transformaciones funcionales torch.func (grad, vmap). La Fase 38 puede revisitarlo.

Siguiente: theory/03-compile-and-distributed.md — el pipeline de compile + survey distribuido.