hiyouga/LlamaFactory · error · RuntimeError

NpuRMSNormKernel requires torch_npu.

Error message

NpuRMSNormKernel requires torch_npu.

What it means

The npu_fused_rmsnorm plugin depends on torch_npu for torch_npu.npu_rms_norm. If the import failed at module load, check_deps() raises RuntimeError chained to the stored ImportError when the plugin is applied.

Source

Thrown at src/llamafactory/v1/plugins/model_plugins/kernels/ops/rms_norm/npu_rms_norm.py:171

        "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."""

    @staticmethod
    def check_device() -> None:
        current = get_current_accelerator().type
        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 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.

View on GitHub (pinned to f28afaf635)

Solutions

  1. Install the torch_npu build matching torch and CANN; verify `import torch_npu` succeeds
  2. Read the chained __cause__ ImportError for the precise broken symbol/module
  3. Source CANN set_env.sh in the container/launch environment

Example fix

# before
# RuntimeError: NpuRMSNormKernel requires torch_npu.
# after
pip install torch_npu==<matching> && source /usr/local/Ascend/ascend-toolkit/set_env.sh
Defensive patterns

Strategy: validation

Validate before calling

try:
    import torch_npu  # noqa: F401
except ImportError:
    kernels = [k for k in kernels if k != "npu_fused_rmsnorm"]

Prevention

When it happens

Trigger: Applying npu_fused_rmsnorm in an environment where `import torch_npu` raises — wheel missing, torch version mismatch, or CANN toolkit environment not sourced.

Common situations: Ascend containers with mismatched torch/torch_npu/CANN triples; pip upgrading torch and silently breaking torch_npu's pinned ABI.

Related errors


AI-assisted analysis of hiyouga/LlamaFactory@f28afaf635 (2026-08-14). Data as JSON: /api/errors/971a9c5a2231c6e8. Report an issue: GitHub.