第一步：创建虚拟环境：

conda create -n deepspeed python=3.10

第二步：进入虚拟环境，安装Pytorch 2.3.1

# CUDA 12.1
conda install pytorch==2.3.1 torchvision==0.18.1 torchaudio==2.3.1 pytorch-cuda=12.1 -c pytorch -c nvidia

第三步：查看机器的cuda版本，并记录下来，操作方法：在虚拟环境中执行：

~$ CUDA_VISIBLE_DEVICES=0 python -c "import torch; print(torch.cuda.get_device_capability())"
(8, 9)
~$ CUDA_VISIBLE_DEVICES=0 python -c "import torch; print(torch.cuda.get_device_properties(torch.device('cuda')))"
_CudaDeviceProperties(name='NVIDIA A800 80GB PCIe', major=8, minor=0, total_memory=81050MB, multi_processor_count=108)
~$ CUDA_VISIBLE_DEVICES=0 python -c "import torch; print(torch.cuda.get_arch_list())"
['sm_50', 'sm_60', 'sm_61', 'sm_70', 'sm_75', 'sm_80', 'sm_86', 'sm_37', 'sm_90', 'compute_37']
 

对应命令如下：

CUDA_VISIBLE_DEVICES=0 python -c "import torch; print(torch.cuda.get_device_capability())"CUDA_VISIBLE_DEVICES=0 python -c "import torch; print(torch.cuda.get_device_properties(torch.device('cuda')))"CUDA_VISIBLE_DEVICES=0 python -c "import torch; print(torch.cuda.get_arch_list())"

第四步：克隆deepspeed仓库

git clone https://github.com/microsoft/DeepSpeed.git

第五步：安装DeepSpeed

TORCH_CUDA_ARCH_LIST="8.9" DS_BUILD_CPU_ADAM=1 DS_BUILD_UTILS=1 pip install . --global-option="build_ext" --global-option="-j8" --no-cache -v --disable-pip-version-check 2>&1 | tee build.log

上述中TORCH_CUDA_ARCH_LIST="8.9"代表着使用的cuda版本，需要根据实际情况设置

如果需要使用 CPU Offload 优化器参数，设置参数 DS_BUILD_CPU_ADAM=1；
如果需要使用 NVMe Offload，设置参数 DS_BUILD_UTILS=1：

最后安装成功， 安装成功后可以通过ds_report来查看对应的安装报告。

可能遇到的问题：error: can't find Rust compiler

解决办法：

# 参考：https://anaconda.org/conda-forge/rust
安装命令：conda install conda-forge::rust