From a5946ff1f07ab57405c0ec050ec2a16ea2c1b6c5 Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Sat, 5 Jul 2025 09:21:18 -0400 Subject: [PATCH] build fa2 from source for base image with torch2.6 and cu124 (#2867) --- .github/workflows/base.yml | 6 ++++-- docker/Dockerfile-base | 4 ++++ 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/.github/workflows/base.yml b/.github/workflows/base.yml index 966bd2f5b..1b9862d85 100644 --- a/.github/workflows/base.yml +++ b/.github/workflows/base.yml @@ -5,11 +5,13 @@ on: branches: - "main" paths: - - 'Dockerfile-base' + - 'docker/Dockerfile-base' + - 'docker/Dockerfile-uv-base' - '.github/workflows/base.yml' pull_request: paths: - - 'Dockerfile-base' + - 'docker/Dockerfile-base' + - 'docker/Dockerfile-uv-base' - '.github/workflows/base.yml' workflow_dispatch: diff --git a/docker/Dockerfile-base b/docker/Dockerfile-base index 52201f276..df4240325 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.6.0" ] && [ "$CUDA" = "124" ] ; then \ + FLASH_ATTENTION_FORCE_BUILD="TRUE" pip3 install --no-build-isolation flash-attn==2.8.0.post2; \ + fi