安装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)
)
'''