Masked move modeling: el preentrenamiento y la receta compartida

"""What every training run in this project does the same way, whatever it is training.

``loop.py`` (next-move prediction) and ``mmm.py`` (masked move modeling) differ in exactly one
place: how a batch becomes a loss. Everything around that step — the optimizer's parameter
groups, the learning-rate schedule, resuming from a checkpoint, writing one, choosing the
device, probing ``torch.compile`` and sending metrics to MLflow — is identical, so it lives here
and both loops call it. Duplicating it would mean two recipes drifting apart, and a run is only
comparable with another one when the recipe is the same.

``RunConfig`` is the shared half of both YAML configs; each loop's config adds what only it
needs (the decoder preset, or the masking parameters).
"""

from __future__ import annotations

import logging
import math
import os
import time
from collections.abc import Callable, Iterator
from contextlib import nullcontext
from datetime import UTC, datetime
from pathlib import Path
from typing import Any, Literal

import torch
from torch import nn
from torch.utils.data import DataLoader

from rukh import paths
from rukh.config import BaseConfig
from rukh.train.checkpoint import load_checkpoint, restore, save_checkpoint
from rukh.train.schedule import lr_at

log = logging.getLogger(__name__)

Batch = tuple[torch.Tensor, torch.Tensor]

class RunConfig(BaseConfig):
    """The half of a training config that does not depend on what is being trained."""

    tokens_dir: str = "data/tokens/uci"
    block: int = 200
    batch_size: int = 64
    grad_accum: int = 4  # 256 effective sequences
    lr: float = 6e-4
    min_lr_ratio: float = 0.1
    warmup: int = 1000
    max_steps: int = 20000
    weight_decay: float = 0.1
    betas: tuple[float, float] = (0.9, 0.95)
    grad_clip: float = 1.0
    precision: Literal["bf16", "fp32"] = "bf16"
    compile: bool = True
    eval_every: int = 500
    eval_batches: int = 50
    ckpt_every: int = 1000
    out_dir: str = "checkpoints"
    seed: int = 42
    run_name: str | None = None
    unique_run_name: bool = True
    """Append a timestamp to ``run_name``: a second run must not overwrite the ``step-*.pt``
    series the ``TrainingReplay`` of the course reads."""
    workers: int = 0  # DataLoader workers; 0 keeps everything in the main process
    log_every: int = 10  # optimizer steps between training metrics

    @property
    def default_name(self) -> str:
        """Run name when the YAML gives none; each loop names its runs after what it trains."""
        return "run"

    def check(self) -> None:
        """Refuse a config no loop could honour."""
        if min(self.max_steps, self.grad_accum, self.batch_size) < 1:
            raise ValueError("max_steps, grad_accum and batch_size must be positive")

def pick_device() -> str:
    """``RUKH_DEVICE`` if set, else CUDA when available, else CPU."""
    env = os.environ.get("RUKH_DEVICE")
    if env:
        return env
    return "cuda" if torch.cuda.is_available() else "cpu"


def param_groups(model: nn.Module, weight_decay: float) -> list[dict[str, Any]]:
    """Decoupled weight decay on matrices only: norms and biases are left alone."""
    decay = [p for p in model.parameters() if p.requires_grad and p.dim() >= 2]
    no_decay = [p for p in model.parameters() if p.requires_grad and p.dim() < 2]
    return [
        {"params": decay, "weight_decay": weight_decay},
        {"params": no_decay, "weight_decay": 0.0},
    ]


def set_lr(optimizer: torch.optim.Optimizer, step: int, cfg: RunConfig) -> float:
    """Apply the scheduled learning rate of ``step`` to every group and return it."""
    lr = lr_at(step, cfg)
    for group in optimizer.param_groups:
        group["lr"] = lr
    return lr

def forever(loader: DataLoader[Batch]) -> Iterator[Batch]:
    """Repeat a loader for as many steps as the schedule asks for."""
    while True:
        yield from loader


def autocast_for(cfg: RunConfig, where: torch.device) -> tuple[Any, bool]:
    """``(context, enabled)``: bf16 autocast on CUDA when asked, a no-op context otherwise."""
    use_bf16 = cfg.precision == "bf16" and where.type == "cuda"
    if not use_bf16:
        return nullcontext(), False
    return torch.autocast(device_type="cuda", dtype=torch.bfloat16), True

def maybe_compile(
    model: nn.Module,
    enabled: bool,
    sample: torch.Tensor | None = None,
    probe: Callable[[nn.Module], None] | None = None,
) -> nn.Module:
    """``torch.compile`` the model when asked; a failure is a warning, never a stopped run.

    ``torch.compile`` is lazy: on Windows without MSVC it only fails when the first forward
    reaches Inductor. ``sample`` (one batch of the training shape, used as input and target by
    the decoder's ``forward``) or a ``probe`` that runs one step itself forces that compilation
    here, where it can still fall back to eager, and Dynamo's own error suppression covers any
    later recompilation for a different shape.
    """
    if not enabled:
        return model
    try:
        import torch._dynamo as dynamo

        dynamo.config.suppress_errors = True
        compiled = torch.compile(model)
        if probe is not None:
            probe(compiled)
            model.zero_grad(set_to_none=True)
        elif sample is not None:
            _, loss = compiled(sample, sample)
            if loss is not None:
                loss.backward()
            model.zero_grad(set_to_none=True)
        return compiled
    except Exception as exc:  # noqa: BLE001 - compilation backends fail in many ways
        log.warning("torch.compile is unavailable, training eagerly: %s", exc)
        model.zero_grad(set_to_none=True)
        return model

def run_dir(cfg: RunConfig, resume: Path | None = None) -> Path:
    """Where this run writes: the resumed run's folder, or ``out_dir/<run name>``.

    The name carries a timestamp unless ``unique_run_name`` is off, because two runs of the same
    config would otherwise write the same ``step-*.pt`` files and the second would quietly
    overwrite the checkpoint series of the first.
    """
    if resume is not None:
        return Path(resume).resolve().parent
    name = cfg.run_name or cfg.default_name
    if cfg.unique_run_name:
        name = f"{name}-{datetime.now(UTC):%Y%m%d-%H%M%S}"
    return paths.resolve(cfg.out_dir) / name


def skip_batches(batches: Iterator[Batch], count: int) -> int:
    """Wind the batch stream forward ``count`` batches and return how many were skipped.

    A resumed run must not start again at the first window of the first epoch: it would train
    twice on the same games while the schedule believes it is halfway. Windows are memmap slices,
    so winding forward is cheap compared with a step, and it is logged because it is not free.
    """
    if count <= 0:
        return 0
    started = time.perf_counter()
    for _ in range(count):
        next(batches)
    log.info("skipped %d batches in %.1f s to resume", count, time.perf_counter() - started)
    return count

def load_resume(
    resume: Path,
    model: nn.Module,
    optimizer: torch.optim.Optimizer,
    where: torch.device,
) -> tuple[int, float, str | None]:
    """Restore weights, optimizer and RNG; return ``(start_step, best_val, mlflow run id)``."""
    payload = load_checkpoint(resume, map_location=where)
    start_step = restore(payload, model, optimizer)
    recorded = payload.get("best_val")
    best_val = float(recorded) if isinstance(recorded, int | float) else math.inf
    previous = payload.get("run_id")
    run_id = str(previous) if isinstance(previous, str) and previous else None
    log.info("resumed %s at step %d (mlflow run %s)", resume, start_step, run_id or "new")
    return start_step, best_val, run_id


def write_checkpoint(
    path: Path,
    *,
    step: int,
    model: nn.Module,
    optimizer: torch.optim.Optimizer | None,
    cfg: BaseConfig,
    model_cfg: BaseConfig,
    vocab_hash: str | None,
    manifest_sha: str | None,
    best_val: float,
    run_id: str | None,
) -> Path:
    """One checkpoint with the provenance every ``rukh`` run records."""
    from rukh.tracking import git_sha

    return save_checkpoint(
        path,
        step=step,
        model=model,
        optimizer=optimizer,
        cfg=cfg.model_dump(mode="json"),
        model_cfg=model_cfg.model_dump(mode="json"),
        vocab_hash=vocab_hash,
        data_manifest_sha=manifest_sha,
        git_sha=git_sha(),
        best_val=None if math.isinf(best_val) else best_val,
        run_id=run_id,
    )

def log_metrics(metrics: dict[str, float], step: int) -> None:
    """Send metrics to the active MLflow run; a tracking failure never stops training."""
    import mlflow

    clean = {key: value for key, value in metrics.items() if not math.isnan(value)}
    if not clean:
        return
    try:
        mlflow.log_metrics(clean, step=step)
    except Exception as exc:  # noqa: BLE001 - tracking is not worth a lost run
        log.warning("could not log metrics at step %d: %s", step, exc)

"""The decoder's training loop: next-token prediction on a packed stream of move tokens.

The recipe follows ``docs/spec/02`` (component 1): AdamW (0.9/0.95, weight decay 0.1 applied
only to matrices), learning rate 6e-4 with 1 000 warmup steps and a cosine decay, an effective
batch of ``batch_size * grad_accum`` sequences, bf16 autocast on CUDA, gradient clipping at 1.0
and optional ``torch.compile``. Everything the run needs to be reproducible (config, seed, git
SHA, vocabulary hash, data manifest hash) goes to MLflow and into every checkpoint. All of that
machinery lives in ``rukh.train.common``, shared with the encoder's masked-move loop; what is
left here is the part that is specific to predicting the next move.

Resuming is meant to be indistinguishable from never having stopped: the batch stream is wound
forward past the windows the first half of the run already saw, and the MLflow run id travels in
the checkpoint so the curve carries on in the same run instead of starting a second one.

Two throughput numbers are logged because they answer different questions: ``tokens_per_s``
counts every position in the window (what the GPU actually processed, comparable across runs)
and ``real_tokens_per_s`` counts only the non-``<pad>`` targets (what the model learned from).
The training loss is accumulated token-weighted rather than as a mean of means, so a
micro-batch with fewer real tokens does not count as much as a full one.
"""

from __future__ import annotations

import logging
import math
import time
from contextlib import nullcontext
from pathlib import Path
from typing import Any, Literal

import torch
from torch import nn
from torch.utils.data import DataLoader

from rukh import paths
from rukh.models import DecoderConfig, MoveDecoder, preset
from rukh.tokenize.loader import IGNORE_INDEX, PackedDataset, make_loader

from rukh.train.checkpoint import BEST_NAME, read_manifest_sha, read_vocab_hash, step_name
from rukh.train.common import (
    Batch,
    RunConfig,
    autocast_for,
    forever,
    load_resume,
    log_metrics,
    maybe_compile,
    param_groups,
    pick_device,
    run_dir,
    set_lr,
    skip_batches,
    write_checkpoint,
)

log = logging.getLogger(__name__)

__all__ = [
    "Batch",
    "TrainConfig",
    "evaluate",
    "forever",
    "log_metrics",
    "maybe_compile",
    "param_groups",
    "pick_device",
    "run_dir",
    "skip_batches",
    "train",
]


class TrainConfig(RunConfig):
    """Everything one decoder training run needs; unknown keys in the YAML are an error."""

    preset: Literal["tiny", "small", "medium"] = "small"
    model: DecoderConfig | None = None  # overrides the preset when given

    @property
    def default_name(self) -> str:
        return self.preset

    def decoder(self) -> DecoderConfig:
        """The decoder config of this run: the preset (or ``model``) with ``block`` applied."""
        base = self.model if self.model is not None else preset(self.preset)
        return base.model_copy(update={"block": self.block})

def train(cfg: TrainConfig, resume: Path | None = None, device: str | None = None) -> Path:
    """Train a ``MoveDecoder`` and return the path of the final checkpoint."""
    cfg.check()
    torch.manual_seed(cfg.seed)
    where = torch.device(device or pick_device())
    tokens_dir = paths.resolve(cfg.tokens_dir)
    model_cfg = cfg.decoder()

    train_set = PackedDataset(tokens_dir / "train", block=cfg.block)
    val_set = PackedDataset(tokens_dir / "val", block=cfg.block)
    if train_set.info.vocab_size != model_cfg.vocab_size:
        raise ValueError(
            f"{tokens_dir / 'train'} has vocab_size {train_set.info.vocab_size}, "
            f"the model expects {model_cfg.vocab_size}"
        )
    train_loader = make_loader(train_set, cfg.batch_size, seed=cfg.seed, workers=cfg.workers)
    val_loader = make_loader(
        val_set, cfg.batch_size, seed=cfg.seed, workers=0, shuffle=False, drop_last=False
    )
    if not len(train_loader):
        raise ValueError(f"{tokens_dir / 'train'} has fewer than {cfg.batch_size} windows")

    model = MoveDecoder(model_cfg).to(where)
    optimizer = torch.optim.AdamW(param_groups(model, cfg.weight_decay), lr=cfg.lr, betas=cfg.betas)
    start_step = 0
    best_val = math.inf
    run_id: str | None = None
    if resume is not None:
        start_step, best_val, run_id = load_resume(resume, model, optimizer, where)

    autocast, use_bf16 = autocast_for(cfg, where)
    warmup = torch.ones((cfg.batch_size, cfg.block), dtype=torch.long, device=where)
    with autocast:  # compile under the same precision the loop will use
        runnable = maybe_compile(model, cfg.compile, warmup)
    runnable.train()

    out_dir = run_dir(cfg, resume)
    out_dir.mkdir(parents=True, exist_ok=True)
    vocab_hash = read_vocab_hash(tokens_dir / "train")
    manifest_sha = read_manifest_sha(paths.data_dir() / "raw" / "manifest.json")
    tokens_per_step = cfg.batch_size * cfg.grad_accum * cfg.block
    batches = forever(train_loader)
    skip_batches(batches, start_step * cfg.grad_accum)
    final = out_dir / step_name(cfg.max_steps)

    from rukh.tracking import start_run

    params = {
        **cfg.model_dump(mode="json"),
        "vocab_hash": vocab_hash,
        "data_manifest_sha": manifest_sha,
        "device": str(where),
        "num_params": model.num_params(),
    }
    with start_run(out_dir.name, params, tags={"preset": cfg.preset}, run_id=run_id) as run:
        log.info("run %s in %s on %s", run.info.run_id, out_dir, where)
        this_run = str(run.info.run_id)

        def save(path: Path, step: int) -> Path:
            return write_checkpoint(
                path,
                step=step,
                model=model,
                optimizer=optimizer,
                cfg=cfg,
                model_cfg=model_cfg,
                vocab_hash=vocab_hash,
                manifest_sha=manifest_sha,
                best_val=best_val,
                run_id=this_run,
            )

        clock = time.perf_counter()
        for step in range(start_step, cfg.max_steps):
            lr = set_lr(optimizer, step, cfg)
            optimizer.zero_grad(set_to_none=True)
            # Token-weighted, on the device: one synchronisation per step instead of one per
            # micro-batch, and a micro-batch with fewer real tokens weighs less in the mean.
            loss_sum = torch.zeros((), device=where)
            real_tokens = torch.zeros((), device=where)
            for _ in range(cfg.grad_accum):
                x, y = next(batches)
                x, y = x.to(where), y.to(where)
                with autocast:
                    _, loss = runnable(x, y)
                assert loss is not None
                (loss / cfg.grad_accum).backward()
                tokens = (y != IGNORE_INDEX).sum()
                loss_sum += loss.detach().float() * tokens
                real_tokens += tokens
            grad_norm = float(nn.utils.clip_grad_norm_(model.parameters(), cfg.grad_clip))
            optimizer.step()

            done = step + 1
            last = done == cfg.max_steps
            if done % cfg.log_every == 0 or last:
                counted = float(real_tokens.item())
                total = float(loss_sum.item()) / counted if counted else math.nan
                elapsed = max(time.perf_counter() - clock, 1e-9)
                steps = min(cfg.log_every, done - start_step)
                log_metrics(
                    {
                        "train/loss": total,
                        "lr": lr,
                        "grad_norm": grad_norm,
                        "tokens_per_s": tokens_per_step * steps / elapsed,
                        "real_tokens_per_s": counted * steps / elapsed,
                    },
                    step=done,
                )
                clock = time.perf_counter()
            if done % cfg.eval_every == 0 or last:
                val_loss, top1 = evaluate(
                    runnable, val_loader, cfg.eval_batches, where, autocast if use_bf16 else None
                )
                log_metrics({"val/loss": val_loss, "val/top1": top1}, step=done)
                log.info("step %d  val/loss %.4f  val/top1 %.4f", done, val_loss, top1)
                if not math.isnan(val_loss) and val_loss < best_val:
                    best_val = val_loss
                    save(out_dir / BEST_NAME, done)
                clock = time.perf_counter()
            if done % cfg.ckpt_every == 0 or last:
                final = save(out_dir / step_name(done), done)
                clock = time.perf_counter()
    return final

"""Masked move modeling: BERT's pretraining objective moved from words to chess moves.

A fraction of the moves of a game are hidden and the bidirectional encoder has to put them back
from both sides of the hole. The 80/10/10 recipe is BERT's and it is kept for the same reason:
if every hidden position were replaced by ``<mask>``, the model would only ever see ``<mask>``
during pretraining and never during fine-tuning, and it would learn a representation that only
works when a token is missing. Ten per cent of random tokens force it to distrust what it reads
(a move that is *there* may still be wrong) and ten per cent left untouched force it to build a
representation of every position, not only of the marked ones.

Two things differ from the decoder's loop, and only two:

* control tokens are never hidden. ``<bos>``, the two Elo tokens, the result token and
  ``<eos>`` are the *condition* of the game, not the signal; predicting the result from a full
  game would be free and would teach nothing about chess.
* the ignored label is ``-100``, not the decoder's ``0``. ``0`` is ``<pad>`` in both schemes, so
  it cannot double as "not predicted" without the loss losing the ability to tell the two apart;
  "nothing to predict" needs a value outside every vocabulary. See
  ``rukh.models.encoder.MMM_IGNORE_INDEX``.

Everything else — optimizer, schedule, checkpoints, resume, MLflow — is ``rukh.train.common``,
the same machinery the decoder uses, so the two runs are comparable.
"""

from __future__ import annotations

import logging
import math
import time
from contextlib import nullcontext
from pathlib import Path
from typing import Any, Literal

import torch
from pydantic import Field, model_validator
from torch import Tensor, nn
from torch.nn import functional as F
from torch.utils.data import DataLoader

from rukh import paths
from rukh.config import BaseConfig
from rukh.models import EncoderConfig, PositionEncoder
from rukh.models.encoder import MMM_IGNORE_INDEX
from rukh.models.squares import CONTROL_IDS as SQUARE_CONTROL_IDS
from rukh.models.squares import MASK_ID as SQUARE_MASK_ID
from rukh.tokenize.loader import PackedDataset, make_loader
from rukh.tokenize.uci_vocab import MASK_ID as MOVE_MASK_ID
from rukh.tokenize.uci_vocab import SPECIALS, elo_tokens
from rukh.train.checkpoint import BEST_NAME, read_manifest_sha, read_vocab_hash, step_name
from rukh.train.common import (
    Batch,
    RunConfig,
    autocast_for,
    forever,
    load_resume,
    log_metrics,
    maybe_compile,
    param_groups,
    pick_device,
    run_dir,
    set_lr,
    skip_batches,
    write_checkpoint,
)

log = logging.getLogger(__name__)

Scheme = Literal["moves", "squares"]

MOVE_CONTROL_IDS: frozenset[int] = frozenset(
    range(len(SPECIALS) + len(elo_tokens("w")) + len(elo_tokens("b")))
)
"""The 62 ids of the UCI vocabulary that are not moves: specials, results and the Elo bins.

They are a prefix of the vocabulary by construction (see ``uci_vocab.build_vocab``), which is
what lets a random replacement simply draw above them.
"""


def control_ids(scheme: Scheme) -> frozenset[int]:
    """Token ids that are never hidden, for one input scheme."""
    return MOVE_CONTROL_IDS if scheme == "moves" else SQUARE_CONTROL_IDS


def mask_id(scheme: Scheme) -> int:
    """The ``<mask>`` token of one input scheme (id 3 in the P1 vocabulary, 1 in squares)."""
    return MOVE_MASK_ID if scheme == "moves" else SQUARE_MASK_ID

class MaskingConfig(BaseConfig):
    """How much to hide and how: BERT's 15 % at 80/10/10."""

    prob: float = 0.15
    mask_ratio: float = 0.8
    random_ratio: float = 0.1
    keep_ratio: float = 0.1
    seed: int = 0

    @model_validator(mode="after")
    def _check(self) -> MaskingConfig:
        if not 0.0 < self.prob <= 1.0:
            raise ValueError(f"prob must be in (0, 1], got {self.prob}")
        ratios = (self.mask_ratio, self.random_ratio, self.keep_ratio)
        if any(ratio < 0.0 for ratio in ratios):
            raise ValueError("mask_ratio, random_ratio and keep_ratio must be non-negative")
        if abs(sum(ratios) - 1.0) > 1e-9:
            raise ValueError(f"mask/random/keep must add up to 1, got {sum(ratios)}")
        return self


def masking_generator(seed: int, device: torch.device | str = "cpu") -> torch.Generator:
    """A generator for ``apply_masking``: one per run, so the masking is seeded but varies."""
    generator = torch.Generator(device=device)
    generator.manual_seed(seed)
    return generator

def apply_masking(
    batch: Tensor,
    cfg: MaskingConfig,
    vocab: int,
    scheme: Scheme = "moves",
    generator: torch.Generator | None = None,
) -> tuple[Tensor, Tensor]:
    """Hide part of ``batch``; return ``(inputs, labels)`` of the same shape.

    ``labels`` is the original token where a position was selected and ``MMM_IGNORE_INDEX``
    everywhere else, so the loss scores exactly the selected positions and nothing else. Of the
    selected ones, ``mask_ratio`` become ``<mask>``, ``random_ratio`` become another *move*
    (never a control token) and ``keep_ratio`` are left as they were.

    ``generator`` makes the draw reproducible; without one the global torch RNG is used, which
    the training loop seeds once from ``RunConfig.seed``.
    """
    control = control_ids(scheme)
    if vocab <= max(control):
        raise ValueError(f"a vocabulary of {vocab} tokens is all control tokens")
    device = batch.device
    is_control = torch.isin(batch, torch.tensor(sorted(control), device=device))
    draw = torch.rand(batch.shape, generator=generator, device=device)
    selected = (draw < cfg.prob) & ~is_control
    labels = torch.where(selected, batch, torch.full_like(batch, MMM_IGNORE_INDEX))

    decision = torch.rand(batch.shape, generator=generator, device=device)
    inputs = batch.clone()
    inputs[selected & (decision < cfg.mask_ratio)] = mask_id(scheme)
    random_upto = cfg.mask_ratio + cfg.random_ratio
    replace = selected & (decision >= cfg.mask_ratio) & (decision < random_upto)
    if bool(replace.any()):
        # Draw above the control block: a random *move*, never a fake Elo or result token.
        low = max(control) + 1
        noise = torch.randint(
            low, vocab, (int(replace.sum()),), generator=generator, device=device, dtype=batch.dtype
        )
        inputs[replace] = noise
    return inputs, labels

class MmmConfig(RunConfig):
    """One masked-move pretraining run; unknown keys in the YAML are an error.

    ``tokens_dir`` is a packed *move* stream, the very one the decoder trains on, so the two
    models see the same games. The ``squares`` scheme has no pack in P1 (positions live in a
    parquet of FENs, not in a token stream), so it is fine-tuned from its labels in
    ``rukh.train.heads`` rather than pretrained here.
    """

    input: Literal["moves"] = "moves"
    model: EncoderConfig | None = None  # overrides the defaults when given
    masking: MaskingConfig = Field(default_factory=MaskingConfig)

    @property
    def default_name(self) -> str:
        return f"encoder-{self.input}"

    def encoder(self) -> EncoderConfig:
        """The encoder config of this run: ``model`` (or the defaults) with ``block`` applied."""
        base = self.model if self.model is not None else EncoderConfig()
        return base.model_copy(update={"block": self.block, "input": self.input})

def masked_step(
    encoder: PositionEncoder,
    body: nn.Module,
    inputs: Tensor,
    labels: Tensor,
    attention_mask: Tensor | None = None,
) -> tuple[Tensor, Tensor]:
    """Logits and masked-move loss for one batch.

    ``body`` is the encoder or its compiled twin; the head and the loss stay eager because
    ``torch.compile`` wraps ``forward`` and nothing else.
    """
    logits = encoder.mlm_head(body(inputs, attention_mask))
    loss = F.cross_entropy(
        logits.reshape(-1, logits.shape[-1]),
        labels.reshape(-1),
        ignore_index=MMM_IGNORE_INDEX,
    )
    return logits, loss

@torch.no_grad()
def evaluate_mmm(
    encoder: PositionEncoder,
    body: nn.Module,
    loader: DataLoader[Batch],
    batches: int,
    device: torch.device,
    masking: MaskingConfig,
    vocab: int,
    scheme: Scheme = "moves",
    autocast: Any = None,
) -> tuple[float, float]:
    """Validation loss and accuracy over the hidden positions of at most ``batches`` batches.

    The validation masking is drawn from a generator reseeded here, so every evaluation of a run
    (and of the next run) hides the same positions and the curve compares like with like.
    """
    was_training = body.training
    body.eval()
    generator = masking_generator(masking.seed, device)
    loss_sum = 0.0
    weighted = 0
    hits = 0
    counted = 0
    for index, (x, _) in enumerate(loader):
        if index >= batches:
            break
        x = x.to(device)
        inputs, labels = apply_masking(x, masking, vocab, scheme, generator)
        with autocast if autocast is not None else nullcontext():
            logits, loss = masked_step(encoder, body, inputs, labels, encoder.padding_mask(x))
        scored = labels != MMM_IGNORE_INDEX
        tokens = int(scored.sum())
        if torch.isfinite(loss) and tokens:
            loss_sum += loss.float().item() * tokens
            weighted += tokens
        hits += int((logits.argmax(dim=-1) == labels)[scored].sum())
        counted += tokens
    if was_training:
        body.train()
    return (loss_sum / weighted if weighted else math.nan, hits / counted if counted else 0.0)

def train_mmm(cfg: MmmConfig, resume: Path | None = None, device: str | None = None) -> Path:
    """Pretrain a ``PositionEncoder`` with masked move modeling; return the last checkpoint."""
    cfg.check()
    torch.manual_seed(cfg.seed)
    where = torch.device(device or pick_device())
    tokens_dir = paths.resolve(cfg.tokens_dir)
    model_cfg = cfg.encoder()

    train_set = PackedDataset(tokens_dir / "train", block=cfg.block)
    val_set = PackedDataset(tokens_dir / "val", block=cfg.block)
    if train_set.info.vocab_size != model_cfg.tokens:
        raise ValueError(
            f"{tokens_dir / 'train'} has vocab_size {train_set.info.vocab_size}, "
            f"the model expects {model_cfg.tokens}"
        )
    train_loader = make_loader(train_set, cfg.batch_size, seed=cfg.seed, workers=cfg.workers)
    val_loader = make_loader(
        val_set, cfg.batch_size, seed=cfg.seed, workers=0, shuffle=False, drop_last=False
    )
    if not len(train_loader):
        raise ValueError(f"{tokens_dir / 'train'} has fewer than {cfg.batch_size} windows")

    model = PositionEncoder(model_cfg).to(where)
    vocab = model_cfg.tokens
    optimizer = torch.optim.AdamW(param_groups(model, cfg.weight_decay), lr=cfg.lr, betas=cfg.betas)
    start_step = 0
    best_val = math.inf
    run_id: str | None = None
    if resume is not None:
        start_step, best_val, run_id = load_resume(resume, model, optimizer, where)

    autocast, use_bf16 = autocast_for(cfg, where)
    generator = masking_generator(cfg.masking.seed, where)

    def probe(compiled: nn.Module) -> None:
        """One masked step of the training shape, to force compilation here and not mid-run.

        With the padding mask the loop actually passes: probing with ``attention_mask=None``
        compiles a graph the loop never runs and buys a recompilation at step 1, which is the
        very cost this probe exists to pay up front.
        """
        warm = torch.full((cfg.batch_size, cfg.block), max(MOVE_CONTROL_IDS) + 1, device=where)
        labels = torch.full_like(warm, MMM_IGNORE_INDEX)
        labels[:, 0] = warm[:, 0]
        _, loss = masked_step(model, compiled, warm, labels, model.padding_mask(warm))
        loss.backward()

    with autocast:  # compile under the same precision the loop will use
        runnable = maybe_compile(model, cfg.compile, probe=probe)
    runnable.train()

    out_dir = run_dir(cfg, resume)
    out_dir.mkdir(parents=True, exist_ok=True)
    vocab_hash = read_vocab_hash(tokens_dir / "train")
    manifest_sha = read_manifest_sha(paths.data_dir() / "raw" / "manifest.json")
    tokens_per_step = cfg.batch_size * cfg.grad_accum * cfg.block
    batches = forever(train_loader)
    skip_batches(batches, start_step * cfg.grad_accum)
    final = out_dir / step_name(cfg.max_steps)

    from rukh.tracking import start_run

    params = {
        **cfg.model_dump(mode="json"),
        "vocab_hash": vocab_hash,
        "data_manifest_sha": manifest_sha,
        "device": str(where),
        "num_params": model.num_params(),
    }
    tags = {"objective": "mmm", "input": cfg.input}
    with start_run(out_dir.name, params, tags=tags, run_id=run_id) as run:
        log.info("run %s in %s on %s", run.info.run_id, out_dir, where)
        this_run = str(run.info.run_id)

        def save(path: Path, step: int) -> Path:
            return write_checkpoint(
                path,
                step=step,
                model=model,
                optimizer=optimizer,
                cfg=cfg,
                model_cfg=model_cfg,
                vocab_hash=vocab_hash,
                manifest_sha=manifest_sha,
                best_val=best_val,
                run_id=this_run,
            )

        clock = time.perf_counter()
        for step in range(start_step, cfg.max_steps):
            lr = set_lr(optimizer, step, cfg)
            optimizer.zero_grad(set_to_none=True)
            # Weighted by hidden positions, on the device: a micro-batch where the draw hid
            # fewer moves must not weigh as much as one where it hid many.
            loss_sum = torch.zeros((), device=where)
            hidden_tokens = torch.zeros((), device=where)
            for _ in range(cfg.grad_accum):
                x, _ = next(batches)
                x = x.to(where)
                inputs, labels = apply_masking(x, cfg.masking, vocab, cfg.input, generator)
                with autocast:
                    _, loss = masked_step(model, runnable, inputs, labels, model.padding_mask(x))
                (loss / cfg.grad_accum).backward()
                tokens = (labels != MMM_IGNORE_INDEX).sum()
                loss_sum += loss.detach().float() * tokens
                hidden_tokens += tokens
            grad_norm = float(nn.utils.clip_grad_norm_(model.parameters(), cfg.grad_clip))
            optimizer.step()

            done = step + 1
            last = done == cfg.max_steps
            if done % cfg.log_every == 0 or last:
                counted = float(hidden_tokens.item())
                total = float(loss_sum.item()) / counted if counted else math.nan
                elapsed = max(time.perf_counter() - clock, 1e-9)
                steps = min(cfg.log_every, done - start_step)

                log_metrics(
                    {
                        "train/loss": total,
                        "lr": lr,
                        "grad_norm": grad_norm,
                        "tokens_per_s": tokens_per_step * steps / elapsed,
                        "masked_tokens_per_s": counted * steps / elapsed,
                    },
                    step=done,
                )

                clock = time.perf_counter()
            if done % cfg.eval_every == 0 or last:
                val_loss, top1 = evaluate_mmm(
                    model,
                    runnable,
                    val_loader,
                    cfg.eval_batches,
                    where,
                    cfg.masking,
                    vocab,
                    cfg.input,
                    autocast if use_bf16 else None,
                )
                log_metrics({"val/loss": val_loss, "val/top1": top1}, step=done)
                log.info("step %d  val/loss %.4f  val/top1 %.4f", done, val_loss, top1)
                if not math.isnan(val_loss) and val_loss < best_val:
                    best_val = val_loss
                    save(out_dir / BEST_NAME, done)
                clock = time.perf_counter()
            if done % cfg.ckpt_every == 0 or last:
                final = save(out_dir / step_name(done), done)
                clock = time.perf_counter()
    return final

# The bidirectional `PositionEncoder` (8 layers, d=384, 6 heads, 15,052,800 parameters)
# pretrained with masked move modeling on the very same packed UCI stream the decoder trains on,
# so the two models can be compared on the same games. Roughly a third of `small`'s size, so the
# 12 000 steps below are well under an hour on a 5090.
input: moves
tokens_dir: data/tokens/uci
block: 200
batch_size: 96
grad_accum: 3 # 288 effective sequences of 200 tokens
lr: 5.0e-4
min_lr_ratio: 0.1
warmup: 1000
max_steps: 12000
weight_decay: 0.1
betas: [0.9, 0.95]
grad_clip: 1.0
precision: bf16
compile: true
eval_every: 500
eval_batches: 50
ckpt_every: 1000
out_dir: checkpoints
seed: 42
run_name: encoder-mmm
# A timestamp is appended: a second run must not overwrite the step-*.pt series.
unique_run_name: true
workers: 4
log_every: 10
masking:
  # BERT's recipe, unchanged: 15 % of the *moves* are hidden, and of those 80 % become <mask>,
  # 10 % become another move and 10 % are left alone. Control tokens are never hidden.
  prob: 0.15
  mask_ratio: 0.8
  random_ratio: 0.1
  keep_ratio: 0.1
  seed: 0

train_app = typer.Typer(
    help="Train a model: the decoder by default, a subcommand for the encoder.",
    invoke_without_command=True,
)
app.add_typer(train_app, name="train")

@train_app.callback(invoke_without_command=True)
def train_cmd(
    ctx: typer.Context,
    config: Annotated[
        Path | None,
        typer.Option(
            "--config", exists=True, dir_okay=False, readable=True, help="Training YAML config."
        ),
    ] = None,
    model_preset: Annotated[
        str | None, typer.Option("--preset", help="Override the preset: tiny, small or medium.")
    ] = None,
    resume: Annotated[
        Path | None,
        typer.Option(
            "--resume", exists=True, dir_okay=False, readable=True, help="Checkpoint to continue."
        ),
    ] = None,
    max_steps: Annotated[
        int | None, typer.Option("--max-steps", help="Override max_steps from the config.")
    ] = None,
) -> None:
    """Train a MoveDecoder from a packed token stream, logging the run to MLflow.

    ``rukh train --config ...`` is the decoder, exactly as it always was; the subcommands train
    the other models (``rukh train encoder``).
    """
    from rukh.config import load_yaml
    from rukh.models import PRESETS
    from rukh.train import TrainConfig, train

    if ctx.invoked_subcommand is not None:
        return
    if config is None:
        typer.echo("error: --config is required (see rukh train --help)", err=True)
        raise typer.Exit(code=2)
    cfg = load_yaml(config, TrainConfig)
    if model_preset is not None:
        if model_preset not in PRESETS:
            typer.echo(f"error: --preset must be one of {', '.join(PRESETS)}", err=True)
            raise typer.Exit(code=2)
        cfg = cfg.model_copy(update={"preset": model_preset, "model": None})
    if max_steps is not None:
        cfg = cfg.model_copy(update={"max_steps": max_steps})
    logging.basicConfig(level=logging.INFO, format="%(asctime)s  %(message)s")
    try:
        checkpoint = train(cfg, resume=resume)
    except (FileNotFoundError, ValueError) as exc:
        typer.echo(f"error: {exc}", err=True)
        raise typer.Exit(code=1) from exc
    typer.echo(f"preset:     {cfg.preset}")
    typer.echo(f"steps:      {cfg.max_steps}")
    typer.echo(f"checkpoint: {checkpoint}")

@train_app.command("encoder")
def train_encoder_cmd(
    config: Annotated[
        Path,
        typer.Option(
            "--config", exists=True, dir_okay=False, readable=True, help="Training YAML config."
        ),
    ],
    resume: Annotated[
        Path | None,
        typer.Option(
            "--resume", exists=True, dir_okay=False, readable=True, help="Checkpoint to continue."
        ),
    ] = None,
    max_steps: Annotated[
        int | None, typer.Option("--max-steps", help="Override max_steps from the config.")
    ] = None,
) -> None:
    """Pretrain the PositionEncoder with masked move modeling, logging the run to MLflow."""
    from rukh.config import load_yaml
    from rukh.train import MmmConfig, train_mmm

    cfg = load_yaml(config, MmmConfig)
    if max_steps is not None:
        cfg = cfg.model_copy(update={"max_steps": max_steps})
    logging.basicConfig(level=logging.INFO, format="%(asctime)s  %(message)s")
    try:
        checkpoint = train_mmm(cfg, resume=resume)
    except (FileNotFoundError, ValueError) as exc:
        typer.echo(f"error: {exc}", err=True)
        raise typer.Exit(code=1) from exc
    typer.echo(f"input:      {cfg.input}")
    typer.echo(
        f"masking:    {cfg.masking.prob:.0%} at "
        f"{cfg.masking.mask_ratio:.0%}/{cfg.masking.random_ratio:.0%}/"
        f"{cfg.masking.keep_ratio:.0%}"
    )
    typer.echo(f"steps:      {cfg.max_steps}")
    typer.echo(f"checkpoint: {checkpoint}")

VOCAB = 128  # above the 62 control ids, so there are real moves to hide
BLOCK = 16
FIRST_MOVE = max(MOVE_CONTROL_IDS) + 1
TOY = EncoderConfig(vocab_size=VOCAB, n_layer=2, n_head=2, d_model=32, block=BLOCK, dropout=0.0)
START = "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq -"


def moves_batch(rows: int = 50, cols: int = 200, seed: int = 0) -> torch.Tensor:
    """``rows * cols`` move tokens, none of them a control token."""
    generator = torch.Generator().manual_seed(seed)
    return torch.randint(FIRST_MOVE, 2030, (rows, cols), generator=generator)

def test_the_80_10_10_split_holds_over_ten_thousand_tokens() -> None:
    batch = moves_batch()
    assert batch.numel() == 10_000
    cfg = MaskingConfig()
    inputs, labels = apply_masking(batch, cfg, 2030, generator=masking_generator(7))

    selected = labels != MMM_IGNORE_INDEX
    picked = int(selected.sum())
    assert abs(picked / batch.numel() - cfg.prob) < 0.02
    masked = int((inputs == MASK_ID).sum())
    kept = int((selected & (inputs == batch)).sum())
    replaced = picked - masked - kept
    assert abs(masked / picked - cfg.mask_ratio) < 0.04
    assert abs(replaced / picked - cfg.random_ratio) < 0.04
    assert abs(kept / picked - cfg.keep_ratio) < 0.04
    # Nothing outside the selection was touched.
    assert torch.equal(inputs[~selected], batch[~selected])


def test_labels_are_minus_100_exactly_where_nothing_was_hidden() -> None:
    batch = moves_batch(rows=8, cols=32, seed=1)
    inputs, labels = apply_masking(batch, MaskingConfig(), 2030, generator=masking_generator(3))
    selected = labels != MMM_IGNORE_INDEX
    assert torch.equal(labels[selected], batch[selected])  # the label is the original token
    assert set(labels[~selected].tolist()) == {MMM_IGNORE_INDEX}
    assert inputs.shape == labels.shape == batch.shape
    # A kept 10 % looks untouched in the input but still carries a label: that is the point.
    assert bool((selected & (inputs == batch)).any())

def test_control_tokens_are_never_hidden() -> None:
    tok = UciTokenizer()
    games = [
        tok.encode_game("e2e4 c7c5 g1f3 d7d6 d2d4", 1800, 1900, "1-0"),
        tok.encode_game("d2d4 g8f6 c2c4 e7e6", 2100, 2000, "0-1"),
    ]
    width = max(len(game) for game in games)
    batch = torch.tensor([game + [0] * (width - len(game)) for game in games])
    inputs, labels = apply_masking(
        batch, MaskingConfig(prob=1.0), len(tok), generator=masking_generator(5)
    )
    control = torch.isin(batch, torch.tensor(sorted(MOVE_CONTROL_IDS)))
    assert torch.equal(inputs[control], batch[control])  # <bos>, Elo, result, <eos>, <pad>
    assert set(labels[control].tolist()) == {MMM_IGNORE_INDEX}
    assert torch.equal(labels[~control], batch[~control])  # every move is hidden at prob 1
    assert len(MOVE_CONTROL_IDS) == 62 and sorted(MOVE_CONTROL_IDS) == list(range(62))


def test_a_random_replacement_is_always_a_real_move() -> None:
    batch = moves_batch(rows=20, cols=100, seed=2)
    only_random = MaskingConfig(prob=1.0, mask_ratio=0.0, random_ratio=1.0, keep_ratio=0.0)
    inputs, _ = apply_masking(batch, only_random, 2030, generator=masking_generator(11))
    assert int(inputs.min()) >= FIRST_MOVE  # never a fake Elo, result or <mask>
    assert int(inputs.max()) < 2030
    assert not torch.equal(inputs, batch)