Files

700 lines
24 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
#!/usr/bin/env python3
"""Aggregate and gate preregistered Round 08 forward-training results."""
from __future__ import annotations
import argparse
import copy
import hashlib
import json
import math
import statistics
from pathlib import Path
from typing import Any, Iterable
PROTOCOL_ID = "llm-atlas-k3-attnres-forward-training-v1"
PARENT_PROTOCOL_ID = "llm-atlas-k3-attnres-gradient-scale-v1"
METRICS = ("spike_contrast", "peak_normalized")
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser()
parser.add_argument("--manifest", type=Path, required=True)
parser.add_argument("--formal", type=Path, action="append", required=True)
parser.add_argument("--replay", type=Path, required=True)
parser.add_argument("--reference-dir", type=Path, required=True)
parser.add_argument("--aggregate-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 canonical_sha256(value: Any) -> str:
return hashlib.sha256(
json.dumps(
value, ensure_ascii=False, sort_keys=True, separators=(",", ":")
).encode()
).hexdigest()
def file_sha256(path: Path) -> str:
digest = hashlib.sha256()
with path.open("rb") as handle:
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
digest.update(chunk)
return digest.hexdigest()
def mean(values: Iterable[float]) -> float:
return statistics.fmean(values)
def read_result(path: Path, expected_protocol: str) -> dict[str, Any]:
value = json.loads(path.read_text())
if value.get("protocol_id") != expected_protocol:
raise RuntimeError(f"protocol mismatch: {path}")
expected = value.get("canonical_sha256_without_self")
payload = {
key: item
for key, item in value.items()
if key != "canonical_sha256_without_self"
}
if not isinstance(expected, str) or canonical_sha256(payload) != expected:
raise RuntimeError(f"canonical self-hash mismatch: {path}")
return value
def exactly_one(values: list[dict[str, Any]], step: int) -> dict[str, Any]:
matches = [value for value in values if value["step"] == step]
if len(matches) != 1:
raise RuntimeError(f"step {step} missing or duplicated")
return matches[0]
def spectrum_metrics(
diagnostic: dict[str, Any],
spike_layers: tuple[int, ...],
epsilon: float,
) -> dict[str, Any]:
values = [
float(value)
for value in diagnostic["activation_grad_rms_by_block"]
]
if len(values) != 32:
raise RuntimeError("activation-gradient spectrum must have 32 layers")
if any(not math.isfinite(value) or value <= epsilon for value in values):
raise RuntimeError("activation-gradient spectrum is non-finite/non-positive")
spike_indices = {layer - 1 for layer in spike_layers}
spike_values = [
value for index, value in enumerate(values) if index in spike_indices
]
reference_values = [
value for index, value in enumerate(values) if index not in spike_indices
]
spike_mean = mean(spike_values)
reference_mean = mean(reference_values)
global_mean = mean(values)
contrast = spike_mean / reference_mean
peak = max(values) / global_mean
if any(
not math.isfinite(value) or value <= epsilon
for value in (spike_mean, reference_mean, contrast, peak)
):
raise RuntimeError("derived spike metric is non-finite/non-positive")
ordered = sorted(range(32), key=lambda index: (-values[index], index))
return {
"values": values,
"normalized": [value / global_mean for value in values],
"spike_mean": spike_mean,
"reference_mean": reference_mean,
"global_mean": global_mean,
"spike_contrast": contrast,
"peak_normalized": peak,
"peak_layer_1based": ordered[0] + 1,
"top_five_layers_1based": [index + 1 for index in ordered[:5]],
}
def final_bpc(value: dict[str, Any], step: int) -> float:
result = float(exactly_one(value["evaluations"], step)["bits_per_byte"])
if not math.isfinite(result):
raise RuntimeError("final validation BPC is non-finite")
return result
def stable_environment(value: dict[str, Any]) -> dict[str, Any]:
keys = (
"cublas_workspace_config",
"deterministic_algorithms",
"autocast",
"compile",
)
return {key: value["environment"][key] for key in keys}
def pairing_checks(
run: dict[str, Any], reference: dict[str, Any]
) -> dict[str, bool]:
manifest_fields = (
"formal_schedule_sha256",
"validation_tensor_sha256",
"diagnostic_tensor_sha256",
"input_gate_tensor_hashes",
)
checks = {
"seed": run["seed"] == reference["seed"],
"architecture": (
run["architecture"] == reference["architecture"] == "block"
),
"depth": run["depth"] == reference["depth"] == 32,
"steps": run["steps"] == reference["steps"] == 8000,
"batch_size": run["batch_size"] == reference["batch_size"] == 32,
"initial_public_parameters": (
run["hashes"]["initial_public_parameters"]
== reference["hashes"]["initial_public_parameters"]
),
"initial_mixer_parameters": (
run["hashes"]["initial_mixer_parameters"]
== reference["hashes"]["initial_mixer_parameters"]
),
"model_topology": run["model"] == reference["model"],
"optimizer_hyperparameters": (
run["optimizer"] == reference["optimizer"]
),
"scientific_environment": (
stable_environment(run) == stable_environment(reference)
),
}
for field in manifest_fields:
checks[f"manifest.{field}"] = (
run["manifest"][field] == reference["manifest"][field]
)
return checks
def scientific_replay_payload(value: dict[str, Any]) -> dict[str, Any]:
payload = copy.deepcopy(value)
for key in (
"run_kind",
"timing",
"canonical_sha256_without_self",
"parent_runner_canonical_sha256",
):
payload.pop(key, None)
payload["manifest"].pop("path", None)
payload["study_manifest"].pop("path", None)
payload["environment"] = stable_environment(value)
return payload
def quality_gate(
variant_runs: dict[int, dict[str, Any]],
references: dict[int, dict[str, Any]],
*,
step: int,
per_seed_maximum: float,
mean_maximum: float,
) -> dict[str, Any]:
per_seed = {}
for seed, run in sorted(variant_runs.items()):
variant_bpc = final_bpc(run, step)
reference_bpc = final_bpc(references[seed], step)
delta = variant_bpc - reference_bpc
per_seed[str(seed)] = {
"variant_bpc": variant_bpc,
"reference_bpc": reference_bpc,
"delta_bpc": delta,
"passed": delta <= per_seed_maximum,
}
mean_delta = mean(item["delta_bpc"] for item in per_seed.values())
per_seed_passed = all(item["passed"] for item in per_seed.values())
mean_passed = mean_delta <= mean_maximum
return {
"passed": per_seed_passed and mean_passed,
"passed_checks": (
sum(item["passed"] for item in per_seed.values())
+ int(mean_passed)
),
"required_checks": 4,
"per_seed_maximum": per_seed_maximum,
"mean_maximum": mean_maximum,
"mean_delta_bpc": mean_delta,
"mean_passed": mean_passed,
"per_seed": per_seed,
}
def variant_effect(
variant: str,
runs: dict[int, dict[str, Any]],
references: dict[int, dict[str, Any]],
metrics_by_cell: dict[tuple[str, int, int], dict[str, Any]],
*,
step: int,
threshold: float,
quality: dict[str, Any],
) -> dict[str, Any]:
cells = []
for seed in sorted(runs):
candidate = metrics_by_cell[(variant, seed, step)]
reference = metrics_by_cell[("learned_reference", seed, step)]
for metric in METRICS:
reference_value = reference[metric]
candidate_value = candidate[metric]
relative_drop = (
reference_value - candidate_value
) / reference_value
cells.append(
{
"seed": seed,
"metric": metric,
"reference": reference_value,
"variant": candidate_value,
"relative_drop": relative_drop,
"passed": relative_drop >= threshold,
}
)
attenuation_passed = all(cell["passed"] for cell in cells)
return {
"variant": variant,
"threshold": threshold,
"passed_cells": sum(cell["passed"] for cell in cells),
"required_cells": len(cells),
"attenuation_passed": attenuation_passed,
"quality": quality,
"material_response_passed": (
attenuation_passed and quality["passed"]
),
"cells": cells,
}
def interaction_map(
metrics_by_cell: dict[tuple[str, int, int], dict[str, Any]],
seeds: tuple[int, ...],
steps: tuple[int, ...],
) -> dict[str, Any]:
cells = []
for step in steps:
for seed in seeds:
reference = metrics_by_cell[
("learned_reference", seed, step)
]
group6 = metrics_by_cell[
("uniform_group_6_forward", seed, step)
]
group7 = metrics_by_cell[
("uniform_group_7_forward", seed, step)
]
joint = metrics_by_cell[
("uniform_groups_6_7_forward", seed, step)
]
for metric in METRICS:
ref = reference[metric]
effects = {
"group6": math.log(ref / group6[metric]),
"group7": math.log(ref / group7[metric]),
"groups6_7": math.log(ref / joint[metric]),
}
residual = (
effects["groups6_7"]
- effects["group6"]
- effects["group7"]
)
cells.append(
{
"step": step,
"seed": seed,
"metric": metric,
"log_effects": effects,
"interaction_residual": residual,
"relative_drops": {
"group6": (ref - group6[metric]) / ref,
"group7": (ref - group7[metric]) / ref,
"groups6_7": (ref - joint[metric]) / ref,
},
}
)
summaries = []
for step in steps:
for metric in METRICS:
selected = [
cell
for cell in cells
if cell["step"] == step and cell["metric"] == metric
]
residuals = [
cell["interaction_residual"] for cell in selected
]
summaries.append(
{
"step": step,
"metric": metric,
"mean_interaction_residual": mean(residuals),
"minimum": min(residuals),
"maximum": max(residuals),
}
)
return {
"definition": "I67=ln(Xref/X67)-ln(Xref/X6)-ln(Xref/X7)",
"interpretation": (
"descriptive cross-run log-attenuation residual from three "
"independently trained variants; not a causal interaction"
),
"cells": cells,
"summaries": summaries,
}
def environment_metadata(value: dict[str, Any]) -> dict[str, Any]:
return {
key: value["environment"].get(key)
for key in ("gpu", "torch", "cuda", "compute_capability")
}
def write_hashed(path: Path, value: dict[str, Any]) -> None:
value["canonical_sha256_without_self"] = canonical_sha256(value)
path.parent.mkdir(parents=True, exist_ok=True)
path.write_text(
json.dumps(value, ensure_ascii=False, indent=2, sort_keys=True) + "\n"
)
def main() -> None:
args = parse_args()
manifest = json.loads(args.manifest.read_text())
if (
manifest["protocol_id"] != PROTOCOL_ID
or manifest["status"] != "frozen-before-model-output"
):
raise RuntimeError("manifest is not the frozen Round 08 contract")
variants = tuple(manifest["variants"].keys())
seeds = tuple(manifest["formal_seeds"])
steps = tuple(manifest["diagnostic_steps"])
primary_step = manifest["primary_step"]
epsilon = manifest["thresholds"]["positive_denominator_epsilon"]
spike_layers = tuple(manifest["spike_layers_1based"])
expected_cells = {(variant, seed) for variant in variants for seed in seeds}
if len(args.formal) != len(expected_cells):
raise RuntimeError("formal path count does not match the 4×3 matrix")
runs: dict[tuple[str, int], dict[str, Any]] = {}
run_paths: dict[tuple[str, int], Path] = {}
pairing: dict[str, Any] = {}
references: dict[int, dict[str, Any]] = {}
reference_paths: dict[int, Path] = {}
for seed in seeds:
path = args.reference_dir / (
f"formal-depth-32-block-seed-{seed}.json"
)
references[seed] = read_result(path, PARENT_PROTOCOL_ID)
reference_paths[seed] = path
for path in args.formal:
value = read_result(path, PROTOCOL_ID)
identity = (value["variant"], value["seed"])
if identity in runs:
raise RuntimeError(f"duplicate formal cell: {identity}")
if (
value["run_kind"] != "formal"
or value["steps"] != manifest["formal_steps"]
or not value["forward_intervention"]["passed"]
):
raise RuntimeError(f"invalid formal cell: {path}")
runs[identity] = value
run_paths[identity] = path
if set(runs) != expected_cells:
raise RuntimeError("formal matrix identities do not match manifest")
for (variant, seed), value in sorted(runs.items()):
checks = pairing_checks(value, references[seed])
if not all(checks.values()):
raise RuntimeError(
f"historical reference pairing failed: "
f"{variant}/{seed}: {checks}"
)
pairing[f"{variant}:{seed}"] = {
"passed": True,
"checks": checks,
"run_environment": environment_metadata(value),
"reference_environment": environment_metadata(references[seed]),
"metadata_equal": (
environment_metadata(value)
== environment_metadata(references[seed])
),
}
metadata_warnings = [
{
"cell": cell,
"message": (
"GPU/version metadata differs from the historical paired "
"reference; frozen scientific-environment fields still match"
),
"run_environment": item["run_environment"],
"reference_environment": item["reference_environment"],
}
for cell, item in pairing.items()
if not item["metadata_equal"]
]
replay = read_result(args.replay, PROTOCOL_ID)
replay_contract = manifest["replay"]
if (
replay["run_kind"] != "replay"
or replay["variant"] != replay_contract["variant"]
or replay["seed"] != replay_contract["seed"]
or replay["steps"] != manifest["formal_steps"]
or not replay["forward_intervention"]["passed"]
):
raise RuntimeError("invalid replay identity/audit")
formal_primary = runs[
(replay_contract["variant"], replay_contract["seed"])
]
formal_payload = scientific_replay_payload(formal_primary)
replay_payload = scientific_replay_payload(replay)
replay_exact = formal_payload == replay_payload
if not replay_exact:
raise RuntimeError("primary formal/replay scientific payload mismatch")
metrics_by_cell: dict[tuple[str, int, int], dict[str, Any]] = {}
for seed, reference in references.items():
for step in steps:
metrics_by_cell[("learned_reference", seed, step)] = (
spectrum_metrics(
exactly_one(reference["diagnostics"], step),
spike_layers,
epsilon,
)
)
for (variant, seed), value in runs.items():
if tuple(item["step"] for item in value["diagnostics"]) != steps:
raise RuntimeError(f"diagnostic schedule drift: {variant}/{seed}")
for step in steps:
metrics_by_cell[(variant, seed, step)] = spectrum_metrics(
exactly_one(value["diagnostics"], step),
spike_layers,
epsilon,
)
runs_by_variant = {
variant: {seed: runs[(variant, seed)] for seed in seeds}
for variant in variants
}
qualities = {
variant: quality_gate(
variant_runs,
references,
step=primary_step,
per_seed_maximum=manifest["thresholds"][
"final_bpc_delta_per_seed_maximum"
],
mean_maximum=manifest["thresholds"][
"final_bpc_delta_mean_maximum"
],
)
for variant, variant_runs in runs_by_variant.items()
}
effects = {
variant: variant_effect(
variant,
variant_runs,
references,
metrics_by_cell,
step=primary_step,
threshold=manifest["thresholds"]["material_relative_drop"],
quality=qualities[variant],
)
for variant, variant_runs in runs_by_variant.items()
}
primary = effects[manifest["primary_variant"]]
if primary["attenuation_passed"] and primary["quality"]["passed"]:
status = (
"forward_training_attenuation_established_within_reduced_protocol"
)
elif primary["attenuation_passed"]:
status = "quality_guard_failed"
elif primary["quality"]["passed"]:
status = "attenuation_not_established"
else:
status = "attenuation_and_quality_failed"
secondary = {
variant: (
"secondary_material_response"
if effect["material_response_passed"]
else "secondary_response_not_established"
)
for variant, effect in effects.items()
if variant != manifest["primary_variant"]
}
interaction = interaction_map(metrics_by_cell, seeds, steps)
trajectories = []
final_spectra = []
for variant in ("learned_reference",) + variants:
for seed in seeds:
for step in steps:
record = metrics_by_cell[(variant, seed, step)]
reference = metrics_by_cell[
("learned_reference", seed, step)
]
trajectories.append(
{
"variant": variant,
"seed": seed,
"step": step,
"spike_mean": record["spike_mean"],
"reference_mean": record["reference_mean"],
"spike_contrast": record["spike_contrast"],
"peak_normalized": record["peak_normalized"],
"relative_drop": {
metric: (
reference[metric] - record[metric]
)
/ reference[metric]
for metric in METRICS
},
}
)
final = metrics_by_cell[(variant, seed, primary_step)]
final_spectra.append(
{
"variant": variant,
"seed": seed,
**final,
}
)
input_files = {
"manifest": {
"path": str(args.manifest),
"sha256": file_sha256(args.manifest),
},
"formal": [
{
"variant": variant,
"seed": seed,
"path": str(run_paths[(variant, seed)]),
"sha256": file_sha256(run_paths[(variant, seed)]),
}
for variant, seed in sorted(runs)
],
"references": [
{
"seed": seed,
"path": str(reference_paths[seed]),
"sha256": file_sha256(reference_paths[seed]),
}
for seed in seeds
],
"replay": {
"path": str(args.replay),
"sha256": file_sha256(args.replay),
},
}
aggregate = {
"schema_version": 1,
"protocol_id": PROTOCOL_ID,
"status": status,
"scope": (
"depth-32 reduced Block AttnRes train-time architecture "
"ablation; not a real Kimi-K3 checkpoint result"
),
"primary_step": primary_step,
"spike_layers_1based": list(spike_layers),
"thresholds": manifest["thresholds"],
"input_files": input_files,
"historical_pairing": pairing,
"metadata_warnings": metadata_warnings,
"replay": {
"passed": replay_exact,
"scientific_payload_sha256": canonical_sha256(formal_payload),
"excluded": [
"run_kind",
"timing",
"self hashes",
"manifest path strings",
"GPU/version metadata",
],
},
"primary": primary,
"secondary_status": secondary,
"effects": effects,
"interaction": interaction,
"trajectories": trajectories,
"final_spectra": final_spectra,
"processed_target_bytes": manifest["new_target_bytes"],
"historical_reference_target_bytes": (
manifest["historical_reference_target_bytes"]
),
"reporting_boundary": (
"C can change through spike-window numerator and the 27-layer "
"reference denominator; layers 26-28 are intervened but belong "
"to the denominator."
),
}
write_hashed(args.aggregate_output, aggregate)
compact = {
"schema_version": 1,
"protocol_id": PROTOCOL_ID,
"status": status,
"primary_step": primary_step,
"spike_layers_1based": list(spike_layers),
"thresholds": manifest["thresholds"],
"primary": primary,
"secondary_status": secondary,
"effects": effects,
"interaction": interaction,
"trajectories": trajectories,
"final_spectra": final_spectra,
"replay": aggregate["replay"],
"metadata_warnings": metadata_warnings,
"processed_target_bytes": manifest["new_target_bytes"],
"reporting_boundary": aggregate["reporting_boundary"],
"aggregate_sha256": aggregate["canonical_sha256_without_self"],
}
write_hashed(args.compact_output, compact)
reproduction = {
"schema_version": 1,
"protocol_id": PROTOCOL_ID,
"passed": replay_exact,
"formal_variant": replay_contract["variant"],
"seed": replay_contract["seed"],
"formal_file_sha256": file_sha256(
run_paths[
(replay_contract["variant"], replay_contract["seed"])
]
),
"replay_file_sha256": file_sha256(args.replay),
"scientific_payload_sha256": canonical_sha256(formal_payload),
"excluded_fields": aggregate["replay"]["excluded"],
}
write_hashed(args.reproduction_output, reproduction)
print(
json.dumps(
{
"status": status,
"primary_attenuation": {
"passed_cells": primary["passed_cells"],
"required_cells": primary["required_cells"],
},
"primary_quality": {
"passed_checks": primary["quality"]["passed_checks"],
"required_checks": primary["quality"]["required_checks"],
},
"replay_exact": replay_exact,
"aggregate": str(args.aggregate_output),
"compact": str(args.compact_output),
},
ensure_ascii=False,
indent=2,
)
)
if __name__ == "__main__":
main()