scaffold: model (OLMo2-ish + ZBP partition), trainer (DDP/config), data shards, bench

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01GkgLsACEF6CCP7EUfA5fZe

1 files changed, 96 insertions, 0 deletions

diff --git a/src/zbp_scaling/model.py b/src/zbp_scaling/model.py
new file mode 100644
index 0000000..61ea4de
--- /dev/null
+++ b/src/zbp_scaling/model.py
@@ -0,0 +1,96 @@
+"""OLMo2-style decoder LM with the ZBP physical/digital partition (part-1 'mixscoreqk,ffn' convention):
+physical = score-space attention core (per head x row) + whole SwiGLU FFN branch (per token);
+digital = embeddings, RMSNorms outside blocks, q/k/v/proj linears, score product, readout head."""
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .zbp import ZBPBlock, ZBPConfig
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, d, eps=1e-5):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(d))
+        self.eps = eps
+
+    def forward(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight
+
+
+class _Core(nn.Module):
+    """softmax(QK^T/sqrt(dh)) V from concatenated qkv [B,T,3d] -> [B,T,d]; explicit (forward-AD friendly)."""
+
+    def __init__(self, d, heads):
+        super().__init__()
+        self.d, self.h = d, heads
+
+    def forward(self, qkv):
+        B, T, _ = qkv.shape
+        q, k, v = qkv.split(self.d, dim=-1)
+        dh = self.d // self.h
+        q = q.view(B, T, self.h, dh).transpose(1, 2)
+        k = k.view(B, T, self.h, dh).transpose(1, 2)
+        v = v.view(B, T, self.h, dh).transpose(1, 2)
+        att = (q @ k.transpose(-2, -1)) / math.sqrt(dh)
+        mask = torch.ones(T, T, dtype=torch.bool, device=qkv.device).tril()
+        att = att.masked_fill(~mask, float("-inf")).softmax(-1)
+        return (att @ v).transpose(1, 2).reshape(B, T, self.d)
+
+
+class _SwiGLU(nn.Module):
+    """RMS -> (W1 x) * SiLU(W3 x) -> W2, one per-token physical branch (identity skip outside)."""
+
+    def __init__(self, d, hidden, scale):
+        super().__init__()
+        self.norm = RMSNorm(d)
+        self.w1 = nn.Linear(d, hidden, bias=False)
+        self.w3 = nn.Linear(d, hidden, bias=False)
+        self.w2 = nn.Linear(hidden, d, bias=False)
+        with torch.no_grad():
+            self.w2.weight.mul_(scale)
+
+    def forward(self, x):
+        h = self.norm(x)
+        return self.w2(self.w1(h) * F.silu(self.w3(h)))
+
+
+class Block(nn.Module):
+    def __init__(self, d, heads, hidden, cfg, bp, scale, li):
+        super().__init__()
+        self.norm_a = RMSNorm(d)
+        self.qkv = nn.Linear(d, 3 * d, bias=False)
+        core = ZBPBlock(_Core(d, heads), cfg, batch_dims=1, name=f"L{li}.core")
+        core.score_probe = {"d": d, "heads": heads, "dk": d, "dv": d, "window": None, "qk": True}
+        self.core = core
+        self.proj = nn.Linear(d, d, bias=False)
+        with torch.no_grad():
+            self.proj.weight.mul_(scale)
+        self.ffn = ZBPBlock(_SwiGLU(d, hidden, scale), cfg, skip="identity", batch_dims=2, name=f"L{li}.ffn")
+
+    def forward(self, x):
+        x = x + self.proj(self.core(self.qkv(self.norm_a(x))))
+        return self.ffn(x)
+
+
+class ScalingLM(nn.Module):
+    def __init__(self, vocab, d, layers, heads, seq_len, ffn_mult="8/3", cfg=None):
+        super().__init__()
+        cfg = cfg or ZBPConfig(mode="bp")
+        bp = cfg.replace(mode="bp")
+        mult = eval(str(ffn_mult)) if isinstance(ffn_mult, str) else ffn_mult
+        hidden = int(round(mult * d / 64) * 64)
+        scale = 1 / math.sqrt(2 * layers)
+        self.tok = nn.Embedding(vocab, d)
+        self.pos = nn.Embedding(seq_len, d)
+        nn.init.normal_(self.tok.weight, std=0.02); nn.init.normal_(self.pos.weight, std=0.02)
+        self.blocks = nn.Sequential(*[Block(d, heads, hidden, cfg, bp, scale, i) for i in range(layers)])
+        self.norm_f = RMSNorm(d)
+        self.head = nn.Linear(d, vocab, bias=False)
+        self.seq_len = seq_len
+
+    def forward(self, idx):
+        B, T = idx.shape
+        x = self.tok(idx) + self.pos(torch.arange(T, device=idx.device))[None]
+        return self.head(self.norm_f(self.blocks(x)))