Data vs Model Parallelism

Why Distributed Training is Necessary

Training a 70B or 405B parameter model requires trillions of FLOP calculations and hundreds of gigabytes of VRAM.

A single GPU (e.g. 80GB NVIDIA H100) cannot hold the model, let alone train it in a reasonable timeframe.

Distributed Training partitions workloads across clusters of hundreds or thousands of GPUs.

┌──────────────────────────┬──────────────────────────┐
│ 1. DATA PARALLELISM (DDP)│ 2. MODEL PARALLELISM     │
├──────────────────────────┼──────────────────────────┤
│ Model fits on 1 GPU.     │ Model is TOO BIG for 1 GPU!│
│ Replicate model on N GPUs│ Split model weights      │
│ Split DATA BATCH across  │ across N GPUs (Tensor or │
│ GPUs. Sync gradients!    │ Pipeline Parallelism).   │
└──────────────────────────┴──────────────────────────┘

1. Distributed Data Parallelism (DDP)

Used when the model fits inside a single GPU VRAM, but dataset training speed needs acceleration across $N$ GPUs.

  GPU 0: [ Full Model Copy ] ──► Process Batch 0 ──► Gradients 0 ──┐
  GPU 1: [ Full Model Copy ] ──► Process Batch 1 ──► Gradients 1 ──┼──► [ ALLREDUCE ] ──► Average Gradients ──► Step!
  GPU 2: [ Full Model Copy ] ──► Process Batch 2 ──► Gradients 2 ──┘

1. Replicate identical model weights on all $N$ GPUs.
2. Split global data batch into $N$ distinct mini-batches (e.g. 32 samples per GPU).
3. Execute forward and backward passes independently on each GPU.
4. Execute an AllReduce Collective Communication step: GPUs communicate over high-speed NVLink / InfiniBand to average gradients across all workers.
5. All GPUs execute identical optimizer weight updates!

2. Model Parallelism (Tensor & Pipeline)

Used when a single model is too large to fit in 1 GPU VRAM.

A. Tensor Parallelism (Megatron-LM Style - Shoeybi et al., 2019)

Splits individual weight matrices across multiple GPUs within the same layer.

For a matrix multiplication $Y = X W$:

• Column-Parallel Linear: Split $W$ vertically into $[W_1 \mid W_2]$. GPU 0 computes $X W_1$, GPU 1 computes $X W_2$.
• Row-Parallel Linear: Split $W$ horizontally into $\begin{bmatrix} W_1 \ W_2 \end{bmatrix}$. GPU 0 computes $X_1 W_1$, GPU 1 computes $X_2 W_2$, followed by an AllReduce sum.

Used within a single GPU server node connected by high-speed NVLink ($900\text{ GB/s}$).

B. Pipeline Parallelism (GPipe - Huang et al., 2019)

Splits layers sequentially across GPUs:

• GPU 0 holds Layers 1 to 20.
• GPU 1 holds Layers 21 to 40.
• GPU 2 holds Layers 41 to 60.

To avoid idle GPU waiting time (Pipeline Bubbles), 1F1B (One Forward, One Backward) scheduling splits micro-batches so GPUs process forward and backward activations concurrently.

┌──────────────────────────┬──────────────────────────┬──────────────────────────┐
│ PARALLELISM TYPE         │ SPLITS WHAT?             │ COMMUNICATION NETWORK    │
├──────────────────────────┼──────────────────────────┼──────────────────────────┤
│ Data Parallelism (DDP)   │ Data Mini-Batches        │ AllReduce (Gradients)    │
│ Tensor Parallelism (TP)  │ Weight Matrices          │ Intra-Node (NVLink)      │
│ Pipeline Parallelism (PP)│ Model Layers             │ Inter-Node (InfiniBand)  │
└──────────────────────────┴──────────────────────────┴──────────────────────────┘

3D Parallelism (The Frontier Standard)

Training 500B parameter frontier LLMs combines all three techniques simultaneously:

$$\text{Total GPUs} = \text{Data Parallel (DP)} \times \text{Tensor Parallel (TP)} \times \text{Pipeline Parallel (PP)}$$

For example, 1,024 GPUs configured as $\text{DP}=16 \times \text{TP}=8 \times \text{PP}=8$.

Say this out loud

Distributed Data Parallelism (DDP) replicates model weights across GPUs, splitting data batches and synchronizing gradients via AllReduce. Model Parallelism splits large model parameters across GPUs when models exceed 1 GPU VRAM capacity, using Tensor Parallelism intra layer via matrix splitting and Pipeline Parallelism inter layer via sequential layer assignment.

Followups to expect

1. What is Ring-AllReduce? An efficient AllReduce algorithm where GPUs transfer gradient blocks in a logical ring, achieving optimal network bandwidth utilization independent of GPU count $N$.
2. Why is Tensor Parallelism limited to intra-node GPUs? Tensor Parallelism requires high-frequency AllReduce communications after every single Multi-Head Attention layer, requiring low-latency NVLink interconnects available only within a single physical server chassis.