test: AdamW PyTorch parity + checkpoint round-trip + real training
Acceptance tests (GPU-gated not(no_cuda), run on dash5): - adamw_parity_dump.rs + adamw_parity.py: build the tiny model with fixed init, run N AdamW steps on a fixed batch, dump the loss trajectory + final params; the Python side rebuilds the identical model and runs torch.optim.AdamW with matched lr/wd/betas/eps, comparing trajectory + final params within rtol. - checkpoint_roundtrip.rs: train a few steps, save, load into a fresh model with a DIFFERENT init, assert identical logits/loss on a fixed input. - real_training.rs (#[ignore], --release): train on TinyStories for a bounded budget; assert loss drops substantially and print greedy samples. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
180
crates/xtrain-train/tests/adamw_parity.py
Normal file
180
crates/xtrain-train/tests/adamw_parity.py
Normal file
@@ -0,0 +1,180 @@
|
||||
#!/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")
|
||||
Reference in New Issue
Block a user