AI ENGVisual Encyclopedia

MODULE 4: DISTRIBUTED INFRASTRUCTURE · SCENE 14

Copy everything, average the blame

Data parallelism clones the model per GPU and all-reduces gradients every step.

GPU RANK 0LOCAL BATCH GRADIENTGPU RANK 1LOCAL BATCH GRADIENTGPU RANK 2LOCAL BATCH GRADIENTGPU RANK 3LOCAL BATCH GRADIENT

Data Parallelism replicates model weights across GPUs and averages gradients via All-Reduce every step.

TECHNICAL BREAKDOWNModule 4: Distributed Infrastructure & 3D Parallelism

Data Parallelism & ZeRO State Sharding

Data Parallelism (DP) replicates model parameters across multiple GPUs, sending unique micro-batches to each rank and averaging computed gradients across ranks before the optimizer step.

Naive Data Parallelism

Every GPU keeps a full copy of model weights, gradients, and optimizer states. Scale throughput by increasing global batch size = micro_batch * DP_ranks.

ZeRO-1, ZeRO-2, ZeRO-3 Memory Sharding

ZeRO-1: Shards Adam optimizer states across DP ranks (saves 4x memory). ZeRO-2: Shards gradients across ranks (saves 2x further). ZeRO-3: Shards model parameters across ranks (all-gathers weights on-the-fly per layer).

Communication Overlap

Gradient All-Reduce and All-Gather calls are piped concurrently alongside backward pass GEMM operations to maximize GPU compute unit utilization.

MATHEMATICAL FORMULATION · ZERO-3 MEMORY PER GPU
Memory_ZeRO3 = ( 16 N / DP_ranks ) + Activation_Memory

Dividing total model static memory evenly across all data-parallel workers.

REAL-WORLD PRODUCTION ENGINEERING
  • PyTorch FSDP (Fully Sharded Data Parallel) is the open-source implementation of ZeRO-3 built into modern training stacks.
← 13 · The state you pay to keep
15 · The ring that averages the world →