Implementation:Deepspeedai DeepSpeed Setup

Knowledge Sources	DeepSpeed
Domains	Build_System, Installation, Package_Management
Last Updated	2026-02-09 00:00 GMT

Overview

The setup.py file configures DeepSpeed package installation, dependency resolution, and C++/CUDA operator pre-compilation.

Description

DeepSpeed's setup.py orchestrates the complete package build and installation process using setuptools. It manages Python dependencies across multiple optional feature sets (1bit compression, autotuning, sparse attention, etc.), detects hardware capabilities (CUDA, ROCm, CPU architecture), and conditionally pre-compiles C++/CUDA operators based on environment variables and system compatibility. The script generates version metadata including git commit information and hardware configuration details that are embedded into the installed package for runtime validation.

The setup process supports flexible build modes: JIT-only installation (default on Linux) where operators compile on first use, or ahead-of-time compilation (default on Windows) where all compatible operators are pre-built during installation. Environment variables like DS_BUILD_OPS and operator-specific DS_BUILD_{OP_NAME} flags provide fine-grained control over which operators to compile. The script also handles platform-specific concerns such as Windows symlink creation, ROCm/CUDA version detection, and compute capability targeting.

Usage

Use setup.py during package installation via pip install or python -m build. For development, install in editable mode with pip install -e . to enable code changes without reinstalling. Control operator pre-compilation by setting DS_BUILD_OPS=1 to build all compatible ops, or use specific flags like DS_BUILD_CPU_ADAM=1 to build individual operators. For cross-compilation or custom GPU architectures, set TORCH_CUDA_ARCH_LIST to target specific compute capabilities. The setup script automatically detects hardware and adjusts build configuration, but advanced users can override detection with environment variables.

Code Reference

Source Location

• Repository: DeepSpeed
• File: setup.py

Signature

# Main setup call
setup(
    name='deepspeed',
    version=version_str,
    description='DeepSpeed library',
    author='DeepSpeed Team',
    install_requires=install_requires,
    extras_require=extras_require,
    packages=find_packages(include=['deepspeed', 'deepspeed.*']),
    scripts=scripts,
    ext_modules=ext_modules,
    cmdclass=cmdclass
)

# Key helper functions
def fetch_requirements(path: str) -> List[str]:
    """Load requirements from file."""

def is_env_set(key: str) -> bool:
    """Check if environment variable is set and non-empty."""

def get_env_if_set(key: str, default: Any = "") -> Any:
    """Get environment variable or default if unset/empty."""

def command_exists(cmd: str) -> bool:
    """Check if shell command is available."""

def op_envvar(op_name: str) -> str:
    """Get environment variable name for operator (e.g., DS_BUILD_CPU_ADAM)."""

def op_enabled(op_name: str) -> bool:
    """Check if operator building is enabled."""

Import

# Installation from PyPI (end users)
pip install deepspeed

# Installation from source with default settings (JIT mode on Linux)
pip install .

# Installation with all operators pre-compiled
DS_BUILD_OPS=1 pip install .

# Installation with specific operators
DS_BUILD_CPU_ADAM=1 DS_BUILD_TRANSFORMER=1 pip install .

# Development installation (editable mode)
pip install -e .

# Installation with extra dependencies
pip install deepspeed[autotuning,sparse_attn]

I/O Contract

Environment Variables (Inputs)

Variable	Type	Default	Description
DS_BUILD_OPS	int	0 (Linux), 1 (Windows)	Global flag to enable/disable operator pre-compilation
DS_BUILD_{OP_NAME}	int	DS_BUILD_OPS	Per-operator build control (e.g., DS_BUILD_CPU_ADAM, DS_BUILD_TRANSFORMER)
DS_BUILD_STRING	str	+{git_hash}	Version suffix for distribution builds (e.g., .dev20240101)
DS_ENABLE_NINJA	int	0	Enable ninja build system for faster compilation
TORCH_CUDA_ARCH_LIST	str	Auto-detected	Target GPU compute capabilities (e.g., "7.0;8.0;8.6")
DS_SKIP_CUDA_CHECK	int	0	Skip CUDA version compatibility validation
CUDA_HOME	str	Auto-detected	Path to CUDA toolkit
ROCM_HOME	str	Auto-detected	Path to ROCm installation

Generated Files (Outputs)

File	Content	Purpose
deepspeed/git_version_info_installed.py	version, git_hash, git_branch, installed_ops, accelerator_name, torch_info	Runtime version and build metadata
build.txt	Build string (e.g., .dev20240101)	Version suffix for distribution packages
dist/*.whl	Python wheel file	Distributable package
deepspeed/ops/csrc/ (Windows)	Copy of csrc/	Windows-specific build structure
deepspeed/ops/op_builder/ (Windows)	Copy of op_builder/	Windows-specific build structure

Package Metadata (Outputs)

Field	Value	Description
name	deepspeed	Package name on PyPI
version	{version_str}	Semantic version with optional suffix
install_requires	List[str]	Core dependencies (torch, numpy, etc.)
extras_require	Dict[str, List[str]]	Optional feature dependencies
packages	List[str]	Python packages to install
ext_modules	List[Extension]	Compiled C++/CUDA extensions
scripts	List[str]	Command-line tools (deepspeed, ds_report, etc.)

Usage Examples

Standard Installation

# Install from PyPI with JIT compilation
pip install deepspeed

# Install from source with default settings
git clone https://github.com/microsoft/DeepSpeed.git
cd DeepSpeed
pip install .

# Install with all dependencies
pip install deepspeed[all]

Pre-Compile All Compatible Operators

# Build all operators during installation
DS_BUILD_OPS=1 pip install .

# Verify installed operators
python -c "from deepspeed.git_version_info import installed_ops; print(installed_ops)"

Selective Operator Compilation

# Only build CPU Adam and Transformer operators
DS_BUILD_CPU_ADAM=1 DS_BUILD_TRANSFORMER=1 pip install .

# Build all ops except spatial inference
DS_BUILD_OPS=1 DS_BUILD_SPATIAL_INFERENCE=0 pip install .

# Disable specific op when DS_BUILD_OPS=1
DS_BUILD_OPS=1 DS_BUILD_FP_QUANTIZER=0 pip install .

Target Specific GPU Architectures

# Build for Pascal and Volta (compute capabilities 6.0, 7.0)
TORCH_CUDA_ARCH_LIST="6.0;7.0" DS_BUILD_OPS=1 pip install .

# Build for Ampere and newer with PTX for forward compatibility
TORCH_CUDA_ARCH_LIST="8.0;8.6;9.0+PTX" DS_BUILD_OPS=1 pip install .

# Alternative format with spaces
TORCH_CUDA_ARCH_LIST="6.0 7.0 7.5 8.0 8.6" DS_BUILD_OPS=1 pip install .

Development Installation

# Editable install for development
pip install -e .[dev]

# Build with debug symbols
CFLAGS="-g" DS_BUILD_OPS=1 pip install -e .

# Enable verbose CUDA compilation
DS_DEBUG_CUDA_BUILD=1 DS_BUILD_OPS=1 pip install -e .

Distribution Build

# Create wheel for distribution
DS_BUILD_STRING=".dev20240209" python -m build --no-isolation

# Build wheel with all ops pre-compiled
DS_BUILD_OPS=1 DS_BUILD_STRING="" python -m build --no-isolation

# Result: dist/deepspeed-{version}-py3-none-any.whl

Installation with Optional Features

# Install with 1-bit compression support
pip install deepspeed[1bit]

# Install with autotuning capabilities
pip install deepspeed[autotuning]

# Install with sparse attention
pip install deepspeed[sparse_attn]

# Install with inference optimizations
pip install deepspeed[inf]

# Combine multiple features
pip install deepspeed[autotuning,sparse_attn,triton]

Windows Installation

# Windows builds default to DS_BUILD_OPS=1
# Run from Administrator console for symlink creation
python setup.py build_ext
python -m build

# Or use the provided batch script
build_win.bat

Check Installation Details

# Check installed version and build info
import deepspeed
print(f"DeepSpeed version: {deepspeed.__version__}")

from deepspeed.git_version_info import (
    version, git_hash, git_branch,
    installed_ops, accelerator_name, torch_info
)

print(f"Git: {git_branch}@{git_hash}")
print(f"Accelerator: {accelerator_name}")
print(f"Torch info: {torch_info}")
print(f"Pre-compiled ops: {[k for k, v in installed_ops.items() if v]}")

Troubleshooting Installation

# Skip CUDA version check (use with caution)
DS_SKIP_CUDA_CHECK=1 pip install .

# Disable ninja for debugging
DS_ENABLE_NINJA=0 DS_BUILD_OPS=1 pip install .

# Force CPU-only build
CUDA_HOME="" pip install .

# Verbose output
pip install . -v

# Clean previous build artifacts
pip uninstall deepspeed
rm -rf build dist *.egg-info
pip install .

Related Pages

• OpBuilder - Base class for building operators
• All_Ops - Registry of all available operators
• Accelerator - Hardware abstraction layer
• CPU_Adam - Example of a compiled operator
• Installation - Installation guide
• Environment_Variables - Complete env var reference

Installation Workflow

The setup.py script follows this execution flow:

1. Dependency Resolution (lines 81-125):

• Loads base requirements from requirements/requirements.txt
• Configures extras_require for optional features (1bit, autotuning, sparse_attn, etc.)
• Adds platform-specific dependencies (nvidia-ml-py for CUDA, cupy for compression)
• Creates [all] extra with union of all optional dependencies

2. Build Configuration (lines 149-154):

• Determines BUILD_OP_DEFAULT: 1 on Windows, 0 on Linux
• Reads DS_BUILD_OPS environment variable
• Sets up custom build_ext command with optional ninja support

3. Operator Selection (lines 182-202):

• Iterates through ALL_OPS registry from op_builder/all_ops.py
• For each operator:
  • Checks is_compatible() for hardware/software requirements
  • Reads DS_BUILD_{OP_NAME} environment variable
  • If enabled and compatible, adds builder() to ext_modules list
  • For ROCm, calls hipify_extension() to convert CUDA to HIP code

4. Version Generation (lines 206-291):

• Reads base version from version.txt
• Appends git commit hash or DS_BUILD_STRING
• Detects torch version, CUDA/HIP version, NCCL version, bf16 support
• Writes deepspeed/git_version_info_installed.py with metadata

5. Windows-Specific Setup (lines 225-235):

• Copies csrc/, op_builder/, accelerator/ into deepspeed/ops/
• Creates symlinks for build structure
• Uses MANIFEST_win.in for package data

6. Setup Invocation (lines 311-335):

• Calls setuptools.setup() with all configuration
• Installs Python packages via find_packages()
• Compiles ext_modules (C++/CUDA operators)
• Installs command-line scripts (deepspeed, ds_report, etc.)

Operator Build Variables

Each operator has a corresponding DS_BUILD_{OP_NAME} environment variable:

Operator	Environment Variable	Default
cpu_adam	DS_BUILD_CPU_ADAM	DS_BUILD_OPS
cpu_adagrad	DS_BUILD_CPU_ADAGRAD	DS_BUILD_OPS
fused_adam	DS_BUILD_FUSED_ADAM	DS_BUILD_OPS
fused_lamb	DS_BUILD_FUSED_LAMB	DS_BUILD_OPS
sparse_attn	DS_BUILD_SPARSE_ATTN	DS_BUILD_OPS
transformer	DS_BUILD_TRANSFORMER	DS_BUILD_OPS
stochastic_transformer	DS_BUILD_STOCHASTIC_TRANSFORMER	DS_BUILD_OPS
async_io	DS_BUILD_AIO	DS_BUILD_OPS
utils	DS_BUILD_UTILS	DS_BUILD_OPS
quantizer	DS_BUILD_QUANTIZER	DS_BUILD_OPS
transformer_inference	DS_BUILD_TRANSFORMER_INFERENCE	DS_BUILD_OPS
spatial_inference	DS_BUILD_SPATIAL_INFERENCE	DS_BUILD_OPS

See op_builder/all_ops.py for the complete list of operators and their BUILD_VAR definitions.

Extra Dependencies

The extras_require dictionary defines optional feature groups:

Extra	Purpose	Key Dependencies
1bit	1-bit compression using cupy	cupy-cuda{version} or cupy-rocm-{version}
1bit_mpi	1-bit compression with MPI	mpi4py, cupy
autotuning	Automatic hyperparameter tuning	ray, pandas, scikit-learn
autotuning_ml	ML-based autotuning	torch, transformers
sparse_attn	Sparse attention patterns	triton
sparse	Sparse tensor pruning	scipy, numpy
inf	Inference optimizations	py-cpuinfo
sd	Stable diffusion support	diffusers, transformers, accelerate
triton	Triton kernel compilation	triton
deepcompile	Deep compilation optimizations	torch
dev	Development tools	pytest, pre-commit, mypy
readthedocs	Documentation building	sphinx, myst-parser
all	All optional dependencies	Union of all extras

Platform-Specific Behavior

Linux/Unix:

• Default: DS_BUILD_OPS=0 (JIT compilation mode)
• Uses bash commands for git operations
• Installs full script set: deepspeed, ds_ssh, ds_bench, ds_elastic, etc.

Windows:

• Default: DS_BUILD_OPS=1 (pre-compile operators)
• Requires Administrator privileges for symlink creation
• Copies source directories into deepspeed/ops/ for build
• Uses MANIFEST_win.in for package data inclusion
• Installs limited script set: deepspeed.bat, ds_report.bat, ds, ds_report
• Uses Windows-compatible build commands

ROCm/HIP:

• Auto-detects ROCm version from torch.version.hip
• Calls hipify_extension() on compatible operators
• Sets PYTORCH_ROCM_ARCH environment variable
• Uses cupy-rocm instead of cupy-cuda for compression

Build Metadata

The generated git_version_info_installed.py contains:

# Example generated file
version='0.13.0+a44fb58'
git_hash='a44fb58'
git_branch='master'
installed_ops={
    'cpu_adam': True,
    'cpu_adagrad': True,
    'fused_adam': True,
    'transformer': True,
    'sparse_attn': False,
    'async_io': True,
    # ... all operators
}
accelerator_name='cuda'
torch_info={
    'version': '2.1',
    'bf16_support': True,
    'cuda_version': '12.1',
    'nccl_version': '2.19',
    'hip_version': '0.0'
}

This metadata enables runtime validation and helps diagnose version mismatches between build-time and runtime environments.