Browse Source

glm53-flash-nvfp4-4p67: dedicated image with packed_modules_mapping fix

First ds4 boot of the 4p67 quant died after full weight load with
KeyError layers.0.self_attn.in_proj_qkvbfg_a.weight_scale: 4p67 is the
first glm5next checkpoint whose attention projections are quantized
(per-shard MXFP8), and vLLM fuses them (in_proj_qkvbfg_a /
fused_qkv_a_proj / wk_weights_proj) — the fused linears only get a
quant method if the model class exposes packed_modules_mapping (read by
model_loader/utils.py, consumed by _resolve_quant_algo strategy 2).

Overlay image = lab:local + patch_4p67.py (class attr on
Glm5NextForCausalLM + Glm5NextMTP; anchors self-assert at build).
Chose this over porting the fork glm5next package (restructured,
hard-imports the amd64-only b12x backend) and over bf16-dequant loader
shims (more code, loses the MXFP8 GEMMs the card ships).
Alexander Korolev 1 month ago
parent
commit
30268d382b

+ 22 - 0
docker-images/glm53-flash-nvfp4-4p67/Dockerfile

@@ -0,0 +1,22 @@
+# syntax=docker/dockerfile:1.7
+# GLM-5.3-Flash NVFP4-4p67 (local-inference-lab fidelity quant) on DGX Spark.
+#
+# FROM the proven lab image (day-0 base + NoPE-MLA + FujitsuPolycom patches —
+# which already carry the W4A16_NVFP4 / MIXED_PRECISION / MXFP8 modelopt
+# methods this quant needs) and add ONE structural fix: glm5next model-class
+# `packed_modules_mapping` declarations (model.py + mtp.py). 4p67 is the
+# first glm5next checkpoint with QUANTIZED attention projections; without
+# the shard mapping the fused KDA/MLA/indexer linears load unquantized and
+# weight loading dies on `in_proj_qkvbfg_a.weight_scale`.
+# Provenance: mirrors the fork dev/jovian-judgement resolution path
+# (local-inference-lab/vllm, Apache-2.0) without porting the restructured
+# package (which hard-depends on the amd64-only b12x backend).
+
+ARG BASE_IMAGE
+FROM ${BASE_IMAGE}
+
+COPY patch_4p67.py /tmp/patch_4p67.py
+
+RUN set -eux; \
+    python3 /tmp/patch_4p67.py; \
+    rm -f /tmp/patch_4p67.py

+ 24 - 0
docker-images/glm53-flash-nvfp4-4p67/build.sh

@@ -0,0 +1,24 @@
+#!/usr/bin/env bash
+# Build the SM121 GLM-5.3-Flash NVFP4-4p67 image: lab image + the
+# packed_modules_mapping model-class patch (see Dockerfile header).
+#
+# Generic driver: forwards every KEY=VAL pin in versions.env as a --build-arg
+# to the Dockerfile. TAG is set by the sm121 launcher to the recipe's
+# docker_image (manual build: TAG=glm53-flash-nvfp4-4p67:local bash build.sh).
+set -euo pipefail
+cd "$(dirname "$0")"
+TAG="${TAG:-glm53-flash-nvfp4-4p67:local}"
+
+ARGS=()
+while read -r line; do
+    case "$line" in
+        ''|\#*) continue ;;
+    esac
+    key="${line%%=*}"; val="${line#*=}"
+    val="${val%\"}"; val="${val#\"}"   # versions.env values may be shell-quoted
+    [ -n "$key" ] && ARGS+=(--build-arg "${key}=${val}")
+done < versions.env
+
+echo "=== Building ${TAG} ==="
+DOCKER_BUILDKIT=1 docker build "${ARGS[@]}" --progress=plain -t "${TAG}" -f Dockerfile . 2>&1
+echo "=== Done: ${TAG} ==="

+ 61 - 0
docker-images/glm53-flash-nvfp4-4p67/patch_4p67.py

@@ -0,0 +1,61 @@
+#!/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"],
+    }
+'''
+
+TARGETS = [
+    (
+        f"{B}/model.py",
+        "class Glm5NextForCausalLM(\n"
+        "    nn.Module, HasInnerState, SupportsPP, MixtureOfExperts, IsHybrid\n"
+        "):\n",
+    ),
+    (
+        f"{B}/mtp.py",
+        "class Glm5NextMTP(nn.Module, DeepseekV2MixtureOfExperts):\n",
+    ),
+]
+
+for path, anchor in TARGETS:
+    src = open(path).read()
+    if "packed_modules_mapping" in src:
+        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}")
+print("patch_4p67 OK")

+ 3 - 0
docker-images/glm53-flash-nvfp4-4p67/versions.env

@@ -0,0 +1,3 @@
+# Authoritative build inputs for the GLM-5.3-Flash NVFP4-4p67 image.
+# Chain root is digest-pinned in ../glm53-flash-nvfp4-lab/versions.env.
+BASE_IMAGE="glm53-flash-nvfp4-lab:local"

+ 17 - 14
recipes/glm53-flash-nvfp4-4p67/recipe.toml

@@ -14,19 +14,22 @@
 #
 # IMAGE: the card says "should be used with local-inference-lab/vllm branch
 # dev/jovian-judgement" — every image they publish is amd64-only (RTX PRO
-# 6000 world). Our glm53-flash-nvfp4-lab:local base is the same day-0 family
-# and verifiably carries the whole required surface (W4A16_NVFP4 linear+MoE
-# methods, MIXED_PRECISION config, ModelOptMxFp8LinearMethod all present;
-# the FujitsuPolycom patches ride on top), so this quant runs on the SAME
-# image as the lab quant — zero extra build. Verified-by-inspection
-# 2026-08-29; first boot must pass the token-loop sanity check.
+# 6000 world, and the fork's model.py hard-imports the b12x backend). Our
+# day-0 lab base already carries the W4A16/MIXED_PRECISION/MXFP8 modelopt
+# methods; the ONE thing 4p67 additionally needs is fused-projection shard
+# resolution: it is the first glm5next checkpoint with quantized attention
+# projections, and without `packed_modules_mapping` on the model classes the
+# fused KDA/MLA/indexer linears load unquantized (KeyError
+# in_proj_qkvbfg_a.weight_scale after ~12 min of weight loading). So:
+# glm53-flash-nvfp4-4p67:local = lab image + patch_4p67.py (5-line class
+# attr on model.py + mtp.py, mirrors the fork's resolution path).
 #
-# Boot markers (engine log): weight load ~93 GiB/node (≈700 s) →
-# "GPU KV cache size: 1,164,369 tokens" → "Application startup complete"
-# with the MTP speculator attached. The lab-quant "[quantprobe] ...
-# algo=MXFP8" line is NOT expected: this checkpoint's MTP experts are
-# W4A16_NVFP4. Watch SpecDecoding acceptance lines (target: len ~3, same
-# MTP-3 draft path as the lab quant).
+# Boot markers: weight load ~93 GiB/node (≈700 s) → "GPU KV cache size:
+# 1,164,369 tokens" → "Application startup complete" with the MTP speculator
+# attached. quantprobe lines here show algo=MXFP8 for attention projections
+# and W4A16_NVFP4 for routed experts (lab quant showed algo=None for attn —
+# that difference is the whole point of this image). Watch SpecDecoding
+# acceptance (MTP-3, target len ~3).
 #
 # Test request: chat completions need "chat_template_kwargs":{"reasoning_effort"
 # :"low"} — omitted/other values select MAX reasoning effort (empty content on
@@ -37,8 +40,8 @@
 
 engine = "vllm"
 name = "GLM-5.3-Flash NVFP4-4p67 (vLLM, 2x Spark, MTP-3)"
-docker_image = "glm53-flash-nvfp4-lab:local"
-docker_image_dir = "glm53-flash-nvfp4-lab"
+docker_image = "glm53-flash-nvfp4-4p67:local"
+docker_image_dir = "glm53-flash-nvfp4-4p67"
 model = "local-inference-lab/GLM-5.3-Flash-NVFP4-4p67"
 required_nodes = 2
 tensor_parallel_size = 2