Skip to content

How to Install DeepSpeed

by Tim Hopper · Markdown
uv Scientific Python verified with uv 0.12

PyPI serves only a source distribution for DeepSpeed, so pip install deepspeed compiles instead of downloading a wheel. Prebuilt wheels are on the Astral GPU indexes, one index per CUDA version, and they arrive with the CUDA ops already compiled.

The index removes the build, not the toolkit. DeepSpeed runs $CUDA_HOME/bin/nvcc while importing, so a GPU machine still needs either a toolkit or DS_IGNORE_CUDA_DETECTION=1. See Why Installing GPU Python Packages Is So Complicated for background.

Pick an install path

Check your platform

  • Platform: Linux (x86_64). Windows has partial support through WSL2. No macOS GPU support. The index wheels are manylinux_2_27_x86_64 only, with no aarch64 build, so Grace Hopper and other ARM hosts need a source build.
  • PyTorch: installed before DeepSpeed on every path. The index wheels pin it (torch==2.11.*); the source build imports it in setup.py.
  • Toolkit: needed for a source build, and for import deepspeed on a GPU machine unless DS_IGNORE_CUDA_DETECTION=1 is set.
  • System libraries: libaio-dev for the async I/O op used by ZeRO-Infinity and NVMe offloading (apt install libaio-dev on Debian/Ubuntu). No wheel bundles it, so ds_report flags async_io as unavailable until it is installed.

Install from the Astral GPU index

Each index covers one CUDA version, and each wheel targets one PyTorch version. Both sit in a PEP 440 local version segment rather than a platform tag, so deepspeed resolves like any other dependency:

deepspeed-0.19.2+cu.12.8.torch.2.11-cp312-cp312-manylinux_2_27_x86_64.whl

Route torch and deepspeed to indexes on the same CUDA line. Pick the CUDA version your driver supports (cu126, cu128, cu129, cu130, or cu132), then add both:

uv add torch --index pytorch-cu128=https://download.pytorch.org/whl/cu128
uv add deepspeed --index astral-cu128=https://wheels.astral.sh/simple/cu128/
$ uv add deepspeed --index astral-cu128=https://wheels.astral.sh/simple/cu128/
Resolved 46 packages in 5.31s
Downloading deepspeed (186.9MiB)
Prepared 15 packages in 5.43s
Installed 15 packages in 109ms
 + deepspeed==0.19.2+cu.12.8.torch.2.11
 + einops==0.8.2
 + hjson==3.1.0
 ...

Each wheel declares its PyTorch as a dependency (torch==2.11.*), so the resolver picks the build matching the project’s torch. Routing torch to a CUDA-specific index pins the other half of the pair.

Warning

Routing torch is not optional. The Astral indexes serve no torch, so an unrouted torch comes from PyPI, whose CUDA build changes between releases: PyPI ships PyTorch 2.11 as a CUDA 13 build. It satisfies torch==2.11.*, so pairing it with a cu128 wheel resolves and installs cleanly, then fails against CUDA 12 shared libraries at runtime.

Installing deepspeed as a dependency of something else (Megatron, Axolotl, a training framework) skips this routing, because [tool.uv.sources] applies only to declared dependencies. Declare deepspeed directly too.

Configure the index in pyproject.toml

uv add --index writes both indexes and the source mappings:

[tool.uv.sources]
torch = { index = "pytorch-cu128" }
deepspeed = { index = "astral-cu128" }

[[tool.uv.index]]
name = "pytorch-cu128"
url = "https://download.pytorch.org/whl/cu128"
explicit = true   # add by hand; uv add omits it

[[tool.uv.index]]
name = "astral-cu128"
url = "https://wheels.astral.sh/simple/cu128/"
explicit = true

A named index outranks PyPI for every package it carries, so deepspeed comes from Astral. The cu128 index also carries vllm, flash-attn, mmcv, and others; explicit = true keeps it from silently supplying those. uv’s PyTorch guide recommends the flag.

Important

explicit = true without a matching [tool.uv.sources] entry removes the index from the pool and resolves deepspeed from PyPI instead. uv lock reports no error, and the lock records a PyPI source distribution, so uv sync starts the source build this section exists to avoid. Add both or neither.

Import on a GPU machine without a toolkit

DeepSpeed builds a compatibility table for every op at import. One of those probes runs $CUDA_HOME/bin/nvcc -V whenever torch.cuda.is_available() is true, which stops the import on a GPU machine carrying no toolkit:

$ uv run python -c "import deepspeed"
deepspeed.ops.op_builder.builder.MissingCUDAException: CUDA_HOME does not exist,
unable to compile CUDA op(s)

Set DS_IGNORE_CUDA_DETECTION=1 to skip the probe:

DS_IGNORE_CUDA_DETECTION=1 uv run python -c "import deepspeed; print(deepspeed.__version__)"

The variable only skips a check on whether ops could be JIT-compiled. The wheel’s ops are already compiled, so they load and run either way. An op the wheel omits still needs a toolkit to JIT-compile, so keep the toolkit if you rely on one.

The probe never fires on a CPU-only machine, where torch.cuda.is_available() is false. A container image that installs cleanly during a CPU build can still fail on first import once scheduled onto a GPU.

The nvidia-cuda-nvcc-cu12 wheel does not substitute for a toolkit here: it ships ptxas but no nvcc binary, so pointing CUDA_HOME at it fails the same probe with FileNotFoundError.

Confirm the ops arrived precompiled

ds_report prints which ops are installed against which are merely available for JIT compilation. On an index wheel the interesting column is installed:

$ DS_IGNORE_CUDA_DETECTION=1 uv run ds_report
op name ................ installed .. compatible
fused_adam ............. [YES] ...... [OKAY]
cpu_adam ............... [YES] ...... [OKAY]
transformer_inference .. [YES] ...... [OKAY]
async_io ............... [NO] ....... [NO]
...
deepspeed info ................... 0.19.2+cu.12.8.torch.2.11, b919284, HEAD
torch cuda version ............... 12.8
nvcc version ..................... [FAIL] cannot find CUDA_HOME
deepspeed wheel compiled w. ...... torch 2.11, cuda 12.8

installed [YES] is the outcome that DS_BUILD_OPS=1 buys on a source build after several minutes of compilation. Constructing a FusedAdam returns immediately rather than pausing to JIT-compile, and the nvcc version [FAIL] line is expected on this path.

Check the index coverage

Coverage for deepspeed 0.19.2, by index:

Index PyTorch Python
cu126 2.7-2.12 3.9-3.14
cu128 2.7-2.11 3.9-3.14
cu129 2.8-2.11 3.9-3.14
cu130 2.9-2.12 3.10-3.14
cu132 2.12 3.10-3.14

Coverage is uneven within each row: Python 3.9 appears only against PyTorch 2.7 and 2.8, and 3.14 only against 2.10 and newer. The cu118, cu121, and cu124 indexes carry no deepspeed at all, and the newest release on the index trails the newest on PyPI. wheels.astral.sh is the current manifest.

Build from source

Outside the index’s CUDA and PyTorch matrix, DeepSpeed builds from the PyPI source distribution. This path needs a C++ compiler, the CUDA toolkit with nvcc on PATH, and CUDA_HOME set to the toolkit root, because setup.py checks for the toolkit while generating metadata.

DeepSpeed’s setup.py imports torch at the top level, so PyTorch must be present before installation. The --no-build-isolation flag tells the installer to use the current environment’s torch instead of creating a clean build environment:

uv pip install deepspeed --no-build-isolation

This installs the Python package without compiling any CUDA kernels. Ops are compiled at first use via JIT, which adds a one-time delay (seconds to minutes depending on the op) the first time DeepSpeed runs a training job.

Pre-compile ops at install time

To avoid JIT compilation delays at runtime, set DS_BUILD_OPS=1 to compile all compatible ops during installation:

DS_BUILD_OPS=1 uv pip install deepspeed --no-build-isolation

This requires nvcc on PATH and a working C++ compiler. The build takes several minutes.

To compile only specific ops, use individual environment variables instead:

Variable Op
DS_BUILD_CPU_ADAM CPU Adam optimizer
DS_BUILD_FUSED_ADAM Fused Adam (CUDA)
DS_BUILD_AIO Async I/O for NVMe offload
DS_BUILD_TRANSFORMER_INFERENCE Transformer inference kernels
DS_BUILD_SPARSE_ATTN Sparse attention

Set any of these to 1 to pre-compile that op. For example, to compile only the fused Adam optimizer:

DS_BUILD_FUSED_ADAM=1 pip install deepspeed --no-build-isolation

Add a source build to a uv project

For projects managed with uv using uv add and uv sync, use extra-build-dependencies to inject torch into the isolated build environment. The match-runtime = true option ensures the build uses the same torch version the project resolves at runtime:

[project]
dependencies = ["deepspeed", "torch"]

[tool.uv.extra-build-dependencies]
deepspeed = [{ requirement = "torch", match-runtime = true }]

Then run uv sync as normal. uv handles build isolation and torch injection automatically.

To pre-compile ops during the build, pass environment variables with extra-build-variables:

[tool.uv.extra-build-variables]
deepspeed = { DS_BUILD_OPS = "1" }

Install with conda-forge or pixi

DeepSpeed is available on conda-forge, though the version may lag behind PyPI. The conda-forge build handles CUDA toolkit dependencies through the solver:

pixi add deepspeed

For more on when conda-based tools are the better choice for GPU workloads, see uv vs pixi vs conda for Scientific Python.

Verify the installation

Run one training step to confirm the package imports, its ops load, and the engine drives a backward pass:

check_deepspeed.py
import torch, deepspeed

model = torch.nn.Linear(64, 64)
config = {
    "train_batch_size": 8,
    "optimizer": {"type": "AdamW", "params": {"lr": 1e-3}},
    "zero_optimization": {"stage": 1},
    "bf16": {"enabled": True},
}
engine, _, _, _ = deepspeed.initialize(
    model=model, model_parameters=model.parameters(), config=config
)
loss = engine(torch.randn(8, 64, device=engine.device, dtype=torch.bfloat16)).sum()
engine.backward(loss)
engine.step()
print("training step OK", deepspeed.__version__)
$ DS_IGNORE_CUDA_DETECTION=1 uv run python check_deepspeed.py
training step OK 0.19.2+cu.12.8.torch.2.11

deepspeed.initialize expects the distributed environment variables a launcher normally sets. Run the script under deepspeed check_deepspeed.py, or set RANK=0 LOCAL_RANK=0 WORLD_SIZE=1 MASTER_ADDR=127.0.0.1 MASTER_PORT=29500 for a single-process check.

Troubleshooting

CUDA_HOME does not exist, unable to compile CUDA op(s). On a prebuilt wheel this fires at import, from the op compatibility probe; set DS_IGNORE_CUDA_DETECTION=1. On a source build it fires at metadata generation, before any op is compiled, and needs a real toolkit: export CUDA_HOME=/usr/local/cuda. NVIDIA’s CUDA devel Docker images set this automatically, but slim Python images do not.

FileNotFoundError: .../bin/nvcc at import. CUDA_HOME points at a directory holding no nvcc. The nvidia-cuda-nvcc-cu12 wheel is the common cause, since it ships ptxas alone. Unset CUDA_HOME and use DS_IGNORE_CUDA_DETECTION=1, or point it at a full toolkit.

Undefined symbol errors from a DeepSpeed .so at import. The wheel was built against a different PyTorch than the one installed. Compare deepspeed wheel compiled w. in ds_report against torch.__version__, and route torch to an index on the same CUDA line as the DeepSpeed index.

ModuleNotFoundError: No module named 'torch' during install. The installer is building the PyPI source distribution rather than downloading a wheel, and the isolated build has no torch. Point it at an index, or pass --no-build-isolation with torch already installed.

Learn More

RuntimeError: ninja is not available at runtime. DeepSpeed’s JIT compilation uses ninja as its build backend. Install it with pip install ninja or apt install ninja-build.

libaio.h: No such file or directory when building the async I/O op. Install the development headers: apt install libaio-dev on Debian/Ubuntu, or yum install libaio-devel on RHEL/CentOS.

CUDA version mismatch errors. The CUDA toolkit version used to compile ops must be compatible with the CUDA version PyTorch was built against. Check uv run python -c "import torch; print(torch.version.cuda)" and ensure nvcc --version reports a compatible version.

error: invalid command 'bdist_wheel' during install. The wheel package is missing from the environment. Run pip install wheel first, then retry. This happens on minimal base images that don’t ship wheel by default.

DS_BUILD_OPS=1 fails on a machine without a GPU. Pre-compilation requires CUDA headers and a GPU-compatible toolchain even if no physical GPU is present. On CPU-only machines, take the wheel from the index, or skip DS_BUILD_OPS and let ops JIT-compile on the GPU machine at runtime.

Related

Handbook articles:

External resources:

Last updated on
How to Install CUDA PyTorch Wheels Without a GPUHow to Install Flash-Attention