Skip to content

English · Español

Lab 01 — Ensamblar Mini-GPT

Lee theory/01-transformer-block.md y theory/03-tied-embeddings-and-lm-head.md. No consultes solutions/.

Objetivo

Cablea tus embeddings de la Fase 13, RoPE de la Fase 16, MHA de la Fase 15 y el TransformerBlock del Lab 00 en una sola clase MiniGPT de extremo a extremo. Ejecuta un forward sobre el ejemplo canónico de 8 tokens de gramática verbal y verifica la forma de la salida y que la información fluye desde la entrada hasta la salida.

Setup

src/minimodel/mini_gpt.py es el nuevo módulo. Configuración bloqueada (de PHASE_17_PLAN.md §1):

config = MiniGPTConfig(
    d_model=64,
    n_heads=4,
    n_layers=2,
    d_ff=256,
    vocab_size=64,
    context_len=32,
)

Estos no son números mágicos — son los números bloqueados que se usan para el resto del proyecto. Una vez escribas mini_gpt.py, no cambies estos defaults sin un commit phase: revise 17-*.

Tareas

Tarea 1 — dataclass MiniGPTConfig

@dataclass(frozen=True)
class MiniGPTConfig:
    d_model: int
    n_heads: int
    n_layers: int
    d_ff: int
    vocab_size: int
    context_len: int

    def __post_init__(self):
        # Validation: d_model % n_heads == 0, etc.
        ...

Añade validación: d_model % n_heads == 0, todas las dimensiones positivas, context_len >= 1. Una mala config debe lanzar en la construcción, no más tarde.

Tarea 2 — clase MiniGPT

class MiniGPT:
    def __init__(self, config: MiniGPTConfig):
        self.config = config
        self.E = Parameter(np.random.randn(config.vocab_size, config.d_model) * 0.02)
        self.blocks = [
            TransformerBlock(config.d_model, config.n_heads, config.d_ff)
            for _ in range(config.n_layers)
        ]
        self.ln_final = LayerNorm(config.d_model)
        # NOTE: no separate LM head — tied with self.E.

    def __call__(self, tokens: NDArray[np.int64]) -> NDArray[np.float64]:
        """tokens: (T,) int64 → logits: (T, vocab_size) float64."""
        h = self.E[tokens]
        for block in self.blocks:
            h = block(h)
        h = self.ln_final(h)
        logits = h @ self.E.T
        return logits

Restricciones:

  • Tied embedding: la cabeza LM reutiliza self.E. Nunca crees self.W_LM como un parámetro separado.
  • Sin softmax dentro de forward — devuelve logits.
  • Acepta tokens como array int64 de forma (T,). Las llamadas con batch ((B, T)) son una extensión opcional; haz primero la forma de secuencia única.

Tarea 3 — ejecutar sobre la secuencia canónica de 8 tokens

El corpus de la Fase 12 te dio un tokenizer. Úsalo para codificar:

<bos> I work , you work , he

Eso son 8 tokens (incluyendo <bos>). Pásalos por Mini-GPT. La salida debe tener forma (8, 64).

model = MiniGPT(config)
tokens = tokenizer.encode("<bos> I work , you work , he")
assert len(tokens) == 8
logits = model(tokens)
assert logits.shape == (8, 64)

Ahora un sanity-check: imprime los top-5 tokens predichos en la posición 7 (la posición tras he — donde el modelo predeciría el siguiente token):

top5 = np.argsort(logits[-1])[-5:][::-1]
print([tokenizer.decode([t]) for t in top5])

Esperado: basura aleatoria. El modelo no está entrenado. El Lab 01 verifica mecanismo, no capacidad. El entrenamiento es la Fase 18 — momento en que esta misma línea debería sacar algo como works cerca de los primeros.

Tarea 4 — cordura: el forward depende del prefijo completo

Verifica que cambiar el token 0 cambia los logits en la posición 7. Específicamente:

tokens_a = tokenizer.encode("<bos> I work , you work , he")
tokens_b = tokens_a.copy()
tokens_b[0] = some_other_token_id  # change <bos> to something else
logits_a = model(tokens_a)
logits_b = model(tokens_b)
# Output at position 7 must differ in some component.
assert not np.allclose(logits_a[7], logits_b[7])

Este es un test débil, pero pilla un bug habitual: olvidarse de cablear el embedding o de encadenar a través de todos los bloques.

Tarea 5 — comparación RoPE-vs-sin-RoPE

Ejecuta el mismo forward con RoPE deshabilitado (reemplaza RoPE por identidad) y con RoPE habilitado. La salida debería diferir. Captura la norma L2 de la diferencia para al menos una posición.

Esto no es un test que puedas suspender — sólo confirma que RoPE está en la ruta. El Lab 03 irá más allá con tests de causalidad.

Mediciones a capturar

  • Wall-clock del paso forward para la secuencia de 8 tokens en la CPU de Borja. Debería ser ~10–50 ms.
  • Forma de salida: (8, 64).
  • Top-5 predicciones en la posición 7 — capturadas (serán aleatorias).
  • Diferencia L2 entre forwards con RoPE habilitado y deshabilitado.

Guarda en experiments/<date>-phase-17-mini-gpt-assembly/manifest.json.

Aceptación

  • src/minimodel/mini_gpt.py existe con MiniGPTConfig y MiniGPT.
  • MiniGPTConfig valida las entradas.
  • Tied embedding: MiniGPT tiene exactamente una matriz de parámetros de forma (vocab_size, d_model). (Inspecciona vía model.E.shape.)
  • El forward sobre la secuencia canónica de 8 tokens produce forma (8, 64).
  • Cambiar el token 0 cambia los logits en la posición 7.
  • Sin softmax dentro del forward (logits devueltos en crudo).
  • No se usa la librería transformers. No PyTorch.

Trampas a esperar

  • Desatado accidental. self.W_LM = self.E.copy() crea un parámetro diferente. No hagas esto. Usa self.E.T directamente en el forward.
  • Olvidar el LayerNorm final. Los transformers Pre-LN añaden un LN tras el último bloque, antes de la cabeza LM. Saltárselo significa que la cabeza LM ve residuales sin normalizar — funciona matemáticamente pero es la arquitectura equivocada. El Lab 02 lo pillará en el inventario de parámetros (te faltarán 128 parámetros).
  • Off-by-one con context_len vs T. context_len es un máximo. La longitud real de la secuencia \(T\) puede ser cualquier cosa hasta ese tope. No hagas padding hasta context_len — sólo la tabla RoPE necesita abarcar hasta context_len.
  • Devolver probabilidades en lugar de logits. El modelo devuelve logits. El softmax es trabajo de la pérdida (Fase 18) y del sampler (Fase 21). Si haces np.exp(logits) / ... dentro del modelo, deshazlo.

Siguiente: 02-parameter-inventory.md