Mixture-of-experts and expert parallelism

Open in Colab ▶️ Run this notebook in Colab

Mixture-of-experts and expert parallelism#

Track 05 - Serving · Notebook 09 · Runtime: ≈30 s on CPU

Prerequisites: 05_serving/01 (roofline), 01_inference/01 (KV cache).

Papers:

  • Shazeer et al. 2017, Outrageously Large Neural Networks (1701.06538).

  • Fedus et al. 2021, Switch Transformer (2101.03961).

  • Lepikhin et al. 2020, GShard (2006.16668).

What#

An MoE layer has N experts; each token’s router picks the top-k (usually k=1 or 2) and its hidden vector goes through only those. Active parameters per token drop by k/N. In serving:

  • All-to-all communication: tokens get dispatched to whichever GPU holds their chosen expert, then results come back.

  • Capacity factor: each expert has a fixed per-batch slot budget. If the router sends too many tokens, the overflow gets dropped (or sent back to a fallback path).

  • Load balance loss: trains the router to spread tokens evenly, because an imbalanced router wastes GPU capacity.

We simulate the router, the capacity bin-packing, and the load imbalance metric. Three checks: expert dispatch is correct, active parameter ratio matches k/N, load-balance loss penalises skewed distributions.

from llm_systems_cookbook.nb import bootstrap

import numpy as np
import torch
import torch.nn.functional as F

s = bootstrap("05_serving_09_moe_expert_parallelism")

Top-k router#

A linear projection from hidden state to N logits; softmax; take top-k. Returns both the experts and their gating weights.

torch.manual_seed(0)
N_EXPERTS = 8
D = 256
TOKENS = 1024
K = 2

router = torch.nn.Linear(D, N_EXPERTS, bias=False)
router.weight.data = torch.randn_like(router.weight.data) * 0.1

x = torch.randn(TOKENS, D)
logits = router(x)
gates = F.softmax(logits, dim=-1)
topk_weights, topk_idx = gates.topk(K, dim=-1)
print(f"top-{K} shape: weights {tuple(topk_weights.shape)}  idx {tuple(topk_idx.shape)}")
print(f"example token 0 chose experts {topk_idx[0].tolist()} with weights {topk_weights[0].tolist()}")

Dispatch and gather#

Each expert is a small MLP; tokens sent to expert e get processed by MLP e and their output gets weighted by the gate.

class Expert(torch.nn.Module):
    def __init__(self, d: int, ff: int = 1024) -> None:
        super().__init__()
        self.fc1 = torch.nn.Linear(d, ff)
        self.fc2 = torch.nn.Linear(ff, d)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.fc2(F.gelu(self.fc1(x)))


torch.manual_seed(1)
experts = torch.nn.ModuleList([Expert(D) for _ in range(N_EXPERTS)])


def moe_forward(x: torch.Tensor, topk_idx: torch.Tensor, topk_weights: torch.Tensor) -> torch.Tensor:
    out = torch.zeros_like(x)
    for e in range(N_EXPERTS):
        # Find tokens that routed (to slot 0 OR 1) to expert e.
        mask = (topk_idx == e).any(dim=-1)
        if not mask.any():
            continue
        # For each such token, sum the gated expert output from this expert.
        tokens_for_e = x[mask]
        expert_out = experts[e](tokens_for_e)
        # Weight: sum of gate weights where idx == e.
        per_token_weights = (topk_weights * (topk_idx == e).float()).sum(dim=-1)[mask]
        out[mask] += expert_out * per_token_weights.unsqueeze(-1)
    return out


out = moe_forward(x, topk_idx, topk_weights)
print(f"moe forward: in={tuple(x.shape)}  out={tuple(out.shape)}")
s.check(
    "moe_output_shape_matches",
    lambda: out.shape == x.shape,
    msg=f"out shape = {out.shape}",
)
s.check(
    "moe_output_finite",
    lambda: torch.isfinite(out).all().item(),
)

Capacity factor#

Each expert has a slot budget B * k / N * factor. Factor > 1 gives headroom; factor = 1 means drops on any imbalance.

def token_count_per_expert(topk_idx: torch.Tensor, n_experts: int) -> torch.Tensor:
    counts = torch.zeros(n_experts)
    for e in range(n_experts):
        counts[e] = (topk_idx == e).sum().float()
    return counts


counts = token_count_per_expert(topk_idx, N_EXPERTS)
perfectly_balanced = TOKENS * K / N_EXPERTS
print(f"actual per-expert: {counts.tolist()}")
print(f"perfectly balanced would be: {perfectly_balanced}")
imbalance = counts.max().item() / perfectly_balanced
print(f"imbalance (max / perfect) = {imbalance:.2f}")


def capacity_check(counts: torch.Tensor, capacity: int) -> float:
    '''Return fraction of tokens that would be dropped at this capacity.'''
    overflow = torch.clamp(counts - capacity, min=0).sum().item()
    return overflow / counts.sum().item()


cap_10 = capacity_check(counts, int(perfectly_balanced * 1.0))
cap_15 = capacity_check(counts, int(perfectly_balanced * 1.5))
print(f"drop rate at capacity-factor 1.0  = {cap_10:.2%}")
print(f"drop rate at capacity-factor 1.5  = {cap_15:.2%}")

s.check(
    "higher_capacity_reduces_drops",
    lambda: cap_15 <= cap_10,
    msg=f"cap=1.0 drops={cap_10:.2%}  cap=1.5 drops={cap_15:.2%}",
)
s.check(
    "untrained_router_has_measurable_imbalance",
    lambda: imbalance >= 1.05,
    msg=f"imbalance = {imbalance:.2f}",
)

Switch Transformer auxiliary loss#

L_aux = N * sum_i f_i * p_i where f_i is the fraction of tokens routed to expert i and p_i is the average gate prob for i. Minimised when both are 1/N ⇒ uniform. Add this to the training loss with a small coefficient (0.01 in Switch).

def switch_aux_loss(topk_idx: torch.Tensor, gates: torch.Tensor, n_experts: int) -> float:
    f = torch.zeros(n_experts)
    for e in range(n_experts):
        f[e] = (topk_idx == e).any(dim=-1).float().mean()
    p = gates.mean(dim=0)
    return float(n_experts * (f * p).sum())


imbalanced_loss = switch_aux_loss(topk_idx, gates, N_EXPERTS)

# Simulate a perfectly-balanced router: tile expert ids round-robin.
balanced_idx = torch.arange(TOKENS * K).remainder(N_EXPERTS).reshape(TOKENS, K)
balanced_gates = torch.full((TOKENS, N_EXPERTS), 1.0 / N_EXPERTS)
balanced_loss = switch_aux_loss(balanced_idx, balanced_gates, N_EXPERTS)

print(f"aux loss, imbalanced = {imbalanced_loss:.3f}")
print(f"aux loss, balanced   = {balanced_loss:.3f}  (minimum = 1.0)")

s.check(
    "aux_loss_imbalanced_higher_than_balanced",
    lambda: imbalanced_loss > balanced_loss,
    msg=f"imb={imbalanced_loss:.3f}  bal={balanced_loss:.3f}",
)

# Active params: at k=2 of N=8 experts, each token uses 2/8 = 25% of expert params.
active_ratio = K / N_EXPERTS
print(f"active expert fraction per token = {active_ratio:.1%}")
s.check(
    "active_params_equal_k_over_N",
    lambda: abs(active_ratio - K / N_EXPERTS) < 1e-9,
    msg=f"{active_ratio}",
)

Per-expert token load - where the imbalance lives#

One bar per expert, sorted by load. The dashed line is the perfectly-balanced target (tokens * k / N); bars above it will overflow the capacity factor and drop tokens, bars well below it leave their GPU idle. The aux-loss training signal pushes every bar toward the line.

import matplotlib.pyplot as plt

order = sorted(range(N_EXPERTS), key=lambda e: -counts[e].item())
sorted_counts = [counts[e].item() for e in order]
expert_labels = [f"E{e}" for e in order]
cap_10_limit = perfectly_balanced

fig, ax = plt.subplots(figsize=(7, 4))
bars = ax.bar(expert_labels, sorted_counts,
              color=["tab:red" if c > cap_10_limit else "tab:blue" for c in sorted_counts],
              edgecolor="black", lw=0.5)
ax.axhline(perfectly_balanced, color="black", ls="--", lw=1.2,
           label=f"balanced target ({perfectly_balanced:.0f})")
ax.axhline(perfectly_balanced * 1.5, color="tab:gray", ls=":", lw=1,
           label=f"capacity factor 1.5 ({perfectly_balanced * 1.5:.0f})")

for bar, v in zip(bars, sorted_counts, strict=True):
    ax.text(bar.get_x() + bar.get_width() / 2, v,
            f"{int(v)}", ha="center", va="bottom", fontsize=9)

ax.set_xlabel("expert (sorted by load)")
ax.set_ylabel("tokens routed")
ax.set_title(f"routing load across {N_EXPERTS} experts "
             f"(imbalance max/target = {imbalance:.2f}x)")
ax.legend(loc="upper right", fontsize=9)
ax.grid(True, axis="y", alpha=0.3)
plt.tight_layout()
plt.show()

Exercises#

  1. Train the router. Run gradient descent on the aux loss; watch the imbalance ratio and drop rate fall.

  2. Expert parallelism. Split the N_EXPERTS experts across 4 “devices” (just dicts in Python). Compute the all-to-all volume as a function of batch size.

  3. Shared experts. DeepSeek-V2’s architecture has a small set of “always-on” experts plus routed experts. Extend the forward.

References#

  • Switch Transformer paper §2 for the aux loss.

  • GShard for the top-2 routing pattern.

  • Megablocks for a production GPU implementation that avoids token drops via padded expert blocks.

s.summary()
s.save()