File size: 3,338 Bytes
a20151e | 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 | import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import csv
import numpy as np
# ------------------ Load model ------------------
model_name = "Qwen/Qwen2-0.5B"
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float32)
# ------------------ Detect architecture --------
print(model)
# Compute norms
def compute_norms(A: torch.Tensor):
print(A.shape)
norms = {}
norms["1"] = torch.norm(A, p=1).item()
norms["inf"] = torch.norm(A, p=float("inf")).item()
norms["2"] = torch.linalg.svdvals(A).max().item()
norms["F"] = torch.norm(A, p="fro").item()
norms["*"] = torch.linalg.svdvals(A).sum().item()
norms["max"] = A.abs().max().item()
norms["2,1"] = torch.norm(A, dim=0, p=2).sum().item()
norms["2,1,T"] = torch.norm(A.t(), dim=0, p=2).sum().item()
return norms
def dict_of_lists(norm_list):
result = {k: [] for k in norm_list[0].keys()}
for d in norm_list:
for k, v in d.items():
result[k].append(v)
return result
# Store norms and ratios
all_norms_Q, all_norms_K, norm_ratios = [], [], []
for layer in model.model.layers:
W_q = layer.self_attn.q_proj.weight.detach()
W_k = layer.self_attn.k_proj.weight.detach()
if W_k.shape[0] != W_q.shape[0]:
repeat_factor = W_q.shape[0] // W_k.shape[0]
W_k = W_k.repeat(repeat_factor, 1)
norms_q = compute_norms(W_q)
norms_k = compute_norms(W_k)
all_norms_Q.append(norms_q)
all_norms_K.append(norms_k)
ratio_dict = {k: norms_q[k] / norms_k[k] for k in norms_q.keys()}
norm_ratios.append(ratio_dict)
Q_dict = dict_of_lists(all_norms_Q)
K_dict = dict_of_lists(all_norms_K)
R_dict = dict_of_lists(norm_ratios)
sections = [
(["1", "inf", "2", "F"], ["1", "inf", "2", "F"]),
(["*", "max", "2,1", "2,1,T"], ["*", "max", "2,1", "2,1,T"])
]
with open("qwen2.5_norms.csv", "w", newline="") as f:
writer = csv.writer(f)
for idx, (main_norms, keys) in enumerate(sections):
# Header 1
header1 = ["Layer"]
for norm in main_norms:
header1.extend([norm, "", ""])
writer.writerow(header1)
# Header 2
header2 = [""]
for _ in main_norms:
header2.extend(["Q", "K", "Q/K"])
writer.writerow(header2)
# Per-layer rows
for i in range(len(all_norms_Q)):
row = [i + 1]
for key in keys:
row.extend([
round(Q_dict[key][i], 4),
round(K_dict[key][i], 4),
round(R_dict[key][i], 4)
])
writer.writerow(row)
# Mean and Std
mean_row = ["Mean and Std"]
std_row = [""]
for key in keys:
q_vals = np.array(Q_dict[key])
k_vals = np.array(K_dict[key])
r_vals = np.array(R_dict[key])
mean_row.extend([
round(q_vals.mean(), 4),
round(k_vals.mean(), 4),
round(r_vals.mean(), 4)
])
std_row.extend([
round(q_vals.std(), 4),
round(k_vals.std(), 4),
round(r_vals.std(), 4)
])
writer.writerow(mean_row)
writer.writerow(std_row)
|