# nnsight v0.6.0 Release Notes

## 🚀 Highlights

### Seamless Custom Code on NDIF

The biggest change in this release is a new serialization format that enables your local packages to work seamlessly on NDIF—even if they aren't installed on the server.

**How it works:** nnsight now serializes functions and classes *by value* (source code) rather than by reference. This means your local packages are sent along with your request and rebuilt on the server—even if they aren't installed on NDIF.

```python
from nnsight import ndif, LanguageModel
import mymodule

# Register your module for remote execution
ndif.register(mymodule)

model = LanguageModel("meta-llama/Llama-3.1-70B")

with model.trace("Hello world", remote=True):
    result = mymodule.my_analysis_function(model).save()
```

**What gets auto-registered:** Code you write yourself is sent automatically—you don't need to do anything. Specifically:

- **Your script and its local imports.** Any code defined in the file you're running, plus any `.py` files it imports from the same directory or project, are sent with the request.
- **Editable/dev-installed packages.** If you ran `pip install -e .` on a package, nnsight detects it isn't a regular site-package and sends it too.

**What needs `ndif.register()`:** Packages installed normally via `pip install` (i.e. living in `site-packages`) are assumed to already exist on the server and are *not* sent. If you have a pip-installed package that NDIF doesn't have, call `ndif.register(pkg)` to explicitly include it.

**Python version flexibility:** The new serialization format enables Python 3.9+ clients to work with NDIF regardless of the server's Python version.

A real example of where this matters: [nnterp](https://github.com/ndif-team/nnterp) is a library built on nnsight that standardizes transformer interfaces across model families. NDIF doesn't have nnterp installed, but that doesn't matter—just register it:

```python
from nnterp import StandardizedTransformer
from nnsight import ndif
import nnterp

ndif.register(nnterp)

model = StandardizedTransformer("meta-llama/Llama-3.1-70B")

with model.trace("hello", remote=True):
    layer_5_output = model.layers_output[5]
    model.layers_output[10] = layer_5_output
```

This decouples library development from server deployment. nnterp can ship new features and fixes without waiting for NDIF to update its installation. You always run the version you have locally.

---

### vLLM Integration

The vLLM integration has been significantly expanded. nnsight now supports the full range of vLLM deployment configurations — from a single GPU all the way to multi-node tensor parallelism — with the same tracing API.

**Single GPU:**

```python
from nnsight.modeling.vllm import VLLM

model = VLLM("meta-llama/Llama-3.1-8B", dispatch=True)

with model.trace("The Eiffel Tower is in the city of", temperature=0.0):
    hidden = model.model.layers[16].output[0].save()
    logits = model.logits.output.save()
```

**Multi-GPU (tensor parallelism):**

```python
model = VLLM("meta-llama/Llama-3.1-8B", tensor_parallel_size=2, dispatch=True)

with model.trace("The Eiffel Tower is in the city of", temperature=0.0):
    hidden = model.model.layers[16].output[0].save()
```

Tensor parallelism is handled transparently — intervention code always sees complete, unsharded tensors. The `VLLMBatcher` automatically gathers sharded tensors before your code runs and re-shards them afterward.

**Ray distributed executor:**

```python
model = VLLM(
    "meta-llama/Llama-3.1-8B",
    tensor_parallel_size=2,
    distributed_executor_backend="ray",
    dispatch=True,
)
```

Just pass `distributed_executor_backend="ray"` and nnsight handles the rest. Ray workers use the same intervention pipeline as local multiprocessing — mediators, batch groups, and saved values all work identically across backends.

**Multi-node with Ray:**

For multi-node inference (TP workers on different machines), set the `RAY_ADDRESS` environment variable to point at an existing Ray cluster:

```bash
export RAY_ADDRESS="head-node:6379"
```

```python
model = VLLM(
    "meta-llama/Llama-3.1-70B",
    tensor_parallel_size=8,
    distributed_executor_backend="ray",
    dispatch=True,
)

with model.trace("Hello world", temperature=0.0):
    hidden = model.model.layers[40].output[0].save()
```

nnsight automatically joins the cluster as a driver-only node (no GPUs consumed on the client machine) and places workers across available nodes. If no existing cluster is found and `RAY_ADDRESS` is not set, a fresh local Ray cluster is started instead.

**Async mode with streaming:**

Pass `mode="async"` to get an async streaming interface powered by vLLM's `AsyncLLM`. Interventions are set up in the trace block, then tokens stream back as they're generated:

```python
model = VLLM("meta-llama/Llama-3.1-8B-Instruct", dispatch=True, mode="async")

with model.trace(prompt, temperature=0.7, max_tokens=256) as tracer:
    model.model.layers[16].output[0][:] += steering_vector

async for output in tracer.backend():
    if output.outputs:
        print(output.outputs[0].text)
```

This makes it easy to build interactive applications like chat interfaces with real-time token streaming while still applying interventions on every forward pass. Async mode works with all executor backends — single GPU, multiprocessing TP, and Ray.

For the full architectural details, see the [vLLM integration README](https://github.com/ndif-team/nnsight/blob/main/src/nnsight/modeling/vllm/README.md). For a runnable Docker-based multi-node example, see [`examples/multi_node_with_ray/`](https://github.com/ndif-team/nnsight/blob/main/src/nnsight/modeling/vllm/examples/multi_node_with_ray/). For demos including an async chat interface with SAE-based steering, see [nnsight-vllm-demos](https://github.com/ndif-team/nnsight-vllm-demos).

---

### For-Loop Iteration

`tracer.iter` now supports standard Python `for` loops as an alternative to `with` blocks. The `for` loop version is faster because it avoids the source extraction, AST parsing, and compilation that the `with` block requires—the loop body runs inline in the worker thread.

```python
# Preferred: for-loop (faster, no code capture overhead)
with model.generate("Hello", max_new_tokens=5) as tracer:
    logits = list().save()
    for step in tracer.iter[:]:
        logits.append(model.lm_head.output[0][-1].argmax(dim=-1))

# Bounded slices and step index work the same way
with model.generate("Hello", max_new_tokens=5) as tracer:
    hs = list().save()
    for idx in tracer.iter[:3]:
        hs.append(model.transformer.h[idx].output[0])

# Conditionals per step
with model.generate("Hello", max_new_tokens=5) as tracer:
    for step in tracer.iter[:]:
        if step == 2:
            model.transformer.h[0].output[0][:] = 0
```

The `with` block version still works and behaves identically:

```python
# Still supported
with model.generate("Hello", max_new_tokens=5) as tracer:
    logits = list().save()
    with tracer.iter[:]:
        logits.append(model.lm_head.output[0][-1].argmax(dim=-1))
```

All iteration types are supported: unbounded (`iter[:]`), bounded slices (`iter[1:3]`), single indices (`iter[0]`), and lists (`iter[[0, 2, 4]]`).

---

### NDIF Utilities

New utilities to check NDIF status and debug your setup:

```python
from nnsight import ndif

# Compare local vs remote package versions
ndif.compare()

# See which models are currently available on NDIF
ndif.status()

# Check if a specific model is running
ndif.is_model_running("meta-llama/Llama-3.1-70B")
```

`ndif.compare()` displays a table highlighting package version discrepancies between your local setup and NDIF, with critical packages (nnsight, transformers, torch) flagged prominently. `ndif.status()` lists all models currently deployed on NDIF along with their status, and `ndif.is_model_running()` lets you check a specific model before submitting a job.

---

### Cleaner Exception Tracebacks

Exceptions inside traces now show clean, readable stack traces without nnsight internals cluttering the output:

```
Traceback (most recent call last):
  File "my_script.py", line 9, in <module>
    with model.trace("Hello world"):
  File "my_script.py", line 11, in <module>
    output = model.transformer.h[999].output.save()

IndexError: list index out of range
```

For debugging nnsight internals, run your script with the `-d` or `--debug` flag to include the full stack trace:

```bash
python -d my_script.py
```

---

### Smarter Trace vs Invoker Detection

Previously, nnsight determined whether to use invokers by checking if you passed positional arguments to `.trace()`:

```python
# Old behavior:
with model.trace("Hello world"):  # Pos arg → single invoker
    ...

with model.trace() as tracer:  # No pos args → wait for sub-invokers
    with tracer.invoke("hello"):
        ...
```

**The problem:** If you only passed keyword arguments, nnsight would incorrectly wait for sub-invokers:

```python
# This would ERROR in older versions:
with model.trace(input_ids=my_ids):
    output = model.output.save()  # "can't access .output while not interleaving"
```

**Now:** The batching logic detects whether arguments affect the batch size. Arguments like `input_ids` are recognized as "real" inputs and create an implicit invoker. Configuration-only kwargs (like `max_new_tokens`) correctly allow sub-invokers:

```python
# Now works correctly:
with model.trace(input_ids=my_ids):
    hidden = model.transformer.h[0].output[0].save()

# Still supports invokers for generation config:
with model.generate(max_new_tokens=10) as tracer:
    with tracer.invoke("hello"):
        ...
    with tracer.invoke("goodbye"):
        ...
```

---

## Other Changes

- **`python -c` support:** nnsight now works when running code via `python -c "..."`. Previously, this would fail because nnsight couldn't capture the source code. Now you can run quick experiments directly from the command line:

  ```bash
  python -c "
  from nnsight import LanguageModel
  model = LanguageModel('openai-community/gpt2')
  with model.trace('Hello world'):
      hidden = model.transformer.h[0].output[0].save()
  print(hidden.shape)
  "
  ```

- **Multiple NNsight wrappers on the same model:** You can now wrap the same PyTorch model with multiple `NNsight` or `LanguageModel` instances without breaking tracing. Previously, wrapping the same model twice would remove the first wrapper's hooks. Now hooks from multiple wrappers coexist, and each wrapper remains fully functional:

  ```python
  from nnsight import NNsight
  
  model1 = NNsight(my_pytorch_model)
  model2 = NNsight(my_pytorch_model)  # Same underlying model
  
  # Both wrappers work independently
  with model1.trace(input):
      out1 = model1.layer.output.save()
  
  with model2.trace(input):
      out2 = model2.layer.output.save()
  
  # model1 still works after model2 wrapped the same model
  with model1.trace(input):
      out3 = model1.layer.output.save()
  ```

- **VisionLanguageModel:** New `VisionLanguageModel` class for tracing and intervening on vision-language models like LLaVA, Qwen2-VL, and other HuggingFace VLMs. Extends `LanguageModel` with an `AutoProcessor` that handles both text tokenization and image preprocessing. Text-only inputs fall back to the standard tokenization path, and batching across invokes handles `pixel_values` alongside `input_ids`.

  ```python
  from nnsight import VisionLanguageModel
  from PIL import Image

  model = VisionLanguageModel(
      "llava-hf/llava-interleave-qwen-0.5b-hf",
      device_map="auto",
      dispatch=True,
  )
  img = Image.open("photo.jpg")

  with model.trace("<image>\nDescribe this image", images=[img]):
      hidden = model.model.language_model.layers[-1].output.save()

  # Generation works the same way
  with model.generate("<image>\nDescribe this image", images=[img], max_new_tokens=50):
      output = model.generator.output.save()

  # Text-only still works
  with model.trace("Hello world"):
      hidden = model.model.language_model.layers[-1].output.save()
  ```

- **DiffusionModel:** First-class support for diffusion pipelines. The new `DiffusionModel` class wraps any `diffusers.DiffusionPipeline`—UNet-based (Stable Diffusion) and transformer-based (Flux, DiT) alike—so you can trace, intervene on, and iterate over denoising steps with the same API as language models. `.trace()` defaults to a single denoising step for fast exploration; `.generate()` runs the full pipeline. With `dispatch=False`, only lightweight config files are downloaded and the model is created with meta tensors—no GPU memory used until the first trace. This also means diffusion models on NDIF are coming soon.

  Batching across multiple prompts is fully supported via invokers. The `DiffusionBatcher` automatically handles the three batch-size scenarios that arise in diffusion pipelines—plain prompts, multiple images per prompt, and classifier-free guidance (which doubles the batch)—so each invoke sees only its own slice of the activations.

  ```python
  from nnsight import DiffusionModel

  sd = DiffusionModel("stabilityai/stable-diffusion-2-1")

  # Quick single-step trace
  with sd.trace("A cat"):
      denoiser_out = sd.unet.output.save()

  # Full generation with step-by-step access
  with sd.generate("A cat", num_inference_steps=50) as tracer:
      denoiser_outputs = list().save()
      for step in tracer.iter[:]:
          denoiser_outputs.append(sd.unet.output[0].clone())

  # Works with transformer-based pipelines too
  flux = DiffusionModel("black-forest-labs/FLUX.1-schnell")
  with flux.trace("A cat"):
      transformer_out = flux.transformer.output.save()
  ```

- **Memory leak fixes:** Fixed reference loops that could cause memory leaks when wrapping models. The interleaver now uses weak references to avoid circular references, and hooks are properly cleaned up when the wrapper is garbage collected.

- **Compressed results:** Results from NDIF are now compressed with zstandard for smaller download sizes and faster transfers
---

### Performance: 2.4-3.9x Faster Traces

Trace overhead has been dramatically reduced through a series of optimizations targeting both the trace setup path and the per-hook execution path. On a 12-layer MLP benchmark (CPU), nnsight traces are **2.4-3.9x faster** than v0.5.15:


| Scenario | v0.5.15 | v0.6.0 | Speedup |
|----------|---------|---------|---------|
| Empty trace (no saves) | 1,196 us | 308 us | **3.9x** |
| 1 `.save()` call | 1,370 us | 474 us | **2.9x** |
| 12 `.save()` calls | 1,697 us | 716 us | **2.4x** |

The fixed setup cost (source extraction, AST parsing, compilation, thread creation) dropped from ~1,100 us to ~210 us. The marginal cost per `.save()` dropped from ~42 us to ~34 us.


**Where the savings come from:**

- **Always-on trace caching:** Source lines, AST nodes, and compiled code objects are now cached automatically. The first trace at a given call site pays the full compilation cost; subsequent calls skip source extraction, AST parsing, and compilation entirely. The `TRACE_CACHING` config is deprecated (caching is now always enabled).
- **Persistent pymount:** The `.save()` and `.stop()` methods are now mounted once at import and never unmounted, eliminating expensive `PyType_Modified()` calls that previously invalidated all Python type caches on every trace enter/exit.
- **Removed `torch._dynamo.disable` wrappers:** The `@torch._dynamo.disable` decorator on hook functions added unnecessary `set_eval_frame` C calls on every module forward. Since the hooks use closures, thread synchronization, and mutable state that dynamo can't compile anyway, removing the decorator is safe and saves ~4 C calls per hook invocation.
- **Batched `PyFrame_LocalsToFast`:** Cross-invoker variable sharing previously called the CPython `PyFrame_LocalsToFast` C API once per variable. Now all variables are batched into a single dict update and synced with one call.
- **Filtered globals copy:** When starting an intervention thread, nnsight previously copied the entire module `__globals__` dict. Now it only copies the global names actually referenced by the intervention's bytecode (`co_names`).

**How nnsight overhead compares to bare PyTorch hooks:**


The remaining overhead beyond raw PyTorch hooks comes from nnsight's feature set: thread-based deferred execution, source code extraction, automatic batching, cross-invoke variable sharing, and the event-driven mediator protocol. This overhead is constant regardless of model size, so it becomes negligible for real models where the forward pass dominates.


**Other performance-related changes:**
- **`.save()` outside trace raises error:** Since pymount is now persistent, calling `.save()` outside a trace context now raises a `RuntimeError` with a clear message instead of silently succeeding.
- **Cache clearing:** `Globals.cache.clear()` is available to clear all cached source, AST, and code objects if needed.

---

<details>
<summary><b>Advanced: Local Serialization Testing</b></summary>

You can test if your code will serialize correctly for NDIF without actually submitting to the remote server using `remote='local'`:

```python
from nnsight import LanguageModel

model = LanguageModel("openai-community/gpt2")

# Simulates NDIF serialization locally
with model.trace("Hello world", remote='local'):
    hidden = model.transformer.h[0].output[0].save()

print("Success! Your code will work on NDIF.")
```

This is useful for catching serialization issues before submitting expensive remote jobs.

</details>

---

## Breaking Changes

- **Removed deprecated v0.4 API.** The following items from the `nnsight` namespace have been removed. They were deprecated in v0.5.0 and issued warnings or errors on use:
  - **Functions:** `nnsight.apply()`, `nnsight.log()`, `nnsight.local()`, `nnsight.cond()`, `nnsight.iter()`, `nnsight.stop()`, `nnsight.trace()`
  - **Type wrappers:** `nnsight.list`, `nnsight.dict`, `nnsight.int`, `nnsight.float`, `nnsight.str`, `nnsight.bool`, `nnsight.tuple`, `nnsight.bytes`, `nnsight.complex`, `nnsight.bytearray` — use the standard Python builtins directly
  - **The `trace=` parameter** to `.trace()` and `.generate()` — call without a `with` context instead of passing `trace=False`
  - **`obj.stop()`** on arbitrary objects — use `tracer.stop()` instead

- **Newly deprecated (will be removed in a future version):**
  - **`model.iter`, `model.all()`, `model.next()`** — use `tracer.iter`, `tracer.all()`, `tracer.next()` instead
  - **`with tracer.iter[...]:` (with-block syntax)** — use `for step in tracer.iter[...]:` instead, which is faster and avoids code capture overhead

- If you have custom model classes that implement their own `_prepare_input` or `_batch` methods, you may need to update them to match the new signature. See the `Batchable` class in `nnsight/intervention/batching.py` for the updated interface.
