experiment: implement AttnRes forward training runner
This commit is contained in:
wuyang
@@ -0,0 +1,58 @@
|
||||
# Attention Residuals train-time forward intervention
|
||||
|
||||
This directory implements preregistered protocol
|
||||
`llm-atlas-k3-attnres-forward-training-v1`:
|
||||
|
||||
- `research/K3_ATTNRES_FORWARD_TRAINING_SCOPING.md`
|
||||
- `research/K3_ATTNRES_FORWARD_TRAINING_PROTOCOL.md`
|
||||
- `research/K3_ATTNRES_FORWARD_TRAINING_GROK_REVIEW.md`
|
||||
|
||||
It is a depth-32 reduced Block AttnRes architecture ablation. It is not a
|
||||
Kimi-K3 checkpoint forward pass and does not claim to recover unpublished
|
||||
Figure 5 telemetry.
|
||||
|
||||
## Frozen environment
|
||||
|
||||
```text
|
||||
Python /home/wuyang/.pyenv/versions/3.10.14/envs/navi-router-cu128/bin/python
|
||||
PyTorch 2.11.0+cu128
|
||||
GPU NVIDIA GeForce RTX 5090
|
||||
CUBLAS_WORKSPACE_CONFIG=:4096:8
|
||||
maximum concurrency 2
|
||||
```
|
||||
|
||||
## Pre-result gates
|
||||
|
||||
The checked-in gate artifacts must pass before formal output:
|
||||
|
||||
```bash
|
||||
CUBLAS_WORKSPACE_CONFIG=:4096:8 \
|
||||
/home/wuyang/.pyenv/versions/3.10.14/envs/navi-router-cu128/bin/python \
|
||||
experiments/k3/attnres_forward/verify.py step-zero \
|
||||
--cache-dir /home/wuyang/.cache/llm-atlas/k3-attnres-gradient-scale-v1 \
|
||||
--parent-manifest experiments/k3/attnres_gradient/manifest.json \
|
||||
--study-manifest experiments/k3/attnres_forward/manifest.json \
|
||||
--output experiments/k3/attnres_forward/results/gates/step-zero.json
|
||||
```
|
||||
|
||||
`learned_reference` is smoke-only. Its 20-step result is compared with a
|
||||
fresh parent Round 05 smoke using `verify.py smoke-compare`.
|
||||
|
||||
## Formal matrix
|
||||
|
||||
```bash
|
||||
/home/wuyang/.pyenv/versions/3.10.14/envs/navi-router-cu128/bin/python \
|
||||
experiments/k3/attnres_forward/run_matrix.py \
|
||||
--python /home/wuyang/.pyenv/versions/3.10.14/envs/navi-router-cu128/bin/python \
|
||||
--cache-dir /home/wuyang/.cache/llm-atlas/k3-attnres-gradient-scale-v1 \
|
||||
--parent-manifest experiments/k3/attnres_gradient/manifest.json \
|
||||
--study-manifest experiments/k3/attnres_forward/manifest.json \
|
||||
--output-dir experiments/k3/attnres_forward/results/raw \
|
||||
--phase all \
|
||||
--concurrency 2
|
||||
```
|
||||
|
||||
This runs 12 formal cells and one full replay. The analyzer reads all cells,
|
||||
the frozen historical paired references, and generates the only authoritative
|
||||
status, interaction map, and website compact artifact.
|
||||
|
||||
@@ -0,0 +1,684 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Aggregate and gate preregistered Round 08 forward-training 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, Iterable
|
||||
|
||||
|
||||
PROTOCOL_ID = "llm-atlas-k3-attnres-forward-training-v1"
|
||||
PARENT_PROTOCOL_ID = "llm-atlas-k3-attnres-gradient-scale-v1"
|
||||
METRICS = ("spike_contrast", "peak_normalized")
|
||||
|
||||
|
||||
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("--reference-dir", type=Path, required=True)
|
||||
parser.add_argument("--aggregate-output", type=Path, required=True)
|
||||
parser.add_argument("--compact-output", type=Path, required=True)
|
||||
parser.add_argument("--reproduction-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 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 mean(values: Iterable[float]) -> float:
|
||||
return statistics.fmean(values)
|
||||
|
||||
|
||||
def read_result(path: Path, expected_protocol: str) -> dict[str, Any]:
|
||||
value = json.loads(path.read_text())
|
||||
if value.get("protocol_id") != expected_protocol:
|
||||
raise RuntimeError(f"protocol mismatch: {path}")
|
||||
expected = value.get("canonical_sha256_without_self")
|
||||
payload = {
|
||||
key: item
|
||||
for key, item in value.items()
|
||||
if key != "canonical_sha256_without_self"
|
||||
}
|
||||
if not isinstance(expected, str) or canonical_sha256(payload) != expected:
|
||||
raise RuntimeError(f"canonical self-hash mismatch: {path}")
|
||||
return value
|
||||
|
||||
|
||||
def exactly_one(values: list[dict[str, Any]], step: int) -> dict[str, Any]:
|
||||
matches = [value for value in values if value["step"] == step]
|
||||
if len(matches) != 1:
|
||||
raise RuntimeError(f"step {step} missing or duplicated")
|
||||
return matches[0]
|
||||
|
||||
|
||||
def spectrum_metrics(
|
||||
diagnostic: dict[str, Any],
|
||||
spike_layers: tuple[int, ...],
|
||||
epsilon: float,
|
||||
) -> dict[str, Any]:
|
||||
values = [
|
||||
float(value)
|
||||
for value in diagnostic["activation_grad_rms_by_block"]
|
||||
]
|
||||
if len(values) != 32:
|
||||
raise RuntimeError("activation-gradient spectrum must have 32 layers")
|
||||
if any(not math.isfinite(value) or value <= epsilon for value in values):
|
||||
raise RuntimeError("activation-gradient spectrum is non-finite/non-positive")
|
||||
spike_indices = {layer - 1 for layer in spike_layers}
|
||||
spike_values = [
|
||||
value for index, value in enumerate(values) if index in spike_indices
|
||||
]
|
||||
reference_values = [
|
||||
value for index, value in enumerate(values) if index not in spike_indices
|
||||
]
|
||||
spike_mean = mean(spike_values)
|
||||
reference_mean = mean(reference_values)
|
||||
global_mean = mean(values)
|
||||
contrast = spike_mean / reference_mean
|
||||
peak = max(values) / global_mean
|
||||
if any(
|
||||
not math.isfinite(value) or value <= epsilon
|
||||
for value in (spike_mean, reference_mean, contrast, peak)
|
||||
):
|
||||
raise RuntimeError("derived spike metric is non-finite/non-positive")
|
||||
ordered = sorted(range(32), key=lambda index: (-values[index], index))
|
||||
return {
|
||||
"values": values,
|
||||
"normalized": [value / global_mean for value in values],
|
||||
"spike_mean": spike_mean,
|
||||
"reference_mean": reference_mean,
|
||||
"global_mean": global_mean,
|
||||
"spike_contrast": contrast,
|
||||
"peak_normalized": peak,
|
||||
"peak_layer_1based": ordered[0] + 1,
|
||||
"top_five_layers_1based": [index + 1 for index in ordered[:5]],
|
||||
}
|
||||
|
||||
|
||||
def final_bpc(value: dict[str, Any], step: int) -> float:
|
||||
result = float(exactly_one(value["evaluations"], step)["bits_per_byte"])
|
||||
if not math.isfinite(result):
|
||||
raise RuntimeError("final validation BPC is non-finite")
|
||||
return result
|
||||
|
||||
|
||||
def stable_environment(value: dict[str, Any]) -> dict[str, Any]:
|
||||
keys = (
|
||||
"cublas_workspace_config",
|
||||
"deterministic_algorithms",
|
||||
"autocast",
|
||||
"compile",
|
||||
)
|
||||
return {key: value["environment"][key] for key in keys}
|
||||
|
||||
|
||||
def pairing_checks(
|
||||
run: dict[str, Any], reference: dict[str, Any]
|
||||
) -> dict[str, bool]:
|
||||
manifest_fields = (
|
||||
"formal_schedule_sha256",
|
||||
"validation_tensor_sha256",
|
||||
"diagnostic_tensor_sha256",
|
||||
"input_gate_tensor_hashes",
|
||||
)
|
||||
checks = {
|
||||
"seed": run["seed"] == reference["seed"],
|
||||
"architecture": (
|
||||
run["architecture"] == reference["architecture"] == "block"
|
||||
),
|
||||
"depth": run["depth"] == reference["depth"] == 32,
|
||||
"steps": run["steps"] == reference["steps"] == 8000,
|
||||
"batch_size": run["batch_size"] == reference["batch_size"] == 32,
|
||||
"initial_public_parameters": (
|
||||
run["hashes"]["initial_public_parameters"]
|
||||
== reference["hashes"]["initial_public_parameters"]
|
||||
),
|
||||
"initial_mixer_parameters": (
|
||||
run["hashes"]["initial_mixer_parameters"]
|
||||
== reference["hashes"]["initial_mixer_parameters"]
|
||||
),
|
||||
"model_topology": run["model"] == reference["model"],
|
||||
"optimizer_hyperparameters": (
|
||||
run["optimizer"] == reference["optimizer"]
|
||||
),
|
||||
"scientific_environment": (
|
||||
stable_environment(run) == stable_environment(reference)
|
||||
),
|
||||
}
|
||||
for field in manifest_fields:
|
||||
checks[f"manifest.{field}"] = (
|
||||
run["manifest"][field] == reference["manifest"][field]
|
||||
)
|
||||
return checks
|
||||
|
||||
|
||||
def scientific_replay_payload(value: dict[str, Any]) -> dict[str, Any]:
|
||||
payload = copy.deepcopy(value)
|
||||
for key in (
|
||||
"run_kind",
|
||||
"timing",
|
||||
"canonical_sha256_without_self",
|
||||
"parent_runner_canonical_sha256",
|
||||
):
|
||||
payload.pop(key, None)
|
||||
payload["manifest"].pop("path", None)
|
||||
payload["study_manifest"].pop("path", None)
|
||||
payload["environment"] = stable_environment(value)
|
||||
return payload
|
||||
|
||||
|
||||
def quality_gate(
|
||||
variant_runs: dict[int, dict[str, Any]],
|
||||
references: dict[int, dict[str, Any]],
|
||||
*,
|
||||
step: int,
|
||||
per_seed_maximum: float,
|
||||
mean_maximum: float,
|
||||
) -> dict[str, Any]:
|
||||
per_seed = {}
|
||||
for seed, run in sorted(variant_runs.items()):
|
||||
variant_bpc = final_bpc(run, step)
|
||||
reference_bpc = final_bpc(references[seed], step)
|
||||
delta = variant_bpc - reference_bpc
|
||||
per_seed[str(seed)] = {
|
||||
"variant_bpc": variant_bpc,
|
||||
"reference_bpc": reference_bpc,
|
||||
"delta_bpc": delta,
|
||||
"passed": delta <= per_seed_maximum,
|
||||
}
|
||||
mean_delta = mean(item["delta_bpc"] for item in per_seed.values())
|
||||
per_seed_passed = all(item["passed"] for item in per_seed.values())
|
||||
mean_passed = mean_delta <= mean_maximum
|
||||
return {
|
||||
"passed": per_seed_passed and mean_passed,
|
||||
"passed_checks": (
|
||||
sum(item["passed"] for item in per_seed.values())
|
||||
+ int(mean_passed)
|
||||
),
|
||||
"required_checks": 4,
|
||||
"per_seed_maximum": per_seed_maximum,
|
||||
"mean_maximum": mean_maximum,
|
||||
"mean_delta_bpc": mean_delta,
|
||||
"mean_passed": mean_passed,
|
||||
"per_seed": per_seed,
|
||||
}
|
||||
|
||||
|
||||
def variant_effect(
|
||||
variant: str,
|
||||
runs: dict[int, dict[str, Any]],
|
||||
references: dict[int, dict[str, Any]],
|
||||
metrics_by_cell: dict[tuple[str, int, int], dict[str, Any]],
|
||||
*,
|
||||
step: int,
|
||||
threshold: float,
|
||||
quality: dict[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
cells = []
|
||||
for seed in sorted(runs):
|
||||
candidate = metrics_by_cell[(variant, seed, step)]
|
||||
reference = metrics_by_cell[("learned_reference", seed, step)]
|
||||
for metric in METRICS:
|
||||
reference_value = reference[metric]
|
||||
candidate_value = candidate[metric]
|
||||
relative_drop = (
|
||||
reference_value - candidate_value
|
||||
) / reference_value
|
||||
cells.append(
|
||||
{
|
||||
"seed": seed,
|
||||
"metric": metric,
|
||||
"reference": reference_value,
|
||||
"variant": candidate_value,
|
||||
"relative_drop": relative_drop,
|
||||
"passed": relative_drop >= threshold,
|
||||
}
|
||||
)
|
||||
attenuation_passed = all(cell["passed"] for cell in cells)
|
||||
return {
|
||||
"variant": variant,
|
||||
"threshold": threshold,
|
||||
"passed_cells": sum(cell["passed"] for cell in cells),
|
||||
"required_cells": len(cells),
|
||||
"attenuation_passed": attenuation_passed,
|
||||
"quality": quality,
|
||||
"material_response_passed": (
|
||||
attenuation_passed and quality["passed"]
|
||||
),
|
||||
"cells": cells,
|
||||
}
|
||||
|
||||
|
||||
def interaction_map(
|
||||
metrics_by_cell: dict[tuple[str, int, int], dict[str, Any]],
|
||||
seeds: tuple[int, ...],
|
||||
steps: tuple[int, ...],
|
||||
) -> dict[str, Any]:
|
||||
cells = []
|
||||
for step in steps:
|
||||
for seed in seeds:
|
||||
reference = metrics_by_cell[
|
||||
("learned_reference", seed, step)
|
||||
]
|
||||
group6 = metrics_by_cell[
|
||||
("uniform_group_6_forward", seed, step)
|
||||
]
|
||||
group7 = metrics_by_cell[
|
||||
("uniform_group_7_forward", seed, step)
|
||||
]
|
||||
joint = metrics_by_cell[
|
||||
("uniform_groups_6_7_forward", seed, step)
|
||||
]
|
||||
for metric in METRICS:
|
||||
ref = reference[metric]
|
||||
effects = {
|
||||
"group6": math.log(ref / group6[metric]),
|
||||
"group7": math.log(ref / group7[metric]),
|
||||
"groups6_7": math.log(ref / joint[metric]),
|
||||
}
|
||||
residual = (
|
||||
effects["groups6_7"]
|
||||
- effects["group6"]
|
||||
- effects["group7"]
|
||||
)
|
||||
cells.append(
|
||||
{
|
||||
"step": step,
|
||||
"seed": seed,
|
||||
"metric": metric,
|
||||
"log_effects": effects,
|
||||
"interaction_residual": residual,
|
||||
"relative_drops": {
|
||||
"group6": (ref - group6[metric]) / ref,
|
||||
"group7": (ref - group7[metric]) / ref,
|
||||
"groups6_7": (ref - joint[metric]) / ref,
|
||||
},
|
||||
}
|
||||
)
|
||||
summaries = []
|
||||
for step in steps:
|
||||
for metric in METRICS:
|
||||
selected = [
|
||||
cell
|
||||
for cell in cells
|
||||
if cell["step"] == step and cell["metric"] == metric
|
||||
]
|
||||
residuals = [
|
||||
cell["interaction_residual"] for cell in selected
|
||||
]
|
||||
summaries.append(
|
||||
{
|
||||
"step": step,
|
||||
"metric": metric,
|
||||
"mean_interaction_residual": mean(residuals),
|
||||
"minimum": min(residuals),
|
||||
"maximum": max(residuals),
|
||||
}
|
||||
)
|
||||
return {
|
||||
"definition": "I67=ln(Xref/X67)-ln(Xref/X6)-ln(Xref/X7)",
|
||||
"interpretation": (
|
||||
"descriptive cross-run log-attenuation residual from three "
|
||||
"independently trained variants; not a causal interaction"
|
||||
),
|
||||
"cells": cells,
|
||||
"summaries": summaries,
|
||||
}
|
||||
|
||||
|
||||
def environment_metadata(value: dict[str, Any]) -> dict[str, Any]:
|
||||
return {
|
||||
key: value["environment"].get(key)
|
||||
for key in ("gpu", "torch", "cuda", "compute_capability")
|
||||
}
|
||||
|
||||
|
||||
def write_hashed(path: Path, value: dict[str, Any]) -> None:
|
||||
value["canonical_sha256_without_self"] = canonical_sha256(value)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text(
|
||||
json.dumps(value, ensure_ascii=False, indent=2, sort_keys=True) + "\n"
|
||||
)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
manifest = json.loads(args.manifest.read_text())
|
||||
if (
|
||||
manifest["protocol_id"] != PROTOCOL_ID
|
||||
or manifest["status"] != "frozen-before-model-output"
|
||||
):
|
||||
raise RuntimeError("manifest is not the frozen Round 08 contract")
|
||||
variants = tuple(manifest["variants"].keys())
|
||||
seeds = tuple(manifest["formal_seeds"])
|
||||
steps = tuple(manifest["diagnostic_steps"])
|
||||
primary_step = manifest["primary_step"]
|
||||
epsilon = manifest["thresholds"]["positive_denominator_epsilon"]
|
||||
spike_layers = tuple(manifest["spike_layers_1based"])
|
||||
expected_cells = {(variant, seed) for variant in variants for seed in seeds}
|
||||
if len(args.formal) != len(expected_cells):
|
||||
raise RuntimeError("formal path count does not match the 4×3 matrix")
|
||||
|
||||
runs: dict[tuple[str, int], dict[str, Any]] = {}
|
||||
run_paths: dict[tuple[str, int], Path] = {}
|
||||
pairing: dict[str, Any] = {}
|
||||
references: dict[int, dict[str, Any]] = {}
|
||||
reference_paths: dict[int, Path] = {}
|
||||
for seed in seeds:
|
||||
path = args.reference_dir / (
|
||||
f"formal-depth-32-block-seed-{seed}.json"
|
||||
)
|
||||
references[seed] = read_result(path, PARENT_PROTOCOL_ID)
|
||||
reference_paths[seed] = path
|
||||
|
||||
for path in args.formal:
|
||||
value = read_result(path, PROTOCOL_ID)
|
||||
identity = (value["variant"], value["seed"])
|
||||
if identity in runs:
|
||||
raise RuntimeError(f"duplicate formal cell: {identity}")
|
||||
if (
|
||||
value["run_kind"] != "formal"
|
||||
or value["steps"] != manifest["formal_steps"]
|
||||
or not value["forward_intervention"]["passed"]
|
||||
):
|
||||
raise RuntimeError(f"invalid formal cell: {path}")
|
||||
runs[identity] = value
|
||||
run_paths[identity] = path
|
||||
if set(runs) != expected_cells:
|
||||
raise RuntimeError("formal matrix identities do not match manifest")
|
||||
|
||||
for (variant, seed), value in sorted(runs.items()):
|
||||
checks = pairing_checks(value, references[seed])
|
||||
if not all(checks.values()):
|
||||
raise RuntimeError(
|
||||
f"historical reference pairing failed: "
|
||||
f"{variant}/{seed}: {checks}"
|
||||
)
|
||||
pairing[f"{variant}:{seed}"] = {
|
||||
"passed": True,
|
||||
"checks": checks,
|
||||
"run_environment": environment_metadata(value),
|
||||
"reference_environment": environment_metadata(references[seed]),
|
||||
"metadata_equal": (
|
||||
environment_metadata(value)
|
||||
== environment_metadata(references[seed])
|
||||
),
|
||||
}
|
||||
|
||||
replay = read_result(args.replay, PROTOCOL_ID)
|
||||
replay_contract = manifest["replay"]
|
||||
if (
|
||||
replay["run_kind"] != "replay"
|
||||
or replay["variant"] != replay_contract["variant"]
|
||||
or replay["seed"] != replay_contract["seed"]
|
||||
or replay["steps"] != manifest["formal_steps"]
|
||||
or not replay["forward_intervention"]["passed"]
|
||||
):
|
||||
raise RuntimeError("invalid replay identity/audit")
|
||||
formal_primary = runs[
|
||||
(replay_contract["variant"], replay_contract["seed"])
|
||||
]
|
||||
formal_payload = scientific_replay_payload(formal_primary)
|
||||
replay_payload = scientific_replay_payload(replay)
|
||||
replay_exact = formal_payload == replay_payload
|
||||
if not replay_exact:
|
||||
raise RuntimeError("primary formal/replay scientific payload mismatch")
|
||||
|
||||
metrics_by_cell: dict[tuple[str, int, int], dict[str, Any]] = {}
|
||||
for seed, reference in references.items():
|
||||
for step in steps:
|
||||
metrics_by_cell[("learned_reference", seed, step)] = (
|
||||
spectrum_metrics(
|
||||
exactly_one(reference["diagnostics"], step),
|
||||
spike_layers,
|
||||
epsilon,
|
||||
)
|
||||
)
|
||||
for (variant, seed), value in runs.items():
|
||||
if tuple(item["step"] for item in value["diagnostics"]) != steps:
|
||||
raise RuntimeError(f"diagnostic schedule drift: {variant}/{seed}")
|
||||
for step in steps:
|
||||
metrics_by_cell[(variant, seed, step)] = spectrum_metrics(
|
||||
exactly_one(value["diagnostics"], step),
|
||||
spike_layers,
|
||||
epsilon,
|
||||
)
|
||||
|
||||
runs_by_variant = {
|
||||
variant: {seed: runs[(variant, seed)] for seed in seeds}
|
||||
for variant in variants
|
||||
}
|
||||
qualities = {
|
||||
variant: quality_gate(
|
||||
variant_runs,
|
||||
references,
|
||||
step=primary_step,
|
||||
per_seed_maximum=manifest["thresholds"][
|
||||
"final_bpc_delta_per_seed_maximum"
|
||||
],
|
||||
mean_maximum=manifest["thresholds"][
|
||||
"final_bpc_delta_mean_maximum"
|
||||
],
|
||||
)
|
||||
for variant, variant_runs in runs_by_variant.items()
|
||||
}
|
||||
effects = {
|
||||
variant: variant_effect(
|
||||
variant,
|
||||
variant_runs,
|
||||
references,
|
||||
metrics_by_cell,
|
||||
step=primary_step,
|
||||
threshold=manifest["thresholds"]["material_relative_drop"],
|
||||
quality=qualities[variant],
|
||||
)
|
||||
for variant, variant_runs in runs_by_variant.items()
|
||||
}
|
||||
primary = effects[manifest["primary_variant"]]
|
||||
if primary["attenuation_passed"] and primary["quality"]["passed"]:
|
||||
status = (
|
||||
"forward_training_attenuation_established_within_reduced_protocol"
|
||||
)
|
||||
elif primary["attenuation_passed"]:
|
||||
status = "quality_guard_failed"
|
||||
elif primary["quality"]["passed"]:
|
||||
status = "attenuation_not_established"
|
||||
else:
|
||||
status = "attenuation_and_quality_failed"
|
||||
secondary = {
|
||||
variant: (
|
||||
"secondary_material_response"
|
||||
if effect["material_response_passed"]
|
||||
else "secondary_response_not_established"
|
||||
)
|
||||
for variant, effect in effects.items()
|
||||
if variant != manifest["primary_variant"]
|
||||
}
|
||||
interaction = interaction_map(metrics_by_cell, seeds, steps)
|
||||
|
||||
trajectories = []
|
||||
final_spectra = []
|
||||
for variant in ("learned_reference",) + variants:
|
||||
for seed in seeds:
|
||||
for step in steps:
|
||||
record = metrics_by_cell[(variant, seed, step)]
|
||||
reference = metrics_by_cell[
|
||||
("learned_reference", seed, step)
|
||||
]
|
||||
trajectories.append(
|
||||
{
|
||||
"variant": variant,
|
||||
"seed": seed,
|
||||
"step": step,
|
||||
"spike_mean": record["spike_mean"],
|
||||
"reference_mean": record["reference_mean"],
|
||||
"spike_contrast": record["spike_contrast"],
|
||||
"peak_normalized": record["peak_normalized"],
|
||||
"relative_drop": {
|
||||
metric: (
|
||||
reference[metric] - record[metric]
|
||||
)
|
||||
/ reference[metric]
|
||||
for metric in METRICS
|
||||
},
|
||||
}
|
||||
)
|
||||
final = metrics_by_cell[(variant, seed, primary_step)]
|
||||
final_spectra.append(
|
||||
{
|
||||
"variant": variant,
|
||||
"seed": seed,
|
||||
**final,
|
||||
}
|
||||
)
|
||||
|
||||
input_files = {
|
||||
"manifest": {
|
||||
"path": str(args.manifest),
|
||||
"sha256": file_sha256(args.manifest),
|
||||
},
|
||||
"formal": [
|
||||
{
|
||||
"variant": variant,
|
||||
"seed": seed,
|
||||
"path": str(run_paths[(variant, seed)]),
|
||||
"sha256": file_sha256(run_paths[(variant, seed)]),
|
||||
}
|
||||
for variant, seed in sorted(runs)
|
||||
],
|
||||
"references": [
|
||||
{
|
||||
"seed": seed,
|
||||
"path": str(reference_paths[seed]),
|
||||
"sha256": file_sha256(reference_paths[seed]),
|
||||
}
|
||||
for seed in seeds
|
||||
],
|
||||
"replay": {
|
||||
"path": str(args.replay),
|
||||
"sha256": file_sha256(args.replay),
|
||||
},
|
||||
}
|
||||
aggregate = {
|
||||
"schema_version": 1,
|
||||
"protocol_id": PROTOCOL_ID,
|
||||
"status": status,
|
||||
"scope": (
|
||||
"depth-32 reduced Block AttnRes train-time architecture "
|
||||
"ablation; not a real Kimi-K3 checkpoint result"
|
||||
),
|
||||
"primary_step": primary_step,
|
||||
"spike_layers_1based": list(spike_layers),
|
||||
"thresholds": manifest["thresholds"],
|
||||
"input_files": input_files,
|
||||
"historical_pairing": pairing,
|
||||
"replay": {
|
||||
"passed": replay_exact,
|
||||
"scientific_payload_sha256": canonical_sha256(formal_payload),
|
||||
"excluded": [
|
||||
"run_kind",
|
||||
"timing",
|
||||
"self hashes",
|
||||
"manifest path strings",
|
||||
"GPU/version metadata",
|
||||
],
|
||||
},
|
||||
"primary": primary,
|
||||
"secondary_status": secondary,
|
||||
"effects": effects,
|
||||
"interaction": interaction,
|
||||
"trajectories": trajectories,
|
||||
"final_spectra": final_spectra,
|
||||
"processed_target_bytes": manifest["new_target_bytes"],
|
||||
"historical_reference_target_bytes": (
|
||||
manifest["historical_reference_target_bytes"]
|
||||
),
|
||||
"reporting_boundary": (
|
||||
"C can change through spike-window numerator and the 27-layer "
|
||||
"reference denominator; layers 26-28 are intervened but belong "
|
||||
"to the denominator."
|
||||
),
|
||||
}
|
||||
write_hashed(args.aggregate_output, aggregate)
|
||||
|
||||
compact = {
|
||||
"schema_version": 1,
|
||||
"protocol_id": PROTOCOL_ID,
|
||||
"status": status,
|
||||
"primary_step": primary_step,
|
||||
"spike_layers_1based": list(spike_layers),
|
||||
"thresholds": manifest["thresholds"],
|
||||
"primary": primary,
|
||||
"secondary_status": secondary,
|
||||
"effects": effects,
|
||||
"interaction": interaction,
|
||||
"trajectories": trajectories,
|
||||
"final_spectra": final_spectra,
|
||||
"replay": aggregate["replay"],
|
||||
"processed_target_bytes": manifest["new_target_bytes"],
|
||||
"reporting_boundary": aggregate["reporting_boundary"],
|
||||
"aggregate_sha256": aggregate["canonical_sha256_without_self"],
|
||||
}
|
||||
write_hashed(args.compact_output, compact)
|
||||
|
||||
reproduction = {
|
||||
"schema_version": 1,
|
||||
"protocol_id": PROTOCOL_ID,
|
||||
"passed": replay_exact,
|
||||
"formal_variant": replay_contract["variant"],
|
||||
"seed": replay_contract["seed"],
|
||||
"formal_file_sha256": file_sha256(
|
||||
run_paths[
|
||||
(replay_contract["variant"], replay_contract["seed"])
|
||||
]
|
||||
),
|
||||
"replay_file_sha256": file_sha256(args.replay),
|
||||
"scientific_payload_sha256": canonical_sha256(formal_payload),
|
||||
"excluded_fields": aggregate["replay"]["excluded"],
|
||||
}
|
||||
write_hashed(args.reproduction_output, reproduction)
|
||||
print(
|
||||
json.dumps(
|
||||
{
|
||||
"status": status,
|
||||
"primary_attenuation": {
|
||||
"passed_cells": primary["passed_cells"],
|
||||
"required_cells": primary["required_cells"],
|
||||
},
|
||||
"primary_quality": {
|
||||
"passed_checks": primary["quality"]["passed_checks"],
|
||||
"required_checks": primary["quality"]["required_checks"],
|
||||
},
|
||||
"replay_exact": replay_exact,
|
||||
"aggregate": str(args.aggregate_output),
|
||||
"compact": str(args.compact_output),
|
||||
},
|
||||
ensure_ascii=False,
|
||||
indent=2,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,39 @@
|
||||
{
|
||||
"canonical_sha256_without_self": "99aeefc0ba35c199725ed7af377450ecbb5cc5f1aabdda9ac2345a30f03b444f",
|
||||
"excluded_fields": [
|
||||
"protocol wrapper fields",
|
||||
"timing",
|
||||
"self hash",
|
||||
"parent runner self hash",
|
||||
"study manifest"
|
||||
],
|
||||
"field_checks": {
|
||||
"architecture": true,
|
||||
"batch_size": true,
|
||||
"depth": true,
|
||||
"diagnostics": true,
|
||||
"environment": true,
|
||||
"evaluations": true,
|
||||
"gradient_gate": true,
|
||||
"hashes": true,
|
||||
"manifest": true,
|
||||
"model": true,
|
||||
"optimizer": true,
|
||||
"seed": true,
|
||||
"steps": true,
|
||||
"target_bytes_seen": true,
|
||||
"training_history": true
|
||||
},
|
||||
"gate": "empty-selector-parent-equivalence",
|
||||
"parent_file": "/home/wuyang/Code/K3/experiments/k3/attnres_forward/results/gates/parent-smoke.json",
|
||||
"passed": true,
|
||||
"protocol_id": "llm-atlas-k3-attnres-forward-training-v1",
|
||||
"schema_version": 1,
|
||||
"wrapper_file": "/home/wuyang/Code/K3/experiments/k3/attnres_forward/results/gates/wrapper-smoke.json",
|
||||
"wrapper_identity": {
|
||||
"forward_audit": true,
|
||||
"parent_protocol": true,
|
||||
"protocol": true,
|
||||
"variant": true
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,190 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Run the frozen Round 08 matrix with at most two isolated processes."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
|
||||
VARIANTS = (
|
||||
"uniform_group_6_forward",
|
||||
"uniform_group_7_forward",
|
||||
"uniform_groups_6_7_forward",
|
||||
"uniform_group_7_mlp_forward",
|
||||
)
|
||||
SEEDS = (2026073001, 2026073002, 2026073003)
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--python", type=Path, required=True)
|
||||
parser.add_argument("--cache-dir", type=Path, required=True)
|
||||
parser.add_argument("--parent-manifest", type=Path, required=True)
|
||||
parser.add_argument("--study-manifest", type=Path, required=True)
|
||||
parser.add_argument("--output-dir", type=Path, required=True)
|
||||
parser.add_argument(
|
||||
"--phase", choices=("formal", "replay", "all"), default="all"
|
||||
)
|
||||
parser.add_argument("--concurrency", type=int, default=2)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def cell_output(
|
||||
output_dir: Path, variant: str, seed: int, run_kind: str
|
||||
) -> Path:
|
||||
return output_dir / (
|
||||
f"{run_kind}-{variant}-seed-{seed}.json"
|
||||
)
|
||||
|
||||
|
||||
def command_for(
|
||||
args: argparse.Namespace, variant: str, seed: int, run_kind: str
|
||||
) -> list[str]:
|
||||
runner = Path(__file__).resolve().parent / "train.py"
|
||||
return [
|
||||
str(args.python),
|
||||
str(runner),
|
||||
"--variant",
|
||||
variant,
|
||||
"--study-manifest",
|
||||
str(args.study_manifest),
|
||||
"--run-kind",
|
||||
run_kind,
|
||||
"--architecture",
|
||||
"block",
|
||||
"--depth",
|
||||
"32",
|
||||
"--seed",
|
||||
str(seed),
|
||||
"--cache-dir",
|
||||
str(args.cache_dir),
|
||||
"--manifest",
|
||||
str(args.parent_manifest),
|
||||
"--output",
|
||||
str(cell_output(args.output_dir, variant, seed, run_kind)),
|
||||
]
|
||||
|
||||
|
||||
def validate_manifest(args: argparse.Namespace) -> None:
|
||||
manifest = json.loads(args.study_manifest.read_text())
|
||||
if (
|
||||
manifest["status"] != "frozen-before-model-output"
|
||||
or tuple(manifest["variants"]) != VARIANTS
|
||||
or tuple(manifest["formal_seeds"]) != SEEDS
|
||||
or manifest["concurrency_maximum"] != 2
|
||||
):
|
||||
raise RuntimeError("study manifest matrix/concurrency drift")
|
||||
if args.concurrency < 1 or args.concurrency > 2:
|
||||
raise ValueError("the frozen protocol permits one or two processes")
|
||||
|
||||
|
||||
def run_cells(
|
||||
args: argparse.Namespace,
|
||||
cells: list[tuple[str, int, str]],
|
||||
) -> None:
|
||||
args.output_dir.mkdir(parents=True, exist_ok=True)
|
||||
for variant, seed, run_kind in cells:
|
||||
output = cell_output(args.output_dir, variant, seed, run_kind)
|
||||
if output.exists():
|
||||
raise FileExistsError(
|
||||
f"refusing to overwrite existing result: {output}"
|
||||
)
|
||||
|
||||
environment = dict(os.environ)
|
||||
environment["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
|
||||
pending = list(cells)
|
||||
running: list[dict[str, Any]] = []
|
||||
completed = 0
|
||||
while pending or running:
|
||||
while pending and len(running) < args.concurrency:
|
||||
variant, seed, run_kind = pending.pop(0)
|
||||
command = command_for(args, variant, seed, run_kind)
|
||||
process = subprocess.Popen(command, env=environment)
|
||||
running.append(
|
||||
{
|
||||
"identity": (variant, seed, run_kind),
|
||||
"process": process,
|
||||
"started": time.monotonic(),
|
||||
}
|
||||
)
|
||||
print(
|
||||
json.dumps(
|
||||
{
|
||||
"event": "cell_started",
|
||||
"variant": variant,
|
||||
"seed": seed,
|
||||
"run_kind": run_kind,
|
||||
"pid": process.pid,
|
||||
"active": len(running),
|
||||
"remaining": len(pending),
|
||||
},
|
||||
sort_keys=True,
|
||||
),
|
||||
flush=True,
|
||||
)
|
||||
time.sleep(1)
|
||||
survivors = []
|
||||
for item in running:
|
||||
return_code = item["process"].poll()
|
||||
if return_code is None:
|
||||
survivors.append(item)
|
||||
continue
|
||||
variant, seed, run_kind = item["identity"]
|
||||
elapsed = time.monotonic() - item["started"]
|
||||
if return_code != 0:
|
||||
for survivor in survivors:
|
||||
survivor["process"].terminate()
|
||||
for survivor in running:
|
||||
if survivor is not item and survivor not in survivors:
|
||||
survivor["process"].terminate()
|
||||
raise RuntimeError(
|
||||
f"cell failed: {variant}/{seed}/{run_kind}: {return_code}"
|
||||
)
|
||||
completed += 1
|
||||
print(
|
||||
json.dumps(
|
||||
{
|
||||
"event": "cell_completed",
|
||||
"variant": variant,
|
||||
"seed": seed,
|
||||
"run_kind": run_kind,
|
||||
"elapsed_seconds": elapsed,
|
||||
"completed": completed,
|
||||
"total": len(cells),
|
||||
},
|
||||
sort_keys=True,
|
||||
),
|
||||
flush=True,
|
||||
)
|
||||
running = survivors
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
validate_manifest(args)
|
||||
formal = [
|
||||
(variant, seed, "formal")
|
||||
for variant in VARIANTS
|
||||
for seed in SEEDS
|
||||
]
|
||||
replay = [
|
||||
("uniform_groups_6_7_forward", 2026073001, "replay")
|
||||
]
|
||||
cells = (
|
||||
formal
|
||||
if args.phase == "formal"
|
||||
else replay
|
||||
if args.phase == "replay"
|
||||
else formal + replay
|
||||
)
|
||||
run_cells(args, cells)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,442 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Run one preregistered Round 08 train-time uniform-forward cell."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib.util
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
PROTOCOL_ID = "llm-atlas-k3-attnres-forward-training-v1"
|
||||
PARENT_PROTOCOL_ID = "llm-atlas-k3-attnres-gradient-scale-v1"
|
||||
VARIANTS = {
|
||||
"learned_reference": (),
|
||||
"uniform_group_6_forward": tuple(range(40, 48)),
|
||||
"uniform_group_7_forward": tuple(range(48, 56)),
|
||||
"uniform_groups_6_7_forward": tuple(range(40, 56)),
|
||||
"uniform_group_7_mlp_forward": (49, 51, 53, 55),
|
||||
}
|
||||
FORMAL_VARIANTS = tuple(name for name in VARIANTS if name != "learned_reference")
|
||||
EXPECTED_SOURCE_COUNTS = {
|
||||
**{40: 6},
|
||||
**{index: 7 for index in range(41, 49)},
|
||||
**{index: 8 for index in range(49, 56)},
|
||||
}
|
||||
|
||||
|
||||
def load_parent_module() -> Any:
|
||||
path = Path(__file__).resolve().parents[1] / "attnres_gradient" / "train.py"
|
||||
spec = importlib.util.spec_from_file_location("k3_attnres_round05_train", path)
|
||||
if spec is None or spec.loader is None:
|
||||
raise RuntimeError(f"cannot import Round 05 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_module()
|
||||
ACTIVE_VARIANT = "learned_reference"
|
||||
LAST_MODEL: ForwardInterventionLanguageModel | None = None
|
||||
LAST_OPTIMIZER: torch.optim.Optimizer | None = None
|
||||
|
||||
|
||||
def extract_wrapper_argument(name: str) -> str:
|
||||
try:
|
||||
index = sys.argv.index(name)
|
||||
except ValueError as error:
|
||||
raise ValueError(f"missing required wrapper argument: {name}") from error
|
||||
if index + 1 >= len(sys.argv):
|
||||
raise ValueError(f"missing value for wrapper argument: {name}")
|
||||
value = sys.argv[index + 1]
|
||||
del sys.argv[index : index + 2]
|
||||
return value
|
||||
|
||||
|
||||
def argument_value(name: str, default: str | None = None) -> str | None:
|
||||
try:
|
||||
index = sys.argv.index(name)
|
||||
except ValueError:
|
||||
return default
|
||||
if index + 1 >= len(sys.argv):
|
||||
raise ValueError(f"missing value for argument: {name}")
|
||||
return sys.argv[index + 1]
|
||||
|
||||
|
||||
def parameter_names_for_indices(indices: tuple[int, ...]) -> tuple[str, ...]:
|
||||
names = []
|
||||
for index in indices:
|
||||
names.extend(
|
||||
(
|
||||
f"mixers.{index}.query",
|
||||
f"mixers.{index}.key_norm.weight",
|
||||
)
|
||||
)
|
||||
return tuple(names)
|
||||
|
||||
|
||||
class ForwardInterventionLanguageModel(parent.GradientLanguageModel):
|
||||
"""Round 05 model with one frozen selector and parameter-free uniform mixers."""
|
||||
|
||||
def __init__(self, architecture: str):
|
||||
super().__init__(architecture)
|
||||
global LAST_MODEL
|
||||
if architecture != "block":
|
||||
raise ValueError("Round 08 only permits the block architecture")
|
||||
if ACTIVE_VARIANT not in VARIANTS:
|
||||
raise ValueError(f"unknown Round 08 variant: {ACTIVE_VARIANT}")
|
||||
self.forward_variant = ACTIVE_VARIANT
|
||||
self.selected_indices = tuple(VARIANTS[ACTIVE_VARIANT])
|
||||
self.selected_set = frozenset(self.selected_indices)
|
||||
self.forward_calls = 0
|
||||
self.depth_visits = [0] * len(self.mixers)
|
||||
self.output_visits = 0
|
||||
self.source_counts: dict[int, set[int]] = {
|
||||
index: set() for index in range(len(self.mixers))
|
||||
}
|
||||
self.uniform_weight_max_abs_error = 0.0
|
||||
selected_names = parameter_names_for_indices(self.selected_indices)
|
||||
named_parameters = dict(self.named_parameters())
|
||||
self.selected_initial_tensors = {
|
||||
name: named_parameters[name].detach().cpu().clone()
|
||||
for name in selected_names
|
||||
}
|
||||
self.gradient_hook_calls = {
|
||||
name: 0
|
||||
for name in named_parameters
|
||||
if name.startswith("mixers.") or name.startswith("output_mixer.")
|
||||
}
|
||||
self._gradient_hooks = []
|
||||
for name, parameter in named_parameters.items():
|
||||
if name not in self.gradient_hook_calls:
|
||||
continue
|
||||
|
||||
def count_hook(
|
||||
gradient: torch.Tensor, *, parameter_name: str = name
|
||||
) -> torch.Tensor:
|
||||
self.gradient_hook_calls[parameter_name] += 1
|
||||
return gradient
|
||||
|
||||
self._gradient_hooks.append(parameter.register_hook(count_hook))
|
||||
LAST_MODEL = self
|
||||
|
||||
def mix(
|
||||
self,
|
||||
mixer_index: int,
|
||||
sources: list[torch.Tensor],
|
||||
capture: bool,
|
||||
) -> tuple[torch.Tensor, dict[str, Any] | None]:
|
||||
self.depth_visits[mixer_index] += 1
|
||||
self.source_counts[mixer_index].add(len(sources))
|
||||
if mixer_index not in self.selected_set:
|
||||
return self.mixers[mixer_index](sources, capture)
|
||||
|
||||
values = torch.stack(sources, dim=0)
|
||||
logits = torch.zeros(
|
||||
values.shape[0],
|
||||
values.shape[1],
|
||||
values.shape[2],
|
||||
dtype=torch.float32,
|
||||
device=values.device,
|
||||
)
|
||||
weights = torch.softmax(logits, dim=0)
|
||||
expected = torch.tensor(
|
||||
1.0 / len(sources), dtype=weights.dtype, device=weights.device
|
||||
)
|
||||
error = (weights - expected).abs().max().detach().cpu().item()
|
||||
self.uniform_weight_max_abs_error = max(
|
||||
self.uniform_weight_max_abs_error, error
|
||||
)
|
||||
output = torch.einsum(
|
||||
"nbt,nbtd->btd", weights, values.float()
|
||||
).to(values.dtype)
|
||||
if not capture:
|
||||
return output, None
|
||||
entropy = -(weights * torch.log(weights.clamp_min(1e-30))).sum(dim=0)
|
||||
return output, {
|
||||
"mean_weights": weights.mean(dim=(1, 2)).detach().cpu().tolist(),
|
||||
"entropy_mean": entropy.mean().detach().cpu().item(),
|
||||
"sources": len(sources),
|
||||
}
|
||||
|
||||
def forward(
|
||||
self, input_ids: torch.Tensor, capture: bool = False
|
||||
) -> tuple[torch.Tensor, parent.ActivationTrace | None]:
|
||||
if not self.selected_indices:
|
||||
return super().forward(input_ids, capture)
|
||||
|
||||
self.forward_calls += 1
|
||||
embedded = self.embed(input_ids)
|
||||
trace = parent.ActivationTrace([], [], [], [], []) if capture else None
|
||||
completed = [embedded]
|
||||
partial: torch.Tensor | None = None
|
||||
mixer_index = 0
|
||||
for block in self.blocks:
|
||||
for branch_index in range(2):
|
||||
sources = completed + ([] if partial is None else [partial])
|
||||
branch_input, weights = self.mix(
|
||||
mixer_index, sources, capture
|
||||
)
|
||||
mixer_index += 1
|
||||
if branch_index == 0:
|
||||
branch_output = block.attention(
|
||||
block.attention_norm(branch_input)
|
||||
)
|
||||
else:
|
||||
branch_output = block.mlp(block.mlp_norm(branch_input))
|
||||
branch_for_residual = branch_output.float()
|
||||
partial = (
|
||||
branch_for_residual
|
||||
if partial is None
|
||||
else partial + branch_for_residual
|
||||
)
|
||||
if trace is not None:
|
||||
trace.layer_input_rms.append(parent.rms(branch_input))
|
||||
trace.branch_output_rms.append(parent.rms(branch_output))
|
||||
trace.stream_state_rms.append(parent.rms(partial))
|
||||
trace.depth_weights.append(weights or {})
|
||||
if branch_index == 1:
|
||||
partial.retain_grad()
|
||||
trace.block_outputs.append(partial)
|
||||
if mixer_index % parent.round04.SUBLAYERS_PER_BLOCK == 0:
|
||||
completed.append(partial)
|
||||
partial = None
|
||||
if partial is not None or len(completed) != parent.BLOCK_GROUPS + 1:
|
||||
raise RuntimeError("Round 08 Block AttnRes aggregation failed")
|
||||
if self.output_mixer is None:
|
||||
raise RuntimeError("Round 08 output mixer missing")
|
||||
self.output_visits += 1
|
||||
hidden, output_weights = self.output_mixer(completed, capture)
|
||||
if trace is not None:
|
||||
trace.output_weights = output_weights
|
||||
normalized = self.final_norm(hidden)
|
||||
logits = F.linear(normalized, self.token_embedding.weight)
|
||||
return logits, trace
|
||||
|
||||
|
||||
def tensor_exact(left: torch.Tensor, right: torch.Tensor) -> bool:
|
||||
return (
|
||||
left.dtype == right.dtype
|
||||
and tuple(left.shape) == tuple(right.shape)
|
||||
and torch.equal(left.detach().cpu(), right.detach().cpu())
|
||||
)
|
||||
|
||||
|
||||
def build_intervention_audit(
|
||||
model: ForwardInterventionLanguageModel,
|
||||
optimizer: torch.optim.Optimizer,
|
||||
study_manifest: dict[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
selected = tuple(model.selected_indices)
|
||||
selected_names = set(parameter_names_for_indices(selected))
|
||||
mixer_parameters = {
|
||||
name: parameter
|
||||
for name, parameter in model.named_parameters()
|
||||
if name.startswith("mixers.") or name.startswith("output_mixer.")
|
||||
}
|
||||
optimizer_parameters = {
|
||||
parameter
|
||||
for group in optimizer.param_groups
|
||||
for parameter in group["params"]
|
||||
}
|
||||
selected_parameter_checks = {}
|
||||
for name in sorted(selected_names):
|
||||
parameter = mixer_parameters[name]
|
||||
selected_parameter_checks[name] = {
|
||||
"gradient_hook_calls": model.gradient_hook_calls[name],
|
||||
"in_optimizer_param_group": parameter in optimizer_parameters,
|
||||
"optimizer_state_present": parameter in optimizer.state,
|
||||
"final_equals_initial": tensor_exact(
|
||||
parameter, model.selected_initial_tensors[name]
|
||||
),
|
||||
}
|
||||
unselected_parameter_checks = {}
|
||||
for name, parameter in sorted(mixer_parameters.items()):
|
||||
if name in selected_names:
|
||||
continue
|
||||
unselected_parameter_checks[name] = {
|
||||
"gradient_hook_calls": model.gradient_hook_calls[name],
|
||||
"in_optimizer_param_group": parameter in optimizer_parameters,
|
||||
"optimizer_state_present": parameter in optimizer.state,
|
||||
}
|
||||
|
||||
source_counts = {
|
||||
str(index): sorted(values)
|
||||
for index, values in model.source_counts.items()
|
||||
}
|
||||
selected_source_gate = {
|
||||
str(index): (
|
||||
source_counts[str(index)]
|
||||
== [study_manifest["selected_source_counts"][str(index)]]
|
||||
)
|
||||
for index in selected
|
||||
}
|
||||
visit_gate = (
|
||||
all(value == model.forward_calls for value in model.depth_visits)
|
||||
and model.output_visits == model.forward_calls
|
||||
)
|
||||
selected_parameter_gate = all(
|
||||
check["gradient_hook_calls"] == 0
|
||||
and check["in_optimizer_param_group"]
|
||||
and not check["optimizer_state_present"]
|
||||
and check["final_equals_initial"]
|
||||
for check in selected_parameter_checks.values()
|
||||
)
|
||||
unselected_parameter_gate = all(
|
||||
check["gradient_hook_calls"] > 0
|
||||
and check["in_optimizer_param_group"]
|
||||
and check["optimizer_state_present"]
|
||||
for check in unselected_parameter_checks.values()
|
||||
)
|
||||
expected_selected = tuple(
|
||||
study_manifest["variants"]
|
||||
.get(model.forward_variant, {"selected_depth_indices": []})[
|
||||
"selected_depth_indices"
|
||||
]
|
||||
)
|
||||
selector_gate = (
|
||||
selected == expected_selected
|
||||
and 64 not in selected
|
||||
and selected_source_gate == {
|
||||
str(index): True for index in selected
|
||||
}
|
||||
)
|
||||
threshold = study_manifest["thresholds"][
|
||||
"uniform_weight_max_abs_error"
|
||||
]
|
||||
uniform_gate = model.uniform_weight_max_abs_error <= threshold
|
||||
passed = (
|
||||
visit_gate
|
||||
and selector_gate
|
||||
and selected_parameter_gate
|
||||
and unselected_parameter_gate
|
||||
and uniform_gate
|
||||
)
|
||||
return {
|
||||
"passed": passed,
|
||||
"variant": model.forward_variant,
|
||||
"selected_depth_indices": list(selected),
|
||||
"output_mixer_selected": False,
|
||||
"forward_calls": model.forward_calls,
|
||||
"depth_visit_counts": model.depth_visits,
|
||||
"output_visit_count": model.output_visits,
|
||||
"visit_gate": visit_gate,
|
||||
"source_counts_by_depth_index": source_counts,
|
||||
"selected_source_count_checks": selected_source_gate,
|
||||
"selector_gate": selector_gate,
|
||||
"uniform_weight_max_abs_error": model.uniform_weight_max_abs_error,
|
||||
"uniform_weight_threshold": threshold,
|
||||
"uniform_weight_gate": uniform_gate,
|
||||
"selected_parameters": selected_parameter_checks,
|
||||
"selected_parameter_reachability_gate": selected_parameter_gate,
|
||||
"unselected_parameters": unselected_parameter_checks,
|
||||
"unselected_parameter_reachability_gate": unselected_parameter_gate,
|
||||
"semantics": (
|
||||
"selected depth mixers use parameter-free constant-zero logits "
|
||||
"with the parent softmax+einsum arithmetic kernel"
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def rewrite_result(
|
||||
output_path: Path,
|
||||
study_manifest_path: Path,
|
||||
study_manifest: dict[str, Any],
|
||||
) -> None:
|
||||
if LAST_MODEL is None or LAST_OPTIMIZER is None:
|
||||
raise RuntimeError("runner capture state missing")
|
||||
result = json.loads(output_path.read_text())
|
||||
parent_self_hash = result.pop("canonical_sha256_without_self")
|
||||
if result["protocol_id"] != PARENT_PROTOCOL_ID:
|
||||
raise RuntimeError("parent runner protocol drift")
|
||||
result["schema_version"] = 2
|
||||
result["protocol_id"] = PROTOCOL_ID
|
||||
result["parent_protocol_id"] = PARENT_PROTOCOL_ID
|
||||
result["variant"] = ACTIVE_VARIANT
|
||||
result["parent_runner_canonical_sha256"] = parent_self_hash
|
||||
result["study_manifest"] = {
|
||||
"path": str(study_manifest_path),
|
||||
"file_sha256": parent.file_sha256(study_manifest_path),
|
||||
"status": study_manifest["status"],
|
||||
}
|
||||
result["forward_intervention"] = build_intervention_audit(
|
||||
LAST_MODEL, LAST_OPTIMIZER, study_manifest
|
||||
)
|
||||
result["canonical_sha256_without_self"] = parent.canonical_sha256(result)
|
||||
temporary = output_path.with_suffix(output_path.suffix + ".round08.tmp")
|
||||
temporary.write_text(
|
||||
json.dumps(result, ensure_ascii=False, indent=2, sort_keys=True) + "\n"
|
||||
)
|
||||
os.replace(temporary, output_path)
|
||||
if not result["forward_intervention"]["passed"]:
|
||||
raise RuntimeError(
|
||||
f"forward intervention audit failed: "
|
||||
f"{result['forward_intervention']}"
|
||||
)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
global ACTIVE_VARIANT, LAST_OPTIMIZER
|
||||
variant = extract_wrapper_argument("--variant")
|
||||
study_manifest_path = Path(
|
||||
extract_wrapper_argument("--study-manifest")
|
||||
).resolve()
|
||||
if variant not in VARIANTS:
|
||||
raise ValueError(f"unknown variant: {variant}")
|
||||
run_kind = argument_value("--run-kind", "formal")
|
||||
if run_kind in ("formal", "replay") and variant not in FORMAL_VARIANTS:
|
||||
raise ValueError("learned_reference is smoke-only")
|
||||
if argument_value("--architecture") != "block":
|
||||
raise ValueError("Round 08 requires --architecture block")
|
||||
if argument_value("--depth") != "32":
|
||||
raise ValueError("Round 08 requires --depth 32")
|
||||
if run_kind == "replay" and variant != "uniform_groups_6_7_forward":
|
||||
raise ValueError("the frozen replay uses the primary joint variant")
|
||||
|
||||
study_manifest = json.loads(study_manifest_path.read_text())
|
||||
if (
|
||||
study_manifest["protocol_id"] != PROTOCOL_ID
|
||||
or study_manifest["status"] != "frozen-before-model-output"
|
||||
):
|
||||
raise ValueError("study manifest is not the frozen Round 08 contract")
|
||||
expected = tuple(
|
||||
study_manifest["variants"]
|
||||
.get(variant, {"selected_depth_indices": []})[
|
||||
"selected_depth_indices"
|
||||
]
|
||||
)
|
||||
if expected != VARIANTS[variant]:
|
||||
raise ValueError("study manifest selector drift")
|
||||
|
||||
output_value = argument_value("--output")
|
||||
if output_value is None:
|
||||
raise ValueError("--output is required")
|
||||
output_path = Path(output_value).resolve()
|
||||
ACTIVE_VARIANT = variant
|
||||
parent.GradientLanguageModel = ForwardInterventionLanguageModel
|
||||
|
||||
original_adamw = torch.optim.AdamW
|
||||
|
||||
def capture_adamw(*args: Any, **kwargs: Any) -> torch.optim.Optimizer:
|
||||
global LAST_OPTIMIZER
|
||||
LAST_OPTIMIZER = original_adamw(*args, **kwargs)
|
||||
return LAST_OPTIMIZER
|
||||
|
||||
torch.optim.AdamW = capture_adamw # type: ignore[assignment]
|
||||
try:
|
||||
parent.main()
|
||||
finally:
|
||||
torch.optim.AdamW = original_adamw # type: ignore[assignment]
|
||||
rewrite_result(output_path, study_manifest_path, study_manifest)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -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()
|
||||
@@ -169,7 +169,9 @@ uniform。选中分支用同 dtype 的 constant-zero logits 和同一个
|
||||
- CE byte-exact；
|
||||
- activation-gradient spectrum byte-exact；
|
||||
- validation metrics byte-exact；
|
||||
- selected capture weights 等于 `1/N`，max absolute error `≤ 1e-12`；
|
||||
- selected 的**归约前 weight tensor** 等于 FP32 `1/N`，max absolute error
|
||||
`≤ 1e-12`；capture 的 `mean_weights` 因 FP32 大规模 mean 可有约 `1e-8` 的归约舍入，
|
||||
但必须与父 learned capture summary byte-exact；
|
||||
- 若任何跨 variant byte-exact 比较失败，hard-fail；不得在结果后改成容差 gate。
|
||||
|
||||
这个负控制只约束初始化；训练开始后 forward 必须允许分化。
|
||||
|
||||
Reference in New Issue
Block a user