2.7 Building a GPT From Scratch (nanoGPT-style)¶

from dataclasses import dataclass

@dataclass
class GPTConfig:
    # --- Sequence / vocabulary ---
    block_size: int = 256     # max context length T (positions the model can see)
    vocab_size: int = 65      # size of the token vocabulary V

    # --- Model shape ---
    n_layer: int = 6          # number of Transformer blocks stacked
    n_head:  int = 6          # attention heads per block (must divide n_embd)
    n_embd:  int = 384        # residual-stream width (a.k.a. d_model)

    # --- Regularization / numerics ---
    dropout: float = 0.0      # dropout prob (0.0 for small/clean data; 0.1+ to regularize)
    bias:    bool  = False    # use bias terms in Linear / LayerNorm? (False = modern default)

import torch
import torch.nn as nn
import torch.nn.functional as F

class LayerNorm(nn.Module):
    """LayerNorm with an optional bias. PyTorch's built-in does not allow bias=None."""
    def __init__(self, ndim, bias):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(ndim))               # gamma, scale
        self.bias   = nn.Parameter(torch.zeros(ndim)) if bias else None  # beta, shift

    def forward(self, x):
        # Normalize over the last dim only. eps keeps us safe when variance ~ 0.
        return F.layer_norm(x, self.weight.shape, self.weight, self.bias, eps=1e-5)

class CausalSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        assert config.n_embd % config.n_head == 0, "n_embd must be divisible by n_head"
        self.n_head = config.n_head
        self.n_embd = config.n_embd
        self.dropout = config.dropout

        # One projection produces Q, K, V together: output width = 3 * n_embd.
        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
        # Output projection W_O: mixes the per-head results back into the stream.
        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
        self.attn_dropout  = nn.Dropout(config.dropout)   # on the attention weights
        self.resid_dropout = nn.Dropout(config.dropout)   # on the block output

        # Use PyTorch's fused (FlashAttention-style) kernel when available.
        self.flash = hasattr(F, "scaled_dot_product_attention")
        if not self.flash:
            # Fallback: precompute a (1,1,T,T) causal mask buffer for manual attention.
            mask = torch.tril(torch.ones(config.block_size, config.block_size))
            self.register_buffer("bias", mask.view(1, 1, config.block_size, config.block_size))

    def forward(self, x):
        B, T, C = x.size()            # batch, sequence length, embedding dim (= n_embd)

        # Project once, then split into the three roles along the feature axis.
        q, k, v = self.c_attn(x).split(self.n_embd, dim=2)   # each (B, T, C)

        # Reshape (B, T, C) -> (B, n_head, T, head_dim) so heads run in parallel.
        hd = C // self.n_head
        q = q.view(B, T, self.n_head, hd).transpose(1, 2)    # (B, nh, T, hd)
        k = k.view(B, T, self.n_head, hd).transpose(1, 2)
        v = v.view(B, T, self.n_head, hd).transpose(1, 2)

        if self.flash:
            # Fused kernel: applies the causal mask, scaling, softmax, and dropout.
            y = F.scaled_dot_product_attention(
                q, k, v, attn_mask=None,
                dropout_p=self.dropout if self.training else 0.0,
                is_causal=True,
            )
        else:
            # Manual path — identical math, written out so you can read it.
            att = (q @ k.transpose(-2, -1)) * (1.0 / (k.size(-1) ** 0.5))   # (B,nh,T,T)
            att = att.masked_fill(self.bias[:, :, :T, :T] == 0, float("-inf"))
            att = F.softmax(att, dim=-1)
            att = self.attn_dropout(att)
            y = att @ v                                                     # (B,nh,T,hd)

        # Re-assemble heads: (B, nh, T, hd) -> (B, T, C). contiguous() before view.
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        return self.resid_dropout(self.c_proj(y))            # output projection + dropout

class MLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.c_fc   = nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias)  # expand
        self.gelu   = nn.GELU()                                                       # nonlinearity
        self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias)  # contract
        self.dropout = nn.Dropout(config.dropout)

    def forward(self, x):
        x = self.c_fc(x)      # (B, T, 4*n_embd)
        x = self.gelu(x)      # smooth, differentiable ReLU-like activation
        x = self.c_proj(x)    # back down to (B, T, n_embd)
        return self.dropout(x)

class Block(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.ln_1 = LayerNorm(config.n_embd, bias=config.bias)
        self.attn = CausalSelfAttention(config)
        self.ln_2 = LayerNorm(config.n_embd, bias=config.bias)
        self.mlp  = MLP(config)

    def forward(self, x):
        x = x + self.attn(self.ln_1(x))   # pre-norm attention, residual add
        x = x + self.mlp(self.ln_2(x))    # pre-norm MLP,       residual add
        return x

class GPT(nn.Module):
    def __init__(self, config):
        super().__init__()
        assert config.vocab_size is not None and config.block_size is not None
        self.config = config

        self.transformer = nn.ModuleDict(dict(
            wte  = nn.Embedding(config.vocab_size, config.n_embd),   # token embeddings
            wpe  = nn.Embedding(config.block_size, config.n_embd),   # learned positions
            drop = nn.Dropout(config.dropout),
            h    = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
            ln_f = LayerNorm(config.n_embd, bias=config.bias),       # final layernorm
        ))
        # Language-model head: residual stream -> vocabulary logits.
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)

        # WEIGHT TYING: share the token-embedding matrix with the output head.
        # wte maps id -> vector; lm_head maps vector -> id-logits. Same V×n_embd
        # matrix (transposed), so we tie them: saves V*n_embd params and helps.
        self.transformer.wte.weight = self.lm_head.weight

        # Initialize all weights (see _init_weights), then apply a special
        # scaled init to the residual output projections (GPT-2 trick).
        self.apply(self._init_weights)
        for pn, p in self.named_parameters():
            if pn.endswith("c_proj.weight"):
                # Scale residual-path projections by 1/sqrt(2 * n_layer). Explained below.
                torch.nn.init.normal_(p, mean=0.0, std=0.02 / (2 * config.n_layer) ** 0.5)

    def _init_weights(self, module):
        # Linear and Embedding weights ~ N(0, 0.02^2); biases zeroed. The 0.02 std
        # is the GPT-2 convention — small enough to keep early activations tame.
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, idx, targets=None):
        device = idx.device
        B, T = idx.size()
        assert T <= self.config.block_size, f"sequence length {T} > block_size"
        pos = torch.arange(0, T, dtype=torch.long, device=device)   # (T,) position ids

        # 1) Embed tokens and positions, then sum into the residual stream.
        tok_emb = self.transformer.wte(idx)    # (B, T, n_embd)
        pos_emb = self.transformer.wpe(pos)    # (T, n_embd), broadcasts over batch
        x = self.transformer.drop(tok_emb + pos_emb)

        # 2) Run the stack of Transformer blocks.
        for block in self.transformer.h:
            x = block(x)
        x = self.transformer.ln_f(x)           # final layernorm

        # 3) Project to logits and, if training, compute the loss.
        if targets is not None:
            logits = self.lm_head(x)           # (B, T, vocab_size) — all positions
            # Cross-entropy over the vocabulary, flattening batch & time together.
            # ignore_index=-1 lets us mask out padding positions in the targets.
            loss = F.cross_entropy(
                logits.view(-1, logits.size(-1)),
                targets.view(-1),
                ignore_index=-1,
            )
        else:
            # Inference: we only need the LAST position's logits to predict next token.
            logits = self.lm_head(x[:, [-1], :])  # (B, 1, vocab_size) — saves compute
            loss = None
        return logits, loss

import numpy as np

# --- Build a character-level vocabulary from the raw text ---
with open("input.txt", "r", encoding="utf-8") as f:
    text = f.read()
chars = sorted(list(set(text)))            # e.g. 65 unique characters
vocab_size = len(chars)
stoi = {ch: i for i, ch in enumerate(chars)}   # string -> int
itos = {i: ch for i, ch in enumerate(chars)}   # int    -> string
encode = lambda s: [stoi[c] for c in s]        # text -> list[int]
decode = lambda l: "".join(itos[i] for i in l) # list[int] -> text

# Encode the entire corpus once, then split 90/10 into train/val.
data = torch.tensor(encode(text), dtype=torch.long)
n = int(0.9 * len(data))
train_data, val_data = data[:n], data[n:]

def get_batch(split, block_size, batch_size, device):
    """Sample `batch_size` random contiguous windows of length block_size.
    x is tokens [i : i+T]; y is the SAME window shifted by one: [i+1 : i+1+T].
    So y[t] is the next-token target for x[t] — the whole autoregressive labelling."""
    d = train_data if split == "train" else val_data
    ix = torch.randint(len(d) - block_size, (batch_size,))   # random start offsets
    x = torch.stack([d[i     : i + block_size]     for i in ix])  # (B, T)
    y = torch.stack([d[i + 1 : i + 1 + block_size] for i in ix])  # (B, T), shifted
    return x.to(device), y.to(device)

# --- One time: tokenize the corpus into a flat uint16 file ---
# uint16 is enough for any vocab < 65536 (GPT-2's 50257 fits); halves file size vs int32.
import tiktoken                                  # pip install tiktoken
enc = tiktoken.get_encoding("gpt2")
ids = np.array(enc.encode_ordinary(text), dtype=np.uint16)
ids.tofile("train.bin")                          # do the same for the val split

# --- At train time: identical contract to get_batch, backed by the file ---
def get_batch_memmap(split, block_size, batch_size, device):
    # Reopen the memmap every call. Keeping one np.memmap object alive across
    # many reads leaks memory (the pages are never released); reopening is cheap.
    d = np.memmap(f"{split}.bin", dtype=np.uint16, mode="r")
    ix = torch.randint(len(d) - block_size, (batch_size,))
    x = torch.stack([torch.from_numpy(d[i     : i + block_size    ].astype(np.int64)) for i in ix])
    y = torch.stack([torch.from_numpy(d[i + 1 : i + 1 + block_size].astype(np.int64)) for i in ix])
    return x.to(device, non_blocking=True), y.to(device, non_blocking=True)

def configure_optimizer(model, weight_decay, lr, betas, device_type):
    # Collect trainable params, then split by tensor dimensionality.
    params = {pn: p for pn, p in model.named_parameters() if p.requires_grad}
    decay_params   = [p for p in params.values() if p.dim() >= 2]   # matmuls, embeddings
    nodecay_params = [p for p in params.values() if p.dim() <  2]   # biases, LN gains
    optim_groups = [
        {"params": decay_params,   "weight_decay": weight_decay},
        {"params": nodecay_params, "weight_decay": 0.0},
    ]
    # 'fused' AdamW is a faster CUDA path when available.
    use_fused = (device_type == "cuda")
    extra = dict(fused=True) if use_fused else dict()
    return torch.optim.AdamW(optim_groups, lr=lr, betas=betas, **extra)

import math

# --- Hyperparameters (a small but real configuration) ---
device = "cuda" if torch.cuda.is_available() else "cpu"
config = GPTConfig(block_size=256, vocab_size=vocab_size,
                   n_layer=6, n_head=6, n_embd=384, dropout=0.2, bias=False)
batch_size  = 64
max_iters   = 5000
eval_iter   = 250
grad_clip   = 1.0
learning_rate = 3e-4
warmup_iters  = 200       # linear LR warmup
lr_decay_iters = max_iters
min_lr        = 3e-5      # cosine floor (~lr/10)

model = GPT(config).to(device)
optimizer = configure_optimizer(model, weight_decay=0.1, lr=learning_rate,
                                betas=(0.9, 0.95), device_type=device)

def get_lr(it):
    """Linear warmup then cosine decay to min_lr — the standard LLM schedule."""
    if it < warmup_iters:
        return learning_rate * (it + 1) / warmup_iters
    if it > lr_decay_iters:
        return min_lr
    ratio = (it - warmup_iters) / (lr_decay_iters - warmup_iters)   # in [0, 1]
    coeff = 0.5 * (1.0 + math.cos(math.pi * ratio))                 # cosine 1 -> 0
    return min_lr + coeff * (learning_rate - min_lr)

@torch.no_grad()
def estimate_loss(eval_batches=50):
    """Average loss over several batches of train and val — a stable metric."""
    out = {}
    model.eval()
    for split in ("train", "val"):
        losses = torch.zeros(eval_batches)
        for k in range(eval_batches):
            X, Y = get_batch(split, config.block_size, batch_size, device)
            _, loss = model(X, Y)
            losses[k] = loss.item()
        out[split] = losses.mean().item()
    model.train()
    return out

# --- Train ---
model.train()
for it in range(max_iters):
    # Set this step's learning rate on every param group.
    lr = get_lr(it)
    for g in optimizer.param_groups:
        g["lr"] = lr

    # Periodically report train/val loss (catches overfitting early).
    if it % eval_iter == 0 or it == max_iters - 1:
        losses = estimate_loss()
        print(f"step {it:5d} | train {losses['train']:.4f} | val {losses['val']:.4f} | lr {lr:.2e}")

    # One optimization step.
    X, Y = get_batch("train", config.block_size, batch_size, device)
    logits, loss = model(X, Y)            # forward: compute loss
    optimizer.zero_grad(set_to_none=True) # clear stale grads (None is faster than 0)
    loss.backward()                       # backward: autograd fills .grad
    torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)  # tame loss spikes
    optimizer.step()                      # AdamW update

# --- Setup, once, right after building the model ---
device_type = "cuda" if device.startswith("cuda") else "cpu"
use_bf16 = (device_type == "cuda" and torch.cuda.is_bf16_supported())
ctx = torch.autocast(device_type=device_type, dtype=torch.bfloat16, enabled=use_bf16)

torch.set_float32_matmul_precision("high")  # allow TF32 for whatever fp32 matmuls remain
model = torch.compile(model)                # PyTorch 2.x graph capture + kernel fusion

grad_accum_steps = 8                        # tokens/step = batch_size * block_size * this

# --- Inner step, replacing the single forward/backward above ---
optimizer.zero_grad(set_to_none=True)
for _micro in range(grad_accum_steps):
    X, Y = get_batch("train", config.block_size, batch_size, device)
    with ctx:                                  # bf16 forward; fp32 master weights
        logits, loss = model(X, Y)
        loss = loss / grad_accum_steps         # MEAN over micro-batches, not sum
    loss.backward()                            # gradients accumulate into .grad
torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)   # clip the FULL gradient
optimizer.step()

def save_checkpoint(path, model, optimizer, it, config):
    """Save everything needed to resume training bit-continuously."""
    ckpt = {
        "model":            model.state_dict(),
        "optimizer":        optimizer.state_dict(),
        "iter_num":         it,
        "config":           config,
        "torch_rng_state":  torch.get_rng_state(),
        "cuda_rng_state":   torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None,
    }
    torch.save(ckpt, path)

def load_checkpoint(path, model, optimizer, device):
    """Restore model, optimizer, and RNG state; return the iter to resume AT."""
    # PyTorch >= 2.6 defaults to weights_only=True, which refuses to unpickle the
    # GPTConfig dataclass we stored alongside the tensors. Pass weights_only=False
    # for checkpoints YOU produced; never for a file downloaded from the internet
    # (unpickling executes arbitrary code). The alternative is to save config as a
    # plain dict, or to allow-list the class via torch.serialization.add_safe_globals.
    ckpt = torch.load(path, map_location=device, weights_only=False)
    model.load_state_dict(ckpt["model"])
    optimizer.load_state_dict(ckpt["optimizer"])
    # RNG state must be a CPU ByteTensor regardless of map_location, hence .cpu().
    torch.set_rng_state(ckpt["torch_rng_state"].cpu())
    if ckpt["cuda_rng_state"] is not None and torch.cuda.is_available():
        torch.cuda.set_rng_state_all(ckpt["cuda_rng_state"])
    return ckpt["iter_num"] + 1   # resume at the NEXT step, not the saved one

# --- Wire it into the loop ---
resume   = False               # flip to True to resume from ckpt_path
ckpt_path = "ckpt.pt"

start_it = load_checkpoint(ckpt_path, model, optimizer, device) if resume else 0

model.train()
for it in range(start_it, max_iters):
    lr = get_lr(it)
    for g in optimizer.param_groups:
        g["lr"] = lr

    if it % eval_iter == 0 or it == max_iters - 1:
        losses = estimate_loss()
        print(f"step {it:5d} | train {losses['train']:.4f} | val {losses['val']:.4f} | lr {lr:.2e}")
        save_checkpoint(ckpt_path, model, optimizer, it, config)   # checkpoint on eval steps

    X, Y = get_batch("train", config.block_size, batch_size, device)
    logits, loss = model(X, Y)
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
    optimizer.step()

@torch.no_grad()
def generate(model, idx, max_new_tokens, temperature=1.0, top_k=None, top_p=None):
    """Autoregressively extend idx (B, T) by max_new_tokens.
       temperature: >1 flattens the distribution (more random), <1 sharpens it.
       top_k: keep only the k highest-prob tokens before sampling.
       top_p: nucleus sampling — keep the smallest set of tokens whose cumulative
              probability exceeds p."""
    model.eval()
    for _ in range(max_new_tokens):
        # 1) Crop context to block_size — the model cannot attend beyond it.
        idx_cond = idx if idx.size(1) <= model.config.block_size \
                   else idx[:, -model.config.block_size:]

        # 2) Forward pass; take logits at the LAST position only.
        logits, _ = model(idx_cond)            # (B, 1, vocab_size)
        logits = logits[:, -1, :] / temperature  # (B, vocab_size), temperature-scaled

        # 3) Optional top-k filtering: zero out everything below the k-th best logit.
        if top_k is not None:
            v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
            logits[logits < v[:, [-1]]] = -float("inf")

        # 4) Convert to probabilities.
        probs = F.softmax(logits, dim=-1)      # (B, vocab_size)

        # 5) Optional top-p (nucleus) filtering on the probabilities.
        if top_p is not None:
            sorted_probs, sorted_idx = torch.sort(probs, descending=True, dim=-1)
            cumprobs = torch.cumsum(sorted_probs, dim=-1)
            # Mask tokens once cumulative prob has exceeded p (keep the boundary token).
            remove = cumprobs - sorted_probs > top_p
            sorted_probs[remove] = 0.0
            sorted_probs /= sorted_probs.sum(dim=-1, keepdim=True)   # renormalize
            # Scatter the filtered probs back to vocabulary order.
            probs = torch.zeros_like(probs).scatter_(1, sorted_idx, sorted_probs)

        # 6) Sample one token from the (filtered, renormalized) distribution.
        next_id = torch.multinomial(probs, num_samples=1)   # (B, 1)

        # 7) Append and repeat. The new token becomes part of next step's context.
        idx = torch.cat((idx, next_id), dim=1)
    return idx


# --- Generate from a single newline as the prompt ---
context = torch.tensor([[stoi["\n"]]], dtype=torch.long, device=device)
out_ids = generate(model, context, max_new_tokens=500, temperature=0.8, top_k=40)
print(decode(out_ids[0].tolist()))

from typing import Optional

# ═══════════════════════════════════════════════════════════════════════════
# Reused verbatim from Chapter 2.5 (Positional Encoding)
# ═══════════════════════════════════════════════════════════════════════════

def build_rope_cache(seq_len: int, head_dim: int, base: float = 10000.0,
                     device=None, dtype=torch.float32):
    """Precompute cos/sin tables for RoPE (Llama / rotate_half convention).
    Returns cos, sin of shape (seq_len, head_dim)."""
    assert head_dim % 2 == 0, "RoPE needs an even head dimension."
    half = head_dim // 2
    inv_freq = 1.0 / (base ** (torch.arange(0, half, device=device).float() / half))
    pos = torch.arange(seq_len, device=device).float()
    freqs = torch.outer(pos, inv_freq)                 # (seq_len, d/2)
    emb = torch.cat([freqs, freqs], dim=-1)             # (seq_len, d)
    return emb.cos().to(dtype), emb.sin().to(dtype)

def rotate_half(x: torch.Tensor) -> torch.Tensor:
    """Map (x1, x2) -> (-x2, x1) where x1,x2 are the two halves of the last dim."""
    half = x.shape[-1] // 2
    x1, x2 = x[..., :half], x[..., half:]
    return torch.cat([-x2, x1], dim=-1)

def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    """Apply RoPE to x of shape (B, n_head, T, head_dim); cos, sin: (T, head_dim)."""
    cos = cos[None, None, :, :]   # (1, 1, T, head_dim) — broadcasts over batch & heads
    sin = sin[None, None, :, :]
    return x * cos + rotate_half(x) * sin

# ═══════════════════════════════════════════════════════════════════════════
# Reused verbatim from Chapter 2.6 (The Transformer Block)
# ═══════════════════════════════════════════════════════════════════════════

class RMSNorm(nn.Module):
    """Root Mean Square LayerNorm (Zhang & Sennrich, 2019). No bias, no mean subtraction."""
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        rms = x.pow(2).mean(dim=-1, keepdim=True).add(self.eps).sqrt()
        return (x / rms) * self.weight

class SwiGLUFFN(nn.Module):
    """FFN(x) = W_down( Swish(W_gate x) * W_up x ). hidden_dim defaults to
    8/3 * dim rounded up to a multiple of 64, matching Llama's convention."""
    def __init__(self, dim: int, hidden_dim: Optional[int] = None,
                 bias: bool = False, dropout: float = 0.0):
        super().__init__()
        if hidden_dim is None:
            hidden_dim = int(8 * dim / 3)
            hidden_dim = 64 * ((hidden_dim + 63) // 64)
        self.w_gate = nn.Linear(dim, hidden_dim, bias=bias)
        self.w_up   = nn.Linear(dim, hidden_dim, bias=bias)
        self.w_down = nn.Linear(hidden_dim, dim, bias=bias)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        gate = F.silu(self.w_gate(x))
        up   = self.w_up(x)
        return self.dropout(self.w_down(gate * up))

class ModernCausalSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        assert config.n_embd % config.n_head == 0
        self.n_head = config.n_head
        self.n_embd = config.n_embd
        self.head_dim = config.n_embd // config.n_head
        self.dropout = config.dropout

        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
        self.resid_dropout = nn.Dropout(config.dropout)

    def forward(self, x, cos, sin):
        # cos, sin: (T, head_dim), already sliced to this sequence length by the caller.
        B, T, C = x.size()
        q, k, v = self.c_attn(x).split(self.n_embd, dim=2)          # each (B, T, C)

        # Reshape into heads FIRST: (B, T, C) -> (B, n_head, T, head_dim).
        q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2)  # (B, nh, T, hd)
        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)

        # THEN rotate — RoPE must be applied to q and k only, AFTER the head
        # split. Rotating the (B,T,C) tensor before splitting into heads would
        # pair dimensions across head boundaries instead of within a single
        # head's head_dim block — a silent correctness bug, not a crash.
        # v is never rotated: it carries content, not position.
        q = apply_rope(q, cos, sin)
        k = apply_rope(k, cos, sin)
        # Under bf16/autocast, cos/sin must match q/k's dtype. Either build the
        # cache directly in the model's working dtype, or do the safer:
        #   q = apply_rope(q.float(), cos, sin).to(q.dtype)

        y = F.scaled_dot_product_attention(
            q, k, v, attn_mask=None,
            dropout_p=self.dropout if self.training else 0.0,
            is_causal=True,
        )
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        return self.resid_dropout(self.c_proj(y))


class ModernBlock(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.norm_1 = RMSNorm(config.n_embd)
        self.attn   = ModernCausalSelfAttention(config)
        self.norm_2 = RMSNorm(config.n_embd)
        self.mlp    = SwiGLUFFN(config.n_embd, bias=config.bias, dropout=config.dropout)

    def forward(self, x, cos, sin):
        x = x + self.attn(self.norm_1(x), cos, sin)
        x = x + self.mlp(self.norm_2(x))
        return x

class ModernGPT(nn.Module):
    def __init__(self, config):
        super().__init__()
        assert config.vocab_size is not None and config.block_size is not None
        head_dim = config.n_embd // config.n_head
        assert head_dim % 2 == 0, "RoPE needs an even head_dim"   # 384/6=64: OK
        self.config = config

        self.transformer = nn.ModuleDict(dict(
            wte  = nn.Embedding(config.vocab_size, config.n_embd),   # NO wpe: RoPE replaces it
            drop = nn.Dropout(config.dropout),
            h    = nn.ModuleList([ModernBlock(config) for _ in range(config.n_layer)]),
            ln_f = RMSNorm(config.n_embd),
        ))
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)

        # WEIGHT TYING — unchanged from the baseline GPT.
        self.transformer.wte.weight = self.lm_head.weight

        # Build the RoPE cache ONCE, for the full block_size, and register it
        # as a non-persistent buffer: it is cheap to recompute, so we don't
        # want it bloating every checkpoint.
        cos, sin = build_rope_cache(config.block_size, head_dim)
        self.register_buffer("rope_cos", cos, persistent=False)
        self.register_buffer("rope_sin", sin, persistent=False)

        self.apply(self._init_weights)
        for pn, p in self.named_parameters():
            # Still exactly 2 residual adds per block (attn, mlp), so the
            # 1/sqrt(2*n_layer) factor from the baseline is unchanged.
            if pn.endswith("c_proj.weight") or pn.endswith("w_down.weight"):
                torch.nn.init.normal_(p, mean=0.0, std=0.02 / (2 * config.n_layer) ** 0.5)

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, idx, targets=None):
        B, T = idx.size()
        assert T <= self.config.block_size, f"sequence length {T} > block_size"

        x = self.transformer.drop(self.transformer.wte(idx))   # (B, T, n_embd) — no pos_emb add
        cos = self.rope_cos[:T]     # slice the cache to the current sequence length
        sin = self.rope_sin[:T]

        for block in self.transformer.h:
            x = block(x, cos, sin)
        x = self.transformer.ln_f(x)

        if targets is not None:
            logits = self.lm_head(x)
            loss = F.cross_entropy(
                logits.view(-1, logits.size(-1)),
                targets.view(-1),
                ignore_index=-1,
            )
        else:
            logits = self.lm_head(x[:, [-1], :])
            loss = None
        return logits, loss

model = ModernGPT(config).to(device)   # same GPTConfig: block_size=256, n_layer=6,
                                        # n_head=6, n_embd=384, dropout=0.2, bias=False
optimizer = configure_optimizer(model, weight_decay=0.1, lr=learning_rate,
                                betas=(0.9, 0.95), device_type=device)
# ... identical training loop as before, just with this model/optimizer ...

print(sum(p.numel() for p in model.parameters()))   # sanity-check the param count

model = AutoModelForCausalLM.from_pretrained(name, attn_implementation="sdpa")

from transformers import AutoModelForCausalLM

m = AutoModelForCausalLM.from_pretrained("HuggingFaceTB/SmolLM2-135M")
print(m)
# LlamaForCausalLM(
#   (model): LlamaModel(
#     (embed_tokens): Embedding(...)
#     (layers): ModuleList(
#       (0-N): LlamaDecoderLayer(
#         (self_attn): LlamaAttention(...)
#         (mlp): LlamaMLP(...)
#         (input_layernorm): LlamaRMSNorm(...)
#         (post_attention_layernorm): LlamaRMSNorm(...)
#       )
#     )
#     (norm): LlamaRMSNorm(...)
#   )
#   (lm_head): Linear(...)
# )

python -c "import transformers, os; print(os.path.dirname(transformers.__file__))"