Efim Beshmenev
463 lines
18 KiB
Python
463 lines
18 KiB
Python
#!/usr/bin/env python3
|
|
"""Bounded numerical parity check for PyTorch and native transformer inference."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import math
|
|
import struct
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import zlib
|
|
from pathlib import Path
|
|
from typing import Sequence
|
|
|
|
import numpy as np
|
|
|
|
import train_transformer as training
|
|
|
|
|
|
PROBE_INPUT_MAGIC = b"SZTPRB01"
|
|
PROBE_OUTPUT_MAGIC = b"SZTPOU01"
|
|
PROBE_VERSION = 1
|
|
PROBE_INPUT_HEADER = struct.Struct("<8sIIIII")
|
|
PROBE_OUTPUT_HEADER = struct.Struct("<8sIII")
|
|
MAXIMUM_CASES = 64
|
|
OUTPUT_FLOAT_COUNT = 45
|
|
|
|
OUTPUT_NAMES = (
|
|
"improvement_logit",
|
|
"improvement_probability",
|
|
"expected_defect_gain",
|
|
"improvement_probability_variance",
|
|
"expected_defect_gain_variance",
|
|
*(f"plane_probability[{index}]" for index in range(training.FACE_COUNT)),
|
|
*(f"plane_probability_variance[{index}]" for index in range(training.FACE_COUNT)),
|
|
*(f"move_probability[{index}]" for index in range(training.MOVE_COUNT)),
|
|
*(f"move_probability_variance[{index}]" for index in range(training.MOVE_COUNT)),
|
|
*(f"scale_probability[{index}]" for index in range(training.SCALE_COUNT)),
|
|
*(f"scale_probability_variance[{index}]" for index in range(training.SCALE_COUNT)),
|
|
)
|
|
|
|
|
|
def parse_args(argv: Sequence[str]) -> argparse.Namespace:
|
|
repository = Path(__file__).resolve().parents[1]
|
|
parser = argparse.ArgumentParser(
|
|
description=(
|
|
"Compare a serialized SZTRNK01 model through CPU FP32 PyTorch and "
|
|
"the native TransformerRanker implementation."
|
|
)
|
|
)
|
|
parser.add_argument("--repository", default=str(repository))
|
|
parser.add_argument(
|
|
"--model", default="results/search/neural/transformer/current.sztf"
|
|
)
|
|
parser.add_argument(
|
|
"--native-probe",
|
|
help="TransformerRankerSelfTest executable containing --parity-probe",
|
|
)
|
|
parser.add_argument("--cases", type=int, default=8)
|
|
parser.add_argument("--absolute-tolerance", type=float, default=2.0e-4)
|
|
parser.add_argument("--relative-tolerance", type=float, default=2.0e-4)
|
|
parser.add_argument("--timeout-seconds", type=float, default=30.0)
|
|
parser.add_argument("--torch-threads", type=int, default=1)
|
|
parser.add_argument("--device", choices=("cpu", "cuda"), default="cpu")
|
|
return parser.parse_args(argv)
|
|
|
|
|
|
def resolve_path(repository: Path, value: str) -> Path:
|
|
path = Path(value).expanduser()
|
|
if not path.is_absolute():
|
|
path = repository / path
|
|
return path.resolve()
|
|
|
|
|
|
def find_native_probe(repository: Path, explicit: str | None) -> Path:
|
|
if explicit:
|
|
path = resolve_path(repository, explicit)
|
|
if not path.is_file():
|
|
raise FileNotFoundError(f"native probe does not exist: {path}")
|
|
return path
|
|
candidates = (
|
|
repository
|
|
/ "build/msbuild/bin/x64/Release/TransformerRankerSelfTest.exe",
|
|
repository / "build/vs2026/Release/transformer_ranker_selftest.exe",
|
|
repository / "build/TransformerRankerSelfTest.exe",
|
|
)
|
|
for candidate in candidates:
|
|
if candidate.is_file():
|
|
return candidate.resolve()
|
|
raise FileNotFoundError(
|
|
"native probe was not found; build transformer_ranker_selftest or pass "
|
|
"--native-probe"
|
|
)
|
|
|
|
|
|
def validate_options(args: argparse.Namespace) -> None:
|
|
if not 1 <= args.cases <= MAXIMUM_CASES:
|
|
raise ValueError(f"--cases must be in 1..{MAXIMUM_CASES}")
|
|
for name in ("absolute_tolerance", "relative_tolerance"):
|
|
value = float(getattr(args, name))
|
|
if not math.isfinite(value) or value < 0.0:
|
|
raise ValueError(f"--{name.replace('_', '-')} must be finite and non-negative")
|
|
if not math.isfinite(args.timeout_seconds) or not 1.0 <= args.timeout_seconds <= 300.0:
|
|
raise ValueError("--timeout-seconds must be in 1..300")
|
|
if not 1 <= args.torch_threads <= 64:
|
|
raise ValueError("--torch-threads must be in 1..64")
|
|
|
|
|
|
def load_model(
|
|
path: Path, device: training.torch.device
|
|
) -> tuple[list[training.RankerMember], str]:
|
|
model_bytes = path.read_bytes()
|
|
expected_float_count = training.expected_payload_float_count()
|
|
expected_payload_bytes = expected_float_count * 4
|
|
expected_size = training.MODEL_HEADER_BYTES + expected_payload_bytes
|
|
if len(model_bytes) != expected_size:
|
|
raise ValueError(
|
|
f"model size mismatch: {len(model_bytes)} != {expected_size}"
|
|
)
|
|
if model_bytes[:8] != training.MODEL_MAGIC:
|
|
raise ValueError("model magic is invalid")
|
|
|
|
format_version, feature_version, objective_version, approved = struct.unpack_from(
|
|
"<4I", model_bytes, 8
|
|
)
|
|
dimensions = struct.unpack_from("<9I", model_bytes, 24)
|
|
expected_dimensions = (
|
|
training.FACE_COUNT,
|
|
training.FACE_FEATURES,
|
|
training.GLOBAL_FEATURES,
|
|
training.MODEL_WIDTH,
|
|
training.ATTENTION_HEADS,
|
|
training.LAYER_COUNT,
|
|
training.FEED_FORWARD_WIDTH,
|
|
training.ENSEMBLE_SIZE,
|
|
training.TOPOLOGY_COUNT,
|
|
)
|
|
if (
|
|
format_version != training.MODEL_FORMAT_VERSION
|
|
or feature_version != training.FEATURE_FORMAT_VERSION
|
|
or objective_version != training.OBJECTIVE_VERSION
|
|
or approved != 1
|
|
or dimensions != expected_dimensions
|
|
or struct.unpack_from("<I", model_bytes, 60)[0] != 0
|
|
or not any(model_bytes[72:104])
|
|
or struct.unpack_from("<QQ", model_bytes, 144) != (0, 0)
|
|
):
|
|
raise ValueError("model header versions, dimensions, or approval are invalid")
|
|
|
|
validation = struct.unpack_from("<4f", model_bytes, 104)
|
|
if not all(math.isfinite(value) for value in validation):
|
|
raise ValueError("model validation metadata is non-finite")
|
|
if not all(0.0 <= validation[index] <= 1.0 for index in range(3)):
|
|
raise ValueError("model validation probability metrics are out of range")
|
|
|
|
payload_bytes, payload_crc, header_crc, float_count = struct.unpack_from(
|
|
"<QIIQ", model_bytes, 120
|
|
)
|
|
if payload_bytes != expected_payload_bytes or float_count != expected_float_count:
|
|
raise ValueError("model payload accounting is invalid")
|
|
header = bytearray(model_bytes[: training.MODEL_HEADER_BYTES])
|
|
struct.pack_into("<I", header, 132, 0)
|
|
if zlib.crc32(header) & 0xFFFFFFFF != header_crc:
|
|
raise ValueError("model header CRC-32 mismatch")
|
|
payload = model_bytes[training.MODEL_HEADER_BYTES :]
|
|
if zlib.crc32(payload) & 0xFFFFFFFF != payload_crc:
|
|
raise ValueError("model payload CRC-32 mismatch")
|
|
|
|
values = np.frombuffer(payload, dtype="<f4").copy()
|
|
if values.size != expected_float_count or not np.isfinite(values).all():
|
|
raise ValueError("model payload contains invalid float values")
|
|
members: list[training.RankerMember] = []
|
|
cursor = 0
|
|
with training.torch.no_grad():
|
|
for _ in range(training.ENSEMBLE_SIZE):
|
|
member = training.RankerMember()
|
|
for tensor in training.member_tensors(member):
|
|
count = tensor.numel()
|
|
source = training.torch.from_numpy(
|
|
values[cursor : cursor + count].reshape(tuple(tensor.shape))
|
|
)
|
|
tensor.copy_(source)
|
|
cursor += count
|
|
members.append(member.eval().to(device))
|
|
if cursor != values.size:
|
|
raise AssertionError(f"model tensor cursor mismatch: {cursor} != {values.size}")
|
|
return members, hashlib.sha256(model_bytes).hexdigest()
|
|
|
|
|
|
def make_cases(case_count: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
|
|
# Every value is an integer divided by 64, so both sides receive the same
|
|
# exactly representable float32 bits. The bounded formula covers signs,
|
|
# magnitudes, face positions and several topology embeddings.
|
|
sample = np.arange(case_count, dtype=np.int64)[:, None, None]
|
|
face = np.arange(training.FACE_COUNT, dtype=np.int64)[None, :, None]
|
|
component = np.arange(training.FACE_FEATURES, dtype=np.int64)[None, None, :]
|
|
face_features = (
|
|
((sample * 37 + face * 17 + component * 13) % 129) - 64
|
|
).astype(np.float32) / np.float32(64.0)
|
|
global_component = np.arange(training.GLOBAL_FEATURES, dtype=np.int64)[None, :]
|
|
global_features = (
|
|
(
|
|
(np.arange(case_count, dtype=np.int64)[:, None] * 29
|
|
+ global_component * 11)
|
|
% 129
|
|
)
|
|
- 64
|
|
).astype(np.float32) / np.float32(64.0)
|
|
topology = (
|
|
(np.arange(case_count, dtype=np.uint32) * np.uint32(17) + np.uint32(42))
|
|
% np.uint32(training.TOPOLOGY_COUNT)
|
|
)
|
|
return (
|
|
np.ascontiguousarray(face_features, dtype="<f4"),
|
|
np.ascontiguousarray(global_features, dtype="<f4"),
|
|
np.ascontiguousarray(topology, dtype="<u4"),
|
|
)
|
|
|
|
|
|
def write_probe_input(
|
|
path: Path,
|
|
face: np.ndarray,
|
|
global_features: np.ndarray,
|
|
topology: np.ndarray,
|
|
) -> None:
|
|
case_count = int(face.shape[0])
|
|
payload = bytearray(
|
|
PROBE_INPUT_HEADER.pack(
|
|
PROBE_INPUT_MAGIC,
|
|
PROBE_VERSION,
|
|
case_count,
|
|
training.FACE_COUNT,
|
|
training.FACE_FEATURES,
|
|
training.GLOBAL_FEATURES,
|
|
)
|
|
)
|
|
for index in range(case_count):
|
|
payload.extend(struct.pack("<I", int(topology[index])))
|
|
payload.extend(face[index].tobytes(order="C"))
|
|
payload.extend(global_features[index].tobytes(order="C"))
|
|
path.write_bytes(payload)
|
|
|
|
|
|
def read_probe_output(path: Path, expected_cases: int) -> tuple[np.ndarray, np.ndarray]:
|
|
data = path.read_bytes()
|
|
if len(data) < PROBE_OUTPUT_HEADER.size:
|
|
raise ValueError("native probe output is truncated")
|
|
magic, version, case_count, float_count = PROBE_OUTPUT_HEADER.unpack_from(data)
|
|
if (
|
|
magic != PROBE_OUTPUT_MAGIC
|
|
or version != PROBE_VERSION
|
|
or case_count != expected_cases
|
|
or float_count != OUTPUT_FLOAT_COUNT
|
|
):
|
|
raise ValueError("native probe output header is invalid")
|
|
record_size = 4 + OUTPUT_FLOAT_COUNT * 4
|
|
if len(data) != PROBE_OUTPUT_HEADER.size + case_count * record_size:
|
|
raise ValueError("native probe output has trailing or truncated data")
|
|
finite = np.empty(case_count, dtype=np.bool_)
|
|
output = np.empty((case_count, OUTPUT_FLOAT_COUNT), dtype=np.float32)
|
|
offset = PROBE_OUTPUT_HEADER.size
|
|
for index in range(case_count):
|
|
flag = struct.unpack_from("<I", data, offset)[0]
|
|
if flag not in (0, 1):
|
|
raise ValueError("native probe emitted an invalid finite flag")
|
|
finite[index] = bool(flag)
|
|
offset += 4
|
|
output[index] = np.frombuffer(
|
|
data, dtype="<f4", count=OUTPUT_FLOAT_COUNT, offset=offset
|
|
)
|
|
offset += OUTPUT_FLOAT_COUNT * 4
|
|
if not np.isfinite(output).all():
|
|
raise ValueError("native probe emitted NaN or infinity")
|
|
return finite, output
|
|
|
|
|
|
@training.torch.inference_mode()
|
|
def pytorch_reference(
|
|
members: list[training.RankerMember],
|
|
face: np.ndarray,
|
|
global_features: np.ndarray,
|
|
topology: np.ndarray,
|
|
device: training.torch.device,
|
|
) -> np.ndarray:
|
|
face_tensor = training.torch.from_numpy(face).to(device)
|
|
global_tensor = training.torch.from_numpy(global_features).to(device)
|
|
topology_tensor = training.torch.from_numpy(topology.astype(np.int64)).to(device)
|
|
probabilities = []
|
|
gains = []
|
|
plane_probabilities = []
|
|
move_probabilities = []
|
|
scale_probabilities = []
|
|
for member in members:
|
|
value, plane, move, scale = member(
|
|
face_tensor, global_tensor, topology_tensor
|
|
)
|
|
probabilities.append(training.torch.sigmoid(value[:, 0]))
|
|
gains.append(value[:, 1])
|
|
plane_probabilities.append(training.torch.softmax(plane, dim=-1))
|
|
move_probabilities.append(training.torch.softmax(move, dim=-1))
|
|
scale_probabilities.append(training.torch.softmax(scale, dim=-1))
|
|
|
|
def aggregate(values: list[training.torch.Tensor]):
|
|
stacked = training.torch.stack(values)
|
|
mean = stacked.mean(dim=0)
|
|
variance = ((stacked - mean.unsqueeze(0)) ** 2).mean(dim=0)
|
|
return mean, variance
|
|
|
|
probability, probability_variance = aggregate(probabilities)
|
|
gain, gain_variance = aggregate(gains)
|
|
plane, plane_variance = aggregate(plane_probabilities)
|
|
move, move_variance = aggregate(move_probabilities)
|
|
scale, scale_variance = aggregate(scale_probabilities)
|
|
improvement_logit = training.torch.logit(
|
|
probability.clamp(1.0e-6, 1.0 - 1.0e-6)
|
|
)
|
|
result = training.torch.cat(
|
|
(
|
|
improvement_logit[:, None],
|
|
probability[:, None],
|
|
gain[:, None],
|
|
probability_variance[:, None],
|
|
gain_variance[:, None],
|
|
plane,
|
|
plane_variance,
|
|
move,
|
|
move_variance,
|
|
scale,
|
|
scale_variance,
|
|
),
|
|
dim=1,
|
|
)
|
|
if result.shape[1] != OUTPUT_FLOAT_COUNT or not result.isfinite().all():
|
|
raise RuntimeError("PyTorch reference emitted invalid output")
|
|
return result.cpu().numpy()
|
|
|
|
|
|
def report_comparison(
|
|
native: np.ndarray,
|
|
reference: np.ndarray,
|
|
absolute_tolerance: float,
|
|
relative_tolerance: float,
|
|
) -> bool:
|
|
difference = np.abs(native.astype(np.float64) - reference.astype(np.float64))
|
|
tolerance = absolute_tolerance + relative_tolerance * np.abs(
|
|
reference.astype(np.float64)
|
|
)
|
|
max_absolute_flat = int(np.argmax(difference))
|
|
absolute_case, absolute_field = np.unravel_index(
|
|
max_absolute_flat, difference.shape
|
|
)
|
|
relative = difference / np.maximum(
|
|
np.maximum(np.abs(native), np.abs(reference)).astype(np.float64), 1.0e-12
|
|
)
|
|
max_relative_flat = int(np.argmax(relative))
|
|
relative_case, relative_field = np.unravel_index(
|
|
max_relative_flat, relative.shape
|
|
)
|
|
ratio = difference / np.maximum(tolerance, np.finfo(np.float64).tiny)
|
|
max_ratio_flat = int(np.argmax(ratio))
|
|
ratio_case, ratio_field = np.unravel_index(max_ratio_flat, ratio.shape)
|
|
|
|
print(
|
|
"max_abs_diff="
|
|
f"{difference[absolute_case, absolute_field]:.9e} "
|
|
f"case={absolute_case} field={OUTPUT_NAMES[absolute_field]} "
|
|
f"native={native[absolute_case, absolute_field]:.9e} "
|
|
f"pytorch={reference[absolute_case, absolute_field]:.9e}"
|
|
)
|
|
print(
|
|
"max_rel_diff="
|
|
f"{relative[relative_case, relative_field]:.9e} "
|
|
f"case={relative_case} field={OUTPUT_NAMES[relative_field]} "
|
|
f"native={native[relative_case, relative_field]:.9e} "
|
|
f"pytorch={reference[relative_case, relative_field]:.9e}"
|
|
)
|
|
print(
|
|
"max_tolerance_ratio="
|
|
f"{ratio[ratio_case, ratio_field]:.9e} "
|
|
f"case={ratio_case} field={OUTPUT_NAMES[ratio_field]}"
|
|
)
|
|
return bool(np.all(difference <= tolerance))
|
|
|
|
|
|
def main(argv: Sequence[str] | None = None) -> int:
|
|
args = parse_args(argv or sys.argv[1:])
|
|
validate_options(args)
|
|
repository = Path(args.repository).expanduser().resolve()
|
|
model_path = resolve_path(repository, args.model)
|
|
native_probe = find_native_probe(repository, args.native_probe)
|
|
if not model_path.is_file():
|
|
raise FileNotFoundError(f"model does not exist: {model_path}")
|
|
if args.device == "cuda" and not training.torch.cuda.is_available():
|
|
raise RuntimeError("--device cuda requested but CUDA PyTorch is unavailable")
|
|
|
|
training.torch.set_num_threads(args.torch_threads)
|
|
training.torch.set_float32_matmul_precision("highest")
|
|
if training.torch.cuda.is_available():
|
|
training.torch.backends.cuda.matmul.allow_tf32 = False
|
|
training.torch.backends.cudnn.allow_tf32 = False
|
|
device = training.torch.device(args.device)
|
|
members, model_digest = load_model(model_path, device)
|
|
face, global_features, topology = make_cases(args.cases)
|
|
reference = pytorch_reference(
|
|
members, face, global_features, topology, device
|
|
)
|
|
|
|
with tempfile.TemporaryDirectory(prefix="szilassi_transformer_parity_") as temp:
|
|
input_path = Path(temp) / "input.bin"
|
|
output_path = Path(temp) / "native.bin"
|
|
write_probe_input(input_path, face, global_features, topology)
|
|
command = [
|
|
str(native_probe),
|
|
"--parity-probe",
|
|
str(model_path),
|
|
str(input_path),
|
|
str(output_path),
|
|
]
|
|
completed = subprocess.run(
|
|
command,
|
|
cwd=repository,
|
|
check=False,
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=args.timeout_seconds,
|
|
)
|
|
if completed.returncode != 0:
|
|
raise RuntimeError(
|
|
f"native probe failed with exit code {completed.returncode}: "
|
|
f"{completed.stderr.strip() or completed.stdout.strip()}"
|
|
)
|
|
finite, native = read_probe_output(output_path, args.cases)
|
|
if not finite.all():
|
|
bad_cases = np.flatnonzero(~finite).tolist()
|
|
raise RuntimeError(f"native inference fell back for cases: {bad_cases}")
|
|
|
|
print(f"model_sha256={model_digest}")
|
|
print(
|
|
f"device={device.type} cases={args.cases} "
|
|
f"compared_floats={native.size} torch_threads={args.torch_threads}"
|
|
)
|
|
passed = report_comparison(
|
|
native,
|
|
reference,
|
|
args.absolute_tolerance,
|
|
args.relative_tolerance,
|
|
)
|
|
print(
|
|
f"tolerance=atol:{args.absolute_tolerance:.9e},"
|
|
f"rtol:{args.relative_tolerance:.9e} status={'PASS' if passed else 'FAIL'}"
|
|
)
|
|
return 0 if passed else 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
raise SystemExit(main())
|
|
except (OSError, ValueError, RuntimeError, subprocess.TimeoutExpired) as exc:
|
|
print(f"Transformer parity check failed: {exc}", file=sys.stderr)
|
|
raise SystemExit(2) from exc
|