Merge pull request #3639 from AI-Hypercomputer:chengnuojin-no-exp5

PiperOrigin-RevId: 899212298

Author: Google-ML-Automation

tests/unit/custom_mesh_and_rule_test.py

@@ -114,8 +114,10 @@ def compare_sharding_jsons(json1: dict, model1_name: str, json2: dict, model2_na
 # Requires JAX TPU support to generate the simulated TPU topology.
 @pytest.mark.cpu_only
 @pytest.mark.tpu_backend
-@pytest.mark.parametrize("model_name, topology, num_slice", TEST_CASES)
-def test_sharding_dump_for_model(model_name: str, topology: str, num_slice: str) -> None:
+@pytest.mark.parametrize("model_name, topology, num_slice, custom_mesh_and_rule, overrides", TEST_CASES)
+def test_sharding_dump_for_model(
+    model_name: str, topology: str, num_slice: str, custom_mesh_and_rule: str, overrides: tuple
+) -> None:
   """
   Test sharding configurations from train_compile.get_shaped_inputs.
   This test verifies that the sharding configurations for various models and topologies remain consistent with golden files.
@@ -132,9 +134,16 @@ def test_sharding_dump_for_model(model_name: str, topology: str, num_slice: str)
       "enable_nnx=False",
       "pure_nnx_decoder=False",
   ]
+  if custom_mesh_and_rule:
+    params.append(f"custom_mesh_and_rule={custom_mesh_and_rule}")
+  if overrides:
+    params.extend(overrides)
 
   root_dir = "tests/utils/sharding_info"
-  base_path = os.path.join(root_dir, model_name, topology, f"slice_{num_slice}")
+  rule_name = f"rule_{custom_mesh_and_rule}" if custom_mesh_and_rule else "rule_default"
+  if overrides:
+    rule_name += "_" + "_".join(overrides)
+  base_path = os.path.join(root_dir, model_name, topology, f"slice_{num_slice}", rule_name)
 
   named_json_path = os.path.join(base_path, "named_shardings.json")
   logical_json_path = os.path.join(base_path, "logical_shardings.json")
@@ -206,12 +215,16 @@ def test_sharding_dump_for_model(model_name: str, topology: str, num_slice: str)
 
 @pytest.fixture(
     scope="module",
-    params=[pytest.param(case, id=f"{case[0]}-{case[1]}-{case[2]}") for case in TEST_CASES],
+    params=[pytest.param(case, id=f"{case[0]}-{case[1]}-{case[2]}-{case[3]}-{''.join(case[4])}") for case in TEST_CASES],
 )
 def abstract_state_and_shardings(request):
   """Pytest fixture to set up model, config, and generate abstract state once per test case."""
-  model_name, topology, num_slice = request.param
-  print(f"Testing model: {model_name}, topology: {topology}, num_slices: {num_slice}", flush=True)
+  model_name, topology, num_slice, custom_mesh_and_rule, overrides = request.param
+  print(
+      f"Testing model: {model_name}, topology: {topology}, num_slices: {num_slice}, "
+      "rule: {custom_mesh_and_rule}, overrides: {overrides}",
+      flush=True,
+  )
   params = [
       "/deps/MaxText/tests/unit/sharding_compare_test",
       get_test_config_path(),
@@ -223,6 +236,10 @@ def abstract_state_and_shardings(request):
       "enable_nnx=False",
       "pure_nnx_decoder=False",
   ]
+  if custom_mesh_and_rule:
+    params.append(f"custom_mesh_and_rule={custom_mesh_and_rule}")
+  if overrides:
+    params.extend(overrides)
   config = pyconfig.initialize(params)
   validate_config(config)
 
@@ -245,7 +262,16 @@ def abstract_state_and_shardings(request):
   # Get logical shardings from maxtext_utils
   logical_shardings = maxtext_utils.get_logical_annotations(config, topology_mesh, init_state_fn)
 
-  return model_name, topology, num_slice, abstract_state, state_mesh_shardings, logical_shardings
+  return (
+      model_name,
+      topology,
+      num_slice,
+      custom_mesh_and_rule,
+      overrides,
+      abstract_state,
+      state_mesh_shardings,
+      logical_shardings,
+  )
 
 
 @pytest.mark.cpu_only
@@ -257,9 +283,16 @@ class TestGetAbstractState:
   def test_get_abstract_state_sharding(self, abstract_state_and_shardings):  # pylint: disable=redefined-outer-name
     """Tests that get_abstract_state returns a state with the correct abstract structure and compares sharding."""
 
-    model_name, topology, num_slice, abstract_state, state_mesh_shardings, logical_shardings = (
-        abstract_state_and_shardings
-    )
+    (
+        model_name,
+        topology,
+        num_slice,
+        custom_mesh_and_rule,
+        overrides,
+        abstract_state,
+        state_mesh_shardings,
+        logical_shardings,
+    ) = abstract_state_and_shardings
 
     assert hasattr(abstract_state, "params")
     assert hasattr(abstract_state, "opt_state")
@@ -268,7 +301,10 @@ def test_get_abstract_state_sharding(self, abstract_state_and_shardings):  # pyl
     assert param_leaf.dtype == jnp.float32
 
     root_dir = "tests/utils/sharding_info"  # Or your target directory
-    base_path = os.path.join(root_dir, model_name, topology, f"slice_{num_slice}")
+    rule_name = f"rule_{custom_mesh_and_rule}" if custom_mesh_and_rule else "rule_default"
+    if overrides:
+      rule_name += "_" + "_".join(overrides)
+    base_path = os.path.join(root_dir, model_name, topology, f"slice_{num_slice}", rule_name)
     os.makedirs(base_path, exist_ok=True)  # Ensure directory exists for saving actual
 
     error_messages = []

tests/unit/sharding_compare_test.py

@@ -58,57 +58,68 @@
 flags.DEFINE_string("model_name", None, "Specific model name to dump.")
 flags.DEFINE_string("topology", None, "Specific topology to dump.")
 flags.DEFINE_string("num_slice", None, "Specific number of slices to dump.")
-
-
-def run_single_dump(model_name: str, topology: str, num_slice: str) -> None:
-  """Generate sharding json file for one specific model, topology and slice."""
-  subprocess.run(
-      [
-          "python3",
-          "-m",
-          "tests.utils.sharding_dump",
-          get_test_config_path(),
-          f"compile_topology={topology}",
-          f"compile_topology_num_slices={num_slice}",
-          f"model_name={model_name}",
-          "weight_dtype=float32",
-          "log_config=false",
-          "debug_sharding=true",
-      ],
-      check=True,
-  )
+flags.DEFINE_string("custom_mesh_and_rule", None, "Specific custom_mesh_and_rule to dump.")
+
+
+def run_single_dump(model_name: str, topology: str, num_slice: str, custom_mesh_and_rule: str, overrides: tuple) -> None:
+  """Generate sharding json file for one specific model, topology, slice and rule."""
+  args = [
+      "python3",
+      "-m",
+      "tests.utils.sharding_dump",
+      get_test_config_path(),
+      f"compile_topology={topology}",
+      f"compile_topology_num_slices={num_slice}",
+      f"model_name={model_name}",
+      "weight_dtype=float32",
+      "log_config=false",
+      "debug_sharding=true",
+  ]
+  if custom_mesh_and_rule:
+    args.append(f"custom_mesh_and_rule={custom_mesh_and_rule}")
+  if overrides:
+    args.extend(overrides)
+  subprocess.run(args, check=True)
 
 
 def main(argv: Sequence[str]) -> None:
-  """Generate json files for every combination of model, topology and slices."""
+  """Generate json files for every combination of model, topology, slices and rule."""
   if FLAGS.model_name and FLAGS.topology and FLAGS.num_slice:
-    cases_to_run = [(FLAGS.model_name, FLAGS.topology, FLAGS.num_slice)]
+    cmr = FLAGS.custom_mesh_and_rule if FLAGS.custom_mesh_and_rule is not None else ""
+    # We do not natively support overrides via CLI FLAGS. To test explicit cases,
+    # rely on the predefined TEST_CASES.
+    cases_to_run = [(FLAGS.model_name, FLAGS.topology, FLAGS.num_slice, cmr, ())]
     print(
         "Running specific case from command line: "
-        f"Model={FLAGS.model_name}, Topology={FLAGS.topology}, NumSlice={FLAGS.num_slice}"
+        f"Model={FLAGS.model_name}, Topology={FLAGS.topology}, NumSlice={FLAGS.num_slice}, Rule={cmr}"
     )
-  elif FLAGS.model_name or FLAGS.topology or FLAGS.num_slice:
+  elif FLAGS.model_name or FLAGS.topology or FLAGS.num_slice or FLAGS.custom_mesh_and_rule:
     print("Error: To specify a single test case, --model_name, --topology, and --num_slice must all be provided.")
     return
   else:
     cases_to_run = TEST_CASES
     print(f"Running all {len(TEST_CASES)} predefined test cases.")
 
   total = len(cases_to_run)
-  for i, (model_name, topology, num_slice) in enumerate(cases_to_run):
-    print(f"\n[{i+1}/{total}] Processing: {model_name} | {topology} | Slice {num_slice}")
-
-    base_path = Path(f"{MAXTEXT_REPO_ROOT}/tests/utils/sharding_info/{model_name}/" f"{topology}/slice_{num_slice}/")
+  for i, (model_name, topology, num_slice, custom_mesh_and_rule, overrides) in enumerate(cases_to_run):
+    rule_name = f"rule_{custom_mesh_and_rule}" if custom_mesh_and_rule else "rule_default"
+    if overrides:
+      rule_name += "_" + "_".join(overrides)
+    print(f"\n[{i+1}/{total}] Processing: {model_name} | {topology} | Slice {num_slice} | Rule {rule_name}")
+
+    base_path = Path(
+        f"{MAXTEXT_REPO_ROOT}/tests/utils/sharding_info/{model_name}/" f"{topology}/slice_{num_slice}/{rule_name}/"
+    )
     json_path_named = base_path / "named_shardings.json"
     json_path_logical = base_path / "logical_shardings.json"
 
     if json_path_named.exists() and json_path_logical.exists():
       print("  -> Sharding files already exist. Regenerating to overwrite.")
 
     try:
-      run_single_dump(model_name, topology, str(num_slice))
+      run_single_dump(model_name, topology, str(num_slice), custom_mesh_and_rule, overrides)
     except subprocess.CalledProcessError:
-      print(f"!!! FAILED: {model_name} {topology} {num_slice}")
+      print(f"!!! FAILED: {model_name} {topology} {num_slice} {custom_mesh_and_rule} overrides={overrides}")
 
 
 if __name__ == "__main__":