From 743dadd23b575aaa270db1eb5babfe901533f559 Mon Sep 17 00:00:00 2001 From: NuojCheng Date: Fri, 24 Jul 2026 18:42:30 +0000 Subject: [PATCH] add axis inference logic and disable EP flags when EP is not active --- .../configs/custom_mesh_and_rule/cp-as-ep.yml | 1 - .../configs/custom_mesh_and_rule/ep-as-cp.yml | 1 - src/maxtext/configs/types.py | 88 ++++++++++++++++++- tests/unit/moe_test.py | 2 + tests/unit/pyconfig_test.py | 67 ++++++++++++++ tests/unit/train_compile_test.py | 3 + 6 files changed, 159 insertions(+), 3 deletions(-) diff --git a/src/maxtext/configs/custom_mesh_and_rule/cp-as-ep.yml b/src/maxtext/configs/custom_mesh_and_rule/cp-as-ep.yml index 99486cc2b9..56786d6e75 100644 --- a/src/maxtext/configs/custom_mesh_and_rule/cp-as-ep.yml +++ b/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 diff --git a/src/maxtext/configs/custom_mesh_and_rule/ep-as-cp.yml b/src/maxtext/configs/custom_mesh_and_rule/ep-as-cp.yml index a7f5e37281..3b8d303c0d 100644 --- a/src/maxtext/configs/custom_mesh_and_rule/ep-as-cp.yml +++ b/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 diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index f9ab879d6d..41543179b0 100644 --- a/src/maxtext/configs/types.py +++ b/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 diff --git a/tests/unit/moe_test.py b/tests/unit/moe_test.py index 8eb3bd70c9..8d30be8248 100644 --- a/tests/unit/moe_test.py +++ b/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, diff --git a/tests/unit/pyconfig_test.py b/tests/unit/pyconfig_test.py index d3afcb612e..67893becb5 100644 --- a/tests/unit/pyconfig_test.py +++ b/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() diff --git a/tests/unit/train_compile_test.py b/tests/unit/train_compile_test.py index 61d1d88dc1..fb736956aa 100644 --- a/tests/unit/train_compile_test.py +++ b/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",