B13 interp-on-checkpoints methodology

status: not-started
owner: unassigned
created: 2026-06-13
parent_plan: docs/PLAN.md §3 Tier B (B13)
related_prds: sae-timeline-viewer.md (visualization sibling — ships the chart, this PRD ships the infra producing the data)

PRD — Interpretability across a training-run timeline

Goal

Replay the shipped interp tools (posttrainllm sae, posttrainllm memit, posttrainllm rome, posttrainllm patch, posttrainllm causal-trace) across the multi-checkpoint history a single posttrainllm train run produces, so we can see when a feature emerges, not just that a final model has it. Pythia (Biderman et al. 2023) and OLMo (Groeneveld et al. 2024) made this the standard small-model interp protocol; no competitor ships it at posttrainllm’s “every-byte-of-code-here” scale.

The /sae-timeline.astro viewer (already shipped) consumes the output of this PRD’s tooling. Without this PRD, the viewer renders empty.

Why now

• E8 train-time eval hook shipped already (writes <step>-evals.jsonl). This is the interp analogue — replay shipped interp probes per saved checkpoint instead of per-eval task.
• Save-every-N is the natural hook point. Already implemented for the eval-hook PRD; we attach to the same callback list.
• Decisive value for A1 specialist work: training a tool-caller and not knowing when features emerge means we’re working blind. With this PRD, we can pinpoint “tool-call refusal” feature appearance.

Scope — in

• Sources/TinyGPT/InterpReplay.swift — orchestrator. Walks a history directory, loads each checkpoint, runs the requested probe, writes one JSONL row per (probe, layer, checkpoint).
• New CLI: posttrainllm interp-replay <history-dir> --probe {sae,memit, rome,patch,causal-trace} [--layer L | --layers SPEC] --out timeline.jsonl
• Row schema: {step: Int, ckpt_hash: String, probe: String, layer: Int, metric: String, value: Double, extra: [String: Any]}. The extra blob carries per-probe specifics (SAE MSE + L0, MEMIT per-fact residual, etc.). The shared metric axis lets the timeline viewer plot any probe.
• Default behavior: when only --out is given (no probe), replay all probes that have inputs (a .sae config, a --memit-facts file, etc.) — turns this into “give me the full interp timeline for this run”.
• Wire into posttrainllm train’s optional --interp-every N hook (same shape as --eval-every). Non-blocking; skip if previous interp pass is still running.

Scope — out

• New interp probes. This PRD only orchestrates the shipped ones. New probes (e.g. concept bottlenecks, activation patching at scale) go in their own PRDs.
• Live overlay of in-progress checkpoint data while training. The shipped train viewer (C10) already shows real-time loss; the interp side is fast enough as post-hoc.
• Cross-run timeline — comparing the same feature across two separate training runs. Useful but adds 200 LOC of UI; defer.

Files to touch

Don’t touch

• Individual probe CLIs (Sources/TinyGPT/Sae.swift, Memit.swift, etc.) — extend, don’t fork. The orchestrator calls them as library functions, not as subprocesses.
• /sae-timeline.astro — separate PRD (already shipped). Output format here must match what the viewer expects.

Acceptance criteria

• posttrainllm interp-replay <history> --probe sae --out timeline.jsonl on a 5-checkpoint history produces a JSONL with 5 × n_layers rows.
• --interp-every 1000 integrated into posttrainllm train doesn’t block the training loop more than 100 ms per checkpoint (the heavy interp work happens in a separate process).
• /sae-timeline.astro drag-drops the produced JSONL and renders MSE + L0 curves correctly.
• Smoke script interp-replay-smoke.sh passes in CI on the shipped 22M-class checkpoint.
• docs/interpretability.md invocation example reproduces.

Reference patterns

• Sources/TinyGPT/RunLmEval.swift — the post-checkpoint multi-task orchestrator pattern. Same shape: load checkpoint, run probe, write row, repeat.
• posttrainllm sae’s existing .sae sidecar — the per-checkpoint output artifact this PRD batches.
• Pythia paper — the methodology is theirs; cite. Their interp-on-checkpoints findings (feature emergence + reversal points) are the kinds of analysis this enables.

Open questions

• Whether to pre-materialize SAE features per checkpoint (a few MB per checkpoint × 50 checkpoints = manageable) or compute on-demand in the viewer. Recommendation: pre-materialize; the viewer’s drag-drop UX requires it.
• Default --layers selection — every layer (expensive) vs every 4th (Pythia’s choice). Recommendation: every 4th + 0 + last; override with --layers all.