ros/LlamaFactory
1
0
Fork 0
mirror of https://github.com/hiyouga/LlamaFactory.git synced 2026-02-01 20:23:37 +00:00
Code Issues Packages Projects Releases Wiki Activity

[model] efficient tuning for gpt-oss (#9354)

Browse Source 
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
thulyubh22
2025-12-23 16:28:38 +08:00
committed by GitHub
parent 1f1f5a7d1b
commit 7901b2f32e
3 changed files with 18 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -29,6 +29,15 @@ logger = logging.get_logger(__name__)
def configure_attn_implementation(config: "PretrainedConfig", model_args: "ModelArguments") -> None:
if getattr(config, "model_type", None) == "gpt_oss":
from transformers.integrations.hub_kernels import load_and_register_kernel
flash_attn3_kernel = "kernels-community/vllm-flash-attn3"
load_and_register_kernel(flash_attn3_kernel)
setattr(config, "_attn_implementation", flash_attn3_kernel)
setattr(config, "_attn_implementation_internal", flash_attn3_kernel)
model_args.flash_attn = flash_attn3_kernel
return
from transformers.utils import is_flash_attn_2_available
if getattr(config, "model_type", None) == "gemma2":