ndif-team/nnsight

The nnsight package enables interpreting and manipulating the internals of deep learned models.

nnsight – 解释与编辑 PyTorch 模型的内部结构

是什么 – nnsight 是一个可通过 pip install nnsight 安装的 Python 库，允许您查看任何 PyTorch 模型（如 GPT‑2、LLaMA）的内部状态，读取任意层的隐藏状态张量，实时修改它们，计算中间张量的梯度，甚至对模型进行永久性编辑。它可在您自己的 GPU/CPU 上本地运行，也可在 ND 研究所的远程基础设施上运行，适用于非常大的模型。

为何重要 – 现代基础模型（如 GPT‑2、LLaMA 等）是黑箱。希望理解模型为何做出特定预测、测试因果假设或原型化模型编辑技术的研究人员，需要一种无需重写模型代码即可访问和操作内部激活的清晰方法。nnsight 提供了一个高级、Python 风格的 API，抽象掉了钩子注入和追踪的样板代码。

核心功能（如 README 所述）

功能	如何使用	你将获得
激活访问	`with model.trace(prompt): hidden = model.transformer.h[5].output[0].save()`	包含给定提示下第 5 层隐藏状态的实际张量。
就地干预	在追踪块内使用 `model.transformer.h[0].output[0][:] = 0`	前向传播继续使用修改后的激活，可用于测试因果效应。
中间张量的梯度	`with loss.backward(): grad = hs.grad.save()`	任意张量（如隐藏状态）相对于您定义的损失函数的梯度。
批量调用	`with tracer.invoke(prompt): …`	并行运行多个提示；每个调用顺序执行，但通过 `.save()` 共享值。
分步控制生成	`with model.generate(prompt, max_new_tokens=5) as tracer: for step in tracer.iter[:]: …`	可在特定步骤进行干预的自回归生成。
模型编辑	`with model.edit() as edited: edited.transformer.h[0].output[0][:] = 0`	生成一个新 `LanguageModel` 实例，永久包含编辑内容，原始模型保持不变。
仅形状扫描	`with model.scan(prompt): dim = nnsight.save(layer.output.shape[-1])`	无需完整前向传播即可获取张量形状。
缓存与会话	`cache = tracer.cache()` 或 `with model.session() as s: …`	在多个追踪中重用先前捕获的激活，提高效率。
远程执行	`CONFIG.set_default_api_key(..); model = LanguageModel('meta-llama/Meta-Llama-3.1-8B'); with model.trace(..., remote=True): …`	在 NDIF 的云服务上运行追踪代码，适用于本地无法容纳的模型。
vLLM 集成	`from nnsight.modeling.vllm import VLLM; model = VLLM('gpt2', ...)`	保持相同追踪 API 的同时实现高性能推理。
任意 PyTorch 模型	`NNsight(net)` 其中 `net` 是任意 `torch.nn.Module`	相同的追踪/干预工具适用于非 Transformer 模型（如简单前馈网络）。

快速入门示例（来自 README）

from nnsight import LanguageModel

model = LanguageModel('openai-community/gpt2', device_map='auto', dispatch=True)

with model.trace('The Eiffel Tower is in the city of'):
    # 将第一层的激活置零
    model.transformer.h[0].output[0][:] = 0
    # 保存最终隐藏状态和 logits
    hidden = model.transformer.h[-1].output[0].save()
    logits = model.output.save()

print(model.tokenizer.decode(logits.logits.argmax(dim=-1)[0]))

该代码片段展示了加载模型、打开追踪、干预某一层并获取最终预测的过程。

典型用例

机制可解释性 – 检查注意力头或 MLP 块如何贡献于某个 token 的预测。
因果探测 – 将激活置零、添加噪声或替换，以测试信息流的假设。
模型编辑研究 – 创建可逆编辑（例如，“在‘埃菲尔铁塔’之后总是输出‘巴黎’”）。
调试自定义架构 – 对任意 torch.nn.Module 使用相同 API，验证前向传播行为。
高效批量实验 – 使用 invoke/session 机制在单次前向传播中运行多个提示。

限制与注意事项（如 README 所述）

执行顺序很重要 – 在追踪中必须按模型实际执行顺序访问模块；否则会触发 OutOfOrderError 死锁。
基于线程的同步 – 该库在独立工作线程中运行您的追踪代码；值仅在调用 .save() 后才可用。
无界迭代器 – tracer.iter[:] 永不返回，因此后续代码不会运行，除非放在单独的 invoke 块中。
远程执行需要 API 密钥，且模型必须在 NDIF 平台上可用。
vLLM 集成仅限 vLLM 支持的模型（目前为 Transformer 风格语言模型）。

如何进一步了解

文档网站 – https://www.nnsight.net
论文 – NNsight and NDIF: Democratizing Access to Foundation Model Internals (arXiv 2407.14561)
Discord 与论坛 – README 中链接的社区支持与讨论频道。
Colab 入门指南 – 用于动手探索的交互式笔记本。

总结（TL;DR）

nnsight 为研究人员提供了一种简洁、Python 风格的方法，用于追踪、读取、修改和 永久编辑 任何 PyTorch 模型的隐藏状态，支持批量处理、梯度提取、远程执行以及 vLLM 等高性能后端。它是现代基础模型机制研究的真正、面向研究的工具。