| |
| 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')] |
|
|
