nnsight-remote-interpretability

nnsight：透明访问神经网络内部结构

nnsight (/ɛn.saɪt/) 使研究人员能够解读和操控任何 PyTorch 模型的内部结构，并具备独特能力：通过 NDIF，在小型模型上本地运行或在大型模型（70B+）上远程运行相同的代码。

GitHub：ndif-team/nnsight 论文：NNsight and NDIF: Democratizing Access to Foundation Model Internals（ICLR 2025）

核心价值

编写一次，随处运行：相同的可解释性代码既可在本地的 GPT-2 上运行，也可在远程的 Llama-3.1-405B 上运行。只需切换 remote=True。

# Local execution (small model)
with model.trace("Hello world"):
    hidden = model.transformer.h[5].output[0].save()

# Remote execution (massive model) - same code!
with model.trace("Hello world", remote=True):
    hidden = model.model.layers[40].output[0].save()

何时使用 nnsight

在以下情况中使用 nnsight：

• 在本地 GPU 无法运行的大型模型（70B、405B）上进行可解释性实验
• 使用任意 PyTorch 架构（transformers、Mamba、自定义模型）
• 执行多 token 生成干预
• 在不同提示词之间共享激活值
• 无需重新实现即可访问完整的模型内部结构

以下情况可考虑替代工具：

• 希望模型间 API 保持一致 → 使用 TransformerLens
• 需要声明式、可共享的干预 → 使用 pyvene
• 正在训练 SAE → 使用 SAELens
• 只在本地使用小型模型 → TransformerLens 可能更简单

安装

# Basic installation
pip install nnsight

# For vLLM support
pip install "nnsight[vllm]"

如需远程执行 NDIF，请在 login.ndif.us 注册并获取 API 密钥。

核心概念

LanguageModel 封装器

from nnsight import LanguageModel

# Load model (uses HuggingFace under the hood)
model = LanguageModel("openai-community/gpt2", device_map="auto")

# For larger models
model = LanguageModel("meta-llama/Llama-3.1-8B", device_map="auto")

Tracing 上下文

trace 上下文管理器支持延迟执行——操作会被收集到计算图中：

from nnsight import LanguageModel

model = LanguageModel("gpt2", device_map="auto")

with model.trace("The Eiffel Tower is in") as tracer:
    # Access any module's output
    hidden_states = model.transformer.h[5].output[0].save()

    # Access attention patterns
    attn = model.transformer.h[5].attn.attn_dropout.input[0][0].save()

    # Modify activations
    model.transformer.h[8].output[0][:] = 0  # Zero out layer 8

    # Get final output
    logits = model.output.save()

# After context exits, access saved values
print(hidden_states.shape)  # [batch, seq, hidden]

Proxy 对象

在 trace 中，模块访问会返回记录操作的 Proxy 对象：

with model.trace("Hello"):
    # These are all Proxy objects - operations are deferred
    h5_out = model.transformer.h[5].output[0]  # Proxy
    h5_mean = h5_out.mean(dim=-1)              # Proxy
    h5_saved = h5_mean.save()                   # Save for later access

工作流 1：激活值分析

分步说明

from nnsight import LanguageModel
import torch

model = LanguageModel("gpt2", device_map="auto")

prompt = "The capital of France is"

with model.trace(prompt) as tracer:
    # 1. Collect activations from multiple layers
    layer_outputs = []
    for i in range(12):  # GPT-2 has 12 layers
        layer_out = model.transformer.h[i].output[0].save()
        layer_outputs.append(layer_out)

    # 2. Get attention patterns
    attn_patterns = []
    for i in range(12):
        # Access attention weights (after softmax)
        attn = model.transformer.h[i].attn.attn_dropout.input[0][0].save()
        attn_patterns.append(attn)

    # 3. Get final logits
    logits = model.output.save()

# 4. Analyze outside context
for i, layer_out in enumerate(layer_outputs):
    print(f"Layer {i} output shape: {layer_out.shape}")
    print(f"Layer {i} norm: {layer_out.norm().item():.3f}")

# 5. Find top predictions
probs = torch.softmax(logits[0, -1], dim=-1)
top_tokens = probs.topk(5)
for token, prob in zip(top_tokens.indices, top_tokens.values):
    print(f"{model.tokenizer.decode(token)}: {prob.item():.3f}")

检查清单

• 使用 LanguageModel 封装器加载模型
• 使用 trace 上下文执行操作
• 对需要在上下文结束后使用的值调用 .save()
• 在上下文外访问已保存的值
• 使用 .shape、.norm() 等进行分析

工作流 2：激活值修补

分步说明

from nnsight import LanguageModel
import torch

model = LanguageModel("gpt2", device_map="auto")

clean_prompt = "The Eiffel Tower is in"
corrupted_prompt = "The Colosseum is in"

# 1. Get clean activations
with model.trace(clean_prompt) as tracer:
    clean_hidden = model.transformer.h[8].output[0].save()

# 2. Patch clean into corrupted run
with model.trace(corrupted_prompt) as tracer:
    # Replace layer 8 output with clean activations
    model.transformer.h[8].output[0][:] = clean_hidden

    patched_logits = model.output.save()

# 3. Compare predictions
paris_token = model.tokenizer.encode(" Paris")[0]
rome_token = model.tokenizer.encode(" Rome")[0]

patched_probs = torch.softmax(patched_logits[0, -1], dim=-1)
print(f"Paris prob: {patched_probs[paris_token].item():.3f}")
print(f"Rome prob: {patched_probs[rome_token].item():.3f}")

系统化修补扫描

def patch_layer_position(layer, position, clean_cache, corrupted_prompt):
    """Patch single layer/position from clean to corrupted."""
    with model.trace(corrupted_prompt) as tracer:
        # Get current activation
        current = model.transformer.h[layer].output[0]

        # Patch only specific position
        current[:, position, :] = clean_cache[layer][:, position, :]

        logits = model.output.save()

    return logits

# Sweep over all layers and positions
results = torch.zeros(12, seq_len)
for layer in range(12):
    for pos in range(seq_len):
        logits = patch_layer_position(layer, pos, clean_hidden, corrupted)
        results[layer, pos] = compute_metric(logits)

工作流 3：使用 NDIF 远程执行

无需本地 GPU，即可在大型模型上运行相同的实验。

分步说明

from nnsight import LanguageModel

# 1. Load large model (will run remotely)
model = LanguageModel("meta-llama/Llama-3.1-70B")

# 2. Same code, just add remote=True
with model.trace("The meaning of life is", remote=True) as tracer:
    # Access internals of 70B model!
    layer_40_out = model.model.layers[40].output[0].save()
    logits = model.output.save()

# 3. Results returned from NDIF
print(f"Layer 40 shape: {layer_40_out.shape}")

# 4. Generation with interventions
with model.trace(remote=True) as tracer:
    with tracer.invoke("What is 2+2?"):
        # Intervene during generation
        model.model.layers[20].output[0][:, -1, :] *= 1.5

    output = model.generate(max_new_tokens=50)

NDIF 设置

1. 在 login.ndif.us 注册
2. 获取 API 密钥
3. 设置环境变量或将密钥传递给 nnsight：

import os
os.environ["NDIF_API_KEY"] = "your_key"

# Or configure directly
from nnsight import CONFIG
CONFIG.API_KEY = "your_key"

NDIF 上可用的模型

• Llama-3.1-8B、70B、405B
• DeepSeek-R1 模型
• 各种开放权重模型（请访问 ndif.us 查看当前列表）

工作流 4：跨提示词共享激活值

在单个 trace 中共享不同输入之间的激活值。

from nnsight import LanguageModel

model = LanguageModel("gpt2", device_map="auto")

with model.trace() as tracer:
    # First prompt
    with tracer.invoke("The cat sat on the"):
        cat_hidden = model.transformer.h[6].output[0].save()

    # Second prompt - inject cat's activations
    with tracer.invoke("The dog ran through the"):
        # Replace with cat's activations at layer 6
        model.transformer.h[6].output[0][:] = cat_hidden
        dog_with_cat = model.output.save()

# The dog prompt now has cat's internal representations

工作流 5：基于梯度的分析

在反向传播过程中访问梯度。

from nnsight import LanguageModel
import torch

model = LanguageModel("gpt2", device_map="auto")

with model.trace("The quick brown fox") as tracer:
    # Save activations and enable gradient
    hidden = model.transformer.h[5].output[0].save()
    hidden.retain_grad()

    logits = model.output

    # Compute loss on specific token
    target_token = model.tokenizer.encode(" jumps")[0]
    loss = -logits[0, -1, target_token]

    # Backward pass
    loss.backward()

# Access gradients
grad = hidden.grad
print(f"Gradient shape: {grad.shape}")
print(f"Gradient norm: {grad.norm().item():.3f}")

注意：vLLM 或远程执行不支持访问梯度。

常见问题与解决方案

问题：不同模型的模块路径不同

# GPT-2 structure
model.transformer.h[5].output[0]

# LLaMA structure
model.model.layers[5].output[0]

# Solution: Check model structure
print(model._model)  # See actual module names

问题：忘记保存

# WRONG: Value not accessible outside trace
with model.trace("Hello"):
    hidden = model.transformer.h[5].output[0]  # Not saved!

print(hidden)  # Error or wrong value

# RIGHT: Call .save()
with model.trace("Hello"):
    hidden = model.transformer.h[5].output[0].save()

print(hidden)  # Works!

问题：远程执行超时

# For long operations, increase timeout
with model.trace("prompt", remote=True, timeout=300) as tracer:
    # Long operation...

问题：保存大量激活值时的内存问题

# Only save what you need
with model.trace("prompt"):
    # Don't save everything
    for i in range(100):
        model.transformer.h[i].output[0].save()  # Memory heavy!

    # Better: save specific layers
    key_layers = [0, 5, 11]
    for i in key_layers:
        model.transformer.h[i].output[0].save()

问题：vLLM 的梯度限制

# vLLM doesn't support gradients
# Use standard execution for gradient analysis
model = LanguageModel("gpt2", device_map="auto")  # Not vLLM

核心 API 参考

方法/属性	用途
model.trace(prompt, remote=False)	启动 tracing 上下文
proxy.save()	保存值，以便在 trace 结束后访问
proxy[:]	对 proxy 进行切片/索引（赋值会执行修补）
tracer.invoke(prompt)	在 trace 中添加提示词
model.generate(...)	使用干预生成内容
model.output	模型最终输出的 logits
model._model	底层 HuggingFace 模型

与其他工具的比较

特性	nnsight	TransformerLens	pyvene
支持任意架构	是	仅限 Transformers	是
远程执行	是（NDIF）	否	否
API 一致	否	是	是
延迟执行	是	否	否
原生支持 HuggingFace	是	重新实现	是
可共享配置	否	否	是

参考文档

有关详细的 API 文档、教程和高级用法，请参阅 references/ 文件夹：

文件	内容
references/README.md	概览和快速入门指南
references/api.md	LanguageModel、追踪和代理对象的完整 API 参考
references/tutorials.md	本地和远程可解释性的分步教程

外部资源

教程

• 入门指南
• 功能概览
• 远程执行
• 应用教程

官方文档

• 官方文档
• NDIF 信息
• 社区论坛

论文

• NNsight 和 NDIF 论文 - Fiotto-Kaufman 等人（ICLR 2025）

架构支持

nnsight 适用于任何 PyTorch 模型：

• Transformer：GPT-2、LLaMA、Mistral 等
• 状态空间模型：Mamba
• 视觉模型：ViT、CLIP
• 自定义架构：任何 nn.Module

关键是了解模块结构，以便访问正确的组件。

相关技能

• transformer-lens-interpretability：使用 TransformerLens 检查和操作 Transformer 内部结构，为机制可解释性研究提供指导，内容包括 HookP…
• sparse-autoencoder-training：使用 SAELens 训练和分析稀疏自编码器（SAE），将神经网络激活分解为可解释的内容……
• slime-rl-training：使用 slime（一个 Megatron+SGLang 框架）进行大模型 RL 后训练的指导。