利用baukit获取LLM中间层输入输出、修改中间层输出

安装baukit

pip install git+https://github.com/davidbau/baukit

使用方法

1. 获取中间层输入输出

以llama-7b为例，用法见注释：

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from baukit import TraceDict

# 加载模型和tokenizer
model = AutoModelForCausalLM.from_pretrained(model_path,torch_dtype=torch.float16).to(device)
tokenizer = AutoTokenizer.from_pretrained(model_path)

# 将prompt转为token
input_tokens = tokenizer(prompt, return_tensors='pt').to(device)

# 希望获取输入输出的中间层的名字。baukit会在经过这些层时注册钩子并获取输入输出。
hook_layers = [f'model.layers.{l}.self_attn.o_proj' for l in range(n_layers)]

'''
retain_input=True：保留中间层输入
retain_output=True：保留中间层输出
'''
with TraceDict(model, layers=hook_layers, retain_input=True, retain_output=True) as rep:
    model(**input_tokens)
    before_attn=rep[head_hook_layers[5]].input # 第五层attention的O矩阵之前的输入
    after_attn=rep[head_hook_layers[5]].output # 第五层attention的O矩阵之后的输出
    after_attn=rep['model.layers.5.self_attn.q_proj'].output # 会报错，因为head_hook_layers里不包含self_attn.q_proj
    print(before_attn.shape) # torch.Size([1, 28, 4096])，维度是[bsz, num_tokens, dim_model]
    print(after_attn.shape) # torch.Size([1, 28, 4096])

2. 修改中间层输出

注意，修改后会影响模型的前传。即，模型获取某一层的输出ouput，然后将output传给自定义的函数进行修改，然后返回一个新的output传给模型的下一层。

''' 需要定义一个修改输出的函数。这个函数重点是里边那个add_perturbation，
按照baukit的规定，它（add_perturbation）必须接受output和layer_name两个参数：
output参数（模型中间层的输出）；layer_name（表示当前前传到模型的哪一模块了）。

而外面封装的这个wrap_func的参数可以自己随便定义，只要最终返回里边这个add_perturbation即可。
这里给出的用法是：wrap_func传进去了一个edit_layer参数（int），这个参数起到了指定修改特定的哪一层的作用。
外面封装一个参数起到了关键字参数的作用，即维函数的某一个参数指定一个值，以便baukit在修改中间层时使用
（原本的add_perturbation只接受output和layer_name，可能不便实现某些想要的修改效果）。
'''
def wrap_func(edit_layer, perturbation, device, idx=-1):
    def add_perturbation(output, layer_name):
        current_layer = int(layer_name.split(".")[2])
        if current_layer == edit_layer: # 遍历到edit_layer
            if isinstance(output, tuple):
                # output[0].shape [1,12,4096] [1, num_token, dim_hidden]
                output[0][:, idx] += 0.5* perturbation.to(device)
                return output
            else:
                return output
        else:
            return output
    return add_act


edit_layer=5
# 这里用了edit_func将修改函数封装，并传入了edit_layer
intervention_fn = wrap_func(edit_layer, perturbation, model.device)

# 希望利用baukit修改的中间层的名字。baukit会在经过这些层时修改输出。
hook_layers = [f'model.layers.{l}.self_attn.o_proj' for l in range(n_layers)]

# 使用baukit时，将以上的edit_func传给edit_output参数
with TraceDict(model, layers=hook_layers, edit_output=intervention_fn):
	model(**input_tokens)

注：llama 7b的结构：

print(model)
'''
LlamaForCausalLM(
  (model): LlamaModel(
    (embed_tokens): Embedding(32000, 4096, padding_idx=0)
    (layers): ModuleList(
      (0-31): 32 x LlamaDecoderLayer(
        (self_attn): LlamaAttention(
          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (k_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (v_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (o_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (rotary_emb): LlamaRotaryEmbedding()
        )
        (mlp): LlamaMLP(
          (gate_proj): Linear(in_features=4096, out_features=11008, bias=False)
          (up_proj): Linear(in_features=4096, out_features=11008, bias=False)
          (down_proj): Linear(in_features=11008, out_features=4096, bias=False)
          (act_fn): SiLUActivation()
        )
        (input_layernorm): LlamaRMSNorm()
        (post_attention_layernorm): LlamaRMSNorm()
      )
    )
    (norm): LlamaRMSNorm()
  )
  (lm_head): Linear(in_features=4096, out_features=32000, bias=False)
)
'''