experiment: implement AttnRes forward training runner
This commit is contained in:
wuyang
@@ -0,0 +1,289 @@
|
||||
#!/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 = {
|
||||
field: parent_result[field] == wrapper_result[field]
|
||||
for field in fields
|
||||
}
|
||||
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(),
|
||||
"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()
|
||||
Reference in New Issue
Block a user