-
Notifications
You must be signed in to change notification settings - Fork 1
/
preprocessing.py
189 lines (149 loc) · 11 KB
/
preprocessing.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
import pandas as pd
from tqdm.notebook import tqdm
import json, torch
from transformers import AutoModelWithLMHead, AutoTokenizer
from transformers_pruning_new import AutoTokenizer, GPT2Tokenizer, LlamaForCausalLM
from transformers_pruning_new import GPTNeoForCausalLM, AutoModelForCausalLM, AutoModelWithLMHead
import os
from argparse import ArgumentParser
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# It is used to compute the effects of attention heads on perplexity using both magnitude and gradient as baselines.
def parse_args():
"""Parses the command line arguments."""
parser = ArgumentParser()
parser.add_argument("-", "--model_list", nargs="+", default=[])
return parser.parse_args()
if __name__ == "__main__":
args = parse_args()
model_names = args.model_list
weights_norms = ["magnitude_l2_structured"]
head_contributions_file = "./model/head_contributions.json"
# the file containing the importance scores for the heads of the models
if os.path.exists(head_contributions_file):
head_contributions = json.load(open(head_contributions_file, "r"))
else:
head_contributions = {}
for model_name in model_names:
print(model_name)
biases_scores = "./output/head_scores_" + str(args.model_list[0].replace("/", "_")) + ".csv"
biases = None
if os.path.exists(biases_scores):
df = pd.read_csv(biases_scores)
biases = []
for bias in ["gender_and_sex", "race_ethnicity", "religion", "sexual_orientation", "nationality"]:
if bias in df["Group"].unique():
biases += [bias]
if model_name in ["gpt2","distilgpt2"]:
model = AutoModelWithLMHead.from_pretrained("./saved_models/cached_models/" + model_name).to(device)
elif model_name in ["EleutherAI/gpt-neo-125M", "EleutherAI/gpt-neo-1.3B"]:
model = GPTNeoForCausalLM.from_pretrained("./saved_models/cached_models/" + model_name).to(device)
elif model_name in ["EleutherAI/gpt-j-6B"]:
model = AutoModelForCausalLM.from_pretrained("./saved_models/cached_models/" + model_name, revision="float16",torch_dtype=torch.float16,).to(device)
elif model_name in ["meta-llama/Llama-2-7b-chat-hf"]:
model = LlamaForCausalLM.from_pretrained("./saved_models/cached_models/" + model_name, revision="float16",torch_dtype=torch.float16,).bfloat16().to(device)
model_configs = json.load(open("./model/models_config.json", "r"))
num_heads, num_layers = model_configs[model_name]["num_heads"], model_configs[model_name]["num_layers"]
head_dim, max_length = model_configs[model_name]["head_dim"], model_configs[model_name]["max_length"]
if model_name in ["EleutherAI/gpt-neo-125M", "EleutherAI/gpt-neo-1.3B"]:
tokenizer = GPT2Tokenizer.from_pretrained("./saved_models/cached_tokenizers/" + model_name, padding_side="left")
else:
tokenizer = AutoTokenizer.from_pretrained("./saved_models/cached_tokenizers/" + model_name, padding_side="left")
if model_name in head_contributions.keys():
# if the json files already has the scores for a model, we load them and add the new scores to them
model_dict = head_contributions[model_name]
else:
model_dict = {}
max_length=int(max_length/4)
# this is just to make make sure to make the code run for different models
if biases is not None:
for group in biases:
model_dict["cont_to_" + group + "_bias"], model_dict["cont_to_ppl"] = [], []
for head_id in range(num_layers * num_heads):
# the effects on bias and perplexity are computed for each head of each layer based on the validation dataset
model_dict["cont_to_" + group + "_bias"].append(df[(df["Head id"] == head_id + 1) & (df["Group"] == group) & (df["Split"] == "valid")]["Bias"].mean())
model_dict["cont_to_ppl"].append(df[(df["Head id"] == head_id + 1) & (df["Split"] == "valid")]["PPL"].mean())
for weights_norm in weights_norms:
# for different norms, we compute the scores for each head of each layer
if weights_norm not in model_dict.keys():
model_dict[weights_norm] = []
for head_id in range(num_layers * num_heads):
layer_id = int(head_id/num_heads)
start_idx = head_dim*(head_id - layer_id*num_heads)
end_idx = head_dim*(head_id - (layer_id*num_heads) + 1)
head_weights = []
if model_name in ["EleutherAI/gpt-neo-125M", "EleutherAI/gpt-neo-1.3B"]:
for name, para in model.named_parameters():
if name in ["transformer.h." + str(layer_id) + ".attn.attention.k_proj.weight", "transformer.h." + str(layer_id) + ".attn.attention.v_proj.weight", "transformer.h." + str(layer_id) + ".attn.attention.q_proj.weight","transformer.h." + str(layer_id) + ".attn.attention.out_proj.weight"]:
head_weights.append(para)
head_attn_weight = torch.cat((torch.flatten(head_weights[0][start_idx:end_idx,:]),torch.flatten(head_weights[1][start_idx:end_idx,:]),torch.flatten(head_weights[2][start_idx:end_idx,:])))
head_attn_bias = torch.randn(0).to(device) #any empty tensor because this model does not have a bias parameter for the attention heads
head_proj_weight = torch.flatten(head_weights[3][:,start_idx:end_idx])
if model_name in ["EleutherAI/gpt-j-6B"]:
for name, para in model.named_parameters():
if name in ["transformer.h." + str(layer_id) + ".attn.k_proj.weight", "transformer.h." + str(layer_id) + ".attn.v_proj.weight", "transformer.h." + str(layer_id) + ".attn.q_proj.weight","transformer.h." + str(layer_id) + ".attn.out_proj.weight"]:
head_weights.append(para)
head_attn_weight = torch.cat((torch.flatten(head_weights[0][start_idx:end_idx,:]),torch.flatten(head_weights[1][start_idx:end_idx,:]),torch.flatten(head_weights[2][start_idx:end_idx,:])))
head_attn_bias = torch.randn(0).to(device) #any empty tensor because this model does not have a bias parameter for the attention heads
head_proj_weight = torch.flatten(head_weights[3][:,start_idx:end_idx])
if model_name in ["meta-llama/Llama-2-7b-chat-hf"]:
for name, para in model.named_parameters():
if name in ["model.layers." + str(layer_id) + ".self_attn.q_proj.weight", "model.layers." + str(layer_id) + ".self_attn.k_proj.weight", "model.layers." + str(layer_id) + ".self_attn.v_proj.weight","model.layers." + str(layer_id) + ".self_attn.o_proj.weight"]:
head_weights.append(para)
head_attn_weight = torch.cat((torch.flatten(head_weights[0][start_idx:end_idx,:]),torch.flatten(head_weights[1][start_idx:end_idx,:]),torch.flatten(head_weights[2][start_idx:end_idx,:])))
head_attn_bias = torch.randn(0).to(device) #any empty tensor because this model does not have a bias parameter for the attention heads
head_proj_weight = torch.flatten(head_weights[3][:,start_idx:end_idx])
if model_name in ["gpt2", "distilgpt2"]:
for name, para in model.named_parameters():
if name in ["transformer.h." + str(layer_id) + ".attn.c_attn.weight", "transformer.h." + str(layer_id) + ".attn.c_attn.bias", "transformer.h." + str(layer_id) + ".attn.c_proj.weight"]:
head_weights.append(para)
head_attn_weight = torch.cat((torch.flatten(head_weights[0][:, start_idx:end_idx]),torch.flatten(head_weights[0][:, start_idx + head_dim*num_heads: end_idx + head_dim*num_heads]), torch.flatten(head_weights[0][:, start_idx + 2*head_dim*num_heads : end_idx + 2*head_dim*num_heads])))
head_attn_bias = torch.cat((torch.flatten(head_weights[1][start_idx:end_idx]),torch.flatten(head_weights[1][start_idx + head_dim*num_heads : end_idx + head_dim*num_heads]), torch.flatten(head_weights[1][start_idx + 2*head_dim*num_heads : end_idx + 2*head_dim*num_heads])))
head_proj_weight = torch.flatten(head_weights[2][start_idx:end_idx,:])
if weights_norm == "magnitude_l2_structured":
head_weights_norm = torch.norm(torch.cat((head_attn_weight, head_attn_bias, head_proj_weight)), p=2)
model_dict[weights_norm].append(head_weights_norm.tolist())
# we normalize the scores for each head
for score in weights_norms:
absolute_max = (max(abs(max(model_dict[score])), abs(min(model_dict[score]))))
for head_id in range(num_layers * num_heads):
model_dict[score][head_id] /= absolute_max
if 'mask_gradient_l2_structured' not in model_dict.keys():
# if we have the gradient scores already for a model, we don't need to compute them again.
# the gradient scores are computed for all models. Llama 2 gives an error because of the way it is implemented in hugingface.
stride = 512
torch.manual_seed(1)
names = []
with open("./model/wikitext-2-raw-v1_valid.txt", 'r') as fp:
for line in fp:
x = line
names.append(x)
encodings = tokenizer("".join(names) , return_tensors="pt")
seq_len = encodings.input_ids.size(1)
head_importance = torch.zeros(num_layers, num_heads).to(device)
head_mask = torch.ones(num_layers, num_heads).to(device)
head_mask.requires_grad_(requires_grad=True)
prev_end_loc = 0
heads_gradients = torch.zeros([num_layers * num_heads]).to(device)
for begin_loc in tqdm(range(0, seq_len, stride)):
end_loc = min(begin_loc + max_length, seq_len)
trg_len = end_loc - prev_end_loc # may be different from stride on last loop
input_ids = encodings.input_ids[:, begin_loc:end_loc].to(device)
target_ids = input_ids.clone()
target_ids[:, :-trg_len] = -100
outputs = model(input_ids, labels=target_ids, head_mask=head_mask) ## head_mask does not exist in the llama model implementation, so it gives an error
outputs.loss.backward()
head_importance += head_mask.grad.abs().detach()
prev_end_loc = end_loc
if end_loc == seq_len:
break
# Layerwise importance normalization
exponent = 2
norm_by_layer = torch.pow(torch.pow(head_importance, exponent).sum(-1), 1 / exponent)
head_importance /= norm_by_layer.unsqueeze(-1) + 1e-20
head_importance = head_importance.reshape(num_layers * num_heads).tolist()
absolute_max = (max(abs(max(head_importance)), abs(min(head_importance))))
for head_id in range(num_layers * num_heads):
head_importance[head_id] /= absolute_max
model_dict['mask_gradient_l2_structured'] = head_importance
head_contributions[model_name] = model_dict
json.dump(head_contributions, open("./model/head_contributions.json", "w"))