DeepSpeed

Part I: The mental model

deepspeed launcher        one process per GPU, sets RANK/LOCAL_RANK/WORLD_SIZE
      |
      v
your train.py             plain PyTorch model + a ds_config dict/JSON
      |
      v
deepspeed.initialize()    wraps model -> DeepSpeedEngine (deepspeed/runtime/engine.py)
      |                    reads zero_optimization.stage, offload, fp16/bf16
      v
ZeRO partitioning         optimizer state (st.1), + gradients (st.2), + params (st.3)
      |                    each shard lives on exactly one data-parallel rank
      v
engine(batch)             forward: all-gather each layer's params just in time (st.3)
      |
      v
engine.backward(loss)     backward: reduce-scatter gradients, re-partition params
      |
      v
engine.step()             optimizer step on the local shard (maybe on CPU/NVMe)
      |                    loss scaling, grad clipping, grad accumulation handled here
      v
collectives (NCCL)        all-gather, reduce-scatter; offload copies over PCIe/NVMe

The one-sentence identity: DeepSpeed is a training engine that wraps an ordinary PyTorch model and removes the redundant memory that data parallelism duplicates on every GPU, by partitioning the optimizer state, the gradients, and the parameters across the group and gathering each piece back only for the moment it is needed. Ordinary data parallelism replicates everything and communicates only gradients. That is simple and fast, but it caps model size at whatever fits on one device, and for a modern model the optimizer state alone is several times larger than the parameters. ZeRO keeps the single-program simplicity of data parallelism while paying the memory cost of only one Nth of the state per rank, where N is the data-parallel degree.

Two consequences follow, and they are the whole reason the library exists. First, memory stops being the wall. A model that needs 120 GB of optimizer, gradient, and parameter state under plain data parallelism needs under 2 GB per GPU under ZeRO-3 on 64 devices, and when even that is too much, ZeRO-Offload and ZeRO-Infinity spill the partitions to CPU DRAM and then to NVMe, trading bandwidth for capacity along the memory hierarchy. Second, the model code barely changes. You do not rewrite your transformer into column-parallel and row-parallel layers the way a tensor-parallel framework demands. You hand DeepSpeed a plain model and a JSON config, and the engine applies the partitioning from outside. Everything here is verified against the library as it stands in mid 2026, and the project moves quickly, so where a file path or a config key is likely to have shifted I say so and stay at the level of the concept.

One honest orientation before the details. The idea at the center of ZeRO-3, partition the parameters and gather them per layer on the fly, is exactly the idea PyTorch later absorbed as FSDP, and FSDP2 is what torchtitan builds on. Reading DeepSpeed is the shortest path to understanding where that whole family of techniques came from and why it works. If you want the single-device transformer these ideas are applied to, the language models from scratch class builds one end to end, and the broader theory of who-shards-what lives in the parallel computing class.

Part II: Using it

DeepSpeed is a Linux-and-NVIDIA-GPU library first (with ROCm and some CPU and accelerator backends). It installs from PyPI on top of an existing PyTorch, and it wants a CUDA toolkit present because most of its speed comes from custom C++ and CUDA operators it compiles for your machine:

pip install deepspeed

# sanity check: which ops are compatible with this box, and which are
# prebuilt vs will be JIT-compiled on first use
ds_report

ds_report is the first thing to run and the most useful diagnostic in the project. It prints a table of every operator (fused Adam, CPU Adam, the transformer kernels, the inference kernels, and so on), whether your toolchain can build it, and whether it was precompiled at install time or will be just-in-time compiled the first time you touch it. If CPU Adam shows as incompatible, ZeRO-Offload will fall back or fail, and ds_report tells you before a six-hour run does.

The rest of DeepSpeed is driven by two things: a config, and the deepspeed.initialize call that consumes it. The config is a plain dictionary (usually stored as a JSON file, but a Python dict works identically). A minimal ZeRO-3 config with optimizer offload to CPU and parameter offload to NVMe looks like this:

ds_config = {
    "train_micro_batch_size_per_gpu": 4,
    "gradient_accumulation_steps": 8,
    "optimizer": {"type": "AdamW", "params": {"lr": 3e-4}},
    "bf16": {"enabled": True},
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {"device": "cpu", "pin_memory": True},
        "offload_param": {"device": "nvme", "nvme_path": "/local_nvme"},
        "overlap_comm": True,
        "contiguous_gradients": True,
        "stage3_prefetch_bucket_size": 5e8,
        "stage3_param_persistence_threshold": 1e6,
        "stage3_gather_16bit_weights_on_model_save": True,
    },
}

The training code is deliberately close to plain PyTorch, with three method calls replacing the usual four. You do not call loss.backward(), optimizer.step(), and optimizer.zero_grad() yourself. The engine owns those, because it has to interleave gradient reduction, loss scaling, gradient accumulation, and clipping around them:

import deepspeed

model = MyTransformer()          # an ordinary nn.Module

model_engine, optimizer, dataloader, _ = deepspeed.initialize(
    model=model,
    model_parameters=model.parameters(),
    training_data=train_dataset,
    config=ds_config,            # dict or path to ds_config.json
)

for batch in dataloader:
    loss = model_engine(batch)       # forward
    model_engine.backward(loss)      # backward + gradient reduce-scatter
    model_engine.step()              # optimizer step + zero grad (at accum boundary)

deepspeed.initialize returns a four-tuple: the wrapped engine (which is itself an nn.Module you call to run forward), the optimizer it built or adopted, a dataloader if you passed a dataset, and an LR scheduler if the config asked for one. Note that model_engine.step() is called every iteration, but the engine only performs the real optimizer update at gradient-accumulation boundaries. On the intervening micro-steps it just accumulates. That is why train_batch_size in the config equals train_micro_batch_size_per_gpu times gradient_accumulation_steps times the world size, and DeepSpeed will refuse to start if you specify all three and they do not multiply out.

You launch with the deepspeed command, which is a process launcher analogous to torchrun. It reads a GPU count or a multi-node hostfile and spawns one rank per GPU:

# single node, 8 GPUs
deepspeed --num_gpus=8 train.py --deepspeed --deepspeed_config ds_config.json

# multiple nodes: hostfile lists "hostname slots=N" lines
deepspeed --hostfile=hostfile train.py --deepspeed_config ds_config.json

Now the mistakes people make. First, the batch-size identity above trips everyone once. If you set train_batch_size, train_micro_batch_size_per_gpu, and gradient_accumulation_steps to values that are inconsistent with the world size, DeepSpeed asserts at init. The fix is to specify at most two of the three and let the engine derive the rest. Second, for very large models under ZeRO-3 you must construct the model under a deepspeed.zero.Init() context, otherwise every rank allocates the full model before partitioning and you are back to the out-of-memory you were trying to avoid:

import deepspeed

with deepspeed.zero.Init(config_dict_or_path=ds_config):
    model = VeryLargeTransformer()   # params are sharded as they are created,
                                     # so the full model never lands on one GPU

Third, saving a ZeRO-3 checkpoint the naive way gives you shards, not a model. Each rank holds a partition, so model_engine.save_checkpoint(dir) writes sharded state that only DeepSpeed can resume. To get a single consolidated file you either set stage3_gather_16bit_weights_on_model_save and call the 16-bit save helper, or run the zero_to_fp32.py conversion script the engine drops next to the checkpoint. Fourth, and most common with offload: NVMe offload needs a real fast local NVMe path and the async I/O operator built, and putting nvme_path on a network filesystem turns a training run into a disk-latency benchmark.

Finally, note that you rarely write this loop by hand in practice. Hugging Face Transformers and Accelerate both integrate DeepSpeed, so passing a ZeRO config to the Trainer gets you all of the above without the explicit initialize call. Reading the raw API first is still the right way to understand what that integration is doing under you.

Part III: The ZeRO memory math, derived

Everything about ZeRO follows from one accounting exercise, so it is worth doing carefully rather than quoting. Take a model with Ψ parameters trained in mixed precision with the Adam family of optimizers, which is the standard setup. Count the bytes that must live in memory for one optimizer step, per parameter:

The three optimizer rows sum to 12 bytes per parameter. The ZeRO paper calls that multiplier K = 12. So the total training memory for the model state, ignoring activations, is (2 + 2 + K)Ψ = 16Ψ bytes. The crucial observation is that the 12 bytes of optimizer state, three quarters of the total, are Adam bookkeeping that is only ever touched during the optimizer step, and under plain data parallelism every GPU holds an identical copy of all of it. That is the redundancy ZeRO removes, in three stages, by partitioning across the N data-parallel ranks:

The last column is the canonical worked example from the ZeRO paper: a 7.5-billion-parameter model on 64 GPUs. Baseline needs 16 × 7.5 = 120 GB per GPU, which fits on nothing. ZeRO-1 partitions only the 12 bytes of optimizer state, giving 4 × 7.5 + (12 × 7.5)/64 ≈ 31.4 GB. ZeRO-2 also partitions the 2 bytes of gradients, giving 2 × 7.5 + (14 × 7.5)/64 ≈ 16.6 GB, a number that now fits on a commodity accelerator. ZeRO-3 partitions the parameters too, and the whole 16Ψ divides by N cleanly, giving (16 × 7.5)/64 ≈ 1.9 GB. As N grows, ZeRO-3 memory goes to zero, so the model you can train is limited by aggregate memory across all GPUs rather than the memory of any one, which is the entire reason trillion-parameter training became possible.

The natural objection is that all this gathering and scattering must cost communication. It costs less than you would guess. Standard data-parallel training does one all-reduce of the gradients per step, which a ring implementation realizes as a reduce-scatter followed by an all-gather, moving 2Ψ bytes total. ZeRO-1 and ZeRO-2 keep exactly that 2Ψ volume: they reduce-scatter gradients to the rank that owns each shard, that rank updates its slice of parameters, and an all-gather redistributes the updated parameters. Same traffic as ordinary data parallelism, a fraction of the memory. ZeRO-3 is the only stage that costs more. Parameters now have to be all-gathered once in the forward pass and again in the backward pass, and gradients reduce-scattered, so the volume is 3Ψ, which is 1.5 times the baseline. The trade is a 50 percent increase in communication for an N-fold reduction in memory, and on a fast interconnect with good compute overlap that trade is almost always worth taking.

communication volume per step, relative to plain data parallel (= 2Ψ)

   data parallel   ZeRO-1   ZeRO-2   ZeRO-3
        1.0x         1.0x     1.0x     1.5x
                    <--- same traffic -->   <- 1.5x, all-gather params twice ->

One more subtlety worth stating, because it explains why ZeRO-3 is not free even when overlap hides the bandwidth. Gathering a layer's parameters just before it runs means the parameters exist in full for only a brief window, so peak memory depends on how many layers you allow to be live at once (stage3_max_live_parameters) and how far ahead you prefetch (stage3_prefetch_bucket_size). Tune those too aggressively and you recreate the memory pressure you were escaping. Tune them too conservatively and prefetch cannot hide the all-gather latency. That tension is the essence of operating ZeRO-3 well.

Part IV: When it is the right tool

DeepSpeed is the right tool when memory is your binding constraint and you want to keep a mostly plain PyTorch model. That covers pretraining or full fine-tuning of a model whose optimizer state will not fit under plain data parallelism, training on a single node (or even a single GPU) a model that is nominally too big for it via offload, and any workflow already sitting inside the Hugging Face ecosystem where a ZeRO config is a few lines in a Trainer argument. It is also the reference implementation of the offload idea. If you want to train a 13B model on one consumer GPU by spilling optimizer state to CPU, this is the library that pioneered that path.

The honest alternatives, named plainly. PyTorch FSDP and FSDP2 are the same core idea as ZeRO-3, absorbed into PyTorch itself, and if you want a native, forkable stack that composes sharding with tensor and pipeline parallelism over DTensor, torchtitan built on FSDP2 is cleaner to read and own. Megatron-LM is the tool when you need maximum model-FLOP utilization on large NVIDIA clusters and you accept tensor parallelism that rewrites your model into parallel layers. The two are not even really rivals in practice. The largest public models, Megatron-Turing NLG 530B and BLOOM 176B, were trained with Megatron-DeepSpeed, which uses Megatron for tensor parallelism and DeepSpeed for ZeRO data parallelism and pipelining, the so-called 3D parallelism. For inference rather than training, DeepSpeed-Inference exists and is covered below, but the momentum in open-source serving has moved to vLLM and similar engines built around paged KV caches, so reach for those first unless you specifically need ZeRO-Inference to fit a model that will not otherwise load.

The architecture-shaped warning is about where the offload bandwidth lands. Every stage of ZeRO trades communication for memory, and every level of offload trades a slower link for more capacity. GPU memory is fed by NVLink, CPU memory by PCIe, and NVMe by PCIe plus a storage controller, each an order of magnitude slower than the last. Offloading optimizer state to CPU is often nearly free because the optimizer step is a small fraction of step time and overlaps with compute. Offloading parameters to NVMe under ZeRO-Infinity can be the difference between running and not running, but it can also dominate step time if the NVMe is slow or the working set does not fit in the CPU cache tier above it:

memory hierarchy, fastest and smallest at the top

   HBM on GPU        ~TB/s      <- params live here only while a layer runs
        |  NVLink / PCIe
   CPU DRAM          ~tens GB/s <- ZeRO-Offload puts optimizer state here
        |  PCIe + controller
   local NVMe        ~GB/s      <- ZeRO-Infinity spills params/optimizer here

   rule: push state down only as far as capacity forces you,
         because each step down costs an order of magnitude of bandwidth

The failure mode is not an error. It is a run that technically works and is ten times slower than it should be because the offload target cannot keep the GPUs fed. As with mapping tensor parallelism across slow links, the symptom is step time, and the profiler is how you find it.

Part V: The full life of one ZeRO-3 training step

The specimen is one iteration of a ZeRO-3 model with optimizer offload to CPU, launched with the config from Part II. Most of the machinery below is orchestrated by the DeepSpeedEngine in deepspeed/runtime/engine.py and the stage-3 implementation under deepspeed/runtime/zero/. Exact file names in that package have been reorganized more than once, so I name the roles and note where I am confident of a filename.

Stage 1: the launcher and initialize

The deepspeed command (its runner lives in deepspeed/launcher/) reads the GPU count or hostfile, sets up the environment, and execs one copy of your train.py per GPU with RANK, LOCAL_RANK, and WORLD_SIZE populated, exactly like torchrun. From there the program is SPMD, the same code on every rank. Your call to deepspeed.initialize parses the config, initializes the distributed backend through DeepSpeed's communication wrapper (deepspeed/comm/, a thin layer over torch.distributed and NCCL), and constructs the engine appropriate to the configured ZeRO stage. For stage 3 that is the stage-3 engine, which immediately partitions the model's parameters so each rank keeps only its slice.

Stage 2: parameter partitioning and the hooks

The heart of ZeRO-3 is that every parameter carries a partition. If the model was built under deepspeed.zero.Init(), the parameters were sharded as they were created and the full tensor never existed on one GPU. Otherwise the engine partitions them now. The stage-3 machinery (the parameter partitioning and offload logic, historically in files like partition_parameters.py and a parameter offload coordinator in deepspeed/runtime/zero/) then registers forward and backward hooks on every submodule. Those hooks are the trick that turns a partitioned model into a runnable one. Before a submodule runs, its pre-forward hook all-gathers the full parameters for just that module. After it runs, its post-forward hook releases them back to the partitioned state. Parameters below stage3_param_persistence_threshold in size are left resident permanently, because gathering tiny tensors costs more in latency than it saves in memory.

Stage 3: forward with just-in-time all-gather

When you call model_engine(batch), execution walks the module tree as usual, but each module's pre-forward hook fires an all-gather across the data-parallel group to reconstruct that module's parameters, the module computes, and the post-forward hook frees them again. The engine prefetches. While layer L computes, the all-gather for layer L+1 is already in flight on a separate stream, so the communication hides under the compute of the previous layer. This is the same overlap discipline that makes FSDP fast, and it is why overlap_comm and the prefetch-bucket settings matter so much. The result is a normal forward pass and a normal loss, computed as if the full model were resident, while at no instant is more than a couple of layers' worth of parameters actually materialized.

Stage 4: backward with reduce-scatter

model_engine.backward(loss) runs autograd, but with the same hook discipline mirrored. Each module's parameters are re-gathered for its backward, gradients are computed locally, and then, instead of the all-reduce that plain data parallelism would do, ZeRO reduce-scatters the gradients so that each rank ends up holding only the gradient shard for the parameters it owns. The full gradient for a parameter never exists on more than the ranks that momentarily held it. Gradients are accumulated into contiguous buffers when contiguous_gradients is set, which avoids fragmentation and makes the reduce-scatter a single large collective rather than many small ones. If gradient accumulation is configured, backward on the intervening micro-steps simply accumulates locally and skips the reduction until the accumulation boundary, which is why the communication volume is per effective batch, not per micro-batch.

Stage 5: the optimizer step, possibly on the CPU

model_engine.step() at an accumulation boundary is where the offload story pays off. Each rank owns a shard of the gradients and a shard of the optimizer state. With offload_optimizer.device = "cpu", the fp32 master parameters, the Adam momentum, and the variance live in pinned CPU memory, so the gradient shard is copied over PCIe to the host, and the Adam update runs on the CPU using DeepSpeed's hand-written DeepSpeedCPUAdam, a SIMD-vectorized C++ optimizer that keeps the CPU step from becoming the bottleneck. The updated fp32 parameters are cast back to bf16 and copied to the GPU shard. Before the parameters are used again in the next forward, ZeRO-3's all-gather makes the updated values visible to every rank. Loss scaling for fp16, gradient clipping by global norm (which requires reducing partial norms across the shards), and the LR schedule advance all happen inside this one step call. Nothing in your training loop hints at any of it.

Stage 6: what actually crossed the wire

Tally the traffic for the step. In the forward pass, one all-gather per module to reconstruct parameters. In the backward pass, one all-gather per module again, plus one reduce-scatter of gradients. That is the 3Ψ of Part III, 1.5 times plain data parallelism, spread across many small collectives that the prefetcher overlaps with compute. On top of the NCCL traffic, with CPU offload there are PCIe copies of the gradient and parameter shards in and out of host memory, and with NVMe offload there are asynchronous disk reads and writes staged through the async I/O operator. The reshardable checkpoint, if the interval fires, is written by each rank as its own shard, exactly the way the model is partitioned. That closes the loop of one step, from a plain engine(batch) call down to collectives, PCIe copies, and possibly the NVMe, and back up to an updated model with no rank ever having held it whole.

Part VI: Internals deep dives

Deep dive: stages 1 and 2 versus stage 3

The three ZeRO stages are not three variations on one algorithm. They split into two genuinely different implementations, which is why the code base has a combined stage-1-and-2 path and a separate stage-3 path. Stages 1 and 2 leave the parameters replicated. Every rank still holds a full copy of the model weights, so the forward and backward passes are ordinary data-parallel passes with no parameter communication at all. The only cleverness is at the boundaries: gradients are reduce-scattered so each rank keeps one shard (stage 2), and the optimizer state is partitioned so each rank updates only its shard (stages 1 and 2), after which an all-gather republishes the updated parameters. Because parameters are never gathered mid-pass, stages 1 and 2 add zero communication over plain data parallelism and are the safe default whenever the replicated parameters and 2 bytes of gradient still fit. Stage 2 is often the sweet spot. It cuts memory to 2Ψ + 14Ψ/N with no communication penalty.

Stage 3 is the qualitative jump, because it partitions the parameters themselves and therefore has to gather them per layer during both forward and backward, as traced above. That is what costs the extra 0.5x communication, and it is why stage 3 lives in its own file with its own hook machinery, prefetcher, and live- parameter accounting. The practical decision rule is simple. Use the lowest stage whose memory fits, because each stage upward buys memory with either communication or complexity, and stage 2 buys a great deal of memory for free. Stage 3 is for when even the replicated parameters will not fit, which is exactly the regime of the largest models.

Deep dive: ZeRO-Offload and ZeRO-Infinity

ZeRO-Offload extends the partitioning idea down one level of the memory hierarchy. Built originally on ZeRO-2, it moves the optimizer state and the fp32 gradients off the GPU into CPU DRAM and, critically, moves the optimizer computation there too. The insight is a workload-placement argument. The forward and backward passes are compute-heavy and belong on the GPU, but the Adam update is memory-bound and cheap in FLOPs, so running it on the CPU where the state already lives avoids shuttling 12 bytes per parameter back and forth. The reason this is not ruinously slow is DeepSpeedCPUAdam, a fused, AVX-vectorized CPU implementation of Adam that is fast enough to overlap with the next iteration's forward. With ZeRO-Offload a single GPU with a generous host can train models in the low tens of billions of parameters that would otherwise need a small cluster.

ZeRO-Infinity goes further and is built on ZeRO-3. It adds NVMe as a third tier below CPU DRAM, so parameters, gradients, and optimizer state can all spill to disk, and it adds the machinery that makes that tolerable: an asynchronous offload engine that prefetches partitions up the hierarchy before they are needed, memory-centric tiling that breaks individual layers too large for GPU memory into pieces processed in sequence, and a bandwidth-aware partitioning that keeps the collectives efficient. The point of ZeRO-Infinity is that the trainable model size stops being bounded by GPU memory or even by aggregate GPU memory, and becomes bounded by the total capacity of GPU plus CPU plus NVMe across the cluster, which is where the trillion-parameter headline comes from. The cost is real. NVMe bandwidth is a thousand times below HBM, so ZeRO-Infinity is the tool of last resort when nothing else fits, and its job is to make an impossible run merely slow rather than impossible.

Deep dive: the engine and the config system

The DeepSpeedEngine in deepspeed/runtime/engine.py is the object your training loop actually talks to, and it is worth understanding as the coordinator that hides everything else. It subclasses nn.Module, so calling it runs the forward pass, but it also owns the optimizer, the LR scheduler, the loss scaler for fp16, the gradient accumulator, and the checkpoint logic, and it dispatches to whichever ZeRO stage, offload target, and pipeline configuration the config selected. The config itself is the second half of the design, and it is deliberately declarative. A single JSON object with keys like zero_optimization, fp16, bf16, optimizer, scheduler, gradient_clipping, and activation_checkpointing describes the entire training regime, and it is validated against typed config classes in deepspeed/runtime/config.py at init. The payoff of config-as-data is that the same model code trains under ZeRO-1, ZeRO-2, ZeRO-3, or ZeRO-Infinity by editing a JSON file, with no change to the loop, which is exactly the composability that lets a researcher sweep memory strategies the way they sweep learning rates.

The engine also owns DeepSpeed's activation checkpointing, exposed through deepspeed.checkpointing, which goes beyond the stock PyTorch version by optionally partitioning the saved activations across model-parallel ranks, offloading them to CPU, and packing them into contiguous memory to fight fragmentation. These are the same memory-versus-recompute trade-offs discussed in the FlashAttention write-up, applied at the granularity of whole layers rather than the attention kernel.

Deep dive: pipeline parallelism

DeepSpeed's pipeline parallelism is a separate axis from ZeRO, and the two compose. You express the model as a sequence of stages by wrapping it in a PipelineModule, which takes either an nn.Sequential or a list of LayerSpec descriptors and splits the layers into contiguous stages placed on the pipeline-parallel ranks:

from deepspeed.pipe import PipelineModule, LayerSpec

layers = [
    LayerSpec(EmbeddingLayer, vocab, dim),
    *[LayerSpec(TransformerBlock, dim, heads) for _ in range(n_layers)],
    LayerSpec(LMHead, dim, vocab),
]
net = PipelineModule(
    layers=layers,
    num_stages=4,
    partition_method="parameters",   # balance stages by parameter count
    loss_fn=cross_entropy,
)
engine, _, _, _ = deepspeed.initialize(model=net, config=ds_config)

The LayerSpec indirection matters. It records how to build each layer without building it, so a stage constructs only the layers assigned to its rank and the full model is never instantiated anywhere, the same lazy-construction trick as zero.Init. TiedLayerSpec handles weight tying across stages, such as sharing the embedding and the output projection. The pipeline engine runs micro-batches through the stages on a 1F1B-style schedule (one forward, one backward, interleaved) so that the pipeline bubble shrinks as the number of micro-batches grows, and it moves activations and their gradients between adjacent stages with point-to-point sends and receives rather than collectives. Only the last stage computes the loss, which is the honest cost of pipelining: it breaks the illusion that every rank runs the same program. Combined with ZeRO-1 for the optimizer state and Megatron for tensor parallelism, this pipeline forms the third dimension of the 3D parallelism used to train the largest public models.

Deep dive: the fused kernels and op_builder

A surprising amount of DeepSpeed is C++ and CUDA, and the way it manages that native code is itself a design worth studying. The kernel sources live under csrc/, their Python bindings under deepspeed/ops/, and the build logic in a top-level op_builder/ package. Each operator has a builder class (for example a CPU-Adam builder, a fused-Adam builder, a transformer builder, an inference builder) that knows its sources, its compile flags, and how to check whether the host toolchain can build it. Operators can be precompiled at install time by setting DS_BUILD_OPS=1, or, by default, they are just-in-time compiled the first time your program imports them, which is why the first run after install pauses to build. The headline kernels are FusedAdam and DeepSpeedCPUAdam for the optimizer, a fused transformer layer used to set BERT training-speed records, a block-sparse attention module, and the inference kernels discussed next. The lesson of op_builder is that a Python library can ship serious hand-written CUDA without forcing every user to precompile it, by treating compilation as a lazy, per-operator, capability-checked step, which is precisely what ds_report exposes.

Deep dive: DeepSpeed-Inference

Training and inference are different problems, and DeepSpeed grew a separate inference path, entered through deepspeed.init_inference rather than deepspeed.initialize. It does three things. It shards the model across GPUs with tensor parallelism for latency, using an injection policy that recognizes known architectures and slices their attention and MLP weights automatically. It replaces the model's transformer modules with fused inference kernels through kernel injection, which is what replace_with_kernel_inject turns on. And it offers INT8 quantization paths for throughput. A minimal call looks like this, with a note that the argument names have changed across versions:

import deepspeed, torch

engine = deepspeed.init_inference(
    model,
    tensor_parallel={"tp_size": 2},   # older API used mp_size=2
    dtype=torch.float16,
    replace_with_kernel_inject=True,
)
logits = engine(input_ids)

Related but distinct is ZeRO-Inference, which applies the offload idea to inference: it streams a huge model's weights layer by layer from CPU or NVMe through a single GPU, trading latency for the ability to run a model far larger than GPU memory when throughput matters more than per-token speed. It is worth being honest about the landscape here. For most interactive serving, the paged-KV- cache engines like vLLM have become the default, and DeepSpeed-Inference is most compelling when you are already in the DeepSpeed ecosystem or when ZeRO-Inference's offload is the only way to fit the model at all. The later DeepSpeed-FastGen work brought continuous-batching ideas into the same family, but the general trajectory of the field has favored the dedicated serving engines.

Part VII: Reading the repository

The repository is large, but the load-bearing parts are a small subset. Read in this order. Paths reflect the layout in mid 2026 and the runtime/zero internals have been reorganized before, so treat the deeper filenames as roles.

Stage 0, orientation. Read the top-level README.md and run ds_report against your box. Then open deepspeed/__init__.py to see that the public surface is essentially initialize, init_inference, and the zero submodule. Question: what four things does initialize return, and which of them did you provide versus which did the engine build?

Stage 1, the engine. Read deepspeed/runtime/engine.py, following one path from forward to backward to step. Questions: where does gradient accumulation decide to skip the reduction, where does loss scaling live, and how does the engine choose which ZeRO stage implementation to construct?

Stage 2, ZeRO stages 1 and 2. Read the combined stage-1-and-2 optimizer under deepspeed/runtime/zero/. It is the gentler of the two ZeRO paths because parameters stay replicated. Questions: where is the gradient reduce-scatter, and where does the partitioned optimizer state get updated and then all-gathered back?

Stage 3, ZeRO stage 3. Read the stage-3 code and the parameter-partitioning module. This is the hardest and most rewarding reading in the repo. Questions: which hooks all-gather a module's parameters, where are they freed, what does stage3_param_persistence_threshold control, and how does the prefetcher decide what to gather next?

Stage 4, offload and the CPU optimizer. Read the parameter and optimizer offload coordinators and the DeepSpeedCPUAdam binding in deepspeed/ops/adam/, then glance at its C++ source under csrc/. Questions: what gets copied over PCIe and when, and why is the CPU Adam SIMD implementation the thing that makes offload viable?

Stage 5, pipeline and kernels. Read deepspeed/runtime/pipe/ for the PipelineModule, the engine, and the schedule, then skim op_builder/ to understand how a kernel goes from csrc/ source to a callable op. Questions: how does LayerSpec avoid building the whole model, and how does the schedule decide the order of forwards and backwards?

Where not to start: the inference and module_inject trees, the MoE code, and the many accelerator backends are all worthwhile later but will overwhelm a first read. Meet the ZeRO stages and the engine first, because everything else in the library orbits them.

Part VIII: Hands-on labs

Labs 1 through 3 need a single GPU. Labs 4 and 5 want two or more. Log formats and a few config keys drift with the project's fast pace, so read the current docs alongside these.

Lab 1: ds_report and the three method calls. Concept: the engine replaces backward, step, and zero_grad.

pip install deepspeed
ds_report
deepspeed --num_gpus=1 train.py --deepspeed --deepspeed_config ds_config.json

Start from any small PyTorch training script (the DeepSpeed examples repo has a CIFAR one). Convert it: replace the model, optimizer, and loop with deepspeed.initialize and the engine(batch), engine.backward(loss), engine.step() trio. Confirm it trains, then read ds_report and note which ops were JIT-compiled on the first step.

Lab 2: watch memory fall across the stages. Concept: the Part III table, empirically.

# flip only zero_optimization.stage between runs: 0, 1, 2, 3
# after a few steps, print peak allocated memory
import torch
print("peak GB:", torch.cuda.max_memory_allocated() / 1e9)

Run the same model at stage 0 (plain data parallel), then 1, 2, and 3, logging peak memory each time. Watch it drop in roughly the proportions the memory table predicts. This is the fastest way to make the ZeRO math stop being abstract.

Lab 3: offload a model that does not fit. Concept: ZeRO-Offload and the CPU optimizer.

ds_config["zero_optimization"] = {
    "stage": 2,
    "offload_optimizer": {"device": "cpu", "pin_memory": True},
}

Take a model that just barely OOMs at stage 2 on your GPU, add the optimizer offload above, and watch it fit. Compare step time with and without offload and observe how small the penalty is when the optimizer step overlaps with the next forward. Then try stage: 3 with offload_param to cpu and note the larger model you can now hold.

Lab 4: construct a large model under zero.Init. Concept: partitioned construction.

with deepspeed.zero.Init(config_dict_or_path=ds_config):
    model = BigTransformer(layers=48, dim=8192)   # never fully on one GPU
engine, *_ = deepspeed.initialize(model=model, config=ds_config)

Build a model whose full parameter set exceeds one GPU's memory, first without the context (observe the OOM at construction), then inside zero.Init (observe that it succeeds because parameters are sharded as they are created). This is the trick that makes stage 3 usable for genuinely huge models.

Lab 5: pipeline a Sequential model. Concept: stages, LayerSpec, and micro-batches.

from deepspeed.pipe import PipelineModule
net = PipelineModule(layers=my_sequential, num_stages=2, loss_fn=loss_fn)
engine, *_ = deepspeed.initialize(model=net, config=ds_config)

Wrap a plain nn.Sequential in a PipelineModule across two GPUs. Vary gradient_accumulation_steps, which sets the micro-batch count, and watch throughput improve as the pipeline bubble shrinks with more micro-batches. Note that only the last stage reports the loss.

Part IX: Questions and model answers

Understanding checks. Answer aloud before reading.

1. What is DeepSpeed, in one sentence?

A training engine that wraps a plain PyTorch model and removes the memory that data parallelism redundantly replicates on every GPU, by partitioning the optimizer state, gradients, and parameters across the data-parallel group and gathering each piece back only when it is needed.

2. Where does the 16Ψ come from, and which part does ZeRO attack first?

In mixed-precision Adam training, per parameter, you store 2 bytes of fp16 weights, 2 bytes of fp16 gradient, and 12 bytes of fp32 optimizer state (a master copy plus momentum plus variance). That is 16 bytes, so a model with Ψ parameters needs 16Ψ bytes of state. The 12 bytes of optimizer state are the largest and most redundant piece, so ZeRO-1 partitions those first.

3. What exactly does each ZeRO stage partition?

Stage 1 partitions the optimizer state only. Stage 2 also partitions the gradients. Stage 3 also partitions the parameters. Per-GPU memory goes from 16Ψ at baseline, to 4Ψ + 12Ψ/N, to 2Ψ + 14Ψ/N, to 16Ψ/N.

4. Why do stages 1 and 2 add no communication over plain data parallelism, but stage 3 adds 50 percent?

Stages 1 and 2 keep parameters replicated, so the forward and backward passes need no parameter communication, and the gradient reduce-scatter plus parameter all-gather move the same 2Ψ bytes as a plain all-reduce. Stage 3 partitions parameters, so it must all-gather them once in forward and again in backward, plus reduce-scatter gradients, which is 3Ψ, or 1.5 times the baseline.

5. How is a partitioned stage-3 model runnable at all?

Forward and backward hooks on every submodule all-gather that module's full parameters just before it computes and free them right after, with the next module's all-gather prefetched to overlap with the current module's compute. Only a couple of layers' worth of parameters are materialized at any instant.

6. What does deepspeed.zero.Init do and why is it necessary?

It is a context manager that partitions parameters as they are constructed, so a model too large for one GPU is never fully instantiated on any rank. Without it, stage 3 builds the full model first and then partitions, which reintroduces the very out-of-memory you are trying to avoid.

7. Why does the training loop call backward and step on the engine instead of on the tensors and optimizer?

Because the engine has to interleave gradient reduction, loss scaling, gradient accumulation, and clipping around those operations, and under offload the step may run on the CPU. Handing those responsibilities to the engine is what lets the same loop run unchanged across every ZeRO stage and offload target.

8. What is ZeRO-Offload, and why is DeepSpeedCPUAdam central to it?

ZeRO-Offload moves the optimizer state and the optimizer computation off the GPU into CPU memory, keeping the compute-heavy forward and backward on the GPU. It works because the Adam update is memory-bound and cheap in FLOPs, but only if the CPU can do it fast enough, which is why DeepSpeed ships a SIMD-vectorized C++ Adam so the CPU step overlaps with the next iteration rather than stalling it.

9. What does ZeRO-Infinity add, and what does it cost?

It extends stage 3 to spill parameters, gradients, and optimizer state to NVMe below CPU memory, with an asynchronous prefetching offload engine and tiling for oversized layers, so trainable model size is bounded by total GPU plus CPU plus NVMe capacity rather than GPU memory. The cost is NVMe bandwidth, orders of magnitude below HBM, so it is the tool of last resort that makes an impossible run merely slow.

10. How does DeepSpeed relate to PyTorch FSDP?

ZeRO-3 and FSDP are the same core idea, partition the parameters and gather them per layer on the fly. FSDP is that idea absorbed into PyTorch itself, and FSDP2 is what native stacks like torchtitan build on. DeepSpeed is where the idea, and the offload extensions, were pioneered.

11. How does pipeline parallelism differ from ZeRO, and do they compose?

Pipeline parallelism splits the layer stack across ranks and passes activations stage to stage with point-to-point sends, whereas ZeRO shards the state of a data-parallel replica. They are orthogonal axes and compose, and combined with Megatron tensor parallelism they form the 3D parallelism used for the largest models.

12. Why is ds_report the first thing to run?

Most of DeepSpeed's speed is in custom C++ and CUDA operators that are compiled for your machine, either at install or just in time. ds_report tells you which ops your toolchain can build and which are already compiled, so you learn that, for example, CPU Adam is unavailable before a long offload run fails instead of after.

13. When would you reach for Megatron-LM or vLLM instead of DeepSpeed?

Megatron-LM when you need maximum model-FLOP utilization on large NVIDIA clusters and accept tensor parallelism that rewrites your model, though in practice the two are combined as Megatron-DeepSpeed. vLLM and similar engines for interactive inference, where paged KV caches and continuous batching now dominate. DeepSpeed wins when memory is the binding training constraint and you want to keep a mostly plain model.

14. A ZeRO-Infinity run technically works but is ten times too slow. What is the first suspect?

The NVMe offload target cannot keep the GPUs fed, either because the path is on a slow or network disk rather than fast local NVMe, or because the working set overflows the CPU tier above it so the prefetcher cannot hide the disk latency. The symptom is step time, not an error, and the profiler is how you confirm it.

Part X: Design lessons

Redundancy is the enemy of scale. The single observation behind the whole library is that data parallelism duplicates identical optimizer state on every GPU. Naming that redundancy and partitioning it away is the entire ZeRO idea. Wherever a system replicates the same large state across nodes for convenience, there is usually a ZeRO-shaped win waiting.

Treat memory as a hierarchy, not a single pool. ZeRO-Offload and ZeRO-Infinity work because GPU HBM, CPU DRAM, and NVMe are tiers with different bandwidth and capacity, and state can be placed on the slowest tier that still keeps the compute fed. This is the cache hierarchy, virtual memory, and tiered storage insight applied to model training, and the discipline is the same: push data down only as far as capacity forces you.

Put the compute where the data lives. Running the memory-bound Adam step on the CPU because the optimizer state already sits in CPU memory avoids moving 12 bytes per parameter across PCIe twice. Deciding placement by where the operands live, rather than defaulting everything to the accelerator, is a general heterogeneous-computing lesson.

Configuration as data decouples strategy from code. A single JSON object selects the ZeRO stage, the offload targets, the precision, and the schedule, so the same model trains under wildly different memory strategies with no code change. Sweeping memory strategy becomes as cheap as sweeping a hyperparameter, which is the same payoff that declarative config buys everywhere.

Ship serious native code without forcing precompilation. The op_builder and ds_report design lets a Python library carry hand-written CUDA that is compiled lazily, per operator, only where the toolchain supports it. Making the capability check a first-class, inspectable command is why users can diagnose kernel problems before a run rather than during one.

Part XI: Papers and further reading

Nearly every idea in this chapter has a primary source, and the ZeRO line of papers in particular reads almost like the code. Where this site covers the same idea in depth, the companion link points there.

1. Rajbhandari et al., ZeRO, Memory Optimizations Toward Training Trillion Parameter Models, 2019. The paper this whole chapter derives, the three partitioning stages, the 16Ψ accounting, and the 7.5B-on-64-GPUs worked example of Part III. The who-shards-what theory is covered in the parallel computing class on this site.
2. Ren et al., ZeRO-Offload, Democratizing Billion-Scale Model Training, 2021. The workload-placement argument for moving optimizer state and the Adam update to the CPU, including the DeepSpeedCPUAdam design.
3. Rajbhandari et al., ZeRO-Infinity, Breaking the GPU Memory Wall for Extreme Scale Deep Learning, 2021. Adds the NVMe tier, the asynchronous offload engine, and memory-centric tiling, so trainable model size is bounded by total capacity rather than GPU memory.
4. Micikevicius et al., Mixed Precision Training, 2017. The source of the fp16 weights, the fp32 master copy, and the loss scaling that give the memory table its 2 and 4 byte rows. Worked through in the mixed precision note on this site.
5. Kingma and Ba, Adam, A Method for Stochastic Optimization, 2014. The momentum and variance state that make the optimizer multiplier K equal 12 and hand ZeRO its biggest target. Derived in the Adam note.
6. Shoeybi et al., Megatron-LM, Training Multi-Billion Parameter Language Models Using Model Parallelism, 2019. The tensor-parallel partner in 3D parallelism, covered in the Megatron-LM walkthrough.
7. Harlap et al., PipeDream, Fast and Efficient Pipeline Parallel DNN Training, 2018. The origin of the 1F1B schedule that DeepSpeed's pipeline engine runs to shrink the bubble.
8. Smith et al., Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, 2022. The flagship 3D-parallelism run that combined Megatron tensor parallelism with DeepSpeed ZeRO and pipelining.
9. BigScience Workshop, BLOOM, A 176B-Parameter Open-Access Multilingual Language Model, 2022. The other large public model trained with Megatron-DeepSpeed, and a detailed account of running that stack in practice.
10. Zhao et al., PyTorch FSDP, Experiences on Scaling Fully Sharded Data Parallel, 2023. ZeRO-3 absorbed into PyTorch itself, whose FSDP2 descendant is what the torchtitan walkthrough builds on.
11. Aminabadi et al., DeepSpeed Inference, Enabling Efficient Inference of Transformer Models at Unprecedented Scale, 2022. The kernel-injection and tensor-parallel serving path of the inference deep dive, plus the ZeRO-Inference offload idea.
12. Holmes et al., DeepSpeed-FastGen, High-throughput Text Generation for LLMs via MII and DeepSpeed-Inference, 2024. The later serving work in the same family, best read against the paged-KV engines in the vLLM walkthrough.

Part XII: Final takeaway

If the single-device transformer that all of this partitions is the gap, the language models from scratch class builds one end to end, and the parallel computing class covers the who-shards-what theory in general. When you want the PyTorch-native descendant of ZeRO-3, read the torchtitan chapter, and you will recognize the same all-gather-per-layer discipline wearing different names. Then come back and read deepspeed/runtime/engine.py once more, and the three-line training loop will read like plain PyTorch, which is the whole point.