diff --git a/requirements.txt b/requirements.txt index cd5171ebd..718896783 100644 --- a/requirements.txt +++ b/requirements.txt @@ -12,7 +12,7 @@ fire PyYAML>=6.0 requests datasets>=2.15.0 -flash-attn==2.3.3 +flash-attn==2.5.5 sentencepiece wandb einops diff --git a/setup.py b/setup.py index 4a949a60d..40dd0a668 100644 --- a/setup.py +++ b/setup.py @@ -68,7 +68,7 @@ setup( dependency_links=dependency_links, extras_require={ "flash-attn": [ - "flash-attn==2.5.0", + "flash-attn==2.5.5", ], "fused-dense-lib": [ "fused-dense-lib @ git+https://github.com/Dao-AILab/flash-attention@v2.3.3#subdirectory=csrc/fused_dense_lib",