diff options
Diffstat (limited to 'ep_run')
| -rw-r--r-- | ep_run/casc_eq_train.py | 12 |
1 files changed, 10 insertions, 2 deletions
diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py index 5a3e640..ad82b5e 100644 --- a/ep_run/casc_eq_train.py +++ b/ep_run/casc_eq_train.py @@ -40,6 +40,7 @@ ap.add_argument('--zloss', type=float, default=0.0) # z-loss coefficient on ap.add_argument('--kretry', type=int, default=0) # >0: on drift-reject, RETRY the batch once with this many fb rounds (diag B: K8 converges the marginal batches) instead of dropping it ap.add_argument('--bf_late', type=float, default=0.0) # >0: raise beta_floor to this value from step --bf_late_at (late-training SNR fix; dose-response 2026-07-10) ap.add_argument('--bf_late_at', type=int, default=25000) +ap.add_argument('--bf16', action='store_true') # cast model to bf16 (E-accumulation + tok_sigma stay fp32) — the x0.5 cost lever, GATE before production ap.add_argument('--dtop_every', type=int, default=1) # 1 = exact (DEFAULT, BP-parity); 2 = fast mode (~20% cheaper, ~4% CE tax at high lr) ap.add_argument('--gate_every', type=int, default=200) # in-training cos(EP,BP) telemetry; <=0 = fully BP-free (no bp_gate at all) ap.add_argument('--gate_govern', action='store_true') # let gate cos adjust K/bscale (default: observe-only => training control is BP-free) @@ -173,6 +174,13 @@ if args.resume: if _ck.get('lnf') is not None and not isinstance(ln_f, nn.Identity): ln_f.load_state_dict(_ck['lnf']) start_step = int(_ck.get('step', 0)) print(f'[resume] loaded {args.resume} at step {start_step}', flush=True) +if args.bf16: + for _m in (tok, pos, blocks): + _m.to(torch.bfloat16) + if not isinstance(ln_f, nn.Identity): ln_f.to(torch.bfloat16) + if args.untie: + with torch.no_grad(): W_out.data = W_out.data.to(torch.bfloat16) + print('[bf16] model cast to bfloat16 (E-accum + sigma stay fp32)', flush=True) if args.opt == 'muon': from muon import build_hybrid opt, sched = build_hybrid(blocks, all_params, args.lr, args.muon_lr, args.warmup, @@ -220,7 +228,7 @@ def free_states_graphed(x): @torch.no_grad() def tok_sigma(iters=8): """top singular value of tok.weight (power iteration on the raw matrix).""" - W = W_out if args.untie else tok.weight + W = (W_out if args.untie else tok.weight).float() v = torch.randn(W.shape[1], device=dev); v /= v.norm() sig = 1.0 for _ in range(iters): @@ -312,7 +320,7 @@ def ep_step(x, y): return free_ce, beta_t, GOV['K'], False GOV['drift'] = drift E = 0.0 - for z, o in zip(zp, last_outs): E = E + 0.5 * ((z.detach() - o) ** 2).sum() + 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) obj = E / (NBT * beta_t) + obj_loss(readout(zp[-1].detach()).reshape(-1, vocab), y.reshape(-1)) gs = torch.autograd.grad(obj, all_params, allow_unused=True) gn = 0.0 |