diff options
Diffstat (limited to 'ep_run/casc_eq_train.py')
| -rw-r--r-- | ep_run/casc_eq_train.py | 19 |
1 files changed, 16 insertions, 3 deletions
diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py index c77b92c..86a8e06 100644 --- a/ep_run/casc_eq_train.py +++ b/ep_run/casc_eq_train.py @@ -29,6 +29,10 @@ ap.add_argument('--compile', action='store_true') # torch.compile each blo ap.add_argument('--sig_every', type=int, default=25) # tok-sigma refresh interval (amortized) ap.add_argument('--beta_floor', type=float, default=0.0) # >0: floor beta_t (anti finite-beta SNR collapse at depth) ap.add_argument('--beta_fixed', action='store_true') # disable sig^2 schedule, hold beta_t = args.beta constant +ap.add_argument('--beta_cos_min', type=float, default=0.0) # >0: cosine-descend beta from --beta to this over + # --steps (tracks the sinking ceiling by progress; + # bypasses sigma-scaling AND floor). Set --beta to the + # early value (below early ceiling ~0.12). ap.add_argument('--cosine', action='store_true') # warmup then cosine decay to lr_min_ratio*lr over --steps (long runs) ap.add_argument('--lr_min_ratio', type=float, default=0.1) ap.add_argument('--qk_norm', action='store_true') # RMS-norm q,k per head before scores (OLMo2-style; bounds logits, analog-friendly) @@ -472,9 +476,18 @@ def ep_step(x, y): if SIG0 is None: SIG0 = args.sig0 if args.sig0 > 0 else sig beta_t = args.beta * GOV['bscale'] * (SIG0 * SIG0) / max(sig * sig, 1e-9) if args.beta_fixed: beta_t = args.beta * GOV['bscale'] - 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) + if args.beta_cos_min > 0: + # SCHEDULED beta descent (07-20): the ceiling sinks ~1/sigma^2 as training sharpens the + # model (sigma 242->473); no fixed beta stays under it (endgame skip-stall). Descend beta + # by progress like LR — early large (below the high early ceiling), late small (below the + # sunk endgame ceiling). Skips sigma-scaling AND floor entirely. Conservative = safe: CE + # is flat across the in-corridor band, so undershoot costs nothing, overshoot skips. + prog = min(GOV.get('step', 0) / max(args.steps, 1), 1.0) + beta_t = args.beta_cos_min + 0.5 * (args.beta - args.beta_cos_min) * (1 + math.cos(math.pi * prog)) + else: + 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) if args.beta_ride > 1.0: # ride-v2(a): floor jumps must not compose with a pre-charged cap — rescale cap so the # EFFECTIVE beta is continuous across any floor change (the 0.09-at-20k bug, RESULT 37) |