diff --git a/src/llamafactory/model/loader.py b/src/llamafactory/model/loader.py index 2c64c504b..01648e768 100644 --- a/src/llamafactory/model/loader.py +++ b/src/llamafactory/model/loader.py @@ -217,9 +217,9 @@ def load_model( "You are try to using future feature about kernels, please note that this feature " "is not supported for all models. If get any error, please disable this feature, or report the issue." ) - from ..v1.plugins.model_plugins.kernels.interface import apply_default_kernels + from ..v1.plugins.model_plugins.kernels.interface import apply_v1_kernels - model = apply_default_kernels(model, include_kernels=model_args.use_v1_kernels) + model = apply_v1_kernels(model, use_v1_kernels=model_args.use_v1_kernels) trainable_params, all_param = count_parameters(model) if is_trainable: diff --git a/src/llamafactory/v1/plugins/model_plugins/kernels/ops/mlp/npu_fused_moe.py b/src/llamafactory/v1/plugins/model_plugins/kernels/ops/mlp/npu_fused_moe.py index 20c86a2dc..756ffe84b 100644 --- a/src/llamafactory/v1/plugins/model_plugins/kernels/ops/mlp/npu_fused_moe.py +++ b/src/llamafactory/v1/plugins/model_plugins/kernels/ops/mlp/npu_fused_moe.py @@ -29,15 +29,21 @@ import torch.nn.functional as F try: import torch_npu -except ImportError: - pass +except ImportError as exc: + _TORCH_NPU_IMPORT_ERROR = exc +else: + _TORCH_NPU_IMPORT_ERROR = None from ......accelerator.helper import DeviceType, get_current_accelerator +from ......utils.logging import get_logger from ......utils.packages import is_transformers_version_greater_than from ......utils.types import HFModel from ...base import BaseKernel, KernelPlugin +logger = get_logger(__name__) + + class GmmFunction(torch.autograd.Function): """Custom autograd function for NPU Grouped Matrix Multiplication (GMM).""" @@ -49,7 +55,7 @@ class GmmFunction(torch.autograd.Function): ctx: Context object to save tensors for backward pass. x (Tensor): Input tensor. weight (Tensor): Weight tensor. - group_list (list): List of group sizes. + group_list (Tensor): Number of tokens assigned to each expert. Returns: Tensor: The result of the grouped matrix multiplication. @@ -174,14 +180,14 @@ class HybridGmmFunction(torch.autograd.Function): return (None, *grad_x_list, *grad_w_list) -class NpuMoeFused: - """Container for NPU fused MoE forward functions.""" +class NpuMoeFusedV4: + """Container for Transformers v4 NPU fused MoE forward functions.""" @staticmethod - def npu_moe_experts_forward( + def stacked_experts_forward( self, hidden_states: torch.Tensor, routing_weights: torch.Tensor, router_indices: torch.Tensor ) -> torch.Tensor: - """Forward pass for MoE experts using NPU fused operations. + """Forward pass for Transformers v4 MoE experts using NPU fused operations. Args: self: The MoE layer instance. @@ -197,7 +203,9 @@ class NpuMoeFused: permuted_hidden_states, row_ids_map = torch_npu.npu_moe_token_permute( hidden_states, router_indices.to(torch.int32) ) - tokens_per_expert = torch.histc(router_indices, bins=self.num_experts, min=0, max=self.num_experts) + tokens_per_expert = torch.histc( + router_indices.float(), bins=self.num_experts, min=0, max=self.num_experts + ).long() intermediate_hidden_states = GmmFunction.apply(permuted_hidden_states, self.gate_up_proj, tokens_per_expert) intermediate_activations = torch_npu.npu_swiglu(intermediate_hidden_states, dim=-1) output = GmmFunction.apply(intermediate_activations, self.down_proj, tokens_per_expert) @@ -206,61 +214,33 @@ class NpuMoeFused: return next_states @staticmethod - def npu_moe_sparse_block_forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - r"""Forward pass for sparse MoE block using NPU optimization. + def stacked_sparse_block_forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + r"""Forward pass for Transformers v4 sparse MoE block using NPU optimization. Args: self: The MoE sparse block instance. hidden_states (Tensor): Input hidden states. Returns: - Tensor: The routed output. + tuple: A tuple containing the routed output and router logits. """ batch_size = hidden_states.shape[0] hidden_states = hidden_states.reshape(-1, self.hidden_size) router_logits = self.gate(hidden_states) - routing_weights = torch.nn.functional.softmax(router_logits, dim=-1, dtype=torch.float) + routing_weights = F.softmax(router_logits, dim=-1, dtype=torch.float) routing_weights, router_indices = torch.topk(routing_weights, self.top_k, dim=-1) routing_weights = routing_weights / routing_weights.sum(dim=-1, keepdim=True) routing_weights = routing_weights.to(hidden_states.dtype) hidden_states = hidden_states.reshape(batch_size, -1, self.hidden_size) routed_out = self.experts(hidden_states, routing_weights, router_indices) - return routed_out + return routed_out, router_logits @staticmethod - def npu_moe_experts_v5_forward( - self, hidden_states: torch.Tensor, top_k_index: torch.Tensor, top_k_weights: torch.Tensor - ) -> torch.Tensor: - """Forward pass for Transformers v5+ MoE experts using NPU fused operations. - - Transformers v5 stores expert weights in F.linear layout: - gate_up_proj: [num_experts, 2 * intermediate_dim, hidden_dim] - down_proj: [num_experts, hidden_dim, intermediate_dim] - The NPU grouped matmul path expects matmul layout, so both weights are transposed. - """ - hidden_states = hidden_states.reshape(-1, self.hidden_dim) - permuted_hidden_states, row_ids_map = torch_npu.npu_moe_token_permute( - hidden_states, top_k_index.to(torch.int32) - ) - tokens_per_expert = torch.histc(top_k_index.float(), bins=self.num_experts, min=0, max=self.num_experts).long() - - gate_up_proj = self.gate_up_proj.transpose(1, 2) - down_proj = self.down_proj.transpose(1, 2) - intermediate_hidden_states = GmmFunction.apply(permuted_hidden_states, gate_up_proj, tokens_per_expert) - intermediate_activations = torch_npu.npu_swiglu(intermediate_hidden_states, dim=-1) - output = GmmFunction.apply(intermediate_activations, down_proj, tokens_per_expert) - return torch_npu.npu_moe_token_unpermute(output, row_ids_map, probs=top_k_weights) - - -class Qwen3NpuMoeFused: - """Container for Qwen3 NPU fused MoE forward functions.""" - - @staticmethod - def qwen3moe_sparse_moe_block_forward(self, hidden_states: torch.Tensor): - """Forward pass for Qwen3 sparse MoE block using NPU fused operations. + def sparse_block_forward(self, hidden_states: torch.Tensor): + """Forward pass for a Transformers v4 list-backed sparse MoE block using NPU fused operations. Args: - self: The Qwen3 MoE block instance. + self: The sparse MoE block instance. hidden_states (Tensor): Input hidden states. Returns: @@ -304,33 +284,90 @@ class Qwen3NpuMoeFused: next_states = next_states.view(batch_size, sequence_length, -1) return next_states, router_logits + @staticmethod + def shared_sparse_block_forward(self, hidden_states: torch.Tensor): + """Forward pass for a Transformers v4 sparse MoE block with a shared expert.""" + next_states, router_logits = NpuMoeFusedV4.sparse_block_forward(self, hidden_states) -# moe patch config mapping -if is_transformers_version_greater_than("5.0.0"): - kernel_moe_mapping = { - "Qwen3MoeForCausalLM": { - "Qwen3MoeExperts": NpuMoeFused.npu_moe_experts_v5_forward, - }, - "Qwen3VLMoeForConditionalGeneration": { - "Qwen3VLMoeTextExperts": NpuMoeFused.npu_moe_experts_v5_forward, - }, - "Qwen3_5MoeForCausalLM": { - "Qwen3_5MoeExperts": NpuMoeFused.npu_moe_experts_v5_forward, - }, - "Qwen3_5MoeForConditionalGeneration": { - "Qwen3_5MoeExperts": NpuMoeFused.npu_moe_experts_v5_forward, - }, - } -else: - kernel_moe_mapping = { - "Qwen3MoeForCausalLM": { - "Qwen3MoeSparseMoeBlock": Qwen3NpuMoeFused.qwen3moe_sparse_moe_block_forward, - }, - "Qwen3VLMoeForConditionalGeneration": { - "Qwen3VLMoeTextExperts": NpuMoeFused.npu_moe_experts_forward, - "Qwen3VLMoeTextSparseMoeBlock": NpuMoeFused.npu_moe_sparse_block_forward, - }, - } + shared_expert_output = self.shared_expert(hidden_states) + shared_expert_output = F.sigmoid(self.shared_expert_gate(hidden_states)) * shared_expert_output + next_states = next_states + shared_expert_output + return next_states, router_logits + + +class NpuMoeFusedV5: + """Container for Transformers v5 NPU fused MoE forward functions.""" + + @staticmethod + def experts_forward( + self, hidden_states: torch.Tensor, top_k_index: torch.Tensor, top_k_weights: torch.Tensor + ) -> torch.Tensor: + """Forward pass for Transformers v5+ MoE experts using NPU fused operations. + + Transformers v5 stores expert weights in F.linear layout: + gate_up_proj: [num_experts, 2 * intermediate_dim, hidden_dim] + down_proj: [num_experts, hidden_dim, intermediate_dim] + The NPU grouped matmul path expects matmul layout, so both weights are transposed. + """ + hidden_states = hidden_states.reshape(-1, self.hidden_dim) + permuted_hidden_states, row_ids_map = torch_npu.npu_moe_token_permute( + hidden_states, top_k_index.to(torch.int32) + ) + tokens_per_expert = torch.histc(top_k_index.float(), bins=self.num_experts, min=0, max=self.num_experts).long() + + gate_up_proj = self.gate_up_proj.transpose(1, 2) + down_proj = self.down_proj.transpose(1, 2) + intermediate_hidden_states = GmmFunction.apply(permuted_hidden_states, gate_up_proj, tokens_per_expert) + intermediate_activations = torch_npu.npu_swiglu(intermediate_hidden_states, dim=-1) + output = GmmFunction.apply(intermediate_activations, down_proj, tokens_per_expert) + return torch_npu.npu_moe_token_unpermute(output, row_ids_map, probs=top_k_weights) + + +_V4_MODEL_TYPE_TO_PATCHES = { + "qwen3_moe": { + "Qwen3MoeSparseMoeBlock": NpuMoeFusedV4.sparse_block_forward, + }, + "qwen3_next": { + "Qwen3NextSparseMoeBlock": NpuMoeFusedV4.shared_sparse_block_forward, + }, + "qwen3_omni_moe": { + "Qwen3OmniMoeThinkerTextSparseMoeBlock": NpuMoeFusedV4.sparse_block_forward, + "Qwen3OmniMoeTalkerTextSparseMoeBlock": NpuMoeFusedV4.shared_sparse_block_forward, + }, + "qwen3_omni_moe_thinker": { + "Qwen3OmniMoeThinkerTextSparseMoeBlock": NpuMoeFusedV4.sparse_block_forward, + }, + "qwen3_vl_moe": { + "Qwen3VLMoeTextExperts": NpuMoeFusedV4.stacked_experts_forward, + "Qwen3VLMoeTextSparseMoeBlock": NpuMoeFusedV4.stacked_sparse_block_forward, + }, +} + +_V5_MODEL_TYPE_TO_PATCHES = { + "qwen3_moe": { + "Qwen3MoeExperts": NpuMoeFusedV5.experts_forward, + }, + "qwen3_next": { + "Qwen3NextExperts": NpuMoeFusedV5.experts_forward, + }, + "qwen3_omni_moe": { + "Qwen3OmniMoeThinkerTextExperts": NpuMoeFusedV5.experts_forward, + "Qwen3OmniMoeTalkerTextExperts": NpuMoeFusedV5.experts_forward, + }, + "qwen3_omni_moe_thinker": { + "Qwen3OmniMoeThinkerTextExperts": NpuMoeFusedV5.experts_forward, + }, + "qwen3_vl_moe": { + "Qwen3VLMoeTextExperts": NpuMoeFusedV5.experts_forward, + }, + "qwen3_5_moe": { + "Qwen3_5MoeExperts": NpuMoeFusedV5.experts_forward, + }, +} + +_MODEL_TYPE_TO_PATCHES = ( + _V5_MODEL_TYPE_TO_PATCHES if is_transformers_version_greater_than("5.0.0") else _V4_MODEL_TYPE_TO_PATCHES +) @KernelPlugin("npu_fused_moe").register() @@ -343,6 +380,17 @@ class NpuFusedMoEKernel(BaseKernel): if current != DeviceType.NPU: raise RuntimeError(f"NpuFusedMoEKernel requires NPU, current accelerator is {current}.") + @staticmethod + def check_deps() -> None: + if _TORCH_NPU_IMPORT_ERROR is not None: + raise RuntimeError("NpuFusedMoEKernel requires torch_npu.") from _TORCH_NPU_IMPORT_ERROR + + @staticmethod + def _get_patch_forward(model_type: str, module: torch.nn.Module): + """Return the version-specific NPU forward function for a matched MoE module.""" + model_patches = _MODEL_TYPE_TO_PATCHES.get(model_type, {}) + return model_patches.get(module.__class__.__name__) + @staticmethod def _apply(**kwargs) -> HFModel: """Applies the NPU fused MoE kernel to the model. @@ -352,27 +400,21 @@ class NpuFusedMoEKernel(BaseKernel): Returns: HFModel: The model with patched MoE forward functions. - - Raises: - ValueError: If the model is not provided. - RuntimeError: If dependencies are not met. """ - model = kwargs.get("model", None) + model = kwargs["model"] - archs = getattr(model.config, "architectures", None) or [] - target_moe_mapping = None - for arch in archs: - if arch in kernel_moe_mapping: - target_moe_mapping = kernel_moe_mapping[arch] - break - - if target_moe_mapping is None: + model_type = getattr(model.config, "model_type", None) + if model_type not in _MODEL_TYPE_TO_PATCHES: return model + patched_count = 0 for module in model.modules(): - class_name = module.__class__.__name__ - if class_name in target_moe_mapping: - new_forward_func = target_moe_mapping[class_name] - module.forward = types.MethodType(new_forward_func, module) + patch_forward = NpuFusedMoEKernel._get_patch_forward(model_type, module) + if patch_forward is not None: + module.forward = types.MethodType(patch_forward, module) + patched_count += 1 + + if patched_count: + logger.info_rank0(f"Applied NPU fused MoE kernel to {patched_count} modules for model type: {model_type}.") return model diff --git a/src/llamafactory/v1/plugins/model_plugins/kernels/ops/mlp/npu_swiglu.py b/src/llamafactory/v1/plugins/model_plugins/kernels/ops/mlp/npu_swiglu.py index e48de755b..ae0ad7c96 100644 --- a/src/llamafactory/v1/plugins/model_plugins/kernels/ops/mlp/npu_swiglu.py +++ b/src/llamafactory/v1/plugins/model_plugins/kernels/ops/mlp/npu_swiglu.py @@ -20,20 +20,24 @@ Init Phase: """ -import re import types import torch from ......accelerator.helper import DeviceType, get_current_accelerator +from ......utils.logging import get_logger from ......utils.types import HFModel from ...base import BaseKernel, KernelPlugin +logger = get_logger(__name__) + try: import torch_npu -except ImportError: - pass +except ImportError as exc: + _TORCH_NPU_IMPORT_ERROR = exc +else: + _TORCH_NPU_IMPORT_ERROR = None def npu_swiglu_forward(self, hidden_state): @@ -51,81 +55,69 @@ def npu_swiglu_forward(self, hidden_state): ) -def _npu_swiglu_glm4_forward(self, hidden_states): - """SwiGLU forward pass for GLM4 on NPU. - - Args: - self: The GLM4 MLP layer instance. - hidden_states (Tensor): Input hidden states. - - Returns: - Tensor: Output of SwiGLU. - """ - up_states = self.gate_up_proj(hidden_states) - gate, up_states = up_states.chunk(2, dim=-1) - return self.down_proj(torch_npu.npu_swiglu(torch.cat((gate, up_states), dim=-1), dim=-1)) - - -def _npu_swiglu_gemma3ntext_forward(self, hidden_states): - """SwiGLU forward pass for Gemma3nText on NPU. - - Args: - self: The Gemma3nText MLP layer instance. - hidden_states (Tensor): Input hidden states. - - Returns: - Tensor: Output of SwiGLU. - """ - gate_proj = self.gate_proj(hidden_states) - if self.activation_sparsity > 0.0: - gate_proj = self._gaussian_topk(gate_proj) - down_proj = self.down_proj( - torch_npu.npu_swiglu(torch.cat((gate_proj, self.up_proj(hidden_states)), dim=-1), dim=-1) - ) - return down_proj +_MODEL_TYPE_TO_PATCHES = { + "qwen3": { + "Qwen3MLP": npu_swiglu_forward, + }, + "qwen3_moe": { + "Qwen3MoeMLP": npu_swiglu_forward, + }, + "qwen3_next": { + "Qwen3NextMLP": npu_swiglu_forward, + }, + "qwen3_omni_moe": { + "Qwen3OmniMoeThinkerTextMLP": npu_swiglu_forward, + "Qwen3OmniMoeMLP": npu_swiglu_forward, + "Qwen3OmniMoeTalkerTextMLP": npu_swiglu_forward, + "Qwen3OmniMoeCode2WavMlp": npu_swiglu_forward, + }, + "qwen3_omni_moe_thinker": { + "Qwen3OmniMoeThinkerTextMLP": npu_swiglu_forward, + }, + "qwen3_vl": { + "Qwen3VLTextMLP": npu_swiglu_forward, + }, + "qwen3_vl_moe": { + "Qwen3VLMoeTextMLP": npu_swiglu_forward, + }, + "qwen3_5": { + "Qwen3_5MLP": npu_swiglu_forward, + }, + "qwen3_5_moe": { + "Qwen3_5MoeMLP": npu_swiglu_forward, + }, +} @KernelPlugin("npu_fused_swiglu").register() class NpuSwiGluKernel(BaseKernel): """NPU Kernel for fused SwiGLU activation.""" - # just support apply to the following module layers - expect_modules = frozenset( - { - "Qwen3VLMoeTextMLP", - "Qwen3VLTextMLP", - "Qwen3OmniMoeThinkerTextMLP", - "Qwen3OmniMoeMLP", - "Qwen3OmniMoeTalkerTextMLP", - "Qwen3OmniMoeCode2WavMlp", - "Qwen3NextMLP", - "Qwen3MoeMLP", - "Qwen3MLP", - "Qwen2MLP", - "Qwen2MoeMLP", - "Qwen2_5_VLMLP", - "Qwen2_5OmniMLP", - "Llama4TextMLP", - "LlamaMLP", - "Glm4MLP", - "Glm4MoeMLP", - "Glm4vMoeTextMLP", - "Gemma3MLP", - "Gemma2MLP", - "Gemma3nTextMLP", - "Phi3MLP", - "DeepseekV2MLP", - "DeepseekV3MLP", - "SeedOssMLP", - } - ) - @staticmethod def check_device() -> None: current = get_current_accelerator().type if current != DeviceType.NPU: raise RuntimeError(f"NpuSwiGluKernel requires NPU, current accelerator is {current}.") + @staticmethod + def check_deps() -> None: + if _TORCH_NPU_IMPORT_ERROR is not None: + raise RuntimeError("NpuSwiGluKernel requires torch_npu.") from _TORCH_NPU_IMPORT_ERROR + + @staticmethod + def _get_patch_forward(model_type: str, module: torch.nn.Module): + """Return the NPU forward function for a matched SwiGLU MLP module.""" + model_patches = _MODEL_TYPE_TO_PATCHES.get(model_type, {}) + patch_forward = model_patches.get(module.__class__.__name__) + if patch_forward is None: + return None + + config = getattr(module, "config", None) + if getattr(config, "hidden_act", None) != "silu": + return None + + return patch_forward + @staticmethod def _apply(**kwargs) -> "HFModel": """Applies the NPU fused SwiGLU kernel to the model. @@ -135,31 +127,21 @@ class NpuSwiGluKernel(BaseKernel): Returns: HFModel: The model with patched SwiGLU forward functions. - - Raises: - ValueError: If the model is not provided. - RuntimeError: If dependencies are not met. """ - model = kwargs.get("model", None) + model = kwargs["model"] - # Mapping of specific mlp modules to their corresponding kernel implementations - kernel_mapping = { - "Glm4MLP": _npu_swiglu_glm4_forward, - "Glm4vTextMLP": _npu_swiglu_glm4_forward, - "Phi3MLP": _npu_swiglu_glm4_forward, - "Gemma3nTextMLP": _npu_swiglu_gemma3ntext_forward, - } + model_type = getattr(model.config, "model_type", None) + if model_type not in _MODEL_TYPE_TO_PATCHES: + return model - swiglu_pattern = re.compile("MLP", re.IGNORECASE) - for name, module in model.named_modules(): - # Match any module whose class name contains "MLP" - if ( - re.search(swiglu_pattern, module.__class__.__name__) - and module.__class__.__name__ in NpuSwiGluKernel.expect_modules - ): - # Bind function as an instance method to preserve `self` semantics - # and replace the original forward - kernel_func = kernel_mapping.get(module.__class__.__name__, npu_swiglu_forward) - module.forward = types.MethodType(kernel_func, module) + patched_count = 0 + for module in model.modules(): + patch_forward = NpuSwiGluKernel._get_patch_forward(model_type, module) + if patch_forward is not None: + module.forward = types.MethodType(patch_forward, module) + patched_count += 1 + + if patched_count: + logger.info_rank0(f"Applied NPU SwiGLU kernel to {patched_count} modules for model type: {model_type}.") return model diff --git a/src/llamafactory/v1/plugins/model_plugins/kernels/ops/rms_norm/npu_rms_norm.py b/src/llamafactory/v1/plugins/model_plugins/kernels/ops/rms_norm/npu_rms_norm.py index 3d06abc34..0e960b6c7 100644 --- a/src/llamafactory/v1/plugins/model_plugins/kernels/ops/rms_norm/npu_rms_norm.py +++ b/src/llamafactory/v1/plugins/model_plugins/kernels/ops/rms_norm/npu_rms_norm.py @@ -20,54 +20,32 @@ Init Phase: """ -import re import types import torch import torch.nn.functional as F from ......accelerator.helper import DeviceType, get_current_accelerator +from ......utils.logging import get_logger from ......utils.types import HFModel from ...base import BaseKernel, KernelPlugin +logger = get_logger(__name__) + try: import torch_npu -except ImportError: - pass - - -def _should_use_residual_rmsnorm(module): - """Detect whether the module uses residual RMSNorm parameterization. - - Residual RMSNorm uses ``scale = 1.0 + weight`` where weight is initialized to 0, - while standard RMSNorm uses ``scale = weight`` where weight is initialized to 1. - - Args: - module (nn.Module): The RMSNorm module to check. - - Returns: - bool: ``True`` if the module uses residual parameterization, ``False`` otherwise. - - .. note:: - This must follow the module's forward semantics. Do not infer it from trained - weight values because standard RMSNorm weights can also be close to zero. - """ - residual_rmsnorm_classes = { - "Qwen3_5RMSNorm", - "Qwen3_5MoeRMSNorm", - "Qwen3NextRMSNorm", - } - - class_name = module.__class__.__name__ - return class_name in residual_rmsnorm_classes +except ImportError as exc: + _TORCH_NPU_IMPORT_ERROR = exc +else: + _TORCH_NPU_IMPORT_ERROR = None def npu_rms_norm_forward(self, hidden_states): """NPU forward implementation for standard RMSNorm. Args: - self (nn.Module): The RMSNorm module instance with ``weight`` and ``variance_epsilon``. + self (nn.Module): The RMSNorm module instance with ``weight`` and either ``variance_epsilon`` or ``eps``. hidden_states (Tensor): Input hidden states tensor. Returns: @@ -75,48 +53,108 @@ def npu_rms_norm_forward(self, hidden_states): """ _eps = getattr(self, "variance_epsilon", None) or getattr(self, "eps", 1e-6) - if hasattr(self, "weight") and self.weight is not None: - if getattr(self, "_npu_use_residual_rmsnorm", False): - effective_weight = 1.0 + self.weight.float() - else: - effective_weight = self.weight.float() - else: - effective_weight = None + weight = getattr(self, "weight", None) + if weight is None: + raise RuntimeError(f"{self.__class__.__name__} has no RMSNorm weight for NPU RMSNorm kernel.") - if effective_weight is not None: - return torch_npu.npu_rms_norm(hidden_states, effective_weight.to(hidden_states.dtype), epsilon=_eps)[0] - else: - return torch_npu.npu_rms_norm(hidden_states, self.weight, epsilon=_eps)[0] + effective_weight = weight.float() + + return torch_npu.npu_rms_norm(hidden_states, effective_weight.to(hidden_states.dtype), epsilon=_eps)[0] + + +def npu_residual_rms_norm_forward(self, hidden_states): + """NPU forward implementation for residual RMSNorm. + + Residual RMSNorm uses ``scale = 1.0 + weight`` where ``weight`` is initialized + to 0 in the original transformers implementation. + + Args: + self (nn.Module): The residual RMSNorm module with ``weight`` and either ``variance_epsilon`` or ``eps``. + hidden_states (Tensor): Input hidden states tensor. + + Returns: + Tensor: Normalized tensor consistent with residual RMSNorm behavior. + """ + _eps = getattr(self, "variance_epsilon", None) or getattr(self, "eps", 1e-6) + + weight = getattr(self, "weight", None) + if weight is None: + raise RuntimeError(f"{self.__class__.__name__} has no RMSNorm weight for NPU RMSNorm kernel.") + + effective_weight = 1.0 + weight.float() + + return torch_npu.npu_rms_norm(hidden_states, effective_weight.to(hidden_states.dtype), epsilon=_eps)[0] def npu_gated_rms_norm_forward(self, hidden_states, gate=None): """NPU forward implementation for Gated RMSNorm with high-precision FP32 computation. This function performs RMSNorm and gated SiLU multiplication in FP32 for numerical - stability. Unlike standard RMSNorm, Gated RMSNorm in Qwen3.5 uses standard - parameterization (``scale = weight`` where weight is initialized to 1), so the - residual weight adjustment (``1.0 + weight``) is not applied here. + stability. The supported gated RMSNorm modules use ``scale = weight`` with weight + initialized to 1, unlike the residual RMSNorm variants that use ``1.0 + weight``. Args: self (nn.Module): The Gated RMSNorm module instance. hidden_states (Tensor): Input hidden states tensor. - gate (Tensor, optional): Gate tensor for SiLU activation. Defaults to ``None``. + gate (Tensor): Gate tensor for SiLU activation. Returns: Tensor: Output tensor cast back to the original input dtype. + + Raises: + ValueError: If the gate tensor is not provided. """ + if gate is None: + raise ValueError(f"{self.__class__.__name__} requires a gate tensor for NPU Gated RMSNorm.") + input_dtype = hidden_states.dtype hidden_states = hidden_states.to(torch.float32) _eps = getattr(self, "variance_epsilon", None) or getattr(self, "eps", 1e-6) hidden_states = torch_npu.npu_rms_norm(hidden_states, self.weight.float(), epsilon=_eps)[0] - - if gate is not None: - hidden_states = hidden_states * F.silu(gate.to(torch.float32)) + hidden_states = hidden_states * F.silu(gate.to(torch.float32)) return hidden_states.to(input_dtype) +_MODEL_TYPE_TO_PATCHES = { + "qwen3": { + "Qwen3RMSNorm": npu_rms_norm_forward, + }, + "qwen3_moe": { + "Qwen3MoeRMSNorm": npu_rms_norm_forward, + }, + "qwen3_next": { + "Qwen3NextRMSNorm": npu_residual_rms_norm_forward, + "Qwen3NextRMSNormGated": npu_gated_rms_norm_forward, + }, + "qwen3_omni_moe": { + "Qwen3OmniMoeThinkerTextRMSNorm": npu_rms_norm_forward, + "Qwen3OmniMoeTextRMSNorm": npu_rms_norm_forward, + "Qwen3OmniMoeRMSNorm": npu_rms_norm_forward, + "Qwen3OmniMoeCode2WavRMSNorm": npu_rms_norm_forward, + }, + "qwen3_omni_moe_thinker": { + "Qwen3OmniMoeThinkerTextRMSNorm": npu_rms_norm_forward, + "Qwen3OmniMoeTextRMSNorm": npu_rms_norm_forward, + }, + "qwen3_vl": { + "Qwen3VLTextRMSNorm": npu_rms_norm_forward, + }, + "qwen3_vl_moe": { + "Qwen3VLMoeTextRMSNorm": npu_rms_norm_forward, + }, + "qwen3_5": { + "Qwen3_5RMSNorm": npu_residual_rms_norm_forward, + "Qwen3_5RMSNormGated": npu_gated_rms_norm_forward, + }, + "qwen3_5_moe": { + "Qwen3_5MoeRMSNorm": npu_residual_rms_norm_forward, + "Qwen3_5MoeRMSNormGated": npu_gated_rms_norm_forward, + }, +} + + @KernelPlugin("npu_fused_rmsnorm").register() class NpuRMSNormKernel(BaseKernel): """NPU kernel wrapper for RMSNorm that applies the replacement within a model.""" @@ -127,34 +165,45 @@ class NpuRMSNormKernel(BaseKernel): if current != DeviceType.NPU: raise RuntimeError(f"NpuRMSNormKernel requires NPU, current accelerator is {current}.") + @staticmethod + def check_deps() -> None: + if _TORCH_NPU_IMPORT_ERROR is not None: + raise RuntimeError("NpuRMSNormKernel requires torch_npu.") from _TORCH_NPU_IMPORT_ERROR + + @staticmethod + def _get_patch_forward(model_type: str, module: torch.nn.Module): + """Return the NPU forward function for a matched RMSNorm module.""" + model_patches = _MODEL_TYPE_TO_PATCHES.get(model_type, {}) + return model_patches.get(module.__class__.__name__) + @staticmethod def _apply(**kwargs) -> "HFModel": """Iterate the model and apply NPU-optimized forward to matched RMSNorm modules. - Matches modules whose class name contains "RMSNorm" (case-insensitive) and binds - the appropriate NPU-optimized forward function as an instance method via - ``types.MethodType`` to replace the original ``forward``. + Matches modules configured for the current model type, then binds the corresponding + NPU-optimized forward function as an instance method via ``types.MethodType`` to + replace the original ``forward``. Args: **kwargs: Keyword arguments containing the model. Returns: HFModel: The model with NPU fused RMSNorm. - - Raises: - RuntimeError: If ``torch_npu`` is not available. - ValueError: If the model is not provided. """ - model = kwargs.get("model") + model = kwargs["model"] - rms_norm_pattern = re.compile("RMSNorm", re.IGNORECASE) + model_type = getattr(model.config, "model_type", None) + if model_type not in _MODEL_TYPE_TO_PATCHES: + return model - for _, module in model.named_modules(): - if re.search(rms_norm_pattern, module.__class__.__name__): - if "Gated" in module.__class__.__name__: - module.forward = types.MethodType(npu_gated_rms_norm_forward, module) - else: - module._npu_use_residual_rmsnorm = _should_use_residual_rmsnorm(module) - module.forward = types.MethodType(npu_rms_norm_forward, module) + patched_count = 0 + for module in model.modules(): + patch_forward = NpuRMSNormKernel._get_patch_forward(model_type, module) + if patch_forward is not None: + module.forward = types.MethodType(patch_forward, module) + patched_count += 1 + + if patched_count: + logger.info_rank0(f"Applied NPU RMSNorm kernel to {patched_count} modules for model type: {model_type}.") return model diff --git a/src/llamafactory/v1/plugins/model_plugins/kernels/ops/rope/npu_rope.py b/src/llamafactory/v1/plugins/model_plugins/kernels/ops/rope/npu_rope.py index 3ee39e1d7..adb177a74 100644 --- a/src/llamafactory/v1/plugins/model_plugins/kernels/ops/rope/npu_rope.py +++ b/src/llamafactory/v1/plugins/model_plugins/kernels/ops/rope/npu_rope.py @@ -20,7 +20,7 @@ Init Phase: """ -import sys +import importlib import torch @@ -34,16 +34,18 @@ logger = get_logger(__name__) try: import torch_npu -except ImportError: - pass +except ImportError as exc: + _TORCH_NPU_IMPORT_ERROR = exc +else: + _TORCH_NPU_IMPORT_ERROR = None def _apply_npu_rotary_emb(q, k, cos, sin): """Apply NPU-accelerated rotary embedding with automatic Partial RoPE detection. - This function automatically detects whether to use Partial RoPE or Full RoPE - based on the dimension ratio between ``cos/sin`` and ``q/k`` tensors, ensuring - compatibility with future model versions without hardcoding. + Partial RoPE is detected when the ``cos/sin`` width is smaller than the ``q/k`` + head dimension. The leading rotary dimensions are transformed and any trailing + dimensions are passed through unchanged. Args: q (Tensor): Query tensor. @@ -61,14 +63,14 @@ def _apply_npu_rotary_emb(q, k, cos, sin): q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:] k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:] - q_embed = torch_npu.npu_rotary_mul(q_rot, cos, sin).to(q.dtype) - k_embed = torch_npu.npu_rotary_mul(k_rot, cos, sin).to(k.dtype) + q_embed = torch_npu.npu_rotary_mul(q_rot, cos, sin, "half").to(q.dtype) + k_embed = torch_npu.npu_rotary_mul(k_rot, cos, sin, "half").to(k.dtype) q_embed = torch.cat([q_embed, q_pass], dim=-1) k_embed = torch.cat([k_embed, k_pass], dim=-1) else: - q_embed = torch_npu.npu_rotary_mul(q, cos, sin).to(q.dtype) - k_embed = torch_npu.npu_rotary_mul(k, cos, sin).to(k.dtype) + q_embed = torch_npu.npu_rotary_mul(q, cos, sin, "half").to(q.dtype) + k_embed = torch_npu.npu_rotary_mul(k, cos, sin, "half").to(k.dtype) return q_embed, k_embed @@ -84,44 +86,43 @@ def _apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): k (Tensor): Key tensor. cos (Tensor): Cosine part of embedding. sin (Tensor): Sine part of embedding. - position_ids (Tensor, optional): Position IDs. Defaults to ``None``. + position_ids (Tensor | int, optional): Ignored Transformers v4 position IDs, or the Transformers v5 + ``unsqueeze_dim`` when supplied as the fifth positional argument. unsqueeze_dim (int): Dimension to unsqueeze cos and sin. Defaults to 1. Returns: tuple[Tensor, Tensor]: The embedded query and key tensors ``(q_embed, k_embed)``. """ + # In transformers v5, the fifth positional argument is ``unsqueeze_dim``. + if isinstance(position_ids, int): + unsqueeze_dim = position_ids + cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) return _apply_npu_rotary_emb(q, k, cos, sin) -def _apply_multimodal_rotary_pos_emb_qwen25_vl(q, k, cos, sin, mrope_section, unsqueeze_dim=1): - """Apply Rotary Position Embedding with multimodal sections (Qwen2-VL) on NPU. - - This function supports Partial RoPE for multimodal inputs with automatic dimension - detection, ensuring compatibility with future model versions. - - Args: - q (Tensor): Query tensor. - k (Tensor): Key tensor. - cos (Tensor): Cosine part of embedding. - sin (Tensor): Sine part of embedding. - mrope_section (list[int]): Multimodal RoPE section sizes. - unsqueeze_dim (int): Dimension to unsqueeze cos and sin. Defaults to 1. - - Returns: - tuple[Tensor, Tensor]: The embedded query and key tensors ``(q_embed, k_embed)``. - """ - mrope_section = mrope_section * 2 - cos = torch.cat([m[i % 3] for i, m in enumerate(cos.split(mrope_section, dim=-1))], dim=-1).unsqueeze( - unsqueeze_dim - ) - sin = torch.cat([m[i % 3] for i, m in enumerate(sin.split(mrope_section, dim=-1))], dim=-1).unsqueeze( - unsqueeze_dim +def _default_rope_patch(module_type: str): + return ( + ( + f"transformers.models.{module_type}.modeling_{module_type}", + (("apply_rotary_pos_emb", _apply_rotary_pos_emb),), + ), ) - return _apply_npu_rotary_emb(q, k, cos, sin) + +_MODEL_TYPE_TO_PATCHES = { + "qwen3": _default_rope_patch("qwen3"), + "qwen3_moe": _default_rope_patch("qwen3_moe"), + "qwen3_next": _default_rope_patch("qwen3_next"), + "qwen3_omni_moe": _default_rope_patch("qwen3_omni_moe"), + "qwen3_omni_moe_thinker": _default_rope_patch("qwen3_omni_moe"), + "qwen3_vl": _default_rope_patch("qwen3_vl"), + "qwen3_vl_moe": _default_rope_patch("qwen3_vl_moe"), + "qwen3_5": _default_rope_patch("qwen3_5"), + "qwen3_5_moe": _default_rope_patch("qwen3_5_moe"), +} @KernelPlugin("npu_fused_rope").register() @@ -135,50 +136,58 @@ class NpuRoPEKernel(BaseKernel): raise RuntimeError(f"NpuRoPEKernel requires NPU, current accelerator is {current}.") @staticmethod - def _apply(**kwargs) -> "HFModel": - """Apply RoPE acceleration by monkey-patching ``apply_rotary_pos_emb``. + def check_deps() -> None: + if _TORCH_NPU_IMPORT_ERROR is not None: + raise RuntimeError("NpuRoPEKernel requires torch_npu.") from _TORCH_NPU_IMPORT_ERROR - Iterates through the model's modules to find attention layers, identifies - the module where they are defined, and replaces the original - ``apply_rotary_pos_emb`` function in that module's namespace with the - NPU-accelerated version. + @staticmethod + def _apply_model_patches(model_type: str) -> int: + patches = _MODEL_TYPE_TO_PATCHES.get(model_type) + if patches is None: + return 0 + + patched_count = 0 + for module_name, replacements in patches: + try: + target_module = importlib.import_module(module_name) + except Exception as e: + logger.warning_rank0_once(f"Failed to import {module_name} for NPU RoPE kernel: {e}") + continue + + for target_function_name, replacement in replacements: + if not hasattr(target_module, target_function_name): + logger.warning_rank0_once(f"{module_name} has no {target_function_name}, skip NPU RoPE patch.") + continue + + if getattr(target_module, target_function_name) is replacement: + continue + + setattr(target_module, target_function_name, replacement) + patched_count += 1 + + return patched_count + + @staticmethod + def _apply(**kwargs) -> "HFModel": + """Apply RoPE acceleration by monkey-patching rotary embedding functions. + + Selects the target transformers modeling module from ``model.config.model_type`` + and replaces its rotary embedding helper with the NPU-accelerated version. Args: **kwargs: Keyword arguments containing the model. Returns: HFModel: The model with patched RoPE functions. - - Raises: - RuntimeError: If ``torch_npu`` is not available. - ValueError: If the model is not provided. """ - model = kwargs.get("model", None) + model = kwargs["model"] - _modules = set() - for module in model.modules(): - if "Attention" in module.__class__.__name__: - module_name = module.__class__.__module__ - if module_name in _modules: - continue - try: - target_module = sys.modules[module_name] - if hasattr(target_module, "apply_rotary_pos_emb"): - if getattr(target_module, "apply_rotary_pos_emb") is not _apply_rotary_pos_emb: - setattr(target_module, "apply_rotary_pos_emb", _apply_rotary_pos_emb) - _modules.add(module_name) - if hasattr(target_module, "apply_multimodal_rotary_pos_emb"): - if ( - getattr(target_module, "apply_multimodal_rotary_pos_emb") - is not _apply_multimodal_rotary_pos_emb_qwen25_vl - ): - setattr( - target_module, - "apply_multimodal_rotary_pos_emb", - _apply_multimodal_rotary_pos_emb_qwen25_vl, - ) - _modules.add(module_name) - except Exception as e: - logger.warning_rank0_once(f"Failed to apply RoPE kernel to module {module_name}: {e}") + model_type = getattr(model.config, "model_type", None) + if model_type not in _MODEL_TYPE_TO_PATCHES: + return model + + patched_count = NpuRoPEKernel._apply_model_patches(model_type) + if patched_count: + logger.info_rank0(f"Applied NPU RoPE kernel to {patched_count} functions for model type: {model_type}.") return model diff --git a/tests_v1/plugins/model_plugins/test_kernel_plugin.py b/tests_v1/plugins/model_plugins/test_kernel_plugin.py index 29822477f..3b4f141e5 100644 --- a/tests_v1/plugins/model_plugins/test_kernel_plugin.py +++ b/tests_v1/plugins/model_plugins/test_kernel_plugin.py @@ -25,18 +25,19 @@ def _apply_kernel(rank) -> None: setattr(mock_device, "type", "npu") mock_get_accelerator.return_value = mock_device - # reload kernel modules to respect mocked accelerator - for k in list(sys.modules.keys()): - if k.startswith("llamafactory.v1.plugins.model_plugins.kernels"): - del sys.modules[k] - - from llamafactory.v1.plugins.model_plugins.kernels.interface import apply_kernels - model = AutoModelForCausalLM.from_pretrained("llamafactory/tiny-random-qwen3") original_rmsnorm_forward = model.model.layers[0].input_layernorm.forward original_swiglu_forward = model.model.layers[0].mlp.forward - model = apply_kernels(model=model, config={"name": "npu_fused_rmsnorm"}) + with patch.dict(sys.modules, {"torch_npu": MagicMock()}): + # Reload kernel modules so dependency checks use the mocked NPU environment. + for k in list(sys.modules.keys()): + if k.startswith("llamafactory.v1.plugins.model_plugins.kernels"): + del sys.modules[k] + + from llamafactory.v1.plugins.model_plugins.kernels.interface import apply_kernels + + model = apply_kernels(model=model, config={"name": "npu_fused_rmsnorm"}) assert model.model.layers[0].input_layernorm.forward.__func__ is not original_rmsnorm_forward.__func__ assert model.model.layers[0].mlp.forward.__func__ is original_swiglu_forward.__func__ @@ -48,18 +49,19 @@ def _apply_all_kernels(rank) -> None: setattr(mock_device, "type", "npu") mock_get_accelerator.return_value = mock_device - # reload kernel modules to respect mocked accelerator - for k in list(sys.modules.keys()): - if k.startswith("llamafactory.v1.plugins.model_plugins.kernels"): - del sys.modules[k] - - from llamafactory.v1.plugins.model_plugins.kernels.interface import apply_kernels - model = AutoModelForCausalLM.from_pretrained("llamafactory/tiny-random-qwen3") original_rmsnorm_forward = model.model.layers[0].input_layernorm.forward original_swiglu_forward = model.model.layers[0].mlp.forward - model = apply_kernels(model=model, config={"name": "auto"}) + with patch.dict(sys.modules, {"torch_npu": MagicMock()}): + # Reload kernel modules so dependency checks use the mocked NPU environment. + for k in list(sys.modules.keys()): + if k.startswith("llamafactory.v1.plugins.model_plugins.kernels"): + del sys.modules[k] + + from llamafactory.v1.plugins.model_plugins.kernels.interface import apply_kernels + + model = apply_kernels(model=model, config={"name": "auto"}) assert model.model.layers[0].input_layernorm.forward.__func__ is not original_rmsnorm_forward.__func__ assert model.model.layers[0].mlp.forward.__func__ is not original_swiglu_forward.__func__