From 185c0dc51e9fbeda3a36e84bf659af0d105c80e0 Mon Sep 17 00:00:00 2001 From: Yuren Hao Date: Mon, 20 Jul 2026 06:33:50 -0500 Subject: =?UTF-8?q?RESULT=2053:=20=CE=B2=E6=95=99=E6=9D=A1=E9=87=8D?= =?UTF-8?q?=E5=BB=BA=20=E2=80=94=20Codex=20=CF=81=E6=8E=A2=E9=92=88?= =?UTF-8?q?=E6=B5=8B=E5=87=BA=E7=9C=9F=E5=AE=9E=E4=B8=8A=E7=95=8C=CE=B2*?= =?UTF-8?q?=E2=89=880.7@s35000(0.4=E6=94=B6=E6=95=9B/0.7=E5=8F=91=E6=95=A3?= =?UTF-8?q?);=20v1=E7=82=B8@0.09=E5=A4=84=CF=81=3D0.73=E8=BF=9C=E6=9C=AA?= =?UTF-8?q?=E8=B6=8A=E7=95=8C=3D=E6=8E=A7=E5=88=B6=E5=99=A8=E8=87=AA?= =?UTF-8?q?=E7=82=B8=E9=9D=9E=CE=B2=E4=B8=8D=E7=A8=B3;=20=E5=9B=BA?= =?UTF-8?q?=E5=AE=9A=CE=B2=E6=B0=B8=E4=B8=8D=E7=82=B8(=E9=81=87=E4=B8=8B?= =?UTF-8?q?=E6=B2=89=E4=B8=8A=E7=95=8C=E2=86=92skip=E9=98=B2=E5=BE=A1,plai?= =?UTF-8?q?n2=E6=94=B6=E5=AE=9811527skip=E4=BB=8D3.3437=E5=AE=8C=E8=B5=9B)?= =?UTF-8?q?,=E5=8F=AA=E6=9C=89=E6=8E=A7=E5=88=B6=E5=99=A8=E7=82=B8;=20?= =?UTF-8?q?=E7=95=8C=E5=86=85CE=E5=B9=B3=E5=9D=A6+=E5=81=8F=E7=BD=AE?= =?UTF-8?q?=E2=88=9D=CE=B2+=E4=B8=8A=E7=95=8C=E4=B8=8B=E6=B2=89=E2=86=92?= =?UTF-8?q?=E6=94=BF=E7=AD=96=3D=E5=9B=BA=E5=AE=9A=E5=B0=8F=CE=B2,?= =?UTF-8?q?=E6=8E=A7=E5=88=B6=E5=99=A8=E7=BA=BF=E6=8B=86=E9=99=A4(?= =?UTF-8?q?=E9=94=99=E8=AF=AF=E4=B8=94=E5=8D=B1=E9=99=A9=E7=9A=84=E8=BD=B4?= =?UTF-8?q?,=E7=94=A8=E6=88=B7=E5=88=A4=E5=AE=9A=E6=88=90=E7=AB=8B);=201e-?= =?UTF-8?q?3=E5=85=A8=E7=A8=8B=E9=AA=8C=E8=AF=81=E5=B7=B2=E6=8E=92?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_014FAPDWQ49M5Ye3NpTndTpn --- ep_run/probe_cx2_rho.py | 316 ++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 316 insertions(+) create mode 100644 ep_run/probe_cx2_rho.py (limited to 'ep_run/probe_cx2_rho.py') diff --git a/ep_run/probe_cx2_rho.py b/ep_run/probe_cx2_rho.py new file mode 100644 index 0000000..7c9f58e --- /dev/null +++ b/ep_run/probe_cx2_rho.py @@ -0,0 +1,316 @@ +#!/usr/bin/env python3 +"""Read-only Cascade-EP nudged-relaxation rho probe. + +This is deliberately standalone: the trainer executes training at import time. The +model definitions, free-state construction, and relax sweep below are copied from +casc_eq_train.py. No optimizer step or parameter update is performed. +""" + +import gc +import math +import os +import pickle +from pathlib import Path + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F + + +ROOT = Path(__file__).resolve().parent +CHECKPOINT = ROOT / "runs/fw72m_plain2_s35000.pt" +DATA_DIR = ROOT / "data/fineweb_edu" +BETAS = [1e-3, 3e-3, 1e-2, 3e-2, 9e-2, 0.2, 0.4, 0.7, 1.0] + +# The requested run is CPU-only in this environment. Keeping this explicit also +# makes it impossible for the probe to select GPU2. +DEVICE = torch.device("cpu") +torch.set_num_threads(min(32, os.cpu_count() or 1)) +torch.manual_seed(1) + + +class RMSNorm(nn.Module): + def __init__(self, C, eps=1e-6): + super().__init__() + self.g = nn.Parameter(torch.ones(C)) + self.eps = eps + + def forward(self, x): + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.g + + +class SwiGLU(nn.Module): + def __init__(self, C): + super().__init__() + h = ((8 * C // 3) + 63) // 64 * 64 + self.w1 = nn.Linear(C, h, bias=False) + self.w3 = nn.Linear(C, h, bias=False) + self.w2 = nn.Linear(h, C, bias=False) + + def forward(self, x): + return self.w2(F.silu(self.w1(x)) * self.w3(x)) + + +class Olmo2Attn(nn.Module): + def __init__(self, C, H, T): + super().__init__() + self.H, self.hd = H, C // H + self.qkv = nn.Linear(C, 3 * C, bias=False) + self.proj = nn.Linear(C, C, bias=False) + self.qn, self.kn = RMSNorm(C), RMSNorm(C) + inv = 1.0 / (500000.0 ** (torch.arange(0, self.hd, 2).float() / self.hd)) + fr = torch.outer(torch.arange(T).float(), inv) + self.register_buffer("rc", fr.cos(), persistent=False) + self.register_buffer("rs", fr.sin(), persistent=False) + + def rope(self, x): + x1, x2 = x[..., ::2], x[..., 1::2] + c, s = self.rc[None, None], self.rs[None, None] + return torch.stack((x1 * c - x2 * s, x1 * s + x2 * c), dim=-1).flatten(-2) + + def forward(self, x): + B, T, C = x.shape + q, k, v = self.qkv(x).split(C, dim=2) + q, k = self.qn(q), self.kn(k) + q = self.rope(q.view(B, T, self.H, self.hd).transpose(1, 2)) + k = self.rope(k.view(B, T, self.H, self.hd).transpose(1, 2)) + v = v.view(B, T, self.H, self.hd).transpose(1, 2) + y = F.scaled_dot_product_attention(q, k, v, is_causal=True) + return self.proj(y.transpose(1, 2).contiguous().view(B, T, C)) + + +class Olmo2Block(nn.Module): + def __init__(self, C, H, T): + super().__init__() + self.attn = Olmo2Attn(C, H, T) + self.ff = SwiGLU(C) + self.na, self.nf = RMSNorm(C), RMSNorm(C) + + def forward(self, z, mask=None): + z = z + self.na(self.attn(z)) + return z + self.nf(self.ff(z)) + + +def fmt(value): + if value is None: + return "NA" + return f"{value:.8e}" + + +print(f"checkpoint={CHECKPOINT}", flush=True) +print(f"device={DEVICE} torch_threads={torch.get_num_threads()}", flush=True) +ck = torch.load(CHECKPOINT, map_location=DEVICE, weights_only=False) +cfg = ck["config"] +L, C, H, T = (int(cfg[k]) for k in ("L", "C", "H", "T")) +vocab = int(ck["tok"]["weight"].shape[0]) +T2_MAX = int(cfg.get("kmax", 8)) +DTOP_EVERY = int(cfg.get("dtop_every", 1)) +GETA = float(cfg.get("geta", 1.0)) +AMP = bool(cfg.get("amp", False) and DEVICE.type == "cuda") +# Below this relative state-change scale, the observed float32 residuals are a +# numerical fixed-point floor and rho becomes noise/noise (the trainer documents +# the same failure mode). Preserve raw GOV rho, but do not call noise a ceiling. +RHO_NOISE_FLOOR = 1e-6 + +tok = nn.Embedding(vocab, C).to(DEVICE) +pos = nn.Embedding(T, C).to(DEVICE) # loaded for checkpoint parity; OLMo2 does not use it +blocks = nn.ModuleList([Olmo2Block(C, H, T) for _ in range(L)]).to(DEVICE) +ln_f = RMSNorm(C).to(DEVICE) +W_out = nn.Parameter(torch.empty(vocab, C, device=DEVICE)) + +tok.load_state_dict(ck["tok"]) +pos.load_state_dict(ck["pos"]) +blocks.load_state_dict(ck["blocks"]) +ln_f.load_state_dict(ck["lnf"]) +with torch.no_grad(): + W_out.copy_(ck["wout"].to(DEVICE)) + +# Restore the checkpoint optimizer state exactly as the trainer does, then discard +# it: this verifies full checkpoint compatibility without ever taking a train step. +all_params = ( + list(tok.parameters()) + + list(blocks.parameters()) + + list(ln_f.parameters()) + + [W_out] +) +from muon import build_hybrid + +opt, sched = build_hybrid( + blocks, + all_params, + float(cfg["lr"]), + float(cfg.get("muon_lr", 0.02)), + int(cfg["warmup"]), + muon_mom=float(cfg.get("muon_mom", 0.95)), + adam_b1=float(cfg.get("adam_b1", 0.9)), + total_steps=(int(cfg["steps"]) if cfg.get("cosine", False) else 0), + lr_min_ratio=float(cfg.get("lr_min_ratio", 0.1)), +) +opt.load_state_dict(ck["opt"]) +print( + f"loaded_step={int(ck.get('step', 0))} model_state=yes optimizer_state=yes " + f"L={L} C={C} H={H} T={T} vocab={vocab}", + flush=True, +) +del opt, sched +ck.pop("opt", None) +gc.collect() + +for module in (tok, pos, blocks, ln_f): + module.eval() + +mask = torch.triu(torch.full((T, T), float("-inf"), device=DEVICE), 1) + + +def emb(x): + return tok(x) + + +def readout(z): + return ln_f(z) @ W_out.t() + + +def obj_loss(logits2d, y1d): + return F.cross_entropy(logits2d, y1d) + + +def free_states_graphed(x): + """Exact trainer free-state construction.""" + with torch.no_grad(): + z0 = emb(x) + ins, outs, zs = [], [], [] + prev = z0 + with torch.autocast("cuda", dtype=torch.bfloat16, enabled=AMP): + for block in blocks: + i = prev.detach().requires_grad_(True) + o = block(i, mask) + ins.append(i) + outs.append(o) + zs.append(o.detach().float()) + prev = zs[-1] + return z0, zs, ins, outs + + +def relax_probe(z0, zs, ins, outs, y, beta, K, x): + """Trainer relax(), plus a read-only copy of every per-sweep residual.""" + d = [None] * L + geta_l = GETA + residuals = [] + + def forces(refresh_top): + if refresh_top or d[L - 1] is None: + zc = zs[L - 1].detach().requires_grad_(True) + ce = obj_loss(readout(zc).reshape(-1, vocab), y.reshape(-1)) + nbt_loc = zc.shape[0] * zc.shape[1] + g = torch.autograd.grad(ce, zc)[0] + d[L - 1] = (-beta * nbt_loc * g).detach() + for layer in range(L - 2, -1, -1): + d[layer] = torch.autograd.grad( + outs[layer + 1], + ins[layer + 1], + grad_outputs=d[layer + 1].to(outs[layer + 1].dtype), + )[0].detach().float() + + def rebuild(last): + nonlocal ins, outs + prev = z0 + n_ins, n_outs = [], [] + rnum = rden = 0.0 + g_eff = 1.0 if last else geta_l + with torch.autocast("cuda", dtype=torch.bfloat16, enabled=AMP): + for layer in range(L): + if last and layer == 0: + i = emb(x) + else: + i = prev.detach().requires_grad_(True) + o = blocks[layer](i, mask) + znew = o.detach().float() + d[layer] + mixed = znew if g_eff >= 1.0 else ( + zs[layer] + g_eff * (znew - zs[layer]) + ) + with torch.no_grad(): + rnum += float((mixed - zs[layer]).norm()) + rden += float(zs[layer].norm()) + zs[layer] = mixed + n_ins.append(i) + n_outs.append(o) + prev = zs[layer] + ins, outs = n_ins, n_outs + return rnum / max(rden, 1e-9) + + for k in range(K): + forces(k % DTOP_EVERY == 0) + residuals.append(rebuild(k + 1 == K)) + + rho = None + if len(residuals) >= 2 and residuals[-2] > 1e-12: + rho = residuals[-1] / residuals[-2] + gov = {"res": residuals[-1], "rho": rho, "kuse": K} + return residuals, gov + + +# One deterministic, fixed batch from the same binary loader as the trainer. B=1 +# is sufficient because the trainer's nbt_loc factor cancels CE's batch averaging. +train_data = np.memmap(DATA_DIR / "train.bin", dtype=np.uint16, mode="r") +batch_gen = torch.Generator().manual_seed(1 * 7919 + 11) +offset = int(torch.randint(len(train_data) - T - 1, (1,), generator=batch_gen).item()) +x = torch.from_numpy(train_data[offset : offset + T].astype(np.int64))[None].to(DEVICE) +y = torch.from_numpy(train_data[offset + 1 : offset + T + 1].astype(np.int64))[None].to(DEVICE) +print( + f"batch=fineweb_edu/train.bin offset={offset} B=1 T={T} " + f"T2_max={T2_MAX} dtop_every={DTOP_EVERY} geta={GETA} amp={AMP}", + flush=True, +) +print( + f"convergence_rule=numerical_fixed_point(res<={RHO_NOISE_FLOOR:g}) or " + "finite(last_above_floor_rho)<1 (fixed-K trainer has no tolerance stop)", + flush=True, +) + +results = [] +for beta in BETAS: + z0, zs, ins, outs = free_states_graphed(x) + try: + residuals, gov = relax_probe(z0, zs, ins, outs, y, beta, T2_MAX, x) + rhos = [ + (None if i == 0 or residuals[i - 1] <= 1e-12 else residuals[i] / residuals[i - 1]) + for i in range(len(residuals)) + ] + final_res = gov["res"] + raw_gov_rho = gov["rho"] + meaningful = [ + rhos[i] + for i in range(1, len(rhos)) + if rhos[i] is not None + and (residuals[i - 1] > RHO_NOISE_FLOOR or residuals[i] > RHO_NOISE_FLOOR) + ] + rho = meaningful[-1] if meaningful else raw_gov_rho + reached_floor = any(r <= RHO_NOISE_FLOOR for r in residuals) + converged = bool( + math.isfinite(final_res) + and rho is not None + and math.isfinite(rho) + and (reached_floor or rho < 1.0) + ) + print( + f"TRACE beta={beta:g} residuals=[{','.join(fmt(v) for v in residuals)}] " + f"rhos=[{','.join(fmt(v) for v in rhos)}] " + f"raw_final_gov_rho={fmt(raw_gov_rho)}", + flush=True, + ) + except (RuntimeError, FloatingPointError) as exc: + final_res, rho, converged = float("nan"), None, False + print(f"TRACE beta={beta:g} relaxation_error={type(exc).__name__}:{exc}", flush=True) + + results.append((beta, final_res, converged, rho)) + print( + f"RESULT beta={beta:g} final_res={fmt(final_res)} " + f"converged={'yes' if converged else 'no'} rho_at_convergence={fmt(rho)}", + flush=True, + ) + del z0, zs, ins, outs + gc.collect() + +beta_star = next((beta for beta, _, converged, rho in results if (not converged) or (rho is not None and rho >= 1.0)), None) +print(f"BETA_STAR {beta_star if beta_star is not None else 'NONE'}", flush=True) -- cgit v1.2.3