Prune the attention_heads based on the probing results. Still not supporting reover from base. Only support Llama-2-7b-chat-hf Args: args (_type_): _description_ model (_type_): _description_ model_base (_type_, optional): _description_. Defaults to None. device
(
args, model, model_base=None, device=torch.device("cuda:0"), top_k_heads=10
)
| 2066 | |
| 2067 | |
| 2068 | def prune_attention_head( |
| 2069 | args, model, model_base=None, device=torch.device("cuda:0"), top_k_heads=10 |
| 2070 | ): |
| 2071 | """Prune the attention_heads based on the probing results. Still not supporting reover from base. Only support Llama-2-7b-chat-hf |
| 2072 | |
| 2073 | Args: |
| 2074 | args (_type_): _description_ |
| 2075 | model (_type_): _description_ |
| 2076 | model_base (_type_, optional): _description_. Defaults to None. |
| 2077 | device (_type_, optional): _description_. Defaults to torch.device("cuda:0"). |
| 2078 | |
| 2079 | Raises: |
| 2080 | ValueError: _description_ |
| 2081 | """ |
| 2082 | |
| 2083 | layers = model.model.layers |
| 2084 | k = top_k_heads |
| 2085 | print("Pruning top {} attention heads".format(k)) |
| 2086 | |
| 2087 | # find the top-k attention heads in probing results based on the value in the probing_result |
| 2088 | if args.model == "llama2-7b-chat-hf": |
| 2089 | with open("data/probing_result_7b.json", "r") as f: |
| 2090 | # read json file to dict |
| 2091 | probing_result = json.load(f) |
| 2092 | count = sum(value == 1.0 for value in probing_result.values()) |
| 2093 | if k <= count: |
| 2094 | top_k_heads_full = heapq.nlargest( |
| 2095 | 132, probing_result, key=probing_result.get |
| 2096 | ) |
| 2097 | top_k_heads = random.sample(top_k_heads_full, k) |
| 2098 | elif k <= len(probing_result): |
| 2099 | top_k_heads = heapq.nlargest(k, probing_result, key=probing_result.get) |
| 2100 | else: |
| 2101 | raise ValueError("k is larger than the number of attention heads") |
| 2102 | |
| 2103 | extracted_numbers = [ |
| 2104 | list(map(int, re.findall(r"\d+", head))) for head in top_k_heads |
| 2105 | ] |
| 2106 | |
| 2107 | for head in extracted_numbers: |
| 2108 | block_id = head[0] |
| 2109 | head_id = head[1] |
| 2110 | layer = layers[block_id] |
| 2111 | subset = find_layers(layer) |
| 2112 | for name in ["self_attn.k_proj", "self_attn.v_proj", "self_attn.q_proj"]: |
| 2113 | W = subset[name].weight.data |
| 2114 | W_metric = torch.zeros_like(W) |
| 2115 | W_metric[:, head_id * 128 : (head_id + 1) * 128] = 1 |
| 2116 | W_mask = W_metric == 1 |
| 2117 | subset[name].weight.data[W_mask] = 0 |
| 2118 | name = "self_attn.o_proj" |
| 2119 | W = subset[name].weight.data |
| 2120 | W_metric = torch.zeros_like(W) |
| 2121 | W_metric[head_id * 128 : (head_id + 1) * 128, :] = 1 |
| 2122 | W_mask = W_metric == 1 |
| 2123 | subset[name].weight.data[W_mask] = 0 |
| 2124 | elif args.model == "llama2-13b-chat-hf": |
| 2125 | with open("data/probing_result_13b.json", "r") as f: |