mirror of
https://github.com/ggml-org/llama.cpp.git
synced 2026-09-24 13:37:01 +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 <[email protected]> Co-authored-by: Piotr Wilkin <[email protected]>
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
|
||||
if (src.find("<tool_call>") != 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");
|
||||
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)
|
||||
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.
|
||||
|
||||
+2
-33
@@ -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"]
|
||||
|
||||
|
||||
+86
-2
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user