Pipeline parallelism
Compare kernel, parallelism and disaggregated-serving choices using memory, communication and latency constraints.
THE IDEA
Pipeline parallelism
Pipeline parallelism assigns consecutive layers to different devices (stages). Activations flow from stage to stage and several microbatches are in flight at once. Ideal speed-up is limited by two things: the bubble of idle time while the pipeline fills and drains, (p − 1)/(m + p − 1) of the schedule for p stages and m microbatches, and imbalance, because the slowest stage sets the pace for all.
Each stage holds a slice of the layers and passes activations to the next. At the start and end some stages are idle (the bubble), and one slow stage sets the pace of the whole pipeline.
Predict first. Does one slow stage affect the overall pipeline rate?
Chapter context
Scaling serving changes how work is placed and moved. An operation can be mathematically unchanged while its memory traffic changes dramatically. Device count alone does not describe the performance of the resulting system.
TRY IT YOURSELF
Part 1. Do it hands-on
Compute the bubble and the effect of a slow stage by hand, then write the pipeline timing model.
Each stage holds a slice of the layers and passes activations to the next. At the start and end some stages are idle (the bubble), and one slow stage sets the pace of the whole pipeline.
- To do: 4 pipeline stages and 8 microbatches in flight: idle fraction (p−1)/(m+p−1). Type it (3 decimals)
- To do: Stage times are 10, 10, 20 and 10 ms. In steady state the pipeline accepts a new microbatch every max(stage time). Microbatches per second?
- To do: Re-cut the same 50 ms of work into four equal stages of 12.5 ms. Microbatches per second?
Part 2. Do it in code
- bubble_fraction(stages, microbatches): (p − 1) / (m + p − 1).
- total_time(stage_ms, microbatches): sum(stage_ms) + (microbatches − 1) · max(stage_ms) (a synchronous pipeline in steady state).
- steady_rate(stage_ms): microbatches per second once full: 1000 / max(stage_ms).
- balance(stage_ms): re-cut the same total work into equal stages: a list of the same length with every entry sum/len.
- bottleneck(stage_ms): the index of the slowest stage (the first on ties).
- utilization(stage_ms, microbatches): useful work (sum(stage_ms) · microbatches) divided by (len(stage_ms) · total_time).
Stuck, or something not working? You can come back to it any time.
CHECK YOUR UNDERSTANDING
Does one slow stage affect the overall pipeline rate?
Chapter reference: math, worked example & code
This section belongs to the whole chapter, so it is the same on every lesson in it and may cover a different lesson's topic. The code for the experiment you just ran is inside the experiment itself.
Math and notation
A weight-traffic lower-bound estimate is
This assumes the stated bytes really traverse the bottleneck bandwidth. KV transfer adds roughly bytes_to_move/network_bandwidth before protocol and synchronization overhead. Batching and caching can change the bytes per produced token.
Worked example
A fictional 14-GB dense weight scan across a sustained 700-GB/s memory path takes at least 14/700=0.02 seconds, giving a 50-step/s bandwidth ceiling under the one-scan-per-step assumption. Real execution includes other memory traffic and work. A 512-MiB KV transfer over 32 GiB/s has a payload-only lower bound of 1/64 second.
Reference code
The following example uses the Python standard library and does not call a remote service.
from math import isclose
assert isclose(14/700,.02)
assert isclose(1/(14/700),50)
transfer=(512*1024**2)/(32*1024**3)
assert isclose(transfer,1/64)
Python examples are provided for study and local execution; this page does not execute Python.
Common mistakes
Peak device bandwidth is not guaranteed sustained application bandwidth. Exact attention can be computed with different memory layouts. Additional GPUs can add communication bottlenecks.