| 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 |
|
|
