Upload scripts/model_dna.py with huggingface_hub

scripts/model_dna.py

@@ -0,0 +1,100 @@
+# model_dna.py - DNA-DiskChat-0.5b: DiskChat backbone + PKM table stored as CODON rows.
+# v14 scheme: rows persist as codons (uint8) + shared learned codebook; training keeps an
+# fp32 shadow (STE) and re-encodes touched rows each step ("mutation").
+import math, torch
+from torch import nn
+import torch.nn.functional as F
+
+G=4; KCODE=256  # codon group + codewords: 512 dims -> 128 codons -> 128 B/row (75% off int8)
+
+def diag_scan(g,u,C=16):
+    B,T,D=g.shape; pad=(C-T%C)%C
+    if pad: g=F.pad(g,(0,0,0,pad),value=1.0); u=F.pad(u,(0,0,0,pad),value=0.0)
+    Tp=T+pad; NC=Tp//C; g=g.view(B,NC,C,D); u=u.view(B,NC,C,D)
+    s=torch.zeros(B,NC,D,device=g.device,dtype=g.dtype); P=torch.ones_like(s); loc=[]; Ps=[]
+    for t in range(C):
+        s=g[:,:,t,:]*s+u[:,:,t,:]; loc.append(s); P=P*g[:,:,t,:]; Ps.append(P)
+    s_local=torch.stack(loc,2); Pcum=torch.stack(Ps,2); Gc=Pcum[:,:,-1,:]; sl=s_local[:,:,-1,:]
+    carry=torch.zeros(B,D,device=g.device,dtype=g.dtype); car=[]
+    for c in range(NC):
+        car.append(carry); carry=Gc[:,c,:]*carry+sl[:,c,:]
+    car=torch.stack(car,1); return (s_local+Pcum*car[:,:,None,:]).view(B,Tp,D)[:,:T,:]
+
+class FastBlock(nn.Module):
+    def __init__(self,d,ff):
+        super().__init__(); self.n1=nn.LayerNorm(d); self.proj=nn.Linear(d,4*d,bias=False); self.o=nn.Linear(d,d,bias=False)
+        self.n2=nn.LayerNorm(d); self.up=nn.Linear(d,2*ff,bias=False); self.down=nn.Linear(ff,d,bias=False); self.decay=nn.Parameter(torch.full((d,),2.0))
+    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))
+        a,b=self.up(self.n2(x)).chunk(2,-1); return x+self.down(F.silu(b)*a)
+
+class STEDecode(torch.autograd.Function):
+    """fwd: decoded codon rows. bwd: pass grad straight through to shadow rows (STE)."""
+    @staticmethod
+    def forward(ctx,shadow_rows,decoded_rows):
+        return decoded_rows
+    @staticmethod
+    def backward(ctx,gout):
+        return gout,gout   # STE into shadow AND through decode into the codebook
+
+class DnaChat(nn.Module):
+    def __init__(self,vocab=32000,d=512,layers=16,ff=1536,a=2048,b=2048,chunk=16):
+        super().__init__(); self.vocab=vocab; self.d=d; self.layers=layers; self.ff=ff; self.a=a; self.b=b; self.chunk=chunk
+        self.ncod=d//G; self.nrows=a*b
+        self.embed=nn.Embedding(vocab,d); nn.init.normal_(self.embed.weight,std=.02)
+        self.blocks=nn.ModuleList([FastBlock(d,ff) for _ in range(layers)]); self.norm=nn.LayerNorm(d)
+        self.ra=nn.Linear(d,a,bias=False); self.rb=nn.Linear(d,b,bias=False)
+        self.ca=nn.Parameter(torch.randn(a,d)*.02); self.cb=nn.Parameter(torch.randn(b,d)*.02)
+        self.head2=nn.Linear(d,vocab,bias=False); nn.init.normal_(self.head2.weight,std=.02)
+        # DNA table (v16 ZERO-SHADOW + v17): genotype only; POSITIONAL codebooks (reading frames):
+        # 4 codebooks cycled by codon position -> 4x decode capacity at SAME 128 B/row rate.
+        self.codebook=nn.Parameter(torch.randn(4,KCODE,G)*.05)
+        self.register_buffer('codons',torch.randint(0,KCODE,(self.nrows,self.ncod),dtype=torch.uint8))
+        frames=torch.arange(self.ncod)%4
+        self.register_buffer('frames',frames)                       # [ncod]
+        # v17 UltraMem-style implicit expansion: stored 128 B -> effective 2x value dim
+        self.expand=nn.Linear(d,2*d,bias=False); self.contract=nn.Linear(2*d,d,bias=False)
+    @torch.no_grad()
+    def mutate_rows(self,idx,new_rows,chunk=65536):
+        """v17 mutation, CHUNKED to bound cdist memory (fixes OOM at large cadences)."""
+        for s0 in range(0,idx.size(0),chunk):
+            sl=slice(s0,min(s0+chunk,idx.size(0)))
+            x=new_rows[sl].view(-1,self.ncod,G)
+            out=torch.empty(x.size(0),self.ncod,dtype=torch.uint8,device=x.device)
+            for fr in range(4):
+                mask=(self.frames==fr)
+                xf=x[:,mask,:].reshape(-1,G)
+                d2=torch.cdist(xf.bfloat16(),self.codebook[fr].bfloat16())
+                out[:,mask]=d2.argmin(-1).view(x.size(0),-1).to(torch.uint8)
+            self.codons[idx[sl]]=out
+    def decode_rows(self,idx):
+        cod=self.codons[idx].long()                                 # [n,ncod]
+        fr=self.frames.unsqueeze(0).expand_as(cod)                  # [n,ncod]
+        dec=self.codebook[fr,cod].view(-1,self.d).to(getattr(self,'_hdtype',self.codebook.dtype))   # dtype-safe positional decode
+        leaf=dec.detach().clone().requires_grad_(True)
+        self._leaf=leaf; self._leaf_idx=idx
+        return dec + (leaf - leaf.detach())
+    def route(self,h):
+        self._hdtype=h.dtype
+        xa,xb=self.ra(h),self.rb(h); va,ia=xa.topk(2,-1); vb,ib=xb.topk(2,-1)
+        s0=va[:,0]+vb[:,0]; sa=va[:,1]+vb[:,0]; sb=va[:,0]+vb[:,1]; use=sa>=sb
+        k0=ia[:,0]*self.b+ib[:,0]; k1=torch.where(use,ia[:,1]*self.b+ib[:,0],ia[:,0]*self.b+ib[:,1]); s1=torch.where(use,sa,sb)
+        alpha=torch.softmax(torch.stack((s0,s1),1),1).unsqueeze(-1)
+        rows=torch.stack((k0,k1),1)                                 # [N,2]
+        dec=self.decode_rows(rows.reshape(-1)).view(-1,2,self.d)
+        ext=(dec*alpha.to(dec.dtype)).sum(1)
+        ext=self.contract(F.silu(self.expand(ext)).view(ext.size(0),2,self.d).flatten(1))  # v17 implicit expansion
+        code=(self.ca[ia[:,0]]+self.cb[ib[:,0]]).to(xa.dtype)
+        pa,pb=xa.softmax(-1).mean(0),xb.softmax(-1).mean(0); balance=.5*(self.a*pa.square().sum()+self.b*pb.square().sum())
+        return ext+code,balance,rows
+    def features(self,ids):
+        B,T=ids.shape; x=self.embed(ids)
+        for blk in self.blocks: x=blk(x,self.chunk)
+        mem,bal,rows=self.route(x.reshape(B*T,self.d)); return self.norm(x+mem.reshape(B,T,self.d)),bal,rows
+    def config(self): return {k:getattr(self,k) for k in ('vocab','d','layers','ff','a','b')}
+
+def counts(m):
+    ctrl=sum(p.numel() for n,p in m.named_parameters())
+    tbl_bytes=m.nrows*m.ncod  # codon bytes on disk
+    return ctrl,tbl_bytes