Merge pull request #4523 from AI-Hypercomputer:chengnuojin-train-eval-mesh

PiperOrigin-RevId: 952183577

Author: Google-ML-Automation

src/maxtext/configs/base.yml

@@ -616,6 +616,9 @@ input_data_sharding_logical_axes: ['activation_embed_and_logits_batch', 'activat
 # Determines which physical axis plays the role of context parallelism for input data processing and load balancing
 # only supports "context" or "expert" (when custom_mesh_and_rule=ep-as-cp)
 context_sharding: "context"
+# Customized mesh and logical rules for evaluation (e.g. "ep-as-cp")
+custom_mesh_and_rule_for_eval: ""
+logical_axis_rules_for_eval: []
 
 # sharding tolerance: float between 0.0 and 1.0 representing the allowed percentage of non-sharded parameters.
 sharding_tolerance: 0.02

src/maxtext/configs/types.py

@@ -1064,6 +1064,10 @@ class HardwareAndMesh(BaseModel):
       CustomRule.DEFAULT,
       description="Customized mesh and logical rules for granularity.",
   )
+  custom_mesh_and_rule_for_eval: CustomRule = Field(
+      CustomRule.DEFAULT,
+      description="Customized mesh and logical rules for evaluation.",
+  )
   allow_split_physical_axes: bool = Field(False, description="Allow splitting physical axes for device mesh creation.")
   enable_nnx: bool = Field(True, description="Whether to use NNX for model definition.")
   optimize_mesh_for_tpu_v6e: bool = Field(False, description="Apply transformations to the mesh for TPU v6e.")
@@ -1080,6 +1084,9 @@ class LayoutAndSharding(BaseModel):
   """Configuration for data and model sharding rules."""
 
   logical_axis_rules: Any = Field([], description="Rules for mapping logical axes to physical mesh axes.")
+  logical_axis_rules_for_eval: Any = Field(
+      [], description="Rules for mapping logical axes to physical mesh axes during evaluation."
+  )
   data_sharding: Any = Field([], description="Sharding for input data.")
   context_sharding: str = Field("context", description="Physical axis name for context parallelism.")
   input_data_sharding_logical_axes: list[str] = Field(
@@ -2697,31 +2704,46 @@ def validate_num_moe_emb_chunks(self):
             f"Got use_gmm_v2={self.use_gmm_v2}, use_ring_of_experts={self.use_ring_of_experts}."
         )
 
+  @staticmethod
+  def _load_mesh_config_from_yaml(rule_value: str) -> dict:
+    """Helper to load and parse custom mesh YAML configurations."""
+    custom_mesh_path = os.path.join(
+        os.path.dirname(os.path.abspath(__file__)),
+        "custom_mesh_and_rule",
+        f"{rule_value}.yml",
+    )
+
+    if not os.path.exists(custom_mesh_path):
+      # ValueError is more semantically correct for validation errors than NotImplementedError
+      raise ValueError(f"Custom mesh config file not found at {custom_mesh_path}")
+
+    # Explicitly setting encoding removes the need for the pylint disable comment
+    with open(custom_mesh_path, "r", encoding="utf-8") as f:
+      return yaml.safe_load(f) or {}
+
   @model_validator(mode="after")
   def set_derived_and_validate_values(self) -> "MaxTextConfig":
     """
     Computes all derived values and runs all cross-field validations after initial parsing.
     This logic is ported from the legacy pyconfig_deprecated.py system and adapted for Pydantic.
     """
+    # Handle primary custom mesh and rule
     if self.custom_mesh_and_rule is not CustomRule.DEFAULT:
-      custom_mesh_path = os.path.join(
-          os.path.dirname(os.path.abspath(__file__)),
-          "custom_mesh_and_rule",
-          f"{self.custom_mesh_and_rule.value}.yml",
-      )
-      if os.path.exists(custom_mesh_path):
-        with open(custom_mesh_path, "r") as f:  # pylint: disable=unspecified-encoding
-          custom_mesh_config = yaml.safe_load(f)
-          if "mesh_axes" in custom_mesh_config:
-            self.mesh_axes = custom_mesh_config["mesh_axes"]
-          if "logical_axis_rules" in custom_mesh_config:
-            self.logical_axis_rules = custom_mesh_config["logical_axis_rules"]
-          if "data_sharding" in custom_mesh_config:
-            self.data_sharding = custom_mesh_config["data_sharding"]
-          if "context_sharding" in custom_mesh_config:
-            self.context_sharding = custom_mesh_config["context_sharding"]
-      else:
-        raise NotImplementedError(f"Custom mesh config file not found at {custom_mesh_path}")
+      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"):
+        if field in mesh_config:
+          setattr(self, field, mesh_config[field])
+
+    # 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
+      self.custom_mesh_and_rule_for_eval = self.custom_mesh_and_rule
+      self.logical_axis_rules_for_eval = self.logical_axis_rules
+    else:
+      eval_config = self._load_mesh_config_from_yaml(self.custom_mesh_and_rule_for_eval.value)
+      self.logical_axis_rules_for_eval = eval_config.get("logical_axis_rules", self.logical_axis_rules)
 
     # A. SET RUN NAME AND PATHS
     # If run_name is not set, generate one from the JOBSET_NAME environment variable (if available)

src/maxtext/layers/attention_op.py

@@ -75,7 +75,7 @@
 from maxtext.layers.initializers import variable_to_logically_partitioned
 from maxtext.layers.quantizations import AqtQuantization as Quant
 from maxtext.utils import max_utils
-from maxtext.utils.sharding import logical_to_mesh_axes, maybe_shard_with_pspec
+from maxtext.utils.sharding import logical_to_mesh_axes, maybe_shard_with_pspec, get_logical_axis_rules_from_config
 import numpy as np
 # pylint: disable=line-too-long, g-doc-args, g-doc-return-or-yield, bad-continuation, g-inconsistent-quotes
 # pytype: disable=attribute-error
@@ -628,7 +628,7 @@ def maybe_create_nnx(einsum, *args):
       self.AqtEinsum_3 = jnp.einsum
 
   def _logical_to_mesh_axes(self, logical_name):
-    logical_rules = None if self.config.using_pipeline_parallelism else self.config.logical_axis_rules
+    logical_rules = get_logical_axis_rules_from_config(self.config)
     return logical_to_mesh_axes(logical_name, mesh=self.mesh, rules=logical_rules)
 
   def check_attention_inputs(self, query: Array, key: Array | KVTensor, value: Array | KVTensor) -> None:

src/maxtext/layers/moe.py

@@ -43,7 +43,7 @@
 from maxtext.utils import max_utils
 from maxtext.utils import maxtext_utils
 from maxtext.utils.sharding import create_sharding, maybe_shard_with_logical, maybe_shard_with_pspec
-from maxtext.utils.sharding import logical_to_mesh_axes, remove_expert_from_partition_spec
+from maxtext.utils.sharding import logical_to_mesh_axes, remove_expert_from_partition_spec, get_logical_axis_rules_from_config
 import numpy as np
 import qwix
 from qwix.contrib.sparsity import sparsity_module
@@ -644,7 +644,7 @@ def _maybe_shard_with_logical(self, inputs, logical_name):
     )
 
   def _logical_to_mesh_axes(self, logical_name):
-    logical_rules = None if self.config.using_pipeline_parallelism else self.config.logical_axis_rules
+    logical_rules = get_logical_axis_rules_from_config(self.config)
     return logical_to_mesh_axes(logical_name, mesh=self.mesh, rules=logical_rules)
 
   def _maybe_shard_with_pspec(self, inputs, pspec: jax.sharding.PartitionSpec | None):

src/maxtext/models/deepseek.py

@@ -42,6 +42,7 @@
 from maxtext.utils import max_utils
 from maxtext.utils.sharding import create_sharding
 from maxtext.utils.sharding import maybe_shard_with_logical
+from maxtext.utils.sharding import get_logical_axis_rules_from_config
 
 import transformers
 
@@ -189,7 +190,7 @@ def with_logical_constraint(self, x):
         shard_mode=self.config.shard_mode,
         debug_sharding=self.config.debug_sharding,
         extra_stack_level=1,
-        rules=self.config.logical_axis_rules,
+        rules=get_logical_axis_rules_from_config(self.config),
     )
 
   def dropout_op(self, x, deterministic):

src/maxtext/models/qwen3.py

@@ -33,7 +33,7 @@
 
 from maxtext.common.common_types import AttentionType, Config, DType, Array, BATCH, EMBED, MODEL_MODE_TRAIN, LENGTH, MODEL_MODE_AUTOREGRESSIVE
 from maxtext.common.common_types import KV_BATCH, KV_HEAD
-from maxtext.utils.sharding import logical_to_mesh_axes
+from maxtext.utils.sharding import logical_to_mesh_axes, get_logical_axis_rules
 from maxtext.layers import attentions
 from maxtext.layers import initializers as max_initializers
 from maxtext.layers import moe
@@ -612,7 +612,7 @@ def __call__(
     # mixed_qkvz: (B, S, H_k, 2*D_k + 2*D_v*V_per_K)
     mixed_qkvz = qkvz.reshape(new_shape_qkvz)
     if self.mesh is not None:
-      logical_rules = None if self.config.using_pipeline_parallelism else self.config.logical_axis_rules
+      logical_rules = get_logical_axis_rules()
       qkvz_pspec = logical_to_mesh_axes((KV_BATCH, None, KV_HEAD, None), mesh=self.mesh, rules=logical_rules)
       qkvz_sharding = jax.sharding.NamedSharding(self.mesh, qkvz_pspec)
       mixed_qkvz = jax.lax.with_sharding_constraint(mixed_qkvz, qkvz_sharding)
@@ -845,7 +845,7 @@ def extract_state(c_in, v_len):
           compute_dtype=cfg.dtype,
       )
     elif self.mesh is not None:
-      logical_rules = self.config.logical_axis_rules
+      logical_rules = get_logical_axis_rules()
       recurrent_state_arg = (
           recurrent_state
           if recurrent_state is not None

src/maxtext/trainers/pre_train/train.py

@@ -667,7 +667,8 @@ def training_loop_iteration(
 
   # Unpack immutable_data
   config = immutable_data["config"]  # for helpers
-  logical_axis_rules = immutable_data["logical_axis_rules"]
+  logical_axis_rules_for_train = immutable_data["logical_axis_rules_for_train"]
+  logical_axis_rules_for_eval = immutable_data["logical_axis_rules_for_eval"]
   shard_optimizer_over_data = immutable_data["shard_optimizer_over_data"]
   shard_mode = immutable_data["shard_mode"]
   eval_interval = immutable_data["eval_interval"]
@@ -696,7 +697,7 @@ def training_loop_iteration(
     else:
       step_rng_args = ()
     with maybe_record_goodput(recorder, GoodputEvent.STEP, step):
-      with jax.set_mesh(mesh), nn_partitioning.axis_rules(logical_axis_rules):
+      with jax.set_mesh(mesh), nn_partitioning.axis_rules(logical_axis_rules_for_train):
         if shard_optimizer_over_data and isinstance(model, nn.Module):
           state = sharding.maybe_shard_with_name(state, state_mesh_shardings, shard_mode)
         state, metrics = p_train_step(state, example_batch, *step_rng_args)
@@ -728,10 +729,12 @@ def training_loop_iteration(
     # pylint: disable=not-callable
     for eval_batch in eval_data_iterator:
       # Shard input eval data
-      eval_batch = jax.device_put(eval_batch, sharding.get_input_data_sharding(config, mesh))
+      eval_batch = jax.device_put(
+          eval_batch, sharding.get_input_data_sharding(config, mesh, rules=config.logical_axis_rules_for_eval)
+      )
       if 0 < eval_steps <= eval_step_count:
         break
-      with jax.set_mesh(mesh), nn_partitioning.axis_rules(logical_axis_rules):
+      with jax.set_mesh(mesh), nn_partitioning.axis_rules(logical_axis_rules_for_eval):
         eval_metrics = p_eval_step(state, eval_batch, *step_rng_args)
       eval_step_time_delta = datetime.datetime.now() - last_eval_step_completion
       last_eval_step_completion = datetime.datetime.now()
@@ -861,7 +864,8 @@ def train_loop(config, recorder, state=None):
 
   immutable_data = {
       "config": config,
-      "logical_axis_rules": config.logical_axis_rules,
+      "logical_axis_rules_for_train": config.logical_axis_rules,
+      "logical_axis_rules_for_eval": config.logical_axis_rules_for_eval,
       "shard_optimizer_over_data": config.shard_optimizer_over_data,
       "shard_mode": config.shard_mode,
       "steps": config.steps,

src/maxtext/utils/sharding.py

@@ -20,7 +20,7 @@
 from pathlib import Path
 
 from flax import linen as nn, nnx
-from flax.core.spmd import get_logical_axis_rules
+from flax.core.spmd import get_logical_axis_rules as flax_get_logical_axis_rules
 import jax
 from jax.core import Tracer
 from jax.sharding import NamedSharding, PartitionSpec as P, reshard
@@ -40,17 +40,33 @@ def clear_input_shardings_dump():
   _ACTIVATION_SHARDINGS_DUMP.clear()
 
 
-def get_input_data_sharding(config, mesh):
+def get_input_data_sharding(config, mesh, rules=None):
   """Get the input data sharding for the model"""
+  if rules is None:
+    rules = config.logical_axis_rules
   if config.enable_diloco:
-    data_sharding = create_sharding(
-        mesh, ["diloco"] + config.input_data_sharding_logical_axes, rules=config.logical_axis_rules
-    )
+    data_sharding = create_sharding(mesh, ["diloco"] + config.input_data_sharding_logical_axes, rules=rules)
   else:
-    data_sharding = create_sharding(mesh, config.input_data_sharding_logical_axes, rules=config.logical_axis_rules)
+    data_sharding = create_sharding(mesh, config.input_data_sharding_logical_axes, rules=rules)
   return data_sharding
 
 
+def get_logical_axis_rules():
+  """Get the logical axis rules for the flax model"""
+  return flax_get_logical_axis_rules()
+
+
+def get_logical_axis_rules_from_config(config):
+  """Get the logical axis rules from the config.
+  When we use pipeline parallelism or in eval step, we may use
+  a different logical axis rule than the one in config.
+  Plan to deprecate (b/536927795) and use get_logical_axis_rules instead.
+  """
+  if config.using_pipeline_parallelism or config.eval_interval != -1:
+    return None
+  return config.logical_axis_rules
+
+
 def _get_sharding_desc(inputs, extra_stack_level):
   """Get the inputs sharding description using inspect module"""
   frame = inspect.currentframe()

src/maxtext/utils/train_utils.py

@@ -191,13 +191,14 @@ def jit_train_and_eval_step(
   if config.enable_diloco:
     train_step_partial = functools.partial(train_step, model, config, state_mesh_shardings, params_shardings)
     train_step = diloco.build_diloco_train_step(config, train_step_partial, mesh=mesh)
-  data_sharding = sharding.get_input_data_sharding(config, mesh)
+  data_sharding_for_train = sharding.get_input_data_sharding(config, mesh, rules=config.logical_axis_rules)
+  data_sharding_for_eval = sharding.get_input_data_sharding(config, mesh, rules=config.logical_axis_rules_for_eval)
   p_train_step = jit_train_step(
-      config, model, state, state_mesh_shardings, data_sharding, train_step, params_shardings, mesh=mesh
+      config, model, state, state_mesh_shardings, data_sharding_for_train, train_step, params_shardings, mesh=mesh
   )
   p_eval_step = None
   if eval_data_iterator:
-    p_eval_step = jit_eval_step(config, model, state_mesh_shardings, data_sharding, eval_step)
+    p_eval_step = jit_eval_step(config, model, state_mesh_shardings, data_sharding_for_eval, eval_step)
 
   return p_train_step, p_eval_step
 

tests/integration/train_tests.py

@@ -677,6 +677,35 @@ def test_gpu_ring_attention_with_packing(self):
       )
     train_main(thd_ring_attention)
 
+  @pytest.mark.integration_test
+  @pytest.mark.tpu_only
+  def test_fractional_eval_batch_size(self):
+    frac_eval = [  # tests Zero-1 optimizer sharding with gradient accumulation
+        None,
+        get_test_config_path(),
+        f"base_output_directory={self._base_output_directory}",
+        "run_name=runner_test",
+        f"dataset_path={self.dataset_path}",
+        "steps=5",
+        "enable_checkpointing=False",
+        "enable_goodput_recording=False",
+        "dataset_type=synthetic",
+        "remat_policy=minimal",
+        "per_device_batch_size=1",
+        "ici_expert_parallelism=-1",
+        "ici_fsdp_parallelism=1",
+        "use_ring_of_experts=True",
+        "model_name=deepseek3-test",
+        "eval_per_device_batch_size=0.25",
+        "eval_interval=3",
+        "eval_steps=1",
+        "custom_mesh_and_rule_for_eval=ep-as-cp",
+        "override_model_config=true",
+        "base_num_decoder_layers=7",
+        rf"tokenizer_path={os.path.join(MAXTEXT_ASSETS_ROOT, 'tokenizers', 'tokenizer.llama2')}",
+    ]
+    train_main(frac_eval)
+
 
 if __name__ == "__main__":
   absltest.main()