Unable to install flash attention in docker

[shivance]

Hi Team,

First of all thank you so much for building such an amazing project. Keep up the great work!!!

I am trying to install flash attention in my docker container but I am getting a build error like this:

Dockerfile

FROM nvidia/cuda:12.2.0-base-ubuntu22.04

ENV PYTHON_VERSION=3.10
ENV PYTHONDONTWRITEBYTECODE=1
ENV PYTHONUNBUFFERED=1

RUN apt-get -y update \
    && apt-get -qq install --no-install-recommends software-properties-common curl \
    && add-apt-repository ppa:deadsnakes/ppa \
    && apt-get -y update \
    && apt-get -qq install --no-install-recommends python${PYTHON_VERSION} python${PYTHON_VERSION}-dev python3-pip python3-distutils \
    && rm -rf /var/lib/apt/lists/*

RUN curl -sS https://bootstrap.pypa.io/get-pip.py | python${PYTHON_VERSION}

RUN ln -s -f /usr/bin/python${PYTHON_VERSION} /usr/bin/python3 && \
    ln -s -f /usr/bin/python${PYTHON_VERSION} /usr/bin/python && \
    ln -s -f /usr/bin/pip3 /usr/bin/pip

RUN /usr/bin/python${PYTHON_VERSION} -m pip install --upgrade pip

RUN pip install ninja packaging setuptools
RUN MAX_JOBS=4 pip install flash-attn --no-build-isolation

WORKDIR /opt/ml

# install dependencies
COPY requirements.txt /opt/ml/
RUN /usr/bin/python${PYTHON_VERSION} -m pip install -r requirements.txt
...

requirements.txt

# web service
flask==3.0.3
gunicorn==22.0.0

# base ml libraries
transformers==4.44.2
accelerate==0.30.0
torch==2.3.1

# lint
ruff==0.6.2

# misc
editdistance==0.6.2
ninja
packaging

Build Error

------                                                                                                                                                                                                                                    
 > [container_name  7/12] RUN MAX_JOBS=4 pip install flash-attn --no-build-isolation:                                                                                                                                         
0.650 Collecting flash-attn                                                                                                                                                                                                               
0.664   Downloading flash_attn-2.6.3.tar.gz (2.6 MB)                                                                                                                                                                                      
0.689      ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 2.6/2.6 MB 131.9 MB/s eta 0:00:00                                                                                                                                                     
1.510   Preparing metadata (setup.py): started                                                                                                                                                                                            
1.669   Preparing metadata (setup.py): finished with status 'error'                                                                                                                                                                       
1.674   error: subprocess-exited-with-error                                                                                                                                                                                               
1.674                                                                                                                                                                                                                                     
1.674   × python setup.py egg_info did not run successfully.                                                                                                                                                                              
1.674   │ exit code: 1
1.674   ╰─> [6 lines of output]
1.674       Traceback (most recent call last):
1.674         File "<string>", line 2, in <module>
1.674         File "<pip-setuptools-caller>", line 34, in <module>
1.674         File "/tmp/pip-install-t6_ft_n3/flash-attn_773f20acd3d9454b880a9ba4eb746aeb/setup.py", line 21, in <module>
1.674           import torch
1.674       ModuleNotFoundError: No module named 'torch'
1.674       [end of output]
1.674   
1.674   note: This error originates from a subprocess, and is likely not a problem with pip.
1.732 error: metadata-generation-failed
1.732 
1.732 × Encountered error while generating package metadata.
1.732 ╰─> See above for output.
1.732 
1.732 note: This is an issue with the package mentioned above, not pip.
1.732 hint: See above for details.
------
failed to solve: process "/bin/sh -c MAX_JOBS=4 pip install flash-attn --no-build-isolation" did not complete successfully: exit code: 1