Install PyTorch with CUDA 11 from source via pip method

CUDA 11 has appeared a while, but we haven’t seen binary PyTorch installation with CUDA 11 in official PyTorch website yet. So I decided to build PyTorch from source with CUDA 11. However, in the, the recommended method to install related dependencies is “conda install”. Since we don’t know if “conda install -c pytorch magma-cuda110” is available, and I prefer pip install method as personal taste too, so I tried to install other dependencies via pip, build magma with cuda 11 and install it first, and then build pytorch wheel package from source, and finally pip install the package in my ubuntu 20.04 system.

In the following I assume that we already have CUDA 11 and cudnn installed in an ubuntu 20.04 Linux system with Nvidia GPU (not too old one, with compute capability >= 6.0). The steps are as follows.

  1. Run the following pip install command as system admin / root (or sudo pip install …):
    pip install numpy ninja pyyaml mkl mkl-include setuptools cmake cffi typing_extensions future six requests

  2. Go to website to download the latest magma-2.5.3.tar.gz. In a download working directory, run
    tar xzf magma-2.5.3.tar.gz
    to untar unzip it, and get a sub directory magma-2.5.3.
    cd magma-2.5.3
    and follow README or
    to build magma, and install it as root. The default install location is /usr/local/magma. Go to /etc/ and create a file named magma.conf there with one line text “/usr/local/magma/lib”, and run command “ldconfig” as root. Note that build magma-2.5.3 is basically to run
    make lib
    make test
    make sparse-lib
    make sparse-test
    However, with cuda 11, the commands
    make sparse-lib
    make sparse-test
    simply fail, and
    make lib
    has an issue too. Follow to fix the issue. Ignore make sparse-lib and make sparse-test. make lib and make test are all we need.

  3. Finally, follow, get the PyTorch source
    git clone --recursive
    cd pytorch
    git submodule sync
    git submodule update --init --recursive
    Instead of run “python install”, we run
    python bdist_wheel
    to build a python wheel package. If everything runs successfully, we will get a *.whl file in a directory named dist.
    cd dist
    pip install *.whl
    (or if not root, sudo pip install *.whl)
    That’s it.