QK-norm: replace nn.MHA with explicit SDPA attn + --qk_norm (OLMo2-style, analog-friendly); cancel epoch, insert 8-run qk validation, stage roadmap TinyStories-epoch->FineWeb-Edu->OLMo2

3 files changed, 94 insertions, 12 deletions

diff --git a/docs/campaign/CASCADE_ABLATION_PLAN.md b/docs/campaign/CASCADE_ABLATION_PLAN.md
index 279b2fe..585ecc1 100644
--- a/docs/campaign/CASCADE_ABLATION_PLAN.md
+++ b/docs/campaign/CASCADE_ABLATION_PLAN.md
@@ -231,3 +231,39 @@ they carry the dynamics paper + the two-stage-recipe science; D1 takes over the
   demo (task #15). BP twin epoch DEFERRED (no free GPU; parity already sealed so it is nice-to-have).
 - Next: generation samples at checkpoints; BP-twin epoch when a GPU frees; then scale-up corpus
   decision (FineWeb-Edu vs OLMo2/Dolma) for the larger model.
+
+## ROADMAP PIVOT (2026-07-10 03:2x, user directive): QK-norm inserted; staged scale-up.
+**User: cancel the full epoch (done — killed epoch_ep_bf3e4); insert a QK-norm version after the
+current 3-seed; then stages TinyStories-full-epoch -> FineWeb-Edu -> OLMo2.**
+
+**Why QK-norm:** RMS-normalize q,k per head before the scores (OLMo2/Llama-style). It BOUNDS the
+attention logits, attacking the SAME root cause as the beta-floor (sig_tok growth -> logit blowup ->
+finite-beta SNR collapse) but structurally. Analog-friendly (my analysis): it's divisive
+normalization (mature analog/neuromorphic primitive), its Jacobian is symmetric (does NOT worsen the
+PAR/non-reciprocity wall), it's feedforward (no digital root-finder / no adjoint), and it REUSES the
+softmax current-normalization circuitry (reuse doctrine, no tapeout). Bonus analog wins: bounds the
+input range of the analog softmax exp device; reduces sig-growth so relaxation is more robust.
+Analog-preferred alternative to A/B in E-tier: tanh logit soft-cap (tanh is a native analog transfer
+function -- possibly cheaper than the norm's square-sum+divide).
+
+**Code:** nn.MultiheadAttention replaced by explicit CausalSelfAttn (SDPA-backed, fast) in BOTH
+trainers; `--qk_norm` flag (RMS-norm over head_dim w/ learnable per-dim gain). Smoke: EP+qk_norm
+cos=1.0000, 40.06M preserved, 2.49 it/s, SDPA works in the fb backward (fb is first-order, no
+double-backward needed). Also added `--cosine` (warmup->cosine to 0.1x lr) for the long runs.
+
+**QK-norm validation matrix (8 runs, L12 C512, 4000 steps, launched on GPU1):**
+  - qk_bp_s1/s2/s3  = BP + qk_norm (new reference with the new block)
+  - qk_ep_bf_s1/s2/s3 = EP + qk_norm + beta_floor 3e-4 (PARITY test vs qk_bp)
+  - qk_ep_nf_s1/s2  = EP + qk_norm, NO beta_floor (ANALOG test: does qk_norm ALONE hold cos, letting
+    us DROP the beta-floor? un-floored non-qk collapsed to cos 0.896 by step 4000 -- see RESULT 2).
+Decision: (1) qk_ep_bf ~ qk_bp => parity preserved with qk_norm. (2) if qk_ep_nf ALSO holds cos~1 and
+matches => qk_norm supersedes the beta-floor (fewer knobs, cleaner analog story). Watcher qk_watch.sh
+fires at the early analog read (nf step 2500) or all-done.
+
+**STAGED SCALE-UP (after qk_norm validates):**
+  Stage 1: TinyStories FULL EPOCH (58,800 steps, 361M tok) with the validated qk_norm recipe + cosine
+           -> the "neng kan" generation demo (task #15).
+  Stage 2: FineWeb-Edu (real corpus, 32-50k tokenizer, ~150-300M params) -- best small-LM quality.
+  Stage 3: OLMo2 / Dolma recipe -- fully-open reproducible baseline for the paper/collaborators.
+  EP scaling knobs carried forward: beta_floor (or qk_norm if it supersedes), possibly double-sided
+  nudge at larger scale (cancels O(beta) Taylor bias). $20k/run (Rain) ~ few-B tokens/run.
diff --git a/ep_run/casc_bp_train.py b/ep_run/casc_bp_train.py
index 7ee84e8..c056fed 100644
--- a/ep_run/casc_bp_train.py
+++ b/ep_run/casc_bp_train.py
@@ -19,6 +19,7 @@ ap.add_argument('--muon_lr', type=float, default=0.02)
 ap.add_argument('--tok_init', type=float, default=0.0)  # >0: init tok/pos std (GPT-standard 0.02)
 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)
 args = ap.parse_args()
 torch.manual_seed(args.seed)
 dev = 'cuda' if torch.cuda.is_available() else 'cpu'
@@ -33,22 +34,44 @@ def get_batch(split):
     y = torch.stack([torch.from_numpy(data[i + 1:i + 1 + args.T].astype(np.int64)) for i in ix])
     return x.to(dev), y.to(dev)
 
+class CausalSelfAttn(nn.Module):
+    """explicit MHA (SDPA-backed) so we can QK-norm q,k per head before the scores."""
+    def __init__(self, C, H, qk_norm=False):
+        super().__init__()
+        self.H, self.hd, self.qk_norm = H, C // H, qk_norm
+        self.qkv = nn.Linear(C, 3 * C)
+        self.proj = nn.Linear(C, C)
+        if qk_norm:
+            self.q_g = nn.Parameter(torch.ones(self.hd))
+            self.k_g = nn.Parameter(torch.ones(self.hd))
+    def forward(self, x):
+        B, T, C = x.shape
+        q, k, v = self.qkv(x).split(C, dim=2)
+        q = q.view(B, T, self.H, self.hd).transpose(1, 2)
+        k = k.view(B, T, self.H, self.hd).transpose(1, 2)
+        v = v.view(B, T, self.H, self.hd).transpose(1, 2)
+        if self.qk_norm:  # RMS-norm over head_dim (OLMo2-style), learnable per-dim gain
+            q = q * torch.rsqrt(q.pow(2).mean(-1, keepdim=True) + 1e-6) * self.q_g
+            k = k * torch.rsqrt(k.pow(2).mean(-1, keepdim=True) + 1e-6) * self.k_g
+        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        return self.proj(y.transpose(1, 2).contiguous().view(B, T, C))
+
 class Block(nn.Module):
-    def __init__(self, C, H):
+    def __init__(self, C, H, qk_norm=False):
         super().__init__()
         self.ln1, self.ln2 = nn.LayerNorm(C), nn.LayerNorm(C)
-        self.attn = nn.MultiheadAttention(C, H, batch_first=True)
+        self.attn = CausalSelfAttn(C, H, qk_norm)
         self.ff = nn.Sequential(nn.Linear(C, 4 * C), nn.GELU(), nn.Linear(4 * C, C))
-    def forward(self, z, mask):
-        h = self.ln1(z); a, _ = self.attn(h, h, h, attn_mask=mask, need_weights=False)
-        z = z + a; return z + self.ff(self.ln2(z))
+    def forward(self, z, mask=None):
+        z = z + self.attn(self.ln1(z))
+        return z + self.ff(self.ln2(z))
 
 tok = nn.Embedding(vocab, args.C).to(dev)
 pos = nn.Embedding(args.T, args.C).to(dev)
 if args.tok_init > 0:
     with torch.no_grad():
         tok.weight.normal_(0, args.tok_init); pos.weight.normal_(0, args.tok_init)
-blocks = nn.ModuleList([Block(args.C, args.H) for _ in range(args.L)]).to(dev)
+blocks = nn.ModuleList([Block(args.C, args.H, args.qk_norm) for _ in range(args.L)]).to(dev)
 mask = torch.triu(torch.full((args.T, args.T), float('-inf'), device=dev), 1)
 params = list(tok.parameters()) + list(pos.parameters()) + list(blocks.parameters())
 if args.opt == 'muon':
diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py
index ccaa204..f6ae71b 100644
--- a/ep_run/casc_eq_train.py
+++ b/ep_run/casc_eq_train.py
@@ -30,6 +30,7 @@ ap.add_argument('--beta_floor', type=float, default=0.0) # >0: floor beta_t (ant
 ap.add_argument('--beta_fixed', action='store_true')     # disable sig^2 schedule, hold beta_t = args.beta constant
 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)
 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)
@@ -47,22 +48,44 @@ def get_batch(split):
     y = torch.stack([torch.from_numpy(data[i + 1:i + 1 + args.T].astype(np.int64)) for i in ix])
     return x.to(dev), y.to(dev)
 
+class CausalSelfAttn(nn.Module):
+    """explicit MHA (SDPA-backed) so we can QK-norm q,k per head before the scores."""
+    def __init__(self, C, H, qk_norm=False):
+        super().__init__()
+        self.H, self.hd, self.qk_norm = H, C // H, qk_norm
+        self.qkv = nn.Linear(C, 3 * C)
+        self.proj = nn.Linear(C, C)
+        if qk_norm:
+            self.q_g = nn.Parameter(torch.ones(self.hd))
+            self.k_g = nn.Parameter(torch.ones(self.hd))
+    def forward(self, x):
+        B, T, C = x.shape
+        q, k, v = self.qkv(x).split(C, dim=2)
+        q = q.view(B, T, self.H, self.hd).transpose(1, 2)
+        k = k.view(B, T, self.H, self.hd).transpose(1, 2)
+        v = v.view(B, T, self.H, self.hd).transpose(1, 2)
+        if self.qk_norm:  # RMS-norm over head_dim (OLMo2-style), learnable per-dim gain
+            q = q * torch.rsqrt(q.pow(2).mean(-1, keepdim=True) + 1e-6) * self.q_g
+            k = k * torch.rsqrt(k.pow(2).mean(-1, keepdim=True) + 1e-6) * self.k_g
+        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        return self.proj(y.transpose(1, 2).contiguous().view(B, T, C))
+
 class Block(nn.Module):
-    def __init__(self, C, H):
+    def __init__(self, C, H, qk_norm=False):
         super().__init__()
         self.ln1, self.ln2 = nn.LayerNorm(C), nn.LayerNorm(C)
-        self.attn = nn.MultiheadAttention(C, H, batch_first=True)
+        self.attn = CausalSelfAttn(C, H, qk_norm)
         self.ff = nn.Sequential(nn.Linear(C, 4 * C), nn.GELU(), nn.Linear(4 * C, C))
-    def forward(self, z, mask):
-        h = self.ln1(z); a, _ = self.attn(h, h, h, attn_mask=mask, need_weights=False)
-        z = z + a; return z + self.ff(self.ln2(z))
+    def forward(self, z, mask=None):
+        z = z + self.attn(self.ln1(z))
+        return z + self.ff(self.ln2(z))
 
 tok = nn.Embedding(vocab, args.C).to(dev)
 pos = nn.Embedding(args.T, args.C).to(dev)
 if args.tok_init > 0:
     with torch.no_grad():
         tok.weight.normal_(0, args.tok_init); pos.weight.normal_(0, args.tok_init)
-blocks = nn.ModuleList([Block(args.C, args.H) for _ in range(args.L)]).to(dev)
+blocks = nn.ModuleList([Block(args.C, args.H, args.qk_norm) for _ in range(args.L)]).to(dev)
 if args.compile:
     try:
         for i in range(args.L): blocks[i] = torch.compile(blocks[i], mode='reduce-overhead')