zero
**DeepSpeed ZeRO**
**What is ZeRO?**
ZeRO (Zero Redundancy Optimizer) is a memory optimization technique that partitions optimizer states, gradients, and parameters across GPUs.
**ZeRO Stages**
**Memory Distribution**
| Stage | What's Sharded | Memory Reduction |
|-------|----------------|------------------|
| ZeRO-1 | Optimizer states | 4x |
| ZeRO-2 | + Gradients | 8x |
| ZeRO-3 | + Parameters | Linear with GPUs |
**Stage Comparison**
```
Standard DDP:
GPU 0: Params + Grads + OptStates (full copy)
GPU 1: Params + Grads + OptStates (full copy)
ZeRO-3:
GPU 0: Params0 + Grads0 + OptStates0
GPU 1: Params1 + Grads1 + OptStates1
(Gather parameters when needed, shard after)
```
**DeepSpeed Configuration**
**Basic ZeRO-2 Config**
```json
{
"zero_optimization": {
"stage": 2,
"offload_optimizer": {"device": "cpu"},
"contiguous_gradients": true,
"overlap_comm": true
},
"bf16": {"enabled": true},
"train_batch_size": 32,
"gradient_accumulation_steps": 4
}
```
**ZeRO-3 Config**
```json
{
"zero_optimization": {
"stage": 3,
"offload_param": {"device": "cpu"},
"offload_optimizer": {"device": "cpu"},
"stage3_gather_16bit_weights_on_model_save": true
}
}
```
**CPU/NVMe Offload**
Offload optimizer states and optionally weights to CPU RAM or NVMe:
- Enables training 10x larger models
- Trade-off: slower due to CPU↔GPU transfers
- Use for memory-limited scenarios
**Usage with Hugging Face**
```python
from transformers import TrainingArguments
args = TrainingArguments(
deepspeed="ds_config.json",
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
bf16=True,
)
```
**When to Use Each Stage**
| Scenario | Recommended |
|----------|-------------|
| 7B on 4x A100 40GB | ZeRO-2 |
| 13B on 4x A100 40GB | ZeRO-2 + CPU offload |
| 70B on 8x A100 80GB | ZeRO-3 |
| 70B on 4x RTX 4090 | ZeRO-3 + CPU offload |
**Alternatives**
| Library | Highlights |
|---------|------------|
| DeepSpeed | Most features, Microsoft |
| FSDP | PyTorch native, simpler |
| FairScale | Meta, FSDP precursor |
| Megatron-LM | NVIDIA, large-scale |