Pipeline Parallelism › Model Partitioning into Stages

The first design decision in pipeline parallelism is how to divide the model into stages. A stage is a contiguous sequence of layers that runs on a single device. The boundary between stages is called a pipeline cut point, and it determines both the computational load per device and the volume of data that must be communicated between devices at each boundary.

Choosing where to place cut points involves balancing several competing factors. If the stages are imbalanced in computational cost, the slowest stage becomes the bottleneck: faster stages finish quickly but cannot proceed until they receive either activations (during the forward pass) or gradients (during the backward pass) from the slow stage. The goal is to partition layers such that each stage takes approximately the same wall-clock time per sample. This is called load balancing, and it is more subtle than it might initially appear. Two stages with the same number of layers can have very different computational costs if one contains attention layers operating on long sequences and another contains feed-forward layers with large hidden dimensions. Similarly, the first stage that includes the embedding lookup and the last stage that includes the output projection and softmax may have compute profiles that differ from the middle transformer blocks.

For a transformer model with $L$ layers and $D$ devices, the simplest strategy assigns $L/D$ consecutive layers to each device. This achieves perfect load balance if every layer is identical in compute cost, which is approximately true for standard transformer blocks. In practice, the embedding layer and the output projection layer may have different costs than interior transformer blocks, and some architectures include routing layers (as in mixture-of-experts models) that behave very differently from standard dense layers. A profiling-driven partition, where you measure the actual wall-clock time per layer on representative inputs and then use a dynamic programming algorithm to find the partition that minimizes the maximum stage duration, is more reliable than a uniform assignment. The problem of finding the optimal partition is formally equivalent to a bin-packing variant, and it can be solved exactly in polynomial time given per-layer timing measurements.

The communication cost between stages is determined by the size of the activation tensor passed at each cut point. For a transformer with sequence length $T$, batch size $B$, and hidden dimension $H$, the tensor passed between stages has shape $(B, T, H)$, so the communication volume per forward pass is:

where:

• $B$: micro-batch size (number of sequences per micro-batch)
• $T$: sequence length in tokens
• $H$: hidden dimension (model width)
• $\text{bytes\_per\_element}$: 2 for bfloat16, 4 for float32

For a GPT-3-scale model with $H = 12288$, $T = 2048$, and $B = 4$, using bfloat16, the communication tensor has size $4 \times 2048 \times 12288 \times 2 = 201$ MB per stage boundary. This must flow across the interconnect linking the two pipeline stages. On NVLink, which provides roughly 600 GB/s of bidirectional bandwidth between GPUs in the same node, this transfer takes about 0.33 milliseconds. On InfiniBand HDR (200 Gb/s), the same transfer takes closer to 8 milliseconds. The compute time per stage for this same configuration is several hundred milliseconds on a single A100. The ratio of communication to compute time determines whether the pipeline is compute-bound or communication-bound: on NVLink it is clearly compute-bound (the interconnect is fast enough to overlap communication with computation), while on slower networks the communication overhead can become a meaningful fraction of total time.

This is typically much smaller than the communication volume in data parallelism, which involves gradient tensors proportional to total parameter count. A GPT-3-scale model with 175 billion parameters requires 350 GB of gradients (in bfloat16) per all-reduce operation in data parallelism. Compared to 201 MB per pipeline boundary per micro-batch, the pipeline communication overhead is orders of magnitude smaller. Pipeline stage-to-stage communication is also point-to-point: only two devices are involved at each boundary, and only one tensor flows per forward pass and one tensor flows per backward pass. This contrasts with the all-reduce patterns of data parallelism, which involve all devices simultaneously.

Each device in a pipeline stores the parameters and buffers for its assigned layers, plus the activation tensors needed for the backward pass. The parameter memory per device scales with the total model size divided by the number of stages:

where:

• $P_{\text{total}}$: total number of model parameters (in bytes, accounting for numerical precision)
• $D$: number of pipeline stages

This fraction is the primary motivation for pipeline parallelism: models that cannot fit on one device can be spread across multiple devices, with each device using only $1/D$ of the total parameter memory. A model with 175 billion parameters in 16-bit precision requires roughly 350 GB of memory, which exceeds the capacity of any single GPU available as of 2024. Spreading it across 8 pipeline stages reduces per-device parameter memory to about 44 GB, which fits comfortably on modern 80 GB A100 or H100 GPUs. With 16 stages, per-device parameter memory drops to about 22 GB.

Beyond raw parameter storage, the optimizer state memory must also be considered. With the Adam optimizer, each parameter requires two additional floating-point values (the first and second moment estimates). When using mixed-precision training (model weights in bfloat16 but optimizer states in float32), the optimizer state memory is larger than the parameter memory: each parameter requires 4 bytes for the float32 master weight, 4 bytes for the first moment, and 4 bytes for the second moment, totaling 12 bytes per parameter in optimizer state alone. For a stage holding $P_{\text{total}}/D$ parameters, the total memory requirement including optimizer state is:

where the factor of 2 comes from the bfloat16 model parameters and the factor of 12 comes from float32 optimizer states. For a 175 billion parameter model split across 8 stages, this gives roughly $(14 \times 350 \text{ GB}) / 8 \approx 613 \text{ GB} / 8 \approx 76.6 \text{ GB}$ per stage, which is tighter but still feasible on 80 GB GPUs.

The activation memory requires more careful analysis. During the forward pass, each stage must retain its intermediate activations until the corresponding backward pass arrives. In a naive sequential pipeline, stage $k$ must hold its activations in memory from the moment it finishes the forward pass until the backward pass through all later stages has completed and the gradient has been returned to stage $k$. This can be a long wait, especially for early stages in a deep pipeline. The key challenge therefore extends beyond the per-stage parameter and optimizer state memory, which are fixed and predictable, to the activation memory, which depends on the schedule and can vary by an order of magnitude depending on implementation choices.