What Llama changed in the Transformer: RoPE, RMSNorm, SwiGLU

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 json
import math
import os
import random
import time

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 HfApi, get_safetensors_metadata, hf_hub_download
from safetensors.torch import load_file, save_file
from tfs import decoding, training
from tfs.comet import Comet
from tfs.model import Transformer, positional_encoding
from tokenizers import Tokenizer

# sacreBLEU, the reference implementation of BLEU and chrF
if importlib.util.find_spec("sacrebleu") is None:
    pip = [sys.executable, "-m", "pip", "install", "-q"]
    subprocess.check_call(pip + ["sacrebleu==2.6.0"])
import sacrebleu

SMOKE = env.smoke()  # CI runs the lesson on a small slice; you train in full
device = "cuda" if torch.cuda.is_available() else "cpu"
assert device == "cuda" or SMOKE, "This lesson trains ten models: switch to a GPU."
if device == "cuda":
    print("GPU:", torch.cuda.get_device_name(0))

LLAMA = "unsloth/Meta-Llama-3.1-8B-Instruct"  # Meta's Llama 3.1 8B, no license gate
with open(hf_hub_download(LLAMA, "config.json")) as f:
    llama = json.load(f)

FIELDS = ["num_hidden_layers", "hidden_size", "num_attention_heads", "num_key_value_heads",
          "intermediate_size", "hidden_act", "rms_norm_eps", "rope_theta", "vocab_size"]
table(pd.DataFrame({"setting": FIELDS, "value": [str(llama[name]) for name in FIELDS]}))

metadata = get_safetensors_metadata(LLAMA)
rows, total, biases = [], 0, []
for file_metadata in metadata.files_metadata.values():
    for name, tensor in file_metadata.tensors.items():
        total += math.prod(tensor.shape)
        if name.endswith("bias"):
            biases.append(name)
        if ".layers.0." in name:
            short = name.removeprefix("model.layers.0.")
            rows.append({"tensor": short, "shape": str(tensor.shape)})
print(f"parameters, counted from the headers: {total / 1e9:.2f} billion")
print("bias tensors:", len(biases))
table(pd.DataFrame(rows))

tokenizer_path = hf_hub_download("xmarva/tfs-bpe-en-ru-16k", "tokenizer.json")
tokenizer = Tokenizer.from_file(tokenizer_path)
splits = {}
for name in ["train", "validation", "test"]:
    path = hf_hub_download("xmarva/tfs-tatoeba-en-ru", f"data/{name}.parquet",
                           repo_type="dataset")
    splits[name] = pd.read_parquet(path)
train, valid, test = splits["train"], splits["validation"], splits["test"]
if SMOKE:
    train = train.sample(3_000, random_state=SEED)
    valid = valid.head(300)
    test = test.sample(300, random_state=SEED)

VOCAB_SIZE = tokenizer.get_vocab_size()
PAD = tokenizer.token_to_id("<pad>")
BOS = tokenizer.token_to_id("<bos>")
EOS = tokenizer.token_to_id("<eos>")
print("train:", len(train), "validation:", len(valid), "test:", len(test))

def encode_all(sentences):
    return [encoding.ids for encoding in tokenizer.encode_batch(list(sentences))]


train_src, train_tgt = encode_all(train.en), encode_all(train.ru)
valid_src, valid_tgt = encode_all(valid.en), encode_all(valid.ru)

MAX_LEN = 64  # lesson 3 left out the few longer pairs
keep = []
for i in range(len(train_src)):
    if len(train_src[i]) <= MAX_LEN and len(train_tgt[i]) + 1 <= MAX_LEN:
        keep.append(i)

special = {"pad_id": PAD, "bos_id": BOS, "eos_id": EOS}
train_batches = training.make_batches(train_src, train_tgt, keep, 6_000, **special)
valid_rows = range(len(valid_src))
valid_batches = training.make_batches(valid_src, valid_tgt, valid_rows, 6_000, **special)
print("training batches:", len(train_batches), "validation batches:", len(valid_batches))

def attention(q, k, v, mask=None):
    scores = q @ k.transpose(-2, -1) / math.sqrt(q.shape[-1])
    if mask is not None:
        scores = scores.masked_fill(~mask, float("-inf"))  # hidden positions get no weight
    return scores.softmax(dim=-1) @ v


q, k, v = (torch.randn(8, 4, 40, 64, device=device) for _ in range(3))
causal = torch.ones(40, 40, dtype=torch.bool, device=device).tril()[None, None]
ours = attention(q, k, v, causal)
fused = F.scaled_dot_product_attention(q, k, v, attn_mask=causal)
assert torch.allclose(ours, fused, atol=1e-5)
print("largest difference:", (ours - fused).abs().max().item())

def residual(x, sublayer, norm, dropout, pre_norm):
    if pre_norm:
        return x + dropout(sublayer(norm(x)))  # normalize the copy, add to the stream
    return norm(x + dropout(sublayer(x)))  # 2017: add, then normalize the stream

class RMSNorm(nn.Module):
    def __init__(self, d_model, eps=1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(d_model))  # g; there is no b

    def forward(self, x):
        x32 = x.float()
        mean_square = x32.pow(2).mean(dim=-1, keepdim=True)
        normed = x32 * torch.rsqrt(mean_square + self.eps)  # rsqrt: 1 / sqrt
        return (normed * self.weight).to(x.dtype)


x = torch.randn(8, 16, 256)
ours = RMSNorm(256)(x)
theirs = nn.RMSNorm(256, eps=1e-6)(x)  # PyTorch's own, since version 2.4
assert torch.allclose(ours, theirs, atol=1e-6)
print("RMSNorm agrees with torch.nn.RMSNorm")

z = torch.linspace(-5, 3, 400)
fig, ax = plt.subplots(figsize=(6.4, 2.8))
ax.plot(z, torch.relu(z), color=COLORS["muted"], label="ReLU")
ax.plot(z, F.silu(z), color=COLORS["enc"], label="SiLU(z) = z·σ(z)")
ax.axhline(0, color=COLORS["line"], linewidth=0.8)
ax.set_xlabel("z")
ax.set_title("ReLU and SiLU")
ax.legend()
plt.show()
print("smallest value of SiLU:", round(F.silu(z).min().item(), 3))

class SwiGLU(nn.Module):
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.w_gate = nn.Linear(d_model, d_ff, bias=False)
        self.w_up = nn.Linear(d_model, d_ff, bias=False)
        self.w_down = nn.Linear(d_ff, d_model, bias=False)

    def forward(self, x):
        return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x))


class FeedForward(nn.Module):  # lesson 2's, for comparison
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.linear1 = nn.Linear(d_model, d_ff)
        self.linear2 = nn.Linear(d_ff, d_model)

    def forward(self, x):
        return self.linear2(torch.relu(self.linear1(x)))


def swiglu_width(d_ff, multiple=8):
    return multiple * round(2 * d_ff / 3 / multiple)


def parameters(module):
    return sum(p.numel() for p in module.parameters())


print("ReLU, d_ff = 1024:", parameters(FeedForward(256, 1024)))
print(f"SwiGLU, d_ff = {swiglu_width(1024)}:", parameters(SwiGLU(256, swiglu_width(1024))))

def rope_angles(max_len, d_k, base=10_000):
    theta = base ** (-torch.arange(0, d_k, 2) / d_k)  # one speed per pair: (d_k / 2,)
    positions = torch.arange(max_len)
    return positions[:, None] * theta[None, :]  # the angle m·θ_i: (max_len, d_k / 2)


d_k = 256 // 4  # our heads
wavelengths = 2 * math.pi / rope_angles(2, d_k)[1]  # tokens per full turn, pair by pair
llama_wavelengths = 2 * math.pi / rope_angles(2, 128, base=llama["rope_theta"])[1]

fig, ax = plt.subplots(figsize=(6.4, 3))
label = "ours: d_k = 64, b = 10 000"
ax.plot(wavelengths, "o", color=COLORS["enc"], markersize=4, label=label)
label = f"Llama 3.1: d_k = 128, b = {llama['rope_theta']:,.0f}"
ax.plot(llama_wavelengths, ".", color=COLORS["att"], markersize=3, label=label)
ax.axhline(MAX_LEN, color=COLORS["dec"], linestyle=(0, (3, 3)), linewidth=1)
ax.text(40, MAX_LEN * 1.4, "our longest training sentence", fontsize=9, color=COLORS["dec"])
ax.set_yscale("log")
ax.set_xlabel("pair i")
ax.set_ylabel("tokens per full turn")
ax.set_title("Every pair rotates at its own speed")
ax.legend(fontsize=9)
plt.show()

def apply_rope(x, angles):
    # x: (B, H, T, d_k); angles: (T, d_k / 2), one angle per position and pair
    cos, sin = angles.cos().to(x.dtype), angles.sin().to(x.dtype)
    even, odd = x[..., 0::2], x[..., 1::2]
    rotated = torch.empty_like(x)
    rotated[..., 0::2] = even * cos - odd * sin
    rotated[..., 1::2] = even * sin + odd * cos
    return rotated

q, k = torch.randn(1, 1, 30, 64), torch.randn(1, 1, 30, 64)
angles = rope_angles(100, 64)
scores_here = apply_rope(q, angles[:30]) @ apply_rope(k, angles[:30]).transpose(-1, -2)
scores_later = apply_rope(q, angles[7:37]) @ apply_rope(k, angles[7:37]).transpose(-1, -2)
assert torch.allclose(scores_here, scores_later, atol=1e-4)

# Meta's way (llama/model.py): pairs as complex numbers, times e^(i·angle)
q_complex = torch.view_as_complex(q.reshape(1, 1, 30, 32, 2))
rotation = torch.polar(torch.ones(30, 32), angles[:30])  # e^(i·m·θ_i)
meta = torch.view_as_real(q_complex * rotation).flatten(-2)
assert torch.allclose(meta, apply_rope(q, angles[:30]), atol=1e-6)
print("shift-invariant, and equal to the complex-number version")

class Attention(nn.Module):
    def __init__(self, d_model, n_heads, fused=True):
        super().__init__()
        self.n_heads, self.d_k, self.fused = n_heads, d_model // n_heads, fused
        self.w_q = nn.Linear(d_model, d_model)
        self.w_k = nn.Linear(d_model, d_model)
        self.w_v = nn.Linear(d_model, d_model)
        self.w_o = nn.Linear(d_model, d_model)

    def split_heads(self, x):
        batch, length, _ = x.shape
        return x.view(batch, length, self.n_heads, self.d_k).transpose(1, 2)

    def keys_values(self, x, angles=None):
        k, v = self.split_heads(self.w_k(x)), self.split_heads(self.w_v(x))
        if angles is not None:
            k = apply_rope(k, angles)
        return k, v

    def forward(self, x_query, x_key_value, mask, angles=None, cache=None):
        q = self.split_heads(self.w_q(x_query))
        if angles is not None:
            q = apply_rope(q, angles)
        if cache is not None and cache.get("fixed"):  # cross-attention, once (section 11)
            k, v = cache["k"], cache["v"]
        else:
            k, v = self.keys_values(x_key_value, angles)
            if cache is not None:  # self-attention: the cached past plus the new ones
                if "k" in cache:
                    k = torch.cat([cache["k"], k], dim=2)
                    v = torch.cat([cache["v"], v], dim=2)
                cache["k"], cache["v"] = k, v
        if self.fused:
            out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
        else:
            out = attention(q, k, v, mask)
        batch, _, length, _ = out.shape
        return self.w_o(out.transpose(1, 2).reshape(batch, length, -1))

class Block(nn.Module):
    def __init__(self, d_model, n_heads, d_ff, dropout, cross, setup):
        super().__init__()
        self.pre_norm = setup["norm"] != "post"
        Norm = RMSNorm if setup["norm"] == "rms" else nn.LayerNorm
        fused = setup.get("fused", True)
        self.self_attention = Attention(d_model, n_heads, fused)
        self.cross_attention = Attention(d_model, n_heads, fused) if cross else None
        if setup["ffn"] == "swiglu":
            self.feed_forward = SwiGLU(d_model, swiglu_width(d_ff))
        else:
            self.feed_forward = FeedForward(d_model, d_ff)
        self.norm1, self.norm2 = Norm(d_model), Norm(d_model)
        if cross:
            self.norm3 = Norm(d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask, angles, memory=None, memory_mask=None, cache=None):
        def sub(x, sublayer, norm):
            return residual(x, sublayer, norm, self.dropout, self.pre_norm)

        own = cache["self"] if cache else None
        x = sub(x, lambda h: self.self_attention(h, h, mask, angles, own), self.norm1)
        if self.cross_attention is None:  # an encoder block
            return sub(x, self.feed_forward, self.norm2)
        other = cache["cross"] if cache else None

        def cross(h):
            return self.cross_attention(h, memory, memory_mask, cache=other)

        x = sub(x, cross, self.norm2)
        return sub(x, self.feed_forward, self.norm3)

class Seq2Seq(nn.Module):
    def __init__(self, setup, vocab_size, d_model=256, n_heads=4, d_ff=1024, n_layers=3,
                 dropout=0.1, pad_id=0, max_len=256):
        super().__init__()
        self.setup, self.d_model, self.pad_id = setup, d_model, pad_id
        self.embedding = nn.Embedding(vocab_size, d_model)
        nn.init.normal_(self.embedding.weight, mean=0.0, std=d_model**-0.5)
        if setup["pos"] == "sinusoidal":
            self.register_buffer("pe", positional_encoding(max_len, d_model))
        else:
            self.register_buffer("angles", rope_angles(max_len, d_model // n_heads))
        args = (d_model, n_heads, d_ff, dropout)
        self.encoder = nn.ModuleList([Block(*args, False, setup) for _ in range(n_layers)])
        self.decoder = nn.ModuleList([Block(*args, True, setup) for _ in range(n_layers)])
        self.dropout = nn.Dropout(dropout)
        if setup["norm"] != "post":  # pre-norm: one more norm at the end of each stack
            Norm = RMSNorm if setup["norm"] == "rms" else nn.LayerNorm
            self.encoder_norm, self.decoder_norm = Norm(d_model), Norm(d_model)

    def embed(self, ids, start=0):
        x = self.embedding(ids) * math.sqrt(self.d_model)
        if self.setup["pos"] == "sinusoidal":
            x = x + self.pe[start : start + ids.shape[1]]
        return self.dropout(x)

    def rope(self, length, start=0):
        if self.setup["pos"] != "rope":
            return None
        return self.angles[start : start + length]

    def encode(self, src):
        mask = (src != self.pad_id)[:, None, None, :]  # (B, 1, 1, S): which keys are real
        x = self.embed(src)
        for layer in self.encoder:
            x = layer(x, mask, self.rope(src.shape[1]))
        return self.encoder_norm(x) if self.setup["norm"] != "post" else x

    def decode(self, tgt, memory, src):
        length = tgt.shape[1]
        causal = torch.ones(length, length, dtype=torch.bool, device=tgt.device).tril()
        mask = causal[None, None] & (tgt != self.pad_id)[:, None, None, :]
        memory_mask = (src != self.pad_id)[:, None, None, :]
        x = self.embed(tgt)
        for layer in self.decoder:
            x = layer(x, mask, self.rope(length), memory, memory_mask)
        if self.setup["norm"] != "post":
            x = self.decoder_norm(x)
        return x @ self.embedding.weight.T

    def forward(self, src, tgt):
        return self.decode(tgt, self.encode(src), src)

VARIANTS = {
    "2017": {"norm": "post", "ffn": "relu", "pos": "sinusoidal"},
    "+ pre-norm": {"norm": "pre", "ffn": "relu", "pos": "sinusoidal"},
    "+ RMSNorm": {"norm": "rms", "ffn": "relu", "pos": "sinusoidal"},
    "+ SwiGLU": {"norm": "rms", "ffn": "swiglu", "pos": "sinusoidal"},
    "+ RoPE": {"norm": "rms", "ffn": "swiglu", "pos": "rope"},
}


def new_model(setup, seed=SEED):
    torch.manual_seed(seed)  # the same starting weights wherever the shapes allow
    return Seq2Seq(setup, VOCAB_SIZE, pad_id=PAD).to(device)


rows = []
for name, setup in VARIANTS.items():
    rows.append({"model": name, "parameters": parameters(new_model(setup))})
table(pd.DataFrame(rows))

MODEL_REPO = "xmarva/tfs-transformer-en-ru"
with open(hf_hub_download(MODEL_REPO, "config.json")) as f:
    CONFIG = json.load(f)
lesson3 = Transformer(**CONFIG).to(device).eval()
lesson3.load_state_dict(load_file(hf_hub_download(MODEL_REPO, "model.safetensors")))

rebuilt = new_model(VARIANTS["2017"]).eval()
rebuilt.load_state_dict(lesson3.state_dict())  # every name must match, nothing left over
src_ids, tgt_in, _ = (t.to(device) for t in valid_batches[0])
with torch.no_grad():
    difference = (lesson3(src_ids, tgt_in) - rebuilt(src_ids, tgt_in)).abs().max().item()
assert difference < 1e-3
print("largest difference from lesson 3's model:", difference)

def ms_per_step(model, n_steps=30):
    optimizer = training.adam(model)
    run = dict(pad_id=PAD, device=device)
    training.train_steps(model, optimizer, train_batches, lambda step: 1e-4, 5, **run)
    if device == "cuda":
        torch.cuda.synchronize()
    start = time.perf_counter()
    training.train_steps(model, optimizer, train_batches, lambda step: 1e-4, n_steps, **run)
    if device == "cuda":
        torch.cuda.synchronize()
    return (time.perf_counter() - start) / n_steps * 1000


TIMING_STEPS = 30 if not SMOKE else 2
manual = ms_per_step(new_model({**VARIANTS["2017"], "fused": False}), TIMING_STEPS)
fused = ms_per_step(new_model(VARIANTS["2017"]), TIMING_STEPS)
print(f"lesson 2's attention: {manual:.0f} ms per step; fused: {fused:.0f} ms per step")

PEAK = 3e-3  # lesson 3's experiment
EXPERIMENT_STEPS = 1_500 if not SMOKE else 10
experiments = [
    ("2017, no warm-up", VARIANTS["2017"], lambda step: PEAK),
    ("2017, warm-up over 1 000 steps", VARIANTS["2017"],
     lambda step: PEAK * min(1.0, step / 1_000)),
    ("pre-norm, no warm-up", VARIANTS["+ pre-norm"], lambda step: PEAK),
]
curves = {}
for name, setup, schedule in experiments:
    random.seed(SEED)
    model = new_model(setup)
    optimizer = training.adam(model)
    run = dict(pad_id=PAD, device=device)
    curves[name] = training.train_steps(model, optimizer, train_batches, schedule,
                                        EXPERIMENT_STEPS, **run)

fig, ax = plt.subplots(figsize=(6.4, 3.2))
colors = [COLORS["dec"], COLORS["muted"], COLORS["enc"]]
for (name, losses), color in zip(curves.items(), colors):
    smooth = pd.Series(losses).rolling(25, min_periods=1).mean()
    ax.plot(smooth, color=color, label=name)
ax.set_xlabel("step")
ax.set_ylabel("training loss (smoothed)")
ax.set_title("A high learning rate from the first step")
ax.legend(fontsize=9)
plt.show()

ABLATION_EPOCHS = 5 if not SMOKE else 1
ABLATION_PEAK, ABLATION_WARMUP = 1e-3, 1_000


def ablation_schedule(step):
    return ABLATION_PEAK * min(step / ABLATION_WARMUP, math.sqrt(ABLATION_WARMUP / step))


def train_model(setup, epochs, schedule, seed=SEED, checkpoint=None):
    random.seed(seed)
    model = new_model(setup, seed)
    optimizer = training.adam(model)
    history, done, best = [], 0, float("inf")
    start = time.perf_counter()
    for epoch in range(1, epochs + 1):
        losses = training.train_steps(model, optimizer, train_batches, schedule,
                                      len(train_batches), pad_id=PAD, first_step=done + 1,
                                      device=device)
        done += len(losses)
        loss = training.validation_loss(model, valid_batches, pad_id=PAD, device=device)
        minutes = (time.perf_counter() - start) / 60
        history.append({"epoch": epoch, "validation loss": loss, "minutes": minutes})
        if checkpoint is not None and loss < best:  # keep the best weights
            best = loss
            save_file(model.state_dict(), checkpoint)
    return model, history

RUNS = [(name, setup, SEED) for name, setup in VARIANTS.items()]
RUNS.append(("2017, another seed", VARIANTS["2017"], SEED + 1))

ladder_models, ladder_history = {}, {}
for name, setup, seed in RUNS:
    model, history = train_model(setup, ABLATION_EPOCHS, ablation_schedule, seed)
    ladder_models[name], ladder_history[name] = model, history
    last = history[-1]
    loss, minutes = round(last["validation loss"], 3), round(last["minutes"], 1)
    print(name, "| validation loss", loss, "|", minutes, "min")

fig, ax = plt.subplots(figsize=(6.4, 3.4))
colors = [COLORS["muted"], COLORS["data"], COLORS["dec"], COLORS["enc"], COLORS["att"],
          COLORS["line"]]
for (name, history), color in zip(ladder_history.items(), colors):
    epochs = [record["epoch"] for record in history]
    losses = [record["validation loss"] for record in history]
    style = (0, (3, 3)) if "seed" in name else "-"
    ax.plot(epochs, losses, marker="o", markersize=3, color=color, linestyle=style,
            label=name)
ax.set_xlabel("epoch")
ax.set_ylabel("validation loss")
ax.set_title("Six models, one budget")
ax.legend(fontsize=8)
plt.show()

comet = Comet(device)
# (source, translation, reference, score from unbabel-comet 2.2.7), as in lesson 4
PACKAGE_SCORES = [
    ("The researcher said she was exhausted.", "Исследователь сказал, что она устала.",
     "Исследовательница сказала, что вымоталась.", 0.873125),
    ("Tom forgot to lock the door.", "Tom forgot to lock the door.",
     "Том забыл запереть дверь.", 0.614086),
]
for source, translation, reference, theirs in PACKAGE_SCORES:
    ours = comet.score([source], [translation], [reference])[0]
    assert abs(ours - theirs) < 1e-4
print("tfs.comet gives the package's scores")

sample = test.sample(min(5_000, len(test)), random_state=SEED).reset_index(drop=True)
rows = []
for name, model in ladder_models.items():
    hypotheses = decoding.greedy(model, tokenizer, list(sample.en))
    sample[name] = hypotheses
    sample[f"COMET {name}"] = comet.score(sample.en, hypotheses, sample.ru)
    rows.append(
        {
            "model": name,
            "validation loss": ladder_history[name][-1]["validation loss"],
            "BLEU": sacrebleu.corpus_bleu(hypotheses, [list(sample.ru)]).score,
            "chrF": sacrebleu.corpus_chrf(hypotheses, [list(sample.ru)]).score,
            "COMET": sample[f"COMET {name}"].mean(),
        }
    )
ladder_scores = pd.DataFrame(rows)
table(ladder_scores, digits=3)

def bootstrap_against(name, baseline="2017", draws=1_000):
    rng = np.random.default_rng(SEED)
    a = sample[f"COMET {name}"].to_numpy()
    b = sample[f"COMET {baseline}"].to_numpy()
    differences = []
    for _ in range(draws):
        rows_drawn = rng.integers(0, len(sample), size=len(sample))  # with replacement
        differences.append(a[rows_drawn].mean() - b[rows_drawn].mean())
    return np.array(differences)


rows = []
for name in ladder_models:
    if name == "2017":
        continue
    differences = bootstrap_against(name)
    low, high = np.percentile(differences, [2.5, 97.5])
    rows.append({"model − 2017": name, "COMET difference": differences.mean(),
                 "95% of draws from": low, "to": high,
                 "better in": f"{(differences > 0).mean():.0%}"})
table(pd.DataFrame(rows), digits=4)

SEED_RUNS = "xmarva/tfs-seed-runs-en-ru"
path = hf_hub_download(SEED_RUNS, "data/runs.parquet", repo_type="dataset")
seed_runs = pd.read_parquet(path)
columns = ["setup", "seed", "validation loss", "COMET"]

ours = ladder_scores.set_index("model")
here = []  # this notebook's own runs: (setup, seed, name in the ladder)
here.append(("2017", SEED, "2017"))
here.append(("2017", SEED + 1, "2017, another seed"))
here.append(("+ RoPE", SEED, "+ RoPE"))
rows = []
for setup, seed, name in here:
    rows.append({"setup": setup, "seed": seed, "COMET": ours.loc[name, "COMET"],
                 "validation loss": ours.loc[name, "validation loss"]})
other_seeds = seed_runs[seed_runs.epochs == 5][columns]
short_runs = pd.concat([pd.DataFrame(rows), other_seeds], ignore_index=True)

designs = [("2017", COLORS["muted"]), ("+ RoPE", COLORS["att"])]
fig, axes = plt.subplots(1, 2, figsize=(6.4, 2.9))
for ax, metric in zip(axes, ["COMET", "validation loss"]):
    for k, (setup, color) in enumerate(designs):
        values = short_runs[short_runs.setup == setup][metric].to_numpy()
        offsets = np.linspace(-0.12, 0.12, len(values))  # close values stay visible
        ax.scatter(k + offsets, values, color=color, s=30, zorder=3)
    ax.set_xticks([0, 1], ["2017", "all four changes"])
    ax.set_xlim(-0.6, 1.6)
axes[0].set_title("COMET, five runs of each design")
axes[1].set_title("validation loss")
plt.show()

def design(setup):
    return "2017" if setup == "2017" else "all four changes"


rows = []
for setup, runs in short_runs.groupby("setup", sort=False):
    losses = runs["validation loss"]
    rows.append({"design": design(setup), "runs": len(runs),
                 "COMET from": runs.COMET.min(), "to": runs.COMET.max(),
                 "validation loss from": losses.min(), "to ": losses.max()})
table(pd.DataFrame(rows), digits=4)

def cache_bytes_per_token(layers, kv_heads, head_dim, bytes_per_number=2):
    return 2 * layers * kv_heads * head_dim * bytes_per_number  # a key and a value


head_dim = llama["hidden_size"] // llama["num_attention_heads"]
layers = llama["num_hidden_layers"]
with_gqa = cache_bytes_per_token(layers, llama["num_key_value_heads"], head_dim)
without_gqa = cache_bytes_per_token(layers, llama["num_attention_heads"], head_dim)
print(f"Llama 3.1 8B in bfloat16: {with_gqa / 1024:.0f} KiB per token")
print(f"with a key/value head per query head: {without_gqa / 1024:.0f} KiB per token")
print(f"a full context of 128K tokens: {with_gqa * 128 * 1024 / 2**30:.0f} GiB")

def decode_step(model, token, memory, src, caches, position):
    memory_mask = (src != model.pad_id)[:, None, None, :]
    x = model.embed(token, start=position)  # (B, 1, d_model): only the newest token
    angles = model.rope(1, start=position)
    for layer, cache in zip(model.decoder, caches):
        x = layer(x, None, angles, memory, memory_mask, cache)  # no mask: all keys are past
    if model.setup["norm"] != "post":
        x = model.decoder_norm(x)
    return (x @ model.embedding.weight.T)[:, -1]


@torch.no_grad()
def greedy_cached(model, sentences, max_tokens=60, batch_size=256):
    model.eval()
    order = sorted(range(len(sentences)), key=lambda i: len(sentences[i]))
    translations = [None] * len(sentences)
    for start in range(0, len(order), batch_size):
        rows_here = order[start : start + batch_size]
        ids = [tokenizer.encode(sentences[i]).ids for i in rows_here]
        src = training.pad_batch(ids, PAD).to(device)
        memory = model.encode(src)
        caches = []
        for layer in model.decoder:  # cross-attention: keys and values once per sentence
            k, v = layer.cross_attention.keys_values(memory)
            caches.append({"self": {}, "cross": {"fixed": True, "k": k, "v": v}})
        token = torch.full((len(rows_here), 1), BOS, device=device)
        written = []
        done = torch.zeros(len(rows_here), dtype=torch.bool, device=device)
        for position in range(max_tokens):
            scores = decode_step(model, token, memory, src, caches, position)
            next_ids = scores.argmax(dim=-1)
            next_ids[done] = PAD  # a finished sentence only collects padding
            written.append(next_ids)
            done |= next_ids == EOS
            if done.all():
                break
            token = next_ids[:, None]
        written = torch.stack(written, dim=1).tolist()
        for row, i in zip(written, rows_here):
            ids = []
            for token_id in row:
                if token_id in (EOS, PAD):
                    break
                ids.append(token_id)
            translations[i] = tokenizer.decode(ids)
    return translations

model = ladder_models["+ RoPE"]
sentences = list(dict.fromkeys(sample.en))

start = time.perf_counter()
plain = decoding.greedy(model, tokenizer, sentences)
plain_seconds = time.perf_counter() - start
start = time.perf_counter()
cached = greedy_cached(model, sentences)
cached_seconds = time.perf_counter() - start

same = sum(a == b for a, b in zip(plain, cached))
print(f"same translation: {same} of {len(sentences)}")
print(f"without a cache: {plain_seconds:.1f} s, with: {cached_seconds:.1f} s")

FINAL = VARIANTS["+ RoPE"]
FINAL_EPOCHS = 20 if not SMOKE else 1
CHECKPOINT = WORK_DIR / "llama-style.safetensors"

final, final_history = train_model(FINAL, FINAL_EPOCHS, training.paper_schedule,
                                   checkpoint=str(CHECKPOINT))
final.load_state_dict(load_file(CHECKPOINT, device=device))
final.eval()
best = min(record["validation loss"] for record in final_history)
lesson3_loss = training.validation_loss(lesson3, valid_batches, pad_id=PAD, device=device)
print("best validation loss:", round(best, 3))
print("lesson 3's model:", round(lesson3_loss, 3))

sentences = list(dict.fromkeys(test.en))
outputs = {
    "lesson 3": dict(zip(sentences, decoding.beam_search(lesson3, tokenizer, sentences))),
    "this lesson": dict(zip(sentences, decoding.beam_search(final, tokenizer, sentences))),
}
rows = []
for name, translations in outputs.items():
    test[name] = test.en.map(translations)
    test[f"COMET {name}"] = comet.score(test.en, test[name], test.ru)
    rows.append(
        {
            "model": name,
            "BLEU": sacrebleu.corpus_bleu(list(test[name]), [list(test.ru)]).score,
            "chrF": sacrebleu.corpus_chrf(list(test[name]), [list(test.ru)]).score,
            "COMET": test[f"COMET {name}"].mean(),
        }
    )
table(pd.DataFrame(rows), digits=3)

rng = np.random.default_rng(SEED)
a = test["COMET this lesson"].to_numpy()
b = test["COMET lesson 3"].to_numpy()
differences = []
for _ in range(1_000):
    rows_drawn = rng.integers(0, len(test), size=len(test))
    differences.append(a[rows_drawn].mean() - b[rows_drawn].mean())
low, high = np.percentile(differences, [2.5, 97.5])
print(f"COMET, this lesson − lesson 3: {np.mean(differences):+.4f}")
print(f"95% of the draws between {low:+.4f} and {high:+.4f}")

rows = [{"model": "lesson 3", "design": "2017", "validation loss": lesson3_loss,
         "COMET": test["COMET lesson 3"].mean()},
        {"model": "this lesson", "design": "all four changes", "validation loss": best,
         "COMET": test["COMET this lesson"].mean()}]
for _, run in seed_runs[seed_runs.epochs == 20].iterrows():
    rows.append({"model": f"seed {run.seed}", "design": design(run.setup),
                 "validation loss": run["validation loss"], "COMET": run.COMET})
table(pd.DataFrame(rows).sort_values(["design", "COMET"]), digits=4)

path = hf_hub_download(SEED_RUNS, "data/test.parquet", repo_type="dataset")
per_sentence = pd.read_parquet(path)
per_sentence = per_sentence.loc[test.index]  # the same pairs, in the order of our test set
assert (per_sentence.en.to_numpy() == test.en.to_numpy()).all()
a = test["COMET this lesson"].to_numpy()
for name in ["2017 s43, 20 epochs", "2017 s44, 20 epochs"]:
    b = per_sentence[f"COMET {name}"].to_numpy()
    rng = np.random.default_rng(SEED)
    differences = []
    for _ in range(1_000):
        rows_drawn = rng.integers(0, len(test), size=len(test))
        differences.append(a[rows_drawn].mean() - b[rows_drawn].mean())
    low, high = np.percentile(differences, [2.5, 97.5])
    seed = name.removesuffix(", 20 epochs")
    print(f"this lesson − {seed}: {np.mean(differences):+.4f}, "
          f"95% of the draws between {low:+.4f} and {high:+.4f}")