| """VRS Foundation LLM - self-contained model definition (no external package needed).
|
|
|
| Everything required to load and run the model lives in this one file:
|
| config, sampling, and the transformer (GQA + RoPE + RMSNorm + SwiGLU + KV-cache).
|
| """
|
|
|
| from __future__ import annotations
|
|
|
| import math
|
| from dataclasses import dataclass
|
|
|
| import torch
|
| import torch.nn as nn
|
| import torch.nn.functional as F
|
|
|
|
|
| @dataclass
|
| class ModelConfig:
|
| vocab_size: int = 8192
|
| d_model: int = 512
|
| n_layers: int = 8
|
| n_heads: int = 8
|
| n_kv_heads: int = 4
|
| d_ff: int = 1408
|
| max_seq_len: int = 512
|
| rope_theta: float = 10000.0
|
| rms_norm_eps: float = 1e-6
|
| dropout: float = 0.0
|
| tie_embeddings: bool = True
|
|
|
| @property
|
| def head_dim(self) -> int:
|
| return self.d_model // self.n_heads
|
|
|
|
|
|
|
|
|
|
|
| @torch.no_grad()
|
| def apply_repetition_penalty(logits, prev_ids, penalty):
|
| if penalty == 1.0:
|
| return logits
|
| for b in range(prev_ids.size(0)):
|
| ids = torch.unique(prev_ids[b])
|
| vals = logits[b, ids]
|
| logits[b, ids] = torch.where(vals > 0, vals / penalty, vals * penalty)
|
| return logits
|
|
|
|
|
| @torch.no_grad()
|
| def sample_next(logits, temperature=0.8, top_k=50, top_p=None):
|
| if temperature <= 0.0:
|
| return logits.argmax(dim=-1, keepdim=True)
|
| logits = logits / temperature
|
| if top_k is not None:
|
| k = min(top_k, logits.size(-1))
|
| kth = torch.topk(logits, k, dim=-1).values[:, -1, None]
|
| logits = logits.masked_fill(logits < kth, float("-inf"))
|
| probs = F.softmax(logits, dim=-1)
|
| if top_p is not None:
|
| sp, si = torch.sort(probs, descending=True, dim=-1)
|
| cum = torch.cumsum(sp, dim=-1)
|
| remove = cum - sp > top_p
|
| sp = sp.masked_fill(remove, 0.0)
|
| sp = sp / sp.sum(dim=-1, keepdim=True)
|
| choice = torch.multinomial(sp, 1)
|
| return si.gather(-1, choice)
|
| return torch.multinomial(probs, 1)
|
|
|
|
|
|
|
|
|
|
|
| def build_rope_cache(head_dim, max_seq_len, theta):
|
| inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim))
|
| t = torch.arange(max_seq_len, dtype=torch.float32)
|
| freqs = torch.outer(t, inv_freq)
|
| emb = torch.cat((freqs, freqs), dim=-1)
|
| return emb.cos(), emb.sin()
|
|
|
|
|
| def rotate_half(x):
|
| x1, x2 = x.chunk(2, dim=-1)
|
| return torch.cat((-x2, x1), dim=-1)
|
|
|
|
|
| def apply_rope(x, cos, sin):
|
| cos = cos[None, None, :, :]
|
| sin = sin[None, None, :, :]
|
| return (x * cos) + (rotate_half(x) * sin)
|
|
|
|
|
| def repeat_kv(x, n_rep):
|
| if n_rep == 1:
|
| return x
|
| B, n_kv, T, hd = x.shape
|
| return x[:, :, None, :, :].expand(B, n_kv, n_rep, T, hd).reshape(B, n_kv * n_rep, T, hd)
|
|
|
|
|
| class Attention(nn.Module):
|
| def __init__(self, cfg):
|
| super().__init__()
|
| self.n_heads, self.n_kv_heads = cfg.n_heads, cfg.n_kv_heads
|
| self.n_rep = cfg.n_heads // cfg.n_kv_heads
|
| self.hd = cfg.head_dim
|
| self.q_proj = nn.Linear(cfg.d_model, cfg.n_heads * self.hd, bias=False)
|
| self.k_proj = nn.Linear(cfg.d_model, cfg.n_kv_heads * self.hd, bias=False)
|
| self.v_proj = nn.Linear(cfg.d_model, cfg.n_kv_heads * self.hd, bias=False)
|
| self.o_proj = nn.Linear(cfg.n_heads * self.hd, cfg.d_model, bias=False)
|
| self.q_norm = nn.RMSNorm(self.hd, eps=cfg.rms_norm_eps)
|
| self.k_norm = nn.RMSNorm(self.hd, eps=cfg.rms_norm_eps)
|
| self.dropout = cfg.dropout
|
|
|
| def forward(self, x, cos, sin, cache=None):
|
| B, T, _ = x.shape
|
| q = self.q_proj(x).view(B, T, self.n_heads, self.hd).transpose(1, 2)
|
| k = self.k_proj(x).view(B, T, self.n_kv_heads, self.hd).transpose(1, 2)
|
| v = self.v_proj(x).view(B, T, self.n_kv_heads, self.hd).transpose(1, 2)
|
| q = apply_rope(self.q_norm(q), cos, sin)
|
| k = apply_rope(self.k_norm(k), cos, sin)
|
| new_cache = None
|
| if cache is not None:
|
| pk, pv = cache
|
| if pk is not None:
|
| k = torch.cat([pk, k], dim=2)
|
| v = torch.cat([pv, v], dim=2)
|
| new_cache = (k, v)
|
| k = repeat_kv(k, self.n_rep)
|
| v = repeat_kv(v, self.n_rep)
|
| is_causal = q.size(2) > 1
|
| out = F.scaled_dot_product_attention(q, k, v, is_causal=is_causal,
|
| dropout_p=self.dropout if self.training else 0.0)
|
| out = out.transpose(1, 2).contiguous().view(B, T, -1)
|
| return self.o_proj(out), new_cache
|
|
|
|
|
| class SwiGLU(nn.Module):
|
| def __init__(self, cfg):
|
| super().__init__()
|
| self.w1 = nn.Linear(cfg.d_model, cfg.d_ff, bias=False)
|
| self.w3 = nn.Linear(cfg.d_model, cfg.d_ff, bias=False)
|
| self.w2 = nn.Linear(cfg.d_ff, cfg.d_model, bias=False)
|
|
|
| def forward(self, x):
|
| return self.w2(F.silu(self.w1(x)) * self.w3(x))
|
|
|
|
|
| class Block(nn.Module):
|
| def __init__(self, cfg):
|
| super().__init__()
|
| self.attn_norm = nn.RMSNorm(cfg.d_model, eps=cfg.rms_norm_eps)
|
| self.attn = Attention(cfg)
|
| self.ffn_norm = nn.RMSNorm(cfg.d_model, eps=cfg.rms_norm_eps)
|
| self.ffn = SwiGLU(cfg)
|
|
|
| def forward(self, x, cos, sin, cache=None):
|
| h, new_cache = self.attn(self.attn_norm(x), cos, sin, cache)
|
| x = x + h
|
| x = x + self.ffn(self.ffn_norm(x))
|
| return x, new_cache
|
|
|
|
|
| class GPT(nn.Module):
|
| def __init__(self, cfg: ModelConfig):
|
| super().__init__()
|
| self.cfg = cfg
|
| self.embed = nn.Embedding(cfg.vocab_size, cfg.d_model)
|
| self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.n_layers)])
|
| self.norm = nn.RMSNorm(cfg.d_model, eps=cfg.rms_norm_eps)
|
| self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
|
| if cfg.tie_embeddings:
|
| self.lm_head.weight = self.embed.weight
|
| cos, sin = build_rope_cache(cfg.head_dim, cfg.max_seq_len, cfg.rope_theta)
|
| self.register_buffer("rope_cos", cos, persistent=False)
|
| self.register_buffer("rope_sin", sin, persistent=False)
|
|
|
| @torch.no_grad()
|
| def generate(self, idx, max_new_tokens=200, temperature=0.8, top_k=50,
|
| top_p=None, repetition_penalty=1.0, eos_id=None):
|
| self.eval()
|
| B, T = idx.shape
|
| x = self.embed(idx)
|
| cos, sin = self.rope_cos[:T], self.rope_sin[:T]
|
| caches = [(None, None)] * len(self.blocks)
|
| for i, block in enumerate(self.blocks):
|
| x, caches[i] = block(x, cos, sin, caches[i])
|
| logits = self.lm_head(self.norm(x[:, -1:, :]))
|
| out, pos = idx, T
|
| for _ in range(max_new_tokens):
|
| step = apply_repetition_penalty(logits[:, -1, :], out, repetition_penalty)
|
| nxt = sample_next(step, temperature, top_k, top_p)
|
| out = torch.cat([out, nxt], dim=1)
|
| if eos_id is not None and bool((nxt == eos_id).all()):
|
| break
|
| if pos >= self.cfg.max_seq_len:
|
| break
|
| x = self.embed(nxt)
|
| cos, sin = self.rope_cos[pos:pos + 1], self.rope_sin[pos:pos + 1]
|
| for i, block in enumerate(self.blocks):
|
| x, caches[i] = block(x, cos, sin, caches[i])
|
| logits = self.lm_head(self.norm(x))
|
| pos += 1
|
| return out
|
|
|
