用途
nnsight 用于解释和操控任意 PyTorch 模型的内部结构。其核心价值在于“一次编写，随处运行”：同一段解释性代码既可在本地小模型（如 GPT-2）上运行，也可通过 NDIF 远程在超大模型（70B、405B）上运行，只需切换 remote=True。
主要能力
- 支持任意 PyTorch 架构，包括 Transformers（GPT-2、LLaMA、Mistral 等）、状态空间模型（Mamba）、视觉模型（ViT、CLIP）以及自定义 nn.Module。
- 通过 trace 上下文管理器实现延迟执行，操作被收集为计算图。
- 使用代理对象（Proxy）记录对模块输出的操作，并通过 .save() 保存值供上下文外访问。
- 支持激活分析、激活修补、跨提示激活共享、多 token 生成干预以及基于梯度的分析。
- 通过 NDIF 远程执行，可在无本地 GPU 的情况下访问大模型内部。
使用方式
安装：
pip install nnsight
# vLLM 支持
pip install "nnsight[vllm]"
远程 NDIF 执行需在 login.ndif.us 注册获取 API 密钥，并设置环境变量或通过 CONFIG 配置。
基本用法示例：
from nnsight import LanguageModel
model = LanguageModel("openai-community/gpt2", device_map="auto")
with model.trace("The Eiffel Tower is in") as tracer:
hidden_states = model.transformer.h[5].output[0].save()
logits = model.output.save()
print(hidden_states.shape)
远程执行只需添加 remote=True：
with model.trace("Hello world", remote=True):
hidden = model.model.layers[40].output[0].save()
适用场景
- 在本地 GPU 无法容纳的大模型（70B、405B）上运行解释性实验。
- 对任意 PyTorch 架构进行内部访问，无需重新实现模型。
- 执行多 token 生成干预。
- 在不同提示之间共享激活。
限制
- 梯度访问不支持 vLLM 或远程执行。
- 不同模型的模块路径不同（如 GPT-2 为 model.transformer.h[i]，LLaMA 为 model.model.layers[i]），需确认模型结构。
- 在 trace 上下文外访问未调用 .save() 的值会出错。
- 远程执行可能超时，可通过 timeout 参数调整。
- 保存大量激活会占用较多内存，建议只保存需要的层。
- 若需要跨模型一致的 API，可考虑 TransformerLens；若需要声明式、可共享的干预，可考虑 pyvene；若训练 SAE，可考虑 SAELens。
参考
GitHub: ndif-team/nnsight；论文：NNsight and NDIF: Democratizing Access to Foundation Model Internals (ICLR 2025)。详细 API 文档、教程见 references/ 文件夹。