File size: 8,123 Bytes
c9721e1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 | """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())
|