"""Mixture of Experts From Scratch — YOUR implementation. Fill in every `# TODO`. Signatures, docstrings and validation contracts are given; the mechanism is yours. pytest test_lab.py -v # red until you implement LAB_MODULE=solution pytest test_lab.py -v # the reference (must be green) python solution.py # the worked example ORDER OF WORK 1. softmax, logsumexp, matvec, relu (primitives — max-subtraction!) 2. router_logits, route_token, route_batch (the router) 3. load_balance_loss, router_z_loss, expert_utilization (the aux losses) 4. expert_capacity, apply_capacity (dropping and padding) 5. make_expert, expert_forward, moe_forward, total_training_loss 6. moe_parameter_counts, expert_parallel_comm_bytes, comm_seconds 7. simulate_collapse (the failure mode) THE MONEY TEST is `test_router_collapses_without_an_auxiliary_loss`, paired with `test_auxiliary_loss_prevents_collapse`. Rich-get-richer is the DEFAULT behaviour of an MoE router, not an exotic failure — you must actively prevent it. THE CALIBRATION POINT to aim for: perfectly balanced routing gives a load-balance loss of EXACTLY 1.0. If your implementation does not produce 1.0 there, the formula is wrong. DETERMINISM: seeded `random.Random(seed)` only. Ties in the router break by expert index. """ import math import random 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. """ # TODO: implement. See WARMUP.md and the docstring above. raise NotImplementedError def logsumexp(xs): """log(sum(exp(x))), computed stably. Used by the router z-loss.""" # TODO: implement. See WARMUP.md and the docstring above. raise NotImplementedError def matvec(matrix, vec): """matrix (rows x cols) @ vec (cols) -> (rows).""" # TODO: implement. See WARMUP.md and the docstring above. raise NotImplementedError def relu(xs): # TODO: implement. See WARMUP.md and the docstring above. raise NotImplementedError 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. """ # TODO: implement. See WARMUP.md and the docstring above. raise NotImplementedError 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. """ # TODO: implement. See WARMUP.md and the docstring above. raise NotImplementedError 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) """ # TODO: implement. See WARMUP.md and the docstring above. raise NotImplementedError 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. """ # TODO: implement. See WARMUP.md and the docstring above. raise NotImplementedError 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. """ # TODO: implement. See WARMUP.md and the docstring above. raise NotImplementedError 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. """ # TODO: implement. See WARMUP.md and the docstring above. raise NotImplementedError 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. """ # TODO: implement. See WARMUP.md and the docstring above. raise NotImplementedError 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. """ # TODO: implement. See WARMUP.md and the docstring above. raise NotImplementedError def make_expert(d_model, d_ff, seed): """A tiny two-matrix FFN with deterministic pseudo-random weights.""" # TODO: implement. See WARMUP.md and the docstring above. raise NotImplementedError def expert_forward(expert, token): """out = W_out @ relu(W_in @ token).""" # TODO: implement. See WARMUP.md and the docstring above. raise NotImplementedError 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. """ # TODO: implement. See WARMUP.md and the docstring above. raise NotImplementedError 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. """ # TODO: implement. See WARMUP.md and the docstring above. raise NotImplementedError 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 """ # TODO: implement. See WARMUP.md and the docstring above. raise NotImplementedError 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. """ # TODO: implement. See WARMUP.md and the docstring above. raise NotImplementedError 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.""" # TODO: implement. See WARMUP.md and the docstring above. raise NotImplementedError 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. """ # TODO: implement. See WARMUP.md and the docstring above. raise NotImplementedError if __name__ == "__main__": print("Fill in the TODOs, then compare against: python solution.py")