AI Engineer Exhaustive
MediumPlusDeepSeek-V3

Load Balancing

DeepSeek-V3

Medium

Auxiliary-loss-free load balancing via bias

Independent study note. Written from the public paper and official code. This is not TensorTonic Plus and does not reproduce their exercises, starter code, or tests. For the official version, subscribe on TensorTonic.

Overview

Sparse MoE layers waste capacity when a few experts absorb most tokens and starve the rest. The usual fix is an auxiliary load-balance loss on the router, but a large coefficient fights the language-modeling objective and a small one fails to keep experts busy. DeepSeek-V3 therefore keeps routing selection and routing weights on different scores. A per-expert bias b_i is added only when choosing the top-K_r experts; the value that actually multiplies each expert output stays the unbiased affinity s_{i,t}. The paper (Wang et al. 2024a, as used in §2.1) reports that this split balances expert load during the 14.8T-token pre-train without the quality hit of a strong auxiliary loss. A tiny sequence-wise balance term remains as a safety net against collapse inside a single sequence.

How it works

Token-to-expert affinities are sigmoid scores of the hidden state against each expert centroid (see the router note). For load-aware selection, those scores are shifted by a vector \mathbf{b}\in\mathbb{R}^{N_r} that is not a learned Linear bias:

g'_{i,t}=\begin{cases} s_{i,t}, & s_{i,t}+b_i\in\operatorname{Topk}\bigl(\{s_{j,t}+b_j\}_{j=1}^{N_r},K_r\bigr),\\ 0, & \text{otherwise.} \end{cases}

The gating values that scale expert outputs are still g_{i,t}=g'_{i,t}/\sum_j g'_{j,t} from the unbiased s_{i,t}. After each training step the implementation counts how many tokens in the global batch landed on expert i. If that count is above the fair share K_r T_{\text{batch}}/N_r, b_i is decreased by a step \gamma; if below, b_i is increased by \gamma. Overloaded experts become harder to pick next step; idle ones become easier. The paper sets \gamma=0.001 for the first 14.3T tokens and freezes \gamma=0 for the last 500B.

The complementary sequence-wise loss is the usual Switch-style product of fractional assignment and mean probability, but with a very small coefficient:

\mathcal{L}_{\mathrm{Bal}}=\alpha\sum_{i=1}^{N_r}f_i P_i,\qquad\alpha=10^{-4}.

Here f_i is the fraction of tokens in a sequence that selected expert i (normalized so a perfectly balanced sequence has f_i=1), and P_i is the mean of the renormalized affinities s'_{i,t}=s_{i,t}/\sum_j s_{j,t} over that sequence. The paper is explicit that this term is only there to stop extreme per-sequence imbalance; batch-level balance is the job of \mathbf{b}.

Together with node-limited routing (at most M=4 nodes per token), the bias rule is why V3 trains and serves without token dropping.

Official code

The released inference graph lives in inference/model.py. Class Gate stores self.bias as an nn.Parameter of shape (n_routed_experts,) only when args.dim == 7168 (the 671B width). The forward pass computes scores, keeps original_scores, then does scores = scores + self.bias before grouped top-k. Weights are gathered from original_scores, matching equation (16). The \gamma update and \mathcal{L}_{\mathrm{Bal}} are training-time; they do not appear in this inference file. Hyper-parameters for the 671B checkpoint are in inference/configs/config_671B.json.

Watch-outs

Sources