Data Parallelism

Table of Contents

1. Learning Objectives
2. Prerequisites
3. Notation
4. Core Intuition
5. Single-Device Training Recap
6. Data Parallelism Formulation
7. Gradient Aggregation and AllReduce
8. Synchronous Distributed SGD
9. Communication Complexity
10. Scaling Efficiency and Amdahl's Law
11. Effective Batch Size and Learning Rate
12. Worked Examples
13. Connection to the Broader Curriculum
14. Common Pitfalls and Misconceptions
15. Research Perspective
16. Summary of Takeaways
17. Exercises

Learning Objectives

After reading this chapter, you should be able to:

1. Define data parallelism: replicate model, partition data across $K$ devices.
2. Derive synchronous gradient averaging: $\bar{\mathbf{g}} = \frac{1}{K}\sum_{k=1}^K \mathbf{g}_k$.
3. Explain AllReduce as the communication primitive for gradient aggregation.
4. Analyze communication volume $O(K \cdot |\boldsymbol{\theta}|)$ vs. ring AllReduce $O(|\boldsymbol{\theta}|)$.
5. Apply Amdahl's law to parallel training efficiency.
6. Relate effective batch size to learning rate scaling rules.

Prerequisites

• Gradient Descent — SGD update
• Backpropagation — gradient computation

Notation

• $K$ — Number of parallel workers (GPUs)
• $\mathcal{B}_k$ — Mini-batch on worker $k$
• $\mathbf{g}_k$ — Local gradient on worker $k$
• $\bar{\mathbf{g}}$ — Averaged global gradient
• $\boldsymbol{\theta}$ — Shared model parameters

Core Intuition

Training large models requires multiple GPUs. Data parallelism is the simplest strategy: replicate the full model on each device, split the mini-batch across devices, compute local gradients, average gradients, and apply synchronized update.

This is equivalent to Gradient Descent on a larger effective batch — but requires communication (AllReduce) that can become the bottleneck at scale.

Series context. Volume IV, Chapter 19 (Parallelism). Orthogonal to Flash Attention (memory) and Quantization (precision).

Single-Device Training Recap

Definition 1 (SGD Update).

Each step: forward + backward on batch $\mathcal{B}$, then update.

Data Parallelism Formulation

Definition 2 (Data Parallel Setup). $K$ devices, each holds copy of $\boldsymbol{\theta}$. Mini-batch $\mathcal{B}$ split into $\mathcal{B}_1, \ldots, \mathcal{B}_K$ with $|\mathcal{B}_k| = |\mathcal{B}|/K$.

Algorithm:

1. Each device $k$ computes local gradient $\mathbf{g}_k = \frac{1}{|\mathcal{B}_k|}\sum_{i \in \mathcal{B}_k} \nabla \mathcal{L}_i$
2. AllReduce: compute average $\bar{\mathbf{g}} = \frac{1}{K}\sum_k \mathbf{g}_k$
3. Each device updates: $\boldsymbol{\theta} \leftarrow \boldsymbol{\theta} - \alpha \bar{\mathbf{g}}$

Theorem 1 (Equivalence). Synchronized data parallel SGD with averaged gradients is equivalent to SGD on full batch $\mathcal{B} = \cup_k \mathcal{B}_k$ (assuming deterministic ops).

Proof. $\bar{\mathbf{g}} = \frac{1}{K}\sum_k \frac{1}{|\mathcal{B}_k|}\sum_{i \in \mathcal{B}_k} \nabla \mathcal{L}_i = \frac{1}{|\mathcal{B}|}\sum_{i \in \mathcal{B}} \nabla \mathcal{L}_i$. $\blacksquare$

Gradient Aggregation and AllReduce

Definition 3 (AllReduce). Collective operation: each of $K$ processes holds vector $\mathbf{v}_k$; after AllReduce, all processes hold $\sum_k \mathbf{v}_k$ (or average).

Definition 4 (Ring AllReduce). Arrange $K$ GPUs in ring; pass gradient chunks in $K-1$ steps. Bandwidth-optimal: each GPU sends/receives $2(K-1)/K \approx 2$ × gradient size total.

Proposition 1. Naive AllReduce: $O(K \cdot |\boldsymbol{\theta}|)$ data transferred. Ring AllReduce: $O(|\boldsymbol{\theta}|)$ per device regardless of $K$ (bandwidth-limited).

Synchronous Distributed SGD

Definition 5 (Synchronous Training). All devices wait for slowest before update — straggler problem.

Definition 6 (Asynchronous Training). Devices update independently with stale gradients — convergence complications.

Production LLM training uses synchronous data parallelism with gradient accumulation for large effective batches (Scaling Laws).

Communication Complexity

Definition 7 (Communication Time).

for ring AllReduce (factor 2 for send + receive).

Definition 8 (Compute Time).

Proposition 2. Training is communication-bound when $T_{\text{comm}} > T_{\text{compute}}$ — common for small models on fast GPUs or large $K$.

Scaling Efficiency and Amdahl's Law

Definition 9 (Strong Scaling Efficiency).

where $T_K$ is time with $K$ devices.

Theorem 2 (Amdahl's Law). If fraction $p$ of work is parallelizable:

Sequential overhead (communication, I/O) limits speedup.

Proposition 3. For LLM training with large batches, $p \approx 1$ and $\eta > 0.9$ achievable up to hundreds of GPUs with optimized AllReduce.

Effective Batch Size and Learning Rate

Definition 10 (Effective Batch Size). $B_{\text{eff}} = K \times |\mathcal{B}_k| \times \text{accumulation steps}$.

Proposition 4 (Linear Scaling Rule). When increasing $B_{\text{eff}}$ by factor $K$, scale learning rate $\alpha \leftarrow K \alpha$ (up to stability limit) — preserves SGD dynamics approximately.

Proposition 5 (Square Root Scaling). Alternative: $\alpha \leftarrow \sqrt{K}\,\alpha$ — more conservative, better for large $K$.

See Gradient Descent and Scaling Laws.

Worked Examples

Example 1: AllReduce Volume

$|\boldsymbol{\theta}| = 7 \times 10^9$, FP32, $K = 8$: gradient size 28 GB; ring AllReduce $\approx$ 56 GB total transferred across network.

Example 2: Effective Batch

$K=64$, local batch 2, accumulation 4: $B_{\text{eff}} = 512$.

Connection to the Broader Curriculum

• Gradient Descent — base algorithm
• Scaling Laws — compute-optimal batch sizes
• Backpropagation — local gradient computation
• Quantization — quantized AllReduce

Common Pitfalls and Misconceptions

Pitfall 1: Assuming linear speedup to arbitrary $K$.

Pitfall 2: Ignoring straggler effects in synchronous training.

Pitfall 3: Not scaling learning rate with batch size.

Pitfall 4: Confusing data parallelism with model/tensor parallelism.

Research Perspective

Distributed data parallelism underpins virtually all large-scale neural network training. Ring AllReduce (Baidu, 2017) established bandwidth-optimal gradient aggregation. Horovod (Sergeev & Balso, 2018) standardized collective communication across heterogeneous clusters. Megatron-LM and subsequent systems combine data parallelism with tensor and pipeline parallelism for models that exceed single-device memory. ZeRO (Rajbhandari et al., 2020) shards optimizer states across devices, reducing per-GPU memory while preserving the data-parallel abstraction. Gradient compression and error-feedback methods remain active research directions for bandwidth-constrained environments.

Summary of Takeaways

• Data parallel — Replicate model, split batch
• Gradient — Average via AllReduce
• Ring AllReduce — $O(\lVert \boldsymbol{\theta} \rVert)$ bandwidth
• Efficiency — Amdahl limits speedup
• Batch scaling — Effective batch $K \cdot — \mathcal{B}_k —$

Exercises

Exercise 1. Prove Theorem 1.

Exercise 2. Derive ring AllReduce communication volume.

Exercise 3. Compute $T_{\text{comm}}$ vs $T_{\text{compute}}$ for given config.

Exercise 4. Amdahl: if 5% sequential, max speedup at $K=100$?

Exercise 5. Linear vs. sqrt learning rate scaling.

Exercise 6. Compare data vs. model parallelism for 70B model.

Exercise 7. Gradient accumulation equivalence.

Exercise 8. Connect to Scaling Laws compute budget.