From 1675ca54f3c5f3418829bdcfeae2baa42bf752d3 Mon Sep 17 00:00:00 2001 From: wuyang <5700876+banisherwy@user.noreply.gitee.com> Date: Thu, 30 Jul 2026 07:53:12 +0800 Subject: [PATCH] research: lock AttnRes gradient runner --- experiments/k3/attnres_gradient/README.md | 57 ++ .../k3/attnres_gradient/build_dataset.py | 206 +++++ experiments/k3/attnres_gradient/manifest.json | 167 ++++ experiments/k3/attnres_gradient/train.py | 795 ++++++++++++++++++ .../K3_ATTNRES_GRADIENT_SCALE_PROTOCOL.md | 10 +- 5 files changed, 1233 insertions(+), 2 deletions(-) create mode 100644 experiments/k3/attnres_gradient/README.md create mode 100644 experiments/k3/attnres_gradient/build_dataset.py create mode 100644 experiments/k3/attnres_gradient/manifest.json create mode 100644 experiments/k3/attnres_gradient/train.py diff --git a/experiments/k3/attnres_gradient/README.md b/experiments/k3/attnres_gradient/README.md new file mode 100644 index 0000000..3ad13c5 --- /dev/null +++ b/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. diff --git a/experiments/k3/attnres_gradient/build_dataset.py b/experiments/k3/attnres_gradient/build_dataset.py new file mode 100644 index 0000000..d3765d8 --- /dev/null +++ b/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() diff --git a/experiments/k3/attnres_gradient/manifest.json b/experiments/k3/attnres_gradient/manifest.json new file mode 100644 index 0000000..c57e6c7 --- /dev/null +++ b/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" + } +} diff --git a/experiments/k3/attnres_gradient/train.py b/experiments/k3/attnres_gradient/train.py new file mode 100644 index 0000000..f95e7c4 --- /dev/null +++ b/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() diff --git a/research/K3_ATTNRES_GRADIENT_SCALE_PROTOCOL.md b/research/K3_ATTNRES_GRADIENT_SCALE_PROTOCOL.md index d0100e4..c587376 100644 --- a/research/K3_ATTNRES_GRADIENT_SCALE_PROTOCOL.md +++ b/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