add TORCH_CUDA_ARCH_LIST arg

2024-07-15 12:06:43 +08:00
parent 7b1a20d0b8
commit 3fc430582e
1 changed files with 1 additions and 0 deletions
--- a/finetune/docker-compose_update.yml
+++ b/finetune/docker-compose_update.yml
@@ -31,6 +31,7 @@ services:
        USE_XPU: 0
        CUDA: cu121
        CUDA_ARCH_LIST: "80;86;89;90" # for RTX 4090, all : "80;86;89;90"
+        TORCH_CUDA_ARCH_LIST: "8.0;8.6;8.9,9.0"
        SETUPTOOLS_VERSION: "69.5.1"
        DCUTLASS_NVCC_ARCHS: "80;86;89;90;90a" # 90a for H100 GPU 89:GeForce RTX 4090 
        DEEPSPEED_VERSION: "master"