Add MP-01 (9500) 100M base: weights, config, tokenizer, model code, model card and eval results
11e2c7c verified | """VBQ transformer: Llama-style decoder (RMSNorm, RoPE, GQA, QK-norm, SwiGLU, tied embeddings).""" | |
| from __future__ import annotations | |
| import math | |
| from dataclasses import dataclass, asdict | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| class VBQConfig: | |
| vocab_size: int = 32768 | |
| d_model: int = 640 | |
| n_layers: int = 12 | |
| n_heads: int = 10 | |
| n_kv_heads: int = 5 | |
| mlp_hidden: int = 1664 | |
| max_seq_len: int = 2048 | |
| rope_theta: float = 10000.0 | |
| norm_eps: float = 1e-5 | |
| qk_norm: bool = True | |
| tie_embeddings: bool = True | |
| # Mixture of experts (0 = dense SwiGLU). Top-1 routing, Switch-style. | |
| n_experts: int = 0 | |
| capacity_factor: float = 1.25 | |
| def head_dim(self) -> int: | |
| return self.d_model // self.n_heads | |
| def n_params(self) -> int: | |
| d, hd = self.d_model, self.head_dim | |
| attn = d * d * 2 + d * (self.n_kv_heads * hd) * 2 | |
| mlp = 3 * d * self.mlp_hidden * max(1, self.n_experts) + (d * self.n_experts if self.n_experts else 0) | |
| per_layer = attn + mlp + 2 * d + (2 * hd if self.qk_norm else 0) | |
| emb = self.vocab_size * d * (1 if self.tie_embeddings else 2) | |
| return emb + self.n_layers * per_layer + d | |
| def to_dict(self): | |
| return asdict(self) | |
| class RMSNorm(nn.Module): | |
| def __init__(self, dim: int, eps: float = 1e-5): | |
| super().__init__() | |
| self.eps = eps | |
| self.weight = nn.Parameter(torch.ones(dim)) | |
| def forward(self, x): | |
| return F.rms_norm(x, (x.shape[-1],), self.weight.to(x.dtype), self.eps) | |
| def build_rope(seq_len: int, head_dim: int, theta: float, device=None): | |
| inv = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device, dtype=torch.float32) / head_dim)) | |
| t = torch.arange(seq_len, device=device, dtype=torch.float32) | |
| freqs = torch.outer(t, inv) | |
| return freqs.cos(), freqs.sin() # [T, hd/2] | |
| def apply_rope(x, cos, sin): | |
| # x: [B, H, T, hd]; half-split rotation | |
| x1, x2 = x.chunk(2, dim=-1) | |
| cos = cos[None, None, : x.shape[2]].to(x.dtype) | |
| sin = sin[None, None, : x.shape[2]].to(x.dtype) | |
| return torch.cat((x1 * cos - x2 * sin, x2 * cos + x1 * sin), dim=-1) | |
| class Attention(nn.Module): | |
| def __init__(self, c: VBQConfig): | |
| super().__init__() | |
| self.c = c | |
| hd = c.head_dim | |
| self.wq = nn.Linear(c.d_model, c.n_heads * hd, bias=False) | |
| self.wk = nn.Linear(c.d_model, c.n_kv_heads * hd, bias=False) | |
| self.wv = nn.Linear(c.d_model, c.n_kv_heads * hd, bias=False) | |
| self.wo = nn.Linear(c.n_heads * hd, c.d_model, bias=False) | |
| if c.qk_norm: | |
| self.q_norm = RMSNorm(hd, c.norm_eps) | |
| self.k_norm = RMSNorm(hd, c.norm_eps) | |
| def forward(self, x, cos, sin): | |
| B, T, _ = x.shape | |
| c, hd = self.c, self.c.head_dim | |
| q = self.wq(x).view(B, T, c.n_heads, hd).transpose(1, 2) | |
| k = self.wk(x).view(B, T, c.n_kv_heads, hd).transpose(1, 2) | |
| v = self.wv(x).view(B, T, c.n_kv_heads, hd).transpose(1, 2) | |
| if c.qk_norm: | |
| q, k = self.q_norm(q), self.k_norm(k) | |
| q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin) | |
| if c.n_kv_heads != c.n_heads: | |
| rep = c.n_heads // c.n_kv_heads | |
| k = k[:, :, None].expand(B, c.n_kv_heads, rep, T, hd).reshape(B, c.n_heads, T, hd) | |
| v = v[:, :, None].expand(B, c.n_kv_heads, rep, T, hd).reshape(B, c.n_heads, T, hd) | |
| y = F.scaled_dot_product_attention(q, k, v, is_causal=True) | |
| return self.wo(y.transpose(1, 2).reshape(B, T, c.n_heads * hd)) | |
| class MLP(nn.Module): | |
| def __init__(self, c: VBQConfig): | |
| super().__init__() | |
| self.gate_up = nn.Linear(c.d_model, 2 * c.mlp_hidden, bias=False) | |
| self.down = nn.Linear(c.mlp_hidden, c.d_model, bias=False) | |
| def forward(self, x): | |
| g, u = self.gate_up(x).chunk(2, dim=-1) | |
| return self.down(F.silu(g) * u) | |
| class MoE(nn.Module): | |
| """Top-1 mixture of SwiGLU experts (Switch Transformer routing). | |
| Fixed shapes (compile-friendly): every expert gets a buffer of | |
| capacity = ceil(capacity_factor * tokens / experts) slots while training | |
| (tokens past it skip the MLP and keep only the residual); at eval every | |
| token fits. The chosen expert's output is scaled by p / p.detach(), which | |
| is 1 in value, so a model upcycled from identical copies computes exactly | |
| what the dense model did, while the router still gets gradients. | |
| Returns (y, aux) with aux = Switch load-balancing loss (1.0 = balanced). | |
| """ | |
| def __init__(self, c: VBQConfig): | |
| super().__init__() | |
| self.c = c | |
| E, d, h = c.n_experts, c.d_model, c.mlp_hidden | |
| self.router = nn.Linear(d, E, bias=False) | |
| self.gate_up = nn.Parameter(torch.empty(E, 2 * h, d)) | |
| self.down = nn.Parameter(torch.empty(E, d, h)) | |
| nn.init.normal_(self.gate_up, 0.0, 0.02) | |
| nn.init.normal_(self.down, 0.0, 0.02) | |
| def forward(self, x): | |
| shape = x.shape | |
| x = x.reshape(-1, shape[-1]) | |
| N, d = x.shape | |
| E = self.c.n_experts | |
| probs = F.softmax(self.router(x).float(), dim=-1) # [N, E] | |
| gate, idx = probs.max(-1) # [N] | |
| onehot = F.one_hot(idx, E) # [N, E] | |
| pos = ((onehot.cumsum(0) - 1) * onehot).sum(-1) # slot within its expert | |
| C = math.ceil(self.c.capacity_factor * N / E) if self.training else N | |
| keep = pos < C | |
| slot = torch.where(keep, idx * C + pos, torch.full_like(pos, E * C)) # overflow -> dump slot | |
| buf = x.new_zeros(E * C + 1, d).index_put((slot,), x) | |
| hbuf = buf[: E * C].view(E, C, d) | |
| g, u = torch.bmm(hbuf, self.gate_up.to(x.dtype).transpose(1, 2)).chunk(2, dim=-1) | |
| y = torch.bmm(F.silu(g) * u, self.down.to(x.dtype).transpose(1, 2)) # [E, C, d] | |
| y = torch.cat([y.reshape(E * C, d), y.new_zeros(1, d)])[slot] | |
| scale = (gate / gate.detach()) * keep | |
| y = y * scale.to(y.dtype).unsqueeze(-1) | |
| aux = E * (onehot.float().mean(0) * probs.mean(0)).sum() | |
| return y.reshape(shape), aux | |
| class Block(nn.Module): | |
| def __init__(self, c: VBQConfig): | |
| super().__init__() | |
| self.attn_norm = RMSNorm(c.d_model, c.norm_eps) | |
| self.attn = Attention(c) | |
| self.mlp_norm = RMSNorm(c.d_model, c.norm_eps) | |
| self.mlp = MoE(c) if c.n_experts else MLP(c) | |
| def forward(self, x, cos, sin): | |
| x = x + self.attn(self.attn_norm(x), cos, sin) | |
| return x + self.ffn(self.mlp_norm(x))[0] | |
| def ffn(self, x): | |
| """(output, load-balancing loss) for dense and MoE blocks alike.""" | |
| y = self.mlp(x) | |
| return y if isinstance(y, tuple) else (y, x.new_zeros((), dtype=torch.float32)) | |
| class FusedLinearCE(torch.autograd.Function): | |
| """Cross-entropy over a huge vocab without ever materialising full logits. | |
| Gradients w.r.t. hidden states and the (tied) output matrix are computed chunk by chunk | |
| during the forward pass, so the backward pass is a cheap scale. Saves ~4 GB at 32k vocab. | |
| """ | |
| def forward(ctx, h, w, targets, chunk: int, z_loss: float): | |
| N = h.shape[0] | |
| wb = w.to(h.dtype) | |
| grad_h = torch.empty_like(h) | |
| grad_w = torch.zeros_like(w, dtype=torch.float32) | |
| total = torch.zeros((), device=h.device, dtype=torch.float32) | |
| for i in range(0, N, chunk): | |
| hc, tc = h[i : i + chunk], targets[i : i + chunk] | |
| logits = (hc @ wb.t()).float() | |
| lse = torch.logsumexp(logits, dim=-1) | |
| tgt = logits.gather(1, tc[:, None]).squeeze(1) | |
| total += (lse - tgt).sum() | |
| p = torch.exp(logits - lse[:, None]) | |
| if z_loss: | |
| total += z_loss * (lse * lse).sum() | |
| p.mul_(1.0 + 2.0 * z_loss * lse[:, None]) | |
| p.scatter_add_(1, tc[:, None], torch.full_like(tc[:, None], -1, dtype=p.dtype)) | |
| p.mul_(1.0 / N) | |
| pb = p.to(h.dtype) | |
| grad_h[i : i + chunk] = pb @ wb | |
| grad_w += (pb.t() @ hc).float() | |
| ctx.save_for_backward(grad_h, grad_w) | |
| return total / N | |
| def backward(ctx, g): | |
| grad_h, grad_w = ctx.saved_tensors | |
| return grad_h * g.to(grad_h.dtype), grad_w * g, None, None, None | |
| class VBQ(nn.Module): | |
| def __init__(self, c: VBQConfig): | |
| super().__init__() | |
| self.c = c | |
| self.tok_emb = nn.Embedding(c.vocab_size, c.d_model) | |
| self.blocks = nn.ModuleList(Block(c) for _ in range(c.n_layers)) | |
| self.final_norm = RMSNorm(c.d_model, c.norm_eps) | |
| self.lm_head = None if c.tie_embeddings else nn.Linear(c.d_model, c.vocab_size, bias=False) | |
| cos, sin = build_rope(c.max_seq_len, c.head_dim, c.rope_theta) | |
| self.register_buffer("rope_cos", cos, persistent=False) | |
| self.register_buffer("rope_sin", sin, persistent=False) | |
| self.apply(self._init) | |
| for n, p in self.named_parameters(): # scaled init on residual output projections | |
| if n.endswith("wo.weight") or n.endswith("down.weight"): | |
| nn.init.normal_(p, 0.0, 0.02 / math.sqrt(2 * c.n_layers)) | |
| def _init(m): | |
| if isinstance(m, (nn.Linear, nn.Embedding)): | |
| nn.init.normal_(m.weight, 0.0, 0.02) | |
| def out_weight(self): | |
| return self.tok_emb.weight if self.lm_head is None else self.lm_head.weight | |
| def hidden(self, idx): | |
| return self.hidden_aux(idx)[0] | |
| def hidden_aux(self, idx): | |
| """Final hidden states and the mean MoE load-balancing loss (0 for dense).""" | |
| return self._blocks(self.tok_emb(idx)) | |
| def hidden_patch(self, idx, k: int): | |
| """Patch-level training (Shao et al. 2024): idx [B, P*k]; each group of k tokens enters as the mean of its | |
| token embeddings, so the transformer runs on P positions. Same weights as token-level training.""" | |
| B, L = idx.shape | |
| return self._blocks(self.tok_emb(idx).view(B, L // k, k, -1).mean(2)) | |
| def _blocks(self, x): | |
| aux = x.new_zeros((), dtype=torch.float32) | |
| for blk in self.blocks: | |
| x = x + blk.attn(blk.attn_norm(x), self.rope_cos, self.rope_sin) | |
| y, a = blk.ffn(blk.mlp_norm(x)) | |
| x = x + y | |
| aux = aux + a | |
| return self.final_norm(x), aux / len(self.blocks) | |
| def forward(self, idx, targets=None, z_loss: float = 0.0, ce_chunk: int = 4096): | |
| h = self.hidden(idx) | |
| if targets is None: | |
| return (h @ self.out_weight.to(h.dtype).t()).float() | |
| loss = FusedLinearCE.apply(h.reshape(-1, h.shape[-1]), self.out_weight, targets.reshape(-1), ce_chunk, z_loss) | |
| return loss | |
| def generate(self, idx, max_new_tokens=128, temperature=0.8, top_k=50, top_p=0.95, eos_id=None, repetition_penalty=1.0): | |
| self.eval() | |
| for _ in range(max_new_tokens): | |
| ctx = idx[:, -self.c.max_seq_len :] | |
| logits = self(ctx)[:, -1] | |
| if repetition_penalty != 1.0: | |
| for tok in set(ctx[0].tolist()): | |
| logits[0, tok] = logits[0, tok] / repetition_penalty if logits[0, tok] > 0 else logits[0, tok] * repetition_penalty | |
| if temperature <= 0: | |
| nxt = logits.argmax(-1, keepdim=True) | |
| else: | |
| logits = logits / temperature | |
| if top_k: | |
| v, _ = torch.topk(logits, min(top_k, logits.shape[-1])) | |
| logits[logits < v[:, [-1]]] = -float("inf") | |
| probs = F.softmax(logits, dim=-1) | |
| if top_p < 1.0: | |
| sp, si = torch.sort(probs, descending=True) | |
| cum = sp.cumsum(-1) | |
| sp[cum - sp > top_p] = 0 | |
| probs = torch.zeros_like(probs).scatter_(1, si, sp) | |
| probs = probs / probs.sum(-1, keepdim=True) | |
| nxt = torch.multinomial(probs, 1) | |
| idx = torch.cat([idx, nxt], dim=1) | |
| yield int(nxt.item()) | |
| if eos_id is not None and int(nxt.item()) == eos_id: | |
| break | |
| def generate_cached(self, idx, max_new_tokens=128, temperature=0.8, top_k=50, top_p=0.95, eos_id=None, repetition_penalty=1.0): | |
| """KV-cached sampling. Yields token ids one by one. idx: [1, T] prompt.""" | |
| self.eval() | |
| c = self.c | |
| idx = idx[:, -(c.max_seq_len - 1):] | |
| stop_ids = set() if eos_id is None else ({eos_id} if isinstance(eos_id, int) else set(eos_id)) | |
| caches: list = [None] * c.n_layers | |
| seen = set(idx[0].tolist()) | |
| def forward(ids, pos): | |
| T = ids.shape[1] | |
| x = self.tok_emb(ids) | |
| cos, sin = self.rope_cos[pos : pos + T], self.rope_sin[pos : pos + T] | |
| for i, blk in enumerate(self.blocks): | |
| a = blk.attn | |
| h = blk.attn_norm(x) | |
| q = a.wq(h).view(1, T, c.n_heads, c.head_dim).transpose(1, 2) | |
| k = a.wk(h).view(1, T, c.n_kv_heads, c.head_dim).transpose(1, 2) | |
| v = a.wv(h).view(1, T, c.n_kv_heads, c.head_dim).transpose(1, 2) | |
| if c.qk_norm: | |
| q, k = a.q_norm(q), a.k_norm(k) | |
| q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin) | |
| if caches[i] is not None: | |
| k, v = torch.cat([caches[i][0], k], 2), torch.cat([caches[i][1], v], 2) | |
| caches[i] = (k, v) | |
| rep = c.n_heads // c.n_kv_heads | |
| kk = k.repeat_interleave(rep, dim=1) if rep > 1 else k | |
| vv = v.repeat_interleave(rep, dim=1) if rep > 1 else v | |
| y = F.scaled_dot_product_attention(q, kk, vv, is_causal=(T > 1)) | |
| x = x + a.wo(y.transpose(1, 2).reshape(1, T, c.n_heads * c.head_dim)) | |
| x = x + blk.ffn(blk.mlp_norm(x))[0] | |
| return (self.final_norm(x[:, -1]) @ self.out_weight.to(x.dtype).t()).float() | |
| logits = forward(idx, 0) | |
| pos = idx.shape[1] | |
| for _ in range(max_new_tokens): | |
| if pos >= c.max_seq_len: | |
| break | |
| lg = logits[0].clone() | |
| if repetition_penalty != 1.0: | |
| for tok in seen: | |
| lg[tok] = lg[tok] / repetition_penalty if lg[tok] > 0 else lg[tok] * repetition_penalty | |
| if temperature <= 0: | |
| nxt = int(lg.argmax()) | |
| else: | |
| lg = lg / temperature | |
| if top_k: | |
| v, _ = torch.topk(lg, min(top_k, lg.numel())) | |
| lg[lg < v[-1]] = -float("inf") | |
| probs = F.softmax(lg, dim=-1) | |
| if top_p < 1.0: | |
| sp, si = torch.sort(probs, descending=True) | |
| sp[sp.cumsum(-1) - sp > top_p] = 0 | |
| probs = torch.zeros_like(probs).scatter_(0, si, sp) | |
| probs = probs / probs.sum() | |
| nxt = int(torch.multinomial(probs, 1)) | |
| yield nxt | |
| if nxt in stop_ids: | |
| break | |
| seen.add(nxt) | |
| logits = forward(torch.tensor([[nxt]], device=idx.device), pos) | |
| pos += 1 | |