Tensor parallelism

Read it in Chapter 4 · checkpoint +50 XP →

Split each layer's matrices across GPUs (Megatron-style). Huge models fit, but every layer needs collectives.

Tensor parallelism (TP) shards the weight matrices inside each layer. In the Megatron-LM scheme, the first MLP matmul is split by columns and the second by rows, so a layer needs only one all-reduce in forward and one in backward. With attention handled the same way, that's ~4 activation all-reduces per transformer layer per micro-batch.

Trade-offs

• ✅ Weights, gradients, optimizer state and activations shrink by the TP degree.
• ❌ Communication is frequent, latency-sensitive and on the critical path. It mostly can't be hidden.
• ❌ Shards get narrow: hidden/TP too small means small GEMMs and lower Tensor Core efficiency.

Rule: keep TP inside the NVLink domain (TP ≤ 8 on HGX servers). TP=16 across InfiniBand & RoCE is a classic way to waste a cluster.

Use it in a mission

Related

Sources & further reading

• Megatron-LM (Shoeybi et al., 2019)
• Distributed Training (T. Tsugriy) — parallelism strategies, collectives, memory

Written for LazyLab in our own words; sources are cited for depth and verification.