File size: 4,320 Bytes
8020c75 | 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 | from __future__ import annotations
import hashlib
import json
from collections import Counter
from pathlib import Path
import numpy as np
import onnx
import onnxruntime as ort
from onnxruntime.transformers.float16 import convert_float_to_float16
from onnxruntime.quantization.onnx_model import ONNXModel
ROOT = Path(__file__).resolve().parent
OUTPUT_DIR = ROOT / "output"
REPORT_DIR = ROOT / "reports"
GRAPH_NAMES = ("decoder_prefill", "decoder_step")
def sha256(path: Path) -> str:
digest = hashlib.sha256()
with path.open("rb") as handle:
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
digest.update(chunk)
return digest.hexdigest()
def make_inputs(session: ort.InferenceSession) -> dict[str, np.ndarray]:
rng = np.random.default_rng(20260716)
feeds: dict[str, np.ndarray] = {}
for value in session.get_inputs():
shape = [257 if dimension == "past_len" else dimension for dimension in value.shape]
if value.name == "vision_embeds":
feeds[value.name] = rng.normal(0, 0.2, shape).astype(np.float32)
elif value.name == "token_one_hot":
token = np.zeros(shape, dtype=np.float32)
token[0, 0, 4] = 1
feeds[value.name] = token
elif value.name == "position_ids":
feeds[value.name] = np.array([[257]], dtype=np.int32)
else:
feeds[value.name] = rng.normal(0, 0.02, shape).astype(np.float32)
return feeds
def compare(reference: Path, candidate: Path) -> dict:
expected_session = ort.InferenceSession(
str(reference), providers=["CPUExecutionProvider"]
)
actual_session = ort.InferenceSession(
str(candidate), providers=["CPUExecutionProvider"]
)
feeds = make_inputs(expected_session)
expected = expected_session.run(None, feeds)
actual = actual_session.run(None, feeds)
differences = [
float(np.max(np.abs(expected_value - actual_value)))
for expected_value, actual_value in zip(expected, actual)
]
expected_token = int(expected[0][0, -1].argmax())
actual_token = int(actual[0][0, -1].argmax())
return {
"max_abs": max(differences),
"logits_max_abs": differences[0],
"cache_max_abs": max(differences[1:]),
"reference_top_token": expected_token,
"candidate_top_token": actual_token,
"top_token_matches": expected_token == actual_token,
}
def main() -> None:
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
REPORT_DIR.mkdir(parents=True, exist_ok=True)
report = {
"description": (
"Complete Baberu decoder converted to FP16 internal tensors with FP32 "
"public inputs and outputs; no layer, hidden-size, vocabulary, or "
"generation-logic reduction"
),
"graphs": {},
}
for name in GRAPH_NAMES:
source = OUTPUT_DIR / f"{name}_fp32.onnx"
destination = OUTPUT_DIR / f"{name}_fp16.onnx"
if not source.exists():
raise SystemExit(f"Missing {source}. Run export_decoder_fp32.py first.")
print(f"Converting {source.name} -> {destination.name}", flush=True)
model = convert_float_to_float16(onnx.load(source), keep_io_types=True)
sortable = ONNXModel(model)
sortable.topological_sort()
model = sortable.model
model.producer_name = "vibe-manga-baberu-webgpu-fp16"
model.producer_version = "1"
onnx.checker.check_model(model)
onnx.save(model, destination)
operators = Counter(node.op_type for node in model.graph.node)
parity = compare(source, destination)
if not parity["top_token_matches"]:
raise RuntimeError(f"{destination.name}: synthetic top-token parity failed")
report["graphs"][name] = {
"source": source.name,
"destination": destination.name,
"bytes": destination.stat().st_size,
"sha256": sha256(destination),
"operators": dict(sorted(operators.items())),
"fp32_cpu_parity": parity,
}
destination = REPORT_DIR / "decoder-fp16-report.json"
destination.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
print(json.dumps(report, indent=2))
if __name__ == "__main__":
main()
|