wuyang/llm-atlas
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

experiment: implement AttnRes local path runner

Browse Source
This commit is contained in:
wuyang
2026-07-30 12:53:22 +08:00
parent 6911efc6e7
commit 39a9ad6215
3 changed files with 1461 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
experiments/k3/attnres_local_path/README.md
+30 -2
View File
@@ -18,5 +18,33 @@ GPU NVIDIA GeForce RTX 5090
CUBLAS_WORKSPACE_CONFIG=:4096:8 CUBLAS_WORKSPACE_CONFIG=:4096:8
``` ```
The runner and exact commands will be added after this preregistration The preregistration was committed as `6911efc` before the runner or any result
manifest is committed. No result file may precede that commit. 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.
experiments/k3/attnres_local_path/analyze.py
+498
View File
@@ -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()
experiments/k3/attnres_local_path/train.py
+933
View File
@@ -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()