423A35C7/LLaMA-Factory
1
0
Fork 0
mirror of https://github.com/hiyouga/LLaMA-Factory.git synced 2025-12-15 03:10:35 +08:00
Code Issues Packages Projects Releases Wiki Activity

clean code

Browse Source 
Former-commit-id: 2ed8270112
This commit is contained in:
hiyouga
2024-06-13 01:58:16 +08:00
parent 7366647b43
commit 833aa324c2
4 changed files with 17 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
src/llamafactory/model/model_utils/attention.py
View File

@@ -1,7 +1,8 @@
from typing import TYPE_CHECKING
from transformers.utils import is_flash_attn_2_available, is_torch_sdpa_available
from ...extras.logging import get_logger
from ...extras.packages import is_flash_attn2_available, is_sdpa_available
if TYPE_CHECKING:
@@ -21,13 +22,13 @@ def configure_attn_implementation(config: "PretrainedConfig", model_args: "Model
requested_attn_implementation = "eager"
elif model_args.flash_attn == "sdpa":
if not is_sdpa_available():
if not is_torch_sdpa_available():
logger.warning("torch>=2.1.1 is required for SDPA attention.")
return
requested_attn_implementation = "sdpa"
elif model_args.flash_attn == "fa2":
if not is_flash_attn2_available():
if not is_flash_attn_2_available():
logger.warning("FlashAttention-2 is not installed.")
return