File size: 5,471 Bytes
093306a | 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 | """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):
# Interleaved real/imaginary coordinates; RoPE computed in float32.
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)
# Scale residual projections with depth, after common random initialization.
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)
|