Open-source Interpretability NNsight NDIF

What the intervention graph separates from the runtime

NNsight is an open-source extension of PyTorch that introduces deferred remote execution. NDIF, the National Deep Inference Fabric, is the inference service that executes those requests on shared GPUs and pretrained models. The enabling object is the intervention graph, which the authors describe as an architecture developed to decouple experimental design from model runtime (Fiotto-Kaufman et al., arXiv:2407.14561v4).

The usual pattern in internals research is to insert experiment code between steps of the network: read activations, patch them, collect gradients, edit parameters, or train a small auxiliary model. Classical practice does that with per-experiment hooks or a custom network wrapper. That ties the experiment to a particular deployment, so each lab inherits weight loading, sharding, and storage. The NNsight pattern stores the experiment, then interleaves it later with a model that someone else has already loaded.

Formally, the underlying computation graph is a bipartite directed acyclic graph C = (V, A, E): variable nodes V hold objects, apply nodes A are operations, and edges run only between those two types. In NNsight, apply nodes correspond to PyTorch modules. Their definition gives each apply node one outgoing edge; an appendix argues this is equivalent, for their use, to the many-output apply nodes in the Theano-style graph of Al-Rfou et al. (2016).

An intervention is not a second copy of the network. A single user edit is a component C′ = (V′, A′, E′) plus two attachment sets. Getters, G ⊆ V × A′, feed a model variable into an intervention apply node. Setters, S ⊆ V′ × A, write an intervention variable back into a model apply node. The intervention graph is the union of those components. Interleaving them with C is the experiment. The combined graph is valid only if, for every getter edge (v_i, a′_j) and setter edge (v′_k, a_l), there is no directed path from a_l to v_i and no directed path from v′_k to a′_j. That rule is what keeps the augmented graph acyclic.

Because the edit is a graph, it can be stored as JSON, versioned, optimized, and sent to a remote process. The three requirements the authors set against hook-based hosting are separation of experiment code from engineering code, expressiveness of local PyTorch, and safe co-tenancy on one loaded model. NNsight is their answer to expressiveness: interventions are Python and PyTorch inside a tracing context. NDIF is their answer to co-tenancy.

Prerequisites

You need the NNsight client and a service that already hosts the checkpoint. The paper points to documentation and tutorials at nnsight.net, client code at github.com/ndif-team/nnsight, and service code at github.com/ndif-team/ndif. This article does not establish current install commands, API-key header formats, or the live model catalog; Figure 4 of the NDIF paper shows API-key handling in the frontend, and the later European deployment used token configuration, but neither source prints the request schema.

Remote submission requires the model’s Hugging Face meta version: parameters as placeholders that carry shape and dtype, not the full checkpoint. LanguageModel uses that configuration to build the local module tree. Provider-side access control on that download is how NDIF restricts clients to models the host is allowed to serve. If the provider will not release the configuration, the client cannot form the request. That is an authorization gate, not a step to bypass.

The tracing context is PyTorch notation. The paper states that NNsight wraps all 217 fundamental PyTorch tensor operations, wraps modules so that inputs, outputs, and gradients are visible, and can load architectures distributed through Hugging Face. The authors claim that any experiment writable with standard PyTorch hooks can be written in NNsight. That is their expressiveness claim, not an independent audit. It does not mean the shared service will expose raw hosted parameters or keep your edits in the shared checkpoint; those limits are below.

For a model the public service does not host, the supported alternative in the paper is to run the NDIF code yourself and preload one shared instance. Experiments are still graphs sent to that instance. Self-hosting removes the public catalog constraint; it does not restore the pattern of one customized model process per experiment.

Run the intervention without hosting a customized model

The listing below is the paper’s remote example, not a result re-run for this article. On meta-llama/Meta-Llama-3.1-8B, it sets three coordinates of the layer-16 MLP down-projection input at the last position and prints the next-token decode. The authors report the printout lie for the prompt “The truth is the”, as a demonstration that those three edits invert the completion in that snippet. The indices are part of the example. They are not a neuron atlas, and nothing here transfers them to another checkpoint or layer.

1. Bind an authorized, hosted checkpoint to a local meta tree. Import LanguageModel, pass a Hugging Face model id, and let NNsight wrap each submodule in an Envoy. The Envoy is the local handle for a module you do not have to instantiate as a private full-weight fork. Appendix examples also name Llama 3.1 70B and 405B ids; the clean remote listing in the paper’s Figure 3 is the 8B id below. Whether any of these are loaded on a given NDIF deployment is outside this source.
2. Open a tracing context and defer execution. with lm.trace(prompt, remote=True) both supplies inputs and selects the remote backend. Inside the block, tensor values are proxies. The paper is explicit that execution does not start until the context completes, and that the same graph can instead be executed locally if remote=True is omitted.
3. Write the PyTorch intervention on Envoy proxies. Reading .input or .output traces a getter. Assigning into a slice traces setters. The three-line body below is the paper’s equivalent of a forward pre-hook that writes 10 into mlp.input[:, -1, neurons]. The same style covers activation patching: their appendix listing copies one residual-stream vector over another inside lm.model.layers[5] for a two-prompt batch, using subject-token indices 5 and 6, and they describe that local listing as changing a “Paris” completion to “Rome”. Add remote=True only as the trace flag from Figure 3; do not assume the appendix patching cell already contains it.
4. Mark returns with .save(), and use a session when one graph is not enough. out = lm.output.save() attaches a lock so the node value is not freed when its listener count hits zero. Unsaved proxies are not materialized for the main program. Gradients use the same deferral: accessing .grad builds a protocol node that injects the backward value, and their examples call .save() on that proxy and invoke backward() inside the trace. A tensor that does not require grad will not yield a useful gradient; one tiny-model listing sets requires_grad = True inside the trace for that reason. For several forwards that must share values, a LoRA loop or a remote linear probe, open session(remote=True) and put each trace inside it. The session ships the passes as one request so intermediate values need not be downloaded and uploaded again. Their LoRA sketch creates new remote parameters and an optimizer inside the session; it is an illustration of remote parameter creation, not an in-place rewrite of the hosted checkpoint. The extracted listing calls LORA(module, dim, lora_dim) where dim is never assigned, so it is not a paste-ready trainer.
5. Exit the context and read only saved values. On exit, NNsight finalizes the graph. With remote=True, the client serializes it to the service’s custom JSON and sends the request. After the service reports completion, saved results are pulled from the object store and inserted back into the local graph. Code after the block then uses ordinary eager PyTorch, as in out["logits"][:, -1].argmax() followed by the tokenizer decode.

from nnsight import LanguageModel
model_id = "meta-llama/Meta-Llama-3.1-8B"
lm = LanguageModel(model_id)
mlp = lm.model.layers[16].mlp.down_proj
neurons = [394, 5490, 8929]
prompt = "The truth is the"
with lm.trace(prompt, remote=True):
    mlp.input[:, -1, neurons] = 10
    out = lm.output.save()
last = out["logits"][:, -1].argmax()
prediction = lm.tokenizer.decode(last)
print(prediction)

That snippet is quoted from Figure 3 of arXiv:2407.14561v4. It specifies the model id, layer, neuron indices, prompt, and remote flag. It does not specify decoding hyperparameters beyond the library default used in the figure, and this article does not re-execute it.

How the shared runtime applies the graph

Figure 4 of the same paper is the remote path. The client builds the intervention graph, serializes it, and sends it to the NDIF frontend over the WebSocket and HTTP request path. The backend keeps model instances on dedicated GPU nodes. A router hands the request to the head shard of the requested model through Ray. For a model the size of Llama 3.1 405B, weights are already split by tensor parallelism. Shard 0 forwards the request to the other shards. Every shard deserializes the full intervention graph and executes it, but each shard holds only its slice of parameters. A Torch NCCL head coordinates the distributed forward. Saved results are gathered at shard 0, written to the frontend object store, and the shard’s WebSocket client tells the local client to pull them.

Interleaving is temporary hooks, not a resident customized model. On execution, NNsight adds PyTorch hooks only on modules whose .input or .output the trace touched. Those hooks are the boundary between the intervention graph and the original computation graph. Subgraphs run when the corresponding hook fires, including more than once if the module runs more than once. After the pass, the hooks are removed. Users therefore simulate an intervention on the shared instance without permanently modifying its original computation graph. Parameters are not directly readable; the paper says users interact with a copy. That is why many experiments can share one preloaded process: the experiment arrives as data, and the process does not keep each user’s hooks or weight edits.

Sharding does not mean the intervention sees only a shard-local fragment. When the graph meets a distributed tensor, the described NDIF path records placement, gathers a full tensor, injects that full tensor into the intervention graph, then re-shards using the saved placement. Only the head shard talks to the user. The intervention can therefore be written as if the activation were whole, at the cost of a gather inside the cluster. The paper presents this as a memory-efficiency compromise, not as a claim that full-tensor gathers are free.

Black-box multi-tenant servers are a different split. The related-work section names vLLM, Hugging Face text-generation-inference, and TensorRT-LLM as systems that share hosted models but do not permit inspection or modification. LoRA-serving systems share many adapters; the authors distinguish NNsight’s finer-grained activation, gradient, and graph edits from that adapter interface. Garçon, Anthropic’s internal system, is described as remote customization by arbitrary-code hooks, which the authors say makes co-tenancy unsafe, so each user gets a private model instance. NDIF’s safety argument is the opposite: the graph is the allowed edit, and arbitrary code is not executed as the hook body on the shared process.

Co-tenancy in the evaluated implementation is sequential. Appendix B says multiple users can use the same instance without reloading parameters. The concurrent-user test, on Llama-3.1-8B hosted on a 48 GB RTX 6000 Ada, queued users and ran one forward pass each. Prompts were up to 24 tokens and each request saved one uniformly chosen layer output. From 1 to 100 concurrent users, median response time rose approximately linearly and variance increased. The paper does not print the per-point seconds in text. It explicitly says true parallel execution in one batch was future work: nodes would record batch slices, and execution would extract those slices while preserving gradients. Do not design an experiment that assumes your tokens are already fused with another user’s batch.

On cost, Section 4 compares exclusive HPC with remote NDIF on the same four NVIDIA H100 PCIe 82 GB GPUs, CUDA 12.3, and a 24-core Xeon Gold 6342. The intervention is activation patching on one batch of 32 Indirect Object Identification examples. The paper prints ± intervals without defining them in the table caption, so they are left as printed.

For OPT models from 125 million to 66 billion parameters, HPC setup time in Table 2 rises from 0.635±0.035 s to 23.697±0.079 s. NDIF setup stays at 0.149±0.008 s (125M) and 0.423±0.095 s (66B), which the authors attribute to preloading. Patching request time on NDIF is higher than local HPC at both ends: 0.541±0.123 s versus 0.021±0.001 s at 125M, and 2.739±0.154 s versus 1.889±0.040 s at 66B. Their interpretation, in that setting, is that remote execution still improves runtime for models with 3 billion parameters and more, because the loading term dominates a roughly constant communication overhead. For Llama 3.1, Table 4 reports HPC weight-load time of 5.991±0.192 s at 8B and 43.617±1.118 s at 70B, against NDIF setup of 0.518±0.005 s and 0.695±0.056 s. Table 3’s patching times remain higher remotely (0.998±0.134 s versus 0.333±0.013 s at 8B; 2.369±0.117 s versus 1.992±0.056 s at 70B). The gain they emphasize is not a faster patch; it is not paying the load on every exclusive allocation.

Against Petals, they used private instances on one RTX A6000 (49 GB), CUDA 12.5, a 12-core Ryzen 9 5900X, and about 60 MB/s between client and server. Standard remote inference, returning final hidden states, was comparable. Activation patching was not: Petals had to return intermediate hidden states for a local edit and send them back, while NDIF computed the patch and a metric such as logit difference on the server and returned the metric. The authors report that NDIF significantly outperformed Petals on that intervention task. Figure 6 uses n = 128 and says marker widths are comparable to 95% confidence intervals; the prose does not print the Petals seconds.

On exclusive HPC, without NDIF, Table 1 finds NNsight’s load and patching times comparable to baukit and pyvene on GPT2-XL, Gemma-7B, and Llama-3.1-8B, and slower to load than those libraries only in the sense that TransformerLens took about three times longer to load, which the authors attribute to weight preprocessing. Library choice on a machine you already occupy is therefore mostly API and workflow, not a large runtime gap in that table. Remote execution is the part that removes per-experiment hosting.

Checks before treating a remote result as the intervention you wrote

Confirm the graph, not just the printed token. The paper’s checks are local and structural.

Use the scan flag and validation path before a large remote pass. Scanning runs an initial pass in PyTorch FakeTensorMode and fills Envoys with fake tensors so shapes, dtypes, and device placement are visible while you trace. Validation then applies the intervention to those fake tensors as they are added, which catches shape and dtype errors before the real forward. The paper does not print the keyword spelling of the scan flag in the extracted pages; it only names “a scan flag”. FakeTensor checking is not a numerical replication of the remote model.

Check acyclicity against the getter and setter rule above. For every getter edge (v_i, a′_j) and setter edge (v′_k, a_l), there must be no directed path from a_l to v_i and no directed path from v′_k to a′_j (Fiotto-Kaufman et al., arXiv:2407.14561v4). The second condition forbids a directed path from the setter variable v′_k to the getter apply node a′_j. A graph that contains either of those directed paths is invalid. That failure is in the intervention graph, not a remote-service outage.

Confirm every tensor you will index after the block was passed through .save(). A missing save is a local graph mistake: the service can finish and still leave the main program without that activation. Conversely, saving full residual streams turns NDIF back toward the Petals cost model. The design win depends on computing what you can on the server and returning a small result.

Separate a bad graph from a bad remote checkpoint by executing the same trace locally on a model you can host, or at least by passing FakeTensor validation, before setting remote=True. Local success does not imply the remote checkpoint, shard map, or dtype matches. It does tell you the intervention graph itself was executable PyTorch.

For a causal claim, record the intervention the graph actually encodes: paper and version, model id, layer or module, tensor slice, and the value written. The 8B neuron listing and the layer-5 patching listing are causal edits in the API sense, the forward is modified and the next-token distribution is what the authors read out. They are not, by themselves, a circuit result that identifies those units under controls. The activation-patching benchmark times are runtimes, not IOI accuracy. Do not read Table 1 or Table 2 as an effect size on model behavior.

Failure cases and limits

Eager PyTorch habits fail inside the trace. Indexing a proxy as if the forward had already run, or printing an unsaved activation from inside the block and expecting a numeric tensor, fights deferred execution. Debugging that depends on real values belongs after .save(), or on a local run.

Closed-weight models are out of scope for this mechanism. The discussion states that the approach does not provide access to proprietary systems such as GPT-4 or Claude. The authors say a vendor could in principle adopt the same split while keeping parameters secret, because clients never receive the weights, but that is a feasibility remark, not a deployed product in the paper. Commercial chat APIs that return only tokens are not a substitute: the introduction’s complaint is exactly that those APIs withhold intermediate activations and gradients.

If the checkpoint is not preloaded on the service you can reach, remote=True does not conjure it. You are back to hosting, unless you deploy NDIF and preload that model once for many graphs (Fiotto-Kaufman et al., arXiv:2407.14561v4). Authorization failures look the same from the client side as a missing meta download.

Small models are a poor fit for the remote path in the authors’ own OPT measurements (Fiotto-Kaufman et al., arXiv:2407.14561v4, Table 2). At 125 million parameters, patching runtime is 0.021±0.001 s on HPC and 0.541±0.123 s on NDIF, so the patch itself is slower remotely. Setup time is 0.635±0.035 s on HPC and 0.149±0.008 s on NDIF; on the printed means, 0.635 − 0.149 = 0.486 s of setup is avoided. Use remote execution when the alternative is loading a model large enough that startup dominates, not as a default for every local hook.

Large saves and multi-GPU deployments have failed in a second, independent stack. Heithoff et al., arXiv:2508.10553v1 (14 Aug 2025), report a European NDIF-compatible service, eDIF, not a re-benchmark of the NDIF cluster in the ICLR paper. They followed the NDIF deployment guide, hosted GPT-2 small (124 million parameters, 1× RTX A6000), DeepSeek-R1-Distill-Llama-8B (8.03 billion, 2×), and DeepSeek-R1-Distill-Llama-70B (70.6 billion, 5×) on eight 48 GB RTX A6000 GPUs, and had 16 selected researchers run probing, tracing, and interventions, including activation patching and causal tracing, from 12 May to 25 June 2025. NNsight traces stayed local; inference was remote, which is the same split. They also report the operational costs that split does not erase: activation downloads sometimes lasting up to 30 minutes, intermittent execution interruptions and model freezes, crashes when models spanned GPUs and touched shared parameters, and GPUs that could not be shared across different loaded models, so unused memory on a card stayed stranded. Request volume rose above 700 in a single day near the end of the window, with most requests processed. Feedback ratings in that paper should not be generalized: the discussion says feedback participation was relatively low. Those incidents are feasibility-study observations on that deployment. They are not NDIF service-level measurements, and they are not a retraction of the intervention-graph design.

Queueing is a limit even when the graph is valid. Under the NDIF load-test implementation, concurrent users wait on separate forwards of one shared 8B instance (Fiotto-Kaufman et al., arXiv:2407.14561v4, Appendix D.2: Llama-3.1-8B on a 48 GB RTX 6000 Ada, one forward pass per queued user). A long trace, a backward pass, or a session of many steps occupies that queue. The paper does not report fairness policy, timeouts, or quotas.

Gradient and training graphs have extra failure points that the forward demo hides (Fiotto-Kaufman et al., arXiv:2407.14561v4, Appendix B.1 and Code Examples 5–8). backward() must be invoked inside the deferred graph, the saved object is .grad, and a module created with requires_grad_(False) stays silent unless the trace turns grad on. Optimizer steps in a session run on the session’s own graph, sequentially with the traces, on parameters that exist remotely. If those parameters are never returned, the local process does not hold the trained probe.

Unknowns that should stay unknown: this article uses only arXiv:2407.14561v4, dated 1 Apr 2025 in the PDF stamp, and arXiv:2508.10553v1, dated 14 Aug 2025. It does not establish the current NDIF or eDIF model list, authentication flow, queue policy, or whether the proposed parallel batch co-tenancy later shipped. It does not re-run the listings, and it does not treat the authors’ next-token demonstrations in those listings as replicated effects.

Editorial Team

Editorial Team

Publisher description: The J-Space Editorial Team curates and synthesizes cutting-edge research in mechanistic interpretability, AI alignment, and cognitive science. With a commitment to precision, empirical rigor, and clear exposition, the team distills complex technical developments into accessible yet authoritative analyses for researchers, engineers, and policymakers working at the frontier of artificial intelligence.