Skip to content

English · Español

04 — Recorrido del grafo de autograd: un backward a través del mini-GPT, cinta de funciones expuesta

Hacemos un paso forward y backward del mini-GPT de la Fase 17 en PyTorch, y enumeramos la cinta (tape) que autograd construye. Cada operación deja un Function nodo en el grafo; al hacer .backward(), autograd recorre la cinta en orden inverso. Aquí lo vemos en código real y en lista de pasos.

Este archivo extiende theory/02-autograd-engine.md con la cinta de funciones completa para un backward a través del mini-GPT de la Fase 17. Es un ejercicio en profundidad de lectura del grafo de autograd de PyTorch a una escala lo bastante pequeña como para enumerarla.

Setup

Coge el mini-GPT de la Fase 17 como un torch.nn.Module:

  • \(d_\text{model} = 64\), \(n_\text{heads} = 4\), \(d_h = 16\), \(n_\text{layers} = 2\), \(d_\text{ff} = 256\).
  • Secuencia única: \(T = 8\) tokens.
  • Batch: \(B = 1\).

Forward de un batch:

logits = model(input_ids)        # (1, 8, V)
loss = F.cross_entropy(
    logits.view(-1, V),          # (8, V)
    target_ids.view(-1),         # (8,)
)
loss.backward()

El grafo de autograd para loss.backward() es el que vamos a recorrer.

La cinta del forward — operación a operación

Cada torch.Tensor.grad_fn registra la operación que lo produjo. Recorriendo el grafo en orden forward (de arriba abajo):

  1. input_ids → Embedding.forward → EmbeddingBackward. Forma de salida: (1, 8, 64).
  2. Suma broadcast del position-embed → AddBackward. Forma de salida: (1, 8, 64).

Para cada uno de los \(L = 2\) bloques:

  1. LayerNorm.forward (pre-attn) → NativeLayerNormBackward. Salida (1, 8, 64).
  2. Proyección Q (Linear(64, 64)) → AddmmBackward. Salida (1, 8, 64).
  3. Proyección K → AddmmBackward. Salida (1, 8, 64).
  4. Proyección V → AddmmBackward. Salida (1, 8, 64).
  5. Reshape Q a (1, 4, 8, 16) → ViewBackward. (Sin movimiento real de datos; sólo metadata.)
  6. Reshape K, V igual → 2× ViewBackward.
  7. Q @ K^T (torch.matmul) → BmmBackward. Salida (1, 4, 8, 8).
  8. Escala por 1/sqrt(d_h) = 1/4 → MulBackward. Salida (1, 4, 8, 8).
  9. Suma de máscara causal (broadcasting -inf triangular superior) → AddBackward.
  10. softmax → SoftmaxBackward. Salida (1, 4, 8, 8).
  11. attn @ V → BmmBackward. Salida (1, 4, 8, 16).
  12. Reshape de vuelta a (1, 8, 64) → ViewBackward.
  13. Proyección de salida (Linear(64, 64)) → AddmmBackward. Salida (1, 8, 64).
  14. Suma residual (entrada del paso 2 o del bloque previo) → AddBackward.
  15. LayerNorm.forward (pre-FFN) → NativeLayerNormBackward.
  16. Primer Linear de FFN (Linear(64, 256)) → AddmmBackward. Salida (1, 8, 256).
  17. GELU → GeluBackward. Salida (1, 8, 256).
  18. Segundo Linear de FFN (Linear(256, 64)) → AddmmBackward. Salida (1, 8, 64).
  19. Suma residual → AddBackward.

Para el bloque 2, repite los pasos 3-21.

Tras ambos bloques:

  1. LayerNorm final → NativeLayerNormBackward. Salida (1, 8, 64).
  2. LM head (Linear(64, V=512)) → AddmmBackward. Salida (1, 8, 512).

Para la pérdida:

  1. view(-1, V) → ViewBackward. Salida (8, 512).
  2. cross_entropy(logits, targets) se descompone en:
    • log_softmax: → LogSoftmaxBackward. Salida (8, 512).
    • NLLLoss: → NllLossBackward. Salida () (escalar).

Así que la cinta del forward para un mini-GPT de 2 capas con \(T = 8, V = 512\) tiene aproximadamente 41 funciones de backward (~20 por bloque + ~5 fuera). Cada una es un nodo en el grafo de autograd.

Lo que hace loss.backward() — nodo a nodo

loss.backward() invoca engine.execute(...), que:

  1. Empieza con grad_output = torch.ones_like(loss) (es decir, 1.0 para la pérdida escalar).
  2. Recorre el grafo en orden topológico inverso, empezando desde loss.grad_fn = NllLossBackward.
  3. Para cada nodo, llama a su método backward con el grad_output actual, recupera los grad_input(s), los acumula en el .grad de los tensores hoja correspondientes (parámetros / entradas marcadas requires_grad=True).

Recorrido en orden inverso (leer de abajo arriba):

NllLossBackward(grad=1.0) → ∂loss/∂logits = (softmax - one_hot) / N, shape (8, 512)
LogSoftmaxBackward → (the cross-entropy gradient is folded in)
ViewBackward → (8, 512) reshaped back to (1, 8, 512)
AddmmBackward (LM head):
    ∂loss/∂W_lm_head = h_final^T @ (softmax - one_hot)/N, shape (V, 64)
    ∂loss/∂b_lm_head = sum_T (softmax - one_hot)/N, shape (V,)
    ∂loss/∂h_final = (softmax - one_hot)/N @ W_lm_head, shape (1, 8, 64)
NativeLayerNormBackward (final LN):
    gradient w.r.t. input + gradient w.r.t. scale + gradient w.r.t. bias
... (block 2's full chain in reverse) ...
... (block 1's full chain in reverse) ...
AddBackward (position-embed) → grad gets routed to both inputs
EmbeddingBackward → updates embedding rows for the input tokens

Cada nodo *Backward sabe exactamente qué transformación de gradiente aplicar porque capturó las entradas relevantes del forward (tensores guardados) durante el forward.

Tensores guardados — qué guarda cada nodo

Un detalle: los nodos de autograd guardan justo lo suficiente del estado del forward para calcular el backward. Por ejemplo:

  • MulBackward(a, b) guarda tanto a como b (porque \(\partial(ab)/\partial a = b\) los necesita ambos).
  • AddBackward no guarda nada (porque \(\partial(a+b)/\partial a = 1\) no necesita entradas).
  • SoftmaxBackward guarda la salida (porque el Jacobiano del softmax se expresa en términos de softmax(x), no x).
  • GeluBackward guarda la entrada (porque la derivada de GELU es función de x).
  • AddmmBackward guarda las entradas (input, weight, bias) — las tres.

Para el mini-GPT, la memoria total de tensores guardados es aproximadamente:

\[\text{saved} \approx \sum_\text{forward nodes} \text{(input + maybe output size)} \times s\]

Para \(T = 8, d_\text{model} = 64, V = 512, L = 2\):

  • Salida de embedding guardada en el paso 1: \(8 \cdot 64 \cdot 4 = 2\) KiB.
  • Salidas de LayerNorm (×6 a través de bloques): \(8 \cdot 64 \cdot 4 \cdot 6 = 12\) KiB.
  • Intermedios de attention (salidas QKV, scores, pesos de attn, salida de attn): \(\approx 50\) KiB.
  • Intermedios de FFN: \(\approx 30\) KiB.
  • Salida del LM head: \(8 \cdot 512 \cdot 4 = 16\) KiB.

Total: \(\approx 130\) KiB. La memoria de activación para un forward a \(T = 8\). Esto coincide aproximadamente con el tamaño del KV-cache de la Fase 22 (theory/05-mini-gpt-memory-worked-example.md) — ambos son \(O(T \cdot d_\text{model} \cdot L)\), sólo con constantes distintas.

Para entrenamiento, las activaciones dominan la memoria (no los pesos, no el KV cache). A escala GPT-3, ésta es la razón por la que existe el gradient checkpointing. A escala §A13, las activaciones son 130 KiB y no necesitamos checkpointing — pero la metodología es idéntica.

Visualizando el grafo

PyTorch tiene torchviz:

from torchviz import make_dot
graph = make_dot(loss, params=dict(model.named_parameters()))
graph.render("mini_gpt_backward")  # produces a .pdf

La salida es un DAG con ~50 nodos (40 de backward + parámetros + tensores de entrada). A escala §A13, el grafo es legible en una página. A escala GPT-2 (12 capas), es un muro.

Gotchas comunes que la cinta expone

  1. Las ops in-place rompen la cinta. Si haces h.add_(residual) en lugar de h = h + residual, el nodo de autograd para el h original se sobrescribe. Si alguien aguas abajo de h necesitaba el valor original, el backward falla con RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation. El break 00-... de la Fase 25 provoca esto a propósito.
  2. Detach trunca la cinta. h.detach() devuelve un tensor sin grad_fn. Útil para "quiero este valor pero no propagar gradientes a través de él" (por ejemplo, la Q-network objetivo en RL). Usado incorrectamente, los gradientes desaparecen en silencio.
  3. with torch.no_grad(): no registra ningún grad_fn. Usado en inferencia. Si se envuelve alrededor de un paso de entrenamiento, no se computan gradientes y .backward() lanza error.
  4. Retain graph. Por defecto, .backward() libera los tensores guardados tras el backward. Si llamas a .backward() dos veces sobre la misma pérdida, la segunda llamada falla. Usa loss.backward(retain_graph=True) para mantenerlos — pero sólo si de verdad lo necesitas.

Por qué esto importa

En la Fase 25 ya has implementado autograd a mano dos veces: versión escalar (Fase 7), versión tensorial (Fase 8). El autograd de PyTorch es simplemente la versión de calidad de producción de lo que ya escribiste. Recorrer la cinta del mini-GPT muestra:

  • El grafo es finito, de tamaño acotado e inspeccionable.
  • Cada concepto mapea a algo que has escrito.
  • La "magia de autograd" siempre fue matemáticas; la magia es la contabilidad.

Para el agente tutor de gramática de la Fase 32 (en la propia Fase 32, no escribimos un nuevo autograd; usamos el de PyTorch). Conocer la cinta te permite debugéar "por qué mi pérdida personalizada no entrena" leyendo el grafo en lugar de adivinando.

Cita

Paszke, A. et al. (2017). Automatic differentiation in PyTorch. NeurIPS 2017 Autodiff Workshop. https://openreview.net/forum?id=BJJsrmfCZ — el paper original de autograd de PyTorch. La sección 3 describe la cinta; la sección 4 discute el grafo dinámico. El recorrido del mini-GPT de arriba es una instancia del diseño que describen.

Recap en un párrafo

Un forward a través del mini-GPT de 2 capas a \(T = 8\) produce un grafo de autograd de ~40 nodos; .backward() recorre el grafo en reversa, propagando \(\partial L / \partial t\) desde cada nodo a sus entradas usando el Jacobiano local. Cada nodo guarda justo lo suficiente del estado del forward para su backward; a escala §A13, la memoria total de tensores guardados es ~130 KiB. Las ops in-place, .detach(), torch.no_grad(), y el ciclo de vida del tensor guardado son los cuatro mecanismos que cambian el comportamiento del grafo. El ejercicio de break de la Fase 25 demuestra el peligro de in-place sobre este grafo exacto.

Cross-refs: theory/02-autograd-engine.md (la maquinaria de autograd en general), Fase 7-8 (las versiones hechas a mano), README.md de la Fase 17 (la spec del mini-GPT), Fase 22 theory/05-... (memoria de activación vs KV cache).