用途
nnsight 面向神经网络可解释性研究，允许用户以统一代码访问、检查和修改任意 PyTorch 模型的内部激活。其核心特点是通过 NDIF 远程执行，在本地无法承载的大模型（如 70B、405B）上运行同一份解释性实验代码。
主要能力
- 支持任意 PyTorch 架构，包括 Transformer 类模型（GPT-2、LLaMA、Mistral）、状态空间模型（Mamba）、视觉模型（ViT、CLIP）及自定义 nn.Module。
- 通过 LanguageModel 封装加载 HuggingFace 模型。
- trace 上下文延迟执行：trace 内的操作被记录为计算图，退出后统一执行。
- Proxy 对象记录模块访问与运算；调用 .save() 可在 trace 外部读取结果。
- 可读取或赋值修改任意模块输出、注意力权重，实现激活置零、激活修补（activation patching）等干预。
- 支持 tracer.invoke() 在同一 trace 内处理多个提示，实现跨提示激活共享。
- 支持 model.generate() 在生成过程中进行干预。
- 支持梯度分析（retain_grad 与 backward），但 vLLM 和远程执行不支持梯度。
- 远程执行通过 remote=True 切换，使用 NDIF 平台运行大模型。
使用方式
安装：
pip install nnsight
vLLM 支持：
pip install "nnsight[vllm]"
远程 NDIF 执行需在 login.ndif.us 注册获取 API key，并设置环境变量或通过 CONFIG.API_KEY 配置。
典型本地流程：
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()
model.transformer.h[8].output[0][:] = 0
logits = model.output.save()
远程执行仅在 trace 中增加 remote=True，其余代码保持一致。
适用场景
- 在本地 GPU 无法容纳的大模型上做可解释性实验。
- 对任意 PyTorch 架构进行激活分析、注意力模式提取和梯度分析。
- 多 token 生成过程中的干预实验。
- 跨提示的激活注入与比较。
- 系统性激活修补扫描。
限制
- 不具备跨模型一致的 API；不同模型需自行确认模块路径（如 GPT-2 的 transformer.h 与 LLaMA 的 model.layers）。
- 梯度访问不支持 vLLM 和远程执行。
- 远程执行可能超时，需要显式设置 timeout。
- 保存大量激活会占用较多内存，需按需保存。
- 远程执行依赖 NDIF 平台及其可用模型列表（官方提到 Llama-3.1 系列、DeepSeek-R1 等，具体以 ndif.us 为准）。
- 不支持声明式、可分享的干预配置。
与替代工具的比较
- TransformerLens：API 跨模型一致，但仅支持 Transformer 且无远程执行。
- pyvene：支持声明式可分享干预，但无远程执行与延迟执行。
官方提供 references/ 目录下的 README、api.md、tutorials.md 文档，以及 NDIF 社区论坛和 ICLR 2025 论文。