scripts/model_dna_peer.py

# model_dna_peer.py - DNA-DiskChat with a TERNARY controller + PEER expert FFN.
# Backbone: the fast DNA FastBlock recurrence (chunk-parallel) but with BitNet b1.58
# ternary weights; the dense FFN is replaced by a product-key expert (PEER) bank so
# controller capacity grows toward ~500M while active FLOPs stay ~flat. The codon
# memory table / router are inherited unchanged from model_dna.DnaChat.
import sys; sys.path.insert(0, '/root/dna')
import torch, torch.nn as nn, torch.nn.functional as F
from model_dna import DnaChat, diag_scan, counts

def tern(w):                              # BitNet b1.58 ternary STE (weight-only)
    s = w.abs().mean().clamp_min(1e-5)
    wt = (w / s).round().clamp_(-1, 1) * s
    return w + (wt - w).detach()

class BitLinear(nn.Module):
    def __init__(self, i, o):
        super().__init__(); self.weight = nn.Parameter(torch.empty(o, i)); nn.init.normal_(self.weight, std=0.02)
    def forward(self, x): return F.linear(x, tern(self.weight))

class PeerFFN(nn.Module):
    """Product-key expert retrieval FFN. E=nk^2 rank-1 experts/head; retrieve top-k.
    out = sum_k softmax(score)_k * Wdown_k * silu(Wup_k . x).  Active FLOPs ~ topk*d."""
    def __init__(self, d, nk=176, dk=128, topk=8, heads=1):
        super().__init__()
        self.d, self.nk, self.dk, self.topk, self.h, self.E = d, nk, dk, topk, heads, nk * nk
        self.q = nn.Linear(d, heads * 2 * dk, bias=False)
        self.Ka = nn.Parameter(torch.randn(heads, nk, dk) * 0.02)
        self.Kb = nn.Parameter(torch.randn(heads, nk, dk) * 0.02)
        self.up = nn.Embedding(self.E, d)        # Wup   [E,d]
        self.down = nn.Embedding(self.E, d)      # Wdown [E,d]
        nn.init.normal_(self.up.weight, std=0.02); nn.init.normal_(self.down.weight, std=0.02 / (d ** 0.5))
    def forward(self, x):                        # x: [N, d]
        N = x.size(0); k = self.topk
        q = self.q(x).view(N, self.h, 2, self.dk)
        qa, qb = q[:, :, 0], q[:, :, 1]          # [N,h,dk]
        sa = torch.einsum('nhd,hkd->nhk', qa, self.Ka)   # [N,h,nk]
        sb = torch.einsum('nhd,hkd->nhk', qb, self.Kb)
        va, ia = sa.topk(k, -1); vb, ib = sb.topk(k, -1)  # [N,h,k]
        cs = va.unsqueeze(-1) + vb.unsqueeze(-2)          # [N,h,k,k]
        ci = ia.unsqueeze(-1) * self.nk + ib.unsqueeze(-2)
        cs = cs.reshape(N, self.h, -1); ci = ci.reshape(N, self.h, -1)
        tv, ti = cs.topk(k, -1)                            # [N,h,k]
        eid = torch.gather(ci, -1, ti)                     # [N,h,k] expert ids
        gate = torch.softmax(tv, -1)                       # [N,h,k]
        u = self.up(eid); dn = self.down(eid)              # [N,h,k,d]
        hid = F.silu(torch.einsum('nhkd,nd->nhk', u, x))   # [N,h,k]
        out = torch.einsum('nhk,nhkd->nd', gate * hid, dn) # [N,d]
        return out / self.h

class FastBlockTP(nn.Module):
    def __init__(self, d, nk, dk, topk, heads):
        super().__init__()
        self.n1 = nn.LayerNorm(d); self.proj = BitLinear(d, 4 * d); self.o = BitLinear(d, d)
        self.decay = nn.Parameter(torch.full((d,), 2.0))
        self.n2 = nn.LayerNorm(d); self.peer = PeerFFN(d, nk, dk, topk, heads)
    def forward(self, x, C=16):
        k, v, r, g = self.proj(self.n1(x)).chunk(4, -1)
        g = torch.sigmoid(g + self.decay); u = (1 - g) * torch.tanh(k)
        s = diag_scan(g, u, C); x = x + self.o(torch.sigmoid(r) * s * torch.sigmoid(v))
        B, T, d = x.shape
        return x + self.peer(self.n2(x).reshape(B * T, d)).reshape(B, T, d)

class DnaPeer(DnaChat):
    def __init__(self, vocab=32000, d=512, layers=16, ff=1536, a=2048, b=2048,
                 nk=176, dk=128, topk=8, pheads=1, chunk=16):
        super().__init__(vocab, d, layers, ff, a, b, chunk)
        self.nk, self.dk, self.topk, self.pheads = nk, dk, topk, pheads
        self.blocks = nn.ModuleList([FastBlockTP(d, nk, dk, topk, pheads) for _ in range(layers)])
    def config(self):
        c = super().config(); c.update(nk=self.nk, dk=self.dk, topk=self.topk, pheads=self.pheads); return c

def peer_counts(m):
    ctrl = sum(p.numel() for p in m.parameters())
    active = 0
    for blk in m.blocks:
        active += blk.proj.weight.numel() + blk.o.weight.numel()          # ternary recurrence
        active += blk.peer.q.weight.numel() + m.topk * 2 * m.d            # router + topk experts
    active += m.embed.weight.numel()                                       # tied head
    active += m.ra.weight.numel() + m.rb.weight.numel()                    # memory router
    return ctrl, active