#!/usr/bin/env python3 """Train one frozen residual variant for the reduced Attention Residuals study.""" from __future__ import annotations import argparse import hashlib import json import math import os import platform import statistics 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-reduced-v1" ARCHITECTURES = ("baseline", "full", "block") EXPECTED_SEEDS = (2026073001, 2026073002, 2026073003) EVAL_STEPS = (0, 100, 250, 500, 1000, 1500, 2000) VOCABULARY = 256 CONTEXT = 256 LAYERS = 16 SUBLAYERS = LAYERS * 2 BLOCKS = 8 SUBLAYERS_PER_BLOCK = SUBLAYERS // BLOCKS D_MODEL = 192 HEADS = 6 D_HEAD = D_MODEL // HEADS D_FF = 768 RMS_EPS = 1e-6 PEAK_LR = 3e-4 MIN_LR = 3e-5 WARMUP_STEPS = 100 WEIGHT_DECAY = 0.1 BETAS = (0.9, 0.95) ADAM_EPS = 1e-8 GRAD_CLIP = 1.0 def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() parser.add_argument("--architecture", choices=ARCHITECTURES, required=True) parser.add_argument("--seed", type=int, required=True) parser.add_argument("--steps", type=int, default=2000) 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") return parser.parse_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_json_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() header = f"{value.dtype}|{tuple(value.shape)}|".encode() return header + value.view(torch.uint8).numpy().tobytes() def 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 window_start(seed: int, step: int, row: int, corpus_length: int) -> int: payload = "\0".join( [PROTOCOL_ID, "train-window", str(seed), str(step), str(row)] ).encode() value = int.from_bytes(hashlib.sha256(payload).digest()[:8], "big") return value % (corpus_length - (CONTEXT + 1)) class ByteCorpus: def __init__(self, cache_dir: Path, manifest: dict[str, Any], device: torch.device): self.device = device self.train = np.memmap(cache_dir / "train.bin", dtype=np.uint8, mode="r") self.validation = np.memmap( cache_dir / "validation.bin", dtype=np.uint8, mode="r" ) self.validation_starts = manifest["windows"]["validation_starts"] self.diagnostic_starts = manifest["windows"]["diagnostic_starts"] def training_batch( self, seed: int, step: int, batch_size: int ) -> tuple[torch.Tensor, torch.Tensor]: rows = np.empty((batch_size, CONTEXT + 1), dtype=np.int64) for row in range(batch_size): start = window_start(seed, step, row, len(self.train)) rows[row] = self.train[start : start + CONTEXT + 1] tensor = torch.from_numpy(rows).to(self.device, non_blocking=False) return tensor[:, :-1], tensor[:, 1:] def fixed_batch( self, starts: list[int], begin: int, end: int ) -> tuple[torch.Tensor, torch.Tensor]: chosen = starts[begin:end] rows = np.empty((len(chosen), CONTEXT + 1), dtype=np.int64) for row, start in enumerate(chosen): rows[row] = self.validation[start : start + CONTEXT + 1] tensor = torch.from_numpy(rows).to(self.device, non_blocking=False) return tensor[:, :-1], tensor[:, 1:] class RMSNorm(nn.Module): def __init__(self, dimension: int): super().__init__() self.weight = nn.Parameter(torch.ones(dimension)) def forward(self, value: torch.Tensor) -> torch.Tensor: normalized = value.float() * torch.rsqrt( value.float().square().mean(dim=-1, keepdim=True) + RMS_EPS ) return normalized.to(value.dtype) * self.weight class CausalAttention(nn.Module): def __init__(self): super().__init__() self.qkv = nn.Linear(D_MODEL, 3 * D_MODEL, bias=False) self.o_proj = nn.Linear(D_MODEL, D_MODEL, bias=False) mask = torch.triu(torch.ones(CONTEXT, CONTEXT, dtype=torch.bool), diagonal=1) self.register_buffer("causal_mask", mask, persistent=False) def forward(self, value: torch.Tensor) -> torch.Tensor: batch, sequence, _ = value.shape qkv = self.qkv(value).view(batch, sequence, 3, HEADS, D_HEAD) query, key, content = qkv.unbind(dim=2) query = query.transpose(1, 2) key = key.transpose(1, 2) content = content.transpose(1, 2) scores = torch.matmul(query, key.transpose(-1, -2)).float() / math.sqrt(D_HEAD) scores = scores.masked_fill( self.causal_mask[:sequence, :sequence], float("-inf") ) probabilities = torch.softmax(scores, dim=-1).to(query.dtype) mixed = torch.matmul(probabilities, content) mixed = mixed.transpose(1, 2).contiguous().view(batch, sequence, D_MODEL) return self.o_proj(mixed) class SwiGLU(nn.Module): def __init__(self): super().__init__() self.gate = nn.Linear(D_MODEL, D_FF, bias=False) self.up = nn.Linear(D_MODEL, D_FF, bias=False) self.down = nn.Linear(D_FF, D_MODEL, bias=False) def forward(self, value: torch.Tensor) -> torch.Tensor: return self.down(F.silu(self.gate(value)) * self.up(value)) class TransformerBlock(nn.Module): def __init__(self): super().__init__() self.attention_norm = RMSNorm(D_MODEL) self.attention = CausalAttention() self.mlp_norm = RMSNorm(D_MODEL) self.mlp = SwiGLU() class DepthMixer(nn.Module): def __init__(self): super().__init__() self.query = nn.Parameter(torch.zeros(D_MODEL)) self.key_norm = RMSNorm(D_MODEL) def forward( self, sources: list[torch.Tensor], capture: bool = False ) -> tuple[torch.Tensor, dict[str, Any] | None]: values = torch.stack(sources, dim=0) keys = self.key_norm(values) logits = torch.einsum("d,nbtd->nbt", self.query, keys.float()) weights = torch.softmax(logits, dim=0) output = torch.einsum("nbt,nbtd->btd", weights, values.float()).to( values.dtype ) if not capture: return output, None entropy = -(weights * torch.log(weights.clamp_min(1e-30))).sum(dim=0) return output, { "mean_weights": weights.mean(dim=(1, 2)).detach().cpu().tolist(), "entropy_mean": entropy.mean().detach().cpu().item(), "sources": len(sources), } @dataclass class TraceAccumulator: 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 ReducedLanguageModel(nn.Module): def __init__(self, architecture: str): super().__init__() self.architecture = architecture self.token_embedding = nn.Embedding(VOCABULARY, D_MODEL) self.position_embedding = nn.Embedding(CONTEXT, D_MODEL) self.blocks = nn.ModuleList([TransformerBlock() for _ in range(LAYERS)]) self.final_norm = RMSNorm(D_MODEL) if architecture == "baseline": self.mixers = nn.ModuleList() self.output_mixer = None else: self.mixers = nn.ModuleList([DepthMixer() for _ in range(SUBLAYERS)]) self.output_mixer = DepthMixer() self.reset_parameters() def reset_parameters(self) -> None: for module in self.modules(): if isinstance(module, nn.Embedding): nn.init.normal_(module.weight, mean=0.0, std=0.02) elif isinstance(module, nn.Linear): nn.init.normal_(module.weight, mean=0.0, std=0.02) elif isinstance(module, RMSNorm): nn.init.ones_(module.weight) scaled = 0.02 / math.sqrt(2 * LAYERS) for block in self.blocks: nn.init.normal_(block.attention.o_proj.weight, mean=0.0, std=scaled) nn.init.normal_(block.mlp.down.weight, mean=0.0, std=scaled) for mixer in self.mixers: nn.init.zeros_(mixer.query) nn.init.ones_(mixer.key_norm.weight) if self.output_mixer is not None: nn.init.zeros_(self.output_mixer.query) nn.init.ones_(self.output_mixer.key_norm.weight) def embed(self, input_ids: torch.Tensor) -> torch.Tensor: positions = torch.arange(input_ids.shape[1], device=input_ids.device) return self.token_embedding(input_ids) + self.position_embedding(positions)[None] def forward( self, input_ids: torch.Tensor, capture: bool = False ) -> tuple[torch.Tensor, TraceAccumulator | None]: embedded = self.embed(input_ids) trace = ( TraceAccumulator([], [], [], []) 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: trace.layer_input_rms.append(rms(mlp_input)) trace.branch_output_rms.append(rms(mlp_output)) trace.stream_state_rms.append(rms(hidden)) elif self.architecture == "full": sources = [embedded] mixer_index = 0 for block in self.blocks: attention_input, weights = self.mixers[mixer_index](sources, capture) mixer_index += 1 attention_output = block.attention(block.attention_norm(attention_input)) sources.append(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(torch.stack(sources, dim=0)) ) trace.depth_weights.append(weights or {}) mlp_input, weights = self.mixers[mixer_index](sources, capture) mixer_index += 1 mlp_output = block.mlp(block.mlp_norm(mlp_input)) sources.append(mlp_output) if trace is not None: trace.layer_input_rms.append(rms(mlp_input)) trace.branch_output_rms.append(rms(mlp_output)) trace.stream_state_rms.append( rms(torch.stack(sources, dim=0)) ) trace.depth_weights.append(weights or {}) assert self.output_mixer is not None hidden, output_weights = self.output_mixer(sources, capture) if trace is not None: trace.output_weights = output_weights 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)) partial = ( branch_output if partial is None else partial + branch_output ) 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 mixer_index % SUBLAYERS_PER_BLOCK == 0: completed.append(partial) partial = None assert partial is None assert len(completed) == BLOCKS + 1 assert self.output_mixer is not None 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 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 def cross_entropy(logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor: return F.cross_entropy(logits.float().view(-1, VOCABULARY), targets.view(-1)) @torch.no_grad() def evaluate( model: ReducedLanguageModel, corpus: ByteCorpus, starts: list[int], 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(starts, begin, end) with torch.autocast(device_type="cuda", dtype=torch.bfloat16): logits, _ = model(inputs) loss = F.cross_entropy( logits.float().view(-1, VOCABULARY), targets.view(-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 percentile(values: list[float], quantile: float) -> float: return float(np.quantile(np.asarray(values, dtype=np.float64), quantile)) def core_parameter_gradient_rms(model: ReducedLanguageModel) -> list[float]: values = [] for block in model.blocks: sum_square = 0.0 count = 0 for parameter in block.parameters(): if parameter.grad is None: continue gradient = parameter.grad.detach().float() sum_square += gradient.square().sum().detach().cpu().item() count += gradient.numel() values.append(math.sqrt(sum_square / count)) return values def diagnostic( model: ReducedLanguageModel, corpus: ByteCorpus, window_count: int, ) -> 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) loss = cross_entropy(logits, targets) loss.backward() gradients = core_parameter_gradient_rms(model) assert trace is not None return { "loss_nats": loss.detach().cpu().item(), "bits_per_byte": loss.detach().cpu().item() / math.log(2), "layer_input_rms": trace.layer_input_rms, "branch_output_rms": trace.branch_output_rms, "stream_state_rms": trace.stream_state_rms, "core_parameter_grad_rms_by_block": gradients, "depth_weights": trace.depth_weights, "output_weights": trace.output_weights, } def parameter_inventory(model: ReducedLanguageModel) -> 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.") ) embedding = model.token_embedding.weight.numel() + model.position_embedding.weight.numel() return { "total": total, "core": total - mixer, "mixer": mixer, "embedding": embedding, } def main() -> None: args = parse_args() if not torch.cuda.is_available(): raise RuntimeError("CUDA is required by the frozen protocol") if args.run_kind != "smoke" and args.seed not in EXPECTED_SEEDS: raise ValueError(f"formal/replay seed is not preregistered: {args.seed}") 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["dataset"]["revision"] != ( "b08601e04326c79dfdd32d625aee71d232d685c3" ): raise ValueError("dataset revision mismatch") corpus = ByteCorpus(args.cache_dir, manifest, device) model = ReducedLanguageModel(args.architecture).to(device) initial_common_hash = state_hash(model, include_mixers=False) initial_mixer_hash = ( state_hash(model, include_mixers=True) if args.architecture != "baseline" else None ) inventory = parameter_inventory(model) decay_parameters: list[nn.Parameter] = [] no_decay_parameters: list[nn.Parameter] = [] for parameter in model.parameters(): if parameter.ndim >= 2: decay_parameters.append(parameter) 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 EVAL_STEPS if step <= args.steps) | {0, args.steps} ) evaluations = [ { "step": 0, **evaluate( model, corpus, corpus.validation_starts, args.validation_windows, args.eval_batch_size, ), } ] 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, corpus.validation_starts, args.validation_windows, args.eval_batch_size, ), } ) model.train() training_peak_allocated = torch.cuda.max_memory_allocated() training_peak_reserved = torch.cuda.max_memory_reserved() diagnostic_result = diagnostic(model, corpus, args.diagnostic_windows) final_common_hash = state_hash(model, include_mixers=False) final_mixer_hash = ( state_hash(model, include_mixers=True) if args.architecture != "baseline" else None ) timing = { "warmup_steps_excluded": args.timing_warmup, "measured_steps": len(step_times), "mean_ms": statistics.fmean(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, "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" ], }, "model": { "layers": LAYERS, "sublayers": SUBLAYERS, "blocks_for_block_attnres": BLOCKS, "sublayers_per_attnres_block": SUBLAYERS_PER_BLOCK, "d_model": D_MODEL, "heads": HEADS, "d_head": D_HEAD, "d_ff": D_FF, "context": CONTEXT, "vocabulary": VOCABULARY, "parameters": inventory, }, "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_common_parameters": initial_common_hash, "initial_mixer_parameters": initial_mixer_hash, "final_common_parameters": final_common_hash, "final_mixer_parameters": final_mixer_hash, }, "evaluations": evaluations, "training_history": training_history, "diagnostic": diagnostic_result, "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", "compile": False, }, } result["canonical_sha256_without_self"] = canonical_json_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), "architecture": args.architecture, "seed": args.seed, "steps": args.steps, "final_bpc": evaluations[-1]["bits_per_byte"], "initial_common_hash": initial_common_hash, "final_common_hash": final_common_hash, "canonical_sha256": result["canonical_sha256_without_self"], "timing": timing, }, ensure_ascii=False, indent=2, ) ) if __name__ == "__main__": main()