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
einsumsabiendo 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 contraautograd. - Reconstruir el esqueleto del
Learnerde miniai (with_cbs,run_cbs, excepcionesCancel*) 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ón | Qué reemplaza | Cómo se escribe | Quién hace el trabajo |
|---|---|---|---|
| Python puro | — | tres for | el intérprete, elemento por elemento |
| Numba | el loop más interno (sobre k) | @njit def dot(a,b) sobre arrays de NumPy | código máquina compilado en la primera llamada |
| Broadcasting | los dos loops internos | c[i] = (a[i,:,None] * b).sum(dim=0) | el kernel en C de PyTorch, una fila por vez |
einsum / @ | los tres loops | torch.einsum('ik,kj->ij', a, b) o a @ b | BLAS: una rutina de matmul optimizada por décadas |
| CUDA | los tres loops, en paralelo | kernel 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.
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:
-
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 es (en float32). En el bucle más interno sobre , la lectura de 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 , fijado e incrementando , la dirección salta de fila en fila: . ¡Cada paso salta bytes! Si , 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. - Intensidad Aritmética y el Modelo Roofline: Multiplicar dos matrices de dimensiones y requiere exactamente operaciones de punto flotante (un producto y una suma por celda). La intensidad aritmética es el cociente entre operaciones y bytes transferidos: 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 . Con , el procesador está el 99.9% del tiempo ocioso esperando datos de memoria (memory-bandwidth bound).
- 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 y 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 exactamente veces y cada valor de exactamente veces: Con bloques estándar de , la intensidad sube a ; y al llevarla a registros de los Tensor Cores mediante micro-tiles de , 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:
- 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). - 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. - 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.
| Tensor | Shape | Gradiente | Shape 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 , donde la entrada es , los pesos y el bias , la salida es .
Sea la función de costo escalar final. Asumimos conocido el gradiente de la loss respecto a la salida de la capa: (representado como out.g).
Escribiendo la componente escalar , derivamos componente a componente mediante la regla de la cadena multivariable:
- Gradiente respecto a la entrada : Dimensionalmente: , idéntico al shape de .
- Gradiente respecto a los pesos : Dimensionalmente: , idéntico al shape de . La suma sobre el batch queda absorbida en el producto escalar interno de las columnas de y las columnas de .
-
Gradiente respecto al sesgo :
En código PyTorch:
out.g.sum(0).
Por qué autograd nunca materializa la matriz Jacobiana completa: La matriz jacobiana tendría tamaño . Para un batch moderado de y , esa matriz tendría elementos (¡274 Gigabytes de VRAM para una sola capa!). En reversa, autograd evalúa directamente el Vector-Jacobian Product (VJP): proyecta el vector cotangente a través de usando un único GEMM sin instanciar jamás la hiper-matriz jacobiana.
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:
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íntoma | Causa probable | Qué 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
- 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.
- 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). - 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. - 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 contraautogradporque hace exactamente lo mismo. - Un framework de entrenamiento robusto se construye con callbacks + excepciones de control de flujo, no con
ifs acumulados dentro del loop principal.with_cbsenvuelve cada fase conbefore_*/after_*y atrapa laCancel*correspondiente; hastapredictybackwardson callbacks. - Un LR finder es un callback de veinte líneas que multiplica el lr en cada batch y lanza
CancelFitExceptioncuando la loss explota — elLearnerno sabe que existe. - 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
Learnercon callbacks en vivo. - fastai/course22p2 — los notebooks de esta unidad:
01_matmul,03_backpropy09_learner; el último tiene elLearnercompleto 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.