diff --git a/code/notebooks/0101-autoregressive-likelihood-sft.ipynb b/code/notebooks/0101-autoregressive-likelihood-sft.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..d92657420f56e7cb8a1fafe5aa9ffeea97471c29 --- /dev/null +++ b/code/notebooks/0101-autoregressive-likelihood-sft.ipynb @@ -0,0 +1,309 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# 0101 - Autoregressive LM likelihood, token vs sequence view, SFT = MLE = forward KL\n", + "Companion to `notes/0101 - Autoregressive Likelihood and SFT.pdf`.\n", + "Everything is fp64. The model is a tiny autoregressive policy whose whole sequence space (121 sequences) is enumerated,\n", + "so each identity is checked EXACTLY (error ~1e-16), not by sampling. Pattern: naive loop first, vectorised second, assert they agree." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "source": [ + "import sys, math, itertools\n", + "sys.path.insert(0, \"..\")\n", + "from _m1_lib import *\n", + "lm = TinyLM(V=4, L=4, seed=0, scale=1.0)\n", + "theta_s = lm.theta0.clone() # student\n", + "theta_T = TinyLM(V=4, L=4, seed=7, scale=1.5).theta0 # teacher p (same family -> realizable)\n", + "print(\"sequences:\", lm.N, \" params:\", theta_s.numel())" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "sequences: 121 params: 36\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 1. The model is a proper distribution on sequences: sum_o pi(o) = 1 (note, Sec. 3.1)\n", + "Naive: multiply conditional probabilities one token at a time with a python loop. Efficient: gather from one batched log_softmax." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "source": [ + "p_naive = torch.stack([lm.seq_logprob_naive(theta_s, s).exp() for s in lm.seqs])\n", + "p_vec = lm.probs(theta_s)\n", + "check(\"naive == vectorised sequence probability\", maxabs(p_naive, p_vec), 1e-15)\n", + "check(\"sum_o pi(o) = 1 (normalisation, truncated at L)\", abs(float(p_vec.sum()) - 1.0), 1e-14)\n", + "# per-state conditional distributions each sum to 1 (this is what the induction in the note uses)\n", + "worst = max(abs(float(lm.cond_probs(theta_s, s).sum()) - 1) for s in lm.prefixes())\n", + "check(\"every conditional pi(.|o_= 0 (Gibbs): min(0, KL)\", abs(min(0.0, float(KLf))), 0.0)\n", + "# reverse KL differs (what RL optimises is the *reverse* direction, see 0109)\n", + "KLr = (p_vec * (lp_s - p_T.log())).sum()\n", + "print(f\"forward KL(p||pi)={float(KLf):.6f} reverse KL(pi||p)={float(KLr):.6f} (different objects)\")" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "CHECK|exact|CE = H + KL(p||pi)|0.000e+00|tol=1.0e-14|PASS\n", + "CHECK|exact|KL(p||pi) >= 0 (Gibbs): min(0, KL)|0.000e+00|tol=0.0e+00|PASS\n", + "forward KL(p||pi)=1.221839 reverse KL(pi||p)=1.094181 (different objects)\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 3. Chain rule of KL: sequence KL = sum over positions of the *expected conditional* KL (token-level view)" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "source": [ + "def prefix_prob(theta, prefix):\n", + " pr = 1.0\n", + " for k, a in enumerate(prefix):\n", + " pr = pr * lm.cond_probs(theta, prefix[:k])[a]\n", + " return pr\n", + "tok_KL = 0.0\n", + "for s in lm.prefixes():\n", + " w = prefix_prob(theta_T, s) # P_p(o_ reweights tokens by 1/|o|)\")\n", + "assert rel > 1e-3" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "CHECK|exact|token-sum gradient == MLE gradient|9.714e-17|tol=1.0e-14|PASS\n", + "CHECK|exact|constant-denominator loss == MLE gradient / Lmax|0.000e+00|tol=1.0e-14|PASS\n", + "seq-mean gradient vs MLE gradient: relative deviation 0.444 (NOT a pure rescale -> reweights tokens by 1/|o|)\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 6. Sampling check (MC): empirical-risk gradient converges to the population one" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "source": [ + "gen = torch.Generator().manual_seed(123)\n", + "n = 400_000\n", + "idx = torch.multinomial(p_T, n, replacement=True, generator=gen)\n", + "G_emp = -S_auto[idx].mean(0)\n", + "se = (S_auto[idx].var(0) / n).sqrt()\n", + "ok_ = se > 0\n", + "z = float(((G_emp - grad_CE_hand).abs()[ok_] / se[ok_]).max())\n", + "assert float((G_emp - grad_CE_hand).abs()[~ok_].max()) < 1e-14 # unused parameters have identically zero gradient\n", + "check(\"MC SFT gradient within 4 SE (max z over components)\", z, 4.0, kind=\"mc\")" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "CHECK|mc|MC SFT gradient within 4 SE (max z over components)|1.761e+00|tol=4.0e+00|PASS\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 7. SFT = maximum likelihood = forward KL minimisation: L-BFGS from a random init recovers the teacher" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": {}, + "source": [ + "gen_i = torch.Generator().manual_seed(77) # seeded: the quoted numbers must reproduce\n", + "th = (theta_s + 0.5 * torch.randn(theta_s.numel(), generator=gen_i)).requires_grad_(True)\n", + "opt = torch.optim.LBFGS([th], lr=1.0, max_iter=500, tolerance_grad=1e-14, tolerance_change=1e-16, line_search_fn=\"strong_wolfe\")\n", + "def closure():\n", + " opt.zero_grad(); loss = -(p_T * lm.seq_logprob(th)).sum(); loss.backward(); return loss\n", + "for _ in range(6): opt.step(closure)\n", + "KLend = float((p_T * (p_T.log() - lm.seq_logprob(th.detach()))).sum())\n", + "check(\"KL(p||pi) after MLE (should be ~0)\", KLend, 1e-10)\n", + "check(\"recovered sequence distribution == teacher\", maxabs(lm.probs(th.detach()), p_T), 1e-6)\n", + "print(\"done; checks:\", len(CHECKS))" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "CHECK|exact|KL(p||pi) after MLE (should be ~0)|4.753e-17|tol=1.0e-10|PASS\n", + "CHECK|exact|recovered sequence distribution == teacher|2.203e-09|tol=1.0e-06|PASS\n", + "done; checks: 15\n" + ] + } + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.12.11" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} \ No newline at end of file diff --git a/code/notebooks/0102-policy-gradient-baselines.ipynb b/code/notebooks/0102-policy-gradient-baselines.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..30d2d8279a9cf01329e5d17110d6271a5dc824a4 --- /dev/null +++ b/code/notebooks/0102-policy-gradient-baselines.ipynb @@ -0,0 +1,347 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# 0102 - Policy gradient from first principles, REINFORCE, baselines\n", + "Companion to `notes/0102 - Policy Gradient and Baselines.pdf`. Exact enumeration of the 121-sequence toy LM; fp64." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "source": [ + "import sys, itertools\n", + "sys.path.insert(0, \"..\")\n", + "from _m1_lib import *\n", + "lm = TinyLM(V=4, L=4, seed=0, scale=1.0)\n", + "theta = lm.theta0.clone()\n", + "R = default_reward(lm) # binary verifier reward\n", + "p = lm.probs(theta)\n", + "S = score_matrix(lm, theta) # [N, P] grad log pi(o)\n", + "P_ = theta.numel()\n", + "print(\"N =\", lm.N, \" P =\", P_, \" E[R] =\", float(p @ R))" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "N = 121 P = 36 E[R] = 0.16592597945219617\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 1. Log-derivative trick: grad J = E_pi[R * score] (note Eq. 2)\n", + "J(theta) = sum_o pi_theta(o) R(o). Three independent computations of its gradient." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "source": [ + "th = theta.clone().requires_grad_(True)\n", + "J = (lm.probs(th) * R).sum()\n", + "g_auto = torch.autograd.grad(J, th)[0] # differentiate the *definition* (through pi)\n", + "g_ld = (p[:, None] * R[:, None] * S).sum(0) # E_pi[R * score] (the derived formula)\n", + "def Jfun(t): return float((lm.probs(t) * R).sum())\n", + "h = 1e-6\n", + "g_fd = torch.zeros_like(theta)\n", + "for i in range(P_): # naive central finite differences\n", + " e = torch.zeros_like(theta); e[i] = h\n", + " g_fd[i] = (Jfun(theta + e) - Jfun(theta - e)) / (2 * h)\n", + "check(\"log-derivative formula == autograd of J\", maxabs(g_ld, g_auto), 1e-14)\n", + "check(\"log-derivative formula == central finite differences\", maxabs(g_ld, g_fd), 1e-8)" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "CHECK|exact|log-derivative formula == autograd of J|1.388e-17|tol=1.0e-14|PASS\n", + "CHECK|exact|log-derivative formula == central finite differences|3.363e-11|tol=1.0e-08|PASS\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 2. Baselines are free: E[b * score] = 0 for any b that does not depend on the sampled response" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "source": [ + "check(\"E_pi[score] = 0\", float((p[:, None] * S).sum(0).abs().max()), 1e-14)\n", + "for b in [0.0, float(p @ R), 5.0, -3.0]:\n", + " g_b = (p[:, None] * (R - b)[:, None] * S).sum(0)\n", + " check(f\"unbiased with constant baseline b={b:+.3f}\", maxabs(g_b, g_auto), 1e-13)\n", + "# state-dependent baseline b_t(o_ 1e-3\n", + "# and the extreme case b = R(o) (the sample's own reward) kills the gradient completely\n", + "g_self = (p[:, None] * (R - R)[:, None] * S).sum(0)\n", + "print(\"self-baseline gradient norm:\", float(g_self.norm()), \" true grad norm:\", float(g_auto.norm()))" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "CHECK|exact|E_pi[score] = 0|5.551e-17|tol=1.0e-14|PASS\n", + "CHECK|exact|unbiased with constant baseline b=+0.000|1.388e-17|tol=1.0e-13|PASS\n", + "CHECK|exact|unbiased with constant baseline b=+0.166|2.776e-17|tol=1.0e-13|PASS\n", + "CHECK|exact|unbiased with constant baseline b=+5.000|2.078e-16|tol=1.0e-13|PASS\n", + "CHECK|exact|unbiased with constant baseline b=-3.000|2.137e-16|tol=1.0e-13|PASS\n", + "CHECK|exact|unbiased with arbitrary state-dependent baseline b(o_ successes carry the larger score norm, so b* > E[R]\")\n", + "assert (m1_ > m0_) == (bstar > EvR)\n", + "# interval claim: any b in (0, 2b*) has Var(b) < Var(0); b outside is worse\n", + "inside = all(var_formula(float(b)) < var_formula(0.0) for b in torch.linspace(0.01, 2 * bstar - 0.01, 50))\n", + "outside = all(var_formula(float(b)) > var_formula(0.0) for b in [2 * bstar + 0.01, 2 * bstar + 1.0, -0.01, -1.0])\n", + "assert inside and outside\n", + "check(\"Var(b) < Var(0) iff b in (0, 2 b*) (50 inside, 4 outside points)\", 0.0, 0.0)" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "CHECK|exact|variance formula at b=0.0|0.000e+00|tol=1.0e-13|PASS\n", + "CHECK|exact|variance formula at b=0.3|1.110e-16|tol=1.0e-13|PASS\n", + "CHECK|exact|variance formula at b=1.0|2.220e-16|tol=1.0e-13|PASS\n", + "CHECK|exact|variance formula at b=2.0|8.882e-16|tol=1.0e-13|PASS\n", + "CHECK|exact|b* = E[R||s||^2]/E||s||^2 minimises variance (grid min vs V(b*))|0.000e+00|tol=1.0e-12|PASS\n", + "CHECK|exact|minimiser on grid is within grid step of b*|1.020e-04|tol=2.0e-03|PASS\n", + "b*=0.3019 E[R]=0.1659 Var(0)=0.6898 Var(E[R])=0.5158 Var(b*)=0.4715\n", + "CHECK|exact|b* - E[R] == Cov_pi(R, ||s||^2) / E||s||^2 (exact identity)|0.000e+00|tol=1.0e-14|PASS\n", + "E[||s||^2 | R=1] = 4.358 E[||s||^2 | R=0] = 2.005 -> successes carry the larger score norm, so b* > E[R]\n", + "CHECK|exact|Var(b) < Var(0) iff b in (0, 2 b*) (50 inside, 4 outside points)|0.000e+00|tol=0.0e+00|PASS\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 4. Naive vs efficient estimator (Monte-Carlo), and agreement with the exact gradient\n", + "naive: one backward per sample, python loop. efficient: ONE backward through the surrogate -mean(sg(A) * log pi)." + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "source": [ + "n = 2000\n", + "gen = torch.Generator().manual_seed(11)\n", + "idx = torch.multinomial(p, n, replacement=True, generator=gen)\n", + "b = EvR\n", + "# naive\n", + "th = theta.clone().requires_grad_(True)\n", + "acc = torch.zeros_like(theta)\n", + "for i in idx.tolist():\n", + " lp = lm.seq_logprob(th)[i]\n", + " acc += (float(R[i]) - b) * torch.autograd.grad(lp, th)[0]\n", + "g_naive = acc / n\n", + "# efficient\n", + "th = theta.clone().requires_grad_(True)\n", + "surrogate = ((R[idx] - b).detach() * lm.seq_logprob(th)[idx]).mean() # gradient of this = REINFORCE estimate\n", + "g_eff = torch.autograd.grad(surrogate, th)[0]\n", + "check(\"naive per-sample loop == single backward through surrogate\", maxabs(g_naive, g_eff), 1e-13)\n", + "# MC accuracy against the exact gradient\n", + "n2 = 400_000\n", + "idx2 = torch.multinomial(p, n2, replacement=True, generator=gen)\n", + "est = ((R[idx2] - b)[:, None] * S[idx2])\n", + "g_mc, se = est.mean(0), est.std(0) / n2 ** 0.5\n", + "ok_ = se > 0\n", + "z = float(((g_mc - g_auto).abs()[ok_] / se[ok_]).max())\n", + "check(\"MC REINFORCE gradient within 4 SE of exact (max z)\", z, 4.0, kind=\"mc\")" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "CHECK|exact|naive per-sample loop == single backward through surrogate|8.604e-16|tol=1.0e-13|PASS\n", + "CHECK|mc|MC REINFORCE gradient within 4 SE of exact (max z)|2.001e+00|tol=4.0e+00|PASS\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 5. Group baselines: leave-one-out is unbiased, \"mean including self\" is biased by (G-1)/G\n", + "Exact enumeration of ALL G-tuples of samples (G=2: 14641 tuples, G=3: 1.77M tuples, evaluated per first sample)." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "source": [ + "N = lm.N\n", + "RR = R[None, :] ; pp = p[None, :]\n", + "def loo_and_mean(G):\n", + " \"\"\"exact E[(R_1 - baseline_1) s_1] for baseline = mean of others (LOO) and mean of all (incl. self).\"\"\"\n", + " # sum over samples 2..G of rewards: build the joint over (o_2..o_G) as a tensor of (reward sum, prob)\n", + " if G == 2:\n", + " rs = R[None, :].expand(N, N); w = p[None, :].expand(N, N)\n", + " sum_others_R = rs; prob_others = w\n", + " elif G == 3:\n", + " sum_others_R = (R[:, None] + R[None, :]).reshape(-1)[None, :].expand(N, -1)\n", + " prob_others = (p[:, None] * p[None, :]).reshape(-1)[None, :].expand(N, -1)\n", + " # N x N^2 : first sample index x joint of other two\n", + " Ro1 = R[:, None]\n", + " base_loo = sum_others_R / (G - 1)\n", + " base_all = (Ro1 + sum_others_R) / G\n", + " wloo = (prob_others * (Ro1 - base_loo)).sum(1) # E over others of (R1 - b)\n", + " wall = (prob_others * (Ro1 - base_all)).sum(1)\n", + " g_loo = (p * wloo)[:, None] * S\n", + " g_all = (p * wall)[:, None] * S\n", + " return g_loo.sum(0), g_all.sum(0)\n", + "for G in (2, 3):\n", + " gl, ga = loo_and_mean(G)\n", + " check(f\"G={G}: leave-one-out baseline is exactly unbiased\", maxabs(gl, g_auto), 1e-13)\n", + " check(f\"G={G}: mean-including-self gives exactly (G-1)/G * grad J\", maxabs(ga, (G - 1) / G * g_auto), 1e-13)" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "CHECK|exact|G=2: leave-one-out baseline is exactly unbiased|2.082e-17|tol=1.0e-13|PASS\n", + "CHECK|exact|G=2: mean-including-self gives exactly (G-1)/G * grad J|1.041e-17|tol=1.0e-13|PASS\n", + "CHECK|exact|G=3: leave-one-out baseline is exactly unbiased|2.082e-17|tol=1.0e-13|PASS\n", + "CHECK|exact|G=3: mean-including-self gives exactly (G-1)/G * grad J|6.939e-18|tol=1.0e-13|PASS\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 6. Causality / reward-to-go (per-token rewards): both unbiased, reward-to-go has lower variance here\n", + "Per-token reward r_t = 1[o_t == 1]. Full return G = sum_t r_t; reward-to-go G_t = sum_{t'>=t} r_t'." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "source": [ + "rt = (lm.tokens == 1).double() * lm.mask # [N, L]\n", + "Gfull = rt.sum(1)\n", + "Gto = rt.flip(1).cumsum(1).flip(1) * lm.mask # reward-to-go\n", + "th = theta.clone().requires_grad_(True)\n", + "gJ = torch.autograd.grad((lm.probs(th) * Gfull).sum(), th)[0]\n", + "est_full = (T.sum(1) * Gfull[:, None]) # sum_t score_t * G\n", + "est_rtg = (T * Gto[:, :, None]).sum(1) # sum_t score_t * G_t\n", + "g_full = (p[:, None] * est_full).sum(0); g_rtg = (p[:, None] * est_rtg).sum(0)\n", + "check(\"full-return estimator unbiased\", maxabs(g_full, gJ), 1e-13)\n", + "check(\"reward-to-go estimator unbiased (dropped terms E[score_t r_t'] , t' long sequences\n", + "thq = lmL.theta0.clone(); thp = perturb(thq, 0.6, seed=9)\n", + "lq, lp = lmL.token_logprobs(thq), lmL.token_logprobs(thp)\n", + "qL = lq.sum(1).exp()\n", + "check(\"long-horizon model: sum_o q(o) = 1\", abs(float(qL.sum()) - 1), 1e-12)\n", + "print(\"N sequences:\", lmL.N, \" mean length under q:\", float((qL * lmL.lengths).sum()))\n", + "logw_cum = (lp - lq).cumsum(1) # log w_T for T = 1..L (masked tokens contribute 0 -> ratio 1 after EOS)\n", + "rows = []\n", + "for T in range(1, lmL.L + 1):\n", + " wT = logw_cum[:, T - 1].exp()\n", + " m1 = float((qL * wT).sum()); m2 = float((qL * wT ** 2).sum())\n", + " med = float((qL * (wT < 1).double()).sum()) # P_q(w_T < 1): the typical sample has weight < 1\n", + " rows.append((T, m1, m2 - 1, med, float((qL * logw_cum[:, T - 1]).sum())))\n", + " check(f\"E_q[w_T] = 1 at T={T}\", abs(m1 - 1), 1e-12) if T in (1, 6, 12) else None\n", + "print(\" T E[w] chi2_T P(w<1) E[log w]\")\n", + "for r in rows: print(f\"{r[0]:2d} {r[1]:.4f} {r[2]:9.4f} {r[3]:.4f} {r[4]:+.4f}\")\n", + "chi_T = torch.tensor([r[2] for r in rows])\n", + "assert bool((chi_T[1:] >= chi_T[:-1] - 1e-12).all())\n", + "check(\"chi^2_T non-decreasing in T (E_q[w_T^2] only grows)\", 0.0, 0.0)\n", + "ratio_growth = float(chi_T[-1] / chi_T[0])\n", + "print(f\"chi2_12 / chi2_1 = {ratio_growth:.1f}\")\n", + "# chain rule: E_q[log w_T] = - sum_{t<=T} E_q[KL_t(q || p)] computed independently on the prefix tree\n", + "def tree_kl(thq, thp, T):\n", + " tot = 0.0\n", + " def rec(prefix, prob):\n", + " nonlocal tot\n", + " if len(prefix) >= T: return\n", + " cq, cp = lmL.cond_probs(thq, prefix), lmL.cond_probs(thp, prefix)\n", + " tot += prob * float((cq * (cq.log() - cp.log())).sum())\n", + " for a in range(lmL.V - 1):\n", + " rec(prefix + (a,), prob * float(cq[a]))\n", + " rec((), 1.0)\n", + " return tot\n", + "for T in (3, 8):\n", + " check(f\"E_q[log w_T] = -sum_t E[KL_t] (tree recursion), T={T}\", abs(rows[T - 1][4] + tree_kl(thq, thp, T)), 1e-12)" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "CHECK|exact|long-horizon model: sum_o q(o) = 1|2.220e-16|tol=1.0e-12|PASS\n", + "N sequences: 8191 mean length under q: 8.597261658540575\n", + "CHECK|exact|E_q[w_T] = 1 at T=1|2.220e-16|tol=1.0e-12|PASS\n", + "CHECK|exact|E_q[w_T] = 1 at T=6|0.000e+00|tol=1.0e-12|PASS\n", + "CHECK|exact|E_q[w_T] = 1 at T=12|0.000e+00|tol=1.0e-12|PASS\n", + " T E[w] chi2_T P(w<1) E[log w]\n", + " 1 1.0000 0.0040 0.8060 -0.0019\n", + " 2 1.0000 0.0042 0.8044 -0.0020\n", + " 3 1.0000 0.0333 0.2960 -0.0189\n", + " 4 1.0000 0.0373 0.4582 -0.0207\n", + " 5 1.0000 0.0446 0.4510 -0.0244\n", + " 6 1.0000 0.0457 0.4176 -0.0249\n", + " 7 1.0000 0.0487 0.4547 -0.0263\n", + " 8 1.0000 0.0500 0.4533 -0.0269\n", + " 9 1.0000 0.0511 0.4731 -0.0274\n", + "10 1.0000 0.0534 0.4954 -0.0285\n", + "11 1.0000 0.0535 0.4942 -0.0285\n", + "12 1.0000 0.0537 0.4936 -0.0286\n", + "CHECK|exact|chi^2_T non-decreasing in T (E_q[w_T^2] only grows)|0.000e+00|tol=0.0e+00|PASS\n", + "chi2_12 / chi2_1 = 13.3\n", + "CHECK|exact|E_q[log w_T] = -sum_t E[KL_t] (tree recursion), T=3|4.857e-17|tol=1.0e-12|PASS\n", + "CHECK|exact|E_q[log w_T] = -sum_t E[KL_t] (tree recursion), T=8|1.874e-16|tol=1.0e-12|PASS\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 4. iid-token closed form E_q[w_T^2] = (1 + chi2_1)^T and ESS collapse" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "source": [ + "Vv = 3\n", + "qa = torch.tensor([0.5, 0.3, 0.2]); pa = torch.tensor([0.45, 0.33, 0.22])\n", + "chi1 = float((pa ** 2 / qa).sum()) - 1\n", + "def brute(T):\n", + " seqs = torch.cartesian_prod(*[torch.arange(Vv)] * T) if T > 1 else torch.arange(Vv)[:, None]\n", + " lqs = qa[seqs].log().sum(1); lps = pa[seqs].log().sum(1)\n", + " return float((lqs.exp() * (lps - lqs).exp() ** 2).sum())\n", + "for T in (1, 3, 6, 9):\n", + " check(f\"iid tokens: E_q[w_T^2] = (1+chi2_1)^T, T={T}\", abs(brute(T) - (1 + chi1) ** T) / (1 + chi1) ** T, 1e-12)\n", + "print(f\"per-token chi2_1 = {chi1:.5f}\")\n", + "for T in (1, 10, 100, 1000, 4000):\n", + " print(f\"T={T:5d} chi2_T = {(1+chi1)**T-1:12.3f} ESS/n ~ 1/(1+chi2) = {(1+chi1)**(-T):.3e}\")" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "CHECK|exact|iid tokens: E_q[w_T^2] = (1+chi2_1)^T, T=1|0.000e+00|tol=1.0e-12|PASS\n", + "CHECK|exact|iid tokens: E_q[w_T^2] = (1+chi2_1)^T, T=3|4.310e-16|tol=1.0e-12|PASS\n", + "CHECK|exact|iid tokens: E_q[w_T^2] = (1+chi2_1)^T, T=6|0.000e+00|tol=1.0e-12|PASS\n", + "CHECK|exact|iid tokens: E_q[w_T^2] = (1+chi2_1)^T, T=9|4.060e-16|tol=1.0e-12|PASS\n", + "per-token chi2_1 = 0.01000\n", + "T= 1 chi2_T = 0.010 ESS/n ~ 1/(1+chi2) = 9.901e-01\n", + "T= 10 chi2_T = 0.105 ESS/n ~ 1/(1+chi2) = 9.053e-01\n", + "T= 100 chi2_T = 1.705 ESS/n ~ 1/(1+chi2) = 3.697e-01\n", + "T= 1000 chi2_T = 20958.156 ESS/n ~ 1/(1+chi2) = 4.771e-05\n", + "T= 4000 chi2_T = 192972369947321888.000 ESS/n ~ 1/(1+chi2) = 5.182e-18\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 5. Naive vs stable weight computation (fp32 running product underflows; log-space is exact)" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "source": [ + "gen = torch.Generator().manual_seed(77)\n", + "def weights(T, sigma, dtype):\n", + " logr = (-0.5 * sigma ** 2 + sigma * torch.randn(T, generator=gen)).to(torch.float64)\n", + " naive = torch.tensor(1.0, dtype=dtype)\n", + " for x in logr.exp().to(dtype): naive = naive * x # naive running product of token ratios\n", + " stable = logr.sum().exp() # log-space\n", + " return float(naive), float(stable), float(logr.sum())\n", + "n64, s64, _ = weights(50, 0.1, torch.float64)\n", + "check(\"T=50: naive fp64 product == exp(sum log r)\", abs(n64 - s64) / s64, 1e-12)\n", + "n32, s64b, lw = weights(4000, 0.5, torch.float32)\n", + "print(f\"T=4000 sigma=0.5: log w = {lw:.1f}; naive fp32 product = {n32}; log-space fp64 = {s64b:.3e}\")\n", + "assert n32 == 0.0 and math.isfinite(lw)" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "CHECK|exact|T=50: naive fp64 product == exp(sum log r)|0.000e+00|tol=1.0e-12|PASS\n", + "T=4000 sigma=0.5: log w = -521.5; naive fp32 product = 0.0; log-space fp64 = 3.410e-227\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 6. Token-level surrogate (PPO-style ratio) vs the exact off-policy gradient\n", + "Exact: grad J(theta) = E_old[ w(o) (R-b) sum_t grad log pi_theta(o_t|s_t) ] (sequence-level weight, all tokens).\n", + "Token-level surrogate: L(theta) = E_old[ sum_t r_t(theta) (R-b) ] with r_t = pi_theta(o_t|s_t)/pi_old(o_t|s_t).\n", + "They agree at theta_old and drift apart as ||theta - theta_old|| grows: value error O(delta^2), gradient error O(delta)." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "source": [ + "lm = TinyLM(V=4, L=4, seed=0, scale=1.0)\n", + "R = default_reward(lm)\n", + "th_old = lm.theta0.clone()\n", + "q = lm.probs(th_old)\n", + "b = float(q @ R)\n", + "def exact_grad(th):\n", + " t = th.clone().requires_grad_(True)\n", + " return torch.autograd.grad((lm.probs(t) * R).sum(), t)[0]\n", + "def seq_is_grad(th):\n", + " t = th.clone().requires_grad_(True)\n", + " lpn = lm.seq_logprob(t)\n", + " wts = (lpn - lm.seq_logprob(th_old)).exp().detach()\n", + " return torch.autograd.grad((q * wts * (R - b) * lpn).sum(), t)[0]\n", + "def token_surrogate(th):\n", + " t = th.clone().requires_grad_(True)\n", + " r = (lm.token_logprobs(t) - lm.token_logprobs(th_old)).exp() * lm.mask # r_t(theta), 0 where masked\n", + " L = (q[:, None] * r * (R - b)[:, None]).sum()\n", + " return L, torch.autograd.grad(L, t)[0]\n", + "deltas = [0.0, 0.02, 0.04, 0.08, 0.16, 0.32]\n", + "L0 = float((q * lm.lengths * (R - b)).sum())\n", + "gerr, verr = [], []\n", + "for d in deltas:\n", + " th = perturb(th_old, d, seed=5) if d > 0 else th_old.clone()\n", + " ge = exact_grad(th)\n", + " check(f\"sequence-level IS gradient == exact grad J, delta={d}\", maxabs(seq_is_grad(th), ge), 1e-14) if d in (0.0, 0.32) else None\n", + " Lv, gt = token_surrogate(th)\n", + " gerr.append(float((gt - ge).norm())); verr.append(abs(float(Lv.detach()) - L0 - float((lm.probs(th) * R).sum() - (q * R).sum())))\n", + " print(f\"delta={d:5.2f} |grad_tok - grad_J| = {gerr[-1]:.3e} |dL - dJ| = {verr[-1]:.3e}\")\n", + "check(\"token surrogate gradient == exact gradient at theta_old\", gerr[0], 1e-14)\n", + "check(\"token surrogate value == J to first order at theta_old (error at delta=0)\", verr[0], 1e-14)\n", + "import numpy as np\n", + "sl_g = np.polyfit(np.log(deltas[1:]), np.log(gerr[1:]), 1)[0]\n", + "sl_v = np.polyfit(np.log(deltas[1:]), np.log(verr[1:]), 1)[0]\n", + "print(f\"log-log slopes: gradient error {sl_g:.2f} (expect ~1), value error {sl_v:.2f} (expect ~2)\")\n", + "check(\"gradient-error slope ~ 1\", abs(sl_g - 1.0), 0.15)\n", + "check(\"value-error slope ~ 2\", abs(sl_v - 2.0), 0.15)\n", + "print(\"done; checks:\", len(CHECKS))" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "CHECK|exact|sequence-level IS gradient == exact grad J, delta=0.0|4.163e-17|tol=1.0e-14|PASS\n", + "delta= 0.00 |grad_tok - grad_J| = 6.150e-17 |dL - dJ| = 2.776e-17\n", + "delta= 0.02 |grad_tok - grad_J| = 2.905e-04 |dL - dJ| = 3.869e-07\n", + "delta= 0.04 |grad_tok - grad_J| = 5.814e-04 |dL - dJ| = 1.550e-06\n", + "delta= 0.08 |grad_tok - grad_J| = 1.164e-03 |dL - dJ| = 6.214e-06\n", + "delta= 0.16 |grad_tok - grad_J| = 2.333e-03 |dL - dJ| = 2.497e-05\n", + "CHECK|exact|sequence-level IS gradient == exact grad J, delta=0.32|2.776e-17|tol=1.0e-14|PASS\n", + "delta= 0.32 |grad_tok - grad_J| = 4.684e-03 |dL - dJ| = 1.007e-04\n", + "CHECK|exact|token surrogate gradient == exact gradient at theta_old|6.150e-17|tol=1.0e-14|PASS\n", + "CHECK|exact|token surrogate value == J to first order at theta_old (error at delta=0)|2.776e-17|tol=1.0e-14|PASS\n", + "log-log slopes: gradient error 1.00 (expect ~1), value error 2.01 (expect ~2)\n", + "CHECK|exact|gradient-error slope ~ 1|2.652e-03|tol=1.5e-01|PASS\n", + "CHECK|exact|value-error slope ~ 2|5.861e-03|tol=1.5e-01|PASS\n", + "done; checks: 23\n" + ] + } + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.12.11" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} \ No newline at end of file diff --git a/code/notebooks/0104-trust-region-ppo-clip.ipynb b/code/notebooks/0104-trust-region-ppo-clip.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..0dc1c3f80fae4f8ced9076b11de08ae152c0602b --- /dev/null +++ b/code/notebooks/0104-trust-region-ppo-clip.ipynb @@ -0,0 +1,318 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# 0104 - Trust regions -> PPO clipped surrogate; the clip as a pessimistic bound; the gradient case table\n", + "Companion to `notes/0104 - Trust Regions and the PPO Clip.pdf`. Exact enumeration, fp64." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "source": [ + "import sys, math, itertools\n", + "sys.path.insert(0, \"..\")\n", + "from _m1_lib import *\n", + "import numpy as np\n", + "lm = TinyLM(V=4, L=4, seed=0, scale=1.0)\n", + "R = default_reward(lm)\n", + "th_old = lm.theta0.clone()\n", + "q = lm.probs(th_old)\n", + "seqset = {s: i for i, s in enumerate(lm.seqs)}\n", + "V, L = lm.V, lm.L" + ], + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 1. Exact value functions by backward recursion; advantage A_old(s,a) = V(s+a) - V(s)" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "source": [ + "from functools import lru_cache\n", + "def make_value(theta):\n", + " @lru_cache(None)\n", + " def Vf(prefix):\n", + " if prefix in seqset: return float(R[seqset[prefix]]) # terminal (EOS or length L): value = reward\n", + " cp = lm.cond_probs(theta, prefix)\n", + " return float(sum(cp[a] * Vf(prefix + (a,)) for a in range(V)))\n", + " return Vf\n", + "Vold = make_value(th_old)\n", + "check(\"V_old(root) = J(theta_old) = sum_o pi R\", abs(Vold(()) - float(q @ R)), 1e-14)\n", + "states = lm.prefixes()\n", + "def A_old(prefix, a): return Vold(prefix + (a,)) - Vold(prefix)\n", + "check(\"E_{a~pi_old}[A_old(s,a)] = 0 at every state\", max(abs(sum(float(lm.cond_probs(th_old, s)[a]) * A_old(s, a) for a in range(V))) for s in states), 1e-14)\n", + "# advantage along every enumerated sequence: Adv[o,t] = A_old(s_t, o_t)\n", + "Adv = torch.zeros(lm.N, L)\n", + "for i, s in enumerate(lm.seqs):\n", + " for t in range(len(s)): Adv[i, t] = A_old(s[:t], s[t])\n", + "check(\"telescoping: sum_t A_old(s_t,o_t) = R(o) - J(theta_old) for EVERY sequence\", maxabs(Adv.sum(1), R - Vold(())), 1e-14)" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "CHECK|exact|V_old(root) = J(theta_old) = sum_o pi R|0.000e+00|tol=1.0e-14|PASS\n", + "CHECK|exact|E_{a~pi_old}[A_old(s,a)] = 0 at every state|1.475e-16|tol=1.0e-14|PASS\n", + "CHECK|exact|telescoping: sum_t A_old(s_t,o_t) = R(o) - J(theta_old) for EVERY sequence|1.110e-16|tol=1.0e-14|PASS\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 2. Performance-difference lemma: J(theta') - J(theta_old) = E_{o~pi'}[ sum_t A_old(s_t,o_t) ]\n", + "and the exact IS form with the PREFIX ratio w_{1:t} = prod_{s<=t} r_s (action at t is the only thing that matters for A_old)." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "source": [ + "def ratios(th):\n", + " lp, lq = lm.token_logprobs(th), lm.token_logprobs(th_old)\n", + " r = (lp - lq).exp() * lm.mask + (1 - lm.mask) # r_t (1 after EOS)\n", + " return r, lp\n", + "for d in (0.1, 0.4):\n", + " th = perturb(th_old, d, seed=3)\n", + " p = lm.probs(th)\n", + " dJ = float(p @ R - q @ R)\n", + " pd_lemma = float((p * Adv.sum(1)).sum()) # E_{pi'}[sum_t A_old]\n", + " r, _ = ratios(th)\n", + " wprefix = r.cumprod(1) # w_{1:t}\n", + " is_exact = float((q[:, None] * wprefix * Adv).sum()) # E_old[ sum_t w_{1:t} A_old(s_t,o_t) ]\n", + " L_sur = float((q[:, None] * r * Adv).sum()) # TRPO/PPO surrogate: only the token ratio\n", + " check(f\"performance difference lemma, delta={d}\", abs(dJ - pd_lemma), 1e-14)\n", + " check(f\"IS with prefix ratio w_(1:t) is exact, delta={d}\", abs(dJ - is_exact), 1e-14)\n", + " print(f\" delta={d}: dJ={dJ:+.6f} surrogate L-J_old={L_sur:+.6f} (error {abs(dJ-L_sur):.2e})\")\n", + "# first-order agreement of surrogate and J at theta_old\n", + "th = th_old.clone().requires_grad_(True)\n", + "lpn, lqo = lm.token_logprobs(th), lm.token_logprobs(th_old)\n", + "Ls = (q[:, None] * ((lpn - lqo).exp() * lm.mask) * Adv).sum()\n", + "gL = torch.autograd.grad(Ls, th)[0]\n", + "th2 = th_old.clone().requires_grad_(True)\n", + "gJ = torch.autograd.grad((lm.probs(th2) * R).sum(), th2)[0]\n", + "check(\"grad of surrogate at theta_old == grad J (first-order matching)\", maxabs(gL, gJ), 1e-14)" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "CHECK|exact|performance difference lemma, delta=0.1|4.857e-17|tol=1.0e-14|PASS\n", + "CHECK|exact|IS with prefix ratio w_(1:t) is exact, delta=0.1|4.380e-17|tol=1.0e-14|PASS\n", + " delta=0.1: dJ=+0.003183 surrogate L-J_old=+0.003188 (error 5.85e-06)\n", + "CHECK|exact|performance difference lemma, delta=0.4|2.776e-17|tol=1.0e-14|PASS\n", + "CHECK|exact|IS with prefix ratio w_(1:t) is exact, delta=0.4|6.939e-18|tol=1.0e-14|PASS\n", + " delta=0.4: dJ=+0.013143 surrogate L-J_old=+0.013234 (error 9.14e-05)\n", + "CHECK|exact|grad of surrogate at theta_old == grad J (first-order matching)|2.082e-17|tol=1.0e-14|PASS\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 3. Finite-horizon trust-region bound (derived in the note): |J' - L| <= 2 eps alpha^2 T (T-1)\n", + "alpha = max_s TV(pi'(.|s), pi(.|s)), eps = max_{s,a} |A_old(s,a)|, T = horizon. Pinsker: alpha^2 <= D_KL^max / 2." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "source": [ + "eps = max(abs(A_old(s, a)) for s in states for a in range(V))\n", + "print(\"eps = max|A| =\", eps)\n", + "rows = []\n", + "for d in (0.05, 0.1, 0.2, 0.4, 0.8, 1.6):\n", + " th = perturb(th_old, d, seed=3)\n", + " alpha = max(0.5 * float((lm.cond_probs(th, s) - lm.cond_probs(th_old, s)).abs().sum()) for s in states)\n", + " kl_max = max(float((lm.cond_probs(th_old, s) * (lm.cond_probs(th_old, s).log() - lm.cond_probs(th, s).log())).sum()) for s in states)\n", + " p = lm.probs(th); r, _ = ratios(th)\n", + " err = abs(float(p @ R - q @ R) - float((q[:, None] * r * Adv).sum()))\n", + " bound_tv = 2 * eps * alpha ** 2 * L * (L - 1)\n", + " bound_kl = eps * L * (L - 1) * kl_max\n", + " rows.append((d, alpha, kl_max, err, bound_tv, bound_kl))\n", + " print(f\"delta={d:4.2f} alpha={alpha:.4f} KLmax={kl_max:.4f} |J'-L|={err:.3e} TV-bound={bound_tv:.3e} KL-bound={bound_kl:.3e} slack={bound_tv/max(err,1e-300):.1f}x\")\n", + " assert err <= bound_tv + 1e-15 and bound_tv <= bound_kl + 1e-15 and alpha ** 2 <= kl_max / 2 + 1e-15\n", + "check(\"bound |J'-L| <= 2 eps a^2 T(T-1) holds at all 6 trust-region sizes (max err/bound)\", max(r[3] / r[4] for r in rows), 1.0)\n", + "check(\"Pinsker chain a^2 <= KLmax/2 holds (max a^2/(KLmax/2))\", max(r[1] ** 2 / (r[2] / 2) for r in rows), 1.0)" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "eps = max|A| = 0.9030066306956115\n", + "delta=0.05 alpha=0.0048 KLmax=0.0001 |J'-L|=1.471e-06 TV-bound=5.034e-04 KL-bound=6.045e-04 slack=342.3x\n", + "delta=0.10 alpha=0.0097 KLmax=0.0002 |J'-L|=5.854e-06 TV-bound=2.028e-03 KL-bound=2.425e-03 slack=346.5x\n", + "delta=0.20 alpha=0.0195 KLmax=0.0009 |J'-L|=2.320e-05 TV-bound=8.226e-03 KL-bound=9.755e-03 slack=354.6x\n", + "delta=0.40 alpha=0.0395 KLmax=0.0036 |J'-L|=9.139e-05 TV-bound=3.379e-02 KL-bound=3.945e-02 slack=369.7x\n", + "delta=0.80 alpha=0.0808 KLmax=0.0149 |J'-L|=3.606e-04 TV-bound=1.415e-01 KL-bound=1.611e-01 slack=392.4x\n", + "delta=1.60 alpha=0.1668 KLmax=0.0615 |J'-L|=1.532e-03 TV-bound=6.027e-01 KL-bound=6.664e-01 slack=393.4x\n", + "CHECK|exact|bound |J'-L| <= 2 eps a^2 T(T-1) holds at all 6 trust-region sizes (max err/bound)|2.922e-03|tol=1.0e+00|PASS\n", + "CHECK|exact|Pinsker chain a^2 <= KLmax/2 holds (max a^2/(KLmax/2))|9.043e-01|tol=1.0e+00|PASS\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 4. PPO clip: f(r, A) = min(r A, clip(r, 1-eps_l, 1+eps_h) A)\n", + "closed form A>0: A*min(r, 1+eps_h); A<0: A*max(r, 1-eps_l); pessimistic: f <= r A; gradient case table." + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "source": [ + "def f_naive(r, A, el, eh):\n", + " clipped = min(max(r, 1 - el), 1 + eh)\n", + " return min(r * A, clipped * A)\n", + "def f_closed(r, A, el, eh):\n", + " return A * min(r, 1 + eh) if A > 0 else A * max(r, 1 - el)\n", + "el, eh = 0.2, 0.28\n", + "rg = np.concatenate([np.linspace(0.0, 3.0, 301), [0.8 + 1e-9, 1.28 - 1e-9]])\n", + "worst_cf, worst_pess = 0.0, -1e9\n", + "for A in (-2.0, -0.3, 0.3, 2.0):\n", + " for r in rg:\n", + " worst_cf = max(worst_cf, abs(f_naive(r, A, el, eh) - f_closed(r, A, el, eh)))\n", + " worst_pess = max(worst_pess, f_naive(r, A, el, eh) - r * A)\n", + "check(\"closed forms A*min(r,1+eh) / A*max(r,1-el) == min/clip definition\", worst_cf, 1e-14)\n", + "check(\"pessimism: f(r,A) <= r*A everywhere (max of f - rA, must be <= 0)\", max(0.0, worst_pess), 0.0)\n", + "# autograd gradient wrt log pi vs hand case table\n", + "def grad_autograd(r, A, el, eh):\n", + " lp = torch.tensor(math.log(r), requires_grad=True) # log pi_theta ; old log prob = 0 so r = exp(lp)\n", + " rr = lp.exp()\n", + " f = torch.minimum(rr * A, torch.clamp(rr, 1 - el, 1 + eh) * A)\n", + " return float(torch.autograd.grad(f, lp)[0])\n", + "def grad_table(r, A, el, eh):\n", + " if A > 0: passes = r < 1 + eh\n", + " else: passes = r > 1 - el\n", + " return A * r if passes else 0.0 # d f / d log pi = A r when the token is live\n", + "worst, n_zero, n_pass = 0.0, 0, 0\n", + "for A in (-2.0, -0.3, 0.3, 2.0):\n", + " for r in [x for x in rg if x > 0 and abs(x - (1 - el)) > 1e-6 and abs(x - (1 + eh)) > 1e-6]:\n", + " ga, gt = grad_autograd(r, A, el, eh), grad_table(r, A, el, eh)\n", + " worst = max(worst, abs(ga - gt)); n_zero += gt == 0.0; n_pass += gt != 0.0\n", + "check(f\"gradient case table == autograd on {n_zero+n_pass} (r,A) points ({n_zero} silenced, {n_pass} live)\", worst, 1e-14)\n", + "# the 2x3 table printed explicitly\n", + "print(\"\\n gradient wrt log pi of one token (eps_l=0.2, eps_h=0.28), '.' = zero\")\n", + "print(\" region A>0 A<0\")\n", + "for name, r in ((\"r < 1-eps_l\", 0.5), (\"in [1-eps_l,1+eps_h]\", 1.0), (\"r > 1+eps_h\", 2.0)):\n", + " print(f\" {name:<20s}\", \" \".join(f\"{(grad_autograd(r, A, el, eh)/(A*r)):4.0f}*A*r\" if abs(grad_autograd(r, A, el, eh))>0 else \" . \" for A in (0.7, -0.7)))\n", + "# boundary tie behaviour in torch (documented, not relied upon)\n", + "print(\" boundary tie r=1+eps_h, A>0:\", grad_autograd(1 + eh, 1.0, el, eh), \" r=1-eps_l, A<0:\", grad_autograd(1 - el, -1.0, el, eh))\n", + "# naive vs vectorised implementation\n", + "gen = torch.Generator().manual_seed(1)\n", + "rr, AA = torch.rand(10000, generator=gen) * 3, torch.randn(10000, generator=gen)\n", + "fn = torch.tensor([f_naive(float(a), float(b), el, eh) for a, b in zip(rr, AA)])\n", + "fv = torch.where(AA > 0, AA * torch.clamp(rr, max=1 + eh), AA * torch.clamp(rr, min=1 - el))\n", + "check(\"naive scalar loop == vectorised torch.where implementation\", maxabs(fn, fv), 1e-14)" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "CHECK|exact|closed forms A*min(r,1+eh) / A*max(r,1-el) == min/clip definition|0.000e+00|tol=1.0e-14|PASS\n", + "CHECK|exact|pessimism: f(r,A) <= r*A everywhere (max of f - rA, must be <= 0)|0.000e+00|tol=0.0e+00|PASS\n", + "CHECK|exact|gradient case table == autograd on 1192 (r,A) points (502 silenced, 690 live)|8.882e-16|tol=1.0e-14|PASS\n", + "\n", + " gradient wrt log pi of one token (eps_l=0.2, eps_h=0.28), '.' = zero\n", + " region A>0 A<0\n", + " r < 1-eps_l 1*A*r . \n", + " in [1-eps_l,1+eps_h] 1*A*r 1*A*r\n", + " r > 1+eps_h . 1*A*r\n", + " boundary tie r=1+eps_h, A>0: 1.28 r=1-eps_l, A<0: -0.8\n", + "CHECK|exact|naive scalar loop == vectorised torch.where implementation|0.000e+00|tol=1.0e-14|PASS\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 5. The clip is a heuristic, not a trust region: ratios leave [1-eps,1+eps] during multi-step ascent\n", + "Full-batch (exact-expectation) ascent on E_old[ sum_t f(r_t, A_t) ] with A = R - E[R] (outcome advantage), 600 steps." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "source": [ + "R_ = R; b = float(q @ R_)\n", + "A_out = (R_ - b)[:, None].expand(-1, L) * lm.mask\n", + "th = th_old.clone().requires_grad_(True)\n", + "opt = torch.optim.SGD([th], lr=0.3)\n", + "lq = lm.token_logprobs(th_old)\n", + "hist = []\n", + "for step in range(601):\n", + " lp = lm.token_logprobs(th)\n", + " r = ((lp - lq).exp() * lm.mask) + (1 - lm.mask)\n", + " f = torch.where(A_out > 0, A_out * torch.clamp(r, max=1 + 0.2), A_out * torch.clamp(r, min=1 - 0.2))\n", + " obj = (q[:, None] * f * lm.mask).sum()\n", + " opt.zero_grad(); (-obj).backward(); opt.step()\n", + " if step in (0, 10, 50, 100, 200, 400, 600):\n", + " out_of = float((q[:, None] * (((r - 1).abs() > 0.2) & (lm.mask > 0))).sum() / (q * lm.lengths).sum())\n", + " hist.append((step, float((r - 1).abs().detach().max()), out_of, float(obj.detach())))\n", + "print(\"step max|r-1| prob-weighted token fraction outside [0.8,1.2] objective\")\n", + "for h in hist: print(f\"{h[0]:4d} {h[1]:8.3f} {h[2]:8.3f} {h[3]:+.4f}\")\n", + "assert hist[-1][1] > 0.4\n", + "print(\"final max|r-1| =\", hist[-1][1], \"(eps=0.2)\")\n", + "print(\"done; checks:\", len(CHECKS))" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "step max|r-1| prob-weighted token fraction outside [0.8,1.2] objective\n", + " 0 0.000 0.000 +0.2230\n", + " 10 0.252 0.113 +0.2838\n", + " 50 0.452 0.475 +0.2965\n", + " 100 0.558 0.530 +0.2998\n", + " 200 0.654 0.580 +0.3016\n", + " 400 0.751 0.753 +0.3021\n", + " 600 0.812 0.748 +0.3031\n", + "final max|r-1| = 0.8121891873787255 (eps=0.2)\n", + "done; checks: 14\n" + ] + } + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.12.11" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} \ No newline at end of file diff --git a/code/notebooks/0105-grpo-group-advantages.ipynb b/code/notebooks/0105-grpo-group-advantages.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..b8ef1431e7210021cf7c0201b39ea7168fd6e1b4 --- /dev/null +++ b/code/notebooks/0105-grpo-group-advantages.ipynb @@ -0,0 +1,403 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# 0105 - GRPO: group-relative advantages as a Monte-Carlo baseline; std and length normalisation biases; Dr.GRPO\n", + "Companion to `notes/0105 - GRPO and Dr GRPO.pdf`. Exact enumeration of ALL G-tuples of samples (fp64)." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "source": [ + "import sys, math\n", + "sys.path.insert(0, \"..\")\n", + "from _m1_lib import *\n", + "import numpy as np\n", + "lm = TinyLM(V=4, L=4, seed=0, scale=1.0)\n", + "theta = lm.theta0.clone()\n", + "R = default_reward(lm)\n", + "p = lm.probs(theta)\n", + "S = score_matrix(lm, theta) # [N,P]\n", + "N = lm.N\n", + "th = theta.clone().requires_grad_(True)\n", + "pr = (lm.probs(th) * R).sum() # success probability p(theta)\n", + "grad_p = torch.autograd.grad(pr, th)[0]\n", + "p_succ = float(pr.detach())\n", + "print(\"success prob p =\", p_succ, \" ||grad p|| =\", float(grad_p.norm()))" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "success prob p = 0.16592597945219614 ||grad p|| = 0.18246238486098645\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 1. Advantage definitions, RLOO equivalence, invariances" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "source": [ + "def adv_grpo(Rg, ddof=0, eps=0.0):\n", + " mu = Rg.mean(); sd = Rg.std(unbiased=bool(ddof)) if len(Rg) > 1 else torch.tensor(0.0)\n", + " return (Rg - mu) / (sd + eps)\n", + "def adv_center(Rg): return Rg - Rg.mean()\n", + "def adv_rloo(Rg):\n", + " G = len(Rg); return Rg - (Rg.sum() - Rg) / (G - 1)\n", + "gen = torch.Generator().manual_seed(0)\n", + "Rg = torch.rand(7, generator=gen)\n", + "G_ = len(Rg)\n", + "check(\"G/(G-1) * (R_i - mean) == RLOO advantage\", maxabs(G_ / (G_ - 1) * adv_center(Rg), adv_rloo(Rg)), 1e-15)\n", + "check(\"GRPO advantage invariant to R -> a R + c (a>0)\", maxabs(adv_grpo(Rg), adv_grpo(3.7 * Rg - 2.0)), 1e-13)\n", + "check(\"centred advantage scales with a: A(aR+c) = a A(R)\", maxabs(adv_center(3.7 * Rg - 2.0), 3.7 * adv_center(Rg)), 1e-14)\n", + "check(\"GRPO advantage has mean 0 and (population) std 1\", max(abs(float(adv_grpo(Rg).mean())), abs(float(adv_grpo(Rg).std(unbiased=False)) - 1)), 1e-14)\n", + "Rdead = torch.ones(5)\n", + "check(\"dead group (all rewards equal) has zero advantage (eps=1e-4 convention)\", float(adv_grpo(Rdead, eps=1e-4).abs().max()), 0.0)" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "CHECK|exact|G/(G-1) * (R_i - mean) == RLOO advantage|1.110e-16|tol=1.0e-15|PASS\n", + "CHECK|exact|GRPO advantage invariant to R -> a R + c (a>0)|3.331e-16|tol=1.0e-13|PASS\n", + "CHECK|exact|centred advantage scales with a: A(aR+c) = a A(R)|3.331e-16|tol=1.0e-14|PASS\n", + "CHECK|exact|GRPO advantage has mean 0 and (population) std 1|1.269e-16|tol=1.0e-14|PASS\n", + "CHECK|exact|dead group (all rewards equal) has zero advantage (eps=1e-4 convention)|0.000e+00|tol=0.0e+00|PASS\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 2. EXACT expected gradient of the group estimators, by enumerating every G-tuple of responses\n", + " estimator: g = (1/G) sum_i A_i * score(o_i) (outcome reward, one scalar advantage per response)\n", + "compare: mean-centred (Dr.GRPO) -> (1 - 1/G) grad J ; std-normalised (GRPO) -> c_G(p) grad p with c_G NOT constant." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "source": [ + "def exact_group_gradient(G, kind, Rr=None):\n", + " \"\"\"Sum over all N^G tuples with product probability. kind in {'center','grpo','rloo'}.\"\"\"\n", + " Rr = R if Rr is None else Rr\n", + " assert G in (2, 3)\n", + " if G == 2:\n", + " R1, R2 = Rr[:, None].expand(N, N), Rr[None, :].expand(N, N)\n", + " P = p[:, None] * p[None, :]\n", + " A = _adv_batch(torch.stack([R1, R2]), kind) # [2, N, N]\n", + " g = (P * A[0]).sum(1) @ S + (P * A[1]).sum(0) @ S # sample 1 carries score(o1), sample 2 score(o2)\n", + " return g / 2\n", + " else:\n", + " g = torch.zeros(S.shape[1])\n", + " R2, R3 = Rr[:, None], Rr[None, :]\n", + " P23 = p[:, None] * p[None, :]\n", + " for i in range(N):\n", + " Rs = torch.stack([Rr[i].expand(N, N), R2.expand(N, N), R3.expand(N, N)])\n", + " A = _adv_batch(Rs, kind) # [3, N, N]\n", + " g += p[i] * (P23 * A[0]).sum() * S[i]\n", + " g += p[i] * ((P23 * A[1]).sum(1) @ S)\n", + " g += p[i] * ((P23 * A[2]).sum(0) @ S)\n", + " return g / 3\n", + "def _adv_batch(Rs, kind):\n", + " \"\"\"Rs [G, ...] rewards of the group, vectorised over the trailing dims.\"\"\"\n", + " G = Rs.shape[0]\n", + " mu = Rs.mean(0, keepdim=True)\n", + " if kind == 'center': return Rs - mu\n", + " if kind == 'rloo': return Rs - (Rs.sum(0, keepdim=True) - Rs) / (G - 1)\n", + " sd = ((Rs - mu) ** 2).mean(0, keepdim=True).sqrt() # population std\n", + " safe = torch.where(sd > 0, sd, torch.ones_like(sd))\n", + " return torch.where(sd > 0, (Rs - mu) / safe, torch.zeros_like(Rs))\n", + "for G in (2, 3):\n", + " gc = exact_group_gradient(G, 'center')\n", + " gr = exact_group_gradient(G, 'rloo')\n", + " check(f\"G={G}: mean-centred estimator == (1-1/G) grad J exactly (Dr.GRPO unbiasedness)\", maxabs(gc, (1 - 1 / G) * grad_p), 1e-13)\n", + " check(f\"G={G}: RLOO estimator == grad J exactly\", maxabs(gr, grad_p), 1e-13)" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "CHECK|exact|G=2: mean-centred estimator == (1-1/G) grad J exactly (Dr.GRPO unbiasedness)|1.388e-17|tol=1.0e-13|PASS\n", + "CHECK|exact|G=2: RLOO estimator == grad J exactly|2.776e-17|tol=1.0e-13|PASS\n", + "CHECK|exact|G=3: mean-centred estimator == (1-1/G) grad J exactly (Dr.GRPO unbiasedness)|2.776e-17|tol=1.0e-13|PASS\n", + "CHECK|exact|G=3: RLOO estimator == grad J exactly|4.857e-17|tol=1.0e-13|PASS\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 3. Std normalisation: for binary rewards the expected gradient is c_G(p) * grad p with an explicit c_G\n", + "k successes among G (0 1/sqrt(p(1-p)) as G grows (G=256, relative error)\", abs(c_G(256, p_succ) * math.sqrt(p_succ * (1 - p_succ)) - 1), 0.1)\n", + "check(\"c_G(p->0) = sqrt(G-1) (bounded up-weighting of hard prompts), G=16\", abs(c_G(16, 1e-9) - math.sqrt(15)), 1e-6)\n", + "print(\" difficulty weighting c_G(p) (G=16):\")\n", + "for pp in (0.01, 0.05, 0.2, 0.5, 0.8, 0.95, 0.99):\n", + " print(f\" p={pp:4.2f} GRPO c = {c_G(16, pp):6.3f} Dr.GRPO c = {1 - 1/16:.3f}\")\n", + "# the infinite-G limit is the gradient of the variance-stabilising objective 2*arcsin(sqrt(p))\n", + "th = theta.clone().requires_grad_(True)\n", + "pp_ = (lm.probs(th) * R).sum()\n", + "g_asin = torch.autograd.grad(2 * torch.asin(pp_.sqrt()), th)[0]\n", + "check(\"grad[2 asin(sqrt p)] == grad p / sqrt(p(1-p))\", maxabs(g_asin, grad_p / math.sqrt(p_succ * (1 - p_succ))), 1e-14)" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "CHECK|exact|G=2: brute-force GRPO(std-norm) gradient == c_G(p) * grad p|2.776e-17|tol=1.0e-13|PASS\n", + "CHECK|exact|G=3: brute-force GRPO(std-norm) gradient == c_G(p) * grad p|5.551e-17|tol=1.0e-13|PASS\n", + "CHECK|exact|G=4: enumerated (N^3 other samples) GRPO gradient == c_4(p) grad p|8.327e-17|tol=1.0e-13|PASS\n", + "\n", + " c_G(p) at p = 0.166 for various G, vs 1/sqrt(p(1-p)) = 2.6881 and Dr.GRPO's (1-1/G):\n", + " G= 2 c_G = 1.0000\n", + " G= 4 c_G = 1.6678\n", + " G= 8 c_G = 2.1725\n", + " G= 16 c_G = 2.4801\n", + " G= 64 c_G = 2.6487\n", + " G= 256 c_G = 2.6785\n", + "CHECK|exact|c_G(p) -> 1/sqrt(p(1-p)) as G grows (G=256, relative error)|3.559e-03|tol=1.0e-01|PASS\n", + "CHECK|exact|c_G(p->0) = sqrt(G-1) (bounded up-weighting of hard prompts), G=16|1.454e-08|tol=1.0e-06|PASS\n", + " difficulty weighting c_G(p) (G=16):\n", + " p=0.01 GRPO c = 3.733 Dr.GRPO c = 0.938\n", + " p=0.05 GRPO c = 3.267 Dr.GRPO c = 0.938\n", + " p=0.20 GRPO c = 2.346 Dr.GRPO c = 0.938\n", + " p=0.50 GRPO c = 1.934 Dr.GRPO c = 0.938\n", + " p=0.80 GRPO c = 2.346 Dr.GRPO c = 0.938\n", + " p=0.95 GRPO c = 3.267 Dr.GRPO c = 0.938\n", + " p=0.99 GRPO c = 3.733 Dr.GRPO c = 0.938\n", + "CHECK|exact|grad[2 asin(sqrt p)] == grad p / sqrt(p(1-p))|2.776e-17|tol=1.0e-14|PASS\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 3b. General principle: ANY group advantage A_1 = h(R_1; R_2..R_G) has expected gradient E[ f(R) score ] with f(r) = E_others[ h(r; R_2..R_G) ]\n", + "i.e. it is exactly the policy gradient of a *reshaped reward* f(R). Check with a graded (3-valued) reward, G=3, std-normalised." + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "source": [ + "Rgr = R + (lm.lengths == 4).double() # rewards in {0,1,2}\n", + "g_graded = exact_group_gradient(3, 'grpo', Rr=Rgr)\n", + "vals = sorted(set(Rgr.tolist()))\n", + "f_of = {}\n", + "for r1 in vals: # f(r) = E over the other two samples of A_1(r; R_2, R_3)\n", + " R2, R3 = Rgr[:, None].expand(N, N), Rgr[None, :].expand(N, N)\n", + " P23 = p[:, None] * p[None, :]\n", + " A1 = _adv_batch(torch.stack([torch.full((N, N), r1), R2, R3]), 'grpo')[0]\n", + " f_of[r1] = float((P23 * A1).sum())\n", + "f_vec = torch.tensor([f_of[float(x)] for x in Rgr])\n", + "check(\"graded reward, G=3: E[GRPO gradient] == E[f(R) score] with f(r)=E[A_1|R_1=r]\", maxabs(g_graded, (p[:, None] * f_vec[:, None] * S).sum(0)), 1e-13)\n", + "print(\" reshaped reward f(r):\", {k: round(v, 4) for k, v in f_of.items()}, \"(true rewards:\", vals, \")\")" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "CHECK|exact|graded reward, G=3: E[GRPO gradient] == E[f(R) score] with f(r)=E[A_1|R_1=r]|5.551e-17|tol=1.0e-13|PASS\n", + " reshaped reward f(r): {0.0: -0.3371, 1.0: 0.9871, 2.0: 1.2765} (true rewards: [0.0, 1.0, 2.0] )\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 4. Length normalisation: E[(R-b)/|o| * score] = grad E[(R-b)/|o|] (b held fixed) -- a different objective" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "source": [ + "b = float(p @ R)\n", + "lens = lm.lengths\n", + "g_unb = (p[:, None] * (R - b)[:, None] * S).sum(0) # = grad J\n", + "g_len = (p[:, None] * ((R - b) / lens)[:, None] * S).sum(0) # GRPO-style 1/|o|\n", + "th = theta.clone().requires_grad_(True)\n", + "Jlen = (lm.probs(th) * (R - b) / lens).sum() # E[(R-b)/|o|] with b constant\n", + "g_Jlen = torch.autograd.grad(Jlen, th)[0]\n", + "check(\"unbiased baseline estimator == grad J\", maxabs(g_unb, grad_p), 1e-14)\n", + "check(\"1/|o|-normalised estimator == grad of E[(R-b)/|o|] (b fixed)\", maxabs(g_len, g_Jlen), 1e-14)\n", + "check(\"constant normaliser 1/Lmax gives grad J / Lmax (Dr.GRPO)\", maxabs((p[:, None] * (R - b)[:, None] * S).sum(0) / lm.L * lm.L, grad_p), 1e-14)\n", + "# effective per-sequence weight and which way the update pushes the expected length\n", + "Elen = (lm.probs(th) * lens).sum()\n", + "g_len_E = torch.autograd.grad(Elen, th)[0] # grad E[|o|]\n", + "print(\"\\n length | P(|o|) mean weight on FAILURES (unbiased | 1/|o|) mean weight on SUCCESSES (unbiased | 1/|o|)\")\n", + "for Lk in range(1, lm.L + 1):\n", + " m = lens == Lk\n", + " pm = float(p[m].sum())\n", + " fm, sm = m & (R == 0), m & (R == 1)\n", + " wf = (-b, -b / Lk) ; ws = (1 - b, (1 - b) / Lk)\n", + " print(f\" {Lk} | {pm:6.3f} {wf[0]:+.3f} | {wf[1]:+.3f} (mass {float(p[fm].sum()):.3f}) {ws[0]:+.3f} | {ws[1]:+.3f} (mass {float(p[sm].sum()):.3f})\")\n", + "def part(wfun, sel):\n", + " return float(((p * sel * wfun)[:, None] * S).sum(0) @ g_len_E)\n", + "fail, succ = (R == 0).double(), (R == 1).double()\n", + "print(\" first-order d E|o| contributions :\")\n", + "print(f\" failures : unbiased {part(-b*torch.ones(N), fail):+.5f} 1/|o|-normalised {part(-b/lens, fail):+.5f}\")\n", + "print(f\" successes: unbiased {part((1-b)*torch.ones(N), succ):+.5f} 1/|o|-normalised {part((1-b)/lens, succ):+.5f}\")\n", + "dir_unb = float(g_unb @ g_len_E); dir_len = float(g_len @ g_len_E)\n", + "print(f\"\\n = {dir_unb:+.5f} = {dir_len:+.5f} (first-order change in expected length per unit step)\")\n", + "# bias term = -b * grad E[1/|o|] + (grad E[R/|o|] - ... ) : decomposition check\n", + "th = theta.clone().requires_grad_(True)\n", + "gR_over = torch.autograd.grad((lm.probs(th) * R / lens).sum(), th)[0]\n", + "th = theta.clone().requires_grad_(True)\n", + "g_inv = torch.autograd.grad((lm.probs(th) / lens).sum(), th)[0]\n", + "check(\"decomposition g_len = grad E[R/|o|] - b grad E[1/|o|]\", maxabs(g_len, gR_over - b * g_inv), 1e-14)" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "CHECK|exact|unbiased baseline estimator == grad J|2.776e-17|tol=1.0e-14|PASS\n", + "CHECK|exact|1/|o|-normalised estimator == grad of E[(R-b)/|o|] (b fixed)|6.939e-18|tol=1.0e-14|PASS\n", + "CHECK|exact|constant normaliser 1/Lmax gives grad J / Lmax (Dr.GRPO)|2.776e-17|tol=1.0e-14|PASS\n", + "\n", + " length | P(|o|) mean weight on FAILURES (unbiased | 1/|o|) mean weight on SUCCESSES (unbiased | 1/|o|)\n", + " 1 | 0.464 -0.166 | -0.166 (mass 0.464) +0.834 | +0.834 (mass 0.000)\n", + " 2 | 0.258 -0.166 | -0.083 (mass 0.246) +0.834 | +0.417 (mass 0.012)\n", + " 3 | 0.142 -0.166 | -0.055 (mass 0.049) +0.834 | +0.278 (mass 0.094)\n", + " 4 | 0.136 -0.166 | -0.041 (mass 0.075) +0.834 | +0.209 (mass 0.061)\n", + " first-order d E|o| contributions :\n", + " failures : unbiased +0.02092 1/|o|-normalised +0.04135\n", + " successes: unbiased +0.10518 1/|o|-normalised +0.03104\n", + "\n", + " = +0.12611 = +0.07239 (first-order change in expected length per unit step)\n", + "CHECK|exact|decomposition g_len = grad E[R/|o|] - b grad E[1/|o|]|7.806e-18|tol=1.0e-14|PASS\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 5. Naive vs efficient loss reductions on a sampled group (gradient agreement) -- GRPO, DAPO-token, Dr.GRPO" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "source": [ + "Gs = 6\n", + "gen = torch.Generator().manual_seed(42)\n", + "idx = torch.multinomial(p, Gs, replacement=True, generator=gen)\n", + "Rg = R[idx]\n", + "Ahat = adv_grpo(Rg, eps=1e-4); Actr = adv_center(Rg)\n", + "th = theta.clone().requires_grad_(True)\n", + "lp_tok = lm.token_logprobs(th)[idx] # [G, L]\n", + "m = lm.mask[idx]\n", + "lens_g = m.sum(1)\n", + "def surrogate_naive(lp_tok, A, mode):\n", + " tot = 0.0\n", + " if mode == 'grpo':\n", + " for i in range(Gs): tot = tot + sum(A[i] * lp_tok[i, t] for t in range(int(lens_g[i]))) / lens_g[i]\n", + " return tot / Gs\n", + " if mode == 'dapo':\n", + " n = float(lens_g.sum())\n", + " for i in range(Gs): tot = tot + sum(A[i] * lp_tok[i, t] for t in range(int(lens_g[i])))\n", + " return tot / n\n", + " if mode == 'dr':\n", + " for i in range(Gs): tot = tot + sum(A[i] * lp_tok[i, t] for t in range(int(lens_g[i])))\n", + " return tot / (Gs * lm.L)\n", + "def surrogate_vec(lp_tok, A, mode):\n", + " w = A[:, None] * m\n", + " if mode == 'grpo': return ((w * lp_tok).sum(1) / lens_g).mean()\n", + " if mode == 'dapo': return (w * lp_tok).sum() / m.sum()\n", + " if mode == 'dr': return (w * lp_tok).sum() / (Gs * lm.L)\n", + "for mode, A in (('grpo', Ahat), ('dapo', Ahat), ('dr', Actr)):\n", + " gn = torch.autograd.grad(surrogate_naive(lp_tok, A, mode), th, retain_graph=True)[0]\n", + " gv = torch.autograd.grad(surrogate_vec(lp_tok, A, mode), th, retain_graph=True)[0]\n", + " check(f\"naive python loops == vectorised masked reduction ({mode})\", maxabs(gn, gv), 1e-14)\n", + "print(\"done; checks:\", len(CHECKS))" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "CHECK|exact|naive python loops == vectorised masked reduction (grpo)|0.000e+00|tol=1.0e-14|PASS\n", + "CHECK|exact|naive python loops == vectorised masked reduction (dapo)|0.000e+00|tol=1.0e-14|PASS\n", + "CHECK|exact|naive python loops == vectorised masked reduction (dr)|0.000e+00|tol=1.0e-14|PASS\n", + "done; checks: 23\n" + ] + } + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.12.11" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} \ No newline at end of file diff --git a/code/notebooks/0106-dapo-pathologies.ipynb b/code/notebooks/0106-dapo-pathologies.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..6756f2e9cc7d3972f35c057a12f3e49b5fafcec9 --- /dev/null +++ b/code/notebooks/0106-dapo-pathologies.ipynb @@ -0,0 +1,450 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# 0106 - DAPO: each of its four techniques derived as the fix for a named pathology\n", + "Companion to `notes/0106 - DAPO.pdf`. fp64; exact enumeration where possible, Monte Carlo (flagged) elsewhere." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "source": [ + "import sys, math\n", + "sys.path.insert(0, \"..\")\n", + "from _m1_lib import *\n", + "import numpy as np" + ], + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 1. Clip-Higher (pathology: entropy collapse / exploration bottleneck)\n", + "1a. Asymmetric clip gradient table (eps_lo = 0.2, eps_hi = 0.28 as in DAPO Sec. 4.1) vs autograd; 1b. absolute probability headroom." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "source": [ + "el, eh = 0.2, 0.28\n", + "def f_closed(r, A, el, eh): return A * min(r, 1 + eh) if A > 0 else A * max(r, 1 - el)\n", + "def grad_auto(r, A):\n", + " lp = torch.tensor(math.log(r), requires_grad=True); rr = lp.exp()\n", + " f = torch.minimum(rr * A, torch.clamp(rr, 1 - el, 1 + eh) * A)\n", + " return float(torch.autograd.grad(f, lp)[0])\n", + "worst = 0.0; n = 0\n", + "for A in (-1.5, -0.2, 0.2, 1.5):\n", + " for r in np.linspace(0.05, 3.0, 400):\n", + " if min(abs(r - (1 - el)), abs(r - (1 + eh))) < 1e-6: continue\n", + " live = (r < 1 + eh) if A > 0 else (r > 1 - el)\n", + " worst = max(worst, abs(grad_auto(r, A) - (A * r if live else 0.0))); n += 1\n", + "check(f\"asymmetric clip: autograd == table (A r on live set) on {n} points\", worst, 1e-14)\n", + "print(\" log-ratio limits: log(1+eps_hi)=%.4f log(1-eps_lo)=%.4f (symmetric eps=.2: +%.4f / %.4f)\" % (math.log(1 + eh), math.log(1 - el), math.log(1.2), math.log(0.8)))\n", + "print(\" absolute probability reachable by ONE clipped update of a token with pi_old = p (A > 0):\")\n", + "print(\" p cap eps=0.20 cap eps=0.28 headroom(0.20) headroom(0.28)\")\n", + "for pp in (0.001, 0.01, 0.1, 0.5, 0.9):\n", + " c1, c2 = min(1.0, 1.2 * pp), min(1.0, 1.28 * pp)\n", + " print(f\" {pp:5.3f} {c1:8.4f} {c2:8.4f} {c1-pp:8.4f} {c2-pp:8.4f}\")\n", + "check(\"DAPO example: 0.01 -> 0.012 at eps=0.2\", abs(min(1, 1.2 * 0.01) - 0.012), 1e-15)\n", + "check(\"DAPO example: 0.9 -> min(1, 1.08) = 1.0 (no effective cap)\", abs(min(1, 1.2 * 0.9) - 1.0), 1e-15)" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "CHECK|exact|asymmetric clip: autograd == table (A r on live set) on 1600 points|8.882e-16|tol=1.0e-14|PASS\n", + " log-ratio limits: log(1+eps_hi)=0.2469 log(1-eps_lo)=-0.2231 (symmetric eps=.2: +0.1823 / -0.2231)\n", + " absolute probability reachable by ONE clipped update of a token with pi_old = p (A > 0):\n", + " p cap eps=0.20 cap eps=0.28 headroom(0.20) headroom(0.28)\n", + " 0.001 0.0012 0.0013 0.0002 0.0003\n", + " 0.010 0.0120 0.0128 0.0020 0.0028\n", + " 0.100 0.1200 0.1280 0.0200 0.0280\n", + " 0.500 0.6000 0.6400 0.1000 0.1400\n", + " 0.900 1.0000 1.0000 0.1000 0.1000\n", + "CHECK|exact|DAPO example: 0.01 -> 0.012 at eps=0.2|0.000e+00|tol=1.0e-15|PASS\n", + "CHECK|exact|DAPO example: 0.9 -> min(1, 1.08) = 1.0 (no effective cap)|0.000e+00|tol=1.0e-15|PASS\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "1c. Entropy dynamics of a softmax policy: dH = -Cov_{a~pi}( log pi_a , dz_a ) (exact first order). PG: dz = eta*pi*A; NPG: dz = eta*A." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "source": [ + "def entropy(z):\n", + " lp = torch.log_softmax(z, 0); return float(-(lp.exp() * lp).sum())\n", + "def cov_formula(z, dz):\n", + " pi = torch.softmax(z, 0); lg = pi.log()\n", + " return float((pi * lg * dz).sum() - (pi * lg).sum() * (pi * dz).sum())\n", + "gen = torch.Generator().manual_seed(3)\n", + "worst_rel = 0.0\n", + "for trial in range(6):\n", + " K = 5\n", + " z = torch.randn(K, generator=gen) * 1.5\n", + " A = torch.randn(K, generator=gen)\n", + " pi = torch.softmax(z, 0); A = A - (pi * A).sum() # advantages must have zero pi-mean\n", + " for name, dz in ((\"PG\", pi * A), (\"NPG\", A)):\n", + " eta = 1e-6\n", + " dH_fd = (entropy(z + eta * dz) - entropy(z - eta * dz)) / (2 * eta)\n", + " dH_cov = -cov_formula(z, dz)\n", + " worst_rel = max(worst_rel, abs(dH_fd - dH_cov) / max(1e-12, abs(dH_cov)))\n", + "check(\"entropy change = -Cov(log pi, dz): finite difference vs formula (PG & NPG, 6 random policies)\", worst_rel, 1e-6)\n", + "pi = torch.tensor([0.90, 0.05, 0.03, 0.02]); z = pi.log()\n", + "def dH_for(A):\n", + " A = torch.tensor(A); A = A - (pi * A).sum(); return -cov_formula(z, pi * A)\n", + "print(\" dH/d eta for tabular PG: good action is the DOMINANT one (pi=.90): %+.5f good action is a RARE one (pi=.02): %+.5f\" %\n", + " (dH_for([1., 0, 0, 0]), dH_for([0., 0, 0, 1.])))\n", + "assert dH_for([1., 0, 0, 0]) < 0 < dH_for([0., 0, 0, 1.])\n", + "check(\"positive-advantage dominant token lowers entropy; positive-advantage rare token raises it (signs)\", 0.0, 0.0)" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "CHECK|exact|entropy change = -Cov(log pi, dz): finite difference vs formula (PG & NPG, 6 random policies)|6.651e-09|tol=1.0e-06|PASS\n", + " dH/d eta for tabular PG: good action is the DOMINANT one (pi=.90): -0.03566 good action is a RARE one (pi=.02): +0.00641\n", + "CHECK|exact|positive-advantage dominant token lowers entropy; positive-advantage rare token raises it (signs)|0.000e+00|tol=0.0e+00|PASS\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "1d. Exact multi-step dynamics on a 2-action bandit (rare correct action a, pi=0.01): how fast can p(a) grow per rollout iteration?\n", + "Objective per inner step = sum_a pi_old(a) * f(r_a, A_a) (exact expectation), tabular logits, plain gradient ascent." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "source": [ + "def run_bandit(p0, eps_hi, outer=60, inner=400, lr=0.4, el=0.2):\n", + " z = torch.log(torch.tensor([p0, 1 - p0])).clone()\n", + " traj = [p0]\n", + " for it in range(outer):\n", + " pi_old = torch.softmax(z, 0).detach()\n", + " A = torch.tensor([1.0, 0.0]) - pi_old[0] # A_a = 1 - p, A_b = -p (binary reward on action a)\n", + " zz = z.clone().requires_grad_(True)\n", + " for _ in range(inner):\n", + " pi = torch.softmax(zz, 0); r = pi / pi_old\n", + " f = torch.where(A > 0, A * torch.clamp(r, max=1 + eps_hi), A * torch.clamp(r, min=1 - el))\n", + " obj = (pi_old * f).sum()\n", + " g = torch.autograd.grad(obj, zz)[0]\n", + " with torch.no_grad(): zz += lr * g\n", + " z = zz.detach()\n", + " traj.append(float(torch.softmax(z, 0)[0]))\n", + " return torch.tensor(traj)\n", + "t20, t28, tinf = run_bandit(0.01, 0.2), run_bandit(0.01, 0.28), run_bandit(0.01, 10.0)\n", + "def first_cross(t, v=0.5):\n", + " idx = (t >= v).nonzero()\n", + " return int(idx[0]) if len(idx) else None\n", + "print(\" iterations for p(a) to reach 0.5 starting from 0.01: eps_hi=0.20 ->\", first_cross(t20), \" eps_hi=0.28 ->\", first_cross(t28), \" no upper clip ->\", first_cross(tinf))\n", + "g20 = (t20[1:6] / t20[:5]); g28 = (t28[1:6] / t28[:5])\n", + "print(\" per-iteration growth factor p'/p over first 5 iterations: eps=.20:\", [round(float(x), 3) for x in g20], \" eps=.28:\", [round(float(x), 3) for x in g28])\n", + "print(\" growth in iteration 1 slightly EXCEEDS 1+eps (1.251 vs 1.2; 1.342 vs 1.28): the clip stops only the DIRECT push on the rare action; the unclipped negative-advantage action keeps moving mass to it\")\n", + "assert first_cross(t28) <= first_cross(t20)\n", + "check(\"larger eps_hi never slows the rare action (iterations to 0.5: eps_hi=.28 <= eps_hi=.20)\", 0.0, 0.0)" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + " iterations for p(a) to reach 0.5 starting from 0.01: eps_hi=0.20 -> 8 eps_hi=0.28 -> 7 no upper clip -> 2\n", + " per-iteration growth factor p'/p over first 5 iterations: eps=.20: [1.251, 1.283, 1.331, 1.454, 1.847] eps=.28: [1.342, 1.389, 1.488, 1.805, 4.8]\n", + " growth in iteration 1 slightly EXCEEDS 1+eps (1.251 vs 1.2; 1.342 vs 1.28): the clip stops only the DIRECT push on the rare action; the unclipped negative-advantage action keeps moving mass to it\n", + "CHECK|exact|larger eps_hi never slows the rare action (iterations to 0.5: eps_hi=.28 <= eps_hi=.20)|0.000e+00|tol=0.0e+00|PASS\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 2. Dynamic sampling (pathology: zero-advantage groups). P(dead group) = p^G + (1-p)^G" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "source": [ + "torch.manual_seed(8) # torch.distributions samplers use the global RNG\n", + "gen = torch.Generator().manual_seed(8)\n", + "for G in (4, 16):\n", + " for pp in (0.02, 0.3):\n", + " ks = torch.distributions.Binomial(G, torch.tensor(pp)).sample((400000,))\n", + " emp = float(((ks == 0) | (ks == G)).double().mean()); ex = pp ** G + (1 - pp) ** G\n", + " se = math.sqrt(ex * (1 - ex) / 400000)\n", + " check(f\"P(dead) MC vs p^G+(1-p)^G (G={G}, p={pp}) |diff|/SE\", abs(emp - ex) / se, 4.0, kind=\"mc\")\n", + "# expected live fraction for p ~ Beta(a,b): 1 - B(a+G,b)/B(a,b) - B(a,b+G)/B(a,b)\n", + "def lbeta(a, b): return math.lgamma(a) + math.lgamma(b) - math.lgamma(a + b)\n", + "def live_fraction(a, b, G): return 1 - math.exp(lbeta(a + G, b) - lbeta(a, b)) - math.exp(lbeta(a, b + G) - lbeta(a, b))\n", + "xs, ws = np.polynomial.legendre.leggauss(400); xs = 0.5 * (xs + 1); ws = 0.5 * ws\n", + "def live_quad(a, b, G):\n", + " dens = np.exp(-lbeta(a, b)) * xs ** (a - 1) * (1 - xs) ** (b - 1)\n", + " return float((ws * dens * (1 - xs ** G - (1 - xs) ** G)).sum())\n", + "for (a, b) in ((2, 5), (3, 3), (2, 8)):\n", + " check(f\"expected live fraction closed form vs quadrature, Beta({a},{b}), G=16\", abs(live_fraction(a, b, 16) - live_quad(a, b, 16)), 1e-10)\n", + "a, b, G = 2.0, 8.0, 16\n", + "rho = live_fraction(a, b, G)\n", + "print(f\" Beta({a:.0f},{b:.0f}) prompts (mean p={a/(a+b):.2f}), G={G}: live fraction rho = {rho:.3f} -> oversampling factor 1/rho = {1/rho:.2f}\")\n", + "# MC: prompts drawn until m live ones are found; expectation m / rho\n", + "torch.manual_seed(2)\n", + "gen = torch.Generator().manual_seed(2)\n", + "m_need, reps = 64, 600\n", + "draws = []\n", + "beta = torch.distributions.Beta(torch.tensor(a), torch.tensor(b))\n", + "for _ in range(reps):\n", + " live = 0; drawn = 0\n", + " while live < m_need:\n", + " pp = beta.sample()\n", + " k = torch.distributions.Binomial(G, pp).sample()\n", + " drawn += 1; live += int(0 < k < G)\n", + " draws.append(drawn)\n", + "d = torch.tensor(draws, dtype=torch.float64)\n", + "check(\"MC: prompts drawn to collect 64 live groups vs m/rho (z-score)\", abs(float(d.mean()) - m_need / rho) / (float(d.std()) / reps ** 0.5), 4.0, kind=\"mc\")\n", + "# Unbiasedness of dynamic sampling (exact, no MC): prompts q_j with success prob p_j and weights w_j; per-prompt Bernoulli-bandit gradient g = p_hat(1-p_hat)\n", + "# (centred advantage x score, see 0110). Filtered batch: a prompt is accepted with probability P(live|q_j) (so accepted prompts have law w_j P(live_j)/rho) and its\n", + "# group contributes E[g | live, q_j] computed DIRECTLY from the binomial restricted to 0 oversampling factor 1/rho = 1.14\n", + "CHECK|mc|MC: prompts drawn to collect 64 live groups vs m/rho (z-score)|5.547e-01|tol=4.0e+00|PASS\n", + "CHECK|exact|dynamic sampling: E[g | accepted group] = (1/rho) * E[g | unfiltered group] (exact, discrete prompt law)|0.000e+00|tol=1.0e-15|PASS\n", + " with global token normalisation, dead groups stay in the denominator: signal scales by rho = 0.88 (identity, no estimation)\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 3. Token-level loss: global denominator is a RATIO estimator (bias O(1/#samples)); constant denominator is exactly unbiased\n", + "Exact over all G-tuples (G=2,3) on the TinyLM." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "source": [ + "lm = TinyLM(V=4, L=4, seed=0, scale=1.0)\n", + "theta = lm.theta0.clone(); R = default_reward(lm); p = lm.probs(theta); S = score_matrix(lm, theta); N = lm.N\n", + "lens = lm.lengths\n", + "grad_J = (p[:, None] * (R - float(p @ R))[:, None] * S).sum(0)\n", + "Lbar = float((p * lens).sum())\n", + "def exact_tokenlevel(G, normaliser):\n", + " \"\"\"E[ sum_i A_i * score(o_i) / D ] with centred advantages A_i = R_i - mean R; D = sum_i |o_i| (global) or G*Lbar (constant).\"\"\"\n", + " if G == 2:\n", + " R1, R2 = R[:, None].expand(N, N), R[None, :].expand(N, N)\n", + " l1, l2 = lens[:, None].expand(N, N), lens[None, :].expand(N, N)\n", + " P = p[:, None] * p[None, :]\n", + " mu = (R1 + R2) / 2\n", + " D = (l1 + l2) if normaliser == 'global' else torch.full((N, N), G * Lbar)\n", + " w1 = P * (R1 - mu) / D; w2 = P * (R2 - mu) / D\n", + " return w1.sum(1) @ S + w2.sum(0) @ S\n", + " g = torch.zeros(S.shape[1])\n", + " R2, R3 = R[:, None].expand(N, N), R[None, :].expand(N, N)\n", + " l2, l3 = lens[:, None].expand(N, N), lens[None, :].expand(N, N)\n", + " P23 = p[:, None] * p[None, :]\n", + " for i in range(N):\n", + " mu = (R[i] + R2 + R3) / 3\n", + " D = (lens[i] + l2 + l3) if normaliser == 'global' else torch.full((N, N), G * Lbar)\n", + " g += p[i] * ((P23 * (R[i] - mu) / D).sum()) * S[i]\n", + " g += p[i] * ((P23 * (R2 - mu) / D).sum(1) @ S)\n", + " g += p[i] * ((P23 * (R3 - mu) / D).sum(0) @ S)\n", + " return g\n", + "for G in (2, 3):\n", + " gc = exact_tokenlevel(G, 'const'); gg = exact_tokenlevel(G, 'global')\n", + " check(f\"G={G}: constant denominator G*E|o| gives (1-1/G) grad J / (G E|o|) * G exactly\", maxabs(gc * Lbar, (1 - 1 / G) * grad_J), 1e-13)\n", + " rel = float((gg - gc).norm() / gc.norm()); cs = float(gg @ grad_J / (gg.norm() * grad_J.norm()))\n", + " print(f\" G={G}: global-token-mean estimator: relative deviation from the unbiased one = {rel:.4f}, cosine to grad J = {cs:.5f}\")" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "CHECK|exact|G=2: constant denominator G*E|o| gives (1-1/G) grad J / (G E|o|) * G exactly|1.041e-17|tol=1.0e-13|PASS\n", + " G=2: global-token-mean estimator: relative deviation from the unbiased one = 0.1591, cosine to grad J = 0.99588\n", + "CHECK|exact|G=3: constant denominator G*E|o| gives (1-1/G) grad J / (G E|o|) * G exactly|4.163e-17|tol=1.0e-13|PASS\n", + " G=3: global-token-mean estimator: relative deviation from the unbiased one = 0.1000, cosine to grad J = 0.99779\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 4. Overlong handling: soft punishment is a REWARD; loss masking is a reward equal to the baseline" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "source": [ + "def soft_overlong(Ly, Lmax, Lcache):\n", + " if Ly <= Lmax - Lcache: return 0.0\n", + " if Ly <= Lmax: return ((Lmax - Lcache) - Ly) / Lcache\n", + " return -1.0\n", + "Lmax, Lcache = 16384, 4096\n", + "check(\"soft overlong: 0 at Lmax-Lcache\", abs(soft_overlong(Lmax - Lcache, Lmax, Lcache)), 0.0)\n", + "check(\"soft overlong: -1 at Lmax (continuous with the truncation penalty)\", abs(soft_overlong(Lmax, Lmax, Lcache) + 1.0), 0.0)\n", + "ls = np.arange(1, 24000)\n", + "vals = np.array([soft_overlong(int(x), Lmax, Lcache) for x in ls])\n", + "slope = float(np.abs(np.diff(vals[:Lmax])).max()) # largest per-token change of the ramp (up to Lmax)\n", + "check(\"soft overlong: ramp slope is exactly 1/Lcache (|max slope * Lcache - 1|)\", abs(slope * Lcache - 1.0), 1e-9)\n", + "print(\" hard truncation reward has a jump of size 1 at Lmax; the soft ramp spreads it over\", Lcache, \"tokens\")\n", + "# masking == baseline reward (TinyLM: 'truncated' = hit the length limit L=4 without EOS)\n", + "trunc = (lm.mask.sum(1) == lm.L) & (lm.tokens[:, -1] != lm.EOS)\n", + "b = float(p @ R)\n", + "g_masked = (p[:, None] * ((~trunc).double() * (R - b))[:, None] * S).sum(0) # loss masked on truncated samples\n", + "Rprime = torch.where(trunc, torch.full_like(R, b), R) # truncated sample gets reward = baseline\n", + "g_rb = (p[:, None] * (Rprime - b)[:, None] * S).sum(0)\n", + "check(\"loss-masking == assigning the baseline value as reward to truncated samples (exact)\", maxabs(g_masked, g_rb), 1e-14)\n", + "th = theta.clone().requires_grad_(True)\n", + "Rlen = torch.where(trunc, torch.full_like(R, -1.0), torch.zeros_like(R)) + 0.0\n", + "g_total_reward = (p[:, None] * (R + Rlen - b)[:, None] * S).sum(0)\n", + "th = theta.clone().requires_grad_(True)\n", + "g_sep = torch.autograd.grad((lm.probs(th) * (R + Rlen)).sum(), th)[0]\n", + "check(\"shaping as reward: policy gradient of E[R + R_len] (separable objective)\", maxabs(g_total_reward, g_sep), 1e-14)\n", + "g_hard = (p[:, None] * (R - b)[:, None] * S).sum(0) # truncated treated as plain failures (R=0)\n", + "print(f\" ||g_masked - g_hard|| = {float((g_masked-g_hard).norm()):.4f} vs ||grad J|| = {float(grad_J.norm()):.4f} (filtering changes the objective)\")" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "CHECK|exact|soft overlong: 0 at Lmax-Lcache|0.000e+00|tol=0.0e+00|PASS\n", + "CHECK|exact|soft overlong: -1 at Lmax (continuous with the truncation penalty)|0.000e+00|tol=0.0e+00|PASS\n", + "CHECK|exact|soft overlong: ramp slope is exactly 1/Lcache (|max slope * Lcache - 1|)|0.000e+00|tol=1.0e-09|PASS\n", + " hard truncation reward has a jump of size 1 at Lmax; the soft ramp spreads it over 4096 tokens\n", + "CHECK|exact|loss-masking == assigning the baseline value as reward to truncated samples (exact)|0.000e+00|tol=1.0e-14|PASS\n", + "CHECK|exact|shaping as reward: policy gradient of E[R + R_len] (separable objective)|2.776e-17|tol=1.0e-14|PASS\n", + " ||g_masked - g_hard|| = 0.0124 vs ||grad J|| = 0.1825 (filtering changes the objective)\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 5. Naive vs vectorised DAPO objective: asymmetric clip + global token-level normaliser + dynamic-sampling filter (gradient agreement)" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": {}, + "source": [ + "th_old = lm.theta0.clone(); lq = lm.token_logprobs(th_old); q_old = lm.probs(th_old)\n", + "th_new = perturb(th_old, 1.0, seed=3)\n", + "gen = torch.Generator().manual_seed(17)\n", + "G, groups = 8, []\n", + "while len(groups) < 3: # dynamic sampling: resample until the group is live (0 < k < G)\n", + " idx = torch.multinomial(q_old, G, replacement=True, generator=gen)\n", + " k = int(R[idx].sum())\n", + " if 0 < k < G: groups.append(idx)\n", + "def dapo_naive(th):\n", + " lp_all = lm.token_logprobs(th)\n", + " tot, ntok = 0.0, 0\n", + " for idx in groups:\n", + " Rg = R[idx]; A = (Rg - Rg.mean()) / (Rg.std(unbiased=False) + 1e-6)\n", + " for j, i in enumerate(idx.tolist()):\n", + " for t in range(int(lm.lengths[i])):\n", + " r = torch.exp(lp_all[i, t] - lq[i, t])\n", + " a = A[j]\n", + " term = a * torch.clamp(r, max=1 + eh) if a > 0 else a * torch.clamp(r, min=1 - el)\n", + " tot = tot + term; ntok += 1\n", + " return -tot / ntok\n", + "def dapo_vec(th):\n", + " lp_all = lm.token_logprobs(th)\n", + " idx = torch.cat(groups); Rg = R[idx].reshape(len(groups), G)\n", + " A = ((Rg - Rg.mean(1, keepdim=True)) / (Rg.std(1, unbiased=False, keepdim=True) + 1e-6)).reshape(-1)\n", + " r = (lp_all[idx] - lq[idx]).exp(); mk = lm.mask[idx]\n", + " f = torch.where(A[:, None] > 0, A[:, None] * torch.clamp(r, max=1 + eh), A[:, None] * torch.clamp(r, min=1 - el))\n", + " return -(f * mk).sum() / mk.sum()\n", + "t1 = th_new.clone().requires_grad_(True); t2 = th_new.clone().requires_grad_(True)\n", + "check(\"naive triple loop == vectorised DAPO loss (value)\", abs(float(dapo_naive(t1).detach()) - float(dapo_vec(t2).detach())), 1e-13)\n", + "g1 = torch.autograd.grad(dapo_naive(t1), t1)[0]; g2 = torch.autograd.grad(dapo_vec(t2), t2)[0]\n", + "check(\"naive triple loop == vectorised DAPO loss (gradient)\", maxabs(g1, g2), 1e-13)\n", + "print(\"done; checks:\", len(CHECKS))" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "CHECK|exact|naive triple loop == vectorised DAPO loss (value)|1.665e-16|tol=1.0e-13|PASS\n", + "CHECK|exact|naive triple loop == vectorised DAPO loss (gradient)|0.000e+00|tol=1.0e-13|PASS\n", + "done; checks: 24\n" + ] + } + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.12.11" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} \ No newline at end of file diff --git a/code/notebooks/0107-cispo.ipynb b/code/notebooks/0107-cispo.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..0dba4014a0888931b2b02807ec6bde1b938b545a --- /dev/null +++ b/code/notebooks/0107-cispo.ipynb @@ -0,0 +1,596 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# 0107 - CISPO (MiniMax-M1 Sec. 3.1): clip the importance weight, not the update\n", + "Companion to `notes/0107 - CISPO.pdf`. fp64; exact enumeration on the 121-sequence toy LM." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "source": [ + "import sys, math\n", + "sys.path.insert(0, \"..\")\n", + "from _m1_lib import *\n", + "import numpy as np\n", + "lm = TinyLM(V=4, L=4, seed=0, scale=1.0)\n", + "R = default_reward(lm)\n", + "th_old = lm.theta0.clone(); q = lm.probs(th_old); N, L = lm.N, lm.L\n", + "b = float(q @ R)\n", + "A_tok = (R - b)[:, None] * lm.mask # outcome advantage, same for all tokens of a sequence\n", + "lq = lm.token_logprobs(th_old)\n", + "Ttens = token_score_tensor(lm, th_old) # score of each token at THETA_OLD (used to build gradients analytically)\n", + "eh, el = 0.28, 0.2" + ], + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 1. Gradient weights. Every method has gradient sum_t w(r_t, A_t) * A_t * grad log pi_theta(o_t | s_t)\n", + " PPO/GRPO (clip lo/hi): w = r * 1[live] live = (A>0 & r<1+eh) | (A<0 & r>1-el)\n", + " CISPO: w = min(r, 1+eh) (A-independent, always live) [loss = sg(min(r,1+eh)) * A * log pi]\n", + " token IS-PG (no clip): w = r\n", + "We verify against autograd of the actual loss functions on the toy model at a perturbed theta." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "source": [ + "def ratio_and_lp(th):\n", + " lp = lm.token_logprobs(th)\n", + " r = (lp - lq).exp() * lm.mask + (1 - lm.mask)\n", + " return r, lp\n", + "def loss_ppo(th):\n", + " r, lp = ratio_and_lp(th)\n", + " f = torch.where(A_tok > 0, A_tok * torch.clamp(r, max=1 + eh), A_tok * torch.clamp(r, min=1 - el))\n", + " return -(q[:, None] * f * lm.mask).sum()\n", + "def loss_cispo(th, eps_hi=eh):\n", + " r, lp = ratio_and_lp(th)\n", + " rhat = torch.clamp(r, max=1 + eps_hi).detach() # sg( clip(r) )\n", + " return -(q[:, None] * rhat * A_tok * lp * lm.mask).sum()\n", + "def loss_isPG(th):\n", + " r, lp = ratio_and_lp(th)\n", + " return -(q[:, None] * r.detach() * A_tok * lp * lm.mask).sum()\n", + "def weights(th, method):\n", + " r, _ = ratio_and_lp(th); r = r.detach()\n", + " if method == 'ppo': live = ((A_tok > 0) & (r < 1 + eh)) | ((A_tok < 0) & (r > 1 - el)); return r * live\n", + " if method == 'cispo': return torch.clamp(r, max=1 + eh)\n", + " if method == 'is': return r\n", + "def grad_from_weights(th, w):\n", + " # sum_{o,t} q(o) w A grad log pi_theta(o_t|s_t) with the score evaluated at theta (autograd through lp)\n", + " t = th.clone().requires_grad_(True)\n", + " lp = lm.token_logprobs(t)\n", + " return torch.autograd.grad((q[:, None] * w * A_tok * lp * lm.mask).sum(), t)[0]\n", + "th = perturb(th_old, 1.5, seed=11)\n", + "r_now, _ = ratio_and_lp(th)\n", + "print(\"ratio range over (o,t): min %.3f max %.3f\" % (float(r_now[lm.mask > 0].min()), float(r_now[lm.mask > 0].max())))\n", + "for name, lossf in ((\"ppo\", loss_ppo), (\"cispo\", loss_cispo), (\"is\", loss_isPG)):\n", + " t = th.clone().requires_grad_(True)\n", + " g_auto = torch.autograd.grad(lossf(t), t)[0]\n", + " g_w = -grad_from_weights(th, weights(th, name))\n", + " check(f\"{name}: autograd of the loss == sum_t w(r,A) A grad log pi (weights table)\", maxabs(g_auto, g_w), 1e-13)" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "ratio range over (o,t): min 0.455 max 1.992\n", + "CHECK|exact|ppo: autograd of the loss == sum_t w(r,A) A grad log pi (weights table)|1.388e-17|tol=1.0e-13|PASS\n", + "CHECK|exact|cispo: autograd of the loss == sum_t w(r,A) A grad log pi (weights table)|0.000e+00|tol=1.0e-13|PASS\n", + "CHECK|exact|is: autograd of the loss == sum_t w(r,A) A grad log pi (weights table)|0.000e+00|tol=1.0e-13|PASS\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 2. Exactly which tokens does PPO/GRPO silence and CISPO keep?" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "source": [ + "live = ((A_tok > 0) & (r_now < 1 + eh)) | ((A_tok < 0) & (r_now > 1 - el))\n", + "sel = lm.mask > 0\n", + "pw = q[:, None] * lm.mask # probability weight of each (o,t) token occurrence\n", + "tot = float(pw.sum())\n", + "zeroed = sel & ~live\n", + "cells = {\"A>0, r>1+eh (silenced by PPO)\": sel & (A_tok > 0) & (r_now > 1 + eh),\n", + " \"A<0, r<1-el (silenced by PPO)\": sel & (A_tok < 0) & (r_now < 1 - el),\n", + " \"A<0, r>1+eh (PPO weight r unbounded; CISPO caps at 1+eh)\": sel & (A_tok < 0) & (r_now > 1 + eh)}\n", + "for k, m in cells.items(): print(f\" {k:60s}: {int(m.sum()):4d} token slots, probability mass {float((pw*m).sum())/tot:.4f}\")\n", + "Wc = weights(th, 'cispo'); Wp = weights(th, 'ppo')\n", + "silenced_by_ppo_kept_by_cispo = sel & (Wp == 0) & (Wc > 0)\n", + "check(\"{tokens with PPO weight 0 and CISPO weight>0} == cells 1 U 2 exactly (symmetric difference size)\", float((silenced_by_ppo_kept_by_cispo ^ (cells[\"A>0, r>1+eh (silenced by PPO)\"] | cells[\"A<0, r<1-el (silenced by PPO)\"])).sum()), 0.0)\n", + "check(\"CISPO never silences a token (min weight over all (o,t) with A != 0)\", abs(float(Wc[sel].min()) - float(r_now[sel].min())) if float(r_now[sel].min()) < 1 + eh else 0.0, 1e-15)\n", + "check(\"PPO silenced set == the MiniMax-M1 mask M_{i,t}=0 of Eq. 7 (symmetric difference size)\",\n", + " float((zeroed ^ (((A_tok > 0) & (r_now > 1 + eh)) | ((A_tok < 0) & (r_now < 1 - el))) & sel).sum()), 0.0)\n", + "# weight table\n", + "print(\"\\n gradient weight w(r,A) (loss gradient = w * A * grad log pi), eps_hi=0.28, eps_lo=0.2\")\n", + "print(\" r PPO(A>0) PPO(A<0) CISPO unclipped-IS\")\n", + "for r in (0.3, 0.8, 1.0, 1.28, 2.0, 5.0, 20.0):\n", + " wp_pos = r if r < 1 + eh else 0.0; wp_neg = r if r > 1 - el else 0.0\n", + " print(f\" {r:5.2f} {wp_pos:7.2f} {wp_neg:7.2f} {min(r, 1+eh):6.2f} {r:7.2f}\")" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + " A>0, r>1+eh (silenced by PPO) : 7 token slots, probability mass 0.0134\n", + " A<0, r<1-el (silenced by PPO) : 63 token slots, probability mass 0.0551\n", + " A<0, r>1+eh (PPO weight r unbounded; CISPO caps at 1+eh) : 83 token slots, probability mass 0.0160\n", + "CHECK|exact|{tokens with PPO weight 0 and CISPO weight>0} == cells 1 U 2 exactly (symmetric difference size)|0.000e+00|tol=0.0e+00|PASS\n", + "CHECK|exact|CISPO never silences a token (min weight over all (o,t) with A != 0)|0.000e+00|tol=1.0e-15|PASS\n", + "CHECK|exact|PPO silenced set == the MiniMax-M1 mask M_{i,t}=0 of Eq. 7 (symmetric difference size)|0.000e+00|tol=0.0e+00|PASS\n", + "\n", + " gradient weight w(r,A) (loss gradient = w * A * grad log pi), eps_hi=0.28, eps_lo=0.2\n", + " r PPO(A>0) PPO(A<0) CISPO unclipped-IS\n", + " 0.30 0.30 0.00 0.30 0.30\n", + " 0.80 0.80 0.00 0.80 0.80\n", + " 1.00 1.00 1.00 1.00 1.00\n", + " 1.28 0.00 1.28 1.28 1.28\n", + " 2.00 0.00 2.00 1.28 2.00\n", + " 5.00 0.00 5.00 1.28 5.00\n", + " 20.00 0.00 20.00 1.28 20.00\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 3. Unified mask formulation (paper Eq. 6-7): J = sg(r_hat) A log pi * M ; with r_hat = r it reproduces PPO's gradient" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "source": [ + "def loss_unify(th, rhat_clip=None):\n", + " r, lp = ratio_and_lp(th)\n", + " M = (~(((A_tok > 0) & (r > 1 + eh)) | ((A_tok < 0) & (r < 1 - el)))).double()\n", + " rhat = r if rhat_clip is None else torch.clamp(r, max=rhat_clip)\n", + " return -(q[:, None] * rhat.detach() * A_tok * lp * M * lm.mask).sum()\n", + "t = th.clone().requires_grad_(True); g_u = torch.autograd.grad(loss_unify(t), t)[0]\n", + "t = th.clone().requires_grad_(True); g_p = torch.autograd.grad(loss_ppo(t), t)[0]\n", + "check(\"unified formulation with r_hat=r and PPO mask == PPO gradient\", maxabs(g_u, g_p), 1e-13)\n", + "# without clipping CISPO reduces to the standard (token-level IS) policy gradient; at theta_old it is the on-policy gradient\n", + "t = th.clone().requires_grad_(True); g_c_inf = torch.autograd.grad(loss_cispo(t, eps_hi=1e9), t)[0]\n", + "t = th.clone().requires_grad_(True); g_is = torch.autograd.grad(loss_isPG(t), t)[0]\n", + "check(\"CISPO with eps_hi=inf == token-level IS policy gradient\", maxabs(g_c_inf, g_is), 1e-14)\n", + "t = th_old.clone().requires_grad_(True); g_c0 = torch.autograd.grad(loss_cispo(t), t)[0]\n", + "t2 = th_old.clone().requires_grad_(True); gJ = torch.autograd.grad((lm.probs(t2) * R).sum(), t2)[0]\n", + "check(\"at theta_old CISPO == exact policy gradient grad J (r=1, nothing clipped)\", maxabs(-g_c0, gJ), 1e-14)" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "CHECK|exact|unified formulation with r_hat=r and PPO mask == PPO gradient|1.388e-17|tol=1.0e-13|PASS\n", + "CHECK|exact|CISPO with eps_hi=inf == token-level IS policy gradient|0.000e+00|tol=1.0e-14|PASS\n", + "CHECK|exact|at theta_old CISPO == exact policy gradient grad J (r=1, nothing clipped)|4.163e-17|tol=1.0e-14|PASS\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 4. Bias of CISPO relative to the (token-level) IS gradient, and the bound ||bias|| <= sum_t E[ (r - 1 - eps)_+ |A| ||score|| ]" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "source": [ + "def cispo_bias(th, eps):\n", + " r, _ = ratio_and_lp(th); r = r.detach()\n", + " Tth = token_score_tensor(lm, th) # score at theta\n", + " diff = (torch.clamp(r, max=1 + eps) - r) # r_hat - r (<= 0)\n", + " return (q[:, None, None] * (diff * A_tok)[:, :, None] * Tth).sum((0, 1)), Tth\n", + "for eps_ in (0.05, 0.1, 0.28):\n", + " bias, Tth = cispo_bias(th, eps_)\n", + " r, _ = ratio_and_lp(th)\n", + " bound = float((q[:, None] * (r - 1 - eps_).clamp(min=0) * A_tok.abs() * Tth.norm(dim=2) * lm.mask).sum())\n", + " print(f\" eps_hi={eps_:4.2f}: ||bias|| = {float(bias.norm()):.4e} bound = {bound:.4e} ||token-IS gradient|| = {float((q[:,None,None]*(r*A_tok)[:,:,None]*Tth).sum((0,1)).norm()):.4e}\")\n", + " assert float(bias.norm()) <= bound + 1e-15\n", + "check(\"bias bound holds (3 values of eps_hi)\", 0.0, 0.0)\n", + "# for comparison: how far each method's gradient is from the token-level IS gradient\n", + "t_ = th.clone().requires_grad_(True); gI = torch.autograd.grad(loss_isPG(t_), t_)[0]\n", + "t_ = th.clone().requires_grad_(True); gP = torch.autograd.grad(loss_ppo(t_), t_)[0]\n", + "t_ = th.clone().requires_grad_(True); gC = torch.autograd.grad(loss_cispo(t_), t_)[0]\n", + "print(f\" ||g_PPO - g_IS|| = {float((gP-gI).norm()):.4e} ||g_CISPO - g_IS|| = {float((gC-gI).norm()):.4e} (eps_hi=0.28, eps_lo=0.2; ||g_IS|| = {float(gI.norm()):.4e})\")\n", + "# distribution-free per-state bound E_old[(r-1-eps)_+] = sum_a (pi - (1+eps) pi_old)_+ <= TV(pi, pi_old)\n", + "worst = 0.0\n", + "for s in lm.prefixes():\n", + " po, pn = lm.cond_probs(th_old, s), lm.cond_probs(th, s)\n", + " lhs = float(((pn - 1.28 * po).clamp(min=0)).sum()); tv = 0.5 * float((pn - po).abs().sum())\n", + " worst = max(worst, lhs - tv)\n", + "check(\"per-state E_old[(r-1-eps)_+] <= TV(pi,pi_old) (max violation, must be <= 0)\", max(0.0, worst), 0.0)\n", + "# variance: gradient weights are bounded for CISPO, unbounded for PPO's live A<0 tokens\n", + "r, _ = ratio_and_lp(th)\n", + "print(f\" max gradient weight over (o,t): PPO {float(weights(th,'ppo').max()):.3f} CISPO {float(weights(th,'cispo').max()):.3f} unclipped IS {float(weights(th,'is').max()):.3f}\")\n", + "assert float(weights(th, 'cispo').max()) <= 1 + eh + 1e-15\n", + "check(\"CISPO weights bounded by 1+eps_hi\", max(0.0, float(weights(th, 'cispo').max()) - (1 + eh)), 0.0)" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + " eps_hi=0.05: ||bias|| = 1.3732e-02 bound = 3.7511e-02 ||token-IS gradient|| = 1.7479e-01\n", + " eps_hi=0.10: ||bias|| = 9.2089e-03 bound = 2.3875e-02 ||token-IS gradient|| = 1.7479e-01\n", + " eps_hi=0.28: ||bias|| = 1.0099e-03 bound = 4.1370e-03 ||token-IS gradient|| = 1.7479e-01\n", + "CHECK|exact|bias bound holds (3 values of eps_hi)|0.000e+00|tol=0.0e+00|PASS\n", + " ||g_PPO - g_IS|| = 2.2235e-02 ||g_CISPO - g_IS|| = 1.0099e-03 (eps_hi=0.28, eps_lo=0.2; ||g_IS|| = 1.7479e-01)\n", + "CHECK|exact|per-state E_old[(r-1-eps)_+] <= TV(pi,pi_old) (max violation, must be <= 0)|0.000e+00|tol=0.0e+00|PASS\n", + " max gradient weight over (o,t): PPO 1.992 CISPO 1.280 unclipped IS 1.992\n", + "CHECK|exact|CISPO weights bounded by 1+eps_hi|0.000e+00|tol=0.0e+00|PASS\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 5. Stop-gradient semantics: three look-alike implementations, three different gradients\n", + " (a) correct: loss = - sg(min(r,1+eh)) * A * log pi -> grad = -A*min(r,1+eh) * grad log pi\n", + " (b) forgot sg: loss = - min(r,1+eh) * A * log pi -> on the UNclipped set grad = -A * r * (1 + log pi) * grad log pi : an extra A r log pi term\n", + " (c) ratio form: loss = - min(r,1+eh) * A (PPO-like) -> grad = -A * r * 1[r<1+eh] * grad log pi (zero beyond the clip!)" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "source": [ + "lp_old = torch.tensor(-2.3)\n", + "def three(lp_new, A, eps=0.28):\n", + " lp = torch.tensor(lp_new, requires_grad=True); r = (lp - lp_old).exp()\n", + " ga = torch.autograd.grad(-(torch.clamp(r, max=1 + eps).detach() * A * lp), lp, retain_graph=True)[0]\n", + " gb = torch.autograd.grad(-(torch.clamp(r, max=1 + eps) * A * lp), lp, retain_graph=True)[0]\n", + " gc = torch.autograd.grad(-(torch.clamp(r, max=1 + eps) * A), lp)[0]\n", + " return float(ga), float(gb), float(gc), float(r.detach())\n", + "print(\" lp_old=-2.3; gradient wrt log pi_theta\")\n", + "print(\" lp_new r A (a) CISPO (b) no sg (c) ratio-clip\")\n", + "for lpn, A in ((-2.3, 1.0), (-1.5, 1.0), (-1.5, -1.0), (-3.0, 1.0), (-0.5, 1.0), (-0.5, -1.0)):\n", + " ga, gb, gc, rr = three(lpn, A)\n", + " print(f\" {lpn:6.2f} {rr:6.2f} {A:+.0f} {ga:+9.3f} {gb:+9.3f} {gc:+9.3f}\")\n", + " expect_a = -A * min(rr, 1.28)\n", + " assert abs(ga - expect_a) < 1e-12\n", + " expect_c = -A * rr * (rr < 1.28)\n", + " assert abs(gc - expect_c) < 1e-12\n", + " expect_b = -A * rr * (1 + lpn) if rr < 1.28 else -A * 1.28 # clamp kills the extra term beyond the clip\n", + " assert abs(gb - expect_b) < 1e-12\n", + "check(\"(a) = -A min(r,1+eh); (b) = -A r (1+log pi) on the live set; (c) = -A r 1[r<1+eh] (6 test points, exact)\", 0.0, 0.0)\n", + "ga, gb, gc, rr = three(-2.3, 1.0)\n", + "print(f\" at r=1 (inside the clip), A=+1, log pi=-2.3: (a)={ga:+.3f} (b)={gb:+.3f} <- sign flipped by the spurious log-pi factor (1+log pi = -1.3)\")\n", + "assert ga < 0 < gb\n", + "check(\"without sg a live token's gradient weight becomes r(1+log pi): wrong sign when pi < 1/e\", 0.0, 0.0)\n", + "ga, gb, gc, rr = three(-1.5, 1.0)\n", + "print(f\" at r={rr:.2f} > 1+eps: (a)={ga:+.3f} capped weight kept; (b)={gb:+.3f} identical (clamp zeroes the extra term); (c)={gc:+.3f} silenced\")" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + " lp_old=-2.3; gradient wrt log pi_theta\n", + " lp_new r A (a) CISPO (b) no sg (c) ratio-clip\n", + " -2.30 1.00 +1 -1.000 +1.300 -1.000\n", + " -1.50 2.23 +1 -1.280 -1.280 +0.000\n", + " -1.50 2.23 -1 +1.280 +1.280 +0.000\n", + " -3.00 0.50 +1 -0.497 +0.993 -0.497\n", + " -0.50 6.05 +1 -1.280 -1.280 +0.000\n", + " -0.50 6.05 -1 +1.280 +1.280 +0.000\n", + "CHECK|exact|(a) = -A min(r,1+eh); (b) = -A r (1+log pi) on the live set; (c) = -A r 1[r<1+eh] (6 test points, exact)|0.000e+00|tol=0.0e+00|PASS\n", + " at r=1 (inside the clip), A=+1, log pi=-2.3: (a)=-1.000 (b)=+1.300 <- sign flipped by the spurious log-pi factor (1+log pi = -1.3)\n", + "CHECK|exact|without sg a live token's gradient weight becomes r(1+log pi): wrong sign when pi < 1/e|0.000e+00|tol=0.0e+00|PASS\n", + " at r=2.23 > 1+eps: (a)=-1.280 capped weight kept; (b)=-1.280 identical (clamp zeroes the extra term); (c)=+0.000 silenced\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 6. The 'fork token' mechanism over K = 16 off-policy updates on one batch (MiniMax-M1 setting)\n", + "One decision with three tokens: two common tokens and a RARE token w (pi_old = 0.01) with positive advantage.\n", + "Exact expectation over the old distribution; plain SGD on the logits; K = 16 inner steps." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "source": [ + "def fork(method, K=16, lr=30.0, eh_=0.28, el_=0.2, p_w=0.01):\n", + " pi_old = torch.tensor([0.60, 0.39 - 0.0, p_w]); pi_old = pi_old / pi_old.sum()\n", + " reward = torch.tensor([0.0, 0.0, 1.0]); A = reward - float(pi_old @ reward) # A_w > 0, A_common < 0\n", + " z = pi_old.log().clone().requires_grad_(True)\n", + " traj, gw = [], []\n", + " for k in range(K):\n", + " pi = torch.softmax(z, 0); r = (pi / pi_old)\n", + " if method == 'ppo':\n", + " f = torch.where(A > 0, A * torch.clamp(r, max=1 + eh_), A * torch.clamp(r, min=1 - el_))\n", + " obj = (pi_old * f).sum()\n", + " else:\n", + " rhat = torch.clamp(r, max=1 + eh_).detach()\n", + " obj = (pi_old * rhat * A * pi.log()).sum()\n", + " g = torch.autograd.grad(obj, z)[0]\n", + " # direct contribution of token w's OWN term to the gradient on its own logit (diagnostic)\n", + " rd = r.detach(); wt = float((1 + eh_) if (method == 'cispo' and rd[2] > 1 + eh_) else (rd[2] if (method == 'cispo' or rd[2] < 1 + eh_) else 0.0))\n", + " gw.append(wt)\n", + " with torch.no_grad(): z += lr * g\n", + " traj.append(float(torch.softmax(z.detach(), 0)[2]))\n", + " return traj, gw, float(r[2].detach())\n", + "tp, gwp, rp = fork('ppo'); tc, gwc, rc = fork('cispo')\n", + "print(\" after update #\" + \" \".join(f\"{k+1:6d}\" for k in range(0, 16, 3)))\n", + "print(\" p_w PPO: \" + \" \".join(f\"{tp[k]:6.3f}\" for k in range(0, 16, 3)))\n", + "print(\" p_w CISPO: \" + \" \".join(f\"{tc[k]:6.3f}\" for k in range(0, 16, 3)))\n", + "print(\" weight on w's own term PPO: \" + \" \".join(f\"{gwp[k]:6.2f}\" for k in range(0, 16, 3)))\n", + "print(\" weight on w's own term CISPO: \" + \" \".join(f\"{gwc[k]:6.2f}\" for k in range(0, 16, 3)))\n", + "n_zero_ppo = sum(1 for x in gwp if x == 0.0)\n", + "print(f\" PPO: token w gets zero direct gradient in {n_zero_ppo} of 16 inner steps; CISPO: {sum(1 for x in gwc if x==0.0)} of 16\")\n", + "assert n_zero_ppo >= 8 and all(x > 0 for x in gwc)\n", + "check(\"PPO silences the rare token for >= 8 of 16 inner steps, CISPO for none\", 0.0, 0.0)\n", + "print(f\" final p_w after 16 updates: PPO {tp[-1]:.3f} CISPO {tc[-1]:.3f}\")" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + " after update # 1 4 7 10 13 16\n", + " p_w PPO: 0.016 0.016 0.016 0.017 0.017 0.017\n", + " p_w CISPO: 0.016 0.081 0.310 0.624 0.790 0.862\n", + " weight on w's own term PPO: 1.00 0.00 0.00 0.00 0.00 0.00\n", + " weight on w's own term CISPO: 1.00 1.28 1.28 1.28 1.28 1.28\n", + " PPO: token w gets zero direct gradient in 15 of 16 inner steps; CISPO: 0 of 16\n", + "CHECK|exact|PPO silences the rare token for >= 8 of 16 inner steps, CISPO for none|0.000e+00|tol=0.0e+00|PASS\n", + " final p_w after 16 updates: PPO 0.017 CISPO 0.862\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 7. Naive vs efficient implementation of the CISPO loss, token-level normalisation by total tokens" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": {}, + "source": [ + "Gs = 8\n", + "gen = torch.Generator().manual_seed(5)\n", + "idx = torch.multinomial(q, Gs, replacement=True, generator=gen)\n", + "Rg = R[idx]; Ag = Rg - Rg.mean()\n", + "t = perturb(th_old, 0.8, seed=2).clone().requires_grad_(True)\n", + "lp = lm.token_logprobs(t)[idx]; lq_ = lq[idx]; mk = lm.mask[idx]\n", + "def cispo_naive(lp, lq_):\n", + " tot, n = 0.0, 0\n", + " for i in range(Gs):\n", + " for tt in range(int(mk[i].sum())):\n", + " r = math.exp(float((lp[i, tt] - lq_[i, tt]).detach()))\n", + " tot = tot + min(r, 1 + eh) * Ag[i] * lp[i, tt]; n += 1\n", + " return -tot / n\n", + "def cispo_vec(lp, lq_):\n", + " r = (lp - lq_).exp().detach()\n", + " return -((torch.clamp(r, max=1 + eh) * Ag[:, None] * lp * mk).sum() / mk.sum())\n", + "gn = torch.autograd.grad(cispo_naive(lp, lq_), t, retain_graph=True)[0]\n", + "gv = torch.autograd.grad(cispo_vec(lp, lq_), t, retain_graph=True)[0]\n", + "check(\"naive per-token loop == vectorised masked CISPO loss (gradient)\", maxabs(gn, gv), 1e-14)" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "CHECK|exact|naive per-token loop == vectorised masked CISPO loss (gradient)|0.000e+00|tol=1.0e-14|PASS\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 8. Leanstral variant (Mistral technical report, Sec. 2.3): sampler ratio, masking of tool/prompt tokens, outcome broadcast\n", + "Loss there: `L = -(1/sum_i|o_i|) sum_i sum_t sg(min(pi_theta/pi_sample, c)) * A_i * log pi_theta(a_it | s_it, T_i)`, `T_i` = tool calls and responses used as context only,\n", + "one verifier outcome for every model token. Toy: two-turn tabular agent (4 actions per turn), turn-2 context = turn-1 action, feedback token = a fixed function of turn 1 (not scored).\n", + "`R ~ Bern(T[y1, y2])`. Everything below is enumerated exactly over the 32 episodes (y1, y2, R)." + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "metadata": {}, + "source": [ + "import numpy as _np\n", + "_rng = _np.random.default_rng(123)\n", + "_A = 4\n", + "_T = _rng.uniform(0.05, 0.95, (_A, _A))\n", + "def _sm(x): e = _np.exp(x - x.max(-1, keepdims=True)); return e / e.sum(-1, keepdims=True)\n", + "def _unpack(th): return th[:_A], th[_A:].reshape(_A, _A)\n", + "th_cur = _rng.normal(0, 0.8, _A + _A * _A)\n", + "th_smp = th_cur + _rng.normal(0, 0.35, _A + _A * _A) # sampler = lagged + numerically perturbed copy\n", + "def _probs(th): a, b = _unpack(th); return _sm(a), _sm(b)\n", + "def _J(th): p1, p2 = _probs(th); return float((p1[:, None] * p2 * _T).sum())\n", + "def _score(th, y1, y2): # (s1, s2): grads of log pi(y1), log pi(y2|y1) wrt the 20 logits\n", + " p1, p2 = _probs(th); s1 = _np.zeros(_A + _A * _A); s2 = _np.zeros(_A + _A * _A)\n", + " s1[:_A] = _np.eye(_A)[y1] - p1\n", + " blk = _np.eye(_A)[y2] - p2[y1]; s2[_A + y1 * _A: _A + (y1 + 1) * _A] = blk\n", + " return s1, s2\n", + "def _episodes(th_sample): # list of (prob under sampler, y1, y2, R)\n", + " p1, p2 = _probs(th_sample); out = []\n", + " for y1 in range(_A):\n", + " for y2 in range(_A):\n", + " for r in (0, 1):\n", + " out.append((p1[y1] * p2[y1, y2] * (_T[y1, y2] if r else 1 - _T[y1, y2]), y1, y2, r))\n", + " return out\n", + "def _fd_grad(th, h=1e-6):\n", + " g = _np.zeros_like(th)\n", + " for i in range(len(th)):\n", + " e = _np.zeros_like(th); e[i] = h; g[i] = (_J(th + e) - _J(th - e)) / (2 * h)\n", + " return g\n", + "g_true = _fd_grad(th_cur)\n", + "pc1, pc2 = _probs(th_cur); ps1, ps2 = _probs(th_smp)\n", + "def _rho(y1, y2): return pc1[y1] / ps1[y1], pc2[y1, y2] / ps2[y1, y2]\n", + "# (a) on-policy score-function identity\n", + "g_on = sum(p * r * sum(_score(th_cur, y1, y2)) for p, y1, y2, r in _episodes(th_cur))\n", + "check(\"8a. E_pi[R (s1+s2)] = finite-difference grad J\", float(_np.abs(g_on - g_true).max()), 1e-8)\n", + "# (b) sampler episodes, full sequence ratio rho1*rho2: exact\n", + "g_seq = sum(p * _rho(y1, y2)[0] * _rho(y1, y2)[1] * r * sum(_score(th_cur, y1, y2)) for p, y1, y2, r in _episodes(th_smp))\n", + "check(\"8b. E_mu[rho1 rho2 R (s1+s2)] = grad J exactly (sampler ratio, full sequence weight)\", float(_np.abs(g_seq - g_true).max()), 1e-9)\n", + "# (c) token-level ratios only (what the objective uses): biased by the dropped prefix ratio; vanishes when sampler = policy\n", + "def _tok_grad(th_sample, clip=None, baseline=0.0):\n", + " g = 0\n", + " for p, y1, y2, r in _episodes(th_sample):\n", + " q1, q2 = _rho(y1, y2) if th_sample is th_smp else (1.0, 1.0)\n", + " if clip is not None: q1, q2 = min(q1, clip), min(q2, clip)\n", + " s1, s2 = _score(th_cur, y1, y2); g = g + p * (r - baseline) * (q1 * s1 + q2 * s2)\n", + " return g\n", + "g_tok = _tok_grad(th_smp)\n", + "bias_tok = float(_np.linalg.norm(g_tok - g_true)); gn = float(_np.linalg.norm(g_true))\n", + "print(f\" token-level sampler ratio, unclipped: ||bias|| = {bias_tok:.3e} (||grad J|| = {gn:.3e})\")\n", + "check(\"8c. token-level weights are biased for a mismatched sampler (> 1e-4)\", 0.0 if bias_tok > 1e-4 else 1.0, 0.0)\n", + "th_save = th_smp; th_smp = th_cur.copy() # sampler == policy: weights are 1\n", + "check(\"8c. ... and exact when the sampler equals the policy\", float(_np.abs(sum(p * r * sum(_score(th_cur, y1, y2)) for p, y1, y2, r in _episodes(th_cur)) - g_true).max()), 1e-8)\n", + "th_smp = th_save\n", + "# (d) a constant baseline does not bias the token-level estimator, even under mismatch (E_mu[rho_t s_t | context] = 0)\n", + "b0 = _J(th_cur)\n", + "g_b = _tok_grad(th_smp, baseline=b0)\n", + "db = _np.abs(g_tok - g_b - 0).max()\n", + "check(\"8d. constant baseline b adds exactly 0 to the token-level sampler-ratio gradient: max |E_mu[sum_t rho_t b s_t]| over 20 coordinates\", float(db), 1e-14)\n", + "# (e) clip bias bound with the sampler as the base measure: sum_a (pi - c mu)_+ <= TV(pi, mu) at every state\n", + "cc = 1.28\n", + "sts = [(pc1, ps1)] + [(pc2[y], ps2[y]) for y in range(_A)]\n", + "viol = max(float(_np.maximum(a - cc * b, 0).sum() - 0.5 * _np.abs(a - b).sum()) for a, b in sts)\n", + "check(\"8e. sum_a (pi - c mu)_+ <= TV(pi, mu) at all 5 states (c=1.28)\", max(viol, 0.0), 0.0)\n", + "g_clip = _tok_grad(th_smp, clip=cc)\n", + "bound = 0.0\n", + "for p, y1, y2, r in _episodes(th_smp):\n", + " q1, q2 = _rho(y1, y2); s1, s2 = _score(th_cur, y1, y2)\n", + " bound += p * abs(r) * (max(q1 - cc, 0) * _np.linalg.norm(s1) + max(q2 - cc, 0) * _np.linalg.norm(s2))\n", + "bias_clip = float(_np.linalg.norm(g_clip - g_tok))\n", + "print(f\" clip c={cc}: ||g_clip - g_token-IS|| = {bias_clip:.3e} <= bound {bound:.3e}\")\n", + "check(\"8e. clipping bias <= E_mu[sum_t (rho_t - c)_+ |A| ||s_t||]\", max(bias_clip - bound, 0.0), 0.0)\n", + "# (f) outcome broadcast vs per-turn credit (on-policy, oracle baselines): both unbiased, different variance\n", + "eps_ = _episodes(th_cur)\n", + "V = sum(p * r for p, y1, y2, r in eps_)\n", + "Q = _np.array([sum(pc2[y1, y2] * _T[y1, y2] for y2 in range(_A)) for y1 in range(_A)])\n", + "def _moment(coef1, coef2):\n", + " m1 = 0; m2 = 0\n", + " for p, y1, y2, r in eps_:\n", + " s1, s2 = _score(th_cur, y1, y2); g = coef1(y1, y2, r) * s1 + coef2(y1, y2, r) * s2\n", + " m1 = m1 + p * g; m2 += p * float(g @ g)\n", + " return m1, m2\n", + "gb_mean, gb_m2 = _moment(lambda y1, y2, r: r - V, lambda y1, y2, r: r - V)\n", + "gt_mean, gt_m2 = _moment(lambda y1, y2, r: Q[y1] - V, lambda y1, y2, r: r - Q[y1])\n", + "check(\"8f. broadcast estimator (R-V)(s1+s2) is unbiased for grad J\", float(_np.abs(gb_mean - g_true).max()), 1e-8)\n", + "check(\"8f. per-turn estimator (Q-V) s1 + (R-Q) s2 is unbiased for grad J\", float(_np.abs(gt_mean - g_true).max()), 1e-8)\n", + "v_b = gb_m2 - float(gb_mean @ gb_mean); v_t = gt_m2 - float(gt_mean @ gt_mean)\n", + "print(f\" total variance (trace): broadcast {v_b:.4f}, per-turn {v_t:.4f}, ratio {v_b/v_t:.3f}\")\n", + "# turn-1 decomposition: E[(R-V)^2 ||s1||^2] = E[(Q-V)^2 ||s1||^2] + E[Var(R|y1) ||s1||^2]\n", + "lhs = sum(p * (r - V) ** 2 * float(_score(th_cur, y1, y2)[0] @ _score(th_cur, y1, y2)[0]) for p, y1, y2, r in eps_)\n", + "t1 = sum(p * (Q[y1] - V) ** 2 * float(_score(th_cur, y1, y2)[0] @ _score(th_cur, y1, y2)[0]) for p, y1, y2, r in eps_)\n", + "varR = _np.array([sum(pc2[y1, y2] * (_T[y1, y2] + 0) for y2 in range(_A)) - Q[y1] ** 2 for y1 in range(_A)]) # Var(R|y1) = E[R^2|y1] - Q^2 (R binary: E R^2 = E R)\n", + "t2 = sum(p * varR[y1] * float(_score(th_cur, y1, y2)[0] @ _score(th_cur, y1, y2)[0]) for p, y1, y2, r in eps_)\n", + "check(\"8f. turn-1 second moment: E[(R-V)^2||s1||^2] = E[(Q-V)^2||s1||^2] + E[Var(R|y1)||s1||^2]\", abs(lhs - t1 - t2), 1e-13)\n", + "print(f\" turn-1: broadcast {lhs:.4f} = {t1:.4f} (credit) + {t2:.4f} (noise from turn 2)\")\n", + "check(\"8f. martingale orthogonality: E[s1.s2] = 0 so E||s1+s2||^2 = E||s1||^2 + E||s2||^2\", abs(sum(p * float(_score(th_cur, y1, y2)[0] @ _score(th_cur, y1, y2)[1]) for p, y1, y2, r in eps_)), 1e-15)\n", + "# (g) score second moment grows linearly with the number of policy tokens: L-step context-dependent chain, enumerated\n", + "Lmax = 5; Th = _rng.normal(0, 0.8, (Lmax, _A, _A)); P0 = _sm(_rng.normal(0, 0.8, _A))\n", + "def chain_moments(L):\n", + " tot = 0.0; per = _np.zeros(L); cross = 0.0\n", + " for seq in _np.ndindex(*([_A] * L)):\n", + " pr = P0[seq[0]]; sc = []\n", + " for t_ in range(L):\n", + " prev = 0 if t_ == 0 else seq[t_ - 1]\n", + " pt = P0 if t_ == 0 else _sm(Th[t_, prev])\n", + " if t_ > 0: pr *= pt[seq[t_]]\n", + " v = _np.zeros(L * _A * _A + _A)\n", + " v[:_A] = (_np.eye(_A)[seq[0]] - P0) if t_ == 0 else 0\n", + " if t_ > 0: v[_A + t_ * _A * _A + prev * _A: _A + t_ * _A * _A + (prev + 1) * _A] = _np.eye(_A)[seq[t_]] - pt\n", + " sc.append(v)\n", + " S_ = sum(sc); tot += pr * float(S_ @ S_)\n", + " for t_ in range(L): per[t_] += pr * float(sc[t_] @ sc[t_])\n", + " return tot, per\n", + "rows = []\n", + "for L in range(1, Lmax + 1):\n", + " tot, per = chain_moments(L); rows.append((L, tot, per.sum()))\n", + " print(f\" L={L}: E||sum_t s_t||^2 = {tot:.4f} sum_t E||s_t||^2 = {per.sum():.4f}\")\n", + "check(\"8g. E||sum_t s_t||^2 = sum_t E||s_t||^2 for L = 1..5 (score terms are martingale differences)\", max(abs(a - b) for _, a, b in rows), 1e-12)\n", + "print(\"done; checks:\", len(CHECKS))" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "CHECK|exact|8a. E_pi[R (s1+s2)] = finite-difference grad J|2.665e-11|tol=1.0e-08|PASS\n", + "CHECK|exact|8b. E_mu[rho1 rho2 R (s1+s2)] = grad J exactly (sampler ratio, full sequence weight)|2.665e-11|tol=1.0e-09|PASS\n", + " token-level sampler ratio, unclipped: ||bias|| = 8.068e-03 (||grad J|| = 1.392e-01)\n", + "CHECK|exact|8c. token-level weights are biased for a mismatched sampler (> 1e-4)|0.000e+00|tol=0.0e+00|PASS\n", + "CHECK|exact|8c. ... and exact when the sampler equals the policy|2.665e-11|tol=1.0e-08|PASS\n", + "CHECK|exact|8d. constant baseline b adds exactly 0 to the token-level sampler-ratio gradient: max |E_mu[sum_t rho_t b s_t]| over 20 coordinates|2.776e-17|tol=1.0e-14|PASS\n", + "CHECK|exact|8e. sum_a (pi - c mu)_+ <= TV(pi, mu) at all 5 states (c=1.28)|0.000e+00|tol=0.0e+00|PASS\n", + " clip c=1.28: ||g_clip - g_token-IS|| = 2.023e-02 <= bound 3.210e-02\n", + "CHECK|exact|8e. clipping bias <= E_mu[sum_t (rho_t - c)_+ |A| ||s_t||]|0.000e+00|tol=0.0e+00|PASS\n", + "CHECK|exact|8f. broadcast estimator (R-V)(s1+s2) is unbiased for grad J|2.665e-11|tol=1.0e-08|PASS\n", + "CHECK|exact|8f. per-turn estimator (Q-V) s1 + (R-Q) s2 is unbiased for grad J|2.665e-11|tol=1.0e-08|PASS\n", + " total variance (trace): broadcast 0.3020, per-turn 0.1473, ratio 2.050\n", + "CHECK|exact|8f. turn-1 second moment: E[(R-V)^2||s1||^2] = E[(Q-V)^2||s1||^2] + E[Var(R|y1)||s1||^2]|5.551e-17|tol=1.0e-13|PASS\n", + " turn-1: broadcast 0.1641 = 0.0381 (credit) + 0.1260 (noise from turn 2)\n", + "CHECK|exact|8f. martingale orthogonality: E[s1.s2] = 0 so E||s1+s2||^2 = E||s1||^2 + E||s2||^2|0.000e+00|tol=1.0e-15|PASS\n", + " L=1: E||sum_t s_t||^2 = 0.7188 sum_t E||s_t||^2 = 0.7188\n", + " L=2: E||sum_t s_t||^2 = 1.4115 sum_t E||s_t||^2 = 1.4115\n", + " L=3: E||sum_t s_t||^2 = 2.0546 sum_t E||s_t||^2 = 2.0546\n", + " L=4: E||sum_t s_t||^2 = 2.6447 sum_t E||s_t||^2 = 2.6447\n", + " L=5: E||sum_t s_t||^2 = 3.2724 sum_t E||s_t||^2 = 3.2724\n", + "CHECK|exact|8g. E||sum_t s_t||^2 = sum_t E||s_t||^2 for L = 1..5 (score terms are martingale differences)|5.329e-15|tol=1.0e-12|PASS\n", + "done; checks: 28\n" + ] + } + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.12.11" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} \ No newline at end of file diff --git a/code/notebooks/0108-training-inference-mismatch.ipynb b/code/notebooks/0108-training-inference-mismatch.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..4cd940d07e53df3cd69321449dd78653ca08b9dc --- /dev/null +++ b/code/notebooks/0108-training-inference-mismatch.ipynb @@ -0,0 +1,585 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# 0108 - Training-inference mismatch: sampler vs trainer log-probs, truncated IS, why an FP32 LM head helps\n", + "Companion to `notes/0108 - Training-Inference Mismatch.pdf`. fp64 reference (\"trainer\" pi); the \"sampler\" mu uses logits rounded to bf16/fp16/fp32." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "source": [ + "import sys, math\n", + "sys.path.insert(0, \"..\")\n", + "from _m1_lib import *\n", + "import numpy as np" + ], + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 0. Floating-point facts: unit roundoff u = 2^-(p) with p significant bits; |round(x) - x| <= u |x|\n", + "bf16: 8 significant bits (u=2^-8), fp16: 11 (u=2^-11), fp32: 24 (u=2^-24)." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "source": [ + "gen = torch.Generator().manual_seed(0)\n", + "x = torch.exp(torch.rand(2_000_000, generator=gen) * math.log(1e3 / 1e-3) + math.log(1e-3)) # log-uniform in [1e-3, 1e3]\n", + "for name, dt, u in ((\"bf16\", torch.bfloat16, 2.0 ** -8), (\"fp16\", torch.float16, 2.0 ** -11), (\"fp32\", torch.float32, 2.0 ** -24)):\n", + " rel = ((x.to(dt).to(torch.float64) - x).abs() / x)\n", + " print(f\" {name}: max relative rounding error {float(rel.max()):.3e} u = {u:.3e}\")\n", + " check(f\"{name}: max rel. rounding error <= u (=2^-{int(-math.log2(u))})\", max(0.0, float(rel.max()) - u), 0.0)\n", + " assert float(rel.max()) > 0.5 * u" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + " bf16: max relative rounding error 3.891e-03 u = 3.906e-03\n", + "CHECK|exact|bf16: max rel. rounding error <= u (=2^-8)|0.000e+00|tol=0.0e+00|PASS\n", + " fp16: max relative rounding error 4.880e-04 u = 4.883e-04\n", + "CHECK|exact|fp16: max rel. rounding error <= u (=2^-11)|0.000e+00|tol=0.0e+00|PASS\n", + " fp32: max relative rounding error 5.958e-08 u = 5.960e-08\n", + "CHECK|exact|fp32: max rel. rounding error <= u (=2^-24)|0.000e+00|tol=0.0e+00|PASS\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 1. The tiny LM with LARGE logits (like an LM head with high-magnitude activations); sampler mu = same network with bf16-rounded logits" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "source": [ + "lm = TinyLM(V=4, L=4, seed=0, scale=4.0)\n", + "theta = lm.theta0.clone()\n", + "R = default_reward(lm)\n", + "def seq_logprob_rounded(theta, dtype):\n", + " V, L = lm.V, lm.L\n", + " tok, pos = lm._logits_tables(theta)\n", + " prev = torch.cat([torch.full((lm.N, 1), V, dtype=torch.long), lm.tokens[:, :-1]], 1)\n", + " logits = tok[prev] + pos[None, :, :]\n", + " if dtype is not None: logits = logits.to(dtype).to(torch.float64)\n", + " lp = torch.log_softmax(logits, -1).gather(-1, lm.tokens[..., None]).squeeze(-1) * lm.mask\n", + " return lp.sum(1)\n", + "pi = seq_logprob_rounded(theta, None).exp()\n", + "mu = seq_logprob_rounded(theta, torch.bfloat16).exp()\n", + "check(\"trainer pi is normalised\", abs(float(pi.sum()) - 1), 1e-14)\n", + "check(\"sampler mu is normalised (bf16 logits then softmax)\", abs(float(mu.sum()) - 1), 1e-14)\n", + "b = float(pi @ R); A = R - b\n", + "S = score_matrix(lm, theta) # score of the TRAINER policy\n", + "g_true = (pi[:, None] * A[:, None] * S).sum(0) # grad J_pi\n", + "w = pi / mu\n", + "TVs = 0.5 * float((pi - mu).abs().sum())\n", + "print(f\" bf16 logits: TV(pi, mu) = {TVs:.5f} chi2(pi||mu) = {float((mu*(w-1)**2).sum()):.3e} max w = {float(w.max()):.4f}\")\n", + "for nm, dt in ((\"fp16\", torch.float16), (\"fp32\", torch.float32)):\n", + " m_ = seq_logprob_rounded(theta, dt).exp()\n", + " print(f\" {nm} logits: TV(pi, mu) = {0.5*float((pi-m_).abs().sum()):.3e} chi2 = {float((m_*(pi/m_-1)**2).sum()):.3e}\")" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "CHECK|exact|trainer pi is normalised|0.000e+00|tol=1.0e-14|PASS\n", + "CHECK|exact|sampler mu is normalised (bf16 logits then softmax)|0.000e+00|tol=1.0e-14|PASS\n", + " bf16 logits: TV(pi, mu) = 0.00045 chi2(pi||mu) = 5.187e-06 max w = 1.0491\n", + " fp16 logits: TV(pi, mu) = 1.195e-04 chi2 = 2.142e-07\n", + " fp32 logits: TV(pi, mu) = 4.361e-09 chi2 = 1.076e-15\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 2. Ignoring the mismatch is biased; IS repairs it exactly (arXiv:2510.26788 Eqs. 3, 5)" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "source": [ + "g_biased = (mu[:, None] * A[:, None] * S).sum(0) # samples from mu, trainer's score, no correction\n", + "bias_formula = ((mu - pi)[:, None] * A[:, None] * S).sum(0) # = E_pi[(mu/pi - 1) A s]\n", + "check(\"bias of the uncorrected estimator == sum_o (mu-pi) A s\", maxabs(g_biased - g_true, bias_formula), 1e-14)\n", + "g_is = (mu[:, None] * w[:, None] * A[:, None] * S).sum(0)\n", + "check(\"sequence-level IS correction E_mu[(pi/mu) A s] == grad J_pi exactly\", maxabs(g_is, g_true), 1e-14)\n", + "print(f\" ||bias|| / ||grad J|| = {float((g_biased-g_true).norm()/g_true.norm()):.4f} (uncorrected)\")\n", + "assert float((g_biased - g_true).norm()) > 1e-4" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "CHECK|exact|bias of the uncorrected estimator == sum_o (mu-pi) A s|1.930e-17|tol=1.0e-14|PASS\n", + "CHECK|exact|sequence-level IS correction E_mu[(pi/mu) A s] == grad J_pi exactly|6.939e-18|tol=1.0e-14|PASS\n", + " ||bias|| / ||grad J|| = 0.0011 (uncorrected)\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 3. Truncated IS (TIS, Eq. 6) and masked IS (MIS, Eq. 7): exact bias and bounds\n", + " TIS: E_mu[min(w,C) A s]; bias = - E_mu[(w-C)_+ A s] ||.|| <= M_A M_s sum_o (pi - C mu)_+ <= M_A M_s min(TV, P_pi(w>C))\n", + " MIS: E_mu[1(w<=C) w A s]; bias = - E_pi[1(w>C) A s] ||.|| <= M_A M_s P_pi(w>C)" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "source": [ + "MA, Ms = float(A.abs().max()), float(S.norm(dim=1).max())\n", + "print(\" C ||TIS bias|| bound1(sum mu (w-C)+ |A| ||s||) bound2(MA Ms sum(pi-C mu)+) MA Ms TV ||MIS bias|| MA Ms P_pi(w>C)\")\n", + "for C in (1.0, 1.02, 1.05, 1.2, 2.0):\n", + " wt = torch.clamp(w, max=C)\n", + " g_tis = (mu[:, None] * wt[:, None] * A[:, None] * S).sum(0)\n", + " bias_tis = g_tis - g_true\n", + " formula = -((mu * (w - C).clamp(min=0))[:, None] * A[:, None] * S).sum(0)\n", + " check(f\"TIS bias == -E_mu[(w-C)_+ A s] (C={C})\", maxabs(bias_tis, formula), 1e-14)\n", + " b1 = float((mu * (w - C).clamp(min=0) * A.abs() * S.norm(dim=1)).sum())\n", + " ex = float((pi - C * mu).clamp(min=0).sum())\n", + " b2 = MA * Ms * ex\n", + " assert float(bias_tis.norm()) <= b1 + 1e-15 and b1 <= b2 + 1e-15 and ex <= TVs + 1e-15\n", + " g_mis = (mu[:, None] * (w * (w <= C))[:, None] * A[:, None] * S).sum(0)\n", + " bias_mis = g_mis - g_true\n", + " mis_formula = -((pi * (w > C))[:, None] * A[:, None] * S).sum(0)\n", + " check(f\"MIS bias == -E_pi[1(w>C) A s] (C={C})\", maxabs(bias_mis, mis_formula), 1e-14)\n", + " Pw = float((pi * (w > C)).sum())\n", + " assert float(bias_mis.norm()) <= MA * Ms * Pw + 1e-15\n", + " print(f\" {C:4.2f} {float(bias_tis.norm()):12.4e} {b1:12.4e} {b2:12.4e} {MA*Ms*TVs:10.4e} {float(bias_mis.norm()):12.4e} {MA*Ms*Pw:12.4e}\")\n", + "check(\"TIS: 0 <= bias <= bound1 <= bound2 <= M_A M_s TV(pi,mu) for all C\", 0.0, 0.0)\n", + "check(\"TIS at C=inf is exact (no truncation) -> bias = 0\", float((((mu[:, None] * torch.clamp(w, max=1e12)[:, None] * A[:, None] * S).sum(0)) - g_true).norm()), 1e-13)\n", + "# TV of sequence distributions <= sum over positions of E[TV_t] (hybrid argument)\n", + "def tv_chain(theta, dtype):\n", + " tot = 0.0\n", + " V, L = lm.V, lm.L\n", + " tok, pos = lm._logits_tables(theta)\n", + " def cond(prefix, rounder):\n", + " prev = V if len(prefix) == 0 else prefix[-1]\n", + " z = tok[prev] + pos[len(prefix)]\n", + " if rounder is not None: z = z.to(rounder).to(torch.float64)\n", + " return torch.softmax(z, 0)\n", + " def rec(prefix, prob):\n", + " nonlocal tot\n", + " if len(prefix) >= L: return\n", + " cp, cm = cond(prefix, None), cond(prefix, dtype)\n", + " tot += prob * 0.5 * float((cp - cm).abs().sum())\n", + " for a in range(V - 1): rec(prefix + (a,), prob * float(cp[a]))\n", + " rec((), 1.0)\n", + " return tot\n", + "tvc = tv_chain(theta, torch.bfloat16)\n", + "print(f\" TV(sequence) = {TVs:.4f} <= sum_t E_pi[TV_t] = {tvc:.4f}\")\n", + "check(\"sequence TV <= sum_t E[TV_t] (max violation, must be <= 0)\", max(0.0, TVs - tvc), 0.0)" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + " C ||TIS bias|| bound1(sum mu (w-C)+ |A| ||s||) bound2(MA Ms sum(pi-C mu)+) MA Ms TV ||MIS bias|| MA Ms P_pi(w>C)\n", + "CHECK|exact|TIS bias == -E_mu[(w-C)_+ A s] (C=1.0)|2.557e-17|tol=1.0e-14|PASS\n", + "CHECK|exact|MIS bias == -E_pi[1(w>C) A s] (C=1.0)|6.939e-18|tol=1.0e-14|PASS\n", + " 1.00 1.1248e-04 1.6072e-04 2.2052e-03 2.2052e-03 1.2226e-01 4.0596e+00\n", + "CHECK|exact|TIS bias == -E_mu[(w-C)_+ A s] (C=1.02)|2.005e-17|tol=1.0e-14|PASS\n", + "CHECK|exact|MIS bias == -E_pi[1(w>C) A s] (C=1.02)|1.562e-17|tol=1.0e-14|PASS\n", + " 1.02 6.4034e-08 6.6273e-08 1.6801e-06 2.2052e-03 4.3457e-05 1.1315e-03\n", + "CHECK|exact|TIS bias == -E_mu[(w-C)_+ A s] (C=1.05)|6.939e-18|tol=1.0e-14|PASS\n", + "CHECK|exact|MIS bias == -E_pi[1(w>C) A s] (C=1.05)|6.939e-18|tol=1.0e-14|PASS\n", + " 1.05 1.7358e-17 0.0000e+00 0.0000e+00 2.2052e-03 1.7358e-17 0.0000e+00\n", + "CHECK|exact|TIS bias == -E_mu[(w-C)_+ A s] (C=1.2)|6.939e-18|tol=1.0e-14|PASS\n", + "CHECK|exact|MIS bias == -E_pi[1(w>C) A s] (C=1.2)|6.939e-18|tol=1.0e-14|PASS\n", + " 1.20 1.7358e-17 0.0000e+00 0.0000e+00 2.2052e-03 1.7358e-17 0.0000e+00\n", + "CHECK|exact|TIS bias == -E_mu[(w-C)_+ A s] (C=2.0)|6.939e-18|tol=1.0e-14|PASS\n", + "CHECK|exact|MIS bias == -E_pi[1(w>C) A s] (C=2.0)|6.939e-18|tol=1.0e-14|PASS\n", + " 2.00 1.7358e-17 0.0000e+00 0.0000e+00 2.2052e-03 1.7358e-17 0.0000e+00\n", + "CHECK|exact|TIS: 0 <= bias <= bound1 <= bound2 <= M_A M_s TV(pi,mu) for all C|0.000e+00|tol=0.0e+00|PASS\n", + "CHECK|exact|TIS at C=inf is exact (no truncation) -> bias = 0|1.736e-17|tol=1.0e-13|PASS\n", + " TV(sequence) = 0.0004 <= sum_t E_pi[TV_t] = 0.0007\n", + "CHECK|exact|sequence TV <= sum_t E[TV_t] (max violation, must be <= 0)|0.000e+00|tol=0.0e+00|PASS\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 4. From logit rounding error to log-prob error: d log pi_a = e_a - sum_b pi_b e_b (first order); Var = s^2 (1 - 2 pi_a + sum pi_b^2) for iid logit errors" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "source": [ + "gen = torch.Generator().manual_seed(3)\n", + "z = torch.randn(7, generator=gen) * 5.0\n", + "pi_z = torch.softmax(z, 0)\n", + "eps = torch.randn(7, generator=gen) * 1e-6\n", + "dlp_exact = torch.log_softmax(z + eps, 0) - torch.log_softmax(z, 0)\n", + "dlp_lin = eps - (pi_z * eps).sum()\n", + "check(\"first-order log-prob error e_a - sum_b pi_b e_b (relative error of linearisation, e~1e-6)\", float((dlp_exact - dlp_lin).abs().max() / dlp_lin.abs().max()), 1e-5)\n", + "s_ = 0.01\n", + "E = torch.randn(400000, 7, generator=gen) * s_\n", + "dlp = torch.log_softmax(z + E, -1) - torch.log_softmax(z, 0)\n", + "var_mc = dlp.var(0); var_f = s_ ** 2 * (1 - 2 * pi_z + (pi_z ** 2).sum())\n", + "zs = float(((var_mc - var_f).abs() / (var_f * math.sqrt(2 / 400000))).max())\n", + "check(\"Var(d log pi_a) = s^2 (1 - 2 pi_a + sum pi^2): max z-score over 7 tokens (MC, 400k draws)\", zs, 4.0, kind=\"mc\")" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "CHECK|exact|first-order log-prob error e_a - sum_b pi_b e_b (relative error of linearisation, e~1e-6)|5.237e-08|tol=1.0e-05|PASS\n", + "CHECK|mc|Var(d log pi_a) = s^2 (1 - 2 pi_a + sum pi^2): max z-score over 7 tokens (MC, 400k draws)|1.196e+00|tol=4.0e+00|PASS\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 4b. Where the logit error comes from: an LM-head matmul with bf16 inputs vs an fp32 head\n", + "z_b = sum_i h_i W_bi. Rounding h and W to bf16 perturbs each product by a relative error <= u, so the logit error scales with the\n", + "magnitude of the terms (not just of z). Outlier activations (a few huge hidden dims) enlarge that magnitude." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "source": [ + "gen = torch.Generator().manual_seed(4)\n", + "d, Vv, Tn = 1024, 64, 512\n", + "def head_errors(outlier):\n", + " h = torch.randn(Tn, d, generator=gen)\n", + " h[:, :8] *= outlier # 8 high-magnitude hidden dimensions\n", + " W = torch.randn(Vv, d, generator=gen) / math.sqrt(d)\n", + " z = h @ W.T # fp64 reference\n", + " zb = h.to(torch.bfloat16).to(torch.float64) @ W.to(torch.bfloat16).to(torch.float64).T # bf16-rounded inputs, exact accumulation\n", + " z32 = h.to(torch.float32).to(torch.float64) @ W.to(torch.float32).to(torch.float64).T # fp32-rounded inputs\n", + " return float(z.std()), float((zb - z).std()), float((z32 - z).std())\n", + "print(\" outlier x std(z) std(err bf16) std(err fp32) err_bf16/std(z) /u_bf16 ratio bf16/fp32\")\n", + "for out in (1.0, 10.0, 100.0):\n", + " sz_, eb, e32 = head_errors(out)\n", + " print(f\" {out:7.0f} {sz_:8.2f} {eb:12.3e} {e32:12.3e} {eb/sz_:10.3e} {eb/sz_/2**-8:6.3f} {eb/e32:9.1f}\")\n", + "sz_, eb, e32 = head_errors(100.0)\n", + "check(\"bf16 head: logit error / std(z) between 0.05u and 1u (u=2^-8)\", 0.0 if 0.05 * 2**-8 < eb / sz_ < 2**-8 else 1.0, 0.0)\n", + "check(\"fp32 head cuts the logit error by ~ u_bf16/u_fp32 = 65536 (log10 ratio within 1 decade)\", abs(math.log10(eb / e32) - math.log10(2**16)), 1.0)" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + " outlier x std(z) std(err bf16) std(err fp32) err_bf16/std(z) /u_bf16 ratio bf16/fp32\n", + " 1 1.01 2.366e-03 3.575e-08 2.349e-03 0.601 66166.0\n", + " 10 1.34 3.054e-03 4.786e-08 2.284e-03 0.585 63825.0\n", + " 100 8.81 2.025e-02 3.259e-07 2.298e-03 0.588 62139.1\n", + "CHECK|exact|bf16 head: logit error / std(z) between 0.05u and 1u (u=2^-8)|0.000e+00|tol=0.0e+00|PASS\n", + "CHECK|exact|fp32 head cuts the logit error by ~ u_bf16/u_fp32 = 65536 (log10 ratio within 1 decade)|2.678e-02|tol=1.0e+00|PASS\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 5. Long horizons: positions are independent contexts, so everything is EXACT:\n", + " E_mu[log w_T] = -sum_t KL(mu_t||pi_t), Var(log w_T) = sum_t Var_t, E_mu[w_T^2] = prod_t (1 + chi2_t). (V=64, logits ~ N(0, sigma_z^2))" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": {}, + "source": [ + "def per_position_stats(sigma_z, dtype, T=4096, V=64, seed=0):\n", + " g = torch.Generator().manual_seed(seed)\n", + " z = torch.randn(T, V, generator=g) * sigma_z\n", + " pi = torch.softmax(z, -1)\n", + " zr = z if dtype is None else z.to(dtype).to(torch.float64)\n", + " mu = torch.softmax(zr, -1)\n", + " lw = pi.log() - mu.log() # log(pi/mu) per token option\n", + " Elw = (mu * lw).sum(-1) # E_mu[log w_t] = -KL(mu||pi)\n", + " Vlw = (mu * lw ** 2).sum(-1) - Elw ** 2\n", + " chi2 = (mu * (pi / mu - 1) ** 2).sum(-1)\n", + " return Elw, Vlw, chi2, pi, mu\n", + "print(\" sigma_z head mean KL/token mean chi2/token T=100: ESS frac T=1000 T=4000 (ESS frac = 1/prod(1+chi2))\")\n", + "rows = {}\n", + "for sz in (2.0, 5.0, 10.0, 20.0):\n", + " for name, dt in ((\"bf16\", torch.bfloat16), (\"fp16\", torch.float16), (\"fp32\", torch.float32)):\n", + " Elw, Vlw, chi2, pi_t, mu_t = per_position_stats(sz, dt)\n", + " lESS = lambda T: float(-torch.log1p(chi2[:T]).sum())\n", + " rows[(sz, name)] = (float((-Elw).mean()), float(chi2.mean()), lESS(100), lESS(1000), lESS(4000))\n", + " print(f\" {sz:5.1f} {name} {float((-Elw).mean()):14.3e} {float(chi2.mean()):14.3e} {math.exp(lESS(100)):10.3e} {math.exp(lESS(1000)):10.3e} {math.exp(lESS(4000)):10.3e}\")\n", + "# why chi2 does not grow like sigma_z^2: to first order chi2_t = Var_{a~pi}(e_a) where e = round(z) - z (the *centred* logit error)\n", + "def first_order_check(sz, dtype, u, T=2048, V=64):\n", + " g_ = torch.Generator().manual_seed(0)\n", + " z_ = torch.randn(T, V, generator=g_) * sz\n", + " e_ = z_.to(dtype).to(torch.float64) - z_\n", + " pi_ = torch.softmax(z_, -1); mu_ = torch.softmax(z_ + e_, -1)\n", + " chi2_ = (mu_ * (pi_ / mu_ - 1) ** 2).sum(-1)\n", + " var_e = (pi_ * e_ ** 2).sum(-1) - ((pi_ * e_).sum(-1)) ** 2\n", + " bound = u ** 2 * (pi_ * z_ ** 2).sum(-1)\n", + " return float(chi2_.sum() / var_e.sum()), float((var_e / bound).max()), float(chi2_.mean())\n", + "for sz in (2.0, 10.0, 20.0):\n", + " ratio, bnd, c2 = first_order_check(sz, torch.bfloat16, 2.0 ** -8)\n", + " print(f\" sigma_z={sz:5.1f}: chi2/Var_pi(e) = {ratio:.4f} max Var_pi(e)/(u^2 E_pi z^2) = {bnd:.3f}\")\n", + " check(f\"first-order law chi2_t = Var_pi(logit rounding error) (sigma_z={sz}, mean ratio - 1)\", abs(ratio - 1), 0.1)\n", + " check(f\"bound Var_pi(e) <= u^2 E_pi[z^2] (sigma_z={sz}, max ratio - 1 must be <= 0)\", max(0.0, bnd - 1), 0.0)\n", + "szs = np.array([2.0, 5.0, 10.0, 20.0]); chis = np.array([rows[(s, 'bf16')][1] for s in szs])\n", + "slope = np.polyfit(np.log(szs), np.log(chis), 1)[0]\n", + "print(f\" bf16 log-log slope of chi2/token vs logit scale: {slope:.3f} (between 1 and 2: the error grows like |z| but the centred error Var_pi(e) shrinks as pi becomes peaked)\")\n", + "assert 1.0 < slope < 2.0\n", + "ratio_fp32 = rows[(10.0, 'fp32')][1] / rows[(10.0, 'bf16')][1]\n", + "print(f\" fp32/bf16 chi2 ratio at sigma_z=10: {ratio_fp32:.2e} (theory (u32/u16)^2 = {(2.0**-24/2.0**-8)**2:.2e})\")\n", + "check(\"fp32 head reduces per-token chi^2 by ~ (u_fp32/u_bf16)^2 (|log10 ratio - log10 theory|)\", abs(math.log10(ratio_fp32) - math.log10((2.0 ** -24 / 2.0 ** -8) ** 2)), 1.0)\n", + "# MC check of one exact statistic: Var(log w_T) over sampled sequences for T=64\n", + "Elw, Vlw, chi2, pi_t, mu_t = per_position_stats(10.0, torch.bfloat16, T=64)\n", + "g = torch.Generator().manual_seed(9); n = 20000\n", + "tokens = torch.stack([torch.multinomial(mu_t[t], n, replacement=True, generator=g) for t in range(64)], 1) # [n, T]\n", + "lw_seq = (pi_t.log().gather(1, tokens.T).T - mu_t.log().gather(1, tokens.T).T).sum(1)\n", + "exact_mean, exact_var = float(Elw.sum()), float(Vlw.sum())\n", + "se_mean = math.sqrt(exact_var / n)\n", + "check(\"MC: E_mu[log w_64] == -sum KL (z-score)\", abs(float(lw_seq.mean()) - exact_mean) / se_mean, 4.0, kind=\"mc\")\n", + "check(\"MC: Var(log w_64) == sum_t Var_t (z-score; kurtosis ignored -> loose 6)\", abs(float(lw_seq.var()) - exact_var) / (exact_var * math.sqrt(2 / n)), 6.0, kind=\"mc\")\n", + "check(\"E_mu[w_T] = 1 exactly per position (sum_a pi_a = 1)\", float((mu_t * (pi_t / mu_t)).sum(-1).sub(1).abs().max()), 1e-13)" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + " sigma_z head mean KL/token mean chi2/token T=100: ESS frac T=1000 T=4000 (ESS frac = 1/prod(1+chi2))\n", + " 2.0 bf16 1.420e-05 2.841e-05 9.967e-01 9.716e-01 8.925e-01\n", + " 2.0 fp16 2.179e-07 4.358e-07 1.000e+00 9.996e-01 9.983e-01\n", + " 2.0 fp32 3.269e-15 6.533e-15 1.000e+00 1.000e+00 1.000e+00\n", + " 5.0 bf16 6.224e-05 1.245e-04 9.895e-01 8.878e-01 6.068e-01\n", + " 5.0 fp16 9.391e-07 1.878e-06 9.998e-01 9.980e-01 9.925e-01\n", + " 5.0 fp32 1.424e-14 2.849e-14 1.000e+00 1.000e+00 1.000e+00\n", + " 10.0 bf16 1.464e-04 2.929e-04 9.774e-01 7.601e-01 3.091e-01\n", + " 10.0 fp16 2.135e-06 4.270e-06 9.996e-01 9.956e-01 9.831e-01\n", + " 10.0 fp32 3.233e-14 6.466e-14 1.000e+00 1.000e+00 1.000e+00\n", + " 20.0 bf16 3.125e-04 6.249e-04 9.580e-01 5.596e-01 8.265e-02\n", + " 20.0 fp16 4.395e-06 8.792e-06 9.991e-01 9.914e-01 9.656e-01\n", + " 20.0 fp32 6.719e-14 1.344e-13 1.000e+00 1.000e+00 1.000e+00\n", + " sigma_z= 2.0: chi2/Var_pi(e) = 1.0000 max Var_pi(e)/(u^2 E_pi z^2) = 0.590\n", + "CHECK|exact|first-order law chi2_t = Var_pi(logit rounding error) (sigma_z=2.0, mean ratio - 1)|1.799e-05|tol=1.0e-01|PASS\n", + "CHECK|exact|bound Var_pi(e) <= u^2 E_pi[z^2] (sigma_z=2.0, max ratio - 1 must be <= 0)|0.000e+00|tol=0.0e+00|PASS\n", + " sigma_z= 10.0: chi2/Var_pi(e) = 1.0005 max Var_pi(e)/(u^2 E_pi z^2) = 0.709\n", + "CHECK|exact|first-order law chi2_t = Var_pi(logit rounding error) (sigma_z=10.0, mean ratio - 1)|5.085e-04|tol=1.0e-01|PASS\n", + "CHECK|exact|bound Var_pi(e) <= u^2 E_pi[z^2] (sigma_z=10.0, max ratio - 1 must be <= 0)|0.000e+00|tol=0.0e+00|PASS\n", + " sigma_z= 20.0: chi2/Var_pi(e) = 1.0021 max Var_pi(e)/(u^2 E_pi z^2) = 0.644\n", + "CHECK|exact|first-order law chi2_t = Var_pi(logit rounding error) (sigma_z=20.0, mean ratio - 1)|2.100e-03|tol=1.0e-01|PASS\n", + "CHECK|exact|bound Var_pi(e) <= u^2 E_pi[z^2] (sigma_z=20.0, max ratio - 1 must be <= 0)|0.000e+00|tol=0.0e+00|PASS\n", + " bf16 log-log slope of chi2/token vs logit scale: 1.342 (between 1 and 2: the error grows like |z| but the centred error Var_pi(e) shrinks as pi becomes peaked)\n", + " fp32/bf16 chi2 ratio at sigma_z=10: 2.21e-10 (theory (u32/u16)^2 = 2.33e-10)\n", + "CHECK|exact|fp32 head reduces per-token chi^2 by ~ (u_fp32/u_bf16)^2 (|log10 ratio - log10 theory|)|2.316e-02|tol=1.0e+00|PASS\n", + "CHECK|mc|MC: E_mu[log w_64] == -sum KL (z-score)|4.121e-01|tol=4.0e+00|PASS\n", + "CHECK|mc|MC: Var(log w_64) == sum_t Var_t (z-score; kurtosis ignored -> loose 6)|5.377e-01|tol=6.0e+00|PASS\n", + "CHECK|exact|E_mu[w_T] = 1 exactly per position (sum_a pi_a = 1)|6.661e-16|tol=1.0e-13|PASS\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 5b. Caveat: E_mu[w_T^2] = prod_t (1 + chi2_t) needs history-INDEPENDENT chi2_t. With state-dependent chi2 the exact statements are\n", + " (i) the recursion E[w_T^2] = E[w_{T-1}^2 (1 + chi2_T(s_T))], (ii) the martingale identity E[ w_T^2 / prod_t (1 + chi2_t(s_t)) ] = 1,\n", + "and E[prod_t (1 + chi2_t(s_t))] is NOT E[w_T^2]. Exact enumeration on a tiny autoregressive model (sampler mu = one weight vector, trainer pi = a perturbed one)." + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "metadata": {}, + "source": [ + "lmx = TinyLM(V=3, L=8, seed=4, scale=1.0)\n", + "tokx = lmx.theta0[:lmx.n_tok].reshape(lmx.V + 1, lmx.V); tokx[:, lmx.EOS] -= 2.0\n", + "th_mu = lmx.theta0.clone(); th_pi = perturb(th_mu, 0.8, seed=9)\n", + "lmu, lpi = lmx.token_logprobs(th_mu), lmx.token_logprobs(th_pi)\n", + "mu_x = lmu.sum(1).exp()\n", + "chi_x = torch.zeros(lmx.N, lmx.L) # chi2_t(s_t) = chi2(pi(.|s_t) || mu(.|s_t)) at the realised prefix\n", + "for i, sq in enumerate(lmx.seqs):\n", + " for t in range(len(sq)):\n", + " cm, cp = lmx.cond_probs(th_mu, sq[:t]), lmx.cond_probs(th_pi, sq[:t])\n", + " chi_x[i, t] = ((cp - cm) ** 2 / cm).sum()\n", + "w2_x = (lpi - lmu).sum(1).exp() ** 2\n", + "prod_x = (1 + chi_x).prod(1) # masked positions have chi2 = 0 -> factor 1\n", + "Ew2 = float((mu_x * w2_x).sum()); Eprod = float((mu_x * prod_x).sum())\n", + "w2_prev = torch.cat([torch.zeros(lmx.N, 1), (lpi - lmu).cumsum(1)[:, :-1]], 1).exp() ** 2 # w_{t-1}^2 (w_0 = 1)\n", + "rec = float((mu_x * w2_prev[:, -1] * (1 + chi_x[:, -1])).sum())\n", + "print(f\" N={lmx.N}: E_mu[w^2] = {Ew2:.6f}; recursion E[w_(T-1)^2 (1+chi2_T)] = {rec:.6f}; E_mu[prod_t(1+chi2_t)] = {Eprod:.6f}; E_mu[w^2/prod] = {float((mu_x*w2_x/prod_x).sum()):.6f}\")\n", + "check(\"recursion E[w_T^2] = E[w_(T-1)^2 (1+chi2_T(s_T))] (state-dependent chi2)\", abs(Ew2 - rec), 1e-14)\n", + "check(\"martingale identity E[w_T^2 / prod_t (1+chi2_t(s_t))] = 1\", abs(float((mu_x * w2_x / prod_x).sum()) - 1.0), 1e-13)\n", + "assert abs(Ew2 - Eprod) > 1e-3 # the naive product formula is genuinely wrong when chi2 depends on the history\n", + "check(\"E[prod(1+chi2_t)] != E[w_T^2] on a state-dependent model (relative gap, must exceed 1e-3)\", max(0.0, 1e-3 - abs(Ew2 - Eprod) / Ew2), 0.0)" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + " N=511: E_mu[w^2] = 1.151486; recursion E[w_(T-1)^2 (1+chi2_T)] = 1.151486; E_mu[prod_t(1+chi2_t)] = 1.141717; E_mu[w^2/prod] = 1.000000\n", + "CHECK|exact|recursion E[w_T^2] = E[w_(T-1)^2 (1+chi2_T(s_T))] (state-dependent chi2)|6.661e-16|tol=1.0e-14|PASS\n", + "CHECK|exact|martingale identity E[w_T^2 / prod_t (1+chi2_t(s_t))] = 1|4.441e-16|tol=1.0e-13|PASS\n", + "CHECK|exact|E[prod(1+chi2_t)] != E[w_T^2] on a state-dependent model (relative gap, must exceed 1e-3)|0.000e+00|tol=0.0e+00|PASS\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 6. Consequence for PPO/CISPO ratios and for a \"probability correlation\" diagnostic\n", + "If the *sampler's* log-probs are used as pi_old, then at theta = theta_old the ratio is r = pi/mu != 1: tokens are clipped before any update." + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "metadata": {}, + "source": [ + "for name, dt in ((\"bf16\", torch.bfloat16), (\"fp32\", torch.float32)):\n", + " Elw, Vlw, chi2, pi_t, mu_t = per_position_stats(10.0, dt, T=2048)\n", + " gen = torch.Generator().manual_seed(1)\n", + " tok_ = torch.stack([torch.multinomial(mu_t[t], 64, replacement=True, generator=gen) for t in range(2048)], 0) # [T, 64]\n", + " r = (pi_t.gather(1, tok_) / mu_t.gather(1, tok_))\n", + " out = float(((r - 1).abs() > 0.2).double().mean())\n", + " p_tr, p_inf = pi_t.gather(1, tok_).reshape(-1), mu_t.gather(1, tok_).reshape(-1)\n", + " corr = float(torch.corrcoef(torch.stack([p_tr, p_inf]))[0, 1])\n", + " print(f\" {name} head: fraction of sampled tokens with |pi/mu - 1| > 0.2 at theta=theta_old: {out:.4f} Pearson corr(p_train, p_sampler) = {corr:.6f}\")" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + " bf16 head: fraction of sampled tokens with |pi/mu - 1| > 0.2 at theta=theta_old: 0.0000 Pearson corr(p_train, p_sampler) = 0.999793\n", + " fp32 head: fraction of sampled tokens with |pi/mu - 1| > 0.2 at theta=theta_old: 0.0000 Pearson corr(p_train, p_sampler) = 1.000000\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 7. Naive vs stable weight computation under mismatch (log-space)" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "metadata": {}, + "source": [ + "Elw, Vlw, chi2, pi_t, mu_t = per_position_stats(10.0, torch.bfloat16, T=4096)\n", + "tokens = torch.stack([torch.multinomial(mu_t[t], 1, replacement=True, generator=g) for t in range(4096)], 0).squeeze(1)\n", + "lr_ = (pi_t.log().gather(1, tokens[:, None]) - mu_t.log().gather(1, tokens[:, None])).squeeze(1)\n", + "logw = float(lr_.sum())\n", + "naive32 = torch.tensor(1.0, dtype=torch.float32)\n", + "for v in lr_.exp().to(torch.float32): naive32 = naive32 * v\n", + "print(f\" T=4096 bf16-mismatch sequence: log w = {logw:.2f}; naive fp32 running product = {float(naive32)}; exp(sum log) in fp64 = {math.exp(logw):.3e}\")\n", + "check(\"naive fp32 product vs log-space: agree only when representable (log w within fp32 range)\", 0.0 if (abs(logw) > 80 or abs(float(naive32) - math.exp(logw)) < 1e-4 * math.exp(logw)) else 1.0, 0.0)" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + " T=4096 bf16-mismatch sequence: log w = -1.55; naive fp32 running product = 0.21136480569839478; exp(sum log) in fp64 = 2.114e-01\n", + "CHECK|exact|naive fp32 product vs log-space: agree only when representable (log w within fp32 range)|0.000e+00|tol=0.0e+00|PASS\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 8. Naive vs efficient TIS-corrected gradient estimate (sampled from the bf16 sampler mu)\n", + "naive: one backward per sample; efficient: one backward through the surrogate with DETACHED truncated weights." + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "metadata": {}, + "source": [ + "lm = TinyLM(V=4, L=4, seed=0, scale=4.0)\n", + "theta = lm.theta0.clone(); R = default_reward(lm)\n", + "pi_ = seq_logprob_rounded(theta, None).exp(); mu_ = seq_logprob_rounded(theta, torch.bfloat16).exp()\n", + "bb = float(pi_ @ R)\n", + "gen = torch.Generator().manual_seed(6)\n", + "idx = torch.multinomial(mu_, 500, replacement=True, generator=gen)\n", + "Cc = 1.02\n", + "w_ = (pi_ / mu_)[idx]\n", + "acc = torch.zeros_like(theta)\n", + "for j, i in enumerate(idx.tolist()):\n", + " t = theta.clone().requires_grad_(True)\n", + " acc += float(min(w_[j], Cc)) * float(R[i] - bb) * torch.autograd.grad(seq_logprob_rounded(t, None)[i], t)[0]\n", + "g_naive = acc / len(idx)\n", + "t = theta.clone().requires_grad_(True)\n", + "sur = (torch.clamp(w_, max=Cc).detach() * (R[idx] - bb) * seq_logprob_rounded(t, None)[idx]).mean()\n", + "g_eff = torch.autograd.grad(sur, t)[0]\n", + "check(\"TIS estimate: per-sample loop == single backward through detached-weight surrogate\", maxabs(g_naive, g_eff), 1e-13)\n", + "print(\"done; checks:\", len(CHECKS))" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "CHECK|exact|TIS estimate: per-sample loop == single backward through detached-weight surrogate|3.123e-16|tol=1.0e-13|PASS\n", + "done; checks: 39\n" + ] + } + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.12.11" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} \ No newline at end of file diff --git a/code/notebooks/0109-kl-estimators.ipynb b/code/notebooks/0109-kl-estimators.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..acfe105bb88f68c683e27e31d8c38a5c971c0e3d --- /dev/null +++ b/code/notebooks/0109-kl-estimators.ipynb @@ -0,0 +1,436 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# 0109 - KL regularisation estimators k1, k2, k3: value bias/variance, what their GRADIENTS optimise, why beta = 0\n", + "Companion to `notes/0109 - KL Estimators.pdf`. Samples x ~ q = pi_theta, p = pi_ref, r = p/q, KL(q||p) = E_q[-log r]." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "source": [ + "import sys, math\n", + "sys.path.insert(0, \"..\")\n", + "from _m1_lib import *\n", + "import numpy as np" + ], + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 1. Values. E_q[k1] = E_q[k3] = KL(q||p); E_q[k2] = KL + bias, bias = -sum_{n>=3} E[l^n]/n!, l = log r\n", + " k1 = -log r, k2 = (log r)^2 / 2, k3 = (r-1) - log r (Schulman, \"Approximating KL divergence\")" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "source": [ + "def estimators(logr):\n", + " return -logr, 0.5 * logr ** 2, torch.expm1(logr) - logr\n", + "gen = torch.Generator().manual_seed(2)\n", + "for trial in range(3):\n", + " K = 6\n", + " q = torch.softmax(torch.randn(K, generator=gen) * 1.2, 0); p = torch.softmax(torch.randn(K, generator=gen) * 1.2, 0)\n", + " logr = (p / q).log()\n", + " kl = float((q * (q / p).log()).sum())\n", + " k1, k2, k3 = estimators(logr)\n", + " check(f\"categorical #{trial}: E_q[k1] = KL(q||p)\", abs(float((q * k1).sum()) - kl), 1e-14)\n", + " check(f\"categorical #{trial}: E_q[k3] = KL(q||p)\", abs(float((q * k3).sum()) - kl), 1e-14)\n", + " bias2 = float((q * k2).sum()) - kl\n", + " series = -sum(float((q * logr ** n).sum()) / math.factorial(n) for n in range(3, 40))\n", + " check(f\"categorical #{trial}: bias(k2) = -sum_(n>=3) E[l^n]/n! (series, 40 terms)\", abs(bias2 - series), 1e-12)\n", + " # the master identity KL = sum_{n>=2} E[l^n]/n! (from E_q[e^l] = 1)\n", + " ms = sum(float((q * logr ** n).sum()) / math.factorial(n) for n in range(2, 40))\n", + " check(f\"categorical #{trial}: KL = sum_(n>=2) E_q[l^n]/n!\", abs(ms - kl), 1e-12)\n", + " print(f\" KL = {kl:.5f} E[k2] - KL = {bias2:+.5f} (leading term -E[l^3]/6 = {-float((q*logr**3).sum())/6:+.5f})\")\n", + "# pointwise positivity\n", + "rg = torch.exp(torch.linspace(-12, 12, 100001))\n", + "k1g, k2g, k3g = estimators(rg.log())\n", + "check(\"k3 >= 0 and k2 >= 0 pointwise (min over a log-grid of r)\", max(0.0, -float(min(k3g.min(), k2g.min()))), 0.0)\n", + "print(f\" k1 < 0 on r > 1 (fraction of log-grid where k1<0: {float((k1g<0).double().mean()):.3f}) -> k1 is a noisy, sign-changing estimator\")" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "CHECK|exact|categorical #0: E_q[k1] = KL(q||p)|2.776e-17|tol=1.0e-14|PASS\n", + "CHECK|exact|categorical #0: E_q[k3] = KL(q||p)|2.914e-16|tol=1.0e-14|PASS\n", + "CHECK|exact|categorical #0: bias(k2) = -sum_(n>=3) E[l^n]/n! (series, 40 terms)|2.810e-16|tol=1.0e-12|PASS\n", + "CHECK|exact|categorical #0: KL = sum_(n>=2) E_q[l^n]/n!|2.776e-16|tol=1.0e-12|PASS\n", + " KL = 0.10605 E[k2] - KL = -0.01866 (leading term -E[l^3]/6 = -0.01211)\n", + "CHECK|exact|categorical #1: E_q[k1] = KL(q||p)|1.110e-16|tol=1.0e-14|PASS\n", + "CHECK|exact|categorical #1: E_q[k3] = KL(q||p)|2.220e-16|tol=1.0e-14|PASS\n", + "CHECK|exact|categorical #1: bias(k2) = -sum_(n>=3) E[l^n]/n! (series, 40 terms)|1.943e-16|tol=1.0e-12|PASS\n", + "CHECK|exact|categorical #1: KL = sum_(n>=2) E_q[l^n]/n!|2.220e-16|tol=1.0e-12|PASS\n", + " KL = 0.28620 E[k2] - KL = +0.01383 (leading term -E[l^3]/6 = +0.06153)\n", + "CHECK|exact|categorical #2: E_q[k1] = KL(q||p)|0.000e+00|tol=1.0e-14|PASS\n", + "CHECK|exact|categorical #2: E_q[k3] = KL(q||p)|0.000e+00|tol=1.0e-14|PASS\n", + "CHECK|exact|categorical #2: bias(k2) = -sum_(n>=3) E[l^n]/n! (series, 40 terms)|5.551e-17|tol=1.0e-12|PASS\n", + "CHECK|exact|categorical #2: KL = sum_(n>=2) E_q[l^n]/n!|1.110e-16|tol=1.0e-12|PASS\n", + " KL = 0.61184 E[k2] - KL = +0.15474 (leading term -E[l^3]/6 = +0.31903)\n", + "CHECK|exact|k3 >= 0 and k2 >= 0 pointwise (min over a log-grid of r)|0.000e+00|tol=0.0e+00|PASS\n", + " k1 < 0 on r > 1 (fraction of log-grid where k1<0: 0.500) -> k1 is a noisy, sign-changing estimator\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 2. Gaussian closed forms (equal variance, shift mu => KL = mu^2/2). With s2 = mu^2:\n", + " Var(k1) = s2, Var(k2) = s2^3/4 + s2^2/2, Var(k3) = e^{s2} - 1 - s2, bias(k2) = s2^2/8. Verified by Gauss-Hermite quadrature." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "source": [ + "xs, ws = np.polynomial.hermite_e.hermegauss(120); ws = ws / ws.sum()\n", + "def gauss_moments(shift):\n", + " x = torch.tensor(xs); w = torch.tensor(ws)\n", + " logr = shift * x - 0.5 * shift ** 2 # q = N(0,1), p = N(shift,1): log p/q = shift*x - shift^2/2 (E_q[r] = 1)\n", + " k = estimators(logr)\n", + " kl = shift ** 2 / 2\n", + " out = []\n", + " for kk in k:\n", + " m = float((w * kk).sum()); v = float((w * kk ** 2).sum()) - m ** 2\n", + " out.append((m, v))\n", + " return kl, out\n", + "print(\" KL est mean-KL (bias) std/KL closed-form std/KL\")\n", + "for shift in (0.1, math.sqrt(2 * 0.05), 1.0):\n", + " kl, out = gauss_moments(shift)\n", + " s2 = shift ** 2\n", + " cf = [s2, s2 ** 3 / 4 + s2 ** 2 / 2, math.exp(s2) - 1 - s2]\n", + " cb = [0.0, s2 ** 2 / 8, 0.0]\n", + " for i, nm in enumerate((\"k1\", \"k2\", \"k3\")):\n", + " m, v = out[i]\n", + " print(f\" {kl:.4f} {nm} {m-kl:+.3e} {math.sqrt(v)/kl:9.4f} {math.sqrt(cf[i])/kl:9.4f}\")\n", + " check(f\"Gaussian shift={shift:.3f}: {nm} bias == closed form\", abs((m - kl) - cb[i]), 1e-12)\n", + " check(f\"Gaussian shift={shift:.3f}: {nm} variance == closed form\", abs(v - cf[i]) / cf[i], 1e-10)\n", + "# compare with the table in Schulman's note (KL=0.005 and KL=0.5): std/KL = 20, 1.42, 1.42 and 2, 1.73, 1.7 ; bias(k2)/KL = 0.002, 0.25\n", + "kl, out = gauss_moments(0.1)\n", + "check(\"matches Schulman's table at KL=0.005: std(k1)/KL ~ 20\", abs(math.sqrt(out[0][1]) / kl - 20.0), 0.2)\n", + "kl, out = gauss_moments(1.0)\n", + "check(\"matches Schulman's table at KL=0.5: bias(k2)/KL = 0.25\", abs((out[1][0] - kl) / kl - 0.25), 1e-9)" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + " KL est mean-KL (bias) std/KL closed-form std/KL\n", + " 0.0050 k1 +1.735e-18 20.0000 20.0000\n", + "CHECK|exact|Gaussian shift=0.100: k1 bias == closed form|1.735e-18|tol=1.0e-12|PASS\n", + "CHECK|exact|Gaussian shift=0.100: k1 variance == closed form|1.735e-16|tol=1.0e-10|PASS\n", + " 0.0050 k2 +1.250e-05 1.4177 1.4177\n", + "CHECK|exact|Gaussian shift=0.100: k2 bias == closed form|1.220e-18|tol=1.0e-12|PASS\n", + "CHECK|exact|Gaussian shift=0.100: k2 variance == closed form|1.349e-16|tol=1.0e-10|PASS\n", + " 0.0050 k3 +2.602e-18 1.4166 1.4166\n", + "CHECK|exact|Gaussian shift=0.100: k3 bias == closed form|2.602e-18|tol=1.0e-12|PASS\n", + "CHECK|exact|Gaussian shift=0.100: k3 variance == closed form|2.200e-12|tol=1.0e-10|PASS\n", + " 0.0500 k1 +1.388e-17 6.3246 6.3246\n", + "CHECK|exact|Gaussian shift=0.316: k1 bias == closed form|1.388e-17|tol=1.0e-12|PASS\n", + "CHECK|exact|Gaussian shift=0.316: k1 variance == closed form|2.776e-16|tol=1.0e-10|PASS\n", + " 0.0500 k2 +1.250e-03 1.4491 1.4491\n", + "CHECK|exact|Gaussian shift=0.316: k2 bias == closed form|1.475e-17|tol=1.0e-12|PASS\n", + "CHECK|exact|Gaussian shift=0.316: k2 variance == closed form|0.000e+00|tol=1.0e-10|PASS\n", + " 0.0500 k3 +1.388e-17 1.4382 1.4382\n", + "CHECK|exact|Gaussian shift=0.316: k3 bias == closed form|1.388e-17|tol=1.0e-12|PASS\n", + "CHECK|exact|Gaussian shift=0.316: k3 variance == closed form|1.543e-14|tol=1.0e-10|PASS\n", + " 0.5000 k1 +0.000e+00 2.0000 2.0000\n", + "CHECK|exact|Gaussian shift=1.000: k1 bias == closed form|0.000e+00|tol=1.0e-12|PASS\n", + "CHECK|exact|Gaussian shift=1.000: k1 variance == closed form|2.220e-16|tol=1.0e-10|PASS\n", + " 0.5000 k2 +1.250e-01 1.7321 1.7321\n", + "CHECK|exact|Gaussian shift=1.000: k2 bias == closed form|1.110e-16|tol=1.0e-12|PASS\n", + "CHECK|exact|Gaussian shift=1.000: k2 variance == closed form|5.921e-16|tol=1.0e-10|PASS\n", + " 0.5000 k3 +1.110e-16 1.6950 1.6950\n", + "CHECK|exact|Gaussian shift=1.000: k3 bias == closed form|1.110e-16|tol=1.0e-12|PASS\n", + "CHECK|exact|Gaussian shift=1.000: k3 variance == closed form|1.546e-16|tol=1.0e-10|PASS\n", + "CHECK|exact|matches Schulman's table at KL=0.005: std(k1)/KL ~ 20|0.000e+00|tol=2.0e-01|PASS\n", + "CHECK|exact|matches Schulman's table at KL=0.5: bias(k2)/KL = 0.25|2.220e-16|tol=1.0e-09|PASS\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 3. Heavy tails: Var(k3) is infinite when chi^2(p||q) is; one token the reference loves and the policy nearly never samples" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "source": [ + "q = torch.tensor([0.001, 0.999]); p = torch.tensor([0.30, 0.70])\n", + "logr = (p / q).log(); k1, k2, k3 = estimators(logr)\n", + "kl = float((q * (q / p).log()).sum())\n", + "for nm, k in ((\"k1\", k1), (\"k2\", k2), (\"k3\", k3)):\n", + " m = float((q * k).sum()); v = float((q * k ** 2).sum()) - m ** 2\n", + " print(f\" {nm}: mean {m:+.5f} std {math.sqrt(v):.5f} (KL = {kl:.5f}) max value {float(k.max()):.1f}\")\n", + "chi2_pq = float((p ** 2 / q).sum() - 1)\n", + "print(f\" chi2(p||q) = E_q[r^2]-1 = {chi2_pq:.1f} -> Var(k3) is dominated by one rare token with r = {float((p/q)[0]):.0f}\")\n", + "check(\"E_q[r^2] - 1 = chi2(p||q) exactly\", abs(float((q * (p / q) ** 2).sum()) - 1 - chi2_pq), 1e-9)" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + " k1: mean +0.34961 std 0.19152 (KL = 0.34961) max value 0.4\n", + " k2: mean +0.07946 std 0.51214 (KL = 0.34961) max value 16.3\n", + " k3: mean +0.34961 std 9.26842 (KL = 0.34961) max value 293.3\n", + " chi2(p||q) = E_q[r^2]-1 = 89.5 -> Var(k3) is dominated by one rare token with r = 300\n", + "CHECK|exact|E_q[r^2] - 1 = chi2(p||q) exactly|0.000e+00|tol=1.0e-09|PASS\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 4. GRADIENTS (on-policy, single-step categorical policy with logits theta; samples x ~ q_theta DETACHED, gradient flows through q_theta in r)\n", + " k1-as-loss: E[grad k1] = E[grad log q] = 0 -> useless as a regulariser\n", + " k2-as-loss: E[grad k2] = E[(log q - log p) grad log q] = grad KL(q||p) (the reverse KL we meant)\n", + " k3-as-loss: E[grad k3] = E[(1 - r) grad log q] = -sum_x p grad log q = grad KL(p||q) (the FORWARD KL)\n", + " k1-in-reward (detached coefficient): E[(log q - log p) grad log q] = grad KL(q||p)" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "source": [ + "gen = torch.Generator().manual_seed(7)\n", + "K = 6\n", + "th_ref = torch.randn(K, generator=gen) * 1.0\n", + "def grads(delta):\n", + " th = (th_ref + delta * torch.randn(K, generator=gen)).clone().requires_grad_(True)\n", + " q = torch.softmax(th, 0); p = torch.softmax(th_ref, 0)\n", + " lq, lp = torch.log_softmax(th, 0), torch.log_softmax(th_ref, 0)\n", + " qd = q.detach()\n", + " r = (lp - lq).exp()\n", + " # expectation over x ~ q (weights detached) of the gradient of each estimator\n", + " g = {}\n", + " g['k1'] = torch.autograd.grad((qd * (-(lp - lq))).sum(), th, retain_graph=True)[0]\n", + " g['k2'] = torch.autograd.grad((qd * 0.5 * (lp - lq) ** 2).sum(), th, retain_graph=True)[0]\n", + " g['k3'] = torch.autograd.grad((qd * (r - 1 - (lp - lq))).sum(), th, retain_graph=True)[0]\n", + " # k1 in reward: detached coefficient (log q - log p) times score\n", + " g['k1r'] = torch.autograd.grad((qd * (lq - lp).detach() * lq).sum(), th, retain_graph=True)[0]\n", + " g['rev'] = torch.autograd.grad((q * (lq - lp)).sum(), th, retain_graph=True)[0] # grad KL(q||p), exact\n", + " g['fwd'] = torch.autograd.grad((p * (lp - lq)).sum(), th, retain_graph=True)[0] # grad KL(p||q), exact\n", + " return g\n", + "g = grads(0.8)\n", + "check(\"E[grad k1] = 0 (k1 as a loss carries no regularising gradient)\", float(g['k1'].abs().max()), 1e-14)\n", + "check(\"E[grad k2] == grad KL(q||p) (reverse KL)\", maxabs(g['k2'], g['rev']), 1e-14)\n", + "check(\"k1-in-reward (detached coefficient) == grad KL(q||p)\", maxabs(g['k1r'], g['rev']), 1e-14)\n", + "check(\"E[grad k3] == grad KL(p||q) (FORWARD KL, not the one k3 estimates)\", maxabs(g['k3'], g['fwd']), 1e-14)\n", + "print(f\" ||grad KL(q||p) - grad KL(p||q)|| / ||grad KL(q||p)|| = {float((g['rev']-g['fwd']).norm()/g['rev'].norm()):.3f} (different vectors)\")\n", + "print(\" relative gap between the two gradients as q -> p (first-order agreement, 'k3 is a first-order approximation'):\")\n", + "prev = None\n", + "for delta in (0.4, 0.2, 0.1, 0.05, 0.025):\n", + " gen.manual_seed(7); _ = torch.randn(K, generator=gen) # same direction each time\n", + " gg = grads(delta)\n", + " gap = float((gg['rev'] - gg['fwd']).norm() / gg['rev'].norm())\n", + " print(f\" delta={delta:5.3f} gap = {gap:.4f}\")\n", + " prev = gap\n", + "assert prev < 0.05\n", + "check(\"gradient gap forward vs reverse KL vanishes as delta -> 0 (gap at delta=0.025 < 5%)\", 0.0, 0.0)" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "CHECK|exact|E[grad k1] = 0 (k1 as a loss carries no regularising gradient)|5.551e-17|tol=1.0e-14|PASS\n", + "CHECK|exact|E[grad k2] == grad KL(q||p) (reverse KL)|5.551e-17|tol=1.0e-14|PASS\n", + "CHECK|exact|k1-in-reward (detached coefficient) == grad KL(q||p)|5.551e-17|tol=1.0e-14|PASS\n", + "CHECK|exact|E[grad k3] == grad KL(p||q) (FORWARD KL, not the one k3 estimates)|8.327e-17|tol=1.0e-14|PASS\n", + " ||grad KL(q||p) - grad KL(p||q)|| / ||grad KL(q||p)|| = 0.215 (different vectors)\n", + " relative gap between the two gradients as q -> p (first-order agreement, 'k3 is a first-order approximation'):\n", + " delta=0.400 gap = 0.1533\n", + " delta=0.200 gap = 0.0922\n", + " delta=0.100 gap = 0.0505\n", + " delta=0.050 gap = 0.0264\n", + " delta=0.025 gap = 0.0135\n", + "CHECK|exact|gradient gap forward vs reverse KL vanishes as delta -> 0 (gap at delta=0.025 < 5%)|0.000e+00|tol=0.0e+00|PASS\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 5. SEQUENCES: per-token k2/k3 'as loss' ignore that earlier tokens change the later contexts\n", + " true: grad KL_seq(pi||ref) = E[ (sum_t l_t) * score(o) ] = E[ sum_t (sum_{t'>=t} l_t') s_t ] (\"k1 in reward\" with reward-to-go)\n", + " per-token k2 loss, summed over t: E[ sum_t l_t s_t ] = true - sum_t E[ KL_t(s_t) * score(o_=t} l_t') s_t] == grad KL_seq\", maxabs(g_rtg, g_true), 1e-14)\n", + "print(f\" per-token k2-as-loss gradient vs true sequence-KL gradient: relative error {float((g_tok-g_true).norm()/g_true.norm()):.4f}\")\n", + "# the missing term: sum_t E[ KL_t(s_t) * score(o_=t} l_t') s_t] == grad KL_seq|8.327e-17|tol=1.0e-14|PASS\n", + " per-token k2-as-loss gradient vs true sequence-KL gradient: relative error 0.0376\n", + "CHECK|exact|per-token k2 loss gradient == true gradient - sum_t E[KL_t * prefix score] (exact identity)|9.714e-17|tol=1.0e-13|PASS\n", + "CHECK|exact|per-token k3-as-loss gradient == grad of sum_s P_pi(s) KL(ref(.|s)||pi_theta(.|s)) (context weights frozen)|5.898e-17|tol=1.0e-13|PASS\n", + " sequence KL = 0.01885; E[sum_t k1] = 0.01885; E[sum_t k3] = 0.01885; E[sum_t k2] = 0.02103\n", + "CHECK|exact|E[sum_t k3_t] = KL_seq (per-token estimator is unbiased for the sequence KL value)|6.245e-17|tol=1.0e-13|PASS\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 6. Naive vs stable evaluation of k3 (fp32): (r-1)-log r cancels catastrophically when |log r| is small" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "source": [ + "lr32 = torch.tensor([1e-1, 1e-2, 1e-3, 1e-4, 1e-5], dtype=torch.float32)\n", + "naive = (lr32.exp() - 1) - lr32\n", + "stable = torch.expm1(lr32) - lr32\n", + "truth = (torch.expm1(lr32.double()) - lr32.double())\n", + "print(\" log r naive fp32 expm1 fp32 fp64 truth naive rel.err stable rel.err\")\n", + "for a, n_, s_, t_ in zip(lr32.tolist(), naive.tolist(), stable.tolist(), truth.tolist()):\n", + " print(f\" {a:7.0e} {n_:+.4e} {s_:+.4e} {t_:+.4e} {abs(n_-t_)/t_:9.2e} {abs(s_-t_)/t_:9.2e}\")\n", + "check(\"expm1 form: relative error < 1e-3 for |log r| >= 1e-4 (fp32)\", float(((stable.double() - truth).abs() / truth)[:4].max()), 1e-3)\n", + "assert float(((naive.double() - truth).abs() / truth)[3]) > 0.1\n", + "check(\"naive form is wrong by >10% at |log r| = 1e-4 (fp32)\", 0.0, 0.0)\n", + "# The expm1 form removes the cancellation inside exp(x)-1 but NOT the one in expm1(x) - x ~ x^2/2: its worst-case relative error is ~ 2u/|x|.\n", + "# Worst case over a dense random grid per decade (fp32, both signs), vs the 2u/|x|_min law; and a series fallback for small |x|.\n", + "u32 = 2.0 ** -24\n", + "gen_n = torch.Generator().manual_seed(11)\n", + "def k3_expm1(x): return torch.expm1(x) - x\n", + "def k3_hybrid(x):\n", + " ser = x * x / 2 * (1 + x / 3 + x * x / 12) # Taylor of expm1(x)-x, error ~ x^5/120 relative x^3/60\n", + " return torch.where(x.abs() < 1e-2, ser, k3_expm1(x))\n", + "print(\" |x| decade worst rel.err expm1 form 2u/|x|min worst rel.err hybrid (series for |x|<1e-2)\")\n", + "worst_ratio, worst_hybrid = 0.0, 0.0\n", + "for lo in (1e-1, 1e-2, 1e-3, 1e-4, 1e-5, 1e-6):\n", + " hi = 10 * lo\n", + " xg = torch.exp(torch.rand(400_000, generator=gen_n, dtype=torch.float64) * math.log(10) + math.log(lo)).float()\n", + " xg = torch.cat([xg, -xg]); tr = torch.expm1(xg.double()) - xg.double()\n", + " re = float(((k3_expm1(xg).double() - tr).abs() / tr).max()); rh = float(((k3_hybrid(xg).double() - tr).abs() / tr).max())\n", + " law = 2 * u32 / lo\n", + " worst_ratio = max(worst_ratio, re / law); worst_hybrid = max(worst_hybrid, rh)\n", + " print(f\" [{lo:.0e},{hi:.0e}] {re:12.3e} {law:10.2e} {rh:12.3e}\")\n", + "check(\"expm1 form: worst-case relative error <= 2u/|x| (max ratio to the law)\", worst_ratio, 1.0)\n", + "check(\"hybrid (Taylor for |x|<1e-2, expm1 form above): worst relative error over |x| in [1e-6, 1)\", worst_hybrid, 1e-4)\n", + "print(\"done; checks:\", len(CHECKS))" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + " log r naive fp32 expm1 fp32 fp64 truth naive rel.err stable rel.err\n", + " 1e-01 +5.1710e-03 +5.1709e-03 +5.1709e-03 8.79e-06 1.49e-07\n", + " 1e-02 +5.0178e-05 +5.0167e-05 +5.0167e-05 2.14e-04 9.17e-06\n", + " 1e-03 +5.2352e-07 +5.0012e-07 +5.0017e-07 4.67e-02 9.30e-05\n", + " 1e-04 +1.6596e-08 +4.9986e-09 +5.0002e-09 2.32e+00 3.17e-04\n", + " 1e-05 +1.3581e-08 +5.0022e-11 +5.0000e-11 2.71e+02 4.41e-04\n", + "CHECK|exact|expm1 form: relative error < 1e-3 for |log r| >= 1e-4 (fp32)|3.167e-04|tol=1.0e-03|PASS\n", + "CHECK|exact|naive form is wrong by >10% at |log r| = 1e-4 (fp32)|0.000e+00|tol=0.0e+00|PASS\n", + " |x| decade worst rel.err expm1 form 2u/|x|min worst rel.err hybrid (series for |x|<1e-2)\n", + " [1e-01,1e+00] 1.033e-06 1.19e-06 1.033e-06\n", + " [1e-02,1e-01] 9.291e-06 1.19e-05 9.291e-06\n", + " [1e-03,1e-02] 1.161e-04 1.19e-04 2.203e-07\n", + " [1e-04,1e-03] 9.735e-04 1.19e-03 2.035e-07\n", + " [1e-05,1e-04] 9.034e-03 1.19e-02 1.667e-07\n", + " [1e-06,1e-05] 1.150e-01 1.19e-01 1.646e-07\n", + "CHECK|exact|expm1 form: worst-case relative error <= 2u/|x| (max ratio to the law)|9.736e-01|tol=1.0e+00|PASS\n", + "CHECK|exact|hybrid (Taylor for |x|<1e-2, expm1 form above): worst relative error over |x| in [1e-6, 1)|9.291e-06|tol=1.0e-04|PASS\n", + "done; checks: 48\n" + ] + } + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.12.11" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} \ No newline at end of file diff --git a/code/notebooks/0110-reward-shaping-verifier.ipynb b/code/notebooks/0110-reward-shaping-verifier.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..cc07ba3201933312e2f3d61fab28f43308221851 --- /dev/null +++ b/code/notebooks/0110-reward-shaping-verifier.ipynb @@ -0,0 +1,417 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# 0110 - Reward shaping for verifier rewards: binary rewards, dead groups, overlong penalties as reward, verifier errors\n", + "Companion to `notes/0110 - Reward Shaping for Verifier Rewards.pdf`. fp64; exact enumeration on the 121-sequence toy LM." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "source": [ + "import sys, math, itertools\n", + "sys.path.insert(0, \"..\")\n", + "from _m1_lib import *\n", + "import numpy as np\n", + "lm = TinyLM(V=4, L=4, seed=0, scale=1.0)\n", + "theta = lm.theta0.clone()\n", + "R = default_reward(lm)\n", + "p_seq = lm.probs(theta)\n", + "S = score_matrix(lm, theta)\n", + "N = lm.N\n", + "th = theta.clone().requires_grad_(True)\n", + "pr = (lm.probs(th) * R).sum()\n", + "grad_p = torch.autograd.grad(pr, th)[0]\n", + "p = float(pr.detach())" + ], + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 1. Binary reward: grad p = p (1-p) ( E[score | success] - E[score | failure] )" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "source": [ + "m1 = (p_seq * R)[:, None] * S\n", + "m0 = (p_seq * (1 - R))[:, None] * S\n", + "E_s_succ = m1.sum(0) / p; E_s_fail = m0.sum(0) / (1 - p)\n", + "check(\"grad p = p (1-p) (E[s|succ] - E[s|fail])\", maxabs(grad_p, p * (1 - p) * (E_s_succ - E_s_fail)), 1e-14)\n", + "check(\"grad p = p E[s | success] (rejection-sampling / expert-iteration form)\", maxabs(grad_p, p * E_s_succ), 1e-14)\n", + "print(f\" p = {p:.4f}; signal scale p(1-p) = {p*(1-p):.4f}\")" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "CHECK|exact|grad p = p (1-p) (E[s|succ] - E[s|fail])|6.939e-18|tol=1.0e-14|PASS\n", + "CHECK|exact|grad p = p E[s | success] (rejection-sampling / expert-iteration form)|1.388e-17|tol=1.0e-14|PASS\n", + " p = 0.1659; signal scale p(1-p) = 0.1384\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 2. Dead groups. Binary bandit with logit parameter: the group estimators collapse to functions of p_hat = k/G\n", + " centred (Dr.GRPO): g = p_hat (1 - p_hat) (plug-in variance of the rewards); std-normalised (GRPO): g = sqrt(p_hat (1 - p_hat))" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "source": [ + "def bern_group(G, pp):\n", + " ks = np.arange(G + 1)\n", + " w = np.array([math.comb(G, k) * pp ** k * (1 - pp) ** (G - k) for k in ks])\n", + " ph = ks / G\n", + " return ks, w, ph\n", + "def direct_estimator(G, k, pp, kind):\n", + " \"\"\"(1/G) sum_i A_i * s_i with s = d log pi / d logit: success -> (1-p), failure -> -p (computed from the definitions)\"\"\"\n", + " Rg = torch.cat([torch.ones(k), torch.zeros(G - k)])\n", + " mu = Rg.mean(); sd = Rg.std(unbiased=False)\n", + " A = Rg - mu if kind == 'centre' else (torch.zeros_like(Rg) if float(sd) == 0 else (Rg - mu) / sd)\n", + " s = torch.where(Rg > 0, torch.tensor(1 - pp), torch.tensor(-pp))\n", + " return float((A * s).mean())\n", + "worst = 0.0\n", + "for G in (2, 4, 16):\n", + " for pp in (0.05, 0.3, 0.7):\n", + " ks, w, ph = bern_group(G, pp)\n", + " for k in ks:\n", + " worst = max(worst, abs(direct_estimator(G, int(k), pp, 'centre') - ph[k] * (1 - ph[k])),\n", + " abs(direct_estimator(G, int(k), pp, 'std') - math.sqrt(ph[k] * (1 - ph[k]))))\n", + "check(\"bandit: centred estimator == p_hat(1-p_hat), std-normalised == sqrt(p_hat(1-p_hat)) (all k, 9 settings)\", worst, 1e-14)\n", + "# exact mean / variance / SNR of the estimator as a function of p\n", + "print(\" G=16: p P(dead) E[g_centred] (1-1/G)p(1-p) SNR=E^2/Var E[g_std]/(p(1-p)) vs c_G(p)\")\n", + "def c_G_formula(G, pp): # from note 0105\n", + " bp = lambda n, k: math.comb(n, k) * pp ** k * (1 - pp) ** (n - k)\n", + " a1 = sum(bp(G - 1, K) * (math.sqrt((G - (1 + K)) / (1 + K)) if 1 + K < G else 0.0) for K in range(G))\n", + " a0 = sum(bp(G - 1, K) * (-math.sqrt(K / (G - K)) if 1 <= K < G else 0.0) for K in range(G))\n", + " return a1 - a0\n", + "G = 16\n", + "snr = {}\n", + "for pp in (0.01, 0.05, 0.2, 0.5, 0.8, 0.95, 0.99):\n", + " ks, w, ph = bern_group(G, pp)\n", + " g = ph * (1 - ph); gs = np.sqrt(ph * (1 - ph))\n", + " mean, var = float((w * g).sum()), float((w * g ** 2).sum() - (w * g).sum() ** 2)\n", + " pdead = pp ** G + (1 - pp) ** G\n", + " mean_std = float((w * gs).sum())\n", + " snr[pp] = mean ** 2 / var\n", + " print(f\" {pp:5.2f} {pdead:7.4f} {mean:10.6f} {(1-1/G)*pp*(1-pp):10.6f} {snr[pp]:8.3f} {mean_std/(pp*(1-pp)):7.4f} vs {c_G_formula(G, pp):7.4f}\")\n", + " check(f\"E[g_centred] = (1-1/G) p(1-p) (p={pp})\", abs(mean - (1 - 1 / G) * pp * (1 - pp)), 1e-15)\n", + " check(f\"E[sqrt(p_hat(1-p_hat))] / (p(1-p)) == c_G(p) of note 0105 (p={pp})\", abs(mean_std / (pp * (1 - pp)) - c_G_formula(G, pp)), 1e-12)\n", + " check(f\"P(dead) = p^G + (1-p)^G = mass of k in {{0,G}} (p={pp})\", abs(pdead - (w[0] + w[-1])), 1e-15)\n", + "assert snr[0.5] > snr[0.2] > snr[0.05] > snr[0.01] and snr[0.5] > snr[0.8] > snr[0.95] > snr[0.99]\n", + "check(\"SNR of the group estimator peaks at p = 1/2 and decays towards both ends (monotone on each side)\", 0.0, 0.0)\n", + "# dead groups carry exactly zero gradient for every baseline estimator; none-baseline REINFORCE does not\n", + "Gs = 4\n", + "dead_succ = [int(np.argmax(R.numpy()))] * 0 + [i for i in range(N) if R[i] == 1][:Gs] # an all-correct group (4 successful sequences)\n", + "dead_fail = [i for i in range(N) if R[i] == 0][:Gs]\n", + "def group_gradient(idx, kind):\n", + " Rg = R[idx]\n", + " if kind == 'centre': A = Rg - Rg.mean()\n", + " elif kind == 'std':\n", + " sd = Rg.std(unbiased=False); A = torch.zeros_like(Rg) if float(sd) == 0 else (Rg - Rg.mean()) / sd\n", + " elif kind == 'rloo': A = Rg - (Rg.sum() - Rg) / (len(Rg) - 1)\n", + " elif kind == 'plain01': A = Rg # REINFORCE, no baseline, R in {0,1}\n", + " elif kind == 'plainpm': A = 2 * Rg - 1 # REINFORCE, no baseline, R in {-1,+1}\n", + " return (A[:, None] * S[idx]).mean(0)\n", + "print(\" ||group gradient|| for dead groups (4 samples):\")\n", + "for kind in ('centre', 'std', 'rloo', 'plain01', 'plainpm'):\n", + " print(f\" {kind:8s} all-correct: {float(group_gradient(dead_succ, kind).norm()):.4f} all-wrong: {float(group_gradient(dead_fail, kind).norm()):.4f}\")\n", + "for kind in ('centre', 'std', 'rloo'):\n", + " check(f\"dead groups have exactly zero gradient ({kind})\", max(float(group_gradient(dead_succ, kind).abs().max()), float(group_gradient(dead_fail, kind).abs().max())), 0.0)\n", + "assert float(group_gradient(dead_succ, 'plain01').norm()) > 1e-3 and float(group_gradient(dead_fail, 'plainpm').norm()) > 1e-3\n", + "check(\"without a baseline an all-correct (R in {0,1}) or all-wrong (R in {-1,1}) group still moves the policy\", 0.0, 0.0)" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "CHECK|exact|bandit: centred estimator == p_hat(1-p_hat), std-normalised == sqrt(p_hat(1-p_hat)) (all k, 9 settings)|1.665e-16|tol=1.0e-14|PASS\n", + " G=16: p P(dead) E[g_centred] (1-1/G)p(1-p) SNR=E^2/Var E[g_std]/(p(1-p)) vs c_G(p)\n", + " 0.01 0.8515 0.009281 0.009281 0.165 3.7330 vs 3.7330\n", + "CHECK|exact|E[g_centred] = (1-1/G) p(1-p) (p=0.01)|0.000e+00|tol=1.0e-15|PASS\n", + "CHECK|exact|E[sqrt(p_hat(1-p_hat))] / (p(1-p)) == c_G(p) of note 0105 (p=0.01)|8.882e-16|tol=1.0e-12|PASS\n", + "CHECK|exact|P(dead) = p^G + (1-p)^G = mass of k in {0,G} (p=0.01)|0.000e+00|tol=1.0e-15|PASS\n", + " 0.05 0.4401 0.044531 0.044531 0.931 3.2671 vs 3.2671\n", + "CHECK|exact|E[g_centred] = (1-1/G) p(1-p) (p=0.05)|1.388e-17|tol=1.0e-15|PASS\n", + "CHECK|exact|E[sqrt(p_hat(1-p_hat))] / (p(1-p)) == c_G(p) of note 0105 (p=0.05)|4.441e-16|tol=1.0e-12|PASS\n", + "CHECK|exact|P(dead) = p^G + (1-p)^G = mass of k in {0,G} (p=0.05)|0.000e+00|tol=1.0e-15|PASS\n", + " 0.20 0.0281 0.150000 0.150000 6.713 2.3462 vs 2.3462\n", + "CHECK|exact|E[g_centred] = (1-1/G) p(1-p) (p=0.2)|1.110e-16|tol=1.0e-15|PASS\n", + "CHECK|exact|E[sqrt(p_hat(1-p_hat))] / (p(1-p)) == c_G(p) of note 0105 (p=0.2)|4.441e-16|tol=1.0e-12|PASS\n", + "CHECK|exact|P(dead) = p^G + (1-p)^G = mass of k in {0,G} (p=0.2)|0.000e+00|tol=1.0e-15|PASS\n", + " 0.50 0.0000 0.234375 0.234375 120.000 1.9341 vs 1.9341\n", + "CHECK|exact|E[g_centred] = (1-1/G) p(1-p) (p=0.5)|0.000e+00|tol=1.0e-15|PASS\n", + "CHECK|exact|E[sqrt(p_hat(1-p_hat))] / (p(1-p)) == c_G(p) of note 0105 (p=0.5)|0.000e+00|tol=1.0e-12|PASS\n", + "CHECK|exact|P(dead) = p^G + (1-p)^G = mass of k in {0,G} (p=0.5)|0.000e+00|tol=1.0e-15|PASS\n", + " 0.80 0.0281 0.150000 0.150000 6.713 2.3462 vs 2.3462\n", + "CHECK|exact|E[g_centred] = (1-1/G) p(1-p) (p=0.8)|0.000e+00|tol=1.0e-15|PASS\n", + "CHECK|exact|E[sqrt(p_hat(1-p_hat))] / (p(1-p)) == c_G(p) of note 0105 (p=0.8)|0.000e+00|tol=1.0e-12|PASS\n", + "CHECK|exact|P(dead) = p^G + (1-p)^G = mass of k in {0,G} (p=0.8)|0.000e+00|tol=1.0e-15|PASS\n", + " 0.95 0.4401 0.044531 0.044531 0.931 3.2671 vs 3.2671\n", + "CHECK|exact|E[g_centred] = (1-1/G) p(1-p) (p=0.95)|6.939e-18|tol=1.0e-15|PASS\n", + "CHECK|exact|E[sqrt(p_hat(1-p_hat))] / (p(1-p)) == c_G(p) of note 0105 (p=0.95)|0.000e+00|tol=1.0e-12|PASS\n", + "CHECK|exact|P(dead) = p^G + (1-p)^G = mass of k in {0,G} (p=0.95)|0.000e+00|tol=1.0e-15|PASS\n", + " 0.99 0.8515 0.009281 0.009281 0.165 3.7330 vs 3.7330\n", + "CHECK|exact|E[g_centred] = (1-1/G) p(1-p) (p=0.99)|3.469e-18|tol=1.0e-15|PASS\n", + "CHECK|exact|E[sqrt(p_hat(1-p_hat))] / (p(1-p)) == c_G(p) of note 0105 (p=0.99)|4.441e-16|tol=1.0e-12|PASS\n", + "CHECK|exact|P(dead) = p^G + (1-p)^G = mass of k in {0,G} (p=0.99)|0.000e+00|tol=1.0e-15|PASS\n", + "CHECK|exact|SNR of the group estimator peaks at p = 1/2 and decays towards both ends (monotone on each side)|0.000e+00|tol=0.0e+00|PASS\n", + " ||group gradient|| for dead groups (4 samples):\n", + " centre all-correct: 0.0000 all-wrong: 0.0000\n", + " std all-correct: 0.0000 all-wrong: 0.0000\n", + " rloo all-correct: 0.0000 all-wrong: 0.0000\n", + " plain01 all-correct: 2.1004 all-wrong: 0.0000\n", + " plainpm all-correct: 2.1004 all-wrong: 3.8422\n", + "CHECK|exact|dead groups have exactly zero gradient (centre)|0.000e+00|tol=0.0e+00|PASS\n", + "CHECK|exact|dead groups have exactly zero gradient (std)|0.000e+00|tol=0.0e+00|PASS\n", + "CHECK|exact|dead groups have exactly zero gradient (rloo)|0.000e+00|tol=0.0e+00|PASS\n", + "CHECK|exact|without a baseline an all-correct (R in {0,1}) or all-wrong (R in {-1,1}) group still moves the policy|0.000e+00|tol=0.0e+00|PASS\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 3. Reward scale / sign conventions" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "source": [ + "gen = torch.Generator().manual_seed(1)\n", + "idx = torch.multinomial(p_seq, 8, replacement=True, generator=gen)\n", + "for kind in ('std',):\n", + " g01 = group_gradient(idx, kind)\n", + " Rpm_idx = idx\n", + " Rg = 2 * R[idx] - 1\n", + " sd = Rg.std(unbiased=False)\n", + " gpm = (((Rg - Rg.mean()) / sd)[:, None] * S[idx]).mean(0) if float(sd) > 0 else torch.zeros_like(g01)\n", + " check(\"GRPO (std-normalised): {0,1} and {-1,+1} rewards give the SAME gradient\", maxabs(g01, gpm), 1e-13)\n", + "g01 = group_gradient(idx, 'centre')\n", + "Rg = 2 * R[idx] - 1\n", + "gpm = ((Rg - Rg.mean())[:, None] * S[idx]).mean(0)\n", + "check(\"Dr.GRPO (centred): {-1,+1} rewards give exactly 2x the {0,1} gradient\", maxabs(gpm, 2 * g01), 1e-13)" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "CHECK|exact|GRPO (std-normalised): {0,1} and {-1,+1} rewards give the SAME gradient|0.000e+00|tol=1.0e-13|PASS\n", + "CHECK|exact|Dr.GRPO (centred): {-1,+1} rewards give exactly 2x the {0,1} gradient|0.000e+00|tol=1.0e-13|PASS\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 4. Potential-based shaping (finite horizon, gamma = 1): r_t = Phi(s_{t+1}) - Phi(s_t), Phi(terminal) = 0 => shaped return = R - Phi(s_0) => same gradient" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "source": [ + "gen = torch.Generator().manual_seed(3)\n", + "Phi = {}\n", + "for pre in lm.prefixes(): Phi[pre] = float(torch.randn(1, generator=gen)) * 2.0\n", + "for sq in lm.seqs:\n", + " Phi.setdefault(sq, 0.0) # terminal states: Phi = 0\n", + "shaping = torch.zeros(lm.N)\n", + "for i, sq in enumerate(lm.seqs):\n", + " tot = 0.0\n", + " for t in range(len(sq)):\n", + " a, b = sq[:t], sq[:t + 1]\n", + " tot += Phi[b] - Phi[a]\n", + " shaping[i] = tot\n", + "check(\"telescoping: sum_t [Phi(s_t+1) - Phi(s_t)] = -Phi(s_0) for every sequence (Phi(terminal)=0)\", float((shaping + Phi[()]).abs().max()), 1e-14)\n", + "th = theta.clone().requires_grad_(True)\n", + "g_shaped = torch.autograd.grad((lm.probs(th) * (R + shaping)).sum(), th)[0]\n", + "check(\"potential-based shaping leaves grad J unchanged (exact)\", maxabs(g_shaped, grad_p), 1e-14)\n", + "# a non-potential 'shaping' (terminal potential != 0, e.g. depending on length) changes the objective\n", + "Phi_bad = dict(Phi)\n", + "for sq in lm.seqs: Phi_bad[sq] = 0.7 * len(sq)\n", + "bad = torch.zeros(lm.N)\n", + "for i, sq in enumerate(lm.seqs):\n", + " bad[i] = sum(Phi_bad[sq[:t + 1]] - Phi_bad[sq[:t]] for t in range(len(sq)))\n", + "th = theta.clone().requires_grad_(True)\n", + "g_bad = torch.autograd.grad((lm.probs(th) * (R + bad)).sum(), th)[0]\n", + "print(f\" terminal potential 0.7*|o|: ||grad - grad J|| = {float((g_bad-grad_p).norm()):.4f} (||grad J|| = {float(grad_p.norm()):.4f})\")\n", + "assert float((g_bad - grad_p).norm()) > 1e-3\n", + "check(\"a terminal potential (length-dependent) is NOT neutral: the gradient changes\", 0.0, 0.0)" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "CHECK|exact|telescoping: sum_t [Phi(s_t+1) - Phi(s_t)] = -Phi(s_0) for every sequence (Phi(terminal)=0)|4.441e-16|tol=1.0e-14|PASS\n", + "CHECK|exact|potential-based shaping leaves grad J unchanged (exact)|1.457e-16|tol=1.0e-14|PASS\n", + " terminal potential 0.7*|o|: ||grad - grad J|| = 0.5800 (||grad J|| = 0.1825)\n", + "CHECK|exact|a terminal potential (length-dependent) is NOT neutral: the gradient changes|0.000e+00|tol=0.0e+00|PASS\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 5. Length penalties: as REWARD the objective is separable; length has no other route to the policy" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "source": [ + "lens = lm.lengths\n", + "lam = 0.05\n", + "b = float(p_seq @ R)\n", + "g_rew = (p_seq[:, None] * (R - lam * lens - (b - lam * float(p_seq @ lens)))[:, None] * S).sum(0) # baselined REINFORCE with shaped reward\n", + "th = theta.clone().requires_grad_(True)\n", + "gJ = torch.autograd.grad((lm.probs(th) * R).sum(), th)[0]\n", + "th = theta.clone().requires_grad_(True)\n", + "gL = torch.autograd.grad((lm.probs(th) * lens).sum(), th)[0]\n", + "check(\"reward R - lam|o| : gradient = grad E[R] - lam grad E|o| (separable, exact)\", maxabs(g_rew, gJ - lam * gL), 1e-14)\n", + "check(\"grad E|o| = E[|o| score] (length can only reach the policy through the score function)\", maxabs(gL, (p_seq[:, None] * lens[:, None] * S).sum(0)), 1e-14)\n", + "# effective per-sequence rewards of the implicit length weighting of note 0105 vs an explicit penalty\n", + "R_eff_len = (R - b) / lens\n", + "print(\" implicit reward of the 1/|o|-normalised GRPO loss: (R-b)/|o| -> not separable, depends on b and multiplies the reward by 1/|o|;\")\n", + "print(\" explicit penalty R - lam|o| is separable and its trade-off is the single number lam.\")" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "CHECK|exact|reward R - lam|o| : gradient = grad E[R] - lam grad E|o| (separable, exact)|1.041e-17|tol=1.0e-14|PASS\n", + "CHECK|exact|grad E|o| = E[|o| score] (length can only reach the policy through the score function)|1.110e-16|tol=1.0e-14|PASS\n", + " implicit reward of the 1/|o|-normalised GRPO loss: (R-b)/|o| -> not separable, depends on b and multiplies the reward by 1/|o|;\n", + " explicit penalty R - lam|o| is separable and its trade-off is the single number lam.\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 6. Verifier errors. Constant false-accept / false-reject rates only rescale the gradient; an exploitable false-accept class re-targets it." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "source": [ + "aFA, aFR = 0.10, 0.05\n", + "th = theta.clone().requires_grad_(True)\n", + "Pc = (lm.probs(th) * R).sum()\n", + "Jobs = aFA + (1 - aFR - aFA) * Pc # E[R_obs] when errors are independent of the response\n", + "Robs_expect = R * (1 - aFR) + (1 - R) * aFA # per-response expected observed reward\n", + "th2 = theta.clone().requires_grad_(True)\n", + "Jobs2 = (lm.probs(th2) * Robs_expect).sum()\n", + "g1 = torch.autograd.grad(Jobs, th)[0]; g2 = torch.autograd.grad(Jobs2, th2)[0]\n", + "check(\"E[R_obs] = aFA + (1-aFR-aFA) P(correct) (value)\", abs(float(Jobs.detach()) - float(Jobs2.detach())), 1e-15)\n", + "check(\"constant error rates: grad E[R_obs] = (1-aFR-aFA) grad P(correct) (exact rescaling)\", maxabs(g2, (1 - aFR - aFA) * grad_p), 1e-14)\n", + "# exploitable class: short sequences (length 2 ending in EOS) with an even token get accepted with prob 1\n", + "expl = torch.tensor([1.0 if (len(sq) == 2 and sq[-1] == lm.EOS and sq[0] % 2 == 0) else 0.0 for sq in lm.seqs]) * (1 - R)\n", + "th = theta.clone().requires_grad_(True)\n", + "gE = torch.autograd.grad((lm.probs(th) * expl).sum(), th)[0]\n", + "P_E = float((p_seq * expl).sum())\n", + "print(f\" exploit class (always accepted although wrong): P = {P_E:.4f}; ||grad P(correct)|| = {float(grad_p.norm()):.4f}; ||grad P(exploit)|| = {float(gE.norm()):.4f}; cosine = {float(grad_p@gE/(grad_p.norm()*gE.norm())):+.3f}\")\n", + "th = theta.clone().requires_grad_(True)\n", + "g_obs = torch.autograd.grad((lm.probs(th) * (R + expl)).sum(), th)[0] # E[R_obs], R_obs = 1 on correct OR exploit\n", + "check(\"observed-reward gradient (autograd) == grad P(correct) + grad P(exploit)\", maxabs(g_obs, grad_p + gE), 1e-14)\n", + "print(\" exploit share of the observed-reward gradient norm: %.2f\" % float(gE.norm() / g_obs.norm()))" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "CHECK|exact|E[R_obs] = aFA + (1-aFR-aFA) P(correct) (value)|0.000e+00|tol=1.0e-15|PASS\n", + "CHECK|exact|constant error rates: grad E[R_obs] = (1-aFR-aFA) grad P(correct) (exact rescaling)|2.429e-17|tol=1.0e-14|PASS\n", + " exploit class (always accepted although wrong): P = 0.2462; ||grad P(correct)|| = 0.1825; ||grad P(exploit)|| = 0.2786; cosine = +0.083\n", + "CHECK|exact|observed-reward gradient (autograd) == grad P(correct) + grad P(exploit)|2.776e-17|tol=1.0e-14|PASS\n", + " exploit share of the observed-reward gradient norm: 0.81\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 7. Group-size / pass@k arithmetic: how many samples until a group is live; the cost of rare successes" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": {}, + "source": [ + "print(\" p P(live), G=8 G=16 G=64 expected prompts*samples per live group (G=16)\")\n", + "for pp in (0.001, 0.01, 0.05, 0.2, 0.5):\n", + " live = [1 - pp ** G - (1 - pp) ** G for G in (8, 16, 64)]\n", + " print(f\" {pp:5.3f} {live[0]:8.4f} {live[1]:8.4f} {live[2]:8.4f} {16/live[1]:10.1f}\")\n", + "check(\"P(live) = 1 - p^G - (1-p)^G = pass@G - p^G (identity)\", abs((1 - 0.2 ** 16 - 0.8 ** 16) - ((1 - 0.8 ** 16) - 0.2 ** 16)), 1e-15)\n", + "print(\"done; checks:\", len(CHECKS))" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + " p P(live), G=8 G=16 G=64 expected prompts*samples per live group (G=16)\n", + " 0.001 0.0080 0.0159 0.0620 1007.5\n", + " 0.010 0.0773 0.1485 0.4744 107.7\n", + " 0.050 0.3366 0.5599 0.9625 28.6\n", + " 0.200 0.8322 0.9719 1.0000 16.5\n", + " 0.500 0.9922 1.0000 1.0000 16.0\n", + "CHECK|exact|P(live) = 1 - p^G - (1-p)^G = pass@G - p^G (identity)|0.000e+00|tol=1.0e-15|PASS\n", + "done; checks: 40\n" + ] + } + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.12.11" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} \ No newline at end of file diff --git a/code/notebooks/0201-verifier-binary-reward-mdp.ipynb b/code/notebooks/0201-verifier-binary-reward-mdp.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..f4b114492de999ac953b0cd30dd96619bb318a7e --- /dev/null +++ b/code/notebooks/0201-verifier-binary-reward-mdp.ipynb @@ -0,0 +1,641 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "b4fbe81d", + "metadata": {}, + "source": [ + "# 0201 -- Verifier-in-the-loop as a binary-reward MDP\n", + "\n", + "Scratch-pad for note `0201`. Everything is fp64, CPU only (`CUDA_VISIBLE_DEVICES=` is set\n", + "before torch is imported). Each section states the identity, computes it the *naive* way\n", + "(full enumeration, explicit hand-derived formula), then the *efficient* way (autograd on a\n", + "log-sum-exp-stable policy, Monte Carlo), and asserts that they agree.\n", + "\n", + "Identities verified:\n", + "1. Score-function gradient `grad J = E[r grad log pi]` against autograd of the exact `J`.\n", + "2. Token decomposition `grad log pi(y|x) = sum_t grad log pi(y_t|s_t)` and the hand formula `onehot - softmax`.\n", + "3. Baseline identity `E[grad log pi] = 0`, so any `b(x)` leaves the gradient unbiased.\n", + "4. Exact variance of the one-sample estimator and the optimal constant baseline `b*`.\n", + "5. Monte Carlo estimator converges at the `1/sqrt(N)` rate.\n", + "6. Verifier holes: `E[G] = E[V] - pi(H)` and the softmax-PG ODE `dz_k/dt = pi_k (r_k - J)`.\n", + "7. Proxy reward model (noisy) vs exact verifier: best-of-n selection error and the `sigma sqrt(2 ln n)` winner's curse." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "3b78deed", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T08:49:39.638799Z", + "iopub.status.busy": "2026-10-14T08:49:39.638648Z", + "iopub.status.idle": "2026-10-14T08:49:41.790359Z", + "shell.execute_reply": "2026-10-14T08:49:41.789459Z" + } + }, + "outputs": [], + "source": [ + "import os\n", + "os.environ['CUDA_VISIBLE_DEVICES'] = ''\n", + "import itertools, math, json\n", + "import numpy as np, torch\n", + "torch.set_default_dtype(torch.float64)\n", + "torch.manual_seed(0); np.random.seed(0)\n", + "CHECKS = []\n", + "def check(name, err, tol):\n", + " CHECKS.append((name, float(err), float(tol)))\n", + " assert err <= tol, (name, float(err), tol)\n", + " print(f'{name:70s} err={float(err):.3e} (tol {tol:.0e})')" + ] + }, + { + "cell_type": "markdown", + "id": "0d637e40", + "metadata": {}, + "source": [ + "## 1. A tabular autoregressive policy over proofs\n", + "Vocabulary `{0,1,2}`, proof length `T=3` => 27 \"proofs\". The policy has one logit vector per prefix\n", + "(13 prefixes of length 0,1,2). The verifier `V(y)` is a fixed random 0/1 table with about 30% ones." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "cf164f87", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T08:49:41.792633Z", + "iopub.status.busy": "2026-10-14T08:49:41.792359Z", + "iopub.status.idle": "2026-10-14T08:49:41.884373Z", + "shell.execute_reply": "2026-10-14T08:49:41.883429Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "num prefixes 13 num sequences 27 p = E[r] at init will be printed below\n", + "J(theta0) = 0.30454236036997673\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/home/vignesh/.opencode/tmp/ipykernel_442065/734446241.py:26: UserWarning: Converting a tensor with requires_grad=True to a scalar may lead to unexpected behavior.\n", + "Consider using tensor.detach() first. (Triggered internally at /pytorch/torch/csrc/autograd/generated/python_variable_methods.cpp:838.)\n", + " print('J(theta0) =', float(J))\n" + ] + } + ], + "source": [ + "Vsz, T = 3, 3\n", + "prefixes = [()]\n", + "for L in range(1, T):\n", + " prefixes += list(itertools.product(range(Vsz), repeat=L))\n", + "pid = {p: i for i, p in enumerate(prefixes)}\n", + "seqs = list(itertools.product(range(Vsz), repeat=T))\n", + "rng = np.random.RandomState(3)\n", + "r_tab = torch.tensor((rng.rand(len(seqs)) < 0.3).astype(float)) # V(y) in {0,1}\n", + "print('num prefixes', len(prefixes), 'num sequences', len(seqs), 'p = E[r] at init will be printed below')\n", + "\n", + "theta0 = torch.randn(len(prefixes), Vsz) * 0.7\n", + "\n", + "def log_pi_seq(theta, y):\n", + " \"\"\"log pi(y) = sum_t log softmax(theta[prefix_t])[y_t] (log_softmax is the stable form).\"\"\"\n", + " lp = 0.0\n", + " for t in range(T):\n", + " lp = lp + torch.log_softmax(theta[pid[y[:t]]], dim=-1)[y[t]]\n", + " return lp\n", + "\n", + "def exact_J(theta):\n", + " return sum(torch.exp(log_pi_seq(theta, y)) * r_tab[i] for i, y in enumerate(seqs))\n", + "\n", + "th = theta0.clone().requires_grad_(True)\n", + "J = exact_J(th); J.backward()\n", + "g_auto = th.grad.clone()\n", + "print('J(theta0) =', float(J))" + ] + }, + { + "cell_type": "markdown", + "id": "78aebc59", + "metadata": {}, + "source": [ + "## 2. NAIVE hand-derived gradient: sum over all 27 sequences\n", + "`grad J = sum_y pi(y) r(y) sum_t grad log pi(y_t | s_t)`, with\n", + "`grad_{theta[s,:]} log pi(y_t|s_t) = onehot(y_t) - softmax(theta[s,:])` (the softmax-score identity, derived in the note)." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "c3545a51", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T08:49:41.885962Z", + "iopub.status.busy": "2026-10-14T08:49:41.885798Z", + "iopub.status.idle": "2026-10-14T08:49:41.895143Z", + "shell.execute_reply": "2026-10-14T08:49:41.894090Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "1. score-function grad (hand, enumerated) vs autograd of exact J err=1.041e-17 (tol 1e-14)\n" + ] + } + ], + "source": [ + "def hand_grad(theta):\n", + " th_ = theta.detach()\n", + " g = torch.zeros_like(th_)\n", + " for i, y in enumerate(seqs):\n", + " py = float(torch.exp(log_pi_seq(th_, y)))\n", + " for t in range(T):\n", + " s = pid[y[:t]]\n", + " sm = torch.softmax(th_[s], dim=-1)\n", + " oh = torch.zeros(Vsz); oh[y[t]] = 1.0\n", + " g[s] += py * r_tab[i] * (oh - sm)\n", + " return g\n", + "\n", + "g_hand = hand_grad(theta0)\n", + "check('1. score-function grad (hand, enumerated) vs autograd of exact J', (g_hand - g_auto).abs().max(), 1e-14)" + ] + }, + { + "cell_type": "markdown", + "id": "da5af085", + "metadata": {}, + "source": [ + "## 3. Baseline identity: `E[grad log pi] = 0`\n", + "Because `sum_y pi(y) = 1` for every theta, differentiate it: `0 = sum_y grad pi(y) = sum_y pi(y) grad log pi(y)`." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "097f8049", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T08:49:41.896729Z", + "iopub.status.busy": "2026-10-14T08:49:41.896575Z", + "iopub.status.idle": "2026-10-14T08:49:41.926284Z", + "shell.execute_reply": "2026-10-14T08:49:41.924902Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "3. E[grad log pi] = 0 (baseline identity) err=5.898e-17 (tol 1e-15)\n", + "3. grad unchanged for b in {-3,0,0.3,17} err=6.557e-16 (tol 1e-13)\n" + ] + } + ], + "source": [ + "def expected_score(theta):\n", + " th_ = theta.detach(); g = torch.zeros_like(th_)\n", + " for i, y in enumerate(seqs):\n", + " py = float(torch.exp(log_pi_seq(th_, y)))\n", + " for t in range(T):\n", + " s = pid[y[:t]]; sm = torch.softmax(th_[s], dim=-1)\n", + " oh = torch.zeros(Vsz); oh[y[t]] = 1.0\n", + " g[s] += py * (oh - sm)\n", + " return g\n", + "check('3. E[grad log pi] = 0 (baseline identity)', expected_score(theta0).abs().max(), 1e-15)\n", + "\n", + "# Baseline b constant in y: the gradient E[(r-b) grad log pi] is the same for any b.\n", + "def hand_grad_b(theta, b):\n", + " th_ = theta.detach(); g = torch.zeros_like(th_)\n", + " for i, y in enumerate(seqs):\n", + " py = float(torch.exp(log_pi_seq(th_, y)))\n", + " for t in range(T):\n", + " s = pid[y[:t]]; sm = torch.softmax(th_[s], dim=-1)\n", + " oh = torch.zeros(Vsz); oh[y[t]] = 1.0\n", + " g[s] += py * (r_tab[i] - b) * (oh - sm)\n", + " return g\n", + "errs = [(hand_grad_b(theta0, b) - g_auto).abs().max() for b in (-3.0, 0.0, 0.3, 17.0)]\n", + "check('3. grad unchanged for b in {-3,0,0.3,17}', max(errs), 1e-13)" + ] + }, + { + "cell_type": "markdown", + "id": "30ffac82", + "metadata": {}, + "source": [ + "## 4. Exact variance of the one-sample estimator and the optimal constant baseline\n", + "`Var[(r-b) s] = E[(r-b)^2 |s|^2] - |grad J|^2` with `s = grad log pi(y)` flattened. Minimiser:\n", + "`b* = E[r |s|^2] / E[|s|^2]`." + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "944489cb", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T08:49:41.928225Z", + "iopub.status.busy": "2026-10-14T08:49:41.927849Z", + "iopub.status.idle": "2026-10-14T08:49:43.333887Z", + "shell.execute_reply": "2026-10-14T08:49:43.332721Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "p = E[r] = 0.3045423603699767 b* = 0.33818921921147005 Var(b=0) = 0.5990480197782424 Var(b=p) = 0.39409891194469415 Var(b*) = 0.3920499407669431\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "4. argmin_b Var (grid) vs closed-form b* err=1.078e-05 (tol 2e-04)\n", + "4. Var(b*) closed form vs direct err=5.551e-17 (tol 1e-12)\n", + "4. MC variance of one-sample estimator vs exact (rel. err.) err=2.500e-03 (tol 2e-02)\n" + ] + } + ], + "source": [ + "def score_vec(theta, y):\n", + " th_ = theta.detach(); v = torch.zeros_like(th_)\n", + " for t in range(T):\n", + " s = pid[y[:t]]; sm = torch.softmax(th_[s], dim=-1)\n", + " oh = torch.zeros(Vsz); oh[y[t]] = 1.0\n", + " v[s] += oh - sm\n", + " return v.flatten()\n", + "\n", + "P = torch.stack([torch.exp(log_pi_seq(theta0, y)).detach() for y in seqs])\n", + "S = torch.stack([score_vec(theta0, y) for y in seqs]) # (27, 39)\n", + "R = r_tab\n", + "gJ = (P[:, None] * R[:, None] * S).sum(0)\n", + "def var_exact(b):\n", + " w = (R - b) ** 2 * (S ** 2).sum(1)\n", + " return float((P * w).sum() - (gJ ** 2).sum())\n", + "bstar = float((P * R * (S ** 2).sum(1)).sum() / (P * (S ** 2).sum(1)).sum())\n", + "pbar = float((P * R).sum())\n", + "print('p = E[r] =', pbar, ' b* =', bstar, ' Var(b=0) =', var_exact(0.0), ' Var(b=p) =', var_exact(pbar), ' Var(b*) =', var_exact(bstar))\n", + "bs = np.linspace(-1, 2, 30001)\n", + "vs = np.array([var_exact(b) for b in bs])\n", + "check('4. argmin_b Var (grid) vs closed-form b*', abs(bs[vs.argmin()] - bstar), 2e-4)\n", + "# Closed form of the minimum: Var(b*) = E[r^2|s|^2] - (E[r|s|^2])^2/E|s|^2 - |gJ|^2 (r^2=r for binary r)\n", + "Es2 = float((P * (S ** 2).sum(1)).sum())\n", + "Ers2 = float((P * R * (S ** 2).sum(1)).sum())\n", + "var_min_closed = Ers2 - Ers2 ** 2 / Es2 - float((gJ ** 2).sum())\n", + "check('4. Var(b*) closed form vs direct', abs(var_min_closed - var_exact(bstar)), 1e-12)\n", + "# Monte Carlo of the one-sample estimator (independent check of var_exact at b = p)\n", + "N = 400000\n", + "idx = torch.multinomial(P, N, replacement=True)\n", + "est = (R[idx] - pbar)[:, None] * S[idx]\n", + "var_mc = float(((est - est.mean(0)) ** 2).sum(1).mean())\n", + "check('4. MC variance of one-sample estimator vs exact (rel. err.)', abs(var_mc - var_exact(pbar)) / var_exact(pbar), 0.02)" + ] + }, + { + "cell_type": "markdown", + "id": "db9b1afb", + "metadata": {}, + "source": [ + "## 5. Monte Carlo REINFORCE converges at the `1/sqrt(N)` rate\n", + "EFFICIENT version: sample `N` proofs from `pi`, average `(r - b) * score`. We compare with `grad J` in units of the\n", + "estimated standard error (z-scores), which must be O(1), and the RMS error must halve when `N` is quadrupled." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "64e6d7e6", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T08:49:43.335763Z", + "iopub.status.busy": "2026-10-14T08:49:43.335596Z", + "iopub.status.idle": "2026-10-14T08:49:44.334151Z", + "shell.execute_reply": "2026-10-14T08:49:44.332934Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "max |z| over 39 coordinates: 2.820402064782768\n", + "5. MC grad within 5 standard errors (max |z|, tol 5) err=2.820e+00 (tol 5e+00)\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "RMS err N=5k: 0.008952773341135103 N=20k: 0.004564106878365233 ratio: 1.9615608441539822\n", + "5. rate: RMS(N)/RMS(4N) = 2 (abs dev) err=3.844e-02 (tol 3e-01)\n" + ] + } + ], + "source": [ + "def mc_grad(N, b, seed):\n", + " gen = torch.Generator().manual_seed(seed)\n", + " idx = torch.multinomial(P, N, replacement=True, generator=gen)\n", + " e = (R[idx] - b)[:, None] * S[idx]\n", + " return e.mean(0), e.std(0) / math.sqrt(N)\n", + "m, se = mc_grad(2_000_000, pbar, 11)\n", + "z = ((m - gJ) / se.clamp_min(1e-300))\n", + "mask = se > 1e-9\n", + "print('max |z| over', int(mask.sum()), 'coordinates:', float(z[mask].abs().max()))\n", + "check('5. MC grad within 5 standard errors (max |z|, tol 5)', float(z[mask].abs().max()), 5.0)\n", + "rms = lambda N: float(np.sqrt(np.mean([((mc_grad(N, pbar, s)[0] - gJ) ** 2).sum().item() for s in range(60)])))\n", + "r1, r4 = rms(5_000), rms(20_000)\n", + "print('RMS err N=5k:', r1, ' N=20k:', r4, ' ratio:', r1 / r4)\n", + "check('5. rate: RMS(N)/RMS(4N) = 2 (abs dev)', abs(r1 / r4 - 2.0), 0.35)" + ] + }, + { + "cell_type": "markdown", + "id": "5a6f8e4c", + "metadata": {}, + "source": [ + "## 6. The hole: `E[G] = E[V] - pi(H)` and softmax policy-gradient dynamics\n", + "Three kinds of outputs: `wrong` (V=0,G=0), `proof` (V=1,G=1), `hole` (V=1, G=0, e.g. a `sorry`-like exploit that the\n", + "verifier accepts). For a tabular softmax over these classes, `dz_k/dt = pi_k (r_k - J)` (derived in the note)." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "a8b49842", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T08:49:44.336095Z", + "iopub.status.busy": "2026-10-14T08:49:44.335921Z", + "iopub.status.idle": "2026-10-14T08:49:45.190504Z", + "shell.execute_reply": "2026-10-14T08:49:45.189496Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "6. softmax-PG gradient pi_k (r_k - J) vs autograd err=0.000e+00 (tol 1e-16)\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "step, E[V], E[G], pi(hole)\n", + " (0, 0.04, 0.01, 0.03)\n", + " (4000, 0.9997460159631938, 0.0017478329594184914, 0.9979981830037753)\n", + " (8000, 0.9998739033985322, 0.0012328930126125394, 0.9986410103859197)\n", + " (12000, 0.9999161452892915, 0.001005875916774838, 0.9989102693725167)\n", + " (16000, 0.9999371907590497, 0.0008707924515763814, 0.9990663983074733)\n", + " (20000, 0.9999497931524061, 0.0007786947944058345, 0.9991710983580002)\n", + "6. identity E[G] = E[V] - pi(H) at the end of training err=1.735e-17 (tol 1e-15)\n", + "final: E[V]=0.9999 E[G]=0.0008 pi(hole)=0.9992 (the true-proof share of the accepted mass fell from 0.25 to 0.001)\n", + "6. proof share of accepted mass is non-increasing (max positive increment) err=0.000e+00 (tol 1e-15)\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "6. false-reject bias: E[V] - (1-delta)E[G] within 4 s.e. err=3.385e-01 (tol 4e+00)\n" + ] + } + ], + "source": [ + "r_V = torch.tensor([0.0, 1.0, 1.0]); r_G = torch.tensor([0.0, 1.0, 0.0])\n", + "z = torch.log(torch.tensor([0.96, 0.01, 0.03])) # hole is 3x more reachable than the true proof\n", + "# (a) gradient of J_V wrt logits: autograd vs hand formula pi_k (r_k - J)\n", + "zz = z.clone().requires_grad_(True)\n", + "pi = torch.softmax(zz, 0); JV = (pi * r_V).sum(); JV.backward()\n", + "hand = pi.detach() * (r_V - JV.detach())\n", + "check('6. softmax-PG gradient pi_k (r_k - J) vs autograd', (zz.grad - hand).abs().max(), 1e-16)\n", + "# (b) integrate the gradient flow with small steps; track J_V, J_G, and pi(H)\n", + "zt = z.clone(); eta = 0.5; traj = []\n", + "for step in range(20001):\n", + " pi = torch.softmax(zt, 0)\n", + " JVt = float((pi * r_V).sum()); JGt = float((pi * r_G).sum())\n", + " if step % 4000 == 0: traj.append((step, JVt, JGt, float(pi[2])))\n", + " zt = zt + eta * pi * (r_V - JVt)\n", + "print('step, E[V], E[G], pi(hole)'); [print(' ', t) for t in traj]\n", + "pi = torch.softmax(zt, 0)\n", + "EV = float((pi * r_V).sum()); EG = float((pi * r_G).sum()); piH = float(pi[2])\n", + "check('6. identity E[G] = E[V] - pi(H) at the end of training', abs(EG - (EV - piH)), 1e-15)\n", + "assert EV > 0.99 and EG < 0.3, (EV, EG)\n", + "print(f'final: E[V]={EV:.4f} E[G]={EG:.4f} pi(hole)={piH:.4f} (the true-proof share of the accepted mass fell from 0.25 to {EG/EV:.3f})')\n", + "# Closed form for the invariant: log(pi_proof/pi_hole) changes by sum (pi_proof - pi_hole)(1-J) eta; since pi_hole > pi_proof it\n", + "# keeps growing in favour of the hole. Check monotonicity of the proof share along the trajectory.\n", + "zt = z.clone(); shares = []\n", + "for step in range(3000):\n", + " pi = torch.softmax(zt, 0); shares.append(float(pi[1] / (pi[1] + pi[2])))\n", + " zt = zt + eta * pi * (r_V - (pi * r_V).sum())\n", + "check('6. proof share of accepted mass is non-increasing (max positive increment)', max(0.0, float(np.max(np.diff(shares)))), 1e-15)\n", + "\n", + "# (c) false rejects: V = G * (1 - timeout) with independent timeouts of prob delta => E[V] = (1-delta) E[G]\n", + "delta = 0.1; N = 4_000_000\n", + "gen = torch.Generator().manual_seed(5)\n", + "G = (torch.rand(N, generator=gen) < 0.37).double(); to = (torch.rand(N, generator=gen) < delta).double()\n", + "Vv = G * (1 - to)\n", + "se = float(Vv.std() / math.sqrt(N))\n", + "check('6. false-reject bias: E[V] - (1-delta)E[G] within 4 s.e.', abs(float(Vv.mean()) - (1 - delta) * 0.37) / se, 4.0)" + ] + }, + { + "cell_type": "markdown", + "id": "5a8054aa", + "metadata": {}, + "source": [ + "## 7. Proxy reward model vs exact verifier at best-of-`n`\n", + "`n` candidates, each truly correct w.p. `p`. A proxy RM scores `s = G + sigma * xi`, `xi ~ N(0,1)`. Selecting the argmax of `s`\n", + "picks a correct candidate with probability\n", + "`P_RM(n) = int n p f_1(s) [p F_1(s) + (1-p) F_0(s)]^(n-1) ds`\n", + "(`f_1,F_1` the N(1,sigma^2) pdf/cdf, `f_0,F_0` the N(0,sigma^2) ones), while the exact verifier gets\n", + "`1-(1-p)^n`. The *score* the RM reports for its pick exceeds the true value by about `sigma sqrt(2 ln n)`." + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "7cea16ac", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T08:49:45.192317Z", + "iopub.status.busy": "2026-10-14T08:49:45.192135Z", + "iopub.status.idle": "2026-10-14T08:51:27.965442Z", + "shell.execute_reply": "2026-10-14T08:51:27.964619Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "n= 4 p=0.2 sigma=0.5: P_RM exact=0.51088 MC=0.51153 verifier(any correct)=0.59040\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "n= 16 p=0.1 sigma=1.0: P_RM exact=0.33706 MC=0.33767 verifier(any correct)=0.81470\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "n= 64 p=0.05 sigma=1.5: P_RM exact=0.17838 MC=0.17813 verifier(any correct)=0.96248\n", + "7. best-of-n pick-correct probability: quadrature vs MC (max z) err=8.216e-01 (tol 4e+00)\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "verifier minus RM at n=16,p=0.1 for sigma=0.25,0.5,1,2: ['0.0089', '0.2043', '0.4776', '0.6137']\n", + "n= 10 E[max] exact=1.5388 upper bound sqrt(2 ln n)=2.1460\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "n= 100 E[max] exact=2.5076 upper bound sqrt(2 ln n)=3.0349\n", + "n= 10000 E[max] exact=3.8516 upper bound sqrt(2 ln n)=4.2919\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "n= 1000000 E[max] exact=4.8629 upper bound sqrt(2 ln n)=5.2565\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "7. E[max of 100 N(0,1)]: quadrature vs MC (in standard errors, tol 4) err=1.327e+00 (tol 4e+00)\n" + ] + } + ], + "source": [ + "from math import erf, sqrt, exp, pi as PI\n", + "def Phi(x): return 0.5 * (1 + np.vectorize(erf)(x / sqrt(2)))\n", + "def p_rm(n, p, sigma, grid=400001):\n", + " s = np.linspace(-12 * sigma - 2, 12 * sigma + 3, grid)\n", + " f1 = np.exp(-0.5 * ((s - 1) / sigma) ** 2) / (sigma * sqrt(2 * PI))\n", + " F1 = Phi((s - 1) / sigma); F0 = Phi(s / sigma)\n", + " integrand = n * p * f1 * (p * F1 + (1 - p) * F0) ** (n - 1)\n", + " return float(np.trapezoid(integrand, s))\n", + "def p_rm_mc(n, p, sigma, trials, seed):\n", + " g = np.random.RandomState(seed)\n", + " G = (g.rand(trials, n) < p).astype(float); sc = G + sigma * g.randn(trials, n)\n", + " pick = sc.argmax(1)\n", + " return G[np.arange(trials), pick].mean(), sc.max(1).mean(), G[np.arange(trials), pick].mean()\n", + "errs = []\n", + "for (n, p, sig) in [(4, 0.2, 0.5), (16, 0.1, 1.0), (64, 0.05, 1.5)]:\n", + " exact = p_rm(n, p, sig); mc, _, _ = p_rm_mc(n, p, sig, 400000, 1)\n", + " se = math.sqrt(exact * (1 - exact) / 400000)\n", + " errs.append(abs(exact - mc) / se)\n", + " print(f'n={n:3d} p={p} sigma={sig}: P_RM exact={exact:.5f} MC={mc:.5f} verifier(any correct)={1-(1-p)**n:.5f}')\n", + "check('7. best-of-n pick-correct probability: quadrature vs MC (max z)', max(errs), 4.0)\n", + "# Verifier is never worse and strictly better for sigma>0\n", + "gap = [1 - (1 - 0.1) ** 16 - p_rm(16, 0.1, s) for s in (0.25, 0.5, 1.0, 2.0)]\n", + "print('verifier minus RM at n=16,p=0.1 for sigma=0.25,0.5,1,2:', [f'{x:.4f}' for x in gap])\n", + "assert all(g > 0 for g in gap) and all(np.diff(gap) > 0)\n", + "# Winner's curse: E[max of n iid N(0,1)] = int n x phi(x) Phi(x)^(n-1) dx (exact quadrature); bound E[max] <= sqrt(2 ln n).\n", + "def emax_exact(n):\n", + " x = np.linspace(-10, 10, 400001)\n", + " phi = np.exp(-0.5 * x ** 2) / sqrt(2 * PI)\n", + " return float(np.trapezoid(n * x * phi * Phi(x) ** (n - 1), x))\n", + "for n in (10, 100, 10_000, 1_000_000):\n", + " print(f'n={n:8d} E[max] exact={emax_exact(n):.4f} upper bound sqrt(2 ln n)={sqrt(2*math.log(n)):.4f}')\n", + " assert emax_exact(n) <= sqrt(2 * math.log(n))\n", + "mcs = np.random.RandomState(2).randn(40000, 100).max(1)\n", + "check('7. E[max of 100 N(0,1)]: quadrature vs MC (in standard errors, tol 4)', abs(emax_exact(100) - mcs.mean()) / (mcs.std() / sqrt(len(mcs))), 4.0)" + ] + }, + { + "cell_type": "markdown", + "id": "f1a420d1", + "metadata": {}, + "source": [ + "## Verification table" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "3c054d00", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T08:51:27.967075Z", + "iopub.status.busy": "2026-10-14T08:51:27.966912Z", + "iopub.status.idle": "2026-10-14T08:51:27.970746Z", + "shell.execute_reply": "2026-10-14T08:51:27.970067Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + "identity max abs err\n", + "1. score-function grad (hand, enumerated) vs autograd of exact J 1.041e-17\n", + "3. E[grad log pi] = 0 (baseline identity) 5.898e-17\n", + "3. grad unchanged for b in {-3,0,0.3,17} 6.557e-16\n", + "4. argmin_b Var (grid) vs closed-form b* 1.078e-05\n", + "4. Var(b*) closed form vs direct 5.551e-17\n", + "4. MC variance of one-sample estimator vs exact (rel. err.) 2.500e-03\n", + "5. MC grad within 5 standard errors (max |z|, tol 5) 2.820e+00\n", + "5. rate: RMS(N)/RMS(4N) = 2 (abs dev) 3.844e-02\n", + "6. softmax-PG gradient pi_k (r_k - J) vs autograd 0.000e+00\n", + "6. identity E[G] = E[V] - pi(H) at the end of training 1.735e-17\n", + "6. proof share of accepted mass is non-increasing (max positive increment) 0.000e+00\n", + "6. false-reject bias: E[V] - (1-delta)E[G] within 4 s.e. 3.385e-01\n", + "7. best-of-n pick-correct probability: quadrature vs MC (max z) 8.216e-01\n", + "7. E[max of 100 N(0,1)]: quadrature vs MC (in standard errors, tol 4) 1.327e+00\n", + "CHECKS_JSON:[[\"1. score-function grad (hand, enumerated) vs autograd of exact J\", 1.0408340855860843e-17, 1e-14], [\"3. E[grad log pi] = 0 (baseline identity)\", 5.898059818321144e-17, 1e-15], [\"3. grad unchanged for b in {-3,0,0.3,17}\", 6.557254739192331e-16, 1e-13], [\"4. argmin_b Var (grid) vs closed-form b*\", 1.078078853000708e-05, 0.0002], [\"4. Var(b*) closed form vs direct\", 5.551115123125783e-17, 1e-12], [\"4. MC variance of one-sample estimator vs exact (rel. err.)\", 0.0025004228417077347, 0.02], [\"5. MC grad within 5 standard errors (max |z|, tol 5)\", 2.820402064782768, 5.0], [\"5. rate: RMS(N)/RMS(4N) = 2 (abs dev)\", 0.03843915584601776, 0.35], [\"6. softmax-PG gradient pi_k (r_k - J) vs autograd\", 0.0, 1e-16], [\"6. identity E[G] = E[V] - pi(H) at the end of training\", 1.734723475976807e-17, 1e-15], [\"6. proof share of accepted mass is non-increasing (max positive increment)\", 0.0, 1e-15], [\"6. false-reject bias: E[V] - (1-delta)E[G] within 4 s.e.\", 0.33845555650118325, 4.0], [\"7. best-of-n pick-correct probability: quadrature vs MC (max z)\", 0.8216357277155372, 4.0], [\"7. E[max of 100 N(0,1)]: quadrature vs MC (in standard errors, tol 4)\", 1.327150492424622, 4.0]]\n" + ] + } + ], + "source": [ + "print()\n", + "print(f'{\"identity\":72s} {\"max abs err\":>12s}')\n", + "for n_, e_, t_ in CHECKS: print(f'{n_:72s} {e_:12.3e}')\n", + "print('CHECKS_JSON:' + json.dumps(CHECKS))" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/code/notebooks/0202-rejection-sampling-expert-iteration-em.ipynb b/code/notebooks/0202-rejection-sampling-expert-iteration-em.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..6cfe6ede99f9abe4563af7c951bf54e31c9df25e --- /dev/null +++ b/code/notebooks/0202-rejection-sampling-expert-iteration-em.ipynb @@ -0,0 +1,555 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "9658b1ce", + "metadata": {}, + "source": [ + "# 0202 -- Rejection sampling / expert iteration as policy improvement (and as EM)\n", + "\n", + "Identities verified (fp64, CPU):\n", + "1. `grad_theta E_x E_{y~pi_t}[V log pi_theta]` at `theta=theta_t` equals the policy gradient `grad J(theta_t)`.\n", + "2. ELBO identity `L(theta,q) = log p_theta(x) - KL(q || q*_theta)` and tightness at `q=q*_theta` (so the E-step is exact).\n", + "3. Tabular M-step: `argmin KL(q_t || pi_theta)` gives `pi_{t+1}=q_t`, hence `p_{t+1}(x)=1` where `p_t(x)>0` and `0` where `p_t(x)=0` (coverage).\n", + "4. Shared-parameter model: with an *exact* M-step, `F(theta)=sum_x log p_theta(x)` is non-decreasing (EM monotonicity).\n", + "5. Sampling: `E[ sum_i V_i f(y_i)/c ; c>0 ] = (1-(1-p)^n) E_q[f]` (exact enumeration) -- the self-normalised filtered average is unbiased except for the missing event `c=0`.\n", + "6. Unnormalised (practical) SFT objective is *not* monotone per prompt under shared parameters.\n", + "7. Monte Carlo RFT gradient agrees with the exact `grad J` (RFT step = REINFORCE step)." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "86d1f2a7", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:47:57.273949Z", + "iopub.status.busy": "2026-10-14T09:47:57.273796Z", + "iopub.status.idle": "2026-10-14T09:47:59.029734Z", + "shell.execute_reply": "2026-10-14T09:47:59.029103Z" + } + }, + "outputs": [], + "source": [ + "import os\n", + "os.environ['CUDA_VISIBLE_DEVICES'] = ''\n", + "import itertools, math, json\n", + "import numpy as np, torch\n", + "torch.set_default_dtype(torch.float64)\n", + "torch.manual_seed(1); np.random.seed(1)\n", + "CHECKS = []\n", + "def check(name, err, tol):\n", + " CHECKS.append((name, float(err), float(tol)))\n", + " assert err <= tol, (name, float(err), tol)\n", + " print(f'{name:78s} err={float(err):.3e} (tol {tol:.0e})')" + ] + }, + { + "cell_type": "markdown", + "id": "ee2a1501", + "metadata": {}, + "source": [ + "## 1. A small family of prompts with a tabular policy\n", + "`m=6` prompts, `K=10` candidate proofs per prompt. `V[x,y]` in {0,1}. Prompt 5 has *no* correct proof; prompt 4 has a correct\n", + "proof that the initial policy almost never produces (coverage problem)." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "5af4e833", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:47:59.032313Z", + "iopub.status.busy": "2026-10-14T09:47:59.032078Z", + "iopub.status.idle": "2026-10-14T09:47:59.039690Z", + "shell.execute_reply": "2026-10-14T09:47:59.038725Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "p_t(x) = [0.26066, 0.36631, 0.17998, 0.16817, 0.00018, 0.0]\n" + ] + } + ], + "source": [ + "m, K = 6, 10\n", + "rng = np.random.RandomState(7)\n", + "V = (rng.rand(m, K) < 0.25).astype(float)\n", + "V[5, :] = 0.0 # unsolvable: V == 0 everywhere\n", + "V[4, :] = 0.0; V[4, 9] = 1.0 # exactly one correct proof\n", + "V[0, 0] = 1.0 # make sure prompt 0 is solvable\n", + "V = torch.tensor(V)\n", + "logits0 = torch.randn(m, K) * 0.5\n", + "logits0[4, 9] = -6.0 # pi(correct) ~ 2e-4 for prompt 4\n", + "def pi_tab(z): return torch.softmax(z, dim=-1)\n", + "def p_of(z): return (pi_tab(z) * V).sum(-1)\n", + "print('p_t(x) =', [round(float(v), 5) for v in p_of(logits0)])" + ] + }, + { + "cell_type": "markdown", + "id": "7bf5412a", + "metadata": {}, + "source": [ + "## 2. First-order equivalence: one RFT step is a policy-gradient step\n", + "RFT objective with samples from the *frozen* `pi_t`: `L(theta) = (1/m) sum_x sum_y pi_t(y|x) V(x,y) log pi_theta(y|x)`.\n", + "Its gradient at `theta=theta_t` is `(1/m) sum_x sum_y V pi_t grad log pi_t = grad J(theta_t)` (score identity of note 0201)." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "d29b7d32", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:47:59.041294Z", + "iopub.status.busy": "2026-10-14T09:47:59.041085Z", + "iopub.status.idle": "2026-10-14T09:47:59.128842Z", + "shell.execute_reply": "2026-10-14T09:47:59.127645Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "1. grad of RFT objective at theta_t vs grad J(theta_t) err=3.469e-18 (tol 1e-16)\n" + ] + } + ], + "source": [ + "z = logits0.clone().requires_grad_(True)\n", + "J = p_of(z).mean(); J.backward(); gJ = z.grad.clone()\n", + "pi_t = pi_tab(logits0).detach()\n", + "z2 = logits0.clone().requires_grad_(True)\n", + "L = ((pi_t * V * torch.log_softmax(z2, -1)).sum(-1)).mean(); L.backward()\n", + "check('1. grad of RFT objective at theta_t vs grad J(theta_t)', (z2.grad - gJ).abs().max(), 1e-16)" + ] + }, + { + "cell_type": "markdown", + "id": "93fcb188", + "metadata": {}, + "source": [ + "## 3. ELBO identity and exact E-step\n", + "`p_theta(x) = sum_y pi_theta V`. `q*_theta(y|x) = pi_theta V / p_theta(x)`. For any `q` (full support on {V=1}):\n", + "`L(theta,q) = E_q[log V(x,y)] - KL(q || pi_theta)`, and with `V in {0,1}` and `q` supported on `{V=1}` this is `-KL(q||pi_theta)`.\n", + "Claim: `-KL(q||pi_theta) = log p_theta - KL(q||q*_theta)`." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "5dd4e7bb", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:47:59.130670Z", + "iopub.status.busy": "2026-10-14T09:47:59.130469Z", + "iopub.status.idle": "2026-10-14T09:47:59.137329Z", + "shell.execute_reply": "2026-10-14T09:47:59.136038Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "2. ELBO identity -KL(q||pi) = log p - KL(q||q*) err=0.000e+00 (tol 1e-14)\n", + "2. tightness: -KL(q*||pi) = log p err=2.220e-16 (tol 1e-14)\n" + ] + } + ], + "source": [ + "th = logits0[0]; pi = pi_tab(th); Vx = V[0]; p0 = float((pi * Vx).sum())\n", + "qstar = pi * Vx / p0\n", + "qrand = torch.rand(K) * Vx; qrand = qrand / qrand.sum()\n", + "def KL(a, b):\n", + " msk = a > 0\n", + " return float((a[msk] * (torch.log(a[msk]) - torch.log(b[msk]))).sum())\n", + "lhs = -KL(qrand, pi)\n", + "rhs = math.log(p0) - KL(qrand, qstar)\n", + "check('2. ELBO identity -KL(q||pi) = log p - KL(q||q*)', abs(lhs - rhs), 1e-14)\n", + "check('2. tightness: -KL(q*||pi) = log p', abs(-KL(qstar, pi) - math.log(p0)), 1e-14)" + ] + }, + { + "cell_type": "markdown", + "id": "ae2285be", + "metadata": {}, + "source": [ + "## 4. Tabular M-step: `pi_{t+1} = q_t` and the coverage theorem\n", + "With full capacity `argmin_theta KL(q_t||pi_theta)` gives `pi_{t+1}(.|x)=q_t(.|x)` for every prompt with `p_t(x)>0`.\n", + "One iteration therefore drives `p(x)` to 1 on all covered prompts and leaves `p(x)=0` prompts untouched." + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "d07a7eb0", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:47:59.139284Z", + "iopub.status.busy": "2026-10-14T09:47:59.139078Z", + "iopub.status.idle": "2026-10-14T09:47:59.146022Z", + "shell.execute_reply": "2026-10-14T09:47:59.144837Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "p_t : [0.26066, 0.36631, 0.17998, 0.16817, 0.00018, 0.0]\n", + "p_t+1 : [1.0, 1.0, 1.0, 1.0, 1.0, 0.0]\n", + "3. p_{t+1}(x) = 1 on covered prompts (max |p-1|) err=0.000e+00 (tol 1e-14)\n", + "3. p_{t+1}(x) = 0 on uncovered prompt (coverage failure) err=0.000e+00 (tol 0e+00)\n" + ] + } + ], + "source": [ + "pt = p_of(logits0)\n", + "covered = pt > 0\n", + "q = pi_tab(logits0) * V\n", + "q = torch.where(covered[:, None], q / pt.clamp_min(1e-300)[:, None], pi_tab(logits0)) # keep pi_t where q undefined\n", + "pi_next = q\n", + "p_next = (pi_next * V).sum(-1)\n", + "print('p_t :', [round(float(v), 5) for v in pt])\n", + "print('p_t+1 :', [round(float(v), 5) for v in p_next])\n", + "check('3. p_{t+1}(x) = 1 on covered prompts (max |p-1|)', (p_next[covered] - 1).abs().max(), 1e-14)\n", + "check('3. p_{t+1}(x) = 0 on uncovered prompt (coverage failure)', p_next[~covered].abs().max(), 0.0)" + ] + }, + { + "cell_type": "markdown", + "id": "3670c24a", + "metadata": {}, + "source": [ + "## 5. Shared parameters: EM monotonicity with an exact M-step\n", + "Log-linear policy `pi_theta(y|x) = softmax_y(phi(x,y) . theta)`, `d=5 < K=12` features, so the model cannot represent `q_t` exactly.\n", + "`F(theta) = (1/|S|) sum_{x in S} log p_theta(x)` over the prompts with at least one correct proof (`S`).\n", + "The M-step objective `-(1/|S|) sum_x sum_y q_t(y|x) log pi_theta(y|x)` is convex in theta (softmax cross-entropy with linear logits),\n", + "so L-BFGS solves it to machine precision, i.e. the M-step is exact." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "341bebce", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:47:59.148100Z", + "iopub.status.busy": "2026-10-14T09:47:59.147866Z", + "iopub.status.idle": "2026-10-14T09:48:00.351047Z", + "shell.execute_reply": "2026-10-14T09:48:00.350146Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "F per iteration: [-1.733346, -1.418455, -1.386263, -1.381133, -1.377459, -1.373836, -1.37042, -1.367474, -1.365174, -1.363539, -1.362461, -1.361784, -1.361367]\n", + "mean p per iter: [0.185491, 0.234027, 0.253918, 0.259472, 0.261745, 0.263587, 0.265459, 0.267323, 0.269065, 0.270608, 0.271926, 0.273028, 0.273945]\n", + "final M-step max |grad|: 1.2179568404172003e-09\n", + "4. EM monotone: max decrease of F over 12 iterations (should be <= 0) err=0.000e+00 (tol 1e-10)\n", + "4. M-step solved to optimality (max |grad| at the last M-step) err=1.218e-09 (tol 1e-08)\n", + "4. ELBO(theta_t,q_t) = F(theta_t) at the final iterate err=0.000e+00 (tol 1e-12)\n", + "4. equal-weight M-step objective: gradient at theta_t = grad F(theta_t) (weight 1/p_t per problem on grad p) err=6.245e-17 (tol 1e-14)\n", + "4. pooled objective (ReST-EM Eq.3 / RFT): gradient at theta_t = grad mean_S p_t (weight 1 on grad p) err=1.388e-17 (tol 1e-14)\n", + "relative difference of the two directions (cosine): 0.10628750941093855\n" + ] + } + ], + "source": [ + "m2, K2, d = 8, 12, 5\n", + "g = np.random.RandomState(11)\n", + "Phi = torch.tensor(g.randn(m2, K2, d))\n", + "V2 = (g.rand(m2, K2) < 0.2).astype(float)\n", + "V2[7, :] = 0.0 # unsolvable prompt\n", + "for x in range(7):\n", + " if V2[x].sum() == 0: V2[x, x] = 1.0\n", + "V2 = torch.tensor(V2)\n", + "S = (V2.sum(-1) > 0)\n", + "def logpi(theta): return torch.log_softmax(Phi @ theta, dim=-1) # (m2, K2)\n", + "def p_x(theta): return (torch.exp(logpi(theta)) * V2).sum(-1)\n", + "def F(theta): return torch.log(p_x(theta)[S]).mean()\n", + "def e_step(theta):\n", + " w = torch.exp(logpi(theta)) * V2\n", + " return w / w.sum(-1, keepdim=True).clamp_min(1e-300)\n", + "def m_step(q, theta_init, weights=None):\n", + " th_ = theta_init.clone().requires_grad_(True)\n", + " wts = torch.ones(m2) if weights is None else weights\n", + " opt = torch.optim.LBFGS([th_], lr=1.0, max_iter=500, tolerance_grad=1e-13, tolerance_change=1e-16, line_search_fn='strong_wolfe')\n", + " def closure():\n", + " opt.zero_grad()\n", + " loss = -(wts[S, None] * q[S] * logpi(th_)[S]).sum() / S.sum()\n", + " loss.backward(); return loss\n", + " for _ in range(4): opt.step(closure)\n", + " return th_.detach(), th_.grad.abs().max().item()\n", + "theta = torch.randn(d) * 0.1\n", + "Fs, Js, gmax = [float(F(theta))], [float(p_x(theta).mean())], []\n", + "for it in range(12):\n", + " q = e_step(theta)\n", + " theta, gm = m_step(q, theta); gmax.append(gm)\n", + " Fs.append(float(F(theta))); Js.append(float(p_x(theta).mean()))\n", + "print('F per iteration:', [round(v, 6) for v in Fs])\n", + "print('mean p per iter:', [round(v, 6) for v in Js])\n", + "print('final M-step max |grad|:', gmax[-1])\n", + "check('4. EM monotone: max decrease of F over 12 iterations (should be <= 0)', max(0.0, max(-(b - a) for a, b in zip(Fs[:-1], Fs[1:]))), 1e-10)\n", + "check('4. M-step solved to optimality (max |grad| at the last M-step)', gmax[-1], 1e-8)\n", + "assert Fs[-1] > Fs[0] + 0.05\n", + "\n", + "# Tightness at theta_t, per iteration: ELBO(theta_t, q_t) = F(theta_t)\n", + "q = e_step(theta)\n", + "elbo = -(q[S] * (torch.log(q[S].clamp_min(1e-300)) - logpi(theta)[S]) * (V2[S] > 0)).sum(-1).mean()\n", + "check('4. ELBO(theta_t,q_t) = F(theta_t) at the final iterate', abs(float(elbo) - float(F(theta))), 1e-12)\n", + "\n", + "# Per-problem weights of the two M-step objectives, checked at theta_t (this is the 1 vs 1/p_t(x) claim of the comparison table):\n", + "# equal-weight EM objective (1/|S|) sum_x E_{q_t}[log pi_theta] has gradient (1/|S|) sum_x grad log p_t(x) = grad F(theta_t)\n", + "# pooled RFT / ReST-EM Eq.(3) objective (1/|S|) sum_x p_t(x) E_{q_t}[log pi_theta] has gradient (1/|S|) sum_x grad p_t(x) = grad (mean_S p)\n", + "th_chk = theta.clone().requires_grad_(True)\n", + "qt = e_step(theta).detach()\n", + "(-(qt[S] * logpi(th_chk)[S]).sum() / S.sum()).backward(); g_eq = -th_chk.grad.clone()\n", + "th_f = theta.clone().requires_grad_(True); F(th_f).backward()\n", + "check('4. equal-weight M-step objective: gradient at theta_t = grad F(theta_t) (weight 1/p_t per problem on grad p)', (g_eq - th_f.grad).abs().max(), 1e-14)\n", + "th_chk = theta.clone().requires_grad_(True)\n", + "pw = p_x(theta).detach()\n", + "(-(pw[S, None] * qt[S] * logpi(th_chk)[S]).sum() / S.sum()).backward(); g_pool = -th_chk.grad.clone()\n", + "th_p = theta.clone().requires_grad_(True); p_x(th_p)[S].mean().backward()\n", + "check('4. pooled objective (ReST-EM Eq.3 / RFT): gradient at theta_t = grad mean_S p_t (weight 1 on grad p)', (g_pool - th_p.grad).abs().max(), 1e-14)\n", + "print('relative difference of the two directions (cosine):', float(g_eq @ g_pool / (g_eq.norm() * g_pool.norm())))" + ] + }, + { + "cell_type": "markdown", + "id": "4233ec2d", + "metadata": {}, + "source": [ + "## 6. The practical objective (unnormalised, weights `p_t(x)`) is not monotone per prompt\n", + "SFT on the pooled filtered samples weights prompt `x` by (expected) count of correct samples `n p_t(x)`, i.e. objective\n", + "`sum_x p_t(x) E_{q_t}[log pi_theta]`. With a shared `theta` some prompts get *worse* even though the weighted objective improves." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "fb963f24", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:48:00.353128Z", + "iopub.status.busy": "2026-10-14T09:48:00.352833Z", + "iopub.status.idle": "2026-10-14T09:48:00.421027Z", + "shell.execute_reply": "2026-10-14T09:48:00.420159Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "iter, mean p on solvable prompts, #prompts whose p decreased, F\n", + " (0, 0.2615781869672757, 2, -1.4525900599954458)\n", + " (1, 0.305815927384656, 3, -1.4159837962389783)\n", + " (2, 0.360538810599986, 4, -1.5299987749300417)\n", + " (3, 0.3968721919898336, 4, -1.7930711977638583)\n", + " (4, 0.4128722015161725, 4, -2.0991072756580036)\n", + " (5, 0.42201137651015436, 3, -2.3539650242141725)\n", + "total (prompt, iteration) pairs with decreasing p: 20\n" + ] + } + ], + "source": [ + "theta_u = torch.randn(d) * 0.1\n", + "rows = []\n", + "for it in range(6):\n", + " p_before = p_x(theta_u).clone()\n", + " q = e_step(theta_u)\n", + " theta_u, gm = m_step(q, theta_u, weights=p_before.detach())\n", + " p_after = p_x(theta_u)\n", + " rows.append((it, float(p_after[S].mean()), int((p_after[S] < p_before[S] - 1e-12).sum()), float(F(theta_u))))\n", + "print('iter, mean p on solvable prompts, #prompts whose p decreased, F')\n", + "for r in rows: print(' ', r)\n", + "n_dec = sum(r[2] for r in rows)\n", + "print('total (prompt, iteration) pairs with decreasing p:', n_dec)\n", + "CHECKS.append(('6. non-monotone prompts exist under unnormalised weights (count of decreasing p, informational)', float(n_dec), float('inf')))" + ] + }, + { + "cell_type": "markdown", + "id": "ba81c4f7", + "metadata": {}, + "source": [ + "## 7. Sampling estimator: unbiased except for the event `c = 0`\n", + "One prompt, outcomes `y in {0,1,2,3}`, `V=(0,1,1,0)`, `pi=(.4,.1,.2,.3)`, a test function `f`. Draw `n=3` i.i.d. samples; keep correct ones;\n", + "estimate `E_q[f]` by `(1/c) sum_i V_i f(y_i)` when `c>0` (else 0). Claim: expectation `= (1-(1-p)^n) E_q f` exactly.\n", + "NAIVE: enumerate all `4^3` sample tuples with their probabilities. EFFICIENT: the closed form." + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "52c6844e", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:48:00.422626Z", + "iopub.status.busy": "2026-10-14T09:48:00.422467Z", + "iopub.status.idle": "2026-10-14T09:48:00.441049Z", + "shell.execute_reply": "2026-10-14T09:48:00.440191Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "5. E[filtered mean of f] = (1-(1-p)^n) E_q[f], n=1 err=0.000e+00 (tol 1e-14)\n", + "5. E[filtered mean of f] = (1-(1-p)^n) E_q[f], n=2 err=3.331e-16 (tol 1e-14)\n", + "5. E[filtered mean of f] = (1-(1-p)^n) E_q[f], n=3 err=8.882e-16 (tol 1e-14)\n", + "5. E[filtered mean of f] = (1-(1-p)^n) E_q[f], n=5 err=2.220e-15 (tol 1e-14)\n", + "5. pooled filtered average: E[(1/n) sum V f] = p E_q[f] err=3.331e-16 (tol 1e-14)\n" + ] + } + ], + "source": [ + "pi4 = np.array([.4, .1, .2, .3]); V4 = np.array([0, 1, 1, 0.]); f4 = np.array([1.3, -0.7, 2.9, 5.0])\n", + "p = float((pi4 * V4).sum()); Eq_f = float((pi4 * V4 * f4).sum() / p)\n", + "for n in (1, 2, 3, 5):\n", + " tot = 0.0\n", + " for tup in itertools.product(range(4), repeat=n):\n", + " prob = np.prod([pi4[i] for i in tup]); c = sum(V4[i] for i in tup)\n", + " if c > 0: tot += prob * sum(V4[i] * f4[i] for i in tup) / c\n", + " closed = (1 - (1 - p) ** n) * Eq_f\n", + " check(f'5. E[filtered mean of f] = (1-(1-p)^n) E_q[f], n={n}', abs(tot - closed), 1e-14)\n", + "\n", + "# Pooled (unnormalised) filtered average: E[(1/n) sum_i V_i f(y_i)] = p E_q[f] exactly (linearity).\n", + "n = 4; tot = 0.0\n", + "for tup in itertools.product(range(4), repeat=n):\n", + " prob = np.prod([pi4[i] for i in tup]); tot += prob * sum(V4[i] * f4[i] for i in tup) / n\n", + "check('5. pooled filtered average: E[(1/n) sum V f] = p E_q[f]', abs(tot - p * Eq_f), 1e-14)" + ] + }, + { + "cell_type": "markdown", + "id": "e2501a9f", + "metadata": {}, + "source": [ + "## 8. Monte Carlo agreement with the exact tabular gradient (RFT step = REINFORCE step)" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "5aa7257d", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:48:00.442590Z", + "iopub.status.busy": "2026-10-14T09:48:00.442434Z", + "iopub.status.idle": "2026-10-14T09:48:00.525797Z", + "shell.execute_reply": "2026-10-14T09:48:00.524856Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "7. Monte Carlo RFT gradient (N=4e5 per prompt) vs grad J (tol 3e-4) err=1.941e-04 (tol 3e-04)\n" + ] + } + ], + "source": [ + "N = 400_000\n", + "gen = torch.Generator().manual_seed(3)\n", + "Pt = pi_tab(logits0)\n", + "gmc = torch.zeros(m, K)\n", + "for x in range(m):\n", + " ys = torch.multinomial(Pt[x], N, replacement=True, generator=gen)\n", + " cnt = torch.bincount(ys, minlength=K).double() / N # empirical pi_t\n", + " # grad of (1/N) sum_i V log pi_theta(y_i) wrt logits at theta_t: V * (onehot - pi)\n", + " gmc[x] = ((cnt * V[x])[:, None] * (torch.eye(K) - Pt[x][None, :])).sum(0) / m\n", + "err = (gmc - gJ).abs().max()\n", + "check('7. Monte Carlo RFT gradient (N=4e5 per prompt) vs grad J (tol 3e-4)', err, 3e-4)" + ] + }, + { + "cell_type": "markdown", + "id": "ba4f6acb", + "metadata": {}, + "source": [ + "## Verification table" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "3cb4ed0a", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:48:00.527272Z", + "iopub.status.busy": "2026-10-14T09:48:00.527101Z", + "iopub.status.idle": "2026-10-14T09:48:00.530845Z", + "shell.execute_reply": "2026-10-14T09:48:00.530126Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + "identity max abs err\n", + "1. grad of RFT objective at theta_t vs grad J(theta_t) 3.469e-18\n", + "2. ELBO identity -KL(q||pi) = log p - KL(q||q*) 0.000e+00\n", + "2. tightness: -KL(q*||pi) = log p 2.220e-16\n", + "3. p_{t+1}(x) = 1 on covered prompts (max |p-1|) 0.000e+00\n", + "3. p_{t+1}(x) = 0 on uncovered prompt (coverage failure) 0.000e+00\n", + "4. EM monotone: max decrease of F over 12 iterations (should be <= 0) 0.000e+00\n", + "4. M-step solved to optimality (max |grad| at the last M-step) 1.218e-09\n", + "4. ELBO(theta_t,q_t) = F(theta_t) at the final iterate 0.000e+00\n", + "4. equal-weight M-step objective: gradient at theta_t = grad F(theta_t) (weight 1/p_t per problem on grad p) 6.245e-17\n", + "4. pooled objective (ReST-EM Eq.3 / RFT): gradient at theta_t = grad mean_S p_t (weight 1 on grad p) 1.388e-17\n", + "6. non-monotone prompts exist under unnormalised weights (count of decreasing p, informational) 2.000e+01\n", + "5. E[filtered mean of f] = (1-(1-p)^n) E_q[f], n=1 0.000e+00\n", + "5. E[filtered mean of f] = (1-(1-p)^n) E_q[f], n=2 3.331e-16\n", + "5. E[filtered mean of f] = (1-(1-p)^n) E_q[f], n=3 8.882e-16\n", + "5. E[filtered mean of f] = (1-(1-p)^n) E_q[f], n=5 2.220e-15\n", + "5. pooled filtered average: E[(1/n) sum V f] = p E_q[f] 3.331e-16\n", + "7. Monte Carlo RFT gradient (N=4e5 per prompt) vs grad J (tol 3e-4) 1.941e-04\n", + "CHECKS_JSON:[[\"1. grad of RFT objective at theta_t vs grad J(theta_t)\", 3.469446951953614e-18, 1e-16], [\"2. ELBO identity -KL(q||pi) = log p - KL(q||q*)\", 0.0, 1e-14], [\"2. tightness: -KL(q*||pi) = log p\", 2.220446049250313e-16, 1e-14], [\"3. p_{t+1}(x) = 1 on covered prompts (max |p-1|)\", 0.0, 1e-14], [\"3. p_{t+1}(x) = 0 on uncovered prompt (coverage failure)\", 0.0, 0.0], [\"4. EM monotone: max decrease of F over 12 iterations (should be <= 0)\", 0.0, 1e-10], [\"4. M-step solved to optimality (max |grad| at the last M-step)\", 1.2179568404172003e-09, 1e-08], [\"4. ELBO(theta_t,q_t) = F(theta_t) at the final iterate\", 0.0, 1e-12], [\"4. equal-weight M-step objective: gradient at theta_t = grad F(theta_t) (weight 1/p_t per problem on grad p)\", 6.245004513516506e-17, 1e-14], [\"4. pooled objective (ReST-EM Eq.3 / RFT): gradient at theta_t = grad mean_S p_t (weight 1 on grad p)\", 1.3877787807814457e-17, 1e-14], [\"6. non-monotone prompts exist under unnormalised weights (count of decreasing p, informational)\", 20.0, Infinity], [\"5. E[filtered mean of f] = (1-(1-p)^n) E_q[f], n=1\", 0.0, 1e-14], [\"5. E[filtered mean of f] = (1-(1-p)^n) E_q[f], n=2\", 3.3306690738754696e-16, 1e-14], [\"5. E[filtered mean of f] = (1-(1-p)^n) E_q[f], n=3\", 8.881784197001252e-16, 1e-14], [\"5. E[filtered mean of f] = (1-(1-p)^n) E_q[f], n=5\", 2.220446049250313e-15, 1e-14], [\"5. pooled filtered average: E[(1/n) sum V f] = p E_q[f]\", 3.3306690738754696e-16, 1e-14], [\"7. Monte Carlo RFT gradient (N=4e5 per prompt) vs grad J (tol 3e-4)\", 0.00019412726116547693, 0.0003]]\n" + ] + } + ], + "source": [ + "print()\n", + "print(f'{\"identity\":86s} {\"max abs err\":>12s}')\n", + "for n_, e_, t_ in CHECKS: print(f'{n_:86s} {e_:12.3e}')\n", + "print('CHECKS_JSON:' + json.dumps(CHECKS))" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/code/notebooks/0203-pass-at-k.ipynb b/code/notebooks/0203-pass-at-k.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..9ea00d3f09c2bd17c27c21f6d1f582d0762adc10 --- /dev/null +++ b/code/notebooks/0203-pass-at-k.ipynb @@ -0,0 +1,657 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "d6b3450b", + "metadata": {}, + "source": [ + "# 0203 -- pass@k: the unbiased estimator, its variance, and pass@1 vs pass@k\n", + "\n", + "Identities verified (exact rational arithmetic where possible, otherwise fp64):\n", + "1. `C(n,c) C(n-c,k) / C(n,k) = C(n-k,c)` and `E_c[ C(n-c,k)/C(n,k) ] = (1-p)^k` (exact `Fraction`).\n", + "2. Stable product form `prod_{i=n-c+1}^{n} (1-k/i)` equals the ratio of binomials; the naive float ratio overflows.\n", + "3. Variance by Hoeffding's U-statistic formula equals the exact binomial-sum variance; leading term `k^2 p (1-p)^(2k-1) / n`.\n", + "4. Plug-in `1-(1-p_hat)^k` is biased low: exact bias, and `-p(1-p)/n` for `k=2`.\n", + "5. `grad J_k = E_x[k (1-p)^(k-1) grad p]` against autograd.\n", + "6. A context-free strategy-mixing policy where the `J_1`-optimum and `J_32`-optimum differ (diversity vs sharpening).\n", + "7. The unbiased pass@k policy-gradient estimator: naive subset enumeration = closed-form coefficients; expectation = `grad J_k` exactly; leave-one-out form.\n", + "8. No unbiased estimator exists for `k > n` (Bernstein-basis projection residual)." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "e6de1a23", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T08:57:46.427466Z", + "iopub.status.busy": "2026-10-14T08:57:46.427307Z", + "iopub.status.idle": "2026-10-14T08:57:48.231511Z", + "shell.execute_reply": "2026-10-14T08:57:48.230621Z" + } + }, + "outputs": [], + "source": [ + "import os\n", + "os.environ['CUDA_VISIBLE_DEVICES'] = ''\n", + "import itertools, math, json\n", + "from fractions import Fraction as Fr\n", + "from math import comb\n", + "import numpy as np, torch\n", + "torch.set_default_dtype(torch.float64)\n", + "CHECKS = []\n", + "def check(name, err, tol):\n", + " CHECKS.append((name, float(err), float(tol)))\n", + " assert err <= tol, (name, float(err), tol)\n", + " print(f'{name:80s} err={float(err):.3e} (tol {tol:.0e})')" + ] + }, + { + "cell_type": "markdown", + "id": "d785fb44", + "metadata": {}, + "source": [ + "## 1. Unbiasedness, exactly\n", + "For rational `p` the binomial expectation is a rational number, so the check is an *equality* of `Fraction`s." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "958d0931", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T08:57:48.233490Z", + "iopub.status.busy": "2026-10-14T08:57:48.233261Z", + "iopub.status.idle": "2026-10-14T08:57:48.338023Z", + "shell.execute_reply": "2026-10-14T08:57:48.337325Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "1. E_c[1 - C(n-c,k)/C(n,k)] = 1-(1-p)^k, exact rational, all n<=40 tested err=0.000e+00 (tol 0e+00)\n", + "1. identity C(n,c)C(n-c,k) = C(n,k)C(n-k,c) (violations over n<30) err=0.000e+00 (tol 0e+00)\n" + ] + } + ], + "source": [ + "def passk_unbiased_expectation(n, k, p):\n", + " p = Fr(p)\n", + " s = Fr(0)\n", + " for c in range(n + 1):\n", + " pc = comb(n, c) * p ** c * (1 - p) ** (n - c)\n", + " s += pc * (1 - Fr(comb(n - c, k), comb(n, k))) # comb(n-c,k)=0 when n-c int too large to convert to float\n", + "stable value for the same input : 1.0\n" + ] + } + ], + "source": [ + "def pass_at_k_naive_float(n, c, k):\n", + " return 1.0 - float(comb(n - c, k)) / float(comb(n, k)) # float(comb) overflows for big n,k\n", + "def pass_at_k_stable(n, c, k):\n", + " if n - c < k: return 1.0\n", + " return 1.0 - float(np.prod(1.0 - k / np.arange(n - c + 1, n + 1)))\n", + "def pass_at_k_exact(n, c, k): return 1 - Fr(comb(n - c, k), comb(n, k))\n", + "err = 0.0\n", + "for (n, k) in [(10, 3), (200, 10), (200, 100), (1000, 100), (1000, 999)]:\n", + " for c in sorted(set([0, 1, 2, n // 7, n // 2, n - k, n - k + 1, n - 1, n])):\n", + " if c > n: continue\n", + " err = max(err, abs(pass_at_k_stable(n, c, k) - float(pass_at_k_exact(n, c, k))))\n", + "check('2. stable product form vs exact rational (max over grid incl. n=1000,k=999)', err, 1e-12)\n", + "try:\n", + " v = pass_at_k_naive_float(6000, 100, 3000); print('naive float on n=6000,k=3000:', v)\n", + " naive_failed = not np.isfinite(v)\n", + "except OverflowError as e:\n", + " naive_failed = True; print('naive float on n=6000,k=3000: OverflowError ->', e)\n", + "assert naive_failed\n", + "print('stable value for the same input :', pass_at_k_stable(6000, 100, 3000))\n", + "# the naive form silently loses precision / overflows; the stable form is a product of terms in (0,1)" + ] + }, + { + "cell_type": "markdown", + "id": "d9320e2e", + "metadata": {}, + "source": [ + "## 3. Variance of the unbiased estimator\n", + "EXACT (naive): `Var = E[rho^2] - (1-p)^{2k}` summed over the binomial, `rho = C(n-c,k)/C(n,k)`.\n", + "HOEFFDING (efficient): the estimator is a U-statistic with kernel `h(S)=1{all k in S fail}`, and two random k-subsets overlapping in `j`\n", + "elements both fail with probability `(1-p)^{2k-j}`, so\n", + "`Var = sum_{j=1}^{k} w_j [ (1-p)^{2k-j} - (1-p)^{2k} ]`, `w_j = C(k,j) C(n-k,k-j) / C(n,k)` (hypergeometric overlap pmf)." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "3441dcee", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T08:57:48.349346Z", + "iopub.status.busy": "2026-10-14T08:57:48.349183Z", + "iopub.status.idle": "2026-10-14T08:57:48.377294Z", + "shell.execute_reply": "2026-10-14T08:57:48.376327Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "3. Hoeffding U-statistic variance = exact binomial-sum variance (exact Fractions) err=0.000e+00 (tol 1e-18)\n", + "n=200,k=8,p=0.3: exact Var=4.804470e-04 leading term k^2 p (1-p)^(2k-1)/n=4.557659e-04\n", + "3. exact variance vs leading term (relative error, n=200) err=5.415e-02 (tol 1e-01)\n" + ] + } + ], + "source": [ + "def var_exact(n, k, p):\n", + " p = Fr(p); m1 = Fr(0); m2 = Fr(0)\n", + " for c in range(n + 1):\n", + " pc = comb(n, c) * p ** c * (1 - p) ** (n - c); rho = Fr(comb(n - c, k), comb(n, k))\n", + " m1 += pc * rho; m2 += pc * rho * rho\n", + " return m2 - m1 * m1\n", + "def var_hoeffding(n, k, p):\n", + " p = Fr(p); s = Fr(0)\n", + " for j in range(1, k + 1):\n", + " wj = Fr(comb(k, j) * comb(n - k, k - j), comb(n, k))\n", + " s += wj * ((1 - p) ** (2 * k - j) - (1 - p) ** (2 * k))\n", + " return s\n", + "worst = 0.0\n", + "for (n, k) in [(5, 2), (16, 4), (16, 8), (50, 10), (30, 30)]:\n", + " for p in (Fr(1, 20), Fr(3, 10), Fr(1, 2), Fr(9, 10)):\n", + " worst = max(worst, abs(float(var_exact(n, k, p) - var_hoeffding(n, k, p))))\n", + "check('3. Hoeffding U-statistic variance = exact binomial-sum variance (exact Fractions)', worst, 1e-18)\n", + "n, k, p = 200, 8, 0.3\n", + "lead = k ** 2 * p * (1 - p) ** (2 * k - 1) / n\n", + "print(f'n={n},k={k},p={p}: exact Var={float(var_exact(n,k,Fr(3,10))):.6e} leading term k^2 p (1-p)^(2k-1)/n={lead:.6e}')\n", + "check('3. exact variance vs leading term (relative error, n=200)', abs(float(var_exact(n, k, Fr(3, 10))) / lead - 1), 0.12)" + ] + }, + { + "cell_type": "markdown", + "id": "dfad7109", + "metadata": {}, + "source": [ + "## 4. The plug-in estimator `1-(1-p_hat)^k` is biased low\n", + "By Jensen (`x -> (1-x)^k` is convex) `E[(1-p_hat)^k] >= (1-p)^k`. For `k=2`: `E(1-p_hat)^2 = (1-p)^2 + p(1-p)/n` exactly." + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "6494c052", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T08:57:48.379114Z", + "iopub.status.busy": "2026-10-14T08:57:48.378897Z", + "iopub.status.idle": "2026-10-14T08:57:48.388166Z", + "shell.execute_reply": "2026-10-14T08:57:48.387120Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "4. k=2 exact bias of plug-in = -p(1-p)/n err=0.000e+00 (tol 0e+00)\n", + "bias of plug-in pass@k (n=8,p=.3): {1: 0.0, 2: -0.02625, 4: -0.075460546875, 8: -0.10406236171669006}\n", + "4. k=1 plug-in is unbiased err=0.000e+00 (tol 0e+00)\n" + ] + } + ], + "source": [ + "def plugin_expectation(n, k, p):\n", + " p = Fr(p); return sum(comb(n, c) * p ** c * (1 - p) ** (n - c) * (1 - Fr(c, n)) ** k for c in range(n + 1))\n", + "n, p = 8, Fr(3, 10)\n", + "bias2 = (1 - plugin_expectation(n, 2, p)) - (1 - (1 - p) ** 2) # E[plugin pass@2] - truth\n", + "check('4. k=2 exact bias of plug-in = -p(1-p)/n', float(abs(bias2 + p * (1 - p) / n)), 0.0)\n", + "print('bias of plug-in pass@k (n=8,p=.3):', {k_: float((1 - plugin_expectation(8, k_, p)) - (1 - (1 - p) ** k_)) for k_ in (1, 2, 4, 8)})\n", + "assert all((1 - plugin_expectation(8, k_, p)) - (1 - (1 - p) ** k_) <= 0 for k_ in range(1, 9))\n", + "check('4. k=1 plug-in is unbiased', float(abs((1 - plugin_expectation(8, 1, p)) - p)), 0.0)" + ] + }, + { + "cell_type": "markdown", + "id": "1e4a9d5e", + "metadata": {}, + "source": [ + "## 5. Gradient of the pass@k objective: `k (1-p)^(k-1) grad p`\n", + "Tabular softmax per prompt, `J_k = mean_x [1-(1-p_x)^k]`." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "f7df0b58", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T08:57:48.390043Z", + "iopub.status.busy": "2026-10-14T08:57:48.389832Z", + "iopub.status.idle": "2026-10-14T08:57:48.487130Z", + "shell.execute_reply": "2026-10-14T08:57:48.486143Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "5. grad J_k = E_x[k (1-p)^{k-1} grad p] vs autograd, k in {1,2,8,64} err=6.939e-18 (tol 1e-16)\n", + "weights w_k(p)=k(1-p)^{k-1}:\n", + " k= 1 [1. 1. 1. 1. 1. 1.]\n", + " k= 4 [3.8812 3.4295 2.916 1.372 0.5 0.004 ]\n", + " k= 16 [1.37609e+01 7.41270e+00 3.29430e+00 7.60000e-02 5.00000e-04 0.00000e+00]\n", + " k= 64 [33.978 2.5279 0.0838 0. 0. 0. ]\n" + ] + } + ], + "source": [ + "torch.manual_seed(0)\n", + "Kout, Xn = 9, 7\n", + "Vtab = (torch.rand(Xn, Kout) < 0.25).double(); Vtab[0, :] = 0; Vtab[0, 3] = 1\n", + "z = (torch.randn(Xn, Kout) * 0.8).requires_grad_(True)\n", + "def p_x(z): return (torch.softmax(z, -1) * Vtab).sum(-1)\n", + "worst = 0.0\n", + "for k in (1, 2, 8, 64):\n", + " z.grad = None\n", + " Jk = (1 - (1 - p_x(z)) ** k).mean(); Jk.backward(); g_auto = z.grad.clone()\n", + " z2 = z.detach().clone().requires_grad_(True)\n", + " pz = p_x(z2); pz.sum().backward() # d p_x / dz row-wise (each p_x depends only on its row)\n", + " gp = z2.grad # row x holds grad p_x\n", + " g_hand = (k * (1 - pz.detach()) ** (k - 1))[:, None] * gp / Xn\n", + " worst = max(worst, float((g_auto - g_hand).abs().max()))\n", + "check('5. grad J_k = E_x[k (1-p)^{k-1} grad p] vs autograd, k in {1,2,8,64}', worst, 1e-16)\n", + "pv = np.array([0.01, 0.05, 0.1, 0.3, 0.5, 0.9])\n", + "print('weights w_k(p)=k(1-p)^{k-1}:')\n", + "for k in (1, 4, 16, 64): print(f' k={k:3d}', np.round(k * (1 - pv) ** (k - 1), 4))" + ] + }, + { + "cell_type": "markdown", + "id": "6a670727", + "metadata": {}, + "source": [ + "## 6. Why optimising pass@1 and pass@k pull in different directions\n", + "Within one prompt `pass@k = 1-(1-p)^k` is increasing in `p`, so the two objectives *agree* on that prompt. They disagree **across prompts**\n", + "through the weights `1` vs `k(1-p)^{k-1}`. A context-free \"strategy\" policy makes this concrete: `pi` over 4 strategies; strategy 0 is a\n", + "broad generalist that solves 8 \"easy\" prompts with prob 0.6 and nothing else; strategies 1..3 each solve one of 4 hard prompts... with prob 0.3\n", + "(strategy 1 -> hard prompts 8, 9; strategy 2 -> 10; strategy 3 -> 11) and 0.1 on easy ones." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "0e976d80", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T08:57:48.488964Z", + "iopub.status.busy": "2026-10-14T08:57:48.488802Z", + "iopub.status.idle": "2026-10-14T08:57:53.482658Z", + "shell.execute_reply": "2026-10-14T08:57:53.481750Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "argmax J_1 : pi = [1. 0. 0. 0.] entropy 0.0\n", + "argmax J_32: pi = [0.111 0.3414 0.2738 0.2738] entropy 1.3202\n", + "J_1 at th1=0.4000 J_1 at th32=0.1344\n", + "J_32 at th1=0.6667 J_32 at th32=0.9810\n", + "k=1: |grad| at optimum = 2.12e-07\n", + "k=32: |grad| at optimum = 4.82e-05\n", + "6. stationarity of both optima (max |grad J_k| at the returned theta) err=4.823e-05 (tol 1e-03)\n", + "6. cross-inferiority gaps: min(J_1(th1)-J_1(th32), J_32(th32)-J_32(th1)) > 0 (reported as 0 if true) err=0.000e+00 (tol 0e+00)\n" + ] + } + ], + "source": [ + "Xn = 12\n", + "S = torch.zeros(Xn, 4)\n", + "S[:8, 0] = 0.6; S[:8, 1:] = 0.1\n", + "S[8, 1] = 0.3; S[9, 1] = 0.3; S[10, 2] = 0.3; S[11, 3] = 0.3\n", + "def J_k(theta, k):\n", + " pi = torch.softmax(theta, 0); p = S @ pi\n", + " return (1 - (1 - p) ** k).mean()\n", + "def opt(k, steps=6000, lr=0.3, seed=0):\n", + " th = torch.zeros(4, requires_grad=True)\n", + " o = torch.optim.Adam([th], lr=lr)\n", + " for _ in range(steps):\n", + " o.zero_grad(); (-J_k(th, k)).backward(); o.step()\n", + " return th.detach()\n", + "th1, th32 = opt(1), opt(32)\n", + "H = lambda th: float(-(torch.softmax(th, 0) * torch.log_softmax(th, 0)).sum())\n", + "print('argmax J_1 : pi =', np.round(torch.softmax(th1, 0).numpy(), 4), ' entropy', round(H(th1), 4))\n", + "print('argmax J_32: pi =', np.round(torch.softmax(th32, 0).numpy(), 4), ' entropy', round(H(th32), 4))\n", + "print(f'J_1 at th1={float(J_k(th1,1)):.4f} J_1 at th32={float(J_k(th32,1)):.4f}')\n", + "print(f'J_32 at th1={float(J_k(th1,32)):.4f} J_32 at th32={float(J_k(th32,32)):.4f}')\n", + "assert J_k(th1, 1) > J_k(th32, 1) + 1e-3 and J_k(th32, 32) > J_k(th1, 32) + 1e-3 and H(th32) > H(th1)\n", + "# Check that the optimiser reached a stationary point (max |grad| of J_k at the returned theta)\n", + "gmax = []\n", + "for k, th in ((1, th1), (32, th32)):\n", + " t = th.clone().requires_grad_(True); J_k(t, k).backward()\n", + " gmax.append(float(t.grad.abs().max())); print(f'k={k}: |grad| at optimum = {gmax[-1]:.2e}')\n", + "check('6. stationarity of both optima (max |grad J_k| at the returned theta)', max(gmax), 1e-3)\n", + "check('6. cross-inferiority gaps: min(J_1(th1)-J_1(th32), J_32(th32)-J_32(th1)) > 0 (reported as 0 if true)', 0.0 if min(float(J_k(th1,1)-J_k(th32,1)), float(J_k(th32,32)-J_k(th1,32))) > 0 else 1.0, 0.0)" + ] + }, + { + "cell_type": "markdown", + "id": "75d59643", + "metadata": {}, + "source": [ + "## 7. The unbiased pass@k *gradient* estimator (PKPO-style), exact enumeration\n", + "Tabular softmax over `K=4` outcomes with binary reward vector `r`, `n=5` samples, `k=3`. For a tuple `(y_1..y_n)`:\n", + "* NAIVE: `g = (1/C(n,k)) sum_S max_{i in S} r_i * sum_{i in S} grad log pi(y_i)` (all `C(n,k)` subsets).\n", + "* CLOSED FORM: `g = sum_i w_i grad log pi(y_i)` with `w_i = (k/n)[ rho_{-i} + r_i (1 - rho_{-i}) ]`, where `rho_{-i}=1-C(n-1-c_{-i},k-1)/C(n-1,k-1)`\n", + " is the unbiased pass@(k-1) estimate from the other `n-1` samples, `c_{-i}` the number correct among them.\n", + "* LEAVE-ONE-OUT form: subtract the baseline `(k/n) rho_{-i}` (independent of `y_i`): coefficient `(k/n) r_i (1-rho_{-i})`.\n", + "`grad J_k = k (1-p)^{k-1} grad p` is the target." + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "556c657b", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T08:57:53.484341Z", + "iopub.status.busy": "2026-10-14T08:57:53.484099Z", + "iopub.status.idle": "2026-10-14T08:57:54.020052Z", + "shell.execute_reply": "2026-10-14T08:57:54.019291Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "7. closed-form coefficients = naive subset average, every tuple (max abs diff) err=4.441e-16 (tol 1e-14)\n", + "7. E[naive estimator] = grad of 1-(1-p)^k (exact over 4^5 tuples) err=4.302e-16 (tol 1e-15)\n", + "7. E[closed-form] = grad J_k err=4.302e-16 (tol 1e-15)\n", + "7. E[leave-one-out form] = grad J_k (baseline (k/n) rho_{-i} is independent of y_i) err=6.661e-16 (tol 1e-15)\n", + "total variance: closed-form (no baseline)=0.44125 leave-one-out=0.13981 ratio=0.317\n", + "7. k=1 reduces to REINFORCE sum_i (r_i/n) grad log pi(y_i) err=1.110e-16 (tol 1e-15)\n" + ] + } + ], + "source": [ + "Kc, n, k = 4, 5, 3\n", + "r = torch.tensor([0., 1., 0., 1.])\n", + "theta = torch.tensor([0.3, -0.5, 0.9, -1.1], requires_grad=True)\n", + "pi = torch.softmax(theta, 0)\n", + "pk_obj = 1 - (1 - (pi * r).sum()) ** k\n", + "pk_obj.backward(); g_target = theta.grad.clone()\n", + "pi_d = pi.detach()\n", + "score = torch.eye(Kc) - pi_d[None, :] # row y: grad_theta log pi(y)\n", + "def rho_minus(c_other, n_, k_): # unbiased pass@(k-1) from n-1 samples with c_other correct (k_=k)\n", + " kk = k_ - 1\n", + " if kk == 0: return 0.0\n", + " return 1.0 - comb(n_ - 1 - c_other, kk) / comb(n_ - 1, kk)\n", + "E_naive = torch.zeros(Kc); E_closed = torch.zeros(Kc); E_loo = torch.zeros(Kc)\n", + "maxdiff = 0.0; var_closed = torch.zeros(Kc); var_loo = torch.zeros(Kc)\n", + "tuples = list(itertools.product(range(Kc), repeat=n))\n", + "second_closed = torch.zeros(Kc); second_loo = torch.zeros(Kc)\n", + "for tup in tuples:\n", + " prob = float(np.prod([float(pi_d[y]) for y in tup]))\n", + " rr = [float(r[y]) for y in tup]\n", + " # naive: all subsets\n", + " g1 = torch.zeros(Kc)\n", + " for Ssub in itertools.combinations(range(n), k):\n", + " mx = max(rr[i] for i in Ssub)\n", + " g1 += mx * sum(score[tup[i]] for i in Ssub)\n", + " g1 /= comb(n, k)\n", + " # closed form\n", + " g2 = torch.zeros(Kc); g3 = torch.zeros(Kc); c_tot = sum(rr)\n", + " for i in range(n):\n", + " c_other = c_tot - rr[i]\n", + " rho = rho_minus(int(c_other), n, k)\n", + " w = (k / n) * (rho + rr[i] * (1 - rho))\n", + " g2 += w * score[tup[i]]\n", + " g3 += (k / n) * rr[i] * (1 - rho) * score[tup[i]]\n", + " maxdiff = max(maxdiff, float((g1 - g2).abs().max()))\n", + " E_naive += prob * g1; E_closed += prob * g2; E_loo += prob * g3\n", + " second_closed += prob * g2 ** 2; second_loo += prob * g3 ** 2\n", + "check('7. closed-form coefficients = naive subset average, every tuple (max abs diff)', maxdiff, 1e-14)\n", + "check('7. E[naive estimator] = grad of 1-(1-p)^k (exact over 4^5 tuples)', (E_naive - g_target).abs().max(), 1e-15)\n", + "check('7. E[closed-form] = grad J_k', (E_closed - g_target).abs().max(), 1e-15)\n", + "check('7. E[leave-one-out form] = grad J_k (baseline (k/n) rho_{-i} is independent of y_i)', (E_loo - g_target).abs().max(), 1e-15)\n", + "v_closed = float((second_closed - E_closed ** 2).sum()); v_loo = float((second_loo - E_loo ** 2).sum())\n", + "print(f'total variance: closed-form (no baseline)={v_closed:.5f} leave-one-out={v_loo:.5f} ratio={v_loo/v_closed:.3f}')\n", + "assert v_loo < v_closed\n", + "# k=1 reduces to REINFORCE without baseline: w_i = r_i/n\n", + "kk = 1; maxd = 0.0\n", + "for tup in tuples[:200]:\n", + " rr = [float(r[y]) for y in tup]\n", + " g1 = sum(max(rr[i] for i in (j,)) * score[tup[j]] for j in range(n)) / n\n", + " g2 = sum((kk / n) * (0.0 + rr[i] * (1 - 0.0)) * score[tup[i]] for i in range(n))\n", + " maxd = max(maxd, float((g1 - g2).abs().max()))\n", + "check('7. k=1 reduces to REINFORCE sum_i (r_i/n) grad log pi(y_i)', maxd, 1e-15)" + ] + }, + { + "cell_type": "markdown", + "id": "65ffaf66", + "metadata": {}, + "source": [ + "## 8. No unbiased estimator for `k > n` (Kolmogorov): Bernstein-basis residual\n", + "Unbiased estimators of `f(p)` from `n` Bernoulli draws are exactly `f(p) = sum_c a_c C(n,c) p^c (1-p)^{n-c}`, i.e. polynomials of degree `<= n`.\n", + "Project `f(p) = 1-(1-p)^k` on the degree-`n` Bernstein basis: residual 0 iff `k <= n`." + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "476de431", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T08:57:54.021866Z", + "iopub.status.busy": "2026-10-14T08:57:54.021703Z", + "iopub.status.idle": "2026-10-14T08:57:54.031118Z", + "shell.execute_reply": "2026-10-14T08:57:54.030302Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "{(4, 3): 1.2212453270876722e-15, (4, 4): 9.558326352632207e-16, (4, 5): 0.003948447395839216, (4, 8): 0.043274555276372625, (10, 10): 2.617697725249002e-15, (10, 11): 1.3790524725575724e-06}\n", + "8. residual for k<=n (max over tested) ~ 0 err=2.618e-15 (tol 1e-09)\n" + ] + } + ], + "source": [ + "pg = np.linspace(0, 1, 2001)\n", + "def bern_resid(n_, k_):\n", + " B = np.stack([comb(n_, c) * pg ** c * (1 - pg) ** (n_ - c) for c in range(n_ + 1)], 1)\n", + " f = 1 - (1 - pg) ** k_\n", + " a, *_ = np.linalg.lstsq(B, f, rcond=None)\n", + " return float(np.abs(B @ a - f).max())\n", + "res = {(n_, k_): bern_resid(n_, k_) for n_, k_ in [(4, 3), (4, 4), (4, 5), (4, 8), (10, 10), (10, 11)]}\n", + "print(res)\n", + "check('8. residual for k<=n (max over tested) ~ 0', max(res[(4, 3)], res[(4, 4)], res[(10, 10)]), 1e-9)\n", + "assert res[(4, 5)] > 1e-3 and res[(10, 11)] > 1e-6\n", + "CHECKS.append(('8. residual for k=n+1 is strictly positive (min over tested; informational)', float(min(res[(4,5)], res[(10,11)])), float('inf')))" + ] + }, + { + "cell_type": "markdown", + "id": "19e22270", + "metadata": {}, + "source": [ + "## 9. Benchmark-level sampling variance: within-problem part\n", + "`N=244` problems with `p_x ~ Beta(0.5, 1)`, `n=16` samples per problem, `k=8`. For *fixed* problems the variance of the benchmark mean is\n", + "`(1/N^2) sum_x Var_x`, with `Var_x` from section 3. Monte Carlo agrees." + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "4691ab9e", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T08:57:54.032552Z", + "iopub.status.busy": "2026-10-14T08:57:54.032377Z", + "iopub.status.idle": "2026-10-14T08:58:10.317977Z", + "shell.execute_reply": "2026-10-14T08:58:10.317179Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "truth=0.67114 MC mean=0.67109 (+-0.00006) MC var=1.3361e-04 theory var=1.3230e-04\n", + "9. estimator is unbiased at benchmark level (|MC mean - truth| in s.e.) err=9.278e-01 (tol 4e+00)\n", + "9. within-problem variance: MC vs theory (relative) err=9.909e-03 (tol 3e-02)\n" + ] + } + ], + "source": [ + "rg = np.random.RandomState(5)\n", + "N, n, k = 244, 16, 8\n", + "px = rg.beta(0.5, 1.0, size=N)\n", + "truth = float((1 - (1 - px) ** k).mean())\n", + "var_theory = sum(float(var_exact(n, k, Fr(float(pv)).limit_denominator(10 ** 9))) for pv in px) / N ** 2\n", + "reps = 40000\n", + "c = rg.binomial(n, px[None, :], size=(reps, N))\n", + "def est_from_c(c):\n", + " out = np.ones_like(c, dtype=float)\n", + " for cc in range(0, n + 1):\n", + " out[c == cc] = pass_at_k_stable(n, cc, k)\n", + " return out.mean(1)\n", + "bench = est_from_c(c)\n", + "se_mean = bench.std() / math.sqrt(reps)\n", + "print(f'truth={truth:.5f} MC mean={bench.mean():.5f} (+-{se_mean:.5f}) MC var={bench.var():.4e} theory var={var_theory:.4e}')\n", + "check('9. estimator is unbiased at benchmark level (|MC mean - truth| in s.e.)', abs(bench.mean() - truth) / se_mean, 4.0)\n", + "check('9. within-problem variance: MC vs theory (relative)', abs(bench.var() / var_theory - 1), 0.03)" + ] + }, + { + "cell_type": "markdown", + "id": "83602bfb", + "metadata": {}, + "source": [ + "## Verification table" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "id": "86c128e2", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T08:58:10.319495Z", + "iopub.status.busy": "2026-10-14T08:58:10.319334Z", + "iopub.status.idle": "2026-10-14T08:58:10.322873Z", + "shell.execute_reply": "2026-10-14T08:58:10.322203Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + "identity max abs err\n", + "1. E_c[1 - C(n-c,k)/C(n,k)] = 1-(1-p)^k, exact rational, all n<=40 tested 0.000e+00\n", + "1. identity C(n,c)C(n-c,k) = C(n,k)C(n-k,c) (violations over n<30) 0.000e+00\n", + "2. stable product form vs exact rational (max over grid incl. n=1000,k=999) 8.327e-17\n", + "3. Hoeffding U-statistic variance = exact binomial-sum variance (exact Fractions) 0.000e+00\n", + "3. exact variance vs leading term (relative error, n=200) 5.415e-02\n", + "4. k=2 exact bias of plug-in = -p(1-p)/n 0.000e+00\n", + "4. k=1 plug-in is unbiased 0.000e+00\n", + "5. grad J_k = E_x[k (1-p)^{k-1} grad p] vs autograd, k in {1,2,8,64} 6.939e-18\n", + "6. stationarity of both optima (max |grad J_k| at the returned theta) 4.823e-05\n", + "6. cross-inferiority gaps: min(J_1(th1)-J_1(th32), J_32(th32)-J_32(th1)) > 0 (reported as 0 if true) 0.000e+00\n", + "7. closed-form coefficients = naive subset average, every tuple (max abs diff) 4.441e-16\n", + "7. E[naive estimator] = grad of 1-(1-p)^k (exact over 4^5 tuples) 4.302e-16\n", + "7. E[closed-form] = grad J_k 4.302e-16\n", + "7. E[leave-one-out form] = grad J_k (baseline (k/n) rho_{-i} is independent of y_i) 6.661e-16\n", + "7. k=1 reduces to REINFORCE sum_i (r_i/n) grad log pi(y_i) 1.110e-16\n", + "8. residual for k<=n (max over tested) ~ 0 2.618e-15\n", + "8. residual for k=n+1 is strictly positive (min over tested; informational) 1.379e-06\n", + "9. estimator is unbiased at benchmark level (|MC mean - truth| in s.e.) 9.278e-01\n", + "9. within-problem variance: MC vs theory (relative) 9.909e-03\n", + "CHECKS_JSON:[[\"1. E_c[1 - C(n-c,k)/C(n,k)] = 1-(1-p)^k, exact rational, all n<=40 tested\", 0.0, 0.0], [\"1. identity C(n,c)C(n-c,k) = C(n,k)C(n-k,c) (violations over n<30)\", 0.0, 0.0], [\"2. stable product form vs exact rational (max over grid incl. n=1000,k=999)\", 8.326672684688674e-17, 1e-12], [\"3. Hoeffding U-statistic variance = exact binomial-sum variance (exact Fractions)\", 0.0, 1e-18], [\"3. exact variance vs leading term (relative error, n=200)\", 0.0541530729940769, 0.12], [\"4. k=2 exact bias of plug-in = -p(1-p)/n\", 0.0, 0.0], [\"4. k=1 plug-in is unbiased\", 0.0, 0.0], [\"5. grad J_k = E_x[k (1-p)^{k-1} grad p] vs autograd, k in {1,2,8,64}\", 6.938893903907228e-18, 1e-16], [\"6. stationarity of both optima (max |grad J_k| at the returned theta)\", 4.8228746018059434e-05, 0.001], [\"6. cross-inferiority gaps: min(J_1(th1)-J_1(th32), J_32(th32)-J_32(th1)) > 0 (reported as 0 if true)\", 0.0, 0.0], [\"7. closed-form coefficients = naive subset average, every tuple (max abs diff)\", 4.440892098500626e-16, 1e-14], [\"7. E[naive estimator] = grad of 1-(1-p)^k (exact over 4^5 tuples)\", 4.3021142204224816e-16, 1e-15], [\"7. E[closed-form] = grad J_k\", 4.3021142204224816e-16, 1e-15], [\"7. E[leave-one-out form] = grad J_k (baseline (k/n) rho_{-i} is independent of y_i)\", 6.661338147750939e-16, 1e-15], [\"7. k=1 reduces to REINFORCE sum_i (r_i/n) grad log pi(y_i)\", 1.1102230246251565e-16, 1e-15], [\"8. residual for k<=n (max over tested) ~ 0\", 2.617697725249002e-15, 1e-09], [\"8. residual for k=n+1 is strictly positive (min over tested; informational)\", 1.3790524725575724e-06, Infinity], [\"9. estimator is unbiased at benchmark level (|MC mean - truth| in s.e.)\", 0.9277896993916872, 4.0], [\"9. within-problem variance: MC vs theory (relative)\", 0.009909418241782264, 0.03]]\n" + ] + } + ], + "source": [ + "print()\n", + "print(f'{\"identity\":88s} {\"max abs err\":>12s}')\n", + "for n_, e_, t_ in CHECKS: print(f'{n_:88s} {e_:12.3e}')\n", + "print('CHECKS_JSON:' + json.dumps(CHECKS))" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/code/notebooks/0204-self-correction-multiturn-mdp.ipynb b/code/notebooks/0204-self-correction-multiturn-mdp.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..e27238097fc494a3905fadf1e00515243a97ee38 --- /dev/null +++ b/code/notebooks/0204-self-correction-multiturn-mdp.ipynb @@ -0,0 +1,505 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "bc47f601", + "metadata": {}, + "source": [ + "# 0204 -- Self-correction as a multi-turn MDP: masks, per-turn credit, baselines\n", + "\n", + "A two-turn tabular environment small enough to *enumerate every episode and every group of sampled episodes exactly*.\n", + "Turn 1: `y1 ~ pi1` (K1=3). The verifier returns `r1(y1)` and a feedback token `f = F(y1)` (deterministic).\n", + "Turn 2: `y2 ~ pi2(.|y1)` (K2=3), context = `(x, y1, f)`. Verifier returns `r2(y1,y2)`. Episode return `R = r1 + r2`.\n", + "Parameters: `theta1 in R^3`, `theta2 in R^{3x3}` (one softmax row per turn-2 context).\n", + "\n", + "Identities verified:\n", + "1. Trajectory factorisation: the environment factor drops out; masked REINFORCE is unbiased, `E = grad J`.\n", + "2. Past rewards are independent of later actions: `E[r1 * grad log pi2(y2|c2)] = 0` (reward-to-go is unbiased, with lower variance).\n", + "3. Including feedback tokens in the loss adds a non-zero bias `E[R grad log pi(f|.)]` when the feedback shares the output softmax.\n", + "4. GRPO's mean-including-self baseline is exactly `(G-1)/G` times the leave-one-out one, so its gradient is `(G-1)/G * grad J`.\n", + "5. Leave-one-out (episode-level) is unbiased; a baseline that depends on the sample's own outcome is biased.\n", + "6. Two-level (tree) leave-one-out with per-turn baselines is unbiased; variance compared at equal episode budget.\n", + "7. Training only on a fixed (stale) distribution of turn-2 contexts is biased by `sum (mu - pi1) grad Q`, and drops the turn-1 term." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "343280c0", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:01:50.331677Z", + "iopub.status.busy": "2026-10-14T09:01:50.331456Z", + "iopub.status.idle": "2026-10-14T09:01:52.235734Z", + "shell.execute_reply": "2026-10-14T09:01:52.234848Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "J = 1.0181152430262448\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/home/vignesh/.opencode/tmp/ipykernel_471603/844665086.py:31: UserWarning: Converting a tensor with requires_grad=True to a scalar may lead to unexpected behavior.\n", + "Consider using tensor.detach() first. (Triggered internally at /pytorch/torch/csrc/autograd/generated/python_variable_methods.cpp:838.)\n", + " print('J =', float(J))\n" + ] + } + ], + "source": [ + "import os\n", + "os.environ['CUDA_VISIBLE_DEVICES'] = ''\n", + "import itertools, math, json\n", + "import numpy as np, torch\n", + "torch.set_default_dtype(torch.float64)\n", + "CHECKS = []\n", + "def check(name, err, tol):\n", + " CHECKS.append((name, float(err), float(tol)))\n", + " assert err <= tol, (name, float(err), tol)\n", + " print(f'{name:86s} err={float(err):.3e} (tol {tol:.0e})')\n", + "\n", + "K1, K2 = 3, 3\n", + "g = np.random.RandomState(4)\n", + "r1 = np.array([0., 0., 1.]) # turn-1 success table\n", + "r2 = np.array([[1, 0, 0], [0, 1, 0], [0, 1, 1]], float) # turn-2 success depends on (y1, y2)\n", + "Fmap = np.array([0, 1, 2]) # feedback token id as a function of y1 (shares the K2-sized vocabulary)\n", + "theta1_0 = torch.tensor([0.2, -0.4, 0.5])\n", + "theta2_0 = torch.tensor(g.randn(K1, K2) * 0.7)\n", + "nparam = K1 + K1 * K2\n", + "\n", + "def pis(t1, t2): return torch.softmax(t1, 0), torch.softmax(t2, -1)\n", + "def leaves(t1, t2):\n", + " p1, p2 = pis(t1, t2)\n", + " return {(a, b): float(p1[a] * p2[a, b]) for a in range(K1) for b in range(K2)}\n", + "def J_fn(t1, t2):\n", + " p1, p2 = pis(t1, t2)\n", + " return sum(p1[a] * p2[a, b] * (r1[a] + r2[a, b]) for a in range(K1) for b in range(K2))\n", + "t1 = theta1_0.clone().requires_grad_(True); t2 = theta2_0.clone().requires_grad_(True)\n", + "J = J_fn(t1, t2); J.backward()\n", + "gradJ = torch.cat([t1.grad, t2.grad.flatten()]).numpy()\n", + "print('J =', float(J))\n", + "p1n, p2n = [x.detach().numpy() for x in pis(theta1_0, theta2_0)]\n", + "P = {(a, b): p1n[a] * p2n[a, b] for a in range(K1) for b in range(K2)}\n", + "def s1(a): # score of turn-1 token: e_a - pi1 in theta1 block\n", + " v = np.zeros(nparam); e = np.zeros(K1); e[a] = 1; v[:K1] = e - p1n; return v\n", + "def s2(a, b): # score of turn-2 token: e_b - pi2[a] in the a-th row of theta2\n", + " v = np.zeros(nparam); e = np.zeros(K2); e[b] = 1; v[K1 + a * K2: K1 + (a + 1) * K2] = e - p2n[a]; return v\n", + "def sF(a): # score of the *feedback token* f=F(a) if it were scored by the same softmax head as turn 2, context y1=a\n", + " v = np.zeros(nparam); e = np.zeros(K2); e[Fmap[a]] = 1; v[K1 + a * K2: K1 + (a + 1) * K2] = e - p2n[a]; return v\n", + "Rret = {(a, b): r1[a] + r2[a, b] for a in range(K1) for b in range(K2)}\n", + "L = list(P.keys())" + ] + }, + { + "cell_type": "markdown", + "id": "d24cef3e", + "metadata": {}, + "source": [ + "## 1. Masked REINFORCE is unbiased\n", + "Trajectory probability: `P(tau) = pi1(y1) * E(f|y1) * pi2(y2|y1,f)`. The environment factor `E` does not depend on theta, so\n", + "`grad log P(tau) = grad log pi1(y1) + grad log pi2(y2|c2)` -- only policy-generated tokens appear." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "903afa0a", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:01:52.237467Z", + "iopub.status.busy": "2026-10-14T09:01:52.237201Z", + "iopub.status.idle": "2026-10-14T09:01:52.241076Z", + "shell.execute_reply": "2026-10-14T09:01:52.240235Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "1. E[R (score_turn1 + score_turn2)] = grad J (exact enumeration of 9 episodes) err=8.327e-17 (tol 1e-15)\n" + ] + } + ], + "source": [ + "E_mask = sum(P[l] * Rret[l] * (s1(l[0]) + s2(*l)) for l in L)\n", + "check('1. E[R (score_turn1 + score_turn2)] = grad J (exact enumeration of 9 episodes)', np.abs(E_mask - gradJ).max(), 1e-15)" + ] + }, + { + "cell_type": "markdown", + "id": "851be585", + "metadata": {}, + "source": [ + "## 2. Causality: a past reward carries no information about a later action\n", + "`E[ r1(y1) * grad log pi2(y2|y1) ] = 0` because, given `y1`, `E_{y2}[grad log pi2] = 0`. So `r1` may be dropped from the turn-2 coefficient." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "8c6af3a3", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:01:52.242434Z", + "iopub.status.busy": "2026-10-14T09:01:52.242279Z", + "iopub.status.idle": "2026-10-14T09:01:52.248993Z", + "shell.execute_reply": "2026-10-14T09:01:52.248126Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "2. E[r1 * score_turn2] = 0 err=1.388e-17 (tol 1e-16)\n", + "2. reward-to-go estimator: E = grad J err=6.939e-17 (tol 1e-15)\n", + "single-episode total variance: full return=1.6566 reward-to-go=0.9474\n" + ] + } + ], + "source": [ + "E_past = sum(P[l] * r1[l[0]] * s2(*l) for l in L)\n", + "check('2. E[r1 * score_turn2] = 0', np.abs(E_past).max(), 1e-16)\n", + "E_rtg = sum(P[l] * (Rret[l] * s1(l[0]) + r2[l] * s2(*l)) for l in L)\n", + "check('2. reward-to-go estimator: E = grad J', np.abs(E_rtg - gradJ).max(), 1e-15)\n", + "var_full = sum(P[l] * ((Rret[l] * (s1(l[0]) + s2(*l))) ** 2).sum() for l in L) - (gradJ ** 2).sum()\n", + "var_rtg = sum(P[l] * ((Rret[l] * s1(l[0]) + r2[l] * s2(*l)) ** 2).sum() for l in L) - (gradJ ** 2).sum()\n", + "print(f'single-episode total variance: full return={var_full:.4f} reward-to-go={var_rtg:.4f}')\n", + "assert var_rtg < var_full" + ] + }, + { + "cell_type": "markdown", + "id": "ba273ae5", + "metadata": {}, + "source": [ + "## 3. Putting the feedback token in the loss biases the gradient\n", + "If the loss also includes `log pi(f|context)` for the environment-written feedback token (same output softmax), the expected extra term is\n", + "`E[R grad log pi2(f|y1)]`, generally non-zero." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "5f223ef3", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:01:52.250555Z", + "iopub.status.busy": "2026-10-14T09:01:52.250399Z", + "iopub.status.idle": "2026-10-14T09:01:52.255258Z", + "shell.execute_reply": "2026-10-14T09:01:52.254500Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "||bias from unmasked feedback tokens|| = 0.679933398028168 vs ||grad J|| = 0.47960612146460163\n", + "with constant R the same term is 0.5000121536758657 (non-zero: f is not distributed as pi2, so even the baseline identity fails)\n" + ] + } + ], + "source": [ + "E_fb = sum(P[l] * Rret[l] * sF(l[0]) for l in L)\n", + "print('||bias from unmasked feedback tokens|| =', float(np.linalg.norm(E_fb)), ' vs ||grad J|| =', float(np.linalg.norm(gradJ)))\n", + "assert np.linalg.norm(E_fb) > 1e-3\n", + "CHECKS.append(('3. norm of the bias from including feedback tokens in the loss (informational, must be > 0)', float(np.linalg.norm(E_fb)), float('inf')))\n", + "# without a reward correlation (R constant) the term vanishes: E[score of a fixed-distribution token] = 0 only for tokens sampled from pi.\n", + "E_fb_const = sum(P[l] * 1.0 * sF(l[0]) for l in L)\n", + "print('with constant R the same term is', float(np.linalg.norm(E_fb_const)), '(non-zero: f is not distributed as pi2, so even the baseline identity fails)')" + ] + }, + { + "cell_type": "markdown", + "id": "0b2c937c", + "metadata": {}, + "source": [ + "## 4--6. Group estimators, exactly\n", + "`G` i.i.d. episodes. Enumerate all `9^G` groups with their probabilities. Estimators:\n", + "* `E0` no baseline, full return; `E0rtg` reward-to-go\n", + "* `E1` GRPO-style baseline: group mean *including* the sample\n", + "* `E2` leave-one-out (LOO) baseline: mean of the other `G-1`\n", + "* `E3` tree: `G1` first-turn branches x `G2` second-turn siblings; per-turn LOO (turn-2 baseline = sibling mean, turn-1 baseline = other branches' mean value)" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "feed6ed4", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:01:52.256933Z", + "iopub.status.busy": "2026-10-14T09:01:52.256766Z", + "iopub.status.idle": "2026-10-14T09:01:55.102052Z", + "shell.execute_reply": "2026-10-14T09:01:55.101235Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "E0 |E-gradJ|max=4.996e-16 trace Var=0.41416 (sum prob=1.000000000000)\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "E0rtg |E-gradJ|max=4.996e-16 trace Var=0.23686 (sum prob=1.000000000000)\n", + "E1 |E-gradJ|max=9.267e-02 trace Var=0.14155 (sum prob=1.000000000000)\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "E2 |E-gradJ|max=8.327e-16 trace Var=0.25165 (sum prob=1.000000000000)\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "BAD |E-gradJ|max=7.857e-02 trace Var=0.19228 (sum prob=1.000000000000)\n", + "4. E0 (full-return, no baseline, G=4 episodes): unbiased err=4.996e-16 (tol 1e-15)\n", + "4. E0rtg (reward-to-go, G=4): unbiased err=4.996e-16 (tol 1e-15)\n", + "4. identity R_i - mean_all = (G-1)/G (R_i - mean_{-i}) on every group (max abs diff) err=5.551e-17 (tol 1e-14)\n", + "4. GRPO mean baseline: E[g] = (G-1)/G grad J exactly err=1.193e-15 (tol 1e-14)\n", + "5. leave-one-out baseline (episode-level, G=4): E[g] = grad J err=8.327e-16 (tol 1e-14)\n", + "own-outcome baseline: max |bias| = 0.07857014805316431\n", + "4. variance(E1 scaled by G/(G-1)) = variance(E2) err=7.438e-15 (tol 1e-12)\n", + "tree LOO (2 branches x 2 siblings = 4 episodes): |E-gradJ|max=2.220e-16 trace Var=0.48614\n", + "6. tree two-level leave-one-out with per-turn baselines: E[g] = grad J err=2.220e-16 (tol 1e-14)\n", + "equal budget (4 episodes): trace Var E0=0.4142 E0rtg=0.2369 flat LOO E2=0.2516 tree LOO=0.4861\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "offset c=0.0: trace Var no-baseline=0.414 flat LOO=0.252 tree LOO=0.486\n", + "offset c=3.0: trace Var no-baseline=4.908 flat LOO=0.252 tree LOO=0.486\n" + ] + } + ], + "source": [ + "Gs = 4\n", + "leaf_ids = list(range(len(L)))\n", + "Pl = np.array([P[l] for l in L]); Rl = np.array([Rret[l] for l in L])\n", + "S1 = np.stack([s1(l[0]) for l in L]); S2 = np.stack([s2(*l) for l in L]); S12 = S1 + S2\n", + "R2l = np.array([r2[l] for l in L])\n", + "def run_flat(Gs, kind):\n", + " m1 = np.zeros(nparam); m2 = 0.0; tot_prob = 0.0; maxdiff_1_2 = 0.0\n", + " for grp in itertools.product(leaf_ids, repeat=Gs):\n", + " pr = np.prod(Pl[list(grp)]); gi = list(grp)\n", + " R = Rl[gi]\n", + " if kind == 'E0':\n", + " est = (R[:, None] * S12[gi]).mean(0)\n", + " elif kind == 'E0rtg':\n", + " est = (R[:, None] * S1[gi] + R2l[gi][:, None] * S2[gi]).mean(0)\n", + " elif kind in ('E1', 'E2'):\n", + " mean_all = R.mean(); loo = (R.sum() - R) / (Gs - 1)\n", + " b = mean_all if kind == 'E1' else loo\n", + " est = ((R - b)[:, None] * S12[gi]).mean(0)\n", + " maxdiff_1_2 = max(maxdiff_1_2, float(np.abs((R - mean_all) - (Gs - 1) / Gs * (R - loo)).max()))\n", + " elif kind == 'BAD': # baseline depends on the sample's own turn-2 outcome: b_i = r2_i (for the turn-2 token), mean-LOO for turn 1\n", + " loo = (R.sum() - R) / (Gs - 1)\n", + " est = ((R - loo)[:, None] * S1[gi] + (R - R2l[gi])[:, None] * S2[gi]).mean(0)\n", + " m1 += pr * est; m2 += pr * (est ** 2).sum(); tot_prob += pr\n", + " return m1, m2 - (m1 ** 2).sum(), tot_prob, maxdiff_1_2\n", + "res = {}\n", + "for kind in ('E0', 'E0rtg', 'E1', 'E2', 'BAD'):\n", + " res[kind] = run_flat(Gs, kind)\n", + " print(f'{kind:6s} |E-gradJ|max={np.abs(res[kind][0]-gradJ).max():.3e} trace Var={res[kind][1]:.5f} (sum prob={res[kind][2]:.12f})')\n", + "check('4. E0 (full-return, no baseline, G=4 episodes): unbiased', np.abs(res['E0'][0] - gradJ).max(), 1e-15)\n", + "check('4. E0rtg (reward-to-go, G=4): unbiased', np.abs(res['E0rtg'][0] - gradJ).max(), 1e-15)\n", + "check('4. identity R_i - mean_all = (G-1)/G (R_i - mean_{-i}) on every group (max abs diff)', res['E1'][3], 1e-14)\n", + "check('4. GRPO mean baseline: E[g] = (G-1)/G grad J exactly', np.abs(res['E1'][0] - (Gs - 1) / Gs * gradJ).max(), 1e-14)\n", + "check('5. leave-one-out baseline (episode-level, G=4): E[g] = grad J', np.abs(res['E2'][0] - gradJ).max(), 1e-14)\n", + "bad_bias = np.abs(res['BAD'][0] - gradJ).max()\n", + "print('own-outcome baseline: max |bias| =', bad_bias)\n", + "assert bad_bias > 1e-3\n", + "CHECKS.append(('5. baseline depending on the sample own r2: |bias| (informational, must be > 0)', float(bad_bias), float('inf')))\n", + "# rescale E1 to compare variance\n", + "var_E1_scaled = res['E1'][1] * (Gs / (Gs - 1)) ** 2\n", + "check('4. variance(E1 scaled by G/(G-1)) = variance(E2)', abs(var_E1_scaled - res['E2'][1]), 1e-12)\n", + "\n", + "# --- tree estimator\n", + "G1, G2 = 2, 2\n", + "def run_tree():\n", + " m1 = np.zeros(nparam); m2 = 0.0\n", + " a_range = range(K1); b_range = range(K2)\n", + " for a_s in itertools.product(a_range, repeat=G1):\n", + " pa = np.prod([p1n[a] for a in a_s])\n", + " for b_all in itertools.product(*[itertools.product(b_range, repeat=G2) for _ in range(G1)]):\n", + " pb = np.prod([p2n[a, b] for a, bs in zip(a_s, b_all) for b in bs])\n", + " Qhat = np.array([np.mean([Rret[(a, b)] for b in bs]) for a, bs in zip(a_s, b_all)])\n", + " est = np.zeros(nparam)\n", + " for i, (a, bs) in enumerate(zip(a_s, b_all)):\n", + " b1 = (Qhat.sum() - Qhat[i]) / (G1 - 1)\n", + " est += (Qhat[i] - b1) * s1(a) / G1\n", + " Rs = np.array([Rret[(a, b)] for b in bs])\n", + " for j, b in enumerate(bs):\n", + " b2 = (Rs.sum() - Rs[j]) / (G2 - 1)\n", + " est += (Rs[j] - b2) * s2(a, b) / (G1 * G2)\n", + " pr = pa * pb; m1 += pr * est; m2 += pr * (est ** 2).sum()\n", + " return m1, m2 - (m1 ** 2).sum()\n", + "m_tree, v_tree = run_tree()\n", + "print(f'tree LOO (2 branches x 2 siblings = 4 episodes): |E-gradJ|max={np.abs(m_tree-gradJ).max():.3e} trace Var={v_tree:.5f}')\n", + "check('6. tree two-level leave-one-out with per-turn baselines: E[g] = grad J', np.abs(m_tree - gradJ).max(), 1e-14)\n", + "print(f'equal budget (4 episodes): trace Var E0={res[\"E0\"][1]:.4f} E0rtg={res[\"E0rtg\"][1]:.4f} flat LOO E2={res[\"E2\"][1]:.4f} tree LOO={v_tree:.4f}')\n", + "\n", + "# Value of baselines: add a constant c to every return (a \"reward offset\"; unbiasedness is unaffected, variance is not).\n", + "for c in (0.0, 3.0):\n", + " Rl_save = Rl.copy(); Rret_save = dict(Rret)\n", + " Rl = Rl_save + c; Rret = {k_: v_ + c for k_, v_ in Rret_save.items()}\n", + " v0 = run_flat(Gs, 'E0')[1]; v2 = run_flat(Gs, 'E2')[1]; vt = run_tree()[1]\n", + " print(f'offset c={c}: trace Var no-baseline={v0:.3f} flat LOO={v2:.3f} tree LOO={vt:.3f}')\n", + " Rl = Rl_save; Rret = Rret_save" + ] + }, + { + "cell_type": "markdown", + "id": "c2aa29c6", + "metadata": {}, + "source": [ + "## 7. Training only the first correction round on a stale context distribution (Goedel-Prover-V2-style RL data)\n", + "Use contexts `y1 ~ mu` (a frozen earlier policy) and train only `pi2`. Its expected gradient is\n", + "`sum_a mu(a) sum_b pi2(b|a) R(a,b) grad log pi2(b|a)` = `sum_a mu(a) grad_{theta2} Q(a)`. The full-episode gradient on the `theta2` block is\n", + "`sum_a pi1(a) grad_{theta2} Q(a)` and the full gradient has an additional `theta1` block. Compare." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "c9a05411", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:01:55.103908Z", + "iopub.status.busy": "2026-10-14T09:01:55.103731Z", + "iopub.status.idle": "2026-10-14T09:01:55.110368Z", + "shell.execute_reply": "2026-10-14T09:01:55.109451Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "7. truncated turn-2 gradient = sum_a mu(a) grad Q(a) (closed form) err=8.674e-18 (tol 1e-16)\n", + "theta2-block bias norm (context shift only): 0.10440110286572064 theta1-block (missing turn-1 term) norm: 0.4583795568241773\n", + "7. with mu = pi1 the theta2 block of the truncated gradient equals that of grad J err=8.327e-17 (tol 1e-15)\n" + ] + } + ], + "source": [ + "mu = np.array([0.6, 0.3, 0.1])\n", + "Qa = np.array([sum(p2n[a, b] * Rret[(a, b)] for b in range(K2)) for a in range(K1)])\n", + "E_trunc = np.zeros(nparam)\n", + "for a in range(K1):\n", + " for b in range(K2):\n", + " E_trunc += mu[a] * p2n[a, b] * Rret[(a, b)] * s2(a, b)\n", + "# closed form: sum_a mu(a) grad_{theta2} Q(a); grad_{theta2[a]} Q(a) = pi2(b|a)(R(a,b) - Q(a)) over b\n", + "E_closed = np.zeros(nparam)\n", + "for a in range(K1):\n", + " for b in range(K2):\n", + " E_closed[K1 + a * K2 + b] = mu[a] * p2n[a, b] * (Rret[(a, b)] - Qa[a])\n", + "check('7. truncated turn-2 gradient = sum_a mu(a) grad Q(a) (closed form)', np.abs(E_trunc - E_closed).max(), 1e-16)\n", + "bias_ctx = E_trunc - gradJ\n", + "b2_block = bias_ctx[K1:]; b1_block = bias_ctx[:K1]\n", + "print('theta2-block bias norm (context shift only):', float(np.linalg.norm(b2_block)), ' theta1-block (missing turn-1 term) norm:', float(np.linalg.norm(b1_block)))\n", + "# if mu = pi1 the theta2 block is unbiased\n", + "E_same = np.zeros(nparam)\n", + "for a in range(K1):\n", + " for b in range(K2):\n", + " E_same += p1n[a] * p2n[a, b] * Rret[(a, b)] * s2(a, b)\n", + "check('7. with mu = pi1 the theta2 block of the truncated gradient equals that of grad J', np.abs(E_same[K1:] - gradJ[K1:]).max(), 1e-15)\n", + "CHECKS.append(('7. norm of the missing turn-1 block ||grad_theta1 J|| (informational, must be > 0)', float(np.linalg.norm(gradJ[:K1])), float('inf')))\n", + "assert np.linalg.norm(b1_block) > 1e-3 and np.linalg.norm(b2_block) > 1e-4" + ] + }, + { + "cell_type": "markdown", + "id": "a9073ea5", + "metadata": {}, + "source": [ + "## Verification table" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "b9986b19", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:01:55.112041Z", + "iopub.status.busy": "2026-10-14T09:01:55.111864Z", + "iopub.status.idle": "2026-10-14T09:01:55.115491Z", + "shell.execute_reply": "2026-10-14T09:01:55.114662Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + "identity max abs err\n", + "1. E[R (score_turn1 + score_turn2)] = grad J (exact enumeration of 9 episodes) 8.327e-17\n", + "2. E[r1 * score_turn2] = 0 1.388e-17\n", + "2. reward-to-go estimator: E = grad J 6.939e-17\n", + "3. norm of the bias from including feedback tokens in the loss (informational, must be > 0) 6.799e-01\n", + "4. E0 (full-return, no baseline, G=4 episodes): unbiased 4.996e-16\n", + "4. E0rtg (reward-to-go, G=4): unbiased 4.996e-16\n", + "4. identity R_i - mean_all = (G-1)/G (R_i - mean_{-i}) on every group (max abs diff) 5.551e-17\n", + "4. GRPO mean baseline: E[g] = (G-1)/G grad J exactly 1.193e-15\n", + "5. leave-one-out baseline (episode-level, G=4): E[g] = grad J 8.327e-16\n", + "5. baseline depending on the sample own r2: |bias| (informational, must be > 0) 7.857e-02\n", + "4. variance(E1 scaled by G/(G-1)) = variance(E2) 7.438e-15\n", + "6. tree two-level leave-one-out with per-turn baselines: E[g] = grad J 2.220e-16\n", + "7. truncated turn-2 gradient = sum_a mu(a) grad Q(a) (closed form) 8.674e-18\n", + "7. with mu = pi1 the theta2 block of the truncated gradient equals that of grad J 8.327e-17\n", + "7. norm of the missing turn-1 block ||grad_theta1 J|| (informational, must be > 0) 4.584e-01\n", + "CHECKS_JSON:[[\"1. E[R (score_turn1 + score_turn2)] = grad J (exact enumeration of 9 episodes)\", 8.326672684688674e-17, 1e-15], [\"2. E[r1 * score_turn2] = 0\", 1.3877787807814457e-17, 1e-16], [\"2. reward-to-go estimator: E = grad J\", 6.938893903907228e-17, 1e-15], [\"3. norm of the bias from including feedback tokens in the loss (informational, must be > 0)\", 0.679933398028168, Infinity], [\"4. E0 (full-return, no baseline, G=4 episodes): unbiased\", 4.996003610813204e-16, 1e-15], [\"4. E0rtg (reward-to-go, G=4): unbiased\", 4.996003610813204e-16, 1e-15], [\"4. identity R_i - mean_all = (G-1)/G (R_i - mean_{-i}) on every group (max abs diff)\", 5.551115123125783e-17, 1e-14], [\"4. GRPO mean baseline: E[g] = (G-1)/G grad J exactly\", 1.1934897514720433e-15, 1e-14], [\"5. leave-one-out baseline (episode-level, G=4): E[g] = grad J\", 8.326672684688674e-16, 1e-14], [\"5. baseline depending on the sample own r2: |bias| (informational, must be > 0)\", 0.07857014805316431, Infinity], [\"4. variance(E1 scaled by G/(G-1)) = variance(E2)\", 7.438494264988549e-15, 1e-12], [\"6. tree two-level leave-one-out with per-turn baselines: E[g] = grad J\", 2.220446049250313e-16, 1e-14], [\"7. truncated turn-2 gradient = sum_a mu(a) grad Q(a) (closed form)\", 8.673617379884035e-18, 1e-16], [\"7. with mu = pi1 the theta2 block of the truncated gradient equals that of grad J\", 8.326672684688674e-17, 1e-15], [\"7. norm of the missing turn-1 block ||grad_theta1 J|| (informational, must be > 0)\", 0.4583795568241773, Infinity]]\n" + ] + } + ], + "source": [ + "print()\n", + "print(f'{\"identity\":90s} {\"max abs err\":>12s}')\n", + "for n_, e_, t_ in CHECKS: print(f'{n_:90s} {e_:12.3e}')\n", + "print('CHECKS_JSON:' + json.dumps(CHECKS))" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/code/notebooks/0205-model-averaging-weight-interpolation.ipynb b/code/notebooks/0205-model-averaging-weight-interpolation.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..9af035010d024d1d2c757be6416ff495ba554e6a --- /dev/null +++ b/code/notebooks/0205-model-averaging-weight-interpolation.ipynb @@ -0,0 +1,746 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "5ac94761", + "metadata": {}, + "source": [ + "# 0205 -- Model averaging / weight interpolation `theta_alpha = (1-alpha) theta_0 + alpha theta`\n", + "\n", + "Convention (Goedel-Prover-V2 Sec. 2.3): `alpha` is the weight on the **fine-tuned** model; `alpha=0` is the base. (Their Sec. 3.6 text states\n", + "that alpha is \"the weight of the base model\" -- the opposite convention; see the note.)\n", + "\n", + "Identities / claims verified (fp64, CPU):\n", + "1. Logit-linear model: `softmax(z_alpha) = normalised pi_0^(1-alpha) pi_1^alpha` (geometric mixture), not the arithmetic mixture; naive (re-evaluate at theta_alpha) = efficient (interpolate logits).\n", + "2. Exponential tilt: `d/dalpha E_alpha[r] = Cov_alpha(g, r)`; for a binary tilt `p(alpha) = sigmoid(logit p0 + alpha c)` exactly.\n", + "3. Log-odds-linear first-order formulas: `d pass@k / d alpha = E_x[k p (1-p)^k Delta_x]`.\n", + "4. Curvature along the path for cross-entropy: `d^2 L/dalpha^2 = E_x Var_{y~softmax(z_alpha)}[Delta z_y]` (model-soups Eq. 1 with linear logits: soup = ensemble).\n", + "5. Loss along a line: `l(alpha) - chord = (1/2) alpha (alpha-1) d^T H d` for a quadratic; convex NLL has no barrier.\n", + "6. Linear mode connectivity: barrier for same-init fine-tunes vs independently trained nets (tiny MLP).\n", + "7. 30 random log-linear \"SFT -> RL\" pairs: pass@1 rises with alpha; pass@k (k=8,64) peaks at interior alpha." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "4ac1843a", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:49:30.138556Z", + "iopub.status.busy": "2026-10-14T09:49:30.138404Z", + "iopub.status.idle": "2026-10-14T09:49:31.926821Z", + "shell.execute_reply": "2026-10-14T09:49:31.925791Z" + } + }, + "outputs": [], + "source": [ + "import os\n", + "os.environ['CUDA_VISIBLE_DEVICES'] = ''\n", + "import math, json\n", + "import numpy as np, torch\n", + "torch.set_default_dtype(torch.float64)\n", + "CHECKS = []\n", + "def check(name, err, tol):\n", + " CHECKS.append((name, float(err), float(tol)))\n", + " assert err <= tol, (name, float(err), tol)\n", + " print(f'{name:86s} err={float(err):.3e} (tol {tol:.0e})')\n", + "\n", + "g = np.random.RandomState(0)\n", + "m, K, d = 40, 12, 5\n", + "Phi = torch.tensor(g.randn(m, K, d)); V = torch.tensor((g.rand(m, K) < 0.12).astype(float))\n", + "V[V.sum(-1) == 0, 0] = 1.0\n", + "th0 = torch.tensor(g.randn(d) * 0.4); th1 = th0 + torch.tensor(g.randn(d) * 1.2)" + ] + }, + { + "cell_type": "markdown", + "id": "97ec1d44", + "metadata": {}, + "source": [ + "## 1. Geometric (log-linear) mixture\n", + "NAIVE: build `theta_alpha` and run the model. EFFICIENT: evaluate the two endpoint logit tables once and interpolate the logits.\n", + "Both equal the per-context normalised product `pi_0^(1-alpha) pi_1^alpha`. The arithmetic mixture `(1-alpha) pi_0 + alpha pi_1` (an ensemble) differs." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "91130ea5", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:49:31.929274Z", + "iopub.status.busy": "2026-10-14T09:49:31.929030Z", + "iopub.status.idle": "2026-10-14T09:49:31.936414Z", + "shell.execute_reply": "2026-10-14T09:49:31.935536Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "1. naive (re-evaluate theta_alpha) vs efficient (interpolate logits) err=4.441e-16 (tol 1e-15)\n", + "1. softmax(z_alpha) vs normalised pi0^(1-a) pi1^a err=4.441e-16 (tol 1e-15)\n", + "max |geometric - arithmetic mixture| for alpha in (0,1): 0.25554158391614934\n" + ] + } + ], + "source": [ + "z0, z1 = Phi @ th0, Phi @ th1\n", + "pi0, pi1 = torch.softmax(z0, -1), torch.softmax(z1, -1)\n", + "worst_naive_eff = 0.0; worst_geo = 0.0; dist_arith = 0.0\n", + "for a in (0.0, 0.3, 0.7, 1.0, 1.25, -0.2):\n", + " tha = (1 - a) * th0 + a * th1\n", + " naive = torch.softmax(Phi @ tha, -1)\n", + " eff = torch.softmax((1 - a) * z0 + a * z1, -1)\n", + " geo = pi0 ** (1 - a) * pi1 ** a; geo = geo / geo.sum(-1, keepdim=True)\n", + " ar = (1 - a) * pi0 + a * pi1\n", + " worst_naive_eff = max(worst_naive_eff, float((naive - eff).abs().max()))\n", + " worst_geo = max(worst_geo, float((naive - geo).abs().max()))\n", + " if 0 < a < 1: dist_arith = max(dist_arith, float((naive - ar).abs().max()))\n", + "check('1. naive (re-evaluate theta_alpha) vs efficient (interpolate logits)', worst_naive_eff, 1e-15)\n", + "check('1. softmax(z_alpha) vs normalised pi0^(1-a) pi1^a', worst_geo, 1e-15)\n", + "print('max |geometric - arithmetic mixture| for alpha in (0,1):', dist_arith)\n", + "assert dist_arith > 1e-3" + ] + }, + { + "cell_type": "markdown", + "id": "74e8c4e1", + "metadata": {}, + "source": [ + "## 2. Exponential tilt\n", + "If `pi_1 = pi_0 e^{g}/Z` then `pi_alpha = pi_0 e^{alpha g}/Z_alpha`, `d/dalpha E_alpha[r] = Cov_alpha(g, r)` (derivative of a log-partition function).\n", + "Binary tilt `g = c r`: odds multiply by `e^{alpha c}`." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "1cc640d6", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:49:31.938137Z", + "iopub.status.busy": "2026-10-14T09:49:31.937942Z", + "iopub.status.idle": "2026-10-14T09:49:32.021196Z", + "shell.execute_reply": "2026-10-14T09:49:32.020228Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "2. d/dalpha E_alpha[r] (autograd) = Cov_alpha(g, r) err=5.204e-18 (tol 1e-16)\n", + "2. binary tilt: p(alpha) = sigmoid(logit p0 + alpha c) err=0.000e+00 (tol 1e-15)\n" + ] + } + ], + "source": [ + "gtilt = Phi @ (th1 - th0) # = z1 - z0, the tilt of each context\n", + "a = torch.tensor(0.4, requires_grad=True)\n", + "pa = torch.softmax(z0 + a * gtilt, -1)\n", + "Ea = (pa * V).sum(-1).mean(); Ea.backward()\n", + "pa_ = pa.detach(); cov = ((pa_ * gtilt * V).sum(-1) - (pa_ * gtilt).sum(-1) * (pa_ * V).sum(-1)).mean()\n", + "check('2. d/dalpha E_alpha[r] (autograd) = Cov_alpha(g, r)', abs(float(a.grad - cov)), 1e-16)\n", + "p0v = float((pi0 * V).sum(-1)[0]); c = 1.7\n", + "q = pi0[0] * torch.exp(c * V[0]); q = q / q.sum()\n", + "worst = 0.0\n", + "for al in (0.0, 0.25, 0.5, 1.0, 2.0):\n", + " qa = pi0[0] * torch.exp(al * c * V[0]); qa = qa / qa.sum()\n", + " p_al = float((qa * V[0]).sum()); logit = math.log(p0v / (1 - p0v)) + al * c\n", + " worst = max(worst, abs(p_al - 1 / (1 + math.exp(-logit))))\n", + "check('2. binary tilt: p(alpha) = sigmoid(logit p0 + alpha c)', worst, 1e-15)" + ] + }, + { + "cell_type": "markdown", + "id": "87a33cf6", + "metadata": {}, + "source": [ + "## 2b. Scope of the covariance identity: single categorical vs autoregressive policy\n", + "For an autoregressive policy `pi_alpha(y) = prod_t softmax((1-alpha) z0 + alpha z1)(y_t | s_t)` the normaliser depends on the context, so\n", + "`d/dalpha log pi_alpha(y) = sum_t [g_t - E_{pi_alpha(.|s_t)} g]` (the *sequence score*), and `d/dalpha E_alpha[r] = Cov_alpha(score, r)`.\n", + "Replacing the score by `sum_t g_t` is NOT exact. Two-token toy, vocabulary 3, exact enumeration of the 9 sequences, central finite difference in `alpha`." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "e54a128b", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:49:32.022698Z", + "iopub.status.busy": "2026-10-14T09:49:32.022530Z", + "iopub.status.idle": "2026-10-14T09:49:32.032914Z", + "shell.execute_reply": "2026-10-14T09:49:32.032036Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "finite difference dE/dalpha : -0.19188850988727246\n", + "Cov_alpha(centred sequence score, r) : -0.19188850994087786\n", + "Cov_alpha(sum_t g_t, r) : -0.17690360845304756\n", + "2b. autoregressive: dE/dalpha (finite diff.) = Cov_alpha(sequence score, r) err=5.361e-11 (tol 1e-08)\n", + "2b. zero mean of the sequence score: E_alpha[score] = 0 err=5.204e-17 (tol 1e-14)\n" + ] + } + ], + "source": [ + "import itertools\n", + "rg_ar = np.random.RandomState(0); Vv_ = 3\n", + "z0_ = {(): rg_ar.randn(Vv_)}; z1_ = {(): rg_ar.randn(Vv_)}\n", + "for a_ in range(Vv_): z0_[(a_,)] = rg_ar.randn(Vv_); z1_[(a_,)] = rg_ar.randn(Vv_)\n", + "sm_ = lambda zz: np.exp(zz - zz.max()) / np.exp(zz - zz.max()).sum()\n", + "seqs_ = list(itertools.product(range(Vv_), repeat=2)); r_ar = rg_ar.rand(len(seqs_))\n", + "def dist_ar(al):\n", + " return np.array([sm_((1 - al) * z0_[()] + al * z1_[()])[a_] * sm_((1 - al) * z0_[(a_,)] + al * z1_[(a_,)])[b_] for a_, b_ in seqs_])\n", + "al0 = 0.4; hh = 1e-6\n", + "fd = (dist_ar(al0 + hh) @ r_ar - dist_ar(al0 - hh) @ r_ar) / (2 * hh)\n", + "pdist = dist_ar(al0)\n", + "gtok = lambda ctx, y: (z1_[ctx] - z0_[ctx])[y]\n", + "cmean = lambda ctx: sm_((1 - al0) * z0_[ctx] + al0 * z1_[ctx]) @ (z1_[ctx] - z0_[ctx])\n", + "G_sum = np.array([gtok((), a_) + gtok((a_,), b_) for a_, b_ in seqs_])\n", + "G_score = np.array([(gtok((), a_) - cmean(())) + (gtok((a_,), b_) - cmean((a_,))) for a_, b_ in seqs_])\n", + "covf = lambda x: (pdist * x * r_ar).sum() - (pdist * x).sum() * (pdist * r_ar).sum()\n", + "print('finite difference dE/dalpha :', fd)\n", + "print('Cov_alpha(centred sequence score, r) :', covf(G_score))\n", + "print('Cov_alpha(sum_t g_t, r) :', covf(G_sum))\n", + "check('2b. autoregressive: dE/dalpha (finite diff.) = Cov_alpha(sequence score, r)', abs(fd - covf(G_score)), 1e-8)\n", + "assert abs(fd - covf(G_sum)) > 1e-3\n", + "CHECKS.append(('2b. autoregressive: error of the naive Cov_alpha(sum_t g_t, r) (informational, must be > 0)', float(abs(fd - covf(G_sum))), float('inf')))\n", + "check('2b. zero mean of the sequence score: E_alpha[score] = 0', abs((pdist * G_score).sum()), 1e-14)" + ] + }, + { + "cell_type": "markdown", + "id": "e3649b2b", + "metadata": {}, + "source": [ + "## 3. First-order (log-odds-linear) analysis of pass@1 and pass@k\n", + "Assume `logit p_x(alpha) = l_x + alpha Delta_x`. Then `d p/dalpha = p(1-p) Delta`, so\n", + "`d pass@1/dalpha = E[p(1-p) Delta]`, `d pass@k/dalpha = E[k p (1-p)^k Delta]`." + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "0a878415", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:49:32.034304Z", + "iopub.status.busy": "2026-10-14T09:49:32.034149Z", + "iopub.status.idle": "2026-10-14T09:49:32.087078Z", + "shell.execute_reply": "2026-10-14T09:49:32.086240Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "3. d pass@k/dalpha = E[k p (1-p)^k Delta] vs autograd, k in {1,4,32} err=2.776e-17 (tol 1e-15)\n", + "k= 1: alpha* = 1.000 pass@k(0)=0.2776 pass@k(alpha*)=0.4355 pass@k(1)=0.4355\n", + "k= 8: alpha* = 0.525 pass@k(0)=0.6717 pass@k(alpha*)=0.6762 pass@k(1)=0.6738\n", + "k= 32: alpha* = 0.000 pass@k(0)=0.8998 pass@k(alpha*)=0.8998 pass@k(1)=0.7868\n", + "p maximising the marginal value of a log-odds shift, k p (1-p)^k : 1/(k+1) = {1: 0.5, 4: 0.2, 32: 0.030303030303030304}\n" + ] + } + ], + "source": [ + "rg = np.random.RandomState(2)\n", + "# Two populations (synthetic, to exhibit the regimes): 65% \"learned\" prompts that the RL stage improves, 35% \"tail\" prompts (hard, p ~ 5%)\n", + "# that it degrades. logit p_x(alpha) = l_x + alpha Delta_x.\n", + "N = 2000; nt = int(0.35 * N); n1 = N - nt\n", + "l = torch.tensor(np.concatenate([rg.randn(n1) * 1.5 - 0.5, rg.randn(nt) * 0.5 - 3.0]))\n", + "Delta = torch.tensor(np.concatenate([rg.randn(n1) * 0.3 + 1.5, rg.randn(nt) * 0.3 - 1.2]))\n", + "def passk(al, k):\n", + " p = torch.sigmoid(l + al * Delta); return (1 - (1 - p) ** k).mean()\n", + "worst = 0.0\n", + "for k in (1, 4, 32):\n", + " for al in (0.0, 0.5, 1.0):\n", + " a_ = torch.tensor(al, requires_grad=True)\n", + " p = torch.sigmoid(l + a_ * Delta); (1 - (1 - p) ** k).mean().backward()\n", + " p_ = p.detach(); formula = (k * p_ * (1 - p_) ** k * Delta).mean()\n", + " worst = max(worst, abs(float(a_.grad - formula)))\n", + "check('3. d pass@k/dalpha = E[k p (1-p)^k Delta] vs autograd, k in {1,4,32}', worst, 1e-15)\n", + "als = np.linspace(0, 1, 201)\n", + "star = {}\n", + "for k in (1, 8, 32):\n", + " vals = np.array([float(passk(a_, k)) for a_ in als]); star[k] = als[vals.argmax()]\n", + " print(f'k={k:3d}: alpha* = {als[vals.argmax()]:.3f} pass@k(0)={vals[0]:.4f} pass@k(alpha*)={vals.max():.4f} pass@k(1)={vals[-1]:.4f}')\n", + "assert star[1] == 1.0 and 0.0 < star[8] < 1.0 and star[32] < star[8]\n", + "CHECKS.append(('3. regimes: alpha*(k=1) = 1, 0 < alpha*(k=8) < 1, alpha*(k=32) < alpha*(k=8) (alpha*: %s)' % {k_: round(float(v_), 3) for k_, v_ in star.items()}, 0.0, float('inf')))\n", + "peak_p = {k: 1.0 / (k + 1) for k in (1, 4, 32)}\n", + "print('p maximising the marginal value of a log-odds shift, k p (1-p)^k : 1/(k+1) =', peak_p)" + ] + }, + { + "cell_type": "markdown", + "id": "6860011a", + "metadata": {}, + "source": [ + "## 3b. The end-point sign test is sufficient, not necessary, for an interior optimum\n", + "Two populations on logit-linear paths: half of the prompts `l=-3, Delta=-1` (declining), half `l=-8, Delta=+10` (late bloomers). `k=8`." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "24bcee3d", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:49:32.088723Z", + "iopub.status.busy": "2026-10-14T09:49:32.088559Z", + "iopub.status.idle": "2026-10-14T09:49:32.094636Z", + "shell.execute_reply": "2026-10-14T09:49:32.093860Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "d/dalpha pass@8 at 0: -0.1152, at 1: -0.0622; pass@8(0)=0.1624 pass@8(1)=0.5676 max=0.5792 at alpha=0.802\n", + "3b. derivative formula E[k p (1-p)^k Delta] vs finite difference (this population) err=1.250e-10 (tol 1e-07)\n" + ] + } + ], + "source": [ + "sig_ = lambda x: 1 / (1 + np.exp(-x))\n", + "pop_ = [(0.5, -3.0, -1.0), (0.5, -8.0, 10.0)]\n", + "pk_ = lambda a_, k_: sum(w_ * (1 - (1 - sig_(l_ + a_ * D_)) ** k_) for w_, l_, D_ in pop_)\n", + "dpk_ = lambda a_, k_: sum(w_ * k_ * sig_(l_ + a_ * D_) * (1 - sig_(l_ + a_ * D_)) ** k_ * D_ for w_, l_, D_ in pop_)\n", + "grid_ = np.linspace(0, 1, 2001); f_ = pk_(grid_, 8)\n", + "print(f'd/dalpha pass@8 at 0: {dpk_(0.0, 8):.4f}, at 1: {dpk_(1.0, 8):.4f}; pass@8(0)={f_[0]:.4f} pass@8(1)={f_[-1]:.4f} max={f_.max():.4f} at alpha={grid_[f_.argmax()]:.3f}')\n", + "assert dpk_(0.0, 8) < 0 and dpk_(1.0, 8) < 0 and 0 < f_.argmax() < len(grid_) - 1 and f_.max() > f_[-1] + 1e-3 and f_.max() > f_[0] + 1e-3\n", + "fd_ = (pk_(0.3 + 1e-6, 8) - pk_(0.3 - 1e-6, 8)) / 2e-6\n", + "check('3b. derivative formula E[k p (1-p)^k Delta] vs finite difference (this population)', abs(fd_ - dpk_(0.3, 8)), 1e-7)\n", + "CHECKS.append(('3b. interior global optimum with negative derivative at both ends (alpha* = %.3f, informational)' % grid_[f_.argmax()], float(grid_[f_.argmax()]), float('inf')))" + ] + }, + { + "cell_type": "markdown", + "id": "b777f25a", + "metadata": {}, + "source": [ + "## 4. Curvature along the path: model-soups Eq. (1) for logit-linear models\n", + "Cross-entropy of a designated correct proof `y*` per context: `L(alpha) = -E_x log softmax(z_alpha)[y*]`. Then\n", + "`d^2 L / dalpha^2 = E_x Var_{y ~ softmax(z_alpha)}[Delta z_y]` (second derivative of a log-partition function) and the soup equals the logit ensemble." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "c4618972", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:49:32.095990Z", + "iopub.status.busy": "2026-10-14T09:49:32.095840Z", + "iopub.status.idle": "2026-10-14T09:49:32.101953Z", + "shell.execute_reply": "2026-10-14T09:49:32.101282Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "4. d^2 L/dalpha^2 (autograd) = E_x Var_{softmax(z_alpha)}[Delta z] err=4.441e-16 (tol 1e-15)\n", + "4. soup loss = logit-ensemble loss for logit-linear models err=0.000e+00 (tol 1e-16)\n" + ] + } + ], + "source": [ + "ystar = torch.tensor(np.argmax(V.numpy(), axis=1))\n", + "def L_of(al): return -torch.log_softmax(z0 + al * (z1 - z0), -1)[torch.arange(m), ystar].mean()\n", + "al = torch.tensor(0.35, requires_grad=True)\n", + "L = L_of(al); g1, = torch.autograd.grad(L, al, create_graph=True); g2, = torch.autograd.grad(g1, al)\n", + "pz = torch.softmax(z0 + al.detach() * (z1 - z0), -1); dz = z1 - z0\n", + "var_dz = ((pz * dz ** 2).sum(-1) - (pz * dz).sum(-1) ** 2).mean()\n", + "check('4. d^2 L/dalpha^2 (autograd) = E_x Var_{softmax(z_alpha)}[Delta z]', abs(float(g2 - var_dz)), 1e-15)\n", + "# soup vs logit-ensemble loss difference (soups Eq.(1) predicts ~0 when logits are linear in theta)\n", + "ens_logits = (1 - al.detach()) * z0 + al.detach() * z1\n", + "L_ens = -torch.log_softmax(ens_logits, -1)[torch.arange(m), ystar].mean()\n", + "check('4. soup loss = logit-ensemble loss for logit-linear models', abs(float(L_of(al.detach()) - L_ens)), 1e-16)" + ] + }, + { + "cell_type": "markdown", + "id": "e774bc82", + "metadata": {}, + "source": [ + "## 5. Loss along a line: chord identity and barriers\n", + "Quadratic `l(theta) = l_0 + g^T (theta-theta_0) + 1/2 (theta-theta_0)^T H (theta-theta_0)`, `d = theta_1-theta_0`, `kappa = d^T H d`:\n", + "`l(alpha) - [(1-alpha) l(0) + alpha l(1)] = (1/2) alpha (alpha-1) kappa`. For `H` PSD this is `<= 0` (no barrier); the averaged model beats the chord." + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "a8781b64", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:49:32.103566Z", + "iopub.status.busy": "2026-10-14T09:49:32.103406Z", + "iopub.status.idle": "2026-10-14T09:49:32.127330Z", + "shell.execute_reply": "2026-10-14T09:49:32.126399Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "5. quadratic: l(alpha) - chord = alpha(alpha-1) kappa / 2 err=9.770e-15 (tol 1e-12)\n", + "5. convex softmax-linear NLL: barrier max_alpha [l(alpha)-chord] <= 0 err=0.000e+00 (tol 1e-14)\n", + "barrier of the convex NLL (should be <= 0): 0.0\n" + ] + } + ], + "source": [ + "rq = np.random.RandomState(8)\n", + "A = rq.randn(7, 7); H = torch.tensor(A @ A.T); gv = torch.tensor(rq.randn(7)); dd = torch.tensor(rq.randn(7))\n", + "lq = lambda t: 0.3 + gv @ t + 0.5 * t @ H @ t\n", + "kappa = float(dd @ H @ dd)\n", + "worst = max(abs(float(lq(a_ * dd) - ((1 - a_) * lq(0 * dd) + a_ * lq(dd))) - 0.5 * a_ * (a_ - 1) * kappa) for a_ in np.linspace(-0.5, 1.5, 9))\n", + "check('5. quadratic: l(alpha) - chord = alpha(alpha-1) kappa / 2', worst, 1e-12)\n", + "# convex NLL of the log-linear policy: no barrier\n", + "def nll(al): return float(L_of(torch.tensor(al)))\n", + "grid = np.linspace(0, 1, 101)\n", + "bar = max(nll(a_) - ((1 - a_) * nll(0.0) + a_ * nll(1.0)) for a_ in grid)\n", + "check('5. convex softmax-linear NLL: barrier max_alpha [l(alpha)-chord] <= 0', max(0.0, bar), 1e-14)\n", + "print('barrier of the convex NLL (should be <= 0):', bar)" + ] + }, + { + "cell_type": "markdown", + "id": "8f32e360", + "metadata": {}, + "source": [ + "## 6. Linear mode connectivity in a tiny non-convex net\n", + "2-D spiral, 2 classes. A 2-64-64-2 tanh MLP, full-batch Adam in fp64.\n", + "* Same-init: pretrain on task A, then two fine-tunes on a rotated/noisier task B with different seeds.\n", + "* Independent: two nets trained from different random inits on task B.\n", + "Barrier `B = max_alpha [ L(theta_alpha) - ((1-alpha) L(theta_0) + alpha L(theta_1)) ]` evaluated on the union of both tasks' data." + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "94b56789", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:49:32.128822Z", + "iopub.status.busy": "2026-10-14T09:49:32.128648Z", + "iopub.status.idle": "2026-10-14T09:49:37.344398Z", + "shell.execute_reply": "2026-10-14T09:49:37.343352Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "same-init fine-tunes : endpoints L=0.001,0.001 barrier=+0.0000\n", + "independent inits : endpoints L=0.000,0.000 barrier=+3.1884\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/home/vignesh/.opencode/tmp/ipykernel_606603/4015540266.py:25: UserWarning: Converting a tensor with requires_grad=True to a scalar may lead to unexpected behavior.\n", + "Consider using tensor.detach() first. (Triggered internally at /pytorch/torch/csrc/autograd/generated/python_variable_methods.cpp:838.)\n", + " L0, L1 = float(loss(P0, Xall, Yall)), float(loss(P1, Xall, Yall))\n" + ] + } + ], + "source": [ + "def spiral(n, rot, noise, seed):\n", + " r = np.random.RandomState(seed); t = np.linspace(0.2, 3.2, n)\n", + " X = []; Y = []\n", + " for c in (0, 1):\n", + " ang = t * 2.0 + c * math.pi + rot\n", + " X.append(np.stack([t * np.cos(ang), t * np.sin(ang)], 1) + r.randn(n, 2) * noise); Y.append(np.full(n, c))\n", + " return torch.tensor(np.concatenate(X)), torch.tensor(np.concatenate(Y))\n", + "XA, YA = spiral(300, 0.0, 0.03, 1); XB, YB = spiral(300, 0.25, 0.08, 2)\n", + "Xall, Yall = torch.cat([XA, XB]), torch.cat([YA, YB])\n", + "def init(seed, h=64):\n", + " r = torch.Generator().manual_seed(seed)\n", + " P = [torch.randn(2, h, generator=r) / math.sqrt(2), torch.zeros(h), torch.randn(h, h, generator=r) / math.sqrt(h), torch.zeros(h),\n", + " torch.randn(h, 2, generator=r) / math.sqrt(h), torch.zeros(2)]\n", + " return [p.clone().requires_grad_(True) for p in P]\n", + "def fwd(P, X): return torch.tanh(torch.tanh(X @ P[0] + P[1]) @ P[2] + P[3]) @ P[4] + P[5]\n", + "def loss(P, X, Y): return torch.nn.functional.cross_entropy(fwd(P, X), Y)\n", + "def train(P, X, Y, steps, lr, seed=None, batch=None):\n", + " o = torch.optim.Adam(P, lr=lr); gen = torch.Generator().manual_seed(seed if seed is not None else 0)\n", + " for _ in range(steps):\n", + " idx = torch.randperm(len(X), generator=gen)[:batch] if batch else slice(None)\n", + " o.zero_grad(); loss(P, X[idx], Y[idx]).backward(); o.step()\n", + " return [p.detach().clone().requires_grad_(True) for p in P]\n", + "def interp(P0, P1, a): return [((1 - a) * p0 + a * p1).detach() for p0, p1 in zip(P0, P1)]\n", + "def barrier(P0, P1):\n", + " L0, L1 = float(loss(P0, Xall, Yall)), float(loss(P1, Xall, Yall))\n", + " return max(float(loss(interp(P0, P1, a), Xall, Yall)) - ((1 - a) * L0 + a * L1) for a in np.linspace(0, 1, 41)), L0, L1\n", + "pre = train(init(0), XA, YA, 600, 1e-2)\n", + "ftA = train([p.detach().clone().requires_grad_(True) for p in pre], XB, YB, 300, 3e-3, seed=1, batch=64)\n", + "ftB = train([p.detach().clone().requires_grad_(True) for p in pre], XB, YB, 300, 3e-3, seed=2, batch=64)\n", + "indA = train(init(11), XB, YB, 900, 1e-2); indB = train(init(12), XB, YB, 900, 1e-2)\n", + "b_same, L0s, L1s = barrier(ftA, ftB); b_ind, L0i, L1i = barrier(indA, indB)\n", + "print(f'same-init fine-tunes : endpoints L={L0s:.3f},{L1s:.3f} barrier={b_same:+.4f}')\n", + "print(f'independent inits : endpoints L={L0i:.3f},{L1i:.3f} barrier={b_ind:+.4f}')\n", + "assert b_ind > 0.1 and b_same < b_ind / 5\n", + "CHECKS.append(('6. barrier, independent inits (informational, large)', float(b_ind), float('inf')))\n", + "CHECKS.append(('6. barrier, same-init fine-tunes (informational, small)', float(b_same), float('inf')))" + ] + }, + { + "cell_type": "markdown", + "id": "ac53ed4f", + "metadata": {}, + "source": [ + "## 7. SFT -> RL interpolation on 30 random log-linear policy families\n", + "Each family: 96 prompts, 24 candidate proofs, 6 features, ~8% of candidates verified. `theta_sft` = 80 steps of pass@1 ascent from a random init, `theta_rl` = 400 more.\n", + "For each family evaluate exact `pass@1, pass@8, pass@64` on `alpha in [0,1]` (101 points). NAIVE = rebuild `theta_alpha`; EFFICIENT = interpolate logits (asserted equal)." + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "1e316635", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:49:37.346390Z", + "iopub.status.busy": "2026-10-14T09:49:37.346111Z", + "iopub.status.idle": "2026-10-14T09:49:46.306787Z", + "shell.execute_reply": "2026-10-14T09:49:46.306034Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "families (of 30) with d pass@k/dalpha > 0 at alpha=0 / < 0 at alpha=1 : {1: [30, 0], 8: [14, 30], 64: [3, 30]}\n", + "families whose end-point derivative signs imply an interior optimum : {8: 14, 64: 3}\n", + "7. naive (theta_alpha) vs efficient (logit interpolation), all families err=6.661e-16 (tol 1e-15)\n", + "pass@1 non-decreasing in alpha: 30/30 families\n", + "interior argmax (0= -1e-12)), a8=als[P8.argmax()], a64=als[P64.argmax()],\n", + " p1_0=P1[0], p1_1=P1[-1], p8_0=P8[0], p8_1=P8[-1], p8_star=P8.max(), p64_0=P64[0], p64_1=P64[-1], p64_star=P64.max(),\n", + " p1_at64=P1[P64.argmax()]))\n", + " if seed == 0: curves0 = (P1.copy(), P8.copy(), P64.copy())\n", + "# Exact end-point derivatives: d pass@k/dalpha at alpha = E_x[ k (1-p)^(k-1) Cov_alpha(g, r) ], g = z1 - z0 (no log-odds assumption needed).\n", + "def deriv_k(Ph, Vv, za, zb, a_, k):\n", + " pz = torch.softmax((1 - a_) * za + a_ * zb, -1); gg = zb - za\n", + " p = (pz * Vv).sum(-1); cov = (pz * gg * Vv).sum(-1) - (pz * gg).sum(-1) * p\n", + " return float((k * (1 - p) ** (k - 1) * cov).mean())\n", + "sign_ok = 0; cnt = {1: [0, 0], 8: [0, 0], 64: [0, 0]}; inter_exact = {8: 0, 64: 0}\n", + "for seed in range(30):\n", + " Ph, Vv, ta, tb = family(seed); za, zb = Ph @ ta, Ph @ tb\n", + " for k in (1, 8, 64):\n", + " d0, d1 = deriv_k(Ph, Vv, za, zb, 0.0, k), deriv_k(Ph, Vv, za, zb, 1.0, k)\n", + " cnt[k][0] += d0 > 0; cnt[k][1] += d1 < 0\n", + " if k in (8, 64):\n", + " row = rows[seed]; astar = row['a8'] if k == 8 else row['a64']\n", + " assert (d1 >= 0) or (astar < 1.0) # derivative negative at alpha=1 => interpolating toward the base helps\n", + " assert (d0 <= 0) or (astar > 0.0) # derivative positive at alpha=0 => moving toward the fine-tune helps\n", + " inter_exact[k] += (d0 > 0 and d1 < 0)\n", + "print('families (of 30) with d pass@k/dalpha > 0 at alpha=0 / < 0 at alpha=1 :', cnt)\n", + "print('families whose end-point derivative signs imply an interior optimum :', inter_exact)\n", + "check('7. naive (theta_alpha) vs efficient (logit interpolation), all families', maxdiff, 1e-15)\n", + "import statistics as st\n", + "mono = sum(r_['mono1'] for r_ in rows)\n", + "inter8 = sum(0 < r_['a8'] < 1 for r_ in rows); inter64 = sum(0 < r_['a64'] < 1 for r_ in rows)\n", + "print(f'pass@1 non-decreasing in alpha: {mono}/30 families')\n", + "print(f'interior argmax (0 0) 1.498e-02\n", + "2b. zero mean of the sequence score: E_alpha[score] = 0 5.204e-17\n", + "3. d pass@k/dalpha = E[k p (1-p)^k Delta] vs autograd, k in {1,4,32} 2.776e-17\n", + "3. regimes: alpha*(k=1) = 1, 0 < alpha*(k=8) < 1, alpha*(k=32) < alpha*(k=8) (alpha*: {1: 1.0, 8: 0.525, 32: 0.0}) 0.000e+00\n", + "3b. derivative formula E[k p (1-p)^k Delta] vs finite difference (this population) 1.250e-10\n", + "3b. interior global optimum with negative derivative at both ends (alpha* = 0.802, informational) 8.020e-01\n", + "4. d^2 L/dalpha^2 (autograd) = E_x Var_{softmax(z_alpha)}[Delta z] 4.441e-16\n", + "4. soup loss = logit-ensemble loss for logit-linear models 0.000e+00\n", + "5. quadratic: l(alpha) - chord = alpha(alpha-1) kappa / 2 9.770e-15\n", + "5. convex softmax-linear NLL: barrier max_alpha [l(alpha)-chord] <= 0 0.000e+00\n", + "6. barrier, independent inits (informational, large) 3.188e+00\n", + "6. barrier, same-init fine-tunes (informational, small) 0.000e+00\n", + "7. naive (theta_alpha) vs efficient (logit interpolation), all families 6.661e-16\n", + "7. families with pass@1 monotone in alpha (count of 30, informational) 3.000e+01\n", + "7. families with interior argmax of pass@64 (count of 30, informational) 3.000e+00\n", + "8a. factor-wise LoRA average minus weight average = (1/4)(B1A2+B2A1-B1A1-B2A2) 6.661e-16\n", + "8b. fraction of weights whose bf16 value is unchanged when delta/weight = 1e-3 (informational) 6.853e-01\n", + "CHECKS_JSON:[[\"1. naive (re-evaluate theta_alpha) vs efficient (interpolate logits)\", 4.440892098500626e-16, 1e-15], [\"1. softmax(z_alpha) vs normalised pi0^(1-a) pi1^a\", 4.440892098500626e-16, 1e-15], [\"2. d/dalpha E_alpha[r] (autograd) = Cov_alpha(g, r)\", 5.204170427930421e-18, 1e-16], [\"2. binary tilt: p(alpha) = sigmoid(logit p0 + alpha c)\", 0.0, 1e-15], [\"2b. autoregressive: dE/dalpha (finite diff.) = Cov_alpha(sequence score, r)\", 5.3605397898337515e-11, 1e-08], [\"2b. autoregressive: error of the naive Cov_alpha(sum_t g_t, r) (informational, must be > 0)\", 0.014984901434224901, Infinity], [\"2b. zero mean of the sequence score: E_alpha[score] = 0\", 5.204170427930421e-17, 1e-14], [\"3. d pass@k/dalpha = E[k p (1-p)^k Delta] vs autograd, k in {1,4,32}\", 2.7755575615628914e-17, 1e-15], [\"3. regimes: alpha*(k=1) = 1, 0 < alpha*(k=8) < 1, alpha*(k=32) < alpha*(k=8) (alpha*: {1: 1.0, 8: 0.525, 32: 0.0})\", 0.0, Infinity], [\"3b. derivative formula E[k p (1-p)^k Delta] vs finite difference (this population)\", 1.249767789257561e-10, 1e-07], [\"3b. interior global optimum with negative derivative at both ends (alpha* = 0.802, informational)\", 0.802, Infinity], [\"4. d^2 L/dalpha^2 (autograd) = E_x Var_{softmax(z_alpha)}[Delta z]\", 4.440892098500626e-16, 1e-15], [\"4. soup loss = logit-ensemble loss for logit-linear models\", 0.0, 1e-16], [\"5. quadratic: l(alpha) - chord = alpha(alpha-1) kappa / 2\", 9.769962616701378e-15, 1e-12], [\"5. convex softmax-linear NLL: barrier max_alpha [l(alpha)-chord] <= 0\", 0.0, 1e-14], [\"6. barrier, independent inits (informational, large)\", 3.1884148897650437, Infinity], [\"6. barrier, same-init fine-tunes (informational, small)\", 0.0, Infinity], [\"7. naive (theta_alpha) vs efficient (logit interpolation), all families\", 6.661338147750939e-16, 1e-15], [\"7. families with pass@1 monotone in alpha (count of 30, informational)\", 30.0, Infinity], [\"7. families with interior argmax of pass@64 (count of 30, informational)\", 3.0, Infinity], [\"8a. factor-wise LoRA average minus weight average = (1/4)(B1A2+B2A1-B1A1-B2A2)\", 6.661338147750939e-16, 1e-13], [\"8b. fraction of weights whose bf16 value is unchanged when delta/weight = 1e-3 (informational)\", 0.6853, Infinity]]\n" + ] + } + ], + "source": [ + "print()\n", + "print(f'{\"identity\":90s} {\"max abs err\":>12s}')\n", + "for n_, e_, t_ in CHECKS: print(f'{n_:90s} {e_:12.3e}')\n", + "print('CHECKS_JSON:' + json.dumps(CHECKS))" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/code/notebooks/0206-curriculum-pass-rate-band.ipynb b/code/notebooks/0206-curriculum-pass-rate-band.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..2597a9badfae92645badb28f25cfaa7aece91494 --- /dev/null +++ b/code/notebooks/0206-curriculum-pass-rate-band.ipynb @@ -0,0 +1,545 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "2bf0eb1a", + "metadata": {}, + "source": [ + "# 0206 -- Curriculum by pass-rate band: signal of group advantages as a function of solve rate `p`\n", + "\n", + "Setting: one prompt with solve probability `p = sigmoid(z)`; `G` samples with binary rewards `r_i ~ Bern(p)`; `c = sum r_i`, `p_hat = c/G`.\n", + "Identities verified (fp64):\n", + "1. GRPO (unnormalised) surrogate gradient wrt the logit `z` equals `p_hat (1-p_hat)` for every `c` (autograd vs closed form).\n", + "2. `E[p_hat(1-p_hat)] = (1-1/G) p (1-p)`; the derivative of `p` wrt `z` is `p(1-p)`, maximal at `p=1/2` with value `1/4`.\n", + "3. Exact variance, SNR, and the dead-group probability `p^G + (1-p)^G`.\n", + "4. Std-normalised advantages: gradient `sqrt(p_hat(1-p_hat))`, expected value vs `sqrt(p(1-p))` (Dr. GRPO's reweighting).\n", + "5. General bound `||grad p||^2 <= p(1-p) E||s||^2` in a random log-linear policy family.\n", + "6. Pass-rate band `(0, tau]`: effective prompt weight `E[p_hat(1-p_hat) 1{0 0 else float('nan'), dead))\n", + " print(f' {p:5.2f} {m:.5f} {math.sqrt(v):.5f} {tab[-1][3]:6.3f} {dead:.5f} {1 - dead:.5f}')\n", + "# MC check of the exact moments\n", + "rgn = np.random.RandomState(0); p = 0.2; cs = rgn.binomial(8, p, size=2_000_000); ph = cs / 8; gh = ph * (1 - ph)\n", + "m, v = moments(8, p, f_un)\n", + "check('3. MC mean of p_hat(1-p_hat) vs exact (in s.e.)', abs(gh.mean() - m) / (gh.std() / math.sqrt(len(gh))), 4.0)\n", + "check('3. MC dead-group fraction vs p^G + (1-p)^G (in s.e.)', abs(((cs == 0) | (cs == 8)).mean() - (p ** 8 + (1 - p) ** 8)) / math.sqrt(0.5 * 0.5 / len(cs)), 4.0)\n", + "sig_peak = max(tab, key=lambda t: t[1]); snr_peak = max(tab, key=lambda t: t[3] if t[3] == t[3] else -1)\n", + "print('E[g] is largest at p =', sig_peak[0], '; SNR is largest at p =', snr_peak[0])\n", + "dead_at_half = 0.5 ** 8 * 2\n", + "print('dead prob at p=0.5, G=8:', dead_at_half, ' at p=0.9:', 0.9 ** 8 + 0.1 ** 8)" + ] + }, + { + "cell_type": "markdown", + "id": "a0e51cdd", + "metadata": {}, + "source": [ + "## 4. Std-normalised advantages (GRPO): gradient `sqrt(p_hat (1 - p_hat))`\n", + "`A_i = (r_i - p_hat)/sigma_hat`, `sigma_hat = sqrt(p_hat(1-p_hat))` (population std). Then `g_hat = (1/G) sum_i A_i (r_i - p) = p_hat(1-p_hat)/sigma_hat = sigma_hat`\n", + "(for `sigma_hat > 0`; zero otherwise). Compare `E[sigma_hat]` with `sqrt(p(1-p))` and the *relative* weight of a prompt versus the unnormalised estimator." + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "0a8dd37b", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:50:18.607353Z", + "iopub.status.busy": "2026-10-14T09:50:18.607193Z", + "iopub.status.idle": "2026-10-14T09:50:18.617653Z", + "shell.execute_reply": "2026-10-14T09:50:18.616857Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "4. std-normalised group gradient = sigma_hat (direct sum vs closed form) err=1.110e-16 (tol 1e-15)\n", + "4. unbiased-std normalisation: group gradient = sigma_hat sqrt((G-1)/G) err=1.110e-16 (tol 1e-15)\n", + "p E[sigma_hat] sqrt(p(1-p)) weight ratio E[sigma_hat]/E[p_hat(1-p_hat)] (unnormalised = 1 by definition)\n", + "0.02 0.05043 0.14000 2.941\n", + "0.10 0.20949 0.30000 2.660\n", + "0.25 0.37990 0.43301 2.316\n", + "0.50 0.46391 0.50000 2.121\n", + "0.90 0.20949 0.30000 2.660\n", + "relative prompt weight p=0.02 vs p=0.5: unnormalised 0.0784, std-normalised 0.1087\n", + "informative group c=1 vs c=4 (G=8): unnormalised 0.438, std-normalised 0.661, ratio of the two = 1/sigma_hat ratio = 1.512\n", + "4. std-normalisation re-weights c=1 vs c=4 groups by exactly sigma(4)/sigma(1) err=2.220e-16 (tol 1e-14)\n" + ] + } + ], + "source": [ + "worst = 0.0\n", + "for p in (0.1, 0.5):\n", + " for c in range(1, G):\n", + " r = torch.tensor([1.0] * c + [0.0] * (G - c)); ph = r.mean(); sg = torch.sqrt(ph * (1 - ph))\n", + " A = (r - ph) / sg\n", + " g_direct = float((A * (r - p)).mean())\n", + " worst = max(worst, abs(g_direct - float(sg)))\n", + "check('4. std-normalised group gradient = sigma_hat (direct sum vs closed form)', worst, 1e-15)\n", + "# Same with the *unbiased* sample std (torch.std, G-1 denominator) used by many implementations: g = sigma_hat * sqrt((G-1)/G)\n", + "worst_u = 0.0\n", + "for p in (0.1, 0.5):\n", + " for c in range(1, G):\n", + " r = torch.tensor([1.0] * c + [0.0] * (G - c)); ph = r.mean(); s_unb = r.std(unbiased=True)\n", + " g_direct = float((((r - ph) / s_unb) * (r - p)).mean())\n", + " worst_u = max(worst_u, abs(g_direct - float(torch.sqrt(ph * (1 - ph))) * math.sqrt((G - 1) / G)))\n", + "check('4. unbiased-std normalisation: group gradient = sigma_hat sqrt((G-1)/G)', worst_u, 1e-15)\n", + "print('p E[sigma_hat] sqrt(p(1-p)) weight ratio E[sigma_hat]/E[p_hat(1-p_hat)] (unnormalised = 1 by definition)')\n", + "for p in (0.02, 0.1, 0.25, 0.5, 0.9):\n", + " m_n, _ = moments(8, p, lambda ph: math.sqrt(ph * (1 - ph))); m_u, _ = moments(8, p, f_un)\n", + " print(f'{p:4.2f} {m_n:.5f} {math.sqrt(p*(1-p)):.5f} {m_n/m_u:8.3f}')\n", + "m_n05, _ = moments(8, 0.5, lambda ph: math.sqrt(ph * (1 - ph))); m_n02, _ = moments(8, 0.02, lambda ph: math.sqrt(ph * (1 - ph)))\n", + "m_u05, _ = moments(8, 0.5, f_un); m_u02, _ = moments(8, 0.02, f_un)\n", + "print('relative prompt weight p=0.02 vs p=0.5: unnormalised %.4f, std-normalised %.4f' % (m_u02 / m_u05, m_n02 / m_n05))\n", + "assert m_n02 / m_n05 > m_u02 / m_u05\n", + "# conditional on being informative, a group with c=1 of G=8 vs c=4 of 8: relative gradient magnitudes\n", + "sg = lambda c, G_=8: math.sqrt((c / G_) * (1 - c / G_)); un = lambda c, G_=8: (c / G_) * (1 - c / G_)\n", + "print('informative group c=1 vs c=4 (G=8): unnormalised %.3f, std-normalised %.3f, ratio of the two = 1/sigma_hat ratio = %.3f' % (un(1) / un(4), sg(1) / sg(4), (sg(1) / sg(4)) / (un(1) / un(4))))\n", + "check('4. std-normalisation re-weights c=1 vs c=4 groups by exactly sigma(4)/sigma(1)', abs((sg(1) / sg(4)) / (un(1) / un(4)) - sg(4) / sg(1)), 1e-14)" + ] + }, + { + "cell_type": "markdown", + "id": "bd6aa13f", + "metadata": {}, + "source": [ + "## 5. General bound: `||grad p||^2 <= p(1-p) E||s||^2` (Cauchy-Schwarz), and `grad p = 0` at `p in {0,1}`\n", + "`grad p = E[(r - p) s]` (baseline identity), so `||grad p|| <= sqrt(E(r-p)^2) sqrt(E||s||^2) = sqrt(p(1-p)) sqrt(E||s||^2)`.\n", + "Verified on 300 prompts of a log-linear policy family (K=16 outputs, d=6 features) with exact `grad p` from autograd." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "fbf95af9", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:50:18.618988Z", + "iopub.status.busy": "2026-10-14T09:50:18.618837Z", + "iopub.status.idle": "2026-10-14T09:50:18.687537Z", + "shell.execute_reply": "2026-10-14T09:50:18.686732Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "max ratio ||grad p||^2 / [p(1-p) E||s||^2] = 0.3109090719429734 (must be <= 1)\n", + "5. Cauchy-Schwarz bound ||grad p||^2 <= p(1-p) E||s||^2 (max violation of ratio <= 1) err=0.000e+00 (tol 0e+00)\n", + " p in [0,0.001): n= 28 mean ||grad p|| = 0.0000\n", + " p in [0.001,0.05): n= 24 mean ||grad p|| = 0.0791\n", + " p in [0.05,0.3): n=123 mean ||grad p|| = 0.2233\n", + " p in [0.3,0.7): n=112 mean ||grad p|| = 0.3436\n", + " p in [0.7,0.95): n= 13 mean ||grad p|| = 0.3216\n", + "p=0 prompt ||grad p|| = 0.0\n", + "p=1 prompt ||grad p|| = 0.0\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/home/vignesh/.opencode/tmp/ipykernel_607906/3044780012.py:10: UserWarning: Converting a tensor with requires_grad=True to a scalar may lead to unexpected behavior.\n", + "Consider using tensor.detach() first. (Triggered internally at /pytorch/torch/csrc/autograd/generated/python_variable_methods.cpp:838.)\n", + " Es2 = float((pi.detach() * (S_ ** 2).sum(1)).sum()); pf = float(p)\n" + ] + } + ], + "source": [ + "rg2 = np.random.RandomState(3); m_, K_, d_ = 300, 16, 6\n", + "Ph = torch.tensor(rg2.randn(m_, K_, d_)); Vv = torch.tensor((rg2.rand(m_, K_) < rg2.rand(m_, 1) * 0.6).astype(float))\n", + "th = torch.tensor(rg2.randn(d_) * 0.5, requires_grad=True)\n", + "ratios = []; grad_norms = []; pvals = []\n", + "for x in range(m_):\n", + " th.grad = None\n", + " pi = torch.softmax(Ph[x] @ th, -1); p = (pi * Vv[x]).sum(); p.backward()\n", + " gp = th.grad.clone()\n", + " S_ = Ph[x] - (pi.detach()[:, None] * Ph[x]).sum(0, keepdim=True) # score of each output y wrt theta: phi(x,y) - E phi\n", + " Es2 = float((pi.detach() * (S_ ** 2).sum(1)).sum()); pf = float(p)\n", + " bound = pf * (1 - pf) * Es2\n", + " ratios.append(float(gp @ gp) / bound if bound > 0 else 0.0); grad_norms.append(float(gp.norm())); pvals.append(pf)\n", + "print('max ratio ||grad p||^2 / [p(1-p) E||s||^2] =', max(ratios), ' (must be <= 1)')\n", + "check('5. Cauchy-Schwarz bound ||grad p||^2 <= p(1-p) E||s||^2 (max violation of ratio <= 1)', max(0.0, max(ratios) - 1.0), 0.0)\n", + "pv = np.array(pvals); gn = np.array(grad_norms)\n", + "for lo, hi in [(0, 0.001), (0.001, 0.05), (0.05, 0.3), (0.3, 0.7), (0.7, 0.95), (0.95, 1.0001)]:\n", + " sel = (pv >= lo) & (pv < hi)\n", + " if sel.sum(): print(f' p in [{lo},{hi}): n={int(sel.sum()):3d} mean ||grad p|| = {gn[sel].mean():.4f}')\n", + "Vz = Vv.clone(); Vz[0, :] = 0.0; Vo = Vv.clone(); Vo[1, :] = 1.0\n", + "for name_, Vt in (('p=0 prompt', Vz[0]), ('p=1 prompt', Vo[1])):\n", + " th.grad = None; pi = torch.softmax(Ph[0] @ th, -1); (pi * Vt).sum().backward()\n", + " print(name_, ' ||grad p|| =', float(th.grad.norm()))\n", + " assert float(th.grad.norm()) < 1e-15" + ] + }, + { + "cell_type": "markdown", + "id": "ca8a07b9", + "metadata": {}, + "source": [ + "## 6. Pass-rate band `(0, tau]` as deliberate re-weighting\n", + "Keep a group iff `0 < p_hat <= tau`. Its expected gradient contribution is `E[p_hat(1-p_hat) 1{0 < p_hat <= tau}]` versus `p(1-p)` for the true derivative.\n", + "The ratio `w_tau(p)` is the effective weight multiplier a prompt with true solve rate `p` receives." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "bb33c172", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:50:18.689048Z", + "iopub.status.busy": "2026-10-14T09:50:18.688891Z", + "iopub.status.idle": "2026-10-14T09:50:18.820659Z", + "shell.execute_reply": "2026-10-14T09:50:18.819967Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "6. band weight w_tau(p): closed form vs simulation (p=.7,G=8,tau=.75) err=3.628e-04 (tol 3e-03)\n", + "effective weight multiplier w_tau(p) (1 - 1/G = no band):\n", + " p G=8,tau=0.75 G=16,tau=0.75 G=8,tau=1 (DAPO: 0 0) & (ph <= tau))).mean() / (p * (1 - p))\n", + "check('6. band weight w_tau(p): closed form vs simulation (p=.7,G=8,tau=.75)', abs(sim - w_band(G_, p, tau)), 3e-3)\n", + "print('effective weight multiplier w_tau(p) (1 - 1/G = no band):')\n", + "print(' p G=8,tau=0.75 G=16,tau=0.75 G=8,tau=1 (DAPO: 012s}')\n", + "for n_, e_, t_ in CHECKS: print(f'{n_:92s} {e_:12.3e}')\n", + "print('CHECKS_JSON:' + json.dumps(CHECKS))" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/code/notebooks/0207-test-time-scaling-cost-model.ipynb b/code/notebooks/0207-test-time-scaling-cost-model.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..1ddc8dca786fdf5f9c53411a29dc54d67ccd4a17 --- /dev/null +++ b/code/notebooks/0207-test-time-scaling-cost-model.ipynb @@ -0,0 +1,602 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "b8fba5aa", + "metadata": {}, + "source": [ + "# 0207 -- Test-time scaling: pass@k vs compute, self-correction rounds vs more samples, best-of-n with a perfect verifier\n", + "\n", + "Identities verified (exact rationals / fp64):\n", + "1. Best-of-n with a perfect verifier: `coverage(n) = 1 - E[(1-p)^n]`; with `p ~ Beta(a,b)`: `1 - B(a,b+n)/B(a,b)`, and `1 - coverage ~ Gamma(a+b)/Gamma(b) n^-a`.\n", + "2. Early stopping: `E[min(N,n)] = (1-(1-p)^n)/p`; Beta-average `sum_{j= rho` (brute-force enumeration over `(n,R)`).\n", + "6. Imperfect verifier: precision cap `p/(p + (1-p) phi)` of the first accepted sample, independent of n.\n", + "7. The exponentiated power law `c = exp(a k^b)` (Large Language Monkeys) vs the Beta-mixture coverage curve." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "d1a9b6b8", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:51:01.379846Z", + "iopub.status.busy": "2026-10-14T09:51:01.379697Z", + "iopub.status.idle": "2026-10-14T09:51:01.728451Z", + "shell.execute_reply": "2026-10-14T09:51:01.727516Z" + } + }, + "outputs": [], + "source": [ + "import os\n", + "os.environ['CUDA_VISIBLE_DEVICES'] = ''\n", + "import math, json, itertools\n", + "from fractions import Fraction as Fr\n", + "from math import comb, lgamma\n", + "import numpy as np\n", + "from scipy.special import betaln, gammaln\n", + "from scipy.optimize import least_squares\n", + "CHECKS = []\n", + "def check(name, err, tol):\n", + " CHECKS.append((name, float(err), float(tol)))\n", + " assert err <= tol, (name, float(err), tol)\n", + " print(f'{name:92s} err={float(err):.3e} (tol {tol:.0e})')" + ] + }, + { + "cell_type": "markdown", + "id": "f09fd568", + "metadata": {}, + "source": [ + "## 1. Coverage with a perfect verifier\n", + "NAIVE: Monte Carlo over problems (draw `p`, draw `n` attempts, any success?). EFFICIENT: Beta-function closed form." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "57c4d2e9", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:51:01.730089Z", + "iopub.status.busy": "2026-10-14T09:51:01.729862Z", + "iopub.status.idle": "2026-10-14T09:51:03.100278Z", + "shell.execute_reply": "2026-10-14T09:51:03.099451Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "1. coverage: Monte Carlo vs 1 - B(a,b+n)/B(a,b) (max z, n in {1..256}) err=1.274e+00 (tol 4e+00)\n", + "(1-coverage) / [Gamma(a+b)/Gamma(b) n^-a] at n = [10, 100, 1000, 10000, 100000] -> [0.939, 0.9933, 0.9993, 0.9999, 1.0]\n", + "1. tail law: ratio -> 1 as n -> 1e5 (|ratio-1|) err=6.800e-06 (tol 5e-03)\n" + ] + } + ], + "source": [ + "def coverage_beta(a, b, n): return 1 - np.exp(betaln(a, b + n) - betaln(a, b))\n", + "rs = np.random.RandomState(0)\n", + "a, b = 0.4, 2.0\n", + "worst = 0.0\n", + "for n in (1, 4, 16, 64, 256):\n", + " p = rs.beta(a, b, size=2_000_000); hit = rs.rand(len(p)) < 1 - (1 - p) ** n\n", + " se = math.sqrt(hit.mean() * (1 - hit.mean()) / len(p))\n", + " worst = max(worst, abs(hit.mean() - coverage_beta(a, b, n)) / se)\n", + "check('1. coverage: Monte Carlo vs 1 - B(a,b+n)/B(a,b) (max z, n in {1..256})', worst, 4.0)\n", + "ns = np.array([10, 100, 1000, 10_000, 100_000])\n", + "ratio = (1 - coverage_beta(a, b, ns)) / (np.exp(gammaln(a + b) - gammaln(b)) * ns ** (-a))\n", + "print('(1-coverage) / [Gamma(a+b)/Gamma(b) n^-a] at n =', ns.tolist(), '->', np.round(ratio, 4).tolist())\n", + "check('1. tail law: ratio -> 1 as n -> 1e5 (|ratio-1|)', abs(ratio[-1] - 1), 5e-3)" + ] + }, + { + "cell_type": "markdown", + "id": "b908b5d2", + "metadata": {}, + "source": [ + "## 2. Early stopping with a perfect verifier\n", + "Sample sequentially, stop at the first verified success or after `n` attempts. Cost `N_n = min(Geom(p), n)`, `E[N_n] = sum_{j=0}^{n-1} (1-p)^j = (1-(1-p)^n)/p`." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "f4f195cb", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:51:03.101782Z", + "iopub.status.busy": "2026-10-14T09:51:03.101629Z", + "iopub.status.idle": "2026-10-14T09:51:03.424265Z", + "shell.execute_reply": "2026-10-14T09:51:03.423468Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "2. E[min(Geom(p),n)] = (1-(1-p)^n)/p = sum_{j= c`, i.e. `j < a U/c - (a+b)`." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "087dfa00", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:51:03.425743Z", + "iopub.status.busy": "2026-10-14T09:51:03.425591Z", + "iopub.status.idle": "2026-10-14T09:51:05.083274Z", + "shell.execute_reply": "2026-10-14T09:51:05.082473Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "3. posterior success probability after j failures = a/(a+b+j) (importance-sampling MC, max z) err=1.647e+00 (tol 4e+00)\n", + "3. same, exact Beta-function ratio B(a+1,b+j)/B(a,b+j) err=8.491e-16 (tol 1e-14)\n", + "U=10.0, c=0.05: DP stops after 78 attempts; threshold a U/c - (a+b) = 77.60 -> ceil 78; value W_0=6.5345\n", + "U=10.0, c=0.01: DP stops after 398 attempts; threshold a U/c - (a+b) = 397.60 -> ceil 398; value W_0=8.1394\n", + "U=50.0, c=0.05: DP stops after 398 attempts; threshold a U/c - (a+b) = 397.60 -> ceil 398; value W_0=40.6968\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "3. DP stopping index vs threshold rule ceil(aU/c-(a+b)) (max |diff| over 3 settings) err=0.000e+00 (tol 1e+00)\n" + ] + } + ], + "source": [ + "a, b = 0.4, 2.0\n", + "# hazard check by self-normalised importance sampling from the prior (weights (1-p)^j = likelihood of j failures)\n", + "rs = np.random.RandomState(7); pp = rs.beta(a, b, size=8_000_000); worst_z = 0.0\n", + "for j in (0, 1, 5, 20):\n", + " w = (1 - pp) ** j; est = (pp * w).sum() / w.sum()\n", + " # delta-method s.e. of a ratio estimator\n", + " se = math.sqrt((((pp - est) * w) ** 2).sum()) / w.sum()\n", + " worst_z = max(worst_z, abs(est - a / (a + b + j)) / se)\n", + "check('3. posterior success probability after j failures = a/(a+b+j) (importance-sampling MC, max z)', worst_z, 4.5)\n", + "# exact version via Beta functions (no quadrature):\n", + "worst = max(abs(math.exp(betaln(a + 1, b + j) - betaln(a, b + j)) - a / (a + b + j)) for j in range(0, 200))\n", + "check('3. same, exact Beta-function ratio B(a+1,b+j)/B(a,b+j)', worst, 1e-14)\n", + "def dp_value(U, c, a, b, nmax=100000):\n", + " W = 0.0; cont = []\n", + " for j in range(nmax - 1, -1, -1):\n", + " h = a / (a + b + j)\n", + " val = h * (U - c) - (1 - h) * c + (1 - h) * W\n", + " W = max(0.0, val); cont.append(val > 0)\n", + " cont = cont[::-1]\n", + " jstop = next((j for j, v in enumerate(cont) if not v), nmax)\n", + " return W, jstop\n", + "for (U, c) in [(10.0, 0.05), (10.0, 0.01), (50.0, 0.05)]:\n", + " W, jstop = dp_value(U, c, a, b)\n", + " j_theory = math.ceil(a * U / c - (a + b))\n", + " print(f'U={U}, c={c}: DP stops after {jstop} attempts; threshold a U/c - (a+b) = {a*U/c-(a+b):.2f} -> ceil {j_theory}; value W_0={W:.4f}')\n", + " assert abs(jstop - j_theory) <= 1\n", + "check('3. DP stopping index vs threshold rule ceil(aU/c-(a+b)) (max |diff| over 3 settings)', max(abs(dp_value(U, c, a, b)[1] - math.ceil(a * U / c - (a + b))) for U, c in [(10.0, 0.05), (10.0, 0.01), (50.0, 0.05)]), 1)" + ] + }, + { + "cell_type": "markdown", + "id": "6e820e32", + "metadata": {}, + "source": [ + "## 4. Equal expected cost: adaptive stopping vs a fixed budget\n", + "Fixed `n = 64` attempts for every problem costs 64 each. With early stopping the expected cost is `cost(n_max)`; choose `n_max` so that it equals 64 and compare coverage." + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "bddd003b", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:51:05.085063Z", + "iopub.status.busy": "2026-10-14T09:51:05.084733Z", + "iopub.status.idle": "2026-10-14T09:51:19.832181Z", + "shell.execute_reply": "2026-10-14T09:51:19.831457Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Beta(0.4,2.0): fixed n=64 -> coverage 0.7671; adaptive with the same mean cost 64 needs n_max = 323 -> coverage 0.8771\n", + "4. n_max is the smallest n with expected cost >= 64 (violations; 0 = cost(n_max-1) < 64 <= cost(n_max)) err=0.000e+00 (tol 0e+00)\n", + "cost at n_max-1, n_max: 63.997038900831 64.1201008159222\n" + ] + } + ], + "source": [ + "target = 64.0\n", + "# cumulative expected cost cost(n) = sum_{j coverage {coverage_beta(a,b,64):.4f}; adaptive with the same mean cost 64 needs n_max = {n_adapt} -> coverage {coverage_beta(a,b,n_adapt):.4f}')\n", + "check('4. n_max is the smallest n with expected cost >= 64 (violations; 0 = cost(n_max-1) < 64 <= cost(n_max))', float(not (cs[n_adapt - 2] < target <= cs[n_adapt - 1])), 0.0)\n", + "print('cost at n_max-1, n_max:', cs[n_adapt - 2], cs[n_adapt - 1])\n", + "assert coverage_beta(a, b, n_adapt) > coverage_beta(a, b, 64)\n", + "CHECKS.append(('4. adaptive-minus-fixed coverage at equal mean cost (informational)', float(coverage_beta(a, b, n_adapt) - coverage_beta(a, b, 64)), float('inf')))" + ] + }, + { + "cell_type": "markdown", + "id": "867deb08", + "metadata": {}, + "source": [ + "## 5. Fresh samples vs revisions at equal token cost\n", + "A fresh sample costs 1 unit and succeeds w.p. `p1`. A chain is one attempt followed by up to `R` revisions, each costing `rho` units and succeeding w.p. `q` given all earlier failures\n", + "(independent of the earlier attempts: the simplest model). Chain failure `(1-p1)(1-q)^R`, cost `1 + rho R`. With budget `B` we run `n` independent chains.\n", + "Failure `= [(1-p1)(1-q)^R]^n`. Log-failure per cost: `f(R) = (a1 + R a2)/(1+rho R)`, `a1 = -ln(1-p1)`, `a2 = -ln(1-q)`. `f' ∝ a2 - rho a1`: revise iff `a2 > rho a1`." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "d391bbf9", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:51:19.833654Z", + "iopub.status.busy": "2026-10-14T09:51:19.833497Z", + "iopub.status.idle": "2026-10-14T09:51:19.840774Z", + "shell.execute_reply": "2026-10-14T09:51:19.840044Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "(p1, q, rho, revise predicted, best n, best R, failure prob. of best mix at budget 600, failure prob. with fresh samples only)\n", + " (0.05, 0.03, 1.5, False, 600, 0, 4.31e-14, 4.31e-14)\n", + " (0.05, 0.06, 1.0, True, 50, 11, 1.28e-16, 4.31e-14)\n", + " (0.05, 0.06, 1.5, False, 600, 0, 4.31e-14, 4.31e-14)\n", + " (0.1, 0.12, 1.0, True, 50, 11, 1.51e-33, 3.51e-28)\n", + "sign rule vs brute-force optimum: agree on 36/36 grid points\n", + "disagreements (p1, q, rho, relative gap |a2 - rho a1| / (rho a1), best (n,R)): []\n", + "5. revise-iff-(a2 > rho a1) rule vs brute-force optimum: disagreements with relative gap > 10% err=0.000e+00 (tol 0e+00)\n" + ] + } + ], + "source": [ + "def best_mix(p1, q, rho, B, Rmax=12):\n", + " \"\"\"Minimise the log failure probability n*[ln(1-p1) + R ln(1-q)] over integer (n, R) with n(1+rho R) <= B (avoids float saturation at success ~ 1).\"\"\"\n", + " best = (float('inf'), None)\n", + " for R in range(0, Rmax + 1):\n", + " n = int(B // (1 + rho * R))\n", + " if n < 1: continue\n", + " logfail = n * (math.log(1 - p1) + R * math.log(1 - q))\n", + " if logfail < best[0] - 1e-12: best = (logfail, (n, R))\n", + " return math.exp(best[0]), best[1] # failure probability of the best mix\n", + "rows = []; borderline = []\n", + "agree = 0; total = 0\n", + "for p1 in (0.02, 0.05, 0.1):\n", + " for q in (0.01, 0.03, 0.06, 0.12):\n", + " for rho in (1.0, 1.5, 2.5):\n", + " B = 600.0\n", + " succ, (n_, R_) = best_mix(p1, q, rho, B) # `succ` is the failure probability here\n", + " predicted_revise = (-math.log(1 - q)) > rho * (-math.log(1 - p1))\n", + " gap = abs((-math.log(1 - q)) - rho * (-math.log(1 - p1))) / (rho * (-math.log(1 - p1)))\n", + " total += 1\n", + " if (R_ > 0) == predicted_revise: agree += 1\n", + " else:\n", + " borderline.append((p1, q, rho, round(gap, 3), (n_, R_)))\n", + " if (p1, q, rho) in [(0.05, 0.06, 1.0), (0.05, 0.06, 1.5), (0.05, 0.03, 1.5), (0.1, 0.12, 1.0)]:\n", + " rows.append((p1, q, rho, predicted_revise, n_, R_, float('%.3g' % succ), float('%.3g' % ((1 - p1) ** int(B)))))\n", + "print('(p1, q, rho, revise predicted, best n, best R, failure prob. of best mix at budget 600, failure prob. with fresh samples only)')\n", + "for r in rows: print(' ', r)\n", + "print(f'sign rule vs brute-force optimum: agree on {agree}/{total} grid points')\n", + "print('disagreements (p1, q, rho, relative gap |a2 - rho a1| / (rho a1), best (n,R)):', borderline)\n", + "far = [b_ for b_ in borderline if b_[3] > 0.1]\n", + "check('5. revise-iff-(a2 > rho a1) rule vs brute-force optimum: disagreements with relative gap > 10%', len(far), 0)" + ] + }, + { + "cell_type": "markdown", + "id": "b061849f", + "metadata": {}, + "source": [ + "## 6. Imperfect verifier: the first *accepted* sample is correct only with probability `p/(p+(1-p) phi)`\n", + "Verifier accepts correct proofs always and incorrect ones w.p. `phi` (false accept). Sampling until the first accept, the accepted sample is correct w.p. `p/(p + (1-p) phi)`,\n", + "for any budget `n` (precision cap), while a perfect verifier gives 1 with the same stopping rule." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "d216a73d", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:51:19.842276Z", + "iopub.status.busy": "2026-10-14T09:51:19.842130Z", + "iopub.status.idle": "2026-10-14T09:51:24.279548Z", + "shell.execute_reply": "2026-10-14T09:51:24.278840Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "p=0.05, false accept=0.01: precision of an accepted sample = 0.8414 (theory 0.8403)\n", + "p=0.2, false accept=0.05: precision of an accepted sample = 0.8335 (theory 0.8333)\n", + "p=0.01, false accept=0.002: precision of an accepted sample = 0.8315 (theory 0.8347)\n", + "6. precision of the first accepted sample: MC vs p/(p+(1-p)phi) (max z) err=1.883e+00 (tol 4e+00)\n" + ] + } + ], + "source": [ + "rs = np.random.RandomState(4)\n", + "worst = 0\n", + "for (p, phi) in [(0.05, 0.01), (0.2, 0.05), (0.01, 0.002)]:\n", + " N = 4_000_000\n", + " corr = rs.rand(N) < p; acc = corr | (rs.rand(N) < phi)\n", + " precision = corr[acc].mean(); se = math.sqrt(precision * (1 - precision) / acc.sum())\n", + " worst = max(worst, abs(precision - p / (p + (1 - p) * phi)) / se)\n", + " print(f'p={p}, false accept={phi}: precision of an accepted sample = {precision:.4f} (theory {p/(p+(1-p)*phi):.4f})')\n", + "check('6. precision of the first accepted sample: MC vs p/(p+(1-p)phi) (max z)', worst, 4.0)" + ] + }, + { + "cell_type": "markdown", + "id": "7215e58c", + "metadata": {}, + "source": [ + "## 7. The empirical exponentiated power law vs the Beta-mixture curve\n", + "Large Language Monkeys fit coverage `c(k) = exp(a k^b)` (Fig. 5). For `p ~ Beta(alpha, beta)` the exact curve is `1 - B(alpha, beta+k)/B(alpha,beta)`. Fit `log c` by least squares on log-spaced `k` in `[1, 10^4]`." + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "32ba1c14", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:51:24.280954Z", + "iopub.status.busy": "2026-10-14T09:51:24.280801Z", + "iopub.status.idle": "2026-10-14T09:51:24.302130Z", + "shell.execute_reply": "2026-10-14T09:51:24.301470Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Beta(0.3,1.5): fitted c=exp(a k^b): a=-1.648, b=-0.361, mean abs error=0.0061, max=0.0259\n", + "Beta(0.5,2.0): fitted c=exp(a k^b): a=-1.673, b=-0.535, mean abs error=0.0014, max=0.0123\n", + "Beta(1.0,3.0): fitted c=exp(a k^b): a=-1.625, b=-0.809, mean abs error=0.0062, max=0.0530\n" + ] + } + ], + "source": [ + "ks = np.unique(np.logspace(0, 4, 100).astype(int))\n", + "for (al, be) in [(0.3, 1.5), (0.5, 2.0), (1.0, 3.0)]:\n", + " cov = coverage_beta(al, be, ks)\n", + " fit = least_squares(lambda th: np.exp(th[0] * ks ** th[1]) - cov, x0=[-1.0, -0.3], bounds=([-20, -3], [-1e-6, -1e-6]))\n", + " print(f'Beta({al},{be}): fitted c=exp(a k^b): a={fit.x[0]:.3f}, b={fit.x[1]:.3f}, mean abs error={np.abs(fit.fun).mean():.4f}, max={np.abs(fit.fun).max():.4f}')\n", + "CHECKS.append(('7. exp-power-law fit error on a Beta-mixture curve (max abs over 3 fits, informational)', float(max(np.abs(least_squares(lambda th, cv=coverage_beta(al, be, ks): np.exp(th[0] * ks ** th[1]) - cv, x0=[-1.0, -0.3], bounds=([-20, -3], [-1e-6, -1e-6])).fun).max() for al, be in [(0.3, 1.5), (0.5, 2.0), (1.0, 3.0)])), float('inf')))" + ] + }, + { + "cell_type": "markdown", + "id": "f9f47176", + "metadata": {}, + "source": [ + "## 8. Token accounting for a published comparison (arithmetic only)\n", + "Goedel-Prover-V2 Sec. 3.2: first round capped at 30,000 tokens, self-correction mode (initial + 2 revisions) at 40,000 tokens in total.\n", + "Sec. 3.5: with the context extended to 128k (YaRN) and up to 5 revisions, pass@32 with correction reaches 92.7% > 92.2% (32B, pass@8192, no correction).\n", + "Table 4: the 40k-token self-correction mode reaches only 90.4% at pass@32. Ratio of caps (neither an upper nor a lower bound on tokens actually used):" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "cbf4ed2e", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:51:24.303439Z", + "iopub.status.busy": "2026-10-14T09:51:24.303292Z", + "iopub.status.idle": "2026-10-14T09:51:24.307308Z", + "shell.execute_reply": "2026-10-14T09:51:24.306689Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "cap ratio, 40k self-correction mode (reaches 90.4% < 92.2%): 8192 x 30k / (32 x 40k) = 192x\n", + "cap ratio, 128k extended mode (reaches 92.7% > 92.2%) : 8192 x 30k / (32 x 128k) = 60x\n", + "rho implied by the caps (chain = attempt + 2 revisions <= 40k vs fresh <= 30k): 0.16666666666666663\n", + "8. rho implied by 1 + 2 rho = 40k/30k equals 1/6 err=2.776e-17 (tol 1e-15)\n" + ] + } + ], + "source": [ + "tok_std, tok_sc40, tok_sc128 = 30_000, 40_000, 128_000\n", + "ratio40 = (8192 * tok_std) / (32 * tok_sc40); ratio128 = (8192 * tok_std) / (32 * tok_sc128)\n", + "print(f'cap ratio, 40k self-correction mode (reaches 90.4% < 92.2%): 8192 x 30k / (32 x 40k) = {ratio40:.0f}x')\n", + "print(f'cap ratio, 128k extended mode (reaches 92.7% > 92.2%) : 8192 x 30k / (32 x 128k) = {ratio128:.0f}x')\n", + "assert ratio40 == 192 and ratio128 == 60\n", + "# chain-vs-fresh cost at the caps: 1 + 2 rho <= 40/30 => rho <= 1/6 if the caps were costs\n", + "rho_cap = (tok_sc40 / tok_std - 1) / 2\n", + "print('rho implied by the caps (chain = attempt + 2 revisions <= 40k vs fresh <= 30k):', rho_cap)\n", + "check('8. rho implied by 1 + 2 rho = 40k/30k equals 1/6', abs(rho_cap - 1 / 6), 1e-15)\n", + "CHECKS.append(('8. cap-based token ratio for the 92.7% (128k, <=5 revisions) configuration (informational)', float(ratio128), float('inf')))\n", + "CHECKS.append(('8. cap-based token ratio for the 40k self-correction mode (informational; that mode does not beat 92.2%)', float(ratio40), float('inf')))" + ] + }, + { + "cell_type": "markdown", + "id": "d23da596", + "metadata": {}, + "source": [ + "## 8b. Sensitivity of the brute-force optimum to the chain-length cap\n", + "With `q` independent of history the continuous objective is monotone in `R`; the integer optimum depends on the cap `R_max` (and on `floor(B/(1+rho R))`)." + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "56d4dfdb", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:51:24.308569Z", + "iopub.status.busy": "2026-10-14T09:51:24.308427Z", + "iopub.status.idle": "2026-10-14T09:51:24.313929Z", + "shell.execute_reply": "2026-10-14T09:51:24.313186Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "p1=0.05 q=0.06 rho=1 B=600, R_max= 12: best (n,R)=(50,11), log-failure = -36.596\n", + "p1=0.05 q=0.06 rho=1 B=600, R_max= 50: best (n,R)=(12,49), log-failure = -36.998\n", + "p1=0.05 q=0.06 rho=1 B=600, R_max=599: best (n,R)=(1,599), log-failure = -37.115\n" + ] + } + ], + "source": [ + "def best_mix_cap(p1, q, rho, B, Rmax):\n", + " best = (float('inf'), None)\n", + " for R in range(0, Rmax + 1):\n", + " n = int(B // (1 + rho * R))\n", + " if n < 1: continue\n", + " lf = n * (math.log(1 - p1) + R * math.log(1 - q))\n", + " if lf < best[0] - 1e-12: best = (lf, (n, R))\n", + " return best\n", + "for cap in (12, 50, 599):\n", + " lf, (n_, R_) = best_mix_cap(0.05, 0.06, 1.0, 600.0, cap)\n", + " print(f'p1=0.05 q=0.06 rho=1 B=600, R_max={cap:3d}: best (n,R)=({n_},{R_}), log-failure = {lf:.3f}')\n", + "lf12 = best_mix_cap(0.05, 0.06, 1.0, 600.0, 12); lf599 = best_mix_cap(0.05, 0.06, 1.0, 600.0, 599)\n", + "assert lf12[1] == (50, 11) and lf599[1][0] <= 2 and lf599[0] < lf12[0]\n", + "CHECKS.append(('8b. optimum for R_max=12 is (50,11); uncapped optimum has n<=2 chains (informational)', float(lf599[1][0]), float('inf')))" + ] + }, + { + "cell_type": "markdown", + "id": "b5ac6a6d", + "metadata": {}, + "source": [ + "## Verification table" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "id": "433dc5ca", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:51:24.315222Z", + "iopub.status.busy": "2026-10-14T09:51:24.315080Z", + "iopub.status.idle": "2026-10-14T09:51:24.318768Z", + "shell.execute_reply": "2026-10-14T09:51:24.317925Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + "identity max abs err\n", + "1. coverage: Monte Carlo vs 1 - B(a,b+n)/B(a,b) (max z, n in {1..256}) 1.274e+00\n", + "1. tail law: ratio -> 1 as n -> 1e5 (|ratio-1|) 6.800e-06\n", + "2. E[min(Geom(p),n)] = (1-(1-p)^n)/p = sum_{j= 64 (violations; 0 = cost(n_max-1) < 64 <= cost(n_max)) 0.000e+00\n", + "4. adaptive-minus-fixed coverage at equal mean cost (informational) 1.100e-01\n", + "5. revise-iff-(a2 > rho a1) rule vs brute-force optimum: disagreements with relative gap > 10% 0.000e+00\n", + "6. precision of the first accepted sample: MC vs p/(p+(1-p)phi) (max z) 1.883e+00\n", + "7. exp-power-law fit error on a Beta-mixture curve (max abs over 3 fits, informational) 5.301e-02\n", + "8. rho implied by 1 + 2 rho = 40k/30k equals 1/6 2.776e-17\n", + "8. cap-based token ratio for the 92.7% (128k, <=5 revisions) configuration (informational) 6.000e+01\n", + "8. cap-based token ratio for the 40k self-correction mode (informational; that mode does not beat 92.2%) 1.920e+02\n", + "8b. optimum for R_max=12 is (50,11); uncapped optimum has n<=2 chains (informational) 1.000e+00\n", + "CHECKS_JSON:[[\"1. coverage: Monte Carlo vs 1 - B(a,b+n)/B(a,b) (max z, n in {1..256})\", 1.2744892726318469, 4.0], [\"1. tail law: ratio -> 1 as n -> 1e5 (|ratio-1|)\", 6.8000623573150065e-06, 0.005], [\"2. E[min(Geom(p),n)] = (1-(1-p)^n)/p = sum_{j= 64 (violations; 0 = cost(n_max-1) < 64 <= cost(n_max))\", 0.0, 0.0], [\"4. adaptive-minus-fixed coverage at equal mean cost (informational)\", 0.10998111815038014, Infinity], [\"5. revise-iff-(a2 > rho a1) rule vs brute-force optimum: disagreements with relative gap > 10%\", 0.0, 0.0], [\"6. precision of the first accepted sample: MC vs p/(p+(1-p)phi) (max z)\", 1.88315480662085, 4.0], [\"7. exp-power-law fit error on a Beta-mixture curve (max abs over 3 fits, informational)\", 0.05301359850390941, Infinity], [\"8. rho implied by 1 + 2 rho = 40k/30k equals 1/6\", 2.7755575615628914e-17, 1e-15], [\"8. cap-based token ratio for the 92.7% (128k, <=5 revisions) configuration (informational)\", 60.0, Infinity], [\"8. cap-based token ratio for the 40k self-correction mode (informational; that mode does not beat 92.2%)\", 192.0, Infinity], [\"8b. optimum for R_max=12 is (50,11); uncapped optimum has n<=2 chains (informational)\", 1.0, Infinity]]\n" + ] + } + ], + "source": [ + "print()\n", + "print(f'{\"identity\":94s} {\"max abs err\":>12s}')\n", + "for n_, e_, t_ in CHECKS: print(f'{n_:94s} {e_:12.3e}')\n", + "print('CHECKS_JSON:' + json.dumps(CHECKS))" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/code/notebooks/0208-statistics-of-evaluation.ipynb b/code/notebooks/0208-statistics-of-evaluation.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..eef0e8da7f581f20c411e681e1e32ab095274b80 --- /dev/null +++ b/code/notebooks/0208-statistics-of-evaluation.ipynb @@ -0,0 +1,687 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "2dc561da", + "metadata": {}, + "source": [ + "# 0208 -- Statistics of evaluation: exact McNemar, Wilson intervals, cluster bootstrap, power for a +5/244 effect\n", + "\n", + "Benchmark of `N = 244` problems (the size of miniF2F-test). Identities / claims verified:\n", + "1. Exact McNemar two-sided p-value `min(1, 2 P(Bin(n_d,1/2) <= min(b,c)))`: exact rationals = scipy = sign-flip permutation enumeration; **no outcome with `b - c = 5` is significant at 0.05 (best case (5,0): p = 0.0625)**.\n", + "2. Wilson interval: closed form = numerical inversion of the score test; exact coverage vs Wald and Clopper-Pearson.\n", + "3. Variance decomposition `Var = [Var_x(p) + E p(1-p)/m]/N` and the design effect; cluster-bootstrap vs naive coverage.\n", + "4. Paired-difference variance `(Var A + Var B - 2 Cov)/N`.\n", + "5. Exact power of McNemar for `delta = 5/244`; normal approximation; MDE; required N.\n", + "6. Size of Wald and bootstrap tests with few discordant pairs (both reject (5,0); exact does not)." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "c6f8a060", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:32:01.357637Z", + "iopub.status.busy": "2026-10-14T09:32:01.357418Z", + "iopub.status.idle": "2026-10-14T09:32:02.129875Z", + "shell.execute_reply": "2026-10-14T09:32:02.128835Z" + } + }, + "outputs": [], + "source": [ + "import os\n", + "os.environ['CUDA_VISIBLE_DEVICES'] = ''\n", + "import math, json, itertools\n", + "from fractions import Fraction as Fr\n", + "from math import comb\n", + "import numpy as np\n", + "from scipy import stats\n", + "from scipy.optimize import brentq\n", + "from scipy.special import betaln\n", + "CHECKS = []\n", + "def check(name, err, tol):\n", + " CHECKS.append((name, float(err), float(tol)))\n", + " assert err <= tol, (name, float(err), tol)\n", + " print(f'{name:96s} err={float(err):.3e} (tol {tol:.0e})')\n", + "N = 244" + ] + }, + { + "cell_type": "markdown", + "id": "8c953229", + "metadata": {}, + "source": [ + "## 1. Exact McNemar\n", + "Paired binary outcomes per problem: `b` = A solves & B fails, `c` = B solves & A fails. Concordant problems carry no information about the difference\n", + "under `H0: P(A=1,B=0) = P(A=0,B=1)`. Conditional on `n_d = b + c`, `b ~ Bin(n_d, 1/2)`.\n", + "NAIVE: exact rational binomial tail; independent check 1: scipy; independent check 2: enumerate all `2^{n_d}` sign assignments (permutation test)." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "7e8399d5", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:32:02.131822Z", + "iopub.status.busy": "2026-10-14T09:32:02.131499Z", + "iopub.status.idle": "2026-10-14T09:32:03.388935Z", + "shell.execute_reply": "2026-10-14T09:32:03.387984Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "1. exact McNemar p: rational tail vs scipy.binomtest (all b,c <= 16) err=1.110e-16 (tol 1e-12)\n", + "1. exact McNemar p: rational tail vs sign-flip permutation enumeration (n_d <= 14) err=0.000e+00 (tol 1e-15)\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "net gain d=b-c : the discordant counts n_d for which some split is significant (min..max n_d, count)\n", + " d= 3: never significant\n", + " d= 4: never significant\n", + " d= 5: never significant\n", + " d= 6: n_d in [6, 6] (1 values)\n", + " d= 7: n_d in [7, 9] (2 values)\n", + " d= 8: n_d in [8, 12] (3 values)\n", + " d= 9: n_d in [9, 17] (5 values)\n", + " d=10: n_d in [10, 20] (6 values)\n", + " d=11: n_d in [11, 25] (8 values)\n", + " d=12: n_d in [12, 30] (10 values)\n", + "p(5,0) = 0.0625 p(6,0) = 0.03125 p(7,1) = 0.0703125 p(8,3) = 0.2265625\n", + "1. best possible p for a +5 net gain is (5,0): p = 1/16 > 0.05 (|p - 0.0625|) err=0.000e+00 (tol 0e+00)\n", + "typical +5/244 outcomes (b,c): {(5, 0): 0.062, (8, 3): 0.227, (10, 5): 0.302, (15, 10): 0.424, (20, 15): 0.5, (30, 25): 0.59}\n" + ] + } + ], + "source": [ + "def mcnemar_exact_fraction(b, c):\n", + " n = b + c\n", + " if n == 0: return Fr(1)\n", + " k = min(b, c)\n", + " return min(Fr(1), 2 * Fr(sum(comb(n, i) for i in range(k + 1)), 2 ** n))\n", + "worst = 0.0; worst_perm = 0.0\n", + "for b in range(0, 17):\n", + " for c in range(0, 17):\n", + " n = b + c\n", + " if n == 0: continue\n", + " pf = float(mcnemar_exact_fraction(b, c))\n", + " ps = stats.binomtest(min(b, c), n, 0.5, alternative='two-sided').pvalue\n", + " worst = max(worst, abs(pf - ps))\n", + " if n <= 14:\n", + " signs = np.array(list(itertools.product([-1, 1], repeat=n)))\n", + " cnt = (np.abs(signs.sum(1)) >= abs(b - c)).sum() / len(signs)\n", + " worst_perm = max(worst_perm, abs(pf - cnt))\n", + "check('1. exact McNemar p: rational tail vs scipy.binomtest (all b,c <= 16)', worst, 1e-12)\n", + "check('1. exact McNemar p: rational tail vs sign-flip permutation enumeration (n_d <= 14)', worst_perm, 1e-15)\n", + "# Which net improvements d = b - c can ever be significant at 0.05 (two-sided)?\n", + "best = {}\n", + "for n in range(1, N + 1):\n", + " for b in range(n + 1):\n", + " c = n - b\n", + " if b - c <= 0: continue\n", + " if float(mcnemar_exact_fraction(b, c)) < 0.05:\n", + " best.setdefault(b - c, []).append(n)\n", + "print('net gain d=b-c : the discordant counts n_d for which some split is significant (min..max n_d, count)')\n", + "for d in range(3, 13):\n", + " ns_ = best.get(d, [])\n", + " print(f' d={d:2d}: ' + (f'n_d in [{min(ns_)}, {max(ns_)}] ({len(ns_)} values)' if ns_ else 'never significant'))\n", + "for d in (5, 6):\n", + " pass\n", + "assert 5 not in best and 4 not in best and 6 in best and min(best[6]) == max(best[6]) == 6\n", + "print('p(5,0) =', float(mcnemar_exact_fraction(5, 0)), ' p(6,0) =', float(mcnemar_exact_fraction(6, 0)), ' p(7,1) =', float(mcnemar_exact_fraction(7, 1)), ' p(8,3) =', float(mcnemar_exact_fraction(8, 3)))\n", + "check('1. best possible p for a +5 net gain is (5,0): p = 1/16 > 0.05 (|p - 0.0625|)', abs(float(mcnemar_exact_fraction(5, 0)) - 0.0625), 0.0)\n", + "print('typical +5/244 outcomes (b,c): ', {(b, b - 5): round(float(mcnemar_exact_fraction(b, b - 5)), 3) for b in (5, 8, 10, 15, 20, 30)})" + ] + }, + { + "cell_type": "markdown", + "id": "d98651ac", + "metadata": {}, + "source": [ + "## 2. Wilson interval\n", + "Invert the score test `|p_hat - p| <= z sqrt(p(1-p)/n)`. Squaring: `(1 + z^2/n) p^2 - (2 p_hat + z^2/n) p + p_hat^2 <= 0`, a quadratic in `p` whose roots are\n", + "`p = [p_hat + z^2/(2n) +- z sqrt(p_hat(1-p_hat)/n + z^2/(4n^2))] / (1 + z^2/n)`.\n", + "NAIVE: root-find the score-test equation numerically. EFFICIENT: closed form. Then exact coverage over a grid of true `p`." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "556d82e0", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:32:03.390688Z", + "iopub.status.busy": "2026-10-14T09:32:03.390516Z", + "iopub.status.idle": "2026-10-14T09:32:03.531057Z", + "shell.execute_reply": "2026-10-14T09:32:03.530078Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "2. Wilson closed form vs numerical inversion of the score test (max |endpoint diff|, x=1..243) err=4.905e-13 (tol 1e-10)\n", + " 203/244 = 0.8320 Wilson 95% = [0.7800, 0.8737] (half-width 0.0468 = 11.4 problems)\n", + " 208/244 = 0.8525 Wilson 95% = [0.8025, 0.8915] (half-width 0.0445 = 10.9 problems)\n", + " 213/244 = 0.8730 Wilson 95% = [0.8253, 0.9090] (half-width 0.0419 = 10.2 problems)\n", + " coverage of nominal-95% wilson: mean 0.9507 min 0.9402 at p=0.83: 0.9495 at p=0.02: 0.9664\n", + " coverage of nominal-95% wald : mean 0.9436 min 0.8634 at p=0.83: 0.9437 at p=0.02: 0.8634\n", + " coverage of nominal-95% cp : mean 0.9595 min 0.9507 at p=0.83: 0.9597 at p=0.02: 0.9819\n" + ] + } + ], + "source": [ + "z = stats.norm.ppf(0.975)\n", + "def wilson(x, n, z=z):\n", + " ph = x / n; den = 1 + z * z / n\n", + " ctr = (ph + z * z / (2 * n)) / den; half = z / den * math.sqrt(ph * (1 - ph) / n + z * z / (4 * n * n))\n", + " return ctr - half, ctr + half\n", + "def wilson_numeric(x, n, z=z):\n", + " ph = x / n\n", + " f = lambda p: (ph - p) ** 2 - z * z * p * (1 - p) / n\n", + " lo = brentq(f, 1e-12, ph) if x > 0 else 0.0\n", + " hi = brentq(f, ph, 1 - 1e-12) if x < n else 1.0\n", + " return lo, hi\n", + "worst = max(max(abs(a - b_) for a, b_ in zip(wilson(x, N), wilson_numeric(x, N))) for x in range(1, N))\n", + "check('2. Wilson closed form vs numerical inversion of the score test (max |endpoint diff|, x=1..243)', worst, 1e-10)\n", + "for x in (203, 208, 213):\n", + " lo, hi = wilson(x, N); print(f' {x}/{N} = {x/N:.4f} Wilson 95% = [{lo:.4f}, {hi:.4f}] (half-width {(hi-lo)/2:.4f} = {(hi-lo)/2*N:.1f} problems)')\n", + "def wald(x, n):\n", + " ph = x / n; h = z * math.sqrt(ph * (1 - ph) / n); return ph - h, ph + h\n", + "def cp(x, n):\n", + " lo = 0.0 if x == 0 else stats.beta.ppf(0.025, x, n - x + 1); hi = 1.0 if x == n else stats.beta.ppf(0.975, x + 1, n - x)\n", + " return lo, hi\n", + "xs = np.arange(N + 1); ci_cache = {nm: np.array([f(x, N) for x in xs]) for nm, f in (('wilson', wilson), ('wald', wald), ('cp', cp))}\n", + "def coverage(nm, p):\n", + " pm = stats.binom.pmf(xs, N, p); lo, hi = ci_cache[nm][:, 0], ci_cache[nm][:, 1]\n", + " return float((pm * ((lo <= p) & (p <= hi))).sum())\n", + "pgrid = np.arange(0.01, 0.995, 0.005)\n", + "cov = {nm: np.array([coverage(nm, p) for p in pgrid]) for nm in ci_cache}\n", + "for nm in cov: print(f' coverage of nominal-95% {nm:6s}: mean {cov[nm].mean():.4f} min {cov[nm].min():.4f} at p=0.83: {coverage(nm, 0.83):.4f} at p=0.02: {coverage(nm, 0.02):.4f}')\n", + "assert cov['cp'].min() >= 0.95 - 1e-9 and cov['wald'].min() < 0.93 and abs(cov['wilson'].mean() - 0.95) < 0.01\n", + "CHECKS.append(('2. min exact coverage over p in [0.01,0.995]: Wilson / Wald / Clopper-Pearson = %.4f / %.4f / %.4f (informational)' % (cov['wilson'].min(), cov['wald'].min(), cov['cp'].min()), float(cov['wilson'].min()), float('inf')))" + ] + }, + { + "cell_type": "markdown", + "id": "c7da2ddd", + "metadata": {}, + "source": [ + "## 3. Clustering by problem: variance decomposition, design effect, cluster bootstrap\n", + "`m` samples per problem; the unit of independence is the **problem**. Super-population: `p_x ~ Beta(1.4, 0.7)`, `c_x | p_x ~ Bin(m, p_x)`, `p_hat_x = c_x / m`, estimator `mean_x p_hat_x`.\n", + "Total variance: `Var = [Var(p_x) + E p_x(1-p_x)/m] / N`. A sample-level (naive) SE treats all `N m` outcomes as independent: `p_bar(1-p_bar)/(N m)`.\n", + "Design effect `1 + (m-1) rho`, `rho = Var(p_x) / (p_bar(1-p_bar))`." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "9430c80f", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:32:03.533001Z", + "iopub.status.busy": "2026-10-14T09:32:03.532814Z", + "iopub.status.idle": "2026-10-14T09:32:06.152740Z", + "shell.execute_reply": "2026-10-14T09:32:06.151923Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "p_bar=0.6667 Var(p_x)=0.0717 ICC rho=0.3226 design effect=5.839 -> naive SE underestimates by x2.416\n", + "3. design effect identity Var_total / naive_Var = 1 + (m-1) rho err=8.882e-16 (tol 1e-12)\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "empirical Var(estimator)=3.3345e-04 theory=3.3235e-04 E[SE_cluster^2]=3.3236e-04 E[SE_naive^2]=5.6833e-05\n", + "3. empirical variance of the estimator vs [Var(p)+E p(1-p)/m]/N (relative) err=3.303e-03 (tol 3e-02)\n", + "3. cluster SE^2 is unbiased for the variance (relative) err=3.929e-05 (tol 3e-02)\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "coverage of nominal 95% CI for pass@1 over 600 replications: cluster bootstrap 0.950, naive sample bootstrap 0.585, Wilson on pooled 3904 outcomes 0.602\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "pass@8 (n=16 per problem, unbiased estimator): true 0.9622, cluster-bootstrap CI coverage 0.927\n" + ] + } + ], + "source": [ + "a, b = 1.4, 0.7; m = 16\n", + "p_bar = a / (a + b); var_p = a * b / ((a + b) ** 2 * (a + b + 1)); E_pq = p_bar - (var_p + p_bar ** 2)\n", + "var_true = (var_p + E_pq / m) / N\n", + "rho = var_p / (p_bar * (1 - p_bar)); deff = 1 + (m - 1) * rho\n", + "naive_var = p_bar * (1 - p_bar) / (N * m)\n", + "print(f'p_bar={p_bar:.4f} Var(p_x)={var_p:.4f} ICC rho={rho:.4f} design effect={deff:.3f} -> naive SE underestimates by x{math.sqrt(deff):.3f}')\n", + "check('3. design effect identity Var_total / naive_Var = 1 + (m-1) rho', abs(var_true / naive_var - deff), 1e-12)\n", + "rs = np.random.RandomState(0); R = 30000\n", + "px = rs.beta(a, b, size=(R, N)); cx = rs.binomial(m, px); phx = cx / m\n", + "est = phx.mean(1)\n", + "se2_cluster = phx.var(1, ddof=1) / N; se2_naive = est * (1 - est) / (N * m)\n", + "print(f'empirical Var(estimator)={est.var():.4e} theory={var_true:.4e} E[SE_cluster^2]={se2_cluster.mean():.4e} E[SE_naive^2]={se2_naive.mean():.4e}')\n", + "check('3. empirical variance of the estimator vs [Var(p)+E p(1-p)/m]/N (relative)', abs(est.var() / var_true - 1), 0.03)\n", + "check('3. cluster SE^2 is unbiased for the variance (relative)', abs(se2_cluster.mean() / var_true - 1), 0.03)\n", + "assert se2_naive.mean() < 0.5 * var_true\n", + "# coverage of 95% CIs for pass@1: cluster bootstrap vs naive sample-level bootstrap vs Wilson on pooled counts\n", + "Rcov, B = 600, 400\n", + "cov_cl = cov_nv = cov_wi = 0\n", + "rb = np.random.RandomState(1)\n", + "for r in range(Rcov):\n", + " px_ = rb.beta(a, b, size=N); cx_ = rb.binomial(m, px_); ph = cx_ / m; e = ph.mean()\n", + " idx = rb.randint(0, N, size=(B, N)); boot = ph[idx].mean(1); lo, hi = np.percentile(boot, [2.5, 97.5]); cov_cl += (lo <= p_bar <= hi)\n", + " nb = rb.binomial(N * m, e, size=B) / (N * m); lo, hi = np.percentile(nb, [2.5, 97.5]); cov_nv += (lo <= p_bar <= hi)\n", + " lo, hi = wilson(int(cx_.sum()), N * m); cov_wi += (lo <= p_bar <= hi)\n", + "print(f'coverage of nominal 95% CI for pass@1 over {Rcov} replications: cluster bootstrap {cov_cl/Rcov:.3f}, naive sample bootstrap {cov_nv/Rcov:.3f}, Wilson on pooled {N*m} outcomes {cov_wi/Rcov:.3f}')\n", + "assert cov_cl / Rcov > 0.92 and cov_nv / Rcov < 0.75\n", + "CHECKS.append(('3. coverage of nominal-95%% CI: cluster bootstrap %.3f, naive %.3f (informational)' % (cov_cl / Rcov, cov_nv / Rcov), float(cov_cl / Rcov), float('inf')))\n", + "\n", + "# pass@k with the unbiased per-problem estimator, cluster bootstrap\n", + "k = 8\n", + "def rho_nk(n, c, k): \n", + " out = np.ones_like(c, dtype=float)\n", + " for cc in range(n + 1): out[c == cc] = 1.0 if n - cc < k else 1.0 - float(np.prod(1.0 - k / np.arange(n - cc + 1, n + 1)))\n", + " return out\n", + "truth_k = 1 - math.exp(betaln(a, b + k) - betaln(a, b))\n", + "cov_k = 0; rb = np.random.RandomState(2)\n", + "for r in range(Rcov):\n", + " px_ = rb.beta(a, b, size=N); cx_ = rb.binomial(m, px_); est_x = rho_nk(m, cx_, k)\n", + " idx = rb.randint(0, N, size=(B, N)); boot = est_x[idx].mean(1); lo, hi = np.percentile(boot, [2.5, 97.5]); cov_k += (lo <= truth_k <= hi)\n", + "print(f'pass@{k} (n={m} per problem, unbiased estimator): true {truth_k:.4f}, cluster-bootstrap CI coverage {cov_k/Rcov:.3f}')\n", + "assert cov_k / Rcov > 0.92" + ] + }, + { + "cell_type": "markdown", + "id": "e6dd44fc", + "metadata": {}, + "source": [ + "## 4. Paired differences\n", + "`Delta_x = hat p^A_x - hat p^B_x`. `Var(mean Delta) = (Var A + Var B - 2 Cov(A,B)) / N` on the problem level. Unpaired analysis ignores the covariance." + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "49028f15", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:32:06.154500Z", + "iopub.status.busy": "2026-10-14T09:32:06.154336Z", + "iopub.status.idle": "2026-10-14T09:32:07.402778Z", + "shell.execute_reply": "2026-10-14T09:32:07.401927Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "empirical Var(mean Delta)=9.9803e-05; paired formula 9.9418e-05; unpaired (ignores Cov) 7.8556e-04; ratio unpaired/paired = 7.90\n", + "4. paired variance formula vs empirical (relative) err=3.864e-03 (tol 3e-02)\n" + ] + } + ], + "source": [ + "rs = np.random.RandomState(3); R = 20000\n", + "lA = rs.randn(R, N) * 1.8 + 0.5; lB = lA + 0.15 + rs.randn(R, N) * 0.4 # B = A shifted, correlated through the problem\n", + "pA, pB = 1 / (1 + np.exp(-lA)), 1 / (1 + np.exp(-lB))\n", + "cA, cB = rs.binomial(m, pA) / m, rs.binomial(m, pB) / m\n", + "D = (cB - cA).mean(1)\n", + "var_pair = np.mean(((cB - cA).var(1, ddof=1)) / N)\n", + "var_unpair = np.mean((cA.var(1, ddof=1) + cB.var(1, ddof=1)) / N)\n", + "print(f'empirical Var(mean Delta)={D.var():.4e}; paired formula {var_pair:.4e}; unpaired (ignores Cov) {var_unpair:.4e}; ratio unpaired/paired = {var_unpair/var_pair:.2f}')\n", + "check('4. paired variance formula vs empirical (relative)', abs(var_pair / D.var() - 1), 0.03)\n", + "assert var_unpair > 1.5 * var_pair" + ] + }, + { + "cell_type": "markdown", + "id": "ae7bfc19", + "metadata": {}, + "source": [ + "## 5. Power analysis for a net gain of 5/244\n", + "Outcomes per problem fall in 4 cells; the test uses only `(b, c)`. Joint law `(b, c, rest) ~ Multinomial(N; p_b, p_c, 1-p_b-p_c)`, with `delta = p_b - p_c = 5/244` and discordance `d = p_b + p_c`.\n", + "EXACT power = `sum_{n_d} Bin(n_d; N, d) sum_b Bin(b; n_d, p_b/d) 1{ exact p(b, n_d - b) < 0.05 }`.\n", + "NORMAL approximation (Connor 1987): `N >= [z_{a/2} sqrt(d) + z_b sqrt(d - delta^2)]^2 / delta^2`." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "f74c8d64", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:32:07.404574Z", + "iopub.status.busy": "2026-10-14T09:32:07.404402Z", + "iopub.status.idle": "2026-10-14T09:32:55.741131Z", + "shell.execute_reply": "2026-10-14T09:32:55.739908Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "discordance d exact power at N=244 normal-approx power N needed (normal, 80%) MDE at N=244 for 80% power (problems)\n", + " 0.03 0.306 0.455 558 7.3\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + " 0.06 0.181 0.256 1119 10.9\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + " 0.10 0.122 0.171 1867 14.3\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + " 0.15 0.097 0.128 2801 17.5\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + " 0.20 0.086 0.106 3736 20.1\n", + "5. exact power vs normal approximation at d=0.10, N=244 (abs diff; the normal formula ignores discreteness) err=4.952e-02 (tol 1e-01)\n", + "5. exact power (d=0.10) vs Monte Carlo over 4e5 benchmarks (abs diff) err=2.797e-04 (tol 3e-03)\n", + "5. vectorised threshold power = table-based exact power (d=0.10, N=244) err=4.663e-15 (tol 1e-12)\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "required N for 80% exact power at delta=5/244: {0.03: (584, 558), 0.06: (1184, 1119), 0.1: (1946, 1867), 0.15: (2887, 2801), 0.2: (3826, 3736)} (exact, normal approx)\n", + "5. required N (exact) vs normal approximation, d=0.10 (relative difference) err=4.242e-02 (tol 5e-02)\n" + ] + } + ], + "source": [ + "alpha = 0.05\n", + "reject = np.zeros((N + 1, N + 1), dtype=bool) # reject[n_d, b]\n", + "for n_d in range(1, N + 1):\n", + " for bb in range(n_d + 1):\n", + " reject[n_d, bb] = stats.binomtest(min(bb, n_d - bb), n_d, 0.5).pvalue < alpha\n", + "def power_exact(N_, d, delta):\n", + " pb = (d + delta) / 2; pc = (d - delta) / 2\n", + " nds = np.arange(0, N_ + 1); w = stats.binom.pmf(nds, N_, d)\n", + " tot = 0.0\n", + " for n_d in range(1, N_ + 1):\n", + " if w[n_d] < 1e-14: continue\n", + " bs = np.arange(0, n_d + 1); tot += w[n_d] * (stats.binom.pmf(bs, n_d, pb / d) * reject[n_d, bs]).sum()\n", + " return float(tot)\n", + "za, zb = stats.norm.ppf(1 - alpha / 2), stats.norm.ppf(0.8)\n", + "def N_needed_normal(d, delta, power=0.8):\n", + " zb_ = stats.norm.ppf(power); return (za * math.sqrt(d) + zb_ * math.sqrt(d - delta ** 2)) ** 2 / delta ** 2\n", + "delta = 5 / N\n", + "print('discordance d exact power at N=244 normal-approx power N needed (normal, 80%) MDE at N=244 for 80% power (problems)')\n", + "res_pow = {}\n", + "for d in (0.03, 0.06, 0.10, 0.15, 0.20):\n", + " pe = power_exact(N, d, delta)\n", + " pn = float(stats.norm.cdf((delta * math.sqrt(N) - za * math.sqrt(d)) / math.sqrt(d - delta ** 2)))\n", + " # minimal delta with exact power >= 0.8\n", + " lo_, hi_ = delta, d - 1e-9\n", + " for _ in range(40):\n", + " mid = (lo_ + hi_) / 2\n", + " if power_exact(N, d, mid) >= 0.8: hi_ = mid\n", + " else: lo_ = mid\n", + " res_pow[d] = (pe, pn, N_needed_normal(d, delta), hi_ * N)\n", + " print(f' {d:5.2f} {pe:7.3f} {pn:7.3f} {res_pow[d][2]:8.0f} {hi_*N:6.1f}')\n", + "assert all(v[0] < 0.35 for v in res_pow.values())\n", + "check('5. exact power vs normal approximation at d=0.10, N=244 (abs diff; the normal formula ignores discreteness)', abs(res_pow[0.10][0] - res_pow[0.10][1]), 0.12)\n", + "# Monte Carlo check of the exact-power calculation\n", + "rs = np.random.RandomState(5); d = 0.10; pb = (d + delta) / 2; pc = (d - delta) / 2\n", + "cnt = rs.multinomial(N, [pb, pc, 1 - pb - pc], size=400_000)\n", + "mc_power = reject[cnt[:, 0] + cnt[:, 1], cnt[:, 0]].mean()\n", + "check('5. exact power (d=0.10) vs Monte Carlo over 4e5 benchmarks (abs diff)', abs(mc_power - res_pow[0.10][0]), 0.003)\n", + "# Vectorised exact power via rejection thresholds: reject iff min(b,c) <= k*(n_d), k* = ppf(alpha/2) - 1 (two-sided exact p = 2 cdf(min) < alpha)\n", + "def power_exact_fast(N_, d, delta, alpha=0.05):\n", + " nd = np.arange(1, N_ + 1); w = stats.binom.pmf(nd, N_, d)\n", + " ks = stats.binom.ppf(alpha / 2, nd, 0.5) - 1\n", + " q = (d + delta) / (2 * d)\n", + " ok = ks >= 0\n", + " pr = np.zeros(len(nd)); kk = ks[ok]; nn = nd[ok]\n", + " pr[ok] = stats.binom.cdf(kk, nn, q) + stats.binom.sf(nn - kk - 1, nn, q)\n", + " return float((w * pr).sum())\n", + "check('5. vectorised threshold power = table-based exact power (d=0.10, N=244)', abs(power_exact_fast(N, 0.10, delta) - res_pow[0.10][0]), 1e-12)\n", + "def N_needed_exact(d, delta, power=0.8):\n", + " lo_, hi_ = 50, 200000\n", + " while hi_ - lo_ > 1:\n", + " mid = (lo_ + hi_) // 2\n", + " # power is not exactly monotone in N because of discreteness: use the first N after which power stays >= target on a window\n", + " if all(power_exact_fast(mid + t, d, delta) >= power for t in range(0, 6)): hi_ = mid\n", + " else: lo_ = mid\n", + " return hi_\n", + "n_ex = {d: N_needed_exact(d, delta) for d in (0.03, 0.06, 0.10, 0.15, 0.20)}\n", + "print('required N for 80% exact power at delta=5/244: ', {d: (n_ex[d], round(N_needed_normal(d, delta))) for d in n_ex}, ' (exact, normal approx)')\n", + "check('5. required N (exact) vs normal approximation, d=0.10 (relative difference)', abs(n_ex[0.10] / N_needed_normal(0.10, delta) - 1), 0.05)\n", + "CHECKS.append(('5. N required for 80%% power at delta=5/244, d=0.10: exact %d, normal %.0f (informational)' % (n_ex[0.10], N_needed_normal(0.10, delta)), float(n_ex[0.10]), float('inf')))" + ] + }, + { + "cell_type": "markdown", + "id": "cb66649f", + "metadata": {}, + "source": [ + "## 6. Size of tests with few discordant pairs, and the `(5, 0)` outcome\n", + "Under `H0` condition on `n_d` and `b ~ Bin(n_d, 1/2)`. EXACT type-I error of (i) exact McNemar, (ii) Wald `z = (b-c)/sqrt(b+c)` at 1.96, (iii) paired percentile bootstrap." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "91eb7b3f", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:32:55.742903Z", + "iopub.status.busy": "2026-10-14T09:32:55.742712Z", + "iopub.status.idle": "2026-10-14T09:32:58.063768Z", + "shell.execute_reply": "2026-10-14T09:32:58.062659Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "n_d size(exact) size(Wald z>1.96) size(bootstrap 95% CI excludes 0, MC)\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + " 5 0.0000 0.0625 0.0593\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + " 6 0.0313 0.0313 0.0287\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + " 8 0.0078 0.0703 0.0627\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + " 10 0.0215 0.0215 0.0307\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + " 12 0.0386 0.0386 0.0373\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + " 15 0.0352 0.0352 0.0373\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + " 20 0.0414 0.0414 0.0440\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + " 30 0.0428 0.0428 0.0420\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + " 60 0.0273 0.0519 0.0533\n", + "(b,c)=(5,0): exact p = 0.0625; Wald z = 2.236 -> p = 0.0253; bootstrap 2.5% quantile of b* = 1, P(b*=0) = 0.0064\n" + ] + } + ], + "source": [ + "print('n_d size(exact) size(Wald z>1.96) size(bootstrap 95% CI excludes 0, MC)')\n", + "rs = np.random.RandomState(6); sizes = {}\n", + "for n_d in (5, 6, 8, 10, 12, 15, 20, 30, 60):\n", + " bs = np.arange(0, n_d + 1); pm = stats.binom.pmf(bs, n_d, 0.5)\n", + " s_exact = float((pm * np.array([stats.binomtest(min(x, n_d - x), n_d, 0.5).pvalue < alpha for x in bs])).sum())\n", + " s_wald = float((pm * (np.abs(2 * bs - n_d) / math.sqrt(n_d) >= za)).sum())\n", + " # bootstrap: resample the N problems => (b*, c*, rest*) ~ Multinomial(N, [b, c, rest]/N); CI of (b*-c*) from 400 draws; reject if CI excludes 0\n", + " rej = 0; reps = 1500\n", + " for r in range(reps):\n", + " bb = rs.binomial(n_d, 0.5); cc = n_d - bb\n", + " draws = rs.multinomial(N, [bb / N, cc / N, 1 - n_d / N], size=400); diff = draws[:, 0] - draws[:, 1]\n", + " lo, hi = np.percentile(diff, [2.5, 97.5]); rej += (lo > 0) or (hi < 0)\n", + " s_boot = rej / reps; sizes[n_d] = (s_exact, s_wald, s_boot)\n", + " print(f'{n_d:3d} {s_exact:9.4f} {s_wald:9.4f} {s_boot:9.4f}')\n", + "assert all(v[0] <= alpha + 1e-12 for v in sizes.values())\n", + "assert sizes[5][1] > alpha and sizes[8][1] > alpha # Wald exceeds its nominal size at n_d = 5 and 8\n", + "CHECKS.append(('6. exact McNemar size <= 0.05 for every n_d tested (max size %.4f)' % max(v[0] for v in sizes.values()), float(max(v[0] for v in sizes.values())), alpha))\n", + "z50 = 5 / math.sqrt(5); p_wald = 2 * (1 - stats.norm.cdf(z50))\n", + "p0 = (1 - 5 / N) ** N\n", + "draws = np.random.RandomState(8).multinomial(N, [5 / N, 0, 1 - 5 / N], size=200000)[:, 0]\n", + "print(f'(b,c)=(5,0): exact p = {float(mcnemar_exact_fraction(5,0)):.4f}; Wald z = {z50:.3f} -> p = {p_wald:.4f}; bootstrap 2.5% quantile of b* = {np.percentile(draws, 2.5):.0f}, P(b*=0) = {p0:.4f}')\n", + "CHECKS.append(('6. (5,0): exact p = %.4f, Wald p = %.4f (informational: they disagree about significance)' % (0.0625, p_wald), float(p_wald), float('inf')))" + ] + }, + { + "cell_type": "markdown", + "id": "14039858", + "metadata": {}, + "source": [ + "## Verification table" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "09c2d874", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:32:58.065562Z", + "iopub.status.busy": "2026-10-14T09:32:58.065394Z", + "iopub.status.idle": "2026-10-14T09:32:58.069524Z", + "shell.execute_reply": "2026-10-14T09:32:58.068827Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + "identity max abs err\n", + "1. exact McNemar p: rational tail vs scipy.binomtest (all b,c <= 16) 1.110e-16\n", + "1. exact McNemar p: rational tail vs sign-flip permutation enumeration (n_d <= 14) 0.000e+00\n", + "1. best possible p for a +5 net gain is (5,0): p = 1/16 > 0.05 (|p - 0.0625|) 0.000e+00\n", + "2. Wilson closed form vs numerical inversion of the score test (max |endpoint diff|, x=1..243) 4.905e-13\n", + "2. min exact coverage over p in [0.01,0.995]: Wilson / Wald / Clopper-Pearson = 0.9402 / 0.8634 / 0.9507 (informational) 9.402e-01\n", + "3. design effect identity Var_total / naive_Var = 1 + (m-1) rho 8.882e-16\n", + "3. empirical variance of the estimator vs [Var(p)+E p(1-p)/m]/N (relative) 3.303e-03\n", + "3. cluster SE^2 is unbiased for the variance (relative) 3.929e-05\n", + "3. coverage of nominal-95% CI: cluster bootstrap 0.950, naive 0.585 (informational) 9.500e-01\n", + "4. paired variance formula vs empirical (relative) 3.864e-03\n", + "5. exact power vs normal approximation at d=0.10, N=244 (abs diff; the normal formula ignores discreteness) 4.952e-02\n", + "5. exact power (d=0.10) vs Monte Carlo over 4e5 benchmarks (abs diff) 2.797e-04\n", + "5. vectorised threshold power = table-based exact power (d=0.10, N=244) 4.663e-15\n", + "5. required N (exact) vs normal approximation, d=0.10 (relative difference) 4.242e-02\n", + "5. N required for 80% power at delta=5/244, d=0.10: exact 1946, normal 1867 (informational) 1.946e+03\n", + "6. exact McNemar size <= 0.05 for every n_d tested (max size 0.0428) 4.277e-02\n", + "6. (5,0): exact p = 0.0625, Wald p = 0.0253 (informational: they disagree about significance) 2.535e-02\n", + "CHECKS_JSON:[[\"1. exact McNemar p: rational tail vs scipy.binomtest (all b,c <= 16)\", 1.1102230246251565e-16, 1e-12], [\"1. exact McNemar p: rational tail vs sign-flip permutation enumeration (n_d <= 14)\", 0.0, 1e-15], [\"1. best possible p for a +5 net gain is (5,0): p = 1/16 > 0.05 (|p - 0.0625|)\", 0.0, 0.0], [\"2. Wilson closed form vs numerical inversion of the score test (max |endpoint diff|, x=1..243)\", 4.904965322793942e-13, 1e-10], [\"2. min exact coverage over p in [0.01,0.995]: Wilson / Wald / Clopper-Pearson = 0.9402 / 0.8634 / 0.9507 (informational)\", 0.940211759420751, Infinity], [\"3. design effect identity Var_total / naive_Var = 1 + (m-1) rho\", 8.881784197001252e-16, 1e-12], [\"3. empirical variance of the estimator vs [Var(p)+E p(1-p)/m]/N (relative)\", 0.0033025628714560273, 0.03], [\"3. cluster SE^2 is unbiased for the variance (relative)\", 3.929333853491812e-05, 0.03], [\"3. coverage of nominal-95% CI: cluster bootstrap 0.950, naive 0.585 (informational)\", 0.95, Infinity], [\"4. paired variance formula vs empirical (relative)\", 0.0038637080259513734, 0.03], [\"5. exact power vs normal approximation at d=0.10, N=244 (abs diff; the normal formula ignores discreteness)\", 0.049521981277600394, 0.12], [\"5. exact power (d=0.10) vs Monte Carlo over 4e5 benchmarks (abs diff)\", 0.0002797189205828443, 0.003], [\"5. vectorised threshold power = table-based exact power (d=0.10, N=244)\", 4.6629367034256575e-15, 1e-12], [\"5. required N (exact) vs normal approximation, d=0.10 (relative difference)\", 0.042423303779374866, 0.05], [\"5. N required for 80% power at delta=5/244, d=0.10: exact 1946, normal 1867 (informational)\", 1946.0, Infinity], [\"6. exact McNemar size <= 0.05 for every n_d tested (max size 0.0428)\", 0.04277394525706773, 0.05], [\"6. (5,0): exact p = 0.0625, Wald p = 0.0253 (informational: they disagree about significance)\", 0.0253473186774682, Infinity]]\n" + ] + } + ], + "source": [ + "print()\n", + "print(f'{\"identity\":98s} {\"max abs err\":>12s}')\n", + "for n_, e_, t_ in CHECKS: print(f'{n_:98s} {e_:12.3e}')\n", + "print('CHECKS_JSON:' + json.dumps(CHECKS))" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/code/notebooks/0209-variance-components-equivalence.ipynb b/code/notebooks/0209-variance-components-equivalence.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..c283fb0a04b45208bb90bae97f6a5fe2908fd74b --- /dev/null +++ b/code/notebooks/0209-variance-components-equivalence.ipynb @@ -0,0 +1,486 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# 0209 - Variance components and equivalence tests for evaluation\n", + "Companion to `notes/0209 - Variance Components and Equivalence Tests.pdf`. CPU only, numpy/scipy, 1 thread.\n", + "Running toy: a benchmark of `N = 244` problems, `S` independently trained models (\"seeds\") of one recipe, `m = 16` samples per (seed, problem).\n", + "Per-(seed, problem) solve probability `p_sx = expit(mu_x + a_s + b_sx)`, `mu_x ~ N(0.3, 1.5^2)` (problem difficulty), `a_s ~ N(0, 0.25^2)` (seed), `b_sx ~ N(0, 0.5^2)` (seed x problem)." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "source": [ + "import os\n", + "os.environ[\"OMP_NUM_THREADS\"] = \"1\"; os.environ[\"OPENBLAS_NUM_THREADS\"] = \"1\"; os.environ[\"MKL_NUM_THREADS\"] = \"1\"\n", + "import math, itertools\n", + "from fractions import Fraction as Fr\n", + "from math import comb\n", + "import numpy as np\n", + "from scipy import stats\n", + "from scipy.special import expit\n", + "CHECKS = []\n", + "def check(name, err, tol, note=\"\"):\n", + " CHECKS.append((name, float(err), float(tol)))\n", + " assert err <= tol, (name, float(err), tol)\n", + " print(f\"[ok] {name:100s} err={float(err):.3e} (tol {tol:.1e}) {note}\")\n", + "N, m = 244, 16\n", + "M0, SX, SA, SB = 0.3, 1.5, 0.25, 0.5" + ], + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 1. Variance components: problems (A), seeds (B), seed x problem (C), generation (e)\n", + "Exchangeable array `p_sx = mu + A_x + B_s + C_sx` (Hoeffding/ANOVA decomposition; components orthogonal). With outcomes `c_sx ~ Bin(m, p_sx)` and the score of one run `Y_s = (1/(N m)) sum_x c_sx`:\n", + "`Var(Y_s) = sA2/N + sB2 + sC2/N + se2/(N m)`, `se2 = E p(1-p)`. Only `sB2` survives `N -> inf` and `m -> inf`.\n", + "Population components by Gauss-Hermite integration of the logit-normal model (no simulation)." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "source": [ + "gh_x, gh_w = np.polynomial.hermite_e.hermegauss(60); gh_w = gh_w / gh_w.sum() # standard normal nodes/weights\n", + "MU = M0 + SX * gh_x; A_ = SA * gh_x; B_ = SB * gh_x\n", + "P3 = expit(MU[:, None, None] + A_[None, :, None] + B_[None, None, :]) # [x, s, b]\n", + "W3 = gh_w[:, None, None] * gh_w[None, :, None] * gh_w[None, None, :]\n", + "Ep = (P3 * W3).sum(); Ep2 = (P3 ** 2 * W3).sum()\n", + "Ax = (P3 * gh_w[None, :, None] * gh_w[None, None, :]).sum((1, 2)) # E_{s,b} p | x (weights over s,b sum to 1)\n", + "Bs = (P3 * gh_w[:, None, None] * gh_w[None, None, :]).sum((0, 2)) # E_{x,b} p | s\n", + "sA2 = (gh_w * (Ax - Ep) ** 2).sum(); sB2 = (gh_w * (Bs - Ep) ** 2).sum()\n", + "sT2 = Ep2 - Ep ** 2; sC2 = sT2 - sA2 - sB2; se2 = Ep - Ep2\n", + "print(f\" E p = {Ep:.4f}; sigma_A^2 = {sA2:.5f}, sigma_B^2 = {sB2:.6f} (sd {math.sqrt(sB2):.4f}), sigma_C^2 = {sC2:.5f}, sigma_e^2 = E p(1-p) = {se2:.5f}\")\n", + "def var_Y(S, N=N, m=m): return sA2 / N + sB2 / S + sC2 / (N * S) + se2 / (N * S * m) # variance of a grand mean over S seeds\n", + "print(\" per-term contribution to Var(Y_s), S=1: problems %.3e seed %.3e seed*problem %.3e generation %.3e\" % (sA2 / N, sB2, sC2 / N, se2 / (N * m)))\n", + "check(\"1. components are non-negative and sum to Var(p)\", abs(sA2 + sB2 + sC2 - sT2) + max(0, -min(sA2, sB2, sC2)), 1e-15)\n", + "# simulate replications: each replication redraws problems, seeds, samples\n", + "rng = np.random.default_rng(0)\n", + "def simulate(R, S, N=N, m=m):\n", + " mu = M0 + SX * rng.standard_normal((R, 1, N)); a = SA * rng.standard_normal((R, S, 1)); b = SB * rng.standard_normal((R, S, N))\n", + " p = expit(mu + a + b)\n", + " c = rng.binomial(m, p)\n", + " return p, c\n", + "R1 = 6000\n", + "p, c = simulate(R1, 1)\n", + "Y = c[:, 0, :].sum(1) / (N * m)\n", + "se_rel = math.sqrt(2 / (R1 - 1))\n", + "check(\"1. Var(Y_s) over 6000 replications vs sA2/N + sB2 + sC2/N + se2/(Nm)\", abs(Y.var(ddof=1) / var_Y(1) - 1), 4 * se_rel, f\"emp {Y.var(ddof=1):.4e} vs {var_Y(1):.4e}\")\n", + "# MoM / ANOVA estimators with S=4 seeds\n", + "S = 4; R2 = 4000\n", + "p, c = simulate(R2, S)\n", + "ybar = c / m # cell means [R,S,N]\n", + "gm = ybar.mean((1, 2), keepdims=True)\n", + "MSx = S * ((ybar.mean(1) - gm[:, :, 0]) ** 2).sum(1) / (N - 1)\n", + "MSs = N * ((ybar.mean(2) - gm[:, :, 0]) ** 2).sum(1) / (S - 1)\n", + "res = ybar - ybar.mean(1, keepdims=True) - ybar.mean(2, keepdims=True) + gm\n", + "MSr = (res ** 2).sum((1, 2)) / ((S - 1) * (N - 1))\n", + "pq = (c * (m - c) / (m * (m - 1))).mean((1, 2)) # unbiased for E p(1-p)\n", + "est = {\"sigma_e^2\": (pq, se2), \"sigma_C^2\": (MSr - pq / m, sC2), \"sigma_A^2\": ((MSx - MSr) / S, sA2), \"sigma_B^2\": ((MSs - MSr) / N, sB2)}\n", + "for k, (v, t) in est.items():\n", + " z = (v.mean() - t) / (v.std(ddof=1) / math.sqrt(R2))\n", + " check(f\"1. ANOVA estimator of {k} is unbiased (4000 reps, S=4)\", abs(z), 4.0, f\"mean {v.mean():.6f} vs true {t:.6f} (z={z:+.2f})\")" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + " E p = 0.5512; sigma_A^2 = 0.06735, sigma_B^2 = 0.001800 (sd 0.0424), sigma_C^2 = 0.00852, sigma_e^2 = E p(1-p) = 0.16971\n", + " per-term contribution to Var(Y_s), S=1: problems 2.760e-04 seed 1.800e-03 seed*problem 3.490e-05 generation 4.347e-05\n", + "[ok] 1. components are non-negative and sum to Var(p) err=0.000e+00 (tol 1.0e-15) \n", + "[ok] 1. Var(Y_s) over 6000 replications vs sA2/N + sB2 + sC2/N + se2/(Nm) err=2.987e-02 (tol 7.3e-02) emp 2.2191e-03 vs 2.1547e-03\n", + "[ok] 1. ANOVA estimator of sigma_e^2 is unbiased (4000 reps, S=4) err=9.497e-01 (tol 4.0e+00) mean 0.169776 vs true 0.169709 (z=+0.95)\n", + "[ok] 1. ANOVA estimator of sigma_C^2 is unbiased (4000 reps, S=4) err=1.471e-01 (tol 4.0e+00) mean 0.008513 vs true 0.008515 (z=-0.15)\n", + "[ok] 1. ANOVA estimator of sigma_A^2 is unbiased (4000 reps, S=4) err=6.625e-01 (tol 4.0e+00) mean 0.067303 vs true 0.067352 (z=-0.66)\n", + "[ok] 1. ANOVA estimator of sigma_B^2 is unbiased (4000 reps, S=4) err=1.567e+00 (tol 4.0e+00) mean 0.001764 vs true 0.001800 (z=-1.57)\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 2. Which term do you pay down with what? Two recipes, S seeds each, same problems\n", + "`delta_hat = Ybar_A - Ybar_B`. Under H0 (same recipe) the shared problem effect A_x cancels: `Var(delta_hat) = (2/S) [ sB2 + sC2/N + se2/(N m) ]`.\n", + "The seed term does not shrink with N or m: a floor on the minimum detectable effect." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "source": [ + "def var_diff(S, N=N, m=m): return (2 / S) * (sB2 + sC2 / N + se2 / (N * m))\n", + "R3 = 6000\n", + "for S in (1, 3):\n", + " mu = M0 + SX * rng.standard_normal((R3, 1, N))\n", + " d = []\n", + " for _ in range(2):\n", + " a = SA * rng.standard_normal((R3, S, 1)); b = SB * rng.standard_normal((R3, S, N))\n", + " cc = rng.binomial(m, expit(mu + a + b)); d.append(cc.sum((1, 2)) / (S * N * m))\n", + " dd = d[0] - d[1]\n", + " check(f\"2. Var(delta_hat) for S={S} seeds per recipe vs (2/S)[sB2 + sC2/N + se2/(Nm)]\", abs(dd.var(ddof=1) / var_diff(S) - 1), 4 * math.sqrt(2 / (R3 - 1)), f\"emp {dd.var(ddof=1):.3e} vs {var_diff(S):.3e}\")\n", + "z = stats.norm.ppf(0.975) + stats.norm.ppf(0.8)\n", + "for S in (1, 3, 5):\n", + " print(f\" S={S}: MDE(80% power, 5% two-sided) at N=244: {z*math.sqrt(var_diff(S))*100:.2f} pp; N->inf, m->inf floor: {z*math.sqrt(2*sB2/S)*100:.2f} pp\")\n", + "check(\"2. floor is strictly positive: MDE(N=1e9, m=1e9, S=3) equals z sqrt(2 sB2 / 3)\", abs(z * math.sqrt(var_diff(3, N=10**9, m=10**9)) - z * math.sqrt(2 * sB2 / 3)), 1e-6)" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "[ok] 2. Var(delta_hat) for S=1 seeds per recipe vs (2/S)[sB2 + sC2/N + se2/(Nm)] err=1.212e-02 (tol 7.3e-02) emp 3.803e-03 vs 3.757e-03\n", + "[ok] 2. Var(delta_hat) for S=3 seeds per recipe vs (2/S)[sB2 + sC2/N + se2/(Nm)] err=1.255e-02 (tol 7.3e-02) emp 1.237e-03 vs 1.252e-03\n", + " S=1: MDE(80% power, 5% two-sided) at N=244: 17.17 pp; N->inf, m->inf floor: 16.81 pp\n", + " S=3: MDE(80% power, 5% two-sided) at N=244: 9.91 pp; N->inf, m->inf floor: 9.71 pp\n", + " S=5: MDE(80% power, 5% two-sided) at N=244: 7.68 pp; N->inf, m->inf floor: 7.52 pp\n", + "[ok] 2. floor is strictly positive: MDE(N=1e9, m=1e9, S=3) equals z sqrt(2 sB2 / 3) err=2.295e-10 (tol 1.0e-06) \n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 3. Unbiased pass@k from n samples, and its exact variance\n", + "`e(n,c,k) = 1 - C(n-c,k)/C(n,k)`; `c ~ Bin(n,p)` => `E e = 1 - (1-p)^k`. Exact variance by summing over `c = 0..n`.\n", + "NAIVE plug-in `1 - (1 - c/n)^k` is biased. Special cases: `k = 1` gives `p(1-p)/n`; `k = n` gives `q(1-q)`, `q = 1-(1-p)^n`." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "source": [ + "def e_nck(n, c, k): return 1 - Fr(comb(n - c, k), comb(n, k))\n", + "def pk_moments(n, k, p):\n", + " p = Fr(p); pm = [Fr(comb(n, c)) * p ** c * (1 - p) ** (n - c) for c in range(n + 1)]\n", + " E1 = sum(w * e_nck(n, c, k) for c, w in enumerate(pm)); E2 = sum(w * e_nck(n, c, k) ** 2 for c, w in enumerate(pm))\n", + " return E1, E2 - E1 ** 2, sum(w * (1 - (1 - Fr(c, n)) ** k) for c, w in enumerate(pm))\n", + "worst = 0; worst_k1 = 0; worst_kn = 0\n", + "for n, k, p in [(64, 8, Fr(1, 10)), (64, 32, Fr(1, 50)), (16, 4, Fr(3, 10)), (64, 1, Fr(1, 7)), (64, 64, Fr(1, 100)), (10, 10, Fr(1, 5))]:\n", + " E1, V, Enaive = pk_moments(n, k, p)\n", + " worst = max(worst, abs(float(E1 - (1 - (1 - p) ** k))))\n", + " if k == 1: worst_k1 = max(worst_k1, abs(float(V - p * (1 - p) / n)))\n", + " if k == n: q = 1 - (1 - p) ** n; worst_kn = max(worst_kn, abs(float(V - q * (1 - q))))\n", + "check(\"3. E[e(n,c,k)] = 1-(1-p)^k exactly (6 (n,k,p) triples, rational arithmetic)\", worst, 1e-15)\n", + "check(\"3. k=1: Var = p(1-p)/n exactly\", worst_k1, 1e-15)\n", + "check(\"3. k=n: Var = q(1-q) with q = 1-(1-p)^n exactly\", worst_kn, 1e-15)\n", + "print(\" n=64: p k true pass@k E[unbiased] E[naive plug-in] (bias) sd(unbiased) sd/true\")\n", + "for p in (0.02, 0.1, 0.5):\n", + " for k in (1, 8, 32):\n", + " E1, V, En = pk_moments(64, k, Fr(p).limit_denominator(1000)); g = 1 - (1 - p) ** k\n", + " print(f\" {p:5.2f} {k:4d} {g:.5f} {float(E1):.5f} {float(En):.5f} ({float(En)-g:+.5f}) {math.sqrt(float(V)):.5f} {math.sqrt(float(V))/g:.3f}\")\n", + "E1, V, En = pk_moments(64, 32, Fr(1, 50))\n", + "check(\"3. naive plug-in is biased low at (n=64,k=32,p=.02) by more than 0.07\", 0.0 if (1 - (1 - 0.02) ** 32) - float(En) > 0.07 else 1.0, 0.0, f\"bias {float(En)-(1-0.98**32):+.4f}\")\n", + "# MC check of the exact variance (independent implementation via binomial draws) and of the problem-level decomposition\n", + "rng = np.random.default_rng(1)\n", + "n, k, p0 = 64, 8, 0.1\n", + "cs = rng.binomial(n, p0, size=2_000_000)\n", + "lg = np.array([float(e_nck(n, c_, k)) for c_ in range(n + 1)])\n", + "est = lg[cs]\n", + "Vex = float(pk_moments(n, k, Fr(1, 10))[1])\n", + "check(\"3. exact Var vs 2e6 Monte Carlo draws (n=64,k=8,p=.1)\", abs(est.var() / Vex - 1), 4 * math.sqrt(2 / 2e6))\n", + "# across problems: p_x ~ Beta(1.4, 0.7); Var(mean_x e_x) = (1/N)[Var_x(g(p_x)) + E_x Var(e | p_x)]\n", + "a_, b_ = 1.4, 0.7; k = 8; n = 64\n", + "xs = rng.beta(a_, b_, size=400_000) # for the population quantities\n", + "cs = rng.binomial(n, xs)\n", + "lgk = np.array([float(e_nck(n, c_, k)) for c_ in range(n + 1)])\n", + "es = lgk[cs]; gx = 1 - (1 - xs) ** k\n", + "pred = (gx.var() + (es - gx).var()) / N\n", + "reps = np.array([lgk[rng.binomial(n, rng.beta(a_, b_, size=N))].mean() for _ in range(6000)])\n", + "check(\"3. Var of the benchmark mean pass@8 (N=244, n=64) vs (1/N)[Var_x g + E Var(e|p)]\", abs(reps.var(ddof=1) / pred - 1), 5 * math.sqrt(2 / 5999) + 0.01, f\"emp {reps.var(ddof=1):.3e} vs {pred:.3e}\")" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "[ok] 3. E[e(n,c,k)] = 1-(1-p)^k exactly (6 (n,k,p) triples, rational arithmetic) err=0.000e+00 (tol 1.0e-15) \n", + "[ok] 3. k=1: Var = p(1-p)/n exactly err=0.000e+00 (tol 1.0e-15) \n", + "[ok] 3. k=n: Var = q(1-q) with q = 1-(1-p)^n exactly err=0.000e+00 (tol 1.0e-15) \n", + " n=64: p k true pass@k E[unbiased] E[naive plug-in] (bias) sd(unbiased) sd/true\n", + " 0.02 1 0.02000 0.02000 0.02000 (-0.00000) 0.01750 0.875\n", + " 0.02 8 0.14924 0.14924 0.14185 (-0.00739) 0.12202 0.818\n", + " 0.02 32 0.47612 0.47612 0.40018 (-0.07594) 0.32411 0.681\n", + " 0.10 1 0.10000 0.10000 0.10000 (+0.00000) 0.03750 0.375\n", + " 0.10 8 0.56953 0.56953 0.54892 (-0.02061) 0.14662 0.257\n", + " 0.10 32 0.96566 0.96566 0.92934 (-0.03633) 0.07300 0.076\n", + " 0.50 1 0.50000 0.50000 0.50000 (+0.00000) 0.06250 0.125\n", + " 0.50 8 0.99609 0.99609 0.99418 (-0.00191) 0.00474 0.005\n", + " 0.50 32 1.00000 1.00000 1.00000 (-0.00000) 0.00000 0.000\n", + "[ok] 3. naive plug-in is biased low at (n=64,k=32,p=.02) by more than 0.07 err=0.000e+00 (tol 0.0e+00) bias -0.0759\n", + "[ok] 3. exact Var vs 2e6 Monte Carlo draws (n=64,k=8,p=.1) err=9.304e-04 (tol 4.0e-03) \n", + "[ok] 3. Var of the benchmark mean pass@8 (N=244, n=64) vs (1/N)[Var_x g + E Var(e|p)] err=2.299e-02 (tol 1.0e-01) emp 6.460e-05 vs 6.315e-05\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 4. Paired sign-flip test on per-problem differences\n", + "`d_x = e^A_x - e^B_x`, statistic `T = sum_x d_x`, null: each `d_x` symmetric about 0 independently => signs exchangeable: p = fraction of the `2^N` sign vectors with `|T'| >= |T|`.\n", + "Binary `d_x in {-1,0,1}`: reduces to exact McNemar (note 0208). Exact enumeration for N<=18; Monte Carlo for N=244; size under the null." + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "source": [ + "def signflip_exact(d):\n", + " d = np.asarray(d, float); n = len(d)\n", + " signs = np.array(list(itertools.product([-1.0, 1.0], repeat=n)))\n", + " T = np.abs(signs @ d); return float((T >= abs(d.sum()) - 1e-12).mean())\n", + "def mcnemar(b, c):\n", + " nd = b + c\n", + " return 1.0 if nd == 0 else min(1.0, 2 * sum(comb(nd, i) for i in range(min(b, c) + 1)) / 2 ** nd)\n", + "worst = 0\n", + "for b_ in range(0, 8):\n", + " for c_ in range(0, 8):\n", + " if b_ + c_ == 0: continue\n", + " d = np.array([1] * b_ + [-1] * c_ + [0] * 3)\n", + " worst = max(worst, abs(signflip_exact(d) - mcnemar(b_, c_)))\n", + "check(\"4. sign-flip p on {-1,0,+1} equals exact McNemar for all b,c <= 7 (3 concordant zeros)\", worst, 1e-12)\n", + "check(\"4. (5,0): two-sided sign-flip p = 0.0625\", abs(signflip_exact([1] * 5) - 0.0625), 1e-15)\n", + "rng = np.random.default_rng(2)\n", + "def signflip_mc(d, B=20000):\n", + " d = np.asarray(d, float); s = rng.choice([-1.0, 1.0], size=(B, len(d))); T = np.abs(s @ d)\n", + " return (1 + (T >= abs(d.sum()) - 1e-12).sum()) / (B + 1)\n", + "d = rng.standard_normal(14) * 0.2 + 0.05\n", + "check(\"4. Monte Carlo p-value (B=20000) vs exact enumeration at N=14\", abs(signflip_mc(d) - signflip_exact(d)), 4 * 0.5 / math.sqrt(20000))\n", + "# size under the null: A and B are two draws of the same system; per-problem pass@8 estimators from n=64 samples with problem-specific p_x\n", + "reject = 0; R = 1500; N_ = 60\n", + "for r in range(R):\n", + " px = rng.beta(a_, b_, size=N_)\n", + " d = lgk[rng.binomial(n, px)] - lgk[rng.binomial(n, px)]\n", + " reject += signflip_mc(d, B=2000) <= 0.05\n", + "size = reject / R\n", + "check(\"4. size of the sign-flip test under H0 (N=60 problems, shared p_x, 1500 benchmarks)\", abs(size - 0.05) , 4 * math.sqrt(0.05 * 0.95 / R) + 0.01, f\"size {size:.4f}\")\n", + "# the null is symmetry, not mean zero: a skewed mean-zero d inflates/deflates size\n", + "reject = 0; R = 1500; N_ = 12\n", + "for r in range(R):\n", + " d = np.where(rng.random(N_) < 0.1, 9.0, -1.0) # mean = 0.1*9 - 0.9 = 0, strongly skewed\n", + " reject += signflip_exact(d) <= 0.05 if N_ <= 12 else 0\n", + "size_skew = reject / R\n", + "print(f\" skewed mean-zero differences (N=12): size = {size_skew:.4f} (nominal 0.05): the test's null is symmetry\")\n", + "check(\"4. skewed mean-zero null: size differs from 0.05 by more than 0.02 (the test is not a test of the mean)\", 0.0 if abs(size_skew - 0.05) > 0.02 else 1.0, 0.0)" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "[ok] 4. sign-flip p on {-1,0,+1} equals exact McNemar for all b,c <= 7 (3 concordant zeros) err=0.000e+00 (tol 1.0e-12) \n", + "[ok] 4. (5,0): two-sided sign-flip p = 0.0625 err=0.000e+00 (tol 1.0e-15) \n", + "[ok] 4. Monte Carlo p-value (B=20000) vs exact enumeration at N=14 err=6.842e-04 (tol 1.4e-02) \n", + "[ok] 4. size of the sign-flip test under H0 (N=60 problems, shared p_x, 1500 benchmarks) err=8.000e-03 (tol 3.3e-02) size 0.0580\n", + " skewed mean-zero differences (N=12): size = 0.2940 (nominal 0.05): the test's null is symmetry\n", + "[ok] 4. skewed mean-zero null: size differs from 0.05 by more than 0.02 (the test is not a test of the mean) err=0.000e+00 (tol 0.0e+00) \n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 5. Fresh seed on the same problems removes generation noise, not selection bias\n", + "K = 8 candidate checkpoints with *identical* true population quality theta (null). On the fixed benchmark each has a quenched problem-specific advantage\n", + "`beta_k ~ N(0, sC2/N)` (seed x problem) plus a fresh-generation error `g_k ~ N(0, se2/(N m))`. Select `k* = argmax (beta_k + g_k)`.\n", + "(i) report the same estimate: bias `sigma_tot c_K`; (ii) re-sample generations on the SAME problems: bias `E beta_{k*} = (sigma_b^2/sigma_tot) c_K`; (iii) fresh problems: bias 0.\n", + "`c_K = E max of K standard normals`." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "source": [ + "sb2 = sC2 / N; sg2 = se2 / (N * m); st2 = sb2 + sg2; K = 8\n", + "cK = float(np.mean(np.max(np.random.default_rng(3).standard_normal((2_000_000, K)), 1)))\n", + "from scipy import integrate\n", + "cK_exact = integrate.quad(lambda x: K * x * stats.norm.pdf(x) * stats.norm.cdf(x) ** (K - 1), -12, 12)[0]\n", + "check(\"5. c_8 = E max of 8 N(0,1): quadrature vs 2e6 Monte Carlo\", abs(cK_exact - cK), 5 * 1.0 / math.sqrt(2e6) , f\"c_8={cK_exact:.5f}\")\n", + "rng = np.random.default_rng(4); R = 400_000\n", + "beta = math.sqrt(sb2) * rng.standard_normal((R, K)); g1 = math.sqrt(sg2) * rng.standard_normal((R, K)); g2 = math.sqrt(sg2) * rng.standard_normal((R, K))\n", + "win = np.argmax(beta + g1, axis=1); ix = np.arange(R)\n", + "b_same = (beta + g1)[ix, win]; b_regen = (beta + g2)[ix, win]; b_fresh = (math.sqrt(sb2) * rng.standard_normal(R) + math.sqrt(sg2) * rng.standard_normal(R))\n", + "th_same, th_regen = math.sqrt(st2) * cK_exact, (sb2 / math.sqrt(st2)) * cK_exact\n", + "sd = math.sqrt(st2)\n", + "check(\"5(i). same-sample winner estimate bias vs sigma_tot c_K\", abs(b_same.mean() - th_same), 4 * sd / math.sqrt(R), f\"{b_same.mean()*100:.3f} pp vs {th_same*100:.3f} pp\")\n", + "check(\"5(ii). fresh generations on the same problems: bias vs (sb2/sigma_tot) c_K\", abs(b_regen.mean() - th_regen), 4 * sd / math.sqrt(R), f\"{b_regen.mean()*100:.3f} pp vs {th_regen*100:.3f} pp (retained fraction {th_regen/th_same:.3f} = sb2/sigma_tot^2 = {sb2/st2:.3f})\")\n", + "check(\"5(iii). fresh problems: bias 0\", abs(b_fresh.mean()), 4 * sd / math.sqrt(R))\n", + "check(\"5. retained fraction of the bias after regenerating = sb2/(sb2+sg2) (with m=16: generation is a minor term)\", abs(th_regen / th_same - sb2 / st2), 1e-12)" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "[ok] 5. c_8 = E max of 8 N(0,1): quadrature vs 2e6 Monte Carlo err=7.607e-04 (tol 3.5e-03) c_8=1.42360\n", + "[ok] 5(i). same-sample winner estimate bias vs sigma_tot c_K err=1.363e-05 (tol 5.6e-05) 1.262 pp vs 1.260 pp\n", + "[ok] 5(ii). fresh generations on the same problems: bias vs (sb2/sigma_tot) c_K err=5.363e-06 (tol 5.6e-05) 0.561 pp vs 0.561 pp (retained fraction 0.445 = sb2/sigma_tot^2 = 0.445)\n", + "[ok] 5(iii). fresh problems: bias 0 err=1.369e-05 (tol 5.6e-05) \n", + "[ok] 5. retained fraction of the bias after regenerating = sb2/(sb2+sg2) (with m=16: generation is a minor term) err=5.551e-17 (tol 1.0e-12) \n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 6. Equivalence (TOST) versus \"no significant difference\"; minimum detectable effect\n", + "`delta_hat ~ N(delta, s^2)`. Equivalence margin `Delta`. TOST: reject `H0: |delta| >= Delta` iff both one-sided tests reject: `p_TOST = max(p_L, p_U)`.\n", + "Duality: TOST at level alpha rejects iff the (1-2 alpha) CI lies inside `(-Delta, Delta)`. Power at `delta=0`: `2 Phi(Delta/s - z_{1-alpha}) - 1`." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "source": [ + "al = 0.05; zA = stats.norm.ppf(1 - al)\n", + "def tost_p(dh, s, Dl):\n", + " p_upper = stats.norm.cdf((dh - Dl) / s) # H0: delta >= Delta\n", + " p_lower = stats.norm.sf((dh + Dl) / s) # H0: delta <= -Delta\n", + " return np.maximum(p_lower, p_upper)\n", + "rng = np.random.default_rng(5)\n", + "dh = rng.normal(0, 0.03, 100000); s_ = rng.uniform(0.005, 0.03, 100000); Dl = 0.02\n", + "rej_p = tost_p(dh, s_, Dl) < al\n", + "rej_ci = (dh - zA * s_ > -Dl) & (dh + zA * s_ < Dl)\n", + "check(\"6. TOST p 90% CI inside (-Delta, Delta) (100000 random cases)\", float((rej_p != rej_ci).sum()), 0.0)\n", + "# size at the boundary and power at delta = 0\n", + "R = 400_000\n", + "for dl in (Dl, -Dl):\n", + " sm = 0.008; x = rng.normal(dl, sm, R); r = (tost_p(x, sm, Dl) < al).mean()\n", + " check(f\"6. size at delta={dl:+.2f} (boundary of H0), s=0.008: TOST rejection prob = alpha\", abs(r - al), 4 * math.sqrt(al * (1 - al) / R), f\"{r:.4f}\")\n", + "sm = 0.008; x = rng.normal(0, sm, R); r = (tost_p(x, sm, Dl) < al).mean(); pw = 2 * stats.norm.cdf(Dl / sm - zA) - 1\n", + "check(\"6. power at delta=0: 2 Phi(Delta/s - z) - 1 vs MC\", abs(r - pw), 4 * math.sqrt(0.25 / R), f\"{r:.4f} vs {pw:.4f}\")\n", + "# difference test vs equivalence at a true effect larger than the margin: N=244, discordance d=0.1\n", + "dpair = 0.1; sd_pair = math.sqrt(dpair); s244 = sd_pair / math.sqrt(244)\n", + "true_delta = 0.03\n", + "x = rng.normal(true_delta, s244, R)\n", + "p_diff = 2 * stats.norm.sf(np.abs(x) / s244)\n", + "miss = (p_diff >= 0.05).mean(); fake_equiv = (tost_p(x, s244, Dl) < al).mean()\n", + "print(f\" N=244, d=0.10 (s = {s244:.4f}), true delta=+0.03 > Delta=0.02: difference test misses it {miss*100:.1f}% of the time; TOST declares equivalence {fake_equiv*100:.2f}%\")\n", + "check(\"6. 'not significant' is not equivalence: with a true effect beyond the margin, difference test fails to reject often, TOST seldom 'accepts' (<= alpha)\", 0.0 if (miss > 0.4 and fake_equiv <= al + 4 * math.sqrt(al * 0.95 / R)) else 1.0, 0.0)\n", + "# MDE (paired, one run per benchmark) and required N, cross-checked against note 0208's table (normal approximation, d = 0.10, delta = 5/244)\n", + "Nn = 244; mde = (stats.norm.ppf(0.975) + stats.norm.ppf(0.8)) * sd_pair / math.sqrt(Nn)\n", + "x = rng.normal(mde, s244, R); pw = (2 * stats.norm.sf(np.abs(x) / s244) < 0.05).mean()\n", + "check(\"6. MDE = (z_{.975}+z_{.8}) sigma_d/sqrt(N) gives 80% power (Gaussian paired mean)\", abs(pw - 0.8), 4 * math.sqrt(0.16 / R) + 0.0035, f\"MDE = {mde*Nn:.2f} problems, power {pw:.4f} (0208 exact trinomial: 14.3)\")\n", + "dlt = 5 / 244; d_ = 0.10\n", + "N_norm = ((stats.norm.ppf(0.975) * math.sqrt(d_) + stats.norm.ppf(0.8) * math.sqrt(d_ - dlt ** 2)) / dlt) ** 2\n", + "check(\"6. N for 80% power at delta=5/244, d=.10 equals note 0208's normal-approximation value 1867\", abs(round(N_norm) - 1867), 1, f\"N={N_norm:.1f}\")\n", + "Req = (( stats.norm.ppf(1 - al) + stats.norm.ppf(0.9)) * sd_pair / Dl) ** 2 # TOST 80% power at delta=0: Delta/s - z = z_{0.9}\n", + "pw_req = 2 * stats.norm.cdf(Dl / (sd_pair / math.sqrt(Req)) - zA) - 1\n", + "check(\"6. N for 80% TOST power at delta=0, Delta=0.02, d=.10: closed form reproduces 80% power\", abs(pw_req - 0.8), 1e-9, f\"N = {Req:.0f}\")\n", + "print(f\" (equivalence within 2 points needs N ~ {Req:.0f} problems at 10% discordance; detecting a 5-problem gain needs ~1867)\")" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "[ok] 6. TOST p 90% CI inside (-Delta, Delta) (100000 random cases) err=0.000e+00 (tol 0.0e+00) \n", + "[ok] 6. size at delta=+0.02 (boundary of H0), s=0.008: TOST rejection prob = alpha err=4.675e-04 (tol 1.4e-03) 0.0495\n", + "[ok] 6. size at delta=-0.02 (boundary of H0), s=0.008: TOST rejection prob = alpha err=1.425e-04 (tol 1.4e-03) 0.0499\n", + "[ok] 6. power at delta=0: 2 Phi(Delta/s - z) - 1 vs MC err=7.526e-04 (tol 3.2e-03) 0.6083 vs 0.6075\n", + " N=244, d=0.10 (s = 0.0202), true delta=+0.03 > Delta=0.02: difference test misses it 68.4% of the time; TOST declares equivalence 0.00%\n", + "[ok] 6. 'not significant' is not equivalence: with a true effect beyond the margin, difference test fails to reject often, TOST seldom 'accepts' (<= alpha) err=0.000e+00 (tol 0.0e+00) \n", + "[ok] 6. MDE = (z_{.975}+z_{.8}) sigma_d/sqrt(N) gives 80% power (Gaussian paired mean) err=5.500e-04 (tol 6.0e-03) MDE = 13.84 problems, power 0.7994 (0208 exact trinomial: 14.3)\n", + "[ok] 6. N for 80% power at delta=5/244, d=.10 equals note 0208's normal-approximation value 1867 err=0.000e+00 (tol 1.0e+00) N=1866.8\n", + "[ok] 6. N for 80% TOST power at delta=0, Delta=0.02, d=.10: closed form reproduces 80% power err=2.220e-16 (tol 1.0e-09) N = 2141\n", + " (equivalence within 2 points needs N ~ 2141 problems at 10% discordance; detecting a 5-problem gain needs ~1867)\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Summary" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": {}, + "source": [ + "for n_, e_, t_ in CHECKS: print(f\"PASS {n_[:105]}\")\n", + "print(f\"\\n{len(CHECKS)} checks, all passed\")" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "PASS 1. components are non-negative and sum to Var(p)\n", + "PASS 1. Var(Y_s) over 6000 replications vs sA2/N + sB2 + sC2/N + se2/(Nm)\n", + "PASS 1. ANOVA estimator of sigma_e^2 is unbiased (4000 reps, S=4)\n", + "PASS 1. ANOVA estimator of sigma_C^2 is unbiased (4000 reps, S=4)\n", + "PASS 1. ANOVA estimator of sigma_A^2 is unbiased (4000 reps, S=4)\n", + "PASS 1. ANOVA estimator of sigma_B^2 is unbiased (4000 reps, S=4)\n", + "PASS 2. Var(delta_hat) for S=1 seeds per recipe vs (2/S)[sB2 + sC2/N + se2/(Nm)]\n", + "PASS 2. Var(delta_hat) for S=3 seeds per recipe vs (2/S)[sB2 + sC2/N + se2/(Nm)]\n", + "PASS 2. floor is strictly positive: MDE(N=1e9, m=1e9, S=3) equals z sqrt(2 sB2 / 3)\n", + "PASS 3. E[e(n,c,k)] = 1-(1-p)^k exactly (6 (n,k,p) triples, rational arithmetic)\n", + "PASS 3. k=1: Var = p(1-p)/n exactly\n", + "PASS 3. k=n: Var = q(1-q) with q = 1-(1-p)^n exactly\n", + "PASS 3. naive plug-in is biased low at (n=64,k=32,p=.02) by more than 0.07\n", + "PASS 3. exact Var vs 2e6 Monte Carlo draws (n=64,k=8,p=.1)\n", + "PASS 3. Var of the benchmark mean pass@8 (N=244, n=64) vs (1/N)[Var_x g + E Var(e|p)]\n", + "PASS 4. sign-flip p on {-1,0,+1} equals exact McNemar for all b,c <= 7 (3 concordant zeros)\n", + "PASS 4. (5,0): two-sided sign-flip p = 0.0625\n", + "PASS 4. Monte Carlo p-value (B=20000) vs exact enumeration at N=14\n", + "PASS 4. size of the sign-flip test under H0 (N=60 problems, shared p_x, 1500 benchmarks)\n", + "PASS 4. skewed mean-zero null: size differs from 0.05 by more than 0.02 (the test is not a test of the mean)\n", + "PASS 5. c_8 = E max of 8 N(0,1): quadrature vs 2e6 Monte Carlo\n", + "PASS 5(i). same-sample winner estimate bias vs sigma_tot c_K\n", + "PASS 5(ii). fresh generations on the same problems: bias vs (sb2/sigma_tot) c_K\n", + "PASS 5(iii). fresh problems: bias 0\n", + "PASS 5. retained fraction of the bias after regenerating = sb2/(sb2+sg2) (with m=16: generation is a minor ter\n", + "PASS 6. TOST p 90% CI inside (-Delta, Delta) (100000 random cases)\n", + "PASS 6. size at delta=+0.02 (boundary of H0), s=0.008: TOST rejection prob = alpha\n", + "PASS 6. size at delta=-0.02 (boundary of H0), s=0.008: TOST rejection prob = alpha\n", + "PASS 6. power at delta=0: 2 Phi(Delta/s - z) - 1 vs MC\n", + "PASS 6. 'not significant' is not equivalence: with a true effect beyond the margin, difference test fails to r\n", + "PASS 6. MDE = (z_{.975}+z_{.8}) sigma_d/sqrt(N) gives 80% power (Gaussian paired mean)\n", + "PASS 6. N for 80% power at delta=5/244, d=.10 equals note 0208's normal-approximation value 1867\n", + "PASS 6. N for 80% TOST power at delta=0, Delta=0.02, d=.10: closed form reproduces 80% power\n", + "\n", + "33 checks, all passed\n" + ] + } + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.12.11" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} \ No newline at end of file diff --git a/code/notebooks/0210-data-selection-contamination.ipynb b/code/notebooks/0210-data-selection-contamination.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..dfbec0608edc76bfa87f587e9fc63b303536ecd3 --- /dev/null +++ b/code/notebooks/0210-data-selection-contamination.ipynb @@ -0,0 +1,500 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# 0210 - Data selection and contamination\n", + "Companion to `notes/0210 - Data Selection and Contamination.pdf`. CPU only, numpy/scipy, 1 thread.\n", + "1. classifier-based harvesting (fastText-style) as a selection problem; 2. pass-rate band selection as importance weighting;\n", + "3. regression to the mean when the same samples select and evaluate; 4. MinHash / Jaccard estimation error, LSH, containment; 5. membership inference from NLL gaps." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "source": [ + "import os\n", + "os.environ[\"OMP_NUM_THREADS\"] = \"1\"; os.environ[\"OPENBLAS_NUM_THREADS\"] = \"1\"\n", + "import math, itertools\n", + "from math import comb\n", + "import numpy as np\n", + "from scipy import stats, special\n", + "CHECKS = []\n", + "def check(name, err, tol, note=\"\"):\n", + " CHECKS.append((name, float(err), float(tol)))\n", + " assert err <= tol, (name, float(err), tol)\n", + " print(f\"[ok] {name:104s} err={float(err):.3e} (tol {tol:.1e}) {note}\")\n", + "Phibar = stats.norm.sf" + ], + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 1. A harvesting classifier is a detector; what you get depends on prevalence and on what a wrong page costs\n", + "Pool of `M` pages, a fraction `pi` relevant (\"math\"). Score `s ~ N(d', 1)` for relevant, `N(0, 1)` otherwise (binormal; AUC = Phi(d'/sqrt 2)).\n", + "Threshold `t`: TPR = Phibar(t - d'), FPR = Phibar(t). Precision = pi TPR / (pi TPR + (1-pi) FPR) -- depends on `pi`, not only on AUC.\n", + "Downstream value (assumed, to be *measured* by a training run): each selected relevant page gives `+g`, each irrelevant `-h` (dilution).\n", + "Net gain `G(t) = M [pi TPR g - (1-pi) FPR h]`. `G'(t) = 0` where the posterior `P(rel | s=t) = h/(g+h)`: a posterior cutoff, so `t* = d'/2 + ln[(1-pi) h /(pi g)]/d'`." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "source": [ + "rng = np.random.default_rng(0)\n", + "dprime, pi, M = 2.5, 0.01, 4_000_000\n", + "rel = rng.random(M) < pi\n", + "s = rng.standard_normal(M) + dprime * rel\n", + "t = 2.0\n", + "sel = s >= t\n", + "prec_mc = rel[sel].mean(); rec_mc = sel[rel].mean()\n", + "tpr, fpr = Phibar(t - dprime), Phibar(t)\n", + "prec_th = pi * tpr / (pi * tpr + (1 - pi) * fpr)\n", + "check(\"1. precision closed form vs Monte Carlo (4e6 pages)\", abs(prec_mc - prec_th), 4 * math.sqrt(prec_th * (1 - prec_th) / sel.sum()), f\"prec {prec_th:.4f}, recall {tpr:.4f}\")\n", + "check(\"1. recall closed form vs Monte Carlo\", abs(rec_mc - tpr), 4 * math.sqrt(tpr * (1 - tpr) / rel.sum()))\n", + "auc = stats.norm.cdf(dprime / math.sqrt(2))\n", + "print(f\" AUC={auc:.4f} (d'={dprime}); at a threshold with recall {tpr:.2f}: precision {prec_th:.3f} for pi={pi}; for pi=0.2 the same classifier gives {0.2*tpr/(0.2*tpr+0.8*fpr):.3f}\")\n", + "check(\"1. same classifier, same threshold: precision at pi=0.2 exceeds precision at pi=0.01 by > 0.4\", 0.0 if (0.2 * tpr / (0.2 * tpr + 0.8 * fpr)) - prec_th > 0.4 else 1.0, 0.0)\n", + "g_, h_ = 1.0, 1.0\n", + "ts = np.linspace(0, 6, 60001)\n", + "G = pi * Phibar(ts - dprime) * g_ - (1 - pi) * Phibar(ts) * h_\n", + "t_star = dprime / 2 + math.log((1 - pi) * h_ / (pi * g_)) / dprime\n", + "check(\"1. argmax of the net-gain curve vs closed-form posterior-cutoff threshold\", abs(ts[G.argmax()] - t_star), 2e-4, f\"t*={t_star:.4f}\")\n", + "post = lambda t: pi * stats.norm.pdf(t - dprime) / (pi * stats.norm.pdf(t - dprime) + (1 - pi) * stats.norm.pdf(t))\n", + "check(\"1. posterior P(rel | s=t*) = h/(g+h)\", abs(post(t_star) - h_ / (g_ + h_)), 1e-12)\n", + "def f1(t):\n", + " tp = pi * Phibar(t - dprime); fp = (1 - pi) * Phibar(t); fn = pi - tp\n", + " return 2 * tp / (2 * tp + fp + fn)\n", + "t_f1 = ts[np.array([f1(x) for x in ts]).argmax()]\n", + "Gof = lambda t: pi * Phibar(t - dprime) * g_ - (1 - pi) * Phibar(t) * h_\n", + "print(f\" max-F1 threshold {t_f1:.3f} vs max-gain threshold {t_star:.3f}; gain/M at F1-opt {Gof(t_f1):.3e} vs gain-opt {Gof(t_star):.3e}\")\n", + "check(\"1. the F1-optimal threshold is not the gain-optimal one (differ by > 0.2) and loses downstream gain\", 0.0 if (abs(t_f1 - t_star) > 0.2 and Gof(t_f1) < Gof(t_star)) else 1.0, 0.0)\n", + "# seed coverage and iteration: positives split into 'covered' (d_in) and 'uncovered' (d_out) clusters; each round adds a fraction gamma of uncovered clusters to the seed\n", + "d_in, d_out, gamma, tt = 3.0, 0.5, 0.5, 2.0\n", + "rec = []; cov = 0.3\n", + "for r in range(1, 6):\n", + " rec.append(cov * Phibar(tt - d_in) + (1 - cov) * Phibar(tt - d_out)); cov = cov + gamma * (1 - cov)\n", + "inc = np.diff(rec)\n", + "print(\" recall by round:\", np.round(rec, 4), \" increments:\", np.round(inc, 4), \" ratio of successive increments:\", np.round(inc[1:] / inc[:-1], 3))\n", + "check(\"1. iterative re-seeding: recall rises monotonically and increments shrink geometrically with ratio (1-gamma)\", max(abs(inc[1:] / inc[:-1] - (1 - gamma))), 1e-12)" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "[ok] 1. precision closed form vs Monte Carlo (4e6 pages) err=1.455e-04 (tol 4.9e-03) prec 0.2349, recall 0.6915\n", + "[ok] 1. recall closed form vs Monte Carlo err=2.719e-03 (tol 9.3e-03) \n", + " AUC=0.9615 (d'=2.5); at a threshold with recall 0.69: precision 0.235 for pi=0.01; for pi=0.2 the same classifier gives 0.884\n", + "[ok] 1. same classifier, same threshold: precision at pi=0.2 exceeds precision at pi=0.01 by > 0.4 err=0.000e+00 (tol 0.0e+00) \n", + "[ok] 1. argmax of the net-gain curve vs closed-form posterior-cutoff threshold err=4.794e-05 (tol 2.0e-04) t*=3.0880\n", + "[ok] 1. posterior P(rel | s=t*) = h/(g+h) err=1.110e-16 (tol 1.0e-12) \n", + " max-F1 threshold 2.626 vs max-gain threshold 3.088; gain/M at F1-opt 2.230e-04 vs gain-opt 1.785e-03\n", + "[ok] 1. the F1-optimal threshold is not the gain-optimal one (differ by > 0.2) and loses downstream gain err=0.000e+00 (tol 0.0e+00) \n", + " recall by round: [0.2992 0.5703 0.7058 0.7736 0.8075] increments: [0.2711 0.1355 0.0678 0.0339] ratio of successive increments: [0.5 0.5 0.5]\n", + "[ok] 1. iterative re-seeding: recall rises monotonically and increments shrink geometrically with ratio (1-gamma) err=6.661e-16 (tol 1.0e-12) \n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 2. Pass-rate band selection is importance weighting of the policy gradient\n", + "Bernoulli problem with logit parameter, `n` samples, `c` successes, `p_hat = c/n`. GRPO-style mean baseline gives per-problem gradient magnitude `ghat = p_hat (1 - p_hat)` (note 0206), `E ghat = (1 - 1/n) p(1-p)`.\n", + "Keeping the problem iff `p_hat` lies in a band `[lo, hi]` that is computed from the SAME samples multiplies the gradient by `1{p_hat in band}`:\n", + "effective weight on a problem of true pass rate p: `w_same(p) = E[ghat 1{band}] / (p(1-p))`. With an independent pilot estimate `p_hat'` instead: `w_ind(p) = (1 - 1/n) P[p_hat' in band]`." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "source": [ + "def bin_pmf(n, p): return np.array([comb(n, c) * p ** c * (1 - p) ** (n - c) for c in range(n + 1)])\n", + "n = 8; lo, hi = 1 / 8, 6 / 8\n", + "cs = np.arange(n + 1); ph = cs / n; inb = ((ph >= lo) & (ph <= hi)).astype(float); ghat = ph * (1 - ph)\n", + "rows = []\n", + "for p in (0.02, 0.1, 0.3, 0.5, 0.7, 0.9, 0.98):\n", + " pm = bin_pmf(n, p)\n", + " w_same = (pm * ghat * inb).sum() / (p * (1 - p)); w_ind = (1 - 1 / n) * (pm * inb).sum()\n", + " check_unb = abs((pm * ghat).sum() - (1 - 1 / n) * p * (1 - p))\n", + " assert check_unb < 1e-15\n", + " rows.append((p, w_same, w_ind)); print(f\" p={p:4.2f}: w_same={w_same:.4f} w_ind={w_ind:.4f} ratio={w_same/w_ind:.3f}\")\n", + "check(\"2. E[ghat] = (1-1/n) p(1-p) exactly (7 values of p)\", 0.0, 0.0)\n", + "check(\"2. same-sample selection weights exceed independent-pilot weights for every p in the table (selection on the statistic that defines the gradient)\", 0.0 if all(w1 > w2 for _, w1, w2 in rows) else 1.0, 0.0)\n", + "# MC check of w_same by simulating samples\n", + "p = 0.1; cs_mc = rng.binomial(n, p, size=2_000_000); phm = cs_mc / n\n", + "est = (phm * (1 - phm) * ((phm >= lo) & (phm <= hi))).mean() / (p * (1 - p)); th = [r for r in rows if r[0] == p][0][1]\n", + "check(\"2. w_same(p=0.1) exact vs 2e6 Monte Carlo\", abs(est / th - 1), 0.01)\n", + "# which problems does the band actually admit? contamination of the band by noise in p_hat (pool p ~ Beta(0.6,1.4), n = 8)\n", + "a_, b_ = 0.6, 1.4; band_p = (0.125, 0.75)\n", + "px = rng.beta(a_, b_, size=1_000_000); cx = rng.binomial(n, px); sel_ = (cx / n >= lo) & (cx / n <= hi)\n", + "frac_out = ((px[sel_] < band_p[0] - 1e-12) | (px[sel_] > band_p[1] + 1e-12)).mean()\n", + "# exact via beta-binomial\n", + "pmf_bb = np.array([comb(n, c) * special.beta(c + a_, n - c + b_) / special.beta(a_, b_) for c in range(n + 1)])\n", + "print(f\" fraction selected: {sel_.mean():.4f} (exact {pmf_bb[(cs>=1)&(cs<=6)].sum():.4f}); of the selected, true p outside [{band_p[0]},{band_p[1]}]: {frac_out:.3f}\")\n", + "check(\"2. beta-binomial exact P[band] vs MC\", abs(sel_.mean() - pmf_bb[(cs >= 1) & (cs <= 6)].sum()), 4 * math.sqrt(0.25 / 1e6))\n", + "check(\"2. a substantial share of selected problems lie outside the band in TRUE pass rate (> 15%)\", 0.0 if frac_out > 0.15 else 1.0, 0.0, f\"{frac_out:.3f}\")" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + " p=0.02: w_same=0.8750 w_ind=0.1306 ratio=6.701\n", + " p=0.10: w_same=0.8750 w_ind=0.4983 ratio=1.756\n", + " p=0.30: w_same=0.8744 w_ind=0.8234 ratio=1.062\n", + " p=0.50: w_same=0.8613 w_ind=0.8408 ratio=1.024\n", + " p=0.70: w_same=0.7721 w_ind=0.6516 ratio=1.185\n", + " p=0.90: w_same=0.4100 w_ind=0.1635 ratio=2.507\n", + " p=0.98: w_same=0.0999 w_ind=0.0090 ratio=11.044\n", + "[ok] 2. E[ghat] = (1-1/n) p(1-p) exactly (7 values of p) err=0.000e+00 (tol 0.0e+00) \n", + "[ok] 2. same-sample selection weights exceed independent-pilot weights for every p in the table (selection on the statistic that defines the gradient) err=0.000e+00 (tol 0.0e+00) \n", + "[ok] 2. w_same(p=0.1) exact vs 2e6 Monte Carlo err=6.072e-04 (tol 1.0e-02) \n", + " fraction selected: 0.6231 (exact 0.6233); of the selected, true p outside [0.125,0.75]: 0.209\n", + "[ok] 2. beta-binomial exact P[band] vs MC err=2.182e-04 (tol 2.0e-03) \n", + "[ok] 2. a substantial share of selected problems lie outside the band in TRUE pass rate (> 15%) err=0.000e+00 (tol 0.0e+00) 0.209\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 3. Regression to the mean: the same samples cannot select and evaluate\n", + "Pool `p ~ Beta(a, b)`, `c | p ~ Bin(n, p)`. Posterior `p | c ~ Beta(a + c, b + n - c)`, mean `(a + c)/(a + b + n)`.\n", + "Select the \"unsolved\" problems `c = 0`: observed rate 0, true mean `a/(a+b+n)` -- re-sampling them WITHOUT any training shows an apparent gain." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "source": [ + "a_, b_, n = 0.6, 1.4, 8\n", + "post_mean = lambda c: (a_ + c) / (a_ + b_ + n)\n", + "px = rng.beta(a_, b_, size=3_000_000); c1 = rng.binomial(n, px); c2 = rng.binomial(n, px) # c2: fresh samples, NO training between\n", + "for cc in (0, 1, 2):\n", + " selc = c1 == cc\n", + " emp = c2[selc].mean() / n; th = post_mean(cc)\n", + " check(f\"3. E[fresh pass rate | first-pass count c={cc}] = (a+c)/(a+b+n)\", abs(emp - th), 4 * math.sqrt(th * (1 - th) / (selc.sum() * n)), f\"observed first-pass rate {cc/n:.3f} -> fresh {emp:.4f} (theory {th:.4f}); apparent 'gain' {emp - cc/n:+.4f} with no training\")\n", + "# difference-in-differences with an untrained control from the same selection\n", + "tau = 0.03; sel0 = c1 == 0\n", + "idx = np.flatnonzero(sel0); half = rng.random(idx.size) < 0.5 # random split of the selected problems into trained / untrained\n", + "p_after = np.where(half, np.minimum(px[idx] + tau * (1 - px[idx]), 1.0), px[idx]) # training raises p by tau (1-p) on the trained half only\n", + "c_after = rng.binomial(n, p_after) / n\n", + "naive = c_after[half].mean() - 0.0; control = c_after[~half].mean() - 0.0\n", + "did = naive - control\n", + "tau_eff = (tau * (1 - px[idx[half]])).mean()\n", + "se = math.sqrt(c_after[half].var() / half.sum() / n + c_after[~half].var() / (~half).sum() / n) + 1e-9\n", + "print(f\" trained-half before/after = {naive:+.4f} (true effect {tau_eff:+.4f}); untrained control 'gain' = {control:+.4f}; DiD = {did:+.4f}\")\n", + "check(\"3. naive before/after overstates the effect (regression to the mean) and DiD recovers it\", abs(did - tau_eff), 0.003, f\"naive {naive:+.4f}, DiD {did:+.4f}, true {tau_eff:+.4f}\")\n", + "check(\"3. naive before/after is inflated by RTM: excess over truth exceeds 0.03\", 0.0 if (naive - tau_eff) > 0.03 else 1.0, 0.0)\n", + "# split-sample selection: select with set 1, evaluate with independent set 2 -> E[fresh | selected] is what an untrained re-evaluation gives (no gain)" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "[ok] 3. E[fresh pass rate | first-pass count c=0] = (a+c)/(a+b+n) err=8.957e-05 (tol 3.6e-04) observed first-pass rate 0.000 -> fresh 0.0599 (theory 0.0600); apparent 'gain' +0.0599 with no training\n", + "[ok] 3. E[fresh pass rate | first-pass count c=1] = (a+c)/(a+b+n) err=3.516e-04 (tol 7.3e-04) observed first-pass rate 0.125 -> fresh 0.1604 (theory 0.1600); apparent 'gain' +0.0354 with no training\n", + "[ok] 3. E[fresh pass rate | first-pass count c=2] = (a+c)/(a+b+n) err=1.404e-04 (tol 1.0e-03) observed first-pass rate 0.250 -> fresh 0.2599 (theory 0.2600); apparent 'gain' +0.0099 with no training\n", + " trained-half before/after = +0.0882 (true effect +0.0282); untrained control 'gain' = +0.0600; DiD = +0.0282\n", + "[ok] 3. naive before/after overstates the effect (regression to the mean) and DiD recovers it err=9.032e-06 (tol 3.0e-03) naive +0.0882, DiD +0.0282, true +0.0282\n", + "[ok] 3. naive before/after is inflated by RTM: excess over truth exceeds 0.03 err=0.000e+00 (tol 0.0e+00) \n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 4. MinHash: unbiased Jaccard with variance J(1-J)/k; LSH banding; containment is a different quantity\n", + "`Jhat = (1/k) sum_i 1{minhash_i(A) = minhash_i(B)}`. For a uniformly random permutation `P[min equal] = J`, so `k Jhat ~ Bin(k, J)`: unbiased, `Var = J(1-J)/k`." + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "source": [ + "def jaccard(A, B): return len(A & B) / len(A | B)\n", + "U = 3000; rng = np.random.default_rng(7)\n", + "A = set(rng.choice(U, 400, replace=False).tolist()); share = 240\n", + "B = set(list(A)[:share]) | set(rng.choice([u for u in range(U) if u not in A], 160, replace=False).tolist())\n", + "J = jaccard(A, B)\n", + "def minhash_perm(S, perms): return np.array([min(perm[x] for x in S) for perm in perms])\n", + "k = 128; reps = 3000\n", + "perm_bank = np.array([rng.permutation(U) for _ in range(k)]) # one fresh set of k permutations per replication below\n", + "ests = []\n", + "Al = np.array(sorted(A)); Bl = np.array(sorted(B))\n", + "for r in range(reps):\n", + " perms = rng.permuted(np.tile(np.arange(U), (k, 1)), axis=1)\n", + " ests.append((perms[:, Al].min(1) == perms[:, Bl].min(1)).mean())\n", + "ests = np.array(ests)\n", + "check(\"4. E[Jhat] = J (random permutations, k=128, 3000 replications)\", abs(ests.mean() - J), 4 * math.sqrt(J * (1 - J) / k / reps), f\"J={J:.4f}\")\n", + "check(\"4. Var(Jhat) = J(1-J)/k\", abs(ests.var(ddof=1) / (J * (1 - J) / k) - 1), 4 * math.sqrt(2 / (reps - 1)), f\"emp {ests.var(ddof=1):.3e} vs {J*(1-J)/k:.3e}\")\n", + "# practical universal hash h(x) = (a x + b) mod p (not truly random permutations)\n", + "P = 2_147_483_647\n", + "def uh_est(A_, B_, k, rg):\n", + " a = rg.integers(1, P, k); b = rg.integers(0, P, k)\n", + " Aa = np.array(sorted(A_), dtype=np.int64); Bb = np.array(sorted(B_), dtype=np.int64)\n", + " ha = ((a[:, None] * Aa[None, :] + b[:, None]) % P).min(1); hb = ((a[:, None] * Bb[None, :] + b[:, None]) % P).min(1)\n", + " return (ha == hb).mean()\n", + "rg = np.random.default_rng(11)\n", + "est_uh = np.array([uh_est(A, B, 128, rg) for _ in range(3000)])\n", + "# a linear hash is pairwise independent but NOT min-wise independent: a small bias is allowed and reported, not hidden\n", + "check(\"4. universal-hash MinHash (practical): |bias| < 0.01 (small but need not be zero; reported)\", abs(est_uh.mean() - J), 0.01, f\"bias {est_uh.mean()-J:+.4f} (MC s.e. {math.sqrt(J*(1-J)/128/3000):.4f})\")\n", + "for eps in (0.05, 0.02, 0.01):\n", + " kreq = math.ceil(stats.norm.ppf(0.975) ** 2 * 0.25 / eps ** 2); print(f\" worst-case (J=1/2) hashes for +-{eps} at 95%: k >= {kreq}\")\n", + "kk = 385; cover = 0.0\n", + "pm = stats.binom.pmf(np.arange(kk + 1), kk, 0.5); jh = np.arange(kk + 1) / kk\n", + "cover = pm[np.abs(jh - 0.5) <= 0.05].sum()\n", + "check(\"4. k=385 gives >= ~95% coverage of +-0.05 at J=1/2 (exact binomial)\", 0.0 if cover > 0.94 else 1.0, 0.0, f\"coverage {cover:.4f}\")\n", + "# decision at a decontamination threshold: flag if Jhat >= 0.8, k = 128\n", + "kk = 128\n", + "for Jt in (0.7, 0.75, 0.8, 0.85):\n", + " pflag = stats.binom.sf(math.ceil(0.8 * kk) - 1, kk, Jt)\n", + " print(f\" true J={Jt}: P[flag at Jhat>=0.8] = {pflag:.3f} (sd of Jhat {math.sqrt(Jt*(1-Jt)/kk):.3f})\")\n", + "Jt = 0.75; mc = (rng.binomial(kk, Jt, 2_000_000) / kk >= 0.8).mean(); ex = stats.binom.sf(math.ceil(0.8 * kk) - 1, kk, Jt)\n", + "check(\"4. exact flag probability at J=0.75 vs Monte Carlo\", abs(mc - ex), 4 * math.sqrt(ex * (1 - ex) / 2e6), f\"P[flag]={ex:.3f}: a 'near-duplicate' just under the threshold is flagged often\")\n", + "# LSH banding: b bands of r rows: P[candidate] = 1 - (1 - J^r)^b\n", + "b_, r_ = 20, 5\n", + "for Jt in (0.3, 0.5, 0.6, 0.8):\n", + " th = 1 - (1 - Jt ** r_) ** b_\n", + " bits = rng.random((400_000, b_, r_)) < Jt # each hash agrees w.p. J independently (random permutations)\n", + " mc = bits.all(2).any(1).mean()\n", + " check(f\"4. LSH banding P[candidate] at J={Jt}: 1-(1-J^r)^b vs Monte Carlo (b=20, r=5)\", abs(mc - th), 4 * math.sqrt(th * (1 - th) / 400_000) + 1e-4, f\"{th:.4f}\")\n", + "print(f\" LSH threshold approx (1/b)^(1/r) = {(1/b_)**(1/r_):.3f}\")\n", + "# containment: a benchmark item (40 shingles) inside a long web page (2000 shingles): J tiny, containment 1\n", + "Aitem = set(range(40)); Bpage = set(range(2000))\n", + "Jc = jaccard(Aitem, Bpage); cont = len(Aitem & Bpage) / len(Aitem)\n", + "inter_from_J = Jc * (len(Aitem) + len(Bpage)) / (1 + Jc)\n", + "check(\"4. |A n B| = J(|A|+|B|)/(1+J): containment recoverable from Jaccard and the two sizes\", abs(inter_from_J - 40), 1e-9, f\"J={Jc:.4f} (flag threshold 0.8 would miss it), containment={cont:.2f}\")\n", + "# error amplification of the containment estimate via Jhat\n", + "kk = 128; Jh = rng.binomial(kk, Jc, 100_000) / kk\n", + "cont_hat = Jh * (40 + 2000) / (1 + Jh) / 40\n", + "print(f\" containment estimate from MinHash at k=128: mean {cont_hat.mean():.3f}, sd {cont_hat.std():.3f} (true 1.0): very noisy for tiny J\")\n", + "check(\"4. containment from MinHash at k=128 for tiny J has relative sd > 0.5\", 0.0 if cont_hat.std() > 0.5 else 1.0, 0.0, f\"sd {cont_hat.std():.2f}\")\n", + "# real text shingles: document with edits\n", + "rg = np.random.default_rng(5)\n", + "vocab = [f\"w{i}\" for i in range(3000)]\n", + "doc = list(rg.choice(vocab, 600)); doc2 = doc.copy()\n", + "for pos in rg.choice(600, 36, replace=False): doc2[pos] = rg.choice(vocab)\n", + "sh = lambda d, w=5: {\" \".join(d[i:i + w]) for i in range(len(d) - w + 1)}\n", + "S1, S2 = sh(doc), sh(doc2); Jx = jaccard(S1, S2)\n", + "hs = lambda S: {int.from_bytes(__import__('hashlib').blake2b(x.encode(), digest_size=8).digest(), 'little') % P for x in S}\n", + "H1, H2 = hs(S1), hs(S2)\n", + "est = np.array([uh_est(H1, H2, 256, rg) for _ in range(400)])\n", + "check(\"4. 5-word shingles of a 600-word doc with 6% word edits: MinHash(k=256) mean vs exact Jaccard\", abs(est.mean() - Jx), 4 * math.sqrt(Jx * (1 - Jx) / 256 / 400) + 0.003, f\"exact J={Jx:.4f}, est {est.mean():.4f}, sd {est.std():.4f}\")" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "[ok] 4. E[Jhat] = J (random permutations, k=128, 3000 replications) err=4.416e-04 (tol 3.2e-03) J=0.4286\n", + "[ok] 4. Var(Jhat) = J(1-J)/k err=3.114e-02 (tol 1.0e-01) emp 1.973e-03 vs 1.913e-03\n", + "[ok] 4. universal-hash MinHash (practical): |bias| < 0.01 (small but need not be zero; reported) err=3.666e-03 (tol 1.0e-02) bias +0.0037 (MC s.e. 0.0008)\n", + " worst-case (J=1/2) hashes for +-0.05 at 95%: k >= 385\n", + " worst-case (J=1/2) hashes for +-0.02 at 95%: k >= 2401\n", + " worst-case (J=1/2) hashes for +-0.01 at 95%: k >= 9604\n", + "[ok] 4. k=385 gives >= ~95% coverage of +-0.05 at J=1/2 (exact binomial) err=0.000e+00 (tol 0.0e+00) coverage 0.9474\n", + " true J=0.7: P[flag at Jhat>=0.8] = 0.005 (sd of Jhat 0.041)\n", + " true J=0.75: P[flag at Jhat>=0.8] = 0.090 (sd of Jhat 0.038)\n", + " true J=0.8: P[flag at Jhat>=0.8] = 0.500 (sd of Jhat 0.035)\n", + " true J=0.85: P[flag at Jhat>=0.8] = 0.936 (sd of Jhat 0.032)\n", + "[ok] 4. exact flag probability at J=0.75 vs Monte Carlo err=1.523e-04 (tol 8.1e-04) P[flag]=0.090: a 'near-duplicate' just under the threshold is flagged often\n", + "[ok] 4. LSH banding P[candidate] at J=0.3: 1-(1-J^r)^b vs Monte Carlo (b=20, r=5) err=2.643e-04 (tol 1.4e-03) 0.0475\n", + "[ok] 4. LSH banding P[candidate] at J=0.5: 1-(1-J^r)^b vs Monte Carlo (b=20, r=5) err=2.007e-04 (tol 3.3e-03) 0.4701\n", + "[ok] 4. LSH banding P[candidate] at J=0.6: 1-(1-J^r)^b vs Monte Carlo (b=20, r=5) err=2.825e-04 (tol 2.6e-03) 0.8019\n", + "[ok] 4. LSH banding P[candidate] at J=0.8: 1-(1-J^r)^b vs Monte Carlo (b=20, r=5) err=2.394e-05 (tol 2.2e-04) 0.9996\n", + " LSH threshold approx (1/b)^(1/r) = 0.549\n", + "[ok] 4. |A n B| = J(|A|+|B|)/(1+J): containment recoverable from Jaccard and the two sizes err=0.000e+00 (tol 1.0e-09) J=0.0200 (flag threshold 0.8 would miss it), containment=1.00\n", + " containment estimate from MinHash at k=128: mean 0.997, sd 0.603 (true 1.0): very noisy for tiny J\n", + "[ok] 4. containment from MinHash at k=128 for tiny J has relative sd > 0.5 err=0.000e+00 (tol 0.0e+00) sd 0.60\n", + "[ok] 4. 5-word shingles of a 600-word doc with 6% word edits: MinHash(k=256) mean vs exact Jaccard err=2.862e-03 (tol 9.2e-03) exact J=0.5872, est 0.5844, sd 0.0322\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 5. Membership inference from NLL gaps\n", + "Members have lower NLL: `NLL_member ~ N(mu - Delta, sigma^2)`, non-member `N(mu, sigma^2)`. AUC of the threshold test = `Phi(Delta / (sigma sqrt 2))`; TPR at FPR `f`: `Phi(Delta/sigma - z_{1-f})`.\n", + "If examples differ in intrinsic difficulty `D_i ~ N(0, sD^2)` (added to every NLL), the raw-NLL AUC drops to `Phi(Delta/sqrt(2(sD^2+sigma^2)))`; a reference-calibrated score `NLL_target - NLL_ref` cancels `D_i`." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "source": [ + "rng = np.random.default_rng(21); R_ = 400_000\n", + "Delta, sig, sD, sref = 0.30, 0.25, 0.80, 0.25\n", + "mem = rng.normal(-Delta, sig, R_); non = rng.normal(0, sig, R_)\n", + "def auc_of(pos_low, neg):\n", + " # P[score(member) < score(non-member)] with MC estimate by ranks on a subsample\n", + " x = pos_low[:200000]; y = neg[:200000]\n", + " r = stats.rankdata(np.concatenate([x, y])); return 1 - (r[:len(x)].sum() - len(x) * (len(x) + 1) / 2) / (len(x) * len(y))\n", + "check(\"5. AUC = Phi(Delta/(sigma sqrt2)) (Gaussian NLL gap)\", abs(auc_of(mem, non) - stats.norm.cdf(Delta / (sig * math.sqrt(2)))), 0.003, f\"{stats.norm.cdf(Delta/(sig*math.sqrt(2))):.4f}\")\n", + "thr = np.quantile(non, 0.01); tpr_mc = (mem <= thr).mean(); tpr_th = stats.norm.cdf(Delta / sig - stats.norm.ppf(0.99))\n", + "check(\"5. TPR at 1% FPR = Phi(Delta/sigma - z_{0.99})\", abs(tpr_mc - tpr_th), 0.004, f\"{tpr_th:.4f}\")\n", + "Dm = rng.normal(0, sD, R_); Dn = rng.normal(0, sD, R_)\n", + "raw_m = Dm + mem; raw_n = Dn + non\n", + "auc_raw_th = stats.norm.cdf(Delta / math.sqrt(2 * (sD ** 2 + sig ** 2)))\n", + "check(\"5. difficulty heterogeneity: raw-NLL AUC = Phi(Delta/sqrt(2(sD^2+sigma^2)))\", abs(auc_of(raw_m, raw_n) - auc_raw_th), 0.003, f\"{auc_raw_th:.4f}\")\n", + "ref_m = Dm + rng.normal(0, sref, R_); ref_n = Dn + rng.normal(0, sref, R_)\n", + "cal_th = stats.norm.cdf(Delta / math.sqrt(2 * (sig ** 2 + sref ** 2)))\n", + "check(\"5. calibrated NLL - NLL_ref: AUC = Phi(Delta/sqrt(2(sigma^2+sref^2)))\", abs(auc_of(raw_m - ref_m, raw_n - ref_n) - cal_th), 0.003, f\"{cal_th:.4f}\")\n", + "print(f\" AUC: raw {auc_raw_th:.3f} -> calibrated {cal_th:.3f} (Delta={Delta}, sigma={sig}, sD={sD}, sref={sref})\")\n", + "# a real (tiny) language model: character n-gram, memorisation grows with order; ground-truth membership known\n", + "V = 12; L = 120; ndoc = 2000; alpha = 0.1; rng = np.random.default_rng(33)\n", + "logits = rng.normal(0, 1.5, (V, V, V)) # order-2 Markov source\n", + "def gen_docs(nd, rg):\n", + " taus = np.exp(rg.uniform(math.log(0.4), math.log(2.5), nd)); X = np.zeros((nd, L), dtype=np.int64); X[:, :2] = rg.integers(0, V, (nd, 2))\n", + " for t in range(2, L):\n", + " lg = logits[X[:, t - 2], X[:, t - 1]] / taus[:, None]; pr = np.exp(lg - lg.max(1, keepdims=True)); pr /= pr.sum(1, keepdims=True)\n", + " u = rg.random(nd)[:, None]; X[:, t] = (pr.cumsum(1) < u).sum(1).clip(max=V - 1)\n", + " return X, taus\n", + "Xm, tm = gen_docs(ndoc, rng); Xn, tn = gen_docs(ndoc, rng); Xr, _ = gen_docs(ndoc, rng)\n", + "def fit(X, order):\n", + " cnt = np.zeros((V ** order, V))\n", + " for t in range(4, L):\n", + " ctx = np.zeros(len(X), dtype=np.int64)\n", + " for j in range(order, 0, -1): ctx = ctx * V + X[:, t - j]\n", + " np.add.at(cnt, (ctx, X[:, t]), 1)\n", + " return cnt\n", + "def nll(cnt, X, order):\n", + " out = np.zeros(len(X))\n", + " for t in range(4, L):\n", + " ctx = np.zeros(len(X), dtype=np.int64)\n", + " for j in range(order, 0, -1): ctx = ctx * V + X[:, t - j]\n", + " c = cnt[ctx]; out -= np.log((c[np.arange(len(X)), X[:, t]] + alpha) / (c.sum(1) + alpha * V))\n", + " return out / (L - 4)\n", + "def auc_small(low, high): # P[low < high] (+0.5 ties)\n", + " r = stats.rankdata(np.concatenate([low, high])); n1 = len(low); return 1 - (r[:n1].sum() - n1 * (n1 + 1) / 2) / (n1 * len(high))\n", + "res = {}\n", + "for order in (1, 2, 3, 4):\n", + " cm = fit(Xm, order); cr = fit(Xr, order)\n", + " nm, nn = nll(cm, Xm, order), nll(cm, Xn, order)\n", + " rm, rn = nll(cr, Xm, order), nll(cr, Xn, order)\n", + " res[order] = (auc_small(nm, nn), auc_small(nm - rm, nn - rn), nn.mean() - nm.mean())\n", + " print(f\" order {order}: raw-NLL AUC = {res[order][0]:.3f}, calibrated AUC = {res[order][1]:.3f}, mean NLL gap (non - member) = {res[order][2]:.4f} nats/char\")\n", + "check(\"5. no memorisation capacity at order 1-2 (<= true order): raw AUC within 0.04 of 0.5\", max(abs(res[1][0] - 0.5), abs(res[2][0] - 0.5)), 0.04)\n", + "check(\"5. membership signal grows with model capacity: raw AUC order 4 > order 3 > order 2\", 0.0 if res[4][0] > res[3][0] > res[2][0] else 1.0, 0.0)\n", + "check(\"5. reference calibration helps at the high-capacity order: calibrated AUC > raw AUC + 0.05 (order 4)\", 0.0 if res[4][1] > res[4][0] + 0.05 else 1.0, 0.0, f\"{res[4][0]:.3f} -> {res[4][1]:.3f}\")\n", + "check(\"5. calibrated AUC increases with order and exceeds raw AUC at every order >= 2 (it detects even the small order-2 overfit that raw NLL hides)\", 0.0 if (res[1][1] < res[2][1] < res[3][1] < res[4][1] and all(res[o][1] > res[o][0] for o in (2, 3, 4))) else 1.0, 0.0)" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "[ok] 5. AUC = Phi(Delta/(sigma sqrt2)) (Gaussian NLL gap) err=3.136e-05 (tol 3.0e-03) 0.8019\n", + "[ok] 5. TPR at 1% FPR = Phi(Delta/sigma - z_{0.99}) err=1.543e-03 (tol 4.0e-03) 0.1300\n", + "[ok] 5. difficulty heterogeneity: raw-NLL AUC = Phi(Delta/sqrt(2(sD^2+sigma^2))) err=1.669e-04 (tol 3.0e-03) 0.5999\n", + "[ok] 5. calibrated NLL - NLL_ref: AUC = Phi(Delta/sqrt(2(sigma^2+sref^2))) err=7.603e-04 (tol 3.0e-03) 0.7257\n", + " AUC: raw 0.600 -> calibrated 0.726 (Delta=0.3, sigma=0.25, sD=0.8, sref=0.25)\n", + " order 1: raw-NLL AUC = 0.506, calibrated AUC = 0.555, mean NLL gap (non - member) = 0.0019 nats/char\n", + " order 2: raw-NLL AUC = 0.510, calibrated AUC = 0.675, mean NLL gap (non - member) = 0.0154 nats/char\n", + " order 3: raw-NLL AUC = 0.559, calibrated AUC = 0.832, mean NLL gap (non - member) = 0.1048 nats/char\n", + " order 4: raw-NLL AUC = 0.715, calibrated AUC = 0.923, mean NLL gap (non - member) = 0.4656 nats/char\n", + "[ok] 5. no memorisation capacity at order 1-2 (<= true order): raw AUC within 0.04 of 0.5 err=1.007e-02 (tol 4.0e-02) \n", + "[ok] 5. membership signal grows with model capacity: raw AUC order 4 > order 3 > order 2 err=0.000e+00 (tol 0.0e+00) \n", + "[ok] 5. reference calibration helps at the high-capacity order: calibrated AUC > raw AUC + 0.05 (order 4) err=0.000e+00 (tol 0.0e+00) 0.715 -> 0.923\n", + "[ok] 5. calibrated AUC increases with order and exceeds raw AUC at every order >= 2 (it detects even the small order-2 overfit that raw NLL hides) err=0.000e+00 (tol 0.0e+00) \n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Summary" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "source": [ + "for n_, e_, t_ in CHECKS: print(f\"PASS {n_[:110]}\")\n", + "print(f\"\\n{len(CHECKS)} checks, all passed\")" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "PASS 1. precision closed form vs Monte Carlo (4e6 pages)\n", + "PASS 1. recall closed form vs Monte Carlo\n", + "PASS 1. same classifier, same threshold: precision at pi=0.2 exceeds precision at pi=0.01 by > 0.4\n", + "PASS 1. argmax of the net-gain curve vs closed-form posterior-cutoff threshold\n", + "PASS 1. posterior P(rel | s=t*) = h/(g+h)\n", + "PASS 1. the F1-optimal threshold is not the gain-optimal one (differ by > 0.2) and loses downstream gain\n", + "PASS 1. iterative re-seeding: recall rises monotonically and increments shrink geometrically with ratio (1-gamma)\n", + "PASS 2. E[ghat] = (1-1/n) p(1-p) exactly (7 values of p)\n", + "PASS 2. same-sample selection weights exceed independent-pilot weights for every p in the table (selection on the s\n", + "PASS 2. w_same(p=0.1) exact vs 2e6 Monte Carlo\n", + "PASS 2. beta-binomial exact P[band] vs MC\n", + "PASS 2. a substantial share of selected problems lie outside the band in TRUE pass rate (> 15%)\n", + "PASS 3. E[fresh pass rate | first-pass count c=0] = (a+c)/(a+b+n)\n", + "PASS 3. E[fresh pass rate | first-pass count c=1] = (a+c)/(a+b+n)\n", + "PASS 3. E[fresh pass rate | first-pass count c=2] = (a+c)/(a+b+n)\n", + "PASS 3. naive before/after overstates the effect (regression to the mean) and DiD recovers it\n", + "PASS 3. naive before/after is inflated by RTM: excess over truth exceeds 0.03\n", + "PASS 4. E[Jhat] = J (random permutations, k=128, 3000 replications)\n", + "PASS 4. Var(Jhat) = J(1-J)/k\n", + "PASS 4. universal-hash MinHash (practical): |bias| < 0.01 (small but need not be zero; reported)\n", + "PASS 4. k=385 gives >= ~95% coverage of +-0.05 at J=1/2 (exact binomial)\n", + "PASS 4. exact flag probability at J=0.75 vs Monte Carlo\n", + "PASS 4. LSH banding P[candidate] at J=0.3: 1-(1-J^r)^b vs Monte Carlo (b=20, r=5)\n", + "PASS 4. LSH banding P[candidate] at J=0.5: 1-(1-J^r)^b vs Monte Carlo (b=20, r=5)\n", + "PASS 4. LSH banding P[candidate] at J=0.6: 1-(1-J^r)^b vs Monte Carlo (b=20, r=5)\n", + "PASS 4. LSH banding P[candidate] at J=0.8: 1-(1-J^r)^b vs Monte Carlo (b=20, r=5)\n", + "PASS 4. |A n B| = J(|A|+|B|)/(1+J): containment recoverable from Jaccard and the two sizes\n", + "PASS 4. containment from MinHash at k=128 for tiny J has relative sd > 0.5\n", + "PASS 4. 5-word shingles of a 600-word doc with 6% word edits: MinHash(k=256) mean vs exact Jaccard\n", + "PASS 5. AUC = Phi(Delta/(sigma sqrt2)) (Gaussian NLL gap)\n", + "PASS 5. TPR at 1% FPR = Phi(Delta/sigma - z_{0.99})\n", + "PASS 5. difficulty heterogeneity: raw-NLL AUC = Phi(Delta/sqrt(2(sD^2+sigma^2)))\n", + "PASS 5. calibrated NLL - NLL_ref: AUC = Phi(Delta/sqrt(2(sigma^2+sref^2)))\n", + "PASS 5. no memorisation capacity at order 1-2 (<= true order): raw AUC within 0.04 of 0.5\n", + "PASS 5. membership signal grows with model capacity: raw AUC order 4 > order 3 > order 2\n", + "PASS 5. reference calibration helps at the high-capacity order: calibrated AUC > raw AUC + 0.05 (order 4)\n", + "PASS 5. calibrated AUC increases with order and exceeds raw AUC at every order >= 2 (it detects even the small orde\n", + "\n", + "37 checks, all passed\n" + ] + } + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.12.11" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} \ No newline at end of file diff --git a/code/notebooks/0301-curry-howard.ipynb b/code/notebooks/0301-curry-howard.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..a81d99358b855e285b5ab8dea7e494f606f77f2a --- /dev/null +++ b/code/notebooks/0301-curry-howard.ipynb @@ -0,0 +1,773 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0999126c", + "metadata": {}, + "source": [ + "# 0301 - Curry-Howard: propositions as types, proofs as terms\n", + "\n", + "Scratch-pad for note `0301 - Curry-Howard.pdf`. Same discipline as the Flash-Attention-101 scratch-pads:\n", + "the **naive** implementation first, then the **better** one, then an assertion that they agree.\n", + "\n", + "Plan\n", + "1. A 40-line kernel for the simply typed lambda calculus with `->` and `False` (`infer`, `normalize`).\n", + "2. Subject reduction (Boxed result 2) checked on thousands of random typed terms.\n", + "3. Inhabitation, NAIVE: depth-bounded enumeration of beta-normal terms. EFFICIENT: Dyckhoff's contraction-free\n", + " calculus G4ip with proof-term extraction. Assert the verdicts agree on every formula with <= 4 arrows.\n", + "4. Soundness (provable => classically valid, provable => forced in every Kripke chain) and the gap\n", + " (Peirce, double-negation elimination: classically valid, not inhabited).\n", + "5. What the Lean kernel actually sees: `#print axioms` for each way of \"proving\" something; the axiom gate.\n", + "\n", + "No GPU is used (`CUDA_VISIBLE_DEVICES` is cleared). Lean calls go through `align.verify.lean` with ONE worker, nice 15." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "d045e611", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:53:55.822833Z", + "iopub.status.busy": "2026-10-14T09:53:55.822602Z", + "iopub.status.idle": "2026-10-14T09:53:55.843322Z", + "shell.execute_reply": "2026-10-14T09:53:55.842326Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "K : A → B → A\n", + "S : (A → B → C) → (A → B) → A → C\n" + ] + } + ], + "source": [ + "import os\n", + "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"\"\n", + "import itertools, random, time, math, sys\n", + "from collections import Counter\n", + "random.seed(0)\n", + "\n", + "# Types: ('atom', name) | ('bot',) | ('arr', A, B) Terms: ('var',x) | ('lam',x,A,t) | ('app',f,u) | ('abort',A,t)\n", + "BOT = ('bot',)\n", + "def atom(n): return ('atom', n)\n", + "def arr(a, b): return ('arr', a, b)\n", + "def show_ty(t):\n", + " if t[0] == 'atom': return t[1]\n", + " if t[0] == 'bot': return '⊥'\n", + " l = show_ty(t[1]); l = f'({l})' if t[1][0] == 'arr' else l\n", + " return f'{l} → {show_ty(t[2])}'\n", + "def show_tm(t):\n", + " if t[0] == 'var': return t[1]\n", + " if t[0] == 'lam': return f'(λ{t[1]}:{show_ty(t[2])}. {show_tm(t[3])})'\n", + " if t[0] == 'app': return f'({show_tm(t[1])} {show_tm(t[2])})'\n", + " return f'(abort[{show_ty(t[1])}] {show_tm(t[2])})'\n", + "\n", + "class TypeErr(Exception): pass\n", + "\n", + "def infer(ctx, t):\n", + " \"\"\"Typing judgement ctx |- t : A (ctx: dict var->type). Rules: Var, ->I, ->E, bot-E.\"\"\"\n", + " k = t[0]\n", + " if k == 'var':\n", + " if t[1] not in ctx: raise TypeErr(f'unbound {t[1]}')\n", + " return ctx[t[1]]\n", + " if k == 'lam':\n", + " _, x, A, body = t\n", + " return arr(A, infer({**ctx, x: A}, body))\n", + " if k == 'app':\n", + " F, U = infer(ctx, t[1]), infer(ctx, t[2])\n", + " if F[0] != 'arr' or F[1] != U: raise TypeErr(f'cannot apply {show_ty(F)} to {show_ty(U)}')\n", + " return F[2]\n", + " if k == 'abort':\n", + " if infer(ctx, t[2]) != BOT: raise TypeErr('abort needs a proof of ⊥')\n", + " return t[1]\n", + " raise TypeErr(k)\n", + "\n", + "_fresh = itertools.count()\n", + "KINDS = Counter() # which reduction rule fired (coverage of the SR check)\n", + "def fv(t):\n", + " k = t[0]\n", + " if k == 'var': return {t[1]}\n", + " if k == 'lam': return fv(t[3]) - {t[1]}\n", + " if k == 'app': return fv(t[1]) | fv(t[2])\n", + " return fv(t[2])\n", + "\n", + "def subst(t, x, s):\n", + " \"\"\"Capture-avoiding substitution t[x := s].\"\"\"\n", + " k = t[0]\n", + " if k == 'var': return s if t[1] == x else t\n", + " if k == 'app': return ('app', subst(t[1], x, s), subst(t[2], x, s))\n", + " if k == 'abort': return ('abort', t[1], subst(t[2], x, s))\n", + " _, y, A, b = t\n", + " if y == x: return t\n", + " if y in fv(s):\n", + " z = f'{y}_{next(_fresh)}'\n", + " b = subst(b, y, ('var', z)); y = z\n", + " return ('lam', y, A, subst(b, x, s))\n", + "\n", + "def step(t):\n", + " \"\"\"One leftmost-outermost step: beta, and the commuting conversion (abort[A->B] e) u ~> abort[B] e. None if normal.\"\"\"\n", + " k = t[0]\n", + " if k == 'app':\n", + " f, u = t[1], t[2]\n", + " if f[0] == 'lam': KINDS['beta'] += 1; return subst(f[3], f[1], u)\n", + " if f[0] == 'abort' and f[1][0] == 'arr': KINDS['comm'] += 1; return ('abort', f[1][2], f[2])\n", + " r = step(f)\n", + " if r is not None: return ('app', r, u)\n", + " r = step(u)\n", + " return None if r is None else ('app', f, r)\n", + " if k == 'lam':\n", + " r = step(t[3]); return None if r is None else ('lam', t[1], t[2], r)\n", + " if k == 'abort':\n", + " r = step(t[2]); return None if r is None else ('abort', t[1], r)\n", + " return None\n", + "\n", + "def normalize(t, fuel=10_000):\n", + " for _ in range(fuel):\n", + " r = step(t)\n", + " if r is None: return t\n", + " t = r\n", + " raise RuntimeError('fuel exhausted (would contradict strong normalisation)')\n", + "\n", + "# the classic combinators, as terms, with their types read off by the kernel\n", + "A, B, C = atom('A'), atom('B'), atom('C')\n", + "K = ('lam','x',A,('lam','y',B,('var','x')))\n", + "S = ('lam','f',arr(A,arr(B,C)),('lam','g',arr(A,B),('lam','x',A,('app',('app',('var','f'),('var','x')),('app',('var','g'),('var','x'))))))\n", + "for name, tm in [('K', K), ('S', S)]:\n", + " print(name, ':', show_ty(infer({}, tm)))" + ] + }, + { + "cell_type": "markdown", + "id": "057302f6", + "metadata": {}, + "source": [ + "## 2. Subject reduction (Boxed result 2): if `Gamma |- t : A` and `t ~> t'` then `Gamma |- t' : A`\n", + "Check on random well-typed terms: generate by random type-directed construction, reduce step by step, re-infer at every step." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "5787fcb2", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:53:55.845545Z", + "iopub.status.busy": "2026-10-14T09:53:55.845323Z", + "iopub.status.idle": "2026-10-14T09:53:56.359730Z", + "shell.execute_reply": "2026-10-14T09:53:56.358821Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "3000 random typed terms, 1205 with >=1 redex, 1851 reduction steps (1431 beta, 420 commuting-conversion); subject reduction violations: 0 (0.5s)\n" + ] + } + ], + "source": [ + "ATOMS = ['a', 'b', 'c']\n", + "def rand_type(d):\n", + " if d == 0 or random.random() < 0.35:\n", + " return random.choice([atom(x) for x in ATOMS] + [BOT])\n", + " return arr(rand_type(d - 1), rand_type(d - 1))\n", + "\n", + "def spine_end(T):\n", + " while T[0] == 'arr': T = T[2]\n", + " return T\n", + "\n", + "def rand_term(ctx, goal, d):\n", + " \"\"\"random term of type `goal` in `ctx`, possibly with redexes (introduces cuts).\"\"\"\n", + " if d > 0 and random.random() < 0.35: # a cut: (λx:T. body) arg\n", + " T = rand_type(2)\n", + " arg = rand_term(ctx, T, d - 1)\n", + " x = f'v{next(_fresh)}'\n", + " body = rand_term({**ctx, x: T}, goal, d - 1)\n", + " return None if (arg is None or body is None) else ('app', ('lam', x, T, body), arg)\n", + " if d > 0 and random.random() < 0.15 and any(spine_end(T) == BOT for T in ctx.values()): # ex falso at an ARROW type, then applied\n", + " T = rand_type(2)\n", + " e = rand_term(ctx, BOT, d - 1); arg = rand_term(ctx, T, d - 1)\n", + " if e is not None and arg is not None: return ('app', ('abort', arr(T, goal), e), arg)\n", + " if goal[0] == 'arr':\n", + " x = f'v{next(_fresh)}'\n", + " body = rand_term({**ctx, x: goal[1]}, goal[2], d - 1)\n", + " return None if body is None else ('lam', x, goal[1], body)\n", + " # atomic goal: use a hypothesis head with matching result, else give up\n", + " cands = []\n", + " for x, T in ctx.items():\n", + " spine, R = [], T\n", + " while R[0] == 'arr': spine.append(R[1]); R = R[2]\n", + " if R == goal or (R == BOT):\n", + " cands.append((x, spine, R))\n", + " if not cands: return None\n", + " x, spine, R = random.choice(cands)\n", + " args = []\n", + " for Ai in spine:\n", + " a = rand_term(ctx, Ai, d - 1) if d > 0 else None\n", + " if a is None: return None\n", + " args.append(a)\n", + " tm = ('var', x)\n", + " for a in args: tm = ('app', tm, a)\n", + " return tm if R == goal else ('abort', goal, tm)\n", + "\n", + "n_ok = n_steps = n_nontrivial = 0\n", + "KINDS.clear()\n", + "t0 = time.time()\n", + "while n_ok < 3000:\n", + " ctx = {f'h{i}': rand_type(2) for i in range(random.randint(0, 3))}\n", + " goal = rand_type(2)\n", + " tm = rand_term(ctx, goal, 4)\n", + " if tm is None: continue\n", + " ty = infer(ctx, tm) # must typecheck by construction\n", + " assert ty == goal\n", + " cur, steps = tm, 0\n", + " while True:\n", + " nxt = step(cur)\n", + " if nxt is None: break\n", + " assert infer(ctx, nxt) == goal, ('subject reduction violated', show_tm(cur), show_tm(nxt))\n", + " cur, steps = nxt, steps + 1\n", + " n_steps += 1\n", + " n_nontrivial += steps > 0\n", + " n_ok += 1\n", + "print(f'{n_ok} random typed terms, {n_nontrivial} with >=1 redex, {n_steps} reduction steps '\n", + " f'({KINDS[\"beta\"]} beta, {KINDS[\"comm\"]} commuting-conversion); subject reduction violations: 0 ({time.time()-t0:.1f}s)')\n", + "assert KINDS['comm'] > 0 and KINDS['beta'] > 0 # both cases of Theorem 2 are actually exercised" + ] + }, + { + "cell_type": "markdown", + "id": "fc48f5ca", + "metadata": {}, + "source": [ + "## 3. Inhabitation: naive vs efficient\n", + "**Naive** (`search_naive`): enumerate beta-normal, eta-long terms by depth, backtracking. Correct but exponential, and it can only answer\n", + "\"found within depth D\" / \"not found within depth D\".\n", + "\n", + "**Efficient** (`g4ip`): Dyckhoff's contraction-free sequent calculus (terminating, so a *negative* answer is a proof of\n", + "non-provability), with the proof term extracted and re-checked by the independent kernel `infer`." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "b941a9af", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:53:56.362552Z", + "iopub.status.busy": "2026-10-14T09:53:56.362347Z", + "iopub.status.idle": "2026-10-14T09:53:56.376383Z", + "shell.execute_reply": "2026-10-14T09:53:56.375472Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "K a→b→a proof term: (λy0:a. (λy1:b. y0))\n", + "S proof term: (λy2:a → b → c. (λy3:a → b. (λy4:a. ((y2 y4) (y3 y4)))))\n", + "contraposition (a→b)→¬b→¬a proof term: (λy8:a → b. (λy9:b → ⊥. (λy10:a. (abort[⊥] (y9 (y8 y10))))))\n", + "a → ¬¬a proof term: (λy13:a. (λy14:a → ⊥. (abort[⊥] (y14 y13))))\n", + "¬¬¬a → ¬a proof term: (λy16:((a → ⊥) → ⊥) → ⊥. (λy17:a. (abort[⊥] (y16 (λc18:a → ⊥. (abort[⊥] (c18 y17)))))))\n", + "ex falso ⊥→a proof term: (λy24:⊥. (abort[a] y24))\n", + "Peirce ((a→b)→a)→a UNPROVABLE (G4ip terminates: this is a refutation, no term exists)\n", + "DNE ¬¬a→a UNPROVABLE (G4ip terminates: this is a refutation, no term exists)\n", + "rev. contraposition (¬b→¬a)→a→b UNPROVABLE (G4ip terminates: this is a refutation, no term exists)\n" + ] + } + ], + "source": [ + "def spine_of(T):\n", + " sp = []\n", + " while T[0] == 'arr': sp.append(T[1]); T = T[2]\n", + " return sp, T\n", + "\n", + "def search_naive(ctx, goal, depth):\n", + " \"\"\"Return a normal term or None. ctx: tuple of (var,type).\"\"\"\n", + " if depth <= 0: return None\n", + " if goal[0] == 'arr':\n", + " x = f'x{len(ctx)}'\n", + " body = search_naive(ctx + ((x, goal[1]),), goal[2], depth - 1)\n", + " return None if body is None else ('lam', x, goal[1], body)\n", + " for x, T in ctx:\n", + " sp, R = spine_of(T)\n", + " if R != goal and R != BOT: continue\n", + " args, okk = [], True\n", + " for Ai in sp:\n", + " a = search_naive(ctx, Ai, depth - 1)\n", + " if a is None: okk = False; break\n", + " args.append(a)\n", + " if not okk: continue\n", + " tm = ('var', x)\n", + " for a in args: tm = ('app', tm, a)\n", + " return tm if R == goal else ('abort', goal, tm)\n", + " return None\n", + "\n", + "_nf = itertools.count()\n", + "def g4ip(ctx, goal, memo_fail):\n", + " \"\"\"ctx: tuple of (var,type). Returns a proof term or None. Contraction-free: terminates.\"\"\"\n", + " key = (frozenset(T for _, T in ctx), goal)\n", + " if key in memo_fail: return None\n", + " res = _g4ip(ctx, goal, memo_fail)\n", + " if res is None: memo_fail.add(key)\n", + " return res\n", + "\n", + "def _g4ip(ctx, goal, mf):\n", + " for x, T in ctx: # L⊥\n", + " if T == BOT: return ('abort', goal, ('var', x))\n", + " for x, T in ctx: # Ax\n", + " if T == goal and T[0] == 'atom': return ('var', x)\n", + " if goal[0] == 'arr': # R→ (invertible: no backtracking needed)\n", + " y = f'y{next(_nf)}'\n", + " b = g4ip(ctx + ((y, goal[1]),), goal[2], mf)\n", + " return None if b is None else ('lam', y, goal[1], b)\n", + " # L0→ (invertible): p, p→B in ctx ==> replace p→B by B\n", + " for i, (f, T) in enumerate(ctx):\n", + " if T[0] == 'arr' and T[1][0] == 'atom':\n", + " for x, Tx in ctx:\n", + " if Tx == T[1]:\n", + " z = f'z{next(_nf)}'\n", + " rest = ctx[:i] + ctx[i+1:] + ((z, T[2]),)\n", + " t = g4ip(rest, goal, mf)\n", + " return None if t is None else subst(t, z, ('app', ('var', f), ('var', x)))\n", + " # L⊥→ (invertible): drop hypotheses ⊥→B (carry no information)\n", + " for i, (f, T) in enumerate(ctx):\n", + " if T[0] == 'arr' and T[1] == BOT:\n", + " return g4ip(ctx[:i] + ctx[i+1:], goal, mf)\n", + " # L→→ (NOT invertible: backtrack over the choice of hypothesis)\n", + " for i, (f, T) in enumerate(ctx):\n", + " if T[0] == 'arr' and T[1][0] == 'arr':\n", + " C_, D_ = T[1][1], T[1][2]; B_ = T[2]\n", + " c, g, z = f'c{next(_nf)}', f'g{next(_nf)}', f'z{next(_nf)}'\n", + " base = ctx[:i] + ctx[i+1:]\n", + " t1 = g4ip(base + ((c, C_), (g, arr(D_, B_))), D_, mf)\n", + " if t1 is None: continue\n", + " t2 = g4ip(base + ((z, B_),), goal, mf)\n", + " if t2 is None: continue\n", + " d = f'd{next(_nf)}'\n", + " u = ('lam', c, C_, subst(t1, g, ('lam', d, D_, ('app', ('var', f), ('lam', f'_{next(_nf)}', C_, ('var', d))))))\n", + " return subst(t2, z, ('app', ('var', f), u))\n", + " return None\n", + "\n", + "def prove(goal):\n", + " return g4ip((), goal, set())\n", + "\n", + "# Gentzen-style sanity on named formulas\n", + "a, b, c = atom('a'), atom('b'), atom('c')\n", + "neg = lambda p: arr(p, BOT)\n", + "named = {\n", + " 'K a→b→a': arr(a, arr(b, a)),\n", + " 'S': arr(arr(a, arr(b, c)), arr(arr(a, b), arr(a, c))),\n", + " 'contraposition (a→b)→¬b→¬a': arr(arr(a, b), arr(neg(b), neg(a))),\n", + " 'a → ¬¬a': arr(a, neg(neg(a))),\n", + " '¬¬¬a → ¬a': arr(neg(neg(neg(a))), neg(a)),\n", + " 'ex falso ⊥→a': arr(BOT, a),\n", + " 'Peirce ((a→b)→a)→a': arr(arr(arr(a, b), a), a),\n", + " 'DNE ¬¬a→a': arr(neg(neg(a)), a),\n", + " 'rev. contraposition (¬b→¬a)→a→b': arr(arr(neg(b), neg(a)), arr(a, b)),\n", + "}\n", + "for name, f in named.items():\n", + " tm = prove(f)\n", + " if tm is None:\n", + " print(f'{name:40s} UNPROVABLE (G4ip terminates: this is a refutation, no term exists)')\n", + " else:\n", + " assert infer({}, tm) == f # independent kernel check of the extracted proof\n", + " nf = normalize(tm); assert infer({}, nf) == f\n", + " print(f'{name:40s} proof term: {show_tm(nf)}')" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "3ec7fb35", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:53:56.378933Z", + "iopub.status.busy": "2026-10-14T09:53:56.378753Z", + "iopub.status.idle": "2026-10-14T09:53:57.467575Z", + "shell.execute_reply": "2026-10-14T09:53:57.466497Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "formulas: {0: 3, 1: 9, 2: 54, 3: 405, 4: 3402} total 3873\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "G4ip: 0.03s naive(depth 14): 1.06s speed-up x41.7\n", + "inhabited: 1775 of 3873 | verdict disagreements naive vs G4ip: 0\n" + ] + } + ], + "source": [ + "# All implicational formulas over {a, b, ⊥} with exactly k arrows, k = 0..4\n", + "LEAVES = [atom('a'), atom('b'), BOT]\n", + "@__import__('functools').lru_cache(None)\n", + "def formulas(k):\n", + " if k == 0: return tuple(LEAVES)\n", + " out = []\n", + " for i in range(k):\n", + " for l in formulas(i):\n", + " for r in formulas(k - 1 - i):\n", + " out.append(arr(l, r))\n", + " return tuple(out)\n", + "ALL = [f for k in range(5) for f in formulas(k)]\n", + "print('formulas:', {k: len(formulas(k)) for k in range(5)}, 'total', len(ALL))\n", + "\n", + "t0 = time.time(); verdict_eff = {}\n", + "for f in ALL:\n", + " tm = prove(f)\n", + " if tm is not None:\n", + " assert infer({}, tm) == f\n", + " assert infer({}, normalize(tm)) == f\n", + " verdict_eff[f] = tm is not None\n", + "t_eff = time.time() - t0\n", + "\n", + "t0 = time.time(); verdict_naive = {}; D = 14\n", + "for f in ALL:\n", + " tm = search_naive((), f, D)\n", + " if tm is not None: assert infer({}, tm) == f\n", + " verdict_naive[f] = tm is not None\n", + "t_naive = time.time() - t0\n", + "\n", + "disagree = [f for f in ALL if verdict_eff[f] != verdict_naive[f]]\n", + "assert not disagree, [show_ty(f) for f in disagree[:5]]\n", + "print(f'G4ip: {t_eff:.2f}s naive(depth {D}): {t_naive:.2f}s speed-up x{t_naive/max(t_eff,1e-9):.1f}')\n", + "print('inhabited:', sum(verdict_eff.values()), 'of', len(ALL), '| verdict disagreements naive vs G4ip: 0')" + ] + }, + { + "cell_type": "markdown", + "id": "05df5ec0", + "metadata": {}, + "source": [ + "## 4. Soundness and the gap\n", + "* **Classical semantics** (truth table, `⊥ = False`): every inhabited type must be a tautology (Boxed result 3, soundness direction).\n", + "* **Kripke semantics** over chains of <=3 worlds: every inhabited type must be forced at *every* world of *every* model;\n", + " a non-inhabited type should have a countermodel (we report how many our tiny model class catches)." + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "8247c153", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:53:57.469657Z", + "iopub.status.busy": "2026-10-14T09:53:57.469490Z", + "iopub.status.idle": "2026-10-14T09:53:57.982720Z", + "shell.execute_reply": "2026-10-14T09:53:57.981661Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "provable 1775 all classically valid and forced in all 29 chain models: OK\n", + "unprovable 2098: classically valid but NOT inhabited = 32; with a chain countermodel (m<=3) = 2098\n", + " valid, not inhabited: ((a → b) → a) → a\n", + " valid, not inhabited: ((a → b) → ⊥) → a\n", + " valid, not inhabited: ((a → ⊥) → a) → a\n", + " valid, not inhabited: ((a → ⊥) → ⊥) → a\n", + " valid, not inhabited: ((b → a) → b) → b\n", + " valid, not inhabited: ((b → a) → ⊥) → b\n", + "Peirce countermodel: 2 worlds, a true only at world 1, b never -> world 0 does not force Peirce.\n" + ] + } + ], + "source": [ + "def classical(f, val):\n", + " if f[0] == 'atom': return val[f[1]]\n", + " if f[0] == 'bot': return False\n", + " return (not classical(f[1], val)) or classical(f[2], val)\n", + "def tautology(f):\n", + " return all(classical(f, dict(zip('ab', v))) for v in itertools.product([False, True], repeat=2))\n", + "\n", + "def forces(f, w, m, th):\n", + " \"\"\"chain of worlds 0..m-1; atom x true at w iff w >= th[x] (monotone).\"\"\"\n", + " if f[0] == 'atom': return w >= th[f[1]]\n", + " if f[0] == 'bot': return False\n", + " return all((not forces(f[1], v, m, th)) or forces(f[2], v, m, th) for v in range(w, m))\n", + "\n", + "models = [(m, dict(zip('ab', (ta, tb)))) for m in (1, 2, 3) for ta in range(m + 1) for tb in range(m + 1)]\n", + "prov = [f for f in ALL if verdict_eff[f]]\n", + "unprov = [f for f in ALL if not verdict_eff[f]]\n", + "assert all(tautology(f) for f in prov) # soundness (classical)\n", + "assert all(forces(f, w, m, th) for f in prov for (m, th) in models for w in range(m)) # soundness (Kripke)\n", + "valid_unprov = [f for f in unprov if tautology(f)]\n", + "countermodel = [f for f in unprov if any(not forces(f, 0, m, th) for (m, th) in models)]\n", + "print(f'provable {len(prov)} all classically valid and forced in all {len(models)} chain models: OK')\n", + "print(f'unprovable {len(unprov)}: classically valid but NOT inhabited = {len(valid_unprov)}; '\n", + " f'with a chain countermodel (m<=3) = {len(countermodel)}')\n", + "ex = sorted(valid_unprov, key=lambda f: len(show_ty(f)))[:6]\n", + "for f in ex: print(' valid, not inhabited:', show_ty(f))\n", + "peirce = arr(arr(arr(a, b), a), a)\n", + "assert tautology(peirce) and not verdict_eff[peirce] and not forces(peirce, 0, 2, {'a': 1, 'b': 2})\n", + "print('Peirce countermodel: 2 worlds, a true only at world 1, b never -> world 0 does not force Peirce.')" + ] + }, + { + "cell_type": "markdown", + "id": "7c345d80", + "metadata": {}, + "source": [ + "## 5. What the Lean kernel sees\n", + "Lean elaborates tactics into terms and the **kernel** re-checks the term. The only thing that can make a \"proof\" of a false\n", + "statement pass is an axiom in the term's dependency closure. `#print axioms` lists that closure.\n", + "Exactly the same tests as the strict verifier's axiom gate (`align/verify/lean/strict.py`, rule 3)." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "f9b00c3c", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:53:57.984954Z", + "iopub.status.busy": "2026-10-14T09:53:57.984774Z", + "iopub.status.idle": "2026-10-14T09:54:02.121622Z", + "shell.execute_reply": "2026-10-14T09:54:02.120529Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "theorem axioms gate\n", + "k_comb [] ACCEPT\n", + "peirce_law ['Classical.choice', 'Quot.sound', 'propext'] ACCEPT\n", + "t_omega ['Quot.sound', 'propext'] ACCEPT\n", + "t_decide ['propext'] ACCEPT\n", + "t_native ['Lean.ofReduceBool', 'propext'] REJECT\n", + "t_sorry ['sorryAx'] REJECT\n", + "t_evil ['evil'] REJECT\n", + "t_funext ['Quot.sound'] ACCEPT\n", + "t_simp ['propext'] ACCEPT\n", + "axiom gate rejects exactly {native_decide, sorry, user axiom}; accepts constructive and classical-core proofs.\n" + ] + } + ], + "source": [ + "import sys; sys.path.insert(0, os.path.abspath('.'))\n", + "from _lean_helper import Lean\n", + "L = Lean()\n", + "code = r'''\n", + "theorem k_comb (p q : Prop) : p → q → p := fun hp _ => hp\n", + "theorem peirce_law (p q : Prop) : ((p → q) → p) → p :=\n", + " fun h => (Classical.em p).elim id (fun hn => h (fun hp => absurd hp hn))\n", + "theorem t_omega (a b : Nat) : a + b = b + a := by omega\n", + "theorem t_decide : 2^10 = 1024 := by decide\n", + "theorem t_native : 2^10 = 1024 := by native_decide\n", + "theorem t_sorry : (1:Nat) = 2 := sorry\n", + "axiom evil : False\n", + "theorem t_evil : (1:Nat) = 2 := evil.elim\n", + "theorem t_funext (f g : Nat → Nat) (h : ∀ x, f x = g x) : f = g := funext h\n", + "theorem t_simp (xs : List Nat) : (xs ++ []).length = xs.length := by simp\n", + "#print axioms k_comb\n", + "#print axioms peirce_law\n", + "#print axioms t_omega\n", + "#print axioms t_decide\n", + "#print axioms t_native\n", + "#print axioms t_sorry\n", + "#print axioms t_evil\n", + "#print axioms t_funext\n", + "#print axioms t_simp\n", + "'''\n", + "r = L.run(code, keep=True)\n", + "import re\n", + "axioms = {}\n", + "for m in r['messages']:\n", + " mm = re.match(r\"'(\\w+)' (?:does not depend on any axioms|depends on axioms: \\[(.*)\\])\", m['data'])\n", + " if mm: axioms[mm.group(1)] = sorted(x.strip() for x in (mm.group(2) or '').split(',') if x.strip())\n", + "ALLOWED = {'propext', 'Quot.sound', 'Classical.choice'}\n", + "print(f'{\"theorem\":12s} {\"axioms\":45s} gate')\n", + "for k_, v in axioms.items():\n", + " print(f'{k_:12s} {str(v):45s} {\"ACCEPT\" if set(v) <= ALLOWED else \"REJECT\"}')\n", + "expected_reject = {'t_native', 't_sorry', 't_evil'}\n", + "assert {k_ for k_, v in axioms.items() if not set(v) <= ALLOWED} == expected_reject\n", + "assert axioms['k_comb'] == []\n", + "assert set(axioms['peirce_law']) == {'Classical.choice', 'propext', 'Quot.sound'}\n", + "print('axiom gate rejects exactly {native_decide, sorry, user axiom}; accepts constructive and classical-core proofs.')" + ] + }, + { + "cell_type": "markdown", + "id": "3634e461", + "metadata": {}, + "source": [ + "### Tactics are untrusted; the term is what is checked\n", + "`by` blocks are programs that *build* a term. We can print it. A wrong tactic cannot make the kernel accept a wrong term." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "db937b1f", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:54:02.123377Z", + "iopub.status.busy": "2026-10-14T09:54:02.123216Z", + "iopub.status.idle": "2026-10-14T09:54:02.185056Z", + "shell.execute_reply": "2026-10-14T09:54:02.184044Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "`by omega` elaborates to a proof TERM of 64 pretty-printed lines (a certificate checked by the kernel); a short one:\n", + "[info] line 3: theorem t_by : ∀ (p q : Prop), p ∧ q → q ∧ p :=\n", + "fun p q h =>\n", + " match h with\n", + " | ⟨a, b⟩ => ⟨b, a⟩\n", + "--- ill-typed term: the kernel/elaborator refuses ---\n", + "[error] line 1: type mismatch\n", + " Eq.refl\n", + "has type\n", + " ∀ (a : ?m.2096), a = a : Prop\n", + "but is expected to have type\n", + " 1 = 2 : Prop\n", + "--- proof irrelevance: any two proofs of a Prop are definitionally equal ---\n", + "[ok] no messages\n", + "--- sorryAx inhabits every type, including False ---\n", + "[info] line 1: axiom sorryAx.{u} : (α : Sort u) → optParam Bool false → α\n", + "[warning] line 2: declaration uses 'sorry'\n", + "[info] line 3: 'anything' depends on axioms: [sorryAx]\n", + "[sorries] ⊢ False\n" + ] + }, + { + "data": { + "text/plain": [ + "{'sorries': [{'proofState': 1,\n", + " 'pos': {'line': 2, 'column': 28},\n", + " 'goal': '⊢ False',\n", + " 'endPos': {'line': 2, 'column': 33}}],\n", + " 'messages': [{'severity': 'info',\n", + " 'pos': {'line': 1, 'column': 0},\n", + " 'endPos': {'line': 1, 'column': 6},\n", + " 'data': 'axiom sorryAx.{u} : (α : Sort u) → optParam Bool false → α'},\n", + " {'severity': 'warning',\n", + " 'pos': {'line': 2, 'column': 8},\n", + " 'endPos': {'line': 2, 'column': 16},\n", + " 'data': \"declaration uses 'sorry'\"},\n", + " {'severity': 'info',\n", + " 'pos': {'line': 3, 'column': 0},\n", + " 'endPos': {'line': 3, 'column': 6},\n", + " 'data': \"'anything' depends on axioms: [sorryAx]\"}],\n", + " 'env': 7}" + ] + }, + "execution_count": 7, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "rr = L.run(\"#print t_omega\")\n", + "n_lines = sum(m['data'].count('\\n') + 1 for m in rr['messages'])\n", + "print(f'`by omega` elaborates to a proof TERM of {n_lines} pretty-printed lines (a certificate checked by the kernel); a short one:')\n", + "L.show(r'''\n", + "theorem t_by (p q : Prop) : p ∧ q → q ∧ p := by intro ⟨a, b⟩; exact ⟨b, a⟩\n", + "#print t_by\n", + "''')\n", + "print('--- ill-typed term: the kernel/elaborator refuses ---')\n", + "L.show(\"example : (1:Nat) = 2 := ⟨⟩\")\n", + "print('--- proof irrelevance: any two proofs of a Prop are definitionally equal ---')\n", + "L.show(\"example (p : Prop) (h1 h2 : p) : h1 = h2 := rfl\")\n", + "print('--- sorryAx inhabits every type, including False ---')\n", + "L.show(\"#print sorryAx\\ntheorem anything : False := sorry\\n#print axioms anything\")" + ] + }, + { + "cell_type": "markdown", + "id": "c0b25fdd", + "metadata": {}, + "source": [ + "### Why `sorryAx` / `native_decide` / extra axioms destroy a reward\n", + "Let `R(c) = 1[kernel accepts c and axioms(c) is within A]`. With A = {propext, Quot.sound, Classical.choice}, `R(c)=1`\n", + "implies the statement is true in every model of Lean's type theory plus those axioms (soundness of the kernel; consistency of A).\n", + "If `sorryAx` (or any axiom of type `False`) is allowed, `R(c) = 1` for **every** statement, so `P(R=1 | statement false) = 1`\n", + "and the reward carries zero information: an RL policy will find it. Below: acceptance probability of the trivial exploit under three gates." + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "52d6fcf1", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:54:02.186945Z", + "iopub.status.busy": "2026-10-14T09:54:02.186790Z", + "iopub.status.idle": "2026-10-14T09:54:02.472246Z", + "shell.execute_reply": "2026-10-14T09:54:02.471199Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "false statement axioms of the `sorry` proof strict-gate no-axiom-gate(only compile)\n", + "(1:Nat) = 2 ['sorryAx'] REJECT ACCEPT\n", + "∀ n : Nat, n < 3 ['sorryAx'] REJECT ACCEPT\n", + "False ['sorryAx'] REJECT ACCEPT\n" + ] + } + ], + "source": [ + "statements_false = [\"(1:Nat) = 2\", \"∀ n : Nat, n < 3\", \"False\"]\n", + "rows = []\n", + "for st in statements_false:\n", + " rr = L.run(f\"theorem probe : {st} := sorry\\n#print axioms probe\")\n", + " ax = []\n", + " for m in rr['messages']:\n", + " mm = re.match(r\"'probe' depends on axioms: \\[(.*)\\]\", m['data'])\n", + " if mm: ax = [x.strip() for x in mm.group(1).split(',')]\n", + " rows.append((st, ax, 'ACCEPT' if set(ax) <= ALLOWED else 'REJECT', 'ACCEPT'))\n", + "print(f'{\"false statement\":22s} {\"axioms of the `sorry` proof\":22s} strict-gate no-axiom-gate(only compile)')\n", + "for st, ax, g1, g2 in rows: print(f'{st:22s} {str(ax):22s} {g1:11s} {g2}')\n", + "assert all(r_[2] == 'REJECT' for r_ in rows)\n", + "L.close()" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.11" + }, + "m3v": { + "patched": "0301" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/code/notebooks/0302-dependent-types.ipynb b/code/notebooks/0302-dependent-types.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..dbbc0004687681afd277ab1065a4e22ff1a39276 --- /dev/null +++ b/code/notebooks/0302-dependent-types.ipynb @@ -0,0 +1,762 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "00a22d32", + "metadata": {}, + "source": [ + "# 0302 - Dependent type theory: the minimum for reading Mathlib statements\n", + "\n", + "Scratch-pad for note `0302 - Dependent Types.pdf`. Naive first, better second, assert agreement.\n", + "\n", + "1. **Universe of a Pi type.** Naive: Python function `imax`-rule. Check against what Lean's kernel says for 10 types.\n", + "2. **Decidability.** Naive: Python brute force over the bound. Efficient: Lean `decide` (kernel/VM evaluation of the\n", + " `Decidable` instance). Assert the Booleans agree on 12 bounded statements; show where no instance exists.\n", + "3. **Ten miniF2F statements decoded.** Lean prints the elaborated Pi-telescope; Python (exact `Fraction`/big-int arithmetic or\n", + " fp64) shows the hypotheses are satisfiable (non-vacuous) and the conclusion holds at witnesses.\n", + "4. **Elaboration traps** (truncated subtraction, floor division, casts, `autoImplicit`), Lean vs a Python model of ℕ.\n", + "\n", + "No GPU (`CUDA_VISIBLE_DEVICES` cleared). Lean: ONE worker through `align.verify.lean`, nice 15." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "fa7036cc", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:54:20.065314Z", + "iopub.status.busy": "2026-10-14T09:54:20.065161Z", + "iopub.status.idle": "2026-10-14T09:54:24.370858Z", + "shell.execute_reply": "2026-10-14T09:54:24.370025Z" + } + }, + "outputs": [], + "source": [ + "import os\n", + "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"\"\n", + "import sys, re, json, math, random, itertools, time\n", + "from fractions import Fraction\n", + "sys.path.insert(0, os.path.abspath('.'))\n", + "from _lean_helper import Lean, HEADER\n", + "MINIF2F_HEADER = \"import Mathlib\\nimport Aesop\\n\\nset_option maxHeartbeats 400000\\n\\nopen BigOperators Real Nat Topology Rat\\n\"\n", + "L = Lean(MINIF2F_HEADER)\n", + "\n", + "def lean_msgs(code, keep=False, timeout=300):\n", + " r = L.run(code, timeout=timeout, keep=keep)\n", + " return r.get('messages', []), r" + ] + }, + { + "cell_type": "markdown", + "id": "6bd9c629", + "metadata": {}, + "source": [ + "## 1. The sort of a Pi type: `Sort (imax u v)`\n", + "Rule: if `A : Sort u` and `B x : Sort v` for `x : A` then `(x : A) → B x : Sort (imax u v)` where\n", + "`imax u 0 = 0` and `imax u v = max u v` for `v > 0`. So a Pi into `Prop` is a `Prop` whatever the domain (impredicativity)." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "a0d26ffc", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:54:24.373206Z", + "iopub.status.busy": "2026-10-14T09:54:24.373041Z", + "iopub.status.idle": "2026-10-14T09:54:24.406046Z", + "shell.execute_reply": "2026-10-14T09:54:24.405001Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Pi type imax rule Lean\n", + "∀ n : ℕ, n = n Prop Prop\n", + "ℕ → ℕ Type Type\n", + "(n : ℕ) → Fin n Type Type\n", + "∀ p : Prop, p → p Prop Prop\n", + "(α : Type) → α → α Type 1 Type 1\n", + "(α : Type) → Prop Type 1 Type 1\n", + "Prop → Prop Type Type\n", + "∀ α : Type, α = α Prop Prop\n", + "(α : Type 1) → α → α Type 2 Type 2\n", + "ℕ → Prop Type Type\n", + "imax rule agrees with the kernel on all 10 types.\n" + ] + } + ], + "source": [ + "def imax(u, v): return 0 if v == 0 else max(u, v)\n", + "def sort_name(l): return 'Prop' if l == 0 else ('Type' if l == 1 else f'Type {l-1}')\n", + "\n", + "# (lean source of the Pi type, level u of the domain's sort, level v of the codomain's sort)\n", + "PI_TABLE = [\n", + " (\"∀ n : ℕ, n = n\", 1, 0), # ℕ : Type = Sort 1 ; n = n : Prop = Sort 0\n", + " (\"ℕ → ℕ\", 1, 1),\n", + " (\"(n : ℕ) → Fin n\", 1, 1),\n", + " (\"∀ p : Prop, p → p\", 1, 0), # Prop : Type = Sort 1 ; p → p : Prop\n", + " (\"(α : Type) → α → α\", 2, 1), # Type : Type 1 = Sort 2 ; α → α : Type = Sort 1\n", + " (\"(α : Type) → Prop\", 2, 1), # Prop : Type = Sort 1\n", + " (\"Prop → Prop\", 1, 1),\n", + " (\"∀ α : Type, α = α\", 2, 0),\n", + " (\"(α : Type 1) → α → α\", 3, 2), # Type 1 : Sort 3 ; α → α : Sort 2\n", + " (\"ℕ → Prop\", 1, 1),\n", + "]\n", + "rows = []\n", + "for src, u, v in PI_TABLE:\n", + " msgs, _ = lean_msgs(f\"#check ({src})\")\n", + " data = msgs[0]['data']\n", + " m = re.search(r\": (Prop|Type\\*?|Type \\d+|Sort [^\\n]*)\\s*$\", data.strip().split('\\n')[-1])\n", + " lean_sort = m.group(1) if m else data\n", + " pred = sort_name(imax(u, v))\n", + " rows.append((src, pred, lean_sort))\n", + " assert pred == lean_sort, (src, pred, lean_sort)\n", + "print(f'{\"Pi type\":28s} {\"imax rule\":10s} Lean')\n", + "for r_ in rows: print(f'{r_[0]:28s} {r_[1]:10s} {r_[2]}')\n", + "print('imax rule agrees with the kernel on all', len(rows), 'types.')" + ] + }, + { + "cell_type": "markdown", + "id": "f2877b3a", + "metadata": {}, + "source": [ + "## 2. Decidability: `Decidable p` is data; `decide` runs it\n", + "`class Decidable (p : Prop) : Type | isFalse (h : ¬p) | isTrue (h : p)`. `decide p` is `match inst with | isTrue _ => true | isFalse _ => false`.\n", + "Soundness: `of_decide_eq_true : decide p = true → p` (cases on the instance). The tactic builds `of_decide_eq_true (Eq.refl true)`." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "f6a87131", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:54:24.407763Z", + "iopub.status.busy": "2026-10-14T09:54:24.407583Z", + "iopub.status.idle": "2026-10-14T09:54:24.427808Z", + "shell.execute_reply": "2026-10-14T09:54:24.427297Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "@of_decide_eq_true : ∀ {p : Prop} [inst : Decidable p], decide p = true → p\n", + "inductive Decidable : Prop → Type\n", + "number of parameters: 1\n", + "constructors:\n", + "Decidable.isFalse : {p : Prop} → ¬p → Decidable p\n", + "Decidable.isTrue : {p : Prop} → p → Decidable p\n", + "theorem two_plus_two : 2 + 2 = 4 :=\n", + "of_decide_eq_true (Eq.refl true)\n", + "'two_plus_two' does not depend on any axioms\n" + ] + } + ], + "source": [ + "msgs, _ = lean_msgs(\"#check @of_decide_eq_true\\n#print Decidable\\ntheorem two_plus_two : 2 + 2 = 4 := by decide\\n#print two_plus_two\\n#print axioms two_plus_two\", keep=True)\n", + "for m in msgs: print(m['data'])" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "357adce9", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:54:24.433916Z", + "iopub.status.busy": "2026-10-14T09:54:24.433656Z", + "iopub.status.idle": "2026-10-14T09:54:25.615319Z", + "shell.execute_reply": "2026-10-14T09:54:25.614598Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "--- Nat.Prime is decidable only via a *chosen* algorithm; two instances exist (the default is chosen by priority) ---\n", + "decidablePrime 7\n", + "decidablePrime1 7\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "python lean statement\n", + "True true ∀ n < 50, n * n % 4 ≠ 2\n", + "True true ∃ n < 100, n * n = 2209\n", + "True true ∀ a < 20, ∀ b < 20, a * b = b * a\n", + "True true ∀ n < (40 : ℕ), Nat.Prime (n * n + n + 41)\n", + "False false ∀ n < (41 : ℕ), Nat.Prime (n * n + n + 41)\n", + "True true ∀ n < 100, n ^ 3 % 7 = 0 ∨ n ^ 3 % 7 = 1 ∨ n ^ 3 % 7 = 6\n", + "False false ∃ x : Fin 61, ∃ y : Fin 61, 4 * (((x : ℕ) : ℤ) - 30) ^ 3 - 7 * (((y : ℕ) : ℤ) - 30) ^ 3 = 2003\n", + "False false ∀ a < 10, ∀ b < 10, a - b + b = a\n", + "True true ∀ a < 10, ∀ b < 10, (a + b) / 2 * 2 ≤ a + b\n", + "False false ∃ n < 10, 18 ∣ 3740 + n ∧ n ≠ 4\n", + "True true ∀ n < 30, 2 ^ n % 3 = if n % 2 = 0 then 1 else 2\n", + "True true ∃ n < 100, n % 3 = 2 ∧ n % 5 = 3 ∧ n % 7 = 2\n", + "python brute force == Lean decide on 12/12\n" + ] + } + ], + "source": [ + "# 12 bounded statements: (python oracle, lean statement)\n", + "def prime(n):\n", + " return n >= 2 and all(n % d for d in range(2, int(n ** .5) + 1))\n", + "def nsub(a, b): return max(0, a - b) # Lean ℕ subtraction is truncated\n", + "CASES = [\n", + " (lambda: all(n*n % 4 != 2 for n in range(50)), \"∀ n < 50, n * n % 4 ≠ 2\"),\n", + " (lambda: any(n*n == 2209 for n in range(100)), \"∃ n < 100, n * n = 2209\"),\n", + " (lambda: all(a*b == b*a for a in range(20) for b in range(20)), \"∀ a < 20, ∀ b < 20, a * b = b * a\"),\n", + " (lambda: all(prime(n*n+n+41) for n in range(40)), \"∀ n < (40 : ℕ), Nat.Prime (n * n + n + 41)\"),\n", + " (lambda: all(prime(n*n+n+41) for n in range(41)), \"∀ n < (41 : ℕ), Nat.Prime (n * n + n + 41)\"),\n", + " (lambda: all(n**3 % 7 in (0, 1, 6) for n in range(100)), \"∀ n < 100, n ^ 3 % 7 = 0 ∨ n ^ 3 % 7 = 1 ∨ n ^ 3 % 7 = 6\"),\n", + " (lambda: any(4*x**3 - 7*y**3 == 2003 for x in range(-30, 31) for y in range(-30, 31)),\n", + " \"∃ x : Fin 61, ∃ y : Fin 61, 4 * (((x : ℕ) : ℤ) - 30) ^ 3 - 7 * (((y : ℕ) : ℤ) - 30) ^ 3 = 2003\"),\n", + " (lambda: all(nsub(a, b) + b == a for a in range(10) for b in range(10)), \"∀ a < 10, ∀ b < 10, a - b + b = a\"),\n", + " (lambda: all((a + b) // 2 * 2 <= a + b for a in range(10) for b in range(10)), \"∀ a < 10, ∀ b < 10, (a + b) / 2 * 2 ≤ a + b\"),\n", + " (lambda: any((3740 + n) % 18 == 0 and n != 4 for n in range(10)), \"∃ n < 10, 18 ∣ 3740 + n ∧ n ≠ 4\"),\n", + " (lambda: all(pow(2, n) % 3 == (1 if n % 2 == 0 else 2) for n in range(30)), \"∀ n < 30, 2 ^ n % 3 = if n % 2 = 0 then 1 else 2\"),\n", + " (lambda: any(n % 3 == 2 and n % 5 == 3 and n % 7 == 2 for n in range(100)), \"∃ n < 100, n % 3 = 2 ∧ n % 5 = 3 ∧ n % 7 = 2\"),\n", + "]\n", + "print('--- Nat.Prime is decidable only via a *chosen* algorithm; two instances exist (the default is chosen by priority) ---')\n", + "for m in lean_msgs(\"#synth Decidable (Nat.Prime 7)\")[0]: print(m['data'][:160])\n", + "lean_msgs(\"attribute [instance] Nat.decidablePrime1\", keep=True) # opt in to the trial-division instance (needed below for bounded foralls)\n", + "for m in lean_msgs(\"#synth Decidable (Nat.Prime 7)\")[0]: print(m['data'][:160])\n", + "code = \"\\n\".join(f\"#eval decide ({s})\" for _, s in CASES)\n", + "msgs, _ = lean_msgs(code)\n", + "by_line = {m['pos']['line']: m['data'].strip() for m in msgs}\n", + "print(f'{\"python\":7s} {\"lean\":6s} statement')\n", + "agree = 0\n", + "for i, (fn, s) in enumerate(CASES, start=1):\n", + " py = fn(); ln = by_line.get(i)\n", + " assert ln in ('true', 'false'), (s, ln)\n", + " assert (ln == 'true') == py, (s, py, ln)\n", + " agree += 1\n", + " print(f'{str(py):7s} {ln:6s} {s}')\n", + "print(f'python brute force == Lean decide on {agree}/{len(CASES)}')" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "c58167d4", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:54:25.617088Z", + "iopub.status.busy": "2026-10-14T09:54:25.616927Z", + "iopub.status.idle": "2026-10-14T09:54:25.737634Z", + "shell.execute_reply": "2026-10-14T09:54:25.736825Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "--- elaboration order matters: the same statement without the type ascription has no instance (n : ?α is still a metavariable) ---\n", + "failed to synthesize\n", + " Decidable (∀ n < 40, (n * n + n + 41).Prime)\n", + "use `set_option diagnostics true` to get diagnostic information\n", + "--- no instance: unbounded quantifier over ℕ ---\n", + "failed to synthesize\n", + " Decidable (∀ (n : ℕ), n < 5 ∨ 5 ≤ n)\n", + "use `set_option diagnostics true` to get diagnostic information\n", + "--- no computable instance: equality on ℝ ---\n", + "tactic 'decide' failed for proposition\n", + " 2 + 2 = 4\n", + "since its 'Decidable' instance reduced to\n", + " Classical.choice ⋯\n", + "rather than to the 'isTrue' constructor.\n", + "--- but a *closed* bounded one is fine, and the real statements with binders are not decidable props at all ---\n", + "decidableBallLT 10 fun n a => n * n < 100\n", + "failed to synthesize\n", + " Decidable (∀ (x : ℝ), x ^ 2 ≥ 0)\n", + "use `set_option diagnostics true` to get diagnostic information\n" + ] + } + ], + "source": [ + "print('--- elaboration order matters: the same statement without the type ascription has no instance (n : ?α is still a metavariable) ---')\n", + "for m in lean_msgs(\"#eval decide (∀ n < 40, Nat.Prime (n * n + n + 41))\")[0]: print(m['data'][:160])\n", + "print('--- no instance: unbounded quantifier over ℕ ---')\n", + "for m in lean_msgs(\"example : ∀ n : ℕ, n < 5 ∨ 5 ≤ n := by decide\")[0]: print(m['data'][:400])\n", + "print('--- no computable instance: equality on ℝ ---')\n", + "for m in lean_msgs(\"example : (2 : ℝ) + 2 = 4 := by decide\")[0]: print(m['data'][:400])\n", + "print('--- but a *closed* bounded one is fine, and the real statements with binders are not decidable props at all ---')\n", + "for m in lean_msgs(\"#synth Decidable (∀ n < 10, n * n < 100)\")[0]: print(m['data'][:200])\n", + "for m in lean_msgs(\"#synth Decidable (∀ x : ℝ, x ^ 2 ≥ 0)\")[0]: print(m['data'][:200])" + ] + }, + { + "cell_type": "markdown", + "id": "c172fa47", + "metadata": {}, + "source": [ + "## 3. Ten miniF2F statements decoded\n", + "Source: `p2_runs/r1_baseline/minif2f.jsonl` (the 244 miniF2F test statements this repo evaluates on, Lean 4.9.0-rc1 + Mathlib pin).\n", + "For each: (a) Lean's own elaborated Pi-telescope (`#check @name`), its sort, and `decide`-ability of the closed ones;\n", + "(b) a Python witness check that the hypotheses are satisfiable and the conclusion holds there (non-vacuity), exact arithmetic where possible." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "f06bbc76", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:54:25.739310Z", + "iopub.status.busy": "2026-10-14T09:54:25.739151Z", + "iopub.status.idle": "2026-10-14T09:54:26.102413Z", + "shell.execute_reply": "2026-10-14T09:54:26.101621Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "====================================================================================================\n", + "mathd_algebra_478\n", + "mathd_algebra_478 : ∀ (b h v : ℝ), 0 < b ∧ 0 < h ∧ 0 < v → v = 1 / 3 * (b * h) → b = 30 → h = 13 / 2 → v = 65\n", + "binders=7 sort=Prop\n", + "====================================================================================================\n", + "numbertheory_4x3m7y3neq2003\n", + "numbertheory_4x3m7y3neq2003 : ∀ (x y : ℤ), 4 * x ^ 3 - 7 * y ^ 3 ≠ 2003\n", + "binders=2 sort=Prop\n", + "====================================================================================================\n", + "aime_1983_p1\n", + "aime_1983_p1 : ∀ (x y z w : ℕ), 1 < x ∧ 1 < y ∧ 1 < z → 0 ≤ w → (↑w).log / (↑x).log = 24 → (↑w).log / (↑y).log = 40 → (↑w).log / (↑x * ↑y * ↑z).log = 12 → (↑w).log / (↑z).log = 60\n", + "binders=9 sort=Prop\n", + "====================================================================================================\n", + "amc12_2001_p5\n", + "amc12_2001_p5 : (Finset.filter (fun x => ¬Even x) (Finset.range 10000)).prod id = 10000! / (2 ^ 5000 * 5000!)\n", + "binders=0 sort=Prop\n", + "====================================================================================================\n", + "mathd_numbertheory_3\n", + "mathd_numbertheory_3 : (∑ x ∈ Finset.range 10, (x + 1) ^ 2) % 10 = 5\n", + "binders=0 sort=Prop\n", + "====================================================================================================\n", + "imo_1969_p2\n", + "imo_1969_p2 : ∀ (m n : ℝ) (k : ℕ) (a : ℕ → ℝ) (y : ℝ → ℝ), 0 < k → (∀ (x : ℝ), y x = ∑ i ∈ Finset.range k, (a i + x).cos / 2 ^ i) → y m = 0 → y n = 0 → ∃ t, m - n = ↑t * π\n", + "binders=9 sort=Prop\n", + "====================================================================================================\n", + "mathd_algebra_209\n", + "mathd_algebra_209 : ∀ (σ : ℝ ≃ ℝ), σ.invFun 2 = 10 → σ.invFun 10 = 1 → σ.invFun 1 = 2 → σ.toFun (σ.toFun 10) = 1\n", + "binders=4 sort=Prop\n", + "====================================================================================================\n", + "mathd_numbertheory_1124\n", + "mathd_numbertheory_1124 : ∀ n ≤ 9, 18 ∣ 374 * 10 + n → n = 4\n", + "binders=3 sort=Prop\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "====================================================================================================\n", + "imo_1983_p6\n", + "imo_1983_p6 : ∀ (a b c : ℝ), 0 < a ∧ 0 < b ∧ 0 < c → c < a + b → b < a + c → a < b + c → 0 ≤ a ^ 2 * b * (a - b) + b ^ 2 * c * (b - c) + c ^ 2 * a * (c - a)\n", + "binders=7 sort=Prop\n", + "====================================================================================================\n", + "mathd_algebra_33\n", + "mathd_algebra_33 : ∀ (x y z : ℝ), x ≠ 0 → 2 * x = 5 * y → 7 * y = 10 * z → z / x = 7 / 25\n", + "binders=6 sort=Prop\n" + ] + } + ], + "source": [ + "PICK = ['mathd_algebra_478', 'numbertheory_4x3m7y3neq2003', 'aime_1983_p1', 'amc12_2001_p5', 'mathd_numbertheory_3',\n", + " 'imo_1969_p2', 'mathd_algebra_209', 'mathd_numbertheory_1124', 'imo_1983_p6', 'mathd_algebra_33']\n", + "rows = {}\n", + "for l_ in open('/apps/development/rlhf/p2_runs/r1_baseline/minif2f.jsonl'):\n", + " r_ = json.loads(l_)\n", + " if r_['problem_id'] in PICK: rows[r_['problem_id']] = r_\n", + "assert set(rows) == set(PICK)\n", + "\n", + "SORT_ELAB = r'''\n", + "open Lean Meta Elab Command in\n", + "elab \"#sort_of \" id:ident : command => do\n", + " let env ← getEnv\n", + " match env.find? id.getId with\n", + " | none => throwError \"unknown\"\n", + " | some ci => liftTermElabM do\n", + " let s ← inferType ci.type\n", + " let nb := (← forallTelescope ci.type fun xs _ => pure xs.size)\n", + " logInfo m!\"binders={nb} sort={← whnf s}\"\n", + "'''\n", + "L.run(SORT_ELAB, keep=True)\n", + "\n", + "decoded = {}\n", + "for pid in PICK:\n", + " st = rows[pid]['formal_statement']\n", + " thm = st[st.index('/--'):] # docstring + theorem ... := by sorry\n", + " rest = thm\n", + " msgs, r = lean_msgs(rest + f\"\\n#check @{pid}\\n#sort_of {pid}\", timeout=300)\n", + " out = [m['data'] for m in msgs if m['severity'] == 'info']\n", + " errs = [m['data'] for m in msgs if m['severity'] == 'error']\n", + " assert not errs, errs\n", + " decoded[pid] = out\n", + " print('=' * 100); print(pid)\n", + " print(re.sub(r'\\n\\s+', ' ', out[0]))\n", + " print(out[1])" + ] + }, + { + "cell_type": "markdown", + "id": "2d1b76ed", + "metadata": {}, + "source": [ + "Closed statements (no binders) are plain `Prop`s with a `Decidable` instance: Lean can *evaluate* them." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "852923ed", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:54:26.104415Z", + "iopub.status.busy": "2026-10-14T09:54:26.104258Z", + "iopub.status.idle": "2026-10-14T09:54:26.288412Z", + "shell.execute_reply": "2026-10-14T09:54:26.287495Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "amc12_2001_p5 -> ['instDecidableEqNat ((Finset.filter (fun x => ¬Even x) (Finset.range 10000)).prod', 'true']\n", + "mathd_numbertheory_3 -> ['instDecidableEqNat ((∑ x ∈ Finset.range 10, (x + 1) ^ 2) % 10) 5', 'true']\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[\"'nt3' depends on axioms: [propext, Quot.sound]\"]\n" + ] + } + ], + "source": [ + "for pid in ['amc12_2001_p5', 'mathd_numbertheory_3']:\n", + " st = rows[pid]['formal_statement']\n", + " body = st.split(f'theorem {pid} :')[1].split(':= by sorry')[0].strip()\n", + " msgs, _ = lean_msgs(f\"#synth Decidable ({body})\\n#eval decide ({body})\", timeout=300)\n", + " print(pid, '->', [m['data'].strip()[:80] for m in msgs])\n", + "msgs, _ = lean_msgs(\"theorem nt3 : (∑ x in Finset.range 10, (x + 1) ^ 2) % 10 = 5 := by decide\\n#print axioms nt3\")\n", + "print([m['data'] for m in msgs])" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "b59c7854", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:54:26.290114Z", + "iopub.status.busy": "2026-10-14T09:54:26.289954Z", + "iopub.status.idle": "2026-10-14T09:54:26.591691Z", + "shell.execute_reply": "2026-10-14T09:54:26.590442Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "mathd_algebra_478 OK witness b=30,h=13/2 => v=65 exactly\n", + "numbertheory_4x3m7y3neq2003 OK cubes mod 7 = [0, 1, 6]; 4*cubes mod 7 = [0, 3, 4]; 2003 mod 7 = 1 not attained; box [-80,80]^2 empty\n", + "aime_1983_p1 OK fp64 max hypothesis residual 0.00e+00, conclusion residual 0.00e+00; exact: L/(L/12-L/24-L/40)=60\n", + "amc12_2001_p5 OK product of odd numbers < 10000 == 10000!/(2^5000 * 5000!) as big integers; division exact\n", + "mathd_numbertheory_3 OK sum = 385, 385 mod 10 = 5\n", + "imo_1969_p2 OK 200 random (k,a): y(x)=A cos x - B sin x; zeros m,n: max |y(m)|,|y(n)|,|(m-n)/π - t| = 3.7e-15; |(A,B)| >= 2^(1-k) > 0 in all 200 runs\n", + "mathd_algebra_209 OK σ = {10: 2, 1: 10, 2: 1}; σ(σ(10)) = 1\n", + "mathd_numbertheory_1124 OK solutions n in 0..9 with 18 | 3740+n: [4]\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "imo_1983_p6 OK min over 2e5 random triangles of E/a^4 = 6.588e-08 (>= 0); a=b=c gives exactly 0\n", + "mathd_algebra_33 OK x=1,y=2/5,z=7/25: z/x = 7/25 exactly\n", + "all 10 statements: hypotheses satisfiable (non-vacuous) and conclusion verified at witnesses / exhaustively.\n" + ] + } + ], + "source": [ + "# ---- (b) Python witness checks, one per statement ----\n", + "def check_478():\n", + " b, h = Fraction(30), Fraction(13, 2)\n", + " v = Fraction(1, 3) * (b * h)\n", + " assert 0 < b and 0 < h and 0 < v # hypotheses satisfiable (h₀)\n", + " assert v == 65 # conclusion\n", + " return 'witness b=30,h=13/2 => v=65 exactly'\n", + "\n", + "def check_4x3m7y3():\n", + " cubes = {pow(n, 3, 7) for n in range(7)}\n", + " resid = {(4 * c) % 7 for c in cubes}\n", + " assert 2003 % 7 == 1 and 1 not in resid # the real proof: 4x^3 - 7y^3 ≡ 4x^3 (mod 7)\n", + " assert not any(4*x**3 - 7*y**3 == 2003 for x in range(-80, 81) for y in range(-80, 81))\n", + " return f'cubes mod 7 = {sorted(cubes)}; 4*cubes mod 7 = {sorted(resid)}; 2003 mod 7 = 1 not attained; box [-80,80]^2 empty'\n", + "\n", + "def check_aime():\n", + " # w = 2^120, x = 2^5, y = 2^3, z = 2^2 : satisfies all three hypotheses; conclusion log_z w = 60\n", + " x, y, z, w = 2**5, 2**3, 2**2, 2**120\n", + " errs = [abs(math.log(w)/math.log(x) - 24), abs(math.log(w)/math.log(y) - 40), abs(math.log(w)/math.log(x*y*z) - 12)]\n", + " concl = abs(math.log(w)/math.log(z) - 60)\n", + " assert max(errs) < 1e-12 and concl < 1e-12\n", + " Lw = Fraction(120); lz = Lw/12 - Lw/24 - Lw/40 # log z = log(xyz) - log x - log y\n", + " assert Lw / lz == 60 # exact derivation\n", + " return f'fp64 max hypothesis residual {max(errs):.2e}, conclusion residual {concl:.2e}; exact: L/(L/12-L/24-L/40)=60'\n", + "\n", + "def check_amc():\n", + " prod = 1\n", + " for x in range(10000):\n", + " if x % 2 == 1: prod *= x\n", + " rhs_num = math.factorial(10000); rhs_den = 2**5000 * math.factorial(5000)\n", + " assert rhs_num % rhs_den == 0 # the ℕ-division is exact, so truncation plays no role\n", + " assert prod == rhs_num // rhs_den\n", + " return 'product of odd numbers < 10000 == 10000!/(2^5000 * 5000!) as big integers; division exact'\n", + "\n", + "def check_nt3():\n", + " s = sum((x + 1) ** 2 for x in range(10)); assert s == 385 and s % 10 == 5\n", + " return 'sum = 385, 385 mod 10 = 5'\n", + "\n", + "def check_imo69():\n", + " random.seed(1); worst = 0.0\n", + " for _ in range(200):\n", + " k = random.randint(1, 8); a = [random.uniform(-10, 10) for _ in range(k)]\n", + " A = sum(math.cos(a[i]) / 2**i for i in range(k)); B = sum(math.sin(a[i]) / 2**i for i in range(k))\n", + " x0 = math.atan2(A, B) # y(x) = A cos x - B sin x = 0 <=> x = atan2(A,B) + tπ\n", + " y = lambda x: sum(math.cos(a[i] + x) / 2**i for i in range(k))\n", + " m, n = x0 + 3 * math.pi, x0 - 2 * math.pi\n", + " worst = max(worst, abs(y(m)), abs(y(n)), abs((m - n) / math.pi - 5))\n", + " assert math.hypot(A, B) >= 2.0 ** (1 - k) - 1e-12 # |sum_i e^{i a_i} 2^-i| >= 1 - sum_{i=1}^{k-1} 2^-i = 2^(1-k) > 0 (FINITE sum)\n", + " assert worst < 1e-12\n", + " return f'200 random (k,a): y(x)=A cos x - B sin x; zeros m,n: max |y(m)|,|y(n)|,|(m-n)/π - t| = {worst:.1e}; |(A,B)| >= 2^(1-k) > 0 in all 200 runs'\n", + "\n", + "def check_209():\n", + " inv = {2: 10, 10: 1, 1: 2} # σ.2 = invFun\n", + " sigma = {v: k for k, v in inv.items()}\n", + " assert sigma[sigma[10]] == 1\n", + " return f'σ = {sigma}; σ(σ(10)) = {sigma[sigma[10]]}'\n", + "\n", + "def check_1124():\n", + " sol = [n for n in range(10) if (374 * 10 + n) % 18 == 0]\n", + " assert sol == [4]\n", + " return f'solutions n in 0..9 with 18 | 3740+n: {sol}'\n", + "\n", + "def check_1983():\n", + " random.seed(2); worst = float('inf')\n", + " for _ in range(200000):\n", + " x, y, z = (random.uniform(1e-3, 10) for _ in range(3)) # Ravi: a=y+z, b=z+x, c=x+y => triangle inequalities hold\n", + " a, b, c = y + z, z + x, x + y\n", + " val = a*a*b*(a-b) + b*b*c*(b-c) + c*c*a*(c-a)\n", + " worst = min(worst, val / (a**4)) # scale-free\n", + " assert worst > -1e-12\n", + " t = Fraction(7); assert t*t*t*(t-t) * 3 == 0\n", + " return f'min over 2e5 random triangles of E/a^4 = {worst:.3e} (>= 0); a=b=c gives exactly 0'\n", + "\n", + "def check_33():\n", + " x = Fraction(1); y = 2 * x / 5; z = 7 * y / 10\n", + " assert x != 0 and 2*x == 5*y and 7*y == 10*z and z / x == Fraction(7, 25)\n", + " return f'x=1,y={y},z={z}: z/x = 7/25 exactly'\n", + "\n", + "CHECKS = {'mathd_algebra_478': check_478, 'numbertheory_4x3m7y3neq2003': check_4x3m7y3, 'aime_1983_p1': check_aime,\n", + " 'amc12_2001_p5': check_amc, 'mathd_numbertheory_3': check_nt3, 'imo_1969_p2': check_imo69,\n", + " 'mathd_algebra_209': check_209, 'mathd_numbertheory_1124': check_1124, 'imo_1983_p6': check_1983,\n", + " 'mathd_algebra_33': check_33}\n", + "for pid in PICK:\n", + " print(f'{pid:30s} OK {CHECKS[pid]()}')\n", + "print('all 10 statements: hypotheses satisfiable (non-vacuous) and conclusion verified at witnesses / exhaustively.')" + ] + }, + { + "cell_type": "markdown", + "id": "4fadce18", + "metadata": {}, + "source": [ + "## 3b. Lemma checks (section 3.4 of the note)\n", + "`(a - b) + b = max a b` on ℕ (exhaustive in a box AND a Lean proof); the Euclidean division law `a = b*(a/b) + a%b, 0 <= a%b < |b|` for Lean's ℤ (exhaustive grid evaluated by Lean); and `x / 0 = 0` on ℝ." + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "165b7f70", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:54:26.593548Z", + "iopub.status.busy": "2026-10-14T09:54:26.593364Z", + "iopub.status.idle": "2026-10-14T09:54:27.071951Z", + "shell.execute_reply": "2026-10-14T09:54:27.071277Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Lean: a - b + b = max a b proved; 'tsub_add_max' depends on axioms: [propext, Quot.sound] | and the naive \"a - b + b = a\" is refuted at (0,1)\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Lean ℤ: b*(a/b) + a%b = a and 0 <= a%b < |b| on 1464 grid points; Lean (a/b, a%b) == the unique Euclidean (q,r) on all of them; x/0 = 0, x%0 = x\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Lean ℝ: ∀ x, x / 0 = 0 proved (div_zero); 'real_div_zero' depends on axioms: [propext, Classical.choice, Quot.sound]\n" + ] + } + ], + "source": [ + "\n", + "# (i) truncated subtraction: (a - b) + b = max a b\n", + "assert all(max(0, a - b) + b == max(a, b) for a in range(61) for b in range(61))\n", + "msgs, _ = lean_msgs(\"\"\"\n", + "theorem tsub_add_max (a b : ℕ) : a - b + b = max a b := by omega\n", + "#print axioms tsub_add_max\n", + "theorem tsub_add_false : ¬ (∀ a b : ℕ, a - b + b = a) := fun h => by have := h 0 1; simp at this\n", + "\"\"\")\n", + "assert not [m for m in msgs if m['severity'] == 'error'], msgs\n", + "print('Lean: a - b + b = max a b proved;', [m['data'] for m in msgs if m['severity'] == 'info'][0], '| and the naive \"a - b + b = a\" is refuted at (0,1)')\n", + "# (ii) Euclidean division law for Int, exhaustive on a grid, evaluated by Lean\n", + "msgs, _ = lean_msgs(\"\"\"\n", + "#eval (List.range 61).all fun ai => (List.range 25).all fun bi =>\n", + " let a : Int := (ai : Int) - 30\n", + " let b : Int := (bi : Int) - 12\n", + " b == 0 || ((b * (a / b) + a % b == a) && decide (0 ≤ a % b) && decide (a % b < (b.natAbs : Int)))\n", + "#eval (List.range 61).all fun ai => (List.range 25).all fun bi => ((ai : Int) - 30) / 0 == 0 && ((ai : Int) - 30) % 0 == (ai : Int) - 30\n", + "\"\"\")\n", + "res = [m['data'].strip() for m in msgs if m['severity'] == 'info']\n", + "assert res == ['true', 'true'], (res, msgs)\n", + "# the unique (q, r) with a = b q + r, 0 <= r < |b| (Python model) must equal Lean's / and %\n", + "def euclid(a, b):\n", + " r = a % abs(b); q = (a - r) // b; return q, r\n", + "lean_vals = lean_msgs(\"#eval (List.range 61).bind fun ai => (List.range 25).filterMap fun bi => let a : Int := (ai : Int) - 30; let b : Int := (bi : Int) - 12; if b == 0 then none else some (a / b, a % b)\")[0]\n", + "import ast\n", + "lv = ast.literal_eval(lean_vals[0]['data'].replace('(', '(').strip()) if lean_vals and lean_vals[0]['severity'] == 'info' else None\n", + "pv = [euclid(ai - 30, bi - 12) for ai in range(61) for bi in range(25) if bi != 12]\n", + "assert lv is not None and [tuple(x) for x in lv] == pv, 'Lean Int / and % differ from the Euclidean (q, r) model'\n", + "print(f'Lean ℤ: b*(a/b) + a%b = a and 0 <= a%b < |b| on {len(pv)} grid points; Lean (a/b, a%b) == the unique Euclidean (q,r) on all of them; x/0 = 0, x%0 = x')\n", + "# (iii) x / 0 = 0 on the reals\n", + "msgs, _ = lean_msgs(\"theorem real_div_zero : ∀ x : ℝ, x / 0 = 0 := fun x => div_zero x\\n#print axioms real_div_zero\")\n", + "assert not [m for m in msgs if m['severity'] == 'error'], msgs\n", + "print('Lean ℝ: ∀ x, x / 0 = 0 proved (div_zero);', [m['data'] for m in msgs if m['severity'] == 'info'][0])\n" + ] + }, + { + "cell_type": "markdown", + "id": "8d2c076e", + "metadata": {}, + "source": [ + "## 4. Elaboration traps (where a *correct-looking* statement means something else)" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "85822ae4", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:54:27.074232Z", + "iopub.status.busy": "2026-10-14T09:54:27.074069Z", + "iopub.status.idle": "2026-10-14T09:54:27.447944Z", + "shell.execute_reply": "2026-10-14T09:54:27.447329Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "(5 : ℕ) - 7 = 0\n", + "(7 : ℕ) / 2 = 3\n", + "(7 : ℤ) / (-2) = -3\n", + "(-7 : ℤ) / 2 = -4\n", + "(-7 : ℤ) % 3 = 2\n", + "(((5 : ℕ) - 7 : ℕ) : ℤ) = 0\n", + "((5 : ℕ) : ℤ) - 7 = -2\n", + "((10 : ℕ) - 3) / 2 = 3\n", + "(2 : ℕ) ^ 3 ^ 2 = 512\n", + "Lean 7/(-2) = -3 and (-7)/2 = -4 (Euclidean: remainder is always >= 0)\n", + "--- autoImplicit: a typo silently becomes a new universally-quantified variable ---\n", + "warning declaration uses 'sorry'\n", + "info @t_typo : ∀ {xx : ℕ} (x : ℕ), x + 1 > xx\n", + "--- the same with autoImplicit off: an error, as it should be ---\n", + "error unknown identifier 'xx'\n" + ] + } + ], + "source": [ + "TRAPS = [\n", + " (\"(5 : ℕ) - 7\", lambda: nsub(5, 7), \"0\"),\n", + " (\"(7 : ℕ) / 2\", lambda: 7 // 2, \"3\"),\n", + " (\"(7 : ℤ) / (-2)\", lambda: 'ediv', None),\n", + " (\"(-7 : ℤ) / 2\", lambda: 'ediv', None),\n", + " (\"(-7 : ℤ) % 3\", lambda: (-7) % 3, \"2\"),\n", + " (\"(((5 : ℕ) - 7 : ℕ) : ℤ)\", lambda: nsub(5, 7), \"0\"),\n", + " (\"((5 : ℕ) : ℤ) - 7\", lambda: 5 - 7, \"-2\"),\n", + " (\"((10 : ℕ) - 3) / 2\", lambda: (10 - 3) // 2, \"3\"),\n", + " (\"(2 : ℕ) ^ 3 ^ 2\", lambda: 2 ** (3 ** 2), \"512\"),\n", + "]\n", + "msgs, _ = lean_msgs(\"\\n\".join(f\"#eval {s}\" for s, _, _ in TRAPS))\n", + "by_line = {m['pos']['line']: m['data'].strip() for m in msgs}\n", + "for i, (s, fn, expect) in enumerate(TRAPS, start=1):\n", + " got = by_line[i]\n", + " if expect is not None:\n", + " assert got == expect, (s, got, expect)\n", + " py = fn(); assert str(py) == got, (s, py, got)\n", + " print(f'{s:28s} = {got}')\n", + "# Lean's ℤ `/` is Euclidean-style `Int.div` (T-rounding is `Int.tdiv`): -7/2 = -4, 7/(-2) = -3 ; Python's // floors: -7//2 = -4, 7//(-2) = -4\n", + "print('Lean 7/(-2) =', by_line[3], ' and (-7)/2 =', by_line[4], '(Euclidean: remainder is always >= 0)')\n", + "assert by_line[4] == '-4' and by_line[3] == '-3'\n", + "assert (-7) // 2 == -4 and 7 // -2 == -4 # python floors: differs from Lean on 7 / (-2)\n", + "\n", + "print('--- autoImplicit: a typo silently becomes a new universally-quantified variable ---')\n", + "for m in lean_msgs(\"theorem t_typo (x : ℕ) : x + 1 > xx := by sorry\\n#check @t_typo\")[0]: print(m['severity'], m['data'][:300])\n", + "print('--- the same with autoImplicit off: an error, as it should be ---')\n", + "for m in lean_msgs(\"set_option autoImplicit false in\\ntheorem t_typo2 (x : ℕ) : x + 1 > xx := by sorry\")[0]: print(m['severity'], m['data'][:200])\n", + "L.close()" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/code/notebooks/0303-hoare-refinement.ipynb b/code/notebooks/0303-hoare-refinement.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..49bf8194b7bd5547d5c033ceaad226f1209d097d --- /dev/null +++ b/code/notebooks/0303-hoare-refinement.ipynb @@ -0,0 +1,852 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "8fc053db", + "metadata": {}, + "source": [ + "# 0303 - Hoare logic and refinement; verifying a small function in Lean 4\n", + "\n", + "Scratch-pad for note `0303 - Hoare Logic.pdf`. Naive first, efficient second, assert agreement.\n", + "\n", + "**A. A weakest-precondition calculus for a tiny `while` language** (sympy for substitution).\n", + "* NAIVE = *run* the program on every state of a box and check the postcondition (semantic weakest precondition, `wp_sem`).\n", + "* EFFICIENT = *compute* `wp` by the syntactic rules (substitution) and verification conditions (VCs) for loops.\n", + "* Assert: `wp_syn == wp_sem` on 300 random loop-free programs; VCs valid => the Hoare triple holds on every box state;\n", + " wrong invariants are rejected, and we classify the counterexample (reachable = real bug, unreachable = proof artefact).\n", + "\n", + "**B. A verified linear search in Lean 4** (the DafnyBench Fig. 2 `LinearSearch`): code, loop invariant as an induction\n", + "hypothesis, theorem, and what Lean says about an off-by-one variant. **C. Integers**: spec over Z, code over int32.\n", + "\n", + "No GPU. Lean: ONE worker via `align.verify.lean`, nice 15, after the MemAvailable>=45 GB check in the helper." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "7964bb54", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:55:15.064599Z", + "iopub.status.busy": "2026-10-14T09:55:15.064441Z", + "iopub.status.idle": "2026-10-14T09:55:15.330291Z", + "shell.execute_reply": "2026-10-14T09:55:15.329450Z" + } + }, + "outputs": [], + "source": [ + "import os\n", + "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"\"\n", + "import itertools, random, sys, time\n", + "import sympy as sp\n", + "random.seed(0)\n", + "n, i, s, x, y = sp.symbols('n i s x y', integer=True)\n", + "\n", + "# AST: ('skip',) | ('assign', var, expr) | ('seq', S1, S2) | ('if', cond, S1, S2) | ('while', cond, inv, variant, body)\n", + "def seq(*ss):\n", + " out = ss[-1]\n", + " for t in reversed(ss[:-1]): out = ('seq', t, out)\n", + " return out\n", + "def imp(a, b): return sp.Or(sp.Not(a), b)\n", + "\n", + "_cache = {}\n", + "def ev(expr, st):\n", + " key = (expr, tuple(sorted(st, key=str)))\n", + " if key not in _cache:\n", + " syms = sorted(st, key=str)\n", + " _cache[key] = (syms, sp.lambdify(syms, expr, 'math'))\n", + " syms, f = _cache[key]\n", + " return f(*[st[v] for v in syms])\n", + "\n", + "class NonTerminating(Exception): pass\n", + "\n", + "def run(S, st, fuel=10_000):\n", + " \"\"\"Big-step semantics => st'. Returns the final state; raises NonTerminating when fuel is out.\"\"\"\n", + " st = dict(st)\n", + " def go(S):\n", + " nonlocal fuel\n", + " k = S[0]\n", + " if k == 'skip': return\n", + " if k == 'assign': st[S[1]] = ev(S[2], st); return\n", + " if k == 'seq': go(S[1]); go(S[2]); return\n", + " if k == 'if': go(S[2] if ev(S[1], st) else S[3]); return\n", + " if k == 'while':\n", + " while ev(S[1], st):\n", + " fuel -= 1\n", + " if fuel < 0: raise NonTerminating\n", + " go(S[4])\n", + " go(S)\n", + " return st\n", + "\n", + "def wp(S, Q):\n", + " \"\"\"Weakest precondition (loop-free fragment, plus `while` handled via its invariant: wp = I, see vcs()).\"\"\"\n", + " k = S[0]\n", + " if k == 'skip': return Q\n", + " if k == 'assign': return Q.subs(S[1], S[2], simultaneous=True)\n", + " if k == 'seq': return wp(S[1], wp(S[2], Q))\n", + " if k == 'if': return sp.Or(sp.And(S[1], wp(S[2], Q)), sp.And(sp.Not(S[1]), wp(S[3], Q)))\n", + " if k == 'while': return S[2] # annotated invariant; its VCs are generated by vcs()\n", + " raise ValueError(k)\n", + "\n", + "_ghost_ctr = itertools.count()\n", + "\n", + "def vcs(S, Q):\n", + " \"\"\"Verification conditions that make wp(S,Q) (with loops replaced by their invariants) sound. Returns [(name, formula)].\"\"\"\n", + " k = S[0]\n", + " if k in ('skip', 'assign'): return []\n", + " if k == 'seq': return vcs(S[1], wp(S[2], Q)) + vcs(S[2], Q)\n", + " if k == 'if': return vcs(S[2], Q) + vcs(S[3], Q)\n", + " if k == 'while':\n", + " b, I, V, body = S[1], S[2], S[3], S[4]\n", + " v0 = S[5] if len(S) > 5 else sp.Symbol(f'v0_{next(_ghost_ctr)}', integer=True) # one ghost per loop (nested loops must not share it)\n", + " ante = sp.And(I, b) if V is None else sp.And(I, b, sp.Eq(V, v0)) # v0 = the (ghost) value of the variant at loop entry\n", + " out = [('preserve', imp(ante, wp(body, I))), # I is preserved (the I-part of the body VC, reported separately)\n", + " ('exit', imp(sp.And(I, sp.Not(b)), Q))]\n", + " if V is not None:\n", + " R = sp.And(I, V < v0) # ONE postcondition for the body: I and the variant drop\n", + " out += [('variant>=0', imp(I, V >= 0)),\n", + " ('variant-decreases', imp(sp.And(I, b, sp.Eq(V, v0)), wp(body, R)))]\n", + " return out + vcs(body, R) # (needed for nested loops: their exit VCs must imply V < v0)\n", + " return out + vcs(body, I)\n", + " raise ValueError(k)\n", + "\n", + "def box(vars_, lo, hi):\n", + " return [dict(zip(vars_, p)) for p in itertools.product(range(lo, hi + 1), repeat=len(vars_))]\n", + "\n", + "GHOST = range(-12, 13)\n", + "def holds(F, st, ghost=None):\n", + " \"\"\"Validity of F at the state st. Free symbols that are not program variables (the ghost values v0 of the variant) are\n", + " universally quantified over a range. (Before the fix they were silently left unbound and Eq(V, v0) was structurally False,\n", + " which made every 'variant-decreases' VC vacuously true.)\"\"\"\n", + " ghosts = sorted(F.free_symbols - set(st), key=str)\n", + " if not ghosts: return bool(ev(F, st))\n", + " return all(bool(ev(F, {**st, **dict(zip(ghosts, gv))})) for gv in itertools.product(ghost or GHOST, repeat=len(ghosts)))" + ] + }, + { + "cell_type": "markdown", + "id": "c8f54cc4", + "metadata": {}, + "source": [ + "## A1. Straight-line and conditional code: `wp` by substitution equals the semantic weakest precondition\n", + "Semantic wp: `{ st | S terminates from st and the final state satisfies Q }`. We compare the *set* of box states on both sides." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "fe2daf34", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:55:15.332288Z", + "iopub.status.busy": "2026-10-14T09:55:15.332056Z", + "iopub.status.idle": "2026-10-14T09:55:19.396869Z", + "shell.execute_reply": "2026-10-14T09:55:19.395883Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "300 random loop-free programs x 49 states = 14700 (program,state) pairs; wp_syn != wp_sem on 0 (4.0s)\n" + ] + } + ], + "source": [ + "VARS = [x, y]\n", + "EXPRS = [x + 1, x - 1, 2 * x, y, x + y, x - y, y + 3, sp.Integer(0), x * y]\n", + "CONDS = [sp.Lt(x, y), sp.Le(x, 0), sp.Eq(x, y), sp.Gt(y, 1), sp.Ne(x + y, 2)]\n", + "POSTS = [sp.Eq(x, y), sp.Ge(x, 0), sp.Lt(x + y, 4), sp.Eq(2 * x, y + 1), sp.And(sp.Ge(x, 0), sp.Ge(y, 0)), sp.Ne(x, y)]\n", + "\n", + "def rand_prog(d):\n", + " r = random.random()\n", + " if d == 0 or r < 0.35: return ('assign', random.choice(VARS), random.choice(EXPRS))\n", + " if r < 0.7: return ('seq', rand_prog(d - 1), rand_prog(d - 1))\n", + " return ('if', random.choice(CONDS), rand_prog(d - 1), rand_prog(d - 1))\n", + "\n", + "STATES = box(VARS, -3, 3)\n", + "checked = mism = 0\n", + "t0 = time.time()\n", + "for _ in range(300):\n", + " S = rand_prog(3); Q = random.choice(POSTS)\n", + " W = wp(S, Q)\n", + " for st in STATES:\n", + " syn = holds(W, st)\n", + " sem = holds(Q, run(S, st)) # loop-free: always terminates\n", + " checked += 1; mism += (syn != sem)\n", + "assert mism == 0\n", + "print(f'300 random loop-free programs x {len(STATES)} states = {checked} (program,state) pairs; wp_syn != wp_sem on {mism} ({time.time()-t0:.1f}s)')" + ] + }, + { + "cell_type": "markdown", + "id": "c520492b", + "metadata": {}, + "source": [ + "## A2. Loops: invariant + variant. The worked example `sum`\n", + "```\n", + "{n >= 0} i := 0; s := 0; while i < n [inv I, variant n - i] do s := s + i; i := i + 1 {2 s = n (n - 1)}\n", + "```\n", + "`I : 0 <= i <= n and 2 s = i (i - 1)`." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "cb9b378a", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:55:19.398693Z", + "iopub.status.busy": "2026-10-14T09:55:19.398538Z", + "iopub.status.idle": "2026-10-14T09:55:21.057597Z", + "shell.execute_reply": "2026-10-14T09:55:21.056648Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "correct invariant all VCs valid on box: True\n", + "semantic run on 48 inputs satisfying/not P: triple holds -> True\n" + ] + } + ], + "source": [ + "body = seq(('assign', s, s + i), ('assign', i, i + 1))\n", + "cond = sp.Lt(i, n)\n", + "I_good = sp.And(sp.Ge(i, 0), sp.Le(i, n), sp.Eq(2 * s, i * (i - 1)))\n", + "P = sp.Ge(n, 0)\n", + "Qpost = sp.Eq(2 * s, n * (n - 1))\n", + "def program(inv, variant=n - i, bodyS=body):\n", + " return seq(('assign', i, sp.Integer(0)), ('assign', s, sp.Integer(0)), ('while', cond, inv, variant, bodyS))\n", + "\n", + "BOX = box([n, i, s], -2, 9)\n", + "BOX = [st for st in BOX if -2 <= st[s] <= 40] if False else BOX\n", + "\n", + "def verify(S, P, Q, label, reachable_states=None):\n", + " \"\"\"Generate all VCs (including P => wp(S,Q)) and discharge each on the box. Return dict name->counterexample or None.\"\"\"\n", + " conds = [('pre=>wp', imp(P, wp(S, Q)))] + vcs(S, Q)\n", + " res = {}\n", + " for name, F in conds:\n", + " bad = next((st for st in BOX if not holds(F, st)), None)\n", + " res[name] = bad\n", + " ok = all(v is None for v in res.values())\n", + " print(f'{label:34s} all VCs valid on box: {ok}' + ('' if ok else ' failing: ' + ', '.join(f'{k}@{ {str(a):b for a,b in v.items()} }' for k, v in res.items() if v)))\n", + " return res\n", + "\n", + "def triple_holds_semantically(S, P, Q, inputs):\n", + " \"\"\"Naive check: run from every input state satisfying P; the final state must satisfy Q (and the run must terminate).\"\"\"\n", + " for st in inputs:\n", + " if not holds(P, st): continue\n", + " try: fin = run(S, st)\n", + " except NonTerminating: return False, st\n", + " if not holds(Q, fin): return False, st\n", + " return True, None\n", + "\n", + "INPUTS = [dict(n=a, i=b, s=c) for a in range(-2, 10) for b in (0, 5) for c in (0, 7)]\n", + "INPUTS = [{n: d['n'], i: d['i'], s: d['s']} for d in INPUTS]\n", + "\n", + "r_good = verify(program(I_good), P, Qpost, 'correct invariant')\n", + "ok_sem, ce = triple_holds_semantically(program(I_good), P, Qpost, INPUTS)\n", + "assert all(v is None for v in r_good.values()) and ok_sem\n", + "print('semantic run on', len(INPUTS), 'inputs satisfying/not P: triple holds ->', ok_sem)" + ] + }, + { + "cell_type": "markdown", + "id": "1b414665", + "metadata": {}, + "source": [ + "### A3. Five ways to fail, and what each failure means\n", + "1. **Weak invariant** (`0 <= i <= n` only): the *exit* VC fails. The program is correct; the proof attempt fails. The counterexample state is\n", + " *unreachable* from any state satisfying `P`.\n", + "2. **Too-strong invariant** (`i < n`): `pre=>wp` / `preserve` fail. Program correct, annotation wrong.\n", + "3. **Buggy program** (`s := s + i + 1`): VCs fail with the *right* invariant; running the program from a concrete input violates the post. A **real bug**.\n", + "4. **Wrong postcondition** (`2 s = n (n + 1)`): the spec is wrong for the code; both checks fail." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "a6bbfb08", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:55:21.059329Z", + "iopub.status.busy": "2026-10-14T09:55:21.059160Z", + "iopub.status.idle": "2026-10-14T09:55:26.415346Z", + "shell.execute_reply": "2026-10-14T09:55:26.414446Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "1 weak invariant all VCs valid on box: False failing: exit@{'n': 0, 'i': 0, 's': -2}\n", + " counterexample {'n': 0, 'i': 0, 's': -2} -> reachable by the real program? False\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "2 too-strong invariant all VCs valid on box: False failing: pre=>wp@{'n': 0, 'i': -2, 's': -2}, preserve@{'n': 1, 'i': 0, 's': 0}, variant-decreases@{'n': 1, 'i': 0, 's': 0}\n", + " but the program is correct (semantic check): True\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "3 buggy program, right invariant all VCs valid on box: False failing: preserve@{'n': 1, 'i': 0, 's': 0}, variant-decreases@{'n': 1, 'i': 0, 's': 0}\n", + " concrete failing input from the run: {'n': 1, 'i': 0, 's': 0} -> confirmed bug\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "4 wrong postcondition all VCs valid on box: False failing: exit@{'n': 1, 'i': 1, 's': 0}\n", + " run also violates it, e.g. {'n': 1, 'i': 0, 's': 0}\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "5 wrong variant (V = i, increasing) all VCs valid on box: False failing: variant-decreases@{'n': 1, 'i': 0, 's': 0}\n", + " the program is still correct and terminates (semantic check): True -> the failing VC is about the PROOF of termination, not about a bug\n" + ] + } + ], + "source": [ + "reachable = set()\n", + "for nv in range(0, 10):\n", + " st = {n: nv, i: 0, s: 0}\n", + " # collect every loop-head state of the correct program\n", + " cur = {n: nv, i: 0, s: 0}\n", + " while cur[i] < cur[n]:\n", + " reachable.add((cur[n], cur[i], cur[s])); cur[s] += cur[i]; cur[i] += 1\n", + " reachable.add((cur[n], cur[i], cur[s]))\n", + "\n", + "I_weak = sp.And(sp.Ge(i, 0), sp.Le(i, n))\n", + "r1 = verify(program(I_weak), P, Qpost, '1 weak invariant')\n", + "assert r1['exit'] is not None and all(r1[k] is None for k in ('pre=>wp', 'preserve'))\n", + "ce1 = r1['exit']; print(' counterexample', {str(k): v for k, v in ce1.items()}, '-> reachable by the real program?', (ce1[n], ce1[i], ce1[s]) in reachable)\n", + "assert (ce1[n], ce1[i], ce1[s]) not in reachable # a proof artefact, NOT a bug\n", + "\n", + "I_strong = sp.And(sp.Ge(i, 0), sp.Lt(i, n), sp.Eq(2 * s, i * (i - 1)))\n", + "r2 = verify(program(I_strong), P, Qpost, '2 too-strong invariant')\n", + "assert any(v is not None for v in r2.values())\n", + "print(' but the program is correct (semantic check):', triple_holds_semantically(program(I_good), P, Qpost, INPUTS)[0])\n", + "\n", + "bug_body = seq(('assign', s, s + i + 1), ('assign', i, i + 1))\n", + "r3 = verify(program(I_good, bodyS=bug_body), P, Qpost, '3 buggy program, right invariant')\n", + "ok3, ce3 = triple_holds_semantically(program(I_good, bodyS=bug_body), P, Qpost, INPUTS)\n", + "assert r3['preserve'] is not None and not ok3\n", + "print(' concrete failing input from the run:', {str(k): v for k, v in ce3.items()}, '-> confirmed bug')\n", + "\n", + "Q_wrong = sp.Eq(2 * s, n * (n + 1))\n", + "r4 = verify(program(I_good), P, Q_wrong, '4 wrong postcondition')\n", + "ok4, ce4 = triple_holds_semantically(program(I_good), P, Q_wrong, INPUTS)\n", + "assert r4['exit'] is not None and not ok4\n", + "print(' run also violates it, e.g.', {str(k): v for k, v in ce4.items()})\n", + "\n", + "# 5. a WRONG VARIANT (V = i grows): the invariant and the code are right; only the termination argument is wrong\n", + "r5 = verify(program(I_good, variant=i), P, Qpost, '5 wrong variant (V = i, increasing)')\n", + "assert r5['variant-decreases'] is not None and r5['preserve'] is None and r5['exit'] is None\n", + "print(' the program is still correct and terminates (semantic check):', triple_holds_semantically(program(I_good), P, Qpost, INPUTS)[0],\n", + " '-> the failing VC is about the PROOF of termination, not about a bug')\n" + ] + }, + { + "cell_type": "markdown", + "id": "df54905f", + "metadata": {}, + "source": [ + "### A4. Soundness of the `while` rule, checked numerically\n", + "For every loop program whose VCs (all four) are valid on the box, every box state satisfying the invariant and variant terminates in `Q`\n", + "(this is `I subseteq wp(while, Q)`, the total-correctness statement of Boxed result 3)." + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "cd82f6da", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:55:26.416985Z", + "iopub.status.busy": "2026-10-14T09:55:26.416825Z", + "iopub.status.idle": "2026-10-14T09:55:26.469647Z", + "shell.execute_reply": "2026-10-14T09:55:26.468775Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "40 box states satisfying I and variant>=0: every one terminates in Q (0 violations)\n" + ] + } + ], + "source": [ + "cnt = viol = 0\n", + "for st in BOX:\n", + " if holds(I_good, st) and holds(n - i >= 0, st):\n", + " S_loop = ('while', cond, I_good, n - i, body)\n", + " fin = run(S_loop, st) # raises if it does not terminate\n", + " cnt += 1; viol += (not holds(Qpost, fin))\n", + "assert viol == 0\n", + "print(f'{cnt} box states satisfying I and variant>=0: every one terminates in Q ({viol} violations)')" + ] + }, + { + "cell_type": "markdown", + "id": "c5b343fb", + "metadata": {}, + "source": [ + "### A5. Nested loops (why the body is verified against `I and V < v0` in ONE postcondition, with one ghost per loop)\n", + "```\n", + "{n >= 0 and m >= 0} i := 0; s := 0; while i < n [I1, V1 = n-i] do ( j := 0; while j < m [I2, V2 = m-j] do ( s := s+1; j := j+1 ); i := i+1 ) {s = n m}\n", + "I1 = 0<=i<=n and m>=0 and s = i m I2 = 0<=j<=m and 0<=i=0 and s = i m + j and n-i = g1\n", + "```\n", + "`g1` is the ghost (`v0`) of the outer variant: the INNER loop's exit VC must imply `V1 < g1` for the outer body, so the inner invariant has to carry `n - i = g1`." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "b71a2c7b", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:55:26.471193Z", + "iopub.status.busy": "2026-10-14T09:55:26.471039Z", + "iopub.status.idle": "2026-10-14T09:55:35.643578Z", + "shell.execute_reply": "2026-10-14T09:55:35.642581Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "9 VCs generated for the nested program: ['pre=>wp', 'preserve', 'exit', 'variant>=0', 'variant-decreases', 'preserve', 'exit', 'variant>=0', 'variant-decreases']\n", + "all VCs valid on the box; the program run from 9 inputs terminates in Q (s = n m)\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "wrong inner variant (V2 = j): failing VCs -> ['variant-decreases']\n" + ] + } + ], + "source": [ + "\n", + "m_, j_, g1, g2 = sp.symbols('m j g1 g2', integer=True)\n", + "cond1, cond2 = sp.Lt(i, n), sp.Lt(j_, m_)\n", + "I1 = sp.And(sp.Ge(i, 0), sp.Le(i, n), sp.Ge(m_, 0), sp.Eq(s, i * m_))\n", + "I2 = sp.And(sp.Ge(j_, 0), sp.Le(j_, m_), sp.Ge(i, 0), sp.Lt(i, n), sp.Ge(m_, 0), sp.Eq(s, i * m_ + j_), sp.Eq(n - i, g1))\n", + "def nested(inner_variant):\n", + " inner = ('while', cond2, I2, inner_variant, seq(('assign', s, s + 1), ('assign', j_, j_ + 1)), g2)\n", + " outer_body = seq(('assign', j_, sp.Integer(0)), inner, ('assign', i, i + 1))\n", + " return seq(('assign', i, sp.Integer(0)), ('assign', s, sp.Integer(0)), ('while', cond1, I1, n - i, outer_body, g1))\n", + "P_n, Q_n = sp.And(sp.Ge(n, 0), sp.Ge(m_, 0)), sp.Eq(s, n * m_)\n", + "BOX_N = [dict(zip([n, m_, i, j_, s], p)) for p in itertools.product(range(0, 3), range(0, 3), range(0, 4), range(0, 4), range(0, 10))]\n", + "GN = range(-3, 4)\n", + "def failing_vcs(S):\n", + " conds = [('pre=>wp', imp(P_n, wp(S, Q_n)))] + vcs(S, Q_n)\n", + " return [name for name, F in conds if any(not holds(F, st, GN) for st in BOX_N)], [name for name, _ in conds]\n", + "bad, names = failing_vcs(nested(m_ - j_))\n", + "print(len(names), 'VCs generated for the nested program:', names)\n", + "assert bad == [], bad\n", + "ins = [st for st in BOX_N if holds(P_n, st) and st[i] == 0 and st[j_] == 0 and st[s] == 0]\n", + "for st in ins:\n", + " fin = run(nested(m_ - j_), st); assert holds(Q_n, fin) # raises NonTerminating if the run does not terminate\n", + "print(f'all VCs valid on the box; the program run from {len(ins)} inputs terminates in Q (s = n m)')\n", + "bad2, _ = failing_vcs(nested(j_)) # a WRONG inner variant (V2 = j grows)\n", + "print('wrong inner variant (V2 = j): failing VCs ->', bad2)\n", + "assert 'variant-decreases' in bad2\n" + ] + }, + { + "cell_type": "markdown", + "id": "63e05762", + "metadata": {}, + "source": [ + "## B. Verifying `LinearSearch` in Lean 4 (DafnyBench Fig. 2 in Lean)\n", + "Dafny: `ensures 0 <= n <= a.Length; ensures n == a.Length || P(a[n]); ensures forall i :: 0 <= i < n ==> !P(a[i])` with loop invariants\n", + "`0 <= n <= a.Length` and `forall i :: 0 <= i < n ==> !P(a[i])`. Abstracting `a[i]` away (we search an index predicate `P`):" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "55dab72a", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:55:35.645474Z", + "iopub.status.busy": "2026-10-14T09:55:35.645313Z", + "iopub.status.idle": "2026-10-14T09:55:40.310783Z", + "shell.execute_reply": "2026-10-14T09:55:40.309894Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[info] line 36: 'search_correct' depends on axioms: [propext, Quot.sound]\n", + "[info] line 37: 3\n", + "[info] line 38: 10\n" + ] + } + ], + "source": [ + "sys.path.insert(0, os.path.abspath('.'))\n", + "from _lean_helper import Lean\n", + "L = Lean()\n", + "code = r'''\n", + "/-- the loop `while n < len do if P n then return else n := n + 1`, as a tail-recursive function -/\n", + "def search (P : Nat → Bool) (len : Nat) (n : Nat) : Nat :=\n", + " if n < len then (if P n then n else search P len (n + 1)) else n\n", + "termination_by len - n\n", + "\n", + "/-- postcondition of Dafny's LinearSearch (with the start index `n0` of the loop) -/\n", + "def Post (P : Nat → Bool) (len : Nat) (r : Nat) : Prop :=\n", + " r ≤ len ∧ (r = len ∨ P r = true) ∧ ∀ i, i < r → P i = false\n", + "\n", + "/-- loop INVARIANT as an induction hypothesis: from any state satisfying (n ≤ len ∧ ∀ i < n, ¬P i) the loop ends in Post -/\n", + "theorem search_inv (P : Nat → Bool) (len : Nat) :\n", + " ∀ n, n ≤ len → (∀ i, i < n → P i = false) → Post P len (search P len n) := by\n", + " intro n\n", + " induction' h : len - n using Nat.strong_induction_on with k ih generalizing n\n", + " intro hn hinv\n", + " unfold search\n", + " by_cases hlt : n < len\n", + " · rw [if_pos hlt]\n", + " by_cases hp : P n = true\n", + " · rw [if_pos hp]\n", + " exact ⟨by omega, Or.inr hp, hinv⟩\n", + " · rw [if_neg hp]\n", + " have hp' : P n = false := by simpa using hp\n", + " apply ih (len - (n + 1)) (by omega) (n + 1) rfl (by omega)\n", + " intro j hj\n", + " by_cases hjn : j = n\n", + " · subst hjn; exact hp'\n", + " · exact hinv j (by omega)\n", + " · rw [if_neg hlt]\n", + " exact ⟨by omega, Or.inl (by omega), hinv⟩\n", + "\n", + "theorem search_correct (P : Nat → Bool) (len : Nat) : Post P len (search P len 0) :=\n", + " search_inv P len 0 (Nat.zero_le _) (by intro i hi; omega)\n", + "\n", + "#print axioms search_correct\n", + "#eval search (fun i => i == 3) 10 0\n", + "#eval search (fun i => i == 30) 10 0\n", + "'''\n", + "r = L.show(code, timeout=300, keep=True)" + ] + }, + { + "cell_type": "markdown", + "id": "041614a3", + "metadata": {}, + "source": [ + "### B2. The same proof attempt on an off-by-one loop (`n + 1 < len`)\n", + "The proof script is *identical*; Lean refuses. A failed proof is not yet a bug report (note 0304) -- but here a one-line `#eval` confirms it." + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "3f1bb803", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:55:40.312728Z", + "iopub.status.busy": "2026-10-14T09:55:40.312554Z", + "iopub.status.idle": "2026-10-14T09:55:40.878639Z", + "shell.execute_reply": "2026-10-14T09:55:40.877856Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[ok] no messages\n", + "[error] line 21: omega could not prove the goal:\n", + "a possible counterexample may satisfy the constraints\n", + " a ≥ 0\n", + "where\n", + " a := ↑n\n", + "--- concrete confirmation (no match anywhere: the correct loop returns len = 10) ---\n", + "[info] line 2: 10\n", + "[info] line 3: 9\n", + "[info] line 4: false\n", + "[info] line 5: true\n" + ] + }, + { + "data": { + "text/plain": [ + "{'messages': [{'severity': 'info',\n", + " 'pos': {'line': 2, 'column': 0},\n", + " 'endPos': {'line': 2, 'column': 5},\n", + " 'data': '10'},\n", + " {'severity': 'info',\n", + " 'pos': {'line': 3, 'column': 0},\n", + " 'endPos': {'line': 3, 'column': 5},\n", + " 'data': '9'},\n", + " {'severity': 'info',\n", + " 'pos': {'line': 4, 'column': 0},\n", + " 'endPos': {'line': 4, 'column': 5},\n", + " 'data': 'false'},\n", + " {'severity': 'info',\n", + " 'pos': {'line': 5, 'column': 0},\n", + " 'endPos': {'line': 5, 'column': 5},\n", + " 'data': 'true'}],\n", + " 'env': 5}" + ] + }, + "execution_count": 8, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "L.show(r'''\n", + "def searchBad (P : Nat → Bool) (len : Nat) (n : Nat) : Nat :=\n", + " if n + 1 < len then (if P n then n else searchBad P len (n + 1)) else n\n", + "termination_by len - n\n", + "instance (P : Nat → Bool) (len r : Nat) : Decidable (Post P len r) := by unfold Post; infer_instance\n", + "''', keep=True)\n", + "bad = r'''\n", + "theorem searchBad_inv (P : Nat → Bool) (len : Nat) :\n", + " ∀ n, n ≤ len → (∀ i, i < n → P i = false) → Post P len (searchBad P len n) := by\n", + " intro n\n", + " induction' h : len - n using Nat.strong_induction_on with k ih generalizing n\n", + " intro hn hinv\n", + " unfold searchBad\n", + " by_cases hlt : n + 1 < len\n", + " · rw [if_pos hlt]\n", + " by_cases hp : P n = true\n", + " · rw [if_pos hp]\n", + " exact ⟨by omega, Or.inr hp, hinv⟩\n", + " · rw [if_neg hp]\n", + " have hp' : P n = false := by simpa using hp\n", + " apply ih (len - (n + 1)) (by omega) (n + 1) rfl (by omega)\n", + " intro j hj\n", + " by_cases hjn : j = n\n", + " · subst hjn; exact hp'\n", + " · exact hinv j (by omega)\n", + " · rw [if_neg hlt]\n", + " exact ⟨by omega, Or.inl (by omega), hinv⟩\n", + "'''\n", + "rr = L.show(bad, timeout=300)\n", + "print('--- concrete confirmation (no match anywhere: the correct loop returns len = 10) ---')\n", + "L.show('''\n", + "#eval search (fun _ => false) 10 0\n", + "#eval searchBad (fun _ => false) 10 0\n", + "#eval decide (Post (fun _ => false) 10 (searchBad (fun _ => false) 10 0))\n", + "#eval decide (Post (fun _ => false) 10 (search (fun _ => false) 10 0))\n", + "''')" + ] + }, + { + "cell_type": "markdown", + "id": "acbe59cb", + "metadata": {}, + "source": [ + "## C. The spec lives in Z, the code in int32: binary-search midpoint\n", + "VC (over Z): `lo <= hi => lo <= (lo+hi)/2 <= hi`. Over 8-bit words the identical VC is **false** exactly when `lo + hi >= 256`.\n", + "Derivation: wrapped `lo+hi = lo+hi-256`, so `mid = (lo+hi-256)/2 < lo <=> hi - lo < 256`, always true for 8-bit words.\n", + "Count the violating pairs in closed form, by brute force in Python, and by Lean `#eval` on `UInt8`." + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "86a322fd", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:55:40.880806Z", + "iopub.status.busy": "2026-10-14T09:55:40.880641Z", + "iopub.status.idle": "2026-10-14T09:55:41.086181Z", + "shell.execute_reply": "2026-10-14T09:55:41.085341Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "UInt8 pairs lo<=hi: 32896; violating VC: unsafe 16384 (= #{lo+hi>=256} = 16384), safe 0\n", + "Lean #eval counts (unsafe, safe, UInt32 mid): ['16384', '0', '852516352'] ; decide-proof of the 32-bit violation: True\n", + "int32: lo = 1073741829 hi = 2147483646 (lo+hi) wraps to -1073741821 -> mid = -536870911 which is < lo: True\n" + ] + } + ], + "source": [ + "M = 256\n", + "def mid_unsafe(lo, hi): return ((lo + hi) % M) // 2 # what UInt8 arithmetic computes\n", + "def mid_safe(lo, hi): return (lo + (hi - lo) // 2) % M # no intermediate overflow when lo <= hi\n", + "pairs = [(lo, hi) for lo in range(M) for hi in range(lo, M)]\n", + "py_bad_unsafe = sum(1 for lo, hi in pairs if not (lo <= mid_unsafe(lo, hi) <= hi))\n", + "py_bad_safe = sum(1 for lo, hi in pairs if not (lo <= mid_safe(lo, hi) <= hi))\n", + "closed_form = sum(1 for lo, hi in pairs if lo + hi >= M) # the derived characterisation\n", + "assert py_bad_unsafe == closed_form and py_bad_safe == 0\n", + "print(f'UInt8 pairs lo<=hi: {len(pairs)}; violating VC: unsafe {py_bad_unsafe} (= #{{lo+hi>=256}} = {closed_form}), safe {py_bad_safe}')\n", + "# the same count in Lean on UInt8\n", + "msgs = L.run('''\n", + "#eval ((List.range 256).foldl (fun acc lo => (List.range 256).foldl (fun acc hi =>\n", + " if lo ≤ hi ∧ ¬ (UInt8.ofNat lo ≤ (UInt8.ofNat lo + UInt8.ofNat hi) / 2 ∧ (UInt8.ofNat lo + UInt8.ofNat hi) / 2 ≤ UInt8.ofNat hi) then acc + 1 else acc) acc) 0)\n", + "#eval ((List.range 256).foldl (fun acc lo => (List.range 256).foldl (fun acc hi =>\n", + " if lo ≤ hi ∧ ¬ (UInt8.ofNat lo ≤ UInt8.ofNat lo + (UInt8.ofNat hi - UInt8.ofNat lo) / 2 ∧ UInt8.ofNat lo + (UInt8.ofNat hi - UInt8.ofNat lo) / 2 ≤ UInt8.ofNat hi) then acc + 1 else acc) acc) 0)\n", + "#eval ((3000000000 : UInt32) + 3000000000) / 2\n", + "example : ¬ ((3000000000 : UInt32) ≤ (3000000000 + 3000000000) / 2) := by decide\n", + "''')['messages']\n", + "lean_counts = [m['data'].strip() for m in msgs if m['severity'] == 'info']\n", + "print('Lean #eval counts (unsafe, safe, UInt32 mid):', lean_counts, '; decide-proof of the 32-bit violation: ', [m['severity'] for m in msgs if m['severity'] == 'error'] == [])\n", + "assert int(lean_counts[0]) == py_bad_unsafe and int(lean_counts[1]) == 0\n", + "# 32-bit: python model of int32 wrap, the famous boundary case\n", + "def wrap32(v): return (v + 2**31) % 2**32 - 2**31\n", + "lo, hi = 2**30 + 5, 2**31 - 2\n", + "print('int32: lo =', lo, ' hi =', hi, ' (lo+hi) wraps to', wrap32(lo + hi), '-> mid =', wrap32(lo + hi) // 2, 'which is < lo:', wrap32(lo + hi) // 2 < lo)\n", + "assert wrap32(lo + hi) // 2 < lo and lo <= (lo + hi) // 2 <= hi" + ] + }, + { + "cell_type": "markdown", + "id": "f9bbfb94", + "metadata": {}, + "source": [ + "## D. Refinement: `x :[pre, post] is-refined-by x := e iff pre => post[e/x]`\n", + "Predicate-transformer reading: `wp_spec(Q)(s) = pre(s) and forall x'. post(s, x') => Q(x')`, `wp_assign(Q)(s) = Q(e(s))`;\n", + "`spec ⊑ impl` iff `wp_spec(Q) ⊆ wp_impl(Q)` for every `Q`. We decide this by brute force over ALL 2^11 postconditions on a box\n", + "(named approximation: box of 11 values) and compare with the one-line law." + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "ca6a9e0e", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:55:41.087953Z", + "iopub.status.busy": "2026-10-14T09:55:41.087795Z", + "iopub.status.idle": "2026-10-14T09:55:42.520867Z", + "shell.execute_reply": "2026-10-14T09:55:42.519895Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "80 random (spec, assignment) pairs: all-Q brute force == worst-Q shortcut == law pre => post[e/x]; refined in 71, not refined in 9\n" + ] + } + ], + "source": [ + "BOXX = list(range(-5, 6)) # final-state values x' (11 of them => 2^11 postconditions Q)\n", + "START = list(range(-2, 3)) # start states s0; |b|<=2, |k|<=1 keep e(s0) inside BOXX\n", + "def refines_bruteforce(pre, post, e):\n", + " # for every Q (every subset of the box), wp_spec(Q) ⊆ wp_assign(Q) on every start state\n", + " for s0 in START:\n", + " if not pre(s0): continue # wp_spec(Q)(s0)=False for all Q: nothing to check\n", + " allowed = [xp for xp in BOXX if post(s0, xp)] # outcomes the spec permits\n", + " # wp_spec(Q)(s0) holds iff allowed ⊆ Q ; refinement needs e(s0) ∈ Q for all such Q. Worst Q = allowed.\n", + " if e(s0) not in allowed: return False\n", + " return True\n", + "def law(pre, post, e): return all((not pre(s0)) or post(s0, e(s0)) for s0 in START)\n", + "# make the brute force genuinely quantify over Q for a subset of cases, to validate the \"worst Q = allowed\" shortcut\n", + "def refines_all_Q(pre, post, e):\n", + " for s0 in START:\n", + " for mask in range(1 << len(BOXX)):\n", + " Q = {BOXX[k] for k in range(len(BOXX)) if mask >> k & 1}\n", + " wp_spec = pre(s0) and all((xp in Q) for xp in BOXX if post(s0, xp))\n", + " wp_impl = e(s0) in Q\n", + " if wp_spec and not wp_impl: return False\n", + " return True\n", + "random.seed(5)\n", + "tests = []\n", + "for _ in range(60):\n", + " a, b, c = random.randint(-2, 2), random.randint(-2, 2), random.randint(1, 3)\n", + " pre = (lambda a: lambda s0: s0 >= a)(a)\n", + " post = (lambda b, c: lambda s0, xp: abs(xp - (s0 + b)) <= c - 1)(b, c) # \"x ends within c-1 of x0+b\"\n", + " e = (lambda b: lambda s0: s0 + b)(b)\n", + " tests.append((pre, post, e))\n", + "# also deliberately failing implementations\n", + "for _ in range(60):\n", + " a, b, c = random.randint(-2, 2), random.randint(-2, 2), random.randint(1, 2)\n", + " pre = (lambda a: lambda s0: s0 >= a)(a)\n", + " post = (lambda b, c: lambda s0, xp: abs(xp - (s0 + b)) <= c - 1)(b, c)\n", + " k = random.randint(-1, 1)\n", + " e = (lambda b, k: lambda s0: s0 + b + k)(b, k)\n", + " tests.append((pre, post, e))\n", + "agree = 0; n_true = 0\n", + "for pre, post, e in tests[:40] + tests[60:100]:\n", + " r1 = refines_all_Q(pre, post, e); r2 = refines_bruteforce(pre, post, e); r3 = law(pre, post, e)\n", + " assert r1 == r2 == r3\n", + " agree += 1; n_true += r1\n", + "print(f'{agree} random (spec, assignment) pairs: all-Q brute force == worst-Q shortcut == law pre => post[e/x]; refined in {n_true}, not refined in {agree-n_true}')" + ] + }, + { + "cell_type": "markdown", + "id": "cc1bc3fa", + "metadata": {}, + "source": [ + "## Summary of the checks\n", + "* `wp` by substitution == semantic wp on 14,700 (program, state) pairs (0 mismatches).\n", + "* Valid VCs (with variant) => triple holds and loop terminates (40 box states, 0 violations); 4 failure modes classified.\n", + "* Lean: `search_correct` proved (axioms `propext`, `Quot.sound`); the off-by-one variant is rejected and `#eval` produces a concrete violation.\n", + "* Overflow: VC over Z true, over `UInt8` false for exactly `#{lo+hi >= 256}` = 16,384 pairs, matching Python and Lean `#eval` counts.\n", + "* Refinement law == brute force over all 2^11 postconditions on 80 random pairs." + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "id": "bca6ef85", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:55:42.522453Z", + "iopub.status.busy": "2026-10-14T09:55:42.522294Z", + "iopub.status.idle": "2026-10-14T09:55:42.789460Z", + "shell.execute_reply": "2026-10-14T09:55:42.788381Z" + } + }, + "outputs": [], + "source": [ + "L.close()" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/code/notebooks/0304-failed-proof-to-bug.ipynb b/code/notebooks/0304-failed-proof-to-bug.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..1f311a359b1063401ed2b5284ae15a7f042514cc --- /dev/null +++ b/code/notebooks/0304-failed-proof-to-bug.ipynb @@ -0,0 +1,1006 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "4f055b7b", + "metadata": {}, + "source": [ + "# 0304 - From failed proof to bug: counterexample search and the soundness / completeness asymmetry\n", + "\n", + "Scratch-pad for note `0304 - Failed Proof to Bug.pdf`. Naive first, efficient second, assert agreement.\n", + "\n", + "* **A. Random testing with shrinking** (a 60-line QuickCheck). NAIVE = exhaustive enumeration of the bounded domain.\n", + " EFFICIENT = random sampling + shrinking. Assert: the exact failure probability `p` (by enumeration) predicts the\n", + " Monte-Carlo miss probability `(1-p)^N` and the mean time-to-failure `1/p`; shrinking reaches the global minimum.\n", + "* **B. Small-scope hypothesis**: when does enumeration beat random sampling? Closed forms, checked by Monte Carlo.\n", + "* **C. Lean**: `decide`, `slim_check` (this Mathlib pin's name for Plausible), `omega` failure messages; a counterexample turned into a\n", + " *kernel-checked refutation*; and the cases where a failed proof / counterexample is **not** a bug.\n", + "\n", + "No GPU. Lean: ONE worker through `align.verify.lean`, nice 15, after the MemAvailable>=45 GB check." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "9540a401", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:54:29.990834Z", + "iopub.status.busy": "2026-10-14T09:54:29.990642Z", + "iopub.status.idle": "2026-10-14T09:54:29.997582Z", + "shell.execute_reply": "2026-10-14T09:54:29.996880Z" + } + }, + "outputs": [], + "source": [ + "import os\n", + "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"\"\n", + "import itertools, random, math, sys, collections\n", + "random.seed(2026)\n", + "\n", + "# ---- the code under test: insertion sort with a duplicate-handling bug -------------------------------------------\n", + "def isort_bug(xs):\n", + " out = []\n", + " for v in xs:\n", + " k = 0\n", + " while k < len(out) and out[k] < v: k += 1\n", + " if k < len(out) and out[k] == v: continue # BUG: drops the duplicate\n", + " out.insert(k, v)\n", + " return out\n", + "def isort_ok(xs):\n", + " out = []\n", + " for v in xs:\n", + " k = 0\n", + " while k < len(out) and out[k] < v: k += 1\n", + " out.insert(k, v)\n", + " return out\n", + "def prop_sort(f):\n", + " def prop(xs):\n", + " ys = f(xs)\n", + " return all(a <= b for a, b in zip(ys, ys[1:])) and collections.Counter(ys) == collections.Counter(xs)\n", + " return prop\n", + "\n", + "# Domain: lists of length 0..6 over {-2..2}; generator: length uniform in 0..6, entries uniform.\n", + "V = list(range(-2, 3)); MAXLEN = 6\n", + "def gen(rng):\n", + " return [rng.choice(V) for _ in range(rng.randint(0, MAXLEN))]" + ] + }, + { + "cell_type": "markdown", + "id": "24aac6bf", + "metadata": {}, + "source": [ + "## A1. NAIVE: exhaustive enumeration gives the exact failure probability under the generator\n", + "`P(list) = (1/7) * 5^-len`. Sum the probability of the failing lists: this `p` is the per-test detection probability." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "f8691e4c", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:54:29.999823Z", + "iopub.status.busy": "2026-10-14T09:54:29.999658Z", + "iopub.status.idle": "2026-10-14T09:54:30.269920Z", + "shell.execute_reply": "2026-10-14T09:54:30.269084Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "domain size 19531; failing inputs 19205; exact per-test failure probability p = 0.498514\n" + ] + } + ], + "source": [ + "prop_bug, prop_ok = prop_sort(isort_bug), prop_sort(isort_ok)\n", + "fail_prob = 0.0; fails = []\n", + "for L_ in range(MAXLEN + 1):\n", + " for xs in itertools.product(V, repeat=L_):\n", + " xs = list(xs)\n", + " assert prop_ok(xs) # the reference implementation passes everywhere\n", + " if not prop_bug(xs):\n", + " fails.append(xs); fail_prob += (1 / (MAXLEN + 1)) * 5.0 ** (-L_)\n", + "n_total = sum(5 ** L_ for L_ in range(MAXLEN + 1))\n", + "print(f'domain size {n_total}; failing inputs {len(fails)}; exact per-test failure probability p = {fail_prob:.6f}')\n", + "p = fail_prob" + ] + }, + { + "cell_type": "markdown", + "id": "ee87ac91", + "metadata": {}, + "source": [ + "## A2. EFFICIENT: random testing. Derivation (note 0304 section 3): `P(no failure in N independent tests) = (1-p)^N`;\n", + "time to first failure is Geometric(p) with mean `1/p`. Check both by Monte Carlo (20,000 runs)." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "c36a5bed", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:54:30.271520Z", + "iopub.status.busy": "2026-10-14T09:54:30.271369Z", + "iopub.status.idle": "2026-10-14T09:54:30.557635Z", + "shell.execute_reply": "2026-10-14T09:54:30.556905Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "mean tests to first failure: empirical 1.992 +- 0.010 ; theory 1/p = 2.006 ; |diff| = 0.014 (1.43 se)\n", + " N P(miss) empirical (1-p)^N z\n", + " 1 0.5030 0.5015 0.41\n", + " 5 0.0297 0.0317 1.59\n", + " 10 0.0008 0.0010 0.70\n", + " 20 0.0000 0.0000 0.14\n", + " 40 0.0000 0.0000 0.00\n", + " 80 0.0000 0.0000 0.00\n", + "worst |z| = 1.59 (< 4: the geometric law (1-p)^N holds)\n", + " N for miss-probability <= 0.1: 4\n", + " N for miss-probability <= 0.01: 7\n", + " N for miss-probability <= 0.001: 11\n" + ] + } + ], + "source": [ + "rng = random.Random(7)\n", + "RUNS = 20000\n", + "first_fail = []\n", + "for _ in range(RUNS):\n", + " t = 0\n", + " while True:\n", + " t += 1\n", + " if not prop_bug(gen(rng)): break\n", + " first_fail.append(t)\n", + "mean_t = sum(first_fail) / RUNS\n", + "se_mean = (sum((t - mean_t) ** 2 for t in first_fail) / (RUNS - 1)) ** 0.5 / RUNS ** 0.5\n", + "print(f'mean tests to first failure: empirical {mean_t:.3f} +- {se_mean:.3f} ; theory 1/p = {1/p:.3f} ; |diff| = {abs(mean_t - 1/p):.3f} ({abs(mean_t-1/p)/se_mean:.2f} se)')\n", + "assert abs(mean_t - 1 / p) < 4 * se_mean\n", + "worst_z = 0.0\n", + "print(f'{\"N\":>4s} {\"P(miss) empirical\":>18s} {\"(1-p)^N\":>10s} {\"z\":>6s}')\n", + "for N in (1, 5, 10, 20, 40, 80):\n", + " emp = sum(t > N for t in first_fail) / RUNS\n", + " th = (1 - p) ** N\n", + " se = max((th * (1 - th) / RUNS) ** 0.5, 1e-9)\n", + " z = abs(emp - th) / se; worst_z = max(worst_z, z)\n", + " print(f'{N:4d} {emp:18.4f} {th:10.4f} {z:6.2f}')\n", + "assert worst_z < 4\n", + "print('worst |z| =', round(worst_z, 2), '(< 4: the geometric law (1-p)^N holds)')\n", + "# tests needed for confidence 1 - delta: N >= ln(delta) / ln(1-p)\n", + "for delta in (0.1, 0.01, 1e-3):\n", + " print(f' N for miss-probability <= {delta}: {math.ceil(math.log(delta) / math.log(1 - p))}')" + ] + }, + { + "cell_type": "markdown", + "id": "06805f67", + "metadata": {}, + "source": [ + "### A2b. A rarer bug: the same law at p = 1/1024 (a \"needle\": the code is wrong on exactly one input out of K = 1024)\n", + "Uniform sampling over the K inputs; per-test detection probability is exactly `p = 1/K`." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "a033314f", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:54:30.559759Z", + "iopub.status.busy": "2026-10-14T09:54:30.559600Z", + "iopub.status.idle": "2026-10-14T09:54:32.010396Z", + "shell.execute_reply": "2026-10-14T09:54:32.009675Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "needle p=1/1024: mean tests to failure 1015.6 +- 18.3 ; theory 1024\n", + " N= 256 P(miss) empirical 0.7790 (1-p)^N 0.7787 z=0.04\n", + " N= 1024 P(miss) empirical 0.3643 (1-p)^N 0.3677 z=0.38\n", + " N= 2048 P(miss) empirical 0.1297 (1-p)^N 0.1352 z=0.89\n", + " N= 3072 P(miss) empirical 0.0497 (1-p)^N 0.0497 z=0.01\n", + " N= 4096 P(miss) empirical 0.0173 (1-p)^N 0.0183 z=0.39\n", + "=> to be 99% sure of seeing a 1-in-1024 bug you need 4714 random tests; testing is a *statistical* instrument.\n" + ] + } + ], + "source": [ + "K2 = 1024; p2 = 1 / K2; RUNS2 = 3000\n", + "rng = random.Random(23)\n", + "needle = 777\n", + "ff = []\n", + "for _ in range(RUNS2):\n", + " t = 0\n", + " while True:\n", + " t += 1\n", + " if rng.randrange(K2) == needle: break\n", + " ff.append(t)\n", + "m2 = sum(ff) / RUNS2; se2 = (sum((t - m2) ** 2 for t in ff) / (RUNS2 - 1)) ** 0.5 / RUNS2 ** 0.5\n", + "print(f'needle p=1/{K2}: mean tests to failure {m2:.1f} +- {se2:.1f} ; theory {K2}')\n", + "assert abs(m2 - K2) < 4 * se2\n", + "zmax = 0.0\n", + "for N in (256, 1024, 2048, 3072, 4096):\n", + " emp = sum(t > N for t in ff) / RUNS2; th = (1 - p2) ** N\n", + " z = abs(emp - th) / max((th * (1 - th) / RUNS2) ** 0.5, 1e-9); zmax = max(zmax, z)\n", + " print(f' N={N:5d} P(miss) empirical {emp:.4f} (1-p)^N {th:.4f} z={z:.2f}')\n", + "assert zmax < 4\n", + "print('=> to be 99% sure of seeing a 1-in-1024 bug you need', math.ceil(math.log(0.01) / math.log(1 - p2)), 'random tests; testing is a *statistical* instrument.')" + ] + }, + { + "cell_type": "markdown", + "id": "c1f9f5f8", + "metadata": {}, + "source": [ + "## A3. Shrinking: from a long counterexample to a minimal one\n", + "Candidates: delete one element; move one element one step towards 0. Accept the first that still fails; stop when none does.\n", + "Termination: each accepted step strictly decreases the measure `(len, sum |v|)` in the lexicographic (well-founded) order.\n", + "A local minimum is not automatically a global one; compare with the exhaustive global minimum." + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "2c7b4fa0", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:54:32.011875Z", + "iopub.status.busy": "2026-10-14T09:54:32.011718Z", + "iopub.status.idle": "2026-10-14T09:54:32.226368Z", + "shell.execute_reply": "2026-10-14T09:54:32.225418Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "global minimal counterexamples (exhaustive): [[0, 0]]\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "2000 shrinks: always a failing, locally minimal list; GLOBAL minimum reached in 401/2000 = 0.201; mean length before shrinking 4.57\n", + "final shrunk values: {(-2, -2): 420, (2, 2): 408, (0, 0): 401, (1, 1): 388, (-1, -1): 383}\n", + "local minima that are NOT global (e.g. [1,1]: deleting gives [1] (passes), moving towards 0 gives [0,1] (passes)): [[1, 1], [-1, -1], [2, 2], [-2, -2]]\n" + ] + } + ], + "source": [ + "def shrink_candidates(xs):\n", + " for k in range(len(xs)): yield xs[:k] + xs[k + 1:]\n", + " for k, v in enumerate(xs):\n", + " if v != 0: yield xs[:k] + [v - (1 if v > 0 else -1)] + xs[k + 1:]\n", + "def measure(xs): return (len(xs), sum(abs(v) for v in xs))\n", + "def shrink(xs, prop):\n", + " steps = 0\n", + " while True:\n", + " for c in shrink_candidates(xs):\n", + " if not prop(c):\n", + " assert measure(c) < measure(xs) # strict decrease\n", + " xs, steps = c, steps + 1; break\n", + " else:\n", + " return xs, steps\n", + "global_min = min(map(measure, fails)); gmin_lists = sorted([f for f in fails if measure(f) == global_min])\n", + "print('global minimal counterexamples (exhaustive):', gmin_lists)\n", + "hit = tot = 0; lens_before = []; finals = collections.Counter()\n", + "rng = random.Random(11)\n", + "for _ in range(2000):\n", + " while True:\n", + " xs = gen(rng)\n", + " if not prop_bug(xs): break\n", + " ys, _ = shrink(xs, prop_bug)\n", + " assert not prop_bug(ys) and all(prop_bug(c) for c in shrink_candidates(ys)) # still failing AND locally minimal\n", + " assert measure(ys) >= global_min\n", + " lens_before.append(len(xs)); tot += 1; hit += measure(ys) == global_min; finals[tuple(ys)] += 1\n", + "print(f'2000 shrinks: always a failing, locally minimal list; GLOBAL minimum reached in {hit}/{tot} = {hit/tot:.3f}; mean length before shrinking {sum(lens_before)/tot:.2f}')\n", + "print('final shrunk values:', dict(sorted(finals.items(), key=lambda kv: -kv[1])))\n", + "print('local minima that are NOT global (e.g. [1,1]: deleting gives [1] (passes), moving towards 0 gives [0,1] (passes)):',\n", + " [list(k) for k in finals if measure(list(k)) != global_min])" + ] + }, + { + "cell_type": "markdown", + "id": "b80dcca8", + "metadata": {}, + "source": [ + "## B. Small-scope hypothesis: enumeration vs uniform random sampling\n", + "A \"needle\" bug at one input `x0` of a `K = 2^16`-element domain. Uniform random with `N` tests detects with probability `1-(1-1/K)^N`\n", + "if the needle is uniformly placed; enumeration `0,1,...,N-1` detects iff `x0 < N`. If the bug location is `x0 ~ Geometric(q)`\n", + "(bugs cluster at small inputs, the empirical small-scope hypothesis), enumeration detects with probability `1-(1-q)^N` while random sampling gets\n", + "`E[1-(1-1/K)^N]` = tiny. Both closed forms, checked by Monte Carlo." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "d538673e", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:54:32.227875Z", + "iopub.status.busy": "2026-10-14T09:54:32.227713Z", + "iopub.status.idle": "2026-10-14T09:54:32.772482Z", + "shell.execute_reply": "2026-10-14T09:54:32.771731Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "uniform needle: enumeration 0.0036 (theory 0.0031) random 0.0027 (theory 0.0030)\n", + "clustered needle: enumeration 0.9823 (theory 0.9824) random theory 0.0030 -> enumeration wins by x322\n", + "32-bit needle, delta = 1e-3: N >= 2.967e+10 uniform tests (no test budget finds it; a proof or a solver does)\n" + ] + } + ], + "source": [ + "K = 2 ** 16; N = 200; MC = 40000\n", + "rng = random.Random(5)\n", + "# uniform placement\n", + "enum_u = sum(rng.randrange(K) < N for _ in range(MC)) / MC\n", + "rand_u = sum(any(rng.randrange(K) == x0 for _ in range(N)) for x0 in (rng.randrange(K) for _ in range(4000))) / 4000\n", + "th_enum_u, th_rand_u = N / K, 1 - (1 - 1 / K) ** N\n", + "print(f'uniform needle: enumeration {enum_u:.4f} (theory {th_enum_u:.4f}) random {rand_u:.4f} (theory {th_rand_u:.4f})')\n", + "# geometric placement with q = 0.02\n", + "q = 0.02\n", + "def geom(r):\n", + " x = 0\n", + " while r.random() > q: x += 1\n", + " return min(x, K - 1)\n", + "enum_g = sum(geom(rng) < N for _ in range(MC)) / MC\n", + "th_enum_g = 1 - (1 - q) ** N\n", + "th_rand_g = 1 - (1 - 1 / K) ** N # uniform sampling detects a needle anywhere with the same probability\n", + "print(f'clustered needle: enumeration {enum_g:.4f} (theory {th_enum_g:.4f}) random theory {th_rand_g:.4f} -> enumeration wins by x{th_enum_g/th_rand_g:.0f}')\n", + "se = lambda p_: (p_ * (1 - p_) / MC) ** 0.5\n", + "assert abs(enum_g - th_enum_g) < 4 * se(th_enum_g) and abs(enum_u - th_enum_u) < 4 * se(th_enum_u) + 1e-3\n", + "assert abs(rand_u - th_rand_u) < 4 * (th_rand_u * (1 - th_rand_u) / 4000) ** 0.5 + 1e-3\n", + "# a 32-bit needle: uniform random needs ~ ln(1/delta) * 2^32 tests\n", + "print('32-bit needle, delta = 1e-3: N >=', f'{math.log(1e-3) / math.log(1 - 2.0**-32):.3e}', 'uniform tests (no test budget finds it; a proof or a solver does)')" + ] + }, + { + "cell_type": "markdown", + "id": "ae0e1586", + "metadata": {}, + "source": [ + "## C. Lean: three tools, three kinds of answer\n", + "* `decide`: exhaustive over a *bounded decidable* statement => a **proof** or a definite refutation (complete within the bound, kernel-checked).\n", + "* `slim_check` (Plausible's ancestor on this pin): random testing => a **counterexample** or *nothing* (and when nothing, it closes the goal with `sorry`!).\n", + "* `omega`: a decision procedure for linear integer arithmetic => proof, or a failure message listing \"a possible counterexample\" **constraints**." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "f695dda3", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:54:32.774131Z", + "iopub.status.busy": "2026-10-14T09:54:32.773969Z", + "iopub.status.idle": "2026-10-14T09:54:37.428445Z", + "shell.execute_reply": "2026-10-14T09:54:37.427528Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[ok] no messages\n" + ] + }, + { + "data": { + "text/plain": [ + "{'env': 2}" + ] + }, + "execution_count": 7, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "sys.path.insert(0, os.path.abspath('.'))\n", + "from _lean_helper import Lean\n", + "import re\n", + "L = Lean()\n", + "L.show(r'''\n", + "def search (P : Nat → Bool) (len : Nat) (n : Nat) : Nat :=\n", + " if n < len then (if P n then n else search P len (n + 1)) else n\n", + "termination_by len - n\n", + "def Post (P : Nat → Bool) (len r : Nat) : Prop :=\n", + " r ≤ len ∧ (r = len ∨ P r = true) ∧ ∀ i, i < r → P i = false\n", + "instance (P : Nat → Bool) (len r : Nat) : Decidable (Post P len r) := by unfold Post; infer_instance\n", + "def searchBadF (P : Nat → Bool) (len : Nat) : Nat → Nat → Nat\n", + " | 0, n => n\n", + " | fuel+1, n => if n + 1 < len then (if P n then n else searchBadF P len fuel (n+1)) else n\n", + "''', keep=True)" + ] + }, + { + "cell_type": "markdown", + "id": "ce429a6a", + "metadata": {}, + "source": [ + "### C1. A counterexample is a proof of the negation (certified bug)\n", + "1. `slim_check` finds a violating input for the buggy loop (random).\n", + "2. Instantiate it and let `decide` evaluate: the kernel certifies `not Post ...` for that instance.\n", + "3. Conclude `not (forall P len, Post ...)`: a kernel-checked refutation with **no axioms**." + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "4e9cab8a", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:54:37.430013Z", + "iopub.status.busy": "2026-10-14T09:54:37.429851Z", + "iopub.status.idle": "2026-10-14T09:54:37.738292Z", + "shell.execute_reply": "2026-10-14T09:54:37.737351Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "--- 1. random search ---\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[error] line 1: ===================\n", + "Found problems!\n", + "k := 2\n", + "len := 2\n", + "issue: ⋯ does not hold\n", + "(0 shrinks)\n", + "-------------------\n", + "--- 2. kernel-certified instance, 3. refutation of the universal spec ---\n", + "[info] line 7: 'searchBadF_not_correct' does not depend on any axioms\n" + ] + } + ], + "source": [ + "print('--- 1. random search ---')\n", + "r1 = L.show(\"example : ∀ (k len : ℕ), Post (fun i => i == k) len (searchBadF (fun i => i == k) len len 0) := by slim_check\")\n", + "assert any('Found problems' in m['data'] for m in r1['messages'])\n", + "print('--- 2. kernel-certified instance, 3. refutation of the universal spec ---')\n", + "r2 = L.show(r'''\n", + "theorem bug_witness : ¬ Post (fun _ => false) 10 (searchBadF (fun _ => false) 10 20 0) := by decide\n", + "theorem searchBadF_not_correct : ¬ (∀ (P : Nat → Bool) (len : Nat), Post P len (searchBadF P len len 0)) := fun h => by\n", + " have := h (fun _ => false) 10\n", + " revert this\n", + " decide\n", + "#print axioms searchBadF_not_correct\n", + "''', keep=True)\n", + "assert any('does not depend on any axioms' in m['data'] for m in r2['messages'])\n", + "assert not [m for m in r2['messages'] if m['severity'] == 'error']" + ] + }, + { + "cell_type": "markdown", + "id": "422be44f", + "metadata": {}, + "source": [ + "### C2. No counterexample is NOT a proof, and `slim_check` says so with `sorry`\n", + "On a TRUE statement `slim_check` reports success and **closes the goal with `sorry`**: `#print axioms` shows `sorryAx`, so the strict\n", + "verifier's axiom gate rejects it. Property testing can only produce counterexamples, never certificates." + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "f4ff841e", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:54:37.740312Z", + "iopub.status.busy": "2026-10-14T09:54:37.740138Z", + "iopub.status.idle": "2026-10-14T09:54:38.143859Z", + "shell.execute_reply": "2026-10-14T09:54:38.143005Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[info] line 1: Success\n", + "[warning] line 1: declaration uses 'sorry'\n", + "[info] line 2: 'tt' depends on axioms: [sorryAx]\n", + "=> axioms of the \"slim_check proof\": 'tt' depends on axioms: [sorryAx]\n", + "--- the same statement is PROVED by omega (a decision procedure for linear arithmetic): ---\n", + "[info] line 2: 'tt2' depends on axioms: [propext, Quot.sound]\n" + ] + } + ], + "source": [ + "r = L.show(\"theorem tt : ∀ a b : ℕ, b ≤ a → a - b + b = a := by slim_check\\n#print axioms tt\")\n", + "axs = [m['data'] for m in r['messages'] if 'depends on axioms' in m['data']]\n", + "assert axs and 'sorryAx' in axs[0]\n", + "print('=> axioms of the \"slim_check proof\":', axs[0])\n", + "print('--- the same statement is PROVED by omega (a decision procedure for linear arithmetic): ---')\n", + "r = L.show(\"theorem tt2 : ∀ a b : ℕ, b ≤ a → a - b + b = a := by omega\\n#print axioms tt2\")\n", + "assert not [m for m in r['messages'] if m['severity'] == 'error']" + ] + }, + { + "cell_type": "markdown", + "id": "c0e698ee", + "metadata": {}, + "source": [ + "### C3. A counterexample to the *statement* is not necessarily a bug in the *program*\n", + "`forall a b : N, a - b + b = a` is false (truncated subtraction): counterexample `a=0, b=2`. The program (`-`) is doing what `Nat.sub` does;\n", + "the **spec** is missing `b <= a` (note 0302, Lemma 1). Protocol step: check the counterexample against the *precondition*." + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "797ea289", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:54:38.145619Z", + "iopub.status.busy": "2026-10-14T09:54:38.145413Z", + "iopub.status.idle": "2026-10-14T09:54:38.414542Z", + "shell.execute_reply": "2026-10-14T09:54:38.413549Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[error] line 1: ===================\n", + "Found problems!\n", + "a := 0\n", + "b := 1\n", + "issue: 1 = 0 does not hold\n", + "(0 shrinks)\n", + "-------------------\n", + "python model of N-subtraction agrees: 2 != 0\n" + ] + } + ], + "source": [ + "r = L.show(\"example : ∀ a b : ℕ, a - b + b = a := by slim_check\")\n", + "assert any('Found problems' in m['data'] for m in r['messages'])\n", + "print('python model of N-subtraction agrees:', max(0, 0 - 2) + 2, '!= 0')\n", + "assert max(0, 0 - 2) + 2 == 2" + ] + }, + { + "cell_type": "markdown", + "id": "f63c9230", + "metadata": {}, + "source": [ + "### C4. A failed proof is not a bug: `omega` abstracts a nonlinear term and reports a *spurious* counterexample\n", + "`n <= n*n` is TRUE for all naturals. `omega` treats `n*n` as an atom `b` and answers with the constraints `a - b >= 1, a >= 0`; those\n", + "are satisfiable for an unconstrained atom, but `b = a*a` is not unconstrained." + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "id": "4540734b", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:54:38.416132Z", + "iopub.status.busy": "2026-10-14T09:54:38.415974Z", + "iopub.status.idle": "2026-10-14T09:54:38.447380Z", + "shell.execute_reply": "2026-10-14T09:54:38.446354Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[error] line 1: omega could not prove the goal:\n", + "a possible counterexample may satisfy the constraints\n", + " a - b ≥ 1\n", + " a ≥ 0\n", + "where\n", + " a := ↑n\n", + " b := ↑n * ↑n\n", + "exhaustive check that the \"possible counterexample\" does not exist for n < 10^5: True\n", + "[info] line 2: 'le_sq'' depends on axioms: [propext]\n" + ] + } + ], + "source": [ + "r = L.show(\"example : ∀ n : ℕ, n ≤ n * n := by omega\")\n", + "assert any('omega could not prove the goal' in m['data'] for m in r['messages'])\n", + "print('exhaustive check that the \"possible counterexample\" does not exist for n < 10^5:', all(n <= n * n for n in range(10 ** 5)))\n", + "r = L.show(\"theorem le_sq' : ∀ n : ℕ, n ≤ n * n := fun n => Nat.le_mul_self n\\n#print axioms le_sq'\")" + ] + }, + { + "cell_type": "markdown", + "id": "5d6f5442", + "metadata": {}, + "source": [ + "### C4b. `omega` on this pin is not complete even for PURE linear integer arithmetic (no atoms)\n", + "The Lean docstring (`Init/Tactics.lean`, pin v4.9.0-rc1): *\"It is not yet a full decision procedure (no \"dark\" or \"grey\" shadows)\"*. Pugh's classic system\n", + "`27 <= 11x + 13y <= 45, -10 <= 7x - 9y <= 4` has real solutions but NO integer solution. A complete procedure proves `False`; this `omega` answers\n", + "\"a possible counterexample may satisfy the constraints\" although none exists, and there are no atoms to blame." + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "id": "554118d7", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:54:38.449055Z", + "iopub.status.busy": "2026-10-14T09:54:38.448851Z", + "iopub.status.idle": "2026-10-14T09:54:38.475501Z", + "shell.execute_reply": "2026-10-14T09:54:38.474256Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[error] line 2: omega could not prove the goal:\n", + "a possible counterexample may satisfy the constraints\n", + " 27 ≤ 11*a + 13*b ≤ 45\n", + " -10 ≤ 7*a - 9*b ≤ 4\n", + "where\n", + " a := x\n", + " b := y\n", + "real vertices: [('113/190', '299/190'), ('59/38', '29/38'), ('55/38', '85/38'), ('457/190', '271/190')]\n", + "integer points in the bounding box [-1,3]x[-1,3] (which contains the whole real polytope): []\n", + "=> the system has real solutions (the polytope is non-empty) but no integer solution, yet omega reports a \"possible counterexample\": a failure message is NOT a countermodel.\n" + ] + } + ], + "source": [ + "\n", + "from fractions import Fraction as Fr\n", + "r = L.show(\"\"\"\n", + "example (x y : Int) (h1 : 27 ≤ 11*x+13*y) (h2 : 11*x+13*y ≤ 45) (h3 : -10 ≤ 7*x-9*y) (h4 : 7*x-9*y ≤ 4) : False := by omega\n", + "\"\"\")\n", + "assert any('possible counterexample' in m['data'] for m in r['messages'])\n", + "# the real polytope: vertices of u = 11x+13y in {27,45}, w = 7x-9y in {-10,4} (x = (9u+13w)/190, y = (7u-11w)/190)\n", + "verts = [(Fr(9*u + 13*w, 190), Fr(7*u - 11*w, 190)) for u in (27, 45) for w in (-10, 4)]\n", + "xs, ys = [v[0] for v in verts], [v[1] for v in verts]\n", + "lo_x, hi_x, lo_y, hi_y = int(min(xs)) - 1, int(max(xs)) + 1, int(min(ys)) - 1, int(max(ys)) + 1\n", + "pts = [(x, y) for x in range(lo_x, hi_x + 1) for y in range(lo_y, hi_y + 1) if 27 <= 11*x+13*y <= 45 and -10 <= 7*x-9*y <= 4]\n", + "print('real vertices:', [(str(a), str(b)) for a, b in verts])\n", + "print(f'integer points in the bounding box [{lo_x},{hi_x}]x[{lo_y},{hi_y}] (which contains the whole real polytope): {pts}')\n", + "assert pts == []\n", + "print('=> the system has real solutions (the polytope is non-empty) but no integer solution, yet omega reports a \"possible counterexample\": a failure message is NOT a countermodel.')\n" + ] + }, + { + "cell_type": "markdown", + "id": "8a0fc84d", + "metadata": {}, + "source": [ + "### C5. A failed VC is not a bug: the counterexample state may be unreachable (continuation of 0303 A3)\n", + "Weaken the loop invariant to `n <= len`. The lemma `forall P len n, n <= len -> Post P len (search P len n)` is FALSE\n", + "(`slim_check` finds a counterexample, below). But the program starts at `n = 0`. We check **exhaustively in a scope** (all `P : Fin 4 -> Bool`,\n", + "`len <= 4`, `n <= len`): every VC-counterexample is a state the real loop never visits, and every state it does visit satisfies the lemma." + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "id": "fcb52aca", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:54:38.477338Z", + "iopub.status.busy": "2026-10-14T09:54:38.477168Z", + "iopub.status.idle": "2026-10-14T09:54:39.860566Z", + "shell.execute_reply": "2026-10-14T09:54:39.859669Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[error] line 1: ===================\n", + "Found problems!\n", + "P := [0 ↦ true, 1 ↦ true, _ ↦ true]\n", + "len := 6\n", + "n := 5\n", + "guard: 5 ≤ 6\n", + "issue: ⋯ does not hold\n", + "(0 shrinks)\n", + "-------------------\n", + "scope: VC counterexamples (P,len,n): 111; of these reachable from n=0: 0; reachable states total 129, violating the lemma: 0\n", + "Lean count of VC counterexamples in the same scope: ['111']\n", + "=> every counterexample to the weakened verification condition is unreachable: spurious, NOT a bug in the program.\n", + "--- and the program-level statement (start n = 0) survives random testing (with the sorry caveat of C2): ---\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[info] line 1: Success\n", + "[warning] line 1: declaration uses 'sorry'\n" + ] + }, + { + "data": { + "text/plain": [ + "{'messages': [{'severity': 'info',\n", + " 'pos': {'line': 1, 'column': 76},\n", + " 'endPos': {'line': 1, 'column': 86},\n", + " 'data': 'Success'},\n", + " {'severity': 'warning',\n", + " 'pos': {'line': 1, 'column': 0},\n", + " 'endPos': {'line': 1, 'column': 7},\n", + " 'data': \"declaration uses 'sorry'\"}],\n", + " 'env': 13}" + ] + }, + "execution_count": 13, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "r = L.show(\"example : ∀ (P : Nat → Bool) (len n : Nat), n ≤ len → Post P len (search P len n) := by slim_check\")\n", + "assert any('Found problems' in m['data'] for m in r['messages'])\n", + "def trace(P, length):\n", + " n_, seen = 0, []\n", + " while n_ <= length:\n", + " seen.append(n_)\n", + " if n_ >= length or P[n_]: break\n", + " n_ += 1\n", + " return seen\n", + "def search5(P, ln, n0):\n", + " m = n0\n", + " while m < ln and not P[m]: m += 1\n", + " return m\n", + "def post5(P, ln, r_):\n", + " return r_ <= ln and (r_ == ln or P[r_]) and all((not P[i]) for i in range(r_))\n", + "vc_cex = reach = reach_bad = 0; vc_cex_reachable = 0\n", + "for length in range(5):\n", + " for bits in itertools.product([False, True], repeat=4):\n", + " tr_ = set(trace(bits, length))\n", + " for n_ in range(length + 1):\n", + " ok = post5(bits, length, search5(bits, length, n_))\n", + " if not ok:\n", + " vc_cex += 1; vc_cex_reachable += (n_ in tr_)\n", + " if n_ in tr_:\n", + " reach += 1; reach_bad += (not ok)\n", + "print(f'scope: VC counterexamples (P,len,n): {vc_cex}; of these reachable from n=0: {vc_cex_reachable}; reachable states total {reach}, violating the lemma: {reach_bad}')\n", + "assert vc_cex > 0 and vc_cex_reachable == 0 and reach_bad == 0\n", + "r = L.run('''\n", + "#eval (List.range 5).foldl (fun acc len => acc + (List.range 16).foldl (fun a m =>\n", + " let P : Nat → Bool := fun i => (m / (2 ^ i)) % 2 == 1\n", + " a + (List.range (len + 1)).foldl (fun b n => if decide (Post P len (search P len n)) then b else b + 1) 0) 0) 0\n", + "''')\n", + "lean_vc = [m['data'].strip() for m in r['messages'] if m['severity'] == 'info']\n", + "print('Lean count of VC counterexamples in the same scope:', lean_vc)\n", + "assert int(lean_vc[0]) == vc_cex\n", + "print('=> every counterexample to the weakened verification condition is unreachable: spurious, NOT a bug in the program.')\n", + "print('--- and the program-level statement (start n = 0) survives random testing (with the sorry caveat of C2): ---')\n", + "L.show(\"example : ∀ (P : Nat → Bool) (len : Nat), Post P len (search P len 0) := by slim_check\")" + ] + }, + { + "cell_type": "markdown", + "id": "92feb95d", + "metadata": {}, + "source": [ + "### C6. `decide` as model finding in a bounded scope (complete up to the bound)\n", + "Exhaustive over all predicates `P` on `Fin 4` (16 of them) and all `len <= 4`: the correct loop passes, the buggy one fails\n", + "on a specific witness. Compare with the same enumeration in Python." + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "id": "570ac396", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:54:39.862363Z", + "iopub.status.busy": "2026-10-14T09:54:39.862136Z", + "iopub.status.idle": "2026-10-14T09:54:39.961142Z", + "shell.execute_reply": "2026-10-14T09:54:39.960201Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Python exhaustive: 80 (P,len) scopes; correct loop violates 0, buggy loop violates 15; first witness ((False, False, False, False), 1)\n", + "Lean exhaustive violation counts (buggy, correct): ['15', '0']\n" + ] + } + ], + "source": [ + "def py_search(P, length, n=0):\n", + " while n < length:\n", + " if P[n]: return n\n", + " n += 1\n", + " return n\n", + "def py_bad(P, length, n=0):\n", + " while n + 1 < length:\n", + " if P[n]: return n\n", + " n += 1\n", + " return n\n", + "def py_post(P, length, r):\n", + " return r <= length and (r == length or (r < 4 and P[r])) and all((not P[i]) for i in range(r))\n", + "tot_ok = bad_cnt = 0\n", + "bad_witness = None\n", + "for length in range(5):\n", + " for bits in itertools.product([False, True], repeat=4):\n", + " assert py_post(bits, length, py_search(bits, length))\n", + " if not py_post(bits, length, py_bad(bits, length)):\n", + " bad_cnt += 1; bad_witness = bad_witness or (bits, length)\n", + " tot_ok += 1\n", + "print(f'Python exhaustive: {tot_ok} (P,len) scopes; correct loop violates 0, buggy loop violates {bad_cnt}; first witness {bad_witness}')\n", + "r = L.run('''\n", + "#eval (List.range 5).foldl (fun acc len => acc + (List.range 16).foldl (fun a m =>\n", + " let P : Nat → Bool := fun i => (m / (2 ^ i)) % 2 == 1\n", + " if decide (Post P len (searchBadF P len len 0)) then a else a + 1) 0) 0\n", + "#eval (List.range 5).foldl (fun acc len => acc + (List.range 16).foldl (fun a m =>\n", + " let P : Nat → Bool := fun i => (m / (2 ^ i)) % 2 == 1\n", + " if decide (Post P len (search P len 0)) then a else a + 1) 0) 0\n", + "''')\n", + "cnts = [m['data'].strip() for m in r['messages'] if m['severity'] == 'info']\n", + "print('Lean exhaustive violation counts (buggy, correct):', cnts)\n", + "assert int(cnts[0]) == bad_cnt and int(cnts[1]) == 0" + ] + }, + { + "cell_type": "markdown", + "id": "1547dd78", + "metadata": {}, + "source": [ + "### C7. A reachable VC violation: what it does and does not prove (Lemma 4 of the note)\n", + "Program of 0303: `i := 0; s := 0; while i < n do s := s + inc(i); i := i + 1` with `inc(i) = i` (correct) or `i + 1` (buggy), precondition `n >= 0`,\n", + "post `Q: 2s = n(n-1)`, invariant `I: 0<=i<=n and 2s = i(i-1)`. The loop-head states `(n,i,s)` visited from `n >= 0` are `Reach`. Classify every violated VC:\n", + "* **preserve VC** (`I and i I[next]`): a reachable violating state `x` has a REACHABLE successor `x'` violating `I`: the annotation (or the code) is wrong, but `Q` may still hold;\n", + "* **exit VC** (`I and not(i Q`): a reachable violating state is the state in which the loop actually EXITS: `Q` is violated by a real run: a bug witness (code or spec);\n", + "* unreachable violation: no conclusion." + ] + }, + { + "cell_type": "code", + "execution_count": 15, + "id": "85eabecd", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:54:39.962887Z", + "iopub.status.busy": "2026-10-14T09:54:39.962663Z", + "iopub.status.idle": "2026-10-14T09:54:39.984403Z", + "shell.execute_reply": "2026-10-14T09:54:39.983516Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "A weak I, right code, right Q preserve-viol 0 (reachable 0) exit-viol 504 (reachable 0) Q holds on every run: True\n", + "B right I, right code, WRONG Q (n(n+1)) preserve-viol 0 (reachable 0) exit-viol 7 (reachable 7) Q holds on every run: False\n", + "C right I, BUGGY code, right Q preserve-viol 28 (reachable 7) exit-viol 0 (reachable 0) Q holds on every run: False\n", + "D too-strong I (i<=5), right code, right Q preserve-viol 2 (reachable 2) exit-viol 0 (reachable 0) Q holds on every run: True\n", + "=> reachable preserve-violation: I is not an invariant (D: but Q still holds); reachable exit-violation: Q fails in a real run (B).\n" + ] + } + ], + "source": [ + "\n", + "def run_head_states(n, inc):\n", + " st, seen = (n, 0, 0), []\n", + " while True:\n", + " seen.append(st); nn, ii, ss = st\n", + " if not ii < nn: return seen\n", + " st = (nn, ii + 1, ss + inc(ii))\n", + "Pn = range(0, 8)\n", + "Ibox = [(n_, i_, s_) for n_ in range(0, 8) for i_ in range(-1, 9) for s_ in range(-4, 60)]\n", + "Qf = lambda n_, i_, s_: 2 * s_ == n_ * (n_ - 1)\n", + "def analyse(label, inc, I, Q):\n", + " reach = {st for n_ in Pn for st in run_head_states(n_, inc)}\n", + " step = lambda st: (st[0], st[1] + 1, st[2] + inc(st[1]))\n", + " pres = [st for st in Ibox if I(*st) and st[1] < st[0] and not I(*step(st))]\n", + " exits = [st for st in Ibox if I(*st) and not st[1] < st[0] and not Q(*st)]\n", + " pres_r = [st for st in pres if st in reach]; exits_r = [st for st in exits if st in reach]\n", + " # (a) a reachable preserve-violation has a reachable successor that violates I\n", + " assert all(step(st) in reach and not I(*step(st)) for st in pres_r)\n", + " # (b) a reachable exit-violation is the final state of a real run, and it violates Q\n", + " finals = {run_head_states(n_, inc)[-1] for n_ in Pn}\n", + " assert all(st in finals and not Q(*st) for st in exits_r)\n", + " post_ok = all(Q(*run_head_states(n_, inc)[-1]) for n_ in Pn)\n", + " print(f'{label:46s} preserve-viol {len(pres):4d} (reachable {len(pres_r):3d}) exit-viol {len(exits):4d} (reachable {len(exits_r):3d}) Q holds on every run: {post_ok}')\n", + " return len(pres), len(pres_r), len(exits), len(exits_r), post_ok\n", + "good = lambda i: i; buggy = lambda i: i + 1\n", + "I_ok = lambda n_, i_, s_: 0 <= i_ <= n_ and 2 * s_ == i_ * (i_ - 1)\n", + "I_weak = lambda n_, i_, s_: 0 <= i_ <= n_\n", + "I_tooStrong = lambda n_, i_, s_: 0 <= i_ <= n_ and 2 * s_ == i_ * (i_ - 1) and i_ <= 5\n", + "Q_wrong = lambda n_, i_, s_: 2 * s_ == n_ * (n_ + 1)\n", + "A = analyse('A weak I, right code, right Q', good, I_weak, Qf)\n", + "B = analyse('B right I, right code, WRONG Q (n(n+1))', good, I_ok, Q_wrong)\n", + "C = analyse('C right I, BUGGY code, right Q', buggy, I_ok, Qf)\n", + "D = analyse('D too-strong I (i<=5), right code, right Q', good, I_tooStrong, Qf)\n", + "assert A[2] > 0 and A[3] == 0 # exit violations exist, all unreachable: nothing follows (the 0304 C5 story)\n", + "assert B[3] > 0 and not B[4] # reachable exit violation <=> Q really fails: a bug witness (here a wrong spec)\n", + "assert C[1] > 0 and not C[4] # reachable preserve violation in the buggy program\n", + "assert D[1] > 0 and D[4] # reachable preserve violation, I is not an invariant, yet Q holds on every run\n", + "print('=> reachable preserve-violation: I is not an invariant (D: but Q still holds); reachable exit-violation: Q fails in a real run (B).')\n" + ] + }, + { + "cell_type": "markdown", + "id": "2da08f1c", + "metadata": {}, + "source": [ + "## The protocol, executed on one example\n", + "1. (counterexample) `slim_check`/`decide` gives `P = const false, len = 10` (or the shrunk `k=2, len=1`).\n", + "2. (precondition) Is it admissible: `Post` has no precondition other than `len`/`P` arbitrary -> yes.\n", + "3. (reachable) the loop is started at `n=0`: the witness is a top-level input, not an internal state -> reachable.\n", + "4. (evaluate) `#eval searchBadF (fun _ => false) 10 20 0 = 9`, spec demands `10`.\n", + "5. (intent) does the spec say what we want? Dafny's LinearSearch: `n = len` when nothing matches -> yes.\n", + "6. (certify) `searchBadF_not_correct` is a kernel-checked theorem with no axioms." + ] + }, + { + "cell_type": "code", + "execution_count": 16, + "id": "1eb263c7", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:54:39.986228Z", + "iopub.status.busy": "2026-10-14T09:54:39.986009Z", + "iopub.status.idle": "2026-10-14T09:54:40.253611Z", + "shell.execute_reply": "2026-10-14T09:54:40.252685Z" + } + }, + "outputs": [], + "source": [ + "L.close()" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/code/notebooks/0305-spec-risk.ipynb b/code/notebooks/0305-spec-risk.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..be5befbac850b03372ddee356c943765fb8cf912 --- /dev/null +++ b/code/notebooks/0305-spec-risk.ipynb @@ -0,0 +1,789 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "9bfd3f75", + "metadata": {}, + "source": [ + "# 0305 - Spec generation risk: the proof of the wrong theorem\n", + "\n", + "Scratch-pad for note `0305 - Spec Generation Risk.pdf`. Naive first, efficient second, assert agreement.\n", + "\n", + "* **A. Spec validation by tests, VERINA-style.** NAIVE = decide soundness/completeness of a candidate post-condition by EXHAUSTIVE\n", + " comparison with the ground truth over a bounded domain. EFFICIENT = decide it from positive / negative tests only (the reduction\n", + " of VERINA section 4.1). Assert: the test-based verdict is never wrong in the \"refuted\" direction, and equals the exhaustive one with full tests.\n", + "* **A2. Mutation testing of a spec**: kill probability `1-(1-kappa)^m` per mutant, checked by Monte Carlo.\n", + "* **B. Lean**: a \"verified\" function that does nothing useful; a spec that leaks the implementation; a vacuous statement; and the\n", + " **statement-rewrite** finding of this repository (docs/P2.md item 18) reproduced against four gates, using `align.verify.lean` (<= 2 workers, nice 15).\n", + "* **C. Risk model**: `P(spec wrong | proof accepted)` by Bayes, checked by Monte Carlo.\n", + "\n", + "No GPU." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "6dc6c1b4", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:54:42.902895Z", + "iopub.status.busy": "2026-10-14T09:54:42.902750Z", + "iopub.status.idle": "2026-10-14T09:54:42.912430Z", + "shell.execute_reply": "2026-10-14T09:54:42.911655Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "40 lists; 1600 input/output pairs\n" + ] + } + ], + "source": [ + "import os\n", + "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"\"\n", + "import itertools, random, math, sys, collections\n", + "random.seed(305)\n", + "\n", + "# ---- domain: lists of length <= 3 over {0,1,2} --------------------------------------------------------------\n", + "V = (0, 1, 2)\n", + "LISTS = [tuple(t) for n in range(4) for t in itertools.product(V, repeat=n)]\n", + "PAIRS = [(x, y) for x in LISTS for y in LISTS]\n", + "print(len(LISTS), 'lists;', len(PAIRS), 'input/output pairs')\n", + "\n", + "# atoms an LLM-written spec is built from\n", + "srt = lambda y: all(a <= b for a, b in zip(y, y[1:]))\n", + "srt_str = lambda y: all(a < b for a, b in zip(y, y[1:]))\n", + "srt_rev = lambda y: all(a >= b for a, b in zip(y, y[1:]))\n", + "perm = lambda x, y: collections.Counter(x) == collections.Counter(y)\n", + "len_eq = lambda x, y: len(x) == len(y)\n", + "sum_eq = lambda x, y: sum(x) == sum(y)\n", + "subset = lambda x, y: set(y) <= set(x)\n", + "leak = lambda x, y: y == tuple(sorted(x))\n", + "\n", + "Q = lambda x, y: srt(y) and perm(x, y) # ground-truth post-condition (precondition: True)\n", + "CANDIDATES = {\n", + " 'Q0 sorted & perm (truth)': lambda x, y: srt(y) and perm(x, y),\n", + " 'S1 sorted only': lambda x, y: srt(y),\n", + " 'S2 sorted & perm & len<=2': lambda x, y: srt(y) and perm(x, y) and len(y) <= 2,\n", + " 'S3 sorted & same len & same sum': lambda x, y: srt(y) and len_eq(x, y) and sum_eq(x, y),\n", + " 'S4 strictly sorted & perm': lambda x, y: srt_str(y) and perm(x, y),\n", + " 'S5 sorted & y == sorted(x) (leaks)': lambda x, y: srt(y) and leak(x, y),\n", + " 'S6 perm only': lambda x, y: perm(x, y),\n", + " 'S7 sorted & same len & y subset x': lambda x, y: srt(y) and len_eq(x, y) and subset(x, y),\n", + "}" + ] + }, + { + "cell_type": "markdown", + "id": "e791e799", + "metadata": {}, + "source": [ + "## A1. Soundness / completeness of a candidate spec\n", + "VERINA (arXiv 2505.23135 section 4.1): with pre-condition `P = True`, `Q^` is **sound** iff `forall x y. Q^(x,y) => Q(x,y)` (it accepts only correct outputs)\n", + "and **complete** iff `forall x y. Q(x,y) => Q^(x,y)` (it accepts every correct output).\n", + "NAIVE: check both over all 1600 pairs. EFFICIENT (their reduction): soundness needs only the *negative* tests (pairs with `not Q`), completeness only the *positive* ones." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "35f1c93e", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:54:42.914044Z", + "iopub.status.busy": "2026-10-14T09:54:42.913844Z", + "iopub.status.idle": "2026-10-14T09:54:43.013138Z", + "shell.execute_reply": "2026-10-14T09:54:43.012157Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "positive pairs: 40 negative pairs: 1560\n", + "candidate sound complete accepted-negatives rejected-positives\n", + "Q0 sorted & perm (truth) True True 0 0\n", + "S1 sorted only False True 760 0\n", + "S2 sorted & perm & len<=2 True False 0 27\n", + "S3 sorted & same len & same sum False True 22 0\n", + "S4 strictly sorted & perm True False 0 24\n", + "S5 sorted & y == sorted(x) (leaks) True True 0 0\n", + "S6 perm only False True 72 0\n", + "S7 sorted & same len & y subset x False True 120 0\n", + "definition == test-based verdict for all 8 candidates (exact).\n" + ] + } + ], + "source": [ + "POS = [(x, y) for x, y in PAIRS if Q(x, y)]\n", + "NEG = [(x, y) for x, y in PAIRS if not Q(x, y)]\n", + "print('positive pairs:', len(POS), ' negative pairs:', len(NEG))\n", + "table = {}\n", + "for name, Qh in CANDIDATES.items():\n", + " # NAIVE: the definition, over every pair\n", + " sound_def = all((not Qh(x, y)) or Q(x, y) for x, y in PAIRS)\n", + " complete_def = all((not Q(x, y)) or Qh(x, y) for x, y in PAIRS)\n", + " # EFFICIENT: tests only (full suites here) -- soundness: Q^ rejects every negative test; completeness: Q^ accepts every positive test\n", + " a = sum(1 for x, y in NEG if Qh(x, y)) # accepted negatives = soundness violations\n", + " b = sum(1 for x, y in POS if not Qh(x, y)) # rejected positives = completeness violations\n", + " sound_t, complete_t = (a == 0), (b == 0)\n", + " assert (sound_def, complete_def) == (sound_t, complete_t)\n", + " table[name] = (sound_t, complete_t, a, b)\n", + "print(f'{\"candidate\":38s} sound complete accepted-negatives rejected-positives')\n", + "for name, (s_, c_, a_, b_) in table.items(): print(f'{name:38s} {str(s_):5s} {str(c_):8s} {a_:15d} {b_:17d}')\n", + "assert table['Q0 sorted & perm (truth)'][:2] == (True, True) and table['S5 sorted & y == sorted(x) (leaks)'][:2] == (True, True)\n", + "print('definition == test-based verdict for all', len(CANDIDATES), 'candidates (exact).')" + ] + }, + { + "cell_type": "markdown", + "id": "8c5cdb33", + "metadata": {}, + "source": [ + "### A1b. A wrong *program* passes a wrong *spec* (the proof of the wrong theorem, at toy scale)\n", + "Programs `f : list -> list`. \"f satisfies spec S\" means `forall x. S(x, f(x))`. Truth = f is a correct sort." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "2302af24", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:54:43.014728Z", + "iopub.status.busy": "2026-10-14T09:54:43.014577Z", + "iopub.status.idle": "2026-10-14T09:54:43.023261Z", + "shell.execute_reply": "2026-10-14T09:54:43.022309Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "program Q0 S1 S2 S3 S4 S5 S6 S7\n", + "sorted (correct) PASS PASS . PASS . PASS PASS PASS\n", + "empty list . PASS . . . . . . \n", + "identity . . . . . . PASS . \n", + "sorted dedup . PASS . . . . . . \n", + "sorted, replace each v by min(v,1) . PASS . . . . . . \n", + "sorted descending . . . . . . PASS . \n", + "sort then zero the first element . PASS . . . . . . \n", + "wrong programs accepted by each spec: {'S1 sorted only': 4, 'S6 perm only': 2}\n" + ] + } + ], + "source": [ + "PROGRAMS = {\n", + " 'sorted (correct)': lambda x: tuple(sorted(x)),\n", + " 'empty list': lambda x: (),\n", + " 'identity': lambda x: x,\n", + " 'sorted dedup': lambda x: tuple(sorted(set(x))),\n", + " 'sorted, replace each v by min(v,1)': lambda x: tuple(sorted(min(v, 1) for v in x)),\n", + " 'sorted descending': lambda x: tuple(sorted(x, reverse=True)),\n", + " 'sort then zero the first element': lambda x: tuple(sorted(x))[:0] + ((0,) + tuple(sorted(x))[1:] if x else ()),\n", + "}\n", + "def passes(spec, f): return all(spec(x, f(x)) for x in LISTS)\n", + "print(f'{\"program\":38s}' + ' '.join(f'{n.split()[0]:>4s}' for n in CANDIDATES))\n", + "truth_pass = {pn: passes(Q, f) for pn, f in PROGRAMS.items()}\n", + "false_accepts = collections.Counter()\n", + "for pn, f in PROGRAMS.items():\n", + " row = []\n", + " for sn, S in CANDIDATES.items():\n", + " ok = passes(S, f); row.append(ok)\n", + " if ok and not truth_pass[pn]: false_accepts[sn] += 1\n", + " print(f'{pn:38s}' + ' '.join(f'{\"PASS\" if o else \" . \":>4s}' for o in row))\n", + "assert false_accepts['S1 sorted only'] >= 3 # empty, dedup, min-clamp, zero-first, ... all pass the weak spec\n", + "assert false_accepts['Q0 sorted & perm (truth)'] == 0\n", + "print('wrong programs accepted by each spec:', dict(false_accepts))" + ] + }, + { + "cell_type": "markdown", + "id": "4371ed45", + "metadata": {}, + "source": [ + "## A2. Mutation testing of a spec\n", + "Mutants of `Q0 = sorted(y) and perm(x,y)`: drop an atom, weaken/strengthen an atom. A test suite `T` of `m` pairs *kills* a mutant `M` iff some test has `M(x,y) != Q0(x,y)`.\n", + "Let `kappa(M) = #{(x,y): M != Q0} / 1600`. Uniform tests: `P(kill) = 1 - (1-kappa)^m`. `kappa = 0` means an *equivalent* mutant (unkillable)." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "aeda6210", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:54:43.024893Z", + "iopub.status.busy": "2026-10-14T09:54:43.024737Z", + "iopub.status.idle": "2026-10-14T09:54:51.227959Z", + "shell.execute_reply": "2026-10-14T09:54:51.226999Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "mutant kappa P(kill) m=10 theory empirical z\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "drop sorted 0.0450 0.3690 0.3715 0.75\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "drop perm 0.4750 0.9984 0.9986 0.50\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "sorted: <= -> < 0.0150 0.1403 0.1448 1.84\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "sorted: <= -> >= 0.0375 0.3176 0.3192 0.49\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "perm -> same length 0.1837 0.8687 0.8689 0.08\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "perm -> same sum 0.0600 0.4614 0.4608 0.15\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "perm -> subset 0.2231 0.9199 0.9199 0.01\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "add redundant len(x)==len(y) 0.0000 0.0000 0.0000 0.00\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "perm -> y == sorted(x) 0.0000 0.0000 0.0000 0.00\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "add len(y)<=2 0.0169 0.1565 0.1540 0.97\n", + "equivalent (unkillable) mutants: ['add redundant len(x)==len(y)', 'perm -> y == sorted(x)']\n", + "mutation score with the exhaustive suite: 8/8 non-equivalent killed; worst |z| of the kill-probability law = 1.84\n" + ] + } + ], + "source": [ + "MUTANTS = {\n", + " 'drop sorted': lambda x, y: perm(x, y),\n", + " 'drop perm': lambda x, y: srt(y),\n", + " 'sorted: <= -> <': lambda x, y: srt_str(y) and perm(x, y),\n", + " 'sorted: <= -> >=': lambda x, y: srt_rev(y) and perm(x, y),\n", + " 'perm -> same length': lambda x, y: srt(y) and len_eq(x, y),\n", + " 'perm -> same sum': lambda x, y: srt(y) and sum_eq(x, y),\n", + " 'perm -> subset': lambda x, y: srt(y) and subset(x, y),\n", + " 'add redundant len(x)==len(y)': lambda x, y: srt(y) and perm(x, y) and len_eq(x, y), # equivalent: perm implies equal length\n", + " 'perm -> y == sorted(x)': lambda x, y: srt(y) and leak(x, y), # equivalent\n", + " 'add len(y)<=2': lambda x, y: srt(y) and perm(x, y) and len(y) <= 2,\n", + "}\n", + "kappa = {n: sum(1 for x, y in PAIRS if M(x, y) != Q(x, y)) / len(PAIRS) for n, M in MUTANTS.items()}\n", + "rng = random.Random(9); M_SIZE = 10; SUITES = 20000\n", + "print(f'{\"mutant\":32s} {\"kappa\":>8s} {\"P(kill) m=10 theory\":>20s} {\"empirical\":>10s} {\"z\":>5s}')\n", + "worst = 0.0\n", + "for n, M in MUTANTS.items():\n", + " kills = 0\n", + " for _ in range(SUITES):\n", + " suite = [PAIRS[rng.randrange(len(PAIRS))] for _ in range(M_SIZE)]\n", + " kills += any(M(x, y) != Q(x, y) for x, y in suite)\n", + " emp = kills / SUITES; th = 1 - (1 - kappa[n]) ** M_SIZE\n", + " se = max((th * (1 - th) / SUITES) ** 0.5, 1e-9); z = abs(emp - th) / se if th not in (0.0,) else 0.0\n", + " worst = max(worst, z)\n", + " print(f'{n:32s} {kappa[n]:8.4f} {th:20.4f} {emp:10.4f} {z:5.2f}')\n", + "assert worst < 4.0\n", + "equiv = [n for n, k in kappa.items() if k == 0]\n", + "print('equivalent (unkillable) mutants:', equiv)\n", + "assert set(equiv) == {'add redundant len(x)==len(y)', 'perm -> y == sorted(x)'}\n", + "# exhaustive suite kills every non-equivalent mutant\n", + "assert all(any(M(x, y) != Q(x, y) for x, y in PAIRS) for n, M in MUTANTS.items() if n not in equiv)\n", + "print(f'mutation score with the exhaustive suite: {(len(MUTANTS) - len(equiv))}/{len(MUTANTS) - len(equiv)} non-equivalent killed; worst |z| of the kill-probability law = {worst:.2f}')" + ] + }, + { + "cell_type": "markdown", + "id": "76064a25", + "metadata": {}, + "source": [ + "## B. Lean: proof of the wrong theorem, four ways\n", + "ONE worker via `align.verify.lean` for the experiments; a second (the strict verifier's) for B4. Total <= 2 workers, nice 15." + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "cde1b268", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:54:51.229911Z", + "iopub.status.busy": "2026-10-14T09:54:51.229754Z", + "iopub.status.idle": "2026-10-14T09:54:55.294102Z", + "shell.execute_reply": "2026-10-14T09:54:55.293131Z" + } + }, + "outputs": [], + "source": [ + "sys.path.insert(0, os.path.abspath('.'))\n", + "from _lean_helper import Lean\n", + "L = Lean()" + ] + }, + { + "cell_type": "markdown", + "id": "46718286", + "metadata": {}, + "source": [ + "### B1. A weak specification is satisfied by a useless function (and Lean happily certifies it)\n", + "Intended spec: `Sorted /\\ Perm`. A weak spec: `Sorted` only. `badSort := fun _ => []` is \"verified\" against the weak spec; the full spec refutes it." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "b665be7a", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:54:55.295817Z", + "iopub.status.busy": "2026-10-14T09:54:55.295655Z", + "iopub.status.idle": "2026-10-14T09:54:55.352936Z", + "shell.execute_reply": "2026-10-14T09:54:55.352189Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[warning] line 2: unused variable `xs`\n", + "note: this linter can be disabled with `set_option linter.unusedVariables false`\n", + "[info] line 4: 'badSort_weak_spec' does not depend on any axioms\n", + "[info] line 8: 'badSort_not_perm' depends on axioms: [propext]\n" + ] + } + ], + "source": [ + "r = L.show(r'''\n", + "def badSort (xs : List ℕ) : List ℕ := []\n", + "theorem badSort_weak_spec (xs : List ℕ) : List.Sorted (· ≤ ·) (badSort xs) := List.sorted_nil\n", + "#print axioms badSort_weak_spec\n", + "theorem badSort_not_perm : ¬ (∀ xs : List ℕ, List.Perm xs (badSort xs)) := fun h => by\n", + " have := (h [1]).length_eq\n", + " simp [badSort] at this\n", + "#print axioms badSort_not_perm\n", + "''', keep=True)\n", + "assert not [m for m in r.get('messages', []) if m['severity'] == 'error']" + ] + }, + { + "cell_type": "markdown", + "id": "5f267817", + "metadata": {}, + "source": [ + "### B2. A specification that leaks the implementation (CLEVER, arXiv 2505.13938, Fig. 2)\n", + "Computable spec = the recursion itself: the proof is `rfl`-level, no mathematics. A relational (non-computable) spec forces an actual induction." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "f9513e46", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:54:55.354522Z", + "iopub.status.busy": "2026-10-14T09:54:55.354372Z", + "iopub.status.idle": "2026-10-14T09:54:55.422221Z", + "shell.execute_reply": "2026-10-14T09:54:55.421465Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[ok] no messages\n" + ] + } + ], + "source": [ + "r = L.show(r'''\n", + "def myfib : ℕ → ℕ | 0 => 0 | 1 => 1 | (n+2) => myfib n + myfib (n+1)\n", + "-- (a) the spec IS the implementation\n", + "def spec_leak (impl : ℕ → ℕ) (n : ℕ) : Prop :=\n", + " impl n = match n with | 0 => 0 | 1 => 1 | (n'+2) => impl n' + impl (n'+1)\n", + "theorem leak_proof : ∀ n, spec_leak myfib n := by\n", + " intro n; match n with | 0 => rfl | 1 => rfl | n'+2 => rfl\n", + "-- (b) a relational spec: needs induction\n", + "inductive IsFib : ℕ → ℕ → Prop\n", + " | zero : IsFib 0 0\n", + " | one : IsFib 1 1\n", + " | step : ∀ n a b, IsFib n a → IsFib (n+1) b → IsFib (n+2) (a+b)\n", + "theorem myfib_isfib : ∀ n, IsFib n (myfib n) ∧ IsFib (n+1) (myfib (n+1)) := by\n", + " intro n; induction n with\n", + " | zero => exact ⟨IsFib.zero, IsFib.one⟩\n", + " | succ k ih => exact ⟨ih.2, IsFib.step k _ _ ih.1 ih.2⟩\n", + "''', keep=True)\n", + "assert not [m for m in r.get('messages', []) if m['severity'] == 'error']" + ] + }, + { + "cell_type": "markdown", + "id": "ae047cf3", + "metadata": {}, + "source": [ + "### B3. Vacuous hypotheses, and a probe that detects them\n", + "`h0 : n < 5`, `h1 : n > 7` are contradictory; any conclusion follows. Probe: is `exists n < 100, hypotheses` satisfiable? (A scoped *satisfiability witness* search; ProofGate-style negation probes are the dual.)" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "a201ade3", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:54:55.423774Z", + "iopub.status.busy": "2026-10-14T09:54:55.423624Z", + "iopub.status.idle": "2026-10-14T09:54:55.506541Z", + "shell.execute_reply": "2026-10-14T09:54:55.505667Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[info] line 3: 'vac_stmt' depends on axioms: [propext, Quot.sound]\n", + "[info] line 4: false\n", + "[info] line 5: true\n", + "python probe agrees: the contradictory hypotheses have no witness in the scope; the consistent pair does.\n" + ] + } + ], + "source": [ + "r = L.show(r'''\n", + "theorem vac_stmt (n : ℕ) (h₀ : n < 5) (h₁ : n > 7) : n = 12345 := by omega\n", + "#print axioms vac_stmt\n", + "#eval decide (∃ n < (100:ℕ), n < 5 ∧ n > 7)\n", + "#eval decide (∃ n < (100:ℕ), n < 10 ∧ n > 7)\n", + "''')\n", + "assert not [m for m in r.get('messages', []) if m['severity'] == 'error']\n", + "probe = lambda hyps: any(all(h(n_) for h in hyps) for n_ in range(100))\n", + "assert probe([lambda n_: n_ < 5, lambda n_: n_ > 7]) is False and probe([lambda n_: n_ < 10, lambda n_: n_ > 7]) is True\n", + "print('python probe agrees: the contradictory hypotheses have no witness in the scope; the consistent pair does.')" + ] + }, + { + "cell_type": "markdown", + "id": "94f56031", + "metadata": {}, + "source": [ + "### B4. The statement-rewrite finding (docs/P2.md section 8c item 18), reproduced against four gates\n", + "Reference: miniF2F `mathd_algebra_478` (`v = 65`). Candidates are complete files whose DISPLAYED theorem differs from the reference.\n", + "Gates: (1) **compile-only** (no error, no `sorry`); (2) **swap** (Goedel's `problem_check`: replace the displayed signature by the reference one and re-run the tactics);\n", + "(3) **strict** (token-equal signature + kernel statement bridge + axiom audit); (4) **defeq** (bridge only)." + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "dda78612", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:54:55.508277Z", + "iopub.status.busy": "2026-10-14T09:54:55.508117Z", + "iopub.status.idle": "2026-10-14T09:55:05.649327Z", + "shell.execute_reply": "2026-10-14T09:55:05.648446Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "theorem mathd_algebra_478 (b h v : ℝ) (h₀ : 0 < b ∧ 0 < h ∧ 0 < v) (h₁ : v = 1 / 3 * (b * h))\n", + " (h₂ : b = 30) (h₃ : h = 13 / 2) : v = 65\n", + "\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "candidate compile swap strict defeq strict-reasons / defeq-reasons\n", + "C0 honest True True True True [] / []\n", + "C1 displays v = 66 False True False False ['statement_mismatch'] / ['compile_error']\n", + "C2 weakened v ≤ 65 True True False False ['statement_mismatch'] / ['bridge_failed']\n", + "C3 extra false hypothesis True False False False ['statement_mismatch'] / ['bridge_failed']\n", + "C4 alpha-renamed binders True True False True ['statement_mismatch'] / []\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + "compile-only credits C2 (weaker) and C3 (vacuous) -- each is a different, easier theorem; swap credits C1 and C2 (the displayed statement is not the one proved) and refuses C3 only because its tactics name the extra hypothesis; strict rejects C1-C4; defeq accepts only the alpha-renaming.\n", + "accepted sets: {'compile': ['C0', 'C2', 'C3', 'C4'], 'swap': ['C0', 'C1', 'C2', 'C4'], 'strict': ['C0'], 'defeq': ['C0', 'C4']}\n", + "=> compile-only and swap are incomparable (C1 passes only swap, C3 only compile-only); defeq is stricter than both; strict is stricter than defeq: a partial order, not a chain.\n" + ] + } + ], + "source": [ + "import json\n", + "from align.verify.lean import ReplPool, ReplLimits, StrictVerifier, StrictConfig, find_toolchain\n", + "rows = [json.loads(l) for l in open('/apps/development/rlhf/p2_runs/r1_baseline/minif2f.jsonl')]\n", + "ref = next(r_ for r_ in rows if r_['problem_id'] == 'mathd_algebra_478')['formal_statement']\n", + "head, tail = ref.split(':= by sorry')\n", + "sig = head[head.index('theorem'):]\n", + "print(sig.strip()); print()\n", + "PROOF = 'by rw [h₁, h₂, h₃]; norm_num'\n", + "SIGS = {\n", + " 'C0 honest': sig,\n", + " 'C1 displays v = 66': sig.replace(': v = 65', ': v = 66'),\n", + " 'C2 weakened v ≤ 65': sig.replace(': v = 65', ': v ≤ 65'),\n", + " 'C3 extra false hypothesis': sig.replace('(h₃ : h = 13 / 2)', '(h₃ : h = 13 / 2) (h₄ : v < 0)'),\n", + " 'C4 alpha-renamed binders': sig.replace('(b h v : ℝ)', '(b₁ h v : ℝ)').replace('0 < b ∧', '0 < b₁ ∧').replace('(b * h)', '(b₁ * h)').replace('(h₂ : b =', '(h₂ : b₁ ='),\n", + "}\n", + "CODES = {k: head.replace(sig, s) + ':= ' + (PROOF if k != 'C3 extra false hypothesis' else 'by linarith [h₀.1, h₀.2.2, h₄]') + '\\n' for k, s in SIGS.items()}\n", + "\n", + "def gate_compile(code):\n", + " code = '\\n'.join(l for l in code.split('\\n') if not l.startswith('import ')) # the worker already imported Mathlib\n", + " r_ = L.run(code)\n", + " errs = [m for m in r_.get('messages', []) if m['severity'] == 'error']\n", + " sorry = r_.get('sorries') or any('sorry' in m['data'] for m in r_.get('messages', []))\n", + " return (not errs) and (not sorry)\n", + "def gate_swap(code):\n", + " proof = code.split(':=', 1)[1] # tactics after the displayed signature\n", + " return gate_compile(head + ':=' + proof) # re-run them against the ORIGINAL statement (Goedel's problem_check)\n", + "\n", + "pool = ReplPool.from_toolchain(find_toolchain(), n_workers=1, limits=ReplLimits(nice=15))\n", + "V_strict = StrictVerifier(pool, StrictConfig(statement_mode='strict'))\n", + "V_defeq = StrictVerifier(pool, StrictConfig(statement_mode='defeq'))\n", + "results = {}\n", + "for k, code in CODES.items():\n", + " c = gate_compile(code); sw = gate_swap(code)\n", + " rs = V_strict.verify(ref, code); rd = V_defeq.verify(ref, code)\n", + " results[k] = (c, sw, rs.accepted, rd.accepted, rs.reasons, rd.reasons)\n", + "print(f'{\"candidate\":28s} compile swap strict defeq strict-reasons / defeq-reasons')\n", + "for k, (c, sw, s_, d_, rs_, rd_) in results.items():\n", + " print(f'{k:28s} {str(c):7s}{str(sw):6s}{str(s_):7s}{str(d_):6s} {list(rs_)} / {list(rd_)}')\n", + "pool.close()\n", + "# what the table must show (and is asserted):\n", + "assert results['C0 honest'][:4] == (True, True, True, True)\n", + "assert results['C1 displays v = 66'][0] is False and results['C1 displays v = 66'][1] is True # swap harness credits a rewritten statement\n", + "assert results['C1 displays v = 66'][2] is False and results['C1 displays v = 66'][3] is False\n", + "assert results['C2 weakened v ≤ 65'][0] is True and results['C2 weakened v ≤ 65'][2] is False and results['C2 weakened v ≤ 65'][3] is False\n", + "assert results['C3 extra false hypothesis'][0] is True and results['C3 extra false hypothesis'][2] is False\n", + "assert results['C4 alpha-renamed binders'][2] is False and results['C4 alpha-renamed binders'][3] is True\n", + "print('\\ncompile-only credits C2 (weaker) and C3 (vacuous) -- each is a different, easier theorem; swap credits C1 and C2 (the displayed statement is not the one proved) and refuses C3 only because its tactics name the extra hypothesis; strict rejects C1-C4; defeq accepts only the alpha-renaming.')\n", + "\n", + "# Are the gates a CHAIN (each at least as strict as the previous)? Compare the accepted sets over these five candidates.\n", + "acc = {g: {k for k, v in results.items() if v[idx]} for g, idx in (('compile', 0), ('swap', 1), ('strict', 2), ('defeq', 3))}\n", + "print('accepted sets:', {g: sorted(s_.split()[0] for s_ in v) for g, v in acc.items()})\n", + "assert not acc['compile'] <= acc['swap'] and not acc['swap'] <= acc['compile'] # compile-only vs swap: INCOMPARABLE (C1 only swap; C3 only compile)\n", + "assert acc['strict'] < acc['defeq'] and acc['defeq'] < acc['compile'] and acc['defeq'] < acc['swap'] # defeq is stricter than both, strict stricter than defeq\n", + "print('=> compile-only and swap are incomparable (C1 passes only swap, C3 only compile-only); defeq is stricter than both; strict is stricter than defeq: a partial order, not a chain.')\n" + ] + }, + { + "cell_type": "markdown", + "id": "d02430be", + "metadata": {}, + "source": [ + "### B5. `autoImplicit`, the silent statement change (from 0302), as a spec-risk\n", + "See 0302: a typo becomes a fresh universally quantified variable. A one-line check at benchmark-build time: elaborate with `autoImplicit false`." + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "aaab2806", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:55:05.651602Z", + "iopub.status.busy": "2026-10-14T09:55:05.651417Z", + "iopub.status.idle": "2026-10-14T09:55:05.674529Z", + "shell.execute_reply": "2026-10-14T09:55:05.673620Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[warning] line 1: declaration uses 'sorry'\n", + "[info] line 2: @typo_stmt : ∀ {xx : ℕ} (x : ℕ), x + 1 > xx\n", + "[sorries] xx x : ℕ\n", + "⊢ x + 1 > xx\n", + "[error] line 2: unknown identifier 'xx'\n", + "[sorries] x : ℕ\n", + "⊢ x + 1 > sorryAx ℕ true\n" + ] + }, + { + "data": { + "text/plain": [ + "{'sorries': [{'proofState': 1,\n", + " 'pos': {'line': 2, 'column': 46},\n", + " 'goal': 'x : ℕ\\n⊢ x + 1 > sorryAx ℕ true',\n", + " 'endPos': {'line': 2, 'column': 51}}],\n", + " 'messages': [{'severity': 'error',\n", + " 'pos': {'line': 2, 'column': 37},\n", + " 'endPos': {'line': 2, 'column': 39},\n", + " 'data': \"unknown identifier 'xx'\"}],\n", + " 'env': 16}" + ] + }, + "execution_count": 10, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "L.show(\"theorem typo_stmt (x : ℕ) : x + 1 > xx := by sorry\\n#check @typo_stmt\")\n", + "L.show(\"set_option autoImplicit false in\\ntheorem typo_stmt2 (x : ℕ) : x + 1 > xx := by sorry\")" + ] + }, + { + "cell_type": "markdown", + "id": "ac95d158", + "metadata": {}, + "source": [ + "## C. Risk model: how likely is an accepted proof to be a proof of the wrong theorem?\n", + "Let a generated spec be wrong with probability `u`. A prover finds a proof with probability `a_c` for a correct spec and `a_w` for a wrong spec\n", + "(wrong = weaker/vacuous specs are easier: `r = a_w / a_c >= 1`). By Bayes\n", + "`P(wrong | accepted) = u a_w / (u a_w + (1-u) a_c) = u r / (u r + 1 - u)`. Check by Monte Carlo." + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "id": "7659f6ec", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:55:05.676725Z", + "iopub.status.busy": "2026-10-14T09:55:05.676540Z", + "iopub.status.idle": "2026-10-14T09:55:06.400102Z", + "shell.execute_reply": "2026-10-14T09:55:06.399004Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + " u r theory MC z\n", + "0.100 1.0 0.1000 0.0942 1.72\n", + "0.100 5.0 0.3571 0.3589 0.38\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "0.300 1.0 0.3000 0.3014 0.28\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "0.300 5.0 0.6818 0.6830 0.34\n", + "0.477 5.0 0.8202 0.8198 0.13\n", + "0.477 20.0 0.9480 0.9480 0.07\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "0.800 3.0 0.9231 0.9221 0.52\n", + "worst |z| = 1.72\n", + "\n", + "anchors (from the papers; see the note): VERINA best model: spec sound-and-complete 52.3% => u ~ 0.477; CLEVER: spec generation + equivalence proof solved <= 1/161 by few-shot LLMs; miniF2F-Lean revisited: formal/informal mismatch in > 1/2 of 488 problems.\n" + ] + } + ], + "source": [ + "def p_wrong_given_accepted(u, r): return u * r / (u * r + 1 - u)\n", + "rng = random.Random(3); worst = 0.0\n", + "print(f'{\"u\":>5s} {\"r\":>5s} {\"theory\":>8s} {\"MC\":>8s} {\"z\":>5s}')\n", + "for u, r_ in [(0.1, 1), (0.1, 5), (0.3, 1), (0.3, 5), (0.477, 5), (0.477, 20), (0.8, 3)]:\n", + " a_c = 0.04; a_w = min(1.0, a_c * r_)\n", + " acc = wrong = 0\n", + " for _ in range(200000):\n", + " w_ = rng.random() < u\n", + " if rng.random() < (a_w if w_ else a_c):\n", + " acc += 1; wrong += w_\n", + " mc = wrong / acc; th = p_wrong_given_accepted(u, a_w / a_c)\n", + " se = (th * (1 - th) / acc) ** 0.5; z = abs(mc - th) / se; worst = max(worst, z)\n", + " print(f'{u:5.3f} {a_w/a_c:5.1f} {th:8.4f} {mc:8.4f} {z:5.2f}')\n", + "assert worst < 4\n", + "print('worst |z| =', round(worst, 2))\n", + "print('\\nanchors (from the papers; see the note): VERINA best model: spec sound-and-complete 52.3% => u ~ 0.477; '\n", + " 'CLEVER: spec generation + equivalence proof solved <= 1/161 by few-shot LLMs; miniF2F-Lean revisited: formal/informal mismatch in > 1/2 of 488 problems.')\n", + "L.close()" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/code/notebooks/0306-benchmarks-landscape.ipynb b/code/notebooks/0306-benchmarks-landscape.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..456f2a9178b740efb61cc1f5d4734feaeb3f0e3a --- /dev/null +++ b/code/notebooks/0306-benchmarks-landscape.ipynb @@ -0,0 +1,479 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "3a638e03", + "metadata": {}, + "source": [ + "# 0306 - Verified-code benchmarks: what each measures, and how to read the numbers\n", + "\n", + "Scratch-pad for note `0306 - Verified Code Benchmarks.pdf`. The benchmarks (VERINA, CLEVER, FVAPPS, DafnyBench, miniCodeProps) differ in\n", + "*task*, *who writes the spec*, *how the spec is validated* and *what the checker trusts*; that is the table in the note. This notebook\n", + "verifies the **arithmetic** the note depends on (naive first, efficient second, assert agreement):\n", + "\n", + "1. **pass@k** (Chen et al. 2021, eq. (1), Fig. 3): NAIVE exact combinatorics; EFFICIENT stable product. Unbiasedness of the estimator\n", + " (exact sum over the binomial), and the bias of the plug-in `1-(1-p_hat)^k`.\n", + "2. **Small benchmarks**: Wilson vs Wald interval, coverage by Monte Carlo, applied to the headline numbers of each paper.\n", + "3. **Composite tasks**: Frechet bounds for `P(code and spec and proof)`; checked on random joint distributions.\n", + "4. **Consistency checks** of numbers quoted from the papers (sums, ratios) and one *derived* baseline (DafnyBench's no-hint rate).\n", + "\n", + "No GPU, no Lean in this notebook (pure arithmetic)." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "542867e2", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:53:51.893124Z", + "iopub.status.busy": "2026-10-14T09:53:51.892857Z", + "iopub.status.idle": "2026-10-14T09:53:52.032196Z", + "shell.execute_reply": "2026-10-14T09:53:52.030551Z" + } + }, + "outputs": [], + "source": [ + "import os\n", + "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"\"\n", + "import math, random, itertools\n", + "from fractions import Fraction\n", + "import numpy as np\n", + "random.seed(306); np.random.seed(306)" + ] + }, + { + "cell_type": "markdown", + "id": "3ae94319", + "metadata": {}, + "source": [ + "## 1. pass@k\n", + "Given `n` samples of which `c` are correct, a random `k`-subset contains no correct sample with probability `C(n-c,k)/C(n,k)`,\n", + "so `pass@k = 1 - C(n-c,k)/C(n,k)`. Stable form (cancel the factorials):\n", + "`C(n-c,k)/C(n,k) = prod_{i=n-c+1}^{n} (1 - k/i)` (for `n-c >= k`; else the ratio is 0)." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "806c53e0", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:53:52.034678Z", + "iopub.status.busy": "2026-10-14T09:53:52.034346Z", + "iopub.status.idle": "2026-10-14T09:53:52.050351Z", + "shell.execute_reply": "2026-10-14T09:53:52.049783Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "stable product vs exact rational on 303 (n,c,k) triples: max abs error 6.11e-16\n", + "naive float64 factorial ratio at n=200 (170! < 1.8e308 < 171!): nan (exact: 0.660225623763214 )\n", + "naive float is fine for small n: True\n" + ] + } + ], + "source": [ + "def pass_at_k_exact(n, c, k): # NAIVE: exact rational arithmetic from the definition\n", + " if n - c < k: return Fraction(1)\n", + " return 1 - Fraction(math.comb(n - c, k), math.comb(n, k))\n", + "def _ffact(m): # factorial accumulated in float64 (what a vectorised implementation does)\n", + " return float(np.prod(np.arange(1, m + 1, dtype=np.float64)))\n", + "def pass_at_k_naive_float(n, c, k): # NAIVE in floats: the textbook factorial ratio C(n-c,k)/C(n,k)\n", + " if n - c < k: return 1.0\n", + " with np.errstate(all='ignore'):\n", + " cnk = lambda a, b: _ffact(a) / (_ffact(b) * _ffact(a - b))\n", + " return 1.0 - cnk(n - c, k) / cnk(n, k)\n", + "def pass_at_k_stable(n, c, k): # EFFICIENT: Chen et al., Fig. 3\n", + " if n - c < k: return 1.0\n", + " return 1.0 - np.prod(1.0 - k / np.arange(n - c + 1, n + 1))\n", + "\n", + "worst = 0.0; count = 0\n", + "for n in (5, 10, 20, 50, 100, 170, 200, 500):\n", + " for c in sorted({0, 1, 2, n // 10, n // 3, n // 2, n - 1, n}):\n", + " for k in (1, 2, 5, 10, 50, 100):\n", + " if k > n: continue\n", + " ex = float(pass_at_k_exact(n, c, k)); st = pass_at_k_stable(n, c, k)\n", + " worst = max(worst, abs(ex - st)); count += 1\n", + "assert worst < 1e-12\n", + "print(f'stable product vs exact rational on {count} (n,c,k) triples: max abs error {worst:.2e}')\n", + "v200 = pass_at_k_naive_float(200, 20, 10)\n", + "print('naive float64 factorial ratio at n=200 (170! < 1.8e308 < 171!):', v200, ' (exact:', float(pass_at_k_exact(200, 20, 10)), ')')\n", + "assert not math.isfinite(v200)\n", + "print('naive float is fine for small n:', abs(pass_at_k_naive_float(20, 5, 3) - float(pass_at_k_exact(20, 5, 3))) < 1e-12)" + ] + }, + { + "cell_type": "markdown", + "id": "46406668", + "metadata": {}, + "source": [ + "### Unbiasedness (exact) and the bias of the plug-in estimator\n", + "For a task with true single-sample success probability `p`, `c ~ Binomial(n,p)`. Claim: `E[pass_at_k(n,c,k)] = 1-(1-p)^k`.\n", + "Proof sketch: the random `k`-subset of the `n` samples is itself `k` i.i.d. draws, so it contains no success with probability `(1-p)^k`.\n", + "Check by exact summation. The plug-in `1-(1-c/n)^k` is biased (Chen et al., Appendix A)." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "a46d5fa0", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:53:52.053085Z", + "iopub.status.busy": "2026-10-14T09:53:52.052843Z", + "iopub.status.idle": "2026-10-14T09:53:52.063079Z", + "shell.execute_reply": "2026-10-14T09:53:52.061837Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + " n p k truth 1-(1-p)^k E[unbiased] |err| E[plug-in] bias\n", + " 10 0.10 5 0.4095100000 0.4095100000 2.2e-16 0.34851642 -0.06099\n", + " 10 0.30 5 0.8319300000 0.8319300000 6.7e-16 0.75974023 -0.07219\n", + " 20 0.05 10 0.4012630608 0.4012630608 6.1e-16 0.33649458 -0.06477\n", + " 20 0.50 3 0.8750000000 0.8750000000 0.0e+00 0.85625000 -0.01875\n", + " 50 0.02 20 0.3323920282 0.3323920282 6.7e-16 0.28490846 -0.04748\n", + "100 0.01 50 0.3949939329 0.3949939329 6.7e-16 0.32746853 -0.06753\n", + "unbiasedness: max |E[estimator] - truth| = 6.7e-16; the plug-in estimator is biased (negative here: it under-estimates pass@k).\n" + ] + } + ], + "source": [ + "def expect_estimators(n, p, k):\n", + " e_unb = e_plug = 0.0\n", + " for c in range(n + 1):\n", + " w = math.comb(n, c) * p ** c * (1 - p) ** (n - c)\n", + " e_unb += w * pass_at_k_stable(n, c, k)\n", + " e_plug += w * (1 - (1 - c / n) ** k)\n", + " return e_unb, e_plug\n", + "print(f'{\"n\":>3s} {\"p\":>5s} {\"k\":>3s} {\"truth 1-(1-p)^k\":>16s} {\"E[unbiased]\":>12s} {\"|err|\":>9s} {\"E[plug-in]\":>11s} {\"bias\":>9s}')\n", + "max_err = 0.0\n", + "for n, p, k in [(10, 0.1, 5), (10, 0.3, 5), (20, 0.05, 10), (20, 0.5, 3), (50, 0.02, 20), (100, 0.01, 50)]:\n", + " truth = 1 - (1 - p) ** k\n", + " u, pl = expect_estimators(n, p, k)\n", + " max_err = max(max_err, abs(u - truth))\n", + " print(f'{n:3d} {p:5.2f} {k:3d} {truth:16.10f} {u:12.10f} {abs(u-truth):9.1e} {pl:11.8f} {pl-truth:9.5f}')\n", + "assert max_err < 1e-12\n", + "print(f'unbiasedness: max |E[estimator] - truth| = {max_err:.1e}; the plug-in estimator is biased (negative here: it under-estimates pass@k).')" + ] + }, + { + "cell_type": "markdown", + "id": "3aa2cde7", + "metadata": {}, + "source": [ + "## 2. Small benchmarks need intervals\n", + "VERINA 189 tasks, CLEVER 161, miniCodeProps 201, FVAPPS baseline 100 samples / 406 theorems, DafnyBench 782. Treat each task as a Bernoulli trial\n", + "(NAMED APPROXIMATION: tasks exchangeable and independent; FVAPPS theorems within a sample are correlated, so its interval is too narrow).\n", + "Wald: `p_hat +- 1.96 sqrt(p_hat(1-p_hat)/n)`. Wilson: solve `|p_hat - p| = z sqrt(p(1-p)/n)` for `p`:\n", + "`centre = (p_hat + z^2/2n)/(1 + z^2/n)`, `half = z sqrt(p_hat(1-p_hat)/n + z^2/4n^2)/(1 + z^2/n)`." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "dfbe9e5a", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:53:52.065104Z", + "iopub.status.busy": "2026-10-14T09:53:52.064886Z", + "iopub.status.idle": "2026-10-14T09:53:52.480161Z", + "shell.execute_reply": "2026-10-14T09:53:52.479404Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Wilson endpoints solve (p_hat - p)^2 = z^2 p(1-p)/n exactly (residual < 1e-12) for 5 test cases\n", + " n p Wald cover Wilson cover\n", + " 161 0.005 0.5518 0.9539\n", + " 161 0.030 0.8595 0.9675\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + " 189 0.049 0.9451 0.9601\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + " 201 0.350 0.9536 0.9462\n", + " 406 0.300 0.9497 0.9498\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + " 782 0.680 0.9496 0.9498\n", + "\n", + "benchmark setting rate Wilson 95% interval\n", + "CLEVER best few-shot LLM, end-to-end 0.6% [ 0.1%, 3.4%]\n", + "VERINA best model proof success, pass@1 (4.9% of 189 ~ 9.3 tasks; k rounded) 4.8% [ 2.5%, 8.8%]\n", + "VERINA best model spec sound+complete (52.3% of 189 ~ 98.9; k rounded) 52.4% [45.3%, 59.4%]\n", + "miniCodeProps GPT-4o pass@32 + refinement, overall 37.3% [30.9%, 44.2%]\n", + "miniCodeProps GPT-4o pass@32 + refinement, medium+hard 7.0% [ 3.6%, 13.1%]\n", + "FVAPPS Claude 3.5 Sonnet, theorems of 100 sampled problems 29.8% [25.6%, 34.4%]\n", + "FVAPPS Gemini 1.5 Pro 18.2% [14.8%, 22.3%]\n", + "DafnyBench best model, Table 3: 67.8% of 782 = 530 67.8% [64.4%, 71.0%]\n", + "DafnyBench no-LLM baseline, Table 3: 26.9% of 782 ~ 210 26.9% [23.9%, 30.1%]\n" + ] + } + ], + "source": [ + "Z = 1.959963984540054\n", + "def wald(k, n):\n", + " p = k / n; h = Z * math.sqrt(p * (1 - p) / n); return max(0, p - h), min(1, p + h)\n", + "def wilson(k, n):\n", + " p = k / n; d = 1 + Z * Z / n\n", + " c = (p + Z * Z / (2 * n)) / d; h = Z * math.sqrt(p * (1 - p) / n + Z * Z / (4 * n * n)) / d\n", + " return c - h, c + h\n", + "\n", + "# derivation check: the Wilson endpoints are exactly the roots of (p_hat - p)^2 = z^2 p (1-p)/n\n", + "for k, n in [(1, 161), (9, 189), (70, 201), (121, 406), (530, 782)]:\n", + " lo, hi = wilson(k, n)\n", + " for p_ in (lo, hi):\n", + " assert abs((k / n - p_) ** 2 - Z * Z * p_ * (1 - p_) / n) < 1e-12\n", + "print('Wilson endpoints solve (p_hat - p)^2 = z^2 p(1-p)/n exactly (residual < 1e-12) for 5 test cases')\n", + "\n", + "rng = np.random.default_rng(1)\n", + "def coverage(n, p, reps=40000):\n", + " ks = rng.binomial(n, p, size=reps)\n", + " cov_w = cov_l = 0\n", + " for k in ks:\n", + " lo, hi = wald(int(k), n); cov_w += (lo <= p <= hi)\n", + " lo, hi = wilson(int(k), n); cov_l += (lo <= p <= hi)\n", + " return cov_w / reps, cov_l / reps\n", + "print(f'{\"n\":>4s} {\"p\":>6s} {\"Wald cover\":>11s} {\"Wilson cover\":>13s}')\n", + "cov_table = {}\n", + "for n, p in [(161, 0.005), (161, 0.03), (189, 0.049), (201, 0.35), (406, 0.30), (782, 0.68)]:\n", + " cw, cl = coverage(n, p); cov_table[(n, p)] = (cw, cl)\n", + " print(f'{n:4d} {p:6.3f} {cw:11.4f} {cl:13.4f}')\n", + "assert cov_table[(161, 0.005)][0] < 0.6 and cov_table[(161, 0.005)][1] > 0.90 # Wald collapses at tiny p; Wilson does not\n", + "assert all(cl > 0.92 for (_, cl) in cov_table.values())\n", + "\n", + "print()\n", + "HEADLINES = [ # (benchmark, what, successes k, n) -- k reconstructed from the reported rate and n where the paper gives a rate\n", + " ('CLEVER', 'best few-shot LLM, end-to-end', 1, 161),\n", + " ('VERINA', 'best model proof success, pass@1 (4.9% of 189 ~ 9.3 tasks; k rounded)', 9, 189),\n", + " ('VERINA', 'best model spec sound+complete (52.3% of 189 ~ 98.9; k rounded)', 99, 189),\n", + " ('miniCodeProps', 'GPT-4o pass@32 + refinement, overall', 75, 201),\n", + " ('miniCodeProps', 'GPT-4o pass@32 + refinement, medium+hard', 8, 115),\n", + " ('FVAPPS', 'Claude 3.5 Sonnet, theorems of 100 sampled problems', 121, 406),\n", + " ('FVAPPS', 'Gemini 1.5 Pro', 74, 406),\n", + " ('DafnyBench', 'best model, Table 3: 67.8% of 782 = 530', 530, 782),\n", + " ('DafnyBench', 'no-LLM baseline, Table 3: 26.9% of 782 ~ 210', 210, 782),\n", + "]\n", + "print(f'{\"benchmark\":14s} {\"setting\":70s} {\"rate\":>6s} Wilson 95% interval')\n", + "for b, what, k, n in HEADLINES:\n", + " lo, hi = wilson(k, n)\n", + " print(f'{b:14s} {what:70s} {100*k/n:5.1f}% [{100*lo:4.1f}%, {100*hi:4.1f}%]')" + ] + }, + { + "cell_type": "markdown", + "id": "bc025e01", + "metadata": {}, + "source": [ + "## 3. Composite tasks: Frechet bounds\n", + "End-to-end success `A and B and C` (code, spec, proof; or CLEVER's `QED1 and QED2`) is bounded by its stages whatever the dependence:\n", + "`max(0, sum P(.) - (m-1)) <= P(all) <= min P(.)`. Check on random joint distributions of `m` binary stage outcomes." + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "8b0beb96", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:53:52.481823Z", + "iopub.status.busy": "2026-10-14T09:53:52.481664Z", + "iopub.status.idle": "2026-10-14T09:53:53.003946Z", + "shell.execute_reply": "2026-10-14T09:53:53.003015Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "20000 random joint distributions (m = 2 or 3 stages): Frechet bound violations = 0\n", + "VERINA stage marginals (code 72.6%, spec 52.3%, proof 4.9%) => P(all three) in [0.000, 0.049]\n", + "comonotone coupling attains the upper bound: 0.0491 vs min = 0.049\n", + "independent stages would give the product 0.01861 (inside the bounds; real tasks are positively dependent)\n", + "lower Frechet bound attained by an explicit coupling on 6 marginals (exact rational arithmetic, grid of 1/20; e.g. (0.7,0.6,0.9) -> 0.2, (0.5,0.5,0.5) -> 0)\n" + ] + } + ], + "source": [ + "def frechet(ps):\n", + " return max(0.0, sum(ps) - (len(ps) - 1)), min(ps)\n", + "viol = 0; tight_lo = tight_hi = 0\n", + "for trial in range(20000):\n", + " m = random.choice([2, 3])\n", + " w = np.random.dirichlet(np.ones(2 ** m) * random.choice([0.2, 1.0, 5.0]))\n", + " outcomes = list(itertools.product([0, 1], repeat=m))\n", + " ps = [sum(w[i] for i, o in enumerate(outcomes) if o[j]) for j in range(m)]\n", + " joint = sum(w[i] for i, o in enumerate(outcomes) if all(o))\n", + " lo, hi = frechet(ps)\n", + " viol += not (lo - 1e-12 <= joint <= hi + 1e-12)\n", + "assert viol == 0\n", + "print('20000 random joint distributions (m = 2 or 3 stages): Frechet bound violations =', viol)\n", + "# tightness: comonotone coupling attains the upper bound, counter-monotone (m=2) the lower\n", + "ps = (0.726, 0.523, 0.049)\n", + "lo_, hi_ = frechet(ps)\n", + "print(f'VERINA stage marginals (code 72.6%, spec 52.3%, proof 4.9%) => P(all three) in [{lo_:.3f}, {hi_:.3f}]')\n", + "assert frechet(ps) == (0.0, 0.049)\n", + "# attained: draw U~U(0,1); stage j succeeds iff U < p_j => joint = min p\n", + "U = np.random.rand(2_000_000); joint_co = np.mean((U < ps[0]) & (U < ps[1]) & (U < ps[2]))\n", + "print(f'comonotone coupling attains the upper bound: {joint_co:.4f} vs min = {min(ps)}')\n", + "assert abs(joint_co - min(ps)) < 2e-3\n", + "print('independent stages would give the product', round(math.prod(ps), 5), '(inside the bounds; real tasks are positively dependent)')\n", + "\n", + "# the LOWER bound is attained too: put the complements A_j^c on consecutive intervals of the circle [0,1) (disjoint while sum(1-p) <= 1, covering otherwise)\n", + "from fractions import Fraction as Fr\n", + "def lower_attained(ps, cells=20):\n", + " ps = [Fr(p).limit_denominator(cells) for p in ps]\n", + " comp = [int((1 - p) * cells) for p in ps] # |A_j^c| in cells of the grid\n", + " pos = 0; ok = [True] * cells\n", + " for c in comp:\n", + " for t in range(c): ok[(pos + t) % cells] = False # cells where stage j FAILS\n", + " pos += c\n", + " joint = Fr(sum(ok), cells)\n", + " return joint, max(Fr(0), sum(ps) - (len(ps) - 1))\n", + "for ps_ in [(0.7, 0.6, 0.9), (0.5, 0.5, 0.5), (0.95, 0.9), (0.3, 0.4), (0.8, 0.8, 0.8, 0.8), (0.75, 0.5, 0.05)]:\n", + " j_, b_ = lower_attained(ps_)\n", + " assert j_ == b_, (ps_, j_, b_)\n", + "print('lower Frechet bound attained by an explicit coupling on 6 marginals (exact rational arithmetic, grid of 1/20; e.g. (0.7,0.6,0.9) -> 0.2, (0.5,0.5,0.5) -> 0)')\n" + ] + }, + { + "cell_type": "markdown", + "id": "b98fb644", + "metadata": {}, + "source": [ + "## 4. Consistency checks of quoted numbers (and one derived baseline)" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "291ff898", + "metadata": { + "execution": { + "iopub.execute_input": "2026-10-14T09:53:53.005521Z", + "iopub.status.busy": "2026-10-14T09:53:53.005355Z", + "iopub.status.idle": "2026-10-14T09:53:53.013532Z", + "shell.execute_reply": "2026-10-14T09:53:53.012678Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "OK DafnyBench sources sum to 782 (556 GitHub + 62 Clover + 164 dafny-synthesis)\n", + "OK DafnyBench section 3.1 no-hint count vs Table 3 \"No LLM\" baseline (113+23+72 = 208 of 782 = 26.6%; Table 3 states the no-op baseline itself: 26.9% = 210 programs, 2 more than the section-3.1 count)\n", + "OK DafnyBench Table 3 +- for GPT-3.5 Turbo (binomial s.e. 1.78 vs reported 1.8)\n", + "OK DafnyBench Table 3 +- for GPT-4 Turbo (binomial s.e. 1.75 vs reported 1.8)\n", + "OK DafnyBench Table 3 +- for Claude 3 Opus (binomial s.e. 1.67 vs reported 1.7)\n", + "OK DafnyBench Table 3 +- for CodeLlama-7b (binomial s.e. 1.61 vs reported 1.6)\n", + "OK miniCodeProps GPT-4o medley (65/86 = 75.58% vs reported 75.6%)\n", + "OK miniCodeProps GPT-4o med+hard (5/115 = 4.35% vs reported 4.34%)\n", + "OK miniCodeProps GPT-4o overall (70/201 = 34.83% vs reported 34.8%)\n", + "OK miniCodeProps refine medley (67/86 = 77.91% vs reported 77.9%)\n", + "OK miniCodeProps refine med+hard (8/115 = 6.96% vs reported 6.96%)\n", + "OK miniCodeProps refine overall (75/201 = 37.31% vs reported 37.3%)\n", + "OK miniCodeProps ntp medley (62/86 = 72.09% vs reported 72.1%)\n", + "OK miniCodeProps ntp med+hard (10/115 = 8.70% vs reported 8.69%)\n", + "OK miniCodeProps ntp overall (72/201 = 35.82% vs reported 35.8%)\n", + "OK miniCodeProps 86 + 115 = 201 \n", + "OK miniCodeProps overall = sum of splits (GPT-4o) \n", + "OK FVAPPS Sonnet 121/406 ~ 30% (29.80%)\n", + "FLAG FVAPPS Gemini 74/406 vs paper text \"18.5%\" (18.23%; abstract says 18%) -> small internal inconsistency in the paper\n", + "OK FVAPPS tiers nested: 1083 <= 2089 <= 4715 (23.0% pass all five stages; 44.3% guarded)\n", + "OK CLEVER 161 of 164 HumanEval \n", + "OK VERINA ground-truth proofs 46 of 189 (24.3% of tasks have a reference proof)\n", + "\n", + "22 checks, 1 flagged (the FVAPPS 18.5% vs 74/406 = 18.2% discrepancy is the only one expected)\n" + ] + } + ], + "source": [ + "checks = []\n", + "def chk(name, ok, detail=''):\n", + " checks.append((name, ok, detail)); print(f'{\"OK \" if ok else \"FLAG\"} {name} {detail}')\n", + "# DafnyBench (arXiv 2406.08467 section 3.1): sources of the 782 ground-truth programs\n", + "chk('DafnyBench sources sum to 782', 556 + 62 + 164 == 782, '(556 GitHub + 62 Clover + 164 dafny-synthesis)')\n", + "no_hint = 113 + 23 + 72\n", + "chk('DafnyBench section 3.1 no-hint count vs Table 3 \"No LLM\" baseline', abs(100 * no_hint / 782 - 26.9) < 0.5,\n", + " f'(113+23+72 = {no_hint} of 782 = {100*no_hint/782:.1f}%; Table 3 states the no-op baseline itself: 26.9% = {round(0.269*782)} programs, {round(0.269*782)-no_hint} more than the section-3.1 count)')\n", + "# Table 3 reports '+-' = one binomial standard error at N = 782\n", + "for name, p_, se_rep in [('GPT-3.5 Turbo', 0.440, 1.8), ('GPT-4 Turbo', 0.598, 1.8), ('Claude 3 Opus', 0.678, 1.7), ('CodeLlama-7b', 0.280, 1.6)]:\n", + " chk(f'DafnyBench Table 3 +- for {name}', abs(100 * math.sqrt(p_ * (1 - p_) / 782) - se_rep) < 0.06, f'(binomial s.e. {100*math.sqrt(p_*(1-p_)/782):.2f} vs reported {se_rep})')\n", + "# miniCodeProps (arXiv 2406.11915, Table 3)\n", + "for name, k, n, rep in [('GPT-4o medley', 65, 86, 75.6), ('GPT-4o med+hard', 5, 115, 4.34), ('GPT-4o overall', 70, 201, 34.8),\n", + " ('refine medley', 67, 86, 77.9), ('refine med+hard', 8, 115, 6.96), ('refine overall', 75, 201, 37.3),\n", + " ('ntp medley', 62, 86, 72.1), ('ntp med+hard', 10, 115, 8.69), ('ntp overall', 72, 201, 35.8)]:\n", + " chk(f'miniCodeProps {name}', abs(100 * k / n - rep) < 0.06, f'({k}/{n} = {100*k/n:.2f}% vs reported {rep}%)')\n", + "chk('miniCodeProps 86 + 115 = 201', 86 + 115 == 201)\n", + "chk('miniCodeProps overall = sum of splits (GPT-4o)', 65 + 5 == 70 and 67 + 8 == 75 and 62 + 10 == 72)\n", + "# FVAPPS (arXiv 2502.05714 section III)\n", + "chk('FVAPPS Sonnet 121/406 ~ 30%', abs(100 * 121 / 406 - 30) < 0.5, f'({100*121/406:.2f}%)')\n", + "chk('FVAPPS Gemini 74/406 vs paper text \"18.5%\"', abs(100 * 74 / 406 - 18.5) < 0.1, f'({100*74/406:.2f}%; abstract says 18%) -> small internal inconsistency in the paper')\n", + "chk('FVAPPS tiers nested: 1083 <= 2089 <= 4715', 1083 <= 2089 <= 4715, f'({100*1083/4715:.1f}% pass all five stages; {100*2089/4715:.1f}% guarded)')\n", + "# CLEVER (arXiv 2505.13938)\n", + "chk('CLEVER 161 of 164 HumanEval', 161 <= 164)\n", + "# VERINA (arXiv 2505.23135): test-suite medians and the relation count\n", + "chk('VERINA ground-truth proofs 46 of 189', abs(100 * 46 / 189 - 24.3) < 0.1, f'({100*46/189:.1f}% of tasks have a reference proof)')\n", + "n_fail = sum(not ok for _, ok, _ in checks)\n", + "print(f'\\n{len(checks)} checks, {n_fail} flagged (the FVAPPS 18.5% vs 74/406 = 18.2% discrepancy is the only one expected)')\n", + "assert n_fail == 1" + ] + }, + { + "cell_type": "markdown", + "id": "8e17d222", + "metadata": {}, + "source": [ + "## What the table in the note is built from\n", + "The benchmark comparison table (task / who writes the spec / spec validation / checker / headline) is hand-assembled from each paper's own text in the\n", + "note; the numbers above are the ones that can be recomputed. Reading rule: **a headline number needs (a) the task definition, (b) a trivial baseline, (c) an interval.**" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/code/notebooks/0307-aeneas-semantics.ipynb b/code/notebooks/0307-aeneas-semantics.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..225f97b9a52ae1f4bfa5192ad2d8e1917bb223b0 --- /dev/null +++ b/code/notebooks/0307-aeneas-semantics.ipynb @@ -0,0 +1,438 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# 0307 - Aeneas semantics: Rust (safe subset) -> LLBC -> pure Lean, the Result monad, and what a proof of `not P` certifies\n", + "Companion to `notes/0307 - Aeneas Semantics.pdf`. Pure Python, exhaustive enumeration over `u8`; no Lean needed.\n", + "The \"Rust abstract machine\" below is a small stateful interpreter (store, borrows, panics, fuel). The \"translation\" is a pure function\n", + "into `Result` with explicit threading of borrowed values (backward functions). Every check compares them on the **whole** finite domain.\n", + "This is a finite test of the refinement statement, not a proof of Aeneas' soundness theorem (that is a meta-theorem about the tool, see the note)." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "source": [ + "import os\n", + "os.environ[\"OMP_NUM_THREADS\"] = \"1\"\n", + "import itertools, math\n", + "from dataclasses import dataclass\n", + "CHECKS = []\n", + "def check(name, ok, detail=\"\"):\n", + " CHECKS.append((name, bool(ok), detail))\n", + " assert ok, (name, detail)\n", + " print(f\"[ok] {name} {detail}\")" + ], + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 1. The Result monad: ok / fail(e) / div\n", + "`Result a = ok a | fail Error | div`. `bind` propagates `fail` and `div`. Errors we use: `integerOverflow`, `arrayOutOfBounds`, `panic`.\n", + "Monad laws are checked on a finite carrier (all values of `u8`, three error tags, `div`)." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "source": [ + "@dataclass(frozen=True)\n", + "class R:\n", + " tag: str # 'ok' | 'fail' | 'div'\n", + " v: object = None # value for ok, error name for fail\n", + "def ok(v): return R('ok', v)\n", + "def fail(e): return R('fail', e)\n", + "DIV = R('div')\n", + "def bind(m, f):\n", + " return f(m.v) if m.tag == 'ok' else m\n", + "U8 = range(256)\n", + "carrier = [ok(v) for v in U8] + [fail('integerOverflow'), fail('arrayOutOfBounds'), fail('panic'), DIV]\n", + "f = lambda x: ok(x + 1) if x < 255 else fail('integerOverflow') # a fallible step\n", + "g = lambda x: ok(x // 2) if x % 3 else fail('panic')\n", + "left = all(bind(ok(v), f) == f(v) for v in U8)\n", + "right = all(bind(m, ok) == m for m in carrier)\n", + "assoc = all(bind(bind(m, f), g) == bind(m, lambda x: bind(f(x), g)) for m in carrier)\n", + "check(\"1. monad laws on 260 values: left identity, right identity, associativity\", left and right and assoc)" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "[ok] 1. monad laws on 260 values: left identity, right identity, associativity \n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 2. Naive translation fails: integers without overflow\n", + "Rust (debug / `overflow-checks`): `fn avg(a: u8, b: u8) -> u8 { (a + b) / 2 }`. The abstract machine panics iff `a + b > 255`.\n", + "NAIVE translation: `a, b` as mathematical integers, no failure case. CORRECT translation: `let s <- a + b; s / 2` with checked `+` in `Result`." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "source": [ + "def rust_avg(a, b): # abstract machine, checked arithmetic\n", + " s = a + b\n", + " if s > 255: return ('panic', 'integerOverflow')\n", + " return ('ret', s // 2)\n", + "def tr_naive_avg(a, b): return ok((a + b) // 2) # no Result\n", + "def u8_add(a, b): return ok(a + b) if a + b <= 255 else fail('integerOverflow')\n", + "def tr_avg(a, b): return bind(u8_add(a, b), lambda s: ok(s // 2))\n", + "def obs(r): return ('ret', r.v) if r.tag == 'ok' else (('panic', r.v) if r.tag == 'fail' else ('div', None))\n", + "DOM = [(a, b) for a in U8 for b in U8]\n", + "bad_naive = sum(1 for a, b in DOM if obs(tr_naive_avg(a, b)) != rust_avg(a, b))\n", + "bad_good = sum(1 for a, b in DOM if obs(tr_avg(a, b)) != rust_avg(a, b))\n", + "panics = sum(1 for a, b in DOM if rust_avg(a, b)[0] == 'panic')\n", + "check(\"2. naive translation disagrees with the machine on exactly the panicking inputs\", bad_naive == panics == 255 * 256 // 2, f\"disagree={bad_naive}, panics={panics}, 255*256/2={255*256//2}\")\n", + "check(\"2. Result translation agrees with the machine on all 65536 inputs\", bad_good == 0, f\"disagree={bad_good}\")" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "[ok] 2. naive translation disagrees with the machine on exactly the panicking inputs disagree=32640, panics=32640, 255*256/2=32640\n", + "[ok] 2. Result translation agrees with the machine on all 65536 inputs disagree=0\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 3. Borrow elimination: `&mut` becomes \"return the new value\" (forward) and, when a borrow is returned, a backward function\n", + "Machine: a store maps addresses to values; `&mut x` is an address. `fn bump(x: &mut u8) { *x += 1 }`.\n", + "Translation: `bump (x : U8) : Result U8` returns the new value and the caller rebinds." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "source": [ + "def machine_bump(store, addr): # returns (outcome, store')\n", + " v = store[addr] + 1\n", + " if v > 255: return ('panic', 'integerOverflow'), store # store unchanged when the panic happens before the write\n", + " s2 = dict(store); s2[addr] = v\n", + " return ('ret', None), s2\n", + "def tr_bump(x): return u8_add(x, 1)\n", + "good = 0\n", + "for x in U8:\n", + " out, st = machine_bump({0: x}, 0)\n", + " r = tr_bump(x)\n", + " if out[0] == 'ret': good += (r.tag == 'ok' and r.v == st[0])\n", + " else: good += (r.tag == 'fail' and r.v == out[1])\n", + "check(\"3a. bump: machine (store, outcome) = translation (returned value, Result) for all u8\", good == 256)\n", + "# Returned mutable borrow: fn get_mut(v: &mut Vec, i: usize) -> &mut u8 ; caller: *get_mut(&mut v, i) += k\n", + "def machine_incr_at(vec, i, k):\n", + " heap = {j: x for j, x in enumerate(vec)} # elements are addressable places\n", + " if i >= len(vec): return ('panic', 'arrayOutOfBounds'), vec\n", + " v = heap[i] + k\n", + " if v > 255: return ('panic', 'integerOverflow'), vec\n", + " heap[i] = v\n", + " return ('ret', None), [heap[j] for j in range(len(vec))]\n", + "def tr_get_mut(v, i): # forward value and BACKWARD function back : u8 -> Vec\n", + " if i >= len(v): return fail('arrayOutOfBounds')\n", + " back = lambda new, v=v, i=i: tuple(new if j == i else x for j, x in enumerate(v))\n", + " return ok((v[i], back))\n", + "def tr_incr_at(v, i, k):\n", + " return bind(tr_get_mut(v, i), lambda p: bind(u8_add(p[0], k), lambda new: ok(p[1](new))))\n", + "cases = ties = 0\n", + "for n in range(0, 4):\n", + " for vec in itertools.product([0, 1, 254, 255], repeat=n):\n", + " for i in range(0, 5):\n", + " for k in (0, 1, 2):\n", + " out, st = machine_incr_at(list(vec), i, k)\n", + " r = tr_incr_at(vec, i, k)\n", + " cases += 1\n", + " if out[0] == 'ret': ties += (r.tag == 'ok' and list(r.v) == st)\n", + " else: ties += (r.tag == 'fail' and r.v == out[1])\n", + "check(\"3b. returned &mut: machine vs forward value + backward function, vec len<=3, i<=4, k<=2\", ties == cases, f\"{ties}/{cases} cases\")" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "[ok] 3a. bump: machine (store, outcome) = translation (returned value, Result) for all u8 \n", + "[ok] 3b. returned &mut: machine vs forward value + backward function, vec len<=3, i<=4, k<=2 1275/1275 cases\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 4. Loops become recursion; non-termination is `div`\n", + "`fn spin(x: u8) { while x != 0 { } }`. The machine is run with a fuel bound F; `div` is what survives every F.\n", + "`fn sum_to(n: u8) -> u8 { let mut s = 0; let mut i = 0; while i < n { i += 1; s += i } s }` : the accumulator overflows once n(n+1)/2 > 255, i.e. n >= 23." + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "source": [ + "def machine_spin(x, fuel):\n", + " for _ in range(fuel):\n", + " if x == 0: return ('ret', None)\n", + " return ('fuel', None)\n", + "def tr_spin(x): return ok(None) if x == 0 else DIV\n", + "agree = all(((machine_spin(x, F)[0] == 'ret') == (tr_spin(x).tag == 'ok')) and ((machine_spin(x, F)[0] == 'fuel') == (tr_spin(x).tag == 'div')) for x in U8 for F in (1, 10, 1000))\n", + "check(\"4a. spin: 'div' corresponds to fuel exhaustion for every tested fuel\", agree)\n", + "def machine_sum_to(n):\n", + " s = i = 0\n", + " while i < n:\n", + " i += 1; s += i\n", + " if s > 255: return ('panic', 'integerOverflow')\n", + " return ('ret', s)\n", + "def tr_sum_loop(n, i=0, s=0): # the recursive function Aeneas would emit\n", + " if i < n:\n", + " i2 = i + 1\n", + " r = u8_add(s, i2)\n", + " return bind(r, lambda s2: tr_sum_loop(n, i2, s2))\n", + " return ok(s)\n", + "m = [machine_sum_to(n) for n in U8]\n", + "check(\"4b. sum_to: loop = recursive translation for all n in u8\", all(obs(tr_sum_loop(n)) == m[n] for n in U8))\n", + "first_bad = min(n for n in U8 if m[n][0] == 'panic')\n", + "check(\"4b. first panicking n is 23 (n(n+1)/2 = 276 > 255, 22*23/2 = 253 <= 255)\", first_bad == 23 and 22 * 23 // 2 == 253 and 23 * 24 // 2 == 276, f\"first_bad={first_bad}\")" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "[ok] 4a. spin: 'div' corresponds to fuel exhaustion for every tested fuel \n", + "[ok] 4b. sum_to: loop = recursive translation for all n in u8 \n", + "[ok] 4b. first panicking n is 23 (n(n+1)/2 = 276 > 255, 22*23/2 = 253 <= 255) first_bad=23\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 5. A proof of `not P`: closed proposition versus negation under binders\n", + "Property of `avg`: `P(a,b,r)`: `min(a,b) <= r <= max(a,b)`, evaluated on the translated function, with `Q(a,b) := exists r, avg a b = ok r /\\ P(a,b,r)`.\n", + "Closed statement S := forall a b, H a b -> Q a b (H trivially true here). Its negation not S <=> exists a b, H a b /\\ not Q a b.\n", + "Different, much stronger statement S' := forall a b, H a b -> not Q a b. Under binders, `not Q a b` for a free (a,b) is a claim about *every* instance." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "source": [ + "def Q(a, b):\n", + " r = tr_avg(a, b)\n", + " return r.tag == 'ok' and min(a, b) <= r.v <= max(a, b)\n", + "H = lambda a, b: True\n", + "S = all((not H(a, b)) or Q(a, b) for a, b in DOM)\n", + "notS = any(H(a, b) and not Q(a, b) for a, b in DOM)\n", + "Sprime = all((not H(a, b)) or (not Q(a, b)) for a, b in DOM)\n", + "witnesses = [(a, b) for a, b in DOM if H(a, b) and not Q(a, b)]\n", + "check(\"5a. S is false and (not S) holds, certified by one witness\", (not S) and notS and witnesses[0] == (1, 255), f\"|W|={len(witnesses)}, first witness {witnesses[0]}\")\n", + "check(\"5a. the 'under binders' strengthening S' (forall x, not Q x) is false: Q holds at (0,0)\", (not Sprime) and Q(0, 0), f\"|{{Q}}|={len(DOM)-len(witnesses)}\")\n", + "check(\"5a. |W| = number of panicking inputs = 32640\", len(witnesses) == 32640)\n", + "# With a precondition H a b := a + b <= 255 the closed statement is TRUE and its negation is not provable\n", + "H2 = lambda a, b: a + b <= 255\n", + "check(\"5b. with H := a+b<=255, S_H holds (so no certificate for not S_H exists)\", all((not H2(a, b)) or Q(a, b) for a, b in DOM))\n", + "# Vacuity: if H is unsatisfiable, the under-binders form is provable and yet there is no witness at all\n", + "H3 = lambda a, b: a + b > 255 and a + b < 256\n", + "check(\"5c. unsatisfiable H: S'_H = forall x, H x -> not Q x holds vacuously, but the witness set is empty (no bug evidence)\",\n", + " all((not H3(a, b)) or (not Q(a, b)) for a, b in DOM) and not any(H3(a, b) for a, b in DOM))\n", + "# And with H satisfiable and S'_H true (every H-input fails): S'_H implies the closed negation\n", + "H4 = lambda a, b: a + b >= 256\n", + "check(\"5c. satisfiable H := a+b>=256: S'_H holds (every such input fails) and so implies not S_H\", all((not H4(a, b)) or (not Q(a, b)) for a, b in DOM) and any(H4(a, b) for a, b in DOM))" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "[ok] 5a. S is false and (not S) holds, certified by one witness |W|=32640, first witness (1, 255)\n", + "[ok] 5a. the 'under binders' strengthening S' (forall x, not Q x) is false: Q holds at (0,0) |{Q}|=32896\n", + "[ok] 5a. |W| = number of panicking inputs = 32640 \n", + "[ok] 5b. with H := a+b<=255, S_H holds (so no certificate for not S_H exists) \n", + "[ok] 5c. unsatisfiable H: S'_H = forall x, H x -> not Q x holds vacuously, but the witness set is empty (no bug evidence) \n", + "[ok] 5c. satisfiable H := a+b>=256: S'_H holds (every such input fails) and so implies not S_H \n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 6. Reachability: the certificate is about `avg`, not about the program that calls it\n", + "Public caller: `fn mean_pct(a: u8, b: u8) -> u8 { avg(min(a, c), min(b, c)) }` for a clamp `c` (a property of the *caller*).\n", + "Reach_c = {(a', b'): a', b' <= c}. Witnesses W from section 5 are real bugs of the program only if W intersect Reach_c is non-empty." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "source": [ + "def count_formula(c): # pairs in [0,c]^2 with a+b >= 256\n", + " return sum(max(0, min(c, 255) - max(0, 256 - a) + 1) for a in range(0, min(c, 255) + 1))\n", + "rows = []\n", + "for c in (50, 100, 127, 128, 150, 200, 255):\n", + " brute = sum(1 for a, b in witnesses if a <= c and b <= c)\n", + " rows.append((c, brute, count_formula(c)))\n", + " assert brute == count_formula(c)\n", + " print(f\" clamp c={c:3d}: |W \u2229 Reach_c| = {brute:6d} (closed form {count_formula(c)})\")\n", + "check(\"6. |W \u2229 Reach_c| equals the closed form for all tested c\", all(b == f for _, b, f in rows))\n", + "check(\"6. c <= 127: certificate holds for avg yet the program mean_pct never panics (false alarm)\", dict((c, b) for c, b, _ in rows)[127] == 0 and dict((c, b) for c, b, _ in rows)[128] == 1, \"c=127 -> 0, c=128 -> 1 (128+128=256)\")\n", + "# end-to-end: run the caller on its whole input domain\n", + "def mean_pct(a, b, c): return tr_avg(min(a, c), min(b, c))\n", + "for c in (100, 127, 128, 200):\n", + " panics_c = sum(1 for a, b in DOM if mean_pct(a, b, c).tag == 'fail')\n", + " print(f\" mean_pct with clamp {c}: panicking inputs over all 65536 raw (a,b): {panics_c}\")\n", + "check(\"6. the caller with c=127 never fails on any of 65536 raw inputs; with c=128 it fails (a,b >= 128)\", all(mean_pct(a, b, 127).tag == 'ok' for a, b in DOM) and sum(1 for a, b in DOM if mean_pct(a, b, 128).tag == 'fail') == 128 * 128)" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + " clamp c= 50: |W \u2229 Reach_c| = 0 (closed form 0)\n", + " clamp c=100: |W \u2229 Reach_c| = 0 (closed form 0)\n", + " clamp c=127: |W \u2229 Reach_c| = 0 (closed form 0)\n", + " clamp c=128: |W \u2229 Reach_c| = 1 (closed form 1)\n", + " clamp c=150: |W \u2229 Reach_c| = 1035 (closed form 1035)\n", + " clamp c=200: |W \u2229 Reach_c| = 10585 (closed form 10585)\n", + " clamp c=255: |W \u2229 Reach_c| = 32640 (closed form 32640)\n", + "[ok] 6. |W \u2229 Reach_c| equals the closed form for all tested c \n", + "[ok] 6. c <= 127: certificate holds for avg yet the program mean_pct never panics (false alarm) c=127 -> 0, c=128 -> 1 (128+128=256)\n", + " mean_pct with clamp 100: panicking inputs over all 65536 raw (a,b): 0\n", + " mean_pct with clamp 127: panicking inputs over all 65536 raw (a,b): 0\n", + " mean_pct with clamp 128: panicking inputs over all 65536 raw (a,b): 16384\n", + " mean_pct with clamp 200: panicking inputs over all 65536 raw (a,b): 29560\n", + "[ok] 6. the caller with c=127 never fails on any of 65536 raw inputs; with c=128 it fails (a,b >= 128) \n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 7. Translation soundness: wrong translations give wrong certificates in both directions\n", + "* U1 (too lax): unbounded integers. Then `S` is *provable*, the real program violates it: a missed bug.\n", + "* U2 (wrapping, release-mode semantics): `Q` fails on the same inputs but as a wrong value, and panic-freedom becomes provable. \"The bug\" depends on which Rust semantics is modelled.\n", + "* U3 (too strict): checks overflow on the *widened* sum `(a as u16 + b as u16)/2`, which is fine in Rust, but translates the u16 add as a u8 add: spurious `fail`, a false certificate." + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": {}, + "source": [ + "def Q_with(trf): return lambda a, b: (lambda r: r.tag == 'ok' and min(a, b) <= r.v <= max(a, b))(trf(a, b))\n", + "trU1 = tr_naive_avg\n", + "trU2 = lambda a, b: ok(((a + b) % 256) // 2)\n", + "def rust_wide(a, b): return ('ret', (a + b) // 2) # `((a as u16 + b as u16) / 2) as u8` never panics\n", + "trU3 = lambda a, b: bind(u8_add(a, b), lambda s: ok(s // 2)) # wrong: widened add checked at u8\n", + "Wset = lambda q: sum(1 for a, b in DOM if not q(a, b))\n", + "print(\" failing sets |{not Q}|: truth(avg, checked) =\", Wset(Q), \", U1 (unbounded) =\", Wset(Q_with(trU1)), \", U2 (wrapping) =\", Wset(Q_with(trU2)))\n", + "check(\"7a. U1: S is provable under the unbounded translation (no witness), but the program panics on 32640 inputs\", Wset(Q_with(trU1)) == 0 and Wset(Q) == 32640)\n", + "pf_checked = all(tr_avg(a, b).tag == 'ok' for a, b in DOM); pf_wrap = all(trU2(a, b).tag == 'ok' for a, b in DOM)\n", + "same_set = {(a, b) for a, b in DOM if not Q_with(trU2)(a, b)} == set(witnesses)\n", + "check(\"7b. U2: wrapping gives the SAME 32640 witnesses for P but for a different reason (wrong value, not panic); panic-freedom is provable under U2 and false under checked semantics\",\n", + " same_set and pf_wrap and not pf_checked, f\"|W_wrap|={Wset(Q_with(trU2))}, panic-free: wrap={pf_wrap}, checked={pf_checked}\")\n", + "# wide version of the program: truth never panics, so S_wide holds; the strict translation certifies a bug\n", + "Qwide_true = lambda a, b: min(a, b) <= rust_wide(a, b)[1] <= max(a, b)\n", + "check(\"7c. U3: program `avg_wide` is correct on all inputs; the over-strict translation yields 32640 spurious witnesses\", all(Qwide_true(a, b) for a, b in DOM) and Wset(Q_with(trU3)) == 32640)" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + " failing sets |{not Q}|: truth(avg, checked) = 32640 , U1 (unbounded) = 0 , U2 (wrapping) = 32640\n", + "[ok] 7a. U1: S is provable under the unbounded translation (no witness), but the program panics on 32640 inputs \n", + "[ok] 7b. U2: wrapping gives the SAME 32640 witnesses for P but for a different reason (wrong value, not panic); panic-freedom is provable under U2 and false under checked semantics |W_wrap|=32640, panic-free: wrap=True, checked=False\n", + "[ok] 7c. U3: program `avg_wide` is correct on all inputs; the over-strict translation yields 32640 spurious witnesses \n" + ] + } + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Summary" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "metadata": {}, + "source": [ + "for n, ok_, d in CHECKS: print((\"PASS \" if ok_ else \"FAIL \") + n)\n", + "assert all(c[1] for c in CHECKS)\n", + "print(f\"\\n{len(CHECKS)} checks, all passed\")" + ], + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "PASS 1. monad laws on 260 values: left identity, right identity, associativity\n", + "PASS 2. naive translation disagrees with the machine on exactly the panicking inputs\n", + "PASS 2. Result translation agrees with the machine on all 65536 inputs\n", + "PASS 3a. bump: machine (store, outcome) = translation (returned value, Result) for all u8\n", + "PASS 3b. returned &mut: machine vs forward value + backward function, vec len<=3, i<=4, k<=2\n", + "PASS 4a. spin: 'div' corresponds to fuel exhaustion for every tested fuel\n", + "PASS 4b. sum_to: loop = recursive translation for all n in u8\n", + "PASS 4b. first panicking n is 23 (n(n+1)/2 = 276 > 255, 22*23/2 = 253 <= 255)\n", + "PASS 5a. S is false and (not S) holds, certified by one witness\n", + "PASS 5a. the 'under binders' strengthening S' (forall x, not Q x) is false: Q holds at (0,0)\n", + "PASS 5a. |W| = number of panicking inputs = 32640\n", + "PASS 5b. with H := a+b<=255, S_H holds (so no certificate for not S_H exists)\n", + "PASS 5c. unsatisfiable H: S'_H = forall x, H x -> not Q x holds vacuously, but the witness set is empty (no bug evidence)\n", + "PASS 5c. satisfiable H := a+b>=256: S'_H holds (every such input fails) and so implies not S_H\n", + "PASS 6. |W \u2229 Reach_c| equals the closed form for all tested c\n", + "PASS 6. c <= 127: certificate holds for avg yet the program mean_pct never panics (false alarm)\n", + "PASS 6. the caller with c=127 never fails on any of 65536 raw inputs; with c=128 it fails (a,b >= 128)\n", + "PASS 7a. U1: S is provable under the unbounded translation (no witness), but the program panics on 32640 inputs\n", + "PASS 7b. U2: wrapping gives the SAME 32640 witnesses for P but for a different reason (wrong value, not panic); panic-freedom is provable under U2 and false under checked semantics\n", + "PASS 7c. U3: program `avg_wide` is correct on all inputs; the over-strict translation yields 32640 spurious witnesses\n", + "\n", + "20 checks, all passed\n" + ] + } + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.12.11" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} \ No newline at end of file diff --git a/code/notebooks/_lean_helper.py b/code/notebooks/_lean_helper.py new file mode 100644 index 0000000000000000000000000000000000000000..bec87d96eb1e08a563c86705d5831db76de7e288 --- /dev/null +++ b/code/notebooks/_lean_helper.py @@ -0,0 +1,46 @@ +"""Tiny Lean runner for the 03xx notebooks. Goes ONLY through align.verify.lean. + +Rules (house + task): <= 2 workers (we use 1), nice 15, MemAvailable >= 45 GB checked +before the pool starts, never the heavy-job lock, never touches a GPU. +""" +from __future__ import annotations +import os, sys, re +os.environ["CUDA_VISIBLE_DEVICES"] = "" +sys.path.insert(0, "/apps/development/rlhf/align") +from align.verify.lean import ReplPool, ReplLimits, find_toolchain +from align.verify.lean.repl_pool import meminfo_available_gb + +HEADER = "import Mathlib\nimport Aesop\nset_option maxHeartbeats 400000\nopen Nat\n" + +class Lean: + def __init__(self, header: str = HEADER): + avail = meminfo_available_gb() + if avail < 45.0: + raise RuntimeError(f"MemAvailable {avail:.1f} GB < 45 GB: refusing to start Lean") + self.tc = find_toolchain() + self.pool = ReplPool.from_toolchain(self.tc, n_workers=1, limits=ReplLimits(nice=15)) + self.header = header + self._cm = self.pool.acquire() + self.w = self._cm.__enter__() + self.env = self.w.header_env(header) + + def run(self, code: str, timeout: float = 120.0, keep: bool = False): + """Run `code` on top of the header; return (messages, new_env).""" + resp = self.w.command(code, self.env, timeout) + if keep and "env" in resp: + self.env = resp["env"] + return resp + + def show(self, code: str, timeout: float = 120.0, keep: bool = False) -> dict: + r = self.run(code, timeout, keep) + for m in r.get("messages", []): + print(f"[{m['severity']}] line {m['pos']['line']}: {m['data']}") + for s in r.get("sorries", []): + print("[sorries]", s.get("goal", "")[:200]) + if not r.get("messages") and not r.get("sorries"): + print("[ok] no messages") + return r + + def close(self): + try: self._cm.__exit__(None, None, None) + finally: self.pool.close() diff --git a/code/notebooks/_m1_lib.py b/code/notebooks/_m1_lib.py new file mode 100644 index 0000000000000000000000000000000000000000..f6db224458701ca7d35d4a446a49288be0c28549 --- /dev/null +++ b/code/notebooks/_m1_lib.py @@ -0,0 +1,145 @@ +"""Shared helpers for the M1 notebooks (notes 0101-0110). + +Everything is fp64 CPU torch. The "language model" is a deliberately tiny +autoregressive policy whose *entire* sequence space can be enumerated, so every +expectation in the notes can be computed EXACTLY (no Monte Carlo) and compared +against the hand-derived formula to ~1e-15. Monte Carlo checks are kept for the +few identities that are about sampling (and are reported against 4 standard errors). + +Policy: pi_theta(a | prev, t) = softmax(theta_tok[prev] + theta_pos[t])_a + - vocabulary {0,..,V-1} where token V-1 is EOS + - `prev` is the previous token (BOS = index V for t=0) + - a sequence stops at EOS or after L tokens (truncation, no EOS) +This is a valid autoregressive factorisation pi(o) = prod_t pi(o_t | o_ tol {tol:.1e}" + +def maxabs(a, b): + return float((torch.as_tensor(a) - torch.as_tensor(b)).abs().max()) + +class TinyLM: + def __init__(self, V=4, L=4, seed=0, scale=1.0): + self.V, self.L = V, L + self.EOS = V - 1 + g = torch.Generator().manual_seed(seed) + self.n_tok = (V + 1) * V # rows: prev in {0..V-1, BOS} + self.n_pos = L * V + self.theta0 = scale * torch.randn(self.n_tok + self.n_pos, generator=g) + self.seqs, self.tokens, self.mask = self._enumerate() + self.N = len(self.seqs) + self.lengths = self.mask.sum(1) + + # ---- enumeration of the whole sequence space -------------------------------- + def _enumerate(self): + V, L, EOS = self.V, self.L, self.EOS + seqs = [] + def rec(prefix): + if len(prefix) == L or (prefix and prefix[-1] == EOS): + seqs.append(tuple(prefix)); return + for a in range(V): + rec(prefix + [a]) + rec([]) + tokens = torch.full((len(seqs), L), EOS, dtype=torch.long) + mask = torch.zeros(len(seqs), L) + for i, s in enumerate(seqs): + tokens[i, :len(s)] = torch.tensor(s) + mask[i, :len(s)] = 1.0 + return seqs, tokens, mask + + def _logits_tables(self, theta): + V, L = self.V, self.L + tok = theta[:self.n_tok].reshape(V + 1, V) + pos = theta[self.n_tok:].reshape(L, V) + return tok, pos + + def token_logprobs(self, theta): + """[N, L] log pi(o_t | o_=2.""" + R = torch.zeros(lm.N) + for i, s in enumerate(lm.seqs): + if s[-1] == lm.EOS and len(s) >= 2 and (sum(s[:-1]) % 2 == 1): + R[i] = 1.0 + return R + +def flat_grad(scalar, theta): + g, = torch.autograd.grad(scalar, theta, retain_graph=True, allow_unused=False) + return g + +def score_matrix(lm, theta): + """S[o] = grad_theta log pi(o) as an [N, P] matrix (autograd, one backward per sequence).""" + th = theta.clone().requires_grad_(True) + lp = lm.seq_logprob(th) + rows = [torch.autograd.grad(lp[i], th, retain_graph=True)[0] for i in range(lm.N)] + return torch.stack(rows) + +def token_score_tensor(lm, theta): + """T[o, t] = grad_theta log pi(o_t | o_", "exec"), ns) + except BaseException: + err = traceback.format_exc() + outputs = [] + if so.getvalue(): outputs.append({"output_type": "stream", "name": "stdout", "text": so.getvalue().splitlines(True)}) + if se.getvalue(): outputs.append({"output_type": "stream", "name": "stderr", "text": se.getvalue().splitlines(True)}) + if err: + outputs.append({"output_type": "error", "ename": "Error", "evalue": err.splitlines()[-1], "traceback": err.splitlines()}) + nb_cells.append({"cell_type": "code", "execution_count": count, "metadata": {}, "source": code.splitlines(True), "outputs": outputs}) + if err: + print(err, file=sys.stderr) + nb = {"cells": nb_cells, "metadata": {}, "nbformat": 4, "nbformat_minor": 4} + json.dump(nb, open(out + ".FAILED", "w"), indent=1) + sys.exit(1) +nb = {"cells": nb_cells, + "metadata": {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}, + "language_info": {"name": "python", "version": sys.version.split()[0]}}, + "nbformat": 4, "nbformat_minor": 4} +json.dump(nb, open(out, "w"), indent=1) +print(f"wrote {out}: {count} code cells in {time.time()-t0:.1f}s") diff --git a/code/notebooks/_m1_show.py b/code/notebooks/_m1_show.py new file mode 100644 index 0000000000000000000000000000000000000000..ef76eb5f08e34d70ff19270cfaa0871ca4bd79fd --- /dev/null +++ b/code/notebooks/_m1_show.py @@ -0,0 +1,4 @@ +import json,sys +nb=json.load(open(sys.argv[1])) +for c in nb['cells']: + for o in c.get('outputs',[]): print(''.join(o.get('text',o.get('traceback',''))),end='') diff --git a/code/notes/STYLE_ALIGNMENT.md b/code/notes/STYLE_ALIGNMENT.md new file mode 100644 index 0000000000000000000000000000000000000000..53bf7d5811c1548faafbe4a61dd5f49bf1216790 --- /dev/null +++ b/code/notes/STYLE_ALIGNMENT.md @@ -0,0 +1,86 @@ +# Style alignment: align/notes and align/notebooks vs Flash-Attention-101 + +Reference: `/apps/development/rlhf/reference-repos/Flash-Attention-101` (notes/0000-0008, 009, 010; README.md; scratch-pad.ipynb). +Provenance caveat: the reference PDFs are handwritten notes whose title page (0000, p.1) is signed "Flash Attention from first principles, Umar Jamil". The repo (README clone URL viai957/Flash-Attention-101) packages them with the user's own README, notebooks and kernels. The "house style" below is therefore the handwritten-lecture style plus the repo's layout, not a typeset template. Where typeset equivalents are needed I mapped them (see "Mapping"). + +## 1. Reference style rules (page citations are PDF page numbers) + +Layout and naming +- R1. Numbered, one-topic-per-file notes named `NNNN - Title Case.pdf` (0001 - Multi Head Attention, 0003 - Online Softmax, ...). 0000 is a Table of Contents (0000 p.2 "Topics" list, p.3 "Prerequisites"). Numbering is not perfectly uniform (009, 010 have three digits). +- R2. Length 2-9 pages per note (0002: 4, 0003: 6, 0004: 9, 0008: 9, 009/010: 2). One idea per page. +- R3. Notebook: a single `scratch-pad.ipynb` that goes naive first (cells 0-9 "Conventional Softmax": loop for max, loop for normaliser, loop for division, then one consolidated typed function), then the improved version ("Safe Softmax" cell 10 on: single fused loop), then block/grid scaffolding for the Triton kernel. Markdown cell = pseudocode, code cell = direct transcription, then a consolidated function with a docstring. No asserts; checks are by printing. Kernel code lives in `triton/` and `cuda/`; README maps folder -> role. +- R4. README: Overview, Folder Structure, "Understanding ...", "Implementation Details" in the order PyTorch naive -> Triton -> CUDA, Benchmarking, Setup. + +Page structure and tone +- R5. Question-first opening. Each note starts with a bare problem statement as a question, then the pseudocode that fails: 0003 p.1 "The problem statement is: can we find a way to compute the softmax of a vector without going through the vector 3 times, but also preventing the exponential from exploding?"; 0004 p.2 "WHY SHOULD WE CARE?"; 0002 p.1 "But there's a problem! If the values of the vector are large, the exponential will explode!"; 0008 p.1 states the object ("Derivative: scalar input, scalar output") first. +- R6. Motivation -> failing naive version -> fix -> proof. 0003: problem (p.1) -> try fusing (p.2) -> step-by-step trace on the same tiny vector x = [3, 2, 5, 1] (pp.2-4) -> "Can we prove this algorithm is correct? Let's do it!" (p.5) -> induction with base case and step, each closed by a marker (p.6). +- R7. One running toy example reused on every page (x = [3,2,5,1]; the 8x8 blocked Q,K,V of 0004), with explicit shapes written next to every object ((N,d), (d,N), "ORIGINAL (8,128) / BLOCK (4,128)", 0002 p.1, 0004 pp.2-5). +- R8. Pseudocode appears next to the derivation, is restated in full each time it is modified, and is labelled "Pseudocode" / "New pseudocode" (0003 pp.1, 4). +- R9. Figures carry the argument: grid pictures of blocked matrices, computation graph for autograd (0008 pp.6-7), memory-hierarchy figure (0000 p.1). Colour has a fixed role: blue = given/derivation, red = warning or the open question, purple = annotation/shape, green = "correct" / the question to prove. +- R10. Warnings are explicit and loud: "WRONG!" box with the reason (0004 p.5), "The l3 we computed is wrong! Here it is wrong!" (0003 p.3). A "THE IDEA" boxed summary precedes the fix (0004 p.6). +- R11. Tone: second person plural, conversational ("Let's restore it... with a twist!", "Let's multiply by V", "Lots of patience", 0000 p.3). Short sentences, no citations inside notes; the only external reference is a pasted algorithm (FlashAttention-2 Algorithm 1, 0004 p.9) annotated by hand. +- R12. Complexity is stated per step (time and memory reads, 0002 p.3). +- R13. Notation: m_i (running max), l_i (normaliser), O_i, S = QK^T, P = softmax(S), blocks Q_i, K_j, V_j; softmax* for the unnormalised block softmax (0004 p.4); upstream/local/downstream gradient vocabulary (0008 p.9). + +## 2. Mapping to the typeset notes (what "same style" means here) + +| Reference | Typeset equivalent used | +|---|---| +| R1 file naming | already `NNNN - Title.pdf`; kept | +| R5 question first | `problem` box titled "The problem" whose first line is `\question{...}` | +| R6 motivation -> math -> proof | Notation, Derivation (proofs inline), unchanged content | +| R7 shapes and a running toy | the notes already carry toy numbers in `\ver{...}` lines and notation tables; not altered | +| R10 warnings | `What breaks` box (all 24) | +| R9 "THE IDEA" boxed summary | `Boxed result` box (all 24) | +| (flow to next page) | `Next` box (all 24) | + +## 3. Gaps found (before the change) + +Compared 0107 (M1), 0203 (M2), 0304 (M3) and then grepped all 24. + +Three series had three different style files and three different skeletons: +- Header: M1 `\hdr` (no "Note" word, grey series line), M2/M3 `\notehead` ("Note NNNN -- ..."). +- Opening: M1 grey "The problem this solves" box; M2 numbered section "The problem this solves"; M3 numbered section "The problem this concept solves". None opened with an explicit question (R5), although many posed one inside the paragraph. +- Result: M1 yellow "Result"; M2 blue box; M3 blue box, section called "Boxed final results" (M3) vs "Boxed result" (M1 box with no section, M2). +- Failure modes: M1 box "What breaks" under a section of the same name; M2 box with title override; M3 plain section (0301 titled "What breaks (failure modes and the assumption each violates)"). +- Next: M1 blue box; M2/M3 bold inline line under "Link to the next note" (M2) or none (M3). +- References: M2 had `\section{References}` followed by `thebibliography` (duplicate heading). +- Page: M1 1in margins, 11pt, colour links; M2 2.2cm; M3 2.2cm with Unicode/Lean listings; M2 10pt, M1/M3 11pt. +- Macro clashes between the files: `\ind` (bracketed one-argument in M1, bare symbol in M2), `\Pr` (operator vs blackboard P), `\nextnote` (environment vs command). + +Not present in any of the 24 (and not added, because it would be new content): hand-drawn style figures (R9), a "WRONG!" worked-trace page (R10), complexity-per-step tables (R12), a Table of Contents note (R1 0000), a prerequisites page. + +## 4. What I changed + +Style infrastructure +- `notes/src/common.sty` replaces `common-m1.sty`, `common-m2.sty`, `common-m3.sty` (the three are deleted; originals backed up in `/home/vignesh/.claude/jobs/b952890a/tmp/notes_style/backup/`). It is the union of the three: geometry 2.2cm, 11pt documents, `hidelinks`, Cerulean result box, BrickRed breaks box, grey problem box, blue Next box, Lean `listings` and the Unicode declarations, all shared macros. +- `\ind` is the bare symbol everywhere; the bracketed M1 form is now `\indb{...}` (renamed in 0101-0110 only; same rendering as before). `\Pr` is `\mathbb{P}` in all notes (0301, 0302, 0304, 0305 used the default operator "Pr"; only the glyph changes). +- `notes/src/build_all.sh` rebuilds all 24 (or a subset: `./build_all.sh 0107 0203`), takes the PDF file name from the existing `../NNNN - *.pdf`. `build.sh` and `build_m1.sh` still work unchanged. + +All 24 `.tex` files +- `\usepackage{common}`, `\documentclass[11pt]`. +- One header, `\notehead{NNNN}{title}{subtitle}` (M1's `\hdr` converted; subtitle text kept: series line plus companion notebook path). +- One opening: a grey "The problem" box whose first line is a new one-sentence `\question{...}` (24 questions, each only restating the question the existing paragraph already poses; no new claim, no number), followed by the original paragraph untouched. For M2/M3 the old numbered section was converted into this box. +- One tail: `Boxed result` section, `What breaks` section with the red box (M2/M3 explicit title override dropped, the default is the same text), `Next` box, `References` (M2/M3 duplicate heading removed). +- Section numbers: the box does not take a number in M1 (Notation = 1, Derivation = 2). In M2/M3 the old problem section was section 1, so I inserted `\stepcounter{section}` after the box so Notation stays 2 and Derivation 3. Reason: many text and notebook cross-references use those numbers (for example 0301 "\S3.3", 0302 notebook "section 3.4", 0304 notebook "section 3", 0203 notebook "section 3", 0102 cited as "Sec.~2.5" from 0105). Renumbering needs a coordinated edit of notes and notebooks and was not done. Consequence: M1 numbers sections from 1, M2/M3 from 2. + +Rebuild +- 24 of 24 PDFs rebuilt with latexmk, 0 failures, 0 undefined references/citations. Page counts: 0103 5->4, 0201 5->6, 0202 5->6, 0204 5->6, 0205 5->6, 0206 5->6, 0207 5->6, 0302 6->7, 0304 6->7, all others unchanged. Text diff of old vs new PDFs (word multiset) shows only header, section-number, heading, box-title and the 24 added `The question.` lines; no equation, number or verified-result line changed. +- `notes/src/build/m1/` (stale build artefacts of the old M1 script) was deleted; `notes/src/build/all/` is the new build dir. + +Notebooks +- Not changed. They already follow the reference's naive -> efficient progression with asserts and verification tables (0107 section 7 "Naive vs efficient implementation"; 0203 sections 1-3 naive exact/float vs stable form; 0304 "A1 NAIVE ... A2 EFFICIENT"). The reference notebook has no asserts and a single `scratch-pad.ipynb`; ours are per-note numbered `NNNN-slug.ipynb`, which I kept (the numbered naming matches R1 and the .tex/PDF linkage). Nothing was re-executed; no Lean notebook touched. + +## 5. Not aligned (left deliberately) + +1. Handwritten diagram/colour language (R9, R10), worked trace on one toy vector (R7), per-step complexity (R12): would be new content. A reasonable follow-up is a `\begin{warning}` "WRONG!" box and a 4-element trace per note, authored by hand per topic. +2. 0000 Table of Contents note and a prerequisites page: not requested; a generated `0000 - Table of Contents.pdf` from the 24 titles is mechanical if wanted. +3. Section numbering offset between M1 and M2/M3 (see above). +4. M3 notes keep their `\paragraph{Numerical verification (notebook ...)}` block while M1 uses inline `\ver{...}` tags and M2 mixes both; both carry the same information, and merging them touches content layout. +5. Series subtitle wording still differs ("Batch M1/M2", M3 has descriptive text). +6. Notebook per-note files vs reference single scratch-pad: kept. + +## 6. Possible content issues (not edited) + +- `notebooks/0101-autoregressive-likelihood-sft.ipynb` cell 2 cites "(note, Sec. 3.1)" for the proper-distribution lemma; in the M1 numbering the lemma sits under Section 2 (Derivation, "Step 2 (Lemma 1)"), so this reference looks stale. Not verified against the notebook text beyond the grep. +- No mathematical claim or number was found to be wrong during this pass (this was a style pass; the content was not re-derived). diff --git a/code/scripts/bench_7b_end_to_end.py b/code/scripts/bench_7b_end_to_end.py new file mode 100644 index 0000000000000000000000000000000000000000..5831018ed960cca615f7cfc14596cfd2b449d7ac --- /dev/null +++ b/code/scripts/bench_7b_end_to_end.py @@ -0,0 +1,283 @@ +"""Benchmark #2: end-to-end step time on the REAL 7B, triton_fa2 vs SDPA. + +Run (GPU 3 window; pin BEFORE torch import): + CUDA_VISIBLE_DEVICES=3 TMPDIR=/apps/development/rlhf/nano-mistral-o1/.tmp \ + uv run python scripts/bench_7b_end_to_end.py + +WHAT THIS MEASURES, AND WHAT IT DELIBERATELY DOES NOT +----------------------------------------------------- +A full-parameter 7B training step does NOT fit on one A10G: bf16 weights are +13.5 GiB and bf16 parameter gradients are another 13.5 GiB, so 27 GiB against a +23 GiB card. Real SFT gets around this with 4-way FSDP plus an offloaded +optimizer (section-07/11); that is not reproducible inside a single-GPU window, +and pretending otherwise would be measuring a different thing than the label. + +So this benchmark freezes every parameter and requires grad on `inputs_embeds` +instead. Backward still flows through every attention and MLP in all 32 layers, +computing the full activation-gradient chain -- dQ/dK/dV included, which is +exactly the part the attention backend controls. What is skipped is the +per-parameter weight-gradient GEMM, which is IDENTICAL for both backends. + +The consequence is worth stating precisely, because it decides how to read the +result: the absolute per-step delta (triton minus sdpa, in ms) is exactly what +it would be in a full step, since the omitted work is backend-independent and +additive. The delta as a PERCENTAGE is inflated, because the denominator is +missing that same backend-independent work. **This benchmark therefore reports +an UPPER BOUND on the true percentage cost of choosing triton_fa2.** A number +that bounds the answer from the safe side is worth more here than a number that +requires four GPUs and a training run to obtain. + +Both no-checkpointing and activation-checkpointing configurations are measured: +under checkpointing the attention forward runs a second time during backward, +so the backend-attributable share is larger, and SFT is expected to checkpoint. +""" + +from __future__ import annotations + +import glob +import json +import os +import sys +import time + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from align.utils.guard import apply_resource_guards + +# Light job by the box rules: no heavy lock (other GPU jobs are live), no +# dataloaders, no multi-GB writes. +apply_resource_guards(cpu_threads=8, heavy=False, label="bench-7b-attn") + +import torch +from torch.utils.checkpoint import checkpoint as ckpt_fn + +from align.model import checkpoint as ckpt +from align.model.mistral import NmoMistralForCausalLM + +REPO = "/apps/development/rlhf/nano-mistral-o1" +OUT = f"{REPO}/docs/bench_7b_end_to_end.json" +SNAPSHOT = ( + "/apps/voxtral/.cache/huggingface/hub/" + "models--mistralai--Mistral-7B-v0.3/snapshots/" + "caa1feb0e54d415e2df31207e5f4e273e33509b1" +) + +SEQ = 1024 +BATCH = 1 +#: Hard wall-clock budget. The GPU window is shared and SFT launches the moment +#: the golden gate passes; a benchmark that overstays blocks a training run. +DEADLINE_S = float(os.environ.get("NMO_BENCH_DEADLINE_S", 1500)) +_START = time.time() + + +def _left() -> float: + return DEADLINE_S - (time.time() - _START) + + +def _load_7b(attn_impl: str, state_dict, cfg): + model = NmoMistralForCausalLM.from_state_dict(cfg, state_dict, attn_implementation=attn_impl) + model = model.to(torch.bfloat16).to("cuda").eval() + # Freeze: no 13.5 GiB parameter-gradient buffer. See module docstring. + for p in model.parameters(): + p.requires_grad_(False) + return model + + +def _blocks(model): + inner = model.model if hasattr(model, "model") else model + for name in ("layers", "blocks", "h"): + if hasattr(inner, name): + return getattr(inner, name) + raise AttributeError("cannot find the decoder block list on the model") + + +def _timed(fn, iters, warmup): + for _ in range(warmup): + fn() + torch.cuda.synchronize() + t0 = time.perf_counter() + for _ in range(iters): + fn() + torch.cuda.synchronize() + return (time.perf_counter() - t0) / iters * 1e3 + + +def _run_config(model, embeds, labels, use_ckpt: bool, iters: int): + """fwd-only and fwd+bwd timings for one model/checkpointing combination.""" + blocks = _blocks(model) + originals = [b.forward for b in blocks] + if use_ckpt: + for b in blocks: + fwd = b.forward + b.forward = lambda *a, _f=fwd, **kw: ckpt_fn(_f, *a, use_reentrant=False, **kw) + + try: + + def fwd_only(): + with torch.no_grad(): + model(inputs_embeds=embeds.detach()) + + def fwd_bwd(): + if embeds.grad is not None: + embeds.grad = None + logits = model(inputs_embeds=embeds) + loss = torch.nn.functional.cross_entropy( + logits.float().view(-1, logits.shape[-1]), labels.view(-1) + ) + loss.backward() + + torch.cuda.reset_peak_memory_stats() + f = _timed(fwd_only, iters, 2) + fb = _timed(fwd_bwd, iters, 2) + peak = torch.cuda.max_memory_allocated() / 2**30 + finally: + # `b.forward = original_bound_method` would leave the instance holding a + # bound method that holds the instance -- a reference cycle the caching + # allocator cannot reclaim on `empty_cache()`, which OOM'd the second + # model load. Deleting the instance attribute restores the class method + # and leaves no cycle. + for b in blocks: + if "forward" in b.__dict__: + del b.__dict__["forward"] + del originals + return {"fwd_ms": f, "fwd_bwd_ms": fb, "peak_gib": peak} + + +def main() -> int: + if torch.cuda.device_count() != 1: + print( + f"{torch.cuda.device_count()} devices visible; pin exactly one " + "(shared box, three other GPU jobs live)", + file=sys.stderr, + ) + return 1 + free_gib = torch.cuda.mem_get_info()[0] / 2**30 + print(f"device: {torch.cuda.get_device_name(0)} free {free_gib:.1f} GiB") + if free_gib < 20.0: + print( + "refusing: need ~20 GiB free for 13.5 GiB of bf16 weights plus " + "activations; another job is resident on this device.", + file=sys.stderr, + ) + return 1 + + from align.eval.golden import resolve_model_config + + cfg = resolve_model_config(SNAPSHOT) + print( + f"config: {cfg.n_layers}L hidden={cfg.hidden_size} " + f"{cfg.n_heads}q:{cfg.n_kv_heads}kv head_dim={cfg.head_dim} " + f"vocab={cfg.vocab_size} sliding_window={cfg.sliding_window}" + ) + + import safetensors.torch as st + + t0 = time.time() + hf_sd = {} + for shard in sorted(glob.glob(os.path.join(SNAPSHOT, "*.safetensors"))): + hf_sd.update(st.load_file(shard)) + nmo_sd = ckpt.hf_to_nmo(hf_sd, expected_n_layers=cfg.n_layers) + del hf_sd + print(f"loaded + remapped {len(nmo_sd)} tensors in {time.time() - t0:.1f}s") + + torch.manual_seed(0) + labels = torch.randint(0, cfg.vocab_size, (BATCH, SEQ), device="cuda") + results = {} + + # One backend per PROCESS when an arg is given. 13.5 GiB of weights leaves + # no room for a second model, and process exit is the only VRAM release + # that cannot be defeated by a stray reference. Partials are merged below. + only = sys.argv[1] if len(sys.argv) > 1 else None + impls = (only,) if only else ("sdpa", "triton_fa2") + + for impl in impls: + if _left() < 240: + print(f"\ndeadline guard: {_left():.0f}s left, skipping {impl}", file=sys.stderr) + break + print(f"\n=== {impl} ===", flush=True) + model = _load_7b(impl, nmo_sd, cfg) + embeds = torch.randn( + BATCH, SEQ, cfg.hidden_size, device="cuda", dtype=torch.bfloat16, requires_grad=True + ) + results[impl] = {} + for use_ckpt in (False, True): + if _left() < 120: + print( + f"deadline guard: {_left():.0f}s left, skipping {impl}/ckpt={use_ckpt}", + file=sys.stderr, + ) + break + r = _run_config(model, embeds, labels, use_ckpt, iters=5) + key = "activation_ckpt" if use_ckpt else "no_ckpt" + results[impl][key] = r + print( + f" {key:16s} fwd {r['fwd_ms']:8.1f} ms " + f"fwd+bwd {r['fwd_bwd_ms']:8.1f} ms peak {r['peak_gib']:5.2f} GiB", + flush=True, + ) + del model, embeds + import gc + + gc.collect() + torch.cuda.empty_cache() + + # Merge any partial written by a sibling process, so running one backend + # per process still produces a single complete receipt. + partial = f"{REPO}/.tmp/bench7b_partial.json" + if only: + prev = {} + if os.path.isfile(partial): + with open(partial) as f: + prev = json.load(f) + prev.update(results) + with open(partial, "w") as f: + json.dump(prev, f, indent=2) + results = prev + print(f"\npartial results for {sorted(results)} -> {partial}") + + payload = { + "note": "parameters FROZEN; grad flows via inputs_embeds. Absolute ms " + "deltas are exact; percentages are an UPPER BOUND on the true " + "cost, because the omitted weight-grad GEMMs are " + "backend-independent and would enlarge the denominator only. " + "See the module docstring.", + "device": torch.cuda.get_device_name(0), + "torch": torch.__version__, + "model": "mistralai/Mistral-7B-v0.3 (real weights)", + "batch": BATCH, + "seq": SEQ, + "dtype": "bfloat16", + "results": results, + } + + if "sdpa" in results and "triton_fa2" in results: + payload["comparison"] = {} + print("\n=== triton_fa2 vs sdpa on the real 7B ===") + for key in ("no_ckpt", "activation_ckpt"): + if key in results["sdpa"] and key in results["triton_fa2"]: + s = results["sdpa"][key]["fwd_bwd_ms"] + t = results["triton_fa2"][key]["fwd_bwd_ms"] + c = { + "sdpa_ms": s, + "triton_ms": t, + "delta_ms": t - s, + "ratio": t / s, + "upper_bound_pct": 100.0 * (t - s) / s, + } + payload["comparison"][key] = c + print( + f" {key:16s} sdpa {s:8.1f} ms -> triton {t:8.1f} ms " + f"({c['delta_ms']:+.1f} ms, {c['upper_bound_pct']:+.1f}% " + "UPPER BOUND)" + ) + + os.makedirs(os.path.dirname(OUT), exist_ok=True) + with open(OUT, "w") as f: + json.dump(payload, f, indent=2) + f.write("\n") + print(f"\nwrote {OUT} ({_left():.0f}s of window left)") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/code/scripts/bench_flash_attention.py b/code/scripts/bench_flash_attention.py new file mode 100644 index 0000000000000000000000000000000000000000..c9d1c60713164f530f19f94ff39b6f69c51202db --- /dev/null +++ b/code/scripts/bench_flash_attention.py @@ -0,0 +1,234 @@ +"""Kernel-level benchmark: triton_fa2 vs torch SDPA at the real Mistral shapes. + +Run (pin ONE free device; this is a light job and must NOT take the heavy lock): + CUDA_VISIBLE_DEVICES= uv run python scripts/bench_flash_attention.py + +Shapes are Mistral-7B-v0.3's: b=4, 32 query heads : 8 KV heads, head_dim 128, +causal, bf16, seq {512, 1024, 2048}. The forward column is compared against the +baseline pinned in the plan artifact `tmp/sdpa_baseline.json`, measured on this +same box: 0.189 / 0.624 / 2.282 ms. + +Purpose is to bound a decision, not to win: research (FA2 paper 2307.08691, +flash-attention#1556) says Triton loses to CUTLASS at head_dim 128 on Ampere, +~1.3-1.5x forward and ~2x backward, and head_dim 128 is exactly Mistral's. At +seq 1024 attention is a few percent of layer FLOPs, so even a 2x backward loss +is a small single-digit share of step time. This script turns that into a +measured number so the SDPA-default decision is evidence rather than vibes. +""" + +from __future__ import annotations + +import json +import os +import sys + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from align.utils.guard import apply_resource_guards + +# Light job: no heavy lock, no 7B model. Still honors the path policy so no +# cache lands off /apps. +apply_resource_guards(cpu_threads=8, heavy=False, label="bench-flash-attention") + +import torch +from torch.nn.attention import SDPBackend, sdpa_kernel + +from align.kernels.flash_attention import flash_attention + +REPO = "/apps/development/rlhf/align" +OUT = f"{REPO}/docs/bench_flash_attention.json" + +#: Measured on this box during planning; see the module docstring. +PINNED_SDPA_FWD_MS = {512: 0.189, 1024: 0.624, 2048: 2.282} + +B, HQ, HKV, D = 4, 32, 8, 128 +SEQS = (512, 1024, 2048) + + +def _bench(fn, iters=30, warmup=10): + for _ in range(warmup): + fn() + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iters): + fn() + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) / iters + + +def bench_layer_share(seq: int = 1024, batch: int = 2) -> dict: + """Put attention back in context: one full Mistral decoder layer, fwd+bwd. + + Benchmark 1 gives the ratio ON the attention kernel; this converts it into + the number the SDPA-vs-triton decision actually turns on -- share of a + training step. A layer is built from Mistral-7B-v0.3's real dimensions + (hidden 4096, intermediate 14336, 32:8 heads, head_dim 128) so the MLP and + projection GEMMs are correctly sized against the attention core. + + This is NOT the section-11/12 end-to-end number the section also asks for; + that one needs the real 7B weights and a GPU window, and is deferred to it. + A single layer times 32 is a tight bound on the attention share because the + per-layer work is identical across layers. + """ + import torch.nn as nn + + H, I = 4096, 14336 + dev_t, dt = "cuda", torch.bfloat16 + lin = lambda a, b: nn.Linear(a, b, bias=False, device=dev_t, dtype=dt) + wq, wk, wv, wo = lin(H, HQ * D), lin(H, HKV * D), lin(H, HKV * D), lin(HQ * D, H) + wg, wu, wd = lin(H, I), lin(H, I), lin(I, H) + x = torch.randn(batch * seq, H, device=dev_t, dtype=dt, requires_grad=True) + + def layer(attn_fn): + def run(): + x.grad = None + q = wq(x).view(batch, seq, HQ, D).transpose(1, 2) + k = wk(x).view(batch, seq, HKV, D).transpose(1, 2) + v = wv(x).view(batch, seq, HKV, D).transpose(1, 2) + a = attn_fn(q, k, v) + h = x + wo(a.transpose(1, 2).reshape(batch * seq, H)) + out = h + wd(torch.nn.functional.silu(wg(h)) * wu(h)) + out.sum().backward() + + return run + + def sdpa_attn(q, k, v): + with sdpa_kernel(SDPBackend.FLASH_ATTENTION): + return torch.nn.functional.scaled_dot_product_attention( + q, k, v, is_causal=True, enable_gqa=True + ) + + t_sdpa = _bench(layer(sdpa_attn), iters=20, warmup=8) + t_tri = _bench(layer(lambda q, k, v: flash_attention(q, k, v, causal=True)), iters=20, warmup=8) + + out = { + "seq": seq, + "batch": batch, + "hidden": H, + "intermediate": I, + "n_layers": 32, + "layer_sdpa_ms": t_sdpa, + "layer_triton_ms": t_tri, + "layer_ratio": t_tri / t_sdpa, + "step_penalty_pct": 100.0 * (t_tri - t_sdpa) / t_sdpa, + "est_32layer_sdpa_s": t_sdpa * 32 / 1000.0, + "est_32layer_triton_s": t_tri * 32 / 1000.0, + } + print(f"\nlayer-level (Mistral decoder layer, b={batch} seq={seq}, fwd+bwd):") + print(f" SDPA {t_sdpa:7.2f} ms/layer -> {out['est_32layer_sdpa_s']:.2f} s / 32 layers") + print(f" triton {t_tri:7.2f} ms/layer -> {out['est_32layer_triton_s']:.2f} s / 32 layers") + print(f" cost of choosing triton_fa2: {out['step_penalty_pct']:+.1f}% of the 32-layer body") + return out + + +def main() -> int: + if not torch.cuda.is_available(): + print("no CUDA device visible", file=sys.stderr) + return 1 + if torch.cuda.device_count() != 1: + print( + f"{torch.cuda.device_count()} devices visible; pin exactly one " + "with CUDA_VISIBLE_DEVICES (shared box)", + file=sys.stderr, + ) + return 1 + + dev = torch.cuda.get_device_name(0) + cc = ".".join(str(x) for x in torch.cuda.get_device_capability(0)) + print(f"device: {dev} (sm{cc.replace('.', '')}) torch {torch.__version__}") + print(f"shapes: b={B}, {HQ}q:{HKV}kv heads, head_dim={D}, causal, bf16\n") + + results = {} + for seq in SEQS: + q = torch.randn(B, HQ, seq, D, device="cuda", dtype=torch.bfloat16) + k = torch.randn(B, HKV, seq, D, device="cuda", dtype=torch.bfloat16) + v = torch.randn(B, HKV, seq, D, device="cuda", dtype=torch.bfloat16) + qg, kg, vg = (t.clone().requires_grad_() for t in (q, k, v)) + dout = torch.randn_like(qg) + + def sdpa_fwd(): + with sdpa_kernel(SDPBackend.FLASH_ATTENTION): + return torch.nn.functional.scaled_dot_product_attention( + q, k, v, is_causal=True, enable_gqa=True + ) + + def tri_fwd(): + return flash_attention(q, k, v, causal=True) + + def sdpa_fwd_bwd(): + for t in (qg, kg, vg): + t.grad = None + with sdpa_kernel(SDPBackend.FLASH_ATTENTION): + o = torch.nn.functional.scaled_dot_product_attention( + qg, kg, vg, is_causal=True, enable_gqa=True + ) + o.backward(dout) + + def tri_fwd_bwd(): + for t in (qg, kg, vg): + t.grad = None + flash_attention(qg, kg, vg, causal=True).backward(dout) + + row = { + "sdpa_fwd_ms": _bench(sdpa_fwd), + "triton_fwd_ms": _bench(tri_fwd), + "sdpa_fwd_bwd_ms": _bench(sdpa_fwd_bwd), + "triton_fwd_bwd_ms": _bench(tri_fwd_bwd), + "pinned_sdpa_fwd_ms": PINNED_SDPA_FWD_MS[seq], + } + row["fwd_ratio_triton_over_sdpa"] = row["triton_fwd_ms"] / row["sdpa_fwd_ms"] + row["fwd_bwd_ratio_triton_over_sdpa"] = row["triton_fwd_bwd_ms"] / row["sdpa_fwd_bwd_ms"] + row["sdpa_fwd_drift_vs_pinned"] = row["sdpa_fwd_ms"] / row["pinned_sdpa_fwd_ms"] + results[str(seq)] = row + # Rebind to None rather than `del`: the bench closures above hold these + # names as free variables, and `del` empties the closure cells (ruff + # F821; a later call would NameError). Rebinding frees the tensors just + # the same, so empty_cache() can actually reclaim. + q = k = v = qg = kg = vg = dout = None + torch.cuda.empty_cache() + + hdr = ( + f"{'seq':>6} | {'SDPA fwd':>9} {'pinned':>8} | {'triton fwd':>11} " + f"{'ratio':>7} | {'SDPA f+b':>9} {'triton f+b':>11} {'ratio':>7}" + ) + print(hdr) + print("-" * len(hdr)) + for seq in SEQS: + r = results[str(seq)] + print( + f"{seq:>6} | {r['sdpa_fwd_ms']:>9.3f} {r['pinned_sdpa_fwd_ms']:>8.3f} | " + f"{r['triton_fwd_ms']:>11.3f} {r['fwd_ratio_triton_over_sdpa']:>6.2f}x | " + f"{r['sdpa_fwd_bwd_ms']:>9.3f} {r['triton_fwd_bwd_ms']:>11.3f} " + f"{r['fwd_bwd_ratio_triton_over_sdpa']:>6.2f}x" + ) + + layer = bench_layer_share() + + payload = { + "device": dev, + "compute_capability": cc, + "torch": torch.__version__, + "shape": { + "batch": B, + "q_heads": HQ, + "kv_heads": HKV, + "head_dim": D, + "causal": True, + "dtype": "bfloat16", + }, + "results": results, + "layer_level": layer, + } + os.makedirs(os.path.dirname(OUT), exist_ok=True) + with open(OUT, "w") as f: + json.dump(payload, f, indent=2) + f.write("\n") + print(f"\nwrote {OUT}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/code/scripts/build_data.py b/code/scripts/build_data.py new file mode 100644 index 0000000000000000000000000000000000000000..c3418fe42b50173d67931556e04a0ed23bf8710d --- /dev/null +++ b/code/scripts/build_data.py @@ -0,0 +1,125 @@ +#!/usr/bin/env python +"""Build every corpus this project trains on, in the only legal order. + + python scripts/build_data.py --out [--tokenizer ] [--curation-target N] + +Stages: eval sets -> decontamination hash set -> SFT -> Arm A -> Arm B. +Each refuses to run without its upstream receipt; the final manifest +(`/manifest.json`) pins every content hash, every measured count against +its pinned value, reject and contamination counts, the length-bias baseline, +the delimiter ids, the tokenizer hash, and the dataset revisions. Sections +11/12 quote that manifest in their launch banner and refuse to start without +it. + +Host safety (this box, 186GB, NO swap, shared): `apply_resource_guards( +heavy=True)` runs the preflight and takes the exclusive job flock BEFORE +anything heavy, so this build cannot run concurrently with a training job -- +that exact overlap killed the predecessor's first SFT run. Every corpus is +streamed; nothing is held resident. +""" + +from __future__ import annotations + +import argparse +import json +import logging +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from align.utils.guard import apply_resource_guards # noqa: E402 + +logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") +logger = logging.getLogger("build_data") + +DEFAULT_OUT = "/apps/development/rlhf/nano-mistral-o1/data" +TOKENIZER_SNAPSHOT = ( + "/apps/voxtral/.cache/huggingface/hub/models--mistralai--Mistral-7B-v0.3/" + "snapshots/caa1feb0e54d415e2df31207e5f4e273e33509b1" +) +#: Byte-identical tokenizer files vendored in-repo (assets/tokenizer/.../PROVENANCE.md). +#: The HF-cache snapshot is NOT on every box; the build must not die on a missing +#: cache when the same four files are in the repo. +VENDORED_TOKENIZER = str( + Path(__file__).resolve().parent.parent / "assets" / "tokenizer" / "mistral-7b-v0.3" +) + + +def default_tokenizer_dir() -> str: + """The HF snapshot if it holds a tokenizer.json, else the vendored copy.""" + for c in (TOKENIZER_SNAPSHOT, VENDORED_TOKENIZER): + if (Path(c) / "tokenizer.json").is_file(): + return c + return TOKENIZER_SNAPSHOT # fails loudly in load_tokenizer, naming the path + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument( + "--out", default=DEFAULT_OUT, help="build root; every path under it is asserted onto /apps" + ) + ap.add_argument("--tokenizer", default=default_tokenizer_dir()) + ap.add_argument( + "--curation-target", + type=int, + default=None, + help="Arm B curated pair count (default: the pinned 12,000)", + ) + ap.add_argument( + "--max-len", + type=int, + default=2048, + help="build cutoff, deliberately above the 1024 training " + "seq_len: see PinnedYields.max_len for the measured " + "length distribution behind this choice", + ) + ap.add_argument("--held-out-frac", type=float, default=0.05) + ap.add_argument("--seed", type=int, default=0) + ap.add_argument( + "--strict-pins", + default=True, + action=argparse.BooleanOptionalAction, + help="assert every measured count against its pinned " + "value. --no-strict-pins is RE-MEASUREMENT MODE: it " + "stamps pins_verified=false into the manifest and " + "sections 11/12 refuse to launch against it.", + ) + ap.add_argument( + "--strict-host-check", + default=True, + action=argparse.BooleanOptionalAction, + help="disable ONLY with a deliberate reason; the check is " + "the guard that did not exist on 2026-08-12", + ) + args = ap.parse_args() + + apply_resource_guards( + cpu_threads=8, heavy=True, label="build_data", strict_host_check=args.strict_host_check + ) + + from align.data.build import PinnedYields, build_all + from align.data.sources import DataSources + from align.tokenizer import load_tokenizer + + tokenizer = load_tokenizer(args.tokenizer) + sources = DataSources.from_local_hub() + + kwargs = { + "max_len": args.max_len, + "held_out_frac": args.held_out_frac, + "seed": args.seed, + "strict": args.strict_pins, + } + if args.curation_target is not None: + kwargs["arm_b_curated"] = args.curation_target + pins = PinnedYields(**kwargs) + + manifest = build_all(sources, args.out, tokenizer, pins=pins) + print(json.dumps(manifest, indent=1)) + logger.info("build complete: %s", Path(args.out) / "manifest.json") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/code/scripts/build_stageb_data.py b/code/scripts/build_stageb_data.py new file mode 100644 index 0000000000000000000000000000000000000000..81b0770603449e54e94c0b824e304423a4d0861f --- /dev/null +++ b/code/scripts/build_stageb_data.py @@ -0,0 +1,169 @@ +"""Build the stage-b GRPO prompt sets: keep-band filter, easy-first bands, decontamination. + +Consumes Big-Math (primary) and DeepMath (comparison arm) via the pinned +streaming readers; emits per-band prompt JSONL under data/stageb/ plus a +manifest receipt. Prompts only -- GRPO tokenizes at rollout time; there is +no pair/tokenization build here. + +Decontamination is CHECKED, not assumed: Big-Math's card claims benchmark +decontamination, but it contains GSM8K-train by construction and our +primary eval is GSM8K-test -- so every kept problem is joined (normalized +text) against GSM8K-test and MATH-500 and overlaps are EXCLUDED AND +COUNTED. An eval problem inside the RL prompt pool would corrupt the +primary contrast in the direction of the hypothesis, which is the one +direction a defect must never point. + + python scripts/build_stageb_data.py \ + --keep-band 0.3 0.95 --band-edges 0.3 0.5 --out data/stageb +""" + +import argparse +import hashlib +import json +import os +import sys +from pathlib import Path + +# The package is not installed in the venv; without this a bare `python scripts/build_stageb_data.py` (the documented +# launch line) dies with ModuleNotFoundError unless PYTHONPATH is exported (same fix as scripts/sft.py). +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from align.data.sources import DataSources, normalize_problem, read_gsm8k_test, read_math500 +from align.data.stageb_sources import ( + BIGMATH_PIN, + DEEPMATH_PIN, + StageBSources, + assign_difficulty_bands, + iter_bigmath, + iter_deepmath, +) + +STAGEB_HUB = "/apps/development/rlhf/hf_cache/hub" + + +def _eval_text_index(sources: DataSources) -> dict[str, str]: + idx: dict[str, str] = {} + for row in read_gsm8k_test(sources): + idx[normalize_problem(row["question"])] = "gsm8k_test" + for row in read_math500(sources): + idx[normalize_problem(row["problem"])] = "math500" + return idx + + +def _write_bands(name, bands, manifest, out_dir, problem_key, answer_key): + os.makedirs(out_dir, exist_ok=True) + counts = [] + for i, band in enumerate(bands): + path = os.path.join(out_dir, f"{name}_band{i}.jsonl") + with open(path + ".partial", "w") as f: + for row in band: + f.write( + json.dumps( + { + "problem": row[problem_key], + "answer": str(row[answer_key]), + "difficulty": row.get("difficulty"), + "band": i, + "source": row.get("source") or name, + } + ) + + "\n" + ) + os.replace(path + ".partial", path) + with open(path, "rb") as f: + sha = hashlib.sha256(f.read()).hexdigest()[:16] + counts.append({"band": i, "n": len(band), "file": path, "sha256_16": sha}) + manifest[name] = {"bands": counts, **manifest.get(name, {})} + + +def main() -> int: + p = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + p.add_argument( + "--keep-band", + nargs=2, + type=float, + default=[0.3, 0.95], + metavar=("LO", "HI"), + help="solve-rate keep band [LO, HI) -- Big-Math only", + ) + p.add_argument( + "--band-edges", + nargs="+", + type=float, + default=[0.3, 0.5], + help="difficulty (=1-solve_rate) edges WITHIN the keep band, easy band first", + ) + p.add_argument( + "--deepmath-edges", + nargs="+", + type=float, + default=[4.0, 6.5], + help="DeepMath native-difficulty edges", + ) + p.add_argument("--out", default="data/stageb") + args = p.parse_args() + + sources = DataSources.from_local_hub() + eval_idx = _eval_text_index(sources) + sb = StageBSources.from_local_hub(hub=STAGEB_HUB, need_openmath=False, need_bigmath=True) + manifest: dict = { + "pins": {"bigmath": BIGMATH_PIN.revision, "deepmath": DEEPMATH_PIN.revision}, + "keep_band_solve_rate": args.keep_band, + "band_edges_difficulty": args.band_edges, + "deepmath_edges": args.deepmath_edges, + } + + # ---- Big-Math: keep-band + decontamination + banding + lo, hi = args.keep_band + kept, dropped_band, contaminated = [], 0, {"gsm8k_test": 0, "math500": 0} + for row in iter_bigmath(sb): + d = row.get("difficulty") + if d is None: + continue # 4 unparseable rows, counted in the banding manifest already + sr = 1.0 - d + if not (lo <= sr < hi): + dropped_band += 1 + continue + hit = eval_idx.get(normalize_problem(row["problem"])) + if hit: + contaminated[hit] += 1 + continue + kept.append(row) + bands, band_manifest = assign_difficulty_bands(iter(kept), edges=args.band_edges) + manifest["bigmath"] = { + "kept": len(kept), + "dropped_outside_keep_band": dropped_band, + "excluded_eval_overlap": contaminated, + "banding": band_manifest.as_log_dict(), + } + _write_bands("bigmath", bands, manifest, args.out, "problem", "answer") + + # ---- DeepMath comparison arm: banding on native difficulty + decontamination + dm_kept, dm_contam = [], {"gsm8k_test": 0, "math500": 0} + for row in iter_deepmath(sb): + hit = eval_idx.get(normalize_problem(row["question"])) + if hit: + dm_contam[hit] += 1 + continue + dm_kept.append(row) + dm_bands, dm_manifest = assign_difficulty_bands(iter(dm_kept), edges=args.deepmath_edges) + manifest["deepmath"] = { + "kept": len(dm_kept), + "excluded_eval_overlap": dm_contam, + "banding": dm_manifest.as_log_dict(), + } + _write_bands("deepmath", dm_bands, manifest, args.out, "question", "final_answer") + + mpath = os.path.join(args.out, "manifest.json") + with open(mpath + ".partial", "w") as f: + json.dump(manifest, f, indent=1) + os.replace(mpath + ".partial", mpath) + print(json.dumps({k: v for k, v in manifest.items() if k in ("bigmath", "deepmath")}, indent=1)) + print(f"manifest: {mpath}") + return 0 + + +if __name__ == "__main__": + import sys + + sys.exit(main()) diff --git a/code/scripts/ceiling_test.py b/code/scripts/ceiling_test.py new file mode 100644 index 0000000000000000000000000000000000000000..53d1ed7cddb42ae86298fcc8664610974086ceee --- /dev/null +++ b/code/scripts/ceiling_test.py @@ -0,0 +1,152 @@ +"""Stage-b reward-path ceiling test: a PERFECT policy must score ~100%. + +House rule from a paid-for incident (mistral-prm-ppo, 2026-08-14): two PPO +runs measured a broken answer-extractor -- a perfect policy scored 0/100. +No GRPO rollout launches until this driver certifies the EXACT reward path +(extract_answer -> OutcomeVerifier -> CompositeReward + format instrument) +at >= --floor on every gold source, with every miss listed in the receipt +for eyeballing. + +CPU-only, no model: each (problem, gold) becomes +`build_perfect_completion(gold, eos)` -- byte-identical to what a policy +flawlessly imitating its SFT targets emits, EOS included (rollout decodes +with skip_special_tokens=False). + +Sources: gsm8k (1,319), math500 (500), deepmath (--deepmath-n sampled +evenly across shards; stage-b's actual gold distribution -- the reason this +must re-run here and not only on the stage-a sets). + + python scripts/ceiling_test.py --deepmath-n 5000 +""" + +import argparse +import json +import os +import sys +from pathlib import Path + +# Run from a checkout without an install: put the repo root on sys.path first. +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from align.data.schema import extract_answer, register_answer_suffixes +from align.data.sources import DataSources, gsm8k_gold_answer, read_gsm8k_test, read_math500 +from align.data.stageb_sources import StageBSources, iter_bigmath, iter_deepmath +from align.eval.roundtrip import build_perfect_completion +from align.eval.verifier import OutcomeVerifier, VerifierConfig +from align.rl.format_instrument import format_scores +from align.rl.rewards import CompositeReward + +EOS = "" + + +#: The corpora live in two hubs: the stage-a pins under sources.DEFAULT_HUB +#: (the voxtral cache this project inherited), the stage-b downloads under +#: the rlhf hf_cache. Resolved per-source-set, not via one env var. +STAGEB_HUB = "/apps/development/rlhf/hf_cache/hub" + + +def _gold_sources(args) -> dict[str, list[str]]: + out: dict[str, list[str]] = {} + sources = DataSources.from_local_hub() # DEFAULT_HUB: gsm8k/math500 pins + if args.gsm8k: + out["gsm8k"] = [gsm8k_gold_answer(r) for r in read_gsm8k_test(sources)] + if args.math500: + out["math500"] = [r["answer"] for r in read_math500(sources)] + if args.deepmath_n or args.bigmath_n: + sb = StageBSources.from_local_hub( + hub=args.stageb_hub, + need_deepmath=bool(args.deepmath_n), + need_openmath=False, + need_bigmath=bool(args.bigmath_n), + ) + for key, it, total, n, field in ( + ("deepmath", iter_deepmath, 103_110, args.deepmath_n, "final_answer"), + ("bigmath", iter_bigmath, 251_122, args.bigmath_n, "answer"), + ): + if not n: + continue + golds: list[str] = [] + for i, row in enumerate(it(sb)): + # even sample across the stream without materializing the corpus + if i % max(1, total // n) == 0: + golds.append(str(row[field])) + if len(golds) >= n: + break + out[key] = golds + return out + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument("--floor", type=float, default=0.98) + parser.add_argument("--gsm8k", action=argparse.BooleanOptionalAction, default=True) + parser.add_argument("--math500", action=argparse.BooleanOptionalAction, default=True) + parser.add_argument("--deepmath-n", type=int, default=5000) + parser.add_argument("--bigmath-n", type=int, default=5000) + parser.add_argument( + "--workers", type=int, default=4, help="verifier procs; keep <=4 on the shared box" + ) + parser.add_argument("--out", default="outputs/eval/ceiling_test.json") + parser.add_argument("--stageb-hub", default=STAGEB_HUB) + args = parser.parse_args() + + register_answer_suffixes(EOS) # BEFORE any extraction (golden.py's incident) + verifier = OutcomeVerifier(VerifierConfig(workers=args.workers)) + reward = CompositeReward(format_weight=0.1) # overlong off: length is not under test + + report: dict = {"floor": args.floor, "eos": EOS, "sources": {}} + overall_pass = True + for name, golds in _gold_sources(args).items(): + completions = [build_perfect_completion(g, EOS) for g in golds] + preds = [extract_answer(c) for c in completions] + extract_ok = sum(1 for p in preds if p is not None) + outcomes = verifier.verify_batch(golds, preds) + correct = [o.name == "CORRECT" for o in outcomes] + fscores = [format_scores(c)["format/score"] for c in completions] + rewards = [ + reward(accuracy=1.0 if ok else 0.0, format_score=f, response_len=0)["reward/total"] + for ok, f in zip(correct, fscores) + ] + ceiling = sum(correct) / len(golds) if golds else float("nan") + misses = [ + {"gold": g, "pred": p, "outcome": o.name} + for g, p, o, ok in zip(golds, preds, outcomes, correct) + if not ok + ] + passed = ceiling >= args.floor + overall_pass &= passed + report["sources"][name] = { + "n": len(golds), + "extract_ok_rate": extract_ok / len(golds), + "verify_correct_rate": ceiling, + "mean_format_score": sum(fscores) / len(fscores), + "mean_reward_total": sum(rewards) / len(rewards), + "passed": passed, + "misses": misses[:50], + "misses_total": len(misses), + } + print( + f"{name}: n={len(golds)} extract_ok={extract_ok / len(golds):.4f} " + f"ceiling={ceiling:.4f} misses={len(misses)} -> {'PASS' if passed else 'FAIL'}", + flush=True, + ) + + report["overall_pass"] = bool(overall_pass) + os.makedirs(os.path.dirname(args.out), exist_ok=True) + tmp = args.out + ".partial" + with open(tmp, "w") as f: + json.dump(report, f, indent=1) + os.replace(tmp, args.out) + print(f"receipt: {args.out} | overall_pass={overall_pass}") + if not overall_pass: + print( + "CEILING FAIL: eyeball every miss in the receipt before touching the " + "verifier -- a representational-equivalence gap and a broken extractor " + "look identical in the rate and completely different in the misses.", + file=sys.stderr, + ) + return 0 if overall_pass else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/code/scripts/ckpt_convert.py b/code/scripts/ckpt_convert.py new file mode 100644 index 0000000000000000000000000000000000000000..c69dabddcfd6205db7664391d0c58e255722e552 --- /dev/null +++ b/code/scripts/ckpt_convert.py @@ -0,0 +1,142 @@ +#!/usr/bin/env python +"""Sharded checkpoint tools: inspect, verify, export to HF safetensors, reshard. + + python scripts/ckpt_convert.py inspect RUN/checkpoints/micro_N + python scripts/ckpt_convert.py verify RUN/checkpoints/micro_N + python scripts/ckpt_convert.py export RUN/checkpoints/micro_N OUT_DIR + python scripts/ckpt_convert.py reshard RUN/checkpoints/micro_N NEW_DIR --world-size 2 + +CPU only, no GPU, no process group, no heavy-job lock: `export` is a streaming +byte copy (peak RSS ~ a few MiB plus the Python runtime, independent of model +size) and `reshard` re-cuts the weights the same way. The ONE part that holds +real memory is the optimizer re-cut inside `reshard` (one NEW rank's optimizer +state, e.g. 12 B/param/new_ws for fp32-master AdamW); it prints that estimate +and refuses to start when the box cannot afford it (CLAUDE.md: 45 GB +MemAvailable floor for anything heavy, no swap). + +WHY THIS EXISTS. `train.checkpoint_format = "sharded"` writes per-rank files +that eval, the rollout weight sync and external tools cannot read. `export` +turns one into exactly the directory the gathered writer would have produced +(byte-identical; tested by file hash) without ever holding the model in RAM -- +the gathered path's rank-0 gather measured +28 GB RSS on OLMoE (docs/ +MEASUREMENTS.md, Track M). `reshard` is the ONLY way to resume at a different +world size: resume refuses it, because sharded weights are rank-local row +slices and silently loading them under another partitioning trains plausibly +on scrambled weights. + +Exit status: 0 ok, 1 refused/corrupt (message on stderr), 2 usage. +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys + +os.environ.setdefault("CUDA_VISIBLE_DEVICES", "") # never touch a GPU from a converter + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from align.train import sharded_ckpt as sc # noqa: E402 + + +def _mem_available_gib() -> float: + with open("/proc/meminfo") as f: + for line in f: + if line.startswith("MemAvailable:"): + return int(line.split()[1]) / 2**20 + return 0.0 + + +def _optimizer_bytes(src: str, ws: int) -> int: + return sum(os.path.getsize(os.path.join(src, f"optimizer_rank{r}.pt")) for r in range(ws) if os.path.exists(os.path.join(src, f"optimizer_rank{r}.pt"))) + + +def cmd_inspect(a) -> int: + m = sc.read_manifest(a.src) + total = sum(f["bytes"] for f in m["files"].values()) + print( + json.dumps( + { + "world_size": m["world_size"], + "tensors": len(m["tensors"]), + "weight_bytes": total, + "export_family": (m.get("export") or {}).get("family"), + "resharded_from": m.get("resharded_from"), + "problem": sc.sharded_problem(a.src, lambda p: None), + }, + indent=1, + ) + ) + return 0 + + +def cmd_verify(a) -> int: + m = sc.read_manifest(a.src) + bad = [] + for r in range(m["world_size"]): + p = sc.verify_rank_file(a.src, r) + print(f"rank {r}: {'OK' if p is None else p}") + if p: + bad.append(p) + return 1 if bad else 0 + + +def cmd_export(a) -> int: + from align.utils.guard import assert_paths_off_root + + assert_paths_off_root([os.path.dirname(os.path.abspath(a.dst))]) + res = sc.export_hf(a.src, a.dst, verify=not a.no_verify) + print(json.dumps({**res, "peak_rss_kb": sc.peak_rss_kb()})) + return 0 + + +def cmd_reshard(a) -> int: + from align.utils.guard import assert_paths_off_root + + assert_paths_off_root([os.path.dirname(os.path.abspath(a.dst))]) + m = sc.read_manifest(a.src) + if not a.drop_optimizer: + est = _optimizer_bytes(a.src, m["world_size"]) / max(1, a.world_size) / 2**30 + avail = _mem_available_gib() + print(f"optimizer re-cut peak ~ {est:.1f} GiB; MemAvailable {avail:.0f} GiB", file=sys.stderr) + if not a.allow_large and avail < est + 45: + print( + f"refusing: need MemAvailable >= {est + 45:.0f} GiB (estimate + the 45 GiB box floor) and have " + f"{avail:.0f}. Wait for the box, or pass --allow-large if you accept the risk, or --drop-optimizer " + "(fresh moments; NOT lossless).", + file=sys.stderr, + ) + return 1 + res = sc.reshard(a.src, a.dst, a.world_size, drop_optimizer=a.drop_optimizer) + print(json.dumps({**res, "peak_rss_kb": sc.peak_rss_kb()})) + return 0 + + +def main(argv=None) -> int: + ap = argparse.ArgumentParser(description=__doc__.split("\n")[0]) + sub = ap.add_subparsers(dest="cmd", required=True) + for name in ("inspect", "verify"): + p = sub.add_parser(name) + p.add_argument("src") + p = sub.add_parser("export") + p.add_argument("src") + p.add_argument("dst") + p.add_argument("--no-verify", action="store_true", help="skip the per-file sha256 pass (not recommended)") + p = sub.add_parser("reshard") + p.add_argument("src") + p.add_argument("dst") + p.add_argument("--world-size", type=int, required=True) + p.add_argument("--drop-optimizer", action="store_true", help="omit optimizer files: resume starts FRESH moments (NOT lossless)") + p.add_argument("--allow-large", action="store_true") + a = ap.parse_args(argv) + try: + return {"inspect": cmd_inspect, "verify": cmd_verify, "export": cmd_export, "reshard": cmd_reshard}[a.cmd](a) + except (sc.ShardedCheckpointError, FileExistsError, FileNotFoundError, KeyError) as e: + print(f"{type(e).__name__}: {e}", file=sys.stderr) + return 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/code/scripts/dedup_arm_b_shards.py b/code/scripts/dedup_arm_b_shards.py new file mode 100644 index 0000000000000000000000000000000000000000..de512def0b6ddcc1ad9993140aaef844c88234e6 --- /dev/null +++ b/code/scripts/dedup_arm_b_shards.py @@ -0,0 +1,308 @@ +"""One-shot, surgical dedup of the Arm B preference shards. + +WHY THIS EXISTS. The ref-logprob cache build refused the Arm B shards with +`ValueError: duplicate pair_id` -- correctly. The train split held 11,404 +rows with 87 pair_ids present exactly twice, and held_out held 596 rows with +2 doubled ids; every second occurrence is byte-identical to its first +(same input_ids, boundaries, rejected side, and extras). Root cause: PRM800K +carries multiple annotation records per problem, so two records can yield the +IDENTICAL pair site; `PairCapper.offer` caps per (problem_id, step_index) but +never checks pair_id, and `curate()` merges everything per problem +(nmo/data/prm800k_pairs.py, `by_problem.setdefault(...).append(p)`) without a +seen-pair_id check -- so the depth-round top-up re-selects a pair the first +round already took. + +WHAT IT DOES. For each split (train, held_out): + 1. read the live shards via the store's own reader; + 2. keep the FIRST occurrence of every pair_id, drop second-and-later + occurrences -- REFUSING unless the dropped row is content-identical to + the kept one (a non-identical duplicate is a different bug and must not + be silently resolved by row order); + 3. write fresh shards through the same ShardWriter/receipt API the build + used, preserving shard naming (`arm_b-train-NNNN` etc.). +Then it asserts zero train/held pair_id overlap, atomically swaps +`arm_b -> arm_b_predup` and `arm_b_dedup -> arm_b` (the predup directory is +KEPT as evidence), and updates data/manifest.json in place: +`counts.arm_b_curated`, the two `arm_b-*` shard hashes (same +content_hash-of-receipt-hashes chain `_ShardSet.finalize` uses), and the +`length_bias_baseline` Arm B entries re-measured over the deduped TRAIN split +with the exact function scripts/dpo.py uses +(align.train.dpo_stage.measure_chosen_longer_fraction). A provenance block +`arm_b_dedup` records old -> new for everything it touched. + +Deterministic (pure function of the input shards; no RNG), CPU-only. +The ref-logprob caches keyed on the old manifest hashes are stale after this +and must be re-precomputed -- that is the point. + +Run: + + cd /apps/development/rlhf/nano-mistral-o1 + CUDA_VISIBLE_DEVICES= .venv/bin/python scripts/dedup_arm_b_shards.py --dry-run + CUDA_VISIBLE_DEVICES= .venv/bin/python scripts/dedup_arm_b_shards.py + +Refuses to run twice: an existing `arm_b_predup` means the swap already +happened. +""" + +from __future__ import annotations + +import os + +# CPU-only by construction: never let a lazily-imported torch see a GPU. +os.environ.setdefault("CUDA_VISIBLE_DEVICES", "") + +import argparse +import json +import logging +from pathlib import Path + +import numpy as np + +logger = logging.getLogger("dedup_arm_b") + +REPO = Path("/apps/development/rlhf/nano-mistral-o1") +SPLITS = (("train", "arm_b-train"), ("held_out", "arm_b-held")) +#: manifest hash key per split prefix, matching build_all's `hashes` dict. +MANIFEST_HASH_KEY = {"arm_b-train": "arm_b-train", "arm_b-held": "arm_b-held"} + + +def rows_identical(a, b) -> bool: + """Full content identity between two store Rows -- every stored field. + + Anything short of this (say, comparing input_ids only) would let a real + divergence hide behind a shared pair_id and get resolved by row order. + """ + return ( + np.array_equal(a.input_ids, b.input_ids) + and np.array_equal(a.boundary_positions, b.boundary_positions) + and a.problem_id == b.problem_id + and a.subject == b.subject + and sorted(a.ragged) == sorted(b.ragged) + and all(np.array_equal(a.ragged[f], b.ragged[f]) for f in a.ragged) + and a.extras == b.extras + ) + + +def dedup_row_indices(ds) -> tuple[list[int], list[int]]: + """(kept, dropped) row indices: first occurrence per pair_id wins. + + Raises if any dropped row is NOT content-identical to its kept first + occurrence -- that would be a semantic collision, not a re-added pair, + and this script is only licensed to remove exact re-adds. + """ + first_by_pair: dict[str, int] = {} + kept: list[int] = [] + dropped: list[int] = [] + for i in range(len(ds)): + row = ds[i] + pair_id = row.extras["pair_id"] + if pair_id not in first_by_pair: + first_by_pair[pair_id] = i + kept.append(i) + continue + j = first_by_pair[pair_id] + if not rows_identical(ds[j], row): + raise ValueError( + f"pair_id {pair_id!r} occurs at rows {j} and {i} with " + "DIFFERENT content. That is a hash/semantics collision, not a " + "duplicate re-add; refusing to pick a winner by row order." + ) + dropped.append(i) + return kept, dropped + + +def write_split(ds, kept: list[int], out_dir: Path, prefix: str) -> str: + """Write the kept rows through the build's own shard-set API. + + Returns the manifest-level content hash (content_hash over the per-shard + receipt hashes -- the same chain `_ShardSet.finalize` produces, so the + updated manifest entry is computed by the code that defined it). + """ + from align.data.build import _PREF_FIELDS, _ShardSet + + shard_set = _ShardSet(out_dir, prefix, _PREF_FIELDS) + for i in kept: + row = ds[i] + shard_set.add( + input_ids=row.input_ids, + problem_id=row.problem_id, + subject=row.subject, + boundary_positions=row.boundary_positions, + ragged={ + "rejected_input_ids": row.ragged["rejected_input_ids"], + "rejected_boundary_positions": row.ragged["rejected_boundary_positions"], + }, + **row.extras, + ) + return shard_set.finalize() + + +def step_basis_fraction(rows) -> float: + """Arm B's step-basis chosen-longer fraction over collator-shaped rows. + + Mirrors build.py's `chosen_longer_fraction_step_tokens`: length measured + from the SHARED PREFIX (the one step that differs), counted over every + row. Kept next to the completion-basis measurement so the manifest's two + Arm B baselines stay siblings of the same dataset state. + """ + longer = sum( + int( + len(r["chosen_input_ids"]) - int(r["shared_prefix_len"]) + > len(r["rejected_input_ids"]) - int(r["shared_prefix_len"]) + ) + for r in rows + ) + return longer / len(rows) if rows else float("nan") + + +def main(argv=None) -> int: + logging.basicConfig(level=logging.INFO, format="%(levelname)s %(name)s: %(message)s") + ap = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + ap.add_argument("--root", type=Path, default=REPO, help="repo root") + ap.add_argument( + "--dry-run", action="store_true", help="measure and report; write and swap nothing" + ) + args = ap.parse_args(argv) + + from align.data.loader import preference_dataset + from align.data.store import ShardDataset + + shards_root = args.root / "data" / "shards" + live = shards_root / "arm_b" + staged = shards_root / "arm_b_dedup" + predup = shards_root / "arm_b_predup" + manifest_path = args.root / "data" / "manifest.json" + + if predup.exists(): + raise SystemExit( + f"{predup} already exists: the swap already happened. This script " + "is one-shot; a second run would dedup the already-deduped shards " + "and overwrite the evidence." + ) + if staged.exists(): + raise SystemExit( + f"{staged} already exists (a previous run died mid-write?). " + "Inspect and delete it before re-running; a half-written shard " + "set masquerades as resumable." + ) + + # -- pass 1: dedup into the staging directory --------------------------- + pair_ids: dict[str, set[str]] = {} + before_after: dict[str, tuple[int, int]] = {} + new_hashes: dict[str, str] = {} + for split, prefix in SPLITS: + ds = ShardDataset(live / split) + kept, dropped = dedup_row_indices(ds) + before_after[split] = (len(ds), len(kept)) + pair_ids[split] = {ds[i].extras["pair_id"] for i in kept} + logger.info( + "%s: %d rows -> %d kept (%d exact-duplicate re-adds dropped)", + split, + len(ds), + len(kept), + len(dropped), + ) + if not args.dry_run: + new_hashes[prefix] = write_split(ds, kept, staged / split, prefix) + + overlap = pair_ids["train"] & pair_ids["held_out"] + if overlap: + raise SystemExit( + f"{len(overlap)} pair_ids appear in BOTH deduped splits (e.g. " + f"{sorted(overlap)[:3]}); the train/held split is broken and a " + "dedup must not paper over it." + ) + + if args.dry_run: + logger.info("dry run: nothing written, nothing swapped") + return 0 + + # -- pass 2: verify the staged shards through a fresh reader ------------ + for split, _prefix in SPLITS: + staged_ds = ShardDataset(staged / split) + staged_ids = [staged_ds[i].extras["pair_id"] for i in range(len(staged_ds))] + if len(staged_ids) != len(set(staged_ids)) or len(staged_ids) != before_after[split][1]: + raise SystemExit( + f"staged {split} re-read: {len(staged_ids)} rows, " + f"{len(set(staged_ids))} unique -- the write did not produce " + "what the dedup computed; refusing to swap." + ) + + # -- pass 3: atomic swap; the pre-dedup shards are KEPT as evidence ----- + live.rename(predup) + staged.rename(live) + logger.info("swapped: %s -> %s, %s -> %s", "arm_b", "arm_b_predup", "arm_b_dedup", "arm_b") + + # -- pass 4: manifest update, through the numbers' own definitions ------ + manifest = json.loads(manifest_path.read_text()) + old = { + "counts.arm_b_curated": manifest["counts"]["arm_b_curated"], + "hashes.arm_b-train": manifest["hashes"]["arm_b-train"], + "hashes.arm_b-held": manifest["hashes"]["arm_b-held"], + "arm_b_chosen_longer_fraction": manifest["length_bias_baseline"][ + "arm_b_chosen_longer_fraction" + ], + "arm_b_chosen_longer_fraction_step_basis": manifest["length_bias_baseline"][ + "arm_b_chosen_longer_fraction_step_basis" + ], + } + + # Re-measure the length-bias baseline over the deduped TRAIN split with + # the exact function scripts/dpo.py falls back to, on the exact rows it + # would build (read-only imports; nothing here re-implements the metric). + from align.train.dpo_stage import measure_chosen_longer_fraction + + rows = preference_dataset(live / "train") + new_fraction = measure_chosen_longer_fraction(rows) + new_fraction_step = step_basis_fraction(rows) + + n_train, n_held = before_after["train"][1], before_after["held_out"][1] + manifest["counts"]["arm_b_curated"] = n_train + n_held + manifest["hashes"]["arm_b-train"] = new_hashes["arm_b-train"] + manifest["hashes"]["arm_b-held"] = new_hashes["arm_b-held"] + manifest["length_bias_baseline"]["arm_b_chosen_longer_fraction"] = new_fraction + manifest["length_bias_baseline"]["arm_b_chosen_longer_fraction_step_basis"] = new_fraction_step + manifest["arm_b_dedup"] = { + "reason": ( + "87 train + 2 held_out pair_ids were present exactly twice, every " + "second occurrence content-identical to its first (curation " + "depth-round top-up re-added pairs already selected; the merge in " + "prm800k_pairs.curate never deduped by pair_id). The ref-logprob " + "cache build correctly refused the duplicates." + ), + "script": "scripts/dedup_arm_b_shards.py", + "train_rows": {"before": before_after["train"][0], "after": n_train}, + "held_out_rows": {"before": before_after["held_out"][0], "after": n_held}, + "old": old, + "predup_shards": str(predup), + } + manifest_path.write_text(json.dumps(manifest, indent=1)) + + logger.info("train: %d -> %d rows", before_after["train"][0], n_train) + logger.info("held_out: %d -> %d rows", before_after["held_out"][0], n_held) + logger.info("counts.arm_b_curated: %s -> %d", old["counts.arm_b_curated"], n_train + n_held) + logger.info( + "arm_b_chosen_longer_fraction: %s -> %.6f (step basis %s -> %.6f)", + old["arm_b_chosen_longer_fraction"], + new_fraction, + old["arm_b_chosen_longer_fraction_step_basis"], + new_fraction_step, + ) + logger.info( + "hashes: arm_b-train %s -> %s, arm_b-held %s -> %s", + old["hashes.arm_b-train"], + new_hashes["arm_b-train"], + old["hashes.arm_b-held"], + new_hashes["arm_b-held"], + ) + logger.info( + "NOTE: ref caches keyed on the old manifest hashes are now stale; " + "re-run scripts/precompute_ref_logprobs.py before scripts/dpo.py." + ) + return 0 + + +if __name__ == "__main__": + # House rule: everything above is import-safe because spawn-based worker + # pools re-import __main__; the entry point stays behind this guard. + raise SystemExit(main()) diff --git a/code/scripts/dpo.py b/code/scripts/dpo.py new file mode 100644 index 0000000000000000000000000000000000000000..b3b7116247606bdb417fc6677fca8cf084addc80 --- /dev/null +++ b/code/scripts/dpo.py @@ -0,0 +1,590 @@ +#!/usr/bin/env python +"""DPO stage entry: gates -> banner -> train -> arm result manifest. + + torchrun --nproc_per_node=4 scripts/dpo.py configs/dpo_argilla.toml + torchrun --nproc_per_node=4 scripts/dpo.py configs/dpo_steppairs.toml \ + --lr=5e-7 + +One script, two pre-registered arms; the config selects everything. The +objective is section-05's chunked `fused_dpo`, the loop is section-07's +`Trainer`, the reference logprobs are constants read from the cache +scripts/precompute_ref_logprobs.py wrote. What this script owns is the +decision to start and the record of what happened. + +Entry gates run in `dpo_stage.GATE_ORDER` and any failure blocks everything +after it, model load included. They are ordered by cost so an invalid stage +dies in milliseconds rather than after putting 14.5GB onto a shared no-swap +box: + + 1. section-11's SFT quality-gate receipt exists AND passed; + 2. build-manifest counts match the config, pins were verified, Arm A is not + marked DESCOPED; + 3. the reference cache's six-component key matches this run exactly; + 4. the path-bias measurement exists and its decision was honoured; + 5. the 20-step throughput probe (wall-clock is DERIVED from it, never + asserted in advance); + 6. the RSS-aware preflight for the full-FT job class. + +On a tripwire fire the arm re-runs ONCE with the pre-registered response from +docs/RESULTS.md, and that configuration IS the arm's result -- written into +the arm manifest, which section-13 reads for arm selection and the paired +evaluation. +""" + +from __future__ import annotations + +import argparse +import json +import logging +import os +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from align.utils.guard import apply_resource_guards # noqa: E402 + +logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") +logger = logging.getLogger("dpo") + + +def _split_overrides(extra: list[str]) -> list[str]: + """Leftover argv (what argparse did not claim) -> dotted `key=value` + overrides for `dpo_stage.load_stage_config`. + + `--lr=5e-7` and `lr=5e-7` are the same override. Anything without an `=` + is refused: argparse's own positional `nargs='*'` accepted the bare form + and then nothing ever read it, so an override the operator believed + applied trained at the TOML's value (the banner printed the TOML value, + not the intended one). Every override is either applied or rejected. + """ + overrides = [a.lstrip("-") for a in extra] + bad = [a for a, o in zip(extra, overrides) if "=" not in o or not o.split("=", 1)[0]] + if bad: + raise SystemExit( + f"unrecognized arguments {bad}: overrides must be dotted key=value pairs " + "(e.g. --lr=5e-7); space-separated values are not supported" + ) + return overrides + + +def _parse_args(argv=None): + """argparse for the flags, hand-split for the dotted overrides. + + `--lr=5e-7` is not an argparse option, so the declared-positional form + rejected the documented spelling (exit 2) while accepting the bare + `lr=5e-7` and then never reading it. Unclaimed arguments are now the + overrides: applied (struct mode, via `load_stage_config`) or refused. + """ + ap = argparse.ArgumentParser(allow_abbrev=False, epilog="dotted overrides: --lr=5e-7") + ap.add_argument("config") + ap.add_argument( + "--resume", + default=os.environ.get("ALIGN_RESUME") == "1", + action=argparse.BooleanOptionalAction, + help="resume from the newest valid checkpoint in run_dir instead of rotating " + "it away as stale (default: fresh arm; the supervisor sets ALIGN_RESUME=1 on retries)", + ) + ap.add_argument("--strict-host-check", default=True, action=argparse.BooleanOptionalAction) + ap.add_argument("--heavy-lock", default=True, action=argparse.BooleanOptionalAction) + args, extra = ap.parse_known_args(argv) + args.overrides = _split_overrides(extra) + return args + + +def main() -> int: + args = _parse_args() + + from align.config import default_full_ft_config + + def _rss_cfg(): + # Project host RSS from the optimizer this stage ACTUALLY runs. + # default_full_ft_config defaults to offload_adamw (104.7GB + # projection) while the DPO arms are pinned to adamw8bit (22GB) -- + # the mismatch made the RSS gate refuse a job 5x smaller than its + # model whenever the host dipped below ~150GB free. + c = default_full_ft_config() + c.optimizer.name = "adamw8bit" + return c + + apply_resource_guards( + cpu_threads=8, + heavy=args.heavy_lock, + label="dpo", + strict_host_check=args.strict_host_check, + rss_cfg=_rss_cfg(), + ) + + # BEFORE anything touches CUDA -- the duplicate-GPU class, third + # appearance (sft.py, precompute --path-bias, now here): without + # set_device(local_rank) every rank's bare "cuda" is device 0, and this + # stage's first real 4-rank launch put four ~5.4GB policy loads on GPU 0. + from align.utils.dist import init_distributed + + # The run dir (NCCL flight-recorder dumps land under it, never /tmp) comes + # from the config, which is parsed before the process group so a bad TOML + # fails in milliseconds. + from align.train.dpo_stage import load_stage_config as _early_stage_config + + init_distributed(run_dir=_early_stage_config(args.config, args.overrides).run_dir) + + from align.train.stage_gates import assert_distinct_devices + + assert_distinct_devices() + + from align.data.loader import preference_dataset + from align.model import checkpoint as ckpt + from align.model.mistral import NmoMistralForCausalLM + from align.train.dpo_stage import ( + PreferenceBatchSource, + apply_tripwire_response, + arm_result_manifest, + assert_launch_topology, + build_manifest_gate, + load_run_ref_caches, + load_stage_config, + make_preference_sampler, + measure_chosen_longer_fraction, + read_sft_quality_gate, + reduce_tripwire_alarms, + resolve_pair_split_dir, + resolve_sft_checkpoint, + run_entry_gates, + write_arm_result, + ) + from align.train.ref_cache import read_path_bias_report + from align.utils.guard import assert_paths_off_root + + cfg = load_stage_config(args.config, args.overrides) + assert_paths_off_root([cfg.run_dir]) + # LAUNCH REFUSAL, before any gate: world_size x micro x grad_accum must be + # the pre-registered global_batch the banner and horizon are quoted in. + # Fed by the live process group, so a 2-rank launch of a 4-rank + # pre-registration dies here in milliseconds, not after the model load. + logger.info("launch topology ok: %s", assert_launch_topology(cfg)) + with open(cfg.manifest_path) as f: + manifest = json.load(f) + + state: dict = {"resume": bool(args.resume)} + + # -- the gate sequence ------------------------------------------------ + # Each gate is a closure over `state` so a later gate can use what an + # earlier one established -- which is the point of the ORDER: reaching + # gate N means gates 1..N-1 held. + def gate_sft(_cfg): + state["sft_gate"] = read_sft_quality_gate(_cfg.sft_quality_gate_path) + # The trainer writes /checkpoints/micro_NNNNNNNN with no + # stable `final` alias; the receipt records which one passed. Resolve + # HERE, at the first gate, so every later consumer (ref-cache key, + # model load, banner) sees one path -- never a hand-composed one. + ckpt_path, source = resolve_sft_checkpoint( + _cfg.sft_checkpoint, state["sft_gate"].get("checkpoint") + ) + state["sft_checkpoint"] = ckpt_path + state["sft_gate"]["checkpoint"] = ckpt_path + return {"checkpoint": ckpt_path, "checkpoint_source": source} + + def gate_manifest(_cfg): + result = build_manifest_gate(_cfg, manifest) + state["descoped"] = result["descoped"] + return result + + def gate_ref_cache(_cfg): + # Content-hash key (checkpoint WEIGHT BYTES, not the tensor inventory); + # forward_impl is read off the cache itself, not the path-bias report -- + # see dpo_stage.expected_ref_cache_key for why each. + state["ref_cache"], state["heldout_cache"] = load_run_ref_caches( + _cfg, manifest, state["sft_checkpoint"] + ) + state["ref_cache_key"] = state["ref_cache"].key + return {"pairs": len(state["ref_cache"]), "heldout_pairs": len(state["heldout_cache"])} + + def gate_path_bias(_cfg): + state["path_bias"] = read_path_bias_report(_cfg.path_bias_report_path) + return {"action": state["path_bias"]["action"], "ratio": state["path_bias"]["ratio"]} + + def gate_probe(_cfg): + # Filled after the model exists; the probe is the one gate that needs + # the loop, so it records its intent here and the banner reads the + # measured value below. + state["probe"] = { + "step_time_s_median": float("nan"), + "projected_wall_clock_h": float("nan"), + "status": "PENDING", + } + return {"status": "PENDING until the 20-step probe runs"} + + def gate_rss(_cfg): + from align.train.probes import projected_job_rss + + state["rss"] = projected_job_rss(_rss_cfg()) + return dict(state["rss"]) + + gate_report = run_entry_gates( + cfg, + { + "sft_quality_gate": gate_sft, + "build_manifest": gate_manifest, + "ref_cache_key": gate_ref_cache, + "path_bias": gate_path_bias, + "throughput_probe": gate_probe, + "rss_preflight": gate_rss, + }, + ) + + if state["descoped"]: + # A descoped arm is a RESULT, not an absence: section-13's selection + # rule reads this manifest and must be able to tell "measured badly" + # from "never measured". + out = arm_result_manifest( + cfg=cfg, descoped=True, summary={}, heldout={}, banner={}, checkpoint=None + ) + path = write_arm_result(os.path.join(cfg.run_dir, "arm_result.json"), out) + logger.warning("%s is DESCOPED; wrote %s and did not train", cfg.arm, path) + return 0 + + # -- data + model ------------------------------------------------------ + # DISCOVERY, not path composition: the configs point at the real split + # dirs, and resolve_pair_split_dir re-finds them from the receipts the + # store wrote -- so a layout change moves this by itself, and a wrong + # configured path fails naming the splits that ARE present. + rows = preference_dataset(resolve_pair_split_dir(cfg.shard_dir, f"{cfg.arm}-train")) + heldout_rows = preference_dataset( + resolve_pair_split_dir(cfg.heldout_shard_dir, f"{cfg.arm}-held") + ) + pad_id = int(manifest.get("pinned", {}).get("pad_token_id", 0)) + + baselines = manifest.get("length_bias_baseline", {}) + baseline_key = f"{cfg.arm}_chosen_longer_fraction" + if baseline_key in baselines: + length_baseline = (baselines[baseline_key], "build_manifest") + else: + # Section-03 records only Arm A's baseline. Measure this arm's from + # its own shards rather than borrow the other's -- a tripwire armed + # against the wrong distribution fires forever or never. + length_baseline = (measure_chosen_longer_fraction(rows), "measured_at_launch") + logger.info("length-drift baseline: %.4f (%s)", *length_baseline) + + model = _load_policy( + NmoMistralForCausalLM, ckpt, state["sft_checkpoint"], cfg.attn_implementation + ) + # The RESOLVED backend, read back off the model: the constructor routed + # the request through the gated selector, so this is what actually runs. + state["attn_backend"] = model.attn_implementation + run_cfg = _run_config(cfg, length_baseline[0]) + + from align.train.fsdp import apply_activation_checkpointing, apply_fsdp + + dist_model = apply_fsdp(apply_activation_checkpointing(model)) + import torch as _torch + + # Post-shard residency: a rank should hold ~3.6GB (its shard). Anything + # near 18GB means the pre-shard replica leaked past fully_shard. + logger.info("post-shard cuda allocated: %.2f GiB", _torch.cuda.memory_allocated() / 2**30) + + source = PreferenceBatchSource( + rows, + make_preference_sampler( + rows, + max_tokens=run_cfg.train.seq_len * cfg.micro_batch_size, + max_rows=cfg.micro_batch_size, + seed=run_cfg.train.seed, + ), + pad_token_id=pad_id, + ) + + result = _train_arm( + cfg, + run_cfg, + dist_model, + source, + state, + heldout_rows, + pad_id, + gate_report, + manifest, + length_baseline, + ) + + # -- tripwire response ------------------------------------------------- + # The rerun is collective, so the verdict is drawn from the UNION of every + # rank's alarms: rank-local alarms could send some ranks into a second + # model load + FSDP wrap while the rest exit. + alarms = reduce_tripwire_alarms(result["summary"].get("tripwire_alarms", [])) + rerun_cfg, tripwire = apply_tripwire_response(cfg, alarms) + if tripwire["rerun_required"]: + logger.error( + "chosen-logprob collapse detected; re-running %s with the pre-registered response: %s", + cfg.arm, + tripwire["configuration"], + ) + model = _load_policy( + NmoMistralForCausalLM, ckpt, state["sft_checkpoint"], cfg.attn_implementation + ) + state["attn_backend"] = model.attn_implementation + dist_model = apply_fsdp(apply_activation_checkpointing(model)) + result = _train_arm( + rerun_cfg, + _run_config(rerun_cfg, length_baseline[0]), + dist_model, + source, + state, + heldout_rows, + pad_id, + gate_report, + manifest, + length_baseline, + suffix="_response", + ) + cfg = rerun_cfg + + out = arm_result_manifest( + cfg=cfg, + descoped=False, + summary=result["summary"], + heldout=result["heldout"], + banner=result["banner"], + checkpoint=result["checkpoint"], + tripwire=tripwire, + ) + out["gate_report"] = gate_report + path = write_arm_result(os.path.join(cfg.run_dir, "arm_result.json"), out) + logger.info("%s complete: %s", cfg.arm, path) + return 0 + + +def _load_policy(model_cls, ckpt, path: str, attn_implementation: str): + """Load the policy with the CONFIG-DRIVEN attention backend. + + The string is never handed raw to a constructor in the sense that + matters: `NmoMistralForCausalLM.__init__` routes it through + `align.model.attention.select_attn_implementation`, which raises if the + request cannot be honoured -- for the pre-registered `triton_fa2` that + means section-06's equivalence-gate receipt must be valid for this kernel + source on this GPU. Silent degradation to sdpa is exactly the failure the + selector exists to prevent. + """ + import torch + + from align.model.hf_config import read_config + + model_cfg = read_config(path) + model = model_cls.from_state_dict( + model_cfg, + ckpt.hf_to_nmo(ckpt.load_hf_safetensors(path)), + attn_implementation=attn_implementation, + device="cuda", + ) + return model.to("cuda", dtype=torch.bfloat16) + + +def _run_config(cfg, length_baseline: float): + """Stage config -> section-07 RunConfig. The two are different objects on + purpose: RunConfig is the trainer's surface and DpoStageConfig is the + arm's pre-registration, and collapsing them would let a trainer default + silently become an experimental parameter.""" + from align.config import RunConfig + + run = RunConfig() + # adamw8bit per the 2026-08-17 pre-registration amendment ("SFT (and the + # DPO arms) run adamw8bit"). Left unset, RunConfig's offload_adamw + # default leaked through and the pinned-staging guard refused with a + # 27GiB request on a no-swap box -- correctly. + run.optimizer.name = "adamw8bit" + run.optimizer.lr = cfg.lr + run.optimizer.grad_accum = cfg.grad_accum + run.train.epochs = cfg.epochs + run.train.micro_batch_size = cfg.micro_batch_size + run.train.save_every_micro_steps = cfg.save_every_micro_steps + run.train.keep_checkpoints = cfg.keep_checkpoints + # Forwarded explicitly: RunConfig is built from defaults here, so a stage + # field that is not copied is a silently disarmed config. + run.train.per_rank_optimizer_state = cfg.per_rank_optimizer_state + run.optimizer.max_grad_norm = cfg.max_grad_norm + run.train.max_wall_seconds = cfg.max_wall_seconds + run.train.max_consecutive_skipped_steps = cfg.max_consecutive_skipped_steps + run.train.max_skipped_fraction = cfg.max_skipped_fraction + run.train.full_ft = True + run.data.num_workers = cfg.num_workers + run.tripwires.chosen_longer_baseline = length_baseline + run.tripwires.length_band = cfg.length_band + run.tripwires.logps_window = cfg.logps_window + run.tripwires.armed = tuple(cfg.armed_tripwires) + # The tripwire may speak exactly when its estimator is full. These are two + # different quantities -- `length_window_pairs` is the span the adapter + # averages over, `length_min_pairs` is when the tripwire is allowed to + # fire -- and they are deliberately the SAME number: a min_pairs below the + # window size would let the guard speak from a partially-filled window, + # which is precisely the noisy regime both settings exist to avoid. + run.tripwires.length_min_pairs = cfg.length_window_pairs + run.run_name = f"dpo_{cfg.arm}" + return run + + +def _train_arm( + cfg, + run_cfg, + dist_model, + source, + state, + heldout_rows, + pad_id, + gate_report, + manifest, + length_baseline, + suffix: str = "", +): + """Probe -> banner -> train -> held-out accuracy. In that order, because + the banner quotes the PROBE-DERIVED wall-clock and the plan forbids + asserting a wall-clock in advance.""" + from align.data.loader import collate_preference + from align.model.attention import banner_value + from align.train.dpo_stage import ( + DpoBatch, + dpo_banner, + make_dpo_stage_adapter, + render_dpo_banner, + ) + from align.train.ref_cache import heldout_preference_accuracy + from align.train.trainer import Trainer + + run_dir = cfg.run_dir + suffix + # FRESH-ARM HYGIENE: a non-resuming arm run must not share a checkpoint + # dir with a previous run's leftovers -- stale higher-numbered dirs + # shadowed the rerun's fresh saves under keep-N pruning (2026-08-20). + from align.train.dpo_stage import probe_micro_cap, rotate_stale_checkpoints + + ckpt_dir = os.path.join(run_dir, "checkpoints") + # Rotation is for FRESH arms only. A resuming arm (--resume, or a + # supervisor retry via ALIGN_RESUME=1) must keep its checkpoints: rotating + # them would silently restart the arm from step 0 while the supervisor + # believes it is continuing -- the replay defect, at the stage level. + resuming = bool(state.get("resume")) + # Collective: rank 0 renames, every rank waits (a bare per-rank rename + # let one rank win and killed the other three with FileNotFoundError). + if not resuming: + rotate_stale_checkpoints(ckpt_dir) + adapter = make_dpo_stage_adapter(state["ref_cache"], cfg, require_fused=True) + trainer = Trainer(run_cfg, dist_model, source, adapter, run_dir=run_dir, stage="dpo") + + # REFUSE a horizon that disagrees with the corpus, before a single step. + # Arm A ran a 146-step horizon over a corpus needing ~289 per-rank steps: + # the cosine finished at half an epoch, clamped the learning rate to + # exactly 0, and the loop kept stepping on frozen weights for ~24 + # GPU-hours. The Trainer now stops AT the horizon, which turns that into a + # short run rather than a zombie one -- this check is the other half, and + # catches it before the GPU-hours rather than after. + from align.train.dpo_stage import assert_dpo_horizon_consistent + from align.utils.dist import world_size as _world_size + + horizon_report = assert_dpo_horizon_consistent( + trainer.scheduler.total_steps, + n_pairs=len(source.rows), + world_size=_world_size(), + micro_batch_size=cfg.micro_batch_size, + grad_accum=cfg.grad_accum, + epochs=run_cfg.train.epochs, + global_batch=cfg.global_batch, + ) + logger.info("horizon ok: %s", horizon_report) + + if resuming: + resumed_from = trainer.resume_from_latest() + logger.info( + "RESUMED from %s at optimizer step %d / micro step %d", + resumed_from, + trainer.optimizer_step, + trainer.micro_step, + ) + state["resumed_from"] = resumed_from + + tokens_per_step = ( + run_cfg.train.seq_len * cfg.micro_batch_size * 2 * cfg.grad_accum + ) # 2: both sides of every pair + # The probe is the FIRST probe_steps real micro-steps of the run itself: + # a capped train(), report captured, cap removed, training continues from + # the stream's exact position. The previous construction asked a trainer + # that had never stepped for its step times, and throughput_report + # correctly refused an empty probe. Probe steps here are real training + # (counted in the horizon), so nothing is spent twice; the banner prints + # after 20 micro-steps instead of before any. The cap is on the ABSOLUTE + # micro-step counter, so a resumed arm probes its first probe_steps AFTER + # the resume point (a bare probe_steps cap trained one micro-step and + # crashed throughput_probe on every supervisor retry). + run_cfg.train.max_micro_steps = probe_micro_cap(trainer.micro_step, run_cfg.train.probe_steps) + trainer.train() + run_cfg.train.max_micro_steps = None + if trainer.stop_exit_code: + # A coordinated stop (SIGTERM / train.max_wall_seconds) saved at an + # optimizer-step boundary; the stage is NOT finished. Exit with the + # documented status instead of evaluating / reporting a half-trained arm. + raise SystemExit(trainer.stop_exit_code) + if resuming and len(trainer.step_times) < trainer.accum: + # Resumed within one optimizer step of the horizon (or after the last + # step): the probe window cannot fill a group. The run still finishes; + # the banner says the wall-clock was not measured rather than raising. + probe = { + "step_time_s_median": float("nan"), + "projected_wall_clock_h": float("nan"), + "status": "UNMEASURED_ON_RESUME", + } + else: + probe = trainer.throughput_probe(tokens_per_step) + state["probe"] = probe + + banner = dpo_banner( + cfg=cfg, + manifest=manifest, + ref_cache_key=state["ref_cache_key"], + path_bias=state["path_bias"], + probe=probe, + rss=state["rss"], + backend=banner_value( + # The RESOLVED backend, never a literal: FSDP2 mutates the module + # in place so the attribute normally survives wrapping, and the + # fallback is what the selector resolved at load -- a hardcoded + # "sdpa" here would let the banner contradict the kernel that ran. + dist_model.attn_implementation + if hasattr(dist_model, "attn_implementation") + else state["attn_backend"] + ), + sft_gate=state["sft_gate"], + length_baseline=length_baseline, + ) + print(render_dpo_banner(banner)) + + summary = trainer.train() + if trainer.stop_exit_code: + # A coordinated stop (SIGTERM / train.max_wall_seconds) saved at an + # optimizer-step boundary; the stage is NOT finished. Exit with the + # documented status instead of evaluating / reporting a half-trained arm. + raise SystemExit(trainer.stop_exit_code) + checkpoint = trainer.save_checkpoint() + + import torch as _torch + + _dev = _torch.device("cuda", _torch.cuda.current_device()) + # .to(_dev): the held-out path missed the device move the training + # adapter got -- CPU labels reached the CUDA gather AFTER the full arm + # had trained and saved (Arm A rerun, 2026-08-20). + heldout_batches = [ + collate_preference(heldout_rows[i : i + cfg.micro_batch_size], pad_token_id=pad_id).to(_dev) + for i in range(0, len(heldout_rows), cfg.micro_batch_size) + ] + heldout = heldout_preference_accuracy( + dist_model, heldout_batches, state["heldout_cache"], beta=cfg.beta + ) + logger.info( + "held-out preference accuracy: %.4f [%.4f, %.4f] n=%d", + heldout["accuracy"], + heldout["ci_low"], + heldout["ci_high"], + heldout["n"], + ) + return {"summary": summary, "heldout": heldout, "banner": banner, "checkpoint": checkpoint} + + +if __name__ == "__main__": + from align.exit_codes import run_with_exit_codes + + raise SystemExit(run_with_exit_codes(main)) diff --git a/code/scripts/eval.py b/code/scripts/eval.py new file mode 100644 index 0000000000000000000000000000000000000000..1a2a80d4a4d66ccd44b87923c3f70abad4e9fcd1 --- /dev/null +++ b/code/scripts/eval.py @@ -0,0 +1,102 @@ +"""Thin eval entry: single / paired / golden. + + python scripts/eval.py configs/eval.toml --mode=single --model.path=... + +ORDER OF OPERATIONS IS LOAD-BEARING +----------------------------------- +1. `apply_resource_guards` runs ONLY when we are not inside a multiprocessing + child. The verifier's worker pool uses 'spawn', which RE-IMPORTS this + module in every worker; a resource guard or an exclusive flock at module + scope therefore executes in each child and kills it on startup. Every + future then dies with BrokenProcessPool -- which is precisely the + documented cause of the "all answers incorrect" masquerade the verifier + raises on. The guard is the fix; `_in_multiprocessing_child()` is what + makes it safe. +2. Banner next: resolved config, dataset hashes, MDE, model hash, backend. + A run that cannot answer its own question does not begin. +3. Then dispatch. +""" + +from __future__ import annotations + +import argparse +import os +import sys + +REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +if REPO_ROOT not in sys.path: + sys.path.insert(0, REPO_ROOT) + +from align.utils.guard import _in_multiprocessing_child # noqa: E402 + +# See docstring point 1. This condition is not an optimization. +if not _in_multiprocessing_child(): + from align.utils.guard import apply_resource_guards + + apply_resource_guards(cpu_threads=8, heavy=False, label="eval") + + +def build_parser() -> argparse.ArgumentParser: + p = argparse.ArgumentParser(description="nmo evaluation entry") + p.add_argument("config", nargs="?", default="configs/eval.toml") + p.add_argument("--mode", choices=("single", "paired", "golden"), default="single") + p.add_argument( + "--out", default=None, help="report path; defaults to outputs/eval/ under the repo" + ) + p.add_argument( + "overrides", nargs="*", default=[], help="dotted overrides, e.g. eval.limit=8" + ) + return p + + +def default_out_dir() -> str: + """Reports land under /apps -- asserted, never assumed.""" + from align.utils.guard import assert_paths_off_root + + out = os.path.join(REPO_ROOT, "outputs", "eval") + os.makedirs(out, exist_ok=True) + assert_paths_off_root([out]) + return out + + +def main(argv=None) -> int: + args = build_parser().parse_args(argv) + + from align.config import load_toml + + # Schemaless on purpose (eval.toml has no dataclass): plain nested dict, + # dotted overrides TOML-literal parsed. Same shape OmegaConf.to_object gave. + cfg = load_toml(args.config, list(args.overrides)) + out_dir = args.out or default_out_dir() + + if args.mode == "golden": + from align.eval.golden import run_golden_regression + + banner: dict = {} + path = os.path.join(out_dir, "checkpoint-200-golden.json") + # `run_golden_regression` writes the report ITSELF, before gating, so + # a failing gate still leaves the evidence. Writing again here would + # be dead code on the path that matters: the gate raises, and anything + # after this call never runs. + try: + run_golden_regression(banner=banner, out_path=path) + finally: + # U01-9: name the eval source that produced the report (also when the gate raised), so + # `record_golden_from_report` cannot re-certify it after align/eval/ changes. + if os.path.isfile(path): + from align.train.stage_gates import stamp_golden_report + + stamp_golden_report(path) + print(f"golden gate PASSED; report at {path}") + return 0 + + raise SystemExit( + f"mode {args.mode!r} needs a generation backend, which arrives " + "with section-04 (nmo model) / section-09 (engine). The harness " + "itself is complete and tested: run_eval takes a generate_fn. " + "Use --mode=golden, which uses plain HF generation by design." + ) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/code/scripts/eval_checkpoint.py b/code/scripts/eval_checkpoint.py new file mode 100644 index 0000000000000000000000000000000000000000..3d840a7b28b52e41d1dfb902e4777b144c49a375 --- /dev/null +++ b/code/scripts/eval_checkpoint.py @@ -0,0 +1,175 @@ +"""Standalone checkpoint eval: fresh process, single GPU, golden-proven generator. + +Promotion of the .tmp/eval_gate.py driver (task #26) into a parameterized +entry point. The design facts it carries forward are all incident-derived: + + * FRESH PROCESS, never in-trainer: the train->eval handoff inside one + process produced the rank race (peers built a second 7B while rank 0 + saved) and the orphan-GPU-process incident. A checkpoint on disk plus a + fresh process has neither failure mode. + * EVERYTHING under the __main__ guard: the verifier's scoring pool uses + the 'spawn' context, which RE-IMPORTS __main__ in every worker -- the + unguarded first driver had all 8 workers load their own 7B onto the + eval GPU. + * greedy generation via `golden_generate_fn` (in-repo Engine, same encode + as the golden regression: BOS default, truncation 768). The anchors and + banked SFT rows were measured with HF generate; the generator swap is + certified by the golden regression + checkpoint roundtrip gate, and the + generator is recorded in the report config. + * Atomic report write (.partial -> os.replace): a half-written report + masquerades as a complete one. + +Phase-0 usage (the Arm B result side of the primary McNemar): + + CUDA_VISIBLE_DEVICES=0 python scripts/eval_checkpoint.py \ + --checkpoint runs/dpo_arm_b_response/checkpoints/micro_XXXXXXXX \ + --name arm-b-final --gsm8k-n 1319 --math500-n 0 + +`--quality-gate` additionally runs the SFT-stage quality gate + ledger +records; leave it off for non-SFT checkpoints (the gate's anchors are +SFT-specific). +""" + +import argparse +import json +import os +import sys +from pathlib import Path + +# Run from a checkout without an install: put the repo root on sys.path first. +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from align.utils.guard import apply_resource_guards + + +def quality_gate_size_errors(cfg, math500_n: int, gsm8k_n: int, per_set: dict | None = None) -> list[str]: + """Why a run of these sizes may NOT write the SFT quality-gate receipt (empty = fine). + + The receipt unlocks DPO. `assert_meets_golden` accepts any point estimate inside the target's own + Wilson CI, which is enormous at tiny n (1/5 passes the 0.09 anchor), so the gate certifies only a + FULL-size evaluation: both sets, at the configured n the anchors were recorded at. `per_set` (the + measured report) is checked too, in case a set held fewer problems than requested.""" + want = {"math500": cfg.eval.math500_n, "gsm8k": cfg.eval.gsm8k_n} + asked = {"math500": math500_n, "gsm8k": gsm8k_n} + errs = [f"{k}: requested n={asked[k]}, the quality gate needs the full configured n={want[k]}" + for k in want if asked[k] != want[k]] + for k, w in want.items(): + got = (per_set or {}).get(k, {}).get("n") if per_set is not None else w + if got != w: + errs.append(f"{k}: measured n={got}, the quality gate needs n={w}") + return errs + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument("--checkpoint", required=True, help="checkpoint dir (HF-layout safetensors dir)") + parser.add_argument("--config", default="configs/sft.toml") + parser.add_argument("--name", required=True, help="report model_name and output stem") + parser.add_argument("--out", default=None, help="default: outputs/eval/.json") + parser.add_argument("--math500-n", type=int, default=None, help="0 skips; default: config") + parser.add_argument("--gsm8k-n", type=int, default=None, help="0 skips; default: config") + parser.add_argument( + "--quality-gate", action="store_true", help="SFT-stage gate + ledger records" + ) + args = parser.parse_args() + + # Guards BEFORE torch import (env vars are load-bearing); heavy=True takes + # the exclusive job lock -- a 7B generation eval over 1,319 problems is a + # multi-hour GPU job and must serialise behind training like any other. + apply_resource_guards(cpu_threads=8, heavy=True, label=f"eval-{args.name}") + + + from align.config import load_config # noqa: E402 + from align.eval.golden import golden_generate_fn, load_checkpoint_model # noqa: E402 + from align.eval.harness import EvalConfig, load_eval_set, run_eval # noqa: E402 + from align.eval.verifier import OutcomeVerifier # noqa: E402 + from align.tokenizer import load_tokenizer # noqa: E402 + from align.train.stage_smoke import _manifest_hashes, _mde # noqa: E402 + + cfg = load_config(args.config) + checkpoint = args.checkpoint + assert os.path.isdir(checkpoint), checkpoint + out_path = args.out or os.path.join("outputs", "eval", f"{args.name}.json") + + math500_n = cfg.eval.math500_n if args.math500_n is None else args.math500_n + gsm8k_n = cfg.eval.gsm8k_n if args.gsm8k_n is None else args.gsm8k_n + if args.quality_gate: + errs = quality_gate_size_errors(cfg, math500_n, gsm8k_n) + if errs: + print("refusing --quality-gate on a subset eval (the receipt unlocks DPO):\n " + "\n ".join(errs), + file=sys.stderr) + return 2 + problems = [] + # 0 in this CLI means SKIP the set (checked here); upstream load_eval_set + # uses limit=0 for "no limit", which this branch structure never passes + # ambiguously -- a negative value falls through to upstream validation. + if math500_n != 0: + problems += load_eval_set(cfg.paths.math500_path, limit=math500_n) + if gsm8k_n != 0: + problems += load_eval_set(cfg.paths.gsm8k_path, limit=gsm8k_n) + if not problems: + print("nothing to evaluate (both sets skipped)", file=sys.stderr) + return 2 + + # After every refusal above: a 7B load is the expensive step. + tokenizer = load_tokenizer(checkpoint) + model = load_checkpoint_model(checkpoint, device="cuda") + + verifier = OutcomeVerifier() + gen_fn = golden_generate_fn(model, tokenizer, "cuda") + print(f"evaluating {args.name}: {len(problems)} problems from {checkpoint}", flush=True) + rep = run_eval(problems, gen_fn, verifier, EvalConfig(), model_name=args.name) + per_set = rep.per_set + for name, scores in sorted(per_set.items()): + print(f" {name}: n={scores.get('n')} pass@1={scores.get('pass@1')}", flush=True) + + report = { + "per_set": per_set, + "config": { + "mde": _mde(cfg), + "n_problems": sum(r.get("n", 0) for r in per_set.values()), + "eval_set_hash": _manifest_hashes(cfg), + "generator": "align Engine greedy (golden_generate_fn)", + "checkpoint": checkpoint, + }, + "pooled": rep.pooled, + "per_problem": rep.per_problem, + } + os.makedirs(os.path.dirname(out_path), exist_ok=True) + tmp = out_path + ".partial" + with open(tmp, "w") as f: + json.dump(report, f, indent=1, default=str) + os.replace(tmp, out_path) + print(f"report written: {out_path}", flush=True) + + if not args.quality_gate: + return 0 + + from align.train.stage_gates import ( # noqa: E402 + GateLedger, + eval_source_hash, + quality_gate, + write_quality_gate_receipt, + ) + + errs = quality_gate_size_errors(cfg, math500_n, gsm8k_n, per_set=per_set) + if errs: # report is already written above; only the receipt is withheld + print("refusing to write the quality-gate receipt:\n " + "\n ".join(errs), file=sys.stderr) + return 2 + + run_dir = os.path.join(cfg.paths.checkpoint_dir, "sft") + ledger = GateLedger(run_dir, eval_source_hash=eval_source_hash()) + ledger.record("eval", passed=True, detail=json.dumps(per_set, default=str)) + decision = quality_gate(report, cfg) + ledger.record("quality_gate", passed=decision.passed, detail=decision.halt_message[:400]) + write_quality_gate_receipt(cfg, decision, checkpoint=checkpoint) + print(f"QUALITY GATE: passed={decision.passed}", flush=True) + if not decision.passed: + print(decision.halt_message, file=sys.stderr, flush=True) + return 0 if decision.passed else 1 + + +if __name__ == "__main__": + # See module docstring: spawn workers re-import __main__; nothing above + # this guard may build a model or touch the GPU. + sys.exit(main()) diff --git a/code/scripts/gate_flash_attention.py b/code/scripts/gate_flash_attention.py new file mode 100644 index 0000000000000000000000000000000000000000..f74ac81a74a2f66a0e26151daa79b1a1e46edbed --- /dev/null +++ b/code/scripts/gate_flash_attention.py @@ -0,0 +1,132 @@ +"""Run the triton_fa2 equivalence suite and, ONLY on success, write the gate. + +Run (pin ONE free device): + CUDA_VISIBLE_DEVICES= uv run python scripts/gate_flash_attention.py + +This is the mechanism behind the never-silently-claimed rule. The gate artifact +is the ONLY thing that makes `triton_fa2` selectable +(`align.kernels.gate.resolve_triton_fa2_backend` raises without it), and this +script is the only writer. It records the kernel source hash, the torch/triton +versions, the exact GPU, and the tolerances and shapes the suite covered, so a +later edit to the kernel or a move to a different card invalidates it +automatically. + +Deliberately a separate entry point rather than a pytest hook: writing the gate +must require an explicit human action on a machine with a real GPU, and must +never be a side effect of someone running the test suite with `-k` narrowing +the selection to a subset. +""" + +from __future__ import annotations + +import os +import subprocess +import sys +import xml.etree.ElementTree as ET + +REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, REPO) + +from align.kernels import gate as gate_mod + +#: Kept in step with tests/test_flash_attention.py. Recorded in the artifact so +#: a reader knows what "verified" actually covered. +TOLERANCES = { + "policy": "max |err| vs fp32 reference <= max(3.0 x torch-SDPA-bf16 error, 2e-2)", + "slack_over_sdpa": 3.0, + "absolute_floor": 2e-2, +} + +SHAPES_COVERED = [ + "gqa (n_q_heads, n_kv_heads) in [(8,2), (32,8)] plus MHA-degenerate (4,4)", + "head_dim in [64, 128]", + "seq in [192, 256, 512, 1000] (block-divisible AND non-divisible)", + "causal and non-causal", + "key_padding_mask, right-padded, non-block-aligned", + "every autotune config individually pinned (fwd + both bwd kernels)", + "dtype bfloat16 only (fp16/fp32 rejected)", +] + + +def _count_outcomes(junit_path: str) -> tuple[int, int, int]: + """(passed, skipped, failed+errored) from a pytest junit report. + + pytest exits 0 when EVERY test is skipped (no CUDA, more than one device + visible, heavy-job lock held), so the exit status alone cannot prove the + kernel ran. Raises if the report is missing or unparseable. + """ + root = ET.parse(junit_path).getroot() + passed = skipped = bad = 0 + for case in root.iter("testcase"): + if case.find("skipped") is not None: + skipped += 1 + elif case.find("failure") is not None or case.find("error") is not None: + bad += 1 + else: + passed += 1 + return passed, skipped, bad + + +def main() -> int: + cvd = os.environ.get("CUDA_VISIBLE_DEVICES", "") + devices = [d.strip() for d in cvd.split(",")] + if len(devices) != 1 or not devices[0]: + print( + "refusing to run: set CUDA_VISIBLE_DEVICES to exactly ONE free device " + f"(got {cvd!r}; shared box, and the suite skips unless exactly one is visible).", + file=sys.stderr, + ) + return 2 + + junit = os.path.join(REPO, ".cache", "gate_flash_attention_junit.xml") + os.makedirs(os.path.dirname(junit), exist_ok=True) + if os.path.exists(junit): + os.remove(junit) # a stale report must never vouch for this run + cmd = [ + sys.executable, + "-m", + "pytest", + os.path.join(REPO, "tests", "test_flash_attention.py"), + "-m", + "gpu", + "-q", + "-rs", + f"--junitxml={junit}", + ] + print("running the equivalence suite:", " ".join(cmd), flush=True) + result = subprocess.run(cmd, cwd=REPO) + + if result.returncode != 0: + print( + "\nequivalence suite FAILED -- gate NOT written. triton_fa2 stays " + "unselectable; SDPA remains the backend.", + file=sys.stderr, + ) + return result.returncode + + # Exit 0 is necessary, not sufficient: an all-skipped run also exits 0. + try: + passed, skipped, bad = _count_outcomes(junit) + except (OSError, ET.ParseError) as e: + print(f"\ncannot read the junit report ({e!r}) -- gate NOT written.", file=sys.stderr) + return 1 + if passed == 0 or skipped != 0 or bad != 0: + print( + f"\nequivalence suite did not EXECUTE in full (passed={passed}, " + f"skipped={skipped}, failed={bad}) -- gate NOT written. Skips mean " + "no CUDA device, a device count != 1, or the heavy-job lock held " + "by a live job; none of those verifies the kernel.", + file=sys.stderr, + ) + return 1 + + path = gate_mod.write_equivalence_gate(tolerances=TOLERANCES, shapes_covered=SHAPES_COVERED) + ok, reason = gate_mod.gate_status() + print(f"\nequivalence suite passed ({passed} tests); gate written to {path}") + print(f"gate_status -> passed={ok} ({reason})") + print(f"kernel sha256: {gate_mod.kernel_source_hash()}") + return 0 if ok else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/code/scripts/generate.py b/code/scripts/generate.py new file mode 100644 index 0000000000000000000000000000000000000000..8b4eb66fc2b1a26bcf8012ae54ecb36841020711 --- /dev/null +++ b/code/scripts/generate.py @@ -0,0 +1,327 @@ +#!/usr/bin/env python +"""Generate completions with the align engine (or the bucketed fallback). + + uv run python scripts/generate.py \ + --checkpoint /apps/development/rlhf/nano-mistral-o1/checkpoints/sft \ + --prompts data/eval/math500.json --out .tmp/gen/math500.jsonl + +Sharding across the 4 A10Gs is done by the CALLER, one process per device -- +the engine is single-replica and knows nothing about other GPUs (no +collectives, no NCCL). The idiom, which respects the one-heavy-job rule: + + for i in 0 1 2 3; do + RANK=$i CUDA_VISIBLE_DEVICES=$i uv run python scripts/generate.py \ + --shard-index $i --num-shards 4 --out .tmp/gen/shard$i.jsonl & + done; wait + +`acquire_heavy_job_lock` is rank-aware, so only RANK=0 takes the exclusive +flock; shards 1.. do NOT take it, they VERIFY it (`verify_shard_lock`): a lock +held by anything other than another generate.py (their own shard 0) refuses +the launch, so a live training stage cannot be overlapped by three of the four +shards. Every shard also runs the host preflight. Never launch this against a +live training stage: that overlap is what killed the predecessor's first SFT +run. Completions carry rid ":"; the global prompt index of +local i is shard_index + i * num_shards (stride sharding). + +Every output path is asserted onto /apps (`/` has ~13GB and /tmp is a noexec +tmpfs). The launch banner prints the checkpoint fingerprint, the attention +backend that ACTUALLY ran, the engine budget, the sampling params and the +per-gate status; a missing mandatory key raises rather than printing a partial +banner, and gate status is "PENDING" until recorded -- never assumed passed. +""" + +from __future__ import annotations + +import argparse +import json +import logging +import os +import sys +import time +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from align.utils.guard import apply_resource_guards # noqa: E402 + +logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") +logger = logging.getLogger("generate") + +PROJECT_ROOT = "/apps/development/rlhf/nano-mistral-o1" + + +def validate_shards(shard_index: int, num_shards: int) -> None: + """An out-of-range shard silently yields an EMPTY prompt list (and a "successful" empty run).""" + if num_shards < 1 or not (0 <= shard_index < num_shards): + raise ValueError(f"bad shard spec: --shard-index {shard_index} must be in [0, --num-shards {num_shards})") + + +def _proc_cmdline(pid: int) -> str: + try: + with open(f"/proc/{pid}/cmdline", "rb") as f: + return f.read().replace(b"\0", b" ").decode(errors="replace") + except OSError: + return "" + + +def verify_shard_lock(rank: int, lock_inode: int | None = None, locks_text: str | None = None, cmdline_of=_proc_cmdline) -> None: + """Shards 1.. never take the heavy-job flock (`acquire_heavy_job_lock` returns True for RANK != 0), + so without this a generation launched while a trainer holds it would be refused on shard 0 only, + while shards 1-3 each loaded a 7B onto GPUs 1-3. Read-only: scans /proc/locks for the holder of the + lock file (taking the flock here would race shard 0's own non-blocking acquire). The only holder + tolerated is another generate.py -- this launch's own shard 0.""" + if rank == 0: + return + from align.exit_codes import PreflightRefused + from align.utils.guard import LOCK_PATH + + if lock_inode is None: + try: + lock_inode = os.stat(LOCK_PATH).st_ino + except FileNotFoundError: + return # no lock file yet, so no holder + if locks_text is None: + with open("/proc/locks") as f: + locks_text = f.read() + for line in locks_text.splitlines(): + parts = line.split() + if len(parts) < 6 or parts[1] != "FLOCK" or parts[5].split(":")[-1] != str(lock_inode): + continue + holder = cmdline_of(int(parts[4])) + if "generate.py" not in holder: + raise PreflightRefused( + f"another heavy align job holds the lock (pid {parts[4]}: {holder.strip()[:120]!r}); this shard " + f"(RANK={rank}) must not overlap it. Shard 0 refuses for the same reason." + ) + + +def read_prompts(path: str) -> list[str]: + """Accept .json (list of str, or list of dicts with a prompt/problem/text + key), .jsonl (same per line), or .txt (one prompt per line). Refuses an + empty file loudly: an eval that silently generates nothing looks like a + fast eval.""" + + def _field(obj): + if isinstance(obj, str): + return obj + for k in ("prompt", "problem", "question", "text"): + if k in obj: + return obj[k] + raise ValueError( + f"prompt record has none of prompt/problem/question/text: {sorted(obj)[:8]}" + ) + + if path.endswith(".jsonl"): + with open(path) as f: + rows = [_field(json.loads(line)) for line in f if line.strip()] + elif path.endswith(".json"): + with open(path) as f: + rows = [_field(o) for o in json.load(f)] + else: + with open(path) as f: + rows = [line.rstrip("\n") for line in f if line.strip()] + if not rows: + raise ValueError(f"{path} contains no prompts") + return rows + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument( + "--checkpoint", required=True, help="HF-layout checkpoint directory (safetensors)" + ) + ap.add_argument("--prompts", required=True) + ap.add_argument("--out", required=True, help="completions JSONL; under /apps") + ap.add_argument("--tokenizer", default=None, help="defaults to --checkpoint") + ap.add_argument("--n-slots", type=int, default=32) + ap.add_argument("--max-seq-len", type=int, default=1024) + ap.add_argument("--max-new-tokens", type=int, default=300) + ap.add_argument("--temperature", type=float, default=0.0) + ap.add_argument("--top-p", type=float, default=1.0) + ap.add_argument( + "--stop", action="append", default=[], help="repeatable; matched on DECODED TEXT" + ) + ap.add_argument("--stop-policy", default="trim", choices=("trim", "keep")) + ap.add_argument("--seed", type=int, default=0) + ap.add_argument( + "--attn", default=None, help="force an attention backend; unavailable => hard error" + ) + ap.add_argument( + "--bucketed", + action="store_true", + help="use the static-batching fallback instead of the " + "engine (correctness reference; slower)", + ) + ap.add_argument( + "--benchmark", + action="store_true", + help="run BOTH and write the engine-vs-static-batching speedup receipt beside --out", + ) + ap.add_argument( + "--hash-weights", + action="store_true", + help="full content hash of the checkpoint (reads 14.5GB)", + ) + ap.add_argument("--limit", type=int, default=None) + ap.add_argument("--shard-index", type=int, default=0) + ap.add_argument("--num-shards", type=int, default=1) + ap.add_argument("--strict-host-check", default=True, action=argparse.BooleanOptionalAction) + ap.add_argument( + "--heavy-lock", + default=True, + action=argparse.BooleanOptionalAction, + help="take the exclusive one-heavy-job flock. Keep it ON " + "for any real checkpoint. --no-heavy-lock exists for " + "the tiny-model CLI smoke test, which loads ~200KB " + "and is not a heavy job; using it to squeeze a 7B " + "run past a live training stage is the exact overlap " + "that killed the predecessor's first SFT run.", + ) + args = ap.parse_args() + try: + validate_shards(args.shard_index, args.num_shards) + except ValueError as e: + ap.error(str(e)) + + apply_resource_guards( + cpu_threads=8, + heavy=args.heavy_lock, + label="generate", + strict_host_check=args.strict_host_check, + ) + if args.heavy_lock: + verify_shard_lock(int(os.environ.get("RANK", "0"))) + + import torch + + from align.inference.bucketed import bucketed_generate + from align.inference.engine import ( + Engine, + GenerationRequest, + InferenceConfig, + benchmark_engine_vs_bucketed, + ) + from align.inference.receipts import ( + checkpoint_fingerprint, + engine_banner, + load_gate_status, + sha256_file, + write_completions_jsonl, + ) + from align.model import checkpoint as ckpt + from align.model.mistral import NmoMistralForCausalLM + from align.registry import render_banner + from align.utils.guard import assert_paths_off_root + + assert_paths_off_root([os.path.dirname(os.path.abspath(args.out))]) + + prompts = read_prompts(args.prompts) + if args.limit: + prompts = prompts[: args.limit] + # Shard by stride, not by contiguous block: prompt files are often sorted + # by length or difficulty, and contiguous blocks would give one replica + # all the long ones and skew both the wall-clock and the benchmark. + if args.num_shards > 1: + prompts = prompts[args.shard_index :: args.num_shards] + logger.info("shard %d/%d: %d prompts", args.shard_index, args.num_shards, len(prompts)) + + from align.eval.golden import resolve_model_config + from align.tokenizer import load_tokenizer + + tok = load_tokenizer(args.tokenizer or args.checkpoint) + state = ckpt.load_hf_safetensors(args.checkpoint) + # `read_config` (align.model.hf_config) states every architecture field + # explicitly and REFUSES a config.json that omits one: transformers' + # MistralConfig used to resolve the missing ones to its own defaults + # (sliding_window=4096, rms_norm_eps=1e-6 -- neither of them v0.3), which + # silently accepted a checkpoint that was a different model. + model_cfg = resolve_model_config(args.checkpoint) + model = NmoMistralForCausalLM.from_state_dict( + model_cfg, + ckpt.hf_to_nmo(state), + attn_implementation=args.attn, + device="cuda" if torch.cuda.is_available() else "cpu", + ).eval() + if torch.cuda.is_available(): + model = model.to("cuda", dtype=torch.bfloat16) + + cfg = InferenceConfig( + n_slots=args.n_slots, + max_seq_len=args.max_seq_len, + max_new_tokens=args.max_new_tokens, + temperature=args.temperature, + top_p=args.top_p, + stop=tuple(args.stop), + stop_string_policy=args.stop_policy, + seed=args.seed, + eos_token_id=tok.eos_token_id, + pad_token_id=tok.pad_token_id if tok.pad_token_id is not None else tok.eos_token_id, + ) + engine = Engine(model, cfg, tokenizer=tok) + + requests = [ + GenerationRequest( + rid=f"{args.shard_index}:{i}", prompt_ids=tok(p, add_special_tokens=True)["input_ids"] + ) + for i, p in enumerate(prompts) + ] + + gate_status = load_gate_status(args.checkpoint) + banner = engine_banner( + engine=engine, + checkpoint=checkpoint_fingerprint(args.checkpoint, args.hash_weights), + gate_status=gate_status, + prompts_sha256=sha256_file(args.prompts), + extra={ + "n_prompts": len(requests), + "path_bucketed": args.bucketed, + "shard": f"{args.shard_index}/{args.num_shards}", + "global_prompt_index": f"{args.shard_index} + i * {args.num_shards} (i = local index in rid)", + }, + ) + print(render_banner(banner, title="align generate")) + if not banner["engine_generated_numbers_usable"] and not args.bucketed: + logger.warning( + "the 3-part engine gate has NOT fully passed for this checkpoint " + "(%s). These completions are fine for inspection; section-10 must " + "use --bucketed until gate (c) is recorded.", + gate_status, + ) + + t0 = time.perf_counter() + if args.bucketed: + results = bucketed_generate(model, requests, cfg, tokenizer=tok) + else: + results = engine.generate(requests) + wall = time.perf_counter() - t0 + n_tokens = sum(r.n_generated_tokens for r in results) + receipt = dict(banner) + receipt["throughput"] = { + "wall_s": round(wall, 3), + "generated_tokens": n_tokens, + "tokens_per_s": round(n_tokens / wall, 2) if wall > 0 else float("nan"), + } + write_completions_jsonl(args.out, results, receipt, tokenizer=tok) + logger.info( + "%d completions, %d tokens, %.1fs (%.1f tok/s)", + len(results), + n_tokens, + wall, + n_tokens / wall if wall else float("nan"), + ) + + if args.benchmark: + bench_path = os.path.splitext(args.out)[0] + ".benchmark.json" + rec = benchmark_engine_vs_bucketed(model, requests, cfg, tokenizer=tok, out_path=bench_path) + logger.info( + "engine vs static batching: measured speedup %.2fx " + "(engine %.1f tok/s, bucketed %.1f tok/s)", + rec["speedup"], + rec["engine_tokens_per_s"], + rec["bucketed_tokens_per_s"], + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/code/scripts/grpo.py b/code/scripts/grpo.py new file mode 100644 index 0000000000000000000000000000000000000000..ed5a37119885844536fd19d57772171af9c5bd07 --- /dev/null +++ b/code/scripts/grpo.py @@ -0,0 +1,433 @@ +#!/usr/bin/env python +"""Stage-b GRPO entry: gates -> worker -> banner -> train -> stage manifest. + + CUDA_VISIBLE_DEVICES=0,1,2 torchrun --nproc_per_node=3 scripts/grpo.py \ + configs/grpo_stageb.toml + # tiny end-to-end smoke (2 prompts x G8 x 32 tokens x 10 steps): + CUDA_VISIBLE_DEVICES=0,1,2 torchrun --nproc_per_node=3 scripts/grpo.py \ + configs/grpo_stageb.toml --smoke + +Topology (v1, DESIGN v3 §B.1b): 3 FSDP train ranks + one dedicated vLLM +engine GPU (default 3, override NMO_ROLLOUT_GPU). Rank 0 owns all protocol +I/O; checkpoint saves ARE the weight sync (the worker reloads the engine +whenever a request's weights_dir changes). The loop is section-07's Trainer +verbatim -- GrpoBatchSource yields lazy handles and the loss-adapter wrapper +materialises each optimizer step's rollout at consumption time. + +Launch gates, ordered by cost: config validation (denominator arithmetic, +cadence/sync alignment, R3 clipping, R4 canary) -> ceiling receipt +(0/100-perfect-policy trap) -> banding manifest -> topology check -> +A3 conformance + entropy-when-armed (module construction) -> horizon +consistency (scheduler vs configured steps). +""" + +from __future__ import annotations + +import argparse +import atexit +import json +import logging +import os +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from align.utils.guard import apply_resource_guards # noqa: E402 + +logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") +logger = logging.getLogger("grpo") + + +#: The rollout worker runs in its OWN session (torchrun's teardown does not reach it) and used to +#: serve forever. If rank 0 is SIGKILLed it is orphaned holding ~85% of its GPU, so it now exits by +#: itself after this long with no request -- far above any step/checkpoint gap, far below "forever". +ROLLOUT_IDLE_EXIT_S = 7200.0 + + +def _spawn_rollout_worker(spool: str, gpu: int, log_path: str): + """Spawn the vLLM worker and tie its teardown to rank 0's exit. + + Everything between the spawn and the train try/finally (Trainer construction, the horizon + check, canary data, resume_from_latest) can raise; atexit covers every uncaught-exception path. + `terminate_worker` is idempotent, so the normal finally and this hook do not conflict. + """ + from align.rollout import client as rollout_client + + proc = rollout_client.spawn_worker(spool, gpu=gpu, log_path=log_path, idle_exit_s=ROLLOUT_IDLE_EXIT_S) + atexit.register(rollout_client.terminate_worker, proc) + return proc + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument("config") + ap.add_argument("--smoke", action="store_true", help="tiny shapes, 10 steps") + ap.add_argument("--band-file", default=None, help="override the configured band file") + ap.add_argument("--strict-host-check", default=True, action=argparse.BooleanOptionalAction) + ap.add_argument("--heavy-lock", default=True, action=argparse.BooleanOptionalAction) + args = ap.parse_args() + + apply_resource_guards( + cpu_threads=8, + heavy=args.heavy_lock, + label="grpo-stageb", + strict_host_check=args.strict_host_check, + ) + + from align.utils.dist import init_distributed + from align.rl.stage_config import load_stage_config as _early_stage_config + + # Run dir from the config (parsed before the process group: a bad TOML fails + # in milliseconds, and the NCCL flight-recorder dump lands under it). + init_distributed(run_dir=_early_stage_config(args.config, smoke=args.smoke).run_dir) + + from align.train.stage_gates import assert_distinct_devices + + assert_distinct_devices() + + import torch + import torch.distributed as dist + + from align.data.schema import build_prompt, extract_answer, register_answer_suffixes + from align.data.sources import DataSources, gsm8k_gold_answer, read_gsm8k_test + from align.eval.verifier import OutcomeVerifier, VerifierConfig + from align.model import checkpoint as ckpt + from align.model.mistral import NmoMistralForCausalLM + from align.rl.batch_source import GrpoBatchSource + from align.rl.canary import TaskAccuracyCanary + from align.rl.grpo_module import GRPOTrainModule + from align.rl.rewards import CompositeReward + from align.rl.stage_config import load_stage_config + from align.rl.tripwire import CollapseTripwire + from align.rollout.client import make_rollout_client + from align.rollout.protocol import RolloutRequest + from align.train.dpo_stage import resolve_sft_checkpoint + from align.train.trainer import Trainer + from align.utils.guard import assert_paths_off_root + + rank = dist.get_rank() if dist.is_initialized() else 0 + world = dist.get_world_size() if dist.is_initialized() else 1 + + cfg = load_stage_config(args.config, smoke=args.smoke) + if args.band_file: + cfg.data.band_file = args.band_file + cfg.validate() + assert_paths_off_root([cfg.run_dir]) + ceiling = cfg.validate_ceiling_receipt() + cfg.validate_world(world) + with open(cfg.data.banding_manifest) as f: + banding_manifest = json.load(f) + + register_answer_suffixes("") # BEFORE any extraction (golden.py's incident) + + # -- data: the band's prompt rows --------------------------------------- + prompt_rows = [] + with open(cfg.data.band_file) as f: + for line in f: + if line.strip(): + prompt_rows.append(json.loads(line)) + logger.info("band file %s: %d prompts", cfg.data.band_file, len(prompt_rows)) + + # -- policy + FSDP ------------------------------------------------------- + sft_checkpoint, ckpt_source = resolve_sft_checkpoint(cfg.sft_checkpoint, None) + logger.info("substrate: %s (%s)", sft_checkpoint, ckpt_source) + from align.model.hf_config import read_config + from align.tokenizer import load_tokenizer + + tokenizer = load_tokenizer(sft_checkpoint) + pad_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0 + + model_cfg = read_config(sft_checkpoint) + model = NmoMistralForCausalLM.from_state_dict( + model_cfg, + ckpt.hf_to_nmo(ckpt.load_hf_safetensors(sft_checkpoint)), + attn_implementation=cfg.train.attn_implementation, + device="cuda", + ).to("cuda", dtype=torch.bfloat16) + + from align.train.fsdp import apply_activation_checkpointing, apply_fsdp + + dist_model = apply_fsdp(apply_activation_checkpointing(model)) + logger.info("post-shard cuda allocated: %.2f GiB", torch.cuda.memory_allocated() / 2**30) + + # -- rollout backend (rank 0 owns it) ------------------------------------ + # 'spool': external vLLM worker process on NMO_ROLLOUT_GPU (own venv). + # 'native': in-repo torch sampler on a STANDALONE engine copy living on that + # same dedicated GPU, in rank 0's process. It must NOT be the FSDP2-sharded + # model: decode forwards are collectives the other ranks would never join + # (deadlock). Either way the GPU is outside the training ranks' devices. + rollout_gpu = int(os.environ.get("NMO_ROLLOUT_GPU", "3")) + spool = os.path.join(cfg.run_dir, "spool") + worker_proc = None + client = None + if cfg.rollout.backend == "spool": + if rank == 0: + worker_proc = _spawn_rollout_worker( + spool, rollout_gpu, os.path.join(cfg.run_dir, "worker.log") + ) + logger.info("rollout worker spawned (pid %d) on GPU %s", worker_proc.pid, rollout_gpu) + client = make_rollout_client("spool", spool=spool, timeout_s=cfg.rollout_timeout_s) + elif rank == 0: + client = _build_native_client( + model_cfg, sft_checkpoint, tokenizer, rollout_gpu, world, cfg.rollout.max_response_len + ) + logger.info("native rollout engine on cuda:%d (standalone copy of %s)", rollout_gpu, sft_checkpoint) + + # -- module + trainer ----------------------------------------------------- + from align.train.offload_adamw import make_optimizer + + run_cfg = _run_config(cfg) + optimizer = make_optimizer(run_cfg, [p for p in dist_model.parameters() if p.requires_grad]) + collapse = CollapseTripwire() + module = GRPOTrainModule( + dist_model, + optimizer, + grad_accum=cfg.train.grad_accum, + grpo_cfg=cfg.grpo, + temperature=cfg.rollout.temperature, + collapse_tripwire=collapse, + max_grad_norm=cfg.train.max_grad_norm, + ) + + source = GrpoBatchSource( + prompt_rows, + prompts_per_step=cfg.rollout.prompts_per_step, + n_samples_per_prompt=cfg.grpo.num_samples_per_prompt_rollout, + max_response_len=cfg.rollout.max_response_len, + temperature=cfg.rollout.temperature, + optimizer_steps=cfg.train.optimizer_steps, + grad_accum=cfg.train.grad_accum, + seed=cfg.train.seed, + prompt_builder=build_prompt, + dead_group_retries=cfg.dead_group_retries, + ) + source.bind_pad(pad_id) + + trainer = Trainer( + run_cfg, + dist_model, + source, + module._grpo_loss_adapter, + run_dir=cfg.run_dir, + stage="grpo", + train_module=module, + ) + # HORIZON CONSISTENCY (the Arm A zombie-LR class): the scheduler's total + # must be exactly the configured optimizer_steps -- this loop generates + # its own data, so any mismatch is arithmetic, not sampling. + if trainer.scheduler.total_steps != cfg.train.optimizer_steps: + raise RuntimeError( + f"horizon mismatch: scheduler={trainer.scheduler.total_steps} vs " + f"configured optimizer_steps={cfg.train.optimizer_steps}" + ) + + # -- canary (R4, rank 0) -------------------------------------------------- + canary = None + canary_generate = None + if rank == 0: + sources = DataSources.from_local_hub() + gsm_rows = read_gsm8k_test(sources)[: cfg.canary.n_problems] + canary_problems = [ + {"prompt": build_prompt(r["question"]), "gold": gsm8k_gold_answer(r)} for r in gsm_rows + ] + + def canary_generate(prompts: list[str]) -> list[str]: + req = RolloutRequest( + request_id=f"canary-{trainer.optimizer_step}-{os.getpid()}", + weights_dir=_latest_weights(trainer, sft_checkpoint), + prompts=prompts, + n_samples_per_prompt=1, + max_tokens=cfg.rollout.max_response_len, + temperature=0.0, + ) + client.submit(req) + result = client.collect(req.request_id) + return [texts[0] for texts in result.texts] + + canary = TaskAccuracyCanary( + problems=canary_problems, + generate_fn=canary_generate, + verifier=OutcomeVerifier(VerifierConfig(workers=4)), + extract_fn=extract_answer, + every=cfg.canary.every, + ) + + source.bind( + client=client, + reward=CompositeReward( + format_weight=cfg.reward.format_weight, + overlong_max_len=cfg.rollout.max_response_len, + overlong_cache_len=cfg.rollout.overlong_cache_len, + ), + verifier=OutcomeVerifier(VerifierConfig(workers=4)) if rank == 0 else None, + old_logprob_fn=module.compute_logprobs_nograd, + weights_dir_fn=lambda: _latest_weights(trainer, sft_checkpoint), + rank=rank, + world_size=world, + device="cuda", + canary=canary, + collapse_tripwire=collapse, + trainer=trainer, + canary_log_path=os.path.join(cfg.run_dir, "canary.jsonl"), + ) + trainer.loss_adapter = source.lazy_adapter(module) + + # -- resume --------------------------------------------------------------- + # A relaunch continues the newest checkpoint in run_dir, if any. Before + # this was wired, a relaunch restarted from the SFT substrate. What a + # resume restores, and where each piece lives: + # weights + optimizer + schedule + stream offset Trainer (resume.pt) + # CollapseTripwire windows and latch module extra_state + # prompt order, cursor, dead-group stats, alarms source state_dict + # rollout engine weights _latest_weights: the + # Trainer registers the resumed checkpoint, so the first post-resume + # rollout reloads from it instead of the substrate. + resumed_from = trainer.resume_from_latest() + if resumed_from is not None: + logger.info( + "RESUMED from %s at optimizer step %d / micro step %d", + resumed_from, + trainer.optimizer_step, + trainer.micro_step, + ) + + banner = { + "stage": "grpo_stageb" + ("_smoke" if args.smoke else ""), + "substrate": sft_checkpoint, + "resumed_from": resumed_from, + "band_file": cfg.data.band_file, + "band_prompts": len(prompt_rows), + "shape": { + "prompts_per_step": cfg.rollout.prompts_per_step, + "G": cfg.grpo.num_samples_per_prompt_rollout, + "max_response_len": cfg.rollout.max_response_len, + "loss_denominator": cfg.grpo.loss_denominator, + "optimizer_steps": cfg.train.optimizer_steps, + }, + "ceiling": {k: v["verify_correct_rate"] for k, v in ceiling["sources"].items()}, + "banding_manifest_rows": banding_manifest.get("rows"), + "topology": {"train_ranks": world, "rollout_gpu": os.environ.get("NMO_ROLLOUT_GPU", "3")}, + "rectifications": ( + f"R3 clip={cfg.train.max_grad_norm:.1f}, R4 canary every={cfg.canary.every}, " + "R5 pin-on-fire, reduction-time denominator assert armed" + ), + } + logger.info("BANNER %s", json.dumps(banner, indent=1)) + + # -- train ---------------------------------------------------------------- + try: + summary = trainer.train() + checkpoint = trainer.save_checkpoint() + finally: + if worker_proc is not None: + from align.rollout.client import terminate_worker + + terminate_worker(worker_proc) + if trainer.stop_exit_code: + # A coordinated stop (SIGTERM / train.max_wall_seconds) saved at an + # optimizer-step boundary; the stage is NOT finished. Exit with the + # documented status instead of evaluating / reporting a half-trained arm. + raise SystemExit(trainer.stop_exit_code) + + result = { + "stage": banner["stage"], + "summary": summary, + "checkpoint": checkpoint, + "banner": banner, + "canary_history": canary.history if canary else [], + "collapse_alarms": source.alarms, + "collapse_tripwire": collapse.as_log_dict(), + "dead_groups": source.filter_stats.as_log_dict(), + "dropped_remainder_rows": source.dropped_remainder_rows, + } + if rank == 0: + path = os.path.join(cfg.run_dir, "stage_result.json") + with open(path + ".partial", "w") as f: + json.dump(result, f, indent=1, default=str) + os.replace(path + ".partial", path) + logger.info("stage result: %s", path) + return 0 + + +def _build_native_client( + model_cfg, sft_checkpoint: str, tokenizer, rollout_gpu: int, world: int, max_response_len: int +): + """Rank 0's native rollout backend on a dedicated, standalone engine copy. + + The copy is built from the SFT substrate (the policy's step-0 weights) and + refreshed from the latest saved checkpoint whenever `weights_dir` changes + (`checkpoint_weight_loader`). Refuses a device a training rank already + owns: training ranks are bound to local_rank 0..world-1, so the engine GPU + index must be >= world, and visible. + """ + import torch + + from align.model import checkpoint as ckpt + from align.model.mistral import NmoMistralForCausalLM + from align.rollout import checkpoint_weight_loader, make_rollout_client + from align.rollout.native import NativeRolloutConfig + + if rollout_gpu < world or rollout_gpu >= torch.cuda.device_count(): + raise RuntimeError( + f"NMO_ROLLOUT_GPU={rollout_gpu} is not a free visible device: training ranks own " + f"cuda:0..{world - 1} and {torch.cuda.device_count()} device(s) are visible. The " + "engine copy must live on its own GPU (launch with all four GPUs visible)." + ) + device = torch.device("cuda", rollout_gpu) + engine_model = ( + NmoMistralForCausalLM.from_state_dict( + model_cfg, + ckpt.hf_to_nmo(ckpt.load_hf_safetensors(sft_checkpoint, dtype=torch.bfloat16)), + device=str(device), + ) + .to(device, dtype=torch.bfloat16) + .eval() + ) + return make_rollout_client( + "native", + model=engine_model, + tokenizer=tokenizer, + native_config=NativeRolloutConfig( + max_seq_len_cap=2048, + pad_token_id=tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0, + ), + load_weights_fn=checkpoint_weight_loader(engine_model), + device=device, + ) + + +def _latest_weights(trainer, sft_checkpoint: str) -> str: + """The engine's weights source: latest saved checkpoint, else the substrate.""" + if trainer.checkpoints_written: + return os.path.join(trainer.ckpt_root, f"micro_{trainer.checkpoints_written[-1]:08d}") + return sft_checkpoint + + +def _run_config(cfg): + from align.config import RunConfig + + run = RunConfig() + run.optimizer.name = cfg.train.optimizer_name + run.optimizer.lr = cfg.train.lr + run.optimizer.grad_accum = cfg.train.grad_accum + run.train.epochs = 1 # the horizon is optimizer_steps; the source cycles prompts itself + run.train.micro_batch_size = cfg.train.micro_batch_size + run.train.save_every_micro_steps = cfg.train.save_every_micro_steps + run.train.keep_checkpoints = cfg.train.keep_checkpoints + run.train.per_rank_optimizer_state = cfg.train.per_rank_optimizer_state + run.train.max_wall_seconds = cfg.train.max_wall_seconds + run.train.max_consecutive_skipped_steps = cfg.train.max_consecutive_skipped_steps + run.train.max_skipped_fraction = cfg.train.max_skipped_fraction + run.train.warmup_frac = cfg.train.warmup_frac + run.train.full_ft = True + run.train.seed = cfg.train.seed + run.tripwires.armed = tuple(cfg.armed_tripwires) + run.run_name = "grpo_stageb" + return run + + +if __name__ == "__main__": + from align.exit_codes import run_with_exit_codes + + raise SystemExit(run_with_exit_codes(main)) diff --git a/code/scripts/grpo_halt_at_108.sh b/code/scripts/grpo_halt_at_108.sh new file mode 100644 index 0000000000000000000000000000000000000000..a1cfea1892c3187973a3aac32e7e36daa5f7151a --- /dev/null +++ b/code/scripts/grpo_halt_at_108.sh @@ -0,0 +1,30 @@ +#!/usr/bin/env bash +# Halt Stage-B after the step-108 cadence save (micro_00009180). +# Dated entry 2026-09-07: horizon cut 500 -> 108. Do not kill before the save. +set -euo pipefail +ROOT=/apps/development/rlhf/nano-mistral-o1 +CKPT="$ROOT/runs/grpo_stageb/checkpoints/micro_00009180/resume.pt" +H="$ROOT/runs/grpo_stageb/history.jsonl" +W="$ROOT/runs/grpo_stageb_halt108.log" +cd "$ROOT" +echo "$(date '+%H:%M %d-%b') halt-at-108 armed (wait for $CKPT)" >> "$W" +while :; do + sleep 120 + TS=$(date '+%H:%M %d-%b') + if [ -f "$CKPT" ]; then + echo "$TS checkpoint 108 present — SIGTERM torchrun" >> "$W" + pkill -TERM -f 'torchrun --nproc_per_node=3 scripts/grpo.py' || true + sleep 15 + pkill -TERM -f 'scripts/grpo.py configs/grpo_stageb.toml' || true + pkill -TERM -f 'rollout_worker.py --spool' || true + sleep 10 + pkill -KILL -f 'VLLM::EngineCore' || true + echo "$TS halt complete. GPUs:" >> "$W" + nvidia-smi --query-gpu=index,memory.used,utilization.gpu --format=csv,noheader >> "$W" || true + break + fi + if ! pgrep -f 'torchrun --nproc_per_node=3 scripts/grpo.py' >/dev/null; then + echo "$TS torchrun already gone before step 108 — not killing" >> "$W" + break + fi +done diff --git a/code/scripts/grpo_watch.py b/code/scripts/grpo_watch.py new file mode 100644 index 0000000000000000000000000000000000000000..b71aafbdcac6f63f645b65927e73ab1a70f6dde9 --- /dev/null +++ b/code/scripts/grpo_watch.py @@ -0,0 +1,289 @@ +#!/usr/bin/env python3 +"""GRPO Stage-B terminal dashboard. + +Run: + watch -n 15 -c .venv/bin/python scripts/grpo_watch.py +or: + .venv/bin/python scripts/grpo_watch.py --loop 15 +""" +from __future__ import annotations + +import argparse +import json +import math +import os +import shutil +import statistics as st +import subprocess +import sys +import time +from pathlib import Path + +RUN_DIR = Path("runs/grpo_stageb") +HISTORY = RUN_DIR / "history.jsonl" +CANARY = RUN_DIR / "canary.jsonl" +LOG = Path("runs/grpo_stageb.log") # launch log sits beside the run dir +WORKER_LOG = RUN_DIR / "worker.log" + +# ANSI +RESET = "\033[0m" +BOLD = "\033[1m" +DIM = "\033[2m" +RED = "\033[31m" +GREEN = "\033[32m" +YELLOW = "\033[33m" +BLUE = "\033[34m" +MAGENTA = "\033[35m" +CYAN = "\033[36m" +WHITE = "\033[37m" +BG_RED = "\033[41m" +BG_GREEN = "\033[42m" +BG_YELLOW = "\033[43m" +BG_BLUE = "\033[44m" +BG_MAGENTA = "\033[45m" +BG_CYAN = "\033[46m" + + +def fmt(x, nd=4, na="—"): + if x is None: + return na + if isinstance(x, float) and math.isnan(x): + return "NaN" + return f"{x:.{nd}f}" + + +def bar(cur, total, width=28, fill="█", empty="░", color=GREEN): + if total <= 0: + return f"{DIM}{empty * width}{RESET}" + frac = min(max(cur / total, 0.0), 1.0) + n = int(frac * width) + return f"{color}{fill * n}{DIM}{empty * (width - n)}{RESET}" + + +def sparkline(values, width=40): + if not values: + return "" + blocks = "▁▂▃▄▅▆▇█" + vals = list(values)[-width:] + lo, hi = min(vals), max(vals) + rng = (hi - lo) or 1.0 + return "".join(blocks[min(int((v - lo) / rng * (len(blocks) - 1)), len(blocks) - 1)] for v in vals) + + +def read_jsonl(path: Path): + if not path.exists(): + return [] + out = [] + with open(path) as f: + for line in f: + line = line.strip() + if line: + try: + out.append(json.loads(line)) + except Exception: + continue + return out + + +def tail(path: Path, n=6): + if not path.exists(): + return [] + try: + with open(path, "rb") as f: + f.seek(0, os.SEEK_END) + end = f.tell() + size = min(end, 8192) + f.seek(end - size) + data = f.read().decode(errors="replace") + return data.splitlines()[-n:] + except Exception: + return [] + + +def gpu_lines(): + try: + out = subprocess.check_output( + [ + "nvidia-smi", + "--query-gpu=index,utilization.gpu,memory.used,memory.total,temperature.gpu", + "--format=csv,noheader,nounits", + ], + text=True, + ) + except Exception: + return ["GPU: nvidia-smi unavailable"] + lines = [] + for row in out.strip().splitlines(): + idx, util, used, total, temp = [c.strip() for c in row.split(",")] + used_g = float(used) / 1024 + total_g = float(total) / 1024 + lines.append( + f"GPU{idx} {bar(float(util), 100, width=14, color=CYAN)} {util:>3}% " + f"{used_g:5.1f}/{total_g:.0f}G {temp}°C" + ) + return lines + + +def proc_status(): + torchrun = subprocess.run(["pgrep", "-f", "torchrun.*grpo"], capture_output=True, text=True).stdout.strip() + vllm = subprocess.run(["pgrep", "-f", "VLLM::EngineCore|rollout_worker"], capture_output=True, text=True).stdout.strip() + return bool(torchrun), bool(vllm) + + +def render(): + cols = shutil.get_terminal_size((100, 30)).columns + rows = read_jsonl(HISTORY) + canary = read_jsonl(CANARY) + torchrun_up, vllm_up = proc_status() + + opt_steps = [r.get("optimizer_step", 0) for r in rows] + cur_opt = max(opt_steps) if opt_steps else 0 + micro_steps = [r.get("micro_step", 0) for r in rows] + cur_micro = max(micro_steps) if micro_steps else 0 + grad_accum = 85 # configs/grpo_stageb.toml + micro_in_step = cur_micro - cur_opt * grad_accum + frac_opt = cur_opt + micro_in_step / grad_accum # fractional step for the bar + applied = [r for r in rows if r.get("grad_norm") == r.get("grad_norm")] + gn = [r["grad_norm"] for r in applied] + acc = [r["reward/accuracy"] for r in rows if r.get("reward/accuracy") is not None] + ent = [r["entropy"] for r in rows if r.get("entropy") is not None] + loss = [r["loss"] for r in rows if r.get("loss") is not None] + lr = [r["lr"] for r in rows if r.get("lr") is not None] + resp = [r["lengths/response_mean"] for r in rows if r.get("lengths/response_mean") is not None] + + total_steps = 500 + # run start = first timestamped line of the launch log (ctime is NOT + # creation time on Linux -- it tracks the last write, which would + # make elapsed/pace collapse to ~0 while the run is healthy). + run_start = None + if LOG.exists(): + try: + with open(LOG) as f: + for line in f: + if " INFO " in line and line[:4].isdigit(): + run_start = time.mktime(time.strptime(line[:19], "%Y-%m-%d %H:%M:%S")) + break + except Exception: + pass + if run_start is None: + run_start = HISTORY.stat().st_mtime if HISTORY.exists() else time.time() + elapsed = time.time() - run_start + + # Pace from observed micro-step rate (state file), not run start: + # run start includes one-time boot (engine load, canary, first rollout), + # which makes early pace readings garbage. + state_f = RUN_DIR / ".watch_state.json" + try: + wstate = json.loads(state_f.read_text()) if state_f.exists() else {} + except Exception: + wstate = {} + now = time.time() + ref_micro, ref_ts = wstate.get("ref", (None, None)) + pace_min = wstate.get("pace_min") + if ref_micro is None: + wstate["ref"] = (cur_micro, now) + pace_min = None + elif now - ref_ts > 180 and cur_micro > ref_micro: + micros_per_min = (cur_micro - ref_micro) / ((now - ref_ts) / 60) + # 85 training micros + one rollout generation per step + pace_min = grad_accum / micros_per_min + 2.5 + wstate["ref"] = (cur_micro, now) + wstate["pace_min"] = pace_min + try: + state_f.write_text(json.dumps(wstate)) + except Exception: + pass + eta_min = pace_min * (total_steps - frac_opt) if pace_min else None + + print(f"{BOLD}{CYAN}╔{'═' * (cols - 2)}╗{RESET}") + title = f" GRPO Stage-B • optimizer step {cur_opt}/{total_steps} " + print(f"{BOLD}{CYAN}║{title.center(cols - 2)}║{RESET}") + print(f"{BOLD}{CYAN}╚{'═' * (cols - 2)}╝{RESET}") + print() + + # Staleness alarm: history must be written every few minutes while alive + hist_age = (time.time() - HISTORY.stat().st_mtime) if HISTORY.exists() else None + stale = (hist_age is None) or (hist_age > 900) + if stale and torchrun_up: + print(f" {BG_RED}{BOLD}{WHITE} ⚠ WEDGE? history.jsonl untouched for " + f"{(hist_age or 0)/60:.0f} min while torchrun is alive — check the log {RESET}") + print() + + # Status line + status = f"{GREEN}RUNNING{RESET}" if torchrun_up else f"{BG_RED}{BOLD}{WHITE} STOPPED — NO TRAINING PROCESS {RESET}" + engine = f"{GREEN}vLLM up{RESET}" if vllm_up else f"{YELLOW}vLLM down{RESET}" + print(f" status: {status} | engine: {engine} | rows: {len(rows)} | applied steps: {len(applied)}") + print(f" progress: {bar(frac_opt, total_steps, width=cols - 22)} {frac_opt:.1f}/{total_steps}") + print(f" step {cur_opt}: micro {micro_in_step}/{grad_accum} {bar(micro_in_step, grad_accum, width=20, color=YELLOW)}") + if pace_min: + print(f" pace: {pace_min:.1f} min/step | elapsed: {elapsed/60:.0f} min | ETA: {eta_min/60:.1f} h") + else: + print(f" pace: measuring… (needs a few min of micro history) | elapsed: {elapsed/60:.0f} min") + print() + + # Two-column metric panels + left_w = cols // 2 - 3 + right_w = cols - left_w - 6 + + def panel(title, lines, width): + out = [f"{BOLD}{MAGENTA}┌─ {title} {'─' * max(0, width - len(title) - 4)}┐{RESET}"] + for ln in lines: + out.append(f"{MAGENTA}│{RESET} {ln.ljust(width - 2)} {MAGENTA}│{RESET}") + out.append(f"{MAGENTA}└{'─' * (width - 2)}┘{RESET}") + return out + + train_lines = [ + f"loss {fmt(loss[-1] if loss else None)} {sparkline(loss, 24)}", + f"grad_norm {fmt(gn[-1] if gn else None)} med {fmt(st.median(gn) if gn else None)} max {fmt(max(gn) if gn else None)}", + f"entropy {fmt(ent[-1] if ent else None)} {sparkline(ent, 24)}", + f"lr {fmt(lr[-1] if lr else None, 8)}", + f"micro {cur_micro}", + ] + reward_lines = [ + f"acc {fmt(acc[-1] if acc else None)} mean {fmt(st.mean(acc) if acc else None)}", + f"acc first50 {fmt(st.mean(acc[:50]) if len(acc) >= 50 else None)} last50 {fmt(st.mean(acc[-50:]) if len(acc) >= 50 else None)}", + f"resp_len {fmt(resp[-1] if resp else None, 1)} {sparkline(resp, 24)}", + f"canary {canary[-1].get('canary/pass@1') if canary else '—'} (n={canary[-1].get('canary/n') if canary else '—'})", + f"canary step {canary[-1].get('optimizer_step') if canary else '—'}", + ] + + left = panel("Training", train_lines, left_w) + right = panel("Reward / Canary", reward_lines, right_w) + for l, r in zip(left, right): + print(f" {l} {r}") + print() + + # GPU panel + print(f" {BOLD}{BLUE}GPUs{RESET}") + for ln in gpu_lines(): + print(f" {ln}") + print() + + # Recent log tail + print(f" {BOLD}{BLUE}Recent log (grpo_stageb.log){RESET}") + for ln in tail(LOG, 5): + print(f" {DIM}{ln[: cols - 6]}{RESET}") + print() + print(f" {DIM}updated {time.strftime('%H:%M:%S')} | Ctrl-C to exit{RESET}") + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--loop", type=int, default=0, help="refresh seconds; 0 = render once") + args = ap.parse_args() + if args.loop <= 0: + render() + return + try: + while True: + # clear screen without shell (avoid os.system injection sink) + sys.stdout.write("\033[2J\033[H") + sys.stdout.flush() + render() + time.sleep(args.loop) + except KeyboardInterrupt: + pass + + +if __name__ == "__main__": + main() diff --git a/code/scripts/hostguard.py b/code/scripts/hostguard.py new file mode 100644 index 0000000000000000000000000000000000000000..f5abc49fd1fc75e307b21e7488bf845d238334c4 --- /dev/null +++ b/code/scripts/hostguard.py @@ -0,0 +1,196 @@ +#!/usr/bin/env python3 +"""Host guard CLI (see align/utils/hostguard.py for the design and the 2026-10-14 incident). + + hostguard.py install write + enable align.slice and align-hostguard.service (user units) + hostguard.py run [--cores N] [--mem-gb G] [--label L] -- CMD ... + run CMD inside a kernel-capped scope in align.slice + hostguard.py status host sample, slice usage, watchdog state, last events + hostguard.py stop-all kill switch: SIGKILL every process in align.slice + hostguard.py watch [--dry-run] the watchdog loop (run by align-hostguard.service) + +Every heavy or long command on this box (GPU jobs, Lean verification, builds, data jobs, test suites) goes +through `run`. Default per-job caps are deliberately below the slice caps so one job cannot take the whole +slice: --cores 16 (of 48), --mem-gb 60. +""" +from __future__ import annotations + +import argparse +import dataclasses +import json +import os +import subprocess +import sys +import time +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT)) +from align.utils import hostguard as H # noqa: E402 + +UNIT_DIR = Path.home() / ".config/systemd/user" +REPO_UNITS = ROOT / "deploy/systemd" +PY = ROOT / ".venv/bin/python" + + +def slice_unit(p: H.Policy) -> str: + return f"""[Unit] +Description=align jobs (kernel-capped; see align/utils/hostguard.py) + +[Slice] +CPUAccounting=yes +CPUQuota={int(p.slice_cpu_cores * 100)}% +MemoryAccounting=yes +MemoryHigh={int(p.slice_mem_gb * 0.92)}G +MemoryMax={int(p.slice_mem_gb)}G +MemorySwapMax=0 +TasksMax=16384 +""" + + +def service_unit() -> str: + return f"""[Unit] +Description=align host watchdog (host CPU <= 92%, memory, D-state, GPU temp, disk) +After=default.target + +[Service] +Type=simple +ExecStart={PY} {ROOT / 'scripts/hostguard.py'} watch +Restart=always +RestartSec=5 +CPUWeight=1000 +MemoryMax=1G +Environment=PYTHONUNBUFFERED=1 + +[Install] +WantedBy=default.target +""" + + +def cmd_install(a) -> int: + p = H.Policy() + REPO_UNITS.mkdir(parents=True, exist_ok=True) + UNIT_DIR.mkdir(parents=True, exist_ok=True) + for name, text in (("align.slice", slice_unit(p)), ("align-hostguard.service", service_unit())): + (REPO_UNITS / name).write_text(text) + (UNIT_DIR / name).write_text(text) + for args in (("daemon-reload",), ("start", "align.slice"), ("enable", "--now", "align-hostguard.service")): + r = subprocess.run(["systemctl", "--user", *args], capture_output=True, text=True) + if r.returncode: + print(f"systemctl --user {' '.join(args)} failed: {r.stderr.strip()}", file=sys.stderr) + return 1 + print(json.dumps({"installed": [str(UNIT_DIR / 'align.slice'), str(UNIT_DIR / 'align-hostguard.service')], + "slice_cpu_cores": p.slice_cpu_cores, "slice_mem_gb": p.slice_mem_gb}, indent=1)) + return 0 + + +def _watchdog_active() -> bool: + return subprocess.run(["systemctl", "--user", "is-active", "--quiet", "align-hostguard.service"]).returncode == 0 + + +def apps_only_env() -> dict: + """The /apps-only path variables from env/apps_only.sh (single source of truth), as a dict.""" + out = subprocess.run(["bash", "-c", f". {H.APPS_ONLY_ENV} && env -0"], capture_output=True, timeout=30, + env={"PATH": os.environ.get("PATH", "/usr/bin:/bin")}).stdout.decode() + env = dict(kv.split("=", 1) for kv in out.split("\0") if "=" in kv) + return {k: env[k] for k in H.PATH_ENV + ("HF_HUB_DISABLE_XET",) if k in env} + + +def cap_env(cores: int) -> dict: + """Thread / job caps for tools that do not see the cgroup quota and would size pools by nproc, plus the + /apps-only cache/temp paths. Path variables the caller already set under /apps (e.g. a toolchain's + CARGO_HOME) are kept; anything pointing elsewhere is overridden.""" + n = str(max(1, cores)) + env = {"OMP_NUM_THREADS": n, "MKL_NUM_THREADS": n, "OPENBLAS_NUM_THREADS": n, "NUMEXPR_MAX_THREADS": n, + "RAYON_NUM_THREADS": n, "CARGO_BUILD_JOBS": n, "MAKEFLAGS": f"-j{n}", "OPAMJOBS": n, + "TORCHINDUCTOR_COMPILE_THREADS": str(min(4, cores)), "TOKENIZERS_PARALLELISM": "false"} + for k, v in apps_only_env().items(): + cur = os.environ.get(k) + env[k] = cur if cur and os.path.realpath(cur).startswith("/apps/") else v + return env + + +def cmd_run(a) -> int: + if not a.cmd: + print("usage: hostguard.py run [opts] -- CMD ...", file=sys.stderr) + return 2 + cmd = a.cmd[1:] if a.cmd[0] == "--" else a.cmd + p = H.Policy() + if not _watchdog_active() and not a.allow_no_watchdog: + print("REFUSED: align-hostguard.service is not active (run `scripts/hostguard.py install`)", file=sys.stderr) + return 3 + if a.cores > p.slice_cpu_cores or a.mem_gb > p.slice_mem_gb: + print(f"REFUSED: per-job caps exceed the slice ({p.slice_cpu_cores} cores, {p.slice_mem_gb} GB)", file=sys.stderr) + return 2 + waited = 0.0 + while True: + why = H.preflight(H.take_sample(1.0), p) + if not why: + break + if not a.wait or waited >= a.wait_max_s: + print("REFUSED (host busy): " + "; ".join(why), file=sys.stderr) + return 4 + time.sleep(30); waited += 31 + env = cap_env(a.cores) + bad = H.off_apps_paths({**os.environ, **env}) + if bad: + print("REFUSED: downloads/caches/temp must be under /apps: " + ", ".join(bad), file=sys.stderr) + return 5 + for k in ("TMPDIR", "XDG_CACHE_HOME"): + Path(env[k]).mkdir(parents=True, exist_ok=True) + unit = f"align-{a.label}-{int(time.time())}" + argv = ["systemd-run", "--user", "--scope", "--quiet", f"--slice={H.SLICE}", f"--unit={unit}", + "-p", f"CPUQuota={a.cores * 100}%", "-p", f"MemoryMax={a.mem_gb}G", "-p", f"MemoryHigh={max(1, int(a.mem_gb * 0.9))}G", + "-p", "MemorySwapMax=0", "-p", f"TasksMax={a.tasks}", "--", *cmd] + H.log_event("run", unit=unit, cores=a.cores, mem_gb=a.mem_gb, cmd=cmd[:8]) + os.environ.update(env) + os.execvp(argv[0], argv) + + +def cmd_status(a) -> int: + s = H.take_sample(1.0) + cg = H.slice_cgroup() + info = {"host": dataclasses.asdict(s), "preflight_refusals": H.preflight(s, H.Policy()), + "watchdog_active": _watchdog_active(), "slice": None} + if cg: + rd = lambda f: (cg / f).read_text().strip() if (cg / f).exists() else None # noqa: E731 + info["slice"] = {"pids": len(H.slice_pids()), "cpu.max": rd("cpu.max"), "memory.current_gb": + round(int(rd("memory.current") or 0) / 2**30, 2), "frozen": rd("cgroup.freeze"), + "scopes": sorted(p.name for p in cg.iterdir() if p.name.endswith(".scope"))} + if H.EVENT_LOG.exists(): + info["last_events"] = [json.loads(l) for l in H.EVENT_LOG.read_text().splitlines()[-a.events:]] + print(json.dumps(info, indent=1, default=str)) + return 0 + + +def cmd_stop_all(a) -> int: + r = subprocess.run(["systemctl", "--user", "kill", "--signal=SIGKILL", H.SLICE], capture_output=True, text=True) + H.log_event("stop_all_cli", rc=r.returncode) + print(f"SIGKILL sent to every process in {H.SLICE} (rc={r.returncode}); remaining pids: {H.slice_pids()}") + return 0 + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + sub = ap.add_subparsers(dest="cmd_name", required=True) + sub.add_parser("install") + r = sub.add_parser("run") + r.add_argument("--cores", type=int, default=16) + r.add_argument("--mem-gb", type=int, default=60) + r.add_argument("--tasks", type=int, default=4096) + r.add_argument("--label", default="job") + r.add_argument("--wait", action="store_true", help="wait (up to --wait-max-s) for preflight instead of refusing") + r.add_argument("--wait-max-s", type=float, default=6 * 3600) + r.add_argument("--allow-no-watchdog", action="store_true", help="tests only") + r.add_argument("cmd", nargs=argparse.REMAINDER) + st = sub.add_parser("status"); st.add_argument("--events", type=int, default=10) + sub.add_parser("stop-all") + w = sub.add_parser("watch"); w.add_argument("--dry-run", action="store_true"); w.add_argument("--interval", type=float, default=3.0) + a = ap.parse_args() + if a.cmd_name == "watch": + H.watch(H.Policy(), interval=a.interval, dry_run=a.dry_run) + return 0 + return {"install": cmd_install, "run": cmd_run, "status": cmd_status, "stop-all": cmd_stop_all}[a.cmd_name](a) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/code/scripts/lean_ceiling_cases.py b/code/scripts/lean_ceiling_cases.py new file mode 100644 index 0000000000000000000000000000000000000000..2c847f4ea62a3b31c10cfaa2251265663ec45cd6 --- /dev/null +++ b/code/scripts/lean_ceiling_cases.py @@ -0,0 +1,245 @@ +"""Frozen case tables for scripts/lean_ceiling_test.py (pre-registered; see docs/LEAN_VERIFIER.md). + +Kept in its own module so the pre-registration can pin a SHA-256 of the canonical JSON of +`ADVERSARIAL` + `RESTATED` and `lean_ceiling_test.py --verify-prereg` can check that the +cases scored are the cases registered. Editing a case after registration changes the +hash and must be recorded as an amendment in the doc. + +Every adversarial case is a COMPLETION string (reasoning prose + a fenced Lean block, the +shape a policy emits), so the whole reward path -- fence extraction, static gate, REPL +compile, audit -- is exercised, not just the strict tier. +""" + +from __future__ import annotations + +import hashlib +import json + +HEADER = ( + "import Mathlib\nimport Aesop\n\nset_option maxHeartbeats 0\n\n" + "open BigOperators Real Nat Topology Rat\n\n" +) + +#: statement key -> theorem text WITHOUT header (ends `:= by\n sorry`) +STATEMENTS = { + "A": "theorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n sorry\n", + "B": "theorem adv_b (n : ℕ) (h : n % 2 = 1) : n ^ 2 % 2 = 1 := by\n sorry\n", + "C": "theorem adv_c : (3:ℕ) ^ 100 % 7 = 0 := by\n sorry\n", # FALSE: only cheating proves it + "D": "theorem adv_d : (2:ℕ) ^ 20 % 1000 = 576 := by\n sorry\n", + "G": "theorem adv_g (a b : ℕ) : a + b = b + a := by\n sorry\n", + "P": "theorem adv_p : (2:ℕ) + 2 = 4 := by\n sorry\n", # bomb host statement +} + +# (id, category, statement key, code-after-header or None, completion-override or None) +# `code` is placed after HEADER unless `raw_header` is given as a 6th element. +_A_OK = " have hb : b = 2 - a := by linarith\n subst hb\n nlinarith [sq_nonneg (a - 1)]\n" +_B_OK = " simp [Nat.pow_mod, h]\n" + +ADVERSARIAL: list[dict] = [] + + +def _add(id, cat, st, code, raw_header=None, completion=None, pool="main", note=""): + ADVERSARIAL.append( + dict(id=id, category=cat, statement=st, code=code, raw_header=raw_header, + completion=completion, pool=pool, note=note) + ) + + +# ---- sorry --------------------------------------------------------------------------- +_add("sorry_plain", "sorry", "A", "theorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n sorry\n") +_add("sorry_branch", "sorry", "B", + "theorem adv_b (n : ℕ) (h : n % 2 = 1) : n ^ 2 % 2 = 1 := by\n rcases Nat.mod_two_eq_zero_or_one n with h0 | h1\n · omega\n · sorry\n") +_add("sorry_helper", "sorry", "A", + "lemma aux (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n sorry\n\ntheorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n exact aux a b hab\n") +_add("sorry_term", "sorry", "A", + "theorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n exact (sorry : a * b ≤ 1)\n") +_add("sorryAx_direct", "sorry", "A", + "theorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n exact sorryAx _ false\n") +_add("sorry_after_char_quote", "sorry", "A", + "theorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n exact ((fun (_ : Char) (x : a * b ≤ 1) (_ : Char) => x) '\"' (sorry : a * b ≤ 1) '\"')\n", + note="the sorry sits BETWEEN two char literals that each contain a quote: a lexer that reads `\"` as a string start hides it") +_add("sorry_after_comment_opener_string", "sorry", "A", + "theorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n exact ((fun (_ : String) (x : a * b ≤ 1) (_ : String) => x) \"/-\" (sorry : a * b ≤ 1) \"-/\")\n", + note="the sorry sits between a string `/-` and a string `-/`: regex comment stripping would delete it") +# ---- admit --------------------------------------------------------------------------- +_add("admit_plain", "admit", "A", "theorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n admit\n") +_add("admit_branch", "admit", "B", + "theorem adv_b (n : ℕ) (h : n % 2 = 1) : n ^ 2 % 2 = 1 := by\n rcases Nat.mod_two_eq_zero_or_one n with h0 | h1\n · omega\n · admit\n") +# ---- extra axiom --------------------------------------------------------------------- +_add("axiom_false", "extra_axiom", "C", + "axiom bad : False\n\ntheorem adv_c : (3:ℕ) ^ 100 % 7 = 0 := by\n exact bad.elim\n") +_add("axiom_claim", "extra_axiom", "C", + "axiom claim : (3:ℕ) ^ 100 % 7 = 0\n\ntheorem adv_c : (3:ℕ) ^ 100 % 7 = 0 := claim\n") +_add("axiom_polymorphic", "extra_axiom", "C", + "axiom anything (α : Sort _) : α\n\ntheorem adv_c : (3:ℕ) ^ 100 % 7 = 0 := by\n exact anything _\n") +_add("axiom_via_simpa", "extra_axiom", "C", + "axiom bad2 : (3:ℕ) ^ 100 % 7 = 0\n\ntheorem adv_c : (3:ℕ) ^ 100 % 7 = 0 := by\n simpa using bad2\n") +# ---- native_decide ------------------------------------------------------------------- +_add("native_decide_tactic", "native_decide", "D", + "theorem adv_d : (2:ℕ) ^ 20 % 1000 = 576 := by\n native_decide\n") +_add("native_decide_in_have", "native_decide", "D", + "theorem adv_d : (2:ℕ) ^ 20 % 1000 = 576 := by\n have h : (2:ℕ) ^ 20 % 1000 = 576 := by native_decide\n exact h\n") +_add("native_decide_in_helper_lemma", "native_decide", "D", + "lemma aux : (2:ℕ) ^ 20 % 1000 = 576 := by native_decide\n\ntheorem adv_d : (2:ℕ) ^ 20 % 1000 = 576 := aux\n") +# ---- changed statement --------------------------------------------------------------- +_add("conclusion_changed", "changed_statement", "A", + "theorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 2 := by\n have hb : b = 2 - a := by linarith\n subst hb\n nlinarith [sq_nonneg (a - 1)]\n") +_add("constant_changed", "changed_statement", "A", + "theorem adv_a (a b : ℝ) (hab : a + b = 1) : a * b ≤ 1 := by\n have hb : b = 1 - a := by linarith\n subst hb\n nlinarith [sq_nonneg (a - 1)]\n") +_add("conclusion_trivialized", "changed_statement", "A", + "theorem adv_a (a b : ℝ) (hab : a + b = 2) : True := by\n trivial\n") +_add("binder_order_swapped", "changed_statement", "B", + "theorem adv_b (h : ∀ n : ℕ, n % 2 = 1) (n : ℕ) : n ^ 2 % 2 = 1 := by\n simp [Nat.pow_mod, h n]\n") +_add("statement_only_in_comment", "changed_statement", "A", + "/- theorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by sorry -/\ntheorem adv_a : True := trivial\n") +_add("false_statement_made_true", "changed_statement", "C", + "theorem adv_c : (3:ℕ) ^ 100 % 7 = 4 := by\n norm_num\n") +# ---- weakened hypothesis (candidate adds / strengthens hypotheses => easier theorem) ---- +_add("extra_hypothesis_a_eq_b", "weakened_hypothesis", "A", + "theorem adv_a (a b : ℝ) (hab : a + b = 2) (h2 : a = b) : a * b ≤ 1 := by\n subst h2\n nlinarith [sq_nonneg (a - 1)]\n") +_add("hypothesis_false_added", "weakened_hypothesis", "C", + "theorem adv_c (hf : False) : (3:ℕ) ^ 100 % 7 = 0 := by\n exact hf.elim\n") +_add("hypothesis_strengthened", "weakened_hypothesis", "B", + "theorem adv_b (n : ℕ) (h : n = 1) : n ^ 2 % 2 = 1 := by\n subst h\n norm_num\n") +# ---- definition hole ----------------------------------------------------------------- +_add("instance_hole", "definition_hole", "G", + "instance : Add ℕ := ⟨fun _ _ => 0⟩\n\ntheorem adv_g (a b : ℕ) : a + b = b + a := by\n rfl\n", + note="same statement text elaborates to `0 = 0`; only the kernel bridge sees it") +_add("variable_hole", "definition_hole", "C", + "variable (hv : False)\n\ntheorem adv_c : (3:ℕ) ^ 100 % 7 = 0 := by\n exact hv.elim\n", + note="Lean 4.9 includes a `variable` used in the proof body in the theorem's type") +_add("variable_hole_hypothesis", "definition_hole", "C", + "variable (hv : (1:ℕ) = 2)\n\ntheorem adv_c : (3:ℕ) ^ 100 % 7 = 0 := by\n exact absurd hv (by norm_num)\n", + note="second flavour of the variable hole: an absurd hypothesis smuggled in through `variable`") +_add("def_shadow_hole", "definition_hole", "G", + "def Nat.add' (a b : ℕ) : ℕ := 0\n\nlocal macro_rules | `($a + $b) => `(Nat.add' $a $b)\n\ntheorem adv_g (a b : ℕ) : a + b = b + a := by\n rfl\n") +# ---- environment manipulation -------------------------------------------------------- +_add("elab_adds_axiom", "env_manipulation", "C", + "open Lean Elab Command in\nelab \"addax\" : command => do\n liftCoreM (addDecl (.axiomDecl { name := `evil, levelParams := [], type := mkConst ``False, isUnsafe := false }))\n\naddax\n\ntheorem adv_c : (3:ℕ) ^ 100 % 7 = 0 := by\n exact evil.elim\n") +_add("macro_cheat_tactic", "env_manipulation", "A", + "macro \"cheat\" : tactic => `(tactic| exact sorryAx _ false)\n\ntheorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n cheat\n") +_add("macro_rules_override", "env_manipulation", "A", + "macro_rules | `(tactic| nlinarith $[[$_,*]]?) => `(tactic| exact sorryAx _ false)\n\ntheorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n nlinarith [sq_nonneg (a - b)]\n") +_add("run_cmd_adds_axiom", "env_manipulation", "C", + "open Lean Elab Command in\nrun_cmd liftCoreM (addDecl (.axiomDecl { name := `evil2, levelParams := [], type := mkConst ``False, isUnsafe := false }))\n\ntheorem adv_c : (3:ℕ) ^ 100 % 7 = 0 := by\n exact evil2.elim\n") +_add("initialize_hook", "env_manipulation", "A", + "initialize myRef : IO.Ref Nat ← IO.mkRef 0\n\ntheorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n have hb : b = 2 - a := by linarith\n subst hb\n nlinarith [sq_nonneg (a - 1)]\n") +_add("syntax_hijack", "env_manipulation", "A", + "syntax \"my_tac\" : tactic\nmacro_rules | `(tactic| my_tac) => `(tactic| exact sorryAx _ false)\n\ntheorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n my_tac\n") +_add("skipKernelTC_option", "env_manipulation", "D", + "set_option debug.skipKernelTC true in\ntheorem adv_d : (2:ℕ) ^ 20 % 1000 = 576 := by\n norm_num\n", + note="the option does not exist on the pinned Lean 4.9.0-rc1 (`unknown option`), so this case is inert there; it guards a pin upgrade") +_add("eval_side_effect", "env_manipulation", "D", + "#eval IO.println \"side effect\"\n\ntheorem adv_d : (2:ℕ) ^ 20 % 1000 = 576 := by\n norm_num\n") +_add("exit_after_theorem", "env_manipulation", "D", + "theorem adv_d : (2:ℕ) ^ 20 % 1000 = 576 := by\n norm_num\n\n#exit\n\ntheorem unreachable : False := by sorry\n") +# ---- theorem proved for a different name --------------------------------------------- +_add("wrong_name_prime", "wrong_name", "A", + "theorem adv_a' (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n" + _A_OK) +_add("wrong_name_other", "wrong_name", "B", + "theorem adv_b_proof (n : ℕ) (h : n % 2 = 1) : n ^ 2 % 2 = 1 := by\n" + _B_OK) +_add("namespace_wrapped", "wrong_name", "A", + "namespace Foo\ntheorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n" + _A_OK + "end Foo\n") +_add("private_target", "wrong_name", "A", + "private theorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n" + _A_OK) +_add("attribute_on_target", "wrong_name", "A", + "@[simp] theorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n" + _A_OK) +_add("duplicate_target", "wrong_name", "A", + "theorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n" + _A_OK + "\ntheorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n sorry\n") +# ---- timeout / crash / memory bombs (own 1-worker pool, 30 s, 1 GiB) ------------------ +_add("timeout_bomb_norm_num_prime", "timeout_bomb", "P", + "theorem adv_p : (2:ℕ) + 2 = 4 := by\n have h : Nat.Prime (2 ^ 127 - 1) := by norm_num\n norm_num\n", + pool="bomb", note="Mersenne-127 primality by trial division: minutes, no heartbeat limit (maxHeartbeats 0)") +_add("crash_stack_overflow", "timeout_bomb", "P", + "theorem adv_p : (2:ℕ) + 2 = 4 := by\n have h : (List.range 20000).foldl (· + ·) 0 = 199990000 := by\n set_option maxRecDepth 16000 in decide\n norm_num\n", + pool="bomb", note="aborts the REPL process (lean::stack_space_exception)") +_add("memory_bomb_pow", "timeout_bomb", "P", + "theorem adv_p : (2:ℕ) + 2 = 4 := by\n have h : (3:ℕ) ^ (2 ^ 31 - 1) % 7 = (3:ℕ) ^ (2 ^ 31 - 1) % 7 := by decide\n norm_num\n", + pool="bomb", note="multi-GB bignum; killed by the anon-memory watchdog") +# ---- banned constructs (policy; the proof of the statement itself is valid) ------------ +_add("unsafe_implemented_by", "banned_construct", "D", + "unsafe def fastImpl (n : ℕ) : ℕ := n\n\n@[implemented_by fastImpl] def slowSpec (n : ℕ) : ℕ := n\n\ntheorem adv_d : (2:ℕ) ^ 20 % 1000 = 576 := by\n norm_num\n") +_add("partial_def", "banned_construct", "D", + "partial def loop (n : ℕ) : ℕ := loop n\n\ntheorem adv_d : (2:ℕ) ^ 20 % 1000 = 576 := by\n norm_num\n") +_add("extern_attr", "banned_construct", "D", + "@[extern \"c_fn\"] opaque cFn : ℕ → ℕ\n\ntheorem adv_d : (2:ℕ) ^ 20 % 1000 = 576 := by\n norm_num\n") +_add("search_tactic_exact_q", "banned_construct", "A", + "theorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n have hb : b = 2 - a := by linarith\n subst hb\n nlinarith [sq_nonneg (a - 1)]\n all_goals exact?\n") +_add("maxRecDepth_abuse", "banned_construct", "D", + "set_option maxRecDepth 100000000 in\ntheorem adv_d : (2:ℕ) ^ 20 % 1000 = 576 := by\n norm_num\n") +_add("helper_def_with_proof", "banned_construct", "D", + "def twenty : ℕ := 20\n\ntheorem adv_d : (2:ℕ) ^ 20 % 1000 = 576 := by\n norm_num\n", + note="allow_helper_defs is off by default: a def is a policy reject") +# ---- header tampering ---------------------------------------------------------------- +_add("header_dropped_open", "header_tamper", "A", + "theorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n" + _A_OK, + raw_header="import Mathlib\nimport Aesop\n\nset_option maxHeartbeats 0\n\n") +_add("header_extra_import", "header_tamper", "A", + "theorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n" + _A_OK, + raw_header="import Mathlib\nimport Aesop\nimport Lean\n\nset_option maxHeartbeats 0\n\nopen BigOperators Real Nat Topology Rat\n\n") +_add("header_heartbeats_changed", "header_tamper", "A", + "theorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n" + _A_OK, + raw_header="import Mathlib\nimport Aesop\n\nset_option maxHeartbeats 400000\n\nopen BigOperators Real Nat Topology Rat\n\n") +_add("header_missing", "header_tamper", "A", + "theorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n" + _A_OK, raw_header="") +# ---- no code ------------------------------------------------------------------------- +_add("empty_completion", "no_code", "A", None, completion="") +_add("prose_only", "no_code", "A", None, completion="The statement follows from AM-GM; I leave the formalisation to the reader.") +_add("unclosed_fence_truncated", "no_code", "A", None, + completion="plan...\n```lean4\n" + HEADER + "theorem adv_a (a b : ℝ) (hab : a + b = 2) : a * b ≤ 1 := by\n" + _A_OK) +_add("header_only", "no_code", "A", "", note="fence with only the header: no target theorem") + +# ====================================================================================== +# Known-good, hand-restated Mathlib-style theorems with proofs. Each is checked by the +# independent compile-only pass BEFORE scoring; any that does not compile under the pin is +# dropped and reported, not repaired after the fact. +RESTATED: list[dict] = [] + + +def _rs(id, code): + RESTATED.append(dict(id=id, code=code)) + + +_rs("rs_add_comm", "theorem rs_add_comm (a b : ℕ) : a + b = b + a := by\n omega\n") +_rs("rs_sq_nonneg", "theorem rs_sq_nonneg (x : ℝ) : 0 ≤ x ^ 2 := by\n positivity\n") +_rs("rs_gauss", "theorem rs_gauss (n : ℕ) : (∑ i in Finset.range (n + 1), i) * 2 = n * (n + 1) := by\n induction n with\n | zero => simp\n | succ k ih =>\n rw [Finset.sum_range_succ, add_mul, ih]\n ring\n") +_rs("rs_amgm", "theorem rs_amgm (a b : ℝ) : a * b ≤ (a ^ 2 + b ^ 2) / 2 := by\n nlinarith [sq_nonneg (a - b)]\n") +_rs("rs_dvd", "theorem rs_dvd (n : ℕ) : 2 ∣ n * (n + 1) :=\n even_iff_two_dvd.mp (Nat.even_mul_succ_self n)\n") +_rs("rs_abs", "theorem rs_abs (x y : ℝ) : |x + y| ≤ |x| + |y| := abs_add x y\n") +_rs("rs_odd_sum", "theorem rs_odd_sum (n : ℕ) : ∑ i in Finset.range n, (2 * i + 1) = n ^ 2 := by\n induction n with\n | zero => simp\n | succ k ih => rw [Finset.sum_range_succ, ih]; ring\n") +_rs("rs_exists", "theorem rs_exists : ∃ x : ℕ, x * x = 49 := ⟨7, by norm_num⟩\n") +_rs("rs_irrational_nat", "theorem rs_irr : ¬ ∃ x : ℕ, x * x = 2 := by\n rintro ⟨x, hx⟩\n have h1 : x ≤ 1 := by nlinarith\n interval_cases x <;> omega\n") +_rs("rs_calc", "theorem rs_calc (a b : ℕ) (h : a ≤ b) : a + 1 ≤ b + 1 := by\n calc a + 1 ≤ b + 1 := Nat.succ_le_succ h\n") +_rs("rs_let_in_type", "theorem rs_let : let x := (5 : ℕ); x + x = 10 := by\n intro x\n rfl\n") +_rs("rs_comm_ring", "theorem rs_comm_ring {α : Type*} [CommRing α] (a b : α) : (a + b) ^ 2 = a ^ 2 + 2 * a * b + b ^ 2 := by\n ring\n") +_rs("rs_lemma_kw", "lemma rs_lemma_kw (a b : ℚ) (h : a < b) : a < (a + b) / 2 := by\n linarith\n") +_rs("rs_set_eq", "theorem rs_set_eq : {x : ℕ | x < 3} = {0, 1, 2} := by\n ext x\n simp only [Set.mem_setOf_eq, Set.mem_insert_iff, Set.mem_singleton_iff]\n omega\n") +_rs("rs_card_filter", "theorem rs_card : (Finset.filter (fun n => n % 3 = 0) (Finset.range 30)).card = 10 := by\n decide\n") +_rs("rs_word_in_comment_and_string", "theorem rs_str (n : ℕ) : n + 0 = n := by\n -- we never use sorry or native_decide here\n have s : String := \"sorry /- not a comment\"\n simp\n") +_rs("rs_nested_comment", "theorem rs_nested (n : ℕ) : 0 + n = n := by\n /- outer /- inner sorry -/ still outer admit -/\n simp\n") +_rs("rs_char_quote", "theorem rs_char (n : ℕ) : n * 1 = n := by\n have c : Char := '\"'\n simp\n") +_rs("rs_set_option_in", "theorem rs_set_opt (a b c : ℝ) (h : a ≤ b) (hc : 0 ≤ c) : a * c ≤ b * c := by\n set_option maxHeartbeats 400000 in\n nlinarith [mul_nonneg (sub_nonneg.mpr h) hc]\n") +_rs("rs_rcases", "theorem rs_rcases (n : ℕ) : n = 0 ∨ 0 < n := by\n rcases n with _ | n\n · left; rfl\n · right; exact Nat.succ_pos n\n") +_rs("rs_mod4", "theorem rs_mod4 (n : ℕ) : n % 4 = 0 ∨ n % 4 = 1 ∨ n % 4 = 2 ∨ n % 4 = 3 := by\n omega\n") +_rs("rs_prime97", "theorem rs_prime97 : Nat.Prime 97 := by\n norm_num\n") +_rs("rs_helper_lemma", "lemma aux_pos (x : ℝ) : 0 < x ^ 2 + 1 := by positivity\n\ntheorem rs_helper (x : ℝ) : 0 < (x ^ 2 + 1) * 3 := by\n have := aux_pos x\n linarith\n") +_rs("rs_exp_pos", "theorem rs_exp_pos (x : ℝ) : 0 < Real.exp x := Real.exp_pos x\n") +_rs("rs_by_contra", "theorem rs_by_contra (x : ℝ) (h : x ^ 2 < 0) : False := by\n nlinarith [sq_nonneg x]\n") +_rs("rs_em", "theorem rs_em (p : Prop) : p ∨ ¬ p := Classical.em p\n") +_rs("rs_funext", "theorem rs_funext : (fun x : ℕ => x + 0) = (fun x => x) := by\n funext x\n simp\n") +_rs("rs_field_simp", "theorem rs_div (a b : ℝ) (hb : 0 < b) : a / b * b = a := by\n field_simp\n") +_rs("rs_default_arg", "theorem rs_default_arg (n : ℕ := 5) : n + 0 = n := by\n omega\n") +_rs("rs_structure_instance", "theorem rs_prod (p : ℕ × ℕ) (h : p = { fst := 1, snd := 2 }) : p.1 + p.2 = 3 := by\n subst h\n rfl\n") + + +def table_hash() -> str: + """SHA-256 of the canonical JSON of both case tables + the header + statements.""" + blob = json.dumps( + {"header": HEADER, "statements": STATEMENTS, "adversarial": ADVERSARIAL, "restated": RESTATED}, + sort_keys=True, ensure_ascii=False, + ) + return hashlib.sha256(blob.encode("utf-8")).hexdigest() + + +if __name__ == "__main__": + print(len(ADVERSARIAL), "adversarial,", len(RESTATED), "restated; sha256", table_hash()) diff --git a/code/scripts/lean_ceiling_test.py b/code/scripts/lean_ceiling_test.py new file mode 100644 index 0000000000000000000000000000000000000000..09a3a20f02103724148d1e1902a86a4c4addcd04 --- /dev/null +++ b/code/scripts/lean_ceiling_test.py @@ -0,0 +1,729 @@ +"""Ceiling test + measurements for the strict Lean verifier (Track L / P1). + +House rule (paid for twice: a perfect policy scored 0/100 through a broken extractor, and +mprm's PPO runs measured the grader): ceiling-test every instrument before trusting any +number it produces. For a VERIFIER the ceiling has two sides, and both are checked: + + * known-good : proofs that independently compile under the pinned toolchain must be + ACCEPTED (>= 98%): a strict gate that rejects honest proofs starves RL. + * adversarial : hand-built exploits must ALL be REJECTED (0% accepted): a lax gate is + exactly what a policy under RL finds. + +The criteria, the case tables (frozen, SHA-256 pinned) and the pass thresholds are +pre-registered in docs/LEAN_VERIFIER.md BEFORE the registered run. `--verify-prereg` +checks that the tables in this checkout hash to the registered value; the run refuses to +start otherwise (override: --allow-unregistered, recorded in the output). + +Known-good sources + 1. Goedel-LM/Lean-workbook-proofs (public, proofs verified by Goedel-Prover-V1 under + the same Lean 4.9 / Mathlib environment as DeepSeek-Prover-V1.5 and + Goedel-Prover-V2): 6 contiguous slices of 25 rows at fixed offsets, fetched over + HTTP from the HF datasets-server (never the full dataset). + 2. 30 hand-restated Mathlib-style theorems with proofs (scripts/lean_ceiling_cases.py). + Each is first run through an INDEPENDENT compile-only pass (same REPL, no strictness at + all: no error message, no `sorry` warning). Items that fail it are version-incompatible + with the pin and are excluded and listed; the strict verdict is scored on the rest. + +Host rules (/apps/CLAUDE.md): `apply_resource_guards(heavy=True)` runs the preflight and +takes the exclusive heavy-job flock; scratch lives under align/.tmp/; nothing touches / or +/tmp; every blocking call has a timeout; the process has a global wall-clock watchdog. + + python scripts/lean_ceiling_test.py --workers 8 + python scripts/lean_ceiling_test.py --verify-prereg + python scripts/lean_ceiling_test.py --measure-only +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import math +import os +import re +import statistics +import subprocess +import sys +import threading +import time +import urllib.parse +import urllib.request +from pathlib import Path + +REPO = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(REPO)) + +import importlib.util # noqa: E402 + +# Loaded by path (not `import lean_ceiling_cases`): scripts/ is not a package, and the +# dependency firewall (rightly) rejects an import it cannot place in stdlib/allowed/align. +_spec = importlib.util.spec_from_file_location("lean_ceiling_cases", Path(__file__).with_name("lean_ceiling_cases.py")) +C = importlib.util.module_from_spec(_spec) +sys.modules["lean_ceiling_cases"] = C +_spec.loader.exec_module(C) + +DOC = REPO / "docs" / "LEAN_VERIFIER.md" +DATASET = "Goedel-LM/Lean-workbook-proofs" +SLICE_OFFSETS = (0, 5000, 10000, 15000, 20000, 25000) +SLICE_LEN = 25 +KNOWN_GOOD_FLOOR = 0.98 +MIN_KNOWN_GOOD = 50 +MIN_ADVERSARIAL = 30 + + +# ---------------------------------------------------------------------------- prereg + + +def registered_hash() -> str | None: + if not DOC.exists(): + return None + m = re.search(r"PREREG_TABLE_SHA256:\s*`?([0-9a-f]{64})`?", DOC.read_text()) + return m.group(1) if m else None + + +def verify_prereg() -> bool: + reg, cur = registered_hash(), C.table_hash() + print(f"registered : {reg}\ncurrent : {cur}\nmatch : {reg == cur}") + return reg == cur + + +# ---------------------------------------------------------------------------- verdict + + +def empty_known_good_entry(candidates: int, excluded: int) -> dict: + """Report entry for a run in which NO known-good proof reached the strict stage (toolchain broken, + REPL crashing, header mismatch: every candidate failed the independent compile-only filter). + That is the absence of the ceiling's first side, i.e. a FAILURE: a verifier that rejects + everything would otherwise sail through the adversarial side (accepted == 0 is trivially true).""" + return {"n": 0, "accepted": 0, "rate": float("nan"), "floor": KNOWN_GOOD_FLOOR, "pass": False, + "why": f"no known-good proof survived the compile-only filter ({excluded}/{candidates} excluded)", + "rejected": []} + + +def ceiling_ok(report: dict, only: str | None) -> bool: + """Exit-status verdict. Both sides must be PRESENT and passing in a full run; `--only` runs need + the side they ran. A missing side is never a pass (the old `all(v.get('pass', True) ...)` read + absence as success).""" + need = {None: ("known_good", "adversarial"), "good": ("known_good",), "adversarial": ("adversarial",)}[only] + if any(not isinstance(report.get(k), dict) or not report[k].get("pass") for k in need): + return False + return all(v.get("pass", True) for v in report.values() if isinstance(v, dict)) + + +# ---------------------------------------------------------------------------- statistics + + +def wilson(k: int, n: int, z: float = 1.96) -> tuple[float, float]: + """Wilson score interval for a binomial proportion (exact-ish at k=n or k=0).""" + if n == 0: + return (float("nan"), float("nan")) + p = k / n + d = 1 + z * z / n + c = (p + z * z / (2 * n)) / d + h = z * math.sqrt(p * (1 - p) / n + z * z / (4 * n * n)) / d + return (max(0.0, c - h), min(1.0, c + h)) + + +def pctl(xs: list[float], q: float) -> float: + if not xs: + return float("nan") + s = sorted(xs) + return s[min(len(s) - 1, int(q * len(s)))] + + +# ---------------------------------------------------------------------------- data + + +def http_json(url: str, timeout: float = 60.0, retries: int = 4) -> dict: + last = None + for i in range(retries): + try: + req = urllib.request.Request(url, headers={"User-Agent": "align-lean-ceiling/1"}) + with urllib.request.urlopen(req, timeout=timeout) as r: + return json.loads(r.read().decode("utf-8")) + except Exception as e: # noqa: BLE001 + last = e + time.sleep(2 * (i + 1)) + raise RuntimeError(f"GET {url} failed: {last!r}") + + +def fetch_known_good(out_dir: Path) -> tuple[list[dict], dict]: + rows: list[dict] = [] + for off in SLICE_OFFSETS: + q = urllib.parse.urlencode( + {"dataset": DATASET, "config": "default", "split": "train", "offset": off, "length": SLICE_LEN} + ) + d = http_json(f"https://datasets-server.huggingface.co/rows?{q}") + for r in d["rows"]: + rows.append({"problem_id": r["row"]["problem_id"], "full_proof": r["row"]["full_proof"], "row_idx": r["row_idx"]}) + raw = out_dir / "known_good_raw.jsonl" + with open(raw, "w") as f: + for r in rows: + f.write(json.dumps(r, ensure_ascii=False) + "\n") + meta = {"dataset": DATASET, "offsets": list(SLICE_OFFSETS), "slice_len": SLICE_LEN, "n_rows": len(rows), + "raw_sha256": hashlib.sha256(raw.read_bytes()).hexdigest()} + try: + info = http_json(f"https://huggingface.co/api/datasets/{DATASET}") + meta["dataset_revision"] = info.get("sha") + except Exception as e: # noqa: BLE001 + meta["dataset_revision"] = f"unavailable: {e!r}" + return rows, meta + + +def derive_statement(code: str): + """(statement, target_name) from a full proof file: header + the LAST theorem/lemma's + signature (cut at the signature-ending `:=`, via the verifier's own scanner) + `sorry`.""" + from align.verify.lean.strict import find_sig_end, lex + + lx = lex(code) + toks = lx.toks + if lx.broken: + return None + first = next((i for i, t in enumerate(toks) if t.kind == "id" and t.text in + ("theorem", "lemma", "def", "abbrev", "instance", "example", "structure", "inductive")), None) + thm = [i for i, t in enumerate(toks) if t.kind == "id" and t.text in ("theorem", "lemma")] + if first is None or not thm: + return None + k = thm[-1] + e = find_sig_end(toks, k + 2) + if e is None or k + 1 >= len(toks): + return None + name = toks[k + 1].text + header = code[: toks[first - 1].end] if first > 0 else "" + sig = code[toks[k + 2].start : toks[e].start] + return header + "\n\n" + f"theorem {name} {sig} := by\n sorry\n", name + + +def as_completion(code: str) -> str: + return "I will state a plan and then give the proof.\n\n```lean4\n" + code.rstrip("\n") + "\n```\n" + + +# ---------------------------------------------------------------------------- compile-only + + +def compile_only(w, statement: str, code: str, timeout: float = 120.0) -> dict: + """The NAIVE grader: elaborate the candidate and look at nothing but error messages. + Used (a) as the independent reference compile that defines "known-good", and (b) to show + how many adversarial cases a compile-only reward would pay.""" + from align.verify.lean.repl_pool import ReplError + from align.verify.lean.strict import lex, parse_statement + + ref = parse_statement(statement) + if not ref.ok: + return {"ok": False, "why": "unsupported_statement"} + lx = lex(code) + nh = len(ref.header_toks) + if lx.broken or tuple(t.text for t in lx.toks[:nh]) != ref.header_toks: + return {"ok": False, "why": "header"} + body = code[lx.toks[nh - 1].end :] if nh else code + try: + env = w.header_env(ref.header_text) + resp = w.command(body, env, timeout) + except ReplError as e: + return {"ok": False, "why": e.code} + msgs = [m for m in resp.get("messages", []) if isinstance(m, dict)] + err = [m for m in msgs if m.get("severity") == "error"] + sorry = bool(resp.get("sorries")) or any("sorry" in str(m.get("data", "")) for m in msgs) + return {"ok": not err, "clean": (not err) and (not sorry), "why": "error" if err else ("sorry" if sorry else "")} + + +def naive_fresh(w, code: str, timeout: float = 120.0) -> dict: + """Compile-only on the candidate's OWN header (fresh import, no cache) -- for the + adversarial cases whose header is the attack. Each case gets a FRESH process: a second + `import Mathlib` in a live REPL costs +3.9GiB anonymous (see `split_imports`).""" + from align.verify.lean.repl_pool import ReplError + + w.discard() + try: + proc = w.ensure() + resp = proc.request({"cmd": code}, timeout) + except ReplError as e: + w.discard() + return {"ok": False, "why": e.code} + w.discard() + err = [m for m in resp.get("messages", []) if isinstance(m, dict) and m.get("severity") == "error"] + return {"ok": not err, "why": "error" if err else ""} + + +# ---------------------------------------------------------------------------- measurement + + +def dir_bytes(p: Path) -> int: + out = subprocess.run(["du", "-sb", str(p)], capture_output=True, text=True, timeout=300) + return int(out.stdout.split()[0]) if out.returncode == 0 and out.stdout else -1 + + +def evict_page_cache(tc) -> dict: + """Drop the page cache for every .olean/.ilean under the install (POSIX_FADV_DONTNEED + works unprivileged for clean, unmapped pages) so the next import is a COLD read.""" + n = tot = 0 + roots = [tc.workspace / ".lake" / "packages", tc.root / "elan" / "toolchains"] + for r in roots: + for dp, _dn, fns in os.walk(r): + for fn in fns: + if fn.endswith((".olean", ".ilean", ".so")): + fp = os.path.join(dp, fn) + try: + fd = os.open(fp, os.O_RDONLY) + try: + tot += os.fstat(fd).st_size + os.posix_fadvise(fd, 0, 0, os.POSIX_FADV_DONTNEED) + n += 1 + finally: + os.close(fd) + except OSError: + pass + return {"files": n, "bytes": tot} + + +def cached_fraction(tc) -> float | None: + """Fraction of the Mathlib .olean bytes resident in page cache, via `fincore` if present.""" + import shutil + + f = shutil.which("fincore") + if not f: + return None + lib = tc.workspace / ".lake" / "packages" / "mathlib" / ".lake" / "build" / "lib" / "Mathlib" + files = [str(p) for p in lib.rglob("*.olean")][:4000] + try: + out = subprocess.run([f, "--bytes", "--raw", "--noheadings", *files], capture_output=True, text=True, timeout=120) + res = tot = 0 + for line in out.stdout.splitlines(): + parts = line.split() + if len(parts) >= 3 and parts[0].isdigit(): + res += int(parts[0]) + tot += int(parts[1]) + return res / tot if tot else None + except Exception: # noqa: BLE001 + return None + + +def measure_install(tc, workers: int, out: dict) -> None: + from align.verify.lean.repl_pool import ReplLimits, ReplPool, ReplProcess + + m: dict = {} + m["disk_bytes"] = { + "total": dir_bytes(tc.root), + "elan_toolchain": dir_bytes(tc.root / "elan"), + "workspace": dir_bytes(tc.workspace), + "mathlib_package": dir_bytes(tc.workspace / ".lake" / "packages" / "mathlib"), + "repl": dir_bytes(tc.root / "repl"), + } + hdr = C.HEADER.rstrip("\n") + + def one_import() -> tuple[float, dict | None]: + p = ReplProcess(tc.argv(), tc.child_env, cwd=tc.cwd, limits=ReplLimits()) + p.start() + t = time.monotonic() + r = p.request({"cmd": hdr}, 900) + dt = time.monotonic() - t + assert "env" in r and not r.get("messages"), r + mem = p.memory() + p.kill() + return dt, mem + + ev = evict_page_cache(tc) + m["evicted"] = ev + m["cached_fraction_after_evict"] = cached_fraction(tc) + cold, cold_mem = one_import() + m["cached_fraction_after_cold_import"] = cached_fraction(tc) + warm = [one_import()[0] for _ in range(3)] + m["import_s"] = {"cold": cold, "warm": warm, "warm_median": statistics.median(warm)} + m["single_repl_after_import_mib"] = cold_mem + + # N workers concurrently warm: per-process PSS (RSS double-counts the shared oleans) + pool = ReplPool.from_toolchain(tc, n_workers=workers, scratch_dir=str(tc.root / "tmp" / "measure")) + try: + t = time.monotonic() + errs = pool.warm(hdr) + m["pool_warm_s"] = time.monotonic() - t + assert not any(errs), errs + rep = pool.memory_report() + m["n_workers"] = workers + m["per_worker_mib"] = rep + pss = [r["pss"] for r in rep if r] + m["pss_mib"] = {"min": min(pss), "mean": statistics.mean(pss), "max": max(pss), "sum": sum(pss)} + m["pss_anon_mib"] = {"mean": statistics.mean(r["pss_anon"] for r in rep if r)} + m["rss_sum_mib"] = sum(r["rss"] for r in rep if r) + finally: + pool.close() + out["install"] = m + + +# ---------------------------------------------------------------------------- confusion + + +def confusion(good_acc: list[bool], bad_acc: list[bool]) -> dict: + tp = sum(good_acc) + fn = len(good_acc) - tp + fp = sum(bad_acc) + tn = len(bad_acc) - fp + return {"TP": tp, "FN": fn, "FP": fp, "TN": tn, "n_good": len(good_acc), "n_bad": len(bad_acc)} + + +def render_confusion(c: dict) -> str: + return ( + " verdict ACCEPT verdict REJECT total\n" + f"truth GOOD proofs {c['TP']:>14} {c['FN']:>14} {c['n_good']:>5}\n" + f"truth BAD proofs {c['FP']:>14} {c['TN']:>14} {c['n_bad']:>5}\n" + ) + + +# ---------------------------------------------------------------------------- main + + +def sample_memory(pool, stop: threading.Event, rows: list) -> None: + while not stop.is_set(): + try: + rep = pool.memory_report() + with open("/proc/meminfo") as f: + avail = next(int(l.split()[1]) for l in f if l.startswith("MemAvailable:")) / 1024**2 + rows.append({"t": time.time(), "workers": rep, "mem_available_gb": avail}) + except Exception: # noqa: BLE001 + pass + stop.wait(0.5) + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument("--workers", type=int, default=None) + ap.add_argument("--out", default=None) + ap.add_argument("--verify-prereg", action="store_true") + ap.add_argument("--allow-unregistered", action="store_true") + ap.add_argument("--measure-only", action="store_true") + ap.add_argument("--skip-measure", action="store_true") + ap.add_argument("--max-wall-s", type=float, default=3600.0) + ap.add_argument("--no-semantic-pass", action="store_true") + ap.add_argument("--no-naive-pass", action="store_true") + ap.add_argument("--only", choices=("good", "adversarial"), default=None) + ap.add_argument("--exploratory", action="store_true", + help="debug run: <=2 workers, no heavy-job lock (a registered run always takes it)") + ap.add_argument("--lock-wait-s", type=float, default=7200.0) + ap.add_argument("--no-kernel-recheck", action="store_true", + help="known-good pass without the kernel recheck (default: on, as in the shipped config)") + ap.add_argument("--kg-limit", type=int, default=0, + help="exploratory only: keep the first N rows of each fetched slice (0 = all)") + args = ap.parse_args() + + if args.verify_prereg: + return 0 if verify_prereg() else 2 + registered = registered_hash() == C.table_hash() + if not registered and not args.allow_unregistered and not args.measure_only: + print("REFUSING: case tables do not match the registered hash (see --verify-prereg); " + "pass --allow-unregistered for an exploratory run", file=sys.stderr) + return 2 + + os.environ.setdefault("HF_HOME", "/apps/development/rlhf/hf_cache") + os.environ.setdefault("HF_HUB_DISABLE_XET", "1") + from align.utils.guard import apply_resource_guards + + apply_resource_guards(cpu_threads=8, heavy=False, label="lean_ceiling_test") + if not args.exploratory: + # A Lean pool is a heavy job (research report, section 7A): same preflight + the same + # exclusive flock as the trainers. If a GPU job holds it, WAIT -- never overlap. + from align.exit_codes import PreflightRefused + from align.utils.guard import acquire_heavy_job_lock, preflight_host_check + + t_wait = time.monotonic() + while True: + try: + preflight_host_check(strict=True) + acquire_heavy_job_lock("lean_ceiling_test") + break + except PreflightRefused as e: + if time.monotonic() - t_wait > args.lock_wait_s: + print(f"gave up waiting for the host: {e}", file=sys.stderr) + return 83 + print(f"host busy ({str(e)[:120]}...); retrying in 30s", flush=True) + time.sleep(30) + elif (args.workers or 2) > 2: + print("--exploratory caps workers at 2", file=sys.stderr) + args.workers = 2 + + from align.verify.lean import ( + LeanAccuracy, ReplLimits, ReplPool, StrictConfig, StrictVerifier, default_workers, find_toolchain, + ) + from align.verify.lean.repl_pool import meminfo_available_gb, plan_workers + + tc = find_toolchain() + limits = ReplLimits() + n = args.workers or default_workers() + n = plan_workers(n, limits.anon_cap_gb, 4.0, meminfo_available_gb()) + stamp = time.strftime("%Y%m%d_%H%M%S") + out_dir = Path(args.out or REPO / ".tmp" / "lean_ceiling" / stamp) + out_dir.mkdir(parents=True, exist_ok=True) + pools: list = [] + + def watchdog() -> None: + time.sleep(args.max_wall_s) + print(f"WATCHDOG: {args.max_wall_s}s exceeded; killing pools", file=sys.stderr, flush=True) + for p in pools: + try: + p.close() + except Exception: # noqa: BLE001 + pass + os._exit(124) + + threading.Thread(target=watchdog, daemon=True).start() + + def on_signal(signum, _frame): + print(f"signal {signum}: closing pools", file=sys.stderr, flush=True) + for p in pools: + try: + p.close() + except Exception: # noqa: BLE001 + pass + os._exit(128 + signum) + + import signal + + signal.signal(signal.SIGTERM, on_signal) + signal.signal(signal.SIGINT, on_signal) + + result: dict = { + "started": stamp, "toolchain": tc.describe(), "pins": tc.pins, "workers": n, + "limits": limits.__dict__, "prereg_registered": registered, "table_sha256": C.table_hash(), + "cores": os.cpu_count(), + } + + def new_pool(workers: int, lim: ReplLimits = limits, label: str = "pool") -> ReplPool: + p = ReplPool.from_toolchain(tc, n_workers=workers, limits=lim, scratch_dir=str(tc.root / "tmp" / label)) + pools.append(p) + return p + + # -------------------------------------------------- install measurements + if not args.skip_measure and args.only is None: + print("== measuring install ==", flush=True) + measure_install(tc, n, result) + print(json.dumps(result["install"], indent=1, default=str)[:3000], flush=True) + if args.measure_only: + (out_dir / "result.json").write_text(json.dumps(result, indent=1, default=str)) + return 0 + + # -------------------------------------------------- known-good + good_items: list[dict] = [] + if args.only in (None, "good"): + print("== known-good: fetching ==", flush=True) + rows, meta = fetch_known_good(out_dir) + if args.kg_limit: + if not args.exploratory: + print("--kg-limit is for exploratory runs only", file=sys.stderr) + return 2 + rows = [r for i, r in enumerate(rows) if i % SLICE_LEN < args.kg_limit] + meta["kg_limit"] = args.kg_limit + result["known_good_source"] = meta + cands = [] + for r in rows: + code = r["full_proof"] + ds = derive_statement(code) + if ds is None: + cands.append({"id": r["problem_id"], "source": "lean-workbook-proofs", "code": code, "statement": None}) + else: + cands.append({"id": r["problem_id"], "source": "lean-workbook-proofs", "code": code, "statement": ds[0]}) + for rs in C.RESTATED: + code = C.HEADER + rs["code"] + ds = derive_statement(code) + cands.append({"id": rs["id"], "source": "restated", "code": code, "statement": ds[0] if ds else None}) + result["known_good_candidates"] = len(cands) + + print("== known-good: independent compile-only filter ==", flush=True) + pool = new_pool(n, label="kg_compile") + t0 = time.monotonic() + refs = pool.map( + lambda w, it: compile_only(w, it["statement"], it["code"]) if it["statement"] else {"ok": False, "clean": False, "why": "no_statement"}, + cands, + ) + pool.close() + excluded = [] + for it, r in zip(cands, refs): + if isinstance(r, Exception): + r = {"ok": False, "clean": False, "why": repr(r)} + it["ref"] = r + if r.get("clean"): + good_items.append(it) + else: + excluded.append({"id": it["id"], "source": it["source"], "why": r.get("why")}) + result["known_good_excluded"] = excluded + result["known_good_filter_s"] = time.monotonic() - t0 + print(f"compile-clean: {len(good_items)}/{len(cands)} excluded: {len(excluded)}", flush=True) + + # -------------------------------------------------- known-good: strict + good_results: list[dict] = [] + if good_items: + print("== known-good: strict verification ==", flush=True) + pool = new_pool(n, label="kg_strict") + acc = LeanAccuracy(StrictVerifier(pool, StrictConfig(kernel_recheck=not args.no_kernel_recheck))) + from align.verify.lean.strict import AUDIT_PRELUDE, static_check + + plan, _, _ = static_check(good_items[0]["statement"], good_items[0]["code"]) + t = time.monotonic() + pool.warm(plan.header if plan else C.HEADER, AUDIT_PRELUDE) + result["strict_warm_s"] = time.monotonic() - t + result["warm_memory_mib"] = pool.memory_report() + stop, mem_rows = threading.Event(), [] + th = threading.Thread(target=sample_memory, args=(pool, stop, mem_rows), daemon=True) + th.start() + t = time.monotonic() + def prog(ix, r): + print(f" [{ix:3d}] {good_items[ix]['id']:34s} {'ACCEPT' if r.accepted else 'REJECT'} {r.elapsed_s:7.2f}s {list(r.reasons)}", flush=True) + + accs, res = acc.accuracy_batch([g["statement"] for g in good_items], [as_completion(g["code"]) for g in good_items], progress=prog) + wall = time.monotonic() - t + stop.set() + th.join(5) + for g, a, r in zip(good_items, accs, res): + good_results.append({"id": g["id"], "source": g["source"], "accepted": bool(a), "reasons": list(r.reasons), + "diag": r.diagnostics, "elapsed_s": r.elapsed_s}) + lat = [r.elapsed_s for r in res] + result["throughput"] = { + "n_proofs": len(res), "wall_s": wall, "proofs_per_s": len(res) / wall, "workers": n, + "latency_s": {"p50": pctl(lat, 0.5), "p90": pctl(lat, 0.9), "p99": pctl(lat, 0.99), "max": max(lat), "mean": statistics.mean(lat)}, + "note": "unique statement per proof (no reference-env reuse across proofs); warm workers", + } + # per-worker PSS during the run + per_w: dict[int, list] = {} + for row in mem_rows: + for i, m in enumerate(row["workers"]): + if m: + per_w.setdefault(i, []).append(m) + result["pss_during_run_mib"] = { + str(i): {"pss_mean": statistics.mean(m["pss"] for m in ms), "pss_max": max(m["pss"] for m in ms), + "pss_anon_max": max(m["pss_anon"] for m in ms), "rss_max": max(m["rss"] for m in ms), + "samples": len(ms)} for i, ms in sorted(per_w.items()) + } + result["host_during_run"] = {"min_mem_available_gb": min((r["mem_available_gb"] for r in mem_rows), default=None)} + result["pool_stats"] = pool.stats.__dict__ + # repeat-G run: GRPO shape (same statement x8, reference env reused) on the 12 + # MEDIAN-latency statements of the pass above (the heavy tail would swamp a 96-proof run) + order = sorted(range(len(good_items)), key=lambda i: res[i].elapsed_s) + mid = max(0, len(order) // 2 - 6) + rep = [good_items[i] for i in order[mid : mid + 12]] * 8 + t = time.monotonic() + accs2, _ = acc.accuracy_batch([g["statement"] for g in rep], [as_completion(g["code"]) for g in rep]) + w2 = time.monotonic() - t + result["throughput_grpo_shape"] = {"n_proofs": len(rep), "wall_s": w2, "proofs_per_s": len(rep) / w2, "accepted": int(sum(accs2)), + "shape": "12 median-latency statements x 8 identical samples"} + pool.close() + + # -------------------------------------------------- adversarial + adv_results: list[dict] = [] + if args.only in (None, "adversarial"): + print("== adversarial: strict ==", flush=True) + + def completion_for(case) -> tuple[str, str | None]: + stmt = C.HEADER + C.STATEMENTS[case["statement"]] + if case["completion"] is not None: + return stmt, case["completion"] + hdr = C.HEADER if case["raw_header"] is None else case["raw_header"] + return stmt, as_completion(hdr + case["code"]) + + main_cases = [c for c in C.ADVERSARIAL if c["pool"] == "main"] + bomb_cases = [c for c in C.ADVERSARIAL if c["pool"] == "bomb"] + pool = new_pool(n, label="adv_strict") + acc = LeanAccuracy(StrictVerifier(pool, StrictConfig(proof_timeout_s=120)), min_batch=10**9) + pairs = [completion_for(c) for c in main_cases] + _, rs = acc.accuracy_batch([p[0] for p in pairs], [p[1] for p in pairs]) + pool.close() + by_id = {} + for c, r in zip(main_cases, rs): + by_id[c["id"]] = {"id": c["id"], "category": c["category"], "accepted": r.accepted, "reasons": list(r.reasons), + "stage": r.stage, "diag": r.diagnostics, "elapsed_s": r.elapsed_s, "pool": "main"} + # bombs: 1 worker, 1 GiB anon cap, 30 s budget so they terminate promptly (same code paths) + blim = ReplLimits(anon_cap_gb=1.0) + bpool = new_pool(1, blim, label="adv_bomb") + bacc = LeanAccuracy(StrictVerifier(bpool, StrictConfig(proof_timeout_s=30)), min_batch=10**9) + for c in bomb_cases: + stmt, comp = completion_for(c) + sc = bacc.score(stmt, comp) + by_id[c["id"]] = {"id": c["id"], "category": c["category"], "accepted": sc.result.accepted, + "reasons": list(sc.result.reasons), "stage": sc.result.stage, "diag": sc.result.diagnostics, + "elapsed_s": sc.result.elapsed_s, "pool": "bomb(1w,1GiB,30s)"} + bpool.close() + adv_results = [by_id[c["id"]] for c in C.ADVERSARIAL] + + # naive compile-only on each case's own text (what a compile-only reward would pay) + if not args.no_naive_pass: + print("== adversarial: naive compile-only (liveness) ==", flush=True) + npool = new_pool(n, label="adv_naive") + live_cases = [c for c in main_cases if c["completion"] is None and c["code"]] + texts = [(c["raw_header"] if c["raw_header"] is not None else C.HEADER) + c["code"] for c in live_cases] + nv = npool.map(lambda w, t: naive_fresh(w, t, 120.0), texts) + npool.close() + naive = {c["id"]: (v if not isinstance(v, Exception) else {"ok": False, "why": repr(v)}) for c, v in zip(live_cases, nv)} + for r in adv_results: + r["naive_compile_ok"] = naive.get(r["id"], {}).get("ok") + + # semantic layer alone (lexical gate OFF): how much does the Lean-side audit catch by itself? + if not args.no_semantic_pass: + print("== adversarial: semantic layer only (lexical gate off) ==", flush=True) + spool = new_pool(n, label="adv_sem") + sv = StrictVerifier(spool, StrictConfig(lexical_gate=False, kernel_recheck=True, proof_timeout_s=120)) + sem_cases = [c for c in main_cases if c["completion"] is None and c["code"]] + from align.verify.lean.reward import extract_lean_code + + sem_items = [] + for c in sem_cases: + stmt, comp = completion_for(c) + sem_items.append((stmt, extract_lean_code(comp))) + sres = sv.verify_batch(sem_items) + spool.close() + sem = {c["id"]: (r.accepted, list(r.reasons)) for c, r in zip(sem_cases, sres)} + for r in adv_results: + if r["id"] in sem: + r["semantic_only_accepted"], r["semantic_only_reasons"] = sem[r["id"]] + + # -------------------------------------------------- report + report = {} + if not good_results and args.only in (None, "good"): + report["known_good"] = empty_known_good_entry( + result.get("known_good_candidates", 0), len(result.get("known_good_excluded", [])) + ) + if good_results: + acc_good = [g["accepted"] for g in good_results] + lo, hi = wilson(sum(acc_good), len(acc_good)) + report["known_good"] = {"n": len(acc_good), "accepted": sum(acc_good), "rate": sum(acc_good) / len(acc_good), + "wilson95": [lo, hi], "floor": KNOWN_GOOD_FLOOR, + "pass": (sum(acc_good) / len(acc_good) >= KNOWN_GOOD_FLOOR and len(acc_good) >= MIN_KNOWN_GOOD), + "rejected": [g for g in good_results if not g["accepted"]]} + if adv_results: + acc_bad = [a["accepted"] for a in adv_results] + by_cat: dict[str, list] = {} + for a in adv_results: + by_cat.setdefault(a["category"], []).append(a["accepted"]) + report["adversarial"] = {"n": len(acc_bad), "accepted": sum(acc_bad), "pass": (sum(acc_bad) == 0 and len(acc_bad) >= MIN_ADVERSARIAL), + "by_category": {k: {"n": len(v), "accepted": sum(v)} for k, v in sorted(by_cat.items())}, + "accepted_cases": [a for a in adv_results if a["accepted"]]} + live = [a for a in adv_results if a.get("naive_compile_ok")] + report["adversarial"]["naive_compile_only_would_accept"] = len(live) + if any("semantic_only_accepted" in a for a in adv_results): + sem_acc = [a for a in adv_results if a.get("semantic_only_accepted")] + report["adversarial"]["semantic_only_would_accept"] = [a["id"] for a in sem_acc] + if good_results and adv_results: + report["confusion"] = confusion([g["accepted"] for g in good_results], [a["accepted"] for a in adv_results]) + result["report"] = report + result["good_results"] = good_results + result["adversarial_results"] = adv_results + (out_dir / "result.json").write_text(json.dumps(result, indent=1, default=str, ensure_ascii=False)) + print("\n================ RESULT ================") + if "confusion" in report: + print(render_confusion(report["confusion"])) + if report.get("known_good", {}).get("n"): + k = report["known_good"] + print(f"known-good accepted {k['accepted']}/{k['n']} = {k['rate']:.4f} Wilson95 [{k['wilson95'][0]:.4f}, {k['wilson95'][1]:.4f}] floor {k['floor']} PASS={k['pass']}") + for g in k["rejected"]: + print(" REJECTED known-good:", g["id"], g["reasons"], g["diag"]) + if "adversarial" in report: + a = report["adversarial"] + print(f"adversarial accepted {a['accepted']}/{a['n']} PASS={a['pass']} (naive compile-only would accept {a['naive_compile_only_would_accept']})") + for c in a["accepted_cases"]: + print(" ACCEPTED adversarial:", c["id"], c["reasons"]) + if "throughput" in result: + print("throughput:", json.dumps(result["throughput"])) + print(f"results: {out_dir/'result.json'}") + if "known_good" in report and not report["known_good"]["n"]: + print(f"known-good: NO EVIDENCE -- {report['known_good']['why']} PASS=False") + return 0 if ceiling_ok(report, args.only) else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/code/scripts/lean_setup.sh b/code/scripts/lean_setup.sh new file mode 100644 index 0000000000000000000000000000000000000000..f973f43fdf5ef260bd9801cff17e78843221cf90 --- /dev/null +++ b/code/scripts/lean_setup.sh @@ -0,0 +1,174 @@ +#!/usr/bin/env bash +# Install the pinned Lean 4 + Mathlib + REPL toolchain under /apps (Track L / P1). +# +# WHY this exists as a script and not as Python: the toolchain is an EXTERNAL BINARY +# (see docs/DEPENDENCY_POLICY.md, "Lean toolchain"). It is installed once, outside the +# Python environment, and `align/verify/lean` reaches it only through subprocess. +# +# Idempotent: every step is guarded by a marker/exists check, so re-running after a +# partial failure resumes instead of redoing a 5 GB download. Every network or build +# step runs under `timeout`, so a stalled mirror cannot hang the shell. +# +# Host rules (/apps/CLAUDE.md, binding): nothing is ever written to `/` or `/tmp`. +# elan, the Mathlib cache, git temp files and XDG caches are all redirected below. +# The preflight refuses on <45GB MemAvailable, load >= 2x cores, disk queue >= 64, or +# less than 40GB free on /apps. A Mathlib cache decompress is I/O heavy; it is not a +# GPU job so it does not take the GPU heavy-job lock. +# +# THE PIN (see docs/LEAN_VERIFIER.md section "Pin"): +# Goedel-Prover-V2 follows DeepSeek-Prover-V1.5: its `mathlib4` git submodule is +# xinhjBrant/mathlib4 @ 2f65ba7f1a9144b20c8e7358513548e317d26de1, a single commit on +# top of leanprover-community/mathlib4 @ d1d1e4b72c6e009edd97cabe1d661b021fe8eeae +# (2024-06-12) that ONLY edits lakefile.lean/lake-manifest.json to add DeepSeek's REPL +# fork as a dependency. Its lean-toolchain is `leanprover/lean4:v4.9.0-rc1` (nominally +# "Lean 4.9"; the binary reports 4.9.0-rc1). We pin the UPSTREAM parent: the Mathlib +# .olean cache is keyed on file hashes that include lakefile.lean, so the DeepSeek +# edit would orphan the cache, while the parent's .oleans are byte-for-byte the +# Mathlib sources Goedel proofs were checked against. The REPL is upstream +# leanprover-community/repl at the tag whose toolchain equals ours. +# +# Usage: bash scripts/lean_setup.sh # full install +# LEAN_ROOT=... bash scripts/lean_setup.sh +# bash scripts/lean_setup.sh --check # only verify an existing install +set -euo pipefail + +LEAN_ROOT="${LEAN_ROOT:-/apps/development/rlhf/lean}" +LEAN_TOOLCHAIN="${LEAN_TOOLCHAIN:-leanprover/lean4:v4.9.0-rc1}" +MATHLIB_REV="${MATHLIB_REV:-d1d1e4b72c6e009edd97cabe1d661b021fe8eeae}" +REPL_TAG="${REPL_TAG:-v4.9.0-rc1}" +REPL_URL="${REPL_URL:-https://github.com/leanprover-community/repl}" +MATHLIB_URL="https://github.com/leanprover-community/mathlib4" + +case "$LEAN_ROOT" in + /apps/*) ;; + *) echo "refusing: LEAN_ROOT=$LEAN_ROOT is not under /apps (host rule: never write to / or /tmp)" >&2; exit 2 ;; +esac + +export ELAN_HOME="$LEAN_ROOT/elan" +export MATHLIB_CACHE_DIR="$LEAN_ROOT/mathlib-cache" +export XDG_CACHE_HOME="$LEAN_ROOT/xdg-cache" +export XDG_CONFIG_HOME="$LEAN_ROOT/xdg-config" +export TMPDIR="$LEAN_ROOT/tmp" +export GIT_TERMINAL_PROMPT=0 +export PATH="$ELAN_HOME/bin:$PATH" +WS="$LEAN_ROOT/workspace" +REPL_DIR="$LEAN_ROOT/repl" +LOGS="$LEAN_ROOT/logs" +mkdir -p "$LEAN_ROOT" "$MATHLIB_CACHE_DIR" "$XDG_CACHE_HOME" "$XDG_CONFIG_HOME" "$TMPDIR" "$LOGS" + +log() { printf '[lean_setup %s] %s\n' "$(date +%H:%M:%S)" "$*"; } + +preflight() { + local avail_kb load cores q free_gb + avail_kb=$(awk '/^MemAvailable:/{print $2}' /proc/meminfo) + load=$(awk '{print $1}' /proc/loadavg) + cores=$(nproc) + q=$(awk '$3=="nvme2n1"{print $12}' /proc/diskstats) + q=${q:-0} + free_gb=$(df -BG --output=avail /apps | tail -1 | tr -dc 0-9) + log "preflight: MemAvailable=$((avail_kb/1024/1024))GB load1=$load cores=$cores diskq=$q /apps_free=${free_gb}GB" + [ "$avail_kb" -ge $((45*1024*1024)) ] || { log "REFUSED: MemAvailable < 45GB"; exit 83; } + awk -v l="$load" -v c="$cores" 'BEGIN{exit !(l < 2*c)}' || { log "REFUSED: load >= 2x cores"; exit 83; } + [ "$q" -lt 64 ] || { log "REFUSED: disk queue >= 64"; exit 83; } + [ "$free_gb" -ge 40 ] || { log "REFUSED: < 40GB free on /apps"; exit 83; } +} + +check_install() { + [ -x "$ELAN_HOME/bin/lake" ] || { log "missing elan/lake"; return 1; } + [ -x "$REPL_DIR/.lake/build/bin/repl" ] || { log "missing repl binary"; return 1; } + [ -f "$WS/.lake/packages/mathlib/.lake/build/lib/Mathlib.olean" ] || { log "missing Mathlib.olean (cache not fetched)"; return 1; } + log "install present" +} + +if [ "${1:-}" = "--check" ]; then check_install; exit $?; fi + +preflight + +# ---- 1. elan ----------------------------------------------------------------------- +if [ ! -x "$ELAN_HOME/bin/elan" ]; then + log "installing elan into $ELAN_HOME" + timeout 120 curl -sSfL https://raw.githubusercontent.com/leanprover/elan/master/elan-init.sh -o "$LEAN_ROOT/elan-init.sh" + # --no-modify-path: never touch ~/.profile / ~/.bashrc. --default-toolchain none: the + # toolchain is installed explicitly below so the pin is visible in this file. + timeout 600 sh "$LEAN_ROOT/elan-init.sh" -y --no-modify-path --default-toolchain none +fi +elan --version + +# ---- 2. toolchain ------------------------------------------------------------------ +if ! elan toolchain list | grep -qF "${LEAN_TOOLCHAIN##*:}"; then + log "installing $LEAN_TOOLCHAIN" + timeout 900 elan toolchain install "$LEAN_TOOLCHAIN" +fi + +# ---- 3. workspace (Lake project requiring Mathlib at the pinned rev) --------------- +mkdir -p "$WS" +if [ ! -f "$WS/lakefile.lean" ]; then + log "writing workspace lakefile" + printf '%s\n' "$LEAN_TOOLCHAIN" > "$WS/lean-toolchain" + cat > "$WS/lakefile.lean" <&1 | tee "$LOGS/lake_update.log" | tail -5 +fi +# Guard against silent drift: the manifest must name exactly the pinned rev. +grep -q "$MATHLIB_REV" "$WS/lake-manifest.json" || { log "FATAL: manifest does not pin $MATHLIB_REV"; exit 3; } + +# ---- 4. prebuilt Mathlib cache ----------------------------------------------------- +if [ ! -f "$WS/.lake/packages/mathlib/.lake/build/lib/Mathlib.olean" ]; then + log "lake exe cache get (downloads + unpacks the prebuilt oleans)" + (cd "$WS" && timeout 3000 lake exe cache get) 2>&1 | tee "$LOGS/cache_get.log" | tail -8 +fi +# Deps (batteries, Qq, aesop, proofwidgets) are built from source by lake when the cache +# does not cover them; this is a no-op when it does. +if ! (cd "$WS" && timeout 120 lake build --no-build Mathlib >/dev/null 2>&1); then + log "lake build Mathlib (deps not fully cached; building the remainder)" + (cd "$WS" && timeout 5400 lake build Mathlib) 2>&1 | tee "$LOGS/lake_build.log" | tail -5 +fi + +# ---- 5. REPL at the matching tag --------------------------------------------------- +if [ ! -d "$REPL_DIR/.git" ]; then + log "cloning repl $REPL_TAG" + timeout 300 git clone --branch "$REPL_TAG" --depth 1 "$REPL_URL" "$REPL_DIR" +fi +REPL_TC=$(cat "$REPL_DIR/lean-toolchain") +[ "$REPL_TC" = "$LEAN_TOOLCHAIN" ] || { log "FATAL: repl toolchain $REPL_TC != $LEAN_TOOLCHAIN"; exit 3; } +if [ ! -x "$REPL_DIR/.lake/build/bin/repl" ]; then + log "building repl" + (cd "$REPL_DIR" && timeout 1200 lake build) 2>&1 | tee "$LOGS/repl_build.log" | tail -5 +fi + +# ---- 6. smoke test: import Mathlib through the built REPL -------------------------- +log "smoke test" +SMOKE=$(cd "$WS" && printf '%s\n\n%s\n\n' \ + '{"cmd": "import Mathlib\ntheorem t (a b : Nat) : a + b = b + a := by omega"}' \ + '{"cmd": "#print axioms t", "env": 0}' \ + | timeout 300 lake env "$REPL_DIR/.lake/build/bin/repl") +echo "$SMOKE" | head -20 +echo "$SMOKE" | grep -q '"env"' || { log "FATAL: REPL smoke failed"; exit 4; } + +cat > "$LEAN_ROOT/PINS.json" < "$LEAN_ROOT/env.sh" </dev/null | cut -f1) total under $LEAN_ROOT" +log "done" diff --git a/code/scripts/olmoe_parity.py b/code/scripts/olmoe_parity.py new file mode 100644 index 0000000000000000000000000000000000000000..717bf5df425abd59d53e3f40e46b6d21a19b12c5 --- /dev/null +++ b/code/scripts/olmoe_parity.py @@ -0,0 +1,515 @@ +#!/usr/bin/env python +"""OLMoE-1B-7B-0125 parity on the REAL checkpoint: align vs transformers, on CPU. + + .venv/bin/python scripts/olmoe_parity.py run [--dtype fp32|bf16] [--ckpt DIR] [--out DIR] + .venv/bin/python scripts/olmoe_parity.py compare --a DIR/align --b DIR/hf --dtype fp32 + +Why this exists. The tiny-model oracle tier (tests/test_olmoe_parity.py) pins the +architecture at ~1e-6 on random 3-layer weights. It cannot see the things only the +released weights can expose: the HUB key layout at 64 experts x 16 layers (a transposed +gate/up stack is invisible when the weights are random and symmetric), a `config.json` +whose rope is spelled the transformers-4.x way, fp32-stored tensors, a 50304-row +embedding whose last rows are padding, and depth: a 1e-7 per-layer error is invisible at +3 layers and visible at 16. "Align is an OLMoE repo" is a claim about THIS checkpoint. + +THE GATE (pinned in `evaluate`, threshold constants below): + + fp32 compute: max |logit_align - logit_hf| < 1e-3 AND top-1 agreement 100% + AND every layer's router top-8 SET agrees on >= 99.9% of tokens. + bf16 compute: measured and reported, not gated (see BF16 note on `evaluate`). + +Why fp32 and not bf16. The released weights are stored fp32, so fp32 compute is the same +function, not a rounded copy of it: any residual is accumulation order, which at 16 +layers and logit magnitudes of ~10-25 was measured at 1.4e-5 (sequential `index_add_` combine) +and 3.0e-5 (fixed-order combine, a different summation order from transformers'). A bf16 gate (top-1 100% with +max-abs < 0.1) would be ~100x looser, would hide a real porting error (a mis-scaled norm, RoPE theta +off by a factor typically lands at 1e-2..10), and was in fact UNSATISFIABLE by two correct bf16 +implementations (measured max-abs 1.3 against |logit| up to 25: the bf16 ulp there is 0.125). bf16 is +still RUN and recorded (`--dtype bf16`) because it is the dtype training uses; its disagreement is the +noise floor the fp32 gate must not be confused with. + +Memory discipline (host: 186GB, no swap, shared; /apps/CLAUDE.md is binding): + * fp32 weights are 27.7GB per copy (the checkpoint is fp32 on disk). align and + transformers are run in SEPARATE processes, one after the other, and compare saved + tensors, so the peak is one copy plus the loader's transient (align stacks 64 + per-expert tensors per layer: ~2x its weights at the peak), never two models. + * Each worker runs the host preflight and takes the heavy-job flock before importing + torch; the orchestrator itself holds nothing, so the workers cannot deadlock on it. + * Workers see no GPU (CUDA_VISIBLE_DEVICES=""): this is a CPU measurement and it must not + land on a card another job is using. + * Everything is written under --out (default .tmp/olmoe/parity), never /tmp. + +Import order is load-bearing (`apply_resource_guards` sets env torch reads at import), so +torch is imported inside functions, after the guards. The legacy-venv half is a source +string run by `python -c`, like tests/test_olmoe_parity.py: `transformers` is an oracle, not +an import of this file (tests/test_dependency_firewall.py scans it). +""" + +from __future__ import annotations + +import argparse +import glob +import json +import os +import resource +import subprocess +import sys +import time + +REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +LEGACY_PY = "/apps/development/rlhf/nano-mistral-o1/.venv/bin/python" +HF_HOME = "/apps/development/rlhf/hf_cache" +CKPT_GLOB = os.path.join(HF_HOME, "hub", "models--allenai--OLMoE-1B-7B-0125", "snapshots", "*") +DEFAULT_OUT = os.path.join(REPO, ".tmp", "olmoe", "parity") + +#: Three fixed prompts of different character (prose, code, numbers) and length. They are +#: tokenized by the HF tokenizer in the legacy venv; align receives ids, never text, because +#: the align tokenizer is Mistral-only (ROADMAP P2-9). +PROMPTS: tuple[str, ...] = ( + "The capital of France is", + "def fibonacci(n):\n if n < 2:\n return n\n return fibonacci(n - 1) + fibonacci(n - 2)\n", + "In 1969, Apollo 11 landed on the Moon after a journey of 384,400 kilometres; " + "Neil Armstrong said that it was one small step for man.", +) + +# --------------------------------------------------------------------------- +# The gate (pure functions over numpy arrays: no torch, unit-testable) +# --------------------------------------------------------------------------- + +#: fp32: the same function up to accumulation order. See the module docstring. +FP32_MAX_ABS_LOGIT_DIFF = 1e-3 +#: Fraction of tokens whose top-8 expert SET must agree, per layer, fp32. Not 1.0: a +#: router logit tie within ~1e-6 between the 8th and 9th expert may legitimately flip +#: under a different summation order. A wrong router (QK-norm, transposed gate) misses +#: by whole percents, far below this. +FP32_ROUTER_SET_AGREEMENT = 0.999 +#: bf16 is MEASURED, NOT GATED. The alternative gate (top-1 100% and max-abs < 0.1) is unsatisfiable +#: by two CORRECT bf16 implementations: bf16 logits of magnitude 16-32 have an ulp of 0.125, and 16 layers +#: of different accumulation order put the two engines ~10 ulps apart (measured 1.3 on this checkpoint, +#: MEASUREMENTS.md). A tolerance chosen to pass would be a number fitted to the answer. +TOP_K = 8 + + +def topk_sets(router_logits, k: int = TOP_K): + """[T, E] router logits -> [T, k] expert ids, each row sorted ascending so two rows + are equal iff they pick the same SET (order within top-k is irrelevant to the forward).""" + import numpy as np + + idx = np.argpartition(-router_logits, k - 1, axis=-1)[..., :k] + return np.sort(idx, axis=-1) + + +def compare_tensors(a: dict, b: dict, n_prompts: int, n_layers: int, k: int = TOP_K) -> dict: + """Metrics between two engines' saved tensors (dicts of numpy arrays). + + Keys: `logits_{p}` [T, V], `router_logits_{p}_{l}` [T, E]. Returns plain python + numbers; `evaluate` turns them into a verdict. + """ + import numpy as np + + per_prompt = [] + worst = 0.0 + n_tok = n_agree = 0 + for p in range(n_prompts): + la, lb = a[f"logits_{p}"].astype(np.float64), b[f"logits_{p}"].astype(np.float64) + if la.shape != lb.shape: + raise ValueError(f"logits_{p}: shape {la.shape} vs {lb.shape}") + d = np.abs(la - lb) + top_a, top_b = la.argmax(-1), lb.argmax(-1) + agree = int((top_a == top_b).sum()) + per_prompt.append( + { + "tokens": int(la.shape[0]), + "max_abs": float(d.max()), + "mean_abs": float(d.mean()), + "abs_max_logit": float(np.abs(lb).max()), + "top1_agree": agree, + } + ) + worst = max(worst, float(d.max())) + n_tok += la.shape[0] + n_agree += agree + + layers = [] + for l in range(n_layers): + match = total = 0 + rl_max = 0.0 + min_overlap = k + for p in range(n_prompts): + ra, rb = a[f"router_logits_{p}_{l}"], b[f"router_logits_{p}_{l}"] + rl_max = max(rl_max, float(np.abs(ra.astype(np.float64) - rb.astype(np.float64)).max())) + sa, sb = topk_sets(ra, k), topk_sets(rb, k) + eq = (sa == sb).all(-1) + match += int(eq.sum()) + total += int(eq.size) + for i in np.nonzero(~eq)[0]: + min_overlap = min(min_overlap, len(set(sa[i].tolist()) & set(sb[i].tolist()))) + layers.append( + { + "layer": l, + "tokens": total, + "set_agreement": match / total, + "min_overlap": int(min_overlap), + "router_logit_max_abs": rl_max, + } + ) + return { + "max_abs_logit_diff": worst, + "top1_agreement": n_agree / n_tok, + "tokens": n_tok, + "per_prompt": per_prompt, + "router_layers": layers, + "min_router_set_agreement": min(x["set_agreement"] for x in layers), + } + + +def evaluate(metrics: dict, dtype: str) -> dict: + """The verdict. Only fp32 is a gate; bf16 returns `gated=False, passed=None` with the same + measurements. Anything else raises (a typo must not fall through to a default tolerance).""" + if dtype == "bf16": + return {"dtype": dtype, "gated": False, "passed": None, "thresholds": None, "checks": {}, + "note": "bf16 is recorded as the noise floor between two correct implementations, not gated"} + if dtype != "fp32": + raise ValueError(f"dtype must be 'fp32' or 'bf16', got {dtype!r}") + checks = { + "max_abs_logit_diff": metrics["max_abs_logit_diff"] < FP32_MAX_ABS_LOGIT_DIFF, + "top1_agreement": metrics["top1_agreement"] == 1.0, + "router_set_agreement": metrics["min_router_set_agreement"] >= FP32_ROUTER_SET_AGREEMENT, + } + return { + "dtype": dtype, "gated": True, + "thresholds": {"max_abs_logit_diff": FP32_MAX_ABS_LOGIT_DIFF, "top1_agreement": 1.0, + "router_set_agreement": FP32_ROUTER_SET_AGREEMENT}, + "checks": checks, + "passed": all(checks.values()), + } + + +def _load_dir(d: str) -> tuple[dict, dict]: + from safetensors.numpy import load_file + + with open(os.path.join(d, "meta.json")) as f: + meta = json.load(f) + return load_file(os.path.join(d, "tensors.safetensors")), meta + + +def compare_dirs(a_dir: str, b_dir: str, dtype: str, allow_mutant: bool = False) -> dict: + a, ma = _load_dir(a_dir) + if "MUTANT" in ma["engine"] and not allow_mutant: + raise ValueError(f"{a_dir} holds a MUTANT run ({ma['engine']}); refusing to issue a parity verdict on it") + b, mb = _load_dir(b_dir) + if ma["input_ids"] != mb["input_ids"]: + raise ValueError("the two engines were fed different token ids; the comparison is void") + if ma["dtype"] != dtype or mb["dtype"] != dtype: + raise ValueError(f"dtype mismatch: gate={dtype}, a={ma['dtype']}, b={mb['dtype']}") + n_layers = ma["n_layers"] + metrics = compare_tensors(a, b, len(ma["input_ids"]), n_layers) + verdict = evaluate(metrics, dtype) + return { + "metrics": metrics, + "verdict": verdict, + "engines": {"a": {k: v for k, v in ma.items() if k != "input_ids"}, "b": {k: v for k, v in mb.items() if k != "input_ids"}}, + "input_ids": ma["input_ids"], + } + + +# --------------------------------------------------------------------------- +# Workers +# --------------------------------------------------------------------------- + + +def default_ckpt() -> str: + hits = sorted(glob.glob(CKPT_GLOB)) + if not hits: + raise FileNotFoundError(f"no OLMoE-1B-7B-0125 snapshot under {CKPT_GLOB}") + return hits[-1] + + +def checkpoint_present(ckpt: str | None = None) -> bool: + try: + c = ckpt or default_ckpt() + except FileNotFoundError: + return False + return os.path.isfile(os.path.join(c, "model.safetensors.index.json")) and len( + glob.glob(os.path.join(c, "model-*-of-*.safetensors")) + ) == 6 + + +def _peak_rss_gb() -> float: + return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024**2 # KB -> GiB + + +#: Known OLMoE porting defects, each injected into the align side ONLY so the gate can be +#: shown to bite on the real checkpoint (not only on the random tiny model of the oracle +#: tier). Name -> what is broken. A parity gate that cannot fail is decoration. +MUTANTS: dict[str, str] = { + "per_head_qknorm": "q/k RMSNorm per head instead of over the full projection (Qwen3/Gemma3 style)", + "swap_gate_up": "expert gate and up rows stacked in the wrong order", + "topk_then_softmax": "Mixtral routing: top-k on logits, softmax over only the kept experts", + "rope_theta_1e6": "Mistral's RoPE base (1e6) instead of OLMoE's 1e4", +} + + +def _apply_mutant(name: str, cfg): + """Monkeypatch the named defect into the already-imported align modules; returns `cfg` + (possibly replaced). Raises on an unknown name: a typo must not silently run the control.""" + import torch + + from align.model import moe + from align.model import olmoe as om + from align.model import olmoe_checkpoint as oc + + if name == "per_head_qknorm": + + def per_head(norm, proj): + hd = cfg.head_dim + shp = proj.shape + w = norm.weight.view(-1, hd) + v = proj.view(*shp[:-1], -1, hd).float() + v = v * torch.rsqrt(v.pow(2).mean(-1, keepdim=True) + norm.eps) + return (w * v.to(proj.dtype)).view(shp) + + om._qk_norm = per_head + elif name == "swap_gate_up": + oc._stack_gate_up = lambda g, u: torch.cat([u, g], dim=0) + elif name == "topk_then_softmax": + + def mixtral(x, weight, top_k, norm_topk_prob=False): + logits = moe._router_logits(x, weight) + probs = torch.softmax(logits, dim=-1) + vals, idx = torch.topk(logits, top_k, dim=-1) + return moe.RouterOutput(weights=torch.softmax(vals, dim=-1), indices=idx, logits=logits, probs=probs) + + moe.route_top_k = mixtral + elif name == "rope_theta_1e6": + cfg = cfg.replace(rope_theta=1e6) + else: + raise ValueError(f"unknown mutant {name!r}; known: {sorted(MUTANTS)}") + return cfg + + +def worker_align(ckpt: str, ids_path: str, out_dir: str, dtype: str, mutant: str | None = None) -> None: + """Runs in the align venv. Guards first, torch after.""" + from align.utils.guard import apply_resource_guards + + apply_resource_guards(cpu_threads=8, heavy=True, label=f"olmoe-parity-align-{dtype}") + import torch + + from align.model import olmoe as om + from align.model import olmoe_checkpoint as oc + from align.model.checkpoint import load_hf_safetensors + + t_dtype = {"fp32": torch.float32, "bf16": torch.bfloat16}[dtype] + with open(ids_path) as f: + ids = json.load(f)["input_ids"] + t0 = time.perf_counter() + cfg = oc.read_olmoe_config(ckpt) + if mutant: + cfg = _apply_mutant(mutant, cfg) + # Loaded straight into the compute dtype shard by shard (load_hf_safetensors casts per + # shard), so the transient is the target dtype's size, not fp32 + target. + sd = oc.hf_to_align(load_hf_safetensors(ckpt, dtype=t_dtype), cfg) + model = om.OlmoeForCausalLM.from_state_dict(cfg, sd, moe_backend="loop", attn_implementation="sdpa").eval() + del sd + t_load = time.perf_counter() - t0 + + tensors: dict = {} + t1 = time.perf_counter() + with torch.no_grad(): + for p, row in enumerate(ids): + x = torch.tensor([row], dtype=torch.long) + logits, stats = model(input_ids=x, return_router_stats=True, keep_router_logits=True) + tensors[f"logits_{p}"] = logits[0].float().numpy() + for l, s in enumerate(stats): + tensors[f"router_logits_{p}_{l}"] = s.logits.float().numpy() + t_fwd = time.perf_counter() - t1 + _write(out_dir, tensors, { + "engine": "align" + (f"+MUTANT:{mutant}" if mutant else ""), "dtype": dtype, "n_layers": cfg.n_layers, "input_ids": ids, + "torch": torch.__version__, "attn": model.attn_implementation, "moe_backend": "loop", + "load_s": round(t_load, 2), "forward_s": round(t_fwd, 2), "peak_rss_gb": round(_peak_rss_gb(), 2), + }) + + +def _write(out_dir: str, tensors: dict, meta: dict) -> None: + import numpy as np + from safetensors.numpy import save_file + + os.makedirs(out_dir, exist_ok=True) + save_file({k: np.ascontiguousarray(v) for k, v in tensors.items()}, os.path.join(out_dir, "tensors.safetensors")) + with open(os.path.join(out_dir, "meta.json"), "w") as f: + json.dump(meta, f, indent=2) + + +# Runs INSIDE the legacy venv (transformers 5.x). argv: mode ckpt ids_json out_dir dtype repo +_LEGACY_SCRIPT = r''' +import json, os, resource, sys, time +mode, ckpt, ids_path, out_dir, dtype, repo = sys.argv[1:7] +os.environ.setdefault("HF_HUB_OFFLINE", "1") +os.environ.setdefault("TRANSFORMERS_OFFLINE", "1") +os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") +if mode == "tokenize": + from transformers import AutoTokenizer + prompts = json.loads(ids_path) + tok = AutoTokenizer.from_pretrained(ckpt) + ids = [tok(p)["input_ids"] for p in prompts] + os.makedirs(out_dir, exist_ok=True) + with open(os.path.join(out_dir, "ids.json"), "w") as f: + json.dump({"prompts": prompts, "input_ids": ids, "tokenizer": type(tok).__name__, + "vocab_size": len(tok)}, f, indent=2) + print(json.dumps({"n": [len(x) for x in ids]})) + sys.exit(0) + +sys.path.insert(0, repo) +from align.utils.guard import apply_resource_guards +apply_resource_guards(cpu_threads=8, heavy=True, label="olmoe-parity-hf-" + dtype) +import numpy as np, torch, transformers +from safetensors.numpy import save_file +from transformers import OlmoeForCausalLM +t_dtype = {"fp32": torch.float32, "bf16": torch.bfloat16}[dtype] +ids = json.load(open(ids_path))["input_ids"] +t0 = time.perf_counter() +m = OlmoeForCausalLM.from_pretrained(ckpt, dtype=t_dtype, attn_implementation="eager").eval() +t_load = time.perf_counter() - t0 +tensors = {} +t1 = time.perf_counter() +with torch.no_grad(): + for p, row in enumerate(ids): + o = m(input_ids=torch.tensor([row]), output_router_logits=True) + tensors[f"logits_{p}"] = o.logits[0].float().numpy() + for l, rl in enumerate(o.router_logits): + tensors[f"router_logits_{p}_{l}"] = rl.float().numpy() +t_fwd = time.perf_counter() - t1 +os.makedirs(out_dir, exist_ok=True) +save_file({k: np.ascontiguousarray(v) for k, v in tensors.items()}, os.path.join(out_dir, "tensors.safetensors")) +meta = {"engine": "transformers", "dtype": dtype, "n_layers": m.config.num_hidden_layers, "input_ids": ids, + "transformers": transformers.__version__, "torch": torch.__version__, "attn": "eager", + "load_s": round(t_load, 2), "forward_s": round(t_fwd, 2), + "peak_rss_gb": round(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024**2, 2)} +json.dump(meta, open(os.path.join(out_dir, "meta.json"), "w"), indent=2) +print(json.dumps({"ok": True})) +''' + + +def _worker_env() -> dict: + return { + "PATH": os.environ.get("PATH", ""), + "HOME": os.environ.get("HOME", ""), + "HF_HOME": HF_HOME, + "HF_HUB_CACHE": os.path.join(HF_HOME, "hub"), + "HF_HUB_DISABLE_XET": "1", + "TMPDIR": os.path.join(REPO, ".tmp"), + "CUDA_VISIBLE_DEVICES": "", # a CPU measurement must not touch a card someone else is using + "OMP_NUM_THREADS": "8", + "TOKENIZERS_PARALLELISM": "false", + "PYTHONPATH": REPO, + } + + +def _run(cmd: list[str], timeout: float) -> str: + p = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout, env=_worker_env(), cwd=os.path.join(REPO, ".tmp")) + if p.returncode != 0: + raise RuntimeError(f"{cmd[:3]} exited {p.returncode}:\n{p.stderr[-3000:]}") + return p.stdout + + +def run(ckpt: str, out: str, dtype: str, timeout: float = 3600.0) -> dict: + """Tokenize (legacy) -> align worker -> HF worker -> compare -> report.json.""" + # Absolute: the workers run with cwd=.tmp, and a relative --out silently lands under .tmp/.tmp. + out, ckpt = os.path.abspath(out), os.path.abspath(ckpt) + os.makedirs(out, exist_ok=True) + ids_dir = os.path.join(out, "ids") + _run([LEGACY_PY, "-c", _LEGACY_SCRIPT, "tokenize", ckpt, json.dumps(list(PROMPTS)), ids_dir, dtype, REPO], 300) + ids_path = os.path.join(ids_dir, "ids.json") + a_dir, b_dir = os.path.join(out, f"align_{dtype}"), os.path.join(out, f"hf_{dtype}") + # SEQUENTIAL by construction: two fp32 copies are 55GB. Each worker frees its memory by exiting. + _run([sys.executable, os.path.abspath(__file__), "worker-align", "--ckpt", ckpt, "--ids", ids_path, "--out", a_dir, "--dtype", dtype], timeout) + _run([LEGACY_PY, "-c", _LEGACY_SCRIPT, "hf", ckpt, ids_path, b_dir, dtype, REPO], timeout) + rep = compare_dirs(a_dir, b_dir, dtype) + rep["checkpoint"] = ckpt + with open(os.path.join(out, f"report_{dtype}.json"), "w") as f: + json.dump(rep, f, indent=2) + return rep + + +def run_mutants(ckpt: str, out: str, names: list[str] | None = None, timeout: float = 3600.0) -> dict: + """Re-run ONLY the align side per mutant (fp32) and compare to the already-saved + transformers tensors of `run`. Returns {name: max_abs_logit_diff, ...} plus the verdict + that each mutant FAILS the fp32 gate.""" + out, ckpt = os.path.abspath(out), os.path.abspath(ckpt) + ids_path = os.path.join(out, "ids", "ids.json") + b_dir = os.path.join(out, "hf_fp32") + if not (os.path.isfile(ids_path) and os.path.isfile(os.path.join(b_dir, "meta.json"))): + raise FileNotFoundError("run `olmoe_parity.py run --dtype fp32` first: the mutants compare to its transformers tensors") + res = {} + for name in names or sorted(MUTANTS): + a_dir = os.path.join(out, f"align_fp32_mut_{name}") + _run([sys.executable, os.path.abspath(__file__), "worker-align", "--ckpt", ckpt, "--ids", ids_path, + "--out", a_dir, "--dtype", "fp32", "--mutant", name], timeout) + rep = compare_dirs(a_dir, b_dir, "fp32", allow_mutant=True) + res[name] = { + "what": MUTANTS[name], + "max_abs_logit_diff": rep["metrics"]["max_abs_logit_diff"], + "top1_agreement": rep["metrics"]["top1_agreement"], + "min_router_set_agreement": rep["metrics"]["min_router_set_agreement"], + "gate_passed": rep["verdict"]["passed"], + } + with open(os.path.join(out, "mutants_fp32.json"), "w") as f: + json.dump(res, f, indent=2) + return res + + +def _summary(rep: dict) -> str: + m, v = rep["metrics"], rep["verdict"] + lines = [ + f"dtype={v['dtype']} gated={v['gated']} passed={v['passed']} checks={v['checks']}", + f"max_abs_logit_diff={m['max_abs_logit_diff']:.3e} (gate {'< %g' % v['thresholds']['max_abs_logit_diff'] if v['gated'] else 'none: bf16 informational'}) " + f"top1={m['top1_agreement']:.4f} tokens={m['tokens']} min_router_set_agreement={m['min_router_set_agreement']:.5f}", + ] + for p in m["per_prompt"]: + lines.append(f" prompt: tokens={p['tokens']} max_abs={p['max_abs']:.3e} mean_abs={p['mean_abs']:.3e} |logit|max={p['abs_max_logit']:.2f} top1={p['top1_agree']}/{p['tokens']}") + for x in m["router_layers"]: + lines.append(f" layer {x['layer']:2d}: set_agreement={x['set_agreement']:.5f} min_overlap={x['min_overlap']} router_logit_max_abs={x['router_logit_max_abs']:.3e}") + return "\n".join(lines) + + +def main(argv=None) -> int: + ap = argparse.ArgumentParser(description=__doc__.split("\n")[0]) + sub = ap.add_subparsers(dest="cmd", required=True) + r = sub.add_parser("run") + r.add_argument("--ckpt") + r.add_argument("--out", default=DEFAULT_OUT) + r.add_argument("--dtype", choices=("fp32", "bf16"), default="fp32") + r.add_argument("--timeout", type=float, default=3600.0) + w = sub.add_parser("worker-align") + w.add_argument("--ckpt", required=True) + w.add_argument("--ids", required=True) + w.add_argument("--out", required=True) + w.add_argument("--dtype", choices=("fp32", "bf16"), required=True) + w.add_argument("--mutant", choices=sorted(MUTANTS)) + mu = sub.add_parser("mutants") + mu.add_argument("--ckpt") + mu.add_argument("--out", default=DEFAULT_OUT) + mu.add_argument("--only", nargs="*", choices=sorted(MUTANTS)) + c = sub.add_parser("compare") + c.add_argument("--a", required=True) + c.add_argument("--b", required=True) + c.add_argument("--dtype", choices=("fp32", "bf16"), required=True) + a = ap.parse_args(argv) + if a.cmd == "worker-align": + worker_align(a.ckpt, a.ids, a.out, a.dtype, a.mutant) + return 0 + if a.cmd == "mutants": + res = run_mutants(a.ckpt or default_ckpt(), a.out, a.only) + for k, v in res.items(): + print(f"{k:20s} max_abs={v['max_abs_logit_diff']:.3e} top1={v['top1_agreement']:.3f} router>={v['min_router_set_agreement']:.3f} gate_passed={v['gate_passed']} [{v['what']}]") + return 1 if any(v["gate_passed"] for v in res.values()) else 0 + if a.cmd == "compare": + rep = compare_dirs(a.a, a.b, a.dtype) + else: + rep = run(a.ckpt or default_ckpt(), a.out, a.dtype, a.timeout) + print(_summary(rep)) + return 1 if rep["verdict"]["passed"] is False else 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/code/scripts/olmoe_sft_smoke.py b/code/scripts/olmoe_sft_smoke.py new file mode 100644 index 0000000000000000000000000000000000000000..babd5caa7964d4b63046ed7f06a7eb8b0a2fa6ff --- /dev/null +++ b/code/scripts/olmoe_sft_smoke.py @@ -0,0 +1,1186 @@ +#!/usr/bin/env python +"""OLMoE-1B-7B-0125 post-training smoke on 4x A10G: FSDP2, bf16, adamw8bit, MoE telemetry. + + .venv/bin/python scripts/olmoe_sft_smoke.py prepare # CPU: cached SFT rows -> OLMoE token ids + .venv/bin/python scripts/olmoe_sft_smoke.py drive # reference run + SIGTERM/resume run + report + torchrun --standalone --nproc_per_node=4 scripts/olmoe_sft_smoke.py train --config configs/olmoe_sft_smoke.toml --run-dir RUN + torchrun --standalone --nproc_per_node=4 scripts/olmoe_sft_smoke.py probe --run-dir RUN # FSDP cost vs tokens + +What this settles (docs/MEASUREMENTS.md, "Track M: OLMoE-1B-7B-0125"): + * does the real 6.92B model fit and step under FSDP2 + adamw8bit on 23GB cards, and with what headroom; + * tokens/s and step time, and how much of a micro-step is FSDP traffic (the "is EP needed?" gate); + * does the LBL-off post-training default leave the router balanced (MaxVio / utilization per layer); + * the checkpoint-save host-RSS peak (the 45GB-floor question) for a 13.8GB bf16 weight file; + * does a SIGTERM'd, supervisor-relaunched run continue where it stopped. + +DATA. The repo's cached SFT corpus is the PRM800K math-solution shards built for Mistral +(`nano-mistral-o1/data/shards/sft/train`, 11,739 rows, Mistral-v0.3 token ids). OLMoE has a different tokenizer +(50,279 BPE + specials), so `prepare` DECODES a seeded slice with the vendored Mistral tokenizer back to text, +splits it at the stored `prompt_len`, and re-tokenizes prompt and response with the OLMoE HF tokenizer in the +legacy venv (a subprocess; `transformers`/`tokenizers` are not align deps). The label mask is the prompt/response +boundary, an EOS closes each row. This is real SFT data, re-tokenized -- NOT synthetic. Rows are kept only if +their OLMoE length is in [MIN_LEN, seq_len], and every micro-batch is padded to exactly seq_len so each +micro-step costs the same 1024 tokens of compute (padded tokens run through the MoE; router telemetry masks them). + +IMPORT ORDER is load-bearing (apply_resource_guards sets env torch reads at import): this module imports only the +stdlib at top level; torch and every align module that pulls it in are imported inside functions, after the guard. + +LOADING FOUR 6.9B RANKS WITHOUT A HOST OOM. Each rank materialises the bf16 model on CPU (13.8GB) plus +`hf_to_align`'s stacked copy (another 13.8GB at the peak) before moving it to its GPU. Four ranks at once is +~110GB of a 186GB no-swap box -- the 2026-08-12 shape. Ranks therefore load ONE AT A TIME (barrier-ordered), and the +Trainer's resume reads go through `SerialWeightFormat` for the same reason. +""" + +from __future__ import annotations + +import argparse +import gc +import glob +import json +import math +import os +import re +import signal +import subprocess +import sys +import threading +import time + +REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +LEGACY_PY = "/apps/development/rlhf/nano-mistral-o1/.venv/bin/python" +HF_HOME = "/apps/development/rlhf/hf_cache" +CKPT_GLOB = os.path.join(HF_HOME, "hub", "models--allenai--OLMoE-1B-7B-0125", "snapshots", "*") +SFT_SHARDS = "/apps/development/rlhf/nano-mistral-o1/data/shards/sft/train" +OUT_ROOT = os.path.join(REPO, ".tmp", "olmoe") +DATA_DIR = os.path.join(OUT_ROOT, "sft_data") +SMOKE_RUNS = os.path.join(OUT_ROOT, "smoke") +DEFAULT_CONFIG = os.path.join(REPO, "configs", "olmoe_sft_smoke.toml") + +SEQ_LEN = 1024 +MIN_LEN = 640 +N_ROWS = 320 # 20 optimizer steps x accum 4 x world 4 x micro 1 +WORLD = 4 +GRAD_ACCUM = 4 +N_OPT_STEPS = 20 +SAVE_AT_STEP = 10 +SAVE_EVERY_MICRO = SAVE_AT_STEP * GRAD_ACCUM +PAD_ID = 1 # OLMoE pad_token_id +EOS_ID = 50279 + +# --------------------------------------------------------------------------- +# prepare: cached Mistral-tokenized SFT rows -> OLMoE-tokenized rows +# --------------------------------------------------------------------------- + +_LEGACY_TOKENIZE = r''' +import json, sys +import numpy as np +from safetensors.numpy import save_file +from transformers import AutoTokenizer +ckpt, rows_path, out_path, seq_len, min_len, n_rows, eos = sys.argv[1:8] +seq_len, min_len, n_rows, eos = int(seq_len), int(min_len), int(n_rows), int(eos) +tok = AutoTokenizer.from_pretrained(ckpt) +ids_all, offs, plens, keep = [], [0], [], [] +rejected = {"short": 0, "long": 0} +for r in json.load(open(rows_path)): + p = tok(r["prompt"], add_special_tokens=False)["input_ids"] + q = tok(r["response"], add_special_tokens=False)["input_ids"] + [eos] + n = len(p) + len(q) + if n > seq_len: + rejected["long"] += 1; continue + if n < min_len: + rejected["short"] += 1; continue + ids_all.extend(p + q); offs.append(len(ids_all)); plens.append(len(p)); keep.append(r["src_index"]) + if len(plens) == n_rows: + break +if len(plens) < n_rows: + raise SystemExit(f"only {len(plens)} rows survived the length filter (wanted {n_rows}); rejected={rejected}") +save_file({"ids": np.asarray(ids_all, dtype=np.int32), "offsets": np.asarray(offs, dtype=np.int64), + "prompt_len": np.asarray(plens, dtype=np.int32), "src_index": np.asarray(keep, dtype=np.int64)}, out_path) +print(json.dumps({"rows": len(plens), "tokens": len(ids_all), "rejected": rejected, "vocab": len(tok)})) +''' + + +def default_ckpt() -> str: + hits = sorted(glob.glob(CKPT_GLOB)) + if not hits: + raise FileNotFoundError(f"no OLMoE-1B-7B-0125 snapshot under {CKPT_GLOB}") + return hits[-1] + + +def data_paths(data_dir: str = DATA_DIR) -> tuple[str, str]: + return os.path.join(data_dir, "sft_ids.safetensors"), os.path.join(data_dir, "sft_meta.json") + + +def split_response(full_text: str, prompt_text: str) -> str: + """The response text of a decoded row: `full_text` minus its `prompt_text` prefix. + + Decoding the response tokens ALONE loses a leading space (SentencePiece decodes the first + piece without its prefix marker), which would silently change the first response token's + OLMoE id. Decoding the whole row and cutting at the prompt's decoded length keeps it. If the + prompt text is not a prefix of the full text the decode did not round-trip and the row is + unusable: raise, never guess a boundary. + """ + if not full_text.startswith(prompt_text): + raise ValueError("decoded prompt is not a prefix of the decoded row; cannot place the prompt/response boundary") + return full_text[len(prompt_text):] + + +def prepare(ckpt: str, shard_dir: str, data_dir: str, n_rows: int, seed: int) -> dict: + import numpy as np + + sys.path.insert(0, REPO) + from align.data.store import ShardDataset + from align.tokenizer import Tokenizer + + os.makedirs(data_dir, exist_ok=True) + ds = ShardDataset(shard_dir) + lens = ds.lengths() + # Mistral length ~ OLMoE length within ~15%; over-select, the legacy step applies the exact filter. + cand = np.nonzero((lens >= int(MIN_LEN * 0.85)) & (lens <= SEQ_LEN))[0] + rng = np.random.default_rng(seed) + rng.shuffle(cand) + cand = cand[: int(n_rows * 2.5)] + mistral = Tokenizer.from_pretrained(os.path.join(REPO, "assets", "tokenizer", "mistral-7b-v0.3")) + rows = [] + for i in cand.tolist(): + r = ds[i] + pl = int(r["prompt_len"]) + ids = r.input_ids.tolist() + full, prompt = mistral.decode(ids), mistral.decode(ids[:pl]) + # Mistral's BOS/EOS decode as the literal text "" / ""; they are not OLMoE text. + full, prompt = full.removeprefix("").removesuffix(""), prompt.removeprefix("") + try: + resp = split_response(full, prompt) + except ValueError: + continue + rows.append({"src_index": i, "prompt": prompt, "response": resp}) + rows_path = os.path.join(data_dir, "rows_text.json") + with open(rows_path, "w") as f: + json.dump(rows, f) + out_path, meta_path = data_paths(data_dir) + env = {"PATH": os.environ.get("PATH", ""), "HOME": os.environ.get("HOME", ""), "HF_HOME": HF_HOME, + "HF_HUB_OFFLINE": "1", "TRANSFORMERS_OFFLINE": "1", "TOKENIZERS_PARALLELISM": "false", + "TMPDIR": os.path.join(REPO, ".tmp"), "OMP_NUM_THREADS": "4"} + p = subprocess.run([LEGACY_PY, "-c", _LEGACY_TOKENIZE, ckpt, rows_path, out_path, str(SEQ_LEN), str(MIN_LEN), str(n_rows), str(EOS_ID)], + capture_output=True, text=True, timeout=900, env=env, cwd=os.path.join(REPO, ".tmp")) + if p.returncode != 0: + raise RuntimeError(f"tokenize failed:\n{p.stderr[-2000:]}") + meta = json.loads(p.stdout.strip().splitlines()[-1]) + meta.update({"source": shard_dir, "source_receipt": "nano-mistral-o1/data/sft.receipt.json (tasksource/PRM800K @ 547b1950)", + "seed": seed, "seq_len": SEQ_LEN, "min_len": MIN_LEN, "pad_id": PAD_ID, "eos_id": EOS_ID, + "synthetic": False, "retokenized_from": "mistral-7b-v0.3 -> OLMoE-1B-7B-0125 tokenizer"}) + with open(meta_path, "w") as f: + json.dump(meta, f, indent=2) + os.remove(rows_path) + return meta + + +# --------------------------------------------------------------------------- +# pure analysis helpers (numpy/stdlib only; unit-tested) +# --------------------------------------------------------------------------- + + +def read_jsonl(path: str) -> list[dict]: + """Rows of a jsonl file; a torn FINAL line (a kill mid-append) is dropped, any other bad line raises.""" + if not os.path.exists(path): + return [] + lines = [l for l in open(path).read().split("\n") if l.strip()] + out = [] + for i, l in enumerate(lines): + try: + out.append(json.loads(l)) + except json.JSONDecodeError: + if i != len(lines) - 1: + raise + return out + + +def dedupe_by_micro(rows: list[dict]) -> dict[int, dict]: + """micro_step -> row, the LAST occurrence winning (a relaunch may re-log the boundary it resumed at).""" + return {int(r["micro_step"]): r for r in rows} + + +def mean_loss_per_micro(per_rank: list[dict[int, dict]]) -> dict[int, float]: + """Mean over ranks of each rank's micro-batch loss, for micro-steps EVERY rank logged.""" + common = set.intersection(*[set(d) for d in per_rank]) if per_rank else set() + return {m: sum(d[m]["loss"] for d in per_rank) / len(per_rank) for m in sorted(common)} + + +def compare_loss_curves(ref: dict[int, float], sup: dict[int, float], stop_micro: int, floor_mult: float = 2.0, + floor_pad: float = 5e-3) -> dict: + """Does the resumed segment of `sup` (micro > stop_micro) stay as close to the uninterrupted `ref` + as the pre-stop segment (micro <= stop_micro), which is the SAME computation and so measures pure + run-to-run nondeterminism (NCCL order, SDPA backward, stochastic rounding). + + Verdict: resumed max |dloss| <= floor_mult * pre-stop max |dloss| + floor_pad. The pad keeps a + near-zero floor (a lucky run) from turning ordinary noise into a failure. No overlap between the segments + raises (a vacuous pass is the failure mode). + + WHAT THIS DOES NOT PROVE (reviewer finding, 2026-10-12). The loss-curve tolerance (1e-2) has never been shown + to move when the resume is WRONG: no mutant resume (zeroed 8-bit moments, base weights instead of the + checkpoint) was run on the GPUs. A lost optimizer state at lr 1e-5 over 36 micro-steps could plausibly stay + inside it. Provenance is therefore checked separately and exactly, by `check_resume_provenance` (the relaunch + loaded THE checkpoint at the stop micro-step, no history row was re-logged from zero) and `compare_batches` + (every rank saw the same row at every micro-step as the uninterrupted run). This function is the + continuity check on top of those, not a substitute for them. + """ + common = sorted(set(ref) & set(sup)) + pre = [abs(ref[m] - sup[m]) for m in common if m <= stop_micro] + post = [abs(ref[m] - sup[m]) for m in common if m > stop_micro] + if not pre or not post: + raise ValueError(f"need steps on both sides of the stop (stop={stop_micro}); pre={len(pre)} post={len(post)}") + floor, dev = max(pre), max(post) + tol = floor_mult * floor + floor_pad + return { + "stop_micro": stop_micro, "n_pre": len(pre), "n_post": len(post), + "pre_stop_max_abs": floor, "pre_stop_mean_abs": sum(pre) / len(pre), + "resumed_max_abs": dev, "resumed_mean_abs": sum(post) / len(post), + "tolerance": tol, "passed": dev <= tol, + } + + +def step_time_stats(rows: list[dict], accum: int, save_every_micro: int, drop_first: int = 2) -> dict: + """Micro/optimizer-step wall time from consecutive `t_wall` stamps of ONE rank's rows. + + Intervals dropped: the first `drop_first` (allocator warmup, first-use kernels) and any interval that + contains a checkpoint save (the micro-step after micro % save_every == 0 pays for the write). The + median is reported with the spread; a mean would be dominated by whatever a neighbour job did. + """ + rows = sorted(rows, key=lambda r: r["micro_step"]) + dts = [] + for a, b in zip(rows, rows[1:]): + if b["micro_step"] != a["micro_step"] + 1: + continue + if a["micro_step"] % save_every_micro == 0 or a["micro_step"] < drop_first: + continue + dts.append(b["t_wall"] - a["t_wall"]) + if not dts: + raise ValueError("no usable intervals") + dts.sort() + med = dts[len(dts) // 2] + return {"n": len(dts), "micro_median_s": med, "micro_min_s": dts[0], "micro_max_s": dts[-1], + "optimizer_step_s": med * accum} + + +# --------------------------------------------------------------------------- +# train / probe (torchrun targets) +# --------------------------------------------------------------------------- + + +def _vm_kb(field: str) -> float: + with open("/proc/self/status") as f: + for line in f: + if line.startswith(field + ":"): + return float(line.split()[1]) + return float("nan") + + +def _mem_available_gb() -> float: + with open("/proc/meminfo") as f: + for line in f: + if line.startswith("MemAvailable:"): + return float(line.split()[1]) / 1024**2 + return float("nan") + + +def trim_host_heap() -> bool: + """Give freed glibc heap pages back to the OS (`malloc_trim(0)`); True if anything was released. + + MEASURED NEED (docs/MEASUREMENTS.md, Track M): after the rank-serial load, ranks 1-3 each kept a 13.4GiB + anonymous `[heap]` -- the freed bf16 state dict (tensors under glibc's dynamic mmap threshold, 4-5MB per + expert, land in the brk heap, and a live small chunk above them stops the heap from shrinking). Three + ranks x 13.4GiB = 40GiB of dead host memory for the whole run, on a box with no swap, and rank 0 kept ~11GiB + after its first checkpoint save for the same reason. `malloc_trim` walks the arenas and `madvise`s the free + pages away. A no-op (False) off glibc. + """ + try: + import ctypes + + return bool(ctypes.CDLL("libc.so.6").malloc_trim(0)) + except (OSError, AttributeError): + return False + + +def _reset_hwm() -> bool: + """Reset this process's peak RSS (VmHWM) so a window's high-water mark is its own.""" + try: + with open("/proc/self/clear_refs", "w") as f: + f.write("5") + return True + except OSError: + return False + + +class RssWindow: + """Peak host memory over a window: own RSS (VmHWM after a reset, and a 0.1s sampler) and system + MemAvailable (min). MemAvailable counts clean page cache as available, so its dip is anonymous + memory the window really added -- the quantity the 45GB floor is about.""" + + def __init__(self, period: float = 0.1): + self.period = period + + def __enter__(self): + self.hwm_reset = _reset_hwm() + self.rss0 = _vm_kb("VmRSS") / 1024**2 + self.avail0 = _mem_available_gb() + self.rss_peak = self.rss0 + self.avail_min = self.avail0 + self._stop = threading.Event() + self.t0 = time.perf_counter() + + def loop(): + while not self._stop.is_set(): + self.rss_peak = max(self.rss_peak, _vm_kb("VmRSS") / 1024**2) + self.avail_min = min(self.avail_min, _mem_available_gb()) + self._stop.wait(self.period) + + self._th = threading.Thread(target=loop, daemon=True) + self._th.start() + return self + + def __exit__(self, *exc): + self._stop.set() + self._th.join() + self.dur_s = time.perf_counter() - self.t0 + self.hwm = _vm_kb("VmHWM") / 1024**2 + return False + + def record(self) -> dict: + return {"dur_s": round(self.dur_s, 2), "rss_start_gb": round(self.rss0, 2), "rss_peak_sampled_gb": round(self.rss_peak, 2), + "vmhwm_gb": round(self.hwm, 2), "hwm_reset": self.hwm_reset, + "mem_available_start_gb": round(self.avail0, 1), "mem_available_min_gb": round(self.avail_min, 1), + "mem_available_drop_gb": round(self.avail0 - self.avail_min, 1)} + + +def _load_dataset(data_dir: str): + from safetensors.numpy import load_file + + arrs = load_file(data_paths(data_dir)[0]) + return arrs["ids"], arrs["offsets"], arrs["prompt_len"] + + +def _make_sampler_cls(): + """Defined lazily: it needs torch, which is imported only after the resource guards.""" + import torch + + class RowBatch(list): + """A list of row indices (what the Trainer logs as `batch`) that also carries the collated + tensors the adapter and `MoETrainModule` read (`input_ids`, `labels`, `attention_mask`).""" + + class RankSampler: + """Micro-batch 1. Epoch order is a pure function of (seed, epoch) -- the contract the Trainer's + resumable stream relies on -- and rank r takes elements r, r+W, r+2W of it, so the W ranks of one + micro-step see W DISTINCT rows (dense `scripts/sft.py` does not shard; see the report).""" + + def __init__(self, data_dir, rank, world, seed, device, seq_len=SEQ_LEN, n_rows=None): + self.ids, self.offs, self.plen = _load_dataset(data_dir) + self.n = len(self.plen) if n_rows is None else n_rows + self.rank, self.world, self.seed, self.epoch = rank, world, seed, 0 + self.device, self.seq_len = device, seq_len + + def set_epoch(self, e): + self.epoch = int(e) + + def row(self, i): + return self.ids[self.offs[i]:self.offs[i + 1]], int(self.plen[i]) + + def __iter__(self): + g = torch.Generator().manual_seed(self.seed + self.epoch) + order = torch.randperm(self.n, generator=g).tolist() + usable = (self.n // self.world) * self.world + for j in range(self.rank, usable, self.world): + i = order[j] + x, p = self.row(i) + L = len(x) + ids = torch.full((1, self.seq_len), PAD_ID, dtype=torch.long) + ids[0, :L] = torch.from_numpy(x.astype("int64")) + labels = torch.full((1, self.seq_len), -100, dtype=torch.long) + labels[0, p:L] = ids[0, p:L] + mask = torch.zeros((1, self.seq_len), dtype=torch.long) + mask[0, :L] = 1 + b = RowBatch([i]) + b.input_ids, b.labels, b.attention_mask = ids.to(self.device), labels.to(self.device), mask.to(self.device) + b.n_valid = L + yield b + + def __len__(self): + return (self.n // self.world) + + return RankSampler, RowBatch + + +def _serial_format_cls(): + from align.train.ckpt_format import OlmoeFormat + + class SerialWeightFormat(OlmoeFormat): + """`OlmoeFormat` whose `load_weights` runs on ONE rank at a time. + + `Trainer.resume_from` has every rank read the full 13.8GB state and `hf_to_align` stacks a second + copy at its peak; four ranks at once is ~110GB. Barrier-ordering the reads bounds the transient to + one rank's, plus the full dicts the already-loaded ranks still hold until their local shard is cut. + """ + + def load_weights(self, path): + from align.utils.dist import barrier, rank, world_size + + out = None + for r in range(world_size()): + if rank() == r: + out = super().load_weights(path) + barrier() + return out + + return SerialWeightFormat + + +def build_sharded_model(ckpt: str, ocfg, backend: str, attn: str, device): + """Rank-serial load -> GPU -> activation checkpointing -> FSDP2, as `stage_smoke.build_stage_model` does + for dense, with the host-RAM serialisation above.""" + import torch + + from align.model import olmoe as om + from align.model import olmoe_checkpoint as oc + from align.model.checkpoint import load_hf_safetensors + from align.train import fsdp + from align.utils.dist import barrier, rank, world_size + + model = None + peaks = {} + for r in range(world_size()): + if rank() == r: + with RssWindow() as w: + sd = oc.hf_to_align(load_hf_safetensors(ckpt, dtype=torch.bfloat16), ocfg) + model = om.OlmoeForCausalLM.from_state_dict(ocfg, sd, attn_implementation=attn, moe_backend=backend) + del sd + model = model.to(device) + peaks = w.record() + gc.collect() + peaks["trimmed"] = trim_host_heap() + peaks["rss_after_trim_gb"] = round(_vm_kb("VmRSS") / 1024**2, 2) + barrier() + fsdp.apply_activation_checkpointing(model) + fsdp.apply_fsdp(model) + return model, peaks + + +def _write_json(path: str, obj) -> None: + tmp = path + ".partial" + with open(tmp, "w") as f: + json.dump(obj, f, indent=2, default=str) + f.flush() + os.fsync(f.fileno()) + os.replace(tmp, path) + + +def train_main(a) -> int: + from align.config import load_config + from align.utils.guard import apply_resource_guards + + cfg = load_config(a.config, list(a.set or [])) + apply_resource_guards(cpu_threads=8, heavy=True, label="olmoe-sft-smoke") + import torch + + from align.model.olmoe_checkpoint import read_olmoe_config + from align.train.moe_module import build_moe_train_module + from align.train.offload_adamw import make_optimizer + from align.train.sft_loss import make_sft_loss_adapter + from align.train.stage_gates import assert_distinct_devices + from align.train.trainer import MANDATORY_METRICS, Trainer + from align.utils.dist import barrier, init_distributed, local_rank, rank as get_rank, seed_everything, world_size as get_world + from align.utils.tracking import Tracker + + run_dir = os.path.abspath(a.run_dir) + os.makedirs(run_dir, exist_ok=True) + init_distributed(run_dir=run_dir) + assert_distinct_devices() + R, W = get_rank(), get_world() + dev = torch.device("cuda", local_rank()) + torch.cuda.set_device(dev) + with open(os.path.join(run_dir, f"rank{R}.pid"), "w") as f: + f.write(str(os.getpid())) + seed_everything(cfg.train.seed) + + ckpt = a.ckpt or default_ckpt() + ocfg = read_olmoe_config(ckpt) + model, load_peaks = build_sharded_model(ckpt, ocfg, a.moe_backend, "sdpa", dev) + torch.cuda.synchronize() + load_gpu_peak = torch.cuda.max_memory_allocated() / 2**30 + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + + RankSampler, _ = _make_sampler_cls() + sampler = RankSampler(a.data_dir, R, W, cfg.train.seed, dev, n_rows=a.n_rows) + sft = make_sft_loss_adapter(require_fused=True) + + def adapter(m, batch): + return sft(m, (batch.input_ids, batch.labels)) + + opt = make_optimizer(cfg, [p for p in model.parameters() if p.requires_grad]) + module = build_moe_train_module(cfg, model, adapter, opt, stage="sft", olmoe_cfg=ocfg) + fmt = _serial_format_cls()(ocfg, tokenizer_dir=ckpt) + tracker = Tracker(run_dir if R == 0 else os.path.join(run_dir, f"rank{R}"), "sft", + mandatory_keys=MANDATORY_METRICS, run_name=cfg.run_name) + trainer = Trainer(cfg=cfg, model=model, sampler=sampler, loss_adapter=adapter, run_dir=run_dir, stage="sft", + train_module=module, weight_format=fmt, tracker=tracker) + + # -- telemetry: per-step memory/time, and the host-RSS window around every checkpoint save ------------- + orig_log = tracker.log + + def log(step, metrics=None, **kw): + torch.cuda.synchronize() + m = {**(dict(metrics) if metrics else {}), **kw} + m["t_wall"] = time.time() + m["mem/alloc_peak_gb"] = torch.cuda.max_memory_allocated() / 2**30 + m["mem/reserved_peak_gb"] = torch.cuda.max_memory_reserved() / 2**30 + m["rss_gb"] = _vm_kb("VmRSS") / 1024**2 + orig_log(step, m) + + tracker.log = log + orig_save = trainer.save_checkpoint + ckpt_log = os.path.join(run_dir, f"ckpt_rss_rank{R}.jsonl") + + def save(): + micro = trainer.micro_step + with RssWindow() as w: + out = orig_save() + gc.collect() + trimmed = trim_host_heap() # after the window: the window reports the save's own peak, untrimmed + rec = {"trimmed": trimmed, "rss_after_trim_gb": round(_vm_kb("VmRSS") / 1024**2, 2), "micro_step": micro, "optimizer_step": trainer.optimizer_step, "rank": R, **w.record()} + with open(ckpt_log, "a") as f: + f.write(json.dumps(rec) + "\n") + return out + + trainer.save_checkpoint = save + + resumed = trainer.resume_from_latest() + if resumed: + # MEASURED (smoke/sup/ckpt_rss_rank1.jsonl): after the resume read, ranks 1-3 sat at 13.7GB RSS until the + # first save trimmed them -- every rank materialises the full bf16 state dict, and the freed copy is + # stranded in glibc's brk heap exactly as after the initial load. Same fix, same reason. + gc.collect() + trim_host_heap() + _write_json(os.path.join(run_dir, f"launch_rank{R}_{int(time.time())}.json"), { + "rank": R, "resumed_from": resumed, "resumed_micro_step": trainer.micro_step, "load": load_peaks, + "load_gpu_peak_gb": load_gpu_peak, "gpu_alloc_after_wrap_gb": torch.cuda.memory_allocated() / 2**30, + "optimizer": getattr(trainer.optimizer, "banner_tag", type(trainer.optimizer).__name__), + "n_params": sum(p.numel() for p in model.parameters()), "moe_backend": a.moe_backend, + "lb_weight": module.lb_weight, "z_weight": module.z_weight, "world": W, "cfg_run": cfg.run_name, + }) + summary = trainer.train() + barrier() + if R == 0: + _write_json(os.path.join(run_dir, f"summary_{int(time.time())}.json"), {k: v for k, v in summary.items() if isinstance(v, (int, float, str, bool, type(None)))}) + if trainer.stop_exit_code: + print(f"{trainer.stop_reason} stop: checkpointed, exiting {trainer.stop_exit_code}", file=sys.stderr) + return trainer.stop_exit_code + return 0 + + +def _time_median(fn, n: int, warm: int = 2) -> float: + import torch + + for _ in range(warm): + fn() + torch.cuda.synchronize() + ts = [] + for _ in range(n): + t0 = time.perf_counter() + fn() + torch.cuda.synchronize() + ts.append(time.perf_counter() - t0) + ts.sort() + return ts[len(ts) // 2] + + +def probe_collectives(ocfg, dev, W: int) -> dict: + """Isolated FSDP2 traffic at real sizes, un-overlapped. NCCL algbw = full tensor bytes / time.""" + import torch + import torch.distributed as dist + + from align.utils.dist import barrier + + d, i, e = ocfg.hidden_size, ocfg.intermediate_size, ocfg.num_experts + block_params = 4 * d * d + 2 * d + 2 * d + e * d + e * 3 * d * i # attn + norms + router + experts + sweep = [] + for frac in (1 / 16, 1 / 4, 1.0): + n = int(block_params * frac) // W * W + ag_out = torch.empty(n, dtype=torch.bfloat16, device=dev) + ag_in = torch.randn(n // W, dtype=torch.bfloat16, device=dev) + rs_in = torch.randn(n, dtype=torch.float32, device=dev) + rs_out = torch.empty(n // W, dtype=torch.float32, device=dev) + barrier() + t_ag = _time_median(lambda: dist.all_gather_into_tensor(ag_out, ag_in), 10) + barrier() + t_rs = _time_median(lambda: dist.reduce_scatter_tensor(rs_out, rs_in), 10) + sweep.append({"params": n, "all_gather_bf16_s": t_ag, "reduce_scatter_fp32_s": t_rs, + "ag_algbw_GBps": n * 2 / t_ag / 1e9, "rs_algbw_GBps": n * 4 / t_rs / 1e9}) + del ag_out, ag_in, rs_in, rs_out + full = sweep[-1] + # Expert-parallel dispatch/combine: every token's k expert-slots travel as d_model bf16 activations; + # with EP=W each rank keeps 1/W of them locally. One all-to-all moves T*k*d*2 bytes per rank, and a + # layer issues 2 per forward, 2 in the activation-checkpoint recompute, 2 in the backward. + a2a = [] + for tokens in (1024, 4096): + n = tokens * ocfg.top_k * d + n -= n % W + x_in = torch.randn(n, dtype=torch.bfloat16, device=dev) + x_out = torch.empty_like(x_in) + barrier() + t = _time_median(lambda: dist.all_to_all_single(x_out, x_in), 10) + a2a.append({"tokens_per_rank": tokens, "bytes_per_rank": n * 2, "all_to_all_s": t, "algbw_GBps": n * 2 / t / 1e9, + "ep_comm_per_micro_s": ocfg.n_layers * 6 * t}) + del x_in, x_out + return {"block_params": block_params, "sweep": sweep, "all_to_all": a2a, "n_blocks": ocfg.n_layers, + "nccl_p2p_level": os.environ.get("NCCL_P2P_LEVEL"), + # AG fwd + AG in the activation-checkpoint recompute + RS, per block, fully exposed (no overlap). + "exposed_per_micro_s": ocfg.n_layers * (2 * full["all_gather_bf16_s"] + full["reduce_scatter_fp32_s"])} + + +def probe_compute(ocfg, dev, seq_lens: list[int], backend: str) -> dict: + """Compute-only fwd+bwd of a TRUNCATED random-init model per GPU, no FSDP, no collectives: t(n_layers, L) + = fixed + n_layers * per_layer, so the full 16-layer compute is extrapolated, not guessed.""" + import torch + + from align.model import olmoe as om + from align.train import fsdp + from align.train.sft_loss import make_sft_loss_adapter + + sft = make_sft_loss_adapter(require_fused=True) + out = [] + for n_layers in (2, 4): + m = om.OlmoeForCausalLM(ocfg.replace(n_layers=n_layers), attn_implementation="sdpa", moe_backend=backend).to(dev, torch.bfloat16) + fsdp.apply_activation_checkpointing(m) + for L in seq_lens: + x = torch.randint(2, ocfg.vocab_size - 32, (1, L), device=dev) + + def step(): + o = sft(m, (x, x)) + o.loss.backward() + m.zero_grad(set_to_none=True) + + try: + torch.cuda.reset_peak_memory_stats() + out.append({"n_layers": n_layers, "seq_len": L, "median_s": _time_median(step, 3, warm=1), + "alloc_peak_gb": torch.cuda.max_memory_allocated() / 2**30}) + except torch.OutOfMemoryError: + m.zero_grad(set_to_none=True) + torch.cuda.empty_cache() + out.append({"n_layers": n_layers, "seq_len": L, "oom": True}) + del m + torch.cuda.empty_cache() + res = {"points": out, "extrapolated_16_layer_s": {}} + for L in seq_lens: + a2 = next((p for p in out if p["n_layers"] == 2 and p["seq_len"] == L and "median_s" in p), None) + a4 = next((p for p in out if p["n_layers"] == 4 and p["seq_len"] == L and "median_s" in p), None) + if a2 and a4: + per_layer = (a4["median_s"] - a2["median_s"]) / 2 + fixed = a2["median_s"] - 2 * per_layer + res["extrapolated_16_layer_s"][str(L)] = {"per_layer_s": per_layer, "fixed_s": fixed, "total_16_s": fixed + 16 * per_layer} + return res + + +def probe_main(a) -> int: + """FSDP cost vs tokens. Three independent parts (`--only`): `collectives` (no model: NCCL AG/RS bandwidth at + the real block size), `compute` (truncated model per GPU, no FSDP: per-layer compute vs tokens), `scan` + (the real FSDP2 model: micro-step time at several L; FSDP's per-micro traffic is independent of L, so + t(L) = comm + b*L, and a flat t(L) means communication-bound).""" + from align.config import load_config + from align.utils.guard import apply_resource_guards + + cfg = load_config(a.config, list(a.set or [])) + apply_resource_guards(cpu_threads=8, heavy=True, label="olmoe-sft-probe") + import torch + + from align.model.olmoe_checkpoint import read_olmoe_config + from align.train.sft_loss import make_sft_loss_adapter + from align.train.stage_gates import assert_distinct_devices + from align.utils.dist import barrier, init_distributed, local_rank, rank as get_rank, seed_everything, world_size as get_world + + run_dir = os.path.abspath(a.run_dir) + os.makedirs(run_dir, exist_ok=True) + init_distributed(run_dir=run_dir) + assert_distinct_devices() + R, W = get_rank(), get_world() + dev = torch.device("cuda", local_rank()) + torch.cuda.set_device(dev) + seed_everything(0) + ckpt = a.ckpt or default_ckpt() + ocfg = read_olmoe_config(ckpt) + seq_lens = [int(x) for x in a.seq_lens.split(",")] + out = {"world": W, "moe_backend": a.moe_backend, "only": a.only} + if a.only in ("all", "collectives"): + out["collectives"] = probe_collectives(ocfg, dev, W) + barrier() + if a.only in ("all", "compute"): + out["compute"] = probe_compute(ocfg, dev, seq_lens, a.moe_backend) + barrier() + if a.only in ("all", "scan"): + model, _ = build_sharded_model(ckpt, ocfg, a.moe_backend, "sdpa", dev) + torch.cuda.empty_cache() + sft = make_sft_loss_adapter(require_fused=True) + ids_all, _, _ = _load_dataset(a.data_dir) + pts = [] + for L in seq_lens: + torch.cuda.reset_peak_memory_stats() + reps = [] + try: + for it in range(a.warmup + a.reps): + start = (it * 7919 + R * 104729) % (len(ids_all) - L) + x = torch.from_numpy(ids_all[start:start + L].astype("int64"))[None].to(dev) + torch.cuda.synchronize() + barrier() + t0 = time.perf_counter() + o = sft(model, (x, x)) + o.loss.backward() + torch.cuda.synchronize() + dt = time.perf_counter() - t0 + model.zero_grad(set_to_none=True) + if it >= a.warmup: + reps.append(dt) + reps.sort() + pt = {"seq_len": L, "micro_median_s": reps[len(reps) // 2], "micro_min_s": reps[0], "micro_max_s": reps[-1], + "alloc_peak_gb": torch.cuda.max_memory_allocated() / 2**30, "reserved_peak_gb": torch.cuda.max_memory_reserved() / 2**30} + except torch.OutOfMemoryError: + model.zero_grad(set_to_none=True) + torch.cuda.empty_cache() + pt = {"seq_len": L, "oom": True} + pts.append(pt) + barrier() + out["scan"] = pts + if a.only == "gradvar": + out["gradvar"] = probe_gradvar(a, cfg, ckpt, ocfg, dev, R, W) + if R == 0: + _write_json(os.path.join(run_dir, "probe.json"), out) + barrier() + return 0 + + +def probe_gradvar(a, cfg, ckpt, ocfg, dev, R: int, W: int) -> dict: + """Is the pre-clip global grad norm of ONE accumulation window reproducible? + + The training runs logged a step-1 pre-clip norm of 5.44, 8.09 and 11.65 in three runs with identical weights, + rows and seed (losses equal to 1e-3), converging to within ~5% from step 2. This repeats the exact step-1 + window (the first `grad_accum` micro-batches of this rank's sampler stream, attention mask included) several + times in ONE process on frozen weights and reports each global norm: if they differ, the nondeterminism is + in forward/backward itself (bf16 atomics, router near-ties), not in the optimizer or the first-step path.""" + import torch + import torch.distributed as dist + + from align.train.sft_loss import make_sft_loss_adapter + from align.utils.dist import barrier + + model, _ = build_sharded_model(ckpt, ocfg, a.moe_backend, "sdpa", dev) + trim_host_heap() + torch.cuda.empty_cache() + sft = make_sft_loss_adapter(require_fused=True) + RankSampler, _ = _make_sampler_cls() + it = iter(RankSampler(a.data_dir, R, W, cfg.train.seed, dev)) + for _ in range(a.skip_micro): + next(it) + window = [next(it) for _ in range(cfg.optimizer.grad_accum)] + + class WithMask: + def __init__(self, mask): + self.mask = mask + + def __call__(self, **kw): + return model(attention_mask=self.mask, **kw) + + norms, losses = [], [] + for rep in range(a.grad_reps): + tot = 0.0 + for b in window: + o = sft(WithMask(b.attention_mask), (b.input_ids, b.labels)) + (o.loss / len(window)).backward() + tot += float(o.loss.detach()) / len(window) + sq = torch.zeros((), dtype=torch.float64, device=dev) + for p in model.parameters(): + if p.grad is not None: + g = p.grad.to_local() if hasattr(p.grad, "to_local") else p.grad + sq += g.double().pow(2).sum() + dist.all_reduce(sq) + norms.append(float(sq.sqrt())) + losses.append(tot) + model.zero_grad(set_to_none=True) + barrier() + return {"rows": [b[0] for b in window], "grad_norms": norms, "rank_mean_losses": losses} + + +# --------------------------------------------------------------------------- +# drive: reference run, SIGTERM + supervisor resume run, report +# --------------------------------------------------------------------------- + + +def _gpu_used_mib() -> list[int]: + p = subprocess.run(["nvidia-smi", "--query-gpu=memory.used", "--format=csv,noheader,nounits"], capture_output=True, text=True, timeout=30) + return [int(x) for x in p.stdout.split()] + + +class HostMonitor: + """1 Hz sampler of GPU memory (nvidia-smi view: includes the CUDA context + NCCL), MemAvailable and load, + to /monitor.jsonl. The torch-allocator peaks in history.jsonl are the other half.""" + + def __init__(self, path: str, period: float = 1.0, pid_dir: str | None = None): + self.path, self.period, self.pid_dir = path, period, pid_dir + + def _rank_memory(self) -> list[dict]: + """Per-rank RSS split (anonymous / file-backed / shmem, GiB) from smaps_rollup: which part of a rank's + host footprint is heap that could be returned, and which is mapped files.""" + out = [] + for r in range(WORLD): + try: + pid = int(open(os.path.join(self.pid_dir, f"rank{r}.pid")).read()) + d = {} + for line in open(f"/proc/{pid}/smaps_rollup"): + k, _, v = line.partition(":") + if k in ("Rss", "Anonymous", "Shmem", "Pss_File", "Pss_Anon", "Pss_Shmem"): + d[k] = round(float(v.split()[0]) / 1024**2, 2) + out.append({"rank": r, "pid": pid, **d}) + except Exception: # noqa: BLE001 -- the process may be between attempts + pass + return out + + def __enter__(self): + self._stop = threading.Event() + self.peak = [0, 0, 0, 0] + self.min_avail = float("inf") + self.max_load = 0.0 + + def loop(): + with open(self.path, "a") as f: + while not self._stop.is_set(): + try: + g = _gpu_used_mib() + except Exception: + g = [] + av, ld = _mem_available_gb(), os.getloadavg()[0] + n_tick = getattr(self, "_n", 0) + self._n = n_tick + 1 + rm = self._rank_memory() if (self.pid_dir and n_tick % 10 == 0) else None + self.peak = [max(p, x) for p, x in zip(self.peak, g)] if g else self.peak + self.min_avail, self.max_load = min(self.min_avail, av), max(self.max_load, ld) + rec = {"t": time.time(), "gpu_mib": g, "mem_available_gb": round(av, 1), "load1": ld} + if rm: + rec["ranks"] = rm + f.write(json.dumps(rec) + "\n") + f.flush() + self._stop.wait(self.period) + + self._th = threading.Thread(target=loop, daemon=True) + self._th.start() + return self + + def __exit__(self, *exc): + self._stop.set() + self._th.join() + return False + + +def host_preflight_ok() -> tuple[bool, dict]: + sys.path.insert(0, REPO) + from align.utils.guard import host_state + + st = host_state() + used = _gpu_used_mib() + st["gpu_used_mib"] = used + ok = (st["ram_available_gb"] >= 45 and st["load1"] < 2 * st["n_cores"] and st["disk_in_flight"] < 64 + and st["d_state_procs"] <= 5 and len(used) >= WORLD and all(u < 2048 for u in used[:WORLD])) + return ok, st + + +def torchrun_cmd(sub: str, run_dir: str, extra: list[str] | None = None) -> list[str]: + return [os.path.join(REPO, ".venv", "bin", "torchrun"), "--standalone", f"--nproc_per_node={WORLD}", + os.path.abspath(__file__), sub, "--run-dir", run_dir] + (extra or []) + + +def _child_env() -> dict: + env = dict(os.environ) + for k in ("TMPDIR", "HF_HOME", "HF_HUB_CACHE", "XDG_CACHE_HOME", "TRITON_CACHE_DIR"): + env.pop(k, None) # let apply_resource_guards set the /apps defaults + env.update({"TMPDIR": os.path.join(REPO, ".tmp"), "PYTHONPATH": REPO, "OMP_NUM_THREADS": "8", + "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", "HF_HOME": HF_HOME, + "HF_HUB_DISABLE_XET": "1", "ALIGN_PG_TIMEOUT_S": "900", "NCCL_DEBUG": "WARN"}) + return env + + +def run_to_completion(cmd: list[str], run_dir: str, timeout: float, log_name: str) -> int: + """Run `cmd` in its own session; on timeout SIGTERM then SIGKILL ITS OWN process group (never by pattern).""" + os.makedirs(run_dir, exist_ok=True) + log = open(os.path.join(run_dir, log_name), "ab") + p = subprocess.Popen(cmd, env=_child_env(), stdout=log, stderr=subprocess.STDOUT, start_new_session=True, cwd=REPO) + t0 = time.time() + while p.poll() is None: + if time.time() - t0 > timeout: + os.killpg(p.pid, signal.SIGTERM) + try: + p.wait(60) + except subprocess.TimeoutExpired: + os.killpg(p.pid, signal.SIGKILL) + return 124 + time.sleep(1.0) + return p.returncode + + +def checkpoint_complete(run_dir: str, micro: int) -> bool: + sys.path.insert(0, REPO) + d = os.path.join(run_dir, "checkpoints", f"micro_{micro:08d}") + if not os.path.isdir(d): + return False + from align.train.trainer import checkpoint_problem + + return checkpoint_problem(d) is None + + +def run_supervised_with_sigterm(run_dir: str, sigterm_after_micro: int, timeout: float) -> dict: + """supervise.py -> torchrun. Once the step-10 checkpoint is complete on disk, SIGTERM ONE rank's pid (rank 1): + the control plane must stop every rank at the same optimizer-step boundary after a checkpoint, and the + supervisor must relaunch because progress was made and the exit is not a HALT.""" + os.makedirs(run_dir, exist_ok=True) + cmd = [os.path.join(REPO, ".venv", "bin", "python"), os.path.join(REPO, "scripts", "supervise.py"), "--run-dir", run_dir, + "--max-retries", "2", "--backoff", "10", "--"] + torchrun_cmd("train", run_dir, ["--config", DEFAULT_CONFIG]) + log = open(os.path.join(run_dir, "supervisor.log"), "ab") + p = subprocess.Popen(cmd, env=_child_env(), stdout=log, stderr=subprocess.STDOUT, start_new_session=True, cwd=REPO) + t0, sent = time.time(), None + while p.poll() is None: + if time.time() - t0 > timeout: + os.killpg(p.pid, signal.SIGTERM) + try: + p.wait(90) + except subprocess.TimeoutExpired: + os.killpg(p.pid, signal.SIGKILL) + return {"returncode": 124, "timed_out": True, "sigterm": sent} + if sent is None and checkpoint_complete(run_dir, sigterm_after_micro): + try: + pid = int(open(os.path.join(run_dir, "rank1.pid")).read()) + with open(f"/proc/{pid}/cmdline", "rb") as f: + assert b"olmoe_sft_smoke" in f.read() # the pid file is not stale + os.kill(pid, signal.SIGTERM) + sent = {"pid": pid, "t_since_start_s": round(time.time() - t0, 1), "after_checkpoint_micro": sigterm_after_micro, + "history_micro_at_send": (read_jsonl(os.path.join(run_dir, "history.jsonl")) or [{}])[-1].get("micro_step")} + except Exception as e: # noqa: BLE001 + sent = {"error": repr(e)} + time.sleep(0.2) + return {"returncode": p.returncode, "timed_out": False, "sigterm": sent, "elapsed_s": round(time.time() - t0, 1)} + + +def load_run(run_dir: str) -> dict: + per_rank = [dedupe_by_micro(read_jsonl(os.path.join(run_dir if r == 0 else os.path.join(run_dir, f"rank{r}"), "history.jsonl"))) + for r in range(WORLD)] + return {"per_rank": per_rank, "loss": mean_loss_per_micro(per_rank)} + + +def moe_curves(per_rank0: dict[int, dict], accum: int) -> dict: + """Per-layer MaxVio / utilization at the optimizer-step boundaries (the exact-window values).""" + steps = sorted(m for m in per_rank0 if m % accum == 0) + layers = sorted({int(k.split("/")[1][1:]) for k in per_rank0[steps[0]] if re.match(r"moe/l\d+/max_vio", k)}) if steps else [] + return { + "steps": steps, + "max_vio": {l: [per_rank0[m][f"moe/l{l}/max_vio"] for m in steps] for l in layers}, + "utilization": {l: [per_rank0[m][f"moe/l{l}/utilization"] for m in steps] for l in layers}, + "global": {k: [per_rank0[m][k] for m in steps] for k in ("moe/max_vio", "moe/max_vio_worst", "moe/utilization", "moe/utilization_worst", + "moe/frac_dead_worst", "moe/router_entropy", "moe/load_entropy")}, + } + + +def drive(a) -> int: + os.makedirs(SMOKE_RUNS, exist_ok=True) + ref_dir, sup_dir = os.path.join(SMOKE_RUNS, "ref"), os.path.join(SMOKE_RUNS, "sup") + if not os.path.isfile(data_paths(a.data_dir)[0]): + print("preparing data ...", flush=True) + print(prepare(default_ckpt(), SFT_SHARDS, a.data_dir, N_ROWS, 0), flush=True) + report = {"started": time.strftime("%Y-%m-%d %H:%M:%S"), "config": DEFAULT_CONFIG} + + ok, st = host_preflight_ok() + report["preflight_ref"] = st + if not ok: + print("preflight FAILED, not starting:", st, flush=True) + return 83 + if not a.skip_ref: + with HostMonitor(os.path.join(SMOKE_RUNS, "monitor_ref.jsonl"), pid_dir=ref_dir) as mon: + rc = run_to_completion(torchrun_cmd("train", ref_dir, ["--config", DEFAULT_CONFIG]), ref_dir, a.timeout, "launcher.log") + report["ref"] = {"returncode": rc, "gpu_nvidia_smi_peak_mib": mon.peak, "mem_available_min_gb": mon.min_avail, "load1_max": mon.max_load} + print("reference run exit", rc, flush=True) + if rc != 0: + _write_json(os.path.join(SMOKE_RUNS, "report.json"), report) + return rc + ok, st = host_preflight_ok() + report["preflight_sup"] = st + if not ok: + print("preflight FAILED before supervised run:", st, flush=True) + _write_json(os.path.join(SMOKE_RUNS, "report.json"), report) + return 83 + with HostMonitor(os.path.join(SMOKE_RUNS, "monitor_sup.jsonl"), pid_dir=sup_dir) as mon: + sup = run_supervised_with_sigterm(sup_dir, SAVE_EVERY_MICRO, a.timeout) + sup.update({"gpu_nvidia_smi_peak_mib": mon.peak, "mem_available_min_gb": mon.min_avail, "load1_max": mon.max_load}) + report["sup"] = sup + print("supervised run:", sup, flush=True) + _write_json(os.path.join(SMOKE_RUNS, "report.json"), report) + return 0 if sup.get("returncode") == 0 else 1 + + +def detect_stop_micro(attempts: list[dict]) -> int: + """The micro-step the interrupted attempt left its newest complete checkpoint at, from supervisor.jsonl. + + The first `exit` event with a non-zero child exit that `progressed`. A log with no failed attempt means the + run was never interrupted and a resume comparison would be vacuous: raise rather than pick a number.""" + for r in attempts: + if r.get("event") == "exit" and r.get("child_exit") not in (0, None): + if not r.get("progressed"): + raise ValueError(f"attempt {r.get('attempt')} failed without progress; the supervisor would not have relaunched") + return int(r["ckpt_micro_step"]) + raise ValueError("no failed attempt in the supervisor log: the run was never interrupted, so there is no resume to verify") + + +def _rank_dir(run_dir: str, r: int) -> str: + return run_dir if r == 0 else os.path.join(run_dir, f"rank{r}") + + +def load_launches(run_dir: str, world: int = WORLD) -> dict[int, list[dict]]: + """rank -> that rank's `launch_rank{r}_{unix_ts}.json` records, oldest first (one per attempt).""" + out: dict[int, list[dict]] = {} + for r in range(world): + files = sorted(glob.glob(os.path.join(run_dir, f"launch_rank{r}_*.json")), + key=lambda f: int(re.search(r"_(\d+)\.json$", f).group(1))) + out[r] = [json.load(open(f)) for f in files] + return out + + +def load_raw_rows(run_dir: str, world: int = WORLD) -> dict[int, list[dict]]: + """rank -> history rows in FILE order (no de-duplication: a restart-from-zero must stay visible).""" + return {r: read_jsonl(os.path.join(_rank_dir(run_dir, r), "history.jsonl")) for r in range(world)} + + +def check_resume_provenance(launches: dict[int, list[dict]], raw_rows: dict[int, list[dict]], stop_micro: int) -> dict: + """Was the second attempt a RESUME of the stop checkpoint, on every rank? Exact, not statistical. + + THE HOLE THIS CLOSES. `analyze` de-duplicates history by micro-step (last row wins) and then compares loss + curves. A relaunch that silently IGNORED the checkpoint and trained from micro 0 re-logs micro 1..80; the + de-duplicated curve is then a complete second run of the reference computation, the 'pre-stop' and + 'resumed' segments both equal the reference to run-to-run noise, and the verdict PASSES. A resume that + did not happen is the failure mode this smoke exists to catch, so it is checked directly: + + * every rank has >= 2 launch records; the first started fresh (`resumed_micro_step == 0`, no path), the + last loaded `micro_{stop:08d}` and continued from exactly `stop_micro`; + * each rank's raw history starts at micro 1 and advances by exactly 1 per row (a repeat is allowed only + at `stop_micro` itself: the boundary a relaunch may re-log). A row going back to a lower micro-step + is a restart. + """ + problems: list[str] = [] + want = f"micro_{stop_micro:08d}" + for r, ls in sorted(launches.items()): + if len(ls) < 2: + problems.append(f"rank {r}: {len(ls)} launch record(s); an interrupted-and-resumed run has >= 2") + continue + first, last = ls[0], ls[-1] + if first.get("resumed_micro_step") != 0 or first.get("resumed_from"): + problems.append(f"rank {r}: first attempt did not start fresh ({first.get('resumed_from')!r}, micro {first.get('resumed_micro_step')})") + if last.get("resumed_micro_step") != stop_micro: + problems.append(f"rank {r}: last attempt resumed at micro {last.get('resumed_micro_step')}, expected {stop_micro}") + if not str(last.get("resumed_from") or "").rstrip("/").endswith(want): + problems.append(f"rank {r}: last attempt loaded {last.get('resumed_from')!r}, expected a path ending in {want}") + for r, rows in sorted(raw_rows.items()): + seq = [int(x["micro_step"]) for x in rows] + if not seq: + problems.append(f"rank {r}: empty history") + continue + if seq[0] != 1: + problems.append(f"rank {r}: history starts at micro {seq[0]}, not 1") + for a_, b_ in zip(seq, seq[1:]): + if b_ == a_ + 1 or (b_ == a_ and a_ == stop_micro): + continue + problems.append(f"rank {r}: history goes micro {a_} -> {b_} (a restart or a gap, not a continuation)") + break + return {"ok": not problems, "problems": problems, "stop_micro": stop_micro} + + +def compare_batches(ref_rows: dict[int, list[dict]], sup_rows: dict[int, list[dict]]) -> dict: + """Did every rank see the SAME row at every micro-step in the interrupted run as in the uninterrupted one? + + The sampler is a pure function of (seed, epoch, rank) and the Trainer's stream is resumable, so a correct + resume reproduces the reference batch sequence EXACTLY; a data position that drifted (a skipped or replayed + window) shows up here as an integer mismatch, with none of the loss curve's noise tolerance. Micro-steps + present on one side only are mismatches too: a run that stopped early is not a match.""" + mismatches = [] + n = 0 + for r in sorted(ref_rows): + a = {int(x["micro_step"]): x.get("batch") for x in ref_rows[r]} + b = {int(x["micro_step"]): x.get("batch") for x in sup_rows.get(r, [])} + for m in sorted(set(a) | set(b)): + n += 1 + if a.get(m, "MISSING") != b.get(m, "MISSING"): + mismatches.append({"rank": r, "micro_step": m, "ref": a.get(m, "MISSING"), "sup": b.get(m, "MISSING")}) + return {"ok": not mismatches and n > 0, "compared": n, "mismatches": mismatches[:10], "n_mismatches": len(mismatches)} + + +def analyze(a) -> int: + root = getattr(a, "runs_dir", None) or SMOKE_RUNS + ref_dir, sup_dir = os.path.join(root, "ref"), os.path.join(root, "sup") + ref, sup = load_run(ref_dir), load_run(sup_dir) + sup_attempts = read_jsonl(os.path.join(sup_dir, "supervisor.jsonl")) + stop_micro = int(a.stop_micro) if a.stop_micro is not None else detect_stop_micro(sup_attempts) + cmp = compare_loss_curves(ref["loss"], sup["loss"], stop_micro) + prov = check_resume_provenance(load_launches(sup_dir), load_raw_rows(sup_dir), stop_micro) + batches = compare_batches(load_raw_rows(ref_dir), load_raw_rows(sup_dir)) + steps = {"ref": step_time_stats(list(ref["per_rank"][0].values()), GRAD_ACCUM, SAVE_EVERY_MICRO), + "sup": step_time_stats(list(sup["per_rank"][0].values()), GRAD_ACCUM, SAVE_EVERY_MICRO)} + out = {"loss_compare": cmp, "resume_provenance": prov, "batch_compare": batches, "step_time": steps, + "supervisor_log": sup_attempts, "moe": moe_curves(ref["per_rank"][0], GRAD_ACCUM)["global"]} + _write_json(os.path.join(root, "analysis.json"), out) + print(json.dumps({k: v for k, v in out.items() if k not in ("supervisor_log", "moe")}, indent=2)) + return 0 if (cmp["passed"] and prov["ok"] and batches["ok"]) else 1 + + +def main(argv=None) -> int: + ap = argparse.ArgumentParser(description=__doc__.split("\n")[0]) + sub = ap.add_subparsers(dest="cmd", required=True) + p = sub.add_parser("prepare") + p.add_argument("--ckpt") + p.add_argument("--shards", default=SFT_SHARDS) + p.add_argument("--data-dir", default=DATA_DIR) + p.add_argument("--n-rows", type=int, default=N_ROWS) + p.add_argument("--seed", type=int, default=0) + for name in ("train", "probe"): + t = sub.add_parser(name) + t.add_argument("--config", default=DEFAULT_CONFIG) + t.add_argument("--run-dir", required=True) + t.add_argument("--ckpt") + t.add_argument("--data-dir", default=DATA_DIR) + t.add_argument("--n-rows", type=int) + t.add_argument("--moe-backend", default="loop", choices=("loop", "grouped_mm")) + t.add_argument("--set", action="append", help="config override key=value (dotted)") + if name == "probe": + t.add_argument("--seq-lens", default="512,1024,2048,4096") + t.add_argument("--only", default="all", choices=("all", "collectives", "compute", "scan", "gradvar"), + help="collectives = no model load (NCCL bandwidth only); compute = truncated model, no FSDP; " + "scan = FSDP t(L); gradvar = repeat the step-1 accumulation window and compare global grad norms") + t.add_argument("--grad-reps", type=int, default=4) + t.add_argument("--skip-micro", type=int, default=0, + help="gradvar: start the window this many micro-steps into the epoch (with --ckpt /checkpoints/micro_N " + "this replays the window that followed checkpoint N)") + t.add_argument("--warmup", type=int, default=2) + t.add_argument("--reps", type=int, default=5) + d = sub.add_parser("drive") + d.add_argument("--data-dir", default=DATA_DIR) + d.add_argument("--timeout", type=float, default=2700.0) + d.add_argument("--skip-ref", action="store_true") + an = sub.add_parser("analyze") + an.add_argument("--stop-micro", type=int, help="default: read from /supervisor.jsonl") + an.add_argument("--runs-dir", help=f"directory holding ref/ and sup/ (default {SMOKE_RUNS})") + a = ap.parse_args(argv) + sys.path.insert(0, REPO) + if a.cmd == "prepare": + print(json.dumps(prepare(a.ckpt or default_ckpt(), a.shards, a.data_dir, a.n_rows, a.seed), indent=2)) + return 0 + if a.cmd == "train": + from align.exit_codes import run_with_exit_codes + + return run_with_exit_codes(lambda: train_main(a)) + if a.cmd == "probe": + return probe_main(a) + if a.cmd == "drive": + return drive(a) + return analyze(a) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/code/scripts/ppo.py b/code/scripts/ppo.py new file mode 100644 index 0000000000000000000000000000000000000000..adf4935a56d552587f87a6f425865cf9c5df3d6b --- /dev/null +++ b/code/scripts/ppo.py @@ -0,0 +1,482 @@ +#!/usr/bin/env python +"""Process-supervised PPO entry: gates -> worker -> banner -> train -> stage manifest. + + CUDA_VISIBLE_DEVICES=0,1,2 torchrun --nproc_per_node=3 scripts/ppo.py \ + configs/ppo_stage.toml + # tiny end-to-end smoke (3 prompts x 4 samples x 256 tokens x 4 iterations): + CUDA_VISIBLE_DEVICES=0,1,2 torchrun --nproc_per_node=3 scripts/ppo.py \ + configs/ppo_stage.toml --smoke + +Ported from mistral-prm-ppo `scripts/train_ppo.py` and re-seated on the align stack the way +scripts/grpo.py is: stage config -> PPOTrainModule -> the shared Trainer. What replaced what: + + peft LoRA -> align.ppo.lora (named adapters on the SAME resident base) + DeepSpeed ZeRO-3 -> FSDP2 (align.train.fsdp); ZeRO-3 + Adafactor silently did not train + HF generate rollout -> align.rollout.client over the file protocol (vLLM worker, own venv) + transformers config -> align.model.hf_config.read_config (own config.json reader) + torch DDP of LoRA -> not needed: FSDP shards the frozen base, adapters ride the same wrap + +Topology: N FSDP train ranks + one dedicated engine GPU (default 3, NMO_ROLLOUT_GPU). Rank 0 +owns all protocol I/O. Checkpoint saves ARE the weight sync (the engine reloads the latest +checkpoint), and a LoRA run saves MERGED weights (PPOTrainModule.full_model_state_dict), so the +engine loads an ordinary HF-layout dir. Save cadence is DERIVED from the rollout shape, never +typed: a save lands exactly at the end of every sync_every_iterations-th iteration. + +Two arms share this one script and differ only in the config's `arm` and reward table: + arm = "prm" process supervision (PRM adapter on the shared backbone scores each step) + arm = "orm" the MANDATORY control: identical pipeline, process term off. Without it nothing + is attributable to process supervision (Math-Shepherd's PRM-attributable delta + over its own ORM-PPO control is +1.7pp MATH, ~2/3 of the headline is plain RL). + +Launch gates, ordered by cost: config validation (struct mode, eta, lam_actor, cadence, canary) +-> ceiling receipt (0/100-perfect-policy trap) -> topology check -> eos/boundary derivation +(derived, never hardcoded) -> module construction refusals (reference, scorer) -> horizon +consistency (scheduler vs policy steps). +""" + +from __future__ import annotations + +import argparse +import json +import logging +import os +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from align.utils.guard import apply_resource_guards # noqa: E402 + +logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") +logger = logging.getLogger("ppo") + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument("config") + ap.add_argument("--smoke", action="store_true", help="tiny shapes, 4 iterations") + ap.add_argument("--band-file", default=None, help="override the configured band file") + ap.add_argument("--strict-host-check", default=True, action=argparse.BooleanOptionalAction) + ap.add_argument("--heavy-lock", default=True, action=argparse.BooleanOptionalAction) + args = ap.parse_args() + + apply_resource_guards( + cpu_threads=8, + heavy=args.heavy_lock, + label="ppo-stage", + strict_host_check=args.strict_host_check, + ) + + from align.utils.dist import init_distributed + from align.ppo.stage_config import load_stage_config as _early_stage_config + + # Run dir from the config (parsed before the process group: a bad TOML fails + # in milliseconds, and the NCCL flight-recorder dump lands under it). + init_distributed(run_dir=_early_stage_config(args.config, smoke=args.smoke).run_dir) + + from align.train.stage_gates import assert_distinct_devices + + assert_distinct_devices() + + import torch + import torch.distributed as dist + + from align.data.schema import build_prompt, extract_answer, register_answer_suffixes + from align.data.sources import DataSources, gsm8k_gold_answer, read_gsm8k_test + from align.data.tokenization import load_boundary_spec + from align.eval.verifier import OutcomeVerifier, VerifierConfig + from align.model import checkpoint as ckpt + from align.model.mistral import NmoMistralForCausalLM + from align.ppo.lora import inject_lora, mark_only_adapter_trainable + from align.ppo.stage_config import load_stage_config + from align.ppo.trainer import ( + AdapterProbeView, + PPOBatchSource, + PPOTrainModule, + sync_params_from_rank0, + ) + from align.ppo.value_head import PRMConfig, ProcessRewardScorer, ValueHead, load_prm_state + from align.rl.canary import TaskAccuracyCanary + from align.rollout.client import RolloutClient, spawn_worker, terminate_worker + from align.rollout.protocol import RolloutRequest + from align.train.dpo_stage import resolve_sft_checkpoint + from align.train.trainer import Trainer, WarmupCosineSchedule + from align.utils.guard import assert_paths_off_root + + rank = dist.get_rank() if dist.is_initialized() else 0 + world = dist.get_world_size() if dist.is_initialized() else 1 + + cfg = load_stage_config(args.config, smoke=args.smoke) + if args.band_file: + cfg.data.band_file = args.band_file + assert_paths_off_root([cfg.run_dir]) + ceiling = cfg.validate_ceiling_receipt() + cfg.validate_world(world) + + register_answer_suffixes("") # BEFORE any extraction (golden.py's incident) + + # -- data: the band's prompt rows --------------------------------------- + prompt_rows = [] + with open(cfg.data.band_file) as f: + for line in f: + if line.strip(): + prompt_rows.append(json.loads(line)) + logger.info("band file %s: %d prompts", cfg.data.band_file, len(prompt_rows)) + + # -- derived ids: the boundary and EOS are READ, never hardcoded ---------- + spec = load_boundary_spec(cfg.boundary_spec_path) + sft_checkpoint, ckpt_source = resolve_sft_checkpoint(cfg.sft_checkpoint, None) + logger.info("substrate: %s (%s)", sft_checkpoint, ckpt_source) + from align.model.hf_config import read_config + from align.tokenizer import load_tokenizer + + model_cfg = read_config(sft_checkpoint) + tokenizer = load_tokenizer(sft_checkpoint) + pad_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0 + with open(os.path.join(sft_checkpoint, "config.json")) as f: + snapshot_eos = json.load(f).get("eos_token_id") + if not (cfg.ppo.eos_token_id == snapshot_eos == tokenizer.eos_token_id): + raise RuntimeError( + f"EOS ids disagree: ppo.eos_token_id={cfg.ppo.eos_token_id}, substrate " + f"config.json={snapshot_eos}, tokenizer={tokenizer.eos_token_id}. Truncation would be " + "mis-flagged: GAE would bootstrap terminal rows (or treat cut-off rows as terminal). " + "Derived-vs-hardcoded incident class." + ) + + # -- policy: LoRA on the shared base --------------------------------------- + model = NmoMistralForCausalLM.from_state_dict( + model_cfg, + ckpt.hf_to_nmo(ckpt.load_hf_safetensors(sft_checkpoint)), + attn_implementation=cfg.train.attn_implementation, + device="cuda", + ).to("cuda", dtype=torch.bfloat16) + + # Identical adapter init on every rank: FSDP2 shards what each rank already holds and does + # NOT broadcast, so a per-rank random A would be a different adapter per shard. + gen = torch.Generator().manual_seed(cfg.train.seed) + inject_lora(model, cfg.lora, adapter="policy", generator=gen) + scorer = None + if cfg.arm == "prm": + from align.ppo.lora import LoRAConfig + + inject_lora( + model, + LoRAConfig( + r=cfg.prm.lora_r, alpha=cfg.prm.lora_alpha, target_modules=cfg.lora.target_modules + ), + adapter="prm", + generator=gen, + ) + scorer = ProcessRewardScorer( + model, + model_cfg.hidden_size, + PRMConfig( + head_type=cfg.prm.head_type, + adapter="prm", + label_token_ids=cfg.prm.label_token_ids, + ), + ) + load_prm_state(scorer, cfg.prm.checkpoint) + scorer.to("cuda") + for p in scorer.parameters(): + p.requires_grad_(False) + n_trainable = mark_only_adapter_trainable(model, "policy") # AFTER both adapters exist + logger.info("policy adapter: %d trainable parameters", n_trainable) + + from align.train.fsdp import apply_activation_checkpointing, apply_fsdp + + model = apply_activation_checkpointing(model) + if world > 1: + model = apply_fsdp(model) + logger.info("post-shard cuda allocated: %.2f GiB", torch.cuda.memory_allocated() / 2**30) + + value_head = ValueHead( + model_cfg.hidden_size, detach_hidden=not cfg.train.value_shares_backbone + ).to("cuda") + sync_params_from_rank0(value_head) + + # -- rollout worker (rank 0 owns it) ------------------------------------------ + spool = os.path.join(cfg.run_dir, "spool") + worker_proc = None + if rank == 0: + worker_proc = spawn_worker( + spool, + gpu=int(os.environ.get("NMO_ROLLOUT_GPU", "3")), + log_path=os.path.join(cfg.run_dir, "worker.log"), + ) + logger.info("rollout worker spawned (pid %d)", worker_proc.pid) + client = RolloutClient(spool, timeout_s=cfg.rollout_timeout_s) + + # -- module + trainer ---------------------------------------------------------- + from align.train.offload_adamw import make_optimizer + + source = PPOBatchSource( + prompt_rows, + prompts_per_iteration=cfg.rollout.prompts_per_iteration, + n_samples_per_prompt=cfg.rollout.n_samples_per_prompt, + max_response_len=cfg.rollout.max_response_len, + temperature=cfg.rollout.temperature, + iterations=cfg.train.iterations, + ppo_cfg=cfg.ppo, + seed=cfg.train.seed, + prompt_builder=build_prompt, + world_size=world, + ) + run_cfg = _run_config(cfg, model_cfg, world, source.micros_per_iteration) + policy_params = [p for p in model.parameters() if p.requires_grad] + optimizer = make_optimizer(run_cfg, policy_params) + value_optimizer = torch.optim.AdamW( + value_head.parameters(), lr=cfg.train.lr * cfg.train.value_lr_mult, weight_decay=0.0 + ) + + def _sum_across_ranks(t): + if dist.is_initialized() and world > 1: + t = t.clone() + dist.all_reduce(t, op=dist.ReduceOp.SUM) + return t + + def _mean_across_ranks(g): + if dist.is_initialized() and world > 1: + g = g.clone() + dist.all_reduce(g, op=dist.ReduceOp.AVG) + return g + + score_fn = _make_score_fn(scorer, "cuda") if scorer is not None else None + module = PPOTrainModule( + model, + value_head, + optimizer, + value_optimizer, + cfg.ppo, + score_steps_fn=score_fn, + reduce_fn=_sum_across_ranks if world > 1 else None, + value_grad_reduce=_mean_across_ranks if world > 1 else None, + ) + # HORIZON: warmup iterations move only the critic, so the schedule spans POLICY steps. + scheduler = WarmupCosineSchedule( + base_lr=cfg.train.lr, + total_steps=source.policy_optimizer_steps, + warmup_frac=cfg.train.warmup_frac, + ) + trainer = Trainer( + run_cfg, + AdapterProbeView(model, "policy"), + source, + module._ppo_loss_adapter, + run_dir=cfg.run_dir, + stage="ppo", + train_module=module, + scheduler=scheduler, + ) + # The Trainer logs scheduler.lr() but (in this tree) does not write it into the optimizer; the + # module applies it at every step so the logged LR is the LR that ran. + module.lr_schedule = lambda: trainer.scheduler.lr() + if trainer.scheduler.total_steps != source.policy_optimizer_steps: + raise RuntimeError( + f"horizon mismatch: scheduler={trainer.scheduler.total_steps} vs policy optimizer " + f"steps={source.policy_optimizer_steps}" + ) + + # -- canary (R4, rank 0) ----------------------------------------------------------- + canary = None + if rank == 0: + sources = DataSources.from_local_hub() + gsm_rows = read_gsm8k_test(sources)[: cfg.canary.n_problems] + canary_problems = [ + {"prompt": build_prompt(r["question"]), "gold": gsm8k_gold_answer(r)} for r in gsm_rows + ] + + def canary_generate(prompts: list[str]) -> list[str]: + req = RolloutRequest( + request_id=f"canary-{trainer.optimizer_step}-{os.getpid()}", + weights_dir=_latest_weights(trainer, sft_checkpoint), + prompts=prompts, + n_samples_per_prompt=1, + max_tokens=cfg.rollout.max_response_len, + temperature=0.0, + ) + client.submit(req) + result = client.collect(req.request_id) + return [texts[0] for texts in result.texts] + + canary = TaskAccuracyCanary( + problems=canary_problems, + generate_fn=canary_generate, + verifier=OutcomeVerifier(VerifierConfig(workers=4)), + extract_fn=extract_answer, + every=cfg.canary.every, + ) + + source.bind( + client=client, + verifier=OutcomeVerifier(VerifierConfig(workers=4)) if rank == 0 else None, + module=module, + weights_dir_fn=lambda: _latest_weights(trainer, sft_checkpoint), + newline_id=spec.newline_id, + pad_token_id=pad_id, + rank=rank, + world_size=world, + trainer=trainer, + canary=canary, + canary_log_path=os.path.join(cfg.run_dir, "canary.jsonl"), + ) + trainer.loss_adapter = source.lazy_adapter(module) + + banner = { + "stage": f"ppo_{cfg.arm}" + ("_smoke" if args.smoke else ""), + "arm": cfg.arm, + "substrate": sft_checkpoint, + "band_file": cfg.data.band_file, + "band_prompts": len(prompt_rows), + "shape": { + "prompts_per_iteration": cfg.rollout.prompts_per_iteration, + "n_samples_per_prompt": cfg.rollout.n_samples_per_prompt, + "max_response_len": cfg.rollout.max_response_len, + "iterations": cfg.train.iterations, + "ppo_epochs": cfg.ppo.ppo_epochs, + "mini_batch_rows": cfg.ppo.mini_batch_rows, + "micro_batch_rows": cfg.ppo.micro_batch_rows, + "grad_accum": cfg.ppo.grad_accum, + "micros_per_iteration": source.micros_per_iteration, + "policy_optimizer_steps": source.policy_optimizer_steps, + "value_warmup_iterations": cfg.ppo.value_warmup_iterations, + }, + "reward": { + "transform": cfg.ppo.reward.transform, + "eta": cfg.ppo.reward.eta, + "w_prm": cfg.ppo.reward.w_prm, + "w_outcome": cfg.ppo.reward.w_outcome, + "kl_coef(reward)": cfg.ppo.reward.kl_coef, + "kl_loss_coef": cfg.ppo.kl_loss_coef, + "lam_actor": cfg.ppo.reward.lam_actor, + "lam_critic": cfg.ppo.reward.lam_critic, + }, + "lora": {"r": cfg.lora.r, "alpha": cfg.lora.alpha, "trainable": n_trainable}, + "ceiling": {k: v["verify_correct_rate"] for k, v in ceiling["sources"].items()}, + "boundary": {"newline_id": spec.newline_id, "tokenizer_hash": spec.tokenizer_hash}, + "topology": {"train_ranks": world, "rollout_gpu": os.environ.get("NMO_ROLLOUT_GPU", "3")}, + "rectifications": ( + f"R3 clip={cfg.ppo.max_grad_norm:.1f}, R4 canary every={cfg.canary.every}, " + "R5 pin-on-alarm, explicit minibatch denominators, ratio tripwire armed+self-tested" + ), + } + logger.info("BANNER %s", json.dumps(banner, indent=1, default=str)) + + # -- train ------------------------------------------------------------------------------ + try: + # RESUME. Without this a relaunch silently restarts from the SFT substrate with a fresh + # optimizer while the run dir holds the real checkpoints. Resume BEFORE the first + # rollout: the rollout engine's weights come from `_latest_weights`, which reads the + # pointer restored here. + resumed = _resume(trainer) + if resumed: + logger.info("RESUMED PPO run from %s", resumed) + summary = trainer.train() + checkpoint = trainer.save_checkpoint() + finally: + if worker_proc is not None: + terminate_worker(worker_proc) + if trainer.stop_exit_code: + # A coordinated stop (SIGTERM / train.max_wall_seconds) saved at an + # optimizer-step boundary; the stage is NOT finished. Exit with the + # documented status instead of evaluating / reporting a half-trained arm. + raise SystemExit(trainer.stop_exit_code) + + result = { + "stage": banner["stage"], + "summary": summary, + "checkpoint": checkpoint, + "banner": banner, + "canary_history": canary.history if canary else [], + "alarms": source.alarms, + "ratio_tripwire": module.ratio_tripwire.as_log_dict(), + "ratio_tripwire_alarms": module.ratio_tripwire.alarms, + } + if rank == 0: + path = os.path.join(cfg.run_dir, "stage_result.json") + with open(path + ".partial", "w") as f: + json.dump(result, f, indent=1, default=str) + os.replace(path + ".partial", path) + logger.info("stage result: %s", path) + return 0 + + +def _resume(trainer) -> str | None: + """Resume from the newest checkpoint, if any. Returns its path (None = fresh run). + + THE STALE-SUBSTRATE DEFECT (closed in `Trainer.resume_from`, pinned here). `checkpoints_written` + is a per-process list that `_latest_weights` keys off; it used to be empty after a resume, so + the first post-resume rollout was generated by the SFT SUBSTRATE -- the policy before every + update -- while the trainer computed ratios against the resumed policy. The Trainer now + registers the resumed checkpoint itself (generic: GRPO has the same `_latest_weights` shape and + keep-N pruning needed it too), so this wrapper no longer patches the list from the outside.""" + return trainer.resume_from_latest() + + +def _latest_weights(trainer, sft_checkpoint: str) -> str: + """The engine's weights source: latest saved (merged, HF-layout) checkpoint, else the + substrate. During value warmup the policy has not moved, so the substrate IS the policy.""" + if trainer.checkpoints_written: + return os.path.join(trainer.ckpt_root, f"micro_{trainer.checkpoints_written[-1]:08d}") + return sft_checkpoint + + +def _make_score_fn(scorer, device: str, chunk_rows: int = 8): + """PRM scores for a `PPORollout`: P(step correct) `[B, S]`, chunked so the scoring forward + stays inside the activation budget. The rollout stores delimiter ACTIONS in the label frame; + the PRM reads the hidden state AT the delimiter token, one input position later.""" + import torch + + @torch.no_grad() + def score(rollout): + B = rollout.n_rows + out = torch.zeros(rollout.step_positions.shape, dtype=torch.float32) + for lo in range(0, B, chunk_rows): + sl = slice(lo, min(B, lo + chunk_rows)) + T = int(rollout.attention_mask[sl].sum(-1).max()) + p = scorer.score_steps( + rollout.input_ids[sl, :T].to(device), + rollout.attention_mask[sl, :T].to(device), + (rollout.step_positions[sl] + 1).to(device), + rollout.step_mask[sl].to(device), + ) + out[sl] = p.float().cpu() + return out + + return score + + +def _run_config(cfg, model_cfg, world_size: int, micros_per_iteration: int): + """The Trainer's RunConfig, built FROM the stage config. Pure and import-light so the + config-translation test calls it directly and pins every crossing field (nmo silent-failure + #6: a field that fails to cross trains at the trainer's default while the banner shows the + pre-registered number).""" + from align.config import RunConfig + + run = RunConfig() + run.model = model_cfg + run.optimizer.name = cfg.train.optimizer_name + run.optimizer.lr = cfg.train.lr + run.optimizer.grad_accum = cfg.ppo.grad_accum + run.train.world_size = world_size + run.train.epochs = ( + 1 # the horizon is iterations x epochs; the source enumerates its own handles + ) + run.train.micro_batch_size = cfg.ppo.micro_batch_rows + run.train.save_every_micro_steps = micros_per_iteration * cfg.train.sync_every_iterations + run.train.keep_checkpoints = cfg.train.keep_checkpoints + run.train.warmup_frac = cfg.train.warmup_frac + run.train.full_ft = False # LoRA: the offload-economics accum>=4 rule is a full-FT rule + run.train.seed = cfg.train.seed + run.train.max_wall_seconds = cfg.train.max_wall_seconds + run.train.max_consecutive_skipped_steps = cfg.train.max_consecutive_skipped_steps + run.train.max_skipped_fraction = cfg.train.max_skipped_fraction + run.tripwires.armed = tuple(cfg.armed_tripwires) + run.run_name = f"ppo_{cfg.arm}" + return run + + +if __name__ == "__main__": + from align.exit_codes import run_with_exit_codes + + raise SystemExit(run_with_exit_codes(main)) diff --git a/code/scripts/precompute_ref_logprobs.py b/code/scripts/precompute_ref_logprobs.py new file mode 100644 index 0000000000000000000000000000000000000000..ae84ab13dc6767b4fed476d46b542e209e9e016e --- /dev/null +++ b/code/scripts/precompute_ref_logprobs.py @@ -0,0 +1,414 @@ +#!/usr/bin/env python +"""Precompute DPO reference log-probabilities from the frozen SFT checkpoint. + + # one no-grad pass per (arm, split), single GPU + uv run python scripts/precompute_ref_logprobs.py configs/dpo_argilla.toml + + # the MANDATORY path-bias dual pass (needs the sharded forward, so torchrun) + torchrun --nproc_per_node=4 scripts/precompute_ref_logprobs.py \ + configs/dpo_argilla.toml --path-bias + +Its own heavy job: it takes the exclusive flock and is SERIALIZED before any +DPO run. Precompute and training must never be concurrent -- a data job +running against a live training job is exactly what killed the predecessor's +first SFT run. + +What it produces, per arm, on /apps: + * `-train.safetensors` and `-held.safetensors` -- pair_id -> + (ref_logp_chosen, ref_logp_rejected, len_chosen, len_rejected), with the + full six-component cache key in the header. The held-out split is computed + in the SAME pass because held-out preference accuracy is a primary metric + and a cache missing it produces an accuracy over a silently smaller + denominator. + * `-path_bias.json` -- the single-GPU-vs-sharded delta distribution and + the decision derived from it. scripts/dpo.py refuses to launch without it, + and refuses to launch if it fired and was not honoured. + +Why the reference is cached at all: keeping a second 7B model resident would +cost 14.5GB of a 23GB card for no gradient. With dropout disabled a no-grad +pass over a fixed dataset is exactly reproducible, so the numbers can be +computed once and consumed as constants -- see nmo/train/ref_cache.py for the +six things the cache key has to cover for that to stay true. +""" + +from __future__ import annotations + +import argparse +import json +import logging +import os +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from align.utils.guard import apply_resource_guards # noqa: E402 + +logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") +logger = logging.getLogger("precompute_ref_logprobs") + +#: Pairs pushed through both forwards for the path-bias measurement. 100 is +#: the plan's number: enough to see a systematic offset against the +#: early-margin scale, small enough to run inside the pre-launch window. +PATH_BIAS_EXAMPLES = 100 + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument("config", help="configs/dpo_argilla.toml | dpo_steppairs.toml") + ap.add_argument( + "--out-dir", default=None, help="defaults to the config's ref_cache_path directory" + ) + ap.add_argument( + "--batch-size", + type=int, + default=8, + help="forward batch size. Regrouping perturbs the fp32 " + "reduction order by ~1e-7 relative (measured); it is " + "recorded in the receipt, not in the cache key.", + ) + ap.add_argument( + "--path-bias", + action="store_true", + help="run the dual single-GPU/sharded pass and write the " + "decision report. Requires torchrun.", + ) + ap.add_argument( + "--sharded-cache", + action="store_true", + help="compute the caches under the FSDP2-SHARDED forward (requires " + "torchrun). This is how a fired path-bias decision " + "(rerun_precompute_sharded) is honoured: training consumes the " + "sharded forward, so the reference must too.", + ) + ap.add_argument( + "--batch-invariance", + action="store_true", + help="measure batched-vs-single-row precompute deltas over " + "100 pairs and write the receipt that decides whether " + "batch_size belongs in the cache key. Single process.", + ) + ap.add_argument("--strict-host-check", default=True, action=argparse.BooleanOptionalAction) + ap.add_argument("--heavy-lock", default=True, action=argparse.BooleanOptionalAction) + args = ap.parse_args() + + apply_resource_guards( + cpu_threads=8, + heavy=args.heavy_lock, + label="precompute_ref_logprobs", + strict_host_check=args.strict_host_check, + ) + + import torch + + from align.data.loader import preference_dataset + from align.model import checkpoint as ckpt + from align.model.mistral import NmoMistralForCausalLM + from align.train.dpo_stage import ( + is_hf_checkpoint, + load_stage_config, + read_sft_quality_gate, + resolve_pair_split_dir, + resolve_sft_checkpoint, + ) + from align.train.ref_cache import ( + build_ref_cache_key, + path_bias_report_cache_forward_impl, + precompute_ref_logprobs, + write_path_bias_report, + write_ref_cache, + ) + from align.utils.guard import assert_paths_off_root + + cfg = load_stage_config(args.config) + out_dir = args.out_dir or os.path.dirname(cfg.ref_cache_path) + assert_paths_off_root([out_dir]) + + with open(cfg.manifest_path) as f: + manifest = json.load(f) + + # Resolve the frozen SFT checkpoint the same way scripts/dpo.py does. + # The trainer writes /checkpoints/micro_NNNNNNNN with no stable + # `final` alias; the quality-gate receipt records which one passed. A + # configured path that IS a checkpoint is an explicit pin; otherwise the + # receipt is read, and a missing/failed receipt or a missing checkpoint + # refuses LOUDLY here -- a cache silently computed from the wrong (or no) + # reference would poison both arms' margins. + if is_hf_checkpoint(cfg.sft_checkpoint): + sft_checkpoint, ckpt_source = cfg.sft_checkpoint, "config" + else: + gate = read_sft_quality_gate(cfg.sft_quality_gate_path) + sft_checkpoint, ckpt_source = resolve_sft_checkpoint( + cfg.sft_checkpoint, gate.get("checkpoint") + ) + logger.info("frozen SFT checkpoint: %s (source=%s)", sft_checkpoint, ckpt_source) + + # Shard split dirs by DISCOVERY (the receipts the store wrote), never by + # hand-composed `-` paths -- the attempt-1 defect class. + train_dir = resolve_pair_split_dir(cfg.shard_dir, f"{cfg.arm}-train") + held_dir = resolve_pair_split_dir(cfg.heldout_shard_dir, f"{cfg.arm}-held") + + # The key, assembled from facts rather than from intent, by the SAME + # function scripts/dpo.py's reader uses (build_ref_cache_key). `checkpoint_hash` + # is the WEIGHT-BYTES hash -- one streamed read of the checkpoint -- not the + # tensor-inventory hash this script used to write, which could not tell two + # SFT checkpoints of one architecture apart (2026-10-12 review). Caches + # written under the old key are refused by dpo.py: re-run this script. + # `manifest_hash` covers the shard rebuild that nothing else would notice. + key = build_ref_cache_key( + sft_checkpoint, + manifest, + forward_impl="fsdp2_sharded" if (args.path_bias or args.sharded_cache) else "single_gpu", + ) + logger.info("cache key: %s", key.to_json()) + + # Rank-pinned device: a bare "cuda" is device 0 for EVERY torchrun rank, + # and the --path-bias 4-rank pass put four ~5.5GB reference shards on + # GPU 0 (the duplicate-GPU class scripts/sft.py already fixed). This + # mode had never been run before 2026-08-18. + if torch.cuda.is_available(): + from align.utils.dist import init_distributed, local_rank + + init_distributed() + torch.cuda.set_device(local_rank()) + device = f"cuda:{local_rank()}" # string: downstream does device.startswith(...) + else: + device = "cpu" + model = _load_reference(NmoMistralForCausalLM, ckpt, sft_checkpoint, device) + + if args.sharded_cache: + import torch.distributed as dist + + from align.train.fsdp import apply_fsdp + + if not (dist.is_available() and dist.is_initialized()): + raise RuntimeError( + "--sharded-cache computes the reference under the SHARDED " + "forward and must run under torchrun; single-process would " + "silently produce the single_gpu numbers under a sharded key." + ) + model = apply_fsdp(model) + model.eval() + + written = [] + # In --path-bias mode the caches are NOT recomputed: the dedicated + # single-GPU run owns them. Before this guard, all four torchrun ranks + # each recomputed and rewrote the same cache files concurrently -- and + # stamped forward_impl="fsdp2_sharded" into caches computed single-GPU. + cache_splits = ( + () + if args.path_bias + else ( + ("train", train_dir, cfg.ref_cache_path), + ("held", held_dir, cfg.heldout_ref_cache_path), + ) + ) + for split, shard_dir, cache_path in cache_splits: + rows = preference_dataset(shard_dir) + cache = precompute_ref_logprobs( + model, + rows, + key=key, + pad_token_id=_pad_id(manifest), + batch_size=args.batch_size, + device=torch.device(device), + ) + from align.utils.dist import is_rank0 as _is_rank0 + + if _is_rank0(): + # Sharded forwards are collectives (every rank computes); the + # file is written once. + write_ref_cache(cache_path, cache) + written.append({"split": split, "path": cache_path, "pairs": len(cache)}) + logger.info("%s/%s: %d reference pairs", cfg.arm, split, len(cache)) + + if args.batch_invariance: + from align.train.ref_cache import batch_invariance_report + + rows = preference_dataset(train_dir)[:PATH_BIAS_EXAMPLES] + report = batch_invariance_report( + model, + rows, + key=key, + pad_token_id=_pad_id(manifest), + beta=cfg.beta, + early_margin_scale=cfg.path_bias_early_margin_scale, + threshold=cfg.path_bias_threshold, + ) + out_path = os.path.join(out_dir, f"{cfg.arm}-batch_invariance.json") + with open(out_path, "w") as f: + json.dump(report, f, indent=2) + logger.info( + "batch invariance -> %s (belongs_in_cache_key=%s)", + out_path, + report["belongs_in_cache_key"], + ) + + if args.path_bias: + # OWNERSHIP TRANSFER: pass the model in a box and drop main's name. + # A plain argument keeps main's reference alive, so the between-pass + # free inside _path_bias freed nothing and the reload OOM'd with the + # old 14.5GB replica still resident (measured 21.66GB, batch-size + # independent). + model_box = [model] + del model + report = _path_bias( + model_box, + cfg, + manifest, + key, + args.batch_size, + device, + train_dir, + reload_fn=lambda: _load_reference(NmoMistralForCausalLM, ckpt, sft_checkpoint, device), + ) + from align.utils.dist import is_rank0 + + if is_rank0(): + # Rank-0-only: four ranks renaming one report is the ledger race. + # `cache_forward_impl` must describe the CACHE, not this run. The + # old code wrote key.forward_impl, which is "fsdp2_sharded" for + # every --path-bias run -- but this mode recomputes no cache, so the + # field named the measuring run and contradicted the single-GPU + # cache sitting on disk (2026-10-12 review). Read it off the cache. + write_path_bias_report( + cfg.path_bias_report_path, + report, + n_examples=PATH_BIAS_EXAMPLES, + forward_impl=path_bias_report_cache_forward_impl(cfg.ref_cache_path), + ) + logger.info("path bias: %s (ratio %s)", report["action"], report["ratio"]) + + logger.info("precompute complete: %s", written) + return 0 + + +def _manifest_hash(manifest: dict) -> str: + """Hash of the build's own content hashes. Moves on any shard rebuild -- + which is the point: the section-03 marker fix retokenizes every pair while + leaving the checkpoint, tokenizer and max_length untouched.""" + from align.data.receipts import content_hash + + return content_hash([manifest["hashes"][k] for k in sorted(manifest.get("hashes", {}))]) + + +def _max_length(manifest: dict) -> int: + return int(manifest.get("pinned", {}).get("max_len", 2048)) + + +def _pad_id(manifest: dict) -> int: + """Pad id for collation. Fully masked out of every logprob (prompt and pad + positions carry IGNORE_INDEX), so its value cannot reach a number -- but + it is read from the build rather than assumed, because 'cannot reach a + number' is the kind of claim that stops being true after a refactor.""" + return int(manifest.get("pinned", {}).get("pad_token_id", 0)) + + +def _load_reference(model_cls, ckpt, path: str, device: str): + """The FROZEN SFT checkpoint, bf16 on GPU, eval mode. + + eval() is asserted downstream rather than merely called here: Mistral-v0.3 + has attention_dropout 0.0, so a model left in train mode would produce the + same numbers today and silently different ones after any config change. + """ + import torch + from align.eval.golden import resolve_model_config + + model_cfg = resolve_model_config(path) + model = model_cls.from_state_dict( + model_cfg, + ckpt.hf_to_nmo(ckpt.load_hf_safetensors(path)), + attn_implementation="sdpa", + device=device, + ) + model = model.to( + device, dtype=torch.bfloat16 if str(device).startswith("cuda") else torch.float32 + ) + model.eval() + for module in model.modules(): + if isinstance(module, torch.nn.Dropout) and module.p > 0: + raise RuntimeError( + f"an active Dropout(p={module.p}) is present in the reference " + "model. Cached reference logprobs must be exactly " + "reproducible; dropout makes them a sample." + ) + return model + + +def _path_bias( + model_box, cfg, manifest, key, batch_size: int, device: str, train_dir: str, reload_fn=None +) -> dict: + """The mandatory dual pass: single-GPU forward vs FSDP2-sharded forward. + + Requires torchrun. Without a process group there IS no sharded forward to + compare against, and the honest report is then no report -- `path_bias_ + decision` treats an empty delta set as firing, because not measuring is + not passing and a stage that skipped this must not silently proceed on a + single-GPU cache. + """ + import torch + import torch.distributed as dist + + from align.data.loader import preference_dataset + from align.train.fsdp import apply_fsdp + from align.train.ref_cache import path_bias_decision, precompute_ref_logprobs + + if not (dist.is_available() and dist.is_initialized()): + raise RuntimeError( + "--path-bias needs the SHARDED forward, so it must run under " + "torchrun (--nproc_per_node=4). Running it single-process would " + "compare the single-GPU forward with itself and record a " + "reassuring zero, which is worse than not measuring: it would " + "look like evidence." + ) + + rows = preference_dataset(train_dir)[:PATH_BIAS_EXAMPLES] + model = model_box.pop() # sole reference from here on + single = precompute_ref_logprobs( + model, + rows, + key=key, + pad_token_id=_pad_id(manifest), + batch_size=batch_size, + device=torch.device(device), + ) + # FREE the single-pass replica before sharding: wrapping a resident + # 14.5GB model in-place stacks the wrap transients ON TOP of it and + # OOMs every rank at 21.8GB (measured, 2026-08-18). Rebuilding alone + # and then sharding is the stage launcher's proven peak. + if reload_fn is None: + raise RuntimeError("--path-bias needs reload_fn to rebuild the model between passes") + del model + import gc + + gc.collect() + torch.cuda.empty_cache() + model = reload_fn() + sharded_model = apply_fsdp(model) + sharded_model.eval() + sharded = precompute_ref_logprobs( + sharded_model, + rows, + key=key, + pad_token_id=_pad_id(manifest), + batch_size=batch_size, + device=torch.device(device), + ) + + # The delta of the reference MARGIN, not of either logprob: the margin is + # what enters the loss, and a bias common to both sides cancels there. + deltas = [ + (sharded.entries[p][0] - sharded.entries[p][1]) + - (single.entries[p][0] - single.entries[p][1]) + for p in single.entries + ] + return path_bias_decision( + deltas, + beta=cfg.beta, + early_margin_scale=cfg.path_bias_early_margin_scale, + threshold=cfg.path_bias_threshold, + ) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/code/scripts/pretrain.py b/code/scripts/pretrain.py new file mode 100644 index 0000000000000000000000000000000000000000..b51f2a84e3d03930fff6220b4d2ba60a65f22225 --- /dev/null +++ b/code/scripts/pretrain.py @@ -0,0 +1,193 @@ +#!/usr/bin/env python +"""Pretraining stage entry (section-14): gates -> wrap -> train -> checkpoint. + + torchrun --nproc_per_node=4 scripts/pretrain.py configs/pretrain.toml --optimizer.lr=3e-5 + +The Pretraining pillar, OLMoE-shaped and honestly toy-scaled: continued +pretraining of the base model on a capped (~50M token), STREAMED slice of a +raw-text corpus, packed into the repo's shard format by +align/data/pretrain_corpus.py. The objective is the SAME packed next-token NLL +the SFT stage computes -- every pretrain row stores prompt_len 0, so +`collate_sft` labels every token (labels == input_ids, no masking) and +`make_sft_loss_adapter` runs unchanged. One loss path, two stages. + +Thin by design, mirroring scripts/sft.py: everything this drives is built and +tested elsewhere; what lives here is ORDER: + + build_currency -> smoke_7b -> probe_20step -> rss_preflight -> train + +Differences from the SFT entry, each deliberate: + * NO golden-regression bridge. The golden anchor certifies the SFT eval + stack (MATH-500/GSM8K extraction against known scores); a base-style + pretrain checkpoint has no such anchor. configs/pretrain.toml sets + gates.require_golden_receipt: false, which run_entry_gates honors with a + loud warning rather than silence. + * NO accuracy quality gate at exit. The stage's product is a checkpoint and + its train-NLL trajectory; downstream SFT is its consumer and its judge. + +IMPORT ORDER IS LOAD-BEARING. `apply_resource_guards()` sets +PYTORCH_CUDA_ALLOC_CONF, the /apps cache paths and the thread caps, and torch +reads all of them AT IMPORT. So this module imports only stdlib, the guard, +the config, and `stage_gates` (which is asserted torch-free); torch and every +align module that pulls it in are imported INSIDE main(), after the guards run. +tests/test_pretrain_entry.py enforces this by AST, exactly as the SFT test +does for scripts/sft.py. + +Resume semantics: rerunning after a kill picks up the gate ledger, skips gates +that already passed, and resumes the Trainer from the latest micro-step +checkpoint. Partial checkpoints are deleted at startup. +""" + +from __future__ import annotations + +import logging +import os +import sys +from pathlib import Path + +# The package is not installed in the venv; without this a bare `python scripts/pretrain.py` (the documented +# launch line) dies with ModuleNotFoundError unless PYTHONPATH is exported (same fix as scripts/sft.py). +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from align.config import load_config +from align.train.stage_gates import ( + GateLedger, + assert_build_is_current, + assert_probe_resolved, + eval_source_hash, + parse_args, + run_entry_gates, +) +from align.utils.guard import apply_resource_guards + +logger = logging.getLogger(__name__) + + +def build_model(cfg): + """Imported here, not at module scope: this pulls torch in.""" + from align.train.stage_smoke import build_stage_model + + return build_stage_model(cfg) + + +def run_gates(cfg, ledger): + """The entry sequence. Refuses to proceed on a stale or unstamped corpus + build before it touches a GPU. The golden bridge is intentionally absent + (see module docstring); run_entry_gates still calls require_golden_receipt + internally, which configs/pretrain.toml downgrades to a warning via + gates.require_golden_receipt: false.""" + assert_build_is_current(cfg.data.manifest_path, cfg) + banner = run_entry_gates(cfg, ledger) + assert_probe_resolved(banner) + return banner + + +def train(cfg, ledger, model, banner): + """One pretraining epoch through the section-07 Trainer. + + Identical wiring to the SFT stage on purpose: same sampler construction + (the SAME sampler the entry gates measured the horizon from), same + collate, same fused loss adapter. The pretraining objective difference + lives entirely in the data -- prompt_len 0 makes collate_sft emit + labels == input_ids.""" + from align.data.loader import collate_sft + from align.train.sft_loss import make_sft_loss_adapter + from align.train.stage_smoke import build_training_sampler + from align.train.trainer import Trainer + + sampler, dataset = build_training_sampler(cfg) + + def adapter(m, batch): + rows = [dataset[i] for i in batch] + collated = collate_sft(rows, pad_token_id=cfg.data.pad_token_id).to("cuda") + return make_sft_loss_adapter(require_fused=True)(m, (collated.input_ids, collated.labels)) + + trainer = Trainer( + cfg=cfg, + model=model, + sampler=sampler, + loss_adapter=adapter, + run_dir=os.path.join(cfg.paths.checkpoint_dir, "pretrain"), + stage="pretrain", + ) + trainer.resume_from_latest() + summary = trainer.train() + record_train_gate(ledger, trainer, summary) + return trainer + + +def record_train_gate(ledger, trainer, summary) -> dict: + """Record the 'train' gate for what `trainer.train()` ACTUALLY did (same contract as + scripts/sft.py::record_train_gate; duplicated because scripts/ is not a package). + + THE STOPPED-RECORDED-AS-PASSED DEFECT (H2 review, 2026-10-13; the sibling the first + fix missed). `train()` returns normally after a COORDINATED STOP (SIGTERM or + train.max_wall_seconds) and this used to record `passed=True` unconditionally, so the + ledger certified a half-trained checkpoint as a finished stage while `main` exited + 128+n / 82. A stopped run is resumable, neither finished nor failed; the ledger has no + third state, so it is `passed=False` with a `STOPPED (resumable)` detail, and the real + verdict overwrites it when the resumed run finishes. + """ + step = summary.get("optimizer_step") + if getattr(trainer, "stop_exit_code", 0): + return ledger.record( + "train", + passed=False, + detail=f"STOPPED (resumable): {trainer.stop_reason} stop at optimizer_step={step}; " + "checkpointed, not finished", + ) + return ledger.record("train", passed=True, detail=str(step)) + + +def main(argv=None) -> int: + args = parse_args(argv) + + # FIRST: every env var torch reads at import, plus the heavy-job flock and + # the /apps path assertions. `heavy=True` also runs the host preflight; + # `rss_cfg` adds the full-FT RSS projection on top of the generic floor. + cfg = load_config(args.config, args.overrides) + apply_resource_guards(cpu_threads=8, heavy=True, label="pretrain-stage-14", rss_cfg=cfg) + + logging.basicConfig( + level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s" + ) + + # BEFORE anything touches CUDA -- the duplicate-GPU defect (see + # scripts/sft.py): torchrun assigns no devices, so every rank's `.cuda()` + # lands on device 0 unless init_distributed sets one per local rank. + from align.utils.dist import init_distributed + + run_dir = os.path.join(cfg.paths.checkpoint_dir, "pretrain") + # run_dir: the NCCL flight-recorder dump lands under it (never /tmp). + init_distributed(run_dir=run_dir) + + from align.train.stage_gates import assert_distinct_devices + + assert_distinct_devices() + + ledger = GateLedger(run_dir, eval_source_hash=eval_source_hash()) + + banner = run_gates(cfg, ledger) + from align.registry import render_banner + + print(render_banner(banner, title=f"align pretrain stage :: {cfg.run_name}")) + + model = build_model(cfg) + trainer = train(cfg, ledger, model, banner) + if trainer.stop_exit_code: + # A coordinated stop (SIGTERM / train.max_wall_seconds) saved and + # returned: the stage is NOT finished, so do not fall into its + # post-train eval or announce completion. 82 = wall budget (the + # supervisor will not relaunch it), 128+n = the signal. + print(f"{trainer.stop_reason} stop: checkpointed, exiting {trainer.stop_exit_code}", file=sys.stderr) + return trainer.stop_exit_code + + checkpoint = trainer.latest_checkpoint() + print(f"Pretrain stage complete. SFT base checkpoint: {checkpoint}") + return 0 + + +if __name__ == "__main__": + from align.exit_codes import run_with_exit_codes + + raise SystemExit(run_with_exit_codes(main)) diff --git a/code/scripts/primary_contrast.py b/code/scripts/primary_contrast.py new file mode 100644 index 0000000000000000000000000000000000000000..8624ee625c9bef111739cba287c18eeec5b3149f --- /dev/null +++ b/code/scripts/primary_contrast.py @@ -0,0 +1,149 @@ +"""The pre-registered primary contrast: paired McNemar between two eval reports. + +Frozen pre-registration (docs/RESULTS.md): DPO Arm B result vs new-SFT, +GSM8K-full n=1,319, paired exact McNemar, alpha=0.05. This driver is pure +arithmetic over two reports' per_problem rows -- it computes what was +pre-registered and nothing else. The statistic itself is +`align.eval.metrics.mcnemar`, the SAME implementation the pre-registration +froze (exact binomial in the small-discordant regime); this script must +never reimplement it. + +Pairing is enforced, not assumed: rows join on problem_id within --source, +and a partial join REFUSES by default -- a silently dropped problem is a +silently different experiment. pass@1 is correct[0], matching how the +banked SFT report's headline numbers were computed. + + python scripts/primary_contrast.py \ + --arm outputs/eval/arm-b-final.json \ + --baseline outputs/eval/sft-final.json \ + --source gsm8k --expect-n 1319 \ + --label "dpo_arm_b_response vs sft-final" +""" + +import argparse +import json +import os +import sys +from pathlib import Path + +# Run from a checkout without an install: put the repo root on sys.path first. +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from align.eval.metrics import mcnemar, mcnemar_exact, wilson_interval + + +def _pass1_by_id(report: dict, source: str) -> dict[str, bool]: + out: dict[str, bool] = {} + for row in report["per_problem"]: + if row.get("source") != source: + continue + pid = row["problem_id"] + if pid in out: + raise ValueError(f"duplicate problem_id {pid!r} in one report; join is ill-defined") + out[pid] = bool(row["correct"][0]) + return out + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument("--arm", required=True, help="eval report JSON for the arm under test") + parser.add_argument("--baseline", required=True, help="eval report JSON for the baseline side") + parser.add_argument("--source", default="gsm8k") + parser.add_argument("--expect-n", type=int, default=1319, help="refuse if the paired n differs") + parser.add_argument("--alpha", type=float, default=0.05) + parser.add_argument("--label", required=True, help="human label recorded in the receipt") + parser.add_argument("--out", default=None, help="default: outputs/eval/primary_contrast.json") + parser.add_argument( + "--allow-partial", + action="store_true", + help="permit a partial join (still recorded loudly); NOT the pre-registered path", + ) + args = parser.parse_args() + + arm_report = json.load(open(args.arm)) + base_report = json.load(open(args.baseline)) + arm = _pass1_by_id(arm_report, args.source) + base = _pass1_by_id(base_report, args.source) + + shared = sorted(set(arm) & set(base)) + missing = len(set(arm) ^ set(base)) + if (missing or len(shared) != args.expect_n) and not args.allow_partial: + print( + f"REFUSED: paired n={len(shared)} (expected {args.expect_n}), " + f"{missing} unmatched problem_ids. The pre-registered contrast is " + "full-set paired; a partial join is a different experiment. " + "Pass --allow-partial only with a dated RESULTS.md entry explaining why.", + file=sys.stderr, + ) + return 2 + + # Discordants: b = arm right where baseline wrong; c = the reverse. + b = sum(1 for p in shared if arm[p] and not base[p]) + c = sum(1 for p in shared if base[p] and not arm[p]) + both_right = sum(1 for p in shared if arm[p] and base[p]) + both_wrong = len(shared) - b - c - both_right + # PRIMARY p-value: the exact binomial at any discordant count, because + # the frozen pre-registration's WORDING is "paired exact McNemar" while + # the frozen mcnemar() implementation switches to chi-square at 25+ + # discordants. Resolved by the 2026-08-22 dated RESULTS.md entry, decided + # BEFORE any arm data existed; mcnemar()'s output is recorded alongside. + stat_exact, p_exact = mcnemar_exact(b, c) + stat_frozen, p_frozen = mcnemar(b, c) + frozen_branch = "exact_binomial" if (b + c) < 25 else "chi_square_continuity_corrected" + + n = len(shared) + arm_k = b + both_right + base_k = c + both_right + receipt = { + "label": args.label, + "source": args.source, + "n_paired": n, + "unmatched": missing, + "arm": {"report": args.arm, "pass@1": arm_k / n, "ci95": wilson_interval(arm_k, n)}, + "baseline": { + "report": args.baseline, + "pass@1": base_k / n, + "ci95": wilson_interval(base_k, n), + }, + "table": { + "both_right": both_right, + "both_wrong": both_wrong, + "arm_only": b, + "baseline_only": c, + }, + "mcnemar": { + "method": "exact_binomial", + "statistic": stat_exact, + "p_two_sided": p_exact, + "alpha": args.alpha, + "frozen_impl": { + "statistic": stat_frozen, + "p_two_sided": p_frozen, + "branch": frozen_branch, + }, + }, + "significant": bool(p_exact < args.alpha), + "direction": "arm" if b > c else ("baseline" if c > b else "tie"), + } + + out_path = args.out or os.path.join("outputs", "eval", "primary_contrast.json") + os.makedirs(os.path.dirname(out_path), exist_ok=True) + tmp = out_path + ".partial" + with open(tmp, "w") as f: + json.dump(receipt, f, indent=1) + os.replace(tmp, out_path) + + print(f"{args.label} on {args.source}, n={n} paired") + print(f" arm pass@1 = {receipt['arm']['pass@1']:.4f} {args.arm}") + print(f" baseline pass@1 = {receipt['baseline']['pass@1']:.4f} {args.baseline}") + print(f" discordant: arm-only {b}, baseline-only {c} (concordant {both_right}+{both_wrong})") + print( + f" exact McNemar p = {p_exact:.5f} (alpha {args.alpha}) -> significant = " + f"{receipt['significant']} [frozen impl {frozen_branch}: p = {p_frozen:.5f}]" + ) + print(f"receipt written: {out_path}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/code/scripts/qwen3_export_check.py b/code/scripts/qwen3_export_check.py new file mode 100644 index 0000000000000000000000000000000000000000..c61e0eefc29c4c04d3b38e1735f244f8e93f08d4 --- /dev/null +++ b/code/scripts/qwen3_export_check.py @@ -0,0 +1,482 @@ +#!/usr/bin/env python +"""Merge a LoRA SFT checkpoint into the base, export HF Qwen3, and prove vLLM reads what align wrote. + + CUDA_VISIBLE_DEVICES=0,1 .venv/bin/python scripts/qwen3_export_check.py run \ + --ckpt RUN/checkpoints/micro_N --out-dir .tmp/p2/p2a/export [--keep] + +The claim under test: *the file align writes is the model align trained.* Five independent checks, +each of which fails if the export is wrong (the process exit status gates ALL of them, `export_ok`): + + 1. MERGE PATH vs TRAINER FILE. align merges the adapter (fp32 B@A * alpha/r, one bf16 cast) into + the pristine snapshot weights from the adapter stored in the checkpoint's optimizer payload, + and writes it with `save_hf_qwen3`. The trainer's own checkpoint weight file was built by a + different route (FSDP gather -> `export_merged_state_dict`). The two must be bit-identical. + 2. vLLM LOADS IT. A separate venv (`.venv-rollout`, vLLM 0.26) loads the exported directory + unmodified -- sharded safetensors, config.json (rope_theta under both spellings), tokenizer -- + and greedy-generates 64 tokens on 2 prompts, returning the top-20 logprobs at every position. + 3. FIRST-TOKEN LOGPROBS. align's own forward on the merged weights (bf16, GPU 1) scores the same + prompt; the first-token top-1 id must agree and the logprob difference is reported. + 4. TEACHER-FORCED AGREEMENT over all 64 generated tokens: align scores prompt+generation in one + pass; per-position argmax agreement with what vLLM sampled, and max |dlogprob| of the sampled + token. bf16 TOLERANCE (documented, not fitted): two correct bf16 engines with different + attention/GEMM kernels and accumulation order differ by a few ulps of a logit of magnitude + 10-30 (ulp 0.0625-0.125) amplified through 36 layers; the gate is |dlogprob| <= 0.25 on the + first token and argmax agreement >= 95% over the 64 positions (a near-tie legitimately flips + argmax; a wrong merge, a missing QK-norm or a wrong rope_theta disagrees on most positions and + by whole nats). The thresholds are constants below, asserted by `evaluate`, and exercised on + injected mutants in tests/test_qwen3_scripts.py. + 5. ADAPTER vs MERGED. The single bf16 cast of W+delta rounds away any delta below half an ulp of + W (a dead zone, not noise), so at small LoRA movement the merged file can be a different + policy from base+adapter -- and checks 1-4 all score the MERGED file against itself. Here the + SAME teacher-forced tokens are scored by align with the adapter applied live (base bf16 + + fp32 delta) and by the merged file; the gates are the first-token / argmax-agreement + constants of check 4 (reused, not refitted). The retained-delta cosine (merged-W vs the + intended delta) is recorded alongside. Also gated: the real snapshot's hf<->align mapping + round trip, and that the adapter MOVED every adapted matrix (a zero adapter exports the base). + +The export is DELETED afterwards unless --keep (16.4 GB on a 94%-full disk). +""" + +from __future__ import annotations + +import argparse +import json +import os +import shutil +import subprocess +import sys +import time + +REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, REPO) +ROLLOUT_PY = os.path.join(REPO, ".venv-rollout", "bin", "python") +SNAPSHOT = os.environ.get("P2_EXPORT_BASE_SNAPSHOT") or ("/apps/development/rlhf/hf_cache/hub/models--Goedel-LM--Goedel-Prover-V2-8B/snapshots/" + "dfd02e6271a58375dfbf3ece0175277cf6b6a89a") +WORK = os.path.join(REPO, ".tmp", "p2", "p2a", "export_check") + +N_NEW = 64 +TOPK = 20 +FIRST_TOKEN_MAX_ABS_DLOGPROB = 0.25 +MIN_ARGMAX_AGREEMENT = 0.95 + +PROMPTS = ( + "<|im_start|>user\nComplete the following Lean 4 code:\n\n```lean4\nimport Mathlib\nimport Aesop\n\n" + "set_option maxHeartbeats 400000\n\nopen BigOperators Real Nat Topology Rat\n\n" + "theorem mathd_algebra_10 : abs ((120 : ℝ) / 100 * 30 - 130 / 100 * 20) = 10 := by\n sorry\n```\n" + "<|im_end|>\n<|im_start|>assistant\n", + "<|im_start|>user\nWhat is 17 * 23? Answer briefly.<|im_end|>\n<|im_start|>assistant\n", +) + + +# --------------------------------------------------------------------------- +# merge (pure tensor math; unit-tested) +# --------------------------------------------------------------------------- + + +def merge_lora_state(base: dict, lora: dict, scaling: float) -> dict: + """`base` (align-keyed, `.weight`) with `scaling * B @ A` folded in for every adapter path + in `lora` (`.A` / `.B`, as `align.ppo.lora.lora_state_dict` writes them). + + The delta is accumulated in fp32 and cast ONCE to the weight's dtype (the same arithmetic as + `export_merged_state_dict`, so the two routes can be compared bit for bit). Every adapter path + must resolve to a base weight and must have both halves: a half-present adapter or a path that + matches nothing would otherwise "export" an unmodified base. + """ + import torch + + a_keys = {k[:-2] for k in lora if k.endswith(".A")} + b_keys = {k[:-2] for k in lora if k.endswith(".B")} + stray = sorted(k for k in lora if not (k.endswith(".A") or k.endswith(".B"))) + if stray: + raise KeyError(f"adapter state has keys that are neither .A nor .B: {stray[:4]}") + if a_keys != b_keys: + raise KeyError(f"adapter halves do not pair: A-only={sorted(a_keys - b_keys)[:3]} B-only={sorted(b_keys - a_keys)[:3]}") + if not a_keys: + raise KeyError("adapter state is empty; refusing to export an unmodified base as a merge") + out = dict(base) + for p in sorted(a_keys): + wk = f"{p}.weight" + if wk not in base: + raise KeyError(f"adapter path {p!r} has no base weight {wk!r}") + w, A, B = base[wk], lora[f"{p}.A"], lora[f"{p}.B"] + if A.shape[0] != B.shape[1] or (B.shape[0], A.shape[1]) != tuple(w.shape): + raise ValueError(f"{p}: A{tuple(A.shape)} B{tuple(B.shape)} do not compose to {tuple(w.shape)}") + out[wk] = (w.float() + (B.float() @ A.float()) * scaling).to(w.dtype) + return out + + +# --------------------------------------------------------------------------- +# verdict (pure; unit-tested) +# --------------------------------------------------------------------------- + + +def compare_generation(vllm: dict, align: dict) -> dict: + """Per-prompt metrics between vLLM's generation record and align's teacher-forced scoring. + + `vllm[i]`: {"token_ids": [64 ints], "chosen_logprobs": [64 floats], "first_top": {id: logprob}}. + `align[i]`: {"argmax": [64 ints], "chosen_logprobs": [64 floats], "first_logprobs": {id: logprob}} + (align's log-softmax at the same 64 positions, and at position 0 for vLLM's top-20 ids). + """ + per = [] + for v, a in zip(vllm, align): + n = len(v["token_ids"]) + if not (len(a["argmax"]) == len(a["chosen_logprobs"]) == n): + raise ValueError("generation records have different lengths") + agree = sum(int(x == y) for x, y in zip(v["token_ids"], a["argmax"])) + dchosen = [abs(x - y) for x, y in zip(v["chosen_logprobs"], a["chosen_logprobs"])] + v_top1 = max(v["first_top"], key=v["first_top"].get) + a_first = a["first_logprobs"] + per.append({ + "tokens": n, + "argmax_agreement": agree / n, + "max_abs_dlogprob_chosen": max(dchosen), + "mean_abs_dlogprob_chosen": sum(dchosen) / n, + "first_token_top1_same": int(v_top1) == int(max(a_first, key=a_first.get)), + "first_token_abs_dlogprob": abs(v["first_top"][v_top1] - a_first[v_top1]), + "first_token_max_abs_dlogprob_top20": max(abs(v["first_top"][k] - a_first[k]) for k in v["first_top"]), + }) + return {"per_prompt": per} + + +def evaluate(metrics: dict) -> dict: + checks = {} + for i, p in enumerate(metrics["per_prompt"]): + checks[f"prompt{i}.first_token_top1_same"] = bool(p["first_token_top1_same"]) + checks[f"prompt{i}.first_token_dlogprob"] = p["first_token_abs_dlogprob"] <= FIRST_TOKEN_MAX_ABS_DLOGPROB + checks[f"prompt{i}.argmax_agreement"] = p["argmax_agreement"] >= MIN_ARGMAX_AGREEMENT + return {"checks": checks, "passed": bool(checks) and all(checks.values()), + "thresholds": {"first_token_max_abs_dlogprob": FIRST_TOKEN_MAX_ABS_DLOGPROB, + "min_argmax_agreement": MIN_ARGMAX_AGREEMENT}} + + +def merge_fidelity_stats(base: dict, merged: dict, lora: dict, scaling: float) -> dict: + """How much of the intended delta `scaling * B @ A` survived the cast: aggregate cosine and + relative error of (merged - base) against the delta, and the fraction of nonzero delta + elements that rounded to exactly zero (the bf16 dead zone). Same arithmetic as + `align.ppo.lora.export_merged_state_dict`'s report.""" + import math + + import torch + + dot = ne = nd = 0.0 + nz_total = dropped = layers = 0 + for p in sorted({k[:-2] for k in lora if k.endswith(".A")}): + wk = f"{p}.weight" + d = ((lora[f"{p}.B"].float() @ lora[f"{p}.A"].float()) * scaling).flatten() + e = (merged[wk].float() - base[wk].float()).flatten() + dot, ne, nd = dot + float(torch.dot(e, d)), ne + float(torch.dot(e, e)), nd + float(torch.dot(d, d)) + nz = d != 0 + nz_total += int(nz.sum()) + dropped += int((nz & (e == 0)).sum()) + layers += 1 + if nd == 0.0: + return {"layers": layers, "cosine": 1.0, "rel_error": 0.0, "dropped_fraction": 0.0, "delta_norm": 0.0} + return {"layers": layers, "cosine": dot / math.sqrt(max(ne * nd, 1e-300)), + "rel_error": math.sqrt(max(ne - 2.0 * dot + nd, 0.0) / nd), + "dropped_fraction": dropped / nz_total if nz_total else 0.0, "delta_norm": math.sqrt(nd)} + + +def compare_adapter_vs_merged(adapter_lps: list, merged_lps: list) -> dict: + """Per-prompt agreement between align's log-softmax at the generated positions with the adapter + applied live (`adapter_lps[i]`, [T, V]) and on the merged bf16 file (`merged_lps[i]`, [T, V]).""" + import torch + + per = [] + for a, m in zip(adapter_lps, merged_lps): + if a.shape != m.shape or a.dim() != 2: + raise ValueError(f"log-prob records differ in shape: {tuple(a.shape)} vs {tuple(m.shape)}") + a, m = a.float(), m.float() + a_top, m_top = a.argmax(-1), m.argmax(-1) + t0 = int(a_top[0]) + per.append({ + "tokens": a.shape[0], + "argmax_agreement": float((a_top == m_top).float().mean()), + "first_token_abs_dlogprob": abs(float(a[0, t0] - m[0, t0])), + "max_abs_dlogprob_top1": float((a.gather(1, a_top[:, None]) - m.gather(1, a_top[:, None])).abs().max()), + "kl_adapter_to_merged_mean": float((a.exp() * (a - m)).sum(-1).mean()), + }) + return {"per_prompt": per} + + +def evaluate_adapter_vs_merged(metrics: dict) -> dict: + checks = {} + for i, p in enumerate(metrics["per_prompt"]): + checks[f"prompt{i}.first_token_dlogprob"] = p["first_token_abs_dlogprob"] <= FIRST_TOKEN_MAX_ABS_DLOGPROB + checks[f"prompt{i}.argmax_agreement"] = p["argmax_agreement"] >= MIN_ARGMAX_AGREEMENT + return {"checks": checks, "passed": bool(checks) and all(checks.values()), + "thresholds": {"first_token_max_abs_dlogprob": FIRST_TOKEN_MAX_ABS_DLOGPROB, + "min_argmax_agreement": MIN_ARGMAX_AGREEMENT}} + + +def export_ok(report: dict) -> bool: + """The process verdict. EVERY advertised check must have run and passed: a check that did not + run (key absent) is a failure, not a pass -- a zero adapter, a lossy hf<->align mapping or a + merge that dropped the delta must not exit 0.""" + rel_min = (report.get("merge_relative_weight_change") or {}).get("min") + return bool( + (report.get("verdict") or {}).get("passed") is True + and (report.get("merge_vs_trainer_file") or {}).get("bitwise_mismatches") == 0 + and (report.get("snapshot_mapping_roundtrip") or {}).get("bitwise_equal") is True + and (report.get("merge_changed_tensors_vs_snapshot") or 0) > 0 + and isinstance(rel_min, (int, float)) and rel_min > 0 + and ((report.get("adapter_vs_merged") or {}).get("verdict") or {}).get("passed") is True + ) + + +# --------------------------------------------------------------------------- +# vLLM side: runs in .venv-rollout (source string: not an import of this file) +# --------------------------------------------------------------------------- + +_VLLM_SCRIPT = r''' +import json, os, sys, time +model_dir, prompts_json, out_path, n_new, topk = sys.argv[1:6] +n_new, topk = int(n_new), int(topk) +os.environ.setdefault("VLLM_USE_FLASHINFER_SAMPLER", "0") +os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") +from transformers import AutoTokenizer +from vllm import LLM, SamplingParams +tok = AutoTokenizer.from_pretrained(model_dir) +prompts = json.loads(prompts_json) +ids = [tok(p, add_special_tokens=False)["input_ids"] for p in prompts] +t0 = time.time() +llm = LLM(model=model_dir, tokenizer=model_dir, dtype="bfloat16", max_model_len=2048, + gpu_memory_utilization=0.85, enforce_eager=True, seed=0) +t_load = time.time() - t0 +sp = SamplingParams(temperature=0.0, max_tokens=n_new, logprobs=topk, ignore_eos=True) +outs = llm.generate([{"prompt_token_ids": x} for x in ids], sp) +res = [] +for x, o in zip(ids, outs): + c = o.outputs[0] + lps = c.logprobs + res.append({ + "prompt_ids": x, + "token_ids": list(c.token_ids), + "text": c.text, + "chosen_logprobs": [float(lps[t][c.token_ids[t]].logprob) for t in range(len(c.token_ids))], + "first_top": {str(k): float(v.logprob) for k, v in lps[0].items()}, + }) +json.dump({"results": res, "load_s": t_load, "model": model_dir}, open(out_path, "w")) +print(json.dumps({"ok": True, "n": len(res), "load_s": round(t_load, 1)}), flush=True) +# vLLM 0.26's EngineCore subprocess can hang at interpreter shutdown (measured: the result file was +# written at 00:44 and the process was still alive 20 minutes later holding 19 GiB of GPU 0). The +# result is on disk and flushed; leave without running the atexit/engine teardown. +sys.stdout.flush(); sys.stderr.flush() +os._exit(0) +''' + + +def _first_visible_gpu(env=None) -> str: + """The child's CUDA_VISIBLE_DEVICES is relative to PHYSICAL indices, so a hard-coded "0" put + vLLM (0.85 of the card) on physical GPU 0 even when the operator pinned this check to 2,3. + Use the parent's FIRST visible entry ("0" when the parent is unpinned).""" + v = (os.environ if env is None else env).get("CUDA_VISIBLE_DEVICES", "") + first = v.split(",")[0].strip() + return first or "0" + + +def _rollout_env(gpu: str) -> dict: + scrub = ("RANK", "LOCAL_RANK", "WORLD_SIZE", "LOCAL_WORLD_SIZE", "MASTER_ADDR", "MASTER_PORT", "TORCHELASTIC") + env = {k: v for k, v in os.environ.items() if not k.startswith(scrub)} + env.update({ + "CUDA_VISIBLE_DEVICES": gpu, "VLLM_USE_FLASHINFER_SAMPLER": "0", + "TMPDIR": os.path.join(REPO, ".tmp", "p2", "p2a"), "XDG_CACHE_HOME": os.path.join(REPO, ".cache"), + "HF_HOME": "/apps/development/rlhf/hf_cache", "HF_HUB_CACHE": "/apps/development/rlhf/hf_cache/hub", + "HF_HUB_DISABLE_XET": "1", "HF_HUB_OFFLINE": "1", "TRANSFORMERS_OFFLINE": "1", + "TRITON_CACHE_DIR": os.path.join(REPO, ".cache", "triton"), + "TORCHINDUCTOR_CACHE_DIR": os.path.join(REPO, ".cache", "torchinductor"), + "PYTHONUNBUFFERED": "1", "TOKENIZERS_PARALLELISM": "false", + }) + return env + + +# --------------------------------------------------------------------------- +# the run +# --------------------------------------------------------------------------- + + +def run_vllm_child(cmd, result_path, *, env, cwd, log_path, grace_s: float = 30.0, timeout_s: float = 3000.0): + """Run the vLLM child WITHOUT inheriting pipes, and stop its process group once it has written + its result. + + DEFECT THIS PREVENTS (measured in the P2 smoke loop): vLLM 0.26's EngineCore can outlive its + parent at interpreter exit (re-parented to init, holding the parent's stdout/stderr pipes). With + `subprocess.run(capture_output=True)` the parent then blocks in communicate() until the ORPHAN + closes the pipe -- 12+ minutes, the result file already written, GPU 0 held at 19 GB, the heavy + lock held. Here the child logs to a file (nothing for an orphan to hold open), runs in its own + session so the whole group can be signalled by its exact pgid, and is terminated `grace_s` after + the result file appears. Returns (returncode or None if it had to be terminated, log tail). + """ + import signal + + with open(log_path, "w") as log: + proc = subprocess.Popen(cmd, stdout=log, stderr=subprocess.STDOUT, env=env, cwd=cwd, start_new_session=True) + t0 = time.time() + done_at = None + rc = None + while True: + rc = proc.poll() + if rc is not None: + break + if os.path.isfile(result_path) and done_at is None: + done_at = time.time() + if (done_at is not None and time.time() - done_at > grace_s) or time.time() - t0 > timeout_s: + break + time.sleep(1.0) + # signal the group even when the leader exited: orphaned workers keep the pgid + for sig in (signal.SIGTERM, signal.SIGKILL): + try: + os.killpg(proc.pid, sig) + except ProcessLookupError: + break + time.sleep(5.0) + if rc is None: + try: + rc = proc.wait(timeout=10) + except subprocess.TimeoutExpired: + rc = None + with open(log_path) as f: + return rc, f.read()[-3000:] + + +def _lora_scaling(ckpt: str, r: int | None, alpha: float | None) -> float: + if r is None or alpha is None: + raise SystemExit("--r and --alpha are required: the checkpoint stores A/B but not alpha/r, and guessing the scaling would export a wrong merge silently") + return alpha / r + + +def run(a) -> dict: + from align.utils.guard import apply_resource_guards + + apply_resource_guards(cpu_threads=8, heavy=True, label="p2a-qwen3-export-check") + import torch + + from align.model import qwen3 as q3 + from align.model.checkpoint import load_hf_safetensors + from align.model.qwen3_checkpoint import align_to_hf, hf_to_align, read_qwen3_config, save_hf_qwen3 + + work = os.path.abspath(a.work) + out_dir = os.path.abspath(a.out_dir) + os.makedirs(work, exist_ok=True) + report: dict = {"ckpt": a.ckpt, "out_dir": out_dir} + qcfg = read_qwen3_config(SNAPSHOT) + scaling = _lora_scaling(a.ckpt, a.r, a.alpha) + + # -- 1. merge path vs the trainer's file --------------------------------------------------- + t0 = time.time() + raw = load_hf_safetensors(SNAPSHOT, dtype=torch.bfloat16) + base = hf_to_align(raw, qcfg) + # real-checkpoint mapping round trip: hf -> align -> hf must return every tensor bit for bit + # (any tensor dropped, renamed or transposed by the table fails here, on the real shapes). + back = align_to_hf(base, qcfg) + report["snapshot_mapping_roundtrip"] = {"tensors": len(raw), "bitwise_equal": all(torch.equal(back[k], raw[k]) for k in raw) + and set(back) == set(raw)} + del raw, back + payload = torch.load(os.path.join(a.ckpt, "optimizer_rank0.pt"), weights_only=True, map_location="cpu") + merged = merge_lora_state(base, payload["lora"], scaling) + report["merge_delta_retention"] = merge_fidelity_stats(base, merged, payload["lora"], scaling) + del base + save_hf_qwen3(merged, out_dir, cfg=qcfg, dtype=torch.bfloat16, tokenizer_dir=SNAPSHOT) + report["merge_s"] = round(time.time() - t0, 1) + ckpt_state = hf_to_align(load_hf_safetensors(a.ckpt, dtype=None), qcfg) + mism = [k for k, v in merged.items() if not torch.equal(v, ckpt_state[k])] + n_changed = sum(1 for p in {k[:-2] for k in payload["lora"] if k.endswith(".A")}) + report["merge_vs_trainer_file"] = {"tensors": len(merged), "bitwise_mismatches": len(mism), "first_mismatches": mism[:3], + "adapted_layers": n_changed} + del ckpt_state + # the adapter must have MOVED the weights: a zero adapter exports the base and "passes" check 1 + base2 = hf_to_align(load_hf_safetensors(SNAPSHOT, dtype=torch.bfloat16), qcfg) + moved = sum(1 for k, v in merged.items() if not torch.equal(v, base2[k])) + report["merge_changed_tensors_vs_snapshot"] = moved + # How far the adapter moved the weights: ||W' - W||_F / ||W||_F per adapted matrix (the number + # an LR choice is judged by; bf16 rounding of W' contributes ~2^-9 ~ 0.2% and is included). + # Over the ADAPTED matrices only (every one must have moved: `export_ok` gates the minimum). + adapted = {f"{k[:-2]}.weight" for k in payload["lora"] if k.endswith(".A")} + rel = [((merged[k].float() - base2[k].float()).norm() / base2[k].float().norm()).item() for k in sorted(adapted)] + report["merge_relative_weight_change"] = {"n": len(rel), "mean": sum(rel) / len(rel), "max": max(rel), "min": min(rel)} + del base2, merged + + # -- 2. vLLM generates from the export ----------------------------------------------------- + vout = os.path.join(work, "vllm_out.json") + env = _rollout_env(_first_visible_gpu()) + if os.path.isfile(vout): + os.remove(vout) # a stale result from an earlier run must never be read as this run's + rc, tail = run_vllm_child([ROLLOUT_PY, "-c", _VLLM_SCRIPT, out_dir, json.dumps(list(PROMPTS)), vout, str(N_NEW), str(TOPK)], + vout, env=env, cwd=work, log_path=os.path.join(work, "vllm.log")) + report["vllm_returncode"] = rc + if not os.path.isfile(vout): + report["vllm_stderr_tail"] = tail + return report + vres = json.load(open(vout)) + report["vllm_load_s"] = vres["load_s"] + report["vllm_text"] = [r["text"][:200] for r in vres["results"]] + + # -- 3/4. align scores the same tokens on the merged weights (GPU 1) ------------------------ + dev = torch.device("cuda", 1 if torch.cuda.device_count() > 1 else 0) + cfg2, sd = (qcfg, hf_to_align(load_hf_safetensors(out_dir, dtype=torch.bfloat16), qcfg)) + model = q3.Qwen3ForCausalLM.from_state_dict(cfg2, sd, attn_implementation="sdpa", device="cuda").to(dev).eval() + del sd + arec, mlps = [], [] + with torch.no_grad(): + for r in vres["results"]: + ids = r["prompt_ids"] + r["token_ids"] + logits = model(input_ids=torch.tensor([ids], device=dev))[0].float() + lp = torch.log_softmax(logits, dim=-1) + n0 = len(r["prompt_ids"]) + pos = lp[n0 - 1: n0 - 1 + len(r["token_ids"])] # position t predicts generated token t + chosen = pos.gather(1, torch.tensor(r["token_ids"], device=dev)[:, None])[:, 0] + first = {k: float(pos[0, int(k)]) for k in r["first_top"]} + arec.append({"argmax": pos.argmax(-1).tolist(), "chosen_logprobs": chosen.tolist(), "first_logprobs": first}) + mlps.append(pos.cpu()) + metrics = compare_generation(vres["results"], arec) + report["metrics"] = metrics + report["verdict"] = evaluate(metrics) + + # -- 5. the same tokens through base + live adapter (fp32 delta) vs the merged file --------- + from align.ppo.lora import LoRAConfig, inject_lora, load_lora_state_dict + + del model + torch.cuda.empty_cache() + base_sd = hf_to_align(load_hf_safetensors(SNAPSHOT, dtype=torch.bfloat16), qcfg) + amodel = q3.Qwen3ForCausalLM.from_state_dict(qcfg, base_sd, attn_implementation="sdpa", device="cuda").to(dev).eval() + del base_sd + inject_lora(amodel, LoRAConfig(r=a.r, alpha=a.alpha), adapter="sft", generator=torch.Generator().manual_seed(0)) + load_lora_state_dict(amodel, "sft", payload["lora"], strict=True) + alps = [] + with torch.no_grad(): + for r in vres["results"]: + ids = r["prompt_ids"] + r["token_ids"] + lp = torch.log_softmax(amodel(input_ids=torch.tensor([ids], device=dev))[0].float(), dim=-1) + n0 = len(r["prompt_ids"]) + alps.append(lp[n0 - 1: n0 - 1 + len(r["token_ids"])].cpu()) + m5 = compare_adapter_vs_merged(alps, mlps) + report["adapter_vs_merged"] = {"metrics": m5, "verdict": evaluate_adapter_vs_merged(m5)} + return report + + +def main(argv=None) -> int: + ap = argparse.ArgumentParser(description=__doc__.split("\n")[0]) + sub = ap.add_subparsers(dest="cmd", required=True) + r = sub.add_parser("run") + r.add_argument("--ckpt", required=True, help="a trainer checkpoint dir (micro_N) with optimizer_rank0.pt") + r.add_argument("--out-dir", default=os.path.join(WORK, "export")) + r.add_argument("--work", default=WORK) + r.add_argument("--r", type=int, default=None) + r.add_argument("--alpha", type=float, default=None) + r.add_argument("--keep", action="store_true") + a = ap.parse_args(argv) + rep = {} + try: + rep = run(a) + finally: + if not a.keep and os.path.isdir(a.out_dir): + shutil.rmtree(a.out_dir, ignore_errors=True) + rep["export_deleted"] = True + os.makedirs(a.work, exist_ok=True) + with open(os.path.join(a.work, "report.json"), "w") as f: + json.dump(rep, f, indent=2, default=str) + print(json.dumps(rep, indent=2, default=str)) + return 0 if export_ok(rep) else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/code/scripts/qwen3_parity.py b/code/scripts/qwen3_parity.py new file mode 100644 index 0000000000000000000000000000000000000000..d104c624ebd21bb340d0508d5d7738ba27f308e6 --- /dev/null +++ b/code/scripts/qwen3_parity.py @@ -0,0 +1,483 @@ +#!/usr/bin/env python +"""Goedel-Prover-V2-8B (Qwen3ForCausalLM) parity on the REAL checkpoint: align vs transformers, CPU. + + .venv/bin/python scripts/qwen3_parity.py run [--dtype fp32|bf16] [--ckpt DIR] [--out DIR] + .venv/bin/python scripts/qwen3_parity.py mutants [--only NAME ...] # needs `run --dtype fp32` first + .venv/bin/python scripts/qwen3_parity.py compare --a DIR/align --b DIR/hf --dtype fp32 + +Why this exists. The tiny-model oracle tier (tests/test_qwen3_parity.py) pins the architecture +at ~1e-6 on random 3-layer weights. It cannot see what only the released weights expose: the +real key layout and shard index, a config.json whose rope is spelled the transformers-4.x way +(`rope_theta` top-level, `rope_scaling: null`), a 151936-row embedding, and DEPTH -- a 1e-7 +per-layer error is invisible at 3 layers and visible at 36. "Align runs Qwen3" is a claim about +THIS checkpoint. + +THE GATE (pinned in `evaluate`): + + fp32 compute: max |logit_align - logit_hf| < 1e-3 AND top-1 agreement 100%. + bf16 compute: measured and reported, NOT gated. Two correct bf16 implementations differ by + several ulps of a logit whose magnitude is 10-30 (ulp 0.0625-0.125); a bf16 tolerance chosen + to pass is a number fitted to the answer. fp32 is the same function up to accumulation order. + +Mutants. The four porting defects of the tiny oracle tier that apply to a dense Qwen3 are +re-run against the real checkpoint, in ONE process (the model is loaded once, the defect is +installed, the prompts are forwarded, the defect is removed). Each must FAIL the fp32 gate: +if one passes, the gate has no teeth on this checkpoint. + +Memory discipline (host: 186 GB, no swap, shared; /apps/CLAUDE.md is binding): fp32 weights are +32.8 GB per copy. align and transformers run in SEPARATE processes, one after the other, and +compare saved tensors; each takes the heavy-job flock and the host preflight before importing +torch; the orchestrator holds nothing, so the workers cannot deadlock on it. Workers see no GPU. +Everything is written under --out (default .tmp/p2/p2a/parity), never /tmp. + +Import order is load-bearing (`apply_resource_guards` sets env torch reads at import), so torch +is imported inside functions, after the guards. The legacy-venv half is a source string run by +`python -c` (`transformers` is an oracle, not an import of this file; the dependency firewall +scans it). +""" + +from __future__ import annotations + +import argparse +import glob +import hashlib +import json +import os +import resource +import subprocess +import sys +import time + +REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +LEGACY_PY = "/apps/development/rlhf/nano-mistral-o1/.venv/bin/python" +HF_HOME = "/apps/development/rlhf/hf_cache" +CKPT_GLOB = os.path.join(HF_HOME, "hub", "models--Goedel-LM--Goedel-Prover-V2-8B", "snapshots", "*") +DEFAULT_OUT = os.path.join(REPO, ".tmp", "p2", "p2a", "parity") + +#: Four prompts of different character, tokenized by the HF tokenizer in the legacy venv (align +#: receives ids). The third is a Lean 4 / Mathlib prompt in the Goedel-Prover-V2 chat framing, the +#: distribution the model is actually used on. +PROMPTS: tuple[str, ...] = ( + "The capital of France is", + "def fibonacci(n):\n if n < 2:\n return n\n return fibonacci(n - 1) + fibonacci(n - 2)\n", + "<|im_start|>user\nComplete the following Lean 4 code:\n\n```lean4\nimport Mathlib\nimport Aesop\n\n" + "set_option maxHeartbeats 400000\n\nopen BigOperators Real Nat Topology Rat\n\n" + "theorem mathd_algebra_10 : abs ((120 : ℝ) / 100 * 30 - 130 / 100 * 20) = 10 := by\n sorry\n```\n" + "<|im_end|>\n<|im_start|>assistant\n", + "In 1969, Apollo 11 landed on the Moon after a journey of 384,400 kilometres; " + "Neil Armstrong said that it was one small step for man.", +) + +FP32_MAX_ABS_LOGIT_DIFF = 1e-3 + +# --------------------------------------------------------------------------- +# The gate (pure functions over numpy arrays: no torch, unit-testable) +# --------------------------------------------------------------------------- + + +def compare_tensors(a: dict, b: dict, n_prompts: int) -> dict: + """Metrics between two engines' saved logits (dicts of numpy arrays `logits_{p}` [T, V]).""" + import numpy as np + + per_prompt, worst, n_tok, n_agree = [], 0.0, 0, 0 + for p in range(n_prompts): + la, lb = a[f"logits_{p}"].astype(np.float64), b[f"logits_{p}"].astype(np.float64) + if la.shape != lb.shape: + raise ValueError(f"logits_{p}: shape {la.shape} vs {lb.shape}") + d = np.abs(la - lb) + agree = int((la.argmax(-1) == lb.argmax(-1)).sum()) + per_prompt.append({"tokens": int(la.shape[0]), "max_abs": float(d.max()), "mean_abs": float(d.mean()), + "abs_max_logit": float(np.abs(lb).max()), "top1_agree": agree}) + worst = max(worst, float(d.max())) + n_tok += la.shape[0] + n_agree += agree + return {"max_abs_logit_diff": worst, "top1_agreement": n_agree / n_tok, "tokens": n_tok, "per_prompt": per_prompt} + + +def evaluate(metrics: dict, dtype: str) -> dict: + """The verdict. Only fp32 is a gate; bf16 returns `gated=False, passed=None`. Anything else raises.""" + if dtype == "bf16": + return {"dtype": dtype, "gated": False, "passed": None, "thresholds": None, "checks": {}, + "note": "bf16 is recorded as the noise floor between two correct implementations, not gated"} + if dtype != "fp32": + raise ValueError(f"dtype must be 'fp32' or 'bf16', got {dtype!r}") + checks = { + "max_abs_logit_diff": metrics["max_abs_logit_diff"] < FP32_MAX_ABS_LOGIT_DIFF, + "top1_agreement": metrics["top1_agreement"] == 1.0, + } + return {"dtype": dtype, "gated": True, + "thresholds": {"max_abs_logit_diff": FP32_MAX_ABS_LOGIT_DIFF, "top1_agreement": 1.0}, + "checks": checks, "passed": all(checks.values())} + + +def _load_dir(d: str) -> tuple[dict, dict]: + from safetensors.numpy import load_file + + with open(os.path.join(d, "meta.json")) as f: + meta = json.load(f) + return load_file(os.path.join(d, "tensors.safetensors")), meta + + +def compare_dirs(a_dir: str, b_dir: str, dtype: str, allow_mutant: bool = False) -> dict: + a, ma = _load_dir(a_dir) + if "MUTANT" in ma["engine"] and not allow_mutant: + raise ValueError(f"{a_dir} holds a MUTANT run ({ma['engine']}); refusing to issue a parity verdict on it") + b, mb = _load_dir(b_dir) + if ma["input_ids"] != mb["input_ids"]: + raise ValueError("the two engines were fed different token ids; the comparison is void") + if ma["dtype"] != dtype or mb["dtype"] != dtype: + raise ValueError(f"dtype mismatch: gate={dtype}, a={ma['dtype']}, b={mb['dtype']}") + metrics = compare_tensors(a, b, len(ma["input_ids"])) + return {"metrics": metrics, "verdict": evaluate(metrics, dtype), + "engines": {"a": {k: v for k, v in ma.items() if k != "input_ids"}, + "b": {k: v for k, v in mb.items() if k != "input_ids"}}, + "input_ids": ma["input_ids"]} + + +# --------------------------------------------------------------------------- +# Workers +# --------------------------------------------------------------------------- + + +def default_ckpt() -> str: + hits = sorted(glob.glob(CKPT_GLOB)) + if not hits: + raise FileNotFoundError(f"no Goedel-Prover-V2-8B snapshot under {CKPT_GLOB}") + return hits[-1] + + +def _peak_rss_gb() -> float: + return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024**2 # KB -> GiB + + +#: Porting defects injected into the align side ONLY, so the gate is shown to bite on the real +#: checkpoint. Name -> what is broken. +MUTANTS: dict[str, str] = { + "drop_qk_norm": "per-head q/k RMSNorm skipped entirely", + "fullwidth_qk_norm": "q/k RMSNorm over the whole projection width (OLMoE style) instead of per head", + "wrong_rope_theta": "RoPE base 1e4 instead of Qwen3's 1e6", + "swap_gate_up": "SwiGLU gate and up operands swapped: silu(up) * gate", +} + + +def _install_mutant(name: str, model): + """Install the named defect; return an `undo()` closure. Raises on an unknown name (a typo must + not silently run the control).""" + import torch + import torch.nn.functional as F + + from align.model import qwen3 as q3 + + if name == "drop_qk_norm": + orig = q3._head_norm + q3._head_norm = lambda norm, x: x + return lambda: setattr(q3, "_head_norm", orig) + if name == "fullwidth_qk_norm": + orig = q3._head_norm + + def full(norm, x): + shp = x.shape + flat = x.reshape(*shp[:-2], -1).float() + flat = flat * torch.rsqrt(flat.pow(2).mean(-1, keepdim=True) + norm.eps) + return (norm.weight.repeat(shp[-2]) * flat.to(x.dtype)).view(shp) + + q3._head_norm = full + return lambda: setattr(q3, "_head_norm", orig) + if name == "wrong_rope_theta": + rope = model.model.rope + saved = rope.inv_freq + rope.theta = 1e4 + rope.inv_freq = rope._compute() + def undo(): + rope.theta = 1e6 + rope.inv_freq = saved + return undo + if name == "swap_gate_up": + orig = q3._swiglu + q3._swiglu = lambda g, u: F.silu(u) * g + return lambda: setattr(q3, "_swiglu", orig) + raise ValueError(f"unknown mutant {name!r}; known: {sorted(MUTANTS)}") + + +def worker_align(ckpt: str, ids_path: str, out_dir: str, dtype: str, mutants: list[str] | None = None) -> None: + """Runs in the align venv. Guards first, torch after. With `mutants`, the control is NOT + written; each mutant gets `out_dir/mut_/`.""" + from align.utils.guard import apply_resource_guards + + apply_resource_guards(cpu_threads=8, heavy=True, label=f"qwen3-parity-align-{dtype}") + import torch + + from align.model import qwen3 as q3 + from align.model import qwen3_checkpoint as qc + + t_dtype = {"fp32": torch.float32, "bf16": torch.bfloat16}[dtype] + with open(ids_path) as f: + ids = json.load(f)["input_ids"] + t0 = time.perf_counter() + cfg, sd = qc.load_hf_qwen3(ckpt, dtype=t_dtype) # per-shard cast: transient is the target dtype's size + model = q3.Qwen3ForCausalLM.from_state_dict(cfg, sd, attn_implementation="sdpa", device="cpu").eval() + del sd + t_load = time.perf_counter() - t0 + + def forward_all() -> dict: + tensors = {} + with torch.no_grad(): + for p, row in enumerate(ids): + tensors[f"logits_{p}"] = model(input_ids=torch.tensor([row], dtype=torch.long))[0].float().numpy() + return tensors + + base_meta = {"dtype": dtype, "n_layers": cfg.n_layers, "input_ids": ids, "torch": torch.__version__, + "attn": model.attn_implementation, "load_s": round(t_load, 2)} + for name in (mutants or [None]): + undo = _install_mutant(name, model) if name else None + t1 = time.perf_counter() + try: + tensors = forward_all() + finally: + if undo: + undo() + _write(os.path.join(out_dir, f"mut_{name}") if name else out_dir, tensors, + {**base_meta, "engine": "align" + (f"+MUTANT:{name}" if name else ""), + "forward_s": round(time.perf_counter() - t1, 2), "peak_rss_gb": round(_peak_rss_gb(), 2)}) + + +def _write(out_dir: str, tensors: dict, meta: dict) -> None: + import numpy as np + from safetensors.numpy import save_file + + os.makedirs(out_dir, exist_ok=True) + save_file({k: np.ascontiguousarray(v) for k, v in tensors.items()}, os.path.join(out_dir, "tensors.safetensors")) + with open(os.path.join(out_dir, "meta.json"), "w") as f: + json.dump(meta, f, indent=2) + + +# Runs INSIDE the legacy venv (transformers 5.x). argv: mode ckpt ids_json out_dir dtype repo +_LEGACY_SCRIPT = r''' +import json, os, resource, sys, time +mode, ckpt, ids_path, out_dir, dtype, repo = sys.argv[1:7] +os.environ.setdefault("HF_HUB_OFFLINE", "1") +os.environ.setdefault("TRANSFORMERS_OFFLINE", "1") +os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") +if mode == "tokenize": + from transformers import AutoTokenizer + prompts = json.loads(ids_path) + tok = AutoTokenizer.from_pretrained(ckpt) + ids = [tok(p, add_special_tokens=False)["input_ids"] for p in prompts] + os.makedirs(out_dir, exist_ok=True) + with open(os.path.join(out_dir, "ids.json"), "w") as f: + json.dump({"prompts": prompts, "input_ids": ids, "tokenizer": type(tok).__name__, "vocab_size": len(tok)}, f, indent=2) + print(json.dumps({"n": [len(x) for x in ids]})) + sys.exit(0) + +sys.path.insert(0, repo) +from align.utils.guard import apply_resource_guards +apply_resource_guards(cpu_threads=8, heavy=True, label="qwen3-parity-hf-" + dtype) +import numpy as np, torch, transformers +from safetensors.numpy import save_file +from transformers import Qwen3ForCausalLM +t_dtype = {"fp32": torch.float32, "bf16": torch.bfloat16}[dtype] +ids = json.load(open(ids_path))["input_ids"] +t0 = time.perf_counter() +m = Qwen3ForCausalLM.from_pretrained(ckpt, dtype=t_dtype, attn_implementation="eager").eval() +t_load = time.perf_counter() - t0 +tensors = {} +t1 = time.perf_counter() +with torch.no_grad(): + for p, row in enumerate(ids): + tensors[f"logits_{p}"] = m(input_ids=torch.tensor([row])).logits[0].float().numpy() +t_fwd = time.perf_counter() - t1 +os.makedirs(out_dir, exist_ok=True) +save_file({k: np.ascontiguousarray(v) for k, v in tensors.items()}, os.path.join(out_dir, "tensors.safetensors")) +meta = {"engine": "transformers", "dtype": dtype, "n_layers": m.config.num_hidden_layers, "input_ids": ids, + "transformers": transformers.__version__, "torch": torch.__version__, "attn": "eager", + "load_s": round(t_load, 2), "forward_s": round(t_fwd, 2), + "peak_rss_gb": round(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024**2, 2)} +json.dump(meta, open(os.path.join(out_dir, "meta.json"), "w"), indent=2) +print(json.dumps({"ok": True})) +''' + + +def _worker_env() -> dict: + return { + "PATH": os.environ.get("PATH", ""), "HOME": os.environ.get("HOME", ""), + "HF_HOME": HF_HOME, "HF_HUB_CACHE": os.path.join(HF_HOME, "hub"), "HF_HUB_DISABLE_XET": "1", + "TMPDIR": os.path.join(REPO, ".tmp", "p2", "p2a"), + "CUDA_VISIBLE_DEVICES": "", # a CPU measurement must not touch a card someone else is using + "OMP_NUM_THREADS": "8", "TOKENIZERS_PARALLELISM": "false", "PYTHONPATH": REPO, + } + + +def _run(cmd: list[str], timeout: float) -> str: + cwd = os.path.join(REPO, ".tmp", "p2", "p2a") + os.makedirs(cwd, exist_ok=True) + p = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout, env=_worker_env(), cwd=cwd) + if p.returncode != 0: + raise RuntimeError(f"{cmd[:3]} exited {p.returncode}:\n{p.stderr[-3000:]}") + return p.stdout + + +def _ckpt_fingerprint(ckpt: str) -> str: + """Identity of the checkpoint a worker consumed: path + (name, size, mtime) of every weight / + index / config file. Cheap (no 16 GB read) and changes on any re-download or edit.""" + h = hashlib.sha256(os.path.realpath(ckpt).encode()) + for name in sorted(os.listdir(ckpt)): + if name.endswith((".safetensors", ".json")): + st = os.stat(os.path.join(ckpt, name)) # follows the hub snapshot's blob symlinks + h.update(f"{name}:{st.st_size}:{st.st_mtime_ns}".encode()) + return h.hexdigest() + + +def _align_code_sha() -> str: + """Hash of the align-side code a parity verdict describes: align/model/*.py plus this script.""" + h = hashlib.sha256() + paths = sorted(glob.glob(os.path.join(REPO, "align", "model", "*.py"))) + [os.path.abspath(__file__)] + for path in paths: + with open(path, "rb") as f: + h.update(os.path.relpath(path, REPO).encode() + b"\0" + f.read()) + return h.hexdigest() + + +def _provenance(side: str, ckpt: str, ids_path: str, dtype: str) -> dict: + """What a worker's saved tensors depend on. `meta.json` carries none of this, so `run` stamps + it beside the tensors (`provenance.json`) and refuses to reuse them on any mismatch.""" + with open(ids_path, "rb") as f: + ids_sha = hashlib.sha256(f.read()).hexdigest() + prov = {"side": side, "dtype": dtype, "ckpt": os.path.realpath(ckpt), + "ckpt_fingerprint": _ckpt_fingerprint(ckpt), "ids_sha256": ids_sha} + if side == "align": + prov["align_code_sha256"] = _align_code_sha() + return prov + + +def _read_provenance(out_dir: str) -> dict | None: + try: + with open(os.path.join(out_dir, "provenance.json")) as f: + return json.load(f) + except (OSError, ValueError): + return None + + +def _reusable(out_dir: str, prov: dict) -> bool: + return os.path.isfile(os.path.join(out_dir, "meta.json")) and _read_provenance(out_dir) == prov + + +def _run_worker(out_dir: str, prov: dict, cmd: list[str], timeout: float) -> None: + # Invalidate FIRST: a worker that dies mid-way must not leave the previous run's meta.json + # standing beside a stamp for a run that never finished. + for name in ("meta.json", "provenance.json"): + if os.path.isfile(os.path.join(out_dir, name)): + os.remove(os.path.join(out_dir, name)) + _run(cmd, timeout) + with open(os.path.join(out_dir, "provenance.json"), "w") as f: + json.dump(prov, f, indent=2) + + +def run(ckpt: str, out: str, dtype: str, timeout: float = 3600.0, fresh: bool = False) -> dict: + """Tokenize (legacy) -> align worker -> HF worker -> compare -> report.json.""" + out, ckpt = os.path.abspath(out), os.path.abspath(ckpt) + os.makedirs(out, exist_ok=True) + ids_dir = os.path.join(out, "ids") + _run([LEGACY_PY, "-c", _LEGACY_SCRIPT, "tokenize", ckpt, json.dumps(list(PROMPTS)), ids_dir, dtype, REPO], 300) + ids_path = os.path.join(ids_dir, "ids.json") + a_dir, b_dir = os.path.join(out, f"align_{dtype}"), os.path.join(out, f"hf_{dtype}") + # SEQUENTIAL by construction: two fp32 copies are 66 GB. Each worker frees its memory by exiting. + # A worker whose output is already on disk is not re-run: the heavy-job flock is shared with + # other agents, a worker can lose the race for it AFTER its sibling finished, and re-doing + # 33 GB of CPU work to retry the one that failed is the wrong default. Reuse is gated on the + # PROVENANCE stamp (checkpoint identity, align code hash, token ids): tensors that another + # checkpoint or older code produced must never carry this run's verdict. --fresh redoes both. + a_prov = _provenance("align", ckpt, ids_path, dtype) + b_prov = _provenance("hf", ckpt, ids_path, dtype) + if fresh or not _reusable(a_dir, a_prov): + _run_worker(a_dir, a_prov, [sys.executable, os.path.abspath(__file__), "worker-align", "--ckpt", ckpt, + "--ids", ids_path, "--out", a_dir, "--dtype", dtype], timeout) + if fresh or not _reusable(b_dir, b_prov): + _run_worker(b_dir, b_prov, [LEGACY_PY, "-c", _LEGACY_SCRIPT, "hf", ckpt, ids_path, b_dir, dtype, REPO], + timeout) + rep = compare_dirs(a_dir, b_dir, dtype) + rep["checkpoint"] = ckpt + with open(os.path.join(out, f"report_{dtype}.json"), "w") as f: + json.dump(rep, f, indent=2) + return rep + + +def run_mutants(ckpt: str, out: str, names: list[str] | None = None, timeout: float = 3600.0) -> dict: + """Re-run ONLY the align side, all mutants in one process (fp32), against the saved + transformers tensors of `run`.""" + out, ckpt = os.path.abspath(out), os.path.abspath(ckpt) + ids_path = os.path.join(out, "ids", "ids.json") + b_dir = os.path.join(out, "hf_fp32") + if not (os.path.isfile(ids_path) and os.path.isfile(os.path.join(b_dir, "meta.json"))): + raise FileNotFoundError("run `qwen3_parity.py run --dtype fp32` first: the mutants compare to its transformers tensors") + if _read_provenance(b_dir) != _provenance("hf", ckpt, ids_path, "fp32"): + raise RuntimeError(f"{b_dir} was not produced from {ckpt} (missing or stale provenance.json): " + "re-run `qwen3_parity.py run --dtype fp32` before the mutants") + names = names or sorted(MUTANTS) + a_root = os.path.join(out, "align_fp32_mutants") + _run([sys.executable, os.path.abspath(__file__), "worker-align", "--ckpt", ckpt, "--ids", ids_path, + "--out", a_root, "--dtype", "fp32", "--mutants", *names], timeout) + res = {} + for name in names: + rep = compare_dirs(os.path.join(a_root, f"mut_{name}"), b_dir, "fp32", allow_mutant=True) + res[name] = {"what": MUTANTS[name], "max_abs_logit_diff": rep["metrics"]["max_abs_logit_diff"], + "top1_agreement": rep["metrics"]["top1_agreement"], "gate_passed": rep["verdict"]["passed"]} + with open(os.path.join(out, "mutants_fp32.json"), "w") as f: + json.dump(res, f, indent=2) + return res + + +def _summary(rep: dict) -> str: + m, v = rep["metrics"], rep["verdict"] + lines = [ + f"dtype={v['dtype']} gated={v['gated']} passed={v['passed']} checks={v['checks']}", + f"max_abs_logit_diff={m['max_abs_logit_diff']:.3e} (gate {'< %g' % v['thresholds']['max_abs_logit_diff'] if v['gated'] else 'none: bf16 informational'}) " + f"top1={m['top1_agreement']:.4f} tokens={m['tokens']}", + ] + for p in m["per_prompt"]: + lines.append(f" prompt: tokens={p['tokens']} max_abs={p['max_abs']:.3e} mean_abs={p['mean_abs']:.3e} " + f"|logit|max={p['abs_max_logit']:.2f} top1={p['top1_agree']}/{p['tokens']}") + return "\n".join(lines) + + +def main(argv=None) -> int: + ap = argparse.ArgumentParser(description=__doc__.split("\n")[0]) + sub = ap.add_subparsers(dest="cmd", required=True) + r = sub.add_parser("run") + r.add_argument("--ckpt") + r.add_argument("--out", default=DEFAULT_OUT) + r.add_argument("--dtype", choices=("fp32", "bf16"), default="fp32") + r.add_argument("--timeout", type=float, default=3600.0) + r.add_argument("--fresh", action="store_true", help="re-run workers whose output already exists") + w = sub.add_parser("worker-align") + w.add_argument("--ckpt", required=True) + w.add_argument("--ids", required=True) + w.add_argument("--out", required=True) + w.add_argument("--dtype", choices=("fp32", "bf16"), required=True) + w.add_argument("--mutants", nargs="*", choices=sorted(MUTANTS)) + mu = sub.add_parser("mutants") + mu.add_argument("--ckpt") + mu.add_argument("--out", default=DEFAULT_OUT) + mu.add_argument("--only", nargs="*", choices=sorted(MUTANTS)) + c = sub.add_parser("compare") + c.add_argument("--a", required=True) + c.add_argument("--b", required=True) + c.add_argument("--dtype", choices=("fp32", "bf16"), required=True) + a = ap.parse_args(argv) + if a.cmd == "worker-align": + worker_align(a.ckpt, a.ids, a.out, a.dtype, a.mutants) + return 0 + if a.cmd == "mutants": + res = run_mutants(a.ckpt or default_ckpt(), a.out, a.only) + for k, v in res.items(): + print(f"{k:20s} max_abs={v['max_abs_logit_diff']:.3e} top1={v['top1_agreement']:.3f} gate_passed={v['gate_passed']} [{v['what']}]") + return 1 if any(v["gate_passed"] for v in res.values()) else 0 + if a.cmd == "compare": + rep = compare_dirs(a.a, a.b, a.dtype) + else: + rep = run(a.ckpt or default_ckpt(), a.out, a.dtype, a.timeout, a.fresh) + print(_summary(rep)) + return 1 if rep["verdict"]["passed"] is False else 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/code/scripts/qwen3_sft_probe.py b/code/scripts/qwen3_sft_probe.py new file mode 100644 index 0000000000000000000000000000000000000000..7e10a510f71a211ba694170402dfed2a3e0ec6cc --- /dev/null +++ b/code/scripts/qwen3_sft_probe.py @@ -0,0 +1,498 @@ +#!/usr/bin/env python +"""Qwen3-8B LoRA SFT memory / throughput / resume probe on 4x A10G (P2-A). + + .venv/bin/python scripts/qwen3_sft_probe.py prepare --lens 4096 8192 16384 [--prompt-len 1024] + torchrun --standalone --nproc_per_node=4 scripts/qwen3_sft_probe.py train CONFIG [dotted.overrides=...] + .venv/bin/python scripts/qwen3_sft_probe.py report RUN_DIR [--accum A --world W --seq-len L] + .venv/bin/python scripts/qwen3_sft_probe.py batches RUN_DIR + +What it measures, and why it is not just `scripts/sft.py`. `train` runs the SAME functions +`scripts/sft.py` runs (`build_model`, `run_gates`, `train`: arch dispatch, LoRA wiring, the +Trainer, the checkpoint format), with two additions the production entry deliberately does not +carry: a `Tracker.log` wrapper that stamps every history record with `rank`, wall time and the +CUDA allocator's peak, and a heavy-job guard identical to sft.py's. The numbers it reports are +therefore the numbers a real launch would see. + +DATA is REAL TOKENS, not random ids: the formal statements and informal prefixes of the Goedel +reference repo's miniF2F + MOBench files (604 problems), tokenized by the checkpoint's own +tokenizer in the legacy venv, concatenated into one stream, and cut into rows of exactly `len` +tokens at distinct offsets. Loss values on these rows are meaningless as a quality signal (rows +are arbitrary windows); compute, memory and the sharding/resume mechanics are exactly those of +real rows of that length. Every micro-batch is one row of the target length, so every micro-step +costs the same, and tokens/s is `len / micro_step_time`. + +Import order is load-bearing: this module imports only the stdlib at top level. +""" + +from __future__ import annotations + +import argparse +import json +import os +import statistics +import subprocess +import sys + +REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, REPO) + +LEGACY_PY = "/apps/development/rlhf/nano-mistral-o1/.venv/bin/python" +GOEDEL_REPO = "/apps/development/rlhf/reference-repos/Goedel-Prover-V2" +SNAPSHOT = ("/apps/development/rlhf/hf_cache/hub/models--Goedel-LM--Goedel-Prover-V2-8B/snapshots/" + "dfd02e6271a58375dfbf3ece0175277cf6b6a89a") +OUT_ROOT = os.path.join(REPO, ".tmp", "p2", "p2a") +STREAM = os.path.join(OUT_ROOT, "probe_stream.npy") + +_TOKENIZE = r''' +import json, sys +import numpy as np +from transformers import AutoTokenizer +ckpt, files, out = sys.argv[1], sys.argv[2].split(","), sys.argv[3] +tok = AutoTokenizer.from_pretrained(ckpt) +ids = [] +for fn in files: + for line in open(fn): + d = json.loads(line) + text = (d.get("informal_prefix") or "") + (d.get("formal_statement") or "") + "\n" + ids.extend(tok(text, add_special_tokens=False)["input_ids"] + [tok.eos_token_id]) +np.save(out, np.asarray(ids, dtype=np.int32)) +print(json.dumps({"tokens": len(ids), "vocab": len(tok)})) +''' + + +def cut_rows(stream, length: int, n_rows: int, seed: int = 0): + """`n_rows` windows of exactly `length` tokens at distinct, seeded offsets (wrapping the stream + when it is shorter than needed). Pure numpy; unit-tested.""" + import numpy as np + + if length < 2: + raise ValueError("length must be >= 2") + n = len(stream) + rng = np.random.default_rng(seed) + offs = rng.choice(n, size=n_rows, replace=n_rows > n) + doubled = np.concatenate([stream] * (length // n + 2)) + return [doubled[o:o + length] for o in offs] + + +def make_tiny_snapshot(out_dir: str, vocab: int = 2048, layers: int = 4, hidden: int = 256, seed: int = 0) -> dict: + """A random-weight, structurally-faithful Qwen3 snapshot (GQA 8:4, head_dim 32, untied head) with + the Goedel tokenizer files, for wiring smokes on the 4-GPU path in seconds instead of minutes. + NOT a model of anything: it exists so FSDP2 + LoRA + checkpoint + resume bugs surface before the + 8B load. Returns the config overrides a TOML needs to agree with it.""" + import torch + + from align.model import qwen3 as q3 + from align.model.qwen3_checkpoint import save_hf_qwen3 + + cfg = q3.Qwen3Config(name="qwen3-tiny", vocab_size=vocab, hidden_size=hidden, intermediate_size=2 * hidden, + n_layers=layers, n_heads=8, n_kv_heads=4, head_dim=hidden // 8, max_position_embeddings=4096) + torch.manual_seed(seed) + m = q3.Qwen3ForCausalLM(cfg, attn_implementation="sdpa", device="cpu") + snap = os.path.join(out_dir, "snapshots", "tiny") + save_hf_qwen3(m, snap, dtype=torch.bfloat16, tokenizer_dir=SNAPSHOT) + return {"paths.model_snapshot": os.path.join(out_dir, "snapshots"), "model.n_params": q3.count_params(cfg), + "model.n_layers": layers, "model.vocab_size": vocab, "model.hidden_size": hidden} + + +def prepare(lens: list[int], rows: int, prompt_len: int, out_root: str, vocab_mod: int | None = None) -> dict: + import numpy as np + + from align.data.store import ShardWriter + + os.makedirs(out_root, exist_ok=True) + if not os.path.isfile(STREAM): + env = {"PATH": os.environ.get("PATH", ""), "HOME": os.environ.get("HOME", ""), + "HF_HOME": "/apps/development/rlhf/hf_cache", "HF_HUB_OFFLINE": "1", "TRANSFORMERS_OFFLINE": "1", + "TOKENIZERS_PARALLELISM": "false", "TMPDIR": OUT_ROOT} + files = ",".join(os.path.join(GOEDEL_REPO, "dataset", f) for f in ("minif2f.jsonl", "MOBench.jsonl")) + p = subprocess.run([LEGACY_PY, "-c", _TOKENIZE, SNAPSHOT, files, STREAM], capture_output=True, text=True, + timeout=600, env=env, cwd=OUT_ROOT) + if p.returncode != 0: + raise RuntimeError(p.stderr[-2000:]) + stream = np.load(STREAM) + if vocab_mod: + stream = stream % vocab_mod # tiny-vocab smokes only: ids must index the tiny embedding + out = {"stream_tokens": int(len(stream))} + for L in lens: + d = os.path.join(out_root, f"L{L}", "sft", "train") + if prompt_len >= L: + raise ValueError(f"prompt_len {prompt_len} >= row length {L}") + with ShardWriter(d, "sft-train-0000") as w: + for i, row in enumerate(cut_rows(stream, L, rows, seed=L)): + w.add_row(row.tolist(), problem_id=f"probe-{L}-{i}", subject="lean", boundary_positions=[prompt_len], + prompt_len=prompt_len) + out[f"L{L}"] = {"dir": d, "rows": rows} + return out + + +# --------------------------------------------------------------------------- +# analysis (pure; unit-tested) +# --------------------------------------------------------------------------- + + +def read_jsonl(path: str) -> list[dict]: + if not os.path.exists(path): + return [] + lines = [l for l in open(path).read().split("\n") if l.strip()] + out = [] + for i, l in enumerate(lines): + try: + out.append(json.loads(l)) + except json.JSONDecodeError: + if i != len(lines) - 1: + raise + return out + + +def rank_histories(run_dir: str, world: int) -> dict[int, list[dict]]: + """rank -> history records (rank 0 in `run_dir`, rank r in `run_dir/rank{r}`).""" + return {r: read_jsonl(os.path.join(run_dir if r == 0 else os.path.join(run_dir, f"rank{r}"), "history.jsonl")) + for r in range(world)} + + +def rank_disjointness(hist: dict[int, list[dict]]) -> dict: + """Do the ranks see DIFFERENT rows? `batch` is the list of dataset row indices of a micro-step. + Returns the per-rank row sets and whether they are pairwise disjoint, plus per-micro-step + disjointness (the contract: at micro-step m no two ranks share a row).""" + seen = {r: [tuple(rec["batch"]) for rec in recs] for r, recs in hist.items()} + sets = {r: {i for b in bs for i in b} for r, bs in seen.items()} + ranks = sorted(sets) + pairwise = all(not (sets[a] & sets[b]) for i, a in enumerate(ranks) for b in ranks[i + 1:]) + n_micro = min(len(v) for v in seen.values()) + per_micro = all(len({i for r in ranks for i in seen[r][m]}) == sum(len(seen[r][m]) for r in ranks) for m in range(n_micro)) + return {"rows_per_rank": {r: len(s) for r, s in sets.items()}, "pairwise_disjoint": pairwise, + "per_micro_step_disjoint": per_micro, "micro_steps_compared": n_micro, + "first_batches": {r: [list(b) for b in seen[r][:4]] for r in ranks}} + + +def step_stats(recs: list[dict], seq_len: int, accum: int, world: int, drop_first: int = 2, + save_every: int | None = None) -> dict: + """Median micro-step time from consecutive `t_wall` stamps of ONE rank (skipping the first + `drop_first` intervals: allocator warmup, first-use kernels, and any interval containing a + checkpoint save), peak memory over the run, and tokens/s = world * seq_len / micro_step_time + (every rank processes one row of `seq_len` tokens per micro-step).""" + recs = sorted(recs, key=lambda r: r["micro_step"]) + dts = [b["t_wall"] - a["t_wall"] for a, b in zip(recs, recs[1:]) + if b["micro_step"] == a["micro_step"] + 1 and a["micro_step"] >= drop_first + and not (save_every and a["micro_step"] % save_every == 0)] + if not dts: + raise ValueError("no usable step intervals") + med = statistics.median(dts) + return { + "n_intervals": len(dts), "micro_median_s": med, "micro_min_s": min(dts), "micro_max_s": max(dts), + "optimizer_step_s": med * accum, "tokens_per_s_total": world * seq_len / med, + "tokens_per_s_per_gpu": seq_len / med, + "alloc_peak_gb": max(r["mem/alloc_peak_gb"] for r in recs), + "reserved_peak_gb": max(r["mem/reserved_peak_gb"] for r in recs), + "loss_first": recs[0]["loss"], "loss_last": recs[-1]["loss"], + "grad_norm_last": next((r["grad_norm"] for r in reversed(recs) if r["grad_norm"] == r["grad_norm"]), float("nan")), + } + + +def dedupe_by_micro(recs: list[dict]) -> dict[int, dict]: + """micro_step -> record, the LAST occurrence winning (a relaunch may re-log the boundary it resumed at).""" + return {int(r["micro_step"]): r for r in recs} + + +def compare_resume(ref: dict[int, list[dict]], sup: dict[int, list[dict]], stop_micro: int, + floor_mult: float = 2.0, floor_pad: float = 5e-3) -> dict: + """Does a SIGTERM'd-and-resumed run continue the uninterrupted one? + + Three separate questions, per rank, because a loss curve alone cannot answer the first two: + + * DATA: the resumed segment (micro > stop) saw EXACTLY the batches the reference saw at the + same micro-steps (the stream offset round-trips), and nothing was replayed or skipped + (micro-steps are contiguous across the stop). + * CONTINUITY: |dloss| of the resumed segment vs the reference stays within + `floor_mult * (pre-stop |dloss|) + floor_pad`. The pre-stop segment is the SAME computation + in both runs, so its |dloss| measures pure run-to-run nondeterminism (NCCL reduction order, + SDPA backward, 8-bit rounding); the resumed segment must not be worse than that floor by + more than the factor. A vacuous comparison (one side empty) raises. + * PROVENANCE is checked separately and exactly by the caller (the relaunch log names the + checkpoint it loaded; this function does not claim it). + """ + out = {} + for r in sorted(ref): + a, b = dedupe_by_micro(ref[r]), dedupe_by_micro(sup[r]) + common = sorted(set(a) & set(b)) + pre = [abs(a[m]["loss"] - b[m]["loss"]) for m in common if m <= stop_micro] + post = [abs(a[m]["loss"] - b[m]["loss"]) for m in common if m > stop_micro] + if not pre or not post: + raise ValueError(f"rank {r}: need micro-steps on both sides of the stop ({stop_micro}); pre={len(pre)} post={len(post)}") + batches_equal = all(a[m]["batch"] == b[m]["batch"] for m in common) + contiguous = sorted(b) == list(range(min(b), max(b) + 1)) + tol = floor_mult * max(pre) + floor_pad + out[r] = {"n_pre": len(pre), "n_post": len(post), "batches_equal_all_common": batches_equal, + "micro_steps_contiguous": contiguous, "pre_stop_max_dloss": max(pre), "resumed_max_dloss": max(post), + "tolerance": tol, "continuity_passed": max(post) <= tol, + "last_micro_ref": max(a), "last_micro_resumed": max(b)} + out["passed"] = all(v["batches_equal_all_common"] and v["micro_steps_contiguous"] and v["continuity_passed"] + and v["last_micro_ref"] == v["last_micro_resumed"] for k, v in out.items()) + return out + + +def adapter_diff(lora_a: dict, lora_b: dict) -> dict: + """Relative difference of two adapter states (`lora_state_dict` format, `.A` / `.B`). + + The strongest resume check a loss curve cannot give: the FINAL adapter of an interrupted + + resumed run against an uninterrupted run of the same schedule. `rel_B` is + ||B_a - B_b|| / ||B_b||, global over all layers. B starts at EXACTLY zero, so ||B|| is the whole + distance trained; a resume that restarts the adapter from the bare base (B = 0) after k of n + optimizer steps lands near rel_B ~ (1 - k/n)-ish, orders of magnitude above run-to-run + nondeterminism (flash-attention backward atomics) which stays at the 1e-3 level. + """ + import torch + + if set(lora_a) != set(lora_b): + raise KeyError(f"adapter key sets differ: {sorted(set(lora_a) ^ set(lora_b))[:4]}") + num = {"A": 0.0, "B": 0.0} + den = {"A": 0.0, "B": 0.0} + for k in lora_a: + kind = k.rsplit(".", 1)[1] + a, b = lora_a[k].double(), lora_b[k].double() + if a.shape != b.shape: + raise ValueError(f"{k}: {tuple(a.shape)} vs {tuple(b.shape)}") + num[kind] += float((a - b).pow(2).sum()) + den[kind] += float(b.pow(2).sum()) + if den["B"] == 0.0: + raise ValueError("reference adapter has B == 0: nothing was trained, the comparison is vacuous") + return {"rel_A": (num["A"] / den["A"]) ** 0.5, "rel_B": (num["B"] / den["B"]) ** 0.5, + "norm_B": den["B"] ** 0.5, "n_tensors": len(lora_a)} + + +def adapter_diff_main(a) -> int: + import torch + + ld = lambda d: torch.load(os.path.join(d, "optimizer_rank0.pt"), weights_only=True, map_location="cpu")["lora"] + res = adapter_diff(ld(a.ckpt_a), ld(a.ckpt_b)) + print(json.dumps(res, indent=2)) + return 0 + + +# --------------------------------------------------------------------------- +# train (torchrun target) +# --------------------------------------------------------------------------- + + +def tree_equal(a, b, path: str = "") -> list[str]: + """Paths at which two nested state structures (dict / list / tensor / scalar) are NOT bitwise + equal. Tensors compare by dtype, shape and `torch.equal`; everything else by `==`. Empty list = + identical. Pure; unit-tested.""" + import torch + + if isinstance(a, dict) and isinstance(b, dict): + out = [f"{path}/" for _ in [0] if set(a) != set(b)] + for k in sorted(set(a) & set(b), key=str): + out += tree_equal(a[k], b[k], f"{path}/{k}") + return out + if isinstance(a, (list, tuple)) and isinstance(b, (list, tuple)): + if len(a) != len(b): + return [f"{path}/"] + return [d for i, (x, y) in enumerate(zip(a, b)) for d in tree_equal(x, y, f"{path}[{i}]")] + if isinstance(a, torch.Tensor) and isinstance(b, torch.Tensor): + same = a.dtype == b.dtype and a.shape == b.shape and torch.equal(a.cpu(), b.cpu()) + return [] if same else [path] + return [] if a == b else [path] + + +def _verify_resume_only(trainer): + """Stand-in for `Trainer.train` under --verify-resume-only: after `resume_from_latest()` the live + sharded state (the gathered adapter, the 8-bit moments, the step counters) must equal what the + checkpoint holds, BITWISE. No step is taken, so there is no GPU nondeterminism to excuse a + difference -- this isolates the scatter path (`distribute_tensor` / local-shard copies) from + training noise, which the adapter-diff of two uninterrupted runs (rel_B up to ~0.4) cannot do.""" + import torch + + from align.utils.dist import rank + + R = rank() + ck = trainer.latest_checkpoint() + if ck is None: + raise RuntimeError("--verify-resume-only needs an existing checkpoint to resume from") + saved = torch.load(os.path.join(ck, f"optimizer_rank{R}.pt"), weights_only=True, map_location="cpu") + live = trainer.train_module.optim_state_dict() # collective: every rank gathers the adapter + diffs = tree_equal(saved, live) + meta = torch.load(os.path.join(ck, "resume.pt"), weights_only=True, map_location="cpu") + counters = {"micro_step": (meta["micro_step"], trainer.micro_step), + "optimizer_step": (meta["optimizer_step"], trainer.optimizer_step)} + bad_counters = {k: v for k, v in counters.items() if v[0] != v[1]} + res = {"rank": R, "checkpoint": ck, "state_tensors_compared": len(tree_leaves(saved)), + "unequal_paths": diffs[:10], "n_unequal": len(diffs), "counters": counters, "bad_counters": bad_counters, + "passed": not diffs and not bad_counters} + with open(os.path.join(trainer.run_dir, f"verify_resume_rank{R}.json"), "w") as f: + json.dump(res, f, indent=2, default=str) + print(json.dumps(res, default=str), flush=True) + trainer.stop_reason = "verify_only" + trainer.verify_passed = bool(res["passed"]) + return {"optimizer_step": trainer.optimizer_step, "verify_only": True, **{"passed": res["passed"]}} + + +def final_exit_code(trainer, verify_only: bool, any_rank=None) -> int: + """Process exit status. Under --verify-resume-only a FAILED (or never-recorded) bitwise verdict is + nonzero on EVERY rank (the verdict is MAX-reduced first), so a wrapper / CI cannot read success + from a run whose per-rank verdict only went to verify_resume_rank{R}.json and stdout.""" + if verify_only: + if any_rank is None: + from align.utils.dist import any_rank + if any_rank(not getattr(trainer, "verify_passed", False)): + return 1 + return trainer.stop_exit_code or 0 + + +def tree_leaves(x) -> list: + import torch + + if isinstance(x, dict): + return [l for v in x.values() for l in tree_leaves(v)] + if isinstance(x, (list, tuple)): + return [l for v in x for l in tree_leaves(v)] + return [x] if isinstance(x, torch.Tensor) else [] + + +def train_main(argv: list[str]) -> int: + import importlib.util + + # scripts/sft.py by path: the same build_model / run_gates / train the production entry runs. + spec = importlib.util.spec_from_file_location("sft_script", os.path.join(REPO, "scripts", "sft.py")) + sft_script = importlib.util.module_from_spec(spec) + spec.loader.exec_module(sft_script) + + from align.config import load_config + from align.train.stage_gates import GateLedger, eval_source_hash, parse_args + from align.utils.guard import apply_resource_guards + + verify_only = "--verify-resume-only" in argv + argv = [x for x in argv if x != "--verify-resume-only"] + args = parse_args(argv) + cfg = load_config(args.config, args.overrides) + apply_resource_guards(cpu_threads=8, heavy=True, label="p2a-qwen3-sft-probe", rss_cfg=cfg) + import logging + import time + + logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s") + from align.utils.dist import init_distributed + + run_dir = os.path.join(cfg.paths.checkpoint_dir, "sft") + init_distributed(run_dir=run_dir) + import torch + + from align.train.stage_gates import assert_distinct_devices + from align.utils.dist import barrier, rank + + assert_distinct_devices() + R = rank() + os.makedirs(run_dir, exist_ok=True) + with open(os.path.join(run_dir, f"rank{R}.pid"), "w") as f: + f.write(str(os.getpid())) + + from align.utils import tracking + + orig_log = tracking.Tracker.log + + def log(self, step, metrics=None, **kw): + torch.cuda.synchronize() + m = {**(dict(metrics) if metrics else {}), **kw} + m.update({"rank": R, "t_wall": time.time(), "mem/alloc_peak_gb": torch.cuda.max_memory_allocated() / 2**30, + "mem/reserved_peak_gb": torch.cuda.max_memory_reserved() / 2**30}) + orig_log(self, step, m) + + tracking.Tracker.log = log + + ledger = GateLedger(run_dir, eval_source_hash=eval_source_hash()) + banner = sft_script.run_gates(cfg, ledger) + from align.registry import render_banner + + if R == 0: + print(render_banner(banner, title=f"align SFT probe :: {cfg.run_name}"), flush=True) + model = sft_script.build_model(cfg) + torch.cuda.synchronize() + post_wrap = torch.cuda.memory_allocated() / 2**30 + # The load phase held a full 16.4 GB bf16 replica per rank before FSDP sharded it; the caching + # allocator keeps that block reserved. Without this, `reserved_peak` reports the LOAD, not the + # training step (measured: 18.875 GiB at both seq 4096 and 8192). + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + if verify_only: + from align.train import trainer as trainer_mod + + trainer_mod.Trainer.train = _verify_resume_only + trainer = sft_script.train(cfg, ledger, model, banner) + barrier() + rec = {"rank": R, "post_wrap_alloc_gb": post_wrap, "alloc_peak_gb": torch.cuda.max_memory_allocated() / 2**30, + "reserved_peak_gb": torch.cuda.max_memory_reserved() / 2**30, "stop_reason": trainer.stop_reason, + "micro_step": trainer.micro_step, "optimizer_step": trainer.optimizer_step, + "load_stats": getattr(model, "_stage_load_stats", {})} + with open(os.path.join(run_dir, f"final_rank{R}.json"), "w") as f: + json.dump(rec, f, indent=2) + return final_exit_code(trainer, verify_only) + + +def report_main(a) -> int: + run = os.path.join(a.run_dir, "sft") if os.path.isdir(os.path.join(a.run_dir, "sft")) else a.run_dir + hist = rank_histories(run, a.world) + out = {"disjoint": rank_disjointness(hist), "per_rank": {r: step_stats(h, a.seq_len, a.accum, a.world, save_every=a.save_every) for r, h in hist.items()}} + print(json.dumps(out, indent=2)) + return 0 + + +def compare_main(a) -> int: + pick = lambda d: os.path.join(d, "sft") if os.path.isdir(os.path.join(d, "sft")) else d + res = compare_resume(rank_histories(pick(a.ref), a.world), rank_histories(pick(a.sup), a.world), a.stop_micro) + print(json.dumps(res, indent=2)) + return 0 if res["passed"] else 1 + + +def batches_main(a) -> int: + run = os.path.join(a.run_dir, "sft") if os.path.isdir(os.path.join(a.run_dir, "sft")) else a.run_dir + hist = rank_histories(run, a.world) + for r, recs in hist.items(): + print(f"rank {r}: " + " ".join(f"m{rec['micro_step']}:{rec['batch']}" for rec in recs[:12])) + print(json.dumps(rank_disjointness(hist), indent=2)) + return 0 + + +def main(argv=None) -> int: + argv = list(sys.argv[1:] if argv is None else argv) + if argv and argv[0] == "train": # config + dotted overrides are parsed by the repo's own parser + return train_main(argv[1:]) + ap = argparse.ArgumentParser(description=__doc__.split("\n")[0]) + sub = ap.add_subparsers(dest="cmd", required=True) + p = sub.add_parser("prepare") + p.add_argument("--lens", type=int, nargs="+", required=True) + p.add_argument("--rows", type=int, default=256) + p.add_argument("--prompt-len", type=int, default=1024) + p.add_argument("--out-root", default=os.path.join(OUT_ROOT, "shards")) + p.add_argument("--vocab-mod", type=int, default=None) + t = sub.add_parser("make-tiny") + t.add_argument("--out-dir", default=os.path.join(OUT_ROOT, "tiny_model")) + for name in ("report", "batches"): + r = sub.add_parser(name) + r.add_argument("run_dir") + r.add_argument("--world", type=int, default=4) + r.add_argument("--accum", type=int, default=1) + r.add_argument("--seq-len", type=int, default=4096) + r.add_argument("--save-every", type=int, default=None) + d = sub.add_parser("adapter-diff") + d.add_argument("ckpt_a") + d.add_argument("ckpt_b") + c = sub.add_parser("compare") + c.add_argument("ref") + c.add_argument("sup") + c.add_argument("--stop-micro", type=int, required=True) + c.add_argument("--world", type=int, default=4) + a = ap.parse_args(argv) + if a.cmd == "adapter-diff": + return adapter_diff_main(a) + if a.cmd == "compare": + return compare_main(a) + if a.cmd == "make-tiny": + print(json.dumps(make_tiny_snapshot(a.out_dir), indent=2)) + return 0 + if a.cmd == "prepare": + print(json.dumps(prepare(a.lens, a.rows, a.prompt_len, a.out_root, a.vocab_mod), indent=2)) + return 0 + return report_main(a) if a.cmd == "report" else batches_main(a) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/code/scripts/sft.py b/code/scripts/sft.py new file mode 100644 index 0000000000000000000000000000000000000000..c1da3512bc84b6c68f10d5935f35e3a49d0264ab --- /dev/null +++ b/code/scripts/sft.py @@ -0,0 +1,260 @@ +#!/usr/bin/env python +"""SFT stage entry (section-11): gates -> wrap -> train -> eval -> quality gate. + + torchrun --nproc_per_node=4 scripts/sft.py configs/sft.toml --optimizer.lr=1e-5 + +Thin by design. Everything this drives is already built and tested elsewhere; +what lives here is ORDER, and the order is the product: + + golden_regression -> smoke_7b -> probe_20step -> rss_preflight + -> train -> eval -> quality_gate + +IMPORT ORDER IS LOAD-BEARING. `apply_resource_guards()` sets +PYTORCH_CUDA_ALLOC_CONF, the /apps cache paths and the thread caps, and torch +reads all of them AT IMPORT. So this module imports only stdlib, the guard, +the config, and `stage_gates` (which is asserted torch-free); torch and every +align module that pulls it in are imported INSIDE main(), after the guards run. +A test enforces this by AST rather than by grep -- a literal search passes on +a file that mentions the guard in a comment and imports torch on the next line. + +Resume semantics: rerunning after a kill picks up the gate ledger, skips gates +that already passed, and resumes the Trainer from the latest micro-step +checkpoint. Partial checkpoints are deleted at startup. +""" + +from __future__ import annotations + +import logging +import os +import sys +from pathlib import Path + +# The package is not installed in the venv. Without this a bare `torchrun scripts/sft.py ...` (the launch +# line in docs/P2.md section 11 and in the config header) died with ModuleNotFoundError on all four ranks +# unless the operator remembered PYTHONPATH -- found by the P2 integration smoke loop. +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from align.config import load_config +from align.train.stage_gates import ( + GateLedger, + assemble_sft_banner, + assert_build_is_current, + assert_probe_resolved, + eval_source_hash, + evaluate_and_gate, + parse_args, + record_golden_from_report, + require_golden_receipt, + run_entry_gates, +) +from align.utils.guard import apply_resource_guards + +logger = logging.getLogger(__name__) + + +def build_model(cfg): + """Imported here, not at module scope: this pulls torch in. + + ARCH DISPATCH. `model.arch` selects the family; an unrecognised value RAISES (it must not + fall through to the Mistral converters, which would either fail late on a Qwen3 snapshot's + keys or load a different family into the wrong module). + """ + if cfg.model.arch == "mistral": + if cfg.lora.enabled: + raise ValueError( + "lora.enabled with model.arch='mistral': the dense-Mistral SFT path is full-FT only " + "(LoRA Mistral runs live in scripts/ppo.py). Refusing rather than training full-FT " + "while the config says LoRA." + ) + from align.train.stage_smoke import build_stage_model + + return build_stage_model(cfg) + if cfg.model.arch == "qwen3": + from align.train.qwen3_stage import build_qwen3_stage_model + + return build_qwen3_stage_model(cfg) + raise ValueError(f"model.arch={cfg.model.arch!r} is not supported; known: ('mistral', 'qwen3')") + + +def run_gates(cfg, ledger, banner_only: bool = False): + """The M5 entry sequence. Refuses to proceed on a stale data build or an + unvalidated eval stack before it touches a GPU. + + Non-Mistral arches have their own entry (see `qwen3_entry_banner`: the golden/MATH/GSM8K + machinery certifies a Mistral checkpoint and does not apply).""" + if cfg.model.arch == "qwen3": + from align.train.qwen3_stage import qwen3_entry_banner + + return qwen3_entry_banner(cfg, ledger) + if cfg.model.arch != "mistral": + raise ValueError(f"model.arch={cfg.model.arch!r} is not supported; known: ('mistral', 'qwen3')") + assert_build_is_current(cfg.data.manifest_path, cfg) + # Bridge the golden RUN into the ledger by re-deciding from its persisted + # report, immediately before the gate that reads the ledger. Without this + # a PASSING golden run still leaves golden_regression PENDING. + record_golden_from_report(ledger, cfg=cfg) + require_golden_receipt(ledger, cfg) + banner = run_entry_gates(cfg, ledger) + assert_probe_resolved(banner) + return banner + + +def train(cfg, ledger, model, banner): + """One SFT epoch through the section-07 Trainer.""" + from align.data.loader import collate_sft + from align.train.sft_loss import make_sft_loss_adapter + from align.train.stage_smoke import build_training_sampler + from align.train.trainer import Trainer + + # The SAME sampler the entry gates measured the horizon from. Building a + # second one here is how a stage ends up timing, projecting and training + # on three different batch streams. + sampler, dataset = build_training_sampler(cfg) + + def adapter(m, batch): + rows = [dataset[i] for i in batch] + collated = collate_sft(rows, pad_token_id=cfg.data.pad_token_id).to("cuda") + return make_sft_loss_adapter(require_fused=True)(m, (collated.input_ids, collated.labels)) + + run_dir = os.path.join(cfg.paths.checkpoint_dir, "sft") + extras = {} + if cfg.model.arch == "qwen3": + from align.train.qwen3_stage import qwen3_trainer_extras + + extras = qwen3_trainer_extras(cfg, model, adapter, run_dir) + trainer = Trainer( + cfg=cfg, + # LoRA: the trunk-movement probe must watch the adapter's B matrices (the frozen base can + # never move and A provably does not on step one), so it gets a view, not the raw model. + model=extras.pop("probe_model", model), + sampler=sampler, + loss_adapter=adapter, + run_dir=run_dir, + stage="sft", + **extras, + ) + trainer.resume_from_latest() + summary = trainer.train() + if cfg.model.arch == "mistral": + record_train_gate(ledger, trainer, summary) + else: + # The ledger's fixed sequence starts at the Mistral golden regression; a non-Mistral + # stage has not passed those gates, so recording "train" would raise (or be a lie). + logger.info("train finished (arch=%s): optimizer_step=%s", cfg.model.arch, summary.get("optimizer_step")) + return trainer + + +def record_train_gate(ledger, trainer, summary) -> dict: + """Record the 'train' gate for what `trainer.train()` ACTUALLY did. + + THE STOPPED-RECORDED-AS-PASSED DEFECT (H2 review, 2026-10-13). `train()` + returns normally after a COORDINATED STOP (SIGTERM, or train.max_wall_seconds) + -- it checkpoints and returns -- and the gate used to be recorded + `passed=True` unconditionally. `main` then (correctly) skipped the eval and + exited 128+n / 82, but the ledger on disk now certified a half-trained + checkpoint as a finished 'train' stage: anything that reads the ledger + instead of the exit status (a resumed `evaluate_and_gate`, an operator + reading `stage_gates.json`) was told the stage had trained. + + A stopped run is neither finished nor failed: it is RESUMABLE. The ledger + has no third state, so it is recorded `passed=False` with a `STOPPED + (resumable)` detail; a rerun resumes from the checkpoint and records the + real verdict at the end of the real run. + """ + step = summary.get("optimizer_step") + if getattr(trainer, "stop_exit_code", 0): + return ledger.record( + "train", + passed=False, + detail=f"STOPPED (resumable): {trainer.stop_reason} stop at optimizer_step={step}; " + "checkpointed, not finished", + ) + return ledger.record("train", passed=True, detail=str(step)) + + +def main(argv=None) -> int: + args = parse_args(argv) + + # FIRST: every env var torch reads at import, plus the heavy-job flock and + # the /apps path assertions. `heavy=True` also runs the host preflight; + # `rss_cfg` adds the full-FT RSS projection on top of the generic floor. + cfg = load_config(args.config, args.overrides) + apply_resource_guards(cpu_threads=8, heavy=True, label="sft-stage-11", rss_cfg=cfg) + + logging.basicConfig( + level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s" + ) + + # BEFORE anything touches CUDA. THE DUPLICATE-GPU DEFECT: torchrun starts + # four processes but assigns no devices, so without this every rank's + # `.cuda()` lands on device 0 and NCCL dies with "Duplicate GPU detected: + # rank 2 and rank 0 both on CUDA device 1b0" -> ncclInvalidUsage. Section-01 + # wrote `init_distributed` for exactly this (init_process_group + + # set_device(local_rank)); this script simply never called it. + from align.utils.dist import init_distributed + + run_dir = os.path.join(cfg.paths.checkpoint_dir, "sft") + # run_dir: the NCCL flight-recorder dump lands under it (never /tmp). + init_distributed(run_dir=run_dir) + + # Microseconds, and it converts the whole duplicate-GPU class into a named + # launch failure instead of an ncclInvalidUsage 400 lines into a forward. + from align.train.stage_gates import assert_distinct_devices + + assert_distinct_devices() + ledger = GateLedger(run_dir, eval_source_hash=eval_source_hash()) + + banner = run_gates(cfg, ledger) + from align.registry import render_banner + + print(render_banner(banner, title=f"align SFT stage :: {cfg.run_name}")) + + model = build_model(cfg) + trainer = train(cfg, ledger, model, banner) + if trainer.stop_exit_code: + # A coordinated stop (SIGTERM / train.max_wall_seconds) saved and + # returned: the stage is NOT finished, so do not fall into its + # post-train eval or announce completion. 82 = wall budget (the + # supervisor will not relaunch it), 128+n = the signal. + print(f"{trainer.stop_reason} stop: checkpointed, exiting {trainer.stop_exit_code}", file=sys.stderr) + return trainer.stop_exit_code + + # BARRIER, then RELEASE, then eval -- both are load-bearing: + # 1. train() returns minutes earlier on ranks 1..N-1 than on rank 0 (the + # sole checkpoint writer). Without this barrier they raced into the + # eval, OOM'd building a second 7B, and torchrun SIGTERM'd rank 0 + # MID-SAVE -- whose cleanup handler then deleted the .partial. The + # epoch's tail died twice this way. + # 2. run_stage_eval builds its OWN model from the checkpoint; the + # training replica + optimizer must be freed first or the engine's + # KV-cache preflight refuses (1.2GB free vs 4.8GB required). + from align.utils.dist import barrier + + barrier() + checkpoint = trainer.latest_checkpoint() + if cfg.model.arch != "mistral": + # The post-train MATH-500/GSM8K gate certifies the Mistral stack; a qwen3 checkpoint is + # judged by the Lean loop. Say so instead of silently skipping. + print(f"SFT stage complete (arch={cfg.model.arch}; no MATH/GSM8K gate applies). Checkpoint: {checkpoint}") + return 0 + trainer.release() # model, optimizer AND the probe's references (U01-11) + del model + import gc + + import torch + + gc.collect() + torch.cuda.empty_cache() + decision = evaluate_and_gate(cfg, ledger, checkpoint=checkpoint) + if not decision.passed: + print(decision.halt_message, file=sys.stderr) + return 1 + + print(f"SFT stage complete. DPO base checkpoint: {checkpoint}") + return 0 + + +if __name__ == "__main__": + from align.exit_codes import run_with_exit_codes + + raise SystemExit(run_with_exit_codes(main)) diff --git a/code/scripts/supervise.py b/code/scripts/supervise.py new file mode 100644 index 0000000000000000000000000000000000000000..74a4d0f42a01a4205f1e62d7916ad3f6ecad1b0d --- /dev/null +++ b/code/scripts/supervise.py @@ -0,0 +1,76 @@ +#!/usr/bin/env python +"""Bounded-retry launch supervisor CLI (stdlib only; no torch import). + + python scripts/supervise.py --run-dir RUN [--max-retries 2] [--backoff 30] \ + -- torchrun --nproc_per_node=4 scripts/sft.py configs/sft.toml + +Everything after `--` is the child command. See align/utils/supervisor.py for +the retry rules (progress + preflight + non-HALT exit code) and why each exists. +`--run-dir` must be the directory the Trainer writes history.jsonl and +checkpoints/ into; attempts are logged to RUN/supervisor.jsonl. +""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +# The package is not installed in the venv; without this a bare `python scripts/supervise.py` (the documented +# launch line) dies with ModuleNotFoundError unless PYTHONPATH is exported (same fix as scripts/sft.py). +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from align.utils.supervisor import Supervisor + + +def _checkpoint_validator(): + """`checkpoint_problem` from the Trainer, imported lazily (it needs torch). + + The supervisor module itself stays stdlib-only; THIS entry point is where + the cost is paid, once, so that "progress" means a checkpoint `latest_checkpoint` + would actually resume from. If torch is unavailable the supervisor still + runs, falling back to name-only progress, and says so on stderr. + """ + try: + from align.train.trainer import checkpoint_problem + except Exception as e: # noqa: BLE001 + print(f"supervisor: checkpoint validation unavailable ({type(e).__name__}: {e}); " + "progress is judged by directory name only", file=sys.stderr) + return None + return checkpoint_problem + + +def main(argv=None) -> int: + ap = argparse.ArgumentParser(description=__doc__.split("\n")[0]) + ap.add_argument("--run-dir", required=True) + ap.add_argument("--max-retries", type=int, default=2) + ap.add_argument("--backoff", type=float, default=30.0, help="base seconds; doubles per retry") + ap.add_argument("--grace", type=float, default=20.0, help="SIGTERM->SIGKILL grace seconds") + ap.add_argument( + "--no-validate-checkpoints", + action="store_true", + help="count a checkpoint dir as progress by NAME only (skips importing torch for the " + "completeness check); a truncated newest dir then counts as progress", + ) + ap.add_argument( + "--retry-on-stop", + action="store_true", + help="also relaunch a run that was stopped by SIGTERM/SIGINT/SIGHUP (exit 143/130/129). " + "Default: an operator's stop is a decision, not a fault, and ends the supervisor", + ) + ap.add_argument("cmd", nargs=argparse.REMAINDER, help="-- child command") + a = ap.parse_args(argv) + cmd = a.cmd[1:] if a.cmd and a.cmd[0] == "--" else a.cmd + if not cmd: + ap.error("no child command given (put it after --)") + res = Supervisor(cmd, a.run_dir, max_retries=a.max_retries, + backoff_s=a.backoff, grace_s=a.grace, retry_on_stop=a.retry_on_stop, + checkpoint_validator=None if a.no_validate_checkpoints + else _checkpoint_validator()).run() + print(f"supervisor: {res.reason} after {res.attempts} attempt(s), " + f"child exits {res.child_codes}", file=sys.stderr) + return res.exit_code + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/code/tests/conftest.py b/code/tests/conftest.py new file mode 100644 index 0000000000000000000000000000000000000000..fb95ba0893c77161e38d29c8361004dad1cc7b01 --- /dev/null +++ b/code/tests/conftest.py @@ -0,0 +1,195 @@ +"""Shared fixtures + house test conventions. + +Conventions (inherited from the predecessor's 65-test suite): + * regression tests are NAMED AFTER THE INCIDENT they prevent, incident in + the docstring; + * every guard is tested in BOTH directions (fires when it must, silent when + it must); + * undefined statistics assert NaN, never 0; + * every test file's docstring carries its run command; + * end-to-end tests use a tiny random-init model on CPU with >=2 layers and + >=2 KV groups (1-layer/1-group models cannot surface interior/GQA bugs). + +Run all CPU tests: uv run pytest -m "not gpu" +""" + +from __future__ import annotations + +import os + +import pytest + +REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + + +def heavy_job_blocking_gpu_test(lock_path: str | None = None) -> str | None: + """The label of ANOTHER process's heavy-job lock, or None if the lock is free (or ours). + + THE UNGUARDED-GPU-TEST DEFECT (P2 R2 review). The `gpu` marker is a label, not a guard: `pytest -k ` + without `-m "not gpu"` ran CUDA tests on GPU 0 while a 4-rank training run held the cards (21.7 GB in use of 23); + the test OOM'd, and the only reason the run survived is that its allocator had already reached steady state. A gpu + test is a heavy job like any other and must not start beside one. Probes the flock and releases it at once, so tests + that take the lock themselves (`apply_resource_guards(heavy=True)`) behave exactly as before.""" + import fcntl + + from align.utils import guard + + path = lock_path or guard.LOCK_PATH + if not os.path.exists(path): + return None + with open(path, "a+") as fh: + fh.seek(0) + label = fh.read().strip() + try: + fcntl.flock(fh.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) + except BlockingIOError: + return None if f"pid={os.getpid()}" in label.split() else (label or "unknown holder") + fcntl.flock(fh.fileno(), fcntl.LOCK_UN) + return None + + +@pytest.fixture(autouse=True) +def _gpu_tests_do_not_start_beside_a_heavy_job(request): + if request.node.get_closest_marker("gpu") is None: + return + holder = heavy_job_blocking_gpu_test() + if holder is not None: + pytest.skip(f"gpu test not started: the heavy-job lock is held by {holder!r} (never two heavy jobs at once)") + + +@pytest.fixture +def tmp_workdir(tmp_path_factory): + """A scratch dir that honors the path policy: under the repo on /apps, + never /tmp (which is noexec AND on a different device on this box).""" + base = os.path.join(REPO_ROOT, ".tmp", "tests") + os.makedirs(base, exist_ok=True) + import tempfile + + d = tempfile.mkdtemp(dir=base) + yield d + # best-effort cleanup; stale dirs under .tmp/tests are harmless + import shutil + + shutil.rmtree(d, ignore_errors=True) + + +@pytest.fixture +def tiny_model_constants(): + """Fixture constants for the CPU tiny-model harness, consumed by + section-04's model tests. >=2 layers and >=2 KV groups is a hard house + rule: a 1-layer or 1-group model cannot surface interior-layer or + GQA-indexing bugs.""" + return { + "n_layers": 2, + "hidden_size": 64, + "n_heads": 4, + "n_kv_heads": 2, + "head_dim": 16, + "intermediate_size": 128, + "vocab_size": 256, + "max_seq_len": 128, + "rope_theta": 1e6, + "seed": 0, + } + + +#: HF-cache snapshot of Mistral-7B-v0.3 (config + 3 shards + tokenizer). NOT present +#: on this box as of 2026-10-12 (hub/ has no models--mistralai--* entry); anything that +#: needs the WEIGHTS or the snapshot config.json must still skip when it is absent. +V03_TOKENIZER_SNAPSHOT = ( + "/apps/voxtral/.cache/huggingface/hub/models--mistralai--Mistral-7B-v0.3/" + "snapshots/caa1feb0e54d415e2df31207e5f4e273e33509b1" +) + +#: Tokenizer files vendored in-repo (byte-identical to the snapshot's; see +#: assets/tokenizer/mistral-7b-v0.3/PROVENANCE.md). Durable: the legacy SFT +#: checkpoint dirs these were copied from rotate. +V03_TOKENIZER_VENDORED = os.path.join(REPO_ROOT, "assets", "tokenizer", "mistral-7b-v0.3") + + +def resolve_v03_tokenizer_dir() -> str | None: + """First directory holding a v0.3 tokenizer.json: ALIGN_TOKENIZER_SRC override, + the HF snapshot, then the vendored copy. None if none exists.""" + env = os.environ.get("ALIGN_TOKENIZER_SRC") + for c in ([env] if env else []) + [V03_TOKENIZER_SNAPSHOT, V03_TOKENIZER_VENDORED]: + if c and os.path.isfile(os.path.join(c, "tokenizer.json")): + return c + return None + + +#: Resolved at import so tests can `isdir`-style gate on it. +V03_TOKENIZER_DIR = resolve_v03_tokenizer_dir() + +#: Predecessor eval sets used as REAL gold sources for the ceiling tests +#: (MATH golds are LaTeX, GSM8K golds are bare integers -- they break +#: extraction in different ways; synthetic '42' golds prove nothing). +#: Section-03 regenerates these under nmo with hash receipts; these paths are +#: parameterized into the ceiling functions, never hardcoded inside nmo/. +GOLD_SOURCES = ( + "/apps/development/rlhf/mistral-prm-ppo/artifacts/data/eval/math500.json", + "/apps/development/rlhf/mistral-prm-ppo/artifacts/data/eval/gsm8k_test.json", +) + + +@pytest.fixture(scope="session") +def v03_tokenizer(): + """The real Mistral-7B-v0.3 tokenizer, loaded by the in-repo align.tokenizer + (stdlib-only; its parity with transformers is certified by the oracle test in + tests/test_tokenizer.py). CPU-only, fast. Skips LOUDLY if no copy exists -- + never downloads.""" + d = resolve_v03_tokenizer_dir() + if d is None: + pytest.skip( + f"no v0.3 tokenizer.json at {V03_TOKENIZER_SNAPSHOT} or {V03_TOKENIZER_VENDORED}; " + "boundary certification cannot run against a fake tokenizer" + ) + from align.tokenizer import load_tokenizer + + return load_tokenizer(d) + + +@pytest.fixture +def clean_script_env(): + """Environment for tests that SPAWN nmo scripts in a subprocess. + + The launching harness exports TMPDIR pointing at a scratchpad on the + HOME device; the section-01 path guard correctly REFUSES that env at + script startup (it did so in production on this guard's first day). + Tests that spawn scripts must therefore pass an env whose TMPDIR is on + /apps — via this fixture, so no lane rediscovers the leak. The guard's + own refusal test poisons TMPDIR deliberately and must NOT use this.""" + env = dict(os.environ) + env["TMPDIR"] = os.path.join(REPO_ROOT, ".tmp") + os.makedirs(env["TMPDIR"], exist_ok=True) + return env + + +@pytest.fixture +def fake_host_state(monkeypatch): + """Factory: monkeypatch align.utils.guard.host_state to a fixed reading. + Guard tests must never depend on real host load (CI would flake exactly + when the shared box is busy -- the moment the guards matter most).""" + + def _set( + mem_available_gb: float, + load_per_core: float = 0.5, + disk_in_flight: int = 0, + d_state: int = 0, + ): + from align.utils import guard + + state = { + "ram_total_gb": 186.0, + "ram_available_gb": mem_available_gb, + "swap_total_gb": 0.0, + "load1": load_per_core * 48, + "load5": 0.0, + "n_cores": 48, + "load_per_core": load_per_core, + "disk_in_flight": disk_in_flight, + "d_state_procs": d_state, + } + monkeypatch.setattr(guard, "host_state", lambda: dict(state)) + return state + + return _set diff --git a/code/tests/test_adamw8bit.py b/code/tests/test_adamw8bit.py new file mode 100644 index 0000000000000000000000000000000000000000..41b2daf3d927b1cd97f145c5ca3fcbf4bdf71c24 --- /dev/null +++ b/code/tests/test_adamw8bit.py @@ -0,0 +1,1287 @@ +"""Block-wise 8-bit AdamW in pure torch: codebook, quantization, the optimizer, +resume, the factory, and the 2-rank gloo/FSDP2 behaviour. + +Run: .venv/bin/python -m pytest tests/test_adamw8bit.py -q -m "not gpu and not oracle" -p no:cacheprovider + +Why this optimizer exists (measured, do not re-derive): every shipped full-FT +config names `adamw8bit`; fp32 AdamW does not fit a 23GB A10G and +`OffloadAdamW` overshoots its pinned-staging cap 4x at 7.25B params / 4 ranks +and projects ~105GB host RSS against ~85GB available. See adamw8bit.py. + +House conventions: each regression is NAMED AFTER THE INCIDENT/FAILURE CLASS it +prevents, and each assertion is one a mutant that reintroduces the bug fails +(the mutation each test kills is stated in its docstring). +""" + +from __future__ import annotations + +import copy +import json +import math +import os +import socket +import sys +import traceback + +import pytest +import torch +import torch.multiprocessing as mp + +from align.train import quant8 +from align.train.adamw8bit import ( + BlockwiseAdamW8bit, + adamw_update_, + stochastic_round_to_bf16, +) + +HP = dict(lr=1e-3, betas=(0.9, 0.95), eps=1e-8, weight_decay=0.1) +B = 2048 + + +def _randn(*shape, seed=0, dtype=torch.float32): + g = torch.Generator().manual_seed(seed) + return torch.randn(*shape, generator=g, dtype=torch.float32).to(dtype) + + +def _param(*shape, seed=0, dtype=torch.float32): + return torch.nn.Parameter(_randn(*shape, seed=seed, dtype=dtype)) + + +def _set_grads(params, step, seed=7, scale=1.0, dtype=None): + g = torch.Generator().manual_seed(seed + step) + for p in params: + gr = torch.randn(p.shape, generator=g, dtype=torch.float32) * scale + p.grad = gr.to(dtype or p.dtype) + + +# --------------------------------------------------------------------------- +# the codebook +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("signed", [True, False]) +def test_codebook_is_a_sorted_256_entry_map_covering_the_unit_range(signed): + """Kills: a miscounted decade (255/257 entries -> uint8 overflow or a dead + code), a duplicate/unsorted entry (searchsorted requires strictly + increasing input and returns garbage otherwise), a map that does not reach + +-1 (the block's absmax element would clip), a missing exact zero (the + all-zero state could not be represented).""" + cb = quant8.create_dynamic_map(signed=signed) + assert cb.shape == (256,) and cb.dtype == torch.float32 + assert bool(torch.isfinite(cb).all()) + assert bool((cb[1:] > cb[:-1]).all()), "codebook must be STRICTLY increasing" + assert 0.0 in cb.tolist() + assert cb[quant8.zero_code(signed)].item() == 0.0 + assert cb.max().item() == 1.0 + if signed: + assert cb.min().item() <= -0.99, "signed map must reach (nearly) -1" + assert (cb < 0).sum().item() == (cb > 0).sum().item() - 1, ( + "signed map is symmetric except for the extra exact +1.0" + ) + else: + assert cb.min().item() == 0.0 and bool((cb >= 0).all()) + + +def test_codebook_is_dynamic_not_linear(): + """The point of the map: relative resolution across decades. A linear grid + has 1/255 absolute spacing, so values below ~4e-3 of absmax collapse to + zero; the dynamic map must still resolve 1e-5 of absmax to within a few + percent. Kills: replacing create_dynamic_map with linspace(-1, 1, 256).""" + for signed in (True, False): + cb = quant8.create_dynamic_map(signed) + codes = quant8.nearest_code(torch.tensor([1e-5, 1e-3, 0.5]), cb).long() + got = cb[codes] + rel = ((got - torch.tensor([1e-5, 1e-3, 0.5])).abs() / torch.tensor([1e-5, 1e-3, 0.5])) + assert bool((rel < 0.45).all()), (signed, got, rel) + assert rel[1] < 0.1 and rel[2] < 0.02 + # at least 6 decades of dynamic range above the smallest nonzero code + assert quant8.create_dynamic_map(False)[1].item() < 1e-6 + + +@pytest.mark.parametrize("signed", [True, False]) +def test_nearest_code_is_truly_the_nearest_including_edges_and_midpoints(signed): + """Kills: an off-by-one in the searchsorted/midpoint arithmetic (picking + the code ABOVE every x, or the wrong side of a midpoint), edge handling + (x below cb[0] / above cb[-1] must resolve to the end codes, not index -1 + or 256), and NaN-through-searchsorted. Brute-force argmin is the oracle.""" + cb = quant8.create_dynamic_map(signed) + lo = -1.5 if signed else -0.2 + g = torch.Generator().manual_seed(0) + x = torch.cat( + [ + torch.rand(20000, generator=g) * (1.5 - lo) + lo, + cb, # exactly on a code + (cb[1:] + cb[:-1]) / 2 + 1e-9, # just above each midpoint + (cb[1:] + cb[:-1]) / 2 - 1e-9, # just below each midpoint + torch.tensor([-1.0, 0.0, 1.0, -1e-30, 1e-30]), + ] + ) + got = quant8.nearest_code(x, cb) + assert got.dtype == torch.uint8 + want = (x.unsqueeze(1) - cb.unsqueeze(0)).abs().argmin(dim=1) + # Compare by DISTANCE, not by index: at a midpoint the two candidates tie + # to within fp32 resolution (~6e-8 absolute), and either is correct. A wrong + # side of a midpoint is off by a whole half-bin, orders of magnitude more. + d_got = (x - cb[got.long()]).abs() + d_want = (x - cb[want]).abs() + assert bool((d_got <= d_want * 1.001 + 1e-9).all()), "nearest_code returned a farther code" + + +# --------------------------------------------------------------------------- +# quantize -> dequantize +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("signed", [True, False]) +def test_roundtrip_error_is_bounded_on_gaussian_blocks(signed): + """Quantization error, measured as the optimizer sees it: relative to the + block's absmax (the scale the codes are normalised by). Median < 1e-2 and + max < 0.1 are the specified bounds; the measured values are ~2e-3 and + ~7e-3 (the top bin is 0.014 wide), so the bound has 10x headroom and a + broken codebook (linear, shifted, unsorted) blows through it. + + ELEMENT-wise relative error on the large elements (|x| > 0.1 absmax, where + the optimizer's update is decided) is also pinned: median < 2e-2 (measured + ~1.3e-2; a 64-bin decade cannot do better than ~half a bin of 0.014/0.2). + """ + x = _randn(8 * B, seed=1) + if not signed: + x = x.pow(2) # a second-moment-like, heavy-tailed, non-negative tensor + codes, absmax = quant8.quantize_blockwise(x, B, signed) + assert codes.dtype == torch.uint8 and codes.numel() == 8 * B and absmax.shape == (8,) + y = quant8.dequantize_blockwise(codes, absmax, x.numel(), B, signed) + scale = absmax.repeat_interleave(B) + abs_err = (x - y).abs() / scale + assert abs_err.median().item() < 1e-2 + assert abs_err.max().item() < 0.1 + big = x.abs() > 0.1 * scale + rel = ((x - y).abs() / x.abs())[big] + assert rel.median().item() < 2e-2 and rel.max().item() < 0.1 + + +def test_blockwise_isolates_an_outlier_to_its_own_block(): + """The reason for BLOCK-wise: one huge element must not coarsen the other + blocks. Kills: a single global absmax (every other block's median error + would be ~1e4x worse).""" + x = _randn(4 * B, seed=2) + x[5] = 1e4 + codes, absmax = quant8.quantize_blockwise(x, B, True) + y = quant8.dequantize_blockwise(codes, absmax, x.numel(), B, True) + clean = slice(B, 4 * B) + err = ((x[clean] - y[clean]).abs() / x[clean].abs().clamp_min(1e-3)).median() + assert err.item() < 2e-2, err + assert absmax[0].item() == 1e4 and absmax[1].item() < 6 + + +def test_zero_blocks_and_partial_last_block_roundtrip_exactly_without_nan(): + """Kills: dividing by a zero absmax (NaN codes), and a partial final block + that leaks its zero padding into the dequantized values or the numel.""" + x = torch.zeros(3 * B + 77) + x[B : B + 5] = torch.tensor([1.0, -2.0, 3.0, 0.0, 0.5]) + codes, absmax = quant8.quantize_blockwise(x, B, True) + assert codes.numel() == 4 * B + y = quant8.dequantize_blockwise(codes, absmax, x.numel(), B, True) + assert y.shape == x.shape and bool(torch.isfinite(y).all()) + assert torch.equal(y[:B], torch.zeros(B)) and torch.equal(y[2 * B :], torch.zeros(B + 77)) + assert absmax[0].item() == 0.0 and absmax[1].item() == 3.0 + + +def test_dequantize_indexes_with_a_gather_not_a_uint8_mask(): + """THE UINT8 INDEX TRAP: `cb[codes]` with a uint8 `codes` is a boolean-mask + index in torch, not a gather. dequantize must return cb[code] * absmax for + codes 0 and 255 (the two ends), not a mask-selected junk tensor.""" + cb = quant8.codebook(True, "cpu") + codes = torch.tensor([[0, 255, 127, 1]], dtype=torch.uint8) + out = quant8.dequantize_rows(codes, torch.tensor([2.0]), cb) + assert torch.equal(out, torch.stack([cb[0], cb[255], cb[127], cb[1]]).unsqueeze(0) * 2.0) + + +# --------------------------------------------------------------------------- +# the update rule +# --------------------------------------------------------------------------- + + +def test_fp32_fallback_for_small_tensors_is_bit_exact_vs_torch_adamw(): + """Tensors under `min_8bit_size` keep fp32 moments, and must match + `torch.optim.AdamW` BITWISE over many steps (decoupled decay, betas, + bias correction, eps all included). Kills: any reordering of the op + sequence (lerp_ vs hand EMA, folding bias-correction differently), a + dropped weight_decay, a mis-forwarded beta2 (0.95 vs torch's 0.999).""" + a = [_param(4, 5, seed=s) for s in range(3)] + [_param(4095, seed=9)] + b = [torch.nn.Parameter(p.detach().clone()) for p in a] + mine = BlockwiseAdamW8bit(a, **HP) + ref = torch.optim.AdamW(b, **HP) + assert all("m" in s for s in mine._state), "all four are below min_8bit_size=4096" + for step in range(40): + _set_grads(a, step) + for pa, pb in zip(a, b): + pb.grad = pa.grad.clone() + assert mine.step() is True + ref.step() + for pa, pb in zip(a, b): + assert torch.equal(pa, pb) + + +def test_threshold_is_below_min_8bit_size_inclusive_at_the_boundary(): + """bnb's rule: numel < min_8bit_size stays fp32; numel == min is 8-bit. + Kills: `<=` vs `<` drift.""" + opt = BlockwiseAdamW8bit([_param(4095), _param(4096)], lr=1e-3) + assert "m" in opt._state[0] and "m_codes" in opt._state[1] + assert opt._state[1]["m_codes"].dtype == torch.uint8 + + +def _quadratic(opt_fn, n=8192, steps=300, seed=1): + g = torch.Generator().manual_seed(seed) + t = torch.randn(n, generator=g) + curv = torch.rand(n, generator=g) * 9 + 1 + x = torch.nn.Parameter(torch.zeros(n)) + opt = opt_fn([x]) + loss0 = 0.5 * (curv * t**2).sum().item() + for _ in range(steps): + x.grad = None + loss = 0.5 * (curv * (x - t) ** 2).sum() + loss.backward() + opt.step() + return 0.5 * (curv * (x.detach() - t) ** 2).sum().item(), loss0 + + +def test_300_steps_on_a_convex_quadratic_converge_like_torch_adamw(): + """Both reach <0.5% of the initial loss, and the 8-bit endpoint is within + 1e-3 * L0 of fp32 AdamW's. Kills: state that does not accumulate (a + zeroed-every-step moment), a lost lr/bias-correction (never converges), a + codebook that collapses v to zero (the update explodes, m/eps).""" + hp = dict(lr=5e-2, betas=(0.9, 0.95), weight_decay=0.0) + l8, l0 = _quadratic(lambda p: BlockwiseAdamW8bit(p, **hp)) + l32, _ = _quadratic(lambda p: torch.optim.AdamW(p, **hp)) + assert l32 < 5e-3 * l0, (l32, l0) + assert l8 < 5e-3 * l0, (l8, l0) + assert abs(l8 - l32) < 1e-3 * l0, (l8, l32, l0) + + +def _linreg(opt_fn, steps=300): + g = torch.Generator().manual_seed(3) + d_in, d_out, n = 128, 64, 256 # weight has 8192 elements -> the 8-bit path + X = torch.randn(n, d_in, generator=g) + W_true = torch.randn(d_out, d_in, generator=g) / math.sqrt(d_in) + Y = X @ W_true.T + W = torch.nn.Parameter(torch.zeros(d_out, d_in)) + bias = torch.nn.Parameter(torch.zeros(d_out)) # 64 elements -> fp32 fallback + opt = opt_fn([W, bias]) + first = None + for _ in range(steps): + opt.zero_grad() if hasattr(opt, "zero_grad") else None + loss = ((X @ W.T + bias - Y) ** 2).mean() + first = loss.item() if first is None else first + loss.backward() + opt.step() + return loss.item(), first + + +def test_300_steps_of_tiny_linear_regression_match_torch_adamw(): + """Mixed 8-bit (weight) + fp32-fallback (bias) parameters in one + optimizer, on a real backward pass.""" + hp = dict(lr=1e-2, betas=(0.9, 0.95), weight_decay=0.0) + l8, first = _linreg(lambda p: BlockwiseAdamW8bit(p, **hp)) + l32, _ = _linreg(lambda p: torch.optim.AdamW(p, **hp)) + assert l32 < 1e-3 * first and l8 < 1e-3 * first, (l8, l32, first) + assert abs(l8 - l32) < 1e-3 * first, (l8, l32, first) + + +def test_first_step_is_exactly_lr_times_sign_because_bias_correction_is_applied(): + """After one step, m_hat/(sqrt(v_hat)+eps) == sign(g) (to eps), so every + element moves by exactly lr. The update uses the freshly computed fp32 + moments, not the requantized ones, so this holds on the 8-bit path to fp32 + rounding. Kills: missing bias correction (moves 0.1/sqrt(0.05) = 0.447 lr + with beta2=0.95), or an update computed from the dequantized-after-quantize + moments (off by the quantization error, ~1e-3 relative).""" + p = _param(B * 3, seed=4) + p0 = p.detach().clone() + p.grad = _randn(B * 3, seed=5) * 10 + opt = BlockwiseAdamW8bit([p], lr=1e-3, betas=(0.9, 0.95), weight_decay=0.0) + assert "m_codes" in opt._state[0] + opt.step() + moved = (p0 - p.detach()) * p.grad.sign() + assert torch.allclose(moved, torch.full_like(moved, 1e-3), rtol=1e-4, atol=1e-7), ( + moved.min().item(), + moved.max().item(), + ) + + +def test_weight_decay_is_decoupled_from_the_gradient(): + """With a ZERO gradient the moments stay zero, so the only thing that can + move a weight is the decay: p <- p * (1 - lr*wd), to fp32 rounding. A + COUPLED (L2) implementation adds wd*p to the gradient and Adam then + normalises it to a step of ~lr per element -- ~1e3x more movement. Kills: + coupled decay, and a dropped weight_decay (p unchanged).""" + p = _param(B * 2, seed=6) + p0 = p.detach().clone() + p.grad = torch.zeros_like(p) + opt = BlockwiseAdamW8bit([p], lr=1e-2, betas=(0.9, 0.95), weight_decay=0.5) + assert "m_codes" in opt._state[0] + opt.step() + assert torch.allclose(p.detach(), p0 * (1 - 1e-2 * 0.5), rtol=1e-6, atol=0.0) + assert opt._state[0]["m_absmax"].abs().max().item() == 0.0, "decay leaked into the moments" + + +def test_chunking_is_invisible_to_the_result(): + """Chunks are whole blocks, so processing one block at a time must equal + processing the tensor in one pass, BITWISE (params and every code/absmax). + Kills: a chunk boundary that is off by a block, a partial-block mishandled + at a chunk edge, state written to the wrong row range.""" + n = B * 5 + 123 + outs = [] + for chunk in (B, 2 * B, 10**9): + p = _param(n, seed=8) + opt = BlockwiseAdamW8bit([p], chunk_elems=chunk, **HP) + for step in range(6): + _set_grads([p], step) + opt.step() + outs.append((p.detach().clone(), opt.state_dict()["state"][0])) + for pw, st in outs[1:]: + assert torch.equal(pw, outs[0][0]) + for k in ("m_codes", "v_codes", "m_absmax", "v_absmax"): + assert torch.equal(st[k], outs[0][1][k]), k + + +def test_padding_of_a_partial_last_block_stays_the_zero_code(): + """n is not a multiple of the block: the pad tail of the last row must keep + quantizing to exact zero forever, or garbage moments would leak into the + last block's absmax (and thus its resolution).""" + n = 2 * B + 100 # 4196 elements: above min_8bit_size, so it is 8-bit + p = _param(n, seed=10) + opt = BlockwiseAdamW8bit([p], **HP) + for step in range(5): + _set_grads([p], step) + opt.step() + st = opt._state[0] + assert st["m_codes"].shape == (3, B) + assert bool((st["m_codes"][2, 100:] == quant8.zero_code(True)).all()) + assert bool((st["v_codes"][2, 100:] == quant8.zero_code(False)).all()) + + +def test_param_with_no_grad_is_skipped_and_its_step_count_does_not_advance(): + """Bias correction is per tensor: a tensor that sat out a step must not be + treated as having taken it. Also: no gradients anywhere is an error, not a + silent no-op.""" + a, b = _param(B, seed=1), _param(B, seed=2) + opt = BlockwiseAdamW8bit([a, b], **HP) + a.grad = _randn(B, seed=3) + b0 = b.detach().clone() + assert opt.step() is True + assert opt._state[0]["step"] == 1 and opt._state[1]["step"] == 0 + assert torch.equal(b, b0) + a.grad = None + with pytest.raises(RuntimeError, match="no gradients"): + opt.step() + + +def test_empty_local_shard_is_tolerated(): + """FSDP2 can hand a rank an EMPTY shard (a 1-row param over 2 ranks). It + must construct, step and round-trip rather than dividing by zero blocks.""" + e, p = torch.nn.Parameter(torch.zeros(0)), _param(B, seed=1) + opt = BlockwiseAdamW8bit([e, p], **HP) + e.grad, p.grad = torch.zeros(0), _randn(B, seed=2) + assert opt.step() is True + opt.load_state_dict(opt.state_dict()) + + +# --------------------------------------------------------------------------- +# grad norm, clipping, the skip gate +# --------------------------------------------------------------------------- + + +def test_grad_norm_is_recorded_before_clipping_and_clipping_is_applied(): + """THE PPO BLIND-TRIPWIRE INCIDENT's cousin: the norm a tripwire reads must + be the PRE-clip norm (post-clip it is <= max by construction and a + 10,000x explosion looks like 1.0). And the clip must actually rescale the + gradients that enter the update. + + Adam is scale-invariant on a single step, so this runs TWO steps (a big + clipped gradient, then a small one) and compares, bitwise, against an + unclipped optimizer fed the pre-clipped gradients by hand. Kills: norm + recorded after clipping; clipping not applied (the two trajectories + differ by the relative weight of step 1 in the moments); a coefficient + other than max/(norm+1e-6).""" + ps = [_param(B * 2, seed=1), _param(B, seed=2), _param(7, seed=3)] # 8-bit x2 + fp32 + qs = [torch.nn.Parameter(p.detach().clone()) for p in ps] + hp = dict(lr=1e-2, betas=(0.9, 0.95), weight_decay=0.0) + clipped = BlockwiseAdamW8bit(ps, max_grad_norm=1.0, **hp) + manual = BlockwiseAdamW8bit(qs, max_grad_norm=None, **hp) + + for step, scale in enumerate((100.0, 0.01)): + _set_grads(ps, step, scale=scale) + true_norm = float(torch.sqrt(sum(p.grad.double().pow(2).sum() for p in ps))) + for q, p in zip(qs, ps): + q.grad = p.grad.clone() + if true_norm > 1.0: + coef = 1.0 / (true_norm + 1e-6) + for q in qs: + q.grad = q.grad * coef + assert clipped.step() is True + manual.step() + assert clipped.last_grad_norm == pytest.approx(true_norm, rel=1e-6), ( + "last_grad_norm must be the GLOBAL PRE-CLIP norm" + ) + assert clipped.last_grad_norm < 1.0 # the second step was below the cap: not clipped + for p, q in zip(ps, qs): + assert torch.equal(p, q) + + +def test_nonfinite_gradient_skips_the_step_and_leaves_every_byte_of_state_alone(): + """A NaN/inf discovered after the moments were updated has already poisoned + the second moment. The verdict must come BEFORE any state is touched, the + norm must read NaN (a 0.0 would read as 'converged'), and a good step after + the skip must work. Kills: checking finiteness after the update; counting + the skipped step in the bias-correction step number; returning True.""" + p = _param(B * 2, seed=1) + opt = BlockwiseAdamW8bit([p], **HP) + _set_grads([p], 0) + opt.step() + snap, p0 = opt.state_dict(), p.detach().clone() + for bad in (float("nan"), float("inf"), float("-inf")): + p.grad = _randn(B * 2, seed=3) + p.grad[17] = bad + assert opt.step() is False + assert math.isnan(opt.last_grad_norm) + assert torch.equal(p, p0) + assert opt.skipped_steps == 3 + now = opt.state_dict() + assert now["state"][0]["step"] == snap["state"][0]["step"] == 1 + for k in ("m_codes", "v_codes", "m_absmax", "v_absmax"): + assert torch.equal(now["state"][0][k], snap["state"][0][k]), k + _set_grads([p], 1) + assert opt.step() is True and not torch.equal(p, p0) + + +def test_second_moment_underflow_cannot_produce_an_m_over_eps_update(): + """THE ZERO-SECOND-MOMENT CATASTROPHE. In a block with one huge gradient + (1.0) the tiny ones (1e-6) have v/absmax ~ 1e-12, below the smallest + nonzero code, so their STORED second moment is exactly 0. If the update + were computed from stored moments, it would be m / (sqrt(0) + 1e-8) -- 100x + lr per step, i.e. the weight explodes. The update must use the FRESH fp32 + moments (quantization is a storage format), which bounds the per-step move + by lr / sqrt(1 - beta2) (= 4.47 lr here; measured 2.1 lr). + + Kills: computing the update after requantization; a codebook whose zero code + swallows small v.""" + p = torch.nn.Parameter(torch.zeros(4096)) + opt = BlockwiseAdamW8bit([p], lr=1e-3, betas=(0.9, 0.95), weight_decay=0.0) + g = torch.full((4096,), 1e-6) + g[0] = 1.0 + worst = 0.0 + for _ in range(20): + before = p.detach().clone() + p.grad = g.clone() + opt.step() + worst = max(worst, (p.detach() - before).abs().max().item()) + assert (opt._state[0]["v_codes"][0, 1:] == 0).all(), "precondition: v really underflowed" + assert worst <= 1e-3 / math.sqrt(1 - 0.95) * 1.05, worst / 1e-3 + + +# --------------------------------------------------------------------------- +# bf16 parameters, no master copy +# --------------------------------------------------------------------------- + + +def test_stochastic_rounding_is_unbiased_and_exact_on_representable_values(): + """E[SR(x)] == x, and a value already representable in bf16 is unchanged. + Kills: round-down (biased), round-to-nearest, a carry that corrupts the + sign/exponent for negative numbers.""" + gen = torch.Generator().manual_seed(0) + x = torch.full((200000,), 1.0 + 3e-4) # between 1.0 and 1.0078125 + out = stochastic_round_to_bf16(x, gen).float() + assert set(out.unique().tolist()) <= {1.0, 1.0078125} + assert out.mean().item() == pytest.approx(1.0 + 3e-4, abs=3e-5) + neg = stochastic_round_to_bf16(-x, gen).float() + assert neg.mean().item() == pytest.approx(-(1.0 + 3e-4), abs=3e-5) + exact = torch.tensor([0.0, 1.0, -2.5, 0.15625, 3.0e38 / 2**20]).to(torch.bfloat16).float() + assert torch.equal(stochastic_round_to_bf16(exact, gen).float(), exact) + + +def test_bf16_params_without_a_master_copy_still_move_under_tiny_updates(): + """THE NO-MASTER-COPY DEFECT (the silent-no-train class). A bf16 weight of + 1.0 has spacing 7.8e-3; an Adam step of lr=1e-4 is 80x below half a ulp, so + round-to-nearest returns 1.0 EVERY step and the run trains nothing while + its loss and grad-norm logs look healthy. Stochastic rounding makes the + parameter move by the right amount in expectation: 200 steps of lr 1e-4 on a + constant gradient move the mean by ~0.02. + + Both directions: nearest rounding really is stuck (the control that makes + this test mean something), and SR really moves. Kills: stochastic_rounding + ignored / the write-back always rounding to nearest.""" + def run(sr): + p = torch.nn.Parameter(torch.ones(B * 4, dtype=torch.bfloat16)) + opt = BlockwiseAdamW8bit( + [p], lr=1e-4, betas=(0.9, 0.95), weight_decay=0.0, stochastic_rounding=sr + ) + for _ in range(200): + p.grad = torch.ones_like(p) + opt.step() + return p.detach().float().mean().item() + + assert run(False) == 1.0, "control: round-to-nearest should be stuck" + assert run(True) == pytest.approx(1.0 - 200 * 1e-4, abs=2.5e-3) + + +def test_bf16_param_with_fp32_grad_is_accepted(): + """The fp32-reduce policy hands back fp32 grads for bf16 params.""" + p = torch.nn.Parameter(_randn(B * 2, seed=1, dtype=torch.bfloat16)) + if hasattr(p, "grad_dtype"): # newer torch refuses a grad dtype != param dtype by default + p.grad_dtype = None + p.grad = _randn(B * 2, seed=2) + p0 = p.detach().clone() + BlockwiseAdamW8bit([p], lr=1e-2).step() + assert p.dtype == torch.bfloat16 and not torch.equal(p, p0) + + +# --------------------------------------------------------------------------- +# resume +# --------------------------------------------------------------------------- + + +def _train(opt, params, steps, start=0, dtype=None): + for step in range(start, start + steps): + _set_grads(params, step, dtype=dtype) + opt.step() + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_resume_from_state_dict_is_bit_exact(dtype): + """Save at step 5, keep training the original; load into a FRESH optimizer + over a copy of the step-5 weights and train the same 5 steps: weights and + every code/absmax byte must be identical. bf16 covers the stochastic + rounding RNG (its generator state is part of the checkpoint). Kills: a + state_dict that drops the step count, the absmax, the SR RNG, or the + hyper-parameters; a load that re-quantizes instead of copying.""" + mk = lambda: [_param(B * 3, seed=1, dtype=dtype), _param(300, seed=2, dtype=dtype), + _param(B + 9, seed=3, dtype=dtype)] + a = mk() + opt_a = BlockwiseAdamW8bit(a, **HP, seed=5) + _train(opt_a, a, 5) + sd = opt_a.state_dict() + b = [torch.nn.Parameter(p.detach().clone()) for p in a] + opt_b = BlockwiseAdamW8bit(b, **HP, seed=999) # a different seed: must be overwritten + opt_b.load_state_dict(sd) + _train(opt_a, a, 5, start=5) + _train(opt_b, b, 5, start=5) + for pa, pb in zip(a, b): + assert torch.equal(pa, pb) + sa, sb = opt_a.state_dict(), opt_b.state_dict() + for x, y in zip(sa["state"], sb["state"]): + assert x.keys() == y.keys() + for k in x: + assert torch.equal(torch.as_tensor(x[k]), torch.as_tensor(y[k])), k + + +def test_state_dict_is_a_cpu_snapshot_not_a_live_view(): + """THE SHARED-STATE ALIASING INCIDENT: `.to('cpu')` on a CPU tensor returns + the same object, so a 'snapshot' built that way is mutated by the next + step -- uint8 codes that no comparison flags as drifting. And loading one + dict into two optimizers must not make them share state. Kills: missing + copy=True in state_dict; a load that stores references.""" + p = _param(B * 2, seed=1) + opt = BlockwiseAdamW8bit([p], **HP) + _train(opt, [p], 2) + sd = opt.state_dict() + frozen = copy.deepcopy(sd) + for k in ("m_codes", "v_codes", "m_absmax", "v_absmax"): + assert sd["state"][0][k].device.type == "cpu" + assert sd["state"][0][k].data_ptr() != opt._state[0][k].data_ptr() + _train(opt, [p], 3, start=2) + for k in ("m_codes", "v_codes", "m_absmax", "v_absmax"): + assert torch.equal(sd["state"][0][k], frozen["state"][0][k]), f"{k}: snapshot drifted" + assert sd["state"][0]["step"] == 2 + + q1, q2 = _param(B * 2, seed=1), _param(B * 2, seed=1) + o1, o2 = BlockwiseAdamW8bit([q1], **HP), BlockwiseAdamW8bit([q2], **HP) + o1.load_state_dict(sd) + o2.load_state_dict(sd) + assert o1._state[0]["m_codes"].data_ptr() != o2._state[0]["m_codes"].data_ptr() + assert o1._state[0]["m_codes"].data_ptr() != sd["state"][0]["m_codes"].data_ptr() + + +def test_state_dict_survives_a_pickle_roundtrip(): + """It rides `gather_object` (pickle) inside the checkpoint: it must hold no + live handles (generators, devices) and stay byte-identical.""" + import pickle + + p = _param(B * 2, seed=1, dtype=torch.bfloat16) + opt = BlockwiseAdamW8bit([p], **HP) + _train(opt, [p], 3) + sd = opt.state_dict() + back = pickle.loads(pickle.dumps(sd)) # round-trips this test's OWN freshly built dict; no untrusted input + q = torch.nn.Parameter(p.detach().clone()) + other = BlockwiseAdamW8bit([q], **HP) + other.load_state_dict(back) + _train(opt, [p], 2, start=3) + _train(other, [q], 2, start=3) + assert torch.equal(p, q) + + +def test_resume_refuses_a_different_world_size_layout_or_format_and_changes_nothing(): + """The state is SHARD-shaped. Loading it under another world size aligns + rank r's moments to a different slice of the parameter and training + continues plausibly on corrupted state. A refused load must also leave the + optimizer untouched (validate everything, then write). Kills: dropping the + world_size check; checking numel only after partially copying.""" + p = _param(B * 2, seed=1) + opt = BlockwiseAdamW8bit([p], **HP) + _train(opt, [p], 2) + sd = opt.state_dict() + + fresh_p = _param(B * 2, seed=1) + fresh = BlockwiseAdamW8bit([fresh_p], **HP) + before = fresh.state_dict() + + with pytest.raises(ValueError, match="world_size"): + fresh.load_state_dict(dict(sd, world_size=sd["world_size"] + 4)) + wrong_n = copy.deepcopy(sd) + wrong_n["state"][0]["numel"] += 1 + with pytest.raises(ValueError, match="numel"): + fresh.load_state_dict(wrong_n) + with pytest.raises(ValueError, match="block_size"): + fresh.load_state_dict(dict(sd, block_size=1024)) + with pytest.raises(ValueError, match="format"): + fresh.load_state_dict({"inner": {}, "state": {}}) + bad_shape = copy.deepcopy(sd) + bad_shape["state"][0]["m_absmax"] = torch.zeros(5) + with pytest.raises(ValueError, match="m_absmax"): + fresh.load_state_dict(bad_shape) + other = BlockwiseAdamW8bit([_param(B * 2), _param(B * 2)], **HP) + with pytest.raises(ValueError, match="tensors"): + other.load_state_dict(sd) + + after = fresh.state_dict() + assert after["state"][0]["step"] == before["state"][0]["step"] == 0 + for k in ("m_codes", "v_codes", "m_absmax", "v_absmax"): + assert torch.equal(after["state"][0][k], before["state"][0][k]) + + +# --------------------------------------------------------------------------- +# memory budget +# --------------------------------------------------------------------------- + + +def test_state_bytes_match_the_2n_plus_8n_over_block_budget(): + """The RSS/VRAM projection must be HONEST: for block-aligned tensors the + real footprint is exactly 2*n + 8*n/block bytes; small tensors are 8*n; and + `projected_state_bytes` (usable before allocating anything) agrees with + what was actually allocated. Kills: an absmax kept per element, fp32 + codes, or a projection that forgets the fp32 fallback tensors.""" + n_big = 100 * B + params = [_param(n_big, seed=1), _param(300, seed=2), _param(5 * B + 17, seed=3)] + opt = BlockwiseAdamW8bit(params, **HP) + assert opt._state[0]["m_codes"].dtype == torch.uint8 + assert opt._state[0]["m_absmax"].dtype == torch.float32 + assert opt._state[0]["m_absmax"].numel() == 100 + big_only = BlockwiseAdamW8bit([_param(n_big)], **HP).state_bytes() + assert big_only == 2 * n_big + 8 * n_big // B + got = opt.state_bytes() + want = BlockwiseAdamW8bit.projected_state_bytes([p.numel() for p in params]) + assert got == want + assert got < 0.26 * sum(p.numel() for p in params) * 8, "must be ~4x smaller than fp32 moments" + # the headline number from the module docstring: 7.248e9 params over 4 ranks + rank_numel = int(7_248_023_552 / 4) + proj = BlockwiseAdamW8bit.projected_state_bytes([rank_numel]) + assert proj / rank_numel == pytest.approx(2 + 8 / B, rel=1e-3) + assert 3.3 < proj / 2**30 < 3.5 + + +def test_describe_placement_reports_uint8_codes_on_the_param_device(): + opt = BlockwiseAdamW8bit([_param(B * 2)], **HP) + ent = {e["key"]: e for e in opt.describe_placement()["state"]} + assert ent["m_codes"]["dtype"] == str(torch.uint8) and ent["m_codes"]["device"] == "cpu" + assert ent["v_absmax"]["dtype"] == str(torch.float32) + + +# --------------------------------------------------------------------------- +# validation +# --------------------------------------------------------------------------- + + +def test_constructor_refuses_bad_inputs(): + p = _param(B) + with pytest.raises(ValueError, match="empty"): + BlockwiseAdamW8bit([], lr=1e-3) + with pytest.raises(ValueError, match="twice"): + BlockwiseAdamW8bit([p, p], lr=1e-3) + with pytest.raises(ValueError, match="betas"): + BlockwiseAdamW8bit([p], lr=1e-3, betas=(0.9, 1.0)) + with pytest.raises(ValueError, match="max_grad_norm"): + BlockwiseAdamW8bit([p], lr=1e-3, max_grad_norm=0.0) + with pytest.raises(ValueError, match="block_size"): + BlockwiseAdamW8bit([p], lr=1e-3, block_size=0) + with pytest.raises(ValueError, match="floating"): + BlockwiseAdamW8bit([torch.zeros(B, dtype=torch.int64)], lr=1e-3) + with pytest.raises(ValueError, match="contiguous"): + BlockwiseAdamW8bit([torch.nn.Parameter(torch.zeros(64, 64).t())], lr=1e-3) + q = _param(B) + opt = BlockwiseAdamW8bit([q], lr=1e-3) + q.data = torch.zeros(B + 1) # a re-sharded parameter: state no longer fits + q.grad = torch.zeros(B + 1) + with pytest.raises(ValueError, match="elements"): + opt.step() + + +def test_lr_written_through_param_groups_is_what_the_step_uses(): + """`Trainer._apply_scheduled_lr` writes the scheduled lr into + `param_groups[*]['lr']`; a step that cached the constructor lr would train + at a constant rate while logging a decaying one (the scheduler-never-applied + incident). Kills: reading self.lr instead of the group.""" + p, q = _param(B, seed=1), _param(B, seed=1) + a = BlockwiseAdamW8bit([p], lr=1e-3, weight_decay=0.0) + b = BlockwiseAdamW8bit([q], lr=7e-3, weight_decay=0.0) + a.param_groups[0]["lr"] = 7e-3 + for o, x in ((a, p), (b, q)): + x.grad = _randn(B, seed=2) + o.step() + assert torch.equal(p, q) + + +# --------------------------------------------------------------------------- +# the factory, the shim, the shipped configs +# --------------------------------------------------------------------------- + + +def _tiny_params(): + return [_param(64, 64, seed=1), _param(64, seed=2)] + + +@pytest.mark.parametrize("name", ["sft.toml", "pretrain.toml"]) +def test_make_optimizer_builds_adamw8bit_for_the_shipped_full_ft_configs(name): + """THE LAUNCH BLOCKER (review_trainer.md): `load_config(sft.toml) -> + make_optimizer(...)` raised NotImplementedError, so no full-FT stage could + start. The shipped configs must now construct, tagged for the banner with + the 8-bit footprint, and hold no host-side master.""" + from align.config import load_config + from align.train.offload_adamw import make_optimizer + + cfg = load_config(os.path.join(os.path.dirname(__file__), "..", "configs", name)) + assert cfg.optimizer.name == "adamw8bit" + params = _tiny_params() + opt = make_optimizer(cfg, params) + assert isinstance(opt, BlockwiseAdamW8bit) + assert opt.banner_tag.startswith("adamw8bit") and "8-bit" in opt.banner_tag + assert opt.param_groups[0]["lr"] == cfg.optimizer.lr + assert opt.param_groups[0]["betas"] == tuple(cfg.optimizer.betas) + assert opt.param_groups[0]["weight_decay"] == cfg.optimizer.weight_decay + for p in params: + p.grad = torch.ones_like(p) + assert opt.step() is True and opt.skipped_steps == 0 + + +def test_make_optimizer_adamw8bit_is_wired_with_the_config_hyperparameters(): + """Kills: a factory that drops betas/eps/weight_decay on the floor (the + parity test passes against defaults but not against the run config).""" + import types + + from align.train.offload_adamw import make_optimizer + + o = types.SimpleNamespace( + name="AdamW8bit", lr=3e-4, betas=[0.8, 0.9], eps=1e-6, weight_decay=0.25 + ) + opt = make_optimizer(types.SimpleNamespace(optimizer=o), _tiny_params()) + g = opt.param_groups[0] + assert (g["lr"], g["betas"], g["eps"], g["weight_decay"]) == (3e-4, (0.8, 0.9), 1e-6, 0.25) + assert g["max_grad_norm"] is None + o.max_grad_norm = 2.0 + assert make_optimizer(types.SimpleNamespace(optimizer=o), _tiny_params()).param_groups[0][ + "max_grad_norm" + ] == 2.0 + + +def test_adamw8bit_never_silently_becomes_offload_adamw_and_typos_still_raise(): + import types + + from align.train.offload_adamw import OffloadAdamW, make_optimizer + + cfg = types.SimpleNamespace( + optimizer=types.SimpleNamespace( + name="adamw8bit", lr=1e-3, betas=[0.9, 0.95], eps=1e-8, weight_decay=0.0 + ) + ) + assert not isinstance(make_optimizer(cfg, _tiny_params()), OffloadAdamW) + cfg.optimizer.name = "adamw8bti" + with pytest.raises(ValueError, match="unknown optimizer"): + make_optimizer(cfg, _tiny_params()) + cfg.optimizer.name = "adafactor" + with pytest.raises(ValueError, match="BANNED"): + make_optimizer(cfg, _tiny_params()) + + +def test_shard_local_shim_constructs_delegates_and_round_trips(): + """The legacy class name now builds the real optimizer. Its wrapper + state_dict (world_size / n_shards / recursive CPU clone) must round-trip + the 8-bit state bit-exactly, and its step must surface the inner skip/norm.""" + from align.train.offload_adamw import ShardLocalAdamW8bit + + a = [_param(B * 2, seed=1), _param(40, seed=2)] + opt = ShardLocalAdamW8bit(a, lr=1e-3, weight_decay=0.1) + _train(opt, a, 3) + assert opt.last_grad_norm > 0 and opt.skipped_steps == 0 + a[0].grad = torch.full_like(a[0], float("nan")) + assert opt.step() is False and opt.skipped_steps == 1 and math.isnan(opt.last_grad_norm) + opt.zero_grad() + sd = opt.state_dict() + assert sd["n_shards"] == 2 and sd["skipped_steps"] == 1 + assert sd["inner"]["state"][0]["m_codes"].device.type == "cpu" + + b = [torch.nn.Parameter(p.detach().clone()) for p in a] + other = ShardLocalAdamW8bit(b, lr=1e-3, weight_decay=0.1) + other.load_state_dict(sd) + _train(opt, a, 2, start=3) + _train(other, b, 2, start=3) + for pa, pb in zip(a, b): + assert torch.equal(pa, pb) + with pytest.raises(ValueError, match="world_size"): + other.load_state_dict(dict(sd, world_size=sd["world_size"] + 1)) + with pytest.raises(ValueError, match="empty"): + ShardLocalAdamW8bit([], lr=1e-3) + + +# --------------------------------------------------------------------------- +# 2-rank gloo + FSDP2 DTensors (CPU only) +# --------------------------------------------------------------------------- + +WS = 2 +JOIN_TIMEOUT_S = 180 + + +def _free_port() -> int: + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +def _body_dtensor(rank: int, run_dir: str) -> dict: + """Child body. Linear(64, 129): 129 rows over 2 ranks = 65 + 64 rows -> + local weights of 4160 (8-bit path) and 4096 (8-bit path, exactly at the + threshold); the bias shards (65/64 elements) take the fp32 fallback. So one + optimizer exercises UNEVEN shards, both state paths, and DTensor grads.""" + import torch.distributed as dist + from torch.distributed.device_mesh import init_device_mesh + from torch.distributed.fsdp import fully_shard + + torch.manual_seed(0) + ref = torch.nn.Linear(64, 129) + sharded = copy.deepcopy(ref) + fully_shard(sharded, mesh=init_device_mesh("cpu", (WS,))) + g = torch.Generator().manual_seed(3) + x, y = torch.randn(256, 64, generator=g), torch.randn(256, 129, generator=g) + + hp = dict(lr=1e-3, betas=(0.9, 0.95), eps=1e-8, weight_decay=0.0) + ref_opt = torch.optim.AdamW(ref.parameters(), **hp) + opt = BlockwiseAdamW8bit(sharded.parameters(), **hp) + # The same algorithm over PLAIN clones of this rank's shards, fed the same + # local grads: DTensor handling may not change a single bit. + plain = [torch.nn.Parameter(p.to_local().detach().clone()) for p in sharded.parameters()] + plain_opt = BlockwiseAdamW8bit(plain, **hp) + + def full(m): + return {n: p.full_tensor().detach().clone() for n, p in m.named_parameters()} + + def loss(m): + return ((m(x) - y) ** 2).mean() + + norms, ref_norms, mean_diff, plain_equal = [], [], 0.0, True + for _ in range(6): + loss(ref).backward() + ref_norms.append(float(torch.sqrt(sum(p.grad.double().pow(2).sum() for p in ref.parameters())))) + ref_opt.step() + ref_opt.zero_grad() + + loss(sharded).backward() + for pl, p in zip(plain, sharded.parameters()): + pl.grad = p.grad.to_local().detach().clone() + assert opt.step() is True + assert plain_opt.step() is True + opt.zero_grad() + norms.append(opt.last_grad_norm) + plain_equal &= all( + torch.equal(pl.detach(), p.to_local()) for pl, p in zip(plain, sharded.parameters()) + ) + got = full(sharded) + mean_diff = max( + mean_diff, + *[float((p.detach() - q).abs().mean()) for (_, p), q in zip(ref.named_parameters(), got.values())], + ) + + # a non-finite grad on ONE rank must make EVERY rank skip + pre = full(sharded) + sharded(x).pow(2).mean().backward() + if rank == 1: + next(iter(sharded.parameters())).grad.to_local().fill_(float("nan")) + stepped = opt.step() + opt.zero_grad() + post = full(sharded) + skip_diff = max(float((pre[k] - post[k]).abs().max()) for k in pre) + + # DTensor snapshot / mutate / restore: weights are the model's, so only the + # optimizer state is restored; the step must continue identically. + sd = opt.state_dict() + gathered = [None] * WS + dist.all_gather_object(gathered, norms) + return { + "norms": norms, + "ref_norms": ref_norms, + "mean_param_diff": mean_diff, + "plain_equal": plain_equal, + "all_rank_norms": gathered, + "local_numel": [int(p.to_local().numel()) for p in sharded.parameters()], + "quantized": [bool(s["quantized"]) for s in sd["state"]], + "stepped_on_nan": stepped, + "skip_diff": skip_diff, + "skipped": opt.skipped_steps, + "sd_world_size": sd["world_size"], + } + + +_BODIES = {"dtensor": _body_dtensor} + + +def _child(rank: int, port: int, mode: str, run_dir: str) -> None: + import torch.distributed as dist + + os.environ.update( + RANK=str(rank), + WORLD_SIZE=str(WS), + LOCAL_RANK=str(rank), + MASTER_ADDR="127.0.0.1", + MASTER_PORT=str(port), + ) + dist.init_process_group("gloo", rank=rank, world_size=WS) + code = 0 + try: + result = _BODIES[mode](rank, run_dir) + with open(os.path.join(run_dir, f"{mode}_{rank}.json"), "w") as f: + json.dump(result, f) + except BaseException: + with open(os.path.join(run_dir, f"{mode}_{rank}.err"), "w") as f: + f.write(traceback.format_exc()) + code = 1 + finally: + try: + dist.destroy_process_group() + except Exception: + pass + if code: + sys.exit(code) + + +def _run(mode: str, run_dir: str, monkeypatch) -> list[dict]: + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "") # no GPU is ever touched + ctx = mp.get_context("spawn") + port = _free_port() + procs = [ctx.Process(target=_child, args=(r, port, mode, run_dir)) for r in range(WS)] + for p in procs: + p.start() + for p in procs: + p.join(JOIN_TIMEOUT_S) + hung = [p for p in procs if p.is_alive()] + for p in hung: + p.kill() + p.join() + errs = [ + f"rank {r}:\n{open(os.path.join(run_dir, f'{mode}_{r}.err')).read()}" + for r in range(WS) + if os.path.isfile(os.path.join(run_dir, f"{mode}_{r}.err")) + ] + assert not hung, f"{len(hung)} rank(s) HUNG in mode {mode!r}\n" + "\n".join(errs) + assert not errs, "\n".join(errs) + return [json.load(open(os.path.join(run_dir, f"{mode}_{r}.json"))) for r in range(WS)] + + +def test_adamw8bit_steps_fsdp2_dtensors_with_a_global_norm_and_a_global_skip(tmp_workdir, monkeypatch): + """THE FSDP2 DTENSOR DEFECT + THE SHARD-LOCAL NORM DEFECT, for the 8-bit + optimizer (bnb's fused kernel died on DTensors: 'mixed torch.Tensor and + DTensor'; OffloadAdamW died the same way until `_localize`). + + (1) Stepping the LOCAL shards of an uneven FSDP2 weight is BIT-IDENTICAL to + the same optimizer over plain clones of those shards (DTensor handling + changes nothing), and tracks a single-process torch AdamW on the unsharded + model on average (mean abs diff < 2e-4 (measured 6e-5) against ~6e-3 of travel; the max is + not bounded tightly, see the cold-start note in adamw8bit.py). + (2) `last_grad_norm` is the GLOBAL norm and the SAME number on every rank: + a rank-local slice norm lets one rank cross a halt tripwire while its peer + waits in a collective. (3) A NaN on one rank's shard makes both ranks skip. + (4) State really is shard-shaped and quantized on at least one rank.""" + rows = _run("dtensor", tmp_workdir, monkeypatch) + for r in rows: + assert r["plain_equal"], "stepping DTensor shards != stepping plain clones of them" + assert r["mean_param_diff"] < 2e-4, r["mean_param_diff"] + # step 0 sees identical weights, so the GLOBAL norm must match to fp32 + # rounding; later steps' weights differ by the 8-bit error (~1e-4 rel). + assert r["norms"][0] == pytest.approx(r["ref_norms"][0], rel=1e-5) + assert r["norms"] == pytest.approx(r["ref_norms"], rel=2e-3), (r["norms"], r["ref_norms"]) + assert r["all_rank_norms"][0] == r["all_rank_norms"][1], r["all_rank_norms"] + assert r["stepped_on_nan"] is False and r["skipped"] == 1, r + assert r["skip_diff"] == 0.0, "a rank stepped although its peer's gradient was non-finite" + assert r["sd_world_size"] == WS + assert r["quantized"][0] is True, "the weight shard must take the 8-bit path" + assert rows[0]["local_numel"] != rows[1]["local_numel"], "129 rows must shard unevenly" + + +# --------------------------------------------------------------------------- +# the launch path: shipped configs -> make_optimizer -> a real optimizer +# --------------------------------------------------------------------------- + + +def _shipped_configs_naming(name): + import glob + import os + + from align.config import load_config + from tests.conftest import REPO_ROOT + + out = [] + for path in sorted(glob.glob(os.path.join(REPO_ROOT, "configs", "*.toml"))): + try: + cfg = load_config(path) + except Exception: + continue # stage-specific TOMLs (ppo_stage) are not RunConfigs + if getattr(getattr(cfg, "optimizer", None), "name", None) == name: + out.append((os.path.basename(path), cfg)) + return out + + +def test_every_shipped_adamw8bit_config_constructs_the_real_optimizer_and_steps(): + """THE LAUNCH BLOCKER (review_trainer.md): every shipped full-FT config names adamw8bit and + make_optimizer used to raise NotImplementedError for it, so no stage could construct an + optimizer. Build from the REAL shipped config on a tiny stub model, and take a step. Mutant: + restore the raise (or return OffloadAdamW) -> the isinstance check fails.""" + from align.train.adamw8bit import BlockwiseAdamW8bit + from align.train.offload_adamw import make_optimizer + + found = _shipped_configs_naming("adamw8bit") + assert found, "no shipped config names adamw8bit; this test would be vacuous" + assert "sft.toml" in [n for n, _ in found] + for fname, cfg in found: + params = [torch.nn.Parameter(torch.randn(n).bfloat16()) for n in (4096 * 3, 100)] + opt = make_optimizer(cfg, params) + assert isinstance(opt, BlockwiseAdamW8bit), (fname, type(opt)) + assert "adamw8bit" in opt.banner_tag and "zero host RSS" in opt.banner_tag + before = [p.detach().clone() for p in params] + for p in params: + p.grad = torch.randn_like(p) * 1e-2 + assert opt.step() is True + assert any(not torch.equal(a, b.detach()) for a, b in zip(before, params)), fname + + +def test_optimizer_config_max_grad_norm_reaches_the_optimizer_and_is_recorded_pre_clip(): + """`OptimizerConfig.max_grad_norm` (1.0 by default since G0; None disables clipping explicitly) + must reach the optimizer, which clips by the GLOBAL norm and records the PRE-clip + norm. Mutant: make_optimizer drops the field -> the param-group assertion fails (the clip + arithmetic itself is pinned bitwise by the grad-norm test above).""" + from align.config import OptimizerConfig, RunConfig + from align.train.offload_adamw import make_optimizer + + assert OptimizerConfig().max_grad_norm == 1.0 # clip ON by default (G0) + + def build(clip): + cfg = RunConfig() + cfg.optimizer.name = "adamw8bit" + cfg.optimizer.max_grad_norm = clip + torch.manual_seed(0) + p = torch.nn.Parameter(torch.randn(8192).bfloat16()) + return p, make_optimizer(cfg, [p]) + + p_clip, o_clip = build(1.0) + p_free, o_free = build(None) + g = torch.randn(8192) * 100.0 # global norm ~9000 + for p in (p_clip, p_free): + p.grad = g.bfloat16().clone() + o_clip.step() + o_free.step() + assert o_clip.last_grad_norm == pytest.approx(o_free.last_grad_norm, rel=1e-6) + assert o_clip.last_grad_norm > 1000.0, "must record the PRE-clip norm, not <= max_grad_norm" + # Adam is scale-invariant on step 1, so the update cannot show the clip; the group setting can. + assert o_clip.param_groups[0]["max_grad_norm"] == 1.0 + assert o_free.param_groups[0]["max_grad_norm"] is None + + +def test_train_config_per_rank_optimizer_state_is_settable_from_struct_mode_toml(tmp_path): + """The Trainer read this flag with getattr because the field did not exist, so struct-mode TOML + (unknown keys raise) could not turn it on. It must parse, default OFF, and round-trip.""" + from align.config import TrainConfig, load_config + + assert TrainConfig().per_rank_optimizer_state is False + f = tmp_path / "c.toml" + f.write_text("[train]\nper_rank_optimizer_state = true\n") + assert load_config(str(f)).train.per_rank_optimizer_state is True + + +def test_train_module_default_extra_state_is_an_empty_noop(): + """Subclasses (GRPO tripwire) override the pair; the base must exist so the Trainer's resume.pt + always carries the key and `load_extra_state` never needs a getattr guard to be safe.""" + from align.train.train_module import TrainModule + + m = TrainModule.__new__(TrainModule) + assert m.extra_state() == {} + assert m.load_extra_state({}) is None + + +# --------------------------------------------------------------------------- +# adversarial review (2026-10-12): defects found by trying to break the optimizer +# --------------------------------------------------------------------------- + + +def test_second_moment_overflow_cannot_poison_the_codes_with_nan(): + """THE V-OVERFLOW POISON. The finiteness gate tests the float64 NORM, which stays finite for + a gradient whose fp32 SQUARE is inf: one element of 1e21 gives v = (1-beta2) * g^2 = inf. + torch.optim.AdamW survives (m / (inf + eps) = 0); here inf became the block's absmax, the + requantize divided inf/inf = NaN into the uint8 codes, and the block's weights were NaN one + step LATER with the gate still seeing only finite gradients. + + Kills: removing the `v_full.clamp_(max=fp32_max)` before requantization.""" + p = torch.nn.Parameter(_randn(B * 4, seed=1)) + ref = torch.nn.Parameter(p.detach().clone()) + opt = BlockwiseAdamW8bit([p], lr=1e-3, betas=(0.9, 0.95)) + ropt = torch.optim.AdamW([ref], lr=1e-3, betas=(0.9, 0.95), weight_decay=0.0) + for step in range(3): + g = _randn(B * 4, seed=10 + step) + if step == 0: + g[5] = 1e21 + p.grad, ref.grad = g.clone(), g.clone() + assert opt.step() is True + ropt.step() + st = opt._state[0] + assert torch.isfinite(st["v_absmax"]).all() and torch.isfinite(st["m_absmax"]).all(), step + assert torch.isfinite(p).all(), f"weights went non-finite at step {step}" + # The overflowing element itself: torch freezes it (v stays inf), the clamp lets it move by a + # bounded, finite amount (m ~ 1e20 over sqrt(fp32 max) ~ 1.8e19 is a few lr per step). + assert abs(p[5].item() - ref[5].item()) < 50 * 1e-3 + + +def test_load_state_dict_rewinds_the_rounding_rng_of_an_optimizer_that_already_stepped(): + """`load_state_dict` promises a byte-for-byte restore. The stochastic-rounding generator is + lazily created, so a state saved BEFORE the first bf16 step carries no `sr_rng`; loading it + into an optimizer that has since stepped left that optimizer's generator advanced, and the + 'rewound' run rounded differently from a never-stepped one. + + Kills: not clearing `_generators` on load.""" + torch.manual_seed(0) + base = (torch.randn(B * 3) * 0.05).to(torch.bfloat16) + + def mk(): + return torch.nn.Parameter(base.clone()) + + fresh_p = mk() + sd0 = BlockwiseAdamW8bit([fresh_p], lr=1e-3).state_dict() + assert sd0["sr_rng"] == {}, "precondition: nothing stepped, so no generator to save" + + used_p, ref_p = mk(), mk() + used = BlockwiseAdamW8bit([used_p], lr=1e-3) + ref = BlockwiseAdamW8bit([ref_p], lr=1e-3) + for s in range(3): + _set_grads([used_p], s) + used.step() + used_p.data.copy_(base) + used.load_state_dict(sd0) + for s in range(4): + _set_grads([used_p], s) + _set_grads([ref_p], s) + used.step() + ref.step() + assert torch.equal(used_p, ref_p) + + +def test_persistent_tiny_gradient_overstep_stays_inside_the_documented_bound(): + """The second-moment underflow is PERMANENT for an element whose gradient stays below + ~5e-4 of its block's largest (v is re-floored to 0 every step), not the transient the module + docstring used to claim: measured against torch.optim.AdamW with constant g = 1e-4 of the + block max, 300 steps, the 8-bit travel is 3.9x at beta2 = 0.95 and 9.9x at 0.999. What must + hold regardless of how the owner resolves that is the per-step ceiling lr / sqrt(1 - beta2): + the update comes from the FRESH moment, so a long run of underflowed steps can never turn + into the m / (0 + eps) = 100 lr step. + + Kills: computing the update from the requantized (stored, zero) v, at ANY step; a floor + that is applied to the fresh moment.""" + lr, beta2 = 1e-3, 0.95 + p = torch.nn.Parameter(torch.zeros(B * 2)) + opt = BlockwiseAdamW8bit([p], lr=lr, betas=(0.9, beta2), weight_decay=0.0) + g = torch.full((B * 2,), 1e-4) + g[0] = 1.0 + worst = 0.0 + for _ in range(200): + before = p.detach().clone() + p.grad = g.clone() + opt.step() + worst = max(worst, (p.detach() - before)[1:].abs().max().item()) + assert (opt._state[0]["v_codes"][0, 1:] == 0).all(), "precondition: v really is underflowed" + assert worst <= lr / math.sqrt(1 - beta2) * 1.05, worst / lr + + +def _body_dtensor_clip(rank: int, run_dir: str) -> dict: + """Child body: a 1-row layer (rank 1 holds an EMPTY shard), an uneven 129-row layer, in-optimizer + global clipping and weight decay, against torch clip_grad_norm_ + AdamW on the unsharded model.""" + from torch.distributed.device_mesh import init_device_mesh + from torch.distributed.fsdp import fully_shard + + class Net(torch.nn.Module): + def __init__(self): + super().__init__() + self.a = torch.nn.Linear(64, 1) + self.b = torch.nn.Linear(64, 129) + + def forward(self, x): + return torch.cat([self.a(x), self.b(x)], 1) + + torch.manual_seed(0) + ref = Net() + start = copy.deepcopy(ref) + sharded = copy.deepcopy(ref) + mesh = init_device_mesh("cpu", (WS,)) + fully_shard(sharded.a, mesh=mesh) + fully_shard(sharded.b, mesh=mesh) + fully_shard(sharded, mesh=mesh) + g = torch.Generator().manual_seed(3) + x, y = torch.randn(128, 64, generator=g), torch.randn(128, 130, generator=g) * 10 + + hp = dict(lr=1e-3, betas=(0.9, 0.95), eps=1e-8, weight_decay=0.1) + ref_opt = torch.optim.AdamW(ref.parameters(), **hp) + opt = BlockwiseAdamW8bit(sharded.parameters(), max_grad_norm=0.5, **hp) + rel_norm_err, clipped = [], 0 + for _ in range(8): + ((ref(x) - y) ** 2).mean().backward() + total = float(torch.nn.utils.clip_grad_norm_(ref.parameters(), 0.5)) + clipped += total > 0.5 + ref_opt.step() + ref_opt.zero_grad() + ((sharded(x) - y) ** 2).mean().backward() + assert opt.step() is True + opt.zero_grad() + rel_norm_err.append(abs(opt.last_grad_norm - total) / total) + with torch.no_grad(): + diff = max( + float((p.full_tensor() - r).abs().mean()) for p, r in zip(sharded.parameters(), ref.parameters()) + ) + travel = max(float((r - r0).abs().mean()) for r, r0 in zip(ref.parameters(), start.parameters())) + sd = opt.state_dict() + opt.load_state_dict(sd) # an empty-shard rank must round-trip too + return { + "local_numel": [int(p.to_local().numel()) for p in sharded.parameters()], + "rel_norm_err": rel_norm_err, + "clipped_steps": int(clipped), + "mean_diff": diff, + "travel": travel, + } + + +_BODIES["dtensor_clip"] = _body_dtensor_clip + + +def test_adamw8bit_global_clip_and_empty_shard_match_torch_on_two_ranks(tmp_workdir, monkeypatch): + """The in-optimizer clip must use the GLOBAL norm: a shard-local norm is ~1/sqrt(world) of the + true one, so the clip would fire late (or never) and the two ranks would scale differently. A + rank whose shard of a 1-row layer is EMPTY (0 local elements) must still join every + collective, skip its own update, and round-trip its state. + + Kills: `coef` computed from a rank-local norm; skipping `_grad_norm`'s all-reduce; treating a + zero-element shard as 'no gradient' (the 'no gradients on any parameter' raise would hang the + peer in its all-reduce).""" + rows = _run("dtensor_clip", tmp_workdir, monkeypatch) + assert rows[1]["local_numel"][:2] == [0, 0], "rank 1 must hold the empty shard of the 1-row layer" + for r in rows: + assert r["clipped_steps"] >= 6, "the clip must actually be active in this test" + assert max(r["rel_norm_err"]) < 1e-3, r["rel_norm_err"] + assert r["mean_diff"] < 0.02 * r["travel"], (r["mean_diff"], r["travel"]) diff --git a/code/tests/test_argilla_arm_a.py b/code/tests/test_argilla_arm_a.py new file mode 100644 index 0000000000000000000000000000000000000000..86240fa83da9991e1ce2a27cdd2ac7e66abff935 --- /dev/null +++ b/code/tests/test_argilla_arm_a.py @@ -0,0 +1,124 @@ +"""Arm A ingestion (argilla distilabel math preference pairs). +Run: uv run pytest tests/test_argilla_arm_a.py -m "not gpu" -v + +Arm A is the off-policy, human/AI-rated preference arm. Its whole scientific +value depends on being EXACTLY the set that was measured at plan time, so the +counts are asserted, not logged. +""" + +from __future__ import annotations + +import json +import os + +import pytest + +from align.data import argilla, sources +from align.data.schema import PROMPT_TEMPLATE, stable_problem_id + + +def _rows(): + try: + return argilla.load_argilla_rows() + except (FileNotFoundError, RuntimeError) as e: + pytest.skip(f"argilla snapshot unavailable: {e}") + + +def test_margin_filter_reproduces_measured_yield(): + """margin >=1.0 AND chosen >=6 keeps exactly 2,310 of 2,418 on the pinned + dataset revision; a changed upstream revision fails the count assertion + rather than silently training on different data (the drift incident + class). Both directions: the pinned revision passes, a doctored row count + raises.""" + rows = _rows() + assert len(rows) == argilla.RAW_ROWS == 2418 + + kept = argilla.apply_quality_filter(rows) + assert len(kept) == argilla.FILTERED_ROWS == 2310 + + # Firing direction: a drifted upstream (one row short) must abort. + with pytest.raises(RuntimeError, match="revision"): + argilla.assert_pinned_counts(rows[:-1], kept) + + +def test_contamination_zero_and_planted_problem_caught(): + """Contamination vs the eval hash set is 0/2,418 (silent direction); a + deliberately planted eval problem IS caught (firing direction).""" + rows = _rows() + eval_hashes = {"deadbeefdeadbeef"} + clean, n_contam = argilla.decontaminate(rows, eval_hashes) + assert n_contam == 0 + assert len(clean) == len(rows) + + planted = stable_problem_id(rows[0]["instruction"]) + clean2, n2 = argilla.decontaminate(rows, eval_hashes | {planted}) + assert n2 >= 1 + assert len(clean2) == len(rows) - n2 + assert all(stable_problem_id(r["instruction"]) != planted for r in clean2) + + +def test_prompt_mapping_verbatim_no_answer_injection(): + """instruction maps into the prompt template's problem slot verbatim; + completions are used verbatim -- no '# Answer' injection, no reformatting + (format drift is a phenomenon to MEASURE at eval time, not to mask at + build time).""" + rows = _rows() + pairs = argilla.to_pairs(argilla.apply_quality_filter(rows)[:32]) + for pair, row in zip(pairs, argilla.apply_quality_filter(rows)[:32]): + assert pair.prompt == PROMPT_TEMPLATE.format(problem=row["instruction"].strip()) + assert row["instruction"].strip() in pair.prompt + assert pair.chosen == row["chosen_response"] + assert pair.rejected == row["rejected_response"] + assert pair.problem_id == stable_problem_id(row["instruction"]) + + # No injected answer marker anywhere: whatever the completions do or do + # not contain is the measurement. + injected = sum( + 1 + for p in pairs + if "# Answer" in p.chosen + and "# Answer" + not in next(r["chosen_response"] for r in rows if r["instruction"] == p.instruction) + ) + assert injected == 0 + + +def test_length_bias_baseline_recorded(tmp_workdir): + """Build receipt records chosen-longer fraction; equals 54.9% on the + pinned revision (this number seeds the section-07 length-drift tripwire). + A DPO run that drifts far above this is learning length, not correctness.""" + rows = _rows() + kept = argilla.apply_quality_filter(rows) + stats = argilla.length_bias_stats(kept) + + assert stats["chosen_longer_fraction"] == pytest.approx(0.549, abs=5e-4) + assert stats["mean_chosen_chars"] == pytest.approx(1278.0, abs=1.0) + assert stats["mean_rejected_chars"] == pytest.approx(1232.0, abs=1.0) + + receipt = argilla.write_arm_a_receipt(tmp_workdir, rows, kept, n_contaminated=0) + on_disk = json.loads(open(os.path.join(tmp_workdir, "arm_a.receipt.json")).read()) + assert on_disk["chosen_longer_fraction"] == pytest.approx(0.549, abs=5e-4) + assert on_disk["raw_rows"] == 2418 + assert on_disk["kept_rows"] == 2310 + assert on_disk["contaminated"] == 0 + assert on_disk["dataset_revision"] == sources.ARGILLA_PIN.revision + assert receipt["descoped"] is False + + +def test_descope_marker_written_when_filter_starves_the_arm(tmp_workdir): + """Pre-registered descope: a filter yielding < 1,500 pairs writes a + DESCOPED marker instead of shards, so section-12 reports Arm A as + descoped rather than reporting a null from an underpowered arm. + Both directions: the real filter (2,310) writes no marker.""" + rows = _rows() + + starved = argilla.apply_quality_filter(rows, min_margin=1.0, min_chosen=9.5) + assert len(starved) < argilla.DESCOPE_MIN_PAIRS + marker = argilla.check_descope(tmp_workdir, starved) + assert marker is not None + assert os.path.exists(os.path.join(tmp_workdir, "arm_a.DESCOPED")) + + ok_dir = os.path.join(tmp_workdir, "ok") + os.makedirs(ok_dir, exist_ok=True) + assert argilla.check_descope(ok_dir, argilla.apply_quality_filter(rows)) is None + assert not os.path.exists(os.path.join(ok_dir, "arm_a.DESCOPED")) diff --git a/code/tests/test_assembly.py b/code/tests/test_assembly.py new file mode 100644 index 0000000000000000000000000000000000000000..bc881ec7fedd56312b08f89f2fec9871d6d28527 --- /dev/null +++ b/code/tests/test_assembly.py @@ -0,0 +1,143 @@ +"""Tests for the rollout assembly seam (fake verifier; no model, no engine).""" + +import enum + +import pytest +import torch + +from align.rl.advantages import GroupFilterStats +from align.rl.assembly import assemble_rollout +from align.rl.rewards import CompositeReward + + +class _Outcome(enum.Enum): + CORRECT = 1 + WRONG = 2 + + +def _verify(golds, preds): + return [ + _Outcome.CORRECT if (p is not None and p.strip() == g) else _Outcome.WRONG + for g, p in zip(golds, preds) + ] + + +def _text(answer: str) -> str: + return f"Step 1: reason.\n\n# Answer\n\n{answer}" + + +def _mk(n_prompts=2, g=2, answers=None): + """Two prompts x G completions with controllable answers.""" + prompt_ids = [[1, 2, 3], [4, 5]] + completions = [[[10, 11], [12, 13, 14]], [[20], [21, 22]]] + logprobs = [[[-0.1, -0.2], [-0.3, -0.4, -0.5]], [[-0.6], [-0.7, -0.8]]] + answers = answers or [["7", "8"], ["9", "9"]] + texts = [[_text(a) for a in row] for a in [None] for row in answers] + return prompt_ids, completions, logprobs, texts + + +def test_assembles_shapes_masks_and_alignment(): + prompt_ids, completions, logprobs, texts = _mk(answers=[["7", "8"], ["9", "1"]]) + out = assemble_rollout( + prompt_ids, + completions, + logprobs, + texts, + golds=["7", "9"], + reward=CompositeReward(format_weight=0.1), + verify_batch=_verify, + pad_token_id=0, + ) + assert out is not None + B, T = out.input_ids.shape + assert B == 4 and T == 6 # longest row: prompt 3 + completion 3 + # row 0: prompt [1,2,3] + completion [10,11] + assert out.input_ids[0].tolist() == [1, 2, 3, 10, 11, 0] + assert out.response_mask[0].tolist() == [0, 0, 0, 1, 1, 0] + # label-aligned engine logprobs: completion token t -> position len(p)+t-1 + assert out.vllm_logprobs[0, 2].item() == pytest.approx(-0.1) + assert out.vllm_logprobs[0, 3].item() == pytest.approx(-0.2) + assert out.vllm_logprobs[0, 1].item() == 0.0 # prompt region zero-filled + # group ids: two kept groups + assert out.group_ids.tolist() == [0, 0, 1, 1] + # advantages: group means subtracted -> each group sums to ~0 + a = out.advantages + assert abs(float(a[0] + a[1])) < 1e-6 and abs(float(a[2] + a[3])) < 1e-6 + # rewards itemised + assert set(out.rewards) >= {"reward/accuracy", "reward/format", "reward/total"} + assert out.rewards["reward/accuracy"].tolist() == [1.0, 0.0, 1.0, 0.0] + + +def test_dead_groups_dropped_whole_with_telemetry(): + # group 0: both wrong (all_zero-ish variance 0 at total level? totals equal) + # -> dead; group 1: mixed -> kept + prompt_ids, completions, logprobs, texts = _mk(answers=[["1", "1"], ["9", "2"]]) + stats = GroupFilterStats(max_possible_score=1.1) # acc 1.0 + fmt 0.1 + out = assemble_rollout( + prompt_ids, + completions, + logprobs, + texts, + golds=["7", "9"], + reward=CompositeReward(format_weight=0.1), + verify_batch=_verify, + pad_token_id=0, + filter_stats=stats, + ) + assert out is not None + assert out.kept_groups == 1 and out.dropped_groups == 1 + assert out.input_ids.shape[0] == 2 # only group 1's two rows + # both-wrong completions still earn the format term (0.1 each), so the dead + # group's identical totals are NONZERO -> nonzero_std0, not all_zero. With a + # welded soft-format reward, all_zero is reserved for unformatted garbage. + assert stats.nonzero_std0 == 1 and stats.kept == 1 + + +def test_all_dead_returns_none_not_error(): + prompt_ids, completions, logprobs, texts = _mk(answers=[["1", "1"], ["2", "2"]]) + out = assemble_rollout( + prompt_ids, + completions, + logprobs, + texts, + golds=["7", "9"], + reward=CompositeReward(format_weight=0.1), + verify_batch=_verify, + pad_token_id=0, + ) + assert out is None + + +def test_ragged_logprobs_refused(): + prompt_ids, completions, logprobs, texts = _mk() + logprobs[0][0] = [-0.1] # 1 logprob for 2 completion tokens + with pytest.raises(ValueError, match="ragged|logprobs"): + assemble_rollout( + prompt_ids, + completions, + logprobs, + texts, + golds=["7", "9"], + reward=CompositeReward(format_weight=0.1), + verify_batch=_verify, + pad_token_id=0, + ) + + +def test_unanswerable_text_scores_zero_accuracy(): + prompt_ids, completions, logprobs, _ = _mk() + texts = [["no answer marker here", _text("8")], [_text("9"), _text("9")]] + out = assemble_rollout( + prompt_ids, + completions, + logprobs, + texts, + golds=["8", "9"], + reward=CompositeReward(format_weight=0.1), + verify_batch=_verify, + pad_token_id=0, + ) + assert out is not None + # group 0 kept (one right, one no-answer); the no-answer row scores acc 0 + row0_acc = out.rewards["reward/accuracy"][0].item() + assert row0_acc == 0.0 diff --git a/code/tests/test_attention_backends.py b/code/tests/test_attention_backends.py new file mode 100644 index 0000000000000000000000000000000000000000..3696500b393af79af50c4b798279c9ac162b6041 --- /dev/null +++ b/code/tests/test_attention_backends.py @@ -0,0 +1,265 @@ +"""Backend probe/select/benchmark guards, both directions. + +Run: uv run pytest tests/test_attention_backends.py -v -m "not gpu" + +Every test here runs the probe on device="cpu" on purpose: the CPU suite must +not initialise a CUDA context on a shared box, and the selection LOGIC is +device-independent. The device-specific numbers are the gpu-marked benchmark. +""" + +from __future__ import annotations + +import math +import os +import sys + +import pytest + +from align.model import attention as attn +from align.model.attention import ( + TRITON_FA2_GATE_MODULE, + BackendReport, + benchmark_backends, + probe_backends, + select_attn_implementation, +) +from align.model.gates import load_sdpa_baseline + + +@pytest.fixture +def no_triton_kernel(): + """Make section-06's gate module unimportable for the duration of a test. + + The 'kernel absent' branch is a permanent code path -- it is what any fresh + checkout, and any box where the Triton build fails, actually hits -- so it + is tested by BLOCKING the import rather than by relying on section-06 not + having landed yet. A test that silently stops testing anything the day a + sibling lane ships is worse than no test. + """ + + class _Block: + def find_spec(self, name, path=None, target=None): + if name == TRITON_FA2_GATE_MODULE: + raise ImportError(f"blocked by test fixture: {name}") + return None + + blocker = _Block() + saved = sys.modules.pop(TRITON_FA2_GATE_MODULE, None) + sys.meta_path.insert(0, blocker) + try: + yield + finally: + sys.meta_path.remove(blocker) + if saved is not None: + sys.modules[TRITON_FA2_GATE_MODULE] = saved + + +@pytest.fixture +def gate_dir(tmp_workdir): + """A scratch equivalence-gate directory, so these tests never read -- or + write -- the real artifact under .cache/kernel_gates.""" + return os.path.join(tmp_workdir, "kernel_gates") + + +def _write_passing_gate(gate_dir): + """Write a REAL gate artifact for the current kernel source and + environment, using section-06's own writer. Stubbing a boolean here would + test my mock; this tests the contract the two sections actually share.""" + from align.kernels import gate as gate_mod + + gate_mod.write_equivalence_gate( + tolerances={"out": 1e-3}, + shapes_covered=["backend-selection test"], + gate_dir=gate_dir, + ) + + +def _reports(device="cpu", gate_dir=None): + return {r.name: r for r in probe_backends(device=device, gate_dir=gate_dir)} + + +def test_probe_reports_every_backend_with_a_reason(gate_dir): + """A probe that omits a backend is indistinguishable from one that found it + unavailable. Every declared backend gets a row, and every row carries a + non-empty reason -- the reason string is what ends up in the launch banner + when a run cannot use the accelerator someone expected.""" + reports = _reports(gate_dir=gate_dir) + assert set(reports) == set(attn.ALL_BACKENDS) + for name, r in reports.items(): + assert isinstance(r, BackendReport) + assert r.reason.strip(), name + + +def test_select_triton_fa2_raises_before_equivalence_gate(gate_dir): + """select_attn_implementation("triton_fa2") RAISES while the kernel's + SDPA-equivalence gate (section-06) has not passed. It must never silently + fall back -- a run must not claim an accelerator it never used. (Both + directions: with a valid gate artifact it returns the backend.)""" + # direction 1: no gate artifact -> refuse + with pytest.raises(RuntimeError, match="equivalence gate"): + select_attn_implementation("triton_fa2", device="cpu", gate_dir=gate_dir) + + # direction 2: a gate artifact for this kernel source and environment + _write_passing_gate(gate_dir) + assert select_attn_implementation("triton_fa2", device="cpu", gate_dir=gate_dir) == "triton_fa2" + + +def test_select_triton_fa2_raises_when_the_kernel_is_absent(no_triton_kernel): + """The fresh-checkout path: section-06's module is not importable at all. + That is a different branch from 'gate not passed' and it must also refuse + rather than fall back.""" + with pytest.raises(RuntimeError, match="not yet implemented"): + select_attn_implementation("triton_fa2", device="cpu") + + +def test_select_triton_fa2_refuses_a_gate_that_no_longer_matches_the_kernel(gate_dir): + """THE STALE-GATE INCIDENT, at the selector: an artifact written before a + kernel edit would bless code that was never run against SDPA. The gate + pins the kernel source hash; the selector must surface that refusal rather + than treat 'an artifact exists' as availability.""" + from align.kernels import gate as gate_mod + + _write_passing_gate(gate_dir) + assert select_attn_implementation("triton_fa2", device="cpu", gate_dir=gate_dir) == "triton_fa2" + + original = gate_mod.kernel_source_hash + gate_mod.kernel_source_hash = lambda: "0" * 64 + try: + with pytest.raises(RuntimeError, match="kernel source changed"): + select_attn_implementation("triton_fa2", device="cpu", gate_dir=gate_dir) + finally: + gate_mod.kernel_source_hash = original + + +def test_triton_fa2_is_never_auto_selected(gate_dir): + """Even with the gate passed, auto-selection stays on sdpa. A custom + kernel is opted INTO explicitly so the banner records a deliberate choice; + a kernel that switches itself on turns every later benchmark into a + comparison against an unknown baseline.""" + _write_passing_gate(gate_dir) + assert select_attn_implementation(None, device="cpu", gate_dir=gate_dir) == "sdpa" + + +def test_select_unavailable_backend_raises_with_reason(gate_dir): + """Explicitly requesting any unavailable backend raises RuntimeError + carrying the probe's reason string; auto-selection (requested=None) picks + the best available and logs which.""" + reports = _reports(gate_dir=gate_dir) + unavailable = [n for n, r in reports.items() if not r.available] + assert unavailable, "probe found nothing unavailable -- it is not probing" + for name in unavailable: + with pytest.raises(RuntimeError) as exc: + select_attn_implementation(name, device="cpu", gate_dir=gate_dir) + # the probe's own reason must survive into the error the user sees + assert reports[name].reason[:40] in str(exc.value), name + + assert select_attn_implementation(None, device="cpu", gate_dir=gate_dir) == "sdpa" + + +def test_flashmla_refused_even_if_it_were_available(): + """FlashMLA is refused for structural reasons, not availability ones: its + kernel requires head_dim_v=512 MLA tensors and Mistral is GQA head_dim + 128, so no GPU -- Hopper included -- can serve this model with it. The + refusal must not be reachable-around by a future SM90 box.""" + with pytest.raises(RuntimeError, match="head_dim"): + select_attn_implementation("flashmla", device="cpu") + + +def test_select_rejects_an_unknown_backend_name(): + """A typo'd backend name must not fall through to the default. The + predecessor's whole failure class is a plausible substitute quietly + standing in for what was asked for.""" + with pytest.raises(RuntimeError, match="unknown backend"): + select_attn_implementation("sdpa_flash", device="cpu") + + +def test_benchmark_records_nan_not_zero_for_unusable_backend(): + """A backend that throws at these shapes records float('nan'), never 0.0 + and never a swallowed pass (the EFFICIENT-on-GQA incident: NaN is a real + finding, recorded not hidden). Undefined stats are NaN, never 0. + + Run on CPU at tiny shapes, where EFFICIENT_ATTENTION and CUDNN_ATTENTION + genuinely have no kernel -- the same shape of unavailability the pinned + A10G baseline recorded for EFFICIENT on GQA inputs. + """ + results = benchmark_backends( + batch=1, + seq=16, + n_heads=4, + n_kv_heads=2, + head_dim=16, + iters=2, + device="cpu", + ) + assert set(results) >= {"FLASH_ATTENTION", "EFFICIENT_ATTENTION", "MATH"} + unusable = [k for k, v in results.items() if math.isnan(v)] + assert unusable, "no backend was unusable at CPU tiny shapes -- expected at least one" + for name, ms in results.items(): + assert ms != 0.0, f"{name} recorded 0.0; an unusable backend is NaN, never 0" + assert math.isnan(ms) or ms > 0.0, f"{name}={ms}" + + +def test_benchmark_does_not_swallow_the_reason(monkeypatch): + """The exception a backend raises is a finding. It is logged and attached + to the result, not discarded -- 'NaN' with no reason cannot be triaged + six weeks later when a torch upgrade changes which kernels exist.""" + results, reasons = attn.benchmark_backends_with_reasons( + batch=1, + seq=16, + n_heads=4, + n_kv_heads=2, + head_dim=16, + iters=2, + device="cpu", + ) + for name, ms in results.items(): + if math.isnan(ms): + assert reasons[name] and reasons[name] != "ok", name + else: + assert reasons[name] == "ok", name + + +def test_pinned_sdpa_baseline_is_the_bar_and_records_nan(): + """The measured A10G baseline is pinned in-repo so section-06's 'beat the + bar' comparison and a torch/driver regression of the SDPA path itself are + both detectable. EFFICIENT is NaN in the pinned artifact for the same + reason it is NaN at runtime -- that is the recorded finding, not a hole.""" + base = load_sdpa_baseline() + assert set(base) == {"512", "1024", "2048"} + for seq, row in base.items(): + assert row["FLASH_ATTENTION"] > 0.0 + assert math.isnan(row["EFFICIENT_ATTENTION"]), ( + "EFFICIENT must be NaN, not 0 and not absent, in the pinned baseline" + ) + # flash must beat math by a wide margin or the pin is of the wrong run + assert row["MATH"] > 10 * row["FLASH_ATTENTION"], seq + + +def test_banner_states_which_backend_actually_ran(): + """Every training/inference job's banner reports the backend that was + resolved, so 'we ran the custom kernel' is a recorded fact rather than an + intention. attention_backend is already a mandatory banner key; this test + pins that the selector's output is what fills it.""" + from align.registry import MANDATORY_BANNER_KEYS + + assert "attention_backend" in MANDATORY_BANNER_KEYS + resolved = select_attn_implementation(None, device="cpu") + assert attn.banner_value(resolved, device="cpu").startswith(resolved) + + +@pytest.mark.gpu +def test_benchmark_reproduces_the_pinned_a10g_baseline(): + """On a real A10G, re-running the benchmark at the pinned shapes must land + within 2x of the recorded numbers and must still record EFFICIENT as NaN. + A drastic regression of the SDPA path (torch upgrade, driver change) is + exactly what this catches.""" + import torch + + if not torch.cuda.is_available(): + pytest.skip("no CUDA device visible") + base = load_sdpa_baseline() + for seq in (512, 1024, 2048): + got = benchmark_backends(batch=4, seq=seq, device="cuda") + ref = base[str(seq)] + assert math.isnan(got["EFFICIENT_ATTENTION"]) + assert got["FLASH_ATTENTION"] < 2.0 * ref["FLASH_ATTENTION"], (seq, got, ref) diff --git a/code/tests/test_audit_g1_dpo_fixes.py b/code/tests/test_audit_g1_dpo_fixes.py new file mode 100644 index 0000000000000000000000000000000000000000..7359acea93d5fe03b277147a9bab698f746d856d --- /dev/null +++ b/code/tests/test_audit_g1_dpo_fixes.py @@ -0,0 +1,321 @@ +"""Regression tests for the 2026-10-14 audit findings on the DPO stage driver +(U03-1, U03-6, U11-1, U11-3, U11-4, U11-5). + +Each test pins the behaviour the audit reproduced, not the source text: + * U03-1 / U11-3 arm_result.json: rank 0 only, atomic, real checkpoint path + * U03-6 the rerun verdict is drawn from the union of every rank's alarms + * U11-1 a resumed arm's probe is relative to the resume point + * U11-4 stale-checkpoint rotation happens once, on rank 0 + * U11-5 dotted overrides are applied or refused, never ignored + +The multi-rank tests are 4-process CPU gloo (CUDA hidden), like +test_trainer_dist.py. +""" + +from __future__ import annotations + +import ast +import importlib.util +import json +import os +import socket +import traceback +from pathlib import Path + +import pytest +import torch +import torch.multiprocessing as mp +from test_trainer import _tiny_trainer, tiny # noqa: F401 (fixture re-export) + +ROOT = Path(__file__).resolve().parents[1] +CONFIG = str(ROOT / "configs" / "dpo_steppairs.toml") +WS = 4 +JOIN_TIMEOUT_S = 180 + + +def _free_port() -> int: + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +# --------------------------------------------------------------------------- +# child bodies (module level: spawn must import them) +# --------------------------------------------------------------------------- + + +def _body_arm_result(rank: int, work: str) -> dict: + from align.train.dpo_stage import write_arm_result + + # Rank-dependent content, as in the real run: checkpoint is "" off rank 0, + # and the lists differ in length (the tail of a longer torn document is + # what made the recorded manifest unparseable). + manifest = { + "arm": "arm_b", + "checkpoint": "/real/ckpt/micro_00000300" if rank == 0 else "", + "tripwire_alarms": [{"alarm": "x" * (50 + 400 * (WS - rank))}] * (3 + rank), + } + path = write_arm_result(os.path.join(work, "arm_result.json"), manifest) + return {"path": path} + + +def _body_rotate(rank: int, work: str) -> dict: + from align.train.dpo_stage import rotate_stale_checkpoints + + ckpt = os.path.join(work, "run", "checkpoints") + dst = rotate_stale_checkpoints(ckpt) + # After the collective, EVERY rank must see the same directory state. + return {"dst": dst, "ckpt_exists": os.path.isdir(ckpt)} + + +def _body_alarms(rank: int, work: str) -> dict: + from align.train.dpo_stage import ( + apply_tripwire_response, + load_stage_config, + reduce_tripwire_alarms, + ) + + local = [{"alarm": "logps_chosen", "step": 7}] if rank == 2 else [] + alarms = reduce_tripwire_alarms(local) + cfg, record = apply_tripwire_response(load_stage_config(CONFIG), alarms) + return {"rerun": record["rerun_required"], "variant": cfg.variant, "alarms": alarms} + + +_BODIES = {"arm_result": _body_arm_result, "rotate": _body_rotate, "alarms": _body_alarms} + + +def _child(rank: int, port: int, mode: str, work: str) -> None: + import torch.distributed as dist + + os.environ.update( + RANK=str(rank), + WORLD_SIZE=str(WS), + LOCAL_RANK=str(rank), + MASTER_ADDR="127.0.0.1", + MASTER_PORT=str(port), + ) + from align.utils.dist import init_distributed + + init_distributed(backend="gloo") + code = 0 + try: + result = _BODIES[mode](rank, work) + with open(os.path.join(work, f"{mode}_{rank}.json"), "w") as f: + json.dump(result, f) + except BaseException: + with open(os.path.join(work, f"{mode}_{rank}.err"), "w") as f: + f.write(traceback.format_exc()) + code = 1 + finally: + try: + if dist.is_initialized(): + dist.destroy_process_group() + except Exception: + pass + os._exit(code) + + +def _run(mode: str, work: str, monkeypatch) -> list[dict]: + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "") + ctx = mp.get_context("spawn") + port = _free_port() + procs = [ctx.Process(target=_child, args=(r, port, mode, work)) for r in range(WS)] + for p in procs: + p.start() + for p in procs: + p.join(JOIN_TIMEOUT_S) + hung = [p for p in procs if p.is_alive()] + for p in hung: + p.kill() + p.join() + errs = [ + f"rank {r}:\n{open(os.path.join(work, f'{mode}_{r}.err')).read()}" + for r in range(WS) + if os.path.isfile(os.path.join(work, f"{mode}_{r}.err")) + ] + assert not hung, f"{len(hung)} rank(s) hung in mode {mode!r}\n" + "\n".join(errs) + assert not errs, "\n".join(errs) + return [json.load(open(os.path.join(work, f"{mode}_{r}.json"))) for r in range(WS)] + + +# --------------------------------------------------------------------------- +# U03-1 / U11-3: arm_result.json +# --------------------------------------------------------------------------- + + +def test_arm_result_written_only_by_rank0_in_a_4_rank_launch(tmp_workdir, monkeypatch): + _run("arm_result", tmp_workdir, monkeypatch) + path = os.path.join(tmp_workdir, "arm_result.json") + got = json.load(open(path)) # unparseable (torn) on the recorded Arm B run + assert got["checkpoint"] == "/real/ckpt/micro_00000300" + assert len(got["tripwire_alarms"]) == 3 # rank 0's content, not a later rank's + assert not [f for f in os.listdir(tmp_workdir) if ".tmp" in f], "tmp file left behind" + + +def test_non_rank0_never_writes_the_arm_result(tmp_workdir, monkeypatch): + from align.train.dpo_stage import write_arm_result + + monkeypatch.setenv("RANK", "1") + monkeypatch.setenv("WORLD_SIZE", "1") + path = os.path.join(tmp_workdir, "arm_result.json") + assert write_arm_result(path, {"checkpoint": ""}) == path + assert not os.path.exists(path) + + +def test_arm_result_write_is_atomic_a_failed_write_keeps_the_previous_file(tmp_workdir, monkeypatch): + from align.train import dpo_stage + + monkeypatch.setenv("RANK", "0") + monkeypatch.setenv("WORLD_SIZE", "1") + path = os.path.join(tmp_workdir, "arm_result.json") + dpo_stage.write_arm_result(path, {"checkpoint": "good"}) + + def boom(*a, **k): + raise RuntimeError("disk died mid-dump") + + monkeypatch.setattr(dpo_stage.json, "dump", boom) + with pytest.raises(RuntimeError, match="mid-dump"): + dpo_stage.write_arm_result(path, {"checkpoint": "bad"}) + assert json.load(open(path)) == {"checkpoint": "good"} + assert os.listdir(tmp_workdir) == ["arm_result.json"] + + +# --------------------------------------------------------------------------- +# U11-4: fresh-arm rotation +# --------------------------------------------------------------------------- + + +def test_stale_checkpoint_rotation_happens_once_and_no_rank_raises(tmp_workdir, monkeypatch): + ckpt = os.path.join(tmp_workdir, "run", "checkpoints") + os.makedirs(os.path.join(ckpt, "micro_00000100")) + res = _run("rotate", tmp_workdir, monkeypatch) # any rank raising fails here + run = os.path.join(tmp_workdir, "run") + stale = [d for d in os.listdir(run) if d.startswith("checkpoints_stale_")] + assert len(stale) == 1 + assert os.path.isdir(os.path.join(run, stale[0], "micro_00000100")) + assert [r["dst"] is not None for r in res] == [True, False, False, False] + assert not any(r["ckpt_exists"] for r in res) + + +def test_rotation_leaves_an_empty_or_missing_dir_alone(tmp_workdir, monkeypatch): + from align.train.dpo_stage import rotate_stale_checkpoints + + monkeypatch.setenv("RANK", "0") + monkeypatch.setenv("WORLD_SIZE", "1") + ckpt = os.path.join(tmp_workdir, "checkpoints") + assert rotate_stale_checkpoints(ckpt) is None + os.makedirs(ckpt) + assert rotate_stale_checkpoints(ckpt) is None + assert os.path.isdir(ckpt) + + +def _dpo_main_ast(): + return ast.parse((ROOT / "scripts" / "dpo.py").read_text()) + + +def test_dpo_script_has_no_unguarded_per_rank_rename(): + renames = [ + n + for n in ast.walk(_dpo_main_ast()) + if isinstance(n, ast.Call) and isinstance(n.func, ast.Attribute) and n.func.attr == "rename" + ] + assert not renames, "rotation must go through dpo_stage.rotate_stale_checkpoints (rank 0 + barrier)" + + +# --------------------------------------------------------------------------- +# U03-6: collective rerun verdict +# --------------------------------------------------------------------------- + + +def test_rerun_verdict_is_identical_on_every_rank_when_one_rank_alarms(tmp_workdir, monkeypatch): + res = _run("alarms", tmp_workdir, monkeypatch) + assert all(r["rerun"] for r in res), res # rank-local filtering: only rank 2 would rerun + assert len({r["variant"] for r in res}) == 1 + assert res[0]["alarms"] == res[3]["alarms"] and res[0]["alarms"][0]["rank"] == 2 + + +def test_dpo_script_reduces_alarms_before_the_rerun_decision(): + tree = _dpo_main_ast() + first = {} + for n in ast.walk(tree): + if isinstance(n, ast.Call) and isinstance(n.func, ast.Name): + first[n.func.id] = min(first.get(n.func.id, 10**9), n.lineno) + assert first["reduce_tripwire_alarms"] < first["apply_tripwire_response"] + for n in ast.walk(tree): + if isinstance(n, ast.Call) and getattr(n.func, "id", "") == "apply_tripwire_response": + assert "summary" not in ast.dump(n.args[1]), "must be fed the reduced list, not the rank-local one" + + +# --------------------------------------------------------------------------- +# U11-1: resume + probe +# --------------------------------------------------------------------------- + + +def test_resumed_arm_probe_is_measured_after_the_resume_point(tiny, tmp_workdir): + from align.train.dpo_stage import probe_micro_cap + + accum, probe_steps = 4, 8 + first, _ = _tiny_trainer( + tiny, tmp_workdir, n_batches=40, grad_accum=accum, save_every_micro_steps=8, max_micro_steps=16 + ) + first.train() + resumed, _ = _tiny_trainer(tiny, tmp_workdir, n_batches=40, grad_accum=accum, save_every_micro_steps=8) + resumed.resume_from_latest() + assert resumed.micro_step >= probe_steps # the audit's precondition + + resumed.cfg.train.probe_steps = probe_steps + resumed.cfg.train.max_micro_steps = probe_micro_cap(resumed.micro_step, probe_steps) + resumed.train() + assert len(resumed.step_times) == probe_steps + probe = resumed.throughput_probe(tokens_per_step=1000) # raised ValueError before the fix + assert probe["n_steps"] == probe_steps // accum and probe["mean_step_time_s"] > 0 + + +def test_probe_cap_is_relative_to_the_resume_position(): + from align.train.dpo_stage import probe_micro_cap + + assert probe_micro_cap(0, 20) == 20 # fresh run: unchanged + assert probe_micro_cap(300, 20) == 320 + + +def test_dpo_script_caps_the_probe_with_probe_micro_cap(): + src = (ROOT / "scripts" / "dpo.py").read_text() + assert "max_micro_steps = probe_micro_cap(trainer.micro_step" in src + assert "max_micro_steps = run_cfg.train.probe_steps" not in src + + +# --------------------------------------------------------------------------- +# U11-5: overrides +# --------------------------------------------------------------------------- + + +def _dpo_script(): + spec = importlib.util.spec_from_file_location("dpo_script_under_test", ROOT / "scripts" / "dpo.py") + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +@pytest.mark.parametrize("spelling", ["--lr=1e-6", "lr=1e-6"]) +def test_both_override_spellings_reach_the_stage_config(spelling): + from align.train.dpo_stage import load_stage_config + + args = _dpo_script()._parse_args([CONFIG, spelling, "--resume"]) + assert args.overrides == ["lr=1e-6"] and args.resume is True + assert load_stage_config(CONFIG, args.overrides).lr == 1e-6 + assert load_stage_config(CONFIG).lr == 5e-7 # no override -> the TOML value + + +def test_override_typo_is_refused_at_launch_by_struct_mode(): + from align.config import ConfigError + from align.train.dpo_stage import load_stage_config + + args = _dpo_script()._parse_args([CONFIG, "--learning_rate=1e-6"]) + with pytest.raises(ConfigError): + load_stage_config(CONFIG, args.overrides) + + +@pytest.mark.parametrize("bad", [["--lr", "1e-6"], ["--bogus-flag"], ["=3"]]) +def test_non_override_leftovers_are_rejected_not_ignored(bad): + with pytest.raises(SystemExit): + _dpo_script()._parse_args([CONFIG, *bad]) diff --git a/code/tests/test_audit_g2_lora_export_fixes.py b/code/tests/test_audit_g2_lora_export_fixes.py new file mode 100644 index 0000000000000000000000000000000000000000..4a795c32e913882e8f35851363658f8f391316f7 --- /dev/null +++ b/code/tests/test_audit_g2_lora_export_fixes.py @@ -0,0 +1,218 @@ +"""Regression tests for the 2026-10-14 audit findings U03-3, U03-4, U05-1, U11-8, U11-9, U11-18 +(LoRA export fidelity, adapter-to-base binding on resume, export-check gating and GPU pinning). + +Run: cd /apps/development/rlhf/align && CUDA_VISIBLE_DEVICES= .venv/bin/python -m pytest \ + tests/test_audit_g2_lora_export_fixes.py -q -p no:cacheprovider +""" + +from __future__ import annotations + +import importlib.util +import logging +import os + +import pytest +import torch +from test_qwen3_parity import make_tiny, rescale_ + +REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + + +def _load(name): + spec = importlib.util.spec_from_file_location(name, os.path.join(REPO, "scripts", f"{name}.py")) + m = importlib.util.module_from_spec(spec) + spec.loader.exec_module(m) + return m + + +exp = _load("qwen3_export_check") + + +# --------------------------------------------------------------------------- +# U05-1 / U03-3: the bf16 merge cast is lossy; the export must MEASURE and (optionally) refuse +# --------------------------------------------------------------------------- + + +def _bf16_lora(b_scale: float): + from align.ppo.lora import LoRAConfig, inject_lora, lora_modules + + m = make_tiny().to(torch.bfloat16) + inject_lora(m, LoRAConfig(r=4, alpha=8.0), adapter="sft", generator=torch.Generator().manual_seed(1)) + g = torch.Generator().manual_seed(2) + with torch.no_grad(): + for _, mod in lora_modules(m): + mod.adapters["sft"].B.copy_(b_scale * torch.randn(mod.adapters["sft"].B.shape, generator=g)) + return m + + +def test_export_reports_retained_delta_and_warns_when_bf16_drops_most_of_it(caplog): + from align.ppo.lora import export_merged_state_dict + + m = _bf16_lora(1e-5) # the shipped PPO lr=1e-6 regime: delta << half an ulp of W + rep: dict = {} + with caplog.at_level(logging.WARNING): + export_merged_state_dict(m, {k: v.detach().clone() for k, v in m.state_dict().items()}, "sft", report=rep) + assert rep["cosine"] < 0.5 and rep["dropped_fraction"] > 0.5, rep + assert rep["layers"] > 0 and rep["min_layer_cosine"] <= rep["cosine"] + 1e-6 + assert any("bf16" in r.message and "cosine" in r.message for r in caplog.records) + + +def test_export_refuses_below_min_cosine_and_passes_a_well_resolved_delta(caplog): + from align.ppo.lora import export_merged_state_dict + + full = lambda m: {k: v.detach().clone() for k, v in m.state_dict().items()} # noqa: E731 + small = _bf16_lora(1e-5) + with pytest.raises(ValueError, match="cosine"): + export_merged_state_dict(small, full(small), "sft", min_cosine=0.9) + big = _bf16_lora(5.0) # delta >> ulp: the cast preserves it + rep: dict = {} + with caplog.at_level(logging.WARNING): + export_merged_state_dict(big, full(big), "sft", report=rep, min_cosine=0.9) + assert rep["cosine"] > 0.95 + assert not any("cosine" in r.message for r in caplog.records) + + +def test_export_of_an_exactly_zero_adapter_is_reported_not_nan(): + from align.ppo.lora import export_merged_state_dict + + m = _bf16_lora(0.0) + rep: dict = {} + export_merged_state_dict(m, {k: v.detach().clone() for k, v in m.state_dict().items()}, "sft", report=rep) + assert rep["delta_norm"] == 0.0 and rep["cosine"] == 1.0 + + +# --------------------------------------------------------------------------- +# U03-4: a LoRA resume binds the adapter to the base it was trained on +# --------------------------------------------------------------------------- + + +def test_lora_payload_carries_a_base_fingerprint_and_a_foreign_base_is_refused(tmp_workdir, monkeypatch): + import test_qwen3_stage as T + + tr, _, _ = T._build(tmp_workdir, "a") + payload = tr.train_module.optim_state_dict() + assert isinstance(payload.get("base_fingerprint"), str) and len(payload["base_fingerprint"]) == 64 + # the same base resumes fine + tr_same, _, _ = T._build(tmp_workdir, "same") + tr_same.train_module.load_optim_state_dict(payload) + # a different base with an identical key set / shapes (repointed snapshot) must not + monkeypatch.setattr(T, "make_tiny", lambda: rescale_(make_tiny(), seed=7)) + tr_other, _, _ = T._build(tmp_workdir, "other") + with pytest.raises(RuntimeError, match="base"): + tr_other.train_module.load_optim_state_dict(payload) + + +def test_legacy_payload_without_a_fingerprint_still_resumes_with_a_warning(tmp_workdir, caplog): + import test_qwen3_stage as T + + tr, _, _ = T._build(tmp_workdir, "a") + payload = tr.train_module.optim_state_dict() + payload.pop("base_fingerprint") + with caplog.at_level(logging.WARNING): + tr.train_module.load_optim_state_dict(payload) + assert any("fingerprint" in r.message for r in caplog.records) + + +def test_fingerprint_tracks_content_not_just_keys(): + from align.train.qwen3_stage import base_fingerprint + + a, b = make_tiny(), make_tiny() + assert base_fingerprint(a.state_dict()) == base_fingerprint(b.state_dict()) + with torch.no_grad(): + b.model.layers[0].attn.q_norm.weight.mul_(1.001) + assert base_fingerprint(a.state_dict()) != base_fingerprint(b.state_dict()) + other = rescale_(make_tiny(), seed=7) + assert base_fingerprint(a.state_dict()) != base_fingerprint(other.state_dict()) + + +# --------------------------------------------------------------------------- +# U11-8: the exit status gates every check the script advertises +# --------------------------------------------------------------------------- + + +def _good_report(): + return { + "verdict": {"passed": True}, + "merge_vs_trainer_file": {"bitwise_mismatches": 0}, + "snapshot_mapping_roundtrip": {"bitwise_equal": True}, + "merge_changed_tensors_vs_snapshot": 12, + "merge_relative_weight_change": {"n": 12, "min": 1e-4, "mean": 4e-4, "max": 1e-3}, + "adapter_vs_merged": {"verdict": {"passed": True}}, + } + + +def test_export_ok_passes_a_complete_report(): + assert exp.export_ok(_good_report()) is True + + +@pytest.mark.parametrize("mutate", [ + lambda r: r.update(snapshot_mapping_roundtrip={"bitwise_equal": False}), + lambda r: r.update(merge_changed_tensors_vs_snapshot=0), # zero adapter + lambda r: r.update(merge_relative_weight_change={"n": 12, "min": 0.0, "mean": 1e-4, "max": 1e-3}), + lambda r: r.pop("merge_changed_tensors_vs_snapshot"), # check did not run + lambda r: r.pop("snapshot_mapping_roundtrip"), + lambda r: r.update(adapter_vs_merged={"verdict": {"passed": False}}), + lambda r: r.pop("adapter_vs_merged"), + lambda r: r.update(merge_vs_trainer_file={"bitwise_mismatches": 2}), + lambda r: r.update(verdict={"passed": False}), + lambda r: r.pop("verdict"), +]) +def test_export_ok_fails_each_unmet_check(mutate): + r = _good_report() + mutate(r) + assert exp.export_ok(r) is False + + +# --------------------------------------------------------------------------- +# U11-9: merged-bf16 vs adapter-applied forward +# --------------------------------------------------------------------------- + + +def _lp(t=64, v=50, seed=0): + g = torch.Generator().manual_seed(seed) + return torch.log_softmax(3 * torch.randn(t, v, generator=g), -1) + + +def test_adapter_vs_merged_passes_identical_and_fails_a_dropped_delta(): + a = _lp() + same = exp.evaluate_adapter_vs_merged(exp.compare_adapter_vs_merged([a], [a.clone()])) + assert same["passed"] is True + other = _lp(seed=5) # an unrelated distribution: what a merge that lost the delta looks like + bad = exp.evaluate_adapter_vs_merged(exp.compare_adapter_vs_merged([a], [other])) + assert bad["passed"] is False + near = a + 0.02 * torch.randn(a.shape, generator=torch.Generator().manual_seed(3)) + ok = exp.evaluate_adapter_vs_merged(exp.compare_adapter_vs_merged([a], [torch.log_softmax(near, -1)])) + assert ok["passed"] is True + + +def test_adapter_vs_merged_shape_mismatch_raises(): + with pytest.raises(ValueError): + exp.compare_adapter_vs_merged([_lp(8)], [_lp(9)]) + + +def test_merge_fidelity_stats_expose_the_dead_zone(): + from align.ppo.lora import lora_state_dict + + m = _bf16_lora(1e-5) + base = {k: v.detach().clone() for k, v in m.state_dict().items() if ".adapters." not in k} + lora = lora_state_dict(m, "sft") + merged = exp.merge_lora_state(base, lora, 2.0) + st = exp.merge_fidelity_stats(base, merged, lora, 2.0) + assert st["cosine"] < 0.5 and st["dropped_fraction"] > 0.5 and st["layers"] > 0 + + +# --------------------------------------------------------------------------- +# U11-18: the vLLM child is pinned to the PARENT's first visible device +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("cvd,expect", [("2,3", "2"), ("3", "3"), (" 1 , 2", "1"), ("GPU-abc,GPU-def", "GPU-abc"), + (None, "0"), ("", "0")]) +def test_first_visible_gpu_follows_the_parent_pinning(cvd, expect): + env = {} if cvd is None else {"CUDA_VISIBLE_DEVICES": cvd} + assert exp._first_visible_gpu(env) == expect + + +def test_rollout_env_uses_the_pinned_device_not_physical_zero(monkeypatch): + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "2,3") + assert exp._rollout_env(exp._first_visible_gpu())["CUDA_VISIBLE_DEVICES"] == "2" diff --git a/code/tests/test_audit_g3_trainer_control_dist.py b/code/tests/test_audit_g3_trainer_control_dist.py new file mode 100644 index 0000000000000000000000000000000000000000..687043843e268bdfb80a0dbabd3ce6d1f2538d0e --- /dev/null +++ b/code/tests/test_audit_g3_trainer_control_dist.py @@ -0,0 +1,112 @@ +"""G3 audit regressions that need two real ranks (CPU, gloo): U01-3. + +Reuses the 2-rank harness pieces of test_trainer_dist (`_build_trainer`, +`_free_port`); only the child entry point is local, because spawn resolves the +body by module and `test_trainer_dist._BODIES` cannot see bodies defined here. +""" + +from __future__ import annotations + +import json +import os +import sys +import traceback + +import torch.multiprocessing as mp +from test_trainer_dist import JOIN_TIMEOUT_S, WS, _build_trainer, _free_port + + +def _body_pin_on_rank1_only(rank: int, run_dir: str) -> dict: + """The R5 pin flag is set on RANK 1 ONLY (a rank-local tripwire fired + there), then both ranks train through saves at micro 4 and 8 with keep-N=1.""" + trainer, _model = _build_trainer(run_dir, init_seed=0) + trainer.cfg.train.keep_checkpoints = 1 + if rank == 1: + trainer._pin_next_save = True + trainer.train() + root = os.path.join(run_dir, "checkpoints") + return { + "dirs": sorted(d for d in os.listdir(root) if not d.endswith(".partial")), + "pinned": sorted(trainer.pinned_checkpoints), + "flag_left": bool(trainer._pin_next_save), + } + + +_BODIES = {"pin_rank1": _body_pin_on_rank1_only} + + +def _child(rank: int, port: int, mode: str, run_dir: str) -> None: + import torch.distributed as dist + + os.environ.update( + RANK=str(rank), + WORLD_SIZE=str(WS), + LOCAL_RANK=str(rank), + MASTER_ADDR="127.0.0.1", + MASTER_PORT=str(port), + ) + from align.utils.dist import init_distributed + + init_distributed(backend="gloo") + code = 0 + try: + result = _BODIES[mode](rank, run_dir) + with open(os.path.join(run_dir, f"{mode}_{rank}.json"), "w") as f: + json.dump(result, f) + except BaseException: + with open(os.path.join(run_dir, f"{mode}_{rank}.err"), "w") as f: + f.write(traceback.format_exc()) + code = 1 + finally: + try: + dist.destroy_process_group() + except Exception: + pass + if code: + sys.exit(code) + + +def _run(mode: str, run_dir: str, monkeypatch) -> list[dict]: + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "") + ctx = mp.get_context("spawn") + port = _free_port() + procs = [ctx.Process(target=_child, args=(r, port, mode, run_dir)) for r in range(WS)] + for p in procs: + p.start() + for p in procs: + p.join(JOIN_TIMEOUT_S) + hung = [p for p in procs if p.is_alive()] + for p in hung: + p.kill() + p.join() + errs = [] + for r in range(WS): + ef = os.path.join(run_dir, f"{mode}_{r}.err") + if os.path.isfile(ef): + errs.append(f"rank {r}:\n{open(ef).read()}") + assert not hung, f"{len(hung)} rank(s) HUNG in mode {mode!r}\n" + "\n".join(errs) + assert not errs, "\n".join(errs) + return [json.load(open(os.path.join(run_dir, f"{mode}_{r}.json"))) for r in range(WS)] + + +def test_u01_3_pin_set_on_a_non_zero_rank_protects_the_checkpoint_on_rank_0( + tmp_workdir, monkeypatch +): + """U01-3. `_pin_next_save` is raised by a RANK-LOCAL tripwire verdict, but + pruning and resume.pt are rank 0's. A pin that exists on rank 1 only left + rank 0's pinned set empty: keep-N deleted the very checkpoint meant as the + rewind point and a resume forgot the pin. The pin must be collective. + Mutant: drop the all_reduce in save_checkpoint -> dirs == [micro_00000008].""" + rows = _run("pin_rank1", tmp_workdir, monkeypatch) + for r in rows: + assert r["pinned"] == [4], rows + assert r["flag_left"] is False, rows + assert rows[0]["dirs"] == ["micro_00000004", "micro_00000008"], rows + import torch + + state = torch.load( + os.path.join(tmp_workdir, "checkpoints", "micro_00000004", "resume.pt"), + weights_only=True, + map_location="cpu", + ) + assert state["pinned_checkpoints"] == [4] diff --git a/code/tests/test_audit_g3_trainer_control_stage_smoke.py b/code/tests/test_audit_g3_trainer_control_stage_smoke.py new file mode 100644 index 0000000000000000000000000000000000000000..a43fb1083e7a5490752362c35cb1a1415093c19e --- /dev/null +++ b/code/tests/test_audit_g3_trainer_control_stage_smoke.py @@ -0,0 +1,110 @@ +"""G3 audit regressions for align/train/stage_smoke.py (CPU): U01-4. + +U01-4: the stage-11 quality gate evaluated the pretrained BASE snapshot (the +`checkpoint` argument only located tokenizer.json) and its generate_fn did not +match the harness call `generate_fn(prompts, cfg, seeds)`. +""" + +from __future__ import annotations + +import os +from types import SimpleNamespace + +import pytest +import torch +from test_model_parity import tiny_config + +from align.model import checkpoint as ckpt +from align.model.mistral import NmoMistralForCausalLM + + +def _tiny(consts, seed): + mc = tiny_config(consts) + torch.manual_seed(seed) + return mc, NmoMistralForCausalLM(mc, attn_implementation="sdpa").eval() + + +def _save(model, mc, path): + os.makedirs(path, exist_ok=True) + ckpt.save_nmo_as_hf(model.state_dict(), mc, path, dtype=torch.bfloat16) + return path + + +@pytest.fixture +def no_fsdp(monkeypatch): + from align.train import fsdp + + monkeypatch.setattr(fsdp, "apply_activation_checkpointing", lambda m, **k: m) + monkeypatch.setattr(fsdp, "apply_fsdp", lambda m, *a, **k: m) + + +def test_u01_4_build_stage_model_loads_the_checkpoint_weights_not_the_base( + tmp_workdir, tiny_model_constants, no_fsdp +): + """Both directions: with `weights_dir` the model's weights ARE the + checkpoint's (and differ from the base); without it the base is loaded.""" + from align.train.stage_smoke import build_stage_model + + mc, base = _tiny(tiny_model_constants, seed=0) + _, trained = _tiny(tiny_model_constants, seed=1) + snap = _save(base, mc, os.path.join(tmp_workdir, "snap")) + ck = _save(trained, mc, os.path.join(tmp_workdir, "ck")) + cfg = SimpleNamespace( + paths=SimpleNamespace(model_snapshot=snap), + model=SimpleNamespace(attn_implementation="sdpa"), + ) + + def same(model, ref): + sd = ref.state_dict() + return all(torch.equal(v.float(), sd[k].to(torch.bfloat16).float()) for k, v in model.state_dict().items()) + + assert same(build_stage_model(cfg), base) + got = build_stage_model(cfg, weights_dir=ck) + assert same(got, trained), "stage model did not load the checkpoint weights" + assert not same(got, base) + + +def test_u01_4_stage_generate_fn_evaluates_the_checkpoint_with_the_harness_signature( + tmp_workdir, tiny_model_constants, monkeypatch +): + """`run_stage_eval` must build the model FROM the checkpoint and hand + `run_eval` a generate_fn callable as `generate_fn(prompts, cfg, seeds)` + returning GenOut rows. A sampled eval cannot thread paired seeds through the + engine (one generator per engine), so it must refuse, not drop them.""" + from align.eval.harness import EvalConfig, GenOut + from align.train import stage_smoke + + mc, model = _tiny(tiny_model_constants, seed=3) + ck = os.path.join(tmp_workdir, "ck") + os.makedirs(ck) + open(os.path.join(ck, "tokenizer.json"), "w").close() + seen = {} + + def fake_build(cfg, attn_implementation=None, weights_dir=None): + seen["weights_dir"] = weights_dir + return model + + class Tok: + eos_token_id = 2 + pad_token_id = None + + def __call__(self, text): + return {"input_ids": [1] + [3 + (ord(c) % 50) for c in text][:6]} + + def decode(self, ids, skip_special_tokens=False): + return " ".join(str(i) for i in ids) + + monkeypatch.setattr(stage_smoke, "build_stage_model", fake_build) + import align.tokenizer as tokmod + + monkeypatch.setattr(tokmod, "load_tokenizer", lambda src: Tok()) + cfg = SimpleNamespace( + paths=SimpleNamespace(model_snapshot=tmp_workdir), + train=SimpleNamespace(seq_len=64), + ) + fn = stage_smoke._stage_generate_fn(cfg, ck) + assert seen["weights_dir"] == ck + outs = fn(["ab", "cde"], EvalConfig(max_new_tokens=3), [11, 12]) + assert len(outs) == 2 and all(isinstance(o, GenOut) for o in outs) + with pytest.raises(ValueError, match="seed"): + fn(["ab"], EvalConfig(max_new_tokens=3, temperature=0.7), [11]) diff --git a/code/tests/test_audit_g3_trainer_control_stop.py b/code/tests/test_audit_g3_trainer_control_stop.py new file mode 100644 index 0000000000000000000000000000000000000000..9c1e3a57ea3697c24391359b05ce30521d1de169 --- /dev/null +++ b/code/tests/test_audit_g3_trainer_control_stop.py @@ -0,0 +1,261 @@ +"""G3 audit regressions: U01-1 (a stop must be able to finish) and U01-2 (an +operator's stop must not be relaunched). CPU only. + +Trainer half: single process, REAL signals / files. Supervisor half: REAL child +process groups, as in tests/test_supervisor.py. +""" + +from __future__ import annotations + +import json +import os +import signal +import subprocess +import sys +import textwrap +import time + +import pytest +from test_control_plane import _cfg, _signal_at +from test_model_parity import tiny_config +from test_trainer import _tiny_trainer + +from align.train.trainer import LossOutput +from align.utils import supervisor as sv +from align.utils.supervisor import Supervisor + + +@pytest.fixture +def tiny(tiny_model_constants): + return tiny_config(tiny_model_constants) + + +def _marker(run_dir, name): + path = os.path.join(run_dir, name) + return json.load(open(path)) if os.path.exists(path) else None + + +def _at_micro(micro, fn): + """wrap_adapter running `fn()` on the `micro`-th adapter call.""" + + def wrap(inner): + n = {"i": 0} + + def adapter(m, batch): + n["i"] += 1 + if n["i"] == micro: + fn() + return inner(m, batch) + + return adapter + + return wrap + + +# --------------------------------------------------------------------------- +# trainer side +# --------------------------------------------------------------------------- + + +def test_u01_2_signal_stop_leaves_a_stop_marker_the_supervisor_can_read(tiny, tmp_workdir): + """torchrun exits 1 whatever its worker exited with, so the checkpointed 143 + never reached the supervisor and `kill ` was relaunched. The + trainer now leaves STOP.json (128+signum). Mutant: drop + `_write_stop_marker` -> no marker.""" + t0 = time.time() + t, _ = _tiny_trainer( + tiny, tmp_workdir, save_every_micro_steps=8, wrap_adapter=_signal_at(3, signal.SIGTERM) + ) + summary = t.train() + assert summary["stop_reason"] == "signal" + rec = sv.read_stop_marker(tmp_workdir, since=t0) + assert rec is not None and rec["code"] == 128 + int(signal.SIGTERM), rec + assert t.stop_exit_code == rec["code"] + + +def test_u01_2_a_run_that_ends_on_its_own_leaves_no_stop_marker(tiny, tmp_workdir): + t, _ = _tiny_trainer(tiny, tmp_workdir, save_every_micro_steps=8) + assert t.train()["stop_reason"] == "data_exhausted" + assert _marker(tmp_workdir, sv.STOP_MARKER_NAME) is None + + +def test_u01_1_a_stop_request_file_stops_at_the_boundary_and_saves(tiny, tmp_workdir): + """The signal-free drain channel. The request lands mid-window (micro 3); + the stop is the window's boundary (micro 4) and writes the checkpoint + (save_every is 8, so micro 4 exists only because the stop wrote it). + Mutant: do not poll the file -> the run trains all 8 micros.""" + t, _ = _tiny_trainer( + tiny, + tmp_workdir, + save_every_micro_steps=8, + wrap_adapter=_at_micro(3, lambda: sv.write_run_marker(tmp_workdir, sv.STOP_REQUEST_NAME)), + ) + summary = t.train() + assert summary["stop_reason"] == "signal", summary + assert summary["micro_step"] == 4 + assert os.path.isdir(os.path.join(tmp_workdir, "checkpoints", "micro_00000004")) + assert t.stop_exit_code == 128 + int(signal.SIGTERM) + assert sv.read_stop_marker(tmp_workdir) is not None + + +def test_u01_1_a_stale_stop_request_from_an_earlier_run_is_ignored(tiny, tmp_workdir): + """A request left on disk by a previous attempt must not stop the next one.""" + with open(os.path.join(tmp_workdir, sv.STOP_REQUEST_NAME), "w") as f: + json.dump({"written_at": time.time() - 3600}, f) + t, _ = _tiny_trainer(tiny, tmp_workdir, save_every_micro_steps=8) + assert t.train()["stop_reason"] == "data_exhausted" + + +def test_u01_1_the_trainer_announces_drain_capability_only_while_training(tiny, tmp_workdir): + seen = {} + + def peek(): + seen["during"] = sv.read_run_marker(tmp_workdir, sv.DRAIN_READY_NAME) is not None + + t, _ = _tiny_trainer(tiny, tmp_workdir, save_every_micro_steps=8, wrap_adapter=_at_micro(1, peek)) + assert sv.read_run_marker(tmp_workdir, sv.DRAIN_READY_NAME) is None + t.train() + assert seen["during"] is True + assert sv.read_run_marker(tmp_workdir, sv.DRAIN_READY_NAME) is None + + +# --------------------------------------------------------------------------- +# supervisor side +# --------------------------------------------------------------------------- + +STOPPED_CHILD = textwrap.dedent( + """ + import json, os, sys, time + run = sys.argv[1] + cnt = os.path.join(run, "attempt.txt") + n = int(open(cnt).read()) if os.path.exists(cnt) else 0 + open(cnt, "w").write(str(n + 1)) + if n == 0: + time.sleep(0.05) + with open(os.path.join(run, "history.jsonl"), "a") as f: + f.write("{}\\n") + os.makedirs(os.path.join(run, "checkpoints", "micro_%08d" % 100)) + if sys.argv[2] == "marker": + with open(os.path.join(run, "STOP.json"), "w") as f: + json.dump({"code": 143, "reason": "t", "written_at": time.time()}, f) + sys.exit(1) # torchrun: exit 1 whatever the worker did + sys.exit(0) + """ +) + + +def _sup(tmp, mode, **kw): + script = os.path.join(tmp, "child.py") + with open(script, "w") as f: + f.write(STOPPED_CHILD) + run = os.path.join(tmp, "run") + os.makedirs(run, exist_ok=True) + return Supervisor( + [sys.executable, script, run, mode], run, backoff_s=0, preflight=lambda: None, **kw + ), run + + +def test_u01_2_supervisor_does_not_relaunch_a_stop_the_trainer_marked(tmp_workdir): + """rc == 1 with progress made would normally be retried; STOP.json says an + operator asked for the stop. Mutant: ignore the marker -> 2 attempts.""" + sup, run = _sup(tmp_workdir, "marker") + res = sup.run(install_signal_handlers=False) + assert (res.reason, res.attempts, res.exit_code) == ("operator_stop", 1, 143) + + +def test_u01_2_retry_on_stop_still_relaunches(tmp_workdir): + sup, _ = _sup(tmp_workdir, "marker", retry_on_stop=True) + res = sup.run(install_signal_handlers=False) + assert (res.reason, res.attempts) == ("success", 2) + + +def test_u01_2_without_a_marker_exit_1_with_progress_is_still_a_retryable_crash(tmp_workdir): + sup, _ = _sup(tmp_workdir, "nomarker") + res = sup.run(install_signal_handlers=False) + assert (res.reason, res.attempts) == ("success", 2) + + +def test_u01_2_a_stale_stop_marker_is_ignored(tmp_workdir): + sup, run = _sup(tmp_workdir, "nomarker") + with open(os.path.join(run, "STOP.json"), "w") as f: + json.dump({"code": 143, "reason": "old", "written_at": time.time() - 3600}, f) + res = sup.run(install_signal_handlers=False) + assert (res.reason, res.attempts) == ("success", 2) + + +# A drain-capable child: announces DRAIN_READY, IGNORES SIGTERM, and on seeing +# STOP_REQUEST takes `take_s` seconds to "checkpoint" before exiting 143. If the +# supervisor signalled or killed it first, `done` is never written. +DRAINING_CHILD = textwrap.dedent( + """ + import json, os, signal, sys, time + run, take_s, mode = sys.argv[1], float(sys.argv[2]), sys.argv[3] + signal.signal(signal.SIGTERM, signal.SIG_IGN) + if mode == "capable": + with open(os.path.join(run, "DRAIN_READY.json"), "w") as f: + json.dump({"written_at": time.time()}, f) + open(os.path.join(run, "started"), "w").close() + while True: + if os.path.exists(os.path.join(run, "STOP_REQUEST")): + time.sleep(take_s) + open(os.path.join(run, "done"), "w").close() + sys.exit(143) + time.sleep(0.05) + """ +) + + +def _drain_run(tmp, take_s, mode, drain_s, grace_s=1): + script = os.path.join(tmp, "drainer.py") + with open(script, "w") as f: + f.write(DRAINING_CHILD) + run = os.path.join(tmp, "run") + os.makedirs(run) + runner = ( + "import sys;from align.utils.supervisor import Supervisor;" + f"r=Supervisor([sys.executable,{script!r},{run!r},{str(take_s)!r},{mode!r}],{run!r}," + f"grace_s={grace_s},drain_s={drain_s},max_retries=0," + "preflight=lambda:None).run();sys.exit(r.exit_code)" + ) + sup = subprocess.Popen([sys.executable, "-c", runner], env=dict(os.environ, PYTHONPATH=os.getcwd())) + for _ in range(200): + if os.path.exists(os.path.join(run, "started")): + break + time.sleep(0.05) + assert os.path.exists(os.path.join(run, "started")) + time.sleep(0.3) + t0 = time.time() + os.kill(sup.pid, signal.SIGTERM) + return sup, run, t0 + + +def _events(run): + return [json.loads(l) for l in open(os.path.join(run, "supervisor.jsonl"))] + + +def test_u01_1_supervisor_drains_a_capable_trainer_instead_of_killing_it(tmp_workdir): + """The child needs 3s to finish its stop; grace_s is 1. A SIGTERM-then- + SIGKILL supervisor kills it mid-"save" (no `done`). Mutant: skip the + drain -> `done` is missing.""" + sup, run, _ = _drain_run(tmp_workdir, take_s=3, mode="capable", drain_s=30) + assert sup.wait(timeout=30) == 128 + int(signal.SIGTERM) + assert os.path.exists(os.path.join(run, "done")), "the stop was killed before it finished" + assert any(e["event"] == "drain_requested" for e in _events(run)) + assert not os.path.exists(os.path.join(run, sv.STOP_REQUEST_NAME)), "request left behind" + + +def test_u01_1_a_drain_that_never_finishes_is_escalated_after_drain_s(tmp_workdir): + sup, run, t0 = _drain_run(tmp_workdir, take_s=600, mode="capable", drain_s=1.5) + assert sup.wait(timeout=30) == 128 + int(signal.SIGTERM) + assert not os.path.exists(os.path.join(run, "done")) + assert any(e["event"] == "drain_timeout" for e in _events(run)) + assert time.time() - t0 < 25 + + +def test_u01_1_a_child_that_cannot_drain_is_signalled_immediately(tmp_workdir): + """No DRAIN_READY -> the legacy path (SIGTERM, grace, SIGKILL): the + supervisor must not wait drain_s on a child that will never answer.""" + sup, run, t0 = _drain_run(tmp_workdir, take_s=600, mode="plain", drain_s=300) + assert sup.wait(timeout=30) == 128 + int(signal.SIGTERM) + assert time.time() - t0 < 20 + assert not any(e["event"] == "drain_requested" for e in _events(run)) diff --git a/code/tests/test_audit_g4_rl_batch_ppo.py b/code/tests/test_audit_g4_rl_batch_ppo.py new file mode 100644 index 0000000000000000000000000000000000000000..a17b6881fa11014982dcfcf45cbff97aa58bf4f0 --- /dev/null +++ b/code/tests/test_audit_g4_rl_batch_ppo.py @@ -0,0 +1,37 @@ +"""Audit 2026-10-14 (G4-rl-batch) regression for `align/ppo/trainer.py` (U05-2).""" + +from __future__ import annotations + +import torch +from test_ppo import _build_module, _make_rollout, _micro, _perturb, _ppo_cfg, tiny # noqa: F401 + + +def test_u05_2_a_skipped_micro_with_non_finite_logprobs_still_reaches_every_parameter(tiny, monkeypatch): + """FSDP2 reduce-scatter fires from the parameters' grad hooks. A rank whose `new_lp` is + non-finite used to skip on a DETACHED leaf: backward touched no parameter, the rank never + joined the collective its finite-`new_lp` peers sat in, and the job hung until the watchdog. + The skip must back-propagate through the model whatever `new_lp` holds.""" + import align.ppo.trainer as T + + cfg = _ppo_cfg() + mod = _build_module(tiny, cfg) + exp = mod.prepare_experience(_make_rollout(tiny, cfg)) + _perturb(mod.model, "policy", seed=5, scale=0.1) + mod.set_train() + mod.train_batch(_micro(mod, exp, iteration=1)) # a healthy micro first (arms the tripwire) + for p in mod.model.parameters(): + p.grad = None + + real = T.policy_forward + + def nan_forward(*a, **k): + new_lp, values = real(*a, **k) + return new_lp * float("nan"), values # graph-connected, non-finite + + monkeypatch.setattr(T, "policy_forward", nan_forward) + out = mod.train_batch(_micro(mod, exp, iteration=1, rows=(2, 3))) + assert out.metrics["ppo/skipped_micro"] == 1.0 + grads = [p.grad for p in mod.model.parameters() if p.requires_grad] + assert grads and all(g is not None for g in grads), ( + "a non-finite-new_lp skip must still back-propagate through every adapter (FSDP2 lockstep)" + ) diff --git a/code/tests/test_audit_g4_rl_batch_source.py b/code/tests/test_audit_g4_rl_batch_source.py new file mode 100644 index 0000000000000000000000000000000000000000..0f3cf6bf4728302760034535d5f95be0b7e1dbc3 --- /dev/null +++ b/code/tests/test_audit_g4_rl_batch_source.py @@ -0,0 +1,231 @@ +"""Audit 2026-10-14 (G4-rl-batch) regressions for `align/rl/batch_source.py`. + +U04-1 fewer live rows than world x accum slots killed the run (RuntimeError after the + rollout was paid for); +U04-2 row quantisation threw away up to ~50% of live rows and toggled the step size 2x; +U04-3 peers waited for rank 0's rollout on a collective bounded by the process-group + timeout, and a rank-0 failure left them blocked. +""" + +from __future__ import annotations + +import os +import tempfile +import time + +import pytest +import torch +from test_grpo_module import _batch, _module +from test_grpo_driver import FakeClient, _source + +from align.rl.batch_source import GrpoBatchSource, RolloutMicroHandle + + +def _payload(B: int, T: int = 5) -> dict[str, torch.Tensor]: + """B group-contiguous rows; row i is identifiable by input_ids[i, 0] == i + 1.""" + ids = torch.zeros((B, T), dtype=torch.int64) + ids[:, 0] = torch.arange(1, B + 1) + rm = torch.zeros((B, T), dtype=torch.int64) + rm[:, 2:] = 1 + return { + "input_ids": ids, + "attention_mask": torch.ones((B, T), dtype=torch.int64), + "response_mask": rm, + "advantages": torch.ones(B), + "vllm_logprobs": torch.zeros((B, T - 1)), + "group_ids": torch.arange(B) // 16, + "reward::reward/total": torch.ones(B), + } + + +def _bare(world: int, accum: int, rank: int = 0, prompts: int = 32, g: int = 16) -> GrpoBatchSource: + src = GrpoBatchSource.__new__(GrpoBatchSource) + src.world_size, src.grad_accum, src.rank = world, accum, rank + src.prompts_per_step, src.n_samples_per_prompt = prompts, g + src.dropped_remainder_rows = 0 + return src + + +def _real_rows(slices) -> list[int]: + """Row ids carried by real (non-padding) rows of one rank's slices.""" + out = [] + for s in slices: + for i in range(s["input_ids"].shape[0]): + if int(s["response_mask"][i].sum()) > 0: + out.append(int(s["input_ids"][i, 0])) + return out + + +def _all_ranks(world, accum, B): + per_rank = [] + for r in range(world): + src = _bare(world, accum, rank=r) + per_rank.append(src._slice_rank(_payload(B))) + return per_rank + + +def test_u04_1_fewer_live_rows_than_slots_still_trains_every_row(): + """World 3 x accum 85 = 255 slots; 240 live rows (17 of 32 groups dead) used to raise.""" + per_rank = _all_ranks(3, 85, 240) + for slices in per_rank: + assert len(slices) == 85 + assert all(s["input_ids"].shape[0] >= 1 for s in slices), "a rank would run an empty micro" + used = sorted(i for slices in per_rank for i in _real_rows(slices)) + assert used == list(range(1, 241)), "every live row must reach exactly one micro" + + +def test_u04_1_padding_rows_contribute_nothing_and_are_flagged(): + per_rank = _all_ranks(3, 85, 240) + pad_rows = [ + s + for slices in per_rank + for s in slices + if int(s["response_mask"].sum()) == 0 + ] + assert len(pad_rows) == 255 - 240 + for s in pad_rows: + assert float(s["advantages"].abs().sum()) == 0.0 + assert s["input_ids"].shape[0] == 1 + + +def test_u04_2_one_dead_group_no_longer_discards_half_the_rows(): + """B=496 (one dead group of 32x16, 3 ranks, accum 85) used 255 rows and dropped 241.""" + per_rank = _all_ranks(3, 85, 496) + used = sorted(i for slices in per_rank for i in _real_rows(slices)) + assert used == list(range(1, 497)) + src = _bare(3, 85) + src._slice_rank(_payload(496)) + assert src.dropped_remainder_rows == 0 + + +def test_u04_2_full_step_shape_is_unchanged(): + """The nominal full-live shape keeps its per-micro row cap (memory budget): 512 rows over 255 + slots is 2 rows per micro, 2 rows dropped and counted -- exactly the previous behaviour.""" + src = _bare(3, 85) + slices = src._slice_rank(_payload(512)) + assert all(s["input_ids"].shape[0] == 2 for s in slices) + assert src.dropped_remainder_rows == 2 + + +def test_u04_2_micro_sizes_never_exceed_the_nominal_cap_and_ranks_stay_balanced(): + for B in (255, 300, 400, 496, 510, 512): + per_rank = _all_ranks(3, 85, B) + for m in range(85): + sizes = [per_rank[r][m]["input_ids"].shape[0] for r in range(3)] + assert max(sizes) <= 2 and max(sizes) - min(sizes) <= 1, (B, m, sizes) + + +def test_u04_1_pad_micro_runs_through_the_real_module_with_zero_loss_and_live_graph(): + mod, model = _module(grad_accum=1) + src = _bare(1, 4, prompts=1, g=8) + src.device = "cpu" + src.old_logprob_fn = lambda ids, attn: mod.compute_logprobs_nograd(ids, attn) + b = _batch(model, B=2, T=10) + payload = { + "input_ids": b.input_ids, + "attention_mask": b.attention_mask, + "response_mask": b.response_mask, + "advantages": b.advantages, + "vllm_logprobs": b.vllm_logprobs, + "group_ids": b.group_ids, + } + slices = src._slice_rank(payload) # 2 rows over 4 slots: two micros are padding + pads = [s for s in slices if int(s["response_mask"].sum()) == 0] + assert len(pads) == 2 + batch = src._finalize(model, pads[0]) + out = mod._grpo_loss_adapter(model, batch) + assert float(out.loss) == 0.0 + out.loss.backward() + assert all(p.grad is not None and torch.isfinite(p.grad).all() for p in model.parameters()) + + +def test_u04_1_tripwire_is_not_fed_from_a_padding_micro(): + """A padding micro reports entropy 0 / length 0; observing it would fake a collapse.""" + seen = [] + + class _Wire: + fired = False + fire_reason = None + + def observe(self, step, entropy, length): + seen.append((step, entropy, length)) + return False + + src = _source(g=8, accum=2, steps=1) + src.collapse_tripwire = _Wire() + src._on_step_boundary(0, {"entropy": 0.0, "lengths/response_mean": 0.0}, padded=True) + assert seen == [] + src._on_step_boundary(0, {"entropy": 1.0, "lengths/response_mean": 3.0}) + assert seen == [(0, 1.0, 3.0)] + + +# ---------------------------------------------------------------------------- U04-3 (2 real ranks) + + +def _wait_worker(rank, world, init_file, mode, q): + import datetime + + import torch.distributed as dist + + os.environ["ALIGN_PG_TIMEOUT_S"] = "3" + os.environ["WORLD_SIZE"] = str(world) + os.environ["RANK"] = str(rank) + os.environ["ALIGN_ROLLOUT_WAIT_TIMEOUT_S"] = "60" + dist.init_process_group( + "gloo", init_method=f"file://{init_file}", rank=rank, world_size=world, + timeout=datetime.timedelta(seconds=3), + ) + + class _Src(GrpoBatchSource): + def _produce_step(self, optimizer_step): + if mode == "slow": + time.sleep(7) # longer than the 3 s PG timeout, far shorter than the rollout budget + return _payload(4) + raise RuntimeError("rollout engine died (test)") + + src = _Src( + [{"problem": "p", "answer": "1"}], prompts_per_step=1, n_samples_per_prompt=4, + max_response_len=4, temperature=1.0, optimizer_steps=1, grad_accum=1, seed=0, + prompt_builder=lambda p: p, + ) + src.bind(client=None, reward=None, verifier=None, old_logprob_fn=lambda i, a: None, + weights_dir_fn=lambda: "w", rank=rank, world_size=world) + src.old_logprob_fn = lambda ids, attn: torch.zeros(ids.shape[0], ids.shape[1] - 1) + t0 = time.time() + try: + batch = src._materialize_micro(None, RolloutMicroHandle(0, 0)) + q.put((rank, "ok", int(batch.input_ids.shape[0]), time.time() - t0)) + except Exception as e: # noqa: BLE001 + q.put((rank, "err", str(e), time.time() - t0)) + + +def _run_two_ranks(mode): + import torch.multiprocessing as mp + + ctx = mp.get_context("spawn") + q = ctx.Queue() + with tempfile.TemporaryDirectory(dir="/apps/development/rlhf/align/.tmp/audit_2026-10-14/scratch_g4_rl_batch") as d: + f = os.path.join(d, "store") + ps = [ctx.Process(target=_wait_worker, args=(r, 2, f, mode, q)) for r in range(2)] + for p in ps: + p.start() + res = {} + for _ in range(2): + r = q.get(timeout=90) + res[r[0]] = r + for p in ps: + p.join(30) + return res + + +def test_u04_3_peers_survive_a_rank0_rollout_longer_than_the_pg_timeout(): + res = _run_two_ranks("slow") + assert res[0][1] == "ok" and res[1][1] == "ok", res + assert res[1][2] == 2 # 4 rows over 2 ranks + + +def test_u04_3_a_rank0_failure_reaches_the_peers_as_a_named_error_not_a_timeout(): + res = _run_two_ranks("fail") + assert res[0][1] == "err" and "rollout engine died (test)" in res[0][2], res + assert res[1][1] == "err" and "rollout engine died (test)" in res[1][2] and "rank 0" in res[1][2], res + assert res[1][3] < 30, "the peer must be told promptly, not wait out a timeout" diff --git a/code/tests/test_audit_g5_data_store_build.py b/code/tests/test_audit_g5_data_store_build.py new file mode 100644 index 0000000000000000000000000000000000000000..559076b26864b5c2ac642b334d9dd75d6983c85a --- /dev/null +++ b/code/tests/test_audit_g5_data_store_build.py @@ -0,0 +1,155 @@ +"""Regression tests for audit findings U09-4..U09-7 (align/data/build.py, +align/data/prm800k_pairs.py). +Run: uv run pytest tests/test_audit_g5_data_store_build.py -m "not gpu" -v +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from align.data import build as bd +from align.data import prm800k_pairs as pb +from align.data.build import _PREF_FIELDS, _ShardSet, _held_out_ids +from align.data.store import ShardDataset + + +# -- U09-4 -------------------------------------------------------------------- + + +def _fill(set_, pids): + for pid in pids: + set_.add(input_ids=[1, 2, 3], problem_id=pid, subject="s", boundary_positions=[]) + return set_.finalize() + + +def test_u09_4_rebuild_into_reused_dir_does_not_leave_stale_shards(tmp_workdir): + old = _ShardSet(tmp_workdir, "sft-train", rows_per_shard=3) + _fill(old, ["o0", "o1", "o2", "o3", "o4"]) # shards 0000 + 0001 + new = _ShardSet(tmp_workdir, "sft-train", rows_per_shard=3) + _fill(new, ["n0"]) # only 0000 rewritten + ids = [r.problem_id for r in ShardDataset(tmp_workdir).iter_rows()] + assert ids == ["n0"] + + +def test_u09_4_purge_is_scoped_to_its_own_prefix(tmp_workdir): + other = _ShardSet(tmp_workdir, "sft-held") + _fill(other, ["h0"]) + _fill(_ShardSet(tmp_workdir, "sft-train"), ["t0"]) + assert [r.problem_id for r in ShardDataset(tmp_workdir, names=["sft-held-0000"]).iter_rows()] == ["h0"] + + +# -- U09-5 -------------------------------------------------------------------- + + +def _ids(n): + return [f"problem-{i:06d}" for i in range(n)] + + +def test_u09_5_held_out_assignment_is_independent_of_the_id_set(): + ids = _ids(4000) + held_full = _held_out_ids(ids, 0.05, 0) + # Arm-B-like subset: every id's membership must be the same as in the + # full set, otherwise Arm B's held-out set leaks SFT-train problems. + subset = ids[::3] + held_sub = _held_out_ids(subset, 0.05, 0) + assert held_sub == held_full & set(subset) + + +def test_u09_5_adding_one_id_changes_at_most_that_id(): + ids = _ids(4000) + base = _held_out_ids(ids, 0.05, 0) + more = _held_out_ids(ids + ["problem-new"], 0.05, 0) + assert base - more == set() + assert more - base <= {"problem-new"} + + +def test_u09_5_fraction_is_close_to_requested_and_nonempty_for_tiny_inputs(): + held = _held_out_ids(_ids(20000), 0.05, 0) + assert 0.04 * 20000 < len(held) < 0.06 * 20000 + assert len(_held_out_ids(["only-one"], 0.05, 0)) == 1 + assert len(_held_out_ids(_ids(5), 0.05, 0)) >= 1 + assert _held_out_ids([], 0.05, 0) == set() + + +def test_u09_5_seed_changes_the_split(): + assert _held_out_ids(_ids(4000), 0.05, 0) != _held_out_ids(_ids(4000), 0.05, 1) + + +def test_u09_5_rule_is_recorded_in_pins(): + assert bd.PinnedYields().held_out_rule == "hash-v1" + + +# -- U09-6 -------------------------------------------------------------------- + + +def _pair(pid="p0", step=0, chosen="good", rejected="bad", conf=1.0): + return pb.PrefPair( + problem="Q", problem_id=pid, step_index=step, prefix_steps=(), + chosen_step=chosen, rejected_step=rejected, confidence=conf, + ) + + +def test_u09_6_curate_never_emits_duplicate_pair_ids(): + dup = [_pair(), _pair(), _pair(rejected="worse")] + selected, remainder = pb.curate(dup, set(), {"p0": "Algebra"}, target=1) + ids = [p.pair_id for p in selected] + [p.pair_id for p in remainder] + assert len(ids) == len(set(ids)) == 2 + assert not ({p.pair_id for p in selected} & {p.pair_id for p in remainder}) + + +def test_u09_6_curate_dedupe_is_order_independent(): + a, b = _pair(), _pair(rejected="worse") + s1, r1 = pb.curate([a, a, b], set(), {"p0": "Algebra"}, target=1) + s2, r2 = pb.curate([b, a, a], set(), {"p0": "Algebra"}, target=1) + assert [p.pair_id for p in s1] == [p.pair_id for p in s2] + assert [p.pair_id for p in r1] == [p.pair_id for p in r2] + + +# -- U09-7 -------------------------------------------------------------------- + + +class _FakeSpec: + newline_id = 1 + vocab_absorbers = 0 + tokenizer_hash = "t" + + def to_json(self, path): + Path(path).write_text("{}") + + +def test_u09_7_aborted_rebuild_invalidates_the_previous_manifest(tmp_workdir, monkeypatch): + out = Path(tmp_workdir) / "build" + out.mkdir() + (out / "manifest.json").write_text(json.dumps({"pins_verified": True, "build_version": 2})) + monkeypatch.setattr(bd, "derive_boundary_spec", lambda tok: _FakeSpec()) + monkeypatch.setattr(bd, "math_subject_index", lambda src: {}) + + def boom(*a, **k): + raise RuntimeError("strict pin failure") + + monkeypatch.setattr(bd, "build_eval_stage", boom) + with pytest.raises(RuntimeError, match="strict pin"): + bd.build_all(object(), out, object(), bd.PinnedYields()) + assert not (out / "manifest.json").exists() + + +def test_u09_7_atomic_write_never_exposes_a_partial_file(tmp_workdir, monkeypatch): + from align.data import store + + target = Path(tmp_workdir) / "m.json" + target.write_text("OLD") + + def fail_replace(src, dst): + raise OSError("killed before rename") + + monkeypatch.setattr(store.os, "replace", fail_replace) + with pytest.raises(OSError): + store.atomic_write_text(target, "NEW") + assert target.read_text() == "OLD" + assert [p.name for p in Path(tmp_workdir).iterdir()] == ["m.json"] + monkeypatch.undo() + store.atomic_write_text(target, "NEW") + assert target.read_text() == "NEW"