Los labs del afinado

"""Count how many times the model has seen each Elo header, before blaming the model.

Lab 1 of M4. The claim under test is "the Elo conditioning opens the axis but gives no strength",
which was the conclusion of the previous milestone. Before accepting it, count: the second token of
every packed game *is* White's Elo header, so a histogram of `tokens[starts + 1]` says exactly how
much training each condition got.

Run from the repository root. `--tokens` points at any packed stream:

    uv run python labs/m4/header_histogram.py
    uv run python labs/m4/header_histogram.py --tokens data/tokens-elo/uci/train

The corpus every published model trained on (`data/tokens-v4/uci/train`) starts at `<w1800>` and
has nothing below it: twelve header tokens that exist in the vocabulary, have a row in the
embedding table, and never received a gradient.
"""

import argparse
import collections
import sys
from pathlib import Path

import numpy as np

from rukh.tokenize.uci_vocab import UciTokenizer

# The Windows console is cp1252 by default; ask for UTF-8 before printing the bar chart.
if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(encoding="utf-8", errors="replace")

parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--tokens", default="data/tokens-v4/uci/train", type=Path)
parser.add_argument("--width", default=48, type=int, help="Columns of the widest bar.")
args = parser.parse_args()

tok = UciTokenizer()
tokens = np.load(args.tokens / "tokens.npy", mmap_mode="r")
starts = np.load(args.tokens / "starts.npy")

# Every game is `<bos> <wXXXX> <bXXXX> moves...`, so offset 1 from its start is White's header.
counts = collections.Counter(tokens[starts + 1].tolist())
if not counts:
    raise SystemExit(f"no games in {args.tokens}")

top = max(counts.values())
print(f"{args.tokens}: {len(starts):,} games, {len(tokens):,} tokens\n")
for token_id, n in sorted(counts.items()):
    bar = "#" * max(1, round(args.width * n / top))
    print(f"  {tok.ids[token_id]:>9s}  {n:>9,d}  {bar}")

# The twelve tokens the vocabulary has below 1800, and how many of them the corpus ever used.
below = [tok.vocab[f"<w{elo:04d}>"] for elo in range(600, 1800, 100)]
seen = sum(1 for token_id in below if counts.get(token_id))
print(f"\nheaders below <w1800>: {seen} of {len(below)} ever seen in this corpus")
print("an untrained header is not a header that does not work: it is noise with a row in a table")

"""Check that the LoRA written by hand is LoRA, including against `peft`.

Lab 4 and Lab 5 of M4, in one script and with no checkpoint: it builds its own toy decoder.

Four claims, in order of how easy they are to get wrong:

1. **An untrained adapter is the identity.** `B` starts at zero, so applying LoRA changes nothing
   until a gradient arrives. Any other initialisation moves the weights before the model has
   learned anything, and every comparison against the base starts from a model that is already
   different.
2. **Only the factors train**, and there are exactly `layers x targets x 2 x r x d_model` of them.
3. **Merging is exact.** `W + BA` folded into the weight gives the same logits as the wrapper
   computing it on the fly. That is what lets a 1.6 MB adapter be exported, quantised and served
   as an ordinary model.
4. **It is the same update rule as `peft`'s.** Same base weights, same `A` copied across, same
   optimiser, same batch: the two loss curves have to coincide step by step. If the scaling, the
   side `A` multiplies on, or the zero initialisation of `B` differed, they would separate on the
   first step -- and all three mistakes produce a model that trains and converges anyway.

Run from the repository root:

    uv run python labs/m4/lora_check.py

The fourth check needs the `hf` extra (`uv sync --extra cu128 --extra hf`); without it the script
runs the first three and says so.
"""

import sys

import torch

from rukh.models import DecoderConfig, MoveDecoder, preset
from rukh.models.lora import LoraConfig, apply_lora, lora_modules, merge_lora

if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(encoding="utf-8", errors="replace")

TOY = DecoderConfig(vocab_size=64, n_layer=2, n_head=2, d_model=32, block=16, dropout=0.0)
R, ALPHA, STEPS = 4, 8, 6


def toy(seed: int = 0) -> MoveDecoder:
    torch.manual_seed(seed)
    return MoveDecoder(TOY)


torch.manual_seed(99)
x = torch.randint(1, 64, (4, 16))
y = torch.randint(1, 64, (4, 16))

# --- 1. an untrained adapter is the identity ---------------------------------
model = toy().eval()
with torch.no_grad():
    before = model(x)[0]
trainable = apply_lora(model, LoraConfig(r=R, alpha=ALPHA, targets=("q", "v")))
with torch.no_grad():
    after = model(x)[0]
print(f"1. identity at step 0: max |delta| = {(before - after).abs().max():.2e}")

# --- 2. the parameter count is the formula -----------------------------------
expected = TOY.n_layer * 2 * 2 * R * TOY.d_model
by_grad = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"2. trainable: {trainable:,} == {by_grad:,} == layers*targets*2*r*d_model = {expected:,}")
medium = preset("medium")
real = medium.n_layer * 2 * 2 * 8 * medium.d_model
print(f"   on `medium` with r=8: {real:,} of 115,120,128 = {100 * real / 115_120_128:.3f} %")
print(f"   as float32: {real * 4:,} bytes = {real * 4 / 1e6:.1f} MB")

# --- 3. merging is exact ------------------------------------------------------
torch.manual_seed(7)
with torch.no_grad():
    for _, adapter in lora_modules(model):
        for b in adapter.b:
            b.normal_(std=0.05)
    adapted = model(x)[0].clone()
merge_lora(model)
with torch.no_grad():
    merged = model(x)[0]
print(f"3. merge is exact: max |delta| = {(adapted - merged).abs().max():.2e}")
print(f"   and no wrappers are left: {len(list(lora_modules(model)))} adapters")

# --- 4. the same update rule as `peft` ---------------------------------------
try:
    import peft

    from rukh.models.hf import RukhForCausalLM
except ImportError as exc:
    print(f"4. skipped: needs the `hf` extra ({exc})")
    raise SystemExit(0) from None

ours = toy()
apply_lora(ours, LoraConfig(r=R, alpha=ALPHA, targets=("attn_out",)))
theirs = peft.get_peft_model(
    RukhForCausalLM.from_decoder(toy()),
    peft.LoraConfig(
        r=R, lora_alpha=ALPHA, lora_dropout=0.0, bias="none", target_modules=["attn.proj"]
    ),
)

# `attn.proj` and not `q`/`v`: `peft` cannot express a per-slice adapter on a fused `qkv`, so over
# there `target_modules=["qkv"]` would be ONE A/B pair of 2304 rows against our one per slice.
# Different parameterisations; comparing them would measure that difference, not the code.
ours_by_layer = dict(lora_modules(ours))
with torch.no_grad():
    for name, module in theirs.named_modules():
        if "default" not in getattr(module, "lora_A", {}):
            continue
        key = name.replace("base_model.model.decoder.", "").replace(".base_layer", "")
        module.lora_A["default"].weight.copy_(ours_by_layer[key].a[0])

mine = torch.optim.SGD([p for p in ours.parameters() if p.requires_grad], lr=0.5)
yours = torch.optim.SGD([p for p in theirs.parameters() if p.requires_grad], lr=0.5)
print("4. step   ours          peft          |delta|")
for step in range(STEPS):
    mine.zero_grad()
    a = ours(x, y)[1]
    a.backward()
    mine.step()

    yours.zero_grad()
    b = theirs(x, labels=y).loss
    b.backward()
    yours.step()
    print(f"   {step:>4d}   {a.item():.8f}    {b.item():.8f}    {abs(a.item() - b.item()):.2e}")

"""How many games each pair of conditions would need before its intervals stop overlapping.

Lab 6 of M4. The sweep of this milestone leaves `1500 < 2000 < 2400` unproven, and there are two
very different reasons a comparison can come out unproven: too few games, or no difference to
find. Telling them apart is arithmetic, not opinion.

The ladder plays `n` games per condition and scores them, so the score is a proportion and its
standard error is `sqrt(p (1 - p) / n)`. Two 95 % intervals stop touching when the gap between
the scores is wider than `1.96 * (se1 + se2)`, which for two proportions near a half means

    n  >  3.84 / (delta p)^2

per condition. Put the measured gap in and the answer comes out in games. A pair that needs two
hundred is a pair worth re-running; a pair that needs twenty thousand is a pair where the honest
conclusion is that the difference is not there, and no amount of machine time will change it.

Run from the repository root, on the sweep's own file:

    uv run python labs/m4/games_needed.py
    uv run python labs/m4/games_needed.py --only 1500,2000,2400

``--only`` keeps just the conditions it names, in the order they are given, which is how the
acceptance criterion of this project is read: it asks about `<w1500>`, `<w2000>` and `<w2400>`
and about no other pair.
"""

import argparse
import json
import math
import sys
from pathlib import Path

if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(encoding="utf-8", errors="replace")

Z = 1.959963985  # two-sided 95 %

parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--results", default="artifacts/eval/medium-elo-elo-sweep/results.json")
parser.add_argument("--games", type=int, default=160, help="Games per condition already played.")
parser.add_argument("--only", default="", help="Comma-separated headers to keep, e.g. 1500,2000.")
args = parser.parse_args()

rows = json.loads(Path(args.results).read_text(encoding="utf-8"))["rows"]
if args.only:
    wanted = [int(value) for value in args.only.split(",")]
    by_header = {int(row["header_elo"]): row for row in rows}
    missing = [elo for elo in wanted if elo not in by_header]
    if missing:
        raise SystemExit(f"{args.results} has no condition for {missing}")
    rows = [by_header[elo] for elo in wanted]
if len(rows) < 2:
    raise SystemExit(f"{args.results} has {len(rows)} condition(s); a comparison needs two")

def se(p: float, n: int) -> float:
    return math.sqrt(max(p * (1.0 - p), 1e-12) / n)


def needed(p1: float, p2: float) -> float:
    """Games per condition for the two 95 % intervals to stop touching."""
    gap = abs(p2 - p1)
    if gap <= 0:
        return math.inf
    spread = math.sqrt(p1 * (1 - p1)) + math.sqrt(p2 * (1 - p2))
    return (Z * spread / gap) ** 2

print(f"{args.results}: {len(rows)} conditions, {args.games} games each\n")
print("  pair                 score      gap   separated now   games needed each")
for low, high in zip(rows, rows[1:], strict=False):
    p1, p2 = low["score"], high["score"]
    gap = p2 - p1
    apart = (p1 + Z * se(p1, args.games)) < (p2 - Z * se(p2, args.games))
    n = needed(p1, p2)
    n_text = "never" if not math.isfinite(n) else f"{math.ceil(n):,}"
    print(
        f"  <w{low['header_elo']}> -> <w{high['header_elo']}>  "
        f"{p1:.3f}->{p2:.3f}  {gap:+.3f}   {'yes' if apart else 'no ':<13}   {n_text}"
    )

ends = (rows[0], rows[-1])
p1, p2 = ends[0]["score"], ends[1]["score"]
apart = (p1 + Z * se(p1, args.games)) < (p2 - Z * se(p2, args.games))
print(
    f"\n  ends <w{ends[0]['header_elo']}> -> <w{ends[1]['header_elo']}>: "
    f"{p1:.3f} -> {p2:.3f}, separated now: {'yes' if apart else 'no'}, "
    f"needs {math.ceil(needed(p1, p2)):,} games each"
)
print(
    "\na pair that needs a few hundred games is a measurement problem;\n"
    "a pair that needs tens of thousands is a difference that is not there"
)

"""How many directions a rank-8 adapter really moves, measured with an SVD.

The figure of M4's LoRA section, and the check labs 4 and 5 lean on. "The correction has to pass
through eight dimensions" is easy to write and easy to believe without checking.
The check is one line of linear algebra: the singular values of
``delta W = (alpha / r) B A`` are the lengths of the directions that correction can push along, and
there are exactly ``r`` of them that are not zero. Next to the singular values of ``W`` itself --
seven hundred and sixty-eight of them, all non-zero -- the picture is the whole method.

Run from the repository root:

    uv run python labs/m4/lora_spectrum.py --adapter checkpoints/lora-e4-20260920-160133
"""

import argparse
import json
import sys
from pathlib import Path

import torch

from rukh.models.lora import load_adapter, lora_modules
from rukh.train import load_model

if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(encoding="utf-8", errors="replace")

parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--ckpt", default="checkpoints/medium-v4-20260919-174623/best.pt", type=Path)
parser.add_argument("--adapter", default="checkpoints/lora-e4-20260920-160133", type=Path)
parser.add_argument("--layer", default=0, type=int, help="Block to read the matrices from.")
parser.add_argument("--keep", default=24, type=int, help="Singular values to write out.")
parser.add_argument("--out", default="artifacts/web/lora-spectrum.json", type=Path)
args = parser.parse_args()

model = load_model(args.ckpt)[0].eval()
cfg = load_adapter(model, args.adapter / "adapter.safetensors")
adapters = dict(lora_modules(model))
name = f"blocks.{args.layer}.attn.qkv"
adapter = adapters[name]

# The query slice of the fused projection: the first `d_model` output rows.
start, stop = adapter.slices[0]
with torch.no_grad():
    base = adapter.base.weight[start:stop].float()
    delta = (cfg.scale * (adapter.b[0].float() @ adapter.a[0].float())).float()
    base_sv = torch.linalg.svdvals(base)
    delta_sv = torch.linalg.svdvals(delta)

nonzero = int((delta_sv > delta_sv.max() * 1e-6).sum())
print(f"{name} query slice: W is {tuple(base.shape)}, delta W is {tuple(delta.shape)}")
print(f"  r = {cfg.r}, alpha = {cfg.alpha}, scale = {cfg.scale}")
print(
    f"  non-zero singular values: W {int((base_sv > base_sv.max() * 1e-6).sum())}, "
    f"delta W {nonzero}"
)
print(f"  ||delta W||_F / ||W||_F = {torch.linalg.norm(delta) / torch.linalg.norm(base):.4f}")
print("  first singular values of delta W:", [round(float(v), 4) for v in delta_sv[: cfg.r + 2]])

payload = {
    "adapter": args.adapter.name,
    "module": name,
    "r": cfg.r,
    "alpha": cfg.alpha,
    "scale": cfg.scale,
    "shape": list(base.shape),
    "nonzero_delta": nonzero,
    "relative_norm": float(torch.linalg.norm(delta) / torch.linalg.norm(base)),
    "base": [float(v) for v in base_sv[: args.keep]],
    "delta": [float(v) for v in delta_sv[: args.keep]],
}
args.out.parent.mkdir(parents=True, exist_ok=True)
args.out.write_text(json.dumps(payload, indent=1) + "\n", encoding="utf-8", newline="\n")
print(f"\nwrote {args.out}")