Dockerfile build ERROR: Could not build wheels for fused-dense-lib, which is required to install pyproject.toml-based projects

[li126com]

When I install and compile flash-attention from the source code in the dockerfile, I implemented the following command:

cd flash-attention/csrc/fused_dense_lib && /opt/conda/bin/pip install .

And I met the problem:

WARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv
2128.6 Processing /InternLM/flash-attention/csrc/fused_dense_lib
2128.6   Preparing metadata (setup.py): started
2133.5   Preparing metadata (setup.py): finished with status 'done'
2133.5 Building wheels for collected packages: fused-dense-lib
2133.5   Building wheel for fused-dense-lib (setup.py): started
2138.5   Building wheel for fused-dense-lib (setup.py): finished with status 'error'
2138.5   error: subprocess-exited-with-error
2138.5   
2138.5   × python setup.py bdist_wheel did not run successfully.
2138.5   │ exit code: 1
2138.5   ╰─> [61 lines of output]
2138.5       No CUDA runtime is found, using CUDA_HOME='/usr/local/cuda'
2138.5       running bdist_wheel
2138.5       running build
2138.5       running build_ext
2138.5       building 'fused_dense_lib' extension
2138.5       creating /InternLM/flash-attention/csrc/fused_dense_lib/build
2138.5       creating /InternLM/flash-attention/csrc/fused_dense_lib/build/temp.linux-x86_64-cpython-310
2138.5       Traceback (most recent call last):
2138.5         File "<string>", line 2, in <module>
2138.5         File "<pip-setuptools-caller>", line 34, in <module>
2138.5         File "/InternLM/flash-attention/csrc/fused_dense_lib/setup.py", line 26, in <module>
2138.5           setup(
2138.5         File "/opt/conda/lib/python3.10/site-packages/setuptools/__init__.py", line 107, in setup
2138.5           return distutils.core.setup(**attrs)
2138.5         File "/opt/conda/lib/python3.10/site-packages/setuptools/_distutils/core.py", line 185, in setup
2138.5           return run_commands(dist)
2138.5         File "/opt/conda/lib/python3.10/site-packages/setuptools/_distutils/core.py", line 201, in run_commands
2138.5           dist.run_commands()
2138.5         File "/opt/conda/lib/python3.10/site-packages/setuptools/_distutils/dist.py", line 969, in run_commands
2138.5           self.run_command(cmd)
2138.5         File "/opt/conda/lib/python3.10/site-packages/setuptools/dist.py", line 1234, in run_command
2138.5           super().run_command(command)
2138.5         File "/opt/conda/lib/python3.10/site-packages/setuptools/_distutils/dist.py", line 988, in run_command
2138.5           cmd_obj.run()
2138.5         File "/opt/conda/lib/python3.10/site-packages/wheel/bdist_wheel.py", line 325, in run
2138.5           self.run_command("build")
2138.5         File "/opt/conda/lib/python3.10/site-packages/setuptools/_distutils/cmd.py", line 318, in run_command
2138.5           self.distribution.run_command(command)
2138.5         File "/opt/conda/lib/python3.10/site-packages/setuptools/dist.py", line 1234, in run_command
2138.5           super().run_command(command)
2138.5         File "/opt/conda/lib/python3.10/site-packages/setuptools/_distutils/dist.py", line 988, in run_command
2138.5           cmd_obj.run()
2138.5         File "/opt/conda/lib/python3.10/site-packages/setuptools/_distutils/command/build.py", line 131, in run
2138.5           self.run_command(cmd_name)
2138.5         File "/opt/conda/lib/python3.10/site-packages/setuptools/_distutils/cmd.py", line 318, in run_command
2138.5           self.distribution.run_command(command)
2138.5         File "/opt/conda/lib/python3.10/site-packages/setuptools/dist.py", line 1234, in run_command
2138.5           super().run_command(command)
2138.5         File "/opt/conda/lib/python3.10/site-packages/setuptools/_distutils/dist.py", line 988, in run_command
2138.5           cmd_obj.run()
2138.5         File "/opt/conda/lib/python3.10/site-packages/setuptools/command/build_ext.py", line 84, in run
2138.5           _build_ext.run(self)
2138.5         File "/opt/conda/lib/python3.10/site-packages/setuptools/_distutils/command/build_ext.py", line 345, in run
2138.5           self.build_extensions()
2138.5         File "/opt/conda/lib/python3.10/site-packages/torch/utils/cpp_extension.py", line 843, in build_extensions
2138.5           build_ext.build_extensions(self)
2138.5         File "/opt/conda/lib/python3.10/site-packages/setuptools/_distutils/command/build_ext.py", line 467, in build_extensions
2138.5           self._build_extensions_serial()
2138.5         File "/opt/conda/lib/python3.10/site-packages/setuptools/_distutils/command/build_ext.py", line 493, in _build_extensions_serial
2138.5           self.build_extension(ext)
2138.5         File "/opt/conda/lib/python3.10/site-packages/setuptools/command/build_ext.py", line 246, in build_extension
2138.5           _build_ext.build_extension(self, ext)
2138.5         File "/opt/conda/lib/python3.10/site-packages/setuptools/_distutils/command/build_ext.py", line 548, in build_extension
2138.5           objects = self.compiler.compile(
2138.5         File "/opt/conda/lib/python3.10/site-packages/torch/utils/cpp_extension.py", line 649, in unix_wrap_ninja_compile
2138.5           cuda_post_cflags = unix_cuda_flags(cuda_post_cflags)
2138.5         File "/opt/conda/lib/python3.10/site-packages/torch/utils/cpp_extension.py", line 548, in unix_cuda_flags
2138.5           cflags + _get_cuda_arch_flags(cflags))
2138.5         File "/opt/conda/lib/python3.10/site-packages/torch/utils/cpp_extension.py", line 1780, in _get_cuda_arch_flags
2138.5           arch_list[-1] += '+PTX'
2138.5       IndexError: list index out of range
2138.5       [end of output]
2138.5   
2138.5   note: This error originates from a subprocess, and is likely not a problem with pip.
2138.5   ERROR: Failed building wheel for fused-dense-lib
2138.5   Running setup.py clean for fused-dense-lib
2143.5 Failed to build fused-dense-lib
2143.5 ERROR: Could not build wheels for fused-dense-lib, which is required to install pyproject.toml-based projects

[li126com]

From nvidia/cuda:11.7.1-cudnn8-devel-ubuntu18.04 as base

[li126com]

Finally, I solved this problem successfully by adding the following line in front of flash-attention installation in the dockerfile:

ARG TORCH_CUDA_ARCH_LIST="3.7;5.0;6.0;7.0;7.5;8.0;8.6+PTX"

Refer to:
#344 (comment)
pytorch/extension-cpp#71