-
Notifications
You must be signed in to change notification settings - Fork 2
/
compile_torch.sh
executable file
·144 lines (123 loc) · 3.86 KB
/
compile_torch.sh
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
#!/bin/bash
# simple usage: work_path=/home/yhao/pt ./compile_torch.sh
# make sure you have activated the correct conda environment before running this script
set -e
work_path=${work_path:-/home/yhao/p9}
# clean_install=1 will remove the existing pytorch folder and re-clone it
# if not, it will just update the existing pytorch and dependent packages
clean_install=${clean_install:-0}
# clean_torch=1 will run python setup.py clean to remove previous pytorch build files
clean_torch=${clean_torch:-0}
torch_only=${torch_only:-0}
export MAX_JOBS=256
debug=${debug:-0}
# for specific commit or branch
torch_commit=${torch_commit:-""}
torch_branch=${torch_branch:-"main"}
torch_pull=${torch_pull:-0}
# torchbench installation takes a long time, so it can be disabled
no_torchbench=${no_torchbench:-0}
# disable ROCM when working on servers with NVIDIA GPUs and AMD GPUs
export USE_ROCM=0
export USE_NCCL=1
# export ROCR_VISIBLE_DEVICES=3
# export CUDA_VISIBLE_DEVICES=1
# function to check the return value of the previous command
check_return_value() {
if [ $? -ne 0 ]; then
echo "Error: $1"
exit 1
fi
}
# print configs
echo "work_path: ${work_path}"
echo "clean_install: ${clean_install}"
echo "clean_torch: ${clean_torch}"
echo "torch_only: ${torch_only}"
echo "torch_branch: ${torch_branch}"
echo "torch_commit: ${torch_commit}"
# if you have an error named like version `GLIBCXX_3.4.30' not found, you can add `-c conda-forge` to the following command. And also for your `conda create -n pt_compiled -c conda-forge python=3.10` command
conda install -y magma-cuda121 -c pytorch
conda install -y ccache cmake ninja mkl mkl-include libpng libjpeg-turbo graphviz -c conda-forge
export CMAKE_PREFIX_PATH=${CONDA_PREFIX:-"$(dirname $(which conda))/../"}
cd $work_path
if [ $clean_install -eq 1 ]; then
rm -rf pytorch text vision audio benchmark data
git clone --recursive [email protected]:pytorch/pytorch.git
git clone --recursive [email protected]:pytorch/text.git
git clone --recursive [email protected]:pytorch/data.git
git clone --recursive [email protected]:pytorch/vision.git
git clone --recursive [email protected]:pytorch/audio.git
git clone --recursive [email protected]:pytorch/benchmark.git
fi
cd $work_path/pytorch
git fetch
if [ -n "$torch_commit" ]; then
git checkout $torch_commit
echo "warnging: you are using a specific commit. don't forget to create a new branch if you want to make changes"
else
git checkout $torch_branch
fi
if [ $torch_pull -eq 1 ]; then
git pull
fi
git submodule sync
git submodule update --init --recursive
pip install -r requirements.txt
make triton
if [ $clean_torch -eq 1 ]; then
python setup.py clean
fi
if [ $debug -eq 1 ]; then
debug_prefix="env DEBUG=1"
else
debug_prefix=""
fi
${debug_prefix} python setup.py develop
function notify_finish() {
if command -v notify &>/dev/null; then
notify "PyTorch Compilation is done"
fi
}
if [ $torch_only -eq 1 ]; then
notify_finish
exit 0
fi
function upgrade_pack() {
cd $work_path/$1
git pull
# the data.txt is weird, so we don't use it. leave it here for future reference
# git checkout "$work_path/pytorch/.github/ci_commit_pins/$1.txt"
git submodule sync
git submodule update --init --recursive
pip uninstall -y $1
python setup.py clean
python setup.py install
echo "$1 installation is done"
}
# install torchdata
cd $work_path
upgrade_pack data
# install torchtext
cd $work_path
export CC=$(which gcc)
export CXX=$(which g++)
upgrade_pack text
# install torchvision
export FORCE_CUDA=1
upgrade_pack vision
# install torchaudio
upgrade_pack audio
if [ $no_torchbench -eq 1 ]; then
notify_finish
exit 0
fi
# install torchbench
pip install pyyaml
cd $work_path/benchmark
git pull
git submodule sync
git submodule update --init --recursive
python install.py
echo "torchbench installation is done"
notify_finish