From 0378cb2509d4a198931997176849448f22d47690 Mon Sep 17 00:00:00 2001
From: lingyuzeng <pylyzeng@gmail.com>
Date: Mon, 15 Jul 2024 12:03:24 +0800
Subject: [PATCH] add all arch

---
 finetune/Dockerfile.update | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/finetune/Dockerfile.update b/finetune/Dockerfile.update
index 3fc447b..d50cf7d 100644
--- a/finetune/Dockerfile.update
+++ b/finetune/Dockerfile.update
@@ -111,6 +111,8 @@ python3 -m pip uninstall -y apex
 EOT
 
 # install apex
+ARG TORCH_CUDA_ARCH_LIST="6.0;6.1;6.2;7.0;7.5;8.0;8.6;9.0"
+ENV TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST}
 RUN <<EOT
 #!/bin/bash
 source /opt/conda/etc/profile.d/conda.sh
@@ -427,7 +429,6 @@ pip install flash-attn --no-build-isolation -i https://pypi.org/simple/ --truste
 EOT
 
 # other packages
-ENV TORCH_CUDA_ARCH_LIST="80;86;89;90"
 RUN <<EOT
 #!/bin/bash
 source /opt/conda/etc/profile.d/conda.sh