add dense_fsdp_use_two_stage_all_gather

Author: aireenmei

src/maxtext/configs/base.yml

@@ -656,6 +656,9 @@ check_vma: False
 # Enable ZeRO-1 optimizer sharding over data axis
 shard_optimizer_over_data: false
 
+# when dense MLP weight matrices are sharded on both fsdp and fsdp-transpose axes, use two separate all-gather calls
+dense_fsdp_use_two_stage_all_gather: false
+
 # Unless explicitly specified, the number of TPU slices is automatically determined. It should only be set for
 # disaggregated reinforcement learning workloads using multiple slices. For ahead of time compilation,
 # you should set compile_toplogy_num_slices, which will in turn set this value. For non-TPU environments this is set to 1.

src/maxtext/configs/types.py

@@ -1114,6 +1114,10 @@ class LayoutAndSharding(BaseModel):
       "with auto sharding, megablox kernel, and EP / FSDP parallelisms.",
   )
   shard_optimizer_over_data: bool = Field(False, description="Enable ZeRO-1 optimizer sharding over the data axis.")
+  dense_fsdp_use_two_stage_all_gather: bool = Field(
+      False,
+      description="Use two separate All-Gather calls for dense MLP weights sharded on both FSDP and FSDP-transpose.",
+  )
   internal_compile: bool = Field(
       False,
       description="Use internal_compile to bypass open-source topology mappings.",

src/maxtext/layers/linears.py

@@ -23,7 +23,7 @@
 import jax.numpy as jnp
 
 from jax import lax
-from jax.sharding import NamedSharding, Mesh
+from jax.sharding import NamedSharding, Mesh, PartitionSpec
 from jax.ad_checkpoint import checkpoint_name
 
 from flax import nnx
@@ -38,6 +38,9 @@
 from maxtext.utils import max_logging
 from maxtext.utils import max_utils
 from maxtext.utils.sharding import maybe_shard_with_logical
+from maxtext.utils.sharding import maybe_shard_with_name
+from maxtext.utils.sharding import get_physical_spec_without_axes
+from maxtext.utils.sharding import FSDP_MESH_AXES
 
 
 def _convert_to_activation_function(fn_or_string: str | Callable[..., Any]) -> Callable[..., Any]:
@@ -121,6 +124,9 @@ def __init__(
       shard_mode: ShardMode = ShardMode.AUTO,
       matmul_precision: str = "default",
       parameter_memory_host_offload: bool = False,
+      mesh: Mesh | None = None,
+      use_two_stage_all_gather: bool = False,
+      debug_sharding: bool = False,
       *,  # Following arguments are keyword-only
       rngs: nnx.Rngs = None,
   ):
@@ -140,6 +146,13 @@ def __init__(
       shard_mode: auto or explicit shard mode.
       matmul_precision: Precision for matrix multiplication.
       parameter_memory_host_offload: Determines whether to offload params to host
+      mesh: Mesh of devices and physical axes, needed for two-stage all-gather.
+      use_two_stage_all_gather: when the kernel is sharded on both the fsdp and
+        fsdp_transpose axes, gather the two axes with two separate all-gather
+        calls (separated by an optimization barrier) to avoid the relayout
+        transpose XLA emits for a single combined 2-axis all-gather.
+      debug_sharding: when True, log the logical/physical sharding of the
+        two-stage all-gather constraints to the sharding dump files.
       rngs: RNG state for initialization in nnx.
     """
     self.in_features_shape = canonicalize_tuple(in_features_shape)
@@ -154,6 +167,9 @@ def __init__(
     self.shard_mode = shard_mode
     self.matmul_precision = matmul_precision
     self.parameter_memory_host_offload = parameter_memory_host_offload
+    self.mesh = mesh
+    self.use_two_stage_all_gather = use_two_stage_all_gather
+    self.debug_sharding = debug_sharding
 
     # Parameter initialization
     kernel_shape = self.in_features_shape + self.out_features_shape
@@ -200,6 +216,39 @@ def quant_dot_general(self) -> nnx_wrappers.ToNNX | None:
       return None
     return getattr(self, self._quant_dot_general_name)
 
+  def _maybe_two_stage_all_gather(self, kernel):
+    """Gather a 2D-FSDP-sharded MLP kernel with two single-axis all-gathers.
+
+    When the kernel is sharded on both the `fsdp` and `fsdp_transpose` mesh axes,
+    a single combined 2-axis all-gather forces XLA to materialize an interleave
+    transpose to fix the layout. Splitting into two single-axis gathers separated
+    by an `optimization_barrier` makes each stage produce a contiguous layout, so
+    no transpose is emitted. Mirrors `moe_fsdp_use_two_stage_all_gather`.
+    """
+    if (
+        not self.use_two_stage_all_gather
+        or self.mesh is None
+        or self.mesh.shape.get("fsdp", 1) <= 1
+        or self.mesh.shape.get("fsdp_transpose", 1) <= 1
+    ):
+      return kernel
+
+    # kernel_axes is a plain tuple of logical names; wrap it so the logical-to-physical
+    # lookup treats it as a single spec rather than a pytree of strings.
+    full_logical = PartitionSpec(*self.kernel_axes)
+    # Stage 1 gathers fsdp_transpose, stage 2 gathers the remaining fsdp.
+    stage1 = get_physical_spec_without_axes(full_logical, self.mesh, ("fsdp_transpose",))
+    stage2 = get_physical_spec_without_axes(full_logical, self.mesh, FSDP_MESH_AXES)
+    if stage1.spec == stage2.spec:
+      # Not sharded on both FSDP axes, so a single all-gather is already optimal.
+      return kernel
+
+    shard = functools.partial(maybe_shard_with_name, shard_mode=self.shard_mode, debug_sharding=self.debug_sharding)
+    kernel = shard(kernel, stage1)
+    kernel = jax.lax.optimization_barrier(kernel)
+    kernel = shard(kernel, stage2)
+    return kernel
+
   def __call__(self, inputs: Array, _initializing: bool = False, out_sharding: NamedSharding | None = None) -> Array:
     """Applies a linear transformation to the inputs along multiple dimensions.
 
@@ -230,6 +279,8 @@ def __call__(self, inputs: Array, _initializing: bool = False, out_sharding: Nam
         kernel = jax.device_put(kernel, max_utils.device_space())
       kernel = jnp.asarray(kernel, self.dtype)
 
+    kernel = self._maybe_two_stage_all_gather(kernel)
+
     # out_sharding should be None for auto mesh axis
     if self.shard_mode != ShardMode.EXPLICIT:
       out_sharding = None
@@ -422,6 +473,9 @@ def __init__(
           use_bias=self.use_bias,
           shard_mode=self.config.shard_mode,
           matmul_precision=self.config.matmul_precision,
+          mesh=self.mesh,
+          use_two_stage_all_gather=self.config.dense_fsdp_use_two_stage_all_gather,
+          debug_sharding=self.config.debug_sharding,
           rngs=rngs,
       )
     else:
@@ -438,6 +492,9 @@ def __init__(
             use_bias=self.use_bias,
             shard_mode=self.config.shard_mode,
             matmul_precision=self.config.matmul_precision,
+            mesh=self.mesh,
+            use_two_stage_all_gather=self.config.dense_fsdp_use_two_stage_all_gather,
+            debug_sharding=self.config.debug_sharding,
             rngs=rngs,
         )
         setattr(self, dense_name, module)
@@ -453,6 +510,9 @@ def __init__(
         use_bias=self.use_bias,
         shard_mode=self.config.shard_mode,
         matmul_precision=self.config.matmul_precision,
+        mesh=self.mesh,
+        use_two_stage_all_gather=self.config.dense_fsdp_use_two_stage_all_gather,
+        debug_sharding=self.config.debug_sharding,
         rngs=rngs,
     )
 

src/maxtext/utils/sharding.py

@@ -822,34 +822,50 @@ def get_formatted_sharding_annotations(params, mesh=None):
   return "\n".join(annotation_lines)
 
 
-def remove_fsdp_sharding(sharding_tree):
-  """Recursively traverses the sharding tree to remove fsdp axes."""
+FSDP_MESH_AXES = ("fsdp", "fsdp_transpose")
 
-  def _remove_fsdp_from_partition_spec(named_sharding):
-    """Removes 'fsdp' and 'fsdp_transpose' from a PartitionSpec."""
+
+def remove_mesh_axes_from_partition_spec(pspec, axes_to_remove, dims=None):
+  """Return `pspec` with `axes_to_remove` stripped from the given dims.
+
+  Replacing a mesh axis with `None` in a PartitionSpec instructs JAX to replicate
+  the array data along that physical mesh dimension, so peeling an axis here is
+  what turns a sharding constraint into an all-gather.
+
+  Args:
+    pspec: The PartitionSpec to strip axes from.
+    axes_to_remove: Collection of physical mesh axis names to remove.
+    dims: Dim indices to modify; `None` (the default) modifies every dim.
+
+  Returns:
+    A new PartitionSpec with the requested axes replaced by `None`.
+  """
+  axes_to_remove = set(axes_to_remove)
+  new_spec = []
+  for i, axis in enumerate(pspec):
+    if axis is None or (dims is not None and i not in dims):
+      new_spec.append(axis)
+    elif isinstance(axis, str):
+      new_spec.append(None if axis in axes_to_remove else axis)
+    elif isinstance(axis, (list, tuple)):
+      new_spec.append(tuple(a for a in axis if a not in axes_to_remove) or None)
+    else:
+      raise ValueError(f"Unsupported axis type: {type(axis)}")
+  return jax.sharding.PartitionSpec(*new_spec)
+
+
+def remove_mesh_axes_from_sharding(sharding_tree, axes_to_remove):
+  """Recursively traverses a sharding tree removing `axes_to_remove` from each spec."""
+
+  def _peel(named_sharding):
     if isinstance(named_sharding, jax.sharding.NamedSharding):
-      new_spec = []
-      # Iterate through each axis in the original PartitionSpec.
-      for axis in named_sharding.spec:
-        if axis is None:
-          new_spec.append(None)
-        elif isinstance(axis, str):
-          # If the axis is 'fsdp', replace it with None to signify replication.
-          if axis not in ("fsdp", "fsdp_transpose"):
-            new_spec.append(axis)
-          else:
-            new_spec.append(None)
-        elif isinstance(axis, (list, tuple)):
-          # If the axis is a collection, filter out 'fsdp'.
-          new_axis = [a for a in axis if a not in ("fsdp", "fsdp_transpose")]
-          new_spec.append(tuple(new_axis))
-        else:
-          raise ValueError(f"Unsupported_axis_type: {type(axis)}")
-        # Return a new sharding object with the modified spec.
-      return jax.sharding.NamedSharding(named_sharding.mesh, jax.sharding.PartitionSpec(*new_spec))
+      return jax.sharding.NamedSharding(
+          named_sharding.mesh,
+          remove_mesh_axes_from_partition_spec(named_sharding.spec, axes_to_remove),
+      )
     return named_sharding
 
-  return jax.tree.map(_remove_fsdp_from_partition_spec, sharding_tree)
+  return jax.tree.map(_peel, sharding_tree)
 
 
 def remove_expert_from_partition_spec(pspec, dims_to_peel):
@@ -863,19 +879,31 @@ def remove_expert_from_partition_spec(pspec, dims_to_peel):
   untouched. Avoids needing a separate `activation_batch_no_exp` logical rule that every
   `custom_mesh_and_rule` set would have to redefine.
   """
-  new_spec = list(pspec)
-  for i in dims_to_peel:
-    axis = new_spec[i]
-    if axis is None:
-      continue
-    if isinstance(axis, str):
-      new_spec[i] = None if axis == "expert" else axis
-    elif isinstance(axis, (list, tuple)):
-      filtered = tuple(a for a in axis if a != "expert")
-      new_spec[i] = filtered or None
-    else:
-      raise ValueError(f"Unsupported axis type: {type(axis)}")
-  return jax.sharding.PartitionSpec(*new_spec)
+  return remove_mesh_axes_from_partition_spec(pspec, ("expert",), dims=dims_to_peel)
+
+
+def get_physical_spec_without_axes(full_logical, mesh, axes_to_remove, logical_axis_rules=None):
+  """Resolve `full_logical` to a physical sharding with `axes_to_remove` peeled off.
+
+  Combines the logical-to-physical lookup with an axis peel, producing a target
+  layout for an all-gather over exactly the named mesh axes. Peeling a subset of
+  the FSDP axes (rather than all of them at once) is what lets a 2D-FSDP-sharded
+  weight be gathered in two separate single-axis stages.
+
+  Args:
+    full_logical: A PyTree of logical PartitionSpecs. Note that a bare tuple of
+      logical names is a pytree of strings, not a leaf -- wrap it in a
+      `PartitionSpec` before passing it in.
+    mesh: The JAX device mesh.
+    axes_to_remove: Collection of physical mesh axis names to peel.
+    logical_axis_rules: Rules for converting logical axes to physical mesh axes.
+      Defaults to the ambient rules context.
+
+  Returns:
+    A PyTree of physical `jax.sharding.NamedSharding` objects.
+  """
+  physical = logical_to_mesh_sharding(full_logical, mesh=mesh, rules=logical_axis_rules)
+  return remove_mesh_axes_from_sharding(physical, axes_to_remove)
 
 
 def get_physical_spec_no_fsdp(full_logical, mesh, logical_axis_rules):
@@ -902,11 +930,7 @@ def get_physical_spec_no_fsdp(full_logical, mesh, logical_axis_rules):
     mesh axis.
   """
 
-  # Convert the high-level logical spec to a physical one using default rules.
-  physical = logical_to_mesh_sharding(full_logical, mesh=mesh, rules=logical_axis_rules)
-  # Apply the function to remove the FSDP sharding, defining our target layout.
-  physical_no_fsdp = remove_fsdp_sharding(physical)
-  return physical_no_fsdp
+  return get_physical_spec_without_axes(full_logical, mesh, FSDP_MESH_AXES, logical_axis_rules)
 
 
 def all_gather_over_fsdp(variables, sharding_info, mesh, logical_axis_rules, shard_mode):

tests/unit/linears_test.py

@@ -17,9 +17,11 @@
 import sys
 import unittest
 from flax import nnx
+from flax.linen import partitioning as nn_partitioning
 import jax
 import jax.numpy as jnp
 import numpy as np
+import pytest
 
 from maxtext.layers import linears
 from maxtext.configs import pyconfig
@@ -199,6 +201,67 @@ def test_fused_mlp(self):
     self.assertEqual(outputs.shape, (batch_size, seq_len, in_features))
     self.assertEqual(layer.wi.kernel[...].shape, (in_features, 1, intermediate_dim))
 
+  @pytest.mark.tpu_only
+  def test_dense_fsdp_two_stage_all_gather_tpu_only(self):
+    """`dense_fsdp_use_two_stage_all_gather` is a pure resharding and must not change MLP numerics.
+
+    The flag replaces the single combined 2-axis (fsdp x fsdp_transpose) kernel all-gather with two
+    single-axis gathers separated by an `optimization_barrier`. That changes only the weight's layout,
+    not its values, so on a 2D-FSDP mesh the MLP output must match the default single-all-gather path.
+    Covers both the non-fused (2D `wi`/`wo` kernels) and fused (3D `wi` kernel) MLP paths.
+    """
+    # Imperative skip: calling jax.device_count() in a decorator would force JAX init during collection.
+    if jax.device_count() != 4:
+      self.skipTest("Dense FSDP two-stage all-gather test requires exactly 4 devices")
+
+    in_features = 128
+    intermediate_dim = 256
+    batch_size = 4
+    seq_len = 8
+
+    def build_and_run(use_two_stage, fused_mlp):
+      config_arguments = {
+          "per_device_batch_size": 1.0,
+          "run_name": "test",
+          "enable_checkpointing": False,
+          "max_target_length": seq_len,
+          "dtype": "bfloat16",
+          "fused_mlp": fused_mlp,
+          "ici_fsdp_parallelism": 2,
+          "ici_fsdp_transpose_parallelism": 2,
+          "dense_fsdp_use_two_stage_all_gather": use_two_stage,
+      }
+      argv = [sys.argv[0], get_test_config_path()]
+      cfg = pyconfig.initialize(argv, **config_arguments)
+
+      devices_array = maxtext_utils.create_device_mesh(cfg)
+      mesh = jax.sharding.Mesh(devices_array, cfg.mesh_axes)
+      # Fresh, fixed-seed Rngs yields deterministic, identical weights across both builds.
+      with mesh, nn_partitioning.axis_rules(cfg.logical_axis_rules):
+        layer = linears.MlpBlock(
+            config=cfg,
+            mesh=mesh,
+            in_features=in_features,
+            intermediate_dim=intermediate_dim,
+            rngs=nnx.Rngs(params=0, dropout=1),
+        )
+        inputs = jax.random.uniform(jax.random.PRNGKey(1), (batch_size, seq_len, in_features), dtype=cfg.dtype)
+        return nnx.jit(lambda m, x: m(x))(layer, inputs)
+
+    for fused_mlp in (False, True):
+      with self.subTest(fused_mlp=fused_mlp):
+        reference = build_and_run(use_two_stage=False, fused_mlp=fused_mlp)
+        two_stage = build_and_run(use_two_stage=True, fused_mlp=fused_mlp)
+
+        self.assertEqual(two_stage.shape, (batch_size, seq_len, in_features))
+        self.assertTrue(
+            jnp.allclose(reference, two_stage, rtol=1e-2, atol=1e-2, equal_nan=False),
+            msg=(
+                f"two-stage all-gather changed MLP output (fused_mlp={fused_mlp}): max abs diff="
+                f"{jnp.max(jnp.abs(reference.astype(jnp.float32) - two_stage.astype(jnp.float32)))}"
+            ),
+        )
+
 
 if __name__ == "__main__":
   unittest.main()