Add axis inference logic and disable EP flags when EP is not active

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

@@ -16,7 +16,6 @@
 # parallelism in core MoE part (between EP all2all).
 mesh_axes: ['data', 'stage', 'fsdp', 'context', 'expert']
 data_sharding: [['data', 'stage', 'fsdp', 'context', 'expert']]
-context_sharding: 'context'
 logical_axis_rules: [
                       # ==========================================
                       # Vocabulary Embedding

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

@@ -16,7 +16,6 @@
 # components except core dMoE part (between EP all2all).
 mesh_axes: ['data', 'stage', 'fsdp', 'expert']
 data_sharding: [['data', 'stage', 'fsdp', 'expert']]
-context_sharding: 'expert'
 logical_axis_rules: [
                       # ==========================================
                       # Vocabulary Embedding

src/maxtext/configs/types.py

@@ -2515,6 +2515,67 @@ class DerivedValues(BaseModel):
 # ----------------------------------------------------------------------------
 # Helper Functions
 # ----------------------------------------------------------------------------
+
+
+def _normalize_axes(axes: Any) -> tuple[str, ...]:
+  """Normalize a logical-rule mapping value to a tuple of axis name strings.
+
+  Args:
+    axes: The right-hand side of a logical axis rule entry.  Can be a single
+      string, a list/tuple of strings, or an empty list.
+
+  Returns:
+    A (possibly empty) tuple of physical axis name strings.
+  """
+  if axes is None:
+    return ()
+  if isinstance(axes, str):
+    return (axes,)
+  if isinstance(axes, (list, tuple)):
+    return tuple(a for a in axes if isinstance(a, str))
+  return ()
+
+
+def infer_cp_axes(logical_axis_rules: list) -> tuple[str, ...]:
+  """Infer which physical mesh axis/axes serve as Context Parallelism (CP).
+
+  Scans *logical_axis_rules* for the ``activation_length`` logical axis and
+  returns the physical axis/axes it is mapped to.
+
+  Args:
+    logical_axis_rules: The list of ``[logical_name, physical_axes]`` pairs
+      (the ``logical_axis_rules`` config field).
+
+  Returns:
+    A tuple of physical axis name strings that act as CP.  Empty if the
+    ``activation_length`` logical axis is not found in the rules.
+  """
+  for rule in logical_axis_rules:
+    if rule and len(rule) >= 2 and rule[0] == "activation_length":
+      return _normalize_axes(rule[1])
+  return ()
+
+
+def infer_ep_axes(logical_axis_rules: list) -> tuple[str, ...]:
+  """Infer which physical mesh axis/axes serve as Expert Parallelism (EP).
+
+  Scans *logical_axis_rules* for the ``exp`` logical axis and returns the
+  physical axis/axes it is mapped to.
+
+  Args:
+    logical_axis_rules: The list of ``[logical_name, physical_axes]`` pairs
+      (the ``logical_axis_rules`` config field).
+
+  Returns:
+    A tuple of physical axis name strings that act as EP.  Empty if the
+    ``exp`` logical axis is not found in the rules.
+  """
+  for rule in logical_axis_rules:
+    if rule and len(rule) >= 2 and rule[0] == "exp":
+      return _normalize_axes(rule[1])
+  return ()
+
+
 def get_individual_scales(scale: int) -> tuple[int, int, int, int]:
   """Choose appropriate scales for individual dimensions based on global scale."""
   if scale == 0:
@@ -2735,10 +2796,35 @@ def set_derived_and_validate_values(self) -> "MaxTextConfig":
       mesh_config = self._load_mesh_config_from_yaml(self.custom_mesh_and_rule.value)
 
       # Use setattr to dynamically apply attributes, keeping code compact
-      for field in ("mesh_axes", "logical_axis_rules", "data_sharding", "context_sharding"):
+      for field in ("mesh_axes", "logical_axis_rules", "data_sharding"):
         if field in mesh_config:
           setattr(self, field, mesh_config[field])
 
+      # Infer context_sharding from logical_axis_rules when using custom mesh rules.
+      # Falls back to the default ("context") when no activation_length rule is found.
+      cp_axes = infer_cp_axes(self.logical_axis_rules)
+      if cp_axes:
+        self.context_sharding = cp_axes[0]
+
+    # Infer EP rank from logical_axis_rules and disable incompatible flags when EP rank > 1.
+    ep_axes = infer_ep_axes(self.logical_axis_rules)
+    ep_rank = 1
+    for axis_name in ep_axes:
+      ici_val = getattr(self, f"ici_{axis_name}_parallelism", 1)
+      dcn_val = getattr(self, f"dcn_{axis_name}_parallelism", 1)
+      ep_rank *= ici_val * max(dcn_val, 1)
+    if ep_rank == 1:
+      _ep_disabled_flags = {
+          "use_random_routing": False,
+          "use_ragged_sort": False,
+          "ragged_buffer_factor": -1.0,
+          "use_ring_of_experts": False,
+      }
+      for flag_name, disabled_value in _ep_disabled_flags.items():
+        current = getattr(self, flag_name)
+        if current != disabled_value:
+          raise ValueError(f"When EP rank is 1, {flag_name} must be {disabled_value} (was {current}).")
+
     # Handle eval custom mesh and rule
     if self.custom_mesh_and_rule_for_eval is CustomRule.DEFAULT:
       # Fallback to primary rule if eval is DEFAULT

tests/unit/moe_test.py

@@ -573,6 +573,7 @@ def test_moe_emb_chunking_random_routing(self):
         weight_dtype="bfloat16",
         megablox=False,
         sparse_matmul=True,
+        ici_expert_parallelism=4,
         use_tokamax_gmm=True,
         use_gmm_v2=True,
         num_moe_emb_chunks=4,
@@ -659,6 +660,7 @@ def test_moe_emb_chunking_gmm_v2(self):
         weight_dtype="bfloat16",
         megablox=False,
         sparse_matmul=True,
+        ici_expert_parallelism=4,
         use_tokamax_gmm=True,
         use_gmm_v2=True,
         num_moe_emb_chunks=4,

tests/unit/pyconfig_test.py

@@ -20,6 +20,7 @@
 
 from maxtext.configs import pyconfig
 from maxtext.configs.pyconfig import resolve_config_path, _CONFIG_FILE_MAPPING, _module_from_path
+from maxtext.configs.types import _normalize_axes, infer_cp_axes, infer_ep_axes
 from maxtext.input_pipeline import data_processing_utils
 from maxtext.utils.globals import MAXTEXT_CONFIGS_DIR, MAXTEXT_PKG_DIR
 from tests.utils.test_helpers import get_test_config_path, get_post_train_test_config_path
@@ -332,6 +333,72 @@ def test_eval_start_step_config(self):
     )
     self.assertEqual(config_override.eval_start_step, 50)
 
+  # ------------------------------------------------------------------
+  # Tests for infer_cp_axes / infer_ep_axes and EP rank flag disabling
+  # ------------------------------------------------------------------
+
+  def test_infer_cp_axes_default(self):
+    """No activation_length rule -> empty tuple (falls back to default 'context')."""
+    self.assertEqual(infer_cp_axes([]), ())
+    self.assertEqual(infer_cp_axes([["exp", "expert"]]), ())
+
+  def test_infer_cp_axes_context(self):
+    """activation_length -> ['context'] (standard case)."""
+    rules = [["activation_length", ["context"]]]
+    self.assertEqual(infer_cp_axes(rules), ("context",))
+
+  def test_infer_cp_axes_expert(self):
+    """activation_length -> ['expert'] (ep-as-cp case)."""
+    rules = [["activation_length", ["expert"]]]
+    self.assertEqual(infer_cp_axes(rules), ("expert",))
+
+  def test_infer_ep_axes_default(self):
+    """No exp rule -> empty tuple."""
+    self.assertEqual(infer_ep_axes([]), ())
+
+  def test_infer_ep_axes_single(self):
+    """exp -> 'expert' (standard case)."""
+    rules = [["exp", "expert"]]
+    self.assertEqual(infer_ep_axes(rules), ("expert",))
+
+  def test_infer_ep_axes_multi(self):
+    """exp -> ['context', 'expert'] (cp-as-ep case)."""
+    rules = [["exp", ["context", "expert"]]]
+    self.assertEqual(infer_ep_axes(rules), ("context", "expert"))
+
+  def test_ep_rank_single_axis(self):
+    """EP rank with single expert axis, ICI=4."""
+    ep_axes = infer_ep_axes([["exp", "expert"]])
+    ep_rank = 1
+    ici = {"expert": 4}
+    for ax in ep_axes:
+      ep_rank *= ici.get(ax, 1)
+    self.assertEqual(ep_rank, 4)
+
+  def test_ep_rank_multi_axis(self):
+    """EP rank with context=2, expert=4 -> rank 8."""
+    ep_axes = infer_ep_axes([["exp", ["context", "expert"]]])
+    ep_rank = 1
+    ici = {"context": 2, "expert": 4}
+    for ax in ep_axes:
+      ep_rank *= ici.get(ax, 1)
+    self.assertEqual(ep_rank, 8)
+
+  def test_ep_rank_1_no_disabling(self):
+    """EP rank=1 -> no flags should be affected."""
+    ep_axes = infer_ep_axes([["exp", "expert"]])
+    ep_rank = 1
+    for _ in ep_axes:
+      ep_rank *= 1  # parallelism = 1
+    self.assertEqual(ep_rank, 1)
+
+  def test_normalize_axes_basics(self):
+    """_normalize_axes handles None, str, list, and empty list."""
+    self.assertEqual(_normalize_axes(None), ())
+    self.assertEqual(_normalize_axes("expert"), ("expert",))
+    self.assertEqual(_normalize_axes(["a", "b"]), ("a", "b"))
+    self.assertEqual(_normalize_axes([]), ())
+
 
 if __name__ == "__main__":
   unittest.main()

tests/unit/train_compile_test.py

@@ -434,6 +434,7 @@ def test_moe_megablox_ring_ep_random(self):
             "use_iota_embed=true",
             "compile_topology_num_slices=1",
             "model_name=deepseek3-test",
+            "ici_expert_parallelism=4",
             "sparse_matmul=True",
             "megablox=True",
             "per_device_batch_size=4",
@@ -578,6 +579,7 @@ def test_moe_emb_chunking(self):
             "use_iota_embed=true",
             "compile_topology_num_slices=1",
             "model_name=deepseek3-test",
+            "ici_expert_parallelism=4",
             "sparse_matmul=True",
             "megablox=False",
             "use_tokamax_gmm=True",
@@ -606,6 +608,7 @@ def test_moe_emb_chunking_with_mlp_bias(self):
             "use_iota_embed=true",
             "compile_topology_num_slices=1",
             "model_name=deepseek3-test",
+            "ici_expert_parallelism=4",
             "sparse_matmul=True",
             "megablox=False",
             "use_tokamax_gmm=True",