Distributed Training

1. The Scaling Challenge

Training modern large language models (LLMs) requires distributing computation across multiple accelerators (GPUs/TPUs). For a model with parameters and a dataset of tokens, the total compute is approximately:

For GPT-3 (175B parameters, 300B tokens): FLOPs. At 312 TFLOPS (A100 GPU), this requires GPU-seconds, or about 1000 A100s for 34 days.

2. Parallelism Strategies

3. Data Parallelism

3.1 Basic Data Parallelism

Each GPU holds a complete copy of the model. The batch is split across GPUs:

After local gradient computation, gradients are averaged via AllReduce:

3.2 Communication Analysis

The AllReduce operation requires transferring the full gradient tensor twice (reduce-scatter + all-gather):

where is the number of parameters and is the bytes per parameter (e.g., 2 for FP16).

For a 175B parameter model in FP16: GB per GPU.

3.3 Ring AllReduce

Ring AllReduce achieves optimal bandwidth utilization:

For :

4. ZeRO (Zero Redundancy Optimizer)

4.1 Problem: Memory Redundancy

In standard data parallelism, each GPU stores:

1. Model parameters: bytes
2. Gradients: bytes
3. Optimizer states: bytes (Adam: momentum + variance in FP32 + FP32 master weights)

Total: bytes per GPU.

For 175B params in FP16: TB per GPU (impossible on single A100 with 80GB).

4.2 ZeRO Stages

ZeRO (Rajbhandari et al., 2020) partitions optimizer states, gradients, and parameters across GPUs:

Stage 1 (ZeRO-1): Partition optimizer states

Stage 2 (ZeRO-2): Partition optimizer states + gradients

Stage 3 (ZeRO-3): Partition everything

4.3 ZeRO Memory Formula

5. Tensor Parallelism

5.1 Column Parallel Linear Layer

Split the weight matrix by columns:

Each GPU holds and computes:

The full output is gathered via AllGather:

5.2 Row Parallel Linear Layer

Split the weight matrix by rows:

Each GPU holds and a shard of the input:

The partial results are summed via AllReduce:

5.3 Megatron-LM Style

Megatron-LM combines column and row parallelism for transformer layers:

1. MLP: Column-parallel → Row-parallel
2. Self-attention: Split heads across GPUs
3. Residual connections: AllReduce after each sub-layer

6. Pipeline Parallelism

6.1 GPipe

GPipe (Huang et al., 2019) splits the model into stages and processes micro-batches:

Bubble fraction:

For , the bubble fraction approaches 0.

6.2 1F1B (One Forward One Backward)

Interleaves forward and backward passes:

Bubble fraction:

Key advantage: constant memory instead of in GPipe.

6.3 Pipeline Bubble Analysis

The pipeline bubble represents idle time when GPUs are waiting for activations from previous stages:

where is the time for one micro-batch forward + backward pass.

To minimize bubble:

1. Increase (more micro-batches)
2. Decrease (fewer stages)
3. Use 1F1B schedule to overlap

6.4 Communication Overlap

In pipeline parallelism, only point-to-point communication between adjacent stages:

where is batch size, is sequence length, and is hidden dimension.

7. FSDP (Fully Sharded Data Parallel)

7.1 Overview

FSDP (Zhao et al., 2023) is PyTorch's implementation of ZeRO-3 style partitioning:

1. Initialization: Shard parameters across GPUs
2. Forward pass: AllGather parameters before each layer, discard after
3. Backward pass: AllGather again, compute gradients, reduce-scatter
4. Optimizer step: Each GPU updates only its shard

7.2 FSDP Communication Pattern

7.3 Mixed Precision with FSDP

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import MixedPrecision

mp_policy = MixedPrecision(
    param_dtype=torch.bfloat16,
    reduce_dtype=torch.float32,
    buffer_dtype=torch.float32,
)

model = FSDP(model, mixed_precision=mp_policy)

8. DeepSpeed

8.1 DeepSpeed Features

DeepSpeed (Rasley et al., 2020) provides:

• ZeRO stages 1-3
• Mixed precision training
• Gradient accumulation
• Activation checkpointing
• 3D parallelism (DP + TP + PP)
• DeepSpeed-ML inference optimization

8.2 DeepSpeed Configuration

{
  "train_batch_size": 1024,
  "gradient_accumulation_steps": 8,
  "fp16": {
    "enabled": true,
    "loss_scale": 0,
    "initial_scale_power": 16
  },
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    },
    "offload_param": {
      "device": "cpu",
      "pin_memory": true
    },
    "overlap_comm": true,
    "contiguous_gradients": true,
    "sub_group_size": 1e9,
    "reduce_bucket_size": 5e8,
    "stage3_prefetch_bucket_size": 5e8,
    "stage3_param_persistence_threshold": 1e6,
    "stage3_max_live_parameters": 1e9,
    "stage3_max_reuse_distance": 1e9
  }
}

8.3 ZeRO-Offload

Offload optimizer states and gradients to CPU memory:

This enables training models larger than GPU memory by using CPU as extension.

9. Activation Checkpointing

9.1 Problem: Activation Memory

During training, activations from forward pass must be saved for backward pass:

For large models, this dominates GPU memory.

9.2 Checkpointing Strategy

Only save activations at checkpoint boundaries:

Trade-off: more FLOPs (recompute activations) but times less memory.

9.3 Selective Checkpointing

Checkpoint only attention layers (most memory-intensive):

vs. MLP activations:

10. Performance Analysis

10.1 Scaling Efficiency

where is single-GPU time and is -GPU time.

10.2 Communication Volume Summary

10.3 Optimal Strategy Selection

For model size and GPU memory :

1. : Data parallelism
2. : ZeRO-2 or ZeRO-3
3. : ZeRO-3 + CPU offloading
4. : 3D parallelism (DP + TP + PP)

11. Implementation Example

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from deepspeed import initialize

def setup_ddp(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)

def train_step_ddp(model, batch, optimizer):
    optimizer.zero_grad()
    loss = model(**batch)
    loss.backward()
    optimizer.step()
    return loss

# DeepSpeed integration
model_engine, optimizer, _, _ = initialize(
    model=model,
    model_parameters=model.parameters(),
    config="ds_config.json"
)

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

References

1. Rajbhandari et al. (2020). "ZeRO: Memory Optimizations Toward Training Trillion Parameter Models." SC'20.
2. Huang et al. (2019). "GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism." NeurIPS.
3. Narayanan et al. (2021). "Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM." SC'21.
4. Zhao et al. (2023). "PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel." VLDB.
5. Rasley et al. (2020). "DeepSpeed: System Optimizations Enable Training Deep Learning Models with Over 100 Billion Parameters." KDD.