File size: 2,179 Bytes
80300e5 | 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 | from dataclasses import dataclass, asdict
import torch, torch.nn as nn, torch.nn.functional as F
@dataclass
class EXLLMConfig:
vocab_size:int=1029
max_seq_len:int=128
d_model:int=160
n_layers:int=4
n_heads:int=5
d_ff:int=512
rms_eps:float=1e-5
class RMSNorm(nn.Module):
def __init__(self,dim,eps=1e-5): super().__init__(); self.weight=nn.Parameter(torch.ones(dim)); self.eps=eps
def forward(self,x): return x*torch.rsqrt(x.pow(2).mean(-1,keepdim=True)+self.eps)*self.weight
class Block(nn.Module):
def __init__(self,cfg):
super().__init__(); d=cfg.d_model; self.n_heads=cfg.n_heads; self.head_dim=d//cfg.n_heads
self.norm1=RMSNorm(d,cfg.rms_eps); self.qkv=nn.Linear(d,3*d,bias=False); self.proj=nn.Linear(d,d,bias=False)
self.norm2=RMSNorm(d,cfg.rms_eps); self.fc1=nn.Linear(d,cfg.d_ff,bias=False); self.fc2=nn.Linear(cfg.d_ff,d,bias=False)
def forward(self,x):
b,t,d=x.shape; h=self.norm1(x); qkv=self.qkv(h).view(b,t,3,self.n_heads,self.head_dim).permute(2,0,3,1,4)
q,k,v=qkv[0],qkv[1],qkv[2]; a=F.scaled_dot_product_attention(q,k,v,is_causal=True)
x=x+self.proj(a.transpose(1,2).contiguous().view(b,t,d)); h=self.norm2(x); return x+self.fc2(F.relu(self.fc1(h)))
class EXLLM(nn.Module):
def __init__(self,cfg):
super().__init__(); self.cfg=cfg; self.tok=nn.Embedding(cfg.vocab_size,cfg.d_model); self.pos=nn.Embedding(cfg.max_seq_len,cfg.d_model)
self.blocks=nn.ModuleList([Block(cfg) for _ in range(cfg.n_layers)]); self.norm=RMSNorm(cfg.d_model,cfg.rms_eps); self.lm_head=nn.Linear(cfg.d_model,cfg.vocab_size,bias=False); self.lm_head.weight=self.tok.weight; self.apply(self._init)
def _init(self,m):
if isinstance(m,(nn.Linear,nn.Embedding)): nn.init.normal_(m.weight,0.0,0.02)
def forward(self,idx):
b,t=idx.shape
if t>self.cfg.max_seq_len: raise ValueError('context too long')
p=torch.arange(t,device=idx.device); x=self.tok(idx)+self.pos(p)[None,:,:]
for block in self.blocks: x=block(x)
return self.lm_head(self.norm(x))
def num_parameters(self): return sum(p.numel() for p in self.parameters())
|