From 184979e7415bee1e15a4f791c630183950352a31 Mon Sep 17 00:00:00 2001
From: gberasmus87 <gberasmus87@gmail.com>
Date: Thu, 10 Sep 2026 20:28:50 +1200
Subject: [PATCH] fix(glm4_moe): load compressed-tensors NVFP4 expert
 checkpoints

glm4_moe is the only NVFP4 family that never learned the llm-compressor spelling.
Its expert key pattern matches `weight | weight_scale | weight_scale_2` and its
source spec sets neither `kind_map` nor `global_reciprocal`, so on a
compressed-tensors export -- `gesong2077/GLM-4.5-Air-NVFP4`, whose experts are
`weight_packed` + `weight_scale` + `weight_global_scale` -- the only tensors that
match are the `weight_scale` ones, whose name the two dialects happen to share.
The packed weights and the global scales are missed entirely, so the offload
banks fill with scales and no weights.

The machinery already exists; glm4_moe just is not wired to it.
`Nvfp4ExpertSourceSpec.kind_map` folds `weight_packed` -> `weight` and
`weight_global_scale` -> `weight_scale_2`, and `global_reciprocal` inverts the
quant-side global scale at ingest -- without which the expert dequant is inverted
and the model emits a single repeated token. glm5_next carries exactly this pair
for its own compressed-tensors release; this mirrors it, including selecting the
spec from `quant_method`.

`input_global_scale` deliberately does not match: the routed-expert path is W4A16
and never quantizes activations.

`detect_expert_quant` already recognises these exports, so `parse_config` needs
no change -- `expert_quant` is `nvfp4` either way. The reader was the only gap.

Verified on `gesong2077/GLM-4.5-Air-NVFP4` (compressed-tensors, config_groups
num_bits=4 type=float group_size=16 strategy=tensor_group), against its real
weight index: 69,120 routed-expert tensors, of which 17,280 are
`input_global_scale`. Before, the spec matched 17,280 -- every one a
`weight_scale`, no weights at all. After, it matches all 51,840 non-activation
tensors, folding 17,280 `weight_packed` -> `weight` and 17,280
`weight_global_scale` -> `weight_scale_2`, and still rejects every
`input_global_scale`. The modelopt path is untouched: an export without
`quant_method == "compressed-tensors"` still gets the original spec.
---
 python/freetoken/models/glm4_moe/weight.py | 24 ++++++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/python/freetoken/models/glm4_moe/weight.py b/python/freetoken/models/glm4_moe/weight.py
index 5d4f9fcea..6815ab80a 100644
--- a/python/freetoken/models/glm4_moe/weight.py
+++ b/python/freetoken/models/glm4_moe/weight.py
@@ -35,6 +35,24 @@
     ),
     desc="GLM NVFP4 experts",
 )
+# llm-compressor export (gesong2077/GLM-4.5-Air-NVFP4): weight_packed | weight_scale |
+# weight_global_scale (quant-side global -> reciprocal at ingest). Same treatment glm5_next
+# gives its compressed-tensors release. ``input_global_scale`` (the calibrated W4A4 activation
+# scale) deliberately does not match: the routed-expert path is W4A16 and never quantizes
+# activations.
+_ROUTED_EXPERT_CT_KEY_RE = re.compile(
+    r"^model\.layers\.(?P<layer>\d+)\.mlp\.experts\.(?P<expert>\d+)\."
+    r"(?P<proj>gate_proj|up_proj|down_proj)\."
+    r"(?P<kind>weight_packed|weight_global_scale|weight_scale)$"
+)
+_NVFP4_CT_SOURCE_SPEC = Nvfp4ExpertSourceSpec(
+    key_pattern=_ROUTED_EXPERT_CT_KEY_RE,
+    proj_to_role={"gate_proj": "gate", "up_proj": "up", "down_proj": "down"},
+    layer_to_bank=_NVFP4_SOURCE_SPEC.layer_to_bank,
+    desc="GLM NVFP4 experts (compressed-tensors)",
+    kind_map={"weight_packed": "weight", "weight_global_scale": "weight_scale_2"},
+    global_reciprocal=True,
+)
 
 
 # --------------------------------------------------------------------------------------
@@ -191,6 +209,12 @@ def _iter_resident_weights(reader, config, primary) -> Iterator[tuple[str, torch
 # Routed expert host banks (NVFP4) for the offload cache.
 # --------------------------------------------------------------------------------------
 def nvfp4_expert_spec(model_path: str, config):
+    """The expert reader for this checkpoint's dialect: modelopt names, or the
+    compressed-tensors ones folded back onto them."""
+    quant = getattr(cached_load_hf_config(model_path), "quantization_config", None) or {}
+    get = quant.get if isinstance(quant, dict) else (lambda k, d=None: getattr(quant, k, d))
+    if str(get("quant_method") or "").lower() == "compressed-tensors":
+        return _NVFP4_CT_SOURCE_SPEC
     return _NVFP4_SOURCE_SPEC