← Inference From First Principles

Distributed Inference

GPT-2 split across processes by tensor and pipeline parallelism, exact to the last token and no faster, and the cost model that explains why: what each strategy splits, what it must communicate, what that communication costs, and what a machine has to provide before splitting helps.

GPT-2 was split across six processes, each holding a sixth of every layer’s attention heads and feed-forward width. Every one of the 15 runs of the split code, over 1, 2, 3, 4 and 6 processes, produced the library model’s 25 tokens exactly, and the final logits of the prompt differed from the library’s by at most 9.2e-5. The split model was also slower. One process took 23 milliseconds to produce a token. Six took 54 to 59.

Nothing was wrong with the split. The tokens were right and the arithmetic was shared out evenly. What the six processes had in common was one memory system, and a decode step is a pass over the weights. Eight independent copies of the model, running side by side with nothing to communicate, produced 38 to 45 tokens a second in total, the same as one copy produced alone (40), because one process already streamed the weights at 21 GB/s and the machine’s memory gave no more than about 25 GB/s (between 20 and 30 in my runs) to any number of processes.

That is the chapter’s argument in one measurement. Splitting a model over devices pays when each device brings its own memory bandwidth, or its own memory to hold what does not fit, and costs a fixed price in communication for every layer it cuts. The rest is measuring the price, checking the arithmetic that predicts it, and being clear about which parts of this were measured on this machine and which are modelled for hardware it does not have.

The question

This machine has one GPU, no NCCL, and 24 CPU threads. A chapter on distributed inference cannot, here, measure a multi-GPU system, and the starter chapter did not try: it counted bytes in a simulator and said so. This chapter has three kinds of material and keeps them apart.

First, a cost model: what each strategy splits, how many bytes and messages per generated token it needs, and what they cost on a link described by a latency and a bandwidth. Second, real distributed execution: GPT-2 split by heads and by layers across operating-system processes that talk through PyTorch’s gloo backend, checked for exactness and timed, with the collectives measured on their own. Third, records of what was not run: NCCL, a second GPU, a second node.

The questions are four. Does splitting change the answer? What does the communication cost, and does the standard two-constant model predict it? When does splitting make a token faster, and what did it do here? And, for real model configurations, which strategies fit on which devices and how much latency per collective can a split afford?

Where this sits

Chapter 7 measured this GPU’s memory bandwidth, 303 GB/s, and showed that decoding one token is limited by how fast the weights can be streamed. That is the quantity the cost model divides by. Chapters 14, 21 and 23 gave the size of the KV cache per token and showed that grouped-query attention shrinks it; here the number of KV heads also limits how far attention can be split. Chapter 24 read ten configurations into fingerprints and said every option that changes what is computed belongs in the manifest; the parallel layout (which strategy, over how many devices) will have to join them, and this chapter measures whether it changes the answer. Chapter 27 ran many streams on one engine; the engine here is smaller and its purpose is different.

The chapter also corrects the starter’s accounting. Its simulator charged tensor parallelism (p − 1) times the activation bytes, called data parallelism’s cost “gradient/activation all-reduce” (inference has no gradients, as the starter’s own text says), and treated the other strategies as one activation each. None of those is what a collective moves. A test in the repository compares the starter’s tensor figure of 300 bytes for four devices with the 150 a ring all-reduce puts on the wire per device for the same buffer.

Chapter 29 takes up splitting prefill from decode across machines, which assumes this chapter’s accounting of what moving state costs. Chapter 30 takes up what happens when one of several devices fails.

What each strategy splits

Data parallelism puts a whole copy of the model on each device and gives each request to one of them. Nothing is split and nothing is synchronised: unlike training there are no gradients. What remains is routing, the manifest of Chapter 2, and the choice of replica. It is the only strategy that adds throughput without touching the latency of a request, and it needs every replica to hold the whole model.

Tensor parallelism cuts each layer across devices. The scheme is Megatron-LM’s. The first matrix of the feed-forward block is split by columns, so each device applies the activation function to its own slice with no communication; the second is split by rows, so the devices’ partial outputs add up to the layer’s output. Attention is split by heads: the query, key and value projections by columns, so a head’s whole computation is local, and the output projection by rows. The paper counts the result as two all-reduces in the forward pass of a layer (one after attention, one after the feed-forward), and two more in the backward pass, which inference does not have. The embedding table is split along the vocabulary and the logits are gathered; the paper’s training setting can fuse that gather into the loss, and inference cannot, because it needs the logits.

A decode step of GPT-2 therefore makes 2 × 12 = 24 all-reduces, each of a tensor of 768 values (3,072 bytes in float32), and one all-gather of the vocabulary-split logits (about 100 KB per device at 2 devices). Llama-3.1-8B makes 64 all-reduces of 8 KB each and one gather, per generated token. The number is fixed by the layer count and not by the number of devices.

Pipeline parallelism cuts the model by layers into stages, one per device, and sends the activations from stage to stage. The communication is small: one activation tensor at each of the p − 1 boundaries, plus a token id back. The catch for generation is a dependency that training does not have. A sequence’s next token cannot enter the first stage until the last stage has produced the previous one, so a single sequence visits the stages in turn and only one is busy at a time. GPipe’s bubble, \((K−1)/(M+K−1)\) for \(M\) micro-batches over \(K\) stages and negligible for \(M \ge 4K\) in the authors’ experiments, describes pipelining independent micro-batches of a training step. For decoding, the unit that fills the pipeline is a different sequence, and a round in which every one of \(m\) sequences advances a token takes \(\max(m, p)\) stage-times: with fewer sequences than stages some stages are idle and the round is one trip through the pipeline, and with at least as many the slowest stage is busy throughout.

Expert parallelism places the experts of a mixture-of-experts layer on different devices and sends each token to the devices that own the experts it chose, and back: two all-to-alls per MoE layer. Chapters 19 and 20 built the dispatch and its balancing; here only the communication is measured.

The DeepInfer README at the commit pinned in this book lists, among its supported features, expert parallelism with load balancing (“MoE EP + EPLB”), data parallelism for attention, and context parallelism. Tensor and pipeline parallelism are not stated. Its benchmark section reports 4 GPUs for prefill and 64 for decode, each within one scale-up domain, without naming the interconnect or stating a maximum number of devices, and lists a KV cache system, prefill/decode disaggregation and load balancing across replicas as required for a complete service and not included. That is the authors’ description of a system on hardware I do not have; I did not run it.

A cost model

A message of \(n\) bytes between two devices takes \(\alpha + n/\beta\) seconds: a latency and a bandwidth. A ring all-reduce over \(p\) devices is a reduce-scatter followed by an all-gather, each \(p − 1\) steps, so

\[ T_{\text{all-reduce}}(n, p) = 2(p-1)\,\alpha + \frac{2(p-1)}{p}\,\frac{n}{\beta}, \qquad T_{\text{all-gather}}(n, p) = (p-1)\left(\alpha + \frac{n}{\beta}\right), \]
where \(n\) is the buffer being reduced, or the bytes each device contributes to the gather. Each device puts \(2(p−1)n/p\) bytes on the wire in an all-reduce, which is 150 bytes for a 100-byte buffer on four devices, not the 300 of the starter. These are standard formulas; I did not read a source for them, and the chapter tests them against measurements.

For decode the model needs two more things. A decode step is memory-bound (Chapter 7), so its time is the bytes it must read divided by memory bandwidth. Those bytes are not the bytes the model occupies: the input embedding is stored but a step reads one row of it, while the output head is read in full. The cost model keeps stored bytes, which decide whether a device can hold its share, apart from streamed bytes, which decide how long a step takes. The step on a device is then its streamed weights plus its share of the KV cache, over the bandwidth, plus the communication.

The model needs the shape of a model, and the formula that gives it should be exact or its predictions are suspect. I checked it against the headers of the checkpoints in the local cache. For seven models it matched the number of values in the files exactly: GPT-2 (124,439,808), TinyLlama (1,100,048,384), Phi-3-mini (3,821,079,552), Gemma-2-2B (2,614,341,888), Mistral-7B (7,248,023,552), Llama-3.2-1B (1,235,814,400) and Llama-3.2-3B (3,212,749,824). Qwen3-1.7B (1,720,574,976) and Llama-3.1-8B (8,030,261,248) have formula figures only, because their weights are not cached, and Pythia’s layout has no formula. The first comparison disagreed for three models, and each disagreement was a fact about the checkpoint and not an error of arithmetic. GPT-2’s file also stores its causal-mask buffers, 12 × 1,024² = 12,582,912 values that are not parameters. Mistral’s folder holds a second copy of the weights. And Gemma-2 ties its input and output embeddings by default without its configuration saying so, and has four normalisation layers per block where the Llama layout has two. A formula checked on Llama alone would have passed.

What this machine can do

Before the experiments, what is and is not available, recorded by a command and not assumed:

Capability Here
torch.distributed with the gloo backend (CPU, over TCP) available, after setting USE_LIBUV=0 (the Windows build has no libuv)
NCCL not available in this build; the PyTorch documentation lists NCCL as supporting GPU collectives and Windows as supporting every backend but NCCL
MPI not available
CUDA devices one
gloo reduce_scatter not supported (it raised), as the documentation’s table says
gloo all_to_all_single worked, though the table’s all_to_all row says unsupported

So NCCL, multi-GPU and multi-node runs are NOT_RUN. The distributed code that does run needs no GPU and no accelerator-specific library: it needs a machine with the gloo backend and at least two CPU threads, and its tests skip themselves otherwise, so continuous integration without a multi-GPU machine can still run them.

Does splitting change the answer?

Hypothesis (book hypothesis). Splitting GPT-2 by heads and feed-forward width, or by layers, changes the arithmetic’s order and nothing else, so the generated tokens equal the unsplit model’s.

Test. The lab implements GPT-2’s forward pass in a form that can hold any share of the heads and of the feed-forward width, and any range of layers, with its own KV cache of its heads (GPT2Part). It is compared with the Hugging Face model on a 32-token prompt of book text and 25 greedy tokens. Tensor parallelism ran on 1, 2, 3, 4 and 6 processes, three times each. Pipeline parallelism ran on 1, 2 and 4 stages with 1, 2, 4 and 8 sequences in flight, three times each.

Result. In 15 of 15 tensor-parallel runs the 25 tokens equalled the library’s, every rank agreed, and the largest difference in the final logits of the prompt was 9.2e-5. In every pipeline-parallel run every sequence equalled the library’s greedy tokens. The hypothesis stands for greedy decoding of GPT-2 in float32. Reduction order differs (the logits are not bitwise equal, and Chapter 17 and Chapter 24 explain what a near-tie does with that), so the claim is about these prompts and this precision, not a guarantee; half precision, other models and sampled decoding were not run.

The count of collectives is also a result: each decode step made exactly 25, as the Megatron arithmetic says.

What the collectives cost

Each collective was timed alone between 2, 4 and 8 processes over loopback TCP, across sizes from 4 bytes to 16 MB.

Median time, messages of 64 B to 64 KB 2 processes 4 processes 8 processes
all-reduce 132 to 140 µs 339 to 362 µs 1.59 to 1.68 ms
all-gather 90 to 111 µs 206 to 246 µs 0.88 to 1.13 ms
all-to-all 57 to 73 µs 104 to 157 µs 375 to 599 µs
ring send/receive 52 to 65 µs 57 to 78 µs 117 to 145 µs

For messages this small the time is flat in the size: 64 bytes and 64 KB cost almost the same. That is the regime of decode. The all-reduce of a token’s activation is 3 KB. The cost of a decode collective is its latency, and the bandwidth term is not in it.

Does the two-constant model predict the latency? From the 2-process fit (an intercept of 118.7 µs, which as \(2(p−1)\alpha\) gives \(\alpha = 59\) µs) the model predicts a 4-process all-reduce of 356 µs; the measured flat value is 339 to 362. It predicts 830 µs for 8 processes; the measured value is 1.59 to 1.68 ms, twice as much. The all-to-all follows the model more closely: with \(\alpha\) taken from the 54 µs of a point-to-point message, \((p−1)\alpha\) gives 162 µs for 4 processes (measured 104 to 157) and 378 µs for 8 (measured 375 to 599). So the model predicted the 4-process latency within a few per cent and was wrong by a factor of two at 8, where eight processes on a 24-thread machine are no longer the idle, evenly spread participants the formula assumes. I do not know whether the cause is scheduling, the gloo algorithm or both.

It fails for large messages on this transport. The fitted model underpredicted held-out all-reduces of 4 MB and 16 MB by 86% and 94% at 2 processes, by 65% and 81% at 4, and by 30% and 74% at 8. A finer grid made the failure specific. At 2 processes, the all-reduce took 0.19 to 0.20 ms at 256 KB, 0.41 to 0.52 at 1 MB, then 11.7 to 14.3 ms at 2 MB, and at 4 MB 3.2, 14.3 and 11.9 ms in three repetitions: a jump of about twenty-five times for twice the bytes, and then no regularity. The 4-process times were smooth (1.12 to 1.20 ms at 2 MB, 5.0 to 5.3 at 4 MB). The plateaus near 12 to 14 ms are close to the Windows scheduler tick, which suggests a timer, but I did not look for it. For decode it does not matter, because nothing in a decode step approaches a megabyte. For a prefill that all-reduces a long prompt’s activations it would, and the model would be the wrong tool on this transport.

Splitting the model, and what it did to the step

Tensor-parallel GPT-2, decode steps of a 25-token generation after a 32-token prompt, three runs per size, median step time and the time inside collectives:

Processes Step Time in collectives Step minus collectives
1 23.1 to 23.9 ms 0 23.1 to 23.9 ms
2 25.7 to 30.5 4.4 to 5.3 21.3 to 25.1
3 32.7 to 45.3 8.6 to 16.1 24.1 to 29.2
4 32.8 to 47.5 11.6 to 21.7 21.1 to 25.8
6 54.4 to 58.7 32.6 to 35.2 21.8 to 24.3

Two things are visible. Communication grew with the number of processes: 25 collectives per step cost about 180 to 210 µs each at 2 processes, 460 to 870 µs at 4 and 1.3 to 1.4 ms at 6. And the non-communication time did not fall. Each process held a half, a quarter or a sixth of the weights and took as long over its share as one process took over all of it.

The communication is partly predicted. From the microbenchmarks, 24 all-reduces at their 1 KB time plus one all-gather at 64 KB come to 3.3 ms at 2 processes and 8.4 ms at 4, against 4.4 to 5.3 and 11.6 to 21.7 measured. The measured figure includes the time a process waits in the collective for the slowest of the others, and the model’s does not, so the measured time exceeded the prediction by 35 to 63% at 2 processes and by 37 to 157% at 4. As a prediction of the order of magnitude and the trend it is useful; as a figure to design with, it is optimistic.

Where the head goes matters. With the output head replicated on all four processes instead of split by vocabulary, a step took 64.5 to 69.9 ms against 45.7 to 50.1 ms with the head split. The head is 154 MB of float32 weights, a third of the model, and replicating it quadruples that traffic. That is the same lesson as the main one, arrived at from a different side.

Prefill behaves differently, and the difference is informative. A 32-token prefill took 73 to 76 ms on one process and 48 to 65 ms on two and 40 to 64 ms on four, a speed-up of up to about 1.9 times in the best runs (the ranges overlap, so the claim is loose), before it got worse at six (76 to 98 ms). A prefill multiplies 32 rows at once, so it does more arithmetic per byte of weights and its limit is nearer arithmetic than memory.

Why: the memory system is shared

Hypothesis (book hypothesis). The decode step on this machine is limited by memory bandwidth that all the processes share, so splitting the weights across processes cannot make the step faster, while it adds communication.

Two measurements test it independently of the tensor-parallel runs.

Independent replicas. Each process decoded its own sequence, with nothing to communicate, and all started together. Aggregate tokens per second over three runs: 1 replica, 39.6 to 40.6; 2 replicas, 37.9 to 50.1; 4, 38.9 to 43.4; 6, 44.6 to 50.0; 8, 37.8 to 44.7. Eight times the processes gave about the throughput of one, with each replica’s own rate falling from 40 to 19, 11, 8 and 5 tokens a second. A decode step streams 497.8 MB of weights, and one process did it at 20.8 to 21.5 GB/s.

A memory benchmark. Each process summed its own 512 MiB tensor, simultaneously. The aggregate bandwidth across 2, 4, 6, 8 and 12 processes was 16 to 24, 21 to 26, 19 to 28, 25 to 30 and 22 to 27 GB/s: a plateau of roughly 25 GB/s that no number of processes exceeded, with each process’s share falling as more were added.

The two models predict different things. If every process had its own memory bandwidth, the non-communication time of a step would fall with the share of the weights: 12, 6 and 4 ms at 2, 4 and 6 processes. If the bandwidth is shared, it would stay at about 23 ms. The measured values were 21 to 29 at every size. The hypothesis stands, and the explanation is not circular: it rests on a benchmark that has nothing to do with tensor parallelism and on a prediction, the constant non-communication time, that could have failed.

That is also why a pure-CPU experiment of this kind validates the mechanism (the split is exact; the collectives cost what they cost) and not the speed-up. On a GPU box each device has its own memory, which is where tensor parallelism’s gain comes from. The experiment cannot show that gain, and I am not claiming it.

The pipeline hits the same ceiling

Pipeline-parallel generation of 16 tokens from 32-token prompts, tokens per second over the whole run, prefill included, three runs:

Stages 1 sequence 2 sequences 4 sequences 8 sequences
1 29.5 to 36.1 27.3 to 33.0 33.6 to 36.0 32.8 to 33.2
2 32.9 to 34.9 38.6 to 41.2 34.7 to 40.7 39.5 to 41.7
4 24.9 to 32.3 28.6 to 39.1 33.5 to 38.7 38.2 to 42.4

The model’s prediction for one sequence holds: with one sequence in flight, 2 or 4 stages are no faster than one stage (it visits each in turn). With enough sequences to fill the stages the rate rises, but only to about 40 tokens a second, 1.2 to 1.3 times the single stage, and not to 2 or 4 times. The ceiling is the one the replicas hit. A pipeline of 4 stages with 8 sequences in flight is the best case for the strategy on paper, and the machine’s shared memory is the limit.

Expert dispatch

The all-to-all that carries tokens to experts and back was measured for 768-wide activations, top-2 routing, spread over 2, 4 and 8 processes. For a decode batch of 1 or 8 tokens it took 65 µs at 2 processes, 106 to 121 µs at 4 and 370 to 385 µs at 8: all latency. A decode step of a mixture model pays that twice per MoE layer. At 64 and 512 tokens the times rose to 0.17 to 0.52 ms and 3.6 to 5.5 ms at 4 and 8 processes. At 2 processes the same erratic jumps appeared as in the all-reduce (14.8 ms for 64 tokens, then 1.2 ms for 512), which I have left in the table.

The machine has one GPU, but it has a link: the host to device copy. With pinned memory, a copy of 8 KB took 11.1 µs to the device and 9.2 µs back; 1 MB took 95 and 88 µs (11.0 and 11.9 GB/s); 16 MB took 1.35 and 1.30 ms (12.5 and 12.9 GB/s); 256 MB took 21.8 and 21.7 ms (12.3 and 12.4 GB/s). So the link has a latency of about 10 µs and a bandwidth of about 12.4 GB/s. If two cards of this class exchanged data through the host, each crossing would cost this, and a transfer between them would pay it twice.

Modelled for hardware this machine does not have

The model can say what the measured constants imply for real configurations, and it has to be clear that this part is arithmetic and not measurement. The setting: decode of one sequence at a 4,096-token context in fp16, on a device with the 303 GB/s Chapter 7 measured, under three links. Two are measured here: the gloo loopback (54 µs, 1.2 GB/s, from the ring send/receive) and a transfer through the host (20 µs and 6.2 GB/s, the host-device figures doubled). The third uses the bandwidth the Megatron-LM paper reports for NVSwitch inside a server, 300 GB/s, with a latency of 10 µs that is my assumption and not a reported figure.

Llama-3.1-8B in fp16 needs 14.96 GiB for its weights, which does not fit a 12 GiB card. That is the first thing splitting does: it makes the model runnable.

Strategy, devices Weights per device Step with the gloo link through the host scale-up (assumed latency)
tensor, 2 7.48 GiB 33.2 ms 28.3 ms 26.9 ms
tensor, 4 3.74 GiB 34.5 ms 20.7 ms 16.7 ms
tensor, 8 1.87 GiB 56.1 ms 24.7 ms 15.4 ms
pipeline, 2 7.48 GiB 51.4 ms 51.4 ms 51.3 ms
pipeline, 4 4.23 GiB 51.5 ms 51.4 ms 51.3 ms
pipeline, 8 2.60 GiB 51.8 ms 51.5 ms 51.4 ms

(A single device with the memory to hold it would take 51.3 ms by the same arithmetic, which is the pipeline’s figure: a pipeline reads the whole model once per token whatever the number of stages.)

The pipeline makes a model that does not fit run at the speed a bigger device would, and no faster; its stages’ extra throughput appears only with several sequences in flight. Tensor parallelism reduces the latency of a token, down to a limit set by the link: with the slow gloo-class link, 8 devices are slower than 2 (56 against 33 ms), because 65 collectives times several hundred microseconds exceed the compute they save; with the fast link, 8 devices are faster than 4 (15.4 against 16.7 ms), barely.

How much latency per collective can a split afford? Tensor parallelism over \(p\) devices beats one device that could hold the model if the time per all-reduce is below \(T_1 (1 − 1/p) / 2L\), where \(T_1\) is the single-device step and \(L\) the number of layers (ignoring the final gather). For two devices:

Model All-reduces per token Largest time per all-reduce for a 2-device split to win
GPT-2 24 27 µs
Llama-3.2-1B 32 134 µs
Qwen3-1.7B 56 115 µs
Mistral-7B 64 381 µs
Llama-3.1-8B 64 401 µs

The larger the model, the more a collective can cost before the split stops paying, because the compute it hides is larger while the number of collectives stays at two per layer. The measured 2-process gloo all-reduce, 132 to 140 µs, would not make a 2-way split of GPT-2 pay (it needs under 27 µs) and would for an 8-billion-parameter model on a 303 GB/s device. Whether any real interconnect clears these thresholds is a question for hardware I did not measure.

The number of KV heads limits the split. Tensor parallelism divides the KV cache by splitting its heads. Llama-3.1-8B, Mistral-7B, Llama-3.2-1B and Qwen3-1.7B have 8 KV heads. Past 8 devices the extra devices each need a copy of a KV head, and the cache per device stops falling; a multi-query model (one KV head) cannot shard its cache at all. Chapter 21’s saving in cache size is therefore also a limit on tensor-parallel degree.

What the numbers do not show

  • No GPU was split. The distributed runs are CPU processes sharing a memory system. They show that the splits are exact, what the communication costs through gloo on loopback, and that decode on a shared memory system does not speed up. They do not show how much a GPU box gains.
  • Loopback TCP is not an interconnect. Its latency (tens of microseconds) and bandwidth (about 1.2 GB/s in the larger point-to-point messages) belong to this operating system’s network stack and to gloo. NCCL, NVLink and InfiniBand were not measured.
  • The collective model is validated only for small messages and few processes. It predicted the 4-process latency from the 2-process one, missed the 8-process one by a factor of two, and failed for messages over a megabyte.
  • The modelled table is arithmetic. It uses a roofline estimate of decode, three link parameters, one of which has an assumed latency, and ignores overlap of communication with compute, batching, and the sampling step.
  • One model, float32, greedy. Exactness was checked for GPT-2 in float32 with greedy decoding.
  • Noise. Three runs per cell on a desktop with other work; the tensor-parallel step times have ranges that overlap between neighbouring sizes.
  • Expert parallelism’s compute was not built. Only the all-to-all was measured.

Inference Lab increment

examples/ch28/parallel_lab.py adds:

  • Link, and the collective-time functions p2p_time, allreduce_time, allgather_time, alltoall_time, with allreduce_wire_bytes and fit_line.
  • ModelShape, shape_from_config and plan: parameter counts that match real checkpoints, stored against streamed bytes, KV sharing by head, and the cost of a decode step under data, tensor, pipeline and expert strategies; tp_break_even_allreduce_s; pipeline_round_time.
  • GPT2Part: GPT-2 by any share of the heads and the width, and any range of layers, with Comm, a communicator that counts.
  • run_task: a launcher of gloo process groups, and the tasks that measure collectives, tensor-parallel and pipeline-parallel generation, replicas and memory streaming.
  • Commands shapes, plans, collectives, collectives_fine, membw, tp, pp, dp, ep, link and gates.

To check the machine and run the split model:

python examples/ch28/parallel_lab.py gates
python examples/ch28/parallel_lab.py tp

The rule this adds: a parallel layout is part of the artefact’s identity until it is shown not to change the answer, and a split is justified by what each device brings, not by how many devices there are. The cost model says what to expect, the exactness check says whether the split is the same model, and the measurement says whether the machine provides what the model assumed.

What you should now be able to do

  • Say what each of the four strategies splits and, for a given model, how many collectives of what size one decoded token needs.
  • Write the ring all-reduce and all-gather times in terms of a latency and a bandwidth, and say for which message sizes the latency alone decides.
  • Check a parameter-count formula against a checkpoint’s header, and explain a mismatch.
  • Show that a split model is the same model, and explain why exactness does not imply speed.
  • Decide, for a model and a device, whether splitting is needed to fit it, and what latency per collective a split of it can afford.

Key terms

  • Data parallelism (in inference): replicas of the whole model serving different requests; no gradients, no collectives.
  • Tensor parallelism: each layer split across devices, with collectives inside it.
  • Pipeline parallelism: layers split into stages with activations passed between them.
  • Expert parallelism: experts of a mixture layer placed on different devices, with an all-to-all to move tokens.
  • Collective: a communication in which all devices take part (all-reduce, all-gather, all-to-all).
  • Alpha-beta model: the time of a message as a latency plus bytes over bandwidth.
  • Stored versus streamed bytes: what a device must hold, against what a step must read.
  • Scale-up domain: the group of devices joined by the fastest interconnect.

Reader exercises

  1. Run tp on a machine with several memory channels per socket, or several machines, and see whether the non-communication time starts to fall.
  2. Add context parallelism or sequence parallelism for prefill and measure what changes in the collectives.
  3. Replace gloo with NCCL on a multi-GPU machine, rerun collectives, and fit the two constants.
  4. Overlap the all-reduce of one layer with the compute of the next micro-batch and measure how much of the communication is hidden.
  5. Build the expert-parallel layer from Chapters 19 and 20 on top of the all-to-all and check it for exactness.
  6. Extend plan to speculative decoding (Chapter 17) and to a batch of 32 and see where the memory-bound assumption stops holding.
  7. Find the cause of the 12 to 14 ms plateaus in the 2-process collectives.
  8. Write the manifest fields that record a parallel layout, and a test that fails when two runs with different layouts claim the same identity.

References

Evidence in this chapter

Claim Category
Tensor parallelism over 2 to 6 processes and pipeline parallelism over 1 to 4 stages reproduce GPT-2’s greedy tokens; the largest logit difference is 9.2e-5 Book result (float32, CPU, 3 runs per cell)
A tensor-parallel decode step makes 24 all-reduces and one all-gather Book result; Megatron’s two per layer is Established background
A small gloo all-reduce costs 132 to 140 µs at 2 processes, 339 to 362 µs at 4, 1.59 to 1.68 ms at 8; flat in size up to 64 KB Book result (loopback TCP, this machine)
The two-constant model predicts the 4-process latency from the 2-process one, misses the 8-process one by a factor of two, and fails above 1 MB Book result
Tensor-parallel decode of GPT-2 took 23 ms on one process and 54 to 59 ms on six; the non-communication time stayed at 21 to 29 ms Book result
Eight independent replicas gave about the aggregate throughput of one (38 to 45 against 40 tokens/s) Book result
Memory streaming across 2 to 12 processes plateaued at about 25 GB/s Book result
A pipeline of 2 or 4 stages is no faster than one for a single sequence, and about 1.2 to 1.3 times faster with enough sequences Book result (shared memory ceiling)
The parameter formula matches seven checkpoints exactly Book result
Host-device link: about 10 µs latency and 12.4 GB/s Book result (this machine)
Llama-3.1-8B needs two 12 GiB cards; tensor parallelism over 2 to 8 devices takes 15 to 56 ms a token depending on the link Modelled (roofline, three link parameters, one latency assumed)
The largest time per all-reduce at which a 2-way split pays is 27 µs for GPT-2 and 401 µs for Llama-3.1-8B Modelled
The Megatron and GPipe partitions and bubble formula Reported by the papers; training setting; not reproduced
DeepInfer’s supported parallelism and benchmark hardware Reported by the authors’ README; not run
NCCL, multi-GPU, multi-node, NVLink, GPU tensor or pipeline parallelism NOT_RUN

Implementation status

  • Manuscript: content/books/inference/28-chapter.md
  • Runnable examples: examples/ch28/topology.py (original toy accounting, unchanged), examples/ch28/parallel_lab.py (cost model, shapes, split GPT-2, launcher, experiments)
  • Tests: tests/test_ch28_topology.py (toy), tests/test_ch28_parallel_lab.py (33 tests: the collective formulas and the fit, the starter’s figure against a ring all-reduce, the pipeline round time and the bubble formula, the parameter formula against tiny GPT-2 and Llama models including tied, untied and grouped-query layouts, Gemma and Qwen3 conventions, the checkpoint counter’s handling of buffers and duplicate copies, the KV sharing rule, every strategy’s plan including stored against streamed bytes, the break-even equation, an invalid split refused, the unsplit part against the library, and, when gloo and four CPU threads are available, real multi-process runs of the collectives, tensor parallelism over four ranks and pipeline parallelism over two stages for exactness; no GPU needed)
  • Evidence: research/ch28/sources.md, research/ch28/results.json, research/evidence-register.md
  • Verification command: python -m unittest discover -s tests -v (set CH28_SKIP_DISTRIBUTED=1 to skip the multi-process tests)

The chapter is a draft with draft = false. It makes no claim about any multi-GPU or multi-node system, about NCCL, or about the speed of a split on hardware with separate memories. Chapter 29 asks what it costs to move the state of a request between machines, which is where this chapter’s link arithmetic is used.