Skip to main content

Chapter 3: Distributed Training Techniques

Introduction​

Training LLMs with billions of parameters requires distributing work across many GPUs. This chapter covers the key parallelism strategies.

Data Parallelism​

The Simplest Approach​

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel

# Each GPU processes different data, same model
model = DistributedDataParallel(model)

Limitation: Model must fit on single GPU

Model Parallelism​

Tensor Parallelism​

Split individual layers across GPUs:

[Layer Weight Matrix]
   ↓ Split columns
[GPU 0: W[:, :d/2]] [GPU 1: W[:, d/2:]]

Pipeline Parallelism​

Different layers on different GPUs:

GPU 0: Layers 0-11
GPU 1: Layers 12-23
GPU 2: Layers 24-35
GPU 3: Layers 36-47

ZeRO: Zero Redundancy Optimizer​

Three Stages​

StageWhat's PartitionedMemory Savings
ZeRO-1Optimizer states4x
ZeRO-2+ Gradients8x
ZeRO-3+ ParametersLinear with GPUs

Implementation with DeepSpeed​

import deepspeed

model_engine, optimizer, _, _ = deepspeed.initialize(
    model=model,
    config={
        "zero_optimization": {
            "stage": 3,
            "offload_optimizer": {"device": "cpu"},
            "offload_param": {"device": "cpu"}
        },
        "bf16": {"enabled": True}
    }
)

Training Compute​

Chinchilla Scaling​

Optimal training: approximately 20 tokens per parameter

# Compute-optimal model size
def optimal_model_size(compute_budget):
    # C ≈ 6 * N * D (N=params, D=tokens)
    # Optimal: D ≈ 20N
    return compute_budget / (6 * 20)

Key Takeaways​

Summary
  1. Data parallelism for small models
  2. Model parallelism when model doesn't fit on one GPU
  3. ZeRO dramatically reduces memory requirements
  4. Mixed precision (bf16) speeds training