| """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 |
| |
| 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: |
| |
| y = F.scaled_dot_product_attention(q,kr,vr,dropout_p=0.0,is_causal=False) |
| else: |
| |
| 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 |
| 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()) |
|
|
