GPT from scratch: GPT-2, the KV cache and a model that learns chess

import importlib.util
import subprocess
import sys

TFS = "git+https://github.com/xmarva/transformers-from-scratch@main"
if importlib.util.find_spec("tfs") is None:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", TFS])

from tfs import env
from tfs.viz import table, token_chips, use_course_style

env.quiet()  # no progress bars or environment warnings in the outputs
SEED = 42
env.seed_everything(SEED)
COLORS = use_course_style()
WORK_DIR = env.work_dir()

import importlib.util
import json
import math
import random
import re
import subprocess
import sys
import time
from collections import Counter

# python-chess knows the rules of chess; our models will never be told them
if importlib.util.find_spec("chess") is None:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "chess==1.11.2"])
import chess
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
from huggingface_hub import get_safetensors_metadata, hf_hub_download
from IPython.display import display
from safetensors.torch import load_file
from tfs.viz import heatmap
from transformers import AutoTokenizer
# the reference implementation, for the check
from transformers import GPT2LMHeadModel

SMOKE = env.smoke()  # CI runs the lesson on a small slice; you run it in full
device = "cuda" if torch.cuda.is_available() else "cpu"
assert device == "cuda" or SMOKE, "This lesson trains a model for 25 minutes: use a GPU."
if device == "cuda":
    print("GPU:", torch.cuda.get_device_name(0))

GPT_FAMILY = ["openai-gpt", "gpt2", "gpt2-medium", "gpt2-large", "gpt2-xl"]
rows = []
for name in GPT_FAMILY:
    repo = "openai-community/" + name
    with open(hf_hub_download(repo, "config.json")) as f:
        c = json.load(f)
    stored = sum(get_safetensors_metadata(repo).parameter_count.values())
    rows.append({"model": name, "layers": c["n_layer"], "width": c["n_embd"],
                 "heads": c["n_head"], "context": c["n_positions"],
                 "vocabulary": c["vocab_size"], "numbers stored": stored})
family = pd.DataFrame(rows)
table(family)

def gpt_parameters(layers, d, vocab, context):
    per_layer = 12 * d * d + 13 * d
    embeddings = vocab * d + context * d
    final_norm = 2 * d  # GPT-2 has one more LayerNorm at the end
    return layers * per_layer + embeddings + final_norm


family["by the formula"] = [gpt_parameters(r.layers, r.width, r.vocabulary, r.context)
                            for r in family.itertuples()]
family["gap"] = family["numbers stored"] - family["by the formula"]
family["one mask per layer"] = family["layers"] * family["context"] ** 2
table(family[["model", "numbers stored", "by the formula", "gap", "one mask per layer"]])

fig, ax = plt.subplots(figsize=(6.4, 2.6))
names = family["model"][1:]
for k, r in enumerate(family[1:].itertuples()):
    embed = (r.vocabulary + r.context) * r.width / 1e6
    layers_part = r.layers * (12 * r.width ** 2 + 13 * r.width) / 1e6
    ax.barh(k, embed, color=COLORS["data"], label="embeddings" if k == 0 else None)
    ax.barh(k, layers_part, left=embed, color=COLORS["dec"],
            label="layers" if k == 0 else None)
    share = embed / (embed + layers_part)
    ax.text(embed + layers_part + 20, k, f"embeddings {share:.0%}", va="center",
            color=COLORS["muted"])
ax.set_yticks(range(len(names)), names)
ax.invert_yaxis()
ax.set_xlabel("millions of parameters")
ax.set_xlim(0, 2100)
ax.legend(loc="upper right", frameon=False)
plt.show()

def attention_by_hand(q, k, v):
    length = q.shape[-2]
    scores = q @ k.transpose(-2, -1) / math.sqrt(q.shape[-1])
    causal = torch.ones(length, length, dtype=torch.bool, device=q.device).tril()
    scores = scores.masked_fill(~causal, float("-inf"))
    return scores.softmax(dim=-1) @ v


# batch 1, 2 heads, 5 positions, 4 dimensions
q, k, v = torch.randn(3, 1, 2, 5, 4).unbind(0)
fused = F.scaled_dot_product_attention(q, k, v, is_causal=True)
print("largest difference:", (attention_by_hand(q, k, v) - fused).abs().max().item())

class Attention(nn.Module):
    def __init__(self, config):
        super().__init__()
        d = config["n_embd"]
        self.heads = config["n_head"]
        self.qkv = nn.Linear(d, 3 * d)  # queries, keys and values in one matrix
        self.out = nn.Linear(d, d)

    def split(self, x):
        batch, length, d = x.shape
        return x.view(batch, length, self.heads, d // self.heads).transpose(1, 2)

    def forward(self, x, cache=None):
        q, k, v = self.qkv(x).split(x.shape[-1], dim=-1)
        q, k, v = self.split(q), self.split(k), self.split(v)
        if cache is not None and "k" in cache:  # the keys and values of earlier tokens
            k = torch.cat([cache["k"], k], dim=2)
            v = torch.cat([cache["v"], v], dim=2)
        if cache is not None:
            cache["k"], cache["v"] = k, v
        # one new token may look at everything; several new tokens need the causal mask
        y = F.scaled_dot_product_attention(q, k, v, is_causal=q.shape[2] > 1)
        y = y.transpose(1, 2).reshape(x.shape)
        return self.out(y)

x = torch.linspace(-6, 6, 10001)
print("largest gap, tanh GELU against exact:",
      f"{(F.gelu(x, approximate='tanh') - F.gelu(x)).abs().max().item():.1e}")


class Block(nn.Module):
    def __init__(self, config):
        super().__init__()
        d, eps = config["n_embd"], config["layer_norm_epsilon"]
        self.norm1 = nn.LayerNorm(d, eps=eps)
        self.attention = Attention(config)
        self.norm2 = nn.LayerNorm(d, eps=eps)
        self.up = nn.Linear(d, 4 * d)
        self.down = nn.Linear(4 * d, d)
        # OpenAI's "gelu_new" is the tanh approximation
        self.approximate = "tanh" if config["activation_function"] == "gelu_new" else "none"

    def forward(self, x, cache=None):
        x = x + self.attention(self.norm1(x), cache)
        hidden = F.gelu(self.up(self.norm2(x)), approximate=self.approximate)
        return x + self.down(hidden)

class GPT(nn.Module):
    def __init__(self, config):
        super().__init__()
        d = config["n_embd"]
        self.tokens = nn.Embedding(config["vocab_size"], d)
        self.positions = nn.Embedding(config["n_positions"], d)
        self.blocks = nn.ModuleList(Block(config) for _ in range(config["n_layer"]))
        self.norm = nn.LayerNorm(d, eps=config["layer_norm_epsilon"])

    def forward(self, ids, caches=None, keep_states=False):
        past = 0
        if caches is not None and "k" in caches[0]:
            past = caches[0]["k"].shape[2]  # tokens already in the cache
        positions = torch.arange(past, past + ids.shape[1], device=ids.device)
        x = self.tokens(ids) + self.positions(positions)
        states = [x]
        for i, block in enumerate(self.blocks):
            x = block(x, caches[i] if caches is not None else None)
            states.append(x)
        logits = self.unembed(x)
        return (logits, states) if keep_states else logits

    def unembed(self, x):
        return self.norm(x) @ self.tokens.weight.T  # tied: the input table, read backwards

GPT2 = "openai-community/gpt2"
tokenizer = AutoTokenizer.from_pretrained(GPT2)
with open(hf_hub_download(GPT2, "config.json")) as f:
    config = json.load(f)


def openai_names(config):
    names = {"tokens.weight": "wte.weight", "positions.weight": "wpe.weight",
             "norm.weight": "ln_f.weight", "norm.bias": "ln_f.bias"}
    parts = {"norm1": "ln_1", "attention.qkv": "attn.c_attn",
             "attention.out": "attn.c_proj", "norm2": "ln_2",
             "up": "mlp.c_fc", "down": "mlp.c_proj"}
    for i in range(config["n_layer"]):
        for ours, theirs in parts.items():
            for kind in ("weight", "bias"):
                names[f"blocks.{i}.{ours}.{kind}"] = f"h.{i}.{theirs}.{kind}"
    return names


CONV1D = ("qkv.weight", "out.weight", "up.weight", "down.weight")
weights = load_file(hf_hub_download(GPT2, "model.safetensors"))
state = {}
for ours, theirs in openai_names(config).items():
    tensor = weights[theirs]
    if ours.endswith(CONV1D):
        tensor = tensor.T
    state[ours] = tensor
gpt2 = GPT(config)
gpt2.load_state_dict(state, strict=True)
gpt2 = gpt2.to(device).eval()
left_over = sorted(set(weights) - set(openai_names(config).values()))
print(f"loaded {sum(p.numel() for p in gpt2.parameters()):,} parameters")
example = left_over[0]
print(f"not loaded: {len(left_over)} tensors, e.g. {example}, "
      f"of shape {tuple(weights[example].shape)}")

reference = GPT2LMHeadModel.from_pretrained(GPT2).to(device).eval()
CHECK = ("In 1997 Deep Blue beat Garry Kasparov, the world chess champion, "
         "in a six-game match.")
check = tokenizer(CHECK, return_tensors="pt").to(device)
with torch.no_grad():
    theirs = reference(**check).logits
    ours = gpt2(check["input_ids"])
gap = (ours - theirs).abs().max().item()
print(f"largest difference in the scores: {gap:.1e}")
print(f"the scores themselves reach {theirs.abs().max():.0f}")
assert gap < 1e-2
del reference

def next_token_probabilities(text):
    ids = tokenizer(text, return_tensors="pt")["input_ids"].to(device)
    with torch.no_grad():
        logits = gpt2(ids)
    probs = logits[0, :-1].float().softmax(dim=-1)
    targets = ids[0, 1:]
    chosen = probs[torch.arange(len(targets)), targets]
    return tokenizer.convert_ids_to_tokens(ids[0]), chosen.cpu()


TEXT = ("AlphaGo, a program developed by Google DeepMind, beat Lee Sedol "
        "four games to one in Seoul.")
tokens, chosen = next_token_probabilities(TEXT)
print(f"loss {-chosen.log().mean():.2f}")
print("the probability of each token, given the ones before it:")
labels = [token.replace("Ġ", "␣") for token in tokens[1:]]
heatmap(chosen[None].numpy(), cols=labels, role="dec", vmin=0, vmax=1, annotate=True,
        cell=0.42)
plt.show()

@torch.no_grad()
def generate(model, ids, steps, use_cache=True, temperature=0.0, generator=None):
    caches = [{} for _ in model.blocks] if use_cache else None
    new = ids
    for _ in range(steps):
        logits = model(new if use_cache else ids, caches)[:, -1].float()
        if temperature == 0:
            token = logits.argmax(dim=-1, keepdim=True)
        else:
            probs = (logits / temperature).softmax(dim=-1)
            token = torch.multinomial(probs, 1, generator=generator)
        ids = torch.cat([ids, token], dim=1)
        new = token  # with a cache, only the new token goes in
    return ids


RUY_LOPEZ = "1. e4 e5 2. Nf3 Nc6 3. Bb5"
prompt = tokenizer(RUY_LOPEZ, return_tensors="pt")["input_ids"].to(device)
print(tokenizer.decode(generate(gpt2, prompt, 40)[0]))

with_cache = generate(gpt2, prompt, 40, use_cache=True)
without_cache = generate(gpt2, prompt, 40, use_cache=False)
print("same tokens with and without the cache:", torch.equal(with_cache, without_cache))
assert torch.equal(with_cache, without_cache)

@torch.no_grad()
def time_per_token(use_cache, steps):
    caches = [{} for _ in gpt2.blocks] if use_cache else None
    ids = prompt
    new = prompt
    times = []
    for _ in range(steps):
        if device == "cuda":
            torch.cuda.synchronize()
        start = time.perf_counter()
        logits = gpt2(new if use_cache else ids, caches)[:, -1]
        token = logits.argmax(dim=-1, keepdim=True)
        if device == "cuda":
            torch.cuda.synchronize()
        times.append(time.perf_counter() - start)
        ids = torch.cat([ids, token], dim=1)
        new = token
    return np.array(times) * 1000


STEPS = 60 if SMOKE else 900
time_per_token(True, 5)  # warm up the GPU
slow = time_per_token(False, STEPS)
fast = time_per_token(True, STEPS)
print(f"{STEPS} tokens: {slow.sum() / 1000:.1f} s without the cache, "
      f"{fast.sum() / 1000:.1f} s with it")

fig, ax = plt.subplots(figsize=(6.4, 3.0))
length = prompt.shape[1] + np.arange(STEPS)
ax.plot(length, slow, color=COLORS["line"], label="no cache: the whole text every step")
ax.plot(length, fast, color=COLORS["dec"], label="KV cache: one token every step")
ax.set_xlabel("tokens in the text")
ax.set_ylabel("ms for the next token")
ax.set_ylim(bottom=0)
ax.legend(frameon=False)
plt.show()

def kv_bytes_per_token(layers, d):
    return 2 * layers * d * 2  # keys and values, every layer, 2 bytes per number


for name, layers, d in [("GPT-2", 12, 768), ("GPT-2 XL", 48, 1600), ("GPT-3", 96, 12288)]:
    per_token = kv_bytes_per_token(layers, d)
    print(f"{name:9s} {per_token / 1024:7,.0f} KB per token, "
          f"{per_token * 2048 / 2**30:5.2f} GB for 2,048 tokens")

generator = torch.Generator(device=device).manual_seed(SEED)
START = "The strangest chess game ever played"
start = tokenizer(START, return_tensors="pt")["input_ids"].to(device)
for _ in range(3):
    sample = generate(gpt2, start, 30, temperature=0.8, generator=generator)
    print("·", tokenizer.decode(sample[0]).replace("\n", " "))

@torch.no_grad()
def logit_lens(text):
    ids = tokenizer(text, return_tensors="pt")["input_ids"].to(device)
    _, states = gpt2(ids, keep_states=True)
    rows = []
    for layer, state in enumerate(states):
        probs = gpt2.unembed(state[0, -1]).float().softmax(dim=-1)
        best = probs.argmax().item()
        rows.append((layer, tokenizer.decode(best), probs[best].item()))
    return rows


PROMPTS = ["In 1997, IBM's Deep Blue defeated the world chess champion Garry",
           "Bobby Fischer won the 1972 World Chess Championship against Boris"]
lens = {text: logit_lens(text) for text in PROMPTS}
fig, axes = plt.subplots(1, 2, figsize=(9.6, 4.4), sharey=True)
for ax, text in zip(axes, PROMPTS):
    rows = lens[text]
    layers = [layer for layer, _, _ in rows]
    top = [p for _, _, p in rows]
    ax.barh(layers, top, color=COLORS["dec"], height=0.7)
    for layer, word, p in rows:
        ax.text(p + 0.02, layer, repr(word), va="center", fontsize=9, color=COLORS["ink"])
    ax.set_xlim(0, 1.45)
    ax.set_title("… " + " ".join(text.split()[-5:]), fontsize=10)
    ax.set_xlabel("probability of the top token")
axes[0].set_ylabel("after layer (0 = embeddings)")
axes[0].set_yticks(range(len(lens[PROMPTS[0]])))
axes[0].invert_yaxis()
plt.tight_layout()
plt.show()

LICHESS = "Lichess/standard-chess-games"
SHARD = "data/year=2025/month=01/train-00000-of-00072.parquet"
path = hf_hub_download(LICHESS, SHARD, repo_type="dataset")
lichess = pd.read_parquet(path, columns=["Event", "WhiteElo", "BlackElo", "movetext"])
if SMOKE:
    lichess = lichess.head(100_000)
print(f"{len(lichess):,} games")
print(lichess["movetext"][0][:240], "…")

# "Rated Blitz game", "Blitz swiss https://lichess.org/swiss/…" -> "Blitz"
SPEEDS = r"(UltraBullet|Bullet|Blitz|Rapid|Classical|Correspondence)"
kind = lichess["Event"].str.extract(SPEEDS)[0]
table(kind.value_counts().rename_axis("time control").reset_index(name="games"))

CLOCK = re.compile(r"\{[^}]*\}")  # { [%clk 0:02:59] } and the like
BLACK_NUMBER = re.compile(r"\d+\.\.\. ")  # "1... " before Black's moves
GLYPHS = re.compile(r"[?!]")  # annotations such as ?? and !?
NUMBER = re.compile(r"(\d+)\. ")  # "1. e4" -> "1.e4"
RESULT = re.compile(r" ?(1-0|0-1|1/2-1/2|\*)$")


def clean(movetext):
    text = CLOCK.sub("", movetext)
    text = BLACK_NUMBER.sub("", text)
    text = GLYPHS.sub("", text)
    text = NUMBER.sub(r"\1.", text)
    text = " ".join(text.split())  # one space between moves
    text = RESULT.sub("", text)
    return ";" + text


slow = kind.isin(["Blitz", "Rapid", "Classical"])
strong = (lichess["WhiteElo"] >= 1800) & (lichess["BlackElo"] >= 1800)
games = [clean(movetext) for movetext in lichess["movetext"][slow & strong]]
games = [game for game in games if game.count(" ") >= 10]  # at least five moves each
print(f"{len(games):,} games kept, e.g.\n{games[0][:200]} …")

VOCAB = sorted(set("".join(games)))
print(len(VOCAB), "characters:", "".join(VOCAB))
stoi = {char: i for i, char in enumerate(VOCAB)}
lengths = np.array([len(game) for game in games])
print(f"a game is {np.median(lengths):.0f} characters long in the median, "
      f"{np.percentile(lengths, 90):.0f} at the 90th percentile")

CONTEXT = 512  # characters the model reads at once: most games fit
rng = np.random.default_rng(SEED)
order = rng.permutation(len(games))
test_games = [games[i] for i in order[:4000]]
train_games = [games[i] for i in order[4000:]]

lookup = np.zeros(128, dtype=np.uint8)
for char, i in stoi.items():
    lookup[ord(char)] = i
stream = lookup[np.frombuffer("".join(train_games).encode("ascii"), dtype=np.uint8)]
stream = torch.from_numpy(stream.astype(np.int64))
starts = np.flatnonzero(stream.numpy() == stoi[";"])
starts = starts[starts + CONTEXT + 1 <= len(stream)]  # every window must fit in the stream
print(f"{len(train_games):,} training games, {len(stream) / 1e6:.0f} million characters; "
      f"{len(test_games):,} held out")

LINE = "1. e4 e5 2. Nf3 Nc6 3. Bb5 a6 4. Ba4 Nf6 5. O-O Be7"
pieces = [tokenizer.decode([i]) for i in tokenizer(LINE)["input_ids"]]
token_chips([pieces], labels=["GPT-2 tokens"], max_width=640)
plt.show()
print(f"10 moves, {len(pieces)} tokens")

def moves_of(game):
    return [word.split(".")[-1] for word in game[1:].split(" ")]  # "12.Nf3" -> "Nf3"


def prompt_for(moves, upto, spaced):
    # the game before move number `upto` (counted from 0), ready for that move
    text = ""
    for i in range(upto):
        if i % 2 == 0:
            text += f"{i // 2 + 1}." + (" " if spaced else "")
        text += moves[i] + " "
    if upto % 2 == 0:  # White to move: the move number comes first
        text += f"{upto // 2 + 1}."
    return text


def test_positions(games, n, seed):
    rng = random.Random(seed)
    positions = []
    for game in games:
        moves = moves_of(game)
        upto = rng.randrange(1, min(len(moves), 70))
        if len(prompt_for(moves, upto, spaced=False)) > CONTEXT - 10:
            continue  # longer than our model's context
        board = chess.Board()
        for move in moves[:upto]:
            board.push_san(move)
        positions.append({"moves": moves, "upto": upto, "board": board,
                          "played": moves[upto]})
        if len(positions) == n:
            break
    return positions


positions = test_positions(test_games, 100 if SMOKE else 1000, SEED)
example = positions[0]
moves, upto = example["moves"], example["upto"]
print("GPT-2 gets:    ", repr(prompt_for(moves, upto, spaced=True).rstrip()))
print("our model gets:", repr(";" + prompt_for(moves, upto, spaced=False)))
print("the player played:", example["played"])

def judge(board, move):
    try:
        board.parse_san(move)
        return "legal"
    except chess.InvalidMoveError:
        return "not a move"
    except chess.AmbiguousMoveError:
        return "ambiguous"
    except chess.IllegalMoveError:
        pass
    pseudo_legal = {board.san(m).rstrip("+#") for m in board.pseudo_legal_moves}
    if move.rstrip("+#") in pseudo_legal:
        return "leaves the king in check"
    return "no piece can do that"


def score(make_move, positions):
    verdicts, same_as_player = [], 0
    for position in positions:
        move = make_move(position)
        verdicts.append(judge(position["board"], move))
        same_as_player += move.rstrip("+#") == position["played"].rstrip("+#")
    shares = {}
    for verdict, count in Counter(verdicts).items():
        shares[verdict] = count / len(positions)
    return shares, same_as_player / len(positions)


@torch.no_grad()
def gpt2_move(position):
    prompt = prompt_for(position["moves"], position["upto"], spaced=True).rstrip()
    ids = tokenizer(prompt, return_tensors="pt")["input_ids"].to(device)
    continuation = tokenizer.decode(generate(gpt2, ids, 5)[0, ids.shape[1]:])
    words = continuation.split()
    return words[0] if words else ""


gpt2_shares, gpt2_same = score(gpt2_move, positions)
print(f"GPT-2: {gpt2_shares.get('legal', 0.0):.1%} legal moves, "
      f"{gpt2_same:.1%} the same move as the player")
table(pd.DataFrame({"verdict": list(gpt2_shares), "share": list(gpt2_shares.values())}))

chess_config = {"vocab_size": len(VOCAB), "n_positions": CONTEXT, "n_embd": 512,
                "n_head": 8, "n_layer": 8, "layer_norm_epsilon": 1e-5,
                "activation_function": "gelu_new"}


def init_like_gpt2(model, layers):
    for module in model.modules():
        if isinstance(module, (nn.Linear, nn.Embedding)):
            nn.init.normal_(module.weight, std=0.02)
        if isinstance(module, nn.Linear):
            nn.init.zeros_(module.bias)
    for block in model.blocks:  # the two matrices that add to the residual stream
        nn.init.normal_(block.attention.out.weight, std=0.02 / math.sqrt(2 * layers))
        nn.init.normal_(block.down.weight, std=0.02 / math.sqrt(2 * layers))


torch.manual_seed(SEED)
chess_gpt = GPT(chess_config)
init_like_gpt2(chess_gpt, chess_config["n_layer"])
chess_gpt = chess_gpt.to(device)
print(f"{sum(p.numel() for p in chess_gpt.parameters()) / 1e6:.1f} million parameters")

@torch.no_grad()
def chess_move(model, position, generator=None):
    prompt = ";" + prompt_for(position["moves"], position["upto"], spaced=False)
    ids = torch.tensor([[stoi[char] for char in prompt]], device=device)
    caches = [{} for _ in model.blocks]
    move = ""
    for _ in range(8):  # the longest move, "exd8=Q+", has 7 characters
        logits = model(ids, caches)[0, -1].float()
        if generator is None:
            char = logits.argmax().item()
        else:
            char = torch.multinomial(logits.softmax(dim=-1), 1, generator=generator).item()
        if VOCAB[char] == " ":
            break
        move += VOCAB[char]
        ids = torch.tensor([[char]], device=device)
    return move


def windows(rng, size):
    picked = starts[rng.integers(0, len(starts), size)]
    rows = torch.stack([stream[s:s + CONTEXT + 1] for s in picked]).to(device)
    return rows[:, :-1], rows[:, 1:]  # inputs, and the same characters shifted by one


BATCH, PEAK_LR, WARMUP = 64, 1e-3, 300
STEPS = 30 if SMOKE else 3400
EVAL_EVERY = 10 if SMOKE else 500


def learning_rate(step):
    if step < WARMUP:
        return PEAK_LR * (step + 1) / WARMUP
    progress = (step - WARMUP) / max(1, STEPS - WARMUP)
    cosine = 0.5 * (1 + math.cos(math.pi * progress))
    return PEAK_LR * (0.1 + 0.9 * cosine)  # down to a tenth of the peak


matrices = [p for p in chess_gpt.parameters() if p.dim() == 2]
others = [p for p in chess_gpt.parameters() if p.dim() < 2]
groups = [{"params": matrices, "weight_decay": 0.1},
          {"params": others, "weight_decay": 0.0}]
optimizer = torch.optim.AdamW(groups, lr=PEAK_LR, betas=(0.9, 0.95))
scaler = torch.amp.GradScaler(enabled=device == "cuda")
# the compiler fuses the training passes into fewer GPU kernels; the first steps wait for it
compiled = torch.compile(chess_gpt) if device == "cuda" else chess_gpt
quick_test = positions[:200]

rng = np.random.default_rng(SEED)
curve = []
start = time.time()
for step in range(STEPS):
    for group in optimizer.param_groups:
        group["lr"] = learning_rate(step)
    inputs, targets = windows(rng, BATCH)
    with torch.autocast(device, dtype=torch.float16, enabled=device == "cuda"):
        logits = compiled(inputs)
    loss = F.cross_entropy(logits.float().view(-1, len(VOCAB)), targets.reshape(-1))
    optimizer.zero_grad(set_to_none=True)
    scaler.scale(loss).backward()
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(chess_gpt.parameters(), 1.0)
    scaler.step(optimizer)
    scaler.update()
    if (step + 1) % EVAL_EVERY == 0 or step + 1 == STEPS:
        chess_gpt.eval()
        shares, _ = score(lambda position: chess_move(chess_gpt, position), quick_test)
        chess_gpt.train()
        minutes = (time.time() - start) / 60
        curve.append({"step": step + 1, "minutes": minutes, "loss": loss.item(),
                      "legal": shares.get("legal", 0.0)})
        print(f"step {step + 1:5d}  {minutes:5.1f} min  loss {loss.item():.3f}  "
              f"legal moves {shares.get('legal', 0.0):.1%}")
chess_gpt = chess_gpt.eval()

curve = pd.DataFrame(curve)
fig, (left, right) = plt.subplots(1, 2, figsize=(9.6, 3.0))
left.plot(curve["minutes"], curve["loss"], color=COLORS["dec"], marker="o", markersize=3)
left.set_xlabel("minutes of training")
left.set_ylabel("loss per character")
right.plot(curve["minutes"], curve["legal"], color=COLORS["dec"], marker="o", markersize=3,
           label="our chess GPT")
right.axhline(gpt2_shares.get("legal", 0.0), color=COLORS["line"], linestyle="--",
              label="GPT-2")
right.set_xlabel("minutes of training")
right.set_ylabel("legal moves, 200 positions")
right.set_ylim(0, 1)
right.legend(frameon=False, loc="lower right")
plt.tight_layout()
plt.show()

chess_shares, chess_same = score(lambda p: chess_move(chess_gpt, p), positions)
print(f"our chess GPT: {chess_shares.get('legal', 0.0):.1%} legal moves, "
      f"{chess_same:.1%} the same move as the player")

path = hf_hub_download("adamkarvonen/chess_llms", "lichess_8layers_ckpt_no_optimizer.pt")
checkpoint = torch.load(path, map_location="cpu", weights_only=True)
args = checkpoint["model_args"]
print(args)
seen = checkpoint["iter_num"] * checkpoint["config"]["batch_size"] * args["block_size"]
print(f"characters seen in training: {seen / 1e9:.0f} billion; "
      f"ours: {STEPS * BATCH * CONTEXT / 1e6:.0f} million")

assert "".join(VOCAB) == " #+-.0123456789;=BKNOQRabcdefghx"  # his alphabet, in his order
karvonen_config = {"vocab_size": args["vocab_size"], "n_positions": args["block_size"],
                   "n_embd": args["n_embd"], "n_head": args["n_head"],
                   "n_layer": args["n_layer"], "layer_norm_epsilon": 1e-5,
                   "activation_function": "gelu"}  # the exact GELU
weights = checkpoint["model"]
karvonen = GPT(karvonen_config)
state = karvonen.state_dict()
zeroed = 0
for ours, theirs in openai_names(karvonen_config).items():
    key = "_orig_mod.transformer." + theirs  # torch.compile added "_orig_mod."
    if key in weights:
        state[ours] = weights[key]
    else:  # trained without biases: a zero bias is the same thing
        state[ours] = torch.zeros_like(state[ours])
        zeroed += 1
karvonen.load_state_dict(state, strict=True)
karvonen = karvonen.to(device).eval()
output_layer = weights["_orig_mod.lm_head.weight"]
tied = torch.equal(output_layer, weights["_orig_mod.transformer.wte.weight"])
print(f"{len(weights)} tensors in the file; {zeroed} biases of ours set to zero; "
      f"output layer tied to the embeddings: {tied}")

karvonen_shares, karvonen_same = score(lambda p: chess_move(karvonen, p), positions)
verdicts = ["legal", "leaves the king in check", "no piece can do that", "ambiguous",
            "not a move"]
by_model = {"GPT-2, 124M": gpt2_shares, "ours, 25 minutes": chess_shares,
            "Karvonen's, 72 hours": karvonen_shares}
comparison = pd.DataFrame({"verdict": verdicts})
for name, shares in by_model.items():
    comparison[name] = [shares.get(verdict, 0.0) for verdict in verdicts]
display(table(comparison, digits=3))
print(f"the move the player actually made: GPT-2 {gpt2_same:.1%}, ours {chess_same:.1%}, "
      f"Karvonen's {karvonen_same:.1%}")

def play_against_random(model, seed, max_moves=150):
    rng = random.Random(seed)
    generator = torch.Generator(device=device).manual_seed(seed)
    limit = model.positions.num_embeddings  # the model's context
    board = chess.Board()
    moves, illegal = [], 0
    while not board.is_game_over() and len(moves) < max_moves:
        if board.turn == chess.WHITE:
            if len(prompt_for(moves, len(moves), spaced=False)) > limit - 10:
                return "unfinished", moves, illegal
            position = {"moves": moves, "upto": len(moves)}
            move = None
            for attempt in range(5):
                candidate = chess_move(model, position, generator if attempt > 0 else None)
                if judge(board, candidate) == "legal":
                    move = candidate
                    break
                illegal += 1
            if move is None:
                return "forfeit", moves, illegal
        else:
            move = board.san(rng.choice(list(board.legal_moves)))
        move = board.san(board.parse_san(move))  # the standard spelling, + where it checks
        board.push_san(move)
        moves.append(move)
    if not board.is_game_over():
        return "unfinished", moves, illegal
    return board.result(), moves, illegal


played = []
for name, model in [("ours, 25 minutes", chess_gpt), ("Karvonen's, 72 hours", karvonen)]:
    for seed in range(10 if SMOKE else 100):
        result, moves, illegal = play_against_random(model, seed)
        played.append({"model": name, "result": result, "half-moves": len(moves),
                       "illegal tries": illegal, "game": moves})
played = pd.DataFrame(played)
outcomes = pd.crosstab(played["model"], played["result"])
outcomes["illegal tries per game"] = played.groupby("model")["illegal tries"].mean()
table(outcomes.reset_index())

# White in capitals, Black in small letters; class 0 is an empty square
PIECES = "PNBRQKpnbrqk"


def board_classes(board):
    classes = np.zeros(64, dtype=np.int64)  # square 0 is a1, 7 is h1, 63 is h8
    for square, piece in board.piece_map().items():
        classes[square] = 1 + PIECES.index(piece.symbol())
    return classes


def draw_board(ax, classes, title, truth=None):
    for square in range(64):
        file, rank = square % 8, square // 8
        light = (file + rank) % 2 == 1
        color = COLORS["card"] if light else COLORS["grid"]
        ax.add_patch(plt.Rectangle((file, rank), 1, 1, color=color))
        if truth is not None and classes[square] != truth[square]:
            ax.add_patch(plt.Rectangle((file + 0.04, rank + 0.04), 0.92, 0.92, fill=False,
                                       edgecolor=COLORS["dec"], linewidth=2))
        if classes[square] > 0:
            piece = PIECES[classes[square] - 1]
            ax.text(file + 0.5, rank + 0.47, piece, ha="center", va="center", fontsize=13,
                    weight="bold", family="monospace",
                    color=COLORS["ink"] if piece.isupper() else COLORS["muted"])
    ax.set_xlim(0, 8)
    ax.set_ylim(0, 8)
    ax.set_xticks(np.arange(8) + 0.5, list("abcdefgh"))
    ax.set_yticks(np.arange(8) + 0.5, range(1, 9))
    ax.tick_params(length=0)
    ax.set_aspect("equal")
    ax.grid(False)
    ax.set_title(title, fontsize=10)


wins = played[played["result"] == "1-0"]
if len(wins):
    shortest = wins.loc[wins["half-moves"].idxmin()]
    game = shortest["game"]
    print(shortest["model"] + ":", prompt_for(game, len(game), spaced=True).strip())
    board = chess.Board()
    for move in game:
        board.push_san(move)
    fig, ax = plt.subplots(figsize=(3.2, 3.4))
    draw_board(ax, board_classes(board), f"after {len(game)} half-moves: checkmate")
    plt.show()