File size: 9,169 Bytes
071512e | 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 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 | """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": {}} # fmt: skip
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)
# each IR against torch (float32) on a prompt shorter than its length, so the padding is read past as at run time
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") # the converter reads links to the backbone files, never the readout
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)} # fmt: skip
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())
|