from __future__ import annotations from dataclasses import asdict, dataclass import math import torch from torch import nn from torch.nn import functional as F @dataclass class ModelConfig: vocab_size: int block_size: int = 256 n_layer: int = 3 n_head: int = 4 n_embd: int = 96 ffn_dim: int = 256 rope_theta: float = 10000.0 dropout: float = 0.0 tie_embeddings: bool = True def validate(self) -> None: if self.n_embd % self.n_head: raise ValueError("n_embd must be divisible by n_head") if (self.n_embd // self.n_head) % 2: raise ValueError("RoPE requires an even attention head dimension") if self.block_size < 2: raise ValueError("block_size must be at least 2") class RMSNorm(nn.Module): def __init__(self, width: int, eps: float = 1e-6): super().__init__() self.weight = nn.Parameter(torch.ones(width)) self.eps = eps def forward(self, x: torch.Tensor) -> torch.Tensor: source_dtype = x.dtype normalized = x.float() * torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) + self.eps) return normalized.to(source_dtype) * self.weight.to(source_dtype) class RotaryEmbedding(nn.Module): def __init__(self, head_dim: int, theta: float): super().__init__() inverse_frequencies = 1.0 / (theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)) self.register_buffer("inverse_frequencies", inverse_frequencies, persistent=False) def forward(self, sequence_length: int, device: torch.device, dtype: torch.dtype) -> tuple[torch.Tensor, torch.Tensor]: positions = torch.arange(sequence_length, device=device, dtype=torch.float32) angles = torch.outer(positions, self.inverse_frequencies.to(device=device)) return angles.cos().to(dtype=dtype)[None, None, :, :], angles.sin().to(dtype=dtype)[None, None, :, :] def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: even = x[..., 0::2] odd = x[..., 1::2] result = torch.empty_like(x) result[..., 0::2] = even * cos - odd * sin result[..., 1::2] = even * sin + odd * cos return result class SelfAttention(nn.Module): def __init__(self, config: ModelConfig): super().__init__() self.n_head = config.n_head self.head_dim = config.n_embd // config.n_head self.qkv = nn.Linear(config.n_embd, 3 * config.n_embd, bias=False) self.output = nn.Linear(config.n_embd, config.n_embd, bias=False) self.dropout = config.dropout def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: batch, sequence_length, width = x.shape qkv = self.qkv(x).view(batch, sequence_length, 3, self.n_head, self.head_dim) query, key, value = qkv.unbind(dim=2) query = query.transpose(1, 2) key = key.transpose(1, 2) value = value.transpose(1, 2) query = apply_rope(query, cos, sin) key = apply_rope(key, cos, sin) # PyTorch SDPA selects the available fused backend (including Flash when supported). attended = F.scaled_dot_product_attention( query, key, value, attn_mask=None, dropout_p=self.dropout if self.training else 0.0, is_causal=True, ) attended = attended.transpose(1, 2).contiguous().view(batch, sequence_length, width) return self.output(attended) class SwiGLU(nn.Module): def __init__(self, width: int, hidden_width: int): super().__init__() self.gate = nn.Linear(width, hidden_width, bias=False) self.up = nn.Linear(width, hidden_width, bias=False) self.down = nn.Linear(hidden_width, width, bias=False) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.down(F.silu(self.gate(x)) * self.up(x)) class TransformerBlock(nn.Module): def __init__(self, config: ModelConfig): super().__init__() self.attention_norm = RMSNorm(config.n_embd) self.attention = SelfAttention(config) self.feed_forward_norm = RMSNorm(config.n_embd) self.feed_forward = SwiGLU(config.n_embd, config.ffn_dim) self.residual_dropout = nn.Dropout(config.dropout) def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: x = x + self.residual_dropout(self.attention(self.attention_norm(x), cos, sin)) x = x + self.residual_dropout(self.feed_forward(self.feed_forward_norm(x))) return x class DivCentererLM(nn.Module): def __init__(self, config: ModelConfig): super().__init__() config.validate() self.config = config self.token_embedding = nn.Embedding(config.vocab_size, config.n_embd) self.rotary = RotaryEmbedding(config.n_embd // config.n_head, config.rope_theta) self.layers = nn.ModuleList([TransformerBlock(config) for _ in range(config.n_layer)]) self.final_norm = RMSNorm(config.n_embd) self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False) if config.tie_embeddings: self.lm_head.weight = self.token_embedding.weight self.apply(self._initialize_weights) @staticmethod def _initialize_weights(module: nn.Module) -> None: if isinstance(module, nn.Linear): nn.init.normal_(module.weight, mean=0.0, std=0.02) 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=0.02) def forward(self, input_ids: torch.Tensor, labels: torch.Tensor | None = None) -> tuple[torch.Tensor, torch.Tensor | None]: _, sequence_length = input_ids.shape if sequence_length > self.config.block_size: raise ValueError(f"Sequence length {sequence_length} exceeds block_size={self.config.block_size}") x = self.token_embedding(input_ids) cos, sin = self.rotary(sequence_length, input_ids.device, x.dtype) for layer in self.layers: x = layer(x, cos, sin) logits = self.lm_head(self.final_norm(x)) loss = None if labels is not None: loss = F.cross_entropy( logits.float().reshape(-1, self.config.vocab_size), labels.reshape(-1), ignore_index=-100, ) return logits, loss def parameter_count(self) -> int: return sum(parameter.numel() for parameter in self.parameters()) def config_dict(self) -> dict: return asdict(self.config)