Update installation

Signed-off-by: Cheng-Ping Hsieh <[email protected]>

Author: hsiehjackson

docker/Dockerfile

@@ -19,9 +19,10 @@ WORKDIR /workspace/
 COPY ./requirements.txt .
 RUN pip install --upgrade pip \
     && pip install -r requirements.txt \
+    && pip install flash-attn==2.6.0.post1 --no-build-isolation \
     && pip install git+https://github.com/HazyResearch/flash-attention.git#subdirectory=csrc/rotary \
     && pip install causal-conv1d==1.4.0 \
-    && pip install mamba-ssm==2.2.2
+    && pip install mamba-ssm==2.2.2 
 
 RUN [ "python3", "-c", "import nltk; nltk.download('punkt')"]
 

docker/requirements.txt

@@ -12,6 +12,5 @@ tiktoken
 tenacity
 accelerate
 huggingface_hub==0.23.4
-transformers==4.44.1
-flash-attn==2.6.0.post1
+transformers==4.44.2
 vllm==0.5.4