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

handle return sig change for fa3

Browse Source
This commit is contained in:
Wing Lian
2025-05-18 08:28:52 -07:00
parent b22150751f
commit 323a9cb153
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -644,11 +644,11 @@ class ModelLoader:
def flash_attn_func_v3_wrapper(*args, **kwargs): def flash_attn_func_v3_wrapper(*args, **kwargs):
kwargs.pop("dropout_p", None) kwargs.pop("dropout_p", None)
return flash_attn_func_v3(*args, **kwargs) return flash_attn_func_v3(*args, **kwargs)[0]
def flash_attn_varlen_func_v3_wrapper(*args, **kwargs): def flash_attn_varlen_func_v3_wrapper(*args, **kwargs):
kwargs.pop("dropout_p", None) kwargs.pop("dropout_p", None)
return flash_attn_varlen_func_v3(*args, **kwargs) return flash_attn_varlen_func_v3(*args, **kwargs)[0]
transformers.modeling_flash_attention_utils.flash_attn_func = ( transformers.modeling_flash_attention_utils.flash_attn_func = (
flash_attn_func_v3_wrapper flash_attn_func_v3_wrapper