RegalInvader's picture
Release: Marco-Nano-Instruct 1-, 2- and 3-bit GGUF for phones
e5235fb verified
Raw History Blame Contribute Delete
13.9 kB
#!/usr/bin/env python3
"""T2.7 · GGUF export check: packed ternary students → llama.cpp Q2_0 (D-51: N-bit students → Q2_0 / Q4_0 /
Q3_K), round trip and greedy agreement.
# template: llama.cpp's own conversion of the base model (metadata, tensor names and order)
python scripts/t27_gguf.py build --template T.gguf --out S.gguf --packed <packed_dir> # a student
python scripts/t27_gguf.py build --template T.gguf --out S.gguf --rtn-dense <hf_dir> # RTN Qwen3-0.6B
python scripts/t27_gguf.py compare --gguf S.gguf (--packed <dir> | --rtn-dense <hf_dir>) --out R.json
`compare` runs both models greedy on the same prompt token ids (200 held-out prompts, chat template, ≤512 prompt
tokens, 64 new tokens, end-of-sequence suppressed in both) and reports position-wise token agreement, the first
divergence per prompt, and llama.cpp's own timings. PyTorch is the reference: the packed student dequantised to
bf16 (`load_packed_student`), on CPU. llama.cpp runs CPU-only (`llama-server`, GPU hidden: gen's GPU is busy).
"""
from __future__ import annotations
import argparse
import json
import os
import random
import subprocess
import sys
import time
import urllib.request
from pathlib import Path
import torch
REPO = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(REPO / "src"))
from lowbit.export import gguf as gx # noqa: E402
from lowbit.quant.ternary import QuantSpec, dequant_packed, export_packed # noqa: E402
LLAMA = Path(os.environ.get("MT_LLAMA", "/path/to/llama.cpp")) # gguf-py, the converter (spark: third_party)
# the pegfix build: mainline's chat-output parser returns HTTP 500 on some outputs (invalid UTF-8 from a degraded
# model), which the fork's UTF-8 fix removes; same source base, so the same Q2_0 kernels
SERVER_LIBS = "/path/to/llama.cpp-fork/build-pegfix/bin"
SERVER = Path(SERVER_LIBS) / "llama-server"
HELDOUT = Path("/path/to/data/prompts/v2/heldout.jsonl")
def rtn_dense(hf_dir: str):
"""A dense model with every MLP projection ternarised by round-to-nearest (Q-A, g128): (dense, packed, model)."""
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(hf_dir, dtype=torch.bfloat16)
dense, packed = {}, {}
for n, p in model.named_parameters():
if any(n.endswith(f"mlp.{k}_proj.weight") for k in ("gate", "up", "down")):
pk = export_packed(p.detach().float(), QuantSpec(method="absmean", group_size=128))
packed[n[: -len(".weight")]] = pk
with torch.no_grad():
p.copy_(dequant_packed(pk)) # the PyTorch reference holds exactly the exported weights
else:
dense[n] = p.detach()
return dense, packed, model
def packed_student(packed_dir: str):
"""(plain tensors, packed experts); D-51's N-bit students carry their method per tensor (`read_packed`)."""
from lowbit.model.student import read_packed
sd, packed, _ = read_packed(packed_dir)
return sd, packed
def build(a) -> None:
g = gx.gguf_module(LLAMA)
if a.packed:
dense, packed = packed_student(a.packed)
else:
dense, packed, _ = rtn_dense(a.rtn_dense)
names = {t.name for t in g.GGUFReader(a.template).tensors}
if "output.weight" in names and "lm_head.weight" not in dense: # tied: the converter still writes output
dense["lm_head.weight"] = dense["model.embed_tokens.weight"]
dense = {k: v for k, v in dense.items() if gx.gguf_name(k) in names} # a tied lm_head absent from the GGUF
t0 = time.time()
meta = {}
if a.name: # a release: its own name, and the teacher recorded as its base model
meta = {"general.name": a.name, "general.finetune": "QAT", "general.base_model.count": 1,
"general.base_model.0.name": "Marco Nano Instruct", "general.base_model.0.organization": "AIDC-AI",
"general.base_model.0.repo_url": "https://huggingface.co/AIDC-AI/Marco-Nano-Instruct"}
info = gx.write_gguf(a.template, a.out, gx.student_tensors(dense, packed, g, a.ternary_as, a.dense_as), g, meta)
info["write_s"] = round(time.time() - t0, 1)
if a.ternary_as == "q2_0":
info.update(gx.check_round_trip(a.out, packed, g))
print(json.dumps(info))
def prompts(tok, n: int, max_len: int = 512) -> list[list[int]]:
rows = [json.loads(l) for l in open(HELDOUT)]
random.Random(0).shuffle(rows)
out = []
for r in rows:
text = tok.apply_chat_template([{"role": "user", "content": r["text"]}], add_generation_prompt=True,
tokenize=False)
ids = tok(text, add_special_tokens=False)["input_ids"]
if len(ids) <= max_len:
out.append(list(ids))
if len(out) == n:
return out
return out
@torch.no_grad()
def torch_greedy(model, tok, ps: list[list[int]], new: int, batch: int) -> list[list[int]]:
tok.padding_side = "left"
pad = tok.pad_token_id if tok.pad_token_id is not None else tok.eos_token_id
out = []
for i in range(0, len(ps), batch):
chunk = ps[i:i + batch]
L = max(len(p) for p in chunk)
ids = torch.tensor([[pad] * (L - len(p)) + p for p in chunk])
att = torch.tensor([[0] * (L - len(p)) + [1] * len(p) for p in chunk])
g = model.generate(input_ids=ids, attention_mask=att, max_new_tokens=new, min_new_tokens=new,
do_sample=False, pad_token_id=pad)
out += [row[L:].tolist() for row in g]
print(json.dumps({"torch_greedy": i + len(chunk), "of": len(ps)}), flush=True)
return out
def start_server(gguf: str, threads: int, port: int) -> subprocess.Popen:
env = {**os.environ, "CUDA_VISIBLE_DEVICES": "", "LD_LIBRARY_PATH": SERVER_LIBS}
# --reasoning-format none: plain completions, no think-block parsing (a parser hang stalled one request)
proc = subprocess.Popen([str(SERVER), "-m", gguf, "--port", str(port), "--host", "127.0.0.1", "-ngl", "0",
"-t", str(threads), "-c", "8192", "--parallel", "1", "--no-webui",
"--reasoning-format", "none"],
env=env, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
for _ in range(600):
try:
urllib.request.urlopen(f"http://127.0.0.1:{port}/health", timeout=2)
return proc
except Exception:
if proc.poll() is not None:
raise RuntimeError(f"llama-server exited with {proc.returncode}")
time.sleep(1)
raise RuntimeError("llama-server did not come up")
def stop_server(proc: subprocess.Popen) -> None:
proc.terminate() # it can sit on SIGTERM mid-request; a slow shutdown must not discard results (it did once)
try:
proc.wait(30)
except subprocess.TimeoutExpired:
proc.kill()
proc.wait()
def server_greedy(gguf: str, ps: list[list[int]], new: int, threads: int, port: int,
timeout_s: int = 120) -> tuple[list, dict]:
"""Greedy tokens per prompt; a request over `timeout_s` is recorded as None and the server restarted."""
proc = start_server(gguf, threads, port)
outs, pp, tg, timeouts = [], [], [], 0
try:
for p in ps:
body = {"prompt": p, "n_predict": new, "temperature": 0.0, "top_k": 1, "ignore_eos": True,
"cache_prompt": False, "return_tokens": True, "samplers": ["top_k"]}
req = urllib.request.Request(f"http://127.0.0.1:{port}/completion", json.dumps(body).encode(),
{"Content-Type": "application/json"})
try:
r = json.loads(urllib.request.urlopen(req, timeout=timeout_s).read())
outs.append(r["tokens"])
pp.append(r["timings"]["prompt_per_second"])
tg.append(r["timings"]["predicted_per_second"])
except (TimeoutError, OSError) as ex:
outs.append(None)
timeouts += 1
print(json.dumps({"llama_timeout": len(outs) - 1, "error": repr(ex)[:120]}), flush=True)
stop_server(proc)
proc = start_server(gguf, threads, port)
if len(outs) % 20 == 0:
print(json.dumps({"llama_greedy": len(outs), "of": len(ps), "gen_tok_s": round(tg[-1], 1) if tg else None,
"prompt_tok_s": round(pp[-1], 1) if pp else None}), flush=True)
return outs, {"prompt_tok_s_median": sorted(pp)[len(pp) // 2], "gen_tok_s_median": sorted(tg)[len(tg) // 2],
"threads": threads, "timeouts": timeouts}
finally:
stop_server(proc)
def compare(a) -> None:
from transformers import AutoTokenizer
from lowbit.model import student as st
torch.set_num_threads(a.threads)
if a.packed:
model = st.load_packed_student(a.packed, attn_impl="sdpa", experts_impl="grouped_mm")
tok = AutoTokenizer.from_pretrained(a.tokenizer)
else:
_, _, model = rtn_dense(a.rtn_dense)
tok = AutoTokenizer.from_pretrained(a.rtn_dense)
model.eval()
ps = prompts(tok, a.n)
t0 = time.time()
ref = torch_greedy(model, tok, ps, a.new, a.batch)
torch_s = time.time() - t0
del model
got, timing = server_greedy(a.gguf, ps, a.new, a.threads, a.port)
pairs = [(r, g) for r, g in zip(ref, got) if g is not None] # timed-out requests are reported, not scored
agree = sum(x == y for r, g in pairs for x, y in zip(r, g))
first = [next((i for i, (x, y) in enumerate(zip(r, g)) if x != y), a.new) for r, g in pairs]
res = {"prompts": len(pairs), "new_tokens": a.new, "token_agreement": agree / (len(pairs) * a.new),
"identical_sequences": sum(f == a.new for f in first), "first_divergence_median": sorted(first)[len(first) // 2],
"first_divergence_mean": sum(first) / len(first), "torch_cpu_s": round(torch_s), **timing,
"gguf_bytes": Path(a.gguf).stat().st_size}
Path(a.out).write_text(json.dumps({**res, "torch": ref, "llama": got}))
print(json.dumps(res))
def forced(a) -> None:
"""For every position of PyTorch's greedy continuation, does llama.cpp pick the same next token given the same
context? No cascade after a first difference; llama.cpp's prompt cache makes each step one new token."""
from transformers import AutoTokenizer
ref = json.loads(Path(a.compare).read_text())
ps = prompts(AutoTokenizer.from_pretrained(a.tokenizer), ref["prompts"])
proc = start_server(a.gguf, a.threads, a.port)
hits = n = 0
per_pos = [0] * ref["new_tokens"]
try:
for k, (p, cont) in enumerate(zip(ps, ref["torch"])):
for i, want in enumerate(cont):
body = {"prompt": p + cont[:i], "n_predict": 1, "temperature": 0.0, "top_k": 1, "ignore_eos": True,
"cache_prompt": True, "return_tokens": True, "samplers": ["top_k"]}
req = urllib.request.Request(f"http://127.0.0.1:{a.port}/completion", json.dumps(body).encode(),
{"Content-Type": "application/json"})
got = json.loads(urllib.request.urlopen(req, timeout=120).read())["tokens"][0]
hits += got == want
per_pos[i] += got == want
n += 1
if (k + 1) % 25 == 0:
print(json.dumps({"forced": k + 1, "of": len(ps), "agreement": round(hits / n, 4)}), flush=True)
finally:
stop_server(proc)
res = {"gguf": a.gguf, "positions": n, "forced_top1_agreement": hits / n,
"agreement_first_16": sum(per_pos[:16]) / (16 * len(ps)), "agreement_last_16": sum(per_pos[-16:]) / (16 * len(ps))}
Path(a.out).write_text(json.dumps(res, indent=1))
print(json.dumps(res))
def main() -> None:
p = argparse.ArgumentParser()
sub = p.add_subparsers(dest="op", required=True)
b = sub.add_parser("build")
b.add_argument("--template", required=True)
b.add_argument("--out", required=True)
src = b.add_mutually_exclusive_group(required=True)
src.add_argument("--packed")
src.add_argument("--rtn-dense")
b.add_argument("--ternary-as", choices=["q2_0", "bf16"], default="q2_0",
help="q2_0: each packed tensor in its native type (ternary/int2 Q2_0, int4 Q4_0, int3 Q3_K + Q4_0 "
"down; D-51); bf16: the lossless control")
b.add_argument("--name", help="a release name for general.name (and base-model metadata)")
b.add_argument("--dense-as", choices=["bf16", "q8_0"], default="bf16",
help="attention and embedding: bf16 as trained, or q8_0 (post hoc, as the Q8 teacher; faster on CPU)")
tf = sub.add_parser("forced", help="teacher-forced next-token agreement against a compare run's PyTorch tokens")
tf.add_argument("--gguf", required=True)
tf.add_argument("--compare", required=True, help="the compare run's JSON (its PyTorch greedy tokens)")
tf.add_argument("--tokenizer", required=True)
tf.add_argument("--out", required=True)
tf.add_argument("--threads", type=int, default=16)
tf.add_argument("--port", type=int, default=18190)
c = sub.add_parser("compare")
c.add_argument("--gguf", required=True)
c.add_argument("--out", required=True)
src = c.add_mutually_exclusive_group(required=True)
src.add_argument("--packed")
src.add_argument("--rtn-dense")
c.add_argument("--tokenizer", default="/path/to/data/models/Marco-Nano-Instruct")
c.add_argument("--n", type=int, default=200)
c.add_argument("--new", type=int, default=64)
c.add_argument("--batch", type=int, default=25)
c.add_argument("--threads", type=int, default=16)
c.add_argument("--port", type=int, default=18190)
a = p.parse_args()
{"build": build, "compare": compare, "forced": forced}[a.op](a)
if __name__ == "__main__":
main()