NNsight and NDIF: Democratizing Access to Foundation Model Internals
interpretability;safety;large language models;distributed inference;scalable infrastructure;deferred execution;computation graphs;resource sharing
Yes
No ethics review needed.
1. Do users submit a pair of NNsight request and an input prompt? Or how is the input prompt for exercising an intervention generated?
2. When multiple NNsight requests are submitted, are all the interventions applied to a single model instance for inference, or is a separate model instance created for each request?
3. If a single model instance is instrumented with multiple NNsight requests, how does NDIF ensure that a request that modifies model parameters does not affect other requests?
4. How is KV-cache managed for multiple NNsight requests?
- This paper is focusing on an important problem since tools that enable introspection of model internals are extremely valuable for ML research and applications. 
- Opting for Pytorch-native tools is a great design choice as that significantly reduces development and integration burden for practitioners. 
- Enabling resource sharing via NDIF service is very useful to democratizing access to SOTA models.
The paper introduces two systems NNsight and NDIF that collectively aim to reduce the developer and hardware costs of analyzing and modifying the inference behavior of open-source models. NNsight is an instrumentation framework for PyTorch models, while NDIF is an inference service that enables deferred execution of NNsight instrumentations on a remote shared model deployment. The paper provides evaluation comparing to inference baselines of shared (Petal) and non-shared HPC deployments.
- The paper seems to overclaim, particularly in the title and abstract, that the work applies to foundation models in general, when in reality it only applies to open-source models, since model internals knowledge is required to create interventions. The authors should more carefully scope the claims. 
- The evaluation does not study co-tenancy scenarios where the NDIF is servicing multiple NNsight requests. This is a significant oversight considering that the resource sharing benefits of NDIF is a major contribution of the work. The authors should include results showing not just the performance implications of servicing the multiple NNsight requests but also validating the correctness of the co-tenancy features.
Yes
No ethics review needed.
Could you please share more evaluation results on the performance and scalability of the proposed systems?
* The intervention graph extends the model computational graph and decouples experimental design from model runtime, reducing engineering complexity.
* NDIF effectively shares GPU resources among researchers (co-tenancy), reducing cost and enabling large-scale experiments
* This is a novel idea with great potential benefit to the research community.
This paper introduces NNsight and NDIF, two open systems that provide efficient, transparent access to the internals of large neural networks for research purposes. NNsight extends PyTorch to offer deferred remote execution of intervention graphs, while NDIF serves as a scalable inference service that executes these intervention requests, enabling resource sharing among users. The work addresses challenges like limited access to state-of-the-art models and the significant resource demands of large-scale AI research by sharing resources.
* The evaluation section is limited to the end-to-end performance with little detail on the system optimizations. Specifically, more information is needed to assess the performance and scalability of this system using an increasing number of users.
* While NDIF addresses large-scale experiments on open models, it does not cover closed, proprietary models hosted proprietarily. 
* Lack of comparison with some co-serving systems like S-LoRA (Sheng et al, MLSys 2024) and dLoRA (Wu et al, OSDI 2024).
Yes
No ethics review needed.
It looks like the main technical contribution of this work is the enablement of efficient/virtualised remote interventions on LLMs. Could the authors elaborate on why they believe this contribution is relevant and impactful for the ICLR community?
- The proposed approach involves some new mechanisms that complement the related work.
- By enabling remote execution of the intervention code the authors have demonstrated a superior performance compared to Petals.
This work presents a new PyTorch-compatible library for remotely accessing the internal structures of deep-neural-network-based models such as large language models (LLMs) and altering the way they operate. Using this library, it is possible to produce intervention code which is hooked into the original model to read or replace original model parameters and activations. This intervention code is translated into a directed acyclic intervention graph that augments the computational graph of the original model.  Others have proposed similar, but less flexible mechanisms. For instance, pyvene (Wu et al., 2024) supports dictionary-based intervention definitions instead of the code-based interventions proposed by this work, which can also be transformed into graph-based representations and further optimised, e.g., by TorchScript. Petals (Borzunov et al., 2023), on the other hand, does not support remote execution of the intervention code, requiring it to be executed locally, which limits virtualisation opportunities and leads to additional communication overheads.
An interesting real-life application of the infrastructure built is missing. For instance, the authors could consider demonstrating their system on a large-scale interpretability study or model editing task that would be infeasible without their infrastructure.
Yes
No ethics review needed.
Please refer to the weaknesses and provide more clarifications.
1. The intervention graph architecture and the NNsight system provide a practical and efficient way to decouple experimental and engineering code, making it easier to conduct complex experiments on studying the intermediate representations and gradients of large models.
2. The paper includes a thorough performance evaluation comparing NNsight and NDIF with HPC and Petals, demonstrating the efficiency and effectiveness of the proposed solutions.
3. The performance comparisons between NNsight and other frameworks like baukit, pyvene, TransformerLens show the similar efficiency of NNsight.
4. The remote execution of NNIF provides the easy usage for users.
The paper introduces NNsight and NDIF, a framework and a cloud inference service designed to facilitate research on large-scale AI models, and enable study of the representations and computations learned by large neural networks. This work proposes a method for organizing experiments on very large models, reducing engineering burden, enhancing reproducibility, and enabling low-cost communication with remote models. The NNsight is an open-source implementation of the intervention graph architecture that extends PyTorch to support transparent model interventions without requiring local storage or management of model parameters. And the NDIF is an open-source cloud inference service that supports the NNsight by providing behind-the-scenes user sharing of model instances to reduce the costs of large-scale AI research.
1. While the developed tools of NNsight and NDIF are very useful in scientific study of AI models, there is no founded insights of this work. I'm not sure whether this contribution is enough for the acceptance.
2. Exploiting the DAG representation is not novel. As indicated by the authors, [1] has proposed the DAG concepts and use it to accelerate computation. Besides, the Torch Profiler [2] and torch.jit.trace [3] can help to dissect the model. I do not see any obvious technical challenges to implement the NNsight. Could you elaborate more on such challenges and illustrate how you address them? And how the NNsight is different from the Torch Profiler and torch.jit.trace?
3. For the practical usage like shown in Figure 4, I do not see obvious advantages of NNsight than using hooks of Torch. The required number of code lines of them are similar. Could you also elaborate more about this?

[1] Theano: A python framework for fast computation of mathematical expressions. 2016
[2] PyTorch Profiler. https://pytorch.org/tutorials/recipes/recipes/profiler_recipe.html.
[3] Torch.jit.trace. https://pytorch.org/docs/stable/generated/torch.jit.trace.html.
We create a framework for separating experimental code and model runtime, enabling researchers to run experiments on very large models without local hosting.
@inproceedings{
anonymous2024nnsight,
title={{NN}sight and {NDIF}: Democratizing Access to Foundation Model Internals},
author={Anonymous},
booktitle={Submitted to The Thirteenth International Conference on Learning Representations},
year={2024},
url={https://openreview.net/forum?id=MxbEiFRf39},
note={under review}
}
We introduce NNsight and NDIF, technologies that work in tandem to enable scientific study of the representations and computations learned by very large neural networks. NNsight is an open-source system that extends PyTorch to introduce deferred remote execution. NDIF is a scalable inference service that executes NNsight requests, allowing users to share GPU resources and pretrained models. These technologies are enabled by the intervention graph, an architecture developed to decouple experiment design from model runtime. Together, this framework provides transparent and efficient access to the internals of foundation-size deep neural networks without imposing the cost or complexity of hosting customized models individually. We conduct a quantitative survey of the machine learning literature that reveals a growing gap in the study of the internals of large-scale AI. We demonstrate the design and use of our framework to address this gap by enabling a range of research methods on huge models. Finally, we conduct benchmarks to compare performance with previous approaches.

Code and documentation will be made available open-source.
I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.
I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
interpretability safety large language models distributed inference scalable infrastructure deferred execution computation graphs resource sharing
I certify that there is no acknowledgement section in this submission for double blind review.
/pdf/643714a4cdcd62d3c9cf59398ff72ccbb34b8af6.pdf
infrastructure, software libraries, hardware, systems, etc.
I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.
I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.
/attachment/9df494a527d859b663f97deae037e2d1078e9a57.zip
NNsight and NDIF: Democratizing Access to Foundation Model Internals
ICLR 2025 Conference Submission
ICLR.cc/2025/Conference/Submission