upload model_v5.py

model_v5.py

@@ -0,0 +1,243 @@
+"""
+V5 Model — 200M parametre, V4 mimari + büyütülmüş.
+
+Mimari:
+  - Layer: 18 (V4: 8)
+  - Head: 14 (V4: 10)
+  - Embd: 896 (V4: 640)
+  - Vocab: 32000 (V4: 16000)
+  - Context: 2048 (V4: 512)
+  - Toplam: ~210M parametre
+
+Modern teknikler (V4'ten):
+  - RoPE (real-valued)
+  - RMSNorm
+  - SwiGLU (hidden ~2560)
+  - QK-norm
+  - Logit soft-cap
+  - Tied embeddings
+  - Scaled init
+"""
+
+import math
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+@dataclass
+class GPTConfigV5:
+    block_size: int = 2048
+    vocab_size: int = 32000
+    n_layer: int = 18
+    n_head: int = 14
+    n_embd: int = 896
+    dropout: float = 0.0
+    rope_theta: float = 10000.0
+    logit_softcap: float = 30.0
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+
+    def forward(self, x):
+        dtype = x.dtype
+        x = x.float()
+        rms = torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+        return (self.weight * (x * rms)).to(dtype)
+
+
+def precompute_rope(dim, end, theta=10000.0, device=None):
+    inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, device=device).float() / dim))
+    t = torch.arange(end, device=device, 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_rotary_emb(xq, xk, cos, sin):
+    cos = cos.unsqueeze(0).unsqueeze(0)
+    sin = sin.unsqueeze(0).unsqueeze(0)
+    xq_out = (xq * cos) + (rotate_half(xq) * sin)
+    xk_out = (xk * cos) + (rotate_half(xk) * sin)
+    return xq_out.to(xq.dtype), xk_out.to(xk.dtype)
+
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        assert cfg.n_embd % cfg.n_head == 0
+        self.n_head = cfg.n_head
+        self.n_embd = cfg.n_embd
+        self.head_dim = cfg.n_embd // cfg.n_head
+        self.c_attn = nn.Linear(cfg.n_embd, 3 * cfg.n_embd, bias=False)
+        self.c_proj = nn.Linear(cfg.n_embd, cfg.n_embd, bias=False)
+        self.dropout = cfg.dropout
+        self.q_norm = RMSNorm(self.head_dim)
+        self.k_norm = RMSNorm(self.head_dim)
+
+    def forward(self, x, cos, sin):
+        B, T, C = x.shape
+        q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
+        q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
+        k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
+        v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+        q, k = apply_rotary_emb(q, k, cos[:T], sin[:T])
+        y = F.scaled_dot_product_attention(
+            q, k, v, dropout_p=self.dropout if self.training else 0.0, is_causal=True
+        )
+        y = y.transpose(1, 2).contiguous().view(B, T, C)
+        return self.c_proj(y)
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, dim, hidden_dim=None):
+        super().__init__()
+        if hidden_dim is None:
+            hidden_dim = int(8 * dim / 3)
+            hidden_dim = ((hidden_dim + 255) // 256) * 256
+        self.w1 = nn.Linear(dim, hidden_dim, bias=False)
+        self.w2 = nn.Linear(dim, hidden_dim, bias=False)
+        self.w3 = nn.Linear(hidden_dim, dim, bias=False)
+        self.hidden_dim = hidden_dim
+
+    def forward(self, x):
+        return self.w3(F.silu(self.w1(x)) * self.w2(x))
+
+
+class Block(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.norm1 = RMSNorm(cfg.n_embd)
+        self.attn = CausalSelfAttention(cfg)
+        self.norm2 = RMSNorm(cfg.n_embd)
+        self.mlp = SwiGLU(cfg.n_embd)
+
+    def forward(self, x, cos, sin):
+        x = x + self.attn(self.norm1(x), cos, sin)
+        x = x + self.mlp(self.norm2(x))
+        return x
+
+
+class GPTV5(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.wte = nn.Embedding(cfg.vocab_size, cfg.n_embd)
+        self.h = nn.ModuleList([Block(cfg) for _ in range(cfg.n_layer)])
+        self.norm_f = RMSNorm(cfg.n_embd)
+        self.lm_head = nn.Linear(cfg.n_embd, cfg.vocab_size, bias=False)
+        self.lm_head.weight = self.wte.weight  # tied
+
+        head_dim = cfg.n_embd // cfg.n_head
+        cos, sin = precompute_rope(head_dim, cfg.block_size * 2, cfg.rope_theta)
+        self.register_buffer("rope_cos", cos, persistent=False)
+        self.register_buffer("rope_sin", sin, persistent=False)
+
+        self.apply(self._init_weights)
+        for pn, p in self.named_parameters():
+            if pn.endswith("c_proj.weight") or pn.endswith("w3.weight"):
+                nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * cfg.n_layer))
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+    def num_params(self):
+        return sum(p.numel() for p in self.parameters())
+
+    def forward(self, idx, targets=None):
+        B, T = idx.shape
+        max_t = self.rope_cos.size(0)
+        assert T <= max_t, f"T={T} > rope buffer={max_t}"
+
+        x = self.wte(idx)
+        cos = self.rope_cos[:T]
+        sin = self.rope_sin[:T]
+        for block in self.h:
+            x = block(x, cos, sin)
+        x = self.norm_f(x)
+
+        if targets is not None:
+            logits = self.lm_head(x)
+            if self.cfg.logit_softcap > 0:
+                cap = self.cfg.logit_softcap
+                logits = cap * torch.tanh(logits / cap)
+            loss = F.cross_entropy(
+                logits.view(-1, logits.size(-1)),
+                targets.view(-1),
+                ignore_index=-1,
+            )
+            return logits, loss
+        else:
+            logits = self.lm_head(x[:, [-1], :])
+            if self.cfg.logit_softcap > 0:
+                cap = self.cfg.logit_softcap
+                logits = cap * torch.tanh(logits / cap)
+            return logits, None
+
+    @torch.no_grad()
+    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None,
+                 repetition_penalty=1.15, no_repeat_ngram_size=3,
+                 max_context=None):
+        if max_context is None:
+            max_context = self.rope_cos.size(0)
+        for _ in range(max_new_tokens):
+            idx_cond = idx if idx.size(1) <= max_context else idx[:, -max_context:]
+            logits, _ = self(idx_cond)
+            logits = logits[:, -1, :] / temperature
+
+            if repetition_penalty != 1.0:
+                for b in range(idx.size(0)):
+                    seen = set(idx[b].tolist()[-256:])
+                    for tok in seen:
+                        if logits[b, tok] > 0:
+                            logits[b, tok] /= repetition_penalty
+                        else:
+                            logits[b, tok] *= repetition_penalty
+
+            if no_repeat_ngram_size > 0 and idx.size(1) >= no_repeat_ngram_size:
+                for b in range(idx.size(0)):
+                    tokens = idx[b].tolist()
+                    n = no_repeat_ngram_size
+                    prefix = tuple(tokens[-(n-1):])
+                    banned = set()
+                    for i in range(len(tokens) - n + 1):
+                        if tuple(tokens[i:i+n-1]) == prefix:
+                            banned.add(tokens[i+n-1])
+                    for tok in banned:
+                        logits[b, tok] = -float("inf")
+
+            if top_k is not None:
+                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                logits[logits < v[:, [-1]]] = -float("inf")
+            probs = F.softmax(logits, dim=-1)
+            next_id = torch.multinomial(probs, num_samples=1)
+            idx = torch.cat([idx, next_id], dim=1)
+        return idx
+
+
+if __name__ == "__main__":
+    cfg = GPTConfigV5()
+    m = GPTV5(cfg)
+    print(f"V5: {m.num_params()/1e6:.2f}M param")
+    print(f"  Layer: {cfg.n_layer}, Head: {cfg.n_head}, Embd: {cfg.n_embd}")
+    print(f"  Vocab: {cfg.vocab_size}, Block: {cfg.block_size}")
+    print(f"  SwiGLU hidden: {m.h[0].mlp.hidden_dim}")
+    x = torch.randint(0, cfg.vocab_size, (2, 64))
+    logits, loss = m(x, x)
+    print(f"Forward: loss {loss.item():.4f}")