diff --git a/conversion/minimax.py b/conversion/minimax.py index 4857775cb..6e0b78303 100644 --- a/conversion/minimax.py +++ b/conversion/minimax.py @@ -7,25 +7,149 @@ import torch if TYPE_CHECKING: from torch import Tensor -from .base import ModelBase, TextModel, gguf +from .base import LazyTorchTensor, ModelBase, TextModel, gguf @ModelBase.register("MiniMaxM2ForCausalLM") class MiniMaxM2Model(TextModel): model_arch = gguf.MODEL_ARCH.MINIMAXM2 - _experts_cache: dict[int, dict[str, Tensor]] = {} + _experts_cache: dict[int, dict[str, Tensor]] + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._experts_cache = {} + self._adaptive_retained_by_layer = self._load_adaptive_retained_by_layer() + self._uniform_expert_count = self.find_hparam(["num_local_experts", "num_experts"]) + + def _load_adaptive_retained_by_layer(self) -> dict[int, int] | None: + raw_counts = self.hparams.get("reap_adaptive_retained_experts_by_layer") + if not raw_counts: + if self.hparams.get("reap_adaptive_nonuniform_experts"): + raise ValueError( + "MiniMax-M2 adaptive checkpoint declares reap_adaptive_nonuniform_experts " + "but lacks reap_adaptive_retained_experts_by_layer" + ) + return None + + if not isinstance(raw_counts, dict): + raise ValueError( + "reap_adaptive_retained_experts_by_layer must be a layer->count object, " + f"got {type(raw_counts).__name__}" + ) + + max_experts = self.find_hparam(["num_local_experts", "num_experts"]) + counts: dict[int, int] = {} + for raw_layer, raw_count in raw_counts.items(): + layer = int(raw_layer) + count = int(raw_count) + if layer < 0 or layer >= self.block_count: + raise ValueError( + f"adaptive retained expert metadata has out-of-range layer {layer}; " + f"block_count={self.block_count}" + ) + if count < 1 or count > max_experts: + raise ValueError( + f"adaptive retained expert count for layer {layer} is {count}; " + f"expected 1..{max_experts}" + ) + counts[layer] = count + + missing = sorted(set(range(self.block_count)) - set(counts)) + if missing: + raise ValueError( + "MiniMax-M2 adaptive GGUF export requires counts for every layer; " + f"missing layers={missing[:8]}{'...' if len(missing) > 8 else ''}" + ) + if max(counts.values()) != max_experts: + raise ValueError( + "MiniMax-M2 adaptive config num_local_experts must equal the max retained " + f"layer count for padded GGUF export; got config={max_experts} " + f"metadata_max={max(counts.values())}" + ) + return counts + + def _layer_expert_count(self, bid: int) -> int: + if self._adaptive_retained_by_layer is None: + return self._uniform_expert_count + return self._adaptive_retained_by_layer[bid] + + def _pad_expert_axis(self, data_torch: Tensor, count: int, *, dim: int, fill_value: float = 0.0) -> Tensor: + if count == self._uniform_expert_count: + return data_torch + if data_torch.shape[dim] != count: + raise ValueError( + "MiniMax-M2 adaptive tensor shape does not match metadata: " + f"shape={tuple(data_torch.shape)} expert_dim={dim} " + f"shape_count={data_torch.shape[dim]} metadata_count={count}" + ) + + pad_shape = list(data_torch.shape) + pad_shape[dim] = self._uniform_expert_count - count + if isinstance(data_torch, LazyTorchTensor): + pad = LazyTorchTensor( + meta=LazyTorchTensor.meta_with_dtype_and_shape(data_torch.dtype, tuple(pad_shape)), + args=(data_torch,), + func=lambda source: torch.full( + tuple(pad_shape), + fill_value, + dtype=source.dtype, + device=source.device, + ), + ) + else: + pad = torch.full( + tuple(pad_shape), + fill_value, + dtype=data_torch.dtype, + device=data_torch.device, + ) + return torch.cat([data_torch, pad], dim=dim).contiguous() + + def _add_adaptive_padding_metadata(self): + if self._adaptive_retained_by_layer is None: + return + counts = [self._adaptive_retained_by_layer[layer] for layer in range(self.block_count)] + self.gguf_writer.add_bool("minimax-m2.reap_adaptive_padded_uniform", True) + self.gguf_writer.add_bool("minimax-m2.reap_adaptive_dummy_experts_masked", True) + self.gguf_writer.add_array("minimax-m2.reap_adaptive_retained_experts_by_layer", counts) def set_gguf_parameters(self): super().set_gguf_parameters() 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"])) + self._add_adaptive_padding_metadata() + + def prepare_tensors(self): + super().prepare_tensors() + if self._experts_cache: + pending = {bid: len(cache) for bid, cache in sorted(self._experts_cache.items())} + raise RuntimeError( + "MiniMax-M2 expert export ended with incomplete per-layer expert caches; " + f"pending_tensors_by_layer={pending}" + ) def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None): + if bid is not None and self._adaptive_retained_by_layer is not None: + count = self._layer_expert_count(bid) + if name.endswith(".block_sparse_moe.gate.weight"): + data_torch = self._pad_expert_axis(data_torch, count, dim=0) + elif name.endswith((".block_sparse_moe.e_score_correction.bias", ".block_sparse_moe.e_score_correction_bias")): + # Large finite negative instead of -inf: selection scores are + # sigmoid(logits) + bias with sigmoid in [0, 1], so -1e9 makes a + # padded expert unselectable exactly like -inf would, while + # ggml_validate_row_data (llama-quantize) rejects inf values. + data_torch = self._pad_expert_axis( + data_torch, + count, + dim=0, + fill_value=-1.0e9, + ) + # merge expert weights if 'experts' in name: - n_experts = self.find_hparam(["num_local_experts", "num_experts"]) assert bid is not None + n_experts = self._layer_expert_count(bid) expert_cache = self._experts_cache.setdefault(bid, {}) expert_cache[name] = data_torch @@ -44,6 +168,7 @@ class MiniMaxM2Model(TextModel): del expert_cache[ename] data_torch = torch.stack(datas, dim=0) + data_torch = self._pad_expert_axis(data_torch, n_experts, 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) diff --git a/tests/test_minimax_adaptive_export.py b/tests/test_minimax_adaptive_export.py new file mode 100644 index 000000000..0a694bafc --- /dev/null +++ b/tests/test_minimax_adaptive_export.py @@ -0,0 +1,134 @@ +from __future__ import annotations + +import pytest +import torch + +from conversion.base import LazyTorchTensor, TextModel, gguf +from conversion.minimax import MiniMaxM2Model + + +class DummyWriter: + def __init__(self) -> None: + self.values: dict[str, object] = {} + + def add_bool(self, key: str, value: bool) -> None: + self.values[key] = value + + def add_array(self, key: str, value: list[int]) -> None: + self.values[key] = value + + +def make_minimax_model(counts: dict[int, int], *, max_experts: int = 4) -> MiniMaxM2Model: + model = object.__new__(MiniMaxM2Model) + model.block_count = len(counts) + model.hparams = { + "hidden_size": 3, + "intermediate_size": 2, + "num_hidden_layers": len(counts), + "num_local_experts": max_experts, + "reap_adaptive_nonuniform_experts": True, + "reap_adaptive_retained_experts_by_layer": { + str(layer): count for layer, count in counts.items() + }, + } + model.tensor_map = gguf.get_tensor_name_map(gguf.MODEL_ARCH.MINIMAXM2, model.block_count) + model._experts_cache = {} + model.fuse_gate_up_exps = False + model._adaptive_retained_by_layer = model._load_adaptive_retained_by_layer() + model._uniform_expert_count = model.find_hparam(["num_local_experts", "num_experts"]) + return model + + +def test_minimax_adaptive_export_pads_experts_router_and_bias() -> None: + model = make_minimax_model({0: 2, 1: 4}) + + outputs: list[tuple[str, torch.Tensor]] = [] + for expert_id in range(2): + for weight_name in ("w1", "w2", "w3"): + name = f"model.layers.0.block_sparse_moe.experts.{expert_id}.{weight_name}.weight" + tensor = torch.full((2, 3), expert_id + 1.0) + outputs.extend(model.modify_tensors(tensor, name, 0)) + + assert [name for name, _ in outputs] == [ + "blk.0.ffn_gate_exps.weight", + "blk.0.ffn_down_exps.weight", + "blk.0.ffn_up_exps.weight", + ] + for _, tensor in outputs: + assert tensor.shape == (4, 2, 3) + assert torch.equal(tensor[2:], torch.zeros_like(tensor[2:])) + + gate_outputs = list( + model.modify_tensors( + torch.ones((2, 3)), + "model.layers.0.block_sparse_moe.gate.weight", + 0, + ) + ) + assert len(gate_outputs) == 1 + assert gate_outputs[0][0] == "blk.0.ffn_gate_inp.weight" + assert gate_outputs[0][1].shape == (4, 3) + assert torch.equal(gate_outputs[0][1][2:], torch.zeros_like(gate_outputs[0][1][2:])) + + bias_outputs = list( + model.modify_tensors( + torch.tensor([0.25, 0.5]), + "model.layers.0.block_sparse_moe.e_score_correction.bias", + 0, + ) + ) + assert len(bias_outputs) == 1 + assert bias_outputs[0][0] == "blk.0.exp_probs_b.bias" + assert bias_outputs[0][1].shape == (4,) + # Finite large-negative mask (not -inf): ggml_validate_row_data rejects inf, + # and sigmoid-bounded selection scores make -1e9 equally unselectable. + assert torch.isfinite(bias_outputs[0][1][2:]).all() + assert (bias_outputs[0][1][2:] <= -1.0e8).all() + + +def test_minimax_adaptive_export_records_padding_metadata() -> None: + model = make_minimax_model({0: 2, 1: 4}) + model.gguf_writer = DummyWriter() + + model._add_adaptive_padding_metadata() + + assert model.gguf_writer.values == { + "minimax-m2.reap_adaptive_padded_uniform": True, + "minimax-m2.reap_adaptive_dummy_experts_masked": True, + "minimax-m2.reap_adaptive_retained_experts_by_layer": [2, 4], + } + + +def test_minimax_adaptive_export_materializes_lazy_padding_on_source_device() -> None: + model = make_minimax_model({0: 2, 1: 4}) + source = LazyTorchTensor.from_eager(torch.ones((2, 3))) + + padded = model._pad_expert_axis(source, 2, dim=0, fill_value=-1.0) + eager = LazyTorchTensor.to_eager(padded) + + assert eager.device.type == "cpu" + assert eager.shape == (4, 3) + assert torch.equal(eager[:2], torch.ones((2, 3))) + assert torch.equal(eager[2:], torch.full((2, 3), -1.0)) + + +def test_minimax_adaptive_export_fails_when_metadata_is_incomplete() -> None: + model = object.__new__(MiniMaxM2Model) + model.block_count = 2 + model.hparams = { + "num_local_experts": 4, + "reap_adaptive_nonuniform_experts": True, + "reap_adaptive_retained_experts_by_layer": {"0": 2}, + } + + with pytest.raises(ValueError, match="requires counts for every layer"): + model._load_adaptive_retained_by_layer() + + +def test_minimax_adaptive_export_fails_on_incomplete_layer_cache(monkeypatch: pytest.MonkeyPatch) -> None: + model = make_minimax_model({0: 2, 1: 4}) + model._experts_cache = {0: {"partial": torch.ones((2, 3))}} + monkeypatch.setattr(TextModel, "prepare_tensors", lambda self: None) + + with pytest.raises(RuntimeError, match="incomplete per-layer expert caches"): + model.prepare_tensors()