Self-contained LDT model definition

model.py

@@ -0,0 +1,159 @@
+"""LDT-10M — 10M-param LLaMA-style prose LM, trained from scratch.
+
+Self-contained model definition (no transformers dependency). Load the
+safetensors weights with `LDT.from_pretrained` (below) or the snippet in the
+README.
+
+Architecture (10,284,480 params, tied embeddings):
+  - vocab 12288 (gollem_eval byte-level BPE)
+  - d_model 320, n_layers 5, n_heads 5 (head_dim 64), SwiGLU FFN inter 896
+  - RMSNorm pre-norm, RoPE (base 10000), causal attention, ctx 512
+  - Standard LLaMA (no sliding window)
+"""
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def precompute_rope(dim, max_pos, base=10000.0):
+    freqs = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
+    t = torch.arange(max_pos).float()
+    angles = torch.outer(t, freqs)
+    return torch.polar(torch.ones_like(angles), angles)  # complex
+
+
+def apply_rope(x, freqs_cis, offset=0):
+    B, nh, S, hd = x.shape
+    x = x.view(B, nh, S, hd // 2, 2)
+    xr = x[..., 0].float()
+    xi = x[..., 1].float()
+    fc = freqs_cis[offset:offset + S].to(x.device)
+    fr, fi = fc.real, fc.imag
+    out_r = xr * fr - xi * fi
+    out_i = xr * fi + xi * fr
+    out = torch.stack([out_r, out_i], dim=-1).view(B, nh, S, hd)
+    return out.to(x.dtype)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim, eps=1e-5):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        norm = x.float().pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
+        return (x.float() * norm).to(x.dtype) * self.weight
+
+
+class Attention(nn.Module):
+    def __init__(self, d, n_heads):
+        super().__init__()
+        self.n_heads = n_heads
+        self.head_dim = d // n_heads
+        self.wq = nn.Linear(d, d, bias=False)
+        self.wk = nn.Linear(d, d, bias=False)
+        self.wv = nn.Linear(d, d, bias=False)
+        self.wo = nn.Linear(d, d, bias=False)
+
+    def forward(self, x, freqs_cis, offset=0):
+        B, S, _ = x.shape
+        q = self.wq(x).view(B, S, self.n_heads, self.head_dim).transpose(1, 2)
+        k = self.wk(x).view(B, S, self.n_heads, self.head_dim).transpose(1, 2)
+        v = self.wv(x).view(B, S, self.n_heads, self.head_dim).transpose(1, 2)
+        q = apply_rope(q, freqs_cis, offset)
+        k = apply_rope(k, freqs_cis, offset)
+        out = F.scaled_dot_product_attention(q, k, v)
+        out = out.transpose(1, 2).contiguous().view(B, S, -1)
+        return self.wo(out)
+
+
+class MLP(nn.Module):
+    def __init__(self, d, ff):
+        super().__init__()
+        self.w1 = nn.Linear(d, ff, bias=False)  # gate
+        self.w2 = nn.Linear(d, ff, bias=False)  # up
+        self.w3 = nn.Linear(ff, d, bias=False)  # down
+
+    def forward(self, x):
+        return self.w3(F.silu(self.w1(x)) * self.w2(x))
+
+
+class Block(nn.Module):
+    def __init__(self, d, n_heads, ff):
+        super().__init__()
+        self.ln1 = RMSNorm(d)
+        self.attn = Attention(d, n_heads)
+        self.ln2 = RMSNorm(d)
+        self.ffn = MLP(d, ff)
+
+    def forward(self, x, freqs_cis, offset=0):
+        x = x + self.attn(self.ln1(x), freqs_cis, offset)
+        x = x + self.ffn(self.ln2(x))
+        return x
+
+
+class LDT(nn.Module):
+    def __init__(self, vocab=12288, d=320, n_layers=5, n_heads=5, ff=896, ctx=512):
+        super().__init__()
+        self.vocab = vocab
+        self.ctx = ctx
+        self.d = d
+        self.tok = nn.Embedding(vocab, d)
+        self.blocks = nn.ModuleList([Block(d, n_heads, ff) for _ in range(n_layers)])
+        self.ln_f = RMSNorm(d)
+        self.head = nn.Linear(d, vocab, bias=False)
+        self.head.weight = self.tok.weight  # tied
+        self.freqs_cis = precompute_rope(d // n_heads, ctx)
+
+    @property
+    def tied_weights(self):
+        return ["lm_head"]
+
+    def forward(self, idx, targets=None):
+        B, S = idx.shape
+        h = self.tok(idx)
+        for b in self.blocks:
+            h = b(h, self.freqs_cis)
+        h = self.ln_f(h)
+        logits = self.head(h)
+        if targets is not None:
+            loss = F.cross_entropy(
+                logits[:, :-1].reshape(-1, logits.size(-1)),
+                targets[:, 1:].reshape(-1),
+                ignore_index=-1,
+            )
+            return loss
+        return logits
+
+    @torch.no_grad()
+    def generate(self, idx, max_new_tokens=128, temperature=0.8, top_k=40, seed=0):
+        g = torch.Generator(device=idx.device).manual_seed(seed)
+        for _ in range(max_new_tokens):
+            ctx_in = idx[:, -self.ctx:]
+            logits = self(ctx_in)[:, -1]
+            if temperature and temperature > 0:
+                logits = logits / temperature
+            if top_k:
+                v, _ = torch.topk(logits, top_k, dim=-1)
+                logits[logits < v[:, -1, None]] = float("-inf")
+            p = torch.softmax(logits, dim=-1)
+            nxt = torch.multinomial(p, 1, generator=g)
+            idx = torch.cat([idx, nxt], dim=1)
+        return idx
+
+    @classmethod
+    def from_pretrained(cls, path, device="cpu"):
+        """Load the safetensors weights (tied lm_head re-bound to tok.weight)."""
+        from safetensors.torch import load_file
+
+        m = cls()
+        sd = load_file(path, device=device)
+        # head.weight is tied to tok.weight and is NOT stored separately.
+        sd.pop("head.weight", None)
+        missing, unexpected = m.load_state_dict(sd, strict=False)
+        assert "head.weight" in missing, f"unexpected missing keys: {missing}"
+        assert not unexpected, f"unexpected keys: {unexpected}"
+        m.head.weight = m.tok.weight  # restore the tie
+        return m.to(device).eval()