File size: 6,715 Bytes
9b994a0 | 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 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 | 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)
|