| """ZENITH decoder-only Transformer (Llama-style: RMSNorm, RoPE, SwiGLU, no biases). |
| |
| Parameter names follow the Llama layout so a later GGUF export is a direct |
| name mapping. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import math |
| from typing import Literal |
|
|
| import torch |
| import torch.nn.functional as F |
| from torch import nn |
|
|
| from zenith.config import ModelConfig |
|
|
|
|
| class RMSNorm(nn.Module): |
| def __init__(self, dim: int, eps: float): |
| super().__init__() |
| self.eps = eps |
| self.weight = nn.Parameter(torch.ones(dim)) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| xf = x.float() |
| normed = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + self.eps) |
| return normed.type_as(x) * self.weight |
|
|
|
|
| def rope_tables(head_dim: int, context_len: int, theta: float) -> tuple[torch.Tensor, torch.Tensor]: |
| inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)) |
| angles = torch.outer(torch.arange(context_len, dtype=torch.float32), inv_freq) |
| angles = torch.cat([angles, angles], dim=-1) |
| return angles.cos(), angles.sin() |
|
|
|
|
| def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: |
| |
| half = x.shape[-1] // 2 |
| rotated = torch.cat([-x[..., half:], x[..., :half]], dim=-1) |
| return (x * cos + rotated * sin).type_as(x) |
|
|
|
|
| class Attention(nn.Module): |
| def __init__(self, cfg: ModelConfig): |
| super().__init__() |
| self.n_heads = cfg.n_heads |
| self.head_dim = cfg.d_model // cfg.n_heads |
| self.q_proj = nn.Linear(cfg.d_model, cfg.d_model, bias=False) |
| self.k_proj = nn.Linear(cfg.d_model, cfg.d_model, bias=False) |
| self.v_proj = nn.Linear(cfg.d_model, cfg.d_model, bias=False) |
| self.out_proj = nn.Linear(cfg.d_model, cfg.d_model, bias=False) |
|
|
| def forward(self, x, cos, sin, impl: str) -> torch.Tensor: |
| B, T, C = x.shape |
|
|
| def heads(t): |
| return t.view(B, T, self.n_heads, self.head_dim).transpose(1, 2) |
|
|
| q = apply_rope(heads(self.q_proj(x)), cos, sin) |
| k = apply_rope(heads(self.k_proj(x)), cos, sin) |
| v = heads(self.v_proj(x)) |
| if impl == "sdpa": |
| y = F.scaled_dot_product_attention(q, k, v, is_causal=True) |
| else: |
| scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim) |
| mask = torch.ones(T, T, dtype=torch.bool, device=x.device).tril() |
| scores = scores.masked_fill(~mask, float("-inf")) |
| y = scores.softmax(dim=-1) @ v |
| return self.out_proj(y.transpose(1, 2).reshape(B, T, C)) |
|
|
|
|
| class SwiGLU(nn.Module): |
| def __init__(self, cfg: ModelConfig): |
| super().__init__() |
| self.gate_proj = nn.Linear(cfg.d_model, cfg.ffn_hidden, bias=False) |
| self.up_proj = nn.Linear(cfg.d_model, cfg.ffn_hidden, bias=False) |
| self.down_proj = nn.Linear(cfg.ffn_hidden, cfg.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 Block(nn.Module): |
| def __init__(self, cfg: ModelConfig): |
| super().__init__() |
| self.attn_norm = RMSNorm(cfg.d_model, cfg.norm_eps) |
| self.attn = Attention(cfg) |
| self.mlp_norm = RMSNorm(cfg.d_model, cfg.norm_eps) |
| self.mlp = SwiGLU(cfg) |
|
|
| def forward(self, x, cos, sin, impl: str) -> torch.Tensor: |
| x = x + self.attn(self.attn_norm(x), cos, sin, impl) |
| return x + self.mlp(self.mlp_norm(x)) |
|
|
|
|
| class ZenithLM(nn.Module): |
| def __init__(self, cfg: ModelConfig): |
| super().__init__() |
| if cfg.d_model % cfg.n_heads: |
| raise ValueError("d_model must be divisible by n_heads") |
| self.cfg = cfg |
| self.attn_impl: Literal["sdpa", "reference"] = "sdpa" |
| self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.d_model) |
| self.blocks = nn.ModuleList(Block(cfg) for _ in range(cfg.n_layers)) |
| self.final_norm = RMSNorm(cfg.d_model, cfg.norm_eps) |
| cos, sin = rope_tables(cfg.d_model // cfg.n_heads, cfg.context_len, cfg.rope_theta) |
| self.register_buffer("rope_cos", cos, persistent=False) |
| self.register_buffer("rope_sin", sin, persistent=False) |
| self._init_weights() |
|
|
| def _init_weights(self) -> None: |
| std = self.cfg.init_std |
| for module in self.modules(): |
| if isinstance(module, (nn.Linear, nn.Embedding)): |
| nn.init.normal_(module.weight, mean=0.0, std=std) |
| scaled = std / math.sqrt(2 * self.cfg.n_layers) |
| for block in self.blocks: |
| nn.init.normal_(block.attn.out_proj.weight, mean=0.0, std=scaled) |
| nn.init.normal_(block.mlp.down_proj.weight, mean=0.0, std=scaled) |
|
|
| def lm_head_weight(self) -> torch.Tensor: |
| return self.tok_emb.weight |
|
|
| def num_params(self) -> int: |
| return sum(p.numel() for p in self.parameters()) |
|
|
| def forward(self, idx: torch.Tensor, targets: torch.Tensor | None = None): |
| T = idx.shape[1] |
| if T > self.cfg.context_len: |
| raise ValueError(f"sequence length {T} exceeds context_len {self.cfg.context_len}") |
| cos, sin = self.rope_cos[:T], self.rope_sin[:T] |
| x = self.tok_emb(idx) |
| for block in self.blocks: |
| x = block(x, cos, sin, self.attn_impl) |
| logits = F.linear(self.final_norm(x), self.lm_head_weight()) |
| loss = None |
| if targets is not None: |
| loss = F.cross_entropy(logits.view(-1, logits.size(-1)).float(), targets.reshape(-1), ignore_index=-100) |
| return logits, loss |
|
|
