| """ |
| Stage 2: MODEL ARCHITECTURE. |
| |
| A from-scratch GPT-style decoder-only transformer, built directly in PyTorch |
| (no pretrained weights, no transformers.AutoModel -- this is the part that |
| makes "built an LLM from scratch" an honest claim): |
| - token + positional embeddings |
| - N transformer blocks, each with: multi-head causal self-attention, |
| a feed-forward MLP, and layer norm (pre-norm, like GPT-2) |
| - a final layer norm + linear head tied to the token embedding weights |
| |
| Sizes (see PRESETS below): |
| - "tiny" ~ a few M params -> for a fast CPU dry-run / sanity check |
| - "small" ~ 30M params -> trains in well under an hour on an RTX 3060 |
| - "medium" ~ 125M params -> GPT-1 scale, still fits an RTX 3060 12GB |
| """ |
| import math |
| from dataclasses import dataclass |
|
|
| import torch |
| import torch.nn as nn |
| from torch.nn import functional as F |
|
|
| PRESETS = { |
| "tiny": dict(n_layer=4, n_head=4, n_embd=128, block_size=128), |
| "small": dict(n_layer=6, n_head=6, n_embd=384, block_size=256), |
| "medium": dict(n_layer=12, n_head=12, n_embd=768, block_size=512), |
| } |
|
|
|
|
| @dataclass |
| class GPTConfig: |
| vocab_size: int = 8000 |
| block_size: int = 256 |
| n_layer: int = 6 |
| n_head: int = 6 |
| n_embd: int = 384 |
| dropout: float = 0.1 |
|
|
|
|
| class CausalSelfAttention(nn.Module): |
| def __init__(self, config: GPTConfig): |
| super().__init__() |
| assert config.n_embd % config.n_head == 0 |
| 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) |
| self.proj = nn.Linear(config.n_embd, config.n_embd) |
| self.attn_dropout = nn.Dropout(config.dropout) |
| self.resid_dropout = nn.Dropout(config.dropout) |
| mask = torch.tril(torch.ones(config.block_size, config.block_size)) |
| self.register_buffer("mask", mask.view(1, 1, config.block_size, config.block_size)) |
|
|
| def forward(self, x): |
| B, T, C = x.shape |
| qkv = self.qkv(x) |
| q, k, v = qkv.split(C, dim=2) |
| q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2) |
| k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2) |
| v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2) |
|
|
| att = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim) |
| att = att.masked_fill(self.mask[:, :, :T, :T] == 0, float("-inf")) |
| att = F.softmax(att, dim=-1) |
| att = self.attn_dropout(att) |
| y = att @ v |
| y = y.transpose(1, 2).contiguous().view(B, T, C) |
| return self.resid_dropout(self.proj(y)) |
|
|
|
|
| class MLP(nn.Module): |
| def __init__(self, config: GPTConfig): |
| super().__init__() |
| self.fc = nn.Linear(config.n_embd, 4 * config.n_embd) |
| self.proj = nn.Linear(4 * config.n_embd, config.n_embd) |
| self.dropout = nn.Dropout(config.dropout) |
|
|
| def forward(self, x): |
| return self.dropout(self.proj(F.gelu(self.fc(x)))) |
|
|
|
|
| class Block(nn.Module): |
| def __init__(self, config: GPTConfig): |
| super().__init__() |
| self.ln1 = nn.LayerNorm(config.n_embd) |
| self.attn = CausalSelfAttention(config) |
| self.ln2 = nn.LayerNorm(config.n_embd) |
| self.mlp = MLP(config) |
|
|
| def forward(self, x): |
| x = x + self.attn(self.ln1(x)) |
| x = x + self.mlp(self.ln2(x)) |
| return x |
|
|
|
|
| class GPT(nn.Module): |
| def __init__(self, config: GPTConfig): |
| super().__init__() |
| self.config = config |
| self.tok_emb = nn.Embedding(config.vocab_size, config.n_embd) |
| self.pos_emb = nn.Embedding(config.block_size, config.n_embd) |
| self.drop = nn.Dropout(config.dropout) |
| self.blocks = nn.ModuleList([Block(config) for _ in range(config.n_layer)]) |
| self.ln_f = nn.LayerNorm(config.n_embd) |
| self.head = nn.Linear(config.n_embd, config.vocab_size, bias=False) |
| self.head.weight = self.tok_emb.weight |
|
|
| self.apply(self._init_weights) |
| n_params = sum(p.numel() for p in self.parameters()) |
| print(f"GPT initialized: {n_params/1e6:.2f}M parameters " |
| f"(n_layer={config.n_layer}, n_head={config.n_head}, " |
| f"n_embd={config.n_embd}, vocab_size={config.vocab_size}, " |
| f"block_size={config.block_size})") |
|
|
| def _init_weights(self, module): |
| 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, idx, targets=None): |
| B, T = idx.shape |
| assert T <= self.config.block_size, "sequence longer than block_size" |
| pos = torch.arange(0, T, device=idx.device).unsqueeze(0) |
| x = self.drop(self.tok_emb(idx) + self.pos_emb(pos)) |
| for block in self.blocks: |
| x = block(x) |
| x = self.ln_f(x) |
| logits = self.head(x) |
|
|
| loss = None |
| if targets is not None: |
| loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1) |
| return logits, loss |
|
|
| @torch.no_grad() |
| def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None, top_p=None): |
| self.eval() |
| for _ in range(max_new_tokens): |
| idx_cond = idx[:, -self.config.block_size:] |
| logits, _ = self(idx_cond) |
| logits = logits[:, -1, :] / max(temperature, 1e-5) |
|
|
| if top_k is not None: |
| v, _ = torch.topk(logits, min(top_k, logits.size(-1))) |
| logits[logits < v[:, [-1]]] = float("-inf") |
|
|
| if top_p is not None: |
| sorted_logits, sorted_idx = torch.sort(logits, descending=True) |
| probs = F.softmax(sorted_logits, dim=-1) |
| cum_probs = torch.cumsum(probs, dim=-1) |
| remove = cum_probs > top_p |
| remove[:, 1:] = remove[:, :-1].clone() |
| remove[:, 0] = False |
| sorted_logits[remove] = float("-inf") |
| logits = torch.full_like(logits, float("-inf")).scatter(1, sorted_idx, sorted_logits) |
|
|
| probs = F.softmax(logits, dim=-1) |
| next_id = torch.multinomial(probs, num_samples=1) |
| idx = torch.cat([idx, next_id], dim=1) |
| return idx |
|
|
