"""Original sanitize-time projection fusions; tiny synthetic checkpoint packs. Run under test-supervisor --command. No model downloads or full model loads. """ import os import sys print("weight-fusion reference: importing", file=sys.stderr, flush=True) import hashlib import json from types import SimpleNamespace import mlx.core as mx import mlx.nn as nn import numpy as np from mtplx.models import qwen4_exp as qwen def checkpoint_metadata(root): """Header/lazy-graph audit of the installed checkpoint, never mx.eval.""" import glob from pathlib import Path from contextlib import redirect_stdout config = json.loads((Path(root) / "config.json").read_text())["text_config"] os.environ.update(MTPLX_FUSED_GATE_UP="1", MTPLX_FUSED_GDN_INPROJ="1", MTPLX_FUSED_QSA_QKV="1") weights = {} for path in glob.glob(str(Path(root) / "model*.safetensors")): print(f"weight-fusion headers: {Path(path).name}", file=sys.stderr, flush=True) weights.update(mx.load(path)) before = len(weights) layers = [] for _ in range(config["num_hidden_layers"]): attention = nn.Module() attention.indexer = nn.Module() if config["indexer_n_heads"] else None layers.append(SimpleNamespace(linear_attn=nn.Module(), self_attn=attention, mlp=SimpleNamespace(switch_mlp=SimpleNamespace(down_proj=None), shared_expert=SimpleNamespace(down_proj=None)))) model = SimpleNamespace(layers=layers) with redirect_stdout(sys.stderr): weights = qwen._fuse_gate_up_sanitize(model, weights) weights = qwen._fuse_gdn_in_proj_sanitize(model, weights) weights = qwen._fuse_qsa_qkv_sanitize(model, weights) rows = [[name, list(a.shape), str(a.dtype).removeprefix("mlx.core.")] for name, a in sorted(weights.items())] fusions = {} for i, layer in enumerate(layers): base = f"language_model.model.layers.{i}" for suffix, module in (("mlp.switch_mlp", layer.mlp.switch_mlp), ("mlp.shared_expert", layer.mlp.shared_expert), ("linear_attn.in_proj_fused", getattr(layer.linear_attn, "in_proj_fused", None)), ("self_attn.qkv_fused", getattr(layer.self_attn, "qkv_fused", None))): if module is not None and hasattr(module, "bits"): fusions[f"{base}.{suffix}"] = dict(bits=module.bits, group=module.group_size, splits=getattr(module, "_splits", [])) print(json.dumps(dict(before=before, after=len(weights), fusions=fusions, metadata_sha256=hashlib.sha256(json.dumps(rows, separators=(",", ":")).encode()).hexdigest()), separators=(",", ":")), flush=True) def main(): assert mx.__version__ == "0.32.2" profiles = [(mask, "normal") for mask in range(8)] + [ (7, name) for name in ("qsa_bias", "indexer_mixed", "qkv_mixed", "gdn_mixed", "gdn_missing", "indexer_missing", "no_indexer", "gate_missing_bias", "shared_only", "gdn_2bit")] for mask, profile in profiles: print(f"weight-fusion reference: {mask}/{profile}", file=sys.stderr, flush=True) os.environ.update(MTPLX_FUSED_GATE_UP=str(mask & 1), MTPLX_FUSED_GDN_INPROJ=str((mask >> 1) & 1), MTPLX_FUSED_QSA_QKV=str((mask >> 2) & 1)) weights = {} specs = [("language_model.model.layers.0.mlp.switch_mlp.gate_proj", [2, 4], 320), ("language_model.model.layers.0.mlp.switch_mlp.up_proj", [2, 4], 320), ("language_model.model.layers.0.mlp.shared_expert.gate_proj", [4], 320), ("language_model.model.layers.0.mlp.shared_expert.up_proj", [4], 320)] specs += [(f"language_model.model.layers.0.linear_attn.in_proj_{sub}", [n], 160 if profile == "gdn_2bit" else (640 if profile == "gdn_mixed" and sub == "a" else 320)) for sub, n in (("qkv", 5), ("z", 3), ("b", 1), ("a", 1))] specs += [(f"language_model.model.layers.1.self_attn.{sub}", [n], 640 if (profile == "indexer_mixed" and sub.startswith("indexer")) or (profile == "qkv_mixed" and sub == "k_proj") else 320) for sub, n in (("q_proj", 5), ("k_proj", 2), ("v_proj", 2), ("indexer.index_qk_proj", 3))] for index, (base, prefix, columns) in enumerate(specs): for part, width, dtype in (("weight", columns, mx.uint32), ("scales", 40, mx.bfloat16), ("biases", 40, mx.bfloat16)): shape = (*prefix, width) data = np.arange(np.prod(shape), dtype=np.int32) % 17 + index weights[f"{base}.{part}"] = mx.array(data, dtype=dtype).reshape(shape) remove = { "gdn_missing": ["language_model.model.layers.0.linear_attn.in_proj_a.biases"], "gate_missing_bias": ["language_model.model.layers.0.mlp.switch_mlp.up_proj.biases"], "shared_only": ["language_model.model.layers.0.mlp.switch_mlp.gate_proj.weight"], }.get(profile, []) if profile == "indexer_missing": remove += [f"language_model.model.layers.1.self_attn.indexer.index_qk_proj.{p}" for p in ("weight", "scales", "biases")] for name in remove: del weights[name] if profile == "qsa_bias": weights["language_model.model.layers.1.self_attn.q_proj.bias"] = mx.zeros((5,), mx.bfloat16) # Only the attributes read/replaced by the original sanitize routines. mlp = SimpleNamespace(switch_mlp=SimpleNamespace(down_proj=None), shared_expert=SimpleNamespace(down_proj=None)) gdn, attention = nn.Module(), nn.Module() attention.indexer = None if profile == "no_indexer" else nn.Module() model = SimpleNamespace(layers=[SimpleNamespace(mlp=mlp, linear_attn=gdn), SimpleNamespace(self_attn=attention)]) # Original routines print human progress on stdout; preserve JSONL there. from contextlib import redirect_stdout with redirect_stdout(sys.stderr): weights = qwen._fuse_gate_up_sanitize(model, weights) weights = qwen._fuse_gdn_in_proj_sanitize(model, weights) weights = qwen._fuse_qsa_qkv_sanitize(model, weights) metadata = {} for name, module in (("language_model.model.layers.0.mlp.switch_mlp", mlp.switch_mlp), ("language_model.model.layers.0.mlp.shared_expert", mlp.shared_expert), ("language_model.model.layers.0.linear_attn.in_proj_fused", getattr(gdn, "in_proj_fused", None)), ("language_model.model.layers.1.self_attn.qkv_fused", getattr(attention, "qkv_fused", None))): if module is not None and hasattr(module, "bits"): metadata[name] = dict(bits=module.bits, group=module.group_size, splits=getattr(module, "_splits", [])) ordered = sorted(weights.items()) mx.eval(*[a for _, a in ordered]) print(json.dumps(dict(mask=mask, profile=profile, fusions=metadata, tensors=[ dict(name=name, shape=a.shape, sha256=hashlib.sha256(np.asarray(a.astype(mx.float32)).astype('