Transformer
This commit is contained in:
Efim Beshmenev
@@ -0,0 +1,462 @@
|
||||
#!/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
|
||||
Reference in New Issue
Block a user