Make MoE dispatch/MLP expert-axis batch sharding configurable (fix Mixtral EP throughput)

[gulsumgudukbay]

Description

PR #4007 added 'expert' to activation_batch_moe to fix a DeepSeek MoE throughput regression. That change is applied to the post-dispatch core MoE activations (dispatch_axis/mlp_axis) as well, where the tensor's expert dim E is already
sharded via activation_exp.
For models with few large experts (e.g. Mixtral), mapping the batch dim onto 'expert' too double-maps two tensor dims onto one mesh axis, so GSPMD abandons the expert-local layout and falls back to FSDP-style AllGather+ReduceScatter instead of expert-parallel AllToAll, creatiing a large throughput regression under single-node expert parallelism (ici_expert_parallelism=-1).
This PR adds a config flag moe_dispatch_no_expert_sharding (default false) that selects, for the training dispatch/MLP batch axis only, a new activation_batch_no_exp rule (['data','fsdp','fsdp_transpose'], i.e. without 'expert'). Mixtral-8x7b uses it as true.
This is the per-model knob anticipated in the #4007 review discussion (sharding core MoE activations by the expert physical axes rather than the batch dimension), without changing the default for any other model.

Behavior

• Default false is byte-identical to current main for every model.
• Only mixtral-8x7b opts in; no other config inherits it.
• For mixtral-8x7b, the flag only changes sharding when the expert mesh axis size > 1 (expert parallelism active). When expert is size 1 (TPU/FSDP-primary, or non-expert parallelism GPU) the two axis rules are identical, so those paths are unaffected.

Scope (why only the training dispatch/MLP axes)

• mask_axes is intentionally left on activation_batch_moe: switching it too regresses Mixtral bs=11 (~10,900 -> ~8,400 tok/s/device, train-step all-to-all 5 -> 3).
• The inference-mode dispatch_axis/mlp_axis are left unchanged: the inference path is already expert-parallel with activation_batch_moe (no FSDP fallback) at both prefill (bs=1: all-to-all=290, all-gather=0) and decode (bs=8: all-to-all=193, all-gather=0), so the regression does not occur there. The change is scoped to the training path.

Tests

Measured on 1x MI355X node (8 GPUs), JAX 0.9.1, ici_expert_parallelism=-1, capacity_factor>0:

Model	bs	tok/s/device (before)	tok/s/device (this PR)	train-step all-to-all
Mixtral-8x7b	11	~7,400	~10,900	0 -> 5
Mixtral-8x7b	6	~7,500	~11,400	0 -> 5
DeepSeek-v2-lite-16b	8	~17,800	~17,800 (unchanged)	0 (unchanged)

Mixtral recovers ~47% throughput via restored expert-parallel AllToAll; DeepSeek (flag off by default) is unchanged.

Added a unit regression test test_moe_dispatch_keeps_expert_on_expert_dim (tests/unit/moe_test.py), parametrized over mixtral-8x7b, mixtral-8x22b, and deepseek3-671b: when moe_dispatch_no_expert_sharding is set it asserts the expert dim is sharded by the expert mesh axis and the batch dim is not (guarding the expert-parallel dispatch/MLP sharding). Passing locally.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[codecov]

Codecov Report

❌ Patch coverage is 84.00000% with 4 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/utils/sharding.py	69.23%	2 Missing and 2 partials ⚠️

📢 Thoughts on this report? Let us know!

[i-chaochen]

can we add test into moe_test.py for activation_batch_no_exp and moe_dispatch_batch_axis ? a pure logical check to make sure batch doesn't steal expert

[gulsumgudukbay]

can we add test into moe_test.py for activation_batch_no_exp and moe_dispatch_batch_axis ? a pure logical check to make sure batch doesn't steal expert

you beat me! I just pushed it seconds after this comment

[yeandy]

I tested across many models on MI355 single node. Mixtral regression is resolved without any impact to deepseek-v2-lite.

Seeing slight regression to llama2-7b, not suspecting that this change is the cause.