Tensor, pipeline & sequence parallelism

01 · Concept

Concept

Your cluster has eight accelerators per node, so you set tensor-parallel degree to eight. It is the obvious choice: eight is the node, eight divides everything you checked, and the job launches without complaint. Six hours later the memory profile looks wrong — the key and value projections are taking the same space on every rank instead of an eighth, and the collective time per step is higher than your model predicted. Nothing failed. The framework quietly did the only thing it could.

Fully sharded data parallelism distributes model state but still asks one rank to execute a module’s complete matrix operations after gathering its parameters. When a single layer is too large to compute, when depth exceeds what one device can hold end to end, or when sequences are long enough that positions themselves must be divided, the computation has to be partitioned. Model parallelism is the family of layouts that do that, along three different axes.

Tensor parallelism divides an operation across feature dimensions. For a linear layer $Y=XW$, ranks can own columns of $W$ and produce different output slices, or own rows and sum partial results. Transformer projections are arranged in complementary column- and row-parallel pairs so no gather is needed between adjacent operations, and collectives — all-reduce, all-gather, reduce-scatter — restore whatever the next step requires. The absolute requirement is divisibility: every partitioned dimension must split evenly by the degree, or be padded until it does.

So audit this model dimension by dimension. The padded vocabulary is $248{,}320 = 512\times 485$, splitting cleanly at any power-of-two degree up to 512 — that padding was chosen for exactly this purpose, as lesson 5.3 showed. The FFN intermediate is $17{,}408 = 1024\times 17$, splitting up to 1024; at degree 8 each rank owns a $5120\times 2176$ slice of the up-projection. The 24 query heads divide by 1, 2, 3, 4, 6, 8, 12, and 24. Everything so far comfortably supports degree 8.

Then you reach the KV heads. There are four. Grouped-query attention — the serving decision from lesson 5.6 — means 24 query heads share only 4 key/value heads at head dimension 256, so the K and V projections are $5120\to 1024$ rather than $5120\to 6144$. Four divides by 1, 2, and 4, and stops. At degree 4 the split is exact: each rank owns 6 query heads and 1 KV head, a $5120\to 3072$ slice of the q_proj weight — 6 heads carrying 256 query values and 256 gate values each, because the weight is 12288 wide, not 6144 — a $5120\to 256$ slice each of K and V, and a $1536\to 5120$ row-parallel slice of the output projection. At degree 8 there is no exact split, and frameworks fall back to replicating KV heads across pairs of ranks. The job runs, the mathematics is correct, and you have paid for duplicated K/V storage and additional communication in exchange for nothing.

The Gated DeltaNet layers have their own head structure — 48 value heads and 16 query/key heads, each of dimension 128 — so their clean divisor is 16 for the QK path. Since 48 of the 64 layers are DeltaNet, a layout has to satisfy both structures at once, and again the smallest divisor wins.

Pipeline parallelism partitions depth instead. Consecutive groups of layers become stages; a microbatch flows through stage 0, then stage 1, and so on, with activations crossing stage boundaries and parameters staying put. Here the architecture is unusually cooperative. The 64 layers are not a uniform stack but 16 repeating super-blocks, each three DeltaNet layers followed by one full-attention layer. Any pipeline degree dividing 16 — 2, 4, 8, or 16 — gives stages with identical layer composition. Pipeline degree 4 puts 4 super-blocks, 16 layers, on each stage, and every stage sees exactly 12 DeltaNet layers and 4 attention layers.

Split at a degree that does not divide 16 and the balance breaks: a boundary drawn inside a super-block leaves one stage holding more full-attention layers than another. Since attention and DeltaNet layers have different compute profiles and different memory behaviour with sequence length, unequal composition means unequal stage time, and the slowest stage sets the pace for all of them.

Pipelines also have bubbles — idle periods during fill, drain, and dependency waits. With $P$ stages and $m$ microbatches, the classic schedule’s bubble fraction is about

of ideal time lost. More microbatches shrink it and increase activation state. Interleaved schedules and activation checkpointing (lesson 5.10) trade memory, recomputation, and bubble size against each other.

Sequence parallelism partitions the token dimension. Normalization and position-wise FFNs process local slices happily. Attention needs interaction across positions, so the system exchanges query, key, value, or head partitions; DeepSpeed-Ulysses uses all-to-all to transpose between sequence and head partitions, letting each rank compute complete attention for a subset of heads before transposing back. This axis matters most for a model advertising 262,144 native context, and note that only the 16 attention layers need the exchange — the DeltaNet layers scan positions with a fixed-size recurrent state.

These axes compose into 3D or higher-dimensional layouts: data-parallel groups across replicas, pipeline stages across layer ranges, tensor-parallel ranks within a stage, sequence sharding for long contexts. The product of group sizes must equal the world size, and each rank needs deterministic coordinates in that grid. Composition also complicates checkpointing — the manifest must record which axis partitioned each tensor, and changing topology at resume requires resharding rather than reassigning rank files.

Measure model FLOPs utilization, but also collective time, exposed communication, stage idle time, memory peaks, and stragglers. A high theoretical FLOP count behind long bubbles is not efficiency. Profile at the real sequence and microbatch lengths, because a layout tuned for 4,096-token sequences can behave very differently at 262,144.

The durable map is by axis: data parallelism splits examples, tensor parallelism splits feature computation, pipeline parallelism splits depth, sequence parallelism splits positions. Each makes a larger problem fit by manufacturing coordination. The best topology aligns the heaviest communication with the fastest links, keeps every stage doing useful work, and — the lesson of the 4 KV heads — respects the divisor the architecture actually gives you rather than the one the hardware suggests.