From 438b3611d1538547876c9e7af044b022dc0e53f2 Mon Sep 17 00:00:00 2001 From: Cui-yshoho <73014084+Cui-yshoho@users.noreply.github.com> Date: Tue, 20 Oct 2026 02:17:27 +0800 Subject: [PATCH] feat: adopt HyperParallel unified model setup (#10877) --- .../train/hyper_parallel/loader.py | 162 ++++++++ .../train/hyper_parallel/model_registry.py | 125 ++++++ .../train/hyper_parallel/trainer.py | 306 ++++++++++++-- .../train/hyper_parallel/workflow.py | 65 ++- tests/train/test_hyper_parallel.py | 375 ++++++++++++++++++ 5 files changed, 980 insertions(+), 53 deletions(-) create mode 100644 src/llamafactory/train/hyper_parallel/loader.py create mode 100644 src/llamafactory/train/hyper_parallel/model_registry.py create mode 100644 tests/train/test_hyper_parallel.py diff --git a/src/llamafactory/train/hyper_parallel/loader.py b/src/llamafactory/train/hyper_parallel/loader.py new file mode 100644 index 000000000..53176e060 --- /dev/null +++ b/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 diff --git a/src/llamafactory/train/hyper_parallel/model_registry.py b/src/llamafactory/train/hyper_parallel/model_registry.py new file mode 100644 index 000000000..d7937209d --- /dev/null +++ b/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 [])] diff --git a/src/llamafactory/train/hyper_parallel/trainer.py b/src/llamafactory/train/hyper_parallel/trainer.py index c05099622..dac2f11ae 100644 --- a/src/llamafactory/train/hyper_parallel/trainer.py +++ b/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 - - # Let CustomSeq2SeqTrainer handle everything except ref_model — - # Custom would prepare it with accelerate's fsdp2_prepare_model, - # but we need HP's version instead. + # Keep LlamaFactory's Trainer lifecycle; Hyper owns only the already + # constructed model's distributed layout and checkpoint state. 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_rank = get_cp_rank(hp_args) if self._cp_size > 1 else 0 - self._dp_rank = get_dp_rank(hp_args) if self._cp_size > 1 else get_platform().get_rank() - - # Prepare ref_model with the same CP + EP + HSDP path as the train model. - self.ref_model = ref_model - if self.ref_model is not None: - self.ref_model = self._prepare_model_for_hyper_parallel(self.ref_model) + self._cp_mesh = distributed_setup.mesh_context.cp_mesh + if self._cp_size > 1 and self._cp_mesh is None: + raise RuntimeError("HyperParallel CP requires a cp_mesh from DistributedSetup.") + if self._cp_size > 1: + _enable_cp_causal_lm_loss(self.model, self._cp_mesh) + self._model_shards_cp_inputs = any( + getattr(module, "_hp_cp_shards_inputs", False) for module in self.model.modules() + ) 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.""" - if self._cp_size > 1: - model = cp_prepare_model(model, self.accelerator, self._hp_args) - if self._ep_size > 1: - model = ep_prepare_model(model, self.accelerator, self._hp_args) - return fsdp2_prepare_model(self.accelerator, model, self._hp_args) + """Return a model already prepared by HyperParallel's unified model path.""" + if not isinstance(model, HSDPModule): + raise RuntimeError( + "HyperParallel models must be parallelized during meta initialization before Trainer preparation." + ) + 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 diff --git a/src/llamafactory/train/hyper_parallel/workflow.py b/src/llamafactory/train/hyper_parallel/workflow.py index 7f2efcd90..3d1c8329e 100644 --- a/src/llamafactory/train/hyper_parallel/workflow.py +++ b/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) - - ref_model = None - if finetuning_args.use_asft_loss: - ref_model = create_ref_model(model_args, finetuning_args) + model = load_hyper_parallel_model( + tokenizer, + 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, diff --git a/tests/train/test_hyper_parallel.py b/tests/train/test_hyper_parallel.py new file mode 100644 index 000000000..aa2a2d354 --- /dev/null +++ b/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"] == {}