"""Split one big model over several GPUs and count the time one word takes.

Python standard library only. A toy: one made-up GPU, two links and three models, with every
number below taken from a public file (read 2026-10-19) or computed from those.

Made-up GPU (80 GiB of memory, 2,039 GB/s to read it, 312 TFLOP/s): the memory size is the
80 GiB card of lesson 18 and the other two are the A100 80 GB SXM column of the NVIDIA A100
data sheet (June 2021 revision) that lessons 19 and 20 use. The same sheet lists the two
links: NVLink at 600 GB/s and PCIe Gen4 at 64 GB/s. The sheet gives one number per link and
does not say whether it counts both directions; the toy uses it as written.

Start-up times: NCCL (NVIDIA's library that GPUs use to swap data) keeps its default tuning
constants in src/tuning/cost_model.cc of github.com/NVIDIA/nccl (master, last change to that
file 2026-08-31, commit d551aef). For its ring all-reduce with the small-message protocol it
lists a base latency of 6.6 microseconds and, per step, 0.6 microseconds on NVLink and 1.0 on
PCIe. These are the library's own planning figures, not measurements.

Model shapes: the config.json files of the Llama 3.1 8B and 70B models on Hugging Face, and
Table 3 of the Llama 3 paper (arXiv 2407.21783) for the 405B, as saved next to this file.
Every size is in bytes and every time in seconds unless the name says otherwise.
"""
import json
from math import ceil
from pathlib import Path

HERE = Path(__file__).parent
GIB = 1024**3

MODELS = {  # key -> (name on the page, saved file)
    "8b": ("8B", "llama31_8b_config.json"),
    "70b": ("70B", "llama31_70b_config.json"),
    "405b": ("405B", "llama31_405b_shape.json"),
}
COUNTS = (1, 2, 4, 8, 16)   # GPUs the page offers; the A100 data sheet lists servers of up to 16, and between servers the link is slower
SPLITS = ("slices", "groups")   # tensor parallelism (slices of every layer), pipeline parallelism (groups of layers)
WORDS = (1, 16, 64)         # words moving through the GPUs together (one user, or many users at once)

GPU_GIB = 80                # lesson 18: the 80 GiB card
GPU_BANDWIDTH = 2.039e12    # bytes per second the GPU reads its own memory (A100 80 GB SXM data sheet, Jun21)
GPU_FLOPS = 312e12          # 16-bit multiply-adds per second, dense (same sheet)
USABLE = 0.92               # lesson 18: share of a card a server claims by default
USABLE_BYTES = int(GPU_GIB * GIB * USABLE)

LINKS = {  # key -> (label, bytes per second, NCCL step latency in microseconds)
    "fast": ("Fast link", 600e9, 0.6),   # NVLink on the data sheet; NCCL NVLINK ring LL step
    "slow": ("Slow link", 64e9, 1.0),    # PCIe Gen4 on the data sheet; NCCL PCI ring LL step
}
BASE_US = 6.6               # NCCL ring LL base latency, microseconds, paid once per all-reduce
ALL_REDUCES_PER_LAYER = 2   # Megatron-LM (arXiv 1909.08053): after the attention block and after the feed-forward block
WORD_BYTES = 2              # a number in a word's hidden vector takes 16 bits
CHAT_TOKENS = 8192          # the chat whose attention cache must fit (lesson 18)
ACTIVATION_TOKENS = 2048    # lesson 18's working-memory estimate
OVERHEAD_GIB = 1            # lesson 18's allowance
CACHE_VALUE_BYTES = 2       # lesson 18: the attention cache keeps 16-bit numbers


def load_config(key: str) -> dict:
    return json.loads((HERE / MODELS[key][1]).read_text())


def head_dim(cfg: dict) -> int:
    return cfg.get("head_dim") or cfg["hidden_size"] // cfg["num_attention_heads"]


def param_count(cfg: dict) -> int:
    """Weights of a Llama-style model counted from its shape (lesson 18's formula)."""
    h, ffn = cfg["hidden_size"], cfg["intermediate_size"]
    kv_width = cfg["num_key_value_heads"] * head_dim(cfg)
    per_layer = 2 * h * h + 2 * h * kv_width + 3 * h * ffn + 2 * h
    tables = cfg["vocab_size"] * h * (1 if cfg.get("tie_word_embeddings") else 2)
    return cfg["num_hidden_layers"] * per_layer + tables + h


def weight_bytes(cfg: dict) -> int:
    """All weights at 16 bits: two bytes each."""
    return 2 * param_count(cfg)


def kv_bytes_per_token(cfg: dict) -> int:
    return 2 * cfg["num_hidden_layers"] * cfg["num_key_value_heads"] * head_dim(cfg) * CACHE_VALUE_BYTES


def activation_bytes(cfg: dict) -> int:
    return ACTIVATION_TOKENS * (3 * cfg["intermediate_size"] + 4 * cfg["hidden_size"]) * 2


def message_bytes(cfg: dict) -> int:
    """What the GPUs swap after each split step for one word: that word's hidden vector."""
    return cfg["hidden_size"] * WORD_BYTES


def ceil_div(a: int, b: int) -> int:
    return -(-a // b)


def memory_per_gpu(key: str, n: int) -> dict:
    """Bytes one of n GPUs holds. Weights and the chat's cache are shared out evenly (the cache by
    key/value head, so never into more parts than the model has heads); working memory is not shared."""
    cfg = load_config(key)
    parts = {
        "weights": ceil_div(weight_bytes(cfg), n),
        "cache": ceil_div(kv_bytes_per_token(cfg) * CHAT_TOKENS, min(n, cfg["num_key_value_heads"])),
        "working": activation_bytes(cfg) + OVERHEAD_GIB * GIB,
    }
    parts["total"] = sum(parts.values())
    return parts


def fits(key: str, n: int) -> bool:
    return memory_per_gpu(key, n)["total"] <= USABLE_BYTES


def read_time(key: str, n: int) -> float:
    """Each GPU reads only its 1/n of the weights, so one word's read takes 1/n of the time."""
    return ceil_div(weight_bytes(load_config(key)), n) / GPU_BANDWIDTH


def compute_time(key: str, n: int, words: int = 1) -> float:
    """Two operations per weight for each word, shared over n GPUs."""
    return 2 * param_count(load_config(key)) * words / (n * GPU_FLOPS)


def allreduce_time(size: int, n: int, link: str) -> float:
    """Ring all-reduce of `size` bytes over n GPUs: 2(n-1) steps, each pays the step latency, and each
    GPU sends about 2(n-1)/n times `size` in all (the NCCL tests documentation derives this factor)."""
    if n == 1:
        return 0.0
    _, bandwidth, step_us = LINKS[link]
    startup = (BASE_US + 2 * (n - 1) * step_us) * 1e-6
    return startup + 2 * (n - 1) / n * size / bandwidth


def handoff_time(size: int, link: str) -> float:
    """One message from a group of layers to the next: base latency, one step, and the bytes."""
    _, bandwidth, step_us = LINKS[link]
    return (BASE_US + step_us) * 1e-6 + size / bandwidth


def word_time(key: str, n: int, link: str, split: str = "slices", words: int = 1) -> dict:
    """One step that gives every one of `words` words its next word. Reading and computing overlap,
    so the longer of the two counts; the transfers come after (the toy does not overlap them with
    the reading). A user waits one step per word, so this is also one user's time per word."""
    cfg = load_config(key)
    read, compute = read_time(key, n), compute_time(key, n, words)
    size = message_bytes(cfg) * words
    if split == "slices":
        transfer = ALL_REDUCES_PER_LAYER * cfg["num_hidden_layers"] * allreduce_time(size, n, link)
        busy = max(read, compute)
    elif split == "groups":
        # the groups work one after another, so n groups each read 1/n of the weights in turn
        transfer = (n - 1) * handoff_time(size, link)
        busy = n * max(read, compute)
    else:
        raise ValueError(f"split must be one of {SPLITS}")
    total = busy + transfer
    return {"read": busy, "transfer": transfer, "total": total, "share": transfer / total, "words_per_s": 1 / total}


def best_n(key: str, link: str, words: int = 1, max_n: int = 64, divide_heads: bool = False) -> int:
    """The GPU count (any whole number up to max_n, among those that fit) with the shortest word time.
    With divide_heads, only counts that divide the model's attention heads: slices are cut by head."""
    heads = load_config(key)["num_attention_heads"]
    options = [n for n in range(1, max_n + 1) if fits(key, n) and (not divide_heads or heads % n == 0)]
    return min(options, key=lambda n: word_time(key, n, link, words=words)["total"])


def least_n(key: str) -> int:
    """The fewest GPUs (any whole number) that hold the model with room for working memory."""
    return next(n for n in range(1, 65) if fits(key, n))


def settings():
    for key in MODELS:
        for n in COUNTS:
            for link in LINKS:
                for split in SPLITS:
                    for words in WORDS:
                        yield key, n, link, split, words