wuyang
537 lines
19 KiB
Python
537 lines
19 KiB
Python
#!/usr/bin/env python3
|
|
"""Validate, aggregate, and compact the reduced Attention Residuals study."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import math
|
|
import os
|
|
import statistics
|
|
from pathlib import Path
|
|
from typing import Any, Iterable
|
|
|
|
|
|
PROTOCOL_ID = "llm-atlas-k3-attnres-reduced-v1"
|
|
ARCHITECTURES = ("baseline", "full", "block")
|
|
SEEDS = (2026073001, 2026073002, 2026073003)
|
|
REPLAY_FIELDS = (
|
|
"manifest",
|
|
"model",
|
|
"optimizer",
|
|
"hashes",
|
|
"evaluations",
|
|
"training_history",
|
|
"diagnostic",
|
|
"environment",
|
|
)
|
|
SMOKE_FIELDS = REPLAY_FIELDS
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--formal-dir", type=Path, required=True)
|
|
parser.add_argument("--smoke-dir", type=Path, required=True)
|
|
parser.add_argument("--replay", type=Path, required=True)
|
|
parser.add_argument("--manifest", 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 write_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 elementwise_summary(rows: list[list[float]]) -> dict[str, list[float]]:
|
|
length = len(rows[0])
|
|
if any(len(row) != length for row in rows):
|
|
raise ValueError("array lengths do not match")
|
|
return {
|
|
"mean": [mean(row[index] for row in rows) for index in range(length)],
|
|
"min": [min(row[index] for row in rows) for index in range(length)],
|
|
"max": [max(row[index] for row in rows) for index in range(length)],
|
|
}
|
|
|
|
|
|
def coefficient_of_variation(values: list[float]) -> float:
|
|
average = mean(values)
|
|
variance = mean((value - average) ** 2 for value in values)
|
|
return math.sqrt(variance) / average
|
|
|
|
|
|
def verdict(deltas: list[float]) -> dict[str, Any]:
|
|
average = mean(deltas)
|
|
if all(delta < 0 for delta in deltas) and average <= -0.010:
|
|
label = "directional support in this reduced protocol"
|
|
elif all(delta > 0 for delta in deltas) and average >= 0.010:
|
|
label = "directional concern in this reduced protocol"
|
|
else:
|
|
label = "inconclusive at this budget"
|
|
return {
|
|
"paired_deltas_bpc": deltas,
|
|
"mean_delta_bpc": average,
|
|
"min_delta_bpc": min(deltas),
|
|
"max_delta_bpc": max(deltas),
|
|
"same_direction": len({delta < 0 for delta in deltas}) == 1,
|
|
"threshold_bpc": 0.010,
|
|
"verdict": label,
|
|
}
|
|
|
|
|
|
def average_depth_weights(runs: list[dict[str, Any]]) -> dict[str, Any]:
|
|
rows_by_run = [run["diagnostic"]["depth_weights"] for run in runs]
|
|
layer_count = len(rows_by_run[0])
|
|
if any(len(rows) != layer_count for rows in rows_by_run):
|
|
raise ValueError("depth-weight layer counts differ")
|
|
rows = []
|
|
max_sources = 0
|
|
for layer in range(layer_count):
|
|
source_count = rows_by_run[0][layer]["sources"]
|
|
if any(rows[layer]["sources"] != source_count for rows in rows_by_run):
|
|
raise ValueError("source count differs across seeds")
|
|
weights = [
|
|
mean(rows_by_run[seed_index][layer]["mean_weights"][source]
|
|
for seed_index in range(len(runs)))
|
|
for source in range(source_count)
|
|
]
|
|
entropies = [
|
|
rows_by_run[seed_index][layer]["entropy_mean"]
|
|
for seed_index in range(len(runs))
|
|
]
|
|
rows.append(
|
|
{
|
|
"sublayer": layer + 1,
|
|
"sources": source_count,
|
|
"mean_weights": weights,
|
|
"entropy_mean": mean(entropies),
|
|
"entropy_min": min(entropies),
|
|
"entropy_max": max(entropies),
|
|
}
|
|
)
|
|
max_sources = max(max_sources, source_count)
|
|
|
|
output_rows = [run["diagnostic"]["output_weights"] for run in runs]
|
|
output_source_count = output_rows[0]["sources"]
|
|
output_weights = [
|
|
mean(row["mean_weights"][source] for row in output_rows)
|
|
for source in range(output_source_count)
|
|
]
|
|
return {
|
|
"rows": rows,
|
|
"max_sources": max_sources,
|
|
"output": {
|
|
"sources": output_source_count,
|
|
"mean_weights": output_weights,
|
|
"entropy_mean": mean(row["entropy_mean"] for row in output_rows),
|
|
"entropy_min": min(row["entropy_mean"] for row in output_rows),
|
|
"entropy_max": max(row["entropy_mean"] for row in output_rows),
|
|
},
|
|
}
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
manifest = read_json(args.manifest)
|
|
if manifest["protocol_id"] != PROTOCOL_ID:
|
|
raise ValueError("manifest protocol mismatch")
|
|
|
|
runs: dict[tuple[str, int], dict[str, Any]] = {}
|
|
formal_file_hashes: dict[str, str] = {}
|
|
for seed in SEEDS:
|
|
for architecture in ARCHITECTURES:
|
|
path = args.formal_dir / f"{architecture}-{seed}.json"
|
|
run = read_json(path)
|
|
if run["protocol_id"] != PROTOCOL_ID:
|
|
raise ValueError(f"protocol mismatch: {path}")
|
|
if run["run_kind"] != "formal":
|
|
raise ValueError(f"not a formal run: {path}")
|
|
if run["architecture"] != architecture or run["seed"] != seed:
|
|
raise ValueError(f"cell identity mismatch: {path}")
|
|
if run["steps"] != 2000 or run["batch_size"] != 32:
|
|
raise ValueError(f"formal budget mismatch: {path}")
|
|
if run["target_bytes_seen"] != 16_384_000:
|
|
raise ValueError(f"target byte count mismatch: {path}")
|
|
if run["manifest"]["file_sha256"] != file_sha256(args.manifest):
|
|
raise ValueError(f"manifest file hash mismatch: {path}")
|
|
if run["manifest"]["formal_schedule_sha256"] != (
|
|
manifest["windows"]["formal_schedule_sha256"]
|
|
):
|
|
raise ValueError(f"schedule mismatch: {path}")
|
|
if run["evaluations"][-1]["step"] != 2000:
|
|
raise ValueError(f"missing final evaluation: {path}")
|
|
if any(
|
|
not math.isfinite(value)
|
|
for evaluation in run["evaluations"]
|
|
for value in (
|
|
evaluation["cross_entropy_nats"],
|
|
evaluation["bits_per_byte"],
|
|
)
|
|
):
|
|
raise ValueError(f"non-finite evaluation: {path}")
|
|
if len(run["diagnostic"]["layer_input_rms"]) != 32:
|
|
raise ValueError(f"diagnostic depth mismatch: {path}")
|
|
if len(run["diagnostic"]["core_parameter_grad_rms_by_block"]) != 16:
|
|
raise ValueError(f"gradient depth mismatch: {path}")
|
|
runs[(architecture, seed)] = run
|
|
formal_file_hashes[path.name] = file_sha256(path)
|
|
|
|
common_initial_exact = {}
|
|
for seed in SEEDS:
|
|
hashes = {
|
|
architecture: runs[(architecture, seed)]["hashes"][
|
|
"initial_common_parameters"
|
|
]
|
|
for architecture in ARCHITECTURES
|
|
}
|
|
common_initial_exact[str(seed)] = {
|
|
"hashes": hashes,
|
|
"exact": len(set(hashes.values())) == 1,
|
|
}
|
|
if not common_initial_exact[str(seed)]["exact"]:
|
|
raise ValueError(f"common initialization mismatch for seed {seed}")
|
|
|
|
by_seed = []
|
|
for seed in SEEDS:
|
|
values = {
|
|
architecture: runs[(architecture, seed)]["evaluations"][-1][
|
|
"bits_per_byte"
|
|
]
|
|
for architecture in ARCHITECTURES
|
|
}
|
|
by_seed.append(
|
|
{
|
|
"seed": seed,
|
|
"final_bpc": values,
|
|
"full_minus_baseline": values["full"] - values["baseline"],
|
|
"block_minus_baseline": values["block"] - values["baseline"],
|
|
"block_minus_full": values["block"] - values["full"],
|
|
}
|
|
)
|
|
|
|
final = {
|
|
"by_seed": by_seed,
|
|
"means": {
|
|
architecture: mean(
|
|
runs[(architecture, seed)]["evaluations"][-1]["bits_per_byte"]
|
|
for seed in SEEDS
|
|
)
|
|
for architecture in ARCHITECTURES
|
|
},
|
|
"full_contrast": verdict(
|
|
[row["full_minus_baseline"] for row in by_seed]
|
|
),
|
|
"block_contrast": verdict(
|
|
[row["block_minus_baseline"] for row in by_seed]
|
|
),
|
|
"block_minus_full": {
|
|
"paired_deltas_bpc": [row["block_minus_full"] for row in by_seed],
|
|
"mean_delta_bpc": mean(row["block_minus_full"] for row in by_seed),
|
|
},
|
|
}
|
|
|
|
evaluation_steps = [
|
|
evaluation["step"] for evaluation in runs[("baseline", SEEDS[0])]["evaluations"]
|
|
]
|
|
curves = {}
|
|
for architecture in ARCHITECTURES:
|
|
curve = []
|
|
for index, step in enumerate(evaluation_steps):
|
|
values = [
|
|
runs[(architecture, seed)]["evaluations"][index]["bits_per_byte"]
|
|
for seed in SEEDS
|
|
]
|
|
if any(
|
|
runs[(architecture, seed)]["evaluations"][index]["step"] != step
|
|
for seed in SEEDS
|
|
):
|
|
raise ValueError("evaluation step mismatch")
|
|
curve.append(
|
|
{
|
|
"step": step,
|
|
"mean_bpc": mean(values),
|
|
"min_bpc": min(values),
|
|
"max_bpc": max(values),
|
|
"by_seed": values,
|
|
}
|
|
)
|
|
curves[architecture] = curve
|
|
|
|
timing = {}
|
|
for architecture in ARCHITECTURES:
|
|
cells = [runs[(architecture, seed)]["timing"] for seed in SEEDS]
|
|
timing[architecture] = {
|
|
"mean_step_ms": mean(cell["mean_ms"] for cell in cells),
|
|
"median_step_ms": mean(cell["median_ms"] for cell in cells),
|
|
"p95_step_ms": mean(cell["p95_ms"] for cell in cells),
|
|
"mean_peak_allocated_bytes": mean(
|
|
cell["peak_allocated_bytes"] for cell in cells
|
|
),
|
|
"mean_peak_reserved_bytes": mean(
|
|
cell["peak_reserved_bytes"] for cell in cells
|
|
),
|
|
"by_seed": cells,
|
|
}
|
|
timing["relative_to_baseline"] = {
|
|
architecture: {
|
|
"step_time_ratio": timing[architecture]["mean_step_ms"]
|
|
/ timing["baseline"]["mean_step_ms"],
|
|
"allocated_memory_ratio": timing[architecture][
|
|
"mean_peak_allocated_bytes"
|
|
]
|
|
/ timing["baseline"]["mean_peak_allocated_bytes"],
|
|
}
|
|
for architecture in ("full", "block")
|
|
}
|
|
|
|
parameters = {
|
|
architecture: runs[(architecture, SEEDS[0])]["model"]["parameters"]
|
|
for architecture in ARCHITECTURES
|
|
}
|
|
parameters["mixer_overhead_fraction_of_baseline"] = (
|
|
parameters["full"]["mixer"] / parameters["baseline"]["total"]
|
|
)
|
|
|
|
traces = {}
|
|
gradients = {}
|
|
for architecture in ARCHITECTURES:
|
|
architecture_runs = [runs[(architecture, seed)] for seed in SEEDS]
|
|
traces[architecture] = {
|
|
key: elementwise_summary(
|
|
[run["diagnostic"][key] for run in architecture_runs]
|
|
)
|
|
for key in (
|
|
"layer_input_rms",
|
|
"branch_output_rms",
|
|
"stream_state_rms",
|
|
)
|
|
}
|
|
gradient_rows = [
|
|
run["diagnostic"]["core_parameter_grad_rms_by_block"]
|
|
for run in architecture_runs
|
|
]
|
|
gradients[architecture] = {
|
|
"by_block": elementwise_summary(gradient_rows),
|
|
"cv_by_seed": [
|
|
coefficient_of_variation(row) for row in gradient_rows
|
|
],
|
|
"mean_cv": mean(coefficient_of_variation(row) for row in gradient_rows),
|
|
"first_last_ratio_by_seed": [
|
|
row[0] / row[-1] for row in gradient_rows
|
|
],
|
|
"mean_first_last_ratio": mean(row[0] / row[-1] for row in gradient_rows),
|
|
}
|
|
|
|
mixers = {
|
|
architecture: average_depth_weights(
|
|
[runs[(architecture, seed)] for seed in SEEDS]
|
|
)
|
|
for architecture in ("full", "block")
|
|
}
|
|
|
|
full_branch = traces["full"]["branch_output_rms"]["mean"]
|
|
largest_index = max(range(len(full_branch)), key=full_branch.__getitem__)
|
|
# Full output source 0 is the embedding; branch l is source l+1.
|
|
largest_source_weight = mixers["full"]["output"]["mean_weights"][
|
|
largest_index + 1
|
|
]
|
|
uniform_output_weight = 1 / mixers["full"]["output"]["sources"]
|
|
posthoc = {
|
|
"label": "post-hoc descriptive callout; not a preregistered endpoint",
|
|
"largest_full_branch_sublayer": largest_index + 1,
|
|
"largest_full_branch_rms": full_branch[largest_index],
|
|
"corresponding_final_output_weight": largest_source_weight,
|
|
"uniform_final_output_weight": uniform_output_weight,
|
|
"weight_over_uniform": largest_source_weight / uniform_output_weight,
|
|
}
|
|
|
|
replay = read_json(args.replay)
|
|
formal_replay_source = runs[("block", 2026073001)]
|
|
replay_exact = {
|
|
field: formal_replay_source[field] == replay[field]
|
|
for field in REPLAY_FIELDS
|
|
}
|
|
if not all(replay_exact.values()):
|
|
raise ValueError(f"formal replay mismatch: {replay_exact}")
|
|
|
|
smoke = {}
|
|
for architecture in ARCHITECTURES:
|
|
first_path = args.smoke_dir / f"{architecture}-2026073001-a.json"
|
|
second_path = args.smoke_dir / f"{architecture}-2026073001-b.json"
|
|
first = read_json(first_path)
|
|
second = read_json(second_path)
|
|
exact = {field: first[field] == second[field] for field in SMOKE_FIELDS}
|
|
if not all(exact.values()):
|
|
raise ValueError(f"smoke mismatch for {architecture}: {exact}")
|
|
smoke[architecture] = {
|
|
"fields": exact,
|
|
"all_exact": True,
|
|
"first_sha256": file_sha256(first_path),
|
|
"second_sha256": file_sha256(second_path),
|
|
}
|
|
|
|
reproduction = {
|
|
"schema_version": 1,
|
|
"protocol_id": PROTOCOL_ID,
|
|
"manifest_sha256": file_sha256(args.manifest),
|
|
"formal_files": formal_file_hashes,
|
|
"common_initial_parameters": common_initial_exact,
|
|
"smoke": smoke,
|
|
"formal_replay": {
|
|
"architecture": "block",
|
|
"seed": 2026073001,
|
|
"fields": replay_exact,
|
|
"all_numeric_and_hash_fields_exact": all(replay_exact.values()),
|
|
"timing_exact_required": False,
|
|
"timing_exact_observed": formal_replay_source["timing"] == replay["timing"],
|
|
"formal_file_sha256": formal_file_hashes[
|
|
"block-2026073001.json"
|
|
],
|
|
"replay_file_sha256": file_sha256(args.replay),
|
|
},
|
|
}
|
|
reproduction["canonical_sha256_without_self"] = canonical_sha256(reproduction)
|
|
|
|
analysis = {
|
|
"final_validation": final,
|
|
"evaluation_curves": curves,
|
|
"timing": timing,
|
|
"parameters": parameters,
|
|
"traces": traces,
|
|
"gradients": gradients,
|
|
"mixers": mixers,
|
|
"posthoc": posthoc,
|
|
"interpretation": {
|
|
"primary": (
|
|
"Both Full and Block AttnRes satisfy the preregistered "
|
|
"directional-support rule in this reduced protocol."
|
|
),
|
|
"bounded_depth_pattern": (
|
|
"Block partial-state RMS resets every four residual sublayers; "
|
|
"the complete 32-point vectors are reported."
|
|
),
|
|
"gradient_boundary": (
|
|
"The preregistered core-parameter gradient RMS is not more "
|
|
"uniform for AttnRes here; this metric and scale do not reproduce "
|
|
"the paper's large-model gradient-magnitude result."
|
|
),
|
|
"scope": (
|
|
"Reduced byte-level WikiText-2 mechanism probe; not a K3 "
|
|
"checkpoint run, paper-scale reproduction, benchmark, or "
|
|
"same-FLOP comparison."
|
|
),
|
|
},
|
|
}
|
|
|
|
raw = {
|
|
"schema_version": 1,
|
|
"protocol_id": PROTOCOL_ID,
|
|
"manifest": manifest,
|
|
"provenance": {
|
|
"manifest_file_sha256": file_sha256(args.manifest),
|
|
"formal_file_sha256": formal_file_hashes,
|
|
"reproduction_sha256": reproduction[
|
|
"canonical_sha256_without_self"
|
|
],
|
|
},
|
|
"formal_runs": [
|
|
runs[(architecture, seed)]
|
|
for seed in SEEDS
|
|
for architecture in ARCHITECTURES
|
|
],
|
|
"analysis": analysis,
|
|
"reproduction": reproduction,
|
|
}
|
|
raw["canonical_sha256_without_self"] = canonical_sha256(raw)
|
|
|
|
compact = {
|
|
"schema_version": 1,
|
|
"protocol_id": PROTOCOL_ID,
|
|
"dataset": {
|
|
"repository": manifest["dataset"]["repository"],
|
|
"revision": manifest["dataset"]["revision"],
|
|
"train_bytes": manifest["dataset"]["splits"]["train"][
|
|
"concatenated_bytes"
|
|
],
|
|
"schedule_sha256": manifest["windows"]["formal_schedule_sha256"],
|
|
"validation_sha256": manifest["windows"][
|
|
"validation_tensor_sha256"
|
|
],
|
|
},
|
|
"grid": {
|
|
"architectures": list(ARCHITECTURES),
|
|
"seeds": list(SEEDS),
|
|
"runs": 9,
|
|
"steps_per_run": 2000,
|
|
"target_bytes_per_run": 16_384_000,
|
|
"target_bytes_total": 9 * 16_384_000,
|
|
},
|
|
"final_validation": final,
|
|
"evaluation_curves": curves,
|
|
"timing": timing,
|
|
"parameters": parameters,
|
|
"traces": traces,
|
|
"gradients": gradients,
|
|
"mixers": mixers,
|
|
"posthoc": posthoc,
|
|
"interpretation": analysis["interpretation"],
|
|
"reproduction": reproduction,
|
|
"source_sha256": raw["canonical_sha256_without_self"],
|
|
}
|
|
compact["canonical_sha256_without_self"] = canonical_sha256(compact)
|
|
|
|
write_json(args.reproduction_output, reproduction)
|
|
write_json(args.output, raw)
|
|
write_json(args.compact_output, compact)
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"output": str(args.output),
|
|
"compact_output": str(args.compact_output),
|
|
"reproduction_output": str(args.reproduction_output),
|
|
"raw_sha256": file_sha256(args.output),
|
|
"compact_sha256": file_sha256(args.compact_output),
|
|
"reproduction_sha256": file_sha256(args.reproduction_output),
|
|
"full": final["full_contrast"],
|
|
"block": final["block_contrast"],
|
|
"formal_replay": reproduction["formal_replay"],
|
|
},
|
|
ensure_ascii=False,
|
|
indent=2,
|
|
)
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|