diff --git a/conversion/afmoe.py b/conversion/afmoe.py index 5e66a51da6..122aab03f5 100644 --- a/conversion/afmoe.py +++ b/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) diff --git a/conversion/arctic.py b/conversion/arctic.py index 775cacaab9..1f30ea9bcc 100644 --- a/conversion/arctic.py +++ b/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}") diff --git a/conversion/bailingmoe.py b/conversion/bailingmoe.py index 2c6425cb64..7d267736d8 100644 --- a/conversion/bailingmoe.py +++ b/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): diff --git a/conversion/base.py b/conversion/base.py index 4f4fbca347..d963267d47 100644 --- a/conversion/base.py +++ b/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\d+)", + "xid": r"(?P\d+)", + "w": "(?P" + "|".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) diff --git a/conversion/command_r.py b/conversion/command_r.py index 118565c669..1c7c204244 100644 --- a/conversion/command_r.py +++ b/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}") diff --git a/conversion/deepseek.py b/conversion/deepseek.py index dbfa96c2ee..ce38831f54 100644 --- a/conversion/deepseek.py +++ b/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): diff --git a/conversion/ernie.py b/conversion/ernie.py index aa8a3bc8ee..ee45ddb6b6 100644 --- a/conversion/ernie.py +++ b/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") diff --git a/conversion/exaone.py b/conversion/exaone.py index 1cd2244dbc..9a51c54c02 100644 --- a/conversion/exaone.py +++ b/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): diff --git a/conversion/glm.py b/conversion/glm.py index e28f54574e..528651e86c 100644 --- a/conversion/glm.py +++ b/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): diff --git a/conversion/grovemoe.py b/conversion/grovemoe.py index a8be931cb9..9f39261cc7 100644 --- a/conversion/grovemoe.py +++ b/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}") diff --git a/conversion/hunyuan.py b/conversion/hunyuan.py index f5ac8a4fb7..be507b7e1b 100644 --- a/conversion/hunyuan.py +++ b/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] diff --git a/conversion/kimi_k3.py b/conversion/kimi_k3.py index a5005e15bd..e401074af3 100644 --- a/conversion/kimi_k3.py +++ b/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",))] # `_res_norm.weight` and `_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"] diff --git a/conversion/kimi_linear.py b/conversion/kimi_linear.py index f2e6cda83c..118182a770 100644 --- a/conversion/kimi_linear.py +++ b/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") diff --git a/conversion/laguna.py b/conversion/laguna.py index a90f355ca9..72d54d05b6 100644 --- a/conversion/laguna.py +++ b/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"): diff --git a/conversion/lfm2.py b/conversion/lfm2.py index 70ce45658b..1085634c0a 100644 --- a/conversion/lfm2.py +++ b/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): diff --git a/conversion/llada.py b/conversion/llada.py index 98dc9de95b..2ebea0cd7c 100644 --- a/conversion/llada.py +++ b/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] diff --git a/conversion/llama.py b/conversion/llama.py index 1aced49c54..be535ca830 100644 --- a/conversion/llama.py +++ b/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): diff --git a/conversion/mellum.py b/conversion/mellum.py index 79bc6755cc..3a8f75d7a1 100644 --- a/conversion/mellum.py +++ b/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] diff --git a/conversion/mimo.py b/conversion/mimo.py index ca2ed28ad3..5c43b5ff76 100644 --- a/conversion/mimo.py +++ b/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): diff --git a/conversion/minimax.py b/conversion/minimax.py index d7a00bac95..ba7f7ab1d2 100644 --- a/conversion/minimax.py +++ b/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): diff --git a/conversion/nemotron.py b/conversion/nemotron.py index c46cec1438..e894d53cb9 100644 --- a/conversion/nemotron.py +++ b/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", diff --git a/conversion/olmo.py b/conversion/olmo.py index 1664c30e40..93dfe9d3e8 100644 --- a/conversion/olmo.py +++ b/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] diff --git a/conversion/phi.py b/conversion/phi.py index df4bfe809a..39588634f4 100644 --- a/conversion/phi.py +++ b/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}") diff --git a/conversion/qwen.py b/conversion/qwen.py index ead435455d..bed7e4c8fa 100644 --- a/conversion/qwen.py +++ b/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): diff --git a/conversion/smallthinker.py b/conversion/smallthinker.py index 1b0f79aa3e..fcea9ac16a 100644 --- a/conversion/smallthinker.py +++ b/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"), + )]