Phase 10 - Lesson 04

Preentrenamiento de un Mini GPT (124M Parámetros)

This lesson includes a graded coding exercise that runs in your browser, unlocked with lifetime access.

GPT-2 Small tiene 124 millones de parámetros. Eso representa 12 capas de transformer, 12 cabezales de atención y embeddings de 768 dimensiones. Puedes entrenarlo desde cero en una sola GPU en unas pocas horas. La mayoría de las personas nunca hacen esto. Usan checkpoints preentrenados. Pero si no entrenas uno tú mismo, realmente no entiendes lo que está sucediendo dentro del modelo sobre el cual estás construyendo productos.

Tipo: Construcción Lenguajes: Python (con numpy) Requisitos previos: Fase 10, Lecciones 01-03 (Tokenizers, Building a Tokenizer, Data Pipelines) Tiempo: ~120 minutos

Objetivos de Aprendizaje

  • Implementar la arquitectura completa de GPT-2 (124M parámetros) desde cero: embeddings de tokens, embeddings posicionales, bloques de transformer y el cabezal del modelo de lenguaje
  • Entrenar un modelo GPT en un corpus de texto usando la predicción del siguiente token con pérdida de entropía cruzada
  • Implementar la generación de texto autorregresiva con muestreo de temperatura y filtrado top-k/top-p
  • Monitorear las curvas de pérdida de entrenamiento y validar que el modelo aprenda patrones de lenguaje coherentes

El Problema

Sabes lo que es un transformer. Has leído los diagramas. Puedes recitar "attention is all you need" y dibujar cajas etiquetadas como "Multi-Head Attention" en una pizarra.

Nada de eso significa que entiendas lo que sucede cuando un modelo genera texto.

Hay 124,438,272 parámetros en GPT-2 Small (con compartición de pesos/weight tying). Cada uno de ellos se definió ejecutando un bucle de entrenamiento: paso forward, calcular pérdida, paso backward, actualizar pesos. Doce bloques de transformer. Doce cabezales de atención por bloque. Un espacio de embedding de 768 dimensiones. Un vocabulario de 50,257 tokens. Cada vez que el modelo genera un token, todos los 124 millones de parámetros participan en una única cadena de multiplicación de matrices que toma una secuencia de IDs de tokens y produce una distribución de probabilidad sobre el siguiente token.

Si nunca has construido esto tú mismo, estás trabajando con una caja negra. Puedes usar la API. Puedes hacer ajuste fino. Pero cuando algo sale mal -- cuando el modelo alucina, cuando se repite, cuando se rehúsa a seguir instrucciones -- no tienes un modelo mental de por qué.

Esta lección construye GPT-2 Small desde cero. No en PyTorch. En numpy. Cada multiplicación de matrices es visible. Cada gradiente es calculado por tu código. Verás exactamente cómo 124 millones de números conspiran para predecir la siguiente palabra.

El Concepto

La Arquitectura GPT

GPT es un modelo de lenguaje autorregresivo. "Autorregresivo" significa que genera un token a la vez, cada uno condicionado por todos los tokens anteriores. La arquitectura es una pila de bloques decodificadores de transformer.

Aquí está el grafo de computación completo, desde los IDs de tokens hasta las probabilidades del siguiente token:

  1. Entran los IDs de tokens. Formato/Shape: (batch_size, seq_len).
  2. Búsqueda de embedding de tokens. Cada ID se mapea a un vector de 768 dimensiones. Formato/Shape: (batch_size, seq_len, 768).
  3. Búsqueda de embedding de posiciones. Cada posición (0, 1, 2, ...) se mapea a un vector de 768 dimensiones. Mismo formato.
  4. Sumar embeddings de tokens + embeddings de posiciones.
  5. Pasar a través de 12 bloques de transformer.
  6. Normalización de capa final (Layer normalization).
  7. Proyección lineal al tamaño del vocabulario. Formato/Shape: (batch_size, seq_len, vocab_size).
  8. Softmax para obtener las probabilidades.

Ese es el modelo completo. Sin convoluciones. Sin recurrencia. Solo embeddings, atención, redes feedforward y normalizaciones de capa apilados 12 veces.

graph TD
    A["IDs de Tokens\n(batch, seq_len)"] --> B["Embeddings de Tokens\n(batch, seq_len, 768)"]
    A --> C["Embeddings de Posición\n(batch, seq_len, 768)"]
    B --> D["Sumar"]
    C --> D
    D --> E["Bloque Transformer 1"]
    E --> F["Bloque Transformer 2"]
    F --> G["..."]
    G --> H["Bloque Transformer 12"]
    H --> I["Layer Norm"]
    I --> J["Cabezal Linear\n(768 -> 50257)"]
    J --> K["Softmax\nProbabilidades del siguiente token"]

    style A fill:#1a1a2e,stroke:#e94560,color:#fff
    style B fill:#1a1a2e,stroke:#0f3460,color:#fff
    style C fill:#1a1a2e,stroke:#0f3460,color:#fff
    style D fill:#1a1a2e,stroke:#16213e,color:#fff
    style E fill:#1a1a2e,stroke:#e94560,color:#fff
    style F fill:#1a1a2e,stroke:#e94560,color:#fff
    style H fill:#1a1a2e,stroke:#e94560,color:#fff
    style I fill:#1a1a2e,stroke:#16213e,color:#fff
    style J fill:#1a1a2e,stroke:#0f3460,color:#fff
    style K fill:#1a1a2e,stroke:#51cf66,color:#fff

El Bloque Transformer

Cada uno de los 12 bloques sigue el mismo patrón. Arquitectura pre-norm (GPT-2 usa pre-norm, no post-norm como el transformer original):

  1. LayerNorm
  2. Multi-Head Self-Attention
  3. Conexión residual (sumar entrada de vuelta)
  4. LayerNorm
  5. Red Feed-Forward (MLP)
  6. Conexión residual (sumar entrada de vuelta)

Las conexiones residuales son críticas. Sin ellas, los gradientes se desvanecen antes de llegar al bloque 1 durante la retropropagación (backpropagation). Con ellas, los gradientes pueden fluir directamente desde la pérdida a cualquier capa a través del camino de "salto" (skip path). Es por eso que se pueden apilar 12, 32 o incluso 96 bloques (se rumorea que GPT-4 usa 120).

Atención: El Mecanismo Central

La autoatención (self-attention) permite que cada token mire a cada token anterior y decida cuánta atención prestar a cada uno. Aquí está la matemática.

Para cada posición del token, calcula tres vectores a partir de la entrada:

  • Query (Q): "¿Qué estoy buscando?"
  • Key (K): "¿Qué contengo?"
  • Value (V): "¿Qué información llevo?"
Q = input @ W_q    (768 -> 768)
K = input @ W_k    (768 -> 768)
V = input @ W_v    (768 -> 768)

attention_scores = Q @ K^T / sqrt(d_k)
attention_scores = mask(attention_scores)   # causal mask: -inf for future positions
attention_weights = softmax(attention_scores)
output = attention_weights @ V

La máscara causal es lo que hace que GPT sea autorregresivo. La posición 5 puede prestar atención a las posiciones 0-5 pero no a 6, 7, 8, etc. Esto evita que el modelo "haga trampa" al mirar tokens futuros durante el entrenamiento.

La atención multicabezal (multi-head attention) divide el espacio de 768 dimensiones en 12 cabezales de 64 dimensiones cada uno. Cada cabezal aprende un patrón de atención diferente. Un cabezal puede rastrear relaciones sintácticas (concordancia sujeto-verbo). Otro puede rastrear similitud semántica (sinónimos). Otro puede rastrear proximidad posicional (palabras cercanas). Las salidas de los 12 cabezales se concatenan y proyectan de vuelta a 768 dimensiones.

graph LR
    subgraph MultiHead["Atención Multicabezal (12 cabezales)"]
        direction TB
        I["Entrada (768)"] --> S1["Dividir en 12 cabezales"]
        S1 --> H1["Cabezal 1\n(64 dims)"]
        S1 --> H2["Cabezal 2\n(64 dims)"]
        S1 --> H3["..."]
        S1 --> H12["Cabezal 12\n(64 dims)"]
        H1 --> C["Concat (768)"]
        H2 --> C
        H3 --> C
        H12 --> C
        C --> O["Proyección de Salida\n(768 -> 768)"]
    end

    subgraph SingleHead["Cada Cabezal Computa"]
        direction TB
        Q["Q = X @ W_q"] --> A["scores = Q @ K^T / 8"]
        K["K = X @ W_k"] --> A
        A --> M["Aplicar máscara causal"]
        M --> SM["Softmax"]
        SM --> MUL["weights @ V"]
        V["V = X @ W_v"] --> MUL
    end

    style I fill:#1a1a2e,stroke:#e94560,color:#fff
    style O fill:#1a1a2e,stroke:#e94560,color:#fff
    style Q fill:#1a1a2e,stroke:#0f3460,color:#fff
    style K fill:#1a1a2e,stroke:#0f3460,color:#fff
    style V fill:#1a1a2e,stroke:#0f3460,color:#fff

La división por sqrt(d_k) -- sqrt(64) = 8 -- es el escalamiento. Sin ella, los productos punto crecen mucho para vectores de alta dimensión, empujando a softmax a regiones donde los gradientes son casi cero. Esta fue una de las ideas clave del artículo original "Attention Is All You Need".

KV Cache: Por qué la Inferencia es Rápida

Durante el entrenamiento, procesas toda la secuencia a la vez. Durante la inferencia, generas un token a la vez. Sin optimización, generar el token N requiere volver a calcular la atención para todos los N-1 tokens anteriores. Eso es O(N^2) por token generado, o O(N^3) en total para una secuencia de longitud N.

KV Cache resuelve esto. Después de calcular K y V para cada token, almacénalos. Al generar el token N+1, solo necesitas calcular Q para el nuevo token y buscar los K y V almacenados en caché de todos los tokens anteriores. Esto reduce el costo por token de O(N) a O(1) para el cálculo de K y V. El cálculo del score de atención sigue siendo O(N) porque prestas atención a todas las posiciones anteriores, pero evitas multiplicaciones de matrices redundantes en la entrada.

Para GPT-2 con 12 capas y 12 cabezales, el KV cache almacena 2 (K + V) x 12 capas x 12 cabezales x 64 dims = 18,432 valores por token. Para una secuencia de 1024 tokens, eso es aproximadamente 75MB en FP32. Para Llama 3 405B con 128 capas, el KV cache para una sola secuencia puede superar los 10GB. Es por esto que la inferencia de contexto largo está limitada por la memoria (memory-bound).

Prefill vs Decode: Dos Fases de Inferencia

Cuando envías un prompt a un LLM, la inferencia ocurre en dos fases distintas.

Prefill procesa todo tu prompt en paralelo. Todos los tokens son conocidos, por lo que el modelo puede calcular la atención para todas las posiciones simultáneamente. Esta fase está limitada por el poder computacional (compute-bound) -- la GPU realiza multiplicaciones de matrices a máxima capacidad. Para un prompt de 1000 tokens en una A100, el prefill tarda aproximadamente 20-50ms.

Decode genera tokens uno a la vez. Cada nuevo token depende de todos los tokens anteriores. Esta fase está limitada por la memoria (memory-bound) -- el cuello de botella es leer los pesos del modelo y el KV cache de la memoria de la GPU, no la matemática de matrices en sí. Los núcleos de procesamiento de la GPU permanecen mayormente inactivos esperando las lecturas de memoria. Para GPT-2, cada paso de decodificación toma aproximadamente el mismo tiempo independientemente de cuántos FLOPs requieran las multiplicaciones de matrices, porque el ancho de banda de memoria es la restricción.

Esta distinción importa para los sistemas de producción. El rendimiento (throughput) de prefill escala con el cómputo de la GPU (más FLOPS = prefill más rápido). El rendimiento de decode escala con el ancho de banda de memoria (memoria más rápida = decode más rápido). Por eso la H100 de NVIDIA se enfocó en mejoras en el ancho de banda de la memoria con respecto a la A100 -- acelera directamente la generación de tokens.

graph LR
    subgraph Prefill["Fase 1: Prefill"]
        direction TB
        P1["Prompt completo\n(todos los tokens conocidos)"] --> P2["Computación paralela\n(limitada por cómputo)"] --> P3["Construye el KV Cache"]
    end

    subgraph Decode["Fase 2: Decode"]
        direction TB
        D1["Generar token N"] --> D2["Leer KV Cache\n(limitada por memoria)"] --> D3["Añadir al KV Cache"] --> D4["Generar token N+1"]
        D4 -.->|repetir| D1
    end

    Prefill --> Decode

    style P1 fill:#1a1a2e,stroke:#51cf66,color:#fff
    style P2 fill:#1a1a2e,stroke:#51cf66,color:#fff
    style P3 fill:#1a1a2e,stroke:#51cf66,color:#fff
    style D1 fill:#1a1a2e,stroke:#e94560,color:#fff
    style D2 fill:#1a1a2e,stroke:#e94560,color:#fff
    style D3 fill:#1a1a2e,stroke:#e94560,color:#fff
    style D4 fill:#1a1a2e,stroke:#e94560,color:#fff

El Bucle de Entrenamiento

Entrenar un LLM consiste en la predicción del siguiente token. Dados los tokens [0, 1, 2, ..., N-1], predecir los tokens [1, 2, 3, ..., N]. La función de pérdida es la entropía cruzada entre la distribución de probabilidad predicha del modelo y el siguiente token real.

Un paso de entrenamiento:

  1. Paso Forward: Ejecutar el lote a través de los 12 bloques. Obtener los logits (scores previos a softmax) para cada posición.
  2. Calcular Pérdida: Entropía cruzada entre los logits y los tokens objetivo (la entrada desplazada por una posición).
  3. Paso Backward: Calcular los gradientes para los 124M parámetros usando retropropagación.
  4. Paso del Optimizador: Actualizar los pesos. GPT-2 usa Adam con calentamiento de tasa de aprendizaje (learning rate warmup) y decaimiento de coseno.

El cronograma de tasa de aprendizaje (learning rate schedule) importa más de lo que imaginas. GPT-2 realiza un calentamiento (warmup) desde 0 hasta la tasa de aprendizaje máxima durante los primeros 2,000 pasos, y luego decae siguiendo una curva de coseno. Comenzar con una tasa de aprendizaje alta hace que el modelo diverja. Mantener una tasa alta constante causa oscilaciones al final del entrenamiento. El patrón de warmup-then-decay se utiliza en todos los LLMs importantes.

GPT-2 Small: Los Números

Componente Formato Parâmetros
Embeddings de tokens (50257, 768) 38,597,376
Embeddings de posiciones (1024, 768) 786,432
Atención por bloque (W_q, W_k, W_v, W_out) 4 x (768, 768) 2,359,296
FFN por bloque (up + down) (768, 3072) + (3072, 768) 4,718,592
LayerNorms por bloque (2x) 2 x 768 x 2 3,072
LayerNorm Final 768 x 2 1,536
Total por bloque 7,080,960
Total (12 bloques) 85,054,464 + 39,383,808 = 124,438,272

La proyección de salida (cabezal de logits) comparte pesos con la matriz de embeddings de tokens. Esto se llama compartición de pesos (weight tying) -- reduce la cantidad de parámetros en 38M y mejora el rendimiento porque obliga al modelo a usar el mismo espacio de representación para entrada y salida.

Confrúyelo

Paso 1: Capa de Embedding

Los embeddings de tokens mapean cada uno de los 50,257 tokens posibles a un vector de 768 dimensiones. Los embeddings de posiciones agregan información sobre dónde se encuentra cada token en la secuencia. Ambos se suman.

import numpy as np

class Embedding:
    def __init__(self, vocab_size, embed_dim, max_seq_len):
        self.token_embed = np.random.randn(vocab_size, embed_dim) * 0.02
        self.pos_embed = np.random.randn(max_seq_len, embed_dim) * 0.02

    def forward(self, token_ids):
        seq_len = token_ids.shape[-1]
        tok_emb = self.token_embed[token_ids]
        pos_emb = self.pos_embed[:seq_len]
        return tok_emb + pos_emb

La desviación estándar de 0.02 para la inicialización proviene del artículo de GPT-2. Un valor demasiado grande provocará que los pasos forward iniciales producban valores extremos que desestabilicen el entrenamiento. Un valor demasiado pequeño provocará que las salidas iniciales sean casi idénticas para todas las entradas, haciendo que las señales de gradiente iniciales sean inútiles.

Paso 2: Autoatención con Máscara Causal

Primero, atención de un solo cabezal (single-head attention). La máscara causal establece las posiciones futuras en menos infinito antes de softmax, asegurando que cada posición solo pueda prestar atención a sí misma y a las posiciones anteriores.

def attention(Q, K, V, mask=None):
    d_k = Q.shape[-1]
    scores = Q @ K.transpose(0, -1, -2 if Q.ndim == 4 else 1) / np.sqrt(d_k)
    if mask is not None:
        scores = scores + mask
    weights = np.exp(scores - scores.max(axis=-1, keepdims=True))
    weights = weights / weights.sum(axis=-1, keepdims=True)
    return weights @ V

La implementación de softmax resta el máximo antes de la exponenciación. Sin esto, exp(numero_grande) se desborda a infinito. Este es un truque de estabilidad numérica que no cambia el resultado porque softmax(x - c) = softmax(x) para cualquier constante c.

Paso 3: Atención Multicabezal

Divide la entrada de 768 dimensiones en 12 cabezales de 64 dimensiones cada uno. Cada cabezal calcula la atención de forma independiente. Concatena los resultados y proyecta de vuelta a 768 dimensiones.

class MultiHeadAttention:
    def __init__(self, embed_dim, num_heads):
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        self.W_q = np.random.randn(embed_dim, embed_dim) * 0.02
        self.W_k = np.random.randn(embed_dim, embed_dim) * 0.02
        self.W_v = np.random.randn(embed_dim, embed_dim) * 0.02
        self.W_out = np.random.randn(embed_dim, embed_dim) * 0.02

    def forward(self, x, mask=None):
        batch, seq_len, d = x.shape
        Q = (x @ self.W_q).reshape(batch, seq_len, self.num_heads, self.head_dim).transpose(0, 2, 1, 3)
        K = (x @ self.W_k).reshape(batch, seq_len, self.num_heads, self.head_dim).transpose(0, 2, 1, 3)
        V = (x @ self.W_v).reshape(batch, seq_len, self.num_heads, self.head_dim).transpose(0, 2, 1, 3)

        scores = Q @ K.transpose(0, 1, 3, 2) / np.sqrt(self.head_dim)
        if mask is not None:
            scores = scores + mask
        weights = np.exp(scores - scores.max(axis=-1, keepdims=True))
        weights = weights / weights.sum(axis=-1, keepdims=True)
        attn_out = weights @ V

        attn_out = attn_out.transpose(0, 2, 1, 3).reshape(batch, seq_len, d)
        return attn_out @ self.W_out

El truco de reshape-transpose-reshape es la parte más confusa de la atención multicabezal. Esto es lo que sucede: el tensor de (batch, seq_len, 768) se convierte en (batch, seq_len, 12, 64), luego en (batch, 12, seq_len, 64). Ahora cada uno de los 12 cabezales tiene su propia matriz de (seq_len, 64) para ejecutar la atención. Después de la atención, revertemos el proceso: (batch, 12, seq_len, 64) se convierte en (batch, seq_len, 12, 64) y luego en (batch, seq_len, 768).

Paso 4: Bloque Transformer

Un bloque transformer completo: LayerNorm, atención multicabezal con conexión residual, LayerNorm, red feedforward con conexión residual.

class LayerNorm:
    def __init__(self, dim, eps=1e-5):
        self.gamma = np.ones(dim)
        self.beta = np.zeros(dim)
        self.eps = eps

    def forward(self, x):
        mean = x.mean(axis=-1, keepdims=True)
        var = x.var(axis=-1, keepdims=True)
        return self.gamma * (x - mean) / np.sqrt(var + self.eps) + self.beta


class FeedForward:
    def __init__(self, embed_dim, ff_dim):
        self.W1 = np.random.randn(embed_dim, ff_dim) * 0.02
        self.b1 = np.zeros(ff_dim)
        self.W2 = np.random.randn(ff_dim, embed_dim) * 0.02
        self.b2 = np.zeros(embed_dim)

    def forward(self, x):
        h = x @ self.W1 + self.b1
        h = np.maximum(0, h)  # GELU approximation: ReLU for simplicity
        return h @ self.W2 + self.b2


class TransformerBlock:
    def __init__(self, embed_dim, num_heads, ff_dim):
        self.ln1 = LayerNorm(embed_dim)
        self.attn = MultiHeadAttention(embed_dim, num_heads)
        self.ln2 = LayerNorm(embed_dim)
        self.ffn = FeedForward(embed_dim, ff_dim)

    def forward(self, x, mask=None):
        x = x + self.attn.forward(self.ln1.forward(x), mask)
        x = x + self.ffn.forward(self.ln2.forward(x))
        return x

La red feedforward expande la entrada de 768 dimensiones a 3,072 dimensiones (4x), aplica una no linealidad y luego la proyecta de vuelta a 768. Este patrón de expansión-contracción le da al modelo una representación interna "más amplia" para trabajar en cada posición. GPT-2 usa la activación GELU, pero aquí usamos ReLU por simplicidad -- la diferencia es menor para comprender la arquitectura.

Paso 5: Modelo GPT Completo

Apila 12 bloques de transformer. Agrega la capa de embedding al frente y la proyección de salida atrás.

class MiniGPT:
    def __init__(self, vocab_size=50257, embed_dim=768, num_heads=12,
                 num_layers=12, max_seq_len=1024, ff_dim=3072):
        self.embedding = Embedding(vocab_size, embed_dim, max_seq_len)
        self.blocks = [
            TransformerBlock(embed_dim, num_heads, ff_dim)
            for _ in range(num_layers)
        ]
        self.ln_f = LayerNorm(embed_dim)
        self.vocab_size = vocab_size
        self.embed_dim = embed_dim

    def forward(self, token_ids):
        seq_len = token_ids.shape[-1]
        mask = np.triu(np.full((seq_len, seq_len), -1e9), k=1)

        x = self.embedding.forward(token_ids)
        for block in self.blocks:
            x = block.forward(x, mask)
        x = self.ln_f.forward(x)

        logits = x @ self.embedding.token_embed.T
        return logits

    def count_parameters(self):
        total = 0
        total += self.embedding.token_embed.size
        total += self.embedding.pos_embed.size
        for block in self.blocks:
            total += block.attn.W_q.size + block.attn.W_k.size
            total += block.attn.W_v.size + block.attn.W_out.size
            total += block.ffn.W1.size + block.ffn.b1.size
            total += block.ffn.W2.size + block.ffn.b2.size
            total += block.ln1.gamma.size + block.ln1.beta.size
            total += block.ln2.gamma.size + block.ln2.beta.size
        total += self.ln_f.gamma.size + self.ln_f.beta.size
        return total

Observa la compartición de pesos (weight tying): logits = x @ self.embedding.token_embed.T. La proyección de salida reutiliza la matriz de embeddings de tokens (transpuesta). Esto no es solo un truque de ahorro de parámetros. Significa que el modelo utiliza el mismo espacio vectorial para comprender los tokens (embeddings) y predecirlos (salida).

Paso 6: Bucle de Entrenamiento

Para un entrenamiento real con 124M de parámetros, necesitarías una GPU y PyTorch. Este bucle de entrenamiento demuestra la mecánica en un modelo pequeño que se ejecuta puramente en numpy. Usamos un modelo diminuto (4 capas, 4 cabezales, 128 dimensiones) para que sea manejable.

def cross_entropy_loss(logits, targets):
    batch, seq_len, vocab_size = logits.shape
    logits_flat = logits.reshape(-1, vocab_size)
    targets_flat = targets.reshape(-1)

    max_logits = logits_flat.max(axis=-1, keepdims=True)
    log_softmax = logits_flat - max_logits - np.log(
        np.exp(logits_flat - max_logits).sum(axis=-1, keepdims=True)
    )

    loss = -log_softmax[np.arange(len(targets_flat)), targets_flat].mean()
    return loss


def train_mini_gpt(text, vocab_size=256, embed_dim=128, num_heads=4,
                   num_layers=4, seq_len=64, num_steps=200, lr=3e-4):
    tokens = np.array(list(text.encode("utf-8")[:2048]))
    model = MiniGPT(
        vocab_size=vocab_size, embed_dim=embed_dim, num_heads=num_heads,
        num_layers=num_layers, max_seq_len=seq_len, ff_dim=embed_dim * 4
    )

    print(f"Model parameters: {model.count_parameters():,}")
    print(f"Training tokens: {len(tokens):,}")
    print(f"Config: {num_layers} layers, {num_heads} heads, {embed_dim} dims")
    print()

    for step in range(num_steps):
        start_idx = np.random.randint(0, max(1, len(tokens) - seq_len - 1))
        batch_tokens = tokens[start_idx:start_idx + seq_len + 1]

        input_ids = batch_tokens[:-1].reshape(1, -1)
        target_ids = batch_tokens[1:].reshape(1, -1)

        logits = model.forward(input_ids)
        loss = cross_entropy_loss(logits, target_ids)

        if step % 20 == 0:
            print(f"Step {step:4d} | Loss: {loss:.4f}")

    return model

La pérdida comienza cerca de ln(vocab_size) -- para un vocabulario a nivel de bytes de 256 tokens, eso es ln(256) = 5.55. Un modelo aleatorio asigna la misma probabilidad a cada token. A medida que progresa el entrenamiento, la pérdida disminuye porque el modelo aprende a predecir patrones comunes: "th" después de "t", un espacio después de un punto, etc.

En producción, usarías el optimizador Adam con acumulación de gradientes (gradient accumulation), calentamiento de la tasa de aprendizaje (learning rate warmup) y recorte de gradientes (gradient clipping). El bucle forward-pass-loss-backward-update es idéntico. El optimizador es más sofisticado.

Paso 7: Generación de Texto

La generación utiliza el modelo entrenado para predecir un token a la vez. Cada predicción se muestrea a partir de la distribución de salida (o se toma de manera codiciosa/greedy como el argmax).

def generate(model, prompt_tokens, max_new_tokens=100, temperature=0.8):
    tokens = list(prompt_tokens)
    seq_len = model.embedding.pos_embed.shape[0]

    for _ in range(max_new_tokens):
        context = np.array(tokens[-seq_len:]).reshape(1, -1)
        logits = model.forward(context)
        next_logits = logits[0, -1, :]

        next_logits = next_logits / temperature
        probs = np.exp(next_logits - next_logits.max())
        probs = probs / probs.sum()

        next_token = np.random.choice(len(probs), p=probs)
        tokens.append(next_token)

    return tokens

La temperatura controla la aleatoriedad. Una temperatura de 1.0 usa la distribución original. Una temperatura de 0.5 la agudiza (más determinista -- el modelo elige sus opciones principales más a menudo). Una temperatura de 1.5 la aplana (más aleatoria -- los tokens de baja probabilidad obtienen una mayor oportunidad). Una temperatura de 0.0 es decodificación codiciosa (siempre elige el token de mayor probabilidad).

La ventana tokens[-seq_len:] es necesaria porque el modelo tiene una longitud máxima de contexto (1024 para GPT-2). Una vez superada, debes descartar los tokens más antiguos. Esta es la "ventana de contexto" de la que todos hablan.

Úsalo

Demostración Completa de Entrenamiento y Generación

corpus = """The transformer architecture has revolutionized natural language processing.
Attention mechanisms allow the model to focus on relevant parts of the input.
Self-attention computes relationships between all pairs of positions in a sequence.
Multi-head attention splits the representation into multiple subspaces.
Each attention head can learn different types of relationships.
The feedforward network provides nonlinear transformations at each position.
Residual connections enable gradient flow through deep networks.
Layer normalization stabilizes training by normalizing activations.
Position embeddings give the model information about token ordering.
The causal mask ensures autoregressive generation during training.
Pre-training on large text corpora teaches the model general language understanding.
Fine-tuning adapts the pre-trained model to specific downstream tasks."""

model = train_mini_gpt(corpus, num_steps=200)

prompt = list("The transformer".encode("utf-8"))
output_tokens = generate(model, prompt, max_new_tokens=100, temperature=0.8)
generated_text = bytes(output_tokens).decode("utf-8", errors="replace")
print(f"\nGenerated: {generated_text}")

En un corpus pequeño con un modelo pequeño, el texto generado será semi-coherente en el mejor de los casos. Aprenderá algunos patrones a nivel de bytes del texto de entrenamiento, pero no podrá generalizar de la manera en que lo hace GPT-2 con 40GB de datos de entrenamiento y la arquitectura completa de 124M de parámetros. El punto no es la calidad de la salida. El punto es que puedes rastrear cada paso: búsqueda de embedding, cálculo de atención, transformación feedforward, proyección de logits, softmax y muestreo. Cada operación es visible.

Envíalo

Esta lección produce outputs/prompt-gpt-architecture-analyzer.md -- un prompt que analiza las elecciones de arquitectura en cualquier modelo estilo GPT. Aliméntalo con un model card o reporte técnico y desglosará la asignación de parámetros, el diseño de atención y las decisiones de escalado.

Ejercicios

  1. Modifica el modelo para usar 24 capas y 16 cabezales en lugar de 12/12. Cuenta los parámetros. ¿Cómo se compara duplicar la profundidad frente a duplicar el ancho (dimensión del embedding)?

  2. Implementa la función de activación GELU (GELU(x) = x * 0.5 * (1 + erf(x / sqrt(2)))) y reemplaza la ReLU en la red feedforward. Ejecuta el entrenamiento durante 500 pasos con cada activación y compara la pérdida final.

  3. Agrega un KV cache a la función de generación. Almacena los tensores K y V para cada capa después del primer paso forward, y reutilízalos para los tokens subsiguientes. Mide la aceleración: genera 200 tokens con y sin el caché y compara el tiempo real de reloj.

  4. Implementa el muestreo top-k (solo considera los k tokens de mayor probabilidad) y el muestreo top-p (muestreo de núcleo/nucleus sampling: considera el conjunto más pequeño de tokens cuya probabilidad acumulada supere p). Compara la calidad de la salida a temperatura 0.8 con top-k=50 vs top-p=0.95.

  5. Construye un graficador de la curva de pérdida de entrenamiento. Entrena el modelo durante 1000 pasos y grafica pérdida vs paso. Identifica las tres fases: descenso inicial rápido (aprender bytes comunes), fase media más lenta (aprender patrones de bytes) y meseta (sobreajuste/overfitting en el flujo del corpus pequeño). La forma de esta curva es la misma tanto si estás entrenando un modelo de 128 dimensiones como GPT-4.

Términos Clave

Término Lo que la gente dice Lo que realmente significa
Autorregresivo "Genera una palabra a la vez" Cada token de salida está condicionado por todos los tokens anteriores; el modelo predice P(token_n | token_0, ..., token_{n-1})
Máscara causal "No puede ver el futuro" Una matriz triangular superior de valores -infinito que evita la atención a posiciones futuras durante el entrenamiento
Atención multicabezal "Múltiples patrones de atención" Dividir Q, K, V en cabezales paralelos (por ejemplo, 12 cabezales de 64 dimensiones cada uno para GPT-2) para que cada cabezal aprenda diferentes tipos de relaciones
KV Cache "Caché para velocidad" Almacenar los tensores Key (Llave) y Value (Valor) calculados de tokens anteriores para evitar cálculos redundantes durante la generación autorregresiva
Prefill "Procesar el prompt" La primera fase de inferencia donde todos los tokens del prompt se procesan en paralelo; limitada por cómputo en los FLOPS de la GPU
Decode "Generar tokens" La segunda fase de inferencia donde los tokens se generan de uno en uno; limitada por memoria en el ancho de banda de la GPU
Compartición de pesos (Weight tying) "Compartir embeddings" Utilizar la misma matriz para los embeddings de tokens de entrada y el cabezal de proyección de salida; ahorra 38M de parámetros en GPT-2
Conexión residual "Conexión de salto (Skip connection)" Sumar la entrada directamente a la salida de una subcapa (x + sublayer(x)); permite el flujo de gradientes en redes profundas
Normalización de capa (Layer normalization) "Normalizar activaciones" Normalizar a lo largo de la dimensión de características (features) para media 0 y varianza 1, con parámetros de escala (scale) y sesgo (bias) aprendibles
Pérdida de entropía cruzada "Qué tan incorrectas son las predicciones" -log(probabilidad asignada al siguiente token correcto), promediada sobre todas las posiciones; el objetivo estándar de entrenamiento de LLMs

Lecturas Adicionales

0 lifetime access. Curriculum based on AI Engineering from Scratch by Rohit Ghumare (MIT, used under attribution).