Composing a Parallelism Layout From a Memory and Bandwidth Budget

A cluster of $D$ devices offers four ways to divide a training job: replicate the model and split the batch (data parallelism), split each layer’s tensors (tensor parallelism), split the layers into stages (pipeline parallelism), and shard the replicated state (ZeRO-style partitioning). The dimensions multiply to $D$, so a choice of one constrains the others, and each has a different communication signature: data parallelism all-reduces gradients of the full model size, tensor parallelism all-reduces activations inside every layer, pipeline parallelism sends boundary activations between stages, and sharding gathers parameters.

Choosing the sizes is a constrained optimisation, and the constraints are concrete enough to write down.

The three constraints

Memory. For each device, the sum of sharded parameter storage, optimiser state, gradient buffer, gathered working set, and activation memory must be below capacity with margin for the allocator. The memory accounting supplies the per-parameter terms; the layout determines how each is divided.

Bandwidth. Each dimension’s communication must fit inside the compute time it overlaps with. The relevant bandwidth is not the average link rate but the rate of the slowest link in the process group that dimension uses. A tensor-parallel group spanning two nodes pays cross-node latency on every layer.

Batch. The total batch must be divisible by the data-parallel degree, and the per-device micro-batch must be large enough that the matmuls stay compute-bound. Pipeline depth and micro-batch count then determine the bubble, as the pipeline schedule analysis shows.

The dimensions are conventionally written as a product $d \cdot t \cdot p = D$ for data, tensor, and pipeline degrees, with the shard degree folded into the data dimension because both partition the same replicated axis. Frameworks differ on whether sharding counts separately, but the arithmetic is the same: the data-parallel group of size $d$ is further divided into shard groups.

Dividing the state

Take a model with $N$ parameters, bf16 compute weights, fp32 master weights, fp32 gradients, and fp32 Adam moments, for 18 bytes per parameter before activations.

Pipeline splitting divides parameters and all per-parameter state by $p$, because each stage holds a distinct set of layers. Tensor parallelism divides the parameters within a layer by $t$, but it does not divide the optimiser state for the sharded tensors uniformly: a column-parallel linear layer keeps a full input and a sharded output, so its weight is split by $t$, while the row-parallel layer that follows splits its input and keeps a full output, so its weight is also split by $t$. The state follows the weights, so both are divided.

Sharding divides whatever remains by the shard group size $s$, where $s = d / (\text{replication within the data group})$; in the simplest configuration the whole data dimension is a shard group and $s = d$.

So the per-device state is approximately $18N/(p \cdot t \cdot s)$ bytes, with $s$ equal to the data-parallel degree when full sharding is used. Since $d \cdot t \cdot p = D$ and $s = d$, the expression becomes $18N/D$: the state is divided across every device regardless of how the dimensions are arranged. That is the useful invariant. The arrangement does not change the total, only the communication cost of accessing it.

What the arrangement does change

Since the state division is arrangement-independent, the layout is chosen to minimise communication subject to memory and topology. The communication volumes per step per device, from the earlier analyses, are:

• Data parallelism: gradient all-reduce of $2N$ per device, reduced by sharding to a reduce-scatter plus parameter all-gather of comparable volume.
• Tensor parallelism: two all-reduces of the activation tensor per layer, forward and backward, of size $b \cdot S \cdot H$ each.
• Pipeline parallelism: one boundary tensor of $b \cdot S \cdot H$ per micro-batch per stage boundary, forward and backward.

The tensor-parallel term is per layer and therefore multiplied by the layer count; the pipeline term is per boundary and multiplied by the number of micro-batches. For a model with $L$ layers split into $p$ stages of $L/p$ layers each, the tensor-parallel traffic is proportional to $L \cdot bSH$ while the pipeline traffic is proportional to $M \cdot bSH$ where $M$ is the micro-batch count. Tensor parallelism is the more expensive dimension by roughly the ratio of layers to micro-batches, which for typical configurations is an order of magnitude.

This is the reason tensor parallelism is confined to the fastest interconnect available, usually within a node, and its degree is kept small. It is also why tensor parallelism reduces activation memory: the activations inside a layer are divided by $t$, which is the only dimension that reduces per-layer activation storage.

A worked allocation

Suppose $D = 64$ devices, each with 80 GB, and a model with $N = 70$ billion parameters at 18 bytes per parameter. Total state is 1260 GB, or about 20 GB per device if evenly divided. That leaves roughly 60 GB per device for activations, workspace, and the gathered parameter window, before any margin.

The arrangement follows from the topology. Assume eight devices per node with a fast intra-node fabric and a slower inter-node link. Tensor parallelism is bounded by the intra-node degree, so $t \le 8$; take $t = 8$, keeping every tensor-parallel collective inside a node. The remaining $D/t = 8$ devices are split between pipeline and data. A pipeline of $p = 4$ and data of $d = 2$ is one option; $p = 2$, $d = 4$ is another.

The pipeline choice is set by the bubble, the number of boundary transfers, and activation residency. With $p = 4$ and $M$ micro-batches the idle fraction is $3/(M+3)$, so $M = 32$ gives about 9 percent. With $p = 2$ the fraction is $1/(M+1)$, so $M = 16$ already gives about 6 percent, and the step needs one boundary transfer instead of three. What $p = 2$ costs is residency: each stage holds the layers for 35 billion parameters, so the activation set of a stage spans twice as many layers as a stage of the $p = 4$ layout. Per-device state is the same 19.7 GB in both arrangements, so the decision turns on whether the larger per-stage activation set fits alongside it.

The activation constraint decides. Per-device activation memory scales with the micro-batch size, the sequence length, and $1/t$ for the intra-layer part, and with the number of resident micro-batches $M$ for the pipeline part. Raising $M$ to shrink the bubble raises activation memory linearly, so the two cannot both be pushed. A layout that needs a large $M$ to control the bubble must reduce the micro-batch, which then risks falling below the arithmetic-intensity floor.

def pipeline_bubble(p, m):
    return (p - 1) / (m + p - 1)

def per_device_state(n_params, bytes_per_param, world):
    return n_params * bytes_per_param / world

# 70B parameters, mixed precision with fp32 master weights and Adam moments.
state = per_device_state(70e9, 18, 64)          # ~19.7 GB
for p in (2, 4, 8):
    for m in (8, 16, 32, 64):
        b = pipeline_bubble(p, m)
        print(p, m, round(b, 3))

The output is a table of idle fractions that makes the trade visible before any run is attempted: at $p = 8$, $M = 8$ the pipeline is nearly half idle, and only $M = 64$ brings it under 10 percent.

Where the shard degree sits

Full sharding of the data-parallel group divides the state by $d$, which is already counted in the $18N/D$ invariant. Partial sharding, where only the optimiser state is partitioned and gradients and parameters stay replicated, is the stage-1 configuration described in the ZeRO stage analysis. It is the right choice when the state fits with that division and the extra all-gather of parameters is not worth the saving.

The decision is quantitative: stage 1 costs one parameter all-gather per step; stage 3 costs two per step plus the release machinery. For a model whose state fits under stage 1, the extra gather buys nothing, so stage 1 is correct. For a model that does not fit, the question is whether pipeline splitting or full sharding is cheaper, and the answer depends on the interconnect: sharding’s gathers are of the same order as pipeline boundary transfers, but a pipeline also divides the model into stages that may be unbalanced.

Failure modes of a composed layout

Unbalanced pipeline stages. A stage holding the embedding table has a large parameter count and little compute; a stage holding uniform transformer blocks has the opposite. The pipeline runs at the rate of its slowest stage, and a layout that looks balanced in parameter count can be badly balanced in time.

A tensor-parallel group crossing a slow link. If $t$ exceeds the intra-node device count, every layer’s two all-reduces traverse the inter-node fabric. The traffic is $L$ times the boundary traffic of a pipeline, so the step becomes latency-bound in a way that no amount of pipeline tuning repairs.

A micro-batch below the intensity floor. Pipeline and tensor splitting both reduce the per-device batch. If the resulting matmul has too few rows to amortise the weight load, the step is bandwidth-bound and adding devices reduces efficiency, as the arithmetic intensity discussion predicts.

Checkpoint layouts tied to the arrangement. A sharded checkpoint records the layout it was written under. Changing $p$, $t$, or $s$ between runs requires reshaping, and a checkpoint written under a layout with a different layer-to-stage assignment cannot be loaded by simple concatenation.

The last of these is the one with the longest tail. The layout is not only a runtime configuration; it is part of the format of everything the run produces, and a change to it is a migration rather than a flag.