aggregate speed bench: speed tier hf+sd 1.47x; accuracy tier hf+sd+t80+avg (0.91x, cos 0.94, avg free); compile demoted

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 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014FAPDWQ49M5Ye3NpTndTpn

1 files changed, 9 insertions, 1 deletions

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)