wuyang
443 lines
16 KiB
Python
443 lines
16 KiB
Python
#!/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()
|