import os import copy import time import random import itertools import numpy as np import jax.lax as lax import jax.numpy as jnp from utils import rngmix import matplotlib.pyplot as plt from typing import NamedTuple from collections import defaultdict from flax.core import freeze, unfreeze from scipy.optimize import linear_sum_assignment, minimize from jax import random, tree_util, jit, grad, value_and_grad def compute_objective(A, X, X_prime, Y, Y_prime): A_inv = np.linalg.inv(A) term1 = X - X_prime @ A.T term2 = Y - Y_prime @ A_inv return np.sum(term1**2) + np.sum(term2**2) def compute_gradient(A, X, X_prime, Y, Y_prime): A_inv = np.linalg.inv(A) term1 = -2 * X.T @ X_prime + 2 * A @ X_prime.T @ X_prime term2 = 2 * A_inv.T @ Y_prime.T @ (Y - Y_prime @ A_inv) @ A_inv.T return term1 + term2 def line_search(A, grad, X, X_prime, Y, Y_prime, max_step=1, tau=0.5, c1=1e-4): eta = max_step f_current = compute_objective(A, X, X_prime, Y, Y_prime) grad_norm2 = np.sum(grad**2) n = A.shape[0] while eta > 1e-10: A_new = A - eta * grad if np.linalg.matrix_rank(A_new) < n: eta *= tau continue f_new = compute_objective(A_new, X, X_prime, Y, Y_prime) if f_new <= f_current - c1 * eta * grad_norm2: return eta eta *= tau return 0 @jit def compute_objective_jax(A, X, X_prime, Y, Y_prime, cond_threshold=1e6): cond = jnp.linalg.cond(A) def safe_obj(): A_inv = jnp.linalg.inv(A) term1 = X - X_prime @ A.T term2 = Y - Y_prime @ A_inv return jnp.sum(term1**2) + jnp.sum(term2**2) return lax.cond(cond > cond_threshold, lambda: jnp.inf, safe_obj) compute_value_and_grad_jax = jit(value_and_grad(compute_objective_jax)) def solve_orthogonal(X, X_prime, Y, Y_prime): B = X.T @ X_prime + Y.T @ Y_prime U, _, Vt = np.linalg.svd(B) return U @ Vt def solve_rope(X, X_prime, Y, Y_prime,max_iters=200, tol=1e-16): d = X.shape[1] assert d % 2 == 0, "d must be even." assert X.shape[1] == X_prime.shape[1] == Y.shape[1] == Y_prime.shape[1] def rot(theta): c, s = np.cos(theta), np.sin(theta) return np.array([[c, -s], [s, c]]) def block_cols(j): return [2*j, 2*j+1] def solve_block(Q1blk, Q2blk, K1blk, K2blk): A, B, Ah, Bh = Q1blk, Q2blk, K1blk, K2blk a, ah = np.sum(A*A), np.sum(Ah*Ah) c_const = np.sum(B*B) + np.sum(Bh*Bh) C, Ch = A.T @ B, Ah.T @ Bh t_tr, s_sk = np.trace(C), C[0,1] - C[1,0] th_tr, sh_sk = np.trace(Ch), Ch[0,1] - Ch[1,0] u, v, w = t_tr**2+s_sk**2, th_tr**2+sh_sk**2, t_tr*th_tr+s_sk*sh_sk eps = 1e-18 def phi(t): return max(u*t + v/max(t,eps) + 2*w, 0.0) def gprime(t): denom = np.sqrt(phi(t)) if denom < eps: return a - ah/(t*t) return (a - ah/(t*t)) - (u - v/(t*t)) / denom t0 = np.sqrt((ah+eps)/(a+eps)) t_lo, gp_lo = t0, gprime(t0) if gp_lo < 0.0: t_hi = t_lo for _ in range(max_iters): t_hi *= 2.0 if gprime(t_hi) >= 0.0: break else: t_hi = t_lo for _ in range(max_iters): t_lo *= 0.5 if gprime(t_lo) <= 0.0: break def gval(t): return a*t + ah/max(t,eps) + c_const - 2*np.sqrt(phi(t)) if not (gprime(t_lo) <= 0.0 <= gprime(t_hi)): t_star = min([(t_lo,gval(t_lo)),(t_hi,gval(t_hi))], key=lambda z:z[1])[0] else: for _ in range(max_iters): t_mid = 0.5*(t_lo+t_hi) gp_mid = gprime(t_mid) if abs(gp_mid) < tol or (t_hi-t_lo) <= tol*(1+t_mid): t_star = t_mid; break if gp_mid < 0.0: t_lo = t_mid else: t_hi = t_mid else: t_star = 0.5*(t_lo+t_hi) rho = np.sqrt(max(t_star, eps)) alpha = rho*t_tr + (1/rho)*th_tr beta = rho*s_sk + (1/rho)*sh_sk theta = np.arctan2(beta, alpha) return rho, theta P = np.zeros((d, d)) for j in range(d//2): cols = block_cols(j) rho, theta = solve_block(X[:,cols], X_prime[:,cols], Y[:,cols], Y_prime[:,cols]) P[np.ix_(cols, cols)] = rho * rot(theta) return P def optimize_alignment(A_init, X, X_prime, Y, Y_prime, max_iter=5000): objective_values = [] grad_norms = [] condition_nums = [] def obj_fn(flat_A): A = flat_A.reshape(A_init.shape) obj, grad_val = compute_value_and_grad_jax(jnp.array(A), jnp.array(X), jnp.array(X_prime), jnp.array(Y), jnp.array(Y_prime)) return float(obj), np.array(grad_val).flatten() def callback(flat_A): A = flat_A.reshape(A_init.shape) obj, grad_val = compute_value_and_grad_jax(jnp.array(A), jnp.array(X), jnp.array(X_prime), jnp.array(Y), jnp.array(Y_prime)) grad_norm = jnp.linalg.norm(grad_val, 'fro') cond = jnp.linalg.cond(jnp.array(A)) objective_values.append(float(obj)) grad_norms.append(float(grad_norm)) condition_nums.append(float(cond)) res = minimize(obj_fn, A_init.flatten(), jac=True, method='L-BFGS-B', options={'maxiter': max_iter}, callback=callback) A_opt = res.x.reshape(A_init.shape) return A_opt, objective_values, grad_norms, condition_nums def extract_attention_params(attn): c_attn_kernel = np.array(attn['c_attn']['kernel']) c_attn_bias = np.array(attn['c_attn']['bias']) c_proj_kernel = np.array(attn['c_proj']['kernel']) c_proj_bias = np.array(attn['c_proj']['bias']) query, key, value = np.split(c_attn_kernel, 3, axis=0) query_bias, key_bias, value_bias = np.split(c_attn_bias, 3, axis=0) return query, key, value, query_bias, key_bias, value_bias, c_proj_kernel, c_proj_bias def reshape_attention_weights(query, key, value, query_bias, key_bias, value_bias, out_kernel, num_heads): D = query.shape[1] D_k = D_v = D // num_heads def stack_per_head(tensor, axis=0): return np.stack([ tensor[i * D_k:(i + 1) * D_k, :].T if axis == 0 else tensor[:, i * D_k:(i + 1) * D_k].T for i in range(num_heads) ]) def stack_bias_per_head(bias): return np.stack([bias[i * D_k:(i + 1) * D_k].T for i in range(num_heads)]) W_Q = stack_per_head(query) W_K = stack_per_head(key) W_V = stack_per_head(value) W_O = stack_per_head(out_kernel, axis=1) b_Q = stack_bias_per_head(query_bias) b_K = stack_bias_per_head(key_bias) b_V = stack_bias_per_head(value_bias) return W_Q, b_Q, W_K, b_K, W_V, b_V, W_O def compute_extended_weights(W, b): return np.vstack([W, b.reshape(1, -1)]) def compute_cost_matrix(W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a, W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b, h, activations, alpha=0.5): C = np.zeros((h, h)) for i in range(h): tilde_W_Q_a_i = np.vstack([W_Q_a[i], b_Q_a[i].reshape(1, -1)]) tilde_W_K_a_i = np.vstack([W_K_a[i], b_K_a[i].reshape(1, -1)]) tilde_W_V_a_i = np.vstack([W_V_a[i], b_V_a[i].reshape(1, -1)]) QKT_a_i = tilde_W_Q_a_i @ tilde_W_K_a_i.T VO_a_i = tilde_W_V_a_i @ W_O_a[i] centered_QKT_a_i = QKT_a_i - np.mean(QKT_a_i, axis=1, keepdims=True) for j in range(h): tilde_W_Q_b_j = np.vstack([W_Q_b[j], b_Q_b[j].reshape(1, -1)]) tilde_W_K_b_j = np.vstack([W_K_b[j], b_K_b[j].reshape(1, -1)]) tilde_W_V_b_j = np.vstack([W_V_b[j], b_V_b[j].reshape(1, -1)]) QKT_b_j = tilde_W_Q_b_j @ tilde_W_K_b_j.T VO_b_j = tilde_W_V_b_j @ W_O_b[j] centered_QKT_b_j = QKT_b_j - np.mean(QKT_b_j, axis=1, keepdims=True) cost = 0.5 * np.sum((centered_QKT_a_i - centered_QKT_b_j) ** 2) cost += 0.5 * np.sum((VO_a_i - VO_b_j) ** 2) C[i, j] = cost return C def additive_align_single_head(W_Q_a_i, b_Q_a_i, W_K_a_i, b_K_a_i, W_V_a_i, b_V_a_i, W_O_a_i, W_Q_b_i, b_Q_b_i, W_K_b_i, b_K_b_i, W_V_b_i, b_V_b_i, W_O_b_i, optimize): tilde_W_Q_a_i = compute_extended_weights(W_Q_a_i, b_Q_a_i) tilde_W_K_a_i = compute_extended_weights(W_K_a_i, b_K_a_i) tilde_W_V_a_i = compute_extended_weights(W_V_a_i, b_V_a_i) Y_O_a_i = W_O_a_i.T tilde_W_Q_b_i = compute_extended_weights(W_Q_b_i, b_Q_b_i) tilde_W_K_b_i = compute_extended_weights(W_K_b_i, b_K_b_i) tilde_W_V_b_i = compute_extended_weights(W_V_b_i, b_V_b_i) Y_O_b_i = W_O_b_i.T A_init = solve_orthogonal(tilde_W_Q_a_i, tilde_W_Q_b_i, tilde_W_K_a_i, tilde_W_K_b_i) B_init = solve_orthogonal(Y_O_a_i, Y_O_b_i, tilde_W_V_a_i, tilde_W_V_b_i) if optimize: A, objective_values_A, grad_norms_A, condition_nums_A = optimize_alignment( A_init, tilde_W_Q_a_i, tilde_W_Q_b_i, tilde_W_K_a_i, tilde_W_K_b_i ) B, objective_values_B, grad_norms_B, condition_nums_B = optimize_alignment( B_init, Y_O_a_i, Y_O_b_i, tilde_W_V_a_i, tilde_W_V_b_i ) else: A = A_init B = B_init A_inv = np.linalg.inv(A) B_inv = np.linalg.inv(B) W_Q_aligned = W_Q_b_i @ A.T b_Q_aligned = b_Q_b_i @ A.T W_K_aligned = W_K_b_i @ A_inv b_K_aligned = b_K_b_i @ A_inv W_V_aligned = W_V_b_i @ B_inv b_V_aligned = b_V_b_i @ B_inv W_O_aligned = B @ W_O_b_i aligned_params = { 'query': {'kernel': W_Q_aligned, 'bias': b_Q_aligned}, 'key': {'kernel': W_K_aligned, 'bias': b_K_aligned}, 'value': {'kernel': W_V_aligned, 'bias': b_V_aligned}, 'out': {'kernel': W_O_aligned} } if optimize: return { 'aligned_params': aligned_params, 'metrics_A': { 'objective_values': objective_values_A, 'grad_norms': grad_norms_A, 'condition_nums': condition_nums_A }, 'metrics_B': { 'objective_values': objective_values_B, 'grad_norms': grad_norms_B, 'condition_nums': condition_nums_B } } return {'aligned_params': aligned_params} def _fro2(x): if x.ndim == 1: # vector -> Euclidean norm return float(np.linalg.norm(x)**2) else: # matrix -> Frobenius norm return float(np.linalg.norm(x, 'fro')**2) def rope_apply_alignment(W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b, W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a, h): aligned_params = {} total_pre = 0.0 total_post = 0.0 # --- Pairwise totals --- total_qk_pre = 0.0 total_qk_post = 0.0 total_vo_pre = 0.0 total_vo_post = 0.0 for i in range(h): # ===== Build augmented (kernel+bias row) for A (Q,K) ===== tilde_W_Q_a_i = np.vstack([W_Q_a[i], b_Q_a[i].reshape(1, -1)]) tilde_W_K_a_i = np.vstack([W_K_a[i], b_K_a[i].reshape(1, -1)]) # ===== Build augmented (kernel+bias row) for B (Q,K) ===== tilde_W_Q_b_i = np.vstack([W_Q_b[i], b_Q_b[i].reshape(1, -1)]) tilde_W_K_b_i = np.vstack([W_K_b[i], b_K_b[i].reshape(1, -1)]) # ===== Solve A_i for (Q,K) pair ===== A_i_init = solve_rope(tilde_W_Q_a_i, tilde_W_Q_b_i, tilde_W_K_a_i, tilde_W_K_b_i) A_i = A_i_init # (optionally run a refinement step) # ===== Build augmented (kernel+bias row) for (V,O) pair ===== tilde_W_V_a_i = np.vstack([W_V_a[i], b_V_a[i].reshape(1, -1)]) tilde_W_V_b_i = np.vstack([W_V_b[i], b_V_b[i].reshape(1, -1)]) # For O we transpose and pad with one zero row to match augmented shape Y_O = W_O_a[i].T Y_O_prime = W_O_b[i].T Y_O_padded = np.vstack([Y_O, np.zeros((1, Y_O.shape[1]))]) Y_O_prime_padded = np.vstack([Y_O_prime, np.zeros((1, Y_O_prime.shape[1]))]) # ===== Solve B_i for (V,O) pair ===== B_i_init = solve_rope(tilde_W_V_a_i, tilde_W_V_b_i, Y_O_padded, Y_O_prime_padded) B_i = B_i_init # (optionally run a refinement step) # ===== Apply transforms ===== A_i_inv = np.linalg.inv(A_i) B_i_inv = np.linalg.inv(B_i) # W_Q_aligned = W_Q_b[i] @ A_i.T # b_Q_aligned = b_Q_b[i] @ A_i.T # W_K_aligned = W_K_b[i] @ A_i_inv # b_K_aligned = b_K_b[i] @ A_i_inv # W_V_aligned = W_V_b[i] @ B_i_inv # b_V_aligned = b_V_b[i] @ B_i_inv # W_O_aligned = B_i @ W_O_b[i] W_Q_aligned = W_Q_b[i] @ A_i_inv.T # was A_i.T -> FIX: A_i^{-T} b_Q_aligned = b_Q_b[i] @ A_i_inv.T W_K_aligned = W_K_b[i] @ A_i # was A_i_inv -> FIX: A_i b_K_aligned = b_K_b[i] @ A_i # --- V,O pair: use P for V (right-multiply), and P^{-1} for O (left-multiply) --- W_V_aligned = W_V_b[i] @ B_i # was B_i_inv -> FIX: B_i b_V_aligned = b_V_b[i] @ B_i W_O_aligned = B_i_inv @ W_O_b[i] # was B_i @ W_O_b[i] -> FIX: B_i^{-1} on the left aligned_params[f'head_{i}'] = { 'query': {'kernel': W_Q_aligned, 'bias': b_Q_aligned}, 'key': {'kernel': W_K_aligned, 'bias': b_K_aligned}, 'value': {'kernel': W_V_aligned, 'bias': b_V_aligned}, 'out': {'kernel': W_O_aligned} } # ===== Frobenius^2 BEFORE (a vs raw b) ===== pre_q = _fro2(W_Q_a[i] - W_Q_b[i]) + _fro2(b_Q_a[i] - b_Q_b[i]) pre_k = _fro2(W_K_a[i] - W_K_b[i]) + _fro2(b_K_a[i] - b_K_b[i]) pre_v = _fro2(W_V_a[i] - W_V_b[i]) + _fro2(b_V_a[i] - b_V_b[i]) pre_o = _fro2(W_O_a[i] - W_O_b[i]) # O has no bias in your structure # ===== Frobenius^2 AFTER (a vs aligned b) ===== post_q = _fro2(W_Q_a[i] - W_Q_aligned) + _fro2(b_Q_a[i] - b_Q_aligned) post_k = _fro2(W_K_a[i] - W_K_aligned) + _fro2(b_K_a[i] - b_K_aligned) post_v = _fro2(W_V_a[i] - W_V_aligned) + _fro2(b_V_a[i] - b_V_aligned) post_o = _fro2(W_O_a[i] - W_O_aligned) pre_sum = pre_q + pre_k + pre_v + pre_o post_sum = post_q + post_k + post_v + post_o total_pre += pre_sum total_post += post_sum # ===== Pairwise sums ===== pre_qk = pre_q + pre_k post_qk = post_q + post_k pre_vo = pre_v + pre_o post_vo = post_v + post_o total_qk_pre += pre_qk total_qk_post += post_qk total_vo_pre += pre_vo total_vo_post += post_vo # ===== Per-head print ===== print(f"[Head {i}] Fro^2 pre={pre_sum:.6f} post={post_sum:.6f} improve={pre_sum - post_sum:.6f}") print(f" Q: pre={pre_q:.6f} post={post_q:.6f}") print(f" K: pre={pre_k:.6f} post={post_k:.6f}") print(f" V: pre={pre_v:.6f} post={post_v:.6f}") print(f" O: pre={pre_o:.6f} post={post_o:.6f}") # --- New: pairwise breakdowns --- print(f" [Q,K] pair: pre={pre_qk:.6f} post={post_qk:.6f} improve={pre_qk - post_qk:.6f}") print(f" [V,O] pair: pre={pre_vo:.6f} post={post_vo:.6f} improve={pre_vo - post_vo:.6f}") # ===== Totals ===== print("=== Frobenius^2 (including biases) ===") print(f"Total pre : {total_pre:.6f}") print(f"Total post: {total_post:.6f}") print(f"Total improvement: {total_pre - total_post:.6f} ({0.0 if total_pre==0 else 100.0*(total_pre-total_post)/total_pre:.2f}%)") # --- New: Pairwise totals --- print("=== Pairwise Frobenius^2 (including biases) ===") print(f"[Q,K] total pre : {total_qk_pre:.6f}") print(f"[Q,K] total post: {total_qk_post:.6f}") print(f"[Q,K] improvement: {total_qk_pre - total_qk_post:.6f} ({0.0 if total_qk_pre==0 else 100.0*(total_qk_pre-total_qk_post)/total_qk_pre:.2f}%)") print(f"[V, O] total pre : {total_vo_pre:.6f}") print(f"[V, O] total post: {total_vo_post:.6f}") print(f"[V, O] improvement: {total_vo_pre - total_vo_post:.6f} ({0.0 if total_vo_pre==0 else 100.0*(total_vo_pre-total_vo_post)/total_vo_pre:.2f}%)") return aligned_params def merge_aligned_params(aligned_params, h, D, out_bias_b): query_kernel = np.stack([aligned_params[f'head_{i}']['query']['kernel'] for i in range(h)], axis=1) query_bias = np.stack([aligned_params[f'head_{i}']['query']['bias'] for i in range(h)], axis=0) key_kernel = np.stack([aligned_params[f'head_{i}']['key']['kernel'] for i in range(h)], axis=1) key_bias = np.stack([aligned_params[f'head_{i}']['key']['bias'] for i in range(h)], axis=0) value_kernel = np.stack([aligned_params[f'head_{i}']['value']['kernel'] for i in range(h)], axis=1) value_bias = np.stack([aligned_params[f'head_{i}']['value']['bias'] for i in range(h)], axis=0) out_kernel = np.stack([aligned_params[f'head_{i}']['out']['kernel'] for i in range(h)], axis=0) query_kernel = query_kernel.transpose(1, 2, 0).reshape(-1, D) key_kernel = key_kernel.transpose(1, 2, 0).reshape(-1, D) value_kernel = value_kernel.transpose(1, 2, 0).reshape(-1, D) out_kernel = out_kernel.transpose(2, 0, 1).reshape(D, -1) query_bias = query_bias.reshape(-1) key_bias = key_bias.reshape(-1) value_bias = value_bias.reshape(-1) return { 'c_attn': { 'kernel': jnp.array(np.concatenate([query_kernel, key_kernel, value_kernel], axis=0)), 'bias': jnp.array(np.concatenate([query_bias, key_bias, value_bias], axis=0)), }, 'c_proj': {'kernel': jnp.array(out_kernel),'bias': jnp.array(out_bias_b),} } def align_attention_params(rng, params_a, params_b, layer_idx, config, activation, permute_heads=True, optimize=False, alpha=0.5): num_heads = config.lmc_config.n_head attn_a = params_a['transformer']['h'][str(layer_idx)]['attn'] attn_b = params_b['transformer']['h'][str(layer_idx)]['attn'] query_a, key_a, value_a, query_bias_a, key_bias_a, value_bias_a, out_a, out_bias_a = extract_attention_params(attn_a) query_b, key_b, value_b, query_bias_b, key_bias_b, value_bias_b, out_b, out_bias_b = extract_attention_params(attn_b) W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a = reshape_attention_weights(query_a, key_a, value_a, query_bias_a, key_bias_a, value_bias_a, out_a, num_heads) W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b = reshape_attention_weights(query_b, key_b, value_b, query_bias_b, key_bias_b, value_bias_b, out_b, num_heads) if permute_heads: C = compute_cost_matrix(W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a, W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b, num_heads, activation, alpha) row_ind, col_ind = linear_sum_assignment(C) print("Best Permutation Heads:", col_ind) W_Q_b = [W_Q_b[j] for j in col_ind] b_Q_b = [b_Q_b[j] for j in col_ind] W_K_b = [W_K_b[j] for j in col_ind] b_K_b = [b_K_b[j] for j in col_ind] W_V_b = [W_V_b[j] for j in col_ind] b_V_b = [b_V_b[j] for j in col_ind] W_O_b = [W_O_b[j] for j in col_ind] if optimize: metrics_A_all = {key: [] for key in ['objective_values', 'grad_norms', 'condition_nums']} metrics_B_all = {key: [] for key in ['objective_values', 'grad_norms', 'condition_nums']} aligned_params, return_dict = {}, {} if(config.position_embeddings in ["learnable","sinusoidal"]): for i in range(num_heads): result = additive_align_single_head( W_Q_a[i], b_Q_a[i], W_K_a[i], b_K_a[i], W_V_a[i], b_V_a[i], W_O_a[i], W_Q_b[i], b_Q_b[i], W_K_b[i], b_K_b[i], W_V_b[i], b_V_b[i], W_O_b[i], optimize ) aligned_params[f'head_{i}'] = result['aligned_params'] if optimize: for key in metrics_A_all: metrics_A_all[key].append(result['metrics_A'][key]) metrics_B_all[key].append(result['metrics_B'][key]) return_dict['aligned_params'] = merge_aligned_params(aligned_params, num_heads, query_a.shape[1], out_bias_b) if optimize: return_dict['metrics_A_all'] = metrics_A_all return_dict['metrics_B_all'] = metrics_B_all elif(config.position_embeddings in ["rope"]): aligned_params = rope_apply_alignment(W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b, W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a, num_heads) return return_dict def permute_align_attention_params(rng, params_a, params_b, layer_idx, config,col_ind): num_heads = config.lmc_config.n_head attn_a = params_a['transformer']['h'][layer_idx]['attn'] attn_b = params_b['transformer']['h'][layer_idx]['attn'] query_a, key_a, value_a, query_bias_a, key_bias_a, value_bias_a, out_a, out_bias_a = extract_attention_params(attn_a) query_b, key_b, value_b, query_bias_b, key_bias_b, value_bias_b, out_b, out_bias_b = extract_attention_params(attn_b) W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a = reshape_attention_weights(query_a, key_a, value_a, query_bias_a, key_bias_a, value_bias_a, out_a, num_heads) W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b = reshape_attention_weights(query_b, key_b, value_b, query_bias_b, key_bias_b, value_bias_b, out_b, num_heads) W_Q_b = [W_Q_b[j] for j in col_ind] b_Q_b = [b_Q_b[j] for j in col_ind] W_K_b = [W_K_b[j] for j in col_ind] b_K_b = [b_K_b[j] for j in col_ind] W_V_b = [W_V_b[j] for j in col_ind] b_V_b = [b_V_b[j] for j in col_ind] W_O_b = [W_O_b[j] for j in col_ind] if(config.position_embeddings in ["learnable","sinusoidal"]): aligned_params = additive_apply_alignment(W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b, W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a, num_heads) elif(config.position_embeddings in ["rope"]): aligned_params = rope_apply_alignment(W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b, W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a, num_heads) return merge_aligned_params(aligned_params, num_heads, query_a.shape[1], out_bias_b) def naive_align_attention_params(rng, params_a, params_b, layer_idx, config): num_heads = config.lmc_config.n_head attn_a = params_a['transformer']['h'][layer_idx]['attn'] attn_b = params_b['transformer']['h'][layer_idx]['attn'] query_a, key_a, value_a, query_bias_a, key_bias_a, value_bias_a, out_a, out_bias_a = extract_attention_params(attn_a) query_b, key_b, value_b, query_bias_b, key_bias_b, value_bias_b, out_b, out_bias_b = extract_attention_params(attn_b) W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a = reshape_attention_weights(query_a, key_a, value_a, query_bias_a, key_bias_a, value_bias_a, out_a, num_heads) W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b = reshape_attention_weights(query_b, key_b, value_b, query_bias_b, key_bias_b, value_bias_b, out_b, num_heads) if(config.position_embeddings in ["learnable","sinusoidal"]): aligned_params = additive_apply_alignment(W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b, W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a, num_heads) elif(config.position_embeddings in ["rope"]): aligned_params = rope_apply_alignment(W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b, W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a, num_heads) return merge_aligned_params(aligned_params, num_heads, query_a.shape[1], out_bias_b) def all_matching_attn(rng, params_a, params_b, config): results = {} permutations = list(itertools.permutations(range(config.lmc_config.n_head))) if config.lmc_config.n_head > 4: permutations = random.sample(permutations, 24) for perm in permutations: print("Permutation",perm) temp_params = copy.deepcopy(params_b) for layer_idx in config.lmc_layer_indices: aligned_attention_params = permute_align_attention_params(rng, params_a, params_b, str(layer_idx), config, perm) temp_params['transformer']['h'][str(layer_idx)]['attn'] = aligned_attention_params results[str(perm)] = temp_params return results def weight_matching_attn(rng, params_a, params_b, activation, config): params_dict = {} configurations = [ ("permu_head_init_ortho_no_opt", 'ortho', True, False), ("permu_head_init_ortho_opt", 'ortho', True, True), # ("naive_head_init_ortho_no_opt", 'ortho', False, False), # ("naive_head_init_ortho_opt", 'ortho', False, True), ] for name, init_method, permute_heads, optimize in configurations: aligned_params = copy.deepcopy(params_b) if optimize: layer_to_metrics_A = {} layer_to_metrics_B = {} for layer_idx in config.lmc_layer_indices: if activation is not None: activations_for_layer = activation[layer_idx] else: activations_for_layer = None result = align_attention_params( rng, params_a, aligned_params, layer_idx, config, activations_for_layer, permute_heads=permute_heads, optimize=optimize ) aligned_params['transformer']['h'][str(layer_idx)]['attn'] = result['aligned_params'] if optimize: layer_to_metrics_A[layer_idx] = result['metrics_A_all'] layer_to_metrics_B[layer_idx] = result['metrics_B_all'] total_sum = tree_util.tree_reduce(lambda acc, x: acc + jnp.sum(x), aligned_params, initializer=0) print(f"{name}: {total_sum}, sanity check") params_dict[name] = aligned_params return params_dict # cost_head = copy.deepcopy(params_b) # naive_head = copy.deepcopy(params_b) # for layer_idx in config.lmc_layer_indices: # aligned_attention_params = cost_align_attention_params(rng, params_a, params_b, str(layer_idx), config) # cost_head['transformer']['h'][str(layer_idx)]['attn'] = aligned_attention_params # for layer_idx in config.lmc_layer_indices: # aligned_attention_params = naive_align_attention_params(rng, params_a, params_b, str(layer_idx), config) # naive_head['transformer']['h'][str(layer_idx)]['attn'] = aligned_attention_params # return {"cost_head": cost_head, "naive_head": naive_head}