Mixture-of-experts LLMs from scratch with OLMoE

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 re
import subprocess
import sys
import sysconfig
import time
from collections import OrderedDict
from html import escape
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import pyarrow.parquet as pq
import torch
import torch.nn as nn
import torch.nn.functional as F
from IPython.display import HTML, display
from matplotlib.colors import LinearSegmentedColormap, to_rgb
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from transformers import AutoTokenizer

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"
GPUS = [f"cuda:{i}" for i in range(torch.cuda.device_count())] or ["cpu"]
HALF = torch.float16 if device == "cuda" else torch.float32  # 16-bit matmuls crawl on a CPU
assert device == "cuda" or SMOKE, "This lesson runs a 7B model: use two GPUs"
for name in GPUS:
    if name != "cpu":
        print(name, torch.cuda.get_device_name(name))

# repository, name, released, total / active parameters from the paper or the model card
MOE_MODELS = [
    ("mistralai/Mixtral-8x7B-v0.1", "Mixtral 8x7B", "2023-12", "47B / 13B"),
    ("Qwen/Qwen1.5-MoE-A2.7B", "Qwen1.5-MoE", "2024-03", "14.3B / 2.7B"),
    ("deepseek-ai/DeepSeek-V2-Lite", "DeepSeek-V2-Lite", "2024-05", "15.7B / 2.4B"),
    ("allenai/OLMoE-1B-7B-0924", "OLMoE", "2024-09", "6.9B / 1.3B"),
    ("ibm-granite/granite-3.1-3b-a800m-base", "Granite 3.1 3B", "2024-12", "3.3B / 0.8B"),
    ("deepseek-ai/DeepSeek-V3", "DeepSeek-V3", "2024-12", "671B / 37B"),
    ("Qwen/Qwen3-30B-A3B", "Qwen3-30B-A3B", "2025-04", "30.5B / 3.3B"),
    ("openai/gpt-oss-20b", "gpt-oss-20b", "2025-08", "20.9B / 3.6B"),
    ("openai/gpt-oss-120b", "gpt-oss-120b", "2025-08", "116.8B / 5.1B"),
]


def read_config(repo):
    with open(hf_hub_download(repo, "config.json")) as f:
        return json.load(f)


def first(c, *keys, default=0):
    for key in keys:
        if c.get(key) is not None:
            return c[key]
    return default


def router_order(c):
    kind = c["model_type"]
    if kind in ("granitemoe", "gpt_oss"):
        return "top-k, softmax"
    if c.get("scoring_func") == "sigmoid":
        return "sigmoid + bias, top-k, renormalise"
    renorm = kind == "mixtral" or c.get("norm_topk_prob")
    return "softmax, top-k" + (", renormalise" if renorm else "")


configs = {name: read_config(repo) for repo, name, _, _ in MOE_MODELS}
rows = []
for _, name, released, sizes in MOE_MODELS:
    c = configs[name]
    experts = first(c, "num_local_experts", "num_experts", "n_routed_experts")
    top_k = c["num_experts_per_tok"]
    shared = first(c, "n_shared_experts")
    if c.get("shared_expert_intermediate_size"):
        shared = c["shared_expert_intermediate_size"] // first(c, "moe_intermediate_size")
    rows.append(
        {
            "model": name,
            "released": released,
            "total / active": sizes,
            "layers": c["num_hidden_layers"],
            "dense layers": first(c, "first_k_dense_replace"),
            "experts": experts,
            "per token": top_k,
            "shared": shared,
            "expert width": first(c, "moe_intermediate_size", "intermediate_size"),
            "model width": c["hidden_size"],
            "router": router_order(c),
        }
    )
table(pd.DataFrame(rows))

def softmax_topk(s, k):  # OLMoE, Qwen1.5-MoE, DeepSeek-V2-Lite
    weights, chosen = s.softmax(dim=-1).topk(k, dim=-1)
    return weights, chosen


def softmax_topk_renorm(s, k):  # Mixtral, Qwen3
    weights, chosen = softmax_topk(s, k)
    return weights / weights.sum(dim=-1, keepdim=True), chosen


def topk_softmax(s, k):  # Granite, gpt-oss
    top, chosen = s.topk(k, dim=-1)
    return top.softmax(dim=-1), chosen


def sigmoid_bias(s, k, bias):  # DeepSeek-V3: the bias picks the experts, not their weights
    affinity = s.sigmoid()
    chosen = (affinity + bias).topk(k, dim=-1).indices
    weights = affinity.gather(-1, chosen)
    return weights / weights.sum(dim=-1, keepdim=True), chosen


torch.manual_seed(0)
s = torch.randn(1000, 64) * 2  # scores of 1,000 tokens for 64 experts
results = {
    "softmax, top-k": softmax_topk(s, 8),
    "softmax, top-k, renormalise": softmax_topk_renorm(s, 8),
    "top-k, softmax": topk_softmax(s, 8),
    "sigmoid + bias, top-k, renormalise": sigmoid_bias(s, 8, torch.zeros(64)),
}
reference_w, reference_i = results["top-k, softmax"]
rows = []
for name, (w, i) in results.items():
    rows.append(
        {
            "router": name,
            "same experts as top-k, softmax": bool(torch.equal(i, reference_i)),
            "largest weight difference": f"{(w - reference_w).abs().max().item():.2e}",
            "weights sum to": f"{w.sum(dim=-1).mean().item():.3f}",
        }
    )
table(pd.DataFrame(rows))

class Experts(nn.Module):
    """n SwiGLU networks kept as two stacked weights, one row of experts per tensor."""

    def __init__(self, n, width, hidden):
        super().__init__()
        self.gate_up = nn.Parameter(
            torch.empty(n, 2 * hidden, width)
        )  # gate rows, then up rows
        self.down = nn.Parameter(torch.empty(n, width, hidden))

    def run(self, e, x):
        gate, up = (x @ self.gate_up[e].T).chunk(2, dim=-1)
        return (F.silu(gate) * up) @ self.down[e].T


def dispatch(experts, x, weights, chosen):
    """Run each expert once, on the tokens that chose it; add the results back, weighted."""
    k = chosen.shape[-1]
    pairs = chosen.flatten()
    order = pairs.argsort(stable=True)
    tokens = order // k  # the token each sorted pair came from
    w = weights.flatten()[order]
    counts = pairs.bincount(minlength=experts.gate_up.shape[0]).tolist()
    out = torch.zeros_like(x)
    start = 0
    for e, n in enumerate(counts):
        if n:
            rows = tokens[start : start + n]
            out.index_add_(
                0, rows, (experts.run(e, x[rows]) * w[start : start + n, None]).to(x.dtype)
            )
        start += n
    return out


def dense_run(experts, x, weights, chosen):
    """The same layer computed the wasteful way: every expert on every token."""
    n = experts.gate_up.shape[0]
    gate = torch.zeros(x.shape[0], n, dtype=weights.dtype).scatter(1, chosen, weights)
    outputs = torch.stack(
        [experts.run(e, x) for e in range(n)], dim=1
    )  # tokens x experts x width
    return (outputs * gate[..., None]).sum(dim=1)


torch.manual_seed(0)
toy = Experts(64, 256, 128)
nn.init.normal_(toy.gate_up, std=0.05)
nn.init.normal_(toy.down, std=0.05)
x = torch.randn(512, 256)
weights, chosen = softmax_topk(x @ torch.randn(256, 64) * 0.1, 8)
timings = {}
for name, fn in [
    ("every expert on every token", dense_run),
    ("grouped by expert", dispatch),
]:
    start = time.time()
    with torch.no_grad():
        y = fn(toy, x, weights, chosen)
    timings[name] = (time.time() - start, y)
difference = (
    (timings["grouped by expert"][1] - timings["every expert on every token"][1])
    .abs()
    .max()
)
for name, (seconds, _) in timings.items():
    print(f"{name:28} {seconds * 1000:7.1f} ms")
print("max difference:", f"{difference.item():.2e}")

REPO = "allenai/OLMoE-1B-7B-0924"
olmoe_config = configs["OLMoE"]
olmoe_config.setdefault("rms_norm_eps", 1e-5)  # not in the config: the library default
keys = [
    "hidden_size",
    "num_hidden_layers",
    "num_attention_heads",
    "num_key_value_heads",
    "num_experts",
    "num_experts_per_tok",
    "intermediate_size",
    "vocab_size",
    "norm_topk_prob",
    "tie_word_embeddings",
]
for key in keys:
    print(f"{key:22} {olmoe_config[key]}")
print(f"{'rope_theta':22} {olmoe_config.get('rope_theta')}")

class RMSNorm(nn.Module):
    def __init__(self, dim, eps):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

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


def rope_tables(positions, head_dim, theta):
    exponents = torch.arange(0, head_dim, 2, device=positions.device).float() / head_dim
    angles = positions.float()[..., None] / theta**exponents
    angles = torch.cat([angles, angles], dim=-1)
    return angles.cos()[:, None], angles.sin()[:, None]


def rotate(x, cos, sin):
    half = x.shape[-1] // 2
    rotated = torch.cat([-x[..., half:], x[..., :half]], dim=-1)
    return x * cos.to(x.dtype) + rotated * sin.to(x.dtype)


class Attention(nn.Module):
    def __init__(self, c):
        super().__init__()
        d = c["hidden_size"]
        self.heads = c["num_attention_heads"]
        self.head_dim = d // self.heads
        self.q_proj = nn.Linear(d, d, bias=False)
        self.k_proj = nn.Linear(d, d, bias=False)
        self.v_proj = nn.Linear(d, d, bias=False)
        self.o_proj = nn.Linear(d, d, bias=False)
        self.q_norm = RMSNorm(
            d, c["rms_norm_eps"]
        )  # over the whole projection, before the heads
        self.k_norm = RMSNorm(d, c["rms_norm_eps"])

    def forward(self, x, cos, sin, mask):
        b, t, _ = x.shape
        shape = (b, t, self.heads, self.head_dim)
        q = rotate(self.q_norm(self.q_proj(x)).view(shape).transpose(1, 2), cos, sin)
        k = rotate(self.k_norm(self.k_proj(x)).view(shape).transpose(1, 2), cos, sin)
        v = self.v_proj(x).view(shape).transpose(1, 2)
        out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
        return self.o_proj(out.transpose(1, 2).reshape(b, t, -1))


class MoE(nn.Module):
    def __init__(self, c):
        super().__init__()
        d, n = c["hidden_size"], c["num_experts"]
        self.k = c["num_experts_per_tok"]
        self.gate = nn.Linear(d, n, bias=False)  # the router
        self.experts = Experts(n, d, c["intermediate_size"])
        self.keep = None  # section 5: a mask of the experts a pruned model still has
        self.trace = None  # section 4: a list that collects every token's chosen experts

    def forward(self, x):
        b, t, d = x.shape
        flat = x.reshape(-1, d)
        scores = self.gate(flat).float()
        if self.keep is not None:
            scores = scores.masked_fill(~self.keep, float("-inf"))
        weights, chosen = softmax_topk(scores, self.k)  # no renormalisation in OLMoE
        if self.trace is not None:
            self.trace.append(chosen.to(torch.uint8).cpu())
        return dispatch(self.experts, flat, weights, chosen).view(b, t, d)


class Block(nn.Module):
    def __init__(self, c):
        super().__init__()
        self.input_layernorm = RMSNorm(c["hidden_size"], c["rms_norm_eps"])
        self.self_attn = Attention(c)
        self.post_attention_layernorm = RMSNorm(c["hidden_size"], c["rms_norm_eps"])
        self.mlp = MoE(c)

    def forward(self, x, cos, sin, mask):  # x: the residual stream, in 32-bit floats
        x = x + self.self_attn(self.input_layernorm(x), cos, sin, mask).float()
        return x + self.mlp(self.post_attention_layernorm(x)).float()


class OLMoE(nn.Module):
    def __init__(self, c):
        super().__init__()
        self.config = c
        self.embed_tokens = nn.Embedding(c["vocab_size"], c["hidden_size"])
        self.layers = nn.ModuleList([Block(c) for _ in range(c["num_hidden_layers"])])
        self.norm = RMSNorm(c["hidden_size"], c["rms_norm_eps"])
        self.lm_head = nn.Linear(c["hidden_size"], c["vocab_size"], bias=False)

    def hidden(self, ids, mask):
        c = self.config
        positions = torch.arange(ids.shape[1], device=ids.device)[None]
        x = self.embed_tokens(ids).float()
        tables = {}
        for layer in self.layers:
            where = layer.input_layernorm.weight.device  # layers may sit on different GPUs
            if where not in tables:
                theta = c.get("rope_theta") or 10000.0
                tables[where] = rope_tables(
                    positions.to(where), c["hidden_size"] // c["num_attention_heads"], theta
                )
            x = layer(x.to(where), *tables[where], mask.to(where))
        return self.norm(x.to(self.norm.weight.device))

    def logits(self, h):
        return self.lm_head(h.to(self.lm_head.weight.device)).float()

    def forward(self, ids, mask):
        return self.logits(self.hidden(ids, mask))

def placement(c, gpus):
    n = c["num_hidden_layers"]
    return {i: gpus[i * len(gpus) // n] for i in range(n)}


def checkpoint_files(source):
    if Path(source).is_dir():  # a model saved to disk (the smoke run builds a small one)
        return sorted(str(p) for p in Path(source).glob("*.safetensors"))
    try:
        with open(hf_hub_download(source, "model.safetensors.index.json")) as f:
            names = sorted(set(json.load(f)["weight_map"].values()))
    except Exception:
        names = ["model.safetensors"]
    return [hf_hub_download(source, name) for name in names]


def load_olmoe(source, c, dtype, gpus):
    where = placement(c, gpus)
    with torch.device("meta"):
        model = OLMoE(c)

    def target(key):
        match = re.match(r"layers\.(\d+)\.", key)
        return (
            where[int(match.group(1))]
            if match
            else (gpus[0] if "embed" in key else gpus[-1])
        )

    state, pending = {}, {}
    n = c["num_experts"]
    for path in checkpoint_files(source):
        for key, value in load_file(path).items():
            key = key.removeprefix("model.")
            expert = re.match(
                r"(layers\.\d+\.mlp\.experts)\.(\d+)\.(gate|up|down)_proj\.weight", key
            )
            if expert:  # one expert's matrix: wait for all of them, then stack
                prefix, e, part = expert.group(1), int(expert.group(2)), expert.group(3)
                pending.setdefault(prefix, {})[(e, part)] = value
                parts = pending[prefix]
                if len(parts) == 3 * n:
                    device = target(prefix)
                    gate_up = torch.stack(
                        [
                            torch.cat([parts[(e, "gate")], parts[(e, "up")]])
                            for e in range(n)
                        ]
                    )
                    down = torch.stack([parts[(e, "down")] for e in range(n)])
                    state[prefix + ".gate_up"] = gate_up.to(device, dtype)
                    state[prefix + ".down"] = down.to(device, dtype)
                    del pending[prefix]
                continue
            key = key.replace("experts.gate_up_proj", "experts.gate_up")  # already stacked
            key = key.replace("experts.down_proj", "experts.down")
            state[key] = value.to(target(key), dtype)
    model.load_state_dict(state, strict=True, assign=True)
    return model.eval()

tok = AutoTokenizer.from_pretrained(REPO)
text = (
    "The Rhine rises in the Swiss Alps, flows through Lake Constance and along "
    "the French–German border, and reaches the North Sea through a delta in the "
    "Netherlands."
)
check_ids = tok(text, return_tensors="pt").input_ids
torch.save(check_ids, WORK_DIR / "check_ids.pt")

if SMOKE:  # a small OLMoE with random weights stands in for the real one
    from transformers import OlmoeConfig, OlmoeForCausalLM

    small = OlmoeConfig(
        vocab_size=olmoe_config["vocab_size"],
        hidden_size=64,
        intermediate_size=32,
        num_hidden_layers=2,
        num_attention_heads=4,
        num_key_value_heads=4,
        num_experts=16,
        num_experts_per_tok=4,
        eos_token_id=0,
        pad_token_id=1,
    )
    torch.manual_seed(0)
    reference = OlmoeForCausalLM(small).eval()
    SOURCE = WORK_DIR / "olmoe-small"
    reference.save_pretrained(SOURCE)
    with torch.no_grad():
        theirs = reference(check_ids).logits.float()
    olmoe_config = json.loads((SOURCE / "config.json").read_text())
    olmoe_config.setdefault("rope_theta", 10000.0)
else:
    SOURCE = REPO
    script = f"""
import torch
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(
    "{REPO}", dtype=torch.float16, device_map="auto")
ids = torch.load("{WORK_DIR / "check_ids.pt"}").to("cuda:0")
with torch.no_grad():
    torch.save(model(ids).logits.float().cpu(), "{WORK_DIR / "theirs.pt"}")
"""
    subprocess.run([sys.executable, "-c", script], check=True)
    theirs = torch.load(WORK_DIR / "theirs.pt")

model = load_olmoe(SOURCE, olmoe_config, HALF, GPUS)
t = check_ids.shape[1]
causal = torch.ones(t, t, dtype=torch.bool).tril()
with torch.no_grad():
    ours = model(check_ids.to(GPUS[0]), causal).cpu()
print("tokens:", t)
print("largest logit:", f"{theirs.abs().max().item():.2f}")
print("max difference:", f"{(ours - theirs).abs().max().item():.2e}")
print(
    "same next-token prediction everywhere:",
    bool((ours.argmax(-1) == theirs.argmax(-1)).all()),
)

parts = {"embeddings and output": 0, "attention": 0, "routers": 0, "experts": 0, "norms": 0}
for name, p in model.named_parameters():
    if "experts" in name:
        parts["experts"] += p.numel()
    elif "mlp.gate" in name:
        parts["routers"] += p.numel()
    elif "self_attn" in name and "norm" not in name:
        parts["attention"] += p.numel()
    elif "embed" in name or "lm_head" in name:
        parts["embeddings and output"] += p.numel()
    else:
        parts["norms"] += p.numel()
total = sum(parts.values())
c = olmoe_config
per_expert = 3 * c["hidden_size"] * c["intermediate_size"]
active = (
    total
    - parts["experts"]
    + c["num_hidden_layers"] * c["num_experts_per_tok"] * per_expert
)
table(
    pd.DataFrame(
        {
            "part": list(parts),
            "parameters": list(parts.values()),
            "share": [f"{n / total:.1%}" for n in parts.values()],
        }
    )
)
print(f"used for one token: {active:,} of {total:,} ({active / total:.0%})")

LENGTH = 64 if SMOKE else 512
CALIBRATION, HELD_OUT = (2, 1) if SMOKE else (48, 16)


def windows(texts, n):
    ids = []
    for text in texts:
        ids.extend(tok(text, add_special_tokens=False).input_ids + [tok.eos_token_id])
        if len(ids) >= n * LENGTH:
            break
    ids = ids[: n * LENGTH]
    return torch.tensor(ids).view(n, LENGTH)


def english_texts():
    name = "wikitext-103-raw-v1/test-00000-of-00001.parquet"
    path = hf_hub_download("Salesforce/wikitext", name, repo_type="dataset")
    return ["".join(pd.read_parquet(path)["text"].tolist())]


def python_texts():
    root = Path(sysconfig.get_paths()["stdlib"])
    modules = [
        "argparse",
        "difflib",
        "fractions",
        "statistics",
        "textwrap",
        "heapq",
        "dataclasses",
        "pathlib",
        "calendar",
        "ipaddress",
        "csv",
        "shlex",
        "string",
        "gettext",
        "tokenize",
    ]
    files = [root / f"{m}.py" for m in modules if (root / f"{m}.py").exists()]
    files += [p for p in sorted(root.glob("json/*.py")) + sorted(root.glob("email/*.py"))]
    return [f.read_text(errors="ignore") for f in files]


def maths_texts():
    path = hf_hub_download(
        "openai/gsm8k", "main/test-00000-of-00001.parquet", repo_type="dataset"
    )
    frame = pd.read_parquet(path)
    # the solutions carry calculator annotations, <<48/2=24>>, that no reader sees
    return [
        q + "\n" + re.sub(r"<<[^>]*>>", "", a) + "\n"
        for q, a in zip(frame.question, frame.answer)
    ]


def russian_texts():
    path = hf_hub_download(
        "wikimedia/wikipedia",
        "20231101.ru/train-00002-of-00021.parquet",
        repo_type="dataset",
    )
    texts = (
        pq.ParquetFile(path).read_row_group(0, columns=["text"]).column("text").to_pylist()
    )
    return [t for t in texts if len(t) > 2000]  # whole articles, not stubs


corpora = {
    "English": english_texts(),
    "Python": python_texts(),
    "maths": maths_texts(),
    "Russian": russian_texts(),
}
DOMAINS = list(corpora)
data = {}
for name, texts in corpora.items():
    all_ids = windows(texts, CALIBRATION + HELD_OUT)
    data[name] = {"calibration": all_ids[:CALIBRATION], "held out": all_ids[CALIBRATION:]}
words = {
    "English": "The river flows through the old town and into the sea.",
    "Russian": "Река течёт через старый город и впадает в море.",
}
for language, sentence in words.items():
    print(
        f"{language:8} {len(sentence.split())} words -> "
        f"{len(tok(sentence).input_ids)} tokens"
    )
print({name: tuple(d["calibration"].shape) for name, d in data.items()})

BATCH = 2 if SMOKE else 8
n_layers, n_experts = olmoe_config["num_hidden_layers"], olmoe_config["num_experts"]


def causal_mask(ids):
    t = ids.shape[1]
    return torch.ones(t, t, dtype=torch.bool).tril()


def route(batches):
    """Every token's chosen experts in every layer: tokens x layers x k."""
    for layer in model.layers:
        layer.mlp.trace = []
    with torch.no_grad():
        for i in range(0, len(batches), BATCH):
            chunk = batches[i : i + BATCH]
            model.hidden(chunk.to(GPUS[0]), causal_mask(chunk))
    traces = torch.stack([torch.cat(layer.mlp.trace) for layer in model.layers], dim=1)
    for layer in model.layers:
        layer.mlp.trace = None
    return traces


start = time.time()
traces = {name: route(d["calibration"]) for name, d in data.items()}
print(
    f"routed {sum(len(t) for t in traces.values()):,} tokens in {time.time() - start:.0f} s"
)

# usage[d, l, e]: the share of domain d's tokens that chose expert e in layer l
usage = np.zeros((len(DOMAINS), n_layers, n_experts))
for d, name in enumerate(DOMAINS):
    t = traces[name].long().numpy()
    for layer in range(n_layers):
        usage[d, layer] = np.bincount(t[:, layer].flatten(), minlength=n_experts) / len(t)

preference = usage / usage.sum(axis=0, keepdims=True).clip(
    min=1e-9
)  # share per domain, per expert
shades = LinearSegmentedColormap.from_list("share", [COLORS["grid"], COLORS["att"]])
shown_layers = [0, n_layers // 2, n_layers - 1]
fig, axes = plt.subplots(
    len(shown_layers), 1, figsize=(10, 1.15 * len(shown_layers) * len(DOMAINS) / 2.4)
)
for ax, layer in zip(np.atleast_1d(axes), shown_layers):
    p = preference[:, layer]
    order = np.lexsort((-p.max(axis=0), p.argmax(axis=0)))
    ax.imshow(p[:, order], aspect="auto", cmap=shades, vmin=0, vmax=1)
    ax.set_yticks(range(len(DOMAINS)), DOMAINS)
    ax.set_xticks([])
    ax.set_title(
        f"layer {layer}: one column per expert, its use split by kind of text",
        fontsize=10,
        loc="left",
    )
plt.tight_layout()
plt.show()

threshold = 0.5
specialists = np.stack(
    [
        (preference.argmax(axis=0) == d) & (preference.max(axis=0) > threshold)
        for d in range(len(DOMAINS))
    ]
).sum(axis=2)  # domains x layers
fig, ax = plt.subplots(figsize=(9, 3.4))
bottom = np.zeros(n_layers)
for d, name in enumerate(DOMAINS):
    ax.bar(
        range(n_layers),
        specialists[d],
        bottom=bottom,
        label=name,
        color=[COLORS["enc"], COLORS["att"], COLORS["data"], COLORS["dec"]][d],
    )
    bottom += specialists[d]
ax.set_xlabel("layer")
ax.set_ylabel(f"experts of {n_experts}")
ax.set_title("Experts that get more than half of their use from one kind of text")
ax.legend(frameon=False, ncol=4, loc="upper center", bbox_to_anchor=(0.5, -0.2))
plt.show()

def footprint(u, share):
    """How many experts, taken from the most used down, carry `share` of the routing."""
    ordered = np.sort(u)[::-1]
    return int(np.searchsorted(np.cumsum(ordered) / ordered.sum(), share) + 1)


rows = []
for d, name in enumerate(DOMAINS):
    for share in [0.5, 0.9]:
        counts = [footprint(usage[d, layer], share) for layer in range(n_layers)]
        rows.append(
            {
                "text": name,
                "share of routing": f"{share:.0%}",
                "experts needed, fewest layer": min(counts),
                "median layer": int(np.median(counts)),
                "most layer": max(counts),
            }
        )
table(pd.DataFrame(rows))

samples = {
    "English": "The museum opens at nine and closes at six on weekdays.",
    "Python": "for line in open(path):\n    if line.startswith('#'):\n        continue",
    "maths": "She buys 3 boxes of 12 eggs and uses 7, so 36 - 7 = 29 are left.",
    "Russian": "Музей открывается в девять утра и закрывается в шесть вечера.",
}
palette = [COLORS["enc"], COLORS["att"], COLORS["data"], COLORS["dec"]]


def chip(piece, color):
    r, g, b = (int(255 * v) for v in to_rgb(color))
    return (
        f'<span style="background:rgba({r},{g},{b},0.32);border-radius:4px;padding:1px 3px;'
        f'margin:1px;font-family:monospace;white-space:pre">{escape(piece)}</span>'
    )


lines = ['<div style="line-height:2.1">']
for name, sentence in samples.items():
    ids = tok(sentence, return_tensors="pt").input_ids
    trace = route(ids)
    pieces = [tok.decode([i]).replace("\n", "↵") for i in ids[0].tolist()]
    lines.append(f"<div><b>{name}</b></div>")
    for layer in [0, n_layers - 1]:
        chips = []
        for piece, chosen_experts in zip(pieces, trace[:, layer]):
            liked = int(preference[:, layer, int(chosen_experts[0])].argmax())
            chips.append(chip(piece, palette[liked]))
        lines.append(
            f'<div><span style="opacity:0.6">layer {layer}&nbsp;</span>'
            f'{"".join(chips)}</div>'
        )
legend = "".join(chip(d, palette[i]) for i, d in enumerate(DOMAINS))
caption = '<span style="opacity:0.6">the first expert prefers&nbsp;</span>'
lines.append(f"<div>{caption}{legend}</div></div>")
display(HTML("".join(lines)))

def loss_on(batches):
    total, count = 0.0, 0
    with torch.no_grad():
        for i in range(0, len(batches), BATCH):
            ids = batches[i : i + BATCH].to(GPUS[0])
            logits = model(ids, causal_mask(ids))
            target = ids[:, 1:].to(logits.device)
            total += F.cross_entropy(
                logits[:, :-1].flatten(0, 1), target.flatten(), reduction="sum"
            ).item()
            count += target.numel()
    return total / count


def prune_for(domain, keep):
    d = DOMAINS.index(domain)
    for layer, block in enumerate(model.layers):
        mask = torch.zeros(n_experts, dtype=torch.bool)
        mask[np.argsort(-usage[d, layer])[:keep]] = True
        block.mlp.keep = mask.to(block.mlp.gate.weight.device)


def unprune():
    for block in model.layers:
        block.mlp.keep = None


KEEP = [n_experts, n_experts // 4 * 3, n_experts // 2, n_experts // 4]
TARGETS = ["Python", "Russian"]
results = []
start = time.time()
unprune()
full = {name: loss_on(data[name]["held out"]) for name in DOMAINS}
for target in TARGETS:
    for keep in KEEP[1:]:
        prune_for(target, keep)
        for name in DOMAINS:
            results.append(
                {
                    "pruned for": target,
                    "experts kept": keep,
                    "text": name,
                    "loss": loss_on(data[name]["held out"]),
                }
            )
unprune()
for target in TARGETS:
    for name in DOMAINS:
        results.append(
            {
                "pruned for": target,
                "experts kept": n_experts,
                "text": name,
                "loss": full[name],
            }
        )
results = pd.DataFrame(results)
print(f"{len(results)} evaluations in {time.time() - start:.0f} s")
table(
    results.pivot_table(
        index=["pruned for", "text"], columns="experts kept", values="loss"
    ).reset_index(),
    digits=3,
)

def gigabytes(keep):
    rest = total - parts["experts"]
    return (rest + parts["experts"] * keep / n_experts) * 2 / 1e9  # 16-bit weights


fig, axes = plt.subplots(1, len(TARGETS), figsize=(10, 3.6), sharey=True)
for ax, target in zip(axes, TARGETS):
    chunk = results[results["pruned for"] == target]
    for d, name in enumerate(DOMAINS):
        line = chunk[chunk.text == name].sort_values("experts kept")
        ax.plot(
            line["experts kept"],
            line["loss"],
            marker="o",
            color=palette[d],
            label=name,
            linewidth=2.4 if name == target else 1.2,
        )
    ax.set_xticks(KEEP, [f"{k}\n{gigabytes(k):.1f} GB" for k in KEEP], fontsize=8)
    ax.set_xlabel("experts kept per layer, and the model's size")
    ax.set_title(f"pruned for {target}", fontsize=10)
axes[0].set_ylabel("loss on held-out text")
axes[0].legend(frameon=False, fontsize=8.5)
plt.show()

mixed = torch.cat(
    [traces[name] for name in DOMAINS]
)  # every routed token, of all four kinds
rng = np.random.default_rng(0)
sizes = [1, 2, 4, 8, 16, 32, 64, 128]
touched_real, touched_one = [], []
for b in sizes:
    picks = [mixed[rng.choice(len(mixed), b, replace=False)].long() for _ in range(50)]
    touched_real.append(
        np.mean(
            [
                [len(set(p[:, layer].flatten().tolist())) for layer in range(n_layers)]
                for p in picks
            ]
        )
    )
    one = traces["Python"]
    picks = [one[rng.choice(len(one), b, replace=False)].long() for _ in range(50)]
    touched_one.append(
        np.mean(
            [
                [len(set(p[:, layer].flatten().tolist())) for layer in range(n_layers)]
                for p in picks
            ]
        )
    )
k = olmoe_config["num_experts_per_tok"]
uniform = [n_experts * (1 - (1 - k / n_experts) ** b) for b in sizes]
fig, ax = plt.subplots(figsize=(7.2, 3.6))
ax.plot(
    sizes,
    uniform,
    color=COLORS["line"],
    linestyle="--",
    label="uniform routing, the formula",
)
ax.plot(
    sizes,
    touched_real,
    marker="o",
    color=COLORS["enc"],
    label="OLMoE, tokens of all four kinds",
)
ax.plot(
    sizes, touched_one, marker="o", color=COLORS["att"], label="OLMoE, Python tokens only"
)
ax.set_xscale("log", base=2)
ax.set_xticks(sizes, [str(s) for s in sizes])
ax.set_xlabel("tokens in the batch")
ax.set_ylabel("experts touched per layer")
ax.set_title(f"How many of the {n_experts} experts a batch needs")
ax.legend(frameon=False, fontsize=8.5)
plt.show()

def lru_hit_rate(trace, capacity):
    hits = total = 0
    for layer in range(trace.shape[1]):
        cache = OrderedDict()
        for chosen_experts in trace[:, layer].tolist():
            for e in chosen_experts:
                total += 1
                if e in cache:
                    hits += 1
                    cache.move_to_end(e)
                else:
                    cache[e] = True
                    if len(cache) > capacity:
                        cache.popitem(last=False)
    return hits / total


capacities = [8, 16, 24, 32, 48]
sequence = {
    name: traces[name][: 4 * LENGTH] for name in DOMAINS
}  # four windows in reading order
random_trace = torch.stack(
    [
        torch.stack([torch.randperm(n_experts)[:k] for _ in range(n_layers)])
        for _ in range(4 * LENGTH)
    ]
)
fig, ax = plt.subplots(figsize=(7.2, 3.6))
for d, name in enumerate(DOMAINS):
    ax.plot(
        capacities,
        [lru_hit_rate(sequence[name], c) for c in capacities],
        marker="o",
        color=palette[d],
        label=name,
    )
ax.plot(
    capacities,
    [lru_hit_rate(random_trace, c) for c in capacities],
    linestyle="--",
    color=COLORS["line"],
    label="uniform routing",
)
ax.set_xlabel(f"experts kept in GPU memory per layer, of {n_experts}")
ax.set_ylabel("share of expert uses already in memory")
ax.set_title("A least-recently-used cache of experts, one token at a time")
ax.set_ylim(0, 1)
ax.legend(frameon=False, fontsize=8.5)
plt.show()