423A35C7/LLaMA-Factory
1
0
Fork 0
mirror of https://github.com/hiyouga/LLaMA-Factory.git synced 2025-12-24 15:50:35 +08: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: 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 from transformers.utils import is_flash_attn_2_available
if getattr(config, "model_type", None) == "gemma2": if getattr(config, "model_type", None) == "gemma2":

3
src/llamafactory/model/model_utils/liger_kernel.py
View File

@@ -77,6 +77,9 @@ def apply_liger_kernel(
from liger_kernel.transformers import apply_liger_kernel_to_qwen3 as apply_liger_kernel from liger_kernel.transformers import apply_liger_kernel_to_qwen3 as apply_liger_kernel
elif model_type == "qwen3_moe": elif model_type == "qwen3_moe":
from liger_kernel.transformers import apply_liger_kernel_to_qwen3_moe as apply_liger_kernel from liger_kernel.transformers import apply_liger_kernel_to_qwen3_moe as apply_liger_kernel
elif model_type == "gpt_oss":
# Install manually from https://github.com/Comet0322/Liger-Kernel
from liger_kernel.transformers import apply_liger_kernel_to_gpt_oss as apply_liger_kernel
else: else:
logger.warning_rank0("Current model does not support liger kernel.") logger.warning_rank0("Current model does not support liger kernel.")
return return

4
src/llamafactory/model/model_utils/moe.py
View File

@@ -132,6 +132,10 @@ def add_z3_leaf_module(model: "PreTrainedModel") -> None:
_set_z3_leaf_modules(model, [Qwen3OmniMoeThinkerTextSparseMoeBlock]) _set_z3_leaf_modules(model, [Qwen3OmniMoeThinkerTextSparseMoeBlock])
if model_type == "gpt_oss":
from transformers.models.gpt_oss.modeling_gpt_oss import GptOssMLP
_set_z3_leaf_modules(model, [GptOssMLP])
def configure_moe(config: "PretrainedConfig", model_args: "ModelArguments", is_trainable: bool) -> None: def configure_moe(config: "PretrainedConfig", model_args: "ModelArguments", is_trainable: bool) -> None:
if not is_trainable or not model_args.moe_aux_loss_coef: if not is_trainable or not model_args.moe_aux_loss_coef: