add --no-build-isolation flash-attn

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