python – How can I use the cudatoolkit installed by conda instead of system cuda?

I want to use cuda version >= 11.6 because of the dependency of flash-attention. But my system cuda version is 11.4. So I want to install PyTorch 2.1 and CUDA 11.8 in a new cuda environment.

I did like that:

conda create -n pt21 python==3.10
conda activate pt21
conda install cudatoolkit==11.8
conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia

It succeeded:

$ conda list | grep cuda
cuda-cudart               11.8.89                       0    nvidia
cuda-cupti                11.8.87                       0    nvidia
cuda-libraries            11.8.0                        0    nvidia
cuda-nvrtc                11.8.89                       0    nvidia
cuda-nvtx                 11.8.86                       0    nvidia
cuda-runtime              11.8.0                        0    nvidia
cudatoolkit               11.8.0               h6a678d5_0    defaults
pytorch                   2.1.0           py3.10_cuda11.8_cudnn8.7.0_0    pytorch
pytorch-cuda              11.8                 h7e8668a_5    pytorch
pytorch-mutex             1.0                        cuda    pytorch

But when I installed flash-attention, it still told me:

RuntimeError: FlashAttention is only supported on CUDA 11.6 and above.  Note: make sure nvcc has a supported version by running nvcc -V.

And I tried:

(pt21) $ nvcc -V
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2020 NVIDIA Corporation
Built on Tue_Sep_15_19:10:02_PDT_2020
Cuda compilation tools, release 11.1, V11.1.74
Build cuda_11.1.TC455_06.29069683_0

(pt21) $ which nvcc
/usr/bin/nvcc

So, how can I use the cudatoolkit installed by conda instead of system cuda?

Read more here: Source link