update flash attention for gemma support: (#1368)

This commit is contained in:
Wing Lian
2024-03-06 10:08:54 -05:00
committed by GitHub
parent ed70a08348
commit 58b0d4b0d8
2 changed files with 2 additions and 2 deletions

View File

@@ -12,7 +12,7 @@ fire
PyYAML>=6.0 PyYAML>=6.0
requests requests
datasets>=2.15.0 datasets>=2.15.0
flash-attn==2.3.3 flash-attn==2.5.5
sentencepiece sentencepiece
wandb wandb
einops einops

View File

@@ -68,7 +68,7 @@ setup(
dependency_links=dependency_links, dependency_links=dependency_links,
extras_require={ extras_require={
"flash-attn": [ "flash-attn": [
"flash-attn==2.5.0", "flash-attn==2.5.5",
], ],
"fused-dense-lib": [ "fused-dense-lib": [
"fused-dense-lib @ git+https://github.com/Dao-AILab/flash-attention@v2.3.3#subdirectory=csrc/fused_dense_lib", "fused-dense-lib @ git+https://github.com/Dao-AILab/flash-attention@v2.3.3#subdirectory=csrc/fused_dense_lib",