Update to Transformers v5

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"
+            ]
+        }
+    ]
+}

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"
+            }
+        ]
+    }
+}

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

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,

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