feat: adopt HyperParallel unified model setup (#10877)

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,

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"] == {}