423A35C7/LLaMA-Factory
1
0
Fork 0
mirror of https://github.com/hiyouga/LLaMA-Factory.git synced 2026-08-17 05:25:44 +08:00
Code Issues Packages Projects Releases Wiki Activity

[train] add fa3 supported (#10742)

Browse Source
This commit is contained in:
浮梦
2026-08-13 20:39:01 +08:00
committed by GitHub
parent 0bbe481e6e
commit 199b8873d7
1 changed files with 10 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -82,6 +82,14 @@ def configure_attn_implementation(config: "PretrainedConfig", model_args: "Model
return
requested_attn_implementation = "flash_attention_2"
elif model_args.flash_attn == AttentionFunction.FA3:
from transformers.utils import is_flash_attn_3_available
if not is_flash_attn_3_available():
logger.warning_rank0("FlashAttention-3 is not installed.")
return
requested_attn_implementation = "flash_attention_3"
else:
raise NotImplementedError(f"Unknown attention type: {model_args.flash_attn}")
@@ -109,6 +117,8 @@ def print_attn_implementation(config: "PretrainedConfig") -> None:
if attn_implementation == "flash_attention_2":
logger.info_rank0("Using FlashAttention-2 for faster training and inference.")
elif attn_implementation == "flash_attention_3":
logger.info_rank0("Using FlashAttention-3 for faster training and inference.")
elif attn_implementation == "sdpa":
logger.info_rank0("Using torch SDPA for faster training and inference.")
else: