# How to Install DeepSpeed


[PyPI](https://pydevtools.com/handbook/explanation/what-is-pypi.md) serves only a source distribution for DeepSpeed, so `pip install deepspeed` compiles instead of downloading a [wheel](https://pydevtools.com/handbook/reference/wheel.md). Prebuilt wheels are on the [Astral GPU indexes](https://wheels.astral.sh/), 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](https://pydevtools.com/handbook/explanation/installing-cuda-python-packages.md) for background.

## Pick an install path

- [Install from the Astral GPU index](#install-from-the-astral-gpu-index) for any project on a CUDA version it covers. Start here.
- [Build from source](#build-from-source) when your CUDA and PyTorch pair falls outside the index, or when you need an op the wheel omits.
- [Install with conda-forge or pixi](#install-with-conda-forge-or-pixi) to hand CUDA toolkit matching to a solver.

## 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](https://pydevtools.com/handbook/explanation/versioning-python-packages-semver-calver-and-pep-440.md) 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:

```bash
uv add torch --index pytorch-cu128=https://download.pytorch.org/whl/cu128
uv add deepspeed --index astral-cu128=https://wheels.astral.sh/simple/cu128/
```

```console
$ 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:

```toml
[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](https://docs.astral.sh/uv/guides/integration/pytorch/#installing-gpu-enabled-pytorch-extensions) 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:

```console
$ 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:

```bash
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`:

```console
$ 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`](#pre-compile-ops-at-install-time) 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](https://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:

```sh
uv pip install deepspeed --no-build-isolation
```
```sh
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:

```sh
DS_BUILD_OPS=1 uv pip install deepspeed --no-build-isolation
```
```sh
DS_BUILD_OPS=1 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:

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

### Add a source build to a uv project

For projects managed with [uv](https://pydevtools.com/handbook/reference/uv.md) using `uv add` and `uv sync`, use [`extra-build-dependencies`](https://docs.astral.sh/uv/concepts/projects/config/#augmenting-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:

```toml
[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`:

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

## Install with conda-forge or pixi

DeepSpeed is available on [conda-forge](https://pydevtools.com/handbook/reference/conda-forge.md), though the version may lag behind PyPI. The conda-forge build handles CUDA toolkit dependencies through the solver. Either [pixi](https://pydevtools.com/handbook/reference/pixi.md) or [conda](https://pydevtools.com/handbook/reference/conda.md) can resolve from conda-forge:

```sh
pixi add deepspeed
```
```sh
conda install -c conda-forge deepspeed
```
For more on when conda-based tools are the better choice for GPU workloads, see [uv vs pixi vs conda for Scientific Python](https://pydevtools.com/handbook/explanation/uv-vs-pixi-vs-conda-for-scientific-python.md).

## Verify the installation

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

```python {filename="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__)
```

```console
$ 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.