update flash attention for gemma support: (#1368)

2024-03-06 10:08:54 -05:00
parent ed70a08348
commit 58b0d4b0d8
2 changed files with 2 additions and 2 deletions
--- 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
--- 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",