A router picks experts for each token
Shading is the router's probability for each expert. Outlined cells are the experts that actually run.
Expert computations per token2 of 8
Balance term (1.0 is perfectly even)1.37
Challenges
0 of 2 complete
- Find a token that the router splits almost evenly between two experts: neither gets more than 60% of its weight. (not yet)
- Find the setup with the most uneven load: the highest balance term. (not yet)
A real router: OLMoE-1B-7B
Which of 64 experts a trained model sent each token to, at each of its 16 layers, recorded from the real model.
Loading the recorded outputs...
Total versus active parameters
Same shape as Mixtral 8x7B: 32 layers, model width 4,096, expert hidden size 14,336. Change the experts.
Stored (all experts)46.7B
Used per token (k experts)12.9B
Shared by every token: 1.6B Expert weights: 0.2B per expert per layer x 32 layers
Total parameters46.7B
Active per token12.9B
Weights in memory at 16 bits93 GB
Challenges
0 of 3 complete
- Find settings where total and active parameters are equal. (not yet)
- Store more than 300B parameters while each token uses fewer than 13B. (not yet)
- Find the most experts whose weights still fit in 80 GB of memory at 16 bits. (not yet)
Try this
- Switch between top 1 and top 2 routing and watch the gate weights for a token.
- In the real router, pick Python, shade by tokens of this text, and look through the layers for an expert that takes almost all of it. Then check what it took from the other texts.
- In the calculator, set experts and k to 1 to recover the dense model, then add experts and watch only the stored total grow.
Top-k routing and the balance term in codePython with NumPy
import numpy as np
def softmax(z, axis=-1):
e = np.exp(z - z.max(axis=axis, keepdims=True))
return e / e.sum(axis=axis, keepdims=True)
def show(values):
return " ".join(f"{v:.3f}" for v in values)
# Router scores s = W_r x for one token, over 8 experts.
scores = np.array([0.4, 2.1, -0.3, 1.2, 1.7, -1.0, 0.0, 0.6])
probs = softmax(scores) # the router's probability for each expert
top = np.argsort(-scores)[:2] # top 2: only these experts run
gates = softmax(scores[top]) # softmax over the chosen scores only
print("probabilities:", show(probs))
print("experts that run:", ", ".join(f"E{e + 1}" for e in top))
print("gate weights:", show(gates), "| renormalized:", show(probs[top] / probs[top].sum()))
# Load over a batch: four tokens, four experts, top 1.
scores = np.array([[2.0, 0.1, -0.5, 0.3],
[1.5, 0.2, 0.0, -1.0],
[1.8, -0.2, 0.4, 0.1],
[0.1, 1.2, 0.3, -0.4]])
chosen = scores.argmax(axis=1)
f = np.bincount(chosen, minlength=4) / len(chosen) # share of tokens per expert
P = softmax(scores, axis=1).mean(axis=0) # mean router probability
print("tokens per expert:", np.bincount(chosen, minlength=4).tolist())
print(f"balance term N * sum(f * P) = {4 * np.sum(f * P):.3f} (1.0 is perfectly even)")probabilities: 0.067 0.365 0.033 0.148 0.245 0.016 0.045 0.081
experts that run: E2, E5
gate weights: 0.599 0.401 | renormalized: 0.599 0.401
tokens per expert: [3, 1, 0, 0]
balance term N * sum(f * P) = 1.834 (1.0 is perfectly even)Counting Mixtral's parameters in codePython with the standard library
# Mixtral 8x7B's published shape.
d, d_ff, layers, heads, kv_heads, vocab = 4096, 14336, 32, 32, 8, 32000
def count(experts, k):
expert = 3 * d * d_ff # one SwiGLU expert: three weight matrices
attention = 2 * d * d + 2 * d * kv_heads * (d // heads) # Q and O, then K and V
router, norms = d * experts, 2 * d
shared = layers * (attention + router + norms) + 2 * vocab * d + d
return shared + layers * experts * expert, shared + layers * k * expert
for experts, k in ((1, 1), (8, 2), (64, 2)):
total, active = count(experts, k)
name = f"{experts} expert" + ("s" if experts > 1 else "")
print(f"{name}, top {k}: {total / 1e9:.1f}B stored, "
f"{active / 1e9:.1f}B per token, {2 * total / 1e9:.0f} GB at 16 bits")1 expert, top 1: 7.2B stored, 7.2B per token, 14 GB at 16 bits
8 experts, top 2: 46.7B stored, 12.9B per token, 93 GB at 16 bits
64 experts, top 2: 362.4B stored, 12.9B per token, 725 GB at 16 bits