Gather expert/context co-sharded axes in FSDP weight prefetching (EP support)

[darisoy]

Description

FSDP weight prefetching (prefetch_fsdp_weights=True) did not prefetch most weights under expert parallelism.

With ici_expert_parallelism > 1, the non-routed weights (MLA attention, dense MLP, shared experts) are sharded over ('fsdp', ..., 'expert') on their embed-like dims (embed, embed_attn, q_lora, kv_lora in base.yml). _explicit_fsdp_all_gather only looked for fsdp / fsdp_transpose, so it required fsdp to be the minor-most axis of the dimension, found expert after it, and returned None. That made it fall back to an untagged reshard over fsdp, and XLA then added a second, un-prefetched all-gather over expert right before each consumer.

This PR treats expert, context and context_usp_ulysses as FSDP co-sharding axes when they come after an FSDP axis on the same dimension:

• _fsdp_axes_for_dim(entry): if a dimension has an FSDP axis, returns every co-sharding axis from the first FSDP axis onward.
• _remove_fsdp_from_sharding() strips exactly those axes to get the gathered sharding, and _explicit_fsdp_all_gather() gathers them as one xla_explicit_fsdp-tagged all-gather (e.g. 128-way over ('fsdp','expert')).
• Dimensions with no FSDP axis (the routed-expert dim 'expert') and axes before the first FSDP axis (e.g. 'expert' in ('expert','fsdp')) keep their sharding, so routed experts stay EP-sharded and only gather fsdp.

FSDP-only (expert=1) is unchanged: size-1 axes are dropped before the gather axes are computed, and an EP=1 AoT compile gives the same set of tagged all-gathers (146 total / 144 tagged, identical shapes, axes and dims) before and after this change.

Tests

Unit tests: 3 new tests in tests/unit/fsdp_prefetch_test.py: ('fsdp','expert') co-sharded dim gathered together, routed-expert 'expert' dim preserved, and leading 'expert' in a compound ('expert','fsdp') dim preserved. All 11 tests pass.

AoT compile (DSv3 15L/128E, gf=4x4x4, fsdp=32 x expert=4): tagged FSDP all-gathers go from 32 to 144 (every decoder weight gather), and the 132 un-prefetched in-layer expert gathers are gone.

End-to-end, GhostFish 4x4x4: DSv3 15L, 256 experts, ici_expert_parallelism=4, per_device_batch_size=1, top-k 4, remat_policy=custom; prefetch runs use --xla_tpu_enable_fsdp_latency_hiding_scheduler=true --xla_tpu_explicit_fsdp_dedicated_sparse_core_id=0:

	Baseline (no prefetch)	Prefetch (this PR)
Workload	xid/301769330	xid/301769172
Profile	XProf	XProf
Mean step time	1989.0 ms	1757.6 ms (-11.6%)
Exposed FSDP all-gather wait	384 ms/step	91 ms/step
Peak HBM	23.4 GiB	42.0 GiB

The dedicated SparseCore is needed with EP. Without xla_tpu_explicit_fsdp_dedicated_sparse_core_id, the prefetched gathers queue behind the layer's ragged-all-to-all and other SparseCore work, and the step time is 1955.7 ms (-1.7%) (xid/301803674).

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 and provided workload links above if applicable.
• I have added or updated tests accordingly.
• I have updated documentation accordingly.
• I have included unit tests for new features or bug fixes.
• I have added unit test and integration test coverage for all modified code paths.

[gemini-code-assist]

Code Review

This pull request introduces support for gathering co-sharded axes (such as 'expert' and 'context') alongside FSDP axes during FSDP prefetching in MaxText. It adds helper functions _fsdp_axes_for_dim and _remove_fsdp_from_sharding to identify and strip these co-sharded axes from the sharding specifications, and updates the FSDP all-gather logic accordingly. Additionally, comprehensive unit tests have been added to verify the behavior for co-sharded, routed, and compound dimensions. No review comments were provided, and the implementation looks solid, so there is no additional feedback.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!