Skip to content

v0.7.0

Latest
Latest

Choose a tag to compare

View all tags
@JadenFiotto-Kaufman JadenFiotto-Kaufman released this 05 May 05:31
· 43 commits to main since this release
v0.7.0
884bd1c

nnsight v0.7.0 Release Notes

🚀 Highlights

Lazy Hook Execution

The biggest architectural change in 0.7. Wrapped modules no longer carry permanent forward / pre-forward hooks. Instead, each Mediator registers a one-shot hook only when the worker thread actually accesses that module's .input / .output, and the hook self-removes after firing.

What this means in practice: modules nobody touches pay zero per-forward overhead. Across a representative sample of workloads we see 10–50% trace speedups on top of v0.6's gains, with the largest improvements on real models with sparse instrumentation (a few .save() sites against a 30+ layer transformer).

Three pieces make this work:

  • A sentinel forward hook on every wrapped module. PyTorch fast-paths past hook dispatch when _forward_hooks is empty; the sentinel keeps the dict non-empty so a hook added during a forward pass still fires.
  • add_ordered_hook. Inserts a hook at the right position in PyTorch's hook dict by reading mediator_idx so multi-invoke traces fire hooks in invoke-definition order regardless of which mediator registered first.
  • mediator.hooks cleanup list. Every dynamic hook (one-shot, persistent cache, iter-tracker, backward) registers there. Mediator.remove_hooks() drains the list at session cancel — idempotent, so hooks that already self-removed are no-ops.

The old SkipException machinery is gone with it. Envoy.skip(value) now stuffs __nnsight_skip__ into the forward kwargs and the wrapped forward returns it directly. As a side benefit, multi-invoke .skip(...) calls on the same module now correctly accumulate per-invoke skip values and re-concatenate them on the way out, instead of the first invoke's value clobbering the rest.

For the architecture deep-dive, see docs/developing/lazy-hook-system.md.

nnsight-serve — Single-Model HTTP Server Backed by vLLM

A lightweight FastAPI server that runs a vLLM engine and accepts serialized nnsight traces over HTTP. vLLM does the heavy lifting (continuous batching, paged KV cache, request scheduling); nnsight injects intervention hooks on top.

For users who want a persistent, intervention-capable inference endpoint without running an NDIF cluster.

Start the server:

nnsight-serve Qwen/Qwen3-30B-A3B --port 6677 --tensor-parallel-size 4 --api-key mysecret

Client UX is identical to local nnsight — just pass serve="http://host:port". The client only needs a meta model; no GPU required.

from nnsight.modeling.vllm import VLLM

model = VLLM("Qwen/Qwen3-30B-A3B")  # meta model on the client

with model.trace("The Eiffel Tower is in", serve="http://localhost:6677"):
    hidden = model.model.layers[24].output[0].save()
    logits = model.logits.save()

print(model.tokenizer.decode(logits.argmax(dim=-1)))  # Paris
print(hidden.shape)

Three modes:

  • Blocking (default). .save() values appear in scope after the with block.
  • Non-blocking. blocking=False; multiple traces fly concurrently inside vLLM's engine. Retrieve saves via tracer.collect(): 
    with model.trace("prompt 1", serve=url, blocking=False) as t1:
    out1 = model.logits.save()
with model.trace("prompt 2", serve=url, blocking=False) as t2:
    out2 = model.logits.save()

saves1 = t1.collect()  # blocks until response, returns {"out1": tensor}
saves2 = t2.collect()
  • Multi-invoke. Same as local — nest tracer.invoke(...)s inside one trace.

Optional API-key auth (--api-key on the server, api_key= on the client).

Limitations: vLLM only (no HuggingFace backend yet), no cross-trace tensor references, non-blocking mode can't inject saves into the caller frame (use .collect() to get a dict).

For the full guide and limitations, see src/nnsight/modeling/vllm/serve/README.md.

Initial implementation contributed by @khaiwang.

eproperty — A Stable Extension API for Custom Hookable Values

eproperty is the descriptor that backs every .output / .input / .inputs / .logits / .samples access. In 0.7 it becomes a first-class public API for users to add their own hookable values to a model.

A custom Envoy subclass can now expose new attributes that participate in the same request/swap protocol as the built-ins:

from nnsight import NNsight
from nnsight.intervention.envoy import Envoy
from nnsight.intervention.interleaver import eproperty
from nnsight.intervention.hooks import requires_output

class MyAttnEnvoy(Envoy):
    n_heads = 12

    @eproperty(key="output", description="Per-head attention view")
    @requires_output
    def heads(self): ...

    @heads.preprocess
    def heads(self, value):
        # Reshape so the user sees [batch, n_heads, seq, head_dim].
        B, S, H = value.shape
        return value.view(B, S, self.n_heads, H // self.n_heads).transpose(1, 2)

    @heads.transform
    @staticmethod
    def heads(value):
        # Reshape back to [batch, seq, hidden] before the model continues.
        return value.transpose(1, 2).reshape(value.shape[0], value.shape[2], -1)

with model.trace("Hello"):
    h = model.transformer.h[0].attn.heads     # [B, n_heads, S, head_dim]
    h[:, 4] = 0                                # ablate head 4

Three reshape hooks compose to give you full control over the read/write loop:

Decorator Fires on Purpose
@x.preprocess __get__ Reshape value before the user sees it.
@x.postprocess __set__ Reshape user-supplied value before swapping into the model.
@x.transform After value delivery, before next event Reshape back after in-place edits and swap into the model. Closes the loop when preprocess returned a new object.

Pre-setup decorators (requires_output, requires_input, requires_operation_output, requires_operation_input) handle hook installation. A bare @eproperty() with no setup decorator is also valid for values pushed in from outside any single nn.Module — vLLM's logits / samples work this way, as does tracer.result.

Anything that hosts an eproperty just needs to satisfy the lightweight IEnvoy protocol — an interleaver attribute and an optional path: Optional[str]. Envoy, OperationEnvoy, InterleavingTracer, and VLLM all satisfy it.

Eproperties with a description= show up in the model repr tree, and Generic[T] typing means IDEs see the right return type.

For the full extension guide, see docs/usage/extending.md and docs/developing/eproperty-deep-dive.md.

Custom Envoy Classes Per Module — envoys=

Pair eproperty with the new envoys= kwarg and you can attach per-module-type behavior to a whole model in one declaration:

import torch
from nnsight import LanguageModel

model = LanguageModel(
    "gpt2",
    envoys={torch.nn.Linear: MyLinearEnvoy, "self_attn": MyAttnEnvoy},
)

with model.trace("Hello"):
    n = model.transformer.h[0].mlp.c_fc.normalized.save()  # MyLinearEnvoy
    h = model.transformer.h[0].attn.heads                  # MyAttnEnvoy
    h[:, 4] = 0

Three forms:

Form Behavior
None (default) Every descendant wrapped in the base Envoy.
A single Envoy subclass Every descendant wrapped in that class.
Dict[type | str, Type[Envoy]] Per-module mapping.
  • Type keys match against type(module).__mro__, so {torch.nn.Linear: MyLinearEnvoy} matches every concrete Linear subclass.
  • String keys match a dotted suffix of the envoy path (component-wise, alias-aware via rename). Type keys win over string keys.
  • Subclass-level default. Model wrapper classes (LanguageModel, VLLM, your own NNsight subclass) can set envoys = {...} as a class attribute so end users don't have to pass anything.
  • Propagates down the tree. One declaration covers the whole model — each child Envoy is constructed with the same envoys= value.

Source Tracing Rewrite

module.source exposes intermediate operations inside a module's forward — every call site becomes a hookable provider path. In 0.7 the implementation is split into a global accessor and a per-Envoy wrapper, which fixes a class of correctness bugs around multiple Envoys / Interleavers / sessions touching the same module.

Layer Global (per-module) Per-Envoy wrapper
Module forward SourceAccessor SourceEnvoy
Single call site OperationAccessor OperationEnvoy

The accessors are cached on the module itself as module.__source_accessor__, so the rewrite survives torch.compile re-binding forward, accelerate's hot-swap on dispatch, and meta-tensor weight loading. Multiple Envoys wrapping the same module share the underlying accessors — only the per-Envoy wrappers are duplicated.

Other concrete improvements:

  • .source works under tracer.iter[:]. Operation iteration counters bump in lockstep with parent modules across multi-step generation.
  • Recursive .source is correct under shared accessors. Descending into a called function (...source.attention_interface_0.source.scaled_dot_product_attention_0.output) reuses the cached nested accessor instead of building fresh state per access.
  • Friendlier errors. Calling .source on a sub-module from inside another .source raises a clear ValueError directing you to access the sub-module directly. Calling .source outside a trace gives a friendly message instead of an obscure crash.
  • Cleaner repr. Print a SourceEnvoy to see the rewritten forward with operation names and line numbers; print an OperationEnvoy to see it highlighted in surrounding context.

For the full architecture, see docs/concepts/source-tracing.md and docs/developing/source-accessor-internals.md.

vLLM Compatibility Refresh

The vLLM integration got a substantial cleanup pass:

  • The == version pin is gone. pyproject.toml now declares vllm and triton without strict pins; the ImportError: nnsight requires vLLM version X failure on import nnsight.modeling.vllm is no more. nnsight now tracks current vLLM rather than locking users to one specific point release. Contributed by @gsarti.
  • vLLM 0.19+ tensor-parallel init fix. Distributed init (initialize_model_parallel) is now wrapped in set_current_vllm_config(VllmConfig()) and run outside the init_empty_weights meta context. This fixes the Cannot copy out of meta tensor failure on TP setups against vLLM 0.19+.
  • logits and samples are now epropertys, not WrapperModules. This is a small API change with a big internal cleanup payoff. (Breaking — see Migration below.) 
    # 0.6
with model.trace(prompt) as tracer:
    logits = model.logits.output.save()

# 0.7
with model.trace(prompt) as tracer:
    logits = model.logits.save()
  • generator is gone. Generation outputs flow through the result mechanism instead of a WrapperModule.
  • model.generate(...) is now an alias for model.trace(...) on VLLM, for cross-API portability with LanguageModel. max_new_tokens is rewritten to vLLM's max_tokens.
  • Tensor-parallel correctness fixes:
    • CUDA stream propagation to TP worker threads. Workers no longer write to the default stream while the driver computes on a side stream; this was a source of silent non-determinism on TP setups. Diagnosed and fixed by @khaiwang.
    • TP gather/split now uses add_ordered_hook so multiple invokes' hooks fire in mediator-defined order. (Previously the first invoke's hook would consume the whole batch.)
  • vLLM cache hooks rewritten. Per-request cache capture and async saves work correctly across all four mode/backend combinations (sync/async × multiprocessing/Ray).

engineio SSL Race-Condition Patch

Fixes a ~30–55% WebSocket connection failure rate when talking to NDIF over TLS. The patch is applied automatically at import nnsight (before any socketio imports), so you don't need to do anything.

The underlying problem: python-engineio started its read/write background threads concurrently after the WebSocket connect event, while still mid-handshake. Python's SSL sockets are not thread-safe for simultaneous read+write, and the resulting socket corruption manifested as flaky connections.

The patch serializes the handshake — flushes queued packets and receives the response synchronously in the main thread before starting background threads. Idempotent and a no-op for non-WebSocket transports. See python-socketio#1568 for upstream context.

Diagnosed and contributed by @MichaelRipa.

Frame-Based Root-Trace Detection (drop Globals.stack)

The process-wide Globals.stack counter that decided "is this trace the root or a nested one?" is gone. Whether a trace is the root (filter to saved values only) vs. an inner trace (push everything to the parent) is now determined by inspecting the target frame's locals — looking for __nnsight_tracing_info__ to mean "this frame is another tracer's compiled body."

Why it matters:

  • The counter approach broke under user-defined context-manager wrappers around model.trace() / model.session().
  • It made multi-tenant nnsight-serve workers harder — each request needed its own counter scope.
  • Frame-based detection is robust to both.

save is now mounted lazily on first .save() use, not on every trace enter.

⚠️ Heads-up: the frame-based detection is new code on a hot path. If you encounter a situation where a value you expected to be saved comes back missing, please open an issue — we'd like to track edge cases as they surface.

Async Submit + handle_response Split

For users running NDIF jobs from inside async event loops:

  • RemoteBackend.submit_request and get_response each get an async_* sibling using httpx.AsyncClient. No more thread-blocking on remote submissions in async code.
  • Submit no longer auto-dispatches handle_response. The caller decides when to run side effects, so async/streaming pipelines can yield the initial response (with the assigned job_id) before processing it.

The contract for both sync and async paths is now: submit_request returns the initial ResponseModel with job_id set; the caller invokes handle_response when ready.

Other Improvements

  • Custom context managers around model.trace() / model.session() just work. Tracer.capture(frame=...) now accepts an explicit frame, so wrappers no longer confuse the AST parser into capturing the wrong with block:

    from contextlib import contextmanager

@contextmanager
def with_logging(model, prompt):
    print(f"Tracing on {prompt!r}")
    with model.trace(prompt) as tracer:
        yield tracer
    print("done.")

with with_logging(model, "Hello") as tracer:
    out = model.transformer.h[0].output.save()

  • trace=False one-shot bypass on Envoy-bound methods. For methods auto-discovered through Envoy.__getattr__ (e.g. model.generate(...) on LanguageModel), pass trace=False to skip the implicit .trace(...) capture and call the underlying method directly:

    # No tracing — calls HF's .generate() directly.
output = model.generate(input_ids=ids, max_new_tokens=10, trace=False)

    _prepare_input is intentionally not applied in the trace=False path, since some methods don't expect prepared inputs.

  • Envoy.__setattr__ is symmetric with reads. Writes to a wrapper-claimed attribute (e.g. config) are now mirrored to both the Envoy's __dict__ (so __dict__ short-circuit reads stay coherent) and the underlying _module (so __getattr__ fall-through reads stay consistent). This had been broken since v0.6.2 — code that did model.config = new_config would silently see a stale value on read because model.__dict__ had the old one and __getattr__ was never consulted.

    import nnsight
model = nnsight.LanguageModel("openai-community/gpt2")

# Modify a field on the wrapper.
model.config.use_cache = False

# The underlying HF model now also sees the change.
assert model._module.config.use_cache is False

    Internal config writes now use self.__dict__["config"] = ... so wrapper-side bookkeeping doesn't override the underlying HF model's config.

  • Multi-invoke Envoy.skip(...) accumulates correctly. Calling .skip(value) on the same module from multiple invokes now accumulates per-invoke values and re-concatenates them on the way out, instead of the first invoke's value clobbering the rest.

  • tracer.cache(...) sub-views now scope correctly. CacheDict views built for a module path scope iteration / keys / repr / IPython pretty-print to that path's keys (and nested keys), instead of leaking the parent's full storage. Fixes display weirdness and double-prefix lookup bugs when caching multiple module subtrees.

  • DiffusionModel no longer overrides the user's device_map. Removed the device_map = "balanced" if device_map in ("auto", None) else device_map line. Whatever the caller passes propagates through unchanged, fixing Expected all tensors on cuda:0, got cuda:1 failures when the caller passes an explicit single-device map.

  • MetaMixin.dispatch() is idempotent. Calling twice no longer re-runs the load path.

  • Friendlier errors:

    • New MissedProviderError parent class; OutOfOrderError is now a subclass. Cleaner surface for "you accessed a module out of forward-pass order."
    • LanguageModel: clear error when given a multimodal config (directs to VisionLanguageModel).
    • LanguageModel: clear error for empty tokenized input.
    • Deferred exception envelope. Exceptions raised inside vLLM workers now propagate via a typed envelope so the client sees the original exception class name, not a dynamically-substituted wrapper. Contributed by @khaiwang.

  • Top-level from nnsight import save re-export.

  • Agent-evals harness (tests/agent-evals/) — the documentation benchmark we use to measure whether agents can use, guide, and develop with nnsight given the docs we ship. Adds doc-bundle benchmarking, multiple-choice questions, a Claude Code provider for Max-subscription users, browse mode, and full-bundle study with report + plots.

Documentation

Breaking Changes

  1. vLLM .logits / .samples are now epropertys.
    • Old: model.logits.output.save(), model.samples.output.save().
    • New: model.logits.save(), model.samples.save().
    • model.generator is gone.
  2. Envoy._interleaver is now Envoy.interleaver (no underscore). Subclasses or inspection code accessing the private name need to drop the underscore.
  3. Globals.stack, Globals.enter(), Globals.exit() removed. Anything calling these directly breaks. The replacement is automatic — frame-based detection.
  4. vLLM version pin lifted. Code that pinned to vllm==0.15.1 because of nnsight no longer needs to. nnsight is tested against vLLM ≥ 0.19; very old vLLM versions may break in unrelated ways.
  5. OperationEnvoy lives in intervention/source.py, not intervention/envoy.py. Public surface (.output, .input, .inputs, .source) is the same; imports need updating.
  6. SkipException is gone internally. Custom backends or runtimes that caught it should switch to inspecting kwargs for __nnsight_skip__.
  7. Envoy._fake_inputs / Envoy._fake_output removed. The fake-value bookkeeping for "model didn't execute" errors is replaced by the eproperty's own error path.
  8. Diffusion device_map no longer remapped to "balanced". If you depended on the old behavior, pass device_map="balanced" explicitly.
  9. transformers ≥ 5.0 dropped CLIPTextModel.text_model. Diffusion-related code that walked through .text_model.encoder must drop the .text_model segment.

Migration

Search-and-replace:

Old New
model.logits.output model.logits (vLLM)
model.samples.output model.samples (vLLM)
._interleaver .interleaver
text_encoder.text_model.encoder text_encoder.encoder (transformers ≥ 5)
from nnsight.intervention.envoy import OperationEnvoy from nnsight.intervention.source import OperationEnvoy

If you wrote a custom backend that caught SkipException, switch to inspecting kwargs for __nnsight_skip__. If you have a custom context manager around model.trace(...), you can stop hand-rolling frame walks — Tracer.capture(frame=...) is now part of the public API.

New Contributors

Full Changelog: v0.6.3...v0.7.0

Contributors

gsarti and khaiwang
Assets 2
Loading