$ show ndif-team/nnsight
nnsight is a Python library for interpreting and intervening in neural network models during execution. It provides a tracing API that allows users to access and modify internal activations, compute gradients, and perform interventions like activation patching and steering. The system uses a deferred execution model with thread-based synchronization to run user intervention code alongside model forward passes. Key features include support for Hugging Face models, vLLM integration for high-performance inference, remote execution via NDIF infrastructure, and persistent model editing. The architecture centers around Envoy wrappers for modules, a tracer/interleaver system for coordinating execution, and a batching mechanism for efficient multi-input processing.
The LLM identified 5 files worth highlighting for reuse:
Implements the Envoy base class and eproperty descriptor system that enables hookable properties like .output and .input on wrapped modules. This is a clean implementation of the descriptor protocol that allows for lazy, one-shot PyTorch hooks with minimal overhead. The system is well-abstracted and could be extracted for use in other intervention/tracing systems.
Contains the Interleaver/Mediator core execution management system that handles request/swap/barrier events between model and interventions. This implements a sophisticated thread-based synchronization mechanism for coordinating model execution with user intervention code. The mediator pattern implementation is clean and could be reused in other concurrent systems.
Implements a custom cloudpickle-based serialization system with source-based function serialization for cross-version compatibility. This enables remote execution across different Python versions by serializing function source code rather than bytecode. The system includes a persistent object mechanism for efficient model sharing and demonstrates non-obvious skill in solving version compatibility issues.
Implements the batching mechanism that combines multiple invokes into single forward passes for efficiency. The Batcher/Batchable system handles slicing activations per invoke and supports model-specific batchers for non-standard layouts. This is a well-designed solution for efficient batched model execution with intervention capabilities.
Implements VLLMBatcher that handles tensor parallelism gather/scatter operations and slicing for vLLM integration. This transparently handles sharded tensors in vLLM's tensor parallelism setup, which is a non-trivial problem solved cleanly. The component could be adapted for other tensor parallelism frameworks.