diff options
Diffstat (limited to 'ep_run/casc_eq_train.py')
| -rw-r--r-- | ep_run/casc_eq_train.py | 64 |
1 files 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 |