DeepSpeed is Microsoft's open-source deep learning optimization library that makes distributed training and inference of massive models efficient and accessible. It addresses the fundamental problem of GPU memory limitations — modern LLMs require far more memory than any single GPU can hold — through the Zero Redundancy Optimizer (ZeRO) and a layered ecosystem of parallelism, offloading, and inference tools. The key mental model: ZeRO does not copy; it partitions — each GPU holds only its slice of optimizer states, gradients, or parameters, reducing per-GPU memory proportionally to the number of GPUs while keeping communication overhead manageable.
What This Cheat Sheet Covers
This topic spans 16 focused tables and 103 indexed concepts, 98 flashcards. Below is a complete table-by-table outline of this topic, spanning foundational concepts through advanced details.
A jump-to index of every table row in this cheat sheet.
An interactive map of every table and concept in this topic.
Table 1: ZeRO Optimizer Stages — Core Memory Partitioning
ZeRO (Zero Redundancy Optimizer) is the foundation of DeepSpeed's training efficiency. Each progressive stage partitions more of the model's memory footprint across data-parallel GPUs, with the critical distinction that only Stage 3 partitions the model parameters themselves.
| Technique | Example | Description | |
|---|---|---|---|
"zero_optimization": {"stage": 1, "reduce_bucket_size": 5e8} | • Partitions optimizer states (32-bit weights, Adam first and second moment estimates) across ranks • ~4x memory reduction for Adam • no extra communication vs. baseline | ||
"zero_optimization": {"stage": 2, "overlap_comm": true, "contiguous_gradients": true} | • Partitions optimizer states and gradients • each rank retains only the gradients matching its optimizer slice • same communication volume as Stage 1 — effectively a free upgrade | ||
"zero_optimization": {"stage": 3, "stage3_max_live_parameters": 1e9, "stage3_prefetch_bucket_size": 5e7} | • Partitions optimizer states, gradients, and model parameters • enables 16x+ memory reduction • parameters are all-gathered before use and discarded after each forward/backward pass | ||
"zero_optimization": {"stage": 0} | • ZeRO disabled • standard data parallelism with full replication of all states on every GPU • useful as a correctness baseline | ||
"reduce_bucket_size": 5e8 | • Maximum number of elements reduced/all-reduced at once • trades communication frequency for GPU memory • lower values save memory but increase round trips | ||
"allgather_bucket_size": 5e8 | • Maximum elements all-gathered per call in Stage 2/3 • lower values reduce peak memory at the cost of more communication calls | ||
"overlap_comm": true | • Overlaps gradient reduction with the backward computation pass • hides all-reduce latency • recommended to always enable on Stage 2 and 3. | ||
"contiguous_gradients": true | • Copies gradients into a contiguous buffer during the backward pass to reduce memory fragmentation • recommended for Stage 2 and 3. | ||
"reduce_scatter": true | • Uses reduce-scatter instead of all-reduce for gradient synchronization in Stage 2 • more bandwidth-efficient for large parameter counts |