Tensor, pipeline & sequence parallelism

1

Concept

Fully sharded data parallelism distributes model state but still asks each rank to perform a module’s complete matrix operations after gathering parameters. Very large layers, deep stacks, or extreme sequence lengths may need the computation itself partitioned. Model parallelism is a family of layouts that split different axes.

Tensor parallelism divides a large 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 forms to avoid unnecessary gathers between adjacent operations. Collectives such as all-reduce, all-gather, or reduce-scatter restore the information needed by the next step.

Tensor parallelism reduces per-rank parameter and compute load for a layer, but communication occurs frequently inside every block. It performs best over high-bandwidth, low-latency links, often within one server. Poor placement across a slow network can spend more time combining partial matrices than multiplying them. Head counts, hidden widths, and vocabulary partitions must be divisible or padded according to the chosen layout.

Pipeline parallelism assigns consecutive groups of layers to stages. A microbatch flows through stage 0, then stage 1, and so on. To keep stages busy, the global batch is divided into several microbatches whose forward and backward passes are interleaved. Activations cross stage boundaries, while parameters remain owned by their stages.

The pipeline has bubbles: idle periods during fill, drain, or dependency waits. More microbatches can improve utilization but increase scheduling and activation-state complexity. Uneven layer cost creates imbalance, so partitioning by layer count alone may be poor. Activation checkpointing and interleaved schedules trade memory, recomputation, and bubble size.

Sequence parallelism partitions the token dimension. Some operations, such as normalization and position-wise FFNs, can process local sequence slices. Attention needs interaction across positions, so the system must exchange the right query, key, value, or head partitions. DeepSpeed-Ulysses uses all-to-all communication to transpose between sequence and head partitions, enabling each rank to compute complete attention for a subset of heads before transposing back.

These axes compose into 3D or higher-dimensional parallelism. One deployment may use data-parallel groups across replicas, pipeline stages across layer ranges, tensor-parallel ranks within each stage, and sequence sharding for long contexts. The product of group sizes must match the world size, and each rank needs deterministic coordinates in this grid.

Composition makes batch arithmetic and checkpointing harder. Data-parallel ranks average gradients; tensor ranks jointly own one layer; pipeline ranks own different layers. A checkpoint manifest must know which axis partitioned each tensor. Changing topology at resume requires resharding rather than naïvely assigning rank files.

Measure model FLOPs utilization, but also inspect collective time, exposed communication, stage idle time, memory peaks, and stragglers. A high theoretical FLOP count with long bubbles is not efficient. Profiling should reflect the real sequence and microbatch lengths, since a layout optimized for one workload can fail on another.

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 creating coordination. The best topology aligns the heaviest communication with the fastest links and keeps every stage doing useful work.