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