diff --git a/mergekit/_data/architectures/mixtral.json b/mergekit/_data/architectures/mixtral.json new file mode 100644 index 00000000..adbea69f --- /dev/null +++ b/mergekit/_data/architectures/mixtral.json @@ -0,0 +1,57 @@ +{ + "model_type": "mixtral", + "architectures": [ + "MixtralForCausalLM" + ], + "pre_weights": [ + { + "name": "model.embed_tokens.weight", + "is_embed": true + } + ], + "num_layers_config_key": "num_hidden_layers", + "layer_templates": { + "weights": [ + { + "name": "model.layers.${layer_index}.input_layernorm.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.q_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.k_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.v_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.o_proj.weight" + }, + { + "name": "model.layers.${layer_index}.post_attention_layernorm.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.experts.gate_up_proj" + }, + { + "name": "model.layers.${layer_index}.mlp.experts.down_proj" + }, + { + "name": "model.layers.${layer_index}.mlp.gate.weight" + } + ] + }, + "post_weights": [ + { + "name": "model.norm.weight" + }, + { + "name": "lm_head.weight", + "is_embed": true, + "optional": true, + "tied_names": [ + "model.embed_tokens.weight" + ] + } + ] +} diff --git a/mergekit/_data/architectures/qwen3_moe.json b/mergekit/_data/architectures/qwen3_moe.json new file mode 100644 index 00000000..1c7d8495 --- /dev/null +++ b/mergekit/_data/architectures/qwen3_moe.json @@ -0,0 +1,63 @@ +{ + "model_type": "qwen3_moe", + "architectures": [ + "Qwen3MoeForCausalLM" + ], + "pre_weights": [ + { + "name": "model.embed_tokens.weight", + "is_embed": true + } + ], + "post_weights": [ + { + "name": "model.norm.weight" + }, + { + "name": "lm_head.weight", + "is_embed": true, + "optional": true, + "tied_names": [ + "model.embed_tokens.weight" + ] + } + ], + "num_layers_config_key": "num_hidden_layers", + "layer_templates": { + "weights": [ + { + "name": "model.layers.${layer_index}.input_layernorm.weight" + }, + { + "name": "model.layers.${layer_index}.mlp.experts.down_proj" + }, + { + "name": "model.layers.${layer_index}.mlp.experts.gate_up_proj" + }, + { + "name": "model.layers.${layer_index}.mlp.gate.weight" + }, + { + "name": "model.layers.${layer_index}.post_attention_layernorm.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.k_norm.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.k_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.q_norm.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.q_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.v_proj.weight" + }, + { + "name": "model.layers.${layer_index}.self_attn.o_proj.weight" + } + ] + } +} diff --git a/mergekit/architecture/__init__.py b/mergekit/architecture/__init__.py index c731e58b..5f324835 100644 --- a/mergekit/architecture/__init__.py +++ b/mergekit/architecture/__init__.py @@ -2,7 +2,6 @@ # SPDX-License-Identifier: LGPL-3.0-only import logging -from functools import lru_cache from typing import TYPE_CHECKING, Optional from transformers import PretrainedConfig @@ -20,8 +19,6 @@ from mergekit.architecture.moe_defs import ( AfmoeModuleArchitecture, Glm4MoeModuleArchitecture, - MixtralModuleArchitecture, - Qwen3MoeModuleArchitecture, ) from mergekit.options import MergeOptions @@ -38,21 +35,7 @@ def arch_info_for_config(config: PretrainedConfig) -> Optional[ModelArchitecture raise RuntimeError("More than one architecture in config?") arch_name = config.architectures[0] - if arch_name == MixtralModuleArchitecture.ARCHITECTURE_NAME: - module = MixtralModuleArchitecture.from_config(config) - return ModelArchitecture( - modules={"default": ModuleDefinition(architecture=module)}, - architectures=[arch_name], - model_type="mixtral", - ) - elif arch_name == Qwen3MoeModuleArchitecture.ARCHITECTURE_NAME: - module = Qwen3MoeModuleArchitecture.from_config(config) - return ModelArchitecture( - modules={"default": ModuleDefinition(architecture=module)}, - architectures=[arch_name], - model_type="qwen3_moe", - ) - elif arch_name == AfmoeModuleArchitecture.ARCHITECTURE_NAME: + if arch_name == AfmoeModuleArchitecture.ARCHITECTURE_NAME: module = AfmoeModuleArchitecture.from_config(config) return ModelArchitecture( modules={"default": ModuleDefinition(architecture=module)}, diff --git a/mergekit/architecture/auto.py b/mergekit/architecture/auto.py index 5515b44c..902cfe97 100644 --- a/mergekit/architecture/auto.py +++ b/mergekit/architecture/auto.py @@ -8,12 +8,14 @@ from typing import List, Optional, Tuple import torch +from transformers.initialization import no_init_weights from mergekit.architecture.base import ( ModelArchitecture, ModuleDefinition, WeightInfo, ) +from mergekit.architecture.conversion import can_convert_checkpoint_keys from mergekit.architecture.json_definitions import ( JsonLayerTemplates, JsonModuleArchDef, @@ -50,44 +52,108 @@ def get_transformers_info(model: ModelReference, options: MergeOptions) -> tuple f"Unable to load config for {model.model} - tied/ignored weights will not be detected", exc_info=e, ) - return None, None, None + return None, None, None, None + model_obj = None try: - with torch.device("meta"): - model = auto_cls.from_pretrained( - model.model.path, - revision=model.model.revision, + with torch.device("meta"), no_init_weights(): + model_obj = auto_cls.from_config( + cfg, trust_remote_code=options.trust_remote_code, - device_map="meta", ) except Exception as e: LOG.warning( - f"Unable to load model {model.model} with transformers - tied/ignored weights will not be detected", + f"Unable to instantiate model {model.model} with transformers config", exc_info=e, ) - return None, None, None + if model_obj is None: + try: + with torch.device("meta"): + model_obj = auto_cls.from_pretrained( + model.model.path, + revision=model.model.revision, + trust_remote_code=options.trust_remote_code, + device_map="meta", + ) + except Exception as e: + LOG.warning( + f"Unable to load model {model.model} with transformers - tied/ignored weights will not be detected", + exc_info=e, + ) + return None, None, None, None - ignore_on_save = getattr(model, "_keys_to_ignore_on_save", None) + ignore_on_save = getattr(model_obj, "_keys_to_ignore_on_save", None) if _get_tied_weight_keys is None: LOG.warning( "Unable to get tied weights - incompatible transformers version", ) tied_keys = None else: - tied_keys = _get_tied_weight_keys(model) + tied_keys = _get_tied_weight_keys(model_obj) if ignore_on_save is not None: ignore_on_save = set(ignore_on_save) embed_names = set() - _embed_out = model.get_output_embeddings() - _embed_in = model.get_input_embeddings() - for name, module in model.named_modules(): + _embed_out = model_obj.get_output_embeddings() + _embed_in = model_obj.get_input_embeddings() + for name, module in model_obj.named_modules(): if ( isinstance(module, torch.nn.Embedding) or module == _embed_out or module == _embed_in ): embed_names.add(name + ".weight") - return ignore_on_save, tied_keys, embed_names + tensor_names = list(model_obj.state_dict().keys()) + return ignore_on_save, tied_keys, embed_names, tensor_names + + +def _target_present_in_all_models( + target_name: str, + model_tensor_names: dict[ModelReference, set[str]], + model_types: dict[ModelReference, str], + use_transformers_layout: bool, +) -> bool: + if use_transformers_layout: + return all( + can_convert_checkpoint_keys( + model_types.get(model), + model_tensor_names[model], + target_name, + ) + for model in model_tensor_names + ) + return all(target_name in names for names in model_tensor_names.values()) + + +def _template_present_in_all_models( + full_name_template: str, + num_layers: Optional[int], + model_tensor_names: dict[ModelReference, set[str]], + model_types: dict[ModelReference, str], + use_transformers_layout: bool, +) -> bool: + if "${layer_index}" not in full_name_template: + return _target_present_in_all_models( + full_name_template, + model_tensor_names, + model_types, + use_transformers_layout, + ) + if num_layers is None: + return _target_present_in_all_models( + full_name_template.replace("${layer_index}", "0"), + model_tensor_names, + model_types, + use_transformers_layout, + ) + return all( + _target_present_in_all_models( + full_name_template.replace("${layer_index}", str(layer_idx)), + model_tensor_names, + model_types, + use_transformers_layout, + ) + for layer_idx in range(num_layers) + ) @lru_cache(maxsize=128) @@ -103,10 +169,19 @@ def infer_architecture_info( models = list(models) if base_model is None: base_model = models.pop(0) - all_tensor_names = set().union(*model_tensor_names.values()) - in_all_models = all_tensor_names.intersection(*model_tensor_names.values()) - - ignore_on_save, tied_keys, embed_names = get_transformers_info(base_model, options) + raw_tensor_names = set().union(*model_tensor_names.values()) + ignore_on_save, tied_keys, embed_names, transformer_tensor_names = ( + get_transformers_info(base_model, options) + ) + if transformer_tensor_names: + all_tensor_names = set(transformer_tensor_names) + model_types = { + model: model.config(trust_remote_code=options.trust_remote_code).model_type + for model in model_tensor_names + } + else: + all_tensor_names = raw_tensor_names + model_types = {} module_prefixes = set() module_layer_counts = defaultdict(int) @@ -157,9 +232,20 @@ def infer_architecture_info( f" {repr(prefix or 'default')} with {module_layer_counts[prefix]} layers, {len(module_templates[prefix])} templates, and {len(module_loose_weights[prefix])} loose weights" ) - def _wi(template: str, prefix: str) -> WeightInfo: + def _wi( + template: str, + prefix: str, + num_layers: Optional[int] = None, + ) -> WeightInfo: full_name = prefix + template - optional = (full_name.replace("${layer_index}", "0") not in in_all_models) or ( + present_in_all = _template_present_in_all_models( + full_name, + num_layers, + model_tensor_names, + model_types, + use_transformers_layout=bool(transformer_tensor_names), + ) + optional = (not present_in_all) or ( tied_keys is not None and any(re.search(pat, full_name) for pat in tied_keys) ) @@ -182,7 +268,10 @@ def _wi(template: str, prefix: str) -> WeightInfo: architectures=[], pre_weights=[_wi(t, "") for t in module_loose_weights[prefix]], layer_templates=JsonLayerTemplates( - weights=[_wi(t, "") for t in module_templates[prefix]] + weights=[ + _wi(t, "", num_layers=num_layers) + for t in module_templates[prefix] + ] ), post_weights=[], num_layers_config_key=None, diff --git a/mergekit/architecture/base.py b/mergekit/architecture/base.py index 738fb8f6..6406cc43 100644 --- a/mergekit/architecture/base.py +++ b/mergekit/architecture/base.py @@ -4,7 +4,7 @@ from abc import ABC, abstractmethod from typing import Dict, List, Optional, Tuple -import torch # required for Pydantic to resolve PretrainedConfig's torch.dtype forward reference +import torch # noqa: F401 from pydantic import BaseModel, Field from transformers import PretrainedConfig diff --git a/mergekit/architecture/conversion.py b/mergekit/architecture/conversion.py new file mode 100644 index 00000000..e2173d48 --- /dev/null +++ b/mergekit/architecture/conversion.py @@ -0,0 +1,185 @@ +# Copyright (C) 2025 Arcee AI +# SPDX-License-Identifier: LGPL-3.0-only + +import re +from collections.abc import Callable, Mapping +from typing import Optional, Union + +import torch +from transformers.conversion_mapping import get_checkpoint_conversion_mapping +from transformers.core_model_loading import WeightRenaming + +TensorSource = Union[torch.Tensor, Callable[[], torch.Tensor]] + + +class _TransformMatch: + def __init__( + self, + source_key: str, + renamed_key: str, + source_pattern: str, + source_value: TensorSource, + wildcard_values: tuple[str, ...], + ): + self.source_key = source_key + self.renamed_key = renamed_key + self.source_pattern = source_pattern + self.source_value = source_value + self.wildcard_values = wildcard_values + + +def _natural_sort_key(value: str): + return tuple( + int(part) if part.isdigit() else part for part in re.split(r"([0-9]+)", value) + ) + + +def _iter_transform_matches(transform, available: Mapping[str, TensorSource]): + for source_key in sorted(available, key=_natural_sort_key): + renamed_key, source_pattern = transform.rename_source_key(source_key) + if source_pattern is not None: + yield _TransformMatch( + source_key=source_key, + renamed_key=renamed_key, + source_pattern=source_pattern, + source_value=available[source_key], + wildcard_values=_wildcard_values(source_pattern, source_key), + ) + + +def _expanded_target_keys(renamed_key: str, target_patterns: list[str]) -> set[str]: + if not target_patterns: + return {renamed_key} + first_target = target_patterns[0] + if first_target not in renamed_key: + return {renamed_key} + prefix, _sep, suffix = renamed_key.partition(first_target) + return {prefix + target_pattern + suffix for target_pattern in target_patterns} + + +def _wildcard_values(pattern: str, source_key: str) -> tuple[str, ...]: + if "*" not in pattern: + return () + escaped = re.escape(pattern).replace(r"\*", r"([^.]*)") + match = re.search(escaped, source_key) + if not match: + return () + return match.groups() + + +def _has_complete_wildcard_group(wildcard_values: set[tuple[str, ...]]) -> bool: + return wildcard_values == {()} or len(wildcard_values) > 1 + + +def convert_checkpoint_tensors( + model_type: Optional[str], + source_tensors: Mapping[str, TensorSource], + target_key: str, +) -> Optional[torch.Tensor]: + """Convert checkpoint-layout tensors to a Transformers v5 model-layout tensor. + + Transformers v5 owns the conversion registry. This helper keeps mergekit's + tensor loading code independent from the registry's stateful loader protocol. + """ + if not model_type: + return None + + transforms = get_checkpoint_conversion_mapping(model_type) + if not transforms: + return None + + available: dict[str, TensorSource] = dict(source_tensors) + for transform in transforms: + additions: dict[str, TensorSource] = {} + conversion_layers = set() + for match in _iter_transform_matches(transform, available): + if isinstance(transform, WeightRenaming): + additions[match.renamed_key] = match.source_value + continue + + if target_key in _expanded_target_keys( + match.renamed_key, transform.target_patterns + ): + transform.add_tensor( + target_key=match.renamed_key, + source_key=match.source_key, + source_pattern=match.source_pattern, + future=match.source_value, + ) + conversion_layers.add(match.renamed_key) + + available.update(additions) + + if isinstance(transform, WeightRenaming): + continue + + for conversion_layer in conversion_layers: + converted = transform.convert(conversion_layer) + if target_key in converted: + return converted[target_key] + available.update(converted) + + tensor = available.get(target_key) + if callable(tensor): + tensor = tensor() + return tensor if isinstance(tensor, torch.Tensor) else None + + +def can_convert_checkpoint_keys( + model_type: Optional[str], + source_keys: set[str], + target_key: str, +) -> bool: + """Return whether source checkpoint keys can produce a target model key.""" + if target_key in source_keys: + return True + if not model_type: + return False + + transforms = get_checkpoint_conversion_mapping(model_type) + if not transforms: + return False + + available = set(source_keys) + for transform in transforms: + additions = set() + target_wildcards: dict[str, set[tuple[str, ...]]] = {} + for match in _iter_transform_matches( + transform, {source_key: source_key for source_key in available} + ): + if isinstance(transform, WeightRenaming): + additions.add(match.renamed_key) + elif target_key in _expanded_target_keys( + match.renamed_key, transform.target_patterns + ): + target_wildcards.setdefault(match.source_pattern, set()).add( + match.wildcard_values + ) + + available.update(additions) + if target_key in available: + return True + if target_wildcards and all( + target_wildcards.get(source_pattern) + for source_pattern in transform.source_patterns + ): + wildcard_sets = { + frozenset(target_wildcards[source_pattern]) + for source_pattern in transform.source_patterns + } + if len(wildcard_sets) == 1: + wildcard_values = next(iter(wildcard_sets)) + if _has_complete_wildcard_group(wildcard_values): + return True + return False + else: + return False + + if target_wildcards and len(transform.source_patterns) == 1: + wildcard_values = next(iter(target_wildcards.values())) + if _has_complete_wildcard_group(wildcard_values): + return True + + return False + + return False diff --git a/mergekit/architecture/moe_defs.py b/mergekit/architecture/moe_defs.py index 464314c0..e47cbc6a 100644 --- a/mergekit/architecture/moe_defs.py +++ b/mergekit/architecture/moe_defs.py @@ -12,97 +12,6 @@ ) from mergekit.architecture.json_definitions import NAME_TO_ARCH -MISTRAL_INFO = NAME_TO_ARCH["MistralForCausalLM"][0] -MISTRAL_MODULE_ARCH = MISTRAL_INFO.modules["default"].architecture - - -class MixtralModuleArchitecture(ModuleArchitecture, BaseModel): - ARCHITECTURE_NAME: ClassVar[str] = "MixtralForCausalLM" - num_local_experts: int - - def name(self) -> str: - return "mixtral" - - @classmethod - def from_config(cls, config: PretrainedConfig): - return MixtralModuleArchitecture(num_local_experts=config.num_local_experts) - - def pre_weights(self, config: PretrainedConfig) -> List[WeightInfo]: - return MISTRAL_MODULE_ARCH.pre_weights(config) - - def post_weights(self, config: PretrainedConfig) -> List[WeightInfo]: - return MISTRAL_MODULE_ARCH.post_weights(config) - - def num_layers_config_key(self) -> str: - return MISTRAL_MODULE_ARCH.num_layers_config_key() - - def layer_weights( - self, index: int, config: PretrainedConfig - ) -> Optional[List[WeightInfo]]: - num_experts = self.num_local_experts - prefix = f"model.layers.{index}" - tensor_names = [] - for expert_idx in range(num_experts): - for param in ("w1", "w2", "w3"): - tensor_names.append( - prefix + f".block_sparse_moe.experts.{expert_idx}.{param}.weight" - ) - tensor_names.append(prefix + ".block_sparse_moe.gate.weight") - res = [] - for name in tensor_names: - res.append(WeightInfo(name=name)) - for weight_info in MISTRAL_MODULE_ARCH.layer_weights(index, config): - if ".mlp." in weight_info.name: - continue - res.append(weight_info) - return res - - -QWEN3_INFO = NAME_TO_ARCH["Qwen3ForCausalLM"][0] -QWEN3_MODULE_ARCH = QWEN3_INFO.modules["default"].architecture - - -class Qwen3MoeModuleArchitecture(ModuleArchitecture, BaseModel): - ARCHITECTURE_NAME: ClassVar[str] = "Qwen3MoeForCausalLM" - num_experts: int - - def name(self) -> str: - return "qwen3_moe" - - @classmethod - def from_config(cls, config: PretrainedConfig): - return Qwen3MoeModuleArchitecture(num_experts=config.num_experts) - - def pre_weights(self, config: PretrainedConfig) -> List[WeightInfo]: - return QWEN3_MODULE_ARCH.pre_weights(config) - - def post_weights(self, config: PretrainedConfig) -> List[WeightInfo]: - return QWEN3_MODULE_ARCH.post_weights(config) - - def num_layers_config_key(self) -> str: - return QWEN3_MODULE_ARCH.num_layers_config_key() - - def layer_weights( - self, index: int, config: PretrainedConfig - ) -> Optional[List[WeightInfo]]: - prefix = f"model.layers.{index}" - tensor_names = [] - for expert_idx in range(self.num_experts): - for param in ("up_proj", "gate_proj", "down_proj"): - tensor_names.append( - prefix + f".mlp.experts.{expert_idx}.{param}.weight" - ) - tensor_names.append(prefix + ".mlp.gate.weight") - res = [] - for name in tensor_names: - res.append(WeightInfo(name=name)) - for weight_info in QWEN3_MODULE_ARCH.layer_weights(index, config): - if ".mlp." in weight_info.name: - continue - res.append(weight_info) - return res - - AFMOE_PARTIAL_INFO = NAME_TO_ARCH["_AfmoePartialForCausalLM"][0] AFMOE_PARTIAL_MODULE_ARCH = AFMOE_PARTIAL_INFO.modules["default"].architecture @@ -175,19 +84,12 @@ def layer_weights( if index < config.first_k_dense_replace: return GLM4_MODULE_ARCH.layer_weights(index, config) else: - tensor_names = [] - for expert_idx in range(self.num_experts): - tensor_names.append( - prefix + f".mlp.experts.{expert_idx}.gate_proj.weight" - ) - tensor_names.append( - prefix + f".mlp.experts.{expert_idx}.up_proj.weight" - ) - tensor_names.append( - prefix + f".mlp.experts.{expert_idx}.down_proj.weight" - ) - tensor_names.append(prefix + ".mlp.gate.weight") - tensor_names.append(prefix + ".mlp.gate.e_score_correction_bias") + tensor_names = [ + prefix + ".mlp.experts.gate_up_proj", + prefix + ".mlp.experts.down_proj", + prefix + ".mlp.gate.weight", + prefix + ".mlp.gate.e_score_correction_bias", + ] shared_expert_names = [ (prefix + ".mlp.shared_experts.gate_proj.weight", False), (prefix + ".mlp.shared_experts.up_proj.weight", False), diff --git a/mergekit/evo/actors.py b/mergekit/evo/actors.py index 726d2829..64ea515d 100644 --- a/mergekit/evo/actors.py +++ b/mergekit/evo/actors.py @@ -15,6 +15,7 @@ import ray.util.scheduling_strategies import torch import transformers +from transformers.initialization import no_init_weights from transformers.utils import is_flash_attn_2_available from mergekit.architecture.base import ConfiguredModelArchitecture @@ -32,7 +33,6 @@ from mergekit.evo.genome import InvalidGenotypeError, ModelGenome from mergekit.evo.helpers import _eval_model, evaluate_model, merge_model from mergekit.evo.monkeypatch import ( - NoInit, monkeypatch_lmeval_shuffle, monkeypatch_lmeval_vllm, ) @@ -187,7 +187,7 @@ def _maybe_init_model(self, config: MergeConfiguration): if is_flash_attn_2_available(): model_kwargs["attn_implementation"] = "flash_attention_2" - with NoInit(): + with no_init_weights(): inner_model = ( transformers.AutoModelForCausalLM.from_config( cfg_out, diff --git a/mergekit/evo/monkeypatch.py b/mergekit/evo/monkeypatch.py index 8f541052..c2e631e3 100644 --- a/mergekit/evo/monkeypatch.py +++ b/mergekit/evo/monkeypatch.py @@ -2,7 +2,6 @@ # SPDX-License-Identifier: LGPL-3.0-only -import torch import transformers @@ -95,34 +94,3 @@ def monkeypatch_lmeval_vllm(): lm_eval.models.vllm_causallms.VLLM.AUTO_MODEL_CLASS = ( transformers.AutoModelForCausalLM ) - - -class NoInit: - def __enter__(self): - def noop(*args, **kwargs): - pass - - (k, u, n) = ( - torch.nn.init.kaiming_uniform_, - torch.nn.init.uniform_, - torch.nn.init.normal_, - ) - torch.nn.init.kaiming_uniform_ = noop - torch.nn.init.uniform_ = noop - torch.nn.init.normal_ = noop - - transformers.modeling_utils._init_weights = False - self.funcs = (k, u, n) - - def __exit__(self, *args): - (k, u, n) = self.funcs - ( - torch.nn.init.kaiming_uniform_, - torch.nn.init.uniform_, - torch.nn.init.normal_, - ) = ( - k, - u, - n, - ) - transformers.modeling_utils._init_weights = True diff --git a/mergekit/io/tasks.py b/mergekit/io/tasks.py index 6997cbc3..f0525c03 100644 --- a/mergekit/io/tasks.py +++ b/mergekit/io/tasks.py @@ -9,6 +9,7 @@ import torch from mergekit.architecture import WeightInfo +from mergekit.architecture.conversion import convert_checkpoint_tensors from mergekit.common import ImmutableMap, ModelReference, dtype_from_name from mergekit.graph import Task from mergekit.io.lazy_tensor_loader import LazyTensorLoader @@ -90,17 +91,36 @@ def _resolve_name(self, loader: LazyTensorLoader) -> Optional[str]: return name return None + def _load_converted_tensor( + self, loader: LazyTensorLoader + ) -> Optional[torch.Tensor]: + model_type = self.model.config( + trust_remote_code=LoaderCache().trust_remote_code + ).model_type + source_tensors = { + key: (lambda key=key: loader.get_tensor(key, device=self.device or "cpu")) + for key in loader.index.tensor_paths + } + for target in [self.tensor] + list(self.aliases or []): + tensor = convert_checkpoint_tensors(model_type, source_tensors, target) + if tensor is not None: + return tensor.to(self.device or "cpu") + return None + def execute(self) -> Optional[torch.Tensor]: loader = LoaderCache().get(self.model) name = self._resolve_name(loader) - if not name: + if name: + x = loader.get_tensor(name, device=self.device or "cpu") + else: + x = self._load_converted_tensor(loader) + if x is None: if not self.optional: raise RuntimeError( f"Tensor {self.tensor} required but not present in model {self.model}" ) return None - x = loader.get_tensor(name, device=self.device or "cpu") if self.dtype and (dtype := dtype_from_name(self.dtype)) != x.dtype: x = x.to(dtype=dtype) return x diff --git a/mergekit/moe/common.py b/mergekit/moe/common.py index f529207d..d6937189 100644 --- a/mergekit/moe/common.py +++ b/mergekit/moe/common.py @@ -9,6 +9,7 @@ import transformers from mergekit.architecture import WeightInfo +from mergekit.architecture.conversion import convert_checkpoint_tensors from mergekit.common import ModelReference, dtype_from_name from mergekit.io import LazyTensorLoader, TensorWriter from mergekit.merge import MergeOptions @@ -101,3 +102,23 @@ def copy_tensor_out( tensor.to(dtype=out_dtype), clone=clone, ) + + +def convert_and_save_checkpoint_tensors( + model_type: str, + source_tensors: Dict[str, torch.Tensor], + target_name: str, + writer: TensorWriter, + out_dtype: Optional[torch.dtype] = None, + clone: bool = False, +): + tensor = convert_checkpoint_tensors(model_type, source_tensors, target_name) + if tensor is None: + raise RuntimeError( + f"Unable to convert {len(source_tensors)} source tensors to {target_name}" + ) + writer.save_tensor( + target_name, + tensor.to(dtype=out_dtype).contiguous(), + clone=clone, + ) diff --git a/mergekit/moe/mixtral.py b/mergekit/moe/mixtral.py index 9d89a7f3..92c41ce6 100644 --- a/mergekit/moe/mixtral.py +++ b/mergekit/moe/mixtral.py @@ -8,13 +8,20 @@ import tqdm import transformers -from mergekit.architecture import WeightInfo -from mergekit.architecture.moe_defs import MISTRAL_INFO +from mergekit.architecture.json_definitions import NAME_TO_ARCH from mergekit.moe.arch import MoEOutputArchitecture -from mergekit.moe.common import copy_tensor_out, initialize_io, select_dtype +from mergekit.moe.common import ( + convert_and_save_checkpoint_tensors, + copy_tensor_out, + initialize_io, + noise_and_scale, + select_dtype, +) from mergekit.moe.config import MoEMergeConfig from mergekit.options import MergeOptions +MISTRAL_INFO = NAME_TO_ARCH["MistralForCausalLM"][0] + class MixtralMoE(MoEOutputArchitecture): def name(self) -> str: @@ -81,22 +88,8 @@ def _generate_config( ) return out_cfg - def _remap_weight_name(self, weight: WeightInfo) -> str: - if ".mlp." not in weight.name: - # Everything but MLP is identical to base Mistral - return weight.name - - res = weight.name - for needle, replacement in [ - (".mlp.gate_proj", ".block_sparse_moe.experts.{expert_idx}.w1"), - (".mlp.down_proj", ".block_sparse_moe.experts.{expert_idx}.w2"), - (".mlp.up_proj", ".block_sparse_moe.experts.{expert_idx}.w3"), - ]: - res = res.replace(needle, replacement) - return res - def _router_weight_name(self, layer_idx: int) -> str: - return f"model.layers.{layer_idx}.block_sparse_moe.gate.weight" + return f"model.layers.{layer_idx}.mlp.gate.weight" def write_model( self, @@ -129,21 +122,8 @@ def write_model( MISTRAL_INFO.all_weights(base_cfg), desc="Weights", ): - tensor_name = self._remap_weight_name(weight_info) - if "{expert_idx}" in tensor_name: - for expert_index, expert in enumerate(config.experts): - expert_name = tensor_name.replace("{expert_idx}", str(expert_index)) - expert_loader = loaders.get(expert.source_model) - copy_tensor_out( - weight_info, - expert_loader, - writer, - expert=expert, - out_dtype=out_dtype, - output_name=expert_name, - clone=merge_options.clone_tensors, - is_residual="down_proj" in tensor_name, - ) + if ".mlp." in weight_info.name: + continue else: copy_tensor_out( weight_info, @@ -153,6 +133,49 @@ def write_model( clone=merge_options.clone_tensors, ) + for layer_idx in tqdm.trange(base_cfg.num_hidden_layers, desc="Expert weights"): + layer_prefix = f"model.layers.{layer_idx}.mlp" + gate_up_sources = {} + down_sources = {} + for expert_idx, expert in enumerate(config.experts): + expert_loader = loaders.get(expert.source_model) + for source_param, expert_param in [ + ("gate_proj", "w1"), + ("down_proj", "w2"), + ("up_proj", "w3"), + ]: + source_name = f"model.layers.{layer_idx}.mlp.{source_param}.weight" + tensor = expert_loader.get_tensor(source_name) + tensor = noise_and_scale( + tensor, + expert, + is_residual=source_param == "down_proj", + ) + expert_name = ( + f"{layer_prefix}.experts.{expert_idx}.{expert_param}.weight" + ) + if source_param == "down_proj": + down_sources[expert_name] = tensor + else: + gate_up_sources[expert_name] = tensor + + convert_and_save_checkpoint_tensors( + "mixtral", + gate_up_sources, + f"{layer_prefix}.experts.gate_up_proj", + writer, + out_dtype=out_dtype, + clone=merge_options.clone_tensors, + ) + convert_and_save_checkpoint_tensors( + "mixtral", + down_sources, + f"{layer_prefix}.experts.down_proj", + writer, + out_dtype=out_dtype, + clone=merge_options.clone_tensors, + ) + for layer_idx, weight in enumerate( tqdm.tqdm(router_weights, desc="Router weights") ): diff --git a/mergekit/moe/qwen3.py b/mergekit/moe/qwen3.py index f243b9f0..f74dbd90 100644 --- a/mergekit/moe/qwen3.py +++ b/mergekit/moe/qwen3.py @@ -11,7 +11,12 @@ from mergekit.architecture.json_definitions import NAME_TO_ARCH from mergekit.moe.arch import MoEOutputArchitecture -from mergekit.moe.common import copy_tensor_out, initialize_io, select_dtype +from mergekit.moe.common import ( + convert_and_save_checkpoint_tensors, + initialize_io, + noise_and_scale, + select_dtype, +) from mergekit.moe.config import MoEMergeConfig from mergekit.options import MergeOptions @@ -94,21 +99,7 @@ def write_model( ): tensor_name = weight_info.name if ".mlp." in tensor_name: - for expert_idx, expert in enumerate(config.experts): - expert_name = tensor_name.replace( - ".mlp.", f".mlp.experts.{expert_idx}." - ) - expert_loader = loaders.get(expert.source_model) - copy_tensor_out( - weight_info, - expert_loader, - writer, - expert=expert, - is_residual="down_proj" in tensor_name, - output_name=expert_name, - out_dtype=out_dtype, - clone=merge_options.clone_tensors, - ) + continue else: tensor = base_loader.get_tensor( tensor_name, @@ -124,6 +115,43 @@ def write_model( clone=merge_options.clone_tensors, ) + for layer_idx in tqdm.trange(base_cfg.num_hidden_layers, desc="Expert weights"): + layer_prefix = f"model.layers.{layer_idx}.mlp" + gate_up_sources = {} + down_sources = {} + for expert_idx, expert in enumerate(config.experts): + expert_loader = loaders.get(expert.source_model) + for param in ("gate_proj", "up_proj", "down_proj"): + source_name = f"model.layers.{layer_idx}.mlp.{param}.weight" + tensor = expert_loader.get_tensor(source_name) + tensor = noise_and_scale( + tensor, + expert, + is_residual=param == "down_proj", + ) + expert_name = f"{layer_prefix}.experts.{expert_idx}.{param}.weight" + if param == "down_proj": + down_sources[expert_name] = tensor + else: + gate_up_sources[expert_name] = tensor + + convert_and_save_checkpoint_tensors( + "qwen3_moe", + gate_up_sources, + f"{layer_prefix}.experts.gate_up_proj", + writer, + out_dtype=out_dtype, + clone=merge_options.clone_tensors, + ) + convert_and_save_checkpoint_tensors( + "qwen3_moe", + down_sources, + f"{layer_prefix}.experts.down_proj", + writer, + out_dtype=out_dtype, + clone=merge_options.clone_tensors, + ) + for layer_idx, weight in enumerate( tqdm.tqdm(router_weights, desc="Router weights") ): diff --git a/mergekit/moe/router.py b/mergekit/moe/router.py index 7381739d..d4614730 100644 --- a/mergekit/moe/router.py +++ b/mergekit/moe/router.py @@ -115,9 +115,13 @@ def _do_it(tokenized): elif mode in ("hidden", "hidden_avg", "hidden_last"): kwargs = {} if load_in_4bit: - kwargs["load_in_4bit"] = True - if load_in_8bit: - kwargs["load_in_8bit"] = True + kwargs["quantization_config"] = transformers.BitsAndBytesConfig( + load_in_4bit=True + ) + elif load_in_8bit: + kwargs["quantization_config"] = transformers.BitsAndBytesConfig( + load_in_8bit=True + ) model = AutoModelForCausalLM.from_pretrained( model_ref.model.path, revision=model_ref.model.revision, diff --git a/mergekit/scripts/evolve.py b/mergekit/scripts/evolve.py index 4075b40f..40acd859 100644 --- a/mergekit/scripts/evolve.py +++ b/mergekit/scripts/evolve.py @@ -5,6 +5,7 @@ import logging import os import time +import warnings from typing import List, Optional import click @@ -147,6 +148,12 @@ def main( load_in_4bit: bool, force_population_size: Optional[int], ): + warnings.warn( + "mergekit-evolve is deprecated and is likely to be removed in a future release.", + DeprecationWarning, + stacklevel=2, + ) + config = EvolMergeConfiguration.model_validate( yaml.safe_load(open(genome_config_path, "r", encoding="utf-8")) ) diff --git a/pyproject.toml b/pyproject.toml index 05136aaa..d5710eb7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,10 +18,10 @@ dependencies = [ "accelerate~=1.6.0", "pydantic~=2.10.6", "immutables==0.21", - "transformers>=4.45.2", + "transformers>=5.0,<6.0", "tokenizers>=0.20.1", "huggingface_hub", - "peft", + "peft>=0.18.0", "typing-extensions", "sentencepiece", "protobuf", diff --git a/tests/test_architecture_conversion.py b/tests/test_architecture_conversion.py new file mode 100644 index 00000000..220eb3b3 --- /dev/null +++ b/tests/test_architecture_conversion.py @@ -0,0 +1,292 @@ +import torch +import transformers +from safetensors.torch import save_file + +from mergekit.architecture import arch_info_for_config +from mergekit.architecture.auto import infer_architecture_info +from mergekit.architecture.conversion import ( + can_convert_checkpoint_keys, + convert_checkpoint_tensors, +) +from mergekit.common import ModelReference +from mergekit.options import MergeOptions + + +def test_qwen3_moe_expert_weights_convert_to_v5_layout(): + sources = { + "model.layers.0.mlp.experts.0.gate_proj.weight": torch.full((2, 3), 1.0), + "model.layers.0.mlp.experts.1.gate_proj.weight": torch.full((2, 3), 2.0), + "model.layers.0.mlp.experts.0.up_proj.weight": torch.full((2, 3), 3.0), + "model.layers.0.mlp.experts.1.up_proj.weight": torch.full((2, 3), 4.0), + } + + converted = convert_checkpoint_tensors( + "qwen3_moe", + sources, + "model.layers.0.mlp.experts.gate_up_proj", + ) + + torch.testing.assert_close( + converted, + torch.stack( + [ + torch.cat( + [ + sources["model.layers.0.mlp.experts.0.gate_proj.weight"], + sources["model.layers.0.mlp.experts.0.up_proj.weight"], + ], + dim=0, + ), + torch.cat( + [ + sources["model.layers.0.mlp.experts.1.gate_proj.weight"], + sources["model.layers.0.mlp.experts.1.up_proj.weight"], + ], + dim=0, + ), + ], + dim=0, + ), + ) + + +def test_mixtral_old_checkpoint_names_convert_to_v5_layout(): + sources = { + "model.layers.0.block_sparse_moe.experts.0.w1.weight": torch.full((2, 3), 1.0), + "model.layers.0.block_sparse_moe.experts.1.w1.weight": torch.full((2, 3), 2.0), + "model.layers.0.block_sparse_moe.experts.0.w3.weight": torch.full((2, 3), 3.0), + "model.layers.0.block_sparse_moe.experts.1.w3.weight": torch.full((2, 3), 4.0), + } + + converted = convert_checkpoint_tensors( + "mixtral", + sources, + "model.layers.0.mlp.experts.gate_up_proj", + ) + + torch.testing.assert_close( + converted, + torch.stack( + [ + torch.cat( + [ + sources["model.layers.0.block_sparse_moe.experts.0.w1.weight"], + sources["model.layers.0.block_sparse_moe.experts.0.w3.weight"], + ], + dim=0, + ), + torch.cat( + [ + sources["model.layers.0.block_sparse_moe.experts.1.w1.weight"], + sources["model.layers.0.block_sparse_moe.experts.1.w3.weight"], + ], + dim=0, + ), + ], + dim=0, + ), + ) + + +def test_mixtral_writer_keys_convert_to_v5_layout(): + sources = { + "model.layers.0.mlp.experts.0.w1.weight": torch.full((2, 3), 1.0), + "model.layers.0.mlp.experts.1.w1.weight": torch.full((2, 3), 2.0), + "model.layers.0.mlp.experts.0.w3.weight": torch.full((2, 3), 3.0), + "model.layers.0.mlp.experts.1.w3.weight": torch.full((2, 3), 4.0), + } + + converted = convert_checkpoint_tensors( + "mixtral", + sources, + "model.layers.0.mlp.experts.gate_up_proj", + ) + + torch.testing.assert_close( + converted, + torch.stack( + [ + torch.cat( + [ + sources["model.layers.0.mlp.experts.0.w1.weight"], + sources["model.layers.0.mlp.experts.0.w3.weight"], + ], + dim=0, + ), + torch.cat( + [ + sources["model.layers.0.mlp.experts.1.w1.weight"], + sources["model.layers.0.mlp.experts.1.w3.weight"], + ], + dim=0, + ), + ], + dim=0, + ), + ) + + +def test_key_conversion_requires_complete_single_pattern_groups(): + assert not can_convert_checkpoint_keys( + "mixtral", + {"model.layers.0.block_sparse_moe.experts.0.w2.weight"}, + "model.layers.0.mlp.experts.down_proj", + ) + + +def test_key_conversion_requires_matching_wildcard_groups(): + assert not can_convert_checkpoint_keys( + "mixtral", + { + "model.layers.0.block_sparse_moe.experts.0.w1.weight", + "model.layers.0.block_sparse_moe.experts.1.w1.weight", + "model.layers.0.block_sparse_moe.experts.1.w3.weight", + "model.layers.0.block_sparse_moe.experts.2.w3.weight", + }, + "model.layers.0.mlp.experts.gate_up_proj", + ) + + +def test_key_conversion_requires_complete_multi_pattern_groups(): + assert not can_convert_checkpoint_keys( + "mixtral", + { + "model.layers.0.block_sparse_moe.experts.0.w1.weight", + "model.layers.0.block_sparse_moe.experts.0.w3.weight", + }, + "model.layers.0.mlp.experts.gate_up_proj", + ) + + +def test_key_conversion_supports_later_one_to_many_targets(): + assert can_convert_checkpoint_keys( + "hrm_text", + {"layers.0.mlp.gate_up_proj.weight"}, + "layers.0.mlp.up_proj.weight", + ) + + +def test_tensor_conversion_supports_later_one_to_many_targets(): + converted = convert_checkpoint_tensors( + "hrm_text", + {"layers.0.mlp.gate_up_proj.weight": torch.arange(12).reshape(4, 3)}, + "layers.0.mlp.up_proj.weight", + ) + + torch.testing.assert_close(converted, torch.arange(6, 12).reshape(2, 3)) + + +def test_mixtral_architecture_uses_json_v5_layout(): + cfg = transformers.MixtralConfig( + vocab_size=32, + hidden_size=16, + intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=2, + num_local_experts=2, + num_experts_per_tok=1, + ) + cfg.architectures = ["MixtralForCausalLM"] + + arch = arch_info_for_config(cfg) + names = {w.name for w in arch.all_weights(cfg)} + + assert "model.layers.0.mlp.experts.gate_up_proj" in names + assert "model.layers.0.mlp.experts.down_proj" in names + assert "model.layers.0.block_sparse_moe.experts.0.w1.weight" not in names + + +def test_qwen3_moe_architecture_uses_json_v5_layout(): + cfg = transformers.Qwen3MoeConfig( + vocab_size=32, + hidden_size=16, + intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=2, + num_experts=2, + num_experts_per_tok=1, + ) + cfg.architectures = ["Qwen3MoeForCausalLM"] + + arch = arch_info_for_config(cfg) + names = {w.name for w in arch.all_weights(cfg)} + + assert "model.layers.0.mlp.experts.gate_up_proj" in names + assert "model.layers.0.mlp.experts.down_proj" in names + assert "model.layers.0.mlp.experts.0.gate_proj.weight" not in names + + +def test_auto_inference_uses_transformers_v5_layout_with_old_checkpoint_keys( + tmp_path, +): + cfg = transformers.MixtralConfig( + vocab_size=32, + hidden_size=16, + intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=2, + num_local_experts=2, + num_experts_per_tok=1, + ) + cfg.architectures = ["UnknownMixtralForCausalLM"] + cfg.save_pretrained(tmp_path) + save_file( + { + "model.layers.0.block_sparse_moe.experts.0.w1.weight": torch.zeros(2, 3), + "model.layers.0.block_sparse_moe.experts.1.w1.weight": torch.zeros(2, 3), + "model.layers.0.block_sparse_moe.experts.0.w3.weight": torch.zeros(2, 3), + "model.layers.0.block_sparse_moe.experts.1.w3.weight": torch.zeros(2, 3), + "model.layers.0.block_sparse_moe.experts.0.w2.weight": torch.zeros(2, 3), + "model.layers.0.block_sparse_moe.experts.1.w2.weight": torch.zeros(2, 3), + }, + tmp_path / "model.safetensors", + ) + model_ref = ModelReference.parse(str(tmp_path)) + + arch = infer_architecture_info((model_ref,), model_ref, MergeOptions()) + weights = {w.name: w for w in arch.all_weights(cfg)} + + assert "model.layers.0.mlp.experts.gate_up_proj" in weights + assert "model.layers.0.mlp.experts.down_proj" in weights + assert not weights["model.layers.0.mlp.experts.gate_up_proj"].optional + assert ( + "model.layers.${layer_index}.block_sparse_moe.experts.0.w1.weight" + not in weights + ) + + +def test_auto_inference_marks_template_optional_if_missing_in_any_layer(tmp_path): + cfg = transformers.MixtralConfig( + vocab_size=32, + hidden_size=16, + intermediate_size=32, + num_hidden_layers=2, + num_attention_heads=2, + num_key_value_heads=2, + num_local_experts=2, + num_experts_per_tok=1, + ) + cfg.architectures = ["UnknownMixtralForCausalLM"] + cfg.save_pretrained(tmp_path) + save_file( + { + "model.layers.0.block_sparse_moe.experts.0.w1.weight": torch.zeros(2, 3), + "model.layers.0.block_sparse_moe.experts.1.w1.weight": torch.zeros(2, 3), + "model.layers.0.block_sparse_moe.experts.0.w3.weight": torch.zeros(2, 3), + "model.layers.0.block_sparse_moe.experts.1.w3.weight": torch.zeros(2, 3), + "model.layers.1.block_sparse_moe.experts.0.w1.weight": torch.zeros(2, 3), + "model.layers.1.block_sparse_moe.experts.1.w1.weight": torch.zeros(2, 3), + "model.layers.1.block_sparse_moe.experts.0.w3.weight": torch.zeros(2, 3), + }, + tmp_path / "model.safetensors", + ) + model_ref = ModelReference.parse(str(tmp_path)) + + arch = infer_architecture_info((model_ref,), model_ref, MergeOptions()) + module_arch = arch.modules["default"].architecture + weights = {w.name: w for w in module_arch.layer_weights(0, cfg)} + + assert weights["model.layers.0.mlp.experts.gate_up_proj"].optional