Distributed training overview

Core Idea

Think of it like a restaurant kitchen: one chef (GPU) can only make so many dishes per hour. With10 chefs working together, you serve customers10× faster—if they coordinate properly and don't get in each other's way.

The Fundamental Problem

What Are We Distributing?

When training a model, three resources dominate:

1. Computation: Forward pass ($\mathbf{y} = f(\mathbf{x}; \theta)$) and backward pass ($\nabla_\theta \mathcal{L}$)
2. Memory: Storing parameters $\theta$, activations, gradients, optimizer states
3. Data: The training dataset $\mathcal{D} = \{(\mathbf{x}_i, y_i)\}_{i=1}^N$

Main Paralelization Strategies

1. Data Paralelism

The workflow:

1. Split mini-batch $B$ of size $b$ into $k$ sub-batches: $B_1, B_2, \ldots, B_k$ (one per device)
2. Each device computes gradients on its sub-batch: $\mathbf{g}_i = \frac{1}{|B_i|} \sum_{(\mathbf{x}, y) \in B_i} \nabla_\theta \mathcal{L}(f(\mathbf{x}; \theta), y)$
3. Synchronize gradients across devices (all-reduce operation): $\mathbf{g}_{\text{global}} = \frac{1}{k} \sum_{i=1}^k \mathbf{g}_i$
4. Each device updates its local copy with $\mathbf{g}_{\text{global}}$: $\theta \leftarrow \theta - \eta \mathbf{g}_{\text{global}}$

Why this works: The average of gradients across mini-batches equals the gradient of the full batch (linearity of expectation). It's mathematically equivalent to training on a batch of size $b \times k$ on a single device.

Why this step? Synchronizing ensures all devices have identical parameters after each update. Without sync, devices would diverge and produce different models.

Result: 4× faster training (ideally), but communication overhead reduces actual speedup to ~3.5×.

2. Model Paralelism

Two flavors:

A) Tensor Parallelism (Intra-layer splitting) Split individual layers across devices. For a large matrix multiplication $\mathbf{Y} = \mathbf{X} \mathbf{W}$ where $\mathbf{W} \in \mathbb{R}^{d \times h}$:

$\mathbf{W} = \begin{bmatrix} \mathbf{W}_1 \\ \mathbf{W}_2 \end{bmatrix}, \quad \mathbf{Y} = \mathbf{X} \begin{bmatrix} \mathbf{W}_1 \\ \mathbf{W}_2 \end{bmatrix} = \begin{bmatrix} \mathbf{Y}_1 & \mathbf{Y}_2 \end{bmatrix}$

Device 1 computes $\mathbf{Y}_1 = \mathbf{X} \mathbf{W}_1$, Device 2 computes $\mathbf{Y}_2 = \mathbf{X} \mathbf{W}_2$ in parallel. Then concatenate results.

Why this step? When $\mathbf{W}$ is too large for one GPU's memory (e.g., a 50GB weight matrix), splitting it is the only way to fit the model.

B) Pipeline Parallelism (Inter-layer splitting) Split the model vertically by layers. Device 1 has layers1-10, Device 2 has layers 11-20, etc.

Data flows like an assembly line:

• Batch 1 enters Device 1 (layers 1-10)
• Output of Device 1 → 2 (layers 11-20)
• While Device 2 processes Batch 1, Device 1 starts on Batch 2

Challenge: Pipeline bubles—devices sit idle waiting for data from the previous stage.

Solution: Split across 8 GPUs:

• GPU 0: Layers 0-11
• GPU 1: Layers 12-23
• ...
• GPU 7: Layers 84-95

A forward pass requires 7 device-to-device transfers (expensive!). But it's the only way to run the model.

3. Hybrid Parallelism

Example: Training a 100B parameter model on 64 GPUs (8 nodes × 8 GPUs):

• Split model across 8 GPUs (model parallelism)
• Replicate this 8 times (data parallelism across nodes)
• Total: 8 data-parallel replicas, each using 8-way model parallelism

Key Trade-offs

Communication Overhead

Actual speedup: $S = \frac{T_{\text{compute}}}{T_{\text{compute}}/k + T_{\text{comm}}}$

where $T_{\text{comm}}$ is time spent synchronizing gradients/activations.

Why this matters: If $T_{\text{comm}} \approx T_{\text{compute}}/k$, speedup is only $0.5k$ (50% efficiency). High-bandwidth interconnects (NVLink, InfiniBand) reduce $T_{\text{comm}}$.

Scaling Efficiency

• Linear scaling: Efficiency =1 (rare, only with minimal communication)
• Sub-linear scaling: Efficiency = 0.7-0.9 (typical for well-optimized systems)
• Super-linear scaling: Efficiency > 1 (possible when larger batch fits in cache)

Why efficiency degrades: As $k$ increases, each device does less work ($T_{\text{compute}}/k$ shrinks), but $T_{\text{comm}}$ stays constant or grows (more devices to sync).

Common Implementation Patterns

Synchronous vs. Asynchronous Training

Pros: Mathematically equivalent to single-device training (deterministic) Cons: Speed limited by the slowest device (stragglers hurt performance)

Pros: No waiting for stragglers, better hardware utilization Cons: Stale gradients (worker uses $\theta_t$ but updates $\theta_{t+5}$), can hurt convergence

Gradient Accumulation

1. Accumulate gradients over $k = B/b$ micro-batches: $\mathbf{g}_{\text{accum}} = \frac{1}{k} \sum_{i=1}^k \mathbf{g}_i$
2. Update after $k$ forward-backward passes: $\theta \leftarrow \theta - \eta \mathbf{g}_{\text{accum}}$

Why this works: Gradients are linear—averaging over micro-batches equals the gradient of the full batch. It's equivalent to processing a batch of size $B$, just slower.

Mistakes & Misconceptions

The reality: Communication overhead grows with device count. Beyond a certain point (the "scaling wall"), adding devices gives diminishing returns or even slows things down.

Example: Training a small CNN on 100 GPUs might be slower than on 8 GPUs because gradient sync time dominates computation.

Fix: Match paralelism to model size. Small models: 1-8 GPUs with data parallelism. Large models: 100+ GPUs with hybrid strategies.

The difference:

• Large batch on1 GPU: Process 256 samples → update
• Data parallel on 4 GPUs (batch 64 each): Process 64 samples per GPU → sync update (effective batch 256)

They're equivalent in math but different in implementation. Data parallelism adds communication cost; large batches one GPU add memory pressure.

The nuance: Model parallelism adds cross-device communication every layer (high latency). Data parallelism only syncs at the end of backward pass.

Better strategy: Use data parallelism as much as possible, model parallelism only when necessary. For example, a 10B parameter model might fit on a 40GB GPU with activation checkpointing—avoid model parallelism.

Practical Considerations

Choosing a Strategy

Framework Support

• PyTorch: DistributedDataParallel (data), Tensor Parallelism (via Megatron-LM), PipeDream (pipeline)
• TensorFlow: tf.distribute.MirroredStrategy (data), tf.distribute.TPUStrategy (data on TPUs)
• JAX: pmap (data), xmap (arbitrary sharding)
• DepSpeed: ZeRO stages (optimizer sharding), pipeline paralelism
• Horovod: MPI-based data parallelism (framework-agnostic)

Active Recall Practice

Distributed training is like that for AI models. Training a huge AI brain (like ChatGPT) one computer takes forever. So we split the work across many computers (GPUs). Each computer trains on different pictures or sentences, learns something, and then all computers share what they learned. By combining everyone's answers, the AI learns faster.

The tricky part? The computers need to talk to each other to share answers (like you comparing homework answers), and that talking takes time. If they spend too much time talking, it's slower than just working alone!

• Model parallelism: Split model across devices

• Overhead: Communication cost grows with devices

• Devices: More isn't always better (scaling wall)

• Efficiency: Actual speedup / ideal speedup

• Layers: Pipeline splits by layers

• Hybrid: Combines data + model parallelism

• Yield: Best for100B+ parameter models

• Batch: Gradient accumulation simulates large batches

• Ring all-reduce: Efficient gradient sync topology

• Interconnect: NVLink/InfiniBand critical for model parallelism

• Deterministic: Synchronous SGD matches single-GPU

Connections

• 3.1-Neural-network-basics - forward/backward passes being parallelized
• 3.3.5-Batch-normalization - batch statistics in data parallelism
• 3.9-Gradient-descent-optimizers - optimizer state sharding (ZeRO)
• 3.4.1-Transformers-overview - why GPT-3 needs model parallelism
• 4.2-GPU-acceleration - hardware that enables distribution
• 5.1-Training-at-scale - practical engineering of distributed systems

#flashcards/ai-ml

What is the main goal of distributed training? :: Speed up training and enable models larger than single-device memory by paralelizing computation across multiple GPUs/TPUs.

Concept Map

Hinglish (regional understanding)