nnsight-remote-interpretability

# nnsight: 神经网络内部结构的透明访问 nnsight (/ɛn.saɪt/) 使研究人员能够解释和操作任何 PyTorch 模型的内部结构，并具备通过 NDIF 在本地运行小型模型或远程运行大型模型（70B+）的独特能力。 **GitHub**: [ndif-team/nnsight](https://github.com/ndif-team/nnsight) (730+ 星标) **论文**: [NNsight and NDIF: Democratizing Access to Foundation Model Internals](https://arxiv.org/abs/2407.14561) (ICLR 2025) ## 核心价值主张 **一次编写，随处运行**：相同的可解释性代码可以在本地运行 GPT-2，或远程运行 Llama-3.1-405B。只需切换 `remote=True`。 python # 本地执行 (小型模型) with model.trace("Hello world"): hidden = model.transformer.h[5].output[0].save() # 远程执行 (大型模型) - 代码相同！ with model.trace("Hello world", remote=True): hidden = model.model.layers[40].output[0].save() ## 何时使用 nnsight **在以下情况使用 nnsight：** - 需要在本地 GPU 无法承载的大型模型（70B, 405B）上运行可解释性实验 - 使用任何 PyTorch 架构（Transformer、Mamba、自定义模型） - 执行多 Token 生成干预 - 在不同提示词之间共享激活值 - 无需重新实现即可访问完整的模型内部结构 **在以下情况考虑替代方案：** - 需要跨模型的一致 API → 使用 **TransformerLens** - 需要声明式、可共享的干预 → 使用 **pyvene** - 正在训练 SAE → 使用 **SAELens** - 仅在本地处理小型模型 → **TransformerLens** 可能更简单 ## 安装 bash # 基础安装 pip install nnsight # 支持 vLLM pip install "nnsight[vllm]" 如需远程 NDIF 执行，请在 [login.ndif.us](https://login.ndif.us) 注册以获取 API 密钥。 ## 核心概念 ### LanguageModel 包装器 python from nnsight import LanguageModel # 加载模型 (底层使用 HuggingFace) model = LanguageModel("openai-community/gpt2", device_map="auto") # 处理大型模型 model = LanguageModel("meta-llama/Llama-3.1-8B", device_map="auto") ### 追踪上下文 `trace` 上下文管理器支持延迟执行 - 操作被收集到计算图中： python 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() # 访问注意力 p