| """Decoder-only transformer for the IGLA-Coder pilot. |
| |
| One architecture, two arms that differ only in the linear layers inside the |
| blocks: |
| * fp -- ordinary bf16/fp32 linears; |
| * ternary -- BitNet b1.58 linears: parameter-free RMSNorm on the input, |
| 8-bit absmax activations per token and ternary {-1, 0, +1} |
| absmean weights, both through straight-through estimators. |
| Embeddings, the tied output head and the norms stay full precision in both arms. |
| |
| Layout follows the IGLA-Coder spec (t27 specs/igla/coder/arch.t27): RMSNorm, |
| RoPE, grouped-query attention, SwiGLU, tied embeddings (t27 #1034 P1.1). |
| """ |
| from dataclasses import dataclass, asdict |
| import math |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
|
|
| @dataclass |
| class Config: |
| vocab_size: int = 32000 |
| n_layer: int = 12 |
| n_head: int = 12 |
| n_kv_head: int = 4 |
| d_model: int = 768 |
| d_ff: int = 2048 |
| seq_len: int = 2048 |
| rope_theta: float = 10000.0 |
| ternary: bool = False |
| |
| wgrad_ternary: bool = False |
| attn_lns: float = 0.0 |
| lin8: str = "" |
| lin8_scale: bool = False |
|
|
|
|
| |
| SIZES = { |
| "13m": dict(n_layer=6, n_head=4, n_kv_head=2, d_model=256, d_ff=704), |
| "25m": dict(n_layer=8, n_head=6, n_kv_head=2, d_model=384, d_ff=1024), |
| "52m": dict(n_layer=12, n_head=8, n_kv_head=4, d_model=512, d_ff=1408), |
| "100m": dict(n_layer=12, n_head=12, n_kv_head=4, d_model=768, d_ff=2048), |
| |
| |
| "tc": dict(n_layer=6, n_head=5, n_kv_head=1, d_model=320, d_ff=864), |
| } |
|
|
|
|
| def rms_norm(x, eps=1e-6): |
| return F.rms_norm(x, (x.size(-1),), eps=eps) |
|
|
|
|
| class _TernaryWgrad(torch.autograd.Function): |
| """y = x_q W^T as usual; in the backward pass the weight gradient is taken |
| against activations rounded to three levels per token, so dW = g^T x_t is a |
| signed sum of gradient rows (the per-token scale is O(d) work, not O(d^2)).""" |
|
|
| @staticmethod |
| def forward(ctx, xq, wq, xt): |
| ctx.save_for_backward(wq, xt) |
| ctx.dtypes = (xq.dtype, wq.dtype) |
| return F.linear(xq, wq) |
|
|
| @staticmethod |
| def backward(ctx, gy): |
| wq, xt = ctx.saved_tensors |
| gx = (gy @ wq.to(gy.dtype)).to(ctx.dtypes[0]) |
| gw = gy.reshape(-1, gy.size(-1)).t() @ xt.reshape(-1, xt.size(-1)).to(gy.dtype) |
| return gx, gw.to(ctx.dtypes[1]), None |
|
|
|
|
| class BitLinear(nn.Linear): |
| """BitNet b1.58 linear layer (Ma et al. 2024, training-tips reference code).""" |
| wgrad_ternary = False |
|
|
| def forward(self, x): |
| x = rms_norm(x) |
| s = 127.0 / x.abs().amax(dim=-1, keepdim=True).clamp(min=1e-5) |
| xq = x + ((x * s).round().clamp(-128, 127) / s - x).detach() |
| w = self.weight |
| g = w.abs().mean().clamp(min=1e-5) |
| w = w + ((w / g).round().clamp(-1, 1) * g - w).detach() |
| if self.wgrad_ternary: |
| xd = x.detach() |
| a = xd.abs().mean(dim=-1, keepdim=True).clamp(min=1e-5) |
| return _TernaryWgrad.apply(xq, w, (xd / a).round().clamp(-1, 1) * a) |
| return F.linear(xq, w) |
|
|
|
|
| def _fp_magnitudes(e_bits, m_bits, bias, top_nan): |
| """Non-negative values of a small IEEE-style float (subnormals included).""" |
| vals = {0.0} |
| for m in range(1, 1 << m_bits): |
| vals.add(m / (1 << m_bits) * 2.0 ** (1 - bias)) |
| for e in range(1, (1 << e_bits) - (0 if top_nan else 1)): |
| for m in range(1 << m_bits): |
| if top_nan and e == (1 << e_bits) - 1 and m == (1 << m_bits) - 1: |
| continue |
| vals.add((1 + m / (1 << m_bits)) * 2.0 ** (e - bias)) |
| return sorted(vals) |
|
|
|
|
| def _tnf_magnitudes(trits, m_bits): |
| """Ternary-exponent float packed into 8 binary bits (trinity-fpga |
| research/block/four_families.py levels()): e in [-(3^t)//2, (3^t)//2].""" |
| half = 3 ** trits // 2 |
| return sorted({0.0} | {(1 + m / (1 << m_bits)) * 2.0 ** e |
| for e in range(-half, half + 1) for m in range(1 << m_bits)}) |
|
|
|
|
| FORMATS8 = { |
| "e4m3": _fp_magnitudes(4, 3, 7, top_nan=True), |
| "e5m2": _fp_magnitudes(5, 2, 15, top_nan=False), |
| "tnf8e2": _tnf_magnitudes(2, 3), |
| "tnf8e3": _tnf_magnitudes(3, 2), |
| } |
| _TABLES = {} |
|
|
|
|
| def _table(name, device): |
| key = (name, str(device)) |
| if key not in _TABLES: |
| v = torch.tensor(FORMATS8[name], dtype=torch.float32, device=device) |
| _TABLES[key] = (v, (v[1:] + v[:-1]) / 2) |
| return _TABLES[key] |
|
|
|
|
| def round8(x, name, scale): |
| """Round to the nearest value of an 8-bit format (saturating); optionally after |
| per-tensor scaling so the tensor's max magnitude maps to the format's max.""" |
| v, mid = _table(name, x.device) |
| xf = x.float() |
| s = (v[-1] / xf.abs().amax().clamp(min=1e-30)) if scale else None |
| y = xf * s if scale else xf |
| q = torch.sign(y) * v[torch.bucketize(y.abs(), mid)] |
| return (q / s if scale else q).to(x.dtype) |
|
|
|
|
| class _Fake8(torch.autograd.Function): |
| """Linear layer whose three matmuls take 8-bit operands: x and W forward, |
| the output gradient backward (fp32 master weights stay outside).""" |
|
|
| @staticmethod |
| def forward(ctx, x, w, fwd, bwd, scale): |
| xq, wq = round8(x, fwd, scale), round8(w, fwd, scale) |
| ctx.save_for_backward(xq, wq) |
| ctx.bwd, ctx.scale, ctx.dtypes = bwd, scale, (x.dtype, w.dtype) |
| return F.linear(xq, wq) |
|
|
| @staticmethod |
| def backward(ctx, gy): |
| xq, wq = ctx.saved_tensors |
| g = round8(gy, ctx.bwd, ctx.scale) |
| gx = (g @ wq.to(g.dtype)).to(ctx.dtypes[0]) |
| gw = g.reshape(-1, g.size(-1)).t() @ xq.reshape(-1, xq.size(-1)).to(g.dtype) |
| return gx, gw.to(ctx.dtypes[1]), None, None, None |
|
|
|
|
| class Fake8Linear(nn.Linear): |
| fmt_fwd, fmt_bwd, scale = "e4m3", "e5m2", False |
|
|
| def forward(self, x): |
| return _Fake8.apply(x, self.weight, self.fmt_fwd, self.fmt_bwd, self.scale) |
|
|
|
|
| def lns_round(x, base): |
| """Nearest signed power of `base`, rounded in the log domain, as an 8-bit |
| code would hold it (sign + 7-bit exponent in [-96, 31]); tiny values -> 0. |
| base == 1 returns x unchanged (used to test the attention path).""" |
| if base == 1.0: |
| return x |
| mag = x.abs().float() |
| e = torch.round(torch.log(mag.clamp(min=1e-38)) / math.log(base)) |
| q = torch.sign(x.float()) * torch.pow(base, e.clamp(-96, 31)) |
| return torch.where(e < -96, torch.zeros_like(q), q).to(x.dtype) |
|
|
|
|
| def _lns_q(x, inv_lb, lb): |
| """lns_round for float32 tensors with precomputed 1/log(base), log(base).""" |
| e = torch.round(torch.log(x.abs().clamp(min=1e-38)) * inv_lb) |
| q = torch.sign(x) * torch.exp(e.clamp(-96, 31) * lb) |
| return torch.where(e < -96, torch.zeros_like(q), q) |
|
|
|
|
| def _causal_mask(T, device): |
| return torch.ones(T, T, dtype=torch.bool, device=device).triu(1) |
|
|
|
|
| def _lns_attn_fwd(q, k, v, inv_lb, lb, scale): |
| qs, ks, vs = (_lns_q(t.float(), inv_lb, lb) for t in (q, k, v)) |
| s = (qs @ ks.transpose(-1, -2)) * scale |
| p = torch.softmax(s.masked_fill(_causal_mask(s.size(-1), s.device), float("-inf")), dim=-1) |
| return _lns_q(p, inv_lb, lb) @ vs, qs, ks, vs |
|
|
|
|
| def _lns_attn_bwd(qs, ks, vs, go, inv_lb, lb, scale): |
| s = (qs @ ks.transpose(-1, -2)) * scale |
| p = torch.softmax(s.masked_fill(_causal_mask(s.size(-1), s.device), float("-inf")), dim=-1) |
| ps, gos = _lns_q(p, inv_lb, lb), _lns_q(go.float(), inv_lb, lb) |
| gv = ps.transpose(-1, -2) @ gos |
| gp = gos @ vs.transpose(-1, -2) |
| gs = _lns_q(p * (gp - (gp * p).sum(-1, keepdim=True)) * scale, inv_lb, lb) |
| return gs @ ks, gs.transpose(-1, -2) @ qs, gv |
|
|
|
|
| _lns_attn_fwd_c = torch.compile(_lns_attn_fwd, dynamic=False) |
| _lns_attn_bwd_c = torch.compile(_lns_attn_bwd, dynamic=False) |
|
|
|
|
| class _LNSAttention(torch.autograd.Function): |
| """Causal attention whose six matrix products (two forward, four backward) |
| take operands rounded to signed powers of `base`, as a logarithmic-number |
| datapath would (a product there is an exponent addition). Accumulation and |
| softmax stay exact, i.e. an ideal wide accumulator; P is recomputed. |
| base == 1 skips the rounding (used to test the attention path).""" |
|
|
| @staticmethod |
| def forward(ctx, q, k, v, base): |
| scale = q.size(-1) ** -0.5 |
| lb = math.log(base) if base != 1.0 else 0.0 |
| inv = 1.0 / lb if lb else 0.0 |
| fwd = _lns_attn_fwd_c if q.is_cuda else _lns_attn_fwd |
| if base == 1.0: |
| qs, ks, vs = q.float(), k.float(), v.float() |
| p = _causal_probs(qs, ks, scale) |
| out = p @ vs |
| else: |
| out, qs, ks, vs = fwd(q, k, v, inv, lb, scale) |
| ctx.save_for_backward(qs, ks, vs) |
| ctx.base, ctx.scale, ctx.lb, ctx.inv, ctx.dtype = base, scale, lb, inv, q.dtype |
| return out.to(q.dtype) |
|
|
| @staticmethod |
| def backward(ctx, go): |
| qs, ks, vs = ctx.saved_tensors |
| if ctx.base == 1.0: |
| p = _causal_probs(qs, ks, ctx.scale) |
| g = go.float() |
| gv = p.transpose(-1, -2) @ g |
| gp = g @ vs.transpose(-1, -2) |
| gs = p * (gp - (gp * p).sum(-1, keepdim=True)) * ctx.scale |
| gq, gk = gs @ ks, gs.transpose(-1, -2) @ qs |
| else: |
| bwd = _lns_attn_bwd_c if go.is_cuda else _lns_attn_bwd |
| gq, gk, gv = bwd(qs, ks, vs, go, ctx.inv, ctx.lb, ctx.scale) |
| return gq.to(ctx.dtype), gk.to(ctx.dtype), gv.to(ctx.dtype), None |
|
|
|
|
| def _causal_probs(q, k, scale): |
| s = (q.float() @ k.float().transpose(-1, -2)) * scale |
| return torch.softmax(s.masked_fill(_causal_mask(s.size(-1), s.device), float("-inf")), dim=-1) |
|
|
|
|
| @torch.compiler.disable |
| def lns_attention(q, k, v, base): |
| """Opaque to the model-level compile; the helpers inside are compiled on their own.""" |
| return _LNSAttention.apply(q, k, v, base) |
|
|
|
|
| def rope_cache(seq_len, head_dim, theta, device): |
| inv = 1.0 / theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim) |
| t = torch.arange(seq_len, device=device).float() |
| freqs = torch.outer(t, inv) |
| return freqs.cos(), freqs.sin() |
|
|
|
|
| def apply_rope(x, cos, sin): |
| |
| x1, x2 = x.float().chunk(2, dim=-1) |
| cos, sin = cos[None, :, None, :], sin[None, :, None, :] |
| return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1).type_as(x) |
|
|
|
|
| class Attention(nn.Module): |
| def __init__(self, cfg, Lin): |
| super().__init__() |
| self.lns = cfg.attn_lns |
| self.h, self.kv = cfg.n_head, cfg.n_kv_head |
| self.hd = cfg.d_model // cfg.n_head |
| d = cfg.d_model |
| self.q = Lin(d, self.h * self.hd, bias=False) |
| self.k = Lin(d, self.kv * self.hd, bias=False) |
| self.v = Lin(d, self.kv * self.hd, bias=False) |
| self.o = Lin(self.h * self.hd, d, bias=False) |
|
|
| def forward(self, x, cos, sin, cache=None): |
| B, T, _ = x.shape |
| q = apply_rope(self.q(x).view(B, T, self.h, self.hd), cos, sin).transpose(1, 2) |
| k = apply_rope(self.k(x).view(B, T, self.kv, self.hd), cos, sin).transpose(1, 2) |
| v = self.v(x).view(B, T, self.kv, self.hd).transpose(1, 2) |
| if cache is not None: |
| k, v = cache.update(k, v) |
| |
| causal = T > 1 |
| if self.kv != self.h: |
| k = k.repeat_interleave(self.h // self.kv, dim=1) |
| v = v.repeat_interleave(self.h // self.kv, dim=1) |
| if self.lns > 1 and cache is None: |
| y = lns_attention(q, k, v, self.lns) |
| else: |
| y = F.scaled_dot_product_attention(q, k, v, is_causal=causal) |
| return self.o(y.transpose(1, 2).reshape(B, T, self.h * self.hd)) |
|
|
|
|
| class MLP(nn.Module): |
| def __init__(self, cfg, Lin): |
| super().__init__() |
| self.gate = Lin(cfg.d_model, cfg.d_ff, bias=False) |
| self.up = Lin(cfg.d_model, cfg.d_ff, bias=False) |
| self.down = Lin(cfg.d_ff, cfg.d_model, bias=False) |
|
|
| def forward(self, x): |
| return self.down(F.silu(self.gate(x)) * self.up(x)) |
|
|
|
|
| class Block(nn.Module): |
| def __init__(self, cfg, Lin): |
| super().__init__() |
| self.n1 = nn.RMSNorm(cfg.d_model, eps=1e-6) |
| self.attn = Attention(cfg, Lin) |
| self.n2 = nn.RMSNorm(cfg.d_model, eps=1e-6) |
| self.mlp = MLP(cfg, Lin) |
|
|
| def forward(self, x, cos, sin, cache=None): |
| x = x + self.attn(self.n1(x), cos, sin, cache) |
| return x + self.mlp(self.n2(x)) |
|
|
|
|
| class KVCache: |
| """Per-layer key/value buffers for batched generation of equal-length prompts.""" |
|
|
| def __init__(self, B, kv_heads, head_dim, max_len, device, dtype): |
| self.k = torch.zeros(B, kv_heads, max_len, head_dim, device=device, dtype=dtype) |
| self.v = torch.zeros_like(self.k) |
| self.pos = 0 |
|
|
| def update(self, k, v): |
| T = k.size(2) |
| self.k[:, :, self.pos:self.pos + T] = k |
| self.v[:, :, self.pos:self.pos + T] = v |
| self.pos += T |
| return self.k[:, :, :self.pos], self.v[:, :, :self.pos] |
|
|
|
|
| class GPT(nn.Module): |
| def __init__(self, cfg: Config): |
| super().__init__() |
| self.cfg = cfg |
| Lin = BitLinear if cfg.ternary else Fake8Linear if cfg.lin8 else nn.Linear |
| self.wte = nn.Embedding(cfg.vocab_size, cfg.d_model) |
| self.blocks = nn.ModuleList(Block(cfg, Lin) for _ in range(cfg.n_layer)) |
| self.norm = nn.RMSNorm(cfg.d_model, eps=1e-6) |
| cos, sin = rope_cache(cfg.seq_len * 2, cfg.d_model // cfg.n_head, cfg.rope_theta, "cpu") |
| self.register_buffer("cos", cos, persistent=False) |
| self.register_buffer("sin", sin, persistent=False) |
| self.apply(self._init) |
| fwd8, _, bwd8 = cfg.lin8.partition(":") |
| for m in self.modules(): |
| if isinstance(m, BitLinear): |
| m.wgrad_ternary = cfg.wgrad_ternary |
| if isinstance(m, Fake8Linear): |
| m.fmt_fwd, m.fmt_bwd, m.scale = fwd8, bwd8 or fwd8, cfg.lin8_scale |
| for name, p in self.named_parameters(): |
| if name.endswith(("attn.o.weight", "mlp.down.weight")): |
| nn.init.normal_(p, std=0.02 / math.sqrt(2 * cfg.n_layer)) |
|
|
| @staticmethod |
| def _init(m): |
| if isinstance(m, (nn.Linear, nn.Embedding)): |
| nn.init.normal_(m.weight, std=0.02) |
|
|
| def forward(self, idx, targets=None, caches=None): |
| T = idx.size(1) |
| start = caches[0].pos if caches else 0 |
| cos, sin = self.cos[start:start + T], self.sin[start:start + T] |
| x = self.wte(idx) |
| for i, blk in enumerate(self.blocks): |
| x = blk(x, cos, sin, caches[i] if caches else None) |
| logits = F.linear(self.norm(x), self.wte.weight) |
| if targets is None: |
| return logits |
| loss = F.cross_entropy(logits.float().view(-1, logits.size(-1)), targets.reshape(-1), |
| reduction="none") |
| return loss.view(targets.shape) |
|
|
| def num_params(self): |
| return sum(p.numel() for p in self.parameters()) |
|
|
| def flops_per_token(self, seq_len=None): |
| """Training FLOPs per token: 6 * matmul params + causal attention (Kaplan 2020). |
| |
| Every parameter except the norms takes part in a matmul (the tied |
| embedding matrix does so as the output head).""" |
| c = self.cfg |
| norms = sum(p.numel() for m in self.modules() if isinstance(m, nn.RMSNorm) |
| for p in m.parameters()) |
| return 6 * (self.num_params() - norms) + 6 * c.n_layer * (seq_len or c.seq_len) * c.d_model |
|
|
| @torch.no_grad() |
| def generate(self, prompt, max_new, temperature=0.2, top_p=0.95, cache_dtype=torch.float32, |
| stop_fn=None, check_every=16): |
| """Sample continuations for a batch of equal-length prompts (B, T). |
| |
| stop_fn(tokens) is called every `check_every` steps with the tokens |
| generated so far; returning True ends generation early.""" |
| B, T = prompt.shape |
| c = self.cfg |
| caches = [KVCache(B, c.n_kv_head, c.d_model // c.n_head, T + max_new, prompt.device, |
| cache_dtype) for _ in range(c.n_layer)] |
| logits = self(prompt, caches=caches)[:, -1] |
| out = [] |
| for i in range(max_new): |
| if temperature > 0: |
| probs = torch.softmax(logits.float() / temperature, dim=-1) |
| sp, si = probs.sort(dim=-1, descending=True) |
| sp[sp.cumsum(-1) - sp > top_p] = 0 |
| nxt = si.gather(-1, torch.multinomial(sp / sp.sum(-1, keepdim=True), 1)) |
| else: |
| nxt = logits.argmax(-1, keepdim=True) |
| out.append(nxt) |
| if stop_fn and (i + 1) % check_every == 0 and stop_fn(torch.cat(out, dim=1)): |
| break |
| logits = self(nxt, caches=caches)[:, -1] |
| return torch.cat(out, dim=1) |
|
|
|
|
| def build(size, ternary=False, vocab_size=32000, seq_len=2048, **extra): |
| return GPT(Config(vocab_size=vocab_size, seq_len=seq_len, ternary=ternary, **SIZES[size], |
| **extra)) |
|
|
|
|
| def config_dict(model): |
| return asdict(model.cfg) |
|
|
