Parcourir la source

glm53-flash-nvfp4-4p67: MXFP8-attention support patchset + boot verified on ds4

4p67 breaks the day-0 "attention/shared experts are BF16" assumption: it
stores MXFP8 scales for KDA projections, MLA projections and the indexer.
Five targeted patches (patch_4p67.py, all idempotent + compile-checked):

1. packed_modules_mapping class attrs (CausalLM + ConditionalGeneration +
   MTP) so the ModelOpt MIXED_PRECISION resolver resolves fused linear
   names to the per-shard MXFP8 entries.
2. KDA BF16-strip gate: only null quant_config when the checkpoint truly
   lacks KDA shard scales.
3. Indexer wk_weights_proj quant gate (same pattern).
4. FP8 dequant loader shims bypassed when the target param is quantized
   (4p67 names scales weight_scale, the legacy _inv checks would swallow
   the fp8 weight and drop the scale).
5. DecoderLayer MLA quant_config=None gate -> quantized when checkpoint
   declares MLA scales.
Shared _glm5_quantized() helper matches quantized_layers suffix-agnostically
(HF-, vLLM- and MTP-spelled prefixes all hit).

Also: GMU 0.90 -> 0.82 in lab/4p67/dflash2 recipes. GMU is only a startup
guardrail once KV is pinned; 0.90 demands 109.46 GiB free on GB10 and
driver reserve + post-crash leaks made boots flaky ("54-second death"
class). Boot to serving verified on ds4: 1.16M token KV, chat smoke OK.
kilork il y a 1 mois
Parent
commit
2599e39186

+ 275 - 3
docker-images/glm53-flash-nvfp4-4p67/patch_4p67.py

@@ -35,22 +35,38 @@ MAPPING = '''
     }
 '''
 
-TARGETS = [
+# 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 TARGETS:
+for path, anchor in MAPPING_TARGETS:
     src = open(path).read()
-    if "packed_modules_mapping" in src:
+    # 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)
@@ -58,4 +74,260 @@ for path, anchor in TARGETS:
     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")

+ 4 - 1
recipes/glm53-flash-dflash2-mxfp8/recipe.toml

@@ -62,7 +62,10 @@ args = [
     "--served-model-name", "glm53",
     "--trust-remote-code",
     "--distributed-executor-backend", "mp",
-    "--gpu-memory-utilization", "0.90",
+    # GMU is ONLY a startup guardrail (KV pinned below): 0.90 demands
+    # 109.46 GiB free on GB10's 121.63 and driver reserve + page cache
+    # routinely leave ~108.5 -> flaky boot deaths. 0.82 always fits.
+    "--gpu-memory-utilization", "0.82",
     # Pinned 10 GiB KV; the draft adds 5 small sliding-window layers on top
     "--kv-cache-memory-bytes", "10737418240",
     "--kv-cache-dtype", "fp8_ds_mla",

+ 8 - 2
recipes/glm53-flash-nvfp4-4p67/recipe.toml

@@ -54,8 +54,14 @@ args = [
     "--served-model-name", "glm53",
     "--trust-remote-code",
     "--distributed-executor-backend", "mp",
-    # GMU is superseded by the pinned KV pool; 0.90 as upstream boots with it
-    "--gpu-memory-utilization", "0.90",
+    # GMU is ONLY a startup guardrail here (KV pool is pinned below): 0.90
+    # demands 109.46 GiB free on GB10's 121.63 GiB, but driver reserve +
+    # page cache routinely leave ~108.5 -> flaky
+    # "Free memory ... less than desired" boot deaths (this was the class
+    # of the old "54-second death"), and post-crash reruns see ~106.6 GiB.
+    # 0.82 = 99.7 GiB, below any realistic dip; does NOT change what gets
+    # allocated (weights ~93 + pinned KV 10 + activations).
+    "--gpu-memory-utilization", "0.82",
     # Pinned 10 GiB KV (≈1.16M tokens) — GMU sizing has no headroom at
     # ~93 GiB weights/node; pinning also keeps warmup off the swap path
     "--kv-cache-memory-bytes", "10737418240",

+ 8 - 2
recipes/glm53-flash-nvfp4-lab/recipe.toml

@@ -59,8 +59,14 @@ args = [
     "--served-model-name", "glm53",
     "--trust-remote-code",
     "--distributed-executor-backend", "mp",
-    # GMU is superseded by the pinned KV pool; 0.90 as upstream boots with it
-    "--gpu-memory-utilization", "0.90",
+    # GMU is ONLY a startup guardrail here (KV pool is pinned below): 0.90
+    # demands 109.46 GiB free on GB10's 121.63 GiB, but driver reserve +
+    # page cache routinely leave ~108.5 -> flaky
+    # "Free memory ... less than desired" boot deaths (this was the class
+    # of the old "54-second death"), and post-crash reruns see ~106.6 GiB.
+    # 0.82 = 99.7 GiB, below any realistic dip; does NOT change what gets
+    # allocated (weights ~93 + pinned KV 10 + activations).
+    "--gpu-memory-utilization", "0.82",
     # Pinned 10 GiB KV (≈1.16M tokens) — GMU sizing has no headroom at
     # 92.72 GiB weights/node; pinning also keeps warmup off the swap path
     "--kv-cache-memory-bytes", "10737418240",