423A35C7/LLaMA-Factory
1
0
Fork 0
mirror of https://github.com/hiyouga/LLaMA-Factory.git synced 2025-11-05 02:12:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

[misc] fix import error (#9296)

Browse Source
This commit is contained in:
Yaowei Zheng 2025-10-17 10:54:30 +08:00 committed by GitHub
parent 8c341cbaae
commit a442fa90ad
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 4 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -14,8 +14,6 @@
from typing import TYPE_CHECKING from typing import TYPE_CHECKING
from transformers.utils import is_flash_attn_2_available, is_torch_sdpa_available
from ...extras import logging from ...extras import logging
from ...extras.constants import AttentionFunction from ...extras.constants import AttentionFunction
@ -30,6 +28,8 @@ logger = logging.get_logger(__name__)
def configure_attn_implementation(config: "PretrainedConfig", model_args: "ModelArguments") -> None: def configure_attn_implementation(config: "PretrainedConfig", model_args: "ModelArguments") -> None:
from transformers.utils import is_flash_attn_2_available
if getattr(config, "model_type", None) == "gemma2": if getattr(config, "model_type", None) == "gemma2":
if model_args.flash_attn == AttentionFunction.AUTO or model_args.flash_attn == AttentionFunction.FA2: if model_args.flash_attn == AttentionFunction.AUTO or model_args.flash_attn == AttentionFunction.FA2:
if is_flash_attn_2_available(): if is_flash_attn_2_available():
@ -51,6 +51,8 @@ def configure_attn_implementation(config: "PretrainedConfig", model_args: "Model
requested_attn_implementation = "eager" requested_attn_implementation = "eager"
elif model_args.flash_attn == AttentionFunction.SDPA: elif model_args.flash_attn == AttentionFunction.SDPA:
from transformers.utils import is_torch_sdpa_available
if not is_torch_sdpa_available(): if not is_torch_sdpa_available():
logger.warning_rank0("torch>=2.1.1 is required for SDPA attention.") logger.warning_rank0("torch>=2.1.1 is required for SDPA attention.")
return return