| |
| """Custom loader for Compactbot/swordies-22m. |
| |
| Swordies-22M is a from-scratch BPE GPT (NOT a transformers model). This file |
| reconstructs the architecture from config.json and loads model.safetensors. |
| |
| Usage: |
| from load_model import load_model |
| model = load_model("model.safetensors") |
| logits = model(token_ids) # token_ids: int64 [B, T], vocab 8192 |
| probs = torch.softmax(logits, -1) |
| |
| Tensor layout (57 tensors, F32, weight-tied): |
| tok.weight [8192, 448] (also the lm_head, tied) |
| pos.weight [512, 448] |
| blocks.{0..8}.ln1.w [448] |
| blocks.{0..8}.ln2.w [448] |
| blocks.{0..8}.qkv.weight [448, 1344] (fused q|k|v, no bias) |
| blocks.{0..8}.proj.weight [448, 448] |
| blocks.{0..8}.fc1.weight [448, 1408] |
| blocks.{0..8}.fc2.weight [1408, 448] |
| ln_f.w [448] |
| """ |
| import json, os |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from safetensors import safe_open |
|
|
| VOCAB = 8192 |
| D = 448 |
| L = 9 |
| H = 7 |
| FFN = 1408 |
| SEQ = 512 |
|
|
|
|
| class RMSNorm(nn.Module): |
| def __init__(self, d): |
| super().__init__() |
| self.w = nn.Parameter(torch.ones(d)) |
|
|
| def forward(self, x): |
| return self.w * x * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + 1e-6) |
|
|
|
|
| class Block(nn.Module): |
| def __init__(self, d, h): |
| super().__init__() |
| self.ln1 = RMSNorm(d) |
| self.ln2 = RMSNorm(d) |
| self.qkv = nn.Linear(d, 3 * d, bias=False) |
| self.proj = nn.Linear(d, d, bias=False) |
| self.fc1 = nn.Linear(d, FFN, bias=False) |
| self.fc2 = nn.Linear(FFN, d, bias=False) |
| self.h, self.d = h, d |
|
|
| def forward(self, x): |
| B, T, Dd = x.shape |
| h = self.ln1(x) |
| qkv = self.qkv(h).view(B, T, 3, self.h, Dd // self.h).transpose(2, 1) |
| q, k, v = qkv[:, 0], qkv[:, 1], qkv[:, 2] |
| q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2) |
| att = F.scaled_dot_product_attention(q, k, v, is_causal=True) |
| att = att.transpose(1, 2).reshape(B, T, Dd) |
| x = x + self.proj(att) |
| x = x + self.fc2(F.gelu(self.fc1(self.ln2(x)))) |
| return x |
|
|
|
|
| class SwordiesGPT(nn.Module): |
| def __init__(self): |
| super().__init__() |
| self.tok = nn.Embedding(VOCAB, D) |
| self.pos = nn.Embedding(SEQ, D) |
| self.blocks = nn.ModuleList([Block(D, H) for _ in range(L)]) |
| self.ln_f = RMSNorm(D) |
|
|
| def forward(self, idx, targets=None): |
| B, T = idx.shape |
| x = self.tok(idx) + self.pos(torch.arange(T, device=idx.device)) |
| for b in self.blocks: |
| x = b(x) |
| x = self.ln_f(x) |
| logits = x @ self.tok.weight.t() |
| if targets is not None: |
| return F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1)) |
| return logits |
|
|
|
|
| def load_model(path, device="cpu"): |
| """Load model.safetensors into a SwordiesGPT and return it (eval mode).""" |
| model = SwordiesGPT().to(device) |
| with safe_open(path, framework="pt") as f: |
| state = {k: f.get_tensor(k) for k in f.keys()} |
| missing, unexpected = model.load_state_dict(state, strict=True) |
| model.eval() |
| n = sum(p.numel() for p in model.parameters()) |
| assert n == 22487360, f"param mismatch: {n}" |
| return model |
|
|
|
|
| if __name__ == "__main__": |
| here = os.path.dirname(os.path.abspath(__file__)) |
| m = load_model(os.path.join(here, "model.safetensors")) |
| x = torch.randint(0, VOCAB, (1, 64), dtype=torch.int64) |
| with torch.no_grad(): |
| lg = m(x) |
| print("loaded OK; params =", sum(p.numel() for p in m.parameters())) |
| print("logits shape", tuple(lg.shape), "finite:", bool(torch.isfinite(lg).all())) |
