""" Expert Co-activation Geometry — ZeroGPU 60s 适配版 ==================================================== 核心假设: MoE router 是否自发学到互补性？ co-activated expert pair 的 head_sim < 随机 pair 的 head_sim 执行策略（适配 60s 限制）: Step 1 (CPU) : 预计算并缓存各层的 64×64 head_sim 矩阵 Step 2 (GPU) : 仅做 forward pass + 捕获 routing 索引，存到全局 Step 3 (CPU) : 读取缓存矩阵 + routing 结果，做全部统计分析 Hook 位置（已通过 transformers 5.3.0 源码验证）: layer.mlp.gate.forward 返回 (router_logits, router_scores, router_indices) router_indices shape: [seq_len, top_k=8] → hook output[2] 即为每 token 选中的 expert 索引 """ import torch import numpy as np import json, os, gc from itertools import combinations from scipy import stats import gradio as gr import spaces from huggingface_hub import hf_hub_download from safetensors.torch import load_file from transformers import AutoTokenizer, OlmoeForCausalLM from datasets import load_dataset # ══════════════════════════════════════════════════════════════════ # 全局配置 # ══════════════════════════════════════════════════════════════════ REPO_ID = "allenai/OLMoE-1B-7B-0924" N_LAYERS = 16 N_EXPERTS = 64 TOP_K = 8 HEAD_RATIO = 0.01 PROJ = "down_proj" CACHE_DIR = "/data/olmoe_cache" os.makedirs(CACHE_DIR, exist_ok=True) # 跨调用的全局 routing 结果 _routing_results: dict = {} # ══════════════════════════════════════════════════════════════════ # 工具函数（CPU） # ══════════════════════════════════════════════════════════════════ def get_shard_for_key(key: str): index_path = hf_hub_download( repo_id=REPO_ID, filename="model.safetensors.index.json" ) with open(index_path) as f: index = json.load(f) return index["weight_map"].get(key) def load_expert_weight(layer_idx: int, expert_idx: int, proj: str = PROJ): key = f"model.layers.{layer_idx}.mlp.experts.{expert_idx}.{proj}.weight" shard = get_shard_for_key(key) if shard is None: return None fp = hf_hub_download(repo_id=REPO_ID, filename=shard) sd = load_file(fp, device="cpu") W = sd[key].clone() del sd return W def head_sim_from_U(U1, U2, k_head: int) -> float: b1 = U1[:, :k_head] b2 = U2[:, :k_head] return torch.linalg.svdvals(b1.T @ b2).mean().item() # ══════════════════════════════════════════════════════════════════ # Step 1: 预计算 head_sim 矩阵（CPU） # ══════════════════════════════════════════════════════════════════ def build_headsim_matrix(layer_idx: int, proj: str = PROJ): cache_path = os.path.join(CACHE_DIR, f"headsim_L{layer_idx}_{proj}.npy") if os.path.exists(cache_path): print(f"[Layer {layer_idx}] Cache hit.") return np.load(cache_path) print(f"[Layer {layer_idx}] Loading {N_EXPERTS} expert weights...") Us, r_val = [], None for e in range(N_EXPERTS): W = load_expert_weight(layer_idx, e, proj) if W is None: return None W = W.to(torch.float32) if r_val is None: r_val = min(W.shape) U, _, _ = torch.linalg.svd(W, full_matrices=False) Us.append(U) del W k_head = max(1, int(r_val * HEAD_RATIO)) print(f"[Layer {layer_idx}] r={r_val}, k_head={k_head}. Computing pairs...") mat = np.eye(N_EXPERTS, dtype=np.float32) for i, j in combinations(range(N_EXPERTS), 2): sim = head_sim_from_U(Us[i], Us[j], k_head) mat[i, j] = sim mat[j, i] = sim np.save(cache_path, mat) print(f"[Layer {layer_idx}] Saved.") return mat def step1_run(layers_str: str, proj: str, progress=gr.Progress()) -> str: try: layers = [int(x.strip()) for x in layers_str.split(",") if x.strip().isdigit()] except Exception: return "❌ Invalid layer list." results = [] for i, l in enumerate(layers): progress(i / len(layers), desc=f"Layer {l}...") mat = build_headsim_matrix(l, proj) if mat is None: results.append(f"Layer {l}: ❌ Failed") continue od = mat[np.triu_indices(N_EXPERTS, k=1)] results.append( f"Layer {l}: ✅ mean={od.mean():.5f} std={od.std():.5f} " f"min={od.min():.5f} max={od.max():.5f}" ) return "\n".join(results) # ══════════════════════════════════════════════════════════════════ # Step 2: Forward pass（GPU）— 只抓 routing 索引 # ══════════════════════════════════════════════════════════════════ @spaces.GPU(duration=60) def step2_forward_pass(n_tokens: int, layers_str: str) -> str: global _routing_results _routing_results.clear() try: layers = [int(x.strip()) for x in layers_str.split(",") if x.strip().isdigit()] except Exception: return "❌ Invalid layer list." # 加载模型 print("Loading tokenizer + model...") tokenizer = AutoTokenizer.from_pretrained(REPO_ID) model = OlmoeForCausalLM.from_pretrained( REPO_ID, torch_dtype=torch.float16, low_cpu_mem_usage=True, ).cuda() model.eval() # 准备文本 print("Preparing text...") dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="test", trust_remote_code=True) text = " ".join(t for t in dataset["text"][:100] if t.strip()) tokens = tokenizer(text, return_tensors="pt", truncation=True, max_length=n_tokens) input_ids = tokens["input_ids"].to(model.device) actual_len = input_ids.shape[1] print(f"Tokens: {actual_len}") # 注册 hooks # gate.forward → (router_logits, router_scores, router_indices) # router_indices: [seq_len, top_k=8] raw_log = {l: [] for l in layers} hooks = [] for layer_idx in layers: gate = model.model.layers[layer_idx].mlp.gate def make_hook(lidx): def hook(module, inp, output): # output[2] = router_indices indices = output[2].detach().cpu() raw_log[lidx].append(indices) return hook h = gate.register_forward_hook(make_hook(layer_idx)) hooks.append(h) # Forward print("Forward pass...") with torch.no_grad(): _ = model(input_ids) for h in hooks: h.remove() # 合并结果到全局 for l in layers: if raw_log[l]: _routing_results[l] = torch.cat(raw_log[l], dim=0).numpy() print(f"Layer {l}: {_routing_results[l].shape[0]} tokens captured") else: print(f"Layer {l}: ⚠️ No data!") del model gc.collect() torch.cuda.empty_cache() captured = [l for l in layers if l in _routing_results] not_captured = [l for l in layers if l not in _routing_results] msg = f"✅ Forward pass done. Tokens: {actual_len}\n" msg += f"Layers captured: {captured}\n" if not_captured: msg += f"⚠️ Missing layers: {not_captured}\n" msg += "\nRun Step 3 to compute statistics." return msg # ══════════════════════════════════════════════════════════════════ # Step 3: 统计分析（纯 CPU） # ══════════════════════════════════════════════════════════════════ def step3_statistics(layers_str: str, proj: str) -> str: try: layers = [int(x.strip()) for x in layers_str.split(",") if x.strip().isdigit()] except Exception: return "❌ Invalid layer list." missing_mat = [l for l in layers if not os.path.exists( os.path.join(CACHE_DIR, f"headsim_L{l}_{proj}.npy"))] missing_rout = [l for l in layers if l not in _routing_results] if missing_mat: return f"❌ head_sim matrices missing for layers {missing_mat}. Run Step 1." if missing_rout: return f"❌ Routing data missing for layers {missing_rout}. Run Step 2." lines = [ "=" * 82, "Expert Co-activation Geometry — Statistical Results", "=" * 82, "H1: co-activated pairs have LOWER head_sim than random (complementarity)", "", f"{'Layer':>5} | {'CoAct μ':>9} | {'Rand μ':>9} | " f"{'Δ':>9} | {'z':>7} | {'p (1-tail)':>11} | Verdict", "─" * 82, ] for l in layers: mat = np.load(os.path.join(CACHE_DIR, f"headsim_L{l}_{proj}.npy")) decisions = _routing_results[l] # [n_tokens, 8] # co-activated sims coact_sims = [] for row in decisions: experts = [int(e) for e in row] for i, j in combinations(experts, 2): coact_sims.append(float(mat[i, j])) coact_sims = np.array(coact_sims) # random baseline off_diag = mat[np.triu_indices(N_EXPERTS, k=1)] rand_mean = off_diag.mean() rand_std = off_diag.std() coact_mean = coact_sims.mean() n = len(coact_sims) # z-test, left-tail z = (coact_mean - rand_mean) / (rand_std / np.sqrt(n)) p = float(stats.norm.cdf(z)) delta = coact_mean - rand_mean if p < 0.001 and delta < 0: verdict = "✅ COMPLEMENTARY p<0.001" elif p < 0.05 and delta < 0: verdict = "✅ COMPLEMENTARY p<0.05" elif p > 0.95 and delta > 0: verdict = "❌ REDUNDANT" else: verdict = "➖ NO SIGNAL" lines.append( f"{l:>5} | {coact_mean:>9.5f} | {rand_mean:>9.5f} | " f"{delta:>+9.5f} | {z:>7.2f} | {p:>11.6f} | {verdict}" ) lines.append( f" | CoAct p10={np.percentile(coact_sims,10):.5f} " f"p50={np.percentile(coact_sims,50):.5f} " f"p90={np.percentile(coact_sims,90):.5f} n={n}" ) lines.append( f" | Random p10={np.percentile(off_diag,10):.5f} " f"p50={np.percentile(off_diag,50):.5f} " f"p90={np.percentile(off_diag,90):.5f} " f"n_tokens={decisions.shape[0]}" ) lines.append("") lines += [ "=" * 82, "Δ < 0 & p < 0.05 → Router learned complementarity (research hypothesis confirmed)", "Δ > 0 & p > 0.95 → Router selects redundant experts", "|Δ| ≈ 0 → No geometric preference in routing", "=" * 82, ] return "\n".join(lines) # ══════════════════════════════════════════════════════════════════ # Matrix Summary（随时可查，CPU） # ══════════════════════════════════════════════════════════════════ def matrix_summary(layers_str: str, proj: str) -> str: try: layers = [int(x.strip()) for x in layers_str.split(",") if x.strip().isdigit()] except Exception: return "❌ Invalid layer list." lines = ["head_sim Matrix Summary", "=" * 50] for l in layers: p = os.path.join(CACHE_DIR, f"headsim_L{l}_{proj}.npy") if not os.path.exists(p): lines.append(f"\nLayer {l}: ❌ Not computed") continue mat = np.load(p) od = mat[np.triu_indices(N_EXPERTS, k=1)] upper = np.triu(mat, k=1) idx_max = np.unravel_index(upper.argmax(), mat.shape) upper_masked = np.where(upper > 0, upper, np.inf) idx_min = np.unravel_index(upper_masked.argmin(), mat.shape) lines += [ f"\nLayer {l}:", f" mean={od.mean():.5f} std={od.std():.5f}", f" min={od.min():.5f} max={od.max():.5f}", f" p10={np.percentile(od,10):.5f} " f"p50={np.percentile(od,50):.5f} " f"p90={np.percentile(od,90):.5f}", f" Most similar: E{idx_max[0]:02d} vs E{idx_max[1]:02d} " f"sim={mat[idx_max]:.5f}", f" Most different: E{idx_min[0]:02d} vs E{idx_min[1]:02d} " f"sim={mat[idx_min]:.5f}", ] return "\n".join(lines) # ══════════════════════════════════════════════════════════════════ # Gradio UI # ══════════════════════════════════════════════════════════════════ with gr.Blocks(title="Expert Co-activation Geometry") as demo: gr.Markdown(""" # 🔬 Expert Co-activation Geometry **Research Question**: Does OLMoE's router *systematically* co-activate geometrically *complementary* expert pairs (low head_sim)? | Step | Hardware | What it does | |------|----------|--------------| | 1 | CPU | Precompute & cache 64×64 head_sim matrices | | 2 | GPU (≤60s) | Forward pass on WikiText-2, capture routing indices | | 3 | CPU | Statistical analysis: co-activated vs random | """) with gr.Tabs(): with gr.Tab("⚙️ Step 1 · head_sim Matrices (CPU)"): gr.Markdown("Each layer ≈ 3–5 min. Cached to disk — runs once only.") with gr.Row(): s1_layers = gr.Textbox(value="0,7,15", label="Layers (0–15, comma-separated)") s1_proj = gr.Dropdown(["down_proj","up_proj","gate_proj"], value="down_proj", label="Projection") s1_btn = gr.Button("▶ Compute Matrices", variant="primary") s1_out = gr.Textbox(label="Output", lines=10) s1_btn.click(fn=step1_run, inputs=[s1_layers, s1_proj], outputs=[s1_out]) with gr.Tab("🚀 Step 2 · Forward Pass (GPU ≤60s)"): gr.Markdown( "Loads OLMoE, runs WikiText-2 forward pass, captures routing. \n" "Keep tokens ≤ 128 to stay within 60 s. \n" "Results stored in memory → run Step 3 immediately after." ) with gr.Row(): s2_layers = gr.Textbox(value="0,7,15", label="Layers to hook") s2_ntokens = gr.Slider(64, 256, value=128, step=64, label="Tokens (≤128 recommended for 60s)") s2_btn = gr.Button("▶ Run Forward Pass", variant="primary") s2_out = gr.Textbox(label="Status", lines=10) s2_btn.click(fn=step2_forward_pass, inputs=[s2_ntokens, s2_layers], outputs=[s2_out]) with gr.Tab("📊 Step 3 · Statistics (CPU)"): gr.Markdown("Run immediately after Step 2 (routing data lives in memory).") with gr.Row(): s3_layers = gr.Textbox(value="0,7,15", label="Layers") s3_proj = gr.Dropdown(["down_proj","up_proj","gate_proj"], value="down_proj", label="Projection") s3_btn = gr.Button("▶ Analyse", variant="primary") s3_out = gr.Textbox(label="Results", lines=35) s3_btn.click(fn=step3_statistics, inputs=[s3_layers, s3_proj], outputs=[s3_out]) with gr.Tab("🗂️ Matrix Summary"): gr.Markdown("Inspect cached matrices — no GPU needed.") with gr.Row(): sm_layers = gr.Textbox(value="0,7,15", label="Layers") sm_proj = gr.Dropdown(["down_proj","up_proj","gate_proj"], value="down_proj", label="Projection") sm_btn = gr.Button("▶ Show Summary", variant="secondary") sm_out = gr.Textbox(label="Summary", lines=25) sm_btn.click(fn=matrix_summary, inputs=[sm_layers, sm_proj], outputs=[sm_out]) demo.launch()