1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
|
"""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)
# the score-space estimator materialises chunk*B*H*T^2 score tensors: give the core its own small chunk
core_cfg = cfg if cfg.mode == "bp" else cfg.replace(probe_chunk=min(cfg.probe_chunk or 4, 4))
core = ZBPBlock(_Core(d, heads), core_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)))
|
