xtrain/crates/xtrain-train/tests/adamw_parity.py

#!/usr/bin/env python3
"""AdamW-vs-PyTorch parity (Phase T6).

Loads the model dumped by tests/adamw_parity_dump.rs (config, ids, initial
params, the loss trajectory, and final params), rebuilds the IDENTICAL tiny
transformer in PyTorch from the same initial weights, and runs the SAME number
of `torch.optim.AdamW` steps with matched hyperparameters (lr, weight_decay,
betas, eps) on the same fixed batch. It then compares:

  * the per-step loss trajectory (Rust AdamW vs torch AdamW), and
  * the final parameters,

within a relative tolerance. A correct hand-written AdamW (bias correction +
decoupled weight decay) tracks torch's optimizer step-for-step.

Usage:  python3 adamw_parity.py /tmp/xtrain_adamw
"""
import sys
import os
import math
import torch

DIR = sys.argv[1] if len(sys.argv) > 1 else "/tmp/xtrain_adamw"


def read_vec(name):
    path = os.path.join(DIR, name)
    shape = None
    vals = []
    with open(path) as f:
        for line in f:
            line = line.strip()
            if line.startswith("# shape"):
                shape = [int(x) for x in line.split()[2].split(",") if x]
            elif line:
                vals.append(float(line))
    t = torch.tensor(vals, dtype=torch.float64)
    if shape:
        t = t.reshape(shape)
    return t


def read_cfg():
    cfg = {}
    with open(os.path.join(DIR, "config.txt")) as f:
        for line in f:
            k, v = line.split()
            cfg[k] = v
    return cfg


def read_ids(name):
    with open(os.path.join(DIR, name)) as f:
        return [int(x) for x in f.read().split()]


cfg = read_cfg()
DIM = int(cfg["dim"])
NL = int(cfg["n_layers"])
NH = int(cfg["n_heads"])
HD = int(cfg["head_dim"])
EPS = float(cfg["eps"])
THETA = float(cfg["rope_theta"])
LR = float(cfg["lr"])
WD = float(cfg["wd"])
N_STEPS = int(cfg["n_steps"])

ids = read_ids("ids.txt")
targets = read_ids("targets.txt")
SEQ = len(ids)

NAMES = ["embed"]
for l in range(NL):
    for p in ["attn_norm", "wq", "wk", "wv", "wo",
              "ffn_norm", "w_gate", "w_up", "w_down"]:
        NAMES.append(f"l{l}_{p}")
NAMES += ["final_norm", "lm_head"]

# Load the IDENTICAL initial weights as leaf params (float64 reference).
P = {n: read_vec(f"w0_{n}.txt").clone().requires_grad_(True) for n in NAMES}


def rms_norm(x, gamma):
    ms = x.pow(2).mean(dim=-1, keepdim=True)
    return x * torch.rsqrt(ms + EPS) * gamma


def rope(x):  # x: [seq, nh, hd], position = token index
    half = HD // 2
    out = torch.empty_like(x)
    i = torch.arange(half, dtype=torch.float64)
    freq = THETA ** (-(2.0 * i) / HD)
    pos = torch.arange(SEQ, dtype=torch.float64).reshape(SEQ, 1)
    ang = pos * freq
    c = torch.cos(ang).reshape(SEQ, 1, half)
    s = torch.sin(ang).reshape(SEQ, 1, half)
    x0, x1 = x[..., :half], x[..., half:]
    out[..., :half] = x0 * c - x1 * s
    out[..., half:] = x1 * c + x0 * s
    return out


idx = torch.tensor(ids, dtype=torch.long)
tgt = torch.tensor(targets, dtype=torch.long)
mask = torch.triu(torch.full((SEQ, SEQ), -1.0e9, dtype=torch.float64), diagonal=1)


def forward():
    h = P["embed"][idx]
    for l in range(NL):
        x = rms_norm(h, P[f"l{l}_attn_norm"])
        q = (x @ P[f"l{l}_wq"]).reshape(SEQ, NH, HD)
        k = (x @ P[f"l{l}_wk"]).reshape(SEQ, NH, HD)
        v = (x @ P[f"l{l}_wv"]).reshape(SEQ, NH, HD)
        q = rope(q).transpose(0, 1)
        k = rope(k).transpose(0, 1)
        v = v.transpose(0, 1)
        scale = 1.0 / math.sqrt(HD)
        scores = (q @ k.transpose(-1, -2)) * scale + mask
        probs = torch.softmax(scores, dim=-1)
        out = (probs @ v).transpose(0, 1).reshape(SEQ, DIM)
        h = h + out @ P[f"l{l}_wo"]
        x = rms_norm(h, P[f"l{l}_ffn_norm"])
        act = torch.nn.functional.silu(x @ P[f"l{l}_w_gate"]) * (x @ P[f"l{l}_w_up"])
        h = h + act @ P[f"l{l}_w_down"]
    h = rms_norm(h, P["final_norm"])
    return h @ P["lm_head"]


# Match the Rust optimizer: torch.optim.AdamW with the same lr/wd/betas/eps.
opt = torch.optim.AdamW(list(P.values()), lr=LR, betas=(0.9, 0.999),
                        eps=1e-8, weight_decay=WD)

torch_losses = []
for _ in range(N_STEPS):
    opt.zero_grad()
    logits = forward()
    loss = torch.nn.functional.cross_entropy(logits, tgt, reduction="mean")
    torch_losses.append(loss.item())
    loss.backward()
    opt.step()


def relerr(a, b):
    a, b = a.double(), b.double()
    denom = b.abs().clamp(min=1e-6)
    return ((a - b).abs() / denom).max().item()


rust_losses = read_vec("losses.txt")
print("step  rust_loss     torch_loss    relerr")
worst_loss = 0.0
for i in range(N_STEPS):
    rl, tl = rust_losses[i].item(), torch_losses[i]
    e = abs(rl - tl) / max(abs(tl), 1e-6)
    worst_loss = max(worst_loss, e)
    if i < 5 or i == N_STEPS - 1:
        print(f"{i:4d}  {rl:.6e}  {tl:.6e}  {e:.2e}")
print(f"loss trajectory: worst relerr = {worst_loss:.2e}")

RTOL = 2e-2
worst_p, worst_name = 0.0, ""
fails = []
for n in NAMES:
    ref = read_vec(f"wN_{n}.txt")
    e = relerr(P[n].detach(), ref)
    if e > worst_p:
        worst_p, worst_name = e, n
    if e > RTOL:
        fails.append((n, e))
print(f"final params: {len(NAMES)} checked, worst = {worst_name} @ {worst_p:.2e} (rtol={RTOL})")

if worst_loss > RTOL or fails:
    print("FAIL:")
    if worst_loss > RTOL:
        print(f"  loss trajectory relerr {worst_loss:.3e} > {RTOL}")
    for n, e in fails:
        print(f"  param[{n}]: relerr={e:.3e}")
    sys.exit(1)
print("ADAMW PARITY OK: loss trajectory + final params match torch.optim.AdamW within rtol")