"""Mixture of Experts From Scratch — reference solution. Everything that makes an MoE layer work, and everything that makes it break: * the router: logits -> softmax -> top-k -> renormalized gates * the two auxiliary losses: load balancing, and router z-loss * capacity factor, token dropping, and slot padding * the full forward pass, with a shared expert and a residual path for dropped tokens * expert-parallel communication accounting (the all-to-all that Phase 06 fixes) * router collapse: simulate it, measure it, and show the aux loss preventing it Pure stdlib, deterministic. Run `python solution.py` for a worked example. """ import math import random # -------------------------------------------------------------------------------------- # Small linear-algebra helpers (stdlib lists — the mechanism must stay visible) # -------------------------------------------------------------------------------------- def softmax(xs): """Numerically stable softmax: subtract the max before exponentiating. Without max-subtraction, exp(1000) overflows to inf and the whole layer produces NaN. Router logits DO get large during training — this is not a theoretical concern. """ if not xs: raise ValueError("softmax of an empty vector is undefined") m = max(xs) exps = [math.exp(x - m) for x in xs] s = sum(exps) return [e / s for e in exps] def logsumexp(xs): """log(sum(exp(x))), computed stably. Used by the router z-loss.""" if not xs: raise ValueError("logsumexp of an empty vector is undefined") m = max(xs) return m + math.log(sum(math.exp(x - m) for x in xs)) def matvec(matrix, vec): """matrix (rows x cols) @ vec (cols) -> (rows).""" if not matrix: raise ValueError("empty matrix") if len(matrix[0]) != len(vec): raise ValueError(f"shape mismatch: matrix cols {len(matrix[0])} vs vec {len(vec)}") return [sum(w * v for w, v in zip(row, vec)) for row in matrix] def relu(xs): return [x if x > 0 else 0.0 for x in xs] # ====================================================================================== # 1. The router # ====================================================================================== def router_logits(token, router_weights): """One score per expert. `router_weights` is (n_experts x d_model). The router is TINY — d_model * n_experts parameters, versus 3 * d_model * d_ff for a single expert. It is also the single most fragile component in the whole model. """ return matvec(router_weights, token) def route_token(token, router_weights, top_k=2, renormalize=True): """Return [(expert_id, gate_weight), ...] of length top_k, highest first. `renormalize=True` rescales the chosen k gates to sum to 1 (Switch/Mixtral style). With renormalize=False the gates keep their raw softmax mass, so the layer output is scaled down by however much probability went to the experts you did NOT pick — a subtle and real design choice, not a detail. Ties are broken by expert index so the routing is deterministic. """ n_experts = len(router_weights) if not 1 <= top_k <= n_experts: raise ValueError(f"top_k must be in [1, {n_experts}], got {top_k}") probs = softmax(router_logits(token, router_weights)) order = sorted(range(n_experts), key=lambda i: (-probs[i], i))[:top_k] if renormalize: total = sum(probs[i] for i in order) if total <= 0: raise ValueError("degenerate routing: all selected gates are zero") return [(i, probs[i] / total) for i in order] return [(i, probs[i]) for i in order] def route_batch(tokens, router_weights, top_k=2, renormalize=True): """Route a whole batch. Returns (assignments, gates, all_probs). assignments[t] = [expert ids for token t] gates[t] = [gate weights for token t] all_probs[t] = the FULL softmax over experts (needed by the aux losses) """ assignments, gates, all_probs = [], [], [] for tok in tokens: probs = softmax(router_logits(tok, router_weights)) picked = route_token(tok, router_weights, top_k, renormalize) assignments.append([e for e, _ in picked]) gates.append([g for _, g in picked]) all_probs.append(probs) return assignments, gates, all_probs # ====================================================================================== # 2. The auxiliary losses — what keeps the router from collapsing # ====================================================================================== def load_balance_loss(assignments, all_probs, n_experts): """Switch-Transformer auxiliary loss: L_aux = E * sum_i f_i * P_i f_i = fraction of token-slots routed to expert i (discrete, NO gradient) P_i = mean router PROBABILITY mass on expert i (continuous, differentiable) Multiplying them gives a loss that is DRIVEN by the true imbalance in f but flows gradient through P. Minimized when both are uniform (= 1/E), which gives L_aux = E * E * (1/E) * (1/E) = 1.0. So 1.0 means perfectly balanced, and larger is worse. Without this term MoE routers collapse: one expert is randomly slightly better, so it gets more tokens, so it trains more, so it gets better. Within a few thousand steps you have paid for E experts and are running a dense model. """ if n_experts <= 0: raise ValueError("n_experts must be positive") if not assignments: raise ValueError("cannot compute a load-balance loss over zero tokens") if len(assignments) != len(all_probs): raise ValueError("assignments and all_probs must describe the same tokens") total_slots = sum(len(a) for a in assignments) if total_slots == 0: raise ValueError("no tokens were routed anywhere") f = [0.0] * n_experts for row in assignments: for e in row: if not 0 <= e < n_experts: raise ValueError(f"expert id {e} out of range") f[e] += 1.0 / total_slots P = [0.0] * n_experts for probs in all_probs: if len(probs) != n_experts: raise ValueError("probability row has the wrong width") for i, p in enumerate(probs): P[i] += p / len(all_probs) return n_experts * sum(fi * Pi for fi, Pi in zip(f, P)) def router_z_loss(all_logits): """Penalize large router logits: mean over tokens of (logsumexp(logits))^2. Why this exists: the router softmax is the one place in the model where a handful of logits decide a DISCRETE choice. Nothing stops them growing without bound, and large logits mean (a) numerical trouble in bf16 and (b) a saturated softmax whose gradient vanishes, freezing the routing. The z-loss keeps them small. Introduced in ST-MoE (Zoph et al., 2022) and it is one of the highest-value stability tricks in MoE training. """ if not all_logits: raise ValueError("no logits provided") return sum(logsumexp(row) ** 2 for row in all_logits) / len(all_logits) def expert_utilization(assignments, n_experts): """Diagnostic: what fraction of slots each expert received, plus summary stats. `max_over_mean` is the number to watch on a dashboard. 1.0 is perfect balance; above ~3 means the router is collapsing and you should intervene NOW, not at the next checkpoint. """ if n_experts <= 0: raise ValueError("n_experts must be positive") total = sum(len(a) for a in assignments) if total == 0: raise ValueError("no tokens were routed") counts = [0] * n_experts for row in assignments: for e in row: counts[e] += 1 fracs = [c / total for c in counts] mean = 1.0 / n_experts return { "counts": counts, "fractions": fracs, "dead_experts": sum(1 for c in counts if c == 0), "max_over_mean": max(fracs) / mean, "collapsed": max(fracs) / mean > 3.0, } # ====================================================================================== # 3. Capacity — the batching constraint nobody mentions # ====================================================================================== def expert_capacity(n_tokens, n_experts, top_k, capacity_factor=1.25): """Per-expert buffer size. capacity = capacity_factor * n_tokens * top_k / n_experts Hardware wants FIXED-SIZE tensors, so each expert gets a fixed buffer. If more tokens route to an expert than fit, the overflow is DROPPED. If fewer arrive, the buffer is padded with zeros and you pay for compute you throw away. capacity_factor < 1.0 guarantees dropping even under perfect balance, so it is rejected here. """ if capacity_factor < 1.0: raise ValueError("capacity_factor below 1.0 drops tokens even when balanced") if n_experts <= 0 or top_k <= 0 or n_tokens <= 0: raise ValueError("n_tokens, n_experts and top_k must be positive") return int(capacity_factor * n_tokens * top_k / n_experts) def apply_capacity(assignments, n_experts, capacity): """Enforce per-expert capacity, first-come-first-served by token order. Returns (kept, stats) where kept[t] is the sublist of assignments[t] that fit. THE THING TO INTERNALIZE: dropping is SILENT. A dropped token is not an error; it just skips the FFN and passes through on the residual stream. Your loss gets very slightly worse and nothing logs anything. This is a real, live source of quality loss in production MoE models, and watching drop_rate is routine pre-training work. """ if capacity < 0: raise ValueError("capacity cannot be negative") used = [0] * n_experts kept, dropped = [], 0 for row in assignments: keep_row = [] for e in row: if used[e] < capacity: used[e] += 1 keep_row.append(e) else: dropped += 1 kept.append(keep_row) total_slots = sum(len(a) for a in assignments) padded = sum(max(0, capacity - u) for u in used) return kept, { "capacity_per_expert": capacity, "dropped": dropped, "padded_slots": padded, "drop_rate": dropped / total_slots if total_slots else 0.0, "used": used, # Fraction of the allocated expert buffers that held real work. "buffer_utilization": (sum(used) / (capacity * n_experts) if capacity * n_experts else 0.0), } # ====================================================================================== # 4. The expert and the full layer # ====================================================================================== def make_expert(d_model, d_ff, seed): """A tiny two-matrix FFN with deterministic pseudo-random weights.""" rng = random.Random(seed) scale = 1.0 / math.sqrt(d_model) w_in = [[rng.uniform(-scale, scale) for _ in range(d_model)] for _ in range(d_ff)] w_out = [[rng.uniform(-scale, scale) for _ in range(d_ff)] for _ in range(d_model)] return {"w_in": w_in, "w_out": w_out} def expert_forward(expert, token): """out = W_out @ relu(W_in @ token).""" return matvec(expert["w_out"], relu(matvec(expert["w_in"], token))) def moe_forward(tokens, router_weights, experts, top_k=2, capacity_factor=1.25, shared_expert=None, renormalize=True): """The full MoE layer forward pass. Returns (outputs, diagnostics). Structure, in order: 1. route every token (top-k experts + gates) 2. enforce capacity -> some slots are dropped 3. for each kept slot, run the expert and accumulate gate * expert(token) 4. a shared expert (if present) runs for EVERY token, always 5. a token whose slots were ALL dropped passes through unchanged (residual) Step 5 is why dropping is silent rather than fatal — and why it is dangerous. """ n_experts = len(experts) if n_experts == 0: raise ValueError("need at least one expert") if not tokens: raise ValueError("no tokens to route") assignments, gates, all_probs = route_batch(tokens, router_weights, top_k, renormalize) cap = expert_capacity(len(tokens), n_experts, top_k, capacity_factor) kept, cap_stats = apply_capacity(assignments, n_experts, cap) d_model = len(tokens[0]) outputs, fully_dropped = [], 0 for t, token in enumerate(tokens): acc = [0.0] * d_model got_any = False for e, g in zip(assignments[t], gates[t]): if e in kept[t]: contrib = expert_forward(experts[e], token) acc = [a + g * c for a, c in zip(acc, contrib)] got_any = True if shared_expert is not None: acc = [a + c for a, c in zip(acc, expert_forward(shared_expert, token))] got_any = True if not got_any: # Every slot dropped: the token skips the FFN entirely and rides the # residual. No error, no log line. This is the silent failure. acc = list(token) fully_dropped += 1 outputs.append(acc) diagnostics = { **cap_stats, "fully_dropped_tokens": fully_dropped, "load_balance_loss": load_balance_loss(assignments, all_probs, n_experts), "router_z_loss": router_z_loss( [router_logits(tok, router_weights) for tok in tokens]), "utilization": expert_utilization(assignments, n_experts), } return outputs, diagnostics def total_training_loss(task_loss, diagnostics, aux_weight=0.01, z_weight=1e-3): """The objective actually minimized: L = L_task + alpha * L_balance + gamma * L_z `aux_weight` (alpha ~ 0.01) is one of the most finicky hyperparameters in modern pre-training. Too small and the router collapses. Too large and you damage quality by forcing tokens to experts that are wrong for them. """ if aux_weight < 0 or z_weight < 0: raise ValueError("loss weights cannot be negative") return (task_loss + aux_weight * diagnostics["load_balance_loss"] + z_weight * diagnostics["router_z_loss"]) # ====================================================================================== # 5. Parameter and communication accounting # ====================================================================================== def moe_parameter_counts(d_model, d_ff, n_experts, top_k, shared_experts=0): """Total (memory) vs active (FLOPs) parameters for ONE MoE layer. The whole trade in one function: - parameters scale with n_experts -> HBM - FLOPs scale with top_k -> compute """ if top_k > n_experts: raise ValueError("top_k cannot exceed n_experts") one_expert = 2 * d_model * d_ff # this lab's experts are 2-matrix FFNs router = d_model * n_experts total = router + (n_experts + shared_experts) * one_expert active = router + (top_k + shared_experts) * one_expert return {"total": total, "active": active, "router": router, "sparsity_ratio": total / active} def expert_parallel_comm_bytes(n_tokens, d_model, n_layers, top_k, bytes_per_elem=2): """Bytes crossing the interconnect under naive expert parallelism. Per layer a token must be SENT to its expert's chip and the result SENT BACK, for each of its top_k experts. That is two all-to-all collectives per layer. This is the wall the Flash 2.0 team hit, and the reason Phase 06 exists: pipelined prefill changes the sharding axis so this traffic is hidden behind compute instead of paid up front. """ if n_layers <= 0 or n_tokens <= 0: raise ValueError("n_layers and n_tokens must be positive") per_hop = n_tokens * d_model * bytes_per_elem * top_k return 2 * n_layers * per_hop def comm_seconds(total_bytes, link_bandwidth_bytes_per_s, n_collectives, latency_s_each=5e-6): """Wall-clock cost of that traffic: transfer time plus per-collective latency.""" if link_bandwidth_bytes_per_s <= 0: raise ValueError("bandwidth must be positive") return total_bytes / link_bandwidth_bytes_per_s + n_collectives * latency_s_each # ====================================================================================== # 6. Router collapse — simulate the failure # ====================================================================================== def simulate_collapse(n_experts, n_tokens, n_steps, aux_weight, seed=0, feedback=0.06): """A minimal model of the rich-get-richer dynamic that kills MoE routers. Each step, expert i's logit grows in proportion to the share of tokens it just won (that is the positive feedback: more tokens -> more training -> more attractive). The auxiliary loss pushes in the opposite direction, proportional to how far that expert's share is above uniform. Returns the max/mean utilization ratio at each step. With aux_weight = 0 it runs away; with a healthy weight it stays near 1.0. """ if n_experts < 2: raise ValueError("collapse needs at least 2 experts") if aux_weight < 0: raise ValueError("aux_weight cannot be negative") rng = random.Random(seed) logits = [rng.gauss(0, 0.01) for _ in range(n_experts)] # tiny random asymmetry uniform = 1.0 / n_experts history = [] for _ in range(n_steps): shares = softmax(logits) # Positive feedback: winning tokens makes you more attractive next step. # Aux loss: being above uniform pushes your logit back down. logits = [z + feedback * (s - uniform) - aux_weight * (s - uniform) * n_experts for z, s in zip(logits, shares)] history.append(max(shares) / uniform) return history # ====================================================================================== # Worked example # ====================================================================================== def _demo_tokens(n, d_model, seed=1): rng = random.Random(seed) return [[rng.gauss(0, 1) for _ in range(d_model)] for _ in range(n)] def main(): print("=" * 78) print("MIXTURE OF EXPERTS FROM SCRATCH") print("=" * 78) D_MODEL, D_FF, N_EXPERTS, TOP_K = 8, 16, 8, 2 tokens = _demo_tokens(64, D_MODEL) rng = random.Random(0) router_w = [[rng.gauss(0, 0.5) for _ in range(D_MODEL)] for _ in range(N_EXPERTS)] experts = [make_expert(D_MODEL, D_FF, seed=i) for i in range(N_EXPERTS)] print("\n[1] Routing one token") tok = tokens[0] probs = softmax(router_logits(tok, router_w)) print(" full softmax over 8 experts: " + " ".join(f"{p:.3f}" for p in probs)) print(f" top-{TOP_K} renormalized : {route_token(tok, router_w, TOP_K)}") print(f" top-{TOP_K} raw softmax : " f"{route_token(tok, router_w, TOP_K, renormalize=False)}") print(" -> renormalizing makes the gates sum to 1; raw does not, which silently") print(" scales the layer output down. A real design choice.") print("\n[2] The parameter trade (per layer)") for k in (1, 2, 8): c = moe_parameter_counts(D_MODEL, D_FF, N_EXPERTS, k) print(f" top_k={k}: total={c['total']:6d} active={c['active']:6d} " f"sparsity={c['sparsity_ratio']:.2f}x") print(" top_k = n_experts is a dense model wearing an MoE costume (sparsity 1.0)") print("\n[3] Capacity, dropping and padding") for cf in (1.0, 1.25, 2.0): _, diag = moe_forward(tokens, router_w, experts, TOP_K, capacity_factor=cf) print(f" capacity_factor={cf:4.2f}: cap={diag['capacity_per_expert']:3d} " f"dropped={diag['dropped']:3d} ({diag['drop_rate']:5.1%}) " f"padded={diag['padded_slots']:3d} " f"buffer_util={diag['buffer_utilization']:5.1%}") print(" -> higher capacity_factor drops fewer tokens and wastes more compute.") print(" That trade-off IS the hyperparameter.") print("\n[4] Balanced vs collapsed routing, measured") balanced_assign = [[i % N_EXPERTS] for i in range(64)] collapsed_assign = [[0] for _ in range(64)] uni = [[1.0 / N_EXPERTS] * N_EXPERTS for _ in range(64)] skew = [[0.93] + [0.01] * (N_EXPERTS - 1) for _ in range(64)] print(f" balanced : L_aux={load_balance_loss(balanced_assign, uni, N_EXPERTS):.4f}" f" max/mean={expert_utilization(balanced_assign, N_EXPERTS)['max_over_mean']:.2f}") print(f" collapsed: L_aux={load_balance_loss(collapsed_assign, skew, N_EXPERTS):.4f}" f" max/mean={expert_utilization(collapsed_assign, N_EXPERTS)['max_over_mean']:.2f}" f" dead={expert_utilization(collapsed_assign, N_EXPERTS)['dead_experts']}") print(" L_aux = 1.0 is perfect balance. Larger is worse. This is the dashboard number.") print("\n[5] Router z-loss keeps logits small") for scale in (1.0, 5.0, 20.0): scaled = [[z * scale for z in router_logits(t, router_w)] for t in tokens] print(f" logit scale x{scale:4.1f} -> z_loss={router_z_loss(scaled):9.3f}") print(" Large logits saturate the softmax, the gradient vanishes, routing freezes.") print("\n[6] Router collapse, simulated") for aux in (0.0, 0.002, 0.005, 0.01): h = simulate_collapse(N_EXPERTS, 1024, n_steps=600, aux_weight=aux) verdict = "COLLAPSED" if h[-1] > 3.0 else "healthy" print(f" aux_weight={aux:5.3f}: max/mean after 600 steps = {h[-1]:7.2f} " f"{verdict}") print(" -> with no auxiliary loss the router runs away. This is not hypothetical;") print(" it is the default behaviour you must actively prevent.") print("\n[7] The communication wall (naive expert parallelism)") n_tok, d_model, n_layers = 8192, 8192, 60 b = expert_parallel_comm_bytes(n_tok, d_model, n_layers, top_k=2) secs = comm_seconds(b, link_bandwidth_bytes_per_s=100e9 / 8, n_collectives=2 * n_layers) print(f" {n_tok} tokens, d_model={d_model}, {n_layers} layers, top_k=2") print(f" bytes across the interconnect : {b/1e9:8.1f} GB") print(f" time at 100 Gbps : {secs:8.3f} s") print(" -> seconds of pure network time before a single useful FLOP.") print(" Phase 06 fixes this by sharding LAYERS instead of EXPERTS.") print("\n[8] A full forward pass, with a shared expert") shared = make_expert(D_MODEL, D_FF, seed=999) out, diag = moe_forward(tokens, router_w, experts, TOP_K, capacity_factor=1.25, shared_expert=shared) print(f" outputs: {len(out)} tokens x {len(out[0])} dims") print(f" load-balance loss : {diag['load_balance_loss']:.4f}") print(f" router z-loss : {diag['router_z_loss']:.4f}") print(f" dropped slots : {diag['dropped']} ({diag['drop_rate']:.1%})") print(f" fully-dropped toks: {diag['fully_dropped_tokens']} " f"(0 because the shared expert always runs)") print(f" total loss (task=2.0): " f"{total_training_loss(2.0, diag):.4f}") print("=" * 78) if __name__ == "__main__": main()