add FlashAttention-2 support #38

iMountTai · 2023-08-01T08:42:36Z

This PR add support for FlashAttention-2 and enable FlashAttention-2 training by setting the argument --flash_attn

Usage
Just add the --flash_attn to the launching command. For example:

torchrun --nnodes 1 --nproc_per_node 1 run_clm_pt_with_peft.py \
    --other_arguments
    --flash_attn

scripts/training/run_clm_pt_with_peft.py

scripts/training/run_clm_sft_with_peft.py

ymcui · 2023-08-01T09:18:52Z

TODO: Update wiki w.r.t. this PR. @iMountTai

iMountTai and others added 5 commits August 1, 2023 15:56

add FlashAttention-2 patch

f4eda56

modify some suggestions

da6d390

Merge branch 'ymcui:main' into main

5fd2eb5

add flash-attn option

8eb322b

modify some log information

746ece0

iMountTai requested a review from airaria August 1, 2023 08:58

airaria reviewed Aug 1, 2023

View reviewed changes

scripts/training/run_clm_pt_with_peft.py Show resolved Hide resolved

scripts/training/run_clm_sft_with_peft.py Show resolved Hide resolved

airaria approved these changes Aug 1, 2023

View reviewed changes

Stylistic fixes based on Codacy suggestions

1649777

ymcui merged commit 8ef2788 into ymcui:main Aug 1, 2023
1 check passed

Provide feedback