Distributed Training for LLMs

Training modern LLMs requires distributing computation across hundreds or thousands of GPUs. A single NVIDIA A100 has 80GB of GPU memory — far insufficient for models with hundreds of billions of parameters. This tutorial covers the parallelism strategies, memory optimization techniques, and frameworks that enable large-scale LLM training.

Why Distributed Training is Necessary

Memory Requirements

The memory required to train a model depends on its parameter count, batch size, and optimizer state:

Scaling Laws and Training Time

Parallelism Strategies

Data Parallelism (DP)

The simplest form of distributed training — replicate the model on every GPU and split the data:

GPU 0: Model Copy + Batch 0 -> Gradients_0 --+
GPU 1: Model Copy + Batch 1 -> Gradients_1 --+-> All-Reduce -> Average -> Update
GPU 2: Model Copy + Batch 2 -> Gradients_2 --+
GPU 3: Model Copy + Batch 3 -> Gradients_3 --+

Limitations:

• Each GPU must hold the entire model (parameters + optimizer states + gradients)
• Communication cost scales with model size (all-reduce of gradients)
• Cannot train models larger than single GPU memory

ZeRO (Zero Redundancy Optimizer)

DeepSpeed ZeRO eliminates memory redundancy in data parallelism:

ZeRO Stage 3 Memory Layout:
GPU 0: Optimizer shard 0 + Gradient shard 0 + Parameter shard 0
GPU 1: Optimizer shard 1 + Gradient shard 1 + Parameter shard 1
GPU 2: Optimizer shard 2 + Gradient shard 2 + Parameter shard 2
GPU 3: Optimizer shard 3 + Gradient shard 3 + Parameter shard 3

Fully Sharded Data Parallel (FSDP)

PyTorch's native implementation of ZeRO-3:

import torch
import torch.nn as nn
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import ShardingStrategy

model = nn.TransformerDecoder(
    nn.TransformerDecoderLayer(d_model=4096, nhead=32),
    num_layers=80
)

model = FSDP(
    model,
    sharding_strategy=ShardingStrategy.FULL_SHARD,
    mixed_precision=MixedPrecision(
        param_dtype=torch.bfloat16,
        reduce_dtype=torch.bfloat16,
    ),
    auto_wrap_policy=transformer_auto_wrap_policy,
    device_id=torch.cuda.current_device(),
)

for batch in dataloader:
    loss = model(batch)
    loss.backward()
    optimizer.step()

Tensor Parallelism (TP)

Split individual layers across multiple GPUs:

For a linear layer Y = XW where W is split column-wise:

W = [W_0 | W_1]  (split across 2 GPUs)

GPU 0: Y_0 = X @ W_0
GPU 1: Y_1 = X @ W_1
Y = [Y_0 | Y_1]  (concatenated)

Pipeline Parallelism (PP)

Split model layers across GPU stages:

Stage 0 (GPU 0): Layers 0-19   -> Output_0
Stage 1 (GPU 1): Layers 20-39  -> Output_1
Stage 2 (GPU 2): Layers 40-59  -> Output_2
Stage 3 (GPU 3): Layers 60-79  -> Final Output

3D Parallelism

Modern LLM training combines all three strategies:

Mixed-Precision Training

FP16 and BF16

Loss Scaling

scaler = torch.cuda.amp.GradScaler()

for batch in dataloader:
    optimizer.zero_grad()
    
    with torch.cuda.amp.autocast(dtype=torch.bfloat16):
        outputs = model(batch["input_ids"])
        loss = criterion(outputs, batch["labels"])
    
    scaler.scale(loss).backward()
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    scaler.step(optimizer)
    scaler.update()

Communication Optimization

Gradient Accumulation

accumulation_steps = 8
optimizer.zero_grad()

for i, batch in enumerate(dataloader):
    loss = model(batch) / accumulation_steps
    loss.backward()
    
    if (i + 1) % accumulation_steps == 0:
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        optimizer.zero_grad()

Communication Overhead

Training Frameworks

DeepSpeed

import deepspeed

ds_config = {
    "train_batch_size": 1024,
    "gradient_accumulation_steps": 8,
    "fp16": {"enabled": True, "loss_scale": 0, "initial_scale_power": 16},
    "zero_optimization": {
        "stage": 3,
        "overlap_comm": True,
        "contiguous_gradients": True,
        "reduce_bucket_size": 5e7,
        "stage3_prefetch_bucket_size": 5e7,
        "stage3_param_persistence_threshold": 1e6
    },
    "activation_checkpointing": {
        "partition_activations": True,
        "cpu_checkpointing": False,
        "contiguous_memory_optimization": False
    }
}

model_engine, optimizer, _, _ = deepspeed.initialize(
    model=model,
    config=ds_config,
    model_parameters=model.parameters()
)

Activation Checkpointing

Memory Optimization Techniques

CPU Offloading

ds_config = {
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {"device": "cpu", "pin_memory": True},
        "offload_param": {"device": "cpu", "pin_memory": True}
    }
}

Training Stability

Learning Rate Scheduling

from transformers import get_cosine_schedule_with_warmup

scheduler = get_cosine_schedule_with_warmup(
    optimizer,
    num_warmup_steps=2000,
    num_training_steps=100000,
    min_lr_ratio=0.1
)

Gradient Clipping

max_grad_norm = 1.0
torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)

Practice Exercises

1. ZeRO Stage Comparison: Calculate the memory savings for training a 7B model using ZeRO Stage 1, 2, and 3 on 4 GPUs. What is the memory per GPU for each stage?

2. Pipeline Bubble Analysis: Compute the pipeline efficiency for a 32-stage pipeline with 8, 16, 32, and 64 micro-batches. At what point does efficiency exceed 90%?

3. 3D Parallelism Design: Design a 3D parallelism configuration for a 405B parameter model on 2048 GPUs. Justify your choice of TP, PP, and DP degrees.

4. Communication Analysis: For a 70B model with gradient all-reduce on 256 GPUs, calculate the communication volume and estimated latency assuming 100 Gbps interconnect.

Key Takeaways

What to Learn Next