| """ |
| model.py |
| TurkishLLM Generative V1 Model Architecture (52.5M Parameters) |
| Decoder-Only Transformer with SwiGLU, RoPE, RMSNorm, Weight Tying and Causal Self-Attention. |
| Custom PyTorch Architecture. |
| """ |
|
|
| import math |
| from typing import Optional, Tuple |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| import torch.utils.checkpoint |
|
|
|
|
| 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: torch.Tensor) -> torch.Tensor: |
| norm = torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) |
| return x * norm * self.weight |
|
|
|
|
| class RotaryEmbedding(nn.Module): |
| def __init__(self, dim: int, max_seq_len: int = 1024, base: float = 10000.0): |
| super().__init__() |
| self.dim = dim |
| self.max_seq_len = max_seq_len |
| self.base = base |
| inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim)) |
| self.register_buffer("inv_freq", inv_freq, persistent=False) |
|
|
| def forward(self, x: torch.Tensor, seq_len: int) -> Tuple[torch.Tensor, torch.Tensor]: |
| t = torch.arange(seq_len, device=x.device, dtype=self.inv_freq.dtype) |
| freqs = torch.outer(t, self.inv_freq) |
| emb = torch.cat((freqs, freqs), dim=-1) |
| return emb.cos(), emb.sin() |
|
|
|
|
| def rotate_half(x: torch.Tensor) -> torch.Tensor: |
| x1 = x[..., :x.shape[-1] // 2] |
| x2 = x[..., x.shape[-1] // 2:] |
| return torch.cat((-x2, x1), dim=-1) |
|
|
|
|
| def apply_rotary_pos_emb(q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: |
| cos = cos.unsqueeze(0).unsqueeze(1) |
| sin = sin.unsqueeze(0).unsqueeze(1) |
| q_embed = (q * cos) + (rotate_half(q) * sin) |
| k_embed = (k * cos) + (rotate_half(k) * sin) |
| return q_embed, k_embed |
|
|
|
|
| class CausalSelfAttention(nn.Module): |
| def __init__(self, d_model: int = 576, num_heads: int = 9): |
| super().__init__() |
| assert d_model % num_heads == 0, f"d_model {d_model} must be divisible by num_heads {num_heads}" |
| self.d_model = d_model |
| self.num_heads = num_heads |
| self.head_dim = d_model // num_heads |
| self.q_proj = nn.Linear(d_model, d_model, bias=False) |
| self.k_proj = nn.Linear(d_model, d_model, bias=False) |
| self.v_proj = nn.Linear(d_model, d_model, bias=False) |
| self.out_proj = nn.Linear(d_model, d_model, bias=False) |
|
|
| def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor: |
| B, S, D = x.shape |
| q = self.q_proj(x).view(B, S, self.num_heads, self.head_dim).transpose(1, 2) |
| k = self.k_proj(x).view(B, S, self.num_heads, self.head_dim).transpose(1, 2) |
| v = self.v_proj(x).view(B, S, self.num_heads, self.head_dim).transpose(1, 2) |
|
|
| q, k = apply_rotary_pos_emb(q, k, cos, sin) |
|
|
| if attention_mask is not None: |
| if attention_mask.dim() == 2: |
| mask = attention_mask[:, None, None, :].bool() |
| out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, is_causal=False) |
| else: |
| out = F.scaled_dot_product_attention(q, k, v, attn_mask=attention_mask, is_causal=False) |
| else: |
| out = F.scaled_dot_product_attention(q, k, v, is_causal=True) |
|
|
| out = out.transpose(1, 2).contiguous().view(B, S, D) |
| return self.out_proj(out) |
|
|
|
|
| class SwiGLUFFN(nn.Module): |
| def __init__(self, d_model: int = 576, ffn_dim: int = 1536): |
| super().__init__() |
| self.gate_proj = nn.Linear(d_model, ffn_dim, bias=False) |
| self.up_proj = nn.Linear(d_model, ffn_dim, bias=False) |
| self.down_proj = nn.Linear(ffn_dim, d_model, bias=False) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) |
|
|
|
|
| class TransformerBlock(nn.Module): |
| def __init__(self, d_model: int = 576, num_heads: int = 9, ffn_dim: int = 1536): |
| super().__init__() |
| self.norm1 = RMSNorm(d_model) |
| self.attn = CausalSelfAttention(d_model, num_heads) |
| self.norm2 = RMSNorm(d_model) |
| self.ffn = SwiGLUFFN(d_model, ffn_dim) |
|
|
| def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor: |
| x = x + self.attn(self.norm1(x), cos, sin, attention_mask=attention_mask) |
| x = x + self.ffn(self.norm2(x)) |
| return x |
|
|
|
|
| class TurkishLLMGenerativeV1(nn.Module): |
| """ |
| TurkishLLM Generative V1 Model Architecture (52,508,736 Parameters) |
| - Layers: 12 |
| - Hidden Dim (d_model): 576 |
| - Heads: 9 (Head Dim: 64) |
| - FFN: SwiGLU (dim: 1536) |
| - Max Context: 1024 |
| - Vocab Size (Embedding Capacity): 8192 |
| - Active Tokenizer Vocab Size: 4096 |
| - Normalization: RMSNorm (eps=1e-5) |
| - Positional Embedding: RoPE (base=10000.0) |
| - Weight Tying: lm_head.weight == token_embedding.weight |
| """ |
| def __init__( |
| self, |
| vocab_size: int = 8192, |
| d_model: int = 576, |
| num_layers: int = 12, |
| num_heads: int = 9, |
| ffn_dim: int = 1536, |
| max_seq_len: int = 1024 |
| ): |
| super().__init__() |
| self.vocab_size = vocab_size |
| self.d_model = d_model |
| self.num_layers = num_layers |
| self.num_heads = num_heads |
| self.ffn_dim = ffn_dim |
| self.max_seq_len = max_seq_len |
|
|
| self.token_embedding = nn.Embedding(vocab_size, d_model) |
| self.rotary_emb = RotaryEmbedding(d_model // num_heads, max_seq_len=max_seq_len) |
|
|
| self.blocks = nn.ModuleList([ |
| TransformerBlock(d_model, num_heads, ffn_dim) |
| for _ in range(num_layers) |
| ]) |
|
|
| self.norm = RMSNorm(d_model) |
| self.lm_head = nn.Linear(d_model, vocab_size, bias=False) |
| self.lm_head.weight = self.token_embedding.weight |
|
|
| self.gradient_checkpointing = False |
| self.apply(self._init_weights) |
|
|
| def _init_weights(self, module): |
| if isinstance(module, nn.Linear): |
| torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) |
| if module.bias is not None: |
| torch.nn.init.zeros_(module.bias) |
| elif isinstance(module, nn.Embedding): |
| torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) |
|
|
| def enable_gradient_checkpointing(self, enabled: bool = True): |
| self.gradient_checkpointing = enabled |
|
|
| def get_num_params(self) -> int: |
| return sum(p.numel() for p in self.parameters()) |
|
|
| def forward( |
| self, |
| input_ids: torch.Tensor, |
| labels: Optional[torch.Tensor] = None, |
| attention_mask: Optional[torch.Tensor] = None |
| ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: |
| B, S = input_ids.shape |
| assert S <= self.max_seq_len, f"Sequence length {S} exceeds maximum context {self.max_seq_len}" |
|
|
| x = self.token_embedding(input_ids) |
| cos, sin = self.rotary_emb(x, S) |
|
|
| for block in self.blocks: |
| if self.gradient_checkpointing and self.training: |
| x = torch.utils.checkpoint.checkpoint( |
| block, x, cos, sin, attention_mask, use_reentrant=False |
| ) |
| else: |
| x = block(x, cos, sin, attention_mask=attention_mask) |
|
|
| x = self.norm(x) |
| logits = self.lm_head(x) |
|
|
| loss = None |
| if labels is not None: |
| shift_logits = logits[..., :-1, :].contiguous() |
| shift_labels = labels[..., 1:].contiguous() |
| loss = F.cross_entropy( |
| shift_logits.view(-1, self.vocab_size).float(), |
| shift_labels.view(-1), |
| ignore_index=-100 |
| ) |
|
|
| return logits, loss |
|
|
|
|
| def create_model() -> TurkishLLMGenerativeV1: |
| """Instantiate and verify the exact 52.5M Generative V1 architecture.""" |
| model = TurkishLLMGenerativeV1() |
| param_count = sum(p.numel() for p in model.parameters()) |
| assert param_count == 52_508_736, f"PARAM COUNT MISMATCH! Expected 52,508,736, got {param_count:,}" |
| return model |
|
|
