mirror of
https://github.com/ggml-org/llama.cpp.git
synced 2026-10-08 22:10:37 +02:00
convert: add MiMo-V2.6 support (#29257)
* 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 <sigbjorn.skjaeret@huggingface.co> Co-authored-by: Piotr Wilkin <piotr.wilkin@syndatis.com>
This commit is contained in:
co-authored by
Sigbjørn Skjæret
Piotr Wilkin
parent
a60f9aead0
commit
bfd73a876e
+3
-1
@@ -1212,7 +1212,9 @@ std::optional<common_chat_params> common_chat_try_specialized_template(
|
|||||||
// Qwen3-Coder XML tool calls, also used by Nemotron Nano 3, Qwen3.5 and StepFun-3.5-Flash
|
// Qwen3-Coder XML tool calls, also used by Nemotron Nano 3, Qwen3.5 and StepFun-3.5-Flash
|
||||||
if (src.find("<tool_call>") != std::string::npos &&
|
if (src.find("<tool_call>") != std::string::npos &&
|
||||||
src.find("<function=") != std::string::npos &&
|
src.find("<function=") != std::string::npos &&
|
||||||
src.find("<parameter=") != std::string::npos) {
|
src.find("<parameter=") != std::string::npos &&
|
||||||
|
// Exclude models that don't use \n between tags
|
||||||
|
src.find("'<tool_call><function=' ~ tool_call.name ~ '>'") == std::string::npos) {
|
||||||
LOG_DBG("Using specialized template: Qwen3-Coder\n");
|
LOG_DBG("Using specialized template: Qwen3-Coder\n");
|
||||||
return common_chat_params_init_qwen3_coder(tmpl, params);
|
return common_chat_params_init_qwen3_coder(tmpl, params);
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -776,6 +776,36 @@ class ModelBase:
|
|||||||
raw = torch.cat((s.unsqueeze(-1), qs.to(torch.uint8)), dim=-1)
|
raw = torch.cat((s.unsqueeze(-1), qs.to(torch.uint8)), dim=-1)
|
||||||
return raw.reshape(rows, n_blocks * 17).cpu().numpy()
|
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
|
@staticmethod
|
||||||
def _nvfp4_pack(weight: Tensor, scale: Tensor) -> tuple[np.ndarray, list[int]]:
|
def _nvfp4_pack(weight: Tensor, scale: Tensor) -> tuple[np.ndarray, list[int]]:
|
||||||
"""Repack NVFP4 ModelOpt tensors into ggml super-block layout.
|
"""Repack NVFP4 ModelOpt tensors into ggml super-block layout.
|
||||||
|
|||||||
+2
-33
@@ -2,15 +2,14 @@ from __future__ import annotations
|
|||||||
|
|
||||||
import re
|
import re
|
||||||
from pathlib import Path
|
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
|
import torch
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from torch import Tensor
|
from torch import Tensor
|
||||||
|
|
||||||
from .base import LazyTorchTensor, ModelBase, TextModel, gguf, logger
|
from .base import ModelBase, TextModel, gguf, logger
|
||||||
|
|
||||||
from .kimi_linear import KimiLinearModel
|
from .kimi_linear import KimiLinearModel
|
||||||
|
|
||||||
@@ -104,36 +103,6 @@ class KimiK3Model(TextModel):
|
|||||||
"only the routed experts have a repack path"
|
"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:
|
def _write_mxfp4_experts(self) -> None:
|
||||||
n_experts = self.hparams["num_experts"]
|
n_experts = self.hparams["num_experts"]
|
||||||
|
|
||||||
|
|||||||
+86
-2
@@ -10,7 +10,7 @@ import torch
|
|||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from torch import Tensor
|
from torch import Tensor
|
||||||
|
|
||||||
from .base import MmprojModel, ModelBase, TextModel, gguf
|
from .base import MmprojModel, ModelBase, TextModel, gguf, logger
|
||||||
|
|
||||||
|
|
||||||
@ModelBase.register("MiMoV2FlashForCausalLM", "MiMoV2ForCausalLM")
|
@ModelBase.register("MiMoV2FlashForCausalLM", "MiMoV2ForCausalLM")
|
||||||
@@ -167,6 +167,84 @@ class MimoV2Model(TextModel):
|
|||||||
|
|
||||||
self.gguf_writer.add_nextn_predict_layers(self._n_nextn)
|
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
|
_experts: list[dict[str, Tensor]] | None = None
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
@@ -192,7 +270,7 @@ class MimoV2Model(TextModel):
|
|||||||
bid = new_bid
|
bid = new_bid
|
||||||
|
|
||||||
# process the experts separately
|
# 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"]
|
n_experts = self.hparams["n_routed_experts"]
|
||||||
assert bid is not None
|
assert bid is not None
|
||||||
|
|
||||||
@@ -229,6 +307,10 @@ class MimoV2Model(TextModel):
|
|||||||
if len(experts) > 0:
|
if len(experts) > 0:
|
||||||
raise ValueError(f"Unprocessed experts: {experts}")
|
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.register("MiMoV2ForCausalLM")
|
||||||
@ModelBase.example("XiaomiMiMo/MiMo-V2.5")
|
@ModelBase.example("XiaomiMiMo/MiMo-V2.5")
|
||||||
@@ -382,6 +464,8 @@ class MiMoV2VisionAudioModel(MmprojModel):
|
|||||||
"_codebook.inited",
|
"_codebook.inited",
|
||||||
)
|
)
|
||||||
for name, tensor in state_dict.items():
|
for name, tensor in state_dict.items():
|
||||||
|
if name.startswith("decoder."):
|
||||||
|
continue
|
||||||
if name.endswith(skip_suffixes):
|
if name.endswith(skip_suffixes):
|
||||||
continue
|
continue
|
||||||
if m := codebook_re.match(name):
|
if m := codebook_re.match(name):
|
||||||
|
|||||||
Reference in New Issue
Block a user