#!/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)) # float32 to match the engine's precision: this is an optimizer-trajectory # parity over many steps, so we compare f32 training against an f32 reference # (a float64 reference would diverge purely from precision over the steps). t = torch.tensor(vals, dtype=torch.float32) 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", "q_norm", "k_norm", "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 (float32 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.float32) freq = THETA ** (-(2.0 * i) / HD) pos = torch.arange(SEQ, dtype=torch.float32).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.float32), 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) # Per-head QK-norm (Qwen3-style), before RoPE. q = rms_norm(q, P[f"l{l}_q_norm"]) k = rms_norm(k, P[f"l{l}_k_norm"]) 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.detach().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() # allclose-style: a per-element error is acceptable if it is within rtol *or* # atol (absolute). Weights span very small magnitudes, so a pure relative metric # is misleading on near-zero entries; this matches torch.allclose's semantics. def max_mismatch(a, b, rtol, atol): a, b = a.double(), b.double() err = (a - b).abs() tol = atol + rtol * b.abs() over = err - tol # > 0 only where it exceeds the combined tolerance return over.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 ATOL = 1e-3 worst_over, worst_name, worst_rel = 0.0, "", 0.0 fails = [] for n in NAMES: ref = read_vec(f"wN_{n}.txt") over = max_mismatch(P[n].detach(), ref, RTOL, ATOL) rel = relerr(P[n].detach(), ref) if over > worst_over: worst_over, worst_name, worst_rel = over, n, rel if over > 0.0: fails.append((n, rel, over)) print( f"final params: {len(NAMES)} checked, worst = {worst_name} " f"(relerr {worst_rel:.2e}, tol-overflow {worst_over:.2e}) " f"[rtol={RTOL}, atol={ATOL}]" ) if worst_loss > RTOL or fails: print("FAIL:") if worst_loss > RTOL: print(f" loss trajectory relerr {worst_loss:.3e} > {RTOL}") for n, rel, over in fails: print(f" param[{n}]: relerr={rel:.3e} tol-overflow={over:.3e}") sys.exit(1) print("ADAMW PARITY OK: loss trajectory + final params match torch.optim.AdamW (rtol/atol)")