# Têtu Code 2 — inference code (tetu-AI). License: Apache 2.0, see README.md. from dataclasses import asdict, dataclass, field import math from typing import Optional import torch import torch.nn as nn import torch.nn.functional as F try: from torch.nn.attention.flex_attention import BlockMask, flex_attention except ImportError: BlockMask = flex_attention = None @dataclass class TetuCodeConfig: vocab_size: int = 32768 d_model: int = 768 n_layers: int = 24 n_heads: int = 12 n_kv_heads: int = 4 mlp_mult: int = 4 max_seq_len: int = 65536 window: int = 2048 global_every: int = 4 rope_local: float = 10000.0 rope_global: float = 500000.0 log_base: int = 4096 engram: tuple = ((1, (2, 3)), (8, (3, 4)), (15, (4, 6))) engram_heads: int = 2 engram_slots: int = 393209 engram_dim: int = 128 n_value_embeds: int = 3 struct_sizes: tuple = (16, 16, 4, 2) softcap: float = 30.0 doubt_weight: float = 0.1 eos_id: int = 1 def __post_init__(self): self.engram = tuple(((int(l), tuple(o)) for l, o in self.engram)) self.struct_sizes = tuple(self.struct_sizes) def is_global(self, i): return (i + 1) % self.global_every == 0 def to_dict(self): return asdict(self) def rms(x): return F.rms_norm(x, (x.size(-1),)) def rope_table(n, hd, theta): inv = 1.0 / theta ** (torch.arange(0, hd, 2, dtype=torch.float32) / hd) f = torch.outer(torch.arange(n, dtype=torch.float32), inv) return (f.cos(), f.sin()) def apply_rope(x, cos, sin): x1, x2 = x.chunk(2, dim=-1) return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1).type_as(x) def _ordered(m): num = m.sum(dim=-1, dtype=torch.int32) idx = m.to(torch.int32).argsort(dim=-1, descending=False, stable=True).flip(-1).to(torch.int32) return (num[None, None].contiguous(), idx[None, None].contiguous()) def doc_block_masks(doc, window, block=128): n, sw = (doc.size(0) // block, window // block) def doc_causal(b, h, q, kv): return (q >= kv) & (doc[q] == doc[kv]) def doc_local(b, h, q, kv): return doc_causal(b, h, q, kv) & (q // block - kv // block < sw) i = torch.arange(n, device=doc.device) lo, hi = (doc.view(n, block)[:, 0], doc.view(n, block)[:, -1]) some = (i[:, None] >= i[None]) & (lo[:, None] <= hi[None]) & (lo[None] <= hi[:, None]) full = (i[:, None] > i[None]) & (lo[:, None] == hi[None]) & (lo[None] == hi[:, None]) kn, ki = _ordered(some & ~full) fn, fi = _ordered(full) mk = lambda a, b, f: BlockMask.from_kv_blocks(a, ki, b, fi, BLOCK_SIZE=block, mask_mod=f) return (mk(kn, fn, doc_causal), mk(torch.clamp_max(kn, torch.clamp_min(sw - fn, 1)), torch.clamp_max(fn, sw - 1), doc_local)) def dense_masks(doc, window, block=128): i = torch.arange(doc.size(0), device=doc.device) g = (i[:, None] >= i[None]) & (doc[:, None] == doc[None]) return (g, g & (i[:, None] // block - i[None] // block < window // block)) class Linear(nn.Linear): def __init__(self, i, o, zero=False): super().__init__(i, o, bias=False) with torch.no_grad(): if zero: self.weight.zero_() else: b = math.sqrt(3) * 0.5 / math.sqrt(i) self.weight.uniform_(-b, b) class Engram(nn.Module): P = 2147483647 def __init__(self, cfg: TetuCodeConfig, orders, seed): super().__init__() self.orders, self.K, self.S, self.dm = (tuple(orders), cfg.engram_heads, cfg.engram_slots, cfg.engram_dim) self.n_tab = len(self.orders) * self.K self.pad_id = cfg.vocab_size self.table = nn.Embedding(self.n_tab * self.S, self.dm) nn.init.normal_(self.table.weight, std=1.0) g = torch.Generator().manual_seed(seed) mult = torch.randint(1 << 20, 1 << 30, (self.n_tab, max(self.orders)), generator=g, device=g.device) | 1 self.register_buffer('mult', mult, persistent=False) self.register_buffer('offset', torch.arange(self.n_tab) * self.S, persistent=False) d, de = (cfg.d_model, self.n_tab * self.dm) self.wk = Linear(de, d) self.wv = Linear(de, d, zero=True) self.conv = nn.Conv1d(d, d, 4, groups=d, bias=False) nn.init.zeros_(self.conv.weight) def slots(self, idx): T = idx.size(1) grams = [idx] + [F.pad(idx, (j, 0), value=self.pad_id)[:, :T] for j in range(1, max(self.orders))] out, t = ([], 0) for o in self.orders: for _ in range(self.K): h = torch.zeros_like(idx) for j in range(o): h = (h * 1000003 + (grams[j] + 1) * self.mult[t, j]) % self.P out.append(h % self.S + self.offset[t]) t += 1 return torch.stack(out, dim=-1) def read(self, x, e): k, v = (self.wk(e), self.wv(e)) alpha = torch.sigmoid((rms(x) * rms(k)).sum(-1, keepdim=True) / math.sqrt(x.size(-1))) return (alpha * v, alpha) def forward(self, x, idx): B, T, d = x.shape e = self.table(self.slots(idx)).reshape(B, T, -1).to(x.dtype) s, alpha = self.read(x, e) c = self.conv(F.pad(s.transpose(1, 2), (3, 0))).transpose(1, 2) if not torch.compiler.is_compiling(): self.last_gate = alpha.detach().float().mean() return s + F.silu(c) class Attention(nn.Module): def __init__(self, cfg: TetuCodeConfig, is_global: bool): super().__init__() self.H, self.Hk = (cfg.n_heads, cfg.n_kv_heads) self.hd = cfg.d_model // cfg.n_heads d = cfg.d_model self.wq = Linear(d, self.H * self.hd) self.wk = Linear(d, self.Hk * self.hd) self.wv = Linear(d, self.Hk * self.hd) self.wg = Linear(d, self.H * self.hd) self.wo = Linear(self.H * self.hd, d, zero=True) self.is_global, self.window = (is_global, cfg.window) def qkv(self, x, cos, sin, qs=None, ve=None, lam=None): B, T, _ = x.shape q = self.wq(x).view(B, T, self.H, self.hd).transpose(1, 2) k = self.wk(x).view(B, T, self.Hk, self.hd).transpose(1, 2) v = self.wv(x).view(B, T, self.Hk, self.hd).transpose(1, 2) if ve is not None: v = v + lam * ve.to(v.dtype).view(B, T, self.Hk, self.hd).transpose(1, 2) q, k = (apply_rope(rms(q), cos, sin), apply_rope(rms(k), cos, sin)) if qs is not None: q = q * qs.to(q.dtype)[:, None] return (q, k, v) def forward(self, x, cos, sin, mask, qs=None, ve=None, lam=None): B, T, _ = x.shape q, k, v = self.qkv(x, cos, sin, qs, ve, lam) if mask is not None and (not isinstance(mask, torch.Tensor)): y = flex_attention(q, k, v, block_mask=mask, enable_gqa=True) else: r = self.H // self.Hk k, v = (k.repeat_interleave(r, dim=1), v.repeat_interleave(r, dim=1)) if mask is not None: y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask) elif self.is_global or T <= self.window - 127: y = F.scaled_dot_product_attention(q, k, v, is_causal=True) else: i = torch.arange(T, device=x.device) m = (i[:, None] >= i[None]) & (i[:, None] // 128 - i[None] // 128 < self.window // 128) y = F.scaled_dot_product_attention(q, k, v, attn_mask=m) y = y.transpose(1, 2).reshape(B, T, -1) * torch.sigmoid(self.wg(x)) return self.wo(y) class MLP(nn.Module): def __init__(self, d, h): super().__init__() self.w1, self.w2 = (Linear(d, h), Linear(h, d, zero=True)) def forward(self, x): return self.w2(F.relu(self.w1(x)).square()) class Block(nn.Module): def __init__(self, cfg: TetuCodeConfig, i: int, engram_orders=None): super().__init__() self.attn = Attention(cfg, cfg.is_global(i)) self.mlp = MLP(cfg.d_model, cfg.mlp_mult * cfg.d_model) self.engram = Engram(cfg, engram_orders, seed=1234 + i) if engram_orders else None def forward(self, x, idx, cos, sin, mask, qs, ve, lam): if self.engram is not None: x = x + self.engram(rms(x), idx) x = x + self.attn(rms(x), cos, sin, mask, qs, ve, lam) return x + self.mlp(rms(x)) class DoubtHead(nn.Module): def __init__(self, d): super().__init__() self.net = nn.Sequential(nn.Linear(d + 2, 128), nn.ReLU(), nn.Linear(128, 1)) def forward(self, h, logits): with torch.no_grad(): lp = F.log_softmax(logits.float(), dim=-1) conf = torch.stack([lp.max(-1).values, -(lp.exp() * lp).sum(-1)], dim=-1) / 4 return self.net(torch.cat([rms(h.detach()).float(), conf], dim=-1)).squeeze(-1) @dataclass class Output: logits: torch.Tensor loss: Optional[torch.Tensor] = None lm_loss: Optional[torch.Tensor] = None doubt_loss: Optional[torch.Tensor] = None doubt: Optional[torch.Tensor] = None unlike_loss: Optional[torch.Tensor] = None def unpack_struct(s): s = s.long() return (s & 15, s >> 4 & 15, s >> 8 & 3, s >> 10 & 1) class TetuCode(nn.Module): def __init__(self, cfg: TetuCodeConfig): super().__init__() self.unlike_weight = 1.0 self.ckpt_blocks = False self.head_chunk = 8192 self.cfg = cfg d, L = (cfg.d_model, cfg.n_layers) kv = cfg.n_kv_heads * (d // cfg.n_heads) self.tok_emb = nn.Embedding(cfg.vocab_size, d) nn.init.normal_(self.tok_emb.weight, std=1.0) self.struct_emb = nn.ModuleList((nn.Embedding(n, d) for n in cfg.struct_sizes)) for e in self.struct_emb: nn.init.normal_(e.weight, std=0.3) self.value_embeds = nn.ModuleList((nn.Embedding(cfg.vocab_size, kv) for _ in range(cfg.n_value_embeds))) for e in self.value_embeds: nn.init.normal_(e.weight, std=1.0) n = cfg.n_value_embeds self.ve_layer = {i: i for i in range(n)} | {L - n + i: i for i in range(n)} self.ve_lambda = nn.Parameter(torch.full((L,), 0.5)) eng = dict(cfg.engram) self.blocks = nn.ModuleList((Block(cfg, i, eng.get(i)) for i in range(L))) self.skip_w = nn.Parameter(torch.ones(L // 2)) self.lm_head = Linear(d, cfg.vocab_size, zero=True) self.doubt_head = DoubtHead(d) hd = d // cfg.n_heads for name, theta in (('loc', cfg.rope_local), ('glo', cfg.rope_global)): c, s = rope_table(cfg.max_seq_len, hd, theta) self.register_buffer(f'cos_{name}', c, persistent=False) self.register_buffer(f'sin_{name}', s, persistent=False) def num_params(self): return sum((p.numel() for p in self.parameters())) def compute_params(self): lookup = self.tok_emb.weight.numel() + sum((e.weight.numel() for e in self.value_embeds)) lookup += sum((e.weight.numel() for e in self.struct_emb)) lookup += sum((b.engram.table.weight.numel() for b in self.blocks if b.engram is not None)) return self.num_params() - lookup def param_groups(self): tables, engram, head, scal, muon = ([], [], [], [], []) for n, p in self.named_parameters(): if n.endswith('engram.table.weight'): engram.append(p) elif n.startswith(('tok_emb', 'value_embeds', 'struct_emb')): tables.append(p) elif n.startswith('lm_head'): head.append(p) elif p.dim() == 2 and (not n.startswith('doubt_head')): muon.append(p) else: scal.append(p) return {'tables': tables, 'engram': engram, 'head': head, 'scalars': scal, 'muon': muon} def rope(self, T, pos=None): p = torch.arange(T, device=self.cos_loc.device) if pos is None else pos qs = (torch.log(p.float() + 1) / math.log(self.cfg.log_base)).clamp_min(1.0) return ((self.cos_loc[p], self.sin_loc[p]), (self.cos_glo[p], self.sin_glo[p]), qs) def embed(self, idx, struct=None): x = self.tok_emb(idx) if struct is not None: for e, f in zip(self.struct_emb, unpack_struct(struct)): x = x + e(f) return rms(x) def hidden(self, idx, struct=None, doc=None, pos=None): B, T = idx.shape loc, glo, qs = self.rope(T, pos) x = self.embed(idx, struct) L, skips = (self.cfg.n_layers, []) mg = ml = None if doc is not None: mg, ml = (doc_block_masks if idx.is_cuda else dense_masks)(doc, self.cfg.window) for i, blk in enumerate(self.blocks): if i >= L - L // 2: x = x + self.skip_w[L - 1 - i] * skips.pop() j = self.ve_layer.get(i) ve = self.value_embeds[j](idx) if j is not None else None g = blk.attn.is_global c, s = glo if g else loc args = (x, idx, c, s, mg if g else ml, qs if g else None, ve, self.ve_lambda[i]) if self.ckpt_blocks and torch.is_grad_enabled(): x = torch.utils.checkpoint.checkpoint(blk, *args, use_reentrant=False) else: x = blk(*args) if i < L // 2: skips.append(x) return rms(x) def logits(self, h): z = self.lm_head(h).float() return self.cfg.softcap * torch.tanh(z / self.cfg.softcap) def _head_chunk(self, h, t): z = self.logits(h) lse = torch.logsumexp(z, -1) good, bad = (t >= 0, t <= -200) zt = z.gather(-1, torch.where(good, t, (-t - 200).clamp(0, z.size(-1) - 1))[:, None])[:, 0] loss = ((lse - zt) * good).sum() p_bad = torch.exp(zt - lse).clamp(max=1 - 1e-06) unlike = (-torch.log1p(-p_bad) * bad).sum() with torch.no_grad(): lp = F.log_softmax(z, dim=-1) conf = torch.stack([lp.max(-1).values, -(lp.exp() * lp).sum(-1)], dim=-1) / 4 wrong = (z.argmax(-1) != t).float() return (loss, conf, wrong, unlike) def forward(self, idx, targets=None, struct=None, doc=None, pos=None, chunk=None): chunk = chunk or self.head_chunk h = self.hidden(idx, struct, doc, pos) if targets is None: logits = self.logits(h) return Output(logits=logits, doubt=torch.sigmoid(self.doubt_head(h, logits))) hf, tf = (h.flatten(0, 1), targets.flatten()) parts = [torch.utils.checkpoint.checkpoint(self._head_chunk, hf[i:i + chunk], tf[i:i + chunk], use_reentrant=False) for i in range(0, hf.size(0), chunk)] mask = tf >= 0 n = mask.sum().clamp_min(1) conf = torch.cat([p[1] for p in parts]) wrong = torch.cat([p[2] for p in parts]) doubt_logit = self.doubt_head.net(torch.cat([rms(hf.detach()).float(), conf], dim=-1)).squeeze(-1) bce = F.binary_cross_entropy_with_logits(doubt_logit, wrong, reduction='none') out = Output(logits=None, doubt=None) out.lm_loss = sum((p[0] for p in parts)) / n out.doubt_loss = (bce * mask).sum() / n out.unlike_loss = sum((p[3] for p in parts)) / n out.loss = out.lm_loss + self.cfg.doubt_weight * out.doubt_loss + self.unlike_weight * out.unlike_loss return out def engram_gates(self): return [b.engram.last_gate for b in self.blocks if b.engram is not None and hasattr(b.engram, 'last_gate')]