# nnsight：透明访问神经网络内部结构 nnsight (/ɛn.saɪt/) 使研究人员能够解释和操作任何 PyTorch 模型的内部结构，具有在小型模型上本地运行相同代码或通过 NDIF 在大型模型 (70B+) 上远程运行相同代码的独特功能。 **GitHub**：[ndif-team/nnsight]()（730+ 星）**论文**：[NNsight 和 NDIF：民主化对基金会模型内部的访问]() (ICLR 2025) ## 关键价值主张 **一次编写，随处运行**：相同的可解释性代码可在本地 GPT-2 或远程 Llama-3.1-405B 上运行。只需切换“remote=True”即可。 # 本地执行（小模型）使用 model.trace("Hello world")：隐藏 = model.transformer.h[5].output[0].save() # 远程执行（大规模模型）-相同的代码！使用 model.trace("Hello world", remote=True)：隐藏 = model.model.layers[40].output[0].save() ## 何时使用 nnsight **当您需要时使用 nnsight：** * 在对于本地 GPU 来说太大的模型（70B、405B）上运行可解释性实验 * 适用于任何 PyTorch 架构（变压器、Mamba、自定义模型） * 执行多代币生成干预 * 在不同的提示之间共享激活 * 无需重新实现即可访问完整模型内部结构 **在以下情况下考虑替代方案：** * 您希望跨模型保持一致的 API → 使用 **TransformerLens** * 您需要声明性的、可共享的干预措施 → 使用 **pyvene** * 您正在训练 SAE → 使用 **SAELens** * 您仅在本地使用小模型 → **TransformerLens** 可能更简单 ## 安装 # 基本安装 pip 安装 nnsight # 对于 vLLM 支持 pip install "nnsight[vllm]" 对于远程 NDIF 执行，请在 [login.ndif.us]() 上注册以获取 API 密钥。 ## 核心概念 ### 语言模型包装器从 nnsight 导入 LanguageModel # 加载模型（在底层使用 HuggingFace）模型 = LanguageModel("openai-community/gpt2", device_map="auto") # 对于较大的模型模型 = LanguageModel("meta-llama/Llama-3.1-8B", device_map="auto") ### 追踪上下文 `trace` 上下文管理器启用延迟执行 - 操作被收集到计算图中：从 nnsight 导入 LanguageModel 模型 = LanguageModel("gpt2", device_map="auto") 使用 model.trace("埃菲尔铁塔在") 作为跟踪器： # 访问任何模块的输出 hidden_states = model.transformer.h[5].output[0].save() # 访问注意力模式 attn = model.transformer.h[5].attn.attn_dropout.input[0][0].save() # 修改激活 model.transformer.h[8].output[0][:] = 0 # 将第 8 层清零 # 获取最终输出 logits = model.output.save() # 上下文退出后，访问保存的值 print(hidden_states.shape) # [批量、序列、隐藏] ### 代理对象在 `trace` 内部，模块访问返回记录操作的 Proxy 对象：与 model.trace("Hello"): # 这些都是代理对象 - 操作被延迟 h5_out = model.transformer.h[5].output[0] # 代理 h5_mean = h5_out.mean(dim=-1) # 代理 h5_saved = h5_mean.save() # 保存以供以后访问 ## 工作流程 1：激活分析 ### 一步一步从 nnsight 导入 LanguageModel 进口火炬模型 = LanguageModel("gpt2", device_map="auto") 提示=“法国的首都是” 使用 model.trace(prompt) 作为跟踪器： # 1.从多个层收集激活层输出 = [] for i in range(12): # GPT-2 有 12 层 Layer_out = model.transformer.h[i].output[0].save() layer_outputs.append(layer_out) # 2. 获取注意力模式 attn_patterns = [] 对于范围 (12) 内的 i： # 访问注意力权重（在softmax之后） attn = model.transformer.h[i].attn.attn_dropout.input[0][0].save() attn_patterns.append（attn） # 3. 获取最终的 logits logits = model.output.save() # 4. 分析外部环境对于我，枚举中的layer_out（layer_outputs）： print(f"第 {i} 层输出形状：{layer_out.shape}") print(f"第 {i} 层范数：{layer_out.norm().item():.3f}") # 5. 查找热门预测 probs = torch.softmax(logits[0, -1], 暗淡=-1) top_tokens = probs.topk(5) 对于令牌，zip 中的概率（top_tokens.indices，top_tokens.values）： print(f"{model.tokenizer.decode(token)}: {prob.item():.3f}") ### 清单 * 使用 LanguageModel 包装器加载模型 * 使用跟踪上下文进行操作 * 在上下文之后调用你需要的值`.save()` * 访问上下文之外保存的值 * 使用`.shape`、`.norm()`等进行分析 ## 工作流程 2：激活补丁 ### 一步一步从 nnsight 导入 LanguageModel 进口火炬模型 = LanguageModel("gpt2", device_map="auto") clean_prompt = "埃菲尔铁塔在" baded_prompt = "斗兽场在" # 1. 获得干净的激活使用 model.trace(clean_prompt) 作为跟踪器： clean_hidden = model.transformer.h[8].output[0].save() # 2. 修补损坏的运行使用 model.trace(corrupted_prompt) 作为跟踪器： # 用干净的激活替换第 8 层输出 model.transformer.h[8].output[0][:] = clean_hidden patched_logits = model.output.save() # 3. 比较预测 paris_token = model.tokenizer.encode("巴黎")[0] rome_token = model.tokenizer.encode("罗马")[0] patched_probs = torch.softmax(patched_logits[0, -1], dim=-1) print(f"巴黎概率：{patched_probs[paris_token].item():.3f}") print(f"罗马概率：{patched_probs[rome_token].item():.3f}") ### 系统性修补扫描 def patch_layer_position（层，位置，clean_cache，corrupted_prompt）： """将单层/位置从干净修补到损坏。""" 使用 model.trace(corrupted_prompt) 作为跟踪器： # 获取当前激活状态当前= model.transformer.h[层].输出[0] # 仅修补特定位置当前[:, 位置, :] = clean_cache[层][:, 位置, :] logits = model.output.save() 返回逻辑值 # 扫描所有图层和位置结果 = torch.zeros(12, seq_len) 对于范围（12）内的图层：对于范围内的位置（seq_len）： logits = patch_layer_position(层, pos, clean_hidden, 损坏) 结果[层，位置]=compute_metric(logits) ## 工作流程 3：使用 NDIF 远程执行在没有本地 GPU 的大规模模型上运行相同的实验。 ### 一步一步从 nnsight 导入 LanguageModel # 1.加载大模型（将远程运行）模型 = LanguageModel("meta-llama/Llama-3.1-70B") # 2.同样的代码，只需添加remote=True 使用 model.trace("生命的意义是", remote=True) 作为跟踪器： # 访问 70B 型号的内部结构！ layer_40_out = model.model.layers[40].output[0].save() logits = model.output.save() # 3. NDIF 返回的结果 print(f"第 40 层形状：{layer_40_out.shape}") # 4. 干预一代使用 model.trace(remote=True) 作为跟踪器： with tracer.invoke("2+2 是什么？"): # 在生成过程中进行干预 model.model.layers[20].output[0][:, -1, :] *= 1.5 输出 = model.generate(max_new_tokens=50) ### NDIF 设置 1. 在 [login.ndif.us]() 注册 2. 获取API密钥 3.设置环境变量或传递给nnsight：导入操作系统 os.environ["NDIF_API_KEY"] = "your_key" # 或者直接配置从 nnsight 导入 CONFIG CONFIG.API_KEY =“你的密钥” ### NDIF 上的可用型号 * 美洲驼-3.1-8B、70B、405B * DeepSeek-R1 型号 * 各种开放重量模型（查看 [ndif.us]() 以获取当前列表） ## 工作流程 4：交叉提示激活共享在单个跟踪中的不同输入之间共享激活。从 nnsight 导入 LanguageModel 模型 = LanguageModel("gpt2", device_map="auto") 使用 model.trace() 作为跟踪器： # 第一个提示 with tracer.invoke("猫坐在上面"): cat_hidden = model.transformer.h[6].output[0].save() # 第二个提示 - 注入猫的激活 with tracer.invoke("狗跑过"): # 替换为猫在第 6 层的激活 model.transformer.h[6].output[0][:] = cat_hidden dog_with_cat = model.output.save() # 狗提示现在有猫的内部表示 ## 工作流程 5：基于梯度的分析在向后传递过程中访问梯度。从 nnsight 导入 LanguageModel 进口火炬模型 = LanguageModel("gpt2", device_map="auto") 使用 model.trace("The Quick Brown Fox") 作为追踪器： # 保存激活并启用梯度隐藏 = model.transformer.h[5].output[0].save() 隐藏.retain_grad() logits = 模型.输出 # 计算特定代币的损失 target_token = model.tokenizer.encode("跳跃")[0] 损失 = -logits[0, -1, target_token] # 向后传递 loss.backward() # 访问梯度等级 = 隐藏.grad print(f"渐变形状：{grad.shape}") print(f"梯度范数: {grad.norm().item():.3f}") **注意**：vLLM 或远程执行不支持梯度访问。 ## 常见问题及解决方案 ### 问题：模型之间的模块路径不同 # GPT-2结构 model.transformer.h[5].output[0] # LLaMA 结构 model.model.layers[5].output[0] # 解决方案：检查模型结构 print(model._model) # 查看实际模块名称 ### 问题：忘记保存 # 错误：无法在跟踪之外访问值与 model.trace("Hello"): hide = model.transformer.h[5].output[0] # 未保存！ print(hidden) # 错误或错误值 # 右：调用 .save() 与 model.trace("Hello"): 隐藏 = model.transformer.h[5].output[0].save() 打印（隐藏）# 有效！ ### 问题：远程超时 # 对于长时间操作，增加超时时间使用 model.trace("prompt", remote=True, timeout=300) 作为跟踪器： # 长时间运行... ### 问题：内存中有许多已保存的激活 # 只保存你需要的使用 model.trace(“提示”)： # 不要保存所有内容对于范围（100）内的 i： model.transformer.h[i].output[0].save() # 内存很重！ # 更好：保存特定层键层 = [0, 5, 11] 对于 key_layers 中的 i： model.transformer.h[i].output[0].save() ### 问题：vLLM 梯度限制 # vLLM 不支持梯度 # 使用标准执行进行梯度分析 model = LanguageModel("gpt2", device_map="auto") # 不是 vLLM ## 关键 API 参考方法/属性 |目的 ---|--- `model.trace（提示，远程=False）` |开始追踪上下文 `proxy.save()` |保存值以供跟踪后访问 `代理[:]` |切片/索引代理（分配补丁） `tracer.invoke（提示）` |在跟踪中添加提示 `model.generate(...)` |通过干预产生 `模型.输出` |最终模型输出 logits `模型._模型` |底层 HuggingFace 模型 ## 与其他工具的比较特色| nnsight |变压器镜头 |比文 ---|---|---|--- 任何建筑|是的 |仅变形金刚 |是的远程执行|是 (NDIF) |没有 |否一致的API |没有 |是的 |是的延迟执行 |是的 |没有 |否 HuggingFace 原生 |是的 |重新实施 |是的可共享的配置 |没有 |没有 |是的 ## 参考文档有关详细的 API 文档、教程和高级用法，请参阅“references/”文件夹：文件|内容 ---|--- [参考文献/README.md]() |概述和快速入门指南 [references/api.md]() | LanguageModel、跟踪、代理对象的完整 API 参考 [references/tutorials.md]() |本地和远程可解释性的分步教程 ## 外部资源 ### 教程 * [入门]() * [功能概述]() * [远程执行]() * [应用教程]() ### 官方文档 * [官方文档]() * [NDIF 信息]() * [社区论坛]() ### 论文 * [NNsight 和 NDIF 论文]() \- Fiotto-Kaufman 等人。（ICLR 2025） ## 架构支持 nnsight 适用于任何 PyTorch 模型： * **变形金刚**：GPT-2、LLaMA、Mistral 等。 * **状态空间模型**：Mamba * **视觉模型**：ViT、CLIP * **自定义架构**：任何 nn.Module 关键是了解模块结构以访问正确的组件。

Nnsight 远程可解释性

包含技能