From 0d27e14e4538846e23fc0614ae0a0505c36e6eae Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Thu, 20 Nov 2025 09:04:37 -0500 Subject: [PATCH] Torch 2.9.1 base images (#3268) * update torch 2.9.1 base images * update base dockerfile image check --- .github/workflows/base.yml | 8 ++++---- docker/Dockerfile-base | 2 +- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/.github/workflows/base.yml b/.github/workflows/base.yml index 2e8950dd9..eddce1438 100644 --- a/.github/workflows/base.yml +++ b/.github/workflows/base.yml @@ -57,14 +57,14 @@ jobs: cuda_version: 12.8.1 cudnn_version: "" python_version: "3.11" - pytorch: 2.9.0 + pytorch: 2.9.1 torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX" dockerfile: "Dockerfile-base" - cuda: "130" cuda_version: 13.0.0 cudnn_version: "" python_version: "3.11" - pytorch: 2.9.0 + pytorch: 2.9.1 torch_cuda_arch_list: "9.0+PTX" dockerfile: "Dockerfile-base" # - cuda: "128" @@ -146,14 +146,14 @@ jobs: cuda_version: 12.8.1 cudnn_version: "" python_version: "3.11" - pytorch: 2.9.0 + pytorch: 2.9.1 torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX" dockerfile: "Dockerfile-uv-base" - cuda: "130" cuda_version: 13.0.0 cudnn_version: "" python_version: "3.11" - pytorch: 2.9.0 + pytorch: 2.9.1 torch_cuda_arch_list: "9.0+PTX" dockerfile: "Dockerfile-uv-base" steps: diff --git a/docker/Dockerfile-base b/docker/Dockerfile-base index 25eae4fde..cfd30b851 100644 --- a/docker/Dockerfile-base +++ b/docker/Dockerfile-base @@ -51,7 +51,7 @@ RUN git lfs install --skip-repo && \ pip3 install -U --no-cache-dir pydantic==1.10.10 && \ pip3 cache purge -RUN if [ "$PYTORCH_VERSION" = "2.9.0" ] && [ "$CUDA" = "128" ] ; then \ +RUN if [ "$PYTORCH_VERSION" = "2.9.1" ] && [ "$CUDA" = "128" ] ; then \ wget https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.4.17/flash_attn-2.8.3+cu128torch2.9-cp311-cp311-linux_x86_64.whl; \ pip3 install --no-cache-dir flash_attn-2.8.3+cu128torch2.9-cp311-cp311-linux_x86_64.whl; \ rm flash_attn-2.8.3+cu128torch2.9-cp311-cp311-linux_x86_64.whl; \