diff options
Diffstat (limited to 'ep_run')
| -rw-r--r-- | ep_run/agg_bench.log | 13 | ||||
| -rw-r--r-- | ep_run/agg_bench.py | 51 | ||||
| -rw-r--r-- | ep_run/bp_lm.py | 10 | ||||
| -rw-r--r-- | ep_run/lt_ep_train.py | 17 |
4 files changed, 89 insertions, 2 deletions
diff --git a/ep_run/agg_bench.log b/ep_run/agg_bench.log new file mode 100644 index 0000000..35cb799 --- /dev/null +++ b/ep_run/agg_bench.log @@ -0,0 +1,13 @@ +/home/yurenh2/miniconda3/lib/python3.13/site-packages/torch/autograd/graph.py:865: UserWarning: Attempting to run cuBLAS, but there was no current CUDA context! Attempting to set the primary context... (Triggered internally at /pytorch/aten/src/ATen/cuda/CublasHandlePool.cpp:330.) + return Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass + base: 4.50s/step (res 8.2e-03) + hf: 3.49s/step (res 8.2e-03) + sd: 4.07s/step (res 8.2e-03) + hf+sd: 3.06s/step (res 8.2e-03) + hf+sd+t80: 4.97s/step (res 8.2e-03) + hf+sd+t80+avg: 4.97s/step (res 8.2e-03) + cmp: 4.03s/step (res 8.2e-03) +/home/yurenh2/miniconda3/lib/python3.13/site-packages/torch/_inductor/compile_fx.py:321: UserWarning: TensorFloat32 tensor cores for float32 matrix multiplication available but not enabled. Consider setting `torch.set_float32_matmul_precision('high')` for better performance. + warnings.warn( + FULL(cmp_sdpa): 4.77s/step (res 8.2e-03) +AGG_BENCH_DONE diff --git a/ep_run/agg_bench.py b/ep_run/agg_bench.py new file mode 100644 index 0000000..d44b5a8 --- /dev/null +++ b/ep_run/agg_bench.py @@ -0,0 +1,51 @@ +"""THE AGGREGATE SPEED BENCH: full ep_step wall-time for every combination of the speed levers, on a +quiet local GPU, warm s2000 operator, B24 (production shape). Configs: + base : orig track, t2sel40, eager, manual attn (the historical default) + hf : +holofast + sd : +sdpa (eager relax) + hf+sd : the exact-math pack at t2sel40 + hf+sd+t80 : the accuracy pack (cos 0.89->0.94) + hf+sd+t80+avg : + holoavg (trend/plateau estimator) + cmp : --compile alone (manual attn in graph) + cmp(sdpa)+hf+t80+avg : FULL STACK (flash baked into compiled graph) +Reports median full-step time of 3 (after 1 warmup step each) + the step's res, as parity smoke.""" +import time, torch +import lt_ep_train as L + +torch.manual_seed(0) +blk = L.EQBlock(512, 16, 256, 256, c=1.0, attn_mode='thick'); blk.qknorm = True +ck = torch.load('runs/redx_traj/s2000.pt', map_location=L.dev) +with torch.no_grad(): + for p, w in zip(blk.allp, ck['allp']): + p.copy_(w.to(L.dev)) +blk.track = True +torch.manual_seed(42) +idx, y = L.get_batch('train', 24, 256) + +CFG = [ + ('base', dict(hf=False, sd=False, t2=40, avg=False, cmp=False)), + ('hf', dict(hf=True, sd=False, t2=40, avg=False, cmp=False)), + ('sd', dict(hf=False, sd=True, t2=40, avg=False, cmp=False)), + ('hf+sd', dict(hf=True, sd=True, t2=40, avg=False, cmp=False)), + ('hf+sd+t80', dict(hf=True, sd=True, t2=80, avg=False, cmp=False)), + ('hf+sd+t80+avg', dict(hf=True, sd=True, t2=80, avg=True, cmp=False)), + ('cmp', dict(hf=False, sd=False, t2=40, avg=False, cmp=True)), + ('FULL(cmp_sdpa)', dict(hf=True, sd=True, t2=80, avg=True, cmp=True)), +] + +for name, c in CFG: + blk.holofast, blk.sdpa, blk.holoavg = c['hf'], c['sd'], c['avg'] + blk._cstep = None + if c['cmp']: + _tf = blk.tforce_sdpa if c['sd'] else blk.tforce + blk._cstep = torch.compile(lambda z, xin, _tf=_tf: z + 0.1 * _tf(z, xin)) + ts = [] + for rep in range(4): # rep 0 = warmup (compile/JIT) + torch.cuda.synchronize(); t = time.time() + _, res = L.ep_step(blk, idx, y, 150, 20, 0.1, 0.02, jacreg=0.1, holo=2, hr=0.02, + t1max=300, res_est=1e-4, t2sel=c['t2'], corr_every=1, res_gate=0.0, resreg=0.2) + torch.cuda.synchronize(); ts.append(time.time() - t) + med = sorted(ts[1:])[1] + print(f"{name:>16}: {med:6.2f}s/step (res {res:.1e})", flush=True) +blk._cstep = None +print("AGG_BENCH_DONE", flush=True) diff --git a/ep_run/bp_lm.py b/ep_run/bp_lm.py index c01cd4f..1bd12f8 100644 --- a/ep_run/bp_lm.py +++ b/ep_run/bp_lm.py @@ -13,7 +13,7 @@ def fwd(blk, idx): h1 = F.layer_norm(x, (blk.C,), blk.ln1g, blk.ln1b) h2 = F.layer_norm(x, (blk.C,), blk.ln2g, blk.ln2b) h = x + blk.attn(h1) + (F.gelu(h2 @ blk.fc + blk.fcb, approximate='tanh') @ blk.pj + blk.pjb) - return h @ blk.Wh + return h @ (blk.tok.t() if getattr(blk, 'tie', False) else blk.Wh) def evaluate(blk, nb=8, B=32): @@ -36,6 +36,7 @@ def main(): ap.add_argument('--stdinit', action='store_true') # standard transformer init (EQBlock's is tuned for relaxation) ap.add_argument('--beta2', type=float, default=0.999) ap.add_argument('--sched', choices=['cos', 'const'], default='cos') + ap.add_argument('--tie', action='store_true') # tok/Wh weight tying (standard small-LM trick) ap.add_argument('--log', type=int, default=200) ap.add_argument('--ckpt', type=str, default='runs/bp_lm.pt') cfg = ap.parse_args() @@ -49,6 +50,13 @@ def main(): for W in (blk.WO, blk.pj): W.normal_(0, 0.02 / (2 ** 0.5)) blk.pos.normal_(0, 0.01) + if cfg.tie: # tie head to embedding: Wh := tok^T, single parameter + with torch.no_grad(): + blk.tok.copy_(0.5 * (blk.tok + blk.Wh.t())) + blk.Wh = None # fwd() will use tok.t() when tie is on + blk.tie = True + blk.allp = [p for p in blk.allp if p is not blk.Wh] + blk.allp = blk.block + [] # block already contains tok; Wh dropped opt = torch.optim.AdamW(blk.allp, lr=cfg.lr, weight_decay=cfg.wd, betas=(0.9, cfg.beta2)) if cfg.sched == 'cos': sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, cfg.steps, eta_min=cfg.lr * 0.05) diff --git a/ep_run/lt_ep_train.py b/ep_run/lt_ep_train.py index 2807be3..76bb9bb 100644 --- a/ep_run/lt_ep_train.py +++ b/ep_run/lt_ep_train.py @@ -81,6 +81,20 @@ class EQBlock: def Emem(self, z): return -(F.relu(z @ self.Wm) ** 2).sum() + def tforce_sdpa(self, z, xin): # tforce with fused flash attention (compiled fast path only; + h1 = F.layer_norm(z, (self.C,), self.ln1g, self.ln1b) # grad paths keep the manual attn) + h2 = F.layer_norm(z, (self.C,), self.ln2g, self.ln2b) + B = z.size(0) + q = (h1 @ self.WQ).view(B, self.T, self.H, self.dh).transpose(1, 2) + k = (h1 @ self.WK).view(B, self.T, self.H, self.dh).transpose(1, 2) + v = (h1 @ self.WV).view(B, self.T, self.H, self.dh).transpose(1, 2) + if getattr(self, 'qknorm', False): + q = q * torch.rsqrt(q.pow(2).mean(-1, keepdim=True) + 1e-6) + k = k * torch.rsqrt(k.pow(2).mean(-1, keepdim=True) + 1e-6) + att = F.scaled_dot_product_attention(q, k, v, is_causal=True).transpose(1, 2).reshape(B, self.T, self.C) @ self.WO + ff = F.gelu(h2 @ self.fc + self.fcb, approximate='tanh') @ self.pj + self.pjb + return -(z - xin) + att + ff - self.c * z + def tforce(self, z, xin): # pure thick force (no grad machinery) -> torch.compile h1 = F.layer_norm(z, (self.C,), self.ln1g, self.ln1b) h2 = F.layer_norm(z, (self.C,), self.ln2g, self.ln2b) @@ -523,7 +537,8 @@ def main(): blk._cstep = None if cfg.compile and cfg.attn_mode == 'thick': _ee = cfg.eps - blk._cstep = torch.compile(lambda z, xin: z + _ee * blk.tforce(z, xin)) + _tf = blk.tforce_sdpa if cfg.sdpa else blk.tforce # sdpa baked into the compiled graph, NO ambient flags + blk._cstep = torch.compile(lambda z, xin: z + _ee * _tf(z, xin)) mis = None if cfg.wmis > 0: # fixed fabrication mismatch (same devices all run) gm = torch.Generator().manual_seed(1234) |