nnsight – 窥探并编辑深度学习模型

它是什么
nnsight 是一个 Python 库，允许你追踪任何 PyTorch 模型（例如 Hugging Face Transformers、Diffusers、vLLM）的前向传播，并在模型运行时读取或修改中间张量。
它可以在本地运行，或者通过 NDIF 服务在你自己无法托管的大型模型上运行。

它是如何工作的
你用 nnsight 包装一个模型（例如 TransformersModel），并打开一个 with model.trace(...): 块。在块内，你编写普通的 Python 代码来访问层的 .output，使用 .save() 读取值，或覆盖它们。该库会将这些操作注入到真实的前向传播中，因此模型会使用你编辑后的张量进行计算，而无需手动添加钩子。

主要特性

• 读写内部 – 检查隐藏状态、注意力图、梯度；实时替换它们。
• 生成钩子 – 在每个生成步骤进行干预，收集 logits，或运行 logit-lens。
• 批量追踪 – 使用 tracer.invoke 在一次传递中处理多个提示。
• 梯度捕获 – 计算相对于任何内部张量的梯度。
• 远程执行 – 设置 remote=True 并在 NDIF 托管的模型上运行相同的追踪。
• 支持多种运行时 – 为 Hugging Face、Diffusers、vLLM 提供包装器，以及一个通用的 NNsight 包装器，用于任何 torch.nn.Module。

典型用法

from nnsight import TransformersModel
model = TransformersModel("openai-community/gpt2", dispatch=True)

with model.trace("The Eiffel Tower is in the city of"):
    # 将第一个 transformer 块的输出置零
    model.transformer.h[0].output[:] = 0
    # 保存网络中更深层的隐藏状态
    hidden = model.transformer.h[6].output.save()
    # 保留最终的 logits 供以后使用
    logits = model.output.logits.save()

print(hidden.shape)          # torch.Size([1, 10, 768])
print(logits.argmax(-1))     # 令牌预测

同样的模式适用于生成、梯度提取或远程运行。

谁可能使用它

• 研究人员探测模型内部，进行机制可解释性研究，或构建“logit-lens”可视化。
• 为 LLM 构建调试或编辑工具的开发者。
• 需要对基础模型执行自定义干预的代理（README 甚至列出了将 nnsight 暴露给 Claude、OpenAI Codex 等的方法）。

了解更多

• 完整文档：https://www.nnsight.net
• 设计论文（arXiv 2407.14561）和仓库中的 NNsight.md 实现指南。
• README 中链接的交互式 Colab 演练。

安装

pip install nnsight

---

所有细节均直接取自仓库的 README。