import os import sys import json import torch import numpy as np import matplotlib matplotlib.use('Agg') import matplotlib.pyplot as plt from datasets import load_dataset from transformers import AutoModelForCausalLM, AutoTokenizer def main(): device = "cuda" if torch.cuda.is_available() else "cpu" print(f"[Device] Using device: {device}") # 1. Load model and tokenizer print("[Model] Loading model and tokenizer from HF...") model_id = "SRJ5035/sw_glu_sw_64_16_8_4_xpert_gpt" tokenizer = AutoTokenizer.from_pretrained(model_id, revision="main") model = AutoModelForCausalLM.from_pretrained(model_id, revision="main", trust_remote_code=True).eval().to(device) # 2. Recreate exact validation split print("[Data] Loading dataset and building validation split...") ds = load_dataset("BabyLM-community/BabyLM-2026-Strict-Small") all_text = list(ds['train']['text']) split = int(len(all_text) * 0.95) val_texts = all_text[split:] print(f"[Data] Loaded {len(val_texts)} validation sentences.") # Tokenize print("[Data] Tokenizing validation set...") all_ids = [] for t in val_texts: if t.strip(): encoded = tokenizer.encode(t, add_special_tokens=False) all_ids.extend(encoded) total_tokens = len(all_ids) print(f"[Data] Total tokens in validation set: {total_tokens:,}") B, T = 16, 512 chunk_size = B * T # 8192 num_batches = total_tokens // chunk_size print(f"[Data] Generated {num_batches} batches of size {B}x{T} ({chunk_size} tokens each).") # 3. Setup hooks to capture routing decisions blocks = model.transformer.blocks num_layers = len(blocks) num_blocks = blocks[0].num_blocks # 4 print(f"[Model] Found {num_layers} layers, each with {num_blocks} MoE experts.") # Dict to store statistics per layer stats = {} for l in range(num_layers): stats[l] = { "expert_counts": np.zeros(5, dtype=np.int64), "cooc_matrix": np.zeros((num_blocks, num_blocks), dtype=np.int64), "total_tokens_seen": 0 } # Hook implementation def get_routing_hook(layer_idx): def hook(module, input, output): # topk_token_idx is of shape (num_blocks, k_capacity) topk_token_idx = module.last_topk_indices device = topk_token_idx.device k_capacity = topk_token_idx.size(1) # Since input[0] is x_0 of shape (B, T, D) B_batch, T_batch, D_batch = input[0].size() n_tokens = B_batch * T_batch # Count experts per token token_expert_counts = torch.zeros(n_tokens, dtype=torch.long, device=device) for i in range(num_blocks): token_expert_counts.scatter_add_(0, topk_token_idx[i], torch.ones_like(topk_token_idx[i])) counts = torch.bincount(token_expert_counts, minlength=5).cpu().numpy() stats[layer_idx]["expert_counts"] += counts # Co-occurrence matrix M = torch.zeros(n_tokens, num_blocks, dtype=torch.float, device=device) for i in range(num_blocks): M[topk_token_idx[i], i] = 1.0 cooc = (M.T @ M).cpu().numpy().astype(np.int64) stats[layer_idx]["cooc_matrix"] += cooc stats[layer_idx]["total_tokens_seen"] += n_tokens return hook # Register hooks hooks = [] for l in range(num_layers): h = blocks[l].register_forward_hook(get_routing_hook(l)) hooks.append(h) # 4. Run forward passes print("[Eval] Running forward passes over the validation set...") for idx in range(num_batches): batch_ids = all_ids[idx * chunk_size : (idx + 1) * chunk_size] input_ids = torch.tensor(batch_ids, dtype=torch.long).view(B, T).to(device) with torch.no_grad(): model(input_ids) if (idx + 1) % 10 == 0 or (idx + 1) == num_batches: print(f"[Eval] Processed batch {idx + 1}/{num_batches}") # Remove hooks for h in hooks: h.remove() print("[Eval] Completed forward passes. Aggregating statistics...") # 5. Analysis and Calculation global_expert_counts = np.zeros(5, dtype=np.int64) global_cooc_matrix = np.zeros((num_blocks, num_blocks), dtype=np.int64) global_total_tokens = 0 for l in range(num_layers): global_expert_counts += stats[l]["expert_counts"] global_cooc_matrix += stats[l]["cooc_matrix"] global_total_tokens += stats[l]["total_tokens_seen"] # Convert counts to fractions global_expert_frac = global_expert_counts / global_total_tokens # Calculate mean experts per token total_expert_assignments = sum(i * global_expert_counts[i] for i in range(5)) mean_experts_per_token = total_expert_assignments / global_total_tokens # Co-occurrence fraction global_cooc_frac = global_cooc_matrix / global_total_tokens print("\n" + "="*50) print("ROUTER UTILIZATION SUMMARY") print("="*50) print(f"Overall fraction of tokens with 0 experts: {global_expert_frac[0]:.4%} ({global_expert_counts[0]:,} tokens)") print(f"Overall fraction of tokens with >=1 experts: {sum(global_expert_frac[1:]):.4%}") print(f"Overall fraction of tokens with >=3 experts: {sum(global_expert_frac[3:]):.4%}") print(f"Mean experts per token (Theoretical expect: ~2.0): {mean_experts_per_token:.4f}") print("="*50 + "\n") # Save stats to json os.makedirs("./figures", exist_ok=True) raw_stats = { "global": { "expert_counts": global_expert_counts.tolist(), "expert_fractions": global_expert_frac.tolist(), "mean_experts_per_token": mean_experts_per_token, "cooccurrence_matrix": global_cooc_matrix.tolist(), "cooccurrence_fractions": global_cooc_frac.tolist() }, "per_layer": {} } for l in range(num_layers): layer_total = stats[l]["total_tokens_seen"] layer_frac = stats[l]["expert_counts"] / layer_total layer_cooc_frac = stats[l]["cooc_matrix"] / layer_total raw_stats["per_layer"][f"layer_{l+1}"] = { "expert_counts": stats[l]["expert_counts"].tolist(), "expert_fractions": layer_frac.tolist(), "cooccurrence_matrix": stats[l]["cooc_matrix"].tolist(), "cooccurrence_fractions": layer_cooc_frac.tolist() } with open("router_stats.json", "w") as f: json.dump(raw_stats, f, indent=2) print("[Stats] Saved raw statistics to 'router_stats.json'.") # 6. PLOT GENERATION # Setup publication-ready matplotlib style plt.rcParams.update({ "font.family": "serif", "font.serif": ["Times New Roman", "DejaVu Serif", "Liberation Serif", "serif"], "axes.labelsize": 10, "font.size": 10, "legend.fontsize": 8, "xtick.labelsize": 8, "ytick.labelsize": 8, "figure.titlesize": 11 }) # 6a. Plot 1: Histogram of expert choices fig, ax = plt.subplots(figsize=(4.5, 3.5)) x_vals = np.arange(5) bars = ax.bar(x_vals, global_expert_frac * 100, color='#4f46e5', width=0.6, edgecolor='none') ax.set_xlabel("Number of Experts Selected") ax.set_ylabel("Fraction of Tokens (%)") ax.set_xticks(x_vals) ax.set_xticklabels([str(i) for i in x_vals]) ax.spines['top'].set_visible(False) ax.spines['right'].set_visible(False) # Annotate bars for 0 and >=3 experts for i, bar in enumerate(bars): val = global_expert_frac[i] * 100 if i == 0 or i >= 3: ax.annotate(f"{val:.2f}%", xy=(bar.get_x() + bar.get_width() / 2, val), xytext=(0, 3), textcoords="offset points", ha='center', va='bottom', fontsize=8, fontweight='bold') plt.tight_layout() plt.savefig("./figures/router_util_histogram.pdf", dpi=300) plt.savefig("./figures/router_util_histogram.png", dpi=300) plt.close() print("[Plots] Generated Plot 1: router_util_histogram.pdf/.png") # 6b. Plot 2: Stacked bar chart by layer fig, ax = plt.subplots(figsize=(5.5, 3.5)) layer_indices = np.arange(1, num_layers + 1) stacked_data = np.zeros((5, num_layers)) for l in range(num_layers): l_total = stats[l]["total_tokens_seen"] stacked_data[:, l] = stats[l]["expert_counts"] / l_total colors = ['#ef4444', '#f59e0b', '#10b981', '#3b82f6', '#8b5cf6'] bottoms = np.zeros(num_layers) for i in range(5): ax.bar(layer_indices, stacked_data[i], bottom=bottoms, label=f"{i} experts", color=colors[i], width=0.5) bottoms += stacked_data[i] ax.set_xlabel("Layer Index") ax.set_ylabel("Fraction of Tokens") ax.set_xticks(layer_indices) ax.set_xticklabels([str(i) for i in layer_indices]) ax.spines['top'].set_visible(False) ax.spines['right'].set_visible(False) ax.legend(bbox_to_anchor=(1.02, 1), loc='upper left', frameon=False) plt.tight_layout() plt.savefig("./figures/router_util_by_layer.pdf", dpi=300) plt.savefig("./figures/router_util_by_layer.png", dpi=300) plt.close() print("[Plots] Generated Plot 2: router_util_by_layer.pdf/.png") # 6c. Plot 3: Heatmap of expert co-occurrence fig, ax = plt.subplots(figsize=(4.0, 3.5)) im = ax.imshow(global_cooc_frac * 100, cmap='Blues', vmin=0, vmax=100) cbar = fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04) cbar.ax.set_ylabel("Co-occurrence Fraction (%)", rotation=-90, va="bottom", fontsize=8) cbar.ax.tick_params(labelsize=8) ax.set_xticks(np.arange(num_blocks)) ax.set_yticks(np.arange(num_blocks)) ax.set_xticklabels([f"Exp {i+1}" for i in range(num_blocks)]) ax.set_yticklabels([f"Exp {i+1}" for i in range(num_blocks)]) ax.set_xlabel("Expert Index") ax.set_ylabel("Expert Index") for i in range(num_blocks): for j in range(num_blocks): val = global_cooc_frac[i, j] * 100 text_color = "white" if val > 50 else "black" ax.text(j, i, f"{val:.2f}%", ha="center", va="center", color=text_color, fontsize=8, fontweight='bold') plt.tight_layout() plt.savefig("./figures/expert_cooccurrence_heatmap.pdf", dpi=300) plt.savefig("./figures/expert_cooccurrence_heatmap.png", dpi=300) plt.close() print("[Plots] Generated Plot 3: expert_cooccurrence_heatmap.pdf/.png") if __name__ == "__main__": main()