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

src/maxtext/layers/fsdp_prefetch.py

@@ -56,7 +56,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 = 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 +79,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 = sharding.fsdp_axes_for_dim(axes)
     if not fsdp:
       continue
     if axes[len(axes) - len(fsdp) :] != fsdp:

src/maxtext/utils/sharding.py

@@ -1122,6 +1122,30 @@ def get_formatted_sharding_annotations(params, mesh=None):
 
 FSDP_MESH_AXES = ("fsdp", "fsdp_transpose")
 
+# 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`).
+# Under the `ep-as-dp` custom rule (`custom_mesh_and_rule=ep-as-dp`), non-routed weights are
+# already replicated over `expert`, so `expert` is absent from their PartitionSpec and including
+# it here is a safe no-op, while routed MoE weights have `expert` alone on the `exp` dimension
+# (with no FSDP axis) and keep `expert` sharded.
+FSDP_CO_AXES = 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 over FSDP.
+
+  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 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_mesh_axes_from_partition_spec(pspec, axes_to_remove, dims=None):
   """Return `pspec` with `axes_to_remove` stripped from the given dims.
@@ -1239,6 +1263,35 @@ def _peel(named_sharding):
   return jax.tree.map(_peel, sharding_tree)
 
 
+def remove_fsdp_from_partition_spec(pspec):
+  """Return `pspec` with each dimension's FSDP (and co-sharded) axes removed.
+
+  Unlike `remove_mesh_axes_from_partition_spec(pspec, FSDP_CO_AXES)`, which strips `FSDP_CO_AXES`
+  unconditionally from every dimension (and would strip `'expert'` from the MoE routed expert
+  dimension), this only peels co-sharded axes (`'context'`, `'context_usp_ulysses'`, `'expert'`)
+  on dimensions where they follow an `FSDP_MESH_AXES` axis (`'fsdp'` or `'fsdp_transpose'`).
+  """
+  for i, entry in enumerate(pspec):
+    to_remove = fsdp_axes_for_dim(entry)
+    if to_remove:
+      pspec = remove_mesh_axes_from_partition_spec(pspec, to_remove, dims=(i,))
+  return pspec
+
+
+def remove_fsdp_from_sharding(sharding_tree):
+  """Recursively traverses a sharding tree removing each dimension's FSDP (and co-sharded) axes."""
+
+  def _peel(named_sharding):
+    if isinstance(named_sharding, jax.sharding.NamedSharding):
+      return jax.sharding.NamedSharding(
+          named_sharding.mesh,
+          remove_fsdp_from_partition_spec(named_sharding.spec),
+      )
+    return named_sharding
+
+  return jax.tree.map(_peel, sharding_tree)
+
+
 def remove_expert_from_partition_spec(pspec, dims_to_peel):
   """Return `pspec` with the 'expert' mesh axis removed from the given dim indices.
 

tests/unit/fsdp_prefetch_test.py

@@ -29,6 +29,7 @@
 
 from maxtext.common.common_types import ShardMode
 from maxtext.layers import fsdp_prefetch
+from maxtext.utils import sharding
 
 
 def _abstract_mesh(shape, names):
@@ -68,6 +69,18 @@ def test_per_layer_sharding_drops_layer_axis(self):
     s = fsdp_prefetch._per_layer_sharding(stacked, mesh)  # pylint: disable=protected-access
     self.assertEqual(tuple(s.spec), ("fsdp", "tensor"))
 
+  def test_remove_fsdp_from_sharding(self):
+    mesh = _abstract_mesh((4, 2, 2), ("fsdp", "expert", "tensor"))
+    tree = {
+        "co_sharded": NamedSharding(mesh, PartitionSpec(("fsdp", "expert"), "tensor")),
+        "routed": NamedSharding(mesh, PartitionSpec("expert", "fsdp", None)),
+        "compound": NamedSharding(mesh, PartitionSpec(("expert", "fsdp"), None)),
+    }
+    stripped = sharding.remove_fsdp_from_sharding(tree)
+    self.assertEqual(tuple(stripped["co_sharded"].spec), (None, "tensor"))
+    self.assertEqual(tuple(stripped["routed"].spec), ("expert", None, None))
+    self.assertEqual(tuple(stripped["compound"].spec), ("expert", None))
+
 
 class FsdpPrefetchAllGatherTest(unittest.TestCase):
   """Lowers the FSDP weight all-gather against abstract multi-device meshes (no devices needed)."""
@@ -112,6 +125,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.