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