app.py

"""
Mixture-of-Experts (MoE): Explainer + Interactive Routing Simulator
--------------------------------------------------------------------
A Hugging Face Space (Gradio) that:
  1. Explains how sparse Mixture-of-Experts layers work (gating, top-k
     routing, load balancing, capacity, sparse compute savings).
  2. Lets the user run a live simulation: tokens are embedded, a gating
     network scores them against experts, top-k experts are chosen per
     token (optionally with noisy/capacity-limited routing like
     Switch Transformer / Mixtral / GShard), and the results are
     visualized (routing map, expert load, gate-weight histogram,
     FLOPs comparison).

No trained model is required -- everything is a lightweight numpy
simulation, so the Space is instant to load and free to run on CPU.
"""

import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import gradio as gr

# --------------------------------------------------------------------
# Optional ZeroGPU compatibility shim
# --------------------------------------------------------------------
# This app is pure CPU (numpy/matplotlib) and never needs a GPU. If the
# Space's hardware is set to "CPU basic" (recommended), the `spaces`
# package won't even be installed, and everything below is skipped.
#
# If someone instead deploys this on a "ZeroGPU" hardware tier, that
# runtime requires at least one function decorated with @spaces.GPU to
# be present at startup, or it throws a runtime error. This shim
# satisfies that requirement with a trivial no-op, without ever
# actually consuming GPU quota (the simulation itself never calls it).
try:
    import spaces

    @spaces.GPU
    def _zerogpu_startup_check():
        """No-op so ZeroGPU Spaces detect a valid @spaces.GPU function
        at startup. Never called by the app itself -- this Space does
        not use a GPU for anything."""
        return True

except ImportError:
    # `spaces` isn't installed -> we're on CPU basic (or running
    # locally). Nothing to do.
    pass

# --------------------------------------------------------------------
# Core MoE routing simulation
# --------------------------------------------------------------------

RNG_SEED = 0


def simulate_moe(
    num_tokens: int,
    embed_dim: int,
    num_experts: int,
    top_k: int,
    noise_std: float,
    capacity_factor: float,
    seed: int,
):
    """Simulate one forward pass of a sparse MoE gating layer.

    Returns a dict with routing assignments, gate weights, dropped
    tokens (due to capacity), and everything needed to plot.
    """
    rng = np.random.default_rng(seed)

    # 1. Fake token embeddings (in a real model these come from the
    #    previous transformer block's hidden states).
    #    We give tokens soft "cluster" structure so routing isn't pure
    #    noise -- similar to how semantically similar tokens tend to
    #    prefer the same experts in trained MoEs.
    num_clusters = min(num_experts, max(2, num_experts // 2))
    cluster_centers = rng.normal(0, 3.0, size=(num_clusters, embed_dim))
    cluster_ids = rng.integers(0, num_clusters, size=num_tokens)
    tokens = cluster_centers[cluster_ids] + rng.normal(0, 1.0, size=(num_tokens, embed_dim))

    # 2. Gating network: a single linear layer (token_dim -> num_experts)
    W_gate = rng.normal(0, 1.0 / np.sqrt(embed_dim), size=(embed_dim, num_experts))
    logits = tokens @ W_gate

    # 3. Noisy gating (as in Shazeer et al. 2017 / Switch Transformer)
    #    helps exploration & load balancing during training.
    if noise_std > 0:
        logits = logits + rng.normal(0, noise_std, size=logits.shape)

    gate_probs = _softmax(logits, axis=-1)

    # 4. Top-k expert selection per token
    topk_idx = np.argsort(-gate_probs, axis=-1)[:, :top_k]
    topk_val = np.take_along_axis(gate_probs, topk_idx, axis=-1)
    # re-normalize the chosen top-k weights so they sum to 1 (standard practice)
    topk_val = topk_val / topk_val.sum(axis=-1, keepdims=True)

    # 5. Expert capacity (Switch/GShard style): each expert can only
    #    process `capacity` tokens per batch; overflow tokens are
    #    dropped (their residual/skip connection carries them instead).
    capacity = max(1, int(np.ceil(capacity_factor * num_tokens * top_k / num_experts)))
    expert_fill = np.zeros(num_experts, dtype=int)
    kept_mask = np.ones_like(topk_idx, dtype=bool)

    # process tokens in order, first-come-first-served per expert
    for t in range(num_tokens):
        for k in range(top_k):
            e = topk_idx[t, k]
            if expert_fill[e] < capacity:
                expert_fill[e] += 1
            else:
                kept_mask[t, k] = False  # dropped: over capacity

    # 6. Load per expert (tokens actually routed, i.e. not dropped)
    load = np.zeros(num_experts, dtype=int)
    for t in range(num_tokens):
        for k in range(top_k):
            if kept_mask[t, k]:
                load[topk_idx[t, k]] += 1

    dropped_tokens = int((~kept_mask).sum())
    total_routes = num_tokens * top_k

    # 7. Load-balancing auxiliary loss (Switch Transformer formulation)
    #    L_aux = num_experts * sum_e( f_e * P_e )
    #    f_e = fraction of tokens dispatched to expert e (top-1 proxy)
    top1 = topk_idx[:, 0]
    f = np.bincount(top1, minlength=num_experts) / num_tokens
    P = gate_probs.mean(axis=0)
    aux_loss = num_experts * np.sum(f * P)

    return dict(
        tokens=tokens,
        cluster_ids=cluster_ids,
        gate_probs=gate_probs,
        topk_idx=topk_idx,
        topk_val=topk_val,
        kept_mask=kept_mask,
        load=load,
        capacity=capacity,
        dropped_tokens=dropped_tokens,
        total_routes=total_routes,
        aux_loss=aux_loss,
        num_experts=num_experts,
        top_k=top_k,
        num_tokens=num_tokens,
    )


def _softmax(x, axis=-1):
    x = x - np.max(x, axis=axis, keepdims=True)
    e = np.exp(x)
    return e / np.sum(e, axis=axis, keepdims=True)


# --------------------------------------------------------------------
# FLOPs / compute-savings estimate (dense vs sparse MoE)
# --------------------------------------------------------------------

def compute_flops_comparison(num_experts, top_k, ffn_mult=4, embed_dim=1024):
    """Rough FFN FLOPs-per-token estimate: dense model uses ALL experts'
    worth of parameters for every token; MoE only activates top_k.
    Assumes each 'expert' is sized like one dense FFN block."""
    per_expert_flops = 2 * embed_dim * (embed_dim * ffn_mult) * 2  # up+down proj, x2 for mul-add
    dense_equiv_flops = per_expert_flops * num_experts  # a dense FFN with same total capacity
    moe_flops = per_expert_flops * top_k
    speedup = dense_equiv_flops / moe_flops
    return dense_equiv_flops, moe_flops, speedup


# --------------------------------------------------------------------
# Plotting
# --------------------------------------------------------------------

EXPERT_CMAP = plt.get_cmap("tab20")


def plot_routing_map(sim):
    tokens, topk_idx, kept_mask = sim["tokens"], sim["topk_idx"], sim["kept_mask"]
    # 2D projection via first 2 dims (embed space already low-dim for viz clarity)
    if tokens.shape[1] > 2:
        # simple PCA-free projection: just take first two coordinates
        xy = tokens[:, :2]
    else:
        xy = tokens

    fig, ax = plt.subplots(figsize=(6, 5))
    top1 = topk_idx[:, 0]
    top1_kept = kept_mask[:, 0]

    for e in range(sim["num_experts"]):
        mask = (top1 == e) & top1_kept
        ax.scatter(xy[mask, 0], xy[mask, 1], s=35, color=EXPERT_CMAP(e % 20),
                   label=f"Expert {e}", alpha=0.85, edgecolors="white", linewidths=0.3)

    dropped = ~top1_kept
    if dropped.any():
        ax.scatter(xy[dropped, 0], xy[dropped, 1], s=45, facecolors="none",
                   edgecolors="red", linewidths=1.4, label="Dropped (over capacity)")

    ax.set_title("Token routing map (top-1 expert, colored)", fontsize=11)
    ax.set_xlabel("embedding dim 1")
    ax.set_ylabel("embedding dim 2")
    ax.legend(loc="upper right", fontsize=7, ncol=2, framealpha=0.9)
    fig.tight_layout()
    return fig


def plot_expert_load(sim):
    load = sim["load"]
    capacity = sim["capacity"]
    fig, ax = plt.subplots(figsize=(6, 4))
    colors = [EXPERT_CMAP(e % 20) for e in range(sim["num_experts"])]
    bars = ax.bar(range(sim["num_experts"]), load, color=colors, edgecolor="black", linewidth=0.4)
    ax.axhline(capacity, color="red", linestyle="--", linewidth=1.3, label=f"Capacity = {capacity}")
    ax.set_xlabel("Expert index")
    ax.set_ylabel("Tokens routed")
    ax.set_title("Expert load balance", fontsize=11)
    ax.set_xticks(range(sim["num_experts"]))
    ax.legend(fontsize=8)
    fig.tight_layout()
    return fig


def plot_gate_weights(sim):
    vals = sim["topk_val"].flatten()
    fig, ax = plt.subplots(figsize=(6, 4))
    ax.hist(vals, bins=25, color="#6C5CE7", edgecolor="white", alpha=0.9)
    ax.set_xlabel("Normalized gate weight (top-k, per token)")
    ax.set_ylabel("Count")
    ax.set_title("Gate-weight distribution", fontsize=11)
    fig.tight_layout()
    return fig


# --------------------------------------------------------------------
# Gradio callback
# --------------------------------------------------------------------

def run_simulation(num_tokens, embed_dim, num_experts, top_k, noise_std,
                    capacity_factor, ffn_mult, seed):
    top_k = min(top_k, num_experts)
    sim = simulate_moe(
        num_tokens=int(num_tokens),
        embed_dim=int(embed_dim),
        num_experts=int(num_experts),
        top_k=int(top_k),
        noise_std=float(noise_std),
        capacity_factor=float(capacity_factor),
        seed=int(seed),
    )

    dense_flops, moe_flops, speedup = compute_flops_comparison(
        int(num_experts), int(top_k), ffn_mult=int(ffn_mult), embed_dim=int(embed_dim)
    )

    drop_pct = 100.0 * sim["dropped_tokens"] / sim["total_routes"]
    load = sim["load"]
    load_std = float(np.std(load))
    load_mean = float(np.mean(load))
    cv = load_std / load_mean if load_mean > 0 else 0.0

    summary = f"""
### Routing summary

| Metric | Value |
|---|---|
| Tokens | {sim['num_tokens']} |
| Experts | {sim['num_experts']} |
| Top-k | {sim['top_k']} |
| Expert capacity (per expert) | {sim['capacity']} |
| Dropped routes (over capacity) | {sim['dropped_tokens']} / {sim['total_routes']}  ({drop_pct:.1f}%) |
| Load balance (coefficient of variation, lower=better) | {cv:.3f} |
| Load-balancing aux loss (Switch-style) | {sim['aux_loss']:.4f} |

### Compute comparison (FFN block, per token)

| Model | FFN FLOPs/token | Relative |
|---|---|---|
| Dense (uses all {num_experts} experts) | {dense_flops:,.0f} | 1.0x |
| Sparse MoE (top-{top_k} of {num_experts}) | {moe_flops:,.0f} | {1/speedup:.2f}x |

**→ The sparse MoE layer does ~{speedup:.1f}x less FFN compute per token than an
equivalently-sized dense layer**, while keeping (in a trained model) comparable
total parameter capacity — this is MoE's core value proposition.
"""

    return (
        plot_routing_map(sim),
        plot_expert_load(sim),
        plot_gate_weights(sim),
        summary,
    )


# --------------------------------------------------------------------
# Explanatory markdown
# --------------------------------------------------------------------

INTRO_MD = """
# 🧩 Mixture of Experts (MoE) — Explainer & Interactive Simulator

**Mixture of Experts** is a neural network architecture where, instead of every
input passing through one big dense feed-forward block, a layer is split into
many smaller **expert** sub-networks, and a lightweight **gating (router)
network** decides — *per token* — which one or few experts should process it.

This is how models like **Mixtral 8x7B**, **DeepSeek-MoE**, **Switch
Transformer**, and **GShard** scale to huge parameter counts while keeping
inference compute roughly constant: a model might have hundreds of billions
of *total* parameters, but only activate a few billion of them for any given
token.

## How it works

1. **Experts.** Each MoE layer replaces one dense FFN with `N` independent
   FFN "experts" (`E_0 … E_{N-1}`), each with its own weights.
2. **Gating network.** A small linear layer scores each token against every
   expert: `logits = x @ W_gate`, then `softmax` turns scores into a
   probability distribution over experts.
3. **Top-k routing.** Only the top-`k` highest-scoring experts are actually
   run for that token (commonly k=1 for Switch Transformer, k=2 for Mixtral/
   GShard). The token's output is the weighted sum of the chosen experts'
   outputs, weighted by their (renormalized) gate probabilities.
4. **Capacity limits.** To keep training/inference batched and efficient,
   each expert is only allowed to process a fixed number of tokens per batch
   (`capacity`). Tokens that overflow a full expert are **dropped** for that
   layer (they skip via the residual connection instead).
5. **Load balancing.** If left unconstrained, gating tends to collapse onto
   a few "favorite" experts. An **auxiliary load-balancing loss** penalizes
   uneven routing so all experts stay utilized and trained.

## Why it matters

- **Sparse activation** → you get the representational capacity of a huge
  dense model, but the *compute cost per token* of a much smaller one.
- **Specialization** → different experts can implicitly specialize
  (e.g. by language, syntax pattern, or topic), though in practice this
  specialization is soft and not perfectly interpretable.
- **The catch** → routing imbalance, dropped tokens, communication overhead
  (in distributed training experts often live on different devices), and
  training instability are the main engineering challenges.

---

## 🎛️ Try it yourself below

Tune the simulation parameters and watch: which expert each token gets
routed to, how balanced the load is across experts, the gate-weight
distribution, and the theoretical FLOPs savings vs. an equivalent dense
layer.
"""

FOOTER_MD = """
---
### Notes on this simulation
- Token embeddings are **synthetic** (sampled from a few Gaussian clusters)
  purely to make routing patterns visually interpretable — no real language
  model is involved.
- The gating network is a single random linear layer, matching the
  structure (not the trained weights) of a real MoE router.
- Capacity-based dropping and the auxiliary load-balancing loss follow the
  formulation in the **Switch Transformer** paper (Fedus et al., 2021) and
  **GShard** (Lepikhin et al., 2020).
- This is meant as an intuition-building tool, not a benchmark of any real
  model's routing behavior.
"""

# --------------------------------------------------------------------
# Gradio UI
# --------------------------------------------------------------------

with gr.Blocks(title="Mixture of Experts — Explainer & Simulator") as demo:
    gr.Markdown(INTRO_MD)

    with gr.Row():
        with gr.Column(scale=1):
            num_tokens = gr.Slider(8, 512, value=128, step=8, label="Number of tokens (batch)")
            embed_dim = gr.Slider(2, 64, value=16, step=1, label="Token embedding dim")
            num_experts = gr.Slider(2, 16, value=8, step=1, label="Number of experts")
            top_k = gr.Slider(1, 8, value=2, step=1, label="Top-k (experts activated per token)")
            noise_std = gr.Slider(0.0, 3.0, value=0.5, step=0.1, label="Gating noise (std)")
            capacity_factor = gr.Slider(0.5, 3.0, value=1.25, step=0.05,
                                        label="Capacity factor (>1 = slack, <1 = aggressive dropping)")
            ffn_mult = gr.Slider(1, 8, value=4, step=1, label="FFN expansion ratio (per expert)")
            seed = gr.Slider(0, 9999, value=0, step=1, label="Random seed")
            run_btn = gr.Button("▶ Run simulation", variant="primary")

        with gr.Column(scale=2):
            with gr.Row():
                routing_plot = gr.Plot(label="Routing map")
                load_plot = gr.Plot(label="Expert load")
            gate_plot = gr.Plot(label="Gate weight distribution")
            summary_md = gr.Markdown()

    gr.Markdown(FOOTER_MD)

    inputs = [num_tokens, embed_dim, num_experts, top_k, noise_std,
              capacity_factor, ffn_mult, seed]
    outputs = [routing_plot, load_plot, gate_plot, summary_md]

    run_btn.click(fn=run_simulation, inputs=inputs, outputs=outputs)
    demo.load(fn=run_simulation, inputs=inputs, outputs=outputs)

if __name__ == "__main__":
    try:
        demo.launch(theme=gr.themes.Soft())
    except TypeError:
        # Gradio <6: theme belongs on Blocks(), not launch() -- fall back gracefully
        demo.launch()