| from __future__ import annotations |
|
|
| from typing import Optional |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from torch import Tensor |
|
|
| from config import ModelConfig |
|
|
|
|
| class RMSNorm(nn.Module): |
| def __init__(self, dim: int, eps: float = 1.0e-5): |
| super().__init__() |
| self.weight = nn.Parameter(torch.ones(dim)) |
| self.eps = eps |
|
|
| def forward(self, x: Tensor) -> Tensor: |
| dtype = x.dtype |
| x32 = x.float() |
| rms = x32.pow(2).mean(dim=-1, keepdim=True).add_(self.eps).rsqrt_() |
| out = (x32 * rms).to(dtype) |
| return out * self.weight.to(dtype) |
|
|
|
|
| class RotaryEmbedding(nn.Module): |
| def __init__(self, head_dim: int, max_seq_len: int, theta: float = 10_000.0): |
| super().__init__() |
| self.head_dim = head_dim |
| self.max_seq_len = max_seq_len |
| self.theta = theta |
| self._cached_len: int = 0 |
| self._cos_cache: Optional[Tensor] = None |
| self._sin_cache: Optional[Tensor] = None |
|
|
| def _build_cache(self, seq_len: int, device, dtype): |
| inv_freq = 1.0 / ( |
| self.theta ** (torch.arange(0, self.head_dim, 2, dtype=torch.float32, device=device) / self.head_dim) |
| ) |
| t = torch.arange(seq_len, dtype=torch.float32, device=device) |
| freqs = torch.outer(t, inv_freq) |
| emb = torch.cat([freqs, freqs], dim=-1) |
| self._cos_cache = emb.cos().to(dtype) |
| self._sin_cache = emb.sin().to(dtype) |
| self._cached_len = seq_len |
|
|
| def forward(self, seq_len: int, device, dtype) -> tuple[Tensor, Tensor]: |
| if ( |
| self._cos_cache is None |
| or seq_len > self._cached_len |
| or self._cos_cache.device != device |
| or self._cos_cache.dtype != dtype |
| ): |
| self._build_cache(max(seq_len, self.max_seq_len), device, dtype) |
| return self._cos_cache[:seq_len], self._sin_cache[:seq_len] |
|
|
|
|
| def _rotate_half(x: Tensor) -> Tensor: |
| x1, x2 = x.chunk(2, dim=-1) |
| return torch.cat([-x2, x1], dim=-1) |
|
|
|
|
| def apply_rope(q: Tensor, k: Tensor, cos: Tensor, sin: Tensor) -> tuple[Tensor, Tensor]: |
| cos = cos.unsqueeze(0).unsqueeze(0) |
| sin = sin.unsqueeze(0).unsqueeze(0) |
| q_rot = (q * cos) + (_rotate_half(q) * sin) |
| k_rot = (k * cos) + (_rotate_half(k) * sin) |
| return q_rot, k_rot |
|
|
|
|
| class Attention(nn.Module): |
| def __init__(self, cfg: ModelConfig, layer_idx: int): |
| super().__init__() |
| self.cfg = cfg |
| self.layer_idx = layer_idx |
| self.num_heads = cfg.num_heads |
| self.num_kv_heads = cfg.num_kv_heads |
| self.head_dim = cfg.head_dim |
| self.kv_groups = cfg.kv_groups |
| self.scale = self.head_dim ** -0.5 |
|
|
| h, hd = cfg.hidden_size, self.head_dim |
| self.q_proj = nn.Linear(h, self.num_heads * hd, bias=False) |
| self.k_proj = nn.Linear(h, self.num_kv_heads * hd, bias=False) |
| self.v_proj = nn.Linear(h, self.num_kv_heads * hd, bias=False) |
| self.o_proj = nn.Linear(self.num_heads * hd, h, bias=False) |
|
|
| if cfg.qk_norm: |
| self.q_norm = RMSNorm(hd, eps=cfg.rms_norm_eps) |
| self.k_norm = RMSNorm(hd, eps=cfg.rms_norm_eps) |
| else: |
| self.q_norm = nn.Identity() |
| self.k_norm = nn.Identity() |
|
|
| self.attn_dropout = cfg.attn_dropout |
|
|
| def forward(self, x: Tensor, cos: Tensor, sin: Tensor) -> Tensor: |
| B, S, _ = x.shape |
| q = self.q_proj(x).view(B, S, self.num_heads, self.head_dim) |
| k = self.k_proj(x).view(B, S, self.num_kv_heads, self.head_dim) |
| v = self.v_proj(x).view(B, S, self.num_kv_heads, self.head_dim) |
|
|
| q = self.q_norm(q) |
| k = self.k_norm(k) |
|
|
| q = q.transpose(1, 2) |
| k = k.transpose(1, 2) |
| v = v.transpose(1, 2) |
|
|
| q, k = apply_rope(q, k, cos, sin) |
|
|
| if self.kv_groups > 1: |
| k = k.repeat_interleave(self.kv_groups, dim=1) |
| v = v.repeat_interleave(self.kv_groups, dim=1) |
|
|
| out = F.scaled_dot_product_attention( |
| q, k, v, |
| attn_mask=None, |
| dropout_p=self.attn_dropout if self.training else 0.0, |
| is_causal=True, |
| ) |
| out = out.transpose(1, 2).contiguous().view(B, S, self.num_heads * self.head_dim) |
| return self.o_proj(out) |
|
|
|
|
| class SwiGLU(nn.Module): |
| def __init__(self, cfg: ModelConfig): |
| super().__init__() |
| h, i = cfg.hidden_size, cfg.intermediate_size |
| self.gate_proj = nn.Linear(h, i, bias=False) |
| self.up_proj = nn.Linear(h, i, bias=False) |
| self.down_proj = nn.Linear(i, h, bias=False) |
|
|
| def forward(self, x: Tensor) -> Tensor: |
| return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) |
|
|
|
|
| class Block(nn.Module): |
| def __init__(self, cfg: ModelConfig, layer_idx: int): |
| super().__init__() |
| self.input_norm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps) |
| self.attn = Attention(cfg, layer_idx) |
| self.post_attn_norm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps) |
| self.mlp = SwiGLU(cfg) |
| self.resid_drop = nn.Dropout(cfg.resid_dropout) if cfg.resid_dropout > 0 else nn.Identity() |
|
|
| def forward(self, x: Tensor, cos: Tensor, sin: Tensor) -> Tensor: |
| x = x + self.resid_drop(self.attn(self.input_norm(x), cos, sin)) |
| x = x + self.resid_drop(self.mlp(self.post_attn_norm(x))) |
| return x |
|
|
|
|
| class MarulLLM(nn.Module): |
| def __init__(self, cfg: ModelConfig): |
| super().__init__() |
| self.cfg = cfg |
|
|
| self.embed_tokens = nn.Embedding(cfg.vocab_size, cfg.hidden_size) |
| self.rotary = RotaryEmbedding(cfg.head_dim, cfg.max_seq_len, cfg.rope_theta) |
| self.layers = nn.ModuleList(Block(cfg, i) for i in range(cfg.num_layers)) |
| self.final_norm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps) |
|
|
| if cfg.tie_word_embeddings: |
| self.lm_head = None |
| else: |
| self.lm_head = nn.Linear(cfg.hidden_size, cfg.vocab_size, bias=False) |
|
|
| self.apply(self._init_weights) |
| self._scale_residual_inits() |
|
|
| self.num_params = sum(p.numel() for p in self.parameters()) |
| self.num_params_trainable = sum(p.numel() for p in self.parameters() if p.requires_grad) |
| embed_params = cfg.vocab_size * cfg.hidden_size |
| self.num_params_non_embed = self.num_params - embed_params |
|
|
| def _init_weights(self, module: nn.Module) -> None: |
| std = self.cfg.initializer_range |
| if isinstance(module, nn.Linear): |
| nn.init.normal_(module.weight, mean=0.0, std=std) |
| if module.bias is not None: |
| nn.init.zeros_(module.bias) |
| elif isinstance(module, nn.Embedding): |
| nn.init.normal_(module.weight, mean=0.0, std=std) |
|
|
| def _scale_residual_inits(self) -> None: |
| scale = (2 * self.cfg.num_layers) ** -0.5 |
| with torch.no_grad(): |
| for block in self.layers: |
| block.attn.o_proj.weight.mul_(scale) |
| block.mlp.down_proj.weight.mul_(scale) |
|
|
| def forward( |
| self, |
| input_ids: Tensor, |
| targets: Optional[Tensor] = None, |
| return_logits: bool = True, |
| ) -> tuple[Optional[Tensor], Optional[Tensor]]: |
| B, S = input_ids.shape |
| assert S <= self.cfg.max_seq_len, ( |
| f"dizi uzunluğu {S}, modelin bağlam sınırı {self.cfg.max_seq_len}") |
|
|
| x = self.embed_tokens(input_ids) |
| cos, sin = self.rotary(S, x.device, x.dtype) |
|
|
| for block in self.layers: |
| x = block(x, cos, sin) |
| x = self.final_norm(x) |
|
|
| if self.cfg.tie_word_embeddings: |
| logits = F.linear(x, self.embed_tokens.weight) |
| else: |
| logits = self.lm_head(x) |
|
|
| loss: Optional[Tensor] = None |
| if targets is not None: |
| flat_logits = logits.view(-1, logits.size(-1)) |
| flat_targets = targets.view(-1) |
| ce = F.cross_entropy( |
| flat_logits, flat_targets, ignore_index=-100, reduction="mean" |
| ) |
| loss = ce |
| if self.cfg.z_loss_coef > 0: |
| mask = flat_targets != -100 |
| lse = torch.logsumexp(flat_logits, dim=-1) |
| if mask.any(): |
| z = (lse[mask].float().pow(2)).mean() |
| loss = loss + self.cfg.z_loss_coef * z |
|
|
| if not return_logits and targets is not None: |
| logits = None |
| return logits, loss |
|
|
| @torch.no_grad() |
| def generate( |
| self, |
| input_ids: Tensor, |
| max_new_tokens: int = 128, |
| temperature: float = 0.6, |
| top_k: int = 40, |
| top_p: float = 0.88, |
| repetition_penalty: float = 1.20, |
| no_repeat_ngram_size: int = 4, |
| min_p: float = 0.05, |
| eos_token_id: Optional[int] = None, |
| ) -> Tensor: |
| self.eval() |
| eos = eos_token_id if eos_token_id is not None else self.cfg.eos_token_id |
| out = input_ids.clone() |
| device = out.device |
|
|
| for _ in range(max_new_tokens): |
| ctx = out[:, -self.cfg.max_seq_len:] |
| logits, _ = self.forward(ctx) |
| logits = logits[:, -1, :].float() |
|
|
| if repetition_penalty is not None and repetition_penalty != 1.0: |
| for b in range(out.size(0)): |
| seen = torch.unique(out[b]) |
| vals = logits[b, seen] |
| logits[b, seen] = torch.where( |
| vals > 0, vals / repetition_penalty, vals * repetition_penalty |
| ) |
|
|
| if no_repeat_ngram_size and no_repeat_ngram_size > 0: |
| n = no_repeat_ngram_size |
| if out.size(1) >= n - 1: |
| for b in range(out.size(0)): |
| seq = out[b].tolist() |
| ngrams: dict = {} |
| for i in range(len(seq) - n + 1): |
| prefix = tuple(seq[i : i + n - 1]) |
| ngrams.setdefault(prefix, set()).add(seq[i + n - 1]) |
| curr = tuple(seq[-(n - 1):]) |
| if curr in ngrams: |
| banned = torch.tensor(list(ngrams[curr]), device=device, dtype=torch.long) |
| logits[b, banned] = float("-inf") |
|
|
| if temperature is not None and temperature != 1.0: |
| logits = logits / max(temperature, 1.0e-6) |
|
|
| if top_k is not None and top_k > 0: |
| v, _ = torch.topk(logits, k=min(top_k, logits.size(-1))) |
| logits[logits < v[:, -1:]] = float("-inf") |
|
|
| if min_p is not None and min_p > 0.0: |
| probs_tmp = F.softmax(logits, dim=-1) |
| max_probs, _ = probs_tmp.max(dim=-1, keepdim=True) |
| logits = logits.masked_fill(probs_tmp < (max_probs * min_p), float("-inf")) |
|
|
| if top_p is not None and 0.0 < top_p < 1.0: |
| sorted_logits, sorted_idx = torch.sort(logits, descending=True, dim=-1) |
| probs = F.softmax(sorted_logits, dim=-1) |
| cumprobs = probs.cumsum(dim=-1) |
| mask = cumprobs > top_p |
| mask[..., 1:] = mask[..., :-1].clone() |
| mask[..., 0] = False |
| sorted_logits = sorted_logits.masked_fill(mask, float("-inf")) |
| logits = torch.full_like(logits, float("-inf")).scatter(-1, sorted_idx, sorted_logits) |
|
|
| probs = F.softmax(logits, dim=-1) |
| next_token = torch.multinomial(probs, num_samples=1) |
| out = torch.cat([out, next_token], dim=1) |
|
|
| if eos is not None and (next_token == eos).all(): |
| break |
| return out |
|
|
