Distributed Parallel Training

Tensor Parallelism

We split on the last dimension.

each rank gets part of each layer, transfer all data/activations.

This is easy for MLP, but for attention it’s hard?

https://huggingface.co/docs/text-generation-inference/en/conceptual/tensor_parallelism

For a dense layer $Y=XW$ with $W\in {\mathbb{R}}^{D\times D}$, shard $W$ column-wise across $N$ GPUs: $W=[W_1\mid W_2\mid \dots \mid W_N]$ with $W_i\in {\mathbb{R}}^{D\times D/N}$. Broadcast $X$ to every GPU. GPU $i$ computes $Y_i=X{W}_i\in {\mathbb{R}}^{N\times D/N}$ — its shard of the output. Block shapes: $X$ is $[1\times 1]$, $W$ is $[1\times 4]$, $Y$ is $[1\times 4]$ (for 4-way TP).

The two-layer trick (no intermediate comm). If layer 1 outputs $Y=XW$ and layer 2 computes $Z=YU$, the naïve approach gathers $Y$ from all GPUs then broadcasts. Instead: shard $W$ column-wise (each GPU owns $Y_i$) and shard $U$ row-wise (each GPU owns $U_i$). Then $Z=YU={\sum }_{i=1}^N{Y}_i{U}_i$ Each GPU computes one partial product $Y_i{U}_i$ locally — no communication after $XW$. The $N$ partial sums are combined with a single all_reduce at the end. This is exactly the pattern used in the MLP block of a Transformer (up-proj column-sharded, down-proj row-sharded) and in the QKV→out pattern of attention.

TP is typically capped at the size of a single NVLink domain (8 GPUs on an H100 DGX), since each block still requires an all-reduce.

Source

CS231n 2025 Lec 11 slides ~134–143 (TP as Dim split, column sharding block diagram, 4-way TP two-layer column/row trick, “no need for communication after XW=Y”). 2026 PDF not published — using 2025 fallback.

Related

• Model Parallelism
• Pipeline Parallelism
• FSDP
• MFU