124 lines
7.1 KiB
Python
124 lines
7.1 KiB
Python
"""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()
|