[feat] add biren supa accelerator support (#10844)

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

docker/docker-npu/OVERVIEW.md

@@ -184,6 +184,15 @@ To build an image with Docker Compose without starting a container, use `docker
 - Legacy NPU tags are replaced by the `latest-<910b|a3|950>-<ubuntu|openeuler>` format.
 - Validate the exact driver, firmware, CANN, and SoC combination before production deployment.
 
+## Security Risks
+When running containers with these images, be aware of the following security risks:
+
+- Running as the `root` user: Containers run as the `root` user by default, which can introduce security risks. In production environments, it is recommended to create a non-privileged user to run the application.
+
+- Lack of CPU and memory resource limits: Not setting resource limits may cause a container to consume excessive system resources and affect host performance. It is recommended to use the `--cpus` and `--memory` parameters to limit resource usage.
+
+- Device use with `rwm` permissions: The NPU device is assigned read, write, and `mknod` permissions. Although these are required for functionality, the scope of permissions should be carefully evaluated in security-sensitive environments.
+
 ## License and Disclaimer
 
 LlamaFactory is distributed under the [Apache License 2.0](../../LICENSE).

docker/docker-npu/OVERVIEW.zh.md

@@ -184,6 +184,16 @@ docker compose --profile 950-openeuler up -d
 - 旧式 NPU tag 已由 `latest-<910b|a3|950>-<ubuntu|openeuler>` 格式取代。
 - 正式部署前，请验证具体驱动、固件、CANN 和 SoC 组合的兼容性。
 
+## 安全风险
+
+在运行容器时，需要注意以下安全风险：
+
+- 使用 `root` 用户运行：容器默认以 `root` 用户身份运行，可能带来安全隐患。建议在生产环境中创建非特权用户来运行应用程序。
+
+- 缺少 CPU 和内存资源限制：未设置资源限制可能导致容器消耗过多系统资源，影响宿主机性能。建议使用 `--cpus` 和 `--memory` 参数限制资源使用。
+
+- 设备使用 `rwm` 权限：为 NPU 设备分配了读写和 `mknod` 权限，虽然功能运行需要，但在安全敏感环境中应谨慎评估权限范围。
+
 ## 许可证与免责声明
 
 LlamaFactory 基于 [Apache License 2.0](../../LICENSE) 发布。

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()

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():

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():

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"

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

src/llamafactory/model/patcher.py

@@ -86,12 +86,9 @@ def _check_fla_dependencies() -> None:
 
 
 def patch_qwen3_5_forward_npu(model: "PreTrainedModel") -> None:
-    """Patch for Qwen3.5 models on NPU by importing torch_npu to enable torch.cuda compatibility.
+    """Replace the Qwen3.5 gated delta rule kernel with the NPU-compatible implementation.
 
-    On NPU, torch.cuda operations will fail unless torch_npu is imported.
+    Skip the patch unless an Ascend910-series device is detected and triton_ascend is installed.
-    torch_npu provides compatibility layer that maps torch.cuda calls to NPU operations.
-
-    Also replaces chunk_gated_delta_rule with NPU-compatible implementation.
     """
     import importlib.metadata
 

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():

src/llamafactory/train/hyper_parallel/loader.py

@@ -0,0 +1,162 @@
+# Copyright 2026 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.
+
+"""HyperParallel-specific meta model construction and sharded weight loading."""
+
+import os
+from typing import TYPE_CHECKING
+
+import torch
+from hyper_parallel import init_empty_weights
+from hyper_parallel.integration.llamafactory import parallelize_model
+from transformers import (
+    AutoModelForCausalLM,
+    AutoModelForImageTextToText,
+    AutoModelForSeq2SeqLM,
+    AutoModelForTextToWaveform,
+)
+from transformers.modeling_utils import ContextManagers
+
+from ...extras import logging
+from ...extras.misc import count_parameters
+from ...extras.packages import is_torch_version_greater_than
+from ...model.adapter import init_adapter
+from ...model.loader import _get_init_kwargs, load_config
+from ...model.model_utils.liger_kernel import apply_liger_kernel
+from ...model.model_utils.misc import register_autoclass
+from ...model.patcher import patch_config, patch_model
+
+
+if TYPE_CHECKING:
+    from transformers import PreTrainedModel, PreTrainedTokenizer
+
+    from ...hparams import FinetuningArguments, ModelArguments
+
+
+logger = logging.get_logger(__name__)
+
+
+def _get_no_init_weights():
+    """Resolve the Transformers no-init context across supported versions."""
+    try:
+        from transformers.modeling_utils import no_init_weights
+    except ImportError:
+        from transformers.initialization import no_init_weights
+
+    return no_init_weights
+
+
+def _get_model_class(config):
+    """Select the same Transformers AutoModel class as the native loader."""
+    if type(config) in AutoModelForImageTextToText._model_mapping.keys():
+        return AutoModelForImageTextToText
+    if type(config) in AutoModelForSeq2SeqLM._model_mapping.keys():
+        return AutoModelForSeq2SeqLM
+    if type(config) in AutoModelForTextToWaveform._model_mapping.keys():
+        return AutoModelForTextToWaveform
+    return AutoModelForCausalLM
+
+
+def _validate_conv3d_compatibility(model: "PreTrainedModel", model_args: "ModelArguments") -> None:
+    """Retain the native loader's torch 2.9 Conv3D compatibility guard."""
+    if not is_torch_version_greater_than("2.9.0") or is_torch_version_greater_than("2.10.0"):
+        return
+
+    conv3d_modules = [module for module in model.modules() if isinstance(module, torch.nn.Conv3d)]
+    kt_conv3d_ready = (
+        model_args.use_kt
+        and bool(conv3d_modules)
+        and all(getattr(module, "_kt_conv3d_compatible", False) for module in conv3d_modules)
+    )
+    if conv3d_modules and not kt_conv3d_ready:
+        raise ValueError(
+            "Unsupported torch version detected: torch 2.9.x with Conv3D. "
+            "This combination is known to cause severe performance regression. "
+            "Please downgrade torch to <2.9 or remove Conv3D. "
+            "See https://github.com/pytorch/pytorch/issues/166122"
+        )
+    if kt_conv3d_ready:
+        logger.info_rank0("Using KTransformers instance-scoped Conv3D fallback for torch 2.9.x VLM training.")
+
+
+def load_hyper_parallel_model(
+    tokenizer: "PreTrainedTokenizer",
+    model_args: "ModelArguments",
+    finetuning_args: "FinetuningArguments",
+    distributed_setup,
+    hp_args,
+    is_trainable: bool = False,
+) -> "PreTrainedModel":
+    """Build on meta, apply HyperParallel layouts, then load local weight shards."""
+    init_kwargs = _get_init_kwargs(model_args)
+    config = load_config(model_args)
+    patch_config(config, tokenizer, model_args, init_kwargs, is_trainable)
+    apply_liger_kernel(config, model_args, is_trainable, require_logits=False)
+
+    load_class = _get_model_class(config)
+    no_init_weights = _get_no_init_weights()
+    with ContextManagers([no_init_weights(), init_empty_weights()]):
+        model = load_class.from_config(
+            config,
+            trust_remote_code=model_args.trust_remote_code,
+            dtype=model_args.compute_dtype,
+        )
+
+    if getattr(model.config, "model_type", None) in ["qwen2_5_omni", "qwen3_omni_moe"]:
+        model = getattr(model, "thinker")
+
+    patch_model(model, tokenizer, model_args, is_trainable, add_valuehead=False)
+    register_autoclass(config, model, tokenizer)
+    model = init_adapter(config, model, model_args, finetuning_args, is_trainable)
+
+    if model_args.use_v1_kernels and is_trainable:
+        logger.warning_rank0(
+            "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_v1_kernels
+
+        model = apply_v1_kernels(model, use_v1_kernels=model_args.use_v1_kernels)
+
+    activation_mode = hp_args.activation_mode
+    model = parallelize_model(
+        model,
+        distributed_setup,
+        pretrained_path=None if model_args.train_from_scratch else model_args.model_name_or_path,
+        activation_checkpoint="full" if activation_mode in {"recompute", "swap"} else None,
+        swap_inputs=activation_mode == "swap" and hp_args.activation_swap_inputs,
+    )
+
+    _validate_conv3d_compatibility(model, model_args)
+    if not is_trainable:
+        model.requires_grad_(False)
+        model.eval()
+    else:
+        model.train()
+
+    trainable_params, all_param = count_parameters(model)
+    if is_trainable:
+        param_stats = (
+            f"trainable params: {trainable_params:,} || "
+            f"all params: {all_param:,} || trainable%: {100 * trainable_params / all_param:.4f}"
+        )
+    else:
+        param_stats = f"all params: {all_param:,}"
+    logger.info_rank0(param_stats)
+
+    if model_args.print_param_status and int(os.getenv("LOCAL_RANK", "0")) == 0:
+        for name, param in model.named_parameters():
+            print(f"name: {name}, dtype: {param.dtype}, device: {param.device}, trainable: {param.requires_grad}")
+
+    return model

src/llamafactory/train/hyper_parallel/model_registry.py

@@ -0,0 +1,125 @@
+# Copyright 2026 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.
+
+"""Model-specific HyperParallel plans owned by the LlamaFactory backend."""
+
+from collections.abc import Callable
+from typing import Any
+
+
+ParallelPlan = list[dict[str, Any]]
+ParallelPlanProvider = Callable[[], ParallelPlan]
+
+
+def _get_qwen3_moe_parallel_plan() -> ParallelPlan:
+    return [
+        {
+            "match": "*.self_attn",
+            "when": "cp",
+            "inner_target": "self",
+            "region_dispatch": False,
+            "inner_wrapper": {
+                "_target_": (
+                    "hyper_parallel.models.qwen3_moe.adapter.distributed.context_parallel_async."
+                    "qwen3_moe_async_ulysses_cp_wrapper"
+                )
+            },
+        },
+        {
+            "match": "*.mlp",
+            "when": "ep",
+            "region_dispatch": False,
+            "local_compute_fn": {
+                "_target_": (
+                    "hyper_parallel.models.qwen3_moe.adapter.distributed.expert_parallel.qwen3moe_ep_compute_fn"
+                ),
+                "use_grouped_gemm": True,
+            },
+        },
+    ]
+
+
+def _get_qwen3_vl_moe_parallel_plan() -> ParallelPlan:
+    replicated_vision_wrapper = (
+        "hyper_parallel.models.qwen3_vl_moe.adapter.distributed.context_parallel."
+        "qwen3_vl_moe_replicated_vision_cp_wrapper"
+    )
+    return [
+        *[
+            {
+                "match": match,
+                "when": "cp",
+                "params": {},
+                "inner_target": "self",
+                "region_dispatch": False,
+                "inner_wrapper": {"_target_": replicated_vision_wrapper},
+            }
+            for match in (
+                "model.visual.blocks.*.attn",
+                "model.visual.merger",
+                "model.visual.deepstack_merger_list",
+            )
+        ],
+        {
+            "match": "model.language_model",
+            "when": "cp",
+            "params": {},
+            "inner_target": "self",
+            "region_dispatch": False,
+            "inner_wrapper": {
+                "_target_": (
+                    "hyper_parallel.models.qwen3_vl_moe.adapter.distributed.context_parallel."
+                    "qwen3_vl_moe_text_input_cp_wrapper"
+                )
+            },
+        },
+        {
+            "match": "model.language_model.layers.*.self_attn",
+            "when": "cp",
+            "inner_target": "self",
+            "region_dispatch": False,
+            "inner_wrapper": {
+                "_target_": (
+                    "hyper_parallel.models.qwen3_vl_moe.adapter.distributed.context_parallel."
+                    "qwen3_vl_moe_async_ulysses_cp_wrapper"
+                )
+            },
+        },
+        {
+            "match": "model.language_model.layers.*.mlp",
+            "when": "ep",
+            "region_dispatch": False,
+            "local_compute_fn": {
+                "_target_": (
+                    "hyper_parallel.models.qwen3_vl_moe.adapter.distributed.expert_parallel.qwen3_vl_moe_ep_compute_fn"
+                ),
+                "use_grouped_gemm": True,
+            },
+        },
+    ]
+
+
+MODEL_PARALLEL_PLAN_REGISTRY: dict[str, ParallelPlanProvider] = {
+    "qwen3_moe": _get_qwen3_moe_parallel_plan,
+    "qwen3_vl_moe": _get_qwen3_vl_moe_parallel_plan,
+}
+
+
+def apply_model_parallel_plan(hp_args: Any, model_type: str | None) -> None:
+    """Prepend LlamaFactory defaults while preserving explicit user overrides."""
+    provider = MODEL_PARALLEL_PLAN_REGISTRY.get(model_type or "")
+    if provider is None:
+        return
+
+    hp_args.plan_overrides = [*provider(), *(hp_args.plan_overrides or [])]

src/llamafactory/train/hyper_parallel/trainer.py

@@ -18,15 +18,18 @@ import logging
 import os
 import types
 from contextlib import nullcontext
-from functools import partial
+from functools import partial, wraps
 from typing import Any, Optional
 
 import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+from hyper_parallel.core.optimizer import ChainedOptimizer, get_hyper_optimizer
+from hyper_parallel.data.parallel import shard_batch_for_cp
 from hyper_parallel.integration.llamafactory import (
     HSDPModule,
     HyperParallelArguments,
     export_to_hf_format,
-    fsdp2_prepare_model,
     hsdp_sync_stream,
     load_hsdp_model,
     load_hsdp_optimizer_and_scheduler,
@@ -36,14 +39,6 @@ from hyper_parallel.integration.llamafactory import (
 from hyper_parallel.integration.llamafactory import (
     clip_grad_norm_ as hp_clip_grad_norm_,
 )
-from hyper_parallel.integration.llamafactory.context_parallel import (
-    cp_prepare_model,
-    get_cp_rank,
-    get_dp_rank,
-    shard_inputs_for_cp,
-)
-from hyper_parallel.integration.llamafactory.expert_parallel import ep_prepare_model
-from hyper_parallel.platform import get_platform
 from torch import nn
 
 from ..sft.trainer import CustomSeq2SeqTrainer
@@ -52,6 +47,220 @@ from ..sft.trainer import CustomSeq2SeqTrainer
 logger = logging.getLogger(__name__)
 
 
+class _TrainerOptimizerAdapter(torch.optim.Optimizer):
+    """Expose a Hyper optimizer chain through PyTorch's Trainer protocol."""
+
+    def __init__(self, optimizer: ChainedOptimizer):
+        self.optimizer = optimizer
+        super().__init__(optimizer.param_groups, optimizer.defaults)
+        optimizer.param_groups = self.param_groups
+        self.optimizers_dict = optimizer.optimizers_dict
+
+    def step(self, closure=None):
+        return self.optimizer.step(closure=closure)
+
+    def zero_grad(self, set_to_none: bool = True):
+        return self.optimizer.zero_grad(set_to_none=set_to_none)
+
+    def state_dict(self):
+        return self.optimizer.state_dict()
+
+    def load_state_dict(self, state_dict):
+        return self.optimizer.load_state_dict(state_dict)
+
+
+def _create_hyper_muon_optimizer(model: nn.Module, training_args) -> torch.optim.Optimizer:
+    """Create HyperParallel's distributed Muon and AdamW optimizer chain."""
+    muon_params = []
+    adamw_params = []
+    for name, param in model.named_parameters():
+        if not param.requires_grad:
+            continue
+
+        if param.ndim == 2 and "embed" not in name and "lm_head" not in name:
+            muon_params.append(param)
+        else:
+            adamw_params.append(param)
+
+    optimizer = get_hyper_optimizer(
+        model=model,
+        muon_params=[{"params": muon_params}] if muon_params else [],
+        adamw_params=[{"params": adamw_params}] if adamw_params else [],
+        muon_kwargs={
+            "muon_lr": training_args.learning_rate,
+            "muon_weight_decay": training_args.weight_decay,
+        },
+        adamw_kwargs={
+            "adamw_lr": training_args.learning_rate,
+            "adamw_weight_decay": training_args.weight_decay,
+            "adamw_betas": (training_args.adam_beta1, training_args.adam_beta2),
+            "adamw_eps": training_args.adam_epsilon,
+        },
+    )
+    logger.info(
+        "Using HyperParallel Muon optimizer with %d Muon params and %d AdamW params.",
+        len(muon_params),
+        len(adamw_params),
+    )
+    return _TrainerOptimizerAdapter(optimizer)
+
+
+def _shard_inputs_for_cp(inputs: dict[str, Any], cp_mesh) -> dict[str, Any]:
+    """Adapt HuggingFace batches to Hyper's existing CP batch contract."""
+    inputs = dict(inputs)
+    seq_len = inputs["input_ids"].shape[1]
+    if "position_ids" not in inputs:
+        position_ids = torch.arange(
+            seq_len,
+            device=inputs["input_ids"].device,
+            dtype=torch.long,
+        ).unsqueeze(0)
+        inputs["position_ids"] = position_ids.expand(inputs["input_ids"].shape[0], -1)
+
+    global_attention_mask = inputs.get("attention_mask")
+    global_labels = inputs.get("labels")
+    preserve_attention_mask = (
+        isinstance(global_attention_mask, torch.Tensor)
+        and global_attention_mask.ndim == 2
+        and global_attention_mask.shape[-1] == seq_len
+    )
+    sharded = shard_batch_for_cp(inputs, cp_mesh)
+    if preserve_attention_mask:
+        pad_len = (-seq_len) % (cp_mesh.size() * 2)
+        if pad_len > 0:
+            global_attention_mask = torch.nn.functional.pad(global_attention_mask, (0, pad_len), value=0)
+        sharded["attention_mask"] = global_attention_mask
+    if isinstance(global_labels, torch.Tensor) and global_labels.shape[-1] == seq_len:
+        sharded["labels"] = global_labels
+    return sharded
+
+
+def _pad_inputs_for_model_cp(inputs: dict[str, Any], cp_mesh) -> dict[str, Any]:
+    """Pad only token-aligned fields when a multimodal model shards internally."""
+    inputs = dict(inputs)
+    seq_len = inputs["input_ids"].shape[1]
+    pad_len = (-seq_len) % (cp_mesh.size() * 2)
+    if pad_len == 0:
+        return inputs
+
+    pad_values = {
+        "input_ids": 0,
+        "labels": -100,
+        "attention_mask": 0,
+        "mm_token_type_ids": 0,
+    }
+    for name in ("input_ids", "labels", "attention_mask", "mm_token_type_ids", "position_ids"):
+        value = inputs.get(name)
+        if not isinstance(value, torch.Tensor) or value.shape[-1] != seq_len:
+            continue
+        if name == "position_ids":
+            increment = torch.arange(1, pad_len + 1, device=value.device, dtype=value.dtype)
+            increment = increment.reshape(*([1] * (value.ndim - 1)), pad_len)
+            pad = value[..., -1:] + increment
+            pad = pad.expand(*value.shape[:-1], pad_len)
+        else:
+            pad = value.new_full((*value.shape[:-1], pad_len), pad_values[name])
+        inputs[name] = torch.cat((value, pad), dim=-1)
+    return inputs
+
+
+def _build_cp_shift_labels(
+    labels: torch.Tensor,
+    local_seq_len: int,
+    cp_rank: int,
+    ignore_index: int,
+) -> torch.Tensor:
+    """Align full-sequence causal labels with one rank's local logits."""
+    start = cp_rank * local_seq_len
+    end = start + local_seq_len
+    shift_labels = labels[..., start + 1 : min(end + 1, labels.shape[-1])]
+    if shift_labels.shape[-1] < local_seq_len:
+        pad_shape = (*shift_labels.shape[:-1], local_seq_len - shift_labels.shape[-1])
+        shift_labels = torch.cat(
+            (shift_labels, shift_labels.new_full(pad_shape, ignore_index)),
+            dim=-1,
+        )
+    return shift_labels.contiguous()
+
+
+def _enable_cp_causal_lm_loss(model: nn.Module, cp_mesh) -> None:
+    """Use full labels to preserve causal boundaries and CP-global normalization."""
+    cp_size = cp_mesh.size()
+    cp_rank = cp_mesh.get_local_rank()
+    cp_group = cp_mesh.get_group()
+    original_loss_function = model.loss_function
+
+    @wraps(original_loss_function)
+    def cp_loss_function(
+        logits,
+        labels,
+        vocab_size: int,
+        num_items_in_batch=None,
+        ignore_index: int = -100,
+        shift_labels=None,
+        **kwargs,
+    ):
+        if (
+            shift_labels is not None
+            or not isinstance(logits, torch.Tensor)
+            or not isinstance(labels, torch.Tensor)
+            or logits.ndim < 3
+            or labels.ndim < 2
+            or logits.shape[-2] * cp_size != labels.shape[-1]
+        ):
+            return original_loss_function(
+                logits=logits,
+                labels=labels,
+                vocab_size=vocab_size,
+                num_items_in_batch=num_items_in_batch,
+                ignore_index=ignore_index,
+                shift_labels=shift_labels,
+                **kwargs,
+            )
+
+        local_shift_labels = _build_cp_shift_labels(
+            labels,
+            local_seq_len=logits.shape[-2],
+            cp_rank=cp_rank,
+            ignore_index=ignore_index,
+        ).to(logits.device)
+        local_logits = logits[..., :vocab_size].contiguous().float()
+        local_loss_sum = F.cross_entropy(
+            local_logits.view(-1, vocab_size),
+            local_shift_labels.view(-1),
+            ignore_index=ignore_index,
+            reduction="sum",
+        )
+        if num_items_in_batch is not None:
+            denominator = num_items_in_batch
+            if isinstance(denominator, torch.Tensor):
+                denominator = denominator.to(device=local_loss_sum.device, dtype=local_loss_sum.dtype)
+            # Trainer counts the repeated full labels on every CP rank. Remove
+            # that duplication here; Trainer applies its normal world-size
+            # compensation after this loss function returns.
+            return (
+                local_loss_sum
+                * cp_size
+                / torch.clamp(
+                    torch.as_tensor(denominator, device=local_loss_sum.device),
+                    min=1,
+                )
+            )
+
+        cp_num_items = local_shift_labels.ne(ignore_index).sum()
+        dist.all_reduce(cp_num_items, group=cp_group)
+        denominator = torch.clamp(cp_num_items.to(local_loss_sum.dtype), min=1)
+
+        # Hyper FSDP averages gradients over the flattened DP x CP mesh.
+        backward_loss = local_loss_sum * cp_size / denominator
+        group_loss_sum = local_loss_sum.detach().clone()
+        dist.all_reduce(group_loss_sum, group=cp_group)
+        logging_loss = group_loss_sum / denominator
+        return backward_loss + (logging_loss - backward_loss).detach()
+
+    model.loss_function = cp_loss_function
+
+
 class _CPBatchRepeatedBatchSampler(torch.utils.data.BatchSampler):
     """Repeat logical batches so Accelerate shards CP peers onto the same samples."""
 
@@ -143,16 +352,13 @@ class HyperParallelTrainer(CustomSeq2SeqTrainer):
     def __init__(
         self,
         hp_args: HyperParallelArguments,
+        distributed_setup,
         finetuning_args=None,
         processor=None,
-        ref_model: Optional[nn.Module] = None,
         **kwargs,
     ):
-        self._hp_args = hp_args
+        # Keep LlamaFactory's Trainer lifecycle; Hyper owns only the already
-
+        # constructed model's distributed layout and checkpoint state.
-        # Let CustomSeq2SeqTrainer handle everything except ref_model —
-        # Custom would prepare it with accelerate's fsdp2_prepare_model,
-        # but we need HP's version instead.
         super().__init__(
             finetuning_args=finetuning_args,
             processor=processor,
@@ -164,40 +370,48 @@ class HyperParallelTrainer(CustomSeq2SeqTrainer):
             raise ValueError("HyperParallel trainer requires Accelerate FSDP2 mode to be enabled.")
 
         self._cp_size = hp_args.cp_size
-        self._ep_size = hp_args.ep_size
+        self._cp_mesh = distributed_setup.mesh_context.cp_mesh
-        self._cp_rank = get_cp_rank(hp_args) if self._cp_size > 1 else 0
+        if self._cp_size > 1 and self._cp_mesh is None:
-        self._dp_rank = get_dp_rank(hp_args) if self._cp_size > 1 else get_platform().get_rank()
+            raise RuntimeError("HyperParallel CP requires a cp_mesh from DistributedSetup.")
-
+        if self._cp_size > 1:
-        # Prepare ref_model with the same CP + EP + HSDP path as the train model.
+            _enable_cp_causal_lm_loss(self.model, self._cp_mesh)
-        self.ref_model = ref_model
+        self._model_shards_cp_inputs = any(
-        if self.ref_model is not None:
+            getattr(module, "_hp_cp_shards_inputs", False) for module in self.model.modules()
-            self.ref_model = self._prepare_model_for_hyper_parallel(self.ref_model)
+        )
 
         self._orig_accelerator_clip_grad_norm = self.accelerator.clip_grad_norm_
         self._orig_fsdp2_prepare_model = None
+        self._orig_fsdp2_set_auto_wrap_policy = None
         self._accelerator_patches_active = False
 
     def _prepare_model_for_hyper_parallel(self, model: nn.Module) -> nn.Module:
-        """Apply CP/EP preparation before delegating to HyperParallel FSDP2."""
+        """Return a model already prepared by HyperParallel's unified model path."""
-        if self._cp_size > 1:
+        if not isinstance(model, HSDPModule):
-            model = cp_prepare_model(model, self.accelerator, self._hp_args)
+            raise RuntimeError(
-        if self._ep_size > 1:
+                "HyperParallel models must be parallelized during meta initialization before Trainer preparation."
-            model = ep_prepare_model(model, self.accelerator, self._hp_args)
+            )
-        return fsdp2_prepare_model(self.accelerator, model, self._hp_args)
+        return model
 
     def _activate_accelerator_patches(self) -> None:
-        """Patch Accelerate to use HyperParallel fsdp2_prepare_model and clip_grad_norm_."""
+        """Keep Accelerate from rewrapping the prepared model and patch gradient clipping."""
         if self._accelerator_patches_active:
             return
 
         import accelerate.accelerator as acc_module  # pylint: disable=C0415
 
         self._orig_fsdp2_prepare_model = acc_module.fsdp2_prepare_model
+        fsdp_plugin = self.accelerator.state.fsdp_plugin
+        self._orig_fsdp2_set_auto_wrap_policy = fsdp_plugin.set_auto_wrap_policy
 
         def _hp_fsdp2_prepare_model(accelerator, model):
             return self._prepare_model_for_hyper_parallel(model)
 
+        def _hp_set_auto_wrap_policy(plugin, model):
+            del plugin
+            self._prepare_model_for_hyper_parallel(model)
+
         acc_module.fsdp2_prepare_model = _hp_fsdp2_prepare_model
+        fsdp_plugin.set_auto_wrap_policy = types.MethodType(_hp_set_auto_wrap_policy, fsdp_plugin)
 
         def _hp_clip_grad_norm(accelerator, parameters, max_norm, norm_type=2):
             if getattr(accelerator, "is_fsdp2", False):
@@ -224,16 +438,16 @@ class HyperParallelTrainer(CustomSeq2SeqTrainer):
 
         if self._orig_fsdp2_prepare_model is not None:
             acc_module.fsdp2_prepare_model = self._orig_fsdp2_prepare_model
+        if self._orig_fsdp2_set_auto_wrap_policy is not None:
+            self.accelerator.state.fsdp_plugin.set_auto_wrap_policy = self._orig_fsdp2_set_auto_wrap_policy
         self.accelerator.clip_grad_norm_ = self._orig_accelerator_clip_grad_norm
         self._accelerator_patches_active = False
 
     def _wrap_model(self, model: nn.Module, training: bool = True, dataloader=None) -> nn.Module:
-        """Let Accelerate own FSDP2/HSDP wrapping so optimizer remapping stays correct."""
+        """Keep Trainer from wrapping a model already prepared by HyperParallel."""
         del dataloader
         if isinstance(model, HSDPModule):
             return model
-        if training and getattr(self.accelerator, "is_fsdp2", False):
-            return model
         return super()._wrap_model(model, training=training)
 
     def _get_train_sampler(self, train_dataset=None):
@@ -252,7 +466,7 @@ class HyperParallelTrainer(CustomSeq2SeqTrainer):
             batch_size=batch_size,
             drop_last=drop_last,
             repeat_factor=self._cp_size,
-            logical_group_size=max(1, get_platform().get_world_size() // self._cp_size),
+            logical_group_size=max(1, dist.get_world_size() // self._cp_size),
         )
 
     def _get_cp_dataloader(self, dataset, batch_size: int, shuffle: bool):
@@ -281,7 +495,7 @@ class HyperParallelTrainer(CustomSeq2SeqTrainer):
             drop_last=self.args.dataloader_drop_last,
         )
         logical_batches = len(batch_sampler) // self._cp_size
-        dp_size = max(1, get_platform().get_world_size() // self._cp_size)
+        dp_size = max(1, dist.get_world_size() // self._cp_size)
         logical_length = (
             logical_batches // dp_size if self.args.dataloader_drop_last else _ceil_div(logical_batches, dp_size)
         )
@@ -350,7 +564,10 @@ class HyperParallelTrainer(CustomSeq2SeqTrainer):
         inputs = self._prepare_inputs(inputs)
 
         if self._cp_size > 1:
-            inputs = shard_inputs_for_cp(inputs, self._cp_rank, self._cp_size)
+            if self._model_shards_cp_inputs:
+                inputs = _pad_inputs_for_model_cp(inputs, self._cp_mesh)
+            else:
+                inputs = _shard_inputs_for_cp(inputs, self._cp_mesh)
 
         sync_gradients = getattr(self.accelerator, "sync_gradients", True)
         if isinstance(model, HSDPModule):
@@ -372,16 +589,25 @@ class HyperParallelTrainer(CustomSeq2SeqTrainer):
             )
             loss = loss / accumulation_steps
 
-        self.accelerator.backward(loss)
+        # The loss already accounts for the accumulation window, either via
+        # num_items_in_batch or the explicit division above. Compensate the
+        # additional scaling inside Accelerator.backward() so it happens once.
+        backward_loss = loss * self.accelerator.gradient_accumulation_steps
+        self.accelerator.backward(backward_loss)
 
         if isinstance(model, HSDPModule) and sync_gradients:
             hsdp_sync_stream()
 
         return loss.detach()
 
-    def create_optimizer(self):
+    def create_optimizer(self, *args, **kwargs):
         """Create optimizer and wrap step with SkipDTensorDispatch."""
-        optimizer = super().create_optimizer()
+        if self.optimizer is None and getattr(self.finetuning_args, "use_muon", False):
+            model = args[0] if args else kwargs.get("model", self.model)
+            self.optimizer = _create_hyper_muon_optimizer(model, self.args)
+            optimizer = self.optimizer
+        else:
+            optimizer = super().create_optimizer(*args, **kwargs)
         wrap_optimizer_with_skip_dtensor_dispatch(optimizer)
         return optimizer
 

src/llamafactory/train/hyper_parallel/workflow.py

@@ -23,9 +23,12 @@ from ...extras.logging import get_logger
 from ...extras.misc import calculate_tps
 from ...extras.packages import is_hyper_parallel_available, is_transformers_version_greater_than
 from ...extras.ploting import plot_loss
-from ...model import load_model, load_tokenizer
+from ...model import load_tokenizer
+from ...model.loader import load_config
 from ..sft.metric import ComputeAccuracy, ComputeSimilarity, eval_logit_processor
-from ..trainer_utils import create_modelcard_and_push, create_ref_model
+from ..trainer_utils import create_modelcard_and_push
+from .loader import load_hyper_parallel_model
+from .model_registry import apply_model_parallel_plan
 from .trainer import HyperParallelTrainer
 
 
@@ -38,19 +41,38 @@ if TYPE_CHECKING:
 logger = get_logger(__name__)
 
 
-def _prepare_hp_args(finetuning_args: "FinetuningArguments", model_args: "ModelArguments"):
+def _prepare_hp_args(
+    finetuning_args: "FinetuningArguments",
+    model_args: "ModelArguments",
+    training_args: "Seq2SeqTrainingArguments",
+):
     r"""Load HyperParallel arguments and apply LlamaFactory-side overrides.
 
     When activation optimization is enabled, skip native gradient checkpointing
-    so HP can install its own via ``setup_activation_optimization``.
+    so Hyper's unified model builder can install it before FSDP wrapping.
     """
     if not is_hyper_parallel_available():
         raise ImportError("hyper_parallel is not installed. Please install it with `pip install hyper_parallel`.")
 
+    if finetuning_args.finetuning_type != "full":
+        raise ValueError("HyperParallel currently requires full fine-tuning in LlamaFactory.")
+    if model_args.quantization_bit is not None:
+        raise ValueError("HyperParallel meta initialization does not support quantized LlamaFactory models yet.")
+    if model_args.use_unsloth or model_args.mixture_of_depths is not None:
+        raise ValueError("HyperParallel meta initialization is incompatible with Unsloth and mixture-of-depths.")
+    if finetuning_args.use_asft_loss:
+        raise ValueError("HyperParallel does not yet support the additional ASFT reference model.")
+
     from hyper_parallel.integration.llamafactory import HyperParallelArguments  # pylint: disable=C0415
 
     hp_args = HyperParallelArguments.from_finetuning_args(finetuning_args)
 
+    if hp_args.param_dtype is None:
+        if training_args.bf16:
+            hp_args.param_dtype = "bfloat16"
+        elif training_args.fp16:
+            hp_args.param_dtype = "float16"
+
     if getattr(hp_args, "cp_size", None) != finetuning_args.hyper_parallel_cp_size:
         setattr(hp_args, "cp_size", finetuning_args.hyper_parallel_cp_size)
 
@@ -77,17 +99,28 @@ def run_pt(
     finetuning_args: "FinetuningArguments",
     callbacks: Optional[list["TrainerCallback"]] = None,
 ):
-    hp_args = _prepare_hp_args(finetuning_args, model_args)
+    hp_args = _prepare_hp_args(finetuning_args, model_args, training_args)
+    model_type = getattr(load_config(model_args), "model_type", None)
+    apply_model_parallel_plan(hp_args, model_type)
+    distributed_setup = hp_args.build_distributed_setup()
 
     tokenizer_module = load_tokenizer(model_args)
     tokenizer = tokenizer_module["tokenizer"]
     template = get_template_and_fix_tokenizer(tokenizer, data_args)
     dataset_module = get_dataset(template, model_args, data_args, training_args, stage="pt", **tokenizer_module)
-    model = load_model(tokenizer, model_args, finetuning_args, training_args.do_train)
+    model = load_hyper_parallel_model(
+        tokenizer,
+        model_args,
+        finetuning_args,
+        distributed_setup,
+        hp_args,
+        is_trainable=training_args.do_train,
+    )
     data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
 
     trainer = HyperParallelTrainer(
         hp_args=hp_args,
+        distributed_setup=distributed_setup,
         model=model,
         args=training_args,
         finetuning_args=finetuning_args,
@@ -145,17 +178,23 @@ def run_sft(
     generating_args: "GeneratingArguments",
     callbacks: Optional[list["TrainerCallback"]] = None,
 ):
-    hp_args = _prepare_hp_args(finetuning_args, model_args)
+    hp_args = _prepare_hp_args(finetuning_args, model_args, training_args)
+    model_type = getattr(load_config(model_args), "model_type", None)
+    apply_model_parallel_plan(hp_args, model_type)
+    distributed_setup = hp_args.build_distributed_setup()
 
     tokenizer_module = load_tokenizer(model_args)
     tokenizer = tokenizer_module["tokenizer"]
     template = get_template_and_fix_tokenizer(tokenizer, data_args)
     dataset_module = get_dataset(template, model_args, data_args, training_args, stage="sft", **tokenizer_module)
-    model = load_model(tokenizer, model_args, finetuning_args, training_args.do_train)
+    model = load_hyper_parallel_model(
-
+        tokenizer,
-    ref_model = None
+        model_args,
-    if finetuning_args.use_asft_loss:
+        finetuning_args,
-        ref_model = create_ref_model(model_args, finetuning_args)
+        distributed_setup,
+        hp_args,
+        is_trainable=training_args.do_train,
+    )
 
     data_collator = SFTDataCollatorWith4DAttentionMask(
         template=template,
@@ -192,13 +231,13 @@ def run_sft(
 
     trainer = HyperParallelTrainer(
         hp_args=hp_args,
+        distributed_setup=distributed_setup,
         model=model,
         args=training_args,
         finetuning_args=finetuning_args,
         data_collator=data_collator,
         callbacks=callbacks,
         gen_kwargs=gen_kwargs,
-        ref_model=ref_model,
         **dataset_module,
         **tokenizer_module,
         **metric_module,

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:

src/llamafactory/v1/accelerator/helper.py

@@ -45,6 +45,7 @@ class DeviceType(StrEnum):
     META = "meta"
     MPS = "mps"
     NPU = "npu"
+    SUPA = "supa"
     XPU = "xpu"
 
 

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",),
 }
 
 

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

tests/train/test_hyper_parallel.py

@@ -0,0 +1,375 @@
+# Copyright 2026 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.
+
+from contextlib import nullcontext
+from importlib.util import find_spec
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+
+
+if find_spec("hyper_parallel") is None:
+    pytest.skip("HyperParallel is an optional dependency and is not installed.", allow_module_level=True)
+else:
+    import torch
+    from hyper_parallel.core.optimizer import ChainedOptimizer
+
+    from llamafactory.train.hyper_parallel.loader import load_hyper_parallel_model
+    from llamafactory.train.hyper_parallel.model_registry import apply_model_parallel_plan
+    from llamafactory.train.hyper_parallel.trainer import (
+        HyperParallelTrainer,
+        _build_cp_shift_labels,
+        _create_hyper_muon_optimizer,
+        _pad_inputs_for_model_cp,
+        _shard_inputs_for_cp,
+    )
+    from llamafactory.train.hyper_parallel.workflow import _prepare_hp_args
+
+
+def test_create_hyper_muon_optimizer_preserves_llamafactory_parameter_split():
+    class _Model(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.embed_tokens = torch.nn.Embedding(8, 4)
+            self.proj = torch.nn.Linear(4, 4)
+            self.norm = torch.nn.LayerNorm(4)
+            self.lm_head = torch.nn.Linear(4, 8, bias=False)
+
+    model = _Model()
+    training_args = SimpleNamespace(
+        learning_rate=3.0e-4,
+        weight_decay=0.1,
+        adam_beta1=0.8,
+        adam_beta2=0.95,
+        adam_epsilon=1.0e-8,
+    )
+    expected_optimizer = ChainedOptimizer(
+        model,
+        {"adamw": torch.optim.AdamW(model.parameters(), lr=training_args.learning_rate)},
+    )
+
+    with patch(
+        "llamafactory.train.hyper_parallel.trainer.get_hyper_optimizer",
+        return_value=expected_optimizer,
+    ) as get_optimizer:
+        optimizer = _create_hyper_muon_optimizer(model, training_args)
+
+    assert isinstance(optimizer, torch.optim.Optimizer)
+    assert optimizer.optimizer is expected_optimizer
+    call_kwargs = get_optimizer.call_args.kwargs
+    assert call_kwargs["muon_params"][0]["params"] == [model.proj.weight]
+    assert set(call_kwargs["adamw_params"][0]["params"]) == {
+        model.embed_tokens.weight,
+        model.proj.bias,
+        model.norm.weight,
+        model.norm.bias,
+        model.lm_head.weight,
+    }
+    assert call_kwargs["muon_kwargs"] == {"muon_lr": 3.0e-4, "muon_weight_decay": 0.1}
+    assert call_kwargs["adamw_kwargs"] == {
+        "adamw_lr": 3.0e-4,
+        "adamw_weight_decay": 0.1,
+        "adamw_betas": (0.8, 0.95),
+        "adamw_eps": 1.0e-8,
+    }
+    optimizer.param_groups[0]["lr"] = 1.0e-4
+    assert expected_optimizer.chained_optimizers[0].param_groups[0]["lr"] == 1.0e-4
+
+
+def test_cp_batch_adapter_keeps_global_mask_and_generates_local_positions():
+    class _CPMesh:
+        @staticmethod
+        def size():
+            return 2
+
+        @staticmethod
+        def get_local_rank():
+            return 1
+
+    inputs = {
+        "input_ids": torch.arange(8).unsqueeze(0),
+        "labels": torch.arange(8).unsqueeze(0),
+        "attention_mask": torch.ones(1, 8, dtype=torch.long),
+    }
+
+    sharded = _shard_inputs_for_cp(inputs, _CPMesh())
+
+    torch.testing.assert_close(sharded["input_ids"], torch.arange(4, 8).unsqueeze(0))
+    torch.testing.assert_close(sharded["labels"], inputs["labels"])
+    torch.testing.assert_close(sharded["position_ids"], torch.arange(4, 8).unsqueeze(0))
+    torch.testing.assert_close(sharded["attention_mask"], inputs["attention_mask"])
+
+
+def test_cp_shift_labels_preserve_cross_rank_causal_boundary():
+    labels = torch.tensor([[10, 11, 12, 13, 14, 15, 16, 17]])
+
+    rank0_labels = _build_cp_shift_labels(labels, local_seq_len=4, cp_rank=0, ignore_index=-100)
+    rank1_labels = _build_cp_shift_labels(labels, local_seq_len=4, cp_rank=1, ignore_index=-100)
+
+    torch.testing.assert_close(rank0_labels, torch.tensor([[11, 12, 13, 14]]))
+    torch.testing.assert_close(rank1_labels, torch.tensor([[15, 16, 17, -100]]))
+
+
+def test_model_cp_padding_preserves_multimodal_fields():
+    class _CPMesh:
+        @staticmethod
+        def size():
+            return 2
+
+    pixel_values = torch.randn(6, 12)
+    image_grid_thw = torch.tensor([[1, 2, 3]])
+    inputs = {
+        "input_ids": torch.arange(6).unsqueeze(0),
+        "labels": torch.arange(6).unsqueeze(0),
+        "attention_mask": torch.ones(1, 6, dtype=torch.long),
+        "mm_token_type_ids": torch.zeros(1, 6, dtype=torch.long),
+        "pixel_values": pixel_values,
+        "image_grid_thw": image_grid_thw,
+    }
+
+    padded = _pad_inputs_for_model_cp(inputs, _CPMesh())
+
+    assert padded["input_ids"].shape[-1] == 8
+    torch.testing.assert_close(padded["labels"][..., -2:], torch.full((1, 2), -100))
+    torch.testing.assert_close(padded["attention_mask"][..., -2:], torch.zeros(1, 2, dtype=torch.long))
+    assert padded["pixel_values"] is pixel_values
+    assert padded["image_grid_thw"] is image_grid_thw
+
+
+def test_accelerate_patches_skip_torch_auto_wrap_for_prepared_hyper_model():
+    class _FSDPPlugin:
+        def __init__(self):
+            self.calls = []
+
+        def set_auto_wrap_policy(self, model):
+            self.calls.append(model)
+
+    plugin = _FSDPPlugin()
+    original_clip = object()
+    trainer = object.__new__(HyperParallelTrainer)
+    trainer.accelerator = SimpleNamespace(
+        state=SimpleNamespace(fsdp_plugin=plugin),
+        clip_grad_norm_=original_clip,
+    )
+    trainer._orig_accelerator_clip_grad_norm = original_clip
+    trainer._orig_fsdp2_prepare_model = None
+    trainer._orig_fsdp2_set_auto_wrap_policy = None
+    trainer._accelerator_patches_active = False
+    prepared_models = []
+    trainer._prepare_model_for_hyper_parallel = lambda model: prepared_models.append(model) or model
+
+    trainer._activate_accelerator_patches()
+    model = object()
+    plugin.set_auto_wrap_policy(model)
+    trainer._restore_accelerator_patches()
+
+    assert prepared_models == [model]
+    assert plugin.calls == []
+    assert plugin.set_auto_wrap_policy == trainer._orig_fsdp2_set_auto_wrap_policy
+
+
+def test_training_step_compensates_accelerate_gradient_accumulation_scaling():
+    for model_accepts_loss_kwargs, expected_loss in ((True, 8.0), (False, 4.0)):
+        model = torch.nn.Linear(1, 1, bias=False)
+        torch.nn.init.ones_(model.weight)
+        trainer = object.__new__(HyperParallelTrainer)
+        trainer._cp_size = 1
+        trainer.args = SimpleNamespace(n_gpu=1, gradient_accumulation_steps=2)
+        trainer.current_gradient_accumulation_steps = 2
+        trainer.model_accepts_loss_kwargs = model_accepts_loss_kwargs
+        trainer.compute_loss_func = None
+        trainer._prepare_inputs = lambda inputs: inputs
+        trainer.compute_loss_context_manager = nullcontext
+        trainer.compute_loss = lambda model, inputs, num_items_in_batch=None: model.weight.sum() * 8
+        trainer.accelerator = SimpleNamespace(
+            sync_gradients=True,
+            gradient_accumulation_steps=2,
+            backward=lambda loss: (loss / 2).backward(),
+        )
+
+        logged_loss = trainer.training_step(model, {}, num_items_in_batch=1)
+
+        torch.testing.assert_close(logged_loss, torch.tensor(expected_loss))
+        torch.testing.assert_close(model.weight.grad, torch.full_like(model.weight, expected_loss))
+
+
+def test_hyper_model_loader_uses_meta_then_parallelize_path():
+    config = SimpleNamespace(model_type="test")
+    model = torch.nn.Linear(2, 2, device="meta")
+    model.config = config
+    calls = []
+
+    def _apply_v1_kernels(input_model, use_v1_kernels):
+        assert input_model is model
+        assert use_v1_kernels is True
+        calls.append("v1")
+        return input_model
+
+    def _parallelize_model(input_model, *args, **kwargs):
+        assert input_model is model
+        calls.append("parallelize")
+        return input_model
+
+    class _AutoModel:
+        @staticmethod
+        def from_config(input_config, trust_remote_code=False, dtype=None):
+            assert input_config is config
+            assert trust_remote_code is False
+            assert dtype is torch.bfloat16
+            return model
+
+    model_args = SimpleNamespace(
+        model_name_or_path="checkpoint",
+        train_from_scratch=False,
+        trust_remote_code=False,
+        use_kt=False,
+        use_v1_kernels=True,
+        print_param_status=False,
+        compute_dtype=torch.bfloat16,
+    )
+    finetuning_args = SimpleNamespace(stage="sft")
+    hp_args = SimpleNamespace(activation_mode="none", activation_swap_inputs=True)
+    distributed_setup = object()
+
+    with (
+        patch("llamafactory.train.hyper_parallel.loader._get_init_kwargs", return_value={}),
+        patch("llamafactory.train.hyper_parallel.loader.load_config", return_value=config),
+        patch("llamafactory.train.hyper_parallel.loader.patch_config"),
+        patch("llamafactory.train.hyper_parallel.loader.apply_liger_kernel"),
+        patch("llamafactory.train.hyper_parallel.loader._get_model_class", return_value=_AutoModel),
+        patch("llamafactory.train.hyper_parallel.loader._get_no_init_weights", return_value=nullcontext),
+        patch("llamafactory.train.hyper_parallel.loader.init_empty_weights", return_value=nullcontext()),
+        patch("llamafactory.train.hyper_parallel.loader.patch_model"),
+        patch("llamafactory.train.hyper_parallel.loader.register_autoclass"),
+        patch("llamafactory.train.hyper_parallel.loader.init_adapter", return_value=model),
+        patch(
+            "llamafactory.v1.plugins.model_plugins.kernels.interface.apply_v1_kernels",
+            side_effect=_apply_v1_kernels,
+        ) as apply_v1_kernels,
+        patch(
+            "llamafactory.train.hyper_parallel.loader.parallelize_model",
+            side_effect=_parallelize_model,
+        ) as parallelize,
+        patch("llamafactory.train.hyper_parallel.loader._validate_conv3d_compatibility"),
+    ):
+        result = load_hyper_parallel_model(
+            tokenizer=object(),
+            model_args=model_args,
+            finetuning_args=finetuning_args,
+            distributed_setup=distributed_setup,
+            hp_args=hp_args,
+            is_trainable=True,
+        )
+
+    assert result is model
+    assert calls == ["v1", "parallelize"]
+    apply_v1_kernels.assert_called_once_with(model, use_v1_kernels=True)
+    parallelize.assert_called_once_with(
+        model,
+        distributed_setup,
+        pretrained_path="checkpoint",
+        activation_checkpoint=None,
+        swap_inputs=False,
+    )
+
+
+def _make_hyper_workflow_args(finetuning_type="full"):
+    finetuning_args = SimpleNamespace(
+        hyper_parallel_args=None,
+        hyper_parallel_cp_size=1,
+        hyper_parallel_ep_size=1,
+        hyper_parallel_efsdp_size=1,
+        hyper_parallel_token_dispatcher="all_to_all",
+        finetuning_type=finetuning_type,
+        use_asft_loss=False,
+    )
+    model_args = SimpleNamespace(
+        disable_gradient_checkpointing=False,
+        quantization_bit=None,
+        use_unsloth=False,
+        mixture_of_depths=None,
+    )
+    return finetuning_args, model_args
+
+
+def test_prepare_hp_args_accepts_supported_full_tuning():
+    finetuning_args, model_args = _make_hyper_workflow_args()
+    training_args = SimpleNamespace(bf16=True, fp16=False)
+
+    with patch("hyper_parallel.integration.llamafactory.utils.dist.get_world_size", return_value=1):
+        hp_args = _prepare_hp_args(finetuning_args, model_args, training_args)
+
+    assert hp_args.cp_size == 1
+    assert hp_args.ep_size == 1
+    assert hp_args.param_dtype == "bfloat16"
+
+
+def test_prepare_hp_args_rejects_parameter_efficient_tuning():
+    finetuning_args, model_args = _make_hyper_workflow_args(finetuning_type="lora")
+    training_args = SimpleNamespace(bf16=True, fp16=False)
+
+    try:
+        _prepare_hp_args(finetuning_args, model_args, training_args)
+    except ValueError as exc:
+        assert "requires full fine-tuning" in str(exc)
+    else:
+        raise AssertionError("Expected HyperParallel to reject LoRA tuning.")
+
+
+def test_qwen3_moe_parallel_plan_resolves_through_hyper_builder():
+    from hyper_parallel.integration.llamafactory import HyperParallelArguments
+
+    hp_args = HyperParallelArguments(
+        cp_size=2,
+        ep_size=4,
+        device_type="npu",
+        plan_overrides=[{"match": "*.mlp", "region_dispatch": True}],
+    )
+    apply_model_parallel_plan(hp_args, "qwen3_moe")
+
+    with (
+        patch("hyper_parallel.integration.llamafactory.utils.dist.get_world_size", return_value=8),
+        patch("hyper_parallel.integration.llamafactory.utils.MeshContext.build_meshs"),
+    ):
+        setup = hp_args.build_distributed_setup()
+
+    cp_spec = setup.plan_overrides["*.self_attn"]
+    ep_spec = setup.plan_overrides["*.mlp"]
+    assert cp_spec.inner_wrapper.callable.__name__ == "qwen3_moe_async_ulysses_cp_wrapper"
+    assert ep_spec.local_compute_fn.callable.__name__ == "qwen3moe_ep_compute_fn"
+    assert ep_spec.local_compute_fn.use_grouped_gemm is True
+    assert ep_spec.region_dispatch is True
+    assert not setup.module_replacements
+
+
+def test_qwen3_vl_moe_parallel_plan_only_uses_its_model_adapter():
+    hp_args = SimpleNamespace(plan_overrides=None)
+
+    apply_model_parallel_plan(hp_args, "qwen3_vl_moe")
+
+    targets = str(hp_args.plan_overrides)
+    assert "hyper_parallel.models.qwen3_vl_moe.adapter" in targets
+    assert "hyper_parallel.models.qwen3_moe.adapter" not in targets
+    assert {entry["match"] for entry in hp_args.plan_overrides} == {
+        "model.visual.blocks.*.attn",
+        "model.visual.merger",
+        "model.visual.deepstack_merger_list",
+        "model.language_model",
+        "model.language_model.layers.*.self_attn",
+        "model.language_model.layers.*.mlp",
+    }
+    for entry in hp_args.plan_overrides:
+        if entry["match"] in {"model.visual.merger", "model.visual.deepstack_merger_list"}:
+            assert entry["params"] == {}