| |
| """Compare serialized training diagnostics.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
|
|
| def _compare(left: np.ndarray, right: np.ndarray, atol: float, rtol: float): |
| delta = np.abs(left.astype(np.float64) - right.astype(np.float64)) |
| passed = np.isclose(left, right, atol=atol, rtol=rtol) |
| return { |
| "status": "PASS" if bool(np.all(passed)) else "FAIL", |
| "shape": list(left.shape), |
| "max_abs": float(delta.max()) if delta.size else 0.0, |
| "mean_abs": float(delta.mean()) if delta.size else 0.0, |
| "outside_tolerance": int((~passed).sum()), |
| "elements": int(delta.size), |
| "atol": atol, |
| "rtol": rtol, |
| } |
|
|
|
|
| def main() -> int: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("reference", type=Path) |
| parser.add_argument("candidate", type=Path) |
| parser.add_argument("output", type=Path) |
| args = parser.parse_args() |
| with np.load(args.reference) as left, np.load(args.candidate) as right: |
| stages = { |
| "logits": _compare(left["logits"], right["logits"], 1e-5, 1e-5), |
| "loss": _compare(left["loss"], right["loss"], 1e-5, 1e-5), |
| "gradients": _compare(left["gradients"], right["gradients"], 1e-4, 1e-5), |
| "parameters_after_step": _compare( |
| left["parameters"], right["parameters"], 1e-5, 1e-5 |
| ), |
| } |
| if "history" in left.files and "history" in right.files: |
| stages["history"] = _compare(left["history"], right["history"], 1e-5, 1e-5) |
| if "exp_avg" in left.files and "exp_avg" in right.files: |
| stages["optimizer_exp_avg"] = _compare( |
| left["exp_avg"], right["exp_avg"], 1e-5, 1e-5 |
| ) |
| stages["optimizer_exp_avg_sq"] = _compare( |
| left["exp_avg_sq"], right["exp_avg_sq"], 1e-5, 1e-5 |
| ) |
| stages["optimizer_steps"] = _compare( |
| left["optimizer_steps"], right["optimizer_steps"], 0.0, 0.0 |
| ) |
| report = {"schema_version": 1, "stages": stages} |
| report["overall_status"] = ( |
| "PASS" |
| if all(stage["status"] == "PASS" for stage in stages.values()) |
| else "FAIL" |
| ) |
| args.output.mkdir(parents=True, exist_ok=True) |
| (args.output / "training_step.json").write_text( |
| json.dumps(report, indent=2, sort_keys=True) + "\n" |
| ) |
| lines = [ |
| "# One-step training parity", |
| "", |
| f"Overall status: **{report['overall_status']}**", |
| "", |
| "| Stage | Status | Max abs | Outside tolerance |", |
| "|---|---|---:|---:|", |
| ] |
| for name, stage in stages.items(): |
| lines.append( |
| f"| {name} | {stage['status']} | {stage['max_abs']} | " |
| f"{stage['outside_tolerance']} |" |
| ) |
| (args.output / "training_step.md").write_text("\n".join(lines) + "\n") |
| print(json.dumps({"overall_status": report["overall_status"]}, indent=2)) |
| return 0 if report["overall_status"] == "PASS" else 1 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
