diff --git a/src/llamafactory/extras/env.py b/src/llamafactory/extras/env.py index 2f0de31f5..fb809a587 100644 --- a/src/llamafactory/extras/env.py +++ b/src/llamafactory/extras/env.py @@ -33,6 +33,8 @@ def print_env() -> None: import transformers from transformers.utils import is_torch_cuda_available, is_torch_npu_available + from .misc import is_torch_supa_available + info = OrderedDict( { "`llamafactory` version": VERSION, @@ -46,7 +48,15 @@ def print_env() -> None: } ) - if is_torch_cuda_available(): + if is_torch_supa_available(): + info["PyTorch version"] += " (SUPA)" + info["SUPA type"] = torch.supa.get_device_name() + info["SUPA number"] = torch.supa.device_count() + try: + info["SUPA memory"] = f"{torch.supa.mem_get_info()[1] / (1024**3):.2f}GB" + except Exception: + pass + elif is_torch_cuda_available(): info["PyTorch version"] += " (GPU)" info["GPU type"] = torch.cuda.get_device_name() info["GPU number"] = torch.cuda.device_count() diff --git a/src/llamafactory/extras/misc.py b/src/llamafactory/extras/misc.py index f95ab9f49..ddb8f990e 100644 --- a/src/llamafactory/extras/misc.py +++ b/src/llamafactory/extras/misc.py @@ -15,6 +15,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import functools import gc import os import socket @@ -38,9 +39,25 @@ from transformers.utils.versions import require_version from . import logging -_is_fp16_available = is_torch_npu_available() or is_torch_cuda_available() +@functools.lru_cache +def is_torch_supa_available() -> bool: + """Check if the supa (SUPA PrivateUse1) accelerator is available.""" + try: + import torch_supa # noqa: F401 # registers the "supa" device + + return hasattr(torch, "supa") and torch.supa.is_available() + except Exception: + return False + + +_is_fp16_available = is_torch_npu_available() or is_torch_supa_available() or is_torch_cuda_available() try: - _is_bf16_available = is_torch_bf16_gpu_available() or (is_torch_npu_available() and torch.npu.is_bf16_supported()) + if is_torch_supa_available(): + _is_bf16_available = torch.supa.is_bf16_supported() + else: + _is_bf16_available = is_torch_bf16_gpu_available() or ( + is_torch_npu_available() and torch.npu.is_bf16_supported() + ) except Exception: _is_bf16_available = False @@ -147,6 +164,8 @@ def get_current_device() -> "torch.device": device = "xpu:{}".format(os.getenv("LOCAL_RANK", "0")) elif is_torch_npu_available(): device = "npu:{}".format(os.getenv("LOCAL_RANK", "0")) + elif is_torch_supa_available(): + device = "supa:{}".format(os.getenv("LOCAL_RANK", "0")) elif is_torch_mps_available(): device = "mps:{}".format(os.getenv("LOCAL_RANK", "0")) elif is_torch_cuda_available(): @@ -163,6 +182,8 @@ def get_device_name() -> str: device = "xpu" elif is_torch_npu_available(): device = "npu" + elif is_torch_supa_available(): + device = "supa" elif is_torch_mps_available(): device = "mps" elif is_torch_cuda_available(): @@ -190,6 +211,8 @@ def get_device_count() -> int: return torch.xpu.device_count() elif is_torch_npu_available(): return torch.npu.device_count() + elif is_torch_supa_available(): + return torch.supa.device_count() elif is_torch_mps_available(): return torch.mps.device_count() elif is_torch_cuda_available(): @@ -211,6 +234,8 @@ def get_current_memory() -> tuple[int, int]: return torch.xpu.mem_get_info() elif is_torch_npu_available(): return torch.npu.mem_get_info() + elif is_torch_supa_available(): + return torch.supa.mem_get_info() elif is_torch_mps_available(): return torch.mps.current_allocated_memory(), torch.mps.recommended_max_memory() elif is_torch_cuda_available(): @@ -225,6 +250,8 @@ def get_peak_memory() -> tuple[int, int]: return torch.xpu.max_memory_allocated(), torch.xpu.max_memory_reserved() elif is_torch_npu_available(): return torch.npu.max_memory_allocated(), torch.npu.max_memory_reserved() + elif is_torch_supa_available(): + return torch.supa.max_memory_allocated(), torch.supa.max_memory_reserved() elif is_torch_mps_available(): return torch.mps.current_allocated_memory(), -1 elif is_torch_cuda_available(): @@ -251,7 +278,11 @@ def infer_optim_dtype(model_dtype: Optional["torch.dtype"]) -> "torch.dtype": def is_accelerator_available() -> bool: r"""Check if the accelerator is available.""" return ( - is_torch_xpu_available() or is_torch_npu_available() or is_torch_mps_available() or is_torch_cuda_available() + is_torch_xpu_available() + or is_torch_npu_available() + or is_torch_supa_available() + or is_torch_mps_available() + or is_torch_cuda_available() ) @@ -285,6 +316,8 @@ def torch_gc() -> None: torch.xpu.empty_cache() elif is_torch_npu_available(): torch.npu.empty_cache() + elif is_torch_supa_available(): + torch.supa.empty_cache() elif is_torch_mps_available(): torch.mps.empty_cache() elif is_torch_cuda_available(): diff --git a/src/llamafactory/hparams/parser.py b/src/llamafactory/hparams/parser.py index 1f686d641..4252e2a31 100644 --- a/src/llamafactory/hparams/parser.py +++ b/src/llamafactory/hparams/parser.py @@ -33,7 +33,13 @@ from transformers.utils import is_torch_bf16_gpu_available, is_torch_npu_availab from ..extras import logging from ..extras.constants import CHECKPOINT_NAMES, EngineName -from ..extras.misc import check_dependencies, check_version, get_current_device, is_env_enabled +from ..extras.misc import ( + check_dependencies, + check_version, + get_current_device, + is_env_enabled, + is_torch_supa_available, +) from ..extras.packages import is_mcore_adapter_available, is_megatron_bridge_available from .data_args import DataArguments from .evaluation_args import EvaluationArguments @@ -223,6 +229,20 @@ def _set_env_vars() -> None: # avoid use fork method on NPU devices, see https://github.com/hiyouga/LLaMA-Factory/issues/7447 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + if is_torch_supa_available(): + # register the fused SUDNN operators (torch.ops.sudnn.*) used by the supa kernel path. + # torch_supa itself is auto-loaded on `import torch`, but torch_supa_ext is not. + try: + import torch_supa_ext # noqa: F401 + except Exception: + logger.warning_rank0_once( + "Failed to import torch_supa_ext; fused supa operators (e.g. sudnn.rms_norm_func) " + "will be unavailable and the eager path may be extremely slow." + ) + + # supa uses fork-unsafe device handles; align with the NPU workaround. + os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + def _verify_model_args( model_args: "ModelArguments", @@ -260,6 +280,9 @@ def _check_extra_dependencies( check_version("accelerate-kt", mandatory=True) if model_args.use_unsloth: + if is_torch_supa_available(): + raise ValueError("`use_unsloth` is not supported on supa devices.") + check_version("unsloth", mandatory=True) if model_args.enable_liger_kernel: @@ -496,7 +519,11 @@ def get_train_args(args: dict[str, Any] | list[str] | None = None) -> _TRAIN_CLS raise ValueError("Please use scripts/pissa_init.py to initialize PiSSA in DeepSpeed ZeRO-3.") if finetuning_args.pure_bf16: - if not (is_torch_bf16_gpu_available() or (is_torch_npu_available() and torch.npu.is_bf16_supported())): + if not ( + is_torch_bf16_gpu_available() + or (is_torch_npu_available() and torch.npu.is_bf16_supported()) + or (is_torch_supa_available() and torch.supa.is_bf16_supported()) + ): raise ValueError("This device does not support `pure_bf16`.") if is_deepspeed_zero3_enabled(): diff --git a/src/llamafactory/model/model_utils/attention.py b/src/llamafactory/model/model_utils/attention.py index f3ab25df6..3d50fa2b4 100644 --- a/src/llamafactory/model/model_utils/attention.py +++ b/src/llamafactory/model/model_utils/attention.py @@ -16,6 +16,7 @@ from typing import TYPE_CHECKING from ...extras import logging from ...extras.constants import AttentionFunction +from ...extras.misc import is_torch_supa_available from ...extras.packages import is_torch_version_greater_than @@ -63,7 +64,12 @@ def configure_attn_implementation(config: "PretrainedConfig", model_args: "Model model_args.flash_attn = AttentionFunction.DISABLED if model_args.flash_attn == AttentionFunction.AUTO: - return + if is_torch_supa_available(): + # supa has no FlashAttention build, and AUTO would otherwise be resolved against the + # cuda-masqueraded device (which may pick FA2); pin SDPA, the fastest supported backend. + requested_attn_implementation = "sdpa" + else: + return elif model_args.flash_attn == AttentionFunction.DISABLED: requested_attn_implementation = "eager" diff --git a/src/llamafactory/model/model_utils/liger_kernel.py b/src/llamafactory/model/model_utils/liger_kernel.py index 0eece68f8..1698af63d 100644 --- a/src/llamafactory/model/model_utils/liger_kernel.py +++ b/src/llamafactory/model/model_utils/liger_kernel.py @@ -37,6 +37,10 @@ def apply_liger_kernel( if not is_trainable or not model_args.enable_liger_kernel: return + if get_device_name() == "supa": + # liger kernels can not run on supa currently. + return + model_type = getattr(config, "model_type", None) if model_type == "gemma": from liger_kernel.transformers import apply_liger_kernel_to_gemma as apply_liger_kernel diff --git a/src/llamafactory/model/patcher.py b/src/llamafactory/model/patcher.py index f1c7fc8f5..0b198838a 100644 --- a/src/llamafactory/model/patcher.py +++ b/src/llamafactory/model/patcher.py @@ -480,7 +480,7 @@ def patch_model( prepare_model_for_training(model, model_args) autocast_projector_dtype(model, model_args) add_z3_leaf_module(model) - + if getattr(model.config, "model_type", None) in [ "qwen3_5", "qwen3_5_moe", diff --git a/src/llamafactory/train/callbacks.py b/src/llamafactory/train/callbacks.py index b3826a7a9..b50db5c30 100644 --- a/src/llamafactory/train/callbacks.py +++ b/src/llamafactory/train/callbacks.py @@ -33,7 +33,14 @@ from typing_extensions import override from ..extras import logging from ..extras.constants import TRAINER_LOG, V_HEAD_SAFE_WEIGHTS_NAME, V_HEAD_WEIGHTS_NAME -from ..extras.misc import get_peak_memory, is_env_enabled, is_torch_cuda_available, is_torch_npu_available, use_ray +from ..extras.misc import ( + get_peak_memory, + is_env_enabled, + is_torch_cuda_available, + is_torch_npu_available, + is_torch_supa_available, + use_ray, +) from ..extras.packages import is_safetensors_available @@ -504,6 +511,8 @@ class ModuleProfilerCallback(TrainerCallback): @staticmethod def _get_accelerator(): """Detect available accelerator and return (event_factory, synchronize_fn).""" + if is_torch_supa_available(): + return torch.supa.Event, torch.supa.synchronize if is_torch_cuda_available(): return torch.cuda.Event, torch.cuda.synchronize if is_torch_npu_available(): diff --git a/src/llamafactory/train/tuner.py b/src/llamafactory/train/tuner.py index 97a6ef725..08ac5f87e 100644 --- a/src/llamafactory/train/tuner.py +++ b/src/llamafactory/train/tuner.py @@ -291,6 +291,7 @@ class Worker: RAY_NOSET_VISIBLE_DEVICES_LIST = [ "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES", "RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES", + "RAY_EXPERIMENTAL_NOSET_SUPA_VISIBLE_DEVICES", ] is_ray_noset_visible_devices = any(os.environ.get(env_var, None) for env_var in RAY_NOSET_VISIBLE_DEVICES_LIST) if is_ray_noset_visible_devices: diff --git a/src/llamafactory/v1/accelerator/helper.py b/src/llamafactory/v1/accelerator/helper.py index b83e7f2f5..ed94d7ecc 100644 --- a/src/llamafactory/v1/accelerator/helper.py +++ b/src/llamafactory/v1/accelerator/helper.py @@ -45,6 +45,7 @@ class DeviceType(StrEnum): META = "meta" MPS = "mps" NPU = "npu" + SUPA = "supa" XPU = "xpu" diff --git a/src/llamafactory/v1/plugins/model_plugins/kernels/interface.py b/src/llamafactory/v1/plugins/model_plugins/kernels/interface.py index 33c9b1290..4978f6f4f 100644 --- a/src/llamafactory/v1/plugins/model_plugins/kernels/interface.py +++ b/src/llamafactory/v1/plugins/model_plugins/kernels/interface.py @@ -25,11 +25,13 @@ from .ops.mlp.cuda_fused_moe import CudaFusedMoEKernel # noqa: F401 from .ops.mlp.npu_fused_moe import NpuFusedMoEKernel # noqa: F401 from .ops.mlp.npu_swiglu import NpuSwiGluKernel # noqa: F401 from .ops.rms_norm.npu_rms_norm import NpuRMSNormKernel # noqa: F401 +from .ops.rms_norm.supa_rms_norm import SupaRMSNormKernel # noqa: F401 from .ops.rope.npu_rope import NpuRoPEKernel # noqa: F401 _AUTO_KERNELS = { DeviceType.NPU: ("npu_fused_moe", "npu_fused_rmsnorm", "npu_fused_rope", "npu_fused_swiglu"), + DeviceType.SUPA: ("supa_fused_rmsnorm",), } diff --git a/src/llamafactory/v1/plugins/model_plugins/kernels/ops/rms_norm/supa_rms_norm.py b/src/llamafactory/v1/plugins/model_plugins/kernels/ops/rms_norm/supa_rms_norm.py new file mode 100644 index 000000000..38a970bb1 --- /dev/null +++ b/src/llamafactory/v1/plugins/model_plugins/kernels/ops/rms_norm/supa_rms_norm.py @@ -0,0 +1,158 @@ +# Copyright 2025 the LlamaFactory team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""The definition of the SUPA fused RMSNorm kernel.""" + +import types +from functools import lru_cache + +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__) + + +# Standard-convention RMSNorm classes: ``y = x * rsqrt(mean(x^2) + eps) * weight``. +# The fused ``sudnn.rms_norm_func`` matches this convention, so only these classes are patched. +# Variants with a different weight convention (e.g. Gemma's ``1 + weight``, or gated/residual +# RMSNorm) are intentionally excluded and fall back to the eager path. Add new class names here +# after verifying they follow the standard convention. +_SUPPORTED_RMSNORM_CLASSES = frozenset( + { + "LlamaRMSNorm", + "MistralRMSNorm", + "MixtralRMSNorm", + "Phi3RMSNorm", + "InternLM2RMSNorm", + "GlmRMSNorm", + "DeepseekV2RMSNorm", + "DeepseekV3RMSNorm", + "Qwen2RMSNorm", + "Qwen2MoeRMSNorm", + "Qwen2VLRMSNorm", + "Qwen2_5_VLRMSNorm", + "Qwen2_5_VLTextRMSNorm", + "Qwen3RMSNorm", + "Qwen3MoeRMSNorm", + "Qwen3VLTextRMSNorm", + "Qwen3VLMoeTextRMSNorm", + } +) + + + +@lru_cache +def _is_sudnn_rms_norm_available() -> bool: + """Return whether the fused SUPA RMSNorm operator is registered.""" + try: + import torch_supa # noqa: F401 + import torch_supa_ext # noqa: F401 + except Exception: + return False + + return hasattr(torch.ops, "sudnn") and hasattr(torch.ops.sudnn, "rms_norm_func") + + +class SupaRMSNormFunction(torch.autograd.Function): + """Autograd wrapper around the fused SUPA RMSNorm forward operator.""" + + @staticmethod + def forward(ctx, hidden_states: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + weight = weight.to(dtype=hidden_states.dtype) + bias = torch.zeros_like(weight) + output = torch.ops.sudnn.rms_norm_func(hidden_states, weight, bias, eps, False, False) + ctx.save_for_backward(hidden_states, weight) + ctx.eps = eps + return output + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + hidden_states, weight = ctx.saved_tensors + eps = ctx.eps + hidden_size = hidden_states.shape[-1] + + variance = hidden_states.pow(2).mean(-1, keepdim=True) + rrms = torch.rsqrt(variance + eps) + x_norm = hidden_states * rrms + + grad_input = grad_weight = None + if ctx.needs_input_grad[1]: + grad_weight = (grad_output * x_norm).reshape(-1, hidden_size).sum(0).to(weight.dtype) + + if ctx.needs_input_grad[0]: + grad_norm = grad_output * weight + dot = (grad_norm * hidden_states).sum(-1, keepdim=True) + grad_input = rrms * grad_norm - (rrms.pow(3) / hidden_size) * hidden_states * dot + grad_input = grad_input.to(hidden_states.dtype) + + return grad_input, grad_weight, None + + +def supa_rms_norm_forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Run the fused kernel for standard RMSNorm modules.""" + if hidden_states.device.type != DeviceType.SUPA: + return self._supa_rmsnorm_original_forward(hidden_states) + + return SupaRMSNormFunction.apply(hidden_states, self.weight, self.variance_epsilon) + + +@KernelPlugin("supa_fused_rmsnorm").register() +class SupaRMSNormKernel(BaseKernel): + """SUPA fused RMSNorm kernel for trainable models.""" + + @staticmethod + def check_device() -> None: + current = get_current_accelerator().type + if current != DeviceType.SUPA: + raise RuntimeError(f"SupaRMSNormKernel requires SUPA, current accelerator is {current}.") + + @staticmethod + def _apply(**kwargs) -> HFModel: + model = kwargs["model"] + + if not _is_sudnn_rms_norm_available(): + logger.warning_rank0( + "Fused SUPA RMSNorm (torch.ops.sudnn.rms_norm_func) is unavailable; " + "the eager RMSNorm path will be used. Ensure torch_supa_ext is installed." + ) + return model + + patched_count = 0 + patched_classes = set() + for module in model.modules(): + cls = module.__class__ + if getattr(module, "_supa_rmsnorm_patched", False): + continue + if cls.__name__ not in _SUPPORTED_RMSNORM_CLASSES: + continue + + module._supa_rmsnorm_original_forward = module.forward + module.forward = types.MethodType(supa_rms_norm_forward, module) + module._supa_rmsnorm_patched = True + patched_count += 1 + patched_classes.add(cls) + + if patched_count: + logger.info_rank0( + "Applied fused SUPA RMSNorm to {} modules: {}.".format( + patched_count, ", ".join(sorted(cls.__name__ for cls in patched_classes)) + ) + ) + + return model