第一步:创建虚拟环境:
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