148 lines
5.6 KiB
Python
148 lines
5.6 KiB
Python
"""Kleines Llama-artiges Decoder-only Sprachmodell (RMSNorm, RoPE, SwiGLU, GQA)."""
|
|
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
|
|
dim: int = 256
|
|
n_layers: int = 6
|
|
n_heads: int = 8
|
|
n_kv_heads: int = 4 # GQA: weniger KV-Köpfe als Query-Köpfe
|
|
hidden_dim: int = 688 # SwiGLU-Zwischengröße (~8/3 * dim)
|
|
max_seq_len: int = 256
|
|
rope_theta: float = 10000.0
|
|
norm_eps: float = 1e-5
|
|
dropout: float = 0.0
|
|
|
|
|
|
class RMSNorm(nn.Module):
|
|
def __init__(self, dim, eps=1e-5):
|
|
super().__init__()
|
|
self.eps = eps
|
|
self.weight = nn.Parameter(torch.ones(dim))
|
|
|
|
def forward(self, x):
|
|
return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight
|
|
|
|
|
|
def precompute_rope(head_dim, seq_len, theta):
|
|
freqs = 1.0 / (theta ** (torch.arange(0, head_dim, 2).float() / head_dim))
|
|
t = torch.arange(seq_len).float()
|
|
angles = torch.outer(t, freqs) # (T, hd/2)
|
|
return torch.cos(angles), torch.sin(angles)
|
|
|
|
|
|
def apply_rope(x, cos, sin):
|
|
# x: (B, H, T, hd); Paare (gerade, ungerade) rotieren
|
|
x1, x2 = x[..., 0::2], x[..., 1::2]
|
|
cos, sin = cos[None, None], sin[None, None]
|
|
out = torch.stack((x1 * cos - x2 * sin, x1 * sin + x2 * cos), dim=-1)
|
|
return out.flatten(-2)
|
|
|
|
|
|
class Attention(nn.Module):
|
|
def __init__(self, cfg: ModelConfig):
|
|
super().__init__()
|
|
self.n_heads, self.n_kv = cfg.n_heads, cfg.n_kv_heads
|
|
self.hd = cfg.dim // cfg.n_heads
|
|
self.wq = nn.Linear(cfg.dim, cfg.n_heads * self.hd, bias=False)
|
|
self.wk = nn.Linear(cfg.dim, self.n_kv * self.hd, bias=False)
|
|
self.wv = nn.Linear(cfg.dim, self.n_kv * self.hd, bias=False)
|
|
self.wo = nn.Linear(cfg.n_heads * self.hd, cfg.dim, bias=False)
|
|
self.dropout = cfg.dropout
|
|
|
|
def forward(self, x, cos, sin):
|
|
B, T, _ = x.shape
|
|
q = self.wq(x).view(B, T, self.n_heads, self.hd).transpose(1, 2)
|
|
k = self.wk(x).view(B, T, self.n_kv, self.hd).transpose(1, 2)
|
|
v = self.wv(x).view(B, T, self.n_kv, self.hd).transpose(1, 2)
|
|
q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
|
|
rep = self.n_heads // self.n_kv
|
|
if rep > 1: # KV-Köpfe für GQA wiederholen
|
|
k = k.repeat_interleave(rep, dim=1)
|
|
v = v.repeat_interleave(rep, dim=1)
|
|
y = F.scaled_dot_product_attention(q, k, v, is_causal=True,
|
|
dropout_p=self.dropout if self.training else 0.0)
|
|
return self.wo(y.transpose(1, 2).reshape(B, T, -1))
|
|
|
|
|
|
class SwiGLU(nn.Module):
|
|
def __init__(self, cfg: ModelConfig):
|
|
super().__init__()
|
|
self.w1 = nn.Linear(cfg.dim, cfg.hidden_dim, bias=False) # Gate
|
|
self.w3 = nn.Linear(cfg.dim, cfg.hidden_dim, bias=False) # Up
|
|
self.w2 = nn.Linear(cfg.hidden_dim, cfg.dim, bias=False) # Down
|
|
|
|
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 = RMSNorm(cfg.dim, cfg.norm_eps)
|
|
self.attn = Attention(cfg)
|
|
self.mlp_norm = RMSNorm(cfg.dim, cfg.norm_eps)
|
|
self.mlp = SwiGLU(cfg)
|
|
|
|
def forward(self, x, cos, sin):
|
|
x = x + self.attn(self.attn_norm(x), cos, sin)
|
|
return x + self.mlp(self.mlp_norm(x))
|
|
|
|
|
|
class Llama(nn.Module):
|
|
def __init__(self, cfg: ModelConfig):
|
|
super().__init__()
|
|
self.cfg = cfg
|
|
self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.dim)
|
|
self.layers = nn.ModuleList(Block(cfg) for _ in range(cfg.n_layers))
|
|
self.norm = RMSNorm(cfg.dim, cfg.norm_eps)
|
|
self.lm_head = nn.Linear(cfg.dim, cfg.vocab_size, bias=False)
|
|
self.lm_head.weight = self.tok_emb.weight # Weight Tying
|
|
cos, sin = precompute_rope(cfg.dim // cfg.n_heads, cfg.max_seq_len, cfg.rope_theta)
|
|
self.register_buffer("cos", cos, persistent=False)
|
|
self.register_buffer("sin", sin, persistent=False)
|
|
self.apply(self._init)
|
|
for n, p in self.named_parameters(): # skalierte Init für Residual-Projektionen
|
|
if n.endswith("wo.weight") or n.endswith("w2.weight"):
|
|
nn.init.normal_(p, 0.0, 0.02 / math.sqrt(2 * cfg.n_layers))
|
|
|
|
@staticmethod
|
|
def _init(m):
|
|
if isinstance(m, (nn.Linear, nn.Embedding)):
|
|
nn.init.normal_(m.weight, 0.0, 0.02)
|
|
|
|
def num_params(self):
|
|
return sum(p.numel() for p in self.parameters())
|
|
|
|
def forward(self, idx, targets=None):
|
|
T = idx.shape[1]
|
|
x = self.tok_emb(idx)
|
|
cos, sin = self.cos[:T], self.sin[:T]
|
|
for layer in self.layers:
|
|
x = layer(x, cos, sin)
|
|
logits = self.lm_head(self.norm(x))
|
|
loss = None
|
|
if targets is not None:
|
|
loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
|
|
return logits, loss
|
|
|
|
@torch.no_grad()
|
|
def generate(self, idx, max_new_tokens, temperature=0.8, top_k=50, eos_id=None):
|
|
for _ in range(max_new_tokens):
|
|
ctx = idx[:, -self.cfg.max_seq_len:]
|
|
logits, _ = self(ctx)
|
|
logits = logits[:, -1, :] / max(temperature, 1e-5)
|
|
if top_k:
|
|
v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
|
|
logits[logits < v[:, [-1]]] = -float("inf")
|
|
nxt = torch.multinomial(F.softmax(logits, -1), 1)
|
|
idx = torch.cat([idx, nxt], 1)
|
|
if eos_id is not None and nxt.item() == eos_id:
|
|
break
|
|
return idx
|