NNsight and NDIF: Democratizing Access to Open-Weight Foundation Model Internals

Authors: Jaden Fried Fiotto-Kaufman, Alexander Russell Loftus, Eric Todd, Jannik Brinkmann, Koyena Pal, Dmitrii Troitskii, Michael Ripa, Adam Belfki, Can Rager, Caden Juang, Aaron Mueller, Samuel Marks, Arnab Sen Sharma, Francesca Lucchetti, Nikhil Prakash, Carla E. Brodley, Arjun Guha, Jonathan Bell, Byron C Wallace, David Bau
Primary area: infrastructure, software libraries, hardware, systems, etc.
Keywords: interpretability, safety, large language models, distributed inference, scalable infrastructure, deferred execution, computation graphs, resource sharing

Abstract:
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. The National Deep Inference Fabric (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 experimental design from model runtime. Together, this framework provides transparent and efficient access to the internals of deep neural networks such as very large language models (LLMs) 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, documentation, and tutorials are available at https://nnsight.net/.

Decision: accept (poster)
Reviewer scores: 6, 6, 6, 8

Reviews and discussion:
<thread>
creation date:
1730387251108

thread_type:
official_review

<thread_object>
creator: reviewer_jasr

<content>
summary:
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.

strengths:
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.

weaknesses:
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.

questions:
Please refer to the weaknesses and provide more clarifications.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730642671794

thread_type:
official_review

<thread_object>
creator: reviewer_pjhu

<content>
summary:
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.

strengths:
- 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.

weaknesses:
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.

questions:
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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730721545838

thread_type:
official_review

<thread_object>
creator: reviewer_rzzk

<content>
summary:
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.

strengths:
* 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.

weaknesses:
* 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 discussion with relevant co-serving systems like S-LoRA (Sheng et al, MLSys 2024) and dLoRA (Wu et al, OSDI 2024).

questions:
Could you please share more evaluation results on the performance and scalability of the proposed systems?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730742968606

thread_type:
official_review

<thread_object>
creator: reviewer_rq4u

<content>
summary:
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.

strengths:
- 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.

weaknesses:
- 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.

questions:
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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731715003025

thread_type:
public_comment

<thread_object>
creator: ~neel_nanda1

<content>
title:
Paper Thoughts

comment:
Chiming in as an uninvolved mechanistic interpretability researcher, I think that NDIF & NNsight have been highly valuable contributions to the interpretability community. I think this work is being somewhat underrated by the reviewers, and is highly worthy of acceptance to ICLR.

I consider NDIF the most important contribution. As clearly documented in the paper, there is a lack of work in academia on interpreting frontier language models - this makes sense, since hosting and running these is expensive and technically involved, and rarely makes sense for any given lab to invest in. NDIF enables these costs to be paid centrally, and to enable the work of many researchers. Most interpretability work is impossible without access to the model weights, so it wouldn't make sense for NDIF to try to cover proprietary models. But with highly capable open LLMs like LLaMA 405B, NDIF is able to significantly enable interpretability research on frontier LLMs that would otherwise not be practical. In addition to the scientific importance, as frontier LLMs become more widely used in the world, understanding them better seems vital, and certain projects can only be done on extremely capable models.

I also consider NNsight to be a significant contribution. The needs of interpretability research are fairly different from most ML work, and in my opinion good tooling is crucial. I've seen many projects made or broken by whether they had the right tooling. Speaking as the creator of TransformerLens, a similar library, I think the compiler + context manager style of NNSight introduces some valuable abstractions and enables new workflows. Even if most operations could be done by a skilled operator via PyTorch hooks, being able to do things intuitively, efficiently, and safely adds a lot of value (e.g. PyTorch hooks persist unless explicitly removed, which has introduced a *lot* of bugs into my code... A context manager elegantly solves this). Empirically, I've heard many researchers speak highly of NNsight for enabling their work.

This isn't to say that this paper doesn't also make valuable research contributions itself, by in my opinion enabling and improving further research is more important than how novel or insightful a work feels. It's much more impactful to make 20 good papers 20% better, than to write one good paper yourself. To quote the meta-review for [einops (ICLR 2022 Oral)](https://openreview.net/forum?id=oapKSVM2bcj):
> I would not have met [einops] had the paper not been submitted to ICLR, and hence I am certain it should be accepted, so more can see that we care not just about mathiness, but actually enabling progress in our field.
> The job of a conference like ICLR is to expose researchers and practitioners in machine learning to ideas and techniques that may advance their research and practice. 

Einops is, of course, of interest to a much broader community than just interpretability, but I think the broader point still applies.

(Disclosure: I was not at all involved in this paper, but I do know the authors and likely have some positive bias due to that. No one asked me to write this)

(Note also: I don't feel confident in the norms here, so feel free to ignore this comment if these kinds of thoughts are not welcome! I don't see other researchers doing this, but I figure that as ICLR has made a deliberate choice to allow public comments during the rebuttal process, they likely want this kind of public feedback from uninvolved researchers, so long as it's constructive)
</content>
</thread_object>

</thread>
<thread>
creation date:
1731983062500

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
To All Reviewers

comment:
## Paper Changes
Thank you all for your insightful feedback on the paper’s contribution and presentation. Based on your feedback, we have made the following improvements to the PDF, and we welcome further comments, questions or suggestions:

1. To clarify how multiple concurrent user experiments are processed on NDIF, we have updated both Figure 5 and Section 3.3: Safe co-tenancy to discuss with more detail the meaning of co-tenancy and how we validate the correctness of NDIF’s co-tenancy features. (Reviewer RQ4u)
2. We have updated the related work section to include a discussion of S-LoRA and dLoRA (Reviewer RzZK).
3. We have updated the title to emphasize the fact that our system is built for open-weight models (Reviewers RQ4u and RzZK)
4. We have added a new Figure 8 to Appendix B, as well as a new "NDIF Design", Appendix B.2, to clarify how the NDIF system is implemented (Reviewers RzZK, RQ4u, JASr)
5. We have added a new "Load-Testing NDIF for Concurrent Users" Appendix section D.2 to address questions about stress-testing the system in this scenario (RzZK, RQ4u)

We respond below to two questions brought up by multiple reviewers.

## Access to closed, proprietary models (Rq4u, RzZK)
Two reviewers have cited the fact that our system covers exclusively open-weight models as a weakness. A prerequisite for creating a transparent API for closed models is to develop the technology demonstrating its feasibility. NDIF is the first step in this journey: our current focus is on open-weight models because commercial vendors will want to see the technical feasibility before they would consider a commercial deployment. Lines 526-533 in the paper expand the discussion on this topic.

## Experiment on reliability and Robustness (RzZK, RQ4u):
To address the points about evaluating NDIF for robustness as users increase, we evaluated the performance and scalability of the system as more users submit experiments.

We also respond to each individual reviewer in separate responses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732649890287

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
To All Reviewers

comment:
We thank you again for your work.

To address the points about evaluating NDIF for robustness as users increase, we have added a new Section D.2 in the appendix, where we measure the speed and robustness of NDIF with respect to increases in concurrent user count. Please see the new figure 9.

We also ask any reviewers that have not responded to our rebuttal to let us know whether we have addressed your comments adequately.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper implements a system that abstracts common operations during ML experimentation on deep learning models, such as inspection of and intervention upon activations, collecting gradients, and modifying parameters. The abstraction allows experimental code to become portable and executable over more environments, especially distributed environments where maintaining custom model code is time-intensive and brittle.

Reviewers strongly agreed that the system solves an important problem for the research community. Initial reviewer concerns about the empirical validation and proof-of-real-world applicability were at least partly addressed. Reviewers agreed that, while there could be more experimental validation, the problem being tackled is very important to all ML researchers, while the technical approach is highly interesting and novel.