[Tips] PyTorchをself buildしてinstallする方法

はじめに

東京科学大学(Institute of Science Tokyo)の藤井です。

突然ですが、pip installでインストールできるPyTorchのdistributed backendにはncclなどが指定できますが、mpiを指定してtorch.distributed.new_group(backend='mpi')のようにしたい場合は、どうすれば良いかご存知でしょうか？

基本的にはself buildしか方法はないというのが答えです。
PyTorch公式が配布しているPyTorchのdistributed backendにはmpiを指定することができません。また、最新のCUDA Toolkitや古いCUDA Tookitに依存したPyTorchは公式から配布されていないことが多々あります。

本記事では、配布されているPyTorchでは実験が行えないような場合に、self buildを行いsource codeからPyTorchをinstallする方法について紹介します。
また、今回紹介する方法は、大学や研究機関のスパコンかつコンテナが自由に使えない環境を想定しているためsudoなしで環境構築することが可能です。

インストール

ソースコードの準備

PyTorch v2.6.0をインストールしてみましょう。
まず、以下のように必要なコードを手元に落としてきます。

git clone --recursive https://github.com/pytorch/pytorch
cd pytorch
git checkout v2.6.0
git submodule sync
git submodule update --init --recursive

手元に複数のversionのPyTorchを置いておきたいです。名前が衝突してしまうと思うのですが良い管理方法はありませんか？

以下のようにするのはいかがでしょうか？

git clone --recursive https://github.com/pytorch/pytorch pytorch-v2.5.1

インストールオプションの設定

次にインストールオプションを設定します。
以下に示すのはあくまで一例ですので、こちらを参考に適時変更してください。

export USE_DISTRIBUTED=1
export TORCH_USE_CUDA_DSA=1

export CUDA_VERSION=12.8
export CUDNN_VERSION=9.7.0
export CXX_COMPILER="/usr/bin/c++"
export CXX_FLAGS="-D_GLIBCXX_USE_CXX11_ABI=0 -fabi-version=11 -fvisibility-inlines-hidden -DUSE_PTHREADPOOL -DNDEBUG -DUSE_KINETO -DLIBKINETO_NOROCTRACER -DLIBKINETO_NOXPUPTI=ON -DUSE_FBGEMM -DUSE_PYTORCH_QNNPACK -DUSE_XNNPACK -DSYMBOLICATE_MOBILE_DEBUG_HANDLE -O2 -fPIC -Wall -Wextra -Werror=return-type -Werror=non-virtual-dtor -Werror=bool-operation -Wnarrowing -Wno-missing-field-initializers -Wno-type-limits -Wno-array-bounds -Wno-unknown-pragmas -Wno-unused-parameter -Wno-strict-overflow -Wno-strict-aliasing -Wno-stringop-overflow -Wsuggest-override -Wno-psabi -Wno-error=old-style-cast -Wno-missing-braces -fdiagnostics-color=always -faligned-new -Wno-unused-but-set-variable -Wno-maybe-uninitialized -fno-math-errno -fno-trapping-math -Werror=format -Wno-stringop-overflow"
export LAPACK_INFO=mkl
export PERF_WITH_AVX=1
export PERF_WITH_AVX2=1
export USE_CUDA=ON
export USE_CUDNN=ON
export USE_CUSPARSELT=1
export USE_EXCEPTION_PTR=1
export USE_GFLAGS=OFF
export USE_GLOG=OFF
export USE_GLOO=ON
export USE_MKL=ON
export USE_MKLDNN=ON
export USE_MPI=ON
export USE_NCCL=1
export USE_NNPACK=ON
export USE_OPENMP=ON
export USE_ROCM=OFF
export USE_ROCM_KERNEL_ASSERT=OFF
export TORCH_VERSION=2.6.0

export MAX_JOBS=64

なお、手元の環境のCPUが貧弱な場合はexport MAX_JOBS=8のように減らしてみてください。
インストールに失敗する場合の多くが必要なライブラリが足りていないか、export MAX_JOBSの指定が大きすぎることです。

\tmpの容量が足らないとエラーが出ました。どうすればよいでしょうか？

以下のように環境変数を設定してみてください。

export TMPDIR="/path/to/tmp"
export TMP="/path/to/tmp"

インストールの実行

インストールする前に必要なライブラリへのパスを環境変数に設定しましょう。
moduleファイルが整備されている環境の場合は以下のように設定してください。

module load cuda/12.8
module load cudnn/9.7.0
module load nccl/2.25.1-cuda12.8
module load hpcx/2.21.0

あとは、installをするだけです。以下のコマンドでインストール可能です。

python setup.py develop

インストールが完了するとFinished processing dependencies for torch==2.6.0a0+git1eba9b3のようなメッセージが表示されます。

pip installでインストールした通常のPyTorchのinstall configを確認するにはどうすれば良いでしょうか？

以下のように確認できます。

>>> import torch
>>> torch.__config__.show()
'PyTorch built with:\n  - GCC 11.4\n  - C++ Version: 201703\n  - Intel(R) MKL-DNN v3.5.3 (Git Hash 66f0cb9eb66affd2da3bf5f8d897376f04aae6af)\n  - OpenMP 201511 (a.k.a. OpenMP 4.5)\n  - NNPACK is enabled\n  - CPU capability usage: AVX512\n  - CUDA Runtime 12.8\n  - NVCC architecture flags: -gencode;arch=compute_90,code=sm_90\n  - CuDNN 90.7\n  - Build settings: BUILD_TYPE=Release, COMMIT_SHA=1eba9b3aa3c43f86f4a2c807ac8e12c4a7767340, CUDA_VERSION=12.8, CUDNN_VERSION=9.7.0, CXX_COMPILER=/bin/c++, CXX_FLAGS= -D_GLIBCXX_USE_CXX11_ABI=1 -fvisibility-inlines-hidden -DUSE_PTHREADPOOL -DNDEBUG -DUSE_KINETO -DLIBKINETO_NOROCTRACER -DLIBKINETO_NOXPUPTI=ON -DUSE_FBGEMM -DUSE_PYTORCH_QNNPACK -DUSE_XNNPACK -DSYMBOLICATE_MOBILE_DEBUG_HANDLE -O2 -fPIC -Wall -Wextra -Werror=return-type -Werror=non-virtual-dtor -Werror=range-loop-construct -Werror=bool-operation -Wnarrowing -Wno-missing-field-initializers -Wno-type-limits -Wno-array-bounds -Wno-unknown-pragmas -Wno-unused-parameter -Wno-strict-overflow -Wno-strict-aliasing -Wno-stringop-overflow -Wsuggest-override -Wno-psabi -Wno-error=old-style-cast -Wno-missing-braces -fdiagnostics-color=always -faligned-new -Wno-unused-but-set-variable -Wno-maybe-uninitialized -fno-math-errno -fno-trapping-math -Werror=format -Wno-stringop-overflow, PERF_WITH_AVX=1, PERF_WITH_AVX2=1, TORCH_VERSION=2.6.0, USE_CUDA=1, USE_CUDNN=ON, USE_CUSPARSELT=OFF, USE_EIGEN_FOR_BLAS=ON, USE_EXCEPTION_PTR=1, USE_GFLAGS=OFF, USE_GLOG=OFF, USE_GLOO=ON, USE_MKL=OFF, USE_MKLDNN=ON, USE_MPI=1, USE_NCCL=ON, USE_NNPACK=ON, USE_OPENMP=ON, USE_ROCM=OFF, USE_ROCM_KERNEL_ASSERT=OFF, \n'

おわりに

本記事では PyTorchをsource codeからインストールする方法について紹介しました。
何か特殊な用途で、配布されているPyTorchとは異なるものが必要でしたら、こちらの方法でインストールしてみてください。

普段は以下のようにLLMの開発や、低精度学習に関する研究開発を行っています。
よろしければ、そちらもご覧ください。

https://zenn.dev/kaz20/articles/52fee6a3c9ea3b

https://zenn.dev/tokyotech_lm/articles/d6cb3a8fdfc907