File size: 6,162 Bytes
a962f7c | 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 | import torch.nn as nn
import torch
import torch.nn.functional as F
import math
config_path = r"C:\Users\mehdi\Desktop\Pythonfiles\Projects\Transformers\Alpha\config.json"
class RoPE(nn.Module):
def __init__(self, dim, max_seq_len=2048):
super().__init__()
theta = 10000
freqs = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
positions = torch.arange(max_seq_len)
angles = torch.outer(positions, freqs)
self.register_buffer("cos", angles.cos(), persistent=False)
self.register_buffer("sin", angles.sin(), persistent=False)
def rotate_half(self, x):
x1 = x[..., ::2]
x2 = x[..., 1::2]
return torch.stack((-x2,x1),dim=-1).flatten(-2)
def forward(self, x, start_pos=0):
seq_len = x.shape[-2]
assert start_pos + seq_len <= self.cos.size(0)
cos = self.cos[start_pos : start_pos + seq_len] # We slice from start pos of the inf token, this way the previous Qs and Ks aren't computed
sin = self.sin[start_pos : start_pos + seq_len] # We slice from start pos of the inf token, this way the previous Qs and Ks aren't computed
cos = torch.repeat_interleave(cos, 2, dim=-1)
sin = torch.repeat_interleave(sin, 2, dim=-1)
return x*cos + self.rotate_half(x)*sin
class GQA(nn.Module):
def __init__(self, d_model, num_heads, num_kv_heads):
super().__init__()
#self.attn_dropout_p = dropout
#self.resid_dropout = nn.Dropout(dropout)
self.num_heads = num_heads
self.d_model = d_model
self.num_kv_heads = num_kv_heads
self.W_o = nn.Linear(d_model, d_model)
self.d_k = d_model // num_heads
self.Q = nn.Linear(d_model, d_model)
self.K = nn.Linear(d_model, num_kv_heads * self.d_k)
self.V = nn.Linear(d_model, num_kv_heads * self.d_k)
self.rope = RoPE(self.d_k)
def forward(self, x, cache=None, start_pos=0, use_rope = True,):
batch_size = x.size(0)
Q = self.Q(x)
K = self.K(x)
V = self.V(x)
seq_length = x.shape[1]
Q = Q.view(batch_size, seq_length, self.num_heads, self.d_k)
K = K.view(batch_size, seq_length, self.num_kv_heads, self.d_k)
V = V.view(batch_size, seq_length, self.num_kv_heads, self.d_k)
Q = Q.transpose(1,2)
K = K.transpose(1,2)
V = V.transpose(1,2)
if use_rope:
Q = self.rope(Q, start_pos)
K = self.rope(K, start_pos)
if cache is not None:
K = torch.cat([cache[0], K], dim = 2)
V = torch.cat([cache[1], V], dim = 2)
new_cache = K, V
Q_len = Q.size(2)
K_len = K.size(2)
if Q_len == K_len:
output = F.scaled_dot_product_attention(Q, K, V, is_causal=True, enable_gqa=True)
elif Q_len == 1:
output = F.scaled_dot_product_attention(Q, K, V, enable_gqa=True)
else: # multi-token chunk on top of a cache
mask = torch.ones(Q_len, K_len, dtype=torch.bool, device=Q.device).tril(K_len - Q_len)
output = F.scaled_dot_product_attention(Q, K, V, attn_mask=mask, enable_gqa=True)
output = output.transpose(1, 2).contiguous()
output = output.view(batch_size, seq_length, self.d_model)
output = self.W_o(output)
return output, new_cache
class MLP(nn.Module):
def __init__(self, d_model, d_ff):
super().__init__()
hidden = int(2 * d_ff / 3)
hidden = (hidden + 63) // 64 * 64 # 4096 -> 2752
self.w_gate = nn.Linear(d_model, hidden, bias=False)
self.w_up = nn.Linear(d_model, hidden, bias=False)
self.w_down = nn.Linear(hidden, d_model, bias=False)
#self.dropout = nn.Dropout(dropout)
def forward(self, x):
return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x))
class DecoderBlock(nn.Module):
def __init__(self, d_model, num_heads, d_ff, num_kv_heads):
super().__init__()
self.self_attn = GQA(d_model, num_heads, num_kv_heads)
self.norm1 = nn.LayerNorm(d_model)
self.norm2 = nn.LayerNorm(d_model)
self.mlp = MLP(d_model, d_ff)
def forward(self, x, cache=None, start_pos=0):
normed = self.norm1(x)
attn_out, new_cache = self.self_attn(normed, cache=cache, start_pos=start_pos)
x = x + attn_out
mlp_out = self.mlp(self.norm2(x))
x = x + mlp_out
return x, new_cache
class Transformer(nn.Module):
def __init__(self, d_model, num_heads, vocab_size, d_ff, num_kv_heads, num_decoder_layers):
super().__init__()
self.embedding = nn.Embedding(vocab_size, d_model)
self.decoder = nn.ModuleList([DecoderBlock(d_model, num_heads,
d_ff, num_kv_heads) for _ in range(num_decoder_layers)])
self.final_norm = nn.LayerNorm(d_model)
self.output_proj = nn.Linear(d_model, vocab_size)
for m in self.modules():
if isinstance(m, (nn.Linear, nn.Embedding)):
nn.init.normal_(m.weight, std=0.02)
if isinstance(m, nn.Linear) and m.bias is not None:
nn.init.zeros_(m.bias)
for n, p in self.named_parameters():
if n.endswith("W_o.weight") or n.endswith("w_down.weight"):
nn.init.normal_(p, std=0.02 / math.sqrt(2 * len(self.decoder)))
def forward(self, tokens, caches= None, start_pos=0, return_cache=False):
x = self.embedding(tokens)
new_caches = []
for i, block in enumerate(self.decoder):
x, c = block(x, cache=None if caches is None else caches[i], start_pos=start_pos)
new_caches.append(c)
logits = self.output_proj(self.final_norm(x))
return (logits, new_caches) if return_cache else logits |