research: lock AttnRes spike diagnostic runner
This commit is contained in:
@@ -0,0 +1,40 @@
|
|||||||
|
# Attention Residuals spike-path diagnostics
|
||||||
|
|
||||||
|
This directory implements preregistered protocol
|
||||||
|
`llm-atlas-k3-attnres-spike-path-v1`.
|
||||||
|
|
||||||
|
It is a targeted follow-up to Round 05. It replays the exact depth-32 Block
|
||||||
|
training contract and adds diagnostic-only activation positions, gradient
|
||||||
|
reductions, and same-forward backward-rule interventions. It is not a Kimi K3
|
||||||
|
checkpoint run and does not recover the paper's 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
|
||||||
|
```
|
||||||
|
|
||||||
|
## Step-0 smoke
|
||||||
|
|
||||||
|
```bash
|
||||||
|
CUBLAS_WORKSPACE_CONFIG=:4096:8 \
|
||||||
|
/home/wuyang/.pyenv/versions/3.10.14/envs/navi-router-cu128/bin/python \
|
||||||
|
experiments/k3/attnres_spike/train.py \
|
||||||
|
--run-kind smoke \
|
||||||
|
--seed 2026073001 \
|
||||||
|
--cache-dir /home/wuyang/.cache/llm-atlas/k3-attnres-gradient-scale-v1 \
|
||||||
|
--parent-manifest experiments/k3/attnres_gradient/manifest.json \
|
||||||
|
--manifest experiments/k3/attnres_spike/manifest.json \
|
||||||
|
--output /home/wuyang/.cache/llm-atlas/k3-attnres-spike-path-v1/smoke/seed-2026073001.json
|
||||||
|
```
|
||||||
|
|
||||||
|
Formal cells use `--run-kind formal` and all three preregistered seeds. The
|
||||||
|
independent replay uses `--run-kind replay --seed 2026073001`. Formal and replay
|
||||||
|
runs are fixed to 8,000 steps; smoke performs the complete step-0 diagnostic
|
||||||
|
gate without an optimizer step.
|
||||||
|
|
||||||
|
Raw outputs are copied into `results/raw/` only after training equivalence,
|
||||||
|
forward identity, loss-scale, reduction, and replay gates pass.
|
||||||
@@ -0,0 +1,178 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""Rebuild the Round 06 scoping table from frozen Round 05 raw files."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
import json
|
||||||
|
import math
|
||||||
|
import statistics
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Any, Callable
|
||||||
|
|
||||||
|
|
||||||
|
SEEDS = (2026073001, 2026073002, 2026073003)
|
||||||
|
|
||||||
|
|
||||||
|
def parse_args() -> argparse.Namespace:
|
||||||
|
parser = argparse.ArgumentParser()
|
||||||
|
parser.add_argument(
|
||||||
|
"--raw-dir",
|
||||||
|
type=Path,
|
||||||
|
default=Path(__file__).resolve().parents[1]
|
||||||
|
/ "attnres_gradient"
|
||||||
|
/ "results"
|
||||||
|
/ "raw",
|
||||||
|
)
|
||||||
|
return parser.parse_args()
|
||||||
|
|
||||||
|
|
||||||
|
def pearson(left: list[float], right: list[float]) -> float:
|
||||||
|
left_mean = statistics.fmean(left)
|
||||||
|
right_mean = statistics.fmean(right)
|
||||||
|
numerator = sum(
|
||||||
|
(x - left_mean) * (y - right_mean)
|
||||||
|
for x, y in zip(left, right)
|
||||||
|
)
|
||||||
|
left_square = sum((value - left_mean) ** 2 for value in left)
|
||||||
|
right_square = sum((value - right_mean) ** 2 for value in right)
|
||||||
|
return numerator / math.sqrt(left_square * right_square)
|
||||||
|
|
||||||
|
|
||||||
|
def mixer_metrics(value: dict[str, Any]) -> dict[str, float | int]:
|
||||||
|
weights = value["mean_weights"]
|
||||||
|
sources = value["sources"]
|
||||||
|
return {
|
||||||
|
"sources": sources,
|
||||||
|
"latest": weights[-1],
|
||||||
|
"maximum": max(weights),
|
||||||
|
"normalized_entropy": (
|
||||||
|
1.0
|
||||||
|
if sources == 1
|
||||||
|
else value["entropy_mean"] / math.log(sources)
|
||||||
|
),
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def average(rows: list[dict[str, Any]], key: Callable[[dict[str, Any]], float]) -> float:
|
||||||
|
return statistics.fmean(key(row) for row in rows)
|
||||||
|
|
||||||
|
|
||||||
|
def main() -> None:
|
||||||
|
args = parse_args()
|
||||||
|
rows = []
|
||||||
|
source_hashes = {}
|
||||||
|
for seed in SEEDS:
|
||||||
|
path = (
|
||||||
|
args.raw_dir
|
||||||
|
/ f"formal-depth-32-block-seed-{seed}.json"
|
||||||
|
)
|
||||||
|
run = json.loads(path.read_text())
|
||||||
|
source_hashes[str(seed)] = run["canonical_sha256_without_self"]
|
||||||
|
diagnostic = next(
|
||||||
|
item for item in run["diagnostics"] if item["step"] == 8000
|
||||||
|
)
|
||||||
|
gradients = diagnostic["activation_grad_rms_by_block"]
|
||||||
|
gradient_mean = statistics.fmean(gradients)
|
||||||
|
for layer in range(32):
|
||||||
|
rows.append(
|
||||||
|
{
|
||||||
|
"seed": seed,
|
||||||
|
"layer": layer + 1,
|
||||||
|
"group": layer // 4 + 1,
|
||||||
|
"offset": layer % 4 + 1,
|
||||||
|
"normalized_gradient": gradients[layer] / gradient_mean,
|
||||||
|
"attention": mixer_metrics(
|
||||||
|
diagnostic["depth_weights"][2 * layer]
|
||||||
|
),
|
||||||
|
"mlp": mixer_metrics(
|
||||||
|
diagnostic["depth_weights"][2 * layer + 1]
|
||||||
|
),
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
layer_means = []
|
||||||
|
for layer in range(1, 33):
|
||||||
|
selected = [row for row in rows if row["layer"] == layer]
|
||||||
|
layer_means.append(
|
||||||
|
{
|
||||||
|
"layer": layer,
|
||||||
|
"group": selected[0]["group"],
|
||||||
|
"offset": selected[0]["offset"],
|
||||||
|
"normalized_gradient": average(
|
||||||
|
selected, lambda row: row["normalized_gradient"]
|
||||||
|
),
|
||||||
|
"attention_latest": average(
|
||||||
|
selected, lambda row: row["attention"]["latest"]
|
||||||
|
),
|
||||||
|
"mlp_latest": average(
|
||||||
|
selected, lambda row: row["mlp"]["latest"]
|
||||||
|
),
|
||||||
|
"attention_normalized_entropy": average(
|
||||||
|
selected,
|
||||||
|
lambda row: row["attention"]["normalized_entropy"],
|
||||||
|
),
|
||||||
|
"mlp_normalized_entropy": average(
|
||||||
|
selected, lambda row: row["mlp"]["normalized_entropy"]
|
||||||
|
),
|
||||||
|
"per_seed_normalized_gradient": [
|
||||||
|
row["normalized_gradient"] for row in selected
|
||||||
|
],
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
correlations = {}
|
||||||
|
for label, selected in (
|
||||||
|
("all_layers", rows),
|
||||||
|
("layers_19_28", [
|
||||||
|
row for row in rows if 19 <= row["layer"] <= 28
|
||||||
|
]),
|
||||||
|
):
|
||||||
|
gradients = [row["normalized_gradient"] for row in selected]
|
||||||
|
correlations[label] = {
|
||||||
|
"points": len(selected),
|
||||||
|
"attention_latest": pearson(
|
||||||
|
gradients,
|
||||||
|
[row["attention"]["latest"] for row in selected],
|
||||||
|
),
|
||||||
|
"mlp_latest": pearson(
|
||||||
|
gradients, [row["mlp"]["latest"] for row in selected]
|
||||||
|
),
|
||||||
|
"attention_normalized_entropy": pearson(
|
||||||
|
gradients,
|
||||||
|
[
|
||||||
|
row["attention"]["normalized_entropy"]
|
||||||
|
for row in selected
|
||||||
|
],
|
||||||
|
),
|
||||||
|
"mlp_normalized_entropy": pearson(
|
||||||
|
gradients,
|
||||||
|
[row["mlp"]["normalized_entropy"] for row in selected],
|
||||||
|
),
|
||||||
|
"attention_maximum": pearson(
|
||||||
|
gradients,
|
||||||
|
[row["attention"]["maximum"] for row in selected],
|
||||||
|
),
|
||||||
|
"mlp_maximum": pearson(
|
||||||
|
gradients, [row["mlp"]["maximum"] for row in selected]
|
||||||
|
),
|
||||||
|
}
|
||||||
|
|
||||||
|
print(
|
||||||
|
json.dumps(
|
||||||
|
{
|
||||||
|
"schema_version": 1,
|
||||||
|
"identity": "exploratory Round 05 scoping; not confirmatory Round 06",
|
||||||
|
"source_canonical_sha256": source_hashes,
|
||||||
|
"layers_19_28": layer_means[18:28],
|
||||||
|
"correlations": correlations,
|
||||||
|
},
|
||||||
|
ensure_ascii=False,
|
||||||
|
indent=2,
|
||||||
|
sort_keys=True,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user