diff options
Diffstat (limited to 'ep_run/casc_eq_train.py')
| -rw-r--r-- | ep_run/casc_eq_train.py | 20 |
1 files changed, 18 insertions, 2 deletions
diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py index 5214311..7b9780c 100644 --- a/ep_run/casc_eq_train.py +++ b/ep_run/casc_eq_train.py @@ -51,6 +51,9 @@ ap.add_argument('--data', default='tinystories_bpe') # dataset dir under ep_ ap.add_argument('--sync_check', type=int, default=500) # DDP: verify bitwise param sync every N steps (0=off) ap.add_argument('--ddp_backend', default='nccl', choices=['nccl', 'gloo']) # gloo = correctness tests on shared GPUs ap.add_argument('--ddp_grad_test', action='store_true') # one-step grad equivalence test vs single-GPU big batch, then exit +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('--relax_tol', type=float, default=0.0) # >0: ADAPTIVE relax — sweep until rel. state change < tol # (or --kmax), geta backtracks x0.6 on residual GROWTH (rho>=1 # signal), then one final graphed round. 0 = legacy fixed-K. @@ -341,10 +344,13 @@ def relax(z0, zs, ins, outs, y, beta, K, x): ins, outs = n_ins, n_outs return rnum / max(rden, 1e-9) - if not adaptive: # legacy fixed-K path (bit-identical to before) + if not adaptive: # legacy fixed-K path (bit-identical update semantics) + rlist = [] for k in range(K): forces(k % args.dtop_every == 0) - rebuild(k + 1 == K) + rlist.append(rebuild(k + 1 == K)) + if len(rlist) >= 2 and rlist[-2] > 1e-12: + GOV['rho'] = rlist[-1] / rlist[-2] # per-sweep contraction ratio = live loop-gain meter GOV['kuse'] = K return zs, outs @@ -353,6 +359,8 @@ def relax(z0, zs, ins, outs, y, beta, K, x): forces(k % args.dtop_every == 0) res = rebuild(False) k += 1 + if prev_res is not None and prev_res > 1e-12: + GOV['rho'] = res / prev_res if prev_res is not None and res > prev_res and res > args.relax_tol: geta_l = max(0.2, geta_l * 0.6) # residual GREW: local rho>=1 -> damp harder prev_res = res @@ -396,6 +404,8 @@ def ep_step(x, y): fl = args.beta_floor if args.bf_late > 0.0 and GOV.get('step', 0) >= args.bf_late_at: fl = args.bf_late if fl > 0.0: beta_t = max(beta_t, fl) + beta_t = beta_t * GOV.get('cap', 1.0) # wall-2 loop-gain cap OVERRIDES the floor (the ceiling + # can sit below the floor near the wall; survival first) if args.bsign_rand and torch.rand((), generator=BGEN).item() < 0.5: beta_t = -beta_t z0, zs, ins, outs = free_states_graphed(x) zs_free = [z.clone() for z in zs] @@ -422,6 +432,12 @@ def ep_step(x, y): for p in all_params: p.grad = None return free_ce, beta_t, GOV.get('kuse', GOV['K']), False GOV['drift'] = gdrift + if args.beta_cap_rho > 0 and GOV.get('rho') is not None: + rho_g = ddp_bcast_scalar(GOV['rho']) # rank0's meter rules (identical control on all ranks) + if rho_g > args.beta_cap_rho: + GOV['cap'] = max(GOV.get('cap', 1.0) * 0.8, 0.02) # back beta off under the wall-2 ceiling + elif rho_g < 0.5 * args.beta_cap_rho: + GOV['cap'] = min(GOV.get('cap', 1.0) * 1.02, 1.0) # slow recovery toward schedule E = 0.0 for z, o in zip(zp, last_outs): E = E + 0.5 * ((z.detach().float() - o.float()) ** 2).sum() # fp32 accumulation (bf16-safe; no-op in fp32) if args.est == 'single': |