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

Alternatives