Pipeline and expert parallelism
Read a pipeline-stage timeline and an expert token-count distribution to locate structural idle time, stage imbalance, and routing skew
Before you begin
Follow a linked prerequisite when you cannot explain it from memory. Background skills without links are stated as abilities rather than hidden terminology.
A 48-layer transformer does not fit on one eight-GPU server even after tensor sharding. Engineers assign layers 0–11 to one group of GPUs, layers 12–23 to a second group, and so on. A microbatch must pass through the four groups in order during forward computation and return in reverse order during backward computation. A timeline initially shows the later groups idle because no activation has reached them.
Each layer group is a pipeline stage. The stage sends its output activation to the next stage and later receives an activation gradient from that stage. Splitting the global batch into microbatches allows stage 0 to process microbatch 2 while stage 1 processes microbatch 1. A stage timeline produced with framework ranges and GPU profiler events shows fill, steady execution, drain, transfer waits, and imbalance as distinct horizontal gaps.
A mixture-of-experts layer creates a different scheduling problem. A router selects destination experts for each token. Software packs token activations by destination, performs an all-to-all exchange, runs the local expert matrices, and performs another all-to-all to return outputs. Per-expert token counters and grouped-matrix timings reveal routing skew; the busiest expert or destination rank can delay all peers even when the mean token count looks balanced.
Pipeline parallelism names ordered layer partitioning, while expert parallelism names partitioning of conditionally selected experts. Both methods require a schedule model with concrete microbatch counts, stage times, token distributions, transfer sizes, and memory lifetimes. Average compute alone cannot predict the slowest stage or the hottest expert.
Follow one microbatch through pipeline stages
Partition an ordered model into P stages. Stage 0 receives a microbatch and sends its output activation to stage 1. Each stage continues forward until the final stage produces loss. During backward, activation gradients travel in the reverse direction and each stage computes parameter and input gradients.
One microbatch cannot use all stages at once at the beginning. Stage 1 waits for stage 0, and later stages remain idle. Multiple microbatches allow stage 0 to process a later microbatch while stage 1 processes an earlier one, subject to the selected forward/backward schedule.
Derive the forward-only bubble
Assume equal stage time and M microbatches. A forward-only pipeline completes in M + P−1 stage slots. Each stage performs useful work in M slots, so ideal utilization is M/(M+P−1). The P−1 additional slots represent fill and drain.
The formula is an intuition, not a complete training model. GPipe-style schedules run all forward microbatches and then backward, while 1F1B schedules interleave forward and backward after warm-up. Backward times, transfer, recomputation, and activation storage change the exact bubble and memory behavior.
Balance stage cadence
Steady pipeline throughput cannot exceed the slowest stage cadence. A stage with 25 milliseconds of work limits progress even if every other stage takes 18 milliseconds. Faster stages accumulate idle intervals around the slow one.
Partition by measured forward, backward, recomputation, and transfer time, not parameter or layer count alone. Attention cost can vary by sequence and layer type. Embedding and output layers can have unique operations. Moving a boundary changes activation bytes and memory ownership as well as compute.
Trade microbatches against activation memory
Increasing M amortizes structural fill and drain. A larger microbatch count can also increase the number of activations retained before their backward computation. Schedule and activation checkpointing determine the exact live set.
If global batch is fixed, more microbatches mean smaller microbatch size, which can reduce GEMM efficiency. If microbatch size is fixed, global batch grows. State which quantity remains fixed and validate convergence.
Follow expert dispatch
For a mixture-of-experts layer, a router selects one or more experts for each token. Group tokens by destination rank, exchange token activations through all-to-all, run local expert MLPs, and exchange outputs back to their source positions. Packing and unpacking are part of the path.
The local expert batch is the number of routed tokens for that expert, not the global token count. Sparse routing can produce many small, irregular GEMMs. Grouped execution can combine work, but routing skew and capacity decisions remain.
Model routing imbalance and capacity
Average tokens per expert do not determine the step. The maximum loaded rank or expert can delay the return all-to-all and later layers. Record per-expert and per-rank token distributions, including tails.
A capacity factor allocates room above average. Overflow policy can drop tokens, reroute them, or use another path. Drop, reroute, and fallback policies affect model semantics and systems cost. A performance comparison must keep the routing and quality contract fixed.
Map transfers to topology
Adjacent pipeline stages exchange activation and gradient tensors at microbatch cadence. Expert groups perform all-to-all with broad peer concurrency. Their preferred topology can differ: pipeline traffic values neighbor paths, while expert traffic values bisection bandwidth.
When both dimensions share NICs or network cuts, concurrent operations interfere. Calculate message frequency and bytes for both, then inspect the scheduled overlap on the physical route.
Preserve training semantics
Pipeline schedules determine when a weight version contributes to each microbatch and when accumulated gradients become eligible for an optimizer step. Expert routing determines which parameters process each token and what happens to overflow. Changes that alter these rules can change optimization even when device throughput improves.
Record weight-update cadence, loss scaling, global batch construction, router policy, dropped-token rate, and random state. Compare quality and convergence under the same declared contract.
Worked example: Identifying the stage that limits iteration time
A four-stage training pipeline uses 16 microbatches. Including forward, backward, and transfers under the chosen schedule, stages 0, 1, and 3 have an 18-millisecond cadence while stage 2 requires 25 milliseconds.
- Mark each stage's forward and backward slots for all 16 microbatches. Steady issue cannot proceed faster than stage 2's 25-millisecond cadence, so other stages develop recurring idle intervals beyond structural fill and drain.
- Decompose stage 2. Stage 2 contains one additional attention layer and recomputes a checkpointed multilayer perceptron during backward. The stage's activation receive is only 1 millisecond and is not the main difference.
- Move one MLP block from stage 2 to stage 1. Recalculate both stages' parameter memory, saved activations, recomputation, and boundary activation bytes. Stage 1 remains within memory headroom.
- Measure new cadences: stage 1 rises to 21 milliseconds and stage 2 falls to 20 milliseconds. The maximum falls from 25 to 21 milliseconds, improving steady throughput.
- Compare 8, 16, and 32 microbatches. More microbatches reduce structural bubble fraction but increase retained activation state and reduce per-microbatch matrix size when global batch remains fixed.
- Validate the complete schedule, loss, global batch, checkpoint state, and failure recovery. Stage reassignment changes checkpoint ownership and restore mapping.
Conclusion: Pipeline throughput follows the maximum stage time. Balancing the mean number of layers is not sufficient when their real costs differ.
Common errors
- Calling every idle interval a structural pipeline bubble. Separate fill/drain from imbalance, delayed transfer, host gaps, and scheduler error.
- Balancing stages by parameter or layer count. Executed kernels, sequence shapes, recomputation, and transfer determine time.
- Increasing microbatches without accounting for activation memory, local GEMM size, global batch, and latency.
- Benchmarking expert parallelism with uniform routing only. Production skew and overflow policy determine straggler and quality behavior.
Reference summary
Pipeline parallelism assigns ordered layer ranges to stages. A microbatch’s activation moves from one stage to the next during forward, and its activation gradient moves back during backward. Multiple microbatches permit different stages to execute concurrently, but fill, drain, scheduling dependencies, and unequal stage times create idle intervals.
Expert parallelism assigns different mixture-of-experts weights to ranks. The router determines a destination expert for each selected token. An all-to-all dispatches token activations, local grouped GEMMs execute the experts, and a second all-to-all returns outputs. Routing skew makes the most heavily loaded expert or rank determine the synchronization time.
forward-only pipeline utilization = M / (M + P − 1)- Balance stage time using measured forward, backward, recomputation, and transfer cost. Equal layer counts do not guarantee equal stage time.
- More microbatches amortize fill and drain but can retain more activation state or change global-batch policy.
- Map adjacent stages according to activation and gradient transfer bytes, stage cadence, and available links.
- Measure expert token-count distributions and capacity behavior. Uniform synthetic routing does not represent a skewed production router.
- Group small expert matrices when possible, but include packing, dispatch, and imbalance in the end-to-end result.
Knowledge check
What happens to pipeline utilization and activation memory as microbatch count increases?
Show the answer guide
- Increasing microbatch count amortizes the P−1 forward-only fill and drain slots, so the ideal structural bubble fraction falls under the equal-stage model.
- Many training schedules retain activations until the matching backward work, so more in-flight microbatches can increase peak activation memory.
- Under fixed global batch, more microbatches also shrink each microbatch and can reduce GEMM efficiency; under fixed microbatch size, global batch changes and must be treated as a different workload.
Draw a pipeline schedule
Draw stage-by-stage forward and backward schedules for a four-stage pipeline at 4, 8, and 16 microbatches using one declared schedule. Assign measured or representative stage times and activation bytes.
Evidence to keep: Slot diagrams, structural fill/drain calculation, stage-utilization and peak-live-activation estimates, slowest-stage cadence, and one predicted effect of repartitioning or changing microbatch count.
The parallel mesh composes these dimensions and maps their characteristic communication onto the physical hierarchy.
Sources and further reading
The lesson explanations synthesize primary papers and official references. Follow the originals for proofs, exact APIs, architecture details, and version-sensitive behavior.