| """Build the files a faster backend loads, from a model directory (only read) into a new directory. |
| |
| python -m decisionnode_edge.build openvino path/to/model --out path/to/model-openvino [--lengths 512,2048] |
| python -m decisionnode_edge.build gguf path/to/model --out path/to/model-gguf --llama-cpp path/to/llama.cpp |
| |
| openvino The backbone's text path (the token embedding, the layers, the final norm: input_ids [1, L] to the hidden |
| states [1, L, hidden]) exported with torch.onnx's dynamo exporter in float32 at each static length, then |
| saved as OpenVINO IR with float32 weights (the bf16 weights, exactly); the runtime runs it in bf16. One IR per |
| length, since the export unrolls the delta rule over the length: each takes 15 to 20 minutes, and lengths |
| above a few thousand tokens take hours. Then each IR is run once against torch. Needs the `openvino` extra |
| (openvino, onnx, onnxscript); `--device cuda` (the default when a GPU is present) keeps the weights on the card. |
| gguf llama.cpp's converter (convert_hf_to_gguf.py --no-mtp) on links to the backbone files, Q8_0 by default: 8-bit |
| keeps the answers, 4-bit does not. Needs a llama.cpp checkout and its converter's Python requirements, in the |
| environment of `--python` (default: this one). |
| |
| Both write decisionnode_edge_build.json (the backend, the files, the model's fingerprint, versions, times), which |
| EdgeModel.load reads from `artifacts`. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import importlib.metadata |
| import json |
| import os |
| import shutil |
| import subprocess |
| import sys |
| import time |
|
|
| from .backends import BUILD_FILE, BUILD_FORMAT, fingerprint, link_model_files |
|
|
| DEFAULT_LENGTHS = (512, 2048) |
|
|
|
|
| def versions(*names: str) -> dict: |
| out = {} |
| for name in names: |
| try: |
| out[name] = importlib.metadata.version(name) |
| except importlib.metadata.PackageNotFoundError: |
| out[name] = None |
| return out |
|
|
|
|
| def new_directory(out: str) -> None: |
| if os.path.exists(out): |
| raise SystemExit(f"{out} exists: the build writes a new directory") |
| os.makedirs(out) |
|
|
|
|
| def write_record(out: str, record: dict) -> None: |
| with open(os.path.join(out, BUILD_FILE), "w", encoding="utf-8") as f: |
| json.dump(record, f, indent=2) |
| f.write("\n") |
|
|
|
|
| def model_name(model: str) -> str: |
| with open(os.path.join(model, "decisionnode_edge.json"), encoding="utf-8") as f: |
| return json.load(f)["name"] |
|
|
|
|
| def build_openvino(model: str, out: str, lengths: list[int], device: str | None) -> dict: |
| import numpy as np |
| import openvino as ov |
| import torch |
| from transformers import AutoModelForImageTextToText, AutoTokenizer |
|
|
| device = device or ("cuda" if torch.cuda.is_available() else "cpu") |
| record = {"format": BUILD_FORMAT, "backend": "openvino", "model": model_name(model), "model_fingerprint": fingerprint(model), "files": {}, |
| "weights": "float32 (the bf16 weights, exactly); run in bf16", "export_device": device, |
| "versions": versions("torch", "transformers", "onnx", "onnxscript", "openvino"), "seconds": {}, "check": {}} |
| tokenizer = AutoTokenizer.from_pretrained(model) |
| record["pad_token_id"] = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0 |
|
|
| class LastHidden(torch.nn.Module): |
| def __init__(self, text): |
| super().__init__() |
| self.text = text |
|
|
| def forward(self, input_ids): |
| return self.text(input_ids=input_ids, use_cache=False).last_hidden_state |
|
|
| full = AutoModelForImageTextToText.from_pretrained(model, dtype=torch.float32).to(device).eval() |
| text = LastHidden(full.model.language_model).eval() |
| generator = torch.Generator().manual_seed(0) |
| for length in lengths: |
| work = os.path.join(out, f"onnx-{length}") |
| os.makedirs(work) |
| onnx_file = os.path.join(work, "edge.onnx") |
| ids = torch.randint(10, 200000, (1, length), generator=generator).to(device) |
| t0 = time.perf_counter() |
| with torch.inference_mode(): |
| torch.onnx.export(text, (ids,), onnx_file, dynamo=True, external_data=True, input_names=["input_ids"], output_names=["last_hidden_state"], optimize=True) |
| t1 = time.perf_counter() |
| ir = ov.convert_model(onnx_file) |
| name = f"openvino-{length}.xml" |
| ov.save_model(ir, os.path.join(out, name), compress_to_fp16=False) |
| del ir |
| shutil.rmtree(work) |
| record["files"][str(length)] = name |
| record["seconds"][str(length)] = {"export": round(t1 - t0, 1), "convert": round(time.perf_counter() - t1, 1)} |
| print(json.dumps({"length": length, **record["seconds"][str(length)]}), flush=True) |
|
|
| |
| core = ov.Core() |
| for length in lengths: |
| n = length - 37 |
| ids = torch.randint(10, 200000, (1, n), generator=generator) |
| with torch.inference_mode(): |
| want = text(ids.to(device))[0, -1].float().cpu() |
| compiled = core.compile_model(os.path.join(out, record["files"][str(length)]), "CPU", {"INFERENCE_PRECISION_HINT": "bf16"}) |
| x = np.full((1, length), record["pad_token_id"], dtype=np.int64) |
| x[0, :n] = ids[0].numpy() |
| got = torch.from_numpy(np.asarray(compiled.create_infer_request().infer({0: x})[compiled.output(0)][0, n - 1], dtype=np.float32)) |
| cosine = float(torch.nn.functional.cosine_similarity(got, want, dim=0)) |
| record["check"][str(length)] = {"tokens": n, "cosine_to_torch_float32": round(cosine, 6), "largest_difference": round(float((got - want).abs().max()), 4)} |
| print(json.dumps({"check": length, **record["check"][str(length)]}), flush=True) |
| del compiled |
| if not torch.isfinite(got).all() or cosine < 0.99: |
| raise SystemExit(f"the IR at {length} tokens does not reproduce the backbone (cosine {cosine:.4f}): {out} is incomplete") |
| return record |
|
|
|
|
| def build_gguf(model: str, out: str, llama_cpp: str, outtype: str, python: str) -> dict: |
| converter = os.path.join(llama_cpp, "convert_hf_to_gguf.py") |
| if not os.path.isfile(converter): |
| raise SystemExit(f"{converter} not found: --llama-cpp is the llama.cpp checkout") |
| source = os.path.join(out, "source") |
| os.makedirs(source) |
| link_model_files(model, source) |
| name = f"model-{outtype}.gguf" |
| t0 = time.perf_counter() |
| subprocess.run([python, converter, source, "--no-mtp", "--outfile", os.path.join(out, name), "--outtype", outtype], check=True) |
| shutil.rmtree(source) |
| commit = subprocess.run(["git", "-C", llama_cpp, "rev-parse", "--short", "HEAD"], capture_output=True, text=True).stdout.strip() or None |
| return {"format": BUILD_FORMAT, "backend": "llamacpp", "model": model_name(model), "model_fingerprint": fingerprint(model), "files": {"gguf": name}, |
| "type": outtype, "bytes": os.path.getsize(os.path.join(out, name)), "llama_cpp": commit, "seconds": round(time.perf_counter() - t0, 1)} |
|
|
|
|
| def main(argv: list[str] | None = None) -> int: |
| parser = argparse.ArgumentParser(prog="python -m decisionnode_edge.build", description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| parser.add_argument("backend", choices=("openvino", "gguf")) |
| parser.add_argument("model", help="the model directory (only read)") |
| parser.add_argument("--out", required=True, help="the directory to write (must not exist)") |
| parser.add_argument("--lengths", default=",".join(map(str, DEFAULT_LENGTHS)), help="openvino: the static lengths in tokens (default 512,2048)") |
| parser.add_argument("--device", default=None, help="openvino: where the export runs, cuda or cpu (default: cuda when present)") |
| parser.add_argument("--llama-cpp", help="gguf: the llama.cpp checkout (its convert_hf_to_gguf.py)") |
| parser.add_argument("--type", default="q8_0", choices=("q8_0", "bf16"), help="gguf: the weight type (default q8_0)") |
| parser.add_argument("--python", default=sys.executable, help="gguf: the Python that runs llama.cpp's converter (default: this one)") |
| args = parser.parse_args(argv) |
| for name in ("config.json", "model.safetensors", "decisionnode_edge.json"): |
| if not os.path.isfile(os.path.join(args.model, name)): |
| raise SystemExit(f"{args.model} is not a DecisionNode Edge model directory (no {name})") |
| if args.backend == "gguf" and not args.llama_cpp: |
| raise SystemExit("gguf needs --llama-cpp") |
| new_directory(args.out) |
| if args.backend == "openvino": |
| lengths = sorted({int(x) for x in args.lengths.split(",") if x.strip()}) |
| record = build_openvino(args.model, args.out, lengths, args.device) |
| else: |
| record = build_gguf(args.model, args.out, args.llama_cpp, args.type, args.python) |
| write_record(args.out, record) |
| print(json.dumps({k: v for k, v in record.items() if k != "versions"}), flush=True) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main()) |
|
|
