"""Small Llama-family decoder architecture, implemented with standard PyTorch. No pretrained weights or nonstandard attention algorithms are used. """ from __future__ import annotations from dataclasses import asdict, dataclass import math import torch from torch import nn from torch.nn import functional as F from torch.utils.checkpoint import checkpoint @dataclass(frozen=True) class ModelConfig: vocab_size: int = 32000 hidden_size: int = 1024 num_layers: int = 12 num_heads: int = 16 num_kv_heads: int = 4 intermediate_size: int = 2560 max_position_embeddings: int = 2048 rope_theta: float = 10000.0 rms_norm_eps: float = 1e-5 def validate(self): assert self.hidden_size % self.num_heads == 0 assert self.num_heads % self.num_kv_heads == 0 assert (self.hidden_size // self.num_heads) % 2 == 0 assert self.vocab_size <= 65535, 'u16 packed datasets require <=65535 tokens' assert self.max_position_embeddings > 0 return self def to_dict(self): return asdict(self) class RMSNorm(nn.Module): def __init__(self, width: int, eps: float): super().__init__() self.weight = nn.Parameter(torch.ones(width)) self.eps = eps def forward(self, x): dtype = x.dtype z = x.float() return (z * torch.rsqrt(z.square().mean(-1, keepdim=True) + self.eps)).to(dtype) * self.weight def apply_rope(x, positions, theta): """Rotate pairs of head features, standard RoPE (not interleaved sin/cos).""" dim = x.shape[-1] scale = torch.arange(0, dim, 2, device=x.device, dtype=torch.float32) / dim inv_freq = theta ** (-scale) angles = torch.outer(positions.to(torch.float32), inv_freq) c, s = angles.cos()[None, None, :, :], angles.sin()[None, None, :, :] a, b = x.float()[..., :dim//2], x.float()[..., dim//2:] return torch.cat((a*c - b*s, b*c + a*s), dim=-1).to(dtype=x.dtype) class Attention(nn.Module): def __init__(self, cfg: ModelConfig): super().__init__() self.nh = cfg.num_heads self.nkv = cfg.num_kv_heads self.hd = cfg.hidden_size // cfg.num_heads self.theta = cfg.rope_theta self.q = nn.Linear(cfg.hidden_size, self.nh*self.hd, bias=False) self.k = nn.Linear(cfg.hidden_size, self.nkv*self.hd, bias=False) self.v = nn.Linear(cfg.hidden_size, self.nkv*self.hd, bias=False) self.o = nn.Linear(cfg.hidden_size, cfg.hidden_size, bias=False) def forward(self, x, past=None, use_cache=False): b, s, _ = x.shape offset = 0 if past is None else past[0].size(2) pos = torch.arange(offset, offset+s, device=x.device) q = self.q(x).view(b,s,self.nh,self.hd).transpose(1,2) k = self.k(x).view(b,s,self.nkv,self.hd).transpose(1,2) v = self.v(x).view(b,s,self.nkv,self.hd).transpose(1,2) q = apply_rope(q, pos, self.theta) k = apply_rope(k, pos, self.theta) if past is not None: k = torch.cat((past[0], k), dim=2) v = torch.cat((past[1], v), dim=2) new_cache = (k,v) if use_cache else None # GQA via head repetition is portable to Windows/PyTorch SDPA backends. repeats = self.nh // self.nkv kr = k.repeat_interleave(repeats, dim=1) vr = v.repeat_interleave(repeats, dim=1) if past is None: y = F.scaled_dot_product_attention(q,kr,vr,dropout_p=0.0,is_causal=True) elif s == 1: # One new token can attend to every cached key. y = F.scaled_dot_product_attention(q,kr,vr,dropout_p=0.0,is_causal=False) else: # For chunked cached decoding, q[i] can only see keys through offset+i. allowed = torch.arange(k.size(2),device=x.device)[None,:] <= (offset+torch.arange(s,device=x.device)[:,None]) y = F.scaled_dot_product_attention(q,kr,vr,attn_mask=allowed,dropout_p=0.0,is_causal=False) return self.o(y.transpose(1,2).contiguous().view(b,s,-1)), new_cache class Block(nn.Module): def __init__(self,cfg): super().__init__() self.attn_norm = RMSNorm(cfg.hidden_size,cfg.rms_norm_eps) self.attn = Attention(cfg) self.ffn_norm = RMSNorm(cfg.hidden_size,cfg.rms_norm_eps) self.gate = nn.Linear(cfg.hidden_size,cfg.intermediate_size,bias=False) self.up = nn.Linear(cfg.hidden_size,cfg.intermediate_size,bias=False) self.down = nn.Linear(cfg.intermediate_size,cfg.hidden_size,bias=False) def forward(self,x,past=None,use_cache=False): a,new_past = self.attn(self.attn_norm(x),past,use_cache) x = x + a t = self.ffn_norm(x) x = x + self.down(F.silu(self.gate(t)) * self.up(t)) return x,new_past class Velqora(nn.Module): def __init__(self,cfg:ModelConfig): super().__init__() self.cfg = cfg.validate() self.embed = nn.Embedding(cfg.vocab_size,cfg.hidden_size) self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.num_layers)]) self.norm = RMSNorm(cfg.hidden_size,cfg.rms_norm_eps) self.head = nn.Linear(cfg.hidden_size,cfg.vocab_size,bias=False) self.head.weight = self.embed.weight # tied embeddings self.grad_checkpoint = False self.apply(self._init_weights) @staticmethod def _init_weights(module): if isinstance(module,(nn.Linear,nn.Embedding)): nn.init.normal_(module.weight,mean=0.0,std=0.02) def forward(self,input_ids,labels=None,past=None,use_cache=False): if input_ids.ndim!=2: raise ValueError('input_ids must have shape [batch, seq]') past_len = 0 if past is None else past[0][0].size(2) if input_ids.size(1)+past_len > self.cfg.max_position_embeddings: raise ValueError('Context longer than configured model maximum') x = self.embed(input_ids) new_past = [] for i,block in enumerate(self.blocks): previous = None if past is None else past[i] if self.grad_checkpoint and self.training and not use_cache: x = checkpoint(lambda z, layer=block: layer(z)[0],x,use_reentrant=False) else: x,cache = block(x,previous,use_cache) if use_cache: new_past.append(cache) logits = self.head(self.norm(x)) loss = None if labels is not None: loss = F.cross_entropy(logits.float().reshape(-1,logits.size(-1)), labels.reshape(-1),ignore_index=-100) return logits,loss,(new_past if use_cache else None) @torch.inference_mode() def generate_tokens(self,prompt,max_new_tokens=150,temperature=0.8,top_p=0.9,eot_id=None,seed=None): self.eval() if seed is not None: torch.manual_seed(seed) if prompt.is_cuda: torch.cuda.manual_seed_all(seed) prompt = prompt[:, -self.cfg.max_position_embeddings:] logits,_,cache = self(prompt,use_cache=True) generated=[] for _ in range(min(max_new_tokens,self.cfg.max_position_embeddings-prompt.size(1))): score=logits[:,-1,:].float() if temperature <= 0: nxt=score.argmax(dim=-1,keepdim=True) else: score/=temperature probs=torch.softmax(score,dim=-1) if top_p < 1.0: sorted_probs, sorted_idx = torch.sort(probs,descending=True) mask=(sorted_probs.cumsum(-1)-sorted_probs)>top_p sorted_probs=sorted_probs.masked_fill(mask,0) probs=torch.zeros_like(probs).scatter_(-1,sorted_idx,sorted_probs) probs=probs/probs.sum(-1,keepdim=True) nxt=torch.multinomial(probs,num_samples=1) token=int(nxt.item()); generated.append(token) if eot_id is not None and token==eot_id: break logits,_,cache=self(nxt,past=cache,use_cache=True) return generated def parameter_count(model): return sum(p.numel() for p in model.parameters())