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**
```python
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**
```bash
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**
```python
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 |