| """ |
| src/model.py — GPT-style transformer. You write this. |
| |
| Target architecture (nanoGPT-style): |
| - CausalSelfAttention: multi-head attention with causal mask |
| - MLP: two linear layers with GELU activation |
| - TransformerBlock: attention + MLP with residual connections and LayerNorm |
| - GPT: embedding + N blocks + final projection to vocab |
| |
| Reference: https://github.com/karpathy/nanoGPT/blob/master/model.py |
| Paper: "Attention Is All You Need" (Vaswani et al., 2017) — Sections 3.1–3.4 |
| |
| Start here: |
| 1. Define a GPTConfig dataclass (n_layer, n_head, n_embd, vocab_size, block_size) |
| 2. Write CausalSelfAttention.__init__ and .forward |
| 3. Write MLP |
| 4. Write TransformerBlock |
| 5. Write GPT with token + position embeddings, stack of blocks, LM head |
| |
| Sanity check (paste into a scratch script when done): |
| config = GPTConfig(n_layer=6, n_head=6, n_embd=384, vocab_size=50257, block_size=1024) |
| model = GPT(config) |
| x = torch.randint(0, 50257, (4, 1024)) |
| logits, loss = model(x, x) |
| print(logits.shape) # expect: torch.Size([4, 1024, 50257]) |
| print(sum(p.numel() for p in model.parameters()) / 1e6, "M params") # expect ~85M |
| """ |
|
|
| import math |
| import torch |
| import torch.nn as nn |
| from torch.nn import functional as F |
| from dataclasses import dataclass |
|
|
|
|
| @dataclass |
| class GPTConfig: |
| block_size: int = 1024 |
| vocab_size: int = 50257 |
| n_layer: int = 12 |
| n_head: int = 12 |
| n_embd: int = 768 |
| dropout: float = 0.0 |
| bias: bool = True |
|
|
|
|
| |
| class CausalSelfAttention(nn.Module): |
| def __init__(self, config): |
| super().__init__() |
| assert config.n_embd % config.n_head == 0 |
| self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias) |
| self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias) |
| self.attn_dropout = nn.Dropout(config.dropout) |
| self.resid_dropout = nn.Dropout(config.dropout) |
| self.n_head = config.n_head |
| self.n_embd = config.n_embd |
| self.dropout = config.dropout |
| self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention') |
| if not self.flash: |
| self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size)) |
| .view(1, 1, config.block_size, config.block_size)) |
|
|
| def forward(self, x): |
| B, T, C = x.size() |
| q, k, v = self.c_attn(x).split(self.n_embd, dim=2) |
| k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) |
| q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) |
| v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) |
|
|
| if self.flash: |
| y = torch.nn.functional.scaled_dot_product_attention( |
| q, k, v, attn_mask=None, |
| dropout_p=self.dropout if self.training else 0, |
| is_causal=True) |
| else: |
| att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1))) |
| att = att.masked_fill(self.bias[:,:,: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) |
| y = self.resid_dropout(self.c_proj(y)) |
| return y |
| class MLP(nn.Module): |
|
|
| def __init__(self, config): |
| super().__init__() |
| self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias) |
| self.gelu = nn.GELU() |
| self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias) |
| self.dropout = nn.Dropout(config.dropout) |
| def forward(self, x): |
| x = self.c_fc(x) |
| x = self.gelu(x) |
| x = self.c_proj(x) |
| x = self.dropout(x) |
| return x |
| |
| class Block(nn.Module): |
|
|
| def __init__(self, config): |
| super().__init__() |
| self.ln_1 = nn.LayerNorm(config.n_embd) |
| self.attn = CausalSelfAttention(config) |
| self.ln_2 = nn.LayerNorm(config.n_embd) |
| self.mlp = MLP(config) |
| def forward(self, x): |
| x = x + self.attn(self.ln_1(x)) |
| x = x + self.mlp(self.ln_2(x)) |
| return x |
| |
|
|
| class GPT(nn.Module): |
|
|
| def __init__(self, config): |
| super().__init__() |
| self.transformer = nn.ModuleDict(dict( |
| wte = nn.Embedding(config.vocab_size, config.n_embd), |
| wpe = nn.Embedding(config.block_size, config.n_embd), |
| drop = nn.Dropout(config.dropout), |
| h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]), |
| ln_f = nn.LayerNorm(config.n_embd), |
| )) |
| self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False) |
| def forward(self, idx, targets=None): |
| b, t = idx.size() |
| pos = torch.arange(0, t, dtype=torch.long, device=idx.device) |
|
|
| tok_emb = self.transformer.wte(idx) |
| pos_emb = self.transformer.wpe(pos) |
| x = self.transformer.drop(tok_emb + pos_emb) |
| for block in self.transformer.h: |
| x = block(x) |
| x = self.transformer.ln_f(x) |
|
|
| if targets is not None: |
| logits = self.lm_head(x) |
| loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1)) |
| else: |
| logits = self.lm_head(x[:, [-1], :]) |
| loss = None |
|
|
| return logits, loss |
|
|
| def configure_optimizers(self, weight_decay, learning_rate, betas, device_type): |
| |
| decay = {p for n, p in self.named_parameters() if p.requires_grad and p.dim() >= 2} |
| no_decay = {p for n, p in self.named_parameters() if p.requires_grad and p.dim() < 2} |
| groups = [{"params": list(decay), "weight_decay": weight_decay}, |
| {"params": list(no_decay), "weight_decay": 0.0}] |
| return torch.optim.AdamW(groups, lr=learning_rate, betas=betas) |
