File size: 4,593 Bytes
4fd5eb6 | 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 | import math
from typing import Optional, Tuple
import torch
import torch.nn as nn
import torch.nn.functional as F
VOCAB_SIZE = 32_768
MAX_CONTEXT = 2_048
TRAIN_SEQ_LEN = 512
HIDDEN = 896
LAYERS = 20
HEADS = 14
KV_HEADS = 7
INTERMEDIATE = 2_816
ROPE_THETA = 10_000.0
RMS_EPS = 1e-6
def masked_cross_entropy(logits, labels):
valid = labels.ne(-100)
if int(valid.sum().item()) <= 0:
raise RuntimeError('masked_cross_entropy received zero supervised targets')
return F.cross_entropy(logits.float().reshape(-1, VOCAB_SIZE), labels.reshape(-1), ignore_index=-100)
class RMSNorm(nn.Module):
def __init__(self, d=HIDDEN, eps=RMS_EPS):
super().__init__()
self.weight = nn.Parameter(torch.ones(d))
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):
half = x.shape[-1] // 2
return torch.cat((-x[..., half:], x[..., :half]), dim=-1)
class Rotary(nn.Module):
def __init__(self, head_dim, max_seq=MAX_CONTEXT, theta=ROPE_THETA):
super().__init__()
inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2).float() / head_dim))
t = torch.arange(max_seq, dtype=torch.float32)
freqs = torch.outer(t, inv_freq)
emb = torch.cat((freqs, freqs), dim=-1)
self.register_buffer('cos', emb.cos()[None, :, :], persistent=False)
self.register_buffer('sin', emb.sin()[None, :, :], persistent=False)
def forward(self, q, k):
n = q.shape[-2]
cos = self.cos[:, :n].to(q.device, q.dtype)
sin = self.sin[:, :n].to(q.device, q.dtype)
return q * cos + rotate_half(q) * sin, k * cos + rotate_half(k) * sin
class M31Attention(nn.Module):
def __init__(self):
super().__init__()
assert HIDDEN % HEADS == 0
assert HEADS % KV_HEADS == 0
self.head_dim = HIDDEN // HEADS
self.q = nn.Linear(HIDDEN, HIDDEN, bias=False)
self.k = nn.Linear(HIDDEN, KV_HEADS * self.head_dim, bias=False)
self.v = nn.Linear(HIDDEN, KV_HEADS * self.head_dim, bias=False)
self.o = nn.Linear(HIDDEN, HIDDEN, bias=False)
self.rope = Rotary(self.head_dim, MAX_CONTEXT, ROPE_THETA)
def forward(self, x):
b, t, _ = x.shape
q = self.q(x).view(b, t, HEADS, self.head_dim).transpose(1, 2)
k = self.k(x).view(b, t, KV_HEADS, self.head_dim).transpose(1, 2)
v = self.v(x).view(b, t, KV_HEADS, self.head_dim).transpose(1, 2)
repeat = HEADS // KV_HEADS
k = k.repeat_interleave(repeat, dim=1)
v = v.repeat_interleave(repeat, dim=1)
q, k = self.rope(q, k)
y = F.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=True)
return self.o(y.transpose(1, 2).contiguous().view(b, t, HIDDEN))
class M31Block(nn.Module):
def __init__(self):
super().__init__()
self.n1 = RMSNorm(HIDDEN, RMS_EPS)
self.attn = M31Attention()
self.n2 = RMSNorm(HIDDEN, RMS_EPS)
self.gate = nn.Linear(HIDDEN, INTERMEDIATE, bias=False)
self.up = nn.Linear(HIDDEN, INTERMEDIATE, bias=False)
self.down = nn.Linear(INTERMEDIATE, HIDDEN, bias=False)
def forward(self, x):
x = x + self.attn(self.n1(x))
h = self.n2(x)
h = F.silu(self.gate(h)) * self.up(h)
return x + self.down(h)
class M31Model(nn.Module):
def __init__(self):
super().__init__()
self.embed = nn.Embedding(VOCAB_SIZE, HIDDEN)
self.blocks = nn.ModuleList([M31Block() for _ in range(LAYERS)])
self.norm = RMSNorm(HIDDEN, RMS_EPS)
self.apply(self._init_weights)
nn.init.normal_(self.embed.weight, mean=0.0, std=0.02)
self.num_parameters = sum(p.numel() for p in self.parameters())
if self.num_parameters >= 250_000_000:
raise RuntimeError('Hard parameter ceiling violated.')
@staticmethod
def _init_weights(m):
if isinstance(m, nn.Linear):
nn.init.normal_(m.weight, mean=0.0, std=0.02 / math.sqrt(2 * LAYERS))
elif isinstance(m, nn.Embedding):
pass
def forward(self, input_ids, labels: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
x = self.embed(input_ids)
for block in self.blocks:
x = block(x)
x = self.norm(x)
logits = F.linear(x, self.embed.weight)
loss = masked_cross_entropy(logits, labels) if labels is not None else None
return logits, loss
|