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

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:

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.