From 26205246367d875302f0d12005354d0a772355d9 Mon Sep 17 00:00:00 2001
From: hotwa <pylyzeng@gmail.com>
Date: Fri, 21 Jun 2024 15:15:15 +0800
Subject: [PATCH] update

---
 finetune/Dockerfile                    | 50 +++++++++++++-------------
 finetune/docker-compose_pytorch2.3.yml |  2 +-
 2 files changed, 27 insertions(+), 25 deletions(-)

diff --git a/finetune/Dockerfile b/finetune/Dockerfile
index cb3af52..62f67e0 100644
--- a/finetune/Dockerfile
+++ b/finetune/Dockerfile
@@ -374,31 +374,33 @@ DS_BUILD_OPS=${DS_BUILD_OPS} pip install dist/deepspeed*.whl --force-reinstall
 # ./install.sh ${DEEPSPEED_INSTALL_FLAGS} --hostfile /job/hostfile # ./install.sh --allow_sudo --pip_sudo --no_clean --hostfile /path/to/your/hostfile
 EOT
 
-# RUN <<EOT
-# #!/bin/bash
-# source /opt/conda/etc/profile.d/conda.sh
-# conda activate ${CONDA_ENV_NAME}
-# # install transformers
-# git clone https://github.com/huggingface/transformers ${STAGE_DIR}/transformers
-# cd ${STAGE_DIR}/transformers
-# python3 ./setup.py develop
-# python3 -m pip install -U --no-cache-dir "pydantic<2"
-# # install flash-attn
-# # pip install packaging -i https://pypi.org/simple/ --trusted-host pypi.org
-# pip install flash-attn --no-build-isolation -i https://pypi.org/simple/ --trusted-host pypi.org
-# EOT
+# install transformers and flash-attn
+RUN <<EOT
+#!/bin/bash
+source /opt/conda/etc/profile.d/conda.sh
+conda activate ${CONDA_ENV_NAME}
+# install transformers
+git clone https://github.com/huggingface/transformers ${STAGE_DIR}/transformers
+cd ${STAGE_DIR}/transformers
+python3 ./setup.py develop
+python3 -m pip install -U --no-cache-dir "pydantic<2"
+# install flash-attn
+# pip install packaging -i https://pypi.org/simple/ --trusted-host pypi.org
+pip install flash-attn --no-build-isolation -i https://pypi.org/simple/ --trusted-host pypi.org
+EOT
 
-# RUN <<EOT
-# #!/bin/bash
-# source /opt/conda/etc/profile.d/conda.sh
-# conda activate ${CONDA_ENV_NAME}
-# pip install optimum
-# pip install peft tiktoken \
-#     tqdm matplotlib seaborn numpy pandas scikit-learn diffusers \
-#     huggingface_hub spacy blobfile pycocotools \
-#     xformers open_clip_torch \
-#     zstandard -i https://pypi.org/simple/ --trusted-host pypi.org
-# EOT
+# other packages
+RUN <<EOT
+#!/bin/bash
+source /opt/conda/etc/profile.d/conda.sh
+conda activate ${CONDA_ENV_NAME}
+pip install optimum
+pip install peft tiktoken \
+    tqdm matplotlib seaborn numpy pandas scikit-learn diffusers \
+    huggingface_hub spacy blobfile pycocotools \
+    xformers open_clip_torch \
+    zstandard -i https://pypi.org/simple/ --trusted-host pypi.org
+EOT
 
 # add vscode server
 # RUN <<EOT
diff --git a/finetune/docker-compose_pytorch2.3.yml b/finetune/docker-compose_pytorch2.3.yml
index 45d2bed..6b9c708 100644
--- a/finetune/docker-compose_pytorch2.3.yml
+++ b/finetune/docker-compose_pytorch2.3.yml
@@ -32,7 +32,7 @@ services:
         USE_ROCM: 0
         USE_XPU: 0
         CUDA: cu121
-        CUDA_ARCH_LIST: "80;86;89;90" # for RTX 4090, all : "80;86;89;90"
+        CUDA_ARCH_LIST: "80;86;89;90;90a" # for RTX 4090, all : "80;86;89;90"
         SETUPTOOLS_VERSION: "69.5.1"
         DCUTLASS_NVCC_ARCHS: "90a" # 90a for H100 GPU 89:GeForce RTX 4090 
     volumes: