wuyang
111 lines
3.5 KiB
Python
111 lines
3.5 KiB
Python
#!/usr/bin/env python3
|
|
"""Compare deterministic evidence from two DeepSeek-V2-Lite trace runs."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--first", type=Path, required=True)
|
|
parser.add_argument("--second", type=Path, required=True)
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
parser.add_argument("--captured-at", default=None)
|
|
return parser.parse_args()
|
|
|
|
|
|
def read(path: Path) -> dict[str, Any]:
|
|
return json.loads(path.read_text())
|
|
|
|
|
|
def sha256(path: Path) -> str:
|
|
return hashlib.sha256(path.read_bytes()).hexdigest()
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
first = read(args.first)
|
|
second = read(args.second)
|
|
checks: list[dict[str, Any]] = []
|
|
|
|
def check(name: str, left: Any, right: Any) -> None:
|
|
checks.append({"name": name, "exact": left == right})
|
|
|
|
check("provenance", first["provenance"], second["provenance"])
|
|
check("configuration", first["configuration"], second["configuration"])
|
|
check("prompt_tokenization", first["prompts"], second["prompts"])
|
|
check(
|
|
"initial_hidden_hash",
|
|
first["execution"]["initial_hidden"]["sha256_fp32"],
|
|
second["execution"]["initial_hidden"]["sha256_fp32"],
|
|
)
|
|
check(
|
|
"final_hidden_hash",
|
|
first["execution"]["final_hidden"]["sha256_fp32"],
|
|
second["execution"]["final_hidden"]["sha256_fp32"],
|
|
)
|
|
for left, right in zip(
|
|
first["execution"]["layers"],
|
|
second["execution"]["layers"],
|
|
strict=True,
|
|
):
|
|
layer = left["layer"]
|
|
check(
|
|
f"layer_{layer}_hidden_hash",
|
|
left["hidden_after"]["sha256_fp32"],
|
|
right["hidden_after"]["sha256_fp32"],
|
|
)
|
|
check(f"layer_{layer}_mla_shapes", left["mla"], right["mla"])
|
|
if "routing" in left:
|
|
check(
|
|
f"layer_{layer}_aggregate_load",
|
|
left["routing"]["aggregate_load"],
|
|
right["routing"]["aggregate_load"],
|
|
)
|
|
check(
|
|
f"layer_{layer}_token_routes",
|
|
[
|
|
prompt["token_routes"]
|
|
for prompt in left["routing"]["per_prompt"]
|
|
],
|
|
[
|
|
prompt["token_routes"]
|
|
for prompt in right["routing"]["per_prompt"]
|
|
],
|
|
)
|
|
|
|
captured_at = args.captured_at or datetime.now(timezone.utc).isoformat()
|
|
result = {
|
|
"schema_version": 1,
|
|
"captured_at": captured_at,
|
|
"first": {
|
|
"captured_at": first["captured_at"],
|
|
"sha256": sha256(args.first),
|
|
},
|
|
"second": {
|
|
"captured_at": second["captured_at"],
|
|
"sha256": sha256(args.second),
|
|
},
|
|
"timing_compared": False,
|
|
"checks": checks,
|
|
"exact_checks": sum(item["exact"] for item in checks),
|
|
"total_checks": len(checks),
|
|
"all_exact": all(item["exact"] for item in checks),
|
|
}
|
|
if not result["all_exact"]:
|
|
failed = [item["name"] for item in checks if not item["exact"]]
|
|
raise AssertionError(f"trace mismatch: {failed}")
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
args.output.write_text(json.dumps(result, indent=2) + "\n")
|
|
print(json.dumps(result, indent=2))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|