雑記 - クッキーの日記

参考文献

昨日何があったかというと、jax が native Windows に対応していないために手元の Windows 10 に WSL2 をインストールして CUDA をインストールしたのだが CUDA が動かなかった。公式ドキュメント [1] をよく読むと、以下のようにあった。

5. WSL 2 System Requirements

Ensure you are on the latest WSL Kernel or at least 4.19.121+. Once again we recommend 5.10.16.3 or later for better performance and functional fixes.
If you are on Windows 11, you no longer need to be on Windows Insider Program to use WSL. Refer to Windows11 system requirements.
If you are continuing to use Windows 10, see Windows Insider Preview and Windows 10 Support.

昨日の時点では CUDA が要請する WSL のバージョンの要件を満たしていなかったのだと思う（この日記をかいている時点で既にこのマシンを Windows 11 にアップグレードしてしまったのでバージョンがいくらだったかわからないが）。WSL のバージョンをあげるには Windows 10 のまま Insider Preview Build を導入するか Windows 11 にアップグレードするかのどちらかしなければならないが、おそらく後者の方が楽なので後者にした。そうすると winver の結果が以下となる。昨日は 19xxx だった気がする。覚えていない。

バージョン 21H2 (OS ビルド 22000.593)

Windows 11 にアップグレードしても Windows 10 のときにインストールした WSL はそのまま残っている。そのままだと意味がないのでアップデートする必要がある。自分は Ubuntu-20.04 と WSL を「設定 > アプリ」からアンインストールして入れなおしたがおそらく以下のコマンドを実行すればよい（というかアプリからのアンインストールして入れ直しても Ubuntu-20.04 のインストールに失敗するので [2] の手順を踏む必要があるし、そうしても以下のコマンドまで実行しなければ最新の版にならないようである）。

wsl --update
wsl --status

カーネル バージョン: 5.10.102.1

そのうえで改めて Ubuntu-20.04 ディストリビューションを導入する。

また昨日すっぽかしていたが、そもそもマシンに NVIDIA CUDA on WSL driver [3] を導入しておく必要がある。このマシンには native Windows が GPU を動かすためのドライバは以前から導入されているが、WSL から GPU を動かすためのドライバが導入されていなかったという認識である。ドライバのダウンロード時に指定する Product はタスクマネージャのパフォーマンスタブで GPU を選択すると表示されている。Download Type はどちらを選ぶべきなのかわからないがどちらでもよいようなので Game Ready Driver にする。ドライバをインストール後再起動したうえで [4] にしたがい WSL 上の Ubuntu に CUDA をインストールする。そうすると今度こそ以下がエラーにならない。

$ git clone https://github.com/NVIDIA/cuda-samples.git
$ cd cuda-samples/Samples/5_Domain_Specific/BlackScholes
$ make BlackScholes
$ ./BlackScholes
[./BlackScholes] - Starting...
GPU Device 0: "Pascal" with compute capability 6.1
...

ちなみにホスト Windows 側で PyTorch がこれまで通り実行できるか気になるが少なくとも以下のコマンドは通るようにみえる。

$ python
>>> import torch
>>> print(torch.__version__)
1.11.0+cu113
>>> print(torch.cuda.is_available())
True

次に Ubuntu に cuDNN をインストールする。jax が cuDNN 8.2 を推奨しているような気がするので [5] でダウンロード URL を確認して wget しようとするが、これは 403 エラーになる。

wget https://developer.nvidia.com/compute/machine-learning/cudnn/secure/8.2.0.53/11.3_04222021/Ubuntu20_04-x64/libcudnn8_8.2.0.53-1+cuda11.3_amd64.deb

なのでホスト側でダウンロードしたものを Ubuntu から取り込んでインストールする。

$ cp /mnt/c/Users/c-mih/Downloads/libcudnn8_8.2.0.53-1+cuda11.3_amd64.deb ./
$ sudo dpkg -i libcudnn8_8.2.0.53-1+cuda11.3_amd64.deb

最後に肝心の Python（というより Python は既に入っているので pip）を Ubuntu に導入する。

$ sudo apt update
$ sudo apt upgrade
$ sudo apt install python3-pip

これで Ubuntu に JAX がインストールできる。コマンドの角括弧内は自分が入れた CUDA と cuDNN のバージョンに合ったものを選ぶ。jax の README [6] には CUDA が /usr/local/cuda-X.X にあることを想定しているとあるが上の手順で入れると既にそうなっている。

$ pip install --upgrade pip
$ pip install jax[cuda11_cudnn82] -f https://storage.googleapis.com/jax-releases/jax_releases.html

これできちんと GPU を使用してくれるのか気になるが以下を実行すると GpuDevice とは出てくる。

$ python
Python 3.8.10 (default, Mar 15 2022, 12:22:08)
[GCC 9.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import jax
>>> jax.devices()
[GpuDevice(id=0, process_index=0)]