Design nanoGPT memorization experiments: AdamW vs Muon with clip_xmax monitoring

baseline/experiments/nanogpt_memorization/.gitignore

@@ -0,0 +1,7 @@
+__pycache__/
+.pytest_cache/
+*.pyc
+runs/
+resolved/
+*.pt
+*.bin

baseline/experiments/nanogpt_memorization/RESULTS.md

@@ -0,0 +1,30 @@
+# Execution ledger
+
+Date: 2026-09-14.
+
+## Completed in this preparation session
+
+- Read the repository's current baseline recipe, model, optimizer implementation,
+  spectral-monitor implementation, qualification document, and dependency file.
+- Pinned source commit and Git blob identities in `study.json`.
+- Created this separate experimental-design/measurement folder.
+- Ran `python -m pytest -q tests`: **32 passed in 2.78 seconds**, CPU.
+- Tests cover exact-versus-token recall, teacher-forced-versus-free-running
+  behavior, suffix alignment, exhaustive finite-universe exposure and ties,
+  prefix-search censoring, candidate batching, deterministic probes, exact
+  scheduled presentations, source-drift checks, plan counts, and WW guards.
+
+## Not completed / no results claimed
+
+- No end-to-end memorization trainer or natural/association data adapter has been
+  implemented; the integration contract is in `TRAINER_CONTRACT.md`.
+- No FineWeb-Edu download, canary-injected training, or AdamW/Muon training
+  campaign was run in this session.
+- No actual WeightWatcher call was executed here; WeightWatcher is not installed
+  in this execution environment. The hook reuses the inspected repository code.
+- No pinned-nanoGPT optimizer-step, resume, or target-hardware integration test
+  was executed. The controlled test model is not nanoGPT.
+- No empirically optimal configuration was established; the repository's
+  source-backed recipe is inherited without relabeling it as a frozen optimum.
+- There are no measured optimizer advantages, memorization rates, spectral
+  trends, or significance claims yet.

baseline/experiments/nanogpt_memorization/TRAINER_CONTRACT.md

@@ -0,0 +1,149 @@
+# Training adapter contract and acceptance gates
+
+This is a specification for the remaining integration, not a claim that the
+trainer exists. The baseline model, optimizer, and spectral monitor are reused;
+only the experimental data/sampling and behavioral-evaluation integration change.
+The generated baseline YAML files alone do not implement these interventions.
+
+## 1. Load and verify
+
+Resolve `study.json` with `prepare_plan.py`. Read the source-pinned recipe rather
+than copying hyperparameters by hand. Validate the prepared FineWeb cache with
+`rg_nanogpt_one_head.data.validate_prepared_data(data_dir, baseline_config)`.
+Require its dataset revision, exact token/byte counts, document-disjoint split
+flag, and SHA-256 checks. Do not alter the original `train.bin` or its metadata.
+Store intervention manifests and outputs under a new explicit experiment root.
+
+Use `GPT(GPTConfig(**cfg['model']))` from the pinned model module. Construct
+optimizer handles using `make_optimizer_handles(model, profile)` from the pinned
+optimizer module. Preserve the entire parameter partition and initialization.
+A dedicated training wrapper must explicitly support `muon`: the old dated
+campaign's CLI was configured around AdamW and MuonClip and is not an automatic
+launcher for this study.
+
+## 2. Freeze information and sample controls
+
+Create target identities, canary universes, templates, rule examples, and the
+reserved natural-document pool once with the study data seed. Save tokens,
+source document identities, split membership, masks, and content hashes.
+
+Reject every natural target whose full 16/32/64-token scored suffix is already
+present in the allowed background or another split at the corresponding score
+length. Match held-out natural controls on token length and source domain; the
+paired clean model controls the same target's intrinsic predictability. Disjoint
+document IDs alone do not eliminate duplicate text across documents.
+
+For random-token targets, scan for collisions rather than assuming probability
+zero. For canaries, audit context–suffix pairs and each template-equivalent
+mapping. For associations, rotate held-out entity–template combinations while
+ensuring every template is learned on other entities. `probes.py` currently
+implements random sequences and finite canary enumeration only; the natural
+reservation/decontamination and association/rule generators remain to be built.
+
+Do not infer empirical training frequency from the intended duplication factor.
+A full experimental presentation is a complete context plus target in one
+training record with all scored target tokens included in the loss.
+
+## 3. Paired training stream and exact dose
+
+Use distinct RNG streams for initialization, background batches, injection
+slots, filler, evaluation, and spectral randomization. Pair initialization and
+all training-data streams between optimizers. Hash initial model state and a
+preflight trace of batch IDs to verify pairing.
+
+For each reference-epoch interval, create a `presentation_schedule` over the
+actual available record slots. Its seed is derived deterministically from the
+training seed and interval, with the same derivation for both optimizers.
+Doses are full presentations per interval. The small rounding difference in the
+last interval is handled by its actual slot count, not by changing the dose.
+
+The adapter replaces selected ordinary-background records with intervention
+records. It must preserve the effective 8,192 supervised prediction targets per
+optimizer update. Pack or fill a 257-token record so a 256-token input has 256
+next-token labels; do not leave padding targets unintentionally supervised or
+silently mask different numbers of tokens between conditions. Record how much
+loss comes from context, target, and filler. Do not train a scored target across
+an attention-context boundary.
+
+Log each intervention visit with optimizer update, accumulation index, record
+index, target ID, template, and number of supervised target tokens. At every
+checkpoint compare actual cumulative visits with the planned schedule prefix.
+Zero-dose targets must have zero visits. Equivalent base replacement slots are
+used in the clean control, filled with ordinary background records.
+
+## 4. Optimizer step and checkpoint identity
+
+Reuse baseline learning-rate scheduling, gradient accumulation, global clipping,
+and optimizer-step helpers. Resolve warmup steps using the actual baseline
+trainer's rounding convention, then union these actual warmup endpoints into
+`plan['spectral_steps_before_warmup_union']`. Do not guess that convention in a
+new implementation or rescale the one-epoch schedule to the four-epoch horizon.
+
+An update's logged LR must be the value used for that update. Every behavioral
+and spectral record is keyed to the same post-update model-state hash and
+sampled-token counter. Initialization is update zero. Precision, device, compile
+policy, and dependency versions belong in the run fingerprint.
+
+Save checkpoint state atomically to a temporary file, fsync, and rename. Keep a
+rolling restart checkpoint plus immutable diagnostic states. Reuse the baseline
+checkpoint/RNG implementation and retain model, optimizer handles, LR position,
+CPU/accelerator RNG states, dedicated data RNG, dose counters, and current
+schedule offset. Verify round-trip equality and the next resumed update before
+starting long jobs. Never restart a partial result directory silently.
+
+## 5. Behavioral and spectral hooks
+
+Evaluate only fixed probes; never train on evaluation-generated continuations.
+Use the functions in `metrics.py` for suffix-only NLL, teacher-forced accuracy,
+free-running recall, canary ranking, and fixed-boundary prefix sweeps. Save the
+returned candidate score vectors, not merely rounded ranks.
+
+After saving a checkpoint, call:
+
+```python
+from monitor import monitor_training_state
+
+summary = monitor_training_state(
+    model,
+    run_dir,
+    step=completed_updates,
+    tokens_seen=sampled_target_tokens,
+    reference_tokens=80_000_000,
+    seed=training_seed,
+    fingerprint=run_fingerprint,
+    ww_config=baseline_config['weightwatcher'],
+)
+```
+
+The inherited monitor uses one `fix_fingers='clip_xmax'` analysis on CPU clones,
+with ERG and randomization enabled. It preserves RNG state and binds results to
+the checkpoint hash. Test that enabling diagnostics does not change the next
+training batch, next model update, or optimizer state.
+
+Keep `alpha_raw` and `alpha_clip_xmax` separately. Never substitute a raw fit
+for a failed clipped fit, synthesize ERG values, label a failed fit zero, or use
+WeightWatcher as the behavioral definition of memorization. If strict monitoring
+fails, retain the restart checkpoint and write a failure status before stopping;
+resume diagnostics only after resolving the cause. A finite alpha may still be
+scientifically invalid for a random-like spectrum; the interpretation screen is
+separate from successful software execution.
+
+## 6. Before a pilot may be labeled runnable
+
+The adapter must pass all of the following, beyond the component tests shipped
+here:
+
+- Tiny synthetic overfit: increasing exposure increases exact recall on a
+  deliberately learnable toy cohort; zero-dose rows never enter optimization.
+- Natural/association contamination audit and declared finite-universe audit.
+- Paired initial hashes, batch/slot traces, and exact realized dose counts.
+- Actual pinned nanoGPT AdamW and Muon forward/backward/update preflight.
+- Actual WeightWatcher output schema and RNG/model invariance on target hardware.
+- Full checkpoint round-trip, interruption/resume, and next-update equivalence.
+- Matching masks/target-token counts, context bounds, and final token horizon.
+- Same-state checkpoint agreement between online and offline behavioral metrics.
+
+The test suite in this folder currently checks measurement mathematics, suffix
+alignment, a controlled transition model, scheduling counts, plan resolution,
+source-drift refusal, and monitoring configuration guards. It does not establish
+that these end-to-end integration gates have passed.

baseline/experiments/nanogpt_memorization/metrics.py

@@ -0,0 +1,127 @@
+"""Behavioral measurements: no spectral metric is used to define memorization."""
+from __future__ import annotations
+import math
+from typing import Callable, Sequence
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+
+def recall(generated: Sequence[int], target: Sequence[int]) -> dict:
+    g, y = np.asarray(generated), np.asarray(target)
+    if g.ndim != 1 or y.ndim != 1 or len(y) == 0 or g.shape != y.shape:
+        raise ValueError("Generated and target sequences must be nonempty, 1-D, and equal-length.")
+    match = g == y
+    failures = np.flatnonzero(~match)
+    return {"sequence_exact_match": bool(match.all()),
+            "free_running_token_match": float(match.mean()),
+            "longest_exact_prefix": int(failures[0]) if len(failures) else len(y),
+            "continuation_tokens": len(y)}
+
+
+def exposure(candidate_nll: Sequence[float], target_index: int) -> dict:
+    """Exact finite-universe exposure; smaller NLL ranks first.
+
+    Return tie bounds. Primary exposure is conservative: all-equal scores -> 0.
+    These are full-universe ranks ONLY if every possible candidate was scored.
+    """
+    scores = np.asarray(candidate_nll, dtype=np.float64)
+    if scores.ndim != 1 or not len(scores) or not np.isfinite(scores).all():
+        raise ValueError("All candidate NLL scores must be finite and present.")
+    if not 0 <= target_index < len(scores):
+        raise ValueError("Target index is outside the candidate universe.")
+    value = scores[target_index]
+    rank_min = 1 + int(np.count_nonzero(scores < value))
+    rank_max = int(np.count_nonzero(scores <= value))
+    bits = math.log2(len(scores))
+    return {"candidate_count": len(scores), "rank_min": rank_min, "rank_max": rank_max,
+            "exposure_bits_lower": bits - math.log2(rank_max),
+            "exposure_bits_upper": bits - math.log2(rank_min),
+            "tie_count": rank_max - rank_min + 1}
+
+
+def prefix_compression(context: Sequence[int], target: Sequence[int],
+                       lengths: Sequence[int], generate: Callable[[list[int], int], list[int]]) -> dict:
+    """Sweep suffixes of ONE fixed context; never move the target boundary.
+
+    This is prefix-constrained prompt compression, not optimized adversarial ACR.
+    Failure is censored, not proof that no shorter/adversarial prompt exists.
+    """
+    if len(context) == 0 or len(target) == 0 or len(lengths) == 0:
+        raise ValueError("Context, target, and prefix grid must be nonempty.")
+    if any(int(n) != n for n in lengths):
+        raise ValueError("Prefix lengths must be integers.")
+    grid = sorted(set(int(n) for n in lengths))
+    if grid[0] < 1 or grid[-1] > len(context):
+        raise ValueError("Prefix grid is outside the available context.")
+    rows = []
+    for p in grid:
+        rows.append({"prefix_tokens": p, **recall(generate(list(context[-p:]), len(target)), target)})
+    successes = [r["prefix_tokens"] for r in rows if r["sequence_exact_match"]]
+    shortest = min(successes) if successes else None
+    return {"rows": rows, "shortest_successful_prefix_in_grid": shortest,
+            "prefix_compression_ratio": len(target) / shortest if shortest is not None else None,
+            "search_censored": shortest is None,
+            "prompt_class": "fixed_context_suffixes_only"}
+
+
+@torch.inference_mode()
+def score_continuation(model, prefix: Sequence[int], target: Sequence[int]) -> dict:
+    """Teacher-forced suffix NLL/accuracy and separate free-running recall."""
+    if len(prefix) == 0 or len(target) == 0:
+        raise ValueError("Prefix and continuation must both be nonempty.")
+    if len(prefix) + len(target) - 1 > model.cfg.block_size:
+        raise ValueError("Teacher-forced input would exceed model context.")
+    device = next(model.parameters()).device
+    full = torch.tensor([list(prefix) + list(target)], dtype=torch.long, device=device)
+    previous = model.training
+    model.eval()
+    try:
+        logits, _ = model(full[:, :-1])
+        selected = logits[:, len(prefix) - 1:, :]
+        labels = full[:, len(prefix):]
+        losses = F.cross_entropy(selected.reshape(-1, selected.size(-1)), labels.reshape(-1), reduction="none")
+        nll = float(losses.mean().item())
+        if not math.isfinite(nll):
+            raise RuntimeError("Nonfinite continuation loss; do not report a memorization score.")
+        greedy = model.generate_greedy(full[:, :len(prefix)], len(target))[0, len(prefix):].tolist()
+        return {**recall(greedy, target), "suffix_nll": nll,
+                "suffix_nll_sum": float(losses.sum().item()),
+                "suffix_perplexity": math.exp(nll) if nll < 709 else None,
+                "teacher_forced_token_accuracy": float((selected.argmax(-1) == labels).float().mean().item())}
+    finally:
+        model.train(previous)
+
+
+@torch.inference_mode()
+def rank_canary(model, prefix: Sequence[int], candidates: Sequence[Sequence[int]],
+                target_index: int, batch_size: int = 4) -> dict:
+    """Exhaustively score a small declared universe, not sampled rank estimation."""
+    if len(prefix) == 0 or len(candidates) == 0 or batch_size < 1:
+        raise ValueError("Nonempty prefix/universe and positive batch size are required.")
+    if not 0 <= target_index < len(candidates):
+        raise ValueError("Target index is outside the candidate universe.")
+    length = len(candidates[0])
+    if not length or any(len(c) != length for c in candidates):
+        raise ValueError("The declared universe must use equal nonzero token lengths.")
+    if len({tuple(c) for c in candidates}) != len(candidates):
+        raise ValueError("Candidate universe contains duplicates.")
+    if len(prefix) + length - 1 > model.cfg.block_size:
+        raise ValueError("Canary scoring would exceed context.")
+    device = next(model.parameters()).device
+    previous = model.training
+    model.eval()
+    scores = []
+    try:
+        for start in range(0, len(candidates), batch_size):
+            chunk = candidates[start:start + batch_size]
+            full = torch.tensor([list(prefix) + list(c) for c in chunk], dtype=torch.long, device=device)
+            logits, _ = model(full[:, :-1])
+            logits = logits[:, len(prefix) - 1:, :]
+            target = full[:, len(prefix):]
+            loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), target.reshape(-1), reduction="none")
+            scores.extend(loss.reshape(len(chunk), length).sum(-1).double().cpu().tolist())
+    finally:
+        model.train(previous)
+    return {**exposure(scores, target_index), "candidate_nll_sum": scores,
+            "rank_scope": "entire_declared_finite_universe"}

baseline/experiments/nanogpt_memorization/monitor.py

@@ -0,0 +1,22 @@
+"""Live-model hook reusing the repository's audited spectral implementation."""
+from __future__ import annotations
+from pathlib import Path
+
+
+def monitor_training_state(model, run_dir: str | Path, *, step: int, tokens_seen: int,
+                           reference_tokens: int, seed: int, fingerprint: str, ww_config: dict):
+    if ww_config.get("fix_fingers") != "clip_xmax":
+        raise ValueError("This study requires fix_fingers='clip_xmax'.")
+    for key in ("enabled", "ERG", "randomize", "strict", "require_raw_alpha"):
+        if ww_config.get(key) is not True:
+            raise ValueError(f"This study requires weightwatcher.{key}=true.")
+    if step < 0 or tokens_seen < 0 or reference_tokens < 1 or not fingerprint:
+        raise ValueError("Invalid checkpoint identity.")
+    from rg_nanogpt_one_head.spectral import run_weightwatcher
+    # Upstream creates CPU clones of Q/K/V/O/MLP-in/MLP-out, calls WW once,
+    # preserves CPU/accelerator RNG, binds output to the model-state hash,
+    # and stores alpha, raw_alpha, num_fingers, ERG_gap, traps and rand_distance.
+    # No mutation of weights, clipping of weights, or feedback to optimization.
+    return run_weightwatcher(model, run_dir, step=step, tokens_seen=tokens_seen,
+                             train_tokens=reference_tokens, config=dict(ww_config),
+                             seed=seed, fingerprint=fingerprint)