Bundles plus a two-model harness¶

A comparison of a clean run with an edited run is only worth reading when the two runs line up site by site and differ by the stated edit alone; and in a two-model pipeline, where one model drafts text that a second model reads, an edit to the drafter can reach the critic only through the tokens the critic is given. This entry is a method demonstration rather than the reproduction of one paper; the closest published setup is the clean-versus-ablated comparison common to activation-patching work. It bundles a clean and a zero-ablated run of GPT-2 small and of ResNet-18 and checks every module's difference against two plain hook runs, then records GPT-2 small drafting a sentence and DistilGPT2 judging it as one bundle, steers the drafter, and checks that the recorded tokens, the fed_by token spans and the reported reach of the steer all match a hand-run pipeline. The entry passes: every bundle difference matches the plain hooks to exactly 0, the hand-run tokens agree, the spans point at the right critic tokens, and the controls refuse as declared.

Oracles¶

The rules are the ones registered for this entry before any run.

The first run of this script failed the GPT-2 bundle difference with a largest difference of 57.9: on a fork, zeroing the layer-6 MLP's output also rewrote the value recorded for the MLP's inner dropout module, whose output is the same operation. A post-capture module edit is now recorded as its own replacement operation after the module's last operation, the shape a capture-time edit always had, so the inner module keeps the value a plain hook leaves it. No oracle or control moved.

Controls¶

Two prompts of the same token count but different words are compared by design (the gate checks shapes, not values): that comparison ran and returned 585 rows, and is reported here, not graded.

Artifact and replay¶

The whole bundle (clean, ablated, draft, critique and the steered alternatives) is saved as a portable bundles_harness.tlspec: a directory of 59 files, 367,121,025 bytes in all, with gallery sorted-path/content sha256 635d4e6682a89d9df30f55c2d18ff6f3ad701b2f02ee2a330be7c820ec635409. It holds one member directory per run (records as JSON tables, tensors as safetensors) and no pickled file. The script reloads it and replays the fed_by spans, the prompt divergence and the lm_head comparison from the reloaded copy: all three match the in-memory values.

Basis-grid row¶

Counts are nonblank, noncomment, non-docstring lines of the script's marked regions. The TorchLens region builds the bundles, the ablation, the harness and the steered alternative; the plain region is the hook-based differences and the hand-run two-model pipeline. No count against an original reference implementation was measured.

The stamp¶

{
  "commit": "d9f8f53c9bce262e0e483e3a835c8e342f9402fa",
  "date_utc": "2026-11-05T01:30:58Z",
  "device": "cpu",
  "libraries": {
    "numpy": "2.2.6",
    "torchvision": "0.28.0+cpu",
    "transformers": "5.18.0"
  },
  "machine": "x86_64",
  "platform": "Linux-3.10.0-1160.119.1.el7.x86_64-x86_64-with-glibc2.36",
  "python": "3.10.19",
  "torch": "2.13.0+cpu",
  "torch_threads": 1,
  "torchlens": "2.38.1",
  "wall_seconds": 141.9
}

Vary it¶

Use --steer-layer 4 to steer a different GPT-2 block, --critic gpt2 to make a second GPT-2 the critic, or --draft-tokens 20 for a longer draft. These small examples use a randomly initialized module and need no downloads. The first builds a clean and an ablated run of a nested module, bundles them, and checks the difference at every module against plain hooks that record each module's output as its consumers see it, including the inner module that shares the ablated module's last operation and keeps its own value.

# recipe-test: tool=torchlens expected=True
import torch
from torch import nn
import torchlens as tl


class Nested(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.inp = nn.Linear(8, 16)
        self.mlp = nn.Sequential(nn.Linear(16, 16), nn.ReLU())
        self.head = nn.Linear(16, 4)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.head(self.mlp(torch.relu(self.inp(x))))


torch.manual_seed(0)
model = Nested().eval()
x = torch.randn(2, 8)
clean = tl.trace(model, x, save="all", capture=tl.options.CaptureOptions(replayable=True))
ablated = clean.fork("ablated")
ablated.do(tl.module("mlp"), tl.zero_ablate())
bundle = tl.Bundle({"clean": clean, "ablated": ablated}, baseline="clean")


def plain_outputs(zero_mlp: bool) -> dict[str, torch.Tensor]:
    """Module outputs as their consumers see them, from plain hooks."""
    seen: dict[str, torch.Tensor] = {}
    handles = []
    if zero_mlp:  # registered first, so the recorders below see the zeroed output
        handles.append(model.mlp.register_forward_hook(lambda mod, a, out: torch.zeros_like(out)))
    handles += [
        m.register_forward_hook(lambda mod, a, out, n=n: seen.setdefault(n, out))
        for n, m in model.named_modules()
        if n
    ]
    with torch.no_grad():
        model(x)
    for h in handles:
        h.remove()
    return seen


plain_clean, plain_ablated = plain_outputs(zero_mlp=False), plain_outputs(zero_mlp=True)
RESULT = all(
    torch.equal(
        bundle["ablated"].modules[n].calls[0].outs[0] - bundle["clean"].modules[n].calls[0].outs[0],
        plain_ablated[n] - plain_clean[n],
    )
    for n in plain_clean
)

# recipe-test: tool=torchlens expected=True
rows = bundle.compare("clean", "ablated", metric="relative_l2")
RESULT = any(float(distance) > 0 for site, distance in rows)

The second example records a two-step pipeline as one bundle with a fed_by row whose spans are given by hand, as the gallery script gets them from token_spans='auto'.

# recipe-test: tool=torchlens expected=True
lm = nn.Sequential(nn.Embedding(9, 4), nn.ReLU(), nn.Linear(4, 9))
draft_ids = torch.tensor([[1, 2, 3, 5, 6, 7]])  # prompt [1, 2, 3], then the draft [5, 6, 7]
critique_ids = torch.tensor([[4, 5, 6, 7, 8]])  # the critique prompt quotes the draft
harness = tl.Bundle()
harness.add(tl.trace(lm, draft_ids), names="draft")
harness.add(tl.trace(lm, critique_ids), names="critique")
draft_output = {"stream": "input", "batch_row": 0, "completion": [3, 6]}
tl.io.attach_sidecar(harness["draft"], "torchlens.boundary_facts", draft_output)
harness.relate("fed_by", "critique", "draft", token_spans=[[3, 6, 1, 4]], label="critique")
row = harness.relations[0]
RESULT = row.params["basis"] == "verified" and harness.upstream("critique") == ("draft",)