名称: nnsight-远程可解释性 描述: 提供使用nnsight和可选NDIF远程执行解释和操作神经网络内部的指导。适用于需要在没有本地GPU资源的情况下对大规模模型（70B+）运行可解释性实验，或处理任何PyTorch架构时。 版本: 1.0.0 作者: Orchestra Research 许可证: MIT 标签: [nnsight, NDIF, 远程执行, 机制可解释性, 模型内部] 依赖: [nnsight>=0.5.0, torch>=2.0.0]
nnsight: 透明访问神经网络内部
nnsight (/ɛn.saɪt/) 使研究人员能够解释和操作任何PyTorch模型的内部，通过NDIF支持在本地小模型或远程大规模模型（70B+）上运行相同代码的独特能力。
GitHub: ndif-team/nnsight (730+ stars) 论文: NNsight and NDIF: Democratizing Access to Foundation Model Internals (ICLR 2025)
核心价值主张
编写一次，随处运行: 相同的可解释性代码可在本地的GPT-2或远程的Llama-3.1-405B上运行。只需切换 remote=True。
# 本地执行（小模型）
with model.trace("Hello world"):
hidden = model.transformer.h[5].output[0].save()
# 远程执行（大规模模型）- 相同代码！
with model.trace("Hello world", remote=True):
hidden = model.model.layers[40].output[0].save()
何时使用 nnsight
当您需要时使用 nnsight:
- 在本地GPU无法处理的大规模模型上运行可解释性实验（70B, 405B）
- 处理任何PyTorch架构（transformers, Mamba, 自定义模型）
- 执行多令牌生成干预
- 在不同提示之间共享激活
- 访问完整模型内部而无需重新实现
考虑替代方案时:
- 您希望跨模型一致的API → 使用 TransformerLens
- 您需要声明式、可共享的干预 → 使用 pyvene
- 您正在训练SAEs → 使用 SAELens
- 您仅本地处理小模型 → TransformerLens 可能更简单
安装
# 基本安装
pip install nnsight
# 对于 vLLM 支持
pip install "nnsight[vllm]"
对于远程NDIF执行，请在 login.ndif.us 注册获取API密钥。
核心概念
LanguageModel 包装器
from nnsight import LanguageModel
# 加载模型（内部使用HuggingFace）
model = LanguageModel("openai-community/gpt2", device_map="auto")
# 对于更大模型
model = LanguageModel("meta-llama/Llama-3.1-8B", device_map="auto")
追踪上下文
trace 上下文管理器启用延迟执行 - 操作被收集到计算图中：
from nnsight import LanguageModel
model = LanguageModel("gpt2", device_map="auto")
with model.trace("The Eiffel Tower is in") as tracer:
# 访问任何模块的输出
hidden_states = model.transformer.h[5].output[0].save()
# 访问注意力模式
attn = model.transformer.h[5].attn.attn_dropout.input[0][0].save()
# 修改激活
model.transformer.h[8].output[0][:] = 0 # 清零第8层
# 获取最终输出
logits = model.output.save()
# 上下文退出后，访问保存的值
print(hidden_states.shape) # [批次, 序列, 隐藏]
代理对象
在 trace 内部，模块访问返回代理对象，记录操作：
with model.trace("Hello"):
# 这些全是代理对象 - 操作被延迟
h5_out = model.transformer.h[5].output[0] # 代理
h5_mean = h5_out.mean(dim=-1) # 代理
h5_saved = h5_mean.save() # 保存供后续访问
工作流程 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. 从多层收集激活
layer_outputs = []
for i in range(12): # GPT-2 有12层
layer_out = model.transformer.h[i].output[0].save()
layer_outputs.append(layer_out)
# 2. 获取注意力模式
attn_patterns = []
for i in range(12):
# 访问注意力权重（softmax后）
attn = model.transformer.h[i].attn.attn_dropout.input[0][0].save()
attn_patterns.append(attn)
# 3. 获取最终logits
logits = model.output.save()
# 4. 在上下文外分析
for i, layer_out in enumerate(layer_outputs):
print(f"层 {i} 输出形状: {layer_out.shape}")
print(f"层 {i} 范数: {layer_out.norm().item():.3f}")
# 5. 查找 top 预测
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. 获取 clean 激活
with model.trace(clean_prompt) as tracer:
clean_hidden = model.transformer.h[8].output[0].save()
# 2. 将 clean 修补到 corrupted 运行中
with model.trace(corrupted_prompt) as tracer:
# 用 clean 激活替换第8层输出
model.transformer.h[8].output[0][:] = clean_hidden
patched_logits = model.output.save()
# 3. 比较预测
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 概率: {patched_probs[paris_token].item():.3f}")
print(f"Rome 概率: {patched_probs[rome_token].item():.3f}")
系统化修补扫描
def patch_layer_position(layer, position, clean_cache, corrupted_prompt):
"""将单个层/位置从 clean 修补到 corrupted."""
with model.trace(corrupted_prompt) as tracer:
# 获取当前激活
current = model.transformer.h[layer].output[0]
# 仅修补特定位置
current[:, position, :] = clean_cache[layer][:, position, :]
logits = model.output.save()
return logits
# 扫描所有层和位置
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. 加载大模型（将远程运行）
model = LanguageModel("meta-llama/Llama-3.1-70B")
# 2. 相同代码，只需添加 remote=True
with model.trace("The meaning of life is", remote=True) as tracer:
# 访问 70B 模型的内部！
layer_40_out = model.model.layers[40].output[0].save()
logits = model.output.save()
# 3. 从 NDIF 返回结果
print(f"层 40 形状: {layer_40_out.shape}")
# 4. 带干预的生成
with model.trace(remote=True) as tracer:
with tracer.invoke("What is 2+2?"):
# 在生成期间干预
model.model.layers[20].output[0][:, -1, :] *= 1.5
output = model.generate(max_new_tokens=50)
NDIF 设置
- 在 login.ndif.us 注册
- 获取 API 密钥
- 设置环境变量或传递给 nnsight：
import os
os.environ["NDIF_API_KEY"] = "your_key"
# 或直接配置
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:
# 第一个提示
with tracer.invoke("The cat sat on the"):
cat_hidden = model.transformer.h[6].output[0].save()
# 第二个提示 - 注入猫的激活
with tracer.invoke("The dog ran through the"):
# 用猫的激活替换第6层
model.transformer.h[6].output[0][:] = cat_hidden
dog_with_cat = model.output.save()
# 狗提示现在具有猫的内部表示
工作流程 5: 基于梯度的分析
在反向传播期间访问梯度。
from nnsight import LanguageModel
import torch
model = LanguageModel("gpt2", device_map="auto")
with model.trace("The quick brown fox") as tracer:
# 保存激活并启用梯度
hidden = model.transformer.h[5].output[0].save()
hidden.retain_grad()
logits = model.output
# 计算特定 token 的损失
target_token = model.tokenizer.encode(" jumps")[0]
loss = -logits[0, -1, target_token]
# 反向传播
loss.backward()
# 访问梯度
grad = hidden.grad
print(f"梯度形状: {grad.shape}")
print(f"梯度范数: {grad.norm().item():.3f}")
注意: 梯度访问不支持 vLLM 或远程执行。
常见问题与解决方案
问题: 模型路径不同
# GPT-2 结构
model.transformer.h[5].output[0]
# LLaMA 结构
model.model.layers[5].output[0]
# 解决方案: 检查模型结构
print(model._model) # 查看实际模块名称
问题: 忘记保存
# 错误: 值在 trace 外不可访问
with model.trace("Hello"):
hidden = model.transformer.h[5].output[0] # 未保存！
print(hidden) # 错误或错误值
# 正确: 调用 .save()
with model.trace("Hello"):
hidden = model.transformer.h[5].output[0].save()
print(hidden) # 有效！
问题: 远程超时
# 对于长操作，增加超时
with model.trace("prompt", remote=True, timeout=300) as tracer:
# 长操作...
问题: 保存大量激活时的内存
# 仅保存所需
with model.trace("prompt"):
# 不要保存所有
for i in range(100):
model.transformer.h[i].output[0].save() # 内存消耗高！
# 更好: 保存特定层
key_layers = [0, 5, 11]
for i in key_layers:
model.transformer.h[i].output[0].save()
问题: vLLM 梯度限制
# vLLM 不支持梯度
# 使用标准执行进行梯度分析
model = LanguageModel("gpt2", device_map="auto") # 非 vLLM
关键 API 参考
| 方法/属性 | 目的 |
|---|---|
model.trace(prompt, remote=False) |
启动追踪上下文 |
proxy.save() |
保存值供 trace 后访问 |
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 | 本地和远程可解释性的分步教程 |
外部资源
教程
官方文档
论文
- NNsight and NDIF 论文 - Fiotto-Kaufman et al. (ICLR 2025)
架构支持
nnsight 适用于任何 PyTorch 模型：
- Transformers: GPT-2, LLaMA, Mistral 等
- 状态空间模型: Mamba
- 视觉模型: ViT, CLIP
- 自定义架构: 任何 nn.Module
关键是了解模块结构以访问正确组件。