How to Install DeepSpeed

Install DeepSpeed from the Astral GPU index to get precompiled CUDA operations. PyPI publishes a source distribution whose default Linux installation leaves operations to compile at first use. Installing that Python package is not evidence that its CUDA operations work.

Prerequisites

• Linux x86-64 for the Astral CUDA wheels, with a Python version covered by the chosen PyTorch/wheel pair. Python 3.12 is a practical starting point.
• A supported NVIDIA GPU and a host driver compatible with the selected CUDA runtime.
• A uv project with PyTorch declared directly.

The prebuilt CUDA path does not cover ARM hosts, macOS GPU execution, or every DeepSpeed accelerator. For AMD and other backends, follow DeepSpeed’s accelerator-specific setup guides. Native Windows support is partial; WSL2 uses the Linux stack and requires GPU passthrough.

Install matching CUDA builds

Choose an index pair supported by both the Astral manifest and PyTorch. For CUDA 12.8:

Add these source mappings and explicit indexes to pyproject.toml before adding the dependencies:

[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

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

Then declare both packages directly:

uv add torch deepspeed

Expect a CUDA 12.8 torch build and a deepspeed version containing +cu.12.8.torch. followed by the matching PyTorch minor version. The command updates dependencies, the lockfile, and the environment.

Configure indexes before the first resolve. Changing explicit later does not rewrite sources already recorded in uv.lock.

Both packages must be direct project dependencies. Astral’s wheel metadata constrains the PyTorch minor version, such as torch==2.11.*, but that constraint accepts different CUDA builds of the same PyTorch release. Routing both packages keeps the packaged CUDA runtimes aligned.

Coverage depends on the full Python, platform, PyTorch, and CUDA combination. Check the live manifest rather than assuming every wheel on an index has the same coverage or matches the newest PyPI source release.

Check which operations are precompiled

uv run ds_report

uv run executes the command inside the project’s virtual environment, ensuring all dependencies are installed first. The report distinguishes installed operations from operations compatible with local compilation. Astral’s wheel includes fused_adam; an installed [YES] row means the extension is present, while an execution check is still required.

A precompiled operation needs the matching runtime libraries and driver, but no local nvcc. Other operations may need a CUDA toolkit and C++ compiler to compile at first use. For example, async I/O also needs libaio development files. A missing toolkit in ds_report does not invalidate an already compiled operation.

Current DeepSpeed treats failed operation-compatibility probes as unavailable for compilation, instead of requiring every probe to succeed at import. If an older locked environment fails to import without a toolkit, check its release-specific behavior before applying DS_IGNORE_CUDA_DETECTION; bypassing a probe never supplies missing kernels or a compiler.

Verify a compiled CUDA operation

Save this script on the GPU host. It checks the precompiled fused Adam optimizer directly, without requiring a distributed launcher or BF16-capable hardware:

import torch
import deepspeed
from deepspeed.ops.adam import FusedAdam
from deepspeed.git_version_info import installed_ops

assert torch.cuda.is_available(), 'CUDA is unavailable to PyTorch'
assert torch.version.hip is None, 'This check targets NVIDIA CUDA'
assert installed_ops['fused_adam'], 'The fused Adam operation is not precompiled'
print(deepspeed.__version__, torch.__version__, torch.cuda.get_device_name())
parameter = torch.nn.Parameter(torch.ones(64, device='cuda'))
optimizer = FusedAdam([parameter], lr=0.01)
parameter.sum().backward()
optimizer.step()
torch.cuda.synchronize()
assert parameter.is_cuda and torch.isfinite(parameter).all()
assert (parameter < 1).all()
print('DeepSpeed fused Adam CUDA operation OK')

uv run python check_deepspeed.py

Expect the GPU name and DeepSpeed fused Adam CUDA operation OK. This validates one installed operation; it does not validate ZeRO, CPU offload, every extension, or multi-GPU communication. For distributed setup, use the multi-GPU environment tutorial.

A GPU-less image build can install and import DeepSpeed while accelerator detection selects CPU. Run the CUDA operation check after assigning GPU hardware; a successful image build does not establish GPU functionality.

Build operations from source

If no prebuilt wheel matches, follow DeepSpeed’s source-build instructions. The default Linux source install leaves CUDA operations for JIT compilation; those operations need a full CUDA toolkit with nvcc, a supported C++ compiler, and the build dependencies at runtime.

Precompiling operations requires PyTorch in the build environment. Use --no-build-isolation with PyTorch already installed, or configure uv’s extra-build-dependencies with a runtime-matched torch dependency. DS_BUILD_OPS=1 requests all compatible operations; use per-operation flags such as DS_BUILD_FUSED_ADAM=1 when only one extension is needed.

When compiling on a host without GPUs, explicitly set the deployment architectures through TORCH_CUDA_ARCH_LIST. Device detection during a GPU-less image build cannot identify the deployment GPU, and a runtime-only CUDA image has no compiler. See the CUDA packaging explanation for the distinction between driver, runtime, and toolkit.

Fix a failed operation

• CUDA_HOME or nvcc errors: determine whether the requested operation is installed or being JIT-compiled. A JIT build needs the actual toolkit; setting CUDA_HOME to a runtime-only directory does not provide it.
• Undefined symbols: compare the DeepSpeed wheel’s compiled torch/CUDA versions in ds_report with the installed PyTorch build. Reinstall a matching pair.
• An unavailable offload or inference operation: check that operation’s upstream dependencies and supported hardware. Passing the fused Adam check does not validate other operations.
• No matching distribution: inspect Python, CPU architecture, and CUDA/PyTorch coverage in the manifest. Switch to a covered pair or build from source.