"""Grokking on modular addition (a+b) mod p, p=113. 1-layer decoder-only transformer, d_model=128, 4 heads, d_mlp=512, ReLU, learned positional embeddings, NO LayerNorm. Full-batch AdamW, wd=1.0. """ import torch, numpy as np, time, csv, sys, math, argparse torch.set_num_threads(4) # default; may be overridden by --threads below (machine has a 4-CPU cgroup quota) torch.manual_seed(0) np.random.seed(0) import torch.nn as nn, torch.nn.functional as F p = argparse.ArgumentParser() p.add_argument("--max_steps", type=int, default=30000) p.add_argument("--log_every", type=int, default=50) p.add_argument("--time_budget", type=float, default=42*60, help="seconds") p.add_argument("--resume", action="store_true", help="load model.pt (and opt state from ckpt.pt if present), append to log.csv, continue step numbering") p.add_argument("--ckpt_every", type=int, default=1000) p.add_argument("--threads", type=int, default=4) p.add_argument("--emb_scale", type=float, default=1.0, help="multiply tok/pos embedding init by this") p.add_argument("--eps", type=float, default=1e-8) p.add_argument("--clamp", type=float, default=0.0, help="if >0: after every optimizer step rescale ALL params so total L2 norm == clamp") p.add_argument("--wd", type=float, default=1.0, help="AdamW weight decay") p.add_argument("--beta2", type=float, default=0.98, help="Adam beta2 (beta1 fixed at 0.9)") p.add_argument("--seed", type=int, default=0, help="torch seed for model init (data split stays fixed); 0 = original behaviour") args = p.parse_args() import os torch.set_num_threads(args.threads) if args.seed: torch.manual_seed(args.seed) print(f"threads={torch.get_num_threads()} args={vars(args)}", flush=True) P = 113 D_MODEL, N_HEADS, D_MLP = 128, 4, 512 VOCAB = P + 1 # tokens 0..P-1 and '=' = P SEQ = 3 LR, BETAS, WD = 1e-3, (0.9, args.beta2), args.wd FRAC_TRAIN = 0.3 # ---- data ---- rng = np.random.RandomState(0) pairs = np.array([(a, b) for a in range(P) for b in range(P)]) perm = rng.permutation(len(pairs)) n_train = int(FRAC_TRAIN * len(pairs)) train_idx, test_idx = perm[:n_train], perm[n_train:] if not args.resume: np.savez("split.npz", train_idx=train_idx, test_idx=test_idx) else: sp = np.load("split.npz"); assert (sp["train_idx"] == train_idx).all(), "split mismatch on resume" def make(idx): ab = pairs[idx] x = np.stack([ab[:, 0], ab[:, 1], np.full(len(idx), P)], axis=1) y = (ab[:, 0] + ab[:, 1]) % P return torch.tensor(x, dtype=torch.long), torch.tensor(y, dtype=torch.long) Xtr, Ytr = make(train_idx); Xte, Yte = make(test_idx) # ---- model ---- class Block(nn.Module): def __init__(s): super().__init__() s.qkv = nn.Linear(D_MODEL, 3 * D_MODEL, bias=False) s.o = nn.Linear(D_MODEL, D_MODEL, bias=False) s.fc1 = nn.Linear(D_MODEL, D_MLP) s.fc2 = nn.Linear(D_MLP, D_MODEL) s.register_buffer("mask", torch.tril(torch.ones(SEQ, SEQ)).bool()) def forward(s, x): B, T, C = x.shape q, k, v = s.qkv(x).split(D_MODEL, dim=2) hd = C // N_HEADS q = q.view(B, T, N_HEADS, hd).transpose(1, 2) k = k.view(B, T, N_HEADS, hd).transpose(1, 2) v = v.view(B, T, N_HEADS, hd).transpose(1, 2) att = (q @ k.transpose(-2, -1)) / math.sqrt(hd) att = att.masked_fill(~s.mask[:T, :T], float("-inf")).softmax(-1) h = (att @ v).transpose(1, 2).reshape(B, T, C) x = x + s.o(h) x = x + s.fc2(F.relu(s.fc1(x))) return x class Model(nn.Module): def __init__(s): super().__init__() s.tok = nn.Embedding(VOCAB, D_MODEL) s.pos = nn.Embedding(SEQ, D_MODEL) s.block = Block() s.unembed = nn.Linear(D_MODEL, P, bias=False) def forward(s, x): h = s.tok(x) + s.pos(torch.arange(x.shape[1])) h = s.block(h) return s.unembed(h[:, -1, :]) # logits at last position only model = Model() with torch.no_grad(): model.tok.weight.mul_(args.emb_scale); model.pos.weight.mul_(args.emb_scale) opt = torch.optim.AdamW(model.parameters(), lr=LR, betas=BETAS, weight_decay=WD, eps=args.eps) n_params = sum(q.numel() for q in model.parameters()) print(f"params={n_params} train={len(train_idx)} test={len(test_idx)}", flush=True) def weight_norm(): return math.sqrt(sum((q.detach() ** 2).sum().item() for q in model.parameters())) @torch.no_grad() def evaluate(X, Y): model.eval() logits = model(X) loss = F.cross_entropy(logits, Y).item() acc = (logits.argmax(-1) == Y).float().mean().item() model.train() return loss, acc start_step = 0 if args.resume: model.load_state_dict(torch.load("model.pt")) if os.path.exists("ckpt.pt"): ck = torch.load("ckpt.pt"); opt.load_state_dict(ck["opt"]); start_step = ck["step"] print(f"[resume] loaded model.pt + optimizer state from ckpt.pt at step {start_step}", flush=True) else: last = open("log.csv").read().strip().splitlines()[-1] start_step = int(last.split(",")[0]) print(f"[resume] loaded model.pt; NO optimizer state found -> fresh AdamW; continuing from step {start_step} (last logged)", flush=True) f = open("log.csv", "a", newline=""); w = csv.writer(f) else: f = open("log.csv", "w", newline="") w = csv.writer(f); w.writerow(["step", "train_loss", "test_loss", "train_acc", "test_acc", "weight_norm"]) t0 = time.time() PROBE = True pf = open("probe.csv", "w", newline=""); pw = csv.writer(pf); pw.writerow(["step","rms_grad","rms_mhat","rms_sqrt_vhat","frac_sqrt_vhat_lt_eps","rms_step_over_lr"]) consec = 0 step = start_step stop_reason = "max_steps" def save_ckpt(): torch.save(model.state_dict(), "model.pt") torch.save({"opt": opt.state_dict(), "step": step}, "ckpt.pt") for step in range(start_step + 1, start_step + args.max_steps + 1): logits = model(Xtr) loss = F.cross_entropy(logits, Ytr) opt.zero_grad(set_to_none=True) loss.backward() if PROBE and (step % args.log_every == 0 or step == 1): _n=sum(q.numel() for q in model.parameters()); _g=math.sqrt(sum((q.grad**2).sum().item() for q in model.parameters())/_n) _before=[q.detach().clone() for q in model.parameters()] opt.step() if PROBE and (step % args.log_every == 0 or step == 1): _d=math.sqrt(sum(((q.detach()-b)**2).sum().item() for q,b in zip(model.parameters(),_before))/_n)/LR _t=step; _mh=0.0; _vh=0.0; _small=0 for q in model.parameters(): st=opt.state[q]; _mh+=(st['exp_avg']/(1-BETAS[0]**_t)).pow(2).sum().item(); vh=(st['exp_avg_sq']/(1-BETAS[1]**_t)).sqrt(); _vh+=vh.pow(2).sum().item(); _small+=(vh 0: with torch.no_grad(): tot = torch.sqrt(sum((q ** 2).sum() for q in model.parameters())) for q in model.parameters(): q.mul_(args.clamp / tot) if step % args.ckpt_every == 0: save_ckpt() if step % args.log_every == 0 or step == 1: trl, tra = evaluate(Xtr, Ytr); tel, tea = evaluate(Xte, Yte) wn = weight_norm() w.writerow([step, f"{trl:.6f}", f"{tel:.6f}", f"{tra:.5f}", f"{tea:.5f}", f"{wn:.4f}"]); f.flush() if step == start_step + 100: print(f"[timing] 100 steps in {time.time()-t0:.1f}s -> est {args.max_steps} steps = {(time.time()-t0)/100*args.max_steps/60:.1f} min", flush=True) if step % 500 == 0 or step == 1: print(f"step {step:6d} train_loss {trl:.4f} test_loss {tel:.4f} train_acc {tra:.4f} test_acc {tea:.4f} |w| {wn:.2f} t={time.time()-t0:.0f}s", flush=True) if step % args.log_every == 0: consec = consec + 1 if tea >= 0.99 else 0 if consec >= 200: stop_reason = "test_acc>=99% for 200 logged points"; break if time.time() - t0 > args.time_budget: stop_reason = "time_budget"; break f.close() save_ckpt() print(f"done: step={step} reason={stop_reason} wall={time.time()-t0:.1f}s", flush=True)