解读并操控深度学习模型的内部机制

文档 | GitHub | Discord | 论坛 | Twitter | 论文

---

关于

nnsight 是一个 Python 库，能够帮助用户解读和干预深度学习模型的内部运作。它提供了一个简洁、符合 Python 风格的接口，用于：

• 在前向传播过程中访问任意层的激活值
• 修改激活值以研究因果效应
• 计算中间值的梯度
• 高效地对多个输入执行批量干预

nnsight 最初由东北大学的 NDIF 团队 开发，支持在任何 PyTorch 模型上进行本地运行，同时也可通过 NDIF 基础设施对大型模型进行远程执行。

📖 如需深入了解 nnsight 的内部实现细节（如追踪、交错执行、Envoy 系统等），请参阅 NNsight.md。

---

安装

pip install nnsight

---

代理工具

使用以下方法之一，指导 LLM 代理如何使用 nnsight：

技能库

Claude Code

# 打开 Claude Code 终端
claude

# 添加技能市场（仅需一次）
/plugin marketplace add https://github.com/ndif-team/skills.git

# 安装所有技能
/plugin install nnsight@skills

OpenAI Codex

# 打开 OpenAI Codex 终端
codex

# 安装技能
skill-installer install https://github.com/ndif-team/skills.git

Context7 MCP

此外，您也可以使用 Context7 将最新的 nnsight 文档直接提供给您的 LLM。只需在提示中加入 use context7，或在您的 MCP 客户端配置中进行设置：

{
  "mcpServers": {
    "context7": {
      "url": "https://mcp.context7.com/mcp"
    }
  }
}

有关不同 IDE 的完整安装说明，请参阅 Context7 的 README。

文档文件

您还可以将我们的文档直接添加到代理的上下文中：

• CLAUDE.md — 针对使用 nnsight 的 AI 代理的全面指南
• NNsight.md — nnsight 内部机制的深入技术文档

---

快速入门

from nnsight import LanguageModel

model = LanguageModel('openai-community/gpt2', device_map='auto', dispatch=True)

with model.trace('埃菲尔铁塔位于城市里'):
    # 干预激活值（必须按执行顺序访问！）
    model.transformer.h[0].output[0][:] = 0
    
    # 访问并保存后续层的隐藏状态
    hidden_states = model.transformer.h[-1].output[0].save()
    
    # 获取模型输出
    output = model.output.save()

print(model.tokenizer.decode(output.logits.argmax(dim=-1)[0]))

💡 提示： 对于希望在跟踪结束后仍可访问的值，请务必调用 .save()。否则，这些值会被垃圾回收。您也可以使用 nnsight.save(value) 作为替代方案。

访问激活值

with model.trace("埃菲尔铁塔位于城市里"):
    # 访问注意力输出
    attn_output = model.transformer.h[0].attn.output[0].save()
    
    # 访问 MLP 输出
    mlp_output = model.transformer.h[0].mlp.output.save()

    # 访问任意层的输出（按执行顺序访问）
    layer_output = model.transformer.h[5].output[0].save()
    
    # 访问最终的 logits
    logits = model.lm_head.output.save()

注意： GPT-2 的 Transformer 层会返回元组，其中索引 0 包含隐藏状态。

修改激活值

原地修改

with model.trace("你好"):
    # 将所有激活值置零
    model.transformer.h[0].output[0][:] = 0
    
    # 修改特定位置
    model.transformer.h[0].output[0][:, -1, :] = 0  # 仅最后一个 token

替换

with model.trace("你好"):
    # 向激活值中添加噪声
    hs = model.transformer.h[-1].mlp.output.clone()
    noise = 0.01 * torch.randn(hs.shape)
    model.transformer.h[-1].mlp.output = hs + noise
    
    result = model.transformer.h[-1].mlp.output.save()

使用 Invoker 进行批量处理

在一个前向传播中处理多个输入。每个 invoke 都会在一个独立的工作线程中运行：

• 线程按顺序执行（无竞争条件）
• 每个线程通过 .output、.input 等等待所需值
• Invoke 按照定义的顺序运行
• 跨 invoke 引用有效，因为线程是依次执行的
• 在单个 invoke 内，只能按执行顺序访问模块

with model.trace() as tracer:
    # 第一个 invoke：工作线程 1
    with tracer.invoke("埃菲尔铁塔位于"):
        embeddings = model.transformer.wte.output  # 线程在此处等待
        output1 = model.lm_head.output.save()
    
    # 第二个 invoke：工作线程 2（在线程 1 完成后运行）
    with tracer.invoke("_ _ _ _ _ _"):
        model.transformer.wte.output = embeddings  # 使用线程 1 的值
        output2 = model.lm_head.output.save()

无提示的 Invoke

使用不带参数的 .invoke() 来操作整个批次：

with model.trace() as tracer:
    with tracer.invoke("你好"):
        out1 = model.lm_head.output[:, -1].save()
    
    with tracer.invoke(["世界", "测试"]):
        out2 = model.lm_head.output[:, -1].save()
    
    # 不带参数的 invoke：作用于全部 3 个输入
    with tracer.invoke():
        out_all = model.lm_head.output[:, -1].save()  # 形状：[3, 词汇表]

多 token 生成

使用 .generate() 进行自回归生成：

with model.generate("埃菲尔铁塔位于", max_new_tokens=3) as tracer:
    output = model.generator.output.save()

print(model.tokenizer.decode(output[0]))
# “埃菲尔铁塔位于巴黎市”

遍历生成步骤

with model.generate("你好", max_new_tokens=5) as tracer:
    logits = list().save()
    
    # 遍历所有生成步骤
    for step in tracer.iter[:]:
        logits.append(model.lm_head.output[0][-1].argmax(dim=-1))

print(model.tokenizer.batch_decode(logits))

每步的条件干预

with model.generate("Hello", max_new_tokens=5) as tracer:
    outputs = list().save()
    
    for step_idx in tracer.iter[:]:
        if step_idx == 2:
            model.transformer.h[0].output[0][:] = 0  # 仅在第2步执行

        outputs.append(model.transformer.h[-1].output[0])

⚠️ 注意： tracer.iter[:] 之后的代码永远不会执行！无界迭代器会一直等待更多步骤。请将迭代后的代码放在单独的 tracer.invoke() 中。当使用多个调用时，不要将输入传递给 generate() —— 应该将其传递给第一次调用：

with model.generate(max_new_tokens=3) as tracer:
    with tracer.invoke("Hello"):  # 第一次调用——在此处传递输入
        for step in tracer.iter[:]:
            hidden = model.transformer.h[-1].output.save()
    with tracer.invoke():  # 第二次调用——在生成完成后运行
        final = model.output.save()  # 现在可以正常工作了！

梯度

梯度只能在 张量（而非模块）上访问，并且必须位于 with tensor.backward(): 上下文中：

with model.trace("Hello"):
    hs = model.transformer.h[-1].output[0]
    hs.requires_grad_(True)
    
    logits = model.lm_head.output
    loss = logits.sum()
    
    with loss.backward():
        grad = hs.grad.save()

print(grad.shape)

模型编辑

创建持久化的模型修改：

# 创建编辑后的模型（非破坏性）
with model.edit() as model_edited:
    model.transformer.h[0].output[0][:] = 0

# 原始模型未被改变
with model.trace("Hello"):
    out1 = model.transformer.h[0].output[0].save()

# 编辑后的模型已应用修改
with model_edited.trace("Hello"):
    out2 = model_edited.transformer.h[0].output[0].save()

assert not torch.all(out1 == 0)
assert torch.all(out2 == 0)

扫描（形状推断）

无需运行完整模型即可获取形状。与所有跟踪上下文一样，需要使用 .save() 将值保存到块之外：

import nnsight

with model.scan("Hello"):
    dim = nnsight.save(model.transformer.h[0].output[0].shape[-1])

print(dim)  # 768

缓存激活值

自动缓存模块的输出：

with model.trace("Hello") as tracer:
    cache = tracer.cache()

# 访问缓存的值
layer0_out = cache['model.transformer.h.0'].output
print(cache.model.transformer.h[0].output[0].shape)

会话

为了提高效率，可以将多个跟踪分组在一起：

with model.session() as session:
    with model.trace("Hello"):
        hs1 = model.transformer.h[0].output[0].save()
    
    with model.trace("World"):
        model.transformer.h[0].output[0][:] = hs1  # 使用第一次跟踪中的值
        hs2 = model.transformer.h[0].output[0].save()

远程执行（NDIF）

在 NDIF 的远程基础设施上运行：

from nnsight import CONFIG
CONFIG.set_default_api_key("YOUR_API_KEY")

model = LanguageModel("meta-llama/Meta-Llama-3.1-8B")

with model.trace("Hello", remote=True):
    hidden_states = model.model.layers[-1].output.save()

可在 nnsight.net/status 查看可用模型

vLLM 集成

使用 vLLM 进行高性能推理：

from nnsight.modeling.vllm import VLLM

model = VLLM("gpt2", tensor_parallel_size=1, dispatch=True)

with model.trace("Hello", temperature=0.0, max_tokens=5) as tracer:
    logits = list().save()
    
    for step in tracer.iter[:]:
        logits.append(model.logits.output)

NNsight 适用于任何 PyTorch 模型

可将 NNsight 用于任意 PyTorch 模型：

from nnsight import NNsight
import torch

net = torch.nn.Sequential(
    torch.nn.Linear(5, 10),
    torch.nn.Linear(10, 2)
)

model = NNsight(net)

with model.trace(torch.rand(1, 5)):
    layer1_out = model[0].output.save()
    output = model.output.save()

源码跟踪

访问模块前向传播中的中间操作。.source 会重写前向方法，以钩住所有操作：

# 发现可用的操作
print(model.transformer.h[0].attn.source)
# 显示带有操作名称的前向方法，例如：
#   attention_interface_0 -> 66  attn_output, attn_weights = attention_interface(...)
#   self_c_proj_0         -> 79  attn_output = self.c_proj(attn_output)

# 访问操作的值
with model.trace("Hello"):
    attn_out = model.transformer.h[0].attn.source.attention_interface_0.output.save()

临时模块应用

可以在非正常执行顺序的情况下应用模块：

with model.trace("埃菲尔铁塔位于城市中"):
    # 获取中间隐藏状态
    hidden_states = model.transformer.h[-1].output[0]
    
    # 应用 lm_head 以获得“logit lens”视图
    logits = model.lm_head(model.transformer.ln_f(hidden_states))
    tokens = logits.argmax(dim=-1).save()

print(model.tokenizer.decode(tokens[0]))

---

核心概念

基于线程同步的延迟执行

NNsight 使用 延迟执行 和 基于线程的同步：

1. 代码提取：当你进入 with model.trace(...) 块时，nnsight 会捕获你的代码（通过 AST），并立即退出该块。
2. 线程执行：你的代码会在一个独立的工作线程中运行。
3. 值等待：当你访问 .output 时，线程会 等待 直到模型提供该值。
4. 钩子注入：模型使用 PyTorch 钩子将值提供给等待的线程。

with model.trace("Hello"):
    # 代码在工作线程中运行
    # 线程在此处等待，直到层的输出可用
    hs = model.transformer.h[-1].output[0]
    
    # .save() 标记该值，以便在上下文退出后保留
    hs = hs.save()
    # 或者：hs = nnsight.save(hs)

# 退出后，hs 包含实际的张量
print(hs.shape)  # torch.Size([1, 2, 768])

关键点： 你的代码直接运行。当你访问 .output 时，你得到的是 真实的张量——你的线程只是等待它变得可用。

重要提示： 在一次调用中，必须按照执行顺序访问模块。如果在访问第2层的输出之前就尝试访问第5层的输出，就会导致死锁（因为第2层已经执行完毕）。

关键属性

每个模块都具有这些特殊属性。访问它们会导致工作线程 等待 该值：

属性	描述
.output	模块前向传播的输出（线程等待）
.input	模块的第一个位置参数
.inputs	所有输入，格式为 (args_tuple, kwargs_dict)

注意： .grad 只能在 张量（而非模块）上访问，并且必须位于 with tensor.backward(): 上下文中。

模块层次结构

打印模型以查看其结构：

print(model)
# GPT2LMHeadModel(
#   (transformer): GPT2Model(
#     (h): ModuleList(
#       (0-11): 12 x GPT2Block(
#         (attn): GPT2Attention(...)
#         (mlp): GPT2MLP(...)
#       )
#     )
#   )
#   (lm_head): Linear(...)
# )

---

故障排除

错误	原因	解决方法
OutOfOrderError: Value was missed...	模块访问顺序错误	按照前向传播的执行顺序访问模块
tracer.iter[:] 之后出现 NameError	无界迭代后的代码未运行	对于迭代后的代码，请使用单独的 tracer.invoke()；将输入传递给第一次调用，而不是 generate()
ValueError: Cannot invoke during an active model execution	在使用多次调用时，向 generate() 传递了输入	使用 model.generate(max_new_tokens=N) 而不需输入；将提示词传递给第一次 tracer.invoke("Hello")
ValueError: Cannot return output of Envoy...	跟踪时未提供输入	提供输入：model.trace(input) 或使用 tracer.invoke(input)

更多调试技巧，请参阅 文档。

---

更多资源

• 文档 — 教程、指南和 API 参考
• NNsight.md — 关于 nnsight 内部机制的深度技术文档
• CLAUDE.md — 面向使用 nnsight 的 AI 代理的全面指南
• 性能报告 — 详细的性能分析与基准测试

---

引用

如果您在研究中使用了 nnsight，请引用以下内容：

@article{fiottokaufman2024nnsightndifdemocratizingaccess,
      title={NNsight 和 NDIF： democratizing 访问基础模型内部}, 
      author={Jaden Fiotto-Kaufman 和 Alexander R Loftus 和 Eric Todd 和 Jannik Brinkmann 和 Caden Juang 和 Koyena Pal 和 Can Rager 和 Aaron Mueller 和 Samuel Marks 和 Arnab Sen Sharma 和 Francesca Lucchetti 和 Michael Ripa 和 Adam Belfki 和 Nikhil Prakash 和 Sumeet Multani 和 Carla Brodley 和 Arjun Guha 和 Jonathan Bell 和 Byron Wallace 和 David Bau},
      year={2024},
      eprint={2407.14561},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2407.14561}, 
}

nnsight 快速上手指南

nnsight 是一个用于解释和干预深度学习模型内部机制的 Python 库。它提供了简洁的接口，支持访问任意层的激活值、修改激活值以研究因果效应、计算中间值的梯度，以及高效地批量处理干预操作。

环境准备

• 系统要求：Linux, macOS, 或 Windows (推荐 Linux 以获得最佳兼容性)
• Python 版本：Python 3.8 或更高版本
• 前置依赖：
  • PyTorch (nnsight 将自动处理大部分依赖，但需确保环境中已安装兼容的 PyTorch 版本)
  • Hugging Face Transformers (用于加载预训练语言模型)
• 网络提示：加载预训练模型（如 GPT-2, Llama 等）时需要连接 Hugging Face。国内用户建议配置镜像源或使用代理加速下载。

安装步骤

使用 pip 直接安装最新稳定版：

pip install nnsight

提示：如果下载速度较慢，可以使用国内镜像源加速安装：

pip install nnsight -i https://pypi.tuna.tsinghua.edu.cn/simple

基本使用

以下示例展示如何加载一个预训练模型（GPT-2），并在前向传播过程中拦截和修改特定层的激活值。

1. 加载模型并追踪执行

使用 LanguageModel 类加载模型，并通过 model.trace() 上下文管理器进入追踪模式。

from nnsight import LanguageModel

# 加载模型，device_map='auto' 会自动分配设备，dispatch=True 启用分布式/远程支持
model = LanguageModel('openai-community/gpt2', device_map='auto', dispatch=True)

with model.trace('The Eiffel Tower is in the city of'):
    # 干预激活值：将第 0 层输出的第一个元素（hidden states）全部置零
    # 注意：必须按照执行顺序访问模块
    model.transformer.h[0].output[0][:] = 0
    
    # 获取并保存最后一层的隐藏状态
    # ⚠️ 重要：想要在 trace 块外部访问数据，必须调用 .save()
    hidden_states = model.transformer.h[-1].output[0].save()
    
    # 获取模型最终输出
    output = model.output.save()

# 解码并打印结果
print(model.tokenizer.decode(output.logits.argmax(dim=-1)[0]))

2. 关键概念说明

• 延迟执行 (Deferred Execution)：with model.trace(...) 块内的代码不会立即执行，而是被记录并重写，以便在模型实际前向传播时注入干预逻辑。
• 数据保存 (.save())：由于执行是延迟的，块内产生的张量在块结束后会被销毁。若需在块外使用数据（如打印或后续计算），必须对目标张量调用 .save() 方法。
• 按序访问：在 trace 块中访问模块（如 model.transformer.h[0]）时，必须遵循模型的前向传播顺序，否则可能引发错误。

3. 进阶：批量处理与生成

nnsight 支持通过 invoke 进行批量干预，以及通过 generate 进行自回归生成时的逐步干预。

# 批量处理示例：在不同输入上应用不同的干预
with model.trace() as tracer:
    with tracer.invoke("Hello"):
        out1 = model.lm_head.output[:, -1].save()
    
    with tracer.invoke(["World", "Test"]):
        out2 = model.lm_head.output[:, -1].save()

# 生成过程干预示例
with model.generate("The Eiffel Tower is in", max_new_tokens=3) as tracer:
    output = model.generator.output.save()

print(model.tokenizer.decode(output[0]))

更多高级功能（如梯度计算、模型编辑、远程执行等）请参考官方文档或源码注释。