Curso / fast.ai / Parte 2 · Unidad 1
fast.ai · 10 de 14

fast.ai · Parte 2 · Unidad 1 de 5

De Python puro a un framework de entrenamiento propio

Acá arranca la Parte 2: nada de fastai ni vision_learner — se construye, desde la multiplicación de matrices hasta un framework de callbacks completo, todo lo que la Parte 1 usaba como caja negra. Esta unidad condensa cinco lecciones de código denso en tres ideas centrales.

Antes de leerlección 3, SGD y backprop a mano · PyTorch 3 §1, el grafo de autograd

🎯 Al terminar vas a poder

  • Explicar con tus palabras por qué la misma multiplicación de matrices puede tardar medio segundo o medio milisegundo, y en qué capa (Python, C vectorizado, GPU) vive cada factor.
  • Aplicar las tres reglas de broadcasting a un par de shapes cualquiera y predecir el shape del resultado — o detectar que va a fallar — antes de correrlo.
  • Leer y escribir una expresión einsum sabiendo qué índices se suman y cuáles quedan.
  • Escribir el backward de una capa lineal (lin_grad) con las tres fórmulas y sus shapes, y verificarlo contra autograd.
  • Reconstruir el esqueleto del Learner de miniai (with_cbs, run_cbs, excepciones Cancel*) y escribir un callback nuevo — un LR finder — sin tocar el loop.

🟢 Por qué esta sección se trata distinto

La Parte 2 (Lecciones 9 a 25, "Deep Learning Foundations to Stable Diffusion") no tiene libro de referencia — son ~30 notebooks de código puro construyendo un mini-framework llamado miniai y, con él, Stable Diffusion desde cero. Es mucho más código que prosa. Estas páginas la condensan en 5 unidades temáticas con los ejemplos más ilustrativos, en vez de una página por lección — el objetivo es cubrir el terreno conceptual completo sin reproducir cada notebook línea por línea.

1. Matmul desde cero, y por qué importa la velocidad

El punto de partida es literal: multiplicar dos matrices con tres loops anidados de Python puro, sobre MNIST.

def matmul(a,b):
    (ar,ac),(br,bc) = a.shape,b.shape
    c = torch.zeros(ar, bc)
    for i in range(ar):
        for j in range(bc):
            for k in range(ac): c[i,j] += a[i,k] * b[k,j]
    return c

Sobre una matriz 5×784 por 784×10, esto tarda ~500ms. Son solo 5·784·10 = 39.200 multiplicaciones — un procesador moderno hace miles de millones por segundo — así que el tiempo no está en la aritmética, está en el intérprete: cada c[i,j] += a[i,k] * b[k,j] es media docena de llamadas a métodos de Python, cada una creando un tensor de un elemento. El resto de la unidad es una escalera de optimizaciones, cada una reemplazando un loop por una operación vectorizada, midiendo la mejora real en cada paso:

EscalónQué reemplazaCómo se escribeQuién hace el trabajo
Python puro—tres forel intérprete, elemento por elemento
Numbael loop más interno (sobre k)@njit def dot(a,b) sobre arrays de NumPycódigo máquina compilado en la primera llamada
Broadcastinglos dos loops internosc[i] = (a[i,:,None] * b).sum(dim=0)el kernel en C de PyTorch, una fila por vez
einsum / @los tres loopstorch.einsum('ik,kj->ij', a, b) o a @ bBLAS: una rutina de matmul optimizada por décadas
CUDAlos tres loops, en paralelokernel con @cuda.jit, un thread por celda c[i,j]miles de threads de GPU a la vez

El resultado medido en el notebook: de ~500ms (Python puro) a ~0.5ms (CUDA) — una mejora de aproximadamente 5 millones de veces sobre todo el dataset. La lección no es "usá broadcasting" per se — es que cada capa de abstracción entre vos y el hardware (Python interpretado → C vectorizado → kernels paralelos en GPU) tiene un costo medible y real, y entender dónde vive ese costo es lo que te permite decidir cuándo vale la pena bajar un nivel.

Qué loop deja de correr en Python en cada escalón del matmul Los tres loops anidados del matmul en Python puro. Numba compila el loop interno sobre k; broadcasting reemplaza los dos loops internos con un kernel en C por fila; einsum o @ delegan los tres a BLAS; CUDA reparte los tres entre miles de threads. De 500 ms a 0,5 ms sin cambiar la matemática. LOS TRES LOOPS DEL MATMUL EN PYTHON for i in range(ar): for j in range(bc): for k in range(ac): c[i,j] += a[i,k]*b[k,j] 5×784 por 784×10: 39.200 productos, ~500 ms el tiempo no está en la aritmética: cada += pasa por el intérprete y crea un tensor de 1 elemento ESCALÓN SALEN DE PYTHON QUIÉN HACE EL TRABAJO i j k Python puro tres for el intérprete, elemento por elemento ~500 ms Numba @njit dot(a[i], b[:,j]) código máquina compilado (la 1ª llamada compila) broadcasting (a[i,:,None]*b).sum(0) el kernel en C de PyTorch, una fila por vez einsum / @ 'ik,kj->ij' BLAS: matmul optimizado por décadas CUDA @cuda.jit un thread por c[i,j] miles de threads de GPU a la vez ~0,5 ms ~500 ms → ~0,5 ms por la misma matemática; sobre todo el dataset, 5 millones de veces medir la 2ª llamada de @njit (la 1ª compila) · sincronizar la GPU antes de cronometrar
La escalera del notebook 01, leída como quién ejecuta cada loop. Cada escalón saca un for más del intérprete y se lo da a algo que lo hace en bloque: código compilado, un kernel en C, BLAS o miles de threads. La multiplicación es la misma en los cinco.

Dos detalles de medición que el notebook recalca y que conviene grabarse: la primera llamada a una función con @njit incluye la compilación (puede tardar más que la versión de Python), así que se mide la segunda; y las operaciones en GPU son asincrónicas, de modo que cronometrar sin sincronizar mide cuánto tarda encolar el trabajo, no hacerlo.

⚡ Física de hardware: Strides, líneas de caché y 2D block tiling en GEMM

¿Por qué tres loops anidados en Python tardan 500 ms mientras que un kernel optimizado tarda 0.5 ms? La explicación completa combina la sobrecarga del intérprete con la física de la memoria del silicio:

  1. Layout C-Contiguo y Cache Thrashing: En PyTorch y C, los tensores se disponen en memoria por filas (row-major). La dirección física del elemento Ai,k es baseA+(i⋅ac+k)×4 bytes (en float32). En el bucle más interno sobre k, la lectura de Ai,k tiene stride 1 (elementos contiguos). Cuando la CPU solicita un flotante, carga una línea de caché completa de 64 bytes (16 floats). Los siguientes 15 accesos son cache hits en L1 a velocidad de nanosegundos.
    Sin embargo, para Bk,j, fijado j e incrementando k, la dirección salta de fila en fila: baseB+(k⋅bc+j)×4 bytes. ¡Cada paso salta bc×4 bytes! Si bc=784, cada iteración salta 3.136 bytes en memoria. Como cada salto supera los 64 bytes de la línea de caché, el 100% de los accesos a B son cache misses: la CPU se detiene cientos de ciclos esperando que DRAM cargue 64 bytes para usar uno solo y descartar los otros 15.
  2. Intensidad Aritmética y el Modelo Roofline: Multiplicar dos matrices de dimensiones M×K y K×N requiere exactamente 2MNK operaciones de punto flotante (un producto y una suma por celda). La intensidad aritmética I es el cociente entre operaciones y bytes transferidos: I=FLOPsBytes transferidos=2MNK4(MNK+MNK)=0.25 FLOP/Byte En una GPU como la H100 (3.35 TB/s de ancho de banda HBM3 y casi 2.000 TFLOPs FP16 de cómputo), saturar los Tensor Cores requiere una intensidad aritmética de al menos I≥2000×10123.35×1012≈597 FLOP/Byte. Con I=0.25, el procesador está el 99.9% del tiempo ocioso esperando datos de memoria (memory-bandwidth bound).
  3. 2D Block Tiling en cuBLAS/CUTLASS (SRAM on-chip): La solución de ingeniería en bibliotecas como cuBLAS o Triton es particionar las matrices en bloques bidimensionales (tiles) de tamaño TM×TK y TK×TN que caben en la memoria compartida (Shared Memory / SRAM, con más de 20 TB/s de ancho de banda). Los threads cargan un bloque cooperativamente a SRAM una sola vez, sincronizan con barreras de hardware, y reutilizan cada valor de A exactamente TN veces y cada valor de B exactamente TM veces: Itiling=2TMTNTK4(TMTK+TKTN)=TMTN2(TM+TN) FLOP/Byte Con bloques estándar de 128×128, la intensidad sube a 32 FLOP/Byte; y al llevarla a registros de los Tensor Cores mediante micro-tiles de 16×8, se sobrepasa la barrera del ancho de banda y la GPU alcanza su régimen de cálculo puro (compute-bound).

2. Broadcasting y einsum: las reglas detrás de la escalera

Los dos escalones del medio dependen de una idea que la Parte 1 usaba sin nombrar: operar entre tensores de shapes distintos sin escribir loops ni copiar datos. Broadcasting es un conjunto de reglas (heredadas de NumPy, idénticas en PyTorch) que decide si dos shapes son compatibles y cuál es el shape del resultado. Son exactamente tres:

  1. Alinear desde la derecha. Se comparan las dimensiones de atrás hacia adelante: la última con la última, la anteúltima con la anteúltima. Si un tensor tiene menos dimensiones, se le anteponen unos: (5,) contra (3,5) se trata como (1,5) contra (3,5).
  2. Dos dimensiones son compatibles si son iguales, o si una de ellas vale 1. Cualquier otro caso es error: (3,5) contra (4,5) no se puede.
  3. La dimensión de tamaño 1 se estira hasta el tamaño de la otra, sin copiar memoria — internamente PyTorch le pone stride 0, así que leer "la fila siguiente" vuelve a leer la misma. El resultado toma el máximo en cada posición.
a = torch.randn(5, 784); b = torch.randn(784, 10)
a[0].shape, a[0,:,None].shape          # (784,)  (784, 1): el None agrega un eje de tamaño 1
(a[0,:,None] * b).shape              # (784, 1) × (784, 10) → (784, 10): la columna se estira 10 veces
(a[0,:,None] * b).sum(0).shape       # (10,): la fila 0 del producto a @ b

c = torch.tensor([10.,20,30]); t = torch.ones(3,3)
c[:,None].expand_as(t).stride()      # (1, 0): la segunda dimensión no avanza en memoria

La versión con broadcasting del matmul es la regla 3 aplicada una vez por fila: a[i,:,None] tiene shape (784,1), se estira contra b de (784,10), el producto elemento a elemento es a[i,k]·b[k,j] para todos los k,j a la vez, y .sum(0) colapsa k. Los dos loops internos desaparecieron sin cambiar la matemática. La misma regla es la que hace funcionar x - x.mean(1, keepdim=True) (normalizar por fila), logits + mask en atención, o sumar un bias (n_out,) a una salida (batch, n_out).

⚠️ Cuidado

El broadcasting nunca avisa. torch.ones(5) + torch.ones(5,1) no da un vector de 5: por la regla 1 el primero se vuelve (1,5), por la 3 ambos se estiran, y el resultado es (5,5). Es el clásico bug de una loss que "entrena" pero no baja, porque la diferencia pred - target entre un (n,1) y un (n,) produjo una matriz (n,n) y el .mean() promedió todo contra todo. Cuando algo se comporta raro, imprimí los shapes de las dos partes de la operación, no solo el del resultado.

einsum es la forma declarativa de lo mismo: en vez de posicionar ejes con None y sumarlos con .sum(dim), le das a cada eje una letra y describís qué queda. La regla completa cabe en una frase: todo índice que aparece en las entradas y no aparece en la salida se suma; los índices que aparecen en dos entradas se emparejan (deben tener el mismo tamaño); los que sobreviven a la derecha de -> definen el shape del resultado, en ese orden.

torch.einsum('ik,kj->ij', a, b)       # matmul: k está en las dos entradas y no en la salida → se suma
torch.einsum('ij->ji', m)            # transpuesta: nada se suma, solo cambia el orden
torch.einsum('ij->i', m)             # suma por fila: j desaparece
torch.einsum('ii->i', m)             # diagonal: el índice repetido en una misma entrada la recorre
torch.einsum('ij,ij->', m, m)        # suma de cuadrados: todo se suma, sale un escalar
torch.einsum('bik,bkj->bij', x, y)   # matmul por batch: b se conserva, k se suma
torch.einsum('bhqd,bhkd->bhqk', q, k) # scores de atención: q·k por head, sumando sobre d

La última línea es la razón por la que vale la pena aprenderlo: q @ k.transpose(-1,-2) dice lo mismo, pero hay que pensar qué ejes transponer; 'bhqd,bhkd->bhqk' se lee directamente como "para cada batch y head, un score por par (query, key), contraído sobre la dimensión del embedding". Para lo que no es un matmul limpio (tres tensores, ejes en orden raro), einsum es casi siempre más legible que la cadena de permute/reshape equivalente, y PyTorch lo despacha a las mismas rutinas de BLAS cuando puede.

3. Backprop a mano, capa por capa

Con matmul resuelto, el siguiente paso es implementar el forward y backward pass de una red de dos capas completamente a mano — sin autograd todavía. Las tres piezas del forward son mínimas:

def lin(x, w, b): return x @ w + b        # (n, n_in) @ (n_in, n_out) + (n_out,)  → broadcasting del bias
def relu(x): return x.clamp_min(0.)
def mse(out, targ): return (out[:,0]-targ).pow(2).mean()   # el notebook usa regresión sobre el dígito, adrede

El backward guarda el gradiente de cada tensor en un atributo .g que se le cuelga al propio tensor (un truco del notebook: es lo que .grad hace en PyTorch). Toda la regla de la cadena de una capa lineal está en lin_grad:

def lin_grad(inp, out, w, b):
    # out = inp @ w + b ; out.g ya contiene dL/d(out), con shape (n, n_out)
    inp.g = out.g @ w.t()        # (n, n_out) @ (n_out, n_in) → (n, n_in): el gradiente sigue hacia atrás
    w.g   = inp.t() @ out.g      # (n_in, n) @ (n, n_out)   → (n_in, n_out): mismo shape que w, suma sobre el batch
    b.g   = out.g.sum(0)          # (n_out,): el bias se sumó a cada fila, así que su gradiente es la suma de filas

def forward_and_backward(inp, targ):
    l1 = lin(inp, w1, b1); l2 = relu(l1); out = lin(l2, w2, b2)
    diff = out[:,0]-targ
    loss = diff.pow(2).mean()
    # backward: la regla de la cadena, escrita explícitamente
    out.g = 2.*diff[:,None] / inp.shape[0]   # d/d(out) de mean(diff²): 2·diff/n, con el eje que out tiene
    lin_grad(l2, out, w2, b2)
    l1.g = (l1>0).float() * l2.g   # derivada del ReLU: 1 donde pasó, 0 donde cortó
    lin_grad(inp, l1, w1, b1)

Las tres fórmulas salen de una sola observación: out = inp @ w + b es lineal en cada uno de sus tres argumentos, y el gradiente respecto de un factor de un producto de matrices es el gradiente de la salida multiplicado por el otro factor transpuesto, puesto del lado que hace cerrar los shapes. Por eso no hace falta memorizarlas: si inp.g tiene que ser (n, n_in) y tenés out.g de (n, n_out) y w de (n_in, n_out), la única multiplicación que cierra es out.g @ w.t(). Lo mismo para w.g. Y como el bias se sumó a las n filas por broadcasting, su gradiente acumula las n contribuciones — de ahí el sum(0). En general, cada broadcasting del forward se convierte en una suma en el backward.

TensorShapeGradienteShape del gradiente
inp(n, n_in)out.g @ w.t()(n, n_in)
w(n_in, n_out)inp.t() @ out.g(n_in, n_out)
b(n_out,)out.g.sum(0)(n_out,)
l1 (antes del ReLU)(n, nh)(l1>0).float() * l2.g(n, nh)

Regla de sanidad que vale para siempre: el gradiente de un tensor tiene exactamente el shape del tensor. Si w.g te queda (n_out, n_in), transpusiste el producto.

📐 Derivación rigurosa del gradiente matricial y Vector-Jacobian Products (VJP)

En la capa lineal hacia adelante Y=XW+b, donde la entrada es X∈RN×Din, los pesos W∈RDin×Dout y el bias b∈R1×Dout, la salida es Y∈RN×Dout.

Sea L∈R la función de costo escalar final. Asumimos conocido el gradiente de la loss respecto a la salida de la capa: G=∂L∂Y∈RN×Dout (representado como out.g).

Escribiendo la componente escalar Yi,j=∑m=1DinXi,mWm,j+bj, derivamos componente a componente mediante la regla de la cadena multivariable:

  1. Gradiente respecto a la entrada X: ∂L∂Xi,k=∑j=1Dout∂L∂Yi,j∂Yi,j∂Xi,k=∑j=1DoutGi,jWk,j=∑j=1DoutGi,j(W⊤)j,k⟹∂L∂X=GW⊤ Dimensionalmente: (N×Dout)×(Dout×Din)=(N×Din), idéntico al shape de X.
  2. Gradiente respecto a los pesos W: ∂L∂Wk,j=∑i=1N∂L∂Yi,j∂Yi,j∂Wk,j=∑i=1NXi,kGi,j=∑i=1N(X⊤)k,iGi,j⟹∂L∂W=X⊤G Dimensionalmente: (Din×N)×(N×Dout)=(Din×Dout), idéntico al shape de W. La suma sobre el batch queda absorbida en el producto escalar interno de las columnas de X y las columnas de G.
  3. Gradiente respecto al sesgo b: ∂L∂bj=∑i=1N∂L∂Yi,j∂Yi,j∂bj=∑i=1NGi,j⋅1⟹∂L∂b=∑i=1NGi,: En código PyTorch: out.g.sum(0).

Por qué autograd nunca materializa la matriz Jacobiana completa: La matriz jacobiana J=∂Y∂X tendría tamaño (N⋅Dout)×(N⋅Din). Para un batch moderado de N=256 y D=1024, esa matriz tendría 6.87×1010 elementos (¡274 Gigabytes de VRAM para una sola capa!). En reversa, autograd evalúa directamente el Vector-Jacobian Product (VJP): proyecta el vector cotangente G a través de W⊤ usando un único GEMM sin instanciar jamás la hiper-matriz jacobiana.

Los tres gradientes de una capa lineal y sus shapes Forward: inp (n, n_in) por w (n_in, n_out) más b (n_out,) da out (n, n_out). Backward: inp.g = out.g por w transpuesta, w.g = inp transpuesta por out.g, b.g = suma de out.g por columnas. Cada gradiente tiene el shape de su tensor; el factor transpuesto se pone del lado que cierra los shapes. FORWARD · out = inp @ w + b inp (n, n_in) @ w (n_in, n_out) + b (n_out,) = out (n, n_out) el bias se suma a cada fila por broadcasting n filas del batch · n_in entradas · n_out salidas out.g llega de la loss con el shape de out y de ahí salen los tres gradientes de abajo BACKWARD · tres productos, cada uno con el shape de su tensor inp.g (n, n_in) = out.g (n, n_out) @ w.t() (n_out, n_in) w.g (n_in, n_out) = inp.t() (n_in, n) @ out.g (n, n_out) b.g (n_out,) = out.g (n, n_out) .sum(0) regla que vale para siempre: el gradiente de un tensor tiene exactamente el shape del tensor · el factor transpuesto va del lado que hace cerrar los shapes si w.g te queda (n_out, n_in), transpusiste el producto · el ReLU entre medio solo multiplica por (l1 > 0): 1 donde pasó, 0 donde cortó
lin_grad dibujado con los rectángulos a escala de sus shapes. El gradiente que llega (out.g, ámbar) se multiplica por el otro factor del producto, transpuesto y puesto del lado que cierra las dimensiones. Con eso no hay nada que memorizar: el shape del resultado te dice cómo escribir la fórmula.

Este código se verifica comparándolo contra loss.backward() de PyTorch sobre los mismos tensores — y coincide, porque autograd no hace nada distinto: aplica la misma regla de la cadena, automáticamente. El notebook refactoriza esto en tres pasos progresivos: primero clases con backward() propio por capa (Relu, Lin, Mse), después una clase base Module que factoriza el patrón común, y finalmente la versión con nn.Module real de PyTorch — mostrando en código que la "magia" de .backward() es, literalmente, la misma clase Module escrita a mano un párrafo antes:

class Module():
    def __call__(self, *args):
        self.args = args
        self.out = self.forward(*args)     # guarda entradas y salida: el backward las necesita
        return self.out
    def forward(self): raise Exception('not implemented')
    def backward(self): self.bwd(self.out, *self.args)

class Relu(Module):
    def forward(self, inp): return inp.clamp_min(0.)
    def bwd(self, out, inp): inp.g = (inp>0).float() * out.g

class Lin(Module):
    def __init__(self, w, b): self.w, self.b = w, b
    def forward(self, inp): return inp @ self.w + self.b
    def bwd(self, out, inp):
        inp.g = out.g @ self.w.t()
        self.w.g = inp.t() @ out.g
        self.b.g = out.g.sum(0)

class Model():
    def __init__(self, w1, b1, w2, b2): self.layers = [Lin(w1,b1), Relu(), Lin(w2,b2)]; self.loss = Mse()
    def __call__(self, x, targ):
        for l in self.layers: x = l(x)
        return self.loss(x, targ)
    def backward(self):
        self.loss.backward()
        for l in reversed(self.layers): l.backward()   # el backward recorre las capas al revés

El __call__ que guarda self.args y self.out es el corazón del asunto: el backward de cada capa necesita los valores que vio en el forward (el ReLU necesita saber dónde cortó; la capa lineal necesita inp para calcular w.g). Autograd hace lo mismo guardando esos tensores en el grafo — y por eso el forward de un modelo grande consume memoria en proporción a la profundidad, y por eso torch.no_grad() la libera: sin backward no hay nada que recordar.

4. El framework Learner: callbacks en vez de if/else

La pieza más influyente de esta unidad es el diseño del Learner de miniai — básicamente una versión simplificada de por qué Learner/Trainer de fastai y HF están diseñados como están. El problema que resuelve: un training loop necesita soportar mixed precision, logging, early stopping, schedulers de LR, multi-GPU, etc. — si cada feature se agrega como un if dentro del loop, el código se vuelve inmanejable rápido.

La solución es un sistema de callbacks con puntos de enganche nombrados (before_fit, before_batch, after_batch, after_epoch...) más tres excepciones especiales para control de flujo:

class CancelFitException(Exception): pass
class CancelBatchException(Exception): pass
class CancelEpochException(Exception): pass

class Callback(): order = 0

def run_cbs(cbs, method_nm, learn=None):
    for cb in sorted(cbs, key=attrgetter('order')):     # el order decide quién corre primero
        method = getattr(cb, method_nm, None)
        if method is not None: method(learn)               # si el callback no define ese punto, se lo saltea

class SingleBatchCB(Callback):
    order = 1
    def after_batch(self, learn): raise CancelFitException()

El truco que hace que el Learner quede corto es un decorador, with_cbs: envuelve cada fase (batch, epoch, fit) llamando a before_<fase> antes, after_<fase> después, y atrapando la excepción Cancel<Fase>Exception correspondiente — de modo que cancelar un batch salta al siguiente batch, cancelar una epoch a la siguiente epoch, y cancelar el fit termina todo, sin que ningún if lo diga:

Las tres fases anidadas del Learner y qué atrapa cada excepción Cancel fit contiene epoch contiene batch; cada fase corre before_, el cuerpo, after_ y siempre cleanup_. Dentro del batch: predict, get_loss, backward, step, zero_grad con un hook entre cada uno. CancelBatchException la atrapa with_cbs('batch') y sigue con el batch siguiente; CancelEpochException con la epoch siguiente; CancelFitException termina el fit. fit before_fit after_fit cleanup_fit: corre siempre (finally) epoch before_epoch after_epoch cleanup_epoch batch before_batch after_batch cleanup_batch predict get_loss backward step zero_grad UNA EXCEPCIÓN POR FASE CancelFitException termina el fit (SingleBatchCB, LRFinderCB cuando la loss explota) CancelEpochException salta al after de esa epoch y sigue con la epoch siguiente CancelBatchException with_cbs('batch') la atrapa: se saltan backward y step; sigue el batch siguiente after_predict · after_loss · after_backward · after_step: un punto de enganche entre cada paso del batch run_cbs: corre el método en cada callback que lo define, ordenados por .order; el que no lo define se saltea cancelar un batch salta al batch siguiente, cancelar una epoch a la epoch siguiente, cancelar el fit termina todo: la excepción sube hasta el with_cbs de su nombre y ahí se apaga; el cleanup_ de esa fase corre igual
El Learner de miniai como tres cajas anidadas. Un callback se engancha en cualquiera de los puntos nombrados sin tocar el loop, y para cortar lanza la excepción de la fase que quiere abandonar: with_cbs la atrapa exactamente en ese nivel y el resto sigue como si nada.
class with_cbs:
    def __init__(self, nm): self.nm = nm
    def __call__(self, f):
        def _f(o, *args, **kwargs):
            try:
                o.callback(f'before_{self.nm}')
                f(o, *args, **kwargs)
                o.callback(f'after_{self.nm}')
            except globals()[f'Cancel{self.nm.title()}Exception']: pass   # 'batch' → CancelBatchException
            finally: o.callback(f'cleanup_{self.nm}')                    # corre siempre, cancelado o no
        return _f

class Learner():
    def __init__(self, model, dls, loss_func=F.mse_loss, lr=0.1, cbs=None, opt_func=optim.SGD):
        self.model, self.dls, self.loss_func, self.lr, self.opt_func = model, dls, loss_func, lr, opt_func
        self.cbs = list(cbs or [])

    @with_cbs('batch')
    def _one_batch(self):
        self.predict();  self.callback('after_predict')
        self.get_loss(); self.callback('after_loss')
        if self.training:
            self.backward(); self.callback('after_backward')
            self.step();     self.callback('after_step')
            self.zero_grad()

    @with_cbs('epoch')
    def _one_epoch(self):
        for self.iter, self.batch in enumerate(self.dl): self._one_batch()

    def one_epoch(self, training):
        self.model.train(training)
        self.dl = self.dls.train if training else self.dls.valid
        self._one_epoch()

    @with_cbs('fit')
    def _fit(self, train, valid):
        for self.epoch in self.epochs:
            if train: self.one_epoch(True)
            if valid: torch.no_grad()(self.one_epoch)(False)

    def fit(self, n_epochs=1, train=True, valid=True, cbs=None, lr=None):
        cbs = list(cbs or [])
        for cb in cbs: self.cbs.append(cb)          # callbacks solo para este fit (ej. el LR finder)
        try:
            self.n_epochs, self.epochs = n_epochs, range(n_epochs)
            if lr is None: lr = self.lr
            self.opt = self.opt_func(self.model.parameters(), lr)
            self._fit(train, valid)
        finally:
            for cb in cbs: self.cbs.remove(cb)

    def __getattr__(self, name):
        # predict / get_loss / backward / step / zero_grad no están en el Learner: los provee un callback
        if name in ('predict','get_loss','backward','step','zero_grad'): return partial(self.callback, name)
        raise AttributeError(name)

    def callback(self, method_nm): run_cbs(self.cbs, method_nm, self)

    @property
    def training(self): return self.model.training

class TrainCB(Callback):
    def predict(self, learn):   learn.preds = learn.model(learn.batch[0])
    def get_loss(self, learn):  learn.loss = learn.loss_func(learn.preds, learn.batch[1])
    def backward(self, learn):  learn.loss.backward()
    def step(self, learn):      learn.opt.step()
    def zero_grad(self, learn): learn.opt.zero_grad()

Fijate en __getattr__: hasta los cinco pasos del training loop (predict, get_loss, backward, step, zero_grad) son callbacks. TrainCB es el que los implementa en su versión más simple, y reemplazarlo es cómo miniai agrega mixed precision o entrenamiento distribuido con accelerate — un callback nuevo que define esos cinco métodos, cero cambios en el Learner. (La versión del repo usa fastcore para abreviar — fc.store_attr(), fc.L — y agrega n_inp para modelos con varias entradas; el esqueleto de arriba es el mismo sin esa dependencia, y corre tal cual.)

Usar excepciones para control de flujo (no para errores) es la parte más elegante del diseño: un callback de "encontrar el learning rate óptimo" puede simplemente raise CancelFitException() cuando el loss explota, y el Learner no necesita saber nada de esa lógica — solo atrapa la excepción en el nivel correspondiente. Ese callback, en su versión mínima, es esto:

class LRFinderCB(Callback):
    def __init__(self, lr_mult=1.3): self.lr_mult = lr_mult

    def before_fit(self, learn):
        self.lrs, self.losses = [], []
        self.min = math.inf

    def after_batch(self, learn):
        if not learn.training: raise CancelEpochException()   # la validación no interesa acá: saltearla entera
        self.lrs.append(learn.opt.param_groups[0]['lr'])
        loss = learn.loss.item()
        self.losses.append(loss)
        if loss < self.min: self.min = loss
        if math.isnan(loss) or loss > self.min*3: raise CancelFitException()   # explotó: cortar todo
        for g in learn.opt.param_groups: g['lr'] *= self.lr_mult     # si no, subir el lr y seguir

learn.fit(40, cbs=[lrf := LRFinderCB()])   # 40 epochs "de máximo": el callback corta mucho antes
plt.plot(lrf.lrs, lrf.losses); plt.xscale('log')

Veinte líneas, y no hay una sola referencia a él dentro del loop: multiplica el lr después de cada batch, registra la loss, y cuando la loss supera tres veces el mínimo visto (o se vuelve NaN) lanza CancelFitException, que with_cbs('fit') atrapa. El finally de fit saca el callback de la lista, así que el Learner queda limpio para el entrenamiento real. La versión del repo hace lo mismo con un ExponentialLR de PyTorch en vez del *= a mano, y dibuja el gráfico en cleanup_fit — misma lógica. Con este esqueleto, agregar un progress bar con gráfico en vivo (ProgressCB), tracking de métricas (MetricsCB) o mover tensores a GPU (DeviceCB) es escribir una clase nueva de 5-10 líneas, sin tocar el Learner en absoluto.

El order importa cuando dos callbacks se enganchan al mismo punto y uno depende del otro: ProgressCB declara order = MetricsCB.order+1 porque tiene que imprimir las métricas después de que MetricsCB las calculó en su after_epoch; y DeviceCB tiene que mover el batch a GPU en before_batch antes de que cualquier otro callback lo toque. Cuando un callback ve datos "viejos" o en el dispositivo equivocado, lo primero que hay que mirar es el order.

🔧 El mismo patrón que ya usás en producción

CancelFitException/CancelBatchException resuelven exactamente el mismo problema que el numeric_gate.py de tu gate anti-alucinación: interrumpir un flujo en curso desde un punto de control externo, sin que la lógica principal necesite saber por qué se interrumpió. Y el diseño de callbacks con order explícito (para controlar en qué secuencia se ejecutan cuando hay varios enganchados al mismo punto) es la misma idea que HuggingFace Trainer usa con sus TrainerCallback — no es casualidad, es el mismo problema de ingeniería resuelto por gente que se leyó el mismo código.

5. Errores comunes

Casi todos se diagnostican igual: imprimí shapes antes y después de la operación sospechosa, y en el framework, imprimí qué callbacks corren y en qué orden.

SíntomaCausa probableQué hacer
La loss "entrena" pero no baja, o el resultado tiene un shape enorme ((n,n) donde esperabas (n,)).Broadcasting silencioso: (n,1) contra (n,) se estira a (n,n) por las reglas 1 y 3 de §2, y el .mean() promedia todo contra todo.Alinear shapes a mano antes de operar (targ[:,None] o out[:,0]) y asertar el shape del resultado: assert diff.shape == targ.shape.
RuntimeError: shape mismatch al hacer w -= lr * w.g, o gradientes que "funcionan" pero el entrenamiento diverge.Gradiente transpuesto: escribiste out.g.t() @ inp o w @ out.g en vez de inp.t() @ out.g.Regla de §3: el gradiente tiene el shape del tensor. Chequeá w.g.shape == w.shape para cada parámetro y compará contra autograd con torch.allclose.
El backward a mano coincide con autograd en el primer paso y diverge a partir del segundo.Los .g no se limpian entre pasos — o, en la versión con nn.Module, falta opt.zero_grad() y .grad se acumula.Reasignar (w.g = …, no w.g += …) o poner los gradientes en cero al empezar cada paso; en PyTorch, zero_grad() antes de backward(), que es exactamente lo que hace TrainCB en §4.
Un callback ve la métrica de la epoch anterior, o el batch todavía en CPU.order mal puesto: corre antes que el callback del que depende (MetricsCB, DeviceCB).Declarar order = OtroCB.order + 1 y verificar con un print en run_cbs la secuencia real.
El entrenamiento "se cae" con un traceback de CancelFitException, o al revés: un except Exception se traga la cancelación y el loop sigue.Excepciones de control confundidas con errores: lanzaste Cancel* fuera del contexto que la atrapa (por ejemplo dentro de fit pero fuera de _fit), o envolviste el loop con un except demasiado ancho.Las Cancel* solo se lanzan desde callbacks llamados dentro de la fase que las atrapa; nunca las capturés con except Exception. Si tenés que atrapar errores reales, hacelo después del except CancelFitException.
La versión con Numba/@njit es más lenta que Python puro.Mediste la primera llamada, que incluye la compilación.Llamar una vez para compilar y cronometrar la segunda; en GPU, sincronizar (torch.cuda.synchronize()) antes de leer el reloj.

6. Resumen

  1. Cada capa de abstracción entre Python y el hardware (interpretado → vectorizado → CUDA) tiene un costo medible — la escalera matmul lo cuantifica en millones de veces, y el tiempo del loop en Python está en el intérprete, no en la aritmética.
  2. Broadcasting son tres reglas: alinear desde la derecha, dimensiones iguales o 1, la de tamaño 1 se estira sin copiar. Es lo que permite borrar loops — y lo que produce bugs silenciosos cuando (n,) se encuentra con (n,1).
  3. En einsum, todo índice que aparece en las entradas y no en la salida se suma; los demás definen el shape del resultado. Sirve para leer atención ('bhqd,bhkd->bhqk') sin transponer a mano.
  4. Backprop no es magia: es la regla de la cadena aplicada capa por capa. Para una capa lineal, inp.g = out.g @ w.t(), w.g = inp.t() @ out.g, b.g = out.g.sum(0) — y el gradiente siempre tiene el shape de su tensor. Es verificable contra autograd porque hace exactamente lo mismo.
  5. Un framework de entrenamiento robusto se construye con callbacks + excepciones de control de flujo, no con ifs acumulados dentro del loop principal. with_cbs envuelve cada fase con before_*/after_* y atrapa la Cancel* correspondiente; hasta predict y backward son callbacks.
  6. Un LR finder es un callback de veinte líneas que multiplica el lr en cada batch y lanza CancelFitException cuando la loss explota — el Learner no sabe que existe.
  7. Este patrón de diseño no es específico de miniai — es el mismo que usan fastai, HF Trainer, y sistemas de control externos como un gate de seguridad.

7. Ponete a prueba

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

1. ¿Qué shape tiene torch.randn(4, 1, 6) * torch.randn(5, 1)? ¿Y torch.randn(4, 6) * torch.randn(5, 1)?

La primera da (4, 5, 6): (5,1) se alinea a la derecha como (1,5,1), y cada par es igual o vale 1. La segunda falla: alineando (4,6) con (5,1), la última dimensión cierra (6 contra 1) pero la anterior es 4 contra 5 — ninguna es 1, así que no hay broadcasting posible. (§2)

2. En torch.einsum('bhqd,bhkd->bhqk', q, k), ¿qué índice se suma y qué tamaño tiene que coincidir entre los dos tensores?

Se suma d, el único que aparece en las entradas y no en la salida — es la dimensión del embedding por head, y tiene que tener el mismo tamaño en q y en k. b y h también tienen que coincidir (aparecen en ambas entradas), pero se conservan; q y k pueden ser distintos porque cada uno vive en un solo tensor. (§2)

3. inp es (64, 784), w es (784, 50) y out.g es (64, 50). Escribí w.g y justificalo solo con shapes.

w.g = inp.t() @ out.g: (784, 64) @ (64, 50) → (784, 50), el shape de w. Es la única forma de combinar los dos tensores disponibles que produce ese shape. La suma sobre el batch (los 64 ejemplos) queda implícita en la contracción sobre la dimensión de tamaño 64. (§3)

4. ¿Por qué b.g es out.g.sum(0) y no out.g?

Porque en el forward el bias (n_out,) se sumó a cada una de las n filas por broadcasting: un mismo parámetro contribuyó a n salidas, así que su gradiente acumula las n contribuciones. Regla general: cada broadcasting del forward es una suma en el backward, y el gradiente tiene que tener el shape del parámetro, (n_out,). (§3)

5. Un callback lanza CancelBatchException en after_loss. ¿Qué se saltea y qué sigue corriendo?

Se saltean backward, step, zero_grad y el after_batch de ese batch — with_cbs('batch') atrapa la excepción alrededor de _one_batch. Sí corre cleanup_batch (está en el finally), y el loop de _one_epoch sigue con el batch siguiente como si nada. El fit y la epoch no se enteran. (§4)

6. ¿Qué pasa si escribís ProgressCB con order = 0 en vez de MetricsCB.order + 1?

Con el mismo order, sorted conserva el orden en que se agregaron a la lista; si ProgressCB quedó antes, su after_epoch corre antes que el de MetricsCB y muestra las métricas de la epoch anterior (o vacías en la primera). No es un error visible: es un log desfasado en una epoch, que es el peor tipo de bug. (§4, §5)

8. Para practicar

✏️ Ejercicio 1 — La escalera de matmul en tu CPU

Implementá matmul de tres formas (Python puro, broadcasting, einsum), verificá que dan lo mismo y medí cada una. Tendrías que ver tres órdenes de magnitud entre la primera y la segunda, y que @ gana a todas.

import torch, time
torch.manual_seed(0)
a = torch.randn(5, 784); b = torch.randn(784, 10)

def matmul_py(a, b):
    (ar,ac),(br,bc) = a.shape, b.shape
    c = torch.zeros(ar, bc)
    for i in range(ar):
        for j in range(bc):
            for k in range(ac): c[i,j] += a[i,k] * b[k,j]
    return c

def matmul_bc(a, b):
    (ar,ac),(br,bc) = a.shape, b.shape
    c = torch.zeros(ar, bc)
    for i in range(ar): c[i] = (a[i,:,None] * b).sum(dim=0)   # (784,1) × (784,10) → suma sobre k
    return c

def matmul_es(a, b): return torch.einsum('ik,kj->ij', a, b)

def cronometrar(f, reps=3):
    f(a, b)                                     # calentar
    t0 = time.perf_counter()
    for _ in range(reps): r = f(a, b)
    return r, (time.perf_counter() - t0) / reps * 1e3

ref = a @ b
for nombre, f in [("python", matmul_py), ("broadcast", matmul_bc), ("einsum", matmul_es), ("@", torch.matmul)]:
    r, ms = cronometrar(f)
    print(f"{nombre:>10}: {ms:8.3f} ms   ok={torch.allclose(r, ref, atol=1e-4)}")
Pista / qué esperar

Python puro: cientos de milisegundos. Broadcasting: décimas de milisegundo (los 5 loops externos que quedan son baratos). einsum y @: centésimas o menos. Todos con ok=True. Después probá a = torch.randn(500, 784): la versión de Python se vuelve inusable y la diferencia entre broadcasting (que todavía tiene un loop en Python) y @ se abre mucho más.

✏️ Ejercicio 2 — Backprop a mano contra autograd

Corré el forward/backward de §3 sobre datos aleatorios y compará cada .g con el .grad que calcula PyTorch. Si un allclose da False, mirá primero el shape.

import torch
torch.manual_seed(0)
n, n_in, nh = 64, 784, 50
inp  = torch.randn(n, n_in)
targ = torch.randn(n)
w1 = torch.randn(n_in, nh) / n_in**0.5; b1 = torch.zeros(nh)
w2 = torch.randn(nh, 1) / nh**0.5;      b2 = torch.zeros(1)

def lin(x, w, b): return x @ w + b
def relu(x): return x.clamp_min(0.)
def lin_grad(inp, out, w, b):
    inp.g = out.g @ w.t()
    w.g   = inp.t() @ out.g
    b.g   = out.g.sum(0)

def forward_and_backward(inp, targ):
    l1 = lin(inp, w1, b1); l2 = relu(l1); out = lin(l2, w2, b2)
    diff = out[:,0] - targ
    loss = diff.pow(2).mean()
    out.g = 2.*diff[:,None] / inp.shape[0]
    lin_grad(l2, out, w2, b2)
    l1.g = (l1>0).float() * l2.g
    lin_grad(inp, l1, w1, b1)
    return loss

loss = forward_and_backward(inp, targ)

# lo mismo con autograd, sobre copias de los mismos tensores
def mkgrad(x): return x.clone().requires_grad_(True)
pw1, pb1, pw2, pb2, pinp = map(mkgrad, (w1, b1, w2, b2, inp))
out = lin(relu(lin(pinp, pw1, pb1)), pw2, pb2)
loss2 = (out[:,0] - targ).pow(2).mean()
loss2.backward()

for nombre, mano, auto in [("w1", w1.g, pw1.grad), ("b1", b1.g, pb1.grad), ("w2", w2.g, pw2.grad), ("b2", b2.g, pb2.grad), ("inp", inp.g, pinp.grad)]:
    print(f"{nombre:>4}: shape {tuple(mano.shape)}  coincide={torch.allclose(mano, auto, atol=1e-5)}")
Pista / qué esperar

Cinco líneas con coincide=True y shapes (784, 50), (50,), (50, 1), (1,), (64, 784) — cada gradiente con el shape de su tensor. Después rompelo a propósito: cambiá w.g = inp.t() @ out.g por w.g = (out.g.t() @ inp).t() (sigue siendo correcto — ¿por qué?) y por out.g.t() @ inp (shape transpuesto: el allclose ni siquiera puede comparar). Y sacá el / inp.shape[0] de out.g: los gradientes quedan multiplicados por 64, porque la loss era una media y no una suma.

✏️ Ejercicio 3 — Un LR finder sobre el Learner de §4

Pegá tal cual el bloque de with_cbs + Learner + TrainCB y el LRFinderCB de §4 (con los imports de abajo), y corré esto debajo. Es una regresión lineal de juguete: la gracia es ver al callback cortar el entrenamiento solo, sin que el Learner sepa nada.

import math, torch, torch.nn as nn, torch.nn.functional as F
from functools import partial
from operator import attrgetter
from torch import optim
# … acá van las clases de §4: Cancel*, Callback, run_cbs, with_cbs, Learner, TrainCB, LRFinderCB …

torch.manual_seed(0)
x = torch.rand(1024, 1) * 2 - 1
y = 3*x - 1 + 0.1*torch.randn(1024, 1)      # la "verdad": pendiente 3, ordenada -1

class DLS: pass
def batches(x, y, bs=128): return [(x[i:i+bs], y[i:i+bs]) for i in range(0, len(x), bs)]
dls = DLS(); dls.train = batches(x[:768], y[:768]); dls.valid = batches(x[768:], y[768:])

# 1) buscar el lr: arrancar muy bajo y dejar que el callback lo suba hasta que explote
learn = Learner(nn.Linear(1, 1), dls, cbs=[TrainCB()], lr=1e-4)
lrf = LRFinderCB()
learn.fit(40, cbs=[lrf])
for lr, l in zip(lrf.lrs, lrf.losses): print(f"lr={lr:.2e}  loss={l:.3f}")
print("batches vistos:", len(lrf.lrs), "— callbacks que quedaron:", [type(c).__name__ for c in learn.cbs])

# 2) entrenar de verdad con un lr por debajo del punto donde explotó
learn = Learner(nn.Linear(1, 1), dls, cbs=[TrainCB()], lr=0.5)
learn.fit(5)
print(learn.model.weight.item(), learn.model.bias.item(), learn.loss.item())
Pista / qué esperar

Unas 40 líneas de lr / loss: la loss se queda en ~4-5 mientras el lr es minúsculo, baja con fuerza entre 1e-1 y 1 (llega a ~0.02), y en cuanto el lr pasa de ~2 vuelve a subir y el callback corta — pediste 40 epochs y corrió menos de una. La lista de callbacks al final es solo ['TrainCB']: el finally de fit sacó al finder. El entrenamiento con lr=0.5 termina con pendiente ≈ 3.0, ordenada ≈ −1.0 y loss ≈ 0.01 (el ruido que le pusiste a los datos). Para seguir: escribí un PrintCB con after_epoch que imprima learn.epoch y learn.loss, y probá qué pasa si lanzás CancelEpochException desde before_epoch cuando learn.epoch == 2.

9. Fuentes y para seguir

  • Practical Deep Learning for Coders — Part 2 — las lecciones en video; la 11 es la escalera de matmul, la 13-14 el backprop a mano, y la 15-16 construyen el Learner con callbacks en vivo.
  • fastai/course22p2 — los notebooks de esta unidad: 01_matmul, 03_backprop y 09_learner; el último tiene el Learner completo del que §4 es el esqueleto.
  • Broadcasting semantics — PyTorch docs — las tres reglas de §2 en su forma oficial, con la lista de casos que fallan; una página, léela entera.
  • torch.einsum — PyTorch docs — la referencia de la notación, con la regla de "índices que no aparecen en la salida se suman" y una docena de ejemplos cortos.
  • Broadcasting — NumPy user guide — el origen de las reglas, con las figuras que muestran cómo se estira cada eje; útil si las de PyTorch te quedan demasiado secas.
  • fastai — Callback core — la versión de producción del mismo diseño de §4 (order, Cancel*Exception, puntos de enganche), para ver cómo escala el patrón.