training_reference/dense.py

from dataclasses import dataclass, asdict
import hashlib
import math
import numpy as np
import torch
from torch import nn
from torch.nn import functional as F

@dataclass(frozen=True)
class Config:
    arm: str = 'mha_gated_ffn_router'
    vocab: int = 16384
    width: int = 256
    layers: int = 7
    hidden: int = 896
    q_heads: int = 8
    kv_heads: int = 8
    head_dim: int = 32
    eps: float = 1e-6
    seed: int = 20260907
    backend: str = 'sdpa'
    dropout: float = 0.
    value_variance: bool = True
    variance_delivery: str = 'full_width_direct'
    def validate(self):
        assert self.arm == 'mha_gated_ffn_router'
        assert self.backend in ('sdpa','flash')
        assert self.width==self.q_heads*self.head_dim and self.q_heads==self.kv_heads
        assert self.head_dim%2==0 and self.layers>0
        assert 0<=self.dropout<1
        assert self.variance_delivery=='full_width_direct'

def sphere(x,eps=1e-6):
    dtype=x.dtype
    y=x.double() if dtype==torch.float64 else x.float()
    return (y*torch.rsqrt(y.square().mean(-1,keepdim=True)+eps)).to(dtype)

@dataclass
class Layout:
    mask: torch.Tensor | None
    positions: torch.Tensor
    indices: torch.Tensor
    cu: torch.Tensor
    max_length: int

def prepare_layout(segments,device,backend='sdpa'):
    """CPU metadata boundary: positive contiguous document IDs, padding0."""
    if segments.device.type!='cpu':raise ValueError('prepare metadata on CPU before transfer')
    if segments.ndim!=2 or bool((segments<0).any()):raise ValueError('bad segment IDs')
    b,t=segments.shape
    # Vectorized host metadata; no per-token Python loop or device synchronization.
    raw=segments.numpy()
    valid=raw>0
    starts=valid.copy();starts[:,1:] &= raw[:,1:]!=raw[:,:-1]
    cols=np.arange(t)[None,:]
    origins=np.maximum.accumulate(np.where(starts,cols,0),axis=1)
    positions=torch.from_numpy(np.where(valid,cols-origins,0).astype(np.int64))
    indices=np.flatnonzero(valid.reshape(-1))
    start_offsets=np.flatnonzero(starts.reshape(-1)[indices])
    lengths=np.diff(np.append(start_offsets,len(indices)))
    mask=None
    if backend=='sdpa':
        causal=torch.arange(t)[None,:]<=torch.arange(t)[:,None]
        mask=((segments[:,:,None]==segments[:,None,:]) & (segments[:,:,None]>0) & causal).unsqueeze(1).to(device)
    cu=torch.from_numpy(np.concatenate(([0],np.cumsum(lengths))).astype(np.int32))
    return Layout(mask,positions.to(device),torch.tensor(indices,dtype=torch.long,device=device),cu.to(device),int(max(lengths,default=0)))

class Norm(nn.Module):
    def __init__(self,d,eps):
        super().__init__();self.weight=nn.Parameter(torch.ones(d));self.eps=eps
    def forward(self,x):return sphere(x,self.eps)*self.weight.to(x.dtype)

class Router(nn.Module):
    def __init__(self,c):
        super().__init__();self.eps=c.eps
        self.query=nn.Linear(c.width,16,bias=False)
        self.key=nn.Linear(c.width,16,bias=False)
    def forward(self,current,s1,s2):
        q=self.query(sphere(current,self.eps)).float()
        k1,k2=self.key(s1).float(),self.key(s2).float()
        scores=torch.stack(((q*k1).sum(-1),(q*k2).sum(-1)),dim=-1)/4.
        return scores.softmax(-1)

class DeltaReader(nn.Module):
    def __init__(self,c,out_width):
        super().__init__()
        self.aux1=nn.Linear(c.width,out_width,bias=False)
        self.aux2=nn.Linear(c.width,out_width,bias=False)
        self.router=Router(c)
    def forward(self,current,s1,s2):
        v1,v2=self.aux1(s1),self.aux2(s2)
        p=self.router(current,s1,s2).to(v1.dtype)
        return p[...,0:1]*v1+p[...,1:2]*v2

def spread_features(variance,eps=1e-6):
    """Separate direction and smooth magnitude; zero spread maps exactly to zero."""
    work=variance.double() if variance.dtype==torch.float64 else variance.float()
    r=torch.log1p(work/0.1)
    mean_square=r.square().mean(-1,keepdim=True)
    energy=mean_square+eps
    inverse=torch.rsqrt(energy)
    direction=r*inverse
    # Rationalized difference avoids cancellation and is exactly zero at r=0.
    magnitude=torch.log1p(mean_square/((energy.sqrt()+math.sqrt(eps))*(1+math.sqrt(eps))))
    return torch.cat((direction,magnitude),-1)

class FFN(nn.Module):
    def __init__(self,c,enabled=True):
        super().__init__();self.c=c
        self.base=nn.Linear(c.width,2*c.hidden,bias=False)
        self.down=nn.Linear(c.hidden,c.width,bias=False)
        self.reader=DeltaReader(c,2*c.hidden) if enabled else None
        self.spread_input=nn.Linear(c.width+1,2*c.hidden,bias=False) if c.value_variance else None
    def forward(self,x,current,s1,s2,variance=None):
        pre=self.base(x)
        if self.reader is not None:pre=pre+self.reader(current,s1,s2)
        if self.spread_input is not None:
            pre=pre+self.spread_input(spread_features(variance,self.c.eps).to(x.dtype))
        g,u=pre.chunk(2,-1)
        return self.down(F.dropout(F.silu(g)*u,self.c.dropout,self.training))

class AttentionGate(nn.Module):
    """Head-local elementwise write gate, kept as a compilable fusion boundary."""
    def __init__(self,c):
        super().__init__();self.c=c
        self.proj=nn.Linear(c.width,c.q_heads*c.head_dim,bias=False)
    def forward(self,out,xq):
        b,t,h,d=out.shape
        gate=torch.sigmoid(self.proj(xq).view(b,t,h,d))
        return out*gate

class Block(nn.Module):
    def __init__(self,c,index):
        super().__init__();self.c=c
        self.an=Norm(c.width,c.eps);self.fn=Norm(c.width,c.eps)
        self.qk=nn.Linear(c.width,(c.q_heads+c.kv_heads)*c.head_dim,bias=False)
        self.v=nn.Linear(c.width,c.kv_heads*c.head_dim,bias=False)
        # Head-specific elementwise gate from the same pre-normalized query state.
        # Applied after SDPA and before heads are concatenated/output-projected.
        self.attn_gate=AttentionGate(c)
        self.out=nn.Linear(c.width,c.width,bias=False)
        self.ffn=FFN(c,index>0)
    def project_attention(self,h,cos,sin):
        c=self.c;b,t,d=h.shape
        xq=xv=self.an(h)
        q,k=self.qk(xq).split((c.q_heads*c.head_dim,c.kv_heads*c.head_dim),-1)
        q=q.view(b,t,c.q_heads,c.head_dim);k=k.view(b,t,c.kv_heads,c.head_dim)
        v=self.v(xv).view(b,t,c.kv_heads,c.head_dim)
        def rope(x):
            a,z=x.chunk(2,-1)
            co,si=cos.to(x.dtype),sin.to(x.dtype)
            return torch.cat((a*co-z*si,z*co+a*si),-1)
        q,k=rope(q),rope(k)
        if c.value_variance:
            q,k=F.pad(q,(0,c.head_dim)),F.pad(k,(0,c.head_dim))
            # Keep double for the high precision CPU reference.
            vv=v.double() if v.dtype==torch.float64 else v.float()
            v=torch.cat((v,vv.square().to(v.dtype)),-1)
        return q,k,v,xq
    def finish_attention(self,h,paired,xq):
        c=self.c;b,t,d=h.shape
        if c.value_variance:
            out,second=paired.split(c.head_dim,-1)
            mean=out.double() if out.dtype==torch.float64 else out.float()
            moment=second.double() if second.dtype==torch.float64 else second.float()
            variance=(moment-mean.square()).clamp_min(0).reshape(b,t,c.width)
        else:out=paired;variance=None
        out=self.attn_gate(out,xq).reshape(b,t,d)
        a=sphere(h+self.out(out),c.eps)
        return a,self.fn(a),variance
    def forward(self,h,history,layout,cos,sin):
        c=self.c;b,t,d=h.shape
        q,k,v,xq=self.project_attention(h,cos,sin)
        attention_width=c.head_dim*(2 if c.value_variance else 1)
        if c.backend=='flash':
            from flash_attn import flash_attn_varlen_func
            idx=layout.indices
            qp=q.reshape(-1,c.q_heads,attention_width)[idx];kp=k.reshape(-1,c.kv_heads,attention_width)[idx];vp=v.reshape(-1,c.kv_heads,attention_width)[idx]
            y=flash_attn_varlen_func(qp,kp,vp,layout.cu,layout.cu,layout.max_length,layout.max_length,softmax_scale=c.head_dim**-.5,dropout_p=0.,causal=True,deterministic=True)
            # Scatter the paired moments once; calculate variance after unpacking.
            paired=torch.zeros_like(q.reshape(-1,c.q_heads,attention_width)).index_copy(0,idx,y).view(b,t,c.q_heads,attention_width)
        else:
            paired=F.scaled_dot_product_attention(q.transpose(1,2),k.transpose(1,2),v.transpose(1,2),attn_mask=layout.mask,dropout_p=0.,enable_gqa=False,scale=c.head_dim**-.5).transpose(1,2)
        a,normalized,variance=self.finish_attention(h,paired,xq)
        if variance is not None:
            # A causal document's first token has exactly one Value: variance is
            # identically zero. Do not normalize BF16 moment cancellation noise.
            variance=torch.where((layout.positions>0).unsqueeze(-1),variance,0.)
        if history is None:s1=s2=torch.zeros_like(a)
        else:s1,s2=history
        y=sphere(a+self.ffn(normalized,a,s1,s2,variance),c.eps)
        return y,a

class DenseLM(nn.Module):
    def __init__(self,c):
        super().__init__();c.validate();self.cfg=c
        self.embed=nn.Embedding(c.vocab,c.width)
        self.blocks=nn.ModuleList([Block(c,i) for i in range(c.layers)])
        self.norm=Norm(c.width,c.eps);self.lm_head=nn.Linear(c.width,c.vocab,bias=False)
        self.register_buffer('inv_freq',10000.**(-torch.arange(0,c.head_dim,2).float()/c.head_dim),persistent=True)
        for name,p in self.named_parameters():
            if p.is_meta:continue
            if name.endswith('spread_input.weight'):nn.init.zeros_(p);continue
            if p.ndim==1:nn.init.ones_(p);continue
            seed=(int.from_bytes(hashlib.sha256(name.encode()).digest()[:8],'little')+c.seed)%(2**63-1)
            gen=torch.Generator(device=p.device).manual_seed(seed)
            scale=.02/math.sqrt(2*c.layers) if name.endswith(('down.weight','out.weight')) else .02
            nn.init.normal_(p,std=scale,generator=gen)
    def hidden(self,ids,layout):
        h=sphere(self.embed(ids),self.cfg.eps)
        # Keep residual states and their differences in FP32 to avoid cancellation
        # before RMS normalization; autocast still handles matrix multiplications.
        history=None
        angles=layout.positions[...,None].float()*self.inv_freq
        cos,sin=angles.cos().to(h.dtype).unsqueeze(2),angles.sin().to(h.dtype).unsqueeze(2)
        for i,block in enumerate(self.blocks):
            before=h;h,attention_after=block(h,history,layout,cos,sin)
            if i+1<len(self.blocks):
                d1,d2=attention_after-before,h-attention_after
                history=(sphere(d1,self.cfg.eps),sphere(d2,self.cfg.eps))
        return self.norm(h)
    def forward(self,ids,layout,labels=None):
        logits=self.lm_head(self.hidden(ids,layout))
        if labels is None:return logits
        # labels are explicitly pre-shifted; padding and document boundaries=-100.
        loss=F.cross_entropy(logits.float().flatten(0,1),labels.flatten(),ignore_index=-100,reduction='sum')
        return loss/(labels!=-100).sum().clamp_min(1)

def optimizer(model):
    return torch.optim.AdamW([{'params':[p for p in model.parameters() if p.ndim>=2],'weight_decay':.1},{'params':[p for p in model.parameters() if p.ndim<2],'weight_decay':0.}],lr=8e-4,betas=(.9,.95),eps=1e-8)