File size: 4,834 Bytes
9014085 | 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 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 | """
GPT-S2.5 — 5.11M-param subword language model (from scratch).
Custom GPT-2-style architecture: GQA + SwiGLU + RoPE + RMSNorm, weight-tied head.
This is NOT a transformers-native architecture. It is a self-contained
nn.Module you can load directly:
import torch
from model import Model
m = Model()
sd = torch.load("best.pt", map_location="cpu")["model"] # or load safetensors
m.load_state_dict(sd); m.eval()
# generate (greedy):
with torch.no_grad():
x = torch.tensor([[1]])
for _ in range(60):
logits = m(x[:, -512:])[:, -1, :]
x = torch.cat([x, logits.argmax(-1, keepdim=True)], dim=1)
Architecture (matches config.json exactly):
vocab 8192 (BPE), n_embd 256, 4 layers, 8 q-heads / 2 kv-heads (GQA 4:1),
head_dim 32, SwiGLU FFN intermediate 768, RoPE (base 10000), RMSNorm (eps 1e-6),
pre-norm, no biases, weight-tied head (tok.weight reused as lm_head).
Total: 5,114,112 parameters.
"""
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
VOCAB = 8192
N_EMBD = 256
N_LAYERS = 4
N_Q_HEADS = 8
N_KV_HEADS = 2
HEAD_DIM = 32
INTER = 768
MAX_SEQ = 512
ROPE_BASE = 10000.0
Q_DIM = N_Q_HEADS * HEAD_DIM # 256
KV_DIM = N_KV_HEADS * HEAD_DIM # 64
assert Q_DIM == N_EMBD
assert N_Q_HEADS % N_KV_HEADS == 0
class RMSNorm(nn.Module):
def __init__(self, dim, eps=1e-6):
super().__init__()
self.eps = eps
self.weight = nn.Parameter(torch.ones(dim))
def forward(self, x):
return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight
def precompute_rope(head_dim, max_seq, base):
freqs = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim))
t = torch.arange(max_seq, dtype=torch.float)
ang = torch.outer(t, freqs)
return ang.cos(), ang.sin() # [T, head_dim/2]
def apply_rope(x, cos, sin):
# x: [B, H, T, D]
T = x.size(2)
cos = cos[:T].view(1, 1, T, -1)
sin = sin[:T].view(1, 1, T, -1)
x1, x2 = x[..., 0::2], x[..., 1::2]
out = torch.empty_like(x)
out[..., 0::2] = x1 * cos - x2 * sin
out[..., 1::2] = x2 * cos + x1 * sin
return out
class Attention(nn.Module):
def __init__(self):
super().__init__()
self.wq = nn.Linear(N_EMBD, Q_DIM, bias=False)
self.wk = nn.Linear(N_EMBD, KV_DIM, bias=False)
self.wv = nn.Linear(N_EMBD, KV_DIM, bias=False)
self.wo = nn.Linear(Q_DIM, N_EMBD, bias=False)
def forward(self, x, cos, sin):
B, T, _ = x.shape
q = self.wq(x).view(B, T, N_Q_HEADS, HEAD_DIM).transpose(1, 2)
k = self.wk(x).view(B, T, N_KV_HEADS, HEAD_DIM).transpose(1, 2)
v = self.wv(x).view(B, T, N_KV_HEADS, HEAD_DIM).transpose(1, 2)
q = apply_rope(q, cos, sin)
k = apply_rope(k, cos, sin)
# GQA: repeat kv heads to match q heads
rep = N_Q_HEADS // N_KV_HEADS
k = k.repeat_interleave(rep, dim=1)
v = v.repeat_interleave(rep, dim=1)
y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
y = y.transpose(1, 2).contiguous().view(B, T, Q_DIM)
return self.wo(y)
class MLP(nn.Module):
def __init__(self):
super().__init__()
self.wgate = nn.Linear(N_EMBD, INTER, bias=False)
self.wup = nn.Linear(N_EMBD, INTER, bias=False)
self.wdown = nn.Linear(INTER, N_EMBD, bias=False)
def forward(self, x):
return self.wdown(F.silu(self.wgate(x)) * self.wup(x))
class Block(nn.Module):
def __init__(self):
super().__init__()
self.ln1 = RMSNorm(N_EMBD)
self.attn = Attention()
self.ln2 = RMSNorm(N_EMBD)
self.mlp = MLP()
def forward(self, x, cos, sin):
x = x + self.attn(self.ln1(x), cos, sin)
x = x + self.mlp(self.ln2(x))
return x
class Model(nn.Module):
def __init__(self):
super().__init__()
self.tok = nn.Embedding(VOCAB, N_EMBD)
self.blocks = nn.ModuleList([Block() for _ in range(N_LAYERS)])
self.ln_f = RMSNorm(N_EMBD)
self.cos, self.sin = precompute_rope(HEAD_DIM, MAX_SEQ, ROPE_BASE)
def forward(self, idx, targets=None):
B, T = idx.shape
h = self.tok(idx)
cos, sin = self.cos, self.sin
for blk in self.blocks:
h = blk(h, cos, sin)
h = self.ln_f(h)
logits = F.linear(h, self.tok.weight) # weight-tied head
if targets is not None:
loss = F.cross_entropy(logits.view(-1, VOCAB), targets.view(-1))
return loss
return logits
if __name__ == "__main__":
m = Model()
total = sum(p.numel() for p in m.parameters())
print(f"params = {total} (expected 5114112)")
assert total == 5114112, "param count mismatch"
print("OK")
|