""" Tiny GPT-style transformer (~30M params target). Config: - 6 layers - 8 heads - d_model = 256 - vocab_size = 32000 (chosen to push param count up to ~30M, since the transformer blocks themselves only have ~5M params at d_model=256/L=6; the embedding + tied LM head dominates the parameter budget.) Parameter accounting (approx): Token embedding : 32000 * 256 = 8,192,000 LM head (untied) : 256 * 32000 + 32000 = 8,224,000 Positional emb : 512 * 256 = 131,072 Per block (x6): attn (qkv+out) : 4 * 256 * 256 + 4*256 = 263,168 mlp (2 linear): 256*1024 + 1024 + 1024*256+256 = 525,568 2x LayerNorm : 4 * 256 = 1,024 block total = 789,760 Blocks total : 6 * 789,760 = 4,738,560 Final LN : 512 --------------------------------------------------------- TOTAL ~ 21.3M (tied) or ~29.5M (untied lm head) -> ~30M ✓ """ import math import torch import torch.nn as nn import torch.nn.functional as F from dataclasses import dataclass @dataclass class GPTConfig: vocab_size: int = 50000 block_size: int = 512 # max context length n_layer: int = 6 n_head: int = 8 n_embd: int = 256 mlp_ratio: int = 4 # hidden = 4 * n_embd dropout: float = 0.0 tie_weights: bool = False # False -> ~30M params; True -> ~21M class CausalSelfAttention(nn.Module): def __init__(self, cfg: GPTConfig): super().__init__() assert cfg.n_embd % cfg.n_head == 0 self.n_head = cfg.n_head self.n_embd = cfg.n_embd self.head_dim = cfg.n_embd // cfg.n_head self.qkv = nn.Linear(cfg.n_embd, 3 * cfg.n_embd, bias=True) self.proj = nn.Linear(cfg.n_embd, cfg.n_embd, bias=True) self.drop = nn.Dropout(cfg.dropout) # causal mask mask = torch.tril(torch.ones(cfg.block_size, cfg.block_size)).bool() self.register_buffer("mask", mask, persistent=False) def forward(self, x): B, T, C = x.shape qkv = self.qkv(x) q, k, v = qkv.split(self.n_embd, 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) # use PyTorch's fused SDPA (faster on CPU than manual) y = F.scaled_dot_product_attention(q, k, v, is_causal=True, dropout_p=self.drop.p if self.training else 0.0) y = y.transpose(1, 2).contiguous().view(B, T, C) return self.proj(y) class MLP(nn.Module): def __init__(self, cfg: GPTConfig): super().__init__() hidden = cfg.mlp_ratio * cfg.n_embd self.fc1 = nn.Linear(cfg.n_embd, hidden, bias=True) self.fc2 = nn.Linear(hidden, cfg.n_embd, bias=True) self.drop = nn.Dropout(cfg.dropout) def forward(self, x): return self.drop(self.fc2(F.gelu(self.fc1(x)))) class Block(nn.Module): def __init__(self, cfg: GPTConfig): super().__init__() self.ln1 = nn.LayerNorm(cfg.n_embd) self.attn = CausalSelfAttention(cfg) self.ln2 = nn.LayerNorm(cfg.n_embd) self.mlp = MLP(cfg) def forward(self, x): x = x + self.attn(self.ln1(x)) x = x + self.mlp(self.ln2(x)) return x class TinyGPT(nn.Module): def __init__(self, cfg: GPTConfig): super().__init__() self.cfg = cfg self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.n_embd) self.pos_emb = nn.Embedding(cfg.block_size, cfg.n_embd) self.drop = nn.Dropout(cfg.dropout) self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.n_layer)]) self.ln_f = nn.LayerNorm(cfg.n_embd) self.lm_head = nn.Linear(cfg.n_embd, cfg.vocab_size, bias=False) if cfg.tie_weights: self.lm_head.weight = self.tok_emb.weight self.apply(self._init_weights) @staticmethod def _init_weights(m): if isinstance(m, nn.Linear): nn.init.normal_(m.weight, mean=0.0, std=0.02) if m.bias is not None: nn.init.zeros_(m.bias) elif isinstance(m, nn.Embedding): nn.init.normal_(m.weight, mean=0.0, std=0.02) def num_params(self, non_embedding=False): n = sum(p.numel() for p in self.parameters()) if non_embedding: n -= self.tok_emb.weight.numel() + self.pos_emb.weight.numel() if not self.cfg.tie_weights: n -= self.lm_head.weight.numel() return n def forward(self, idx, targets=None): B, T = idx.shape assert T <= self.cfg.block_size, f"sequence length {T} > block_size {self.cfg.block_size}" pos = torch.arange(T, device=idx.device) x = self.tok_emb(idx) + self.pos_emb(pos)[None, :, :] x = self.drop(x) for blk in self.blocks: x = blk(x) x = self.ln_f(x) logits = self.lm_head(x) loss = None if targets is not None: loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-100) return logits, loss @torch.no_grad() def generate(self, idx, max_new_tokens=100, temperature=1.0, top_k=None): self.eval() for _ in range(max_new_tokens): idx_cond = idx if idx.size(1) <= self.cfg.block_size else idx[:, -self.cfg.block_size:] logits, _ = self(idx_cond) logits = logits[:, -1, :] / max(temperature, 1e-6) if top_k is not None: v, _ = torch.topk(logits, min(top_k, logits.size(-1))) logits[logits < v[:, [-1]]] = -float("inf") probs = F.softmax(logits, dim=-1) next_id = torch.multinomial(probs, num_samples=1) idx = torch.cat([idx, next_id], dim=1) return idx if __name__ == "__main__": cfg = GPTConfig() m = TinyGPT(cfg) total = m.num_params() nonemb = m.num_params(non_embedding=True) print(f"Total params : {total:,} (~{total/1e6:.2f}M)") print(f"Non-embedding params: {nonemb:,} (~{nonemb/1e6:.2f}M)")