Add mesh rules that replicate dense weights across context (cp-as-ep-dense-no-cp, cp-as-ep-dense-fsdp)

docs/guides/optimization/custom_mesh_and_rule.md

@@ -51,6 +51,12 @@ This rule utilizes the `data`, `stage`, `fsdp`, and `expert` axes. Its defining
 
 Similar in philosophy to `ep-as-cp.yml`, this configuration explicitly includes the `context` axis in the mesh layout alongside `data`, `stage`, `fsdp`, and `expert`. While context sharding is mapped to the `context` axis globally, within MoE components, this `context` axis dynamically shifts to perform expert parallelism instead of FSDP. This custom rule supports using CP and EP together.
 
+### `cp-as-ep-dense-no-cp.yml` and `cp-as-ep-dense-fsdp.yml`
+
+Variants of `cp-as-ep.yml` that take the dense weights (`embed`, `embed_attn`, `q_lora`, `kv_lora`: attention and shared-expert projections, the MoE router and the shared-expert gate) off the `context` axis. Under `cp-as-ep` they are sharded over `fsdp` x `context` x `expert`, so every all-gather of them spans all devices and also synchronizes the context ranks, which do unequal attention work under causal context parallelism. `cp-as-ep-dense-no-cp.yml` shards them over `fsdp` x `expert` and `cp-as-ep-dense-fsdp.yml` over `fsdp` only, so those all-gathers run in independent groups of `fsdp` x `expert` (or `fsdp`) devices, at the cost of replicating the weights and their gradients `context`-way (`context` x `expert`-way).
+
+Both rules also map `embed_grad` to the token axes. Left to GSPMD, the gradient of a weight replicated over `context` (a partial sum over every token axis) is all-reduced at full size over all devices and then sliced. With `embed_grad` mapped, `linears.route_dense_weight_grad` reduce-scatters it over all devices instead, as under `cp-as-ep`, and then all-gathers the shard over the replicated axes. The GatedDeltaNet projections (`embed_gdn`), activations, the vocabulary weights and the routed experts are sharded exactly as in `cp-as-ep.yml`, and `RoutedMoE` keeps `("context", "expert")` as its expert-parallel axes.
+
 ### `ep-as-dp.yml`
 
 Different with the rule in `base.yml`, this rule configures expert physical axis to function as data parallelism rather than FSDP. This removes the constraint where FSDPxEP is limited by specific model dimensions, particularly for small tensors such as attention projections. Ultimately, this change benefits large-scale training.

src/maxtext/common/common_types.py

@@ -203,6 +203,8 @@ class CustomRule(enum.Enum):
   DEFAULT = ""
   PURE_FSDP = "pure-fsdp"
   CP_AS_EP = "cp-as-ep"  # Support CP and EP together
+  CP_AS_EP_DENSE_NO_CP = "cp-as-ep-dense-no-cp"  # cp-as-ep with the dense weights on fsdp x expert, not context
+  CP_AS_EP_DENSE_FSDP = "cp-as-ep-dense-fsdp"  # cp-as-ep with the dense weights on fsdp only
   EP_AS_CP = "ep-as-cp"  # Support EP only
   PIPELINE_LARGE_MOE = "pipeline-large-moe"
   FSDP_2D = "2d-fsdp"

src/maxtext/configs/custom_mesh_and_rule/cp-as-ep-dense-fsdp.yml

@@ -0,0 +1,97 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Variant of cp-as-ep.yml whose only difference is that the dense weights are sharded on fsdp only.
+#
+# As cp-as-ep-dense-no-cp.yml, but the dense weights ('embed', 'embed_attn', 'q_lora', 'kv_lora') are also replicated
+# over the expert axis: their all-gathers run in context x expert independent fsdp groups, and their gradients are
+# reduce-scattered over all devices ('embed_grad') and then all-gathered over context x expert. Twice the HBM of
+# cp-as-ep-dense-no-cp.yml for those weights. The GatedDeltaNet projections ('embed_gdn'), activations, the
+# vocabulary weights and the routed experts are as in cp-as-ep, and RoutedMoE uses ("context", "expert") as its
+# expert-parallel axes.
+#
+# This rule uses data, stage, FSDP, context, and expert. Context axis acts as expert
+# parallelism in core MoE part (between EP all2all).
+mesh_axes: ['data', 'stage', 'fsdp', 'context', 'expert']
+data_sharding: [['data', 'stage', 'fsdp', 'context', 'expert']]
+logical_axis_rules: [
+                      # ==========================================
+                      # Vocabulary Embedding
+                      # ==========================================
+                      # Vocab Activations
+                      ['activation_embed_and_logits_batch', ['data', 'stage', 'fsdp', 'expert']],
+                      # Vocabulary tiling flattens [batch, length] into one dimension. It must also be sharded
+                      # over the length axis (context), or every context rank computes the logits of all the
+                      # context shards. context is last because it is the minor part of the flattened index.
+                      ['activation_embed_and_logits_batch_sequence', ['data', 'stage', 'fsdp', 'expert', 'context']],
+                      # Vocab Weights
+                      ['vocab', []],
+                      ['embed_vocab', ['fsdp', 'context', 'expert']],
+                      # ==========================================
+                      # Attention
+                      # ==========================================
+                      # Attention Activations
+                      ['activation_batch_attn', ['data', 'fsdp', 'expert']],
+                      ['activation_heads', []],
+                      ['activation_kv_heads', []],
+                      ['activation_length_attn', ['context']],
+                      ['activation_q_length', ['context']],
+                      ['activation_kv_length', []],
+                      ['activation_embed_attn', []],
+                      ['activation_kv', []],
+                      ['activation_kv_batch', ['data', 'fsdp', 'expert']],
+                      ['activation_kv_head_dim', []],
+                      # Attention Weights
+                      ['heads', []],
+                      ['q_heads', []],
+                      ['kv_heads', []],
+                      ['qkv', []],
+                      ['kv', []],
+                      ['kv_head_dim', []],
+                      ['q_lora', ['fsdp']],
+                      ["q_lora_up_proj", []],
+                      ['kv_lora', ['fsdp']],
+                      ["kv_lora_up_proj", []],
+                      # ==========================================
+                      # Mixture of Experts (MoE)
+                      # ==========================================
+                      # MoE Activations
+                      ['activation_batch_moe', ['data', 'fsdp']],
+                      ['activation_exp', ['context', 'expert']],
+                      # MoE Weights
+                      ['exp', ['context', 'expert']],
+                      ['embed_moe', ['fsdp']],
+                      ['embed_router', ['fsdp']],
+                      # ==========================================
+                      # Standard MLP / Dense Layers / Model Structure
+                      # ==========================================
+                      # Dense Activations
+                      ['activation_mlp', []],
+                      ['activation_batch', ['data', 'fsdp', 'expert']],
+                      ['activation_length', ['context']],
+                      ['activation_norm_length', ['context']],
+                      ['activation_embed', []],
+                      ['activation_stage', 'stage'],
+                      # General Weights
+                      ['mlp', []],
+                      ['layers', 'stage'],
+                      # Dense weights: not sharded on context, see the header.
+                      ['embed', ['fsdp']],
+                      ['embed_attn', ['fsdp']],
+                      # GatedDeltaNet projections (models/qwen3.py) stay sharded as in cp-as-ep.
+                      ['embed_gdn', ['fsdp', 'context', 'expert']],
+                      # The dense weights' gradients are reduce-scattered over every token axis, as in cp-as-ep, then
+                      # gathered over the axes the weights are replicated on (linears.route_dense_weight_grad).
+                      ['embed_grad', ['fsdp', 'context', 'expert']],
+                  ]

src/maxtext/configs/custom_mesh_and_rule/cp-as-ep-dense-no-cp.yml

@@ -0,0 +1,100 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Variant of cp-as-ep.yml whose only difference is that the dense weights are not sharded on the context axis.
+#
+# cp-as-ep shards the dense weights ('embed', 'embed_attn', 'q_lora', 'kv_lora': attention and shared-expert
+# projections, the MoE router and the shared-expert gate) over fsdp x context x expert, so every all-gather of them
+# spans all devices and also synchronizes the context ranks, which do unequal attention work. Here they are sharded
+# over fsdp x expert and replicated over context: those all-gathers run in context-many independent fsdp x expert
+# groups, at the cost of a context-way replication of the weights and their gradients. Their gradients are still
+# reduce-scattered over all devices, as in cp-as-ep ('embed_grad'), and the shard is then all-gathered over context;
+# left to GSPMD, the gradient would be all-reduced at full size over all devices instead. The GatedDeltaNet
+# projections ('embed_gdn'), activations, the vocabulary weights and the routed experts ('exp' over context x expert)
+# are sharded as in cp-as-ep, and RoutedMoE uses ("context", "expert") as its expert-parallel axes.
+#
+# This rule uses data, stage, FSDP, context, and expert. Context axis acts as expert
+# parallelism in core MoE part (between EP all2all).
+mesh_axes: ['data', 'stage', 'fsdp', 'context', 'expert']
+data_sharding: [['data', 'stage', 'fsdp', 'context', 'expert']]
+logical_axis_rules: [
+                      # ==========================================
+                      # Vocabulary Embedding
+                      # ==========================================
+                      # Vocab Activations
+                      ['activation_embed_and_logits_batch', ['data', 'stage', 'fsdp', 'expert']],
+                      # Vocabulary tiling flattens [batch, length] into one dimension. It must also be sharded
+                      # over the length axis (context), or every context rank computes the logits of all the
+                      # context shards. context is last because it is the minor part of the flattened index.
+                      ['activation_embed_and_logits_batch_sequence', ['data', 'stage', 'fsdp', 'expert', 'context']],
+                      # Vocab Weights
+                      ['vocab', []],
+                      ['embed_vocab', ['fsdp', 'context', 'expert']],
+                      # ==========================================
+                      # Attention
+                      # ==========================================
+                      # Attention Activations
+                      ['activation_batch_attn', ['data', 'fsdp', 'expert']],
+                      ['activation_heads', []],
+                      ['activation_kv_heads', []],
+                      ['activation_length_attn', ['context']],
+                      ['activation_q_length', ['context']],
+                      ['activation_kv_length', []],
+                      ['activation_embed_attn', []],
+                      ['activation_kv', []],
+                      ['activation_kv_batch', ['data', 'fsdp', 'expert']],
+                      ['activation_kv_head_dim', []],
+                      # Attention Weights
+                      ['heads', []],
+                      ['q_heads', []],
+                      ['kv_heads', []],
+                      ['qkv', []],
+                      ['kv', []],
+                      ['kv_head_dim', []],
+                      ['q_lora', ['fsdp', 'expert']],
+                      ["q_lora_up_proj", []],
+                      ['kv_lora', ['fsdp', 'expert']],
+                      ["kv_lora_up_proj", []],
+                      # ==========================================
+                      # Mixture of Experts (MoE)
+                      # ==========================================
+                      # MoE Activations
+                      ['activation_batch_moe', ['data', 'fsdp']],
+                      ['activation_exp', ['context', 'expert']],
+                      # MoE Weights
+                      ['exp', ['context', 'expert']],
+                      ['embed_moe', ['fsdp']],
+                      ['embed_router', ['fsdp']],
+                      # ==========================================
+                      # Standard MLP / Dense Layers / Model Structure
+                      # ==========================================
+                      # Dense Activations
+                      ['activation_mlp', []],
+                      ['activation_batch', ['data', 'fsdp', 'expert']],
+                      ['activation_length', ['context']],
+                      ['activation_norm_length', ['context']],
+                      ['activation_embed', []],
+                      ['activation_stage', 'stage'],
+                      # General Weights
+                      ['mlp', []],
+                      ['layers', 'stage'],
+                      # Dense weights: not sharded on context, see the header.
+                      ['embed', ['fsdp', 'expert']],
+                      ['embed_attn', ['fsdp', 'expert']],
+                      # GatedDeltaNet projections (models/qwen3.py) stay sharded as in cp-as-ep.
+                      ['embed_gdn', ['fsdp', 'context', 'expert']],
+                      # The dense weights' gradients are reduce-scattered over every token axis, as in cp-as-ep, then
+                      # gathered over the axes the weights are replicated on (linears.route_dense_weight_grad).
+                      ['embed_grad', ['fsdp', 'context', 'expert']],
+                  ]

src/maxtext/layers/linears.py

@@ -36,6 +36,8 @@
 from maxtext.layers.quantizations import AqtQuantization as Quant
 from maxtext.utils import max_logging
 from maxtext.utils import max_utils
+from maxtext.utils.sharding import get_logical_axis_rules
+from maxtext.utils.sharding import logical_to_mesh_axes
 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
@@ -143,6 +145,67 @@ def _compute_dot_general_nnx(
   )
 
 
+# Logical axis a custom rule set maps to every mesh axis the tokens are sharded over when it keeps the dense weights
+# off some of them (cp-as-ep-dense-no-cp, cp-as-ep-dense-fsdp); see `route_dense_weight_grad`.
+DENSE_GRAD_AXIS = "embed_grad"
+_DENSE_EMBED_AXES = ("embed", "embed_attn")
+
+
+@functools.partial(jax.custom_vjp, nondiff_argnums=(1, 2))
+def _reshard_cotangent(x, via, to):
+  """Identity whose cotangent is constrained to the PartitionSpec `via` and then to `to`."""
+  del via, to
+  return x
+
+
+def _reshard_cotangent_fwd(x, via, to):
+  del via, to
+  return x, None
+
+
+def _reshard_cotangent_bwd(via, to, _, cotangent):
+  # Without the barrier, XLA:TPU folds the reduce-scatter into `via` and the all-gather into `to` back into one
+  # full-size all-reduce.
+  cotangent = jax.lax.optimization_barrier(jax.lax.with_sharding_constraint(cotangent, via))
+  return (jax.lax.with_sharding_constraint(cotangent, to),)
+
+
+_reshard_cotangent.defvjp(_reshard_cotangent_fwd, _reshard_cotangent_bwd)
+
+
+def route_dense_weight_grad(kernel: Array, kernel_axes: tuple[None | str, ...]) -> Array:
+  """Reduce-scatters a dense weight's gradient over every token axis before resharding it to the weight's sharding.
+
+  A weight's gradient is a partial sum over every mesh axis the tokens are sharded on. When the weight is sharded
+  over only some of them (cp-as-ep-dense-no-cp and cp-as-ep-dense-fsdp replicate the dense weights over context),
+  GSPMD reduces that partial sum with a full-size all-reduce over every device and keeps a slice of it. Under a rule
+  set that maps `embed_grad` to the token axes, the backward of every kernel with an `embed` or `embed_attn`
+  dimension instead constrains the gradient to the kernel's axes with that dimension on `embed_grad` (one
+  reduce-scatter over all of the token axes' devices, as under cp-as-ep), and then, across an optimization barrier,
+  to the kernel's own sharding (an all-gather of that shard over the axes the weight is replicated on). The forward
+  is unchanged, and under every rule set that does not map `embed_grad` the kernel is returned as is.
+
+  Args:
+    kernel: The kernel, as the dot product consumes it.
+    kernel_axes: Its logical axes.
+
+  Returns:
+    The kernel.
+  """
+  rules = get_logical_axis_rules()
+  if not rules or not any(rule[0] == DENSE_GRAD_AXIS for rule in rules):
+    return kernel
+  if not any(isinstance(axis, str) and axis in _DENSE_EMBED_AXES for axis in kernel_axes):
+    return kernel
+  mesh = jax.sharding.get_abstract_mesh()
+  if mesh.empty:
+    raise ValueError(f"The logical axis rules map {DENSE_GRAD_AXIS!r}, which needs a mesh set with jax.set_mesh.")
+  via_axes = tuple(DENSE_GRAD_AXIS if isinstance(a, str) and a in _DENSE_EMBED_AXES else a for a in kernel_axes)
+  via = logical_to_mesh_axes(via_axes, mesh, rules=rules)
+  to = logical_to_mesh_axes(tuple(kernel_axes), mesh, rules=rules)
+  return _reshard_cotangent(kernel, via, to)
+
+
 class DenseGeneral(nnx.Module):
   """A linear transformation with flexible axes."""
 
@@ -444,6 +507,7 @@ def __call__(
       kernel = kernel[..., begin:end]
 
     kernel = self._maybe_two_stage_all_gather(kernel)
+    kernel = route_dense_weight_grad(kernel, self.kernel_axes)
 
     # out_sharding should be None for auto mesh axis
     if self.shard_mode != ShardMode.EXPLICIT:

src/maxtext/layers/moe.py

@@ -888,6 +888,7 @@ def __call__(self, inputs: jax.Array, _initializing: bool = False) -> Tuple[jax.
     else:
       kernel = self.kernel[...]
     kernel = jnp.asarray(kernel, self.dtype)
+    kernel = linears.route_dense_weight_grad(kernel, self.kernel_axes)
 
     contract_ind = tuple(range(0, len(norm_axis)))
     output_sharding = (
@@ -1175,8 +1176,12 @@ def __init__(
       if self.mesh is not None and "dcp" in self.mesh.shape and self.mesh.shape["dcp"] > 1:
         expert_axes.append("dcp")
       self._expert_parallelism_name = tuple(expert_axes)
-    elif self.config.custom_mesh_and_rule == ctypes.CustomRule.CP_AS_EP:
-      # when custom mesh and rule is cp-as-ep, context axis is same with expert in MoE component
+    elif self.config.custom_mesh_and_rule in (
+        ctypes.CustomRule.CP_AS_EP,
+        ctypes.CustomRule.CP_AS_EP_DENSE_NO_CP,
+        ctypes.CustomRule.CP_AS_EP_DENSE_FSDP,
+    ):
+      # when custom mesh and rule is cp-as-ep (or a variant of it), context axis is same with expert in MoE component
       self._expert_parallelism_name = ("context", "expert")
     else:
       self._expert_parallelism_name = "expert"