← Back to Table of Contents

Chapter 26 β€” Distributed Training

β€œNo single GPU can train a frontier model. The art of distributed training is splitting computation across thousands of GPUs while minimizing communication overhead.”

GPT-4, LLaMA-3 405B, and DeepSeek-V3 were each trained on thousands of GPUs for months. This chapter covers the parallelism strategies that make this possible.

Why Distribute?

The largest models today have hundreds of billions of parameters. A single GPU can’t hold the model, optimizer states, gradients, and activations simultaneously:

Component LLaMA-3 70B (BF16) LLaMA-3 405B (BF16)
Model weights 140 GB 810 GB
Adam optimizer states 560 GB 3,240 GB
Gradients 140 GB 810 GB
Activations (B=1, T=8192) ~30 GB ~170 GB
Total ~870 GB ~5,030 GB
H100 80GB GPUs needed minimum 11 minimum 63

The Three Dimensions of Parallelism

Side-by-side comparison of Data Parallel, Tensor Parallel, and Pipeline Parallel training strategies
Three parallelism strategies β€” production training combines all three (3D parallelism)
3D Parallelism
Data Parallelism (DP)
Replicate model on each GPU. Split the batch across GPUs. All-reduce gradients after backward pass.
Tensor Parallelism (TP)
Split individual weight matrices across GPUs (column-wise or row-wise). Requires all-reduce within each layer.
Pipeline Parallelism (PP)
Assign groups of layers to different GPUs. Forward activations flow GPU-to-GPU. Reduce idle "bubble" time with micro-batches.

Frontier training combines all three β€” 3D parallelism β€” plus additional techniques like sequence parallelism and expert parallelism (for MoE models).

Data Parallelism (DDP)

The simplest strategy: each GPU has a full copy of the model and processes different data.

DDP All-Reduce Gradient Sync
GPU 0: forward + backward β†’ βˆ‡Wβ‚€
β†’
GPU 1: forward + backward β†’ βˆ‡W₁
β†’
GPU 2: forward + backward β†’ βˆ‡Wβ‚‚
β†’
All-Reduce: βˆ‡W = mean(βˆ‡Wβ‚€, βˆ‡W₁, βˆ‡Wβ‚‚)
β†’
Each GPU updates with same βˆ‡W
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18

# PyTorch DDP
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

dist.init_process_group("nccl")  # NVIDIA Collective Communications Library
local_rank = int(os.environ["LOCAL_RANK"])

model = MyTransformer().to(local_rank)
model = DDP(model, device_ids=[local_rank])

# Training loop is identical β€” DDP handles gradient sync automatically
for batch in dataloader:
    loss = model(batch)
    loss.backward()       # gradients synced via all-reduce
    optimizer.step()
    optimizer.zero_grad()

# Launch: torchrun --nproc_per_node=8 train.py

Limitation: every GPU must hold the full model + optimizer states. For a 70B model, that’s ~870 GB β€” impossible with DDP alone.

FSDP / ZeRO: Sharded Data Parallelism

Fully Sharded Data Parallel (FSDP) β€” PyTorch’s implementation of DeepSpeed ZeRO β€” shards not just data but optimizer states, gradients, and even parameters across GPUs:

ZeRO Stage What is sharded Memory per GPU (70B, 8 GPUs)
Stage 0 (DDP) Nothing β€” full replicas ~870 GB
Stage 1 Optimizer states ~210 GB
Stage 2 + Gradients ~158 GB
Stage 3 / FSDP Full + Parameters ~109 GB 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19

# PyTorch FSDP (Fully Sharded Data Parallel)
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import MixedPrecision

mp_policy = MixedPrecision(
    param_dtype=torch.bfloat16,
    reduce_dtype=torch.bfloat16,
    buffer_dtype=torch.bfloat16,
)

model = FSDP(
    model,
    mixed_precision=mp_policy,
    sharding_strategy=ShardingStrategy.FULL_SHARD,  # ZeRO-3
    auto_wrap_policy=transformer_auto_wrap_policy,
)

# Training loop remains the same
# FSDP all-gathers parameters before forward, shards after backward

The tradeoff: FSDP introduces more communication (all-gather before forward, reduce-scatter after backward) but dramatically reduces per-GPU memory.

Tensor Parallelism (TP)

Tensor parallelism splits individual weight matrices across GPUs. For a linear layer Y = XW:

Column-Parallel Linear
X [B, T, d]
β†’
GPU 0: X Γ— W₁ β†’ Y₁ [B, T, d/2]
β†˜
All-Gather β†’ Y [B, T, d]
GPU 1: X Γ— Wβ‚‚ β†’ Yβ‚‚ [B, T, d/2]
β†—

In a transformer, TP is applied to:

  • Attention: Q, K, V projections split column-wise across heads; output projection split row-wise
  • FFN: first linear split column-wise, second split row-wise

Constraint: TP requires fast interconnect (NVLink/NVSwitch) since it communicates within every layer. Typically used within a single node (8 GPUs).

Input (replicated): B Γ— T Γ— d
Weight shard (GPU k): d Γ— d/TP
Output shard (GPU k): B Γ— T Γ— d/TP
After all-gather: B Γ— T Γ— d

Pipeline Parallelism (PP)

Pipeline parallelism assigns different layers to different GPUs:

Pipeline Parallelism with Micro-batches
GPU 0 β€” Layers 0–7: [ΞΌ1]β†’[ΞΌ2]β†’[ΞΌ3]β†’[ΞΌ4] idle [ΞΌ4']β†’[ΞΌ3']β†’[ΞΌ2']β†’[ΞΌ1']
GPU 1 β€” Layers 8–15: idle [ΞΌ1]β†’[ΞΌ2]β†’[ΞΌ3]β†’[ΞΌ4] [ΞΌ4']β†’[ΞΌ3']β†’[ΞΌ2']β†’[ΞΌ1'] idle
GPU 2 β€” Layers 16–23: idle idle [ΞΌ1]β†’[ΞΌ2]β†’[ΞΌ3]β†’[ΞΌ4] β†’ backward β†’
GPU 3 β€” Layers 24–31: idle idle idle [ΞΌ1]β†’[ΞΌ2]β†’[ΞΌ3]β†’[ΞΌ4] β†’ backward β†’

The pipeline bubble (idle time) is reduced by splitting a batch into micro-batches. Schedules like 1F1B (one forward, one backward) and interleaved PP minimize the bubble.

  • Bubble fraction β‰ˆ (PP_size - 1) / n_microbatches
  • With 4 pipeline stages and 32 micro-batches: ~9% bubble overhead

Communication Patterns

Operation When Used Data Volume Requires
All-Reduce DDP gradient sync 2 Γ— model_size Inter-node OK
All-Gather FSDP param collection, TP column-parallel model_size / world_size Fast interconnect preferred
Reduce-Scatter FSDP gradient sharding model_size / world_size Inter-node OK
Point-to-Point Pipeline parallelism activations activation_size Between adjacent stages
All-to-All MoE expert routing tokens Γ— d Needs careful scheduling

Putting It All Together: 3D Parallelism

A practical configuration for training LLaMA-3 405B on 1024 H100 GPUs:

3D Parallelism Configuration Example
Data Parallelism (DP=16) β€” 16 replicas, each processes different data
Tensor Parallelism (TP=8) β€” within each node (NVLink), splits attention heads + FFN
Pipeline Parallelism (PP=8) β€” across 8 nodes, splits 126 layers into 8 stages
Total GPUs: DP Γ— TP Γ— PP = 16 Γ— 8 Γ— 8 = 1,024 H100s

DeepSpeed Integration

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20

# DeepSpeed ZeRO-3 config
ds_config = {
    "train_batch_size": 2048,
    "train_micro_batch_size_per_gpu": 2,
    "gradient_accumulation_steps": 128,
    "bf16": {"enabled": True},
    "zero_optimization": {
        "stage": 3,
        "offload_param": {"device": "none"},
        "offload_optimizer": {"device": "none"},
        "overlap_comm": True,
        "contiguous_gradients": True,
        "reduce_bucket_size": 5e8,
    },
}

# Initialize with HuggingFace Accelerate
from accelerate import Accelerator
accelerator = Accelerator(deepspeed_plugin=ds_config)
model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader)

What’s Next

Not all accelerators are NVIDIA GPUs. The next chapter surveys ASICs and alternative accelerators β€” TPUs, Gaudi, Groq LPUs, and more.

← Previous: Chapter 25 β€” CUDA & Kernel Development Β· Next: Chapter 27 β€” ASICs & Accelerators β†’

Last updated: April 2026

GitHub Β· Language Modelling β€” The Complete Guide Β· Last updated April 2026