NiccoloN
183 lines
7.7 KiB
Python
183 lines
7.7 KiB
Python
#!/usr/bin/env python3
|
|
"""Save exact attention-tap arrays and quantify the MatMul error sources."""
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import onnx
|
|
from onnx import numpy_helper
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parents[2]
|
|
sys.path.insert(0, str(REPO_ROOT / "validation"))
|
|
|
|
from raptor_validation.onnx_utils import onnx_io # noqa: E402
|
|
from raptor_validation.validate_one import ( # noqa: E402
|
|
parse_pim_simulator_outputs,
|
|
sanitize_output_name,
|
|
)
|
|
|
|
|
|
TAP_NAMES = {
|
|
"v": "/model.10/m/m.0/attn/Split_output_2",
|
|
"raw": "/model.10/m/m.0/attn/MatMul_output_0",
|
|
"scaled": "/model.10/m/m.0/attn/Mul_output_0",
|
|
"rhs": "/model.10/m/m.0/attn/Transpose_1_output_0",
|
|
"c": "/model.10/m/m.0/attn/MatMul_1_output_0",
|
|
}
|
|
ABSOLUTE_TOLERANCE = 1e-3
|
|
RELATIVE_TOLERANCE = 1e-5
|
|
|
|
|
|
def sha256(path):
|
|
digest = hashlib.sha256()
|
|
with Path(path).open("rb") as stream:
|
|
for block in iter(lambda: stream.read(1 << 20), b""):
|
|
digest.update(block)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def metric(actual, expected):
|
|
difference = np.abs(actual.astype(np.float64) - expected.astype(np.float64))
|
|
allowed = ABSOLUTE_TOLERANCE + RELATIVE_TOLERANCE * np.abs(expected.astype(np.float64))
|
|
return {
|
|
"max_abs": float(np.max(difference)),
|
|
"mean_abs": float(np.mean(difference)),
|
|
"rms": float(np.sqrt(np.mean(np.square(difference)))),
|
|
"elements_over_validator_limit": int(np.count_nonzero(difference > allowed)),
|
|
}
|
|
|
|
|
|
def f32_matmul(lhs, rhs):
|
|
return np.matmul(lhs.astype(np.float32), rhs.astype(np.float32)).astype(np.float32)
|
|
|
|
|
|
def f64_matmul(lhs, rhs):
|
|
return np.matmul(lhs.astype(np.float64), rhs.astype(np.float64)).astype(np.float64)
|
|
|
|
|
|
def load_constant(model, output_name):
|
|
for initializer in model.graph.initializer:
|
|
if initializer.name == output_name:
|
|
return float(numpy_helper.to_array(initializer).reshape(-1)[0])
|
|
for node in model.graph.node:
|
|
if output_name not in node.output:
|
|
continue
|
|
for attribute in node.attribute:
|
|
if attribute.name == "value" and attribute.HasField("t"):
|
|
return float(numpy_helper.to_array(attribute.t).reshape(-1)[0])
|
|
raise ValueError(f"could not find ONNX Constant producing {output_name}")
|
|
|
|
|
|
def load_arrays(workspace, model_path):
|
|
model = onnx.load(model_path)
|
|
descriptors = onnx_io(model_path)
|
|
output_descriptors = {name: (index, dtype, shape) for index, name, dtype, shape in descriptors[1]}
|
|
missing = sorted(set(TAP_NAMES.values()) - set(output_descriptors))
|
|
if missing:
|
|
raise ValueError("tap model is missing outputs: " + ", ".join(missing))
|
|
|
|
sim_arrays = parse_pim_simulator_outputs(
|
|
workspace / "simulation" / "out.bin", descriptors[1]
|
|
)
|
|
reference = {}
|
|
simulated = {}
|
|
input_files = {}
|
|
for key, name in TAP_NAMES.items():
|
|
index, _dtype, shape = output_descriptors[name]
|
|
csv_path = workspace / "outputs" / f"output{index}_{sanitize_output_name(name)}.csv"
|
|
reference[key] = np.loadtxt(csv_path, delimiter=",", dtype=np.float32).reshape(shape)
|
|
simulated[key] = np.asarray(sim_arrays[index], dtype=np.float32).reshape(shape)
|
|
input_files[key] = csv_path
|
|
return reference, simulated, input_files
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--workspace", type=Path, required=True,
|
|
help="validator workspace containing inputs, outputs, and simulation")
|
|
parser.add_argument("--model", type=Path, required=True, help="five-output ONNX tap model")
|
|
parser.add_argument("--output-dir", type=Path, required=True,
|
|
help="directory for arrays.npz, metadata.json, and decomposition.json")
|
|
args = parser.parse_args()
|
|
args.output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
model = onnx.load(args.model)
|
|
reference, simulated, source_files = load_arrays(args.workspace, args.model)
|
|
scale = np.float32(load_constant(model, "/model.10/m/m.0/attn/Constant_1_output_0"))
|
|
ref_v, ref_raw, ref_scaled, ref_rhs, ref_c = (reference[key] for key in ("v", "raw", "scaled", "rhs", "c"))
|
|
sim_v, sim_raw, sim_scaled, sim_rhs, sim_c = (simulated[key] for key in ("v", "raw", "scaled", "rhs", "c"))
|
|
|
|
ref_score_transpose = np.swapaxes(ref_scaled, -1, -2)
|
|
sim_score_transpose = np.swapaxes(sim_scaled, -1, -2)
|
|
ref_ss_f32 = f32_matmul(ref_v, ref_rhs)
|
|
sim_ss_f32 = f32_matmul(sim_v, sim_rhs)
|
|
ref_ss_f64 = f64_matmul(ref_v, ref_rhs)
|
|
sim_ss_f64 = f64_matmul(sim_v, sim_rhs)
|
|
ref_split_f32 = (f32_matmul(ref_v, np.swapaxes(ref_raw, -1, -2)) * scale).astype(np.float32)
|
|
sim_split_f32 = (f32_matmul(sim_v, np.swapaxes(sim_raw, -1, -2)) * scale).astype(np.float32)
|
|
ref_split_f64 = f64_matmul(ref_v, np.swapaxes(ref_raw, -1, -2)) * np.float64(scale)
|
|
sim_split_f64 = f64_matmul(sim_v, np.swapaxes(sim_raw, -1, -2)) * np.float64(scale)
|
|
|
|
arrays = {
|
|
**{f"ref_{key}": value for key, value in reference.items()},
|
|
**{f"sim_{key}": value for key, value in simulated.items()},
|
|
"ref_ss_f32": ref_ss_f32,
|
|
"sim_ss_f32": sim_ss_f32,
|
|
"ref_ss_f64": ref_ss_f64,
|
|
"sim_ss_f64": sim_ss_f64,
|
|
"ref_split_f32": ref_split_f32,
|
|
"sim_split_f32": sim_split_f32,
|
|
"ref_split_f64": ref_split_f64,
|
|
"sim_split_f64": sim_split_f64,
|
|
}
|
|
arrays_path = args.output_dir / "arrays.npz"
|
|
np.savez_compressed(arrays_path, **arrays)
|
|
|
|
metrics = {
|
|
"validator_policy": {
|
|
"absolute_tolerance": ABSOLUTE_TOLERANCE,
|
|
"relative_tolerance": RELATIVE_TOLERANCE,
|
|
},
|
|
"scale": float(scale),
|
|
"shape": list(ref_c.shape),
|
|
"tap_differences": {key: metric(simulated[key], reference[key]) for key in TAP_NAMES},
|
|
"rhs_transpose_consistency": metric(ref_rhs, ref_score_transpose),
|
|
"sim_rhs_transpose_consistency": metric(sim_rhs, sim_score_transpose),
|
|
"c_sim_vs_ss_f32": metric(sim_c, ref_ss_f32),
|
|
"c_ref_vs_ss_f32": metric(ref_c, ref_ss_f32),
|
|
"c_sim_vs_simulated_inputs_ss_f32": metric(sim_c, sim_ss_f32),
|
|
"v_drift_only": metric(f32_matmul(sim_v, ref_rhs), ref_ss_f32),
|
|
"rhs_drift_only": metric(f32_matmul(ref_v, sim_rhs), ref_ss_f32),
|
|
"joint_input_drift": metric(sim_ss_f32, ref_ss_f32),
|
|
"scale_reassociation_reference": metric(ref_split_f32, ref_ss_f32),
|
|
"scale_reassociation_simulated": metric(sim_split_f32, sim_ss_f32),
|
|
"reference_accumulation_f32_vs_f64": metric(ref_ss_f32, ref_ss_f64),
|
|
"simulated_accumulation_f32_vs_f64": metric(sim_ss_f32, sim_ss_f64),
|
|
"split_accumulation_reference_f32_vs_f64": metric(ref_split_f32, ref_split_f64),
|
|
"split_accumulation_simulated_f32_vs_f64": metric(sim_split_f32, sim_split_f64),
|
|
}
|
|
decomposition_path = args.output_dir / "decomposition.json"
|
|
decomposition_path.write_text(json.dumps(metrics, indent=2) + "\n", encoding="utf-8")
|
|
|
|
metadata = {
|
|
"model": str(args.model),
|
|
"model_sha256": sha256(args.model),
|
|
"workspace": str(args.workspace),
|
|
"arrays_sha256": sha256(arrays_path),
|
|
"source_sha256": {key: sha256(path) for key, path in source_files.items()},
|
|
"simulator_output_sha256": sha256(args.workspace / "simulation" / "out.bin"),
|
|
"input_sha256": sha256(args.workspace / "inputs" / "in0.csv"),
|
|
"outputs": TAP_NAMES,
|
|
"arrays": {key: {"dtype": str(value.dtype), "shape": list(value.shape)} for key, value in arrays.items()},
|
|
}
|
|
(args.output_dir / "metadata.json").write_text(json.dumps(metadata, indent=2) + "\n", encoding="utf-8")
|
|
print(json.dumps(metrics, indent=2))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|