雑記: [PyTorch on WSL2] fast_transformers.causal_product.causal_product_cuda が import できなかったのが解消した話

手元の環境(CUDA 11.6 + cuDNN v8.2.0 on Ubuntu20.04 on WSL2)では pytorch-fast-transformers を pip でインストールしたところ一部の機能は使用できるようだが fast_transformers.causal_product.causal_product_cuda が import できなかった。手元では以下の 2 点を対応すると解消した。

  • 手元では CUDA が 11.6 であるのに対して CUDA 11.5 用の PyTorch を導入していた(2022 年 5 月 1 日現在 11.6 用の PyTorch はないため)。CUDA を 11.5 に落としてこの食い違いを解消した。
  • pytorch-fast-transformers を pip でインストールするのではなく python setup.py でインストールした。
参考文献

  1. No module named 'fast_transformers.causal_product.causal_product_cpu' · Issue #63 · idiap/fast-transformers · GitHub(2022年5月1日参照).
  2. CUDA Toolkit 11.6 Update 2 Downloads | NVIDIA Developer(2022年5月1日参照).
  3. WSL2 上の Ubuntu での NVIDIA CUDA ツールキット, NVIDIA cuDNN, PyTorch, TensorFlow 2.7 のインストールと動作確認(Windows 上)(2022年5月1日参照).


pytorch-fast-transformers を pip でインストールすると fast_transformers.causal_product.causal_product_cuda の import に失敗する件で、参考文献 [1] の nicolasugrinovic 氏の最新の(4 Mar 2021 の)コメントの手順で解消するかもしれないと考えた。が、今更インストールされている CUDA のバージョンが PyTorch が想定するものより新しいものである点を怒られた。この食い違いは認識していたが PyTorch の利用上怒られないので放置していた。

$ python setup.py build_ext --inplace
...
RuntimeError:
The detected CUDA version (11.6) mismatches the version that was used to compile
PyTorch (11.5). Please make sure to use the same CUDA versions.

なので CUDA を 11.5 にする。再度参考文献 [2] を参照して画面に表示されるコマンドを実行する。しかし、最後のコマンドだけは以下に読み替える。

$ sudo apt -y install cuda-11-5

CUDA 11.5 にパスを通しておく [3]。もっとも、実際には 11.6 にパスを通してあったので 6 を 5 に手編集した。

$ echo 'export CUDA_PATH=/usr/local/cuda-11.5' >> ${HOME}/.bashrc
$ echo 'export LD_LIBRARY_PATH=/usr/local/cuda-11.5/lib64:${LD_LIBRARY_PATH}' >> ${HOME}/.bashrc
$ echo 'export PATH=/usr/local/cuda-11.5/bin:${PATH}' >> ${HOME}/.bashrc

再度以下を実行する。と、怒られず処理が進む。

$ python setup.py build_ext --inplace
$ python setup.py install --user

その後以下を実行すると import できるようだ。

$ python
>>> import fast_transformers.causal_product.causal_product_cuda