research: lock AttnRes gradient runner

experiments/k3/attnres_gradient/README.md

@@ -0,0 +1,57 @@
+# Attention Residuals activation-gradient/depth study
+
+This directory implements preregistered protocol
+`llm-atlas-k3-attnres-gradient-scale-v1`:
+
+- `research/K3_ATTNRES_GRADIENT_DEFINITION_AUDIT.md`
+- `research/K3_ATTNRES_GRADIENT_SCALE_PROTOCOL.md`
+
+It is an independent reduced mechanism experiment. It is not a Kimi K3
+checkpoint forward pass and does not claim to recover the paper's unpublished
+Figure 5 telemetry definition.
+
+## 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
+```
+
+## Build the manifest
+
+```bash
+python experiments/k3/attnres_gradient/build_dataset.py \
+  --cache-dir /home/wuyang/.cache/llm-atlas/k3-attnres-gradient-scale-v1 \
+  --manifest experiments/k3/attnres_gradient/manifest.json
+```
+
+## Run a smoke cell
+
+```bash
+CUBLAS_WORKSPACE_CONFIG=:4096:8 \
+python experiments/k3/attnres_gradient/train.py \
+  --run-kind smoke \
+  --architecture block \
+  --depth 32 \
+  --seed 2026073001 \
+  --cache-dir /home/wuyang/.cache/llm-atlas/k3-attnres-gradient-scale-v1 \
+  --manifest experiments/k3/attnres_gradient/manifest.json \
+  --output /home/wuyang/.cache/llm-atlas/k3-attnres-gradient-scale-v1/smoke-a/depth-32-block.json
+```
+
+Smoke is fixed to 20 steps. Formal and replay runs are fixed to 8,000 steps;
+the runner rejects alternative budgets. The same command uses
+`--run-kind formal` or `--run-kind replay` and omits an explicit `--steps`.
+
+Formal output keys use:
+
+```text
+formal/depth-{16|32}-{baseline|block}-seed-{seed}.json
+replay/depth-32-block-seed-2026073001.json
+```
+
+Raw parquet/binary files and working runs remain in the local cache. The
+manifest, runner, complete result JSON, compact website payload, reproduction
+hashes, protocol, and audit enter the public repository.

experiments/k3/attnres_gradient/build_dataset.py

@@ -0,0 +1,206 @@
+#!/usr/bin/env python3
+"""Freeze the byte-level corpus and window schedule for K3 AttnRes Round 05."""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import os
+import urllib.request
+from pathlib import Path
+from typing import Any
+
+import pyarrow.parquet as pq
+
+
+PROTOCOL_ID = "llm-atlas-k3-attnres-gradient-scale-v1"
+DATASET_REPO = "Salesforce/wikitext"
+DATASET_REVISION = "b08601e04326c79dfdd32d625aee71d232d685c3"
+DATASET_VARIANT = "wikitext-2-raw-v1"
+SPLITS = ("train", "validation", "test")
+SEEDS = (2026073001, 2026073002, 2026073003)
+CONTEXT = 256
+FORMAL_STEPS = 8000
+FORMAL_BATCH = 32
+VALIDATION_WINDOWS = 64
+DIAGNOSTIC_WINDOWS = 16
+GATE_STEPS = (0, 1, 7999)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--cache-dir", type=Path, required=True)
+    parser.add_argument("--manifest", type=Path, required=True)
+    return parser.parse_args()
+
+
+def file_sha256(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        for block in iter(lambda: handle.read(1024 * 1024), b""):
+            digest.update(block)
+    return digest.hexdigest()
+
+
+def atomic_json(path: Path, value: dict[str, Any]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    temporary = path.with_suffix(path.suffix + ".tmp")
+    temporary.write_text(
+        json.dumps(value, ensure_ascii=False, indent=2, sort_keys=True) + "\n"
+    )
+    os.replace(temporary, path)
+
+
+def download(url: str, path: Path) -> None:
+    if path.exists():
+        return
+    path.parent.mkdir(parents=True, exist_ok=True)
+    temporary = path.with_suffix(path.suffix + ".part")
+    request = urllib.request.Request(
+        url,
+        headers={"User-Agent": "llm-atlas-k3-attnres-gradient-scale/1.0"},
+    )
+    with urllib.request.urlopen(request, timeout=120) as response:
+        with temporary.open("wb") as output:
+            while block := response.read(1024 * 1024):
+                output.write(block)
+    os.replace(temporary, path)
+
+
+def hashed_start(fields: list[str], corpus_length: int) -> int:
+    value = int.from_bytes(
+        hashlib.sha256("\0".join(fields).encode()).digest()[:8], "big"
+    )
+    return value % (corpus_length - (CONTEXT + 1))
+
+
+def fixed_window_start(label: str, index: int, corpus_length: int) -> int:
+    return hashed_start([PROTOCOL_ID, label, str(index)], corpus_length)
+
+
+def train_window_start(seed: int, step: int, row: int, corpus_length: int) -> int:
+    return hashed_start(
+        [PROTOCOL_ID, "train-window", str(seed), str(step), str(row)],
+        corpus_length,
+    )
+
+
+def concatenate_split(parquet_path: Path) -> tuple[bytes, int]:
+    table = pq.read_table(parquet_path, columns=["text"])
+    rows = table.column("text").to_pylist()
+    payload = b"".join(((row or "") + "\n").encode("utf-8") for row in rows)
+    return payload, len(rows)
+
+
+def tensor_hash(payload: bytes, starts: list[int]) -> str:
+    digest = hashlib.sha256()
+    for start in starts:
+        digest.update(payload[start : start + CONTEXT + 1])
+    return digest.hexdigest()
+
+
+def main() -> None:
+    args = parse_args()
+    args.cache_dir.mkdir(parents=True, exist_ok=True)
+
+    split_manifest: dict[str, Any] = {}
+    split_bytes: dict[str, bytes] = {}
+    for split in SPLITS:
+        relative = f"{DATASET_VARIANT}/{split}-00000-of-00001.parquet"
+        url = (
+            f"https://huggingface.co/datasets/{DATASET_REPO}/resolve/"
+            f"{DATASET_REVISION}/{relative}"
+        )
+        parquet_path = args.cache_dir / f"{split}.parquet"
+        download(url, parquet_path)
+        payload, rows = concatenate_split(parquet_path)
+        binary_path = args.cache_dir / f"{split}.bin"
+        if not binary_path.exists() or file_sha256(binary_path) != hashlib.sha256(
+            payload
+        ).hexdigest():
+            temporary = binary_path.with_suffix(".bin.tmp")
+            temporary.write_bytes(payload)
+            os.replace(temporary, binary_path)
+        split_bytes[split] = payload
+        split_manifest[split] = {
+            "source_path": relative,
+            "source_url": url,
+            "parquet_bytes": parquet_path.stat().st_size,
+            "parquet_sha256": file_sha256(parquet_path),
+            "rows": rows,
+            "concatenated_bytes": len(payload),
+            "concatenated_sha256": hashlib.sha256(payload).hexdigest(),
+            "binary_path": str(binary_path),
+            "binary_sha256": file_sha256(binary_path),
+        }
+
+    train = split_bytes["train"]
+    validation = split_bytes["validation"]
+    schedule_digest = hashlib.sha256()
+    schedule_cells = 0
+    for seed in SEEDS:
+        for step in range(1, FORMAL_STEPS + 1):
+            for row in range(FORMAL_BATCH):
+                start = train_window_start(seed, step, row, len(train))
+                schedule_digest.update(start.to_bytes(8, "big"))
+                schedule_cells += 1
+
+    validation_starts = [
+        fixed_window_start("validation-window", index, len(validation))
+        for index in range(VALIDATION_WINDOWS)
+    ]
+    diagnostic_starts = [
+        fixed_window_start("diagnostic-window", index, len(validation))
+        for index in range(DIAGNOSTIC_WINDOWS)
+    ]
+    gate_tensor_hashes: dict[str, dict[str, str]] = {}
+    for seed in SEEDS:
+        gate_tensor_hashes[str(seed)] = {}
+        for step in GATE_STEPS:
+            starts = [
+                train_window_start(seed, step, row, len(train))
+                for row in range(FORMAL_BATCH)
+            ]
+            gate_tensor_hashes[str(seed)][str(step)] = tensor_hash(train, starts)
+
+    manifest = {
+        "schema_version": 1,
+        "protocol_id": PROTOCOL_ID,
+        "status": "frozen-before-model-output",
+        "dataset": {
+            "repository": DATASET_REPO,
+            "revision": DATASET_REVISION,
+            "variant": DATASET_VARIANT,
+            "preprocessing": (
+                "parquet row order; (text or empty string) + LF; UTF-8; "
+                "no normalization; vocabulary is raw bytes 0..255"
+            ),
+            "splits": split_manifest,
+        },
+        "windows": {
+            "context": CONTEXT,
+            "target_bytes_per_window": CONTEXT,
+            "seeds": list(SEEDS),
+            "formal_steps": FORMAL_STEPS,
+            "formal_batch": FORMAL_BATCH,
+            "formal_schedule_cells": schedule_cells,
+            "formal_schedule_sha256": schedule_digest.hexdigest(),
+            "validation_starts": validation_starts,
+            "validation_tensor_sha256": tensor_hash(
+                validation, validation_starts
+            ),
+            "diagnostic_starts": diagnostic_starts,
+            "diagnostic_tensor_sha256": tensor_hash(
+                validation, diagnostic_starts
+            ),
+            "gate_steps": list(GATE_STEPS),
+            "gate_training_tensor_sha256": gate_tensor_hashes,
+        },
+    }
+    atomic_json(args.manifest, manifest)
+    print(json.dumps(manifest, ensure_ascii=False, indent=2))
+
+
+if __name__ == "__main__":
+    main()

experiments/k3/attnres_gradient/manifest.json

@@ -0,0 +1,167 @@
+{
+  "dataset": {
+    "preprocessing": "parquet row order; (text or empty string) + LF; UTF-8; no normalization; vocabulary is raw bytes 0..255",
+    "repository": "Salesforce/wikitext",
+    "revision": "b08601e04326c79dfdd32d625aee71d232d685c3",
+    "splits": {
+      "test": {
+        "binary_path": "/home/wuyang/.cache/llm-atlas/k3-attnres-gradient-scale-v1/test.bin",
+        "binary_sha256": "bfe9eb16ab9987fb88bde4ea9a30a00f2a45db01dfc14bad78d05325789c4f12",
+        "concatenated_bytes": 1292014,
+        "concatenated_sha256": "bfe9eb16ab9987fb88bde4ea9a30a00f2a45db01dfc14bad78d05325789c4f12",
+        "parquet_bytes": 732610,
+        "parquet_sha256": "5f1bea067869d04849c0f975a2b29c4ff47d867f484f5010ea5e861eab246d91",
+        "rows": 4358,
+        "source_path": "wikitext-2-raw-v1/test-00000-of-00001.parquet",
+        "source_url": "https://huggingface.co/datasets/Salesforce/wikitext/resolve/b08601e04326c79dfdd32d625aee71d232d685c3/wikitext-2-raw-v1/test-00000-of-00001.parquet"
+      },
+      "train": {
+        "binary_path": "/home/wuyang/.cache/llm-atlas/k3-attnres-gradient-scale-v1/train.bin",
+        "binary_sha256": "0ca7d3e74dbe44564ea5942b85232f1bbcb525c9cd481cd5d28a87ee90e7e9b4",
+        "concatenated_bytes": 10951563,
+        "concatenated_sha256": "0ca7d3e74dbe44564ea5942b85232f1bbcb525c9cd481cd5d28a87ee90e7e9b4",
+        "parquet_bytes": 6357543,
+        "parquet_sha256": "e83889baabc497075506f91975be5fac0d45c5290b6b20582c8cd1e853d0c9f7",
+        "rows": 36718,
+        "source_path": "wikitext-2-raw-v1/train-00000-of-00001.parquet",
+        "source_url": "https://huggingface.co/datasets/Salesforce/wikitext/resolve/b08601e04326c79dfdd32d625aee71d232d685c3/wikitext-2-raw-v1/train-00000-of-00001.parquet"
+      },
+      "validation": {
+        "binary_path": "/home/wuyang/.cache/llm-atlas/k3-attnres-gradient-scale-v1/validation.bin",
+        "binary_sha256": "a42356f6a8ff1d25daf25ec9db49e10a537c265581b61c74604bb63231dee719",
+        "concatenated_bytes": 1148008,
+        "concatenated_sha256": "a42356f6a8ff1d25daf25ec9db49e10a537c265581b61c74604bb63231dee719",
+        "parquet_bytes": 657209,
+        "parquet_sha256": "204929b7ff9d6184953f867dedb860e40aa69c078fc1e54b3baaa8fb28511c4c",
+        "rows": 3760,
+        "source_path": "wikitext-2-raw-v1/validation-00000-of-00001.parquet",
+        "source_url": "https://huggingface.co/datasets/Salesforce/wikitext/resolve/b08601e04326c79dfdd32d625aee71d232d685c3/wikitext-2-raw-v1/validation-00000-of-00001.parquet"
+      }
+    },
+    "variant": "wikitext-2-raw-v1"
+  },
+  "protocol_id": "llm-atlas-k3-attnres-gradient-scale-v1",
+  "schema_version": 1,
+  "status": "frozen-before-model-output",
+  "windows": {
+    "context": 256,
+    "diagnostic_starts": [
+      399861,
+      210983,
+      449025,
+      1098024,
+      323754,
+      152932,
+      1091551,
+      1078021,
+      415985,
+      612288,
+      910624,
+      530272,
+      827285,
+      765798,
+      1086876,
+      1035447
+    ],
+    "diagnostic_tensor_sha256": "21117e31db302b10d67b63f035665dc8f220b879d216ccd12b7d2ba86e7b1716",
+    "formal_batch": 32,
+    "formal_schedule_cells": 768000,
+    "formal_schedule_sha256": "5041e09b167f229248d2462324e8c254b8f5938975f135dcd8192b00a54a4f4e",
+    "formal_steps": 8000,
+    "gate_steps": [
+      0,
+      1,
+      7999
+    ],
+    "gate_training_tensor_sha256": {
+      "2026073001": {
+        "0": "52fdd6885cc2bef8e29cedec8c293e1ea71f63fe3cd0b83f627fe640e19a95c5",
+        "1": "9d0a960595a3f57cd18834d880bde56fa8dcbb0b1cbed26bcc9bbf773a67949c",
+        "7999": "8f4f04a889d1f8c9eb75e917dc7cd6ddef1196bdea94e87466270db7be1289ee"
+      },
+      "2026073002": {
+        "0": "156813ff7ab93736c8dba340711a9633b6c952d0f54df30f277e254705cc8b82",
+        "1": "9ee5a434bdf417d658a1485135d48749e98c86f03239eef86dbd81a9d1309da0",
+        "7999": "bca78ffefa000dc3693a790d65933251646c70facc64b42007d500c19bb90977"
+      },
+      "2026073003": {
+        "0": "8ba13ad55eb54501ec443430446793ef41dea21fd067b2a5b0c84eda202ed0ba",
+        "1": "326a12f07637fe70fd39fa2758b389e91dacc7f2f2ede708b80e09c9d4d1c39d",
+        "7999": "7d6b10c3febfb697444909f695d3a45297789022714f8b71dadfec4d6236aea6"
+      }
+    },
+    "seeds": [
+      2026073001,
+      2026073002,
+      2026073003
+    ],
+    "target_bytes_per_window": 256,
+    "validation_starts": [
+      19785,
+      600807,
+      1029500,
+      319878,
+      652204,
+      1072662,
+      679625,
+      1027264,
+      500250,
+      944134,
+      187313,
+      834295,
+      968556,
+      550645,
+      239009,
+      452519,
+      356354,
+      134015,
+      64555,
+      397632,
+      203140,
+      346032,
+      314812,
+      10817,
+      1141274,
+      807645,
+      417975,
+      870687,
+      377265,
+      635426,
+      597238,
+      805324,
+      12300,
+      264343,
+      84743,
+      596894,
+      188690,
+      992517,
+      854512,
+      427504,
+      94167,
+      296670,
+      760313,
+      912279,
+      1054297,
+      81970,
+      419690,
+      971472,
+      1041491,
+      669963,
+      735537,
+      434513,
+      169153,
+      6229,
+      136413,
+      1098303,
+      400950,
+      457810,
+      659776,
+      911665,
+      909832,
+      532969,
+      555820,
+      1019138
+    ],
+    "validation_tensor_sha256": "f459316f13078a163b47c133511bb7181e05170ab89516e196490113893ce338"
+  }
+}

experiments/k3/attnres_gradient/train.py

@@ -0,0 +1,795 @@
+#!/usr/bin/env python3
+"""Run one preregistered AttnRes activation-gradient/depth experiment cell."""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import importlib.util
+import json
+import math
+import os
+import platform
+import statistics
+import sys
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Iterable
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+PROTOCOL_ID = "llm-atlas-k3-attnres-gradient-scale-v1"
+ARCHITECTURES = ("baseline", "block")
+DEPTHS = (16, 32)
+EXPECTED_SEEDS = (2026073001, 2026073002, 2026073003)
+DIAGNOSTIC_STEPS = (0, 100, 500, 2000, 4000, 8000)
+FORMAL_STEPS = 8000
+SMOKE_STEPS = 20
+CONTEXT = 256
+VOCABULARY = 256
+BLOCK_GROUPS = 8
+PEAK_LR = 3e-4
+MIN_LR = 3e-5
+WARMUP_STEPS = 400
+WEIGHT_DECAY = 0.1
+BETAS = (0.9, 0.95)
+ADAM_EPS = 1e-8
+GRAD_CLIP = 1.0
+
+
+def load_round04_module() -> Any:
+    path = Path(__file__).resolve().parents[1] / "attnres" / "train.py"
+    spec = importlib.util.spec_from_file_location("k3_attnres_round04_train", path)
+    if spec is None or spec.loader is None:
+        raise RuntimeError(f"cannot import Round 04 runner from {path}")
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+round04 = load_round04_module()
+
+
+def configure_round04_globals(depth: int) -> None:
+    round04.PROTOCOL_ID = PROTOCOL_ID
+    round04.LAYERS = depth
+    round04.SUBLAYERS = depth * 2
+    round04.BLOCKS = BLOCK_GROUPS
+    round04.SUBLAYERS_PER_BLOCK = (depth * 2) // BLOCK_GROUPS
+    round04.WARMUP_STEPS = WARMUP_STEPS
+    round04.EVAL_STEPS = DIAGNOSTIC_STEPS
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--architecture", choices=ARCHITECTURES, required=True)
+    parser.add_argument("--depth", type=int, choices=DEPTHS, required=True)
+    parser.add_argument("--seed", type=int, required=True)
+    parser.add_argument("--steps", type=int)
+    parser.add_argument("--batch-size", type=int, default=32)
+    parser.add_argument("--cache-dir", type=Path, required=True)
+    parser.add_argument("--manifest", type=Path, required=True)
+    parser.add_argument("--output", type=Path, required=True)
+    parser.add_argument("--validation-windows", type=int, default=64)
+    parser.add_argument("--diagnostic-windows", type=int, default=16)
+    parser.add_argument("--eval-batch-size", type=int, default=8)
+    parser.add_argument("--timing-warmup", type=int, default=20)
+    parser.add_argument(
+        "--run-kind", choices=("smoke", "formal", "replay"), default="formal"
+    )
+    args = parser.parse_args()
+    expected_steps = SMOKE_STEPS if args.run_kind == "smoke" else FORMAL_STEPS
+    if args.steps is None:
+        args.steps = expected_steps
+    if args.steps != expected_steps:
+        raise ValueError(
+            f"{args.run_kind} must run exactly {expected_steps} steps, got {args.steps}"
+        )
+    if args.batch_size != 32:
+        raise ValueError("the frozen protocol requires batch size 32")
+    if args.validation_windows != 64 or args.diagnostic_windows != 16:
+        raise ValueError("the frozen protocol requires 64 validation / 16 diagnostic windows")
+    return args
+
+
+def configure_determinism(seed: int) -> None:
+    if os.environ.get("CUBLAS_WORKSPACE_CONFIG") != ":4096:8":
+        raise RuntimeError("CUBLAS_WORKSPACE_CONFIG must be :4096:8 before Python starts")
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    torch.use_deterministic_algorithms(True)
+    torch.backends.cudnn.benchmark = False
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cuda.matmul.allow_tf32 = False
+    torch.backends.cudnn.allow_tf32 = False
+    torch.set_float32_matmul_precision("highest")
+
+
+def canonical_sha256(value: Any) -> str:
+    payload = json.dumps(
+        value, ensure_ascii=False, sort_keys=True, separators=(",", ":")
+    ).encode()
+    return hashlib.sha256(payload).hexdigest()
+
+
+def file_sha256(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        for block in iter(lambda: handle.read(1024 * 1024), b""):
+            digest.update(block)
+    return digest.hexdigest()
+
+
+def tensor_bytes(tensor: torch.Tensor) -> bytes:
+    value = tensor.detach().cpu().contiguous()
+    return (
+        f"{value.dtype}|{tuple(value.shape)}|".encode()
+        + value.reshape(-1).view(torch.uint8).numpy().tobytes()
+    )
+
+
+def named_state_hash(
+    model: nn.Module, *, include_mixers: bool | None
+) -> str:
+    digest = hashlib.sha256()
+    for name, tensor in sorted(model.state_dict().items()):
+        is_mixer = name.startswith("mixers.") or name.startswith("output_mixer.")
+        if include_mixers is not None and is_mixer != include_mixers:
+            continue
+        digest.update(name.encode())
+        digest.update(b"\0")
+        digest.update(tensor_bytes(tensor))
+    return digest.hexdigest()
+
+
+def state_structure_hash(
+    model: nn.Module, *, include_mixers: bool | None
+) -> tuple[str, int, int]:
+    digest = hashlib.sha256()
+    tensor_count = 0
+    element_count = 0
+    for name, tensor in sorted(model.state_dict().items()):
+        is_mixer = name.startswith("mixers.") or name.startswith("output_mixer.")
+        if include_mixers is not None and is_mixer != include_mixers:
+            continue
+        digest.update(
+            f"{name}|{tuple(tensor.shape)}|{tensor.dtype}|{tensor.numel()}\n".encode()
+        )
+        tensor_count += 1
+        element_count += tensor.numel()
+    return digest.hexdigest(), tensor_count, element_count
+
+
+def recursive_state_hash(value: Any) -> str:
+    digest = hashlib.sha256()
+
+    def visit(path: str, item: Any) -> None:
+        if torch.is_tensor(item):
+            digest.update(f"{path}|tensor|".encode())
+            digest.update(tensor_bytes(item))
+        elif isinstance(item, dict):
+            digest.update(f"{path}|dict|{len(item)}\n".encode())
+            for key in sorted(item, key=lambda candidate: str(candidate)):
+                visit(f"{path}/{key}", item[key])
+        elif isinstance(item, (list, tuple)):
+            digest.update(f"{path}|sequence|{len(item)}\n".encode())
+            for index, child in enumerate(item):
+                visit(f"{path}/{index}", child)
+        else:
+            digest.update(f"{path}|scalar|{repr(item)}\n".encode())
+
+    visit("root", value)
+    return digest.hexdigest()
+
+
+@dataclass
+class ActivationTrace:
+    block_outputs: list[torch.Tensor]
+    layer_input_rms: list[float]
+    branch_output_rms: list[float]
+    stream_state_rms: list[float]
+    depth_weights: list[dict[str, Any]]
+    output_weights: dict[str, Any] | None = None
+
+
+def rms(value: torch.Tensor) -> float:
+    return value.float().square().mean().sqrt().detach().cpu().item()
+
+
+class GradientLanguageModel(round04.ReducedLanguageModel):
+    """Round 04 trunk with aligned post-MLP activation capture."""
+
+    def forward(
+        self, input_ids: torch.Tensor, capture: bool = False
+    ) -> tuple[torch.Tensor, ActivationTrace | None]:
+        embedded = self.embed(input_ids)
+        trace = ActivationTrace([], [], [], [], []) if capture else None
+
+        if self.architecture == "baseline":
+            hidden = embedded
+            for block in self.blocks:
+                attention_input = hidden
+                attention_output = block.attention(block.attention_norm(attention_input))
+                hidden = hidden + attention_output
+                if trace is not None:
+                    trace.layer_input_rms.append(rms(attention_input))
+                    trace.branch_output_rms.append(rms(attention_output))
+                    trace.stream_state_rms.append(rms(hidden))
+                mlp_input = hidden
+                mlp_output = block.mlp(block.mlp_norm(mlp_input))
+                hidden = hidden + mlp_output
+                if trace is not None:
+                    hidden.retain_grad()
+                    trace.block_outputs.append(hidden)
+                    trace.layer_input_rms.append(rms(mlp_input))
+                    trace.branch_output_rms.append(rms(mlp_output))
+                    trace.stream_state_rms.append(rms(hidden))
+        else:
+            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.mixers[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(rms(branch_input))
+                        trace.branch_output_rms.append(rms(branch_output))
+                        trace.stream_state_rms.append(rms(partial))
+                        trace.depth_weights.append(weights or {})
+                        if branch_index == 1:
+                            partial.retain_grad()
+                            trace.block_outputs.append(partial)
+                    if mixer_index % round04.SUBLAYERS_PER_BLOCK == 0:
+                        completed.append(partial)
+                        partial = None
+            if partial is not None or len(completed) != BLOCK_GROUPS + 1:
+                raise RuntimeError("Block AttnRes aggregation contract failed")
+            if self.output_mixer is None:
+                raise RuntimeError("Block AttnRes output mixer missing")
+            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 cross_entropy(logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
+    return F.cross_entropy(
+        logits.float().reshape(-1, VOCABULARY), targets.reshape(-1)
+    )
+
+
+def learning_rate(step: int, total_steps: int) -> float:
+    if step <= WARMUP_STEPS:
+        return PEAK_LR * step / WARMUP_STEPS
+    progress = (step - WARMUP_STEPS) / max(1, total_steps - WARMUP_STEPS)
+    cosine = 0.5 * (1 + math.cos(math.pi * progress))
+    return MIN_LR + (PEAK_LR - MIN_LR) * cosine
+
+
+@torch.no_grad()
+def evaluate(
+    model: GradientLanguageModel,
+    corpus: Any,
+    window_count: int,
+    eval_batch_size: int,
+) -> dict[str, float]:
+    model.eval()
+    loss_sum = 0.0
+    target_count = 0
+    for begin in range(0, window_count, eval_batch_size):
+        end = min(begin + eval_batch_size, window_count)
+        inputs, targets = corpus.fixed_batch(
+            corpus.validation_starts, begin, end
+        )
+        with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+            logits, _ = model(inputs)
+        loss = F.cross_entropy(
+            logits.float().reshape(-1, VOCABULARY),
+            targets.reshape(-1),
+            reduction="sum",
+        )
+        loss_sum += loss.detach().cpu().item()
+        target_count += targets.numel()
+    nats = loss_sum / target_count
+    return {"cross_entropy_nats": nats, "bits_per_byte": nats / math.log(2)}
+
+
+def mean(values: Iterable[float]) -> float:
+    return statistics.fmean(values)
+
+
+def depth_statistics(values: list[float]) -> dict[str, Any]:
+    average = mean(values)
+    variance = mean((value - average) ** 2 for value in values)
+    quartile = len(values) // 4
+    first = mean(values[:quartile])
+    last = mean(values[-quartile:])
+    ratio = first / last
+    return {
+        "mean": average,
+        "population_cv": math.sqrt(variance) / average,
+        "normalized": [value / average for value in values],
+        "first_quartile_mean": first,
+        "last_quartile_mean": last,
+        "first_to_last_ratio": ratio,
+        "imbalance_abs_log_ratio": abs(math.log(ratio)),
+    }
+
+
+def core_parameter_gradient_rms(model: GradientLanguageModel) -> list[float]:
+    values = []
+    for block in model.blocks:
+        sum_square = 0.0
+        count = 0
+        for parameter in block.parameters():
+            if parameter.grad is None:
+                raise RuntimeError("missing core parameter gradient")
+            gradient = parameter.grad.detach().float()
+            if not torch.isfinite(gradient).all():
+                raise RuntimeError("non-finite core parameter gradient")
+            sum_square += gradient.square().sum().detach().cpu().item()
+            count += gradient.numel()
+        values.append(math.sqrt(sum_square / count))
+    return values
+
+
+def activation_storage_unique(outputs: list[torch.Tensor]) -> bool:
+    pointers = [output.untyped_storage().data_ptr() for output in outputs]
+    return len(pointers) == len(set(pointers))
+
+
+def diagnostic(
+    model: GradientLanguageModel,
+    corpus: Any,
+    window_count: int,
+    *,
+    loss_scale: float = 1.0,
+) -> dict[str, Any]:
+    model.eval()
+    model.zero_grad(set_to_none=True)
+    inputs, targets = corpus.fixed_batch(
+        corpus.diagnostic_starts, 0, window_count
+    )
+    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+        logits, trace = model(inputs, capture=True)
+        unscaled_loss = cross_entropy(logits, targets)
+        loss = unscaled_loss * loss_scale
+    if trace is None or len(trace.block_outputs) != len(model.blocks):
+        raise RuntimeError("aligned activation capture count mismatch")
+    expected_shape = (window_count, CONTEXT, round04.D_MODEL)
+    if any(tuple(output.shape) != expected_shape for output in trace.block_outputs):
+        raise RuntimeError("aligned activation capture shape mismatch")
+    if any(output.dtype != torch.float32 for output in trace.block_outputs):
+        raise RuntimeError("aligned activation capture must use FP32 residual state")
+    if not activation_storage_unique(trace.block_outputs):
+        raise RuntimeError("captured block outputs alias storage")
+    loss.backward()
+
+    activation_grad_rms = []
+    activation_output_rms = []
+    activation_dtypes = []
+    for output in trace.block_outputs:
+        if output.grad is None:
+            raise RuntimeError("captured activation gradient is None")
+        gradient = output.grad.detach().float()
+        if not torch.isfinite(gradient).all():
+            raise RuntimeError("captured activation gradient is non-finite")
+        activation_grad_rms.append(
+            gradient.square().mean().sqrt().detach().cpu().item()
+        )
+        activation_output_rms.append(rms(output))
+        activation_dtypes.append(str(output.dtype))
+    parameter_grad_rms = core_parameter_gradient_rms(model)
+
+    return {
+        "loss_nats": unscaled_loss.detach().cpu().item(),
+        "bits_per_byte": unscaled_loss.detach().cpu().item() / math.log(2),
+        "loss_scale": loss_scale,
+        "capture": {
+            "count": len(trace.block_outputs),
+            "shape": list(expected_shape),
+            "dtypes": activation_dtypes,
+            "all_gradients_finite": True,
+            "all_gradients_present": True,
+            "storage_unique": True,
+            "position": (
+                "post-MLP Transformer-block output; Block AttnRes is captured "
+                "before aggregation-partial reset"
+            ),
+        },
+        "activation_grad_rms_by_block": activation_grad_rms,
+        "activation_grad_statistics": depth_statistics(activation_grad_rms),
+        "activation_output_rms_by_block": activation_output_rms,
+        "activation_output_statistics": depth_statistics(activation_output_rms),
+        "core_parameter_grad_rms_by_block": parameter_grad_rms,
+        "core_parameter_grad_statistics": depth_statistics(parameter_grad_rms),
+        "layer_input_rms_by_sublayer": trace.layer_input_rms,
+        "branch_output_rms_by_sublayer": trace.branch_output_rms,
+        "stream_state_rms_by_sublayer": trace.stream_state_rms,
+        "depth_weights": trace.depth_weights,
+        "output_weights": trace.output_weights,
+    }
+
+
+def loss_scale_gate(
+    model: GradientLanguageModel, corpus: Any, window_count: int
+) -> tuple[dict[str, Any], dict[str, Any]]:
+    base = diagnostic(model, corpus, window_count, loss_scale=1.0)
+    doubled = diagnostic(model, corpus, window_count, loss_scale=2.0)
+    base_values = base["activation_grad_rms_by_block"]
+    doubled_values = doubled["activation_grad_rms_by_block"]
+    ratios = [
+        doubled_value / base_value
+        for base_value, doubled_value in zip(base_values, doubled_values)
+    ]
+    base_stats = base["activation_grad_statistics"]
+    doubled_stats = doubled["activation_grad_statistics"]
+    cv_delta = abs(
+        doubled_stats["population_cv"] - base_stats["population_cv"]
+    )
+    ratio_delta = abs(
+        doubled_stats["first_to_last_ratio"]
+        - base_stats["first_to_last_ratio"]
+    )
+    normalized_max_delta = max(
+        abs(left - right)
+        for left, right in zip(
+            base_stats["normalized"], doubled_stats["normalized"]
+        )
+    )
+    passed = (
+        all(abs(ratio - 2.0) <= 1e-5 for ratio in ratios)
+        and cv_delta <= 1e-6
+        and ratio_delta <= 1e-6
+        and normalized_max_delta <= 1e-6
+    )
+    gate = {
+        "passed": passed,
+        "per_block_scale_ratios": ratios,
+        "max_abs_scale_ratio_error": max(abs(ratio - 2.0) for ratio in ratios),
+        "population_cv_abs_delta": cv_delta,
+        "first_to_last_ratio_abs_delta": ratio_delta,
+        "normalized_spectrum_max_abs_delta": normalized_max_delta,
+        "thresholds": {
+            "scale_ratio_abs": 1e-5,
+            "shape_abs": 1e-6,
+        },
+    }
+    if not passed:
+        raise RuntimeError(f"loss-scale diagnostic gate failed: {gate}")
+    return base, gate
+
+
+def percentile(values: list[float], quantile: float) -> float:
+    return float(np.quantile(np.asarray(values, dtype=np.float64), quantile))
+
+
+def parameter_inventory(model: GradientLanguageModel) -> dict[str, int]:
+    total = sum(parameter.numel() for parameter in model.parameters())
+    mixer = sum(
+        parameter.numel()
+        for name, parameter in model.named_parameters()
+        if name.startswith("mixers.") or name.startswith("output_mixer.")
+    )
+    return {
+        "total": total,
+        "core": total - mixer,
+        "mixer": mixer,
+        "embedding": (
+            model.token_embedding.weight.numel()
+            + model.position_embedding.weight.numel()
+        ),
+    }
+
+
+def model_input_gate_hashes(
+    corpus: Any, manifest: dict[str, Any], seed: int, batch_size: int
+) -> dict[str, str]:
+    values: dict[str, str] = {}
+    for step in manifest["windows"]["gate_steps"]:
+        raw_digest = hashlib.sha256()
+        for row in range(batch_size):
+            start = round04.window_start(seed, step, row, len(corpus.train))
+            raw_digest.update(
+                np.asarray(
+                    corpus.train[start : start + CONTEXT + 1], dtype=np.uint8
+                ).tobytes()
+            )
+        expected_raw_hash = manifest["windows"][
+            "gate_training_tensor_sha256"
+        ][str(seed)][str(step)]
+        if raw_digest.hexdigest() != expected_raw_hash:
+            raise RuntimeError(f"manifest gate tensor mismatch at step {step}")
+        inputs, targets = corpus.training_batch(seed, step, batch_size)
+        digest = hashlib.sha256()
+        digest.update(tensor_bytes(inputs))
+        digest.update(tensor_bytes(targets))
+        values[str(step)] = digest.hexdigest()
+    return values
+
+
+def main() -> None:
+    args = parse_args()
+    if not torch.cuda.is_available():
+        raise RuntimeError("CUDA is required by the frozen protocol")
+    if args.seed not in EXPECTED_SEEDS:
+        raise ValueError(f"seed is not preregistered: {args.seed}")
+    configure_round04_globals(args.depth)
+    configure_determinism(args.seed)
+    device = torch.device("cuda")
+
+    manifest = json.loads(args.manifest.read_text())
+    if manifest["protocol_id"] != PROTOCOL_ID:
+        raise ValueError("manifest protocol mismatch")
+    if manifest["windows"]["formal_steps"] != FORMAL_STEPS:
+        raise ValueError("manifest formal-step mismatch")
+    corpus = round04.ByteCorpus(args.cache_dir, manifest, device)
+
+    model = GradientLanguageModel(args.architecture).to(device)
+    public_structure_hash, public_tensors, public_elements = state_structure_hash(
+        model, include_mixers=False
+    )
+    initial_public_hash = named_state_hash(model, include_mixers=False)
+    initial_mixer_hash = (
+        named_state_hash(model, include_mixers=True)
+        if args.architecture == "block"
+        else None
+    )
+    input_gate_hashes = model_input_gate_hashes(
+        corpus, manifest, args.seed, args.batch_size
+    )
+
+    decay_parameters: list[nn.Parameter] = []
+    no_decay_parameters: list[nn.Parameter] = []
+    for parameter in model.parameters():
+        (decay_parameters if parameter.ndim >= 2 else no_decay_parameters).append(
+            parameter
+        )
+    optimizer = torch.optim.AdamW(
+        [
+            {"params": decay_parameters, "weight_decay": WEIGHT_DECAY},
+            {"params": no_decay_parameters, "weight_decay": 0.0},
+        ],
+        lr=PEAK_LR,
+        betas=BETAS,
+        eps=ADAM_EPS,
+    )
+
+    evaluation_steps = sorted(
+        set(step for step in DIAGNOSTIC_STEPS if step <= args.steps)
+        | {0, args.steps}
+    )
+    evaluations = [
+        {
+            "step": 0,
+            **evaluate(
+                model, corpus, args.validation_windows, args.eval_batch_size
+            ),
+        }
+    ]
+    if args.run_kind == "smoke":
+        initial_diagnostic, gradient_gate = loss_scale_gate(
+            model, corpus, args.diagnostic_windows
+        )
+    else:
+        initial_diagnostic = diagnostic(
+            model, corpus, args.diagnostic_windows
+        )
+        gradient_gate = None
+    diagnostics = [{"step": 0, **initial_diagnostic}]
+    model.zero_grad(set_to_none=True)
+
+    training_history: list[dict[str, float | int]] = []
+    step_times: list[float] = []
+    model.train()
+    for step in range(1, args.steps + 1):
+        lr = learning_rate(step, args.steps)
+        for group in optimizer.param_groups:
+            group["lr"] = lr
+        inputs, targets = corpus.training_batch(args.seed, step, args.batch_size)
+        optimizer.zero_grad(set_to_none=True)
+
+        torch.cuda.synchronize()
+        started = time.perf_counter()
+        with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+            logits, _ = model(inputs)
+            loss = cross_entropy(logits, targets)
+        if not torch.isfinite(loss):
+            raise RuntimeError(f"non-finite loss at step {step}: {loss}")
+        loss.backward()
+        unclipped_norm = torch.nn.utils.clip_grad_norm_(
+            model.parameters(), GRAD_CLIP
+        )
+        optimizer.step()
+        torch.cuda.synchronize()
+        elapsed_ms = (time.perf_counter() - started) * 1000
+
+        if step == args.timing_warmup:
+            torch.cuda.reset_peak_memory_stats()
+        elif step > args.timing_warmup:
+            step_times.append(elapsed_ms)
+        if step == 1 or step % 10 == 0 or step == args.steps:
+            training_history.append(
+                {
+                    "step": step,
+                    "loss_nats": loss.detach().cpu().item(),
+                    "bits_per_byte": loss.detach().cpu().item() / math.log(2),
+                    "learning_rate": lr,
+                    "unclipped_grad_norm": float(unclipped_norm.detach().cpu()),
+                }
+            )
+
+        if step in evaluation_steps and step != 0:
+            evaluations.append(
+                {
+                    "step": step,
+                    **evaluate(
+                        model,
+                        corpus,
+                        args.validation_windows,
+                        args.eval_batch_size,
+                    ),
+                }
+            )
+            diagnostics.append(
+                {
+                    "step": step,
+                    **diagnostic(
+                        model, corpus, args.diagnostic_windows
+                    ),
+                }
+            )
+            model.zero_grad(set_to_none=True)
+            model.train()
+
+    training_peak_allocated = torch.cuda.max_memory_allocated()
+    training_peak_reserved = torch.cuda.max_memory_reserved()
+    final_public_hash = named_state_hash(model, include_mixers=False)
+    final_mixer_hash = (
+        named_state_hash(model, include_mixers=True)
+        if args.architecture == "block"
+        else None
+    )
+    final_full_hash = named_state_hash(model, include_mixers=None)
+    optimizer_hash = recursive_state_hash(optimizer.state_dict())
+    timing = {
+        "warmup_steps_excluded": args.timing_warmup,
+        "measured_steps": len(step_times),
+        "mean_ms": mean(step_times) if step_times else None,
+        "median_ms": statistics.median(step_times) if step_times else None,
+        "p95_ms": percentile(step_times, 0.95) if step_times else None,
+        "peak_allocated_bytes": training_peak_allocated,
+        "peak_reserved_bytes": training_peak_reserved,
+    }
+    result = {
+        "schema_version": 1,
+        "protocol_id": PROTOCOL_ID,
+        "run_kind": args.run_kind,
+        "architecture": args.architecture,
+        "depth": args.depth,
+        "seed": args.seed,
+        "steps": args.steps,
+        "batch_size": args.batch_size,
+        "target_bytes_seen": args.steps * args.batch_size * CONTEXT,
+        "manifest": {
+            "path": str(args.manifest),
+            "file_sha256": file_sha256(args.manifest),
+            "formal_schedule_sha256": manifest["windows"][
+                "formal_schedule_sha256"
+            ],
+            "validation_tensor_sha256": manifest["windows"][
+                "validation_tensor_sha256"
+            ],
+            "diagnostic_tensor_sha256": manifest["windows"][
+                "diagnostic_tensor_sha256"
+            ],
+            "input_gate_tensor_hashes": input_gate_hashes,
+        },
+        "model": {
+            "layers": args.depth,
+            "sublayers": args.depth * 2,
+            "attnres_aggregation_groups": BLOCK_GROUPS,
+            "sublayers_per_attnres_group": args.depth * 2 // BLOCK_GROUPS,
+            "transformer_blocks_per_attnres_group": args.depth // BLOCK_GROUPS,
+            "d_model": round04.D_MODEL,
+            "heads": round04.HEADS,
+            "d_head": round04.D_HEAD,
+            "d_ff": round04.D_FF,
+            "context": CONTEXT,
+            "vocabulary": VOCABULARY,
+            "parameters": parameter_inventory(model),
+        },
+        "optimizer": {
+            "name": "AdamW",
+            "betas": list(BETAS),
+            "epsilon": ADAM_EPS,
+            "weight_decay_ndim_ge_2": WEIGHT_DECAY,
+            "peak_lr": PEAK_LR,
+            "min_lr": MIN_LR,
+            "warmup_steps": WARMUP_STEPS,
+            "grad_clip": GRAD_CLIP,
+        },
+        "hashes": {
+            "initial_public_parameter_structure": public_structure_hash,
+            "initial_public_parameter_tensors": public_tensors,
+            "initial_public_parameter_elements": public_elements,
+            "initial_public_parameters": initial_public_hash,
+            "initial_mixer_parameters": initial_mixer_hash,
+            "final_public_parameters": final_public_hash,
+            "final_mixer_parameters": final_mixer_hash,
+            "final_model_state": final_full_hash,
+            "final_optimizer_state": optimizer_hash,
+        },
+        "evaluations": evaluations,
+        "diagnostics": diagnostics,
+        "training_history": training_history,
+        "gradient_gate": gradient_gate,
+        "timing": timing,
+        "environment": {
+            "python": platform.python_version(),
+            "torch": torch.__version__,
+            "cuda": torch.version.cuda,
+            "gpu": torch.cuda.get_device_name(0),
+            "compute_capability": list(torch.cuda.get_device_capability(0)),
+            "cublas_workspace_config": os.environ["CUBLAS_WORKSPACE_CONFIG"],
+            "deterministic_algorithms": torch.are_deterministic_algorithms_enabled(),
+            "autocast": "cuda-bfloat16-forward-fp32-cross-entropy",
+            "compile": False,
+        },
+    }
+    result["canonical_sha256_without_self"] = canonical_sha256(result)
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    temporary = args.output.with_suffix(args.output.suffix + ".tmp")
+    temporary.write_text(
+        json.dumps(result, ensure_ascii=False, indent=2, sort_keys=True) + "\n"
+    )
+    os.replace(temporary, args.output)
+    print(
+        json.dumps(
+            {
+                "output": str(args.output),
+                "run_kind": args.run_kind,
+                "architecture": args.architecture,
+                "depth": args.depth,
+                "seed": args.seed,
+                "steps": args.steps,
+                "final_bpc": evaluations[-1]["bits_per_byte"],
+                "final_activation_gradient_cv": diagnostics[-1][
+                    "activation_grad_statistics"
+                ]["population_cv"],
+                "canonical_sha256": result["canonical_sha256_without_self"],
+                "timing": timing,
+            },
+            ensure_ascii=False,
+            indent=2,
+        )
+    )
+
+
+if __name__ == "__main__":
+    main()

research/K3_ATTNRES_GRADIENT_SCALE_PROTOCOL.md

@@ -111,6 +111,11 @@ diagnostic 16 windows，分别由标签 `validation-window` / `diagnostic-window
 全部 bias-free linear 与 embedding 初始化为 `N(0,0.02)`。attention output projection 与 MLP
 down projection 的标准差为 `0.02 / sqrt(2×depth)`；普通 RMSNorm 为 1。
 
+数值精度进一步固定为：attention / MLP 线性分支受 BF16 autocast；embedding、Baseline hidden
+residual stream 与 Block aggregation partial 都以 FP32 累加。也就是说，Block 每个 BF16
+branch output 在进入 `partial` 前显式转为 FP32。这样两种结构被捕获的 `h_l` 都是 FP32，
+不会把 residual accumulator 精度差异混进梯度形状对比。
+
 ### 4.2 深度与 Block AttnRes 聚合
 
 | Transformer depth | residual sublayers | aggregation groups | sublayers/group | Transformer blocks/group |
@@ -186,8 +191,9 @@ step 0, 100, 500, 2,000, 4,000, 8,000
 | Baseline | attention residual 与 MLP residual 都完成后的 hidden state |
 | Block | MLP branch 已加入、本 aggregation partial 可能保存/reset **之前**的 partial |
 
-所有 `h_l.shape = [16,256,192]`。实现必须为 diagnostic forward 返回独立引用列表，不允许捕获
-reset 后的零张量，不允许把 8 个 aggregation sources 当成 16 / 32 个 Transformer outputs。
+所有 `h_l.shape = [16,256,192]`、dtype 为 FP32。实现必须为 diagnostic forward 返回独立
+引用列表，不允许捕获 reset 后的零张量，不允许把 8 个 aggregation sources 当成 16 / 32 个
+Transformer outputs。
 
 ### 6.2 diagnostic loss 与 gradient magnitude