0.8 — the pipeline rewrite

[JadenFiotto-Kaufman]

TL;DR

0.8 is a ground-up rewrite of nnsight's execution engine. The mental model is unchanged — open with model.trace(...), then read, edit, and .save() a model's internals as ordinary Python — but the core underneath is new: intervention code and the forward pass now interleave as greenlets instead of coordinating across OS threads. The result is a smaller, simpler, easier-to-reason-about codebase with the same (and in places broader) capabilities.

The public API is largely familiar. A handful of things moved, were deprecated, or now raise where they used to no-op — the migration notes are at the bottom.

---

By the numbers (main 0.7 → 0.8)

	0.7 (main)	0.8
Core source	~18.2k LOC	~11.1k LOC (~39% smaller), 59 modules
Execution model	thread workers + locks/queues	greenlets (cooperative, single-threaded)
Diff	—	388 files, +31.8k / −61.1k
Tests	—	692 passing on CPU (vLLM suite needs a GPU)

Base branch: this targets dev (it was cut from dev, so the diff is exactly the rewrite). Heads-up: main currently has 6 commits not in dev/0.8 that should be reconciled separately before a dev → main release.

---

What changed

Execution: greenlets, not threads

Each intervention block runs in its own worker greenlet (a Mediator) that cooperatively hands control back to the model side whenever it parks on a location. Only one greenlet runs at a time, so there are no locks or queues. One shared Interleaver installs persistent pass-through forward hooks on every module; the worker/model event protocol is VALUE / SWAP / SKIP / BARRIER. Reading a location the model already ran past now raises a precise OutOfOrderError instead of deadlocking.

Model classes

• NNsight(module) — base wrapper for any torch.nn.Module.
• TransformersModel("repo/id", task=...) — the primary HuggingFace class, backed by a transformers.pipeline; supports any task (text, vision, multimodal, audio), pre-loaded module instances, and rename= / envoys= customization.
• DiffusionModel — any diffusers pipeline, UNet- or transformer-based (SD/SDXL, Flux/SD3/DiT), with seed=, per-invoke batching (DiffusionBatcher), and denoising-step iteration.
• VLLM(..., mode="sync"|"async") — vLLM-backed runtime with interventions inside the engine worker (GPU; excluded from CPU CI).
• LanguageModel / VisionLanguageModel — now deprecated thin subclasses that warn on construction; use TransformersModel(task=...).

generate vs pipe

model.generate(...) generates through the model and returns token ids on tracer.result, greedy by default (0.7 read the ids off model.generator.output). model.pipe(...) runs the whole task pipeline and returns its records (decoded text, labels). model.trace(...) runs one forward; model.scan(...) runs one forward under fake tensors for shape inference (no dispatch, no weights).

eproperty — the hookable-value descriptor

.input / .inputs / .output, tracer.result, and VLLM.logits / .samples are all built on the reintroduced @eproperty descriptor (preprocess / .postprocess / .transform / .provide, with description= surfacing a value in the Envoy repr tree). You can declare your own hookable values on a model subclass.

save is guarded

nnsight.save(x) / x.save() now raises if called outside a trace (it used to silently no-op). Idiom for collecting per-step values: save the container once, append raw values (xs = nnsight.save([]); xs.append(...)) — appending x.save() drops values on remote.

Iteration

Target occurrences across a repeated run with for step in tracer.iter[...]: or tracer.all(). Use a bounded iter[:N] so code after the loop still runs. The with tracer.iter[...]: form is deprecated; tracer.next() is gone.

Source tracing, cache, scan, skip, stop

module.source.<op> hooks intermediate operations inside a forward (AST-rewritten, lazy, inert outside a trace, recursively drillable). tracer.cache(...) records many modules at once; module.skip(replacement) bypasses a module; tracer.stop() exits early; model.scan(...) infers shapes under fake tensors.

Batching / invokers

tracer.invoke(prompt) batches multiple prompts into one forward with per-invoke row narrowing (left-pad position_ids, empty-invoke whole-batch, multi-invoke .skip() reassembly, tracer.barrier(n) for cross-invoke sharing).

Remote (NDIF)

• remote=True → RemoteBackend (blocking over one websocket, or blocking=False submit/poll).
• remote="local" → in-process serialize/deserialize/execute dry run — catch serialization issues offline, no key.
• AsyncRemoteBackend — await a job for its saves, async for streamed status.
• model.session(remote=True) bundles multiple traces into one job. Model identity via to_model_key() / from_model_key().

Config

CONFIG.API.HOST / APIKEY / COMPRESS, CONFIG.APP.DEBUG / REMOTE_LOGGING / PYMOUNT. Loaded from ~/.config/nnsight/config.yaml over shipped defaults, then env (NDIF_API_KEY, NDIF_HOST, NNSIGHT_DEBUG); -v/--verbose flips DEBUG (full tracebacks).

Packaging

pip install nnsight is a working local + remote install: transformers, huggingface-hub, and the remote deps are core (grouped in pyproject.toml), with dev / vllm / serve extras. Metadata is populated (keywords, classifiers, URLs → ndif-team/nnsight, license), nnsight.__version__ resolves via importlib.metadata, and the C .save() mount extension builds as an optional step (soft-fails to nnsight.save(value)).

---

Docs & tooling

• NNsight.md — a human-readable design-and-implementation manual (why it exists, how it works).
• docs/ — 96 recipe-style pages (usage / concepts / patterns / errors / gotchas / remote / models / developing), routed for both humans and agents by CLAUDE.md.
• NNsight_Walkthrough.ipynb — a runnable, Colab-ready guided tour (every cell verified on CPU).
• CI — .github/workflows install .[dev] and run the CPU suite (vLLM excluded); agent support via Context7 / the nnsight skill.

---

Testing

CUDA_VISIBLE_DEVICES="" pytest tests/ --ignore=tests/vllm → 692 passed. The vLLM suite (tests/vllm/) needs a CUDA GPU + nnsight[vllm].

---

Performance

Two measurements, two stories.

1. Per-intervention overhead — the greenlet win

The rewrite swaps OS worker threads (a lock + GIL handoff on every park/resume) for greenlets (a cheap in-process switch). Isolated on a compute-free deep stack of Linear(8, 8) layers — so wall-time ≈ the interleaving machinery, not the model — 0.8 cuts the cost of touching a module by 3.5–5×, and (unlike 0.7) keeps it flat as the intervention count grows:

Modules touched	Op	0.7 overhead	0.8 overhead	Speedup
50	read + save	3.0 ms	0.7 ms	4.3×
200	read + save	13.7 ms	2.9 ms	4.7×
500	read + save	43.8 ms	8.2 ms	5.3×
500	edit	68.5 ms	18.8 ms	3.6×

Per touched module: 0.7 ≈ 60–137 µs and rising with depth (thread contention); 0.8 ≈ 14–38 µs and flat. The more you inspect at once — cache()-everything, per-head/source-op tracing, deep models, long generations with per-step edits, big batches — the further 0.8 pulls ahead.

2. End-to-end on a real model — a wash, because compute dominates

On GPT-2 (CPU, torch.set_num_threads(4), 3 passes × 80 runs, 40 for generate), a sparse trace is within a couple percent of 0.7: one forward (~37 ms) dwarfs a handful of interventions' worth of microseconds. Model load/dispatch is meaningfully faster.

Workload	0.7 (main)	0.8	Δ
Raw forward (no nnsight)	37.2 ms	37.3 ms	+0.4%
trace — read + save (1 layer)	39.6 ms	40.7 ms	+2.8%
trace — edit + save (1 layer)	39.7 ms	40.4 ms	+1.6%
generate — 10 tokens	302 ms	310 ms	+2.4%
Model load + dispatch	3.45 s	2.66 s	−23%

Bottom line: 0.8 makes nnsight's own overhead ~3.5–5× cheaper and flat-scaling. For a couple of interventions on a large model you won't notice (the forward is the bottleneck); for intervention-heavy work the greenlet engine is a clear win — and per-trace runtime never regresses.

_{Repro: gpt2 + tiny nn.Sequential stacks, CUDA_VISIBLE_DEVICES="", 4 threads; identical logical workload per version (0.7 tuple .output[0] / m.generator.output, 0.8 bare .output / tracer.result).}

---

Known follow-ups (not in this PR)

Tracked gaps carried from the rewrite (vLLM, multi-invoke batching, eproperty providers, async remote, and diffusion seed/batching are in):

• Streaming/hybrid remote (tracer.local() — run a function back on the client) is not ported.
• Serialization edges: frame objects don't serialize; sourceless functions fall back to cloudpickle silently (no __source__ escape hatch); file-defined remote functions using nonlocal fail at load.
• Multimodal generate batching raises NotImplementedError (single multimodal generate works).
• Dropped extensibility/ergonomics: .pyi stubs, automodel=, arbitrary-method trace wrapping.
• Security (greenfield, absent in 0.7 too): no deserialization allow/deny layer — a deliberate hardening decision for remote servers.

---

Migration quick reference

The changes most likely to touch existing code:

0.7	0.8
LanguageModel(...) / VisionLanguageModel(...)	TransformersModel(repo, task=...) — old names still work, but warn
model.generator.output (generated ids)	tracer.result (generated ids); model.pipe(...) for decoded pipeline records
x.save() / nnsight.save(x) outside a trace (silent no-op)	raises — call it inside the trace
tracer.next() / with tracer.iter[...]:	for step in tracer.iter[...]: / tracer.all()
nnsight.apply/log/cond/iter/session(...), nnsight.list/dict/int/... wrappers	plain Python + model.session()
tracer.barrier(n_participants)	tracer.barrier(n)

See docs/reference/version-history.md for the full mapping.

🤖 Generated with Claude Code

[AdamBelfki3]

🙂

[khaiwang]

0.8 review

I ran the 0.8 branch on real GPUs and against an interpretability-workload
benchmark (56 cells, HF and vLLM-sync backends, semantics compared against an
HF control). Everything below was reproduced on this machine: py3.11, torch
2.11, transformers 5.12.1 for local runs; vllm 0.19.1 with transformers 4.57
for engine runs; 2×A100 for the GPU suites. Repro commands are in each item.

Fixes

Worker-side VLLM(module) construction. At the PR head, every vLLM GPU
worker dies at startup: the runner wraps the already-loaded module with
VLLM(self.model), _load assumes a repo-id string, and the module lands in
EngineArgs ('GPT2LMHeadModel' object has no attribute 'lower'; repro:
tests/vllm/test_tracing.py, TP=1). The PR's CI covers the CPU suite only,
so the vLLM path was never executed on a GPU. Fix: a ready module wraps
directly, and the interleaver= kwarg threads to the Envoy the way
rename/envoys already do instead of leaking into the load path.

Barrier exception attribution. Two blocks meet at a tracer.barrier(2);
the block that resumes after release raises. The exception surfaces inside
the releasing block's greenlet, unwinds it mid-body, and gets recorded on the
releasing mediator: with defer_exceptions (the engine driver mode) the
wrong invoke's request fails while the raising invoke returns truncated work
as success. Both blocks belong to one trace, so this is a diagnosability and
partial-result bug for batched traces; other traces are unaffected. Fix: the
release loop records the exception on the mediator that raised it, keeps
releasing the remaining waiters under a deferring interleaver, and re-raises
otherwise with the raiser's traceback already stashed.

Batched saves on vLLM. The documented batching idiom (a container saved
above the invoke blocks, each invoke writing its slot) returns the container
empty. Two causes. .save() marks object ids, which do not cross the
process boundary, so a serialized mediator arrives in the worker with no
record that its captured variables were saved. And each invoke rides its own
request with its own pickled copy of the container, so the copies come back
one per request and the last one clobbers the rest. Fix: __getstate__
ships the saved names with the mediator, and the sync collect merges
same-name copies element-wise before results push. The merge half is
interim: a per-trace worker scope (design question below) removes the copies
entirely, at which point collect.py deletes whole. The names half is
needed under any design, since saved-ness must always travel by name. Same
commit: the chat-input probe hard-imported transformers.pipelines.base.Chat,
which exists only from the transformers 5 chat-pipeline refactor, so every
text trace on a 4.x environment (what vllm 0.19 pins) died on ImportError;
the probe now feature-detects and treats input as not-chat when the helpers
are absent.

Backward sessions on tensors outside autograd. with t.backward(): on a
tensor from an untracked forward fails today with
RuntimeError: cannot register a hook on a tensor that doesn't require gradient, raised from the block's first .grad read. Inside an engine
worker it is worse: the shipped block's source is registered nowhere, the
nested capture fails, and the construct silently degrades to a bare
t.backward() followed by TypeError: 'NoneType' object does not support the context manager protocol, with the bare backward having already run
outside the trace. Fix: execute refuses up front when the tensor does not
require grad (a requires-grad leaf still opens a session; autograd fires its
accumulation hook), and Mediator.__setstate__ registers the deserialized
block's source in linecache under a unique filename, which 0.7's serializer
also did and which the rewrite dropped. Known limitation of the fix as
committed: the linecache entries are keyed by id(self) and never removed,
so a long-lived worker accumulates one entry per shipped block and an id
reuse can point an old traceback at the wrong source. A weakref.finalize
plus a monotonic counter closes both.

Constructor argument routing. Hardening on top of the minimal fixes: the
two decisions every model constructor makes (ready module vs load arguments,
Envoy kwargs vs load kwargs) each move to one place. Loadable.__init__
dispatches a ready module to _wrap and load arguments to _load, so no
subclass _load can receive a module again; the kwarg split derives from
inspect.signature(Envoy.__init__), so the next Envoy parameter routes with
zero mixin edits. tests/test_construction_routing.py pins both invariants
across lazy construction and dispatch replay. Take it or leave it
independently of the fixes beneath it.

Validation: full CPU suite green on the branch (654 passed, 6 skipped),
single-rank GPU vLLM tracing green, and the benchmark's batched steering
cells match the HF control (top-1 agreement 1.00, total variation 0.000).

Design question 1: a per-trace scope on engine workers

Locally, the invoke blocks of one trace share one namespace and one set of
objects through their frame. On the engine, each invoke's mediator is
pickled separately into its request's extra_args, so every shared object
arrives as a private copy per request. One missing link, four symptoms:

• shared containers come back empty or clobbered (patched by the interim
  merge above);
• a barrier never releases: each request holds its own copy, each copy waits
  for arrivals that go to the other copies, and every request retires with
  ValueError: A barrier was never reached by every block it waits for;
• a variable bound in one invoke and read in another resolves locally
  through the shared frame and is simply absent in the sibling's worker
  scope;
• any identity-based cross-invoke pattern silently operates on copies.

0.7 answered this server-side: requests carried a trace id, the first
arrival's globals became canonical, later arrivals were grafted onto them,
and the canonical values shipped home once when the trace completed. The
0.8 equivalent would be a per-trace scope keyed by trace id, established at
deserialization, with mediator scopes pointing at it.

The scheduler is the second half of the question. Continuous batching gives
no co-residency guarantee: with capacity for one request, invoke A runs and
retires (its parked barrier killed by the dangling check) before invoke B is
ever admitted, so shared state alone cannot make a rendezvous work. Either
a trace's requests are gang-scheduled (admitted together, none retired while
a sibling owes it a rendezvous), or the engine contract restricts invokes to
the schedule-independent subset: disjoint writes during the trace, merged
reads after it. Naming that contract is the decision; the interim merge
implements the restricted subset today.

Design question 2: the gradient contract on inference engines

vLLM's forward runs under torch.inference_mode in both eras. On 0.7,
intervention code ran on its own OS threads outside the thread-local mode,
so torch itself raised on every gradient-surface touch: the 0.7 gaps report
records requires_grad_() failing with RuntimeError. Greenlets run
intervention code on the model's thread, inside the mode. That is what
makes in-place steering work on 0.8, and it also means requires_grad_()
now silently no-ops and gradient misuse fails deep inside the machinery
instead of at the door.

The backward refusal above covers one construct. The general decision: an
engine backend could declare centrally that it has no gradient surface, with
every entry point (backward sessions, requires_grad_, .grad access)
raising the same error, or it could grow a real gradient path by recomputing
the relevant forward segment outside inference mode. Until one of those is
chosen, a forward-only probe that flips requires_grad and expects
gradients later is silently wrong on 0.8 where 0.7 raised.

[JadenFiotto-Kaufman]

👍

[JadenFiotto-Kaufman]

👍

[JadenFiotto-Kaufman]

Hmm

[JadenFiotto-Kaufman]

Would be nice if we had some deduplication for this and Request.deserialize somewhere