Training Models Bigger Than Your GPU: A Complete Guide to DeepSpeed

How Microsoft’s Optimization Library Turned Extreme Scale Deep Learning Into a Practical, Everyday Tool

There is a specific, frustrating moment familiar to anyone who has tried training a genuinely large model: the training script runs, the data loads correctly, and then, seconds into the first forward pass, an accelerator simply runs out of memory before the model has even finished materializing on the device. Modern language models and other large architectures routinely carry parameter counts, optimizer states, and gradient buffers that exceed what a single GPU, or even several GPUs working together, can hold using conventional training approaches. DeepSpeed exists specifically to solve that problem, giving engineers a deep learning optimization library purpose built for scaling distributed PyTorch training and inference well past what naive data parallelism can support on its own.

Rather than treating memory limits as a hardware problem to be solved purely by buying bigger accelerators, DeepSpeed treats it as a systems engineering problem, reducing redundant, replicated training state, offloading model state to slower but far more abundant memory tiers when needed, and coordinating several distinct forms of parallelism together so that a workload can actually fit and run efficiently on the hardware genuinely available, rather than the hardware a team might wish it had.

The Core Problem DeepSpeed Solves

Ordinary data parallel training, the default approach most practitioners reach for first, replicates a model’s full parameters, gradients, and optimizer state onto every single device participating in training. That replication is simple and works well for models that comfortably fit on one device, but it becomes the actual bottleneck once a model’s own memory footprint grows large enough that even a single full copy of its parameters, gradients, and optimizer state together cannot fit on one accelerator at all. Optimizer state in particular tends to be a quietly enormous contributor to this problem; a widely used optimizer such as Adam maintains two additional floating point values per parameter beyond the parameter itself, meaning optimizer state alone can easily dwarf the memory footprint of the raw model weights.

DeepSpeed’s foundational innovation, an approach called ZeRO, or Zero Redundancy Optimizer, directly targets this specific inefficiency. Rather than replicating optimizer states, gradients, and parameters identically across every data parallel process, ZeRO partitions each of these components across the available processes instead, with each device holding only its own assigned shard rather than a full, redundant copy of everything. The result is a dramatic reduction in per device memory consumption without abandoning the fundamentally simple, well understood structure of data parallel training, letting models that would otherwise be entirely impossible to train on a given cluster fit comfortably within it instead.

Understanding ZeRO’s Progressive Stages

ZeRO is generally described in terms of progressive stages, each partitioning a further category of training state beyond what the previous stage already handled, and each offering a further step up in memory savings at the cost of somewhat increased communication overhead.

The first stage partitions optimizer states across data parallel processes, already delivering a meaningful reduction in memory consumption for the specific component that tends to be the largest contributor to overall training memory pressure in the first place.

The second stage extends that same partitioning to gradients as well, meaning neither optimizer state nor gradients are fully replicated on every device, compounding the memory savings considerably further while still keeping full, unpartitioned parameters locally available on every device for the actual forward and backward computation itself.

The third stage goes a step further still, partitioning the model’s parameters themselves across devices in addition to optimizer states and gradients. This stage offers the largest reduction in per device memory footprint of the three, since essentially none of a model’s core training state remains fully replicated anywhere, though it does require devices to communicate and reconstruct the relevant parameter shards on demand during both the forward and backward passes, introducing additional communication traffic that needs to be managed carefully to avoid becoming a new bottleneck in its own right.

Building further on top of this third stage, ZeRO-Infinity extends the same underlying partitioning philosophy beyond GPU memory entirely, allowing the full model state to be offloaded not just across multiple GPUs but out to CPU memory or even NVMe storage when GPU memory alone still is not sufficient. This effectively allows training models whose total memory footprint significantly exceeds the combined memory of every GPU actually available in a given cluster, trading some additional data movement overhead for the ability to train models that would otherwise simply be impossible to fit using GPU memory alone under any partitioning scheme.

A Basic Configuration Example

DeepSpeed is configured primarily through a JSON configuration file passed alongside a training script, keeping the actual training code itself largely unchanged from a standard PyTorch training loop. A representative configuration enabling ZeRO stage two, along with mixed precision training, looks roughly like this:

{
  "train_batch_size": 32,
  "gradient_accumulation_steps": 4,
  "fp16": {
    "enabled": true
  },
  "zero_optimization": {
    "stage": 2,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    },
    "allgather_partitions": true,
    "allgather_bucket_size": 2e8,
    "overlap_comm": true,
    "reduce_scatter": true,
    "reduce_bucket_size": 2e8,
    "contiguous_gradients": true
  },
  "optimizer": {
    "type": "AdamW",
    "params": {
      "lr": 3e-5,
      "betas": [0.9, 0.999],
      "eps": 1e-8,
      "weight_decay": 0.01
    }
  }
}

With a configuration file like this in place, initializing DeepSpeed from within an existing training script generally requires only a modest change to how the model and optimizer are wrapped, rather than a fundamental restructuring of the training loop itself:

import deepspeed
model_engine, optimizer, _, _ = deepspeed.initialize(
    model=model,
    model_parameters=model.parameters(),
    config="ds_config.json",
)
for batch in dataloader:
    outputs = model_engine(batch["input_ids"])
    loss = loss_fn(outputs, batch["labels"])
    model_engine.backward(loss)
    model_engine.step()

The general shape of this loop remains recognizably similar to an ordinary PyTorch training loop, with DeepSpeed’s returned model engine object handling the underlying partitioning, gradient synchronization, and optimizer step logic transparently behind that familiar interface.

Coordinating Multiple Forms of Parallelism

Memory partitioning through ZeRO is only one piece of DeepSpeed’s broader approach to scaling training. The library also supports coordinating several distinct forms of parallelism together, since different bottlenecks call for genuinely different solutions, and the largest, most demanding training runs typically need more than one form of parallelism applied simultaneously.

Data parallelism, the familiar baseline approach, splits a training batch across multiple devices, each holding its own copy of the model, or in DeepSpeed’s case, its own shard of the model’s state under ZeRO, and processing a different portion of a given batch before synchronizing results.

Model parallelism instead splits an individual model’s layers or parameters directly across multiple devices, which becomes necessary once a single layer or a small group of layers is itself too large to fit comfortably on one device regardless of how batch data is distributed.

Pipeline parallelism divides a model into sequential stages distributed across different devices, with different devices processing different micro batches at different pipeline stages concurrently, aiming to keep every device productively busy rather than sitting idle while waiting for an earlier stage in a strictly sequential pipeline to finish.

Ulysses sequence parallelism addresses a distinct, increasingly important bottleneck specific to training on extremely long input sequences, partitioning computation specifically along the sequence dimension itself, which becomes essential once a single sequence’s attention computation alone becomes too memory intensive to fit on one device, a scenario that has become considerably more common as training on multi million token sequences has moved from a research curiosity into genuinely practical, actively supported territory.

Because these different forms of parallelism address genuinely different bottlenecks, the largest scale training runs typically combine several of them together simultaneously, commonly referred to as 3D parallelism when data, model, and pipeline parallelism are combined, extended further still when sequence parallelism is layered in as a fourth dimension for workloads specifically constrained by extremely long sequence lengths.

Inference Support

DeepSpeed’s scope extends beyond training into inference as well, offering model parallelism support alongside inference specific kernels optimized for compatible transformer architectures. Serving an extremely large model efficiently carries its own distinct set of challenges compared with training it, generally prioritizing low latency and high throughput for live serving traffic over the raw training speed that dominates training focused optimization work. DeepSpeed’s inference tooling is built specifically around those serving oriented priorities, allowing models trained at genuinely large scale to also actually be deployed and served efficiently afterward, rather than leaving that equally important second half of a large model’s lifecycle unaddressed.

Recent Developments

DeepSpeed continues to see active, ongoing development well beyond its original ZeRO based foundation. Recent work has introduced support for the Muon optimizer as an alternative to more conventional optimizers such as Adam or AdamW for certain training scenarios. A capability referred to as System DMA support for ZeRO-3 specifically targets AMD GPU hardware, offloading collective communication operations off of compute units entirely in order to improve overlap between computation and communication during training, squeezing additional efficiency out of hardware that might otherwise sit partially idle waiting on communication to complete.

SuperOffload, presented at ASPLOS 2026 and receiving an honorable mention for that conference’s best paper award, targets large scale training specifically on newer superchip class hardware, extending the general offloading philosophy already established by ZeRO-Infinity to newer categories of accelerator hardware. ZenFlow, described as a stall free offloading engine, targets a related but distinct problem, specifically the stalls that can otherwise occur during CPU or NVMe offloaded training when data movement fails to overlap cleanly with ongoing computation, aiming to keep GPUs genuinely busy throughout a training step rather than periodically waiting on slower offload tiers to catch up.

Arctic Long Sequence Training addresses the increasingly important challenge of training efficiently on sequences spanning multiple millions of tokens, an area of genuinely rapid growth in demand as applications involving very long documents, codebases, or extended conversational context become more common. DeepNVMe focuses specifically on making NVMe based input and output scaling more affordable and practical for deep learning applications generally, recognizing that storage bandwidth itself can become a genuine bottleneck at sufficiently large scale even when compute and network resources are otherwise well provisioned.

The project also hosts regular, genuinely open office hours on the last Tuesday of each month, held over Zoom and open to anyone interested in joining to discuss ongoing development plans, upcoming features, or simply to ask questions directly of the team actively building the library.

A Track Record at Genuine Scale

DeepSpeed’s practical impact is perhaps best illustrated by the specific, publicly known large scale models it has been used to train. These include Megatron-Turing NLG at five hundred thirty billion parameters, Jurassic-1 at one hundred seventy eight billion parameters, BLOOM at one hundred seventy six billion parameters, GLM at one hundred thirty billion parameters, xTrimoPGLM at one hundred billion parameters, YaLM at one hundred billion parameters, GPT-NeoX at twenty billion parameters, AlexaTM at twenty billion parameters, Turing NLG at seventeen billion parameters, and METRO-LM at five point four billion parameters.

This list spans work from several genuinely distinct organizations across both industry and research settings, reflecting DeepSpeed’s role as shared, broadly adopted infrastructure underlying a meaningful share of the field’s most prominent large scale training efforts, rather than a tool used narrowly within a single organization for a single specific project.

DeepSpeed itself formed a central part of Microsoft’s broader AI at Scale initiative, aimed specifically at enabling next generation AI capability at genuinely large scale, and its influence extends well beyond models trained directly by Microsoft itself, given how widely it has been adopted as underlying infrastructure across the broader large scale training community.

Working Through Existing Frameworks

Rather than requiring every team to integrate directly against DeepSpeed’s own lower level API, the library is integrated with several widely used open source deep learning frameworks, letting teams already standardized on one of these tools adopt DeepSpeed’s optimizations with comparatively modest changes to existing code.

Direct integrations are documented for Hugging Face’s Transformers library, Hugging Face’s Accelerate library, PyTorch Lightning, MosaicML, Determined, and MMEngine. For a team already building on top of any of these frameworks, this generally means DeepSpeed’s memory and scaling optimizations can be adopted incrementally, often through configuration changes and a small number of code adjustments, rather than requiring an entirely separate training pipeline built from scratch specifically around DeepSpeed.

Installing DeepSpeed

The most direct way to get started with DeepSpeed is through pip, installing the latest released version, which is not tied to any specific PyTorch or CUDA version combination:

pip install deepspeed

DeepSpeed includes a range of C++ and CUDA extensions, generally referred to within the project simply as its operators or ops, which by default are built just in time using PyTorch’s own JIT C++ extension loading mechanism, relying on the ninja build tool to compile and dynamically link them at the moment they are actually needed during a training run.

Before installing DeepSpeed itself, PyTorch needs to already be installed, with a version of 2.0 or later, ideally the current stable release, recommended for full feature support. A working CUDA or ROCm compiler, such as nvcc for NVIDIA hardware or hipcc for AMD hardware, is also required to actually compile these C++, CUDA, or HIP extensions. DeepSpeed is developed and tested primarily against NVIDIA’s Pascal, Volta, Ampere, and Hopper GPU architectures, along with AMD’s MI100 and MI200 accelerators, though hardware falling outside this specific tested list is not necessarily unsupported, simply less thoroughly validated by the core development team directly.

Beyond this core, primarily tested hardware, DeepSpeed has also received contributed support for a range of additional accelerator platforms from outside contributors, including Huawei’s Ascend NPU, Intel’s Gaudi 2 AI accelerator, Intel Xeon processors used specifically as a CPU based training and inference target, Intel’s Data Center GPU Max series, and Tecorigin’s Scalable Data Analytics Accelerator. Support for each of these platforms varies somewhat in its current validation status, with some already validated both by the contributing organization and by the core upstream DeepSpeed team, and others currently validated only by the contributing organization itself pending further upstream confirmation.

After installation, confirming exactly which extensions and operators are actually compatible with a given machine is handled through a dedicated diagnostic command:

ds_report

For anyone who would rather pre install specific extensions ahead of time rather than relying on just in time compilation during an actual training run, or who wants to install pre compiled operators directly through PyPI instead, the project’s advanced installation documentation covers those alternative paths in further detail.

Installing on Windows

Despite deep learning infrastructure of this kind traditionally being strongly Linux centric, a genuinely substantial share of DeepSpeed’s feature set is supported directly on Windows as well, covering both training and inference workloads. Notable exceptions currently not supported on Windows include asynchronous I/O and GPU Direct Storage, the latter of which does not support Windows at a more fundamental platform level regardless of DeepSpeed’s own implementation choices.

Setting up DeepSpeed on Windows involves installing a compatible PyTorch version, such as PyTorch 2.3 or later paired with CUDA 12.1, installing Visual C++ build tools such as the Visual Studio 2022 C++ x64 and x86 build tools, and launching a command console with administrator permissions, which is specifically required for creating certain symlink folders DeepSpeed needs, while also making sure the relevant MSVC build tools are available on the system path, or alternatively launching the dedicated Visual Studio 2022 Developer Command Prompt directly with administrator permissions instead. From there, running a provided batch script builds the actual installable wheel package:

build_win.bat

The resulting wheel file is placed directly in a dist folder, ready to be installed through pip in the ordinary way.

Where to Find Further Documentation

The full range of DeepSpeed documentation, tutorials, and technical blog posts is maintained centrally on the project’s own website. This includes a getting started guide covering first steps for newcomers, detailed documentation of DeepSpeed’s JSON configuration format, generated API reference documentation, a broader collection of tutorials covering specific features and use cases in depth, and ongoing blog posts covering both technical deep dives and project news as it happens.

Supporting the Project’s Infrastructure

Being a genuinely open source project, DeepSpeed’s own continuous integration testing depends directly on external hardware resources being made available to the project. The project’s GPU based continuous integration runs are currently kindly supported by Modal, an AI infrastructure platform covering inference, fine tuning, and batch job workloads among other offerings, which directly funds the underlying hardware used for these ongoing CI runs. This kind of infrastructure support matters considerably for a project maintaining broad, ongoing hardware compatibility testing across genuinely diverse GPU architectures and vendors, since that kind of testing simply cannot happen without real, dedicated hardware being consistently available to actually run it against.

Conclusion

Training models that exceed the memory limits of conventional hardware has traditionally forced engineers to choose between scaling down their ambitions or investing in significantly more infrastructure. DeepSpeed offers a fundamentally different solution by treating memory constraints as a systems engineering challenge rather than a hardware problem. Through intelligent partitioning, selective offloading, and highly coordinated parallelism, it enables models far larger than would otherwise fit within available GPU memory.

At the heart of this approach is the ZeRO family of optimizations, culminating in ZeRO-Infinity, which extends model execution beyond GPU memory by transparently offloading model states to CPU memory and storage. Combined with comprehensive support for data, tensor, pipeline, and sequence parallelism, and seamless integration with many of today’s leading deep learning frameworks, DeepSpeed has become a foundational component of the infrastructure used to train many of the world’s largest and most capable language models.

For engineers who encounter memory limitations long before reaching the practical limits of their hardware, DeepSpeed provides more than a collection of optimization techniques. It offers a mature, battle-tested, and actively evolving framework that transforms hardware constraints from a fundamental barrier into an engineering problem that can often be overcome. Rather than forcing projects to compromise on model scale, DeepSpeed enables them to make far more effective use of the resources they already have.

The repository is available at: https://github.com/deepspeedai/DeepSpeed