File size: 3,622 Bytes
b1fd620 | 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 | import torch
import torch.nn as nn
import torch.nn.functional as F
from tokenizers import Tokenizer
from safetensors.torch import load_model
class RMSNorm(nn.Module):
def __init__(self, dim, eps=1e-6):
super().__init__()
self.weight = nn.Parameter(torch.ones(dim))
self.eps = eps
def forward(self, x):
return x * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps).to(x.dtype) * self.weight
def rotate_half(x):
a, b = x.chunk(2, dim=-1)
return torch.cat((-b, a), dim=-1)
class M31Attention(nn.Module):
def __init__(self, c):
super().__init__()
h = c['hidden_size']; nh = c['num_attention_heads']; nk = c['num_key_value_heads']; d = h // nh
self.h, self.nh, self.nk, self.d = h, nh, nk, d
self.q = nn.Linear(h, nh*d, bias=False)
self.k = nn.Linear(h, nk*d, bias=False)
self.v = nn.Linear(h, nk*d, bias=False)
self.o = nn.Linear(h, h, bias=False)
inv = 1.0 / (c['rope_theta'] ** (torch.arange(0, d, 2).float() / d))
pos = torch.arange(c['max_position_embeddings'], dtype=torch.float)
f = torch.outer(pos, inv); e = torch.cat([f, f], dim=-1)
self.register_buffer('cos', e.cos()[None, None], persistent=False)
self.register_buffer('sin', e.sin()[None, None], persistent=False)
def forward(self, x):
b, s, _ = x.shape
q = self.q(x).view(b, s, self.nh, self.d).transpose(1, 2)
k = self.k(x).view(b, s, self.nk, self.d).transpose(1, 2)
v = self.v(x).view(b, s, self.nk, self.d).transpose(1, 2)
co = self.cos[:, :, :s, :].to(x.device, x.dtype); si = self.sin[:, :, :s, :].to(x.device, x.dtype)
q = q * co + rotate_half(q) * si; k = k * co + rotate_half(k) * si
try:
y = F.scaled_dot_product_attention(q, k, v, is_causal=True, enable_gqa=True)
except Exception:
rep = self.nh // self.nk
y = F.scaled_dot_product_attention(q, k.repeat_interleave(rep, 1), v.repeat_interleave(rep, 1), is_causal=True)
return self.o(y.transpose(1, 2).contiguous().view(b, s, self.h))
class M31MLP(nn.Module):
def __init__(self, c):
super().__init__(); h=c['hidden_size']; i=c['intermediate_size']
self.gate=nn.Linear(h,i,bias=False); self.up=nn.Linear(h,i,bias=False); self.down=nn.Linear(i,h,bias=False)
def forward(self,x): return self.down(F.silu(self.gate(x))*self.up(x))
class M31Block(nn.Module):
def __init__(self,c):
super().__init__(); h=c['hidden_size']
self.n1=RMSNorm(h,c['rms_norm_eps']); self.attn=M31Attention(c); self.n2=RMSNorm(h,c['rms_norm_eps']); self.mlp=M31MLP(c)
def forward(self,x):
x = x + self.attn(self.n1(x))
x = x + self.mlp(self.n2(x))
return x
class M31ForCausalLM(nn.Module):
def __init__(self,c):
super().__init__(); h=c['hidden_size']
self.embed=nn.Embedding(c['vocab_size'],h)
self.blocks=nn.ModuleList([M31Block(c) for _ in range(c['num_hidden_layers'])])
self.norm=RMSNorm(h,c['rms_norm_eps'])
self.lm_head=nn.Linear(h,c['vocab_size'],bias=False)
self.lm_head.weight=self.embed.weight
def forward(self,input_ids):
x=self.embed(input_ids)
for b in self.blocks: x=b(x)
return self.lm_head(self.norm(x))
def load_model_bundle(path='.'):
import json
from pathlib import Path
p=Path(path); c=json.loads((p/'config.json').read_text()); tok=Tokenizer.from_file(str(p/'tokenizer.json'))
m=M31ForCausalLM(c); load_model(m,str(p/'model.safetensors'),strict=True); return m,tok
|