Skip to content

Commit 7ab8470

Browse files
JadenFiotto-KaufmanJadenFiotto-Kaufman
authored
Merge pull request #591 from ndif-team/dev
Dev
2 parents b10985f + 15624dc commit 7ab8470

File tree

9 files changed

+390
-106
lines changed

9 files changed

+390
-106
lines changed

‎NNsight.md‎

Lines changed: 62 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,10 @@ This document provides an overview of the design choices and implementation deta
3535
- [Aliasing](#45-aliasing)
3636
- [Handling Conflicts](#46-handling-conflicts)
3737
- [Dispatching and Updates](#47-dispatching-and-updates)
38+
- [Ad-hoc Module Calls](#48-ad-hoc-module-calls)
39+
- [Device Utilities](#49-device-utilities)
40+
- [Envoy Tree Navigation](#410-envoy-tree-navigation)
41+
- [Accessing the Underlying Module](#411-accessing-the-underlying-module)
3842
5. [Features](#5-features)
3943
- [Saving Values](#51-saving-values)
4044
- [Ad-hoc Module Calls](#52-ad-hoc-module-calls)
@@ -54,8 +58,28 @@ This document provides an overview of the design choices and implementation deta
5458
- [LanguageModel](#63-languagemodel)
5559
- [DiffusionModel](#64-diffusionmodel)
5660
- [vLLM](#65-vllm)
57-
7. [Debugging](#7-debugging) *(coming soon)*
58-
8. [Remote Execution](#8-remote-execution) *(coming soon)*
61+
7. [Debugging](#7-debugging)
62+
- [The Challenge](#71-the-challenge)
63+
- [Exception Reconstruction](#72-exception-reconstruction)
64+
- [Line Number Reconstruction](#73-line-number-reconstruction)
65+
- [Where Exceptions Are Caught](#74-where-exceptions-are-caught)
66+
- [DEBUG Mode](#75-debug-mode)
67+
- [What Users See](#76-what-users-see)
68+
- [Common Exceptions](#77-common-exceptions)
69+
- [Debugging Strategies](#78-debugging-strategies)
70+
8. [Remote Execution](#8-remote-execution)
71+
- [NDIF Overview](#81-ndif-overview)
72+
- [Setup](#82-setup)
73+
- [Basic Remote Execution](#83-basic-remote-execution)
74+
- [Remote Model Parameters](#84-remote-model-parameters)
75+
- [Saving Results](#85-saving-results)
76+
- [Sessions for Remote Execution](#86-sessions-for-remote-execution)
77+
- [Gradients Remotely](#87-gradients-remotely)
78+
- [Python Module Whitelist](#88-python-module-whitelist)
79+
- [Limitations](#89-limitations)
80+
- [Hybrid Execution with tracer.local()](#810-hybrid-execution-with-tracerlocal)
81+
- [Print Statements and Logging](#811-print-statements-and-logging)
82+
- [Implementation Details](#812-implementation-details)
5983
9. [Extending NNsight](#9-extending-nnsight) *(coming soon)*
6084

6185
---
@@ -884,8 +908,6 @@ You can trace into nested function calls:
884908
model.transformer.h[0].attn.source.attention_interface_0.source.some_inner_op.output
885909
```
886910

887-
**Note:** Don't use `.source` on a submodule from within another `.source`. Access the submodule directly instead.
888-
889911
---
890912

891913
### 4.4 Method Delegation and Tracing
@@ -2328,34 +2350,36 @@ vLLM integration provides high-performance inference with NNsight interventions.
23282350
```
23292351
┌─────────────────────────────────────────────────────────┐
23302352
│ User Code │
2331-
│ with model.trace("Hello") as tracer: │
2332-
│ logits = model.logits.output.save() │
2333-
└──────────────────────┬──────────────────────────────────┘
2334-
│
2335-
▼
2353+
│ with model.trace("Hello") as tracer: │
2354+
│ logits = model.logits.output.save() │
2355+
└──────┬─────────────────────────────────────────────────┘
2356+
|
2357+
▼
23362358
┌─────────────────────────────────────────────────────────┐
2337-
│ VLLM Class │
2359+
│ VLLM Class |
2360+
| - Instantiates empty 'meta' model │
23382361
│ - Creates NNsightSamplingParams with serialized │
23392362
│ Mediator attached │
2340-
│ - Sends request to vLLM engine │
2341-
└──────────────────────┬──────────────────────────────────┘
2342-
│
2343-
▼
2344-
┌─────────────────────────────────────────────────────────┐
2345-
│ NNsightGPUModelRunner │
2346-
│ - Deserializes Mediator from SamplingParams │
2347-
│ - Wraps model in NNsight │
2348-
│ - Manages batch groups (flat ↔ unflatten) │
2349-
│ - Runs interleaving at multiple phases │
2350-
└──────────────────────┬──────────────────────────────────┘
2351-
│
2352-
▼
2353-
┌─────────────────────────────────────────────────────────┐
2354-
│ finish_nnsight │
2355-
│ - Lets intervention code interact with final output │
2356-
│ - Collects saved values │
2357-
│ - Handles continuous batching cleanup │
2358-
└─────────────────────────────────────────────────────────┘
2363+
│ - Sends prompts + params to vLLM engine │
2364+
└───────┬───────────────────────────────────────▲─────────┘
2365+
│ |
2366+
| |
2367+
┌───────▼──────────────────────────────────────────────────────────────────┐
2368+
│ NNsightLLMEngine |
2369+
| | - Sends final result back to NNsightGPUModelRunner |
2370+
| | for interleaving and gathering saved values |
2371+
└───────|──────────────────────────────────────────────▲───────────────────┘
2372+
│ |
2373+
| |
2374+
┌───────▼──────────────────────────────────────────────▼──────────────────────────────────────────────────────┐
2375+
| NNsightGPUModelRunner - Pre-wrapped NNsight model │
2376+
│ - Deserializes Mediator from SamplingParams finish_nnsight() │
2377+
│ - Manages batch groups (flat ↔ unflatten) 4.) Lets intervention code interact with final output │
2378+
│ - Runs interleaving at multiple phases - Collects saved values |
2379+
| 1.) Forward Pass - Handles continuous batching cleanup |
2380+
| 2.) Logits |
2381+
| 3.) Sampling |
2382+
└─────────────────────────────────────────────────────────────────────────────────────────────────────────────┘
23592383
```
23602384

23612385
#### Usage
@@ -2561,9 +2585,9 @@ vLLM uses two types of parallel linear layers:
25612585
│ └────┬────┘ └────┬────┘ │
25622586
│ │ │ │
25632587
│ ▼ ▼ │
2564-
│ ┌─────────────────────────────────────────────────────────────┐ │
2565-
│ │ all_gather() - collect all shards │ │
2566-
│ └─────────────────────────────────────────────────────────────┘ │
2588+
│ ┌─────────────────────────────────────────────────────────────┐ │
2589+
│ │ all_gather() - collect all shards │ │
2590+
│ └─────────────────────────────────────────────────────────────┘ │
25672591
│ │ │ │
25682592
│ ▼ ▼ │
25692593
│ ┌───────────────┐ ┌───────────────┐ │
@@ -2572,14 +2596,14 @@ vLLM uses two types of parallel linear layers:
25722596
│ └───────┬───────┘ └───────┬───────┘ │
25732597
│ │ │ │
25742598
│ ▼ ▼ │
2575-
│ ┌─────────────────────────────────────────────────────────────┐ │
2576-
│ │ Intervention code runs (identical on all GPUs) │ │
2577-
│ └─────────────────────────────────────────────────────────────┘ │
2599+
│ ┌─────────────────────────────────────────────────────────────┐ │
2600+
│ │ Intervention code runs (identical on all GPUs) │ │
2601+
│ └─────────────────────────────────────────────────────────────┘ │
25782602
│ │ │ │
25792603
│ ▼ ▼ │
2580-
│ ┌─────────────────────────────────────────────────────────────┐ │
2581-
│ │ split() - re-shard for forward pass │ │
2582-
│ └─────────────────────────────────────────────────────────────┘ │
2604+
│ ┌─────────────────────────────────────────────────────────────┐ │
2605+
│ │ split() - re-shard for forward pass │ │
2606+
│ └─────────────────────────────────────────────────────────────┘ │
25832607
│ │ │ │
25842608
│ ▼ ▼ │
25852609
│ ┌─────────┐ ┌─────────┐ │

‎NNsight_Walkthrough.ipynb‎

Lines changed: 18 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -19,14 +19,23 @@
1919
"\n",
2020
"# **NNsight Walkthrough**\n",
2121
"\n",
22-
"## Interpret and Manipulate the Internals of Deep Learning Models\n",
22+
"## The API for a transparent science on black-box AI\n",
2323
"\n",
24-
"**nnsight** is a Python library that gives you full access to the internals of neural networks during inference. Whether you're running models locally or remotely via [NDIF](https://ndif.us/), nnsight lets you:\n",
24+
"In this era of large-scale deep learning, the most interesting AI models are\n",
25+
"massive black boxes that are hard to run. Ordinary commercial inference service\n",
26+
"APIs let us interact with huge models, but they do not let us access model\n",
27+
"internals.\n",
2528
"\n",
26-
"- **Access activations** at any layer during forward passes\n",
27-
"- **Modify activations** to study causal effects \n",
28-
"- **Compute gradients** with respect to intermediate values\n",
29-
"- **Batch interventions** across multiple inputs efficiently\n",
29+
"The `nnsight` library is different: it provides full access to all neural\n",
30+
"network internals. When using `nnsight` together with a remote service like the\n",
31+
"[National Deep Inference Fabric](https://www.ndif.us)\n",
32+
"(NDIF), it is possible to run complex experiments on huge open models easily\n",
33+
"with fully transparent access.\n",
34+
"\n",
35+
"\n",
36+
"Through NDIF and NNsight, our team wants to enable entire labs and independent researchers alike, as we\n",
37+
"believe a large, passionate, and collaborative community will produce the next\n",
38+
"big insights on this profoundly important field.\n",
3039
"\n",
3140
"This walkthrough will teach you nnsight from the ground up, starting with the core mental model and building to advanced features."
3241
]
@@ -1321,6 +1330,9 @@
13211330
"cell_metadata_filter": "-all",
13221331
"main_language": "python",
13231332
"notebook_metadata_filter": "-all"
1333+
},
1334+
"language_info": {
1335+
"name": "python"
13241336
}
13251337
},
13261338
"nbformat": 4,

‎README.md‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,8 @@ Interpret and manipulate the internals of deep learning models
2323

2424
Originally developed in the [NDIF team](https://ndif.us/) at Northeastern University, nnsight supports local execution on any PyTorch model and remote execution on large models via the NDIF infrastructure.
2525

26+
> 📖 For a deeper technical understanding of nnsight's internals (tracing, interleaving, the Envoy system, etc.), see **[NNsight.md](./NNsight.md)**.
27+
2628
## Installation
2729

2830
```bash
@@ -448,9 +450,11 @@ For more debugging tips, see the [documentation](https://www.nnsight.net).
448450

449451
---
450452

451-
## More Examples
453+
## More Resources
452454

453-
Find more examples and tutorials at [nnsight.net](https://www.nnsight.net)
455+
- **[Documentation](https://www.nnsight.net)** — Tutorials, guides, and API reference
456+
- **[NNsight.md](./NNsight.md)** — Deep technical documentation on nnsight
457+
- **[CLAUDE.md](./CLAUDE.md)** — Comprehensive guide for AI agents working with nnsight
454458

455459
---
456460

‎src/nnsight/config.yaml‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,11 @@
11
API:
22
APIKEY: null
3-
HOST: api.ndif.us
4-
SSL: true
3+
HOST: https://api.ndif.us
54
ZLIB: true
65
APP:
76
CACHE_DIR: ~/.cache/nnsight/
7+
CROSS_INVOKER: true
88
DEBUG: false
9-
REMOTE_LOGGING: true
109
PYMOUNT: true
11-
CROSS_INVOKER: true
10+
REMOTE_LOGGING: true
1211
TRACE_CACHING: false

0 commit comments

Comments
 (0)