File size: 4,584 Bytes
52135a8 | 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 | """Quark-50M model definition — standalone."""
import torch
import torch.nn as nn
import torch.nn.functional as F
from dataclasses import dataclass
@dataclass
class ModelConfig:
vocab_size: int = 16384; d_model: int = 512; n_heads: int = 8
n_kv_heads: int = 4; n_layers: int = 12; d_ff: int = 1408
head_dim: int = 64; max_seq_len: int = 2048; rope_theta: float = 10000.0
rms_eps: float = 1e-5; qkv_bias: bool = False; dropout: float = 0.0
class RMSNorm(nn.Module):
def __init__(self, dim, eps=1e-5):
super().__init__(); self.eps = eps; self.scale = nn.Parameter(torch.ones(dim))
def forward(self, x):
return (x.float() * x.float().pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()).to(x.dtype) * self.scale
class RotaryEmbedding(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))
self.register_buffer("inv_freq", inv_freq, persistent=False); self._build(max_seq_len)
def _build(self, seq_len):
t = torch.arange(seq_len, device=self.inv_freq.device).float()
freqs = torch.outer(t, self.inv_freq); emb = torch.cat([freqs, freqs], dim=-1)
self.register_buffer("cos_cache", emb.cos()[None, None], persistent=False)
self.register_buffer("sin_cache", emb.sin()[None, None], persistent=False); self._max = seq_len
@staticmethod
def _rot(x):
x1, x2 = x.chunk(2, dim=-1); return torch.cat([-x2, x1], dim=-1)
def forward(self, q, k):
T = q.size(2)
if T > self._max: self._build(T)
c, s = self.cos_cache[:,:,:T], self.sin_cache[:,:,:T]
return q*c + self._rot(q)*s, k*c + self._rot(k)*s
class GQA(nn.Module):
def __init__(self, cfg):
super().__init__()
self.n_heads, self.n_kv_heads = cfg.n_heads, cfg.n_kv_heads
self.n_groups, self.head_dim = cfg.n_heads // cfg.n_kv_heads, cfg.head_dim
self.q_proj = nn.Linear(cfg.d_model, cfg.n_heads * cfg.head_dim, bias=cfg.qkv_bias)
self.k_proj = nn.Linear(cfg.d_model, cfg.n_kv_heads * cfg.head_dim, bias=cfg.qkv_bias)
self.v_proj = nn.Linear(cfg.d_model, cfg.n_kv_heads * cfg.head_dim, bias=cfg.qkv_bias)
self.o_proj = nn.Linear(cfg.n_heads * cfg.head_dim, cfg.d_model, bias=False)
self.rope = RotaryEmbedding(cfg.head_dim, cfg.max_seq_len, cfg.rope_theta)
def forward(self, x):
B, T, _ = x.shape
q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
k = self.k_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
v = self.v_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
q, k = self.rope(q, k)
if self.n_groups > 1:
B_r, _, T_r, D_r = k.shape
k = k[:,:,None,:,:].expand(B_r, self.n_kv_heads, self.n_groups, T_r, D_r).reshape(B_r, self.n_heads, T_r, D_r)
v = v[:,:,None,:,:].expand(B_r, self.n_kv_heads, self.n_groups, T_r, D_r).reshape(B_r, self.n_heads, T_r, D_r)
return self.o_proj(F.scaled_dot_product_attention(q, k, v, is_causal=True).transpose(1, 2).contiguous().view(B, T, -1))
class Block(nn.Module):
def __init__(self, cfg):
super().__init__()
self.norm_attn = RMSNorm(cfg.d_model, cfg.rms_eps); self.attn = GQA(cfg)
self.norm_ffn = RMSNorm(cfg.d_model, cfg.rms_eps)
self.gate = nn.Linear(cfg.d_model, cfg.d_ff, bias=False)
self.up = nn.Linear(cfg.d_model, cfg.d_ff, bias=False)
self.down = nn.Linear(cfg.d_ff, cfg.d_model, bias=False)
def forward(self, x):
x = x + self.attn(self.norm_attn(x))
h = self.norm_ffn(x); x = x + self.down(F.silu(self.gate(h)) * self.up(h))
return x
class Quark(nn.Module):
def __init__(self, cfg):
super().__init__(); self.cfg = cfg
self.embed_tokens = nn.Embedding(cfg.vocab_size, cfg.d_model)
self.layers = nn.ModuleList([Block(cfg) for _ in range(cfg.n_layers)])
self.norm = RMSNorm(cfg.d_model, cfg.rms_eps)
self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
self.lm_head.weight = self.embed_tokens.weight
def forward(self, ids, labels=None):
x = self.embed_tokens(ids)
for layer in self.layers: x = layer(x)
logits = self.lm_head(self.norm(x))
loss = None
if labels is not None:
loss = F.cross_entropy(logits.view(-1, self.cfg.vocab_size), labels.view(-1), ignore_index=-100)
return loss, logits
|