wuyang
686 lines
25 KiB
Python
686 lines
25 KiB
Python
#!/usr/bin/env python3
|
||
"""Validate, aggregate, and publish K3 AttnRes Round 05 experiment data."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import hashlib
|
||
import json
|
||
import math
|
||
import os
|
||
import shutil
|
||
import statistics
|
||
from pathlib import Path
|
||
from typing import Any, Iterable
|
||
|
||
|
||
PROTOCOL_ID = "llm-atlas-k3-attnres-gradient-scale-v1"
|
||
ARCHITECTURES = ("baseline", "block")
|
||
DEPTHS = (16, 32)
|
||
SEEDS = (2026073001, 2026073002, 2026073003)
|
||
STEPS = (0, 100, 500, 2000, 4000, 8000)
|
||
FORMAL_STEPS = 8000
|
||
FORMAL_BATCH = 32
|
||
TARGET_BYTES_PER_RUN = 65_536_000
|
||
EXPECTED_TOTAL_TARGET_BYTES = 786_432_000
|
||
SMOKE_COMPARE_FIELDS = (
|
||
"protocol_id",
|
||
"run_kind",
|
||
"architecture",
|
||
"depth",
|
||
"seed",
|
||
"steps",
|
||
"batch_size",
|
||
"target_bytes_seen",
|
||
"manifest",
|
||
"model",
|
||
"optimizer",
|
||
"hashes",
|
||
"evaluations",
|
||
"diagnostics",
|
||
"training_history",
|
||
"gradient_gate",
|
||
"environment",
|
||
)
|
||
REPLAY_COMPARE_FIELDS = tuple(
|
||
field for field in SMOKE_COMPARE_FIELDS if field != "run_kind"
|
||
)
|
||
|
||
|
||
def parse_args() -> argparse.Namespace:
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument("--formal-dir", type=Path, required=True)
|
||
parser.add_argument("--smoke-a-dir", type=Path, required=True)
|
||
parser.add_argument("--smoke-b-dir", type=Path, required=True)
|
||
parser.add_argument("--replay", type=Path, required=True)
|
||
parser.add_argument("--manifest", type=Path, required=True)
|
||
parser.add_argument("--raw-output-dir", type=Path, required=True)
|
||
parser.add_argument("--output", type=Path, required=True)
|
||
parser.add_argument("--compact-output", type=Path, required=True)
|
||
parser.add_argument("--reproduction-output", type=Path, required=True)
|
||
return parser.parse_args()
|
||
|
||
|
||
def read_json(path: Path) -> dict[str, Any]:
|
||
return json.loads(path.read_text())
|
||
|
||
|
||
def file_sha256(path: Path) -> str:
|
||
digest = hashlib.sha256()
|
||
with path.open("rb") as handle:
|
||
for block in iter(lambda: handle.read(1024 * 1024), b""):
|
||
digest.update(block)
|
||
return digest.hexdigest()
|
||
|
||
|
||
def canonical_sha256(value: Any) -> str:
|
||
payload = json.dumps(
|
||
value, ensure_ascii=False, sort_keys=True, separators=(",", ":")
|
||
).encode()
|
||
return hashlib.sha256(payload).hexdigest()
|
||
|
||
|
||
def atomic_json(path: Path, value: dict[str, Any]) -> None:
|
||
path.parent.mkdir(parents=True, exist_ok=True)
|
||
temporary = path.with_suffix(path.suffix + ".tmp")
|
||
temporary.write_text(
|
||
json.dumps(value, ensure_ascii=False, indent=2, sort_keys=True) + "\n"
|
||
)
|
||
os.replace(temporary, path)
|
||
|
||
|
||
def mean(values: Iterable[float]) -> float:
|
||
return statistics.fmean(values)
|
||
|
||
|
||
def require_finite(value: Any, path: str = "root") -> None:
|
||
if isinstance(value, float):
|
||
if not math.isfinite(value):
|
||
raise ValueError(f"non-finite float at {path}")
|
||
elif isinstance(value, dict):
|
||
for key, child in value.items():
|
||
require_finite(child, f"{path}.{key}")
|
||
elif isinstance(value, list):
|
||
for index, child in enumerate(value):
|
||
require_finite(child, f"{path}[{index}]")
|
||
|
||
|
||
def selected(run: dict[str, Any], fields: tuple[str, ...]) -> dict[str, Any]:
|
||
return {field: run[field] for field in fields}
|
||
|
||
|
||
def final_diagnostic(run: dict[str, Any]) -> dict[str, Any]:
|
||
diagnostic = run["diagnostics"][-1]
|
||
if diagnostic["step"] != FORMAL_STEPS:
|
||
raise ValueError("final diagnostic is not step 8000")
|
||
return diagnostic
|
||
|
||
|
||
def validate_run(
|
||
run: dict[str, Any],
|
||
*,
|
||
path: Path,
|
||
depth: int,
|
||
architecture: str,
|
||
seed: int,
|
||
manifest: dict[str, Any],
|
||
manifest_hash: str,
|
||
) -> None:
|
||
if run["protocol_id"] != PROTOCOL_ID or run["run_kind"] != "formal":
|
||
raise ValueError(f"formal protocol/kind mismatch: {path}")
|
||
if (
|
||
run["depth"] != depth
|
||
or run["architecture"] != architecture
|
||
or run["seed"] != seed
|
||
):
|
||
raise ValueError(f"formal identity mismatch: {path}")
|
||
if (
|
||
run["steps"] != FORMAL_STEPS
|
||
or run["batch_size"] != FORMAL_BATCH
|
||
or run["target_bytes_seen"] != TARGET_BYTES_PER_RUN
|
||
):
|
||
raise ValueError(f"formal budget mismatch: {path}")
|
||
if run["manifest"]["file_sha256"] != manifest_hash:
|
||
raise ValueError(f"manifest file hash mismatch: {path}")
|
||
for key in (
|
||
"formal_schedule_sha256",
|
||
"validation_tensor_sha256",
|
||
"diagnostic_tensor_sha256",
|
||
):
|
||
if run["manifest"][key] != manifest["windows"][key]:
|
||
raise ValueError(f"manifest {key} mismatch: {path}")
|
||
if [row["step"] for row in run["evaluations"]] != list(STEPS):
|
||
raise ValueError(f"evaluation steps mismatch: {path}")
|
||
if [row["step"] for row in run["diagnostics"]] != list(STEPS):
|
||
raise ValueError(f"diagnostic steps mismatch: {path}")
|
||
if run["model"]["layers"] != depth:
|
||
raise ValueError(f"model depth mismatch: {path}")
|
||
|
||
for diagnostic in run["diagnostics"]:
|
||
capture = diagnostic["capture"]
|
||
if (
|
||
capture["count"] != depth
|
||
or capture["shape"] != [16, 256, 192]
|
||
or set(capture["dtypes"]) != {"torch.float32"}
|
||
or not capture["all_gradients_finite"]
|
||
or not capture["all_gradients_present"]
|
||
or not capture["storage_unique"]
|
||
):
|
||
raise ValueError(f"activation capture gate mismatch: {path}")
|
||
for field in (
|
||
"activation_grad_rms_by_block",
|
||
"activation_output_rms_by_block",
|
||
"core_parameter_grad_rms_by_block",
|
||
):
|
||
if len(diagnostic[field]) != depth:
|
||
raise ValueError(f"{field} length mismatch: {path}")
|
||
for field in (
|
||
"layer_input_rms_by_sublayer",
|
||
"branch_output_rms_by_sublayer",
|
||
"stream_state_rms_by_sublayer",
|
||
):
|
||
if len(diagnostic[field]) != depth * 2:
|
||
raise ValueError(f"{field} length mismatch: {path}")
|
||
if architecture == "baseline":
|
||
if diagnostic["depth_weights"] or diagnostic["output_weights"] is not None:
|
||
raise ValueError(f"unexpected Baseline mixer trace: {path}")
|
||
else:
|
||
if (
|
||
len(diagnostic["depth_weights"]) != depth * 2
|
||
or diagnostic["output_weights"]["sources"] != 9
|
||
):
|
||
raise ValueError(f"Block mixer trace mismatch: {path}")
|
||
require_finite(run, path.name)
|
||
|
||
|
||
def relative_reduction(baseline: float, block: float) -> float:
|
||
return (baseline - block) / baseline
|
||
|
||
|
||
def depth_verdict(rows: list[dict[str, Any]]) -> dict[str, Any]:
|
||
cv_reductions = [row["relative_cv_reduction"] for row in rows]
|
||
imbalance_reductions = [
|
||
row["relative_imbalance_reduction"] for row in rows
|
||
]
|
||
mean_cv = mean(cv_reductions)
|
||
mean_imbalance = mean(imbalance_reductions)
|
||
support = (
|
||
all(value > 0 for value in cv_reductions)
|
||
and mean_cv >= 0.20
|
||
and all(value > 0 for value in imbalance_reductions)
|
||
and mean_imbalance >= 0.20
|
||
)
|
||
concern = (
|
||
all(value < 0 for value in cv_reductions)
|
||
and -mean_cv >= 0.20
|
||
and all(value < 0 for value in imbalance_reductions)
|
||
and -mean_imbalance >= 0.20
|
||
)
|
||
if support:
|
||
label = "joint directional support at this depth"
|
||
elif concern:
|
||
label = "joint directional concern at this depth"
|
||
else:
|
||
label = "mixed / inconclusive at this depth"
|
||
return {
|
||
"label": label,
|
||
"threshold_relative": 0.20,
|
||
"cv_reductions": cv_reductions,
|
||
"mean_cv_reduction": mean_cv,
|
||
"imbalance_reductions": imbalance_reductions,
|
||
"mean_imbalance_reduction": mean_imbalance,
|
||
"all_cv_improve": all(value > 0 for value in cv_reductions),
|
||
"all_imbalance_improve": all(
|
||
value > 0 for value in imbalance_reductions
|
||
),
|
||
}
|
||
|
||
|
||
def summarize_depth(
|
||
depth: int, runs: dict[tuple[int, str, int], dict[str, Any]]
|
||
) -> dict[str, Any]:
|
||
rows = []
|
||
for seed in SEEDS:
|
||
baseline = runs[(depth, "baseline", seed)]
|
||
block = runs[(depth, "block", seed)]
|
||
baseline_diagnostic = final_diagnostic(baseline)
|
||
block_diagnostic = final_diagnostic(block)
|
||
baseline_activation = baseline_diagnostic[
|
||
"activation_grad_statistics"
|
||
]
|
||
block_activation = block_diagnostic["activation_grad_statistics"]
|
||
baseline_bpc = baseline["evaluations"][-1]["bits_per_byte"]
|
||
block_bpc = block["evaluations"][-1]["bits_per_byte"]
|
||
rows.append(
|
||
{
|
||
"seed": seed,
|
||
"baseline_bpc": baseline_bpc,
|
||
"block_bpc": block_bpc,
|
||
"block_minus_baseline_bpc": block_bpc - baseline_bpc,
|
||
"baseline_activation_grad_mean": baseline_activation["mean"],
|
||
"block_activation_grad_mean": block_activation["mean"],
|
||
"block_to_baseline_activation_grad_mean": (
|
||
block_activation["mean"] / baseline_activation["mean"]
|
||
),
|
||
"baseline_cv": baseline_activation["population_cv"],
|
||
"block_cv": block_activation["population_cv"],
|
||
"relative_cv_reduction": relative_reduction(
|
||
baseline_activation["population_cv"],
|
||
block_activation["population_cv"],
|
||
),
|
||
"baseline_first_to_last_ratio": baseline_activation[
|
||
"first_to_last_ratio"
|
||
],
|
||
"block_first_to_last_ratio": block_activation[
|
||
"first_to_last_ratio"
|
||
],
|
||
"baseline_imbalance": baseline_activation[
|
||
"imbalance_abs_log_ratio"
|
||
],
|
||
"block_imbalance": block_activation[
|
||
"imbalance_abs_log_ratio"
|
||
],
|
||
"relative_imbalance_reduction": relative_reduction(
|
||
baseline_activation["imbalance_abs_log_ratio"],
|
||
block_activation["imbalance_abs_log_ratio"],
|
||
),
|
||
"baseline_parameter_grad_cv": baseline_diagnostic[
|
||
"core_parameter_grad_statistics"
|
||
]["population_cv"],
|
||
"block_parameter_grad_cv": block_diagnostic[
|
||
"core_parameter_grad_statistics"
|
||
]["population_cv"],
|
||
}
|
||
)
|
||
verdict = depth_verdict(rows)
|
||
return {
|
||
"depth": depth,
|
||
"by_seed": rows,
|
||
"means": {
|
||
"baseline_bpc": mean(row["baseline_bpc"] for row in rows),
|
||
"block_bpc": mean(row["block_bpc"] for row in rows),
|
||
"block_minus_baseline_bpc": mean(
|
||
row["block_minus_baseline_bpc"] for row in rows
|
||
),
|
||
"baseline_cv": mean(row["baseline_cv"] for row in rows),
|
||
"block_cv": mean(row["block_cv"] for row in rows),
|
||
"relative_cv_reduction": mean(
|
||
row["relative_cv_reduction"] for row in rows
|
||
),
|
||
"baseline_imbalance": mean(
|
||
row["baseline_imbalance"] for row in rows
|
||
),
|
||
"block_imbalance": mean(
|
||
row["block_imbalance"] for row in rows
|
||
),
|
||
"relative_imbalance_reduction": mean(
|
||
row["relative_imbalance_reduction"] for row in rows
|
||
),
|
||
"block_to_baseline_activation_grad_mean": mean(
|
||
row["block_to_baseline_activation_grad_mean"] for row in rows
|
||
),
|
||
"baseline_parameter_grad_cv": mean(
|
||
row["baseline_parameter_grad_cv"] for row in rows
|
||
),
|
||
"block_parameter_grad_cv": mean(
|
||
row["block_parameter_grad_cv"] for row in rows
|
||
),
|
||
},
|
||
"verdict": verdict,
|
||
}
|
||
|
||
|
||
def compact_cell(run: dict[str, Any]) -> dict[str, Any]:
|
||
diagnostics = []
|
||
for row in run["diagnostics"]:
|
||
compact = {
|
||
"step": row["step"],
|
||
"loss_nats": row["loss_nats"],
|
||
"bits_per_byte": row["bits_per_byte"],
|
||
"activation_grad_rms_by_block": row[
|
||
"activation_grad_rms_by_block"
|
||
],
|
||
"activation_grad_statistics": row[
|
||
"activation_grad_statistics"
|
||
],
|
||
"activation_output_rms_by_block": row[
|
||
"activation_output_rms_by_block"
|
||
],
|
||
"activation_output_statistics": row[
|
||
"activation_output_statistics"
|
||
],
|
||
"core_parameter_grad_rms_by_block": row[
|
||
"core_parameter_grad_rms_by_block"
|
||
],
|
||
"core_parameter_grad_statistics": row[
|
||
"core_parameter_grad_statistics"
|
||
],
|
||
}
|
||
if row["depth_weights"]:
|
||
compact["depth_weights"] = row["depth_weights"]
|
||
compact["output_weights"] = row["output_weights"]
|
||
diagnostics.append(compact)
|
||
return {
|
||
"architecture": run["architecture"],
|
||
"depth": run["depth"],
|
||
"seed": run["seed"],
|
||
"evaluations": run["evaluations"],
|
||
"diagnostics": diagnostics,
|
||
"timing": run["timing"],
|
||
"parameters": run["model"]["parameters"],
|
||
"hashes": run["hashes"],
|
||
}
|
||
|
||
|
||
def main() -> None:
|
||
args = parse_args()
|
||
manifest = read_json(args.manifest)
|
||
manifest_hash = file_sha256(args.manifest)
|
||
if manifest["protocol_id"] != PROTOCOL_ID:
|
||
raise ValueError("manifest protocol mismatch")
|
||
if (
|
||
manifest["windows"]["formal_schedule_cells"] != 768_000
|
||
or manifest["windows"]["formal_steps"] != FORMAL_STEPS
|
||
or manifest["windows"]["formal_batch"] != FORMAL_BATCH
|
||
):
|
||
raise ValueError("manifest schedule budget mismatch")
|
||
|
||
runs: dict[tuple[int, str, int], dict[str, Any]] = {}
|
||
source_paths: dict[str, Path] = {}
|
||
formal_hashes: dict[str, str] = {}
|
||
for depth in DEPTHS:
|
||
for architecture in ARCHITECTURES:
|
||
for seed in SEEDS:
|
||
name = f"depth-{depth}-{architecture}-seed-{seed}.json"
|
||
path = args.formal_dir / name
|
||
run = read_json(path)
|
||
validate_run(
|
||
run,
|
||
path=path,
|
||
depth=depth,
|
||
architecture=architecture,
|
||
seed=seed,
|
||
manifest=manifest,
|
||
manifest_hash=manifest_hash,
|
||
)
|
||
runs[(depth, architecture, seed)] = run
|
||
public_name = f"formal-{name}"
|
||
source_paths[public_name] = path
|
||
formal_hashes[public_name] = file_sha256(path)
|
||
|
||
if sum(run["target_bytes_seen"] for run in runs.values()) != (
|
||
EXPECTED_TOTAL_TARGET_BYTES
|
||
):
|
||
raise ValueError("formal total target-byte budget mismatch")
|
||
|
||
common_initial_exact: dict[str, Any] = {}
|
||
input_gate_exact: dict[str, Any] = {}
|
||
for depth in DEPTHS:
|
||
for seed in SEEDS:
|
||
baseline = runs[(depth, "baseline", seed)]
|
||
block = runs[(depth, "block", seed)]
|
||
public_fields = (
|
||
"initial_public_parameter_structure",
|
||
"initial_public_parameter_tensors",
|
||
"initial_public_parameter_elements",
|
||
"initial_public_parameters",
|
||
)
|
||
exact = all(
|
||
baseline["hashes"][field] == block["hashes"][field]
|
||
for field in public_fields
|
||
)
|
||
gate_exact = (
|
||
baseline["manifest"]["input_gate_tensor_hashes"]
|
||
== block["manifest"]["input_gate_tensor_hashes"]
|
||
)
|
||
key = f"depth-{depth}-seed-{seed}"
|
||
common_initial_exact[key] = {
|
||
"exact": exact,
|
||
"baseline": {
|
||
field: baseline["hashes"][field] for field in public_fields
|
||
},
|
||
"block": {
|
||
field: block["hashes"][field] for field in public_fields
|
||
},
|
||
}
|
||
input_gate_exact[key] = {
|
||
"exact": gate_exact,
|
||
"hashes": baseline["manifest"]["input_gate_tensor_hashes"],
|
||
}
|
||
if not exact or not gate_exact:
|
||
raise ValueError(f"paired equality gate failed: {key}")
|
||
|
||
smoke_exact: dict[str, Any] = {}
|
||
smoke_hashes: dict[str, str] = {}
|
||
for depth in DEPTHS:
|
||
for architecture in ARCHITECTURES:
|
||
name = f"depth-{depth}-{architecture}.json"
|
||
left_path = args.smoke_a_dir / name
|
||
right_path = args.smoke_b_dir / name
|
||
left = read_json(left_path)
|
||
right = read_json(right_path)
|
||
left_selected = selected(left, SMOKE_COMPARE_FIELDS)
|
||
right_selected = selected(right, SMOKE_COMPARE_FIELDS)
|
||
exact = left_selected == right_selected
|
||
if (
|
||
not exact
|
||
or left["run_kind"] != "smoke"
|
||
or left["steps"] != 20
|
||
or not left["gradient_gate"]["passed"]
|
||
):
|
||
raise ValueError(f"smoke gate failed: {name}")
|
||
key = f"depth-{depth}-{architecture}"
|
||
smoke_exact[key] = {
|
||
"exact": exact,
|
||
"compare_sha256": canonical_sha256(left_selected),
|
||
"gradient_gate": left["gradient_gate"],
|
||
}
|
||
for label, path in (("a", left_path), ("b", right_path)):
|
||
public_name = f"smoke-{label}-{name}"
|
||
source_paths[public_name] = path
|
||
smoke_hashes[public_name] = file_sha256(path)
|
||
|
||
replay = read_json(args.replay)
|
||
replay_formal = runs[(32, "block", 2026073001)]
|
||
replay_left = selected(replay_formal, REPLAY_COMPARE_FIELDS)
|
||
replay_right = selected(replay, REPLAY_COMPARE_FIELDS)
|
||
replay_exact = replay_left == replay_right
|
||
if (
|
||
replay["run_kind"] != "replay"
|
||
or replay["depth"] != 32
|
||
or replay["architecture"] != "block"
|
||
or replay["seed"] != 2026073001
|
||
or not replay_exact
|
||
):
|
||
raise ValueError("formal replay gate failed")
|
||
replay_public_name = "replay-depth-32-block-seed-2026073001.json"
|
||
source_paths[replay_public_name] = args.replay
|
||
|
||
depth_summaries = {
|
||
str(depth): summarize_depth(depth, runs) for depth in DEPTHS
|
||
}
|
||
depth_labels = [
|
||
depth_summaries[str(depth)]["verdict"]["label"] for depth in DEPTHS
|
||
]
|
||
if all(
|
||
label == "joint directional support at this depth"
|
||
for label in depth_labels
|
||
):
|
||
overall_verdict = (
|
||
"scale-consistent directional support in this operationalization"
|
||
)
|
||
elif all(
|
||
label == "joint directional concern at this depth"
|
||
for label in depth_labels
|
||
):
|
||
overall_verdict = (
|
||
"scale-consistent directional concern in this operationalization"
|
||
)
|
||
else:
|
||
overall_verdict = "depth-dependent or inconclusive"
|
||
|
||
full = {
|
||
"schema_version": 1,
|
||
"protocol_id": PROTOCOL_ID,
|
||
"manifest": manifest,
|
||
"study": {
|
||
"architectures": list(ARCHITECTURES),
|
||
"depths": list(DEPTHS),
|
||
"seeds": list(SEEDS),
|
||
"diagnostic_steps": list(STEPS),
|
||
"formal_runs": len(runs),
|
||
"formal_target_bytes": EXPECTED_TOTAL_TARGET_BYTES,
|
||
"replay_target_bytes": TARGET_BYTES_PER_RUN,
|
||
"gradient_object": (
|
||
"RMS of d(mean token CE)/d(post-MLP Transformer-block output) "
|
||
"over batch×time×channel"
|
||
),
|
||
},
|
||
"depth_summaries": depth_summaries,
|
||
"overall_verdict": overall_verdict,
|
||
"gates": {
|
||
"common_initial_parameters": common_initial_exact,
|
||
"paired_input_tensors": input_gate_exact,
|
||
"smoke_exact": smoke_exact,
|
||
"replay": {
|
||
"exact": replay_exact,
|
||
"compare_fields": list(REPLAY_COMPARE_FIELDS),
|
||
"formal_compare_sha256": canonical_sha256(replay_left),
|
||
"replay_compare_sha256": canonical_sha256(replay_right),
|
||
"formal_final_model_state": replay_formal["hashes"][
|
||
"final_model_state"
|
||
],
|
||
"replay_final_model_state": replay["hashes"][
|
||
"final_model_state"
|
||
],
|
||
"formal_final_optimizer_state": replay_formal["hashes"][
|
||
"final_optimizer_state"
|
||
],
|
||
"replay_final_optimizer_state": replay["hashes"][
|
||
"final_optimizer_state"
|
||
],
|
||
},
|
||
},
|
||
"runs": {
|
||
f"depth-{depth}-{architecture}-seed-{seed}": run
|
||
for (depth, architecture, seed), run in sorted(runs.items())
|
||
},
|
||
}
|
||
full["canonical_sha256_without_self"] = canonical_sha256(full)
|
||
|
||
compact = {
|
||
"schema_version": 1,
|
||
"protocol_id": PROTOCOL_ID,
|
||
"study": full["study"],
|
||
"manifest_summary": {
|
||
"file_sha256": manifest_hash,
|
||
"dataset_revision": manifest["dataset"]["revision"],
|
||
"train_bytes_sha256": manifest["dataset"]["splits"]["train"][
|
||
"concatenated_sha256"
|
||
],
|
||
"formal_schedule_sha256": manifest["windows"][
|
||
"formal_schedule_sha256"
|
||
],
|
||
"validation_tensor_sha256": manifest["windows"][
|
||
"validation_tensor_sha256"
|
||
],
|
||
"diagnostic_tensor_sha256": manifest["windows"][
|
||
"diagnostic_tensor_sha256"
|
||
],
|
||
},
|
||
"depth_summaries": depth_summaries,
|
||
"overall_verdict": overall_verdict,
|
||
"replay_exact": replay_exact,
|
||
"cells": [
|
||
compact_cell(runs[(depth, architecture, seed)])
|
||
for depth in DEPTHS
|
||
for architecture in ARCHITECTURES
|
||
for seed in SEEDS
|
||
],
|
||
}
|
||
compact["canonical_sha256_without_self"] = canonical_sha256(compact)
|
||
|
||
args.raw_output_dir.mkdir(parents=True, exist_ok=True)
|
||
for public_name, source_path in sorted(source_paths.items()):
|
||
target = args.raw_output_dir / public_name
|
||
temporary = target.with_suffix(target.suffix + ".tmp")
|
||
shutil.copyfile(source_path, temporary)
|
||
os.replace(temporary, target)
|
||
|
||
reproduction = {
|
||
"schema_version": 1,
|
||
"protocol_id": PROTOCOL_ID,
|
||
"manifest": {
|
||
"path": str(args.manifest),
|
||
"sha256": manifest_hash,
|
||
},
|
||
"protocol_sha256": file_sha256(
|
||
Path("research/K3_ATTNRES_GRADIENT_SCALE_PROTOCOL.md")
|
||
),
|
||
"definition_audit_sha256": file_sha256(
|
||
Path("research/K3_ATTNRES_GRADIENT_DEFINITION_AUDIT.md")
|
||
),
|
||
"runner_sha256": file_sha256(
|
||
Path("experiments/k3/attnres_gradient/train.py")
|
||
),
|
||
"analyzer_sha256": file_sha256(Path(__file__)),
|
||
"formal_raw_sha256": formal_hashes,
|
||
"smoke_raw_sha256": smoke_hashes,
|
||
"replay_raw_sha256": {
|
||
replay_public_name: file_sha256(args.replay)
|
||
},
|
||
"formal_target_bytes": EXPECTED_TOTAL_TARGET_BYTES,
|
||
"replay_target_bytes": TARGET_BYTES_PER_RUN,
|
||
"smoke_exact": smoke_exact,
|
||
"replay_exact": {
|
||
"exact": replay_exact,
|
||
"compare_sha256": canonical_sha256(replay_left),
|
||
"final_model_state": replay["hashes"]["final_model_state"],
|
||
"final_optimizer_state": replay["hashes"][
|
||
"final_optimizer_state"
|
||
],
|
||
},
|
||
"aggregate_sha256": full["canonical_sha256_without_self"],
|
||
"compact_sha256": compact["canonical_sha256_without_self"],
|
||
"overall_verdict": overall_verdict,
|
||
}
|
||
reproduction["canonical_sha256_without_self"] = canonical_sha256(
|
||
reproduction
|
||
)
|
||
|
||
atomic_json(args.output, full)
|
||
atomic_json(args.compact_output, compact)
|
||
atomic_json(args.reproduction_output, reproduction)
|
||
print(
|
||
json.dumps(
|
||
{
|
||
"formal_runs": len(runs),
|
||
"formal_target_bytes": EXPECTED_TOTAL_TARGET_BYTES,
|
||
"smoke_exact": all(
|
||
row["exact"] for row in smoke_exact.values()
|
||
),
|
||
"replay_exact": replay_exact,
|
||
"depth_verdicts": {
|
||
depth: depth_summaries[str(depth)]["verdict"]["label"]
|
||
for depth in DEPTHS
|
||
},
|
||
"overall_verdict": overall_verdict,
|
||
"aggregate_sha256": full[
|
||
"canonical_sha256_without_self"
|
||
],
|
||
"compact_sha256": compact[
|
||
"canonical_sha256_without_self"
|
||
],
|
||
"reproduction_sha256": reproduction[
|
||
"canonical_sha256_without_self"
|
||
],
|
||
},
|
||
ensure_ascii=False,
|
||
indent=2,
|
||
)
|
||
)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|