nnsight

The tests/ directory is organized by subsystem coverage:

Specialized Directories

• tests/agent-evals/: LLM agent evaluation suite (run separately from standard pytest).
• tests/performance/: Benchmark scripts (not pytest-driven; see performance.md).
• tests/mymethods/: A mock user package used for serialization testing.

nnsight operates using a deferred execution model where code blocks are captured and compiled for execution in worker threads. The core architecture relies on several key abstractions:

• Deferred Execution: Using with model.trace(...) captures a block of code, compiles it into a function, and executes it in a worker thread.
• Mediators: Each execution (invoke) is managed by a Mediator (a worker thread). Mediators communicate with the main thread using an event protocol including VALUE, SWAP, SKIP, BARRIER, END, and EXCEPTION.
• Interleaver and Hooks: nnsight uses a lazy one-shot hook architecture. This involves a sentinel forward hook, add_ordered_hook for precise placement, and mediator_idx for ordering within the Mediator.hooks lifecycle.
• Envoy and eproperty: An Envoy wraps a torch.nn.Module to expose interception interfaces via eproperty descriptors. Key properties include .output, .input, .inputs, .source, .skip, and .next. You can extend this via custom eproperties like preprocess, postprocess, and transform.
• Batching and Invokers: tracer.invoke(...) acts as a worker thread. When working with combined batches, you may encounter empty invokes and need to manage cross-invoke variable sharing using barrier().
• Source Tracing: The .source mechanism rewrites a module's forward AST to make every call site hookable, utilizing SourceAccessor, OperationAccessor, and SourceEnvoy/OperationEnvoy wrappers.

Version 0.6.0 introduced several high-level features:

• NDIF Serialization: Seamless serialization of local code via cloudpickle by-value. Use nnsight.register(...) for pip-installed packages. Local imports in your script are auto-registered. Python 3.9+ clients are compatible regardless of the NDIF server's Python version.
• vLLM Integration: Supports single/multi-GPU tensor parallelism, Ray distributed executors (single/multi-node), and mode="async" with streaming.
• Iteration: tracer.iter now supports plain for loops, which is faster than using with blocks because it avoids code capture overhead.
• NDIF Utilities: Use nnsight.compare(), nnsight.status(), and nnsight.is_model_running(...) for environment diffing and deployment checks.
• Diagnostics: Tracebacks are cleaner by default. To re-enable internal nnsight tracebacks, run with python -d.
• New Model Support: Includes VisionLanguageModel (e.g., LLaVA, Qwen2-VL) and DiffusionModel (UNet + transformer pipelines with DiffusionBatcher).

The nnsight Patterns Index serves as a cookbook of interpretability recipes. Each recipe provides a minimal working example followed by variations and interpretation tips. These patterns are categorized into looking at activations, modifying activations, comparing runs, using gradients, and accessing attention heads.

Categories of Patterns:

1. Look at activations (Observing internal states)

• logit-lens: Decode what each layer is thinking by applying the final norm and unembedding to every layer's residual.
• attention-patterns: Extract the attention probability matrix from a transformer block using the .source attribute.
• sae-and-auxiliary-modules: Wire a Sparse Autoencoder (SAE) or other auxiliary module into a model to trace through it as a first-class submodule.

2. Modify activations (Intervening on internal states)

• activation-patching: Replace activations from one run into another for causal mediation analysis or IOI-style patching.
• ablation: Zero, mean, or noise ablate specific components, positions, or features to measure changes.
• steering: Add a precomputed direction to the residual stream to push model behavior in a target direction.

3. Compare runs (Batch-wide operations)

• multi-prompt-comparison: Use multiple tracer.invoke(...) calls in a single trace or empty invokes for batch-wide operations.
• attribution-patching: Use a linear approximation of activation patching via corrupt-run gradients and activation differences.

4. Gradients (Backpropagation-based methods)

• gradient-based-attribution: Use with logits.sum().backward(): to compute saliency, integrated gradients, and per-component attribution.

5. Heads (Attention head access)

• per-head-attention: Access individual attention heads via in-trace reshaping or by using a custom Envoy with eproperty.transform.

The NNsight architecture is organized into layered subsystems that manage the lifecycle of a trace from capture to execution.

Key subsystems include:

• Tracer: Captures the body of a with block, parses it via AST, and compiles it into a callable function. Located in src/nnsight/intervention/tracing/.
• Backend: Compiles the function source to a code object, executes it, and routes results. Located in src/nnsight/intervention/backends/.
• Interleaver / Mediator: Coordinates the model's forward pass with one worker thread per invoke. It routes value/swap/skip/barrier events between threads via PyTorch hooks. Located in src/nnsight/intervention/interleaver.py.
• Hook system: Implements lazy, one-shot PyTorch hooks installed on demand by mediators. Located in src/nnsight/intervention/hooks.py.
• Source accessor: Provides AST-based forward injection for in-module operation tracing. Located in src/nnsight/intervention/source.py.
• Envoy / eproperty: The user-facing proxy and the descriptor protocol that ties the subsystems together. Located in src/nnsight/intervention/envoy.py and interleaver.py.
• Batching, serialization, runtimes: Handles multi-invoke batching, dill-based serialization for remote execution, and pluggable model runtimes (e.g., vLLM).

The vLLM integration allows NNsight interventions (observing and modifying intermediate activations) to work with models served via vLLM's high-performance inference engine.

Because vLLM uses a different architecture than standard PyTorch, the integration handles several complexities:

• Process Separation: Intervention code is serialized and transported across process boundaries to vLLM worker processes.
• Flat Tensor Format: vLLM uses a [total_tokens, hidden] format instead of [batch, tokens, hidden]. The integration manages the mapping between these formats.
• Continuous Batching: Handles requests joining and leaving the batch dynamically.
• Tensor Parallelism: Automatically gathers sharded tensors so intervention code can work with complete, unsharded tensors, and re-shards them when returning modified values.
• Phased Execution: Hooks into the forward pass, logit computation, and sampling stages independently.

nnsight uses a specialized hook architecture to move values from a running model into a worker thread. Instead of installing permanent hooks on every module (which is slow), it uses a 'lazy' approach:

1. Sentinel forward hook: Every wrapped module is given a no-op forward hook. This ensures PyTorch always uses the dispatch path that supports adding new hooks during a forward pass.
2. Lazy one-shot hooks: When you access a property like .output or .input, a single forward hook is registered. This hook fires once, delivers the value to the Mediator, and then self-removes.
3. Persistent hooks: Certain features like caches and iteration trackers use permanent hooks that fire on every forward pass.
4. Ordered insertion: When multiple mediators hook the same module, hooks fire in the order they were defined using add_ordered_hook and a mediator_idx attribute.

This architecture allows unhooked modules to have near-zero overhead, as they only carry the constant-time sentinel hook.