The engine

Two halves, and the seam between them is the design. The compiler takes a document and decides everything that needs a model but no activations — an address per site, a width per address, the token ids of every prompt, the order the forwards must run in — and returns a Plan: strings, integers and tuples. An engine walks that plan. What each step means is written once, in engine/steps.py, as plain functions taking an engine; an engine supplies eight members and nothing else.

Results come home by filling the plan in. A step's results is the only mutable thing on it, so every number sits on the node that produced it, and a node's path in the tree is its path on disk.

At a glance #

explain compiles a document against the model's config, with no weights loaded, and prints the plan it produced:

$ uv run causalab-mini explain documents/v2/patching.json
<root>: Plan
  score: Observe  metrics=['iia/match', 'logit_diff/logit_diff']  saves=['iia.json', 'logit_diff.json']
      forward 'original' on 'counterfactual'  (4 rows x 11 tokens)
        read  'v_cf' at block_output[0] pos={index:-1} via 'identity'
      forward 'patched' on 'base'  (4 rows x 11 tokens)
        write 'patch' at block_output[0] pos={index:-1} swap(v_cf) via 'identity'
        read  'logits' at lm_head pos={index:-1} via 'identity'

That tree is the plan. Everything below is what put it there and what walks it.

The pipeline #

One document, all the way through #

The document is the smallest real experiment in the repository: read the residual stream at the last token of a counterfactual prompt, swap it into the same place on the base prompt, score what the model then says. The home page shows it in full; what follows is what happens to it.

1. Compile #

build_spec works in the order of the file. Each site's layer is checked against engine.num_layers, then each site becomes an Address through engine.locate; each site that names units or heads gets its count from engine.width or engine.heads, and an index past the end is refused here rather than on a tensor. A featurizer's d is derived from (model, site), because a traced block may not decide anything from a tensor — including how wide the tensor it is about to rotate is.

Then the forwards. A read in model A may be the operand of a write in model B, so B runs after A; _schedule topologically sorts the (model, input) pairs and refuses a cycle. Here it gives original on counterfactual, then patched on base — the order explain printed.

Within a forward, taps are sorted into forward order by Address.key, and at one address the writes go before the reads, because a read in a model sees that model's writes.

Three checks belong to the compiler because they are decidable without a run:

• An attention pattern swapped in from another prompt must line up key for key — the last axis of attention_probs is the keys of the padded batch, so the two prompts must pad to the same width and have the same number of real tokens row by row. Both masks are in hand before any forward.
• A read that buffers every decode step has a size bound — decode × rows × width × 4 bytes, refused over 256 MiB. All three numbers are known before a model is loaded, so the line is drawn before the memory is held.
• A write's features must be distinct indices inside the featurizer's space.

2. The plan #

A Plan is a dict of named steps, and is itself a step, so plans nest — a sweep point is a nested plan. Four kinds of step sit in one: Featurizers constructs the parameter sets, Observe is one pass (forwards in order, then metrics over what they read), Fit is that pass N times with an optimizer between, and Weights turns fitted parameters into results. Inside an Observe, a Forward holds its token ids and a tuple of Taps, one per address, each with its writes and reads.

Every field is frozen except results. There is no execute on a step: what running one means is the engine's business, and a plan that knew would only work on one engine.

Two things cross between steps. An Observe's outputs publish a read — kept, or reduced with mean or pca — under a name a later step's write can name as its operand via ref. And the featurizers are shared live, which is what lets a fit train the rotation a later step scores with. Both live in the walk's State, which is a local of the walk: never saved, never shipped, gone with the plan. A nested plan gets its own outputs and the same featurizers, so sweep points cannot see each other's values.

3. The walk #

steps.run(engine, step) is the whole executor: a Plan recurses into its children, and each other kind of step goes to its own function. Every one of those functions takes the engine as its first argument and calls it for exactly one thing — running a forward. None of them touches the model.

observe runs the pass, then computes the metrics. A metric's eligible rows are the intersection of two halves: the column half, which the client decided from the data, and the position half, which only the run can know. The scored metrics are returned live, because a fit differentiates them, and recorded detached, because what comes home should not carry a graph.

4. The per-forward frame #

Positions are resolved here — inside the session, in the walk, against engine.tokenizer — and not on the client, so that a text anchor is looked for in the text the model will actually see. One Frame is built per forward and every tap of that forward shares it.

The frame's character map is the expensive half: O(L) decode calls of O(L) work per row, measured at 840 ms for 64 rows of 260 tokens. So it is built only when the pass has a position the document does not already fix — a text anchor, or any cut of the continuation. A pass of bare -1s resolves the same integers on every row and every run, and skips the map.

Before looking for any text in a row, the run decodes row 0's ids with its own tokenizer and compares the string to the one the client recorded when it encoded them. Locally those are one object; remotely they are two, and a disagreement would put every position in range and in the wrong place, silently, because both sides produce integers. One row is a sample, not a proof.

What resolution returns beside the forward is the run's own record: per op, each row's window, the reason it is empty when it is, and the tokens it actually addressed. That is plain tuples of integers and strings, so it comes home in results like a metric does.

5. One forward #

engine.forward(forward, values, featurizers) is the only place an engine touches a model's insides. It returns nothing: it leaves what it read in values, keyed by the read's name. values starts as a copy of the published outputs, so a write's operand may be either a read of this pass or something an earlier step published — and the engine cannot tell the difference and does not need to.

6. Home #

results_of flattens a run into {step path: {name: tensor}} — plain strings and tensors, with none of this package's classes in it — and fill is its inverse. That pair is what crosses back from a server, and it is also what makes the returned plan the same object that was passed in.

Plan.write then writes the save manifest and nothing else: nothing lands on disk that is not listed in a step's saves. A nested plan writes into a directory named by its step name, so a swept point's files land under pos=-1/. On the root, the directory also gets document.json (the experiment, verbatim) and run.json.

$ uv run causalab-mini run documents/v2/patching.json --device-map cpu --out out
out/document.json
out/run.json
out/iia.json
out/logit_diff.json

run.json is the engine, whether it ran remotely, the batch size, the versions of the five packages whose code decides what a block does, a sha256 over this package's own source, and the dtype the model was actually served in. No hostname, no timestamp, no absolute path — a record has to be byte-identical across machines to be worth comparing.

{
 "engine": "NNterpEngine",
 "remote": false,
 "batch_size": null,
 "versions": {
  "torch": "2.14.0",
  "transformers": "5.17.0",
  "nnsight": "0.8.1.dev125+ga8ee93782",
  "nnterp": "1.3.1.dev64+g0e4b401af",
  "pydantic": "2.13.5"
 },
 "causalab_mini": "76c8c9d0c0d1f3e0b8e2eaeb09564107d1c0a70cb9e8b7d6cecab09fdc92d8a1",
 "served_dtype": "torch.float32"
}

The engine contract #

Nine names: eight members, and model, which is an attribute. Nothing in base.py has a body, so an engine inherits no behaviour it did not ask for.

The walk over steps, the fit loop, the metrics and the write algebra are shared, so a second engine is these eight members and no more. A test asserts the set of public names each engine class overrides, on both engines, so a member added quietly fails the suite.

What both engines share below forward #

ops/intervene.py knows nothing about layers, models or documents: positions arrive as integers and activations arrive as tensors. One line is the seam the rest of the project is built around:

write = inverse(do(featurize(x)), err, x)

With the identity featurizer that collapses to "replace the tensor", which is activation patching. Swap in a rotation and the same line is DAS — the error term and the unselected directions come from x, the pre-write value, which is what makes a subspace swap leave the orthogonal complement alone.

gather takes a window per row and, of it, the features the selection names. Uniform windows keep the rectangle, (batch, seq, width) → (batch, w, width); ragged windows cannot and come back flat, (total, width), with the plan's own positions saying where each row begins and ends. Which shape it is is decided by the positions, never by the tensor — and, for a batched run, decided once by the compiler over every row, because a window of a ragged pass's rows can happen to be rectangular. scatter is the inverse, and it clones.

Featurizer arithmetic runs inside intervene.exact, which disables autocast. A server may wrap a whole request in torch.autocast — NDIF does, at the served dtype — and inside one a matmul of two fp32 tensors comes back bf16 while a subtraction stays fp32, so half of a Cayley solve is downcast and linalg.solve refuses the pair. A featurizer declares its own dtype, fp32, and this is where that declaration is kept.

One session #

The nnterp engine's whole request is one model.session(...): not a session per forward, not a session per training step, not a loop outside that opens sessions.

    def execute(self, plan: Plan, remote: bool | str = False, batch_size: int | None = None) -> Plan:
        plan.provenance.update(provenance.record(self, remote, batch_size))
        engine, model = self, self.model
        with model.session(remote=remote):
            # What comes home is plain — strings and tensors. The client
            # already has the plan, so only what fills it in needs the trip,
            # and nothing of ours has to survive the way back. (`remote="local"`
            # never noticed: it does not serialize the way back. FINDINGS §19.)
            home = nnsight.save({})
            steps.run(engine, plan, batch_size=batch_size)
            home.update(plan_module.results_of(plan))
            # A server serves the dtype *it* chose; the document's is only a
            # request. Say what ran, where the run record is.
            home["served_dtype"] = str(model.lm_head.weight.dtype)
        plan.provenance["served_dtype"] = home.pop("served_dtype")
        plan_module.fill(plan, home)
        return plan

remote is the only difference between running here and running on NDIF. There is no second code path: the flag is passed to that one session. What the command line changes besides the flag is how the model was loaded — --engine ndif is this same engine with dispatch=False, a shell with no local weights, because the weights are the server's.

Results come home because the session saves a dict. nnsight.save({}) hands back an object the block fills; locally that is the dict itself, remotely it is the copy that will be shipped home. What goes into it is results_of(plan) — plain strings and tensors by step path. The plan itself does not come back: the client never gave it up, so only what fills it in needs the trip. Nothing of ours ships by value either way — the block's module references and the plan's classes resolve by import on the server, which is why it needs causalab_mini and nnterp installed at the client's versions (FINDINGS §24).

Two things happen at load. nnsight loads weights lazily, on the first trace, by replacing the module — so the loader forces dispatch=True and freezes afterwards, and what is frozen is what runs. The freeze is eval() and requires_grad_(False): the model is an instrument, not a parameter. Without it a fit's backward reaches every weight the patched forward touched and leaves a .grad the size of the model behind. Nothing reads it; the memory is real.

A forward with no decode is one model.trace(...). A forward that generates is one model.generate(...) with tracer.iter walking the steps: prompt-frame taps apply at step 0, the prefill; a step's own taps at that step; an "all" write at every step. Greedy, with EOS held off so the bound holds and the loop never outruns the run.

An interior — one argument of one call inside a module's forward — is resolved at locate time, on the client, against the loaded checkpoint's .source. The match is on the source line rather than the operation's name: nnsight names an operation {callable}_{occurrence} and gives assignments the same namespace as calls, so on transformers 5.17 the attention forward has both an attention_interface_0 (the assignment) and an attention_interface_1 (the call). The needle "attention_interface(" is call-shaped and hits one. A document that cannot be addressed fails at compile time, not inside someone else's process.

The reference engine #

The second engine is a plain AutoModelForCausalLM and register_forward_hook: no session, no envoys, no tracing, no remote. One model(**batch) per forward, with a hook on each module a tap names, doing the reading and the writing on the way out. It exists to test one claim — that everything else is genuinely engine-independent — and it shares engine/steps.py and ops/ unchanged.

What it costs to not be standardized is one function: loading.standardized builds an object carrying nnterp's spellings, so Address.resolve walks it unchanged. Three of the six names are published by transformers and needed no family knowledge at all — the decoder under base_model_prefix, the head under get_output_embeddings, the embeddings under get_input_embeddings. The other three are conventions the engine guesses at, and each guess raises rather than picking a first match. See Standardization for what nnterp is doing there.

What it refuses, by name:

• Every interior. A hook fires at a module boundary; attention_query never crosses one. Reaching it would mean either a patched forward per family or reimplementing, inside the engine, the head reshape and the rotary embedding the address says have already happened. Both move per-family model code back into the engine, which is the cost this project exists to measure.
• remote, in any form, including "local". A hook is a Python callable registered on a module object in this process, and there is nothing to ship. Running here silently would produce correct numbers under a false claim about where they came from.
• Running from a meta shell. nnsight's shell can run remotely; a hooks shell has nowhere to run, so the engine — not the base contract — says so.

Where a tensor sits inside a module's value is nnterp's row to say, as a Selection with get and put, and the hooks engine asks the row rather than deciding once. Unwrapping a tuple's first element is one of those (FirstIfTuple); the model's own output, whose logits field is the tensor, is another (Index("logits")). So logits — the component that on a capping family is not lm_head — reads the same bits on both engines, and the hooks engine sees Gemma-2's cap because it is reading the model's output rather than reproducing it.

One thing is harder here than under tracing. The nnterp engine's taps live and die with the with model.trace(...) block; hooks outlive the forward unless removed, and a leaked handle intervenes on the next forward — a wrong number, not an error, and the next forward is usually the one being scored. The registration loop is wrapped in try/finally, and a test asserts the module's _forward_hooks count is unchanged afterwards.

Batching windows #

batch_size bounds how many rows one model call holds. It is a property of the run and not of the experiment — it bounds memory and moves the last bit, nothing else — so it arrives with execute and is absent from the plan.

It lives in the shared walk, in one function. A window of rows is a whole small pass: the same forwards in the same order over a slice of the rows, with the published values it reads sliced the same way, so a write still meets the operand of its row. Every per-row thing a forward holds is a tuple with one entry per row — ids, mask, each op's anchors and positions — so a window is a slice of each, and positions are absolute indices into the padded width the client fixed once for all rows, so they survive the slice unchanged. An engine is handed a forward that is merely shorter and cannot tell. What the windows read is concatenated back in row order, and everything after that — metrics, eligibility, outputs, saves, a fit's loss — sees one pass.

Neither engine changed to get this, and remote="local" windows identically, because the windowing ships with the walk.

What it does change is the last bit: a GEMM over fewer rows rounds differently. The same document, run whole and at --batch-size 1:

About two parts in 10⁸, which is why run.json records the batch size beside the versions.

It does not free memory inside a fit's update: the loss is the mean over the concatenated rows, so the graph of every window is alive until the backward. A fit's memory knob is pairs. Windowing does bound the held-out pass, which runs under no_grad.

A fit is the same plan, repeatedly #

Nothing about the training loop is a second engine. A Fit holds the rows of every update it will make, already batched, already tokenized, already in the order the seed puts them in — as Observe steps, because a training update is the same pass over different rows. An update is observe over one minibatch, and the only things between updates are an optimizer step and an eval pass.

$ uv run causalab-mini explain documents/v2/das.json
<root>: Plan
  featurizers: Featurizers
      rot: subspace k=8 d=16 cayley seed=0 trained=True
  fit: Fit  10 epochs x 1 update  lr=0.001 objective=((1.0, 'ce'),) params=('rot',)  early_stop='iia' patience=3
      epochs[0][0]: Observe  metrics=['iia/logit_diff', 'ce/cross_entropy']
          forward 'original' on 'counterfactual'  (2 rows x 11 tokens)
            read  'v_cf' at block_output[0] pos={index:-1} via 'rot'
          forward 'patched' on 'base'  (2 rows x 11 tokens)
            write 'patch' at block_output[0] pos={index:-1} swap(v_cf) via 'rot'
            read  'logits' at lm_head pos={index:-1} via 'identity'
      evaluation: Observe  metrics=['iia/logit_diff', 'ce/cross_entropy']  saves=['held_out_iia.json']
          forward 'original' on 'counterfactual'  (2 rows x 9 tokens)
            read  'v_cf' at block_output[0] pos={index:-1} via 'rot'
          forward 'patched' on 'base'  (2 rows x 9 tokens)
            write 'patch' at block_output[0] pos={index:-1} swap(v_cf) via 'rot'
            read  'logits' at lm_head pos={index:-1} via 'identity'
  score: Observe  metrics=['iia/logit_diff', 'ce/cross_entropy']  saves=['iia.json', 'ce.json']
      …
  weights: Weights  names=['rot']  saves=['rot.safetensors']

epochs[0][0], evaluation and score are the same pass over different rows. The Featurizers step above them constructs the parameter sets where the run runs, and that is the whole point: the parameter it draws is the tensor the optimizer will step and the write will use. Drawn on the client it would be a different tensor from the one the block rotates by.

The optimizer is created inside the session, over [featurizers[name].weight for name in step.params]. step.params is the protocol's only trainability declaration, so the optimizer's parameter list is it — the model is frozen and nothing else in the run carries a gradient.

Per epoch: the featurizers go into training mode, each update scores the pass, builds the objective, and steps; then they go into eval mode and the held-out pass runs under no_grad on rows the fit never saw. Early stopping watches one named metric with a mode and a patience. A gate's temperature is annealed geometrically from the first update to the last.

The objective is Σ wᵢ·termᵢ, minimized. The sign of the weight is the direction — a positive weight on a cross-entropy minimizes it, a −1 on a logit difference maximizes the margin — and there is no maximize flag anywhere.

Every Observe records what it scored, including the ones inside a fit, so a run is navigable at any depth: root.steps["fit"].epochs[0][1].results["iia"]. On a long fit that is a real amount of small tensors coming home.

What the tests pin #

486 tests, and these are the ones about the seam.

The cross-engine result reproduces from the command line. The same document through both engines gives logit_diff equal in every decimal place — and these are the four numbers FINDINGS.md records from a different machine:

$ uv run causalab-mini run documents/v2/patching.json --engine nnterp --device-map cpu --out out/nnterp
$ uv run causalab-mini run documents/v2/patching.json --engine hooks  --device-map cpu --out out/hooks
$ jq -c '[.[].value]' out/nnterp/logit_diff.json out/hooks/logit_diff.json
[-0.027417995035648346,0.18085992336273193,0.03514503687620163,-0.19330616295337677]
[-0.027417995035648346,0.18085992336273193,0.03514503687620163,-0.19330616295337677]