Skip to content

English · Español

Lab 00 — Un bloque transformer, a mano

Lee primero theory/01-transformer-block.md y theory/02-ffn-and-activations.md. No consultes solutions/.

Objetivo

Implementa un único bloque transformer Pre-LN en NumPy y verifica el paso forward contra una referencia derivada a mano en una configuración de juguete minúscula. El objetivo es control numérico completo — si no puedes trazar un bloque en un ejemplo de 2 tokens y \(d_\text{model} = 4\), no lo entiendes.

Setup

Un módulo nuevo: src/minimodel/transformer/. Tres archivos:

src/minimodel/transformer/
├── __init__.py
├── layer_norm.py
├── ffn.py
└── block.py

Cada archivo recibe primero un BLUEPRINT.md (ya hiciste este drill en las Fases 7, 8, 13, 15 — mismo patrón). MultiHeadAttention se importa desde el módulo de la Fase 15.

Tareas

Tarea 1 — LayerNorm

En src/minimodel/transformer/layer_norm.py:

class LayerNorm:
    """LayerNorm over the last axis. Parameters: gamma (scale), beta (shift)."""
    def __init__(self, d_model: int, eps: float = 1e-5):
        ...

    def __call__(self, x: NDArray[np.float64]) -> NDArray[np.float64]:
        """x: (..., d_model) → (..., d_model). Normalises over the last axis."""

Restricciones:

  • NumPy puro. El paso backward es vía el tensor de autograd de la Fase 8 — tu LayerNorm debería aceptar y devolver tensores autograd, no arrays crudos. (Si tu interfaz de autograd de la Fase 8 está terminada — confirma al abrir la fase.)
  • Validar formas: el último eje de la entrada debe ser igual a d_model.
  • Calcular \(\mu, \sigma^2\) sobre el último eje, no sobre batch o secuencia.

Property tests (añade a tests/test_phase17_layer_norm.py):

  1. Media cero, varianza unidad tras la normalización. Para entrada \(x \sim \mathcal{N}(0, I)\), LN(x) (con \(\gamma = 1, \beta = 0\)) debería tener \(\approx 0\) de media y \(\approx 1\) de varianza en el último eje.
  2. Identidad bajo \(\gamma = 1, \beta = 0\) sobre entrada ya normalizada.
  3. Recuperación de shift / scale. Pon \(\gamma = 2, \beta = 5\); verifica que la salida tenga media ≈ 5, varianza ≈ 4.

Tarea 2 — FFN

En src/minimodel/transformer/ffn.py:

class FFN:
    """Position-wise feed-forward: GELU(x W1 + b1) W2 + b2."""
    def __init__(self, d_model: int, d_ff: int):
        ...

    def __call__(self, x): ...

Restricciones:

  • Usa el GELU aproximado: 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x**3))).
  • Inicializa \(W_1, W_2\) con Xavier (default de la Fase 10).
  • Sesgos a cero en la inicialización.

Tarea 3 — TransformerBlock

En src/minimodel/transformer/block.py:

class TransformerBlock:
    def __init__(self, d_model: int, n_heads: int, d_ff: int):
        self.ln1 = LayerNorm(d_model)
        self.attn = MultiHeadAttention(d_model, n_heads)  # from Phase 15
        self.ln2 = LayerNorm(d_model)
        self.ffn = FFN(d_model, d_ff)

    def __call__(self, x):
        # Pre-LN: norm → sublayer → residual add
        x = x + self.attn(self.ln1(x))
        x = x + self.ffn(self.ln2(x))
        return x

Tarea 4 — referencia derivada a mano

Elige una configuración minúscula: \(d_\text{model} = 4, n_\text{heads} = 1, d_\text{ff} = 8\), \(T = 2\) tokens. Inicializa todos los parámetros con valores pequeños explícitos que tú escribas (no aleatorios — usa una lista fija de fracciones como \(0{,}1, 0{,}2, \ldots\)). Calcula el paso forward del bloque en papel, paso a paso:

  1. \(\text{LN}_1(x)\) — calcula \(\mu, \sigma\) por fila, normaliza, escala, desplaza.
  2. \(\text{MHA}(\cdot)\) — proyecciones Q, K, V, productos punto, softmax, proyección de salida. (Con un head, esto es un único attention.)
  3. Residual: \(z = x + \text{MHA output}\).
  4. \(\text{LN}_2(z)\) — mismo procedimiento que el paso 1, sobre \(z\).
  5. \(\text{FFN}(\cdot)\) — up-project, GELU, down-project.
  6. Residual final: \(y = z + \text{FFN output}\).

Ahora ejecuta tu implementación NumPy sobre las mismas entradas y parámetros. Comprueba np.allclose(numpy_result, hand_result, atol=1e-5).

Esto es tedioso. Te llevará 1-2 horas. Hazlo igualmente. Pillarás:

  • Bugs de eje en LayerNorm (normalizar sobre el eje equivocado).
  • Errores de aproximación de GELU (usar la forma exacta cuando dijiste aproximada, o viceversa).
  • Bugs de orden residual (x + LN(attn(x)) en lugar de x + attn(LN(x))).
  • Confusión de formas Q/K/V.

Tarea 5 — tests de cordura

Añade a tests/test_phase17_block.py:

  1. Preservación de forma. Entrada del bloque \((T, d_\text{model})\) → salida \((T, d_\text{model})\), misma forma.
  2. No-op-en-cero. Si todos los parámetros son cero (o cercanos), el bloque debería ser casi-identidad (ya que ambas sublayers contribuyen ~0). Verifica \(|\text{block}(x) - x| < 10^{-3}\) para entrada normalizada.
  3. Comprobación de diferenciabilidad. Una verificación de forward + backward + norma del gradiente usando el autograd de la Fase 8. (Si el autograd no está aún cableado a través del bloque, salta y documenta.)

Mediciones a capturar

  • Wall-clock para un forward de un bloque en \((T, d_\text{model}) = (8, 64)\). Debería ser sub-milisegundo en la CPU de Borja.
  • Coincidencia derivado-a-mano vs NumPy: error absoluto máximo < \(10^{-5}\). Guardar como experiments/<date>-phase-17-block-by-hand/diff.csv.
  • Manifest: experiments/<date>-phase-17-block-by-hand/manifest.json según src/utils/seeding.py.

Aceptación

  • layer_norm.py, ffn.py, block.py existen todos y pasan los tests de forma.
  • El ejemplo derivado a mano de 2 tokens, \(d_\text{model}=4\) coincide con NumPy hasta \(10^{-5}\).
  • Los property tests para LayerNorm, FFN y bloque pasan.
  • No se usa PyTorch. No se usa la librería transformers.
  • Las notas del lab capturan al menos un bug pillado por la comprobación derivada a mano.

Trampas a esperar

  • GELU aproximado vs exacto. Estandariza en uno. Si usas el GELU exacto basado en scipy.special.erf de scipy para la referencia hecha a mano y el GELU aproximado basado en tanh en el código, difieren en ~0,0001 y tu chequeo 1e-5 falla. Usa la forma aproximada en ambos.
  • Eje de LayerNorm. x.mean(axis=-1, keepdims=True), no axis=0. El keepdims es obligatorio para el broadcasting.
  • Forma de salida de MHA. Con \(n_\text{heads} = 1\), MHA colapsa a single-head attention, pero la proyección de salida \(W_O\) sigue ejecutándose (es \(d_\text{model} \to d_\text{model}\)). No te la saltes.
  • __call__ vs forward. Tus módulos de las Fases ⅞ usaban forward(). La convención PyTorch es __call__ envolviendo forward. Mantén consistencia con el resto de src/minimodel.

Siguiente: 01-assemble-mini-gpt.md