On this page
A large language model (LLM) with 70 billion weights, the numbers it has learned, takes 131.42 GiB at 16 bits each (a bit is one 0 or 1; a GiB is a little over a billion bytes). One made-up GPU, a graphics chip of the kind that runs AI models, can use 73.6 GiB (92% of its 80 GiB). The model cannot sit on one GPU. Splitting it over several GPUs fixes that and lets each GPU read less, so words come sooner; but twice the GPUs does not give twice the speed.
In one minute
A big LLM needs several GPUs because its weights do not fit in one GPU’s memory, so each GPU holds a share. Think of a recipe book too big for one cook, torn into parts, one part for each cook. All the cooks read their own pages at the same moment, so the dish comes sooner. After every step, though, the cooks must call out their partial results to one another, and calling takes time. More cooks mean fewer pages each but more calling. A faster link between the GPUs shortens the wait.
Explore the numbers
Open the widget. Choose a model size, a number of GPUs and a link, the connection between the GPUs. It shows the memory each GPU needs, the time for one word (a word here is one token, a chunk of text) and how much of that time is waiting. It opens on the 70B model (70 billion weights) on 4 GPUs with the fast link: 34.93 GiB each, 18.94 ms a word (a ms is a thousandth of a second) and 8.7% of that spent waiting. Choose the 405B model on 1 GPU and it says “Does not fit”.
How it works
The weights must fit
To write a word, a GPU reads every weight it holds from its own memory, so the weights must sit there (How much memory does an AI model need?). Fewer bits for each weight is one way to make a model fit (What does quantization trade away?). The other is to share the weights among several GPUs.
Figure 1 sets the 70B model’s 131.42 GiB against the 73.6 GiB one GPU can use. Split four ways, each GPU holds 32.85 GiB. The widget adds a chat’s notes (the attention cache) and working space, which brings it to 34.93 GiB. The 405B model takes 755.96 GiB and needs at least 11 GPUs.
Figure 1. The 70B model’s weights against one GPU’s usable memory, then split four ways.
Slicing every layer
A model is a stack of layers, steps the numbers pass through; the 70B model has 80. The usual split cuts every layer into slices, one for each GPU. This is tensor parallelism. The GPUs work on the same layer at the same moment, and each reads only its share of the weights.
Reading is the slow part, because a chip waits on memory (Why a fast AI chip spends most of its time waiting). One GPU would need 69.20 ms to read the 70B model’s weights, if they fit. Four GPUs each read a quarter, in 17.30 ms, at the same time.
Waiting for each other
A slice gives only part of the answer for a layer. The GPUs must add their parts, and every GPU needs the total. That swap is called an all-reduce. It happens twice in each layer, so 160 times for every word of the 70B model. Each message is small, 16,384 bytes, but a swap has a start-up time whatever the size: about 10 microseconds (a millionth of a second) on 4 GPUs. The message goes round a ring of GPUs, so more GPUs mean more steps and a longer start-up. The 160 swaps cost 1.64 ms.
Figure 2 shows reading shrinking and waiting growing. At 2, 4, 8 and 16 GPUs the waiting takes 3.5%, 8.7%, 21.8% and 47.7% of each word.
Figure 2. One word from the 70B model on the fast link: reading (green), waiting (hatched).
When more GPUs stop helping
Doubling the GPUs halves the reading but not the waiting. The 8B model (8 billion weights, small enough for one GPU) shows where it ends. On the fast link, one user’s words come 1.77, 1.69 and 1.35 times faster with each doubling from 1 GPU to 8, and 0.94 times for the last doubling, to 16: slower. Figure 3 shows words a second.
Figure 3. The 8B model, one user. Among the counts shown (1 to 16), speed peaks at 8 GPUs on both links.
The link decides how long the waiting lasts. For one user the messages are tiny, so the link matters little: the 70B model on 8 GPUs takes 11.06 ms a word on the fast link and 12.02 ms on the slow one. A server that handles many chats makes the messages bigger. With 64 words in each swap, the same split takes 11.54 ms on the fast link and 16.53 ms on the slow one.
Another way to split: groups of layers
The other split gives each GPU a group of whole layers: with four GPUs, the first 20 layers go to the first one, and so on. This is pipeline parallelism. Only one short message passes from group to group, so waiting nearly vanishes: 69.23 ms a word on four GPUs, against 69.20 ms of reading alone. But the groups work one after another. One user’s word takes as long as it would on one big GPU, and three of the four GPUs stand idle at any moment. Groups make a model fit; they do not speed up one chat. They pay off when many chats keep every group busy (a later lesson covers batching).
What this means when you build
- Count the fewest GPUs first. Weights, the chat’s notes and working space must fit in each GPU’s usable memory. The 70B model needs 2 GPUs, the 405B model 11.
- Try fewer bits before more GPUs. A smaller file may need fewer GPUs, at a cost in accuracy (lesson 20).
- Expect less than double. Each doubling of GPUs gives less than twice the speed. Watch the share of time spent waiting.
- Ask what links the GPUs. For one user the link barely shows; for many chats at once it matters.
- Keep layer groups for fitting, not speed. They do not shorten one user’s wait for a word.
Optionalabout 22 min
Go deeper
Optional. For readers who write Python. The code is in the code/ folder; every number below was printed by python3 show_numbers.py, worked out by hand, or comes from a cited source read on 2026-10-19.
Where the toy’s numbers come from
The toy has one made-up GPU, two links and three models. The GPU has the 80 GiB of lesson 18, of which a server claims 92%: 73.6 GiB. It reads its memory at 2,039 GB/s and multiplies at 312 TFLOP/s, the figures of the A100 80 GB SXM column in the NVIDIA A100 data sheet (June 2021) that lessons 19 and 20 use. The same sheet lists NVLink at 600 GB/s, which the toy calls the fast link, and PCIe Gen4 at 64 GB/s, the slow link. It gives one figure for each link and does not say whether the figure counts both directions. NVIDIA’s blog post “NVIDIA Ampere Architecture In-Depth” (Krashinsky et al., 14 May 2020) gives 25 GB/s in each direction for one NVLink and 600 GB/s in all, so the 600 probably adds the two directions. The toy uses the figures as written. If both figures count both directions, halving them changes little for one user: the 70B model on 8 GPUs goes from 11.06 to 11.07 ms on the fast link and from 12.02 to 12.09 ms on the slow one. With 64 words at once the slow link goes from 16.53 to 21.12 ms, 28% more.
The start-up times come from NCCL, NVIDIA’s library for GPU-to-GPU communication. Its file src/tuning/cost_model.cc (master branch, last changed 2026-08-31) lists the library’s default tuning constants: a base latency of 6.6 microseconds for a ring all-reduce with the small-message protocol, plus 0.6 microseconds for each step on NVLink and 1.0 on PCIe. These are planning estimates, not measurements. The model shapes are the config.json files of Llama 3.1 8B and 70B and Table 3 of the Llama 3 paper for the 405B, saved in code/; the weight counts match lesson 18.
The memory sum, by hand
The 70B model on 4 GPUs, in bytes. Weights: 70,553,706,496 x 2 / 4 = 35,276,853,248. The chat’s cache, for 8,192 tokens as in lesson 18: 2 x 80 layers x 8 key/value heads x 128 values x 2 bytes = 327,680 bytes a token, 2,684,354,560 for the chat, shared by key/value head: 671,088,640. Working space is not shared: 486,539,264 for 2,048 tokens of intermediate results plus a 1 GiB allowance, 1,560,281,088 in all. The sum is 37,508,222,976 bytes, or 34.93 GiB, under 73.6. The cache can be shared only by head, so 16 GPUs hold as much of it as 8 (the model has 8 key/value heads). This is memory_per_gpu, with its docstring left out:
def memory_per_gpu(key: str, n: int) -> dict:
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 partsThe fewest GPUs that fit are 1 for the 8B model, 2 for the 70B and 11 for the 405B. The widget offers 16 so the 405B can fit at all: 49.60 GiB each.
The swap, by hand
For the 70B model on 4 GPUs a message is one word’s hidden vector, 8,192 numbers of 2 bytes: 16,384 bytes. A ring all-reduce over n GPUs takes 2(n - 1) steps, 6 here, and each GPU sends 2(n - 1) / n times the message, 1.5 x 16,384 = 24,576 bytes (the NCCL tests documentation gives the time as (S / B) x 2(n - 1) / n, for size S and bandwidth B). On the fast link the start-up is 6.6 + 6 x 0.6 = 10.2 microseconds and sending takes 24,576 / 600 GB/s = 0.041 microseconds, so a swap takes 10.24. Times 160 swaps is 1.64 ms. Start-up is 99.6% of it, so for one user the link’s step latency matters and its speed does not. With 64 words in each swap on 8 GPUs, a swap takes 18.06 microseconds on the fast link and 49.27 on the slow one.
One user, 8B model
Each cell is milliseconds a word:
| GPUs | Fast link | Slow link |
|---|---|---|
| 1 | 7.88 | 7.88 |
| 2 | 4.44 | 4.50 |
| 4 | 2.62 | 2.79 |
| 8 | 1.95 | 2.32 |
| 16 | 2.07 | 2.85 |
Waiting takes 11.3%, 24.9%, 49.4% and 76.2% of a word on the fast link at 2, 4, 8 and 16 GPUs, and 12.4%, 29.4%, 57.5% and 82.7% on the slow link.
Check one cell, 2 GPUs on the fast link. Each GPU reads 8,030,261,248 bytes at 2,039 GB/s: 3.94 ms. A message is 4,096 x 2 = 8,192 bytes. A swap over 2 GPUs takes 2 steps, so its start-up is 6.6 + 2 x 0.6 = 7.8 microseconds and it sends 8,192 bytes, 0.014 microseconds: 7.81 in all. The 8B model has 32 layers, so a word makes 64 swaps, 0.50 ms. The total is 3.94 + 0.50 = 4.44 ms, and waiting is 11.3% of it. Past 8 GPUs the waiting (1.58 ms on the fast link) outgrows the reading (0.49 ms). This sum accepts any whole number of GPUs and ignores the head rule under “What the toy leaves out”. The fastest count for one user is 10 on the fast link and 8 on the slow one for the 8B model, 19 and 15 for the 70B model, 36 and 28 for the 405B. Among the counts the page offers, 8 is best for the 8B model and 16 for the others on both links.
Many words at once, 70B on 8 GPUs
Each cell is milliseconds for one step that gives every word its next word:
| Words | Fast link | Slow link |
|---|---|---|
| 1 | 11.06 | 12.02 |
| 16 | 11.17 | 13.09 |
| 64 | 11.54 | 16.53 |
Waiting takes 21.8%, 22.6% and 25.0% of the step on the fast link at 1, 16 and 64 words, and 28.0%, 33.9% and 47.7% on the slow link. The slow link costs 0.96 ms more at 1 word and 4.99 ms more at 64. Calculation takes two operations for each weight and word. For 64 words on 8 GPUs that is 3.62 ms, still under the 8.65 ms of reading, so the toy counts the reading as the busy time, and for one word on one GPU calculation takes 0.45 ms against 69.20 ms of reading. Real servers also overlap the waiting with other work, which the toy does not.
Groups of layers
Figure 4. One word from a 70B model on four GPUs. Slices read together; groups read in turn.
With groups the toy has each group read in turn: 4 x 17.30 = 69.20 ms, plus 3 hand-offs of 7.2 microseconds and a few bytes each, 69.23 ms in all, 0.03% of it waiting. Only one GPU in four works at any moment. In the toy the four groups run one after another, so a word takes about as long as one GPU reading the whole model, which is why groups do nothing for one chat.
The vLLM documentation (read 2026-10-19) gives the same order of advice: if the model fits on one GPU, distributed inference is probably not needed; if it needs several GPUs in one machine, use tensor parallelism; across machines, combine tensor parallelism inside each machine with pipeline parallelism between them. Its tuning page also names pipeline parallelism for very deep and narrow models, where distributing layers is more efficient than tensor sharding.
What the toy leaves out
- Real systems overlap communication with calculation, so some waiting hides. The toy adds the two.
- A GPU also reads its share of the chat’s cache for every word. The toy reads weights only.
- NCCL picks its method by message size. The toy always uses the small-message ring.
- The A100 data sheet lists servers of up to 16 GPUs on the 600 GB/s link. A model spread over several servers meets a slower link between them. The toy keeps one link for every count, so large counts may look better here than across servers.
- Serving software usually needs the number of GPUs to divide the number of attention heads, the 32, 64 and 128 the three models have, so 19 GPUs is not a usable count. Among the counts that do divide them, the fastest are 8 for the 8B model, 16 for the 70B and 32 for the 405B, on both links. Each model has 8 key/value heads, so past 8 GPUs each GPU keeps a copy of one, which the toy’s cache sum allows for.
- Start-up times vary with load. The constants are planning estimates.
- Training also splits the data across GPUs. This lesson covers only a user waiting for the next word.
Proof it works
PYTHONDONTWRITEBYTECODE=1 python3 -m pytest -q -p no:cacheprovider code/The tests check these claims. The three weight totals match the published ones, and 70B on 4 GPUs adds up by hand to 37,508,222,976 bytes. The fewest GPUs that fit are 1, 2 and 11. Reading time times the GPU count stays constant. A ring all-reduce simulated for 2, 4, 8 and 16 GPUs ends with every GPU holding the sum in 2(n - 1) steps and sends the bytes the formula charges. A word makes 64, 160 and 252 swaps. Waiting grows with the GPU count and with a slower link, and the 8B model is slower on 16 GPUs than on 8 on both links. The best counts are 10 and 8, 19 and 15, and 36 and 28; among the counts that divide the heads they are 8, 16 and 32. Layer groups stay within 2% of one GPU’s read. The slow link’s extra cost is more than 5 times larger at 64 words than at 1. Every number this page works out appears in the code’s output; the rest come from the sources named.
The widget repeats the arithmetic in JavaScript, so its tests keep it honest. Its data block must equal what embed_widget.py writes; every element its script looks up must exist; it must stay under 40 KiB and make no network call. A parity test runs the widget’s number code under Node and compares it with Python for all 180 settings.
What the tests do not show: that a real server’s waiting follows the planning constants, or that the data sheet’s link figures hold in a given machine. To check, time one all-reduce of 16,384 bytes with the NCCL tests on your own hardware.
Where it is used
- Megatron-LM (Shoeybi et al., arXiv 1909.08053, 2019) splits each layer’s attention by heads and its feed-forward block by columns and rows. Its forward path has two all-reduces in each layer, after the attention block and after the feed-forward block, the two per layer this lesson counts. The paper reports 77% of linear scaling when the model grew with the number of GPUs, an 8.3-billion-parameter model on 8-way model parallelism against a 1.2-billion-parameter model on one GPU, in training, not in serving.
- The NCCL tests (NVIDIA, performance notes) give the all-reduce time used above.
- vLLM (documentation above) offers the tensor and pipeline settings the lesson describes.
Skills this lesson builds
- Add up the memory one GPU needs for its share of the weights, the chat’s cache and working space.
- Count the swaps a word needs from a model’s layers, and the bytes in each.
- Estimate a swap’s time from its start-up and its bytes, and read the waiting share from it.
- Explain why a split stops paying off, and what a faster link changes.
- Tell tensor parallelism from pipeline parallelism, and say what each does for one user.
Portfolio project: write a script that takes a model shape, a GPU’s memory and a link’s figures, and prints the fewest GPUs, the memory on each, and the time and waiting share for one word at 1, 2, 4, 8 and 16 GPUs. Add a chart and a test that checks one setting by hand, then compare the waiting with a timed all-reduce on hardware you own.
Check yourself
- Why can the 70B model not run on one GPU with 80 GiB of memory?
Answer
Its weights take 131.42 GiB at 16 bits, and one GPU can use 73.6 GiB. With two GPUs each holds about 65.71 GiB of weights, and with the chat’s cache and working space 68.41 GiB, which fits.
- Four GPUs each read a quarter of the 70B model in 17.30 ms, yet a word takes 18.94 ms. Where does the rest go?
Answer
To waiting. The GPUs swap partial results twice in each of the 80 layers, 160 swaps for each word, and each swap costs about 10 microseconds of start-up. That adds 1.64 ms.
- The 405B model has 126 layers. How many swaps does one word need, and how large is each message when a word’s hidden vector has 16,384 numbers of 2 bytes?
Answer
Two swaps in each layer: 252 for each word. Each message is 16,384 x 2 = 32,768 bytes.
- A teammate says: “Going from 8 to 16 GPUs doubles the speed for each user.” What is wrong with that?
Answer
Doubling the GPUs halves the reading, but the waiting grows. In the toy the 70B model gets 1.34 times faster, from 11.06 to 8.27 ms a word on the fast link, and the 8B model gets slower, from 514 to 483 words a second.
- When do groups of layers help, and when do they not?
Answer
They help a model fit across GPUs with little waiting, and they pay off when many chats keep every group busy. They do not shorten one user’s wait: the 70B model takes 69.23 ms a word on four groups, against 69.20 ms of reading alone, with three GPUs in four idle at any moment.
Further reading
Gentle
- Parallelism and Scaling (vLLM documentation, read 2026-10-19)
- NVIDIA A100 Tensor Core GPU data sheet (NVIDIA; memory, bandwidth and link figures)
Deep
- Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism (Shoeybi et al.)
- The Llama 3 Herd of Models (Meta; Table 3 holds the model shapes) and the model pages for Llama 3.1 8B and Llama 3.1 70B (Meta; login needed for the files)
- NVIDIA Ampere Architecture In-Depth (Krashinsky et al., NVIDIA, 14 May 2020; NVLink figures)
- Performance notes for the NCCL tests (NVIDIA; read 2026-10-19)
- NCCL cost model source (NVIDIA; planning constants, read 2026-10-19)
- Optimization and tuning (vLLM documentation; when to use pipeline parallelism, read 2026-10-19)
Optional
Code
Each file runs on its own. Download it, or read it here.
Show code
"""Keep the numbers and model shapes inside ../widget.html equal to gpu_split.py.
python3 embed_widget.py rewrite the data block in widget.html
python3 embed_widget.py --check exit 1 if widget.html differs (test_widget.py runs the same check)
"""
import json
import re
import sys
from pathlib import Path
import gpu_split as g
WIDGET = Path(__file__).resolve().parent.parent / "widget.html"
DATA_BLOCK = re.compile(r'(<script type="application/json" id="gpu-data">)(.*?)(</script>)', re.S)
def widget_data() -> dict:
models = []
for key, (name, _) in g.MODELS.items():
cfg = g.load_config(key)
models.append({"key": key, "name": name, "layers": cfg["num_hidden_layers"], "hidden": cfg["hidden_size"],
"ffn": cfg["intermediate_size"], "heads": cfg["num_attention_heads"], "kv_heads": cfg["num_key_value_heads"],
"head_dim": cfg.get("head_dim"), "vocab": cfg["vocab_size"], "tied": bool(cfg.get("tie_word_embeddings"))})
return {"models": models, "counts": list(g.COUNTS), "words": list(g.WORDS), "splits": list(g.SPLITS),
"links": [{"key": k, "label": label, "bandwidth": bandwidth, "step_us": step_us} for k, (label, bandwidth, step_us) in g.LINKS.items()],
"gpu": {"gib": g.GPU_GIB, "bandwidth": g.GPU_BANDWIDTH, "flops": g.GPU_FLOPS, "usable": g.USABLE},
"base_us": g.BASE_US, "all_reduces": g.ALL_REDUCES_PER_LAYER, "word_bytes": g.WORD_BYTES, "chat_tokens": g.CHAT_TOKENS,
"act_tokens": g.ACTIVATION_TOKENS, "overhead_gib": g.OVERHEAD_GIB, "cache_value_bytes": g.CACHE_VALUE_BYTES}
def embed(html: str) -> str:
"""html with the data block rewritten from gpu_split.py."""
block = json.dumps(widget_data(), indent=1)
return DATA_BLOCK.sub(lambda m: m.group(1) + "\n" + block + "\n" + m.group(3), html, count=1)
if __name__ == "__main__":
current = WIDGET.read_text(encoding="utf-8")
updated = embed(current)
if "--check" in sys.argv:
sys.exit(0 if updated == current else 1)
WIDGET.write_text(updated, encoding="utf-8")
print("widget.html data block now matches gpu_split.py")Show code
"""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, wordsShow code
{
"_source": "Llama 3 paper (arXiv 2407.21783), Table 3: 126 layers, model dimension 16,384, FFN dimension 53,248, 128 attention heads, 8 key/value heads. The vocabulary size (128,256) and the untied output layer are the same as the 8B and 70B files saved next to this one. The gated 405B config.json is not saved here.",
"hidden_size": 16384,
"intermediate_size": 53248,
"max_position_embeddings": 131072,
"model_type": "llama",
"num_attention_heads": 128,
"num_hidden_layers": 126,
"num_key_value_heads": 8,
"tie_word_embeddings": false,
"torch_dtype": "bfloat16",
"vocab_size": 128256
}Show code
{
"architectures": [
"LlamaForCausalLM"
],
"attention_bias": false,
"attention_dropout": 0.0,
"bos_token_id": 128000,
"eos_token_id": 128001,
"hidden_act": "silu",
"hidden_size": 8192,
"initializer_range": 0.02,
"intermediate_size": 28672,
"max_position_embeddings": 131072,
"mlp_bias": false,
"model_type": "llama",
"num_attention_heads": 64,
"num_hidden_layers": 80,
"num_key_value_heads": 8,
"pretraining_tp": 1,
"rms_norm_eps": 1e-05,
"rope_scaling": {
"factor": 8.0,
"low_freq_factor": 1.0,
"high_freq_factor": 4.0,
"original_max_position_embeddings": 8192,
"rope_type": "llama3"
},
"rope_theta": 500000.0,
"tie_word_embeddings": false,
"torch_dtype": "bfloat16",
"transformers_version": "4.43.0.dev0",
"use_cache": true,
"vocab_size": 128256
}Show code
{
"architectures": [
"LlamaForCausalLM"
],
"attention_bias": false,
"attention_dropout": 0.0,
"bos_token_id": 128000,
"eos_token_id": 128001,
"hidden_act": "silu",
"hidden_size": 4096,
"initializer_range": 0.02,
"intermediate_size": 14336,
"max_position_embeddings": 131072,
"mlp_bias": false,
"model_type": "llama",
"num_attention_heads": 32,
"num_hidden_layers": 32,
"num_key_value_heads": 8,
"pretraining_tp": 1,
"rms_norm_eps": 1e-05,
"rope_scaling": {
"factor": 8.0,
"low_freq_factor": 1.0,
"high_freq_factor": 4.0,
"original_max_position_embeddings": 8192,
"rope_type": "llama3"
},
"rope_theta": 500000.0,
"tie_word_embeddings": false,
"torch_dtype": "bfloat16",
"transformers_version": "4.43.0.dev0",
"use_cache": true,
"vocab_size": 128256
}Show code
"""Print every number the lesson quotes. Run: PYTHONDONTWRITEBYTECODE=1 python3 show_numbers.py"""
from decimal import ROUND_HALF_UP, Decimal
import gpu_split as g
def half_up(x: float, places: str = "0.01") -> str:
"""Round the way the widget's JavaScript toFixed does on an exact tie (0.625 -> 0.63); Python's format rounds ties to even."""
return str(Decimal(x).quantize(Decimal(places), rounding=ROUND_HALF_UP))
def gib(b: int) -> str:
return half_up(b / g.GIB)
def ms(s: float) -> str:
return half_up(s * 1e3)
def halved_link_times(key: str, n: int, words: int) -> dict:
"""Word time in seconds on each link if the data sheet's link rates count both directions, so one direction runs at half."""
saved = dict(g.LINKS)
try:
for link, (label, bandwidth, step_us) in saved.items():
g.LINKS[link] = (label, bandwidth / 2, step_us)
return {link: g.word_time(key, n, link, words=words)["total"] for link in g.LINKS}
finally:
g.LINKS.update(saved)
def main() -> None:
print(f"GPU: {g.GPU_GIB} GiB, {g.USABLE_BYTES / g.GIB:.1f} GiB usable, {g.GPU_BANDWIDTH / 1e9:,.0f} GB/s, {g.GPU_FLOPS / 1e12:.0f} TFLOP/s")
for key, (label, bandwidth, step_us) in g.LINKS.items():
print(f"{label}: {bandwidth / 1e9:.0f} GB/s, {step_us} us a step, base {g.BASE_US} us")
print("\n== Models (16 bits)")
for key, (name, _) in g.MODELS.items():
cfg = g.load_config(key)
print(f"{name}: {g.param_count(cfg):,} weights, {gib(g.weight_bytes(cfg))} GiB, {cfg['num_hidden_layers']} layers, "
f"hidden {cfg['hidden_size']}, message {g.message_bytes(cfg):,} bytes, "
f"{g.ALL_REDUCES_PER_LAYER * cfg['num_hidden_layers']} swaps a word, fewest GPUs {g.least_n(key)}")
print("\n== Memory per GPU in GiB: weights, cache, working, total, fits")
for key, (name, _) in g.MODELS.items():
for n in g.COUNTS:
m = g.memory_per_gpu(key, n)
print(f"{name} x{n}: {gib(m['weights'])} + {gib(m['cache'])} + {gib(m['working'])} = {gib(m['total'])} fits={g.fits(key, n)}")
print("\n== One word, one user, GPUs slice every layer: read ms | wait ms | total ms | wait share | words/s")
for key, (name, _) in g.MODELS.items():
for n in g.COUNTS:
for link, (label, _, _) in g.LINKS.items():
t = g.word_time(key, n, link)
print(f"{name} x{n} {label}: {ms(t['read'])} | {ms(t['transfer'])} | {ms(t['total'])} | {t['share'] * 100:.1f}% | {t['words_per_s']:.1f}")
print("\n== Best whole-number GPU count (any count that fits, up to 64), words at once 1/16/64")
for key, (name, _) in g.MODELS.items():
print(name, {w: {link: g.best_n(key, link, w) for link in g.LINKS} for w in g.WORDS})
print("\n== Best GPU count that divides the model's attention heads (up to 64), one user")
for key, (name, _) in g.MODELS.items():
print(name, f"{g.load_config(key)['num_attention_heads']} heads", {link: g.best_n(key, link, divide_heads=True) for link in g.LINKS})
print("\n== 8B: speed-up of each doubling, by link")
for link, (label, _, _) in g.LINKS.items():
times = [g.word_time("8b", n, link)["total"] for n in g.COUNTS]
print(label, [f"x{a / b:.2f}" for a, b in zip(times, times[1:])])
print("\n== 70B on 8 GPUs, words at once: total ms | wait share")
for words in g.WORDS:
row = [f"{label} {ms(g.word_time('70b', 8, link, words=words)['total'])} | {g.word_time('70b', 8, link, words=words)['share'] * 100:.1f}%"
for link, (label, _, _) in g.LINKS.items()]
print(words, *row)
print("\n== Layer groups (pipeline), one user: total ms | wait share")
for key, (name, _) in g.MODELS.items():
for n in g.COUNTS:
t = g.word_time(key, n, "fast", "groups")
u = g.word_time(key, n, "slow", "groups")
print(f"{name} x{n}: fast {ms(t['total'])} | {t['share'] * 100:.2f}% slow {ms(u['total'])} | {u['share'] * 100:.2f}% one GPU busy at a time: {100 / n:.1f}% of GPUs")
print("\n== If the sheet's link rates count both directions (halve them), 70B on 8 GPUs, total ms: as written -> halved")
for words in (1, 64):
halved = halved_link_times("70b", 8, words)
print(words, {link: f"{ms(g.word_time('70b', 8, link, words=words)['total'])} -> {ms(halved[link])}"
f" ({(halved[link] / g.word_time('70b', 8, link, words=words)['total'] - 1) * 100:.0f}% more)" for link in g.LINKS})
if __name__ == "__main__":
main()Show code
"""Guards for gpu_split.py: memory per GPU, read time, the all-reduce, the best GPU count and the numbers the lesson quotes."""
import ast
import re
from pathlib import Path
import pytest
import gpu_split as g
from show_numbers import gib, half_up, halved_link_times, ms
LESSON = Path(__file__).resolve().parent.parent
PUBLISHED = {"8b": 8_030_261_248, "70b": 70_553_706_496, "405b": 405_853_388_800} # same totals lesson 18 pins
def test_weight_counts_match_the_published_totals():
"""Guards a wrong shape file or counting formula: every size on the page starts from these totals."""
assert {k: g.param_count(g.load_config(k)) for k in g.MODELS} == PUBLISHED
def test_memory_per_gpu_by_hand():
"""Guards the memory sum: weights shared, cache shared by kv head, working memory not shared (70B on 4 GPUs, added up by hand)."""
m = g.memory_per_gpu("70b", 4)
assert m["weights"] == 2 * PUBLISHED["70b"] // 4 == 35_276_853_248
assert m["cache"] == 2 * 80 * 8 * 128 * 2 * 8192 // 4 == 671_088_640
assert m["working"] == 2048 * (3 * 28672 + 4 * 8192) * 2 + g.GIB == 486_539_264 + g.GIB
assert m["total"] == m["weights"] + m["cache"] + m["working"]
# the cache is shared by key/value head, so 16 GPUs hold as much of it as 8 do
assert g.memory_per_gpu("70b", 16)["cache"] == g.memory_per_gpu("70b", 8)["cache"]
def test_which_settings_fit():
"""Guards a wrong 'fits' verdict, the first thing the page shows: 70B needs 2 GPUs, 405B needs 11 and the page offers 16."""
assert [g.fits("8b", n) for n in g.COUNTS] == [True] * 5
assert [g.fits("70b", n) for n in g.COUNTS] == [False, True, True, True, True]
assert [g.fits("405b", n) for n in g.COUNTS] == [False, False, False, False, True]
assert {k: g.least_n(k) for k in g.MODELS} == {"8b": 1, "70b": 2, "405b": 11}
def test_read_time_falls_in_step_with_the_gpu_count():
"""Guards a split that does not shorten the read: each GPU reads 1/n of the weights, so n x the read time stays one GPU's read time."""
assert g.read_time("8b", 1) == 2 * PUBLISHED["8b"] / g.GPU_BANDWIDTH
for key in g.MODELS:
for n in g.COUNTS:
assert g.read_time(key, n) * n == pytest.approx(g.read_time(key, 1), rel=1e-6)
for key, n, link, split, words in g.settings(): # the page calls the busy part 'reading': it must be the longer of the two
assert g.compute_time(key, n, words) < g.read_time(key, n)
def ring_all_reduce(vectors: list[list[int]]):
"""Sum the vectors the way a ring does: n-1 steps of reduce-scatter, then n-1 steps of all-gather. Returns (result per GPU, steps, items sent per GPU)."""
n, size = len(vectors), len(vectors[0])
chunk = size // n
data = [list(v) for v in vectors]
sent = [0] * n
steps = 0
for phase in ("reduce", "gather"):
for step in range(n - 1):
messages = []
for gpu in range(n):
c = (gpu - step) % n if phase == "reduce" else (gpu + 1 - step) % n
messages.append(((gpu + 1) % n, c, data[gpu][c * chunk:(c + 1) * chunk]))
sent[gpu] += chunk
for dest, c, part in messages:
for i, x in enumerate(part):
data[dest][c * chunk + i] = data[dest][c * chunk + i] + x if phase == "reduce" else x
steps += 1
return data, steps, sent
@pytest.mark.parametrize("n", [2, 4, 8, 16])
def test_all_reduce_sums_and_costs_two_n_minus_one_steps(n):
"""Guards the all-reduce story: every GPU ends with the sum of all the partial results, in 2(n-1) steps and 2(n-1)/n x size sent."""
size = n * 3
vectors = [[gpu * 100 + i for i in range(size)] for gpu in range(n)]
result, steps, sent = ring_all_reduce(vectors)
assert all(row == [sum(v[i] for v in vectors) for i in range(size)] for row in result)
assert steps == 2 * (n - 1)
assert set(sent) == {2 * (n - 1) * size // n}
# the time formula charges exactly those steps and those bytes
link = g.LINKS["fast"]
expected = (g.BASE_US + steps * link[2]) * 1e-6 + sent[0] / link[1]
assert g.allreduce_time(size, n, "fast") == pytest.approx(expected, rel=1e-12)
def test_one_gpu_waits_on_no_one():
assert g.allreduce_time(8192, 1, "slow") == 0 and g.word_time("8b", 1, "slow")["transfer"] == 0
def test_a_word_makes_two_swaps_per_layer():
"""Guards the swap count the page states (160 for 70B) and that the wait is that many all-reduces of one hidden vector."""
cfg = g.load_config("70b")
t = g.word_time("70b", 4, "fast")
assert t["transfer"] == pytest.approx(2 * 80 * g.allreduce_time(8192 * 2, 4, "fast"), rel=1e-12)
assert {k: 2 * g.load_config(k)["num_hidden_layers"] for k in g.MODELS} == {"8b": 64, "70b": 160, "405b": 252}
assert g.message_bytes(cfg) == 16_384
assert t["total"] == t["read"] + t["transfer"] and t["share"] == pytest.approx(t["transfer"] / t["total"])
def test_waiting_grows_with_gpus_and_shrinks_with_a_faster_link():
for key in g.MODELS:
shares = [g.word_time(key, n, "fast")["share"] for n in g.COUNTS]
assert shares == sorted(shares)
for n in g.COUNTS[1:]:
assert g.word_time(key, n, "slow")["transfer"] > g.word_time(key, n, "fast")["transfer"]
def test_more_gpus_stop_helping():
"""Guards the lesson's turning point: the small model is faster on 8 GPUs than on 16, on both links, and each doubling gains less."""
for link in g.LINKS:
times = [g.word_time("8b", n, link)["total"] for n in g.COUNTS]
gains = [a / b for a, b in zip(times, times[1:])]
assert gains == sorted(gains, reverse=True) and gains[-1] < 1 < gains[-2]
def test_best_gpu_count_depends_on_the_link():
"""Guards the claim 'a faster link lets more GPUs help': best count per model on each link, any whole number that fits, up to 64."""
best = {k: {link: g.best_n(k, link) for link in g.LINKS} for k in g.MODELS}
assert best == {"8b": {"fast": 10, "slow": 8}, "70b": {"fast": 19, "slow": 15}, "405b": {"fast": 36, "slow": 28}}
assert g.best_n("70b", "slow", words=64) == 14 < g.best_n("70b", "slow")
# among the counts the page offers, the best is the same on both links, so the lesson does not claim otherwise
offered = {k: {link: min(g.COUNTS, key=lambda n: g.word_time(k, n, link)["total"] if g.fits(k, n) else float("inf")) for link in g.LINKS} for k in g.MODELS}
assert offered == {"8b": {"fast": 8, "slow": 8}, "70b": {"fast": 16, "slow": 16}, "405b": {"fast": 16, "slow": 16}}
# slices are cut by attention head, so a usable count divides the heads: 19 for the 70B model (64 heads) is not one
usable = {k: {link: g.best_n(k, link, divide_heads=True) for link in g.LINKS} for k in g.MODELS}
assert usable == {"8b": {"fast": 8, "slow": 8}, "70b": {"fast": 16, "slow": 16}, "405b": {"fast": 32, "slow": 32}}
assert all(g.load_config(k)["num_attention_heads"] % n == 0 for k in g.MODELS for n in usable[k].values())
assert g.load_config("70b")["num_attention_heads"] % best["70b"]["fast"] != 0
def test_layer_groups_fit_the_model_but_do_not_speed_up_one_user():
"""Guards the pipeline claim: for one user a word still passes through every weight once, so the time stays near one GPU's full read."""
for key in g.MODELS:
whole = g.read_time(key, 1)
for n in g.COUNTS:
for link in g.LINKS:
total = g.word_time(key, n, link, "groups")["total"]
assert whole <= total <= 1.02 * whole
assert g.word_time(key, n, link, "groups")["share"] < 0.02
with pytest.raises(ValueError): # a split the toy does not know must fail, not be priced as free
g.word_time("8b", 2, "fast", "diagonal")
def test_many_words_at_once_make_the_link_matter():
"""Guards the 'words at once' setting: the slow link costs 5x more extra time at 64 words than at 1 (70B on 8 GPUs)."""
gap = {w: g.word_time("70b", 8, "slow", words=w)["total"] - g.word_time("70b", 8, "fast", words=w)["total"] for w in g.WORDS}
assert gap[64] > 5 * gap[1] and list(gap.values()) == sorted(gap.values())
def test_no_colour_outside_the_site_palette_in_the_lesson_folder():
"""Guards a stray blue or vermillion from the other palette; this file is the only place the codes are written."""
banned = [bytes(c, "ascii").lower() for c in ("0072B2", "56B4E9", "D55E00")]
for path in LESSON.rglob("*"):
if path.is_file() and path.suffix in {".svg", ".html", ".md", ".py", ".json"} and path.name != Path(__file__).name:
text = path.read_bytes().lower()
assert not any(b in text for b in banned), path.name
def _pct(x: float, places: str = "0.1") -> str:
return half_up(x * 100, places) + "%"
def test_every_number_the_page_and_figures_quote_comes_from_the_code():
"""Guards a lesson that drifts from its code: each figure of the prose, the tables and the SVGs is rebuilt here from
gpu_split and must appear in the file that quotes it. One assert names every fragment that is missing."""
lesson = (LESSON / "lesson.md").read_text()
figs = {p.name: re.sub(r"<[^>]+>", " ", p.read_text()) for p in (LESSON / "figs").glob("*.svg")}
wt = lambda key, n, link="fast", **kw: g.word_time(key, n, link, **kw)
mem4 = g.memory_per_gpu("70b", 4)
whole70 = g.weight_bytes(g.load_config("70b"))
gains = [wt("8b", a)["total"] / wt("8b", b)["total"] for a, b in zip(g.COUNTS, g.COUNTS[1:])]
wps = lambda n, link: str(round(wt("8b", n, link)["words_per_s"]))
in_lesson = [
f"{gib(whole70)} GiB", f"{g.USABLE_BYTES / g.GIB:.1f} GiB", f"{gib(mem4['weights'])} GiB", f"{gib(mem4['total'])} GiB",
f"{gib(g.weight_bytes(g.load_config('405b')))} GiB", f"{g.least_n('405b')} GPUs", f"{gib(g.memory_per_gpu('405b', 16)['total'])} GiB",
f"{ms(g.read_time('70b', 1))} ms", f"{ms(g.read_time('70b', 4))} ms", f"{ms(wt('70b', 4)['transfer'])} ms",
f"{ms(wt('70b', 4)['total'])} ms", f"{g.message_bytes(g.load_config('70b')):,} bytes", f"{g.message_bytes(g.load_config('405b')):,} bytes",
f"{g.ALL_REDUCES_PER_LAYER * g.load_config('70b')['num_hidden_layers']} swaps", f"{g.ALL_REDUCES_PER_LAYER * g.load_config('405b')['num_hidden_layers']} for each word",
f"{mem4['total']:,} bytes", f"{mem4['weights']:,}", f"{mem4['cache']:,}", f"{mem4['working']:,}",
", ".join(_pct(wt("70b", n)["share"]) for n in (2, 4, 8)) + " and " + _pct(wt("70b", 16)["share"]),
f"{gains[0]:.2f}, {gains[1]:.2f} and {gains[2]:.2f} times", f"{gains[3]:.2f} times",
f"{ms(wt('70b', 8)['total'])} ms a word on the fast link and {ms(wt('70b', 8, 'slow')['total'])} ms on the slow one",
f"{ms(wt('70b', 8, words=64)['total'])} ms on the fast link and {ms(wt('70b', 8, 'slow', words=64)['total'])} ms on the slow one",
f"{ms(wt('70b', 4, split='groups')['total'])} ms a word on four GPUs", f"{_pct(wt('70b', 4, split='groups')['share'], '0.01')} of it waiting",
f"{ms(wt('8b', 2)['read'])} ms", f"{ms(wt('8b', 2)['transfer'])} ms", f"{ms(wt('8b', 2)['total'])} ms",
f"{half_up(g.allreduce_time(g.message_bytes(g.load_config('8b')), 2, 'fast') * 1e6)} in all",
f"so a swap takes {half_up(g.allreduce_time(g.message_bytes(g.load_config('70b')), 4, 'fast') * 1e6)}.",
f"{half_up(g.allreduce_time(g.message_bytes(g.load_config('70b')) * 64, 8, 'fast') * 1e6)} microseconds on the fast link and "
f"{half_up(g.allreduce_time(g.message_bytes(g.load_config('70b')) * 64, 8, 'slow') * 1e6)} on the slow one",
_pct((g.BASE_US + 6 * 0.6) / (g.allreduce_time(g.message_bytes(g.load_config("70b")), 4, "fast") * 1e6)), # start-up share of one swap
f"{ms(g.compute_time('70b', 8, 64))} ms", f"{ms(g.read_time('70b', 8))} ms", f"{ms(g.compute_time('70b', 1))} ms",
f"{wt('70b', 8)['total'] / wt('70b', 16)['total']:.2f} times faster, from {ms(wt('70b', 8)['total'])} to {ms(wt('70b', 16)['total'])} ms",
f"from {wps(8, 'fast')} to {wps(16, 'fast')} words a second",
f"{ms(wt('70b', 8, 'slow', words=64)['total'] - wt('70b', 8, words=64)['total'])} ms more at 64",
f"{ms(wt('70b', 8, 'slow')['total'] - wt('70b', 8)['total'])} ms more at 1 word",
"10 on the fast link and 8 on the slow one", "19 and 15", "36 and 28",
"Each model has {0} key/value heads, so past {0} GPUs".format(*{g.load_config(k)["num_key_value_heads"] for k in g.MODELS}),
]
shares = lambda rows: ", ".join(_pct(t["share"]) for t in rows[:-1]) + " and " + _pct(rows[-1]["share"])
for n in g.COUNTS: # the 8B table: one row a count, milliseconds on each link
in_lesson.append(f"| {n} | {ms(wt('8b', n)['total'])} | {ms(wt('8b', n, 'slow')['total'])} |")
for words in g.WORDS: # the words-at-once table
in_lesson.append(f"| {words} | {ms(wt('70b', 8, words=words)['total'])} | {ms(wt('70b', 8, 'slow', words=words)['total'])} |")
for link in g.LINKS: # the waiting shares the tables leave to the sentence below each
in_lesson.append(shares([wt("8b", n, link) for n in g.COUNTS[1:]]))
in_lesson.append(shares([wt("70b", 8, link, words=w) for w in g.WORDS]))
for words in (1, 64): # the link-direction sentence: as written, then with both link rates halved
halved = halved_link_times("70b", 8, words)
slow, slow_half = wt("70b", 8, "slow", words=words)["total"], halved["slow"]
in_lesson.append(f"{ms(wt('70b', 8, words=words)['total'])} to {ms(halved['fast'])} ms on the fast link and from {ms(slow)} to {ms(slow_half)} ms on the slow one"
if words == 1 else f"the slow link goes from {ms(slow)} to {ms(slow_half)} ms, {_pct(slow_half / slow - 1, '1')} more")
if words == 64:
assert 0.25 < slow_half / slow - 1 < 0.30
heads = [g.load_config(k)["num_attention_heads"] for k in g.MODELS]
usable = [g.best_n(k, "fast", divide_heads=True) for k in g.MODELS] # the same on both links, as test_best_gpu_count_depends_on_the_link asserts
in_lesson += [f"the {heads[0]}, {heads[1]} and {heads[2]} the three models have", f"so {g.best_n('70b', 'fast')} GPUs is not a usable count",
f"the fastest are {usable[0]} for the 8B model, {usable[1]} for the 70B and {usable[2]} for the 405B, on both links",
f"they are {usable[0]}, {usable[1]} and {usable[2]}.", f"({g.USABLE * 100:.0f}% of its {g.GPU_GIB} GiB)"]
rows70 = [wt("70b", n) for n in (2, 4, 8, 16)]
in_figs = {
"weights-split.svg": [gib(whole70), f"{g.USABLE_BYTES / g.GIB:.1f}", gib(mem4["weights"])],
"time-per-word.svg": [ms(t["total"]) + " ms" for t in rows70] + [", ".join(_pct(t["share"]) for t in rows70)],
"stops-helping.svg": [f"{wps(8, 'fast')} to {wps(16, 'fast')} and the slow link {wps(8, 'slow')} to {wps(16, 'slow')}"],
"slices-vs-groups.svg": [ms(wt("70b", 4)["total"]), ms(wt("70b", 4, split="groups")["total"]), ms(g.read_time("70b", 4))],
}
missing = [f"lesson.md: {x}" for x in in_lesson if x not in lesson]
missing += [f"{name}: {x}" for name, parts in in_figs.items() for x in parts if x not in figs[name]]
assert not missing, "quoted but not found:\n" + "\n".join(missing)
def _without_docstring(source, node):
"""The lines of one function, minus its docstring."""
lines = source.splitlines()[node.lineno - 1:node.end_lineno]
first = node.body[0]
if isinstance(first, ast.Expr) and isinstance(first.value, ast.Constant) and isinstance(first.value.value, str):
del lines[first.lineno - node.lineno:first.end_lineno - node.lineno + 1]
return "\n".join(lines)
def test_every_python_block_in_the_lesson_is_a_verbatim_excerpt_of_the_source():
"""The lesson shows `memory_per_gpu` with its docstring left out. If the code changes and the page does not,
the reader copies a function that was never tested."""
lesson = (LESSON / "lesson.md").read_text()
source = Path(g.__file__).read_text()
in_source = {n.name: _without_docstring(source, n) for n in ast.parse(source).body if isinstance(n, ast.FunctionDef)}
blocks = re.findall(r"^```python\n(.*?)^```", lesson, flags=re.S | re.M)
shown = [node for block in blocks for node in ast.parse(block).body]
assert {node.name for node in shown} == {"memory_per_gpu"}
for block in blocks:
for node in ast.parse(block).body:
assert _without_docstring(block, node) == in_source[node.name], f"{node.name} differs from gpu_split.py"Show code
"""Guards for ../widget.html: its data block, its size, and that its JavaScript computes what gpu_split.py does."""
import json
import re
import subprocess
import gpu_split as g
from embed_widget import WIDGET, embed
HTML = WIDGET.read_text(encoding="utf-8")
# Reads the number code out of the widget between its marker comments, runs it, and compares with the
# values Python wrote into the fixture: the numbers, the usable memory and compute constants, the fits verdict,
# the bar widths and the chart rows the page draws.
PARITY = r"""
import fs from 'node:fs';
const [fixture, widget] = process.argv.slice(2);
const fx = JSON.parse(fs.readFileSync(fixture, 'utf8'));
const html = fs.readFileSync(widget, 'utf8');
const DATA = JSON.parse(html.match(/<script type="application\/json" id="gpu-data">([\s\S]*?)<\/script>/)[1]);
const js = html.slice(html.indexOf('/* ---------- the numbers'), html.indexOf('/* ---------- the page'));
const api = new Function('DATA', "'use strict';" + js + '; return {params, memory, fits, wordTime, usableBytes, computeTime, view, chart};')(DATA);
const close = (a, b) => Math.abs(a - b) <= 1e-9 * Math.abs(b);
const bad = [];
if (api.usableBytes() !== fx.usable) bad.push(['usableBytes', api.usableBytes(), fx.usable]);
fx.cases.forEach((c) => {
const m = DATA.models.find((x) => x.key === c.key);
const link = DATA.links.find((x) => x.key === c.link);
if (api.params(m) !== c.params) bad.push(['params', c.key]);
const mem = api.memory(m, c.n);
for (const k of ['weights', 'cache', 'working', 'total']) if (mem[k] !== c.memory[k]) bad.push([c.key, c.n, k, mem[k], c.memory[k]]);
if (api.fits(m, c.n) !== c.fits) bad.push([c.key, c.n, 'fits']);
const t = api.wordTime(m, c.n, link, c.split, c.words);
for (const [a, b] of [['read', 'read'], ['transfer', 'transfer'], ['total', 'total'], ['share', 'share'], ['wordsPerS', 'words_per_s']]) {
if (!close(t[a], c.time[b])) bad.push([c.key, c.n, c.link, c.split, c.words, a, t[a], c.time[b]]);
}
if (!close(api.computeTime(m, c.n, c.words), c.compute)) bad.push([c.key, c.n, c.words, 'computeTime']);
const v = api.view(m, c.n, link, c.split, c.words);
if (v.ok !== c.fits) bad.push([c.key, c.n, c.link, c.split, c.words, 'view.ok', v.ok, c.fits]);
if (!close(v.memBar, 100 * Math.min(1, c.memory.total / fx.usable))) bad.push([c.key, c.n, 'memBar', v.memBar, 100 * Math.min(1, c.memory.total / fx.usable)]);
if (!close(v.waitBar, c.fits ? 100 * c.time.share : 0)) bad.push([c.key, c.n, c.link, c.split, c.words, 'waitBar', v.waitBar, c.fits ? 100 * c.time.share : 0]);
// the chart for this model, link, split and words: one row per GPU count, bars against the slowest setting that fits
const same = fx.cases.filter((x) => x.key === c.key && x.link === c.link && x.split === c.split && x.words === c.words);
const top = Math.max(...same.filter((x) => x.fits).map((x) => x.time.total));
const rows = api.chart(m, link, c.split, c.words);
rows.forEach((r, i) => {
const x = same[i];
if (r.n !== x.n || r.ok !== x.fits) bad.push([c.key, c.link, c.split, c.words, 'chart row', r.n, r.ok, x && x.n, x && x.fits]);
else if (x.fits && !(close(r.readBar, 100 * x.time.read / top) && close(r.waitBar, 100 * x.time.transfer / top))) bad.push([c.key, c.link, c.split, c.words, 'chart bars', r.n]);
});
});
console.log(JSON.stringify({ cases: fx.cases.length, bad: bad.slice(0, 5) }));
process.exit(bad.length ? 1 : 0);
"""
def test_widget_data_block_matches_gpu_split():
"""Guards an edit to the constants or model files that was not copied into the widget (run embed_widget.py)."""
assert embed(HTML) == HTML
def test_every_element_the_script_looks_up_exists():
"""Guards a widget that never starts: a lookup of an id the page does not contain throws in the browser."""
page = HTML[:HTML.index("<script")]
lookups = set(re.findall(r"\$\('([\w-]+)'\)", HTML[HTML.index("<script>"):]))
assert lookups and all(f'id="{i}"' in page for i in lookups), sorted(i for i in lookups if f'id="{i}"' not in page)
assert 'id="gpu-data"' in HTML and 'id="gpu-root"' in page
def test_widget_is_small_and_asks_for_nothing_online():
assert len(HTML.encode("utf-8")) < 40 * 1024
for banned in ("http://", "https://", "fetch(", "XMLHttpRequest", "WebSocket", "sendBeacon", "import("):
assert banned not in HTML, banned
assert " src=" not in HTML and "@import" not in HTML
def test_widget_numbers_match_the_python_code(tmp_path):
"""Guards the widget's JavaScript drifting from gpu_split.py: the usable memory and compute constants, memory, fit verdict,
every time, and the bars and chart rows the page draws (they must follow the link and the split) must be the lesson's."""
cases = [{"key": key, "n": n, "link": link, "split": split, "words": words, "params": g.param_count(g.load_config(key)),
"memory": g.memory_per_gpu(key, n), "fits": g.fits(key, n), "time": g.word_time(key, n, link, split, words),
"compute": g.compute_time(key, n, words)}
for key, n, link, split, words in g.settings()]
(tmp_path / "fixture.json").write_text(json.dumps({"usable": g.USABLE_BYTES, "cases": cases}), encoding="utf-8")
(tmp_path / "parity.mjs").write_text(PARITY, encoding="utf-8")
run = subprocess.run(["node", str(tmp_path / "parity.mjs"), str(tmp_path / "fixture.json"), str(WIDGET)],
capture_output=True, text=True, timeout=120)
assert run.returncode == 0, run.stdout + run.stderr
assert json.loads(run.stdout)["cases"] == len(g.MODELS) * len(g.COUNTS) * len(g.LINKS) * len(g.SPLITS) * len(g.WORDS) == 180
assert any(c["fits"] for c in cases) and any(not c["fits"] for c in cases)