Publish Prizma — mirror of github.com/nazmiefearmutcu/Prizma (PRISM-Seq §4 bar + continual-learning Prizma)
5066d15 verified | """ | |
| A clean, properly-implemented decoder-only Transformer baseline (NOT a strawman). | |
| Causal multi-head self-attention via scaled_dot_product_attention, RMSNorm pre-norm, | |
| SwiGLU MLP, learned absolute positions. This is the honest reference Prizma-Seq must match | |
| at equal parameter count. | |
| """ | |
| from __future__ import annotations | |
| from dataclasses import dataclass | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| class TFConfig: | |
| vocab: int = 64 | |
| d_model: int = 128 | |
| n_layers: int = 2 | |
| n_heads: int = 4 | |
| d_ff: int = None # default 8/3 * d_model rounded (SwiGLU keeps ~4x FLOPs) | |
| max_len: int = 1024 | |
| dropout: float = 0.0 | |
| tie_embeddings: bool = True | |
| rope: bool = True # use RoPE (parameter-free) instead of a learned pos-embedding | |
| def __post_init__(self): | |
| if self.d_ff is None: | |
| self.d_ff = int(round(8 / 3 * self.d_model / 8) * 8) | |
| def _tf_rope_cache(T, hd, device, dtype): | |
| inv = 1.0 / (10000 ** (torch.arange(0, hd, 2, device=device, dtype=torch.float32) / hd)) | |
| ang = torch.outer(torch.arange(T, device=device, dtype=torch.float32), inv) | |
| return torch.cos(ang).to(dtype), torch.sin(ang).to(dtype) | |
| def _tf_apply_rope(x, cos, sin): | |
| x1, x2 = x[..., 0::2], x[..., 1::2] | |
| out = torch.empty_like(x) | |
| out[..., 0::2] = x1 * cos - x2 * sin | |
| out[..., 1::2] = x1 * sin + x2 * cos | |
| return out | |
| class RMSNorm(nn.Module): | |
| def __init__(self, d, eps=1e-5): | |
| super().__init__() | |
| self.w = nn.Parameter(torch.ones(d)) | |
| self.eps = eps | |
| def forward(self, x): | |
| n = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) | |
| return n * self.w | |
| class CausalSelfAttention(nn.Module): | |
| def __init__(self, cfg: TFConfig): | |
| super().__init__() | |
| assert cfg.d_model % cfg.n_heads == 0 | |
| self.nh = cfg.n_heads | |
| self.hd = cfg.d_model // cfg.n_heads | |
| self.qkv = nn.Linear(cfg.d_model, 3 * cfg.d_model, bias=False) | |
| self.proj = nn.Linear(cfg.d_model, cfg.d_model, bias=False) | |
| self.drop = cfg.dropout | |
| self.rope = cfg.rope | |
| 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.nh, self.hd).transpose(1, 2) | |
| k = k.view(B, T, self.nh, self.hd).transpose(1, 2) | |
| v = v.view(B, T, self.nh, self.hd).transpose(1, 2) | |
| if self.rope: | |
| cos, sin = _tf_rope_cache(T, self.hd, x.device, x.dtype) | |
| q = _tf_apply_rope(q, cos, sin) | |
| k = _tf_apply_rope(k, cos, sin) | |
| o = F.scaled_dot_product_attention(q, k, v, is_causal=True, | |
| dropout_p=self.drop if self.training else 0.0) | |
| o = o.transpose(1, 2).contiguous().view(B, T, C) | |
| return self.proj(o) | |
| def step(self, x_t, cache): | |
| """x_t:[B,1,C]; cache=(K,V) each [B,nh,t,hd] or None. KV-cached O(t) decode.""" | |
| B, _, C = x_t.shape | |
| q, k, v = self.qkv(x_t).split(C, dim=2) | |
| q = q.view(B, 1, self.nh, self.hd).transpose(1, 2) | |
| k = k.view(B, 1, self.nh, self.hd).transpose(1, 2) | |
| v = v.view(B, 1, self.nh, self.hd).transpose(1, 2) | |
| t = 0 if cache is None else cache[0].shape[2] | |
| if self.rope: | |
| cos, sin = _tf_rope_cache(t + 1, self.hd, x_t.device, x_t.dtype) | |
| q = _tf_apply_rope(q, cos[t:t + 1], sin[t:t + 1]) | |
| k = _tf_apply_rope(k, cos[t:t + 1], sin[t:t + 1]) | |
| if cache is not None: | |
| k = torch.cat([cache[0], k], dim=2) | |
| v = torch.cat([cache[1], v], dim=2) | |
| o = F.scaled_dot_product_attention(q, k, v, is_causal=False) # q attends all cached (<=t) | |
| o = o.transpose(1, 2).reshape(B, 1, C) | |
| return self.proj(o), (k, v) | |
| class SwiGLU(nn.Module): | |
| def __init__(self, cfg: TFConfig): | |
| super().__init__() | |
| self.w1 = nn.Linear(cfg.d_model, cfg.d_ff, bias=False) | |
| self.w2 = nn.Linear(cfg.d_model, cfg.d_ff, bias=False) | |
| self.wo = nn.Linear(cfg.d_ff, cfg.d_model, bias=False) | |
| def forward(self, x): | |
| return self.wo(F.silu(self.w1(x)) * self.w2(x)) | |
| class Block(nn.Module): | |
| def __init__(self, cfg: TFConfig): | |
| super().__init__() | |
| self.n1 = RMSNorm(cfg.d_model) | |
| self.attn = CausalSelfAttention(cfg) | |
| self.n2 = RMSNorm(cfg.d_model) | |
| self.mlp = SwiGLU(cfg) | |
| def forward(self, x): | |
| x = x + self.attn(self.n1(x)) | |
| x = x + self.mlp(self.n2(x)) | |
| return x | |
| def step(self, x_t, cache): | |
| o, cache = self.attn.step(self.n1(x_t), cache) | |
| x_t = x_t + o | |
| return x_t + self.mlp(self.n2(x_t)), cache | |
| class Transformer(nn.Module): | |
| def __init__(self, cfg: TFConfig): | |
| super().__init__() | |
| self.cfg = cfg | |
| self.tok = nn.Embedding(cfg.vocab, cfg.d_model) | |
| self.pos = None if cfg.rope else nn.Embedding(cfg.max_len, cfg.d_model) | |
| self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.n_layers)]) | |
| self.nf = RMSNorm(cfg.d_model) | |
| self.head = nn.Linear(cfg.d_model, cfg.vocab, bias=False) | |
| if cfg.tie_embeddings: | |
| self.head.weight = self.tok.weight | |
| self.apply(self._init) | |
| def _init(self, m): | |
| if isinstance(m, nn.Linear): | |
| nn.init.normal_(m.weight, std=0.02) | |
| if m.bias is not None: | |
| nn.init.zeros_(m.bias) | |
| elif isinstance(m, nn.Embedding): | |
| nn.init.normal_(m.weight, std=0.02) | |
| def forward(self, idx): | |
| B, T = idx.shape | |
| x = self.tok(idx) | |
| if self.pos is not None: | |
| x = x + self.pos(torch.arange(T, device=idx.device))[None] | |
| for blk in self.blocks: | |
| x = blk(x) | |
| return self.head(self.nf(x)) | |
| def init_state(self, batch, device): | |
| return [None for _ in self.blocks] # per-layer KV cache (grows with t) | |
| def step(self, tok, state): | |
| """tok:[B,1] -> (logits[B,1,V], new_state). KV-cached: O(t) compute, O(t) memory at step t.""" | |
| x = self.tok(tok) | |
| if self.pos is not None: | |
| p = 0 if state[0] is None else state[0][0].shape[2] | |
| x = x + self.pos(torch.tensor([p], device=tok.device))[None] | |
| new = [] | |
| for blk, c in zip(self.blocks, state): | |
| x, c2 = blk.step(x, c) | |
| new.append(c2) | |
| return self.head(self.nf(x)), new | |
| if __name__ == "__main__": | |
| from common import param_count | |
| cfg = TFConfig(vocab=64, d_model=128, n_layers=2, n_heads=4) | |
| m = Transformer(cfg) | |
| x = torch.randint(0, 64, (2, 32)) | |
| print("logits", tuple(m(x).shape), "params", param_count(m)) | |