SpackでGPU対応PyTorchをインストールする

毎回調べている気がするのでGPUを使えるPyTorchの構築をSpackで行う際のメモしておく。

必要なパッケージを一つづつ入れていってもいいのだが、コンパイラやcudaの依存関係やインストール時の最適化を簡単に解決してくれるスパコン向けのパッケージマネージャであるSpack (<https://spack.io/>)を利用することにより少ない手間で一貫した環境が構築可能となる。Spackの使い方は最低限のことしか触れないため[公式ドキュメント](<https://spack.readthedocs.io/en/latest/>)や[ABCIの説明ページ](<https://docs.abci.ai/ja/tips/spack/>) などを参照。

Ubuntu 20.04.2 LTS (GNU/Linux 5.8.0-48-generic x86_64) と NVIDIA A100 の環境で構築。
GPUドライバを入れるところだけは管理者権限が必要だが、それ以降は一般ユーザー権限でよい。
管理者が入れて一般ユーザがパッケージを使う場合はspackの[upstream機能](<https://spack.readthedocs.io/en/latest/chain.html>)を使うことになると思う。

GPUドライバが入っており、対応するGCCやclangコンパイラがシステムに入っていれば一般ユーザーで適当なディレクトリを作りその中で

```console
$ git clone https://github.com/spack/spack.git
$ cd spack
$ . share/spack/setup-env.sh
$ spack compiler find
$ spack -v install -j 10 py-torch
```

とするだけで必要なパッケージを全て導入してくれる可能性は高い。ただし、パッケージのバージョンやオプションなど下記で説明するパッケージファイルによって指定することにより、今後の環境構築は楽になる。

## GPUドライバの確認、導入

すでにGPUドライバが導入されていたら不要。
Ubuntuなので下記の方法が最も簡単。

```console
$ sudo add-apt-repository ppa:graphics-drivers/ppa
$ ubuntu-drivers devices
== /sys/devices/pci0000:20/0000:20:03.1/0000:21:00.0 ==
modalias : pci:v000010DEd000020F1sv000010DEsd0000145Fbc03sc02i00
vendor   : NVIDIA Corporation
driver   : nvidia-driver-460 - third-party non-free recommended  # 現環境ではこれが推奨されている
driver   : nvidia-driver-450-server - distro non-free
driver   : nvidia-driver-450 - third-party non-free
driver   : nvidia-driver-460-server - distro non-free
driver   : xserver-xorg-video-nouveau - distro free builtin
```

`ppa:graphics-drivers/ppa` リポジトリを追加し、`ubuntu-drivers devices` コマンドで現環境についているGPUの推奨ドライバを確認。
推奨ドライバがすでに入っていれば次はスキップし、入っていなければ今あるドライバを削除。

```console
$ sudo apt purge nvidia-\*  # 推奨ドライバが入っていればpurgeは不要
$ sudo apt autoremove
```

ドライバのインストール。

```
$ sudo apt install nvidia-driver-460
or
$ sudo apt install aptitude # aptで入らなかったらaptitudeで入れる
$ sudo aptitude install nvidia-driver-460
$ sudo reboot
```

再起動して確認。

```
$ nvidia-smi
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 460.67       Driver Version: 460.67       CUDA Version: 11.2     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  A100-PCIE-40GB      Off  | 00000000:21:00.0 Off |                    0 |
| N/A   29C    P0    32W / 250W |      4MiB / 40536MiB |      0%      Default |
|                               |                      |             Disabled |
+-------------------------------+----------------------+----------------------+
|   1  A100-PCIE-40GB      Off  | 00000000:81:00.0 Off |                    0 |
| N/A   28C    P0    32W / 250W |      4MiB / 40536MiB |      0%      Default |
|                               |                      |             Disabled |
+-------------------------------+----------------------+----------------------+
```

GPUが見えていたらドライバの導入は完了。ここまではUbuntuに特化した説明だが以降はLinuxディストリビューションに関係なく同じ。

## CUDAやコンパイラバージョンの確認

`nvidia-smi` の結果に `CUDA Version: 11.2` とあるのでcudaのバージョンは11.2が導入可能。<https://docs.nvidia.com/cuda/cuda-installation-guide-linux/index.html> から確認できる情報ではUbuntu 20.04ではCUDA 11.2を入れる際にGCCバージョンが9.xである必要がある。 以降はGCCコンパイラを元に説明していくがclangを使いたい場合はコンパイラ部分を変更すればいいはず。

<img src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/53817/b32269be-694b-8673-bc90-50044f0d3613.png" width="600">

## Spackで対応GCCのインストール

Ubuntu 20.04 LTSではデフォルトのコンパイラは

```console
$ which gcc
/usr/bin/gcc

$ gcc --version
gcc (Ubuntu 9.3.0-17ubuntu1~20.04) 9.3.0
Copyright (C) 2019 Free Software Foundation, Inc.
This is free software; see the source for copying conditions.  There is NO
warranty; not even for MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
```

バージョンが9.3.0なのでこのまま使うこともできるが、システムアップデートなどでバージョン更新されると不便なのでSpackを使用してバージョンを固定したGCCコンパイラを導入する。spackのパッケージを入れる適当なディレクトリを作成し、その中でspackをダウンロードして環境を適用する。

```console
$ git clone https://github.com/spack/spack.git
$ cd spack
$ . share/spack/setup-env.sh
$ echo $SPACK_ROOT   # このパスの中にパッケージ類が入る
/home/hoge/spack
```

以降、spackの機能を使う際はログインしたときに `$SPACK_ROOT/share/spack/setup-env.sh` をドットコマンドで読み込む。

最初に使えるコンパイラの登録をする。コンパイラ情報はユーザーのホームディレクトリの下の `~/.spack/linux/compilers.yaml` に作られる。

```console
$ spack find compiler
$ spack compiler list
spack compiler list
==> Available compilers
-- clang ubuntu20.04-x86_64 -------------------------------------
clang@10.0.0

-- gcc ubuntu20.04-x86_64 ---------------------------------------
gcc@9.3.0
```

この段階でのコンパイラー情報は次のようになっている。

```yaml:~/.spack/linux/compilers.yaml
compilers:
- compiler:
    spec: clang@10.0.0
    paths:
      cc: /usr/bin/clang
      cxx: /usr/bin/clang++
      f77:
      fc:
    flags: {}
    operating_system: ubuntu20.04
    target: x86_64
    modules: []
    environment: {}
    extra_rpaths: []
- compiler:
    spec: gcc@9.3.0
    paths:
      cc: /usr/bin/gcc
      cxx: /usr/bin/g++
      f77: /usr/bin/gfortran
      fc: /usr/bin/gfortran
    flags: {}
    operating_system: ubuntu20.04
    target: x86_64
    modules: []
    environment: {}
    extra_rpaths: []
```

次にシステムの`gcc@9.3.0`を使い、spackの`gcc@9.3.0`を構築する。

```console
$ spack info gcc 
$ spack spec gcc@9.3.0%gcc@9.3.0
$ spack -v install -j 8 gcc@9.3.0%gcc@9.3.0
```

- spack infoでgccパッケージについて確認
- spack specでインストールされるパッケージの依存関係やオプションなどを確認。 @9.3.0はバージョン指定
- `%gcc@9.3.0`で構築するコンパイラ指定。別バージョンのgccやclangなどが混在しているときのために明示的に指定している。
- specで問題なければinstall。 `-v`は詳細表示。 `-j 8`は8並列でビルド。

構築までかなりの時間がかかるがインストールが終わったら確認。

```console
$ spack find  # spackで入れたパッケージが全て表示される
==> 24 installed packages
-- linux-ubuntu20.04-zen2 / gcc@9.3.0 ---------------------------
autoconf@2.69        bzip2@1.0.8    gdbm@1.19     help2man@1.47.16  libtool@2.4.6   mpc@1.1.0    perl@5.32.1    tar@1.34
automake@1.16.3      diffutils@3.7  gettext@0.21  libiconv@1.16     libxml2@2.9.10  mpfr@3.1.6   pkgconf@1.7.4  xz@5.2.5
berkeley-db@18.1.40  gcc@9.3.0      gmp@6.2.1     libsigsegv@2.12   m4@1.4.18       ncurses@6.2  readline@8.1   zlib@1.2.11

$ spack find -x # 明示的に指定して入れたパッケージが表示される
==> 1 installed package
-- linux-ubuntu20.04-zen2 / gcc@9.3.0 ---------------------------
gcc@9.3.0
```

インストールしたgccをloadしてみる。

```
$ spack load gcc@9.3.0
$ which gcc
/home/hoge/spack/opt/spack/linux-ubuntu20.04-zen2/gcc-9.3.0/gcc-9.3.0-bbkkmscclawlbad4ds72wh7ouuuvc74u/bin/gcc
$ gcc --version
gcc (Spack GCC) 9.3.0
Copyright (C) 2019 Free Software Foundation, Inc.
This is free software; see the source for copying conditions.  There is NO
warranty; not even for MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
```

コンパイラリストに登録。

```console
$ spack compiler find
$ spack compiler list
```

```yaml:~/.spack/linux/compilers.yaml
- compiler:
    spec: gcc@9.3.0
    paths:
      cc: /home/hoge/spack/opt/spack/linux-ubuntu20.04-zen2/gcc-9.3.0/gcc-9.3.0-bbkkmscclawlbad4ds72wh7ouuuvc74u/bin/gcc
      cxx: /home/hoge/spack/opt/spack/linux-ubuntu20.04-zen2/gcc-9.3.0/gcc-9.3.0-bbkkmscclawlbad4ds72wh7ouuuvc74u/bin/g++
      f77: /home/hoge/spack/opt/spack/linux-ubuntu20.04-zen2/gcc-9.3.0/gcc-9.3.0-bbkkmscclawlbad4ds72wh7ouuuvc74u/bin/gfortran
      fc: /home/hoge/spack/opt/spack/linux-ubuntu20.04-zen2/gcc-9.3.0/gcc-9.3.0-bbkkmscclawlbad4ds72wh7ouuuvc74u/bin/gfortran
    flags: {}
    operating_system: ubuntu20.04
    target: x86_64
    modules: []
    environment: {}
    extra_rpaths: []
```


`~/.spack/linux/compilers.yaml`のパスが上のwhichのものになっていれば完了。

今回のようにシステムのデフォルトがで入れたものと同じ名前、アーキテクチャでうまく上書きされないようなら、一旦 compilers.yaml に登録されている(今はシステムの)`gcc@9.3.0`情報を名前を変えるか削除して登録し直す。

```console
$ spack compiler remove gcc@9.3.0 
$ spack compiler find
$ spack compiler list
```


## SpackでPyTorchのインストール

これからPyTorchのインストールを行うが、パッケージによって必要なオプションがあれば `~/.spack/packages.yaml` に記述しておけば`install`行でオプションを指定する必要はなくなる。このようなファイルを作っておけば別環境でも同じバージョン、オプションを指定した環境が容易に再現可能となる。バージョン指定をしなければ基本的に最新の安定版がインストールされる。指定可能なバージョンやオプションの候補は`spack info`で全て確認できる。

cuda_archに指定するCompute Capabilityは <https://developer.nvidia.com/cuda-gpus> で確認できる。

```yaml:~/.spack/packages.yaml
packages:
  python:
    variants:
      tkinter= True

  py-matplotlib:
    variants:
      movies= True
      backend= tkagg

  openmpi:
    version: [4.0.5]
    compiler: [gcc]

  fftw:
    version: [3.3.9]
    compiler: [gcc]
    variants:
      mpi= True
      openmp= True
      pfft_patches= True
      precision= float,double
    providers:
      mpi: [openmpi]

  cuda:
    version: [11.2.2]

  cudnn:
    version: [8.1.1.33-11.2]


  py-torch:
    variants:
      ### V100, GF 2080 Ti 75, A100 80
      cuda_arch= 80
      distributed= True
      nccl= True
      gloo= True

  py-torch-nvidia-apex:
    variants:
      ### V100, GF 2080 Ti 75, A100 80
      cuda_arch= 80

  all:
    compiler: [gcc]
    providers:
      mpi: [openmpi]
    permissions:
      read: world
      write: user
```

上記のような `packages.yaml` を作っておけば

```console
$ spack spec py-torch%gcc@9.3.0
Input spec
--------------------------------
py-torch%gcc@9.3.0

Concretized
--------------------------------
py-torch@1.8.1%gcc@9.3.0~binary~caffe2+cuda+cudnn+distributed~fbgemm~ffmpeg+gloo~leveldb~lmdb~magma~miopen+mkldnn+nccl~nnpack~opencv+openmp~qnnpack~redis~rocm~tbb~test~xnnpack~zstd cuda_arch=80 arch=linux-ubuntu20.04-zen2
    ^cmake@3.20.0%gcc@9.3.0~doc+ncurses+openssl+ownlibs~qt build_type=Release arch=linux-ubuntu20.04-zen2
         ...
    ^cuda@11.2.2%gcc@9.3.0~dev arch=linux-ubuntu20.04-zen2
         ...
    ^cudnn@8.1.1.33-11.2%gcc@9.3.0 arch=linux-ubuntu20.04-zen2
    ^eigen@3.3.9%gcc@9.3.0~ipo build_type=RelWithDebInfo arch=linux-ubuntu20.04-zen2
    ^gloo@master%gcc@9.3.0~ipo build_type=RelWithDebInfo arch=linux-ubuntu20.04-zen2
    ^nccl@2.8.3-1%gcc@9.3.0+cuda cuda_arch=none arch=linux-ubuntu20.04-zen2
        ... 
  ^python@3.8.8%gcc@9.3.0+bz2+ctypes+dbm~debug+libxml2+lzma~nis~optimizations+pic+pyexpat+pythoncmd+readline+shared+sqlite3+ssl~tix+tkinter~ucs4+uuid+zlib patches=...
                        ^expat@2.2.10%gcc@9.3.0+libbsd arch=linux-ubuntu20.04-zen2
                            ^libbsd@0.10.0%gcc@9.3.0 arch=linux-ubuntu20.04-zen2
                        ^libffi@3.3%gcc@9.3.0 patches=26f26c6f29a7ce9bf370ad3ab2610f99365b4bdd7b82e7c31df41a3370d685c0 arch=linux-ubuntu20.04-zen2
                        ^sqlite@3.35.3%gcc@9.3.0+column_metadata+fts~functions~rtree arch=linux-ubuntu20.04-zen2
                        ^tcl@8.6.11%gcc@9.3.0 arch=linux-ubuntu20.04-zen2
                        ^tk@8.6.10%gcc@9.3.0+xft+xss arch=linux-ubuntu20.04-zen2
                            ^libx11@1.7.0%gcc@9.3.0 arch=linux-ubuntu20.04-zen2
                                ^inputproto@2.3.2%gcc@9.3.0 arch=linux-ubuntu20.04-zen2
                                    ^util-macros@1.19.1%gcc@9.3.0 arch=linux-ubuntu20.04-zen2
                                ^kbproto@1.0.7%gcc@9.3.0 arch=linux-ubuntu20.04-zen2
...
```

のように表示され、cuda, cudnnのバージョンやpythonのtkオプションも指定したものがインストールされることを確認できる。

```console
$ spack -v install -j 20 py-torch%gcc@9.3.0
```

で上記で確認した依存関係を含めてインストール。インストールが完了したら

```console
$ spack find -ldpfv py-torch
```

でオプションや依存関係やインストール場所が表示されるので一応確認しておく。その他、必要なpythonパッケージがあれば導入する。

```
$ spack -v install -j 20 py-numpy%gcc@9.3.0 
$ spack -v install -j 20 py-ipython%gcc@9.3.0 
$ spack -v install -j 20 py-scipy%gcc@9.3.0 
...
```

- spackでloadしたpythonのpipから入れることもできるが、それだとHPC用の最適化オプションが効いていないため、上記のようにspackからいちいちパッケージを導入したもののほうが高速に動作すると思われる
- 逆に言えばspackに用意されていないパッケージはpythonをloadしてpipで入れられる(この場合はpipで入れたパッケージはspackのpythonの下に作られるためシステムのpythonと競合することもない)


## GPU対応Pytorchの確認

py-torchパッケージをloadするとPATHやLD_LIBRARY_PATHに反映され使えるようになる。

```
$ which python3
/usr/bin/python3

$ spack load python py-torch py-ipython
$ which python3 
/home/hoge/spack/opt/spack/linux-ubuntu20.04-zen2/gcc-9.3.0/python-3.8.8-7tijmoltswgxk4jlxn55fpsunnpkj3hq/bin/python3
```

GPUの確認。pytorchからGPUが見えていたら成功。

```
$ ipython
ipython
Python 3.8.8 (default, Apr  2 2021, 17:44:13)
Type 'copyright', 'credits' or 'license' for more information
IPython 7.21.0 -- An enhanced Interactive Python. Type '?' for help.

In [1]: import torch

In [2]: torch.cuda.is_available()
Out[2]: True

In [3]: torch.cuda.get_device_name(0)
Out[3]: 'A100-PCIE-40GB'

In [4]: torch.cuda.get_device_name(1)
Out[4]: 'A100-PCIE-40GB'
```

## Pythonパッケージのactivate

今のままではpythonをloadしただけでは依存関係のないpy-ipythonやpy-matplotlibはいちいちloadしないと使えない。インストールしたpythonパッケージを全部ロードするようなaliasを作ってもいいがspackのactivate機能を使えばpythonをloadしただけで、activateしたパッケージを使えるようになる。

```
$ spack extensions python
$ ==> 49 installed:
-- linux-ubuntu20.04-zen2 / gcc@9.3.0 ---------------------------
llvm@11.1.0           py-h5py@3.2.1              py-numpy@1.20.2           py-pygments@2.6.1         py-threadpoolctl@2.0.0
meson@0.57.1          py-ipython@7.21.0          py-parso@0.8.1            py-pyparsing@2.4.7        py-toml@0.10.2
py-argparse@1.4.0     py-ipython-genutils@0.2.0  py-pexpect@4.7.0          py-python-dateutil@2.8.1  py-torch@1.8.1
...
==> None activated.
```

この段階では50個近くpythonパッケージがあるが直接importする可能性のあるものだけactivateしておけばそれに依存するパッケージも一緒にavtivateされる。

```
$ spack activate py-torch
$ spack activate py-numpy
$ spack activate py-matplotlib
...

$ spack extensions python
==> 36 activated:
-- linux-ubuntu20.04-zen2 / gcc@9.3.0 ---------------------------
py-argparse@1.4.0          py-jedi@0.18.0       py-pickleshare@0.7.5      py-python-dateutil@2.8.1  py-tqdm@4.59.0
py-backcall@0.1.0          py-joblib@1.0.1      py-pillow@7.2.0           py-pyyaml@5.3.1           py-traitlets@5.0.4
...
```

この状態になっていれば次から`spack load python`でこれらのパッケージも使えるようになる。


## キャッシュや不要パッケージの削除

spackでは依存関係も含めて1から構築するためダウンロードしたソースやビルド時に生成された中間ファイルが大量にできるので一通り環境構築が終わったら`spack clean`で不要なキャッシュの削除をしておいてもよい。削除のレベルにいくらかあるので適当に選ぶ。

```
$ spack clean -h
usage: spack clean [-hsdfmpba] ...

remove temporary build files and/or downloaded archives

positional arguments:
  specs               one or more package specs

optional arguments:
  -h, --help          show this help message and exit
  -s, --stage         remove all temporary build stages (default)
  -d, --downloads     remove cached downloads
  -f, --failures      force removal of all install failure tracking markers
  -m, --misc-cache    remove long-lived caches, like the virtual package index
  -p, --python-cache  remove .pyc, .pyo files and __pycache__ folders
  -b, --bootstrap     remove software needed to bootstrap Spack
  -a, --all           equivalent to -sdfmpb
```

以下は不要なものは全部削除している。

```
$ spack clean --all
```


目的のパッケージを構築するために作られる大量の中間パッケージも不要ならば`spack gc` でガベージコレクションできる。ただし、これを行うと次に同様のパッケージを追加するときにまた膨大な時間がかかることになるので注意。

```
spack gc
==> The following packages will be uninstalled:
    -- linux-ubuntu20.04-zen2 / gcc@9.3.0 ---------------------------
    vyivrsu at-spi2-atk@2.34.2            jpk2juk gtkplus@3.24.26     yggsngn m4@1.4.18                         tib7bzf perl-test-needs@0.002005
    rtrmwji at-spi2-core@2.38.0           kzme4an harfbuzz@2.6.8      53hqjji mesa@21.0.0                       3lirbfz perl-try-tiny@0.28
    x7xrjvc atk@2.36.0                    n3h5krl help2man@1.47.16    xmt43zs meson@0.57.1                      5xvxinf perl-uri@1.72
...

==> Do you want to proceed? [y/N] N
==> Aborting uninstallation
```

`spack gc`で削除するパッケージ一覧が出てくるのが、これらを削除したとしても明示的にインストールしたパッケージは動作する。もし消したくないパッケージが一覧にある場合は`spack install パッケージ名`をして`spack find -x`に出てくるようにしておく。

```
$ spack install llvm
$ spack find -x
==> 12 installed packages
-- linux-ubuntu20.04-zen2 / gcc@9.3.0 ---------------------------
gcc@9.3.0    py-argparse@1.4.0  py-h5py@3.2.1      py-matplotlib@3.4.1  py-numpy@1.20.2         py-scipy@1.6.2  python@3.8.8
llvm@11.1.0  py-cython@0.29.22  py-ipython@7.21.0  py-mpi4py@3.0.3      py-scikit-learn@0.24.1  py-torch@1.8.1
```

この一覧に出てくるパッケージとそれに依存するものは削除されなくなる。