Building an LLM from Scratch in PyTorch: The Full Lifecycle Cheatsheet¶

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


class MultiHeadAttention(nn.Module):
    def __init__(self, d_model: int, n_heads: int):
        super().__init__()
        assert d_model % n_heads == 0
        self.n_heads = n_heads
        self.head_dim = d_model // n_heads

        self.q = nn.Linear(d_model, d_model, bias=False)
        self.k = nn.Linear(d_model, d_model, bias=False)
        self.v = nn.Linear(d_model, d_model, bias=False)
        self.out = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, C = x.shape

        q = self.q(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        k = self.k(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        v = self.v(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)

        # Causal mask: each position can only attend to itself and earlier positions
        scale = math.sqrt(self.head_dim)
        scores = (q @ k.transpose(-2, -1)) / scale          # (B, H, T, T)
        causal_mask = torch.triu(
            torch.ones(T, T, device=x.device), diagonal=1
        ).bool()
        scores = scores.masked_fill(causal_mask, float("-inf"))

        attn = F.softmax(scores, dim=-1)
        out = (attn @ v).transpose(1, 2).contiguous().view(B, T, C)
        return self.out(out)


class FeedForward(nn.Module):
    def __init__(self, d_model: int, expansion: int = 4):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(d_model, d_model * expansion),
            nn.GELU(),
            nn.Linear(d_model * expansion, d_model),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)


class TransformerBlock(nn.Module):
    def __init__(self, d_model: int, n_heads: int):
        super().__init__()
        self.norm1 = nn.LayerNorm(d_model)
        self.attn  = MultiHeadAttention(d_model, n_heads)
        self.norm2 = nn.LayerNorm(d_model)
        self.ff    = FeedForward(d_model)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.attn(self.norm1(x))   # pre-norm + residual
        x = x + self.ff(self.norm2(x))
        return x


class TinyGPT(nn.Module):
    def __init__(self, vocab_size: int, d_model: int, n_heads: int,
                 n_layers: int, max_seq_len: int):
        super().__init__()
        self.token_emb = nn.Embedding(vocab_size, d_model)
        self.pos_emb   = nn.Embedding(max_seq_len, d_model)
        self.blocks    = nn.ModuleList([TransformerBlock(d_model, n_heads)
                                        for _ in range(n_layers)])
        self.norm      = nn.LayerNorm(d_model)
        self.head      = nn.Linear(d_model, vocab_size, bias=False)

    def forward(self, idx: torch.Tensor) -> torch.Tensor:
        B, T = idx.shape
        positions = torch.arange(T, device=idx.device)
        x = self.token_emb(idx) + self.pos_emb(positions)
        for block in self.blocks:
            x = block(x)
        return self.head(self.norm(x))          # logits: (B, T, vocab_size)

import torch
import torch.nn.functional as F

# --- 1. Tokeniser (character-level — no libraries needed) ---
text = open("tiny_shakespeare.txt").read()
chars = sorted(set(text))
vocab_size = len(chars)
stoi = {c: i for i, c in enumerate(chars)}
itos = {i: c for i, c in enumerate(chars)}
encode = lambda s: [stoi[c] for c in s]
decode = lambda ids: "".join(itos[i] for i in ids)

data = torch.tensor(encode(text), dtype=torch.long)

# --- 2. Batching ---
def get_batch(data: torch.Tensor, batch_size: int, block_size: int,
              device: str) -> tuple[torch.Tensor, torch.Tensor]:
    ix = torch.randint(len(data) - block_size, (batch_size,))
    x = torch.stack([data[i:i + block_size] for i in ix])
    y = torch.stack([data[i + 1:i + block_size + 1] for i in ix])
    return x.to(device), y.to(device)

# --- 3. Training loop ---
device = "cuda" if torch.cuda.is_available() else "cpu"

model = TinyGPT(
    vocab_size=vocab_size,
    d_model=128,
    n_heads=4,
    n_layers=4,
    max_seq_len=256,
).to(device)

optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)

for step in range(5000):
    x, y = get_batch(data, batch_size=32, block_size=256, device=device)
    logits = model(x)                                  # (B, T, vocab_size)
    loss = F.cross_entropy(
        logits.view(-1, vocab_size),                   # (B*T, vocab_size)
        y.view(-1),                                    # (B*T,)
    )
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if step % 500 == 0:
        print(f"step {step:>5} | loss {loss.item():.4f}")

# --- 4. Generation ---
@torch.no_grad()
def generate(model, prompt: str, max_new_tokens: int = 200) -> str:
    model.eval()
    idx = torch.tensor([encode(prompt)], dtype=torch.long, device=device)
    for _ in range(max_new_tokens):
        idx_cond = idx[:, -256:]                       # crop to max_seq_len
        logits = model(idx_cond)
        next_token_logits = logits[:, -1, :]           # last position only
        probs = F.softmax(next_token_logits / 0.8, dim=-1)  # temperature=0.8
        next_token = torch.multinomial(probs, num_samples=1)
        idx = torch.cat([idx, next_token], dim=1)
    return decode(idx[0].tolist())

print(generate(model, "HAMLET:"))

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


# --- RMSNorm: simpler, faster than LayerNorm ---
class RMSNorm(nn.Module):
    def __init__(self, d_model: int, eps: float = 1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(d_model))
        self.eps = eps

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


# --- SwiGLU: gated FFN used in Llama, PaLM, Mistral ---
class SwiGLU(nn.Module):
    def __init__(self, d_model: int, expansion: int = 4):
        super().__init__()
        hidden = int(d_model * expansion * 2 / 3)  # standard Llama sizing
        self.gate = nn.Linear(d_model, hidden, bias=False)
        self.up   = nn.Linear(d_model, hidden, bias=False)
        self.down = nn.Linear(hidden,  d_model, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.down(F.silu(self.gate(x)) * self.up(x))


# --- RoPE: Rotary Position Embeddings ---
def precompute_rope(head_dim: int, max_seq_len: int,
                    base: int = 10000) -> tuple[torch.Tensor, torch.Tensor]:
    theta = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim))
    t = torch.arange(max_seq_len)
    freqs = torch.outer(t, theta)               # (T, head_dim/2)
    cos = freqs.cos().repeat_interleave(2, dim=-1)
    sin = freqs.sin().repeat_interleave(2, dim=-1)
    return cos, sin                             # each: (T, head_dim)

def apply_rope(x: torch.Tensor, cos: torch.Tensor,
               sin: torch.Tensor) -> torch.Tensor:
    # Rotate pairs of dimensions: [x0, x1] → [x0*cos - x1*sin, x0*sin + x1*cos]
    x_rotated = torch.stack([-x[..., 1::2], x[..., ::2]], dim=-1)
    x_rotated = x_rotated.flatten(-2)
    return x * cos + x_rotated * sin


# --- KV Cache: store past keys/values, avoid recomputing ---
class KVCache:
    def __init__(self):
        self.k_cache: torch.Tensor | None = None
        self.v_cache: torch.Tensor | None = None

    def update(self, k: torch.Tensor, v: torch.Tensor
               ) -> tuple[torch.Tensor, torch.Tensor]:
        if self.k_cache is None:
            self.k_cache, self.v_cache = k, v
        else:
            self.k_cache = torch.cat([self.k_cache, k], dim=2)
            self.v_cache = torch.cat([self.v_cache, v], dim=2)
        return self.k_cache, self.v_cache

    def reset(self):
        self.k_cache = self.v_cache = None


# --- Modern Transformer Block combining all three upgrades ---
class ModernTransformerBlock(nn.Module):
    def __init__(self, d_model: int, n_heads: int, max_seq_len: int):
        super().__init__()
        self.n_heads  = n_heads
        self.head_dim = d_model // n_heads
        self.norm1 = RMSNorm(d_model)
        self.norm2 = RMSNorm(d_model)
        self.q  = nn.Linear(d_model, d_model, bias=False)
        self.k  = nn.Linear(d_model, d_model, bias=False)
        self.v  = nn.Linear(d_model, d_model, bias=False)
        self.o  = nn.Linear(d_model, d_model, bias=False)
        self.ff = SwiGLU(d_model)
        cos, sin = precompute_rope(self.head_dim, max_seq_len)
        self.register_buffer("cos", cos)
        self.register_buffer("sin", sin)

    def forward(self, x: torch.Tensor,
                kv_cache: KVCache | None = None,
                start_pos: int = 0) -> torch.Tensor:
        B, T, C = x.shape
        h = self.n_heads
        d = self.head_dim

        xn = self.norm1(x)                       # compute the norm once, reuse 3×
        q = self.q(xn).view(B, T, h, d).transpose(1, 2)
        k = self.k(xn).view(B, T, h, d).transpose(1, 2)
        v = self.v(xn).view(B, T, h, d).transpose(1, 2)

        # Apply RoPE using each token's ABSOLUTE position. During cached
        # decode you feed one token at a time — without start_pos, every new
        # token would be rotated as position 0 and the cache silently breaks.
        cos = self.cos[start_pos:start_pos + T]
        sin = self.sin[start_pos:start_pos + T]
        q = apply_rope(q, cos, sin)
        k = apply_rope(k, cos, sin)

        # Use KV cache at inference time
        if kv_cache is not None:
            k, v = kv_cache.update(k, v)

        # is_causal only when processing multiple tokens. During cached decode
        # (T == 1) the new token may attend to ALL cached positions — SDPA's
        # is_causal=True aligns the mask top-left, which would wrongly mask
        # everything except the first cached token.
        attn = F.scaled_dot_product_attention(q, k, v, is_causal=(T > 1))
        out = attn.transpose(1, 2).contiguous().view(B, T, C)
        x = x + self.o(out)
        x = x + self.ff(self.norm2(x))
        return x

block = ModernTransformerBlock(d_model=64, n_heads=4, max_seq_len=128).eval()
x = torch.randn(1, 12, 64)

with torch.no_grad():
    full = block(x)                                       # ground truth
    cache = KVCache()
    cached = torch.cat(
        [block(x[:, t:t+1], kv_cache=cache, start_pos=t) for t in range(12)],
        dim=1,
    )

print((full - cached).abs().max().item())   # verified: 2.4e-07 ✓

import torch
from tokenizers import Tokenizer, models, trainers, pre_tokenizers

# --- 1. Train a BPE tokeniser ---
def train_bpe_tokeniser(files: list[str], vocab_size: int = 8192) -> Tokenizer:
    tokenizer = Tokenizer(models.BPE())
    tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
    trainer = trainers.BpeTrainer(
        vocab_size=vocab_size,
        special_tokens=["<|pad|>", "<|bos|>", "<|eos|>"],
    )
    tokenizer.train(files, trainer)
    return tokenizer


# --- 2. Mixed Precision (BF16) + Gradient Accumulation training loop ---
# BF16 has FP32's dynamic range, so it needs NO loss scaler — that's the
# whole reason to prefer it. (The FP16 + GradScaler variant is below.)
def train(model, dataloader, optimizer, n_epochs: int,
          accum_steps: int = 8, device: str = "cuda"):

    model.train()
    optimizer.zero_grad()

    for epoch in range(n_epochs):
        for step, (x, y) in enumerate(dataloader):
            x, y = x.to(device), y.to(device)

            # autocast: runs the forward pass in BF16, master weights stay FP32
            with torch.autocast(device_type=device, dtype=torch.bfloat16):
                logits = model(x)
                loss = torch.nn.functional.cross_entropy(
                    logits.view(-1, logits.size(-1)), y.view(-1)
                )
                loss = loss / accum_steps           # normalise before accumulating

            loss.backward()

            # Only update weights every accum_steps steps
            if (step + 1) % accum_steps == 0:
                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
                optimizer.step()
                optimizer.zero_grad()

                print(f"epoch {epoch} step {step} | loss {loss.item() * accum_steps:.4f}")

        # Checkpoint after every epoch
        torch.save({
            "epoch": epoch,
            "model_state": model.state_dict(),
            "optimizer_state": optimizer.state_dict(),
        }, f"checkpoint_epoch_{epoch}.pt")

scaler = torch.amp.GradScaler(device)        # NOT torch.cuda.amp.GradScaler — deprecated

with torch.autocast(device_type=device, dtype=torch.float16):
    loss = compute_loss(model, x, y) / accum_steps
scaler.scale(loss).backward()

if (step + 1) % accum_steps == 0:
    scaler.unscale_(optimizer)               # unscale BEFORE clipping
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    scaler.step(optimizer)
    scaler.update()
    optimizer.zero_grad()

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


class MoELayer(nn.Module):
    def __init__(self, d_model: int, n_experts: int = 8, top_k: int = 2):
        super().__init__()
        assert top_k <= n_experts
        self.n_experts = n_experts
        self.top_k = top_k

        # One FFN per expert
        self.experts = nn.ModuleList([SwiGLU(d_model) for _ in range(n_experts)])

        # Router: projects each token to a score over all experts
        self.router = nn.Linear(d_model, n_experts, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, C = x.shape
        x_flat = x.view(B * T, C)                  # (N, d_model)

        # --- Routing ---
        router_logits = self.router(x_flat)         # (N, n_experts)
        router_probs  = F.softmax(router_logits, dim=-1)

        # Select top-K experts per token
        top_k_probs, top_k_idx = router_probs.topk(self.top_k, dim=-1)
        top_k_probs = top_k_probs / top_k_probs.sum(dim=-1, keepdim=True)  # renormalise

        # --- Dispatch each token to its selected experts ---
        output = torch.zeros_like(x_flat)
        for expert_idx in range(self.n_experts):
            # Find which tokens selected this expert and at what weight
            mask = (top_k_idx == expert_idx).any(dim=-1)   # (N,)
            if not mask.any():
                continue

            tokens_for_expert = x_flat[mask]               # (n_selected, C)
            weight_idx = (top_k_idx[mask] == expert_idx).nonzero(as_tuple=True)[1]
            weights = top_k_probs[mask].gather(1, weight_idx.unsqueeze(1))

            expert_out = self.experts[expert_idx](tokens_for_expert)
            output[mask] += expert_out * weights           # weighted contribution

        # --- Load balancing auxiliary loss (Switch Transformer form) ---
        # Without this, the router collapses to always using the same 2 experts.
        # f_i = fraction of tokens actually DISPATCHED to expert i (hard top-k)
        # P_i = mean router probability for expert i (this part carries the gradient)
        dispatch = F.one_hot(top_k_idx, self.n_experts).sum(dim=1).float()  # (N, E)
        f = dispatch.mean(0) / self.top_k          # dispatch fractions, sum to 1
        P = router_probs.mean(0)                   # mean soft probabilities
        self.aux_loss = self.n_experts * (f * P).sum()

        return output.view(B, T, C)


# In your training loop, add the auxiliary loss:
# total_loss = main_loss + 0.01 * moe_layer.aux_loss

from torch.utils.data import Dataset, DataLoader


# --- Instruction dataset format ---
class InstructionDataset(Dataset):
    def __init__(self, examples: list[dict], tokenizer, max_len: int = 512):
        self.tokenizer = tokenizer
        self.max_len = max_len
        self.data = []

        for ex in examples:
            # Format: special tokens wrap the instruction/response
            text = (
                f"<|user|>\n{ex['instruction']}\n<|assistant|>\n{ex['response']}<|eos|>"
            )
            tokens = tokenizer.encode(text)
            self.data.append(tokens[:max_len])

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        tokens = self.data[idx]
        x = torch.tensor(tokens[:-1], dtype=torch.long)
        y = torch.tensor(tokens[1:],  dtype=torch.long)
        return x, y


# --- Only compute loss on the response tokens, not the instruction ---
def sft_loss(logits: torch.Tensor, targets: torch.Tensor,
             response_start_idx: int) -> torch.Tensor:
    # Mask out the instruction part — we only want the model to learn
    # to generate the response, not repeat the instruction
    mask = torch.zeros_like(targets, dtype=torch.bool)
    mask[:, response_start_idx:] = True

    logits_flat  = logits.view(-1, logits.size(-1))
    targets_flat = targets.view(-1)
    mask_flat    = mask.view(-1)

    return F.cross_entropy(logits_flat[mask_flat], targets_flat[mask_flat])


# --- LoRA for efficient fine-tuning (add to existing Linear layers) ---
class LoRALinear(nn.Module):
    def __init__(self, base: nn.Linear, r: int = 16, alpha: float = 32):
        super().__init__()
        self.base   = base
        self.lora_A = nn.Parameter(torch.randn(r, base.in_features) * 0.01)
        self.lora_B = nn.Parameter(torch.zeros(base.out_features, r))
        self.scale  = alpha / r

        # Freeze the base weights
        for p in self.base.parameters():
            p.requires_grad_(False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.base(x) + (x @ self.lora_A.T @ self.lora_B.T) * self.scale

examples = [
    {
        "instruction": "Explain gradient descent in simple terms.",
        "response": "Gradient descent is an optimisation algorithm that adjusts model weights by moving in the direction that most reduces the loss function — like rolling a ball downhill to find the lowest point in a valley."
    },
    {
        "instruction": "Write a Python function to reverse a string.",
        "response": "```python\ndef reverse_string(s: str) -> str:\n    return s[::-1]\n```"
    },
]

class RewardModel(nn.Module):
    """
    A transformer that maps (prompt + response) → scalar reward.
    Built on top of a pre-trained / SFT-fine-tuned backbone.
    """
    def __init__(self, backbone: TinyGPT):
        super().__init__()
        self.backbone = backbone
        d_model = backbone.token_emb.embedding_dim
        self.reward_head = nn.Linear(d_model, 1, bias=False)

    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
        # ⚠️ TinyGPT.forward() returns LOGITS (B, T, vocab_size) — calling it
        # here and feeding the result into a d_model→1 head crashes with a
        # shape error. The reward head needs HIDDEN STATES, so run the
        # backbone's internals and stop BEFORE the vocabulary projection:
        b = self.backbone
        positions = torch.arange(input_ids.shape[1], device=input_ids.device)
        h = b.token_emb(input_ids) + b.pos_emb(positions)
        for block in b.blocks:
            h = block(h)
        hidden = b.norm(h)                              # (B, T, d_model)

        # Use the last token's hidden state as the sequence representation
        last_hidden = hidden[:, -1, :]                  # (B, d_model)
        return self.reward_head(last_hidden).squeeze(-1) # (B,) scalar rewards


class RewardDataset(Dataset):
    """Each example is a (prompt, chosen_response, rejected_response) triple."""
    def __init__(self, examples: list[dict], tokenizer, max_len: int = 512):
        self.pairs = []
        for ex in examples:
            chosen   = tokenizer.encode(ex["prompt"] + ex["chosen"])[:max_len]
            rejected = tokenizer.encode(ex["prompt"] + ex["rejected"])[:max_len]
            self.pairs.append((
                torch.tensor(chosen,   dtype=torch.long),
                torch.tensor(rejected, dtype=torch.long),
            ))

    def __len__(self):
        return len(self.pairs)

    def __getitem__(self, idx):
        return self.pairs[idx]


def reward_model_loss(reward_model: RewardModel,
                      chosen_ids: torch.Tensor,
                      rejected_ids: torch.Tensor) -> torch.Tensor:
    """
    Bradley-Terry pairwise preference loss:
    maximise P(chosen > rejected) = σ(r_chosen - r_rejected)
    """
    r_chosen   = reward_model(chosen_ids)
    r_rejected = reward_model(rejected_ids)
    # Negative log-likelihood of chosen being preferred
    return -F.logsigmoid(r_chosen - r_rejected).mean()


# Example human preference data
examples = [
    {
        "prompt":    "<|user|>\nWhat is 2+2?\n<|assistant|>\n",
        "chosen":    "4.",
        "rejected":  "I'm not sure, mathematics is complex.",
    },
    {
        "prompt":    "<|user|>\nExplain photosynthesis.\n<|assistant|>\n",
        "chosen":    "Photosynthesis converts sunlight, CO₂, and water into glucose and oxygen.",
        "rejected":  "Plants do something with light I think.",
    },
]

import torch
import torch.nn.functional as F


def compute_ppo_loss(
    policy_logprobs:    torch.Tensor,   # log probs from the current policy
    old_logprobs:       torch.Tensor,   # log probs from the policy at sampling time
    rewards:            torch.Tensor,   # reward model scores
    ref_logprobs:       torch.Tensor,   # log probs from the frozen SFT reference model
    kl_coeff:           float = 0.1,    # KL penalty weight
    clip_eps:           float = 0.2,    # PPO clipping range
) -> torch.Tensor:

    # --- KL Penalty: penalise deviation from the SFT reference model ---
    kl_penalty = policy_logprobs - ref_logprobs   # per-token KL estimate
    shaped_reward = rewards - kl_coeff * kl_penalty.sum(dim=-1)

    # --- PPO Clipped Surrogate Objective ---
    log_ratio = policy_logprobs.sum(-1) - old_logprobs.sum(-1)
    ratio = log_ratio.exp()                        # probability ratio π_θ / π_old

    # Clip the ratio to [1-ε, 1+ε] to prevent too-large updates
    clipped_ratio = ratio.clamp(1 - clip_eps, 1 + clip_eps)
    ppo_loss = -torch.min(
        ratio * shaped_reward,
        clipped_ratio * shaped_reward,
    ).mean()

    return ppo_loss


def sequence_logprobs(model, full_ids: torch.Tensor,
                      prompt_len: int) -> torch.Tensor:
    """
    Log-prob of each RESPONSE token, conditioned on the full sequence.

    Two details people get wrong constantly:
    1. OFF-BY-ONE: logits at position t predict the token at position t+1.
       Gathering logits[t] with token[t] gives you a meaningless quantity.
       Shift: logits[:, :-1] scored against tokens[:, 1:].
    2. CONDITIONING: run the model on prompt+response, not the response
       alone — the response's probabilities depend on the prompt.
    """
    logits = model(full_ids)                            # (B, T, vocab)
    logprobs = F.log_softmax(logits[:, :-1], dim=-1)    # positions 0..T-2
    targets  = full_ids[:, 1:]                          # tokens    1..T-1
    token_logprobs = logprobs.gather(-1, targets.unsqueeze(-1)).squeeze(-1)
    return token_logprobs[:, prompt_len - 1:]           # response tokens only


def rlhf_training_step(
    policy_model,
    ref_model,          # frozen SFT model — provides the KL baseline
    reward_model,
    prompt_ids: torch.Tensor,
    optimizer,
):
    prompt_len = prompt_ids.shape[1]

    # 1. Sample a response from the current policy
    #    (generate = the Phase 2 sampling loop, wrapped as a method)
    with torch.no_grad():
        response_ids = policy_model.generate(prompt_ids, max_new_tokens=128)
        full_ids = torch.cat([prompt_ids, response_ids], dim=1)

        # 2. Score with reward model
        rewards = reward_model(full_ids)

        # 3. Reference (SFT) log probs for the KL penalty
        ref_logprobs = sequence_logprobs(ref_model, full_ids, prompt_len)

        # 4. "Old" policy log probs, frozen at sampling time
        old_logprobs = sequence_logprobs(policy_model, full_ids, prompt_len)

    # 5. Forward pass with gradient enabled
    policy_logprobs = sequence_logprobs(policy_model, full_ids, prompt_len)

    # 6. Compute PPO loss
    loss = compute_ppo_loss(
        policy_logprobs=policy_logprobs,
        old_logprobs=old_logprobs,
        rewards=rewards,
        ref_logprobs=ref_logprobs,
        kl_coeff=0.1,
    )

    optimizer.zero_grad()
    loss.backward()
    torch.nn.utils.clip_grad_norm_(policy_model.parameters(), 1.0)
    optimizer.step()

    return loss.item(), rewards.mean().item()

# The three numbers to watch during RLHF training:
print(f"PPO loss:      {ppo_loss:.4f}")      # should decrease
print(f"Mean reward:   {mean_reward:.4f}")    # should increase
print(f"KL divergence: {kl_div:.4f}")        # must stay bounded (< 4–6 nats)

# If KL divergence explodes: increase kl_coeff (0.1 → 0.2)
# If reward barely moves: decrease kl_coeff (0.1 → 0.05)
# If reward increases but outputs sound robotic: the reward model is being gamed