mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2026-01-29 21:49:54 +00:00
42 lines
799 B
Bash
42 lines
799 B
Bash
#!/usr/bin/env bash
|
|
|
|
method=${1}
|
|
url=${2}
|
|
filename=${3}
|
|
cuda_version=${4}
|
|
python_version=${5}
|
|
torch_version=${6}
|
|
flags=${@:7}
|
|
|
|
git reset --hard HEAD
|
|
mkdir -p ./all_dist
|
|
source activate base
|
|
conda create -n $python_version -y python=$python_version
|
|
source activate $python_version
|
|
|
|
if [ $1 == "pip" ]
|
|
then
|
|
wget -nc -q -O ./pip_wheels/$filename $url
|
|
pip install ./pip_wheels/$filename
|
|
|
|
elif [ $1 == 'conda' ]
|
|
then
|
|
conda install pytorch==$torch_version cudatoolkit=$cuda_version $flags
|
|
else
|
|
echo Invalid installation method
|
|
exit
|
|
fi
|
|
|
|
if [ $cuda_version == "10.2" ]
|
|
then
|
|
cp -r cub-1.8.0/cub/ colossalai/kernel/cuda_native/csrc/kernels/include/
|
|
fi
|
|
|
|
python setup.py bdist_wheel
|
|
mv ./dist/* ./all_dist
|
|
python setup.py clean
|
|
conda deactivate
|
|
conda env remove -n $python_version
|
|
|
|
|