wuyang
295 lines
10 KiB
Python
295 lines
10 KiB
Python
#!/usr/bin/env python3
|
|
"""Run pre-result Round 08 identity gates."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import importlib.util
|
|
import json
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import torch
|
|
|
|
|
|
def load_runner() -> Any:
|
|
path = Path(__file__).resolve().parent / "train.py"
|
|
spec = importlib.util.spec_from_file_location("k3_attnres_round08_train", path)
|
|
if spec is None or spec.loader is None:
|
|
raise RuntimeError(f"cannot import Round 08 runner from {path}")
|
|
module = importlib.util.module_from_spec(spec)
|
|
sys.modules[spec.name] = module
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
runner = load_runner()
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser()
|
|
subparsers = parser.add_subparsers(dest="command", required=True)
|
|
|
|
step_zero = subparsers.add_parser("step-zero")
|
|
step_zero.add_argument("--cache-dir", type=Path, required=True)
|
|
step_zero.add_argument("--parent-manifest", type=Path, required=True)
|
|
step_zero.add_argument("--study-manifest", type=Path, required=True)
|
|
step_zero.add_argument("--output", type=Path, required=True)
|
|
step_zero.add_argument("--seed", type=int, default=2026073001)
|
|
|
|
smoke = subparsers.add_parser("smoke-compare")
|
|
smoke.add_argument("--parent", type=Path, required=True)
|
|
smoke.add_argument("--wrapper", type=Path, required=True)
|
|
smoke.add_argument("--output", type=Path, required=True)
|
|
return parser.parse_args()
|
|
|
|
|
|
def canonical_sha256(value: Any) -> str:
|
|
return hashlib.sha256(
|
|
json.dumps(
|
|
value, ensure_ascii=False, sort_keys=True, separators=(",", ":")
|
|
).encode()
|
|
).hexdigest()
|
|
|
|
|
|
def tensor_sha256(value: torch.Tensor) -> str:
|
|
return hashlib.sha256(runner.parent.tensor_bytes(value)).hexdigest()
|
|
|
|
|
|
def read_and_verify(path: Path) -> dict[str, Any]:
|
|
value = json.loads(path.read_text())
|
|
expected = value["canonical_sha256_without_self"]
|
|
payload = {
|
|
key: item
|
|
for key, item in value.items()
|
|
if key != "canonical_sha256_without_self"
|
|
}
|
|
if canonical_sha256(payload) != expected:
|
|
raise RuntimeError(f"canonical self-hash failed: {path}")
|
|
return value
|
|
|
|
|
|
def smoke_compare(args: argparse.Namespace) -> None:
|
|
parent_result = read_and_verify(args.parent)
|
|
wrapper_result = read_and_verify(args.wrapper)
|
|
fields = (
|
|
"architecture",
|
|
"depth",
|
|
"seed",
|
|
"steps",
|
|
"batch_size",
|
|
"target_bytes_seen",
|
|
"manifest",
|
|
"model",
|
|
"optimizer",
|
|
"hashes",
|
|
"evaluations",
|
|
"diagnostics",
|
|
"training_history",
|
|
"gradient_gate",
|
|
"environment",
|
|
)
|
|
checks = {}
|
|
for field in fields:
|
|
parent_value = exact_structure(parent_result[field])
|
|
wrapper_value = exact_structure(wrapper_result[field])
|
|
if field == "manifest":
|
|
parent_value.pop("path", None)
|
|
wrapper_value.pop("path", None)
|
|
checks[field] = parent_value == wrapper_value
|
|
wrapper_identity = {
|
|
"protocol": wrapper_result["protocol_id"] == runner.PROTOCOL_ID,
|
|
"parent_protocol": (
|
|
wrapper_result["parent_protocol_id"]
|
|
== runner.PARENT_PROTOCOL_ID
|
|
),
|
|
"variant": wrapper_result["variant"] == "learned_reference",
|
|
"forward_audit": wrapper_result["forward_intervention"]["passed"],
|
|
}
|
|
passed = all(checks.values()) and all(wrapper_identity.values())
|
|
result = {
|
|
"schema_version": 1,
|
|
"protocol_id": runner.PROTOCOL_ID,
|
|
"gate": "empty-selector-parent-equivalence",
|
|
"passed": passed,
|
|
"field_checks": checks,
|
|
"wrapper_identity": wrapper_identity,
|
|
"excluded_fields": [
|
|
"protocol wrapper fields",
|
|
"timing",
|
|
"self hash",
|
|
"parent runner self hash",
|
|
"study manifest",
|
|
],
|
|
"parent_file": str(args.parent.resolve()),
|
|
"wrapper_file": str(args.wrapper.resolve()),
|
|
}
|
|
result["canonical_sha256_without_self"] = canonical_sha256(result)
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
args.output.write_text(
|
|
json.dumps(result, ensure_ascii=False, indent=2, sort_keys=True) + "\n"
|
|
)
|
|
if not passed:
|
|
raise RuntimeError(f"empty-selector parent equivalence failed: {checks}")
|
|
|
|
|
|
def exact_structure(value: Any) -> Any:
|
|
return json.loads(
|
|
json.dumps(value, ensure_ascii=False, sort_keys=True)
|
|
)
|
|
|
|
|
|
def step_zero(args: argparse.Namespace) -> None:
|
|
if not torch.cuda.is_available():
|
|
raise RuntimeError("CUDA is required by the frozen step-zero gate")
|
|
study_manifest = json.loads(args.study_manifest.read_text())
|
|
if study_manifest["protocol_id"] != runner.PROTOCOL_ID:
|
|
raise RuntimeError("study manifest mismatch")
|
|
parent_manifest = json.loads(args.parent_manifest.read_text())
|
|
if parent_manifest["protocol_id"] != runner.PARENT_PROTOCOL_ID:
|
|
raise RuntimeError("parent manifest mismatch")
|
|
|
|
parent = runner.parent
|
|
parent.configure_round04_globals(32)
|
|
device = torch.device("cuda")
|
|
corpus = parent.round04.ByteCorpus(
|
|
args.cache_dir, parent_manifest, device
|
|
)
|
|
inputs, targets = corpus.fixed_batch(
|
|
corpus.diagnostic_starts, 0, 16
|
|
)
|
|
|
|
variants = ("learned_reference",) + tuple(
|
|
study_manifest["variants"].keys()
|
|
)
|
|
observations: dict[str, Any] = {}
|
|
reference_payload: dict[str, Any] | None = None
|
|
for variant in variants:
|
|
parent.configure_determinism(args.seed)
|
|
runner.ACTIVE_VARIANT = variant
|
|
model = runner.ForwardInterventionLanguageModel("block").to(device)
|
|
initial_public = parent.named_state_hash(
|
|
model, include_mixers=False
|
|
)
|
|
initial_mixer = parent.named_state_hash(
|
|
model, include_mixers=True
|
|
)
|
|
with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
|
|
logits, trace = model(inputs, capture=True)
|
|
loss = parent.cross_entropy(logits, targets)
|
|
if trace is None:
|
|
raise RuntimeError("step-zero trace missing")
|
|
evaluation = parent.evaluate(model, corpus, 64, 8)
|
|
diagnostic = parent.diagnostic(model, corpus, 16)
|
|
payload = {
|
|
"initial_public_hash": initial_public,
|
|
"initial_mixer_hash": initial_mixer,
|
|
"logits_sha256": tensor_sha256(logits),
|
|
"loss_nats": loss.detach().cpu().item(),
|
|
"loss_tensor_sha256": tensor_sha256(loss),
|
|
"evaluation": exact_structure(evaluation),
|
|
"diagnostic": exact_structure(diagnostic),
|
|
}
|
|
if reference_payload is None:
|
|
reference_payload = payload
|
|
exact_checks = {
|
|
key: payload[key] == reference_payload[key]
|
|
for key in payload
|
|
}
|
|
selected = tuple(runner.VARIANTS[variant])
|
|
capture_checks = {}
|
|
for index in selected:
|
|
summary = trace.depth_weights[index]
|
|
source_count = summary["sources"]
|
|
capture_checks[str(index)] = {
|
|
"source_count": source_count,
|
|
"expected_source_count": study_manifest[
|
|
"selected_source_counts"
|
|
][str(index)],
|
|
"capture_summary_exact_vs_learned": (
|
|
payload["diagnostic"]["depth_weights"][index]
|
|
== reference_payload["diagnostic"]["depth_weights"][index]
|
|
),
|
|
"passed": (
|
|
source_count
|
|
== study_manifest["selected_source_counts"][str(index)]
|
|
and payload["diagnostic"]["depth_weights"][index]
|
|
== reference_payload["diagnostic"]["depth_weights"][index]
|
|
),
|
|
}
|
|
runtime_uniform_gate = (
|
|
model.uniform_weight_max_abs_error
|
|
<= study_manifest["thresholds"][
|
|
"uniform_weight_max_abs_error"
|
|
]
|
|
)
|
|
observations[variant] = {
|
|
"payload": payload,
|
|
"exact_vs_learned_reference": exact_checks,
|
|
"selected_capture_checks": capture_checks,
|
|
"pre_reduction_uniform_weight_max_abs_error": (
|
|
model.uniform_weight_max_abs_error
|
|
),
|
|
"pre_reduction_uniform_weight_gate": runtime_uniform_gate,
|
|
"passed": (
|
|
all(exact_checks.values())
|
|
and all(
|
|
item["passed"] for item in capture_checks.values()
|
|
)
|
|
and runtime_uniform_gate
|
|
),
|
|
}
|
|
del model, logits, loss, trace
|
|
torch.cuda.empty_cache()
|
|
|
|
passed = all(item["passed"] for item in observations.values())
|
|
result = {
|
|
"schema_version": 1,
|
|
"protocol_id": runner.PROTOCOL_ID,
|
|
"gate": "step-zero-cross-variant-byte-exact",
|
|
"seed": args.seed,
|
|
"passed": passed,
|
|
"variants": observations,
|
|
"parent_manifest_sha256": runner.parent.file_sha256(
|
|
args.parent_manifest
|
|
),
|
|
"study_manifest_sha256": runner.parent.file_sha256(
|
|
args.study_manifest
|
|
),
|
|
"environment": {
|
|
"gpu": torch.cuda.get_device_name(0),
|
|
"torch": torch.__version__,
|
|
"cuda": torch.version.cuda,
|
|
"cublas_workspace_config": os.environ.get(
|
|
"CUBLAS_WORKSPACE_CONFIG"
|
|
),
|
|
},
|
|
}
|
|
result["canonical_sha256_without_self"] = canonical_sha256(result)
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
args.output.write_text(
|
|
json.dumps(result, ensure_ascii=False, indent=2, sort_keys=True) + "\n"
|
|
)
|
|
if not passed:
|
|
failed = [
|
|
name
|
|
for name, value in observations.items()
|
|
if not value["passed"]
|
|
]
|
|
raise RuntimeError(f"step-zero exactness failed: {failed}")
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
if args.command == "step-zero":
|
|
step_zero(args)
|
|
else:
|
|
smoke_compare(args)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|