From 39a9ad6215b3f821a34f9ae9e3425be42026b9af Mon Sep 17 00:00:00 2001 From: wuyang <5700876+banisherwy@user.noreply.gitee.com> Date: Thu, 30 Jul 2026 12:53:22 +0800 Subject: [PATCH] experiment: implement AttnRes local path runner --- experiments/k3/attnres_local_path/README.md | 32 +- experiments/k3/attnres_local_path/analyze.py | 498 ++++++++++ experiments/k3/attnres_local_path/train.py | 933 +++++++++++++++++++ 3 files changed, 1461 insertions(+), 2 deletions(-) create mode 100644 experiments/k3/attnres_local_path/analyze.py create mode 100644 experiments/k3/attnres_local_path/train.py diff --git a/experiments/k3/attnres_local_path/README.md b/experiments/k3/attnres_local_path/README.md index 82b0a64..d788c75 100644 --- a/experiments/k3/attnres_local_path/README.md +++ b/experiments/k3/attnres_local_path/README.md @@ -18,5 +18,33 @@ GPU NVIDIA GeForce RTX 5090 CUBLAS_WORKSPACE_CONFIG=:4096:8 ``` -The runner and exact commands will be added after this preregistration -manifest is committed. No result file may precede that commit. +The preregistration was committed as `6911efc` before the runner or any result +file existed. + +## Step-0 smoke + +```bash +CUBLAS_WORKSPACE_CONFIG=:4096:8 \ +/home/wuyang/.pyenv/versions/3.10.14/envs/navi-router-cu128/bin/python \ + experiments/k3/attnres_local_path/train.py \ + --run-kind smoke \ + --seed 2026073001 \ + --cache-dir /home/wuyang/.cache/llm-atlas/k3-attnres-gradient-scale-v1 \ + --data-manifest experiments/k3/attnres_gradient/manifest.json \ + --parent-manifest experiments/k3/attnres_spike/manifest.json \ + --manifest experiments/k3/attnres_local_path/manifest.json \ + --output /home/wuyang/.cache/llm-atlas/k3-attnres-local-path-v1/smoke/seed-2026073001.json +``` + +The first smoke passed all 14-mode forward-identity, selector, Round 06 endpoint, +initialization-negative-control, parent-learned, and loss-scale gates. Its +canonical content hash is +`f708200f4fb122f61f30b97393839382a2094a71a8ea5d48cf47fb7fa094e69b`. + +Formal cells use the same command with `--run-kind formal`, one of the three +manifest seeds, and a new output path. The independent seed-2026073001 run uses +`--run-kind replay`. + +Only `analyze.py` may calculate the global log gap, sufficiency/restoration +scores, and preregistered gates. The site consumes its frozen aggregate rather +than reimplementing thresholds in TypeScript. diff --git a/experiments/k3/attnres_local_path/analyze.py b/experiments/k3/attnres_local_path/analyze.py new file mode 100644 index 0000000..06b8383 --- /dev/null +++ b/experiments/k3/attnres_local_path/analyze.py @@ -0,0 +1,498 @@ +#!/usr/bin/env python3 +"""Aggregate and gate preregistered Round 07 local-path 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 + + +PROTOCOL_ID = "llm-atlas-k3-attnres-local-path-v1" +METRICS = ("spike_contrast", "peak_normalized") +SUFFICIENCY_MODES = ( + "uniform_group_6_only", + "uniform_group_7_only", + "uniform_groups_6_7_only", + "uniform_group_6_attention_only", + "uniform_group_6_mlp_only", + "uniform_group_7_attention_only", + "uniform_group_7_mlp_only", + "uniform_output_only", + "uniform_depth_all", + "uniform_all", +) +RESTORATION_MODES = ( + "uniform_except_group_6", + "uniform_except_group_7", + "uniform_except_groups_6_7", +) + + +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("--output", type=Path, required=True) + return parser.parse_args() + + +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 canonical_sha256(value: Any) -> str: + payload = json.dumps( + value, ensure_ascii=False, sort_keys=True, separators=(",", ":") + ).encode() + return hashlib.sha256(payload).hexdigest() + + +def read_result(path: Path) -> dict[str, Any]: + value = json.loads(path.read_text()) + if value["protocol_id"] != PROTOCOL_ID: + raise RuntimeError(f"protocol mismatch: {path}") + expected = value["canonical_sha256_without_self"] + without_self = { + key: item + for key, item in value.items() + if key != "canonical_sha256_without_self" + } + if canonical_sha256(without_self) != expected: + raise RuntimeError(f"canonical self-hash mismatch: {path}") + return value + + +def replay_payload(value: dict[str, Any]) -> dict[str, Any]: + cleaned = copy.deepcopy(value) + for key in ("run_kind", "timing", "canonical_sha256_without_self"): + cleaned.pop(key) + return cleaned + + +def primary_metrics(mode: dict[str, Any]) -> dict[str, float]: + stats = mode["positions"]["post_mlp_state"]["reductions"][ + "element_rms" + ]["statistics"] + return {name: float(stats[name]) for name in METRICS} + + +def all_true(values: list[bool]) -> bool: + return len(values) > 0 and all(values) + + +def mixed_status( + by_seed_metric: dict[str, dict[str, dict[str, Any]]], + threshold: float, +) -> dict[str, Any]: + passed = [] + signs = [] + metric_passes = {metric: [] for metric in METRICS} + for seed_values in by_seed_metric.values(): + for metric in METRICS: + score = seed_values[metric]["score"] + cell_passed = score is not None and score >= threshold + passed.append(cell_passed) + metric_passes[metric].append(cell_passed) + if score is not None: + signs.append(1 if score >= 0 else -1) + reasons = [] + if any(passed) and not all(passed): + reasons.append("seed_or_metric_pass_split") + if metric_passes[METRICS[0]] != metric_passes[METRICS[1]]: + reasons.append("metric_direction_split") + if len(set(signs)) > 1: + reasons.append("score_sign_split") + return { + "passed": all_true(passed), + "threshold": threshold, + "required_cells": len(passed), + "passed_cells": sum(passed), + "mixed": bool(reasons), + "mixed_reasons": reasons, + } + + +def main() -> None: + args = parse_args() + manifest = json.loads(args.manifest.read_text()) + if manifest["protocol_id"] != PROTOCOL_ID: + raise RuntimeError("manifest protocol mismatch") + if len(args.formal) != 3: + raise RuntimeError("exactly three formal result paths are required") + + formal_pairs = [(path, read_result(path)) for path in args.formal] + formal_pairs.sort(key=lambda item: item[1]["seed"]) + expected_seeds = manifest["formal_seeds"] + if [value["seed"] for _, value in formal_pairs] != expected_seeds: + raise RuntimeError("formal seeds do not match manifest") + for path, value in formal_pairs: + if value["run_kind"] != "formal" or value["steps"] != 8000: + raise RuntimeError(f"invalid formal cell: {path}") + if not value["round06_equivalence"]["passed"]: + raise RuntimeError(f"parent equivalence failed: {path}") + for diagnostic in value["diagnostics"]: + if not diagnostic["parent_learned_round06_exact"]: + raise RuntimeError(f"parent diagnostic mismatch: {path}") + if diagnostic["local_matrix"] is not None: + for gate in ("forward_identity_gate", "endpoint_exactness"): + if not diagnostic[gate]["passed"]: + raise RuntimeError(f"{gate} failed: {path}") + for mode in manifest["matrix_modes"]: + if not diagnostic["local_matrix"][mode]["selector"][ + "passed" + ]: + raise RuntimeError(f"selector failed: {path}:{mode}") + step0 = value["diagnostics"][0] + if ( + not step0["initialization_negative_control"]["passed"] + or not step0["loss_scale_gate"]["passed"] + ): + raise RuntimeError(f"step-0 control failed: {path}") + + replay = read_result(args.replay) + if ( + replay["run_kind"] != "replay" + or replay["seed"] != expected_seeds[0] + or replay["steps"] != 8000 + ): + raise RuntimeError("invalid replay cell") + replay_exact = ( + replay_payload(formal_pairs[0][1]) == replay_payload(replay) + ) + if not replay_exact: + raise RuntimeError("formal seed1 and replay are not canonical exact") + + thresholds = manifest["thresholds"] + cells = [] + sufficiency_by_mode: dict[str, dict[str, dict[str, Any]]] = { + mode: {} for mode in SUFFICIENCY_MODES + } + restoration_by_mode: dict[str, dict[str, dict[str, Any]]] = { + mode: {} for mode in RESTORATION_MODES + } + + for path, value in formal_pairs: + seed_key = str(value["seed"]) + final = value["diagnostics"][-1]["local_matrix"] + mode_metrics = { + mode: primary_metrics(final[mode]) + for mode in manifest["matrix_modes"] + } + global_metrics = {} + for metric in METRICS: + reference = mode_metrics["detached_learned"][metric] + uniform_all = mode_metrics["uniform_all"][metric] + log_gap = math.log(reference / uniform_all) + relative_drop = (reference - uniform_all) / reference + established = ( + reference + > thresholds["positive_denominator_epsilon"] + and uniform_all + > thresholds["positive_denominator_epsilon"] + and log_gap > 0 + and relative_drop + >= thresholds["global_relative_drop_minimum"] + ) + global_metrics[metric] = { + "reference": reference, + "uniform_all": uniform_all, + "log_gap": log_gap, + "relative_drop": relative_drop, + "established": established, + } + for mode in SUFFICIENCY_MODES: + sufficiency_by_mode[mode][seed_key] = {} + for metric in METRICS: + established = global_metrics[metric]["established"] + score = ( + math.log( + mode_metrics["detached_learned"][metric] + / mode_metrics[mode][metric] + ) + / global_metrics[metric]["log_gap"] + if established + else None + ) + sufficiency_by_mode[mode][seed_key][metric] = { + "score": score, + "metric_value": mode_metrics[mode][metric], + "global_gap_established": established, + } + for mode in RESTORATION_MODES: + restoration_by_mode[mode][seed_key] = {} + for metric in METRICS: + established = global_metrics[metric]["established"] + score = ( + math.log( + mode_metrics[mode][metric] + / mode_metrics["uniform_all"][metric] + ) + / global_metrics[metric]["log_gap"] + if established + else None + ) + restoration_by_mode[mode][seed_key][metric] = { + "score": score, + "metric_value": mode_metrics[mode][metric], + "global_gap_established": established, + } + cells.append( + { + "seed": value["seed"], + "path": str(path), + "file_sha256": file_sha256(path), + "canonical_sha256": value[ + "canonical_sha256_without_self" + ], + "mode_metrics": mode_metrics, + "global": global_metrics, + "interaction_residual": { + metric: ( + 1.0 + - sufficiency_by_mode["uniform_output_only"][ + seed_key + ][metric]["score"] + - sufficiency_by_mode["uniform_depth_all"][ + seed_key + ][metric]["score"] + ) + for metric in METRICS + }, + } + ) + + global_gap_passed = all_true( + [ + cell["global"][metric]["established"] + for cell in cells + for metric in METRICS + ] + ) + sufficiency_gates = { + "groups_6_7": mixed_status( + sufficiency_by_mode["uniform_groups_6_7_only"], + thresholds["groups_6_7_sufficiency_minimum"], + ), + "group_6": mixed_status( + sufficiency_by_mode["uniform_group_6_only"], + thresholds["single_group_material_minimum"], + ), + "group_7": mixed_status( + sufficiency_by_mode["uniform_group_7_only"], + thresholds["single_group_material_minimum"], + ), + "output_half_gap": mixed_status( + sufficiency_by_mode["uniform_output_only"], + thresholds["output_half_gap_minimum"], + ), + } + restoration_gates = { + "groups_6_7": mixed_status( + restoration_by_mode["uniform_except_groups_6_7"], + thresholds["groups_6_7_restoration_minimum"], + ), + "group_6": mixed_status( + restoration_by_mode["uniform_except_group_6"], + thresholds["single_group_material_minimum"], + ), + "group_7": mixed_status( + restoration_by_mode["uniform_except_group_7"], + thresholds["single_group_material_minimum"], + ), + } + + branch_gates = {} + for group in (6, 7): + group_passed = sufficiency_gates[f"group_{group}"]["passed"] + candidates = {} + for branch, sibling in (("attention", "mlp"), ("mlp", "attention")): + branch_mode = f"uniform_group_{group}_{branch}_only" + sibling_mode = f"uniform_group_{group}_{sibling}_only" + checks = [] + margins = [] + for seed in expected_seeds: + seed_key = str(seed) + for metric in METRICS: + left = sufficiency_by_mode[branch_mode][seed_key][ + metric + ]["score"] + right = sufficiency_by_mode[sibling_mode][seed_key][ + metric + ]["score"] + margin = ( + left - right + if left is not None and right is not None + else None + ) + margins.append(margin) + checks.append( + left is not None + and left >= thresholds["branch_material_minimum"] + and margin is not None + and margin + >= thresholds["branch_dominance_margin"] + ) + candidates[branch] = { + "passed": group_passed and all_true(checks), + "group_gate_passed": group_passed, + "passed_cells": sum(checks), + "required_cells": len(checks), + "margins": margins, + } + dominant = [ + branch + for branch, gate in candidates.items() + if gate["passed"] + ] + branch_gates[f"group_{group}"] = { + "passed": len(dominant) == 1, + "dominant_branch": dominant[0] if len(dominant) == 1 else None, + "exploratory_sufficiency_only": True, + "candidates": candidates, + } + + localization_passed = ( + global_gap_passed + and sufficiency_gates["groups_6_7"]["passed"] + and restoration_gates["groups_6_7"]["passed"] + ) + localization_status = ( + "established_at_preregistered_bidirectional_50pct_threshold" + if localization_passed + else ( + "one_sided_evidence_localization_not_established" + if ( + sufficiency_gates["groups_6_7"]["passed"] + != restoration_gates["groups_6_7"]["passed"] + ) + else "not_established_at_preregistered_threshold" + ) + ) + + mode_means = {} + for mode in manifest["matrix_modes"]: + mode_means[mode] = { + metric: statistics.fmean( + cell["mode_metrics"][mode][metric] for cell in cells + ) + for metric in METRICS + } + sufficiency_means = { + mode: { + metric: statistics.fmean( + sufficiency_by_mode[mode][str(seed)][metric]["score"] + for seed in expected_seeds + ) + for metric in METRICS + } + for mode in SUFFICIENCY_MODES + } + restoration_means = { + mode: { + metric: statistics.fmean( + restoration_by_mode[mode][str(seed)][metric]["score"] + for seed in expected_seeds + ) + for metric in METRICS + } + for mode in RESTORATION_MODES + } + + result = { + "schema_version": 1, + "protocol_id": PROTOCOL_ID, + "study_identity": manifest["study_identity"], + "manifest": { + "path": str(args.manifest), + "file_sha256": file_sha256(args.manifest), + }, + "formal_cells": cells, + "replay": { + "path": str(args.replay), + "file_sha256": file_sha256(args.replay), + "canonical_sha256": replay[ + "canonical_sha256_without_self" + ], + "formal_seed1_exact_excluding_run_kind_and_timing": replay_exact, + }, + "scores": { + "sufficiency": sufficiency_by_mode, + "restoration": restoration_by_mode, + }, + "means": { + "mode_metrics": mode_means, + "sufficiency": sufficiency_means, + "restoration": restoration_means, + }, + "gates": { + "all_input_and_parent_gates_passed": True, + "global_gap": { + "passed": global_gap_passed, + "required_cells": 6, + "passed_cells": sum( + cell["global"][metric]["established"] + for cell in cells + for metric in METRICS + ), + }, + "sufficiency": sufficiency_gates, + "restoration": restoration_gates, + "localization": { + "passed": localization_passed, + "status": localization_status, + "requires": ( + "groups 6+7 sufficiency and restoration both >=0.50 " + "for 3/3 seeds and both metrics" + ), + }, + "branch_dominance": branch_gates, + }, + "limitations": [ + "same-forward diagnostic backward-rule sensitivity only", + "reduced byte-level language model, not the Kimi K3 checkpoint", + "effects are non-additive and are not contribution percentages", + "group 7 includes layers 26-28 outside fixed spike set 21-25", + "three-seed threshold gates are not population inference", + ], + } + result["canonical_sha256_without_self"] = canonical_sha256(result) + args.output.parent.mkdir(parents=True, exist_ok=True) + temporary = args.output.with_suffix(args.output.suffix + ".tmp") + temporary.write_text( + json.dumps(result, ensure_ascii=False, indent=2, sort_keys=True) + + "\n" + ) + temporary.replace(args.output) + print( + json.dumps( + { + "output": str(args.output), + "formal_cells": len(cells), + "replay_exact": replay_exact, + "global_gap": global_gap_passed, + "localization": localization_status, + "canonical_sha256": result[ + "canonical_sha256_without_self" + ], + }, + ensure_ascii=False, + indent=2, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/experiments/k3/attnres_local_path/train.py b/experiments/k3/attnres_local_path/train.py new file mode 100644 index 0000000..b648559 --- /dev/null +++ b/experiments/k3/attnres_local_path/train.py @@ -0,0 +1,933 @@ +#!/usr/bin/env python3 +"""Exact-replay Round 06 with preregistered local mixer-path interventions.""" + +from __future__ import annotations + +import argparse +import importlib.util +import json +import math +import os +import platform +import statistics +import sys +import time +from pathlib import Path +from typing import Any + +import torch +import torch.nn as nn + + +PROTOCOL_ID = "llm-atlas-k3-attnres-local-path-v1" +PARENT_PROTOCOL_ID = "llm-atlas-k3-attnres-spike-path-v1" +DATA_PROTOCOL_ID = "llm-atlas-k3-attnres-gradient-scale-v1" +DEPTH = 32 +SEEDS = (2026073001, 2026073002, 2026073003) +FORMAL_STEPS = 8000 +PARENT_DIAGNOSTIC_STEPS = (0, 100, 500, 2000, 4000, 8000) +MATRIX_STEPS = (0, 8000) +MATRIX_MODES = ( + "detached_learned", + "uniform_group_6_only", + "uniform_group_7_only", + "uniform_groups_6_7_only", + "uniform_group_6_attention_only", + "uniform_group_6_mlp_only", + "uniform_group_7_attention_only", + "uniform_group_7_mlp_only", + "uniform_output_only", + "uniform_depth_all", + "uniform_all", + "uniform_except_group_6", + "uniform_except_group_7", + "uniform_except_groups_6_7", +) +TRAIN_BATCH_SIZE = 32 +VALIDATION_WINDOWS = 64 +EVAL_BATCH_SIZE = 8 +TIMING_WARMUP = 20 + + +def load_parent() -> Any: + path = Path(__file__).resolve().parents[1] / "attnres_spike" / "train.py" + spec = importlib.util.spec_from_file_location( + "k3_attnres_spike_parent", path + ) + if spec is None or spec.loader is None: + raise RuntimeError(f"cannot import parent runner from {path}") + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +parent = load_parent() + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument( + "--run-kind", choices=("smoke", "formal", "replay"), required=True + ) + parser.add_argument("--seed", type=int, required=True) + parser.add_argument("--cache-dir", type=Path, required=True) + parser.add_argument("--data-manifest", type=Path, required=True) + parser.add_argument("--parent-manifest", type=Path, required=True) + parser.add_argument("--manifest", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + if args.seed not in SEEDS: + parser.error(f"seed must be one of {SEEDS}") + if args.run_kind == "replay" and args.seed != SEEDS[0]: + parser.error(f"replay seed must be {SEEDS[0]}") + return args + + +class LocalPathLanguageModel(parent.SpikeLanguageModel): + """Parent model with a parallel, audited per-mixer coefficient selector.""" + + def __init__(self, architecture: str, manifest: dict[str, Any]): + super().__init__(architecture) + selector = manifest["selector"] + self.uniform_depth = { + mode: frozenset(indices) + for mode, indices in selector["uniform_depth_indices"].items() + } + self.uniform_output = selector["uniform_output"] + self.expected_uniform_counts = selector["expected_uniform_counts"] + self._active_selector_visits: list[dict[str, Any]] | None = None + self.last_selector_visits: list[dict[str, Any]] | None = None + + def forward( + self, + input_ids: torch.Tensor, + capture: bool = False, + mixer_backward_mode: str = "learned", + ) -> tuple[torch.Tensor, parent.SpikeTrace | None]: + if not capture: + if mixer_backward_mode != "learned": + raise RuntimeError("training/evaluation cannot use an intervention") + return super().forward( + input_ids, capture=False, mixer_backward_mode="learned" + ) + if mixer_backward_mode == "learned": + self.last_selector_visits = None + return super().forward( + input_ids, capture=True, mixer_backward_mode="learned" + ) + if mixer_backward_mode not in MATRIX_MODES: + raise ValueError(f"unknown local matrix mode: {mixer_backward_mode}") + self._active_selector_visits = [] + logits, trace = self._diagnostic_forward( + input_ids, mixer_backward_mode + ) + self.last_selector_visits = self._active_selector_visits + self._active_selector_visits = None + return logits, trace + + def _mix( + self, + mixer: nn.Module, + sources: list[torch.Tensor], + labels: list[str], + *, + mode: str, + mixer_index: int | None, + layer: int | None, + branch: str, + group: int | None, + offset: int | None, + ) -> tuple[torch.Tensor, dict[str, Any]]: + if mode == "learned": + return super()._mix( + mixer, + sources, + labels, + mode=mode, + mixer_index=mixer_index, + layer=layer, + branch=branch, + group=group, + offset=offset, + ) + if self._active_selector_visits is None: + raise RuntimeError("local selector visit log is not active") + + if mixer_index is None: + if ( + layer is not None + or group is not None + or branch != "output" + or offset is not None + ): + raise RuntimeError("invalid output mixer identity") + identity = { + "kind": "output", + "index": 64, + "layer": None, + "group": None, + "branch": "output", + "offset": None, + } + use_uniform = bool(self.uniform_output[mode]) + else: + expected_layer = mixer_index // 2 + 1 + expected_branch = "attention" if mixer_index % 2 == 0 else "mlp" + expected_group = (expected_layer - 1) // 4 + 1 + expected_offset = (expected_layer - 1) % 4 + 1 + if ( + not 0 <= mixer_index < 64 + or layer != expected_layer + or branch != expected_branch + or group != expected_group + or offset != expected_offset + ): + raise RuntimeError("invalid depth mixer identity") + identity = { + "kind": "depth", + "index": mixer_index, + "layer": layer, + "group": group, + "branch": branch, + "offset": offset, + } + use_uniform = mixer_index in self.uniform_depth[mode] + + parent_output, _ = mixer(sources, False) + weights = parent.recompute_weights(mixer, sources) + summary = parent.weight_summary( + weights, + labels, + mixer_index=mixer_index, + layer=layer, + branch=branch, + group=group, + offset=offset, + ) + backward_weights = ( + torch.full_like(weights, 1.0 / len(sources)) + if use_uniform + else weights + ) + values = torch.stack(sources, dim=0) + routed = parent.RoutedSourceBackward.apply( + values, parent_output, backward_weights + ) + self._active_selector_visits.append( + { + **identity, + "sources": len(sources), + "use_uniform": use_uniform, + "coefficient": ( + "uniform" if use_uniform else "detached_learned" + ), + } + ) + return routed, summary + + +def expected_visit_identities() -> list[dict[str, Any]]: + result = [] + for index in range(64): + layer = index // 2 + 1 + result.append( + { + "kind": "depth", + "index": index, + "layer": layer, + "group": (layer - 1) // 4 + 1, + "branch": "attention" if index % 2 == 0 else "mlp", + "offset": (layer - 1) % 4 + 1, + } + ) + result.append( + { + "kind": "output", + "index": 64, + "layer": None, + "group": None, + "branch": "output", + "offset": None, + } + ) + return result + + +def validate_selector_visits( + mode: str, + visits: list[dict[str, Any]] | None, + manifest: dict[str, Any], +) -> dict[str, Any]: + if visits is None or len(visits) != 65: + raise RuntimeError("selector visit count mismatch") + identity_keys = ("kind", "index", "layer", "group", "branch", "offset") + actual_identities = [ + {key: visit[key] for key in identity_keys} for visit in visits + ] + expected_identities = expected_visit_identities() + if actual_identities != expected_identities: + raise RuntimeError("selector identity order mismatch") + if len({(item["kind"], item["index"]) for item in visits}) != 65: + raise RuntimeError("selector identities are not unique") + + expected_depth = set( + manifest["selector"]["uniform_depth_indices"][mode] + ) + expected_output = manifest["selector"]["uniform_output"][mode] + actual_depth = { + item["index"] + for item in visits + if item["kind"] == "depth" and item["use_uniform"] + } + actual_output = visits[-1]["use_uniform"] + if actual_depth != expected_depth or actual_output != expected_output: + raise RuntimeError("selector exact-set mismatch") + uniform_count = sum(int(item["use_uniform"]) for item in visits) + expected_count = manifest["selector"]["expected_uniform_counts"][mode] + if uniform_count != expected_count: + raise RuntimeError("selector uniform census mismatch") + return { + "passed": True, + "visit_count": len(visits), + "identities_unique": True, + "identity_order_sha256": parent.canonical_sha256( + actual_identities + ), + "uniform_indices": [ + item["index"] for item in visits if item["use_uniform"] + ], + "uniform_count": uniform_count, + "expected_uniform_count": expected_count, + "visits_sha256": parent.canonical_sha256(visits), + "visits": visits, + } + + +def run_local_diagnostic( + model: LocalPathLanguageModel, + corpus: Any, + optimizer: torch.optim.Optimizer, + manifest: dict[str, Any], + *, + mode: str, + loss_scale: float = 1.0, +) -> dict[str, Any]: + result = parent.run_diagnostic( + model, corpus, optimizer, mode=mode, loss_scale=loss_scale + ) + result["selector"] = validate_selector_visits( + mode, model.last_selector_visits, manifest + ) + return result + + +def forward_identity_gate(matrix: dict[str, Any]) -> dict[str, Any]: + reference = matrix["detached_learned"]["forward"] + comparisons = {} + for mode in MATRIX_MODES[1:]: + other = matrix[mode]["forward"] + comparisons[mode] = { + "logits_exact": ( + other["logits_sha256"] == reference["logits_sha256"] + ), + "loss_exact": other["loss_nats"] == reference["loss_nats"], + "activations_exact": ( + other["activation_sha256"] + == reference["activation_sha256"] + ), + "mixer_summaries_exact": ( + other["mixer_summary_sha256"] + == reference["mixer_summary_sha256"] + ), + } + if not all(all(checks.values()) for checks in comparisons.values()): + raise RuntimeError("local matrix forward identity failed") + return {"passed": True, "comparisons": comparisons} + + +def spectrum_agreement( + left: dict[str, Any], right: dict[str, Any] +) -> dict[str, Any]: + checks = {} + passed = True + for position in parent.POSITIONS: + left_metric = left["positions"][position]["reductions"][ + "element_rms" + ] + right_metric = right["positions"][position]["reductions"][ + "element_rms" + ] + raw_errors = [ + abs(a - b) / a + for a, b in zip( + left_metric["values"], right_metric["values"] + ) + ] + normalized_errors = [ + abs(a - b) + for a, b in zip( + left_metric["statistics"]["normalized"], + right_metric["statistics"]["normalized"], + ) + ] + item_passed = ( + all( + math.isfinite(value) and value > 0 + for value in left_metric["values"] + ) + and max(raw_errors) <= parent.SPECTRUM_TOLERANCE + and max(normalized_errors) <= parent.SPECTRUM_TOLERANCE + ) + passed = passed and item_passed + checks[position] = { + "passed": item_passed, + "max_raw_relative_error": max(raw_errors), + "max_normalized_absolute_error": max(normalized_errors), + } + return {"passed": passed, "checks": checks} + + +def initialization_negative_control( + parent_learned: dict[str, Any], matrix: dict[str, Any] +) -> dict[str, Any]: + reference = matrix["detached_learned"] + comparisons = { + "parent_learned_vs_detached": spectrum_agreement( + parent_learned, reference + ) + } + for mode in MATRIX_MODES[1:]: + comparisons[mode] = spectrum_agreement(reference, matrix[mode]) + passed = all(item["passed"] for item in comparisons.values()) + if not passed: + raise RuntimeError("initialization negative control failed") + return {"passed": True, "comparisons": comparisons} + + +def without_selector(result: dict[str, Any], rename: str | None = None) -> dict[str, Any]: + cleaned = {key: value for key, value in result.items() if key != "selector"} + if rename is not None: + cleaned["mode"] = rename + return cleaned + + +def endpoint_exactness( + matrix: dict[str, Any], parent_diagnostic: dict[str, Any] +) -> dict[str, Any]: + reference_exact = ( + without_selector(matrix["detached_learned"]) + == parent_diagnostic["modes"]["detached_learned"] + ) + uniform_exact = ( + without_selector( + matrix["uniform_all"], rename="uniform_value_backward" + ) + == parent_diagnostic["modes"]["uniform_value_backward"] + ) + checks = { + "detached_learned_round06_exact": reference_exact, + "uniform_all_round06_exact": uniform_exact, + } + if not all(checks.values()): + raise RuntimeError(f"Round 06 endpoint exactness failed: {checks}") + return {"passed": True, "checks": checks} + + +def run_diagnostic_bundle( + model: LocalPathLanguageModel, + corpus: Any, + optimizer: torch.optim.Optimizer, + manifest: dict[str, Any], + parent_diagnostic: dict[str, Any], + step: int, +) -> dict[str, Any]: + parent_learned = parent.run_diagnostic( + model, corpus, optimizer, mode="learned" + ) + if parent_learned != parent_diagnostic["modes"]["learned"]: + raise RuntimeError("parent learned diagnostic is not Round 06 exact") + result: dict[str, Any] = { + "step": step, + "parent_learned": parent_learned, + "parent_learned_round06_exact": True, + "local_matrix": None, + } + if step not in MATRIX_STEPS: + return result + + matrix = { + mode: run_local_diagnostic( + model, corpus, optimizer, manifest, mode=mode + ) + for mode in MATRIX_MODES + } + result["local_matrix"] = matrix + result["forward_identity_gate"] = forward_identity_gate(matrix) + result["endpoint_exactness"] = endpoint_exactness( + matrix, parent_diagnostic + ) + if step == 0: + result["initialization_negative_control"] = ( + initialization_negative_control(parent_learned, matrix) + ) + doubled = run_local_diagnostic( + model, + corpus, + optimizer, + manifest, + mode="detached_learned", + loss_scale=2.0, + ) + result["loss_scale_gate"] = parent.loss_scale_gate( + matrix["detached_learned"], doubled + ) + model.zero_grad(set_to_none=True) + return result + + +def load_and_verify_inputs( + args: argparse.Namespace, +) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any], Path]: + manifest = json.loads(args.manifest.read_text()) + parent_manifest = json.loads(args.parent_manifest.read_text()) + data_manifest = json.loads(args.data_manifest.read_text()) + repo_root = Path(__file__).resolve().parents[3] + + if manifest["protocol_id"] != PROTOCOL_ID: + raise RuntimeError("Round 07 manifest protocol mismatch") + if parent_manifest["protocol_id"] != PARENT_PROTOCOL_ID: + raise RuntimeError("Round 06 parent manifest protocol mismatch") + if data_manifest["protocol_id"] != DATA_PROTOCOL_ID: + raise RuntimeError("data manifest protocol mismatch") + if manifest["formal_seeds"] != list(SEEDS): + raise RuntimeError("formal seed mismatch") + if manifest["matrix_modes"] != list(MATRIX_MODES): + raise RuntimeError("local matrix mode mismatch") + if ( + manifest["training"]["parent_diagnostic_steps"] + != list(PARENT_DIAGNOSTIC_STEPS) + or manifest["training"]["local_matrix_steps"] != list(MATRIX_STEPS) + or manifest["training"]["steps"] != FORMAL_STEPS + ): + raise RuntimeError("diagnostic/training schedule mismatch") + + parent_artifacts = manifest["parent_artifacts"] + if parent.file_sha256(args.parent_manifest) != parent_artifacts[ + "manifest_sha256" + ]: + raise RuntimeError("Round 06 manifest physical hash mismatch") + if parent.file_sha256(Path(parent.__file__)) != parent_artifacts[ + "runner_sha256" + ]: + raise RuntimeError("Round 06 runner physical hash mismatch") + for name in ("protocol", "scoping"): + path = repo_root / parent_artifacts[f"{name}_path"] + if parent.file_sha256(path) != parent_artifacts[f"{name}_sha256"]: + raise RuntimeError(f"Round 06 {name} physical hash mismatch") + for name in ("protocol", "scoping", "grok_review"): + path = repo_root / manifest["current_artifacts"][f"{name}_path"] + if parent.file_sha256(path) != manifest["current_artifacts"][ + f"{name}_sha256" + ]: + raise RuntimeError(f"Round 07 {name} physical hash mismatch") + for key in ( + "formal_schedule_sha256", + "validation_tensor_sha256", + "diagnostic_tensor_sha256", + ): + if ( + data_manifest["windows"][key] + != parent_artifacts[key] + or parent_manifest["parent_artifacts"][key] + != parent_artifacts[key] + ): + raise RuntimeError(f"frozen data hash mismatch: {key}") + + parent_raw_path = ( + repo_root + / "experiments" + / "k3" + / "attnres_spike" + / "results" + / "raw" + / f"formal-seed-{args.seed}.json" + ) + expected = manifest["round06_expected"][str(args.seed)] + if parent.file_sha256(parent_raw_path) != expected["raw_file_sha256"]: + raise RuntimeError("Round 06 raw physical hash mismatch") + parent_raw = json.loads(parent_raw_path.read_text()) + if ( + parent_raw["canonical_sha256_without_self"] + != expected["canonical_sha256"] + or parent_raw["hashes"]["final_model_state"] + != expected["final_model_state"] + or parent_raw["hashes"]["final_optimizer_state"] + != expected["final_optimizer_state"] + ): + raise RuntimeError("Round 06 raw expected-state mismatch") + return manifest, data_manifest, parent_raw, repo_root + + +def frozen_training_compare( + result: dict[str, Any], parent_raw: dict[str, Any] +) -> dict[str, Any]: + checks = { + "final_model_state": ( + result["hashes"]["final_model_state"] + == parent_raw["hashes"]["final_model_state"] + ), + "final_optimizer_state": ( + result["hashes"]["final_optimizer_state"] + == parent_raw["hashes"]["final_optimizer_state"] + ), + "evaluations": result["evaluations"] == parent_raw["evaluations"], + "training_history": ( + result["training_history"] == parent_raw["training_history"] + ), + } + parent_diagnostics_exact = [] + endpoint_exact = [] + for new, old in zip(result["diagnostics"], parent_raw["diagnostics"]): + parent_diagnostics_exact.append( + new["step"] == old["step"] + and new["parent_learned"] == old["modes"]["learned"] + ) + if new["step"] in MATRIX_STEPS: + endpoint_exact.append(new["endpoint_exactness"]["passed"]) + checks["parent_learned_diagnostics"] = all(parent_diagnostics_exact) + checks["round06_endpoints"] = len(endpoint_exact) == 2 and all( + endpoint_exact + ) + if not all(checks.values()): + raise RuntimeError(f"Round 06 training equivalence failed: {checks}") + return {"passed": True, "checks": checks} + + +def main() -> None: + args = parse_args() + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required") + if os.environ.get("CUBLAS_WORKSPACE_CONFIG") != ":4096:8": + raise RuntimeError("CUBLAS_WORKSPACE_CONFIG must be :4096:8") + manifest, data_manifest, parent_raw, repo_root = load_and_verify_inputs(args) + parent.parent.configure_round04_globals(DEPTH) + parent.parent.configure_determinism(args.seed) + corpus = parent.parent.round04.ByteCorpus( + args.cache_dir, data_manifest, torch.device("cuda") + ) + model = LocalPathLanguageModel("block", manifest).to( + torch.device("cuda") + ) + + initial_public_hash = parent.parent.named_state_hash( + model, include_mixers=False + ) + initial_mixer_hash = parent.parent.named_state_hash( + model, include_mixers=True + ) + public_structure_hash, public_tensors, public_elements = ( + parent.parent.state_structure_hash(model, include_mixers=False) + ) + input_gate_hashes = parent.parent.model_input_gate_hashes( + corpus, data_manifest, args.seed, TRAIN_BATCH_SIZE + ) + + decay_parameters: list[nn.Parameter] = [] + no_decay_parameters: list[nn.Parameter] = [] + for parameter in model.parameters(): + target = decay_parameters if parameter.ndim >= 2 else no_decay_parameters + target.append(parameter) + optimizer = torch.optim.AdamW( + [ + { + "params": decay_parameters, + "weight_decay": parent.parent.WEIGHT_DECAY, + }, + {"params": no_decay_parameters, "weight_decay": 0.0}, + ], + lr=parent.parent.PEAK_LR, + betas=parent.parent.BETAS, + eps=parent.parent.ADAM_EPS, + ) + + parent_by_step = { + item["step"]: item for item in parent_raw["diagnostics"] + } + evaluations = [ + { + "step": 0, + **parent.parent.evaluate( + model, corpus, VALIDATION_WINDOWS, EVAL_BATCH_SIZE + ), + } + ] + diagnostics = [ + run_diagnostic_bundle( + model, + corpus, + optimizer, + manifest, + parent_by_step[0], + 0, + ) + ] + print( + json.dumps( + { + "event": "local_matrix", + "step": 0, + "seed": args.seed, + "modes": len(MATRIX_MODES), + "endpoint_exact": diagnostics[0][ + "endpoint_exactness" + ]["passed"], + }, + sort_keys=True, + ), + flush=True, + ) + + steps = 0 if args.run_kind == "smoke" else FORMAL_STEPS + training_history: list[dict[str, float | int]] = [] + step_times: list[float] = [] + if steps: + model.train() + for step in range(1, steps + 1): + lr = parent.parent.learning_rate(step, steps) + for group in optimizer.param_groups: + group["lr"] = lr + inputs, targets = corpus.training_batch( + args.seed, step, TRAIN_BATCH_SIZE + ) + optimizer.zero_grad(set_to_none=True) + torch.cuda.synchronize() + started = time.perf_counter() + with torch.autocast(device_type="cuda", dtype=torch.bfloat16): + logits, trace = model(inputs) + if trace is not None: + raise RuntimeError( + "training unexpectedly captured a trace" + ) + loss = parent.parent.cross_entropy(logits, targets) + if not torch.isfinite(loss): + raise RuntimeError(f"non-finite loss at step {step}") + loss.backward() + unclipped_norm = torch.nn.utils.clip_grad_norm_( + model.parameters(), parent.parent.GRAD_CLIP + ) + optimizer.step() + torch.cuda.synchronize() + elapsed_ms = (time.perf_counter() - started) * 1000 + if step == TIMING_WARMUP: + torch.cuda.reset_peak_memory_stats() + elif step > TIMING_WARMUP: + step_times.append(elapsed_ms) + if step == 1 or step % 10 == 0 or step == steps: + training_history.append( + { + "step": step, + "loss_nats": loss.detach().cpu().item(), + "bits_per_byte": ( + loss.detach().cpu().item() / math.log(2) + ), + "learning_rate": lr, + "unclipped_grad_norm": float( + unclipped_norm.detach().cpu() + ), + } + ) + if step in PARENT_DIAGNOSTIC_STEPS: + evaluations.append( + { + "step": step, + **parent.parent.evaluate( + model, + corpus, + VALIDATION_WINDOWS, + EVAL_BATCH_SIZE, + ), + } + ) + diagnostic = run_diagnostic_bundle( + model, + corpus, + optimizer, + manifest, + parent_by_step[step], + step, + ) + diagnostics.append(diagnostic) + event = { + "event": "diagnostic", + "step": step, + "seed": args.seed, + "validation_bpc": evaluations[-1]["bits_per_byte"], + "parent_exact": diagnostic[ + "parent_learned_round06_exact" + ], + } + if diagnostic["local_matrix"] is not None: + event["modes"] = len(MATRIX_MODES) + event["endpoint_exact"] = diagnostic[ + "endpoint_exactness" + ]["passed"] + print(json.dumps(event, sort_keys=True), flush=True) + model.train() + + timing = { + "warmup_steps_excluded": TIMING_WARMUP, + "measured_steps": len(step_times), + "mean_ms": ( + statistics.fmean(step_times) if step_times else None + ), + "median_ms": ( + statistics.median(step_times) if step_times else None + ), + "p95_ms": ( + parent.type7_quantile(torch.tensor(sorted(step_times)), 0.95) + if step_times + else None + ), + "peak_allocated_bytes": torch.cuda.max_memory_allocated(), + "peak_reserved_bytes": torch.cuda.max_memory_reserved(), + } + result = { + "schema_version": 1, + "protocol_id": PROTOCOL_ID, + "parent_protocol_id": PARENT_PROTOCOL_ID, + "run_kind": args.run_kind, + "architecture": "block", + "depth": DEPTH, + "seed": args.seed, + "steps": steps, + "batch_size": TRAIN_BATCH_SIZE, + "target_bytes_seen": ( + steps * TRAIN_BATCH_SIZE * parent.CONTEXT + ), + "manifest": { + "path": str(args.manifest), + "file_sha256": parent.file_sha256(args.manifest), + "parent_path": str(args.parent_manifest), + "parent_file_sha256": parent.file_sha256( + args.parent_manifest + ), + "data_path": str(args.data_manifest), + "data_file_sha256": parent.file_sha256(args.data_manifest), + "formal_schedule_sha256": data_manifest["windows"][ + "formal_schedule_sha256" + ], + "validation_tensor_sha256": data_manifest["windows"][ + "validation_tensor_sha256" + ], + "diagnostic_tensor_sha256": data_manifest["windows"][ + "diagnostic_tensor_sha256" + ], + "input_gate_tensor_hashes": input_gate_hashes, + "selector_contract_sha256": parent.canonical_sha256( + manifest["selector"] + ), + }, + "model": { + "layers": DEPTH, + "aggregation_groups": 8, + "blocks_per_group": 4, + "d_model": parent.parent.round04.D_MODEL, + "heads": parent.parent.round04.HEADS, + "d_ff": parent.parent.round04.D_FF, + "parameters": parent.parent.parameter_inventory(model), + }, + "hashes": { + "initial_public_parameter_structure": public_structure_hash, + "initial_public_parameter_tensors": public_tensors, + "initial_public_parameter_elements": public_elements, + "initial_public_parameters": initial_public_hash, + "initial_mixer_parameters": initial_mixer_hash, + "final_public_parameters": parent.parent.named_state_hash( + model, include_mixers=False + ), + "final_mixer_parameters": parent.parent.named_state_hash( + model, include_mixers=True + ), + "final_model_state": parent.parent.named_state_hash( + model, include_mixers=None + ), + "final_optimizer_state": parent.parent.recursive_state_hash( + optimizer.state_dict() + ), + }, + "evaluations": evaluations, + "diagnostics": diagnostics, + "training_history": training_history, + "timing": timing, + "environment": { + "python": platform.python_version(), + "torch": torch.__version__, + "cuda": torch.version.cuda, + "gpu": torch.cuda.get_device_name(0), + "compute_capability": list( + torch.cuda.get_device_capability(0) + ), + "cublas_workspace_config": os.environ[ + "CUBLAS_WORKSPACE_CONFIG" + ], + "deterministic_algorithms": ( + torch.are_deterministic_algorithms_enabled() + ), + "autocast": "cuda-bfloat16-forward-fp32-cross-entropy", + "compile": False, + }, + "artifacts": { + "runner_sha256": parent.file_sha256(Path(__file__)), + "protocol_sha256": parent.file_sha256( + repo_root + / "research" + / "K3_ATTNRES_LOCAL_PATH_PROTOCOL.md" + ), + "scoping_sha256": parent.file_sha256( + repo_root + / "research" + / "K3_ATTNRES_LOCAL_PATH_SCOPING.md" + ), + "grok_review_sha256": parent.file_sha256( + repo_root + / "research" + / "K3_ATTNRES_LOCAL_PATH_GROK_REVIEW.md" + ), + }, + } + result["round06_equivalence"] = ( + frozen_training_compare(result, parent_raw) if steps else None + ) + result["canonical_sha256_without_self"] = parent.canonical_sha256( + result + ) + args.output.parent.mkdir(parents=True, exist_ok=True) + temporary = args.output.with_suffix(args.output.suffix + ".tmp") + temporary.write_text( + json.dumps(result, ensure_ascii=False, indent=2, sort_keys=True) + + "\n" + ) + os.replace(temporary, args.output) + print( + json.dumps( + { + "output": str(args.output), + "run_kind": args.run_kind, + "seed": args.seed, + "steps": steps, + "final_bpc": evaluations[-1]["bits_per_byte"], + "canonical_sha256": result[ + "canonical_sha256_without_self" + ], + "timing": timing, + }, + ensure_ascii=False, + indent=2, + ), + flush=True, + ) + + +if __name__ == "__main__": + main()