File size: 3,008 Bytes
4335e83 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 | """oQ5e (imatrix-enhanced) quant of the NEW merged Nemotron, KEEPING the MTP head.
Run with the omlx tool venv python (has omlx + mlx-lm):
/Users/david/.local/share/uv/tools/omlx/bin/python quant_oq5e_mtp.py
Env overrides: SRC, OUT, OQ_LEVEL (default 5), OUT_SUFFIX (default oq5e).
The imatrix cache is shared across levels (keyed on source signature), so an
oQ4e follow-up pass reuses the calibration for free:
OQ_LEVEL=4 OUT_SUFFIX=oq4e /Users/david/.local/share/uv/tools/omlx/bin/python quant_oq5e_mtp.py
See KENOSISTRON2_QUANT_PLAN.md for the full runbook.
"""
import os
import sys
sys.path.insert(0, "/Users/david/AI")
import nemotron_h_mtp_model as P # standalone copy of the serving patch
from mlx_lm.models import nemotron_h as nh
P.set_mtp_active(True)
ok = P.apply()
print(
"patch applied:", ok,
"| Model.sanitize is ours:",
getattr(nh.Model.__dict__.get("sanitize"), "_omlx_nh_mtp", False),
)
assert ok, "nemotron_h MTP patch failed to apply — mtp.* would be stripped"
from omlx.oq import quantize_oq_streaming
LEVEL = float(os.environ.get("OQ_LEVEL", "5"))
LEVEL = int(LEVEL) if LEVEL == int(LEVEL) else LEVEL
SUFFIX = os.environ.get("OUT_SUFFIX", "oq5e")
SRC = os.environ.get("SRC", "/Users/david/AI/NVIDIA-Nemotron-3-Super-120B-merged2")
OUT = os.environ.get("OUT", f"/Users/david/AI/kenosistron2-{SUFFIX}-mtp")
IMATRIX = os.environ.get("IMATRIX", "/Users/david/AI/kenosistron2_imatrix.oqe")
def cb(phase, pct):
print(f"[{SUFFIX}-mtp] {phase}: {pct:.1f}%", flush=True)
quantize_oq_streaming(
model_path=SRC,
output_path=OUT,
oq_level=LEVEL,
group_size=64,
preserve_mtp=True,
dtype="bfloat16",
enhanced=True, # oQe: imatrix-weighted quantization (v0.5.0)
imatrix_cache_path=IMATRIX, # reused across levels if source unchanged
imatrix_reuse_cache=True,
progress_callback=cb,
)
print("done. ->", OUT)
# --- serving fix: pin the tool parser ---------------------------------------
# The Nemotron-3-Super chat_template contains `<tool_call>` + `tool_call.name`
# but NOT the adjacent `<tool_call>\n<function=` marker, so mlx_lm auto-infers
# the json_tools parser — while the model actually emits the qwen3_coder
# `<function=name><parameter=k>v</parameter></function>` format. Mismatch => tool
# calls never parse (leak as text / empty). Pin the correct parser explicitly;
# tokenizer_config `tool_parser_type` overrides inference (mlx_lm tokenizer_utils).
import json as _json
_tc_path = os.path.join(OUT, "tokenizer_config.json")
try:
_tc = _json.load(open(_tc_path))
if _tc.get("tool_parser_type") != "qwen3_coder":
_tc["tool_parser_type"] = "qwen3_coder"
_json.dump(_tc, open(_tc_path, "w"), indent=2, ensure_ascii=False)
print("pinned tool_parser_type=qwen3_coder in", _tc_path)
else:
print("tool_parser_type already qwen3_coder")
except FileNotFoundError:
print("WARN: no tokenizer_config.json in output — set tool_parser_type manually")
|