|
|
@@ -0,0 +1,333 @@
|
|
|
+#!/usr/bin/env python3
|
|
|
+"""Teach the glm5next model classes their fused projection names.
|
|
|
+
|
|
|
+local-inference-lab/GLM-5.3-Flash-NVFP4-4p67 is the first checkpoint of this
|
|
|
+family that QUANTIZES the individual attention projections (per-shard
|
|
|
+MXFP8 entries `self_attn.q_proj` … in hf_quant_config). vLLM fuses those
|
|
|
+shards into single modules (`in_proj_qkvbfg_a`, `fused_qkv_a_proj`,
|
|
|
+`wk_weights_proj`), and the ModelOpt MIXED_PRECISION config resolves a
|
|
|
+fused module's quant_algo through its shards ONLY via the model class's
|
|
|
+`packed_modules_mapping` (handed to the quant config by
|
|
|
+model_executor/model_loader/utils.py). Without it every fused linear comes
|
|
|
+up UnquantizedLinearMethod and weight loading dies with
|
|
|
+
|
|
|
+ KeyError: 'layers.0.self_attn.in_proj_qkvbfg_a.weight_scale'
|
|
|
+
|
|
|
+Anchors self-assert; idempotent.
|
|
|
+"""
|
|
|
+
|
|
|
+import py_compile
|
|
|
+import sys
|
|
|
+
|
|
|
+B = "/usr/local/lib/python3.12/dist-packages/vllm/models/glm5next/nvidia"
|
|
|
+
|
|
|
+MAPPING = '''
|
|
|
+ # NVFP4-4p67: quantized per-shard attention/shared-expert projections
|
|
|
+ # are fused by vLLM; expose the shard names so the ModelOpt
|
|
|
+ # MIXED_PRECISION config can resolve the fused module's quant_algo
|
|
|
+ # (see model_loader/utils.py packed_modules_mapping pass-through).
|
|
|
+ packed_modules_mapping = {
|
|
|
+ "gate_up_proj": ["gate_proj", "up_proj"],
|
|
|
+ "in_proj_qkvbfg_a": ["q_proj", "k_proj", "v_proj",
|
|
|
+ "b_proj", "f_a_proj", "g_a_proj"],
|
|
|
+ "fused_qkv_a_proj": ["q_a_proj", "kv_a_proj_with_mqa"],
|
|
|
+ "wk_weights_proj": ["wk", "weights_proj"],
|
|
|
+ }
|
|
|
+'''
|
|
|
+
|
|
|
+# CRITICAL: it must live on Glm5NextForConditionalGeneration (the registered
|
|
|
+# class). configure_quant_config reads BOTH hf_to_vllm_mapper and
|
|
|
+# packed_modules_mapping off that class; Glm4v's inherited mapper rewrites
|
|
|
+# the config's quantized_layers keys into vLLM spelling, and a class-level
|
|
|
+# mapping on the inner CausalLM would SHADOW the inherited Glm4v mapper's
|
|
|
+# companion config path. MTP is constructed as a separate root model, so it
|
|
|
+# needs its own copy.
|
|
|
+MAPPING_TARGETS = [
|
|
|
+ (
|
|
|
+ f"{B}/model.py",
|
|
|
+ "class Glm5NextForCausalLM(\n"
|
|
|
+ " nn.Module, HasInnerState, SupportsPP, MixtureOfExperts, IsHybrid\n"
|
|
|
+ "):\n",
|
|
|
+ ),
|
|
|
+ (
|
|
|
+ f"{B}/model.py",
|
|
|
+ "class Glm5NextForConditionalGeneration(\n"
|
|
|
+ " Glm4vForConditionalGeneration, HasInnerState, IsHybrid\n"
|
|
|
+ "):\n",
|
|
|
+ ),
|
|
|
+ (
|
|
|
+ f"{B}/mtp.py",
|
|
|
+ "class Glm5NextMTP(nn.Module, DeepseekV2MixtureOfExperts):\n",
|
|
|
+ ),
|
|
|
+]
|
|
|
+
|
|
|
+for path, anchor in MAPPING_TARGETS:
|
|
|
+ src = open(path).read()
|
|
|
+ # count INSERTED BLOCKS (marker appears once per patch), not the plain
|
|
|
+ # string (which also appears in comments and inherited bases).
|
|
|
+ want = sum(1 for p, _ in MAPPING_TARGETS if p == path)
|
|
|
+ if src.count('"wk_weights_proj": ["wk", "weights_proj"],') >= want:
|
|
|
+ print(f"already patched: {path}")
|
|
|
+ continue
|
|
|
+ n = src.count(anchor)
|
|
|
+ assert n == 1, f"anchor {n}x in {path} — day-0 image drifted, review me"
|
|
|
+ open(path, "w").write(src.replace(anchor, anchor + MAPPING, 1))
|
|
|
+ py_compile.compile(path, doraise=True)
|
|
|
+ print(f"patched: {path}")
|
|
|
+# KDA strip gate. day-0 KDA construction FORCIBLY unquantizes its
|
|
|
+# projections ("BF16 even in fp8 checkpoints") by nulling
|
|
|
+# vllm_config.quant_config around super().__init__ — which is exactly the
|
|
|
+# 4p67 assumption-break: this checkpoint stores MXFP8 scales for
|
|
|
+# self_attn.{q,k,v,b,f_a,g_a}_proj, so the fused linear comes up BF16 and
|
|
|
+# the loader dies routing those scales. Make the strip conditional on the
|
|
|
+# checkpoint genuinely lacking KDA shard scales.
|
|
|
+KDA_STRIP_OLD = ''' saved_quant_config = vllm_config.quant_config
|
|
|
+ vllm_config.quant_config = None
|
|
|
+ super().__init__(config, vllm_config, prefix)
|
|
|
+ vllm_config.quant_config = saved_quant_config
|
|
|
+'''
|
|
|
+
|
|
|
+KDA_STRIP_NEW = ''' saved_quant_config = vllm_config.quant_config
|
|
|
+ _qc = vllm_config.quant_config
|
|
|
+ _kda_quantized = _qc is not None and any(
|
|
|
+ ("%s.%s" % (prefix, _s)) in getattr(_qc, "quantized_layers", {})
|
|
|
+ for _s in (
|
|
|
+ "q_proj", "k_proj", "v_proj",
|
|
|
+ "b_proj", "f_a_proj", "g_a_proj",
|
|
|
+ )
|
|
|
+ )
|
|
|
+ if not _kda_quantized:
|
|
|
+ # day-0 checkpoints keep the KDA projections BF16 (no scales
|
|
|
+ # stored for them); NVFP4-4p67 stores MXFP8 -> keep quant.
|
|
|
+ vllm_config.quant_config = None
|
|
|
+ super().__init__(config, vllm_config, prefix)
|
|
|
+ vllm_config.quant_config = saved_quant_config
|
|
|
+'''
|
|
|
+
|
|
|
+kda_path = f"{B}/kda.py"
|
|
|
+src = open(kda_path).read()
|
|
|
+if "_kda_quantized" in src:
|
|
|
+ print(f"already patched: {kda_path}")
|
|
|
+else:
|
|
|
+ n = src.count(KDA_STRIP_OLD)
|
|
|
+ assert n == 1, f"KDA strip anchor {n}x in {kda_path} — review me"
|
|
|
+ open(kda_path, "w").write(src.replace(KDA_STRIP_OLD, KDA_STRIP_NEW, 1))
|
|
|
+ py_compile.compile(kda_path, doraise=True)
|
|
|
+ print(f"patched: {kda_path}")
|
|
|
+
|
|
|
+# (2) Indexer wk_weights_proj: attention.py hardcodes quant_config=None
|
|
|
+# (legacy path dequantizes FP8 wk to BF16 on load). Gate it the same way as
|
|
|
+# KDA: if the checkpoint declares MXFP8 for indexer.wk / indexer.weights_proj
|
|
|
+# (NVFP4-4p67), build the fused linear QUANTIZED and let the stacked loader
|
|
|
+# route weights+scales.
|
|
|
+IDX_OLD = ''' self.wk_weights_proj = MergedColumnParallelLinear(
|
|
|
+ hidden_size,
|
|
|
+ [self.head_dim, self.n_head],
|
|
|
+ bias=False,
|
|
|
+ quant_config=None,
|
|
|
+'''
|
|
|
+IDX_NEW = ''' _idx_quantized = quant_config is not None and any(
|
|
|
+ ("%s.%s" % (prefix, _s)) in getattr(quant_config, "quantized_layers", {})
|
|
|
+ for _s in ("wk", "weights_proj")
|
|
|
+ )
|
|
|
+ self.wk_weights_proj = MergedColumnParallelLinear(
|
|
|
+ hidden_size,
|
|
|
+ [self.head_dim, self.n_head],
|
|
|
+ bias=False,
|
|
|
+ quant_config=(quant_config if _idx_quantized else None),
|
|
|
+'''
|
|
|
+
|
|
|
+apath = f"{B}/attention.py"
|
|
|
+src = open(apath).read()
|
|
|
+if "_idx_quantized" in src:
|
|
|
+ print(f"already patched: {apath}")
|
|
|
+else:
|
|
|
+ n = src.count(IDX_OLD)
|
|
|
+ assert n == 1, f"indexer anchor {n}x in {apath} — review me"
|
|
|
+ open(apath, "w").write(src.replace(IDX_OLD, IDX_NEW, 1))
|
|
|
+ py_compile.compile(apath, doraise=True)
|
|
|
+ print(f"patched: {apath}")
|
|
|
+
|
|
|
+# (3) Loader shims: legacy FP8 dequant helpers must not eat MXFP8 tensors
|
|
|
+# (4p67 names scales `weight_scale` — the wk shim's is_weight test would
|
|
|
+# swallow the fp8 weight and then drop the scale). Skip the shims when the
|
|
|
+# target fused param is quantized (has a weight_scale param).
|
|
|
+model_path = f"{B}/model.py"
|
|
|
+src = open(model_path).read()
|
|
|
+WK_OLD = '''def _try_load_fp8_indexer_wk(name, tensor, buf, params_dict, loaded_params):
|
|
|
+ if "indexer.wk." not in name or "wk_weights" in name:
|
|
|
+ return False
|
|
|
+'''
|
|
|
+WK_NEW = '''def _try_load_fp8_indexer_wk(name, tensor, buf, params_dict, loaded_params):
|
|
|
+ if "indexer.wk." not in name or "wk_weights" in name:
|
|
|
+ return False
|
|
|
+ if (f"{name.rsplit('.wk.', 1)[0]}.wk_weights_proj.weight_scale"
|
|
|
+ in params_dict):
|
|
|
+ # MXFP8 checkpoint: fused param is quantized, stacked loader owns it.
|
|
|
+ return False
|
|
|
+'''
|
|
|
+ATTN_OLD = ''' target_w = f"{layer_prefix}.{target_base}.weight"
|
|
|
+ target_s = f"{layer_prefix}.{target_base}.weight_scale_inv"
|
|
|
+'''
|
|
|
+ATTN_NEW = ''' target_w = f"{layer_prefix}.{target_base}.weight"
|
|
|
+ target_s = f"{layer_prefix}.{target_base}.weight_scale_inv"
|
|
|
+ if f"{layer_prefix}.{target_base}.weight_scale" in params_dict:
|
|
|
+ # MXFP8 checkpoint: param is quantized, stacked loader owns it.
|
|
|
+ return False
|
|
|
+'''
|
|
|
+if "stacked loader owns it" in src:
|
|
|
+ print(f"already patched: {model_path} (wk shim)")
|
|
|
+else:
|
|
|
+ n = src.count(WK_OLD)
|
|
|
+ assert n == 1, f"wk-shim anchor {n}x in {model_path} — review me"
|
|
|
+ src = src.replace(WK_OLD, WK_NEW, 1)
|
|
|
+ open(model_path, "w").write(src)
|
|
|
+ py_compile.compile(model_path, doraise=True)
|
|
|
+ print(f"patched: {model_path} (wk shim)")
|
|
|
+
|
|
|
+if "MXFP8 checkpoint: param is quantized" in src:
|
|
|
+ print(f"already patched: {model_path} (attn shim)")
|
|
|
+else:
|
|
|
+ src = open(model_path).read()
|
|
|
+ n = src.count(ATTN_OLD)
|
|
|
+ assert n == 1, f"attn-shim anchor {n}x in {model_path} — review me"
|
|
|
+ src = src.replace(ATTN_OLD, ATTN_NEW, 1)
|
|
|
+ open(model_path, "w").write(src)
|
|
|
+ py_compile.compile(model_path, doraise=True)
|
|
|
+ print(f"patched: {model_path} (attn shim)")
|
|
|
+
|
|
|
+def helper(name, uses):
|
|
|
+ h = (
|
|
|
+ "def _glm5_q_layers(quant_config):\n"
|
|
|
+ " # quantized_layers keys are HF-spelled at module-construction\n"
|
|
|
+ " # time and vLLM-spelled after apply_vllm_mapper; test both.\n"
|
|
|
+ " return set(getattr(quant_config, 'quantized_layers', {}) or {})\n"
|
|
|
+ "\n"
|
|
|
+ "\n"
|
|
|
+ "def _glm5_quantized(quant_config, prefix, shards):\n"
|
|
|
+ " layers = _glm5_q_layers(quant_config)\n"
|
|
|
+ " if not layers:\n"
|
|
|
+ " return False\n"
|
|
|
+ " tail = prefix\n"
|
|
|
+ " for lead in (\n"
|
|
|
+ " 'language_model.model.', 'model.language_model.',\n"
|
|
|
+ " 'model.', 'language_model.',\n"
|
|
|
+ " ):\n"
|
|
|
+ " if tail.startswith(lead):\n"
|
|
|
+ " tail = tail[len(lead):]\n"
|
|
|
+ " break\n"
|
|
|
+ " return any(\n"
|
|
|
+ " k.endswith(f'.{tail}.{s}') or k == f'{tail}.{s}'\n"
|
|
|
+ " for s in shards\n"
|
|
|
+ " for k in layers\n"
|
|
|
+ " )\n\n\n"
|
|
|
+ )
|
|
|
+ # place right before the first class that needs it
|
|
|
+ anchor = uses
|
|
|
+ return h, anchor
|
|
|
+
|
|
|
+
|
|
|
+for path, defs_marker, insert_before in [
|
|
|
+ (f"{B}/kda.py", "_glm5_quantized", "class Glm5NextLinearAttention"),
|
|
|
+ (f"{B}/attention.py", "_glm5_quantized", "class Indexer(nn.Module):"),
|
|
|
+]:
|
|
|
+ src = open(path).read()
|
|
|
+ if defs_marker in src:
|
|
|
+ print(f"helper already in {path}")
|
|
|
+ continue
|
|
|
+ idx = src.index(insert_before)
|
|
|
+ helper_src, _ = helper(path, insert_before)
|
|
|
+ src = src[:idx] + helper_src + src[idx:]
|
|
|
+ open(path, "w").write(src)
|
|
|
+ py_compile.compile(path, doraise=True)
|
|
|
+ print(f"helper added: {path}")
|
|
|
+
|
|
|
+# rewrite the two constructor gates to use the both-spellings helper
|
|
|
+GATES = [
|
|
|
+ (
|
|
|
+ f"{B}/kda.py",
|
|
|
+ ''' _kda_quantized = _qc is not None and any(
|
|
|
+ ("%s.%s" % (prefix, _s)) in getattr(_qc, "quantized_layers", {})
|
|
|
+ for _s in (
|
|
|
+ "q_proj", "k_proj", "v_proj",
|
|
|
+ "b_proj", "f_a_proj", "g_a_proj",
|
|
|
+ )
|
|
|
+ )
|
|
|
+''',
|
|
|
+ " _kda_quantized = _glm5_quantized(\n"
|
|
|
+ " _qc, prefix,\n"
|
|
|
+ ' ("q_proj", "k_proj", "v_proj",\n'
|
|
|
+ ' "b_proj", "f_a_proj", "g_a_proj"),\n'
|
|
|
+ " )\n",
|
|
|
+ ),
|
|
|
+ (
|
|
|
+ f"{B}/attention.py",
|
|
|
+ ''' _idx_quantized = quant_config is not None and any(
|
|
|
+ ("%s.%s" % (prefix, _s)) in getattr(quant_config, "quantized_layers", {})
|
|
|
+ for _s in ("wk", "weights_proj")
|
|
|
+ )
|
|
|
+''',
|
|
|
+ " _idx_quantized = _glm5_quantized(\n"
|
|
|
+ ' quant_config, prefix, ("wk", "weights_proj")\n'
|
|
|
+ " )\n"
|
|
|
+ " import logging as _lg\n"
|
|
|
+ ' _lg.getLogger("vllm").warning(\n'
|
|
|
+ ' "4p67 IDX GATE prefix=%s q=%s qc=%s n=%s",\n'
|
|
|
+ ' prefix, _idx_quantized,\n'
|
|
|
+ " type(quant_config).__name__,\n"
|
|
|
+ ' len(getattr(quant_config, "quantized_layers", {}) or {}),\n'
|
|
|
+ " )\n",
|
|
|
+ ),
|
|
|
+]
|
|
|
+for path, old, new in GATES:
|
|
|
+ src = open(path).read()
|
|
|
+ if "_glm5_quantized(" in src.split("def _glm5_quantized", 1)[-1]:
|
|
|
+ print(f"gate already rewritten: {path}")
|
|
|
+ continue
|
|
|
+ n = src.count(old)
|
|
|
+ assert n == 1, f"gate anchor {n}x in {path} — review me"
|
|
|
+ open(path, "w").write(src.replace(old, new, 1))
|
|
|
+ py_compile.compile(path, doraise=True)
|
|
|
+ print(f"gate rewritten: {path}")
|
|
|
+
|
|
|
+# (5) DecoderLayer hardcodes quant_config=None for MLA attention ("BF16 in
|
|
|
+# checkpoint" — again broken by 4p67). Gate it; helper lives in kda.py
|
|
|
+# (always importable; model.py already imports from it? if not, guard).
|
|
|
+MLA_OLD = " quant_config=None, # MLA projections are BF16 in checkpoint\n"
|
|
|
+MLA_NEW = """ quant_config=(
|
|
|
+ quant_config
|
|
|
+ if _glm5_quantized(
|
|
|
+ quant_config,
|
|
|
+ f"{prefix}.self_attn",
|
|
|
+ ("q_a_proj", "kv_a_proj_with_mqa"),
|
|
|
+ )
|
|
|
+ else None
|
|
|
+ ), # 4p67: quantized when checkpoint stores MLA scales
|
|
|
+"""
|
|
|
+mla_path = f"{B}/model.py"
|
|
|
+src = open(mla_path).read()
|
|
|
+if "4p67: quantized when checkpoint stores MLA scales" in src:
|
|
|
+ print(f"already patched: {mla_path} (mla gate)")
|
|
|
+else:
|
|
|
+ n = src.count(MLA_OLD)
|
|
|
+ assert n == 1, f"mla-gate anchor {n}x in {mla_path} — review me"
|
|
|
+ src = src.replace(MLA_OLD, MLA_NEW, 1)
|
|
|
+ if "from .kda import" in src:
|
|
|
+ import re as _re
|
|
|
+
|
|
|
+ src = _re.sub(
|
|
|
+ r"from \.kda import (\w+)",
|
|
|
+ r"from .kda import \1, _glm5_quantized",
|
|
|
+ src,
|
|
|
+ count=1,
|
|
|
+ )
|
|
|
+ else:
|
|
|
+ raise AssertionError(
|
|
|
+ f"no 'from .kda import' in {mla_path} — add import explicitly"
|
|
|
+ )
|
|
|
+ assert "_glm5_quantized" in "\n".join(src.splitlines()[:120]), "import missing"
|
|
|
+ open(mla_path, "w").write(src)
|
|
|
+ py_compile.compile(mla_path, doraise=True)
|
|
|
+ print(f"patched: {mla_path} (mla gate)")
|
|
|
+
|
|
|
+print("patch_4p67 OK")
|
