conversion: refactor moe exp tensor handling

conversion/afmoe.py

@@ -2,12 +2,10 @@ from __future__ import annotations
 
 from typing import Callable, Iterable, TYPE_CHECKING
 
-import torch
-
 if TYPE_CHECKING:
     from torch import Tensor
 
-from .base import ModelBase, gguf
+from .base import MOE_HF_MLP, ModelBase, gguf
 
 from .llama import LlamaModel
 
@@ -46,34 +44,7 @@ class AfmoeModel(LlamaModel):
 
         return super().filter_tensors((name, gen))
 
+    moe_experts = [MOE_HF_MLP._replace(weights=("gate_proj", "up_proj", "down_proj"))]
+
     def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
-        # Handle expert weights - they're already merged in the HF format
-        # process the experts separately
-        if name.find("mlp.experts") != -1:
-            n_experts = self.find_hparam(["num_local_experts", "num_experts"])
-            assert bid is not None
-
-            if self._experts is None:
-                self._experts = [{} for _ in range(self.block_count)]
-
-            self._experts[bid][name] = data_torch
-
-            if len(self._experts[bid]) >= n_experts * 3:
-                # merge the experts into a single 3d tensor
-                for w_name in ["gate_proj", "up_proj", "down_proj"]:
-                    datas: list[Tensor] = []
-
-                    for xid in range(n_experts):
-                        ename_to_retrieve = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight"
-                        datas.append(self._experts[bid][ename_to_retrieve])
-                        del self._experts[bid][ename_to_retrieve]
-
-                    data_torch = torch.stack(datas, dim=0)
-                    merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight"
-                    yield from ModelBase.modify_tensors(self, data_torch, merged_name, bid)
-
-                return
-            else:
-                return
-
         yield from ModelBase.modify_tensors(self, data_torch, name, bid)

conversion/arctic.py

@@ -5,12 +5,10 @@ import sys
 
 from typing import Iterable, TYPE_CHECKING
 
-import torch
-
 if TYPE_CHECKING:
     from torch import Tensor
 
-from .base import ModelBase, SentencePieceTokenTypes, TextModel, gguf, logger
+from .base import MOE_BLOCK_SPARSE_MIXTRAL, ModelBase, SentencePieceTokenTypes, TextModel, gguf, logger
 
 from .llama import LlamaModel
 
@@ -109,7 +107,7 @@ class ArcticModel(TextModel):
         self.gguf_writer.add_vocab_size(hparams["vocab_size"])
         self.gguf_writer.add_rope_dimension_count(hparams["hidden_size"] // hparams["num_attention_heads"])
 
-    _experts: list[dict[str, Tensor]] | None = None
+    moe_experts = [MOE_BLOCK_SPARSE_MIXTRAL._replace(n_expert=("num_local_experts",))]
 
     def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
         n_head = self.hparams["num_attention_heads"]
@@ -120,43 +118,4 @@ class ArcticModel(TextModel):
         if name.endswith("k_proj.weight"):
             data_torch = LlamaModel.permute(data_torch, n_head, n_kv_head)
 
-        # process the experts separately
-        if name.find("block_sparse_moe.experts") != -1:
-            n_experts = self.hparams["num_local_experts"]
-
-            assert bid is not None
-
-            if self._experts is None:
-                self._experts = [{} for _ in range(self.block_count)]
-
-            self._experts[bid][name] = data_torch
-
-            if len(self._experts[bid]) >= n_experts * 3:
-                # merge the experts into a single 3d tensor
-                for wid in ["w1", "w2", "w3"]:
-                    datas: list[Tensor] = []
-
-                    for xid in range(n_experts):
-                        ename = f"model.layers.{bid}.block_sparse_moe.experts.{xid}.{wid}.weight"
-                        datas.append(self._experts[bid][ename])
-                        del self._experts[bid][ename]
-
-                    data_torch = torch.stack(datas, dim=0)
-
-                    merged_name = f"layers.{bid}.feed_forward.experts.{wid}.weight"
-
-                    yield from super().modify_tensors(data_torch, merged_name, bid)
-                return
-            else:
-                return
-
         yield from super().modify_tensors(data_torch, name, bid)
-
-    def prepare_tensors(self):
-        super().prepare_tensors()
-
-        if self._experts is not None:
-            # flatten `list[dict[str, Tensor]]` into `list[str]`
-            experts = [k for d in self._experts for k in d.keys()]
-            if len(experts) > 0:
-                raise ValueError(f"Unprocessed experts: {experts}")

conversion/bailingmoe.py

@@ -7,7 +7,7 @@ import torch
 if TYPE_CHECKING:
     from torch import Tensor
 
-from .base import ModelBase, TextModel, gguf
+from .base import MOE_HF_MLP, ModelBase, TextModel, gguf
 
 
 @ModelBase.register("BailingMoeForCausalLM")
@@ -31,7 +31,7 @@ class BailingMoeModel(TextModel):
         self.gguf_writer.add_expert_shared_count(hparams["num_shared_experts"])
         self.gguf_writer.add_expert_weights_norm(hparams["norm_topk_prob"])
 
-    _experts: list[dict[str, Tensor]] | None = None
+    moe_experts = [MOE_HF_MLP]
 
     @staticmethod
     def permute(weights: Tensor, n_head: int, n_head_kv: int | None):
@@ -60,34 +60,6 @@ class BailingMoeModel(TextModel):
             yield from super().modify_tensors(BailingMoeModel.permute(k, n_head, n_kv_head), self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_K, bid), bid)
             yield from super().modify_tensors(v,self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_V, bid), bid)
             return
-        elif name.find("mlp.experts") != -1:
-            n_experts = self.find_hparam(["num_local_experts", "num_experts"])
-            assert bid is not None
-
-            if self._experts is None:
-                self._experts = [{} for _ in range(self.block_count)]
-
-            self._experts[bid][name] = data_torch
-
-            if len(self._experts[bid]) >= n_experts * 3:
-                # merge the experts into a single 3d tensor
-                for w_name in ["down_proj", "gate_proj", "up_proj"]:
-                    datas: list[Tensor] = []
-
-                    for xid in range(n_experts):
-                        ename = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight"
-                        datas.append(self._experts[bid][ename])
-                        del self._experts[bid][ename]
-
-                    data_torch = torch.stack(datas, dim=0)
-
-                    merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight"
-
-                    new_name = self.map_tensor_name(merged_name)
-
-                    yield from super().modify_tensors(data_torch, new_name, bid)
-
-            return
 
         new_name = self.map_tensor_name(name)
 
@@ -97,15 +69,6 @@ class BailingMoeModel(TextModel):
 
         yield from super().modify_tensors(data_torch, new_name, bid)
 
-    def prepare_tensors(self):
-        super().prepare_tensors()
-
-        if self._experts is not None:
-            # flatten `list[dict[str, Tensor]]` into `list[str]`
-            experts = [k for d in self._experts for k in d.keys()]
-            if len(experts) > 0:
-                raise ValueError(f"Unprocessed experts: {experts}")
-
 
 @ModelBase.register("BailingMoeV2ForCausalLM")
 class BailingMoeV2Model(TextModel):
@@ -138,7 +101,7 @@ class BailingMoeV2Model(TextModel):
         if (nextn_layers := self.hparams.get("num_nextn_predict_layers")) is not None:
             self.gguf_writer.add_nextn_predict_layers(nextn_layers)
 
-    _experts: list[dict[str, Tensor]] | None = None
+    moe_experts = [MOE_HF_MLP]
 
     @classmethod
     def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
@@ -149,44 +112,6 @@ class BailingMoeV2Model(TextModel):
 
         return super().filter_tensors((name, gen))
 
-    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
-        if "mlp.experts" in name:
-            n_experts = self.find_hparam(["num_local_experts", "num_experts"])
-            assert bid is not None
-
-            if self._experts is None:
-                self._experts = [{} for _ in range(self.block_count)]
-
-            self._experts[bid][name] = data_torch
-
-            if len(self._experts[bid]) >= n_experts * 3:
-                # merge the experts into a single 3d tensor
-                for w_name in ["down_proj", "gate_proj", "up_proj"]:
-                    datas: list[Tensor] = []
-
-                    for xid in range(n_experts):
-                        ename = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight"
-                        datas.append(self._experts[bid][ename])
-                        del self._experts[bid][ename]
-
-                    data_torch = torch.stack(datas, dim=0)
-
-                    merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight"
-
-                    yield from super().modify_tensors(data_torch, merged_name, bid)
-            return
-
-        yield from super().modify_tensors(data_torch, name, bid)
-
-    def prepare_tensors(self):
-        super().prepare_tensors()
-
-        if self._experts is not None:
-            # flatten `list[dict[str, Tensor]]` into `list[str]`
-            experts = [k for d in self._experts for k in d.keys()]
-            if len(experts) > 0:
-                raise ValueError(f"Unprocessed experts: {experts}")
-
 
 @ModelBase.register("SarvamMoEForCausalLM", "modeling_sarvam_moe.SarvamMoEForCausalLM")
 class SarvamMoEModel(BailingMoeV2Model):

conversion/base.py

@@ -11,9 +11,10 @@ import os
 import re
 import sys
 from enum import IntEnum
+from functools import lru_cache
 from pathlib import Path
 from hashlib import sha256
-from typing import TYPE_CHECKING, Any, Callable, ContextManager, Iterable, Iterator, Literal, Sequence, TypeVar, cast
+from typing import TYPE_CHECKING, Any, Callable, ContextManager, Iterable, Iterator, Literal, NamedTuple, Sequence, TypeVar, cast
 from itertools import chain
 from transformers import AutoConfig
 
@@ -77,6 +78,50 @@ class ModelType(IntEnum):
     MMPROJ = 2
 
 
+class MoEExpertSpec(NamedTuple):
+    # src/dst are name templates using {bid}, {xid} (src only) and {w}
+    # the order of weights is the order the stacked tensors are written in
+    src: str
+    dst: str
+    weights: tuple[str, ...]
+    n_expert: tuple[str, ...] | Callable[[ModelBase], int] = ("num_local_experts", "num_experts", "n_routed_experts")
+
+
+_MOE_PLACEHOLDER_RE = re.compile(r"\{(bid|xid|w)\}")
+
+
+@lru_cache(maxsize=None)
+def _compile_moe_src(spec: MoEExpertSpec) -> re.Pattern[str]:
+    groups = {
+        "bid": r"(?P<bid>\d+)",
+        "xid": r"(?P<xid>\d+)",
+        "w":   "(?P<w>" + "|".join(re.escape(w) for w in spec.weights) + ")",
+    }
+    parts: list[str] = []
+    pos = 0
+    for m in _MOE_PLACEHOLDER_RE.finditer(spec.src):
+        parts.append(re.escape(spec.src[pos:m.start()]))
+        parts.append(groups[m.group(1)])
+        pos = m.end()
+    parts.append(re.escape(spec.src[pos:]))
+    return re.compile("".join(parts) + r"\Z")
+
+
+MOE_HF_MLP = MoEExpertSpec(
+    src="model.layers.{bid}.mlp.experts.{xid}.{w}.weight",
+    dst="model.layers.{bid}.mlp.experts.{w}.weight",
+    weights=("down_proj", "gate_proj", "up_proj"),
+)
+
+MOE_BLOCK_SPARSE = MoEExpertSpec(
+    src="model.layers.{bid}.block_sparse_moe.experts.{xid}.{w}.weight",
+    dst="model.layers.{bid}.block_sparse_moe.experts.{w}.weight",
+    weights=("w1", "w2", "w3"),
+)
+
+MOE_BLOCK_SPARSE_MIXTRAL = MOE_BLOCK_SPARSE._replace(dst="layers.{bid}.feed_forward.experts.{w}.weight")
+
+
 class ModelBase:
     _model_classes: dict[ModelType, dict[str, type[ModelBase]]] = {
         ModelType.TEXT: {},
@@ -109,6 +154,9 @@ class ModelBase:
     block_count: int
     tensor_map: gguf.TensorNameMap
 
+    # per-expert tensors to stack into 3d tensors, see _stack_moe_experts()
+    moe_experts: Sequence[MoEExpertSpec] = ()
+
     # Mistral format specifics
     is_mistral_format: bool = False
     disable_mistral_community_chat_template: bool = False
@@ -153,6 +201,7 @@ class ModelBase:
         self.fuse_gate_up_exps = fuse_gate_up_exps
         self._gate_exp_buffer: dict[int, Tensor] = {}
         self._up_exp_buffer: dict[int, Tensor] = {}
+        self._moe_buffers: dict[tuple[int, int], dict[str, Tensor]] = {}
         self.hparams = ModelBase.load_hparams(self.dir_model, self.is_mistral_format) if hparams is None else hparams
         self.model_tensors = self.index_tensors(remote_hf_model_id=remote_hf_model_id)
         self.metadata_override = metadata_override
@@ -620,6 +669,53 @@ class ModelBase:
     def set_gguf_parameters(self):
         raise NotImplementedError("set_gguf_parameters() must be implemented in subclasses")
 
+    def moe_expert_specs(self) -> Sequence[MoEExpertSpec]:
+        return self.moe_experts
+
+    def moe_n_expert(self, spec: MoEExpertSpec) -> int:
+        if callable(spec.n_expert):
+            return spec.n_expert(self)
+        return self.find_hparam(spec.n_expert)
+
+    def _stack_moe_experts(self, data_torch: Tensor, name: str) -> list[tuple[str, Tensor, int]] | None:
+        # returns None if name is not a per-expert tensor, or an empty list while the block is incomplete
+        for sid, spec in enumerate(self.moe_expert_specs()):
+            m = _compile_moe_src(spec).match(name)
+            if m is None:
+                continue
+
+            bid = int(m.group("bid"))
+            n_expert = self.moe_n_expert(spec)
+
+            buf = self._moe_buffers.setdefault((sid, bid), {})
+            buf[name] = data_torch
+            if len(buf) < n_expert * len(spec.weights):
+                return []
+
+            del self._moe_buffers[(sid, bid)]
+
+            merged: list[tuple[str, Tensor, int]] = []
+            for w in spec.weights:
+                datas = [buf.pop(spec.src.format(bid=bid, xid=xid, w=w)) for xid in range(n_expert)]
+                merged.append((spec.dst.format(bid=bid, w=w), torch.stack(datas, dim=0), bid))
+
+            if buf:
+                raise ValueError(f"Unexpected experts: {list(buf)}")
+
+            return merged
+
+        return None
+
+    def dispatch_tensor(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        stacked = self._stack_moe_experts(data_torch, name)
+        if stacked is None:
+            yield from self.modify_tensors(data_torch, name, bid)
+            return
+
+        # stacked tensors re-enter modify_tensors() as if stored that way
+        for merged_name, merged_data, merged_bid in stacked:
+            yield from self.modify_tensors(merged_data, merged_name, merged_bid)
+
     def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
         new_name = self.map_tensor_name(name)
 
@@ -923,7 +1019,7 @@ class ModelBase:
                     bid = int(part)
                     break
 
-            for new_name, data_torch in (self.modify_tensors(data_torch, name, bid)):
+            for new_name, data_torch in (self.dispatch_tensor(data_torch, name, bid)):
                 # TODO: why do we squeeze here?
                 # data = data_torch.squeeze().numpy()
                 data = data_torch.numpy()
@@ -1023,6 +1119,10 @@ class ModelBase:
 
                 self.gguf_writer.add_tensor(new_name, data, raw_dtype=data_qtype)
 
+        if self._moe_buffers:
+            experts = [k for buf in self._moe_buffers.values() for k in buf]
+            raise ValueError(f"Unprocessed experts: {experts}")
+
     def set_type(self):
         self.gguf_writer.add_type(gguf.GGUFType.MODEL)
 

conversion/command_r.py

@@ -8,7 +8,7 @@ import torch
 if TYPE_CHECKING:
     from torch import Tensor
 
-from .base import ModelBase, TextModel, gguf, logger
+from .base import MOE_HF_MLP, ModelBase, TextModel, gguf, logger
 
 
 @ModelBase.register("CohereForCausalLM")
@@ -62,16 +62,14 @@ class Cohere2Model(TextModel):
 class Cohere2MoeModel(TextModel):
     model_arch = gguf.MODEL_ARCH.COHERE2MOE
     _n_main_layers: int | None = None
-    _expert_tensor_re = re.compile(
-        r"model\.layers\.(\d+)\.mlp\.experts\.(\d+)\.(down_proj|gate_proj|up_proj)\.weight"
-    )
+
+    moe_experts = [MOE_HF_MLP._replace(n_expert=("num_experts",))]
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         if (n_nextn := int(self.hparams.get("num_nextn_predict_layers", 0) or 0)) > 0 and not self.no_mtp:
             self.block_count += n_nextn
             self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)
-        self._experts: list[dict[str, Tensor]] = [{} for _ in range(self.block_count)]
 
     def _set_vocab_gpt2(self) -> None:
         tokens, toktypes, tokpre = self.get_vocab_base()
@@ -140,38 +138,4 @@ class Cohere2MoeModel(TextModel):
             logger.debug(f"Skipping bias tensor {name!r}.")
             return
 
-        if (m := self._expert_tensor_re.fullmatch(name)) is not None:
-            n_experts = self.hparams["num_experts"]
-            layer_idx = int(m.group(1))
-            assert bid is None or bid == layer_idx
-
-            self._experts[layer_idx][name] = data_torch
-
-            expected = {
-                f"model.layers.{layer_idx}.mlp.experts.{xid}.{w_name}.weight"
-                for xid in range(n_experts)
-                for w_name in ("down_proj", "gate_proj", "up_proj")
-            }
-            if expected.issubset(self._experts[layer_idx]):
-                for w_name in ["down_proj", "gate_proj", "up_proj"]:
-                    datas: list[Tensor] = []
-
-                    for xid in range(n_experts):
-                        ename = f"model.layers.{layer_idx}.mlp.experts.{xid}.{w_name}.weight"
-                        datas.append(self._experts[layer_idx][ename])
-                        del self._experts[layer_idx][ename]
-
-                    data_torch = torch.stack(datas, dim=0)
-                    merged_name = f"model.layers.{layer_idx}.mlp.experts.{w_name}.weight"
-
-                    yield from super().modify_tensors(data_torch, merged_name, layer_idx)
-            return
-
         yield from super().modify_tensors(data_torch, name, bid)
-
-    def prepare_tensors(self):
-        super().prepare_tensors()
-
-        experts = [k for d in self._experts for k in d.keys()]
-        if len(experts) > 0:
-            raise ValueError(f"Unprocessed experts: {experts}")

conversion/deepseek.py

@@ -12,7 +12,7 @@ import torch
 if TYPE_CHECKING:
     from torch import Tensor
 
-from .base import LazyTorchTensor, MmprojModel, ModelBase, TextModel, gguf, logger
+from .base import MOE_HF_MLP, LazyTorchTensor, MmprojModel, ModelBase, TextModel, gguf, logger
 
 from .qwen import QwenModel
 
@@ -158,7 +158,7 @@ class DeepseekModel(TextModel):
         self.gguf_writer.add_expert_count(hparams["n_routed_experts"])
         self.gguf_writer.add_expert_shared_count(hparams["n_shared_experts"])
 
-    _experts: list[dict[str, Tensor]] | None = None
+    moe_experts = [MOE_HF_MLP._replace(n_expert=("n_routed_experts",))]
 
     @staticmethod
     def permute(weights: Tensor, n_head: int, n_head_kv: int | None):
@@ -177,46 +177,8 @@ class DeepseekModel(TextModel):
         if name.endswith(("k_proj.weight", "k_proj.bias")):
             data_torch = DeepseekModel.permute(data_torch, n_head, n_kv_head)
 
-        # process the experts separately
-        if name.find("mlp.experts") != -1:
-            n_experts = self.hparams["n_routed_experts"]
-            assert bid is not None
-
-            if self._experts is None:
-                self._experts = [{} for _ in range(self.block_count)]
-
-            self._experts[bid][name] = data_torch
-
-            if len(self._experts[bid]) >= n_experts * 3:
-                # merge the experts into a single 3d tensor
-                for w_name in ["down_proj", "gate_proj", "up_proj"]:
-                    datas: list[Tensor] = []
-
-                    for xid in range(n_experts):
-                        ename = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight"
-                        datas.append(self._experts[bid][ename])
-                        del self._experts[bid][ename]
-
-                    data_torch = torch.stack(datas, dim=0)
-
-                    merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight"
-
-                    yield from super().modify_tensors(data_torch, merged_name, bid)
-                return
-            else:
-                return
-
         yield from super().modify_tensors(data_torch, name, bid)
 
-    def prepare_tensors(self):
-        super().prepare_tensors()
-
-        if self._experts is not None:
-            # flatten `list[dict[str, Tensor]]` into `list[str]`
-            experts = [k for d in self._experts for k in d.keys()]
-            if len(experts) > 0:
-                raise ValueError(f"Unprocessed experts: {experts}")
-
 
 @ModelBase.register(
     "DeepseekV2ForCausalLM",
@@ -380,7 +342,10 @@ class DeepseekV2Model(TextModel):
             # ref https://github.com/ggml-org/llama.cpp/pull/17945
             self.gguf_writer.add_rope_scaling_yarn_log_mul(0.1 * rope_mscale_all)
 
-    _experts: list[dict[str, Tensor]] | None = None
+    moe_experts = [MOE_HF_MLP._replace(n_expert=("n_routed_experts",))]
+
+    def moe_expert_specs(self):
+        return self.moe_experts if self.merge_expert else ()
 
     def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
         # skip lm_head.weight if tie_word_embeddings is True
@@ -396,35 +361,6 @@ class DeepseekV2Model(TextModel):
             if match and int(match.group(1)) >= block_count:
                 return
 
-        # process the experts separately
-        if self.merge_expert and name.find("mlp.experts") != -1:
-            n_experts = self.hparams["n_routed_experts"]
-            assert bid is not None
-
-            if self._experts is None:
-                self._experts = [{} for _ in range(self.block_count)]
-
-            self._experts[bid][name] = data_torch
-
-            if len(self._experts[bid]) >= n_experts * 3:
-                # merge the experts into a single 3d tensor
-                for w_name in ["down_proj", "gate_proj", "up_proj"]:
-                    datas: list[Tensor] = []
-
-                    for xid in range(n_experts):
-                        ename = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight"
-                        datas.append(self._experts[bid][ename])
-                        del self._experts[bid][ename]
-
-                    data_torch = torch.stack(datas, dim=0)
-
-                    merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight"
-
-                    yield from super().modify_tensors(data_torch, merged_name, bid)
-                return
-            else:
-                return
-
         # note: MLA with the absorption optimization, needs these two split and k_b_proj transposed
         if name.endswith("kv_b_proj.weight"):
             name_kb = name.replace("kv_b_proj", "k_b_proj")
@@ -446,15 +382,6 @@ class DeepseekV2Model(TextModel):
 
         yield from super().modify_tensors(data_torch, name, bid)
 
-    def prepare_tensors(self):
-        super().prepare_tensors()
-
-        if self._experts is not None:
-            # flatten `list[dict[str, Tensor]]` into `list[str]`
-            experts = [k for d in self._experts for k in d.keys()]
-            if len(experts) > 0:
-                raise ValueError(f"Unprocessed experts: {experts}")
-
 
 @ModelBase.register("DeepseekV32ForCausalLM")
 class DeepseekV32Model(DeepseekV2Model):

conversion/ernie.py

@@ -6,12 +6,10 @@ import re
 
 from typing import Callable, Iterable, TYPE_CHECKING
 
-import torch
-
 if TYPE_CHECKING:
     from torch import Tensor
 
-from .base import MmprojModel, ModelBase, TextModel, gguf
+from .base import MOE_HF_MLP, MmprojModel, ModelBase, TextModel, gguf
 
 
 @ModelBase.register("Ernie4_5_ForCausalLM", "Ernie4_5ForCausalLM")
@@ -75,11 +73,8 @@ class Ernie4_5Model(TextModel):
 @ModelBase.register("Ernie4_5_MoeForCausalLM")
 class Ernie4_5MoeModel(Ernie4_5Model):
     model_arch = gguf.MODEL_ARCH.ERNIE4_5_MOE
-    _experts: list[dict[str, Tensor]] | None = None
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self._experts = [{} for _ in range(self.block_count)]
+    moe_experts = [MOE_HF_MLP._replace(weights=("gate_proj", "up_proj", "down_proj"), n_expert=("moe_num_experts",))]
 
     def set_gguf_parameters(self):
         super().set_gguf_parameters()
@@ -119,40 +114,7 @@ class Ernie4_5MoeModel(Ernie4_5Model):
         return super().filter_tensors(item)
 
     def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
-        # process the experts separately
-        if name.find("mlp.experts") != -1:
-            n_experts = self.hparams["moe_num_experts"]
-            assert bid is not None
-
-            if self._experts is None:
-                self._experts = [{} for _ in range(self.block_count)]
-
-            self._experts[bid][name] = data_torch
-
-            if len(self._experts[bid]) >= n_experts * 3:
-                # merge the experts into a single 3d tensor
-                for w_name in ["gate_proj", "up_proj", "down_proj"]:
-                    datas: list[Tensor] = []
-
-                    for xid in range(n_experts):
-                        ename_to_retrieve = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight"
-                        datas.append(self._experts[bid][ename_to_retrieve])
-                        del self._experts[bid][ename_to_retrieve]
-
-                    data_torch = torch.stack(datas, dim=0)
-                    merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight"
-                    yield from super().modify_tensors(data_torch, merged_name, bid)
-        else:
-            yield from ModelBase.modify_tensors(self, data_torch, name, bid)
-
-    def prepare_tensors(self):
-        super().prepare_tensors()
-
-        if self._experts is not None:
-            # flatten `list[dict[str, Tensor]]` into `list[str]`
-            experts = [k for d in self._experts for k in d.keys()]
-            if len(experts) > 0:
-                raise ValueError(f"Unprocessed experts: {experts}")
+        yield from ModelBase.modify_tensors(self, data_torch, name, bid)
 
 
 @ModelBase.register("PaddleOCRVLForConditionalGeneration")

conversion/exaone.py

@@ -10,7 +10,7 @@ import torch
 if TYPE_CHECKING:
     from torch import Tensor
 
-from .base import MmprojModel, ModelBase, TextModel, gguf
+from .base import MOE_HF_MLP, MmprojModel, ModelBase, TextModel, gguf
 from .qwenvl import Qwen2VLVisionModel
 
 
@@ -149,7 +149,7 @@ class ExaoneMoEModel(Exaone4Model):
 
         self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.NONE)
 
-    _experts: list[dict[str, Tensor]] | None = None
+    moe_experts = [MOE_HF_MLP]
 
     def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
         if name.startswith("mtp."):
@@ -172,46 +172,8 @@ class ExaoneMoEModel(Exaone4Model):
                     yield from super().modify_tensors(data_torch, new_name.format(bid=bid), bid)
                 return
 
-        if name.find("mlp.experts") != -1:
-            n_experts = self.find_hparam(["num_local_experts", "num_experts"])
-            assert bid is not None
-
-            if self._experts is None:
-                self._experts = [{} for _ in range(self.block_count)]
-
-            self._experts[bid][name] = data_torch
-
-            if len(self._experts[bid]) >= n_experts * 3:
-                # merge the experts into a single 3d tensor
-                for w_name in ["down_proj", "gate_proj", "up_proj"]:
-                    datas: list[Tensor] = []
-
-                    for xid in range(n_experts):
-                        ename = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight"
-                        datas.append(self._experts[bid][ename])
-                        del self._experts[bid][ename]
-
-                    data_torch = torch.stack(datas, dim=0)
-
-                    merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight"
-
-                    new_name = self.map_tensor_name(merged_name)
-
-                    yield from super().modify_tensors(data_torch, new_name, bid)
-                return
-            else:
-                return
-
         yield from super().modify_tensors(data_torch, name, bid)
 
-    def prepare_tensors(self):
-        super().prepare_tensors()
-        if self._experts is not None:
-            # flatten `list[dict[str, Tensor]]` into `list[str]`
-            experts = [k for d in self._experts for k in d.keys()]
-            if len(experts) > 0:
-                raise ValueError(f"Unprocessed experts: {experts}")
-
 
 @ModelBase.register("Exaone4_5_ForConditionalGeneration")
 class Exaone4_5_TextModel(Exaone4Model):

conversion/glm.py

@@ -9,7 +9,7 @@ import torch
 if TYPE_CHECKING:
     from torch import Tensor
 
-from .base import ModelBase, TextModel, gguf, logger
+from .base import MOE_HF_MLP, ModelBase, TextModel, gguf, logger
 
 from .deepseek import DeepseekV2Model
 
@@ -154,7 +154,7 @@ class Glm4MoeModel(TextModel):
         if (num_nextn_predict_layers := self.hparams.get("num_nextn_predict_layers")) is not None:
             self.gguf_writer.add_nextn_predict_layers(num_nextn_predict_layers)
 
-    _experts: list[dict[str, Tensor]] | None = None
+    moe_experts = [MOE_HF_MLP._replace(n_expert=("n_routed_experts",))]
 
     # note: unlike GLM4V non-MoE, we don't need to permute Q/K here since GLM4V_MOE uses Neox ordering already
     def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
@@ -163,45 +163,8 @@ class Glm4MoeModel(TextModel):
             yield from super().modify_tensors(data_torch, "token_embd.weight", bid)
             return
 
-        # Handle routed experts
-        if name.find("mlp.experts") != -1:
-            n_experts = self.hparams["n_routed_experts"]
-            assert bid is not None
-
-            if self._experts is None:
-                self._experts = [{} for _ in range(self.block_count)]
-
-            self._experts[bid][name] = data_torch
-
-            if len(self._experts[bid]) >= n_experts * 3:
-                # merge the experts into a single 3d tensor
-                for w_name in ["down_proj", "gate_proj", "up_proj"]:
-                    datas: list[Tensor] = []
-
-                    for xid in range(n_experts):
-                        ename = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight"
-                        datas.append(self._experts[bid][ename])
-                        del self._experts[bid][ename]
-
-                    data_torch = torch.stack(datas, dim=0)
-
-                    merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight"
-
-                    yield from super().modify_tensors(data_torch, merged_name, bid)
-                return
-            else:
-                return
-
         yield from super().modify_tensors(data_torch, name, bid)
 
-    def prepare_tensors(self):
-        super().prepare_tensors()
-        if self._experts is not None:
-            # flatten `list[dict[str, Tensor]]` into `list[str]`
-            experts = [k for d in self._experts for k in d.keys()]
-            if len(experts) > 0:
-                raise ValueError(f"Unprocessed experts: {experts}")
-
 
 @ModelBase.register("Glm4MoeLiteForCausalLM")
 class Glm4MoeLiteModel(DeepseekV2Model):

conversion/grovemoe.py

@@ -2,12 +2,10 @@ from __future__ import annotations
 
 from typing import Iterable, TYPE_CHECKING
 
-import torch
-
 if TYPE_CHECKING:
     from torch import Tensor
 
-from .base import ModelBase, TextModel, gguf, logger
+from .base import MOE_HF_MLP, ModelBase, TextModel, gguf, logger
 
 
 @ModelBase.register("GroveMoeForCausalLM", "modeling_grove_moe.GroveMoeForCausalLM")
@@ -26,83 +24,19 @@ class GroveMoeModel(TextModel):
         # FIXME?: Hardcoded https://huggingface.co/inclusionAI/GroveMoE-Inst/blob/c4c69e5970d18907b5e6ddccdfd55176fe292df1/modeling_grove_moe.py#L376
         self.gguf_writer.add_expert_group_scale(0.05)
 
-    _experts: list[dict[str, Tensor]] | None = None
-    _chunk_experts: list[dict[str, Tensor]] | None = None
+    moe_experts = [
+        MOE_HF_MLP,
+        # chunk experts are grouped in pairs, see add_experts_per_group
+        MOE_HF_MLP._replace(
+            src="model.layers.{bid}.mlp.chunk_experts.{xid}.{w}.weight",
+            dst="model.layers.{bid}.mlp.chunk_experts.{w}.weight",
+            n_expert=lambda model: model.find_hparam(["num_local_experts", "num_experts"]) // 2,
+        ),
+    ]
 
     def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
         if name.endswith(".expert_bias"):
             # FIXME?: Unused https://huggingface.co/inclusionAI/GroveMoE-Inst/blob/c4c69e5970d18907b5e6ddccdfd55176fe292df1/modeling_grove_moe.py#L303
             return
 
-        # process the experts separately
-        if name.find("chunk_experts") != -1:
-            n_experts = self.find_hparam(["num_local_experts", "num_experts"]) // 2 # see add_experts_per_group
-            assert bid is not None
-
-            if self._chunk_experts is None:
-                self._chunk_experts = [{} for _ in range(self.block_count)]
-
-            self._chunk_experts[bid][name] = data_torch
-
-            if len(self._chunk_experts[bid]) >= n_experts * 3:
-                # merge the experts into a single 3d tensor
-                for w_name in ["down_proj", "gate_proj", "up_proj"]:
-                    datas: list[Tensor] = []
-
-                    for xid in range(n_experts):
-                        ename = f"model.layers.{bid}.mlp.chunk_experts.{xid}.{w_name}.weight"
-                        datas.append(self._chunk_experts[bid][ename])
-                        del self._chunk_experts[bid][ename]
-
-                    data_torch = torch.stack(datas, dim=0)
-
-                    merged_name = f"model.layers.{bid}.mlp.chunk_experts.{w_name}.weight"
-
-                    yield from super().modify_tensors(data_torch, merged_name, bid)
-                return
-            else:
-                return
-        elif name.find("experts") != -1:
-            n_experts = self.find_hparam(["num_local_experts", "num_experts"])
-            assert bid is not None
-
-            if self._experts is None:
-                self._experts = [{} for _ in range(self.block_count)]
-
-            self._experts[bid][name] = data_torch
-
-            if len(self._experts[bid]) >= n_experts * 3:
-                # merge the experts into a single 3d tensor
-                for w_name in ["down_proj", "gate_proj", "up_proj"]:
-                    datas: list[Tensor] = []
-
-                    for xid in range(n_experts):
-                        ename = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight"
-                        datas.append(self._experts[bid][ename])
-                        del self._experts[bid][ename]
-
-                    data_torch = torch.stack(datas, dim=0)
-
-                    merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight"
-
-                    yield from super().modify_tensors(data_torch, merged_name, bid)
-                return
-            else:
-                return
-
         yield from super().modify_tensors(data_torch, name, bid)
-
-    def prepare_tensors(self):
-        super().prepare_tensors()
-
-        if self._chunk_experts is not None:
-            # flatten `list[dict[str, Tensor]]` into `list[str]`
-            chunk_experts = [k for d in self._chunk_experts for k in d.keys()]
-            if len(chunk_experts) > 0:
-                raise ValueError(f"Unprocessed adjugate experts: {chunk_experts}")
-
-        if self._experts is not None:
-            # flatten `list[dict[str, Tensor]]` into `list[str]`
-            experts = [k for d in self._experts for k in d.keys()]
-            if len(experts) > 0:
-                raise ValueError(f"Unprocessed experts: {experts}")

conversion/hunyuan.py

@@ -6,12 +6,10 @@ import re
 from pathlib import Path
 from typing import Callable, Iterable, TYPE_CHECKING
 
-import torch
-
 if TYPE_CHECKING:
     from torch import Tensor
 
-from .base import MmprojModel, ModelBase, TextModel, gguf, logger
+from .base import MOE_HF_MLP, MmprojModel, ModelBase, TextModel, gguf, logger
 
 from .qwen import QwenModel
 
@@ -108,7 +106,7 @@ class HunYuanMoEModel(TextModel):
             assert alpha == 1000 and base == 10000.0 and dim == 128 and self.hparams["max_position_embeddings"] in [32 * 1024, 256 * 1024] , \
                 "HunYuan dynamic RoPE scaling assumptions changed, please update the logic or context length manually"
 
-    _experts: list[dict[str, Tensor]] | None = None
+    moe_experts = [MOE_HF_MLP]
 
     def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
         if name == "lm_head.weight":
@@ -116,42 +114,8 @@ class HunYuanMoEModel(TextModel):
                 logger.info("Skipping tied output layer 'lm_head.weight'")
                 return
 
-        if name.find("mlp.experts") != -1:
-            n_experts = self.find_hparam(["num_local_experts", "num_experts"])
-            assert bid is not None
-
-            if self._experts is None:
-                self._experts = [{} for _ in range(self.block_count)]
-
-            self._experts[bid][name] = data_torch
-
-            if len(self._experts[bid]) >= n_experts * 3:
-                # merge the experts into a single 3d tensor
-                for w_name in ["down_proj", "gate_proj", "up_proj"]:
-                    datas: list[Tensor] = []
-
-                    for xid in range(n_experts):
-                        ename = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight"
-                        datas.append(self._experts[bid][ename])
-                        del self._experts[bid][ename]
-
-                    data_torch = torch.stack(datas, dim=0)
-                    merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight"
-
-                    yield from super().modify_tensors(data_torch, merged_name, bid)
-                return
-            else:
-                return
-
         yield from super().modify_tensors(data_torch, name, bid)
 
-    def prepare_tensors(self):
-        super().prepare_tensors()
-        if self._experts is not None:
-            experts = [k for d in self._experts for k in d.keys()]
-            if len(experts) > 0:
-                raise ValueError(f"Unprocessed experts: {experts}")
-
 
 @ModelBase.register("HunYuanDenseV1ForCausalLM")
 class HunYuanModel(TextModel):
@@ -432,36 +396,4 @@ class HYV3Model(TextModel):
 
         return name, gen
 
-    _experts: list[dict[str, Tensor]] | None = None
-
-    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
-        # merge the per-expert tensors into stacked 3d tensors
-        if name.startswith("model.layers.") and ".mlp.experts." in name:
-            n_experts = self.find_hparam(["num_local_experts", "num_experts"])
-            assert bid is not None
-
-            if self._experts is None:
-                self._experts = [{} for _ in range(self.block_count)]
-
-            self._experts[bid][name] = data_torch
-
-            if len(self._experts[bid]) >= n_experts * 3:
-                for w_name in ("down_proj", "gate_proj", "up_proj"):
-                    datas: list[Tensor] = []
-                    for xid in range(n_experts):
-                        ename = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight"
-                        datas.append(self._experts[bid][ename])
-                        del self._experts[bid][ename]
-
-                    merged = torch.stack(datas, dim=0)
-                    yield from super().modify_tensors(merged, f"model.layers.{bid}.mlp.experts.{w_name}.weight", bid)
-            return
-
-        yield from super().modify_tensors(data_torch, name, bid)
-
-    def prepare_tensors(self):
-        super().prepare_tensors()
-        if self._experts is not None:
-            experts = [k for d in self._experts for k in d.keys()]
-            if experts:
-                raise ValueError(f"Unprocessed experts: {experts}")
+    moe_experts = [MOE_HF_MLP]

conversion/kimi_k3.py

@@ -10,7 +10,7 @@ import torch
 if TYPE_CHECKING:
     from torch import Tensor
 
-from .base import LazyTorchTensor, ModelBase, TextModel, gguf, logger
+from .base import MOE_BLOCK_SPARSE, LazyTorchTensor, ModelBase, TextModel, gguf, logger
 
 from .kimi_linear import KimiLinearModel
 
@@ -29,7 +29,8 @@ class KimiK3Model(TextModel):
 
     model_arch = gguf.MODEL_ARCH.KIMI_K3
 
-    _experts: list[dict[str, Tensor]] | None = None
+    # w1: gate, w2: down, w3: up
+    moe_experts = [MOE_BLOCK_SPARSE._replace(n_expert=("num_experts",))]
 
     # `<x>_res_norm.weight` and `<x>_res_proj.weight` are only used as their
     # elementwise product, so they are fused into one [n_embd] vector here.
@@ -258,10 +259,6 @@ class KimiK3Model(TextModel):
 
     def prepare_tensors(self):
         super().prepare_tensors()
-        if self._experts is not None:
-            leftover = [k for d in self._experts for k in d.keys()]
-            if leftover:
-                raise ValueError(f"Unprocessed experts: {leftover}")
         if self._res_parts:
             raise ValueError(f"Unpaired attention-residual tensors: {sorted(self._res_parts)}")
         if self._is_mxfp4_packed():
@@ -335,30 +332,6 @@ class KimiK3Model(TextModel):
             yield self.format_tensor_name(tensor_id, bid), data_torch
             return
 
-        # --- routed experts: stack per-expert 2D weights into one 3D tensor ---
-        if ".block_sparse_moe.experts." in name:
-            n_experts = self.hparams["num_experts"]
-            assert bid is not None
-
-            if self._experts is None:
-                self._experts = [{} for _ in range(self.block_count)]
-            self._experts[bid][name] = data_torch
-
-            if len(self._experts[bid]) < n_experts * 3:
-                return
-
-            # w1: gate, w2: down, w3: up
-            for wid, tensor_id in (("w1", gguf.MODEL_TENSOR.FFN_GATE_EXP),
-                                   ("w2", gguf.MODEL_TENSOR.FFN_DOWN_EXP),
-                                   ("w3", gguf.MODEL_TENSOR.FFN_UP_EXP)):
-                datas = []
-                for xid in range(n_experts):
-                    ename = f"model.layers.{bid}.block_sparse_moe.experts.{xid}.{wid}.weight"
-                    datas.append(self._experts[bid].pop(ename))
-                stacked = torch.stack(datas, dim=0)
-                yield from super().modify_tensors(stacked, self.format_tensor_name(tensor_id, bid), bid)
-            return
-
         # --- MLA absorption: split kv_b into k_b (transposed) and v_b ---
         if name.endswith("kv_b_proj.weight"):
             n_head_kv = self.hparams["num_key_value_heads"]

conversion/kimi_linear.py

@@ -7,7 +7,7 @@ import torch
 if TYPE_CHECKING:
     from torch import Tensor
 
-from .base import ModelBase, TextModel, gguf, logger
+from .base import MOE_BLOCK_SPARSE, ModelBase, TextModel, gguf, logger
 
 from .qwen import QwenModel
 
@@ -17,7 +17,8 @@ class KimiLinearModel(TextModel):
     """Kimi-Linear model with hybrid MLA+KDA architecture"""
     model_arch = gguf.MODEL_ARCH.KIMI_LINEAR
 
-    _experts: list[dict[str, Tensor]] | None = None
+    # w1: gate, w2: down, w3: up
+    moe_experts = [MOE_BLOCK_SPARSE]
 
     def set_vocab(self):
         try:
@@ -140,13 +141,6 @@ class KimiLinearModel(TextModel):
         # Routed scaling factor (expert_weights_scale = 2.446 for Kimi)
         self.gguf_writer.add_expert_weights_scale(self.hparams["routed_scaling_factor"])
 
-    def prepare_tensors(self):
-        super().prepare_tensors()
-        if self._experts is not None:
-            experts = [k for d in self._experts for k in d.keys()]
-            if len(experts) > 0:
-                raise ValueError(f"Unprocessed experts: {experts}")
-
     def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
         logger.info(f"Processing {name}: shape before = {tuple(data_torch.shape)}")
 
@@ -178,32 +172,6 @@ class KimiLinearModel(TextModel):
             name = name.rpartition(".dt_bias")[0] + ".dt_proj.bias"
             logger.info("Changed dt_bias to dt_proj.bias")
 
-        # process the experts separately
-        if name.find("block_sparse_moe.experts") != -1:
-            n_experts = self.find_hparam(["num_local_experts", "num_experts"])
-            assert bid is not None
-
-            if self._experts is None:
-                self._experts = [{} for _ in range(self.block_count)]
-
-            self._experts[bid][name] = data_torch
-
-            if len(self._experts[bid]) >= n_experts * 3:
-                # merge the experts into a single 3d tensor
-                # w1: gate, w2: down, w3: up
-                for wid, tname in [("w1", gguf.MODEL_TENSOR.FFN_GATE_EXP),
-                                   ("w2", gguf.MODEL_TENSOR.FFN_DOWN_EXP),
-                                   ("w3", gguf.MODEL_TENSOR.FFN_UP_EXP)]:
-                    datas: list[Tensor] = []
-                    for xid in range(n_experts):
-                        ename = f"model.layers.{bid}.block_sparse_moe.experts.{xid}.{wid}.weight"
-                        datas.append(self._experts[bid][ename])
-                        del self._experts[bid][ename]
-                    data_torch = torch.stack(datas, dim=0)
-                    new_name = self.format_tensor_name(tname, bid)
-                    yield from super().modify_tensors(data_torch, new_name, bid)
-            return
-
         # note: MLA with the absorption optimization, needs these two split and k_b_proj transposed
         if name.endswith("kv_b_proj.weight"):
             name_kb = name.replace("kv_b_proj", "k_b_proj")

conversion/laguna.py

@@ -1,23 +1,21 @@
 from __future__ import annotations
 
-import re
 from collections.abc import Iterable
 from typing import TYPE_CHECKING
 
-import torch
-
 if TYPE_CHECKING:
     from torch import Tensor
 
-from .base import ModelBase, TextModel, gguf, logger
+from .base import MOE_HF_MLP, ModelBase, TextModel, gguf, logger
 
 
 @ModelBase.register("LagunaForCausalLM")
 class LagunaModel(TextModel):
     model_arch = gguf.MODEL_ARCH.LAGUNA
-    _experts: list[dict] | None = None
     _gate_types: list[str] | None = None
 
+    moe_experts = [MOE_HF_MLP._replace(weights=("gate_proj", "up_proj", "down_proj"))]
+
     # --- vocab ---------------------------------------------------------------
 
     def set_vocab(self) -> None:
@@ -168,27 +166,6 @@ class LagunaModel(TextModel):
     # --- tensor handling -----------------------------------------------------
 
     def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
-        # Per-expert MoE weights: model.layers.{bid}.mlp.experts.{xid}.{w}.weight.
-        # Only the NUMBERED per-expert weights are stacked; the router bias
-        # (mlp.experts.e_score_correction_bias) takes the normal mapping path.
-        if re.search(r"mlp\.experts\.\d+\.", name):
-            n_experts = self.find_hparam(["num_local_experts", "num_experts"])
-            assert bid is not None
-            if self._experts is None:
-                self._experts = [{} for _ in range(self.block_count)]
-            self._experts[bid][name] = data_torch
-            needed = [f"model.layers.{bid}.mlp.experts.{x}.{w}.weight"
-                      for x in range(n_experts) for w in ("gate_proj", "up_proj", "down_proj")]
-            if all(e in self._experts[bid] for e in needed):
-                for w_name in ["gate_proj", "up_proj", "down_proj"]:
-                    datas = [self._experts[bid][f"model.layers.{bid}.mlp.experts.{x}.{w_name}.weight"]
-                             for x in range(n_experts)]
-                    stacked = torch.stack(datas, dim=0)
-                    merged = f"model.layers.{bid}.mlp.experts.{w_name}.weight"
-                    yield from TextModel.modify_tensors(self, stacked, merged, bid)
-                self._experts[bid].clear()
-                return
-            return
         # Cross-check the gate projection width against the declared gate type;
         # a mismatch means the weights and config disagree -> fail, do not guess.
         if bid is not None and name.endswith("self_attn.g_proj.weight"):

conversion/lfm2.py

@@ -2,12 +2,10 @@ from __future__ import annotations
 
 from typing import Any, Callable, Iterable, TYPE_CHECKING
 
-import torch
-
 if TYPE_CHECKING:
     from torch import Tensor
 
-from .base import MmprojModel, ModelBase, TextModel, gguf
+from .base import MOE_BLOCK_SPARSE, MmprojModel, ModelBase, TextModel, gguf
 
 from .gemma import ConformerAudioModel
 
@@ -112,8 +110,10 @@ class LFM2MoeModel(TextModel):
         self.gguf_writer.add_vocab_size(self.hparams["vocab_size"])
         self.gguf_writer.add_shortconv_l_cache(self.hparams["conv_L_cache"])
 
-    # cache for experts weights for merging
-    _experts_cache: dict[int, dict[str, Tensor]] = {}
+    moe_experts = [MOE_BLOCK_SPARSE._replace(
+        src="model.layers.{bid}.feed_forward.experts.{xid}.{w}.weight",
+        dst="layers.{bid}.feed_forward.experts.{w}.weight",
+    )]
 
     @classmethod
     def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
@@ -129,41 +129,8 @@ class LFM2MoeModel(TextModel):
         if 'conv.conv' in name:
             data_torch = data_torch.squeeze(1)
 
-        # merge expert weights
-        if 'experts' in name:
-            n_experts = self.find_hparam(["num_local_experts", "num_experts"])
-            assert bid is not None
-
-            expert_cache = self._experts_cache.setdefault(bid, {})
-            expert_cache[name] = data_torch
-            expert_weights = ["w1", "w2", "w3"]
-
-            # not enough expert weights to merge
-            if len(expert_cache) < n_experts * len(expert_weights):
-                return
-
-            for w_name in expert_weights:
-                datas: list[Tensor] = []
-
-                for xid in range(n_experts):
-                    ename = f"model.layers.{bid}.feed_forward.experts.{xid}.{w_name}.weight"
-                    datas.append(expert_cache[ename])
-                    del expert_cache[ename]
-
-                data_torch = torch.stack(datas, dim=0)
-                merged_name = f"layers.{bid}.feed_forward.experts.{w_name}.weight"
-
-                yield from super().modify_tensors(data_torch, merged_name, bid)
-
-            del self._experts_cache[bid]
-            return
-
         yield from super().modify_tensors(data_torch, name, bid)
 
-    def prepare_tensors(self):
-        super().prepare_tensors()
-        assert not self._experts_cache
-
 
 @ModelBase.register("Lfm2VlForConditionalGeneration")
 class LFM2VLModel(MmprojModel):

conversion/llada.py

@@ -2,12 +2,10 @@ from __future__ import annotations
 
 from typing import Iterable, TYPE_CHECKING
 
-import torch
-
 if TYPE_CHECKING:
     from torch import Tensor
 
-from .base import ModelBase, TextModel, gguf
+from .base import MOE_HF_MLP, ModelBase, TextModel, gguf
 
 
 @ModelBase.register("LLaDAModelLM")
@@ -126,47 +124,4 @@ class LLaDAMoEModel(TextModel):
         self.gguf_writer.add_causal_attention(False)
         self.gguf_writer.add_diffusion_shift_logits(False)
 
-    _experts: list[dict[str, Tensor]] | None = None
-
-    # Copied from: Qwen2MoeModel
-    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
-        # process the experts separately
-        if name.find("experts") != -1:
-            n_experts = self.find_hparam(["num_local_experts", "num_experts"])
-            assert bid is not None
-
-            if self._experts is None:
-                self._experts = [{} for _ in range(self.block_count)]
-
-            self._experts[bid][name] = data_torch
-
-            if len(self._experts[bid]) >= n_experts * 3:
-                # merge the experts into a single 3d tensor
-                for w_name in ["down_proj", "gate_proj", "up_proj"]:
-                    datas: list[Tensor] = []
-
-                    for xid in range(n_experts):
-                        ename = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight"
-                        datas.append(self._experts[bid][ename])
-                        del self._experts[bid][ename]
-
-                    data_torch = torch.stack(datas, dim=0)
-
-                    merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight"
-
-                    yield from super().modify_tensors(data_torch, merged_name, bid)
-                return
-            else:
-                return
-
-        yield from super().modify_tensors(data_torch, name, bid)
-
-    # Copied from: Qwen2MoeModel
-    def prepare_tensors(self):
-        super().prepare_tensors()
-
-        if self._experts is not None:
-            # flatten `list[dict[str, Tensor]]` into `list[str]`
-            experts = [k for d in self._experts for k in d.keys()]
-            if len(experts) > 0:
-                raise ValueError(f"Unprocessed experts: {experts}")
+    moe_experts = [MOE_HF_MLP]

conversion/llama.py

@@ -11,7 +11,7 @@ import torch
 if TYPE_CHECKING:
     from torch import Tensor
 
-from .base import ModelBase, TextModel, gguf, logger
+from .base import MOE_BLOCK_SPARSE_MIXTRAL, ModelBase, TextModel, gguf, logger
 
 
 @ModelBase.register(
@@ -192,7 +192,7 @@ class LlamaModel(TextModel):
                     scale  = LlamaModel.permute(scale, n_head, n_kv_head)
         super()._repack_nvfp4(name, weight, scale, scale2, input_scale)
 
-    _experts: list[dict[str, Tensor]] | None = None
+    moe_experts = [MOE_BLOCK_SPARSE_MIXTRAL._replace(n_expert=("num_local_experts",))]
 
     @classmethod
     def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
@@ -261,36 +261,6 @@ class LlamaModel(TextModel):
             if name.endswith(("k_proj.weight", "k_proj.bias")):
                 data_torch = LlamaModel.permute(data_torch, n_head, n_kv_head)
 
-        # process the experts separately
-        if name.find("block_sparse_moe.experts") != -1:
-            n_experts = self.hparams["num_local_experts"]
-
-            assert bid is not None
-
-            if self._experts is None:
-                self._experts = [{} for _ in range(self.block_count)]
-
-            self._experts[bid][name] = data_torch
-
-            if len(self._experts[bid]) >= n_experts * 3:
-                # merge the experts into a single 3d tensor
-                for wid in ["w1", "w2", "w3"]:
-                    datas: list[Tensor] = []
-
-                    for xid in range(n_experts):
-                        ename = f"model.layers.{bid}.block_sparse_moe.experts.{xid}.{wid}.weight"
-                        datas.append(self._experts[bid][ename])
-                        del self._experts[bid][ename]
-
-                    data_torch = torch.stack(datas, dim=0)
-
-                    merged_name = f"layers.{bid}.feed_forward.experts.{wid}.weight"
-
-                    yield from super().modify_tensors(data_torch, merged_name, bid)
-                return
-            else:
-                return
-
         yield from super().modify_tensors(data_torch, name, bid)
 
     def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:
@@ -351,12 +321,6 @@ class LlamaModel(TextModel):
                 logger.info(f"{name + ',':<30} {old_dtype} --> {data_qtype.name}, shape = {shape_str}")
                 self.gguf_writer.add_tensor(name, data, raw_dtype=data_qtype)
 
-        if self._experts is not None:
-            # flatten `list[dict[str, Tensor]]` into `list[str]`
-            experts = [k for d in self._experts for k in d.keys()]
-            if len(experts) > 0:
-                raise ValueError(f"Unprocessed experts: {experts}")
-
 
 @ModelBase.register("ArceeForCausalLM")
 class ArceeModel(LlamaModel):

conversion/mellum.py

@@ -1,13 +1,6 @@
 from __future__ import annotations
 
-from typing import Iterable, TYPE_CHECKING
-
-import torch
-
-if TYPE_CHECKING:
-    from torch import Tensor
-
-from .base import ModelBase, TextModel, gguf, logger
+from .base import MOE_HF_MLP, ModelBase, TextModel, gguf, logger
 
 
 @ModelBase.register("MellumForCausalLM")
@@ -28,34 +21,4 @@ class MellumModel(TextModel):
             self.gguf_writer.add_sliding_window_pattern([t == "sliding_attention" for t in self.hparams["layer_types"]])
             logger.info(f"gguf: sliding window pattern length = {len(self.hparams['layer_types'])}")
 
-    _experts: list[dict[str, Tensor]] | None = None
-
-    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
-        if name.find("experts") != -1:
-            n_experts = self.find_hparam(["num_local_experts", "num_experts"])
-            assert bid is not None
-
-            if self._experts is None:
-                self._experts = [{} for _ in range(self.block_count)]
-
-            self._experts[bid][name] = data_torch
-
-            if len(self._experts[bid]) >= n_experts * 3:
-                for w_name in ["down_proj", "gate_proj", "up_proj"]:
-                    datas: list[Tensor] = []
-
-                    for xid in range(n_experts):
-                        ename = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight"
-                        datas.append(self._experts[bid][ename])
-                        del self._experts[bid][ename]
-
-                    data_torch = torch.stack(datas, dim=0)
-
-                    merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight"
-
-                    yield from super().modify_tensors(data_torch, merged_name, bid)
-                return
-            else:
-                return
-
-        yield from super().modify_tensors(data_torch, name, bid)
+    moe_experts = [MOE_HF_MLP]

conversion/mimo.py

@@ -10,7 +10,7 @@ import torch
 if TYPE_CHECKING:
     from torch import Tensor
 
-from .base import MmprojModel, ModelBase, TextModel, gguf
+from .base import MOE_HF_MLP, MmprojModel, ModelBase, TextModel, gguf
 
 
 @ModelBase.register("MiMoV2FlashForCausalLM", "MiMoV2ForCausalLM")
@@ -166,7 +166,16 @@ class MimoV2Model(TextModel):
 
         self.gguf_writer.add_nextn_predict_layers(self._n_nextn)
 
-    _experts: list[dict[str, Tensor]] | None = None
+    # MTP experts keep their mtp name until modify_tensors() remaps them below
+    moe_experts = [
+        MOE_HF_MLP._replace(weights=("gate_proj", "up_proj", "down_proj"), n_expert=("n_routed_experts",)),
+        MOE_HF_MLP._replace(
+            src="model.mtp.layers.{bid}.mlp.experts.{xid}.{w}.weight",
+            dst="model.mtp.layers.{bid}.mlp.experts.{w}.weight",
+            weights=("gate_proj", "up_proj", "down_proj"),
+            n_expert=("n_routed_experts",),
+        ),
+    ]
 
     @classmethod
     def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
@@ -190,44 +199,8 @@ class MimoV2Model(TextModel):
             name = f"model.layers.{new_bid}.{rest}"
             bid = new_bid
 
-        # process the experts separately
-        if name.find("mlp.experts") != -1:
-            n_experts = self.hparams["n_routed_experts"]
-            assert bid is not None
-
-            if self._experts is None:
-                self._experts = [{} for _ in range(self.block_count)]
-
-            self._experts[bid][name] = data_torch
-
-            if len(self._experts[bid]) >= n_experts * 3:
-                # merge the experts into a single 3d tensor
-                for w_name in ["gate_proj", "up_proj", "down_proj"]:
-                    datas: list[Tensor] = []
-
-                    for xid in range(n_experts):
-                        ename_to_retrieve = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight"
-                        datas.append(self._experts[bid][ename_to_retrieve])
-                        del self._experts[bid][ename_to_retrieve]
-
-                    data_torch = torch.stack(datas, dim=0)
-                    merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight"
-
-                    yield from super().modify_tensors(data_torch, merged_name, bid)
-                return
-            else:
-                return
         yield from super().modify_tensors(data_torch, name, bid)
 
-    def prepare_tensors(self):
-        super().prepare_tensors()
-
-        if self._experts is not None:
-            # flatten `list[dict[str, Tensor]]` into `list[str]`
-            experts = [k for d in self._experts for k in d.keys()]
-            if len(experts) > 0:
-                raise ValueError(f"Unprocessed experts: {experts}")
-
 
 @ModelBase.register("MiMoV2ForCausalLM")
 class MiMoV2VisionAudioModel(MmprojModel):

conversion/minimax.py

@@ -1,13 +1,13 @@
 from __future__ import annotations
 
-from typing import Iterable, Sequence, TYPE_CHECKING
+from typing import Sequence, TYPE_CHECKING
 
 import torch
 
 if TYPE_CHECKING:
     from torch import Tensor
 
-from .base import ModelBase, TextModel, MmprojModel, gguf, logger
+from .base import MOE_BLOCK_SPARSE, MOE_BLOCK_SPARSE_MIXTRAL, ModelBase, TextModel, MmprojModel, gguf, logger
 
 
 @ModelBase.register("MiniMaxText01ForCausalLM")
@@ -80,48 +80,14 @@ class MiniMaxText01Model(TextModel):
 
         self.gguf_writer.add_rope_dimension_count(self.hparams["rotary_dim"])
 
-    _experts: list[dict[str, Tensor]] | None = None
-
-    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
-        # process the experts separately
-        if name.find("block_sparse_moe.experts") != -1:
-            n_experts = self.hparams["num_local_experts"]
-
-            assert bid is not None
-
-            if self._experts is None:
-                self._experts = [{} for _ in range(self.block_count)]
-
-            self._experts[bid][name] = data_torch
-
-            if len(self._experts[bid]) >= n_experts * 3:
-                # merge the experts into a single 3d tensor
-                for wid in ["w1", "w2", "w3"]:
-                    datas: list[Tensor] = []
-
-                    for xid in range(n_experts):
-                        ename = f"model.layers.{bid}.block_sparse_moe.experts.{xid}.{wid}.weight"
-                        datas.append(self._experts[bid][ename])
-                        del self._experts[bid][ename]
-
-                    data_torch = torch.stack(datas, dim=0)
-
-                    merged_name = f"layers.{bid}.feed_forward.experts.{wid}.weight"
-
-                    new_name = self.map_tensor_name(merged_name)
-
-                    yield from super().modify_tensors(data_torch, new_name, bid)
-                return
-            else:
-                return
-
-        yield from super().modify_tensors(data_torch, name, bid)
+    moe_experts = [MOE_BLOCK_SPARSE_MIXTRAL._replace(n_expert=("num_local_experts",))]
 
 
 @ModelBase.register("MiniMaxM2ForCausalLM")
 class MiniMaxM2Model(TextModel):
     model_arch = gguf.MODEL_ARCH.MINIMAXM2
-    _experts_cache: dict[int, dict[str, Tensor]] = {}
+
+    moe_experts = [MOE_BLOCK_SPARSE]
 
     def set_gguf_parameters(self):
         super().set_gguf_parameters()
@@ -129,38 +95,6 @@ class MiniMaxM2Model(TextModel):
         self.gguf_writer.add_expert_feed_forward_length(self.find_hparam(["intermediate_size"]))
         self.gguf_writer.add_rope_dimension_count(self.find_hparam(["rotary_dim"]))
 
-    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None):
-        # merge expert weights
-        if "block_sparse_moe.experts." in name:
-            n_experts = self.find_hparam(["num_local_experts", "num_experts"])
-            assert bid is not None
-
-            expert_cache = self._experts_cache.setdefault(bid, {})
-            expert_cache[name] = data_torch
-            expert_weights = ["w1", "w2", "w3"]
-
-            # not enough expert weights to merge
-            if len(expert_cache) < n_experts * len(expert_weights):
-                return
-
-            for w_name in expert_weights:
-                datas: list[Tensor] = []
-
-                for xid in range(n_experts):
-                    ename = f"model.layers.{bid}.block_sparse_moe.experts.{xid}.{w_name}.weight"
-                    datas.append(expert_cache[ename])
-                    del expert_cache[ename]
-
-                data_torch = torch.stack(datas, dim=0)
-                merged_name = f"model.layers.{bid}.block_sparse_moe.experts.{w_name}.weight"
-                new_name = self.map_tensor_name(merged_name)
-                yield from super().modify_tensors(data_torch, new_name, bid)
-
-            del self._experts_cache[bid]
-            return
-
-        yield from super().modify_tensors(data_torch, name, bid)
-
 
 @ModelBase.register("MiniMaxM3SparseForCausalLM", "MiniMaxM3SparseForConditionalGeneration")
 class MiniMaxM3Model(MiniMaxM2Model):

conversion/nemotron.py

@@ -412,6 +412,11 @@ class NemotronHModel(GraniteHybridModel):
         if not self.is_moe:
             self.gguf_writer.add_add_bos_token(True)
 
+    # the MTP experts are renamed in modify_tensors() before they are stacked,
+    # so this arch keeps its own stacking instead of using moe_experts
+    moe_experts = []
+    _experts: list[dict[str, Tensor]] | None = None
+
     _MTP_SPECIAL_RENAMES = {
         "mtp.layers.0.enorm.weight":           "model.layers.{bid}.enorm.weight",
         "mtp.layers.0.hnorm.weight":           "model.layers.{bid}.hnorm.weight",

conversion/olmo.py

@@ -2,12 +2,10 @@ from __future__ import annotations
 
 from typing import Iterable, TYPE_CHECKING
 
-import torch
-
 if TYPE_CHECKING:
     from torch import Tensor
 
-from .base import ModelBase, TextModel, gguf
+from .base import MOE_HF_MLP, ModelBase, TextModel, gguf
 
 from .llama import LlamaModel
 
@@ -74,47 +72,4 @@ class OlmoeModel(TextModel):
         super().set_gguf_parameters()
         self.gguf_writer.add_layer_norm_rms_eps(1e-5)
 
-    _experts: list[dict[str, Tensor]] | None = None
-
-    # Copied from: Qwen2MoeModel
-    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
-        # process the experts separately
-        if name.find("experts") != -1:
-            n_experts = self.find_hparam(["num_local_experts", "num_experts"])
-            assert bid is not None
-
-            if self._experts is None:
-                self._experts = [{} for _ in range(self.block_count)]
-
-            self._experts[bid][name] = data_torch
-
-            if len(self._experts[bid]) >= n_experts * 3:
-                # merge the experts into a single 3d tensor
-                for w_name in ["down_proj", "gate_proj", "up_proj"]:
-                    datas: list[Tensor] = []
-
-                    for xid in range(n_experts):
-                        ename = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight"
-                        datas.append(self._experts[bid][ename])
-                        del self._experts[bid][ename]
-
-                    data_torch = torch.stack(datas, dim=0)
-
-                    merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight"
-
-                    yield from super().modify_tensors(data_torch, merged_name, bid)
-                return
-            else:
-                return
-
-        yield from super().modify_tensors(data_torch, name, bid)
-
-    # Copied from: Qwen2MoeModel
-    def prepare_tensors(self):
-        super().prepare_tensors()
-
-        if self._experts is not None:
-            # flatten `list[dict[str, Tensor]]` into `list[str]`
-            experts = [k for d in self._experts for k in d.keys()]
-            if len(experts) > 0:
-                raise ValueError(f"Unprocessed experts: {experts}")
+    moe_experts = [MOE_HF_MLP]

conversion/phi.py

@@ -10,7 +10,7 @@ import torch
 if TYPE_CHECKING:
     from torch import Tensor
 
-from .base import MmprojModel, ModelBase, SentencePieceTokenTypes, TextModel, gguf, logger
+from .base import MOE_BLOCK_SPARSE, MmprojModel, ModelBase, SentencePieceTokenTypes, TextModel, gguf, logger
 
 
 @ModelBase.register("PhiForCausalLM")
@@ -339,50 +339,9 @@ class Phi4VisionMmprojModel(MmprojModel):
 class PhiMoeModel(Phi3MiniModel):
     model_arch = gguf.MODEL_ARCH.PHIMOE
 
-    _experts: list[dict[str, Tensor]] | None = None
+    moe_experts = [MOE_BLOCK_SPARSE]
 
     def set_gguf_parameters(self):
         super().set_gguf_parameters()
         self.gguf_writer.add_expert_used_count(self.find_hparam(["num_experts_per_tok", "num_experts_per_token"]))
         self.gguf_writer.add_expert_count(self.find_hparam(["num_local_experts", "num_experts"]))
-
-    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
-        # process the experts separately
-        if name.find("block_sparse_moe.experts") != -1:
-            n_experts = self.find_hparam(["num_local_experts", "num_experts"])
-            assert bid is not None
-
-            if self._experts is None:
-                self._experts = [{} for _ in range(self.block_count)]
-
-            self._experts[bid][name] = data_torch
-
-            if len(self._experts[bid]) >= n_experts * 3:
-                # merge the experts into a single 3d tensor
-                for w_name in ["w1", "w2", "w3"]:
-                    datas: list[Tensor] = []
-
-                    for xid in range(n_experts):
-                        ename = f"model.layers.{bid}.block_sparse_moe.experts.{xid}.{w_name}.weight"
-                        datas.append(self._experts[bid][ename])
-                        del self._experts[bid][ename]
-
-                    data_torch = torch.stack(datas, dim=0)
-
-                    merged_name = f"model.layers.{bid}.block_sparse_moe.experts.{w_name}.weight"
-
-                    yield from super().modify_tensors(data_torch, merged_name, bid)
-                return
-            else:
-                return
-
-        yield from super().modify_tensors(data_torch, name, bid)
-
-    def prepare_tensors(self):
-        super().prepare_tensors()
-
-        if self._experts is not None:
-            # flatten `list[dict[str, Tensor]]` into `list[str]`
-            experts = [k for d in self._experts for k in d.keys()]
-            if len(experts) > 0:
-                raise ValueError(f"Unprocessed experts: {experts}")

conversion/qwen.py

@@ -9,7 +9,7 @@ import torch
 if TYPE_CHECKING:
     from torch import Tensor
 
-from .base import ModelBase, TextModel, gguf, logger
+from .base import MOE_HF_MLP, ModelBase, TextModel, gguf, logger
 
 
 @ModelBase.register("QWenLMHeadModel")
@@ -83,7 +83,7 @@ class Qwen2MoeModel(TextModel):
             self.gguf_writer.add_expert_shared_feed_forward_length(shared_expert_intermediate_size)
             logger.info(f"gguf: expert shared feed forward length = {shared_expert_intermediate_size}")
 
-    _experts: list[dict[str, Tensor]] | None = None
+    moe_experts = [MOE_HF_MLP]
 
     def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
         # handle aggregated expert tensors
@@ -112,45 +112,8 @@ class Qwen2MoeModel(TextModel):
             yield from super().modify_tensors(up, mapped_up, bid)
             return
 
-        if name.find("experts") != -1:
-            n_experts = self.find_hparam(["num_local_experts", "num_experts"])
-            assert bid is not None
-
-            if self._experts is None:
-                self._experts = [{} for _ in range(self.block_count)]
-
-            self._experts[bid][name] = data_torch
-
-            if len(self._experts[bid]) >= n_experts * 3:
-                # merge the experts into a single 3d tensor
-                for w_name in ["down_proj", "gate_proj", "up_proj"]:
-                    datas: list[Tensor] = []
-
-                    for xid in range(n_experts):
-                        ename = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight"
-                        datas.append(self._experts[bid][ename])
-                        del self._experts[bid][ename]
-
-                    data_torch = torch.stack(datas, dim=0)
-
-                    merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight"
-
-                    yield from super().modify_tensors(data_torch, merged_name, bid)
-                return
-            else:
-                return
-
         yield from super().modify_tensors(data_torch, name, bid)
 
-    def prepare_tensors(self):
-        super().prepare_tensors()
-
-        if self._experts is not None:
-            # flatten `list[dict[str, Tensor]]` into `list[str]`
-            experts = [k for d in self._experts for k in d.keys()]
-            if len(experts) > 0:
-                raise ValueError(f"Unprocessed experts: {experts}")
-
 
 @ModelBase.register("Qwen3ForCausalLM", "Qwen3Model")
 class Qwen3Model(Qwen2Model):

conversion/smallthinker.py

@@ -1,13 +1,6 @@
 from __future__ import annotations
 
-from typing import Iterable, TYPE_CHECKING
-
-import torch
-
-if TYPE_CHECKING:
-    from torch import Tensor
-
-from .base import ModelBase, TextModel, gguf, logger
+from .base import MOE_BLOCK_SPARSE, ModelBase, TextModel, gguf, logger
 
 
 @ModelBase.register("SmallThinkerForCausalLM")
@@ -38,45 +31,7 @@ class SmallThinkerModel(TextModel):
                         self.gguf_writer.add_sliding_window(sliding_window)
                     break
 
-    _experts: list[dict[str, Tensor]] | None = None
-
-    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
-        # process the experts separately
-        if name.find("experts") != -1:
-            n_experts = self.hparams.get("moe_num_primary_experts") or self.find_hparam(["num_local_experts", "num_experts"])
-            assert bid is not None
-
-            if self._experts is None:
-                self._experts = [{} for _ in range(self.block_count)]
-
-            self._experts[bid][name] = data_torch
-
-            if len(self._experts[bid]) >= n_experts * 3:
-                # merge the experts into a single 3d tensor
-                for w_name in ["down", "gate", "up"]:
-                    datas: list[Tensor] = []
-
-                    for xid in range(n_experts):
-                        ename = f"model.layers.{bid}.block_sparse_moe.experts.{xid}.{w_name}.weight"
-                        datas.append(self._experts[bid][ename])
-                        del self._experts[bid][ename]
-
-                    data_torch = torch.stack(datas, dim=0)
-
-                    merged_name = f"model.layers.{bid}.block_sparse_moe.experts.{w_name}.weight"
-
-                    yield from super().modify_tensors(data_torch, merged_name, bid)
-                return
-            else:
-                return
-
-        yield from super().modify_tensors(data_torch, name, bid)
-
-    def prepare_tensors(self):
-        super().prepare_tensors()
-
-        if self._experts is not None:
-            # flatten `list[dict[str, Tensor]]` into `list[str]`
-            experts = [k for d in self._experts for k in d.keys()]
-            if len(experts) > 0:
-                raise ValueError(f"Unprocessed experts: {experts}")
+    moe_experts = [MOE_BLOCK_SPARSE._replace(
+        weights=("down", "gate", "up"),
+        n_expert=("moe_num_primary_experts", "num_local_experts", "num_experts"),
+    )]