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

[v1] add models & accelerator (#9579)

Browse Source
This commit is contained in:
Yaowei Zheng
2025-12-08 02:30:25 +08:00
committed by GitHub
parent 739954910a
commit 5744f1ea94
27 changed files with 335 additions and 105 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/llamafactory/v1/plugins/model_plugins/kernels/mlp/npu_fused_moe.py
View File

@@ -18,9 +18,9 @@ import torch
import torch.nn.functional as F
import torch_npu
from .....accelerator.helper import is_torch_npu_available
from .....extras.packages import is_transformers_version_greater_than
from .....extras.types import HFModel
from ....trainer_plugins.distributed.accelerate import is_torch_npu_available
from ..constants import DeviceType, KernelType
from ..registry import MetaMoEKernel

2
src/llamafactory/v1/plugins/model_plugins/kernels/mlp/npu_swiglu.py
View File

@@ -17,8 +17,8 @@ import types
import torch
from .....accelerator.helper import is_torch_npu_available
from .....extras.types import HFModel
from ....trainer_plugins.distributed.accelerate import is_torch_npu_available
from ..constants import DeviceType, KernelType
from ..registry import MetaSwiGluKernel

8
src/llamafactory/v1/plugins/model_plugins/kernels/registry.py
View File

@@ -15,8 +15,8 @@
from abc import ABC, ABCMeta, abstractmethod
from typing import Any, Callable, Optional
from ....accelerator.helper import get_current_accelerator
from ....extras.types import HFModel
from ...trainer_plugins.distributed.accelerate import get_available_accelerator
from .constants import DeviceType, KernelType
@@ -206,7 +206,7 @@ def discover_kernels(model: HFModel = None) -> list[type[MetaKernel]]:
discovered_kernels: list[type[MetaKernel]] = []
# Detect current device type
accelerator = get_available_accelerator()
accelerator = get_current_accelerator()
try:
device_type = DeviceType(accelerator.type)
except ValueError:
@@ -238,11 +238,11 @@ def apply_kernel(model: HFModel, kernel: type[MetaKernel], /, **kwargs) -> "HFMo
model = AutoModelForCausalLM.from_pretrained("qwen/qwen2.5-0.5B")
model = apply_kernel(model, NpuRMSNormKernel)
"""
if issubclass(kernel, MetaKernel) and kernel.device == get_available_accelerator().type:
if issubclass(kernel, MetaKernel) and kernel.device == get_current_accelerator().type:
return kernel.apply(model, **kwargs)
raise ValueError(
f"{kernel} must be a MetaKernel instance, or the kernel don't match the device type. got {kernel.device} and {get_available_accelerator().type} instead."
f"{kernel} must be a MetaKernel instance, or the kernel don't match the device type. got {kernel.device} and {get_current_accelerator().type} instead."
)

2
src/llamafactory/v1/plugins/model_plugins/kernels/rms_norm/npu_rms_norm.py
View File

@@ -14,8 +14,8 @@
import re
import types
from .....accelerator.helper import is_torch_npu_available
from .....extras.types import HFModel
from ....trainer_plugins.distributed.accelerate import is_torch_npu_available
from ..constants import DeviceType, KernelType
from ..registry import MetaRMSNormKernel

2
src/llamafactory/v1/plugins/model_plugins/kernels/rope/npu_rope.py
View File

@@ -16,8 +16,8 @@ import sys
import torch
from .....accelerator.helper import is_torch_npu_available
from .....extras.types import HFModel
from ....trainer_plugins.distributed.accelerate import is_torch_npu_available
from ..constants import DeviceType, KernelType
from ..registry import MetaRoPEKernel