DeepSpeed Explained: ZeRO Stages, Memory Savings, and Common Misconceptions

Forum topic · 小凯 · 2026-03-27

Summary

DeepSpeed is Microsoft's deep learning optimization library that leverages the ZeRO (Zero Redundancy Optimizer) technique and parallelism strategies to make large-scale model training more efficient. This tutorial from the Easy AI series explains how ZeRO shards model parameters, gradients, and optimizer states across GPUs to reduce memory redundancy. It addresses a common misconception: multi-GPU training does not automatically split memory across cards—by default each GPU stores the full model and optimizer states, and multiple GPUs mainly improve training speed. The article compares DeepSpeed ZeRO stages with example per-GPU memory footprints: Stage 0 (baseline, ~15.25 GB), Stage 1 (parameter sharding, ~12.6 GB), Stage 2 (optimizer state sharding, ~10.4 GB, recommended), and Stage 3 (layer-wise model partitioning, ~8.15 GB, maximum savings). Core features covered include mixed-precision training with dynamic loss scaling, combinations of tensor/pipeline/data parallelism, activation recomputation, and memory defragmentation.

DeepSpeed Explained

1. What is DeepSpeed?

1.1 Core Concept

DeepSpeed is a deep learning optimization library developed by Microsoft. Through the ZeRO (Zero Redundancy Optimizer) technique and various parallelism strategies, it makes large-scale model training more efficient and easier to use.

1.2 Key Advantages

| Advantage | Description | |-----------|-------------| | GPU memory savings | Per-GPU memory usage drops significantly as the number of GPUs increases | | Training speedup | Efficient parallelism strategies and memory management | | Model scaling | Supports training of much larger models |

1.3 ZeRO Architecture

Sharded storage of model states eliminates memory redundancy:

| Sharding Type | Description | |---------------|-------------| | Parameter sharding (P) | Model parameters are distributed across different GPUs | | Gradient sharding (G) | Distributed computation and storage of gradients | | Optimizer sharding (O) | Distributed storage and management of optimizer states |

---

2. Common Misconception vs. Reality

2.1 Misconception: Memory is automatically split

❌ Wrong idea:

> "When training on multiple GPUs, memory is automatically split evenly, so each card only bears 1/N of the memory load."

2.2 Reality: Each GPU is independent

✅ Actual situation:

> "Every GPU stores the full model parameters and optimizer states—memory is not automatically split." > "Multiple GPUs are mainly used to increase training speed."

---

3. DeepSpeed ZeRO Stage Comparison

Stage 0 (default)

Every GPU stores the complete model
Per-GPU memory: ~15.25 GB

Stage 1

Parameters stored once, sharded across all GPUs
Per-GPU memory: ~12.6 GB

Stage 2 ⭐ Recommended

Parameters shared, optimizer states distributed
Per-GPU memory: ~10.4 GB
Recommended choice!

Stage 3

Model split by layers; parameters flow on demand
Per-GPU memory: ~8.15 GB
Maximum memory savings

---

4. Core Features

ZeRO technology: sharding of parameters / gradients / optimizer states
Mixed-precision training: automatic loss scaling and dynamic precision management
Parallelism strategies: combinations of tensor, pipeline, and data parallelism
Memory optimization: activation recomputation and memory defragmentation

---

*Source: Easy AI Tutorial Series*