[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 (sharding.FSDP_CO_AXES) when they come after an FSDP axis on the same dimension:

• sharding.fsdp_axes_for_dim(entry): if a dimension has an FSDP axis, returns every co-sharding axis from the first FSDP axis onward.
• sharding.remove_fsdp_from_partition_spec() / sharding.remove_fsdp_from_sharding() strips those axes per dimension 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. Under custom_mesh_and_rule=ep-as-dp, non-routed weights are already replicated over 'expert' (so 'expert' is absent from their PartitionSpec), and routed MoE weights keep 'expert' sharded.

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: 4 new tests in tests/unit/fsdp_prefetch_test.py: remove_fsdp_from_sharding, ('fsdp','expert') co-sharded dim gathered together, routed-expert 'expert' dim preserved, and leading 'expert' in a compound ('expert','fsdp') dim preserved. All 12 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 --xla_tpu_rerun_latency_hiding_scheduler_post_sc_assignment=false:

	Baseline (no prefetch)	Prefetch (EP_AS_FSDP, base.yml)	Prefetch (EP_AS_DP, ep-as-dp.yml)	Prefetch (use_ring_of_experts=True)
Workload	xid/301769330	xid/302206654	xid/302234879	xid/302288817
Profile	XProf	XProf	XProf	XProf
Mean step time	1989.0 ms	1743.0 ms (-12.4%)	1721.7 ms (-13.4%)	1683.9 ms (-15.3%)
Exposed FSDP all-gather wait	384 ms/step	16 ms/step	23 ms/step	23 ms/step
Peak HBM	23.4 GiB	39.9 GiB	40.0 GiB	38.8 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 (or Ring-of-Experts all-gather + reduce-scatter) 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

❌ Patch coverage is 90.47619% with 2 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/utils/sharding.py	89.47%	1 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!