From 9a8b2796ca12e4d4c24717485a635a301aa6d07f Mon Sep 17 00:00:00 2001 From: Yuren Hao Date: Mon, 6 Jul 2026 09:17:36 -0500 Subject: aggregate speed bench: speed tier hf+sd 1.47x; accuracy tier hf+sd+t80+avg (0.91x, cos 0.94, avg free); compile demoted MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Full-ep_step wall times on quiet A6000 (warm s2000, B24), res parity across all 8 configs. compile only 1.12x at this shape (historical 1.46x was a different workload split); FULL cmp_sdpa saves 4% over eager at t80 — not worth the guard complexity. tforce_sdpa added (flash baked into compiled graph, flag-free so grad paths never see SDPA). bp_lm --tie probe in flight. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_014FAPDWQ49M5Ye3NpTndTpn --- ep_run/bp_lm.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) (limited to 'ep_run/bp_lm.py') 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) -- cgit v1.2.3