MoE balance batch size vs expert parallel size

MoE balance batch size vs expert parallel size compares two axes that are often collapsed into one setting.

Balance batch size defines which tokens share one routing statistic. EP size defines how expert weights and routed tokens are split across ranks.

Increasing EP does not automatically make a microbatch loss global. Increasing balance batch size does not change which GPU owns an expert.

Keep the statistical population, process group, expert ownership, and dispatch batch separate. Then calculate how many routed assignments each expert and rank should receive.

MoE balance batch size vs expert parallel size

The common auxiliary loss uses E experts and T tokens. For expert i, f_i is its fraction of hard assignments and P_i is its average router probability.

The Megatron Core loss documentation writes the loss as E times the sum of f_i times P_i.

Balance batch size decides the token set behind f_i and P_i. It may be one sequence, one microbatch, several accumulated microbatches, or a global batch across ranks.

Expert parallel size P decides how E experts are distributed. With equal placement, each rank owns E divided by P experts.

The first choice shapes the router objective. The second shapes memory, all-to-all traffic, local expert GEMMs, and the slowest-rank time.

Fanout's batch-size guide for MoE loss shows why the same routes can score differently at microbatch and global scope.

Calculate assignments per expert and rank

Let T be the tokens in the balance population and k the number of selected experts per token.

The population contains A = T times k routed assignments. Under perfect expert balance, each expert receives A divided by E assignments.

With EP size P and equal expert ownership, each rank owns E divided by P experts. Its expected assignments are therefore A divided by P.

Take E = 64 experts, top-2 routing, and T = 4,096 tokens. The router produces 8,192 assignments, or 128 per expert on average.

At EP = 8, each rank owns eight experts and expects 1,024 assignments. At EP = 32, each rank owns two experts and expects 256.

The expert target stayed at 128 because T, k, and E did not change. EP only regrouped those experts and their assignments onto more devices.

This arithmetic is the baseline, not a balance guarantee. A hot expert can push one rank above A divided by P, and co-located hot experts can compound the skew.

Changing EP does not choose the reduction scope

Distributed code must decide which ranks contribute to f_i. The answer is an implementation process group, not a consequence of the letters EP.

Megatron Core exposes microbatch, sequence, and global auxiliary-loss modes in its MoE guide.

Its current global path gathers token counts over a named reduction group and combines global f_i with local router probabilities. The implementation can decompose the gradient without gathering every token's probability matrix.

That distinction matters when data, tensor, context, and expert parallelism coexist. Several ranks may split one sequence, while several EP groups hold replicated copies of the same experts.

Do not assume an all-reduce over the EP ranks is sufficient. Inspect which axes partition tokens and which axes replicate the expert set in the exact framework version.

Log the reduction group membership with the loss. A scalar called global is ambiguous without the ranks and accumulation steps behind it.

Global counts are cheaper than global token routes

A full router-probability tensor has T times E values. The hard expert-count vector has only E values per MoE layer.

The Demons in the Detail paper uses synchronized expert selection frequencies to build a global-batch loss.

Its point is statistical as well as computational. A global batch mixes more sequences and domains, so experts can specialize while aggregate utilization stays balanced.

Averaging losses computed independently on small microbatches is not equivalent. Each microbatch multiplies its own f_i and P_i before the later average.

Synchronizing E counts is much smaller than gathering T by E router probabilities. The local P_i terms can carry gradients while the shared f_i supplies the wider utilization estimate.

The communication still has a cadence. Updating counts after every microbatch, buffering them across accumulation steps, and using a delayed running estimate create different objectives.

Record that cadence instead of describing all three as a larger balance batch.

Wider EP changes execution noise

At fixed T, increasing P lowers expected assignments per rank from A divided by P. The absolute work per rank falls, but relative noise becomes more visible when counts are small.

One extra assignment is 0.1% of an EP-8 rank's 1,024-assignment mean in the worked example. It is 0.4% of an EP-32 rank's 256-assignment mean.

The real tail depends on expert correlations. If the two experts selected by a token often live on one rank, rank load varies more than an independent assignment model predicts.

Wider EP also reduces experts stored per GPU.

The NVIDIA wide-EP report uses 256 experts across 64 GPUs, or four experts per GPU per layer.

That layout adds aggregate memory bandwidth and reduces expert-weight pressure per GPU. It also sends routed activations across a larger all-to-all domain.

The correct EP size comes from the measured makespan, not the balance-loss scalar. Record per-rank received tokens, local expert GEMM shapes, collective time, and the slowest rank.

Fanout's expert versus tensor parallelism guide derives the communication difference between those layouts.

Capacity uses tokens per expert, not EP size

Capacity-limited routers usually size each expert from T, k, E, and a capacity factor. EP size does not belong in the ideal per-expert count A divided by E.

It does affect how expert capacities aggregate on one rank. A rank with E divided by P experts must reserve buffers for the sum of their routed tokens.

Suppose the capacity factor is 1.25 in the worked example. Each expert reserves up to 1.25 times 128, or 160 assignments.

At EP = 8, eight local experts imply capacity for 1,280 assignments per rank. At EP = 32, two local experts imply 320.

Those are equal fractions of the total routed work, but allocator rounding, padding, and grouped-GEMM alignment can make the physical buffers differ.

A locally hot rank can overflow even when global expert frequencies look healthy. Global balance is a corpus objective; one all-to-all step still waits for its busiest destination.

This is why the paper found value in a small local balancing term beside global balance. The local term can shorten stragglers while the global term preserves specialization.

Expert placement is a separate balancing problem

Router training decides which experts tokens prefer. Expert placement decides which preferences share a GPU.

If two hot experts sit on one rank, every token distribution can be unchanged while physical imbalance worsens.

The NVIDIA wide-EP design uses an expert parallel load balancer to move hot and cold experts across GPUs. It supports static mappings and online changes based on recent traffic.

That mechanism does not change the balance batch used to train the router. It changes the mapping from expert IDs to physical ranks.

Fanout's EPLB explainer covers this placement layer. The router loss, dispatch policy, and expert placement should have separate dashboards.

Also separate training from serving. A global training loss may create useful specialization, while a serving batch from one customer domain still concentrates on a small expert set.

Serving placement must react to that traffic even if the training corpus was globally balanced.

Use a four-axis configuration record

First, record the balance population: sequences, non-padding tokens, top-k assignments, microbatches, and accumulation steps.

Second, record the reduction group: exact data, tensor, context, pipeline, and expert-parallel ranks that contribute counts or probabilities.

Third, record expert ownership: E, EP size, experts per rank, replicas, and any placement changes made by EPLB.

Fourth, record execution: assignments per expert, assignments per rank, dropped or padded tokens, grouped-GEMM sizes, all-to-all time, and step makespan.

Sweep balance batch size while holding EP fixed to study specialization. Sweep EP while replaying the same routes to study systems scaling.

Changing both together makes a quality change impossible to separate from a different communication layout.

The MoE routing walkthrough provides the token-level path behind these measurements.

Balance batch size defines what should be statistically balanced. Expert parallel size defines where that routed work runs. A reliable MoE configuration names both and proves each with its own measurements.