From f5edb1fadf385796c66aca81f229ad69b2bd8801 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Thu, 6 Aug 2026 09:46:18 -0700 Subject: [PATCH] fix(load): GGUFs with quantized MoE router gates failed to load Fix for antirez deepseek-v4-flash dspark drafter GGUFs. Should be narrow as llama.cpp never quantizes router gates. --- CHANGELOG.md | 6 +++ gmlx/loader.py | 16 ++++++- gmlx/modules.py | 60 +++++++++++++++++++++++++ tests/test_dequant_unattachable.py | 71 ++++++++++++++++++++++++++++++ 4 files changed, 152 insertions(+), 1 deletion(-) create mode 100644 tests/test_dequant_unattachable.py diff --git a/CHANGELOG.md b/CHANGELOG.md index a528312..970bf70 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,12 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ## [Unreleased] +### Fixed + +- GGUFs that quantize the MoE router gate (some community DeepSeek quants; + llama.cpp's own quantize leaves it F32) now load: small quantized tensors + on raw-array modules are dequantized to f32 at load instead of erroring. + ## [0.2.2] - 2026-08-06 ### Fixed diff --git a/gmlx/loader.py b/gmlx/loader.py index db80898..7434b77 100644 --- a/gmlx/loader.py +++ b/gmlx/loader.py @@ -41,7 +41,11 @@ from .quantized_sdpa_fix import install_quantized_sdpa_mask_fix from .rope_batch_fix import install_rope_batch_fix from .rotating_cache_fix import install_rotating_cache_fix -from .modules import KQuantEmbedding, install_kquant_modules +from .modules import ( + KQuantEmbedding, + dequantize_unattachable_leaves, + install_kquant_modules, +) from .populate import ( maybe_populate_for_load, start_populate, @@ -2799,6 +2803,11 @@ def _install_and_load( ) # 6. swap leaves with kquant equivalents. + dequant = dequantize_unattachable_leaves(model, hf_weights, hf_kquant_meta) + if dequant: + log(f"[install] dequantized {len(dequant)} raw-array leaves to f32: " + + ", ".join(dequant[:3]) + + (f" (+{len(dequant) - 3} more)" if len(dequant) > 3 else "")) n_replaced = install_kquant_modules( model, hf_kquant_meta, native_fp_wire=native_fp_wire) log(f"[install] replaced {n_replaced} leaves with kquant modules") @@ -3241,6 +3250,11 @@ def load_model( # 6. swap leaves with kquant equivalents. loadlog.stage("installing quantized weights") + dequant = dequantize_unattachable_leaves(model, hf_weights, hf_kquant_meta) + if dequant: + _log(f"[install] dequantized {len(dequant)} raw-array leaves to f32: " + + ", ".join(dequant[:3]) + + (f" (+{len(dequant) - 3} more)" if len(dequant) > 3 else "")) n_replaced = install_kquant_modules( model, hf_kquant_meta, native_fp_wire=native_fp_wire) _log(f"[install] replaced {n_replaced} leaves with kquant modules") diff --git a/gmlx/modules.py b/gmlx/modules.py index ea20803..c4332fe 100644 --- a/gmlx/modules.py +++ b/gmlx/modules.py @@ -1409,6 +1409,66 @@ def _wrap(path: str, module): return len(wrapped) +def dequantize_unattachable_leaves(model: nn.Module, + hf_weights: dict, + hf_kquant_meta: dict[str, str], + max_bytes: int = 64 * 1024 * 1024) -> list[str]: + """Dequantize codec'd weights on modules the installer cannot swap. + + Some modules hold their weight as a raw array used inline in the + forward (deepseek-family MoEGate routers), so there is no leaf to + swap a kquant module into. Dequantize those wire bytes to f32 at + load and drop the codec entry. Tensors over ``max_bytes`` (f32 + size) are left alone so ``install_kquant_modules`` fails loud on + them instead of silently materializing gigabytes of float. + + llama.cpp's quantize never codecs router gates, so its GGUFs never + hit this path; so far only an antirez DeepSeek-V4-Flash dspark + drafter GGUF does. If quantized routers become a pattern in model + GGUFs, this approach should change to quantized execution. + + Returns the handled ``path (codec)`` strings for load logging. + """ + import mlx_kquant as kq + + _switch_linear_types, _ = switch_layer_types() + attachable = [nn.Linear, nn.Embedding, KQuantLinear, KQuantEmbedding, + KQuantSwitchLinear, KQuantMultiLinear, NativeFPSwitchLinear] + attachable.extend(_switch_linear_types) + if MultiLinear is not None: + attachable.append(MultiLinear) + attachable = tuple(attachable) + known_codecs = set(kq.codecs()) + handled: list[str] = [] + + def _visit(path: str, module): + weight_key = f"{path}.weight" + codec = hf_kquant_meta.get(weight_key) + if codec is None or isinstance(module, attachable): + return module + target = getattr(module, "weight", None) + if not isinstance(target, mx.array) or codec not in known_codecs: + return module + if target.size * 4 > max_bytes: + return module + scales_key = f"{path}.scales" + scales = hf_weights.get(scales_key) + if scales is None: + scales = mx.zeros((1,), dtype=mx.uint8) + else: + del hf_weights[scales_key] + deq = kq.dequantize(hf_weights[weight_key], scales, codec, mx.float32) + del hf_weights[weight_key] + hf_weights[weight_key] = deq.reshape(target.shape) + del hf_kquant_meta[weight_key] + handled.append(f"{path} ({codec})") + return module + + leaves = model.leaf_modules() + tree_map_with_path(_visit, leaves, is_leaf=nn.Module.is_module) + return handled + + def install_kquant_modules(model: nn.Module, hf_kquant_meta: dict[str, str], native_fp_wire: bool = False) -> int: diff --git a/tests/test_dequant_unattachable.py b/tests/test_dequant_unattachable.py new file mode 100644 index 0000000..cbf2d34 --- /dev/null +++ b/tests/test_dequant_unattachable.py @@ -0,0 +1,71 @@ +"""Codec'd tensors on raw-array carrier modules dequantize to f32 at load. + +Some community GGUF quantizers codec the MoE router gate (llama.cpp's own +quantize leaves it F32). The deepseek-family MoEGate holds that weight as a +raw array, so no KQuant* module can attach; the loader reconstructs the float +tensor instead of failing. Large unattachable tensors keep failing loud. +""" + +import mlx.core as mx +import mlx.nn as nn +import mlx_kquant as kq +import pytest + +from gmlx.modules import dequantize_unattachable_leaves, install_kquant_modules + + +class RawGate(nn.Module): + """MoEGate-style carrier: bare weight array, no Linear.""" + + def __init__(self, n_experts, dims): + super().__init__() + self.weight = mx.zeros((n_experts, dims)) + + +class Block(nn.Module): + def __init__(self, n_experts, dims): + super().__init__() + self.gate = RawGate(n_experts, dims) + self.proj = nn.Linear(dims, dims, bias=False) + + +def _quantized_block(): + mx.random.seed(0) + model = Block(8, 256) + w = mx.random.normal((8, 256)) + wq, scales = kq.quantize(w, "q8_0") + hf_weights = {"gate.weight": wq, "gate.scales": scales} + meta = {"gate.weight": "q8_0", "proj.weight": "q8_0"} + return model, wq, scales, hf_weights, meta + + +def test_small_raw_leaf_dequantized(): + model, wq, scales, hf_weights, meta = _quantized_block() + handled = dequantize_unattachable_leaves(model, hf_weights, meta) + + assert handled == ["gate (q8_0)"] + # Gate is now a plain f32 tensor with the scales sidecar dropped; the + # Linear leaf stays codec'd for install_kquant_modules. + assert "gate.weight" not in meta + assert meta == {"proj.weight": "q8_0"} + assert "gate.scales" not in hf_weights + deq = hf_weights["gate.weight"] + assert deq.shape == (8, 256) + assert deq.dtype == mx.float32 + ref = kq.dequantize(wq, scales, "q8_0", mx.float32).reshape(8, 256) + assert mx.array_equal(deq, ref) + + n = install_kquant_modules(model, meta) + assert n == 1 # proj swapped; gate left as the raw float carrier + + +def test_large_raw_leaf_still_fails_loud(): + model, _, _, hf_weights, meta = _quantized_block() + handled = dequantize_unattachable_leaves( + model, hf_weights, meta, max_bytes=1024) + + assert handled == [] + assert meta["gate.weight"] == "q8_0" + assert "gate.scales" in hf_weights + with pytest.raises(ValueError, match="no recognized module class"): + install_kquant_modules(model, meta)