| 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) |
|
|
| |
| 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)
|
|
|
