Deep Dive - DeepSpeed

A model can fit in GPU memory for inference and still fail before its first training step. Training also keeps gradients, optimizer moments, master weights, and activations. For mixed-precision Adam, model states alone can consume roughly eight times the low-precision parameter file.

DeepSpeed turns those copies into explicit ownership decisions. Its best-known mechanism, ZeRO (Zero Redundancy Optimizer), partitions optimizer state, gradients, and eventually parameters across data-parallel ranks. The model's math stays the same; where each tensor lives and when it moves change.^{[1]Reference 1ZeRO: Memory Optimizations Toward Training Trillion Parameter Models.https://arxiv.org/abs/1910.02054}

What DeepSpeed takes over

Ordinary PyTorch code owns the model, optimizer, backward call, gradient synchronization, and step. deepspeed.initialize(...) wraps that stack in a DeepSpeedEngine. The engine tracks microsteps, accumulation boundaries, distributed groups, precision, gradient reduction, optimizer stepping, monitoring, and checkpoints.^{[2]Reference 2Training APIhttps://deepspeed.readthedocs.io/en/latest/training.html[3]Reference 3DeepSpeed Source Repositoryhttps://github.com/deepspeedai/DeepSpeed/tree/c455031422641a825588c926455458c82b27b6c0}

The user-facing loop stays small:

That compact API hides a serious contract. DeepSpeed must know global batch arithmetic, gradient-accumulation boundaries, data-parallel groups, parameter ownership, precision, and checkpoint layout. A wrong config can preserve execution while silently changing effective batch size, update frequency, or throughput.

💡 Key insight: DeepSpeed is an execution engine, not a memory-only switch. It owns enough of the training step to reschedule state and communication. That control makes config review and checkpoint testing necessary.

Build the memory ledger first

Let $\Psi$ be parameter count. A common mixed-precision Adam setup keeps low-precision parameters and gradients plus three FP32 arrays: master weights, first moment $m$, and second moment $v$.

This $16\Psi$ estimate is a planning model, not a universal allocator receipt. Optimizer implementation, gradient dtype, flat buffers, fragmentation, temporary gathers, and quantized states can change it. Activations, attention workspaces, and CUDA runtime memory sit outside the table.

For four data-parallel ranks, the idealized per-rank state becomes:

$$M_{Z1} = 4\Psi + \frac{12\Psi}{4} = 7\Psi,$$ $$M_{Z2} = 2\Psi + \frac{14\Psi}{4} = 5.5\Psi,$$ $$M_{Z3} = \frac{16\Psi}{4} = 4\Psi.$$

At stage 1, parameters and gradients remain replicated, costing $4\Psi$, while $12\Psi$ of optimizer and master state is sharded. Stage 2 also partitions the $2\Psi$ gradient block. The third stage shards every listed state. Temporary buffers and live layers make observed peaks higher.

The following copy-runnable estimator makes rank count and parameter count easy to change. It reports decimal gigabytes and leaves activation memory out on purpose.

The three ZeRO stages

Each stage adds one partitioned state type. Higher isn't automatically better. Choose the lowest stage that fits while meeting throughput and recovery targets.^{[4]Reference 4ZeRO Configurationhttps://deepspeed.readthedocs.io/en/stable/zero3.html}

Stage 1 is useful when Adam state dominates memory and full weights still fit. The second stage often gives a strong fine-tuning balance because parameters stay resident while gradients and optimizer state shard. Move to stage 3 when even replicated parameters are too large.

The original DeepSpeed ZeRO tutorial demonstrates a 1.5-billion-parameter GPT-2-style model on eight 32 GB V100 GPUs. It reports 18 GB of Adam state, then reduces that block to 2.25 GB per device with ZeRO-1 by partitioning it eight ways.^{[4]Reference 4ZeRO Configurationhttps://deepspeed.readthedocs.io/en/stable/zero3.html} Treat those numbers as one documented setup, not a multiplier for newer hardware or every optimizer.

Configure behavior, then verify arithmetic

A DeepSpeed JSON file selects precision, batch decomposition, optimizer, and ZeRO behavior. For a four-rank run, the configuration below uses BF16 and ZeRO-2.

The global batch identity is:

$$B_{global} = B_{micro} \times A \times N_{data} = 2 \times 8 \times 4 = 64.$$

$B_{micro}$ is samples per GPU per microstep, $A$ is accumulation steps, and $N_{data}$ is data-parallel ranks. If the launch uses eight ranks without changing another term, global batch doubles. Learning-rate and schedule assumptions may then be wrong even though training runs.

Bucket size controls another tradeoff. Small buckets start communication early but increase launch overhead. Large buckets use links efficiently but need more temporary memory and may finish too late to hide under backward compute. Measure exposed communication in a timeline rather than copying a round number from another model.

ZeRO-3 works one layer at a time

ZeRO-3 stores a parameter shard on each rank during steady state. Before a module computes, ranks all-gather that module's parameter shards into a full working copy. After the module no longer needs the full weights, DeepSpeed can release them and return to shards.^{[4]Reference 4ZeRO Configurationhttps://deepspeed.readthedocs.io/en/stable/zero3.html}

Backward produces gradients for the gathered layer. Reduce-scatter combines those gradients and leaves each rank with the shard it owns. Each rank updates its local optimizer and parameter shard, so no rank needs every Adam moment.

Prefetch and persistence settings decide how aggressively DeepSpeed overlaps future gathers and keeps small or frequently reused parameters resident. Too little prefetch exposes communication. Too much prefetch recreates the memory spike sharding was meant to avoid.

Dynamic modules need explicit boundaries

Stage 3 expects every rank to gather parameters in compatible order. Dynamic control flow, tied weights, or modules whose parameters are accessed outside their own forward() can break that assumption. DeepSpeed provides mechanisms such as external-parameter registration and leaf-module configuration for these cases.^{[4]Reference 4ZeRO Configurationhttps://deepspeed.readthedocs.io/en/stable/zero3.html}

A useful debugging reduction disables dynamic branches and runs the smallest deterministic batch. Compare loss and parameter updates against an unsharded baseline. Restore model features one at a time after the gather order is stable.

Offload trades capacity for a slower tier

ZeRO-Offload moves optimizer state and optimizer computation toward host CPU memory. ZeRO-Infinity extends the hierarchy so parameters and optimizer state can use CPU memory or NVMe with staging and overlap.^{[5]Reference 5ZeRO-Offload: Democratizing Billion-Scale Model Traininghttps://www.usenix.org/conference/atc21/presentation/ren-jie[6]Reference 6ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning.https://arxiv.org/abs/2104.07857}

The configuration can look small:

The performance model isn't small. A step may wait on GPU-to-CPU DMA, CPU Adam work, pinned-memory pressure, NVMe queue depth, filesystem contention, or NUMA placement. Offload can make a model fit while reducing tokens per second enough to miss the training deadline.

Treat the memory hierarchy as a pipeline:

⚠️ Common mistake: “Fits” isn't a throughput result. Record step time, tokens per second, exposed transfer time, CPU utilization, and storage bandwidth after every offload change.

NCCL moves the shards DeepSpeed schedules

DeepSpeed decides which data-parallel group owns a shard, when a bucket is ready, and whether communication should overlap. NCCL executes the GPU collectives on NVIDIA systems. The two projects solve adjacent layers.

ZeRO-2 commonly replaces one replicated gradient all-reduce result with reduce-scatter ownership. ZeRO-3 adds parameter all-gathers around module execution. ZeRO++ adds communication-oriented options such as quantized weight exchange, hierarchical partition groups, and quantized gradients to reduce pressure on slower links.^{[4]Reference 4ZeRO Configurationhttps://deepspeed.readthedocs.io/en/stable/zero3.html}

If a stage change hurts throughput, inspect both sides. DeepSpeed logs and profiler traces show bucket timing and overlap. NCCL traces and network counters show algorithm, route, and achieved communication. Tuning only one layer can misdiagnose a topology problem as a bucket problem.

DeepSpeed is broader than ZeRO

ZeRO is the center of gravity, but the repository contains more systems:

DeepSpeed's built-in PipelineEngine is not a promise that every ZeRO stage composes with every pipeline setup. In the pinned snapshot, it rejects ZeRO-2 and ZeRO-3 for that engine path. Separate Megatron-DeepSpeed integrations and other runtimes have different composition contracts.^{[3]Reference 3DeepSpeed Source Repositoryhttps://github.com/deepspeedai/DeepSpeed/tree/c455031422641a825588c926455458c82b27b6c0}

Feature matrices need source and version context. “DeepSpeed supports pipeline plus ZeRO” is too broad to choose a runtime path. Ask which engine, model wrapper, ZeRO stage, checkpoint format, and tested release combination is involved.

Read initialization into the optimizer

Start at deepspeed/__init__.py. initialize(...) selects the communication backend, initializes distributed state, parses config, and returns the engine plus wrapped optimizer, data loader, and scheduler.^{[3]Reference 3DeepSpeed Source Repositoryhttps://github.com/deepspeedai/DeepSpeed/tree/c455031422641a825588c926455458c82b27b6c0}

Then follow the selected stage:

• Stage 1 or 2: runtime/zero/stage_1_and_2.py flattens groups, creates partitions, registers gradient handling, and coordinates reductions.
• Stage 3: runtime/zero/stage3.py builds sharded optimizer state while partition_parameters.py manages parameter status, gathers, and repartitioning.
• CPU or NVMe offload: runtime/swap_tensor/, asynchronous I/O operators, and CPU optimizer code expose the slower memory tier.
• Pipeline mode: runtime/pipe/engine.py owns stage IDs, neighbors, microbatches, and pipeline-specific reduction scheduling.

The source map prevents a common reading mistake. DeepSpeedEngine.step() isn't where every optimization lives. It delegates to optimizer wrappers, communication layers, hooks, swap engines, and checkpoint utilities selected at initialization.

Checkpoints are part of the runtime design

A ZeRO checkpoint is distributed training state, not automatically one portable model file. Each rank can write shards of parameters, optimizer state, scheduler state, and metadata. World size and parallel layout affect how those shards are interpreted.

DeepSpeed's ZeRO tutorial documents zero_to_fp32.py and get_fp32_state_dict_from_zero_checkpoint for consolidating ZeRO checkpoints.^{[4]Reference 4ZeRO Configurationhttps://deepspeed.readthedocs.io/en/stable/zero3.html} Universal Checkpointing targets resume across different parallel configurations by converting state into a topology-independent form.^{[7]Reference 7Universal Checkpointing with DeepSpeed: A Practical Guide.https://www.deepspeed.ai/tutorials/universal-checkpointing/}

Test four paths before a long run:

1. Save a checkpoint at the intended scale.
2. Kill every worker and restore training.
3. Verify loss, scheduler, optimizer step, and sample counters continue correctly.
4. Export a consolidated model and load it in the evaluation or serving runtime.

A directory existing isn't enough. Checkpoint success means the resumed job produces the expected next update and the exported model matches a reference output within the chosen precision tolerance.

DeepSpeed, FSDP, and Megatron answer different questions

These systems overlap, but their default centers differ.

FSDP and ZeRO-3 share a core idea: parameters, gradients, and optimizer state can remain sharded until computation needs them. API boundaries, scheduling, checkpoint formats, mixed-precision policies, and ecosystem integration differ.

Megatron can be paired with distributed optimizer or sharding choices and DeepSpeed-based integrations. Don't select by project logo. Select from model shape, network topology, framework ownership, checkpoint requirements, team expertise, and measured end-to-end throughput.

Production debugging loop

Begin with a small unsharded reference on fixed data. Record loss, gradient norm, optimizer step, global batch, precision, and a checkpoint round trip. That reference tells you whether a later failure changed math or only systems behavior.

Add one mechanism at a time:

For a hang, find the earliest rank error before reading peer timeout stacks. Loss divergence calls for one-update comparison with identical data and math. Throughput drops need a step-time split across compute, collectives, CPU work, I/O, and idle gaps.

Useful questions stay concrete:

• Did global batch change with world size?
• Are activations, model state, or temporary gathers causing the peak?
• Does every rank cross the same accumulation boundary?
• Are all-gathers hidden under compute or exposed?
• Did offload move the bottleneck to CPU, PCIe, or storage?
• Can the latest checkpoint restore after every worker exits?

Project origin, governance, and license

DeepSpeed began in Microsoft Research and was part of Microsoft's AI at Scale work. The current project lives in the deepspeedai/DeepSpeed organization and uses a Technical Steering Committee (TSC) under AI & Data, a directed fund of The Linux Foundation.^{[8]Reference 8DeepSpeed Project Charter and Governancehttps://github.com/deepspeedai/DeepSpeed/blob/c455031422641a825588c926455458c82b27b6c0/GOVERNANCE.md}

The pinned committer roster shows a cross-organization group rather than one company team:^{[9]Reference 9DeepSpeed TSC Committershttps://github.com/deepspeedai/DeepSpeed/blob/c455031422641a825588c926455458c82b27b6c0/COMMITTERS.md}

The TSC can add or remove committers and governs technical direction. Public contributions use Developer Certificate of Origin sign-off. DeepSpeed code uses Apache License 2.0; project governance places documentation under Creative Commons Attribution 4.0.^{[8]Reference 8DeepSpeed Project Charter and Governancehttps://github.com/deepspeedai/DeepSpeed/blob/c455031422641a825588c926455458c82b27b6c0/GOVERNANCE.md[10]Reference 10DeepSpeed Apache License 2.0https://github.com/deepspeedai/DeepSpeed/blob/c455031422641a825588c926455458c82b27b6c0/LICENSE}

Papers behind the system

DeepSpeed grew through a sequence of systems papers rather than one isolated algorithm.

ZeRO names the partitioning family.^{[1]Reference 1ZeRO: Memory Optimizations Toward Training Trillion Parameter Models.https://arxiv.org/abs/1910.02054} The KDD tutorial describes DeepSpeed as a broader system of training optimizations.^{[11]Reference 11DeepSpeed: System Optimizations Enable Training Deep Learning Models with Over 100 Billion Parametershttps://dl.acm.org/doi/10.1145/3394486.3406703} ZeRO-Offload and ZeRO-Infinity push ownership beyond GPU memory, which is why modern DeepSpeed tuning includes DMA and storage measurements alongside GPU kernels.^{[5]Reference 5ZeRO-Offload: Democratizing Billion-Scale Model Traininghttps://www.usenix.org/conference/atc21/presentation/ren-jie[6]Reference 6ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning.https://arxiv.org/abs/2104.07857}

Research results establish that a design can work under reported conditions. A production decision still needs your model, network, optimizer, precision, sequence length, checkpoint cadence, and failure policy.

Strengths and weaknesses

DeepSpeed's strengths are practical:

• ZeRO stages expose a clear progression from optimizer sharding to full model-state sharding.
• CPU and NVMe offload make capacity a tiered-memory problem.
• Engine integration keeps the user loop compact.
• Broad features cover training, inference, pipeline execution, MoE, profiling, and accelerator backends.
• Framework integrations make ZeRO reachable from common model stacks.
• Public governance and source provide code-level debugging paths.

Its weaknesses come from that same breadth:

• Large config and feature surface creates interactions that aren't all valid together.
• Engine ownership can make framework-native debugging or migration harder.
• ZeRO-3 replaces memory replication with frequent communication and temporary gathers.
• Offload can trade an OOM for an unacceptable step time.
• Sharded checkpoints require explicit recovery and export work.
• JIT-compiled operators and accelerator compatibility add installation risk.
• Documentation examples can outlive their exact hardware and version assumptions.

Choose DeepSpeed when its engine, ZeRO, or offload features solve a measured constraint and your team can own the integration. A smaller native stack can be easier when the model already fits and advanced DeepSpeed features add little.

Build a decision receipt

Run the estimator with your parameter count and data-parallel degree. Add measured activation peak, temporary gather buffers, and a safety margin. Predict the lowest ZeRO stage that fits.

Then run a short controlled experiment for plain DDP, your chosen ZeRO stage, and one fallback stage. Keep model, data order, optimizer, precision, and global batch fixed. Record:

• Maximum allocated and reserved GPU memory per rank
• Tokens per second and step-time distribution
• Collective time split by reduce-scatter and all-gather
• CPU, PCIe, and NVMe use if offload is enabled
• Loss and gradient parity over several updates
• Save, kill, restore, and export results

Finish with one sentence: “We chose stage X because constraint Y improved by Z, while cost W stayed within limit Q.” That receipt is more useful than a copied config because it explains what the system proved.

Takeaways

• DeepSpeed wraps the training step so it can control state ownership, collectives, accumulation, precision, and checkpoints.
• Mixed-precision Adam can require about 16 bytes of model state per parameter before activations.
• ZeRO-1 shards optimizer state, ZeRO-2 adds gradients, and ZeRO-3 adds parameters.
• Stage 3 saves steady-state memory by gathering each working set near layer execution.
• Offload expands capacity but spends CPU, interconnect, and storage bandwidth.
• NCCL moves bytes; DeepSpeed decides groups, buckets, and schedule.
• Checkpoint restore and portable export are runtime requirements, not cleanup tasks.
• Current governance is cross-company under Linux Foundation AI & Data, with Apache-2.0 code.

Mastery quiz