update TORCH_CUDA_ARCH_LIST

This commit is contained in:
2024-07-15 12:25:54 +08:00
parent 3fc430582e
commit 47f460bbc0
2 changed files with 3 additions and 3 deletions

View File

@@ -110,8 +110,8 @@ python3 -m pip uninstall -y torch-tensorrt
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;8.9,9.0"
# install apex TORCH_CUDA_ARCH_LIST all "6.0;6.1;6.2;7.0;7.5;8.0;8.6;8.9;9.0"
ARG TORCH_CUDA_ARCH_LIST="8.0;8.6;8.9;9.0"
ENV TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST}
RUN <<EOT
#!/bin/bash

View File

@@ -31,7 +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"
TORCH_CUDA_ARCH_LIST: "8.0;8.6;8.9;9.0" # all "6.0;6.1;6.2;7.0;7.5;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"