AI ENGVisual Encyclopedia

MODULE 4: DISTRIBUTED INFRASTRUCTURE · SCENE 16

Slice the matrix itself

Tensor parallelism splits every matmul across GPUs and syncs twice per layer.

MEGATRON INTRA-NODE SLICING (NVLINK ONLY)

Column-parallel and Row-parallel matrix splits require 2 All-Reduces per transformer layer.

TECHNICAL BREAKDOWNModule 4: Distributed Infrastructure & 3D Parallelism

Tensor Parallelism (Megatron-LM Intra-Node Slicing)

When a single transformer layer's weight matrices exceed single-GPU VRAM or require ultra-fast GEMM acceleration, Tensor Parallelism (TP) splits linear layers across GPUs within a single NVLink node.

Column-Parallel & Row-Parallel GEMMs

In MLP layers: First linear layer (Gate/Up) is split column-wise: Y_i = X * W_i. Second linear layer (Down) is split row-wise: Z = sum(Y_i * V_i).

All-Reduce Synchronization

A Megatron transformer block requires exactly 2 All-Reduce communications per layer: one after self-attention row-parallel projection, and one after MLP row-parallel projection.

Strict NVLink Domain Constraint

Because TP executes All-Reduces on every single layer forward/backward pass, it requires ultra-low latency intra-node NVLink (900 GB/s on H100), limiting TP size to 8 GPUs per node.

MATHEMATICAL FORMULATION · ROW-PARALLEL MATRIX MULTIPLICATION SUMMATION
Y = ∑_(i=1..TP) ( X_i · W_i ) ==> All-Reduce-Sum( Y )

Row-parallel outputs are local partial sums that require an All-Reduce sum across TP ranks to reconstruct full activation tensors.

REAL-WORLD PRODUCTION ENGINEERING
  • Megatron-LM and vLLM use Tensor Parallelism 2, 4, or 8 to fit 70B+ parameters across 8-GPU HGX nodes.