File size: 5,294 Bytes
8561bc9 | 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 | # Model definition for Panhapich/pre-train-llama.
#
# Llama-style decoder-only transformer (RoPE, grouped-query attention, SwiGLU,
# RMSNorm), trained from scratch. This is not a `transformers`-library model
# class -- load model.safetensors into TextGenerationModel directly:
#
# import json
# from safetensors.torch import load_file
# from modeling_llama_custom import TextGenerationModel
#
# config = json.load(open("config.json"))
# model = TextGenerationModel(**config["model_config"])
# model.load_state_dict(load_file("model.safetensors"))
# model.eval()
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
if not hasattr(nn, "RMSNorm"):
class _RMSNormFallback(nn.Module):
def __init__(self, dim, eps=1e-6):
super().__init__()
self.eps = eps
self.weight = nn.Parameter(torch.ones(dim))
def forward(self, x):
rms = x.pow(2).mean(dim=-1, keepdim=True).add(self.eps).rsqrt()
return x * rms * self.weight
nn.RMSNorm = _RMSNormFallback
class RotaryPositionalEncoding(nn.Module):
def __init__(self, head_dim, max_seq_len, theta=10000.0):
super().__init__()
inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2).float() / head_dim))
t = torch.arange(max_seq_len).float()
freqs = torch.outer(t, inv_freq)
self.register_buffer("cos", torch.cos(freqs), persistent=False)
self.register_buffer("sin", torch.sin(freqs), persistent=False)
def rotate(self, x):
T = x.shape[-2]
cos = self.cos[:T].unsqueeze(0).unsqueeze(0)
sin = self.sin[:T].unsqueeze(0).unsqueeze(0)
x1, x2 = x[..., 0::2], x[..., 1::2]
rotated = torch.stack([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
return rotated.flatten(-2)
class GQA(nn.Module):
def __init__(self, hidden_dim, num_heads, num_kv_heads, dropout=0.1):
super().__init__()
self.num_heads = num_heads
self.num_kv_heads = num_kv_heads
self.n_rep = num_heads // num_kv_heads
self.head_dim = hidden_dim // num_heads
self.q_proj = nn.Linear(hidden_dim, num_heads * self.head_dim)
self.k_proj = nn.Linear(hidden_dim, num_kv_heads * self.head_dim)
self.v_proj = nn.Linear(hidden_dim, num_kv_heads * self.head_dim)
self.out_proj = nn.Linear(num_heads * self.head_dim, hidden_dim)
self.dropout = nn.Dropout(dropout)
def forward(self, q, k, v, mask=None, rope=None):
B, T, _ = q.shape
q = self.q_proj(q).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
k = self.k_proj(k).view(B, T, self.num_kv_heads, self.head_dim).transpose(1, 2)
v = self.v_proj(v).view(B, T, self.num_kv_heads, self.head_dim).transpose(1, 2)
if rope is not None:
q = rope.rotate(q)
k = rope.rotate(k)
if self.n_rep > 1:
k = k.repeat_interleave(self.n_rep, dim=1)
v = v.repeat_interleave(self.n_rep, dim=1)
scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
if mask is not None:
scores = scores.masked_fill(~mask.unsqueeze(1).bool(), float('-inf'))
attn = F.softmax(scores, dim=-1)
attn = self.dropout(attn)
out = attn @ v
out = out.transpose(1, 2).reshape(B, T, -1)
return self.out_proj(out)
class SwiGLU(nn.Module):
def __init__(self, hidden_dim, ff_dim):
super().__init__()
self.gate_proj = nn.Linear(hidden_dim, ff_dim)
self.up_proj = nn.Linear(hidden_dim, ff_dim)
self.down_proj = nn.Linear(ff_dim, hidden_dim)
def forward(self, x):
return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
class DecoderLayer(nn.Module):
def __init__(self, hidden_dim, num_heads, num_kv_heads, dropout=0.1):
super().__init__()
self.self_attn = GQA(hidden_dim, num_heads, num_kv_heads, dropout)
self.mlp = SwiGLU(hidden_dim, 4 * hidden_dim)
self.norm1 = nn.RMSNorm(hidden_dim)
self.norm2 = nn.RMSNorm(hidden_dim)
def forward(self, x, mask=None, rope=None):
out = self.norm1(x)
out = self.self_attn(out, out, out, mask, rope)
x = out + x
out = self.norm2(x)
out = self.mlp(out)
return out + x
class TextGenerationModel(nn.Module):
def __init__(self, num_layers, num_heads, num_kv_heads, hidden_dim,
max_seq_len, vocab_size, dropout=0.1):
super().__init__()
self.rope = RotaryPositionalEncoding(hidden_dim // num_heads, max_seq_len)
self.embedding = nn.Embedding(vocab_size, hidden_dim)
self.decoders = nn.ModuleList([
DecoderLayer(hidden_dim, num_heads, num_kv_heads, dropout)
for _ in range(num_layers)
])
self.norm = nn.RMSNorm(hidden_dim)
self.out = nn.Linear(hidden_dim, vocab_size)
def forward(self, ids, mask=None):
x = self.embedding(ids)
for decoder in self.decoders:
x = decoder(x, mask, self.rope)
x = self.norm(x)
return self.out(x)
def create_causal_mask(seq_len, device):
return torch.tril(torch.ones(seq_len, seq_len, dtype=torch.bool, device=device))
|