新虚拟环境中PyTorch已安装但flash_attn仍提示找不到的解决方法

我正在测试一些模型，为此创建了多个虚拟环境。其中一个环境安装PyTorch的过程很顺利，但在另一个新创建的虚拟环境中，安装flash_attn时提示缺少PyTorch。于是我执行了以下命令：

python3 -m pip install torch

执行结果显示PyTorch已经安装：

Requirement already satisfied: torch in ./ATI/lib/python3.12/site-packages (2.7.1)
Requirement already satisfied: filelock in ./ATI/lib/python3.12/site-packages (from torch) (3.18.0)
Requirement already satisfied: typing-extensions>=4.10.0 in ./ATI/lib/python3.12/site-packages (from torch) (4.14.0)
Requirement already satisfied: setuptools in ./ATI/lib/python3.12/site-packages (from torch) (80.9.0)
Requirement already satisfied: sympy>=1.13.3 in ./ATI/lib/python3.12/site-packages (from torch) (1.14.0)
Requirement already satisfied: networkx in ./ATI/lib/python3.12/site-packages (from torch) (3.5)
Requirement already satisfied: jinja2 in ./ATI/lib/python3.12/site-packages (from torch) (3.1.6)
Requirement already satisfied: fsspec in ./ATI/lib/python3.12/site-packages (from torch) (2025.5.1)
Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.6.77 in ./ATI/lib/python3.12/site-packages (from torch) (12.6.77)
Requirement already satisfied: nvidia-cuda-runtime-cu12==12.6.77 in ./ATI/lib/python3.12/site-packages (from torch) (12.6.77)
Requirement already satisfied: nvidia-cuda-cupti-cu12==12.6.80 in ./ATI/lib/python3.12/site-packages (from torch) (12.6.80)
Requirement already satisfied: nvidia-cudnn-cu12==9.5.1.17 in ./ATI/lib/python3.12/site-packages (from torch) (9.5.1.17)
Requirement already satisfied: nvidia-cublas-cu12==12.6.4.1 in ./ATI/lib/python3.12/site-packages (from torch) (12.6.4.1)
Requirement already satisfied: nvidia-cufft-cu12==11.3.0.4 in ./ATI/lib/python3.12/site-packages (from torch) (11.3.0.4)
Requirement already satisfied: nvidia-curand-cu12==10.3.7.77 in ./ATI/lib/python3.12/site-packages (from torch) (10.3.7.77)
Requirement already satisfied: nvidia-cusolver-cu12==11.7.1.2 in ./ATI/lib/python3.12/site-packages (from torch) (11.7.1.2)
Requirement already satisfied: nvidia-cusparse-cu12==12.5.4.2 in ./ATI/lib/python3.12/site-packages (from torch) (12.5.4.2)
Requirement already satisfied: nvidia-cusparselt-cu12==0.6.3 in ./ATI/lib/python3.12/site-packages (from torch) (0.6.3)
Requirement already satisfied: nvidia-nccl-cu12==2.26.2 in ./ATI/lib/python3.12/site-packages (from torch) (2.26.2)
Requirement already satisfied: nvidia-nvtx-cu12==12.6.77 in ./ATI/lib/python3.12/site-packages (from torch) (12.6.77)
Requirement already satisfied: nvidia-nvjitlink-cu12==12.6.85 in ./ATI/lib/python3.12/site-packages (from torch) (12.6.85)
Requirement already satisfied: nvidia-cufile-cu12==1.11.1.6 in ./ATI/lib/python3.12/site-packages (from torch) (1.11.1.6)
Requirement already satisfied: triton==3.3.1 in ./ATI/lib/python3.12/site-packages (from torch) (3.3.1)
Requirement already satisfied: mpmath<1.4,>=1.1.0 in ./ATI/lib/python3.12/site-packages (from sympy>=1.13.3->torch) (1.3.0)
Requirement already satisfied: MarkupSafe>=2.0 in ./ATI/lib/python3.12/site-packages (from jinja2->torch) (3.0.2)

但再次尝试安装flash_attn时，出现以下报错：

Collecting flash_attn
  Using cached flash_attn-2.8.0.post2.tar.gz (7.9 MB)
  Installing build dependencies ... done
  Getting requirements to build wheel ... error
  error: subprocess-exited-with-error
  
  × Getting requirements to build wheel did not run successfully.
  │ exit code: 1
  ╰─> [20 lines of output]
      Traceback (most recent call last):
        File "/home/mickey/dev/ATI/ATI/lib/python3.12/site-packages/pip/_vendor/pyproject_hooks/_in_process/_in_process.py", line 389, in <module>
          main()
        File "/home/mickey/dev/ATI/ATI/lib/python3.12/site-packages/pip/_vendor/pyproject_hooks/_in_process/_in_process.py", line 373, in main
          json_out["return_val"] = hook(**hook_input["kwargs"])
                                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
        File "/home/mickey/dev/ATI/ATI/lib/python3.12/site-packages/pip/_vendor/pyproject_hooks/_in_process/_in_process.py", line 143, in get_requires_for_build_wheel
          return hook(config_settings)
                 ^^^^^^^^^^^^^^^^^^^^^
        File "/tmp/pip-build-env-nt9zjl19/overlay/lib/python3.12/site-packages/setuptools/build_meta.py", line 331, in get_requires_for_build_wheel
          return self._get_build_requires(config_settings, requirements=[])
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
        File "/tmp/pip-build-env-nt9zjl19/overlay/lib/python3.12/site-packages/setuptools/build_meta.py", line 301, in _get_build_requires
          self.run_setup()
        File "/tmp/pip-build-env-nt9zjl19/overlay/lib/python3.12/site-packages/setuptools/build_meta.py", line 512, in run_setup
          super().run_setup(setup_script=setup_script)
        File "/tmp/pip-build-env-nt9zjl19/overlay/lib/python3.12/site-packages/setuptools/build_meta.py", line 317, in run_setup
          exec(code, locals())
        File "<string>", line 22, in <module>
      ModuleNotFoundError: No module named 'torch'
      [end of output]
  
  note: This error originates from a subprocess, and is likely not a problem with pip.
error: subprocess-exited-with-error

× Getting requirements to build wheel did not run successfully.
│ exit code: 1
╰─> See above for output.

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

• 确认虚拟环境激活状态：执行which python3和which pip，检查输出路径是否指向./ATI/lib/python3.12/目录下的可执行文件，确保操作的是目标虚拟环境，而非系统环境或其他虚拟环境。
• 用当前环境Python调用pip安装：安装flash_attn时，直接使用当前环境的Python来执行pip，避免环境混淆：

  python3 -m pip install flash_attn
  

• 验证PyTorch可导入性：激活环境后，打开Python交互终端，输入import torch，如果报错说明PyTorch安装不完整或存在问题，需重新安装。建议使用PyTorch官方针对对应CUDA版本的安装命令，而非默认的pip install torch。
• 尝试预编译轮子：如果源码编译时依赖检测失败，可以安装对应PyTorch和CUDA版本的预编译flash_attn轮子，跳过编译环节的环境检测。
• 清理缓存后重新安装：执行以下命令清理pip缓存，避免缓存导致的依赖问题：

  python3 -m pip cache purge
  python3 -m pip install flash_attn --no-cache-dir
  

• 检查Python版本兼容性：确认flash_attn是否支持Python 3.12，若不支持，创建Python 3.11或3.10版本的虚拟环境后再尝试安装。

内容的提问来源于stack exchange，提问作者easythrees