用途
nnsight（读作 /ɛn.saɪt/）用于解释和操控任意 PyTorch 模型的内部结构。其核心价值是“写一次，随处运行”：同一份可解释性代码既可在本地小模型（如 GPT-2）上运行，也可通过 NDIF 远程在超大模型（70B、405B）上运行，只需切换 remote=True。
主要能力
- 支持任意 PyTorch 架构，包括 Transformers（GPT-2、LLaMA、Mistral 等）、状态空间模型（Mamba）、视觉模型（ViT、CLIP）以及自定义 nn.Module。
- 通过 trace 上下文延迟执行，将操作收集为计算图，可访问任意模块的输出、注意力模式与最终 logits。
- 可修改激活值（如将某层输出置零）、进行激活修补、跨提示共享激活。
- 支持多 token 生成干预，可在生成过程中修改内部状态。
- 支持基于梯度的分析，可保存激活并保留梯度，在反向传播后访问梯度。
- 通过 NDIF 远程执行，无需本地 GPU 即可访问 70B+ 模型内部。
使用方式
安装：
pip install nnsight
# vLLM 支持
pip install "nnsight[vllm]"
远程 NDIF 执行需在 login.ndif.us 注册获取 API key。
基本用法：使用 LanguageModel 包装模型（底层使用 HuggingFace），在 model.trace(...) 上下文中访问模块输出，对需要后续访问的值调用 .save()，在上下文外读取。
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()
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()
NDIF 配置：
import os
os.environ["NDIF_API_KEY"] = "your_key"
# 或
from nnsight import CONFIG
CONFIG.API_KEY = "your_key"
适用场景
- 在本地 GPU 无法容纳的大模型（70B、405B）上运行可解释性实验。
- 对任意 PyTorch 架构进行内部访问，无需重新实现模型。
- 多 token 生成干预、跨提示激活共享。
- 激活分析、激活修补、系统性修补扫描、基于梯度的分析。
限制与注意事项
- 梯度访问不支持 vLLM 或远程执行。
- 不同模型的模块路径不同（如 GPT-2 为 model.transformer.h[5].output[0]，LLaMA 为 model.model.layers[5].output[0]），需先确认模型结构。
- 忘记调用 .save() 会导致上下文外无法访问值。
- 远程长操作需增加 timeout；保存大量激活会占用较多内存，建议只保存需要的层。
- 若需要跨模型一致的 API，可考虑 TransformerLens；需要声明式、可分享的干预可考虑 pyvene；训练 SAE 可考虑 SAELens；仅本地小模型时 TransformerLens 可能更简单。
参考
GitHub: ndif-team/nnsight；论文：NNsight and NDIF: Democratizing Access to Foundation Model Internals (ICLR 2025)。详细 API、教程见项目 references/ 目录及官方文档。