El decoder a mano: config, bloques y el modelo · Rukh

Qué vas a construir

El modelo entero: tres ficheros, 257 líneas contando el __init__.py, y 159 de tests. Es poco, y ese es el argumento para escribirlo a mano en vez de importar nn.TransformerDecoder: cabe en una sentada, y cuando en M3 el encoder sea «lo mismo sin la máscara causal» sabrás qué frase cambias. Todo está aquí, en el tag p2; si copias los bloques en orden, tienes src/rukh/models/.

`config.py`: donde cada decisión se convierte en un número

Aquí acaba todo lo que la lección anterior argumenta.

"""Configuration of the hand-written GPT decoder and the three course presets."""

from __future__ import annotations

from typing import Literal

from pydantic import model_validator

from rukh.config import BaseConfig

src/rukh/models/config.pylíneas 1-9 · p2

BaseConfig es el modelo de pydantic con extra="forbid" de M0: una clave desconocida es un error. Aquí importa más que en ninguna otra config, porque un n_layers: 16 con la ese de más produciría un modelo de doce capas sin que nadie lo supiera.

class DecoderConfig(BaseConfig):
    """Shape of a ``MoveDecoder``.

    ``d_ff`` defaults to ``4 * d_model`` and ``pos`` picks learned positional embeddings or
    rotary embeddings (RoPE) applied to the queries and keys of every head.
    """

    vocab_size: int = 2030
    n_layer: int = 12
    n_head: int = 8
    d_model: int = 512
    d_ff: int | None = None  # None -> 4 * d_model
    block: int = 200
    dropout: float = 0.0
    pos: Literal["learned", "rope"] = "learned"
    tie_embeddings: bool = True

src/rukh/models/config.pylíneas 12-27 · p2

Las decisiones que no son evidentes:

d_ff: int | None = None: el ancho del MLP casi siempre es 4 * d_model, y None significa «el que toca», para no tener que cambiar dos campos al cambiar de talla.
dropout: float = 0.0, a propósito: con 240 millones de tokens y 39 millones de parámetros no hay escasez de datos contra la que regularizar. El encoder de M3 lo enciende, porque verá muchas menos etiquetas.
pos = "learned" por el argumento de la lección anterior. RoPE queda detrás de una bandera, escrito y probado, para que la decisión cambie algún día por una medida.
block: int = 200 es el contexto, y reaparece en el recorte del prompt, el límite de plies, el exportador y los metadatos del .onnx. Vive aquí porque es del modelo.

    @model_validator(mode="after")
    def _check(self) -> DecoderConfig:
        if self.d_model % self.n_head:
            raise ValueError(f"d_model={self.d_model} is not divisible by n_head={self.n_head}")
        if min(self.vocab_size, self.n_layer, self.n_head, self.d_model, self.block) < 1:
            raise ValueError("vocab_size, n_layer, n_head, d_model and block must be positive")
        if not 0.0 <= self.dropout < 1.0:
            raise ValueError(f"dropout must be in [0, 1), got {self.dropout}")
        if self.d_ff is not None and self.d_ff < 1:
            raise ValueError(f"d_ff must be positive, got {self.d_ff}")
        if self.pos == "rope" and (self.d_model // self.n_head) % 2:
            raise ValueError("rope needs an even head dimension")
        return self

src/rukh/models/config.pylíneas 29-41 · p2

El validador corre después de construir el objeto y cubre las tres formas de pedir un modelo imposible: una anchura que no se reparte entre las cabezas, un número no positivo y RoPE con una dimensión de cabeza impar (la rotación va por parejas). Sin él, d_model=100, n_head=8 fallaría mucho más tarde, dentro de un view(), con un mensaje que no menciona ninguno de tus dos números.

    @property
    def head_dim(self) -> int:
        """Width of one attention head."""
        return self.d_model // self.n_head

    @property
    def ff(self) -> int:
        """Hidden width of the MLP (``d_ff`` or ``4 * d_model``)."""
        return self.d_ff if self.d_ff is not None else 4 * self.d_model

src/rukh/models/config.pylíneas 43-51 · p2

head_dim es 64 en los tres presets. Que sea una propiedad y no un campo impide que existan una d_model, una n_head y una head_dim que se contradigan.

PRESETS: dict[str, DecoderConfig] = {
    "tiny": DecoderConfig(n_layer=6, n_head=4, d_model=256),
    "small": DecoderConfig(),
    "medium": DecoderConfig(n_layer=16, n_head=12, d_model=768),
}

src/rukh/models/config.pylíneas 54-58 · p2

Una línea por preset, porque DecoderConfig ya trae los valores de small. Los tres crecen en profundidad y anchura a la vez, siguiendo la diagonal de las leyes de escala.

def preset(name: str) -> DecoderConfig:
    """A copy of a named preset, so callers may mutate it freely."""
    if name not in PRESETS:
        raise ValueError(f"unknown preset {name!r}; expected one of {', '.join(PRESETS)}")
    return PRESETS[name].model_copy(deep=True)

src/rukh/models/config.pylíneas 61-65 · p2

Sin el model_copy(deep=True), un entrenamiento que sobreescribiera cfg.block modificaría el objeto del diccionario, y el siguiente preset pedido en el mismo proceso vendría contaminado. Hay un test para eso.

`decoder.py`: el modelo

La cabecera declara el contrato del fichero.

"""``MoveDecoder``: a hand-written GPT decoder over the fixed UCI vocabulary.

Pure model code: it knows about tensors and nothing about the data pipeline, so it can be
imported by the training loop, the sampler and the ONNX exporter alike. Blocks are pre-norm
(``x = x + attn(ln1(x))``; ``x = x + mlp(ln2(x))``), attention is causal through
``F.scaled_dot_product_attention(..., is_causal=True)``, the MLP uses GELU and the language
modelling head is tied to the token embedding.
"""

from __future__ import annotations

import math

import torch
from torch import Tensor, nn
from torch.nn import functional as F

from rukh.models.config import DecoderConfig

# ``<pad>`` is id 0 in every scheme; the loss must ignore it (see ``loader.IGNORE_INDEX``).
IGNORE_INDEX = 0

src/rukh/models/decoder.pylíneas 1-21 · p2

«Código de modelo puro: sabe de tensores y nada del pipeline de datos». Por eso el bucle, el muestreador y el exportador pueden importarlo sin arrastrarse unos a otros.

IGNORE_INDEX = 0 es <pad>, el número más caro de olvidar del módulo. Sin él en la entropía cruzada, el modelo aprende con toda confianza que después de cualquier cosa viene <pad>, porque es verdad en los datos: todas las secuencias cortas acaban rellenas.

def rope_tables(
    seq_len: int, head_dim: int, device: torch.device, base: float = 10_000.0
) -> tuple[Tensor, Tensor]:
    """``(cos, sin)`` of shape ``(seq_len, head_dim)`` for rotary position embeddings."""
    inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim))
    angles = torch.outer(torch.arange(seq_len, dtype=torch.float32), inv_freq)
    full = torch.cat([angles, angles], dim=-1)
    return full.cos().to(device), full.sin().to(device)

src/rukh/models/decoder.pylíneas 24-31 · p2

inv_freq reparte 32 frecuencias entre las 64 dimensiones de una cabeza, y el cat duplica los ángulos porque la variante por mitades espera la tabla entera. Todo en float32 aunque se entrene en bf16: siete bits de mantisa no bastan para que la posición 199 tenga la fase que le toca.

def apply_rope(x: Tensor, cos: Tensor, sin: Tensor) -> Tensor:
    """Rotate ``x`` of shape ``(B, H, T, D)`` by the angles of the first ``T`` positions."""
    t = x.shape[-2]
    cos_t = cos[:t].to(dtype=x.dtype).view(1, 1, t, -1)
    sin_t = sin[:t].to(dtype=x.dtype).view(1, 1, t, -1)
    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 34-41 · p2

cos[:t] recorta la tabla a la longitud real, y los ángulos bajan a la precisión del tensor después de haberse calculado bien.

class CausalSelfAttention(nn.Module):
    """Multi-head causal self-attention with a single fused ``qkv`` projection."""

    def __init__(self, cfg: DecoderConfig) -> None:
        super().__init__()
        self.n_head = cfg.n_head
        self.head_dim = cfg.head_dim
        self.dropout = cfg.dropout
        self.qkv = nn.Linear(cfg.d_model, 3 * cfg.d_model)
        self.proj = nn.Linear(cfg.d_model, cfg.d_model)
        self.resid_drop = nn.Dropout(cfg.dropout)

src/rukh/models/decoder.pylíneas 44-54 · p2

Una sola matriz qkv produce consulta, clave y valor. Las cabezas no son parámetros sino una reinterpretación de sus columnas, y por eso ocho cuestan lo que una.

    def forward(self, x: Tensor, cos: Tensor | None = None, sin: Tensor | None = None) -> Tensor:
        batch, seq, _ = x.shape
        q, k, v = self.qkv(x).split(x.shape[-1], dim=2)
        shape = (batch, seq, self.n_head, self.head_dim)
        q = q.view(shape).transpose(1, 2)
        k = k.view(shape).transpose(1, 2)
        v = v.view(shape).transpose(1, 2)
        if cos is not None and sin is not None:
            q = apply_rope(q, cos, sin)
            k = apply_rope(k, cos, sin)
        out = F.scaled_dot_product_attention(
            q, k, v, dropout_p=self.dropout if self.training else 0.0, is_causal=True
        )
        out = out.transpose(1, 2).contiguous().view(batch, seq, -1)
        return self.resid_drop(self.proj(out))

src/rukh/models/decoder.pylíneas 56-70 · p2

El forward de la atención es un baile de formas donde se equivoca todo el mundo la primera vez. split da tres tensores (B, T, d), el view parte la última dimensión en cabezas y el transpose(1, 2) las pone delante del tiempo: (B, H, T, D). Ese orden importa, porque la atención multiplica las dos últimas dimensiones; con (B, T, H, D) mezclaría cabezas y posiciones y el modelo entrenaría igual de contento dando basura. La tabla con cada paso y su forma en small está en PyTorch mínimo, junto con view, transpose y contiguous ejecutados sobre números pequeños.

A la vuelta, el contiguous() es obligatorio: tras un transpose la memoria está en otro orden y view se niega a reinterpretarla. Y el dropout_p se apaga a mano al evaluar, porque la función no consulta self.training.

class Mlp(nn.Module):
    """Position-wise GELU feed-forward network."""

    def __init__(self, cfg: DecoderConfig) -> None:
        super().__init__()
        self.fc = nn.Linear(cfg.d_model, cfg.ff)
        self.proj = nn.Linear(cfg.ff, cfg.d_model)
        self.drop = nn.Dropout(cfg.dropout)

    def forward(self, x: Tensor) -> Tensor:
        return self.drop(self.proj(F.gelu(self.fc(x))))

src/rukh/models/decoder.pylíneas 73-83 · p2

Dos tercios de los parámetros de un bloque. La salida se llama proj a propósito: es el nombre que busca la inicialización escalada.

class Block(nn.Module):
    """One pre-norm transformer block."""

    def __init__(self, cfg: DecoderConfig) -> None:
        super().__init__()
        self.ln1 = nn.LayerNorm(cfg.d_model)
        self.attn = CausalSelfAttention(cfg)
        self.ln2 = nn.LayerNorm(cfg.d_model)
        self.mlp = Mlp(cfg)

    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 86-98 · p2

El bloque pre-norm en dos líneas: la normalización está dentro de cada rama, y el camino recto es la autopista por la que el gradiente llega de la pérdida al embedding sin pasar por doce peajes.

class MoveDecoder(nn.Module):
    """GPT decoder over move tokens: ``forward`` returns ``(logits, loss)``.

    ``logits`` is always ``(B, T, vocab_size)``; ``loss`` is ``None`` unless ``targets`` is
    given, in which case it is the mean cross entropy with ``ignore_index=0`` (``<pad>``).
    """

    def __init__(self, cfg: DecoderConfig | None = None) -> None:
        super().__init__()
        self.cfg = cfg or DecoderConfig()
        self.tokens = nn.Embedding(self.cfg.vocab_size, self.cfg.d_model)
        self.positions = (
            nn.Embedding(self.cfg.block, self.cfg.d_model) if self.cfg.pos == "learned" else None
        )
        self.drop = nn.Dropout(self.cfg.dropout)
        self.blocks = nn.ModuleList(Block(self.cfg) for _ in range(self.cfg.n_layer))
        self.ln_f = nn.LayerNorm(self.cfg.d_model)
        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
        if self.cfg.pos == "rope":
            cos, sin = rope_tables(self.cfg.block, self.cfg.head_dim, torch.device("cpu"))
            self.register_buffer("rope_cos", cos, persistent=False)
            self.register_buffer("rope_sin", sin, persistent=False)
        self.apply(self._init_weights)
        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 101-129 · p2

La tabla de posiciones solo existe si pos == "learned"; con RoPE el atributo queda en None, que es lo que consulta el forward. Tres detalles que se pagan si cambian de orden:

El atado no copia: hace que dos nombres apunten al mismo tensor. Atar después de entrenar tiraría uno de los dos.
Las tablas de RoPE son buffers persistent=False: guardarlas sería guardar algo recalculable y fijar el contexto dentro del fichero de pesos.
La inicialización escalada va la última; al revés, apply borraría el ajuste.

    @staticmethod
    def _init_weights(module: nn.Module) -> None:
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)

src/rukh/models/decoder.pylíneas 131-138 · p2

La inicialización de GPT-2: normal de desviación 0,02 y sesgos a cero. Los LayerNorm ya vienen con ganancia 1 y sesgo 0.

    def num_params(self, non_embedding: bool = True) -> int:
        """Parameter count; ``non_embedding`` drops the learned position table (nanoGPT rule)."""
        total = sum(p.numel() for p in self.parameters())
        if non_embedding and self.positions is not None:
            total -= self.positions.weight.numel()
        return total

src/rukh/models/decoder.pylíneas 140-145 · p2

non_embedding=True resta la tabla de posiciones, la convención de nanoGPT, para comparar modelos con contextos distintos. No resta la de tokens, porque atada también es la capa de salida y sí hace cómputo.

    def forward(self, idx: Tensor, targets: Tensor | None = None) -> tuple[Tensor, Tensor | None]:
        """Logits ``(B, T, V)`` for ``idx`` ``(B, T)`` and, with ``targets``, the scalar loss."""
        seq = idx.shape[-1]
        if seq > self.cfg.block:
            raise ValueError(f"sequence of {seq} tokens is longer than block {self.cfg.block}")
        x = self.tokens(idx)
        cos = sin = None
        if self.positions is not None:
            steps = torch.arange(seq, device=idx.device)
            x = x + self.positions(steps)
        else:
            cos, sin = self.rope_cos, self.rope_sin
        x = self.drop(x)
        for block in self.blocks:
            x = block(x, cos, sin)
        logits = self.lm_head(self.ln_f(x))
        loss = None
        if targets is not None:
            loss = F.cross_entropy(
                logits.reshape(-1, logits.shape[-1]),
                targets.reshape(-1),
                ignore_index=IGNORE_INDEX,
            )
        return logits, loss

src/rukh/models/decoder.pylíneas 147-170 · p2

Las posiciones aprendidas se suman al embedding; RoPE viaja como (cos, sin) hasta la atención, que interpreta None como «sin rotación». La guarda de la primera línea da un ValueError con los dos números en vez de un IndexError inútil. Y la pérdida se calcula dentro del modelo, para que el bucle sea _, loss = model(x, y) y el mismo modelo sirva para entrenar e inferir.

    @torch.no_grad()
    def next_logits(self, idx: Tensor) -> Tensor:
        """Logits of the last step only, ``(B, vocab_size)``, cropped to the block size."""
        cropped = idx[:, -self.cfg.block :]
        logits, _ = self(cropped)
        return logits[:, -1, :]

src/rukh/models/decoder.pylíneas 172-177 · p2

El @torch.no_grad() impide que una partida de doscientos plies acumule un grafo de autograd que nadie va a derivar. El recorte idx[:, -block:] es solo una red de seguridad: se lleva por delante la cabecera, y el muestreador de «Muestrear jugadas» recorta de otra manera.

`init.py`: la puerta del paquete

"""Models written from scratch: the ``MoveDecoder`` GPT and its configuration."""

from rukh.models.config import PRESETS, DecoderConfig, preset
from rukh.models.decoder import IGNORE_INDEX, Block, CausalSelfAttention, Mlp, MoveDecoder

__all__ = [
    "IGNORE_INDEX",
    "PRESETS",
    "Block",
    "CausalSelfAttention",
    "DecoderConfig",
    "Mlp",
    "MoveDecoder",
    "preset",
]

src/rukh/models/__init__.pylíneas 1-15 · p2

Solo reexportes: from rukh.models import MoveDecoder funciona desde cualquier sitio, y el día que el decoder se parta en dos ficheros esta lista es lo único que hay que tocar.

// EjecútaloEl modelo ya se construye

Con los tres ficheros de src/rukh/models/ escritos, construye los presets y haz una pasada:

uv run python -c "from rukh.models import MoveDecoder, preset; m = MoveDecoder(preset('small')); print(m.num_params(non_embedding=False), m.num_params())"
uv run python -c "import torch; from rukh.models import MoveDecoder, preset; m = MoveDecoder(preset('tiny')); print(m(torch.ones(2, 10, dtype=torch.long))[0].shape)"

El primero imprime 38971392 38868992: el total y el mismo sin la tabla de posiciones (200 × 512). El segundo, torch.Size([2, 10, 2030]). Un error de forma suele ser un view o un transpose de la atención.

Los tests, que son el contrato

Quince tests que corren en CPU en menos de un segundo y que M3 heredará casi enteros. Si alguna vez escribes o afinas otro modelo de secuencias, estos son los que merece la pena copiar: formas, causalidad, relleno ignorado, determinismo y unos pasos de gradiente que bajan la pérdida.

"""Tests for rukh.models: shapes, causality, loss masking, RoPE, presets and determinism."""

from __future__ import annotations

import pytest
import torch

from rukh.models import PRESETS, DecoderConfig, MoveDecoder, preset

pytestmark = pytest.mark.unit

TOY = DecoderConfig(vocab_size=64, n_layer=2, n_head=4, d_model=32, block=16)


def toy_model(seed: int = 0, **overrides: object) -> MoveDecoder:
    """A deterministic two-layer decoder in eval mode."""
    torch.manual_seed(seed)
    model = MoveDecoder(TOY.model_copy(update=overrides))
    return model.eval()


def toy_batch(batch: int = 2, seq: int = 8, seed: int = 1) -> torch.Tensor:
    generator = torch.Generator().manual_seed(seed)
    return torch.randint(1, TOY.vocab_size, (batch, seq), generator=generator)

tests/unit/test_decoder.pylíneas 1-24 · p2

Un modelo de juguete: un test que tarda cuatro segundos en construir small es un test que nadie ejecuta al guardar.

def test_logits_have_batch_time_vocab_shape() -> None:
    model = toy_model()
    idx = toy_batch()
    logits, loss = model(idx)
    assert logits.shape == (2, 8, TOY.vocab_size)
    assert loss is None


@pytest.mark.parametrize("pos", ["learned", "rope"])
def test_future_tokens_do_not_change_past_logits(pos: str) -> None:
    model = toy_model(pos=pos)
    idx = toy_batch(seq=12)
    cut = 5
    changed = idx.clone()
    changed[:, cut + 1 :] = (changed[:, cut + 1 :] + 7) % (TOY.vocab_size - 1) + 1
    assert not torch.equal(idx, changed)
    with torch.no_grad():
        base, _ = model(idx)
        other, _ = model(changed)
    assert torch.allclose(base[:, : cut + 1], other[:, : cut + 1], atol=1e-5)
    assert not torch.allclose(base[:, cut + 1 :], other[:, cut + 1 :], atol=1e-5)

tests/unit/test_decoder.pylíneas 27-47 · p2

El segundo es el test del módulo: cambia los tokens posteriores a la posición 5 y exige que los logits anteriores no se muevan y los posteriores sí, con las dos formas de posición, porque RoPE toca la atención. assert not torch.equal(idx, changed) comprueba que el propio test cambia algo, y el not allclose sobre el futuro descarta un modelo que devuelva ceros.

@pytest.mark.parametrize("pos", ["learned", "rope"])
def test_rope_and_learned_share_the_output_shape(pos: str) -> None:
    logits, _ = toy_model(pos=pos)(toy_batch())
    assert logits.shape == (2, 8, TOY.vocab_size)


def test_targets_give_a_scalar_loss_that_ignores_padding() -> None:
    model = toy_model()
    idx = toy_batch()
    targets = toy_batch(seed=2)
    _, loss = model(idx, targets)
    assert loss is not None and loss.ndim == 0 and torch.isfinite(loss)

    padded = targets.clone()
    padded[:, 4:] = 0  # <pad>: ignored by the loss
    _, masked_loss = model(idx, padded)
    _, short_loss = model(idx[:, :4], targets[:, :4])
    assert masked_loss is not None and short_loss is not None
    assert torch.allclose(masked_loss, short_loss, atol=1e-5)

    all_pad = torch.zeros_like(targets)
    _, empty = model(idx, all_pad)
    assert empty is not None and torch.isnan(empty)

tests/unit/test_decoder.pylíneas 50-72 · p2

El de la pérdida enmascarada exige que rellenar la segunda mitad dé la misma pérdida que recortarla, lo que solo pasa si ignore_index también saca esas posiciones del denominador. Y con todo relleno la pérdida es NaN, no cero, que sería una pérdida buenísima; por eso el bucle filtra los NaN.

def test_same_seed_gives_the_same_weights_and_logits() -> None:
    idx = toy_batch()
    with torch.no_grad():
        a, _ = toy_model(seed=3)(idx)
        b, _ = toy_model(seed=3)(idx)
        c, _ = toy_model(seed=4)(idx)
    assert torch.equal(a, b)
    assert not torch.allclose(a, c, atol=1e-4)


def test_next_logits_match_the_last_step_of_forward() -> None:
    model = toy_model()
    idx = toy_batch(seq=9)
    with torch.no_grad():
        full, _ = model(idx)
    assert torch.allclose(model.next_logits(idx), full[:, -1], atol=1e-6)


def test_next_logits_crop_sequences_longer_than_the_block() -> None:
    model = toy_model()
    long_idx = toy_batch(seq=TOY.block + 5)
    assert model.next_logits(long_idx).shape == (2, TOY.vocab_size)
    with pytest.raises(ValueError, match="longer than block"):
        model(long_idx)

tests/unit/test_decoder.pylíneas 75-98 · p2

Una secuencia de 21 tokens con block=16 pasa por next_logits y falla por el forward: los dos comportamientos son deliberados, y escribirlos en un test es lo que distingue una decisión de una casualidad.

def test_embeddings_are_tied_and_can_be_untied() -> None:
    tied = toy_model()
    assert tied.lm_head.weight is tied.tokens.weight
    untied = toy_model(tie_embeddings=False)
    assert untied.lm_head.weight is not untied.tokens.weight


def test_preset_sizes(capsys: pytest.CaptureFixture[str]) -> None:
    counts = {}
    for name in PRESETS:
        model = MoveDecoder(preset(name))
        counts[name] = (model.num_params(), model.num_params(non_embedding=False))
    with capsys.disabled():
        print()
        for name, (non_embedding, total) in counts.items():
            print(f"{name:<7} {non_embedding:>12,} non-embedding  {total:>12,} total")
    assert 38_000_000 <= counts["small"][0] <= 45_000_000
    assert counts["tiny"][0] < counts["small"][0] < counts["medium"][0]


def test_preset_copies_are_independent() -> None:
    first = preset("tiny")
    first.dropout = 0.5
    assert preset("tiny").dropout == 0.0
    assert PRESETS["tiny"].dropout == 0.0
    with pytest.raises(ValueError, match="unknown preset"):
        preset("enormous")

tests/unit/test_decoder.pylíneas 101-127 · p2

El atado se comprueba con is, no con ==: con == pasarían dos copias inicializadas con la misma semilla, que es justo el bug que se quiere cazar. test_preset_sizes imprime las tres cuentas incluso con -q, y usa una horquilla porque la cifra exacta ya la comprueba labs/m2/params.py.

def test_config_rejects_unknown_keys_and_bad_shapes() -> None:
    with pytest.raises(ValueError):
        DecoderConfig(n_layers=3)  # type: ignore[call-arg]
    with pytest.raises(ValueError, match="not divisible"):
        DecoderConfig(d_model=100, n_head=8)
    with pytest.raises(ValueError, match="even head dimension"):
        DecoderConfig(d_model=12, n_head=4, pos="rope")


def test_d_ff_defaults_to_four_times_d_model() -> None:
    assert DecoderConfig(d_model=64, n_head=4).ff == 256
    assert DecoderConfig(d_model=64, n_head=4, d_ff=128).ff == 128
    assert MoveDecoder(TOY.model_copy(update={"d_ff": 48})).blocks[0].mlp.fc.out_features == 48

tests/unit/test_decoder.pylíneas 130-142 · p2

La última línea comprueba que el modelo construido tiene un mlp.fc de 48 salidas: probar que la configuración se usa, no solo la configuración.

def test_a_few_steps_of_gradient_descent_reduce_the_loss() -> None:
    torch.manual_seed(5)
    model = MoveDecoder(TOY)
    idx = toy_batch(seq=8, seed=6)
    targets = torch.roll(idx, -1, dims=1)
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-2)
    losses = []
    for _ in range(10):
        _, loss = model(idx, targets)
        assert loss is not None
        losses.append(loss.detach().item())
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    assert losses[-1] < losses[0]

tests/unit/test_decoder.pylíneas 145-159 · p2

Diez pasos de AdamW tienen que bajar la pérdida. No dice que el modelo sea bueno, sino que los gradientes llegan a todas partes: un .detach() de más no falla ningún otro test y hace que el modelo no aprenda nunca.

// EjecútaloLos quince tests del decoder

El fichero de tests está completo. Pásalo entero:

uv run pytest -m unit -q tests/unit/test_decoder.py

15 passed (dos funciones parametrizadas sobre learned y rope), y encima del resumen las tres líneas de test_preset_sizes, con small en 38,971,392.

Qué has aprendido

src/rukh/models/ entero: una configuración que se niega a describir un modelo imposible, un decoder con atención fusionada, una inicialización que cuenta cuántas veces se escribe en la corriente residual y quince tests que son su contrato. Si labs/m2/params.py (lección 10) no imprime 38 971 392 para small, tienes una idea equivocada de tu propio modelo.

Lo siguiente es hacerlo converger: el bucle, el planificador, los checkpoints con procedencia y las configuraciones de entrenamiento.