| """ |
| Stage 2 — ARCHITECTURE (lecture: Architecture, covered briefly because it is "not the hardest part"). |
| |
| Llama-style decoder-only Transformer: |
| - RMSNorm (pre-norm) : cheaper than LayerNorm, stable |
| - RoPE : position via query/key rotation instead of absolute position embeddings |
| - SwiGLU FFN : gated MLP, standard in modern models |
| - Causal self-attention via F.scaled_dot_product_attention (uses FlashAttention when available) |
| - Weight tying: embedding <-> output head |
| """ |
| from __future__ import annotations |
|
|
| import math |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| from config import ModelConfig |
|
|
|
|
| class RMSNorm(nn.Module): |
| def __init__(self, dim: int, eps: float = 1e-6): |
| super().__init__() |
| self.eps = eps |
| self.weight = nn.Parameter(torch.ones(dim)) |
|
|
| def forward(self, x): |
| x_f = x.float() |
| x_f = x_f * torch.rsqrt(x_f.pow(2).mean(-1, keepdim=True) + self.eps) |
| return x_f.type_as(x) * self.weight |
|
|
|
|
| def rope_cache(seq_len: int, head_dim: int, base: float = 10000.0): |
| inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim)) |
| t = torch.arange(seq_len).float() |
| freqs = torch.outer(t, inv_freq) |
| return freqs.cos(), freqs.sin() |
|
|
|
|
| def apply_rope(x, cos, sin): |
| |
| x1, x2 = x[..., 0::2], x[..., 1::2] |
| cos, sin = cos[None, None], sin[None, None] |
| out = torch.stack((x1 * cos - x2 * sin, x1 * sin + x2 * cos), dim=-1) |
| return out.flatten(-2).type_as(x) |
|
|
|
|
| class Attention(nn.Module): |
| def __init__(self, cfg: ModelConfig): |
| super().__init__() |
| assert cfg.d_model % cfg.n_heads == 0 |
| self.n_heads = cfg.n_heads |
| self.head_dim = cfg.d_model // cfg.n_heads |
| self.qkv = nn.Linear(cfg.d_model, 3 * cfg.d_model, bias=False) |
| self.proj = nn.Linear(cfg.d_model, cfg.d_model, bias=False) |
| self.dropout = cfg.dropout |
|
|
| def forward(self, x, cos, sin): |
| B, T, C = x.shape |
| q, k, v = self.qkv(x).split(C, dim=-1) |
| q, k, v = (t.view(B, T, self.n_heads, self.head_dim).transpose(1, 2) for t in (q, k, v)) |
| q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin) |
| y = F.scaled_dot_product_attention( |
| q, k, v, is_causal=True, dropout_p=self.dropout if self.training else 0.0 |
| ) |
| return self.proj(y.transpose(1, 2).contiguous().view(B, T, C)) |
|
|
|
|
| class SwiGLU(nn.Module): |
| def __init__(self, cfg: ModelConfig): |
| super().__init__() |
| hidden = int(cfg.ffn_mult * cfg.d_model) |
| hidden = 64 * ((hidden + 63) // 64) |
| self.w_gate = nn.Linear(cfg.d_model, hidden, bias=False) |
| self.w_up = nn.Linear(cfg.d_model, hidden, bias=False) |
| self.w_down = nn.Linear(hidden, cfg.d_model, bias=False) |
|
|
| def forward(self, x): |
| return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x)) |
|
|
|
|
| class Block(nn.Module): |
| def __init__(self, cfg: ModelConfig): |
| super().__init__() |
| self.norm1 = RMSNorm(cfg.d_model) |
| self.attn = Attention(cfg) |
| self.norm2 = RMSNorm(cfg.d_model) |
| self.ffn = SwiGLU(cfg) |
| self.drop = nn.Dropout(cfg.dropout) |
|
|
| def forward(self, x, cos, sin): |
| x = x + self.drop(self.attn(self.norm1(x), cos, sin)) |
| x = x + self.drop(self.ffn(self.norm2(x))) |
| return x |
|
|
|
|
| class GPT(nn.Module): |
| def __init__(self, cfg: ModelConfig): |
| super().__init__() |
| self.cfg = cfg |
| self.embed = nn.Embedding(cfg.vocab_size, cfg.d_model) |
| self.blocks = nn.ModuleList(Block(cfg) for _ in range(cfg.n_layers)) |
| self.norm = RMSNorm(cfg.d_model) |
| self.head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False) |
| if cfg.tie_embeddings: |
| self.head.weight = self.embed.weight |
| cos, sin = rope_cache(cfg.context_len, cfg.d_model // cfg.n_heads) |
| self.register_buffer("rope_cos", cos, persistent=False) |
| self.register_buffer("rope_sin", sin, persistent=False) |
| self.apply(self._init) |
| |
| for n, p in self.named_parameters(): |
| if n.endswith("proj.weight") or n.endswith("w_down.weight"): |
| nn.init.normal_(p, 0.0, 0.02 / math.sqrt(2 * cfg.n_layers)) |
|
|
| @staticmethod |
| def _init(m): |
| if isinstance(m, nn.Linear): |
| nn.init.normal_(m.weight, 0.0, 0.02) |
| elif isinstance(m, nn.Embedding): |
| nn.init.normal_(m.weight, 0.0, 0.02) |
|
|
| def num_params(self, non_embedding: bool = True) -> int: |
| n = sum(p.numel() for p in self.parameters()) |
| if non_embedding: |
| n -= self.embed.weight.numel() |
| return n |
|
|
| def forward(self, idx, targets=None, loss_mask=None): |
| B, T = idx.shape |
| assert T <= self.cfg.context_len, f"length {T} > context_len {self.cfg.context_len}" |
| x = self.embed(idx) |
| cos, sin = self.rope_cos[:T], self.rope_sin[:T] |
| for blk in self.blocks: |
| x = blk(x, cos, sin) |
| logits = self.head(self.norm(x)) |
| if targets is None: |
| return logits, None |
| loss = F.cross_entropy(logits.float().view(-1, logits.size(-1)), targets.view(-1), reduction="none") |
| if loss_mask is not None: |
| m = loss_mask.view(-1).float() |
| loss = (loss * m).sum() / m.sum().clamp(min=1) |
| else: |
| loss = loss.mean() |
| return logits, loss |
|
|
| def configure_optimizer(self, lr: float, weight_decay: float, betas=(0.9, 0.95)): |
| |
| decay = [p for p in self.parameters() if p.requires_grad and p.dim() >= 2] |
| no_decay = [p for p in self.parameters() if p.requires_grad and p.dim() < 2] |
| groups = [{"params": decay, "weight_decay": weight_decay}, {"params": no_decay, "weight_decay": 0.0}] |
| fused = torch.cuda.is_available() |
| return torch.optim.AdamW(groups, lr=lr, betas=betas, fused=fused) |
|
|
| @torch.no_grad() |
| def generate(self, idx, max_new_tokens: int, temperature: float = 1.0, top_k: int | None = 50, |
| eos_id: int | None = None, valid_vocab: int | None = None): |
| """ |
| Autoregressive sampling for a batch of equal-length prompts. |
| valid_vocab: tokenizer vocab size — ids above it (spare embedding rows) are never sampled. |
| Returns (new_tokens [B, <=max_new], finished [B]) — tokens after EOS are filled with eos_id. |
| """ |
| B = idx.size(0) |
| finished = torch.zeros(B, dtype=torch.bool, device=idx.device) |
| new = [] |
| for _ in range(max_new_tokens): |
| ctx = idx[:, -self.cfg.context_len:] |
| logits, _ = self(ctx) |
| logits = logits[:, -1, :].float() / max(temperature, 1e-5) |
| if valid_vocab is not None and valid_vocab < logits.size(-1): |
| logits[:, valid_vocab:] = -float("inf") |
| if top_k: |
| v, _ = torch.topk(logits, min(top_k, logits.size(-1))) |
| logits[logits < v[:, [-1]]] = -float("inf") |
| nxt = torch.multinomial(F.softmax(logits, dim=-1), 1) |
| if eos_id is not None: |
| nxt = torch.where(finished[:, None], torch.full_like(nxt, eos_id), nxt) |
| finished |= nxt.squeeze(1) == eos_id |
| idx = torch.cat([idx, nxt], dim=1) |
| new.append(nxt) |
| if eos_id is not None and finished.all(): |
| break |
| return torch.cat(new, dim=1), finished |
|
|
