Update installation

Signed-off-by: Cheng-Ping Hsieh <[email protected]>
NVIDIA · Oct 10, 2024 · e40fae1 · e40fae1
1 parent 0193bdd
commit e40fae1
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/docker/Dockerfile b/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')"]
 

diff --git a/docker/requirements.txt b/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