Spaces:
Sleeping
Sleeping
| import numpy as np | |
| import pandas as pd | |
| import torch | |
| from tqdm import tqdm | |
| class HeadAnalyzer: | |
| def __init__(self, config, low_rank_svd_approximation=False, top_k_svd=-1, device="cpu"): | |
| self.config = config | |
| self.stats_functions = dict(config.stats) | |
| self.w_bins = config.w_bins | |
| self.sv_bins = config.sv_bins | |
| self.use_density = config.use_density | |
| self.device = torch.device(device) | |
| # SVD configuration | |
| self.low_rank_svd_approximation = low_rank_svd_approximation | |
| self.top_k_svd = top_k_svd | |
| # initialize data | |
| self.data = { | |
| weight_type: {"weight_type": weight_type} | |
| for weight_type in config.weight_type | |
| } | |
| def analyze_head(self, head): | |
| W_Q_h, W_K_h, W_QK_h = head["W_Q"], head["W_K"], head["W_QK"] | |
| self.fill_WW(W_Q_h, W_K_h, W_QK_h) | |
| if "W_Q_gram" in self.data and "W_Q_gram" in head: | |
| self.fill_gram("W_Q_gram", head["W_Q_gram"]) | |
| if "W_K_gram" in self.data and "W_K_gram" in head: | |
| self.fill_gram("W_K_gram", head["W_K_gram"]) | |
| if "QK_alignment" in self.data and "QK_alignment" in head: | |
| self.fill_alignment("QK_alignment", head["QK_alignment"]) | |
| # tensors are for a given head, thus are matrices | |
| def fill_WW(self, W_Q_h, W_K_h, W_QK, ov=False): | |
| W_Q_key, W_K_key, W_QK_key = "W_Q", "W_K", "W_QK" | |
| if ov: | |
| W_Q_key, W_K_key, W_QK_key = "W_O", "W_V", "W_OV" | |
| W_Q_vector = W_Q_h.flatten().detach().cpu().numpy() | |
| self.fill_vector(W_Q_key, W_Q_vector) | |
| W_K_vector = W_K_h.flatten().detach().cpu().numpy() | |
| self.fill_vector(W_K_key, W_K_vector) | |
| self.fill_matrix(W_QK_key, W_QK) | |
| def fill_stats(self, weight_name, x_arr): | |
| self.data[weight_name].update( | |
| {k: f(x_arr) for k, f in self.stats_functions.items()} | |
| ) | |
| def fill_scalar(self, weight_name, v): | |
| self.data[weight_name].update({weight_name: v}) | |
| def fill_vector(self, weight_name, x_arr, histo=True, copy=False): | |
| if histo: | |
| h, _ = np.histogram(x_arr, bins=self.w_bins, density=self.use_density) | |
| self.data[weight_name].update({"P_w": h}) | |
| if copy: | |
| self.data[weight_name].update({"x": x_arr.to_numpy()}) | |
| self.fill_stats(weight_name, x_arr) | |
| def fill_gram(self, weight_name, W_gram_tensor): | |
| x_arr = W_gram_tensor.flatten().detach().cpu().numpy() | |
| self.fill_vector(weight_name, x_arr, histo=True, copy=False) | |
| try: | |
| W_gpu = W_gram_tensor.to(self.device) | |
| # gram eigenvalues = σᵢ(W)²; take sqrt to store σᵢ(W) | |
| sv2 = torch.linalg.svdvals(W_gpu) | |
| svd = torch.sqrt(sv2.clamp(min=0)).detach().cpu().numpy() | |
| self.data[weight_name].update({"SVD": svd}) | |
| P_sv, _ = np.histogram(svd, bins=self.sv_bins, density=self.use_density) | |
| self.data[weight_name].update({"P_sv": P_sv}) | |
| except Exception as e: | |
| print(f"Warning: SVD computation failed for {weight_name}: {e}") | |
| self.data[weight_name].update({"SVD": None, "P_sv": None}) | |
| def fill_alignment(self, weight_name, cosines_arr): | |
| # cosines_arr: numpy (d_head,), principal-angle cosines between W_Q and W_K col-spaces | |
| self.fill_stats(weight_name, cosines_arr) | |
| self.data[weight_name].update({"SVD": cosines_arr, "P_sv": None}) | |
| def fill_matrix(self, weight_name, W_tensor): | |
| x_arr = W_tensor.flatten().detach().cpu().numpy() | |
| self.fill_vector(weight_name, x_arr, histo=True, copy=False) | |
| try: | |
| W_gpu = W_tensor.to(self.device) | |
| if self.low_rank_svd_approximation: | |
| _, S, _ = torch.svd_lowrank(W_gpu, q=self.top_k_svd) | |
| d = W_gpu.shape[0] | |
| if len(S) < d: | |
| S_padded = torch.zeros(d, dtype=S.dtype, device=S.device) | |
| S_padded[:len(S)] = S | |
| S = S_padded | |
| else: | |
| _, S, _ = torch.linalg.svd(W_gpu) | |
| svd = S.detach().cpu().numpy() | |
| self.data[weight_name].update({"SVD": svd}) | |
| P_sv, _ = np.histogram(svd, bins=self.sv_bins, density=self.use_density) | |
| self.data[weight_name].update({"P_sv": P_sv}) | |
| except (RuntimeError, Exception) as e: | |
| print(f"Warning: SVD computation failed for {weight_name}: {e}") | |
| self.data[weight_name].update({"SVD": None, "P_sv": None}) | |
| def to_pandas(self): | |
| df = pd.DataFrame([v for v in self.data.values()]) | |
| return df | |
| class LayerHeadContainer: | |
| def __init__(self, layer_idx, config, low_rank_svd_approximation=False, top_k_svd=-1, device="cpu"): | |
| self.layer_idx = layer_idx | |
| self.config = config | |
| self.n_heads = config.n_heads | |
| self.head_dim = config.head_dim | |
| self.d_model = config.d_model | |
| self.device = device | |
| # SVD configuration | |
| self.low_rank_svd_approximation = low_rank_svd_approximation | |
| if low_rank_svd_approximation and top_k_svd == -1: | |
| self.top_k_svd = self.head_dim | |
| else: | |
| self.top_k_svd = top_k_svd | |
| # Create HeadAnalyzer instances with SVD configuration | |
| self.data = [ | |
| HeadAnalyzer(config, low_rank_svd_approximation=self.low_rank_svd_approximation, | |
| top_k_svd=self.top_k_svd, device=device) | |
| for _ in range(self.n_heads) | |
| ] | |
| def analyze_layer(self, input_dict): | |
| # expected shape for W is n_heads, d_head, d_model | |
| weight_types = set(self.config.weight_type) | |
| W_Q_h = input_dict["W_Q"] | |
| W_K_h = input_dict["W_K"] | |
| W_QK_all = torch.bmm( | |
| W_Q_h, # (n_heads, head_dim, d_model) | |
| W_K_h.transpose(1, 2) # (n_heads, d_model, head_dim) | |
| ) # Result: (n_heads, head_dim, head_dim) | |
| W_QK_gpu = W_QK_all.to(self.device) | |
| compute_grams = "W_Q_gram" in weight_types or "W_K_gram" in weight_types | |
| compute_alignment = "QK_alignment" in weight_types | |
| if compute_grams: | |
| W_Q_gram_all = torch.bmm(W_Q_h, W_Q_h.transpose(1, 2)).to(self.device) | |
| W_K_gram_all = torch.bmm(W_K_h, W_K_h.transpose(1, 2)).to(self.device) | |
| if compute_alignment: | |
| _, _, Vh_q = torch.linalg.svd(W_Q_h.to(self.device), full_matrices=False) | |
| _, _, Vh_k = torch.linalg.svd(W_K_h.to(self.device), full_matrices=False) | |
| M_all = torch.bmm(Vh_q, Vh_k.transpose(1, 2)) # (n_heads, d_head, d_head) | |
| cosines_all = torch.linalg.svdvals(M_all).clamp(0, 1).detach().cpu().numpy() | |
| for head_idx in tqdm(range(self.n_heads), desc=f" Layer {self.layer_idx} heads", leave=False): | |
| head_data = { | |
| "W_Q": W_Q_h[head_idx], | |
| "W_K": W_K_h[head_idx], | |
| "W_QK": W_QK_gpu[head_idx], | |
| } | |
| if compute_grams: | |
| head_data["W_Q_gram"] = W_Q_gram_all[head_idx] | |
| head_data["W_K_gram"] = W_K_gram_all[head_idx] | |
| if compute_alignment: | |
| head_data["QK_alignment"] = cosines_all[head_idx] | |
| self.data[head_idx].analyze_head(head_data) | |
| def post_process(self, weight_metrics=None, sv_metrics=None): | |
| """ | |
| Post-process analysis by computing additional metrics. | |
| Args: | |
| weight_metrics: Dictionary of metric functions for weight histograms. | |
| Each function should have signature f(h, centers). | |
| If None, uses normality_metrics. | |
| sv_metrics: Dictionary of metric functions for singular values. | |
| Each function should have signature f(h, svd_array). | |
| If None, uses singular_value_metrics. | |
| """ | |
| if weight_metrics is None: | |
| from transformer_analysis.histogram_utils import normality_metrics | |
| weight_metrics = normality_metrics | |
| if sv_metrics is None: | |
| from transformer_analysis.histogram_utils import singular_value_metrics | |
| sv_metrics = singular_value_metrics | |
| for head in self.data: # loop on heads | |
| centers = (head.w_bins[:-1] + head.w_bins[1:]) / 2 | |
| for h in head.data.values(): # loop on weights associated with head | |
| if h.get("P_w") is not None: | |
| for f_m in weight_metrics.values(): | |
| f_m(h, centers) | |
| if "SVD" in h and h["SVD"] is not None: | |
| svd_array = h["SVD"] | |
| for f_m in sv_metrics.values(): | |
| f_m(h, svd_array) | |
| def to_pandas(self): | |
| df_list = [] | |
| for head_idx in range(self.n_heads): | |
| head_df = self.data[head_idx].to_pandas() | |
| head_df["head"] = head_idx | |
| df_list.append(head_df) | |
| df = pd.concat(df_list, ignore_index=True) | |
| df["layer"] = self.layer_idx | |
| return df | |
| if __name__ == "__main__": | |
| from types import SimpleNamespace | |
| import numpy as np | |
| import torch | |
| test_single_head = False | |
| test_layer = True | |
| config = SimpleNamespace() | |
| config.weight_type = ["W_Q", "W_K", "W_QK", "W_Q_gram", "W_K_gram", "QK_alignment"] | |
| config.stats = {"mean": np.mean, "std": np.std} | |
| config.w_bins = np.linspace( | |
| -2, 2, 201 | |
| ) # low number of bins for easy visual inspection | |
| config.use_density = False | |
| config.n_heads = 32 | |
| config.d_model = 1024 | |
| config.head_dim = 12 | |
| # testing a single head | |
| if test_single_head: | |
| print("Testing single head functionality") | |
| W_Q = torch.randn(config.head_dim, config.d_model) | |
| W_K = torch.randn(config.head_dim, config.d_model) | |
| W_QK = torch.randn(config.d_model, config.d_model) | |
| ha = HeadAnalyzer(config) | |
| head_data = {"W_Q": W_Q, "W_K": W_K, "W_QK": W_QK} | |
| ha.analyze_head(head_data) | |
| df = ha.to_pandas() | |
| print("Mean and variance should be consistent with N(0,1)") | |
| print(df[["mean", "std"]]) | |
| if test_layer: | |
| layer_idx = 37 # random choice | |
| layer = LayerHeadContainer(layer_idx, config) | |
| W_Q = torch.randn(config.n_heads, config.head_dim, config.d_model) | |
| W_K = torch.randn(config.n_heads, config.head_dim, config.d_model) | |
| layer_input = {"W_Q": W_Q, "W_K": W_K} | |
| layer.analyze_layer(layer_input) | |
| df = layer.to_pandas() | |
| print(df.columns) | |
| print(df[["layer", "head", "weight_type", "std", "P_w"]]) | |
| import matplotlib.pyplot as plt | |
| plt.plot(df["SVD"][2]) | |
| # plt.show() | |
| plt.savefig("test.png", dpi=150, bbox_inches="tight") | |
| plt.close() | |
| # %% | |