Distributed Parallelism in Deep Learning

Data Parallelism: Split the Batch

Data parallelism is the simplest and most widely used strategy. Each GPU holds a complete copy of the model and processes a different portion of the training batch. After computing local gradients, GPUs synchronize via all-reduce to average gradients before updating weights.

When to Use Data Parallelism

• Model fits comfortably in single GPU memory
• Need to increase training throughput
• Simple setup with minimal code changes

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

# Initialize process group with NCCL backend
dist.init_process_group(backend='nccl')
local_rank = int(os.environ['LOCAL_RANK'])
torch.cuda.set_device(local_rank)

# Create model and wrap with DDP
model = MyModel().cuda(local_rank)
model = DDP(model, device_ids=[local_rank])

# Training loop - gradients are automatically synchronized
for batch in dataloader:
    inputs, labels = batch.cuda(local_rank), labels.cuda(local_rank)
    loss = criterion(model(inputs), labels)
    loss.backward()  # All-reduce happens here
    optimizer.step()
    optimizer.zero_grad()

Tensor Parallelism: Split Within Layers

When a single layer is too large for one GPU, tensor parallelism splits weight matrices horizontally. Each GPU performs part of the matrix multiplication, then combines results via collective communication.

Column vs Row Parallel

In transformer MLPs, we pair column parallel (first linear) with row parallel (second linear):

1. Column Parallel: Split output dimension, each GPU computes partial features, all-gather to combine
2. Row Parallel: Split input dimension (already partitioned), each GPU computes partial output, all-reduce to sum

This pairing minimizes communication by doing one all-gather and one all-reduce per MLP block instead of two of each.

Pipeline Parallelism: Split by Stages

Pipeline parallelism partitions the model into sequential stages, with each GPU responsible for a subset of layers. Micro-batching keeps GPUs busy by processing multiple smaller batches simultaneously.

The Bubble Problem

In naive pipelining, GPUs sit idle waiting for activations from previous stages. This creates "bubbles" of wasted compute time. The bubble ratio depends on the schedule:

• Naive: ~75% bubble time with 4 stages
• GPipe: (p-1)/m bubble ratio (p=stages, m=micro-batches)
• 1F1B: (p-1)/(m+p-1) bubble ratio, plus lower memory usage

3D Parallelism: Combining All Three

For training the largest models (GPT-3, Llama 70B+), we combine all three strategies. The key insight is to match parallelism type to interconnect bandwidth:

# Megatron-LM style configuration
TP=8   # 8-way tensor parallel within DGX node (NVLink)
PP=4   # 4-way pipeline parallel across nodes
DP=2   # 2-way data parallel for larger batches
# Total: 8 × 4 × 2 = 64 GPUs

Communication Hierarchy

# Training Llama 70B on 64 H100 GPUs (8 nodes × 8 GPUs)
python train.py \
    --tensor-model-parallel-size 8 \   # TP=8 within node (NVLink)
    --pipeline-model-parallel-size 4 \  # PP=4 across nodes
    --data-parallel-size 2 \            # DP=2 for larger batches
    --micro-batch-size 1 \
    --global-batch-size 1024 \
    --num-layers 80 \
    --hidden-size 8192 \
    --num-attention-heads 64

# Total: 8 × 4 × 2 = 64 GPUs

ZeRO: Memory-Efficient Data Parallelism

ZeRO (Zero Redundancy Optimizer) reduces memory footprint by partitioning optimizer states, gradients, and parameters across data-parallel ranks instead of replicating them.

{
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": { "device": "cpu" },
    "offload_param": { "device": "cpu" },
    "overlap_comm": true,
    "contiguous_gradients": true
  }
}

Memory Breakdown

For a model with Φ parameters using Adam optimizer with mixed precision:

ZeRO Stages

• ZeRO-1: Partition optimizer states → 4Φ per GPU (4× memory reduction)
• ZeRO-2: + Partition gradients → 2.5Φ per GPU (8× memory reduction)
• ZeRO-3: + Partition parameters → ~1Φ per GPU (16× memory reduction)

{
  "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": "auto",
    "stage3_prefetch_bucket_size": "auto",
    "stage3_param_persistence_threshold": "auto"
  },
  "fp16": {
    "enabled": true,
    "loss_scale_window": 100
  }
}

Choosing the Right Strategy

Strategy Comparison

Further Reading

• Megatron-LM Paper: Efficient Large-Scale Training
• ZeRO: Memory Optimizations Toward Training Trillion Parameter Models
• GPipe: Easy Scaling with Micro-Batch Pipeline Parallelism
• PyTorch FSDP Tutorial
• DeepSpeed Documentation
• NVIDIA Megatron-LM Repository