Data Parallelism and Gradient Synchronization is the foundational distributed training approach where identical model replicas process different data samples, aggregate gradients across replicas, and synchronously apply updates to maintain training consistency.
Data Distributed Parallel (DDP) in PyTorch
- DDP Architecture: Each GPU runs independent data loader, processes batch, computes gradients. Gradients collected via all-reduce, averaged, applied to local model.
- Backward Hook Integration: PyTorch hooks gradient computation, automatically triggers all-reduce upon backward pass completion. Transparent to user code.
- Communication Overhead: All-reduce requires 2× gradient size bandwidth (send + receive). For 1B parameter models, ~8 GB all-reduce per iteration.
- Synchronous Training: All replicas coordinate at gradient application. Stragglers (slower GPUs) block fastest GPUs, reducing effective throughput (synchronized by slowest device).
ZeRO (Zero Redundancy Optimizer) Stages
- ZeRO Stage 1 (Gradient Partitioning): Gradients partitioned across GPUs. GPU i stores gradient partitions [i×n:(i+1)×n]. Reduces gradient memory by factor of N_gpus.
- ZeRO Stage 2 (Gradient + Optimizer State Partitioning): Optimizer state (momentum, variance) also partitioned. Memory reduction: 4-6x (for Adam: 2 gradient copies + 2 momentum + 2 variance).
- ZeRO Stage 3 (Parameter Partitioning): Model weights themselves partitioned. GPU i stores subset of weights. Requires weight broadcast before forward pass (communication overlapped with computation).
- ZeRO-Offload: Optimizer state offloaded to CPU. Reduces GPU memory but requires PCIe bandwidth for state updates (typically 10-20 GB/s). Viable for CPU-rich systems.
Gradient Compression Techniques
- PowerSGD: Rank-reduced low-rank approximation of gradient matrices. Compresses gradients 10-100x with <1% convergence slowdown. Requires extra computation (SVD).
- 1-bit Adam: Quantize gradients to 1-bit per parameter (sign bit only) with momentum compensation. 32x compression but requires careful learning rate tuning.
- Top-K Sparsification: Only communicate top-K gradient values (largest magnitude). Reduces communication 10-100x for sparse gradient models (certain domains like NLP).
- Error Feedback/Momentum Correction: Quantization error accumulated in momentum buffer, compensated in future updates. Prevents convergence degradation from compression.
All-Reduce Communication Patterns
- Ring All-Reduce: Logical ring of N GPUs, gradients passed sequentially. Bandwidth-efficient (uses full link utilization) but latency = O(N).
- Tree All-Reduce: Binary tree minimizes latency O(log N) but underutilizes bandwidth in over-subscribed networks. Cadence slower than ring for large clusters.
- Hybrid Approaches: Two-level hierarchies combine benefits. Intra-rack tree, inter-rack ring. Typical cluster topology shapes algorithm selection.
- Pipelined All-Reduce: Partition gradients into chunks, stream chunks through reduction pipeline. Overlaps communication phases across multiple GPUs.
Overlap of Backward Pass with All-Reduce
- Bucket-Based Gradient Accumulation: Gradients accumulated in buckets (e.g., 25 MB each). Upon bucket completion, all-reduce triggered immediately (not waiting for full backward pass).
- Pipelined All-Reduce: Multiple all-reduces in-flight concurrently. GPU 0 all-reduces bucket 0 while GPU 1 backward-passes bucket 1, GPU 2 computes bucket 2 forward.
- Communication Cost Amortization: Gradient computation (~70% of backward cost), all-reduce (~20-30%), gradient application (~5%). Overlap hides ~80% of all-reduce latency.
- Network Saturation: Full overlap requires sufficient computation between synchronization points. Bandwidth-limited clusters struggle to hide all-reduce even with pipelining.
Gradient Synchronization and Convergence
- Synchronization Semantics: All replicas must see identical gradient sums before parameter updates. Asynchronous approaches (parameter server) degrade convergence.
- Variance Reduction: Synchronous averaging reduces variance in stochastic gradient. Larger effective batch size (N_gpu × batch_size_per_gpu) → lower gradient variance.
- Learning Rate Scaling: Learning rate typically increased proportionally to batch size. 10x larger batch_size → 10x higher learning rate (with linear scaling rule).
- Communication Cost vs Convergence: Trade-off between communication frequency (more frequent sync) and gradient staleness (less frequent sync). Optimal sync interval depends on model, batch size, cluster size.