Skip to content

Distributed Training

Overview

Distributed Training enables scaling model training across multiple GPUs and nodes by dividing computation, parameters, or data across devices. Essential for training LLMs that don't fit on single GPU memory.

  • Common Strategies: Data Parallelism, Model Parallelism, Pipeline Parallelism
  • Frameworks: PyTorch Distributed, DeepSpeed, FSDP (Fully Sharded Data Parallel)
  • Scaling: 16→128 GPUs with linear speedup possible
  • Challenge: Communication overhead, load balancing, convergence stability
  • Industry Standard: FSDP (PyTorch 2.0+) or DeepSpeed for large models

The Scaling Problem

Single GPU Bottleneck

GPU Memory Breakdown (Training 70B LLaMA):

Model Parameters: 70B (140GB in float32)
Optimizer States: 2 × 70B = 140GB
 - Adam: needs m (momentum) and v (variance)
 - Each is same size as parameters
Gradients: 70B = 140GB
Activations: Depends on batch size, seq_len
Total: 280GB+ (excluding activations!)

Available GPU Memory:
 - A100 80GB: Only fits with 4-bit quantization
 - H100 141GB: Barely fits with 2-8 batch size
 - Problem: Large models need multiple GPUs!

Solution:
 - Shard across GPUs
 - Shard parameters across devices
 - Reduce gradient storage per device
 - This is distributed training!

Scaling Reality

Weak scaling:
 - Add more GPUs
 - Keep per-GPU batch size constant
 - Increase total batch size proportionally
 - Time per epoch stays roughly constant 
 - Wall-clock time: Same per epoch

Strong scaling:
 - Add more GPUs
 - Keep total batch size constant
 - Reduce per-GPU batch size
 - Communication overhead becomes significant 
 - Wall-clock time improvement limited by communication

Typical speedup:
 - 4 GPUs: 3.5x speedup (87.5% efficiency)
 - 8 GPUs: 6.5x speedup (81% efficiency)
 - 16 GPUs: 11x speedup (69% efficiency)
 - 128 GPUs: 60x speedup (47% efficiency)
 - Communication costs grow with device count!

-

Distributed Training Strategies

1. Data Parallelism (Most Common)

Architecture:
- ┌─────────────────────────────────────────────────┐
 - Main Model │
 - ┘
 ↓ ↓ ↓
- ┌────────────┬────────────┬────────────┐
 - GPU 0 │ GPU 1 │ GPU 2 │
 - Batch 1 │ Batch 2 │ Batch 3 │
 - ┴────────────┴────────────┘
 ↓ ↓ ↓
- ┌─────────────────────────────────────┐
 - Compute gradients locally │
 - ┘
 ↓
- ┌─────────────────────────────────────┐
 - All-reduce: Synchronize gradients │
 - Average gradient across GPUs │
 - ┘
 ↓
- ┌─────────────────────────────────────┐
 - SGD step: All GPUs update identically │
 - ┘

Pros:
Simple: Each GPU has full model
Works with existing code
Linear scaling for Weak scaling

Cons:
All-reduce communication every step
Cannot scale to very large models (still need model to fit on GPU)
Bandwidth limited

2. Model Parallelism (Pipeline Parallelism)

Architecture:
GPU 0: Layers 1-8 GPU 1: Layers 9-16 GPU 2: Layers 17-24
 ↓ activations ↓ activations ↓ activations
 - ┘

Forward pass (Pipeline parallelism):
Step 1: GPU 0 processes batch, sends to GPU 1
Step 2: GPU 1 processes, sends to GPU 2
Step 3: GPU 2 computes loss
Step 4: GPU 1 computes gradients (while GPU 0 is idle)
Step 5: GPU 0 computes gradients (while GPU 1/2 idle)

Problem: Pipeline bubble - GPUs sit idle during backward pass!

Batch within step (Micro-batching):
 - Split batch into M micro-batches
 - Pipeline better utilized
 - But reduced gradient accumulation
 - Trade-off: Complexity vs. utilization

Pros:
Model can be larger (distributed across GPUs)
Useful for very large models (1T+ parameters)

Cons:
Complex implementation
Pipeline bubble overhead
Requires careful microbatch tuning
Asynchronous communication difficult

3. Fully Sharded Data Parallel (FSDP)

Architecture: Combine Data Parallelism + Parameter Sharding

Standard Data Parallel (Problem):
 - GPU 0: Full model (140GB) + Batch 1
 - GPU 1: Full model (140GB) + Batch 2
 - GPU 2: Full model (140GB) + Batch 3
 - Total GPU memory: 3 × 140GB = 420GB needed!

FSDP Solution:
 - Shard model parameters across GPUs
 - GPU 0: Model layers 1-8 (20GB) + Batch 1
 - GPU 1: Model layers 9-16 (20GB) + Batch 2
 - GPU 2: Model layers 17-24 (20GB) + Batch 3
 - During forward: all-gather sharded parameters (temporary full model)
 - During backward: all-reduce gradients
 - Total GPU memory: 420GB → 140GB (3x reduction!)

Modern Implementation (PyTorch FSDP):
```python
from torch.distributed.fsdp import FSDP, CPUOffload
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy

# Simple FSDP wrapping
model = FSDP(model, auto_wrap_policy=transformer_auto_wrap_policy)

# With CPU offloading (for even larger models)
model = FSDP(
 model,
 auto_wrap_policy=transformer_auto_wrap_policy,
 cpu_offload=CPUOffload(offload_params=True) # Offload unused params to CPU
)

# Training loop
for epoch in range(num_epochs):
 for batch in dataloader:
 x, y = batch

 # Forward pass (automatically gathers shards)
 output = model(x)
 loss = criterion(output, y)

 # Backward (automatically reduces gradients)
 loss.backward()

 # Optimizer step (updates local shard)
 optimizer.step()
 optimizer.zero_grad()

Pros: Memory efficient: Reduce per-GPU memory 2-16x Communication overhead lower than model parallel Easy integration with existing code DeepSpeed/FSDP: Production-ready

Cons: All-gather overhead during forward pass Requires distributed communication setup Slightly more complex debugging

### Comparison

Strategy Memory Speed Complexity Scalability ───────────────────────────────────────────────────────── Data Parallel (limited) Pipeline Parallel FSDP

Recommendation:

  • <8 GPUs: Data Parallel (simplicity)
  • 8-64 GPUs: FSDP (balanced)
  • 64+ GPUs: Pipeline + FSDP hybrid
-

## Communication Optimization

### The Communication Problem

Training time breakdown (8 GPUs, 70B model):

Computation: 800ms per step All-reduce: 200ms per step (gradient sync) Total: 1000ms per step

Efficiency: 80% compute, 20% communication

As GPUs increase:

  • 16 GPUs: Computation 400ms, All-reduce 400ms (50% each!)
  • 128 GPUs: Computation 50ms, All-reduce 950ms (5% compute!)
  • Communication becomes bottleneck!
### Solutions

#### 1. Gradient Accumulation

```python
accumulation_steps = 4

for step, batch in enumerate(dataloader):
 output = model(batch)
 loss = criterion(output, target)

 # Scale loss to account for accumulation
 (loss / accumulation_steps).backward()

 # Every N steps: synchronize gradients
 if (step + 1) % accumulation_steps == 0:
 optimizer.step()
 optimizer.zero_grad()

# Effect:
# Instead of
# Now
# Communication

2. Gradient Compression

class CompressedGradient:
 """Reduce communication by quantizing gradients"""

 def compress(self, grads, compression_ratio=0.01):
 """Keep only top-1% gradients"""
 # Find magnitude of each gradient
 magnitudes = torch.abs(grads)

 # Keep only top-k by magnitude
 k = max(1, int(grads.numel() * compression_ratio))
 threshold = torch.kthvalue(magnitudes, k)[0]

 # Create mask: keep only large gradients
 mask = magnitudes >= threshold
 compressed = grads * mask
 return compressed, mask

 def decompress(self, compressed, mask):
 """Reconstruct"""
 return compressed * mask

# Usage:
grad_compressor = CompressedGradient()
compressed_grads, mask = grad_compressor.compress(grads, compression_ratio=0.1)
dist.all_reduce(compressed_grads) # Only communicate 10% of gradients!
grads = grad_compressor.decompress(compressed_grads, mask)

3. Communication Overlapping

# Instead of
# Do

# PyTorch enables this automatically with:
class OverlappedDistributedDataParallel(DistributedDataParallel):
 def forward(self, *inputs, **kwargs):
 # Start gradient all-reduce in background
 self.reducer.prepare_for_backward(self._get_ddp_logging_data())

 # Compute while all-reduce happens
 return super().forward(*inputs, **kwargs)

# DeepSpeed does this automatically:
import deepspeed
model, optimizer, _, _ = deepspeed.initialize(
 model=model,
 model_parameters=model.parameters(),
 config=config_dict
)
# All communication overlapped automatically!

-

Convergence and Stability

Large Batch Training Challenges

Problem: Large batch size increases training variance
 - Small batch (32): Noisy gradients, but explore well
 - Large batch (16K across 128 GPUs): Stable gradients, may miss optima

Solution: Warmup + learning rate scaling

Implementation

class LinearWarmupScheduler:
 """Gradually increase LR during warmup"""

 def __init__(self, optimizer, warmup_steps, max_lr):
 self.optimizer = optimizer
 self.warmup_steps = warmup_steps
 self.max_lr = max_lr
 self.step_num = 0

 def step(self):
 self.step_num += 1

 if self.step_num < self.warmup_steps:
 lr = (self.step_num / self.warmup_steps) * self.max_lr
 else:
 lr = self.max_lr

 for param_group in self.optimizer.param_groups:
 param_group['lr'] = lr

# Usage:
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
scheduler = LinearWarmupScheduler(optimizer, warmup_steps=1000, max_lr=1e-3)

for step in range(num_steps):
 loss = train_step()
 loss.backward()
 optimizer.step()
 scheduler.step() # Gradually increase LR

Key Takeaways

Data Parallelism: Simple but memory-hungry FSDP: Industry standard for large-scale training Communication is bottleneck for 64+ GPU training Gradient accumulation reduces communication overhead Large batch requires careful warmup and learning rate tuning

-