#!/usr/bin/env python3 """Replay Round 05 Block training with preregistered spike-path diagnostics.""" 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 torch import torch.nn as nn import torch.nn.functional as F PROTOCOL_ID = "llm-atlas-k3-attnres-spike-path-v1" PARENT_PROTOCOL_ID = "llm-atlas-k3-attnres-gradient-scale-v1" DEPTH = 32 SEEDS = (2026073001, 2026073002, 2026073003) FORMAL_STEPS = 8000 DIAGNOSTIC_STEPS = (0, 100, 500, 2000, 4000, 8000) INTERVENTION_STEPS = (0, 8000) INTERVENTION_MODES = ( "learned", "detached_learned", "uniform_value_backward", ) POSITIONS = ( "pre_attention_input", "attention_branch_output", "post_attention_state", "pre_mlp_input", "mlp_branch_output", "post_mlp_state", ) REDUCTIONS = ( "element_rms", "token_rms_mean", "token_rms_median", "token_rms_p95", "batch_mean_rms", "token_mean_rms", "global_l2", ) SPIKE_LAYERS = (21, 22, 23, 24, 25) CONTEXT = 256 VOCABULARY = 256 DIAGNOSTIC_WINDOWS = 16 VALIDATION_WINDOWS = 64 EVAL_BATCH_SIZE = 8 TRAIN_BATCH_SIZE = 32 TIMING_WARMUP = 20 EPSILON = 1e-30 SPECTRUM_TOLERANCE = 1e-6 SCALE_TOLERANCE = 1e-5 def load_parent() -> Any: path = Path(__file__).resolve().parents[1] / "attnres_gradient" / "train.py" spec = importlib.util.spec_from_file_location( "k3_attnres_gradient_parent", path ) if spec is None or spec.loader is None: raise RuntimeError(f"cannot import parent runner from {path}") module = importlib.util.module_from_spec(spec) sys.modules[spec.name] = module spec.loader.exec_module(module) return module parent = load_parent() def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() parser.add_argument("--run-kind", choices=("smoke", "formal", "replay"), required=True) parser.add_argument("--seed", type=int, required=True) parser.add_argument("--cache-dir", type=Path, required=True) parser.add_argument("--parent-manifest", type=Path, required=True) parser.add_argument("--manifest", type=Path, required=True) parser.add_argument("--output", type=Path, required=True) args = parser.parse_args() if args.seed not in SEEDS: raise ValueError(f"seed is not preregistered: {args.seed}") if args.run_kind == "replay" and args.seed != SEEDS[0]: raise ValueError("the preregistered replay is seed 2026073001") return 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 canonical_sha256(value: Any) -> str: payload = json.dumps( value, ensure_ascii=False, sort_keys=True, separators=(",", ":") ).encode() return hashlib.sha256(payload).hexdigest() def tensors_sha256(values: Iterable[torch.Tensor]) -> str: digest = hashlib.sha256() for index, value in enumerate(values): digest.update(f"{index}\0".encode()) digest.update(parent.tensor_bytes(value)) return digest.hexdigest() def type7_quantile(sorted_values: torch.Tensor, probability: float) -> float: if sorted_values.ndim != 1 or sorted_values.numel() == 0: raise ValueError("Type-7 quantile requires a non-empty flat tensor") h = (sorted_values.numel() - 1) * probability lower = math.floor(h) fraction = h - lower if fraction == 0: value = sorted_values[lower] else: value = sorted_values[lower] + fraction * ( sorted_values[lower + 1] - sorted_values[lower] ) return value.detach().cpu().item() def mean(values: Iterable[float]) -> float: return statistics.fmean(values) def layer_statistics(values: list[float]) -> dict[str, Any]: if len(values) != DEPTH or any(not math.isfinite(value) or value <= 0 for value in values): raise RuntimeError("layer metric must contain 32 finite positive values") base = parent.depth_statistics(values) spike = [values[layer - 1] for layer in SPIKE_LAYERS] rest = [ value for layer, value in enumerate(values, start=1) if layer not in SPIKE_LAYERS ] ordered = sorted( range(1, DEPTH + 1), key=lambda layer: (-values[layer - 1], layer) ) base.update( { "spike_layers": list(SPIKE_LAYERS), "spike_mean": mean(spike), "rest_mean": mean(rest), "spike_contrast": mean(spike) / mean(rest), "peak_layer": ordered[0], "peak_normalized": values[ordered[0] - 1] / base["mean"], "top_five_layers": ordered[:5], "top_five_spike_overlap": len(set(ordered[:5]) & set(SPIKE_LAYERS)), } ) return base def gradient_reductions(gradient: torch.Tensor) -> dict[str, float]: value = gradient.detach().float() if value.shape != (DIAGNOSTIC_WINDOWS, CONTEXT, parent.round04.D_MODEL): raise RuntimeError(f"unexpected gradient shape: {tuple(value.shape)}") if not torch.isfinite(value).all(): raise RuntimeError("non-finite activation gradient") token_rms = value.square().mean(dim=-1).sqrt().reshape(-1) token_rms_sorted = torch.sort(token_rms).values result = { "element_rms": value.square().mean().sqrt().detach().cpu().item(), "token_rms_mean": token_rms.mean().detach().cpu().item(), "token_rms_median": type7_quantile(token_rms_sorted, 0.5), "token_rms_p95": type7_quantile(token_rms_sorted, 0.95), "batch_mean_rms": value.mean(dim=0).square().mean().sqrt().detach().cpu().item(), "token_mean_rms": value.mean(dim=1).square().mean().sqrt().detach().cpu().item(), "global_l2": value.square().sum().sqrt().detach().cpu().item(), } if any(not math.isfinite(metric) or metric <= 0 for metric in result.values()): raise RuntimeError("gradient reduction is non-finite or non-positive") expected_l2 = result["element_rms"] * math.sqrt(value.numel()) relative_error = abs(result["global_l2"] - expected_l2) / expected_l2 if relative_error > 1e-6: raise RuntimeError(f"global L2 algebraic control failed: {relative_error}") return result def rms(value: torch.Tensor) -> float: return value.detach().float().square().mean().sqrt().cpu().item() def weight_summary( weights: torch.Tensor, labels: list[str], *, mixer_index: int | None, layer: int | None, branch: str, group: int | None, offset: int | None, ) -> dict[str, Any]: detached = weights.detach().float() if detached.ndim != 3 or detached.shape[0] != len(labels): raise RuntimeError("mixer weight shape/label mismatch") source_summaries = [] for label, source in zip(labels, detached): flattened = torch.sort(source.reshape(-1)).values source_summaries.append( { "label": label, "mean": flattened.mean().cpu().item(), "p05": type7_quantile(flattened, 0.05), "median": type7_quantile(flattened, 0.5), "p95": type7_quantile(flattened, 0.95), } ) entropy = -( detached * torch.log(detached.clamp_min(1e-30)) ).sum(dim=0).mean().cpu().item() source_count = len(labels) means = [item["mean"] for item in source_summaries] return { "mixer_index": mixer_index, "layer": layer, "branch": branch, "group": group, "offset": offset, "sources": source_count, "source_summaries": source_summaries, "entropy_mean": entropy, "normalized_entropy": 1.0 if source_count == 1 else entropy / math.log(source_count), "max_source_mass": max(means), "latest_source_mass": means[-1], } def recompute_weights( mixer: nn.Module, sources: list[torch.Tensor] ) -> torch.Tensor: with torch.no_grad(): values = torch.stack(sources, dim=0) keys = mixer.key_norm(values) logits = torch.einsum("d,nbtd->nbt", mixer.query, keys.float()) return torch.softmax(logits, dim=0) class RoutedSourceBackward(torch.autograd.Function): @staticmethod def forward( ctx: Any, values: torch.Tensor, parent_output: torch.Tensor, backward_weights: torch.Tensor, ) -> torch.Tensor: ctx.save_for_backward(values, backward_weights) return parent_output @staticmethod def backward( ctx: Any, grad_output: torch.Tensor ) -> tuple[torch.Tensor, None, None]: values, backward_weights = ctx.saved_tensors with torch.enable_grad(): surrogate_values = values.detach().requires_grad_(True) with torch.autocast(device_type="cuda", dtype=torch.bfloat16): surrogate_output = torch.einsum( "nbt,nbtd->btd", backward_weights.detach(), surrogate_values.float(), ).to(surrogate_values.dtype) (grad_values,) = torch.autograd.grad( surrogate_output, surrogate_values, grad_output, retain_graph=False, create_graph=False, ) return grad_values, None, None @dataclass class SpikeTrace: positions: dict[str, list[torch.Tensor]] mixers: list[dict[str, Any]] class SpikeLanguageModel(parent.GradientLanguageModel): """Parent model with a diagnostic-only expanded capture graph.""" def forward( self, input_ids: torch.Tensor, capture: bool = False, mixer_backward_mode: str = "learned", ) -> tuple[torch.Tensor, SpikeTrace | None]: if not capture: if mixer_backward_mode != "learned": raise RuntimeError("training/evaluation cannot use an intervention") return super().forward(input_ids, capture=False) if mixer_backward_mode not in INTERVENTION_MODES: raise ValueError(f"unknown intervention mode: {mixer_backward_mode}") return self._diagnostic_forward(input_ids, mixer_backward_mode) @staticmethod def _capture(trace: SpikeTrace, position: str, value: torch.Tensor) -> None: if not value.requires_grad: raise RuntimeError(f"{position} does not require grad") value.retain_grad() trace.positions[position].append(value) @staticmethod def _source_labels( completed_labels: list[str], has_partial: bool ) -> list[str]: return completed_labels + (["current_group_partial"] if has_partial else []) def _mix( self, mixer: nn.Module, sources: list[torch.Tensor], labels: list[str], *, mode: str, mixer_index: int | None, layer: int | None, branch: str, group: int | None, offset: int | None, ) -> tuple[torch.Tensor, dict[str, Any]]: parent_output, _ = mixer(sources, False) weights = recompute_weights(mixer, sources) summary = weight_summary( weights, labels, mixer_index=mixer_index, layer=layer, branch=branch, group=group, offset=offset, ) if mode == "learned": return parent_output, summary values = torch.stack(sources, dim=0) if mode == "detached_learned": backward_weights = weights else: backward_weights = torch.full_like(weights, 1.0 / len(sources)) routed = RoutedSourceBackward.apply( values, parent_output, backward_weights ) return routed, summary def _diagnostic_forward( self, input_ids: torch.Tensor, mode: str ) -> tuple[torch.Tensor, SpikeTrace]: trace = SpikeTrace({position: [] for position in POSITIONS}, []) embedded = self.embed(input_ids) completed = [embedded] completed_labels = ["embedding"] partial: torch.Tensor | None = None mixer_index = 0 for layer_index, block in enumerate(self.blocks): layer = layer_index + 1 group = layer_index // 4 + 1 offset = layer_index % 4 + 1 attention_sources = completed + ( [] if partial is None else [partial] ) attention_labels = self._source_labels( completed_labels, partial is not None ) attention_input, summary = self._mix( self.mixers[mixer_index], attention_sources, attention_labels, mode=mode, mixer_index=mixer_index, layer=layer, branch="attention", group=group, offset=offset, ) trace.mixers.append(summary) mixer_index += 1 self._capture(trace, "pre_attention_input", attention_input) attention_output = block.attention( block.attention_norm(attention_input) ) self._capture(trace, "attention_branch_output", attention_output) attention_for_residual = attention_output.float() partial = ( attention_for_residual if partial is None else partial + attention_for_residual ) self._capture(trace, "post_attention_state", partial) mlp_sources = completed + [partial] mlp_labels = self._source_labels(completed_labels, True) mlp_input, summary = self._mix( self.mixers[mixer_index], mlp_sources, mlp_labels, mode=mode, mixer_index=mixer_index, layer=layer, branch="mlp", group=group, offset=offset, ) trace.mixers.append(summary) mixer_index += 1 self._capture(trace, "pre_mlp_input", mlp_input) mlp_output = block.mlp(block.mlp_norm(mlp_input)) self._capture(trace, "mlp_branch_output", mlp_output) partial = partial + mlp_output.float() self._capture(trace, "post_mlp_state", partial) if mixer_index % parent.round04.SUBLAYERS_PER_BLOCK == 0: completed.append(partial) completed_labels.append(f"completed_group_{group}") partial = None if partial is not None or len(completed) != 9 or mixer_index != 64: raise RuntimeError("Block aggregation topology mismatch") if self.output_mixer is None: raise RuntimeError("output mixer missing") hidden, summary = self._mix( self.output_mixer, completed, completed_labels, mode=mode, mixer_index=None, layer=None, branch="output", group=None, offset=None, ) trace.mixers.append(summary) normalized = self.final_norm(hidden) logits = F.linear(normalized, self.token_embedding.weight) return logits, trace def validate_trace(trace: SpikeTrace) -> dict[str, Any]: expected_shape = (DIAGNOSTIC_WINDOWS, CONTEXT, parent.round04.D_MODEL) capture = {} for position in POSITIONS: values = trace.positions[position] if len(values) != DEPTH: raise RuntimeError(f"{position} capture count mismatch") if any(tuple(value.shape) != expected_shape for value in values): raise RuntimeError(f"{position} capture shape mismatch") pointers = [value.untyped_storage().data_ptr() for value in values] if len(set(pointers)) != len(pointers): raise RuntimeError(f"{position} contains aliased storage") capture[position] = { "count": len(values), "shape": list(expected_shape), "dtypes": [str(value.dtype) for value in values], "storage_unique": True, "activation_sha256": tensors_sha256(values), "output_rms_by_layer": [rms(value) for value in values], } if len(trace.mixers) != 65: raise RuntimeError("mixer capture count mismatch") return capture def run_diagnostic( model: SpikeLanguageModel, corpus: Any, optimizer: torch.optim.Optimizer, *, mode: str, loss_scale: float = 1.0, ) -> dict[str, Any]: optimizer_before = parent.recursive_state_hash(optimizer.state_dict()) model.eval() model.zero_grad(set_to_none=True) inputs, targets = corpus.fixed_batch( corpus.diagnostic_starts, 0, DIAGNOSTIC_WINDOWS ) with torch.autocast(device_type="cuda", dtype=torch.bfloat16): logits, trace = model( inputs, capture=True, mixer_backward_mode=mode ) unscaled_loss = parent.cross_entropy(logits, targets) loss = unscaled_loss * loss_scale if trace is None: raise RuntimeError("diagnostic trace missing") capture = validate_trace(trace) forward = { "logits_sha256": tensors_sha256([logits]), "activation_sha256": { position: capture[position]["activation_sha256"] for position in POSITIONS }, "mixer_summary_sha256": canonical_sha256(trace.mixers), "loss_nats": unscaled_loss.detach().cpu().item(), } loss.backward() positions: dict[str, Any] = {} for position in POSITIONS: reductions_by_layer = {name: [] for name in REDUCTIONS} for value in trace.positions[position]: if value.grad is None: raise RuntimeError(f"{position} gradient is missing") reductions = gradient_reductions(value.grad) for name in REDUCTIONS: reductions_by_layer[name].append(reductions[name]) positions[position] = { "capture": capture[position], "reductions": { name: { "values": values, "statistics": layer_statistics(values), } for name, values in reductions_by_layer.items() }, "all_gradients_present": True, "all_gradients_finite": True, } model.zero_grad(set_to_none=True) optimizer_after = parent.recursive_state_hash(optimizer.state_dict()) if optimizer_before != optimizer_after: raise RuntimeError("diagnostic mutated optimizer state") return { "mode": mode, "loss_scale": loss_scale, "loss_nats": forward["loss_nats"], "bits_per_byte": forward["loss_nats"] / math.log(2), "forward": forward, "positions": positions, "mixers": trace.mixers, "optimizer_state_before": optimizer_before, "optimizer_state_after": optimizer_after, "optimizer_state_unchanged": True, } def forward_identity_gate(modes: dict[str, Any]) -> dict[str, Any]: learned = modes["learned"]["forward"] comparisons = {} for name in INTERVENTION_MODES[1:]: other = modes[name]["forward"] comparisons[name] = { "logits_exact": other["logits_sha256"] == learned["logits_sha256"], "loss_exact": other["loss_nats"] == learned["loss_nats"], "activations_exact": ( other["activation_sha256"] == learned["activation_sha256"] ), "mixer_summaries_exact": ( other["mixer_summary_sha256"] == learned["mixer_summary_sha256"] ), } passed = all(all(values.values()) for values in comparisons.values()) if not passed: raise RuntimeError(f"forward identity gate failed: {comparisons}") return {"passed": True, "comparisons": comparisons} def initialization_negative_control(modes: dict[str, Any]) -> dict[str, Any]: learned = modes["learned"] comparisons = {} passed = True for mode in INTERVENTION_MODES[1:]: per_position = {} for position in POSITIONS: left = learned["positions"][position]["reductions"]["element_rms"] right = modes[mode]["positions"][position]["reductions"]["element_rms"] raw_errors = [ abs(a - b) / a for a, b in zip(left["values"], right["values"]) ] normalized_errors = [ abs(a - b) for a, b in zip( left["statistics"]["normalized"], right["statistics"]["normalized"], ) ] position_passed = ( all(value > 0 and math.isfinite(value) for value in left["values"]) and max(raw_errors) <= SPECTRUM_TOLERANCE and max(normalized_errors) <= SPECTRUM_TOLERANCE ) passed = passed and position_passed per_position[position] = { "passed": position_passed, "max_raw_relative_error": max(raw_errors), "max_normalized_absolute_error": max(normalized_errors), } comparisons[mode] = per_position if not passed: raise RuntimeError( "initialization negative control failed: " + json.dumps(comparisons, sort_keys=True) ) return {"passed": True, "comparisons": comparisons} def loss_scale_gate(base: dict[str, Any], doubled: dict[str, Any]) -> dict[str, Any]: if base["forward"] != doubled["forward"]: raise RuntimeError("loss scaling unexpectedly changed the forward pass") checks = {} passed = True for position in POSITIONS: checks[position] = {} for reduction in REDUCTIONS: left = base["positions"][position]["reductions"][reduction] right = doubled["positions"][position]["reductions"][reduction] scale_errors = [ abs((b / a) - 2.0) for a, b in zip(left["values"], right["values"]) ] normalized_errors = [ abs(a - b) for a, b in zip( left["statistics"]["normalized"], right["statistics"]["normalized"], ) ] cv_error = abs( left["statistics"]["population_cv"] - right["statistics"]["population_cv"] ) contrast_error = abs( left["statistics"]["spike_contrast"] - right["statistics"]["spike_contrast"] ) item_passed = ( max(scale_errors) <= SCALE_TOLERANCE and max(normalized_errors) <= SPECTRUM_TOLERANCE and cv_error <= SPECTRUM_TOLERANCE and contrast_error <= SPECTRUM_TOLERANCE ) passed = passed and item_passed checks[position][reduction] = { "passed": item_passed, "max_scale_ratio_error": max(scale_errors), "max_normalized_absolute_error": max(normalized_errors), "cv_absolute_error": cv_error, "spike_contrast_absolute_error": contrast_error, } if not passed: raise RuntimeError("loss-scale gate failed") return {"passed": True, "checks": checks} def run_diagnostic_bundle( model: SpikeLanguageModel, corpus: Any, optimizer: torch.optim.Optimizer, step: int, ) -> dict[str, Any]: learned = run_diagnostic( model, corpus, optimizer, mode="learned" ) modes = {"learned": learned} if step in INTERVENTION_STEPS: for mode in INTERVENTION_MODES[1:]: modes[mode] = run_diagnostic( model, corpus, optimizer, mode=mode ) result = {"step": step, "modes": modes} if step in INTERVENTION_STEPS: result["forward_identity_gate"] = forward_identity_gate(modes) if step == 0: result["initialization_negative_control"] = ( initialization_negative_control(modes) ) doubled = run_diagnostic( model, corpus, optimizer, mode="learned", loss_scale=2.0 ) result["loss_scale_gate"] = loss_scale_gate(learned, doubled) model.zero_grad(set_to_none=True) return result def frozen_training_compare( result: dict[str, Any], parent_raw: dict[str, Any], ) -> dict[str, Any]: checks = { "final_model_state": ( result["hashes"]["final_model_state"] == parent_raw["hashes"]["final_model_state"] ), "final_optimizer_state": ( result["hashes"]["final_optimizer_state"] == parent_raw["hashes"]["final_optimizer_state"] ), "evaluations": result["evaluations"] == parent_raw["evaluations"], "training_history": ( result["training_history"] == parent_raw["training_history"] ), } post_mlp_exact = [] for new, old in zip(result["diagnostics"], parent_raw["diagnostics"]): new_values = new["modes"]["learned"]["positions"]["post_mlp_state"][ "reductions" ]["element_rms"]["values"] post_mlp_exact.append( new["step"] == old["step"] and new_values == old["activation_grad_rms_by_block"] ) checks["post_mlp_element_rms_all_steps"] = all(post_mlp_exact) passed = all(checks.values()) if not passed: raise RuntimeError(f"Round 05 training equivalence failed: {checks}") return {"passed": True, "checks": checks} def load_and_verify_inputs( args: argparse.Namespace, ) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any], Path]: manifest = json.loads(args.manifest.read_text()) parent_manifest = json.loads(args.parent_manifest.read_text()) repo_root = Path(__file__).resolve().parents[3] if manifest["protocol_id"] != PROTOCOL_ID: raise RuntimeError("Round 06 manifest protocol mismatch") if parent_manifest["protocol_id"] != PARENT_PROTOCOL_ID: raise RuntimeError("parent manifest protocol mismatch") frozen_fields = { "architecture": "block", "depth": DEPTH, "formal_seeds": list(SEEDS), "positions": list(POSITIONS), } for key, expected_value in frozen_fields.items(): if manifest[key] != expected_value: raise RuntimeError(f"Round 06 manifest field mismatch: {key}") if manifest["training"]["steps"] != FORMAL_STEPS: raise RuntimeError("Round 06 formal-step mismatch") if manifest["training"]["diagnostic_steps"] != list(DIAGNOSTIC_STEPS): raise RuntimeError("Round 06 diagnostic-step mismatch") if manifest["interventions"]["modes"] != list(INTERVENTION_MODES): raise RuntimeError("Round 06 intervention-mode mismatch") if manifest["interventions"]["steps"] != list(INTERVENTION_STEPS): raise RuntimeError("Round 06 intervention-step mismatch") manifest_reductions = ( manifest["reductions"]["confirmatory"] + manifest["reductions"]["exploratory"] + manifest["reductions"]["algebraic_control"] ) if manifest_reductions != list(REDUCTIONS): raise RuntimeError("Round 06 reduction list mismatch") expected_parent_hash = manifest["parent_artifacts"]["manifest_sha256"] if file_sha256(args.parent_manifest) != expected_parent_hash: raise RuntimeError("parent manifest file hash mismatch") if file_sha256(Path(parent.__file__)) != manifest["parent_artifacts"]["runner_sha256"]: raise RuntimeError("parent runner file hash mismatch") parent_protocol_path = ( repo_root / "research" / "K3_ATTNRES_GRADIENT_SCALE_PROTOCOL.md" ) if ( file_sha256(parent_protocol_path) != manifest["parent_artifacts"]["protocol_sha256"] ): raise RuntimeError("parent protocol file hash mismatch") for key in ( "formal_schedule_sha256", "validation_tensor_sha256", "diagnostic_tensor_sha256", ): if parent_manifest["windows"][key] != manifest["parent_artifacts"][key]: raise RuntimeError(f"parent data hash mismatch: {key}") parent_raw_path = ( repo_root / "experiments" / "k3" / "attnres_gradient" / "results" / "raw" / f"formal-depth-32-block-seed-{args.seed}.json" ) expected = manifest["round05_expected"][str(args.seed)] if file_sha256(parent_raw_path) != expected["raw_file_sha256"]: raise RuntimeError("Round 05 raw physical hash mismatch") parent_raw = json.loads(parent_raw_path.read_text()) if parent_raw["canonical_sha256_without_self"] != expected["canonical_sha256"]: raise RuntimeError("Round 05 raw canonical hash mismatch") if parent_raw["hashes"]["final_model_state"] != expected["final_model_state"]: raise RuntimeError("Round 05 expected model-state hash mismatch") if ( parent_raw["hashes"]["final_optimizer_state"] != expected["final_optimizer_state"] ): raise RuntimeError("Round 05 expected optimizer-state hash mismatch") return manifest, parent_manifest, parent_raw, repo_root def main() -> None: args = parse_args() if not torch.cuda.is_available(): raise RuntimeError("CUDA is required") if os.environ.get("CUBLAS_WORKSPACE_CONFIG") != ":4096:8": raise RuntimeError("CUBLAS_WORKSPACE_CONFIG must be :4096:8") manifest, parent_manifest, parent_raw, repo_root = load_and_verify_inputs(args) parent.configure_round04_globals(DEPTH) parent.configure_determinism(args.seed) corpus = parent.round04.ByteCorpus( args.cache_dir, parent_manifest, torch.device("cuda") ) model = SpikeLanguageModel("block").to(torch.device("cuda")) initial_public_hash = parent.named_state_hash( model, include_mixers=False ) initial_mixer_hash = parent.named_state_hash( model, include_mixers=True ) public_structure_hash, public_tensors, public_elements = ( parent.state_structure_hash(model, include_mixers=False) ) input_gate_hashes = parent.model_input_gate_hashes( corpus, parent_manifest, args.seed, TRAIN_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": parent.WEIGHT_DECAY}, {"params": no_decay_parameters, "weight_decay": 0.0}, ], lr=parent.PEAK_LR, betas=parent.BETAS, eps=parent.ADAM_EPS, ) evaluations = [ { "step": 0, **parent.evaluate( model, corpus, VALIDATION_WINDOWS, EVAL_BATCH_SIZE ), } ] diagnostics = [run_diagnostic_bundle(model, corpus, optimizer, 0)] print( json.dumps( { "event": "diagnostic", "step": 0, "seed": args.seed, "post_mlp_cv": diagnostics[0]["modes"]["learned"][ "positions" ]["post_mlp_state"]["reductions"]["element_rms"][ "statistics" ]["population_cv"], }, sort_keys=True, ), flush=True, ) steps = 0 if args.run_kind == "smoke" else FORMAL_STEPS training_history: list[dict[str, float | int]] = [] step_times: list[float] = [] if steps: model.train() for step in range(1, steps + 1): lr = parent.learning_rate(step, steps) for group in optimizer.param_groups: group["lr"] = lr inputs, targets = corpus.training_batch( args.seed, step, TRAIN_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, trace = model(inputs) if trace is not None: raise RuntimeError("training unexpectedly captured a trace") loss = parent.cross_entropy(logits, targets) if not torch.isfinite(loss): raise RuntimeError(f"non-finite loss at step {step}") loss.backward() unclipped_norm = torch.nn.utils.clip_grad_norm_( model.parameters(), parent.GRAD_CLIP ) optimizer.step() torch.cuda.synchronize() elapsed_ms = (time.perf_counter() - started) * 1000 if step == TIMING_WARMUP: torch.cuda.reset_peak_memory_stats() elif step > TIMING_WARMUP: step_times.append(elapsed_ms) if step == 1 or step % 10 == 0 or step == 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 DIAGNOSTIC_STEPS: evaluations.append( { "step": step, **parent.evaluate( model, corpus, VALIDATION_WINDOWS, EVAL_BATCH_SIZE, ), } ) diagnostics.append( run_diagnostic_bundle(model, corpus, optimizer, step) ) post = diagnostics[-1]["modes"]["learned"]["positions"][ "post_mlp_state" ]["reductions"]["element_rms"]["statistics"] print( json.dumps( { "event": "diagnostic", "step": step, "seed": args.seed, "validation_bpc": evaluations[-1][ "bits_per_byte" ], "post_mlp_cv": post["population_cv"], "spike_contrast": post["spike_contrast"], }, sort_keys=True, ), flush=True, ) model.train() timing = { "warmup_steps_excluded": 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": ( type7_quantile(torch.tensor(sorted(step_times)), 0.95) if step_times else None ), "peak_allocated_bytes": torch.cuda.max_memory_allocated(), "peak_reserved_bytes": torch.cuda.max_memory_reserved(), } result = { "schema_version": 1, "protocol_id": PROTOCOL_ID, "parent_protocol_id": PARENT_PROTOCOL_ID, "run_kind": args.run_kind, "architecture": "block", "depth": DEPTH, "seed": args.seed, "steps": steps, "batch_size": TRAIN_BATCH_SIZE, "target_bytes_seen": steps * TRAIN_BATCH_SIZE * CONTEXT, "manifest": { "path": str(args.manifest), "file_sha256": file_sha256(args.manifest), "parent_path": str(args.parent_manifest), "parent_file_sha256": file_sha256(args.parent_manifest), "formal_schedule_sha256": parent_manifest["windows"][ "formal_schedule_sha256" ], "validation_tensor_sha256": parent_manifest["windows"][ "validation_tensor_sha256" ], "diagnostic_tensor_sha256": parent_manifest["windows"][ "diagnostic_tensor_sha256" ], "input_gate_tensor_hashes": input_gate_hashes, }, "model": { "layers": DEPTH, "aggregation_groups": 8, "blocks_per_group": 4, "d_model": parent.round04.D_MODEL, "heads": parent.round04.HEADS, "d_ff": parent.round04.D_FF, "parameters": parent.parameter_inventory(model), }, "optimizer": { "name": "AdamW", "betas": list(parent.BETAS), "epsilon": parent.ADAM_EPS, "weight_decay_ndim_ge_2": parent.WEIGHT_DECAY, "peak_lr": parent.PEAK_LR, "min_lr": parent.MIN_LR, "warmup_steps": parent.WARMUP_STEPS, "grad_clip": parent.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": parent.named_state_hash( model, include_mixers=False ), "final_mixer_parameters": parent.named_state_hash( model, include_mixers=True ), "final_model_state": parent.named_state_hash( model, include_mixers=None ), "final_optimizer_state": parent.recursive_state_hash( optimizer.state_dict() ), }, "evaluations": evaluations, "diagnostics": diagnostics, "training_history": training_history, "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, }, "artifacts": { "runner_sha256": file_sha256(Path(__file__)), "protocol_sha256": file_sha256( repo_root / "research" / "K3_ATTNRES_SPIKE_PROTOCOL.md" ), "scoping_sha256": file_sha256( repo_root / "research" / "K3_ATTNRES_SPIKE_SCOPING.md" ), }, } if steps: result["round05_equivalence"] = frozen_training_compare( result, parent_raw ) else: result["round05_equivalence"] = None 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, "seed": args.seed, "steps": steps, "final_bpc": evaluations[-1]["bits_per_byte"], "canonical_sha256": result[ "canonical_sha256_without_self" ], "timing": timing, }, ensure_ascii=False, indent=2, ), flush=True, ) if __name__ == "__main__": main()