| """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) |
|
|
|
|
| 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) |
| self.w2 = nn.Linear(d, ff, bias=False) |
| self.w3 = nn.Linear(ff, d, bias=False) |
|
|
| 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 |
| 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) |
| |
| 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 |
| return m.to(device).eval() |
