From d82aa8a65d1f1205f6b06f7e62069d76cd11441c Mon Sep 17 00:00:00 2001
From: hotwa <pylyzeng@gmail.com>
Date: Sat, 15 Jun 2024 17:53:11 +0800
Subject: [PATCH] add --no-build-isolation flash-attn

---
 finetune/Dockerfile | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/finetune/Dockerfile b/finetune/Dockerfile
index 562d0e1..13aab77 100644
--- a/finetune/Dockerfile
+++ b/finetune/Dockerfile
@@ -183,20 +183,20 @@ EOT
 # EOT
 
 # PyTorch
-ARG PYTORCH_VERSION=1.13.0
+ARG PYTORCH_VERSION=2.2.1
 ENV PYTORCH_VERSION=${PYTORCH_VERSION}
 RUN <<EOT
 #!/bin/bash
 source /opt/conda/etc/profile.d/conda.sh
 conda activate ${CONDA_ENV_NAME}
-pip install deepspeed torch==${PYTORCH_VERSION} torchvision torchaudio bitsandbytes accelerate transformers optimum
-pip install \
-    torch torchvision torchaudio \
-    pydantic transformers datasets accelerate evaluate peft deepspeed tiktoken \
+pip install torch==${PYTORCH_VERSION} torchvision==0.17.1 torchaudio==${PYTORCH_VERSION} --index-url https://download.pytorch.org/whl/cu121
+pip install --no-build-isolation flash-attn -i https://pypi.org/simple/ --trusted-host pypi.org
+pip install deepspeed bitsandbytes accelerate transformers optimum
+pip install pydantic transformers datasets accelerate evaluate peft deepspeed tiktoken \
     sentencepiece tqdm nltk matplotlib seaborn numpy pandas scikit-learn diffusers \
     huggingface_hub spacy Pillow blobfile requests scipy pycocotools protobuf timm \
-    pyyaml ipython xformers opencv-contrib-python open_clip_torch flash-attn \
-    packaging psutil zstandard
+    pyyaml ipython xformers opencv-contrib-python open_clip_torch \
+    packaging psutil zstandard -i https://pypi.org/simple/ --trusted-host pypi.org
 python -c "import deepspeed; print(deepspeed.__version__)"
 EOT