MoE load balancing loss and batch size

MoE load balancing loss changes with batch size even when the model, tokens, and expert assignments stay fixed. The scope used to estimate expert load changes what the router is punished for.

A small microbatch can make the loss behave like a per-sequence constraint. A global batch asks for balance across a wider mix of sequences, which can leave room for domain specialization.

The choice also affects variance and communication. It should be logged as part of the objective, not treated as an invisible implementation detail.

MoE load balancing loss depends on batch size

The Switch Transformer paper defines a common auxiliary loss as N times the sum of f_i times P_i across N experts.

For expert i, f_i is its fraction of hard top-k assignments. P_i is the average soft router probability for that expert over the chosen balancing scope.

The hard counts describe where tokens went. The probabilities carry gradients back into the router because the discrete expert index is not differentiable.

If f and P are both uniform, the unweighted loss has a floor of 1. Concentrating hard assignments and soft probability on the same experts raises it.

Batch size enters through the averages. Change which tokens share one f and P estimate, and the scalar can change without altering a single token's route.

Fanout's MoE routing explainer follows the full token-to-expert path. The narrow question here is where the balancing average starts and ends.

The same assignments can score 2 or 1

Take four experts and two 100-token sequences. For a clean example, assume the router probabilities match the hard routing fractions.

Sequence A sends half its tokens to expert 0 and half to expert 1. Its f and P vectors are [0.5, 0.5, 0, 0].

Sequence B sends half to expert 2 and half to expert 3. Its vectors are [0, 0, 0.5, 0.5].

For either sequence, the loss is 4 times (0.5 squared plus 0.5 squared), which equals 2. Averaging the two sequence losses still gives 2.

Now combine both sequences before computing f and P. Every expert receives one quarter of the 200 tokens, so both vectors become [0.25, 0.25, 0.25, 0.25].

The global loss is 4 times four times 0.25 squared, which equals 1. Overall utilization is perfectly balanced even though neither sequence uses every expert.

Sequence and global scopes prefer different routing structures. Sampling noise is only one part of the difference.

Small microbatches tighten the constraint

The ACL 2025 paper Demons in the Detail studies the scope explicitly.

Large MoE models often use very few sequences per device microbatch. A loss computed there can become close to sequence-level balancing, especially when one microbatch contains a narrow domain.

Suppose a microbatch contains only code. Local balancing pushes its code tokens toward every expert, even if a subset of experts could specialize in code while other domains use the rest.

Gradient accumulation does not automatically repair this. Averaging losses computed on separate microbatches is not the same as forming f and P over their union.

In the worked example, the average of two local losses is 2. Recomputing the statistics after joining the samples gives 1.

The difference comes from multiplying local averages before aggregation. Once each microbatch produces its own f_i times P_i term, a later mean cannot recover the cross-microbatch cancellation.

Global-batch loss changes the estimate

The Qwen engineering note describes a practical global-batch calculation.

It synchronizes the expert selection frequency f_i across parallel groups, uses that global frequency with local router probabilities, then aggregates loss contributions across microbatches.

The communicated statistic has one value per expert rather than one record per token. That keeps the extra synchronization small compared with gathering all router outputs.

On a 3.4B-total, 0.6B-active model trained on 400B tokens, the paper reports test perplexity falling from 7.383 at Balance BSZ 2 to 7.198 at 512.

The curve improved quickly through 128 and then flattened. That result belongs to the paper's models and data, so 128 is not a universal setting.

The useful conclusion is procedural. Sweep the balancing scope separately from the ordinary optimizer batch size, and record the scope behind every loss value.

Variance is only part of the story

A larger sample gives f_i a lower-variance estimate of corpus-level routing. That matters when a few tokens can move a small microbatch's expert fractions by large steps.

For four experts and eight top-1 tokens, one token changes an expert's observed share by 12.5 percentage points. With 512 tokens, one token changes it by about 0.2 points.

The ACL study ran a shuffle ablation to separate sample count from distribution. It drew a small set from the global token pool so the count stayed local while the domain mix resembled the global batch.

The shuffled result was close to global-batch performance in its reported tests. That supports the idea that narrow microbatch composition, not only sample count, creates the stricter constraint.

A useful diagnostic therefore logs both the number of routed tokens and the domain or sequence composition behind f and P.

Low variance does not guarantee the desired objective. A large batch of one domain can still impose balance inside that domain.

Sequence-level balance is a deliberate option

Current Megatron Core MoE documentation exposes separate microbatch, sequence-level, and global auxiliary-loss modes.

Sequence-level balancing computes expert usage for each sequence. It is stricter than asking a mixed global batch to balance only in aggregate.

That strictness can be useful when every request must exercise a broad expert set or when per-sequence hotspots cause a downstream system problem.

It can also oppose specialization. A code sequence cannot lean heavily on code experts without paying a larger sequence-level penalty, even if prose sequences balance the cluster later.

Do not infer the mode from the scalar's name. Save the reduction axes, sequence grouping, expert-parallel group, and accumulation boundary with the training configuration.

The existing router z-loss guide covers a different control. Z-loss limits router logit scale; it does not choose the population over which expert utilization is balanced.

Global balance does not remove local stragglers

Corpus-level balance is a model objective. Hardware still executes one microbatch and one communication step at a time.

A globally balanced update can contain a locally hot expert. The GPU that owns it may receive more tokens, extend the all-to-all critical path, or hit a capacity limit.

The Qwen study tested a small local balancing term alongside global balance on its 43B-total, 6.6B-active setup.

Its reported step time improved from 1.64 to 1.59 seconds. Average perplexity moved from 5.779 to 5.795, showing a small quality and specialization tradeoff for better local execution.

That pattern suggests two measurements. Use global statistics to judge specialization and corpus balance, then use per-rank token counts and step time to find local bottlenecks.

Fanout's EPLB explainer addresses inference-time placement. Training loss and physical expert placement solve related but separate problems.

Choose scope from the failure you see

Start with per-layer expert counts at sequence, microbatch, global batch, and rank scope. A single global histogram can hide the exact imbalance that slows one step.

If experts specialize poorly and microbatches contain narrow domains, compare microbatch and global-batch loss while holding the data order and coefficient fixed.

If step time has a long tail despite good global balance, inspect per-rank counts. A weak local term, capacity policy, or different dispatch layout may help more than enlarging the global scope.

If the loss is unstable, report routed-token count with its mean and variance. Batch size in sequences can be misleading when sequence lengths differ sharply.

Keep the unweighted auxiliary loss separate from its coefficient and from language-model loss. A quiet weighted term may reflect a small coefficient rather than healthy routing.

The batch size for MoE load balancing loss defines the population that must look balanced. Pick that population from the behavior you want, then verify quality, specialization, and hardware balance separately.