WSL2에서 flash_attn 패키지 설치가 잘 안될때 (cuda version)

배경

요즘 flash attention이 핫하다길레 써보기나 하자 하고 흔한 명령어를 입력했다.

 

문제

pip install flash_attn

Collecting flash_attn
  Downloading flash_attn-2.5.9.post1.tar.gz (2.6 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 2.6/2.6 MB 8.8 MB/s eta 0:00:00
  Preparing metadata (setup.py) ... error
  error: subprocess-exited-with-error

  × python setup.py egg_info did not run successfully.
  │ exit code: 1
  ╰─> [12 lines of output]
      fatal: not a git repository (or any of the parent directories): .git
      Traceback (most recent call last):
        File "<string>", line 2, in <module>
        File "<pip-setuptools-caller>", line 34, in <module>
        File "/tmp/pip-install-leg8d8iq/flash-attn_43ce0d408e194ea8b7d3a74aace58a09/setup.py", line 115, in <module>
          raise RuntimeError(
      RuntimeError: FlashAttention is only supported on CUDA 11.6 and above.  Note: make sure nvcc has a supported version by running nvcc -V.


      torch.__version__  = 2.2.0+cu121


      [end of output]

  note: This error originates from a subprocess, and is likely not a problem with pip.
error: metadata-generation-failed

× Encountered error while generating package metadata.
╰─> See above for output.

note: This is an issue with the package mentioned above, not pip.
hint: See above for details.
(base) lovian@suhome:~$

 

그래 CUDA 버전이 11.6 이상이 필요한데, 나는 그 이하라 안된단다. 그러면서 nvcc -V 명령어를 실행해서 버전을 확인하라고 한다.

 

(base) lovian@suhome:~$ nvcc -V
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2021 NVIDIA Corporation
Built on Thu_Nov_18_09:45:30_PST_2021
Cuda compilation tools, release 11.5, V11.5.119
Build cuda_11.5.r11.5/compiler.30672275_0

 

11.5라고 한다. 그럼 새로 깔면 되겠군.

해결 과정

NVIDIA의 가이드를 따라서 설치 한다.

https://docs.nvidia.com/cuda/wsl-user-guide/index.html

CUDA on WSL

 

용량이 꽤 많은 관계로 설치에 시간이 걸린다.

 

약간의 반전

설치 완료 후 혹시 모르니 쉘을 다시 띄우고, 기분 좋게 flash attention을 설치해본다.

어 똑같다...

nvcc -V 도 똑같다.

당연히 될거라 기대하고 했는데, 안되길레 순간적으로 머리가 마비되었다.

 

추적

다시 정신차리고, nvcc 위치를 찾아보자

$ which nvcc
/usr/bin/nvcc
$ ls -al /usr/bin/nvcc
-rwxr-xr-x 1 root root 59 Dec  2  2021 /usr/bin/nvcc
$ cat /usr/bin/nvcc
#!/bin/sh

exec /usr/lib/nvidia-cuda-toolkit/bin/nvcc "$@"

다행이다. 이제 cuda toolkit이 설치된 위치를 파악하면 뭐가 나올 것 같다.

 

$ ls -al /usr/lib/nvidia-cuda-toolkit/bin
total 46088
drwxr-xr-x 3 root root     4096 Dec 27  2023 .
drwxr-xr-x 5 root root     4096 Dec 27  2023 ..
-rwxr-xr-x 1 root root 35497184 Nov 19  2021 cicc
drwxr-xr-x 2 root root     4096 Dec 27  2023 crt
-rwxr-xr-x 1 root root     2273 Mar 15  2022 g++
-rwxr-xr-x 1 root root     2273 Mar 15  2022 gcc
-rwxr-xr-x 1 root root  5998072 Nov 19  2021 nvcc
lrwxrwxrwx 1 root root       17 Mar 15  2022 nvcc.profile -> /etc/nvcc.profile
-rwxr-xr-x 1 root root  5668896 Nov 19  2021 nvprof

설치된 날짜를 보니 적어도 오늘 설치한 것은 아님을 알 수 있다.

 

nvidia에서 받아서 설치한 녀석은 적어도 이곳에 없으니까 찾아보자

$ dpkg -L cuda-toolkit-12-5
/.
/usr
/usr/local
/usr/local/cuda-12.5
/usr/local/cuda-12.5/version.json
/usr/share
/usr/share/doc
/usr/share/doc/cuda-toolkit-12-5
/usr/share/doc/cuda-toolkit-12-5/changelog.Debian.gz

 

껌이네, /usr/local/cuda-12.5에 설치되는 구만.

설마 버전이 들어간 경로를 그대로 쓰는 초보 같은 짓은 안했을거라 기대하면서 /usr/local/을 살펴보자

$ ls -al /usr/local
total 48
drwxr-xr-x 12 root root 4096 Jul 10 23:48 .
drwxr-xr-x 14 root root 4096 May  2  2023 ..
drwxr-xr-x  2 root root 4096 Jan 16 20:27 bin
lrwxrwxrwx  1 root root   22 Jan 16 20:27 cuda -> /etc/alternatives/cuda
lrwxrwxrwx  1 root root   25 Jan 16 20:27 cuda-12 -> /etc/alternatives/cuda-12
drwxr-xr-x 15 root root 4096 Jul 10 23:55 cuda-12.3
drwxr-xr-x 15 root root 4096 Jul 10 23:50 cuda-12.5
drwxr-xr-x  2 root root 4096 May  2  2023 etc
drwxr-xr-x  2 root root 4096 May  2  2023 games
drwxr-xr-x  2 root root 4096 May  2  2023 include
drwxr-xr-x  3 root root 4096 May  2  2023 lib
lrwxrwxrwx  1 root root    9 May  2  2023 man -> share/man
drwxr-xr-x  2 root root 4096 May  2  2023 sbin
drwxr-xr-x  5 root root 4096 Nov 24  2023 share
drwxr-xr-x  2 root root 4096 May  2  2023 src

/etc/alternatives가 있긴 한데, 결국 /usr/local/cuda가 cuda-12.5를 가리키게 되어 있다.

 

아무래도 이미 있던걸 건드리기에는 좀 찝찝해서 cuda를 먼저 참조하도록 경로를 바꾸기로했다.

다들 아는 $HOME/.profile에 PATH를 추가하면 된다.

if [ -d "/usr/local/cuda/bin" ] ; then
    PATH="/usr/local/cuda/bin:$PATH"
fi

 

이제 쉘에 재접속해서 확인해보자

$ nvcc -V
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2024 NVIDIA Corporation
Built on Thu_Jun__6_02:18:23_PDT_2024
Cuda compilation tools, release 12.5, V12.5.82
Build cuda_12.5.r12.5/compiler.34385749_0

 

당연히 flash attention 설치도 성공.

끝.