fcmeyer's picture
Publish validated MLX 6bit conversion and measurements
e5150a2 verified
Raw History Blame Contribute Delete
16.8 kB
"""Reproducible conversion, release gates, benchmarking, and publication."""
import argparse
import fcntl
import hashlib
import json
import math
import shutil
import statistics
import struct
from pathlib import Path
ROOT = Path(__file__).resolve().parents[1]
REV = "c1093c501cc6460427c0ac634db7223d0450c0ee"
SOURCE = "Nanbeige/Nanbeige4.3-33B"
VARIANTS = {"bf16": None, "4bit": (4, 64), "6bit": (6, 64), "8bit": (8, 64),
"4bit-g32": (4, 32), "4bit-g128": (4, 128), "mixed-4-6-g64": (4, 64)}
def write_json(path, value):
path.parent.mkdir(parents=True, exist_ok=True)
tmp = path.with_suffix(".tmp")
tmp.write_text(json.dumps(value, indent=2, ensure_ascii=False) + "\n")
tmp.replace(path)
def inventory(path):
tensors = {}
for f in sorted(path.glob("*.safetensors")):
with f.open("rb") as stream:
length = struct.unpack("<Q", stream.read(8))[0]
header = json.loads(stream.read(length))
for key, value in header.items():
if key == "__metadata__":
continue
assert key not in tensors, key
tensors[key] = {k: value[k] for k in ("dtype", "shape")}
assert tensors, "No weights"
index = json.loads((path / "model.safetensors.index.json").read_text())
assert set(index["weight_map"]) == set(tensors)
return tensors
def convert(name):
from mlx_lm.convert import convert as mlx_convert
path = ROOT / "artifacts" / name
source = ROOT / ("source" if name == "bf16" else "artifacts/bf16")
options = dict(hf_path=str(source), mlx_path=str(path), dtype="bfloat16")
if VARIANTS[name]:
bits, group = VARIANTS[name]
options.update(quantize=True, q_bits=bits, q_group_size=group, q_mode="affine")
if name.startswith("mixed"):
options["quant_predicate"] = "mixed_4_6"
mlx_convert(**options)
# Preserve upstream tokenizer assets byte-for-byte, including the template.
for pattern in ("tokenizer*", "special_tokens_map.json", "added_tokens.json", "generation_config.json"):
for file in (ROOT / "source").glob(pattern):
if file.is_file():
shutil.copy2(file, path / file.name)
write_json(path / "conversion.json", {"source": SOURCE, "revision": REV,
"variant": name, "options": options})
def validate(name, model_path=None):
import mlx.core as mx
import mlx.nn as nn
from mlx_lm import load, stream_generate
from mlx_lm.sample_utils import make_sampler
from transformers import AutoTokenizer
path = Path(model_path) if model_path else ROOT / "artifacts" / name
weights = inventory(path)
config = json.loads((path / "config.json").read_text())
original = json.loads((ROOT / "source/config.json").read_text())
for key in ("rope_theta", "max_position_embeddings", "vocab_size", "num_hidden_layers", "num_attention_heads", "num_key_value_heads", "hidden_size"):
assert config[key] == original[key], key
if name == "bf16":
assert weights == inventory(ROOT / "source"), "BF16 tensor inventory changed"
assert all(v["dtype"] == "BF16" for v in weights.values())
else:
assert config.get("quantization"), "Missing quantization metadata"
assert all(v["dtype"] in ("BF16", "U32") for v in weights.values())
source_tok = AutoTokenizer.from_pretrained(ROOT / "source")
model, tok = load(str(path))
for text in ("Hello, world!", "请解释量化。", "def add(a, b):\n return a+b", "<think></think>"):
assert source_tok.encode(text) == tok.encode(text)
tools = [{"type": "function", "function": {"name": "weather", "description": "Get weather", "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}}}]
messages = [{"role": "user", "content": "Hello"}, {"role": "assistant", "content": "<think>Consider a greeting.</think>Hello!"}, {"role": "user", "content": "What is the weather in Boston?"}]
for thinking in (True, False):
for preserve in (True, False):
for fmt in ("xml", "json"):
kw = dict(tokenize=False, add_generation_prompt=True, enable_thinking=thinking, preserve_thinking=preserve, tools=tools, tool_call_format=fmt)
assert source_tok.apply_chat_template(messages, **kw) == tok.apply_chat_template(messages, **kw)
smoke = []
for prompt in ("In one sentence, explain why the sky looks blue.", "请用一句中文解释为什么天空是蓝色的。"):
rendered = tok.apply_chat_template([{"role": "user", "content": prompt}], tokenize=False, add_generation_prompt=True, enable_thinking=False)
logits = model(mx.array([tok.encode(rendered)]))
assert bool(mx.all(mx.isfinite(logits)).item()), "Non-finite logits"
del logits
text = "".join(r.text for r in stream_generate(model, tok, rendered, max_tokens=128, sampler=make_sampler(temp=0)))
assert text.strip(), "Empty generation"
smoke.append({"prompt": prompt, "response": text})
print(text, flush=True)
corpus_file = ROOT / "evaluation.json"
corpus = json.loads(corpus_file.read_text())
losses = []
for item in corpus:
ids = tok.encode(item["text"])
x = mx.array([ids])
logits = model(x[:, :-1]).astype(mx.float32)
loss = float(nn.losses.cross_entropy(logits, x[:, 1:], reduction="mean").item())
assert math.isfinite(loss)
losses.append({"category": item["category"], "tokens": len(ids)-1, "nll": loss})
del logits, x
mx.clear_cache()
nll = sum(x["tokens"] * x["nll"] for x in losses) / sum(x["tokens"] for x in losses)
report = {"variant": name, "tensor_count": len(weights), "smoke": smoke,
"corpus_sha256": hashlib.sha256(corpus_file.read_bytes()).hexdigest(),
"losses": losses, "mean_nll": nll, "perplexity": math.exp(nll),
"weight_bytes": sum(f.stat().st_size for f in path.glob("*.safetensors")),
"template_cases_passed": 8, "status": "passed"}
if name != "bf16":
baseline = json.loads((ROOT / "reports/bf16.validation.json").read_text())
report["nll_delta_bf16"] = nll - baseline["mean_nll"]
write_json(ROOT / "reports" / f"{name}.validation.json", report)
print(json.dumps(report, ensure_ascii=False), flush=True)
def benchmark(name):
import mlx.core as mx
from mlx_lm import load, stream_generate
from mlx_lm.sample_utils import make_sampler
model, tok = load(str(ROOT / "artifacts" / name))
tok._eos_token_ids = set()
results = []
for length in (512, 4096, 16384):
mx.random.seed(42)
prompt = mx.random.randint(0, model.args.vocab_size, (length,)).tolist()
trials = []
for trial in range(4):
mx.reset_peak_memory()
for response in stream_generate(model, tok, prompt, max_tokens=256, sampler=make_sampler(temp=0), prefill_step_size=512):
pass
assert response.generation_tokens == 256
row = {k: getattr(response, k) for k in ("prompt_tps", "generation_tps", "peak_memory")}
print(name, length, trial, row, flush=True)
if trial:
trials.append(row)
mx.clear_cache()
results.append({"prompt_tokens": length, "generation_tokens": 256, "trials": trials,
"median": {k: statistics.median(t[k] for t in trials) for k in trials[0]}})
write_json(ROOT / "reports" / f"{name}.benchmark.json", {"hardware": "Apple M5 Max, 128 GB", "kv_cache": "unquantized", "prefill_step_size": 512, "results": results})
def publish(name):
from huggingface_hub import HfApi
path = ROOT / "artifacts" / name
report = json.loads((ROOT / "reports" / f"{name}.validation.json").read_text())
assert report["status"] == "passed"
repo = f"fcmeyer/Nanbeige4.3-33B-MLX-{name}"
conversion = json.loads((path / "conversion.json").read_text())
conversion["options"]["hf_path"] = "source" if name == "bf16" else "artifacts/bf16"
conversion["options"]["mlx_path"] = f"artifacts/{name}"
write_json(path / "conversion.json", conversion)
for f in (ROOT / "reports").glob(f"{name}.*.json"):
if not f.name.endswith(".publication.json"):
shutil.copy2(f, path / f.name)
if (ROOT / "reports/environment.json").exists():
shutil.copy2(ROOT / "reports/environment.json", path / "environment.json")
for f in ("requirements.lock.txt", "evaluation.json"):
shutil.copy2(ROOT / f, path / f)
(path / "scripts").mkdir(exist_ok=True)
shutil.copy2(ROOT / "scripts/release.py", path / "scripts/release.py")
if (ROOT / "reports" / f"{name}.fidelity.json").exists():
shutil.copy2(ROOT / "scripts/fidelity.py", path / "scripts/fidelity.py")
# Earlier initial releases placed this script at the repository root.
# Its relative project paths require the scripts/ layout for reproduction.
if (path / "release.py").exists():
(path / "release.py").unlink()
if (ROOT / "LICENSE").exists():
shutil.copy2(ROOT / "LICENSE", path / "LICENSE")
quant = "Unquantized BF16." if name == "bf16" else f"Affine quantization: {VARIANTS[name][0]} base bits, group size {VARIANTS[name][1]}." + (" MLX-LM mixed_4_6 recipe." if name.startswith("mixed") else "")
performance = "Performance benchmarking is pending."
benchmark_file = ROOT / "reports" / f"{name}.benchmark.json"
if benchmark_file.exists():
benchmark_report = json.loads(benchmark_file.read_text())
performance = "Measured on Apple M5 Max with 128 GB unified memory, unquantized KV cache, 512-token prefill steps, 256 output tokens; one warm-up and three measured trials per length. Prompts are seeded synthetic token sequences. Values are medians.\n\n| Input tokens | Prefill tok/s | Decode tok/s | Peak GB |\n|---:|---:|---:|---:|\n"
for row in benchmark_report["results"]:
med = row["median"]
performance += f"| {row['prompt_tokens']} | {med['prompt_tps']:.2f} | {med['generation_tps']:.2f} | {med['peak_memory']:.2f} |\n"
performance += "\nPeak memory is MLX's reported peak allocation, excluding other applications and OS memory. See `environment.json` for OS/runtime details.\n"
delta = f"NLL difference versus BF16: {report['nll_delta_bf16']:+.6f}." if "nll_delta_bf16" in report else ""
fidelity_text = ""
fidelity_path = ROOT / "reports" / f"{name}.fidelity.json"
if fidelity_path.exists():
fidelity = json.loads(fidelity_path.read_text())
fidelity_text = f"Teacher-forced fidelity versus BF16 on the same {fidelity['scored_tokens']} tokens: mean KL(BF16 || quant) {fidelity['mean_kl_bf16_to_quant_nats']:.6f} nats/token; top-1 token agreement {100*fidelity['top1_agreement']:.2f}%. Lower KL means closer output distributions on these inputs. This does not establish downstream capability equivalence."
card = f'''---
license: apache-2.0
library_name: mlx
pipeline_tag: text-generation
base_model: {SOURCE}
language:
- en
- zh
tags:
- mlx
- nanbeige
---
# Nanbeige4.3-33B MLX {name}
MLX conversion of [{SOURCE}](https://huggingface.co/{SOURCE}), by the Nanbeige Team.
Source revision: `{REV}`. {quant}
## Usage
```python
from mlx_lm import load, generate
model, tokenizer = load("{repo}")
prompt = tokenizer.apply_chat_template(
[{{"role": "user", "content": "Hello!"}}],
tokenize=False, add_generation_prompt=True, enable_thinking=False)
print(generate(model, tokenizer, prompt=prompt, max_tokens=256))
```
The upstream chat template is preserved. `enable_thinking` controls current reasoning;
`preserve_thinking=True` retains previous reasoning for agent workflows. Upstream recommends
`tool_call_format="xml"`; JSON is also supported. Template support does not certify any server's tool parser.
## Validation
Weight files: {report['weight_bytes']:,} bytes. Tensor count: {report['tensor_count']}.
English/Chinese generation, finite smoke logits, tokenizer equivalence, and eight thinking/tool template cases passed.
Small fixed synthetic diagnostic corpus ({sum(x['tokens'] for x in report['losses'])} scored tokens): mean token NLL {report['mean_nll']:.6f}, perplexity {report['perplexity']:.6f}.
{delta}
{fidelity_text}
This is a conversion regression check, not a standardized capability benchmark or a claim of benchmark equivalence.
See the included validation JSON for outputs and per-category measurements.
The source's 524,288-token context setting is retained; full 512K context has not been validated.
Performance results, when included, apply only to the stated hardware and tested lengths.
## Performance
{performance}
## Reproduction
Use the included `requirements.lock.txt`, `conversion.json`, `evaluation.json`, and `scripts/release.py`.
From the repository root:
```sh
uv venv --python 3.12 .venv
uv pip install --python .venv/bin/python -r requirements.lock.txt
.venv/bin/hf download {SOURCE} --revision {REV} --local-dir source
.venv/bin/python scripts/release.py convert bf16
.venv/bin/python scripts/release.py validate bf16
```
For quantized variants, then run `.venv/bin/python scripts/release.py convert {name}` and
`.venv/bin/python scripts/release.py validate {name}`. All quantized variants derive independently from BF16.
## Attribution
Original model: Nanbeige Team, Nanbeige4.3-33B (2026), Apache-2.0.
Conversion and measurements: fcmeyer. No claim is made that this is universally the best quantization.
'''
(path / "README.md").write_text(card)
api = HfApi()
api.create_repo(repo_id=repo, repo_type="model", private=False, exist_ok=True)
api.upload_folder(repo_id=repo, repo_type="model", folder_path=path,
ignore_patterns=[".cache/**"],
delete_patterns=["release.py"],
commit_message=f"Publish validated MLX {name} conversion and measurements")
info = api.model_info(repo, files_metadata=True)
remote = {f.rfilename: f.size for f in info.siblings}
for f in path.iterdir():
if f.is_file():
assert remote.get(f.name) == f.stat().st_size, f.name
write_json(ROOT / "reports" / f"{name}.publication.json", {"repo": repo, "revision": info.sha, "inventory_verified": True})
print(f"PUBLISHED https://huggingface.co/{repo} revision={info.sha}", flush=True)
def verify_remote(name):
"""Reload published config/tokenizer with locally hash-verified remote weights."""
from huggingface_hub import HfApi, snapshot_download
from mlx_lm import load, stream_generate
from mlx_lm.sample_utils import make_sampler
report_path = ROOT / "reports" / f"{name}.publication.json"
report = json.loads(report_path.read_text())
repo, revision = report["repo"], report["revision"]
path = ROOT / "artifacts" / (name + "-remote-check")
snapshot_download(repo, revision=revision, local_dir=path, ignore_patterns=["*.safetensors"])
info = HfApi().model_info(repo, revision=revision, files_metadata=True)
for file in info.siblings:
if not file.rfilename.endswith(".safetensors"):
continue
local = ROOT / "artifacts" / name / file.rfilename
with local.open("rb") as stream:
digest = hashlib.file_digest(stream, "sha256").hexdigest()
assert file.lfs and digest == file.lfs.sha256, file.rfilename
target = path / file.rfilename
if not target.exists():
target.symlink_to(local)
model, tokenizer = load(str(path))
prompt = tokenizer.apply_chat_template([{"role": "user", "content": "Say hello."}], tokenize=False, add_generation_prompt=True, enable_thinking=False)
text = "".join(r.text for r in stream_generate(model, tokenizer, prompt, max_tokens=32, sampler=make_sampler(temp=0)))
assert text.strip()
report.update(remote_reload="passed", weight_sha256_verified=True, reload_response=text,
verification_method="Fresh published metadata/tokenizer; local weights SHA256 matched to remote LFS hashes, avoiding redundant transfer")
write_json(report_path, report)
print(json.dumps(report), flush=True)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("action", choices=["convert", "validate", "benchmark", "publish", "verify_remote"])
parser.add_argument("variant", choices=VARIANTS)
args = parser.parse_args()
if args.action == "publish":
publish(args.variant)
else:
(ROOT / "artifacts").mkdir(exist_ok=True)
with (ROOT / "artifacts/.gpu.lock").open("a") as gpu_lock:
fcntl.flock(gpu_lock, fcntl.LOCK_EX)
globals()[args.action](args.variant)