experiment: implement AttnRes local path runner
This commit is contained in:
wuyang
@@ -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()
|
||||
Reference in New Issue
Block a user