स्टेप 5 · एक पूरा छोटा GPT

import torch, torch.nn as nn, torch.nn.functional as F
from dataclasses import dataclass

@dataclass
class Config:
    vocab_size: int = 8192
    n_layer: int = 6
    n_head: int = 6
    d_model: int = 384
    block_size: int = 256      # max context length
    dropout: float = 0.1

class TinyGPT(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
        self.tok = nn.Embedding(cfg.vocab_size, cfg.d_model)
        self.pos = nn.Embedding(cfg.block_size, cfg.d_model)
        self.drop = nn.Dropout(cfg.dropout)
        self.blocks = nn.ModuleList(
            Block(cfg.d_model, cfg.n_head) for _ in range(cfg.n_layer))
        self.norm = RMSNorm(cfg.d_model)
        self.head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
        self.head.weight = self.tok.weight        # weight tying
        self.apply(self._init)

    def _init(self, m):
        if isinstance(m, nn.Linear):
            nn.init.normal_(m.weight, std=0.02)
        elif isinstance(m, nn.Embedding):
            nn.init.normal_(m.weight, std=0.02)

    def forward(self, idx, targets=None):
        B, T = idx.shape
        pos = torch.arange(T, device=idx.device)
        x = self.drop(self.tok(idx) + self.pos(pos))
        for b in self.blocks:
            x = b(x)
        logits = self.head(self.norm(x))          # (B, T, vocab)

        loss = None
        if targets is not None:
            loss = F.cross_entropy(
                logits.view(-1, logits.size(-1)),
                targets.view(-1))
        return logits, loss

    @torch.no_grad()
    def generate(self, idx, max_new=100, temperature=1.0, top_k=50):
        for _ in range(max_new):
            idx_cond = idx[:, -self.cfg.block_size:]      # crop to context
            logits, _ = self(idx_cond)
            logits = logits[:, -1, :] / temperature       # last position only
            if top_k:
                v, _ = torch.topk(logits, top_k)
                logits[logits < v[:, [-1]]] = -float("inf")
            probs = F.softmax(logits, dim=-1)
            nxt = torch.multinomial(probs, num_samples=1)
            idx = torch.cat((idx, nxt), dim=1)
        return idx

def count(cfg):
    emb  = cfg.vocab_size * cfg.d_model                 # tied with output head
    attn = 4 * cfg.d_model * cfg.d_model                # q, k, v, proj
    ff   = 3 * cfg.d_model * int(8 * cfg.d_model / 3)   # gate, up, down
    per_layer = attn + ff
    total = emb + cfg.n_layer * per_layer
    print(f"embedding  {emb/1e6:.1f}M")
    print(f"per layer  {per_layer/1e6:.1f}M  x {cfg.n_layer}")
    print(f"TOTAL      {total/1e6:.1f}M")
    return total

count(Config())     # ~ 14M parameters