From 71fc655ea2849c76ff3529b350120e739dfff7b7 Mon Sep 17 00:00:00 2001 From: Doruk Arisoy <34922984+darisoy@users.noreply.github.com> Date: Mon, 19 Oct 2026 12:56:54 -0700 Subject: [PATCH] Gather expert/context co-sharded FSDP axes in FSDP weight prefetching With expert parallelism, non-routed weights (attention, dense MLP, shared experts) are sharded over ('fsdp', ..., 'expert') on their embed-like dims. The explicit FSDP all-gather only gathered 'fsdp', so it bailed out to an untagged reshard and XLA inserted a second, un-prefetched gather over 'expert' before each consumer. Gather every FSDP co-sharding axis ('fsdp', 'fsdp_transpose', 'context', 'context_usp_ulysses', 'expert') from the first FSDP axis of a dimension onward as one tagged xla_explicit_fsdp gather. Dimensions without an FSDP axis (the routed-expert dim) and axes before the first FSDP axis keep their sharding. --- src/maxtext/layers/fsdp_prefetch.py | 39 +++++++++++++++++++++++++++-- tests/unit/fsdp_prefetch_test.py | 39 +++++++++++++++++++++++++++++ 2 files changed, 76 insertions(+), 2 deletions(-) diff --git a/src/maxtext/layers/fsdp_prefetch.py b/src/maxtext/layers/fsdp_prefetch.py index 75d34b67e6..04570bc7bc 100644 --- a/src/maxtext/layers/fsdp_prefetch.py +++ b/src/maxtext/layers/fsdp_prefetch.py @@ -43,6 +43,41 @@ from maxtext.utils import max_logging, maxtext_utils_nnx, sharding +# Mesh axes that act as FSDP weight-sharding axes when they follow `fsdp` or `fsdp_transpose` on +# the same weight dimension (e.g. `embed` / `embed_attn` / `q_lora` / `kv_lora` in `base.yml`, +# where `expert` and `context` co-shard non-routed weights alongside `fsdp`). +_FSDP_CO_AXES = sharding.FSDP_MESH_AXES + ("context", "context_usp_ulysses", "expert") + + +def _fsdp_axes_for_dim(entry): + """Returns the mesh axes in a PartitionSpec entry that should be gathered by FSDP prefetch. + + If the dimension contains at least one axis in `FSDP_MESH_AXES`, returns all axes in + `_FSDP_CO_AXES` from the first FSDP axis onward. Dimensions without an FSDP axis (such as the + MoE routed expert dimension `'expert'`) and axes preceding the first FSDP axis (such as + `'expert'` in `('expert', 'fsdp')`) are left untouched. + """ + axes = () if entry is None else (tuple(entry) if isinstance(entry, (tuple, list)) else (entry,)) + first_fsdp = next((i for i, a in enumerate(axes) if a in sharding.FSDP_MESH_AXES), None) + if first_fsdp is None: + return () + return tuple(a for a in axes[first_fsdp:] if a in _FSDP_CO_AXES) + + +def _remove_fsdp_from_sharding(named_sharding: NamedSharding) -> NamedSharding: + """Returns `named_sharding` with each dimension's FSDP (and co-sharded) axes removed.""" + new_spec = [] + for entry in named_sharding.spec: + to_remove = set(_fsdp_axes_for_dim(entry)) + if not to_remove: + new_spec.append(entry) + elif isinstance(entry, str): + new_spec.append(None) + else: + new_spec.append(tuple(a for a in entry if a not in to_remove) or None) + return NamedSharding(named_sharding.mesh, jax.sharding.PartitionSpec(*new_spec)) + + def _apply_sharding_hint(w, mesh, shard_mode, debug_sharding=False): """All-gathers `w` over the FSDP mesh axes by resharding it to its FSDP-stripped sharding. @@ -56,7 +91,7 @@ def _apply_sharding_hint(w, mesh, shard_mode, debug_sharding=False): # In auto mode, pin the source to its FSDP-sharded layout so the gather is not hoisted onto the # stored parameter (which would materialize the whole gathered stack). w = jax.lax.with_sharding_constraint(w, src_sharding) - gathered_sharding = sharding.remove_mesh_axes_from_sharding(src_sharding, sharding.FSDP_MESH_AXES) + gathered_sharding = _remove_fsdp_from_sharding(src_sharding) out = _explicit_fsdp_all_gather(w, mesh, src_sharding.spec, gathered_sharding.spec) if out is not None: return out @@ -79,7 +114,7 @@ def _explicit_fsdp_all_gather(w, mesh, src_spec, gathered_spec): for i, entry in enumerate(spec): axes = () if entry is None else (tuple(entry) if isinstance(entry, (tuple, list)) else (entry,)) axes = tuple(a for a in axes if mesh.shape[a] > 1) - fsdp = tuple(a for a in axes if a in sharding.FSDP_MESH_AXES) + fsdp = _fsdp_axes_for_dim(axes) if not fsdp: continue if axes[len(axes) - len(fsdp) :] != fsdp: diff --git a/tests/unit/fsdp_prefetch_test.py b/tests/unit/fsdp_prefetch_test.py index 7ee0ba3df7..209b6dd8f3 100644 --- a/tests/unit/fsdp_prefetch_test.py +++ b/tests/unit/fsdp_prefetch_test.py @@ -112,6 +112,45 @@ def test_explicit_mode_gathers_multiple_fsdp_axes_of_one_dim(self): self.assertIn("(tensor<3x2xf32>) -> tensor<3x8xf32>", gathers[0]) self.assertIn('xla_explicit_fsdp = "yes"', gathers[0]) + def test_explicit_mode_gathers_co_sharded_expert_axis_with_fsdp(self): + mesh = _abstract_mesh((4, 2), ("fsdp", "expert")) + # Non-routed weights (attention, dense MLP, shared experts) are sharded over ('fsdp', 'expert'). + out, gathers = self._lower_gather(mesh, PartitionSpec(("fsdp", "expert"), None), (16, 6)) + self.assertEqual(out.shape, (16, 6)) + self.assertEqual(tuple(out.sharding.spec), (None, None)) + self.assertEqual(len(gathers), 1) + self.assertIn("all_gather_dim = 0", gathers[0]) + self.assertIn('name = "fsdp"', gathers[0]) + self.assertIn('name = "expert"', gathers[0]) + self.assertIn("(tensor<2x6xf32>) -> tensor<16x6xf32>", gathers[0]) + self.assertIn('xla_explicit_fsdp = "yes"', gathers[0]) + + def test_explicit_mode_preserves_routed_expert_axis(self): + mesh = _abstract_mesh((4, 2), ("fsdp", "expert")) + # MoE routed expert weights have 'expert' on dim 0 and 'fsdp' on dim 1; only 'fsdp' is gathered. + out, gathers = self._lower_gather(mesh, PartitionSpec("expert", "fsdp", None), (4, 16, 6)) + self.assertEqual(out.shape, (4, 16, 6)) + self.assertEqual(tuple(out.sharding.spec), ("expert", None, None)) + self.assertEqual(len(gathers), 1) + self.assertIn("all_gather_dim = 1", gathers[0]) + self.assertIn('name = "fsdp"', gathers[0]) + self.assertNotIn('name = "expert"', gathers[0]) + self.assertIn("(tensor<2x4x6xf32>) -> tensor<2x16x6xf32>", gathers[0]) + self.assertIn('xla_explicit_fsdp = "yes"', gathers[0]) + + def test_explicit_mode_preserves_leading_expert_axis_on_compound_dim(self): + mesh = _abstract_mesh((2, 4), ("expert", "fsdp")) + # Compound ('expert', 'fsdp') on the expert dim (e.g. moe_flat_fsdp_weights) keeps 'expert' sharded. + out, gathers = self._lower_gather(mesh, PartitionSpec(("expert", "fsdp"), None), (16, 6)) + self.assertEqual(out.shape, (16, 6)) + self.assertEqual(tuple(out.sharding.spec), ("expert", None)) + self.assertEqual(len(gathers), 1) + self.assertIn("all_gather_dim = 0", gathers[0]) + self.assertIn('name = "fsdp"', gathers[0]) + self.assertNotIn('name = "expert"', gathers[0]) + self.assertIn("(tensor<2x6xf32>) -> tensor<8x6xf32>", gathers[0]) + self.assertIn('xla_explicit_fsdp = "yes"', gathers[0]) + def test_non_minor_fsdp_axis_falls_back_to_resharding(self): mesh = _abstract_mesh((2, 4), ("tensor", "fsdp")) # 'fsdp' is not the minor-most axis of dim 0, so a tiled gather would not give the target layout.