Save inference parity implementation and evaluation harness
This commit is contained in:
Georg Bauer
@@ -0,0 +1,123 @@
|
||||
"""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('<f4').tobytes()).hexdigest())
|
||||
for name, a in ordered]), separators=(",", ":")), flush=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
if len(sys.argv) == 3 and sys.argv[1] == "--metadata-only":
|
||||
checkpoint_metadata(sys.argv[2])
|
||||
else:
|
||||
assert len(sys.argv) == 1
|
||||
main()
|
||||
Reference in New Issue
Block a user