From 0d691cc2a7b8a10602cafaf88bf3b9f44bd4c878 Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Wed, 23 Apr 2025 14:59:03 -0400 Subject: [PATCH] add base docker image with pytorch 2.7.0 and variant for cuda 12.8 (#2551) * add base docker image with pytorch 2.7.0 and variant for cuda 12.8 * my bash is terrible --- .github/workflows/base.yml | 12 ++++++++++++ docker/Dockerfile-base | 4 ++++ requirements.txt | 1 + 3 files changed, 17 insertions(+) diff --git a/.github/workflows/base.yml b/.github/workflows/base.yml index 30859a374..d6850a0de 100644 --- a/.github/workflows/base.yml +++ b/.github/workflows/base.yml @@ -46,6 +46,18 @@ jobs: python_version: "3.11" pytorch: 2.6.0 torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX" + - cuda: "126" + cuda_version: 12.6.3 + cudnn_version: "" + python_version: "3.11" + pytorch: 2.7.0 + torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX" + - cuda: "128" + cuda_version: 12.6.3 + cudnn_version: "" + python_version: "3.11" + pytorch: 2.7.0 + torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX" - cuda: "128" cuda_version: 12.8.1 cudnn_version: "" diff --git a/docker/Dockerfile-base b/docker/Dockerfile-base index 52201f276..cf1af9682 100644 --- a/docker/Dockerfile-base +++ b/docker/Dockerfile-base @@ -37,3 +37,7 @@ RUN git lfs install --skip-repo && \ pip3 install awscli && \ # The base image ships with `pydantic==1.8.2` which is not working pip3 install -U --no-cache-dir pydantic==1.10.10 + +RUN if [ "$PYTORCH_VERSION" = "2.7.0" ] ; then \ + pip3 install flash-attn==2.7.4.post1; \ + fi diff --git a/requirements.txt b/requirements.txt index 2827c2ca1..02da0abf4 100644 --- a/requirements.txt +++ b/requirements.txt @@ -19,6 +19,7 @@ datasets==3.5.0 deepspeed>=0.15.4 trl==0.16.1 hf_xet==1.0.0 +hqq==0.2.5 optimum==1.16.2 hf_transfer