From 116b05f9b0d51135d95b5d5d673b3bcd39e6a635 Mon Sep 17 00:00:00 2001 From: Tri Dao Date: Thu, 11 Jul 2024 02:53:30 -0700 Subject: [PATCH] [CI] Compile with pytorch 2.4.0.dev20240514 --- .github/workflows/publish.yml | 2 +- flash_attn/__init__.py | 2 +- training/Dockerfile | 4 ++-- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index d3382a312..4e423aee1 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -44,7 +44,7 @@ jobs: # manylinux docker image, but I haven't figured out how to install CUDA on manylinux. os: [ubuntu-20.04] python-version: ['3.8', '3.9', '3.10', '3.11', '3.12'] - torch-version: ['2.0.1', '2.1.2', '2.2.2', '2.3.1', '2.4.0.dev20240512'] + torch-version: ['2.0.1', '2.1.2', '2.2.2', '2.3.1', '2.4.0.dev20240514'] cuda-version: ['11.8.0', '12.2.2'] # We need separate wheels that either uses C++11 ABI (-D_GLIBCXX_USE_CXX11_ABI) or not. # Pytorch wheels currently don't use it, but nvcr images have Pytorch compiled with C++11 ABI. diff --git a/flash_attn/__init__.py b/flash_attn/__init__.py index 3e0c49833..c0b387e54 100644 --- a/flash_attn/__init__.py +++ b/flash_attn/__init__.py @@ -1,4 +1,4 @@ -__version__ = "2.6.0" +__version__ = "2.6.0.post1" from flash_attn.flash_attn_interface import ( flash_attn_func, diff --git a/training/Dockerfile b/training/Dockerfile index b52ba1f61..da27cfebe 100644 --- a/training/Dockerfile +++ b/training/Dockerfile @@ -85,7 +85,7 @@ RUN pip install transformers==4.25.1 datasets==2.8.0 pytorch-lightning==1.8.6 tr RUN pip install git+https://github.com/mlcommons/logging.git@2.1.0 # Install FlashAttention -RUN pip install flash-attn==2.6.0 +RUN pip install flash-attn==2.6.0.post1 # Install CUDA extensions for fused dense -RUN pip install git+https://github.com/HazyResearch/flash-attention@v2.6.0#subdirectory=csrc/fused_dense_lib +RUN pip install git+https://github.com/HazyResearch/flash-attention@v2.6.0.post1#subdirectory=csrc/fused_dense_lib