tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

fa3 doesn't support dropout_p, fix unpatching

Browse Source
This commit is contained in:
Wing Lian
2025-05-18 06:26:08 -07:00
parent a064f1c9b4
commit 8c4bc59bfc
2 changed files with 19 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
src/axolotl/utils/models.py
View File

@@ -642,11 +642,19 @@ class ModelLoader:
flash_attn_varlen_func as flash_attn_varlen_func_v3,
)
def flash_attn_func_v3_wrapper(*args, **kwargs):
kwargs.pop("dropout_p", None)
return flash_attn_func_v3(*args, **kwargs)
def flash_attn_varlen_func_v3_wrapper(*args, **kwargs):
kwargs.pop("dropout_p", None)
return flash_attn_varlen_func_v3(*args, **kwargs)
transformers.modeling_flash_attention_utils.flash_attn_func = (
flash_attn_func_v3
flash_attn_func_v3_wrapper
)
transformers.modeling_flash_attention_utils.flash_attn_varlen_func = (
flash_attn_varlen_func_v3
flash_attn_varlen_func_v3_wrapper
)
LOG.info("Switched to Flash Attention v3")