| """Random-initialized decoder: pre-RMSNorm, RoPE, causal SDPA, SwiGLU, tied head.""" |
| from __future__ import annotations |
| import math |
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
| from torch.utils.checkpoint import checkpoint |
| from .config import ModelConfig |
|
|
|
|
| class RMSNorm(nn.Module): |
| def __init__(self, size, epsilon): |
| super().__init__() |
| self.weight = nn.Parameter(torch.ones(size)) |
| self.epsilon = epsilon |
|
|
| def forward(self, x): |
| y = x.float() |
| y = y * torch.rsqrt(y.square().mean(-1, keepdim=True) + self.epsilon) |
| return y.to(x.dtype) * self.weight |
|
|
|
|
| class Attention(nn.Module): |
| def __init__(self, config): |
| super().__init__() |
| self.heads = config.heads |
| self.head_dim = config.hidden_size // config.heads |
| self.qkv = nn.Linear(config.hidden_size, 3*config.hidden_size, bias=False) |
| self.output = nn.Linear(config.hidden_size, config.hidden_size, bias=False) |
| self.dropout = config.dropout |
| inv_freq = 1.0 / (config.rope_theta ** (torch.arange(0, self.head_dim, 2).float() / self.head_dim)) |
| frequencies = torch.outer(torch.arange(config.context_length).float(), inv_freq) |
| self.register_buffer("cos", frequencies.cos()[None, None], persistent=False) |
| self.register_buffer("sin", frequencies.sin()[None, None], persistent=False) |
|
|
| def rotate(self, x): |
| |
| even, odd = x.float()[..., 0::2], x.float()[..., 1::2] |
| cos, sin = self.cos[:, :, :x.shape[-2]], self.sin[:, :, :x.shape[-2]] |
| return torch.stack((even*cos-odd*sin, even*sin+odd*cos), -1).flatten(-2).to(x.dtype) |
|
|
| def forward(self, x): |
| b, t, h = x.shape |
| q, k, v = self.qkv(x).view(b,t,3,self.heads,self.head_dim).permute(2,0,3,1,4).unbind(0) |
| y = F.scaled_dot_product_attention(self.rotate(q), self.rotate(k), v, is_causal=True, |
| dropout_p=self.dropout if self.training else 0.0) |
| return self.output(y.transpose(1,2).contiguous().view(b,t,h)) |
|
|
|
|
| class Block(nn.Module): |
| def __init__(self, config): |
| super().__init__() |
| h, f = config.hidden_size, config.intermediate_size |
| self.attention_norm = RMSNorm(h, config.rms_epsilon) |
| self.attention = Attention(config) |
| self.ffn_norm = RMSNorm(h, config.rms_epsilon) |
| self.gate = nn.Linear(h,f,bias=False) |
| self.up = nn.Linear(h,f,bias=False) |
| self.down = nn.Linear(f,h,bias=False) |
| self.dropout = nn.Dropout(config.dropout) |
|
|
| def forward(self, x): |
| x = x + self.dropout(self.attention(self.attention_norm(x))) |
| z = self.ffn_norm(x) |
| return x + self.dropout(self.down(F.silu(self.gate(z))*self.up(z))) |
|
|
|
|
| class StasisDecoder(nn.Module): |
| def __init__(self, config: ModelConfig): |
| super().__init__() |
| self.config = config |
| self.embedding = nn.Embedding(config.vocab_size, config.hidden_size) |
| self.blocks = nn.ModuleList(Block(config) for _ in range(config.layers)) |
| self.norm = RMSNorm(config.hidden_size, config.rms_epsilon) |
| self.apply(self._init) |
| |
| for block in self.blocks: |
| nn.init.normal_(block.attention.output.weight, std=0.02/math.sqrt(2*config.layers)) |
| nn.init.normal_(block.down.weight, std=0.02/math.sqrt(2*config.layers)) |
|
|
| @staticmethod |
| def _init(module): |
| if isinstance(module, (nn.Linear, nn.Embedding)): |
| nn.init.normal_(module.weight, mean=0.0, std=0.02) |
|
|
| def forward(self, tokens, targets=None, reduction="mean"): |
| if tokens.ndim != 2 or tokens.shape[1] > self.config.context_length: |
| raise ValueError("Expected [batch, sequence] within configured context length") |
| x = self.embedding(tokens) |
| for block in self.blocks: |
| x = checkpoint(block, x, use_reentrant=False) if self.training and self.config.gradient_checkpointing else block(x) |
| logits = F.linear(self.norm(x), self.embedding.weight) |
| loss = None |
| if targets is not None: |
| if targets.shape != tokens.shape: |
| raise ValueError("Targets must match input shape") |
| loss = F.cross_entropy(logits.float().reshape(-1, self.config.vocab_size), targets.reshape(-1), reduction=reduction) |
| return logits, loss |
|
|
| @torch.no_grad() |
| def generate(self, tokens, max_new_tokens=64, temperature=1.0, top_k=50, eos_token_id=None): |
| if max_new_tokens < 0 or temperature <= 0 or top_k < 0: |
| raise ValueError("Invalid sampling controls") |
| was_training = self.training |
| self.eval() |
| try: |
| for _ in range(max_new_tokens): |
| logits, _ = self(tokens[:, -self.config.context_length:]) |
| logits = logits[:, -1].float() / temperature |
| if top_k: |
| threshold = torch.topk(logits, min(top_k, logits.size(-1))).values[:, -1,None] |
| logits = logits.masked_fill(logits < threshold, -float("inf")) |
| nxt = torch.multinomial(F.softmax(logits, dim=-1), 1) |
| tokens = torch.cat([tokens,nxt], dim=1) |
| if eos_token_id is not None and bool((nxt == eos_token_id).all()): |
| break |
| return tokens |
| finally: |
| self.train(was_training) |
|
|
