File size: 5,617 Bytes
0cf8154 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 | """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) # rotate-half layout
return angles.cos(), angles.sin()
def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
# x: (B, H, T, D); cos/sin: (T, D)
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
|