Distributed Training

Distributed Training

Training Paradigms

Data Parallel (DDP)
Each GPU has full model copy, processes different data:

GPU 0: Model copy → Batch 1 → Gradients
GPU 1: Model copy → Batch 2 → Gradients   → AllReduce → Update
GPU 2: Model copy → Batch 3 → Gradients

Model Parallel
Split model across GPUs:
- Tensor Parallel: Split layers across GPUs
- Pipeline Parallel: Split layers sequentially
- Expert Parallel: Split MoE experts

PyTorch DDP

Basic Setup

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

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

# Wrap model
model = YourModel().to(local_rank)
model = DDP(model, device_ids=[local_rank])

# Use DistributedSampler
sampler = DistributedSampler(dataset)
dataloader = DataLoader(dataset, sampler=sampler)

Launch

torchrun --nproc_per_node=4 train.py

FSDP (Fully Sharded Data Parallel)

Why FSDP?
- DDP requires full model on each GPU
- FSDP shards model parameters, gradients, and optimizer states
- Enables training models larger than single GPU memory

Usage

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

model = FSDP(
    model,
    sharding_strategy=ShardingStrategy.FULL_SHARD,
    mixed_precision=MixedPrecision(
        param_dtype=torch.bfloat16,
        reduce_dtype=torch.bfloat16,
        buffer_dtype=torch.bfloat16,
    ),
)

Comparison
| Method | Model Size Limit | Memory Efficiency | Complexity |
|--------|------------------|-------------------|------------|
| DDP | Single GPU memory | Low | Low |
| FSDP | Multi-GPU combined | High | Medium |
| DeepSpeed ZeRO | Multi-GPU combined | Highest | Medium |

Communication Backends
| Backend | Use Case |
|---------|----------|
| NCCL | GPU-to-GPU (preferred) |
| Gloo | CPU or fallback |
| MPI | HPC environments |