"""Trace L5-N541 and L6 peaks for the 80 full Moby Dick routes. One-command use, from the znou tools directory: python trace_541_full_to_l6.py The protocol matches the project's implicit-resonance regime: GPT-2 Small, the pinned revision, float32, BOS plus every ordinary input token, post-GELU MLP activations, and a maximum-over-position destination readout. """ from __future__ import annotations import csv import json import os from pathlib import Path # Required by deterministic CUDA matrix multiplication. This must be set # before importing torch, matching the canonical znou probe stack. os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8") import torch from transformers import GPT2LMHeadModel, GPT2TokenizerFast MODEL_REVISION = "607a30d783dfa663caf39e06633721c8d4cfcd7e" BOS_TOKEN_ID = 50256 SOURCE_LAYER = 5 TARGET_LAYER = 6 SOURCE_NEURON = 541 SCRIPT_DIR = Path(__file__).resolve().parent REPO_DIR = SCRIPT_DIR.parent L5_FILE = REPO_DIR / "data" / "the_sea_implicit_resonance.json" L6_FILE = Path( r"C:\Users\ursad\Desktop\gamedev\claude\dungeon\data" ) / "the_sea_implicit_resonance_L6.json" OUTPUT_FILE = SCRIPT_DIR / "results" / "n541_full_80_l5_l6_peak_trace.csv" def load_routes(): if not L5_FILE.exists(): raise FileNotFoundError(f"Missing L5 corpus file:\n{L5_FILE}") if not L6_FILE.exists(): raise FileNotFoundError(f"Missing L6 corpus file:\n{L6_FILE}") l5 = json.loads(L5_FILE.read_text(encoding="utf-8")) l6 = json.loads(L6_FILE.read_text(encoding="utf-8")) if len(l5) != len(l6): raise RuntimeError(f"Corpus lengths differ: L5={len(l5)}, L6={len(l6)}") routes = [] for index, (left, right) in enumerate(zip(l5, l6)): if left["sentence"] != right["sentence"]: raise RuntimeError(f"Sentence mismatch at corpus row {index + 1}") if int(left["neuron_id"]) == SOURCE_NEURON: routes.append({ "corpus_row": index + 1, "text": left["sentence"], "expected_l6_winner": int(right["neuron_id"]), }) if len(routes) != 80: raise RuntimeError(f"Expected 80 L5-N541 routes, found {len(routes)}") return routes def top_two(values): result = torch.topk(values, k=2) return { "winner": int(result.indices[0]), "winner_value": float(result.values[0]), "runner_up": int(result.indices[1]), "runner_up_value": float(result.values[1]), "margin": float(result.values[0] - result.values[1]), } def token_context(tokens, position, radius=4): start = max(0, position - radius) end = min(len(tokens), position + radius + 1) pieces = [] for i in range(start, end): token = tokens[i].replace("\n", "\\n") pieces.append(f"[{token}]" if i == position else token) return "".join(pieces) class Runner: def __init__(self, model, tokenizer, device): self.model = model self.tokenizer = tokenizer self.device = device @torch.inference_mode() def run(self, text, ablate_position=None): body_ids = self.tokenizer(text, add_special_tokens=False)["input_ids"] ids = [BOS_TOKEN_ID] + body_ids input_ids = torch.tensor([ids], dtype=torch.long, device=self.device) captured = {} def source_hook(_module, _inputs, output): edited = output.clone() if ablate_position is not None: edited[:, ablate_position, SOURCE_NEURON] = 0.0 captured["l5"] = edited.detach().float().cpu()[0] return edited def target_hook(_module, _inputs, output): captured["l6"] = output.detach().float().cpu()[0] h5 = self.model.transformer.h[SOURCE_LAYER].mlp.act.register_forward_hook( source_hook ) h6 = self.model.transformer.h[TARGET_LAYER].mlp.act.register_forward_hook( target_hook ) try: self.model(input_ids=input_ids, use_cache=False) finally: h5.remove() h6.remove() tokens = [""] + [self.tokenizer.decode([i]) for i in body_ids] return captured["l5"], captured["l6"], tokens def peak_for_neuron(activations, neuron): values = activations[:, neuron] value, position = torch.max(values, dim=0) return float(value), int(position) def main(): torch.backends.cuda.matmul.allow_tf32 = False torch.backends.cudnn.allow_tf32 = False torch.set_float32_matmul_precision("highest") torch.use_deterministic_algorithms(True) routes = load_routes() device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"Device: {device}") print("Loading pinned GPT-2 Small in float32...") tokenizer = GPT2TokenizerFast.from_pretrained( "gpt2", revision=MODEL_REVISION, add_prefix_space=False ) model = GPT2LMHeadModel.from_pretrained( "gpt2", revision=MODEL_REVISION, dtype=torch.float32 ).to(device).eval() runner = Runner(model, tokenizer, device) rows = [] for route_number, route in enumerate(routes, 1): natural_l5, natural_l6, tokens = runner.run(route["text"]) l5_top = top_two(natural_l5.amax(dim=0)) l6_top = top_two(natural_l6.amax(dim=0)) if l5_top["winner"] != SOURCE_NEURON: raise RuntimeError( f"Route {route_number} failed L5 verification: " f"expected N541, measured N{l5_top['winner']}" ) if l6_top["winner"] != route["expected_l6_winner"]: raise RuntimeError( f"Route {route_number} failed L6 verification: expected " f"N{route['expected_l6_winner']}, measured N{l6_top['winner']}" ) l5_value, l5_position = peak_for_neuron(natural_l5, SOURCE_NEURON) l6_value, l6_position = peak_for_neuron(natural_l6, l6_top["winner"]) ablated_l5, ablated_l6, ablated_tokens = runner.run( route["text"], ablate_position=l5_position ) if tokens != ablated_tokens: raise RuntimeError(f"Token mismatch on route {route_number}") ablated_top = top_two(ablated_l6.amax(dim=0)) ablated_value, ablated_position = peak_for_neuron( ablated_l6, ablated_top["winner"] ) original_winner_after_value, original_winner_after_position = peak_for_neuron( ablated_l6, l6_top["winner"] ) rows.append({ "route": route_number, "corpus_row": route["corpus_row"], "text": route["text"], "token_count_including_bos": len(tokens), "l5_541_peak_position": l5_position, "l5_541_peak_token": tokens[l5_position], "l5_541_peak_context": token_context(tokens, l5_position), "l5_541_peak_activation": l5_value, "natural_l6_winner": l6_top["winner"], "natural_l6_runner_up": l6_top["runner_up"], "natural_l6_margin": l6_top["margin"], "natural_l6_peak_position": l6_position, "natural_l6_peak_token": tokens[l6_position], "natural_l6_peak_context": token_context(tokens, l6_position), "natural_l6_peak_activation": l6_value, "l6_minus_l5_peak_position": l6_position - l5_position, "ablated_l6_winner": ablated_top["winner"], "ablated_l6_runner_up": ablated_top["runner_up"], "ablated_l6_margin": ablated_top["margin"], "ablated_l6_peak_position": ablated_position, "ablated_l6_peak_token": tokens[ablated_position], "ablated_l6_peak_context": token_context(tokens, ablated_position), "ablated_l6_peak_activation": ablated_value, "winner_changed": l6_top["winner"] != ablated_top["winner"], "natural_winner_after_ablation_max": original_winner_after_value, "natural_winner_after_ablation_peak_position": ( original_winner_after_position ), "natural_winner_max_delta": original_winner_after_value - l6_value, "natural_winner_at_original_peak_natural": float( natural_l6[l6_position, l6_top["winner"]] ), "natural_winner_at_original_peak_ablated": float( ablated_l6[l6_position, l6_top["winner"]] ), "natural_winner_at_original_peak_delta": float( ablated_l6[l6_position, l6_top["winner"]] - natural_l6[l6_position, l6_top["winner"]] ), }) print( f"{route_number:>2}/80 L5 peak {l5_position:>3} " f"-> L6-N{l6_top['winner']:>4} peak {l6_position:>3} " f"-> ablated N{ablated_top['winner']:>4}" ) OUTPUT_FILE.parent.mkdir(parents=True, exist_ok=True) with OUTPUT_FILE.open("w", newline="", encoding="utf-8") as handle: writer = csv.DictWriter(handle, fieldnames=rows[0].keys()) writer.writeheader() writer.writerows(rows) same_position = sum(r["l6_minus_l5_peak_position"] == 0 for r in rows) adjacent_after = sum(r["l6_minus_l5_peak_position"] == 1 for r in rows) changed = sum(r["winner_changed"] for r in rows) print(f"Saved: {OUTPUT_FILE}") print(f"Same L5/L6 peak position: {same_position}/80") print(f"L6 peak one token later: {adjacent_after}/80") print(f"L6 winner changed after ablation: {changed}/80") if __name__ == "__main__": main()