From cb1c76762e6f8e028c8ed67682917356f935597f Mon Sep 17 00:00:00 2001 From: Yuren Hao Date: Fri, 10 Jul 2026 03:22:47 -0500 Subject: 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 --- docs/campaign/CASCADE_ABLATION_PLAN.md | 36 ++++++++++++++++++++++++++++++++++ ep_run/casc_bp_train.py | 35 +++++++++++++++++++++++++++------ ep_run/casc_eq_train.py | 35 +++++++++++++++++++++++++++------ 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') -- cgit v1.2.3