build: try 28

.github/workflows/build-docker-images-extra.yml

@@ -0,0 +1,42 @@
+name: Build other docker images (scheduled)
+
+on:
+  push:
+    branches:
+      - build_ci_docker_image_extra*
+      - docker-with-fa3
+  schedule:
+    - cron: "17 0 * * *"
+
+concurrency:
+  group: docker-images-extra-builds
+  cancel-in-progress: false
+
+jobs:
+  latest-docker:
+    name: "Latest PyTorch [dev] with FA3"
+    runs-on:
+      group: aws-highcpu-32-priv
+    steps:
+      -
+        name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3
+      -
+        name: Check out code
+        uses: actions/checkout@v4
+      -
+        name: Login to DockerHub
+        uses: docker/login-action@v3
+        with:
+          username: ${{ secrets.DOCKERHUB_USERNAME }}
+          password: ${{ secrets.DOCKERHUB_PASSWORD }}
+      -
+        name: Build and push
+        uses: docker/build-push-action@v5
+        with:
+          context: ./docker/transformers-all-latest-gpu
+          build-args: |
+            REF=main
+            FA3=TRUE
+          push: true
+          tags: huggingface/transformers-all-latest-gpu:fa3

docker/transformers-all-latest-gpu/Dockerfile

@@ -57,6 +57,15 @@ RUN python3 -m pip install --no-cache-dir quanto
 # After using A10 as CI runner, let's run FA2 tests
 RUN [ "$PYTORCH" != "pre" ] && python3 -m pip uninstall -y ninja && python3 -m pip install --no-cache-dir ninja && python3 -m pip install flash-attn --no-cache-dir --no-build-isolation || echo "Don't install FA2 with nightly torch"
 
+ARG FA3=FALSE
+RUN if [ "$FA3" != "FALSE" ]; then \
+        git clone https://github.com/Dao-AILab/flash-attention && \
+        cd flash-attention/hopper && \
+        MAX_JOBS=28 python3 setup.py install; \
+    else \
+        echo "FA3 is not installed."; \
+    fi
+
 # TODO (ydshieh): check this again
 # `quanto` will install `ninja` which leads to many `CUDA error: an illegal memory access ...` in some model tests
 # (`deformable_detr`, `rwkv`, `mra`)