| """ |
| MicroGPT: Minimal PyTorch decoder-only GPT. |
| Same architecture as advance-torch-gpt.py but with tiny config for fast iteration. |
| """ |
|
|
| import math |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
|
|
| class RMSNorm(nn.Module): |
| def __init__(self, dim, eps=1e-5): |
| super().__init__() |
| self.weight = nn.Parameter(torch.ones(dim)) |
| self.eps = eps |
|
|
| def forward(self, x): |
| return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps) * self.weight |
|
|
|
|
| class CausalSelfAttention(nn.Module): |
| def __init__(self, n_embd, n_head, block_size, dropout): |
| super().__init__() |
| assert n_embd % n_head == 0 |
| self.n_head = n_head |
| self.head_dim = n_embd // n_head |
| self.q = nn.Linear(n_embd, n_embd, bias=False) |
| self.k = nn.Linear(n_embd, n_embd, bias=False) |
| self.v = nn.Linear(n_embd, n_embd, bias=False) |
| self.proj = nn.Linear(n_embd, n_embd, bias=False) |
| self.attn_drop = nn.Dropout(dropout) |
| self.resid_drop = nn.Dropout(dropout) |
| self.register_buffer('mask', torch.tril(torch.ones(block_size, block_size)).view(1, 1, block_size, block_size)) |
|
|
| def forward(self, x): |
| B, T, C = x.shape |
| q = self.q(x).view(B, T, self.n_head, self.head_dim).transpose(1, 2) |
| k = self.k(x).view(B, T, self.n_head, self.head_dim).transpose(1, 2) |
| v = self.v(x).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_drop(att) |
| y = att @ v |
| y = y.transpose(1, 2).contiguous().view(B, T, C) |
| return self.resid_drop(self.proj(y)) |
|
|
|
|
| class MLP(nn.Module): |
| def __init__(self, n_embd, dropout): |
| super().__init__() |
| self.fc = nn.Linear(n_embd, 4 * n_embd, bias=False) |
| self.proj = nn.Linear(4 * n_embd, n_embd, bias=False) |
| self.drop = nn.Dropout(dropout) |
|
|
| def forward(self, x): |
| x = F.gelu(self.fc(x)) |
| x = self.proj(x) |
| return self.drop(x) |
|
|
|
|
| class Block(nn.Module): |
| def __init__(self, n_embd, n_head, block_size, dropout): |
| super().__init__() |
| self.ln1 = RMSNorm(n_embd) |
| self.attn = CausalSelfAttention(n_embd, n_head, block_size, dropout) |
| self.ln2 = RMSNorm(n_embd) |
| self.mlp = MLP(n_embd, dropout) |
|
|
| def forward(self, x): |
| x = x + self.attn(self.ln1(x)) |
| x = x + self.mlp(self.ln2(x)) |
| return x |
|
|
|
|
| class MicroGPT(nn.Module): |
| """Minimal decoder-only GPT. Drop-in compatible with advance-torch-gpt.py's GPT class.""" |
|
|
| def __init__(self, vocab_size, block_size, n_layer=2, n_head=4, n_embd=128, dropout=0.1): |
| super().__init__() |
| self.block_size = block_size |
| self.wte = nn.Embedding(vocab_size, n_embd) |
| self.wpe = nn.Embedding(block_size, n_embd) |
| self.blocks = nn.ModuleList([Block(n_embd, n_head, block_size, dropout) for _ in range(n_layer)]) |
| self.ln_f = RMSNorm(n_embd) |
| self.lm_head = nn.Linear(n_embd, vocab_size, bias=False) |
| self.lm_head.weight = self.wte.weight |
| self.apply(self._init_weights) |
|
|
| def _init_weights(self, module): |
| if isinstance(module, nn.Linear): |
| nn.init.normal_(module.weight, mean=0.0, std=0.02) |
| 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 |
| if T > self.block_size: |
| raise ValueError(f'block size exceeded: {T} > {self.block_size}') |
| pos = torch.arange(T, device=idx.device) |
| x = self.wte(idx) + self.wpe(pos) |
| for block in self.blocks: |
| x = block(x) |
| x = self.ln_f(x) |
| logits = self.lm_head(x) |
| loss = None |
| if targets is not None: |
| loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.reshape(-1)) |
| return logits, loss |
|
|
| @torch.no_grad() |
| def generate(self, idx, max_new_tokens, temperature=1.0, top_k=40): |
| self.eval() |
| for _ in range(max_new_tokens): |
| idx_cond = idx[:, -self.block_size:] |
| logits, _ = self(idx_cond) |
| logits = logits[:, -1, :] / temperature |
| 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) |
| idx_next = torch.multinomial(probs, num_samples=1) |
| idx = torch.cat([idx, idx_next], dim=1) |
| return idx |
|
|
| @torch.no_grad() |
| def generate_text(self, idx, max_new_tokens, temperature=1.0, top_k=40): |
| """Alias for generate, returns tokens only.""" |
| return self.generate(idx, max_new_tokens, temperature, top_k) |
