app.py

"""
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()