nnsight

nnsight 包允许解释和操作深度学习模型的内部结构。阅读我们的论文！

安装

通过pip安装此包，运行：

pip install nnsight

示例

这是一个简单示例，我们在本地对gpt2运行nnsight API并保存最后一层的隐藏状态：

from nnsight import LanguageModel

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

with model.trace('The Eiffel Tower is in the city of'):

      hidden_states = model.transformer.h[-1].output[0].save()

      output = model.output.save()

让我们逐步分析。

我们从nnsight模块导入LanguageModel对象，并使用gpt2的huggingface仓库ID 'openai-community/gpt2'创建一个gpt2模型。这接受创建模型的参数，包括device_map来指定运行的设备。

from nnsight import LanguageModel

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

然后，我们通过在模型对象上调用.trace(...)来创建一个追踪上下文块。这表示我们要用我们的提示运行模型。

with model.trace('The Eiffel Tower is in the city of') as tracer:

现在调用.trace(...)实际上并不初始化或运行模型。只有在退出追踪块后，才会加载和运行实际模型。块中的所有操作都是"代理"，本质上创建了我们希望稍后执行的操作图。

在这个上下文中，所有操作/干预将应用于给定提示的处理。

hidden_states = model.transformer.h[-1].output[0].save()

在这一行中，我们说，访问transformer的最后一层model.transformer.h[-1]，访问其输出.output，索引为0 .output[0]，并保存它.save()

有几点需要注意，我们可以通过打印模型来查看模型的模块树。这允许我们知道要访问哪些属性来获取所需的模块。运行print(model)得到：

GPT2LMHeadModel(
  (transformer): GPT2Model(
    (wte): Embedding(50257, 768)
    (wpe): Embedding(1024, 768)
    (drop): Dropout(p=0.1, inplace=False)
    (h): ModuleList(
      (0-11): 12 x GPT2Block(
        (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
        (attn): GPT2Attention(
          (c_attn): Conv1D()
          (c_proj): Conv1D()
          (attn_dropout): Dropout(p=0.1, inplace=False)
          (resid_dropout): Dropout(p=0.1, inplace=False)
        )
        (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
        (mlp): GPT2MLP(
          (c_fc): Conv1D()
          (c_proj): Conv1D()
          (act): NewGELUActivation()
          (dropout): Dropout(p=0.1, inplace=False)
        )
      )
    )
    (ln_f): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
  )
  (lm_head): Linear(in_features=768, out_features=50257, bias=False)
)

.output返回此模块输出的代理。这本质上意味着我们在说，当我们在推理过程中到达这个模块的输出时，抓取它并执行我们在其上定义的任何操作（这些操作也成为代理）。这里有两个操作代理，一个用于获取输出的第0个索引，另一个用于保存输出。我们取第0个索引是因为gpt2 transformer层的输出是一个元组，其中第一个索引是实际的隐藏状态（最后两个索引来自注意力）。我们可以在任何代理上调用.shape来获取值最终的形状。运行print(model.transformer.h[-1].output.shape)返回(torch.Size([1, 10, 768]), (torch.Size([1, 12, 10, 64]), torch.Size([1, 12, 10, 64])))

在处理我们正在构建的干预计算图时，当不再需要代理的值时，其值会被取消引用并销毁。然而，在代理上调用.save()会通知计算图保存此代理的值并永不销毁它，允许我们在生成后访问该值。

退出生成器上下文后，模型将使用指定的参数和干预图运行。output将填充实际输出，hidden_states将包含隐藏值。

print(output)
print(hidden_states)

tensor([[ 464,  412,  733,  417, 8765,  318,  287,  262, 1748,  286, 6342]],
       device='cuda:0')
tensor([[[ 0.0505, -0.1728, -0.1690,  ..., -1.0096,  0.1280, -1.0687],
         [ 8.7494,  2.9057,  5.3024,  ..., -8.0418,  1.2964, -2.8677],
         [ 0.2960,  4.6686, -3.6642,  ...,  0.2391, -2.6064,  3.2263],
         ...,
         [ 2.1537,  6.8917,  3.8651,  ...,  0.0588, -1.9866,  5.9188],
         [-0.4460,  7.4285, -9.3065,  ...,  2.0528, -2.7946,  0.5556],
         [ 6.6286,  1.7258,  4.7969,  ...,  7.6714,  3.0682,  2.0481]]],
       device='cuda:0')

操作

大多数基本操作和torch操作都适用于代理，并被添加到计算图中。

from nnsight import LanguageModel
import torch 

model = LanguageModel('openai-community/gpt2', device_map='cuda')

with model.trace('The Eiffel Tower is in the city of'):

  hidden_states_pre = model.transformer.h[-1].output[0].save()

  hs_sum = torch.sum(hidden_states_pre).save()

  hs_edited = hidden_states_pre + hs_sum

  hs_edited = hs_edited.save()

print(hidden_states_pre)
print(hs_sum)
print(hs_edited)

在这个例子中,我们获取隐藏状态的总和并将其添加到隐藏状态本身(无论出于什么原因)。通过保存各个步骤,我们可以看到值是如何变化的。

tensor([[[ 0.0505, -0.1728, -0.1690,  ..., -1.0096,  0.1280, -1.0687],
         [ 8.7494,  2.9057,  5.3024,  ..., -8.0418,  1.2964, -2.8677],
         [ 0.2960,  4.6686, -3.6642,  ...,  0.2391, -2.6064,  3.2263],
         ...,
         [ 2.1537,  6.8917,  3.8651,  ...,  0.0588, -1.9866,  5.9188],
         [-0.4460,  7.4285, -9.3065,  ...,  2.0528, -2.7946,  0.5556],
         [ 6.6286,  1.7258,  4.7969,  ...,  7.6714,  3.0682,  2.0481]]],
       device='cuda:0')
tensor(501.2957, device='cuda:0')
tensor([[[501.3461, 501.1229, 501.1267,  ..., 500.2860, 501.4237, 500.2270],
         [510.0451, 504.2014, 506.5981,  ..., 493.2538, 502.5920, 498.4279],
         [501.5916, 505.9643, 497.6315,  ..., 501.5348, 498.6892, 504.5219],
         ...,
         [503.4493, 508.1874, 505.1607,  ..., 501.3545, 499.3091, 507.2145],
         [500.8496, 508.7242, 491.9892,  ..., 503.3485, 498.5010, 501.8512],
         [507.9242, 503.0215, 506.0926,  ..., 508.9671, 504.3639, 503.3438]]],
       device='cuda:0')

设置

我们通常不仅想看到计算过程中发生的事情,还想干预和编辑信息流。

from nnsight import LanguageModel
import torch 

model = LanguageModel('openai-community/gpt2', device_map='cuda')

with model.trace('The Eiffel Tower is in the city of') as tracer:

  hidden_states_pre = model.transformer.h[-1].mlp.output.clone().save()

  noise = (0.001**0.5)*torch.randn(hidden_states_pre.shape)

  model.transformer.h[-1].mlp.output = hidden_states_pre + noise

  hidden_states_post = model.transformer.h[-1].mlp.output.save()

print(hidden_states_pre)
print(hidden_states_post)

在这个例子中,我们创建一个噪声张量来添加到隐藏状态中。然后我们添加它,使用赋值=运算符用这些新的带噪声的激活来更新.output的值。

我们可以在结果中看到变化:

tensor([[[ 0.0505, -0.1728, -0.1690,  ..., -1.0096,  0.1280, -1.0687],
         [ 8.7494,  2.9057,  5.3024,  ..., -8.0418,  1.2964, -2.8677],
         [ 0.2960,  4.6686, -3.6642,  ...,  0.2391, -2.6064,  3.2263],
         ...,
         [ 2.1537,  6.8917,  3.8651,  ...,  0.0588, -1.9866,  5.9188],
         [-0.4460,  7.4285, -9.3065,  ...,  2.0528, -2.7946,  0.5556],
         [ 6.6286,  1.7258,  4.7969,  ...,  7.6714,  3.0682,  2.0481]]],
       device='cuda:0')
tensor([[[ 0.0674, -0.1741, -0.1771,  ..., -0.9811,  0.1972, -1.0645],
         [ 8.7080,  2.9067,  5.2924,  ..., -8.0253,  1.2729, -2.8419],
         [ 0.2611,  4.6911, -3.6434,  ...,  0.2295, -2.6007,  3.2635],
         ...,
         [ 2.1859,  6.9242,  3.8666,  ...,  0.0556, -2.0282,  5.8863],
         [-0.4568,  7.4101, -9.3698,  ...,  2.0630, -2.7971,  0.5522],
         [ 6.6764,  1.7416,  4.8027,  ...,  7.6507,  3.0754,  2.0218]]],
       device='cuda:0')

多个令牌生成

当生成多个令牌时,使用.generate(...)和.next()来表示后续的干预应该应用于后续的生成。

这里我们再次使用gpt2生成,但生成三个令牌并保存最后一层的隐藏状态:

from nnsight import LanguageModel

model = LanguageModel('openai-community/gpt2', device_map='cuda')

with model.generate('埃菲尔铁塔位于', max_new_tokens=3) as tracer:
 
  hidden_states1 = model.transformer.h[-1].output[0].save()

  invoker.next()

  hidden_states2 = model.transformer.h[-1].next().output[0].save()

  invoker.next()

  hidden_states3 = model.transformer.h[-1].next().output[0].save()

跨提示干预

干预操作可以跨提示工作!在同一个生成块中使用两个调用,操作可以在它们之间工作。

你可以通过不向.trace/.generate传递提示来做到这一点,而是在创建的tracer对象上调用.invoke(...)。

在这种情况下,我们获取第一个提示"麦迪逊广场花园位于纽约"的令牌嵌入,并用它们替换第二个提示的嵌入。

from nnsight import LanguageModel

model = LanguageModel('openai-community/gpt2', device_map='cuda')

with model.generate(max_new_tokens=3) as tracer:
    
    with tracer.invoke("麦迪逊广场花园位于纽约"):

        embeddings = model.transformer.wte.output

    with tracer.invoke("_ _ _ _ _ _ _ _ _ _"):

        model.transformer.wte.output = embeddings

        output = model.generator.output.save()

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

这会产生:

麦迪逊广场花园位于纽约市。
_ _ _ _ _ _ _ _ _ _纽约市。

我们也可以输入预先保存的嵌入张量,如下所示:

from nnsight import LanguageModel

model = LanguageModel('openai-community/gpt2', device_map='cuda')

with model.generate(max_new_tokens=3) as tracer:
    
    with tracer.invoke("麦迪逊广场花园位于纽约") as invoker:

        embeddings = model.transformer.wte.output.save()

with model.generate(max_new_tokens=3) as tracer:

    with tracer.invoke("_ _ _ _ _ _ _ _ _ _") as invoker:

        model.transformer.wte.output = embeddings.value

临时模块

我们还可以在计算过程中的任何时候应用模型模块树中的模块,即使它是无序的。

from nnsight import LanguageModel
import torch

model = LanguageModel("openai-community/gpt2", device_map='cuda')

with model.generate('埃菲尔铁塔位于') as generator:

  hidden_states = model.transformer.h[-1].output[0]
  hidden_states = model.lm_head(model.transformer.ln_f(hidden_states)).save()
  tokens = torch.softmax(hidden_states, dim=2).argmax(dim=2).save()
        
print(hidden_states)
print(tokens)
print(model.tokenizer.decode(tokens[0]))

这里我们像往常一样获取最后一层的隐藏状态。我们还链式应用model.transformer.ln_f和model.lm_head以便将隐藏状态"解码"到词汇空间。应用softmax然后argmax允许我们将词汇空间的隐藏状态转换为实际的令牌,然后我们可以使用分词器来解码。

输出看起来像:

张量([[[   -36.2874,   -35.0114,   -38.0793,  ...,   -40.5163,   -41.3759,
            -34.9193],
          [  -68.8886,   -70.1562,   -71.8408,  ...,   -80.4195,   -78.2552,
            -71.1206],
          [  -82.2950,   -81.6519,   -83.9941,  ...,   -94.4878,   -94.5194,
            -85.6998],
          ...,
          [-113.8675,  -111.8628,  -113.6634,  ...,  -116.7652,  -114.8267,
           -112.3621],
          [  -81.8531,   -83.3006,   -91.8192,  ...,   -92.9943,   -89.8382,
            -85.6898],
          [-103.9307,  -102.5054,  -105.1563,  ...,  -109.3099,  -110.4195,
           -103.1395]]], 设备='cuda:0')
张量([[  198,    12,   417,  8765,   318,   257,   262,  3504,  7372,  6342]],
       设备='cuda:0')

-埃菲尔铁塔是巴黎市中心

---

更多示例可以在[nnsight.net](https://www.nnsight.net)找到

### 引用

如果您在研究中使用`nnsight`，请使用以下方式进行引用

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