[moe] Good 10T: measure capacity overflow

[dlwh]

Summary

This issue started by asking whether MoE capacity padding was costing enough throughput to justify measuring and possibly lowering it in the good-enough 10T recipe. PR #4052 added capacity-overflow metrics, follow-up analysis found the default 1.25 capacity factor costs roughly 11% throughput at 1e21 scale while lower caps only slightly worsen loss, and a later 1e20 EP=4 comparison showed cap=1.0 was 8.3% faster with only a +0.001 BPB hit. Current conclusion: move the default capacity factor to 1.0, because the speedup appears to outweigh the quality loss on the tested runs.

Helpful links

• Instrumentation PR #4052
• Overflow analysis and recommendation
• 1e20 cap=1.0 vs 1.25 comparison
• Decision to update the default

Description

TL;DR: Measure capacity overflow on the current good-enough 10T candidate so routing overflow is quantified instead of guessed.

Hypothesis or Goal

We want to know whether overflow is materially affecting quality, efficiency, or both on the path we are currently considering.

Links

• Parent sweep: MoE Sweep: Hyperparameters, Routing, Architecture, and Optimizer #3469\n* Gate: [moe] Good 10T gate for #3469 #4013

Results

[github-actions]

🤖 I'm looking into this issue. I'll post a task list shortly with my plan and progress.

[claude]

🤖 Plan:

• Wire report_capacity_overflow=True into Grug MoE model
• Add dropped token count to router_stats / router_metrics
• Log capacity overflow metrics in _summarize_router_metrics
• Write test
• Run pre-commit and tests
• Open PR

PR: #4052

Research

moe_mlp() already supports report_capacity_overflow=True returning a dropped-token scalar, but MoEMLP.__call__ never passes it. The count is computed in _prefix_cap_counts() / _moe_mlp_ep_ring_local() but discarded.

Relevant code

• grug_moe.py#L436 - report_capacity_overflow param
• model.py#L328 - moe_mlp() call without overflow reporting
• model.py#L225 - _summarize_router_metrics (no overflow)

Proposed Fix

Pass report_capacity_overflow=True in MoEMLP.__call__, propagate the dropped count through router_stats -> router_metrics, and log it as train/router/capacity_overflow_* in _summarize_router_metrics.

[ClassicLarry]

🤖 EP=2 capacity factor sweep results (d=1024, 10,670 steps, v5p-8) 3e18 runs:

Run	Cap	C4-en BPB	Macro Loss	Tokens/s
main_baseline_3e_18	1.25 (EP=1)	1.0478	3.190	184,709
h_1024_ep2_cap_095	0.95 (EP=2)	1.0502	3.205	176,419
h_1024_ep2_cap_090	0.9 (EP=2)	1.0503	3.207	180,250
h_1024_ep2_cap_080	0.8 (EP=2)	1.0530	3.218	187,521

Human commentary:
The loss increase from dropping tokens, even at 0.8, is negligible. In fact, the 6% improvement in MFU outweighs the increase in loss.

[ClassicLarry]

Looked into this. Summary:

• At 1e21 scale we are incurring roughly a 11% slowdown due to the 1.25x capacity padding.
• At 3e18 scale, even if we set capacity factor to 0.8, the loss impact is minimal (1.0478 vs 1.053) and is outweighed by the MFU boost.
• As we increase EP, increase DP, or decrease batch size, the risk of overflow will increase. Want to be aware of this, because for large runs we may scale DP faster than we scale batch size. The metric to track that defines how hard balancing will be is balance_count=Kseq_lenbatch_size/data_parallelism.
• There is an asymmetry in how we apply the cap: the last expert on the core always takes the hit first, and the last token in the batch always gets dropped first. Given the loss impact is minimal, not viewing this as a large concern.
• Even at balance_count = 64k (4x less than the 1e21 run), for EP=4 there was zero benefit to having a capacity factor over 1.15. If we had an extreme EP=64, then we would have dropped 7% of assignments at capacity cap=1.0, and 1% of assignments at cap=1.25.

Based on this, we can likely set our capacity factor to 1.1 conservatively, or 1.0 aggressively, to achieve a 6-11% speed-up for a negligible loss hit. We could kickoff 3 1e20 runs to validate this.

Details on our current Capacity Overflow approach:

For a given run, we define the following quantities:

• seq_len=S
• batch_size=B
• num_devices=D
• expert_parallelism=EP
• data_parallelism=DP=-1
• _DEFAULT_EP_CAPACITY_FACTOR=1.25

MFU Impact

If EP=1, then every device gets all 64 experts, and we do not trigger the code path to pad with the capacity factor.
If EP>1, then we pad our ragged dot matrix calc with 1.25x size on the batch dim. This incurs an 11% slowdown when hidden_dim is 2048. Wandb No Padding, Padding. As we grow hidden_dim, and MLP matmul flops become proportionally larger, this will asymptotically approach up to a 25% slowdown.

Our 1e21 run ran at 17% MFU with EP=4. Had we not padded, it would have ran at roughly 19% MFU. So as a rule of thumb, we are looking at roughly a max +-3 MFU gain from tuning capacity factor. For runs smaller than 1e21, we can generally get away with EP=1 on v4s and rely only on sharding the data axis, saving us this ~11% overhead.

Mechanics of how the dropping occurs

Walking through an example given the config for the prior 1e21 run:

• seq_len=4096
• EP=4
• batch_size=512
• num_devices=128 (v4-256)

Then:

• tokens per batch = 4096*512 = 2,097,152
• data parallelism = 128/4 = 32
• The batch is sharded over (expert, data), but then aggregated back over expert. So local_batch for expert balancing is tokens_per_batch/DP = 2,097,152/32 = 65,536.
• Each device holds 64/4 = 16 experts.
• Each device has a cap of 1.25 * 65,536 * K = 327,680 selections.

If the device receives less than 327,680 selections, then everything fits. If it receives more, then the first experts fill up first, using the tokens in the order received. This creates two asymmetries: 1) expert 16 will always take the hit first if the device is overloaded. This means that under imbalance, expert 16 may average substantially below 1x capacity. 2) The last token in the batch will take the hit first. This creates a correlation across layers, where the last token in the batch may get dropped from many layers.

Two factors influence the fit:

1. What is the local batch size. (seq_len*batch_size/DP). Even if experts are balanced in expectation, variance across samples will cause some amount of imbalance.
2. EP. If EP<64, then we can get away with more expert imbalance, where a low volume expert creates slack for a high volume expert.

The risk of overflowing increases when we:

1. decrease the batch size.
2. increase EP
3. increase DP

The 1e21 run was sitting at a 262,144/327,680 threshold. A run of batch_size=128 on a v4-16 with EP=4 would have (4)(128)(4096)/(2) = 1,048,576/1,310,720. The 1e21 run is higher risk because it scaled DP from 2 to 32 (x16), but only scaled batch size from 128 to 512 (4x). In summary, if we scale DP faster than we are scaling batch size, we are increasing the risk of overflow.

This means we want to track several things:

1. Number of tokens dropped per expert per layer. We need expert granularity on this metric, to assess if the 'last in line' expert is getting hit hardest.
2. The hypothetical drop rate at different EP levels and different thresholds. To simulate this, we would psum counts over the expert axis, then for each EP bucket, count the number that exceed the threshold.

In practice how much overflow happens?

https://wandb.ai/marin-community/dial_moe/runs/main_baseline_3e_18_overload_logging?nw=nwuserlarrydial
Here is a 3e18 run where each core makes 65,536 assignments to balance. If it had EP 4, we could still achieve perfect balancing with 1.15 _DEFAULT_EP_CAPACITY_FACTOR. At EP=64, we would have consistently dropped ~1-5% of instances.

[ClassicLarry]

@claude look for mistakes or flaws in the reasoning above.

[claude]

Claude finished @ClassicLarry's task in 6m 16s —— View job

---

🤖 Review of the capacity overflow analysis:

• Verify arithmetic in the 1e21 worked example
• Check MFU impact claims and asymptotic reasoning
• Review sweep results interpretation
• Check overflow mechanics and asymmetry claims
• Verify scaling risk analysis (DP vs batch size)
• Review recommendations

---

Review: Capacity Overflow Analysis

Overall the analysis is thorough and directionally correct. I found a few issues ranging from a clear arithmetic error to some reasoning that deserves more nuance.

1. Per-device vs total capacity (arithmetic error)

"Each device has a cap of 1.25 × 65,536 × K = 327,680 selections"

This is the total EP-group capacity, not per-device. From the code (grug_moe.py L277):

local_capacity = int(math.ceil(capacity_factor * assignments / ep_size))

With EP=4: local_capacity = ceil(1.25 × 262,144 / 4) = 81,920 per device. The total EP-group capacity is 81,920 × 4 = 327,680, which matches your number — but labeled "each device" when it should be "across the EP group." The ratio (262,144 / 327,680 = 80% utilization, 25% headroom) is correct either way, so the downstream reasoning is unaffected.

2. Cap < 1.0 guarantees dropping even under perfect balance

This is an unstated consequence of the sweep, not a mistake, but worth calling out explicitly. With EP=2 and cap=0.80:

• local_capacity = ceil(0.8 × assignments / 2) = 0.4 × assignments
• Expected load per device (balanced routing) = assignments / 2 = 0.5 × assignments
• Guaranteed drop ≈ 20% of assignments per device, even with perfect expert balance

The fact that BPB only moved from 1.0478 → 1.0530 despite a ~20% assignment drop rate is a genuinely useful finding — it says the model is robust to dropped assignments (likely because each token has multiple top-K selections, so losing one doesn't zero out the MoE contribution). But the write-up should make this explicit: cap=0.80 is not "occasionally dropping overflow," it's "systematically dropping ~20% by design."

3. "Asymptotically approach 25% slowdown" overstates

The 1.25× padding applies only to the MoE matmuls (up/gate projection, down projection). Attention, layernorm, embeddings, and non-MoE compute are never padded. In a standard MoE transformer, even as hidden_dim grows, the MoE fraction of total FLOPs is bounded below 100% (attention scales as O(seq² × d) and doesn't get padded). So the full-model asymptotic slowdown is 0.25 × (MoE FLOPs / total FLOPs), which for a typical 2/3 MoE fraction caps around ~17%, not 25%. The "up to 25%" framing is defensible as a theoretical upper bound on the MoE-specific cost, but the practical whole-model slowdown will always be less.

4. "Increasing EP increases overflow risk" needs qualification

This is true going from EP=1 → EP>1 (EP=1 skips the capacity code path entirely, so overflow is impossible). But for EP≥2, the picture is more subtle.

The overflow probability depends on the concentration of per-device load around its mean. With balanced routing, device load is approximately Binomial(assignments, local_experts/num_experts), and the overflow threshold is capacity_factor × assignments / ep_size. The z-score for overflow works out to:

z ≈ (cap_factor - 1) × sqrt(assignments / (ep_size × (1 - 1/ep_size)))

Since assignments = tokens × topk / DP and DP = num_devices / EP, we get assignments / ep_size = tokens × topk / num_devices, which is invariant to EP. So in the Gaussian approximation, the overflow probability is roughly constant across EP values (for EP≥2), because increasing EP simultaneously increases assignments (by shrinking DP) while increasing ep_size — the two effects cancel.

The empirical data from the logged run does show different overflow at different simulated EP levels, but that's simulating different EP from a fixed set of routing decisions (fixed balance_count), rather than accounting for the fact that a real higher-EP run would also have different DP and hence different assignments.

The claim is true when DP is held fixed (i.e., you increase EP by adding more devices), but not when num_devices is held fixed (where increasing EP forces DP down).

5. Sweep comparison framing

The 3e18 sweep table compares an EP=1 baseline against EP=2 runs at different caps. The "6% MFU improvement" is comparing within EP=2 (cap 0.95 → 0.80), not against the EP=1 baseline. Against the baseline, the best EP=2 run (cap 0.80) is only ~1.5% faster (187,521 vs 184,709 tok/s) while having worse loss. This is clear if you read carefully, but someone skimming might think the 6% improvement is relative to EP=1.

6. Minor: balance_count formula

"balance_count = K × seq_len × batch_size / data_parallelism"

This is assignments (total selections per data shard across all experts). It's the right quantity for measuring "how much statistical averaging happens," but note it's per data-shard, not per-device. Per-device relevant load is balance_count / EP. If you're comparing overflow risk across configurations, balance_count / EP = K × seq_len × batch_size / num_devices is the more directly comparable quantity (and is EP-invariant as noted above).

---

Summary

The core conclusions are sound:

• Capacity padding incurs a meaningful MFU cost at EP>1 (~11% at the 1e21 scale)
• The model is surprisingly robust to token dropping (cap 0.80 barely hurts loss)
• Reducing cap to 1.0–1.1 is a reasonable recommendation
• Validating at 1e20 scale before the production run is the right call

The main things to fix/clarify: the per-device vs total capacity arithmetic (point 1), explicitly noting that sub-1.0 cap guarantees systematic dropping (point 2), and qualifying the EP vs overflow relationship (point 4).