research: audit AttnRes gradient scale study
This commit is contained in:
wuyang
@@ -0,0 +1,685 @@
|
||||
#!/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()
|
||||
Reference in New Issue
Block a user