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

check for fa first

Browse Source
This commit is contained in:
Wing Lian
2025-05-18 07:04:48 -07:00
parent 8c4bc59bfc
commit b22150751f
1 changed files with 18 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
tests/conftest.py
View File

@@ -435,6 +435,15 @@ def cleanup_monkeypatches():
Trainer._inner_training_loop # pylint: disable=protected-access Trainer._inner_training_loop # pylint: disable=protected-access
) )
original_trainer_training_step = Trainer.training_step original_trainer_training_step = Trainer.training_step
original_fa_func = None
original_fa_varlen_func = None
if (
importlib.util.find_spec("flash_attn")
and hasattr(transformers.modeling_flash_attention_utils, "flash_attn_func")
and hasattr(
transformers.modeling_flash_attention_utils, "flash_attn_varlen_func"
)
):
original_fa_func = transformers.modeling_flash_attention_utils.flash_attn_func original_fa_func = transformers.modeling_flash_attention_utils.flash_attn_func
original_fa_varlen_func = ( original_fa_varlen_func = (
transformers.modeling_flash_attention_utils.flash_attn_varlen_func transformers.modeling_flash_attention_utils.flash_attn_varlen_func
@@ -449,6 +458,7 @@ def cleanup_monkeypatches():
original_trainer_inner_training_loop original_trainer_inner_training_loop
) )
Trainer.training_step = original_trainer_training_step Trainer.training_step = original_trainer_training_step
if original_fa_func:
transformers.modeling_flash_attention_utils.flash_attn_func = original_fa_func transformers.modeling_flash_attention_utils.flash_attn_func = original_fa_func
transformers.modeling_flash_attention_utils.flash_attn_varlen_func = ( transformers.modeling_flash_attention_utils.flash_attn_varlen_func = (
original_fa_varlen_func original_fa_varlen_func