From bfd73a876ed25c5a760d6b60b8fe12a65115fedf Mon Sep 17 00:00:00 2001 From: AesSedai <7980540+AesSedai@users.noreply.github.com> Date: Tue, 22 Sep 2026 05:38:09 -0700 Subject: [PATCH] convert: add MiMo-V2.6 support (#29257) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * convert: add MiMo-V2.6 support Hoist the K3 mxfp4 conversion repack into base.py so it can be reused Remove decoder from mmproj convert * Update conversion/mimo.py * fix: use autoparser --------- Co-authored-by: Sigbjørn Skjæret Co-authored-by: Piotr Wilkin --- common/chat.cpp | 4 +- conversion/base.py | 30 +++++++++++++++ conversion/kimi_k3.py | 35 +---------------- conversion/mimo.py | 88 ++++++++++++++++++++++++++++++++++++++++++- 4 files changed, 121 insertions(+), 36 deletions(-) diff --git a/common/chat.cpp b/common/chat.cpp index 6c8099cf29..ed1942e153 100644 --- a/common/chat.cpp +++ b/common/chat.cpp @@ -1212,7 +1212,9 @@ std::optional common_chat_try_specialized_template( // Qwen3-Coder XML tool calls, also used by Nemotron Nano 3, Qwen3.5 and StepFun-3.5-Flash if (src.find("") != std::string::npos && src.find("'") == std::string::npos) { LOG_DBG("Using specialized template: Qwen3-Coder\n"); return common_chat_params_init_qwen3_coder(tmpl, params); } diff --git a/conversion/base.py b/conversion/base.py index 6aca7f1d34..9fba5a86b0 100644 --- a/conversion/base.py +++ b/conversion/base.py @@ -776,6 +776,36 @@ class ModelBase: raw = torch.cat((s.unsqueeze(-1), qs.to(torch.uint8)), dim=-1) return raw.reshape(rows, n_blocks * 17).cpu().numpy() + def _mxfp4_expert_tensor(self, loaders: list[tuple[Callable[[], Tensor], Callable[[], Tensor]]]): + """ + One stacked [n_expert, rows, cols] MXFP4 tensor, built lazily. + + gguf_writer holds every added tensor until the final write, so building + this eagerly (like the DeepSeek-V4 path does) keeps every expert in + memory at once. lazy means only the tensor being written is resident. + """ + # meta shapes, so this does not read any weights + rows, packed_cols = loaders[0][0]().shape + n_blocks = (packed_cols * 2) // 32 + byte_shape = (len(loaders), rows, n_blocks * 17) + + def load(fns: list[tuple[Callable[[], Tensor], Callable[[], Tensor]]]) -> np.ndarray: + out = np.empty(byte_shape, dtype=np.uint8) + for eid, (packed_fn, scale_fn) in enumerate(fns): + out[eid] = self.repack_mxfp4_blocks( + LazyTorchTensor.to_eager(packed_fn()), + LazyTorchTensor.to_eager(scale_fn()), + ) + return out + + # loaders goes through args, not the closure, so that `func` matches + # LazyBase's single-argument shape + return gguf.LazyNumpyTensor( + meta=gguf.LazyNumpyTensor.meta_with_dtype_and_shape(np.uint8, byte_shape), + args=(loaders,), + func=load, + ) + @staticmethod def _nvfp4_pack(weight: Tensor, scale: Tensor) -> tuple[np.ndarray, list[int]]: """Repack NVFP4 ModelOpt tensors into ggml super-block layout. diff --git a/conversion/kimi_k3.py b/conversion/kimi_k3.py index d15d1d64bf..70aabb7051 100644 --- a/conversion/kimi_k3.py +++ b/conversion/kimi_k3.py @@ -2,15 +2,14 @@ from __future__ import annotations import re from pathlib import Path -from typing import Callable, Iterable, Iterator, TYPE_CHECKING +from typing import Iterable, Iterator, TYPE_CHECKING -import numpy as np import torch if TYPE_CHECKING: from torch import Tensor -from .base import LazyTorchTensor, ModelBase, TextModel, gguf, logger +from .base import ModelBase, TextModel, gguf, logger from .kimi_linear import KimiLinearModel @@ -104,36 +103,6 @@ class KimiK3Model(TextModel): "only the routed experts have a repack path" ) - def _mxfp4_expert_tensor(self, loaders: list[tuple[Callable[[], Tensor], Callable[[], Tensor]]]): - """ - One stacked [n_expert, rows, cols] MXFP4 tensor, built lazily. - - gguf_writer holds every added tensor until the final write, so building - this eagerly (like the DeepSeek-V4 path does) keeps all ~1.38 TB of - experts in memory. lazy means only the tensor being written is resident. - """ - # meta shapes, so this does not read any weights - rows, packed_cols = loaders[0][0]().shape - n_blocks = (packed_cols * 2) // 32 - byte_shape = (len(loaders), rows, n_blocks * 17) - - def load(fns: list[tuple[Callable[[], Tensor], Callable[[], Tensor]]]) -> np.ndarray: - out = np.empty(byte_shape, dtype=np.uint8) - for eid, (packed_fn, scale_fn) in enumerate(fns): - out[eid] = self.repack_mxfp4_blocks( - LazyTorchTensor.to_eager(packed_fn()), - LazyTorchTensor.to_eager(scale_fn()), - ) - return out - - # loaders goes through args, not the closure, so that `func` matches - # LazyBase's single-argument shape - return gguf.LazyNumpyTensor( - meta=gguf.LazyNumpyTensor.meta_with_dtype_and_shape(np.uint8, byte_shape), - args=(loaders,), - func=load, - ) - def _write_mxfp4_experts(self) -> None: n_experts = self.hparams["num_experts"] diff --git a/conversion/mimo.py b/conversion/mimo.py index 15dbeb7e75..8a2689b965 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 MmprojModel, ModelBase, TextModel, gguf, logger @ModelBase.register("MiMoV2FlashForCausalLM", "MiMoV2ForCausalLM") @@ -167,6 +167,84 @@ class MimoV2Model(TextModel): self.gguf_writer.add_nextn_predict_layers(self._n_nextn) + _MXFP4_EXPERT_RE = re.compile( + r"^model\.layers\.(\d+)\.mlp\.experts\.(\d+)\.(gate|up|down)_proj\.weight$" + ) + _MXFP4_PROJ = { + "gate": gguf.MODEL_TENSOR.FFN_GATE_EXP, + "up": gguf.MODEL_TENSOR.FFN_UP_EXP, + "down": gguf.MODEL_TENSOR.FFN_DOWN_EXP, + } + + def _is_mxfp4_packed(self) -> bool: + quant_config = self.hparams.get("quantization_config") or {} + if quant_config.get("store_dtype") != "mxfp4": + return False + # repack_mxfp4_blocks assumes ggml's 32-element group + block_size = quant_config.get("mxfp4_block_size", 32) + if block_size != 32: + raise NotImplementedError( + f"MXFP4 block size {block_size} is not ggml's QK_MXFP4 (32)") + return True + + def _write_mxfp4_experts(self) -> None: + n_experts = self.hparams["n_routed_experts"] + + # the FP8 half uses `weight_scale_inv` and is left to dequant_model + stray = [n for n in self.model_tensors + if n.endswith(".weight_scale") and not self._MXFP4_EXPERT_RE.match(n.removesuffix("_scale"))] + if stray: + raise NotImplementedError( + f"{len(stray)} MXFP4 tensor(s) outside the routed experts, e.g. {stray[0]!r}; " + "only the routed experts have a repack path" + ) + + # (bid, proj) -> {expert id: (weight name, scale name)} + groups: dict[tuple[int, str], dict[int, tuple[str, str]]] = {} + for name in self.model_tensors: + m = self._MXFP4_EXPERT_RE.match(name) + if m is None: + continue + bid, eid, proj = int(m.group(1)), int(m.group(2)), m.group(3) + scale_name = name + "_scale" + if scale_name not in self.model_tensors: + raise KeyError(f"missing {scale_name} for {name}") + groups.setdefault((bid, proj), {})[eid] = (name, scale_name) + + consumed: list[str] = [] + for (bid, proj), experts in sorted(groups.items()): + missing = [e for e in range(n_experts) if e not in experts] + if missing or len(experts) != n_experts: + raise KeyError( + f"layer {bid} {proj}_proj: {len(experts)} of {n_experts} experts present" + + (f", first missing is {missing[0]}" if missing else "") + ) + + loaders = [] + for eid in range(n_experts): + weight_name, scale_name = experts[eid] + loaders.append((self.model_tensors[weight_name], self.model_tensors[scale_name])) + consumed += [weight_name, scale_name] + + data = self._mxfp4_expert_tensor(loaders) + new_name = self.format_tensor_name(self._MXFP4_PROJ[proj], bid) + shape = gguf.quant_shape_from_byte_shape(data.shape, gguf.GGMLQuantizationType.MXFP4) + logger.info( + f"{new_name}: repacked {n_experts} experts to MXFP4, " + f"shape = {{{', '.join(str(n) for n in reversed(shape))}}}" + ) + self.gguf_writer.add_tensor(new_name, data, raw_dtype=gguf.GGMLQuantizationType.MXFP4) + + for name in consumed: + del self.model_tensors[name] + + def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]: + # not a generator on purpose: base.py chains this with get_tensors(), so the + # tensors used here must be removed from model_tensors before that starts + if self._is_mxfp4_packed(): + self._write_mxfp4_experts() + return () + _experts: list[dict[str, Tensor]] | None = None @classmethod @@ -192,7 +270,7 @@ class MimoV2Model(TextModel): bid = new_bid # process the experts separately - if name.find("mlp.experts") != -1: + if ".mlp.experts." in name and name.endswith(".weight"): n_experts = self.hparams["n_routed_experts"] assert bid is not None @@ -229,6 +307,10 @@ class MimoV2Model(TextModel): if len(experts) > 0: raise ValueError(f"Unprocessed experts: {experts}") + if self._is_mxfp4_packed(): + self._is_mxfp4 = True + self.ftype = gguf.LlamaFileType.MOSTLY_MXFP4_MOE + @ModelBase.register("MiMoV2ForCausalLM") @ModelBase.example("XiaomiMiMo/MiMo-V2.5") @@ -382,6 +464,8 @@ class MiMoV2VisionAudioModel(MmprojModel): "_codebook.inited", ) for name, tensor in state_dict.items(): + if name.startswith("decoder."): + continue if name.endswith(skip_suffixes): continue if m := codebook_re.match(name):