squall/flash-attention
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Don't compile for Pytorch 2.1 on CUDA 12.1 due to nvcc segfaults

Browse Source
This commit is contained in:
Tri Dao 2023-09-17 22:15:38 -07:00
parent c984208ddb
commit 799f56fa90
4 changed files with 10 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
.github/workflows/publish.yml vendored
View File

@ -80,6 +80,11 @@ jobs:
cuda-version: '11.7.1' cuda-version: '11.7.1'
- torch-version: '2.1.0.dev20230731' - torch-version: '2.1.0.dev20230731'
cuda-version: '11.8.0' cuda-version: '11.8.0'
# Pytorch >= 2.1 with nvcc 12.1.0 segfaults during compilation, so
# we only use CUDA 12.2. setup.py as a special case that will
# download the wheel for CUDA 12.2 instead.
- torch-version: '2.1.0.dev20230731'
cuda-version: '12.1.0'
steps: steps:
- name: Checkout - name: Checkout

2
flash_attn/__init__.py
View File

@ -1,4 +1,4 @@
__version__ = "2.2.3.post1" __version__ = "2.2.3.post2"
from flash_attn.flash_attn_interface import ( from flash_attn.flash_attn_interface import (
flash_attn_func, flash_attn_func,

2
setup.py
View File

@ -223,6 +223,8 @@ def get_wheel_url():
# _, cuda_version_raw = get_cuda_bare_metal_version(CUDA_HOME) # _, cuda_version_raw = get_cuda_bare_metal_version(CUDA_HOME)
torch_cuda_version = parse(torch.version.cuda) torch_cuda_version = parse(torch.version.cuda)
torch_version_raw = parse(torch.__version__) torch_version_raw = parse(torch.__version__)
if torch_version_raw.major == 2 and torch_version_raw.minor == 1:
torch_cuda_version = parse("12.2")
python_version = f"cp{sys.version_info.major}{sys.version_info.minor}" python_version = f"cp{sys.version_info.major}{sys.version_info.minor}"
platform_name = get_platform() platform_name = get_platform()
flash_version = get_package_version() flash_version = get_package_version()

7
training/Dockerfile
View File

@ -85,14 +85,11 @@ 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 RUN pip install git+https://github.com/mlcommons/logging.git@2.1.0
# Install FlashAttention # Install FlashAttention
RUN pip install flash-attn==2.2.3.post1 RUN pip install flash-attn==2.2.3.post2
# Install CUDA extensions for cross-entropy, fused dense, layer norm # Install CUDA extensions for cross-entropy, fused dense, layer norm
RUN git clone https://github.com/HazyResearch/flash-attention \ RUN git clone https://github.com/HazyResearch/flash-attention \
&& cd flash-attention && git checkout v2.2.3.post1 \ && cd flash-attention && git checkout v2.2.3.post2 \
&& cd csrc/fused_softmax && pip install . && cd ../../ \
&& cd csrc/rotary && pip install . && cd ../../ \
&& cd csrc/layer_norm && pip install . && cd ../../ \ && cd csrc/layer_norm && pip install . && cd ../../ \
&& cd csrc/fused_dense_lib && pip install . && cd ../../ \ && cd csrc/fused_dense_lib && pip install . && cd ../../ \
&& cd csrc/ft_attention && pip install . && cd ../../ \
&& cd .. && rm -rf flash-attention && cd .. && rm -rf flash-attention