update flash attention for gemma support: (axolotl-ai-cloud#1368)

maximegmd · Mar 6, 2024 · 58b0d4b · 58b0d4b
1 parent ed70a08
commit 58b0d4b
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git 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
@@ -68,7 +68,7 @@ def parse_requirements():
     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/[email protected]#subdirectory=csrc/fused_dense_lib",