Minus-v0.2 / conversion /make_configs.py
GarageCyril's picture
Minus-v0.2: squash-trained LFM2.5-VL-450M + screen-state head, AX650 build with FP32 residual stream
97d19f2 verified
Raw History Blame Contribute Delete
2.86 kB
#!/usr/bin/env python3
"""Simplify exported graphs (static shapes) and write one Pulsar2 5.1 build config per
graph, reproducing the minus-v0.1 recipe (docs/LFM2_VL_AX650_OPTIMIZATION.md, D.5):
- NPU3 (all 3 cores), U16 everywhere,
- FP32 on every RMSNorm chain op (Mul/ReduceMean/Add/Sqrt/Div/Mul): 2 chains per conv
layer, 4 per attention layer (operator, q_layernorm, k_layernorm, ffn) + post's norm,
- FP32 runtime I/O (indices S32), Numpy calibration from gen_calib.py,
- vision: onnxsim off, smooth-quant + FP32 conv bias (AXERA vision-encoder practice).
python make_configs.py <build_dir>
"""
import json, os, sys
import onnx
from onnxsim import simplify
B = sys.argv[1]
meta = json.load(open(f"{B}/export_meta.json"))
graphs = ["vision_encoder_512"] + [f"l{i}_{t}_fused" for i, t in enumerate(meta["layer_types"])] + ["post_d"]
os.makedirs(f"{B}/sim", exist_ok=True)
os.makedirs(f"{B}/configs", exist_ok=True)
summary = {}
for g in graphs:
m = onnx.load(f"{B}/onnx/{g}.onnx")
if g != "vision_encoder_512":
m, ok = simplify(m)
assert ok, g
onnx.save(m, f"{B}/sim/{g}.onnx", save_as_external_data=m.ByteSize() > 1.9e9)
norm_ops = [n.name for n in m.graph.node
if "norm" in n.name.lower() and n.op_type in ("Mul", "ReduceMean", "Add", "Sqrt", "Div", "Pow")]
inputs = [i.name for i in m.graph.input]
is_vision = g == "vision_encoder_512"
cfg = {
"model_type": "ONNX",
"npu_mode": "NPU3",
"target_hardware": "AX650",
"onnx_opt": {"enable_onnxsim": False},
"quant": {
"input_configs": [{"tensor_name": n, "calibration_dataset": f"/data/calib/{g}/{n}.tar",
"calibration_format": "Numpy", "calibration_size": -1} for n in inputs],
"layer_configs": [{"start_tensor_names": ["DEFAULT"], "end_tensor_names": ["DEFAULT"],
"data_type": "U16"}],
"calibration_method": "MinMax",
"precision_analysis": True,
"precision_analysis_method": "EndToEnd",
"conv_bias_data_type": "FP32",
"enable_smooth_quant": is_vision,
},
"input_processors": [{"tensor_name": n, "src_dtype": "S32" if n == "indices" else "FP32"}
for n in inputs],
"compiler": {"check": 0},
}
if norm_ops and not is_vision:
cfg["quant"]["layer_configs"].append({"layer_names": norm_ops, "data_type": "FP32"})
json.dump(cfg, open(f"{B}/configs/{g}.json", "w"), indent=1)
summary[g] = {"nodes": len(m.graph.node), "fp32_norm_ops": len(norm_ops),
"op_types": sorted({n.op_type for n in m.graph.node})}
print(g, summary[g]["nodes"], "nodes,", len(norm_ops), "FP32 norm ops", flush=True)
json.dump(summary, open(f"{B}/configs/summary.json", "w"), indent=1)