yqi0/petitgpt

1from __future__ import annotations2 3from dataclasses import dataclass4import math5 6import torch7import torch.nn as nn8import torch.nn.functional as F9 10 11@dataclass12class GPTConfig:13    vocab_size: int = 3200014    n_layers: int = 3015    d_model: int = 57616    n_heads: int = 917    n_kv_heads: int = 3  # GQA KV heads; == n_heads is plain MHA (pre-GQA checkpoints)18    d_ff: int = 1536  # SwiGLU: ~2.67x d_model (MobileLLM/SmolLM2-135M deep-thin shape)19    max_seq_len: int = 204820    dropout: float = 0.021    tie_embeddings: bool = True22 23    # RoPE (rotary positional embedding)24    rope_theta: float = 10000.025    rope_pct: float = 1.0  # fraction of head_dim to rotate (1.0 = full head_dim)26 27 28CANONICAL_DENSE_PARAMETER_COUNT = 124_635_45629_CANONICAL_PARAMETERIZATION = {30    "vocab_size": 32_000,31    "n_layers": 30,32    "d_model": 576,33    "n_heads": 9,34    "n_kv_heads": 3,35    "d_ff": 1_536,36    "tie_embeddings": True,37}38 39 40def expected_gpt_parameter_count(cfg: GPTConfig) -> int:41    """Derive the unique parameter count for the dense bias-free GPT."""42    integer_fields = {43        "vocab_size": cfg.vocab_size,44        "n_layers": cfg.n_layers,45        "d_model": cfg.d_model,46        "n_heads": cfg.n_heads,47        "n_kv_heads": cfg.n_kv_heads,48        "d_ff": cfg.d_ff,49    }50    for name, value in integer_fields.items():51        if isinstance(value, bool) or not isinstance(value, int) or value <= 0:52            raise ValueError(f"GPTConfig.{name} must be a positive integer")53    if cfg.d_model % cfg.n_heads:54        raise ValueError("GPTConfig.d_model must be divisible by n_heads")55    if cfg.n_heads % cfg.n_kv_heads:56        raise ValueError("GPTConfig.n_heads must be divisible by n_kv_heads")57 58    head_dim = cfg.d_model // cfg.n_heads59    kv_dim = cfg.n_kv_heads * head_dim60    token_matrices = 1 if cfg.tie_embeddings else 261    embeddings = token_matrices * cfg.vocab_size * cfg.d_model62    # q + output projections are d_model x d_model; k and v are d_model x kv_dim (GQA)63    attention = 2 * cfg.d_model * cfg.d_model + 2 * cfg.d_model * kv_dim64    swiglu = 3 * cfg.d_model * cfg.d_ff65    block_norms = 2 * cfg.d_model66    final_norm = cfg.d_model67    return int(embeddings + cfg.n_layers * (attention + swiglu + block_norms) + final_norm)68 69 70def audit_gpt_parameter_count(model: nn.Module, cfg: GPTConfig) -> dict[str, int | bool | str]:71    """Fail fast on implementation/config drift and return manifest metadata."""72    expected = expected_gpt_parameter_count(cfg)73    actual = int(sum(parameter.numel() for parameter in model.parameters()))74    trainable = int(75        sum(parameter.numel() for parameter in model.parameters() if parameter.requires_grad)76    )77    if actual != expected:78        raise RuntimeError(79            "GPT parameter count disagrees with the architecture-derived count: "80            f"actual={actual:,}, expected={expected:,}"81        )82 83    canonical = all(84        getattr(cfg, field) == expected_value85        for field, expected_value in _CANONICAL_PARAMETERIZATION.items()86    )87    if canonical and actual != CANONICAL_DENSE_PARAMETER_COUNT:88        raise RuntimeError(89            "canonical PetitGPT parameter count mismatch: "90            f"actual={actual:,}, expected={CANONICAL_DENSE_PARAMETER_COUNT:,}"91        )92 93    return {94        "status": "passed",95        "counting_method": "unique_parameter_objects_excluding_buffers",96        "actual_total": actual,97        "actual_trainable": trainable,98        "derived_expected_total": expected,99        "canonical_parameterization": canonical,100        "canonical_expected_total": CANONICAL_DENSE_PARAMETER_COUNT,101        "canonical_match": canonical and actual == CANONICAL_DENSE_PARAMETER_COUNT,102    }103 104 105def gpt_config_from_checkpoint_dict(cfg_dict: dict) -> GPTConfig:106    """Rebuild a GPTConfig from a checkpoint's serialized config dict.107 108    Pre-GQA checkpoints carry no n_kv_heads; absence means plain MHA109    (n_kv_heads == n_heads), whose fused-QKV weight layout is unchanged.110    """111    cfg_dict = dict(cfg_dict)112    cfg_dict.setdefault("n_kv_heads", cfg_dict["n_heads"])113    return GPTConfig(**cfg_dict)114 115 116class RMSNorm(nn.Module):117    def __init__(self, dim: int, eps: float = 1e-6):118        super().__init__()119        self.eps = eps120        self.weight = nn.Parameter(torch.ones(dim))121 122    def forward(self, x: torch.Tensor) -> torch.Tensor:123        # x: [B, T, C]124        rms = x.pow(2).mean(dim=-1, keepdim=True).add(self.eps).rsqrt()125        return x * rms * self.weight126 127 128def _rotate_half(x: torch.Tensor) -> torch.Tensor:129    # x: [..., D]. Half-split layout (Llama/GPT-NeoX): pairs are (i, i+D/2),130    # matching the `cat([freqs, freqs])` cos/sin cache below.131    half = x.shape[-1] // 2132    x1 = x[..., :half]133    x2 = x[..., half:]134    return torch.cat((-x2, x1), dim=-1)135 136 137class RotaryEmbedding(nn.Module):138    """Precomputes RoPE cos/sin caches up to max_seq_len."""139 140    def __init__(self, head_dim: int, max_seq_len: int, theta: float = 10000.0, pct: float = 1.0):141        super().__init__()142        if head_dim % 2 != 0:143            raise ValueError(f"RoPE requires even head_dim, got {head_dim}")144        self.head_dim = int(head_dim)145        self.max_seq_len = int(max_seq_len)146        self.theta = float(theta)147        self.pct = float(pct)148 149        rope_dim = int(self.head_dim * self.pct)150        rope_dim = rope_dim - (rope_dim % 2)151        rope_dim = max(0, min(rope_dim, self.head_dim))152        self.rope_dim = rope_dim153 154        if self.rope_dim > 0:155            inv_freq = 1.0 / (156                self.theta ** (torch.arange(0, self.rope_dim, 2).float() / self.rope_dim)157            )158            t = torch.arange(self.max_seq_len, dtype=torch.float32)159            freqs = torch.outer(t, inv_freq)  # [T, rope_dim/2]160            emb = torch.cat([freqs, freqs], dim=-1)  # [T, rope_dim]161            cos = emb.cos()162            sin = emb.sin()163        else:164            cos = torch.empty(self.max_seq_len, 0, dtype=torch.float32)165            sin = torch.empty(self.max_seq_len, 0, dtype=torch.float32)166 167        self.register_buffer("cos_cached", cos, persistent=False)168        self.register_buffer("sin_cached", sin, persistent=False)169 170    def forward(171        self, q: torch.Tensor, k: torch.Tensor, seq_len: int, offset: int = 0172    ) -> tuple[torch.Tensor, torch.Tensor]:173        """Apply RoPE to q,k. q,k: [B, nH, T, Hd].174 175        `offset` is the absolute position of the first token in q,k — nonzero176        during KV-cached incremental decoding, where the new tokens sit at177        positions [offset, offset+seq_len).178        """179        end = offset + seq_len180        if end > self.max_seq_len:181            raise ValueError(182                f"position {end} exceeds max_seq_len={self.max_seq_len} for RoPE cache"183            )184        if self.rope_dim == 0:185            return q, k186 187        cos = self.cos_cached[offset:end].to(dtype=q.dtype, device=q.device)  # [T, rope_dim]188        sin = self.sin_cached[offset:end].to(dtype=q.dtype, device=q.device)  # [T, rope_dim]189        cos = cos.unsqueeze(0).unsqueeze(0)  # [1,1,T,rope_dim]190        sin = sin.unsqueeze(0).unsqueeze(0)191 192        q1, q2 = q[..., : self.rope_dim], q[..., self.rope_dim :]193        k1, k2 = k[..., : self.rope_dim], k[..., self.rope_dim :]194 195        q1 = q1 * cos + _rotate_half(q1) * sin196        k1 = k1 * cos + _rotate_half(k1) * sin197 198        q = torch.cat([q1, q2], dim=-1)199        k = torch.cat([k1, k2], dim=-1)200        return q, k201 202 203class SwiGLU(nn.Module):204    def __init__(self, cfg: GPTConfig):205        super().__init__()206        self.w1 = nn.Linear(cfg.d_model, cfg.d_ff, bias=False)207        self.w3 = nn.Linear(cfg.d_model, cfg.d_ff, bias=False)208        self.w2 = nn.Linear(cfg.d_ff, cfg.d_model, bias=False)209        self.drop = nn.Dropout(cfg.dropout)210 211    def forward(self, x: torch.Tensor) -> torch.Tensor:212        x = F.silu(self.w1(x)) * self.w3(x)213        x = self.w2(x)214        return self.drop(x)215 216 217class CausalSelfAttention(nn.Module):218    def __init__(self, cfg: GPTConfig):219        super().__init__()220        assert cfg.d_model % cfg.n_heads == 0221        assert cfg.n_heads % cfg.n_kv_heads == 0222        self.cfg = cfg223        self.head_dim = cfg.d_model // cfg.n_heads224        self.kv_dim = cfg.n_kv_heads * self.head_dim225 226        # QKV fused: one matmul instead of three. K/V carry n_kv_heads (GQA);227        # n_kv_heads == n_heads is plain MHA with the historical 3*d_model layout.228        self.qkv = nn.Linear(cfg.d_model, cfg.d_model + 2 * self.kv_dim, bias=False)229        # residual branch output projection230        self.proj = nn.Linear(cfg.d_model, cfg.d_model, bias=False)231        self.drop = nn.Dropout(cfg.dropout)232 233        self.rope = RotaryEmbedding(234            head_dim=self.head_dim,235            max_seq_len=cfg.max_seq_len,236            theta=cfg.rope_theta,237            pct=cfg.rope_pct,238        )239 240    @staticmethod241    def _incremental_mask(T: int, past_len: int, device: torch.device) -> torch.Tensor:242        """Bottom-right causal mask [T, past_len+T] (True = attend) for decoding243        T new queries against past_len cached keys plus the new keys."""244        q_pos = past_len + torch.arange(T, device=device)245        k_pos = torch.arange(past_len + T, device=device)246        return k_pos[None, :] <= q_pos[:, None]247 248    def forward(249        self,250        x: torch.Tensor,251        past_kv: tuple[torch.Tensor, torch.Tensor] | None = None,252        use_cache: bool = False,253    ):254        """x: [B, T, C]. With no cache this is byte-identical to a plain causal255        forward and returns the output tensor. With `use_cache` (or a supplied256        `past_kv`) it also returns the updated (k, v) for this layer."""257        B, T, C = x.shape258        past_len = 0 if past_kv is None else past_kv[0].size(2)259        if past_len + T > self.cfg.max_seq_len:260            raise ValueError(261                f"cache length {past_len + T} exceeds max_seq_len={self.cfg.max_seq_len}"262            )263 264        qkv = self.qkv(x)  # [B, T, C + 2*kv_dim]265        q, k, v = qkv.split([C, self.kv_dim, self.kv_dim], dim=-1)266 267        q = q.view(B, T, self.cfg.n_heads, self.head_dim).transpose(1, 2)  # [B,nH,T,Hd]268        k = k.view(B, T, self.cfg.n_kv_heads, self.head_dim).transpose(1, 2)  # [B,nKV,T,Hd]269        v = v.view(B, T, self.cfg.n_kv_heads, self.head_dim).transpose(1, 2)270 271        # RoPE rotates only the new tokens, at their absolute positions.272        q, k = self.rope(q, k, seq_len=T, offset=past_len)273 274        # Prepend cached keys/values (already rotated when they were new). The275        # cache stays un-expanded at n_kv_heads so its memory reflects GQA.276        if past_kv is not None:277            k = torch.cat([past_kv[0], k], dim=2)278            v = torch.cat([past_kv[1], v], dim=2)279        present = (k, v) if use_cache else None280 281        # Expand grouped KV heads to the full head count for attention. KV head g282        # serves query heads [g*rep, (g+1)*rep) — repeat_interleave matches SDPA's283        # enable_gqa grouping (torch >= 2.5), which can replace this someday.284        if self.cfg.n_kv_heads != self.cfg.n_heads:285            rep = self.cfg.n_heads // self.cfg.n_kv_heads286            k = k.repeat_interleave(rep, dim=1)287            v = v.repeat_interleave(rep, dim=1)288 289        dropout_p = float(self.cfg.dropout) if (self.training and self.cfg.dropout > 0) else 0.0290 291        if q.device.type == "cuda":292            if past_len == 0:293                y = F.scaled_dot_product_attention(294                    q, k, v, attn_mask=None, dropout_p=dropout_p, is_causal=True295                )296            else:297                y = F.scaled_dot_product_attention(298                    q,299                    k,300                    v,301                    attn_mask=self._incremental_mask(T, past_len, q.device),302                    dropout_p=dropout_p,303                )304        else:305            scale = 1.0 / math.sqrt(self.head_dim)306            att = torch.matmul(q * scale, k.transpose(-2, -1))  # [B,nH,T,past_len+T]307            if past_len == 0:308                mask = torch.triu(torch.ones((T, T), device=q.device, dtype=torch.bool), diagonal=1)309                att = att.masked_fill(mask, float("-inf"))310            else:311                allow = self._incremental_mask(T, past_len, q.device)  # [T, past_len+T]312                att = att.masked_fill(~allow, float("-inf"))313            att = F.softmax(att, dim=-1)314            if dropout_p > 0.0:315                att = F.dropout(att, p=dropout_p)316            y = torch.matmul(att, v)317 318        y = y.transpose(1, 2).contiguous().view(B, T, C)319        y = self.drop(self.proj(y))320        if use_cache:321            return y, present322        return y323 324 325class Block(nn.Module):326    def __init__(self, cfg: GPTConfig):327        super().__init__()328        self.norm1 = RMSNorm(cfg.d_model)329        self.attn = CausalSelfAttention(cfg)330        self.norm2 = RMSNorm(cfg.d_model)331        self.mlp = SwiGLU(cfg)332 333    def forward(334        self,335        x: torch.Tensor,336        past_kv: tuple[torch.Tensor, torch.Tensor] | None = None,337        use_cache: bool = False,338    ):339        if use_cache or past_kv is not None:340            attn_out, present = self.attn(self.norm1(x), past_kv=past_kv, use_cache=True)341            x = x + attn_out342            x = x + self.mlp(self.norm2(x))343            return x, present344        x = x + self.attn(self.norm1(x))345        x = x + self.mlp(self.norm2(x))346        return x347 348 349class GPT(nn.Module):350    def __init__(self, cfg: GPTConfig):351        super().__init__()352        self.cfg = cfg353 354        self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.d_model)355        self.drop = nn.Dropout(cfg.dropout)356 357        self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.n_layers)])358        self.norm_f = RMSNorm(cfg.d_model)359 360        self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)361        if cfg.tie_embeddings:362            self.lm_head.weight = self.tok_emb.weight363 364        # Base init everywhere...365        self.apply(self._init_weights)366        # ...then scale init ONLY on residual-branch output projections: attn.proj and mlp.w2367        self._init_residual_projections()368 369    def _init_weights(self, m: nn.Module):370        if isinstance(m, (nn.Linear, nn.Embedding)):371            torch.nn.init.normal_(m.weight, mean=0.0, std=0.02)372 373    def _init_residual_projections(self):374        std = 0.02 / math.sqrt(2.0 * float(self.cfg.n_layers))375        for blk in self.blocks:376            torch.nn.init.normal_(blk.attn.proj.weight, mean=0.0, std=std)377            torch.nn.init.normal_(blk.mlp.w2.weight, mean=0.0, std=std)378 379    def forward(380        self,381        input_ids: torch.Tensor,382        past_kv: list[tuple[torch.Tensor, torch.Tensor]] | None = None,383        use_cache: bool = False,384    ):385        """Default call `model(input_ids)` returns logits [B, T, V] — unchanged.386 387        For incremental decoding, pass `use_cache=True` to also get a per-layer388        list of (k, v) tensors, and feed it back as `past_kv` with only the new389        token(s) on the next call. See `generate`.390        """391        B, T = input_ids.shape392        past_len = 0 if past_kv is None else past_kv[0][0].size(2)393        if past_len + T > self.cfg.max_seq_len:394            raise ValueError(395                f"T={T} with cache={past_len} exceeds max_seq_len={self.cfg.max_seq_len}"396            )397        if T < 1:398            raise ValueError("Empty sequence")399        caching = use_cache or (past_kv is not None)400 401        x = self.tok_emb(input_ids)402        x = self.drop(x)403        presents: list[tuple[torch.Tensor, torch.Tensor]] = []404        for i, blk in enumerate(self.blocks):405            layer_past = past_kv[i] if past_kv is not None else None406            if caching:407                x, present = blk(x, past_kv=layer_past, use_cache=True)408                presents.append(present)409            else:410                x = blk(x)411        x = self.norm_f(x)412        logits = self.lm_head(x)413        if caching:414            return logits, presents415        return logits416 417    @staticmethod418    def _sample_token(419        logits: torch.Tensor, temperature: float, top_k: int, top_p: float420    ) -> torch.Tensor:421        """logits: [B, V] -> next token [B, 1]. temperature<=0 is greedy."""422        if temperature <= 0:423            return logits.argmax(dim=-1, keepdim=True)424        logits = logits / temperature425        if top_k and top_k > 0:426            k = min(int(top_k), logits.size(-1))427            thresh = torch.topk(logits, k, dim=-1).values[:, -1, None]428            logits = logits.masked_fill(logits < thresh, float("-inf"))429        if top_p and top_p < 1.0:430            sorted_logits, sorted_idx = torch.sort(logits, descending=True, dim=-1)431            cum = torch.softmax(sorted_logits, dim=-1).cumsum(dim=-1)432            drop_sorted = cum > top_p433            drop_sorted[..., 0] = False434            drop = torch.zeros_like(drop_sorted).scatter(-1, sorted_idx, drop_sorted)435            logits = logits.masked_fill(drop, float("-inf"))436        probs = torch.softmax(logits, dim=-1)437        return torch.multinomial(probs, num_samples=1)438 439    @torch.no_grad()440    def generate(441        self,442        input_ids: torch.Tensor,443        max_new_tokens: int,444        *,445        temperature: float = 1.0,446        top_k: int = 0,447        top_p: float = 1.0,448        eos_id: int | None = None,449    ) -> torch.Tensor:450        """KV-cached incremental decoding. input_ids: [B, T] -> [B, T + n].451 452        Prefills the prompt once, then feeds one new token per step against the453        cache (O(T) forwards of length 1) instead of re-running the full growing454        sequence each step. Stops early if all rows emit `eos_id`.455        """456        was_training = self.training457        self.eval()458        logits, past = self.forward(input_ids, use_cache=True)459        out = input_ids460        for _ in range(int(max_new_tokens)):461            next_tok = self._sample_token(logits[:, -1, :], temperature, top_k, top_p)462            out = torch.cat([out, next_tok], dim=1)463            if eos_id is not None and bool((next_tok.squeeze(1) == eos_id).all()):464                break465            if out.size(1) >= self.cfg.max_seq_len:466                break467            logits, past = self.forward(next_tok, past_kv=past, use_cache=True)468        if was_training:469            self.train()470        return out471