wuyang/llm-atlas
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

research: lock AttnRes gradient runner

Browse Source
This commit is contained in:
wuyang
2026-07-30 07:53:12 +08:00
parent 091a05f0a0
commit 1675ca54f3
5 changed files with 1233 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
experiments/k3/attnres_gradient/README.md
+57
View File
@@ -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
+206
View File
@@ -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
+167
View File
@@ -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
+795
View File
@@ -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()