Skip to content

English · Español

04 — PyTorch como sustrato (primer encuentro)

🇪🇸 Después de 23 fases en NumPy puro, esta página introduce PyTorch. No como caja mágica sino como plomería: un Tensor es un Storage + metadatos de vista, una nn.Module es un registro de parámetros, una pasada hacia adelante es una secuencia de dispatch a kernels. Fase 25 desarma el sistema; Fase 24 lo presenta como herramienta para portar el grammar MiniGPT.

Esta es la primera página de teoría de PyTorch del currículo. El framework se introduce tarde y delgado — para la Fase 24 Borja ya sabe qué hace cada operador (Fases 7–17), cómo corre en GPU (Fase 23) y cómo fluye la memoria a través de él (Fase 22). PyTorch se presenta entonces como la capa de routing que envuelve el sustrato.

El objetivo de esta página es mantener a PyTorch en su sitio: es una herramienta, no una cosmovisión. La Fase 25 abre el dispatcher y los internals de autograd.

El modelo mental de dos líneas

nn.Module  =  parameters registry + forward-method-as-graph
Tensor     =  Storage (bytes) + view metadata (shape, stride, dtype, device, requires_grad)

Todo lo demás en PyTorch decora estas dos ideas. Si te quedas con este resumen de dos líneas en mente, el resto es obvio.

Tensor como Storage + vista

Un tensor PyTorch no es un bloque contiguo de memoria por sí mismo. Es una vista sobre un Storage (que sí es un bloque contiguo de memoria). Múltiples tensores pueden compartir un Storage — así es como funcionan view, transpose, narrow, permute sin copiar.

x = torch.arange(12).reshape(3, 4)
y = x.t()             # transposed view
x.storage() is y.storage()    # True — they share the underlying bytes

Los metadatos de la vista (stride, offset) describen cómo recorrer el storage. Una vista transpuesta tiene strides intercambiados; los bytes son idénticos.

Este es exactamente el modelo ndarray de NumPy. La contribución de PyTorch es añadir device, requires_grad y unas pocas decoraciones — no un modelo de memoria distinto.

nn.Module como registro de parámetros

class TinyMLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(64, 256)
        self.fc2 = nn.Linear(256, 600)   # 600 = grammar MiniGPT vocab size
    def forward(self, x):
        return self.fc2(torch.relu(self.fc1(x)))

Lo que nn.Module realmente añade sobre una clase Python plana:

  • .parameters() recorre recursivamente atributos que sean Parameter (una subclase de Tensor con requires_grad=True por defecto).
  • .state_dict() devuelve {name: Tensor} para serialización.
  • .to(device) mueve recursivamente los parámetros a un dispositivo.
  • .train() / .eval() cambia una flag que algunas capas (Dropout, BatchNorm) leen.

Eso es. No hay magia oculta; es un contenedor ligeramente inteligente para parámetros.

Qué hace realmente model.cuda()

model = TinyMLP()
model.cuda()

La llamada .cuda():

  1. Recorre todos los parámetros y buffers recursivamente.
  2. Para cada uno, asigna un tensor equivalente en la GPU (llama a cudaMalloc).
  3. Copia los bytes del host vía cudaMemcpy.
  4. Reemplaza el parámetro in-place (el nn.Module ahora contiene un tensor en GPU).

La estructura y el código del modelo no cambian — solo la flag device del tensor y la memoria subyacente. La siguiente llamada a forward() despacha cada op al backend CUDA en lugar del backend CPU. Sin nuevo modelo, sin copia del código.

Ese dispatch es el dispatcher, tema de la Fase 25.

La pasada hacia adelante como secuencia de dispatch

Para un batch a través de TinyMLP:

y = model(x)
# Equivalent to:
y = model.fc2(torch.relu(model.fc1(x)))
# Each call dispatches based on (op_name, device, dtype):
#   fc1.weight @ x + fc1.bias   → cuBLAS GEMM kernel (CUDA, fp32)
#   torch.relu                  → eltwise CUDA kernel
#   fc2.weight @ ... + fc2.bias → cuBLAS GEMM kernel

Cada línea es un lanzamiento de kernel (o, con torch.compile, un kernel fusionado que cubre varios). El dispatcher (Fase 25) es lo que escoge qué kernel basándose en (op, device, dtype, layout).

Para el port del MiniGPT de la Fase 17 en la Fase 24: cada np.matmul se convierte en torch.matmul, cada np.exp / np.sum / np.softmax se convierte en torch.softmax. La estructura del código del modelo no cambia. El dispatcher hace el resto.

En qué difiere PyTorch de NumPy

Tres cosas importan para la Fase 24:

  1. Devices. Los tensores viven en un device: cpu, cuda:0, cuda:1. Las operaciones entre tensores de devices distintos dan error. NumPy no tiene este concepto.
  2. Autograd. Los tensores con requires_grad=True registran operaciones en un grafo backward. Llamar a .backward() recorre el grafo y acumula gradientes en .grad. La Fase 25 lo hace explícito; la Fase 24 mayormente vive en inferencia (.eval() + torch.no_grad()), así que el grafo no acumula.
  3. Reglas de promoción de dtype/device. PyTorch promociona silenciosamente fp16 + fp32 → fp32 (configurable). NumPy tiene reglas similares pero los modos de fallo difieren.

Para el port del lab, Borja corre inferencia, sin gradientes, fp32 en CPU primero (case con NumPy bit-exacta), luego fp32 en CUDA (case con NumPy a ~1e-5 por la no-asociatividad de la aritmética fp), luego fp16 en CUDA (case con NumPy a ~1e-2; los tests usan tolerancia más laxa).

El sanity check de byte-equivalencia

Tras portar MiniGPT a PyTorch (src/minimodel/torch_minigpt.py), la validación es:

np_model = load_phase17_numpy_minigpt()
pt_model = load_torch_minigpt_from_same_weights()
x = np.random.randn(2, 32, 64).astype(np.float32)
y_np = np_model(x)
y_pt = pt_model(torch.tensor(x)).numpy()
assert np.allclose(y_np, y_pt, atol=1e-5, rtol=1e-5)

Si esto pasa en fp32 en CPU, el port es fiel. Si falla, el port tiene un bug de orden de capa o de mapeo de pesos. Depura el port, no el framework.

Lo que PyTorch no es (todavía, en la Fase 24)

PyTorch no hace — en el uso de la Fase 24 — ninguna de las siguientes:

  • Entrenar (solo inferencia). La Fase 25 explora autograd; la Fase 24 no.
  • Compilar (sin torch.compile). Modo eager.
  • Distribuir (sin DDP/FSDP). Single-device.
  • Cuantizar (solo fp32 / fp16). Fase 26.
  • Mezclar con autograd custom. La Fase 25 introduce torch.autograd.Function.

El uso de PyTorch en la Fase 24 es el viable mínimo: cargar pesos, definir módulos, correr forward, encajar el kernel custom (lab 03). Nada más.

Dónde engancha PyTorch con el kernel custom

En el lab 03, Borja reemplaza torch.softmax(x, dim=-1) en el LM-head del MiniGPT por una llamada al softmax Triton de theory/03:

class GrammarMiniGPT(nn.Module):
    def __init__(self, ...):
        ...
        self.lm_head = nn.Linear(d, 600)
    def forward(self, x):
        logits = self.lm_head(x)    # shape (B, V) with V=600
        return triton_softmax(logits)   # custom kernel replaces F.softmax

El modelo que rodea el swap no cambia. PyTorch despacha nn.Linear a cuBLAS; el kernel Triton custom corre el softmax. Una línea cambiada; el resto del framework no se ve afectado.

Ese es el contrato que compra la Fase 24: la capacidad de encajar un kernel custom en un modelo PyTorch con precisión quirúrgica. La Fase 25 formalizará cómo PyTorch despacha esa op custom (vía torch.library / registro de ops custom).

Problemas de drill

  1. ¿Por qué x.t().contiguous() reserva memoria? (Pista: la vista transpuesta tiene strides no triviales; hacerla contigua reordena bytes.)
  2. ¿Qué te da model.fc1.weight.storage().data_ptr()? (Pista: la dirección literal HBM/CPU del storage.)
  3. En y = model(x.cuda()); y.cpu(), ¿cuántas copias de memoria ocurrieron? (Pista: dos — H2D para x, D2H para y. Más el trabajo interno CUDA-CUDA, que no cuenta como "copias".)
  4. ¿Por qué fp16 en CUDA da resultados distintos a fp32 en CUDA para el mismo modelo? (Pista: rango dinámico y no-asociatividad de la acumulación.)

Lo que ahora deberías ser capaz de hacer

  1. Enunciar el modelo Tensor = Storage + vista de memoria.
  2. Predecir qué hace .cuda() a nivel de byte.
  3. Portar un módulo NumPy a PyTorch por sustitución mecánica y verificar byte-equivalencia en fp32.
  4. Encajar un kernel custom en el método forward de un nn.Module.
  5. Distinguir qué es "PyTorch el framework" vs "las librerías de kernels a las que despacha" (cuBLAS, cuDNN, ATen, NCCL).

Lo que esta página NO cubre

  • __torch_dispatch__ e internals del dispatcher. Fase 25.
  • Internals del motor de autograd (Function, Variable, recorrido del grafo backward). Fase 25.
  • Registro de ops custom vía torch.library. Fase 25 (brevemente) y Fase 27 (en profundidad para paged attention).
  • torch.compile / Inductor. Fase 25.
  • DataLoader, optimizer, bucle de entrenamiento. La Fase 18 ya enseñó estos en NumPy; las versiones PyTorch se introducen según haga falta en la Fase 25.

Siguiente: lab/00-hello-cuda.md — el chequeo de la toolchain.