From 15ed315b3f43aa35532a5dfea3f9884d63fba71a Mon Sep 17 00:00:00 2001 From: Yuren Hao Date: Sat, 18 Jul 2026 23:12:12 -0500 Subject: =?UTF-8?q?wsync:=20=E5=90=8C=E6=AD=A5=E6=9D=83=E9=87=8D=E6=AD=A5?= =?UTF-8?q?=E6=8E=A5=E5=8F=97(=E7=94=A8=E6=88=B7=E9=A3=8E=E6=A0=BC?= =?UTF-8?q?=E9=87=8D=E8=AE=BE=E8=AE=A1,knee=E6=92=A4=E5=9B=9E)=20=E2=80=94?= =?UTF-8?q?=20opt.step=E5=89=8D=E5=BF=AB=E7=85=A7params+=E5=8A=A8=E9=87=8F?= =?UTF-8?q?,=E4=B8=8B=E6=AD=A5nudged=20relax=E8=BF=87=E5=90=8C=E4=B8=80=5F?= =?UTF-8?q?legal=E9=97=A8,=E9=9D=9E=E6=B3=95=E5=8D=B3p=E2=86=90(p+snap)/2?= =?UTF-8?q?=E6=8A=98=E5=8D=8A=E9=87=8D=E6=94=BE(=E9=A2=84=E7=AE=97?= =?UTF-8?q?=E5=A4=8D=E7=94=A8beta=5Fsync=E8=AF=AD=E4=B9=89),=E6=8A=98?= =?UTF-8?q?=E5=8D=8A=E5=88=B0=E5=BA=95=E6=95=B4=E6=AD=A5=E5=9B=9E=E6=BB=9A?= =?UTF-8?q?+skip;=E9=9B=B6=E6=96=B0=E5=B8=B8=E6=95=B0,=E6=A8=A1=E5=9E=8B?= =?UTF-8?q?=E5=87=BD=E6=95=B0=E9=9B=B6=E6=94=B9=E5=8A=A8;42M=E5=BC=95?= =?UTF-8?q?=E7=88=86=E9=A2=84=E6=A3=80+=E6=AD=BB=E4=BA=A1=E7=AA=97?= =?UTF-8?q?=E5=8F=A3=E9=87=8D=E6=94=BE=E5=B7=B2=E6=8E=92c190=E5=90=8E?= 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/casc_eq_train.py | 64 ++++++++++++++++++++++++++++++++++++------------- 1 file changed, 47 insertions(+), 17 deletions(-) diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py index c1593f3..0cd46ee 100644 --- a/ep_run/casc_eq_train.py +++ b/ep_run/casc_eq_train.py @@ -68,11 +68,14 @@ ap.add_argument('--drift_adapt', type=float, default=0.0) # >0: adaptive drift ap.add_argument('--beta_cap_rho', type=float, default=0.0) # >0: LOOP-GAIN CAP on beta — if per-sweep residual # ratio rho^ exceeds this, bscale *= 0.8 (beta backs off # under the wall-2 ceiling); recovers x1.02 when rho^ low -ap.add_argument('--logit_knee', type=float, default=0.0) # >0: piecewise clamp of attn logits above the - # knee (slope knee_slope) — cuts the switching-regime - # Jacobian of runaway sharp heads (b8-class carriers) - # without touching healthy logits below the knee -ap.add_argument('--knee_slope', type=float, default=0.2) +ap.add_argument('--wsync', type=int, default=0) # >0: SYNCHRONOUS WEIGHT-STEP ACCEPTANCE — snapshot + # params+momentum before each opt.step; next step's + # nudged relax measures the new state through the SAME + # _legal gate (res/rho/drift, no new bounds); illegal -> + # roll back and re-apply the update at half scale + # (p <- (p+snap)/2), up to wsync halvings, then full + # revert + skip. The weight trajectory structurally + # cannot dwell past the ceiling. Zero new constants. ap.add_argument('--cap_floor', type=float, default=0.05) # hard bottom of the rho-cap; 0 = pure ceiling-tracking # (cap follows the measured ceiling all the way down; a # pinned bottom above the true ceiling = disguised wall-2) @@ -216,8 +219,6 @@ class Olmo2Attn(nn.Module): fr = torch.outer(torch.arange(T).float(), inv) self.register_buffer('rc', fr.cos(), persistent=False) self.register_buffer('rs', fr.sin(), persistent=False) - if args.logit_knee > 0: - self.register_buffer('cmask', torch.ones(T, T, dtype=torch.bool).tril(), persistent=False) def rope(self, x): x1, x2 = x[..., ::2], x[..., 1::2] c, s = self.rc[None, None], self.rs[None, None] @@ -229,14 +230,7 @@ class Olmo2Attn(nn.Module): 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) - if args.logit_knee > 0: - kn = args.logit_knee - lg = (q @ k.transpose(-2, -1)) * (self.hd ** -0.5) - lg = torch.where(lg > kn, kn + args.knee_slope * (lg - kn), lg) - lg = lg.masked_fill(~self.cmask[:T, :T], float('-inf')) - y = lg.softmax(-1) @ v - else: - y = F.scaled_dot_product_attention(q, k, v, is_causal=True) + 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): @@ -442,6 +436,12 @@ def dFdtheta(zs, x, y, beta): SIG0 = None BGEN = torch.Generator().manual_seed(args.seed + 990) # separate RNG: sign flips must not shift the data stream GOV = {'K': None, 'bscale': 1.0, 'gema': None, 'drift': 0.0, 'gn': 0.0, 'sig': 0.0} +WSNAP = {'p': None, 'o': None} +def _clone_state(sd): + if torch.is_tensor(sd): return sd.clone() + if isinstance(sd, dict): return {k: _clone_state(v) for k, v in sd.items()} + if isinstance(sd, list): return [_clone_state(v) for v in sd] + return sd def ep_step(x, y): """single-sided EP with a QUALITY-GOVERNED estimator: beta_t = beta0*bscale*sig0^2/sig^2, K = GOV['K'] fb rounds; guard = finiteness + drift + grad-norm sanity only.""" @@ -501,14 +501,38 @@ def ep_step(x, y): # SYNCHRONOUS acceptance (--beta_sync): judge THIS step by THIS step's own relax # telemetry — an out-of-window beta can be attempted but can never COMMIT. if (not math.isfinite(gd)) or (gd > dthr and not args.noguard): return False - if args.beta_sync > 0: + if args.beta_sync > 0 or args.wsync > 0: res_s = ddp_bcast_scalar(GOV.get('res', 0.0)) rho_s = ddp_bcast_scalar(GOV.get('rho') or 0.0) if res_s > 0.02 and rho_s > args.beta_cap_rho: return False return True if not _legal(gdrift): ok_retry = False - if args.beta_sync > 0 and not args.noguard: + if args.wsync > 0 and not args.noguard and WSNAP['p'] is not None: + # SYNCHRONOUS WEIGHT-STEP ACCEPTANCE: this state (= last opt.step's result) + # failed the gate -> the UPDATE was illegal. Halve it in weight space + # (p <- (p+snap)/2, delta implicit) and re-measure the same batch; after + # wsync halvings, revert fully (momentum too) and skip. Same idiom and same + # bounds as beta_sync — no new constants; the trajectory cannot dwell + # past the ceiling. + for _h in range(args.wsync + 1): + with torch.no_grad(): + if _h < args.wsync: + for p, s in zip(all_params, WSNAP['p']): p.copy_((p + s) * 0.5) + else: + for p, s in zip(all_params, WSNAP['p']): p.copy_(s) + opt.load_state_dict(WSNAP['o']) + GOV['skr'] = GOV.get('skr', 0) + 1 + z0, zs, ins, outs = free_states_graphed(x_in) + zs_free = [z.clone() for z in zs] + zp, last_outs = relax(z0, zs, ins, outs, y_in, +beta_t, GOV['K'], x_in, bmask=bmask) + with torch.no_grad(): + drift = _drift(zp, zs_free) + gdrift = ddp_max_scalar(drift) + if _legal(gdrift): + ok_retry = True + break + if not ok_retry and args.beta_sync > 0 and not args.noguard: for _h in range(args.beta_sync): # halve beta, retry the SAME batch beta_t = beta_t * 0.5 GOV['skr'] = GOV.get('skr', 0) + 1 @@ -769,6 +793,12 @@ for step in range(start_step, args.steps + 1): if id(p) in sel and g is not None: p.grad = g.detach().clone() torch.nn.utils.clip_grad_norm_(all_params, 1.0) + if args.wsync > 0: + # snapshot the KNOWN-LEGAL pre-step state (this step's relax passed _legal); + # next step's relax measures the post-step state and can roll back to here. + with torch.no_grad(): + WSNAP['p'] = [p.detach().clone() for p in all_params] + WSNAP['o'] = _clone_state(opt.state_dict()) opt.step(); sched.step(); opt.zero_grad(set_to_none=True) if args.qup_bits > 0: # STAGE-0 HW GATE: finite conductance levels. Snap every weight to an ABSOLUTE -- cgit v1.2.3