model/transformer.py

"""
MiniTransformer — Decoder-only Causal Language Model
──────────────────────────────────────────────────────
Architecture at a glance:

  Input tokens (B, T)
       │
  Token Embedding  [vocab_size → d_model]
       │
  ┌────┴──────────────────────────────────────────────────────┐
  │  RoPE (cos, sin) — computed once, shared across all layers │
  └────┬──────────────────────────────────────────────────────┘
       │
  ┌────┴──────────────────────────┐
  │  TransformerBlock × n_layers   │  (Pre-LN, RMSNorm, GQA + SDPA, SwiGLU)
  └────┬──────────────────────────┘
       │
  Final RMSNorm
       │
  LM Head [d_model → vocab_size]  (weight-tied to token embedding)
       │
  Logits (B, T, vocab_size)

WEIGHT TYING:
  The token embedding matrix (vocab_size × d_model) and the LM head weight
  are shared — the same parameter serves both roles. This:
    1. Saves vocab_size × d_model parameters (~4M for 8192 vocab, 512 dim)
    2. Enforces consistency: the model learns the same representation for
       "what token am I?" and "what token should come next?"
  Standard in GPT-2, Llama, and most open-weight models.

WEIGHT INITIALISATION:
  All Linear and Embedding weights are initialised with N(0, 0.02) following
  GPT-2. This keeps activations in a stable range at the start of training.
  Biases (where they exist) are zeroed.
"""

import torch
import torch.nn as nn

from model.config import ModelConfig
from model.block  import TransformerBlock, RMSNorm
from model.positional_encoding import RotaryEmbedding


class MiniTransformer(nn.Module):
    """
    Decoder-only causal language model.

    Args:
        config: A ModelConfig dataclass instance. All hyperparameters flow
                through this object — never pass raw ints directly.

    Usage:
        config = ModelConfig(vocab_size=8192, d_model=512, n_layers=6)
        model  = MiniTransformer(config)
        logits, kvs = model(token_ids)   # logits: [B, T, vocab_size]
    """

    def __init__(self, config: ModelConfig):
        super().__init__()
        self.config = config

        # ── Embeddings ────────────────────────────────────────────────────────
        self.token_emb = nn.Embedding(config.vocab_size, config.d_model)

        # RoPE: produces (cos, sin) for each position; applied inside attention
        self.rope = RotaryEmbedding(
            dim         = config.head_dim,       # operates per attention-head
            max_seq_len = config.context_len,
            base        = config.rope_base,
        )

        # ── Transformer Layers ───────────────────────────────────────────────
        self.blocks = nn.ModuleList([
            TransformerBlock(config) for _ in range(config.n_layers)
        ])

        # ── Output ────────────────────────────────────────────────────────────
        self.ln_final = RMSNorm(config.d_model, eps=config.norm_eps)
        self.lm_head  = nn.Linear(config.d_model, config.vocab_size, bias=False)

        # Weight tying: lm_head reads the same matrix as token_emb
        if config.tie_weights:
            self.lm_head.weight = self.token_emb.weight

        # Initialise all parameters
        self.apply(self._init_weights)

    # ── Initialisation ────────────────────────────────────────────────────────

    def _init_weights(self, module: nn.Module) -> None:
        """GPT-2 style weight initialisation: N(0, 0.02) for all linear/embedding."""
        if isinstance(module, (nn.Linear, nn.Embedding)):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
        if isinstance(module, nn.Linear) and module.bias is not None:
            nn.init.zeros_(module.bias)

    # ── Forward Pass ──────────────────────────────────────────────────────────

    def forward(
        self,
        x:        torch.Tensor,
        past_kvs: list | None = None,
    ) -> tuple[torch.Tensor, list]:
        """
        Args:
            x:        Token indices, shape [B, T].
            past_kvs: List of (k, v) tuples — one per layer — for KV-cache
                      inference. Pass None during training or the first
                      generation step.

        Returns:
            logits:      Shape [B, T, vocab_size].
            present_kvs: Updated KV-cache list (one tuple per layer).
                         Feed this back as `past_kvs` for the next token.
        """
        B, T = x.shape

        # Token embeddings — no positional embedding added here (RoPE handles it)
        x = self.token_emb(x)                   # [B, T, d_model]

        # Build RoPE tables with correct sequence offset for KV-cache decoding
        start_pos = 0 if past_kvs is None else past_kvs[0][0].shape[2]
        cos, sin = self.rope(T, start_pos=start_pos)     # both [1, 1, T, head_dim]

        # Pass through each transformer block, collecting KV-cache outputs
        present_kvs: list = []
        for i, block in enumerate(self.blocks):
            past_kv = past_kvs[i] if past_kvs is not None else None
            x, present_kv = block(x, cos, sin, past_kv)
            present_kvs.append(present_kv)

        # Final normalisation and language model head
        x      = self.ln_final(x)               # [B, T, d_model]
        logits = self.lm_head(x)                # [B, T, vocab_size]
        return logits, present_kvs

    # ── Helpers ───────────────────────────────────────────────────────────────

    @property
    def num_parameters(self) -> int:
        """Count of trainable parameters."""
        return sum(p.numel() for p in self.parameters() if p.requires_grad)

    def __repr__(self) -> str:
        return (
            f"MiniTransformer | {self.num_parameters:,} params\n"
            f"  {self.config}"
        )