File size: 5,620 Bytes
a198f75 | 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 155 156 157 158 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() |