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

{
  "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

{
  "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

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 |