El decoder: escribirlo a mano

En M1 le enseñaste al aprendiz el alfabeto: cómo se escribe una jugada, y una biblioteca ordenada para leer rápido. Ahora se sienta a leer. Va a pasar por millones de partidas intentando adivinar, jugada a jugada, qué viene después, y cuando termine jugará «de oído»: sabrá qué suele seguir a qué, aunque todavía no sepa por qué. Eso es preentrenar un GPT, la misma etapa por la que pasa cualquier LLM de texto.

Esta lección explica qué se va a escribir y por qué cada pieza está donde está; la siguiente lo escribe entero. Si nunca has usado PyTorch, pasa antes por PyTorch mínimo: tensores, nn.Module, autograd y un bucle de entrenamiento de veinte líneas.

Qué vas a construir

Un GPT de 38 971 392 parámetros que juega al ajedrez, escrito a mano, sin transformers ni nn.TransformerDecoder: doce bloques, ocho cabezas, 512 dimensiones, un vocabulario de 2 030 jugadas UCIUCIDos cosas con el mismo nombre. La notación UCI escribe una jugada como casilla de origen, casilla de destino y promoción opcional (e7e8q): no depende del contexto, por eso es la tokenización por defecto. El protocolo UCI es la forma en que hablamos con Stockfish desde python-chess. y una ventana de contextoVentana de contextoNúmero máximo de tokens que el modelo puede ver a la vez; en Rukh, 200 (el bloque del dataloader). Lo que no cabe se trunca. Como la atención cuesta el cuadrado de la longitud, una ventana cuatro veces mayor es dieciséis veces más cara, y por eso importa que el 99,8 % de las partidas del recorte quepan en 200 tokens UCI (medido sobre 20 000 partidas de enero; por carácter solo cabe el 7,77 %). de 200 tokensTokenUnidad mínima que el modelo lee y escribe. En Rukh, por defecto, un token es una jugada completa en notación UCI (e2e4); en otras tokenizaciones puede ser un carácter o un trozo de texto aprendido por BPE. El modelo nunca ve letras ni tableros: ve identificadores enteros de tokens.. Se entrena desde cero sobre los lotes de M1, en tu GPU, en menos de una hora, y acaba exportado a ONNXONNXFormato abierto para describir el grafo de una red y sus pesos, independiente del framework que la entrenó. Es lo que permite entrenar en PyTorch y ejecutar en el navegador con onnxruntime-web, sin Python ni backend., publicado en Hugging Face y jugando contra ti en una pestaña del navegador, sin servidor.

Es la arquitectura de GPT-2 casi pieza por pieza. Cambia las 2 030 jugadas por 50 000 trozos de palabra y los 200 tokens por unos miles, y tienes un modelo de texto. Y todo el curso modifica esto: el encoder de M3 es este bloque sin la máscara causalMáscara causalRestricción que impide a cada posición mirar hacia delante: antes del softmax, los pesos de las claves futuras se ponen a menos infinito. Sin ella el modelo vería la jugada que tiene que predecir y la pérdida bajaría sin que aprendiese nada. En Rukh la aplica F.scaled_dot_product_attention(..., is_causal=True). Es lo que hace «causal» a un modelo: en M3 el encoder es el mismo bloque con causal=False, y cada posición ve la secuencia entera., el fine-tuning de M4 son estos pesos con un gradiente más pequeño y el DPO de M5 es esta red con otra pérdida.

Al terminar el módulo tendrás:

• src/rukh/models/decoder.py, el modelo completo con sus tests. Se lee entero en una sentada, y por eso se escribe a mano en lugar de importarlo.
• Dos modelos entrenados: tiny (5 309 952 parámetros, 189 segundos) para iterar sin miedo, y small (42 minutos) como modelo del curso.
• Un harness de evaluación que produce una fila de la tabla única: legalidad, top-1 y top-3, puzles y Elo estimadoElo estimadoFuerza del modelo obtenida ajustando la fórmula logística del Elo a los resultados de partidas contra Stockfish limitado a varios niveles. No es un Elo de la FIDE ni de Lichess: es el número que mejor explica esos resultados, y solo significa algo acompañado de su intervalo de confianza y de la lista de rivales. contra Stockfish con intervalo de confianza.
• El modelo en ONNX, fp16 e int8, con una prueba de paridad contra PyTorch.
• Dos visualizaciones: el mapa de atenciónAtenciónOperación que mezcla las posiciones de una secuencia: cada posición emite una consulta (Q), cada una ofrece una clave (K) y un valor (V); el producto escalar entre consulta y claves, escalado por 1/√d y pasado por softmax, da los pesos con los que se promedian los valores. Es la única parte del Transformer donde las posiciones se hablan entre sí. de una partida y el entrenamiento checkpointCheckpointFotografía de un entrenamiento guardada en disco: pesos, estado del optimizador, paso alcanzado, configuración y procedencia (hash del vocabulario, del manifiesto de datos y SHA de git). Sirve para reanudar, para evaluar y para publicar; en Rukh se escribe uno cada 1 000 pasos más el mejor por pérdida de validación. a checkpoint.

Las cifras medidas son las de la máquina de referencia (RTX 5090); las de parámetros y la aritmética de la atención son cuentas que puedes hacer con una calculadora.

Teoría justa

Cada apartado termina en una decisión escrita en DecoderConfig o en configs/train/small.yaml. Si al final puedes dibujar el modelo de memoria y decir qué se rompe al quitar cada pieza, la lección ha cumplido.

Embeddings: por qué 2 030 × 512

El modelo no sabe qué es e2e4. Recibe el entero 47, o el que le tocara en la enumeración de M1, y un entero no sirve: no puedes decir que el 47 se parece al 312, y su magnitud solo refleja el orden alfabético de las casillas.

La solución es una tabla con una fila de números reales por cada id del vocabulario: un embeddingEmbeddingTabla que asigna un vector aprendido a cada id del vocabulario, y por extensión ese vector. En rukh-small la tabla es de 2 030 × 512: cada jugada UCI entra en el modelo como un punto en un espacio de 512 dimensiones, aprendido a la vez que el resto de la red.. En rukh-small son 2 030 filas (1 968 jugadas, 8 tokens especiales y 54 tramos de Elo) por 512 columnas, la anchura del modelo:

        self.tokens = nn.Embedding(self.cfg.vocab_size, self.cfg.d_model)

↗src/rukh/models/decoder.pylínea 111 · p2

Son 1 039 360 parámetros, el 2,7 % del modelo. Empiezan como ruido gaussiano (al principio e2e4 y a1h8 son dos puntos al azar) y se aprenden con el resto de la red. Al final, las jugadas que aparecen en contextos parecidos acaban cerca: g1f3 y b1c3, que desarrollan un caballo, tendrán un producto escalar alto. Nadie se lo dice al modelo; sale de predecir bien la siguiente jugada.

Nadie decidió que el comino y la cúrcuma fueran juntos en el supermercado. Acabaron en el mismo lineal porque quien compra uno suele comprar el otro. El modelo hace eso con 2 030 jugadas y 512 estanterías, con una diferencia: aquí ninguna dimensión significa nada por sí sola, solo existe la cercanía. Es lo que les pasa a las palabras en un LLM, y la misma idea sostiene los embeddings de un buscador semántico o de un RAG, que llegan en M3.

¿Por qué 512? d_model es el presupuesto de información que fluye por el modelo, y los parámetros por bloque crecen con d². Con 512 y 12 capas salen 38 971 392 parámetros y, en esta RTX 5090, 422 000 tokens por segundo: los 1 024 millones de tokens de small son 42 minutos. Con 768 y 16 capas (medium) serían casi el triple de cómputo por token. 512 es lo bastante grande para aprender ajedrez y lo bastante pequeño para repetir el experimento un martes por la tarde.

Posiciones: tabla aprendida o RoPE

Con solo los vectores de tokens, e2e4 e7e5 g1f3 y g1f3 e7e5 e2e4 le llegan idénticas al modelo: la atención, como verás enseguida, es una suma ponderada, y las sumas no tienen orden. Un modelo de lenguajeModelo de lenguajeRed neuronal entrenada para predecir el siguiente token dada la secuencia anterior. Aplicado a partidas de ajedrez, predice la siguiente jugada dadas las jugadas previas; de esa única tarea emergen las reglas y el juego. sin posiciones es una bolsa de palabras muy cara. En ajedrez es aún peor, porque una partida es su orden: las mismas jugadas en otro orden dan otra posición, o ni siquiera son legales.

Opción 1: posiciones aprendidas. Una segunda tabla, de 200 × 512 = 102 400 parámetros, que se suma al embedding del token:

        if self.positions is not None:
            steps = torch.arange(seq, device=idx.device)
            x = x + self.positions(steps)

↗src/rukh/models/decoder.pylíneas 154-156 · p2

La fila 7 es «lo que significa estar en la octava posición». Es lo que hace GPT-2: simplísimo, y el modelo aprende lo que quiera de cada posición. La pega es que no hay fila 201: no cabe una secuencia más larga que la tabla, y las posiciones altas se entrenan menos porque muchas partidas terminan antes.

Opción 2: RoPE. RoPERoPECodificación de posición rotatoria (rotary position embedding): en vez de sumar un vector de posición, rota por pares las dimensiones de las consultas y las claves con un ángulo proporcional a la posición. El producto escalar entre dos posiciones depende entonces solo de su distancia, y no hace falta una tabla de posiciones. En Rukh es la alternativa a las posiciones aprendidas (pos: rope). no suma nada al token: rota los vectores de consulta y de clave dentro de la atención, con un ángulo proporcional a la posición. Si giras un plano un ángulo m·θ para la posición m, el producto escalar entre las posiciones m y n depende solo de m − n. La atención ve distancias relativas («la jugada de hace dos») en vez de índices absolutos, sin parámetros nuevos. La tabla aprendida numera las sillas de una sala; RoPE solo sabe a cuántos asientos está cada persona de cada otra. Por eso la usan casi todos los LLM actuales (Llama, Mistral, Qwen): reescalando sus frecuencias, el contexto se estira más allá del de entrenamiento.

Con head_dim 64 hay 32 pares de dimensiones, cada uno con su frecuencia:

    inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim))

↗src/rukh/models/decoder.pylínea 28 · p2

El primer par gira 1 radián por posición y distingue distancias cortas; el último gira 0,000133 y en 200 posiciones apenas recorre 0,027 radianes, así que solo dice «principio o final de la partida». Es un reloj con 32 manecillas: leyendo todas a la vez se reconstruye la posición. Rukh usa la variante por mitades, la de GPT-NeoX y Llama:

    half = x.shape[-1] // 2
    rotated = torch.cat([-x[..., half:], x[..., :half]], dim=-1)
    return x * cos_t + rotated * sin_t

↗src/rukh/models/decoder.pylíneas 39-41 · p2

Es la rotación de toda la vida, x·cos − y·sin y x·sin + y·cos, para las 32 parejas de golpe.

Rukh entrena con posiciones aprendidas (pos: "learned") y deja RoPE como opción: con el contexto fijo en 200 tokens su ventaja no se aprovecha. Es una decisión por argumento, no por medida.

Atención, con números

La pieza central; lo demás del Transformer es fontanería alrededor. Cada posición produce tres vectores con tres matrices aprendidas (en el código, una sola qkv que se parte en tres):

• Q, la consulta (query): qué estoy buscando.
• K, la clave (key): qué ofrezco yo a quien me busque.
• V, el valor (value): qué información entrego si me eligen.

Piensa en un tablón de anuncios. Cada jugada anterior deja clavada una nota con dos caras: una etiqueta de qué clase de cosa es («aquí se clavó un alfil») y un contenido, lo que eso implica. Cuando le toca elegir la jugada 40, el modelo escribe qué está buscando («¿alguien dejó una pieza sin defender?»), lo compara con todas las etiquetas y se lleva una mezcla de los contenidos, pesada por lo bien que encajaba cada etiqueta. No elige una nota: se lleva un promedio de todas.

Hagámoslo con números, en un modelo de juguete de dos dimensiones por cabeza. Tres jugadas:

posición 1: e2e4      posición 2: e7e5      posición 3: g1f3

Supón que después de sus proyecciones los vectores son:

k1 = (1, 0)   v1 = (2, 0)
k2 = (1, 1)   v2 = (0, 2)
k3 = (0, 1)   v3 = (1, 1)

q3 = (1, 1)        <- la consulta de la tercera jugada

Paso 1, producto escalar. El parecido entre la consulta y cada clave:

q3 · k1 = 1*1 + 1*0 = 1
q3 · k2 = 1*1 + 1*1 = 2
q3 · k3 = 1*0 + 1*1 = 1

Paso 2, escala 1/√d. Se divide por la raíz de la dimensión de la cabeza, aquí √2 = 1,4142:

1 / 1,4142 = 0,7071
2 / 1,4142 = 1,4142
1 / 1,4142 = 0,7071

¿Por qué? Con componentes de varianza 1, el producto escalar de d dimensiones tiene varianza d: con d = 64, los valores se mueven en un rango ocho veces mayor. Y el softmax es exponencial: una diferencia de 8 entre dos puntuaciones da 0,9997 frente a 0,0003. Sin la escala, un modelo ancho empezaría con atenciones casi binarias, y donde el softmax satura el gradiente es cero. Dividir por √d lo mantiene donde todavía tiene pendiente.

Paso 3, softmax. Exponenciar y normalizar:

exp(0,7071) = 2,0281
exp(1,4142) = 4,1133
exp(0,7071) = 2,0281
suma        = 8,1695

pesos = 0,2483   0,5035   0,2483      (suman 1)

Sin la escala habría salido 0,2120 · 0,5761 · 0,2120. Con dos dimensiones la diferencia es pequeña; con 64 es la diferencia entre un reparto y un interruptor.

Paso 4, mezcla. La salida es la media de los valores ponderada por esos pesos:

0,2483 * (2, 0) = (0,4966 , 0     )
0,5035 * (0, 2) = (0      , 1,0070)
0,2483 * (1, 1) = (0,2483 , 0,2483)
                  -------------------
salida de q3    = (0,7449 , 1,2553)

Eso es todo: softmax(QKᵀ/√d)·V, repetido para cada posición, cabeza y capa. En PyTorch es una llamada:

        out = F.scaled_dot_product_attention(
            q, k, v, dropout_p=self.dropout if self.training else 0.0, is_causal=True
        )

↗src/rukh/models/decoder.pylíneas 66-68 · p2

La máscara causal, que aquí no es opcional

La consulta de la posición 3 miró a las claves 1, 2 y 3. Si la de la posición 1 mirase a la 2 y la 3, estaría mirando el futuro: en el tablón, vería clavada la nota que tiene que adivinar.

El entrenamiento es teacher forcing: se le da la partida entera y en cada posición predice el token siguiente. Si la posición t puede leer el t+1, la tarea pasa de «predice» a «copia lo que tienes al lado». La pérdida se desploma en unos cientos de pasos, las curvas parecen un éxito, y al generar de verdad, donde el futuro no existe, el modelo juega al azar.

La máscara causalMáscara causalRestricción que impide a cada posición mirar hacia delante: antes del softmax, los pesos de las claves futuras se ponen a menos infinito. Sin ella el modelo vería la jugada que tiene que predecir y la pérdida bajaría sin que aprendiese nada. En Rukh la aplica F.scaled_dot_product_attention(..., is_causal=True). Es lo que hace «causal» a un modelo: en M3 el encoder es el mismo bloque con causal=False, y cada posición ve la secuencia entera. lo impide: antes del softmax, toda puntuación de una clave posterior a la consulta se pone a menos infinito, y exp(-inf) = 0. Para completar la matriz toma q₁ = (1, 0) y q₂ = (0,5, 1) con las mismas claves y escala; así q₁·k₁ = 1 → 0,71, q₂·k₁ = 0,5 → 0,35 y q₂·k₂ = 1,5 → 1,06, y la tercera fila es la de antes:

             clave 1   clave 2   clave 3
consulta 1     0,71      -inf      -inf      -> pesos  1     0     0
consulta 2     0,35      1,06      -inf      -> pesos  0,33  0,67  0
consulta 3     0,71      1,41      0,71      -> pesos  0,25  0,50  0,25

Una matriz triangular inferior. La primera posición solo se ve a sí misma, y por eso el <bos> de M1 no es decorativo: es el único contexto de la primera jugada. En el código es is_causal=True, y en labs/m2/causal_mask.py la verificarás empíricamente: cambiar un token futuro no debe mover ni un decimal de los logits pasados.

Multi-cabeza: ocho de 64 en vez de una de 512

Una cabeza de atenciónCabeza de atenciónUna atención completa calculada en un subespacio de la representación. rukh-small parte sus 512 dimensiones en 8 cabezas de 64: cada una aprende su propio criterio de parecido y sus propios pesos, y sus salidas se concatenan y se vuelven a proyectar. Mismo coste que una cabeza de 512, ocho patrones distintos por capa. con d = 512 calcula un único reparto de pesos por posición: un solo criterio de parecido por capa. Multi-cabeza parte las 512 dimensiones en 8 trozos de 64, hace la misma operación ocho veces en paralelo y mezcla las ocho salidas con una proyección.

Y no cuesta nada: la matriz qkv es de 512 × 1536 con una cabeza o con ocho; solo cambia cómo se leen sus columnas. A cambio, ocho patrones de atención por capa, 96 en todo el modelo. Vuelve al tablón: son ocho lectores mirándolo a la vez, cada uno con su pregunta (las piezas sin defender, de quién es el turno, qué se jugó justo antes) sobre las mismas notas.

En los Transformers de texto se han encontrado cabezas muy especializadas: una que mira al token anterior, otra que empareja comillas. En ajedrez cabe esperar cosas análogas, pero no hay garantía, y por eso la isla de «Exportar a ONNX» deja elegir capa y cabeza en vez de enseñarte la que confirma la historia. El límite: con demasiadas cabezas cada una se queda sin resolución, y casi todos los modelos mantienen head_dim entre 64 y 128.

El bloque pre-norm y por qué doce capas se dejan entrenar

El forward de un bloque de Rukh son dos sumas:

    def forward(self, x: Tensor, cos: Tensor | None = None, sin: Tensor | None = None) -> Tensor:
        x = x + self.attn(self.ln1(x), cos, sin)
        return x + self.mlp(self.ln2(x))

↗src/rukh/models/decoder.pylíneas 96-98 · p2

Atención y MLP, cada una precedida de una normalización y sumada a su entrada. Son dos ideas.

La conexión residualConexión residualSuma de la entrada de una subcapa a su salida (x + f(x)). La derivada de esa suma es la identidad más la de f, así que el gradiente llega íntegro a las capas de abajo aunque f aporte poco; el flujo de vectores que atraviesa el modelo de principio a fin se llama por eso corriente residual. es el x +. Sin ella, el gradiente que llega a la capa 1 es el producto de doce jacobianas (la jacobiana de un bloque es la tabla de cuánto se mueve cada salida al mover cada entrada, y la regla de la cadena las multiplica una tras otra): si cada una lo encoge a la mitad, llega dividido por 4 096. Con la suma, la derivada de cada bloque es la identidad más la de la subcapa: el gradiente llega íntegro por el camino recto y cada bloque aporta una corrección. Por eso al vector que recorre el modelo se le llama corriente residual: una memoria compartida que cada bloque lee, retoca y devuelve.

El pre-normPre-normOrden de un bloque Transformer en el que la normalización va antes de la subcapa y su salida se suma a la entrada: x = x + attn(ln1(x)). Deja un camino sin obstáculos entre la pérdida y las primeras capas, y es lo que permite apilar doce bloques y entrenarlos sin trucos. La variante contraria (post-norm) es la del artículo original y necesita mucho más cuidado. es dónde va el LayerNorm. El Transformer original lo ponía después de la suma (post-norm), en el camino residual, y cada capa reescalaba lo que pasaba: entrenar exigía warmup largo y bastante suerte. Pre-norm lo mete dentro de la rama y deja libre el camino desde la pérdida hasta el embedding. Es un cambio de dos caracteres que lo vuelve entrenable, y lo usan GPT-2, Llama y casi todos los LLM de hoy.

Queda un detalle. La corriente residual recibe 24 aportaciones (12 bloques × 2 subcapas), y si cada una tiene la varianza de una inicialización normal, el modelo arranca con activaciones enormes. Por eso las proyecciones de salida se inicializan más pequeñas:

        scale = 0.02 / math.sqrt(2 * self.cfg.n_layer)
        for name, param in self.named_parameters():
            if name.endswith("proj.weight"):
                nn.init.normal_(param, mean=0.0, std=scale)

↗src/rukh/models/decoder.pylíneas 126-129 · p2

El truco de GPT-2: divide la desviación por √24 (0,00408), así que las escrituras en la corriente residual empiezan pequeñas y el modelo sube el volumen de las que le sirven. Se aplica por nombre, a todo lo que acaba en proj.weight, y por eso esos módulos no se pueden renombrar a la ligera.

El MLP es la parte aburrida y dos tercios de los parámetros de un bloque: 512 → 2048 → 512 con GELU en medio. La atención mezcla información entre posiciones; el MLP procesa cada posición con lo que la atención le trajo, y se suele describir como la memoria del modelo.

Tied embeddings: la misma tabla para leer y para escribir

Al final hay una capaCapa (layer)Cada uno de los bloques que se apilan para formar la red; en inglés, layer. Un bloque del decoder de Rukh es atención más MLP con sus dos normalizaciones, y small apila doce de 512 dimensiones y ocho cabezas. También se llama capa a una transformación suelta, como la lineal final que convierte el vector de la última posición en logits. que convierte el vector de la última posición en 2 030 logitsLogitsPuntuación bruta que el modelo asigna a cada token del vocabulario antes de convertirla en probabilidades (con softmax). En Rukh hay un logit por jugada UCI posible; la demo pone a menos infinito los de las jugadas ilegales antes de muestrear, y la temperatura y el top-k actúan sobre ellos., uno por jugada. Es una matriz de 512 × 2 030, la traspuesta de la tabla de embeddings, y Rukh no crea otra: usa la misma.

        self.lm_head = nn.Linear(self.cfg.d_model, self.cfg.vocab_size, bias=False)
        if self.cfg.tie_embeddings:
            self.lm_head.weight = self.tokens.weight

↗src/rukh/models/decoder.pylíneas 118-120 · p2

Ahorra 1 039 360 parámetros, un 2,7 %, que además reciben el doble de gradiente porque se entrenan por los dos extremos. Y supone algo interesante: el logit de la jugada j es el producto escalar entre el estado final y el vector con el que se lee j. Predecir una jugada es apuntar hacia el sitio donde vive en el espacio de embeddings. Se impone que «lo que significa e2e4» y «lo que hay que pensar para decir e2e4» sean lo mismo: casi siempre es buena hipótesis, y a veces limita.

Qué has aprendido

Puedes dibujar MoveDecoder de memoria y defender cada pieza: los embeddings de 2 030 × 512, la máscara causal que define la tarea, las ocho cabezas que cuestan lo mismo que una, el pre-norm que deja entrenar doce capas y la tabla atada por los dos extremos. Es la anatomía de cualquier GPT.

Esta lección no ejecuta nada. Su comprobación llega en el Lab 1 de la lección 10, cuando labs/m2/params.py sume los parámetros bloque a bloque y te salgan 38 971 392, que puedes calcular antes con la calculadora. Si no te cuadra, alguna pieza no está donde crees.